Skip to content

Add bounded LLM worker cancellation - #22070

Merged
seyeong-han merged 3 commits into
pytorch:mainfrom
seyeong-han:llm-server/worker-cancellation
Aug 27, 2026
Merged

seyeong-han merged 3 commits into
pytorch:mainfrom
seyeong-han:llm-server/worker-cancellation

Conversation

@seyeong-han

Copy link
Copy Markdown
Contributor

Summary

Add bounded cancellation for the existing process-isolated LLM worker protocol.

  • Negotiate supports_cancel and pass monotonically increasing request IDs over JSONL.
  • Deliver cancellation out of band through an inherited POSIX pipe, so stop() never waits behind the blocking request/response lock.
  • Cooperatively call LLMSession::stop() once at the next decode boundary. Cancelled sessions are marked dirty and reset before reuse.
  • If cooperative cancellation misses its grace period, terminate/kill/reap the worker, fail queued and future work, and report /health as unavailable. Model reload remains the supervisor's responsibility.
  • Keep older and non-POSIX workers compatible; they do not advertise cancellation and use the bounded process-termination fallback.

Review guide:

  1. Python transport and process lifecycle: examples/llm_server/python/worker_client.py.
  2. Async cancellation, escalation, and health: examples/llm_server/python/session_runtime.py, serving_chat.py, and server.py.
  3. Native request controller and dirty-on-cancel semantics: examples/llm_server/cpp/worker_loop.h.
  4. MuseGlimmer custom-loop integration: examples/models/muse-glimmer/runtime/runners/muse_glimmer_worker.cpp.
  5. The remaining changes are focused tests, documentation, and thread linkage for worker targets that include the shared controller.

Cancellation remains token-boundary cooperative: it cannot interrupt active prefill, vision preparation, or an in-progress backend invocation. The process fallback provides the bound in those cases.

Test plan

  • pytest -q examples/llm_server/python/tests (229 passed, 6 skipped)
  • Built and ran test_worker_loop and test_worker_prefill_plan (2/2 CTests passed)
  • Built and linked the MLX muse_glimmer_worker, qwen3_5_moe_worker, and gemma4_31b_worker targets
  • Ran the real loopback Uvicorn/socket disconnect test five consecutive times
  • Ran exact-file Python, C++, CMake, and documentation linters

@pytorch-bot

pytorch-bot Bot commented Aug 24, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22070

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit ca6c74a with merge base cff6f4d (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 24, 2026
@github-actions

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

@metascroy

Copy link
Copy Markdown
Contributor

@claude review this code

@claude

claude Bot commented Aug 25, 2026

Copy link
Copy Markdown

Claude finished @metascroy's task in 9m 7s —— View job


Code Review: Add bounded LLM worker cancellation

  • Read the diff and changed files
  • Review Python transport / process lifecycle (worker_client.py)
  • Review async cancellation / escalation / health (session_runtime.py, serving_chat.py, server.py)
  • Review native request controller (worker_loop.h)
  • Review MuseGlimmer integration + CMake changes
  • Post review feedback

Overall this is a careful, well-documented design. The out-of-band POSIX pipe keeps stop() off the blocking JSONL lock, the request-id monotonicity + last_completed gating is sound, the frame reconstruction endianness matches to_bytes(8, "little"), and the additive supports_cancel negotiation keeps legacy/non-POSIX workers working. Comments are excellent. Test coverage (unit + real loopback socket) is strong. A few things worth confirming before merge, none of them blockers.


Observations

1. Client-side stop-string detection now delivers a real worker cancellation (behavior change worth confirming)

serving_chat.py:189 (_collect_until_stop) and serving_chat.py:607 (on_stop) call self._runtime.stop(). Before this PR WorkerClient.stop() was a no-op; now it writes a cancel frame for the active request. When that frame lands before the worker's own stop-string handling completes, the terminal {"done"} carries "cancelled": true and the session is marked dirty (worker_loop.h:778-781).

The finish_reason is already "stop" in this path so user-visible completion is unaffected, and the worker already dirties string-stop turns — so no warm-resume regression. But GenStats.cancelled will now read True for what is semantically an ordinary request-stop completion. Please confirm nothing downstream treats cancelled as "the client disconnected" (metrics, logging, transcript retention). If cancelled is meant strictly for out-of-band cancellation, a stop-string hit shouldn't set it.

2. stop() caches False permanently on EAGAIN for the active requestworker_client.py:339-341

except BlockingIOError:
    self._cancel_delivery = False
    return False

_cancel_delivery is cached, so a subsequent stop() for the same request returns the cached False and never retries the write. The comment documents this as intentional ("per-request delivery failure"), and with a 64KB pipe buffer vs a single 8-byte frame a spurious EAGAIN is effectively impossible — so this is fine in practice. Flagging only so the intent is explicit: escalation (_cancel_generation) is the recovery path, not a retry.

3. Dropped preactivation frame if a cancel arrives for a not-yet-active request while another is activeworker_loop.h:400-405

if (active_session_ != nullptr) {
  if (request_id == active_request_id_) { cancel_active_locked(); }
  return; // conflicting active IDs are never queued for a later request
}

A frame for request N+1 arriving while N is still active is silently dropped rather than parked in pending_request_id_. This is safe given the Python side serializes (reserve happens under the asyncio lock and the prior request's generate() has returned before the next id is reserved), so the window can't occur with the real client. Worth a one-line note in the comment that this relies on the caller's single-in-flight discipline, since a future concurrent driver would silently lose cancellations here.


Minor / nits

  • session_runtime.py:220-221: healthy handles both a callable and a bool for worker.healthy to accommodate test doubles vs the real property. Slightly surprising; a short comment on why both shapes are accepted would help the next reader.
  • worker_loop.h:585-588: the stop strings are read but not length/count-bounded here. The control plane is the source, so likely fine, but a malicious/buggy caller could pass an unbounded list. Low risk given the trust boundary.
  • CMake: adding find_package(Threads REQUIRED) + Threads::Threads to the four worker targets and the extension_llm_session fallback for older installed packages both look correct and appropriately scoped.

Nice work — the cancellation state machine and its seal/dirty invariants are the tricky part and they read as correct to me. Happy to dig deeper into any specific path if useful.
llm-server/worker-cancellation

@seyeong-han
seyeong-han merged commit 5bd86e5 into pytorch:main Aug 27, 2026
210 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants