|
| 1 | +"""Regenerate the op-noise artefacts on Modal's CPU and sync them back locally. |
| 2 | +
|
| 3 | +`make measure-noise` runs on the local CPU, but the canonical drift check |
| 4 | +(`tests/docs/test_numerical_noise_drift.py` under `make test`) runs on Modal's |
| 5 | +CPU, and a few measurements sit on precision boundaries where local-CPU and |
| 6 | +Modal-CPU round differently. This entrypoint runs `scripts.measure_op_noise` on a |
| 7 | +Modal CPU worker (the same environment as the drift check) and returns the |
| 8 | +regenerated `docs/op_noise_data.json`, `docs/numerical_noise.md`, and every |
| 9 | +`torchwright/ops/*.py` (for the in-place noise-footer edits) so the local tree |
| 10 | +matches what `make test` will measure. |
| 11 | +
|
| 12 | +This is the artefact-sync-back case the CLAUDE.md "Running scripts on GPU" rules |
| 13 | +carve out as the only acceptable reason for a purpose-built `modal_*.py` — and it |
| 14 | +imports `IMAGE` from `modal_image` rather than duplicating it. |
| 15 | +
|
| 16 | + uv run modal run modal_measure_noise.py |
| 17 | +""" |
| 18 | + |
| 19 | +import pathlib |
| 20 | +import subprocess |
| 21 | +import sys |
| 22 | + |
| 23 | +import modal |
| 24 | + |
| 25 | +from modal_image import IMAGE |
| 26 | + |
| 27 | +app = modal.App("torchwright-measure-noise", image=IMAGE) |
| 28 | + |
| 29 | + |
| 30 | +# Match the test shard's container exactly (modal_test.py: gpu a100-80gb, cpu 8, |
| 31 | +# memory 32768). `_measure_all` forces CPU compute, but reproducing the same |
| 32 | +# container is what makes a few precision-boundary measurements (e.g. reciprocal |
| 33 | +# p99) match what the drift check under `make test` measures. |
| 34 | +@app.function(gpu="a100-80gb", cpu=8, memory=32768, timeout=1800) |
| 35 | +def measure() -> dict: |
| 36 | + rc = subprocess.run( |
| 37 | + [sys.executable, "-m", "scripts.measure_op_noise"] |
| 38 | + ).returncode |
| 39 | + if rc != 0: |
| 40 | + raise RuntimeError(f"scripts.measure_op_noise exited {rc}") |
| 41 | + |
| 42 | + from scripts.measure_op_noise import DOCS_JSON, DOCS_MD, REPO_ROOT |
| 43 | + |
| 44 | + out: dict[str, str] = {} |
| 45 | + for p in (DOCS_JSON, DOCS_MD): |
| 46 | + out[str(p.relative_to(REPO_ROOT))] = p.read_text() |
| 47 | + for f in sorted((REPO_ROOT / "torchwright" / "ops").glob("*.py")): |
| 48 | + out[str(f.relative_to(REPO_ROOT))] = f.read_text() |
| 49 | + return out |
| 50 | + |
| 51 | + |
| 52 | +@app.local_entrypoint() |
| 53 | +def main(): |
| 54 | + files = measure.remote() |
| 55 | + root = pathlib.Path(__file__).resolve().parent |
| 56 | + for rel, content in sorted(files.items()): |
| 57 | + (root / rel).write_text(content) |
| 58 | + print(f"wrote {rel} ({len(content)} bytes)") |
0 commit comments