From 91a44efa5c0677b511f7c310bd1b7a50f037fbe9 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Thu, 6 Aug 2026 15:12:26 -0600 Subject: [PATCH] pgw#989: the dynamo mint's `warmup_forward` hour gets a breakdown, and `inductor_compile: 0.0` stops lying MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every published JIT cell reported the same phase table — sdxl w8a8-lora64, L40S, gen-worker 0.93.1, pod wz4g1ya8a4khne, 2026-08-06: {'load': 13.87, 'warmup_forward': 4416.91, 'inductor_compile': 0.0, 'seal_publish': 96.809, 'finalize': 0.155} 97.6 % of the mint under one name, next to a zero named after the work. The zero is not a broken clock. `mint_child` opened the `inductor_compile` frame around `_drain_router`, and a fleet mint arms COLD with no router (gw#587), so that phase measured an empty queue — while every compile ran INLINE inside the warm forwards, i.e. inside `warmup_forward`. The drain is now `router_drain`. `gen_worker.warm_spans` measures the split the bucket was hiding: per warm job and in total, `warm_compile_s` vs `warm_execute_s`, with the inductor partition inside the compile half and an explicit `compile_other_s` residual. The parent re-emits it as `jit_compile` rows (`warm:totals`, `warm:`, `warm_job:`), so an AOT-vs-JIT comparison stays one grouped query. The key set is the JIT one, and that is the reason for a second module rather than reuse: MEASURED on the pin (torch 2.13.0+cu130), the AOT partition prices one `torch.compile` call at 1.104 s of 5.054 s — 22 % — because `AotCodeCompiler.compile` is AOT-only. `_compile.compile_inner` is the JIT total and `PyCodeCache.load_by_key_path` (63 %) is where a JIT compile's time actually goes. Tests: 9 rows at the real seams over a frozen REAL metric delta — the partition reconciles to the total with a named residual, the AOT key set is shown to miss 78 % of the same compile, overlays are never summed in, a failing warm job keeps its seconds, an unmeasurable wall OMITS the ratio rather than reporting "compiled 0 %", and the parent emits/stays silent correctly. Suite: 3302 passed, 37 skipped, 1 xfailed; mypy clean (232 files); ruff clean. --- changelog.d/pgw989.md | 23 +++ src/gen_worker/activity.py | 6 + src/gen_worker/mint_child.py | 32 +++-- src/gen_worker/mint_delegate.py | 74 +++++++++- src/gen_worker/warm_spans.py | 225 +++++++++++++++++++++++++++++ tests/test_warm_spans_pgw989.py | 241 ++++++++++++++++++++++++++++++++ 6 files changed, 587 insertions(+), 14 deletions(-) create mode 100644 changelog.d/pgw989.md create mode 100644 src/gen_worker/warm_spans.py create mode 100644 tests/test_warm_spans_pgw989.py diff --git a/changelog.d/pgw989.md b/changelog.d/pgw989.md new file mode 100644 index 00000000..f1c33648 --- /dev/null +++ b/changelog.d/pgw989.md @@ -0,0 +1,23 @@ +- **pgw#989: a dynamo mint's hour gets a breakdown, and the row that never + measured a compile stops claiming to.** Every published JIT cell reported the + same shape — sdxl w8a8-lora64 on an L40S, 2026-08-06: `{'load': 13.87, + '**warmup_forward': 4416.91**, '**inductor_compile': 0.0**, 'seal_publish': + 96.81, 'finalize': 0.16}`. 97.6 % of the mint under one name, beside a zero + named after the work. Not a broken clock: `mint_child` framed + `_drain_router` as `inductor_compile`, and a fleet mint arms COLD with no + router (gw#587), so that phase measured an empty queue while every compile + ran INLINE inside the warm forwards. The drain is now `router_drain`, which + is what it is. +- The warm plan now carries its own ledger (`gen_worker.warm_spans`): per warm + job and in total, `warm_compile_s` vs `warm_execute_s` — with the inductor + partition inside the compile half and an explicit `compile_other_s` residual, + so a newly-introduced phase shows up as the residual growing rather than as + time silently vanishing (pgw#830's rule, applied to the JIT path). The parent + re-emits it as `jit_compile` rows `warm:totals` / `warm:` / + `warm_job:`, so an AOT-vs-JIT comparison stays one grouped query. +- The key set is the JIT one and that difference is the point: MEASURED on the + pin, the AOT partition prices a `torch.compile` call at **1.10 s of 5.05 s + (22 %)** because `AotCodeCompiler.compile` never runs on this path. + `_compile.compile_inner` is the JIT total, and `PyCodeCache.load_by_key_path` + (63 %) is where the time actually is. `MintReport.mint_phases`, documented as + "empty for the dynamo recipe", now carries this. diff --git a/src/gen_worker/activity.py b/src/gen_worker/activity.py index eff5dcc1..c58647ab 100644 --- a/src/gen_worker/activity.py +++ b/src/gen_worker/activity.py @@ -42,6 +42,7 @@ import psutil from . import progress as progress_mod +from . import warm_spans from .pb import worker_scheduler_pb2 as pb logger = logging.getLogger(__name__) @@ -119,6 +120,11 @@ PHASE_LOAD = "load" PHASE_TRACE_GRAPH = "trace_graph" PHASE_INDUCTOR_COMPILE = "inductor_compile" +# pgw#989: the dynamo mint used to report its router drain under +# PHASE_INDUCTOR_COMPILE, next to a `warmup_forward` row holding every compile +# it ever ran. Re-exported (defined in `warm_spans`, which the mint child can +# import without protobuf) so the vocabulary is enumerable from one module. +PHASE_ROUTER_DRAIN = warm_spans.PHASE_ROUTER_DRAIN PHASE_WARMUP_FORWARD = "warmup_forward" PHASE_SEAL_PUBLISH = "seal_publish" # gw#612: post-proof tail — sibling-lane resolution, publish decision, diff --git a/src/gen_worker/mint_child.py b/src/gen_worker/mint_child.py index a13427e0..5fd666b8 100644 --- a/src/gen_worker/mint_child.py +++ b/src/gen_worker/mint_child.py @@ -69,7 +69,7 @@ import msgspec -from . import worker_goals +from . import warm_spans, worker_goals from .config import load_settings from .mint_process import ( EXIT_BAD_REQUEST, @@ -460,9 +460,16 @@ async def _drain() -> None: def _drive_warm_plan( instance: Any, jobs: Sequence[Any], request: MintRequest, *, proof_only: bool = False, -) -> None: +) -> warm_spans.WarmLedger: """Run the endpoint's OWN warm plan, framed as ``warmup_forward``. + Returns the plan's own cost ledger (pgw#989). On the dynamo recipe these + forwards ARE the compile, so ``warmup_forward`` is 97.6 % of the mint under + a single name — measured beside an ``inductor_compile`` row reading 0.0 s. + The ledger splits it into compile and forward per job. It is measured on + the AOT recipe's proof forward too: that job costs real seconds and an + unmeasured cost is how this one hid. + ``proof_only`` runs ONE job (pgw#984, the AOT recipe); otherwise the whole plan runs (the dynamo recipe, where these forwards ARE the compile). @@ -475,15 +482,18 @@ def _drive_warm_plan( ``(phase=warmup_forward, deterministic)`` while the worker was still calling it ``crashed`` and buying the second pod. """ + ledger = warm_spans.WarmLedger() total = 1 if proof_only else len(jobs) frame(phase="warmup_forward", step=0, total=total) for index, job in enumerate(jobs[:total], start=1): frame(phase="warmup_forward", step=index, total=total, note=job.spec.name) try: - _run_warm_job( - instance, job, dict(request.configs.get(job.spec.name) or {}), - request.execution_lane, origin=mint_identity(request)) + with ledger.job(job.spec.name): + _run_warm_job( + instance, job, + dict(request.configs.get(job.spec.name) or {}), + request.execution_lane, origin=mint_identity(request)) except BaseException as exc: if _is_resource_error(exc) or not isinstance(exc, Exception): raise @@ -492,6 +502,7 @@ def _drive_warm_plan( f"not run — warm job {job.spec.name!r} raised " f"{type(exc).__name__}: {exc}. A cell must not seal for a " f"handler that cannot serve.") from exc + return ledger def _drain_router(pipe: Any, *, poll_s: float = 0.5) -> None: @@ -514,7 +525,7 @@ def _drain_router(pipe: Any, *, poll_s: float = 0.5) -> None: _warm, pending, _failed = router.stats() if pending == 0: return - frame(phase="inductor_compile", note=f"{pending} compile(s) queued") + frame(phase=warm_spans.PHASE_ROUTER_DRAIN, note=f"{pending} compile(s) queued") time.sleep(poll_s) @@ -865,8 +876,9 @@ def _load(_execution_lane: str) -> Tuple[Any, Any, Any]: raise MintChildRefused(f"{mint_identity(request)}: {exc}") from exc miss_before = cc.cache_miss_count(pipe) - _drive_warm_plan(instance, jobs, request) - frame(phase="inductor_compile", note="draining any queued compiles") + ledger = _drive_warm_plan(instance, jobs, request) + frame(phase=warm_spans.PHASE_ROUTER_DRAIN, + note="draining any queued compiles") _drain_router(pipe) if cc.execution_count(pipe) <= 0: @@ -898,6 +910,10 @@ def _load(_execution_lane: str) -> Tuple[Any, Any, Any]: peak_vram_bytes=peak, elapsed_s=time.monotonic() - started, phases=_close_phases(), + # pgw#989: the dynamo recipe's `mint_phases` was documented as "empty — + # no per-graph-class breakdown". It has one; it was just never + # measured. The parent re-emits this exactly as it does the AOT table. + mint_phases=ledger.table(), ) diff --git a/src/gen_worker/mint_delegate.py b/src/gen_worker/mint_delegate.py index 581b7b95..5d1656a0 100644 --- a/src/gen_worker/mint_delegate.py +++ b/src/gen_worker/mint_delegate.py @@ -34,14 +34,14 @@ import tempfile from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Dict, Optional, Tuple +from typing import Any, Dict, Mapping, Optional, Tuple from . import activity as activity_mod from . import aot_resume from . import mint_budget from . import mint_process from . import progress as progress_mod -from .mint_process import MintOutcome, MintRequest +from .mint_process import MintOutcome, MintReport, MintRequest logger = logging.getLogger(__name__) @@ -492,10 +492,16 @@ def _emit_jit_compile( """th#1322: one delegated JIT mint's duration, as typed NUMERIC events. This is the fleet's real JIT compile path — the child arms COLD and drives - the endpoint's own warm plan (gw#587), so its `warmup_forward` + - `inductor_compile` spans ARE "how long does JIT take". Before this the - number lived only in the child's stdout, and a serve pod exposes no logs - (pgw#760), so it was unrecoverable the moment the pod went away. + the endpoint's own warm plan (gw#587), so its `warmup_forward` span IS + "how long does JIT take". Before this the number lived only in the child's + stdout, and a serve pod exposes no logs (pgw#760), so it was unrecoverable + the moment the pod went away. + + pgw#989: that span is also 97.6 % of the mint, so "how long" was as far as + it went. ``report.mint_phases`` now carries the warm plan's own ledger + (:class:`gen_worker.warm_spans.WarmLedger`) and it is emitted here under + `phase=warm:` / `phase=warm_job:` — compile vs forward, and the + inductor partition inside the compile half. Shape matches ``aot_mint_phases`` exactly: `phase=minted` carries the total, `phase=child:` the spans inside it. A mint that did NOT produce a @@ -520,6 +526,7 @@ def _emit_jit_compile( phase=f"child:{name}", duration_ms=int(round(value * 1000)), ) + _emit_warm_ledger(report, head=head) # The child's own elapsed is authoritative when it wrote a report; the # parent's wall clock covers a child that died before writing one (spawn # + run, which is the honest cost of that attempt). @@ -541,6 +548,61 @@ def _emit_jit_compile( exc_info=True) +def _emit_warm_ledger(report: Optional[MintReport], *, head: str) -> None: + """pgw#989: the warm plan's compile-vs-forward split, as numeric events. + + Three row shapes, all under ``jit_compile`` so one grouped query still + covers a JIT mint: + + * ``warm:totals`` — the roll-up, carrying ``warm_compile_s`` in + ``duration_ms``. The one number that says whether a slow mint is a slow + COMPILE. + * ``warm:`` — the inductor partition inside the compile half. + * ``warm_job:`` — one warm forward. A plan whose jobs mostly compile + nothing is paying full forward cost for coverage it already has; that is + a different defect from a slow compile, and only the per-job rows can + tell them apart. + + Silent when the recipe produced no ledger (the AOT path, or a child that + died before the warm loop) — an absent measurement is reported as absence. + """ + table = dict(getattr(report, "mint_phases", None) or {}) \ + if report is not None else {} + totals = dict(table.get("totals") or {}) + if not totals.get("warm_wall_s"): + return + overlays = dict(table.get("overlays") or {}) + activity_mod.emit_event( + activity_mod.KIND_JIT_COMPILE, + f"{head} warm_totals={totals} overlays={overlays}", + phase="warm:totals", + duration_ms=int(round(float(totals.get("warm_compile_s") or 0.0) * 1000)), + ) + for name, seconds in sorted(dict(table.get("phases") or {}).items()): + value = float(seconds or 0.0) + if value <= 0: + continue + activity_mod.emit_event( + activity_mod.KIND_JIT_COMPILE, + f"{head} warm_phase={name} seconds={round(value, 2)}", + phase=f"warm:{name}", + duration_ms=int(round(value * 1000)), + ) + for row in table.get("jobs") or (): + if not isinstance(row, Mapping): + continue + wall = float(row.get("wall_s") or 0.0) + if wall <= 0: + continue + activity_mod.emit_event( + activity_mod.KIND_JIT_COMPILE, + f"{head} warm_job={row.get('job')} wall_s={round(wall, 2)} " + f"compile_s={row.get('compile_s')} execute_s={row.get('execute_s')}", + phase=f"warm_job:{row.get('job')}", + duration_ms=int(round(wall * 1000)), + ) + + def _emit_abort( outcome: MintOutcome, family: str, key: str, attempt: int, ) -> None: diff --git a/src/gen_worker/warm_spans.py b/src/gen_worker/warm_spans.py new file mode 100644 index 00000000..d63df75e --- /dev/null +++ b/src/gen_worker/warm_spans.py @@ -0,0 +1,225 @@ +"""pgw#989: what a dynamo mint's ``warmup_forward`` hour is actually made of. + +The number this deletes +----------------------- +The sdxl w8a8-lora64 mint of 2026-08-06 (cell +``ck5-a53e02a7885f9312fb1fa7df…``, L40S, gen-worker 0.93.1, pod +``wz4g1ya8a4khne``) reported this phase table, and the hub stored it verbatim:: + + {'load': 13.87, 'warmup_forward': 4416.91, 'inductor_compile': 0.0, + 'seal_publish': 96.809, 'finalize': 0.155} + +**97.6 % of the mint is one row, and the row named after the work is zero.** +The zero is not a bug in the clock: ``mint_child`` opens the +``inductor_compile`` frame around ``_drain_router``, and the fleet mint arms +COLD with no router (gw#587), so that phase measures an empty queue. Every +compile happens INLINE inside the warm forwards, i.e. inside +``warmup_forward``, which is therefore a bucket holding compile AND execution +AND everything else with no way to tell them apart. + +A mint hour with one name on it cannot be optimised — you cannot tell a +compile-bound mint from a warm plan that is simply running too many diffusion +steps. This module measures the split. + +How, and why these keys +----------------------- +Same instrument as the AOT path (:mod:`gen_worker.aot_compile_spans`): deltas +of dynamo's process-global ``compilation_time_metrics`` across a span. The KEY +SET is different and that difference is the whole point — the AOT partition +was derived for ``aot_compile`` and **silently under-reports a JIT compile**. +MEASURED on the pin (torch 2.13.0+cu130) over one ``torch.compile`` call, +5.06 s wall:: + + AOT partition total ......... 1.104 s (22 % of the compile) + host_compile_s ............ 0.000 s <- AotCodeCompiler.compile is AOT-only + _compile.compile_inner ...... 5.054 s (99.9 % of the compile) + PyCodeCache.load_by_key_path 3.203 s <- where a JIT compile actually goes + async_compile.wait ...... 2.886 s + +So the JIT total is ``_compile.compile_inner`` — the whole of ``torch.compile``, +front-end through kernel load — and the members below partition it. Anything +they do not claim lands in ``compile_other_s`` rather than vanishing, exactly +as pgw#830 established for the AOT ledger. + +Overlays (reported, never summed into the partition): ``triton_s`` and +``parallel_kernel_cpu_s`` nest inside ``kernel_compile_s``, and +``compile_file`` exceeds the wall because it sums the async-compile workers' +own threads. +""" + +from __future__ import annotations + +import contextlib +import time +from typing import Any, Dict, Iterator, List, Mapping, Tuple + +from . import aot_compile_spans + +#: Bump when the partition changes shape, so a reader never mixes two ledgers. +WARM_SPANS_V = 1 + +#: The honest name for the phase the dynamo mint used to call +#: ``inductor_compile``. It waits out a hot-swap router's queue — which the +#: fleet mint arms empty — and it never once measured a compile. Defined in +#: this leaf module rather than in :mod:`gen_worker.activity` because the mint +#: CHILD needs it and deliberately imports neither protobuf nor psutil; +#: ``activity`` re-exports it, so the phase vocabulary still has one home. +PHASE_ROUTER_DRAIN = "router_drain" + +#: The JIT compile's own total: everything ``torch.compile`` does for one +#: frame, measured by dynamo itself. NOT the AOT ``compile_s`` — that one is a +#: parent's Popen-to-reap wall around a child process, which the inline path +#: has no analogue of. +TOTAL_KEY = "_compile.compile_inner" + +#: Members of ``dynamo_compile_s``. Disjoint on the pin: the containment runs +#: compile_inner > call_user_compiler > fw_compiler_base > compile_fx_inner > +#: GraphLowering.compile_to_module > {codegen, PyCodeCache.load_by_key_path}, +#: so taking the LEAVES (and never their parents) is what keeps the sum honest. +JIT_PARTITION_KEYS: Dict[str, Tuple[str, ...]] = { + # Dynamo's front-end: bytecode tracing, guard construction, variable + # building. Cheap on a diffusion unet and reported anyway, because "the + # front-end is cheap" is a claim that should carry its own number. + "tracing_s": ("bytecode_tracing", "build_guards", "variable_builder_call"), + "graph_passes_s": ( + "_recursive_pre_grad_passes", + "_recursive_joint_graph_passes", + "_recursive_post_grad_passes", + ), + "lowering_s": ("GraphLowering.run",), + "codegen_s": ("GraphLowering.codegen",), + # The JIT counterpart of the AOT path's `host_compile_s`: writing the + # generated module out and waiting for the kernels it names. This is where + # a JIT compile's time actually is (63 % of the probe above). + "kernel_compile_s": ("PyCodeCache.load_by_key_path",), + # Kept so a cell minted through an AOT-shaped inline call is not silently + # unattributed; zero on the ordinary JIT path, which is itself a fact. + "host_compile_s": ("AotCodeCompiler.compile",), +} + +#: Reported under their own names, never summed with the partition — each is +#: known to nest inside a member above, and adding them would inflate +#: "attributed" while leaving the residual unexplained (pgw#830's second bug). +JIT_OVERLAY_KEYS: Dict[str, Tuple[str, ...]] = { + # Inside kernel_compile_s. + "async_wait_s": ("async_compile.wait",), + # Sums the async-compile WORKERS' own CPU, so it legitimately exceeds the + # wall — it prices the parallelism, it does not partition it. + "parallel_kernel_cpu_s": ("compile_file",), + "aot_dispatch_s": ("create_aot_dispatcher_function",), +} + + +def _sum(raw: Mapping[str, float], keys: Tuple[str, ...]) -> float: + return round(sum(float(raw.get(k, 0.0)) for k in keys), 3) + + +def partition(raw: Mapping[str, float]) -> Tuple[Dict[str, float], Dict[str, float]]: + """``(partition, overlays)`` for one span's raw metric delta. + + ``compile_other_s`` is the explicit residual of ``dynamo_compile_s`` and + can be negative only if the pin's containment changed — which is the + signal this shape exists to surface, so it is reported, not clamped. + """ + total = round(float(raw.get(TOTAL_KEY, 0.0)), 3) + members = { + label: _sum(raw, keys) for label, keys in JIT_PARTITION_KEYS.items()} + members["dynamo_compile_s"] = total + members["compile_other_s"] = round( + total - sum(v for k, v in members.items() if k != "dynamo_compile_s"), + 3) + overlays = { + label: value for label, keys in JIT_OVERLAY_KEYS.items() + if (value := _sum(raw, keys)) + } + triton = round(sum( + v for k, v in raw.items() + if "triton" in k.lower() and k != "compile_file"), 3) + if triton: + overlays["triton_s"] = triton + return members, overlays + + +class WarmLedger: + """The warm plan's cost, per job and in total. + + Not a profiler: a wall clock and one metric delta per job. The invariant it + keeps is that ``warm_wall_s = dynamo_compile_s + warm_execute_s`` by + construction, so the mint's biggest row can finally answer "compile, or + forward?" — and a warm plan that is slow for a THIRD reason shows up as + ``warm_execute_s`` growing rather than as time with no name. + """ + + def __init__(self) -> None: + self.jobs: List[Dict[str, Any]] = [] + self._raw: Dict[str, float] = {} + self._wall = 0.0 + + @contextlib.contextmanager + def job(self, name: str) -> Iterator[None]: + """Measure ONE warm forward. Never raises on the telemetry path — the + job's own exception propagates untouched.""" + before = aot_compile_spans.phase_snapshot() + started = time.monotonic() + try: + yield + finally: + wall = time.monotonic() - started + self._wall += wall + try: + _p, _o, raw = aot_compile_spans.phase_delta( + before, aot_compile_spans.phase_snapshot()) + for key, value in raw.items(): + self._raw[key] = round( + self._raw.get(key, 0.0) + float(value), 3) + compile_s = round(float(raw.get(TOTAL_KEY, 0.0)), 3) + self.jobs.append({ + "job": name, + "wall_s": round(wall, 3), + "compile_s": compile_s, + "execute_s": round(wall - compile_s, 3), + }) + except Exception: # noqa: BLE001 — telemetry never fails a mint + self.jobs.append({"job": name, "wall_s": round(wall, 3)}) + + def table(self) -> Dict[str, Any]: + """The flat, emittable ledger. Shaped like the AOT phase table + (``totals`` / ``phases`` / ``overlays`` / per-unit rows) so an + AOT-vs-JIT comparison stays ONE grouped query rather than two + vocabularies.""" + members, overlays = partition(self._raw) + compile_s = members["dynamo_compile_s"] + wall = round(self._wall, 3) + compiled = sum(1 for j in self.jobs if j.get("compile_s", 0.0) > 0.5) + return { + "spans_v": WARM_SPANS_V, + "totals": { + "warm_wall_s": wall, + "warm_compile_s": compile_s, + # The residual, and named as one: whatever the warm plan spent + # NOT compiling — the forwards themselves, the payload builds, + # the tempdirs, the decode. + "warm_execute_s": round(wall - compile_s, 3), + "warm_jobs": len(self.jobs), + # A warm plan whose jobs mostly compile NOTHING is a plan + # paying full forward cost for coverage it already has, which + # is a different defect from a slow compile. + "warm_jobs_compiling": compiled, + # OMITTED, not zeroed, when the wall is unmeasurable: a ratio + # over a zero denominator reported as 0.0 would read "this + # mint compiled nothing", which is the exact class of lie this + # module exists to delete. + **({"compile_fraction": round(compile_s / wall, 4)} + if wall > 0 else {}), + }, + "phases": {k: v for k, v in members.items() + if k != "dynamo_compile_s"}, + "overlays": overlays, + "jobs": list(self.jobs), + } + + +__all__ = [ + "JIT_OVERLAY_KEYS", "JIT_PARTITION_KEYS", "PHASE_ROUTER_DRAIN", + "TOTAL_KEY", "WARM_SPANS_V", "WarmLedger", "partition", +] diff --git a/tests/test_warm_spans_pgw989.py b/tests/test_warm_spans_pgw989.py new file mode 100644 index 00000000..9dae48a8 --- /dev/null +++ b/tests/test_warm_spans_pgw989.py @@ -0,0 +1,241 @@ +"""pgw#989: the dynamo mint's warm hour gets a breakdown, and the phase that +never measured a compile stops claiming to. + +The rows are written against the REAL seams — the real metric keys of the +pinned torch, the real ledger, the real parent emitter — so moving any of them +moves the test. +""" + +from __future__ import annotations + +import sys +import types +from typing import Any, Dict, List + +import pytest + +from gen_worker import warm_spans + + +# --------------------------------------------------------------------------- +# the partition +# --------------------------------------------------------------------------- + +#: One real ``torch.compile`` call's ``compilation_time_metrics`` delta, +#: MEASURED on the pin (torch 2.13.0+cu130, CPU inductor, 5.06 s wall). Frozen +#: here rather than re-measured: the point of the row is the ARITHMETIC over a +#: real key set, and a compile inside a unit test would price CI, not the code. +REAL_DELTA: Dict[str, float] = { + "compile_file": 8.638, + "_compile.compile_inner": 5.054, + "compile_attempt_0": 5.022, + "OutputGraph.call_user_compiler": 4.927, + "create_aot_dispatcher_function": 4.355, + "compile_fx..fw_compiler_base": 4.326, + "compile_fx_inner": 4.163, + "fx_codegen_and_compile": 4.163, + "GraphLowering.compile_to_fn": 4.091, + "GraphLowering.compile_to_module": 4.091, + "PyCodeCache.load_by_key_path": 3.203, + "async_compile.wait": 2.886, + "GraphLowering.codegen": 0.882, + "Scheduler.codegen": 0.828, + "CacheBase.get_system.triton_key": 0.424, + "_recursive_joint_graph_passes": 0.163, + "inductor_codecache_torch_key": 0.112, + "bytecode_tracing": 0.087, + "variable_builder_call": 0.062, + "Scheduler.__init__": 0.045, + "GraphLowering.run": 0.039, + "build_guards": 0.023, + "_recursive_post_grad_passes": 0.016, + "_recursive_pre_grad_passes": 0.004, +} + + +def test_partition_sums_to_the_dynamo_total_with_a_named_residual() -> None: + members, _overlays = warm_spans.partition(REAL_DELTA) + total = members["dynamo_compile_s"] + assert total == REAL_DELTA[warm_spans.TOTAL_KEY] + named = sum(v for k, v in members.items() if k != "dynamo_compile_s") + # A partition, not a sample: the members plus the residual ARE the total. + assert named == pytest.approx(total, abs=0.01) + assert members["compile_other_s"] > 0 + + +def test_the_aot_partition_would_have_missed_this_compile() -> None: + """The reason this module exists rather than reusing the AOT key set. + + ``AotCodeCompiler.compile`` never runs on the JIT path, so the AOT ledger + prices a JIT compile at a fifth of its cost and calls the rest residual. + """ + from gen_worker import aot_compile_spans + + aot_members, _o, _raw = aot_compile_spans.phase_delta({}, REAL_DELTA) + aot_total = sum(aot_members.values()) + jit_members, _ = warm_spans.partition(REAL_DELTA) + + assert aot_members["host_compile_s"] == 0.0 + assert aot_total < 0.3 * jit_members["dynamo_compile_s"] + # ...and the member that recovers it is the JIT kernel load. + assert jit_members["kernel_compile_s"] == pytest.approx(3.203) + + +def test_overlays_are_reported_and_never_summed_into_the_partition() -> None: + members, overlays = warm_spans.partition(REAL_DELTA) + assert overlays["async_wait_s"] == pytest.approx(2.886) + # It legitimately exceeds the wall — it prices the async workers' own CPU. + assert overlays["parallel_kernel_cpu_s"] > members["dynamo_compile_s"] + assert set(overlays) & set(members) == set() + + +# --------------------------------------------------------------------------- +# the ledger +# --------------------------------------------------------------------------- + +class _FakeDynamoUtils: + """Stands in for ``torch._dynamo.utils`` so a row can drive the ledger + without a compile. The ledger reads the module the production code reads; + only the counters are synthetic.""" + + def __init__(self) -> None: + self.compilation_time_metrics: Dict[str, List[float]] = {} + + def add(self, delta: Dict[str, float]) -> None: + for key, value in delta.items(): + self.compilation_time_metrics.setdefault(key, []).append(value) + + +@pytest.fixture() +def fake_dynamo(monkeypatch: pytest.MonkeyPatch) -> _FakeDynamoUtils: + utils = _FakeDynamoUtils() + torch = types.ModuleType("torch") + dynamo = types.ModuleType("torch._dynamo") + torch._dynamo = dynamo # type: ignore[attr-defined] + dynamo.utils = utils # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "torch", torch) + monkeypatch.setitem(sys.modules, "torch._dynamo", dynamo) + monkeypatch.setitem(sys.modules, "torch._dynamo.utils", utils) + return utils + + +def test_ledger_splits_compile_from_forward_per_job( + fake_dynamo: _FakeDynamoUtils, +) -> None: + ledger = warm_spans.WarmLedger() + + with ledger.job("generate/a"): + fake_dynamo.add(REAL_DELTA) + # A second job that hits the in-process cache: real wall, no compile. + with ledger.job("generate/b"): + pass + + table = ledger.table() + totals = table["totals"] + assert totals["warm_jobs"] == 2 + assert totals["warm_jobs_compiling"] == 1 + assert totals["warm_compile_s"] == pytest.approx(5.054) + # The residual is the forwards. It is never negative and never invented. + assert totals["warm_execute_s"] == pytest.approx( + totals["warm_wall_s"] - totals["warm_compile_s"], abs=0.01) + # The synthetic jobs take microseconds, so the ratio has no denominator — + # and it is OMITTED rather than reported as "compiled 0 %". + assert "compile_fraction" not in totals + + rows = {r["job"]: r for r in table["jobs"]} + assert rows["generate/a"]["compile_s"] == pytest.approx(5.054) + assert rows["generate/b"]["compile_s"] == 0.0 + + +def test_an_empty_plan_reports_absence_not_zero_cost() -> None: + table = warm_spans.WarmLedger().table() + assert table["totals"]["warm_wall_s"] == 0.0 + assert table["totals"]["warm_jobs"] == 0 + assert table["jobs"] == [] + + +def test_a_failing_warm_job_still_lands_in_the_ledger( + fake_dynamo: _FakeDynamoUtils, +) -> None: + """The job's exception propagates untouched and its seconds are kept — + a mint that died in the warm plan is exactly the one whose spent minutes + have to reach the hub (pgw#825's rule, applied to this ledger).""" + ledger = warm_spans.WarmLedger() + with pytest.raises(RuntimeError, match="boom"): + with ledger.job("generate/dies"): + fake_dynamo.add({"_compile.compile_inner": 2.0}) + raise RuntimeError("boom") + assert ledger.table()["totals"]["warm_jobs"] == 1 + assert ledger.jobs[0]["compile_s"] == pytest.approx(2.0) + + +# --------------------------------------------------------------------------- +# the phase rename, and the parent's emission +# --------------------------------------------------------------------------- + +def test_the_drain_phase_no_longer_calls_itself_a_compile() -> None: + """RED on master: ``mint_child`` framed ``_drain_router`` — which waits out + a queue the fleet mint arms EMPTY — as ``inductor_compile``, so every + dynamo mint's phase table carried ``'inductor_compile': 0.0`` beside a + ``warmup_forward`` row holding the entire compile.""" + from gen_worker import activity, mint_child + + source = __import__("inspect").getsource(mint_child) + assert '"inductor_compile"' not in source + assert warm_spans.PHASE_ROUTER_DRAIN == "router_drain" + # One definition, re-exported where the phase vocabulary lives. + assert activity.PHASE_ROUTER_DRAIN is warm_spans.PHASE_ROUTER_DRAIN + assert activity.PHASE_INDUCTOR_COMPILE != activity.PHASE_ROUTER_DRAIN + + +def test_parent_emits_the_warm_ledger_as_numeric_rows( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from gen_worker import mint_delegate + from gen_worker.mint_process import MintReport + + emitted: List[Dict[str, Any]] = [] + + def _emit(kind: str, detail: str, **kw: Any) -> None: + emitted.append({"kind": kind, "detail": detail, **kw}) + + monkeypatch.setattr(mint_delegate.activity_mod, "emit_event", _emit) + + ledger = warm_spans.WarmLedger() + ledger.jobs = [ + {"job": "generate", "wall_s": 300.0, "compile_s": 280.0, + "execute_s": 20.0}] + ledger._raw = dict(REAL_DELTA) + ledger._wall = 300.0 + + mint_delegate._emit_warm_ledger( + MintReport(status="minted", mint_phases=ledger.table()), head="h=1") + + phases = [row["phase"] for row in emitted] + assert "warm:totals" in phases + assert "warm:kernel_compile_s" in phases + assert "warm_job:generate" in phases + totals = next(r for r in emitted if r["phase"] == "warm:totals") + # A real wall means the ratio IS computable, and it rides the row. + assert "compile_fraction" in totals["detail"] + # The roll-up's NUMERIC column is the compile half — the number a + # regression hunt groups on, not one buried in prose. + assert totals["duration_ms"] == 5054 + assert all(r["kind"] == "jit_compile" for r in emitted) + + +def test_parent_is_silent_when_the_recipe_produced_no_ledger( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An AOT mint, or a child that died before the warm loop. Absence is + reported as absence — never as a zero-cost warm plan.""" + from gen_worker import mint_delegate + from gen_worker.mint_process import MintReport + + emitted: List[Any] = [] + monkeypatch.setattr( + mint_delegate.activity_mod, "emit_event", + lambda *a, **k: emitted.append(a)) + mint_delegate._emit_warm_ledger(MintReport(status="minted"), head="h=1") + mint_delegate._emit_warm_ledger(None, head="h=1") + assert emitted == []