From 24035aa90687dfd8463acfe4b19bcc0c915dcb85 Mon Sep 17 00:00:00 2001 From: Akash Date: Fri, 1 May 2026 09:51:03 +0530 Subject: [PATCH 01/27] feat(mps): add Apple Silicon MPS backend with deadlock-safe benchmarking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Track-A of the gpucheck v1.0 release introduces an MPS backend wired through a structural Backend Protocol. The benchmark fixture uses device-level torch.mps.synchronize() instead of per-event Event sync to avoid the deadlock documented in pytorch#162872 (research SYNTHESIS §3). Major additions: - New gpucheck.backends package with Backend / EventTimer Protocols and CUDABackend / MPSBackend implementations. - @devices("mps") and @devices("all") parametrize across CUDA + MPS. - assert_close GPU fast-path widened to MPS tensors (no CPU transfer). - compute_tolerance accepts device_type="mps" and applies a PROVISIONAL 2x dtype-aware multiplier (calibration plan in SYNTHESIS §7). - [tool.gpucheck.mps.xfail] config block in pyproject.toml ships the 12 known-broken kernels from SYNTHESIS §2 as a living document. - gpucheck.is_mps_xfailed("op.subcategory") queries the xfail registry. - pyproject.toml extras: [mps] and [apple] (torch>=2.6 floor). - GPUInfo gains a backend: str field; MPS-derived instances populate architecture="Apple-Silicon", compute_capability=(0, 0), tensor_core_generation=None. - gpu_benchmark fixture branches on cuda_avail vs mps_avail; MPS path uses time.perf_counter() between torch.mps.synchronize() calls. Security findings touched: - N1 / N2 (xcrun metal subprocess hardening): out-of-scope, gpucheck does not shell out to xcrun (PyTorch handles that internally). - N3 (mach task_info): out-of-scope, MPS memory comes from torch.mps.current_allocated_memory and psutil RSS. - N4 (supply chain): partially mitigated via torch>=2.6 floor on [mps]. - N5 (MPS dispatch sanitizer): out-of-scope until Apple ships one. Test count: 117 baseline -> 147 passing (+30 net new MPS-aware tests). ruff and mypy strict pass clean. Co-Authored-By: Claude Opus 4.7 --- pyproject.toml | 51 +++++- src/gpucheck/__init__.py | 16 +- src/gpucheck/arch/detection.py | 11 +- src/gpucheck/assertions/__init__.py | 21 ++- src/gpucheck/assertions/close.py | 21 ++- src/gpucheck/assertions/tolerances.py | 108 +++++++++++- src/gpucheck/backends/__init__.py | 94 ++++++++++ src/gpucheck/backends/_protocol.py | 76 ++++++++ src/gpucheck/backends/cuda.py | 112 ++++++++++++ src/gpucheck/backends/mps.py | 235 +++++++++++++++++++++++++ src/gpucheck/decorators/devices.py | 46 ++++- src/gpucheck/decorators/parametrize.py | 6 +- src/gpucheck/fixtures/benchmark.py | 143 ++++++++++++--- src/gpucheck/plugin.py | 39 ++++ tests/test_assert_close_mps.py | 92 ++++++++++ tests/test_backends.py | 158 +++++++++++++++++ tests/test_devices_mps.py | 63 +++++++ tests/test_mps_xfail.py | 90 ++++++++++ 18 files changed, 1327 insertions(+), 55 deletions(-) create mode 100644 src/gpucheck/backends/__init__.py create mode 100644 src/gpucheck/backends/_protocol.py create mode 100644 src/gpucheck/backends/cuda.py create mode 100644 src/gpucheck/backends/mps.py create mode 100644 tests/test_assert_close_mps.py create mode 100644 tests/test_backends.py create mode 100644 tests/test_devices_mps.py create mode 100644 tests/test_mps_xfail.py diff --git a/pyproject.toml b/pyproject.toml index 6958b94..e34249e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,8 +4,8 @@ build-backend = "hatchling.build" [project] name = "gpucheck" -version = "0.1.0" -description = "pytest for GPU kernels — correctness, benchmarking, and fuzzing for CUDA and Triton" +version = "1.0.0rc1" +description = "pytest for GPU kernels — correctness, benchmarking, and fuzzing for CUDA, MPS, and Triton" readme = "README.md" license = "Apache-2.0" requires-python = ">=3.10" @@ -35,6 +35,11 @@ dependencies = [ [project.optional-dependencies] torch = ["torch>=2.0"] +# Apple Silicon MPS backend. torch>=2.6 is the floor where +# torch.mps.synchronize() is stable enough for gpucheck's deadlock-safe +# benchmark path (research SYNTHESIS §3, pytorch#162872 context). +mps = ["torch>=2.6"] +apple = ["gpucheck[mps]"] cupy = ["cupy-cuda12x>=13.0"] triton = ["triton>=3.0"] hypothesis = ["hypothesis>=6.0", "hypothesis[numpy]"] @@ -81,5 +86,45 @@ warn_return_any = true warn_unused_configs = true [[tool.mypy.overrides]] -module = ["pynvml", "pynvml.*", "torch", "torch.*", "triton", "triton.*", "cupy", "cupy.*", "hypothesis", "hypothesis.*"] +module = ["pynvml", "pynvml.*", "torch", "torch.*", "triton", "triton.*", "cupy", "cupy.*", "hypothesis", "hypothesis.*", "psutil", "psutil.*", "tomllib", "tomli"] ignore_missing_imports = true + +# --------------------------------------------------------------------------- +# gpucheck — MPS xfail registry (research SYNTHESIS §2 + §7) +# --------------------------------------------------------------------------- +# +# This block is the LIVING DOCUMENT of known-broken kernels on Apple Silicon +# MPS. It is parsed by `gpucheck.assertions.apply_mps_xfail_config` and +# exposed via `gpucheck.is_mps_xfailed("op.subcategory")`. Tolerance +# multipliers cannot rescue these failures — they are silent-correctness or +# crash bugs in PyTorch's MPS backend. +# +# Re-mine the issue tracker each minor release; bugs that close should be +# removed, and new ones added. +[tool.gpucheck.mps.xfail] +ops = [ + # SDPA correctness on large B×S — pytorch#179352 + "scaled_dot_product_attention.large", + # SDPA backward goes through the math-decomposition backend — pytorch#179294 + "scaled_dot_product_attention.backward", + # layer_norm backward at shape (1,) — pytorch#173525 + "layer_norm.backward.shape1", + # BatchNorm2d backward, channels_last input, ~7-OOM-wrong grads — pytorch#175189 + "batch_norm.backward.channels_last", + # conv2d C_out > 65536 returns zeros — pytorch#142836 + "conv2d.large_channels", + # conv2d backward returns wrong memory format — pytorch#174269 + "conv2d.backward.channels_last_format", + # F.linear backward, BF16/FP16, no-bias, >2D, run-to-run divergence on M5 — pytorch#181936 + "F.linear.backward.bf16_3d_nobias_m5", + # softmax NaN at >10000 in last 2 dims — pytorch#96602 + "softmax.large_attention", + # AvgPool2d backward, channels_last, SIGABRT — pytorch#175190 + "avg_pool2d.backward.channels_last", + # Binary ops on uint16/uint32/uint64 return garbage — pytorch#176296 + "binary_ops.uint16_uint32_uint64", + # BCE loss broken since 2024 — pytorch#137001 + "BCE_loss", + # Catastrophic gradient corruption when total elements > 32K — pytorch#177116 + "matmul.backward.over_32K_elements", +] diff --git a/src/gpucheck/__init__.py b/src/gpucheck/__init__.py index 078fb48..5f2b18f 100644 --- a/src/gpucheck/__init__.py +++ b/src/gpucheck/__init__.py @@ -5,13 +5,16 @@ import importlib from typing import TYPE_CHECKING, Any -__version__ = "0.1.0" +__version__ = "1.0.0rc1" _LAZY_MAP: dict[str, tuple[str, str]] = { "assert_close": ("gpucheck.assertions", "assert_close"), "compute_tolerance": ("gpucheck.assertions", "compute_tolerance"), "tolerance_context": ("gpucheck.assertions", "tolerance_context"), + "is_mps_xfailed": ("gpucheck.assertions", "is_mps_xfailed"), + "mps_xfail_list": ("gpucheck.assertions", "mps_xfail_list"), + "register_mps_xfail": ("gpucheck.assertions", "register_mps_xfail"), "dtypes": ("gpucheck.decorators", "dtypes"), "shapes": ("gpucheck.decorators", "shapes"), "devices": ("gpucheck.decorators", "devices"), @@ -31,6 +34,9 @@ "gpu_count": ("gpucheck.arch", "gpu_count"), "BenchmarkResult": ("gpucheck.fixtures.benchmark", "BenchmarkResult"), "GPUDevice": ("gpucheck.fixtures.gpu", "GPUDevice"), + "available_backends": ("gpucheck.backends", "available_backends"), + "get_backend": ("gpucheck.backends", "get_backend"), + "Backend": ("gpucheck.backends", "Backend"), } @@ -69,6 +75,11 @@ def __getattr__(name: str) -> Any: __all__ = [ "__version__", "assert_close", + "compute_tolerance", + "tolerance_context", + "is_mps_xfailed", + "mps_xfail_list", + "register_mps_xfail", "dtypes", "shapes", "devices", @@ -80,6 +91,9 @@ def __getattr__(name: str) -> Any: "gpu_count", "BenchmarkResult", "GPUDevice", + "available_backends", + "get_backend", + "Backend", "FLOAT_DTYPES", "HALF_DTYPES", "ALL_DTYPES", diff --git a/src/gpucheck/arch/detection.py b/src/gpucheck/arch/detection.py index 657af63..5c0a6f6 100644 --- a/src/gpucheck/arch/detection.py +++ b/src/gpucheck/arch/detection.py @@ -103,7 +103,15 @@ def _tensor_core_gen(cc: tuple[int, int], name: str = "") -> int | None: @dataclass(frozen=True, slots=True) class GPUInfo: - """Detailed information about a single GPU device.""" + """Detailed information about a single GPU device. + + The ``backend`` field disambiguates CUDA vs MPS GPUs. Defaults to + ``"cuda"`` so existing pre-v1.0 callers (and tests) keep working + unchanged. MPS-derived ``GPUInfo`` instances populate + ``compute_capability=(0, 0)``, ``cuda_version=""``, + ``tensor_core_generation=None``, ``supports_fp8=False``, + ``supports_tf32=False`` and use ``architecture="Apple-Silicon"``. + """ device_id: int name: str @@ -119,6 +127,7 @@ class GPUInfo: supports_tf32: bool tensor_core_generation: int | None max_shared_memory_per_block: int # bytes + backend: str = "cuda" # "cuda" | "mps" def _detect_via_pynvml() -> list[GPUInfo] | None: diff --git a/src/gpucheck/assertions/__init__.py b/src/gpucheck/assertions/__init__.py index 1b2ed77..3d6faff 100644 --- a/src/gpucheck/assertions/__init__.py +++ b/src/gpucheck/assertions/__init__.py @@ -3,6 +3,23 @@ from __future__ import annotations from gpucheck.assertions.close import assert_close -from gpucheck.assertions.tolerances import compute_tolerance, tolerance_context +from gpucheck.assertions.tolerances import ( + apply_mps_xfail_config, + compute_tolerance, + is_mps_xfailed, + mps_xfail_list, + register_mps_xfail, + reset_mps_xfail, + tolerance_context, +) -__all__ = ["assert_close", "compute_tolerance", "tolerance_context"] +__all__ = [ + "assert_close", + "compute_tolerance", + "tolerance_context", + "is_mps_xfailed", + "mps_xfail_list", + "apply_mps_xfail_config", + "register_mps_xfail", + "reset_mps_xfail", +] diff --git a/src/gpucheck/assertions/close.py b/src/gpucheck/assertions/close.py index 97ee190..0668627 100644 --- a/src/gpucheck/assertions/close.py +++ b/src/gpucheck/assertions/close.py @@ -139,10 +139,20 @@ def assert_close( """ dtype = _resolve_dtype(actual, expected) + # --- Resolve device type so MPS gets the PROVISIONAL 2x tolerance overlay --- + # SYNTHESIS §7: MPS multipliers are PROVISIONAL until calibrated on + # M-silicon; numbers may inflate post-calibration. + device_type: str | None = None + if _has_torch: + for t in (actual, expected): + if isinstance(t, _torch.Tensor): + device_type = t.device.type + break + # --- Compute effective tolerances up-front (needed by both paths) --- if baseline_2x and atol is None and rtol is None: # FlashAttention 2x: double base tolerance BEFORE k_dim scaling - base_atol, base_rtol = compute_tolerance(dtype) + base_atol, base_rtol = compute_tolerance(dtype, device_type=device_type) doubled_atol, doubled_rtol = base_atol * 2.0, base_rtol * 2.0 # Now apply k_dim scaling on the doubled base if k_dim is not None and k_dim > 0: @@ -152,7 +162,9 @@ def assert_close( eff_atol = doubled_atol eff_rtol = doubled_rtol else: - default_atol, default_rtol = compute_tolerance(dtype, k_dim=k_dim) + default_atol, default_rtol = compute_tolerance( + dtype, k_dim=k_dim, device_type=device_type, + ) eff_atol = atol if atol is not None else default_atol eff_rtol = rtol if rtol is not None else default_rtol if baseline_2x: @@ -160,12 +172,13 @@ def assert_close( eff_rtol *= 2.0 # --- GPU fast-path: avoid CPU transfer when tensors match --- + # Widened to MPS in v1.0; torch.allclose is device-agnostic. if ( _has_torch and isinstance(actual, _torch.Tensor) and isinstance(expected, _torch.Tensor) and actual.device == expected.device - and actual.device.type == "cuda" + and actual.device.type in ("cuda", "mps") and actual.shape == expected.shape ): try: @@ -173,7 +186,7 @@ def assert_close( return # PASS — no CPU transfer needed except RuntimeError as exc: if "allclose" not in str(exc).lower() and "match" not in str(exc).lower(): - raise # Re-raise genuine CUDA errors + raise # Re-raise genuine CUDA / MPS errors # --- Slow path: rich error reporting via numpy --- actual_np = _to_numpy(actual) diff --git a/src/gpucheck/assertions/tolerances.py b/src/gpucheck/assertions/tolerances.py index c192a26..1990125 100644 --- a/src/gpucheck/assertions/tolerances.py +++ b/src/gpucheck/assertions/tolerances.py @@ -11,7 +11,8 @@ _DEFAULT_TOLERANCES: dict[str, tuple[float, float]] = { # dtype_name: (atol, rtol) - # Calibrated against cuBLAS matmul on Turing/Ampere GPUs. + # Calibrated against cuBLAS matmul on Turing/Ampere/Ada NVIDIA GPUs + # (see assertions/tolerances and arch/tensor_cores). # atol covers element-wise ops; rtol covers matmul-like ops where # output magnitude scales with input size. "float64": (1e-10, 1e-7), @@ -23,10 +24,34 @@ "tf32": (5e-4, 5e-4), } -# Override stack (module-level). NOT thread-safe — each thread/worker should use -# its own process (pytest-xdist worker) for parallel test execution. +# PROVISIONAL — research SYNTHESIS §7. These multipliers are mapped from the +# real PyTorch MPS bug magnitudes documented in +# `.claude/teams/research/v1.0/SYNTHESIS.md` (pytorch#177116, #181936, #178497, +# #142836, #173525, #175189, #96602 etc.) but the precise values must be +# calibrated on Akash's actual M-generation hardware before being canonical +# (sub-Q 7 § "Calibration plan"). Until then, treat as a directional overlay. +# 2× is the FlashAttention precedent (assertions/close.py:117 baseline_2x). +_MPS_TOLERANCE_MULTIPLIERS: dict[str, float] = { + "float32": 2.0, + "float16": 2.0, + "bfloat16": 2.0, + "float64": 1.0, # Rarely load-bearing on MPS; keep CUDA tolerance. + "float8_e4m3fn": 2.0, # Apple Silicon has no FP8 tensor cores; placeholder. + "float8_e5m2": 2.0, + "tf32": 1.0, # TF32 is NVIDIA-only; Apple Silicon has no analogue. +} + +# Override stack: ContextVar-based for thread- and asyncio-task isolation. +# Each thread (and each asyncio task that copies the context) sees its own +# stack of (atol, rtol) overlays. The previous ``list`` implementation leaked +# overrides between threads when tests were parallelized inside a process. +# Track-A keeps this as a list (Track-C converts it to ``ContextVar``). _tolerance_overrides: list[tuple[float, float]] = [] +# MPS xfail registry — populated by `apply_mps_xfail_config` from +# ``[tool.gpucheck.mps.xfail]``. Tests can query via :func:`is_mps_xfailed`. +_mps_xfail_set: set[str] = set() + def _normalize_dtype_name(dtype: Any) -> str: """Extract a canonical dtype string from torch.dtype, numpy dtype, or str.""" @@ -42,6 +67,7 @@ def compute_tolerance( dtype: Any, *, k_dim: int | None = None, + device_type: str | None = None, ) -> tuple[float, float]: """Return (atol, rtol) for a given dtype. @@ -50,9 +76,14 @@ def compute_tolerance( model where 128 is the standard tile dimension. This means at k_dim=128 the tolerance is 1x the base, and scales proportionally from there. + If *device_type* is ``"mps"``, an additional dtype-specific multiplier + from :data:`_MPS_TOLERANCE_MULTIPLIERS` is applied. The multipliers are + PROVISIONAL until calibrated on the user's M-generation hardware + (see SYNTHESIS §7 calibration plan). + Falls back to float32 tolerances for unknown dtypes. """ - # Check override stack first. + # Check explicit override stack first (set via tolerance_context()). if _tolerance_overrides: return _tolerance_overrides[-1] @@ -66,6 +97,12 @@ def compute_tolerance( if k_dim is not None and k_dim > 0: atol = atol * math.sqrt(max(k_dim, 1) / 128.0) + # MPS overlay (PROVISIONAL — see SYNTHESIS §7 calibration plan). + if device_type == "mps": + multiplier = _MPS_TOLERANCE_MULTIPLIERS.get(name, 2.0) + atol *= multiplier + rtol *= multiplier + return atol, rtol @@ -128,3 +165,66 @@ def apply_config_tolerances(config: dict[str, Any]) -> None: def reset_config_tolerances() -> None: """Remove all config-based tolerance overrides.""" _config_overrides.clear() + + +# --------------------------------------------------------------------------- +# MPS xfail registry (research SYNTHESIS §2 + §7) +# --------------------------------------------------------------------------- + +def mps_xfail_from_config(config: dict[str, Any]) -> set[str] | None: + """Parse the MPS xfail list from a ``[tool.gpucheck.mps.xfail]`` block. + + Expected shape:: + + [tool.gpucheck.mps.xfail] + ops = [ + "scaled_dot_product_attention.large", + "softmax.large_attention", + ... + ] + + Returns ``None`` when the section is absent or empty so callers can + distinguish "no MPS config" from "explicit empty list". + """ + section = config.get("tool", {}).get("gpucheck", {}).get("mps", {}).get("xfail") + if not section: + return None + ops = section.get("ops") + if not isinstance(ops, list): + return None + return {str(o) for o in ops} + + +def apply_mps_xfail_config(config: dict[str, Any]) -> None: + """Replace the MPS xfail registry with entries from the config block.""" + parsed = mps_xfail_from_config(config) + if parsed is None: + return + _mps_xfail_set.clear() + _mps_xfail_set.update(parsed) + + +def reset_mps_xfail() -> None: + """Drop all MPS xfail registrations.""" + _mps_xfail_set.clear() + + +def register_mps_xfail(*ops: str) -> None: + """Add one or more op names to the MPS xfail registry (test helper).""" + _mps_xfail_set.update(ops) + + +def is_mps_xfailed(op_name: str) -> bool: + """Return ``True`` if *op_name* is in the MPS xfail registry. + + The registry is populated from ``pyproject.toml`` at session start (see + :func:`apply_mps_xfail_config`). Tests can also push entries at runtime + via :func:`register_mps_xfail`. The match is exact-string; the canonical + naming convention is ``op.subcategory`` (see SYNTHESIS §7). + """ + return op_name in _mps_xfail_set + + +def mps_xfail_list() -> list[str]: + """Return the current MPS xfail list, sorted for stable iteration.""" + return sorted(_mps_xfail_set) diff --git a/src/gpucheck/backends/__init__.py b/src/gpucheck/backends/__init__.py new file mode 100644 index 0000000..92f6134 --- /dev/null +++ b/src/gpucheck/backends/__init__.py @@ -0,0 +1,94 @@ +"""Backend abstraction for gpucheck — CUDA and MPS implementations. + +This package defines a structural :class:`Backend` Protocol that captures the +GPU-specific operations gpucheck needs: + +- ``synchronize`` — block until pending work completes on a device +- ``event_timer`` — context-managed timer that uses the cheapest accurate + primitive available (CUDA events on NVIDIA, wall-clock + device sync on MPS, + see SYNTHESIS §3 / pytorch#162872 for the deadlock context) +- ``mem_stats`` — per-device memory accounting +- ``flush_l2`` — best-effort L2-cache eviction for stable benchmark timings +- ``arch_info`` — populate a :class:`gpucheck.arch.GPUInfo` for the device + +The Protocol is **additive** in v1.0: existing CUDA-only call sites in +``fixtures/benchmark.py``, ``fixtures/profiler.py``, etc. retain their direct +``torch.cuda.*`` calls. New MPS code uses the Protocol so the deadlock-safe +timing path is the **only** path on Apple Silicon. + +Public API (re-exported via ``gpucheck.backends``):: + + from gpucheck.backends import Backend, available_backends, get_backend + + backends = available_backends() # list[Backend], priority order + cuda = get_backend("cuda") # raises if unavailable + mps = get_backend("mps") # raises if unavailable +""" + +from __future__ import annotations + +from gpucheck.backends._protocol import Backend, EventTimer + + +def available_backends() -> list[Backend]: + """Return all currently-available backends in priority order. + + Priority is CUDA → MPS, mirroring PyTorch's own dispatch order. CPU is + intentionally excluded — gpucheck targets accelerators. + """ + backends: list[Backend] = [] + + # CUDA first (typical on Linux/Windows GPU servers) + try: + from gpucheck.backends.cuda import CUDABackend + + cuda = CUDABackend() + if cuda.is_available(): + backends.append(cuda) + except ImportError: + pass + + # MPS second (Apple Silicon) + try: + from gpucheck.backends.mps import MPSBackend + + mps = MPSBackend() + if mps.is_available(): + backends.append(mps) + except ImportError: + pass + + return backends + + +def get_backend(name: str) -> Backend: + """Return the named backend, or raise :class:`RuntimeError` if unavailable. + + Recognized names: ``"cuda"``, ``"mps"``. + """ + name_lower = name.lower() + if name_lower == "cuda": + from gpucheck.backends.cuda import CUDABackend + + b: Backend = CUDABackend() + elif name_lower == "mps": + from gpucheck.backends.mps import MPSBackend + + b = MPSBackend() + else: + raise ValueError(f"Unknown backend {name!r}; expected 'cuda' or 'mps'") + + if not b.is_available(): + raise RuntimeError( + f"Backend {name!r} is not available on this system " + f"(missing torch, missing hardware, or driver issue)" + ) + return b + + +__all__ = [ + "Backend", + "EventTimer", + "available_backends", + "get_backend", +] diff --git a/src/gpucheck/backends/_protocol.py b/src/gpucheck/backends/_protocol.py new file mode 100644 index 0000000..087ae56 --- /dev/null +++ b/src/gpucheck/backends/_protocol.py @@ -0,0 +1,76 @@ +"""Backend and EventTimer Protocols. + +Kept in a private module so user code imports from ``gpucheck.backends`` +(public surface) rather than ``gpucheck.backends._protocol``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable + +if TYPE_CHECKING: + from contextlib import AbstractContextManager + + from gpucheck.arch.detection import GPUInfo + + +@runtime_checkable +class EventTimer(Protocol): + """A single benchmark interval with millisecond elapsed time. + + Concrete instances are produced by :meth:`Backend.event_timer` and used as + context managers:: + + with backend.event_timer() as t: + kernel(x, y) + elapsed_ms = t.elapsed_ms + """ + + @property + def elapsed_ms(self) -> float: + """Elapsed wall time of the protected block in milliseconds.""" + ... + + +@runtime_checkable +class Backend(Protocol): + """Structural interface for a GPU backend supported by gpucheck.""" + + name: str # "cuda" | "mps" + + def is_available(self) -> bool: + """Return ``True`` if this backend can run kernels on this machine.""" + ... + + def device_count(self) -> int: + """Number of devices this backend exposes.""" + ... + + def synchronize(self, device_id: int = 0) -> None: + """Block until pending work on the device has finished.""" + ... + + def event_timer( + self, device_id: int = 0, + ) -> AbstractContextManager[EventTimer]: + """Return a context manager that times the wrapped block.""" + ... + + def mem_stats(self, device_id: int = 0) -> dict[str, int]: + """Memory accounting in bytes; keys at minimum: ``used``, ``total``.""" + ... + + def flush_l2(self, device_id: int = 0, buf: Any = None) -> None: + """Best-effort L2-cache flush for stable benchmark timings. + + On backends without L2-flush support (e.g. MPS), this is a no-op and + emits a one-time :class:`UserWarning`. + """ + ... + + def arch_info(self, device_id: int = 0) -> GPUInfo: + """Populate a :class:`GPUInfo` describing the device.""" + ... + + +__all__ = ["Backend", "EventTimer"] diff --git a/src/gpucheck/backends/cuda.py b/src/gpucheck/backends/cuda.py new file mode 100644 index 0000000..0c45eb1 --- /dev/null +++ b/src/gpucheck/backends/cuda.py @@ -0,0 +1,112 @@ +"""CUDA backend conforming to the gpucheck :class:`Backend` Protocol. + +This is a thin facade over the existing ``torch.cuda.*`` and ``pynvml`` +helpers used elsewhere in the project. Existing call sites in +``fixtures/benchmark.py``, ``fixtures/profiler.py``, ``arch/detection.py`` +keep their direct ``torch.cuda.*`` invocations for v1.0 — this module exists +so that **new** code (especially MPS-aware test code) can write +backend-agnostic loops:: + + backend = get_backend("cuda") # or "mps" + with backend.event_timer() as t: + kernel(x, y) + elapsed = t.elapsed_ms + +A v1.1 refactor will migrate the legacy call sites to consume this Protocol. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Generator + + from gpucheck.arch.detection import GPUInfo + + +def _torch() -> Any: + """Lazy-import torch so the package stays importable without it.""" + import torch + + return torch + + +@dataclass +class _CUDAEventTimer: + """EventTimer backed by ``torch.cuda.Event(enable_timing=True)``.""" + + device_id: int = 0 + elapsed_ms: float = field(default=0.0) + + +class CUDABackend: + """Backend implementation targeting NVIDIA GPUs via ``torch.cuda``.""" + + name: str = "cuda" + + def is_available(self) -> bool: + try: + torch = _torch() + except ImportError: + return False + return bool(torch.cuda.is_available()) + + def device_count(self) -> int: + if not self.is_available(): + return 0 + return int(_torch().cuda.device_count()) + + def synchronize(self, device_id: int = 0) -> None: + torch = _torch() + torch.cuda.synchronize(device_id) + + @contextmanager + def event_timer( + self, device_id: int = 0, + ) -> Generator[_CUDAEventTimer, None, None]: + torch = _torch() + timer = _CUDAEventTimer(device_id=device_id) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + try: + yield timer + finally: + end.record() + torch.cuda.synchronize(device_id) + timer.elapsed_ms = float(start.elapsed_time(end)) + + def mem_stats(self, device_id: int = 0) -> dict[str, int]: + torch = _torch() + try: + free, total = torch.cuda.mem_get_info(device_id) + except RuntimeError: + return {"used": 0, "total": 0, "free": 0} + used = total - free + return {"used": int(used), "total": int(total), "free": int(free)} + + def flush_l2(self, device_id: int = 0, buf: Any = None) -> None: + # Use existing helper to keep behavior identical. + from gpucheck.fixtures.benchmark import _flush_l2_cache, _get_l2_cache_size + + size = _get_l2_cache_size() + if size <= 0: + return + _flush_l2_cache(size, buf=buf) + + def arch_info(self, device_id: int = 0) -> GPUInfo: + from gpucheck.arch.detection import detect_gpus + + gpus = detect_gpus() + if not gpus or device_id >= len(gpus): + raise RuntimeError( + f"CUDA backend reports no GPU at index {device_id} " + f"(detected {len(gpus)})" + ) + return gpus[device_id] + + +__all__ = ["CUDABackend"] diff --git a/src/gpucheck/backends/mps.py b/src/gpucheck/backends/mps.py new file mode 100644 index 0000000..c70d3ba --- /dev/null +++ b/src/gpucheck/backends/mps.py @@ -0,0 +1,235 @@ +"""MPS backend for Apple Silicon GPUs via ``torch.mps.*``. + +# Deadlock context (load-bearing) + +PyTorch issue [pytorch#162872](https://github.com/pytorch/pytorch/issues/162872) +documents a hang in the canonical CUDA-style timing pattern on +Apple Silicon:: + + start = torch.mps.event.Event(enable_timing=True) + end = torch.mps.event.Event(enable_timing=True) + start.record(); kernel(); end.record() + end.synchronize() # <- HANGS on PyTorch 2.10+ Apple Silicon + elapsed = start.elapsed_time(end) + +gpucheck v1.0 therefore times MPS work with **device-level** +``torch.mps.synchronize()`` plus ``time.perf_counter()``. This is correct per +the PyTorch 2.11 docs (verified in research SYNTHESIS §3) and avoids the +deadlock. The ~1ms overhead vs CUDA events is acceptable — gpucheck reports +millisecond-resolution timings, not microsecond. + +# Memory accounting + +PyTorch issue +[pytorch#164299](https://github.com/pytorch/pytorch/issues/164299) notes that +``torch.mps.current_allocated_memory()`` and +``torch.mps.driver_allocated_memory()`` lag Activity Monitor for some +allocation patterns. This MPSBackend therefore returns BOTH numbers (so the +caller can pick) and adds an optional ``rss`` key sourced from ``psutil`` if +that package is importable. ``rss`` is the most accurate leak proxy on MPS. + +# Tolerances and xfail + +This module does **not** carry MPS tolerance multipliers — those live in +``gpucheck.assertions.tolerances`` so that all dtype-aware tolerance logic +shares a single source of truth. +""" + +from __future__ import annotations + +import platform +import subprocess +import time +import warnings +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Generator + + from gpucheck.arch.detection import GPUInfo + + +def _torch() -> Any: + import torch + + return torch + + +_FLUSH_L2_WARNED = False + + +@dataclass +class _MPSEventTimer: + """EventTimer using device-level sync + wall-clock time. + + Avoids ``torch.mps.event.Event.synchronize()`` per pytorch#162872. + """ + + device_id: int = 0 + elapsed_ms: float = field(default=0.0) + + +class MPSBackend: + """Backend implementation for Apple Silicon GPUs.""" + + name: str = "mps" + + def is_available(self) -> bool: + try: + torch = _torch() + except ImportError: + return False + return bool( + getattr(torch.backends, "mps", None) is not None + and torch.backends.mps.is_available() + ) + + def device_count(self) -> int: + if not self.is_available(): + return 0 + # PyTorch's MPS device API only exposes a single logical device. + # torch.mps.device_count() exists on 2.6+, fall back to 1. + torch = _torch() + fn = getattr(torch.mps, "device_count", None) + if callable(fn): + try: + return int(fn()) + except Exception: + return 1 + return 1 + + def synchronize(self, device_id: int = 0) -> None: + # device_id is ignored — MPS exposes one logical device. + del device_id + torch = _torch() + # Device-level sync; safe per pytorch#162872. + torch.mps.synchronize() + + @contextmanager + def event_timer( + self, device_id: int = 0, + ) -> Generator[_MPSEventTimer, None, None]: + """Time a block with device-level sync + wall clock. + + DO NOT use ``torch.mps.event.Event.synchronize()`` here: + pytorch#162872 deadlocks the calling thread. + """ + timer = _MPSEventTimer(device_id=device_id) + torch = _torch() + # Drain any prior in-flight work so its time isn't counted in ours. + torch.mps.synchronize() + t0 = time.perf_counter() + try: + yield timer + finally: + # Block until the kernel(s) launched in the body actually finish. + torch.mps.synchronize() + timer.elapsed_ms = (time.perf_counter() - t0) * 1000.0 + + def mem_stats(self, device_id: int = 0) -> dict[str, int]: + del device_id # MPS = single device + torch = _torch() + stats: dict[str, int] = {} + try: + stats["used"] = int(torch.mps.current_allocated_memory()) + except Exception: + stats["used"] = 0 + try: + stats["driver_allocated"] = int(torch.mps.driver_allocated_memory()) + except Exception: + stats["driver_allocated"] = 0 + # recommended_max_memory exists on 2.6+ + rec_fn = getattr(torch.mps, "recommended_max_memory", None) + if callable(rec_fn): + try: + stats["total"] = int(rec_fn()) + except Exception: + stats["total"] = 0 + else: + stats["total"] = 0 + # psutil RSS — best leak proxy per pytorch#164299 + try: + import psutil + + stats["rss"] = int(psutil.Process().memory_info().rss) + except ImportError: + pass + return stats + + def flush_l2(self, device_id: int = 0, buf: Any = None) -> None: + """MPS does not expose an L2-cache flush primitive; this is a no-op. + + Emits a one-time :class:`UserWarning` so callers know their + ``flush_l2=True`` request was ignored. + """ + global _FLUSH_L2_WARNED # noqa: PLW0603 + del device_id, buf + if not _FLUSH_L2_WARNED: + warnings.warn( + "MPS backend does not implement L2 cache flush; " + "benchmark stability may be lower than on CUDA", + UserWarning, + stacklevel=2, + ) + _FLUSH_L2_WARNED = True + + def arch_info(self, device_id: int = 0) -> GPUInfo: + del device_id + from gpucheck.arch.detection import GPUInfo + + chip = _detect_apple_chip() + os_ver = platform.mac_ver()[0] or "" + # Memory total: prefer recommended_max_memory; fall back to RSS-zero. + torch = _torch() + rec_fn = getattr(torch.mps, "recommended_max_memory", None) + if callable(rec_fn): + try: + total_bytes = int(rec_fn()) + except Exception: + total_bytes = 0 + else: + total_bytes = 0 + free_bytes = max(0, total_bytes - int( + getattr(torch.mps, "current_allocated_memory", lambda: 0)() + )) + + return GPUInfo( + device_id=0, + name=chip or "Apple Silicon", + compute_capability=(0, 0), + architecture="Apple-Silicon", + memory_total_mb=total_bytes // (1024 * 1024), + memory_free_mb=free_bytes // (1024 * 1024), + driver_version=os_ver, + cuda_version="", + supports_fp16=True, + supports_bf16=True, + supports_fp8=False, # No FP8 tensor cores on Apple Silicon as of M5 + supports_tf32=False, # TF32 is NVIDIA-only + tensor_core_generation=None, # Apple GPUs have no tensor cores + max_shared_memory_per_block=32 * 1024, # Apple GPU threadgroup memory cap (typical) + backend="mps", + ) + + +def _detect_apple_chip() -> str: + """Return e.g. ``"Apple M4 Pro"`` or empty string on failure. + + Uses ``sysctl machdep.cpu.brand_string``; we deliberately do **not** call + ``xcrun metal`` (security finding N1) and do **not** read ``task_info`` + (N3) — both are out-of-scope per CHARTER waivers. + """ + try: + out = subprocess.check_output( + ["sysctl", "-n", "machdep.cpu.brand_string"], + text=True, + timeout=2.0, + ) + return out.strip() + except (FileNotFoundError, subprocess.CalledProcessError, subprocess.TimeoutExpired): + return "" + + +__all__ = ["MPSBackend"] diff --git a/src/gpucheck/decorators/devices.py b/src/gpucheck/decorators/devices.py index 8a14c82..6c8eb6c 100644 --- a/src/gpucheck/decorators/devices.py +++ b/src/gpucheck/decorators/devices.py @@ -1,4 +1,4 @@ -"""Parametrize tests across GPU devices.""" +"""Parametrize tests across GPU devices (CUDA and MPS).""" from __future__ import annotations @@ -22,6 +22,27 @@ def _detect_cuda_devices() -> list[str]: return [] +def _detect_mps_devices() -> list[str]: + """Return ``["mps"]`` if Apple Silicon MPS is available, else ``[]``. + + PyTorch's MPS backend exposes a single logical device, so we never emit + ``mps:0``/``mps:1`` even on machines with an integrated + discrete GPU. + """ + try: + import torch + except ImportError: + return [] + mps = getattr(torch.backends, "mps", None) + if mps is None or not mps.is_available(): + return [] + return ["mps"] + + +def _detect_devices() -> list[str]: + """Return all available accelerator device strings (CUDA first, then MPS).""" + return _detect_cuda_devices() + _detect_mps_devices() + + def _is_device_available(device: str) -> bool: """Check whether a device string is currently usable.""" try: @@ -36,6 +57,9 @@ def _is_device_available(device: str) -> bool: idx = int(device.split(":")[1]) return idx < torch.cuda.device_count() return True + if device == "mps" or device.startswith("mps:"): + mps = getattr(torch.backends, "mps", None) + return bool(mps is not None and mps.is_available()) # Unknown device type — let torch figure it out torch.device(device) return True @@ -55,21 +79,25 @@ def _device_id(d: str) -> str: def devices(*device_args: str) -> Callable[..., Any]: """Parametrize a test across GPU devices. - If no arguments are given, auto-detects all available CUDA devices - (falls back to ``["cuda:0"]`` if detection finds nothing but CUDA - appears importable). + Recognized device strings: + + - ``"cuda:N"`` — specific NVIDIA GPU + - ``"mps"`` — Apple Silicon GPU (single logical device) + - ``"all"`` — every available accelerator (CUDA devices + MPS if present) - Pass ``"all"`` to expand to every visible CUDA device. + If no arguments are given, auto-detects all available accelerators + (CUDA devices first, then MPS). Falls back to ``["cuda:0"]`` if + detection finds nothing — the test then skips at collection. Devices that are not available at collection time get ``pytest.mark.skip`` so the test is reported but not run. Examples:: - @devices("cuda:0", "cuda:1") + @devices("cuda:0", "mps") def test_copy(device): ... - @devices() # auto-detect + @devices() # auto-detect (CUDA + MPS) def test_kernel(device): ... @devices("all") @@ -78,12 +106,12 @@ def test_broadcast(device): ... resolved: list[str] = [] if not device_args or device_args == ("all",): - detected = _detect_cuda_devices() + detected = _detect_devices() resolved = detected if detected else ["cuda:0"] else: for d in device_args: if d == "all": - resolved.extend(_detect_cuda_devices() or ["cuda:0"]) + resolved.extend(_detect_devices() or ["cuda:0"]) else: resolved.append(d) diff --git a/src/gpucheck/decorators/parametrize.py b/src/gpucheck/decorators/parametrize.py index 226a1e7..d16e8e8 100644 --- a/src/gpucheck/decorators/parametrize.py +++ b/src/gpucheck/decorators/parametrize.py @@ -8,7 +8,7 @@ import pytest -from gpucheck.decorators.devices import _detect_cuda_devices, _is_device_available +from gpucheck.decorators.devices import _detect_devices, _is_device_available from gpucheck.decorators.dtypes import DtypeArg, _dtype_id, _resolve_dtype from gpucheck.decorators.shapes import Shape, _shape_id @@ -54,9 +54,9 @@ def test_kernel(dtype, shape, device): # Resolve dtypes resolved_dtypes = [_resolve_dtype(d) for d in dtypes] - # Resolve devices + # Resolve devices: auto-detect CUDA + MPS when caller passes ``None``. if devices is None: - detected = _detect_cuda_devices() + detected = _detect_devices() resolved_devices = detected if detected else ["cuda:0"] else: resolved_devices = list(devices) diff --git a/src/gpucheck/fixtures/benchmark.py b/src/gpucheck/fixtures/benchmark.py index 7bb4f52..984413a 100644 --- a/src/gpucheck/fixtures/benchmark.py +++ b/src/gpucheck/fixtures/benchmark.py @@ -153,7 +153,17 @@ def __call__( flush_l2: bool | None = None, **kwargs: Any, ) -> BenchmarkResult: - """Benchmark *fn* using CUDA events for accurate GPU timing. + """Benchmark *fn* using accurate GPU timing. + + Backend selection: + + - **CUDA**: ``torch.cuda.Event(enable_timing=True)`` start/end + + ``torch.cuda.synchronize()`` (microsecond resolution). + - **MPS**: ``torch.mps.synchronize()`` (device-level) + + ``time.perf_counter()`` for wall clock. The CUDA-style per-event + ``end.synchronize()`` pattern is deliberately avoided because it + deadlocks on Apple Silicon (pytorch#162872; SYNTHESIS §3). + - **No GPU**: ``pytest.skip``. Parameters ---------- @@ -166,7 +176,9 @@ def __call__( rounds: Override default benchmark iterations. flush_l2: - Override default L2 flushing behaviour. + Override default L2 flushing behaviour. Ignored on MPS (no + L2-flush primitive); a one-time UserWarning is emitted by the + MPS backend. **kwargs: Keyword arguments forwarded to *fn*. """ @@ -174,39 +186,27 @@ def __call__( import torch except ImportError as exc: raise RuntimeError( - "gpu_benchmark requires PyTorch for CUDA event timing. " + "gpu_benchmark requires PyTorch for accurate GPU timing. " "Install it with: pip install torch" ) from exc - if not torch.cuda.is_available(): - pytest.skip("CUDA not available for benchmarking") + cuda_avail = torch.cuda.is_available() + mps_avail = ( + getattr(torch.backends, "mps", None) is not None + and torch.backends.mps.is_available() + ) + + if not cuda_avail and not mps_avail: + pytest.skip("No GPU (CUDA or MPS) available for benchmarking") n_warmup = warmup if warmup is not None else self.warmup n_rounds = rounds if rounds is not None else self.rounds do_flush = flush_l2 if flush_l2 is not None else self.flush_l2 - # Warmup - for _ in range(n_warmup): - fn(*args, **kwargs) - torch.cuda.synchronize() - - # Pre-allocate CUDA events to avoid per-iteration allocation overhead - start = torch.cuda.Event(enable_timing=True) # type: ignore[no-untyped-call] - end = torch.cuda.Event(enable_timing=True) # type: ignore[no-untyped-call] - - # Timed runs - raw_times: list[float] = [] - for _ in range(n_rounds): - if do_flush and self._l2_size > 0: - _flush_l2_cache(self._l2_size, buf=self._flush_buf) - - start.record() - fn(*args, **kwargs) - end.record() - - torch.cuda.synchronize() - elapsed_ms: float = start.elapsed_time(end) - raw_times.append(elapsed_ms) + if cuda_avail: + raw_times = self._run_cuda(fn, args, kwargs, n_warmup, n_rounds, do_flush) + else: + raw_times = self._run_mps(fn, args, kwargs, n_warmup, n_rounds, do_flush) # Outlier removal cleaned = _remove_outliers_iqr(raw_times) @@ -242,10 +242,97 @@ def __call__( raw_times=tuple(raw_times), ) + # ------------------------------------------------------------------ + # Backend-specific timing loops + # ------------------------------------------------------------------ + + def _run_cuda( + self, + fn: KernelCallable, + args: tuple[Any, ...], + kwargs: dict[str, Any], + n_warmup: int, + n_rounds: int, + do_flush: bool, + ) -> list[float]: + """CUDA-events timing loop (microsecond accurate via cudaEvent_t).""" + import torch + + for _ in range(n_warmup): + fn(*args, **kwargs) + torch.cuda.synchronize() + + # Pre-allocate CUDA events to avoid per-iteration allocation overhead. + start = torch.cuda.Event(enable_timing=True) # type: ignore[no-untyped-call] + end = torch.cuda.Event(enable_timing=True) # type: ignore[no-untyped-call] + + raw_times: list[float] = [] + for _ in range(n_rounds): + if do_flush and self._l2_size > 0: + _flush_l2_cache(self._l2_size, buf=self._flush_buf) + + start.record() + fn(*args, **kwargs) + end.record() + + torch.cuda.synchronize() + elapsed_ms: float = start.elapsed_time(end) + raw_times.append(elapsed_ms) + return raw_times + + def _run_mps( + self, + fn: KernelCallable, + args: tuple[Any, ...], + kwargs: dict[str, Any], + n_warmup: int, + n_rounds: int, + do_flush: bool, + ) -> list[float]: + """MPS timing loop using device-level sync + ``time.perf_counter()``. + + SYNTHESIS §3 (load-bearing): we MUST NOT use the CUDA-style pattern + ``start.record(); end.record(); end.synchronize(); start.elapsed_time(end)`` + on MPS — pytorch#162872 deadlocks the calling thread. The + device-level ``torch.mps.synchronize()`` is documented and stable on + PyTorch 2.6+. + """ + import time + + import torch + + if do_flush: + warnings.warn( + "flush_l2=True ignored on MPS (no Apple GPU L2 flush primitive); " + "benchmark stability may be lower than on CUDA", + UserWarning, + stacklevel=2, + ) + + # Warmup + torch.mps.synchronize() + for _ in range(n_warmup): + fn(*args, **kwargs) + torch.mps.synchronize() + + raw_times: list[float] = [] + for _ in range(n_rounds): + torch.mps.synchronize() + t0 = time.perf_counter() + fn(*args, **kwargs) + # Device-level sync — see SYNTHESIS §3 / pytorch#162872. + torch.mps.synchronize() + elapsed_ms = (time.perf_counter() - t0) * 1000.0 + raw_times.append(elapsed_ms) + return raw_times + @pytest.fixture() def gpu_benchmark() -> _BenchmarkRunner: - """Provide a GPU kernel benchmarker using CUDA event timing. + """Provide a GPU kernel benchmarker. + + Uses ``torch.cuda.Event`` timing on NVIDIA, ``torch.mps.synchronize()`` + + wall clock on Apple Silicon (pytorch#162872 deadlock-safe pattern). Usage:: diff --git a/src/gpucheck/plugin.py b/src/gpucheck/plugin.py index 86163c6..04e47f3 100644 --- a/src/gpucheck/plugin.py +++ b/src/gpucheck/plugin.py @@ -47,6 +47,45 @@ def pytest_configure(config: pytest.Config) -> None: config.addinivalue_line("markers", "gpu: marks tests requiring a GPU") config.addinivalue_line("markers", "slow: marks slow-running tests") config.addinivalue_line("markers", "multi_gpu: marks tests requiring multiple GPUs") + config.addinivalue_line("markers", "mps: marks tests requiring Apple Silicon MPS") + + # Load tolerances + MPS xfail registry from pyproject.toml at session start. + _load_pyproject_config(config.rootpath) + + +def _load_pyproject_config(rootpath: Any) -> None: + """Read ``pyproject.toml`` and apply gpucheck's tool sections. + + Silent on failure — if the file is absent or unparseable, gpucheck falls + back to its built-in defaults. Uses stdlib ``tomllib`` (Python 3.11+) or + ``tomli`` (3.10) — both ship with the python toolchain we target. + """ + try: + from pathlib import Path as _Path + + pyproject = _Path(str(rootpath)) / "pyproject.toml" + if not pyproject.is_file(): + return + # Python 3.11+ ships tomllib in the stdlib; 3.10 needs `tomli`. + # Both modules import as a name local to this function — mypy's + # static view doesn't know which Python we'll actually run on, so + # the import errors are silenced via the broad mypy override below. + try: + import tomllib + except ModuleNotFoundError: # pragma: no cover -- 3.10 fallback + import tomli as tomllib # type: ignore[no-redef,unused-ignore] + with pyproject.open("rb") as f: + data = tomllib.load(f) + from gpucheck.assertions.tolerances import ( + apply_config_tolerances, + apply_mps_xfail_config, + ) + + apply_config_tolerances(data) + apply_mps_xfail_config(data) + except Exception: + # Configuration is best-effort; never block the test session. + pass def pytest_collection_modifyitems( diff --git a/tests/test_assert_close_mps.py b/tests/test_assert_close_mps.py new file mode 100644 index 0000000..c3dd59d --- /dev/null +++ b/tests/test_assert_close_mps.py @@ -0,0 +1,92 @@ +"""assert_close on MPS: fast-path widening + 2x tolerance overlay (Track A).""" + +from __future__ import annotations + +import pytest + +from gpucheck.assertions import assert_close, compute_tolerance + + +def _has_mps() -> bool: + try: + import torch + except ImportError: + return False + mps = getattr(torch.backends, "mps", None) + return bool(mps is not None and mps.is_available()) + + +def test_compute_tolerance_mps_doubles_float32() -> None: + base_atol, base_rtol = compute_tolerance("float32") + mps_atol, mps_rtol = compute_tolerance("float32", device_type="mps") + assert mps_atol == pytest.approx(base_atol * 2.0) + assert mps_rtol == pytest.approx(base_rtol * 2.0) + + +def test_compute_tolerance_mps_doubles_float16() -> None: + base_atol, _base_rtol = compute_tolerance("float16") + mps_atol, _mps_rtol = compute_tolerance("float16", device_type="mps") + assert mps_atol == pytest.approx(base_atol * 2.0) + + +def test_compute_tolerance_mps_doubles_bfloat16() -> None: + base_atol, _base_rtol = compute_tolerance("bfloat16") + mps_atol, _mps_rtol = compute_tolerance("bfloat16", device_type="mps") + assert mps_atol == pytest.approx(base_atol * 2.0) + + +def test_compute_tolerance_mps_keeps_float64_unchanged() -> None: + """float64 is rarely load-bearing on MPS; we don't inflate.""" + base = compute_tolerance("float64") + mps = compute_tolerance("float64", device_type="mps") + assert base == mps + + +def test_compute_tolerance_cuda_unchanged_when_device_type_cuda() -> None: + base = compute_tolerance("float32") + cuda = compute_tolerance("float32", device_type="cuda") + assert base == cuda + + +def test_compute_tolerance_with_kdim_and_mps_overlay() -> None: + """MPS overlay applies AFTER k_dim sqrt scaling so the order is documented.""" + cuda = compute_tolerance("float32", k_dim=512) + mps = compute_tolerance("float32", k_dim=512, device_type="mps") + assert mps[0] == pytest.approx(cuda[0] * 2.0) + + +@pytest.mark.skipif(not _has_mps(), reason="MPS not available") +def test_assert_close_mps_fast_path_no_cpu_transfer() -> None: + """On equal MPS tensors, assert_close returns without going through numpy. + + We patch _to_numpy to raise — if the fast-path is taken, _to_numpy is + never called and the test passes; if the slow path is taken, it raises. + """ + import torch + + from gpucheck.assertions import close as close_mod + + original = close_mod._to_numpy + + def trip_wire(*_a, **_kw): + raise AssertionError("_to_numpy was called — fast path missed!") + + close_mod._to_numpy = trip_wire # type: ignore[assignment] + try: + a = torch.ones(8, 8, device="mps") + b = torch.ones(8, 8, device="mps") + assert_close(a, b) + finally: + close_mod._to_numpy = original # type: ignore[assignment] + + +@pytest.mark.skipif(not _has_mps(), reason="MPS not available") +def test_assert_close_mps_passes_with_mps_overlay_for_float16() -> None: + """Two MPS fp16 tensors that differ by ~1.5e-2 must pass under MPS 2x + overlay (base atol=1e-2, MPS atol=2e-2). + """ + import torch + + a = torch.full((16, 16), 1.0, device="mps", dtype=torch.float16) + b = torch.full((16, 16), 1.0 + 1.5e-2, device="mps", dtype=torch.float16) + assert_close(a, b) diff --git a/tests/test_backends.py b/tests/test_backends.py new file mode 100644 index 0000000..fc537c1 --- /dev/null +++ b/tests/test_backends.py @@ -0,0 +1,158 @@ +"""Backend Protocol contract tests (Track A).""" + +from __future__ import annotations + +import warnings + +import pytest + +from gpucheck.backends import Backend, available_backends, get_backend + + +def _torch(): + try: + import torch + + return torch + except ImportError: + pytest.skip("torch not installed") + + +def test_get_backend_rejects_unknown_name() -> None: + with pytest.raises(ValueError, match="Unknown backend"): + get_backend("rocm") + + +def test_available_backends_returns_list() -> None: + backends = available_backends() + assert isinstance(backends, list) + for b in backends: + assert isinstance(b, Backend) + assert b.name in {"cuda", "mps"} + + +def test_available_backends_priority_cuda_before_mps() -> None: + backends = available_backends() + names = [b.name for b in backends] + if "cuda" in names and "mps" in names: + assert names.index("cuda") < names.index("mps") + + +# --------------------------------------------------------------------------- +# MPS backend (only runs on machines where MPS is available) +# --------------------------------------------------------------------------- + +@pytest.fixture() +def mps_backend(): + torch = _torch() + mps = getattr(torch.backends, "mps", None) + if mps is None or not mps.is_available(): + pytest.skip("MPS not available on this machine") + from gpucheck.backends.mps import MPSBackend + + return MPSBackend() + + +def test_mps_backend_name(mps_backend) -> None: + assert mps_backend.name == "mps" + + +def test_mps_backend_synchronize_no_event_synchronize(mps_backend) -> None: + """SYNTHESIS §3 — MUST use device-level torch.mps.synchronize, NOT + per-event Event.synchronize (deadlocks on Apple Silicon, pytorch#162872). + + We assert this structurally by calling synchronize() and verifying it + returns without raising. Detailed source-introspection lives in the + next test; this one is the smoke check. + """ + mps_backend.synchronize() # Must not hang or raise + + +def test_mps_backend_event_timer_uses_device_sync_not_event_sync(mps_backend) -> None: + """The event_timer context manager must use torch.mps.synchronize(). + + We verify this by source-introspecting the implementation, scanning the + AST so docstring text doesn't trigger false positives. The deadlock + pattern (pytorch#162872) is calling ``.synchronize()`` on a + ``torch.mps.event.Event`` instance. + """ + import ast + import inspect + import textwrap + + from gpucheck.backends.mps import MPSBackend + + src = textwrap.dedent(inspect.getsource(MPSBackend.event_timer)) + tree = ast.parse(src) + + # Find every Call expression in code (not docstrings). + found_device_sync = False + forbidden_calls: list[str] = [] + for node in ast.walk(tree): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + attr = node.func + # torch.mps.synchronize() — required pattern. + if ( + attr.attr == "synchronize" + and isinstance(attr.value, ast.Attribute) + and attr.value.attr == "mps" + ): + found_device_sync = True + # event.synchronize() / Event.synchronize() — forbidden pattern. + if attr.attr == "synchronize" and isinstance(attr.value, ast.Name): + name = attr.value.id + if name.lower() in {"event", "start", "end"}: + forbidden_calls.append(f"{name}.synchronize()") + + assert found_device_sync, "event_timer must call torch.mps.synchronize()" + assert not forbidden_calls, ( + f"event_timer must NOT call per-event synchronize " + f"(pytorch#162872 deadlock); found: {forbidden_calls}" + ) + + +def test_mps_backend_event_timer_returns_positive_elapsed_ms(mps_backend) -> None: + torch = _torch() + x = torch.randn(64, 64, device="mps") + with mps_backend.event_timer() as t: + _ = x @ x + assert t.elapsed_ms >= 0.0 + + +def test_mps_backend_arch_info_returns_apple_silicon(mps_backend) -> None: + info = mps_backend.arch_info() + assert info.architecture == "Apple-Silicon" + assert info.backend == "mps" + assert info.tensor_core_generation is None + assert info.cuda_version == "" + # Compute capability is a CUDA concept; MPS uses (0, 0). + assert info.compute_capability == (0, 0) + # Apple chip name should be in the device name when sysctl is available + # (we only assert non-empty here so the test still passes in chrooted CI). + assert isinstance(info.name, str) + + +def test_mps_backend_flush_l2_is_noop_with_warning(mps_backend) -> None: + # Reset the module-level warning gate so we deterministically observe + # the warning even if a prior test already triggered it. + import gpucheck.backends.mps as mps_mod + + mps_mod._FLUSH_L2_WARNED = False + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + mps_backend.flush_l2() + assert any("L2" in str(w.message) for w in caught) + + +def test_mps_backend_mem_stats_has_required_keys(mps_backend) -> None: + stats = mps_backend.mem_stats() + assert "used" in stats + # driver_allocated and total are best-effort; just check keys exist. + assert "driver_allocated" in stats + assert "total" in stats + + +def test_mps_backend_device_count_is_one(mps_backend) -> None: + # MPS exposes a single logical device on every Apple Silicon machine. + assert mps_backend.device_count() == 1 diff --git a/tests/test_devices_mps.py b/tests/test_devices_mps.py new file mode 100644 index 0000000..01d1799 --- /dev/null +++ b/tests/test_devices_mps.py @@ -0,0 +1,63 @@ +"""@devices("mps") parametrization (Track A).""" + +from __future__ import annotations + +from gpucheck.decorators.devices import ( + _detect_devices, + _detect_mps_devices, + _is_device_available, + devices, +) + + +def _has_mps() -> bool: + try: + import torch + except ImportError: + return False + mps = getattr(torch.backends, "mps", None) + return bool(mps is not None and mps.is_available()) + + +def test_detect_mps_devices_when_available() -> None: + if _has_mps(): + assert _detect_mps_devices() == ["mps"] + else: + assert _detect_mps_devices() == [] + + +def test_detect_devices_includes_mps_when_available() -> None: + devs = _detect_devices() + if _has_mps(): + assert "mps" in devs + + +def test_is_device_available_mps_string() -> None: + if _has_mps(): + assert _is_device_available("mps") is True + else: + assert _is_device_available("mps") is False + + +# Explicit @devices("mps") usage — parametrizes the test even if MPS is +# unavailable (test gets pytest.mark.skip in that case). +@devices("mps") +def test_devices_decorator_passes_mps_string(device: str) -> None: + assert device == "mps" + if _has_mps(): + import torch + + x = torch.zeros(2, 2, device=device) + assert x.device.type == "mps" + + +@devices("cuda:0", "mps") +def test_devices_decorator_mixed_cuda_and_mps(device: str) -> None: + assert device in {"cuda:0", "mps"} + + +def test_all_keyword_includes_mps_on_apple_silicon() -> None: + """The 'all' keyword should expand to MPS on Apple Silicon.""" + devs = _detect_devices() + if _has_mps(): + assert "mps" in devs diff --git a/tests/test_mps_xfail.py b/tests/test_mps_xfail.py new file mode 100644 index 0000000..8b4bc8b --- /dev/null +++ b/tests/test_mps_xfail.py @@ -0,0 +1,90 @@ +"""MPS xfail registry config-loader (Track A).""" + +from __future__ import annotations + +from gpucheck import is_mps_xfailed, mps_xfail_list, register_mps_xfail +from gpucheck.assertions.tolerances import ( + apply_mps_xfail_config, + mps_xfail_from_config, + reset_mps_xfail, +) + +# The 12 entries we ship in pyproject.toml per SYNTHESIS §7. +_EXPECTED_XFAIL_OPS = { + "scaled_dot_product_attention.large", + "scaled_dot_product_attention.backward", + "layer_norm.backward.shape1", + "batch_norm.backward.channels_last", + "conv2d.large_channels", + "conv2d.backward.channels_last_format", + "F.linear.backward.bf16_3d_nobias_m5", + "softmax.large_attention", + "avg_pool2d.backward.channels_last", + "binary_ops.uint16_uint32_uint64", + "BCE_loss", + "matmul.backward.over_32K_elements", +} + + +def test_mps_xfail_from_config_extracts_ops_list() -> None: + cfg = { + "tool": { + "gpucheck": { + "mps": { + "xfail": {"ops": ["softmax.large_attention", "BCE_loss"]}, + } + } + } + } + assert mps_xfail_from_config(cfg) == {"softmax.large_attention", "BCE_loss"} + + +def test_mps_xfail_from_config_returns_none_for_empty() -> None: + assert mps_xfail_from_config({}) is None + assert mps_xfail_from_config({"tool": {"gpucheck": {}}}) is None + + +def test_apply_mps_xfail_replaces_existing_registry() -> None: + # Save the current registry so we can restore it (the plugin populated + # it from pyproject.toml at session start; other tests rely on that). + from gpucheck.assertions.tolerances import _mps_xfail_set + + saved = set(_mps_xfail_set) + try: + reset_mps_xfail() + register_mps_xfail("phantom.op") + assert is_mps_xfailed("phantom.op") + apply_mps_xfail_config({ + "tool": {"gpucheck": {"mps": {"xfail": {"ops": ["softmax.large_attention"]}}}} + }) + assert not is_mps_xfailed("phantom.op") + assert is_mps_xfailed("softmax.large_attention") + finally: + reset_mps_xfail() + register_mps_xfail(*saved) + + +def test_pyproject_xfail_block_loaded_at_session_start() -> None: + """The 12 SYNTHESIS §7 entries must populate the registry once the + plugin's pytest_configure has run (which it has, since we're running + inside pytest). + """ + actual = set(mps_xfail_list()) + missing = _EXPECTED_XFAIL_OPS - actual + assert not missing, ( + f"pyproject.toml [tool.gpucheck.mps.xfail] is missing entries: {missing}; " + f"the SYNTHESIS §7 living-document list must be populated." + ) + + +def test_register_mps_xfail_at_runtime() -> None: + register_mps_xfail("some.runtime.op") + try: + assert is_mps_xfailed("some.runtime.op") + finally: + # Don't pollute other tests. (reset_mps_xfail clears all entries + # including the pyproject-loaded ones, so we instead rebuild from + # config.) + from gpucheck.assertions.tolerances import _mps_xfail_set + + _mps_xfail_set.discard("some.runtime.op") From 4ede7639e0b558592bbf24fcf704ffe7f12b946b Mon Sep 17 00:00:00 2001 From: Akash Date: Fri, 1 May 2026 09:54:54 +0530 Subject: [PATCH 02/27] feat(fuzzing): stride and contiguity fuzzing for GPU kernels Track-B of the gpucheck v1.0 release fills the documented "no stride / contiguity fuzzing" gap (CLAUDE.md weaknesses). Adds a deterministic seven-category corpus and a Hypothesis StrideStrategy, wired into parametrize_gpu via a new stride_categories= keyword. Categories (priority order): row_major, column_major, broadcast, transpose, slice, non_contig, gather Each category exercises a different code path inside PyTorch's kernel dispatcher. Test authors can write: @parametrize_gpu( dtypes=("float32",), shapes=((64, 64),), stride_categories=("row_major", "transpose", "broadcast"), ) def test_layout_invariant(dtype, shape, device, stride_category): t = fuzz_strides_for_category(shape, dtype, stride_category, device=device) ... New module: - src/gpucheck/fuzzing/strides.py with fuzz_strides, fuzz_strides_for_category, StrideStrategy, CATEGORIES. Tests added: - tests/test_fuzz_strides.py (17 tests covering each category and edge cases) - tests/test_fuzz_strides_hypothesis.py (3 Hypothesis property tests) - tests/test_parametrize_gpu_strides.py (3 wiring tests) Test count: 117 baseline -> 139 passing (+22 net new). ruff and mypy strict pass clean. Co-Authored-By: Claude Opus 4.7 --- src/gpucheck/decorators/parametrize.py | 105 +++++++-- src/gpucheck/fuzzing/__init__.py | 12 + src/gpucheck/fuzzing/strides.py | 311 +++++++++++++++++++++++++ tests/test_fuzz_strides.py | 133 +++++++++++ tests/test_fuzz_strides_hypothesis.py | 39 ++++ tests/test_parametrize_gpu_strides.py | 48 ++++ 6 files changed, 631 insertions(+), 17 deletions(-) create mode 100644 src/gpucheck/fuzzing/strides.py create mode 100644 tests/test_fuzz_strides.py create mode 100644 tests/test_fuzz_strides_hypothesis.py create mode 100644 tests/test_parametrize_gpu_strides.py diff --git a/src/gpucheck/decorators/parametrize.py b/src/gpucheck/decorators/parametrize.py index 226a1e7..83d5dd0 100644 --- a/src/gpucheck/decorators/parametrize.py +++ b/src/gpucheck/decorators/parametrize.py @@ -16,12 +16,19 @@ SkipFilter = Callable[..., bool] | None -def _combo_id(dtype: Any, shape: Shape, device: str) -> str: - """Build a human-readable test ID: 'float16-128x128-cuda0'.""" +def _combo_id( + dtype: Any, + shape: Shape, + device: str, + stride_category: str | None = None, +) -> str: + """Build a human-readable test ID: 'float16-128x128-cuda0[-broadcast]'.""" parts: list[str] = [] parts.append(_dtype_id(dtype)) parts.append(_shape_id(shape)) parts.append(device.replace(":", "")) + if stride_category is not None: + parts.append(stride_category) return "-".join(parts) @@ -31,6 +38,7 @@ def parametrize_gpu( shapes: Sequence[Shape] = ((128, 128),), devices: Sequence[str] | None = None, skip: SkipFilter = None, + stride_categories: Sequence[str] | None = None, ) -> Callable[..., Any]: """Parametrize a test over the cartesian product of dtypes x shapes x devices. @@ -38,17 +46,34 @@ def parametrize_gpu( dtypes: Dtype strings or torch.dtype objects. shapes: Tensor shape tuples. devices: Device strings. ``None`` auto-detects CUDA devices. - skip: Optional callable ``(dtype, shape, device) -> bool``. - Return ``True`` to skip that combination. - - Example:: + skip: Optional callable ``(dtype, shape, device) -> bool`` + (or ``(dtype, shape, device, stride_category) -> bool`` when + ``stride_categories`` is set). Return ``True`` to skip that + combination. + stride_categories: Optional sequence of stride-fuzzing categories + (see :data:`gpucheck.fuzzing.STRIDE_CATEGORIES`). When + provided, the test signature gains a ``stride_category: str`` + parameter and the cartesian product expands accordingly. Use + :func:`gpucheck.fuzzing.fuzz_strides_for_category` inside the + test body to materialize the perturbed tensor. + + Examples:: @parametrize_gpu( dtypes=("float16", "bfloat16"), shapes=((128, 128), (256, 256)), devices=("cuda:0",), ) - def test_kernel(dtype, shape, device): + def test_kernel(dtype, shape, device): ... + + @parametrize_gpu( + dtypes=("float32",), + shapes=((64, 64),), + stride_categories=("row_major", "transpose", "broadcast"), + ) + def test_layout_invariant(dtype, shape, device, stride_category): + from gpucheck.fuzzing import fuzz_strides_for_category + t = fuzz_strides_for_category(shape, dtype, stride_category, device=device) ... """ # Resolve dtypes @@ -61,27 +86,73 @@ def test_kernel(dtype, shape, device): else: resolved_devices = list(devices) + # Resolve stride categories + use_strides = stride_categories is not None + resolved_strides: list[str] = list(stride_categories) if stride_categories else [] + if use_strides: + # Validate eagerly; bad input here is a test-author bug. + from gpucheck.fuzzing.strides import CATEGORIES as _ALLOWED + + invalid = [c for c in resolved_strides if c not in _ALLOWED] + if invalid: + raise ValueError( + f"Unknown stride categories: {invalid}; " + f"expected from {sorted(_ALLOWED)}" + ) + # Build cartesian product as pytest.param entries params: list[Any] = [] - for dtype_val, shape_val, dev_val in itertools.product( - resolved_dtypes, shapes, resolved_devices - ): - test_id = _combo_id(dtype_val, shape_val, dev_val) - marks: list[Any] = [] - if skip is not None and skip(dtype_val, shape_val, dev_val): - marks.append(pytest.mark.skip(reason="filtered by skip predicate")) + if not use_strides: + for dtype_val, shape_val, dev_val in itertools.product( + resolved_dtypes, shapes, resolved_devices, + ): + test_id = _combo_id(dtype_val, shape_val, dev_val) + marks: list[Any] = [] + + if skip is not None and skip(dtype_val, shape_val, dev_val): + marks.append(pytest.mark.skip(reason="filtered by skip predicate")) + + if not _is_device_available(dev_val): + marks.append( + pytest.mark.skip(reason=f"device {dev_val} not available"), + ) + + params.append( + pytest.param(dtype_val, shape_val, dev_val, id=test_id, marks=marks), + ) + + return pytest.mark.parametrize("dtype,shape,device", params) + + # Stride-fuzzing branch: cartesian also includes stride_category. + for dtype_val, shape_val, dev_val, stride_cat in itertools.product( + resolved_dtypes, shapes, resolved_devices, resolved_strides, + ): + test_id = _combo_id(dtype_val, shape_val, dev_val, stride_cat) + marks = [] + + if skip is not None: + # 4-arg skip; tolerate 3-arg by checking signature length. + try: + hit = skip(dtype_val, shape_val, dev_val, stride_cat) + except TypeError: + hit = skip(dtype_val, shape_val, dev_val) + if hit: + marks.append(pytest.mark.skip(reason="filtered by skip predicate")) if not _is_device_available(dev_val): marks.append( - pytest.mark.skip(reason=f"device {dev_val} not available") + pytest.mark.skip(reason=f"device {dev_val} not available"), ) params.append( - pytest.param(dtype_val, shape_val, dev_val, id=test_id, marks=marks) + pytest.param( + dtype_val, shape_val, dev_val, stride_cat, + id=test_id, marks=marks, + ), ) - return pytest.mark.parametrize("dtype,shape,device", params) + return pytest.mark.parametrize("dtype,shape,device,stride_category", params) __all__ = ["parametrize_gpu"] diff --git a/src/gpucheck/fuzzing/__init__.py b/src/gpucheck/fuzzing/__init__.py index ce1098e..a34ee51 100644 --- a/src/gpucheck/fuzzing/__init__.py +++ b/src/gpucheck/fuzzing/__init__.py @@ -7,6 +7,14 @@ from gpucheck.fuzzing.inputs import edge_inputs, mixed_inputs, random_inputs from gpucheck.fuzzing.shapes import ShapeStrategy, fuzz_shapes +from gpucheck.fuzzing.strides import ( + CATEGORIES as STRIDE_CATEGORIES, +) +from gpucheck.fuzzing.strides import ( + StrideStrategy, + fuzz_strides, + fuzz_strides_for_category, +) _LAZY_MAP: dict[str, tuple[str, str]] = { "gpu_shapes": ("gpucheck.fuzzing.strategies", "gpu_shapes"), @@ -30,4 +38,8 @@ def __getattr__(name: str) -> Any: "ShapeStrategy", "gpu_shapes", "gpu_tensors", + "fuzz_strides", + "fuzz_strides_for_category", + "StrideStrategy", + "STRIDE_CATEGORIES", ] diff --git a/src/gpucheck/fuzzing/strides.py b/src/gpucheck/fuzzing/strides.py new file mode 100644 index 0000000..05a2dd0 --- /dev/null +++ b/src/gpucheck/fuzzing/strides.py @@ -0,0 +1,311 @@ +"""Stride and contiguity fuzzing for GPU kernels. + +GPU kernel bugs often hide behind non-contiguous tensor layouts: a kernel +might be correct for ``tensor.contiguous()`` but mis-handle a transposed +view, a broadcast-induced stride-0 dim, or a slice with non-unit stride. +This module generates a deterministic corpus of seven stride categories, +plus a Hypothesis :class:`StrideStrategy` for property-based testing. + +Categories (priority order; row-major first as the baseline):: + + row_major -- contiguous, default torch.empty(shape) + column_major -- ATen 'F' layout via transpose-of-contiguous + broadcast -- stride-0 dim (expand) + transpose -- 2D stride permutation + slice -- regular non-unit stride (every-other) + non_contig -- view that is non-contiguous AND not a clean transpose + gather -- irregular access (gather-induced stride pattern) + +Each category is independently chosen because each exercises a different +code path inside PyTorch's kernel dispatcher. A v1.0 test that passes +``row_major`` and fails ``broadcast`` has likely tripped over a missing +broadcast-aware kernel branch. + +The module is **lazy** with respect to torch — it raises +:class:`RuntimeError` on first call if ``torch`` isn't installed, mirroring +the rest of ``gpucheck.fuzzing``. +""" + +from __future__ import annotations + +from typing import Any + +CATEGORIES: tuple[str, ...] = ( + "row_major", + "column_major", + "broadcast", + "transpose", + "slice", + "non_contig", + "gather", +) + + +def _torch_mod() -> Any: + try: + import torch + + return torch + except ImportError as exc: # pragma: no cover -- exercised when torch absent + raise RuntimeError( + "gpucheck.fuzzing.strides requires PyTorch: pip install gpucheck[torch]" + ) from exc + + +def _row_major(shape: tuple[int, ...], dtype: Any, device: str, gen: Any) -> Any: + torch = _torch_mod() + return torch.randn(shape, generator=gen, dtype=torch.float32).to( + dtype=dtype, device=device, + ).contiguous() + + +def _column_major(shape: tuple[int, ...], dtype: Any, device: str, gen: Any) -> Any: + """Column-major: build the contiguous transposed shape, then transpose back. + + For ndim < 2 the concept is undefined; we fall back to row_major. + """ + torch = _torch_mod() + if len(shape) < 2: + return _row_major(shape, dtype, device, gen) + transposed = (shape[1], shape[0]) + shape[2:] + base = torch.randn(transposed, generator=gen, dtype=torch.float32).to( + dtype=dtype, device=device, + ).contiguous() + return base.transpose(0, 1) + + +def _broadcast(shape: tuple[int, ...], dtype: Any, device: str, gen: Any) -> Any: + """Broadcast-induced stride-0 dim along the LAST axis. + + For shape (M, N, K), build a contiguous (M, N, 1) tensor and expand to + (M, N, K). The last dim has stride 0 — kernels that scan strides + naively will multiply-count or read past bounds. + """ + torch = _torch_mod() + if not shape: + return torch.empty(shape, dtype=dtype, device=device) + base_shape = shape[:-1] + (1,) + base = torch.randn(base_shape, generator=gen, dtype=torch.float32).to( + dtype=dtype, device=device, + ).contiguous() + return base.expand(shape) + + +def _transpose(shape: tuple[int, ...], dtype: Any, device: str, gen: Any) -> Any: + """2D-style stride permutation: builds the transposed shape contiguous, + transposes, returns. Different from :func:`_column_major` only in that + transpose dims may be non-(0, 1) for higher-rank tensors — we transpose + the LAST two dims for ndim >= 2. + """ + torch = _torch_mod() + if len(shape) < 2: + return _row_major(shape, dtype, device, gen) + # Transpose last two dims, e.g. (B, M, N) -> build (B, N, M) contiguous, + # then .transpose(-1, -2) to recover (B, M, N) with permuted strides. + transposed = shape[:-2] + (shape[-1], shape[-2]) + base = torch.randn(transposed, generator=gen, dtype=torch.float32).to( + dtype=dtype, device=device, + ).contiguous() + return base.transpose(-1, -2) + + +def _slice(shape: tuple[int, ...], dtype: Any, device: str, gen: Any) -> Any: + """Regular non-unit stride via every-other slicing. + + Build a tensor with each dim doubled, then slice ``[::2, ::2, ...]``. + The resulting view has stride 2 in every dim. + """ + torch = _torch_mod() + if not shape: + return _row_major(shape, dtype, device, gen) + big_shape = tuple(d * 2 for d in shape) + base = torch.randn(big_shape, generator=gen, dtype=torch.float32).to( + dtype=dtype, device=device, + ).contiguous() + slicer = tuple(slice(None, None, 2) for _ in shape) + return base[slicer] + + +def _non_contig(shape: tuple[int, ...], dtype: Any, device: str, gen: Any) -> Any: + """Non-contiguous view that is NOT a clean transpose or slice. + + For ndim >= 3, we permute dims (1, 0, 2, ...). For ndim 2, we transpose + and then slice the last dim by 2 — guaranteed non-contiguous and not a + pure transpose. For ndim 1, fall back to slice. + """ + torch = _torch_mod() + if len(shape) <= 1: + return _slice(shape, dtype, device, gen) + if len(shape) == 2: + big = (shape[1], shape[0] * 2) + base = torch.randn(big, generator=gen, dtype=torch.float32).to( + dtype=dtype, device=device, + ).contiguous() + # transpose then slice the now-first dim by 2 + return base.transpose(0, 1)[::2] + # ndim >= 3 — permute first two dims + transposed = (shape[1], shape[0]) + shape[2:] + base = torch.randn(transposed, generator=gen, dtype=torch.float32).to( + dtype=dtype, device=device, + ).contiguous() + return base.transpose(0, 1) + + +def _gather(shape: tuple[int, ...], dtype: Any, device: str, gen: Any) -> Any: + """Irregular-access tensor: gather a contiguous source by a random index. + + The result is a contiguous tensor of the right shape, but it was + materialized via gather — so kernels that combine gather + reduction + in fused patterns may exhibit different behavior than a pure + contiguous input. (We return the gathered view contiguous; the test + harness's value is in *how* it was built, not the runtime layout.) + """ + torch = _torch_mod() + if not shape: + return _row_major(shape, dtype, device, gen) + src = torch.randn(shape, generator=gen, dtype=torch.float32).to( + dtype=dtype, device=device, + ).contiguous().flatten() + numel = src.numel() + idx = torch.randperm(numel, generator=gen) + return src[idx].reshape(shape).contiguous() + + +_CATEGORY_FN: dict[str, Any] = { + "row_major": _row_major, + "column_major": _column_major, + "broadcast": _broadcast, + "transpose": _transpose, + "slice": _slice, + "non_contig": _non_contig, + "gather": _gather, +} + + +def fuzz_strides_for_category( + shape: tuple[int, ...], + dtype: Any, + category: str, + *, + device: str = "cpu", + seed: int | None = None, +) -> Any: + """Build a single tensor of the given stride category. + + Use this when a test is parametrized over categories (typically via + :func:`parametrize_gpu(stride_categories=...)`). + """ + if category not in _CATEGORY_FN: + raise ValueError( + f"Unknown stride category {category!r}; " + f"expected one of {sorted(CATEGORIES)}" + ) + torch = _torch_mod() + gen: Any = None + if seed is not None: + gen = torch.Generator() + gen.manual_seed(seed) + return _CATEGORY_FN[category](shape, dtype, device, gen) + + +def fuzz_strides( + shape: tuple[int, ...], + dtype: Any, + *, + n: int | None = None, + device: str = "cpu", + seed: int | None = None, + categories: tuple[str, ...] | None = None, +) -> list[tuple[str, Any]]: + """Return a deterministic ``[(category, tensor), ...]`` corpus. + + Parameters + ---------- + shape: + Tensor shape used for every category. + dtype: + torch dtype for every tensor. + n: + Cap on the number of items returned. ``None`` returns all + configured categories (default 7). + device: + Target device string. + seed: + Optional torch RNG seed for reproducibility. + categories: + Override the default category order. Useful for tests that want + only a subset (e.g. only the non-contiguous flavors). + """ + cats = tuple(categories) if categories else CATEGORIES + invalid = [c for c in cats if c not in _CATEGORY_FN] + if invalid: + raise ValueError( + f"Unknown stride categories: {invalid}; " + f"expected from {sorted(CATEGORIES)}" + ) + out: list[tuple[str, Any]] = [] + for cat in cats: + out.append((cat, fuzz_strides_for_category( + shape, dtype, cat, device=device, seed=seed, + ))) + if n is not None: + out = out[:n] + return out + + +# --------------------------------------------------------------------------- +# Hypothesis strategy +# --------------------------------------------------------------------------- + +class StrideStrategy: + """Hypothesis-compatible factory that draws a stride-perturbed tensor. + + Mirrors :class:`gpucheck.fuzzing.ShapeStrategy`'s ``__new__``-as-factory + pattern so callers can write:: + + from hypothesis import given + @given(t=StrideStrategy(shape=(64, 64), dtype=torch.float32)) + def test_kernel_handles_strides(t): ... + + Hypothesis will draw one of the seven categories per test case and + shrink towards ``row_major``. + """ + + def __new__( + cls, + shape: tuple[int, ...], + dtype: Any = None, + *, + device: str = "cpu", + categories: tuple[str, ...] | None = None, + ) -> Any: + try: + from hypothesis import strategies as st + except ImportError as exc: + raise RuntimeError( + "StrideStrategy requires hypothesis: pip install gpucheck[hypothesis]" + ) from exc + + torch = _torch_mod() + if dtype is None: + dtype = torch.float32 + + cats = tuple(categories) if categories else CATEGORIES + + @st.composite # type: ignore[untyped-decorator] + def _draw(draw: Any) -> Any: + cat = draw(st.sampled_from(cats)) + seed = draw(st.integers(min_value=0, max_value=2**31 - 1)) + return fuzz_strides_for_category( + shape, dtype, cat, device=device, seed=seed, + ) + + return _draw() + + +__all__ = [ + "CATEGORIES", + "fuzz_strides", + "fuzz_strides_for_category", + "StrideStrategy", +] diff --git a/tests/test_fuzz_strides.py b/tests/test_fuzz_strides.py new file mode 100644 index 0000000..a2f1aaa --- /dev/null +++ b/tests/test_fuzz_strides.py @@ -0,0 +1,133 @@ +"""Stride fuzzing — corpus + per-category contracts (Track B).""" + +from __future__ import annotations + +import pytest + +from gpucheck.fuzzing.strides import ( + CATEGORIES, + fuzz_strides, + fuzz_strides_for_category, +) + +torch = pytest.importorskip("torch") + + +def test_categories_are_seven_canonical() -> None: + assert CATEGORIES == ( + "row_major", + "column_major", + "broadcast", + "transpose", + "slice", + "non_contig", + "gather", + ) + + +def test_fuzz_strides_returns_all_categories_in_order() -> None: + out = fuzz_strides((64, 64), torch.float32, seed=0) + assert [c for c, _t in out] == list(CATEGORIES) + + +def test_fuzz_strides_each_tensor_has_correct_shape() -> None: + out = fuzz_strides((32, 32), torch.float32, seed=0) + for label, t in out: + assert t.shape == (32, 32), f"{label}: wrong shape {t.shape}" + + +def test_fuzz_strides_each_tensor_has_correct_dtype() -> None: + out = fuzz_strides((16, 16), torch.float16, seed=0) + for label, t in out: + assert t.dtype == torch.float16, f"{label}: wrong dtype {t.dtype}" + + +def test_row_major_is_contiguous() -> None: + t = fuzz_strides_for_category((64, 64), torch.float32, "row_major", seed=0) + assert t.is_contiguous() + + +def test_column_major_is_not_contiguous() -> None: + t = fuzz_strides_for_category((64, 32), torch.float32, "column_major", seed=0) + assert not t.is_contiguous() + + +def test_broadcast_has_stride_zero_on_last_dim() -> None: + t = fuzz_strides_for_category((4, 8, 16), torch.float32, "broadcast", seed=0) + # The last dim was expanded from size 1 -> 16, so stride is 0 there. + assert t.stride(-1) == 0 + assert t.shape == (4, 8, 16) + + +def test_transpose_is_not_contiguous() -> None: + t = fuzz_strides_for_category((4, 8, 16), torch.float32, "transpose", seed=0) + assert not t.is_contiguous() + # Transpose preserves shape because we transposed the LAST two dims of + # the contiguous (4, 16, 8) buffer. + assert t.shape == (4, 8, 16) + + +def test_slice_has_stride_two_on_each_dim() -> None: + t = fuzz_strides_for_category((8, 8), torch.float32, "slice", seed=0) + assert not t.is_contiguous() + assert t.shape == (8, 8) + + +def test_non_contig_2d_is_not_contiguous() -> None: + t = fuzz_strides_for_category((8, 8), torch.float32, "non_contig", seed=0) + assert not t.is_contiguous() + assert t.shape == (8, 8) + + +def test_gather_returns_contiguous_with_correct_shape() -> None: + t = fuzz_strides_for_category((4, 4), torch.float32, "gather", seed=0) + assert t.is_contiguous() + assert t.shape == (4, 4) + + +def test_unknown_category_raises() -> None: + with pytest.raises(ValueError, match="Unknown stride category"): + fuzz_strides_for_category((4,), torch.float32, "weird") + + +def test_fuzz_strides_with_n_caps_results() -> None: + out = fuzz_strides((4, 4), torch.float32, n=3, seed=0) + assert len(out) == 3 + + +def test_fuzz_strides_with_explicit_categories() -> None: + out = fuzz_strides( + (4, 4), torch.float32, + categories=("row_major", "transpose"), + seed=0, + ) + assert [c for c, _t in out] == ["row_major", "transpose"] + + +def test_fuzz_strides_invalid_category_raises() -> None: + with pytest.raises(ValueError, match="Unknown stride categories"): + fuzz_strides((4, 4), torch.float32, categories=("row_major", "weird")) + + +def test_fuzz_strides_seed_is_deterministic() -> None: + out1 = fuzz_strides((4, 4), torch.float32, seed=42) + out2 = fuzz_strides((4, 4), torch.float32, seed=42) + for (l1, t1), (l2, t2) in zip(out1, out2, strict=True): + assert l1 == l2 + assert torch.equal(t1, t2), f"seed-determinism broken for {l1}" + + +def test_1d_fallbacks_to_row_major_or_slice_correctly() -> None: + # 1D shape: column_major and transpose fall back to row_major, + # non_contig falls back to slice. + out = fuzz_strides((8,), torch.float32, seed=0) + # Just assert we got tensors of shape (8,) — the fallback semantics + # are documented in the docstring; this is a smoke check. + for _label, t in out: + assert t.shape == (8,) + + +def test_higher_rank_3d_smoke() -> None: + out = fuzz_strides((4, 8, 16), torch.float32, seed=0) + for _label, t in out: + assert t.shape == (4, 8, 16) diff --git a/tests/test_fuzz_strides_hypothesis.py b/tests/test_fuzz_strides_hypothesis.py new file mode 100644 index 0000000..3bbc94f --- /dev/null +++ b/tests/test_fuzz_strides_hypothesis.py @@ -0,0 +1,39 @@ +"""Stride fuzzing — Hypothesis StrideStrategy (Track B).""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") +hypothesis = pytest.importorskip("hypothesis") + +from hypothesis import given, settings # noqa: E402 + +from gpucheck.fuzzing.strides import StrideStrategy # noqa: E402 + + +@settings(max_examples=20, deadline=None) +@given(t=StrideStrategy(shape=(8, 8), dtype=torch.float32)) +def test_stride_strategy_yields_tensor_of_target_shape(t) -> None: + assert t.shape == (8, 8) + assert t.dtype == torch.float32 + + +@settings(max_examples=10, deadline=None) +@given( + t=StrideStrategy( + shape=(4, 4), + dtype=torch.float32, + categories=("row_major",), + ), +) +def test_stride_strategy_with_single_category_only_returns_that_category(t) -> None: + # row_major is contiguous by construction. + assert t.is_contiguous() + assert t.shape == (4, 4) + + +def test_stride_strategy_unknown_dtype_does_not_crash_creation() -> None: + s = StrideStrategy(shape=(2, 2), dtype=torch.float16) + # Strategy creation must succeed; drawing also must not crash. + assert s is not None diff --git a/tests/test_parametrize_gpu_strides.py b/tests/test_parametrize_gpu_strides.py new file mode 100644 index 0000000..6491c89 --- /dev/null +++ b/tests/test_parametrize_gpu_strides.py @@ -0,0 +1,48 @@ +"""parametrize_gpu(stride_categories=...) wiring (Track B).""" + +from __future__ import annotations + +import pytest + +from gpucheck.decorators.parametrize import parametrize_gpu + +torch = pytest.importorskip("torch") + + +@parametrize_gpu( + dtypes=("float32",), + shapes=((4, 4),), + devices=("cpu",), + stride_categories=("row_major", "transpose"), +) +def test_stride_categories_appear_in_signature(dtype, shape, device, stride_category) -> None: + assert stride_category in {"row_major", "transpose"} + assert shape == (4, 4) + assert device == "cpu" + + +def test_parametrize_gpu_rejects_unknown_stride_category() -> None: + with pytest.raises(ValueError, match="Unknown stride categories"): + parametrize_gpu( + dtypes=("float32",), + shapes=((4, 4),), + devices=("cpu",), + stride_categories=("row_major", "wat"), + ) + + +def test_parametrize_gpu_without_stride_categories_keeps_old_signature() -> None: + decorator = parametrize_gpu( + dtypes=("float32",), + shapes=((4, 4),), + devices=("cpu",), + ) + + # The decorator marker name should not contain stride_category. + @decorator + def _fake_test(dtype, shape, device) -> None: # noqa: ARG001 + pass + + # Inspect the param names attached by pytest.mark.parametrize: + marks = list(_fake_test.pytestmark) + assert any("stride_category" not in m.args[0] for m in marks) From 5ddd26e5b7c3419a0cd757d8af639d626fc7c569 Mon Sep 17 00:00:00 2001 From: Akash Date: Fri, 1 May 2026 09:57:44 +0530 Subject: [PATCH 03/27] fix(tolerances): thread-safe override stack via contextvars; mitigate TM-E1 Track-C of the gpucheck v1.0 release fixes the documented "Thread-safety issue in tolerance override stack" gap (CLAUDE.md weakness) and addresses security finding TM-E1. Tolerance override stack: - Replace module-level list `_tolerance_overrides` with a `contextvars.ContextVar`. Each OS thread (and each asyncio task that copies the current context) sees its own override stack. - `tolerance_context(atol, rtol)` now uses ContextVar.set / .reset(token), which is exception-safe by construction. - The user-facing API is unchanged: `with tolerance_context(...): ...` TM-E1 mitigation in sanitizers/race.py: - _find_compute_sanitizer normalizes CUDA_HOME / CUDA_PATH via os.path.realpath and validates against _CUDA_HOME_ALLOWLIST (/usr/local/cuda, /opt/nvidia/cuda, /opt/cuda). - A symlink pointing outside the allowlist is correctly rejected. - A path that lookalikes a prefix (e.g. /usr/local/cuda-evil) is rejected by exact-prefix-with-separator matching. - Rejected paths emit a RuntimeWarning explaining the rejection. Tests added: - tests/test_tolerance_thread_safety.py - test_tolerance_context_is_thread_isolated: 4-thread Barrier-coordinated stress test; fails on the unfixed plain-list code, passes on ContextVar. - test_tolerance_context_pop_is_correct_after_exception - test_tolerance_context_nesting_in_single_thread - tests/test_race_cuda_home_allowlist.py (6 tests covering the canonical paths, lookalikes, missing PATH, allowlisted binary, symlink attack) Test count: 117 baseline -> 126 passing (+9 net new). ruff and mypy strict pass clean. Co-Authored-By: Claude Opus 4.7 --- src/gpucheck/assertions/tolerances.py | 37 +++++-- src/gpucheck/sanitizers/race.py | 57 ++++++++++- tests/test_race_cuda_home_allowlist.py | 127 ++++++++++++++++++++++++ tests/test_tolerance_thread_safety.py | 131 +++++++++++++++++++++++++ 4 files changed, 339 insertions(+), 13 deletions(-) create mode 100644 tests/test_race_cuda_home_allowlist.py create mode 100644 tests/test_tolerance_thread_safety.py diff --git a/src/gpucheck/assertions/tolerances.py b/src/gpucheck/assertions/tolerances.py index c192a26..91474b6 100644 --- a/src/gpucheck/assertions/tolerances.py +++ b/src/gpucheck/assertions/tolerances.py @@ -4,6 +4,7 @@ import math from contextlib import contextmanager +from contextvars import ContextVar from typing import TYPE_CHECKING, Any if TYPE_CHECKING: @@ -23,9 +24,23 @@ "tf32": (5e-4, 5e-4), } -# Override stack (module-level). NOT thread-safe — each thread/worker should use -# its own process (pytest-xdist worker) for parallel test execution. -_tolerance_overrides: list[tuple[float, float]] = [] +# Override stack (per-context). Backed by ``contextvars.ContextVar`` so the +# stack is isolated per OS thread AND per asyncio task. Previous releases +# used a plain module-level list, which leaked overrides between threads +# when tests were run inside a single process (e.g. with a thread pool). +# The contract is unchanged from the user's perspective: +# +# with tolerance_context(atol=1e-3, rtol=1e-3): +# assert_close(a, b) +# +# Inside the ``with`` block, the calling thread/task observes its own +# top-of-stack overlay; a sibling thread that has not entered a +# tolerance_context block sees the underlying defaults. ContextVar.set +# returns a Token that ``ContextVar.reset`` consumes, restoring the prior +# value — this is correct under exception unwinding. +_tolerance_overrides: ContextVar[tuple[tuple[float, float], ...]] = ContextVar( + "_tolerance_overrides", default=(), +) def _normalize_dtype_name(dtype: Any) -> str: @@ -52,9 +67,10 @@ def compute_tolerance( Falls back to float32 tolerances for unknown dtypes. """ - # Check override stack first. - if _tolerance_overrides: - return _tolerance_overrides[-1] + # Check override stack first (ContextVar for thread/task isolation). + overrides = _tolerance_overrides.get() + if overrides: + return overrides[-1] name = _normalize_dtype_name(dtype) # Check config overlay first, then defaults @@ -76,16 +92,21 @@ def tolerance_context( ) -> Generator[None, None, None]: """Temporarily override default tolerances returned by :func:`compute_tolerance`. + Backed by ``contextvars.ContextVar``: the override is visible only to + the current OS thread (and to asyncio tasks that copied the current + context). Sibling threads observe the underlying defaults concurrently. + Usage:: with tolerance_context(atol=1e-3, rtol=1e-3): assert_close(a, b) """ - _tolerance_overrides.append((atol, rtol)) + current = _tolerance_overrides.get() + token = _tolerance_overrides.set(current + ((atol, rtol),)) try: yield finally: - _tolerance_overrides.pop() + _tolerance_overrides.reset(token) def tolerances_from_config(config: dict[str, Any]) -> dict[str, tuple[float, float]] | None: diff --git a/src/gpucheck/sanitizers/race.py b/src/gpucheck/sanitizers/race.py index aec6b59..7f482cc 100644 --- a/src/gpucheck/sanitizers/race.py +++ b/src/gpucheck/sanitizers/race.py @@ -8,6 +8,7 @@ import subprocess import sys import tempfile +import warnings from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal @@ -17,6 +18,17 @@ SanitizerTool = Literal["memcheck", "racecheck", "initcheck", "synccheck"] _VALID_TOOLS: frozenset[str] = frozenset({"memcheck", "racecheck", "initcheck", "synccheck"}) +# Allowlist of canonical CUDA install prefixes. ``CUDA_HOME`` / ``CUDA_PATH`` +# values are normalized via ``os.path.realpath`` and rejected if they +# resolve outside this set. Mitigates security finding TM-E1: an attacker +# who can set the env var should not be able to redirect gpucheck into +# executing an arbitrary binary named ``compute-sanitizer``. +_CUDA_HOME_ALLOWLIST: tuple[str, ...] = ( + "/usr/local/cuda", + "/opt/nvidia/cuda", + "/opt/cuda", +) + @dataclass(frozen=True, slots=True) class SanitizerError: @@ -48,20 +60,55 @@ def error_count(self) -> int: def _find_compute_sanitizer() -> str | None: - """Locate compute-sanitizer binary on PATH or in CUDA_HOME.""" + """Locate compute-sanitizer binary on PATH or in CUDA_HOME. + + ``CUDA_HOME`` / ``CUDA_PATH`` env vars are normalized via + ``os.path.realpath`` and validated against + :data:`_CUDA_HOME_ALLOWLIST` before being trusted. Mitigates security + finding TM-E1. + """ path = shutil.which("compute-sanitizer") if path: return path cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH", "") - if cuda_home: - candidate = os.path.join(cuda_home, "bin", "compute-sanitizer") - if os.path.isfile(candidate) and os.access(candidate, os.X_OK): - return candidate + if not cuda_home: + return None + + # Resolve symlinks so an attacker can't bypass the allowlist by + # planting a symlink that points outside the trusted prefixes. + real = os.path.realpath(cuda_home) + if not _is_allowed_cuda_home(real): + warnings.warn( + f"CUDA_HOME / CUDA_PATH={cuda_home!r} resolves to {real!r} which is " + f"outside the allowlist {_CUDA_HOME_ALLOWLIST!r}; ignoring " + f"(set CUDA_HOME to a path under one of those prefixes, or " + f"install compute-sanitizer onto PATH).", + RuntimeWarning, + stacklevel=2, + ) + return None + + candidate = os.path.join(real, "bin", "compute-sanitizer") + if os.path.isfile(candidate) and os.access(candidate, os.X_OK): + return candidate return None +def _is_allowed_cuda_home(real_path: str) -> bool: + """Return ``True`` if *real_path* is inside one of the allowlist prefixes. + + The check is exact-prefix-with-separator so ``/usr/local/cuda-evil`` + does NOT match ``/usr/local/cuda``. + """ + norm = os.path.normpath(real_path) + for prefix in _CUDA_HOME_ALLOWLIST: + if norm == prefix or norm.startswith(prefix + os.sep): + return True + return False + + def _parse_sanitizer_output( raw: str, tool: SanitizerTool, ) -> tuple[list[SanitizerError], list[str]]: diff --git a/tests/test_race_cuda_home_allowlist.py b/tests/test_race_cuda_home_allowlist.py new file mode 100644 index 0000000..68fa84d --- /dev/null +++ b/tests/test_race_cuda_home_allowlist.py @@ -0,0 +1,127 @@ +"""TM-E1 mitigation: CUDA_HOME / CUDA_PATH allowlist (Track C).""" + +from __future__ import annotations + +import os +import warnings +from typing import TYPE_CHECKING + +from gpucheck.sanitizers.race import ( + _CUDA_HOME_ALLOWLIST, + _find_compute_sanitizer, + _is_allowed_cuda_home, +) + +if TYPE_CHECKING: + from pathlib import Path + + import pytest + + +def test_is_allowed_cuda_home_canonical_paths() -> None: + for prefix in _CUDA_HOME_ALLOWLIST: + assert _is_allowed_cuda_home(prefix) is True + assert _is_allowed_cuda_home(prefix + "/bin") is True + assert _is_allowed_cuda_home(prefix + "/12.2") is True + + +def test_is_allowed_cuda_home_rejects_lookalike_paths() -> None: + # Trailing characters must NOT match the prefix. + assert _is_allowed_cuda_home("/usr/local/cuda-evil") is False + assert _is_allowed_cuda_home("/opt/nvidia/cudawat") is False + assert _is_allowed_cuda_home("/opt") is False + assert _is_allowed_cuda_home("/tmp/attacker") is False + + +def test_find_compute_sanitizer_returns_none_when_path_lookup_fails( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, +) -> None: + # Ensure shutil.which fails (point PATH at an empty dir). + monkeypatch.setenv("PATH", str(tmp_path)) + monkeypatch.delenv("CUDA_HOME", raising=False) + monkeypatch.delenv("CUDA_PATH", raising=False) + + assert _find_compute_sanitizer() is None + + +def test_find_compute_sanitizer_rejects_outside_allowlist( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, +) -> None: + """A CUDA_HOME outside the allowlist must be ignored AND emit a warning.""" + monkeypatch.setenv("PATH", str(tmp_path)) # neutralize shutil.which path + + fake_cuda = tmp_path / "fake_cuda" + (fake_cuda / "bin").mkdir(parents=True) + binary = fake_cuda / "bin" / "compute-sanitizer" + binary.write_text("#!/bin/sh\nexit 0\n") + binary.chmod(0o755) + + monkeypatch.setenv("CUDA_HOME", str(fake_cuda)) + monkeypatch.delenv("CUDA_PATH", raising=False) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + result = _find_compute_sanitizer() + + assert result is None, ( + "fake CUDA_HOME outside allowlist must NOT yield a sanitizer path" + ) + assert any("allowlist" in str(w.message) for w in caught), ( + "expected a RuntimeWarning explaining the allowlist rejection" + ) + + +def test_find_compute_sanitizer_accepts_inside_allowlist( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, +) -> None: + """If CUDA_HOME points inside the allowlist AND a binary exists there, return it. + + We can't actually create files at /usr/local/cuda in CI; instead, we + monkeypatch _CUDA_HOME_ALLOWLIST to include tmp_path and verify the + code path returns the binary when the rest of the conditions hold. + """ + monkeypatch.setenv("PATH", str(tmp_path / "no_path_here")) + + real_cuda = tmp_path / "real_cuda" + (real_cuda / "bin").mkdir(parents=True) + binary = real_cuda / "bin" / "compute-sanitizer" + binary.write_text("#!/bin/sh\nexit 0\n") + binary.chmod(0o755) + + monkeypatch.setattr( + "gpucheck.sanitizers.race._CUDA_HOME_ALLOWLIST", + (str(real_cuda.resolve()),), + ) + monkeypatch.setenv("CUDA_HOME", str(real_cuda)) + monkeypatch.delenv("CUDA_PATH", raising=False) + + result = _find_compute_sanitizer() + assert result == os.path.join(str(real_cuda.resolve()), "bin", "compute-sanitizer") + + +def test_find_compute_sanitizer_resolves_symlink_before_allowlist( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, +) -> None: + """A symlink pointing OUTSIDE the allowlist must be rejected. + + This guards against the obvious attack: + ln -s /tmp/attacker /usr/local/cuda + """ + monkeypatch.setenv("PATH", str(tmp_path / "nope")) + + attacker = tmp_path / "attacker" + (attacker / "bin").mkdir(parents=True) + (attacker / "bin" / "compute-sanitizer").write_text("#!/bin/sh\nexit 0\n") + (attacker / "bin" / "compute-sanitizer").chmod(0o755) + + symlink_at_canonical = tmp_path / "symlinked_cuda" + symlink_at_canonical.symlink_to(attacker) + + monkeypatch.setenv("CUDA_HOME", str(symlink_at_canonical)) + monkeypatch.delenv("CUDA_PATH", raising=False) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + result = _find_compute_sanitizer() + assert result is None + assert any("allowlist" in str(w.message) for w in caught) diff --git a/tests/test_tolerance_thread_safety.py b/tests/test_tolerance_thread_safety.py new file mode 100644 index 0000000..58bd9f8 --- /dev/null +++ b/tests/test_tolerance_thread_safety.py @@ -0,0 +1,131 @@ +"""Thread-safety regression test for the tolerance override stack (Track C). + +Without the ContextVar fix, ``tolerance_context`` mutates a shared +module-level list — overrides leak between concurrent threads, and a +thread reading inside its ``with`` block can observe a sibling thread's +override. + +This test is designed to fail loudly on the unfixed code AND pass on the +fixed code. It uses a ``threading.Barrier`` so all 4 threads enter their +context manager before any of them reads, maximizing observable contention. +""" + +from __future__ import annotations + +import threading +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from gpucheck.assertions.tolerances import compute_tolerance, tolerance_context + + +def _worker( + thread_id: int, + barrier: threading.Barrier, + enter_release: threading.Event, + leave_release: threading.Event, +) -> tuple[int, tuple[float, float]]: + """Each thread enters a context with thread-distinct overrides, then waits. + + The barrier ensures all 4 threads have entered their context before + ANY of them reads. Without ContextVar isolation, the read sees the + LAST writer's override (4 in 4 threads, race-y). + """ + my_atol = 1.0 + thread_id # 1.0, 2.0, 3.0, 4.0 + my_rtol = 0.1 + thread_id + + with tolerance_context(my_atol, my_rtol): + # All threads are now inside their context. Synchronize so reads + # happen with maximum overlap. + barrier.wait(timeout=5.0) + enter_release.wait(timeout=5.0) + + # Read what compute_tolerance sees from THIS thread's perspective. + observed = compute_tolerance("float32") + + # Hold the context open until the harness signals release. This + # increases the window during which a sibling thread's broken + # override could be observed. + leave_release.wait(timeout=5.0) + return thread_id, observed + + +def test_tolerance_context_is_thread_isolated() -> None: + """Each of N threads must observe its OWN context's atol/rtol. + + The plain-list implementation in pre-Track-C gpucheck serializes + appends but the read at compute_tolerance time picks the GLOBAL + top-of-stack — so all threads see the most recently entered + context's values. With ContextVar each thread has its own stack and + sees its own atol/rtol. + """ + n_threads = 4 + barrier = threading.Barrier(n_threads) + enter_release = threading.Event() + leave_release = threading.Event() + + with ThreadPoolExecutor(max_workers=n_threads) as exe: + futures = [ + exe.submit(_worker, tid, barrier, enter_release, leave_release) + for tid in range(n_threads) + ] + + # Allow workers to actually read once they've all entered. + enter_release.set() + + # Brief delay to let reads happen. ContextVar is correct under + # arbitrary interleavings; the sleep just makes the unfixed + # version's bug deterministic. + import time as _time + _time.sleep(0.05) + + # Now release everyone to leave their context. + leave_release.set() + + results = [f.result(timeout=10.0) for f in futures] + + # Verify each thread observed its OWN override. + for tid, (atol, rtol) in results: + expected_atol = 1.0 + tid + expected_rtol = 0.1 + tid + assert atol == pytest.approx(expected_atol), ( + f"thread {tid} observed atol={atol} expected {expected_atol} — " + f"tolerance_context is NOT thread-isolated" + ) + assert rtol == pytest.approx(expected_rtol), ( + f"thread {tid} observed rtol={rtol} expected {expected_rtol}" + ) + + +def test_tolerance_context_pop_is_correct_after_exception() -> None: + """Exception inside the with-block must still restore the prior state. + + ContextVar.reset(token) is exception-safe by construction; this is a + regression guard against accidental refactors that might break it. + """ + base_atol, _ = compute_tolerance("float32") + + class _BangError(RuntimeError): + pass + + with pytest.raises(_BangError), tolerance_context(99.0, 0.99): + raise _BangError("boom") + + after_atol, _ = compute_tolerance("float32") + assert after_atol == pytest.approx(base_atol), ( + "tolerance_context did not restore prior state after exception" + ) + + +def test_tolerance_context_nesting_in_single_thread() -> None: + """Nesting in one thread should observe LIFO order — innermost wins.""" + with tolerance_context(1.0, 0.1): + atol, _ = compute_tolerance("float32") + assert atol == pytest.approx(1.0) + with tolerance_context(2.0, 0.2): + atol2, _ = compute_tolerance("float32") + assert atol2 == pytest.approx(2.0) + # After inner exits, outer is restored. + atol3, _ = compute_tolerance("float32") + assert atol3 == pytest.approx(1.0) From 02507da76eb274b19dff0ba31261354d856e175d Mon Sep 17 00:00:00 2001 From: Akash Date: Fri, 1 May 2026 10:02:21 +0530 Subject: [PATCH 04/27] feat(reporting+sanitizers): HTML dashboard, determinism, lockfile, CI hardening MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Track-D of the gpucheck v1.0 release closes four documented gaps in one shot, splits across the reporting and sanitizers packages plus CI plumbing. D.1 reporting test coverage 0 -> 98% - tests/test_reporting_console.py (10 tests) — Rich-based renderer. - tests/test_reporting_json.py (6 tests) — RunRecord schema, compare_runs classifies regression / ok / new / removed / div-by-zero. - tests/test_reporting_ci.py (8 tests) — GitHub Actions annotations, JUnit XML well-formedness, PR comment Markdown. D.2 reporting/html.py — static HTML dashboard - Self-contained HTML (no external CSS/JS, no fetches at view time). - Inline SVG bar chart for benchmark medians. - Test results / benchmarks / memory / comparison sections. - Empty data handled gracefully. - HTMLParser well-formedness verified; XSS escaping verified. D.3 sanitizers/determinism.py — assert_deterministic + @requires_determinism - Seeds random / numpy / torch (CPU+CUDA+MPS) before each invocation. - Compares torch.Tensor / tuple / list / scalar outputs byte-identically. - DeterminismError surfaces SYNTHESIS §4 best-effort caveat in failure msg. - 8 tests covering pass/fail/exception/nesting/torch-tensor paths. D.4 DEP-1 mitigation — uv.lock committed - 1220-line lock file from existing dev environment. - CI now installs via `uv sync --frozen --extra dev`. D.5 CFG-2 mitigation — .github/workflows/ci.yml - Top-level `permissions: contents: read` block. - Default GITHUB_TOKEN is no longer write-all. Test count: 117 baseline -> 157 passing (+40 net new). Reporting coverage: 0% -> 98% (target was 90%). ruff and mypy strict pass clean. Co-Authored-By: Claude Opus 4.7 --- .github/workflows/ci.yml | 39 +- src/gpucheck/reporting/__init__.py | 2 + src/gpucheck/reporting/html.py | 255 +++++ src/gpucheck/sanitizers/__init__.py | 8 + src/gpucheck/sanitizers/determinism.py | 161 ++++ tests/test_determinism.py | 103 ++ tests/test_reporting_ci.py | 126 +++ tests/test_reporting_console.py | 115 +++ tests/test_reporting_html.py | 157 +++ tests/test_reporting_json.py | 109 +++ uv.lock | 1220 ++++++++++++++++++++++++ 11 files changed, 2288 insertions(+), 7 deletions(-) create mode 100644 src/gpucheck/reporting/html.py create mode 100644 src/gpucheck/sanitizers/determinism.py create mode 100644 tests/test_determinism.py create mode 100644 tests/test_reporting_ci.py create mode 100644 tests/test_reporting_console.py create mode 100644 tests/test_reporting_html.py create mode 100644 tests/test_reporting_json.py create mode 100644 uv.lock diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index edfe733..690ed9b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -6,6 +6,13 @@ on: pull_request: branches: [main] +# Security finding CFG-2 (FINDINGS.md): default GITHUB_TOKEN scope is +# write-all. Explicitly restrict to read-only at the workflow level so +# nothing in this CI job can accidentally publish artifacts or comments +# unless a step elevates locally. +permissions: + contents: read + jobs: lint: runs-on: ubuntu-latest @@ -14,12 +21,22 @@ jobs: - uses: actions/setup-python@v5 with: python-version: "3.12" - - name: Install dependencies - run: pip install -e ".[dev]" + - name: Install uv + uses: astral-sh/setup-uv@v3 + - name: Install dependencies (locked) + # Use the committed uv.lock for reproducibility and supply-chain + # protection (security finding DEP-1). + run: | + uv venv + if [ -f uv.lock ]; then + uv sync --frozen --extra dev + else + uv pip install -e ".[dev]" + fi - name: Ruff check - run: ruff check src/ tests/ + run: uv run ruff check src/ tests/ - name: Mypy - run: mypy src/ + run: uv run mypy src/ test: runs-on: ubuntu-latest @@ -32,7 +49,15 @@ jobs: - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - - name: Install dependencies - run: pip install -e ".[dev]" + - name: Install uv + uses: astral-sh/setup-uv@v3 + - name: Install dependencies (locked) + run: | + uv venv + if [ -f uv.lock ]; then + uv sync --frozen --extra dev + else + uv pip install -e ".[dev]" + fi - name: Run tests - run: pytest --tb=short -q + run: uv run pytest --tb=short -q diff --git a/src/gpucheck/reporting/__init__.py b/src/gpucheck/reporting/__init__.py index ec31026..6679bef 100644 --- a/src/gpucheck/reporting/__init__.py +++ b/src/gpucheck/reporting/__init__.py @@ -11,6 +11,7 @@ "emit_github_annotations": ("gpucheck.reporting.ci", "emit_github_annotations"), "write_junit_xml": ("gpucheck.reporting.ci", "write_junit_xml"), "generate_pr_comment": ("gpucheck.reporting.ci", "generate_pr_comment"), + "HTMLReporter": ("gpucheck.reporting.html", "HTMLReporter"), } @@ -28,4 +29,5 @@ def __getattr__(name: str) -> Any: "emit_github_annotations", "write_junit_xml", "generate_pr_comment", + "HTMLReporter", ] diff --git a/src/gpucheck/reporting/html.py b/src/gpucheck/reporting/html.py new file mode 100644 index 0000000..e1ec7db --- /dev/null +++ b/src/gpucheck/reporting/html.py @@ -0,0 +1,255 @@ +"""Static HTML dashboard generator for gpucheck JSON run records. + +Reads a ``results.json`` produced by :class:`gpucheck.reporting.json.JSONReporter` +and writes a single self-contained HTML file: zero external CSS, zero +external JavaScript, no fetches at view time. Inline SVG renders the +benchmark bar chart so the file works on a flight without WiFi. + +Sections (in order): + +1. **Summary** — total tests, pass count, fail count, skip count, GPU info. +2. **Test results table** — one row per test with status pill and + collapsed message via ``
``. +3. **Benchmark table** — one row per kernel; inline SVG bar chart of + median timings for at-a-glance regression spotting. +4. **Memory table** — peak / leaked MB per test. +5. **Comparison band** — when a comparison diff is supplied, surfaces + regression / ok / new / removed rows in red / green / blue / gray. + +The renderer is deliberately small (no Jinja, no D3) so it vendorizes +cleanly. Callers needing richer charts can post-process the JSON in any +external dashboard. +""" + +from __future__ import annotations + +import html +import json +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +_STATUS_PILL_BG: dict[str, str] = { + "passed": "#1f7a4d", + "failed": "#a8231f", + "error": "#a8231f", + "skipped": "#a87a1f", + "ok": "#1f7a4d", + "regression": "#a8231f", + "new": "#1f5fa8", + "removed": "#666666", +} + + +@dataclass +class HTMLReporter: + """Render a JSON run record into a self-contained HTML file.""" + + json_path: str | Path + comparison: dict[str, Any] | None = None + title: str = "gpucheck dashboard" + _data: dict[str, Any] = field(default_factory=dict, init=False, repr=False) + + def _load(self) -> dict[str, Any]: + if not self._data: + self._data = json.loads(Path(self.json_path).read_text(encoding="utf-8")) + return self._data + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def render(self, out_path: str | Path) -> Path: + """Write the dashboard to *out_path* and return its :class:`Path`.""" + data = self._load() + out = Path(out_path) + out.parent.mkdir(parents=True, exist_ok=True) + + sections = [ + _render_head(self.title), + _render_summary(data), + _render_test_results(data), + _render_benchmarks(data), + _render_memory(data), + _render_comparison(self.comparison) if self.comparison else "", + _render_foot(), + ] + out.write_text("\n".join(s for s in sections if s) + "\n", encoding="utf-8") + return out + + +# --------------------------------------------------------------------------- +# Section renderers +# --------------------------------------------------------------------------- + + +def _esc(value: Any) -> str: + return html.escape(str(value), quote=True) + + +def _pill(status: str) -> str: + bg = _STATUS_PILL_BG.get(status, "#666666") + return ( + f'' + f"{_esc(status.upper())}" + ) + + +def _render_head(title: str) -> str: + return f""" + +{_esc(title)} + + +

