From 1a6450cad242a7621e4b8775f01990b93645852c Mon Sep 17 00:00:00 2001 From: Jiong Gong Date: Wed, 23 Sep 2026 12:53:07 +0000 Subject: [PATCH 1/2] compass(design): anchor model_runner.py citations to symbols, and test them Every design-doc citation into atom/model_engine/model_runner.py is now written `model_runner.py::Class.method`. It no longer carries a line number, so an edit that moves lines leaves it true. The only numbers left are the capture census rows. Those are the measured conversion lines, so they stay. tests/compass/test_model_runner_citations.py reads the design docs and the source by path. It fails, naming the doc and the citation, in three cases: - a cited symbol that the source does not define; - a line number, in any spelling, with no symbol before it in its paragraph or table row; - a line number outside that symbol's ast span. Closes #288 Co-Authored-By: Claude Opus 5.5 (1M context) --- .../design/01_execution_and_time_model.md | 25 +-- .../02_model_runner_and_cost_backend.md | 23 ++- atom/compass/design/03_memory_and_kv_model.md | 51 +++--- .../design/04_model_capture_and_cost_ir.md | 6 +- .../design/07_calibration_toolchain.md | 4 +- atom/compass/design/12_open_items.md | 4 +- .../compass/design/14_speculative_decoding.md | 2 +- atom/compass/design/15_parallelism_support.md | 9 +- tests/compass/test_model_runner_citations.py | 148 ++++++++++++++++++ 9 files changed, 211 insertions(+), 61 deletions(-) create mode 100644 tests/compass/test_model_runner_citations.py diff --git a/atom/compass/design/01_execution_and_time_model.md b/atom/compass/design/01_execution_and_time_model.md index 616b6ad01f..8f51ab0bfb 100644 --- a/atom/compass/design/01_execution_and_time_model.md +++ b/atom/compass/design/01_execution_and_time_model.md @@ -244,8 +244,8 @@ Recorded as **T47**. - Simulation wall-clock cost is higher than Option A's. The >=5x target is stated as negotiable with a bottom line of "faster than real runs". Under saturation the prior design was 0.30x. This must be measured early, not assumed. -- `torch.cuda.set_device` (`model_runner.py:972`, in `_setup_device_and_distributed`) and - `torch.cuda.mem_get_info` (`model_runner.py:1676`, in `_read_device_memory`) are the two +- `torch.cuda.set_device` (in `model_runner.py::ModelRunner._setup_device_and_distributed`) and + `torch.cuda.mem_get_info` (in `model_runner.py::ModelRunner._read_device_memory`) are the two hard GPU dependencies a simulated runner must not inherit. --- @@ -1077,7 +1077,7 @@ The ABC is small and has **no `send_kv` / `recv_kv` verb** to fake Every connector's completion reaches the scheduler through **one** method: ``` -ModelRunner.async_proc_aggregation model_runner.py:3378-3399 +model_runner.py::ModelRunner.async_proc_aggregation -> EngineCore._poll_kv_transfer_progress engine_core.py:485-489 -> Scheduler._update_from_kv_xfer_finished scheduler.py:2989-3053 ``` @@ -1415,8 +1415,8 @@ re-derive. **The seam** - `Config.runner_qualname` — `atom/config.py:1595`; consumed `engine_core.py:129`, `async_proc.py:166-169` -- `ModelRunner.forward(batch: ScheduledBatch) -> ScheduledBatchOutput` — - `model_runner.py:3262-3350` +- `model_runner.py::ModelRunner.forward`, whose signature is + `forward(batch: ScheduledBatch) -> ScheduledBatchOutput` - the RPC boundary — `engine_core.py:386-388` - `ScheduledBatch` fields — `scheduler.py:579-820`; notably `detailed_sqsq` / `detailed_sqsk` / `detailed_sk` at `:790-792`, which are sum(N_Q^2), sum(N_Q * N_KV), @@ -1427,11 +1427,12 @@ re-derive. **Existing simulation-shaped hooks in ATOM** - `--load_dummy {empty,zero,xavier}` — `config.py:1556`, `arg_utils.py:260`, `loader.py:179-227,309-310`, `loading_core.py:266-291` -- meta-device model construction — `RapidServeModelRunner._init_weight_params_on_meta`, - `model_runner.py:4216-4239` -- a working non-allocating runner template — `RapidServeModelRunner` overrides at - `model_runner.py:4245,4259,4266,4272,4288,4296` -- `ModelRunner.dummy_execution()` — `model_runner.py:1191-1231`, shows how to hand-build +- meta-device model construction — + `model_runner.py::RapidServeModelRunner._init_weight_params_on_meta` +- a working non-allocating runner template — `model_runner.py::RapidServeModelRunner`, + which overrides `_build_and_load_model`, `_maybe_warmup`, `_kv_budget_extra_reserve`, + `get_num_blocks`, `allocate_kv_cache` and `forward` +- `model_runner.py::ModelRunner.dummy_execution` shows how to hand-build a `ScheduledBatch` - `ScheduledBatch.is_dummy_run` — `scheduler.py:589,781` - simulated TP (`--fake-eplb`) — `atom/distributed/simulated_tp.py`; explicit precedent @@ -1442,8 +1443,8 @@ re-derive. `tools/parse_trace.py` **Memory sizing (needed because it decides which configurations exist)** -- `ModelRunner.get_num_blocks()` — `model_runner.py:1686-1899`, with its four `torch.cuda` - reads in `_read_device_memory` (`model_runner.py:1666-1684`). Five device readings plus +- `model_runner.py::ModelRunner.get_num_blocks`, with its four `torch.cuda` + reads in `model_runner.py::ModelRunner._read_device_memory`. Five device readings plus arithmetic: `mem_get_info`, `allocated_bytes.all.peak`, `(total - free) - memory_reserved()`, `_estimate_cudagraph_overhead()`, a 2% safety margin, then `min(budget - ..., free)` and `plan_pools`. Consumed diff --git a/atom/compass/design/02_model_runner_and_cost_backend.md b/atom/compass/design/02_model_runner_and_cost_backend.md index 38725dfc5b..54aa7cafbd 100644 --- a/atom/compass/design/02_model_runner_and_cost_backend.md +++ b/atom/compass/design/02_model_runner_and_cost_backend.md @@ -26,8 +26,8 @@ enough that it does not become a maintenance burden against upstream ATOM. | Cut | Where | Replaces | Verdict | |---|---|---|---| -| `ModelRunner.run_model()` | `model_runner.py:2867-3087` | `model(input_ids, positions)` + `compute_logits` | **Rejected.** Returns `(logits, hidden_states)` as real tensors that `postprocess` indexes, samples and gathers logprobs from. Faking it still requires a `[bs, vocab]` device allocation. | -| `ModelRunner.forward()` | `model_runner.py:3262-3350` | the whole forward, sampling included | **Chosen.** Its contract is `ScheduledBatch` in, `ScheduledBatchOutput` out — both pure numpy/list/int, already pickled across the worker boundary. | +| `ModelRunner.run_model()` | `model_runner.py::ModelRunner.run_model` | `model(input_ids, positions)` + `compute_logits` | **Rejected.** Returns `(logits, hidden_states)` as real tensors that `postprocess` indexes, samples and gathers logprobs from. Faking it still requires a `[bs, vocab]` device allocation. | +| `ModelRunner.forward()` | `model_runner.py::ModelRunner.forward` | the whole forward, sampling included | **Chosen.** Its contract is `ScheduledBatch` in, `ScheduledBatchOutput` out — both pure numpy/list/int, already pickled across the worker boundary. | | whole-runner replacement via `Config.runner_qualname` | `config.py:1595` | weights, KV tensors, CUDA graphs, sampling | **Chosen as the delivery mechanism** for the above. | ### Decision @@ -40,12 +40,11 @@ and `async_proc.py:166-169`. It already has two in-tree users — swaps in `RapidServeModelRunner` automatically. **The injection itself requires no ATOM change.** -`RapidServeModelRunner` (`model_runner.py:4195-4710`) is a working template for a +`model_runner.py::RapidServeModelRunner` is a working template for a non-allocating runner already in the tree. It overrides exactly the memory-owning -methods: `_build_and_load_model` (`:4245`), `_maybe_warmup` (`:4259`), -`_kv_budget_extra_reserve` (`:4266`), `get_num_blocks` (`:4272`), -`allocate_kv_cache` (`:4288`), `forward` (`:4296`) — and constructs parameters on meta -through `_init_weight_params_on_meta` (`:4216-4239`). +methods: `_build_and_load_model`, `_maybe_warmup`, `_kv_budget_extra_reserve`, +`get_num_blocks`, `allocate_kv_cache` and `forward` — and constructs parameters on meta +through `model_runner.py::RapidServeModelRunner._init_weight_params_on_meta`. ### The RPC surface that must be honoured @@ -82,8 +81,8 @@ third case — a method that raises — and the table of which names wait. step for step (mean 8.99 s vs a real 9.00 s). 3. **`produces_output()`** (`scheduler.py:823-840`): a pure-middle-chunk prefill batch must return an **empty** `token_ids` list with the same `req_ids`, mirroring - the early `return ScheduledBatchOutput(...)` in `ModelRunner.forward` - (`model_runner.py:3329-3335`). + the early `return ScheduledBatchOutput(...)` in + `model_runner.py::ModelRunner.forward`. Speculative decoding adds `num_rejected` / `num_bonus` sized `batch.total_seqs_num` and `draft_token_ids` shaped `[bs, mtp_k]`. ATOM's existing `synthetic_acceptance_rates` path @@ -118,9 +117,9 @@ missing is a statement of which combination Compass uses, and when. | Piece | Where | What it does | |---|---|---| | `--load_dummy {empty,zero,xavier}` | `config.py:1556`, `arg_utils.py:69,260`, `loader.py:126,234,289-307` | skips the checkpoint read; `empty` leaves params uninitialised, the others fill them with finite values in place | -| `_init_weight_params_on_meta` | `model_runner.py:4216-4239` | wraps `Module.register_parameter` so every `nn.Parameter` is replaced by a meta tensor as it is registered | +| `_init_weight_params_on_meta` | `model_runner.py::RapidServeModelRunner._init_weight_params_on_meta` | wraps `Module.register_parameter` so every `nn.Parameter` is replaced by a meta tensor as it is registered | | `no_init_weights` | `models/utils.py:457-496` | uses `torch.device("meta")` as a **context manager**, so construction itself lands on meta - no transient, no GPU, and it covers buffers. **Currently unused in ATOM.** | -| `RapidServeModelRunner._build_and_load_model` | `model_runner.py:4245` | the override point where a runner declines to load | +| `RapidServeModelRunner._build_and_load_model` | `model_runner.py::RapidServeModelRunner._build_and_load_model` | the override point where a runner declines to load | #### Why `_init_weight_params_on_meta` allocates on the real device, despite its name @@ -146,7 +145,7 @@ CUDA IPC and recomputes RoPE caches locally. **So it is not a bug.** It is a deliberate trade — a one-parameter transient bought in exchange for real buffers and unchanged init branches — and for its use case the trade is clearly right. It avoids what the call site calls *"the transient 2x-weights peak that -OOMs at TP=4"* (`model_runner.py:4250-4254`), which is the thing that mattered there. +OOMs at TP=4"* (in `model_runner.py::RapidServeModelRunner._build_and_load_model`), which is the thing that mattered there. Worth noting for anyone reading that code: the call site says construction *"allocates zero GPU bytes"* while the helper says one parameter is transiently real. Both are true of diff --git a/atom/compass/design/03_memory_and_kv_model.md b/atom/compass/design/03_memory_and_kv_model.md index d42c11bcdc..b45526c1d2 100644 --- a/atom/compass/design/03_memory_and_kv_model.md +++ b/atom/compass/design/03_memory_and_kv_model.md @@ -38,7 +38,7 @@ with exactly one point of GPU contact. | Prefix-cache hash (`compute_hash`, xxhash xxh64 chained with the parent) | `block_manager.py:233-245` | **No** | | Prefix-cache scan and claim (`can_allocate`, `allocate`) | `block_manager.py:469-561`, `:563-597` | **No** | | Prefix publish, deferred until the forward computed the KV (`hash_blocks`) | `block_manager.py:696-771` | **No** | -| **`allocate_kv_cache`** — creating the actual tensors | **`model_runner.py:1901-2105`** | **Yes** | +| **`allocate_kv_cache`** — creating the actual tensors | **`model_runner.py::ModelRunner.allocate_kv_cache`** | **Yes** | `BlockManager.__init__` asserts `num_blocks > 0` (`block_manager.py:77`) and nothing else about the device. @@ -50,7 +50,8 @@ about the device. Two changes only: 1. `allocate_kv_cache` becomes a no-op in the simulated runner (`RapidServeModelRunner` - already does exactly this at `model_runner.py:4288-4293`), so no tensor is created. + already does exactly this in `model_runner.py::RapidServeModelRunner.allocate_kv_cache`), + so no tensor is created. 2. `get_num_blocks` returns a block count produced from a device model rather than from device readings. See D14. @@ -126,10 +127,10 @@ probed per request. Whether that is 0.1 ms or 10 ms is unmeasured. Recorded as * ### Open issues -- `allocate_kv_cache` also registers tensors globally via `set_kv_cache_data` - (`model_runner.py:2060-2065`) and cross-validates expected against actual bytes - (`:2067-2094`). The no-op must keep whatever downstream code reads from that registry - satisfied, or supply a descriptor-shaped stand-in. +- `model_runner.py::ModelRunner.allocate_kv_cache` also registers tensors globally via + `set_kv_cache_data` and cross-validates expected against actual bytes. The no-op must + keep whatever downstream code reads from that registry satisfied, or supply a + descriptor-shaped stand-in. - `BlockManager.hash_block_size = block_size * dcp_world_size` (`block_manager.py:93`) — decode context parallelism changes the hash granularity. Out of scope now; noted so it is not discovered later. @@ -140,28 +141,28 @@ probed per request. Whether that is 0.1 ms or 10 ms is unmeasured. Recorded as * ### Problem -`ModelRunner.get_num_blocks()` (`model_runner.py:1686-1899`) is **five device readings +`model_runner.py::ModelRunner.get_num_blocks` is **five device readings plus arithmetic**: ``` -# in _read_device_memory (:1666-1684), which get_num_blocks calls at :1693 -free, total = torch.cuda.mem_get_info() # :1676 -peak_torch = max(allocated_bytes.all.peak, .all.current) # :1677-1680 -non_torch = max((total - free) - torch.cuda.memory_reserved(), 0) # :1683 -# in get_num_blocks (:1686-1899) -cudagraph_overhead = self._estimate_cudagraph_overhead() # :1695 -safety_margin = int(total * 0.02) # :1696 -budget = int(total * config.gpu_memory_utilization) # :1698 +# in _read_device_memory, which get_num_blocks calls first +free, total = torch.cuda.mem_get_info() +peak_torch = max(allocated_bytes.all.peak, .all.current) +non_torch = max((total - free) - torch.cuda.memory_reserved(), 0) +# in get_num_blocks +cudagraph_overhead = self._estimate_cudagraph_overhead() +safety_margin = int(total * 0.02) +budget = int(total * config.gpu_memory_utilization) available_for_kv = min(budget - (peak_torch + non_torch + cudagraph_overhead + safety_margin) - - self._kv_budget_extra_reserve(total), free) # :1699-1706 + - self._kv_budget_extra_reserve(total), free) plan = plan_pools(self._sub_pool_specs(), available_for_kv, - config.max_num_seqs) # :1721 -num_kvcache_blocks = plan.paged_entries # :1764 -# under PP: all_reduce MIN across stages # :1765-1771 + config.max_num_seqs) +num_kvcache_blocks = plan.paged_entries +# under PP: all_reduce MIN across stages ``` -Line numbers are `model_runner.py` at `feature/atomcompass_new` `75a265a3f`. +`_read_device_memory` is `model_runner.py::ModelRunner._read_device_memory`. Consumed at `engine_core.py:132-145`, which sets `config.num_kvcache_blocks` before the `Scheduler` and `BlockManager` are constructed at `:170`. @@ -225,15 +226,15 @@ is the right thing to model and it is stated here so it is not discovered as a g ### The pipeline minimum is not inert An earlier note here said that under PP the `all_reduce(MIN)` across stages -(`model_runner.py:1764-1771`) is inert, *because every stage computes the same number*, -and that it *still needs a live process group or a stub*. Both halves are wrong — and the -engine's own comment on that reduce has said so all along (`model_runner.py:1759-1760`): +(in `model_runner.py::ModelRunner.get_num_blocks`) is inert, *because every stage computes +the same number*, and that it *still needs a live process group or a stub*. Both halves are +wrong — and the engine's own comment on that reduce has said so all along: *"PP stages compute different block counts; block ids must be valid on every stage's KV tensor, so reduce to the global minimum."* The code said what this document denied. **The stages do not compute the same number.** The five readings *are* identical across stages — nothing in the memory model varies with pipeline rank. The layer count is not: -`_get_total_num_layers` (`model_runner.py:1515-1538`) takes a `get_pp_indices` slice, so +`model_runner.py::ModelRunner._get_total_num_layers` takes a `get_pp_indices` slice, so each stage sizes its pool from the layers it actually holds. Re-derived from ATOM's own partitioner at `feature/atomcompass_new` `92f1fdafe`, over the 64-layer hybrid vendored at `tests/compass/qwen3_5_27b_config.json` — one full-attention layer in four, so 16 of the 64 @@ -268,7 +269,7 @@ holds the most paged layers — which is not in general the last one: at pp = 5 3 of 5, and at pp = 6 it is four stages of the six. **Neither a live process group nor a stub is needed.** The reduce is already guarded by -`torch.distributed.is_initialized()` (`model_runner.py:1765`): with no group it does not +`torch.distributed.is_initialized()` (in `model_runner.py::ModelRunner.get_num_blocks`): with no group it does not run, and with one it runs ATOM's own code unchanged. Building a stub for it is building something nothing asks for. diff --git a/atom/compass/design/04_model_capture_and_cost_ir.md b/atom/compass/design/04_model_capture_and_cost_ir.md index 6f8a56281a..d3d74990e0 100644 --- a/atom/compass/design/04_model_capture_and_cost_ir.md +++ b/atom/compass/design/04_model_capture_and_cost_ir.md @@ -198,9 +198,9 @@ fills and slices, nothing else: | where | lines | conversions | |---|---|---| | `aiter_attention.py in prepare_decode` | 1106, 1115, 1121, 1122, 1123, 1131, 1132 | 10 | -| `model_runner.py in prepare_inputs` | 2468, 2479, 2481 | 4 | -| `model_runner.py in prepare_input_ids` | 510, 513 | 2 | -| `model_runner.py in prepare_sample` | 2564 | 1 | +| `model_runner.py::ModelRunner.prepare_inputs` | 2468, 2479, 2481 | 4 | +| `model_runner.py::tokenIDProcessor.prepare_input_ids` | 510, 513 | 2 | +| `model_runner.py::ModelRunner.prepare_sample` | 2564 | 1 | | `backends.py in _mrope_cpu_view` | 398, 400 | 2 | | `gdn_attn.py in _attach_gdn_decode_metadata` | 1237 | 1 | diff --git a/atom/compass/design/07_calibration_toolchain.md b/atom/compass/design/07_calibration_toolchain.md index 0bb5e263df..ac1d9d6990 100644 --- a/atom/compass/design/07_calibration_toolchain.md +++ b/atom/compass/design/07_calibration_toolchain.md @@ -311,8 +311,8 @@ forward context only the runner establishes, and building it by hand means reimp The constraints point at the *runner*, not the server — a single-process `ModelRunner` **is** the worker, so Phase 1b's placement requirement is satisfied. ATOM already shows how to -drive one without a scheduler: `dummy_execution` (`model_runner.py:1191-1231`) and -`warmup_model` (`:1233-1298`) both fabricate `ScheduledBatch`es by hand. +drive one without a scheduler: `model_runner.py::ModelRunner.dummy_execution` and +`model_runner.py::ModelRunner.warmup_model` both fabricate `ScheduledBatch`es by hand. ``` Phase 0 (device-free) standalone bench (GPU, one process) diff --git a/atom/compass/design/12_open_items.md b/atom/compass/design/12_open_items.md index bb4c449307..44adcdcca0 100644 --- a/atom/compass/design/12_open_items.md +++ b/atom/compass/design/12_open_items.md @@ -97,7 +97,7 @@ M-f `14`; M-g `01` D3.5. | T9 | Declare a row-ordering treatment for decode attention | | **T51** | Enumerate the layer-pattern shapes for Qwen3.8-27B and Kimi-K3; confirm the nested-`Repeat` detector reaches the hierarchical form on both | | **T52** | Root-cause the `TorchDispatchMode` 8-rank hang at `dspark_scheduler.py:264` — gates T5 | -| **T81** | Make the D18 capture *symbolic* on ATOM's real forward, or record that it cannot be. **Done for a decode step, at TP1 and TP2, on the published 27B; not attempted for prefill.** A step whose *width* is a free symbol traces through ATOM's own `ModelRunner`, model classes and staging with `ShapeEnv.replacements` **empty** and no symbol solved anywhere. **The census, by operator family** — a total is what made two earlier attempts look like progress while every symbol sat on a view. At TP1, **5,016 of 12,544** shape entries are not plain integers: **514 on 257 `aiter.gemm_a16w16`**, **304 on the 64 attention calls** (`unified_attention_with_output_base`, `linear_attention_with_output_base`) and **802 on 273 normalisation operators** (`_fused_qk_rmsnorm_group_quant_kernel`, `mean.dim`, `pow`, `rsqrt`). At TP2, **5,297 of 13,107**, the same three families plus **266 on the 133 collectives**. The control is the same file with the width left an integer: 0 of 12,544 and 0 of 13,107. Both passes record **the same operators** — 2,521 at TP1 and 2,662 at TP2 — differing in one call that materialises a symbol where it used to lift a constant, so what changed is what the shape entries say and not which operators ran. The breakdown is only worth the bucket that fails: an operator matching no declared family lands in `unclassified` and the test refuses it, and the classifier matches **whole operator names** — it matched by prefix until 2026-09-22, which silently folded `addmm` and `addbmm` into elementwise and `slice_scatter` into views, so the bucket could not have fired for the cases it exists for. No operator on this step is one of those, so the rows above are unchanged by the repair. **The mechanism: the three sites are symptoms of one line of torch.** `SymInt.__index__` and `SymInt.__int__` are `guard_int`, so every host-side use of the width asks for a number, gets the hint, and records the ask as `Eq(s, hint)`. **16 ATOM lines** convert the width during one traced forward — 20 conversions, all host fills and slices in `prepare_decode` (7 lines), `prepare_inputs` (3), `prepare_input_ids` (2), `prepare_sample`, `_mrope_cpu_view` (2) and `gdn_attn` — which is why closing sites one at a time relocates rather than converges. The capture keeps the conversion, to the hint, which is the count ATOM computed, and replaces the *recording* with its own log of every conversion and the line it happened on. **That log is asserted as a multiset against a declared set** (`EXPECTED_HOST_RESOLUTIONS`), identical at TP1 and TP2 and at both step widths, so a conversion at a seventeenth line fails rather than joining the log unremarked; `__bool__` is untouched, so a branch on a width still guards. **One production line changed**, and only one: `_rows` returns `t.shape[0]` rather than `int(t.shape[0])`, which is the identity wherever the size is real. That is site three, and the only one of the three not reachable from a capture-time substitution. **"A symbolic `CpuGpuBuffer`" is not required, measured.** Site two's `-> 16384` is a buffer *capacity* (`max_model_len // block_size`) being solved against the CPU side's constant, and it exists only because that probe symbolises every dimension of the staged device tensor. Capacities are engine configuration, not a function of the step. With only the width symbolic the two slices carry the same symbol and `copy_to_gpu` dispatches unchanged; building the CPU side fake instead costs **38 `prim.device` calls ATOM never makes** and leaves **every family's non-numeric shape-entry count unchanged** — 5,016 either way, family for family — while the *total* entries rise **12,544 → 12,585**, because those 38 calls carry entries of their own. That ablation's code is **not in the tree**: it is a fake, non-numpy-backed CPU side written into the capture's staged allocators, measured, and removed once it had answered, so the row is a measurement whose artifact has to be rebuilt to re-check. **Site three is also not what it was.** With `_rows` no longer converting, the same probe still reaches `forward_context.py:444 in assert_shape_contract` — another file and a later phase of the step, reached from the `run_model` call at `model_runner.py:3281` with `prepare_inputs` already returned, not fourteen lines after `aiter_attention.py:1115` — because that probe hands the caller a bound that is a *different symbol* from the staged buffer's rows, and ATOM's contract is that the two are one number. Equating them is the contract working. It does not fire when the whole step is one symbol, which is why it is one: a decode step has one query row per sequence, so its token count and its sequence count are the same number for a structural reason. **The symbol is free, shown and not assumed**: the step is traced at two widths, 2 and 8, at **both TP1 and TP2**, and the inventories are identical once the symbol's name is set aside — same operators, same 12,544 entries, same 5,016 non-numeric, same digest, same concrete dimensions by value. That is also what says every dimension that stayed a number is not a width in disguise; there are 23 distinct concrete values at TP1 and every one is the model's or the engine's. Three guards at TP1 and two at TP2, all staging capacities read against the width, none an equality — **and the guard set is hint-dependent**: at every step width but 2 a fourth guard ` + 1 > ` appears, at both TP1 and TP2. It is not a specialisation but a `__bool__` comparison on a live symbol, so it is the positive evidence that `__bool__` was left alive; what it costs is that the applicability carries a *lower* bound at any hint but 2, and that the two-width comparison has one artifact that is not identical between widths. **What the symbol does not survive**: `ScheduledBatch.__init__` compares the staged token array's length against the count, which solves the width by comparison rather than conversion, so the capture builds the batch at the concrete width and rebinds its four count fields afterwards. The symbol enters ATOM's staging, not ATOM's batch constructor. **Still open**: prefill and chunked prefill, speculative decode, any structure where the token count is not the sequence count, and the constructor above — this closes the decode step only. Pinned by `tests/compass/test_capture_real_model.py`, whose four passes are the concrete control, the two probes that record the three sites, and the step-symbol pass that specialises nowhere; the earlier assertions were re-pointed rather than removed. Gates T5 alongside T52. | +| **T81** | Make the D18 capture *symbolic* on ATOM's real forward, or record that it cannot be. **Done for a decode step, at TP1 and TP2, on the published 27B; not attempted for prefill.** A step whose *width* is a free symbol traces through ATOM's own `ModelRunner`, model classes and staging with `ShapeEnv.replacements` **empty** and no symbol solved anywhere. **The census, by operator family** — a total is what made two earlier attempts look like progress while every symbol sat on a view. At TP1, **5,016 of 12,544** shape entries are not plain integers: **514 on 257 `aiter.gemm_a16w16`**, **304 on the 64 attention calls** (`unified_attention_with_output_base`, `linear_attention_with_output_base`) and **802 on 273 normalisation operators** (`_fused_qk_rmsnorm_group_quant_kernel`, `mean.dim`, `pow`, `rsqrt`). At TP2, **5,297 of 13,107**, the same three families plus **266 on the 133 collectives**. The control is the same file with the width left an integer: 0 of 12,544 and 0 of 13,107. Both passes record **the same operators** — 2,521 at TP1 and 2,662 at TP2 — differing in one call that materialises a symbol where it used to lift a constant, so what changed is what the shape entries say and not which operators ran. The breakdown is only worth the bucket that fails: an operator matching no declared family lands in `unclassified` and the test refuses it, and the classifier matches **whole operator names** — it matched by prefix until 2026-09-22, which silently folded `addmm` and `addbmm` into elementwise and `slice_scatter` into views, so the bucket could not have fired for the cases it exists for. No operator on this step is one of those, so the rows above are unchanged by the repair. **The mechanism: the three sites are symptoms of one line of torch.** `SymInt.__index__` and `SymInt.__int__` are `guard_int`, so every host-side use of the width asks for a number, gets the hint, and records the ask as `Eq(s, hint)`. **16 ATOM lines** convert the width during one traced forward — 20 conversions, all host fills and slices in `prepare_decode` (7 lines), `prepare_inputs` (3), `prepare_input_ids` (2), `prepare_sample`, `_mrope_cpu_view` (2) and `gdn_attn` — which is why closing sites one at a time relocates rather than converges. The capture keeps the conversion, to the hint, which is the count ATOM computed, and replaces the *recording* with its own log of every conversion and the line it happened on. **That log is asserted as a multiset against a declared set** (`EXPECTED_HOST_RESOLUTIONS`), identical at TP1 and TP2 and at both step widths, so a conversion at a seventeenth line fails rather than joining the log unremarked; `__bool__` is untouched, so a branch on a width still guards. **One production line changed**, and only one: `_rows` returns `t.shape[0]` rather than `int(t.shape[0])`, which is the identity wherever the size is real. That is site three, and the only one of the three not reachable from a capture-time substitution. **"A symbolic `CpuGpuBuffer`" is not required, measured.** Site two's `-> 16384` is a buffer *capacity* (`max_model_len // block_size`) being solved against the CPU side's constant, and it exists only because that probe symbolises every dimension of the staged device tensor. Capacities are engine configuration, not a function of the step. With only the width symbolic the two slices carry the same symbol and `copy_to_gpu` dispatches unchanged; building the CPU side fake instead costs **38 `prim.device` calls ATOM never makes** and leaves **every family's non-numeric shape-entry count unchanged** — 5,016 either way, family for family — while the *total* entries rise **12,544 → 12,585**, because those 38 calls carry entries of their own. That ablation's code is **not in the tree**: it is a fake, non-numpy-backed CPU side written into the capture's staged allocators, measured, and removed once it had answered, so the row is a measurement whose artifact has to be rebuilt to re-check. **Site three is also not what it was.** With `_rows` no longer converting, the same probe still reaches `forward_context.py:444 in assert_shape_contract` — another file and a later phase of the step, reached from the `run_model` call in `model_runner.py::ModelRunner.forward` with `prepare_inputs` already returned, not fourteen lines after `aiter_attention.py:1115` — because that probe hands the caller a bound that is a *different symbol* from the staged buffer's rows, and ATOM's contract is that the two are one number. Equating them is the contract working. It does not fire when the whole step is one symbol, which is why it is one: a decode step has one query row per sequence, so its token count and its sequence count are the same number for a structural reason. **The symbol is free, shown and not assumed**: the step is traced at two widths, 2 and 8, at **both TP1 and TP2**, and the inventories are identical once the symbol's name is set aside — same operators, same 12,544 entries, same 5,016 non-numeric, same digest, same concrete dimensions by value. That is also what says every dimension that stayed a number is not a width in disguise; there are 23 distinct concrete values at TP1 and every one is the model's or the engine's. Three guards at TP1 and two at TP2, all staging capacities read against the width, none an equality — **and the guard set is hint-dependent**: at every step width but 2 a fourth guard ` + 1 > ` appears, at both TP1 and TP2. It is not a specialisation but a `__bool__` comparison on a live symbol, so it is the positive evidence that `__bool__` was left alive; what it costs is that the applicability carries a *lower* bound at any hint but 2, and that the two-width comparison has one artifact that is not identical between widths. **What the symbol does not survive**: `ScheduledBatch.__init__` compares the staged token array's length against the count, which solves the width by comparison rather than conversion, so the capture builds the batch at the concrete width and rebinds its four count fields afterwards. The symbol enters ATOM's staging, not ATOM's batch constructor. **Still open**: prefill and chunked prefill, speculative decode, any structure where the token count is not the sequence count, and the constructor above — this closes the decode step only. Pinned by `tests/compass/test_capture_real_model.py`, whose four passes are the concrete control, the two probes that record the three sites, and the step-symbol pass that specialises nowhere; the earlier assertions were re-pointed rather than removed. Gates T5 alongside T52. | ### Topic 06 — workload harness contract @@ -222,7 +222,7 @@ M-f `14`; M-g `01` D3.5. | T66 | Measure whether the Class-C runtime constants move with PP degree | | T67 | Measure the step-duration spread across DP ranks, and what padding to `unified_bs` costs | | T78 | **`atom/models/qwen3_5.py:427` and `atom/models/glm4_moe.py:426` block PP at compilation level >= 2.** Opened 2026-09-20 by P0.5. Both declare `"intermediate_tensors": 0` in `dynamic_arg_dims`. `atom/utils/decorators.py:525` raises `ValueError("Unsupported dynamic dimensions ...")` for any non-Tensor argument, and that raise sits **before** the decorator's own `IntermediateTensors` handling at `:533-537`, which marks the token dim of each contained tensor. A PP non-first stage receives its activations inside an `IntermediateTensors` container, so the declaration can never be satisfied. Levels 0 and 1 escape it: `decorators.py:485-488` sets `do_not_compile` for `NO_COMPILATION` and `DYNAMO_AS_IS` - 0 and 1 per `atom/config.py:113-116` - and `:505` returns before the raise is reached. Observed at the default level 3: `-pp 2 --enforce-eager` on Qwen3.8-27B raised `ValueError: Unsupported dynamic dimensions [0] for argument intermediate_tensors with type ` (node 18, container `xiaobizh_n18`, `/workspace/p05-t64/pp2_27b.log:126`, 2026-09-20); the same invocation with `--level 0` reached KV sizing instead. **Two instances, one class:** `grep -rn '"intermediate_tensors": 0' atom/models/` returns exactly these two files, so an upstream fix scoped to the M2/M3 target model alone leaves `glm4_moe.py` broken under PP. Other PP-capable models use a bare `@support_torch_compile` and reach the working path. **Not on M1's path**: `15` D94's PP2 test runs a fake model - `tests/test_pp.py` mocks `aiter.dist.parallel_state` (`:126-128`) and builds a `MockConfig` (`:463-468`), loads no model file, and is in the CPU tier. It blocks any **real-model** PP measurement, so it is a prerequisite for M7 cost accuracy and for T66. Upstream ATOM fix, not Compass; owner unassigned. | -| T79 | **`atom/model_ops/attentions/gdn_attn.py:1329-1331` mixes a PP-local layer count with a global one.** Opened 2026-09-20 by P0.5. `total = runner._get_total_num_layers()` returns the **PP-local** slice under PP>1 (`atom/model_engine/model_runner.py:1523-1529`, via `get_pp_indices`), but line 1330's `num_draft = total - hf_config.num_hidden_layers` subtracts the **global** count, and `runner.num_full_attn` (`gdn_attn.py:159-161`) is global as well. Measured on Qwen3.8-27B (64 layers, `full_attention_interval` 4, `num_full_attn` 16) on node 18, container `xiaobizh_n18`, 2026-09-20, with `python -m atom.examples.simple_inference --model -pp {1,2} --enforce-eager --level 0 --max-num-batched-tokens 512`: PP1 gives `num_draft=0`, `n_full=16`, logs `sub-pool kv: entries=80746, entry_bytes=1056768` and exits `EXIT=0` (`/workspace/p05-t64/ctrl_pp1.log:43` and `:410`); PP2 gives `num_draft=-32`, `n_full=-16`, logs `entries=0, entry_bytes=-1056768, num_kvcache_blocks=0` and exits `EXIT=1` (`/workspace/p05-t64/pp2_27b_l0.log:77-80` and `:160`). An exact sign flip at an identical magnitude, so the mechanism is determined rather than inferred. The two invocations also differ in the visible device set (1 vs 1,2) and in `--max-tokens` (8 vs 64); neither enters the `entry_bytes` computation, which is fixed at startup by the layer counts above. Same M1/M7 split as T78: the fake-model path never reaches this code. Upstream ATOM fix. | +| T79 | **`atom/model_ops/attentions/gdn_attn.py:1329-1331` mixes a PP-local layer count with a global one.** Opened 2026-09-20 by P0.5. `total = runner._get_total_num_layers()` returns the **PP-local** slice under PP>1 (`atom/model_engine/model_runner.py::ModelRunner._get_total_num_layers`, via `get_pp_indices`), but line 1330's `num_draft = total - hf_config.num_hidden_layers` subtracts the **global** count, and `runner.num_full_attn` (`gdn_attn.py:159-161`) is global as well. Measured on Qwen3.8-27B (64 layers, `full_attention_interval` 4, `num_full_attn` 16) on node 18, container `xiaobizh_n18`, 2026-09-20, with `python -m atom.examples.simple_inference --model -pp {1,2} --enforce-eager --level 0 --max-num-batched-tokens 512`: PP1 gives `num_draft=0`, `n_full=16`, logs `sub-pool kv: entries=80746, entry_bytes=1056768` and exits `EXIT=0` (`/workspace/p05-t64/ctrl_pp1.log:43` and `:410`); PP2 gives `num_draft=-32`, `n_full=-16`, logs `entries=0, entry_bytes=-1056768, num_kvcache_blocks=0` and exits `EXIT=1` (`/workspace/p05-t64/pp2_27b_l0.log:77-80` and `:160`). An exact sign flip at an identical magnitude, so the mechanism is determined rather than inferred. The two invocations also differ in the visible device set (1 vs 1,2) and in `--max-tokens` (8 vs 64); neither enters the `entry_bytes` computation, which is fixed at startup by the layer counts above. Same M1/M7 split as T78: the fake-model path never reaches this code. Upstream ATOM fix. | | T82 | **`15` D91 Q2's "no scheduling coupling" is contradicted by the head's decode admission loop.** Opened 2026-09-21 by P0.5. `atom/model_engine/scheduler.py:1761-1763` skips any running sequence whose id is in `_pp_inflight_token_block` while composing a decode batch, and `mark_pp_inflight` (`:2309-2316`) puts exactly the sequences whose token is in flight into that set - its docstring at `:2310-2313` reads "Head: block re-scheduling of seqs whose token is now in flight." So the first stage's batch composition at step N is a function of pipeline in-flight state, which is the coupling D91 Q2 says PP does not add. The mechanism is T64's own finding: the head keeps up to `pp_size` independently scheduled batches in flight, and the block is what stops a sequence decoding against a stale token. The item is the amendment to D91 Q2 and whatever follows for the PP LP model; this row registers it and does not amend the decision. | ### Topics 01, 03, 05 — newly opened diff --git a/atom/compass/design/14_speculative_decoding.md b/atom/compass/design/14_speculative_decoding.md index d852df60b4..c1fa4d1ae3 100644 --- a/atom/compass/design/14_speculative_decoding.md +++ b/atom/compass/design/14_speculative_decoding.md @@ -272,7 +272,7 @@ Draft models add: | Term | Source | |---|---| | **draft weights** | Class A, exact from the draft's own HF config. `_MTP_CONFIG` (`config.py:1090-1096`) names the architecture and the `n_predict` attribute per model type. | -| **draft KV layers** | `ModelRunner._num_draft_kv_layers()` (`model_runner.py:1491-1513`) — **the single source of truth, and it is already called out as such in the code.** A draft with a real layer stack (Eagle3, standalone DSpark) needs one slot per layer; serial MTP declares `num_nextn_predict_layers`. | +| **draft KV layers** | `model_runner.py::ModelRunner._num_draft_kv_layers` — **the single source of truth, and it is already called out as such in the code.** A draft with a real layer stack (Eagle3, standalone DSpark) needs one slot per layer; serial MTP declares `num_nextn_predict_layers`. | Two properties inherited rather than designed, both consequences of `03` D13's decision to run ATOM's real block accounting: diff --git a/atom/compass/design/15_parallelism_support.md b/atom/compass/design/15_parallelism_support.md index 2b1d5f11e5..42c0a50f44 100644 --- a/atom/compass/design/15_parallelism_support.md +++ b/atom/compass/design/15_parallelism_support.md @@ -327,10 +327,11 @@ answered — **T86**. every use site imports `get_ep_group` from `aiter.dist.parallel_state` (`moe.py:599`, `fused_moe/mori_v2_prepare_finalize.py:153,636`, `fused_moe/flydsl_mega_experts.py:186`, `eplb.py:1768`, `models/glm4_moe.py:96`, `models/qwen3_next.py:171`, -`model_runner.py:3580`). Both of ATOM's distributed-init paths — +`model_runner.py::ModelRunner._force_aiter_unreg_capture_for_piecewise`). Both of ATOM's distributed-init paths — `init_pp_aware_dist_env` (`distributed/pp_comm.py:46`) when `pp_size > 1`, aiter's -`init_dist_env` (`aiter/ops/communication.py:22`) otherwise, chosen in -`_setup_device_and_distributed` (`model_runner.py:946`, the branch at `:981`) — end in +`init_dist_env` (`aiter/ops/communication.py:22`) otherwise, chosen by the +`pipeline_parallel_size > 1` branch of +`model_runner.py::ModelRunner._setup_device_and_distributed` — end in aiter's `initialize_model_parallel`, which builds the group at `aiter/dist/parallel_state.py:1926-1945` out of @@ -369,7 +370,7 @@ Under DP-attention `CoreManager` rewrites `dp := dp × tp, tp := 1` before any o `dp × pcp × tp` either way: `init_dist_env` passes `world_size = pp × tp × pcp` with `pp` pinned to 1 (`aiter/ops/communication.py:33-40`) and `init_distributed_environment` multiplies DP back in (`parallel_state.py:1726-1729`); the PP branch computes the same -index itself at `model_runner.py:986-988`. +index itself in `model_runner.py::ModelRunner._setup_device_and_distributed`. **Ranks are contiguous whenever `pp == 1`**, which is every reachable EP configuration. The strided last row is the only non-contiguous case, and ATOM refuses it diff --git a/tests/compass/test_model_runner_citations.py b/tests/compass/test_model_runner_citations.py new file mode 100644 index 0000000000..128379a7c2 --- /dev/null +++ b/tests/compass/test_model_runner_citations.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: MIT +"""Every `model_runner.py` citation in the design documents names a symbol. + +A citation is written `model_runner.py::Class.method`, and the symbol must be +defined in `atom/model_engine/model_runner.py` at this checkout. Any line number +cited into that file must fall inside the `ast` span (`def` to `end_lineno`) of +the nearest symbol cited before it in the same paragraph or table row. A line +number with no symbol before it fails. The line-number spellings recognised are +`model_runner.py:N`, `` `.../model_runner.py`:N ``, `model_runner.py` line N, +a bare `` `:N` ``, `(:N` or `# :N` whose nearest preceding file name is +`model_runner.py` or which follows a backticked name defined there, and a table +cell of line numbers right after a symbol's cell. + +Symbols, not line numbers, keep a citation valid across edits to the file: a +line inserted above a cited function leaves a symbol-only citation green. +""" + +from __future__ import annotations + +import ast +import re +from pathlib import Path + +import pytest + +TREE = Path(__file__).resolve().parents[2] +SOURCE = (TREE / "atom/model_engine/model_runner.py").read_text(encoding="utf-8") +DOCS = { + p.name: p.read_text(encoding="utf-8") + for p in sorted((TREE / "atom/compass/design").glob("*.md")) +} +CENSUS_DOC = "04_model_capture_and_cost_ir.md" + +_NUMBERS = r"\d+(?:-\d+)?(?:, ?\d+(?:-\d+)?)*" +_FILE = r"[\w-]+\.(?:py|md|json|log|sh|yaml|txt)\b" +_TOKEN = re.compile( + rf"model_runner\.py::(?P[\w.]+)(?:` \| (?P{_NUMBERS}) \|)?" + rf"|model_runner\.py`?(?::|,? lines? )(?P{_NUMBERS})" + rf"|(?P{_FILE})" + rf"|(?:`(?P(?!{_FILE})[\w.]+)(?:\(\))?` \()?(?:`|\(|# ):(?P{_NUMBERS})" +) + + +def _spans(source: str) -> dict[str, tuple[int, int]]: + spans = {} + + def walk(node: ast.AST, prefix: str) -> None: + for child in ast.iter_child_nodes(node): + if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + spans[prefix + child.name] = (child.lineno, child.end_lineno) + walk(child, prefix + child.name + ".") + + walk(ast.parse(source), "") + return spans + + +def _last(name: str | None) -> str | None: + return name and name.rsplit(".", 1)[-1] + + +def check(docs: dict[str, str], source: str) -> list[str]: + """One failure string per citation that names no symbol or leaves its span.""" + spans, failures = _spans(source), [] + names = {_last(name) for name in spans} + for doc, text in docs.items(): + owner = None + for segment in re.split(r"\n\s*\n|\n(?=\|)", text): + symbol = None + for m in _TOKEN.finditer(segment): + if m["file"]: + owner = m["file"] + continue + if not m["bare"]: + owner = "model_runner.py" + elif owner != "model_runner.py" and _last(m["beside"]) not in names: + continue + if m["symbol"]: + symbol = m["symbol"] + if symbol not in spans: + failures.append(f"{doc}: {symbol} is not defined") + continue + cited = m["cells"] or m["direct"] or m["bare"] + if not cited: + continue + if symbol is None: + failures.append(f"{doc}: {m[0]!r} names no symbol") + continue + lo, hi = spans.get(symbol, (0, 0)) + for part in re.split(r", ?", cited): + if not all(lo <= int(n) <= hi for n in part.split("-")): + failures.append( + f"{doc}: {part} is outside {symbol} ({lo}-{hi})" + ) + return failures + + +def test_every_citation_names_a_symbol_and_stays_in_its_span(): + assert check(DOCS, SOURCE) == [] + + +def test_a_blank_line_above_a_cited_function_stays_green(): + shifted = SOURCE.replace(" def prepare_inputs(", "\n def prepare_inputs(", 1) + before, after = _spans(SOURCE), _spans(shifted) + assert ( + after["ModelRunner.prepare_inputs"][0] + == before["ModelRunner.prepare_inputs"][0] + 1 + ) + assert check(DOCS, shifted) == [] + + +def test_a_line_moved_into_another_function_fails_by_name(): + docs = dict(DOCS) + docs[CENSUS_DOC] = docs[CENSUS_DOC].replace( + "prepare_sample` | 2564 |", "prepare_sample` | 2468 |" + ) + lo, hi = _spans(SOURCE)["ModelRunner.prepare_sample"] + assert check(docs, SOURCE) == [ + f"{CENSUS_DOC}: 2468 is outside ModelRunner.prepare_sample ({lo}-{hi})" + ] + + +@pytest.mark.parametrize( + "citation, token", + [ + ("`model_runner.py:1234`", "'model_runner.py:1234'"), + ("`atom/model_engine/model_runner.py`:1234", "'model_runner.py`:1234'"), + ("`model_runner.py` line 1234", "'model_runner.py` line 1234'"), + ("`model_runner.py` (`:1234`)", "'`:1234'"), + ("`warmup_model` (`:1234`)", "'`warmup_model` (`:1234'"), + ], + ids=[ + "colon", + "path-outside-backticks", + "line-word", + "bare-colon", + "bare-beside-name", + ], +) +def test_a_line_with_no_symbol_fails_by_name(citation, token): + docs = {**DOCS, CENSUS_DOC: DOCS[CENSUS_DOC] + f"\n\nSee {citation}.\n"} + assert check(docs, SOURCE) == [f"{CENSUS_DOC}: {token} names no symbol"] + + +def test_a_renamed_function_fails_by_name(): + renamed = SOURCE.replace("def prepare_sample(", "def prepare_sample_v2(", 1) + assert check(DOCS, renamed) == [ + f"{CENSUS_DOC}: ModelRunner.prepare_sample is not defined" + ] From bf322e862af441b50d2e346a067eeb16e9733ac6 Mon Sep 17 00:00:00 2001 From: Jiong Gong Date: Wed, 23 Sep 2026 13:25:11 +0000 Subject: [PATCH 2/2] compass(design): pin both span bounds and every citation spelling The moved-citation test now moves one census number into a later function and another into an earlier one, including a number from a comma-separated cell. Removing the upper bound, the lower bound or the list parsing now fails it by name. A bullet is now its own segment, like a table row, so a symbol in one bullet no longer covers a line number in the next. Every other spelling the recogniser accepts is now pinned by a parametrize row: `(:N`, `# :N`, `name()` (`:N`), a later table row and a later bullet. A bare `:N` after another `.py` or `.md` file is ignored. The owner extensions nothing needed are gone (`.json`, `.log`, `.sh`, `.yaml`, `.txt`). Co-Authored-By: Claude Opus 5.5 (1M context) --- atom/compass/design/03_memory_and_kv_model.md | 2 +- tests/compass/test_model_runner_citations.py | 103 ++++++++++-------- 2 files changed, 56 insertions(+), 49 deletions(-) diff --git a/atom/compass/design/03_memory_and_kv_model.md b/atom/compass/design/03_memory_and_kv_model.md index b45526c1d2..5d84985168 100644 --- a/atom/compass/design/03_memory_and_kv_model.md +++ b/atom/compass/design/03_memory_and_kv_model.md @@ -145,7 +145,7 @@ probed per request. Whether that is 0.1 ms or 10 ms is unmeasured. Recorded as * plus arithmetic**: ``` -# in _read_device_memory, which get_num_blocks calls first +# in _read_device_memory, which get_num_blocks calls before any budget arithmetic free, total = torch.cuda.mem_get_info() peak_torch = max(allocated_bytes.all.peak, .all.current) non_torch = max((total - free) - torch.cuda.memory_reserved(), 0) diff --git a/tests/compass/test_model_runner_citations.py b/tests/compass/test_model_runner_citations.py index 128379a7c2..6dd169d8ad 100644 --- a/tests/compass/test_model_runner_citations.py +++ b/tests/compass/test_model_runner_citations.py @@ -2,14 +2,10 @@ """Every `model_runner.py` citation in the design documents names a symbol. A citation is written `model_runner.py::Class.method`, and the symbol must be -defined in `atom/model_engine/model_runner.py` at this checkout. Any line number -cited into that file must fall inside the `ast` span (`def` to `end_lineno`) of -the nearest symbol cited before it in the same paragraph or table row. A line -number with no symbol before it fails. The line-number spellings recognised are -`model_runner.py:N`, `` `.../model_runner.py`:N ``, `model_runner.py` line N, -a bare `` `:N` ``, `(:N` or `# :N` whose nearest preceding file name is -`model_runner.py` or which follows a backticked name defined there, and a table -cell of line numbers right after a symbol's cell. +defined in `atom/model_engine/model_runner.py` at this checkout. A line number +cited into that file, in any spelling `_TOKEN` recognises, must fall inside the +`ast` span (`def` to `end_lineno`) of the nearest symbol cited before it in the +same paragraph, bullet or table row; with no symbol there, it fails. Symbols, not line numbers, keep a citation valid across edits to the file: a line inserted above a cited function leaves a symbol-only citation green. @@ -32,7 +28,7 @@ CENSUS_DOC = "04_model_capture_and_cost_ir.md" _NUMBERS = r"\d+(?:-\d+)?(?:, ?\d+(?:-\d+)?)*" -_FILE = r"[\w-]+\.(?:py|md|json|log|sh|yaml|txt)\b" +_FILE = r"[\w-]+\.(?:py|md)\b" _TOKEN = re.compile( rf"model_runner\.py::(?P[\w.]+)(?:` \| (?P{_NUMBERS}) \|)?" rf"|model_runner\.py`?(?::|,? lines? )(?P{_NUMBERS})" @@ -54,17 +50,13 @@ def walk(node: ast.AST, prefix: str) -> None: return spans -def _last(name: str | None) -> str | None: - return name and name.rsplit(".", 1)[-1] - - def check(docs: dict[str, str], source: str) -> list[str]: """One failure string per citation that names no symbol or leaves its span.""" spans, failures = _spans(source), [] - names = {_last(name) for name in spans} + names = {name.rsplit(".", 1)[-1] for name in spans} for doc, text in docs.items(): owner = None - for segment in re.split(r"\n\s*\n|\n(?=\|)", text): + for segment in re.split(r"\n\s*\n|\n(?=\||\s*[-*] )", text): symbol = None for m in _TOKEN.finditer(segment): if m["file"]: @@ -72,7 +64,9 @@ def check(docs: dict[str, str], source: str) -> list[str]: continue if not m["bare"]: owner = "model_runner.py" - elif owner != "model_runner.py" and _last(m["beside"]) not in names: + elif owner != "model_runner.py" and ( + (m["beside"] or "").rsplit(".", 1)[-1] not in names + ): continue if m["symbol"]: symbol = m["symbol"] @@ -100,49 +94,62 @@ def test_every_citation_names_a_symbol_and_stays_in_its_span(): def test_a_blank_line_above_a_cited_function_stays_green(): shifted = SOURCE.replace(" def prepare_inputs(", "\n def prepare_inputs(", 1) - before, after = _spans(SOURCE), _spans(shifted) - assert ( - after["ModelRunner.prepare_inputs"][0] - == before["ModelRunner.prepare_inputs"][0] + 1 - ) + key = "ModelRunner.prepare_inputs" + assert _spans(shifted)[key][0] == _spans(SOURCE)[key][0] + 1 assert check(DOCS, shifted) == [] def test_a_line_moved_into_another_function_fails_by_name(): - docs = dict(DOCS) - docs[CENSUS_DOC] = docs[CENSUS_DOC].replace( - "prepare_sample` | 2564 |", "prepare_sample` | 2468 |" + census = ( + DOCS[CENSUS_DOC] + .replace("prepare_sample` | 2564 |", "prepare_sample` | 2468 |") + .replace("| 2468, 2479, 2481 |", "| 2468, 2479, 2564 |") ) - lo, hi = _spans(SOURCE)["ModelRunner.prepare_sample"] - assert check(docs, SOURCE) == [ - f"{CENSUS_DOC}: 2468 is outside ModelRunner.prepare_sample ({lo}-{hi})" + spans = _spans(SOURCE) + assert check({**DOCS, CENSUS_DOC: census}, SOURCE) == [ + f"{CENSUS_DOC}: 2564 is outside ModelRunner.prepare_inputs (%d-%d)" + % spans["ModelRunner.prepare_inputs"], + f"{CENSUS_DOC}: 2468 is outside ModelRunner.prepare_sample (%d-%d)" + % spans["ModelRunner.prepare_sample"], ] -@pytest.mark.parametrize( - "citation, token", - [ - ("`model_runner.py:1234`", "'model_runner.py:1234'"), - ("`atom/model_engine/model_runner.py`:1234", "'model_runner.py`:1234'"), - ("`model_runner.py` line 1234", "'model_runner.py` line 1234'"), - ("`model_runner.py` (`:1234`)", "'`:1234'"), - ("`warmup_model` (`:1234`)", "'`warmup_model` (`:1234'"), - ], - ids=[ - "colon", - "path-outside-backticks", - "line-word", - "bare-colon", - "bare-beside-name", - ], -) +NO_SYMBOL = { + "colon": ("`model_runner.py:1234`", "'model_runner.py:1234'"), + "path-outside-backticks": ( + "`atom/model_engine/model_runner.py`:1234", + "'model_runner.py`:1234'", + ), + "line-word": ("`model_runner.py` line 1234", "'model_runner.py` line 1234'"), + "bare-backtick": ("`model_runner.py` (`:1234`)", "'`:1234'"), + "bare-paren": ("`model_runner.py` (:1234)", "'(:1234'"), + "bare-comment": ("`model_runner.py` # :1234", "'# :1234'"), + "bare-beside-name": ("`warmup_model` (`:1234`)", "'`warmup_model` (`:1234'"), + "bare-beside-call": ("`warmup_model()` (`:1234`)", "'`warmup_model()` (`:1234'"), + "later-row": ( + "\n| `model_runner.py::ModelRunner.forward` |\n| `model_runner.py:3300` |", + "'model_runner.py:3300'", + ), + "later-bullet": ( + "\n- `model_runner.py::ModelRunner.forward`\n- `model_runner.py:3300`", + "'model_runner.py:3300'", + ), +} + + +@pytest.mark.parametrize("citation, token", NO_SYMBOL.values(), ids=NO_SYMBOL) def test_a_line_with_no_symbol_fails_by_name(citation, token): docs = {**DOCS, CENSUS_DOC: DOCS[CENSUS_DOC] + f"\n\nSee {citation}.\n"} assert check(docs, SOURCE) == [f"{CENSUS_DOC}: {token} names no symbol"] +@pytest.mark.parametrize("owner", ["engine_core.py", "README.md"]) +def test_a_bare_line_owned_by_another_file_is_not_checked(owner): + tail = f"\n\n`model_runner.py::ModelRunner.forward` and `{owner}` (`:1234`).\n" + assert check({**DOCS, CENSUS_DOC: DOCS[CENSUS_DOC] + tail}, SOURCE) == [] + + def test_a_renamed_function_fails_by_name(): - renamed = SOURCE.replace("def prepare_sample(", "def prepare_sample_v2(", 1) - assert check(DOCS, renamed) == [ - f"{CENSUS_DOC}: ModelRunner.prepare_sample is not defined" - ] + renamed = SOURCE.replace("def prepare_sample(", "def prepare_sample2(", 1) + failure = f"{CENSUS_DOC}: ModelRunner.prepare_sample is not defined" + assert check(DOCS, renamed) == [failure]