refactor: make base install registry-only, move benchmark stack to [gpu] extra

.github/workflows/ci.yml

@@ -18,6 +18,41 @@ jobs:
       - run: pip install ruff==0.15.21
       - run: ruff check src tests
 
+  base-install:
+    name: base install (dataset registry only)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+      # The base install must stay lightweight (pyyaml only) and the dataset
+      # registry importable without the benchmark stack — bitmind-subnet
+      # depends on this contract. This job fails if anyone adds an eager
+      # heavy import to the registry path or a heavy dep to base deps.
+      - run: pip install .
+      - run: |
+          python - <<'EOF'
+          from gasbench.dataset.config import load_benchmark_datasets_from_yaml
+          data = load_benchmark_datasets_from_yaml()
+          assert data.get("image") and data.get("video"), "registry came back empty"
+          print({k: len(v) for k, v in data.items()})
+
+          import gasbench
+          try:
+              gasbench.run_benchmark
+              raise SystemExit("benchmark stack importable in base install — heavy dep leaked into base")
+          except ImportError as e:
+              assert "gasbench[gpu]" in str(e)
+
+          import importlib.metadata as im
+          dists = sorted(d.metadata["Name"].lower() for d in im.distributions())
+          heavy = {"torch", "onnxruntime", "onnxruntime-gpu", "transformers", "timm", "modelscope"}
+          leaked = heavy & set(dists)
+          assert not leaked, f"heavy deps in base install: {leaked}"
+          print("base install contract holds:", dists)
+          EOF
+
   typecheck:
     name: mypy (advisory)
     runs-on: ubuntu-latest

README.md

@@ -29,17 +29,21 @@ To learn how to submit your model to **Bittensor Subnet 34** — including the e
 
 ## Installation
 
-GPU support is installed by default:
+To run benchmarks, install the `gpu` extra (on macOS/Windows this
+automatically substitutes CPU onnxruntime via platform markers):
 
 ```bash
 cd gasbench
-pip install -e .
+pip install -e '.[gpu]'
 ```
 