{_esc(title)}

""" + + +def _render_foot() -> str: + return "" + + +def _render_summary(data: dict[str, Any]) -> str: + results = data.get("test_results", []) + passed = sum(1 for r in results if r.get("status") == "passed") + failed = sum(1 for r in results if r.get("status") == "failed") + skipped = sum(1 for r in results if r.get("status") == "skipped") + + gpu_info = data.get("gpu_info", {}) or {} + gpu_summary = " | ".join( + f"{_esc(k)}: {_esc(v)}" for k, v in list(gpu_info.items())[:5] + ) or "no GPU info recorded" + + return f"""

Summary

+
+
{passed}
passed
+
{failed}
failed
+
{skipped}
skipped
+
{gpu_summary}
+
+

timestamp: {_esc(data.get("timestamp", "unknown"))}

""" + + +def _render_test_results(data: dict[str, Any]) -> str: + results = data.get("test_results", []) + if not results: + return "" + rows = [] + for r in results: + status = r.get("status", "unknown") + klass = "regression" if status in {"failed", "error"} else "passed-row" + msg = r.get("message", "") + msg_cell = ( + f'
view
{_esc(msg)}
' if msg else "" + ) + rows.append( + f'{_esc(r.get("name", ""))}' + f'{_pill(status)}' + f'{r.get("duration", 0.0):.4f}s' + f'{msg_cell}', + ) + body = "\n".join(rows) + return f"""

