Doc-to-LoRA release

src/ctx_to_lora/tracker/tracker.py

@@ -0,0 +1,303 @@
+"""Unified tracking interface combining timing and CUDA memory usage.
+
+Primary API
+-----------
+add_tracker(bound_method, name)
+    Wraps a bound instance method so that each invocation records:
+      - wall-clock duration (seconds) in timer.TIMER_REGISTRY[name]
+      - CUDA peak memory increase (bytes) in cuda_memory_tracker.MEMORY_REGISTRY[name]
+        (only if CUDA + torch available; otherwise memory list may stay empty / absent)
+
+print_tracker_stats(name=None)
+    Convenience printer that delegates to time + memory aggregate printers.
+
+Design
+------
+We implement a single wrapper (instead of nesting the individual timer & memory
+wrappers) to avoid multiple layers of indirection and to ensure the measured
+CUDA memory footprint reflects only the original method's body (excluding the
+separate timing wrapper's slight overhead). The wrapper is idempotent: repeated
+calls to add_tracker on the same method are ignored.
+
+This file depends on sibling modules:
+- tracker.timer
+- tracker.cuda_memory_tracker
+
+Both registries remain the single source of truth; no additional registry is introduced.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from time import perf_counter
+from typing import Any
+
+# Support both package (relative) and direct script execution.
+try:  # Package / normal import path
+    from .cuda_memory_tracker import (  # type: ignore
+        MEMORY_REGISTRY,
+        compute_aggregate_memory_stats,
+        print_aggregate_memory_stats,
+        print_global_memory_stats,
+        reset_memory_trackers,
+        save_memory_stats_csv,
+    )
+    from .timer import (  # type: ignore
+        TIMER_REGISTRY,
+        compute_aggregate_timer_stats,
+        print_aggregate_timer_stats,
+        print_global_timer_stats,
+        reset_timers,
+        save_timer_stats_csv,
+    )
+except Exception:  # pragma: no cover - fallback when executed directly
+    import pathlib
+    import sys
+
+    _this_file = pathlib.Path(__file__).resolve()
+    # project root is two levels up from tracker/ (i.e., .../src)
+    _src_root = _this_file.parents[2]
+    if str(_src_root) not in sys.path:
+        sys.path.insert(0, str(_src_root))
+    try:
+        from ctx_to_lora.tracker.cuda_memory_tracker import (  # type: ignore
+            MEMORY_REGISTRY,
+            compute_aggregate_memory_stats,
+            print_aggregate_memory_stats,
+            print_global_memory_stats,
+            reset_memory_trackers,
+            save_memory_stats_csv,
+        )
+        from ctx_to_lora.tracker.timer import (  # type: ignore
+            TIMER_REGISTRY,
+            compute_aggregate_timer_stats,
+            print_aggregate_timer_stats,
+            print_global_timer_stats,
+            reset_timers,
+            save_timer_stats_csv,
+        )
+    except Exception as e:  # If still failing, raise a clearer error.
+        raise ImportError(
+            f"Failed to import tracking dependencies; ensure project root on PYTHONPATH. Original: {e}"
+        )
+
+try:  # Optional torch import (lazy fallback if unavailable)
+    import torch  # type: ignore
+except Exception:  # pragma: no cover - torch absence path
+    torch = None  # type: ignore
+
+__all__ = [
+    "add_tracker",
+    "compute_tracker_stats",
+    "save_tracker_stats_csv",
+    "print_tracker_stats",
+    "print_global_tracker_stats",
+    "reset_trackers",
+]
+
+
+def _cuda_available() -> bool:
+    return bool(torch is not None and torch.cuda.is_available())
+
+
+def add_tracker(func: Callable, name: str) -> None:
+    """Attach a combined time + CUDA memory tracking wrapper to a bound method.
+
+    Parameters
+    ----------
+    func : Callable
+        A bound instance method (instance.method). Raises ValueError if unbound.
+    name : str
+        Registry key under which metrics are stored.
+    """
+    if not hasattr(func, "__self__") or getattr(func, "__self__") is None:
+        # Permit idempotent re-calls if already wrapped.
+        if getattr(func, "__is_tracker_wrapper__", False):
+            return
+        raise ValueError(
+            "add_tracker expects a bound method: call with instance.method"
+        )
+
+    instance = func.__self__  # underlying object
+    method_name = getattr(func, "__name__", None)
+    if method_name is None:
+        raise ValueError("Cannot determine method name for provided callable")
+
+    existing = getattr(instance, method_name, None)
+    if getattr(
+        existing, "__is_tracker_wrapper__", False
+    ):  # Already wrapped via unified tracker
+        return
+
+    # If already individually wrapped by timer or memory tracker, we still wrap only once more;
+    # future calls to add_tracker will become no-ops.
+    orig_bound = existing if existing is not None else func
+
+    def tracked(*args: Any, **kwargs: Any):  # noqa: D401 - combined wrapper
+        use_cuda = _cuda_available()
+        if use_cuda:
+            torch.cuda.synchronize()
+            torch.cuda.reset_peak_memory_stats()
+            start_alloc = torch.cuda.memory_allocated()
+        start_time = perf_counter()
+        try:
+            return orig_bound(*args, **kwargs)
+        finally:
+            elapsed = perf_counter() - start_time
+            TIMER_REGISTRY.setdefault(name, []).append(elapsed)
+            if use_cuda:
+                torch.cuda.synchronize()
+                peak_alloc = torch.cuda.max_memory_allocated()
+                peak_increase = peak_alloc - start_alloc
+                if peak_increase < 0:  # Safety guard (should not happen)
+                    peak_increase = 0
+                MEMORY_REGISTRY.setdefault(name, []).append(int(peak_increase))
+
+    # Introspection / idempotency markers
+    tracked.__name__ = method_name
+    tracked.__qualname__ = getattr(orig_bound, "__qualname__", method_name)
+    tracked.__doc__ = getattr(orig_bound, "__doc__")
+    tracked.__wrapped__ = orig_bound  # type: ignore[attr-defined]
+    tracked.__is_tracker_wrapper__ = True  # type: ignore[attr-defined]
+    tracked.__is_timer_wrapper__ = True  # type: ignore[attr-defined]
+    tracked.__is_memory_wrapper__ = True  # type: ignore[attr-defined]
+    tracked.__tracker_name__ = name  # type: ignore[attr-defined]
+
+    setattr(instance, method_name, tracked)
+
+
+def compute_tracker_stats(
+    name: str | None = None,
+) -> dict[str, dict[str, Any]] | None:
+    """Compute both timing and memory stats for a given name (or all if None).
+
+    Parameters
+    ----------
+    name : Optional[str]
+        Specific tracker name; if None, computes all.
+
+    Returns
+    -------
+    Optional[Dict[str, Dict[str, Any]]]
+        None if no data, else dict with 'timing' and 'memory' keys containing
+        their respective aggregate statistics.
+    """
+    timer_stats = compute_aggregate_timer_stats(name)
+    memory_stats = compute_aggregate_memory_stats(name)
+
+    if timer_stats is None and memory_stats is None:
+        return None
+
+    return {
+        "timing": timer_stats or {},
+        "memory": memory_stats or {},
+    }
+
+
+def save_tracker_stats_csv(file_path: str, name: str | None = None) -> None:
+    """Save both timing and memory stats to separate CSV files.
+
+    Parameters
+    ----------
+    file_path : str
+        Base path for CSV files. Will create file_path_timing.csv and file_path_memory.csv
+    name : Optional[str]
+        Specific tracker name to export. If None, all trackers are exported.
+    """
+    import os
+
+    os.makedirs(os.path.dirname(file_path), exist_ok=True)
+
+    base_path = os.path.splitext(file_path)[0]
+    timer_path = f"{base_path}_timing.csv"
+    memory_path = f"{base_path}_memory.csv"
+
+    # Save timing stats if available
+    timer_stats = compute_aggregate_timer_stats(name)
+    if timer_stats is not None:
+        save_timer_stats_csv(timer_path, name)
+
+    # Save memory stats if available
+    memory_stats = compute_aggregate_memory_stats(name)
+    if memory_stats is not None:
+        save_memory_stats_csv(memory_path, name)
+
+    # If no data at all, raise an error
+    if timer_stats is None and memory_stats is None:
+        print("No tracking data available to export")
+
+
+def print_tracker_stats(name: str | None = None) -> None:
+    """Print both timing and memory stats for a given name (or all if None).
+
+    Parameters
+    ----------
+    name : Optional[str]
+        Specific tracker name; if None, prints all.
+    """
+    print("[tracker] Timing stats:")
+    print_aggregate_timer_stats(name)
+    print("\n[tracker] CUDA memory stats:")
+    print_aggregate_memory_stats(name)
+
+
+def print_global_tracker_stats() -> None:
+    """Print global aggregate timing and memory stats."""
+    print("[tracker] Global timing stats:")
+    print_global_timer_stats()
+    print("\n[tracker] Global CUDA memory stats:")
+    print_global_memory_stats()
+
+
+def reset_trackers() -> None:
+    """Reset all timer and memory tracking data."""
+    reset_timers()
+    reset_memory_trackers()
+
+
+if __name__ == "__main__":  # Demonstration
+    import random
+    import time
+
+    class Demo:
+        def compute(self, n: int = 25_000) -> int:
+            # CPU-bound work
+            s = 0
+            for i in range(n):
+                s += i * i
+            return s
+
+        def gpu_alloc(self, n: int = 500_000):
+            if not _cuda_available():
+                # Simulate light wait to differentiate timing
+                time.sleep(random.uniform(0.01, 0.05))
+                return None
+            t = torch.empty(n, dtype=torch.float32, device="cuda")
+            t.uniform_()  # ensure usage
+            return t.sum().item()
+
+    demo = Demo()
+
+    add_tracker(demo.compute, "compute")
+    add_tracker(demo.gpu_alloc, "gpu_alloc")
+    # Idempotent re-call
+    add_tracker(demo.compute, "compute")
+
+    for _ in range(5):
+        demo.compute(15_000)
+        demo.gpu_alloc(300_000)
+
+    print_tracker_stats()
+    print("\n--- Global Combined Stats ---\n")
+    print_global_tracker_stats()
+
+    # Demonstrate CSV export
+    print("\n[tracker] Saving stats to CSV files...\n")
+    csv_path = "/tmp/tracker_demo_stats.csv"
+    save_tracker_stats_csv(csv_path)
+    print(f"Exported timing stats to: {csv_path.replace('.csv', '_timing.csv')}")
+    print(f"Exported memory stats to: {csv_path.replace('.csv', '_memory.csv')}")
+
+    print("\n[tracker] Resetting registries...\n")
+    reset_trackers()
+    print_tracker_stats()