|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""Measure Stage-B full-content identity overhead on physical GPU backends. |
| 3 | +
|
| 4 | +This benchmark isolates the cost of the collision-resistant X/y SHA-256 used by |
| 5 | +Hausman sample identity. It compares ordinary PanelOLS/RandomEffects fit time |
| 6 | +against the same fit with only ``_full_content_digest`` replaced by a constant |
| 7 | +stub. Numerical work, low-order audit reductions, estimator setup, and all other |
| 8 | +Stage-B code remain unchanged in the baseline. |
| 9 | +
|
| 10 | +The script is intentionally separate from ``validate_panel_stage_b_gpu.py``: |
| 11 | +that runner remains correctness/provenance-only and its frontend source must not |
| 12 | +acquire inferred timing or speedup fields. |
| 13 | +""" |
| 14 | + |
| 15 | +from __future__ import annotations |
| 16 | + |
| 17 | +import argparse |
| 18 | +import importlib.metadata |
| 19 | +import json |
| 20 | +import platform |
| 21 | +import subprocess |
| 22 | +import time |
| 23 | +from pathlib import Path |
| 24 | + |
| 25 | +import numpy as np |
| 26 | + |
| 27 | +import statgpu.panel._diagnostics as diagnostics |
| 28 | +from statgpu.panel import PanelOLS, RandomEffects |
| 29 | + |
| 30 | + |
| 31 | +def _git_sha() -> str: |
| 32 | + return subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip() |
| 33 | + |
| 34 | + |
| 35 | +def _git_status_porcelain() -> str: |
| 36 | + return subprocess.check_output(["git", "status", "--porcelain"], text=True) |
| 37 | + |
| 38 | + |
| 39 | +def _version(name: str): |
| 40 | + try: |
| 41 | + return importlib.metadata.version(name) |
| 42 | + except importlib.metadata.PackageNotFoundError: |
| 43 | + return None |
| 44 | + |
| 45 | + |
| 46 | +def _parse_scales(value: str): |
| 47 | + out = [] |
| 48 | + for item in value.split(","): |
| 49 | + n_text, k_text = item.strip().lower().split("x", 1) |
| 50 | + n, k = int(n_text), int(k_text) |
| 51 | + if n <= 0 or k <= 0: |
| 52 | + raise ValueError("benchmark scales must be positive NxK pairs") |
| 53 | + out.append((n, k)) |
| 54 | + return out |
| 55 | + |
| 56 | + |
| 57 | +def _sync(backend: str): |
| 58 | + if backend == "cupy": |
| 59 | + import cupy as cp |
| 60 | + |
| 61 | + cp.cuda.Stream.null.synchronize() |
| 62 | + elif backend == "torch": |
| 63 | + import torch |
| 64 | + |
| 65 | + torch.cuda.synchronize() |
| 66 | + |
| 67 | + |
| 68 | +def _to_backend(X, y, entity, backend: str): |
| 69 | + if backend == "cupy": |
| 70 | + import cupy as cp |
| 71 | + |
| 72 | + return ( |
| 73 | + cp.asarray(X), |
| 74 | + cp.asarray(y), |
| 75 | + cp.asarray(entity, dtype=cp.int64), |
| 76 | + ) |
| 77 | + if backend == "torch": |
| 78 | + import torch |
| 79 | + |
| 80 | + return ( |
| 81 | + torch.as_tensor(X, dtype=torch.float64, device="cuda"), |
| 82 | + torch.as_tensor(y, dtype=torch.float64, device="cuda"), |
| 83 | + torch.as_tensor(entity, dtype=torch.int64, device="cuda"), |
| 84 | + ) |
| 85 | + raise ValueError(backend) |
| 86 | + |
| 87 | + |
| 88 | +def _device_arg(backend: str): |
| 89 | + return {"cupy": "cuda", "torch": "torch"}[backend] |
| 90 | + |
| 91 | + |
| 92 | +def _dataset(n: int, k: int, seed: int): |
| 93 | + rng = np.random.default_rng(seed) |
| 94 | + X = rng.normal(size=(n, k)).astype(np.float64) |
| 95 | + beta = np.linspace(0.2, 0.8, k, dtype=np.float64) |
| 96 | + entity = np.arange(n, dtype=np.int64) // 20 |
| 97 | + n_entities = int(entity.max()) + 1 |
| 98 | + alpha = np.linspace(-0.5, 0.5, n_entities, dtype=np.float64)[entity] |
| 99 | + y = X @ beta + alpha + rng.normal(scale=0.2, size=n) |
| 100 | + return X, y.astype(np.float64), entity |
| 101 | + |
| 102 | + |
| 103 | +def _fit(model_name: str, X, y, entity, backend: str): |
| 104 | + device = _device_arg(backend) |
| 105 | + if model_name == "PanelOLS": |
| 106 | + model = PanelOLS(entity_effects=True, cov_type="nonrobust", device=device) |
| 107 | + elif model_name == "RandomEffects": |
| 108 | + model = RandomEffects(device=device) |
| 109 | + else: |
| 110 | + raise ValueError(model_name) |
| 111 | + model.fit(X, y, entity_ids=entity) |
| 112 | + return model |
| 113 | + |
| 114 | + |
| 115 | +def _timed_fit(model_name, X, y, entity, backend, *, disable_digest: bool): |
| 116 | + original = diagnostics._full_content_digest |
| 117 | + if disable_digest: |
| 118 | + diagnostics._full_content_digest = lambda _X, _y: "0" * 64 |
| 119 | + try: |
| 120 | + _sync(backend) |
| 121 | + start = time.perf_counter() |
| 122 | + _fit(model_name, X, y, entity, backend) |
| 123 | + _sync(backend) |
| 124 | + return time.perf_counter() - start |
| 125 | + finally: |
| 126 | + diagnostics._full_content_digest = original |
| 127 | + |
| 128 | + |
| 129 | +def _median(values): |
| 130 | + return float(np.median(np.asarray(values, dtype=np.float64))) |
| 131 | + |
| 132 | + |
| 133 | +def main(): |
| 134 | + parser = argparse.ArgumentParser(description=__doc__) |
| 135 | + parser.add_argument("--out", type=Path, required=True) |
| 136 | + parser.add_argument("--expected-sha", required=True) |
| 137 | + parser.add_argument("--backends", default="cupy,torch") |
| 138 | + parser.add_argument( |
| 139 | + "--scales", |
| 140 | + default="10000x2,100000x2,100000x10,500000x2", |
| 141 | + help="comma-separated NxK pairs", |
| 142 | + ) |
| 143 | + parser.add_argument("--repeats", type=int, default=3) |
| 144 | + args = parser.parse_args() |
| 145 | + |
| 146 | + sha = _git_sha() |
| 147 | + if sha != args.expected_sha: |
| 148 | + raise RuntimeError(f"wrong source head: {sha} != {args.expected_sha}") |
| 149 | + dirty = _git_status_porcelain() |
| 150 | + if dirty.strip(): |
| 151 | + raise RuntimeError( |
| 152 | + "identity benchmark requires a clean working tree; uncommitted changes:\n" |
| 153 | + + dirty |
| 154 | + ) |
| 155 | + if args.repeats < 1: |
| 156 | + raise ValueError("--repeats must be positive") |
| 157 | + |
| 158 | + backends = [x.strip() for x in args.backends.split(",") if x.strip()] |
| 159 | + if not backends or any(x not in {"cupy", "torch"} for x in backends): |
| 160 | + raise ValueError("--backends must contain cupy and/or torch") |
| 161 | + scales = _parse_scales(args.scales) |
| 162 | + |
| 163 | + rows = [] |
| 164 | + for scale_index, (n, k) in enumerate(scales): |
| 165 | + X_np, y_np, entity_np = _dataset(n, k, seed=20260808 + scale_index) |
| 166 | + for backend in backends: |
| 167 | + X, y, entity = _to_backend(X_np, y_np, entity_np, backend) |
| 168 | + for model_name in ("PanelOLS", "RandomEffects"): |
| 169 | + # Warm both paths before measurement to avoid one-time import/ |
| 170 | + # allocator effects being attributed to the digest. |
| 171 | + _timed_fit( |
| 172 | + model_name, X, y, entity, backend, disable_digest=False |
| 173 | + ) |
| 174 | + _timed_fit( |
| 175 | + model_name, X, y, entity, backend, disable_digest=True |
| 176 | + ) |
| 177 | + |
| 178 | + with_digest = [] |
| 179 | + without_digest = [] |
| 180 | + for _ in range(args.repeats): |
| 181 | + with_digest.append( |
| 182 | + _timed_fit( |
| 183 | + model_name, |
| 184 | + X, |
| 185 | + y, |
| 186 | + entity, |
| 187 | + backend, |
| 188 | + disable_digest=False, |
| 189 | + ) |
| 190 | + ) |
| 191 | + without_digest.append( |
| 192 | + _timed_fit( |
| 193 | + model_name, |
| 194 | + X, |
| 195 | + y, |
| 196 | + entity, |
| 197 | + backend, |
| 198 | + disable_digest=True, |
| 199 | + ) |
| 200 | + ) |
| 201 | + |
| 202 | + normal = _median(with_digest) |
| 203 | + baseline = _median(without_digest) |
| 204 | + overhead = normal - baseline |
| 205 | + ratio = normal / baseline if baseline > 0.0 else None |
| 206 | + rows.append( |
| 207 | + { |
| 208 | + "backend": backend, |
| 209 | + "model": model_name, |
| 210 | + "n_samples": n, |
| 211 | + "n_features": k, |
| 212 | + "repeats": args.repeats, |
| 213 | + "with_digest_seconds": normal, |
| 214 | + "without_digest_seconds": baseline, |
| 215 | + "digest_overhead_seconds": overhead, |
| 216 | + "with_over_without_ratio": ratio, |
| 217 | + "with_digest_samples": with_digest, |
| 218 | + "without_digest_samples": without_digest, |
| 219 | + } |
| 220 | + ) |
| 221 | + |
| 222 | + payload = { |
| 223 | + "schema_version": 1, |
| 224 | + "git_sha": sha, |
| 225 | + "working_tree_clean": True, |
| 226 | + "benchmark": "panel_stage_b_full_content_identity_overhead", |
| 227 | + "timing_scope": "end-to-end estimator fit with vs without only the SHA-256 full-content digest", |
| 228 | + "target_scale_source": "PR122 fresh-review performance finding", |
| 229 | + "environment": { |
| 230 | + "python": platform.python_version(), |
| 231 | + "platform": platform.platform(), |
| 232 | + "packages": { |
| 233 | + name: _version(name) |
| 234 | + for name in ("statgpu", "numpy", "cupy", "torch") |
| 235 | + }, |
| 236 | + }, |
| 237 | + "rows": rows, |
| 238 | + } |
| 239 | + args.out.parent.mkdir(parents=True, exist_ok=True) |
| 240 | + args.out.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") |
| 241 | + print(json.dumps(payload, indent=2)) |
| 242 | + print(f"PASS — identity-overhead benchmark recorded: {args.out}") |
| 243 | + |
| 244 | + |
| 245 | +if __name__ == "__main__": |
| 246 | + main() |
0 commit comments