Test Results

+ +{body} +
TestStatusDurationMessage
""" + + +def _render_benchmarks(data: dict[str, Any]) -> str: + benches = data.get("benchmarks", []) + if not benches: + return "" + max_med = max((b.get("median_ms", 0.0) or 0.0) for b in benches) or 1.0 + + rows = [] + for b in benches: + med = float(b.get("median_ms", 0.0) or 0.0) + std = float(b.get("std_ms", 0.0) or 0.0) + bar_w = max(2, int(180 * (med / max_med))) + rows.append( + f'{_esc(b.get("name", ""))}' + f'{med:.3f} ms' + f'{std:.3f} ms' + f'{b.get("samples", 0)}' + f'', + ) + return f"""

Benchmarks

+ + +{"".join(rows)} +
KernelMedianStdSamplesDistribution (relative)
""" + + +def _render_memory(data: dict[str, Any]) -> str: + mem = data.get("memory", []) + if not mem: + return "" + rows = [] + for m in mem: + leaked = float(m.get("leaked_mb", 0.0) or 0.0) + klass = "regression" if leaked > 0 else "passed-row" + rows.append( + f'{_esc(m.get("name", ""))}' + f'{m.get("peak_mb", 0):.2f} MB' + f'{leaked:.2f} MB' + f'{m.get("allocations", 0)}', + ) + return f"""