-For CPU-only (no CUDA):
+The base install is intentionally minimal — just the dataset registry
+(`gasbench.dataset.config` and the bundled YAML configs) with no
+torch/onnxruntime. Use it when you only need dataset definitions
+(e.g. bitmind-subnet's validator cache):
 
 ```bash
-pip install -e .[cpu]
+pip install -e .
 ```
 
 ---

pyproject.toml

@@ -9,6 +9,16 @@ description = "GASBench - ML model benchmark evaluation package"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
+    # Base install is intentionally minimal: just the dataset registry
+    # (gasbench.dataset.config + bundled YAML configs). Consumers that only
+    # need dataset definitions (e.g. bitmind-subnet's validator cache) should
+    # depend on plain `gasbench`. Everything needed to actually run
+    # benchmarks lives in the `gpu` extra: `gasbench[gpu]`.
+    "pyyaml==6.0.3",
+]
+
+[project.optional-dependencies]
+gpu = [
     "numpy==2.4.2",
     "pillow==12.1.1",
     "opencv-python-headless==4.13.0.92",
@@ -26,7 +36,6 @@ dependencies = [
     "pyarrow==24.0.0",
     "pandas==3.0.2",
     "requests==2.32.5",
-    "pyyaml==6.0.3",
     "huggingface-hub==1.5.0",
     "datasets==4.8.5",
     "modelscope==1.36.3",
@@ -40,9 +49,8 @@ dependencies = [
     "einops==0.8.2",
     "open-clip-torch==3.3.0",
 ]
-
-[project.optional-dependencies]
 dev = [
+    "gasbench[gpu]",
     "pytest>=8.0",
     "ruff==0.15.21",
     "mypy>=1.10",

src/gasbench/__init__.py

@@ -7,7 +7,7 @@
 
 Example usage:
     from gasbench import run_benchmark
-    
+
     results = await run_benchmark(
         model_path="path/to/model.onnx",
         modality="image",
@@ -17,11 +17,6 @@
 
 from importlib.metadata import PackageNotFoundError, version
 
-from .benchmarks.image_bench import run_image_benchmark
-from .benchmarks.video_bench import run_video_benchmark
-from .benchmark import run_benchmark, print_benchmark_summary, format_benchmark_summary, save_results_to_json
-from .benchmarks.utils.inference import create_inference_session
-
 try:
     __version__ = version("gasbench")
 except PackageNotFoundError:  # source checkout without an installed dist
@@ -30,10 +25,41 @@
 
 __all__ = [
     "run_image_benchmark",
-    "run_video_benchmark", 
+    "run_video_benchmark",
     "run_benchmark",
     "create_inference_session",
     "print_benchmark_summary",
     "format_benchmark_summary",
     "save_results_to_json",
 ]
+
+# Benchmark exports resolve lazily (PEP 562) so that the base install —
+# which carries only the dataset registry and pyyaml — can import
+# gasbench.dataset.config without pulling in torch/onnxruntime. The full
+# benchmark stack requires the `gpu` extra: pip install gasbench[gpu]
+_LAZY_EXPORTS = {
+    "run_image_benchmark": ("gasbench.benchmarks.image_bench", "run_image_benchmark"),
+    "run_video_benchmark": ("gasbench.benchmarks.video_bench", "run_video_benchmark"),
+    "run_benchmark": ("gasbench.benchmark", "run_benchmark"),
+    "print_benchmark_summary": ("gasbench.benchmark", "print_benchmark_summary"),
+    "format_benchmark_summary": ("gasbench.benchmark", "format_benchmark_summary"),
+    "save_results_to_json": ("gasbench.benchmark", "save_results_to_json"),
+    "create_inference_session": ("gasbench.benchmarks.utils.inference", "create_inference_session"),
+}
+
+
+def __getattr__(name):
+    if name in _LAZY_EXPORTS:
+        import importlib
+
+        module_name, attr = _LAZY_EXPORTS[name]
+        try:
+            module = importlib.import_module(module_name)
+        except ImportError as e:
+            raise ImportError(
+                f"gasbench.{name} requires the benchmark dependencies, which are "
+                f"not part of the base install. Install them with: "
+                f"pip install 'gasbench[gpu]' (original error: {e})"
+            ) from e
+        return getattr(module, attr)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/gasbench/cli.py

@@ -8,7 +8,16 @@
 from pathlib import Path
 
 from . import __version__
-from .benchmark import run_benchmark, print_benchmark_summary, save_results_to_json
+
+try:
+    from .benchmark import run_benchmark, print_benchmark_summary, save_results_to_json
+except ImportError:
+    print(
+        "The gasbench CLI requires the benchmark dependencies, which are not "
+        "part of the base install. Install them with: pip install 'gasbench[gpu]'",
+        file=sys.stderr,
+    )
+    sys.exit(1)
 
 
 def add_common_args(parser):

src/gasbench/dataset/__init__.py

@@ -2,7 +2,7 @@
 
 This package contains modules for:
 - Dataset configuration and discovery
-- Dataset iteration and sampling  
+- Dataset iteration and sampling
 - Dataset caching and persistence
 - Dataset downloading and extraction
 """
@@ -19,9 +19,6 @@
     get_benchmark_size,
     discover_benchmark_datasets,
 )
-from .iterator import DatasetIterator, CACHE_MAX_SAMPLES, GASSTATION_CACHE_MAX_SAMPLES
-from .cache import check_dataset_cache, save_sample_to_cache
-from .download import download_and_extract, DatasetAccessError
 
 __all__ = [
     "BenchmarkDatasetConfig",
@@ -38,7 +35,39 @@
     "CACHE_MAX_SAMPLES",
     "GASSTATION_CACHE_MAX_SAMPLES",
     "check_dataset_cache",
-    "save_sample_to_cache", 
+    "save_sample_to_cache",
     "download_and_extract",
     "DatasetAccessError",
 ]
+
+# The config exports above stay eager: they need nothing beyond pyyaml and
+# are what dataset-registry consumers (base install) use. The
+# iterator/cache/download exports below resolve lazily (PEP 562) because
+# their modules import torch / huggingface_hub / datasets, which only exist
+# in the `gpu` extra: pip install gasbench[gpu]
+_LAZY_EXPORTS = {
+    "DatasetIterator": ("gasbench.dataset.iterator", "DatasetIterator"),
+    "CACHE_MAX_SAMPLES": ("gasbench.dataset.iterator", "CACHE_MAX_SAMPLES"),
+    "GASSTATION_CACHE_MAX_SAMPLES": ("gasbench.dataset.iterator", "GASSTATION_CACHE_MAX_SAMPLES"),
+    "check_dataset_cache": ("gasbench.dataset.cache", "check_dataset_cache"),
+    "save_sample_to_cache": ("gasbench.dataset.cache", "save_sample_to_cache"),
+    "download_and_extract": ("gasbench.dataset.download", "download_and_extract"),
+    "DatasetAccessError": ("gasbench.dataset.download", "DatasetAccessError"),
+}
+
+
+def __getattr__(name):
+    if name in _LAZY_EXPORTS:
+        import importlib
+
+        module_name, attr = _LAZY_EXPORTS[name]
+        try:
+            module = importlib.import_module(module_name)
+        except ImportError as e:
+            raise ImportError(
+                f"gasbench.dataset.{name} requires the benchmark dependencies, "
+                f"which are not part of the base install. Install them with: "
+                f"pip install 'gasbench[gpu]' (original error: {e})"
+            ) from e
+        return getattr(module, attr)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")