Repository navigation
Conversation
KastanDay
added
the
100% truly worth submitting upstream to SGLang
|
@codex review |
|
Codex Review: Didn't find any major issues. You're on a roll. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a6d1b36734
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if status != KVPoll.Bootstrapping: | ||
| return |
There was a problem hiding this comment.
When Mori transfer metadata arrives before the prefill scheduler creates MoriKVSender, _handle_transfer_message() intentionally accepts the absent-room case and calls update_status(..., KVPoll.WaitingForInput). This new guard discards that readiness update; the sender later creates the room as Bootstrapping, but the already-received metadata is never replayed, so MoriKVSender.poll() remains bootstrapping until it times out. Preserve this supported ordering by applying the stored readiness when the sender creates the room, while still preventing post-clear stale updates.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 98ece61104
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if status != KVPoll.Bootstrapping: | ||
| return |
There was a problem hiding this comment.
When destination metadata arrives before the prefill scheduler constructs its sender, this now discards the WaitingForInput transition. The Mooncake and NIXL bootstrap threads issue exactly that transition at mooncake/conn.py:1598 and nixl/conn.py:2450, while their sender constructors do not re-evaluate already-complete metadata after creating the room; only Mori received such a recheck. Consequently, metadata-first requests remain Bootstrapping and eventually time out despite having all destination information, so those backends need the same deferred-readiness handling before non-initial transitions are rejected.
Useful? React with 👍 / 👎.
| if ( | ||
| self.request_status[bootstrap_room] == KVPoll.Failed | ||
| or status == KVPoll.Failed | ||
| ): | ||
| self.request_status[bootstrap_room] = KVPoll.Failed |
There was a problem hiding this comment.
When a completion callback overlaps a failure callback from another worker, heartbeat, or abort thread, this compound transition is not atomic: the completion can observe a non-failed state in this condition, the failure can then store Failed, and the completion can subsequently execute the max(...) branch and write Success. That preserves the KV-corruption race this change is intended to close under concurrent delivery; guard the full read/decision/write transition with a shared lock or otherwise make it atomic.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cf13c6c552
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if ready_rooms is None: | ||
| ready_rooms = set() | ||
| self._metadata_ready_rooms = ready_rooms | ||
| ready_rooms.add(bootstrap_room) |
There was a problem hiding this comment.
When destination metadata arrives after clear_status() but before the same bootstrap_room is reused, this unconditionally recreates the readiness marker even though update_status() correctly ignores the absent room. The next CommonKVSender for that ID calls _mark_room_ready_if_metadata_complete() and immediately advances to WaitingForInput, allowing Mooncake/NIXL/Mori to use metadata and destination pointers belonging to the previous request; repeated room IDs are supported and some benchmark clients always use room 0. Associate readiness with a room generation or retain a tombstone so delayed metadata cannot authorize the next request.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 486eb7e65d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| tombstones = getattr(self, "_cleared_room_tombstones", None) | ||
| if tombstones is None: | ||
| tombstones = set() | ||
| self._cleared_room_tombstones = tombstones | ||
| tombstones.add(bootstrap_room) |
There was a problem hiding this comment.
Every sender or receiver cleanup permanently adds its room ID to _cleared_room_tombstones, while IDs are removed only if that exact room is reused. In normal deployments where requests receive unique or monotonically increasing room IDs, this set therefore grows by one entry per completed or aborted request for the lifetime of each manager, eventually consuming unbounded host memory on long-running, high-throughput servers. Cleared-generation protection needs bounded storage or explicit retirement.
Useful? React with 👍 / 👎.
| if bootstrap_room in getattr( | ||
| self, "_cleared_room_tombstones", set() | ||
| ): | ||
| return |
There was a problem hiding this comment.
When a caller reuses a cleared bootstrap_room and the new decode metadata arrives before the new prefill sender is constructed, this tombstone check discards the new generation's readiness notification. The sender subsequently removes the tombstone while entering Bootstrapping, but _metadata_ready_rooms has no entry to advance it to WaitingForInput, so the request remains bootstrapping until it times out even though its new metadata is already present in transfer_infos. This ordering is specifically supported for first-use rooms and must also work after room reuse.
Useful? React with 👍 / 👎.
|
|
||
| def clear(self) -> None: | ||
| self.kv_mgr.request_status.pop(self.bootstrap_room, None) | ||
| self.kv_mgr.clear_status(self.bootstrap_room) |
There was a problem hiding this comment.
This new method call breaks the pre-existing registered CPU test test_nixl_sender_failure_cleanup.py, whose lightweight kv_mgr is a SimpleNamespace without clear_status. Running PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation now fails here with AttributeError instead of verifying that NixlKVSender.failure_exception() cleans state and raises the recorded transfer error, so the registered unit suite cannot pass until the fixture or cleanup compatibility is updated.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a94801f1d6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if status not in ( | ||
| KVPoll.Bootstrapping, | ||
| KVPoll.WaitingForInput, | ||
| ): |
There was a problem hiding this comment.
When a delayed Mori or NIXL metadata message from an old generation arrives after CommonKVSender.clear(), its handler calls update_status(..., KVPoll.WaitingForInput) while the room is absent, and this allowlist recreates the status entry while the handler also repopulates transfer_infos. If a client later reuses that bootstrap_room, the sender's Bootstrapping update remains at WaitingForInput and it can consume the previous request's destination metadata, defeating cleanup and potentially transferring KV into stale pages. Metadata-before-sender must remain supported, but readiness needs generation or tombstone correlation so readiness arriving after a clear cannot initialize the next generation.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6897a882a3
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| mgr=self.kv_manager, | ||
| bootstrap_addr=_bootstrap_addr(req), | ||
| bootstrap_room=req.bootstrap_room, | ||
| request_id=req.rid, |
There was a problem hiding this comment.
When a normal OpenAI request omits the optional rid, the model gateway injects only the bootstrap host/port/room (sgl-model-gateway/src/routers/http/pd_router.rs:232-285) and dispatches the same rid-less request to both workers; each worker then independently generates a UUID in GenerateReqInput._normalize_single_inputs() (python/sglang/srt/managers/io_struct.py:427-432). Consequently, this decode-side ID differs from the prefill-side req.rid passed in prefill.py:274, so mark_metadata_ready() rejects the metadata for the active generation and the request remains bootstrapping until timeout. Generate a shared ID in the router or base generation matching on a value guaranteed to be common to both workers.
Useful? React with 👍 / 👎.
| if ( | ||
| self.request_status[bootstrap_room] == KVPoll.Failed | ||
| or status == KVPoll.Failed | ||
| ): | ||
| self.request_status[bootstrap_room] = KVPoll.Failed |
There was a problem hiding this comment.
When a room is cleared, reused, and then receives a late completion or failure callback from the old transfer—the ordering this change is intended to handle—update_status() sees that the room exists again and applies the stale status to the new request because callers provide only the room, not its request ID. A stale failure therefore makes the new generation terminally failed, while a stale success can advance it prematurely; the new _room_request_ids guard protects clear_status() and metadata only. Status callbacks need to carry and validate the generation before mutating the active room.
Useful? React with 👍 / 👎.
| request_ids = { | ||
| info.request_id for info in self.transfer_infos[room].values() | ||
| } | ||
| if ( | ||
| len(self.transfer_infos[room]) == required_dst_info_num | ||
| and len(request_ids) == 1 |
There was a problem hiding this comment.
When metadata from a cleared generation arrives before a reused room's sender is created, this handler stores it in transfer_infos even though mark_metadata_ready() only records its request ID separately. If the new generation then reports metadata under new agent names, the stale entry makes the map larger than required_dst_info_num and this equality can never pass; with a reused agent name, a late stale packet can instead overwrite a new destination after the room is already ready because the handler does not check the active status or ID before assignment. The transfer-info table must be partitioned or purged by request ID, not merely checked for ID consensus after mutation.
Useful? React with 👍 / 👎.
| if not clear_status(self.bootstrap_room, self.request_id): | ||
| return |
There was a problem hiding this comment.
When an old receiver calls clear() after its room has been reused, this generation check correctly leaves the new status intact but returns only from CommonKVReceiver.clear(). MoriKVReceiver.clear() (mori/conn.py:1780-1784) then continues into _cleanup_room_tracking() and removes the new receiver's room_to_bootstrap_addr and addr_to_rooms_tracker entries; a subsequent prefill heartbeat failure therefore no longer marks the active request failed and it remains stuck until its waiting timeout. Return the clear result to subclasses or make the Mori cleanup conditional on the same generation check.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6b6feb523a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| is_last_chunk=is_last_chunk, | ||
| prefill_aux_index=aux_index, | ||
| state_indices=state_indices, | ||
| request_id=request_id, |
There was a problem hiding this comment.
When staging is enabled and a bootstrap room is reused after an abort or clear, this generation ID is carried by the queued chunk but _send_chunk_ready() still sends only the room and chunk metadata. A delayed CHUNK_READY from the retired generation is therefore accepted by the decode thread and routed to the new room's DecodeRequest; once enough old writers arrive, handle_chunk_arrived() scatters and releases the new request's staging allocation before its own transfer completes, which can commit stale or incomplete KV data. Include and validate the request generation in the CHUNK_READY protocol, as is already done for terminal status messages.
Useful? React with 👍 / 👎.
| @@ -271,6 +271,7 @@ def create_sender(self, req: Req, num_kv_heads: int) -> bool: | |||
| dest_tp_ranks=dest_tp_ranks, | |||
| pp_rank=self.pp_rank, | |||
| req_has_disagg_prefill_dp_rank=req.disagg_prefill_dp_rank is not None, | |||
There was a problem hiding this comment.
When a request without an explicit rid uses the supported Python --mini-lb path, mini_lb.py forwards copies of the same payload to prefill and decode but only injects bootstrap fields; each server therefore generates an independent UUID for req.rid. Making that UUID the generation key here means decode metadata is stored under the decode UUID while the prefill room is active under a different UUID, so store_transfer_info() rejects or partitions the metadata and the request eventually times out in bootstrapping. The Rust PD router's new rid injection does not cover MiniLB, so every PD dispatcher must inject one shared ID before dual dispatch (or the protocol must use a field already shared by both requests).
Useful? React with 👍 / 👎.
| retired = getattr(self, "_retired_room_generations", {}) | ||
| if request_id and (bootstrap_room, request_id) in retired: | ||
| return None |
There was a problem hiding this comment.
During PD true-retraction rebootstrap, build_rebootstrap_payload() intentionally reuses the original rid and bootstrap_room, and the decode path sends transfer metadata before submitting the HTTP prefill recompute. The original prefill sender has already cleared and retired exactly this tuple, so if the ZMQ metadata arrives before the recompute request constructs the replacement sender, this branch discards the only metadata message as stale; the replacement sender then remains in Bootstrapping until timeout. Rebootstrap needs a fresh generation token, or the tuple must be explicitly reopened before metadata is sent rather than treating every reuse of the same request ID as retired.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bddfc1f901
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| request_id = _generation_request_id(req, is_rebootstrap) | ||
| kv_receiver = kv_receiver_class( | ||
| mgr=self.kv_manager, | ||
| bootstrap_addr=_bootstrap_addr(req), | ||
| bootstrap_room=req.bootstrap_room, | ||
| request_id=request_id, |
There was a problem hiding this comment.
When n > 1, the tokenizer manager independently replaces each sample's RID with random UUIDs on the prefill and decode servers (tokenizer_manager.py:1655 and :1674). Using that RID here as the transfer generation therefore gives the two sides different request_id values, so the prefill rejects the decode metadata as belonging to another generation and the request eventually times out. This affects parallel-sampling requests even through the updated built-in routers; derive the generation from shared deterministic data such as the bootstrap room, or make the per-sample RID expansion deterministic across both workers.
Useful? React with 👍 / 👎.
| return bootstrap_room in self.request_status and ( | ||
| active_request_id == request_id | ||
| or (active_request_id is None and request_id == "") | ||
| ) |
There was a problem hiding this comment.
In a mixed-version prefill/decode deployment, the older peer omits the newly added request-ID frame, which all of the updated parsers explicitly accept as None; however, this strict comparison normalizes that to "" and rejects it whenever the new peer's active generation has a router-injected RID. A new prefill therefore discards metadata from an old decode, while a new decode discards completion notifications from an old prefill, causing every in-flight request across the version boundary to time out. Either negotiate generation support or retain a safe legacy fallback for messages without the optional frame.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ce1a4761eb
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| mgr.agent = NotificationFakeAgent(messages) | ||
| rooms = {int(message.split("_", 1)[0]) for message in messages} | ||
| mgr.request_status = {room: KVPoll.WaitingForInput for room in rooms} | ||
| mgr._room_request_ids = {} |
There was a problem hiding this comment.
Running PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_nixl_backend_basic.py now fails the existing untagged notification cases because this empty mapping makes is_current_generation(room, None) compare None with the normalized empty request ID and reject every message at NixlKVManager.update_transfer_status(). Populate each fixture room with the legacy empty ID (for example, {room: "" for room in rooms}) before overriding room 9 in the generation-specific test; otherwise the registered suite fails for KV, auxiliary, state, and staging notifications.
Useful? React with 👍 / 👎.
| obj.insert( | ||
| Self::RID_KEY.to_string(), | ||
| Value::Array(rids.into_iter().map(Value::from).collect()), |
There was a problem hiding this comment.
For /v1/chat/completions with n > 1 and no caller RID, get_chat_batch_size() returns n, so this inserts an n-element RID array even though SGLang models the request as one base prompt with parallel_sample_num=n. GenerateReqInput._normalize_rid() requires an incoming RID list to have the base batch_size (1 here), and consequently raises ValueError: The specified rids length mismatch with the batch_size for batch processing. before generation. Distinguish parallel sampling from a true input batch and inject one shared scalar RID for this case.
Useful? React with 👍 / 👎.
| digest = hashlib.blake2b( | ||
| f"{base_room}:{base_rid}:{sample_index}".encode("utf-8"), | ||
| digest_size=8, | ||
| person=b"sglang-pd", | ||
| ).digest() | ||
| return int.from_bytes(digest, "big") & ((1 << 63) - 1) |
There was a problem hiding this comment.
When PD workers use follow_bootstrap_room with DP size greater than one and no external DP-rank override, DataParallelController.follow_bootstrap_room_scheduler() selects the worker using bootstrap_room % len(workers). The common-prefix warmup remains on base_room while this unconstrained hash sends most samples to different DP workers, whose independent radix caches do not contain that warmed prefix; parallel sampling therefore recomputes the prompt on roughly all but 1 / dp_size samples and defeats the optimization this path performs. Derive distinct rooms that retain the base room's DP affinity.
Useful? React with 👍 / 👎.
| if ( | ||
| room_to_be_aborted in self.request_status | ||
| self.is_current_generation(room_to_be_aborted, request_id) | ||
| and self.check_status(room_to_be_aborted) != KVPoll.Success |
There was a problem hiding this comment.
During a rolling upgrade, an old decode worker sends the four-frame ABORT format, so request_id is None. If it aborts during handshake before sending transfer metadata, the new prefill manager has not yet called allow_legacy_generation(), causing this generation check to reject the abort even though the room is active; the prefill request and its tree-cache ownership then remain until the bootstrap timeout. The Mooncake handler has the same ordering problem, so legacy capability must be established before metadata arrives or untagged early aborts need a safe compatibility path.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a64d1d6967
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| sample_room = _parallel_sample_bootstrap_room( | ||
| objs[i].bootstrap_room, | ||
| objs[i].rid, | ||
| sample_index, | ||
| getattr(self.server_args, "dp_size", 1), |
There was a problem hiding this comment.
When the prefill and decode deployments have different dp_size values, each TokenizerManager calls this helper with its own local size, so the same base request and sample index almost always produce different bootstrap rooms on the two sides. Heterogeneous prefill/decode DP configurations are supported by the routing and bootstrap code, but parallel-sampling requests in such configurations will now leave every sender and receiver waiting in different rooms until timeout; the shared room must be derived from data that is identical on both workers (or from the prefill topology explicitly communicated to both).
Useful? React with 👍 / 👎.
| @@ -271,6 +271,7 @@ def create_sender(self, req: Req, num_kv_heads: int) -> bool: | |||
| dest_tp_ranks=dest_tp_ranks, | |||
| pp_rank=self.pp_rank, | |||
| req_has_disagg_prefill_dp_rank=req.disagg_prefill_dp_rank is not None, | |||
There was a problem hiding this comment.
For requests routed through experimental/sgl-router, this new generation tag differs between prefill and decode: that router's build_outgoing_body injects only bootstrap_host, bootstrap_port, and bootstrap_room (chat.rs:723-738), while each engine independently assigns a random UUID when rid is absent (io_struct.py:431-432). Consequently the receiver sends one RID and this sender registers another, so the prefill side rejects the transfer metadata as a stale generation and ordinary OpenAI requests time out. The experimental router must inject a shared RID into the body just like the two routers updated in this commit.
Useful? React with 👍 / 👎.
a64d1d6 to
2c57835
Compare
|
Superseded by PR 30, which carries the corruption-relevant core of this change (absorbing Failed status plus a locked status read-modify-write) without the room-generation tagging protocol — the tagging here silently drops untagged abort signals from older peers, which is a fail-open regression, and the branch now conflicts with main. |
1 participant
Why
CommonKVManager.update_status()uses the numeric maximum to advance a transfer room through its lifecycle.KVPoll.Failedis0, so a room that has already failed can be promoted back toTransferringorSuccessby a late completion update.That makes transfer outcome depend on message ordering. Once any participant reports failure, later work for the same room must not make the request eligible for decode.
KV cache corruption mechanism
This race can produce incorrect model output rather than only an error or restart:
KVPoll.Failed.TransferringorSuccess;max(Failed, Success)changes the room back toSuccess.Successas permission to commit the transfer and begin generation. The metadata readiness gate confirms that request metadata arrived, but it does not prove that every KV or model-state page was transferred successfully.The fix is fail-closed: once failure is observed for a room, late completion cannot authorize decode.
What changed
KVPoll.Failedabsorbing when either the current room state or the incoming update is failed.TransferringandSuccessupdates, explicit room reuse, and late failure after cleanup.Scope
This PR enforces the status-state invariant after a failure has been reported. It does not detect transfer failures or wait for native transfer work to stop before pages are released.
The changes are independent and complementary: failure detection, terminal failure state, and transfer quiescence close different corruption paths.
Validation
The registered regression test covers the state transitions directly:
Current-main validation completed with
python3 -m py_compileandgit diff --check. Local pytest execution is blocked on macOS by the Linux/CUDA-onlytritonimport path; GitHub currently reports no checks for this branch.CI States
Latest PR Test (Base): ❌ Missing
run-cilabel -- add it to run CI tests.Latest PR Test (Extra): ❌ Blocked --
run-ciis required first.