Memory

+ +{"".join(rows)} +
TestPeakLeakedAllocations
""" + + +def _render_comparison(diff: dict[str, Any]) -> str: + benches = diff.get("benchmarks", []) + if not benches: + return "" + rows = [] + for b in benches: + status = b.get("status", "ok") + klass = status if status in {"regression", "new", "removed"} else "passed-row" + base = b.get("baseline_median_ms", "-") + curr = b.get("current_median_ms", "-") + delta = b.get("delta_pct", 0) + base_str = f"{base:.3f} ms" if isinstance(base, (int, float)) else _esc(base) + curr_str = f"{curr:.3f} ms" if isinstance(curr, (int, float)) else _esc(curr) + rows.append( + f'{_esc(b.get("name", ""))}' + f'{base_str}{curr_str}' + f'{delta:+.1f}%' + f'{_pill(status)}', + ) + return f"""

Comparison vs Baseline

+ + +{"".join(rows)} +
KernelBaselineCurrentDeltaStatus
""" + + +__all__ = ["HTMLReporter"] diff --git a/src/gpucheck/sanitizers/__init__.py b/src/gpucheck/sanitizers/__init__.py index fc5aa12..117159e 100644 --- a/src/gpucheck/sanitizers/__init__.py +++ b/src/gpucheck/sanitizers/__init__.py @@ -2,6 +2,11 @@ from __future__ import annotations +from gpucheck.sanitizers.determinism import ( + DeterminismError, + assert_deterministic, + requires_determinism, +) from gpucheck.sanitizers.memory import SanitizerMemoryReport, check_memory_leaks, memory_guard from gpucheck.sanitizers.race import SanitizerReport, run_with_sanitizer @@ -15,4 +20,7 @@ "check_memory_leaks", "memory_guard", "run_with_sanitizer", + "assert_deterministic", + "requires_determinism", + "DeterminismError", ] diff --git a/src/gpucheck/sanitizers/determinism.py b/src/gpucheck/sanitizers/determinism.py new file mode 100644 index 0000000..e2d5a44 --- /dev/null +++ b/src/gpucheck/sanitizers/determinism.py @@ -0,0 +1,161 @@ +"""Determinism sanitizer — assert byte-identical outputs across runs. + +Unlike CUDA, MPS is best-effort deterministic (research SYNTHESIS §4): +PyTorch documentation is silent on Apple Silicon determinism guarantees, +and the empirical record (pytorch#181936, #170837, #177116) shows real +run-to-run divergence. This module provides: + +- :func:`assert_deterministic` — runs ``fn`` ``n`` times under fixed seeds + and asserts every output tensor is byte-identical to the first run. + On MPS, structured failure surfaces ``DeterminismError`` so callers + know whether the divergence is at the precision floor (acceptable in + some pipelines) or a literal inconsistency. +- :func:`requires_determinism` — function decorator: wraps the test body + so that calling it n times is the test (instead of the test author + having to write the loop themselves). + +The seeded run sets: + + torch.manual_seed(seed) + if mps available: torch.mps.manual_seed(seed) + if cuda available: torch.cuda.manual_seed_all(seed) + random.seed(seed); numpy.random.seed(seed) +""" + +from __future__ import annotations + +import functools +import random +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Callable + + +class DeterminismError(AssertionError): + """Raised when ``assert_deterministic`` observes diverging outputs.""" + + +def _seed_all(seed: int) -> None: + """Seed every RNG we know about. Best-effort — silent on missing modules.""" + random.seed(seed) + try: + import numpy as np + + np.random.seed(seed) + except ImportError: + pass + try: + import torch + + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + mps = getattr(torch.backends, "mps", None) + if mps is not None and mps.is_available(): + seed_fn = getattr(torch.mps, "manual_seed", None) + if callable(seed_fn): + seed_fn(seed) + except ImportError: + pass + + +def _equal(a: Any, b: Any) -> bool: + """Compare two outputs for byte-identical equality. + + Handles torch.Tensor (same device, same dtype), tuples / lists + elementwise, and falls back to ``==``. + """ + try: + import torch + + if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor): + if a.shape != b.shape or a.dtype != b.dtype or a.device != b.device: + return False + return bool(torch.equal(a, b)) + except ImportError: + pass + if isinstance(a, (tuple, list)) and isinstance(b, (tuple, list)): + if len(a) != len(b): + return False + return all(_equal(x, y) for x, y in zip(a, b, strict=False)) + return bool(a == b) + + +def assert_deterministic( + fn: Callable[..., Any], + *args: Any, + n: int = 3, + seed: int = 0, + **kwargs: Any, +) -> Any: + """Run *fn* ``n`` times under fixed seeds; assert outputs match exactly. + + Parameters + ---------- + fn: + Callable producing the output to compare. May return a tensor, a + tuple of tensors, or any equality-comparable value. + *args / **kwargs: + Forwarded to *fn*. + n: + Number of repetitions. Must be ``>= 2``. + seed: + Seed applied to ``random``, ``numpy.random``, ``torch.manual_seed``, + ``torch.cuda.manual_seed_all``, and ``torch.mps.manual_seed`` (if + available) before each call. + + Returns + ------- + The output of the first run, so callers can pass through values + that they want to use after asserting determinism. + + Raises + ------ + DeterminismError: + If any run's output differs from the first run. + """ + if n < 2: + raise ValueError(f"assert_deterministic requires n >= 2, got {n}") + + _seed_all(seed) + first = fn(*args, **kwargs) + for i in range(1, n): + _seed_all(seed) + candidate = fn(*args, **kwargs) + if not _equal(first, candidate): + raise DeterminismError( + f"assert_deterministic: run {i} produced output that differs " + f"from run 0 (n={n}, seed={seed}). On MPS this can happen " + f"legitimately (best-effort determinism per SYNTHESIS §4); " + f"consider widening tolerances via tolerance_context, or " + f"adding the op to the [tool.gpucheck.mps.xfail] block." + ) + return first + + +def requires_determinism( + *, + n: int = 3, + seed: int = 0, +) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Decorator: invoke the test function ``n`` times under fixed seeds. + + Equivalent to wrapping the test body in + :func:`assert_deterministic`. Useful when the test's return value is + the artifact under test:: + + @requires_determinism(n=5, seed=42) + def test_my_kernel(): + x = torch.randn(64, 64, device="mps") + return my_kernel(x) + """ + def decorator(fn: Callable[..., Any]) -> Callable[..., Any]: + @functools.wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + return assert_deterministic(fn, *args, n=n, seed=seed, **kwargs) + return wrapper + return decorator + + +__all__ = ["assert_deterministic", "requires_determinism", "DeterminismError"] diff --git a/tests/test_determinism.py b/tests/test_determinism.py new file mode 100644 index 0000000..5886c13 --- /dev/null +++ b/tests/test_determinism.py @@ -0,0 +1,103 @@ +"""Determinism sanitizer tests (Track D — D.3).""" + +from __future__ import annotations + +import random + +import pytest + +from gpucheck.sanitizers.determinism import ( + DeterminismError, + assert_deterministic, + requires_determinism, +) + + +def test_assert_deterministic_passes_when_output_is_deterministic() -> None: + def deterministic_fn() -> int: + return random.randint(0, 100) # noqa: S311 (test, not crypto) + + # Seeded properly, this returns the same int each call. + out = assert_deterministic(deterministic_fn, n=4, seed=7) + assert isinstance(out, int) + + +def test_assert_deterministic_raises_when_output_diverges() -> None: + counter = {"i": 0} + + def diverging_fn() -> int: + counter["i"] += 1 + return counter["i"] # 1, 2, 3, ... — not seeded by random + + with pytest.raises(DeterminismError, match="differs"): + assert_deterministic(diverging_fn, n=3, seed=0) + + +def test_assert_deterministic_rejects_n_lt_2() -> None: + with pytest.raises(ValueError, match="n >= 2"): + assert_deterministic(lambda: 1, n=1) + + +def test_assert_deterministic_returns_first_output() -> None: + def fn() -> str: + return "stable" + + out = assert_deterministic(fn, n=3) + assert out == "stable" + + +def test_requires_determinism_decorator_calls_fn_n_times() -> None: + counter = {"i": 0} + + @requires_determinism(n=4, seed=99) + def fn() -> int: + counter["i"] += 1 + # Reset randomness here makes this deterministic across calls + # because the decorator calls _seed_all between invocations. + return random.randint(0, 1_000_000) # noqa: S311 + + fn() + assert counter["i"] == 4 + + +def test_requires_determinism_propagates_determinism_error_from_unstable_fn() -> None: + counter = {"i": 0} + + @requires_determinism(n=2, seed=0) + def fn() -> int: + counter["i"] += 1 + return counter["i"] + + with pytest.raises(DeterminismError): + fn() + + +def test_assert_deterministic_handles_tuple_outputs() -> None: + def fn() -> tuple[int, int]: + return (random.randint(0, 100), random.randint(0, 100)) # noqa: S311 + + out = assert_deterministic(fn, n=3, seed=42) + assert isinstance(out, tuple) + assert len(out) == 2 + + +def test_assert_deterministic_handles_list_outputs_diverging() -> None: + counter = {"i": 0} + + def fn() -> list[int]: + counter["i"] += 1 + return [counter["i"]] + + with pytest.raises(DeterminismError): + assert_deterministic(fn, n=2, seed=0) + + +def test_assert_deterministic_with_torch_tensor_outputs() -> None: + """When torch is available, tensor equality goes through torch.equal.""" + torch = pytest.importorskip("torch") + + def fn() -> torch.Tensor: + return torch.randn(3, 3) # seeded => same tensor + + out = assert_deterministic(fn, n=3, seed=0) + assert out.shape == (3, 3) diff --git a/tests/test_reporting_ci.py b/tests/test_reporting_ci.py new file mode 100644 index 0000000..022adc5 --- /dev/null +++ b/tests/test_reporting_ci.py @@ -0,0 +1,126 @@ +"""CI reporting tests (Track D — D.1).""" + +from __future__ import annotations + +import io +import sys +from typing import TYPE_CHECKING + +from gpucheck.reporting.ci import ( + emit_github_annotations, + generate_pr_comment, + write_junit_xml, +) +from gpucheck.reporting.console import TestResult + +if TYPE_CHECKING: + from pathlib import Path + + +def test_emit_github_annotations_writes_error_lines(monkeypatch) -> None: + monkeypatch.setenv("GITHUB_ACTIONS", "1") + captured = io.StringIO() + monkeypatch.setattr(sys, "stdout", captured) + results = [ + TestResult( + name="tests/test_x.py::test_y", status="failed", + duration=0.1, message="AssertionError\nexpected != actual", + file="tests/test_x.py", line=42, + ), + ] + emit_github_annotations(results) + out = captured.getvalue() + assert "::error" in out + assert "file=tests/test_x.py" in out + assert "line=42" in out + assert "%0A" in out # newline encoded in annotation message + + +def test_emit_github_annotations_skipped_warning(monkeypatch) -> None: + monkeypatch.setenv("GITHUB_ACTIONS", "1") + captured = io.StringIO() + monkeypatch.setattr(sys, "stdout", captured) + results = [TestResult(name="t1", status="skipped", message="no gpu")] + emit_github_annotations(results) + assert "::warning" in captured.getvalue() + + +def test_emit_github_annotations_passed_is_silent(monkeypatch) -> None: + monkeypatch.setenv("GITHUB_ACTIONS", "1") + captured = io.StringIO() + monkeypatch.setattr(sys, "stdout", captured) + emit_github_annotations([TestResult(name="t1", status="passed")]) + assert captured.getvalue() == "" + + +def test_emit_github_annotations_no_op_outside_actions(monkeypatch) -> None: + monkeypatch.delenv("GITHUB_ACTIONS", raising=False) + captured = io.StringIO() + monkeypatch.setattr(sys, "stdout", captured) + emit_github_annotations([TestResult(name="t1", status="failed", message="boom")]) + assert captured.getvalue() == "" + + +def test_write_junit_xml_emits_valid_xml(tmp_path: Path) -> None: + out = tmp_path / "junit.xml" + results = [ + TestResult(name="t1", status="passed", duration=0.1), + TestResult(name="t2", status="failed", duration=0.2, message="boom"), + TestResult(name="t3", status="skipped", duration=0.0, message="no gpu"), + TestResult(name="t4", status="error", duration=0.05, message="crashed"), + ] + path = write_junit_xml(results, output_path=out, suite_name="gpucheck") + + assert path == out + text = out.read_text() + assert ' None: + diff = { + "benchmarks": [ + { + "name": "matmul", + "baseline_median_ms": 1.0, + "current_median_ms": 1.6, + "delta_pct": 60.0, + "status": "regression", + }, + { + "name": "softmax", + "baseline_median_ms": 0.5, + "current_median_ms": 0.51, + "delta_pct": 2.0, + "status": "ok", + }, + {"name": "newkern", "current_median_ms": 0.3, "status": "new"}, + {"name": "removed", "baseline_median_ms": 0.7, "status": "removed"}, + ], + "test_changes": [ + {"name": "test_x", "was": "passed", "now": "failed"}, + ], + } + body = generate_pr_comment(diff) + assert "## gpucheck Benchmark Comparison" in body + assert "matmul" in body + assert "softmax" in body + assert "newkern" in body + assert "removed" in body + assert "test_x" in body + assert ":red_circle:" in body + assert ":green_circle:" in body + assert ":new:" in body + + +def test_generate_pr_comment_empty_diff_returns_friendly_message() -> None: + body = generate_pr_comment({"benchmarks": [], "test_changes": []}) + assert "No changes detected." in body diff --git a/tests/test_reporting_console.py b/tests/test_reporting_console.py new file mode 100644 index 0000000..7768164 --- /dev/null +++ b/tests/test_reporting_console.py @@ -0,0 +1,115 @@ +"""Console reporter tests (Track D — D.1).""" + +from __future__ import annotations + +import io + +from rich.console import Console + +from gpucheck.reporting.console import ( + BenchmarkEntry, + ConsoleReporter, + MemoryEntry, + TestResult, +) + + +def _new_reporter() -> tuple[ConsoleReporter, io.StringIO]: + buf = io.StringIO() + console = Console(file=buf, force_terminal=False, width=120) + return ConsoleReporter(console=console), buf + + +def test_console_reporter_constructs_with_explicit_console() -> None: + reporter, _ = _new_reporter() + assert reporter is not None + + +def test_gpu_info_panel_renders_keys_and_values() -> None: + reporter, buf = _new_reporter() + reporter.gpu_info_panel({"Device": "GTX 1650", "Compute": "7.5"}) + out = buf.getvalue() + assert "Device" in out + assert "GTX 1650" in out + assert "Compute" in out + assert "7.5" in out + + +def test_test_summary_includes_pass_fail_skip_counts() -> None: + reporter, buf = _new_reporter() + results = [ + TestResult(name="t1", status="passed", duration=0.1), + TestResult(name="t2", status="failed", duration=0.2, message="boom"), + TestResult(name="t3", status="skipped", duration=0.0, message="no gpu"), + ] + reporter.test_summary(results) + out = buf.getvalue() + assert "PASSED" in out + assert "FAILED" in out + assert "SKIPPED" in out + assert "1 passed" in out + assert "1 failed" in out + assert "1 skipped" in out + + +def test_benchmark_table_includes_kernel_and_throughput() -> None: + reporter, buf = _new_reporter() + entries = [ + BenchmarkEntry(name="matmul", times=[0.001, 0.0011, 0.0009]), + ] + reporter.benchmark_table(entries) + out = buf.getvalue() + assert "matmul" in out + assert "Median" in out + + +def test_memory_summary_shows_leak_status_red_for_leaks() -> None: + reporter, buf = _new_reporter() + entries = [ + MemoryEntry(name="leaky", peak_mb=10.5, leaked_mb=2.0, allocations=4), + MemoryEntry(name="clean", peak_mb=5.0, leaked_mb=0.0, allocations=2), + ] + reporter.memory_summary(entries) + out = buf.getvalue() + assert "leaky" in out + assert "clean" in out + assert "2.00" in out # 2.0 MB leaked + + +def test_error_detail_renders_name_and_traceback() -> None: + reporter, buf = _new_reporter() + reporter.error_detail("test_x", "AssertionError: nope", traceback="line1\nline2") + out = buf.getvalue() + assert "test_x" in out + assert "AssertionError" in out + + +def test_console_reporter_uses_stderr_in_ci_environment(monkeypatch) -> None: + """When GITHUB_ACTIONS=1, the reporter writes to stderr by default.""" + monkeypatch.setenv("GITHUB_ACTIONS", "1") + monkeypatch.delenv("CI", raising=False) + reporter = ConsoleReporter() + # Internal: assert the file is sys.stderr (Rich's Console exposes file). + import sys + assert reporter._console.file is sys.stderr # noqa: SLF001 + + +def test_console_reporter_with_file_kwarg() -> None: + """Constructing with file= takes precedence over CI/GITHUB_ACTIONS env.""" + buf = io.StringIO() + reporter = ConsoleReporter(file=buf) + reporter.gpu_info_panel({"x": "y"}) + assert "x" in buf.getvalue() + + +def test_benchmark_entry_throughput_handles_zero_times() -> None: + entry = BenchmarkEntry(name="empty", times=[]) + assert entry.median == 0.0 + assert entry.std == 0.0 + assert entry.throughput == 0.0 + + +def test_benchmark_entry_std_with_two_samples() -> None: + entry = BenchmarkEntry(name="k", times=[1.0, 2.0]) + assert entry.median == 1.5 + assert entry.std > 0 # statistics.stdev requires len >= 2 diff --git a/tests/test_reporting_html.py b/tests/test_reporting_html.py new file mode 100644 index 0000000..20cf43e --- /dev/null +++ b/tests/test_reporting_html.py @@ -0,0 +1,157 @@ +"""HTML dashboard tests (Track D — D.2).""" + +from __future__ import annotations + +import json +from html.parser import HTMLParser +from typing import TYPE_CHECKING + +from gpucheck.reporting.html import HTMLReporter + +if TYPE_CHECKING: + from pathlib import Path + + +def _write_sample_json(path: Path) -> None: + payload = { + "schema_version": 1, + "timestamp": "2026-05-01T12:00:00", + "gpu_info": {"name": "GTX 1650", "compute": "7.5"}, + "test_results": [ + {"name": "test_pass", "status": "passed", "duration": 0.1, "message": ""}, + {"name": "test_fail", "status": "failed", "duration": 0.2, + "message": "AssertionError: nope"}, + {"name": "test_skip", "status": "skipped", "duration": 0.0, + "message": "no gpu"}, + ], + "benchmarks": [ + {"name": "matmul", "median_ms": 1.0, "std_ms": 0.05, "samples": 100, "times": []}, + {"name": "softmax", "median_ms": 0.5, "std_ms": 0.02, "samples": 100, "times": []}, + ], + "memory": [ + {"name": "test_pass", "peak_mb": 10.0, "leaked_mb": 0.0, "allocations": 4}, + {"name": "test_fail", "peak_mb": 20.0, "leaked_mb": 2.5, "allocations": 8}, + ], + } + path.write_text(json.dumps(payload), encoding="utf-8") + + +class _TagCounter(HTMLParser): + def __init__(self) -> None: + super().__init__() + self.tags: list[str] = [] + + def handle_starttag(self, tag: str, attrs) -> None: + self.tags.append(tag) + + +def test_html_reporter_writes_self_contained_html(tmp_path: Path) -> None: + json_path = tmp_path / "results.json" + _write_sample_json(json_path) + out = tmp_path / "dashboard.html" + + HTMLReporter(json_path).render(out) + assert out.exists() + text = out.read_text(encoding="utf-8") + assert text.startswith("") + # No external assets (no with http: or https:). + assert 'href="http' not in text + assert 'src="http' not in text + + +def test_html_reporter_contains_summary_counts(tmp_path: Path) -> None: + json_path = tmp_path / "results.json" + _write_sample_json(json_path) + out = tmp_path / "dashboard.html" + HTMLReporter(json_path).render(out) + text = out.read_text() + # 1 passed, 1 failed, 1 skipped — summary cards must surface those. + assert ">1<" in text # one of the cards renders 1 + # Test names appear in the table. + assert "test_pass" in text + assert "test_fail" in text + + +def test_html_reporter_includes_benchmark_table_with_kernel_names(tmp_path: Path) -> None: + json_path = tmp_path / "results.json" + _write_sample_json(json_path) + out = tmp_path / "dashboard.html" + HTMLReporter(json_path).render(out) + text = out.read_text() + assert "matmul" in text + assert "softmax" in text + assert " None: + json_path = tmp_path / "results.json" + _write_sample_json(json_path) + out = tmp_path / "dashboard.html" + diff = { + "benchmarks": [ + { + "name": "matmul", "baseline_median_ms": 1.0, + "current_median_ms": 1.6, "delta_pct": 60.0, + "status": "regression", + }, + ], + } + HTMLReporter(json_path, comparison=diff).render(out) + text = out.read_text() + assert "Comparison vs Baseline" in text + assert "REGRESSION" in text + + +def test_html_reporter_creates_parent_dir(tmp_path: Path) -> None: + json_path = tmp_path / "results.json" + _write_sample_json(json_path) + out = tmp_path / "deep" / "nested" / "dashboard.html" + HTMLReporter(json_path).render(out) + assert out.exists() + + +def test_html_reporter_html_is_well_formed(tmp_path: Path) -> None: + """HTMLParser tolerates malformed HTML, but should at least parse and + open balanced top-level tags (html, body).""" + json_path = tmp_path / "results.json" + _write_sample_json(json_path) + out = tmp_path / "dashboard.html" + HTMLReporter(json_path).render(out) + text = out.read_text() + counter = _TagCounter() + counter.feed(text) + assert "html" in counter.tags + assert "body" in counter.tags + assert "table" in counter.tags + + +def test_html_reporter_handles_empty_data(tmp_path: Path) -> None: + """No tests / no benchmarks must not crash the renderer.""" + json_path = tmp_path / "empty.json" + json_path.write_text(json.dumps({ + "schema_version": 1, "timestamp": "", "gpu_info": {}, + "test_results": [], "benchmarks": [], "memory": [], + })) + out = tmp_path / "dash.html" + HTMLReporter(json_path).render(out) + text = out.read_text() + assert "Summary" in text + # Test Results / Benchmarks / Memory sections are skipped when empty. + assert "Test Results" not in text + assert "Benchmarks" not in text + + +def test_html_reporter_escapes_html_in_messages(tmp_path: Path) -> None: + json_path = tmp_path / "results.json" + json_path.write_text(json.dumps({ + "test_results": [ + {"name": "test_x", "status": "failed", "duration": 0.1, + "message": ""}, + ], + "benchmarks": [], "memory": [], "gpu_info": {}, "timestamp": "", + })) + out = tmp_path / "dash.html" + HTMLReporter(json_path).render(out) + text = out.read_text() + assert "