Add GPU telemetry with fallback and test coverage

patchwork/telemetry/profiler.py

@@ -0,0 +1,154 @@
+"""
+patchwork.telemetry.profiler
+==============================
+VRAM and wall-clock telemetry, kept fully decoupled from graph.py and
+the tools -- nothing in the agent logic needs to know this exists.
+Used by benchmarks/evaluate.py to wrap graph.invoke() calls.
+
+Must never crash when no NVIDIA GPU is present (CI runs on ubuntu-latest
+with no GPU) -- gpu_available=False + peak_vram_mb=None is the correct
+degraded result, not an exception.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+import time
+from collections.abc import Callable
+from functools import wraps
+from typing import Final, ParamSpec, TypeVar
+
+import pynvml  # type: ignore[import-untyped]  # nvidia-ml-py ships no py.typed marker
+from pydantic import BaseModel, Field
+
+logger = logging.getLogger("patchwork.telemetry.profiler")
+
+DEFAULT_POLL_INTERVAL_SEC: Final[float] = 0.05
+DEFAULT_DEVICE_INDEX: Final[int] = 0
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+class TelemetryResult(BaseModel):
+    duration_sec: float = Field(ge=0.0)
+    peak_vram_mb: float | None = None  # None if no GPU / nvml unavailable
+    gpu_available: bool = False
+    error_message: str | None = None  # only for nvml-level failures, not "no GPU"
+
+
+class GPUProfiler:
+    """Context manager. Polls VRAM usage on a background thread while the
+    `with` block runs, tracks the peak. Falls back cleanly to
+    gpu_available=False if no NVIDIA GPU/driver is present -- this is
+    the expected path in CI, not an error condition.
+    """
+
+    def __init__(
+        self,
+        device_index: int = DEFAULT_DEVICE_INDEX,
+        interval: float = DEFAULT_POLL_INTERVAL_SEC,
+    ) -> None:
+        self.device_index = device_index
+        self.interval = interval
+        self.peak_vram_mb: float | None = None
+        self.gpu_available = False
+        self.error_message: str | None = None
+        self.duration_sec: float = 0.0
+
+        self._stop_event = threading.Event()
+        self._thread: threading.Thread | None = None
+        self._start_time: float = 0.0
+
+    def _monitor(self) -> None:
+        try:
+            pynvml.nvmlInit()
+        except pynvml.NVMLError as exc:
+            # no driver / no GPU -- expected on CI, not a crash
+            self.error_message = str(exc)
+            return
+
+        try:
+            handle = pynvml.nvmlDeviceGetHandleByIndex(self.device_index)
+        except pynvml.NVMLError as exc:
+            self.error_message = str(exc)
+            pynvml.nvmlShutdown()
+            return
+
+        self.gpu_available = True
+        self.peak_vram_mb = 0.0
+
+        while not self._stop_event.is_set():
+            try:
+                mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+            except pynvml.NVMLError as exc:
+                # device disappeared mid-run (driver reset, etc) -- stop
+                # polling rather than crash the profiling thread
+                logger.warning(
+                    "gpu_poll_failed",
+                    extra={"event": "gpu_poll_failed", "error": str(exc)},
+                )
+                break
+            used_mb = mem_info.used / (1024 * 1024)
+            self.peak_vram_mb = max(self.peak_vram_mb, used_mb)
+            self._stop_event.wait(self.interval)
+
+        pynvml.nvmlShutdown()
+
+    def __enter__(self) -> GPUProfiler:  # noqa: PYI034  # typing.Self needs 3.11+, project supports 3.10
+        self._stop_event.clear()
+        self._thread = threading.Thread(target=self._monitor, daemon=True)
+        self._thread.start()
+        self._start_time = time.perf_counter()
+        return self
+
+    def __exit__(self, *exc_info: object) -> None:
+        self.duration_sec = time.perf_counter() - self._start_time
+        self._stop_event.set()
+        if self._thread is not None:
+            self._thread.join(timeout=self.interval * 5)
+
+    def result(self) -> TelemetryResult:
+        return TelemetryResult(
+            duration_sec=self.duration_sec,
+            peak_vram_mb=self.peak_vram_mb,
+            gpu_available=self.gpu_available,
+            error_message=self.error_message,
+        )
+
+
+def profile_call(
+    func: Callable[..., T], *args: object, **kwargs: object
+) -> tuple[T, TelemetryResult]:
+    """Runs func(*args, **kwargs) under a GPUProfiler, returns (result, telemetry)."""
+    with GPUProfiler() as profiler:
+        result = func(*args, **kwargs)
+    return result, profiler.result()
+
+
+def profile_vram(func: Callable[P, T]) -> Callable[P, T]:
+    """Transparent decorator: wraps func in a GPUProfiler, logs the
+    telemetry, and returns func's original result unchanged. For call
+    sites that just want VRAM logging without touching their return
+    signature (business logic stays untouched, per the isolation rule).
+    """
+
+    @wraps(func)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        with GPUProfiler() as profiler:
+            result = func(*args, **kwargs)
+        telemetry = profiler.result()
+        logger.info(
+            "profiled_call",
+            extra={
+                "event": "profiled_call",
+                "function": func.__name__,
+                "duration_sec": round(telemetry.duration_sec, 3),
+                "peak_vram_mb": telemetry.peak_vram_mb,
+                "gpu_available": telemetry.gpu_available,
+            },
+        )
+        return result
+
+    return wrapper

scripts/manual_profiler_test.py

@@ -0,0 +1,66 @@
+"""
+scripts/manual_profiler_test.py
+==================================
+tests/test_profiler.py only proves the no-GPU fallback path (that's all
+CI/this sandbox has). This script is the only way to confirm the actual
+positive path -- gpu_available=True and a plausible VRAM number -- since
+that requires a real NVIDIA GPU running real inference.
+
+Usage:
+    python scripts/manual_profiler_test.py
+"""
+
+from __future__ import annotations
+
+from patchwork.graph import build_patchwork_graph, build_structured_llm
+from patchwork.state import create_initial_state
+from patchwork.telemetry.profiler import GPUProfiler, profile_call
+
+BUGGY_SOURCE = "def add(a, b):\n    return a - b\n"
+
+
+def main() -> None:
+    print("Building structured LLM client and graph...")
+    structured_llm = build_structured_llm()
+    graph = build_patchwork_graph(structured_llm)
+    initial = create_initial_state("manual_profiler_target.py", BUGGY_SOURCE)
+
+    print("Running one audit pass wrapped in GPUProfiler...\n")
+    result, telemetry = profile_call(graph.invoke, initial)
+
+    print(f"gpu_available: {telemetry.gpu_available}")
+    print(f"peak_vram_mb:  {telemetry.peak_vram_mb}")
+    print(f"duration_sec:  {round(telemetry.duration_sec, 2)}")
+    print(f"error_message: {telemetry.error_message}")
+
+    if not telemetry.gpu_available:
+        print("\nWARNING: gpu_available is False. Either nvidia-ml-py can't see your")
+        print("GPU, or Ollama is running on CPU. Check `nvidia-smi` and `ollama ps`")
+        print(
+            "(the PROCESSOR column should say 100% GPU) before trusting benchmark numbers."
+        )
+    else:
+        target_max = 3400  # <=3.4GB Q8_0 ceiling from project spec
+        if telemetry.peak_vram_mb is not None and telemetry.peak_vram_mb > target_max:
+            print(
+                f"\nNOTE: peak VRAM ({telemetry.peak_vram_mb:.0f}MB) exceeds your project's"
+            )
+            print(
+                f"target ceiling of {target_max}MB -- worth investigating before benchmarking."
+            )
+
+    print(
+        f"\nSandbox passed: {result['sandbox_result'].passed if result['sandbox_result'] else 'N/A'}"
+    )
+
+    # Second run: raw context manager, for comparing against profile_call's
+    # numbers -- they should be close but not identical (profiling overhead).
+    print("\n--- Second run via raw GPUProfiler context manager ---")
+    with GPUProfiler(interval=0.02) as profiler:
+        graph.invoke(create_initial_state("manual_profiler_target_2.py", BUGGY_SOURCE))
+    r2 = profiler.result()
+    print(f"peak_vram_mb: {r2.peak_vram_mb}, duration_sec: {round(r2.duration_sec, 2)}")
+
+
+if __name__ == "__main__":
+    main()

tests/test_profiler.py

@@ -0,0 +1,100 @@
+"""
+tests/test_profiler.py
+========================
+This sandbox and CI both have no NVIDIA GPU, so every test here runs
+against the real no-GPU fallback path -- which is exactly the path that
+matters most to get right, since it's the one CI will always exercise.
+"""
+
+from __future__ import annotations
+
+import time
+
+from patchwork.telemetry.profiler import (
+    GPUProfiler,
+    TelemetryResult,
+    profile_call,
+    profile_vram,
+)
+
+
+class TestGPUProfilerNoGPU:
+    def test_context_manager_does_not_raise_without_gpu(self) -> None:
+        with GPUProfiler() as profiler:
+            time.sleep(0.05)
+        result = profiler.result()
+        assert isinstance(result, TelemetryResult)
+
+    def test_gpu_available_false_without_driver(self) -> None:
+        with GPUProfiler() as profiler:
+            pass
+        assert profiler.result().gpu_available is False
+
+    def test_peak_vram_none_without_gpu(self) -> None:
+        with GPUProfiler() as profiler:
+            pass
+        assert profiler.result().peak_vram_mb is None
+
+    def test_duration_still_measured_without_gpu(self) -> None:
+        with GPUProfiler(interval=0.01) as profiler:
+            time.sleep(0.1)
+        result = profiler.result()
+        assert result.duration_sec >= 0.1
+
+    def test_error_message_populated_when_nvml_unavailable(self) -> None:
+        with GPUProfiler() as profiler:
+            pass
+        assert profiler.result().error_message is not None
+
+    def test_result_is_pydantic_model_and_json_serializable(self) -> None:
+        with GPUProfiler() as profiler:
+            pass
+        payload = profiler.result().model_dump_json()
+        assert isinstance(payload, str)
+        assert '"gpu_available":false' in payload.replace(" ", "")
+
+
+class TestProfileCall:
+    def test_returns_function_result_unchanged(self) -> None:
+        def add(a: int, b: int) -> int:
+            return a + b
+
+        result, telemetry = profile_call(add, 2, 3)
+        assert result == 5
+        assert isinstance(telemetry, TelemetryResult)
+
+    def test_supports_kwargs(self) -> None:
+        def greet(name: str, greeting: str = "hello") -> str:
+            return f"{greeting}, {name}"
+
+        result, _telemetry = profile_call(greet, name="world", greeting="hi")
+        assert result == "hi, world"
+
+    def test_exception_in_wrapped_function_propagates(self) -> None:
+        def boom() -> None:
+            raise ValueError("intentional")
+
+        try:
+            profile_call(boom)
+            raised = False
+        except ValueError:
+            raised = True
+        assert raised  # profiling must not swallow the caller's own errors
+
+
+class TestProfileVramDecorator:
+    def test_decorated_function_returns_original_result(self) -> None:
+        @profile_vram
+        def multiply(a: int, b: int) -> int:
+            return a * b
+
+        assert multiply(4, 5) == 20
+
+    def test_decorator_preserves_function_name(self) -> None:
+        # without functools.wraps, this would be "wrapper" instead --
+        # matters for debugging/logging where you inspect __name__
+        @profile_vram
+        def my_function() -> int:
+            return 1
+
+        assert my_function.__name__ == "my_function"