feat(client): decode opaque routed expert payloads - #2408
Open
biswapanda wants to merge 12 commits into
Open
Conversation
Cancelling the EnvClient's awaiting task previously orphaned the server-side rollout: ZMQ multiplexes every request over one socket, so abandonment is invisible without an explicit message. On CancelledError for a run, the client now fire-and-forgets a 'cancel' carrying the run's request_id; the pool broker routes it to the worker holding that run (answering unknown/finished targets inline), and the worker cancels the rollout's asyncio task. An aborted run still replies (error='Cancelled') so broker and client accounting stay exact. Callers keep implicit semantics: cancelling the client task is all it takes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- Shield sandbox creation through the id capture: a cancel that aborts the create POST mid-flight left the platform building a sandbox this side never learned the id of, so teardown could not delete it and the sandbox leaked to the reconciler. Youngest-first overload sheds cancel exactly the episodes still provisioning, making this common. - Hold strong references to fire-and-forget cancel tasks: the loop only keeps weak refs, so an unreferenced task could be garbage collected before the cancel was ever sent.
close() could tear down the ZMQ socket before a scheduled fire-and-forget cancel ran, leaving the server-side rollout running. Also apply repo formatting to the touched files.
- Register run tasks in the dispatch loop, not the handler: a cancel scheduled ahead of the run's handler task found no entry and the rollout kept running. ZMQ preserves per-client frame order, so the run frame always lands first and dispatch-time registration is race free. - Tolerate unparseable cancel payloads in the pool broker: the unpack ran on the broker loop with no error handling, so one bad frame tore down every in-flight rollout behind it. - Drain cancel tasks after the receiver stops and loop until empty: a run cancelled in the same shutdown wave could schedule its cancel after a single early drain and lose it to the socket close.
# Conflicts: # verifiers/v1/runtimes/prime.py
biswapanda
changed the title
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Summary
cache_saltfrom sampling extras to the native Generate API fieldDynamo and vLLM do not need to know Prime-RL's routed-expert representation. Verifiers owns interpretation because it consumes the data for trajectory attribution and router replay.
Validation
uv run pytest tests/v1/test_train_client.py tests/v1/test_graph.py -q: 11 passeduv run ruff formatanduv run ruff checkpassed on touched files_decode_routed_expertsconsumed the opaque wire value and returned shape[14, 2, 2], dtypeuint8, prompt start1, and experts[0,1,2,3,4,6,7]Routing Conf. 0.3212and0.3233on steps 1 and 2PRIME_ROUTED_EXPERTS_PATH_EXECUTED, request977f27f8-dd7d-411e-b98a-c5cf4971e22f,OpaquePayload.bytes=184The integration run used vLLM #52723 with merged #53204, Dynamo #13531 plus its existing request/control and worker-metadata dependencies, and Prime-RL #3296.