diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 2e6cfc4f2..2d978f807 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -22,7 +22,7 @@ import torch import torch.distributed as dist import zmq -from sglang.srt.managers.io_struct import sock_recv, sock_send +from sglang.srt.managers.io_struct import sock_recv, sock_send, wrap_as_pickle # -------------------------------------- config base ------------------------------------------ @@ -1440,7 +1440,7 @@ def _create_zmq_rpc_broadcast( except Exception as e: _log(f"[ZmqRpc] error inside handler: {e}") resp = {"result": None, "error": str(e)} - sock_send(sock, resp) + sock_send(sock, wrap_as_pickle(resp)) thread = threading.Thread(target=serve_loop, daemon=True) thread.start() @@ -1479,11 +1479,13 @@ class _ZmqRpcHandle: def call(*args, **kwargs): sock_send( self._socket, - { - "method": method_name, - "args": args, - "kwargs": kwargs, - }, + wrap_as_pickle( + { + "method": method_name, + "args": args, + "kwargs": kwargs, + } + ), ) response = sock_recv(self._socket) if response["error"]: diff --git a/python/sglang/srt/disaggregation/encode_grpc_server.py b/python/sglang/srt/disaggregation/encode_grpc_server.py index 71c3d9729..5abb27473 100644 --- a/python/sglang/srt/disaggregation/encode_grpc_server.py +++ b/python/sglang/srt/disaggregation/encode_grpc_server.py @@ -26,7 +26,7 @@ from sglang.srt.disaggregation.encode_server import ( handle_scheduler_receive_url_request, launch_encoder, ) -from sglang.srt.managers.io_struct import async_sock_send +from sglang.srt.managers.io_struct import async_sock_send, wrap_as_pickle from sglang.srt.managers.schedule_batch import Modality from sglang.srt.server_args import PortArgs, ServerArgs from sglang.srt.utils import random_uuid @@ -96,7 +96,7 @@ class SGLangEncoderServer(SGLangEncoderServicer): "part_idx": request.part_idx, } for socket in self.send_sockets: - await async_sock_send(socket, request_dict) + await async_sock_send(socket, wrap_as_pickle(request_dict)) # gRPC encode is image-only; encoder.encode() requires modality ( diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index f92639a11..6de7918ab 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -50,6 +50,7 @@ from sglang.srt.managers.io_struct import ( async_sock_recv, async_sock_send, sock_send, + wrap_as_pickle, ) from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache @@ -2398,12 +2399,14 @@ class EncoderScheduler: for sock in self.send_sockets: sock_send( sock, - { - "type": "batch_encode", - "modality": modality.name, - "requests": requests, - "enter_time": start, - }, + wrap_as_pickle( + { + "type": "batch_encode", + "modality": modality.name, + "requests": requests, + "enter_time": start, + } + ), ) logger.info(f"Dispatching batch of {len(group)} {modality.name} requests") @@ -2449,7 +2452,7 @@ class EncoderScheduler: req = p.request try: for sock in self.send_sockets: - sock_send(sock, req) + sock_send(sock, wrap_as_pickle(req)) result = await self.encoder.encode_request(req, modality) if not p.future.done(): p.future.set_result(result) @@ -2738,7 +2741,7 @@ class DPDispatcher: ) try: - await async_sock_send(self.dispatch_sockets[rank], request) + await async_sock_send(self.dispatch_sockets[rank], wrap_as_pickle(request)) # An alive-but-stuck worker (NCCL deadlock etc.) wouldn't trip # the watchdog, so bound the wait explicitly. return await asyncio.wait_for(future, timeout=ENCODER_REQ_TIMEOUT) @@ -2787,7 +2790,7 @@ class DPDispatcher: f"dp_rank={rank}, pending={self.pending_counts}" ) try: - await async_sock_send(self.dispatch_sockets[rank], request) + await async_sock_send(self.dispatch_sockets[rank], wrap_as_pickle(request)) return await asyncio.wait_for(future, timeout=ENCODER_REQ_TIMEOUT) except asyncio.TimeoutError: self.pending_futures[rank].pop(key, None) @@ -2828,7 +2831,9 @@ class DPDispatcher: self.req_id_to_rank[req_id] = rank rank_keys.append((rank, req_id)) request_copy = {**request, "req_id": req_id} - await async_sock_send(self.dispatch_sockets[rank], request_copy) + await async_sock_send( + self.dispatch_sockets[rank], wrap_as_pickle(request_copy) + ) futures.append(future) # Concurrent wait → total bounded by eff_timeout, not # dp_size × eff_timeout. @@ -3056,7 +3061,7 @@ async def _dp_worker_handle_request( # pyzmq async send isn't safe for concurrent senders. try: async with send_lock: - await async_sock_send(send_sock, envelope) + await async_sock_send(send_sock, wrap_as_pickle(envelope)) except Exception: logger.error( f"DP worker {dp_rank} failed to send envelope for " @@ -3603,7 +3608,7 @@ async def handle_encode_request(request: dict): ) else: for socket in send_sockets: - sock_send(socket, request) + sock_send(socket, wrap_as_pickle(request)) nbytes, embedding_len, embedding_dim, error_msg, error_code = ( await encoder.encode_request(request, modality) ) @@ -3824,7 +3829,7 @@ async def health_generate(): # Broadcast to other TP ranks so distributed ops stay in sync for socket in send_sockets: - sock_send(socket, dummy_request) + sock_send(socket, wrap_as_pickle(dummy_request)) # Run encode on rank 0 with timeout _, _, _, error_msg, _ = await asyncio.wait_for( diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 1ed7c291a..f438dabed 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -109,6 +109,7 @@ from sglang.srt.utils import ( set_prometheus_multiproc_dir, set_ulimit, ) +from sglang.srt.utils.msgspec_utils import msgspec_to_builtins from sglang.srt.utils.network import get_zmq_socket, is_port_available from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter from sglang.srt.utils.watchdog import SubprocessWatchdog @@ -994,12 +995,14 @@ class Engine(EngineScoreMixin, EngineBase): internal_states = self.loop.run_until_complete( self.tokenizer_manager.get_internal_state() ) - return { - **dataclasses.asdict(self.tokenizer_manager.server_args), - **self._scheduler_init_result.scheduler_infos[0], - "internal_states": internal_states, - "version": __version__, - } + return msgspec_to_builtins( + { + **dataclasses.asdict(self.tokenizer_manager.server_args), + **self._scheduler_init_result.scheduler_infos[0], + "internal_states": internal_states, + "version": __version__, + } + ) def init_weights_update_group( self, diff --git a/python/sglang/srt/entrypoints/grpc_bridge.py b/python/sglang/srt/entrypoints/grpc_bridge.py index 5c56edc03..fa0ab61b8 100644 --- a/python/sglang/srt/entrypoints/grpc_bridge.py +++ b/python/sglang/srt/entrypoints/grpc_bridge.py @@ -15,6 +15,8 @@ from typing import Any, Awaitable, Callable, Dict, List, Optional from pydantic import ValidationError +from sglang.srt.utils.msgspec_utils import msgspec_to_builtins + logger = logging.getLogger(__name__) @@ -383,9 +385,9 @@ class RuntimeHandle: return json.dumps(result, default=str) def get_server_info(self) -> str: - result: Dict[str, Any] = dict(dataclasses.asdict(self.server_args)) + result: Dict[str, Any] = dataclasses.asdict(self.server_args) result.update(self.scheduler_info) - return json.dumps(result, default=str) + return json.dumps(msgspec_to_builtins(result), default=str) def health_check(self) -> bool: from sglang.srt.managers.tokenizer_manager import ServerStatus diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 78a091f7a..9676c9ef4 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -173,6 +173,7 @@ from sglang.srt.utils.json_response import ( dumps_json, orjson_response, ) +from sglang.srt.utils.msgspec_utils import msgspec_to_builtins from sglang.srt.utils.watchdog import SubprocessWatchdog from sglang.utils import get_exception_traceback from sglang.version import __version__ @@ -698,16 +699,18 @@ async def server_info(): server_args = _global_state.tokenizer_manager.server_args # server_args.model_config is not serializable but should be excluded by asdict. - return { - **dataclasses.asdict(server_args), - **_global_state.scheduler_info, - "internal_states": internal_states, - "version": __version__, - # Structured KV-event publisher descriptor for KV-aware routers. - # `None` when publishing is disabled or misconfigured; see - # `ServerArgs.describe_kv_events_publisher` for the precise contract. - "kv_events": server_args.describe_kv_events_publisher(), - } + return msgspec_to_builtins( + { + **dataclasses.asdict(server_args), + **_global_state.scheduler_info, + "internal_states": internal_states, + "version": __version__, + # Structured KV-event publisher descriptor for KV-aware routers. + # `None` when publishing is disabled or misconfigured; see + # `ServerArgs.describe_kv_events_publisher` for the precise contract. + "kv_events": server_args.describe_kv_events_publisher(), + } + ) @app.get("/get_load") @@ -1417,7 +1420,7 @@ async def load_lora_adapter( """Load a new LoRA adapter without re-launching the server.""" result = await _global_state.tokenizer_manager.load_lora_adapter(obj, request) status_code = HTTPStatus.OK if result.success else HTTPStatus.BAD_REQUEST - return ORJSONResponse(result, status_code=status_code) + return ORJSONResponse(msgspec_to_builtins(result), status_code=status_code) @app.api_route("/load_lora_adapter_from_tensors", methods=["POST"]) @@ -1429,7 +1432,7 @@ async def load_lora_adapter_from_tensors( obj, request ) status_code = HTTPStatus.OK if result.success else HTTPStatus.BAD_REQUEST - return ORJSONResponse(result, status_code=status_code) + return ORJSONResponse(msgspec_to_builtins(result), status_code=status_code) @app.api_route("/unload_lora_adapter", methods=["POST"]) @@ -1440,7 +1443,7 @@ async def unload_lora_adapter( """Load a new LoRA adapter without re-launching the server.""" result = await _global_state.tokenizer_manager.unload_lora_adapter(obj, request) status_code = HTTPStatus.OK if result.success else HTTPStatus.BAD_REQUEST - return ORJSONResponse(result, status_code=status_code) + return ORJSONResponse(msgspec_to_builtins(result), status_code=status_code) @app.api_route("/open_session", methods=["GET", "POST"]) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index eb9455dc2..b950dbe16 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -241,6 +241,10 @@ class Envs: SGLANG_LOG_SCHEDULER_STATUS_TARGET = EnvStr("") SGLANG_LOG_SCHEDULER_STATUS_INTERVAL = EnvFloat(60.0) + # IPC + SGLANG_USE_PICKLE_IPC = EnvBool(True) + SGLANG_LOG_PICKLE_IPC_OBJECTS = EnvBool(False) + # SGLang CI SGLANG_IS_IN_CI = EnvBool(False) SGLANG_IS_IN_CI_AMD = EnvBool(False) diff --git a/python/sglang/srt/lora/lora_registry.py b/python/sglang/srt/lora/lora_registry.py index 08b7a2812..4e72efab5 100644 --- a/python/sglang/srt/lora/lora_registry.py +++ b/python/sglang/srt/lora/lora_registry.py @@ -15,16 +15,17 @@ import asyncio from collections import OrderedDict -from dataclasses import dataclass, field, fields from typing import Dict, List, Optional, Union from uuid import NAMESPACE_URL, uuid4, uuid5 +import msgspec +from msgspec.structs import fields + from sglang.srt.utils import ConcurrentCounter from sglang.srt.utils.aio_rwlock import RWLock -@dataclass(frozen=True) -class LoRARef: +class LoRARef(msgspec.Struct, frozen=True, array_like=True): """ Reference record for a LoRA model. @@ -33,7 +34,7 @@ class LoRARef: keys (e.g., radix cache). """ - lora_id: str = field(default_factory=lambda: uuid4().hex) + lora_id: str = msgspec.field(default_factory=lambda: uuid4().hex) lora_name: Optional[str] = None lora_path: Optional[str] = None pinned: Optional[bool] = None diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index 50c49ff35..2353ac51f 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -38,6 +38,8 @@ from sglang.srt.managers.io_struct import ( TokenizedGenerateReqInput, sock_recv, sock_send, + unwrap_from_pickle, + wrap_as_pickle, ) from sglang.srt.managers.load_snapshot import create_load_snapshot_reader from sglang.srt.managers.schedule_batch import Req @@ -238,10 +240,14 @@ class DataParallelController: if refresh_load_budget and self.refresh_load_budget_on_dispatch: self.refresh_load_budget() - req.time_stats = DPControllerReqTimeStats.new_from_obj(req.time_stats) + time_stats = DPControllerReqTimeStats.new_from_obj( + unwrap_from_pickle(req.time_stats) + ) - req.time_stats.set_dp_dispatch_time() + time_stats.set_dp_dispatch_time() + req.time_stats = wrap_as_pickle(time_stats) self.dispatching(req) + req.time_stats = time_stats req.time_stats.set_dp_dispatch_finish_time() def dispatch_batch_generate(self, batch_req: BatchTokenizedGenerateReqInput): @@ -382,11 +388,11 @@ class DataParallelController: connected_clients = 0 while connected_clients < expected_clients: # Wait for client handshake - client_rank = sock_recv(rep_socket).decode() + client_rank = sock_recv(rep_socket) logger.debug(f"Received handshake from node {client_rank}") # Send worker ports to client - sock_send(rep_socket, worker_ports) + sock_send(rep_socket, wrap_as_pickle(worker_ports)) connected_clients += 1 logger.debug( f"Sent worker ports to {connected_clients}/{expected_clients} nodes" @@ -411,7 +417,7 @@ class DataParallelController: while True: # Wait for client handshake try: - client_rank = sock_recv(rep_socket).decode() + client_rank = sock_recv(rep_socket) except Exception: logger.exception( "Failed to recv/decode handshake in reply thread; continue" @@ -420,7 +426,7 @@ class DataParallelController: logger.debug(f"Received handshake from node {client_rank}") # Send worker ports to client - sock_send(rep_socket, worker_ports) + sock_send(rep_socket, wrap_as_pickle(worker_ports)) logger.debug(f"Sent worker ports to node {client_rank}") def _receive_ports_as_client(self, endpoint: str, node_rank: int) -> List[int]: @@ -433,7 +439,7 @@ class DataParallelController: try: # Send handshake with our node rank - sock_send(req_socket, str(node_rank).encode()) + sock_send(req_socket, wrap_as_pickle(str(node_rank))) # Receive worker ports worker_ports = sock_recv(req_socket) diff --git a/python/sglang/srt/managers/embed_types.py b/python/sglang/srt/managers/embed_types.py index 010e18e56..1846f3c39 100644 --- a/python/sglang/srt/managers/embed_types.py +++ b/python/sglang/srt/managers/embed_types.py @@ -12,20 +12,19 @@ # limitations under the License. # ============================================================================== """ -Dataclasses for embedding injection. +Structs for embedding injection. These are placed in a separate module to avoid circular imports between io_struct.py and schedule_batch.py. """ -from dataclasses import dataclass from typing import List +import msgspec import torch -@dataclass -class PositionalEmbeds: +class PositionalEmbeds(msgspec.Struct, array_like=True): """Embeddings to place at specific token positions. Accepts either a list of [1, hidden_dim] tensors or a pre-stacked [N, hidden_dim] tensor. diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 2ea058648..f8fdad62c 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -23,6 +23,8 @@ instead, such as sglang.srt.utils.common. from __future__ import annotations import copy +import logging +import pickle import uuid from array import array from collections import Counter @@ -36,26 +38,29 @@ from typing import ( List, Literal, Optional, + Type, Union, ) +import msgspec +import numpy as np import torch import zmq import zmq.asyncio from pydantic import PlainValidator +from sglang.srt.environ import envs from sglang.srt.lora.lora_registry import LoRARef from sglang.srt.managers.embed_types import PositionalEmbeds -from sglang.srt.managers.schedule_batch import Modality, MultimodalInputs +from sglang.srt.managers.schedule_batch import Modality from sglang.srt.multimodal.mm_utils import has_valid_data -from sglang.srt.observability.req_time_stats import ( - APIServerReqTimeStats, - DPControllerReqTimeStats, - SchedulerReqTimeStats, -) from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.utils import ImageData, VideoData from sglang.srt.utils.field_validators import validate_optional_list_i64_1d_2d +from sglang.srt.utils.msgspec_utils import ( + Base64Bytes, + msgspec_struct_pydantic_core_schema, +) # Handle serialization of Image for pydantic if TYPE_CHECKING: @@ -63,27 +68,46 @@ if TYPE_CHECKING: else: Image = Any - -@dataclass -class BaseReq: - rid: Optional[str] = field(default=None, kw_only=True) - http_worker_ipc: Optional[str] = field(default=None, kw_only=True) +logger = logging.getLogger(__name__) -@dataclass -class BaseBatchReq: - rids: Optional[List[str]] = field(default=None, kw_only=True) - http_worker_ipcs: Optional[List[Optional[str]]] = field(default=None, kw_only=True) +class BaseReq(msgspec.Struct, tag=True, kw_only=True, array_like=True): + """Base for single-request IPC payloads.""" - def regenerate_rids(self): - """Generate new request IDs and return them.""" - self.rids = [uuid.uuid4().hex for _ in range(len(self.rids))] - return self.rids + rid: Optional[str] = None + http_worker_ipc: Optional[str] = None + + @classmethod + def __get_pydantic_core_schema__(cls, source, handler): + return msgspec_struct_pydantic_core_schema(cls, handler) + + +class BaseBatchReq(msgspec.Struct, tag=True, kw_only=True, array_like=True): + """Base for batched IPC payloads.""" + + rids: Optional[List[str]] = None + http_worker_ipcs: Optional[List[Optional[str]]] = None + + @classmethod + def __get_pydantic_core_schema__(cls, source, handler): + return msgspec_struct_pydantic_core_schema(cls, handler) + + +class PickleWrapper(msgspec.Struct, tag=True, array_like=True): + """Wraps an arbitrary Python object as pickle-serialized bytes for msgpack IPC. + + In msgpack mode, fields that carry opaque or non-msgspec-typed payloads + (e.g. multimodal inputs, time stats, customized info) are stored as + PickleWrapper so the outer struct can still be msgpack-encoded. In pickle + mode (_USE_PICKLE_IPC=True), wrap_as_pickle / unwrap_from_pickle are no-ops + and this class is not used on the wire. + """ + + data: bytes # Parameters for a session -@dataclass -class SessionParams: +class SessionParams(msgspec.Struct, kw_only=True, array_like=True): # The session identifier. Used by the scheduler to look up or create the # Session object that groups all requests in a multi-turn conversation. id: Optional[str] = None @@ -741,16 +765,14 @@ class GenerateReqInput: return sub -@dataclass -class TokenizedGenerateReqInput(BaseReq): - # The input text +class TokenizedGenerateReqInput(BaseReq, kw_only=True): input_text: Optional[Union[str, List[Union[str, List[str]]]]] # The input token ids input_ids: Optional[array] # Optional[array[int]] # The input embeds input_embeds: Optional[List[List[float]]] # The multimodal inputs - mm_inputs: Optional[MultimodalInputs] + mm_inputs: Optional[PickleWrapper] # Pickled Optional[MultimodalProcessorOutput] token_type_ids: Optional[List[int]] # The sampling parameters sampling_params: SamplingParams @@ -820,7 +842,10 @@ class TokenizedGenerateReqInput(BaseReq): need_wait_for_mm_inputs: Optional[bool] = None num_items_assigned: Optional[Dict[Modality, List[int]]] = None - mm_data_mooncake: Optional[List[Any]] = None + # Pickled Optional[List[{"url": MultimodalDataInputItem, "modality": Modality}]] + # from MMReceiverBase._extract_url_data. "url" is ImageData.url, + # dict["url"] when present, or the original raw multimodal item. + mm_data_mooncake: Optional[PickleWrapper] = None # Encoder URL snapshot frozen at tokenizer-side dispatch time so that # encoder_idx assignments stay consistent in the scheduler subprocess. # Internal IPC only. @@ -830,11 +855,21 @@ class TokenizedGenerateReqInput(BaseReq): multi_item_delimiter_indices: Optional[List[int]] = None # For observability - time_stats: Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] = None + # Pickled Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] + time_stats: Optional[PickleWrapper] = None + + def wrap_pickle_fields(self): + self.mm_inputs = wrap_as_pickle(self.mm_inputs) + self.mm_data_mooncake = wrap_as_pickle(self.mm_data_mooncake) + self.time_stats = wrap_as_pickle(self.time_stats) + + def unwrap_pickle_fields(self): + self.mm_inputs = unwrap_from_pickle(self.mm_inputs) + self.mm_data_mooncake = unwrap_from_pickle(self.mm_data_mooncake) + self.time_stats = unwrap_from_pickle(self.time_stats) -@dataclass -class BatchTokenizedGenerateReqInput(BaseBatchReq): +class BatchTokenizedGenerateReqInput(BaseBatchReq, kw_only=True): # The batch of tokenized requests batch: List[TokenizedGenerateReqInput] @@ -1081,14 +1116,12 @@ class EmbeddingReqInput: return sub -@dataclass -class TokenizedEmbeddingReqInput(BaseReq): - # The input text +class TokenizedEmbeddingReqInput(BaseReq, kw_only=True): input_text: Optional[Union[str, List[Union[str, List[str]]]]] # The input token ids input_ids: Optional[array] # array[int] # The multimodal inputs - mm_inputs: Optional[MultimodalInputs] + mm_inputs: Optional[PickleWrapper] # Pickled Optional[MultimodalProcessorOutput] # The token type ids token_type_ids: Optional[List[int]] # Dummy sampling params for compatibility @@ -1109,11 +1142,19 @@ class TokenizedEmbeddingReqInput(BaseReq): multi_item_delimiter_indices: Optional[List[int]] = None # For observability - time_stats: Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] = None + # Pickled Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] + time_stats: Optional[PickleWrapper] = None + + def wrap_pickle_fields(self): + self.mm_inputs = wrap_as_pickle(self.mm_inputs) + self.time_stats = wrap_as_pickle(self.time_stats) + + def unwrap_pickle_fields(self): + self.mm_inputs = unwrap_from_pickle(self.mm_inputs) + self.time_stats = unwrap_from_pickle(self.time_stats) -@dataclass -class BatchTokenizedEmbeddingReqInput(BaseBatchReq): +class BatchTokenizedEmbeddingReqInput(BaseBatchReq, kw_only=True): # The batch of tokenized embedding requests batch: List[TokenizedEmbeddingReqInput] @@ -1140,8 +1181,7 @@ CachedTokensDetails = Dict[str, Union[int, str]] FinishReasonDict = Dict[str, Optional[Union[str, int, List[int]]]] -@dataclass -class BatchTokenIDOutput(BaseBatchReq): +class BatchTokenIDOutput(BaseBatchReq, kw_only=True): # The finish reason finished_reasons: List[Optional[FinishReasonDict]] # For incremental decoding @@ -1200,14 +1240,15 @@ class BatchTokenIDOutput(BaseBatchReq): token_steps: Optional[List[List[int]]] = None # Customized info - customized_info: Optional[Dict[str, List[Any]]] = None + customized_info: Optional[PickleWrapper] = None # Detailed breakdown of cached tokens by source (device/host/storage) cached_tokens_details: Optional[List[Optional[CachedTokensDetails]]] = None # DP rank of the scheduler that processed each request dp_ranks: Optional[List[Optional[int]]] = None # For observability - time_stats: Optional[List[SchedulerReqTimeStats]] = None + # Pickled Optional[List[SchedulerReqTimeStats]] + time_stats: Optional[PickleWrapper] = None # Multimodal prompt token counts (image/audio/video). None when not applicable. image_tokens: Optional[List[int]] = None @@ -1222,8 +1263,7 @@ class BatchTokenIDOutput(BaseBatchReq): spec_correct_drafts_histogram: Optional[List[List[int]]] = None -@dataclass -class BatchStrOutput(BaseBatchReq): +class BatchStrOutput(BaseBatchReq, kw_only=True): # The finish reason finished_reasons: List[Optional[FinishReasonDict]] # The output decoded strings @@ -1275,14 +1315,15 @@ class BatchStrOutput(BaseBatchReq): token_steps: Optional[List[List[int]]] = None # Customized info - customized_info: Optional[Dict[str, List[Any]]] = None + customized_info: Optional[PickleWrapper] = None # Detailed breakdown of cached tokens by source (device/host/storage) cached_tokens_details: Optional[List[Optional[CachedTokensDetails]]] = None # DP rank of the scheduler that processed each request dp_ranks: Optional[List[Optional[int]]] = None # For observability - time_stats: Optional[List[SchedulerReqTimeStats]] = None + # Pickled Optional[List[SchedulerReqTimeStats]] + time_stats: Optional[PickleWrapper] = None # Multimodal prompt token counts (image/audio/video). None when not applicable. image_tokens: Optional[List[int]] = None @@ -1297,8 +1338,7 @@ class BatchStrOutput(BaseBatchReq): spec_correct_drafts_histogram: Optional[List[List[int]]] = None -@dataclass -class BatchEmbeddingOutput(BaseBatchReq): +class BatchEmbeddingOutput(BaseBatchReq, kw_only=True): # The finish reason finished_reasons: List[Optional[FinishReasonDict]] # The output embedding @@ -1316,7 +1356,8 @@ class BatchEmbeddingOutput(BaseBatchReq): cached_tokens_details: Optional[List[Optional[CachedTokensDetails]]] = None # For observability - time_stats: Optional[List[SchedulerReqTimeStats]] = None + # Pickled Optional[List[SchedulerReqTimeStats]] + time_stats: Optional[PickleWrapper] = None # Optional pooled hidden states (pre-head transformer output). # Two IPC formats, disambiguated by len vs len(rids): @@ -1325,68 +1366,57 @@ class BatchEmbeddingOutput(BaseBatchReq): pooled_hidden_states: Optional[List[Optional[torch.Tensor]]] = None -@dataclass -class ClearHiCacheReqInput(BaseReq): +class ClearHiCacheReqInput(BaseReq, kw_only=True): pass -@dataclass -class ClearHiCacheReqOutput(BaseReq): +class ClearHiCacheReqOutput(BaseReq, kw_only=True): success: bool -@dataclass -class FlushCacheReqInput(BaseReq): +class FlushCacheReqInput(BaseReq, kw_only=True): timeout_s: Optional[float] = None -@dataclass -class FlushCacheReqOutput(BaseReq): +class FlushCacheReqOutput(BaseReq, kw_only=True): success: bool message: str = "" -@dataclass -class AddExternalCorpusReqInput(BaseReq): +class AddExternalCorpusReqInput(BaseReq, kw_only=True): corpus_id: Optional[str] = None file_path: Optional[str] = None documents: Optional[List[str]] = None token_chunks: Optional[List[List[int]]] = None -@dataclass -class AddExternalCorpusReqOutput(BaseReq): +class AddExternalCorpusReqOutput(BaseReq, kw_only=True): success: bool corpus_id: str = "" message: str = "" loaded_token_count: int = 0 -@dataclass -class RemoveExternalCorpusReqInput(BaseReq): +class RemoveExternalCorpusReqInput(BaseReq, kw_only=True): corpus_id: str -@dataclass -class RemoveExternalCorpusReqOutput(BaseReq): +class RemoveExternalCorpusReqOutput(BaseReq, kw_only=True): success: bool message: str = "" -@dataclass -class ListExternalCorporaReqInput(BaseReq): +class ListExternalCorporaReqInput(BaseReq, kw_only=True): pass -@dataclass -class ListExternalCorporaReqOutput(BaseReq): +class ListExternalCorporaReqOutput(BaseReq, kw_only=True): success: bool - corpus_token_counts: Dict[str, int] = field(default_factory=dict) + corpus_token_counts: Dict[str, int] = msgspec.field(default_factory=dict) message: str = "" -@dataclass -class AttachHiCacheStorageReqInput(BaseReq): +class AttachHiCacheStorageReqInput(BaseReq, kw_only=True): """Dynamically attach (enable) HiCache storage backend at runtime. Note: `hicache_storage_backend_extra_config_json` is a JSON string. It may contain both: @@ -1400,27 +1430,23 @@ class AttachHiCacheStorageReqInput(BaseReq): hicache_write_policy: Optional[str] = None -@dataclass -class AttachHiCacheStorageReqOutput(BaseReq): +class AttachHiCacheStorageReqOutput(BaseReq, kw_only=True): success: bool message: str = "" -@dataclass -class DetachHiCacheStorageReqInput(BaseReq): +class DetachHiCacheStorageReqInput(BaseReq, kw_only=True): """Dynamically detach (disable) HiCache storage backend at runtime.""" pass -@dataclass -class DetachHiCacheStorageReqOutput(BaseReq): +class DetachHiCacheStorageReqOutput(BaseReq, kw_only=True): success: bool message: str = "" -@dataclass -class PauseGenerationReqInput(BaseReq): +class PauseGenerationReqInput(BaseReq, kw_only=True): """ Note that the PauseGenerationRequests is only supported in SGLang Server. abort: Abort and return all requests currently being processed. @@ -1442,8 +1468,7 @@ class PauseGenerationReqInput(BaseReq): mode: Literal["abort", "retract", "in_place"] = "abort" -@dataclass -class ContinueGenerationReqInput(BaseReq): +class ContinueGenerationReqInput(BaseReq, kw_only=True): # Call torch.cuda.empty_cache() before un-pausing. Returns blocks # cached by the PyTorch allocator (left over from transient allocs # during post-weight-update processing) back to the driver before @@ -1452,22 +1477,19 @@ class ContinueGenerationReqInput(BaseReq): torch_empty_cache: bool = True -@dataclass -class TokenizerWorkerRegistrationReq(BaseReq): +class TokenizerWorkerRegistrationReq(BaseReq, kw_only=True): """Sent by each TokenizerWorker on startup to register its IPC name with the router.""" worker_ipc_name: str -@dataclass -class PauseContinueBroadcastReq(BaseReq): +class PauseContinueBroadcastReq(BaseReq, kw_only=True): """Broadcast from router to all workers to set is_pause state.""" is_pause: bool -@dataclass -class UpdateWeightFromDiskReqInput(BaseReq): +class UpdateWeightFromDiskReqInput(BaseReq, kw_only=True): # The model path with the new weights model_path: str # The format to load the weights @@ -1492,16 +1514,14 @@ class UpdateWeightFromDiskReqInput(BaseReq): manifest: Optional[Dict[str, Any]] = None -@dataclass -class UpdateWeightFromDiskReqOutput(BaseReq): +class UpdateWeightFromDiskReqOutput(BaseReq, kw_only=True): success: bool message: str # Number of paused requests during weight sync. num_paused_requests: int = 0 -@dataclass -class UpdateWeightsFromDistributedReqInput(BaseReq): +class UpdateWeightsFromDistributedReqInput(BaseReq, kw_only=True): names: List[str] dtypes: List[str] shapes: List[List[int]] @@ -1519,25 +1539,20 @@ class UpdateWeightsFromDistributedReqInput(BaseReq): torch_empty_cache: bool = False -@dataclass -class UpdateWeightsFromDistributedReqOutput(BaseReq): +class UpdateWeightsFromDistributedReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class UpdateWeightsFromTensorReqInput(BaseReq): - """Update model weights from tensor input. +class UpdateWeightsFromTensorReqInput(BaseReq, kw_only=True): + """Internal IPC request for updating model weights from serialized tensors.""" - - Tensors are serialized for transmission - - Data is structured in JSON for easy transmission over HTTP - """ - - # Accepts both base64 str (from HTTP/JSON, which has no bytes type) and - # raw bytes (from the Python Engine API / MultiprocessingSerializer). - # Normalized to List[bytes] by normalize_serialized_named_tensor_payloads - # in tokenizer_control_mixin before forwarding over scheduler IPC. - serialized_named_tensors: List[Union[str, bytes]] + # Serialized named tensors, normalized to raw MultiprocessingSerializer + # bytes before scheduler IPC. Python Engine callers construct this field + # with bytes directly. FastAPI HTTP callers send base64 strings because JSON + # has no bytes type; the Annotated Base64Bytes marker is used only by the + # msgspec-to-Pydantic schema for the HTTP protocol to decode those strings. + serialized_named_tensors: Annotated[List[bytes], Base64Bytes()] # Optional format specification for loading load_format: Optional[str] = None # Whether to flush the cache after updating weights @@ -1552,14 +1567,12 @@ class UpdateWeightsFromTensorReqInput(BaseReq): torch_empty_cache: bool = False -@dataclass -class UpdateWeightsFromTensorReqOutput(BaseReq): +class UpdateWeightsFromTensorReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class InitWeightsSendGroupForRemoteInstanceReqInput(BaseReq): +class InitWeightsSendGroupForRemoteInstanceReqInput(BaseReq, kw_only=True): # The master address master_address: str # The ports for each rank's communication group @@ -1576,8 +1589,7 @@ class InitWeightsSendGroupForRemoteInstanceReqInput(BaseReq): # Now UpdateWeightsFromIPCReqInput and UpdateWeightsFromIPCReqOutput # are only used by Checkpoint Engine (https://github.com/MoonshotAI/checkpoint-engine) -@dataclass -class UpdateWeightsFromIPCReqInput(BaseReq): +class UpdateWeightsFromIPCReqInput(BaseReq, kw_only=True): # ZMQ socket paths for each device UUID zmq_handles: Dict[str, str] # Whether to flush cache after weight update @@ -1588,20 +1600,17 @@ class UpdateWeightsFromIPCReqInput(BaseReq): torch_empty_cache: bool = False -@dataclass -class UpdateWeightsFromIPCReqOutput(BaseReq): +class UpdateWeightsFromIPCReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class InitWeightsSendGroupForRemoteInstanceReqOutput(BaseReq): +class InitWeightsSendGroupForRemoteInstanceReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class SendWeightsToRemoteInstanceReqInput(BaseReq): +class SendWeightsToRemoteInstanceReqInput(BaseReq, kw_only=True): # The master address master_address: str # The ports for each rank's communication group @@ -1610,27 +1619,23 @@ class SendWeightsToRemoteInstanceReqInput(BaseReq): group_name: str = "weight_send_group" -@dataclass -class SendWeightsToRemoteInstanceReqOutput(BaseReq): +class SendWeightsToRemoteInstanceReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class UpdateExpertBackupReq(BaseReq): +class UpdateExpertBackupReq(BaseReq, kw_only=True): pass -@dataclass -class BackupDramReq(BaseReq): +class BackupDramReq(BaseReq, kw_only=True): rank: int weight_pointer_map: Dict[str, Any] session_id: str buffer_size: int -@dataclass -class InitWeightsUpdateGroupReqInput(BaseReq): +class InitWeightsUpdateGroupReqInput(BaseReq, kw_only=True): # The master address master_address: str # The master port @@ -1645,90 +1650,75 @@ class InitWeightsUpdateGroupReqInput(BaseReq): backend: str = "nccl" -@dataclass -class InitWeightsUpdateGroupReqOutput(BaseReq): +class InitWeightsUpdateGroupReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class DestroyWeightsUpdateGroupReqInput(BaseReq): +class DestroyWeightsUpdateGroupReqInput(BaseReq, kw_only=True): group_name: str = "weight_update_group" -@dataclass -class DestroyWeightsUpdateGroupReqOutput(BaseReq): +class DestroyWeightsUpdateGroupReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class UpdateWeightVersionReqInput(BaseReq): +class UpdateWeightVersionReqInput(BaseReq, kw_only=True): # The new weight version new_version: str # Whether to abort all running requests before updating abort_all_requests: bool = True -@dataclass -class GetWeightsByNameReqInput(BaseReq): +class GetWeightsByNameReqInput(BaseReq, kw_only=True): name: str truncate_size: int = 100 -@dataclass -class GetWeightsByNameReqOutput(BaseReq): +class GetWeightsByNameReqOutput(BaseReq, kw_only=True): parameter: Optional[List[Any]] -@dataclass -class ReleaseMemoryOccupationReqInput(BaseReq): +class ReleaseMemoryOccupationReqInput(BaseReq, kw_only=True): # Optional tags to identify the memory region, which is primarily used for RL # Currently we only support `weights` and `kv_cache` tags: Optional[List[str]] = None -@dataclass -class ReleaseMemoryOccupationReqOutput(BaseReq): +class ReleaseMemoryOccupationReqOutput(BaseReq, kw_only=True): pass -@dataclass -class ResumeMemoryOccupationReqInput(BaseReq): +class ResumeMemoryOccupationReqInput(BaseReq, kw_only=True): # Optional tags to identify the memory region, which is primarily used for RL # Currently we only support `weights` and `kv_cache` tags: Optional[List[str]] = None -@dataclass -class ResumeMemoryOccupationReqOutput(BaseReq): +class ResumeMemoryOccupationReqOutput(BaseReq, kw_only=True): pass -@dataclass -class CheckWeightsReqInput(BaseReq): +class CheckWeightsReqInput(BaseReq, kw_only=True): action: str = "checksum" -@dataclass -class CheckWeightsReqOutput(BaseReq): +class CheckWeightsReqOutput(BaseReq, kw_only=True): success: bool message: str payload: Optional[Dict[str, Any]] = None -@dataclass -class SlowDownReqInput(BaseReq): +class SlowDownReqInput(BaseReq, kw_only=True): forward_sleep_time: Optional[float] -@dataclass -class SlowDownReqOutput(BaseReq): +class SlowDownReqOutput(BaseReq, kw_only=True): pass -@dataclass -class AbortReq(BaseReq): +class AbortReq(BaseReq, kw_only=True): # Whether to abort all requests abort_all: bool = False # The finished reason data (from BaseFinishReason.to_json()) @@ -1741,28 +1731,23 @@ class AbortReq(BaseReq): self.rid = "" -@dataclass -class ActiveRanksOutput(BaseReq): +class ActiveRanksOutput(BaseReq, kw_only=True): status: List[bool] -@dataclass -class GetInternalStateReq(BaseReq): +class GetInternalStateReq(BaseReq, kw_only=True): pass -@dataclass -class GetInternalStateReqOutput(BaseReq): +class GetInternalStateReqOutput(BaseReq, kw_only=True): internal_state: Dict[str, Any] -@dataclass -class SetInternalStateReq(BaseReq): +class SetInternalStateReq(BaseReq, kw_only=True): server_args: Dict[str, Any] -@dataclass -class SetInternalStateReqOutput(BaseReq): +class SetInternalStateReqOutput(BaseReq, kw_only=True): updated: bool server_args: Dict[str, Any] @@ -1772,8 +1757,7 @@ class ProfileReqType(Enum): STOP_PROFILE = 2 -@dataclass -class ProfileReq(BaseReq): +class ProfileReq(BaseReq, kw_only=True): req_type: ProfileReqType = ProfileReqType.START_PROFILE # The output directory output_dir: Optional[str] = None @@ -1800,26 +1784,22 @@ class ProfileReq(BaseReq): profile_stages: Optional[List[str]] = None -@dataclass -class ProfileReqOutput(BaseReq): +class ProfileReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class FreezeGCReq(BaseReq): +class FreezeGCReq(BaseReq, kw_only=True): pass -@dataclass -class ShutdownReq(BaseReq): +class ShutdownReq(BaseReq, kw_only=True): # Broadcast across TP ranks via the normal recv path, so all ranks break # the scheduler loop on the same iteration. pass -@dataclass -class ConfigureLoggingReq(BaseReq): +class ConfigureLoggingReq(BaseReq, kw_only=True): log_requests: Optional[bool] = None log_requests_level: Optional[int] = None log_requests_format: Optional[str] = None @@ -1830,27 +1810,23 @@ class ConfigureLoggingReq(BaseReq): dump_requests_exclude_meta_keys: Optional[List[str]] = None -@dataclass -class OpenSessionReqInput(BaseReq): +class OpenSessionReqInput(BaseReq, kw_only=True): capacity_of_str_len: int session_id: Optional[str] = None streaming: Optional[bool] = None timeout: Optional[float] = None -@dataclass -class CloseSessionReqInput(BaseReq): +class CloseSessionReqInput(BaseReq, kw_only=True): session_id: str -@dataclass -class OpenSessionReqOutput(BaseReq): +class OpenSessionReqOutput(BaseReq, kw_only=True): session_id: Optional[str] success: bool -@dataclass -class HealthCheckOutput(BaseReq): +class HealthCheckOutput(BaseReq, kw_only=True): pass @@ -1860,33 +1836,36 @@ class ExpertDistributionReqType(Enum): DUMP_RECORD = 3 -@dataclass -class ExpertDistributionReq(BaseReq): +class ExpertDistributionReq(BaseReq, kw_only=True): action: ExpertDistributionReqType -@dataclass -class ExpertDistributionReqOutput(BaseReq): +class ExpertDistributionReqOutput(BaseReq, kw_only=True): pass -@dataclass -class Function: +class Function(msgspec.Struct, kw_only=True, array_like=True): description: Optional[str] = None name: Optional[str] = None parameters: Optional[Dict[str, Any]] = None + @classmethod + def __get_pydantic_core_schema__(cls, source, handler): + return msgspec_struct_pydantic_core_schema(cls, handler) -@dataclass -class Tool: + +class Tool(msgspec.Struct, kw_only=True, array_like=True): function: Function type: str = "function" + @classmethod + def __get_pydantic_core_schema__(cls, source, handler): + return msgspec_struct_pydantic_core_schema(cls, handler) -@dataclass -class ParseFunctionCallReq(BaseReq): + +class ParseFunctionCallReq(BaseReq, kw_only=True): text: str # The text to parse. - tools: List[Tool] = field( + tools: List[Tool] = msgspec.field( default_factory=list ) # A list of available function tools (name, parameters, etc.). tool_call_parser: Optional[str] = ( @@ -1894,33 +1873,28 @@ class ParseFunctionCallReq(BaseReq): ) -@dataclass -class SeparateReasoningReqInput(BaseReq): +class SeparateReasoningReqInput(BaseReq, kw_only=True): text: str # The text to parse. reasoning_parser: str # Specify the parser type, e.g., "deepseek-r1". return_blocks: bool = False # If True, also return segmented reasoning blocks. -@dataclass -class VertexGenerateReqInput(BaseReq): +class VertexGenerateReqInput(BaseReq, kw_only=True): instances: List[Dict[str, Any]] parameters: Optional[Dict[str, Any]] = None -@dataclass -class RpcReqInput(BaseReq): +class RpcReqInput(BaseReq, kw_only=True): method: str parameters: Optional[Dict[str, Any]] = None -@dataclass -class RpcReqOutput(BaseReq): +class RpcReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class LoadLoRAAdapterReqInput(BaseReq): +class LoadLoRAAdapterReqInput(BaseReq, kw_only=True): # The name of the lora module to newly loaded. lora_name: str # The path of loading. @@ -1939,8 +1913,7 @@ class LoadLoRAAdapterReqInput(BaseReq): ) -@dataclass -class UnloadLoRAAdapterReqInput(BaseReq): +class UnloadLoRAAdapterReqInput(BaseReq, kw_only=True): # The name of lora module to unload. lora_name: str # The unique identifier for the LoRA adapter, which automatically generated in the `TokenizerManager`. @@ -1953,8 +1926,7 @@ class UnloadLoRAAdapterReqInput(BaseReq): ) -@dataclass -class LoadLoRAAdapterFromTensorsReqInput(BaseReq): +class LoadLoRAAdapterFromTensorsReqInput(BaseReq, kw_only=True): lora_name: str config_dict: Dict[str, Any] serialized_tensors: str @@ -1972,8 +1944,7 @@ class LoadLoRAAdapterFromTensorsReqInput(BaseReq): ) -@dataclass -class LoRAUpdateOutput(BaseReq): +class LoRAUpdateOutput(BaseReq, kw_only=True): success: bool error_message: Optional[str] = None loaded_adapters: Optional[Dict[str, Union[str, LoRARef]]] = None @@ -1989,13 +1960,11 @@ class BlockReqType(Enum): UNBLOCK = 2 -@dataclass -class BlockReqInput(BaseReq): +class BlockReqInput(BaseReq, kw_only=True): req_type: BlockReqType -@dataclass -class MemoryMetrics: +class MemoryMetrics(msgspec.Struct, array_like=True): """Memory breakdown metrics.""" weight_gb: float @@ -2004,16 +1973,14 @@ class MemoryMetrics: token_capacity: int -@dataclass -class SpeculativeMetrics: +class SpeculativeMetrics(msgspec.Struct, array_like=True): """Speculative decoding metrics.""" accept_length: float accept_rate: float -@dataclass -class LoRAMetrics: +class LoRAMetrics(msgspec.Struct, array_like=True): """LoRA adapter pool metrics.""" slots_used: int @@ -2021,8 +1988,7 @@ class LoRAMetrics: utilization: float -@dataclass -class DisaggregationMetrics: +class DisaggregationMetrics(msgspec.Struct, array_like=True): """PD disaggregation metrics.""" mode: str # "prefill", "decode", or "null" @@ -2035,9 +2001,8 @@ class DisaggregationMetrics: kv_transfer_latency_ms: float = 0.0 -@dataclass -class QueueMetrics: - """Detailed queue breakdown.""" +class QueueMetrics(msgspec.Struct, array_like=True): + """Detailed queue info breakdown.""" waiting: int grammar: int @@ -2045,15 +2010,14 @@ class QueueMetrics: retracted: int -@dataclass -class GetLoadsReqInput(BaseReq): +class GetLoadsReqInput(BaseReq, kw_only=True): """Request for /v1/loads endpoint.""" VALID_SECTIONS = frozenset( {"core", "memory", "spec", "lora", "disagg", "queues", "all"} ) - include: List[str] = field(default_factory=lambda: ["all"]) + include: List[str] = msgspec.field(default_factory=lambda: ["all"]) dp_rank: Optional[int] = None def __post_init__(self): @@ -2067,8 +2031,7 @@ class GetLoadsReqInput(BaseReq): ) -@dataclass -class GetLoadsReqOutput(BaseReq): +class GetLoadsReqOutput(BaseReq, kw_only=True): """Per-DP-rank load metrics for /v1/loads endpoint.""" dp_rank: int @@ -2097,55 +2060,33 @@ class GetLoadsReqOutput(BaseReq): queues: Optional[QueueMetrics] = None -@dataclass -class SetInjectDumpMetadataReqInput(BaseReq): +class SetInjectDumpMetadataReqInput(BaseReq, kw_only=True): dump_metadata: Dict[str, Any] -@dataclass -class SetInjectDumpMetadataReqOutput(BaseReq): +class SetInjectDumpMetadataReqOutput(BaseReq, kw_only=True): success: bool -@dataclass -class LazyDumpTensorsReqInput(BaseReq): +class LazyDumpTensorsReqInput(BaseReq, kw_only=True): pass -@dataclass -class LazyDumpTensorsReqOutput(BaseReq): +class LazyDumpTensorsReqOutput(BaseReq, kw_only=True): success: bool -@dataclass -class DumperControlReqInput(BaseReq): +class DumperControlReqInput(BaseReq, kw_only=True): method: str body: Dict[str, Any] -@dataclass -class DumperControlReqOutput(BaseReq): +class DumperControlReqOutput(BaseReq, kw_only=True): success: bool response: List[Dict[str, Any]] error: str = "" -def sock_send(socket: zmq.Socket, obj: Any, flags: int = 0) -> None: - socket.send_pyobj(obj, flags=flags) - - -def sock_recv(socket: zmq.Socket, flags: int = 0) -> Any: - return socket.recv_pyobj(flags=flags) - - -async def async_sock_send(socket: zmq.asyncio.Socket, obj: Any, flags: int = 0) -> None: - await socket.send_pyobj(obj, flags=flags) - - -async def async_sock_recv(socket: zmq.asyncio.Socket, flags: int = 0) -> Any: - return await socket.recv_pyobj(flags=flags) - - # The following request types are either defined in other files, # or not subclasses of BaseReq/BaseBatchReq, so we skip the check for them. _IGNORE_REQ_TYPES_CHECK = ( @@ -2182,12 +2123,12 @@ 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. Kept as an explicit registry so -# opaque usage can be audited and gradually narrowed. +# 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 = ( - TokenizedGenerateReqInput, # mm_data_mooncake: Optional[List[Any]] +_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]] @@ -2201,6 +2142,159 @@ _REQ_TYPES_WITH_OPAQUE_FIELDS = ( SetInjectDumpMetadataReqInput, # dump_metadata: Dict[str, Any] DumperControlReqInput, # body: Dict[str, Any] DumperControlReqOutput, # response: List[Dict[str, Any]] - BatchTokenIDOutput, # customized_info: Optional[Dict[str, List[Any]]] - BatchStrOutput, # customized_info: Optional[Dict[str, List[Any]]] ) + + +def wrap_as_pickle(obj: object) -> object: + if obj is None: + return None + if _USE_PICKLE_IPC: + return obj + return PickleWrapper(pickle.dumps(obj)) + + +def unwrap_from_pickle(obj: Optional[object]) -> Optional[object]: + if obj is None: + return None + if _USE_PICKLE_IPC: + return obj + assert isinstance(obj, PickleWrapper) + return pickle.loads(obj.data) + + +def enc_hook(obj: Any) -> Any: + if isinstance(obj, array): + return (obj.typecode, obj.tobytes()) + elif isinstance(obj, torch.Tensor): + tensor_dtype = str(obj.dtype).removeprefix("torch.") + raw_data = ( + obj.cpu().contiguous().reshape(-1).view(torch.uint8).numpy().tobytes() + ) + return (obj.shape, tensor_dtype, raw_data) + elif isinstance(obj, np.ndarray): + raw_data = np.ascontiguousarray(obj).reshape(-1).view(np.uint8).data + return (obj.shape, obj.dtype.str, raw_data) + elif isinstance(obj, np.floating): + return float(obj) + else: + raise TypeError( + f"Cannot msgpack encode object of type {type(obj)} with enc_hook. " + "Use an explicit PickleWrapper field via wrap_as_pickle(...) for " + "arbitrary payloads, or add a dedicated enc_hook/dec_hook branch " + "for this transport type." + ) + + +def dec_hook(tp: Type, obj: Any) -> Any: + if tp is array: + typecode, raw_data = obj + res = array(typecode) + res.frombytes(raw_data) + return res + elif tp is torch.Tensor: + shape, dtype, data = obj + tensor_dtype = getattr(torch, dtype) + if len(data) == 0: + return torch.empty(shape, dtype=tensor_dtype) + return torch.frombuffer(bytearray(data), dtype=tensor_dtype).reshape(shape) + elif tp is np.ndarray: + shape, dtype, data = obj + return np.frombuffer(data, dtype=np.dtype(dtype)).copy().reshape(shape) + else: + raise TypeError( + f"Cannot msgpack decode object of type {type(obj)} as {tp} with " + "dec_hook. Use an explicit PickleWrapper field via wrap_as_pickle(...) " + "and unwrap_from_pickle(...) for arbitrary payloads, or add a " + "dedicated enc_hook/dec_hook branch for this transport type." + ) + + +_struct_types = tuple( + cls + for cls in BaseReq.__subclasses__() + + BaseBatchReq.__subclasses__() + + [PickleWrapper] +) +# Primitive types that msgpack can serialize directly without PickleWrapper. +# Do not include str here: msgspec rejects a Union containing both str and bytes +# as multiple str-like arms. Top-level strings use PickleWrapper; string fields +# inside typed structs are still decoded by their struct schemas. +_primitive_types = (int, float, bool, bytes) +_all_types = _struct_types + _primitive_types + +_msgpack_encoder = msgspec.msgpack.Encoder(enc_hook=enc_hook) +_msgpack_decoder = msgspec.msgpack.Decoder(Union[_all_types], dec_hook=dec_hook) +_USE_PICKLE_IPC = envs.SGLANG_USE_PICKLE_IPC.get() + + +def hook_custom_types(*new_types: Type): + global _msgpack_decoder, _all_types + _all_types = tuple(dict.fromkeys(_all_types + new_types)) + _msgpack_decoder = msgspec.msgpack.Decoder(Union[_all_types], dec_hook=dec_hook) + + +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." + ) + + +def _maybe_unwrap_pickle(obj: Any) -> Any: + if isinstance(obj, PickleWrapper): + obj = pickle.loads(obj.data) + if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): + logger.info(f"Object of type {type(obj)} is unwrapped from PickleWrapper.") + return obj + + return obj + + +def msgpack_encode(obj: Any) -> bytes: + return _msgpack_encoder.encode(_maybe_wrap_pickle(obj)) + + +def msgpack_decode(data: bytes) -> Any: + return _maybe_unwrap_pickle(_msgpack_decoder.decode(data)) + + +def sock_send(socket: zmq.Socket, obj: Any, flags: int = 0) -> None: + if _USE_PICKLE_IPC: + socket.send_pyobj(obj, flags=flags, protocol=pickle.HIGHEST_PROTOCOL) + return + + socket.send(msgpack_encode(obj), flags=flags) + + +def sock_recv(socket: zmq.Socket, flags: int = 0) -> Any: + if _USE_PICKLE_IPC: + return socket.recv_pyobj(flags=flags) + + data = socket.recv(flags=flags) + return msgpack_decode(data) + + +async def async_sock_send(socket: zmq.asyncio.Socket, obj: Any, flags: int = 0) -> None: + if _USE_PICKLE_IPC: + await socket.send_pyobj(obj, flags=flags, protocol=pickle.HIGHEST_PROTOCOL) + return + + await socket.send(msgpack_encode(obj), flags=flags) + + +async def async_sock_recv(socket: zmq.asyncio.Socket, flags: int = 0) -> Any: + if _USE_PICKLE_IPC: + return await socket.recv_pyobj(flags=flags) + + data = await socket.recv(flags=flags) + return msgpack_decode(data) diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index e27214468..94f589517 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -53,6 +53,8 @@ from sglang.srt.managers.io_struct import ( async_sock_send, sock_recv, sock_send, + unwrap_from_pickle, + wrap_as_pickle, ) from sglang.srt.managers.load_snapshot import ( create_load_snapshot_reader, @@ -122,17 +124,29 @@ def _extract_field_by_index( if field is None: return None + should_wrap_result = field_name in ("customized_info", "time_stats") + if should_wrap_result: + field = unwrap_from_pickle(field) + if field is None: + return None + if isinstance(field, dict): new_field = {} for k, v in field.items(): - new_field[k] = v[index] if len(v) > index else None + if len(v) > index: + new_field[k] = [v[index]] if should_wrap_result else v[index] + else: + new_field[k] = [None] if should_wrap_result else None + if should_wrap_result: + return wrap_as_pickle(new_field) if new_field else None return new_field if check_length: if len(field) <= index: return None - return [field[index]] + new_field = [field[index]] + return wrap_as_pickle(new_field) if should_wrap_result else new_field def _handle_output_by_index(output, i): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 5f3c85b65..21c007ea3 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -266,6 +266,7 @@ from sglang.srt.utils.hf_transformers_utils import ( get_tokenizer, get_tokenizer_from_processor, ) +from sglang.srt.utils.msgspec_utils import msgspec_to_builtins from sglang.srt.utils.numa_utils import get_numa_node_if_available, numa_bind_to_node from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method from sglang.srt.utils.tensor_bridge import use_mlx @@ -3752,7 +3753,7 @@ class Scheduler( # This field is not serializable. ret.pop("model_config", None) - return GetInternalStateReqOutput(internal_state=ret) + return GetInternalStateReqOutput(internal_state=msgspec_to_builtins(ret)) def set_internal_state(self, recv_req: SetInternalStateReq): server_args_dict = recv_req.server_args @@ -3801,7 +3802,7 @@ class Scheduler( server_args.pop("model_config", None) return SetInternalStateReqOutput( updated=if_success, - server_args=server_args, + server_args=msgspec_to_builtins(server_args), ) def save_remote_model(self, **kwargs): diff --git a/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py b/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py index e41538eb2..9cd95e565 100644 --- a/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py +++ b/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py @@ -1,6 +1,5 @@ from __future__ import annotations -import dataclasses import time from dataclasses import dataclass from typing import ( @@ -10,13 +9,14 @@ from typing import ( Optional, ) +import msgspec import zmq from sglang.srt.disaggregation.kv_events import ( EventPublisherFactory, KVEventBatch, ) -from sglang.srt.managers.io_struct import sock_send +from sglang.srt.managers.io_struct import hook_custom_types, sock_send if TYPE_CHECKING: from sglang.srt.distributed.parallel_state_wrapper import ParallelState @@ -26,8 +26,7 @@ if TYPE_CHECKING: class SchedulerStats: ... # type: ignore[no-redef] -@dataclasses.dataclass -class KvMetrics: +class KvMetrics(msgspec.Struct, tag=True, kw_only=True, array_like=True): request_active_slots: int = 0 request_total_slots: int = 0 kv_active_blocks: int = 0 @@ -38,6 +37,9 @@ class KvMetrics: data_parallel_rank: int = 0 +hook_custom_types(KvMetrics) + + @dataclass(kw_only=True, slots=True) class SchedulerKvEventsPublisher: kv_events_config: Optional[str] diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 1e3d33bce..1be8cc088 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -19,6 +19,7 @@ from sglang.srt.managers.io_struct import ( BatchEmbeddingOutput, BatchTokenIDOutput, CachedTokensDetails, + wrap_as_pickle, ) from sglang.srt.managers.schedule_batch import ( BaseFinishReason, @@ -229,7 +230,7 @@ class SchedulerOutputStreamer: BatchEmbeddingOutput( rids=rids, http_worker_ipcs=http_worker_ipcs, - time_stats=time_stats, + time_stats=wrap_as_pickle(time_stats), finished_reasons=finished_reasons, embeddings=embeddings, prompt_tokens=prompt_tokens, @@ -516,7 +517,7 @@ class _GenerationStreamAccumulator: spec_verify_ct=self.spec_verify_ct, spec_num_correct_drafts=self.spec_num_correct_drafts, spec_correct_drafts_histogram=self.spec_correct_drafts_histogram, - time_stats=self.time_stats, + time_stats=wrap_as_pickle(self.time_stats), finished_reasons=self.finished_reasons, decoded_texts=self.decoded_texts, decode_ids=self.decode_ids_list, @@ -549,7 +550,9 @@ class _GenerationStreamAccumulator: output_hidden_states=self.output_hidden_states, routed_experts=self.routed_experts, indexer_topk=self.indexer_topk, - customized_info=self.customized_info, + customized_info=( + wrap_as_pickle(self.customized_info) if self.customized_info else None + ), placeholder_tokens_idx=None, placeholder_tokens_val=None, retraction_counts=self.retraction_counts, diff --git a/python/sglang/srt/managers/scheduler_components/request_receiver.py b/python/sglang/srt/managers/scheduler_components/request_receiver.py index 60ad1b4b4..3e7ff21e1 100644 --- a/python/sglang/srt/managers/scheduler_components/request_receiver.py +++ b/python/sglang/srt/managers/scheduler_components/request_receiver.py @@ -89,6 +89,9 @@ class SchedulerRequestReceiver: recv_reqs = self._broadcast_reqs_across_ranks(recv_reqs) + if self.ps.pp_rank == 0: + self.unwrap_pickle_wrapper(recv_reqs) + recv_reqs = self._apply_mm_receiver(recv_reqs) self._finalize_shm_features(recv_reqs) @@ -195,6 +198,19 @@ class SchedulerRequestReceiver: ) return recv_reqs + def unwrap_pickle_wrapper(self, recv_reqs: Optional[List]) -> None: + if not recv_reqs: + return + + for req in recv_reqs: + if isinstance(req, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput)): + req.unwrap_pickle_fields() + elif isinstance( + req, (BatchTokenizedGenerateReqInput, BatchTokenizedEmbeddingReqInput) + ): + for sub_req in req: + sub_req.unwrap_pickle_fields() + def _apply_mm_receiver(self, recv_reqs: List) -> List: # Process MM requests under EPD-disaggregation mode if ( diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 9dd16b2a3..dae1355a9 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -80,6 +80,7 @@ from sglang.srt.managers.io_struct import ( async_sock_recv, async_sock_send, sock_send, + unwrap_from_pickle, ) from sglang.srt.managers.load_snapshot import create_load_snapshot_reader from sglang.srt.managers.mm_utils import TensorTransportMode, wrap_shm_features @@ -1332,7 +1333,10 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): ): tokenized_obj.time_stats.set_api_server_dispatch_time() tokenized_obj = wrap_shm_features(tokenized_obj) + time_stats = tokenized_obj.time_stats + tokenized_obj.wrap_pickle_fields() self._dispatch_to_scheduler(tokenized_obj) + tokenized_obj.time_stats = time_stats tokenized_obj.time_stats.set_api_server_dispatch_finish_time() def _send_batch_request( @@ -1342,13 +1346,19 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): ], ): """Send a batch of tokenized requests as a single batched request to the scheduler.""" + set_time_batch(tokenized_objs, "set_api_server_dispatch_time") + time_stats = [tokenized_obj.time_stats for tokenized_obj in tokenized_objs] + for tokenized_obj in tokenized_objs: + tokenized_obj.wrap_pickle_fields() + if isinstance(tokenized_objs[0], TokenizedGenerateReqInput): batch_req = BatchTokenizedGenerateReqInput(batch=tokenized_objs) else: batch_req = BatchTokenizedEmbeddingReqInput(batch=tokenized_objs) - set_time_batch(tokenized_objs, "set_api_server_dispatch_time") self._dispatch_to_scheduler(batch_req) + for tokenized_obj, time_stat in zip(tokenized_objs, time_stats): + tokenized_obj.time_stats = time_stat set_time_batch(tokenized_objs, "set_api_server_dispatch_finish_time") def _coalesce_streaming_chunks( @@ -1856,6 +1866,11 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): BatchTokenIDOutput, ], ): + recv_obj.time_stats = unwrap_from_pickle(recv_obj.time_stats) + if isinstance(recv_obj, (BatchStrOutput, BatchTokenIDOutput)): + customized_info = unwrap_from_pickle(recv_obj.customized_info) + else: + customized_info = None pending_notify: dict[str, ReqState] = {} batch_notify_size = self.server_args.batch_notify_size for i, rid in enumerate(recv_obj.rids): @@ -1911,8 +1926,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): meta_info["cached_tokens_details"] = recv_obj.cached_tokens_details[ i ] - if recv_obj.customized_info is not None: - for k, v in recv_obj.customized_info.items(): + if customized_info is not None: + for k, v in customized_info.items(): if k not in state.customized_info_accumulated: state.customized_info_accumulated[k] = [] state.customized_info_accumulated[k].extend(v[i]) diff --git a/python/sglang/srt/sampling/sampling_params.py b/python/sglang/srt/sampling/sampling_params.py index 7f9c8f38c..ce4a7c2d8 100644 --- a/python/sglang/srt/sampling/sampling_params.py +++ b/python/sglang/srt/sampling/sampling_params.py @@ -15,7 +15,7 @@ import logging import math -from typing import Any, Dict, List, Optional, Set, Union +from typing import Dict, List, Optional, Set, Union import msgspec @@ -25,6 +25,17 @@ try: except ImportError: import sre_parse # Python < 3.11 +# JSON-safe value types for custom_params. Must survive msgpack IPC +# without PickleWrapper. After deserialization on the scheduler side, +# Req.__init__ injects "__req__" (a Req object) into the dict in-process; +# that augmented dict is never re-serialized. +_JsonScalar = Union[None, bool, int, float, str] +CustomParamValue = Union[ + _JsonScalar, + List[_JsonScalar], + Dict[str, _JsonScalar], +] + _SAMPLING_EPS = 1e-6 TOP_K_ALL = 1 << 30 @@ -96,7 +107,7 @@ class SamplingParams(msgspec.Struct, kw_only=True, omit_defaults=True): skip_special_tokens: bool = True spaces_between_special_tokens: bool = True no_stop_trim: bool = False - custom_params: Optional[Dict[str, Any]] = None + custom_params: Optional[Dict[str, CustomParamValue]] = None stream_interval: Optional[int] = None logit_bias: Optional[Dict[str, float]] = None sampling_seed: Optional[int] = None diff --git a/python/sglang/srt/utils/msgspec_utils.py b/python/sglang/srt/utils/msgspec_utils.py new file mode 100644 index 000000000..60955b1cc --- /dev/null +++ b/python/sglang/srt/utils/msgspec_utils.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import base64 +import binascii +from typing import Any + +import msgspec +from pydantic_core import core_schema + + +class Base64Bytes: + """Pydantic marker for HTTP JSON base64-encoded bytes fields.""" + + def __get_pydantic_core_schema__(self, source_type: Any, handler): + return core_schema.no_info_before_validator_function( + self._decode_value, + handler(source_type), + ) + + @classmethod + def _decode_value(cls, value: Any) -> Any: + if isinstance(value, str): + try: + return base64.b64decode(value, validate=True) + except binascii.Error as exc: + raise ValueError("Expected base64-encoded bytes") from exc + + if isinstance(value, list): + return [cls._decode_value(item) for item in value] + + if isinstance(value, tuple): + return tuple(cls._decode_value(item) for item in value) + + return value + + +def msgspec_to_builtins(obj: Any) -> Any: + """Recursively convert msgspec structs to dict/list Python 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 isinstance(obj, dict): + return {key: msgspec_to_builtins(value) for key, value in obj.items()} + + if isinstance(obj, list): + return [msgspec_to_builtins(item) for item in obj] + + if isinstance(obj, tuple): + return tuple(msgspec_to_builtins(item) for item in obj) + + if isinstance(obj, set): + return [msgspec_to_builtins(item) for item in obj] + + return obj + + +def msgspec_struct_pydantic_core_schema(cls: type[msgspec.Struct], handler): + fields = {} + for struct_field in msgspec.structs.fields(cls): + field_schema = handler.generate_schema(struct_field.type) + required = ( + struct_field.default is msgspec.NODEFAULT + and struct_field.default_factory is msgspec.NODEFAULT + ) + + if struct_field.default is not msgspec.NODEFAULT: + field_schema = core_schema.with_default_schema( + field_schema, + default=struct_field.default, + ) + elif struct_field.default_factory is not msgspec.NODEFAULT: + field_schema = core_schema.with_default_schema( + field_schema, + default_factory=struct_field.default_factory, + ) + + fields[struct_field.name] = core_schema.typed_dict_field( + field_schema, + required=required, + ) + + typed_dict_schema = core_schema.typed_dict_schema( + fields, + cls_name=cls.__name__, + extra_behavior="ignore", + ref=cls.__name__, + ) + + def build_struct(value): + return value if isinstance(value, cls) else cls(**value) + + dict_to_struct_schema = core_schema.no_info_after_validator_function( + build_struct, + typed_dict_schema, + ) + return core_schema.json_or_python_schema( + json_schema=dict_to_struct_schema, + python_schema=core_schema.union_schema( + [ + core_schema.is_instance_schema(cls), + dict_to_struct_schema, + ], + mode="left_to_right", + ), + ) diff --git a/python/sglang/test/scripted_runtime/http_server.py b/python/sglang/test/scripted_runtime/http_server.py index 476cd1f0f..6d86a8652 100644 --- a/python/sglang/test/scripted_runtime/http_server.py +++ b/python/sglang/test/scripted_runtime/http_server.py @@ -12,7 +12,7 @@ import zmq from sglang.srt.entrypoints.http_server import launch_server from sglang.srt.environ import envs -from sglang.srt.managers.io_struct import sock_recv, sock_send +from sglang.srt.managers.io_struct import sock_recv, sock_send, wrap_as_pickle from sglang.srt.server_args import ServerArgs from sglang.srt.utils.network import get_free_port, get_zmq_socket_on_host from sglang.test.scripted_runtime.io_struct import ( @@ -90,7 +90,7 @@ class ScriptedHttpServer: raise RuntimeError(f"ScriptedHttpServer is dirty: {self._dirty}") fn_path = f"{script_fn.__module__}:{script_fn.__qualname__}" - sock_send(self._socket, RunScript(fn_path=fn_path, args=args)) + sock_send(self._socket, wrap_as_pickle(RunScript(fn_path=fn_path, args=args))) if not self._socket.poll(int(timeout_s * 1000)): if not self._server_process.is_alive(): @@ -117,7 +117,7 @@ class ScriptedHttpServer: fatal_error: Optional[OutOfBandError] = None try: try: - sock_send(self._socket, Shutdown()) + sock_send(self._socket, wrap_as_pickle(Shutdown())) except zmq.ZMQError: pass diff --git a/python/sglang/test/scripted_runtime/scheduler_hook.py b/python/sglang/test/scripted_runtime/scheduler_hook.py index 34cd9d783..58253cfcb 100644 --- a/python/sglang/test/scripted_runtime/scheduler_hook.py +++ b/python/sglang/test/scripted_runtime/scheduler_hook.py @@ -11,7 +11,7 @@ from typing import TYPE_CHECKING, Generator, List, Optional, Tuple import zmq from sglang.srt.environ import envs -from sglang.srt.managers.io_struct import sock_recv, sock_send +from sglang.srt.managers.io_struct import sock_recv, sock_send, wrap_as_pickle from sglang.srt.utils.network import get_zmq_socket from sglang.test.scripted_runtime.background_http_poster import BackgroundHttpPoster from sglang.test.scripted_runtime.context import ScriptedContext @@ -153,7 +153,7 @@ class ScriptedSchedulerHook: socket = get_zmq_socket(ctx_zmq, zmq.PAIR, endpoint, bind=False) try: yield from _drive_engine_through_warmup(self._context) - sock_send(socket, HookReady()) + sock_send(socket, wrap_as_pickle(HookReady())) while True: msg = sock_recv(socket) match msg: @@ -170,10 +170,12 @@ class ScriptedSchedulerHook: except Exception: sock_send( socket, - ScriptFailed(traceback=traceback.format_exc()), + wrap_as_pickle( + ScriptFailed(traceback=traceback.format_exc()) + ), ) else: - sock_send(socket, ScriptSucceeded()) + sock_send(socket, wrap_as_pickle(ScriptSucceeded())) case _: raise ValueError(f"dispatch loop: unknown command {msg!r}") finally: diff --git a/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py b/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py index 0875a0f70..aee5795dd 100644 --- a/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py +++ b/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py @@ -41,6 +41,11 @@ class ScriptedTokenizerRecvProxy: "ScriptedTokenizerRecvProxy.recv_pyobj: blocking recv is not supported" ) + def recv(self, flags: int = 0) -> bytes: + raise NotImplementedError( + "TODO: support ScriptedTokenizerRecvProxy.recv for msgpack IPC" + ) + def wait_until_arrived( self, predicate: Callable[[Any], bool], diff --git a/test/registered/core/test_srt_endpoint.py b/test/registered/core/test_srt_endpoint.py index 646530046..81010b6c5 100644 --- a/test/registered/core/test_srt_endpoint.py +++ b/test/registered/core/test_srt_endpoint.py @@ -31,6 +31,8 @@ from sglang.test.test_utils import ( register_cuda_ci(est_time=134, stage="base-b", runner_config="1-gpu-small") register_amd_ci(est_time=130, suite="stage-b-test-1-gpu-small-amd") +SERVER_ENV = {"SGLANG_USE_PICKLE_IPC": "0"} + class TestSRTEndpoint(CustomTestCase): @classmethod @@ -41,6 +43,7 @@ class TestSRTEndpoint(CustomTestCase): cls.model, cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + env=SERVER_ENV, other_args=( "--enable-custom-logit-processor", "--mem-fraction-static", @@ -469,6 +472,12 @@ class TestSRTEndpoint(CustomTestCase): response = requests.post(self.base_url + "/flush_cache") assert response.status_code == 200 + server_info = requests.get(self.base_url + "/server_info").json() + page_size = server_info.get("page_size") or 1 + + def align_down(num_tokens): + return num_tokens // page_size * page_size + def send_and_check_cached_tokens(input_ids): response = requests.post( self.base_url + "/generate", @@ -483,10 +492,14 @@ class TestSRTEndpoint(CustomTestCase): return response_json["meta_info"]["cached_tokens"] self.assertEqual(send_and_check_cached_tokens(range(0, 100)), 0) - self.assertEqual(send_and_check_cached_tokens(range(0, 10000)), 100) - self.assertEqual(send_and_check_cached_tokens(range(0, 10000)), 9999) - self.assertEqual(send_and_check_cached_tokens(range(0, 1000)), 999) - self.assertEqual(send_and_check_cached_tokens(range(0, 11000)), 10000) + self.assertEqual(send_and_check_cached_tokens(range(0, 10000)), align_down(100)) + self.assertEqual( + send_and_check_cached_tokens(range(0, 10000)), align_down(9999) + ) + self.assertEqual(send_and_check_cached_tokens(range(0, 1000)), align_down(999)) + self.assertEqual( + send_and_check_cached_tokens(range(0, 11000)), align_down(10000) + ) def test_get_server_info(self): response = requests.get(self.base_url + "/server_info") @@ -648,6 +661,7 @@ class TestTokenizeDetokenize(CustomTestCase): cls.model, cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + env=SERVER_ENV, ) cls.tokenizer = get_tokenizer(cls.model) diff --git a/test/registered/unit/entrypoints/test_server_info.py b/test/registered/unit/entrypoints/test_server_info.py index acd3443e3..3f348dc78 100644 --- a/test/registered/unit/entrypoints/test_server_info.py +++ b/test/registered/unit/entrypoints/test_server_info.py @@ -22,10 +22,12 @@ Current coverage: import asyncio import dataclasses +import json import unittest from types import SimpleNamespace from sglang.srt.entrypoints import http_server +from sglang.srt.lora.lora_registry import LoRARef from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -33,7 +35,9 @@ from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=5, suite="base-a-test-cpu") -def _call_server_info_with(server_args: ServerArgs) -> dict: +def _call_server_info_with( + server_args: ServerArgs, internal_states: list[dict] | None = None +) -> dict: """Invoke `http_server.server_info()` against a stub global state. Bypasses the FastAPI HTTP layer (no TestClient): the handler is an @@ -44,7 +48,7 @@ def _call_server_info_with(server_args: ServerArgs) -> dict: """ async def _fake_internal_state(): - return [{"max_req_input_len": 1024}] + return internal_states or [{"max_req_input_len": 1024}] stub_state = SimpleNamespace( tokenizer_manager=SimpleNamespace( @@ -291,6 +295,31 @@ class TestServerInfoExistingFieldsPreserved(CustomTestCase): self.assertIn("kv_events", info) self.assertIsNotNone(info["kv_events"]) + def test_lora_refs_are_json_serializable_dicts(self): + lora_ref = LoRARef( + lora_id="lora-id", + lora_name="adapter", + lora_path="/tmp/adapter", + pinned=True, + ) + args = ServerArgs(model_path="dummy") + args.lora_paths = [lora_ref] + + info = _call_server_info_with( + args, + internal_states=[{"lora_paths": [lora_ref]}], + ) + + expected = { + "lora_id": "lora-id", + "lora_name": "adapter", + "lora_path": "/tmp/adapter", + "pinned": True, + } + self.assertEqual(info["lora_paths"], [expected]) + self.assertEqual(info["internal_states"][0]["lora_paths"], [expected]) + json.dumps(info) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py index bfbd6ec41..55249b948 100644 --- a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py +++ b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py @@ -13,10 +13,11 @@ Covers: """ import asyncio -import dataclasses import unittest from unittest.mock import AsyncMock, MagicMock, Mock +import msgspec + from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel @@ -35,7 +36,7 @@ _NOT_FINISHED = object() # Sentinel: request has not finished yet # Categorised by value shape so that _make_batch_str_output can assign # type-appropriate defaults without hardcoding every field name. # When a field is renamed upstream, the old name simply won't appear in -# dataclasses.fields() and the new name will fall through to the +# msgspec.structs.fields() and the new name will fall through to the # pattern-matching or safe fallback — no test breakage. # --------------------------------------------------------------------------- @@ -145,7 +146,7 @@ def _make_abort_req(rid: str, abort_message: str = "Aborted") -> AbortReq: def _make_batch_str_output(rid: str, finished_reason=None) -> BatchStrOutput: """Create a minimal BatchStrOutput for a single request. - Uses dataclass field introspection so that new or renamed fields in + Uses struct field introspection so that new or renamed fields in BatchStrOutput don't break this test. Only the fields that matter for test logic (rids, finished_reasons, output_strs) are set explicitly; all others receive type-appropriate defaults based on naming patterns. @@ -159,7 +160,7 @@ def _make_batch_str_output(rid: str, finished_reason=None) -> BatchStrOutput: fr = finished_reason kwargs = {} - for f in dataclasses.fields(BatchStrOutput): + for f in msgspec.structs.fields(BatchStrOutput): if f.name == "rids": kwargs[f.name] = [rid] elif f.name == "finished_reasons": @@ -176,8 +177,8 @@ def _make_batch_str_output(rid: str, finished_reason=None) -> BatchStrOutput: kwargs[f.name] = [None] # Fields with class defaults — skip, let the default be used elif ( - f.default is not dataclasses.MISSING - or f.default_factory is not dataclasses.MISSING + f.default is not msgspec.NODEFAULT + or f.default_factory is not msgspec.NODEFAULT ): continue # Unknown required field — provide a safe per-request default.