interleaver: don't let a parked worker or an early stop escape cleanup

[JadenFiotto-Kaufman]

Two findings from the agent stress sweep, items 2 and 3 in planning/next-steps.md. Same shape of problem — something escaping before the cleanup around it finishes — but independent fixes, so one commit each.

1. cancel() leaked a parked worker

Broken: envoy.interleave calls check_dangling_mediators() only on the success path; when the model's forward raises it goes straight to finally: cancel(), which set mediator.worker = None without unwinding. Dropping a reference does not end a greenlet, so the worker stayed parked, holding its frame, the block's Scope, and through it the model. A gc pass does not reach it either — a suspended greenlet's frames are invisible to the cyclic collector. Roughly 0.5 GiB of CUDA held per errored trace, cumulative (repro: tasks/skip-stop-scan-corners/noskills/repro_8.py).

Changed: cancel() throws GreenletExit into any worker still alive, which unwinds it and runs the block's finally blocks. An exception out of one of those warns rather than propagating, so it cannot mask the model error already in flight (check_dangling_mediators is the precedent for telling the two apart). check_dangling_mediators is deliberately still not called on the error path — it raises OutOfOrderError, which would bury the real error.

cancel() is the single choke point for envoy.interleave, tracer.py and the vLLM tracer. On the vLLM client side the workers are serialized onto requests and never started, so alive is False there and nothing changes.

2. handle() abandoned the visit on tracer.stop()

Broken: EarlyStopException subclasses Exception and defer_exceptions is False on a local trace, so the raise in the mediator loop left Interleaver.handle immediately — skipping the rest of ready, the count update, assemble_skip, the observers loop that feeds tracer.cache, and the fragment undo. Two symptoms from one cause:

• stopping at layer 5 with a cache left model.transformer.h.5 out of the cache;
• a stop in one invoke of a batched generate cost every sibling invoke that step, so a stop after 3 steps gave len(picks) == 3 but len(sibling) == 2.

Changed: an early stop is treated as control flow — caught, kept in a local, that mediator's pending cleared, the visit finished, and the exception re-raised just before the return (after undo). Other exceptions keep today's behaviour, and defer_exceptions still records a stop on its mediator rather than raising it, so the vLLM path is unchanged. The loop terminates because a mediator with no pending is not ready.

The three deliberate calls are now stated in docs/usage/stop-and-early-exit.md (two new sections plus two gotcha bullets): the module you stop at is still recorded and served; the workers parked on the same visit are still served; and a stop in one invoke ends the whole batched forward, since it is one shared run. docs/developing/interleaver-internals.md had two sentences describing the old behaviour of cancel and of stop, so both were corrected — the only file I touched outside my assignment, and it is the developer page for the function I changed.

Tested

Environment: nnsight-stress/env with PYTHONPATH at this worktree's src. The optional _c/py_mount extension is not built there, so .save() as a method fails and about half the suite errors; I compiled it into the worktree (gitignored) to get a real run.

New tests, all five verified failing against origin/0.8 and passing here:

• tests/test_interleaving.py — cancel unwinds a still-parked worker (its finally runs); cancel warns rather than raises when the unwind hits a user finally; an early stop finishes the visit (the sibling worker is served and the count is bumped) before it raises.
• tests/test_memory.py — a model that raises mid-forward frees the module afterwards; five such traces do not accumulate. CPU-only, adapted from repro_8.py, which needed CUDA.
• tests/test_language.py::TestEarlyStop — the module stopped at is still in the cache; a stop in one invoke leaves both invokes with the same step count.

Ran clean, per file: test_interleaving, test_memory, test_language, test_batching, test_tracing, test_saving, test_editing, test_backward, test_envoy, test_source, test_fragments, test_modeling, test_vlm, test_diffusion. Both new doc snippets were run verbatim and produce what the page claims. (Running several test files in one pytest invocation errors on transformers.pipeline for unrelated, pre-existing reasons, as does test_serialization on its own — both reproduce identically on origin/0.8.)

Not tested / left alone

• vLLM is not installed in this environment, so tests/vllm/ could not run. The defer_exceptions branch is unchanged by construction: a stop under deferral is still recorded on its mediator and never raised out of the handoff.
• The fragment undo path (tensor parallel) is picked up by the stop fix — undo now runs before the re-raise — but there is no multi-device setup here to exercise it. Per the item note, it is not skipped, just untested.
• No check_dangling_mediators call was added to the error path, per the decision note.

🤖 Generated with Claude Code