diff --git a/.gitignore b/.gitignore index 57b966665..737f90f43 100644 --- a/.gitignore +++ b/.gitignore @@ -283,3 +283,6 @@ test/registered/xpu/test_nvidia_nemotron_3_nano.py artifacts/ .claude/scheduled_tasks.lock .humanize/ + +# Internal, non-published docs +internal-docs/ diff --git a/python/sglang/srt/elastic_ep/expert_backup_client.py b/python/sglang/srt/elastic_ep/expert_backup_client.py index 55c973af7..8b77f7f07 100644 --- a/python/sglang/srt/elastic_ep/expert_backup_client.py +++ b/python/sglang/srt/elastic_ep/expert_backup_client.py @@ -159,12 +159,10 @@ class ExpertBackupClient: param = param.narrow( 0, param.shape[0] // 2, param.shape[0] // 2 ) - server_ptr_list.append(weight_info["weight_ptr"]) + server_ptr_list.append(weight_info.weight_ptr) local_ptr_list.append(param.data_ptr()) - assert ( - param.numel() * param.element_size() == weight_info["byte_size"] - ) - weight_size_list.append(weight_info["byte_size"]) + assert param.numel() * param.element_size() == weight_info.byte_size + weight_size_list.append(weight_info.byte_size) before_transfer = time.time() ret = self.transfer_engine.engine.batch_transfer_sync_read( self.session_id_list[i], diff --git a/python/sglang/srt/elastic_ep/expert_backup_manager.py b/python/sglang/srt/elastic_ep/expert_backup_manager.py index c864dee8c..3ca4a7249 100644 --- a/python/sglang/srt/elastic_ep/expert_backup_manager.py +++ b/python/sglang/srt/elastic_ep/expert_backup_manager.py @@ -9,7 +9,12 @@ import zmq from sglang.srt.configs.load_config import LoadConfig from sglang.srt.configs.model_config import ModelConfig from sglang.srt.environ import envs -from sglang.srt.managers.io_struct import BackupDramReq, sock_recv, sock_send +from sglang.srt.managers.io_struct import ( + BackupDramReq, + ExpertWeightPointer, + sock_recv, + sock_send, +) from sglang.srt.model_loader.loader import DefaultModelLoader, get_model_loader from sglang.srt.model_loader.utils import set_default_torch_dtype from sglang.srt.server_args import ( @@ -128,15 +133,10 @@ class ExpertBackupManager: end_byte = current_byte_offset + byte_size weight_ptr = buffer_base_ptr + current_byte_offset self.continuous_buffer[start_byte:end_byte].copy_(weight_bytes) - self.weight_pointer_map[name] = { - "name": name, - "weight_ptr": weight_ptr, - "shape": weight_info["shape"], - "numel": weight_info["numel"], - "dtype": weight_info["dtype"], - "element_size": weight_info["element_size"], - "byte_size": byte_size, - } + self.weight_pointer_map[name] = ExpertWeightPointer( + weight_ptr=weight_ptr, + byte_size=byte_size, + ) current_byte_offset = end_byte diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index dbc1748be..d7bae9633 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -797,6 +797,11 @@ if os.environ.get("DUMPER_SERVER_PORT") == "reuse": async def _dumper_control_handler(method: str, request: Request): body_bytes = await request.body() body = await request.json() if body_bytes else {} + if not isinstance(body, dict): + return ORJSONResponse( + status_code=400, + content={"error": "Request body must be a JSON object."}, + ) obj = DumperControlReqInput(method=method, body=body) results = await _global_state.tokenizer_manager.dumper_control(obj) if any(not r.success for r in results): diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index b12454003..6ff55e01d 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -235,6 +235,7 @@ class Envs: # IPC SGLANG_USE_PICKLE_IPC = EnvBool(True) + # Log top-level PickleWrapper frames unwrapped on msgpack IPC decode. SGLANG_LOG_PICKLE_IPC_OBJECTS = EnvBool(False) # SGLang CI diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 354660bc8..5029c8b13 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1551,7 +1551,7 @@ class UpdateWeightFromDiskReqInput(BaseReq, kw_only=True): token_step: int = 0 # Whether to flush the cache after updating weights flush_cache: bool = True - # Tensor metadata + # Tensor metadata from the JSON request body, so it is already msgpack-native. manifest: Optional[Dict[str, Any]] = None @@ -1669,9 +1669,17 @@ class UpdateExpertBackupReq(BaseReq, kw_only=True): pass +class ExpertWeightPointer(msgspec.Struct, kw_only=True, array_like=True): + # One expert weight's pointer + byte length in the DRAM backup buffer. + # array_like: the map has tens of thousands of entries, so positional + # encoding drops the repeated field names from the wire. + weight_ptr: int + byte_size: int + + class BackupDramReq(BaseReq, kw_only=True): rank: int - weight_pointer_map: Dict[str, Any] + weight_pointer_map: Dict[str, ExpertWeightPointer] session_id: str buffer_size: int @@ -1718,7 +1726,9 @@ class GetWeightsByNameReqInput(BaseReq, kw_only=True): class GetWeightsByNameReqOutput(BaseReq, kw_only=True): - parameter: Optional[List[Any]] + # A flat List[float] or a per-row List[List[float]]. The union is on the + # element: Union[List[float], List[List[float]]] is invalid msgspec. + parameter: Optional[List[Union[float, List[float]]]] class ReleaseMemoryOccupationReqInput(BaseReq, kw_only=True): @@ -1746,10 +1756,32 @@ class CheckWeightsReqInput(BaseReq, kw_only=True): allow_quant_error: bool = False +# Wire versions of the pydantic ParallelismInfo/ChecksumInfo in +# sglang.srt.utils.weight_checker. Not array_like: the payload is read by field +# name and re-serialized to JSON, so it must stay a {field: value} map. +class ParallelismInfo(msgspec.Struct, kw_only=True): + tp_rank: int + tp_size: int + dp_rank: int + dp_size: int + pp_rank: int + pp_size: int + rank: int + size: int + + +class ChecksumInfo(msgspec.Struct, kw_only=True): + checksums: Dict[str, str] + per_gpu_checksum: str + parallelism_info: ParallelismInfo + + class CheckWeightsReqOutput(BaseReq, kw_only=True): success: bool message: str - payload: Optional[Dict[str, Any]] = None + # One ChecksumInfo per TP rank. The producer wraps the tp==1 result in a + # one-element list so the shape is always a list. + payload: Optional[List[ChecksumInfo]] = None class SlowDownReqInput(BaseReq, kw_only=True): @@ -1782,16 +1814,19 @@ class GetInternalStateReq(BaseReq, kw_only=True): class GetInternalStateReqOutput(BaseReq, kw_only=True): + # A vars() dump of ServerArgs, left untyped because a struct would drift. The + # producer sanitizes it with msgspec_to_builtins so every value is + # msgpack-native. internal_state: Dict[str, Any] class SetInternalStateReq(BaseReq, kw_only=True): - server_args: Dict[str, Any] + # Only numeric scheduler knobs are accepted (see Scheduler.set_internal_state). + server_args: Dict[str, Union[int, float]] class SetInternalStateReqOutput(BaseReq, kw_only=True): updated: bool - server_args: Dict[str, Any] class ProfileReqType(Enum): @@ -1922,13 +1957,16 @@ class SeparateReasoningReqInput(BaseReq, kw_only=True): class VertexGenerateReqInput(BaseReq, kw_only=True): + # Both fields come from the JSON request body, so they are already + # msgpack-native. instances: List[Dict[str, Any]] parameters: Optional[Dict[str, Any]] = None class RpcReqInput(BaseReq, kw_only=True): method: str - parameters: Optional[Dict[str, Any]] = None + # collective_rpc kwargs are flat scalars across all in-tree callers. + parameters: Optional[Dict[str, Union[bool, int, float, str, None]]] = None class RpcReqOutput(BaseReq, kw_only=True): @@ -1970,10 +2008,12 @@ class UnloadLoRAAdapterReqInput(BaseReq, kw_only=True): class LoadLoRAAdapterFromTensorsReqInput(BaseReq, kw_only=True): lora_name: str + # The PEFT adapter_config.json, already JSON — a tighter type would only add + # decode strictness with no benefit. config_dict: Dict[str, Any] serialized_tensors: str pinned: bool = False - added_tokens_config: Optional[Dict[str, Any]] = None + added_tokens_config: Optional[Dict[str, int]] = None lora_id: Optional[str] = None load_format: Optional[str] = None @@ -2006,10 +2046,6 @@ class BlockReqInput(BaseReq, kw_only=True): req_type: BlockReqType -class SetInjectDumpMetadataReqInput(BaseReq, kw_only=True): - dump_metadata: Dict[str, Any] - - class SetInjectDumpMetadataReqOutput(BaseReq, kw_only=True): success: bool @@ -2024,11 +2060,13 @@ class LazyDumpTensorsReqOutput(BaseReq, kw_only=True): class DumperControlReqInput(BaseReq, kw_only=True): method: str + # JSON request body (guarded to be a dict at the /dumper endpoint). body: Dict[str, Any] class DumperControlReqOutput(BaseReq, kw_only=True): success: bool + # JSON-native per-worker response dicts. response: List[Dict[str, Any]] error: str = "" @@ -2068,28 +2106,6 @@ def _check_all_req_types(): _check_all_req_types() -# IPC struct types whose fields still use opaque annotations (Any, Dict[str, Any], -# List[Any], etc.) instead of precise types. Keep these on explicit pickle -# transport until their field schemas are tightened, and keep the registry -# explicit so opaque usage can be audited and gradually narrowed. -# NOTE: GenerateReqInput and EmbeddingReqInput are standalone (not BaseReq/ -# BaseBatchReq subclasses) and are tracked separately. -_REQ_TYPES_WITH_OPAQUE_FIELDS: tuple[Type[msgspec.Struct], ...] = ( - UpdateWeightFromDiskReqInput, # manifest: Optional[Dict[str, Any]] - BackupDramReq, # weight_pointer_map: Dict[str, Any] - GetWeightsByNameReqOutput, # parameter: Optional[List[Any]] - CheckWeightsReqOutput, # payload: Optional[Dict[str, Any]] - GetInternalStateReqOutput, # internal_state: Dict[str, Any] - SetInternalStateReq, # server_args: Dict[str, Any] - SetInternalStateReqOutput, # server_args: Dict[str, Any] - VertexGenerateReqInput, # instances, parameters: Dict[str, Any] - RpcReqInput, # parameters: Optional[Dict[str, Any]] - LoadLoRAAdapterFromTensorsReqInput, # config_dict, added_tokens_config: Dict[str, Any] - SetInjectDumpMetadataReqInput, # dump_metadata: Dict[str, Any] - DumperControlReqInput, # body: Dict[str, Any] - DumperControlReqOutput, # response: List[Dict[str, Any]] -) - def wrap_as_pickle(obj: object) -> object: if obj is None: @@ -2180,19 +2196,13 @@ def hook_custom_types(*new_types: Type): def _maybe_wrap_pickle(obj: Any) -> Any: - if isinstance(obj, _REQ_TYPES_WITH_OPAQUE_FIELDS): - if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): - logger.info(f"Object of type {type(obj)} is wrapped via PickleWrapper.") - return PickleWrapper(pickle.dumps(obj)) - if isinstance(obj, (msgspec.Struct, *_primitive_types)): return obj raise TypeError( f"Cannot serialize object of type {type(obj)} over msgpack IPC. " - "Add a precise msgspec-compatible type, use an explicit PickleWrapper " - "field for the opaque payload, or add the struct to " - "_REQ_TYPES_WITH_OPAQUE_FIELDS with an audit comment." + "Add a precise msgspec-compatible type, or use an explicit PickleWrapper " + "field via wrap_as_pickle(...) for the opaque payload." ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 2667ba014..ea11a15a3 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3823,8 +3823,10 @@ class Scheduler( if info_record is not None: ret["dspark_info_record"] = info_record - # This field is not serializable. + # These fields are not msgpack-serializable (a config object and a bound + # signal handler); no reader consumes them. ret.pop("model_config", None) + ret.pop("custom_sigquit_handler", None) return GetInternalStateReqOutput(internal_state=msgspec_to_builtins(ret)) @@ -3906,13 +3908,7 @@ class Scheduler( get_server_args().override(source="update_server_args", **remaining) logger.info(f"Global server args updated! {get_server_args()=}") - server_args = dict(vars(get_server_args())) - # This field is not serializable. - server_args.pop("model_config", None) - return SetInternalStateReqOutput( - updated=if_success, - server_args=msgspec_to_builtins(server_args), - ) + return SetInternalStateReqOutput(updated=if_success) def save_remote_model(self, **kwargs): self.weight_updater.save_remote_model(kwargs) diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index 49e9a19a7..7013ce04f 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py @@ -8,6 +8,7 @@ from contextlib import contextmanager from dataclasses import dataclass, field from typing import Any, Callable, Dict, Iterator, Optional, Tuple +import msgspec import torch from sglang.srt.constants import ( @@ -18,6 +19,7 @@ from sglang.srt.constants import ( ) from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.managers.io_struct import ( + ChecksumInfo, CheckWeightsReqInput, CheckWeightsReqOutput, DestroyWeightsUpdateGroupReqInput, @@ -289,6 +291,11 @@ class SchedulerWeightUpdaterManager: all_payloads, payload, group=self.tp_cpu_group ) payload = all_payloads + if payload is not None: + # Normalize to one ChecksumInfo per rank so the wire shape is a + # uniform List[ChecksumInfo] (tp==1 becomes a single-element list). + per_rank = payload if isinstance(payload, list) else [payload] + payload = [msgspec.convert(p, ChecksumInfo) for p in per_rank] return CheckWeightsReqOutput( success=True, message="Success.", payload=payload ) diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index 24b608c0c..9cbec3725 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -15,6 +15,7 @@ from sglang.srt.managers.io_struct import ( AddExternalCorpusReqOutput, AttachHiCacheStorageReqInput, AttachHiCacheStorageReqOutput, + ChecksumInfo, CheckWeightsReqInput, CheckWeightsReqOutput, ClearHiCacheReqInput, @@ -77,6 +78,7 @@ from sglang.srt.utils import ( get_bool_env_var, normalize_serialized_named_tensor_payloads, ) +from sglang.srt.utils.msgspec_utils import msgspec_to_builtins from sglang.utils import TypeBasedDispatcher if TYPE_CHECKING: @@ -760,16 +762,15 @@ class TokenizerControlMixin: ranks: Optional[List[Dict]] = None per_engine_checksum: Optional[str] = None if any(r.payload is not None for r in results): - ranks = [] + rank_infos: List[ChecksumInfo] = [] for r in results: - if isinstance(r.payload, list): - ranks.extend(r.payload) - else: - ranks.append(r.payload) + if r.payload is not None: + rank_infos.extend(r.payload) h = hashlib.sha256() - for rank in ranks: - h.update(rank["per_gpu_checksum"].encode()) + for info in rank_infos: + h.update(info.per_gpu_checksum.encode()) per_engine_checksum = h.hexdigest() + ranks = [msgspec_to_builtins(info) for info in rank_infos] return success, message, ranks, per_engine_checksum async def slow_down( diff --git a/python/sglang/srt/utils/msgspec_utils.py b/python/sglang/srt/utils/msgspec_utils.py index 60955b1cc..f848c0028 100644 --- a/python/sglang/srt/utils/msgspec_utils.py +++ b/python/sglang/srt/utils/msgspec_utils.py @@ -2,6 +2,7 @@ from __future__ import annotations import base64 import binascii +import dataclasses from typing import Any import msgspec @@ -35,13 +36,19 @@ class Base64Bytes: def msgspec_to_builtins(obj: Any) -> Any: - """Recursively convert msgspec structs to dict/list Python builtins.""" + """Recursively convert msgspec structs and dataclasses to builtins.""" if isinstance(obj, msgspec.Struct): return { field.name: msgspec_to_builtins(getattr(obj, field.name)) for field in msgspec.structs.fields(type(obj)) } + if dataclasses.is_dataclass(obj) and not isinstance(obj, type): + return { + f.name: msgspec_to_builtins(getattr(obj, f.name)) + for f in dataclasses.fields(obj) + } + if isinstance(obj, dict): return {key: msgspec_to_builtins(value) for key, value in obj.items()} diff --git a/test/registered/spec/dspark/test_dspark_info_dumper.py b/test/registered/spec/dspark/test_dspark_info_dumper.py index 0e948a673..10703ca79 100644 --- a/test/registered/spec/dspark/test_dspark_info_dumper.py +++ b/test/registered/spec/dspark/test_dspark_info_dumper.py @@ -3,6 +3,12 @@ import unittest import torch from sglang.srt.environ import envs +from sglang.srt.managers.io_struct import ( + GetInternalStateReqOutput, + PickleWrapper, + msgpack_decode, + msgpack_encode, +) from sglang.srt.speculative.dspark_components.dspark_observability import ( DecodeStepObservation, DsparkInfoDumper, @@ -12,6 +18,7 @@ from sglang.srt.speculative.dspark_components.dspark_observability import ( resolve_components, resolve_enabled_components, ) +from sglang.srt.utils.msgspec_utils import msgspec_to_builtins from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -394,5 +401,39 @@ class TestReqsAndGpuTiming(CustomTestCase): self.assertGreater(record["target_verify_gpu_ms"], 0.0) +class TestDumpCrossesMsgpackIpc(CustomTestCase): + """Guard the DSpark -> GetInternalStateReqOutput serialization contract. + + `Scheduler.get_internal_state` stores `draft_worker.dump_info_records()` under + `internal_state["dspark_info_record"]`, then ships the struct over the strict + msgpack IPC path (issue #29465). That path has no PickleWrapper fallback: any + value that is not msgpack-native (a numpy scalar, a torch tensor, an + un-converted `msgspec.Struct`) raises at encode time. The dumper's own tests + assert record *values* -- and `assertEqual(np.int64(3), 3)` passes -- so a + scalar that silently became numpy would escape them but fail this round-trip. + """ + + def _real_dump(self): + dumper, clock = make_dumper({"core"}) + dumper.observe_decode_step(make_obs(forward_ct=1)) + clock.advance(0.01) + dumper.observe_decode_step(make_obs(forward_ct=2)) + dumped = dumper.dump() + # DsparkObservability.dump_info_records appends this float onto the raw + # dumper output before the scheduler reads it; mirror the full payload. + dumped["simulate_acc_len"] = 4.0 + return dumped + + def test_real_dump_output_round_trips_natively(self): + internal_state = msgspec_to_builtins( + {"dspark_info_record": self._real_dump(), "max_running_requests": 256} + ) + output = GetInternalStateReqOutput(internal_state=internal_state) + + decoded = msgpack_decode(msgpack_encode(output)) + self.assertNotIsInstance(decoded, PickleWrapper) + self.assertEqual(decoded, output) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py b/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py new file mode 100644 index 000000000..e3d3e95d0 --- /dev/null +++ b/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py @@ -0,0 +1,241 @@ +"""Round-trip coverage for the IPC structs that used to be pickle-wrapped. + +Issue #29465 Task 4 tightened the 13 types in `_REQ_TYPES_WITH_OPAQUE_FIELDS` to +precise msgspec-native annotations and deleted the registry. This test proves +each type now encodes natively over the msgpack IPC path (no `PickleWrapper` +frame) by asserting `msgpack_decode(msgpack_encode(x)) == x`, and guards the +type-specific decisions (the `ExpertWeightPointer` narrowing, the +`CheckWeightsReqOutput` struct mirrors, and the internal-state sanitization). +""" + +import dataclasses +import unittest + +import msgspec + +from sglang.srt.managers import io_struct +from sglang.srt.managers.io_struct import ( + BackupDramReq, + ChecksumInfo, + CheckWeightsReqOutput, + DumperControlReqInput, + DumperControlReqOutput, + ExpertWeightPointer, + GetInternalStateReqOutput, + GetWeightsByNameReqOutput, + LoadLoRAAdapterFromTensorsReqInput, + ParallelismInfo, + RpcReqInput, + SetInternalStateReq, + SetInternalStateReqOutput, + UpdateWeightFromDiskReqInput, + VertexGenerateReqInput, + msgpack_decode, + msgpack_encode, +) +from sglang.srt.model_executor.cuda_graph_config import CudaGraphConfig +from sglang.srt.utils.msgspec_utils import msgspec_to_builtins +from sglang.srt.utils.weight_checker import ChecksumInfo as PydanticChecksumInfo +from sglang.srt.utils.weight_checker import ParallelismInfo as PydanticParallelismInfo +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-c-test-cpu") + + +def _round_trip(obj): + return msgpack_decode(msgpack_encode(obj)) + + +def _double_hop(obj): + # MultiTokenizerRouter and the DP controller re-encode already-decoded + # structs, so a single hop cannot catch re-encode bugs. + return msgpack_decode(msgpack_encode(msgpack_decode(msgpack_encode(obj)))) + + +def _contains_dataclass(obj) -> bool: + if dataclasses.is_dataclass(obj) and not isinstance(obj, type): + return True + if isinstance(obj, dict): + return any(_contains_dataclass(v) for v in obj.values()) + if isinstance(obj, (list, tuple, set)): + return any(_contains_dataclass(v) for v in obj) + return False + + +def _parallelism_info() -> ParallelismInfo: + return ParallelismInfo( + tp_rank=0, tp_size=2, dp_rank=0, dp_size=1, pp_rank=0, pp_size=1, rank=0, size=2 + ) + + +def _checksum_info(tag: str) -> ChecksumInfo: + return ChecksumInfo( + checksums={f"model.layers.{tag}": "deadbeef"}, + per_gpu_checksum="cafef00d", + parallelism_info=_parallelism_info(), + ) + + +# One representative instance per (now-tightened) ex-registry type. The 13th +# entry, SetInjectDumpMetadataReqInput, was deleted as dead code, leaving 12. +REGISTRY_TYPE_INSTANCES = { + "UpdateWeightFromDiskReqInput": UpdateWeightFromDiskReqInput( + model_path="dummy", manifest={"w": [1, 2], "meta": {"k": "v"}} + ), + "BackupDramReq": BackupDramReq( + rank=0, + weight_pointer_map={ + "experts.0.gate_proj": ExpertWeightPointer(weight_ptr=8, byte_size=4), + "experts.1.up_proj": ExpertWeightPointer(weight_ptr=16, byte_size=8), + }, + session_id="session", + buffer_size=1024, + ), + "GetWeightsByNameReqOutput/flat": GetWeightsByNameReqOutput( + parameter=[1.0, 2.5, 3.0] + ), + "GetWeightsByNameReqOutput/nested": GetWeightsByNameReqOutput( + parameter=[[1.0, 2.0], [3.0]] + ), + "GetWeightsByNameReqOutput/none": GetWeightsByNameReqOutput(parameter=None), + "CheckWeightsReqOutput": CheckWeightsReqOutput( + success=True, + message="Success.", + payload=[_checksum_info("0"), _checksum_info("1")], + ), + "GetInternalStateReqOutput": GetInternalStateReqOutput( + internal_state={"a": 1, "b": [1, 2], "c": {"d": "e"}, "f": None} + ), + "SetInternalStateReq": SetInternalStateReq( + server_args={ + "pp_max_micro_batch_size": 4, + "speculative_accept_threshold_acc": 0.5, + } + ), + "SetInternalStateReqOutput": SetInternalStateReqOutput(updated=True), + "VertexGenerateReqInput": VertexGenerateReqInput( + instances=[{"prompt": "hi"}], parameters={"max_tokens": 8} + ), + "RpcReqInput/empty": RpcReqInput(method="collective_rpc", parameters={}), + "RpcReqInput/scalars": RpcReqInput( + method="collective_rpc", + parameters={"flag": True, "n": 1, "ratio": 2.0, "name": "x", "opt": None}, + ), + "RpcReqInput/none": RpcReqInput(method="collective_rpc", parameters=None), + "LoadLoRAAdapterFromTensorsReqInput": LoadLoRAAdapterFromTensorsReqInput( + lora_name="adapter", + config_dict={"r": 8, "lora_alpha": 16, "target_modules": ["q_proj", "v_proj"]}, + serialized_tensors="", + added_tokens_config={"": 32000}, + ), + "DumperControlReqInput": DumperControlReqInput(method="start", body={"k": "v"}), + "DumperControlReqOutput": DumperControlReqOutput( + success=True, response=[{"worker": 0, "ok": True}] + ), +} + +NARROWED_BACKUP_KEYS = ("name", "shape", "numel", "dtype", "element_size") + + +class TestMsgpackIpcRoundtrip(CustomTestCase): + def test_registry_is_empty(self): + # `getattr(..., ())` is deliberate: the acceptance criterion for Task 4 is + # that the symbol is *deleted*, so this asserts its absence rather than + # defensively reading a field. It must survive the symbol removal. + self.assertEqual( + getattr(io_struct, "_REQ_TYPES_WITH_OPAQUE_FIELDS", ()), + (), + ) + + def test_each_type_round_trips_natively(self): + for name, instance in REGISTRY_TYPE_INSTANCES.items(): + with self.subTest(type=name): + encoded = msgpack_encode(instance) + # Natively encoded structs are never wrapped: a PickleWrapper + # frame would decode back to a PickleWrapper, not the type. + self.assertNotIsInstance( + msgpack_decode(encoded), io_struct.PickleWrapper + ) + self.assertEqual(_round_trip(instance), instance) + self.assertEqual(_double_hop(instance), instance) + + def test_backup_dram_req_is_narrowed(self): + # ExpertWeightPointer carries only the two fields the consumer reads; the + # five torch-metadata keys the producer used to send are gone. + field_names = {f.name for f in msgspec.structs.fields(ExpertWeightPointer)} + self.assertEqual(field_names, {"weight_ptr", "byte_size"}) + for dropped in NARROWED_BACKUP_KEYS: + self.assertNotIn(dropped, field_names) + + decoded = _round_trip(REGISTRY_TYPE_INSTANCES["BackupDramReq"]) + pointer = decoded.weight_pointer_map["experts.0.gate_proj"] + self.assertEqual((pointer.weight_ptr, pointer.byte_size), (8, 4)) + + def test_check_weights_mirrors_match_pydantic_models(self): + # Field-parity guard: the msgspec wire structs must not drift from the + # pydantic source of truth in weight_checker. + self.assertEqual( + {f.name for f in msgspec.structs.fields(ParallelismInfo)}, + set(PydanticParallelismInfo.model_fields), + ) + self.assertEqual( + {f.name for f in msgspec.structs.fields(ChecksumInfo)}, + set(PydanticChecksumInfo.model_fields), + ) + + def test_check_weights_multi_rank_payload(self): + # tp>1 sends one ChecksumInfo per rank; the list round-trips and stays a + # {field: value} dict once converted back to builtins for the HTTP body. + instance = REGISTRY_TYPE_INSTANCES["CheckWeightsReqOutput"] + decoded = _round_trip(instance) + self.assertEqual(len(decoded.payload), 2) + as_dict = msgspec_to_builtins(decoded.payload[0]) + self.assertEqual(as_dict["per_gpu_checksum"], "cafef00d") + self.assertIn("tp_rank", as_dict["parallelism_info"]) + + def test_check_weights_producer_conversion(self): + # Mirrors weight_updater.check_weights: WeightChecker returns + # ChecksumInfo.model_dump() (a dict), converted to the msgspec struct via + # msgspec.convert, and the result round-trips as the payload. + pydantic_checksum = PydanticChecksumInfo( + checksums={"model.layers.0": "deadbeef"}, + per_gpu_checksum="cafef00d", + parallelism_info=PydanticParallelismInfo( + tp_rank=0, + tp_size=2, + dp_rank=0, + dp_size=1, + pp_rank=0, + pp_size=1, + rank=0, + size=2, + ), + ) + converted = msgspec.convert(pydantic_checksum.model_dump(), ChecksumInfo) + self.assertEqual(converted.per_gpu_checksum, "cafef00d") + self.assertEqual(converted.parallelism_info.tp_rank, 0) + output = CheckWeightsReqOutput(success=True, message="ok", payload=[converted]) + self.assertEqual(_round_trip(output), output) + + def test_get_internal_state_sanitizes_dataclass(self): + # A live vars(ServerArgs) dump holds a dataclass: cuda_graph_config is an + # Optional[CudaGraphConfig]. The producer sanitizes via msgspec_to_builtins + # so it does not survive onto the wire; materialize CudaGraphConfig + # explicitly. + raw = { + "cuda_graph_config": CudaGraphConfig(), + "max_running_requests": 256, + } + self.assertTrue(_contains_dataclass(raw)) + + sanitized = msgspec_to_builtins(raw) + self.assertFalse(_contains_dataclass(sanitized)) + self.assertIsInstance(sanitized["cuda_graph_config"], dict) + + output = GetInternalStateReqOutput(internal_state=sanitized) + self.assertEqual(_round_trip(output), output) + + +if __name__ == "__main__": + unittest.main()