rpx_benchmark.examples.profile_callable

  1"""Measure a callable on fixed NPZ inputs with warmup and CUDA synchronization.
  2
  3The callable is imported as module:function and receives the NPZ arrays as
  4keyword arguments. It must include the preprocessing/postprocessing intended
  5for measurement. Inputs are loaded once, outside timing. This measures repeated
  6calls on fixed inputs, not accuracy, distinct dataset samples, or throughput.
  7"""
  8
  9from __future__ import annotations
 10
 11import argparse
 12import hashlib
 13import importlib
 14import json
 15import time
 16from pathlib import Path
 17from typing import Any, Callable
 18
 19import numpy as np
 20
 21from ..exceptions import ConfigError
 22from ..profiler import LatencyProfiler, MemoryProfiler, SystemCard
 23
 24
 25def load_callable(name: str) -> Callable[..., Any]:
 26    """Load a model callable without imposing a framework or API client."""
 27    module, separator, attribute = name.partition(":")
 28    if not separator or not module or not attribute:
 29        raise ConfigError("Model must use module:function notation")
 30    fn = getattr(importlib.import_module(module), attribute)
 31    if not callable(fn):
 32        raise ConfigError(f"{name} is not callable")
 33    return fn
 34
 35
 36def measure(
 37    fn: Callable[..., Any],
 38    inputs: dict[str, np.ndarray],
 39    *,
 40    device: str = "cpu",
 41    warmup: int = 10,
 42    repeats: int = 1000,
 43    precision: str = "unspecified",
 44    unit: str = "request",
 45) -> dict[str, Any]:
 46    """Return measured percentiles, raw timings, memory and a system card.
 47
 48    CUDA measurements synchronize all visible devices before and after every
 49    call. CUDA must be available when requested. CPU RSS is a process-lifetime
 50    high-water mark; CUDA values are allocator peaks per visible device.
 51    Precision is a recorded description, never an automatic model conversion.
 52    """
 53    if device not in {"cpu", "cuda"} or warmup < 0 or repeats < 1:
 54        raise ConfigError("device must be cpu/cuda, warmup >= 0, repeats >= 1")
 55    if not inputs or not precision.strip() or not unit.strip():
 56        raise ConfigError("Provide inputs, precision and a measurement unit")
 57    torch: Any = None
 58    cuda_devices: list[int] = []
 59    try:
 60        torch = importlib.import_module("torch")
 61    except ImportError:
 62        if device == "cuda":
 63            raise ConfigError("CUDA profiling requires PyTorch and an available CUDA GPU") from None
 64    if device == "cuda":
 65        if not torch.cuda.is_available():
 66            raise ConfigError("CUDA requested but no CUDA GPU is available")
 67        cuda_devices = list(range(torch.cuda.device_count()))
 68
 69    def sync() -> None:
 70        for index in cuda_devices:
 71            torch.cuda.synchronize(index)
 72
 73    for _ in range(warmup):
 74        sync()
 75        fn(**inputs)
 76        sync()
 77    memory = MemoryProfiler(sample_cuda=False, sample_mps=False).reset()
 78    for index in cuda_devices:
 79        torch.cuda.reset_peak_memory_stats(index)
 80    latency = LatencyProfiler(warmup=0)
 81    for _ in range(repeats):
 82        sync()
 83        started = time.perf_counter()
 84        prediction = fn(**inputs)
 85        sync()
 86        latency.add_sample_seconds(time.perf_counter() - started)
 87        del prediction
 88    peaks = memory.peaks()
 89    cuda_peaks = {
 90        str(index): torch.cuda.max_memory_allocated(index) / (1024**2) for index in cuda_devices
 91    }
 92    card = SystemCard.auto_detect(precision=precision)
 93    # CPU runs should not present a visible but unused GPU as measured hardware.
 94    if device == "cpu":
 95        card.gpu_name, card.gpu_memory_gb, card.gpu_count = "", 0.0, 0
 96    card.input_resolution = "caller_defined; see inputs"
 97    return {
 98        "timing_scope": "synchronized_callable_wall_time; preloaded_inputs",
 99        "device": device,
100        "unit": unit,
101        "warmup_calls": warmup,
102        "measured_calls": repeats,
103        "inputs": {
104            key: {"shape": list(value.shape), "dtype": str(value.dtype)}
105            for key, value in inputs.items()
106        },
107        "latency": latency.percentiles(),
108        "latency_samples_ms": list(latency.samples_ms()),
109        "memory": {
110            "cpu_process_lifetime_peak_mb": peaks["cpu_mb"],
111            "cuda_peak_allocated_mb_by_visible_device": cuda_peaks,
112        },
113        "system_card": card.to_dict(),
114        "notes": [
115            "Warmup excluded. Fixed-input repeats are not distinct samples.",
116            "Precision and callable device placement are caller-controlled.",
117            "API callables include transport time and exclude remote GPU accounting.",
118        ],
119    }
120
121
122def main() -> None:
123    parser = argparse.ArgumentParser(description=__doc__)
124    parser.add_argument("--model", required=True, help="module:function")
125    parser.add_argument("--inputs", type=Path, required=True)
126    parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
127    parser.add_argument("--precision", default="unspecified")
128    parser.add_argument("--unit", default="request", help="frame, clip, pair or request")
129    parser.add_argument("--warmup", type=int, default=10)
130    parser.add_argument("--repeats", type=int, default=1000)
131    parser.add_argument("--output", type=Path, required=True)
132    args = parser.parse_args()
133    with np.load(args.inputs, allow_pickle=False) as payload:
134        inputs = {key: payload[key] for key in payload.files}
135    result = measure(
136        load_callable(args.model),
137        inputs,
138        device=args.device,
139        warmup=args.warmup,
140        repeats=args.repeats,
141        precision=args.precision,
142        unit=args.unit,
143    )
144    result["model_callable"] = args.model
145    result["inputs_sha256"] = hashlib.sha256(args.inputs.read_bytes()).hexdigest()
146    args.output.parent.mkdir(parents=True, exist_ok=True)
147    args.output.write_text(json.dumps(result, indent=2, allow_nan=False) + "\n")
148    print(args.output)
149
150
151if __name__ == "__main__":
152    main()