mate-runtime-cuda 0.1.0__py3-none-any.whl

mate_runtime_cuda/__init__.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from typing import Any
+
+from mate_bench.plugin import PluginManifest
+
+from ._detect import GpuInfo, is_cuda_available, query_gpu
+
+
+class CudaRuntime:
+    name = "cuda"
+    manifest = PluginManifest(requires_mate_bench=">=0.1,<0.2", api_version=1)
+
+    def is_available(self) -> bool:
+        return is_cuda_available()
+
+    def gpu_info(self, index: int = 0) -> dict[str, Any]:
+        info: GpuInfo = query_gpu(index)
+        return {
+            "gpu_vendor": "nvidia",
+            "gpu_name": info.name,
+            "gpu_chip": info.chip,
+            "vram_gb": info.vram_gb,
+            "runtime": info.cuda_version,
+            "driver": info.driver_version,
+            "_gpu_name_raw": info.name_raw,
+            "_gpu_name_known": info.name_known,
+        }

mate_runtime_cuda/_detect.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import re
+import subprocess
+from dataclasses import dataclass
+
+from ._names import chip_from_name, normalize
+
+
+@dataclass
+class GpuInfo:
+    chip: str  # die codename, e.g. "AD102"
+    name: str  # normalized, e.g. "RTX 4090"
+    name_raw: str  # as reported by nvidia-smi, e.g. "NVIDIA GeForce RTX 4090"
+    name_known: bool  # False → unknown model, stored raw + flagged
+    vram_gb: float
+    cuda_version: str  # e.g. "CUDA 12.4"
+    driver_version: str  # e.g. "550.54.15"
+
+
+def is_cuda_available() -> bool:
+    """CUDA is available if nvidia-smi lists at least one GPU."""
+    return bool(_run(["nvidia-smi", "-L"]))
+
+
+def query_gpu(index: int = 0) -> GpuInfo:
+    """Query GPU info for the given device index via nvidia-smi."""
+    rows = _query_smi(
+        ["name", "memory.total", "driver_version", "cuda_version"],
+        units=False,
+    )
+    if not rows:
+        raise RuntimeError("No NVIDIA GPU found. Are the drivers installed?")
+    if index >= len(rows):
+        raise IndexError(f"GPU index {index} out of range ({len(rows)} GPUs found)")
+
+    row = rows[index]
+    raw_name = row[0].strip()
+    vram_mib = _parse_float(row[1])
+    driver = row[2].strip()
+    cuda = row[3].strip()
+
+    name, name_known = normalize(raw_name)
+    chip, chip_known = chip_from_name(name)
+    if not chip_known:
+        chip = name  # fall back to display name as chip key
+
+    return GpuInfo(
+        chip=chip,
+        name=name,
+        name_raw=raw_name,
+        name_known=name_known and chip_known,
+        vram_gb=round(vram_mib / 1024, 1),
+        cuda_version=f"CUDA {cuda}" if cuda else "CUDA unknown",
+        driver_version=driver or "unknown",
+    )
+
+
+# ── internal helpers ──────────────────────────────────────────────────────────
+
+
+def _run(cmd: list[str], timeout: int = 10) -> str:
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
+        return result.stdout
+    except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
+        return ""
+
+
+def _query_smi(fields: list[str], units: bool = False) -> list[list[str]]:
+    """Run nvidia-smi --query-gpu and return parsed rows."""
+    cmd = [
+        "nvidia-smi",
+        f"--query-gpu={','.join(fields)}",
+        "--format=csv,noheader" + ("" if units else ",nounits"),
+    ]
+    out = _run(cmd)
+    if not out.strip():
+        return []
+    return [line.split(",") for line in out.strip().splitlines() if line.strip()]
+
+
+def _parse_float(s: str) -> float:
+    m = re.search(r"[\d.]+", s)
+    return float(m.group()) if m else 0.0

