This commit is contained in:
@@ -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):
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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."
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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()
|
||||||
Reference in New Issue
Block a user