Empty _REQ_TYPES_WITH_OPAQUE_FIELDS on the msgpack IPC path (#29465 Task 4) (#30182)

This commit is contained in:
Jorge António
2026-07-15 14:55:06 -07:00
committed by GitHub
parent 67148447a6
commit 26cb0fcdda
12 changed files with 383 additions and 73 deletions
+3
View File
@@ -283,3 +283,6 @@ test/registered/xpu/test_nvidia_nemotron_3_nano.py
artifacts/ artifacts/
.claude/scheduled_tasks.lock .claude/scheduled_tasks.lock
.humanize/ .humanize/
# Internal, non-published docs
internal-docs/
@@ -159,12 +159,10 @@ class ExpertBackupClient:
param = param.narrow( param = param.narrow(
0, param.shape[0] // 2, param.shape[0] // 2 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()) local_ptr_list.append(param.data_ptr())
assert ( assert param.numel() * param.element_size() == weight_info.byte_size
param.numel() * param.element_size() == weight_info["byte_size"] weight_size_list.append(weight_info.byte_size)
)
weight_size_list.append(weight_info["byte_size"])
before_transfer = time.time() before_transfer = time.time()
ret = self.transfer_engine.engine.batch_transfer_sync_read( ret = self.transfer_engine.engine.batch_transfer_sync_read(
self.session_id_list[i], self.session_id_list[i],
@@ -9,7 +9,12 @@ import zmq
from sglang.srt.configs.load_config import LoadConfig from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.configs.model_config import ModelConfig from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.environ import envs 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.loader import DefaultModelLoader, get_model_loader
from sglang.srt.model_loader.utils import set_default_torch_dtype from sglang.srt.model_loader.utils import set_default_torch_dtype
from sglang.srt.server_args import ( from sglang.srt.server_args import (
@@ -128,15 +133,10 @@ class ExpertBackupManager:
end_byte = current_byte_offset + byte_size end_byte = current_byte_offset + byte_size
weight_ptr = buffer_base_ptr + current_byte_offset weight_ptr = buffer_base_ptr + current_byte_offset
self.continuous_buffer[start_byte:end_byte].copy_(weight_bytes) self.continuous_buffer[start_byte:end_byte].copy_(weight_bytes)
self.weight_pointer_map[name] = { self.weight_pointer_map[name] = ExpertWeightPointer(
"name": name, weight_ptr=weight_ptr,
"weight_ptr": weight_ptr, byte_size=byte_size,
"shape": weight_info["shape"], )
"numel": weight_info["numel"],
"dtype": weight_info["dtype"],
"element_size": weight_info["element_size"],
"byte_size": byte_size,
}
current_byte_offset = end_byte current_byte_offset = end_byte
@@ -797,6 +797,11 @@ if os.environ.get("DUMPER_SERVER_PORT") == "reuse":
async def _dumper_control_handler(method: str, request: Request): async def _dumper_control_handler(method: str, request: Request):
body_bytes = await request.body() body_bytes = await request.body()
body = await request.json() if body_bytes else {} 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) obj = DumperControlReqInput(method=method, body=body)
results = await _global_state.tokenizer_manager.dumper_control(obj) results = await _global_state.tokenizer_manager.dumper_control(obj)
if any(not r.success for r in results): if any(not r.success for r in results):
+1
View File
@@ -235,6 +235,7 @@ class Envs:
# IPC # IPC
SGLANG_USE_PICKLE_IPC = EnvBool(True) SGLANG_USE_PICKLE_IPC = EnvBool(True)
# Log top-level PickleWrapper frames unwrapped on msgpack IPC decode.
SGLANG_LOG_PICKLE_IPC_OBJECTS = EnvBool(False) SGLANG_LOG_PICKLE_IPC_OBJECTS = EnvBool(False)
# SGLang CI # SGLang CI
+52 -42
View File
@@ -1551,7 +1551,7 @@ class UpdateWeightFromDiskReqInput(BaseReq, kw_only=True):
token_step: int = 0 token_step: int = 0
# Whether to flush the cache after updating weights # Whether to flush the cache after updating weights
flush_cache: bool = True 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 manifest: Optional[Dict[str, Any]] = None
@@ -1669,9 +1669,17 @@ class UpdateExpertBackupReq(BaseReq, kw_only=True):
pass 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): class BackupDramReq(BaseReq, kw_only=True):
rank: int rank: int
weight_pointer_map: Dict[str, Any] weight_pointer_map: Dict[str, ExpertWeightPointer]
session_id: str session_id: str
buffer_size: int buffer_size: int
@@ -1718,7 +1726,9 @@ class GetWeightsByNameReqInput(BaseReq, kw_only=True):
class GetWeightsByNameReqOutput(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): class ReleaseMemoryOccupationReqInput(BaseReq, kw_only=True):
@@ -1746,10 +1756,32 @@ class CheckWeightsReqInput(BaseReq, kw_only=True):
allow_quant_error: bool = False 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): class CheckWeightsReqOutput(BaseReq, kw_only=True):
success: bool success: bool
message: str 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): class SlowDownReqInput(BaseReq, kw_only=True):
@@ -1782,16 +1814,19 @@ class GetInternalStateReq(BaseReq, kw_only=True):
class GetInternalStateReqOutput(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] internal_state: Dict[str, Any]
class SetInternalStateReq(BaseReq, kw_only=True): 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): class SetInternalStateReqOutput(BaseReq, kw_only=True):
updated: bool updated: bool
server_args: Dict[str, Any]
class ProfileReqType(Enum): class ProfileReqType(Enum):
@@ -1922,13 +1957,16 @@ class SeparateReasoningReqInput(BaseReq, kw_only=True):
class VertexGenerateReqInput(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]] instances: List[Dict[str, Any]]
parameters: Optional[Dict[str, Any]] = None parameters: Optional[Dict[str, Any]] = None
class RpcReqInput(BaseReq, kw_only=True): class RpcReqInput(BaseReq, kw_only=True):
method: str 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): class RpcReqOutput(BaseReq, kw_only=True):
@@ -1970,10 +2008,12 @@ class UnloadLoRAAdapterReqInput(BaseReq, kw_only=True):
class LoadLoRAAdapterFromTensorsReqInput(BaseReq, kw_only=True): class LoadLoRAAdapterFromTensorsReqInput(BaseReq, kw_only=True):
lora_name: str 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] config_dict: Dict[str, Any]
serialized_tensors: str serialized_tensors: str
pinned: bool = False pinned: bool = False
added_tokens_config: Optional[Dict[str, Any]] = None added_tokens_config: Optional[Dict[str, int]] = None
lora_id: Optional[str] = None lora_id: Optional[str] = None
load_format: Optional[str] = None load_format: Optional[str] = None
@@ -2006,10 +2046,6 @@ class BlockReqInput(BaseReq, kw_only=True):
req_type: BlockReqType req_type: BlockReqType
class SetInjectDumpMetadataReqInput(BaseReq, kw_only=True):
dump_metadata: Dict[str, Any]
class SetInjectDumpMetadataReqOutput(BaseReq, kw_only=True): class SetInjectDumpMetadataReqOutput(BaseReq, kw_only=True):
success: bool success: bool
@@ -2024,11 +2060,13 @@ class LazyDumpTensorsReqOutput(BaseReq, kw_only=True):
class DumperControlReqInput(BaseReq, kw_only=True): class DumperControlReqInput(BaseReq, kw_only=True):
method: str method: str
# JSON request body (guarded to be a dict at the /dumper endpoint).
body: Dict[str, Any] body: Dict[str, Any]
class DumperControlReqOutput(BaseReq, kw_only=True): class DumperControlReqOutput(BaseReq, kw_only=True):
success: bool success: bool
# JSON-native per-worker response dicts.
response: List[Dict[str, Any]] response: List[Dict[str, Any]]
error: str = "" error: str = ""
@@ -2068,28 +2106,6 @@ def _check_all_req_types():
_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: def wrap_as_pickle(obj: object) -> object:
if obj is None: if obj is None:
@@ -2180,19 +2196,13 @@ def hook_custom_types(*new_types: Type):
def _maybe_wrap_pickle(obj: Any) -> Any: 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)): if isinstance(obj, (msgspec.Struct, *_primitive_types)):
return obj return obj
raise TypeError( raise TypeError(
f"Cannot serialize object of type {type(obj)} over msgpack IPC. " f"Cannot serialize object of type {type(obj)} over msgpack IPC. "
"Add a precise msgspec-compatible type, use an explicit PickleWrapper " "Add a precise msgspec-compatible type, or use an explicit PickleWrapper "
"field for the opaque payload, or add the struct to " "field via wrap_as_pickle(...) for the opaque payload."
"_REQ_TYPES_WITH_OPAQUE_FIELDS with an audit comment."
) )
+4 -8
View File
@@ -3823,8 +3823,10 @@ class Scheduler(
if info_record is not None: if info_record is not None:
ret["dspark_info_record"] = info_record 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("model_config", None)
ret.pop("custom_sigquit_handler", None)
return GetInternalStateReqOutput(internal_state=msgspec_to_builtins(ret)) return GetInternalStateReqOutput(internal_state=msgspec_to_builtins(ret))
@@ -3906,13 +3908,7 @@ class Scheduler(
get_server_args().override(source="update_server_args", **remaining) get_server_args().override(source="update_server_args", **remaining)
logger.info(f"Global server args updated! {get_server_args()=}") logger.info(f"Global server args updated! {get_server_args()=}")
server_args = dict(vars(get_server_args())) return SetInternalStateReqOutput(updated=if_success)
# This field is not serializable.
server_args.pop("model_config", None)
return SetInternalStateReqOutput(
updated=if_success,
server_args=msgspec_to_builtins(server_args),
)
def save_remote_model(self, **kwargs): def save_remote_model(self, **kwargs):
self.weight_updater.save_remote_model(kwargs) self.weight_updater.save_remote_model(kwargs)
@@ -8,6 +8,7 @@ from contextlib import contextmanager
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Any, Callable, Dict, Iterator, Optional, Tuple from typing import Any, Callable, Dict, Iterator, Optional, Tuple
import msgspec
import torch import torch
from sglang.srt.constants import ( from sglang.srt.constants import (
@@ -18,6 +19,7 @@ from sglang.srt.constants import (
) )
from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.managers.io_struct import ( from sglang.srt.managers.io_struct import (
ChecksumInfo,
CheckWeightsReqInput, CheckWeightsReqInput,
CheckWeightsReqOutput, CheckWeightsReqOutput,
DestroyWeightsUpdateGroupReqInput, DestroyWeightsUpdateGroupReqInput,
@@ -289,6 +291,11 @@ class SchedulerWeightUpdaterManager:
all_payloads, payload, group=self.tp_cpu_group all_payloads, payload, group=self.tp_cpu_group
) )
payload = all_payloads 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( return CheckWeightsReqOutput(
success=True, message="Success.", payload=payload success=True, message="Success.", payload=payload
) )
@@ -15,6 +15,7 @@ from sglang.srt.managers.io_struct import (
AddExternalCorpusReqOutput, AddExternalCorpusReqOutput,
AttachHiCacheStorageReqInput, AttachHiCacheStorageReqInput,
AttachHiCacheStorageReqOutput, AttachHiCacheStorageReqOutput,
ChecksumInfo,
CheckWeightsReqInput, CheckWeightsReqInput,
CheckWeightsReqOutput, CheckWeightsReqOutput,
ClearHiCacheReqInput, ClearHiCacheReqInput,
@@ -77,6 +78,7 @@ from sglang.srt.utils import (
get_bool_env_var, get_bool_env_var,
normalize_serialized_named_tensor_payloads, normalize_serialized_named_tensor_payloads,
) )
from sglang.srt.utils.msgspec_utils import msgspec_to_builtins
from sglang.utils import TypeBasedDispatcher from sglang.utils import TypeBasedDispatcher
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -760,16 +762,15 @@ class TokenizerControlMixin:
ranks: Optional[List[Dict]] = None ranks: Optional[List[Dict]] = None
per_engine_checksum: Optional[str] = None per_engine_checksum: Optional[str] = None
if any(r.payload is not None for r in results): if any(r.payload is not None for r in results):
ranks = [] rank_infos: List[ChecksumInfo] = []
for r in results: for r in results:
if isinstance(r.payload, list): if r.payload is not None:
ranks.extend(r.payload) rank_infos.extend(r.payload)
else:
ranks.append(r.payload)
h = hashlib.sha256() h = hashlib.sha256()
for rank in ranks: for info in rank_infos:
h.update(rank["per_gpu_checksum"].encode()) h.update(info.per_gpu_checksum.encode())
per_engine_checksum = h.hexdigest() per_engine_checksum = h.hexdigest()
ranks = [msgspec_to_builtins(info) for info in rank_infos]
return success, message, ranks, per_engine_checksum return success, message, ranks, per_engine_checksum
async def slow_down( async def slow_down(
+8 -1
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import base64 import base64
import binascii import binascii
import dataclasses
from typing import Any from typing import Any
import msgspec import msgspec
@@ -35,13 +36,19 @@ class Base64Bytes:
def msgspec_to_builtins(obj: Any) -> Any: 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): if isinstance(obj, msgspec.Struct):
return { return {
field.name: msgspec_to_builtins(getattr(obj, field.name)) field.name: msgspec_to_builtins(getattr(obj, field.name))
for field in msgspec.structs.fields(type(obj)) 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): if isinstance(obj, dict):
return {key: msgspec_to_builtins(value) for key, value in obj.items()} return {key: msgspec_to_builtins(value) for key, value in obj.items()}
@@ -3,6 +3,12 @@ import unittest
import torch import torch
from sglang.srt.environ import envs 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 ( from sglang.srt.speculative.dspark_components.dspark_observability import (
DecodeStepObservation, DecodeStepObservation,
DsparkInfoDumper, DsparkInfoDumper,
@@ -12,6 +18,7 @@ from sglang.srt.speculative.dspark_components.dspark_observability import (
resolve_components, resolve_components,
resolve_enabled_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.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase from sglang.test.test_utils import CustomTestCase
@@ -394,5 +401,39 @@ class TestReqsAndGpuTiming(CustomTestCase):
self.assertGreater(record["target_verify_gpu_ms"], 0.0) 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__": if __name__ == "__main__":
unittest.main() unittest.main()
@@ -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={"<extra>": 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()