simplify running benchmarks

.github/workflows/tests.yaml

@@ -99,7 +99,7 @@ jobs:
       env:
         PYTHONUTF8: 1
       run: >
-        pytest ${{ matrix.pytestargs }} -n=logical --timeout=300 --durations=0
+        pytest ${{ matrix.pytestargs }} -n=logical --dist=worksteal --timeout=300 --durations=0
         --cov --cov-report=xml --cov-report=term  --durations-path=./.github/.test_durations
     - name: upload coverage report
       uses: codecov/codecov-action@v3

dvc/testing/benchmarks/conftest.py

@@ -1 +1,3 @@
 from .fixtures import *  # noqa: F403
+
+pytest_plugins = ["dvc.testing.plugin"]

dvc/testing/benchmarks/fixtures.py

@@ -60,8 +60,8 @@ def dvc_venvs():
 
 
 @pytest.fixture(scope="session")
-def dvc_git_repo(tmp_path_factory, bench_config):
-    url = bench_config.dvc_git_repo
+def dvc_repo(tmp_path_factory, bench_config):
+    url = bench_config.dvc_repo
 
     if os.path.isdir(url):
         return url
@@ -73,8 +73,10 @@ def dvc_git_repo(tmp_path_factory, bench_config):
 
 
 @pytest.fixture(scope="session")
-def dvc_bench_git_repo(tmp_path_factory, bench_config):
-    url = bench_config.dvc_bench_git_repo
+def dvc_bench_repo(tmp_path_factory, bench_config):
+    url = bench_config.dvc_bench_repo
+    if url is None:
+        pytest.skip("--dvc-bench-repo is not set")
 
     if os.path.isdir(url):
         return Path(url)
@@ -90,7 +92,7 @@ def make_dvc_bin(
     dvc_rev,
     dvc_venvs,
     make_dvc_venv,
-    dvc_git_repo,
+    dvc_repo,
     bench_config,
     request,
 ):
