gpu-stack-inspector 0.1.0__py3-none-any.whl

gpu_stack_inspector/probes/metadata_program.py

@@ -0,0 +1,232 @@
+"""The metadata probe: run inside a *target* interpreter as ``python -c <this file>``.
+
+Rules for this file, checked by tests:
+
+- Python 3.8 syntax and the standard library only; targets can be much older than the
+  tool itself.
+- Nothing from the target environment is imported. Package metadata is read from disk.
+- Exactly one JSON document is printed, between BEGIN and END, so whatever a target's
+  sitecustomize or .pth files print cannot be mistaken for the result.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import site
+import sys
+import sysconfig
+from typing import Callable
+
+PROTOCOL = 1
+BEGIN = "GSI-PROBE-BEGIN"
+END = "GSI-PROBE-END"
+
+# Distributions that decide whether and how a Python environment uses an NVIDIA GPU.
+# Only these are reported: a full package list is noise, and can reveal private projects.
+EXACT = frozenset(
+    [
+        "torch", "torchvision", "torchaudio", "triton", "pytorch-triton", "pytorch-triton-rocm",
+        "xformers", "flash-attn", "bitsandbytes", "deepspeed", "vllm", "accelerate",
+        "jax", "jaxlib", "tensorflow", "tensorflow-cpu", "tensorflow-gpu", "tf-nightly", "keras",
+        "cupy", "onnxruntime", "onnxruntime-gpu", "tensorrt", "numpy",
+    ]
+)  # fmt: skip
+PREFIXES = ("nvidia-", "cupy-cuda", "jax-cuda", "tensorrt-")
+MAX_METADATA_BYTES = 64 * 1024
+
+# Libraries in torch/lib that only GPU builds ship (Linux, Windows, macOS names).
+CUDA_LIBRARIES = ("libtorch_cuda.so", "torch_cuda.dll", "libtorch_cuda.dylib")
+HIP_LIBRARIES = ("libtorch_hip.so", "torch_hip.dll")
+
+# torch/version.py lines, old and new style: "cuda = '11.3'", "cuda: Optional[str] = '12.8'",
+# "hip: Optional[str] = None". The file is read as text, never imported.
+VERSION_PY_LINE = re.compile(
+    r"^(?P<key>cuda|hip)\s*(?::[^=\n]*)?=\s*(?:'(?P<single>[^'\n]*)'|\"(?P<double>[^\"\n]*)\"|None)\s*$",
+    re.MULTILINE,
+)
+
+
+def python_facts() -> dict[str, object]:
+    enable_user_site = site.ENABLE_USER_SITE
+    return {
+        "version": "%d.%d.%d" % tuple(sys.version_info[:3]),
+        "implementation": sys.implementation.name,
+        "executable": sys.executable,
+        "prefix": sys.prefix,
+        "base_prefix": getattr(sys, "base_prefix", sys.prefix),
+        "conda": os.path.isdir(os.path.join(sys.prefix, "conda-meta")),
+        "platform": sysconfig.get_platform(),
+        "user_site": None if enable_user_site is None else bool(enable_user_site),
+        "pythonpath_set": bool(os.environ.get("PYTHONPATH")),
+    }
+
+
+def normalize(name: str) -> str:
+    """PEP 503 normalisation: ``Nvidia_CuBLAS.cu12`` → ``nvidia-cublas-cu12``."""
+    return re.sub(r"[-_.]+", "-", name).lower()
+
+
+def wanted(name: str) -> bool:
+    normalized = normalize(name)
+    return normalized in EXACT or normalized.startswith(PREFIXES)
+
+
+def distributions() -> list[dict[str, str]]:
+    """Allowlisted distributions on sys.path, read from metadata without importing them."""
+    try:
+        from importlib import metadata
+    except ImportError:  # Python before 3.8
+        return scan_dist_info(sys.path)
+
+    found = []
+    for dist in metadata.distributions():
+        name = dist.metadata["Name"] if dist.metadata is not None else None
+        version = dist.version
+        if not name or not version or not wanted(name):
+            continue
+        found.append(
+            {"name": normalize(name), "version": version, "location": str(dist.locate_file(""))}
+        )
+    return sorted(found, key=lambda d: (d["name"], d["location"]))
+
+
+def torch_facts() -> dict[str, object] | None:
+    """Static signals about the torch that ``import torch`` would find, without importing it.
+
+    Returns None when torch is not installed.
+    """
+    try:
+        from importlib import metadata
+    except ImportError:  # Python before 3.8: the distributions section still lists torch
+        return None
+    try:
+        dist = metadata.distribution("torch")
+    except metadata.PackageNotFoundError:
+        return None
+    requirements = []
+    for requirement in dist.requires or []:
+        match = re.match(r"[A-Za-z0-9._-]+", requirement)
+        if match and normalize(match.group(0)).startswith("nvidia-"):
+            requirements.append(normalize(match.group(0)))
+    package = str(dist.locate_file("torch"))
+    return {
+        "version": dist.version,
+        "location": str(dist.locate_file("")),
+        "requires_nvidia": sorted(set(requirements)),
+        "libraries": torch_libraries(os.path.join(package, "lib")),
+        "version_py": torch_version_py(os.path.join(package, "version.py")),
+        "conda_records": conda_records(sys.prefix),
+    }
+
+
+def conda_records(prefix: str) -> list[str] | None:
+    """Conda packages that decide the torch build, by file name in conda-meta.
+
+    Names like ``pytorch-1.12.0-py3.8_cuda11.3_cudnn8.3.2_0`` carry the build string, so
+    the (multi-megabyte) records are never opened. None outside conda environments.
+    """
+    try:
+        names = os.listdir(os.path.join(prefix, "conda-meta"))
+    except OSError:
+        return None
+    wanted_prefixes = ("pytorch-", "cudatoolkit-", "cuda-version-", "cuda-runtime-")
+    return sorted(
+        name[: -len(".json")]
+        for name in names
+        if name.endswith(".json") and name.startswith(wanted_prefixes)
+    )
+
+
+def torch_version_py(path: str) -> dict[str, str | None] | None:
+    """The cuda and hip versions torch was built with, read from torch/version.py as text.
+
+    None when the file is missing or names neither (a damaged or unusual install).
+    """
+    try:
+        with open(path, "rb") as handle:
+            text = handle.read(MAX_METADATA_BYTES).decode("utf-8", "replace")
+    except OSError:
+        return None
+    found: dict[str, str | None] = {}
+    for match in VERSION_PY_LINE.finditer(text):
+        value = match.group("single")
+        if value is None:
+            value = match.group("double")
+        found.setdefault(match.group("key"), value)
+    if not found:
+        return None
+    return {"cuda": found.get("cuda"), "hip": found.get("hip")}
+
+
+def torch_libraries(lib_dir: str) -> dict[str, bool] | None:
+    """Whether torch/lib holds CUDA or ROCm libraries; None if the folder cannot be read."""
+    try:
+        names = set(os.listdir(lib_dir))
+    except OSError:
+        return None
+    return {
+        "cuda": any(name in names for name in CUDA_LIBRARIES),
+        "hip": any(name in names for name in HIP_LIBRARIES),
+    }
+
+
+def scan_dist_info(paths: list[str]) -> list[dict[str, str]]:
+    """Fallback without importlib.metadata: read the headers of ``*.dist-info/METADATA``."""
+    found = []
+    for entry in paths:
+        try:
+            folders = sorted(os.listdir(entry or "."))
+        except OSError:
+            continue
+        for folder in folders:
+            if not folder.endswith(".dist-info"):
+                continue
+            headers = read_headers(os.path.join(entry, folder, "METADATA"))
+            name, version = headers.get("name"), headers.get("version")
+            if name and version and wanted(name):
+                found.append({"name": normalize(name), "version": version, "location": entry})
+    return found
+
+
+def read_headers(path: str) -> dict[str, str]:
+    """``Name:``/``Version:`` style headers up to the first blank line."""
+    headers: dict[str, str] = {}
+    try:
+        with open(path, "rb") as handle:
+            text = handle.read(MAX_METADATA_BYTES).decode("utf-8", "replace")
+    except OSError:
+        return headers
+    for line in text.splitlines():
+        if not line.strip():
+            break
+        key, separator, value = line.partition(":")
+        if separator and key.strip().lower() not in headers:
+            headers[key.strip().lower()] = value.strip()
+    return headers
+
+
+def collect(sections: list[tuple[str, Callable[[], object]]]) -> dict[str, object]:
+    """Run each section; an exception is reported and the other sections still run."""
+    errors: list[str] = []
+    result: dict[str, object] = {"protocol": PROTOCOL, "errors": errors}
+    for key, function in sections:
+        try:
+            result[key] = function()
+        except Exception as exc:  # a partial answer is worth more than none
+            result[key] = None
+            errors.append("%s: %s: %s" % (key, type(exc).__name__, exc))
+    return result
+
+
+def main() -> None:
+    result = collect(
+        [("python", python_facts), ("distributions", distributions), ("torch", torch_facts)]
+    )
+    sys.stdout.write(BEGIN + json.dumps(result, ensure_ascii=True, sort_keys=True) + END + "\n")
+    sys.stdout.flush()
+
+
+if __name__ == "__main__":
+    main()

gpu_stack_inspector/probes/runtime_probe.py

@@ -0,0 +1,153 @@
+"""Running the runtime probe in a target interpreter (only with --runtime) and reading it."""
+
+from __future__ import annotations
+
+import functools
+import json
+import tempfile
+from importlib.resources import files
+from pathlib import Path
+from typing import Any
+
+from gpu_stack_inspector.models import (
+    FailureKind,
+    GpuDevice,
+    ProbeFailure,
+    RuntimeFacts,
+    RuntimeStep,
+)
+from gpu_stack_inspector.probes.metadata_probe import python_error_line
+from gpu_stack_inspector.runner import CommandResult, CommandRunner, Outcome, probe_failure
+
+PROBE = "runtime"
+TIMEOUT = 60.0
+"""Importing torch and creating a CUDA context takes a few seconds (2-4 s on the test
+laptop); a driver that hangs must not hang the tool."""
+
+BEGIN, END, STEP = "GSI-RUNTIME-BEGIN", "GSI-RUNTIME-END", "GSI-RUNTIME-STEP"
+_SIGNALS = {
+    -11: "a segmentation fault (SIGSEGV)",
+    139: "a segmentation fault (SIGSEGV)",
+    -6: "an abort (SIGABRT)",
+    134: "an abort (SIGABRT)",
+    -9: "being killed (SIGKILL)",
+    137: "being killed (SIGKILL)",
+    -7: "a bus error (SIGBUS)",
+    135: "a bus error (SIGBUS)",
+}
+
+
+@functools.cache
+def program_source() -> str:
+    resource = files("gpu_stack_inspector.probes").joinpath("runtime_program.py")
+    return resource.read_text(encoding="utf-8")
+
+
+def run_runtime_probe(
+    runner: CommandRunner, path: str, *, timeout: float = TIMEOUT
+) -> RuntimeFacts:
+    """Import torch in the interpreter at ``path`` and use CUDA once. Only with consent."""
+    with tempfile.TemporaryDirectory(prefix="gsi-runtime-") as empty:
+        result = runner.run([path, "-c", program_source()], timeout=timeout, cwd=Path(empty))
+    announced = [
+        line[len(STEP) + 1 :].strip()
+        for line in result.stdout.splitlines()
+        if line.startswith(STEP + " ")
+    ]
+    if result.outcome is Outcome.TIMEOUT:
+        where = f" in step '{announced[-1]}'" if announced else ""
+        detail = (
+            f"no answer within {result.duration:.0f} s{where}; "
+            "CUDA initialisation can hang on a broken driver"
+        )
+        return _died(announced, ProbeFailure(PROBE, FailureKind.TIMEOUT, detail))
+    if result.outcome is Outcome.NONZERO:
+        return _died(announced, _crash(result, announced))
+    failure = probe_failure(PROBE, result)
+    if failure is not None:
+        return RuntimeFacts(failure=failure)
+    return parse_runtime_output(result.stdout)
+
+
+def _crash(result: CommandResult, announced: list[str]) -> ProbeFailure:
+    signal = _SIGNALS.get(result.returncode or 0)
+    how = f"crashed with {signal}" if signal else f"exited with code {result.returncode}"
+    where = f" during '{announced[-1]}'" if announced else ""
+    reason = python_error_line(result.stderr)
+    detail = f"the interpreter {how}{where}" + (f": {reason}" if reason and not signal else "")
+    return ProbeFailure(PROBE, FailureKind.NONZERO_EXIT, detail)
+
+
+def _died(announced: list[str], failure: ProbeFailure) -> RuntimeFacts:
+    """Steps reconstructed from the markers: all finished but the last, which it died in."""
+    steps = tuple(
+        RuntimeStep(
+            name,
+            ok=index < len(announced) - 1,
+            error=None if index < len(announced) - 1 else failure.detail,
+        )
+        for index, name in enumerate(announced)
+    )
+    return RuntimeFacts(steps=steps, failure=failure)
+
+
+def parse_runtime_output(stdout: str) -> RuntimeFacts:
+    start = stdout.find(BEGIN)
+    end = stdout.find(END, start + len(BEGIN)) if start != -1 else -1
+    if start == -1 or end == -1:
+        return RuntimeFacts(failure=_malformed("the interpreter printed no runtime result"))
+    try:
+        document = json.loads(stdout[start + len(BEGIN) : end])
+    except ValueError:
+        return RuntimeFacts(failure=_malformed("the runtime result is not valid JSON"))
+    if not isinstance(document, dict) or document.get("protocol") != 1:
+        return RuntimeFacts(failure=_malformed("unexpected runtime probe protocol"))
+    steps: list[RuntimeStep] = []
+    values: dict[str, Any] = {}
+    for raw in document.get("steps") or []:
+        if not isinstance(raw, dict) or not isinstance(raw.get("step"), str):
+            continue
+        seconds = raw.get("seconds")
+        error = raw.get("error")
+        steps.append(
+            RuntimeStep(
+                raw["step"],
+                ok=raw.get("ok") is True,
+                seconds=float(seconds) if isinstance(seconds, int | float) else None,
+                error=error if isinstance(error, str) else None,
+            )
+        )
+        if raw.get("ok") is True:
+            values[raw["step"]] = raw.get("value")
+    imported = values.get("import torch")
+    imported = imported if isinstance(imported, dict) else {}
+    available = values.get("cuda available")
+    attempted = any(s.step == "tensor operation" for s in steps)
+    return RuntimeFacts(
+        torch_version=_string(imported.get("version")),
+        cuda_version=_string(imported.get("cuda")),
+        cuda_available=available if isinstance(available, bool) else None,
+        devices=_devices(values.get("devices")),
+        tensor_ok=("tensor operation" in values) if attempted else None,
+        steps=tuple(steps),
+    )
+
+
+def _devices(raw: object) -> tuple[GpuDevice, ...]:
+    if not isinstance(raw, list):
+        return ()
+    return tuple(
+        GpuDevice(item["index"], item["name"], None, _string(item.get("capability")))
+        for item in raw
+        if isinstance(item, dict)
+        and isinstance(item.get("index"), int)
+        and isinstance(item.get("name"), str)
+    )
+
+
+def _string(value: object) -> str | None:
+    return value if isinstance(value, str) else None
+
+
+def _malformed(detail: str) -> ProbeFailure:
+    return ProbeFailure(PROBE, FailureKind.MALFORMED_OUTPUT, detail)

gpu_stack_inspector/probes/runtime_program.py

@@ -0,0 +1,101 @@
+"""The runtime probe: run inside a *target* interpreter only with --runtime (explicit consent).
+
+Unlike the metadata probe, this imports torch, creates a CUDA context and runs a tiny
+operation, because only that proves CUDA works. Rules for this file, checked by tests:
+
+- Python 3.8 syntax and the standard library plus torch; nothing else is imported.
+- The workload is tiny (16 floats) and bounded; nothing is written to disk.
+- Each step reports its own failure, so the report can say exactly where CUDA broke.
+- Exactly one JSON document is printed between BEGIN and END.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+import time
+from typing import Any, Callable
+
+PROTOCOL = 1
+BEGIN = "GSI-RUNTIME-BEGIN"
+END = "GSI-RUNTIME-END"
+STEP = "GSI-RUNTIME-STEP"
+"""Printed (and flushed) before each step: if CUDA initialisation kills the interpreter,
+no JSON is printed, and the last marker says which step it died in."""
+ELEMENTS = 16
+
+
+def first_line(exc: BaseException) -> str:
+    text = str(exc).strip().splitlines()
+    return "%s: %s" % (type(exc).__name__, text[0] if text else "")
+
+
+def run_steps() -> dict[str, object]:
+    result: dict[str, object] = {"protocol": PROTOCOL, "steps": []}
+    steps: list[dict[str, object]] = []
+    result["steps"] = steps
+    state: dict[str, Any] = {}
+
+    def step(name: str, action: Callable[[], object]) -> bool:
+        sys.stdout.write("%s %s\n" % (STEP, name))
+        sys.stdout.flush()
+        started = time.monotonic()
+        try:
+            value = action()
+        except BaseException as exc:  # report where it broke, then stop
+            seconds = round(time.monotonic() - started, 3)
+            steps.append({"step": name, "ok": False, "error": first_line(exc), "seconds": seconds})
+            return False
+        seconds = round(time.monotonic() - started, 3)
+        steps.append({"step": name, "ok": True, "value": value, "seconds": seconds})
+        return True
+
+    def import_torch() -> object:
+        import torch  # type: ignore[import-not-found]  # lives in the target, never in the tool
+
+        state["torch"] = torch
+        return {"version": torch.__version__, "cuda": torch.version.cuda}
+
+    def cuda_available() -> object:
+        torch = state["torch"]
+        return bool(torch.cuda.is_available())
+
+    def devices() -> object:
+        torch = state["torch"]
+        count = int(torch.cuda.device_count())
+        found = []
+        for index in range(count):
+            major, minor = torch.cuda.get_device_capability(index)
+            name = str(torch.cuda.get_device_name(index))
+            found.append({"index": index, "name": name, "capability": "%d.%d" % (major, minor)})
+        return found
+
+    def tensor_operation() -> object:
+        torch = state["torch"]
+        x = torch.ones(ELEMENTS, device="cuda")
+        total = (x * 2).sum()
+        torch.cuda.synchronize()
+        value = float(total.cpu().item())
+        if value != 2.0 * ELEMENTS:
+            raise ArithmeticError("expected %s, got %s" % (2.0 * ELEMENTS, value))
+        return value
+
+    if not step("import torch", import_torch):
+        return result
+    if not step("cuda available", cuda_available):
+        return result
+    if steps[-1]["value"] is not True:
+        return result  # nothing more to try: CUDA is not usable from this torch
+    if step("devices", devices):
+        step("tensor operation", tensor_operation)
+    return result
+
+
+def main() -> None:
+    result = run_steps()
+    sys.stdout.write(BEGIN + json.dumps(result, ensure_ascii=True, sort_keys=True) + END + "\n")
+    sys.stdout.flush()
+
+
+if __name__ == "__main__":
+    main()