mate_runtime_cuda/_names.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import re
+
+# (substring to match in normalized name, chip codename)
+# Ordered from most specific to least specific.
+_GPU_CHIP_MAP: list[tuple[str, str]] = [
+    # Blackwell — RTX 50xx
+    ("RTX 5090", "GB202"),
+    ("RTX 5080", "GB203"),
+    ("RTX 5070 Ti", "GB203"),
+    ("RTX 5070", "GB205"),
+    ("RTX 5060 Ti", "GB206"),
+    ("RTX 5060", "GB206"),
+    # Ada Lovelace — RTX 40xx + professional
+    ("RTX 6000 Ada", "AD102"),
+    ("RTX 4090", "AD102"),
+    ("RTX 4080 Super", "AD103"),
+    ("RTX 4080", "AD103"),
+    ("RTX 4070 Ti Super", "AD103"),
+    ("RTX 4070 Ti", "AD104"),
+    ("RTX 4070 Super", "AD104"),
+    ("RTX 4070", "AD104"),
+    ("RTX 4500 Ada", "AD104"),
+    ("RTX 4000 Ada", "AD104"),
+    ("RTX 4060 Ti", "AD106"),
+    ("RTX 4060", "AD107"),
+    # Hopper
+    ("H200", "GH200"),
+    ("H100", "GH100"),
+    # Ampere — data centre
+    ("A100", "GA100"),
+    ("A800", "GA100"),
+    ("RTX A6000", "GA102"),
+    ("RTX A5000", "GA102"),
+    ("RTX A4000", "GA104"),
+    ("RTX A3000", "GA104"),
+    ("RTX A2000", "GA106"),
+    # Ampere — RTX 30xx
+    ("RTX 3090 Ti", "GA102"),
+    ("RTX 3090", "GA102"),
+    ("RTX 3080 Ti", "GA102"),
+    ("RTX 3080", "GA102"),
+    ("RTX 3070 Ti", "GA104"),
+    ("RTX 3070", "GA104"),
+    ("RTX 3060 Ti", "GA104"),
+    ("RTX 3060", "GA106"),
+    ("RTX 3050", "GA107"),
+    # Volta
+    ("Titan V", "GV100"),
+    ("V100", "GV100"),
+    # Turing — RTX 20xx
+    ("RTX 2080 Ti", "TU102"),
+    ("RTX 2080 Super", "TU104"),
+    ("RTX 2080", "TU104"),
+    ("RTX 2070 Super", "TU104"),
+    ("RTX 2070", "TU106"),
+    ("RTX 2060 Super", "TU106"),
+    ("RTX 2060", "TU106"),
+    # Turing — GTX 16xx
+    ("GTX 1660 Ti", "TU116"),
+    ("GTX 1660 Super", "TU116"),
+    ("GTX 1660", "TU116"),
+    ("GTX 1650 Super", "TU116"),
+    ("GTX 1650", "TU117"),
+    # Pascal — GTX 10xx
+    ("Titan Xp", "GP102"),
+    ("GTX 1080 Ti", "GP102"),
+    ("GTX 1080", "GP104"),
+    ("GTX 1070 Ti", "GP104"),
+    ("GTX 1070", "GP104"),
+    ("GTX 1060", "GP106"),
+    ("GTX 1050 Ti", "GP107"),
+    ("GTX 1050", "GP107"),
+]
+
+_STRIP_PREFIX = re.compile(r"^(?:NVIDIA\s+)?(?:GeForce\s+|Tesla\s+|Quadro\s+)?", re.IGNORECASE)
+
+
+def normalize(raw: str) -> tuple[str, bool]:
+    """Return (display_name, is_known_chip).
+
+    Strips vendor/product-line prefixes and looks up the die codename.
+    """
+    display = _STRIP_PREFIX.sub("", raw).strip()
+    _, known = chip_from_name(display)
+    return display, known
+
+
+def chip_from_name(normalized: str) -> tuple[str, bool]:
+    """Return (chip_codename, is_known) for a normalized GPU name."""
+    for substring, chip in _GPU_CHIP_MAP:
+        if substring.lower() in normalized.lower():
+            return chip, True
+    return normalized, False

mate_runtime_cuda-0.1.0.dist-info/METADATA

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: mate-runtime-cuda
+Version: 0.1.0
+Summary: NVIDIA CUDA runtime plugin for mate-bench
+Project-URL: Homepage, https://github.com/T0nd3/mate-bench
+Project-URL: Repository, https://github.com/T0nd3/mate-bench
+Author-email: Benjamin Fäuster <benjamin.faeuster@web.de>
+License: MIT
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: System :: Benchmark
+Requires-Python: >=3.11
+Requires-Dist: mate-bench<0.2,>=0.1

mate_runtime_cuda-0.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+mate_runtime_cuda/__init__.py,sha256=69Z8_O0ibwwHXAVYfb-KJqvWWJbmikY6pqp0SvJRqMM,802
+mate_runtime_cuda/_detect.py,sha256=zCN6mJiDxyYqS8cgJRmss5Zhrr8KpOYFWck9Y4d8Fg0,2711
+mate_runtime_cuda/_names.py,sha256=i72SQVyr1GemtaErcZBPXbTsD2U_xO9bybP14jKIRFs,2746
+mate_runtime_cuda-0.1.0.dist-info/METADATA,sha256=R0UQ7GBEGQqK8XwvtlhgpORZEzcrSmeYXqqqPCW4iPs,606
+mate_runtime_cuda-0.1.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+mate_runtime_cuda-0.1.0.dist-info/entry_points.txt,sha256=dkrl3sdi-_eyTc4dgsMiCJOyFZ9bBxZCEicN9_72IdI,58
+mate_runtime_cuda-0.1.0.dist-info/RECORD,,

mate_runtime_cuda-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

mate_runtime_cuda-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[mate_bench.runtime]
+cuda = mate_runtime_cuda:CudaRuntime