You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Six correctness fixes on the capture path, plus the regression tests for the code that let them ship green. Branched from main (dac94dc); src/ changes are confined to six files.
Everything here came out of a review-and-verify pass over src/dmi. Every finding was reproduced before being filed and independently re-verified before being fixed — that pass refuted 7 other findings and corrected the reasoning on 5 more, which is why this PR is smaller than the audit that produced it.
The fixes
Severity
Fix
Symptom
High
point.py — bound the eager net by staging, not payload alone
Capture silently never delivered, ring reservation leaked permanently, flush_and_wait() returned success
High
adapter.py — disarm hook points on detach_model
Every ordinary forward after a monitored generate wrote unreserved bytes and stole the next step's metadata
Med
generation.py — give the eager decode step its attention mask
Left-padded batches attended their own pad K/V; tokens diverged from model.generate()
Med
generation.py — accept every documented eos_token_id spelling
AttributeError on the list form Qwen/Llama-3 ship
Med
internals.py — refuse duplicate capture chunks
Opaque torch error for normal payloads; silent wrong merge of two TP shards for one-element payloads
Low
clickhouse.py — refuse a database name that escapes its quoting
prefix_get silently read a different table
Two are worth reading the commit message for, because the obvious fix was wrong:
internals.py — a stable sort tiebreak looks like the natural fix and is the worst option: it makes the silent shard merge universal instead of removing it. Reassembly concatenates along the token axis, and two shards are the same tokens rather than more of them, so refusing is the only correct answer.
clickhouse.py — validating database with the existing column rule would reject my-analytics-db, 9lives and défaut, all legal for a backquoted ClickHouse database and all working today. The check refuses only what can break out of the quoting, and a test pins those names still rendering.
The tests
~2,900 lines, closing coverage gaps found by mutation testing — each one a mutation that left the entire suite green. The load-bearing examples:
BackendAdapter.plan_step and attach_model were never entered once by the 1,461-test gate; an unconditional raise at the head of each survived.
install_ring_hooks was never invoked — for spec in []: (a total no-op) left the suite green.
hook_row_basis, a public v1 API, was never called; only facade identity was asserted.
The publisher-lease quarantine handlers were untested, and deleting them lets a writer publish under a stale term after an outcome-unknown claim — the split-brain publish the quarantine exists to prevent.
Where a refusal is pinned, a positive control pins the accepted case too, so the refusal test cannot pass vacuously. Several mutants die only because of those controls.
Verification
CPU gate: 1,548 passed, 0 failed, 0 skipped (1,461 before this work). Order-independence checked across three shuffled seeds.
Every fix is red→green against a reverted copy of src/, with the mutated copy verified to be what actually loaded — this repo has three separate ways for a mutated tree to be silently ignored (pythonpath = ["src"] outranks PYTHONPATH, the __editable__ finder outranks both, and pytest-randomly isn't installed so -p no:randomly is a no-op).
Known limitations, stated rather than buried
The compiled decode path (cuda_graphs=True) still passes no attention mask. A per-step-growing mask changes shape every step and would defeat CUDA-graph capture; the static-max_cache_len alternative could not be verified here without real weights. Left-padded batches are correct on the eager path only — now stated in the docstring.
Detaching costs a recompile._ring_hook_type is a plain int so torch.compile bakes it as a constant, so clearing it invalidates a traced decode graph. That is the price of not corrupting the ring; a recompile-free fix needs the gate to move device-side.
internals.py still groups without shard_rank, so this does not silently make TP runs readable — that would be a contract change and is left for a deliberate decision.
The GPU suite could not serve as a gate. All 10 pre-existing failures trace to huggingface-hub==1.0.0.rc2 is required ... but found huggingface-hub==1.30.0 in the vendored transformers fork, and the suite is flaky run-to-run with the code held constant while another process holds ~15 GiB of the GPU. Notably test_e2e_correctness_hf — the test that would independently catch the decode-mask bug — cannot currently run at all. Restoring it is probably worth more than any single fix here.
Still open, deliberately not in this PR
generation.py:482 (monitoring silently no-ops without an engine) and generation.py:294 (short-circuit or strips one of two kwargs) both need an intent decision about which behaviour was meant, and neither corrupts data. docs/integration-api-v1.md:137 still describes MonitoringConfig as single-field.
Findings T2 (records.py:531 `_validate_cell`) and T3 (records.py:444-476
`_validate_payload_slices`) from the src/dmi polish round: both guards were
dead to the CPU gate. A verifier put `return` at the head of each method and
the whole suite stayed green, and a line trace confirmed 455/460/464/471/476
and 494-497/512-518/520-525 never executed.
Every refusal test also asserts `transport.events == []`, which is the
load-bearing part. The native side does re-check INT32 range
(`CopyLiteralRecordValue` in native/csrc/bindings.cpp), but only at
`push_record_descriptors` -- AFTER `reserve_record` -- so a late native throw
leaves an unmatched ring reservation. What these tests protect is that the
Python guard exists and fires FIRST. The residual the native check does not
cover at all: `True` and `3.7` are silently coerced by `py::cast<int64_t>`,
and a `list` is accepted as INT64_ARRAY.
Two positive controls are deliberate, not filler: the one-open-plus-bounded
slice case is what executes line 471's bounded arithmetic (it alone kills a
`> 1` -> `> 2` mutant), and the exact-bound case pins the boundary as
inclusive (it alone kills `>` -> `>=`).
Mutation-checked against throwaway copies of src: guard -> `return` fails
exactly the 4 T3 tests and exactly the 6 T2 tests; `> 1` -> `> 2` fails 1;
`>` -> `>=` fails 2. No src file changed.
Finding T1: `select_hook_specs`, `apply_hook_selection`,
`hook_belongs_to_pp_rank`, `hook_belongs_to_tp_rank`, `filter_by_pp_rank` and
`filter_by_tp_rank` (hooks/selection.py:113-228) had no behavioral test at
all. A line tracer over the full `-m cpu` run showed only their `def` lines
execute; `select_hook_specs -> return []`, which disables every hook on every
model, left the suite green.
Not dead code: adapters/base.py:168-170 calls all three on every
`attach_model`. Not covered elsewhere either -- the vLLM tests that do assert
on this surface live under third_party/, which pyproject.toml's
`norecursedirs` excludes, so they never run in this repo.
The `and spec.hook_type not in unavailable` clause is load-bearing rather
than defensive: model_shape.py:55 reads `getattr(cfg, "num_experts", 0)`,
which is 0 for Mixtral-style configs that spell it `num_local_experts`, and
`intermediate_dim` falls back to 0 for any non-gpt2 config without
`intermediate_size`/`ffn_dim`.
The stub hook point initialises `enabled` to a sentinel rather than True, so
"left untouched" is provable instead of coincidental.
Mutation-checked: 14 mutants, all killed -- including both always-True and
always-False variants of each rank predicate, `filter_by_pp_rank -> return
specs` and `-> return []`, dropping `enabled = False`, dropping the warning,
and each of the three RuntimeErrors. No src file changed.
Finding T6: the `except BaseException -> _quarantine_locked()` handlers of
`acquire_publisher_lease` and `renew_publisher_lease`
(clickhouse_catalog.py:686-702) were untested. Reducing both to a bare
`return self._leases.acquire(holder)` / `.renew()` left the CPU suite at its
baseline. The existing lease-failure tests all raise taxonomy errors, which
take the `except (CaptureStorageError, ValueError): raise` arm instead, and
the covered quarantine at line 469 is the watermark-INSERT path.
What the gap allows, measured: with the handlers removed, a writer whose
term-2 claim row LANDED but whose connection dropped before the outcome was
known keeps its stale term-1 identity and publishes -- `watermarks = [1]`.
That is the split-brain publish the quarantine exists to prevent.
The renewal test drives `renew_publisher_lease()` directly and says why in
its docstring: `_serial` is an RLock, so the renew inside `publish_snapshot`
is re-entrant within that method's own `except BaseException` at 468-470, and
routing through `publish_snapshot` lets the outer handler mask the inner one
being tested.
Third test is a negative control -- a taxonomy failure must NOT quarantine --
so an over-broad "quarantine on everything" fix cannot pass this file.
Mutation-checked: the handler removal fails both quarantine tests at
`assert writer.publisher_lease is None`; dropping the taxonomy arm fails the
control. Pure fakes, no server, runs under -m cpu. No src file changed.
Finding T5: the sink's CPU-tensor, contiguity, metadata-dtype and
metadata-shape guards (record_adapter.py:261-268) plus the
`persisted != target` durability check (295-298) were all unexercised --
four guards, not the two originally reported. Deleting 261-268 wholesale, or
replacing the durability raise with `pass`, left the CPU suite green; only
reducing `_validate_snapshot` to `return None` failed anything, which is
`_validate_losses` being covered while the accounting line is not.
Why the guards exist, since a reference-format producer cannot trip them:
`CaptureRecordFormat.encode` builds both cells from one `CaptureMetadata` and
already raises on dtype drift. These protect against a FOREIGN producer
writing the same layout name, or an encoder bug. Nothing else in the stack
cross-checks -- native/csrc/reference_python_capture_sink.cpp validates only
slice-internal consistency and never compares the metadata JSON to the slice.
Telling detail: the only existing proof this cross-check matters lives in a
test-only stand-in target in the gpu e2e suite, never in _CapturePackTarget.
The durability patch is applied after `_attach()`, not after construction:
`_attach` re-reads `snapshot()` against its baseline, so an earlier patch
trips "must remain empty before attachment" instead of reaching line 295.
Mutation-checked by rebinding the guards at runtime: all 5 new tests fail,
the 6 pre-existing ones stay green -- independently reproducing the finding.
No src file changed.
Finding T7: `RingTransport._record_cpu_tensor`'s three remaining transport
branches (transport/ring.py:422-443) never executed; this file covered only
SEQ_PREFIX_PACK and SEGMENTED_PACK. An unconditional `raise AssertionError`
at the head of all three left the full suite green.
This is the fallback that must byte-match the CUDA producer, so the clamps
are contractual, not defensive: producer.cu's
`record_producer_prefix_kernel` does `if (rows < 0) rows = 0;` and
`record_producer_chunked_kernel` does `if (bytes < 0) bytes = 0; if (bytes >
input_chunk) bytes = input_chunk;`. Without the CHUNKED clamp, `counts=[6,2]`
on an 8-byte 2-chunk source yields [0,1,2,3,4,5,4,5] -- chunk 1's bytes
duplicated into the record.
Deliberately NOT covered: the `min(byte_source.numel(), ...)` in
PREFIX_STRIP. A verifier showed it is an equivalent mutant -- Python slicing
already clamps, so removing it is byte-identical for every input, and the
surviving `max(0, ...)` carries the semantics. A test claiming to pin it
would assert nothing, which is the fake coverage this round exists to remove.
The -1 parametrization covers `max(0, ...)` instead.
`_entry()`'s signature is generalized to `**spec_args` because
TransportSpec.__post_init__ rejects `feature_bytes` for these three types;
both pre-existing tests call it by keyword and still pass.
Mutation-checked: the three-branch AssertionError now fails 8 cases (was
green); removing both clamps fails the -1 and [6,2] cases. No src changed.
Finding T8: hooks/specs.py:280 never executed, so no CPU test computed a
batched FINAL_LOGITS shape at all -- not merely the over-cap case. The
degenerate returns at 256/259/261/263 and the unknown-type return at 284
were equally untouched; setting `logits_q = 999999` and changing all five
`return []` sites to `return [1234]` both left the suite green.
The cap is reachable, not theoretical: adapters/huggingface/adapter.py:354
does `int(model_inputs.get("logits_to_keep", 0))` with no bound against
q_len, and it reaches `compute_hook_shape` unmodified via the public
`logits_to_keep` kwarg. Without `min(q_len, logits_to_keep)` the meta
declares a larger shape than the tensor actually pushed, and the drain thread
expects more bytes than the ring holds.
The MoE positives are marked plain `cpu` on purpose: the only existing tests
covering them are in test_moe_v1_routing_hooks.py under
`pytest.mark.framework_fork`, which skips without the Transformers fork -- so
those shapes had no coverage in the gate.
Why the neighbouring tests missed this: test_tp_shapes.py:131 drives only the
packed path (batch=0) and asserts just `shape[-1]`, and
test_integration_api_v1.py:275 compares `v1.compute_hook_shape` against
`specs.compute_hook_shape` -- the same object, so it can never catch a body
mutation.
`_VOCAB` is read from the `_cfg()` helper rather than hardcoded.
Mutation-checked: the combined mutant fails exactly the 6 covering tests;
dropping the `min()` fails the cap test. No src file changed.
Finding T9: `create_record_runtime`'s double-activation guard
(engine.py:271-272) and its entire `except BaseException` rollback (341-361)
never executed. Deleting the guard and sabotaging the except body both left
the suite green; the two existing failure tests raise BEFORE the `try` at
line 315.
The original finding's scenario was wrong, and the third test is why. On the
`create_record`/`init`/`start` paths the state reset is pure redundancy --
lines 323-324 already null both fields before `switched = True`, and
`_record_mode = True` is only reached at 338 after `start()` returns, so
`_record_mode=True` with `_ring_transport=None` is unreachable there. Only
`_rt.activate()` raising reaches the handler with all three fields naming the
new record ring, and it is also the only path that shows the second
`deactivate()`; the naive `assert deactivated == [True]` fails there.
Mutation table, each against a throwaway copy of src:
delete the guard (271-272) -> 1 failed (test 1)
sabotage `except BaseException` -> 4 failed (tests 2a-c, 3)
delete the state reset (353-356) -> 1 failed (test 3 ONLY)
delete `record_engine.stop()` -> 3 failed
delete rollback `deactivate()` -> 1 failed (test 3 ONLY)
The two rows marked ONLY are the point: without the activate-failure case
both mutants survive, and the coverage would have proved nothing about the
reset. No src file changed.
Finding T10: `CaptureReader`'s byte-budget eviction disjunct
(reader.py:348-351) was untested -- the file covered count-based eviction and
oversized-entry rejection only. Deleting `or self._footer_cache_bytes +
entry.wire_bytes > self._footer_cache_bytes_limit` left the CPU suite at
baseline, while deleting the whole eviction body or the admission guard each
failed an existing test, localizing the gap to that disjunct.
The branch is live, not dead code behind the oversized-entry guard: that
guard rejects only an entry larger than the ENTIRE budget, whereas this is
cumulative overflow across several individually-admissible entries. With the
shipped defaults (128 packs / 64 MiB) it fires once mean footer size passes
roughly 512 KiB, reachable for packs near `max_pack_records = 10_000`.
The test asserts `_footer_cache_limit >= 3` BEFORE lowering the byte limit,
so count eviction cannot be what it observes -- otherwise it would silently
re-test the neighbouring case. Footer size comes from the existing
`_footer_read_bytes(sealed)` helper rather than a hardcoded constant.
Mutation-checked: the disjunct's removal fails it with `assert 3 == 2`. No
src file changed.
Findings from round 2: `BackendAdapter.plan_step` (base.py:259) and
`attach_model` (base.py:167) were never entered ONCE by the 1508-test gate. A
verifier proved it with the strongest possible probe — an unconditional
`raise` at the head of each function survives the whole suite.
`plan_step` was invisible because the suite's only adapter stub
(tests/test_adapter_protocol.py:111) overrides it, and that override turns out
to be incidental: it exists to record call order, not because anything forces
it. `PlanningAdapter` here deliberately overrides neither, which is the point
of the file. `attach_model` had no call site in tests/ at all -- only a
mention in a docstring.
Two tests carry most of the weight:
test_plan_step_rounds_each_hook_up_to_16_bytes_independently -- three
6-byte hooks must total 48, NOT 18. Alignment is per-hook rather than
per-sum, and this is the only case that kills `align_up_py(nbytes, 16)`
-> `nbytes`. Note 48 is also not align_up(18)=32.
test_attach_model_runs_selection_then_pp_then_tp_filters_and_installs --
at tp_rank=1, is_pp_first=False the survivors must be exactly {Q,
MLP_POST, PATTERN}. round 1 pinned each filter in isolation; nothing
pinned the three-call WIRING, so deleting a call from the chain was free.
The stub hook point starts `enabled` and `_ring_payload` at sentinel strings
rather than True/None, so "left alone" is distinguishable from "installed" and
the not-installed assertions cannot pass vacuously.
12 mutants killed, re-verified after the file was restyled: both unconditional
raises, the alignment mutant, deleting each of the three chain calls, dropping
the empty-shape skip, ignoring actual_q_len, ignoring the per-spec dtype
override, unconditional needs_eager, skipping install_ring_hooks, and
publishing unfiltered specs. No src file changed.
Three round-2 findings in records.py, each mutation-confirmed against the
1508 baseline and each killed by exactly one of these tests:
:308 -- no test exercised a SUCCESSFUL device-gated bind_hook. The
negative branch was pinned (`if True:` is caught); the accepted-gate
branch was not. The reserved flag is RecordReservationItem.needs_reclaim,
and only when true does reserve_record register a pending task reclaim
(drain_thread.cpp:278-281); the gated IDENTITY kernel returns before
copying AND before publishing when the gate blocks (producer.cu:429-430),
so this flag is the ring's only handling of that case. The test asserts
("reserve", ((16, True),)) against the ungated (16, False) already pinned
at line 155.
:347 -- prepare_replay's metadata/plan arity guard never fired. Without
it the body publishes 2 descriptors while _reservation_items reserves 3,
breaking push_record_descriptors' documented ordering contract
(transport/ring.py:391).
:433,:438 -- _validate_entry_output's dtype and input-shape drift
refusals never fired. Realistic input, since a captured plan's entry
outlives the output it was derived from. Note the existing
test_payload_slice_dtype_drift_is_refused pins a DIFFERENT guard (:463)
and does not reach these.
A verifier correction is recorded here because it changes what the first test
means: the original finding said the missing flag lets a record carry stale
payload. It is one step off -- an unpublished gated task with
needs_reclaim=True leaves an unresolved pending reclaim, which
ring_engine_py.cu:511-514 reports as "record flush found incomplete producer
reclaims". The untested property is the safety flag itself.
Every refusal test asserts `transport.events == []`, pinning refusal before
reservation. Kill checks were re-derived independently by monkeypatch
simulation rather than taken on trust. No src file changed.
Round-2 finding, selection.py:101-107: the unknown-token and empty-result
refusals were unexercised, though the path is on the attach line
(adapters/base.py:168 -> selection.py:151 -> :118) and the only existing
reference was a success assertion at tests/test_adapter_protocol.py:385.
The mutant that makes the stakes clear is stronger than the one first filed:
`_HOOK_SELECTIONS.get(token, _ALL_HOOK_TYPES)`, i.e. a typo'd
`hook_selection=` silently selects EVERY hook instead of raising. The suite
stayed green.
The test asserts the message names the offending token AND lists the
available names -- two assertions beyond the obvious, because
`"Available:" in message` alone survives a mutant that emits the label with
an empty list. It also pins propagation through select_hook_specs, so the
attach-path caller is covered rather than just the leaf.
Seven mutants killed, each only by this test. No src file changed.
Round-2 finding, specs.py:121-127. The original filing understated it:
replacing the return value survives, but so does putting `raise
AssertionError` in the BODY -- so this public v1 API is never CALLED by the
CPU gate at all, not merely under-asserted.
The only existing reference is the facade-identity assert at
tests/test_integration_api_v1.py:175, which compares the same object to
itself and therefore can never catch a body change. Its only non-test callers
are in the third_party vLLM adapter, which `norecursedirs` excludes from
collection.
Pins the docstring's contract -- "Logit-shaped payloads are request-scaled;
every other registered shape class is token-scaled" -- plus the unknown-type
refusal. No src file changed.
Round-2 finding, dispatch.py:41-49: no CPU test invoked install_ring_hooks.
The strongest possible mutation proves it -- `for spec in specs:` -> `for
spec in []:`, making the function a total no-op, left the gate at 1508
passed. The only reference in the tree was the facade-identity line at
tests/test_integration_api_v1.py:176.
It needs no CUDA, which is why the gap was closable: install_ring_hooks only
assigns three attributes on `spec.module` and never touches torch.ops.ring
(that is dispatch_producer, a different function), so a plain nn.Identity()
drives it.
`_ring_hook_id` carries spec.layer_no, and per-layer reassembly buckets on
that field (storage/internals.py:_reassemble_per_layer, key index 3) -- so
the stuck-id mutant would collapse every layer into one. That is why the test
asserts the ids are [0, 1, 2] rather than merely present.
Four mutants killed: the no-op loop, `_ring_hook_id = 0`, the removed guard
(AttributeError on None instead of the named RuntimeError), and
`_ring_payload = None`. No src file changed.
Round-2 finding, reader.py:333 -- and the line number is the finding. It was
originally filed at :342, which a verifier REFUTED: :342 sits inside the
double-checked-lock block and is unreachable single-threaded (an
AssertionError there survives the whole gate), so the cited mutation could
not produce the cited LRU consequence. The real hit-path `move_to_end` is
:333, where `-> pass` survives both the full gate and all 141 tests in
test_capture_storage.py + test_capture_summary.py.
Distinct from test_the_footer_cache_evicts_to_stay_inside_its_byte_budget,
committed earlier in this run: that test never re-reads a cached pack, so no
hit-path recency is observable through it. Without :333 the cache degrades
from LRU to FIFO and the hot pack's footer is re-fetched cold on every read
-- the exact cost the cache exists to avoid.
The test asserts `_footer_cache_bytes_limit >= sum(footer_sizes)` up front,
so if the default byte budget ever shrinks it fails loudly instead of
silently degrading into a re-test of byte eviction.
Correcting the verifier on one point: it reported that no existing test kills
`popitem(last=True)` or `:355 -> pass`. Both ARE already covered -- the first
by the byte-budget test from earlier in this run, the second by six tests
including test_estimate_predicts_reads_exactly_without_coalescing. Line 333
is the genuinely uncovered one.
`_RecordingStore` gains a `reads` list additively; `ranges` and every
assertion on it are untouched. No src file changed.
Round-2 finding, engine.py:473-474: mutating `self._auto_batch_group_id += 1`
to `pass` left the gate at 1508 passed. No test referenced the counter on a
real engine -- tests/test_hf_eos_strip.py:47 and
tests/test_e2e_correctness_vs_hf.py:33 both REIMPLEMENT it (in a stand-in
fake and in a docstring), so neither exercises engine.py, and the only real
caller is adapters/huggingface/adapter.py:398, reached solely from
CUDA-gated e2e tests that are not cpu-marked.
A verifier refuted the assumption that this needs hardware:
MonitoringEngine(enable_ring_transport=False) constructs with no ring, no
native host and no CUDA, and the counter returns [0, 1, 2] -- a construction
pattern cpu tests already use (tests/test_monitoring_engine_shutdown.py:22).
Why a stuck counter matters: adapter.py:393-399 bumps the group on every
prefill or batch-size change and mints per-request ids as f"{group}:{i}", so
two successive generate() calls would both emit "0:0", "0:1", ... and the
offload table is MergeTree with no dedup (native/csrc/clickhouse_client.cpp:389),
so those rows collide in one catalog namespace. Pins the documented contract
"returns engine-scoped integers starting at zero"
(docs/integration-api-v1.md:242). No src file changed.
Two rounds over src/dmi. The state file carries the per-finding ledger with
verdicts and dispositions; lessons.md carries what transfers to the next run
-- chiefly that this repo has three separate ways for a mutation harness to
silently test the pristine tree, and that a surviving mutant is
indistinguishable from one that never loaded unless the harness is proven
first against a known-covered mutant.
`database` was the one identifier this class stored raw and then interpolated
exactly like `table`: `_build_select_sql` renders
`FROM {_backtick(db)}.{_backtick(table)}` into every prefilled statement. A
name closing its own quoting rewrote the statement -- `default`.`other` --
commented out the intended table, the WHERE and the ORDER BY, so `prefix_get`
silently read a different table while `table="off`load"` had always been
refused.
Deliberately NOT `_validate_ident`. That is the COLUMN rule,
`[A-Za-z_][A-Za-z0-9_]*`, and it would reject `my-analytics-db`, `9lives` and
`défaut` -- all legal for a backquoted ClickHouse database, all working today,
and all permitted by the documented `database: str` signature. The regression
test pins those four names rendering unchanged, so a later "tighten it to
match the columns" change fails loudly instead of breaking deployments.
Refusing rather than escaping is also deliberate: ClickHouse honours more than
one escape convention inside backquoted identifiers, so a wrong escape would
silently address a DIFFERENT database instead of failing -- the same class of
bug this guards against. Backtick, backslash and control characters are
refused; everything else is passed through as before.
Scope note, since it argues against overrating this: `database` is never
data- or request-influenced in-tree (the only production constructor,
storage/internals.py:_default_reader, leaves it at the "default" literal), and
docs/integration-api-v1.md:1206 already disclaims `custom_select`'s text check
as "not a security boundary". This closes the shape without claiming the class
is hardened against a hostile caller.
Red -> green verified against a reverted copy of src: the injection test fails
before the change, the four legal names pass both before and after.
`generate_greedy_with_monitoring` compared the sampled token against
`eos_token_id` with `!=` and the finished sequence with `==`. Both are correct
only for a scalar: torch returns a plain Python `bool` for `tensor != list` --
no broadcast, no error -- so `.long()` raised
AttributeError: 'bool' object has no attribute 'long'
on the first decode step past `min_new_tokens`, and `.nonzero()` would have
raised on the same bool afterwards.
The list form is not exotic. It is what `generation_config.eos_token_id` holds
for Qwen2.5/Qwen3 and Llama-3, it is what HF's own `generate()` takes, and
line 508 forwards this very argument into `HuggingFaceAdapter`, whose
docstring says it "Accepts ``int``, ``list[int]``, or ``torch.Tensor``". So
the function crashed on a value its own callee documents accepting.
A multi-element TENSOR did not crash but was worse in kind: it broadcast
against the batch dimension, comparing request i against eos id i.
Fixed by normalising once, where `device` is bound, into a 1-D int64 tensor,
then using `torch.isin` at both sites -- elementwise over the batch for every
spelling. int64 because argmax yields int64, so both sides compare in one
dtype.
No behaviour change for a scalar `eos_token_id`: `isin` against a
one-element tensor is identical to the old `!=`, and the in-repo callers
(benchmarks/bench_hf_transport.py, tests/hf_reference_runner.py) all pass
`int(tokenizer.eos_token_id)`. Only the currently-broken list and
multi-element-tensor inputs change, and nothing can depend on those.
The three GPU tests assert the list form against the SCALAR form on the same
scripted model rather than against a hand-computed sequence, so they pin
equivalence instead of transcribing the loop. Red -> green verified against a
reverted copy of src: all three fail before the change. The docstring now
states the accepted spellings.
The eager path admitted a tensor into the ring on `transport_bytes <=
available_capacity()` or `<= payload_cap()`. Both are payload-only: the ring's
real per-step ceiling is min(payload, staging), which is what
RingCapacities.effective_bytes publishes, what native prepare_step and
reserve_record use (`effective_cap = std::min(pcap, scap)`), and what the HF
adapter already warns about ("Effective capacity is staging-limited").
So the safety net accepted exactly the tensor prepare_step had just rejected.
Measured against the real extension with payload=64 MiB, staging=1 MiB:
0.5 MiB: avail 67108864 -> reserve 66584576 -> after flush 67108864
1.0 MiB: avail 67108864 -> reserve 66060288 -> after flush 67108864
2.0 MiB: avail 67108864 -> reserve 65011712 -> after flush 65011712
The drain assembles each flush batch per WHOLE entry and breaks when the entry
does not fit staging (drain_thread.cpp:350,434 -- no split or partial path,
and PinnedStaging is a fixed cudaHostAlloc), so the 2 MiB capture is never
delivered AND its payload reservation is never released: ring capacity shrinks
for the life of the process while flush_and_wait() returns success and exits 0.
Both branches carried the flaw, and which one fires depends on ring occupancy
-- the OVERSIZED path force-flushes first, so a step's first hook takes the
available_capacity branch and a later hook in the same step takes the
payload_cap one. Fixing only one would have left the other live, which is why
there is a test per branch plus a control proving the ring is still used when
staging covers the tensor.
Anything above the effective cap now falls to the existing cpu_direct branch:
slower, but it is currently lost outright, so no correct caller depends on the
old behaviour. `_FakeEagerRingEngine` gains `staging_cap()`, defaulting to the
payload capacity exactly as RingConfig does when pinned staging is left at 0,
so the two pre-existing tests are unaffected.
Red -> green: both new refusal tests fail before the change; 20 passed after.
Prefill was handed `attention_mask`; the decode kwargs were not. HF builds an
all-ones causal mask over the whole cache when none is supplied, so every
decode step let a left-padded row attend to its own pad positions' K/V and the
greedy tokens diverged from model.generate(). Verified previously on a real
model (tiny 4-layer Llama, fp32, sdpa, one left-padded row of 5 pads + 3 real
against one full row):
hf generate : [[66, 80, 80, 80, 80, 80], [36, 67, 63, 21, 107, 63]]
greedy loop : [[66, 80, 107, 82, 52, 52], [36, 67, 63, 21, 107, 63]]
The padded row diverged from the third generated token on; the unpadded row
was identical, and an unpadded control batch matched generate() exactly. Cause
isolated by injecting one variable at a time: supplying the grown mask alone
restored exact parity, correcting position_ids alone changed nothing. So the
mask is the fix and position_ids is left as it was.
The mask grows with the cache -- the caller's prompt columns unchanged, then
all-ones for the generated positions, which are all real. No behaviour change
for unpadded or right-padded input, where the added columns and the implied
all-ones mask agree; only the currently-wrong left-padded case moves.
Not fixed, and now stated in the docstring: the compiled path
(`cuda_graphs=True`) still passes no mask. A per-step-growing mask changes
shape every step and would defeat CUDA-graph capture, and the StaticCache
alternative (a fixed max_cache_len mask mutated in place) could not be
verified here without real weights. Left-padded batches are therefore correct
on the eager path only.
The adaptor's before_forward_manual keeps receiving the caller's prompt mask,
not the grown one: it derives prefill KV offsets from the prompt, which is
what it wants, and that call is unchanged by this commit.
Tests pin the mechanism with a recording stub -- every decode step receives a
mask, its width tracks the cache, prompt padding survives and generated
columns are admitted -- so they need no weights. All three fail before the
change.
`install_ring_hooks` was the only thing that ever wrote `_ring_hook_type` /
`_ring_payload`, and nothing cleared them. `detach_model` restored the prepare
wrapper and emptied `transport._active_specs`, but the HookPoints are
structural members of the model and outlived the call still armed.
That matters because both `generate_with_monitoring` and
`generate_greedy_with_monitoring` call `detach_model` from a `finally`, so
every monitored generate ends with the model armed. The next ordinary forward
on it -- an eval pass, a perplexity computation, another library calling the
same object -- then launches producer kernels with no `prepare_step`
reservation and an empty metadata FIFO. `HookPoint.forward` gates only on
`enabled`, `_ring_hook_type is not None` and `x.is_cuda`; `_active_transport`
stays set until `engine.close()`, so `g_active_engine` is non-null and the
native producer does not early-return; and `hook_no_notify` is documented "No
condition gating. Space is guaranteed by the pre-forward capacity check in
Python" -- a check that did not run.
Measured before this change on a real model: 27 producer dispatches and 29,638
bytes written into the payload ring after detach, while `available_capacity()`
still reported only the previous generate's bytes outstanding. The strays are
worse than lost -- `do_post_processing` pops one TensorMeta per task in FIFO
order, so they consume the NEXT monitored step's metas and every later payload
is paired with the wrong one.
`uninstall_ring_hooks` is the exact inverse of the install: it clears
`_ring_hook_type` (the condition `HookPoint.forward` actually returns on) and
drops the payload reference so a detached model stops pinning the ring buffer.
`enabled` is deliberately left alone -- that carries the hook SELECTION, which
`apply_hook_selection` owns and a re-attach reuses. Unlike install it tolerates
an unbound spec, because it runs from teardown where raising would displace
whatever the caller was already handling.
The cost, stated rather than hidden: `_ring_hook_type` is a plain int
precisely so torch.compile bakes it as a compile-time constant, so clearing it
invalidates a traced decode graph and re-attaching costs a recompile. That is
the price of not corrupting the ring on the next ordinary forward. A
recompile-free variant would need the gate to move device-side; that is a
larger change and is not attempted here.
Tests pin the disarm, the pre-existing teardown that must survive it, the
prepare-wrapper restore, and idempotence (it runs from a `finally`, so a second
call after an error path must not raise). All CPU -- no CUDA needed.
The three `_reassemble_*` helpers grouped a request's chunks by
`(request, layer)` and then called `sorted(chunks)` on `(start_token, tensor)`
pairs. With no sort key, a tie on start token makes Python fall through to
comparing the TENSORS:
* a normal payload raises "Boolean value of Tensor with more than one value
is ambiguous" -- an error naming nothing about the duplicate, the request,
the layer or the shard, so the reader cannot act on it; and
* a ONE-ELEMENT payload compares fine and is silently concatenated, merging
two captures into a single wrong tensor of double the token count.
The silent case is the dangerous one and it is why a stable tiebreak
(`key=lambda c: c[0]`) would have been the wrong fix: it makes that silent
merge universal instead of removing it.
The collision is production-shaped, not hypothetical. `filter_by_tp_rank`
deliberately keeps TP-sharded hooks on every rank, `resolve_shard_rank()`
returns `ctx.tp_rank` for them, and the repo writes all ranks to one table
with `shard_rank` telling the rows apart -- tests/hf_compare_runner.py says so
outright ("Production writes all ranks to one table with shard_rank
distinguishing per-rank rows"), and the repo's own comparison tooling
(tests/hf_comparator.py) picks one shard_rank before grouping, which is
exactly the filter internals.py lacked. Because `get_internal` reassembles
every field eagerly, one colliding sharded act made even the complete,
unsharded hidden_states unreachable.
Reassembly here concatenates along the TOKEN axis, and two shards are the same
tokens rather than more of them, so merging is not a correct answer at any
alignment: refusing is. `_ordered_chunks` now sorts by an explicit key -- so
tensors are never compared under any input -- and raises a named RuntimeError
identifying the act, the request and the duplicated start tokens, and pointing
at shard_rank and shared model_ids as the two causes. RuntimeError rather than
a new type so callers already catching the torch error keep working.
Not changed: grouping still ignores shard_rank, so this does not silently make
TP runs readable. That would be a contract change (docs/integration-api-v1.md
disclaims cross-shard reassembly) and is left for a deliberate decision.
Tests cover both shapes at the per-layer and non-layered sites -- including the
one-element case that previously returned a wrong answer instead of raising --
plus controls proving out-of-order chunks still reassemble in token order and
a single-shard run is still readable.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
A moderate implementation concern and documentation nit remain, and GPU coverage is unavailable.
Pull request overview
This PR fixes six capture-path correctness issues and adds extensive regression coverage across hooks, adapters, generation, storage, and runtime behavior.
Changes:
Enforces staging-aware capture limits and safe hook teardown.
Fixes eager decode masks, EOS handling, duplicate chunks, and database validation.
Adds broad behavioral and lifecycle test coverage.
File summaries
File
Summary
tests/test_tp_shapes.py
Covers hook shape edge cases.
tests/test_storage_internals_duplicate_keys.py
Tests duplicate-chunk refusal.
tests/test_storage_clickhouse.py
Tests database identifier validation.
tests/test_record_runtime.py
Covers record-runtime validation paths.
tests/test_record_cpu_direct.py
Tests CPU-direct transformations.
tests/test_producer_chunked_schema.py
Tests staging-aware eager capture.
tests/test_hook_spec_flags.py
Tests hook metadata flags.
tests/test_hook_selection_filters.py
Tests hook selection and rank filtering.
tests/test_hook_dispatch_install.py
Tests hook binding and teardown.
tests/test_hf_greedy_eos.py
Tests EOS representations and stopping.
tests/test_hf_greedy_decode_mask.py
Tests eager decode masks.
tests/test_hf_adapter_detach.py
Tests adapter cleanup.
tests/test_engine_runtime_api.py
Tests runtime lifecycle and rollback.
tests/test_clickhouse_capture_catalog.py
Tests lease quarantine handling.
tests/test_capture_storage.py
Tests footer caching behavior.
tests/test_capture_record_adapter.py
Tests capture-record validation.
tests/test_adapter_protocol.py
Tests adapter planning and commit flow.
tests/test_adapter_base_pipeline.py
Tests attachment and planning pipelines.
src/dmi/storage/internals.py
Rejects duplicate capture chunks; diagnostic construction should avoid O(n²) scans.
src/dmi/storage/clickhouse.py
Validates database identifiers safely.
src/dmi/hooks/point.py
Enforces effective ring capacity.
src/dmi/hooks/dispatch.py
Adds hook teardown support.
src/dmi/adapters/huggingface/generation.py
Fixes eager masks and EOS handling; documents right-padding limitations.
src/dmi/adapters/huggingface/adapter.py
Disarms hooks during detachment.
.loop/polish-state.md
Records review polish state.
.loop/polish-seen.md
Records reviewed polish items.
.loop/lessons.md
Records review lessons.
Review details
Suppressed comments (2)
src/dmi/adapters/huggingface/generation.py:477
The documented exception is broader than the implementation: the compiled decode path omits attention_mask, and this function always takes logits[:, -1] during prefill, so a right-padded row can both select a pad-position logit and expose right-padding K/V during mask-free decode. The repository's HF callers consistently use left padding (for example, benchmarks/bench_hf_transport.py:813), so please remove “Right-padded” from the unaffected claim or explicitly document right-padded inputs as unsupported.
the eager path (``cuda_graphs=False``). Right-padded or unpadded
batches are unaffected.
src/dmi/storage/internals.py:72
The duplicate case is explicitly expected for TP rows, but starts.count(s) scans the full list for every start, making diagnostic construction O(n²) per request/layer. Long captures with multiple shard rows can spend far more time building the error than refusing the invalid merge; build a frequency map once and derive duplicated from it.
duplicated = sorted({s for s in starts if starts.count(s) > 1})
The v1 doc said "MonitoringConfig currently contains only this schedule",
which stopped being true when `storage_backend` and `capture_sink_config`
were added. Both are exported through `dmi.api.v1`, both are acted on by
`MonitoringEngine`, and two combinations raise at construction -- so a caller
following only this document could not know the fields existed, let alone that
setting them can refuse.
Verified against engine.py:120-133 rather than transcribed: "native" without a
host engine raises, and "capture"/"none" WITH one raises. The fields are
documented elsewhere (docs/capture-storage-design.md, docs/benchmarks.md), so
the two documents disagreed and the v1 contract was the stale one.
The /polish loop's working files are per-run agent bookkeeping, not project
artifacts: nothing in the tree, the build, or CI reads them, and they go
stale once the branch they describe has landed. #133 removes them from
main; dropping them here too so a merge of this branch cannot carry them
back in.
…he mask
The eager decode step sent every row torch.full((B,1), Pmax+step+1):
a k-left-padded row overshot its true position by k+1, and even an
unpadded row by 1 (the CUDA-graph path starts cache_position at Pmax,
so the two paths disagreed by one every step). Prefill is pad-aware
(mask.cumsum(-1)-1), so a row's first decode token belongs at its
real-token count -- exactly what HF derives from the grown decode mask
via decode_mask.cumsum(-1)[:, -1] - 1. Compute next_pos from the mask
once and send next_pos + step per row.
While there, build the widest decode mask once before the loop and
slice it per step instead of a per-step torch.cat, which was O(T^2)
over the decode; values, dtype and device are identical.
The duplicate-chunk refusal tells the caller to 'select a single
shard_rank before reading', but neither get_internal nor LazyInternal
had any such parameter -- the prescribed remedy was unreachable, so
every TP>1 capture was a hard dead end. Give both an optional
shard_rank that filters rows by key[4] before reassembly; the default
None keeps the refusal-on-collision behavior, whose message is now an
instruction the API can actually follow.
HookPoint.forward recomputed min(payload_cap(), staging_cap()) on every
eager invocation -- two pybind crossings per hook forward on getters the
native header documents as not to be called per-step, computed even on
the strip cpu-direct branch that never uses the value, and the fourth
open-coded spelling of RingCapacities.effective_bytes. Name the formula
once as engine.effective_ring_bytes, cache the result on the hook the
first time the eager ring branch needs it, and clear the cache on
(un)install since a re-attach may bind a different engine.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
The duplicate detection is quadratic, and the compiled-generation documentation incorrectly claims right-padded batches are safe without a decode attention mask.
Get a fresh assessment by requesting another Copilot review.
The refusal named exactly the duplicated starts, but found them with
starts.count(s) per start: a long chunked capture paid O(n^2) to build an
error message that only names the repeats. One frequency map now feeds the
same sorted list, so the diagnostic is linear plus the sort.
The new case pins that only the colliding starts are named, sorted -- the
property the O(n^2) set comprehension provided and the fix must keep.
The per-hook cache only moved the burst: the first eager forward still
answered payload_cap()/staging_cap() once for EVERY active hook, and the
native header documents those getters as startup-only. The min now lives on
the RingTransport (engine-lifetime constants, one transport per engine) and
every hook reads it; a re-arm builds a fresh transport, so a new engine's
caps are read rather than the old min. The install/uninstall cache resets go
with it -- there is no per-hook cache left to clear.
The eager-cap test now drives the real RingTransport (counting engine
behind it), so the cache under test is the production one, and pins that a
second arming reads the second engine instead of reusing the first.
The note said right-padded batches were unaffected, but this loop always
takes the last prompt position's logits -- a pad under right padding -- and
the compiled decode step runs with no attention mask, so the pad's K/V stays
visible to decode. Right-padded prompts are unsupported on both paths; only
unpadded is unaffected, and left-padded only on the eager path.
- ring.py imports effective_ring_bytes at module level like every other
consumer; the function-level import bought nothing (dmi.engine has no
transport import to cycle through).
- The eager-cap test uses the real RingTransport with submit_cpu_direct
monkeypatched on the instance, instead of a wrapper class that forwarded
every attribute through __getattr__.
- The chunked-schema fake computes its effective_cap once in __init__
rather than recomputing it behind a property.
No behavior change; 41 passed, 20 skipped (native backend absent).
This filters after prefix_get has already fetched and torch_decoded every shard row. For a TP capture, shard_rank therefore does not reduce ClickHouse traffic or peak memory and can still materialize all ranks before discarding them, defeating the selector and risking an avoidable OOM. Push the rank predicate into the reader query/API so only the selected rows are decoded.
Apply shard rank filtering before eager materialization
src/dmi/storage/internals.py:668
The eager get_internal path has the same late-filter problem: prefix_get((model_id,)) materializes all TP tensors before shard_rank is applied. Selecting one rank should be a storage-level predicate; otherwise large multi-rank captures still incur the full transfer and decode cost even though the result keeps only one shard.
Despite the name and docstring, torch.ones(1, 4) is a one-row, four-element tensor. The old tuple sort would still compare a multi-element tensor and raise the ambiguous-bool error, so this test cannot prove the claimed silent-merge regression for a single-element payload. Use a one-element tensor (or scalar) so the pre-fix code would actually concatenate silently.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this is
Six correctness fixes on the capture path, plus the regression tests for the code that let them ship green. Branched from
main(dac94dc);src/changes are confined to six files.Everything here came out of a review-and-verify pass over
src/dmi. Every finding was reproduced before being filed and independently re-verified before being fixed — that pass refuted 7 other findings and corrected the reasoning on 5 more, which is why this PR is smaller than the audit that produced it.The fixes
point.py— bound the eager net by staging, not payload aloneflush_and_wait()returned successadapter.py— disarm hook points ondetach_modelgeneration.py— give the eager decode step its attention maskmodel.generate()generation.py— accept every documentedeos_token_idspellingAttributeErroron the list form Qwen/Llama-3 shipinternals.py— refuse duplicate capture chunksclickhouse.py— refuse adatabasename that escapes its quotingprefix_getsilently read a different tableTwo are worth reading the commit message for, because the obvious fix was wrong:
internals.py— a stable sort tiebreak looks like the natural fix and is the worst option: it makes the silent shard merge universal instead of removing it. Reassembly concatenates along the token axis, and two shards are the same tokens rather than more of them, so refusing is the only correct answer.clickhouse.py— validatingdatabasewith the existing column rule would rejectmy-analytics-db,9livesanddéfaut, all legal for a backquoted ClickHouse database and all working today. The check refuses only what can break out of the quoting, and a test pins those names still rendering.The tests
~2,900 lines, closing coverage gaps found by mutation testing — each one a mutation that left the entire suite green. The load-bearing examples:
BackendAdapter.plan_stepandattach_modelwere never entered once by the 1,461-test gate; an unconditionalraiseat the head of each survived.install_ring_hookswas never invoked —for spec in []:(a total no-op) left the suite green.hook_row_basis, a public v1 API, was never called; only facade identity was asserted.Where a refusal is pinned, a positive control pins the accepted case too, so the refusal test cannot pass vacuously. Several mutants die only because of those controls.
Verification
src/, with the mutated copy verified to be what actually loaded — this repo has three separate ways for a mutated tree to be silently ignored (pythonpath = ["src"]outranksPYTHONPATH, the__editable__finder outranks both, andpytest-randomlyisn't installed so-p no:randomlyis a no-op).Known limitations, stated rather than buried
cuda_graphs=True) still passes no attention mask. A per-step-growing mask changes shape every step and would defeat CUDA-graph capture; the static-max_cache_lenalternative could not be verified here without real weights. Left-padded batches are correct on the eager path only — now stated in the docstring._ring_hook_typeis a plain int sotorch.compilebakes it as a constant, so clearing it invalidates a traced decode graph. That is the price of not corrupting the ring; a recompile-free fix needs the gate to move device-side.internals.pystill groups withoutshard_rank, so this does not silently make TP runs readable — that would be a contract change and is left for a deliberate decision.huggingface-hub==1.0.0.rc2 is required ... but found huggingface-hub==1.30.0in the vendored transformers fork, and the suite is flaky run-to-run with the code held constant while another process holds ~15 GiB of the GPU. Notablytest_e2e_correctness_hf— the test that would independently catch the decode-mask bug — cannot currently run at all. Restoring it is probably worth more than any single fix here.Still open, deliberately not in this PR
generation.py:482(monitoring silently no-ops without an engine) andgeneration.py:294(short-circuitorstrips one of two kwargs) both need an intent decision about which behaviour was meant, and neither corrupts data.docs/integration-api-v1.md:137still describesMonitoringConfigas single-field.🤖 Generated with Claude Code