engine: bound shutdown drain for mlx generate (#46)

[tlkahn]

Summary

Closes #46 (Stage F slice only; #57 stays open).

Make the engine-thread non-interruptible window under cancel/shutdown explicit, polled, and proven:

1. Prefill remains chunked (MLXD_PREFILL_CHUNK=512) with shutdown/cancel polled before every chunk.
2. Decode polls before every eval step.
3. Seed forward (last prompt token) also polls before running.
4. Residual bound documented and accepted: at most one prefill chunk or one pipelined decode step.
5. GPU tests assert engine release within a 2000 ms drain-style deadline for mid-prefill and mid-decode cancel/shutdown.

Cycle 0 audit

Item	Finding
Prefill/decode polls	Already present at loop tops in handle_generate_real
Seed forward poll	Missing - added
Multi-chunk at chunk=512	Impossible with max_pos=512 - chose R2b (bump tiny_qwen3 max_pos to 2048)
RoPE tables	fwd_rope_freqs_build is head_dim-based, not max_pos-length - safe for R2b
HTTP drain_deadline_ms	Unchanged (bounds socket closure only); engine tests use stricter 2000 ms

Production changes

• engine_gen_hooks_t + engine_set_gen_hooks (test/support; real path only; stub ignores)
• Ordering at each site: poll flags -> hook -> re-poll flags -> work
  • Re-poll ensures a cancel that arrives while a hook is blocked does not run the next mlx op
• Seed-forward cancel/shutdown poll + on_before_seed hook
• Residual-window notes in engine.h and CLAUDE.md

Tests

CPU (test_engine.c):

• test_gen_hooks_default_null
• test_gen_hooks_stub_ignores

GPU (test_engine_gpu.c):

• test_prefill_hook_fires_once_short_prompt
• test_decode_hook_fires_per_step
• test_cancel_mid_decode_within_deadline
• test_shutdown_mid_decode_within_deadline
• test_cancel_mid_prefill_within_deadline
• test_shutdown_mid_prefill_within_deadline
• test_cancel_between_prefill_and_seed
• test_prefill_multi_chunk_hook_positions
• test_cancel_after_first_prefill_chunk
• test_destroy_joins_within_deadline_mid_prefill

Existing B3.11 cancel/shutdown tests retained.

Fixture

• tiny_qwen3{,_sharded,_tied} max_position_embeddings: 512 -> 2048
• tools/gen_tiny_ckpt.c QWEN3 recipe updated so regen stays consistent
• test_oversized_prompt expects 2049/2048

Test plan

• make test (44 passed)
• make tests/test_engine_gpu && ./tests/test_engine_gpu (all engine: bound shutdown drain for non-interruptible mlx ops #46 cases green)
• CI green

[lgsunnyvale]

Code review: PR #125 — engine: bound shutdown drain for mlx generate (#46)

Overall: approve with nits. This is a clean Stage-F slice of #46. The production change is small and well-targeted, the poll contract is explicit, and the GPU test matrix actually pins the cancel/shutdown sites that matter. I do not see a blocking correctness bug.

What works well

• Seed gap closed. On main, prefill/decode already polled; seed forward did not. Cancel arriving in the prefill→decode handoff could still sit inside an unpollable model_forward. The new poll + on_before_seed site closes that.
• poll → hook → re-poll → work is the right contract. Re-poll is load-bearing for barrier-style tests and for any future support hook that blocks. Without it, cancel during a blocked hook would still enter the next mlx op.
• R2b fixture bump is justified. With max_pos == chunk == 512, multi-chunk prefill is unreachable. fwd_rope_freqs_build is head-dim (or null for non-llama3), not max_pos-length, so 512→2048 is safe. Regen recipe updated in lockstep.
• Residual honesty. Separating HTTP drain_deadline_ms (socket closure) from engine-side release, and documenting the residual grain + second-SIGINT escape, matches the real constraint (mlx ops are not preemptible inside Metal).
• Test design. Deterministic condvar barriers beat sleep-and-hope. Coverage across prefill / seed / decode × cancel / shutdown, multi-chunk positions, and destroy-join is the right matrix for this claim.

Findings

1. Residual-window docs omit seed (nit / docs)

CLAUDE.md correctly says polls cover prefill, decode, and seed, but both residual-bound sentences still say only:

at most one chunk or one pipelined decode step

Seed is a third non-interruptible site (1-token forward + logits). Smaller than a prefill chunk, but the bound statement should name it so the three poll sites and the three residual cases stay 1:1.

Suggest aligning engine.h and CLAUDE.md to something like:

residual non-interruptible window is at most one prefill chunk, one seed forward, or one pipelined decode step

2. Deadline tests prove control-plane latency, not mlx residual wall-clock (nit / tests)

MLXD_TEST_DRAIN_BOUND_MS (2000) is measured from cancel/shutdown while the engine is blocked in a hook, i.e. before the next mlx op. After release, re-poll takes the cancel path and almost no GPU work runs.

That correctly proves: "cancel observed at poll site → FINISH_CANCELLED + stream release is fast."

It does not prove: "one 512-token prefill chunk (or one pipelined decode step) on a real checkpoint finishes inside 2000 ms / HTTP drain." Wall-clock residual still scales with model size; second SIGINT remains necessary for that. The PR body is mostly honest about this — please add a one-line comment above MLXD_TEST_DRAIN_BOUND_MS so a future reader does not treat 2000 ms as a bound on residual mlx work.

3. Mid-decode cancel/shutdown should assert zero tokens (suggestion / tests)

Prefill and seed paths assert tokens == 0 + FINISH_CANCELLED. The decode barrier fires before sampling at step 0, so the re-poll contract also implies zero tokens — but test_cancel_mid_decode_within_deadline and test_shutdown_mid_decode_within_deadline only check sole-owner + cancelled done.

Please drain and assert tokens == 0 there too. Cheap, and it locks the re-poll semantics for decode the same way prefill is locked.

4. Missing shutdown twin for seed barrier (nit / tests)

Prefill and decode have cancel and shutdown variants. Seed only has test_cancel_between_prefill_and_seed. A test_shutdown_between_prefill_and_seed would complete the matrix; same shape as the cancel test with engine_signal_shutdown.

5. engine_set_gen_hooks visibility / tearing (nit / hardening)

Documented "not safe concurrent with in-flight generate" is fine for test-only use, and the post/dequeue mailbox mutex gives happens-before for the intended set-then-post pattern.

Two small hardenings worth considering (non-blocking):

• Snapshot engine_gen_hooks_t hooks = eng->gen_hooks; once at the start of handle_generate_real and invoke through the local copy. Makes the "stable for the duration of this generate" contract structural and avoids any torn function-pointer read if someone violates the rule later.
• Optionally assert/document that hooks must not call stream_cancel on the inflight stream either (only the test thread should), to keep the engine-thread contract simple — today the banned list is engine_post / engine_destroy / engine_signal_shutdown only.

6. Destroy-join test name vs claim (nit)

test_destroy_joins_within_deadline_mid_prefill requires a helper thread to barrier_release after 20 ms. It proves destroy+join completes once the hook unblocks and the cancel path runs — not that destroy can recover from a hook (or mlx op) that never returns. The comment in the test is OK; the name is slightly stronger than the proof. Not wrong, just easy to over-read.

Non-issues (checked, no action)

• RoPE / KV sizing under max_pos=2048: freqs are head-dim-based; KV grows with actual seq. Fixture-only config bump without weight rewrite is OK.
• test_http_gen_request.c still has a forced "513"/"512" context-length string — that is a mock injection, not fixture-derived.
• test_oversized_prompt and gen_tiny_ckpt.c QWEN3 recipe updated consistently across dense/tied/sharded.
• Stub path ignoring hooks is correct and CPU-tested.
• CI only runs make test (CPU); GPU coverage stays local/make test-gpu — pre-existing, not introduced here. Still the engine: bound shutdown drain for non-interruptible mlx ops #46 proof lives entirely in test_engine_gpu.c.

Verdict

Approve once findings 1–3 are addressed (1–2 can be comment-only; 3 is a few asserts). Finding 4 is optional symmetry. Finding 5 is optional hardening.

This is a good close for #46's Stage F verification slice; leaving the rest of #57 open is correct.

[tlkahn]

Thanks - agree with the Approve-with-nits calibration. Addressed on 545690a:

1. Residual docs name seed - engine.h and CLAUDE.md residual sentences now list prefill chunk, seed forward, and pipelined decode step (1:1 with the three poll sites).
2. MLXD_TEST_DRAIN_BOUND_MS honesty - comment above the define states it measures control-plane release after cancel/shutdown while blocked in a test hook, not residual mlx wall-clock on a real checkpoint.
3. Decode zero tokens - test_cancel_mid_decode_within_deadline and test_shutdown_mid_decode_within_deadline now drain with tokens == 0 + FINISH_CANCELLED, matching prefill/seed.
4. Shutdown seed twin - added test_shutdown_between_prefill_and_seed (completes cancel/shutdown x prefill/seed/decode).
   5a. Hooks snapshot - handle_generate_real copies eng->gen_hooks once at entry and invokes through the local copy. Covered by test_gen_hooks_stable_for_duration_of_generate (prefill clears hooks mid-flight; seed + decode still fire). Docs on engine_gen_hooks_t / engine_set_gen_hooks note the snapshot.

5b declined as a hard ban: stream_cancel only sets an atomic_bool and is intentionally safe from any thread, including the engine thread. The poll -> hook -> re-poll contract is exactly how a cancel that becomes visible during a blocked (or returning) hook is observed before the next mlx op. The existing hard bans (engine_post / engine_destroy / engine_signal_shutdown) exist because those re-enter the engine control plane from the engine thread. stream_cancel does not. Leaving the ban list as control-plane only; test-style guidance remains "stimulate cancel/shutdown from the test thread for deterministic barriers," not a production hook rule.

6: left test_destroy_joins_within_deadline_mid_prefill name as-is. The helper release is harness necessity (hook blocks forever until released); existing comment already covers that. Not wrong, just easy to over-read.

make test green; ./tests/test_engine_gpu green including the new/tightened cases above. Hard ban list still does not include stream_cancel.