Repository navigation
Conversation
Tighten the types in _REQ_TYPES_WITH_OPAQUE_FIELDS to precise msgspec-native annotations (or delete the dead one), then delete the registry and its _maybe_wrap_pickle branch so every IPC struct encodes natively on the msgpack path.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request tightens IPC struct types from opaque annotations to precise msgspec-native types, removing the _REQ_TYPES_WITH_OPAQUE_FIELDS registry, and adds a comprehensive round-trip test suite for msgpack IPC serialization. The feedback highlights a critical issue in weight_updater.py where msgspec.convert needs from_attributes=True to successfully convert Pydantic models, and suggests expanding the test suite to cover this conversion path.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
0fd53a1 to
74947dc
Compare
| for field in msgspec.structs.fields(type(obj)) | ||
| } | ||
|
|
||
| if dataclasses.is_dataclass(obj) and not isinstance(obj, type): |
There was a problem hiding this comment.
In which senario the object can be dataclass and enum data type?
There was a problem hiding this comment.
Good point. The dataclass branch seems to be needed as internal_state is a vars(ServerArgs) dump. Also cuda_graph_config is a CudaGraphConfig dataclass that can't encode natively without it.
Regarding the enum you are right. Initially I though we needed to handle disaggregation_mode: DisaggregationMode, that said after inspecting again it seems we store disaggregation_mode: str, so we don't need to handle enums at all.
| class DumperControlReqOutput(BaseReq, kw_only=True): | ||
| success: bool | ||
| # JSON-native per-worker response dicts. | ||
| response: List[Dict[str, Any]] |
There was a problem hiding this comment.
What is the real data type for the dictionary, will those type not able to be serialized to bytes?
There was a problem hiding this comment.
In this case, the body is the raw request.json() (afaik, the dumper method endpoint rejects it unless it is a dict). This value is then passed to dumper.configure(**body). Since it comes straight from json.loads, every value is already JSON-native and encodes fine over msgpack.
I think it should stay Dict[str, Any] because configure() takes open-ended kwargs with no fixed schema.
| class DumperControlReqInput(BaseReq, kw_only=True): | ||
| method: str | ||
| # JSON request body (guarded to be a dict at the /dumper endpoint). | ||
| body: Dict[str, Any] |
There was a problem hiding this comment.
What is the real data type for the dictionary, will those type not able to be serialized to bytes?
There was a problem hiding this comment.
response is what handle_request returns (typed -> list[dict]), e.g. get_state()'s config dict, or {} for configure/reset. It's built and re-serialized with json.dumps, so every value is JSON-native and msgpack-safe.
Like body, it stays Dict[str, Any] because the dumper response has no fixed schema.
| "element_size": weight_info["element_size"], | ||
| "byte_size": byte_size, | ||
| } | ||
| self.weight_pointer_map[name] = ExpertWeightPointer( |
There was a problem hiding this comment.
why should we drop the shape, numel, dtype element_size etc info here?
There was a problem hiding this comment.
I couldn't find any consumer that reads them in the codebase. For example, ExpertBackupClient.update_weights uses only weight_ptr and byte_size. Every other field is derivable: name comes from the map key, and shape/numel/dtype/element_size are re-derived from the local parameter tensor (we even assert numel * element_size == byte_size).
So those five were adding only weight on the wire. With one entry per expert weight, ExpertWeightPointer is array_like too, so each entry is just [weight_ptr, byte_size].
|
/tag-run-ci-label |
|
/tag-and-rerun-ci |
Resolve two conflicts introduced by upstream sgl-project#30525 (drop legacy get_loads IPC; move metrics structs to load_snapshot.py): - io_struct.py: apply both deletions — the PR removes SetInjectDumpMetadataReqInput while upstream removes the GetLoads* and metrics structs, so the region between BlockReqInput and SetInjectDumpMetadataReqOutput collapses to empty. - scheduler.py set_internal_state: keep the PR's `return SetInternalStateReqOutput(updated=if_success)` (the server_args field is gone from the struct) on top of upstream's get_server_args().override(...) rename.
|
/tag-and-rerun-ci |
|
/tag-and-rerun-ci |
3 participants
What this does
Closes Task 4 of #29465. The 13 IPC structs in
_REQ_TYPES_WITH_OPAQUE_FIELDShadDict[str, Any]/List[Any]fields, so on the msgpack path (SGLANG_USE_PICKLE_IPC=0)_maybe_wrap_picklewhole-struct-pickled them inside aPickleWrapperframe. This PR tightens eachfield to a precise msgspec-native type (or deletes a dead type), then removes the registry and its
_maybe_wrap_picklebranch. End state: the tuple no longer exists and every IPC struct encodesnatively.
Acceptance criterion (
_REQ_TYPES_WITH_OPAQUE_FIELDSis empty) is met — the symbol is deleted.Per-type changes
12 types tightened, 1 deleted:
UpdateWeightFromDiskReqInputmanifestdocumented as JSON-native; annotation unchangedVertexGenerateReqInputinstances/parametersdocumented as JSON-native; annotation unchangedDumperControlReqInputbodydocumented;isinstance(body, dict)guard added at the/dumperendpointDumperControlReqOutputresponsedocumented as a JSON-nativeList[Dict]GetWeightsByNameReqOutputparameter→Optional[List[Union[float, List[float]]]](element-level union)SetInternalStateReqserver_args→Dict[str, Union[int, float]]RpcReqInputparameters→Optional[Dict[str, Union[bool, int, float, str, None]]]LoadLoRAAdapterFromTensorsReqInputadded_tokens_config→Optional[Dict[str, int]];config_dictkeptDict[str, Any](see note)GetInternalStateReqOutputinternal_statekeptDict[str, Any]; producer sanitizes viamsgspec_to_builtinsSetInternalStateReqOutputserver_argsfield removed (only.updatedis read)BackupDramReqweight_pointer_map→Dict[str, ExpertWeightPointer](payload narrowed — see below)CheckWeightsReqOutputpayload→Optional[List[ChecksumInfo]](new mirror structs)SetInjectDumpMetadataReqInputSupporting changes:
msgspec_to_builtinsnow also converts dataclasses (→ dict), so the liveCudaGraphConfigdataclass in the
vars(ServerArgs)dump sanitizes consistently on both transports instead ofonly on the msgpack path.
get_internal_statealso dropscustom_sigquit_handlerfrom thedump — it can hold a bound signal handler and has no reader.
CheckWeightsReqOutput:ParallelismInfo/ChecksumInfomsgspec mirrors added inio_struct.py(the pydantic models in
weight_checker.pystay as the checker's source of truth). The producernormalizes the
tp>1all_gather_objectlist and thetp==1case to a uniformList[ChecksumInfo]; the consumer reads struct attributes and converts back to builtins, so the/weights_checkerHTTP response shape is unchanged.config_dictonLoadLoRAAdapterFromTensorsReqInputstaysDict[str, Any]: it is the PEFTadapter_config.json, already JSON on both paths, and a tighter type would only add decodestrictness on out-of-tree callers with no consumer benefit.
Deployment / compatibility
These structs cross every sglang zmq socket. Under
SGLANG_USE_PICKLE_IPC=0the wire format changes,so all nodes and processes must run the same version with the same setting (DP inter-node TCP
handshake, PD pairs, routers, elastic-EP).
SGLANG_USE_PICKLE_IPC=1is a kill-switch for framingonly.
Three types change the payload on the pickle path too, so the kill-switch restores framing but
not compatibility — revert is the mitigation for these:
SetInternalStateReqOutput— theserver_argsfield is gone; an old-code peer fails toreconstruct the output.
CheckWeightsReqOutput— the payload isChecksumInfostructs instead of dicts.BackupDramReq— the payload is narrowed (see below).BackupDramReqpayload narrowing:ExpertWeightPointercarries onlyweight_ptrandbyte_size.The producer previously sent five more keys per entry —
name,shape,numel,dtype,element_size— none of which the consumer reads (it derives those from the local parameter tensor,and gets the name from the map key). Because
BackupDramReqonly travels during an elastic-EP weightrefresh, a mixed-version fleet looks healthy until the first refresh, then fails at weight-update
time.
Validation
test/registered/unit/managers/test_msgpack_ipc_roundtrip.py(register_cpu_ci,base-c-test-cpu): per-type native round-trip and double-hop re-encode for all 12 types, theregistry-empty assertion, the
ExpertWeightPointernarrowing, field-parity between theCheckWeightsReqOutputmirrors and the pydantic models, atp>1-shaped payload, and theGetInternalStateReqOutputdataclass sanitization (materializingCudaGraphConfig).test_io_struct.pypasses unchanged.BackupDramReqneeds manual validation before merge. CI has no elastic-EP path, so the wirechange is only unit-covered. TODO before merge: run an elastic-EP backup→refresh cycle on a 2-node
mooncake setup with
SGLANG_USE_PICKLE_IPC=0, confirm weights transfer and inference stays correct,and record the result here. If that environment is unavailable at merge time, land types 1 and 3–13
in this PR (registry down to
(BackupDramReq,), machinery not yet deleted) and finishBackupDramReqplus the deletion in an immediate fast-follow.
Notes
mirror.
SetInjectDumpMetadataReqInputwas added via internal-sync and never wired in-tree; the deletionneeds internal-sync parity. Out-of-tree failure is loud (
ImportError/pickle.loadsAttributeError; the kill-switch does not help, since unpickling needs the class).enc_hook/dec_hookregion and runs contextinto
_maybe_wrap_pickle; expect a small conflict on the registry-branch hunk whichever landssecond — keep both changes. refactor: make time_stats msgpack-native #30005 / refactor: make customized_info msgpack-native (drop PickleWrapper) #29952 / refactor: make mm_data_mooncake msgpack-native (drop PickleWrapper) #29553 are line-number drift only; none touch the
tuple.
Out of scope (separate follow-ups)
continuous_buffer=Nonecrash when a node holds no expert weights(
expert_backup_manager.py).engine.pybase64-encode fix forload_format='flattened_bucket'(out-of-tree Engine callerspassing raw
bytes).SetInjectDumpMetadataReqOutput/LazyDumpTensorspair.-> Optional[torch.Tensor]return annotations in model files.save_sharded_state.py--max-file-size type=strbug.CI States
Latest PR Test (Base): ❌ Run #29377658530
Latest PR Test (Extra): ❌ Run #29377658496