@@ -102,7 +104,7 @@ def make_dvc_bin(
                 pkg = f"dvc[{bench_config.dvc_install_deps}]"
             else:
                 pkg = "dvc"
-            packages = [f"{pkg} @ git+file://{dvc_git_repo}@{dvc_rev}"]
+            packages = [f"{pkg} @ git+file://{dvc_repo}@{dvc_rev}"]
             try:
                 if version.Version(dvc_rev) < version.Version("3.50.3"):
                     packages.append("pygit2==1.14.1")
@@ -186,14 +188,14 @@ def _pull(repo, *args):
 
 
 @pytest.fixture
-def make_dataset(request, bench_config, tmp_dir, dvc_bench_git_repo):
+def make_dataset(request, bench_config, tmp_dir, dvc_bench_repo):
     def _make_dataset(
         dvcfile=False, files=True, cache=False, commit=False, remote=False
     ):
         from dvc.repo import Repo
 
         path = tmp_dir / "dataset"
-        root = dvc_bench_git_repo
+        root = dvc_bench_repo
         src = root / "data" / bench_config.dataset / "dataset"
         src_dvc = src.with_suffix(".dvc")
 
@@ -252,7 +254,7 @@ def _make_project(url, rev=None):
 @pytest.fixture
 def project(bench_config, monkeypatch, make_project):
     rev = bench_config.project_rev
-    url = bench_config.project_git_repo
+    url = bench_config.project_repo
 
     if os.path.isdir(url):
         path = url

dvc/testing/benchmarks/plugin.py

@@ -1,55 +1,74 @@
+import os
+from dataclasses import dataclass, fields
+from pathlib import Path
+from typing import Optional
+
 DEFAULT_DVC_BIN = "dvc"
-DEFAULT_DVC_GIT_REPO = "https://github.com/iterative/dvc"
-DEFAULT_DVC_BENCH_GIT_REPO = "https://github.com/iterative/dvc-bench"
-DEFAULT_PROJECT_GIT_REPO = "https://github.com/iterative/example-get-started"
+DEFAULT_DVC_REPO = os.fspath(Path(__file__).parents[3])
+DEFAULT_PROJECT_REPO = "https://github.com/iterative/example-get-started"
 
 
 def pytest_report_header(config):
     bconf = config.bench_config
-    return f"dvc-bench: (dataset: {bconf.dataset!r}, revs: {bconf.dvc_revs!r})"
+    return f"dvc-bench: {bconf}"
 
 
 def pytest_generate_tests(metafunc):
-    str_revs = metafunc.config.getoption("--dvc-revs")
-    revs = str_revs.split(",") if str_revs else [None]
+    revs = metafunc.config.getoption("--dvc-revs")
+    if not revs:
+        revs = [None]
     if "dvc_rev" in metafunc.fixturenames:
         metafunc.parametrize("dvc_rev", revs, scope="session")
 
 
+@dataclass
 class DVCBenchConfig:
-    def __init__(self):
-        self.dataset = "small"
-        self.dvc_bin = DEFAULT_DVC_BIN
-        self.dvc_revs = None
-        self.dvc_git_repo = DEFAULT_DVC_GIT_REPO
-        self.dvc_install_deps = None
-        self.dvc_bench_git_repo = DEFAULT_DVC_BENCH_GIT_REPO
-        self.project_rev = None
-        self.project_git_repo = DEFAULT_PROJECT_GIT_REPO
+    dataset: str = "tiny"
+    dvc_repo: str = DEFAULT_DVC_REPO
+    dvc_bench_repo: Optional[str] = None
+    project_repo: str = DEFAULT_PROJECT_REPO
+    project_rev: Optional[str] = None
+    dvc_bin: str = DEFAULT_DVC_BIN
+    dvc_revs: Optional[list[str]] = None
+    dvc_install_deps: Optional[str] = None
+
+    def __repr__(self):
+        args = ", ".join(
+            f"{f.name}={val!r}"
+            for f in fields(self)
+            if (val := getattr(self, f.name)) != f.default
+        )
+        return f"{self.__class__.__name__}({args})"
 
 
 def pytest_configure(config):
     config.addinivalue_line(
         "markers",
         "requires(spec): mark a test to run only on versions that satisfy the spec",
     )
+    config.bench_config = DVCBenchConfig(
+        dataset=config.getoption("--dataset"),
+        dvc_repo=config.getoption("--dvc-repo"),
+        dvc_bench_repo=config.getoption("--dvc-bench-repo"),
+        project_repo=config.getoption("--project-repo"),
+        project_rev=config.getoption("--project-rev"),
+        dvc_bin=config.getoption("--dvc-bin"),
+        dvc_revs=config.getoption("--dvc-revs"),
+        dvc_install_deps=config.getoption("--dvc-install-deps"),
+    )
 
-    config.bench_config = DVCBenchConfig()
-    config.bench_config.dataset = config.getoption("--dataset")
-    config.bench_config.dvc_bin = config.getoption("--dvc-bin")
-    config.bench_config.dvc_revs = config.getoption("--dvc-revs")
-    config.bench_config.dvc_git_repo = config.getoption("--dvc-git-repo")
-    config.bench_config.dvc_install_deps = config.getoption("--dvc-install-deps")
-    config.bench_config.dvc_bench_git_repo = config.getoption("--dvc-bench-git-repo")
-    config.bench_config.project_rev = config.getoption("--project-rev")
-    config.bench_config.project_git_repo = config.getoption("--project-git-repo")
+
+def resolve_path(path):
+    if os.path.isdir(path):
+        return os.path.abspath(path)
+    return path
 
 
 def pytest_addoption(parser):
     parser.addoption(
         "--dataset",
         type=str,
-        default="small",
+        default="tiny",
         help="Dataset name to use in tests (e.g. tiny/small/large/mnist/etc)",
     )
 
@@ -62,36 +81,34 @@ def pytest_addoption(parser):
 
     parser.addoption(
         "--dvc-revs",
-        type=str,
+        type=lambda revs: revs.split(","),
         help=("Comma-separated list of DVC revisions to test (overrides `--dvc-bin`)"),
     )
 
     parser.addoption(
-        "--dvc-git-repo",
-        type=str,
-        default=DEFAULT_DVC_GIT_REPO,
+        "--dvc-repo",
+        type=resolve_path,
+        default=DEFAULT_DVC_REPO,
         help="Path or url to dvc git repo",
     )
 
     parser.addoption(
         "--dvc-install-deps",
         type=str,
-        default="",
         help="Comma-separated list of DVC installation packages",
     )
 
     parser.addoption(
-        "--dvc-bench-git-repo",
-        type=str,
-        default=DEFAULT_DVC_BENCH_GIT_REPO,
+        "--dvc-bench-repo",
+        type=resolve_path,
+        default=None,
         help="Path or url to dvc-bench git repo (for loading benchmark datasets)",
     )
 
     parser.addoption("--project-rev", type=str, help="Project revision to test")
-
     parser.addoption(
-        "--project-git-repo",
-        type=str,
-        default=DEFAULT_PROJECT_GIT_REPO,
+        "--project-repo",
+        type=resolve_path,
+        default=DEFAULT_PROJECT_REPO,
         help="Path or url to dvc project",
     )

dvc/testing/plugin.py

@@ -18,3 +18,9 @@ def pytest_configure(config):
     from .benchmarks.plugin import pytest_configure as bench_configure
 
     bench_configure(config)
+
+
+def pytest_report_header(config):
+    from .benchmarks.plugin import pytest_report_header as bench_report_header
+
+    return bench_report_header(config)

pyproject.toml

@@ -159,7 +159,7 @@ namespaces = false
 write_to = "dvc/_dvc_version.py"
 
 [tool.pytest.ini_options]
-addopts = "-ra --cov-config pyproject.toml --dist worksteal --benchmark-skip"
+addopts = "-ra --cov-config pyproject.toml"
 filterwarnings = [
     "error::ResourceWarning",
     "error::pytest.PytestUnraisableExceptionWarning",

tests/__main__.py

@@ -3,4 +3,4 @@
 
     import pytest
 
-    sys.exit(pytest.main(["-v", "-n=logical", *sys.argv[1:]]))
+    sys.exit(pytest.main(["-v", "-n=logical", "--dist=worksteal", *sys.argv[1:]]))