gpu_stack_inspector/reports/__init__.py

@@ -0,0 +1 @@
+"""Renderers: they turn an already redacted snapshot into output and never change it."""

gpu_stack_inspector/reports/json_report.py

@@ -0,0 +1,43 @@
+"""JSON output: the snapshot as plain data, for scripts, CI and bug reports."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import fields, is_dataclass
+from enum import Enum
+from typing import Any
+
+from gpu_stack_inspector.models import FieldComparison, InterpreterFacts, Snapshot
+from gpu_stack_inspector.rules.torch_build import classify_torch
+
+
+def snapshot_to_dict(snapshot: Snapshot) -> dict[str, Any]:
+    """Plain dicts, lists and scalars; enums become their documented string values."""
+    plain = _plain(snapshot)
+    if not isinstance(plain, dict):  # cannot happen: a dataclass becomes a dict
+        raise TypeError("snapshot did not serialise to an object")
+    return plain
+
+
+def render_json(snapshot: Snapshot) -> str:
+    """Indented JSON with a trailing newline. Pass a redacted snapshot."""
+    return json.dumps(snapshot_to_dict(snapshot), indent=2, ensure_ascii=False) + "\n"
+
+
+def _plain(value: Any) -> Any:
+    if isinstance(value, Enum):
+        return value.value
+    if is_dataclass(value) and not isinstance(value, type):
+        result = {f.name: _plain(getattr(value, f.name)) for f in fields(value)}
+        if isinstance(value, InterpreterFacts):
+            # Derived from the facts above, but what readers and scripts look for.
+            result["virtualenv"] = value.virtualenv
+            result["torch_build"] = _plain(classify_torch(value.torch))
+        if isinstance(value, FieldComparison):
+            result["differs"] = value.differs
+        return result
+    if isinstance(value, tuple | list):
+        return [_plain(item) for item in value]
+    if isinstance(value, dict):
+        return {str(key): _plain(item) for key, item in value.items()}
+    return value