| """Sampler-only generation latency measurement and atomic publication.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import asdict, dataclass |
| import hashlib |
| import inspect |
| import json |
| import math |
| import os |
| from pathlib import Path |
| import re |
| import statistics |
| import subprocess |
| import tempfile |
| import time |
| from typing import Callable, Mapping |
|
|
|
|
| EXCLUSIONS = ( |
| "model_and_checkpoint_loading", |
| "first_compilation", |
| "token_decoding", |
| "metrics", |
| "file_io", |
| ) |
| REQUIRED_METADATA = frozenset( |
| { |
| "seed", |
| "dataset", |
| "model", |
| "steps", |
| "gpu_name", |
| "gpu_index", |
| "gpu_compute_capability", |
| "cuda_runtime_version", |
| "pytorch_compiled_cuda_toolkit", |
| "nvidia_driver_version", |
| "precision", |
| "requested_precision", |
| "parameter_precision", |
| "precision_policy", |
| "precision_evidence", |
| "precision_policy_binding", |
| "runtime_code_binding", |
| "synchronization_policy", |
| "environment", |
| "source_commit", |
| "config_sha256", |
| "checkpoint_sha256", |
| "checkpoint_lock_id", |
| "checkpoint_selection", |
| "checkpoint_teacher_family", |
| "adapter_identity", |
| "attempt_id", |
| } |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class TimingResult: |
| mode: str |
| warmups: int |
| repeats: int |
| batch_size: int |
| num_timed_samples: int |
| seconds_per_sample: float | None |
| median_seconds: float |
| standard_deviation_seconds: float |
| standard_deviation_convention: str |
| raw_durations_seconds: tuple[float, ...] |
| exclusions: tuple[str, ...] = EXCLUSIONS |
|
|
|
|
| def _integer(name: str, value: object, *, positive: bool = False) -> int: |
| if type(value) is not int or value < (1 if positive else 0): |
| qualifier = "a positive integer" if positive else "a nonnegative integer" |
| raise ValueError(f"{name} must be {qualifier}") |
| return value |
|
|
|
|
| def benchmark( |
| generate_one: Callable[[], object], |
| synchronize: Callable[[], object], |
| warmups: int = 5, |
| repeats: int = 32, |
| *, |
| clock: Callable[[], float] = time.perf_counter, |
| batch_size: int = 1, |
| mode: str = "primary_latency", |
| ) -> TimingResult: |
| """Warm a loaded sampler, then time only synchronized in-memory calls.""" |
|
|
| warmups = _integer("warmups", warmups) |
| repeats = _integer("repeats", repeats, positive=True) |
| batch_size = _integer("batch_size", batch_size, positive=True) |
| if mode not in {"primary_latency", "throughput"}: |
| raise ValueError("mode must be primary_latency or throughput") |
| if mode == "primary_latency" and batch_size != 1: |
| raise ValueError("primary_latency requires batch size 1") |
| if not callable(generate_one) or not callable(synchronize) or not callable(clock): |
| raise ValueError("generate_one, synchronize, and clock must be callable") |
|
|
| for _ in range(warmups): |
| generate_one() |
|
|
| durations: list[float] = [] |
| for _ in range(repeats): |
| synchronize() |
| started = clock() |
| generate_one() |
| synchronize() |
| stopped = clock() |
| if not all(isinstance(value, (int, float)) and math.isfinite(value) for value in (started, stopped)): |
| raise ValueError("clock samples must be finite numbers") |
| duration = float(stopped - started) |
| if duration < 0 or not math.isfinite(duration): |
| raise ValueError("clock duration must be finite and nonnegative") |
| durations.append(duration) |
|
|
| mean = statistics.fmean(durations) |
| return TimingResult( |
| mode=mode, |
| warmups=warmups, |
| repeats=repeats, |
| batch_size=batch_size, |
| num_timed_samples=repeats * batch_size, |
| seconds_per_sample=mean if mode == "primary_latency" else None, |
| median_seconds=statistics.median(durations), |
| standard_deviation_seconds=statistics.pstdev(durations), |
| standard_deviation_convention="population", |
| raw_durations_seconds=tuple(durations), |
| ) |
|
|
|
|
| def _validate_metadata(metadata: Mapping[str, object]) -> dict[str, object]: |
| missing = sorted(REQUIRED_METADATA - metadata.keys()) |
| if missing: |
| raise ValueError("timing metadata is missing: " + ", ".join(missing)) |
| value = dict(metadata) |
| if type(value["seed"]) is not int or type(value["steps"]) is not int or value["steps"] <= 0: |
| raise ValueError("timing metadata seed/steps are invalid") |
| structured = { |
| "seed", |
| "steps", |
| "gpu_index", |
| "checkpoint_selection", |
| "precision_policy_binding", |
| "runtime_code_binding", |
| } |
| for key in REQUIRED_METADATA - structured: |
| if not isinstance(value[key], str) or not value[key]: |
| raise ValueError(f"timing metadata {key} must be a non-empty string") |
| if type(value["gpu_index"]) is not int or value["gpu_index"] < 0: |
| raise ValueError("timing metadata gpu_index must be a nonnegative integer") |
| if not isinstance(value["checkpoint_selection"], dict) or not value["checkpoint_selection"]: |
| raise ValueError("timing metadata checkpoint_selection must be a non-empty mapping") |
| binding = value["precision_policy_binding"] |
| if not isinstance(binding, dict) or not binding: |
| raise ValueError("timing metadata precision_policy_binding must be a non-empty mapping") |
| code = value["runtime_code_binding"] |
| if not isinstance(code, list) or not code: |
| raise ValueError("timing metadata runtime_code_binding must be a non-empty list") |
| for item in code: |
| if not isinstance(item, dict) or not isinstance(item.get("path"), str): |
| raise ValueError("timing metadata runtime_code_binding entry is invalid") |
| digest = item.get("sha256") |
| if not isinstance(digest, str) or re.fullmatch(r"[0-9a-f]{64}", digest) is None: |
| raise ValueError("timing metadata runtime code SHA-256 is invalid") |
| for key, length in (("source_commit", 40), ("config_sha256", 64), ("checkpoint_sha256", 64)): |
| digest = value[key] |
| if len(digest) != length or any(character not in "0123456789abcdef" for character in digest): |
| raise ValueError(f"timing metadata {key} is not canonical hexadecimal") |
| attempt_id = value["attempt_id"] |
| if len(attempt_id) != 32 or any(character not in "0123456789abcdef" for character in attempt_id): |
| raise ValueError("timing metadata attempt_id is not canonical hexadecimal") |
| return value |
|
|
|
|
| def publish_timing(path: Path, result: TimingResult, metadata: Mapping[str, object]) -> None: |
| """Validate everything before atomically replacing a completed timing result.""" |
|
|
| metadata_document = _validate_metadata(metadata) |
| output = Path(path).absolute() |
| if output.is_symlink() or output.parent.is_symlink(): |
| raise ValueError(f"timing output path is unsafe: {output}") |
| output.parent.mkdir(parents=True, exist_ok=True) |
| payload = { |
| "schema": "dlb-generation-timing-v1", |
| "timing": asdict(result), |
| "metadata": metadata_document, |
| } |
| descriptor, temporary_name = tempfile.mkstemp( |
| prefix=f".{output.name}.", suffix=".partial", dir=output.parent |
| ) |
| temporary = Path(temporary_name) |
| try: |
| with os.fdopen(descriptor, "w", encoding="utf-8") as handle: |
| json.dump(payload, handle, sort_keys=True, separators=(",", ":")) |
| handle.write("\n") |
| handle.flush() |
| os.fsync(handle.fileno()) |
| os.replace(temporary, output) |
| except BaseException: |
| temporary.unlink(missing_ok=True) |
| raise |
|
|
|
|
| def load_metadata(path: Path) -> dict[str, object]: |
| path = Path(path) |
| if not path.is_absolute() or path.is_symlink() or not path.is_file(): |
| raise ValueError(f"benchmark metadata is missing or unsafe: {path}") |
| try: |
| value = json.loads(path.read_text(encoding="utf-8")) |
| except (OSError, UnicodeDecodeError, json.JSONDecodeError) as error: |
| raise ValueError("benchmark metadata is invalid") from error |
| if not isinstance(value, dict): |
| raise ValueError("benchmark metadata must be a mapping") |
| return value |
|
|
|
|
| def _driver_version() -> str: |
| """Return the NVIDIA package driver version, never a CUDA API integer.""" |
|
|
| try: |
| completed = subprocess.run( |
| [ |
| "nvidia-smi", |
| "--query-gpu=driver_version", |
| "--format=csv,noheader,nounits", |
| ], |
| check=True, |
| capture_output=True, |
| text=True, |
| timeout=10, |
| ) |
| except (OSError, subprocess.SubprocessError) as error: |
| raise RuntimeError("cannot query the NVIDIA driver version with nvidia-smi") from error |
| values = {line.strip() for line in completed.stdout.splitlines() if line.strip()} |
| if len(values) != 1: |
| raise RuntimeError("nvidia-smi did not return one consistent NVIDIA driver version") |
| value = next(iter(values)) |
| if re.fullmatch(r"[0-9]+\.[0-9]+(?:\.[0-9]+)*", value) is None: |
| raise RuntimeError("nvidia-smi did not return a semantic NVIDIA driver version") |
| return value |
|
|
|
|
| def _cuda_version(value: object) -> str: |
| if type(value) is not int or value < 1000: |
| raise RuntimeError("cannot query the loaded CUDA runtime version") |
| major = value // 1000 |
| remainder = value % 1000 |
| minor = remainder // 10 |
| patch = remainder % 10 |
| return f"{major}.{minor}" + (f".{patch}" if patch else "") |
|
|
|
|
| def _cuda_runtime_version_from_torch(torch_cuda: object) -> int: |
| cudart_factory = getattr(torch_cuda, "cudart", None) |
| if not callable(cudart_factory): |
| raise RuntimeError("PyTorch does not expose torch.cuda.cudart") |
| cudart = cudart_factory() |
| runtime_query = getattr(cudart, "cudaRuntimeGetVersion", None) |
| if not callable(runtime_query): |
| raise RuntimeError("PyTorch cudart does not expose cudaRuntimeGetVersion") |
| return runtime_query() |
|
|
|
|
| def _cuda_runtime_version_from_ctypes() -> int: |
| """Query libcudart directly when PyTorch does not bind the runtime symbol.""" |
|
|
| import ctypes |
| import ctypes.util |
|
|
| candidates: list[str | None] = [None] |
| discovered = ctypes.util.find_library("cudart") |
| if discovered: |
| candidates.append(discovered) |
| candidates.extend(["libcudart.so", "libcudart.so.12", "libcudart.so.11.0"]) |
|
|
| errors: list[str] = [] |
| for candidate in dict.fromkeys(candidates): |
| label = candidate or "process" |
| try: |
| library = ctypes.CDLL(candidate) if candidate is not None else ctypes.CDLL(None) |
| runtime_query = getattr(library, "cudaRuntimeGetVersion") |
| except (AttributeError, OSError) as error: |
| errors.append(f"{label}: {error}") |
| continue |
| version = ctypes.c_int() |
| try: |
| runtime_query.argtypes = [ctypes.POINTER(ctypes.c_int)] |
| runtime_query.restype = ctypes.c_int |
| status = runtime_query(ctypes.byref(version)) |
| except (AttributeError, TypeError, ValueError) as error: |
| errors.append(f"{label}: {error}") |
| continue |
| if status != 0: |
| errors.append(f"{label}: cudaRuntimeGetVersion returned {status}") |
| continue |
| return version.value |
|
|
| detail = "; ".join(errors[-3:]) |
| suffix = f": {detail}" if detail else "" |
| raise RuntimeError(f"cannot query CUDA runtime version via libcudart{suffix}") |
|
|
|
|
| def _cuda_runtime_version(torch_cuda: object) -> int: |
| errors: list[str] = [] |
| for label, query in ( |
| ("PyTorch cudart", lambda: _cuda_runtime_version_from_torch(torch_cuda)), |
| ("libcudart", _cuda_runtime_version_from_ctypes), |
| ): |
| try: |
| value = query() |
| _cuda_version(value) |
| except (AttributeError, OSError, RuntimeError, TypeError, ValueError) as error: |
| errors.append(f"{label}: {error}") |
| continue |
| return value |
| raise RuntimeError( |
| "cannot query the loaded CUDA runtime version: " + "; ".join(errors) |
| ) |
|
|
|
|
| def cuda_runtime_metadata() -> dict[str, object]: |
| """Collect GPU evidence after initialization but outside every timed interval.""" |
|
|
| import torch |
|
|
| if not torch.cuda.is_available(): |
| raise RuntimeError("generation timing requires a CUDA server") |
| device = torch.cuda.current_device() |
| compiled_toolkit = getattr(torch.version, "cuda", None) |
| if not isinstance(compiled_toolkit, str) or not compiled_toolkit: |
| raise RuntimeError("PyTorch does not report its compiled CUDA toolkit version") |
| try: |
| runtime_version = _cuda_version(_cuda_runtime_version(torch.cuda)) |
| capability = torch.cuda.get_device_capability(device) |
| except (AttributeError, OSError, RuntimeError, TypeError, ValueError) as error: |
| raise RuntimeError("cannot query the loaded CUDA runtime and selected GPU") from error |
| if ( |
| not isinstance(capability, (tuple, list)) |
| or len(capability) != 2 |
| or any(type(part) is not int or part < 0 for part in capability) |
| ): |
| raise RuntimeError("selected GPU compute capability is invalid") |
| return { |
| "gpu_name": str(torch.cuda.get_device_name(device)), |
| "gpu_index": device, |
| "gpu_compute_capability": f"{capability[0]}.{capability[1]}", |
| "cuda_runtime_version": runtime_version, |
| "pytorch_compiled_cuda_toolkit": compiled_toolkit, |
| "nvidia_driver_version": _driver_version(), |
| "synchronization_policy": "torch.cuda.synchronize_before_start_and_after_generate", |
| } |
|
|
|
|
| def _sha256_path(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as handle: |
| for chunk in iter(lambda: handle.read(1024 * 1024), b""): |
| digest.update(chunk) |
| return digest.hexdigest() |
|
|
|
|
| def runtime_code_binding(model: object, generate_one: Callable[[], object]) -> list[dict[str, str]]: |
| """Bind loaded model classes and the executed sampling callable to source bytes.""" |
|
|
| candidates: list[tuple[str, object]] = [("sampling_callable", generate_one), ("model", model)] |
| seen_objects: set[int] = set() |
| cursor = 0 |
| while cursor < len(candidates): |
| _, candidate = candidates[cursor] |
| cursor += 1 |
| identity = id(candidate) |
| if identity in seen_objects: |
| continue |
| seen_objects.add(identity) |
| modules = getattr(candidate, "modules", None) |
| if callable(modules): |
| try: |
| candidates.extend(("loaded_module", module) for module in modules()) |
| except (RuntimeError, TypeError): |
| pass |
| for attribute in ("model", "module", "backbone", "net"): |
| child = getattr(candidate, attribute, None) |
| if child is not None and child is not candidate: |
| candidates.append((f"loaded_{attribute}", child)) |
| if inspect.ismethod(candidate): |
| candidates.append(("bound_method_owner", candidate.__self__)) |
| closure = getattr(candidate, "__closure__", None) |
| if closure: |
| for cell in closure: |
| try: |
| child = cell.cell_contents |
| except ValueError: |
| continue |
| if callable(child): |
| candidates.append(("sampling_closure", child)) |
|
|
| records: dict[tuple[str, str, str], dict[str, str]] = {} |
| for kind, candidate in candidates: |
| target = candidate if inspect.isfunction(candidate) or inspect.ismethod(candidate) else type(candidate) |
| try: |
| source = inspect.getsourcefile(target) |
| except (OSError, TypeError): |
| source = None |
| if not source: |
| continue |
| path = Path(source).resolve() |
| if path.is_symlink() or not path.is_file(): |
| continue |
| digest = _sha256_path(path) |
| symbol = getattr(target, "__qualname__", type(candidate).__qualname__) |
| records[(str(path), digest, symbol)] = { |
| "path": str(path), |
| "sha256": digest, |
| "kind": kind, |
| "symbol": symbol, |
| } |
| if not records: |
| raise ValueError("cannot bind benchmark execution to loaded Python source") |
| return sorted(records.values(), key=lambda item: (item["path"], item["symbol"])) |
|
|
|
|
| def observed_model_precision(model: object) -> str: |
| """Record parameter storage dtype separately from upstream compute autocast.""" |
|
|
| parameters = getattr(model, "parameters", None) |
| if not callable(parameters): |
| raise ValueError("benchmark model does not expose parameters for precision validation") |
| aliases = { |
| "torch.float32": "fp32", |
| "float32": "fp32", |
| "torch.float16": "fp16", |
| "float16": "fp16", |
| "torch.bfloat16": "bf16", |
| "bfloat16": "bf16", |
| } |
| observed = {aliases.get(str(parameter.dtype), str(parameter.dtype)) for parameter in parameters()} |
| if not observed: |
| raise ValueError("benchmark model has no parameters for precision validation") |
| unsupported = observed - {"fp32", "fp16", "bf16"} |
| if unsupported: |
| raise ValueError("benchmark model has unsupported parameter precision: " + ", ".join(sorted(unsupported))) |
| return next(iter(observed)) if len(observed) == 1 else "mixed(" + ",".join(sorted(observed)) + ")" |
|
|
|
|
| def benchmark_and_publish( |
| generate_one: Callable[[], object], |
| *, |
| model: object, |
| output: Path, |
| metadata_path: Path, |
| precision: str, |
| synchronize: Callable[[], object] | None = None, |
| ) -> object: |
| """Server hook used only after a concrete adapter has loaded its real model.""" |
|
|
| import torch |
|
|
| if precision != "author": |
| raise ValueError("server timing accepts only the pinned author precision policy") |
| parameter_precision = observed_model_precision(model) |
| metadata = load_metadata(metadata_path) |
| if ( |
| metadata.get("requested_precision") != "author" |
| or not isinstance(metadata.get("precision"), str) |
| or not isinstance(metadata.get("precision_policy"), str) |
| or not isinstance(metadata.get("precision_evidence"), str) |
| or not isinstance(metadata.get("precision_policy_binding"), dict) |
| ): |
| raise ValueError("benchmark metadata does not bind an authoritative precision policy") |
| synchronize = synchronize or torch.cuda.synchronize |
| latest: object = None |
|
|
| def generate() -> object: |
| nonlocal latest |
| latest = generate_one() |
| return latest |
|
|
| result = benchmark(generate, synchronize, warmups=5, repeats=32, batch_size=1) |
| metadata = { |
| **metadata, |
| "parameter_precision": parameter_precision, |
| "runtime_code_binding": runtime_code_binding(model, generate_one), |
| **cuda_runtime_metadata(), |
| } |
| publish_timing(output, result, metadata) |
| return latest |
|
|
