diff --git a/patchwork/telemetry/__init__.py b/patchwork/telemetry/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/patchwork/telemetry/profiler.py b/patchwork/telemetry/profiler.py new file mode 100644 index 0000000..0f815b2 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/scripts/manual_profiler_test.py b/scripts/manual_profiler_test.py new file mode 100644 index 0000000..ddd7b68 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/tests/test_profiler.py b/tests/test_profiler.py new file mode 100644 index 0000000..8dd5c05 --- /dev/null +++ b/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" \ No newline at end of file