From 65eebd89cb7c5eb2c06bbe2ccd194b3439612aa9 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Wed, 3 Jun 2026 16:36:58 -0700 Subject: [PATCH 01/45] pickle to msgpack migration --- python/sglang/srt/debug_utils/dumper.py | 20 +- .../srt/disaggregation/encode_grpc_server.py | 3 +- .../srt/disaggregation/encode_server.py | 27 +- .../srt/elastic_ep/expert_backup_client.py | 6 +- .../srt/elastic_ep/expert_backup_manager.py | 6 +- python/sglang/srt/entrypoints/engine.py | 6 +- python/sglang/srt/entrypoints/grpc_server.py | 4 +- python/sglang/srt/entrypoints/http_server.py | 404 ++++++++++-- python/sglang/srt/environ.py | 1 + python/sglang/srt/managers/communicator.py | 5 +- .../srt/managers/data_parallel_controller.py | 26 +- .../srt/managers/detokenizer_manager.py | 6 +- python/sglang/srt/managers/embed_types.py | 3 +- python/sglang/srt/managers/io_struct.py | 609 +++++++++--------- .../srt/managers/multi_tokenizer_mixin.py | 76 ++- python/sglang/srt/managers/scheduler.py | 3 +- .../kv_events_publisher.py | 3 +- .../scheduler_components/output_sender.py | 14 +- .../scheduler_components/profiler_manager.py | 2 +- .../scheduler_components/request_receiver.py | 5 +- .../srt/managers/scheduler_input_blocker.py | 8 +- .../srt/managers/tokenizer_control_mixin.py | 8 +- .../sglang/srt/managers/tokenizer_manager.py | 59 +- .../srt/observability/req_time_stats.py | 10 +- python/sglang/srt/observability/trace.py | 11 +- .../unit/utils/test_profile_merger.py | 4 +- 26 files changed, 850 insertions(+), 479 deletions(-) diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 6c7d3f6e570a..edb86aea4b2f 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -1391,6 +1391,7 @@ def _create_zmq_rpc_broadcast( ) -> Optional["_ZmqRpcBroadcast"]: """A general-purpose minimal RPC to support broadcasting executions to multi processes""" import zmq + from sglang.srt.managers.io_struct import (sock_send, sock_recv) rank = _get_rank() world_size = dist.get_world_size() if dist.is_initialized() else 1 @@ -1404,13 +1405,13 @@ def _create_zmq_rpc_broadcast( def serve_loop(): while True: try: - req = sock.recv_pyobj() + req = sock_recv(sock) result = getattr(handler, req["method"])(*req["args"], **req["kwargs"]) resp = {"result": result, "error": None} except Exception as e: _log(f"[ZmqRpc] error inside handler: {e}") resp = {"result": None, "error": str(e)} - sock.send_pyobj(resp) + sock_send(sock, resp) thread = threading.Thread(target=serve_loop, daemon=True) thread.start() @@ -1447,14 +1448,13 @@ def __init__(self, socket, debug_name: str): def __getattr__(self, method_name: str): def call(*args, **kwargs): - self._socket.send_pyobj( - { - "method": method_name, - "args": args, - "kwargs": kwargs, - } - ) - response = self._socket.recv_pyobj() + from sglang.srt.managers.io_struct import (sock_send, sock_recv) + sock_send(self._socket, { + "method": method_name, + "args": args, + "kwargs": kwargs, + }) + response = sock_recv(self._socket) if response["error"]: raise RuntimeError( f"RPC error on {self._debug_name}: {response['error']}" diff --git a/python/sglang/srt/disaggregation/encode_grpc_server.py b/python/sglang/srt/disaggregation/encode_grpc_server.py index 033520093b58..0207fbf6e06c 100644 --- a/python/sglang/srt/disaggregation/encode_grpc_server.py +++ b/python/sglang/srt/disaggregation/encode_grpc_server.py @@ -87,6 +87,7 @@ def __init__( async def Encode( self, request: sglang_encoder_pb2.EncodeRequest, context ) -> sglang_encoder_pb2.EncodeResponse: + from sglang.srt.managers.io_struct import async_sock_send try: request_dict = { "mm_items": list(request.mm_items), @@ -95,7 +96,7 @@ async def Encode( "part_idx": request.part_idx, } for socket in self.send_sockets: - await socket.send_pyobj(request_dict) + await async_sock_send(socket, 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 265f5c6fb6e1..954b9d22ad36 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -40,7 +40,13 @@ ) from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import initialize_dp_attention -from sglang.srt.managers.io_struct import ProfileReq, ProfileReqInput, ProfileReqType +from sglang.srt.managers.io_struct import ( + ProfileReq, + ProfileReqInput, + ProfileReqType, + sock_send, + async_sock_recv, +) from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache from sglang.srt.model_loader import get_model @@ -2303,7 +2309,8 @@ async def _dispatch_group( requests = [p.request for p in group] start = time.time() for sock in self.send_sockets: - sock.send_pyobj( + sock_send( + sock, { "type": "batch_encode", "modality": modality.name, @@ -2355,7 +2362,7 @@ async def _dispatch_per_request( req = p.request try: for sock in self.send_sockets: - sock.send_pyobj(req) + sock_send(sock, req) result = await self.encoder.encode_request(req, modality) if not p.future.done(): p.future.set_result(result) @@ -2395,9 +2402,9 @@ async def run_encoder( ): encoder = MMEncoder(server_args, schedule_path, dist_init_method, rank) while True: - request = await encoder.schedule_socket.recv_pyobj() + request = await async_sock_recv(encoder.schedule_socket) if isinstance(request, ProfileReq): - if request.type == ProfileReqType.START_PROFILE: + if request.req_type == ProfileReqType.START_PROFILE: if encoder.profiler is None: encoder.profiler = EncoderProfiler(encoder.rank) encoder.profiler.start(request) @@ -2527,7 +2534,7 @@ def start_background_send(req_id): ) else: for socket in send_sockets: - socket.send_pyobj(request) + sock_send(socket, request) nbytes, embedding_len, embedding_dim, error_msg, error_code = ( await encoder.encode_request(request, modality) ) @@ -2701,7 +2708,7 @@ async def health_generate(): # Broadcast to other TP ranks so distributed ops stay in sync for socket in send_sockets: - socket.send_pyobj(dummy_request) + sock_send(socket, dummy_request) # Run encode on rank 0 with timeout _, _, _, error_msg, _ = await asyncio.wait_for( @@ -2741,7 +2748,7 @@ async def start_profile_async(obj: Optional[ProfileReqInput] = None): req = ProfileReq(ProfileReqType.START_PROFILE) else: req = ProfileReq( - type=ProfileReqType.START_PROFILE, + req_type=ProfileReqType.START_PROFILE, output_dir=obj.output_dir, start_step=obj.start_step, num_steps=obj.num_steps, @@ -2755,7 +2762,7 @@ async def start_profile_async(obj: Optional[ProfileReqInput] = None): profile_stages=obj.profile_stages, ) for socket in send_sockets: - socket.send_pyobj(req) + sock_send(socket, req) if encoder.profiler is None: encoder.profiler = EncoderProfiler(encoder.rank) ok, msg = encoder.profiler.start(req) @@ -2780,7 +2787,7 @@ async def stop_profile_async(): ) req = ProfileReq(ProfileReqType.STOP_PROFILE) for socket in send_sockets: - socket.send_pyobj(req) + sock_send(socket, req) ok, msg = encoder.profiler.stop() if ok: return Response(content="Stop profiling.\n", status_code=200) diff --git a/python/sglang/srt/elastic_ep/expert_backup_client.py b/python/sglang/srt/elastic_ep/expert_backup_client.py index 1f44813484fd..6b114a699fe4 100644 --- a/python/sglang/srt/elastic_ep/expert_backup_client.py +++ b/python/sglang/srt/elastic_ep/expert_backup_client.py @@ -12,7 +12,7 @@ ) from sglang.srt.environ import envs from sglang.srt.eplb.expert_location import get_global_expert_location_metadata -from sglang.srt.managers.io_struct import UpdateExpertBackupReq +from sglang.srt.managers.io_struct import UpdateExpertBackupReq, sock_send, sock_recv from sglang.srt.server_args import ServerArgs from sglang.srt.utils.network import get_local_ip_auto @@ -65,7 +65,7 @@ def __init__(self, server_args: ServerArgs, model_runner): self.ready_sockets[i].connect( f"tcp://{all_ips[i * get_world_size() // server_args.nnodes]}:{PORT_BASE + i * 2}" ) - self.ready_sockets[i].send_pyobj(UpdateExpertBackupReq()) + sock_send(self.ready_sockets[i], UpdateExpertBackupReq()) self._receive_thread = threading.Thread(target=self._receive_loop, daemon=True) self._receive_thread.start() @@ -73,7 +73,7 @@ def __init__(self, server_args: ServerArgs, model_runner): def _receive_loop(self): cnt = 0 while cnt < self.engine_num: - response = self.recv_list[cnt].recv_pyobj() + response = sock_recv(self.recv_list[cnt]) self.dram_map_list[response.rank] = response.weight_pointer_map self.session_id_list[response.rank] = response.session_id self.buffer_size = max(self.buffer_size, response.buffer_size) diff --git a/python/sglang/srt/elastic_ep/expert_backup_manager.py b/python/sglang/srt/elastic_ep/expert_backup_manager.py index 60afe0be14e7..c864dee8c2bd 100644 --- a/python/sglang/srt/elastic_ep/expert_backup_manager.py +++ b/python/sglang/srt/elastic_ep/expert_backup_manager.py @@ -9,7 +9,7 @@ from sglang.srt.configs.load_config import LoadConfig from sglang.srt.configs.model_config import ModelConfig from sglang.srt.environ import envs -from sglang.srt.managers.io_struct import BackupDramReq +from sglang.srt.managers.io_struct import BackupDramReq, sock_recv, sock_send from sglang.srt.model_loader.loader import DefaultModelLoader, get_model_loader from sglang.srt.model_loader.utils import set_default_torch_dtype from sglang.srt.server_args import ( @@ -62,7 +62,7 @@ def __init__(self, server_args: ServerArgs, port_args: PortArgs): num_ready_clients = 0 while num_ready_clients < server_args.tp_size: - self.recv_from_expert_backup_client.recv_pyobj() + sock_recv(self.recv_from_expert_backup_client) num_ready_clients += 1 back_req = BackupDramReq( @@ -72,7 +72,7 @@ def __init__(self, server_args: ServerArgs, port_args: PortArgs): buffer_size=self.continuous_buffer.numel() * self.continuous_buffer.element_size(), ) - self.send_to_expert_backup_client.send_pyobj(back_req) + sock_send(self.send_to_expert_backup_client, back_req) # Keep the manager subprocess alive until signals signal.pause() diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 3f58ca644c01..8c817c8d3edd 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -80,6 +80,8 @@ UpdateWeightsFromDistributedReqInput, UpdateWeightsFromIPCReqInput, UpdateWeightsFromTensorReqInput, + sock_recv, + sock_send, ) from sglang.srt.managers.multi_tokenizer_mixin import ( MultiTokenizerRouter, @@ -1215,8 +1217,8 @@ def freeze_gc(self): def collective_rpc(self, method: str, **kwargs): obj = RpcReqInput(method=method, parameters=kwargs) - self.send_to_rpc.send_pyobj(obj) - recv_req = self.send_to_rpc.recv_pyobj(zmq.BLOCKY) + sock_send(self.send_to_rpc, obj) + recv_req = sock_recv(self.send_to_rpc, flags=zmq.BLOCKY) assert isinstance(recv_req, RpcReqOutput) assert recv_req.success, recv_req.message diff --git a/python/sglang/srt/entrypoints/grpc_server.py b/python/sglang/srt/entrypoints/grpc_server.py index 00944c276b9e..964ed379fe79 100644 --- a/python/sglang/srt/entrypoints/grpc_server.py +++ b/python/sglang/srt/entrypoints/grpc_server.py @@ -105,7 +105,7 @@ async def start_profile_handler(request): record_shapes = (record_shapes is not False) and env_record_shapes req = ProfileReq( - type=ProfileReqType.START_PROFILE, + req_type=ProfileReqType.START_PROFILE, output_dir=body.get("output_dir"), start_step=body.get("start_step"), num_steps=body.get("num_steps"), @@ -134,7 +134,7 @@ async def start_profile_handler(request): async def stop_profile_handler(request): try: - req = ProfileReq(type=ProfileReqType.STOP_PROFILE) + req = ProfileReq(ProfileReqType.STOP_PROFILE) results = await request_manager.send_communicator_req( req, "profile_communicator", timeout=600.0 ) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index c51dbcd2f808..07f7eca2da89 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -35,6 +35,7 @@ List, Optional, Union, + Literal, ) # Fix a bug of Python threading @@ -109,36 +110,34 @@ from sglang.srt.environ import envs from sglang.srt.function_call.function_call_parser import FunctionCallParser from sglang.srt.managers.io_struct import ( - AbortReq, - AttachHiCacheStorageReqInput, - CheckWeightsReqInput, - CloseSessionReqInput, - ConfigureLoggingReq, - ContinueGenerationReqInput, - DestroyWeightsUpdateGroupReqInput, + CheckWeightsReqInput as CheckWeightsReqInputIpc, + CloseSessionReqInput as CloseSessionReqInputIpc, + ConfigureLoggingReq as ConfigureLoggingReqIpc, + ContinueGenerationReqInput as ContinueGenerationReqInputIpc, + DestroyWeightsUpdateGroupReqInput as DestroyWeightsUpdateGroupReqInputIpc, DumperControlReqInput, EmbeddingReqInput, GenerateReqInput, - GetWeightsByNameReqInput, - InitWeightsSendGroupForRemoteInstanceReqInput, - InitWeightsUpdateGroupReqInput, - LoadLoRAAdapterFromTensorsReqInput, - LoadLoRAAdapterReqInput, - OpenSessionReqInput, + GetWeightsByNameReqInput as GetWeightsByNameReqInputIpc, + InitWeightsSendGroupForRemoteInstanceReqInput as InitWeightsSendGroupForRemoteInstanceReqInputIpc, + InitWeightsUpdateGroupReqInput as InitWeightsUpdateGroupReqInputIpc, + LoadLoRAAdapterFromTensorsReqInput as LoadLoRAAdapterFromTensorsReqInputIpc, + LoadLoRAAdapterReqInput as LoadLoRAAdapterReqInputIpc, + OpenSessionReqInput as OpenSessionReqInputIpc, ParseFunctionCallReq, - PauseGenerationReqInput, + PauseGenerationReqInput as PauseGenerationReqInputIpc, ProfileReqInput, - ReleaseMemoryOccupationReqInput, - ResumeMemoryOccupationReqInput, - SendWeightsToRemoteInstanceReqInput, + ReleaseMemoryOccupationReqInput as ReleaseMemoryOccupationReqInputIpc, + ResumeMemoryOccupationReqInput as ResumeMemoryOccupationReqInputIpc, + SendWeightsToRemoteInstanceReqInput as SendWeightsToRemoteInstanceReqInputIpc, SeparateReasoningReqInput, - SetInternalStateReq, - SlowDownReqInput, - UnloadLoRAAdapterReqInput, - UpdateWeightFromDiskReqInput, - UpdateWeightsFromDistributedReqInput, - UpdateWeightsFromIPCReqInput, - UpdateWeightsFromTensorReqInput, + SetInternalStateReq as SetInternalStateReqIpc, + SlowDownReqInput as SlowDownReqInputIpc, + UnloadLoRAAdapterReqInput as UnloadLoRAAdapterReqInputIpc, + UpdateWeightFromDiskReqInput as UpdateWeightFromDiskReqInputIpc, + UpdateWeightsFromDistributedReqInput as UpdateWeightsFromDistributedReqInputIpc, + UpdateWeightsFromIPCReqInput as UpdateWeightsFromIPCReqInputIpc, + UpdateWeightsFromTensorReqInput as UpdateWeightsFromTensorReqInputIpc, UpdateWeightVersionReqInput, VertexGenerateReqInput, ) @@ -680,12 +679,19 @@ async def get_load(): ] +@dataclasses.dataclass +class SetInternalStateReq: + server_args: Dict[str, Any] + + # example usage: # curl -s -X POST http://localhost:30000/set_internal_state -H "Content-Type: application/json" -d '{"server_args": {"pp_max_micro_batch_size": 8}}' @app.api_route("/set_internal_state", methods=["POST", "PUT"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) async def set_internal_state(obj: SetInternalStateReq, request: Request): - res = await _global_state.tokenizer_manager.set_internal_state(obj) + res = await _global_state.tokenizer_manager.set_internal_state( + SetInternalStateReqIpc(**obj.__dict__) + ) return res @@ -870,6 +876,41 @@ async def clear_hicache_storage_backend(): ) +@dataclasses.dataclass +class AttachHiCacheStorageReqInput: + """Dynamically attach (enable) HiCache storage backend at runtime. + + Note: `hicache_storage_backend_extra_config_json` is a JSON string. It may contain both: + - backend-specific configs (e.g., mooncake master address) + - prefetch-related knobs (prefetch_threshold, prefetch_timeout_*, hicache_storage_pass_prefix_keys) + """ + + hicache_storage_backend: str + hicache_storage_backend_extra_config_json: Optional[str] = None + hicache_storage_prefetch_policy: Optional[str] = None + hicache_write_policy: Optional[str] = None + + def __post_init__(self): + if self.hicache_storage_prefetch_policy is None: + pass + else: + allowed = ["best_effort", "wait_complete", "timeout"] + if self.hicache_storage_prefetch_policy not in allowed: + raise ValueError( + f"Invalid hicache_storage_prefetch_policy: {self.hicache_storage_prefetch_policy!r}. " + f"Expected one of {allowed}." + ) + + if self.hicache_write_policy is None: + return + allowed = ["write_back", "write_through", "write_through_selective"] + if self.hicache_write_policy not in allowed: + raise ValueError( + f"Invalid hicache_write_policy: {self.hicache_write_policy!r}. " + f"Expected one of {allowed}." + ) + + # example usage: # curl -s -X PUT http://127.0.0.1:30000/hicache/storage-backend \ # -H 'Content-Type: application/json' \ @@ -1045,12 +1086,41 @@ async def dump_expert_distribution_record_async(): ) +@dataclasses.dataclass +class UpdateWeightFromDiskReqInput: + # The model path with the new weights + model_path: str + # The format to load the weights + load_format: Optional[str] = None + # Whether to abort all requests before updating weights + abort_all_requests: bool = False + # Optional: Update weight version along with weights + weight_version: Optional[str] = None + # Whether to update weights asynchronously + is_async: bool = False + # Whether to call torch.cuda.empty_cache() during flush + torch_empty_cache: bool = False + # Whether to keep the scheduler paused after weight update + keep_pause: bool = False + # Whether to recapture cuda graph after weight update + recapture_cuda_graph: bool = False + # The trainer step id. Used to know which step's weights are used for sampling. + token_step: int = 0 + # Whether to flush the cache after updating weights + flush_cache: bool = True + # Tensor metadata + manifest: Optional[Dict[str, Any]] = None + + @app.post("/update_weights_from_disk") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weights_from_disk(obj: UpdateWeightFromDiskReqInput, request: Request): """Update the weights from disk inplace without re-launching the server.""" success, message, num_paused_requests = ( - await _global_state.tokenizer_manager.update_weights_from_disk(obj, request) + await _global_state.tokenizer_manager.update_weights_from_disk( + UpdateWeightFromDiskReqInputIpc(**obj.__dict__), + request, + ) ) content = { @@ -1070,6 +1140,22 @@ async def update_weights_from_disk(obj: UpdateWeightFromDiskReqInput, request: R ) +@dataclasses.dataclass +class InitWeightsSendGroupForRemoteInstanceReqInput: + # The master address + master_address: str + # The ports for each rank's communication group + ports: str + # The rank in the communication group + group_rank: int + # The world size + world_size: int + # The group name + group_name: str = "weight_send_group" + # The backend + backend: str = "nccl" + + @app.post("/init_weights_send_group_for_remote_instance") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def init_weights_send_group_for_remote_instance( @@ -1077,7 +1163,8 @@ async def init_weights_send_group_for_remote_instance( ): success, message = ( await _global_state.tokenizer_manager.init_weights_send_group_for_remote_instance( - obj, request + InitWeightsSendGroupForRemoteInstanceReqInputIpc(**obj.__dict__), + request, ) ) content = {"success": success, "message": message} @@ -1087,6 +1174,16 @@ async def init_weights_send_group_for_remote_instance( return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) +@dataclasses.dataclass +class SendWeightsToRemoteInstanceReqInput: + # The master address + master_address: str + # The ports for each rank's communication group + ports: str + # The group name + group_name: str = "weight_send_group" + + @app.post("/send_weights_to_remote_instance") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def send_weights_to_remote_instance( @@ -1094,7 +1191,8 @@ async def send_weights_to_remote_instance( ): success, message = ( await _global_state.tokenizer_manager.send_weights_to_remote_instance( - obj, request + SendWeightsToRemoteInstanceReqInputIpc(**obj.__dict__), + request, ) ) content = {"success": success, "message": message} @@ -1142,6 +1240,22 @@ async def remote_instance_transfer_engine_info(rank: int = None): ) +@dataclasses.dataclass +class InitWeightsUpdateGroupReqInput: + # The master address + master_address: str + # The master port + master_port: int + # The rank offset + rank_offset: int + # The world size + world_size: int + # The group name + group_name: str = "weight_update_group" + # The backend + backend: str = "nccl" + + @app.post("/init_weights_update_group") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def init_weights_update_group( @@ -1149,7 +1263,8 @@ async def init_weights_update_group( ): """Initialize the parameter update group.""" success, message = await _global_state.tokenizer_manager.init_weights_update_group( - obj, request + InitWeightsUpdateGroupReqInputIpc(**obj.__dict__), + request, ) content = {"success": success, "message": message} if success: @@ -1158,6 +1273,11 @@ async def init_weights_update_group( return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) +@dataclasses.dataclass +class DestroyWeightsUpdateGroupReqInput: + group_name: str = "weight_update_group" + + @app.post("/destroy_weights_update_group") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def destroy_weights_update_group( @@ -1165,7 +1285,10 @@ async def destroy_weights_update_group( ): """Destroy the parameter update group.""" success, message = ( - await _global_state.tokenizer_manager.destroy_weights_update_group(obj, request) + await _global_state.tokenizer_manager.destroy_weights_update_group( + DestroyWeightsUpdateGroupReqInputIpc(**obj.__dict__), + request, + ) ) content = {"success": success, "message": message} return ORJSONResponse( @@ -1173,6 +1296,24 @@ async def destroy_weights_update_group( ) +@dataclasses.dataclass +class UpdateWeightsFromTensorReqInput: + # List[Union[str, bytes]] + serialized_named_tensors: List[bytes] + # Optional format specification for loading + load_format: Optional[str] = None + # Whether to flush the cache after updating weights + flush_cache: bool = True + # Whether to abort all requests before updating weights + abort_all_requests: bool = False + # Optional: Update weight version along with weights + weight_version: Optional[str] = None + # Optional: Determine whether to disable updating the draft model + disable_draft_model: Optional[bool] = None + # Whether to call torch.cuda.empty_cache() during flush + torch_empty_cache: bool = False + + @app.post("/update_weights_from_tensor") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weights_from_tensor( @@ -1186,7 +1327,8 @@ async def update_weights_from_tensor( """ success, message = await _global_state.tokenizer_manager.update_weights_from_tensor( - obj, request + UpdateWeightsFromTensorReqInputIpc(**obj.__dict__), + request ) content = {"success": success, "message": message} @@ -1195,6 +1337,25 @@ async def update_weights_from_tensor( ) +@dataclasses.dataclass +class UpdateWeightsFromDistributedReqInput: + names: List[str] + dtypes: List[str] + shapes: List[List[int]] + # The group name + group_name: str = "weight_update_group" + # Whether to flush the cache after updating weights + flush_cache: bool = True + # Whether to abort all requests before updating weights + abort_all_requests: bool = False + # Optional: Update weight version along with weights + weight_version: Optional[str] = None + # Optional format specification for loading + load_format: Optional[str] = None + # Whether to call torch.cuda.empty_cache() during flush + torch_empty_cache: bool = False + + @app.post("/update_weights_from_distributed") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weights_from_distributed( @@ -1203,7 +1364,8 @@ async def update_weights_from_distributed( """Update model parameter from distributed online.""" success, message = ( await _global_state.tokenizer_manager.update_weights_from_distributed( - obj, request + UpdateWeightsFromDistributedReqInputIpc(**obj.__dict__), + request, ) ) @@ -1214,12 +1376,25 @@ async def update_weights_from_distributed( return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) +@dataclasses.dataclass +class UpdateWeightsFromIPCReqInput: + # ZMQ socket paths for each device UUID + zmq_handles: Dict[str, str] + # Whether to flush cache after weight update + flush_cache: bool = True + # Optional: Update weight version along with weights + weight_version: Optional[str] = None + # Whether to call torch.cuda.empty_cache() during flush + torch_empty_cache: bool = False + + @app.post("/update_weights_from_ipc") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Request): """Update the weights from IPC (Inter-Process Communication) for checkpoint-engine integration.""" success, message = await _global_state.tokenizer_manager.update_weights_from_ipc( - obj, request + UpdateWeightsFromIPCReqInputIpc(**obj.__dict__), + request, ) content = {"success": success, "message": message} @@ -1262,12 +1437,21 @@ async def update_weight_version(obj: UpdateWeightVersionReqInput, request: Reque ) +@dataclasses.dataclass +class GetWeightsByNameReqInput: + name: str + truncate_size: int = 100 + + @app.api_route("/get_weights_by_name", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) async def get_weights_by_name(obj: GetWeightsByNameReqInput, request: Request): """Get model parameter by name.""" try: - ret = await _global_state.tokenizer_manager.get_weights_by_name(obj, request) + ret = await _global_state.tokenizer_manager.get_weights_by_name( + GetWeightsByNameReqInputIpc(**obj.__dict__), + request, + ) if ret is None: return _create_error_response("Get parameter by name failed") else: @@ -1276,6 +1460,11 @@ async def get_weights_by_name(obj: GetWeightsByNameReqInput, request: Request): return _create_error_response(e) +@dataclasses.dataclass +class ReleaseMemoryOccupationReqInput: + tags: Optional[List[str]] = None + + @app.api_route("/release_memory_occupation", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) async def release_memory_occupation( @@ -1283,11 +1472,19 @@ async def release_memory_occupation( ): """Release GPU memory occupation temporarily.""" try: - await _global_state.tokenizer_manager.release_memory_occupation(obj, request) + await _global_state.tokenizer_manager.release_memory_occupation( + ReleaseMemoryOccupationReqInputIpc(**obj.__dict__), + request, + ) except Exception as e: return _create_error_response(e) +@dataclasses.dataclass +class ResumeMemoryOccupationReqInput: + tags: Optional[List[str]] = None + + @app.api_route("/resume_memory_occupation", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) async def resume_memory_occupation( @@ -1295,11 +1492,18 @@ async def resume_memory_occupation( ): """Resume GPU memory occupation.""" try: - await _global_state.tokenizer_manager.resume_memory_occupation(obj, request) + await _global_state.tokenizer_manager.resume_memory_occupation( + ResumeMemoryOccupationReqInputIpc(**obj.__dict__), + request, + ) except Exception as e: return _create_error_response(e) +@dataclasses.dataclass +class CheckWeightsReqInput: + action: str = "checksum" + @app.api_route("/weights_checker", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) async def check_weights( @@ -1308,7 +1512,10 @@ async def check_weights( if obj is None: obj = CheckWeightsReqInput() success, message, ranks, per_engine_checksum = ( - await _global_state.tokenizer_manager.check_weights(obj, request) + await _global_state.tokenizer_manager.check_weights( + CheckWeightsReqInputIpc(**obj.__dict__), + request, + ) ) body = {"success": success, "message": message} if ranks is not None: @@ -1318,6 +1525,11 @@ async def check_weights( return ORJSONResponse(body, status_code=200 if success else HTTPStatus.BAD_REQUEST) +@dataclasses.dataclass +class SlowDownReqInput: + forward_sleep_time: Optional[float] + + @app.api_route("/slow_down", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) async def slow_down(obj: SlowDownReqInput, request: Request): @@ -1327,16 +1539,34 @@ async def slow_down(obj: SlowDownReqInput, request: Request): to let it run in full batch size. """ try: - await _global_state.tokenizer_manager.slow_down(obj, request) + await _global_state.tokenizer_manager.slow_down( + SlowDownReqInputIpc(**obj.__dict__), + request, + ) except Exception as e: return _create_error_response(e) +@dataclasses.dataclass +class LoadLoRAAdapterReqInput: + # The name of the lora module to newly loaded. + lora_name: str + # The path of loading. + lora_path: str + # Whether to pin the LoRA adapter in memory. + pinned: bool = False + # The unique identifier for the LoRA adapter, which automatically generated in the `TokenizerManager`. + lora_id: Optional[str] = None + + @app.api_route("/load_lora_adapter", methods=["POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) async def load_lora_adapter(obj: LoadLoRAAdapterReqInput, request: Request): """Load a new LoRA adapter without re-launching the server.""" - result = await _global_state.tokenizer_manager.load_lora_adapter(obj, request) + result = await _global_state.tokenizer_manager.load_lora_adapter( + LoadLoRAAdapterReqInputIpc(**obj.__dict__), + request, + ) if result.success: return ORJSONResponse( @@ -1350,13 +1580,25 @@ async def load_lora_adapter(obj: LoadLoRAAdapterReqInput, request: Request): ) +@dataclasses.dataclass +class LoadLoRAAdapterFromTensorsReqInput: + lora_name: str + config_dict: Dict[str, Any] + serialized_tensors: str + pinned: bool = False + added_tokens_config: Optional[Dict[str, Any]] = None + lora_id: Optional[str] = None + load_format: Optional[str] = None + + @app.api_route("/load_lora_adapter_from_tensors", methods=["POST"]) async def load_lora_adapter_from_tensors( obj: LoadLoRAAdapterFromTensorsReqInput, request: Request ): """Load a new LoRA adapter from tensors without re-launching the server.""" result = await _global_state.tokenizer_manager.load_lora_adapter_from_tensors( - obj, request + LoadLoRAAdapterFromTensorsReqInputIpc(**obj.__dict__), + request, ) if result.success: @@ -1365,11 +1607,22 @@ async def load_lora_adapter_from_tensors( return ORJSONResponse(result, status_code=HTTPStatus.BAD_REQUEST) +@dataclasses.dataclass +class UnloadLoRAAdapterReqInput: + # The name of lora module to unload. + lora_name: str + # The unique identifier for the LoRA adapter, which automatically generated in the `TokenizerManager`. + lora_id: Optional[str] = None + + @app.api_route("/unload_lora_adapter", methods=["POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) async def unload_lora_adapter(obj: UnloadLoRAAdapterReqInput, request: Request): """Load a new LoRA adapter without re-launching the server.""" - result = await _global_state.tokenizer_manager.unload_lora_adapter(obj, request) + result = await _global_state.tokenizer_manager.unload_lora_adapter( + UnloadLoRAAdapterReqInputIpc(**obj.__dict__), + request, + ) if result.success: return ORJSONResponse( @@ -1383,11 +1636,21 @@ async def unload_lora_adapter(obj: UnloadLoRAAdapterReqInput, request: Request): ) +@dataclasses.dataclass +class OpenSessionReqInput: + capacity_of_str_len: int + session_id: Optional[str] = None + streaming: Optional[bool] = None + timeout: Optional[float] = None + @app.api_route("/open_session", methods=["GET", "POST"]) async def open_session(obj: OpenSessionReqInput, request: Request): """Open a session, and return its unique session id.""" try: - session_id = await _global_state.tokenizer_manager.open_session(obj, request) + session_id = await _global_state.tokenizer_manager.open_session( + OpenSessionReqInputIpc(**obj.__dict__), + request, + ) if session_id is None: raise Exception( "Failed to open the session. Check if a session with the same id is still open." @@ -1397,23 +1660,49 @@ async def open_session(obj: OpenSessionReqInput, request: Request): return _create_error_response(e) +@dataclasses.dataclass +class CloseSessionReqInput: + session_id: str + @app.api_route("/close_session", methods=["GET", "POST"]) async def close_session(obj: CloseSessionReqInput, request: Request): """Close the session.""" try: - await _global_state.tokenizer_manager.close_session(obj, request) + await _global_state.tokenizer_manager.close_session( + CloseSessionReqInputIpc(**obj.__dict__), + request, + ) return Response(status_code=200) except Exception as e: return _create_error_response(e) +@dataclasses.dataclass +class ConfigureLoggingReq: + log_requests: Optional[bool] = None + log_requests_level: Optional[int] = None + log_requests_format: Optional[str] = None + log_level: Optional[str] = None + dump_requests_folder: Optional[str] = None + dump_requests_threshold: Optional[int] = None + crash_dump_folder: Optional[str] = None + dump_requests_exclude_meta_keys: Optional[List[str]] = None + + @app.api_route("/configure_logging", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) async def configure_logging(obj: ConfigureLoggingReq, request: Request): """Configure the request logging options.""" - _global_state.tokenizer_manager.configure_logging(obj) + _global_state.tokenizer_manager.configure_logging( + ConfigureLoggingReqIpc(**obj.__dict__) + ) return Response(status_code=200) +@dataclasses.dataclass +class AbortReq: + rid: Optional[str] = None + # Whether to abort all requests + abort_all: bool = False @app.post("/abort_request") @auth_level(AuthLevel.ADMIN_OPTIONAL) @@ -1470,22 +1759,45 @@ async def separate_reasoning_request(obj: SeparateReasoningReqInput, request: Re return ORJSONResponse(content=response_data, status_code=200) +@dataclasses.dataclass +class PauseGenerationReqInput: + mode: Literal["abort", "retract", "in_place"] = "abort" + + def __post_init__(self): + allowed = ["abort", "retract", "in_place"] + if self.mode not in allowed: + raise ValueError( + f"Invalid mode: {self.mode!r}. " f"Expected one of {allowed}." + ) + + + + @app.post("/pause_generation") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def pause_generation(obj: PauseGenerationReqInput, request: Request): """Pause generation.""" - await _global_state.tokenizer_manager.pause_generation(obj) + await _global_state.tokenizer_manager.pause_generation( + PauseGenerationReqInputIpc(**obj.__dict__) + ) return ORJSONResponse( content={"message": "Generation paused successfully.", "status": "ok"}, status_code=200, ) +@dataclasses.dataclass +class ContinueGenerationReqInput: + torch_empty_cache: bool = True + + @app.post("/continue_generation") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def continue_generation(obj: ContinueGenerationReqInput, request: Request): """Continue generation.""" - await _global_state.tokenizer_manager.continue_generation(obj) + await _global_state.tokenizer_manager.continue_generation( + ContinueGenerationReqInputIpc(**obj.__dict__) + ) return ORJSONResponse( content={"message": "Generation continued successfully.", "status": "ok"}, status_code=200, diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 799db26cc676..1bd422803e3b 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -209,6 +209,7 @@ class Envs: SGLANG_LOG_REQUEST_HEADERS = EnvTuple(tuple()) SGLANG_LOG_SCHEDULER_STATUS_TARGET = EnvStr("") SGLANG_LOG_SCHEDULER_STATUS_INTERVAL = EnvFloat(60.0) + SGLANG_LOG_PICKLE_IPC_OBJECTS = EnvBool(False) # SGLang CI SGLANG_IS_IN_CI = EnvBool(False) diff --git a/python/sglang/srt/managers/communicator.py b/python/sglang/srt/managers/communicator.py index 3080f6a755f2..015f0098e2a6 100644 --- a/python/sglang/srt/managers/communicator.py +++ b/python/sglang/srt/managers/communicator.py @@ -4,6 +4,7 @@ import copy from collections import deque from typing import Deque, Generic, List, Optional, TypeVar +from sglang.srt.managers.io_struct import sock_send import zmq @@ -41,7 +42,7 @@ async def queueing_call(self, obj: T): assert self._result_values is None if obj is not None: - self._sender.send_pyobj(obj) + sock_send(self._sender, obj) self._result_event = asyncio.Event() self._result_values = [] @@ -61,7 +62,7 @@ async def watching_call(self, obj): self._result_event = asyncio.Event() if obj is not None: - self._sender.send_pyobj(obj) + sock_send(self._sender, obj) # Capture local refs before await -- after event fires, the first # awakened coroutine clears shared state; later awaiters use local refs. diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index e84a593bf4ad..7ebcae8d4ace 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -36,6 +36,8 @@ ProfileReq, TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, + sock_send, + sock_recv, ) from sglang.srt.managers.load_snapshot import create_load_snapshot_reader from sglang.srt.managers.schedule_batch import Req @@ -205,12 +207,12 @@ def __init__( def send_to_all_workers(self, obj): for i, worker in enumerate(self.workers): if self.status[i]: - worker.send_pyobj(obj) + sock_send(worker, obj) def send_control_message(self, obj): # Send control messages to first worker of tp group for worker in self.workers[:: self.control_message_step]: - worker.send_pyobj(obj) + sock_send(worker, obj) def update_active_ranks(self, ranks: ActiveRanksOutput): self.status = ranks.status @@ -384,7 +386,7 @@ def _broadcast_ports_as_server( logger.debug(f"Received handshake from node {client_rank}") # Send worker ports to client - rep_socket.send_pyobj(worker_ports) + sock_send(rep_socket, worker_ports) connected_clients += 1 logger.debug( f"Sent worker ports to {connected_clients}/{expected_clients} nodes" @@ -418,7 +420,7 @@ def _reply_ports_as_server(self, rep_socket: zmq.Socket, worker_ports: List[int] logger.debug(f"Received handshake from node {client_rank}") # Send worker ports to client - rep_socket.send_pyobj(worker_ports) + sock_send(rep_socket, 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]: @@ -431,10 +433,10 @@ def _receive_ports_as_client(self, endpoint: str, node_rank: int) -> List[int]: try: # Send handshake with our node rank - req_socket.send(str(node_rank).encode()) + sock_send(req_socket, str(node_rank).encode()) # Receive worker ports - worker_ports = req_socket.recv_pyobj() + worker_ports = sock_recv(req_socket) logger.debug(f"Received {len(worker_ports)} worker ports from node 0") return worker_ports except zmq.Again: @@ -597,7 +599,7 @@ def launch_tensor_parallel_group( def maybe_external_dp_rank_routing(self, req: Req): if req.routed_dp_rank is not None: logger.debug(f"Direct routing to DP rank {req.routed_dp_rank}") - self.workers[req.routed_dp_rank].send_pyobj(req) + sock_send(self.workers[req.routed_dp_rank], req) return True return False @@ -608,7 +610,7 @@ def round_robin_scheduler(self, req: Req): while True: if self.status[self.round_robin_counter]: logger.debug(f"Choose worker {self.round_robin_counter}") - self.workers[self.round_robin_counter].send_pyobj(req) + sock_send(self.workers[self.round_robin_counter], req) self.round_robin_counter = (self.round_robin_counter + 1) % len( self.workers ) @@ -626,13 +628,13 @@ def follow_bootstrap_room_scheduler(self, req: Req): "prefill or decode instances; send to the router instead." ) target_rank = req.bootstrap_room % len(self.workers) - self.workers[target_rank].send_pyobj(req) + sock_send(self.workers[target_rank], req) def total_requests_scheduler(self, req: Req): if self.maybe_external_dp_rank_routing(req): return target_worker = self.dp_budget.dispatch(LoadBalanceMethod.TOTAL_REQUESTS) - self.workers[target_worker].send_pyobj(req) + sock_send(self.workers[target_worker], req) def total_tokens_scheduler(self, req: Req): if self.maybe_external_dp_rank_routing(req): @@ -641,14 +643,14 @@ def total_tokens_scheduler(self, req: Req): target_worker = self.dp_budget.dispatch( LoadBalanceMethod.TOTAL_TOKENS, estimated_tokens=estimated_tokens ) - self.workers[target_worker].send_pyobj(req) + sock_send(self.workers[target_worker], req) def event_loop(self): while True: while True: self.soft_watchdog.feed() try: - recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK) + recv_req = sock_recv(self.recv_from_tokenizer, flags=zmq.NOBLOCK) except zmq.ZMQError: break self._request_dispatcher(recv_req) diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index 5f67f1db16ec..122e1620971c 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -34,6 +34,8 @@ BatchTokenIDOutput, ConfigureLoggingReq, FreezeGCReq, + sock_recv, + sock_send, ) from sglang.srt.managers.multi_tokenizer_mixin import MultiHttpWorkerDetokenizerMixin from sglang.srt.observability.cpu_monitor import start_cpu_monitor_thread @@ -144,10 +146,10 @@ def event_loop(self): """The event loop that handles requests""" while True: with self.soft_watchdog.disable(): - recv_obj = self.recv_from_scheduler.recv_pyobj() + recv_obj = sock_recv(self.recv_from_scheduler) output = self._request_dispatcher(recv_obj) if output is not None: - self.send_to_tokenizer.send_pyobj(output) + sock_send(self.send_to_tokenizer, output) self.soft_watchdog.feed() def trim_matched_stop( diff --git a/python/sglang/srt/managers/embed_types.py b/python/sglang/srt/managers/embed_types.py index 48d6859d9c2b..6cfe170d8139 100644 --- a/python/sglang/srt/managers/embed_types.py +++ b/python/sglang/srt/managers/embed_types.py @@ -37,7 +37,8 @@ class PositionalEmbeds: positions: List of positions where embeddings should be injected. """ - embeds: Union[List[torch.Tensor], torch.Tensor] + # Union[List[torch.Tensor], torch.Tensor] + embeds: torch.Tensor positions: List[int] def __post_init__(self): diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 8b62511719aa..15f2bbac43fd 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -25,19 +25,25 @@ from collections import Counter from dataclasses import dataclass, field from enum import Enum -from typing import TYPE_CHECKING, Annotated, Any, Dict, List, Literal, Optional, Union +import logging +import msgspec +import pickle +from typing import TYPE_CHECKING, Annotated, Any, Dict, List, Optional, Union, Type import torch from pydantic import PlainValidator +from zmq import Socket +from zmq.asyncio import Socket as AsyncSocket + +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 BaseFinishReason, Modality from sglang.srt.multimodal.mm_utils import has_valid_data from sglang.srt.observability.req_time_stats import ( - APIServerReqTimeStats, - DPControllerReqTimeStats, SchedulerReqTimeStats, + ReqTimeStatsBase, ) from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.utils import ImageData, VideoData @@ -49,6 +55,7 @@ else: Image = Any +logger = logging.getLogger(__name__) @dataclass class BaseReq(ABC): @@ -73,6 +80,12 @@ def _validate_rid_uniqueness(self): ) +# The BaseReq IPC class for IPC object +class BaseReqIpc(msgspec.Struct, tag=True, kw_only=True): + rid: Optional[str] = None + http_worker_ipc: Optional[str] = None + + @dataclass class BaseBatchReq(ABC): rids: Optional[List[str]] = field(default=None, kw_only=True) @@ -84,8 +97,13 @@ def regenerate_rids(self): return self.rids -@dataclass -class SpeculativeDecodingMetricsMixin: +# The BaseBatchReq IPC class for IPC object +class BaseBatchReqIpc(msgspec.Struct, tag=True, kw_only=True): + rids: Optional[List[str]] = None + http_worker_ipcs: Optional[List[str]] = None + + +class SpeculativeDecodingMetricsMixinReq(BaseBatchReqIpc, kw_only=True): """ Mixin class containing speculative decoding metrics. @@ -725,14 +743,13 @@ def __getitem__(self, i): return sub -@dataclass -class TokenizedGenerateReqInput(BaseReq): +class TokenizedGenerateReqInput(BaseReqIpc, kw_only=True): # The input text - input_text: str + input_text: Optional[str] # str # The input token ids - input_ids: Optional[array[int]] + input_ids: Optional[array] # Optional[array[int]] # The multimodal inputs - mm_inputs: object + mm_inputs: Optional[object] # The sampling parameters sampling_params: SamplingParams # Whether to return the logprobs @@ -742,7 +759,7 @@ class TokenizedGenerateReqInput(BaseReq): # If return logprobs, the number of top logprobs to return at each position. top_logprobs_num: int # If return logprobs, the token id to return logprob for - token_ids_logprob: List[int] + token_ids_logprob: Optional[List[int]] # Whether to stream output stream: bool @@ -757,7 +774,7 @@ class TokenizedGenerateReqInput(BaseReq): return_indexer_topk: bool = False # The input embeds - input_embeds: Optional[Union[List[List[List[float]]], List[List[float]]]] = None + input_embeds: Optional[List[List[Union[float, List[float]]]]] = None # Embedding overrides to place at specific token positions. positional_embed_overrides: Optional[PositionalEmbeds] = None @@ -808,7 +825,7 @@ class TokenizedGenerateReqInput(BaseReq): token_type_ids: Optional[List[int]] = None - need_wait_for_mm_inputs: bool = False + need_wait_for_mm_inputs: Optional[bool] = None num_items_assigned: Optional[Dict[Modality, List[int]]] = None mm_data_mooncake: Optional[List] = None @@ -816,11 +833,11 @@ class TokenizedGenerateReqInput(BaseReq): multi_item_delimiter_indices: Optional[List[int]] = None # For observability - time_stats: Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] = None + # Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] + time_stats: Optional[ReqTimeStatsBase] = None -@dataclass -class BatchTokenizedGenerateReqInput(BaseBatchReq): +class BatchTokenizedGenerateReqInput(BaseBatchReqIpc, kw_only=True): # The batch of tokenized requests batch: List[TokenizedGenerateReqInput] @@ -1041,14 +1058,13 @@ def __getitem__(self, i): return sub -@dataclass -class TokenizedEmbeddingReqInput(BaseReq): +class TokenizedEmbeddingReqInput(BaseReqIpc, kw_only=True): # The input text - input_text: str + input_text: Optional[str] # str # The input token ids - input_ids: array[int] + input_ids: Optional[array] # array[int] # The image inputs - image_inputs: dict + image_inputs: Optional[dict] # The token type ids token_type_ids: List[int] # Dummy sampling params for compatibility @@ -1067,14 +1083,14 @@ class TokenizedEmbeddingReqInput(BaseReq): # Pre-computed delimiter indices for multi-item scoring multi_item_delimiter_indices: Optional[List[int]] = None # For observability - time_stats: Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] = None + # Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] + time_stats: Optional[ReqTimeStatsBase] = None # Whether to return pooled hidden states (pre-head transformer output) return_pooled_hidden_states: bool = False -@dataclass -class BatchTokenizedEmbeddingReqInput(BaseBatchReq): +class BatchTokenizedEmbeddingReqInput(BaseBatchReqIpc, kw_only=True): # The batch of tokenized embedding requests batch: List[TokenizedEmbeddingReqInput] @@ -1088,16 +1104,15 @@ def __iter__(self): return iter(self.batch) -@dataclass -class BatchTokenIDOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin): +class BatchTokenIDOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): # The finish reason - finished_reasons: List[BaseFinishReason] + finished_reasons: List[Optional[Dict[str, Any]]] # List[BaseFinishReason] # For incremental decoding decoded_texts: List[str] - decode_ids: List[array[int]] + decode_ids: List[array] # List[array[int]] read_offsets: List[int] # Only used when `--skip-tokenizer-init` is on - output_ids: Optional[List[array[int]]] + output_ids: Optional[List[array]] # Optional[List[array[int]]] # Detokenization configs skip_special_tokens: List[bool] spaces_between_special_tokens: List[bool] @@ -1110,64 +1125,63 @@ class BatchTokenIDOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin): cached_tokens: List[int] # Logprobs - input_token_logprobs_val: List[float] - input_token_logprobs_idx: List[int] - output_token_logprobs_val: List[float] - output_token_logprobs_idx: List[int] - input_top_logprobs_val: List[List] - input_top_logprobs_idx: List[List] - output_top_logprobs_val: List[List] - output_top_logprobs_idx: List[List] - input_token_ids_logprobs_val: List[List] - input_token_ids_logprobs_idx: List[List] - output_token_ids_logprobs_val: List[List] - output_token_ids_logprobs_idx: List[List] - output_token_entropy_val: List[float] + input_token_logprobs_val: Optional[List[float]] + input_token_logprobs_idx: Optional[List[int]] + output_token_logprobs_val: Optional[List[float]] + output_token_logprobs_idx: Optional[List[int]] + input_top_logprobs_val: Optional[List[List]] + input_top_logprobs_idx: Optional[List[List]] + output_top_logprobs_val: Optional[List[List]] + output_top_logprobs_idx: Optional[List[List]] + input_token_ids_logprobs_val: Optional[List[List]] + input_token_ids_logprobs_idx: Optional[List[List]] + output_token_ids_logprobs_val: Optional[List[List]] + output_token_ids_logprobs_idx: Optional[List[List]] + output_token_entropy_val: Optional[List[float]] # Hidden states - output_hidden_states: List[List[float]] + output_hidden_states: Optional[List[List[float]]] # Per-request routed experts (input + output tokens), shape # (token, layer, top_k). DetokenizerManager encodes to base64 into # BatchStrOutput; on the skip_tokenizer_init path the scheduler sends this # straight to TokenizerManager, which encodes on demand. - routed_experts: List[Optional[torch.Tensor]] + routed_experts: Optional[List[Optional[torch.Tensor]]] - indexer_topk: List[Optional[torch.Tensor]] + indexer_topk: Optional[List[Optional[torch.Tensor]]] # The information of placeholder tokens (e.g., image token) # idx is the index of the token in the prompt after expansion. # val is the length of padded tokens after expansion. - placeholder_tokens_idx: List[Optional[List[int]]] - placeholder_tokens_val: List[Optional[List[int]]] + placeholder_tokens_idx: Optional[List[Optional[List[int]]]] + placeholder_tokens_val: Optional[List[Optional[List[int]]]] # Number of times each request was retracted. - retraction_counts: List[int] + retraction_counts: Optional[List[int]] = None # The trainer step id. Used to know which step's weights are used for sampling. - token_steps: List[List[int]] = None + token_steps: Optional[List[List[int]]] = None # Load for DP balance - load: GetLoadsReqOutput = None + load: Optional[GetLoadsReqOutput] = None # Customized info customized_info: Optional[Dict[str, List[Any]]] = None # Detailed breakdown of cached tokens by source (device/host/storage) cached_tokens_details: Optional[List[Optional[Dict[str, Any]]]] = None # DP rank of the scheduler that processed each request - dp_ranks: Optional[List[int]] = None + dp_ranks: Optional[List[Optional[int]]] = None # For observability time_stats: Optional[List[SchedulerReqTimeStats]] = None -@dataclass -class BatchStrOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin): +class BatchStrOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): # The finish reason - finished_reasons: List[dict] + finished_reasons: List[Optional[Dict[str, Any]]] # List[dict] # The output decoded strings output_strs: List[str] # The token ids - output_ids: Optional[List[int]] + output_ids: Optional[List[array]] # Optional[List[int]] # Token counts prompt_tokens: List[int] @@ -1176,71 +1190,70 @@ class BatchStrOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin): cached_tokens: List[int] # Logprobs - input_token_logprobs_val: List[float] - input_token_logprobs_idx: List[int] - output_token_logprobs_val: List[float] - output_token_logprobs_idx: List[int] - input_top_logprobs_val: List[List] - input_top_logprobs_idx: List[List] - output_top_logprobs_val: List[List] - output_top_logprobs_idx: List[List] - input_token_ids_logprobs_val: List[List] - input_token_ids_logprobs_idx: List[List] - output_token_ids_logprobs_val: List[List] - output_token_ids_logprobs_idx: List[List] - output_token_entropy_val: List[float] + input_token_logprobs_val: Optional[List[float]] + input_token_logprobs_idx: Optional[List[int]] + output_token_logprobs_val: Optional[List[float]] + output_token_logprobs_idx: Optional[List[int]] + input_top_logprobs_val: Optional[List[List]] + input_top_logprobs_idx: Optional[List[List]] + output_top_logprobs_val: Optional[List[List]] + output_top_logprobs_idx: Optional[List[List]] + input_token_ids_logprobs_val: Optional[List[List]] + input_token_ids_logprobs_idx: Optional[List[List]] + output_token_ids_logprobs_val: Optional[List[List]] + output_token_ids_logprobs_idx: Optional[List[List]] + output_token_entropy_val: Optional[List[float]] # Hidden states - output_hidden_states: List[List[float]] + output_hidden_states: Optional[List[List[float]]] # Per-request routed experts, base64-encoded by DetokenizerManager off the # tokenizer hot path. Underlying tensor shape is (token, layer, top_k); # see BatchTokenIDOutput.routed_experts. - routed_experts: List[Optional[str]] + routed_experts: Optional[List[Optional[str]]] - indexer_topk: List[Optional[str]] + indexer_topk: Optional[List[Optional[str]]] # The information of placeholder tokens (e.g., image token) # idx is the index of the token in the prompt after expansion. # val is the length of padded tokens after expansion. - placeholder_tokens_idx: List[Optional[List[int]]] - placeholder_tokens_val: List[Optional[List[int]]] + placeholder_tokens_idx: Optional[List[Optional[List[int]]]] + placeholder_tokens_val: Optional[List[Optional[List[int]]]] # Number of times each request was retracted. - retraction_counts: List[int] + retraction_counts: Optional[List[int]] # The trainer step id. Used to know which step's weights are used for sampling. - token_steps: List[List[int]] = None + token_steps: Optional[List[List[int]]] = None # Load for DP balance - load: GetLoadsReqOutput = None + load: Optional[GetLoadsReqOutput] = None # Customized info customized_info: Optional[Dict[str, List[Any]]] = None # Detailed breakdown of cached tokens by source (device/host/storage) cached_tokens_details: Optional[List[Optional[Dict[str, Any]]]] = None # DP rank of the scheduler that processed each request - dp_ranks: Optional[List[int]] = None + dp_ranks: Optional[List[Optional[int]]] = None # For observability time_stats: Optional[List[SchedulerReqTimeStats]] = None -@dataclass -class BatchEmbeddingOutput(BaseBatchReq): +class BatchEmbeddingOutput(BaseBatchReqIpc, kw_only=True): # The finish reason - finished_reasons: List[BaseFinishReason] + finished_reasons: List[Optional[Dict[str, Any]]] # List[BaseFinishReason] # The output embedding - embeddings: Union[List[List[float]], List[Dict[int, float]]] + embeddings: List[Union[List[float], Dict[int, float]]] # Token counts prompt_tokens: List[int] cached_tokens: List[int] # Placeholder token info - placeholder_tokens_idx: List[Optional[List[int]]] - placeholder_tokens_val: List[Optional[List[int]]] + placeholder_tokens_idx: Optional[List[Optional[List[int]]]] + placeholder_tokens_val: Optional[List[Optional[List[int]]]] # Number of times each request was retracted. - retraction_counts: List[int] + retraction_counts: Optional[List[int]] # Detailed breakdown of cached tokens by source (device/host/storage) cached_tokens_details: Optional[List[Optional[Dict[str, Any]]]] = None @@ -1249,73 +1262,62 @@ class BatchEmbeddingOutput(BaseBatchReq): # Optional pooled hidden states (pre-head transformer output). # Sent as a single stacked tensor to minimize pickle overhead. - pooled_hidden_states: Optional[ - Union[List[Optional[torch.Tensor]], torch.Tensor] - ] = None + # Optional[Union[List[Optional[torch.Tensor]], torch.Tensor]] + pooled_hidden_states: Optional[List[Optional[torch.Tensor]]] = None -@dataclass -class ClearHiCacheReqInput(BaseReq): +class ClearHiCacheReqInput(BaseReqIpc, kw_only=True): pass -@dataclass -class ClearHiCacheReqOutput(BaseReq): +class ClearHiCacheReqOutput(BaseReqIpc, kw_only=True): success: bool -@dataclass -class FlushCacheReqInput(BaseReq): +class FlushCacheReqInput(BaseReqIpc, kw_only=True): timeout_s: Optional[float] = None -@dataclass -class FlushCacheReqOutput(BaseReq): +class FlushCacheReqOutput(BaseReqIpc, kw_only=True): success: bool message: str = "" -@dataclass -class AddExternalCorpusReqInput(BaseReq): +class AddExternalCorpusReqInput(BaseReqIpc, 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(BaseReqIpc, kw_only=True): success: bool corpus_id: str = "" message: str = "" loaded_token_count: int = 0 -@dataclass -class RemoveExternalCorpusReqInput(BaseReq): +class RemoveExternalCorpusReqInput(BaseReqIpc, kw_only=True): corpus_id: str -@dataclass -class RemoveExternalCorpusReqOutput(BaseReq): +class RemoveExternalCorpusReqOutput(BaseReqIpc, kw_only=True): success: bool message: str = "" -@dataclass -class ListExternalCorporaReqInput(BaseReq): +class ListExternalCorporaReqInput(BaseReqIpc, kw_only=True): pass -@dataclass -class ListExternalCorporaReqOutput(BaseReq): +class ListExternalCorporaReqOutput(BaseReqIpc, kw_only=True): success: bool corpus_token_counts: Dict[str, int] = field(default_factory=dict) message: str = "" -@dataclass -class AttachHiCacheStorageReqInput(BaseReq): +# see aslo AttachHiCacheStorageReqInput in the http_server.py +class AttachHiCacheStorageReqInput(BaseReqIpc, 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: @@ -1328,48 +1330,24 @@ class AttachHiCacheStorageReqInput(BaseReq): hicache_storage_prefetch_policy: Optional[str] = None hicache_write_policy: Optional[str] = None - def __post_init__(self): - if self.hicache_storage_prefetch_policy is None: - pass - else: - allowed = ["best_effort", "wait_complete", "timeout"] - if self.hicache_storage_prefetch_policy not in allowed: - raise ValueError( - f"Invalid hicache_storage_prefetch_policy: {self.hicache_storage_prefetch_policy!r}. " - f"Expected one of {allowed}." - ) - - if self.hicache_write_policy is None: - return - allowed = ["write_back", "write_through", "write_through_selective"] - if self.hicache_write_policy not in allowed: - raise ValueError( - f"Invalid hicache_write_policy: {self.hicache_write_policy!r}. " - f"Expected one of {allowed}." - ) - -@dataclass -class AttachHiCacheStorageReqOutput(BaseReq): +class AttachHiCacheStorageReqOutput(BaseReqIpc, kw_only=True): success: bool message: str = "" -@dataclass -class DetachHiCacheStorageReqInput(BaseReq): +class DetachHiCacheStorageReqInput(BaseReqIpc, kw_only=True): """Dynamically detach (disable) HiCache storage backend at runtime.""" pass -@dataclass -class DetachHiCacheStorageReqOutput(BaseReq): +class DetachHiCacheStorageReqOutput(BaseReqIpc, kw_only=True): success: bool message: str = "" -@dataclass -class PauseGenerationReqInput(BaseReq): +class PauseGenerationReqInput(BaseReqIpc, kw_only=True): """ Note that the PauseGenerationRequests is only supported in SGLang Server. abort: Abort and return all requests currently being processed. @@ -1388,7 +1366,7 @@ class PauseGenerationReqInput(BaseReq): recomputed after continue_generation. """ - mode: Literal["abort", "retract", "in_place"] = "abort" + mode: str = "abort" def __post_init__(self): allowed = ["abort", "retract", "in_place"] @@ -1398,8 +1376,7 @@ def __post_init__(self): ) -@dataclass -class ContinueGenerationReqInput(BaseReq): +class ContinueGenerationReqInput(BaseReqIpc, 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 @@ -1408,22 +1385,19 @@ class ContinueGenerationReqInput(BaseReq): torch_empty_cache: bool = True -@dataclass -class TokenizerWorkerRegistration: +class TokenizerWorkerRegistrationReq(BaseReqIpc, kw_only=True): """Sent by each TokenizerWorker on startup to register its IPC name with the router.""" worker_ipc_name: str -@dataclass -class PauseContinueBroadcast: +class PauseContinueBroadcastReq(BaseReqIpc, kw_only=True): """Broadcast from router to all workers to set is_pause state.""" is_pause: bool -@dataclass -class UpdateWeightFromDiskReqInput(BaseReq): +class UpdateWeightFromDiskReqInput(BaseReqIpc, kw_only=True): # The model path with the new weights model_path: str # The format to load the weights @@ -1448,16 +1422,14 @@ class UpdateWeightFromDiskReqInput(BaseReq): manifest: Optional[Dict[str, Any]] = None -@dataclass -class UpdateWeightFromDiskReqOutput(BaseReq): +class UpdateWeightFromDiskReqOutput(BaseReqIpc, kw_only=True): success: bool message: str # Number of paused requests during weight sync. num_paused_requests: Optional[int] = 0 -@dataclass -class UpdateWeightsFromDistributedReqInput(BaseReq): +class UpdateWeightsFromDistributedReqInput(BaseReqIpc, kw_only=True): names: List[str] dtypes: List[str] shapes: List[List[int]] @@ -1475,21 +1447,20 @@ class UpdateWeightsFromDistributedReqInput(BaseReq): torch_empty_cache: bool = False -@dataclass -class UpdateWeightsFromDistributedReqOutput(BaseReq): +class UpdateWeightsFromDistributedReqOutput(BaseReqIpc, kw_only=True): success: bool message: str -@dataclass -class UpdateWeightsFromTensorReqInput(BaseReq): +class UpdateWeightsFromTensorReqInput(BaseReqIpc, kw_only=True): """Update model weights from tensor input. - Tensors are serialized for transmission - Data is structured in JSON for easy transmission over HTTP """ - serialized_named_tensors: List[Union[str, bytes]] + # List[Union[str, bytes]] + serialized_named_tensors: List[bytes] # Optional format specification for loading load_format: Optional[str] = None # Whether to flush the cache after updating weights @@ -1504,14 +1475,12 @@ class UpdateWeightsFromTensorReqInput(BaseReq): torch_empty_cache: bool = False -@dataclass -class UpdateWeightsFromTensorReqOutput(BaseReq): +class UpdateWeightsFromTensorReqOutput(BaseReqIpc, kw_only=True): success: bool message: str -@dataclass -class InitWeightsSendGroupForRemoteInstanceReqInput(BaseReq): +class InitWeightsSendGroupForRemoteInstanceReqInput(BaseReqIpc, kw_only=True): # The master address master_address: str # The ports for each rank's communication group @@ -1528,8 +1497,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(BaseReqIpc, kw_only=True): # ZMQ socket paths for each device UUID zmq_handles: Dict[str, str] # Whether to flush cache after weight update @@ -1540,20 +1508,17 @@ class UpdateWeightsFromIPCReqInput(BaseReq): torch_empty_cache: bool = False -@dataclass -class UpdateWeightsFromIPCReqOutput(BaseReq): +class UpdateWeightsFromIPCReqOutput(BaseReqIpc, kw_only=True): success: bool message: str -@dataclass -class InitWeightsSendGroupForRemoteInstanceReqOutput(BaseReq): +class InitWeightsSendGroupForRemoteInstanceReqOutput(BaseReqIpc, kw_only=True): success: bool message: str -@dataclass -class SendWeightsToRemoteInstanceReqInput(BaseReq): +class SendWeightsToRemoteInstanceReqInput(BaseReqIpc, kw_only=True): # The master address master_address: str # The ports for each rank's communication group @@ -1562,27 +1527,23 @@ class SendWeightsToRemoteInstanceReqInput(BaseReq): group_name: str = "weight_send_group" -@dataclass -class SendWeightsToRemoteInstanceReqOutput(BaseReq): +class SendWeightsToRemoteInstanceReqOutput(BaseReqIpc, kw_only=True): success: bool message: str -@dataclass -class UpdateExpertBackupReq(BaseReq): +class UpdateExpertBackupReq(BaseReqIpc, kw_only=True): pass -@dataclass -class BackupDramReq(BaseReq): +class BackupDramReq(BaseReqIpc, kw_only=True): rank: int weight_pointer_map: Dict[str, Any] session_id: str buffer_size: int -@dataclass -class InitWeightsUpdateGroupReqInput(BaseReq): +class InitWeightsUpdateGroupReqInput(BaseReqIpc, kw_only=True): # The master address master_address: str # The master port @@ -1597,19 +1558,16 @@ class InitWeightsUpdateGroupReqInput(BaseReq): backend: str = "nccl" -@dataclass -class InitWeightsUpdateGroupReqOutput(BaseReq): +class InitWeightsUpdateGroupReqOutput(BaseReqIpc, kw_only=True): success: bool message: str -@dataclass -class DestroyWeightsUpdateGroupReqInput(BaseReq): +class DestroyWeightsUpdateGroupReqInput(BaseReqIpc, kw_only=True): group_name: str = "weight_update_group" -@dataclass -class DestroyWeightsUpdateGroupReqOutput(BaseReq): +class DestroyWeightsUpdateGroupReqOutput(BaseReqIpc, kw_only=True): success: bool message: str @@ -1622,65 +1580,54 @@ class UpdateWeightVersionReqInput(BaseReq): abort_all_requests: bool = True -@dataclass -class GetWeightsByNameReqInput(BaseReq): +class GetWeightsByNameReqInput(BaseReqIpc, kw_only=True): name: str truncate_size: int = 100 -@dataclass -class GetWeightsByNameReqOutput(BaseReq): +class GetWeightsByNameReqOutput(BaseReqIpc, kw_only=True): parameter: list -@dataclass -class ReleaseMemoryOccupationReqInput(BaseReq): +class ReleaseMemoryOccupationReqInput(BaseReqIpc, 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(BaseReqIpc, kw_only=True): pass -@dataclass -class ResumeMemoryOccupationReqInput(BaseReq): +class ResumeMemoryOccupationReqInput(BaseReqIpc, 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(BaseReqIpc, kw_only=True): pass -@dataclass -class CheckWeightsReqInput(BaseReq): +class CheckWeightsReqInput(BaseReqIpc, kw_only=True): action: str = "checksum" -@dataclass -class CheckWeightsReqOutput(BaseReq): +class CheckWeightsReqOutput(BaseReqIpc, kw_only=True): success: bool message: str payload: Optional[Dict] = None -@dataclass -class SlowDownReqInput(BaseReq): +class SlowDownReqInput(BaseReqIpc, kw_only=True): forward_sleep_time: Optional[float] -@dataclass -class SlowDownReqOutput(BaseReq): +class SlowDownReqOutput(BaseReqIpc, kw_only=True): pass -@dataclass -class AbortReq(BaseReq): +class AbortReq(BaseReqIpc, kw_only=True): # Whether to abort all requests abort_all: bool = False # The finished reason data @@ -1693,28 +1640,23 @@ def __post_init__(self): self.rid = "" -@dataclass -class ActiveRanksOutput(BaseReq): +class ActiveRanksOutput(BaseReqIpc, kw_only=True): status: List[bool] -@dataclass -class GetInternalStateReq(BaseReq): +class GetInternalStateReq(BaseReqIpc, kw_only=True): pass -@dataclass -class GetInternalStateReqOutput(BaseReq): +class GetInternalStateReqOutput(BaseReqIpc, kw_only=True): internal_state: Dict[Any, Any] -@dataclass -class SetInternalStateReq(BaseReq): +class SetInternalStateReq(BaseReqIpc, kw_only=True): server_args: Dict[str, Any] -@dataclass -class SetInternalStateReqOutput(BaseReq): +class SetInternalStateReqOutput(BaseReqIpc, kw_only=True): updated: bool server_args: Dict[str, Any] @@ -1750,9 +1692,8 @@ class ProfileReqType(Enum): STOP_PROFILE = 2 -@dataclass -class ProfileReq(BaseReq): - type: ProfileReqType +class ProfileReq(BaseReqIpc, kw_only=True): + req_type: ProfileReqType output_dir: Optional[str] = None start_step: Optional[int] = None num_steps: Optional[int] = None @@ -1766,19 +1707,16 @@ class ProfileReq(BaseReq): profile_stages: Optional[List[str]] = None -@dataclass -class ProfileReqOutput(BaseReq): +class ProfileReqOutput(BaseReqIpc, kw_only=True): success: bool message: str -@dataclass -class FreezeGCReq(BaseReq): +class FreezeGCReq(BaseReqIpc, kw_only=True): pass -@dataclass -class ConfigureLoggingReq(BaseReq): +class ConfigureLoggingReq(BaseReqIpc, kw_only=True): log_requests: Optional[bool] = None log_requests_level: Optional[int] = None log_requests_format: Optional[str] = None @@ -1789,27 +1727,23 @@ class ConfigureLoggingReq(BaseReq): dump_requests_exclude_meta_keys: Optional[List[str]] = None -@dataclass -class OpenSessionReqInput(BaseReq): +class OpenSessionReqInput(BaseReqIpc, 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(BaseReqIpc, kw_only=True): session_id: str -@dataclass -class OpenSessionReqOutput(BaseReq): +class OpenSessionReqOutput(BaseReqIpc, kw_only=True): session_id: Optional[str] success: bool -@dataclass -class HealthCheckOutput(BaseReq): +class HealthCheckOutput(BaseReqIpc, kw_only=True): pass @@ -1819,13 +1753,11 @@ class ExpertDistributionReqType(Enum): DUMP_RECORD = 3 -@dataclass -class ExpertDistributionReq(BaseReq): +class ExpertDistributionReq(BaseReqIpc, kw_only=True): action: ExpertDistributionReqType -@dataclass -class ExpertDistributionReqOutput(BaseReq): +class ExpertDistributionReqOutput(BaseReqIpc, kw_only=True): pass @@ -1852,7 +1784,6 @@ class ParseFunctionCallReq(BaseReq): None # Specify the parser type, e.g. 'llama3', 'qwen25', or 'mistral'. If not specified, tries all. ) - @dataclass class SeparateReasoningReqInput(BaseReq): text: str # The text to parse. @@ -1865,20 +1796,17 @@ class VertexGenerateReqInput(BaseReq): parameters: Optional[dict] = None -@dataclass -class RpcReqInput(BaseReq): +class RpcReqInput(BaseReqIpc, kw_only=True): method: str parameters: Optional[Dict] = None -@dataclass -class RpcReqOutput(BaseReq): +class RpcReqOutput(BaseReqIpc, kw_only=True): success: bool message: str -@dataclass -class LoadLoRAAdapterReqInput(BaseReq): +class LoadLoRAAdapterReqInput(BaseReqIpc, kw_only=True): # The name of the lora module to newly loaded. lora_name: str # The path of loading. @@ -1897,8 +1825,7 @@ def to_ref(self) -> LoRARef: ) -@dataclass -class UnloadLoRAAdapterReqInput(BaseReq): +class UnloadLoRAAdapterReqInput(BaseReqIpc, 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`. @@ -1911,8 +1838,7 @@ def to_ref(self) -> LoRARef: ) -@dataclass -class LoadLoRAAdapterFromTensorsReqInput(BaseReq): +class LoadLoRAAdapterFromTensorsReqInput(BaseReqIpc, kw_only=True): lora_name: str config_dict: Dict[str, Any] serialized_tensors: str @@ -1930,8 +1856,7 @@ def to_ref(self) -> LoRARef: ) -@dataclass -class LoRAUpdateOutput(BaseReq): +class LoRAUpdateOutput(BaseReqIpc, kw_only=True): success: bool error_message: Optional[str] = None loaded_adapters: Optional[Dict[str, LoRARef]] = None @@ -1947,9 +1872,8 @@ class BlockReqType(Enum): UNBLOCK = 2 -@dataclass -class BlockReqInput(BaseReq): - type: BlockReqType +class BlockReqInput(BaseReqIpc, kw_only=True): + req_type: BlockReqType @dataclass @@ -2036,15 +1960,14 @@ class QueueMetrics: retracted: int = field(metadata={"metric": ("gauge", "Retracted requests count")}) -@dataclass -class GetLoadsReqInput(BaseReq): +class GetLoadsReqInput(BaseReqIpc, 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): @@ -2058,45 +1981,36 @@ def __post_init__(self): ) -@dataclass -class GetLoadsReqOutput(BaseReq): +class GetLoadsReqOutput(BaseReqIpc, kw_only=True): """Per-DP-rank load metrics for /v1/loads endpoint.""" - dp_rank: int - timestamp: float + dp_rank: Optional[int] + timestamp: Optional[float] - num_running_reqs: int = field( - metadata={"metric": ("gauge", "Number of running requests")} - ) - num_waiting_reqs: int = field( - metadata={"metric": ("gauge", "Number of waiting requests")} - ) - num_used_tokens: int = field( - metadata={"metric": ("gauge", "Number of tokens in use")} - ) + # metadata={"metric": ("gauge", "Number of running requests")} + num_running_reqs: Optional[int] + # metadata={"metric": ("gauge", "Number of waiting requests")} + num_waiting_reqs: Optional[int] + # metadata={"metric": ("gauge", "Number of tokens in use")} + num_used_tokens: Optional[int] # num_used_tokens + pending prefill tokens (waiting-queue seqlen, incl. # disagg bootstrap/prealloc/transfer queues). Used for DP balance. - num_total_tokens: int = field( - metadata={"metric": ("gauge", "Used tokens plus pending prefill tokens")} - ) - max_total_num_tokens: int = field( - metadata={"metric": ("gauge", "Maximum token capacity")} - ) + # metadata={"metric": ("gauge", "Used tokens plus pending prefill tokens")} + num_total_tokens: Optional[int] + # metadata={"metric": ("gauge", "Maximum token capacity")} + max_total_num_tokens: Optional[int] # FIXME: token_usage is actually max usage across all pools (KV, SWA, mamba), # not just KV token usage. Rename requires API deprecation. - token_usage: float = field(metadata={"metric": ("gauge", "Token pool usage ratio")}) - gen_throughput: float = field( - metadata={"metric": ("gauge", "Generation throughput tokens/sec")} - ) - cache_hit_rate: float = field( - metadata={"metric": ("gauge", "Prefix cache hit rate")} - ) - utilization: float = field( - metadata={"metric": ("gauge", "Overall utilization ratio")} - ) - max_running_requests: int = field( - metadata={"metric": ("gauge", "Maximum running requests capacity")} - ) + # metadata={"metric": ("gauge", "Token pool usage ratio")} + token_usage: Optional[float] + # metadata={"metric": ("gauge", "Generation throughput tokens/sec")} + gen_throughput: Optional[float] + # metadata={"metric": ("gauge", "Prefix cache hit rate")} + cache_hit_rate: Optional[float] + # metadata={"metric": ("gauge", "Overall utilization ratio")} + utilization: Optional[float] + # metadata={"metric": ("gauge", "Maximum running requests capacity")} + max_running_requests: Optional[int] memory: Optional[MemoryMetrics] = None speculative: Optional[SpeculativeMetrics] = None @@ -2105,39 +2019,32 @@ class GetLoadsReqOutput(BaseReq): queues: Optional[QueueMetrics] = None -@dataclass -class WatchLoadUpdateReq(BaseReq): +class WatchLoadUpdateReq(BaseReqIpc, kw_only=True): loads: List[GetLoadsReqOutput] -@dataclass -class SetInjectDumpMetadataReqInput(BaseReq): +class SetInjectDumpMetadataReqInput(BaseReqIpc, kw_only=True): dump_metadata: Dict[str, Any] -@dataclass -class SetInjectDumpMetadataReqOutput(BaseReq): +class SetInjectDumpMetadataReqOutput(BaseReqIpc, kw_only=True): success: bool -@dataclass -class LazyDumpTensorsReqInput(BaseReq): +class LazyDumpTensorsReqInput(BaseReqIpc, kw_only=True): pass -@dataclass -class LazyDumpTensorsReqOutput(BaseReq): +class LazyDumpTensorsReqOutput(BaseReqIpc, kw_only=True): success: bool -@dataclass -class DumperControlReqInput(BaseReq): +class DumperControlReqInput(BaseReqIpc, kw_only=True): method: str body: Dict[str, Any] -@dataclass -class DumperControlReqOutput(BaseReq): +class DumperControlReqOutput(BaseReqIpc, kw_only=True): success: bool response: List[Dict[str, Any]] error: str = "" @@ -2153,10 +2060,15 @@ def _check_all_req_types(): # check its name name = class_type[0] is_io_struct = ( - name.endswith("Req") or name.endswith("Input") or name.endswith("Output") + name.endswith("Req") + or name.endswith("Input") + or name.endswith("Output") + or name.endswith("Ipc") # allow IPC-specific struct naming ) - is_base_req = issubclass(class_type[1], BaseReq) or issubclass( - class_type[1], BaseBatchReq + is_base_req = ( + issubclass(class_type[1], BaseReq) + or issubclass(class_type[1], BaseBatchReq) + or issubclass(class_type[1], msgspec.Struct) # allow IPC objects ) if is_io_struct and not is_base_req: raise ValueError(f"{name} is not a subclass of BaseReq or BaseBatchReq.") @@ -2167,3 +2079,96 @@ def _check_all_req_types(): _check_all_req_types() + + +# Below are msgpack serialization and ipc utils +class PickleWrapperIpc(msgspec.Struct, tag=True): + data: bytes + + +def enc_hook(obj: Any) -> Any: + if isinstance(obj, torch.Tensor): + # encode torch tensor as Tuple(shape, dtype, data) + tensor_dtype = str(obj.dtype).split(".")[-1] # e.g., "float32" + raw_data = obj.flatten().cpu().contiguous().view(torch.uint8).numpy().data + return (obj.shape, tensor_dtype, raw_data) + elif isinstance(obj, array): + return (obj.typecode, obj.tobytes()) + else: + if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): + logger.info(f"Object of type {type(obj)} is encoding with pickle.") + return pickle.dumps(obj) + + +def dec_hook(tp: Type, obj: Any) -> Any: + if tp is torch.Tensor: + # Convert ``obj`` (which should be a ``tuple``) to a torch.Tensor + shape, dtype, data = obj + tensor_dtype = getattr(torch, dtype) + return torch.frombuffer(data, dtype=tensor_dtype).reshape(shape) + elif tp is array: + typecode, raw_data = obj + res = array(typecode) + res.frombytes(raw_data) + return res + else: + if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): + logger.info(f"Object of type {type(obj)} and {str(tp)} is decoding with pickle.") + return pickle.loads(obj) + +_struct_types = tuple( + cls for cls in BaseReqIpc.__subclasses__() + + BaseBatchReqIpc.__subclasses__() + + SpeculativeDecodingMetricsMixinReq.__subclasses__() + + [PickleWrapperIpc] +) +# str primitive types should encode to bytes +_primitive_types = (int, float, bool, bytes) + +_msgpack_encoder = msgspec.msgpack.Encoder(enc_hook=enc_hook) +_msgpack_decoder = msgspec.msgpack.Decoder( + Union[*_struct_types, *_primitive_types], + dec_hook=dec_hook +) + +def _maybe_wrap_pickle(obj: Any) -> PickleWrapperIpc: + # for msgspec.Struct type, we can directly serialize + if isinstance(obj, msgspec.Struct): + return obj + # for primitive types, we can directly serialize + if type(obj) in _primitive_types: + return obj + + if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): + logger.info(f"Object of type {type(obj)} is wrapped via PickleWrapperIpc.") + return PickleWrapperIpc(pickle.dumps(obj)) + +def _maybe_unwrap_pickle(obj: Any) -> Any: + # only PickleWrapperIpc type needs to be unwrapped, + # other types can be directly used + if isinstance(obj, PickleWrapperIpc): + obj = pickle.loads(obj.data) + if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): + logger.info(f"Object of type {type(obj)} is unwrapped from PickleWrapperIpc.") + return obj + + return obj + +def sock_send(socket: Socket, obj: Any, flags=0): + obj = _maybe_wrap_pickle(obj) + socket.send(_msgpack_encoder.encode(obj), flags=flags) + + +def sock_recv(socket: Socket, flags=0) -> Any: + data = socket.recv(flags=flags) + return _maybe_unwrap_pickle(_msgpack_decoder.decode(data)) + + +async def async_sock_send(socket: AsyncSocket, obj: Any, flags=0): + obj = _maybe_wrap_pickle(obj) + await socket.send(_msgpack_encoder.encode(obj), flags=flags) + + +async def async_sock_recv(socket: AsyncSocket, flags=0) -> Any: + data = await socket.recv(flags=flags) + return _maybe_unwrap_pickle(_msgpack_decoder.decode(data)) diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index 5f6cc63091c8..ccab34b93d1f 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -41,14 +41,20 @@ from sglang.srt.managers.io_struct import ( BaseBatchReq, BaseReq, + BaseReqIpc, + BaseBatchReqIpc, BatchEmbeddingOutput, BatchStrOutput, BatchTokenIDOutput, ContinueGenerationReqInput, FreezeGCReq, - PauseContinueBroadcast, + PauseContinueBroadcastReq, PauseGenerationReqInput, - TokenizerWorkerRegistration, + TokenizerWorkerRegistrationReq, + sock_recv, + sock_send, + async_sock_recv, + async_sock_send, ) from sglang.srt.managers.tokenizer_manager import TokenizerManager from sglang.srt.server_args import PortArgs, ServerArgs @@ -93,7 +99,7 @@ def send_output(self, ipc_name: str, output: Any, is_tokenizer: bool = False): if ipc_name not in self._mapping: self._register_ipc_mapping(ipc_name, is_tokenizer=is_tokenizer) - self._mapping[ipc_name].send_pyobj(output) + sock_send(self._mapping[ipc_name], output) def _extract_field_by_index( @@ -324,7 +330,7 @@ def multi_http_worker_event_loop(self: DetokenizerManager): """The event loop that handles requests, for multi multi-http-worker mode""" self.socket_mapping = SocketMapping() while True: - recv_obj = self.recv_from_scheduler.recv_pyobj() + recv_obj = sock_recv(self.recv_from_scheduler) output = self._request_dispatcher(recv_obj) if output is None: continue @@ -332,13 +338,13 @@ def multi_http_worker_event_loop(self: DetokenizerManager): # Fan out the output back to the originating tokenizer worker(s). # In multi-detokenizer mode the upstream MultiDetokenizerRouter may # forward either batched or single requests, so handle both shapes. - if isinstance(recv_obj, BaseBatchReq): + if isinstance(recv_obj, BaseBatchReqIpc): for i, ipc_name in enumerate(recv_obj.http_worker_ipcs): new_output = _handle_output_by_index(output, i) self.socket_mapping.send_output( ipc_name, new_output, is_tokenizer=True ) - elif isinstance(recv_obj, BaseReq): + elif isinstance(recv_obj, BaseReqIpc): self.socket_mapping.send_output( recv_obj.http_worker_ipc, output, is_tokenizer=True ) @@ -394,9 +400,9 @@ def _run_loop(self): async def router_worker_obj(self): """Forward path: workers → scheduler, with pause/continue broadcast.""" while True: - recv_obj = await self.receive_from_worker.recv_pyobj() + recv_obj = await async_sock_recv(self.receive_from_worker) - if isinstance(recv_obj, TokenizerWorkerRegistration): + if isinstance(recv_obj, TokenizerWorkerRegistrationReq): if recv_obj.worker_ipc_name not in self.all_worker_ipcs: self.all_worker_ipcs.add(recv_obj.worker_ipc_name) logger.info( @@ -410,7 +416,7 @@ async def router_worker_obj(self): ): # Broadcast to ALL workers so every worker's is_pause is set is_pause = isinstance(recv_obj, PauseGenerationReqInput) - broadcast = PauseContinueBroadcast(is_pause=is_pause) + broadcast = PauseContinueBroadcastReq(is_pause=is_pause) for ipc_name in self.all_worker_ipcs: self.socket_mapping.send_output(ipc_name, broadcast) # Forward to scheduler rank 0 (it broadcasts to all TP/PP/DP @@ -419,21 +425,21 @@ async def router_worker_obj(self): isinstance(recv_obj, PauseGenerationReqInput) and recv_obj.mode == "abort" ): - await self.send_to_scheduler.send_pyobj(recv_obj) + await async_sock_send(self.send_to_scheduler, recv_obj) continue - await self.send_to_scheduler.send_pyobj(recv_obj) + await async_sock_send(self.send_to_scheduler, recv_obj) async def handle_loop(self): """Backward path: detokenizer → route results to correct worker.""" while True: - recv_obj = await self.recv_from_detokenizer.recv_pyobj() + recv_obj = await async_sock_recv(self.recv_from_detokenizer) await self._distribute_result_to_workers(recv_obj) async def _distribute_result_to_workers(self, recv_obj): - if isinstance(recv_obj, BaseReq): + if isinstance(recv_obj, BaseReqIpc): ipc_names = [recv_obj.http_worker_ipc] - elif isinstance(recv_obj, BaseBatchReq): + elif isinstance(recv_obj, BaseBatchReqIpc): ipc_names = recv_obj.http_worker_ipcs else: raise ValueError(f"Unknown recv_obj type: {type(recv_obj)}") @@ -468,7 +474,7 @@ def _send(self, ipc_name: str, obj: Any) -> None: def event_loop(self): while True: - recv_obj = self.recv_from_scheduler.recv_pyobj() + recv_obj = sock_recv(self.recv_from_scheduler) # FreezeGCReq must freeze every detokenizer process. if isinstance(recv_obj, FreezeGCReq): @@ -477,7 +483,7 @@ def event_loop(self): continue # Single request: route by its own http_worker_ipc. - if isinstance(recv_obj, BaseReq): + if isinstance(recv_obj, BaseReqIpc): assert ( recv_obj.http_worker_ipc is not None ), f"Single req {recv_obj.rid=} missing http_worker_ipc" @@ -485,7 +491,7 @@ def event_loop(self): continue # Batch request. - if isinstance(recv_obj, BaseBatchReq): + if isinstance(recv_obj, BaseBatchReqIpc): # Idle/no-op batch (rids=[]): broadcast to all detokenizers if not recv_obj.rids: for ipc in self.ipc_name_list: @@ -562,18 +568,18 @@ def __init__( ) # Register this worker with the router for pause/continue broadcasting - reg = TokenizerWorkerRegistration(worker_ipc_name=self.tokenizer_ipc_name) - self.send_to_scheduler.send_pyobj(reg) + reg = TokenizerWorkerRegistrationReq(worker_ipc_name=self.tokenizer_ipc_name) + self.send_to_scheduler.send_obj(reg) # Future for awaiting pause/continue broadcast confirmation self._pause_continue_future: Optional[asyncio.Future] = None - # Register PauseContinueBroadcast in the result dispatcher so + # Register PauseContinueBroadcastReq in the result dispatcher so # handle_loop routes it to _handle_pause_continue_broadcast from sglang.utils import TypeBasedDispatcher self._result_dispatcher += TypeBasedDispatcher( - [(PauseContinueBroadcast, self._handle_pause_continue_broadcast)] + [(PauseContinueBroadcastReq, self._handle_pause_continue_broadcast)] ) async def pause_generation(self, obj: PauseGenerationReqInput): @@ -581,7 +587,8 @@ async def pause_generation(self, obj: PauseGenerationReqInput): self._pause_continue_future = loop.create_future() # Send to router which will broadcast to all workers # (router also handles forwarding to scheduler for non-abort modes) - self.send_to_scheduler.send_pyobj(obj) + await self.send_to_scheduler.async_send_obj(obj) + await self._pause_continue_future if obj.mode == "abort": @@ -596,15 +603,15 @@ async def pause_generation(self, obj: PauseGenerationReqInput): async def continue_generation(self, obj: ContinueGenerationReqInput): loop = asyncio.get_event_loop() self._pause_continue_future = loop.create_future() - self.send_to_scheduler.send_pyobj(obj) + await self.send_to_scheduler.async_send_obj(obj) await self._pause_continue_future - def _handle_pause_continue_broadcast(self, obj: PauseContinueBroadcast): + def _handle_pause_continue_broadcast(self, obj: PauseContinueBroadcastReq): """Called from handle_loop when a broadcast arrives from the router.""" loop = asyncio.get_event_loop() loop.create_task(self._apply_pause_continue_broadcast(obj)) - async def _apply_pause_continue_broadcast(self, obj: PauseContinueBroadcast): + async def _apply_pause_continue_broadcast(self, obj: PauseContinueBroadcastReq): """Apply pause/continue state under the condition lock.""" async with self.is_pause_cond: if obj.is_pause: @@ -707,11 +714,22 @@ def write_data_for_multi_tokenizer( class SenderWrapper: - def __init__(self, port_args: PortArgs, send_to_scheduler: zmq.Socket): + def __init__( + self, + port_args: PortArgs, + send_to_scheduler: zmq.Socket, + is_multi_tokenizer: bool = False, + ): self.port_args = port_args self.send_to_scheduler = send_to_scheduler + self.is_multi_tokenizer = is_multi_tokenizer + + def send_obj(self, obj): + if not self.is_multi_tokenizer and hasattr(obj, "http_worker_ipc"): + obj.http_worker_ipc = self.port_args.tokenizer_ipc_name + sock_send(self.send_to_scheduler, obj) - def send_pyobj(self, obj): - if isinstance(obj, BaseReq): + async def async_send_obj(self, obj): + if not self.is_multi_tokenizer and hasattr(obj, "http_worker_ipc"): obj.http_worker_ipc = self.port_args.tokenizer_ipc_name - self.send_to_scheduler.send_pyobj(obj) + await async_sock_send(self.send_to_scheduler, obj) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a21ebc94b75e..606d5b3ee623 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -145,6 +145,7 @@ UpdateWeightsFromDistributedReqInput, UpdateWeightsFromIPCReqInput, UpdateWeightsFromTensorReqInput, + sock_send, ) from sglang.srt.managers.load_snapshot import LoadSnapshot, create_load_snapshot_writer from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors @@ -1555,7 +1556,7 @@ def process_input_requests(self, recv_reqs: List): self.ipc_channels.send_to_tokenizer.send_output(output, recv_req) else: if self.ipc_channels.recv_from_rpc is not None: - self.ipc_channels.recv_from_rpc.send_pyobj(output) + sock_send(self.ipc_channels.recv_from_rpc, output, recv_req) self.flush_wrapper.check_pending() if self.external_corpus_manager is not None: 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 85531e49e28a..9c1a6631197e 100644 --- a/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py +++ b/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py @@ -16,6 +16,7 @@ EventPublisherFactory, KVEventBatch, ) +from sglang.srt.managers.io_struct import sock_send if TYPE_CHECKING: from sglang.srt.distributed.parallel_state_wrapper import ParallelState @@ -85,7 +86,7 @@ def emit_kv_metrics(self): ) if not self.send_metrics_from_scheduler.closed: - self.send_metrics_from_scheduler.send_pyobj(kv_metrics) + sock_send(self.send_metrics_from_scheduler, kv_metrics) def publish_kv_events(self): if not self.enable_kv_cache_events: diff --git a/python/sglang/srt/managers/scheduler_components/output_sender.py b/python/sglang/srt/managers/scheduler_components/output_sender.py index a8eec2fd1283..e6e7cf74737c 100644 --- a/python/sglang/srt/managers/scheduler_components/output_sender.py +++ b/python/sglang/srt/managers/scheduler_components/output_sender.py @@ -2,7 +2,11 @@ import zmq -from sglang.srt.managers.io_struct import BaseBatchReq, BaseReq +from sglang.srt.managers.io_struct import( + BaseBatchReqIpc, + BaseReqIpc, + sock_send +) class SenderWrapper: @@ -11,18 +15,18 @@ def __init__(self, socket: zmq.Socket): def send_output( self, - output: Union[BaseReq, BaseBatchReq], - recv_obj: Optional[Union[BaseReq, BaseBatchReq]] = None, + output: Union[BaseReqIpc, BaseBatchReqIpc], + recv_obj: Optional[Union[BaseReqIpc, BaseBatchReqIpc]] = None, ): if self.socket is None: return if ( - isinstance(recv_obj, BaseReq) + isinstance(recv_obj, BaseReqIpc) and recv_obj.http_worker_ipc is not None and output.http_worker_ipc is None ): # handle communicator reqs for multi-http worker case output.http_worker_ipc = recv_obj.http_worker_ipc - self.socket.send_pyobj(output) + sock_send(self.socket, output) diff --git a/python/sglang/srt/managers/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py index 6e5a58ffaada..05a8b50e03ba 100644 --- a/python/sglang/srt/managers/scheduler_components/profiler_manager.py +++ b/python/sglang/srt/managers/scheduler_components/profiler_manager.py @@ -404,7 +404,7 @@ def _profile_batch_predicate(self, batch: ScheduleBatch): self._start_profile() def _profile(self, recv_req: ProfileReq): - if recv_req.type == ProfileReqType.START_PROFILE: + if recv_req.req_type == ProfileReqType.START_PROFILE: if recv_req.profile_by_stage or recv_req.start_step: return self._init_profile( recv_req.output_dir, diff --git a/python/sglang/srt/managers/scheduler_components/request_receiver.py b/python/sglang/srt/managers/scheduler_components/request_receiver.py index 0c5b83b6701f..16e11eb435cf 100644 --- a/python/sglang/srt/managers/scheduler_components/request_receiver.py +++ b/python/sglang/srt/managers/scheduler_components/request_receiver.py @@ -20,6 +20,7 @@ BatchTokenizedGenerateReqInput, TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, + sock_recv, ) from sglang.srt.managers.mm_utils import ( has_shm_features, @@ -93,7 +94,7 @@ def _pull_raw_reqs(self) -> Optional[List]: try: if self.recv_limit_reached(len(recv_reqs)): break - recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK) + recv_req = sock_recv(self.recv_from_tokenizer, zmq.NOBLOCK) except zmq.ZMQError: break recv_reqs.append(recv_req) @@ -102,7 +103,7 @@ def _pull_raw_reqs(self) -> Optional[List]: try: if self.recv_limit_reached(len(recv_reqs)): break - recv_rpc = self.recv_from_rpc.recv_pyobj(zmq.NOBLOCK) + recv_rpc = sock_recv(self.recv_from_rpc, zmq.NOBLOCK) except zmq.ZMQError: break recv_reqs.append(recv_rpc) diff --git a/python/sglang/srt/managers/scheduler_input_blocker.py b/python/sglang/srt/managers/scheduler_input_blocker.py index b6838ae43180..378c001004d0 100644 --- a/python/sglang/srt/managers/scheduler_input_blocker.py +++ b/python/sglang/srt/managers/scheduler_input_blocker.py @@ -51,10 +51,10 @@ def handle(self, recv_reqs: Optional[List[Any]]): def _handle_recv_req(self, recv_req): if isinstance(recv_req, BlockReqInput): - if recv_req.type == BlockReqType.BLOCK: + if recv_req.req_type == BlockReqType.BLOCK: self._execute_block_req() return [] - elif recv_req.type == BlockReqType.UNBLOCK: + elif recv_req.req_type == BlockReqType.UNBLOCK: self._execute_unblock_req() return [] else: @@ -99,8 +99,8 @@ class _State(Enum): @contextmanager def input_blocker_guard_region(send_to_scheduler): - send_to_scheduler.send_pyobj(BlockReqInput(BlockReqType.BLOCK)) + send_to_scheduler.send_obj(BlockReqInput(BlockReqType.BLOCK)) try: yield finally: - send_to_scheduler.send_pyobj(BlockReqInput(BlockReqType.UNBLOCK)) + send_to_scheduler.send_obj(BlockReqInput(BlockReqType.UNBLOCK)) diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index c9939a1fc93a..c8700a120fdd 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -336,7 +336,7 @@ async def start_profile( ) record_shapes = (record_shapes is not False) and env_record_shapes req = ProfileReq( - type=ProfileReqType.START_PROFILE, + req_type=ProfileReqType.START_PROFILE, output_dir=output_dir, start_step=start_step, num_steps=num_steps, @@ -353,7 +353,7 @@ async def start_profile( async def stop_profile(self: TokenizerManager): self.auto_create_handle_loop() - req = ProfileReq(type=ProfileReqType.STOP_PROFILE) + req = ProfileReq(ProfileReqType.STOP_PROFILE) return await self._execute_profile(req) async def _execute_profile(self: TokenizerManager, req: ProfileReq): @@ -857,7 +857,7 @@ async def open_session( future = asyncio.Future() self.session_futures[obj.session_id] = future - self.send_to_scheduler.send_pyobj(obj) + await self.send_to_scheduler.async_send_obj(obj) try: return await future @@ -869,7 +869,7 @@ async def close_session( obj: CloseSessionReqInput, request: Optional[fastapi.Request] = None, ): - await self.send_to_scheduler.send_pyobj(obj) + await self.send_to_scheduler.async_send_obj(obj) def _update_weight_version_if_provided( self: TokenizerManager, weight_version: Optional[str] diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 2168f053cd0d..218560ba0b6c 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -73,6 +73,7 @@ TokenizedGenerateReqInput, UpdateWeightFromDiskReqInput, UpdateWeightFromDiskReqOutput, + async_sock_recv, ) from sglang.srt.managers.load_snapshot import create_load_snapshot_reader from sglang.srt.managers.mm_utils import TensorTransportMode, wrap_shm_features @@ -357,24 +358,24 @@ def init_tokenizer_and_processor(self): self.async_dynamic_batch_tokenizer = None def init_ipc_channels(self, port_args: PortArgs): + from sglang.srt.managers.multi_tokenizer_mixin import SenderWrapper + context = zmq.asyncio.Context(2) self.recv_from_detokenizer = get_zmq_socket( context, zmq.PULL, port_args.tokenizer_ipc_name, True ) if self.server_args.tokenizer_worker_num == 1: - self.send_to_scheduler = get_zmq_socket( + send_to_scheduler = get_zmq_socket( context, zmq.PUSH, port_args.scheduler_input_ipc_name, True ) + self.send_to_scheduler = SenderWrapper(port_args, send_to_scheduler) else: - from sglang.srt.managers.multi_tokenizer_mixin import SenderWrapper - # Use tokenizer_worker_ipc_name in multi-tokenizer mode send_to_scheduler = get_zmq_socket( context, zmq.PUSH, port_args.tokenizer_worker_ipc_name, False ) - # Make sure that each request carries the tokenizer_ipc_name for response routing - self.send_to_scheduler = SenderWrapper(port_args, send_to_scheduler) + self.send_to_scheduler = SenderWrapper(port_args, send_to_scheduler, True) self.load_snapshot_reader = create_load_snapshot_reader( self.server_args, @@ -1069,15 +1070,15 @@ def _create_tokenized_object( self.fake_bootstrap_room_counter += 1 tokenized_obj = TokenizedGenerateReqInput( - input_text, - input_ids_arr, - mm_inputs, - sampling_params, - obj.return_logprob, - obj.logprob_start_len, - obj.top_logprobs_num, - obj.token_ids_logprob, - obj.stream, + input_text=input_text, + input_ids=input_ids_arr, + mm_inputs=mm_inputs, + sampling_params=sampling_params, + return_logprob=obj.return_logprob, + logprob_start_len=obj.logprob_start_len, + top_logprobs_num=obj.top_logprobs_num, + token_ids_logprob=obj.token_ids_logprob, + stream=obj.stream, rid=obj.rid, http_worker_ipc=obj.http_worker_ipc, bootstrap_host=obj.bootstrap_host, @@ -1117,11 +1118,11 @@ def _create_tokenized_object( ) tokenized_obj = TokenizedEmbeddingReqInput( - input_text, - input_ids_arr, - mm_inputs, - token_type_ids, - sampling_params, + input_text=input_text, + input_ids=input_ids_arr, + image_inputs=mm_inputs, + token_type_ids=token_type_ids, + sampling_params=sampling_params, positional_embed_overrides=positional_embed_overrides, rid=obj.rid, priority=obj.priority, @@ -1253,7 +1254,7 @@ def _send_one_request( ): tokenized_obj.time_stats.set_api_server_dispatch_time() tokenized_obj = wrap_shm_features(tokenized_obj) - self.send_to_scheduler.send_pyobj(tokenized_obj) + self.send_to_scheduler.send_obj(tokenized_obj) tokenized_obj.time_stats.set_api_server_dispatch_finish_time() def _send_batch_request( @@ -1269,7 +1270,7 @@ def _send_batch_request( batch_req = BatchTokenizedEmbeddingReqInput(batch=tokenized_objs) set_time_batch(tokenized_objs, "set_api_server_dispatch_time") - self.send_to_scheduler.send_pyobj(batch_req) + self.send_to_scheduler.send_obj(batch_req) set_time_batch(tokenized_objs, "set_api_server_dispatch_finish_time") def _coalesce_streaming_chunks( @@ -1577,7 +1578,7 @@ def abort_request(self, rid: str = "", abort_all: bool = False): ): return req = AbortReq(rid=rid, abort_all=abort_all) - self.send_to_scheduler.send_pyobj(req) + self.send_to_scheduler.send_obj(req) if self.enable_metrics: # TODO: also use custom_labels from the request self.metrics_collector.observe_one_aborted_request( @@ -1588,7 +1589,7 @@ async def pause_generation(self, obj: PauseGenerationReqInput): async with self.is_pause_cond: self.is_pause = True if obj.mode != "abort": - await self.send_to_scheduler.send_pyobj(obj) + await self.send_to_scheduler.async_send_obj(obj) else: # we are using the model_update_lock to check if there is still on-going requests. while True: @@ -1602,7 +1603,7 @@ async def pause_generation(self, obj: PauseGenerationReqInput): async def continue_generation(self, obj: ContinueGenerationReqInput): async with self.is_pause_cond: self.is_pause = False - await self.send_to_scheduler.send_pyobj(obj) + await self.send_to_scheduler.async_send_obj(obj) self.is_pause_cond.notify_all() async def update_weights_from_disk( @@ -1647,7 +1648,7 @@ def _update_model_path_info(self, model_path: str, load_format: str): async def _wait_for_model_update_from_disk( self, obj: UpdateWeightFromDiskReqInput ) -> Tuple[bool, str]: - self.send_to_scheduler.send_pyobj(obj) + await self.send_to_scheduler.async_send_obj(obj) self.model_update_result = asyncio.Future() if self.server_args.dp_size == 1: result = await self.model_update_result @@ -1687,12 +1688,12 @@ def configure_logging(self, obj: ConfigureLoggingReq): # Let the exception propagate to the caller. # Only legal requests will be sent to scheduler. logging.getLogger().setLevel(obj.log_level.upper()) - self.send_to_scheduler.send_pyobj(obj) + self.send_to_scheduler.send_obj(obj) logging.info(f"Config logging: {obj=}") async def freeze_gc(self): """Send a freeze_gc message to the scheduler first, then freeze locally.""" - self.send_to_scheduler.send_pyobj(FreezeGCReq()) + await self.send_to_scheduler.async_send_obj(FreezeGCReq()) freeze_gc("Tokenizer Manager") return None @@ -1739,7 +1740,7 @@ async def handle_loop(self): """The event loop that handles requests""" while True: with self.soft_watchdog.disable(): - recv_obj = await self.recv_from_detokenizer.recv_pyobj() + recv_obj = await async_sock_recv(self.recv_from_detokenizer) if isinstance( recv_obj, (BatchStrOutput, BatchEmbeddingOutput, BatchTokenIDOutput), @@ -2567,7 +2568,7 @@ def _handle_abort_req(self, recv_obj: AbortReq): state.event.set() def update_active_ranks(self, ranks: ActiveRanksOutput): - self.send_to_scheduler.send_pyobj(ranks) + self.send_to_scheduler.send_obj(ranks) def _handle_open_session_req_output(self, recv_obj): future = self.session_futures.get(recv_obj.session_id) diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index 326aace02765..7a4c448194fb 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/python/sglang/srt/observability/req_time_stats.py @@ -220,10 +220,12 @@ class RequestStage: @dataclass class ReqTimeStatsBase: enable_metrics: bool = False - metrics_collector: Optional[ - Union[SchedulerMetricsCollector, TokenizerMetricsCollector] - ] = None - trace_ctx: Union[TraceReqContext, TraceNullContext] = field( + # use object to avoid msgpack decode build issue, not use Any due to the + # custom object serialized by pickle. + # Optional[Union[SchedulerMetricsCollector, TokenizerMetricsCollector]] + metrics_collector: Optional[object] = None + # Union[TraceReqContext, TraceNullContext] + trace_ctx: object = field( default_factory=TraceNullContext ) disagg_mode: DisaggregationMode = DisaggregationMode.NULL diff --git a/python/sglang/srt/observability/trace.py b/python/sglang/srt/observability/trace.py index a7cc4cfddfa0..b57d394ef0f0 100644 --- a/python/sglang/srt/observability/trace.py +++ b/python/sglang/srt/observability/trace.py @@ -684,11 +684,20 @@ def __del__(self): self.abort(abort_info={"reason": "have unclosed span, auto closed"}) -@dataclass +# Temperary remove the @dataclass decorator to make msgpack treat +# TraceReqContext as a custom class and serialize with pickle like TraceReqContext. +# When the TraceReqContext fixed with msgpack, we can consider to add @dataclass back. class TraceNullContext: tracing_enable: bool = False + def __init__(self, tracing_enable=False): + self.tracing_enable = tracing_enable + def __getattr__(self, name): + # fix the .__dataclass_fields__ is not a dict error + if name.startswith("__") and name.endswith("__"): + raise AttributeError(name) + return self def __call__(self, *args, **kwargs): diff --git a/test/registered/unit/utils/test_profile_merger.py b/test/registered/unit/utils/test_profile_merger.py index eb4cb9bcb906..63fa90778246 100644 --- a/test/registered/unit/utils/test_profile_merger.py +++ b/test/registered/unit/utils/test_profile_merger.py @@ -216,10 +216,10 @@ def test_data_structures_merge_profiles(self): self.assertTrue(req_input.merge_profiles) # Test ProfileReq - req = ProfileReq(type=ProfileReqType.START_PROFILE) + req = ProfileReq(ProfileReqType.START_PROFILE) self.assertFalse(req.merge_profiles) - req = ProfileReq(type=ProfileReqType.START_PROFILE, merge_profiles=True) + req = ProfileReq(ProfileReqType.START_PROFILE, merge_profiles=True) self.assertTrue(req.merge_profiles) def test_integration_parameters(self): From b4b22d1571a94994545205676d8dca0040b83a18 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Wed, 3 Jun 2026 17:32:55 -0700 Subject: [PATCH 02/45] lint --- python/sglang/srt/debug_utils/dumper.py | 19 ++++-- .../srt/disaggregation/encode_grpc_server.py | 1 + .../srt/disaggregation/encode_server.py | 4 +- .../srt/elastic_ep/expert_backup_client.py | 2 +- python/sglang/srt/entrypoints/http_server.py | 64 ++++++++++++++++--- python/sglang/srt/managers/communicator.py | 3 +- .../srt/managers/data_parallel_controller.py | 2 +- python/sglang/srt/managers/embed_types.py | 2 +- python/sglang/srt/managers/io_struct.py | 63 ++++++++++-------- .../srt/managers/multi_tokenizer_mixin.py | 8 +-- .../scheduler_components/output_sender.py | 6 +- .../srt/observability/req_time_stats.py | 4 +- python/sglang/srt/observability/trace.py | 2 +- 13 files changed, 118 insertions(+), 62 deletions(-) diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index edb86aea4b2f..cd0445d100db 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -1391,7 +1391,8 @@ def _create_zmq_rpc_broadcast( ) -> Optional["_ZmqRpcBroadcast"]: """A general-purpose minimal RPC to support broadcasting executions to multi processes""" import zmq - from sglang.srt.managers.io_struct import (sock_send, sock_recv) + + from sglang.srt.managers.io_struct import sock_recv, sock_send rank = _get_rank() world_size = dist.get_world_size() if dist.is_initialized() else 1 @@ -1448,12 +1449,16 @@ def __init__(self, socket, debug_name: str): def __getattr__(self, method_name: str): def call(*args, **kwargs): - from sglang.srt.managers.io_struct import (sock_send, sock_recv) - sock_send(self._socket, { - "method": method_name, - "args": args, - "kwargs": kwargs, - }) + from sglang.srt.managers.io_struct import sock_recv, sock_send + + sock_send( + self._socket, + { + "method": method_name, + "args": args, + "kwargs": kwargs, + }, + ) response = sock_recv(self._socket) if response["error"]: raise RuntimeError( diff --git a/python/sglang/srt/disaggregation/encode_grpc_server.py b/python/sglang/srt/disaggregation/encode_grpc_server.py index 0207fbf6e06c..2e43ad80eff4 100644 --- a/python/sglang/srt/disaggregation/encode_grpc_server.py +++ b/python/sglang/srt/disaggregation/encode_grpc_server.py @@ -88,6 +88,7 @@ async def Encode( self, request: sglang_encoder_pb2.EncodeRequest, context ) -> sglang_encoder_pb2.EncodeResponse: from sglang.srt.managers.io_struct import async_sock_send + try: request_dict = { "mm_items": list(request.mm_items), diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index 954b9d22ad36..2cbf5c1fe6f1 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -44,8 +44,8 @@ ProfileReq, ProfileReqInput, ProfileReqType, - sock_send, async_sock_recv, + sock_send, ) from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache @@ -2316,7 +2316,7 @@ async def _dispatch_group( "modality": modality.name, "requests": requests, "enter_time": start, - } + }, ) logger.info(f"Dispatching batch of {len(group)} {modality.name} requests") diff --git a/python/sglang/srt/elastic_ep/expert_backup_client.py b/python/sglang/srt/elastic_ep/expert_backup_client.py index 6b114a699fe4..6b03188d4a98 100644 --- a/python/sglang/srt/elastic_ep/expert_backup_client.py +++ b/python/sglang/srt/elastic_ep/expert_backup_client.py @@ -12,7 +12,7 @@ ) from sglang.srt.environ import envs from sglang.srt.eplb.expert_location import get_global_expert_location_metadata -from sglang.srt.managers.io_struct import UpdateExpertBackupReq, sock_send, sock_recv +from sglang.srt.managers.io_struct import UpdateExpertBackupReq, sock_recv, sock_send from sglang.srt.server_args import ServerArgs from sglang.srt.utils.network import get_local_ip_auto diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 07f7eca2da89..e319130f0b1c 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -33,9 +33,9 @@ Callable, Dict, List, + Literal, Optional, Union, - Literal, ) # Fix a bug of Python threading @@ -111,33 +111,77 @@ from sglang.srt.function_call.function_call_parser import FunctionCallParser from sglang.srt.managers.io_struct import ( CheckWeightsReqInput as CheckWeightsReqInputIpc, +) +from sglang.srt.managers.io_struct import ( CloseSessionReqInput as CloseSessionReqInputIpc, - ConfigureLoggingReq as ConfigureLoggingReqIpc, +) +from sglang.srt.managers.io_struct import ConfigureLoggingReq as ConfigureLoggingReqIpc +from sglang.srt.managers.io_struct import ( ContinueGenerationReqInput as ContinueGenerationReqInputIpc, +) +from sglang.srt.managers.io_struct import ( DestroyWeightsUpdateGroupReqInput as DestroyWeightsUpdateGroupReqInputIpc, +) +from sglang.srt.managers.io_struct import ( DumperControlReqInput, EmbeddingReqInput, GenerateReqInput, +) +from sglang.srt.managers.io_struct import ( GetWeightsByNameReqInput as GetWeightsByNameReqInputIpc, +) +from sglang.srt.managers.io_struct import ( InitWeightsSendGroupForRemoteInstanceReqInput as InitWeightsSendGroupForRemoteInstanceReqInputIpc, +) +from sglang.srt.managers.io_struct import ( InitWeightsUpdateGroupReqInput as InitWeightsUpdateGroupReqInputIpc, +) +from sglang.srt.managers.io_struct import ( LoadLoRAAdapterFromTensorsReqInput as LoadLoRAAdapterFromTensorsReqInputIpc, +) +from sglang.srt.managers.io_struct import ( LoadLoRAAdapterReqInput as LoadLoRAAdapterReqInputIpc, - OpenSessionReqInput as OpenSessionReqInputIpc, +) +from sglang.srt.managers.io_struct import OpenSessionReqInput as OpenSessionReqInputIpc +from sglang.srt.managers.io_struct import ( ParseFunctionCallReq, +) +from sglang.srt.managers.io_struct import ( PauseGenerationReqInput as PauseGenerationReqInputIpc, +) +from sglang.srt.managers.io_struct import ( ProfileReqInput, +) +from sglang.srt.managers.io_struct import ( ReleaseMemoryOccupationReqInput as ReleaseMemoryOccupationReqInputIpc, +) +from sglang.srt.managers.io_struct import ( ResumeMemoryOccupationReqInput as ResumeMemoryOccupationReqInputIpc, +) +from sglang.srt.managers.io_struct import ( SendWeightsToRemoteInstanceReqInput as SendWeightsToRemoteInstanceReqInputIpc, +) +from sglang.srt.managers.io_struct import ( SeparateReasoningReqInput, - SetInternalStateReq as SetInternalStateReqIpc, - SlowDownReqInput as SlowDownReqInputIpc, +) +from sglang.srt.managers.io_struct import SetInternalStateReq as SetInternalStateReqIpc +from sglang.srt.managers.io_struct import SlowDownReqInput as SlowDownReqInputIpc +from sglang.srt.managers.io_struct import ( UnloadLoRAAdapterReqInput as UnloadLoRAAdapterReqInputIpc, +) +from sglang.srt.managers.io_struct import ( UpdateWeightFromDiskReqInput as UpdateWeightFromDiskReqInputIpc, +) +from sglang.srt.managers.io_struct import ( UpdateWeightsFromDistributedReqInput as UpdateWeightsFromDistributedReqInputIpc, +) +from sglang.srt.managers.io_struct import ( UpdateWeightsFromIPCReqInput as UpdateWeightsFromIPCReqInputIpc, +) +from sglang.srt.managers.io_struct import ( UpdateWeightsFromTensorReqInput as UpdateWeightsFromTensorReqInputIpc, +) +from sglang.srt.managers.io_struct import ( UpdateWeightVersionReqInput, VertexGenerateReqInput, ) @@ -1327,8 +1371,7 @@ async def update_weights_from_tensor( """ success, message = await _global_state.tokenizer_manager.update_weights_from_tensor( - UpdateWeightsFromTensorReqInputIpc(**obj.__dict__), - request + UpdateWeightsFromTensorReqInputIpc(**obj.__dict__), request ) content = {"success": success, "message": message} @@ -1504,6 +1547,7 @@ async def resume_memory_occupation( class CheckWeightsReqInput: action: str = "checksum" + @app.api_route("/weights_checker", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) async def check_weights( @@ -1643,6 +1687,7 @@ class OpenSessionReqInput: streaming: Optional[bool] = None timeout: Optional[float] = None + @app.api_route("/open_session", methods=["GET", "POST"]) async def open_session(obj: OpenSessionReqInput, request: Request): """Open a session, and return its unique session id.""" @@ -1664,6 +1709,7 @@ async def open_session(obj: OpenSessionReqInput, request: Request): class CloseSessionReqInput: session_id: str + @app.api_route("/close_session", methods=["GET", "POST"]) async def close_session(obj: CloseSessionReqInput, request: Request): """Close the session.""" @@ -1698,12 +1744,14 @@ async def configure_logging(obj: ConfigureLoggingReq, request: Request): ) return Response(status_code=200) + @dataclasses.dataclass class AbortReq: rid: Optional[str] = None # Whether to abort all requests abort_all: bool = False + @app.post("/abort_request") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def abort_request(obj: AbortReq, request: Request): @@ -1771,8 +1819,6 @@ def __post_init__(self): ) - - @app.post("/pause_generation") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def pause_generation(obj: PauseGenerationReqInput, request: Request): diff --git a/python/sglang/srt/managers/communicator.py b/python/sglang/srt/managers/communicator.py index 015f0098e2a6..75ccfc6a64f2 100644 --- a/python/sglang/srt/managers/communicator.py +++ b/python/sglang/srt/managers/communicator.py @@ -4,10 +4,11 @@ import copy from collections import deque from typing import Deque, Generic, List, Optional, TypeVar -from sglang.srt.managers.io_struct import sock_send import zmq +from sglang.srt.managers.io_struct import sock_send + T = TypeVar("T") diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index 7ebcae8d4ace..1bf5cab764fc 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -36,8 +36,8 @@ ProfileReq, TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, - sock_send, sock_recv, + sock_send, ) from sglang.srt.managers.load_snapshot import create_load_snapshot_reader from sglang.srt.managers.schedule_batch import Req diff --git a/python/sglang/srt/managers/embed_types.py b/python/sglang/srt/managers/embed_types.py index 6cfe170d8139..1d38e8e76733 100644 --- a/python/sglang/srt/managers/embed_types.py +++ b/python/sglang/srt/managers/embed_types.py @@ -19,7 +19,7 @@ """ from dataclasses import dataclass -from typing import List, Union +from typing import List import torch diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 15f2bbac43fd..a47825c30ac8 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -19,31 +19,30 @@ from __future__ import annotations import copy +import logging +import pickle import uuid from abc import ABC from array import array from collections import Counter from dataclasses import dataclass, field from enum import Enum -import logging -import msgspec -import pickle -from typing import TYPE_CHECKING, Annotated, Any, Dict, List, Optional, Union, Type +from typing import TYPE_CHECKING, Annotated, Any, Dict, List, Optional, Type, Union +import msgspec import torch from pydantic import PlainValidator - from zmq import Socket from zmq.asyncio import Socket as AsyncSocket 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 BaseFinishReason, Modality +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 ( - SchedulerReqTimeStats, ReqTimeStatsBase, + SchedulerReqTimeStats, ) from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.utils import ImageData, VideoData @@ -57,6 +56,7 @@ logger = logging.getLogger(__name__) + @dataclass class BaseReq(ABC): rid: Optional[Union[str, List[str]]] = field(default=None, kw_only=True) @@ -745,9 +745,9 @@ def __getitem__(self, i): class TokenizedGenerateReqInput(BaseReqIpc, kw_only=True): # The input text - input_text: Optional[str] # str + input_text: Optional[str] # str # The input token ids - input_ids: Optional[array] # Optional[array[int]] + input_ids: Optional[array] # Optional[array[int]] # The multimodal inputs mm_inputs: Optional[object] # The sampling parameters @@ -1060,9 +1060,9 @@ def __getitem__(self, i): class TokenizedEmbeddingReqInput(BaseReqIpc, kw_only=True): # The input text - input_text: Optional[str] # str + input_text: Optional[str] # str # The input token ids - input_ids: Optional[array] # array[int] + input_ids: Optional[array] # array[int] # The image inputs image_inputs: Optional[dict] # The token type ids @@ -1106,13 +1106,13 @@ def __iter__(self): class BatchTokenIDOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): # The finish reason - finished_reasons: List[Optional[Dict[str, Any]]] # List[BaseFinishReason] + finished_reasons: List[Optional[Dict[str, Any]]] # List[BaseFinishReason] # For incremental decoding decoded_texts: List[str] - decode_ids: List[array] # List[array[int]] + decode_ids: List[array] # List[array[int]] read_offsets: List[int] # Only used when `--skip-tokenizer-init` is on - output_ids: Optional[List[array]] # Optional[List[array[int]]] + output_ids: Optional[List[array]] # Optional[List[array[int]]] # Detokenization configs skip_special_tokens: List[bool] spaces_between_special_tokens: List[bool] @@ -1177,11 +1177,11 @@ class BatchTokenIDOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): class BatchStrOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): # The finish reason - finished_reasons: List[Optional[Dict[str, Any]]] # List[dict] + finished_reasons: List[Optional[Dict[str, Any]]] # List[dict] # The output decoded strings output_strs: List[str] # The token ids - output_ids: Optional[List[array]] # Optional[List[int]] + output_ids: Optional[List[array]] # Optional[List[int]] # Token counts prompt_tokens: List[int] @@ -1242,7 +1242,7 @@ class BatchStrOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): class BatchEmbeddingOutput(BaseBatchReqIpc, kw_only=True): # The finish reason - finished_reasons: List[Optional[Dict[str, Any]]] # List[BaseFinishReason] + finished_reasons: List[Optional[Dict[str, Any]]] # List[BaseFinishReason] # The output embedding embeddings: List[Union[List[float], Dict[int, float]]] # Token counts @@ -1316,7 +1316,7 @@ class ListExternalCorporaReqOutput(BaseReqIpc, kw_only=True): message: str = "" -# see aslo AttachHiCacheStorageReqInput in the http_server.py +# see also AttachHiCacheStorageReqInput in the http_server.py class AttachHiCacheStorageReqInput(BaseReqIpc, kw_only=True): """Dynamically attach (enable) HiCache storage backend at runtime. @@ -1784,6 +1784,7 @@ class ParseFunctionCallReq(BaseReq): None # Specify the parser type, e.g. 'llama3', 'qwen25', or 'mistral'. If not specified, tries all. ) + @dataclass class SeparateReasoningReqInput(BaseReq): text: str # The text to parse. @@ -2063,12 +2064,12 @@ def _check_all_req_types(): name.endswith("Req") or name.endswith("Input") or name.endswith("Output") - or name.endswith("Ipc") # allow IPC-specific struct naming + or name.endswith("Ipc") # allow IPC-specific struct naming ) is_base_req = ( issubclass(class_type[1], BaseReq) or issubclass(class_type[1], BaseBatchReq) - or issubclass(class_type[1], msgspec.Struct) # allow IPC objects + or issubclass(class_type[1], msgspec.Struct) # allow IPC objects ) if is_io_struct and not is_base_req: raise ValueError(f"{name} is not a subclass of BaseReq or BaseBatchReq.") @@ -2113,11 +2114,15 @@ def dec_hook(tp: Type, obj: Any) -> Any: return res else: if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): - logger.info(f"Object of type {type(obj)} and {str(tp)} is decoding with pickle.") + logger.info( + f"Object of type {type(obj)} and {str(tp)} is decoding with pickle." + ) return pickle.loads(obj) + _struct_types = tuple( - cls for cls in BaseReqIpc.__subclasses__() + cls + for cls in BaseReqIpc.__subclasses__() + BaseBatchReqIpc.__subclasses__() + SpeculativeDecodingMetricsMixinReq.__subclasses__() + [PickleWrapperIpc] @@ -2127,12 +2132,12 @@ def dec_hook(tp: Type, obj: Any) -> Any: _msgpack_encoder = msgspec.msgpack.Encoder(enc_hook=enc_hook) _msgpack_decoder = msgspec.msgpack.Decoder( - Union[*_struct_types, *_primitive_types], - dec_hook=dec_hook + Union[*_struct_types, *_primitive_types], dec_hook=dec_hook ) + def _maybe_wrap_pickle(obj: Any) -> PickleWrapperIpc: - # for msgspec.Struct type, we can directly serialize + # for msgspec.Struct type, we can directly serialize if isinstance(obj, msgspec.Struct): return obj # for primitive types, we can directly serialize @@ -2143,17 +2148,21 @@ def _maybe_wrap_pickle(obj: Any) -> PickleWrapperIpc: logger.info(f"Object of type {type(obj)} is wrapped via PickleWrapperIpc.") return PickleWrapperIpc(pickle.dumps(obj)) + def _maybe_unwrap_pickle(obj: Any) -> Any: # only PickleWrapperIpc type needs to be unwrapped, # other types can be directly used if isinstance(obj, PickleWrapperIpc): - obj = pickle.loads(obj.data) + obj = pickle.loads(obj.data) if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): - logger.info(f"Object of type {type(obj)} is unwrapped from PickleWrapperIpc.") + logger.info( + f"Object of type {type(obj)} is unwrapped from PickleWrapperIpc." + ) return obj return obj + def sock_send(socket: Socket, obj: Any, flags=0): obj = _maybe_wrap_pickle(obj) socket.send(_msgpack_encoder.encode(obj), flags=flags) diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index ccab34b93d1f..3c4e4c65d605 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -40,9 +40,9 @@ from sglang.srt.managers.disagg_service import start_disagg_service from sglang.srt.managers.io_struct import ( BaseBatchReq, + BaseBatchReqIpc, BaseReq, BaseReqIpc, - BaseBatchReqIpc, BatchEmbeddingOutput, BatchStrOutput, BatchTokenIDOutput, @@ -51,10 +51,10 @@ PauseContinueBroadcastReq, PauseGenerationReqInput, TokenizerWorkerRegistrationReq, - sock_recv, - sock_send, async_sock_recv, async_sock_send, + sock_recv, + sock_send, ) from sglang.srt.managers.tokenizer_manager import TokenizerManager from sglang.srt.server_args import PortArgs, ServerArgs @@ -588,7 +588,7 @@ async def pause_generation(self, obj: PauseGenerationReqInput): # Send to router which will broadcast to all workers # (router also handles forwarding to scheduler for non-abort modes) await self.send_to_scheduler.async_send_obj(obj) - + await self._pause_continue_future if obj.mode == "abort": diff --git a/python/sglang/srt/managers/scheduler_components/output_sender.py b/python/sglang/srt/managers/scheduler_components/output_sender.py index e6e7cf74737c..4f63fbb561b7 100644 --- a/python/sglang/srt/managers/scheduler_components/output_sender.py +++ b/python/sglang/srt/managers/scheduler_components/output_sender.py @@ -2,11 +2,7 @@ import zmq -from sglang.srt.managers.io_struct import( - BaseBatchReqIpc, - BaseReqIpc, - sock_send -) +from sglang.srt.managers.io_struct import BaseBatchReqIpc, BaseReqIpc, sock_send class SenderWrapper: diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index 7a4c448194fb..de4e80ded55e 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/python/sglang/srt/observability/req_time_stats.py @@ -225,9 +225,7 @@ class ReqTimeStatsBase: # Optional[Union[SchedulerMetricsCollector, TokenizerMetricsCollector]] metrics_collector: Optional[object] = None # Union[TraceReqContext, TraceNullContext] - trace_ctx: object = field( - default_factory=TraceNullContext - ) + trace_ctx: object = field(default_factory=TraceNullContext) disagg_mode: DisaggregationMode = DisaggregationMode.NULL diff_realtime_monotonic: float = 0.0 diff --git a/python/sglang/srt/observability/trace.py b/python/sglang/srt/observability/trace.py index b57d394ef0f0..5ebf7fd30671 100644 --- a/python/sglang/srt/observability/trace.py +++ b/python/sglang/srt/observability/trace.py @@ -684,7 +684,7 @@ def __del__(self): self.abort(abort_info={"reason": "have unclosed span, auto closed"}) -# Temperary remove the @dataclass decorator to make msgpack treat +# Temporary remove the @dataclass decorator to make msgpack treat # TraceReqContext as a custom class and serialize with pickle like TraceReqContext. # When the TraceReqContext fixed with msgpack, we can consider to add @dataclass back. class TraceNullContext: From 7142327a545e87d7b960d7748b5b7731b00b08cd Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Wed, 3 Jun 2026 18:03:17 -0700 Subject: [PATCH 03/45] support python 3.10 --- python/sglang/srt/managers/io_struct.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index a47825c30ac8..961af7ba5317 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -2129,11 +2129,11 @@ def dec_hook(tp: Type, obj: Any) -> Any: ) # str primitive types should encode to bytes _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[*_struct_types, *_primitive_types], dec_hook=dec_hook -) +_msgpack_decoder = msgspec.msgpack.Decoder(Union[_all_types], dec_hook=dec_hook) def _maybe_wrap_pickle(obj: Any) -> PickleWrapperIpc: From 3527ab9805a764abe661ee37a6c9f870d00784b2 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Thu, 4 Jun 2026 15:48:27 -0700 Subject: [PATCH 04/45] fix the ci --- python/sglang/srt/managers/communicator.py | 14 ++++++- .../srt/managers/data_parallel_controller.py | 2 +- python/sglang/srt/managers/io_struct.py | 37 ++++++++++++++++--- python/sglang/srt/managers/scheduler.py | 28 ++++++++------ .../srt/managers/tokenizer_control_mixin.py | 9 ++++- python/sglang/srt/server_args.py | 10 +++-- python/sglang/test/test_utils.py | 8 ++-- 7 files changed, 80 insertions(+), 28 deletions(-) diff --git a/python/sglang/srt/managers/communicator.py b/python/sglang/srt/managers/communicator.py index 75ccfc6a64f2..58dd0eed6d9f 100644 --- a/python/sglang/srt/managers/communicator.py +++ b/python/sglang/srt/managers/communicator.py @@ -35,6 +35,8 @@ def __init__(self, sender: zmq.Socket, fan_out: int, mode="queueing"): assert mode in ["queueing", "watching"] async def queueing_call(self, obj: T): + from sglang.srt.managers.multi_tokenizer_mixin import SenderWrapper + ready_event = asyncio.Event() if self._result_event is not None or len(self._ready_queue) > 0: self._ready_queue.append(ready_event) @@ -43,7 +45,10 @@ async def queueing_call(self, obj: T): assert self._result_values is None if obj is not None: - sock_send(self._sender, obj) + if isinstance(self._sender, SenderWrapper): + self._sender.send_obj(obj) + else: + sock_send(self._sender, obj) self._result_event = asyncio.Event() self._result_values = [] @@ -57,13 +62,18 @@ async def queueing_call(self, obj: T): return result_values async def watching_call(self, obj): + from sglang.srt.managers.multi_tokenizer_mixin import SenderWrapper + if self._result_event is None: assert self._result_values is None self._result_values = [] self._result_event = asyncio.Event() if obj is not None: - sock_send(self._sender, obj) + if isinstance(self._sender, SenderWrapper): + self._sender.send_obj(obj) + else: + sock_send(self._sender, obj) # Capture local refs before await -- after event fires, the first # awakened coroutine clears shared state; later awaiters use local refs. diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index 3ad13db7c448..47317eada921 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -382,7 +382,7 @@ def _broadcast_ports_as_server( connected_clients = 0 while connected_clients < expected_clients: # Wait for client handshake - client_rank = rep_socket.recv().decode() + client_rank = sock_recv(rep_socket) logger.debug(f"Received handshake from node {client_rank}") # Send worker ports to client diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 961af7ba5317..c06bf3568d26 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -30,6 +30,7 @@ from typing import TYPE_CHECKING, Annotated, Any, Dict, List, Optional, Type, Union import msgspec +import numpy as np import torch from pydantic import PlainValidator from zmq import Socket @@ -45,6 +46,7 @@ SchedulerReqTimeStats, ) from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.srt.server_args import ServerArgs from sglang.srt.utils import ImageData, VideoData from sglang.srt.utils.field_validators import validate_optional_list_i64_1d_2d @@ -1649,16 +1651,25 @@ class GetInternalStateReq(BaseReqIpc, kw_only=True): class GetInternalStateReqOutput(BaseReqIpc, kw_only=True): - internal_state: Dict[Any, Any] + server_args: ServerArgs + last_gen_throughput: float + memory_usage: Dict[str, Union[float, int]] + effective_max_running_requests_per_dp: int + avg_spec_accept_length: Optional[float] = None + step_time_dict: Optional[Dict[str, float]] = None +# The ServerArgs has a lot of fields with different types +# let's use object to make it serialize/deserialize via pickle for now class SetInternalStateReq(BaseReqIpc, kw_only=True): - server_args: Dict[str, Any] + server_args: Dict[str, object] +# The ServerArgs has a lot of fields with different types +# let's use object to make it serialize/deserialize via pickle for now class SetInternalStateReqOutput(BaseReqIpc, kw_only=True): updated: bool - server_args: Dict[str, Any] + server_args: Dict[str, object] @dataclass @@ -2095,10 +2106,19 @@ def enc_hook(obj: Any) -> Any: return (obj.shape, tensor_dtype, raw_data) elif isinstance(obj, array): return (obj.typecode, obj.tobytes()) + elif isinstance(obj, np.floating): + return float(obj) else: if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): logger.info(f"Object of type {type(obj)} is encoding with pickle.") - return pickle.dumps(obj) + try: + return pickle.dumps(obj) + except Exception as e: + if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): + logger.info( + f"Return None since failed to encode object of type {type(obj)} with pickle: {e}" + ) + return None def dec_hook(tp: Type, obj: Any) -> Any: @@ -2117,7 +2137,14 @@ def dec_hook(tp: Type, obj: Any) -> Any: logger.info( f"Object of type {type(obj)} and {str(tp)} is decoding with pickle." ) - return pickle.loads(obj) + try: + return pickle.loads(obj) + except Exception as e: + if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): + logger.info( + f"Return None since failed to decode object of type {type(obj)} with pickle: {e}" + ) + return None _struct_types = tuple( diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index e5953e567a6e..fa81896e5de7 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1561,7 +1561,7 @@ def process_input_requests(self, recv_reqs: List): self.ipc_channels.send_to_tokenizer.send_output(output, recv_req) else: if self.ipc_channels.recv_from_rpc is not None: - sock_send(self.ipc_channels.recv_from_rpc, output, recv_req) + sock_send(self.ipc_channels.recv_from_rpc, output) self.flush_wrapper.check_pending() if self.external_corpus_manager is not None: @@ -3462,9 +3462,9 @@ def flush_cache(self, empty_cache: bool = True): return success def get_internal_state(self, recv_req: GetInternalStateReq): - ret = dict(vars(get_global_server_args())) # vars returns a ref to obj.__dict__ - ret["last_gen_throughput"] = self.metrics_reporter.last_gen_throughput - ret["memory_usage"] = { + server_args = get_global_server_args() + last_gen_throughput = self.metrics_reporter.last_gen_throughput + memory_usage = { "weight": round(self.tp_worker.model_runner.weight_load_mem_usage, 2), "kvcache": round( self.token_to_kv_pool_allocator.get_kvcache().mem_usage, 2 @@ -3472,24 +3472,30 @@ def get_internal_state(self, recv_req: GetInternalStateReq): "token_capacity": int(self.max_total_num_tokens), "graph": round(self.tp_worker.model_runner.graph_mem_usage, 2), } - ret["effective_max_running_requests_per_dp"] = self.max_running_requests + effective_max_running_requests_per_dp = self.max_running_requests + avg_spec_accept_length = None if ( not self.spec_algorithm.is_none() and self.metrics_reporter.spec_total_num_forward_ct > 0 ): - ret["avg_spec_accept_length"] = ( + avg_spec_accept_length = ( self.metrics_reporter.spec_total_num_accept_tokens / self.metrics_reporter.spec_total_num_forward_ct ) + step_time_dict = None if RECORD_STEP_TIME: - ret["step_time_dict"] = self.metrics_reporter.step_time_dict + step_time_dict = self.metrics_reporter.step_time_dict - # This field is not serializable. - ret.pop("model_config", None) - - return GetInternalStateReqOutput(internal_state=ret) + return GetInternalStateReqOutput( + server_args=server_args, + last_gen_throughput=last_gen_throughput, + memory_usage=memory_usage, + effective_max_running_requests_per_dp=effective_max_running_requests_per_dp, + avg_spec_accept_length=avg_spec_accept_length, + step_time_dict=step_time_dict, + ) def set_internal_state(self, recv_req: SetInternalStateReq): server_args_dict = recv_req.server_args diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index c8700a120fdd..4d6dc762ddbc 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import fastapi +import msgspec from sglang.srt.managers.communicator import FanOutCommunicator from sglang.srt.managers.io_struct import ( @@ -791,8 +792,14 @@ async def get_internal_state(self: TokenizerManager) -> List[Dict[Any, Any]]: responses: List[GetInternalStateReqOutput] = ( await self.get_internal_state_communicator(req) ) + results = [] + for response in responses: + res = msgspec.to_builtins(response) + server_args = res.pop("server_args", None) or {} + res = {k: v for k, v in res.items() if v is not None} + results.append(server_args | res) # Many DP ranks - return [res.internal_state for res in responses] + return results async def set_internal_state( self: TokenizerManager, obj: SetInternalStateReq diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 010a4792e268..8b4f55e6030b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -559,10 +559,12 @@ class ServerArgs: enable_lora: Optional[bool] = None enable_lora_overlap_loading: Optional[bool] = None max_lora_rank: Optional[int] = None - lora_target_modules: Optional[Union[set[str], List[str]]] = None - lora_paths: Optional[ - Union[dict[str, str], List[dict[str, str]], List[str], List[LoRARef]] - ] = None + lora_target_modules: Optional[List[str]] = ( + None # Optional[Union[set[str], List[str]]] + ) + lora_paths: Optional[Any] = ( + None # Optional[Union[dict[str, str], List[dict[str, str]], List[str], List[LoRARef]]] + ) max_loaded_loras: Optional[int] = None max_loras_per_batch: int = 8 lora_eviction_policy: str = "lru" diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index 34d2d1e563c2..22936df5c395 100644 --- a/python/sglang/test/test_utils.py +++ b/python/sglang/test/test_utils.py @@ -106,11 +106,11 @@ DEFAULT_DRAFT_MODEL_EAGLE = "lmsys/sglang-EAGLE-llama2-chat-7B" # EAGLE3 model -DEFAULT_TARGET_MODEL_EAGLE3 = "meta-llama/Llama-3.1-8B-Instruct" +DEFAULT_TARGET_MODEL_EAGLE3 = "unsloth/Meta-Llama-3.1-8B-Instruct" DEFAULT_DRAFT_MODEL_EAGLE3 = "lmsys/sglang-EAGLE3-LLaMA3.1-Instruct-8B" # DFLASH model -DEFAULT_TARGET_MODEL_DFLASH = "meta-llama/Llama-3.1-8B-Instruct" +DEFAULT_TARGET_MODEL_DFLASH = "unsloth/Meta-Llama-3.1-8B-Instruct" DEFAULT_DRAFT_MODEL_DFLASH = "z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat" # EAGLE2 with DP-Attention models @@ -118,8 +118,8 @@ DEFAULT_DRAFT_MODEL_EAGLE_DP_ATTN = "Tengyunw/qwen3_30b_moe_eagle3" # Standalone speculative decoding models -DEFAULT_TARGET_MODEL_STANDALONE = "meta-llama/Llama-3.1-8B-Instruct" -DEFAULT_DRAFT_MODEL_STANDALONE = "meta-llama/Llama-3.2-1B-Instruct" +DEFAULT_TARGET_MODEL_STANDALONE = "unsloth/Meta-Llama-3.1-8B-Instruct" +DEFAULT_DRAFT_MODEL_STANDALONE = "unsloth/Meta-Llama-3.2-1B-Instruct" # N-gram speculative decoding models DEFAULT_TARGET_MODEL_NGRAM = "Qwen/Qwen2.5-Coder-7B-Instruct" From 78ccb65ef74f80c1c2ccdbf19eee498fb76f15a0 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Thu, 4 Jun 2026 16:08:02 -0700 Subject: [PATCH 05/45] revert the meta test model --- python/sglang/test/test_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index 22936df5c395..59f221644d0b 100644 --- a/python/sglang/test/test_utils.py +++ b/python/sglang/test/test_utils.py @@ -106,11 +106,11 @@ DEFAULT_DRAFT_MODEL_EAGLE = "lmsys/sglang-EAGLE-llama2-chat-7B" # EAGLE3 model -DEFAULT_TARGET_MODEL_EAGLE3 = "unsloth/Meta-Llama-3.1-8B-Instruct" +DEFAULT_TARGET_MODEL_EAGLE3 = "meta-llama/Meta-Llama-3.1-8B-Instruct" DEFAULT_DRAFT_MODEL_EAGLE3 = "lmsys/sglang-EAGLE3-LLaMA3.1-Instruct-8B" # DFLASH model -DEFAULT_TARGET_MODEL_DFLASH = "unsloth/Meta-Llama-3.1-8B-Instruct" +DEFAULT_TARGET_MODEL_DFLASH = "meta-llama/Meta-Llama-3.1-8B-Instruct" DEFAULT_DRAFT_MODEL_DFLASH = "z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat" # EAGLE2 with DP-Attention models @@ -118,8 +118,8 @@ DEFAULT_DRAFT_MODEL_EAGLE_DP_ATTN = "Tengyunw/qwen3_30b_moe_eagle3" # Standalone speculative decoding models -DEFAULT_TARGET_MODEL_STANDALONE = "unsloth/Meta-Llama-3.1-8B-Instruct" -DEFAULT_DRAFT_MODEL_STANDALONE = "unsloth/Meta-Llama-3.2-1B-Instruct" +DEFAULT_TARGET_MODEL_STANDALONE = "meta-llama/Meta-Llama-3.1-8B-Instruct" +DEFAULT_DRAFT_MODEL_STANDALONE = "meta-llama/Meta-Llama-3.2-1B-Instruct" # N-gram speculative decoding models DEFAULT_TARGET_MODEL_NGRAM = "Qwen/Qwen2.5-Coder-7B-Instruct" From df8e4c335a738ec5f9798d029f68c258357ecf84 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Thu, 4 Jun 2026 16:09:23 -0700 Subject: [PATCH 06/45] fix the model repo --- python/sglang/test/test_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index 59f221644d0b..34d2d1e563c2 100644 --- a/python/sglang/test/test_utils.py +++ b/python/sglang/test/test_utils.py @@ -106,11 +106,11 @@ DEFAULT_DRAFT_MODEL_EAGLE = "lmsys/sglang-EAGLE-llama2-chat-7B" # EAGLE3 model -DEFAULT_TARGET_MODEL_EAGLE3 = "meta-llama/Meta-Llama-3.1-8B-Instruct" +DEFAULT_TARGET_MODEL_EAGLE3 = "meta-llama/Llama-3.1-8B-Instruct" DEFAULT_DRAFT_MODEL_EAGLE3 = "lmsys/sglang-EAGLE3-LLaMA3.1-Instruct-8B" # DFLASH model -DEFAULT_TARGET_MODEL_DFLASH = "meta-llama/Meta-Llama-3.1-8B-Instruct" +DEFAULT_TARGET_MODEL_DFLASH = "meta-llama/Llama-3.1-8B-Instruct" DEFAULT_DRAFT_MODEL_DFLASH = "z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat" # EAGLE2 with DP-Attention models @@ -118,8 +118,8 @@ DEFAULT_DRAFT_MODEL_EAGLE_DP_ATTN = "Tengyunw/qwen3_30b_moe_eagle3" # Standalone speculative decoding models -DEFAULT_TARGET_MODEL_STANDALONE = "meta-llama/Meta-Llama-3.1-8B-Instruct" -DEFAULT_DRAFT_MODEL_STANDALONE = "meta-llama/Meta-Llama-3.2-1B-Instruct" +DEFAULT_TARGET_MODEL_STANDALONE = "meta-llama/Llama-3.1-8B-Instruct" +DEFAULT_DRAFT_MODEL_STANDALONE = "meta-llama/Llama-3.2-1B-Instruct" # N-gram speculative decoding models DEFAULT_TARGET_MODEL_NGRAM = "Qwen/Qwen2.5-Coder-7B-Instruct" From bb64f8bfb1e45abc24e9041df6b8757bb8f22e95 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Thu, 4 Jun 2026 16:28:34 -0700 Subject: [PATCH 07/45] fix logprob type error --- python/sglang/srt/managers/io_struct.py | 56 ++++++++++++------------- 1 file changed, 28 insertions(+), 28 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index b182886f2877..f8a5f457eb91 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1136,22 +1136,22 @@ class BatchTokenIDOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): cached_tokens: List[int] # Logprobs - input_token_logprobs_val: Optional[List[float]] - input_token_logprobs_idx: Optional[List[int]] - output_token_logprobs_val: Optional[List[float]] - output_token_logprobs_idx: Optional[List[int]] - input_top_logprobs_val: Optional[List[List]] - input_top_logprobs_idx: Optional[List[List]] - output_top_logprobs_val: Optional[List[List]] - output_top_logprobs_idx: Optional[List[List]] - input_token_ids_logprobs_val: Optional[List[List]] - input_token_ids_logprobs_idx: Optional[List[List]] - output_token_ids_logprobs_val: Optional[List[List]] - output_token_ids_logprobs_idx: Optional[List[List]] - output_token_entropy_val: Optional[List[float]] + input_token_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] + input_token_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] + output_token_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] + output_token_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] + input_top_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] + input_top_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] + output_top_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] + output_top_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] + input_token_ids_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] + input_token_ids_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] + output_token_ids_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] + output_token_ids_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] + output_token_entropy_val: Optional[List[Optional[float]]] # Hidden states - output_hidden_states: Optional[List[List[float]]] + output_hidden_states: Optional[List[Optional[List[Optional[float]]]]] # Per-request routed experts (input + output tokens), shape # (token, layer, top_k). DetokenizerManager encodes to base64 into @@ -1201,22 +1201,22 @@ class BatchStrOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): cached_tokens: List[int] # Logprobs - input_token_logprobs_val: Optional[List[float]] - input_token_logprobs_idx: Optional[List[int]] - output_token_logprobs_val: Optional[List[float]] - output_token_logprobs_idx: Optional[List[int]] - input_top_logprobs_val: Optional[List[List]] - input_top_logprobs_idx: Optional[List[List]] - output_top_logprobs_val: Optional[List[List]] - output_top_logprobs_idx: Optional[List[List]] - input_token_ids_logprobs_val: Optional[List[List]] - input_token_ids_logprobs_idx: Optional[List[List]] - output_token_ids_logprobs_val: Optional[List[List]] - output_token_ids_logprobs_idx: Optional[List[List]] - output_token_entropy_val: Optional[List[float]] + input_token_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] + input_token_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] + output_token_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] + output_token_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] + input_top_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] + input_top_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] + output_top_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] + output_top_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] + input_token_ids_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] + input_token_ids_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] + output_token_ids_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] + output_token_ids_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] + output_token_entropy_val: Optional[List[Optional[float]]] # Hidden states - output_hidden_states: Optional[List[List[float]]] + output_hidden_states: Optional[List[Optional[List[Optional[float]]]]] # Per-request routed experts, base64-encoded by DetokenizerManager off the # tokenizer hot path. Underlying tensor shape is (token, layer, top_k); From 2f8e99742005d1e6604bfc34bae3c60cd4861798 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Thu, 4 Jun 2026 16:35:43 -0700 Subject: [PATCH 08/45] fix ci for test_data_parallel_controller.py --- .../managers/test_data_parallel_controller.py | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/test/registered/unit/managers/test_data_parallel_controller.py b/test/registered/unit/managers/test_data_parallel_controller.py index 13c0917d69cf..aea06f482818 100644 --- a/test/registered/unit/managers/test_data_parallel_controller.py +++ b/test/registered/unit/managers/test_data_parallel_controller.py @@ -172,14 +172,14 @@ def test_cycles_through_active_workers_in_order(self): ctl.round_robin_scheduler(_req()) # 8 reqs across 4 active workers — 2 each, in round-robin order for i, worker in enumerate(ctl.workers): - self.assertEqual(worker.send_pyobj.call_count, 2, f"worker {i} call count") + self.assertEqual(worker.send.call_count, 2, f"worker {i} call count") def test_first_dispatch_picks_worker_zero(self): ctl = _make_controller(dp_size=4) ctl.round_robin_scheduler(_req()) - ctl.workers[0].send_pyobj.assert_called_once() + ctl.workers[0].send.assert_called_once() for i in (1, 2, 3): - ctl.workers[i].send_pyobj.assert_not_called() + ctl.workers[i].send.assert_not_called() def test_skips_inactive_workers(self): ctl = _make_controller(dp_size=4) @@ -188,16 +188,16 @@ def test_skips_inactive_workers(self): for _ in range(6): ctl.round_robin_scheduler(_req()) # Only workers 0 and 2 are active — should split 6 reqs evenly - self.assertEqual(ctl.workers[0].send_pyobj.call_count, 3) - ctl.workers[1].send_pyobj.assert_not_called() - self.assertEqual(ctl.workers[2].send_pyobj.call_count, 3) - ctl.workers[3].send_pyobj.assert_not_called() + self.assertEqual(ctl.workers[0].send.call_count, 3) + ctl.workers[1].send.assert_not_called() + self.assertEqual(ctl.workers[2].send.call_count, 3) + ctl.workers[3].send.assert_not_called() def test_routed_dp_rank_bypasses_counter(self): """External dp-rank routing must not advance the counter.""" ctl = _make_controller(dp_size=4) ctl.round_robin_scheduler(_req(routed_dp_rank=2)) - ctl.workers[2].send_pyobj.assert_called_once() + ctl.workers[2].send.assert_called_once() self.assertEqual( ctl.round_robin_counter, 0, @@ -205,7 +205,7 @@ def test_routed_dp_rank_bypasses_counter(self): ) # Subsequent round-robin req still lands on worker 0 ctl.round_robin_scheduler(_req()) - ctl.workers[0].send_pyobj.assert_called_once() + ctl.workers[0].send.assert_called_once() class TestFollowBootstrapRoomScheduler(CustomTestCase): @@ -220,7 +220,7 @@ def test_dispatches_by_bootstrap_room_modulo(self): (101, 1), ]: ctl.follow_bootstrap_room_scheduler(_req(bootstrap_room=room)) - ctl.workers[expected_rank].send_pyobj.assert_called() + ctl.workers[expected_rank].send.assert_called() def test_requires_bootstrap_room(self): ctl = _make_controller(dp_size=4) @@ -230,8 +230,8 @@ def test_requires_bootstrap_room(self): def test_routed_dp_rank_bypasses_bootstrap_room(self): ctl = _make_controller(dp_size=4) ctl.follow_bootstrap_room_scheduler(_req(routed_dp_rank=3, bootstrap_room=1)) - ctl.workers[3].send_pyobj.assert_called_once() - ctl.workers[1].send_pyobj.assert_not_called() + ctl.workers[3].send.assert_called_once() + ctl.workers[1].send.assert_not_called() class TestTotalRequestsScheduler(CustomTestCase): @@ -239,9 +239,9 @@ def test_dispatches_to_min_request_worker(self): ctl = _make_controller(dp_size=4) ctl.dp_budget.total_requests = [5, 3, 1, 4] ctl.total_requests_scheduler(_req()) - ctl.workers[2].send_pyobj.assert_called_once() + ctl.workers[2].send.assert_called_once() for i in (0, 1, 3): - ctl.workers[i].send_pyobj.assert_not_called() + ctl.workers[i].send.assert_not_called() self.assertEqual( ctl.dp_budget.total_requests[2], 2, @@ -252,7 +252,7 @@ def test_routed_dp_rank_bypasses_budget(self): ctl = _make_controller(dp_size=4) ctl.dp_budget.total_requests = [5, 3, 1, 4] ctl.total_requests_scheduler(_req(routed_dp_rank=0)) - ctl.workers[0].send_pyobj.assert_called_once() + ctl.workers[0].send.assert_called_once() # DPBudget must not be touched when bypassed self.assertEqual( ctl.dp_budget.total_requests, @@ -275,7 +275,7 @@ def test_total_requests_ignores_status(self): ctl.status[2] = False ctl.total_requests_scheduler(_req()) # Current behaviour: still dispatches to the inactive worker. - ctl.workers[2].send_pyobj.assert_called_once() + ctl.workers[2].send.assert_called_once() if __name__ == "__main__": From 54d2ae45692ae1b9a3ce0086d5e571e6a67459f3 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Thu, 4 Jun 2026 16:43:40 -0700 Subject: [PATCH 09/45] fix unit test for the tokenizer manager --- .../unit/managers/test_tokenizer_manager_rid_cleanup.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 95a8ffffd6d0..d9fe872c86cc 100644 --- a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py +++ b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py @@ -17,6 +17,8 @@ import unittest from unittest.mock import 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 @@ -159,7 +161,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": From 563c9bf12de30f866779b0b0ab3e0da0b8cf37ab Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Fri, 5 Jun 2026 08:05:22 -0700 Subject: [PATCH 10/45] fix ci test --- python/sglang/srt/managers/io_struct.py | 98 ++++++++++++++----- python/sglang/srt/sampling/sampling_params.py | 30 ++++++ 2 files changed, 103 insertions(+), 25 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index f8a5f457eb91..883540ec85d9 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -102,7 +102,7 @@ def regenerate_rids(self): # The BaseBatchReq IPC class for IPC object class BaseBatchReqIpc(msgspec.Struct, tag=True, kw_only=True): rids: Optional[List[str]] = None - http_worker_ipcs: Optional[List[str]] = None + http_worker_ipcs: Optional[List[Optional[str]]] = None class SpeculativeDecodingMetricsMixinReq(BaseBatchReqIpc, kw_only=True): @@ -1136,18 +1136,42 @@ class BatchTokenIDOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): cached_tokens: List[int] # Logprobs - input_token_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] - input_token_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] - output_token_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] - output_token_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] - input_top_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] - input_top_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] - output_top_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] - output_top_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] - input_token_ids_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] - input_token_ids_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] - output_token_ids_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] - output_token_ids_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] + input_token_logprobs_val: Optional[ + List[Optional[List[Optional[Union[float, List[Optional[float]]]]]]] + ] + input_token_logprobs_idx: Optional[ + List[Optional[List[Optional[Union[int, List[Optional[int]]]]]]] + ] + output_token_logprobs_val: Optional[ + List[Optional[List[Optional[Union[float, List[Optional[float]]]]]]] + ] + output_token_logprobs_idx: Optional[ + List[Optional[List[Optional[Union[int, List[Optional[int]]]]]]] + ] + input_top_logprobs_val: Optional[ + List[Optional[List[Optional[Union[float, List[Optional[float]]]]]]] + ] + input_top_logprobs_idx: Optional[ + List[Optional[List[Optional[Union[int, List[Optional[int]]]]]]] + ] + output_top_logprobs_val: Optional[ + List[Optional[List[Optional[Union[float, List[Optional[float]]]]]]] + ] + output_top_logprobs_idx: Optional[ + List[Optional[List[Optional[Union[int, List[Optional[int]]]]]]] + ] + input_token_ids_logprobs_val: Optional[ + List[Optional[List[Optional[Union[float, List[Optional[float]]]]]]] + ] + input_token_ids_logprobs_idx: Optional[ + List[Optional[List[Optional[Union[int, List[Optional[int]]]]]]] + ] + output_token_ids_logprobs_val: Optional[ + List[Optional[List[Optional[Union[float, List[Optional[float]]]]]]] + ] + output_token_ids_logprobs_idx: Optional[ + List[Optional[List[Optional[Union[int, List[Optional[int]]]]]]] + ] output_token_entropy_val: Optional[List[Optional[float]]] # Hidden states @@ -1201,18 +1225,42 @@ class BatchStrOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): cached_tokens: List[int] # Logprobs - input_token_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] - input_token_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] - output_token_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] - output_token_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] - input_top_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] - input_top_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] - output_top_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] - output_top_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] - input_token_ids_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] - input_token_ids_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] - output_token_ids_logprobs_val: Optional[List[Optional[List[Optional[float]]]]] - output_token_ids_logprobs_idx: Optional[List[Optional[List[Optional[int]]]]] + input_token_logprobs_val: Optional[ + List[Optional[List[Optional[Union[float, List[Optional[float]]]]]]] + ] + input_token_logprobs_idx: Optional[ + List[Optional[List[Optional[Union[int, List[Optional[int]]]]]]] + ] + output_token_logprobs_val: Optional[ + List[Optional[List[Optional[Union[float, List[Optional[float]]]]]]] + ] + output_token_logprobs_idx: Optional[ + List[Optional[List[Optional[Union[int, List[Optional[int]]]]]]] + ] + input_top_logprobs_val: Optional[ + List[Optional[List[Optional[Union[float, List[Optional[float]]]]]]] + ] + input_top_logprobs_idx: Optional[ + List[Optional[List[Optional[Union[int, List[Optional[int]]]]]]] + ] + output_top_logprobs_val: Optional[ + List[Optional[List[Optional[Union[float, List[Optional[float]]]]]]] + ] + output_top_logprobs_idx: Optional[ + List[Optional[List[Optional[Union[int, List[Optional[int]]]]]]] + ] + input_token_ids_logprobs_val: Optional[ + List[Optional[List[Optional[Union[float, List[Optional[float]]]]]]] + ] + input_token_ids_logprobs_idx: Optional[ + List[Optional[List[Optional[Union[int, List[Optional[int]]]]]]] + ] + output_token_ids_logprobs_val: Optional[ + List[Optional[List[Optional[Union[float, List[Optional[float]]]]]]] + ] + output_token_ids_logprobs_idx: Optional[ + List[Optional[List[Optional[Union[int, List[Optional[int]]]]]]] + ] output_token_entropy_val: Optional[List[Optional[float]]] # Hidden states diff --git a/python/sglang/srt/sampling/sampling_params.py b/python/sglang/srt/sampling/sampling_params.py index 0744a57fdd82..3b0c0ea7fba2 100644 --- a/python/sglang/srt/sampling/sampling_params.py +++ b/python/sglang/srt/sampling/sampling_params.py @@ -14,6 +14,7 @@ """Sampling parameters for text generation.""" import logging +from dataclasses import dataclass from typing import Any, Dict, List, Optional, Union # sre_parse is deprecated in Python 3.11+, use re._parser instead @@ -28,6 +29,7 @@ logger = logging.getLogger(__name__) +@dataclass class SamplingParams: """ The sampling parameters. @@ -37,6 +39,34 @@ class SamplingParams: for the documentation. """ + max_new_tokens: Optional[int] + stop_strs: Optional[Union[str, List[str]]] + stop_token_ids: Optional[List[int]] + stop_regex_strs: Optional[Union[str, List[str]]] + stop_str_max_len: int + stop_regex_max_len: int + temperature: float + top_p: float + top_k: int + min_p: float + frequency_penalty: float + presence_penalty: float + repetition_penalty: float + min_new_tokens: int + regex: Optional[str] + n: int + json_schema: Optional[str] + ebnf: Optional[str] + structural_tag: Optional[str] + ignore_eos: bool + skip_special_tokens: bool + spaces_between_special_tokens: bool + no_stop_trim: bool + custom_params: Optional[Dict[str, Any]] + stream_interval: Optional[int] + logit_bias: Optional[Dict[str, float]] + sampling_seed: Optional[int] + def __init__( self, max_new_tokens: int = 128, From c250fd02236c2a8624cfbeb92cd85307a25447e1 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Fri, 5 Jun 2026 21:29:43 -0700 Subject: [PATCH 11/45] fix ci test --- python/sglang/srt/disaggregation/encode_server.py | 8 ++++---- python/sglang/srt/entrypoints/grpc_server.py | 2 +- python/sglang/srt/managers/io_struct.py | 2 +- python/sglang/srt/managers/mm_utils.py | 5 +++++ python/sglang/srt/managers/tokenizer_control_mixin.py | 2 +- test/registered/unit/utils/test_profile_merger.py | 4 ++-- 6 files changed, 14 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index 8fc84e687fab..d4d855f9f889 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -2887,7 +2887,7 @@ async def _dp_worker_handle_profile( req = ( ProfileReq(**obj) if obj is not None - else ProfileReq(ProfileReqType.START_PROFILE) + else ProfileReq(req_type=ProfileReqType.START_PROFILE) ) if enc.profiler is None: enc.profiler = EncoderProfiler(dp_rank) @@ -3627,7 +3627,7 @@ async def start_profile_async(obj: Optional[ProfileReqInput] = None): profile_req = None if obj is not None: profile_req = { - "type": ProfileReqType.START_PROFILE, + "req_type": ProfileReqType.START_PROFILE, "output_dir": obj.output_dir, "start_step": obj.start_step, "num_steps": obj.num_steps, @@ -3651,7 +3651,7 @@ async def start_profile_async(obj: Optional[ProfileReqInput] = None): return Response(content="encoder not ready\n", status_code=503) req = None if obj is None: - req = ProfileReq(ProfileReqType.START_PROFILE) + req = ProfileReq(req_type=ProfileReqType.START_PROFILE) else: req = ProfileReq( req_type=ProfileReqType.START_PROFILE, @@ -3697,7 +3697,7 @@ async def stop_profile_async(): return Response( content="profiling not initialized\n", status_code=HTTPStatus.BAD_REQUEST ) - req = ProfileReq(ProfileReqType.STOP_PROFILE) + req = ProfileReq(req_type=ProfileReqType.STOP_PROFILE) for socket in send_sockets: sock_send(socket, req) ok, msg = encoder.profiler.stop() diff --git a/python/sglang/srt/entrypoints/grpc_server.py b/python/sglang/srt/entrypoints/grpc_server.py index 964ed379fe79..320e04a94d8b 100644 --- a/python/sglang/srt/entrypoints/grpc_server.py +++ b/python/sglang/srt/entrypoints/grpc_server.py @@ -134,7 +134,7 @@ async def start_profile_handler(request): async def stop_profile_handler(request): try: - req = ProfileReq(ProfileReqType.STOP_PROFILE) + req = ProfileReq(req_type=ProfileReqType.STOP_PROFILE) results = await request_manager.send_communicator_req( req, "profile_communicator", timeout=600.0 ) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 883540ec85d9..05213925cdfe 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1077,7 +1077,7 @@ class TokenizedEmbeddingReqInput(BaseReqIpc, kw_only=True): # The image inputs image_inputs: Optional[dict] # The token type ids - token_type_ids: List[int] + token_type_ids: Optional[List[int]] # Dummy sampling params for compatibility sampling_params: SamplingParams # Embedding overrides to place at specific token positions. diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 1810d9cd0629..841ecd1ea281 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -1655,6 +1655,11 @@ def wrap_shm_features(obj): item.precomputed_embeddings = _wrap_tensor_or_list( item.precomputed_embeddings ) + # Use pickle to serialize the mm_inputs with SHM pointers, + # so that msgpack doesn't try to serialize the potentially + # large tensors inside mm_inputs + obj.mm_inputs = pickle.dumps(obj.mm_inputs) + return obj diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index 4d6dc762ddbc..844c75a29587 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -354,7 +354,7 @@ async def start_profile( async def stop_profile(self: TokenizerManager): self.auto_create_handle_loop() - req = ProfileReq(ProfileReqType.STOP_PROFILE) + req = ProfileReq(req_type=ProfileReqType.STOP_PROFILE) return await self._execute_profile(req) async def _execute_profile(self: TokenizerManager, req: ProfileReq): diff --git a/test/registered/unit/utils/test_profile_merger.py b/test/registered/unit/utils/test_profile_merger.py index 63fa90778246..e716f85b11a8 100644 --- a/test/registered/unit/utils/test_profile_merger.py +++ b/test/registered/unit/utils/test_profile_merger.py @@ -216,10 +216,10 @@ def test_data_structures_merge_profiles(self): self.assertTrue(req_input.merge_profiles) # Test ProfileReq - req = ProfileReq(ProfileReqType.START_PROFILE) + req = ProfileReq(req_type=ProfileReqType.START_PROFILE) self.assertFalse(req.merge_profiles) - req = ProfileReq(ProfileReqType.START_PROFILE, merge_profiles=True) + req = ProfileReq(req_type=ProfileReqType.START_PROFILE, merge_profiles=True) self.assertTrue(req.merge_profiles) def test_integration_parameters(self): From 77255065b3c6dcdcdbcd13649e3f32a2682adfaf Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Fri, 5 Jun 2026 22:41:03 -0700 Subject: [PATCH 12/45] fix ci test --- python/sglang/srt/managers/io_struct.py | 12 ++++++------ python/sglang/srt/session/session_controller.py | 6 +++--- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 05213925cdfe..5b62a80e0127 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -114,17 +114,17 @@ class SpeculativeDecodingMetricsMixinReq(BaseBatchReqIpc, kw_only=True): """ # Verify count: number of verification forward passes - spec_verify_ct: List[int] + spec_verify_ct: Optional[List[int]] # Accepted drafts: Number of accepted draft tokens during speculative decoding # (strict drafts-only count, excludes the bonus token). - spec_num_correct_drafts: List[int] + spec_num_correct_drafts: Optional[List[int]] # Acceptance histogram: List of lists, where each inner list represents histogram counts. # List index = number of accepted tokens in a step, List value = count of steps with that many accepted tokens. # Example: histogram[0] = 5 means 5 steps with 0 accepted tokens, histogram[3] = 10 means 10 steps with 3 accepted tokens. # Empty list [] when speculative decoding is disabled. - spec_correct_drafts_histogram: List[List[int]] + spec_correct_drafts_histogram: Optional[List[List[int]]] # Parameters for a session @@ -751,7 +751,7 @@ def __getitem__(self, i): class TokenizedGenerateReqInput(BaseReqIpc, kw_only=True): # The input text - input_text: Optional[str] # str + input_text: Optional[Union[str, List[Union[str, List[str]]]]] # str # The input token ids input_ids: Optional[array] # Optional[array[int]] # The multimodal inputs @@ -1071,7 +1071,7 @@ def __getitem__(self, i): class TokenizedEmbeddingReqInput(BaseReqIpc, kw_only=True): # The input text - input_text: Optional[str] # str + input_text: Optional[Union[str, List[Union[str, List[str]]]]] # str # The input token ids input_ids: Optional[array] # array[int] # The image inputs @@ -1303,7 +1303,7 @@ class BatchEmbeddingOutput(BaseBatchReqIpc, kw_only=True): # The finish reason finished_reasons: List[Optional[Dict[str, Any]]] # List[BaseFinishReason] # The output embedding - embeddings: List[Union[List[float], Dict[int, float]]] + embeddings: List[Union[List[float], Dict[int, float], float]] # Token counts prompt_tokens: List[int] cached_tokens: List[int] diff --git a/python/sglang/srt/session/session_controller.py b/python/sglang/srt/session/session_controller.py index 1915c370a540..d6a2a094133f 100644 --- a/python/sglang/srt/session/session_controller.py +++ b/python/sglang/srt/session/session_controller.py @@ -285,10 +285,10 @@ def open(self, recv_req: OpenSessionReqInput) -> OpenSessionReqOutput: session_id = recv_req.session_id if session_id in self.sessions: logger.warning(f"session id {session_id} already exist, cannot open.") - return OpenSessionReqOutput(session_id, False) + return OpenSessionReqOutput(session_id=session_id, success=False) elif session_id is None: logger.warning("session id is None, cannot open.") - return OpenSessionReqOutput(session_id, False) + return OpenSessionReqOutput(session_id=session_id, success=False) else: self.sessions[session_id] = Session( recv_req.capacity_of_str_len, @@ -299,7 +299,7 @@ def open(self, recv_req: OpenSessionReqInput) -> OpenSessionReqOutput: log_info_on_rank0( logger, f"Session opened: {session_id} (active={len(self.sessions)})" ) - return OpenSessionReqOutput(session_id, True) + return OpenSessionReqOutput(session_id=session_id, success=True) def close(self, recv_req: CloseSessionReqInput): session_id = recv_req.session_id From c748f365a66ff476de9aa15e2d8f423d3912dc00 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Sat, 6 Jun 2026 17:24:50 -0700 Subject: [PATCH 13/45] fix ci tests --- python/sglang/srt/managers/io_struct.py | 30 ++++++++++++++----- .../srt/managers/multi_tokenizer_mixin.py | 4 +-- .../scripted_runtime/tokenizer_recv_proxy.py | 6 ++-- 3 files changed, 29 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 5b62a80e0127..7da4fd22c922 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1303,7 +1303,7 @@ class BatchEmbeddingOutput(BaseBatchReqIpc, kw_only=True): # The finish reason finished_reasons: List[Optional[Dict[str, Any]]] # List[BaseFinishReason] # The output embedding - embeddings: List[Union[List[float], Dict[int, float], float]] + embeddings: List[Union[List[Union[float, List[float]]], Dict[int, float], float]] # Token counts prompt_tokens: List[int] cached_tokens: List[int] @@ -2249,21 +2249,37 @@ def _maybe_unwrap_pickle(obj: Any) -> Any: 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: Socket, obj: Any, flags=0): - obj = _maybe_wrap_pickle(obj) - socket.send(_msgpack_encoder.encode(obj), flags=flags) + socket.send(msgpack_encode(obj), flags=flags) def sock_recv(socket: Socket, flags=0) -> Any: data = socket.recv(flags=flags) - return _maybe_unwrap_pickle(_msgpack_decoder.decode(data)) + if isinstance(data, bytes): + return msgpack_decode(data) + else: + # for non-bytes data, we assume it has been decoded + # e.g. ScriptedTokenizerRecvProxy + return data async def async_sock_send(socket: AsyncSocket, obj: Any, flags=0): - obj = _maybe_wrap_pickle(obj) - await socket.send(_msgpack_encoder.encode(obj), flags=flags) + await socket.send(msgpack_encode(obj), flags=flags) async def async_sock_recv(socket: AsyncSocket, flags=0) -> Any: data = await socket.recv(flags=flags) - return _maybe_unwrap_pickle(_msgpack_decoder.decode(data)) + if isinstance(data, bytes): + return msgpack_decode(data) + else: + # for non-bytes data, we assume it has been decoded + # e.g. ScriptedTokenizerRecvProxy + return data diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index c6fe84af8198..a011608c9138 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -756,11 +756,11 @@ def __init__( self.is_multi_tokenizer = is_multi_tokenizer def send_obj(self, obj): - if not self.is_multi_tokenizer and hasattr(obj, "http_worker_ipc"): + if hasattr(obj, "http_worker_ipc"): obj.http_worker_ipc = self.port_args.tokenizer_ipc_name sock_send(self.send_to_scheduler, obj) async def async_send_obj(self, obj): - if not self.is_multi_tokenizer and hasattr(obj, "http_worker_ipc"): + if hasattr(obj, "http_worker_ipc"): obj.http_worker_ipc = self.port_args.tokenizer_ipc_name await async_sock_send(self.send_to_scheduler, obj) diff --git a/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py b/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py index e3fe19c20c2d..a24d75d730f4 100644 --- a/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py +++ b/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py @@ -6,6 +6,8 @@ import zmq +from sglang.srt.managers.io_struct import sock_recv + class ScriptedTokenizerRecvProxy: @@ -13,7 +15,7 @@ def __init__(self, *, underlying: zmq.Socket) -> None: self._underlying = underlying self._buffer: deque = deque() - def recv_pyobj(self, flags: int = 0) -> Any: + def recv(self, flags: int = 0) -> Any: self._drain_underlying() if self._buffer: @@ -49,7 +51,7 @@ def wait_until_arrived( def _drain_underlying(self) -> None: while True: try: - req = self._underlying.recv_pyobj(zmq.NOBLOCK) + req = sock_recv(self._underlying, flags=zmq.NOBLOCK) except zmq.ZMQError: break self._buffer.append(req) From 9023dc6c88b3cc61defa4cbeb0968e88eef96bea Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Sat, 6 Jun 2026 17:57:18 -0700 Subject: [PATCH 14/45] fix ci test --- .../test_tokenizer_recv_proxy.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/test/registered/unit/scripted_runtime/test_tokenizer_recv_proxy.py b/test/registered/unit/scripted_runtime/test_tokenizer_recv_proxy.py index 8adbe40f6471..b5581f3675ad 100644 --- a/test/registered/unit/scripted_runtime/test_tokenizer_recv_proxy.py +++ b/test/registered/unit/scripted_runtime/test_tokenizer_recv_proxy.py @@ -40,7 +40,7 @@ def feed(self, obj: object) -> None: def feed_after_drain_cycles(self, obj: object, *, cycles: int) -> None: self._scheduled.append([cycles, obj]) - def recv_pyobj(self, flags: int = 0) -> object: + def recv(self, flags: int = 0) -> object: if self._ready: return self._ready.popleft() @@ -70,21 +70,21 @@ def test_recv_pyobj_drains_then_pops_fifo(self): underlying.feed(first) underlying.feed(second) - self.assertIs(proxy.recv_pyobj(), first) - self.assertIs(proxy.recv_pyobj(), second) + self.assertIs(proxy.recv(), first) + self.assertIs(proxy.recv(), second) def test_recv_pyobj_empty_noblock_raises_eagain(self): proxy = ScriptedTokenizerRecvProxy(underlying=_FakeUnderlyingSocket()) with self.assertRaises(zmq.ZMQError) as ctx: - proxy.recv_pyobj(zmq.NOBLOCK) + proxy.recv(zmq.NOBLOCK) self.assertEqual(ctx.exception.errno, zmq.EAGAIN) def test_recv_pyobj_empty_blocking_raises_runtime_error(self): proxy = ScriptedTokenizerRecvProxy(underlying=_FakeUnderlyingSocket()) with self.assertRaisesRegex(RuntimeError, "blocking recv is not supported"): - proxy.recv_pyobj() + proxy.recv() class TestScriptedTokenizerRecvProxyWaitUntilArrived(CustomTestCase): @@ -105,7 +105,7 @@ def test_wait_until_arrived_returns_on_first_match_when_buffer_empty(self): proxy.wait_until_arrived(_is_control, timeout_s=1.0) - self.assertIs(proxy.recv_pyobj(), msg) + self.assertIs(proxy.recv(), msg) def test_wait_until_arrived_skips_stale_same_type_object(self): proxy, _, _ = self._proxy_with_stale_control() @@ -120,8 +120,8 @@ def test_wait_until_arrived_returns_on_new_object_after_stale(self): proxy.wait_until_arrived(_is_control, timeout_s=2.0) - self.assertIs(proxy.recv_pyobj(), stale) - self.assertIs(proxy.recv_pyobj(), fresh) + self.assertIs(proxy.recv(), stale) + self.assertIs(proxy.recv(), fresh) def test_wait_until_arrived_rid_predicate_ignores_stale_other_rid(self): underlying = _FakeUnderlyingSocket() @@ -134,8 +134,8 @@ def test_wait_until_arrived_rid_predicate_ignores_stale_other_rid(self): underlying.feed(new) proxy.wait_until_arrived(_is_start_req("new"), timeout_s=1.0) - self.assertIs(proxy.recv_pyobj(), old) - self.assertIs(proxy.recv_pyobj(), new) + self.assertIs(proxy.recv(), old) + self.assertIs(proxy.recv(), new) def test_wait_until_arrived_rid_predicate_skips_stale_same_rid(self): underlying = _FakeUnderlyingSocket() From fad4cb0669b2620d150a16250d3df2ae079b70e0 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Sun, 7 Jun 2026 11:39:22 -0700 Subject: [PATCH 15/45] fix ci tests --- python/sglang/srt/managers/io_struct.py | 28 +++++++++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index fb44ef326e1a..12147152f47b 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1175,7 +1175,19 @@ class BatchTokenIDOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): output_token_entropy_val: Optional[List[Optional[float]]] # Hidden states - output_hidden_states: Optional[List[Optional[List[Optional[float]]]]] + output_hidden_states: Optional[ + List[ + Optional[ + List[ + Optional[ + Union[ + float, List[Optional[Union[float, List[Optional[float]]]]] + ] + ] + ] + ] + ] + ] # Per-request routed experts (input + output tokens), shape # (token, layer, top_k). DetokenizerManager encodes to base64 into @@ -1264,7 +1276,19 @@ class BatchStrOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): output_token_entropy_val: Optional[List[Optional[float]]] # Hidden states - output_hidden_states: Optional[List[Optional[List[Optional[float]]]]] + output_hidden_states: Optional[ + List[ + Optional[ + List[ + Optional[ + Union[ + float, List[Optional[Union[float, List[Optional[float]]]]] + ] + ] + ] + ] + ] + ] # Per-request routed experts, base64-encoded by DetokenizerManager off the # tokenizer hot path. Underlying tensor shape is (token, layer, top_k); From 1379286e4d02374cb3ff444f45076b1a73f3180f Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Sun, 7 Jun 2026 11:46:12 -0700 Subject: [PATCH 16/45] http2 fix --- python/sglang/srt/entrypoints/engine.py | 11 +++++------ python/sglang/srt/entrypoints/http_server.py | 2 +- python/sglang/srt/managers/multi_tokenizer_mixin.py | 2 -- python/sglang/srt/managers/tokenizer_manager.py | 8 ++++++-- python/sglang/srt/server_args.py | 2 +- 5 files changed, 13 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index b567f4920f64..c49e8137ffb8 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -863,15 +863,14 @@ def _launch_subprocesses( for p in detoken_procs: scheduler_init_result.all_child_pids.append(p.pid) - # Init tokenizer manager first, as the bootstrap server is initialized here - if server_args.tokenizer_worker_num == 1: - tokenizer_manager, template_manager = init_tokenizer_manager_func( - server_args, port_args - ) - else: + if server_args.tokenizer_worker_num > 1 or server_args.enable_http2: # Launch multi-tokenizer router tokenizer_manager = MultiTokenizerRouter(server_args, port_args) template_manager = None + else: + tokenizer_manager, template_manager = init_tokenizer_manager_func( + server_args, port_args + ) # Wait for the model to finish loading scheduler_init_result.wait_for_ready() diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 31ce8b4df3e7..8b37bb278e99 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -2576,7 +2576,7 @@ def _setup_and_run_http_server( # workers (forked or spawned) can locate the shared memory # segment created above. envs.SGLANG_GRANIAN_PARENT_PID.set(os.getpid()) - _close_main_process_sockets() + # _close_main_process_sockets() _run_granian_server(server_args) finally: if multi_tokenizer_args_shm is not None: diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index a011608c9138..3c8da710a6cc 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -749,11 +749,9 @@ def __init__( self, port_args: PortArgs, send_to_scheduler: zmq.Socket, - is_multi_tokenizer: bool = False, ): self.port_args = port_args self.send_to_scheduler = send_to_scheduler - self.is_multi_tokenizer = is_multi_tokenizer def send_obj(self, obj): if hasattr(obj, "http_worker_ipc"): diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 9fdaeed7e2cb..55c9d8d3d9ac 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -380,7 +380,11 @@ def init_ipc_channels(self, port_args: PortArgs): self.recv_from_detokenizer = get_zmq_socket( context, zmq.PULL, port_args.tokenizer_ipc_name, True ) - if self.server_args.tokenizer_worker_num == 1: + # for http2, use the connect instead of binding + if ( + self.server_args.tokenizer_worker_num == 1 + and not self.server_args.enable_http2 + ): send_to_scheduler = get_zmq_socket( context, zmq.PUSH, port_args.scheduler_input_ipc_name, True ) @@ -391,7 +395,7 @@ def init_ipc_channels(self, port_args: PortArgs): context, zmq.PUSH, port_args.tokenizer_worker_ipc_name, False ) # Make sure that each request carries the tokenizer_ipc_name for response routing - self.send_to_scheduler = SenderWrapper(port_args, send_to_scheduler, True) + self.send_to_scheduler = SenderWrapper(port_args, send_to_scheduler) self.load_snapshot_reader = create_load_snapshot_reader( self.server_args, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index eda74b609d7e..734cc111a110 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -7995,7 +7995,7 @@ def init_new( else: nccl_port = server_args.nccl_port - if server_args.tokenizer_worker_num == 1: + if server_args.tokenizer_worker_num == 1 and not server_args.enable_http2: tokenizer_worker_ipc_name = None else: tokenizer_worker_ipc_name = ( From 2b331124c68a53982774cfebb4e629cd70823889 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Sun, 7 Jun 2026 12:23:39 -0700 Subject: [PATCH 17/45] fix ci test --- .../scheduler_components/weight_updater.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index 59f6597b3519..12e44c77af53 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py @@ -95,12 +95,14 @@ def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): self.flush_cache_after_weight_update(recv_req) if not success: logger.error(message) - return UpdateWeightFromDiskReqOutput(success, message, 0) + return UpdateWeightFromDiskReqOutput( + success=success, message=message, num_paused_requests=0 + ) def init_weights_update_group(self, recv_req: InitWeightsUpdateGroupReqInput): """Initialize the online model parameter update group.""" success, message = self.tp_worker.init_weights_update_group(recv_req) - return InitWeightsUpdateGroupReqOutput(success, message) + return InitWeightsUpdateGroupReqOutput(success=success, message=message) def destroy_weights_update_group( self, @@ -108,7 +110,7 @@ def destroy_weights_update_group( ): """Destroy the online model parameter update group.""" success, message = self.tp_worker.destroy_weights_update_group(recv_req) - return DestroyWeightsUpdateGroupReqOutput(success, message) + return DestroyWeightsUpdateGroupReqOutput(success=success, message=message) def update_weights_from_distributed( self, @@ -120,7 +122,7 @@ def update_weights_from_distributed( self.flush_cache_after_weight_update(recv_req) else: logger.error(message) - return UpdateWeightsFromDistributedReqOutput(success, message) + return UpdateWeightsFromDistributedReqOutput(success=success, message=message) def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput): """Update the online model parameter from tensors.""" @@ -134,7 +136,7 @@ def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput): else: logger.error(message) torch.distributed.barrier(group=self.tp_cpu_group) - return UpdateWeightsFromTensorReqOutput(success, message) + return UpdateWeightsFromTensorReqOutput(success=success, message=message) def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput): """Update the online model parameter from IPC for checkpoint-engine integration.""" @@ -147,7 +149,7 @@ def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput): if not success: logger.error(message) torch.distributed.barrier(group=self.tp_cpu_group) - return UpdateWeightsFromIPCReqOutput(success, message) + return UpdateWeightsFromIPCReqOutput(success=success, message=message) def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput): parameter = self.tp_worker.get_weights_by_name(recv_req) From 6709987898664e3d1654fe19f0dbb6cabb97b9a0 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Sun, 7 Jun 2026 14:46:17 -0700 Subject: [PATCH 18/45] fix the ci test --- .../scheduler_components/output_streamer.py | 4 +++- .../sglang/srt/managers/tokenizer_manager.py | 21 ++++++++++++++----- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index cac807158563..06db448da3c0 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -217,7 +217,9 @@ def _stream_output_embedding(self, reqs: List[Req]): all_have_phs = all(t is not None for t in phs_list) if all_have_phs: if all(t.shape == phs_list[0].shape for t in phs_list): - stacked_phs = torch.stack(phs_list) + # when stack all tensors, to keep the pooled_hidden_states data type + # wrap the tensor with list to avoid the msgpack issue + stacked_phs = [torch.stack(phs_list)] else: stacked_phs = phs_list else: diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 55c9d8d3d9ac..d499c2c8e871 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1994,11 +1994,22 @@ async def _handle_batch_output( "embedding": recv_obj.embeddings[i], "meta_info": meta_info, } - if ( - recv_obj.pooled_hidden_states is not None - and recv_obj.pooled_hidden_states[i] is not None - ): - out_dict["pooled_hidden_state"] = recv_obj.pooled_hidden_states[i] + if recv_obj.pooled_hidden_states is not None: + # for non-stacked, pooled_hidden_states is a list of tensors + if len(recv_obj.pooled_hidden_states) == len(recv_obj.rids): + if recv_obj.pooled_hidden_states[i] is not None: + out_dict["pooled_hidden_state"] = ( + recv_obj.pooled_hidden_states[i] + ) + # for stacked, pooled_hidden_states is a single tensor list + elif len(recv_obj.pooled_hidden_states) == 1: + if ( + recv_obj.pooled_hidden_states[0] is not None + and recv_obj.pooled_hidden_states[0][i] is not None + ): + out_dict["pooled_hidden_state"] = ( + recv_obj.pooled_hidden_states[0][i] + ) # Set first_token_time on the first output batch. # This is the single write point for first_token_time. From 690a4427e7f9485e88ec30d2edb26d4914de68f0 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Sun, 7 Jun 2026 18:24:13 -0700 Subject: [PATCH 19/45] fix ci test --- .../sglang/srt/arg_groups/speculative_hook.py | 3 ++- python/sglang/srt/managers/io_struct.py | 1 + .../scheduler_components/output_streamer.py | 3 +++ .../sglang/srt/managers/tokenizer_manager.py | 22 ++++++------------- 4 files changed, 13 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index d749efefb0f9..d98e786ad5e6 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -393,7 +393,8 @@ def _handle_eagle_family(server_args: "ServerArgs") -> None: # topk > 1 + page_size > 1 needs the two-pass cascade draft-decode (shared prefix # pass + per-branch expand pass with prefix-tail dup). Only these backends implement # it; flashmla / trtllm_mla / cutlass_mla can't express the per-branch tree, so reject. - _PAGE_TREE_SPEC_BACKENDS = ("flashinfer", "fa3", "triton") + # TODO: @rainj-me, Temp allow fa4 to support the two-pass + _PAGE_TREE_SPEC_BACKENDS = ("flashinfer", "fa3", "triton", "fa4") if ( server_args.speculative_eagle_topk > 1 and server_args.page_size > 1 diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 12147152f47b..0faf38c10a64 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1347,6 +1347,7 @@ class BatchEmbeddingOutput(BaseBatchReqIpc, kw_only=True): # Sent as a single stacked tensor to minimize pickle overhead. # Optional[Union[List[Optional[torch.Tensor]], torch.Tensor]] pooled_hidden_states: Optional[List[Optional[torch.Tensor]]] = None + is_pooled_hidden_states_stacked: bool = False class ClearHiCacheReqInput(BaseReqIpc, kw_only=True): diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 06db448da3c0..079305b8809c 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -213,6 +213,7 @@ def _stream_output_embedding(self, reqs: List[Req]): # (some requests want PHS, others don't) keep the raw list so # positional indexing on the receiver side stays correct. stacked_phs = None + is_pooled_hidden_states_stacked = False if has_phs: all_have_phs = all(t is not None for t in phs_list) if all_have_phs: @@ -220,6 +221,7 @@ def _stream_output_embedding(self, reqs: List[Req]): # when stack all tensors, to keep the pooled_hidden_states data type # wrap the tensor with list to avoid the msgpack issue stacked_phs = [torch.stack(phs_list)] + is_pooled_hidden_states_stacked = True else: stacked_phs = phs_list else: @@ -239,6 +241,7 @@ def _stream_output_embedding(self, reqs: List[Req]): placeholder_tokens_val=None, retraction_counts=retraction_counts, pooled_hidden_states=stacked_phs, + is_pooled_hidden_states_stacked=is_pooled_hidden_states_stacked, ) ) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index d499c2c8e871..22b0fb997189 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1994,22 +1994,14 @@ async def _handle_batch_output( "embedding": recv_obj.embeddings[i], "meta_info": meta_info, } + + pooled_hidden_states = recv_obj.pooled_hidden_states if recv_obj.pooled_hidden_states is not None: - # for non-stacked, pooled_hidden_states is a list of tensors - if len(recv_obj.pooled_hidden_states) == len(recv_obj.rids): - if recv_obj.pooled_hidden_states[i] is not None: - out_dict["pooled_hidden_state"] = ( - recv_obj.pooled_hidden_states[i] - ) - # for stacked, pooled_hidden_states is a single tensor list - elif len(recv_obj.pooled_hidden_states) == 1: - if ( - recv_obj.pooled_hidden_states[0] is not None - and recv_obj.pooled_hidden_states[0][i] is not None - ): - out_dict["pooled_hidden_state"] = ( - recv_obj.pooled_hidden_states[0][i] - ) + if recv_obj.is_pooled_hidden_states_stacked: + # for stacked, pooled_hidden_states is a single tensor list + pooled_hidden_states = recv_obj.pooled_hidden_states[0] + if pooled_hidden_states[i] is not None: + out_dict["pooled_hidden_state"] = pooled_hidden_states[i] # Set first_token_time on the first output batch. # This is the single write point for first_token_time. From 88b7d0fcacb0177e82b4002e89f924c350d5d640 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Mon, 8 Jun 2026 10:49:46 -0700 Subject: [PATCH 20/45] fix ci test --- python/sglang/srt/entrypoints/http_server.py | 10 +++++++--- python/sglang/srt/managers/io_struct.py | 4 ++-- python/sglang/srt/managers/mm_utils.py | 6 ++++++ 3 files changed, 15 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 8b37bb278e99..27e6edbf25cc 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -42,6 +42,7 @@ setattr(threading, "_register_atexit", lambda *args, **kwargs: None) +import msgspec import numpy as np import requests import uvicorn @@ -1629,8 +1630,9 @@ async def load_lora_adapter(obj: LoadLoRAAdapterReqInput, request: Request): LoadLoRAAdapterReqInputIpc(**obj.__dict__), request, ) + result = msgspec.to_builtins(result) - if result.success: + if result.get("success", False): return ORJSONResponse( result, status_code=HTTPStatus.OK, @@ -1662,8 +1664,9 @@ async def load_lora_adapter_from_tensors( LoadLoRAAdapterFromTensorsReqInputIpc(**obj.__dict__), request, ) + result = msgspec.to_builtins(result) - if result.success: + if result.get("success", False): return ORJSONResponse(result, status_code=HTTPStatus.OK) else: return ORJSONResponse(result, status_code=HTTPStatus.BAD_REQUEST) @@ -1685,8 +1688,9 @@ async def unload_lora_adapter(obj: UnloadLoRAAdapterReqInput, request: Request): UnloadLoRAAdapterReqInputIpc(**obj.__dict__), request, ) + result = msgspec.to_builtins(result) - if result.success: + if result.get("success", False): return ORJSONResponse( result, status_code=HTTPStatus.OK, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 0faf38c10a64..eea3bba52dd7 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1075,7 +1075,7 @@ class TokenizedEmbeddingReqInput(BaseReqIpc, kw_only=True): # The input token ids input_ids: Optional[array] # array[int] # The image inputs - image_inputs: Optional[dict] + image_inputs: Optional[object] # The token type ids token_type_ids: Optional[List[int]] # Dummy sampling params for compatibility @@ -1954,7 +1954,7 @@ def to_ref(self) -> LoRARef: class LoRAUpdateOutput(BaseReqIpc, kw_only=True): success: bool error_message: Optional[str] = None - loaded_adapters: Optional[Dict[str, LoRARef]] = None + loaded_adapters: Optional[Dict[str, Union[str, LoRARef]]] = None LoadLoRAAdapterReqOutput = UnloadLoRAAdapterReqOutput = ( diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 841ecd1ea281..1fc824eca72f 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -1660,6 +1660,12 @@ def wrap_shm_features(obj): # large tensors inside mm_inputs obj.mm_inputs = pickle.dumps(obj.mm_inputs) + if hasattr(obj, "image_inputs"): + # Use pickle to serialize the image_inputs + # so that msgpack doesn't try to serialize the potentially + # large tensors inside mm_inputs + obj.image_inputs = pickle.dumps(obj.image_inputs) + return obj From f961a4e67720014d4d8ed879c1779ec88e9bf567 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Mon, 8 Jun 2026 11:16:20 -0700 Subject: [PATCH 21/45] lint --- .../srt/managers/scheduler_components/weight_updater.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index e65c6569bc57..c8945c1a8329 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py @@ -142,7 +142,9 @@ def update_weights_from_distributed( self.flush_cache_after_weight_update(recv_req) else: logger.error(message) - return UpdateWeightsFromDistributedReqOutput(success=success, message=message) + return UpdateWeightsFromDistributedReqOutput( + success=success, message=message + ) def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput): """Update the online model parameter from tensors.""" From 2427aad2aae7a16b6266f8bb80eac010806e79cd Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Mon, 8 Jun 2026 14:33:39 -0700 Subject: [PATCH 22/45] fix ci tests --- python/sglang/srt/entrypoints/http_server.py | 2 +- .../srt/managers/scheduler_components/weight_updater.py | 2 +- python/sglang/srt/managers/tokenizer_control_mixin.py | 7 +++++++ 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 27e6edbf25cc..16da66a6cea7 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -1362,7 +1362,7 @@ async def destroy_weights_update_group( @dataclasses.dataclass class UpdateWeightsFromTensorReqInput: # List[Union[str, bytes]] - serialized_named_tensors: List[bytes] + serialized_named_tensors: List[Union[str, bytes]] # Optional format specification for loading load_format: Optional[str] = None # Whether to flush the cache after updating weights diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index c8945c1a8329..82b71e7c3846 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py @@ -177,7 +177,7 @@ def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput): def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput): parameter = self.tp_worker.get_weights_by_name(recv_req) - return GetWeightsByNameReqOutput(parameter) + return GetWeightsByNameReqOutput(parameter=parameter) def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput): assert ( diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index 844c75a29587..a57688eba865 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -9,6 +9,7 @@ import fastapi import msgspec +import pybase64 from sglang.srt.managers.communicator import FanOutCommunicator from sglang.srt.managers.io_struct import ( @@ -475,6 +476,12 @@ async def update_weights_from_tensor( if obj.abort_all_requests: self.abort_request(abort_all=True) + serialized_named_tensors = [ + pybase64.b64decode(data, validate=True) if isinstance(data, str) else data + for data in obj.serialized_named_tensors + ] + obj.serialized_named_tensors = serialized_named_tensors + async with self.is_pause_cond: is_paused = self.is_pause if is_paused: From 28329d2a14a8ddef668d6cfbd25459091f04345f Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Mon, 8 Jun 2026 15:26:43 -0700 Subject: [PATCH 23/45] fix ci test --- python/sglang/srt/observability/req_time_stats.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index faa5b8932fa0..3921b986f1b5 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/python/sglang/srt/observability/req_time_stats.py @@ -228,7 +228,7 @@ class ReqTimeStatsBase: metrics_collector: Optional[object] = None # Union[TraceReqContext, TraceNullContext] trace_ctx: object = field(default_factory=TraceNullContext) - disagg_mode: DisaggregationMode = DisaggregationMode.NULL + disagg_mode: Optional[DisaggregationMode] = DisaggregationMode.NULL diff_realtime_monotonic: float = 0.0 @classmethod From 1521232ad3804d2fd75e1d7a6e5f26b757bff878 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Mon, 8 Jun 2026 20:49:47 -0700 Subject: [PATCH 24/45] fix ci test --- python/sglang/srt/managers/io_struct.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index ef581d631448..cae75e9213a1 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1522,7 +1522,7 @@ class UpdateWeightFromDiskReqOutput(BaseReqIpc, kw_only=True): class UpdateWeightsFromDistributedReqInput(BaseReqIpc, kw_only=True): names: List[str] - dtypes: List[str] + dtypes: List[bytes] shapes: List[List[int]] # The group name group_name: str = "weight_update_group" @@ -1638,7 +1638,7 @@ class InitWeightsUpdateGroupReqInput(BaseReqIpc, kw_only=True): # The master address master_address: str # The master port - master_port: int + master_port: Union[int, str] # The rank offset rank_offset: int # The world size From ebe1414d956937381fdaf3a1536144b2121e185e Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Tue, 9 Jun 2026 16:32:45 -0700 Subject: [PATCH 25/45] fix ci test --- python/sglang/srt/managers/io_struct.py | 2 +- test/registered/rl/test_update_weights_from_distributed.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index cae75e9213a1..e7a1b8fa9db2 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1522,7 +1522,7 @@ class UpdateWeightFromDiskReqOutput(BaseReqIpc, kw_only=True): class UpdateWeightsFromDistributedReqInput(BaseReqIpc, kw_only=True): names: List[str] - dtypes: List[bytes] + dtypes: List[str] shapes: List[List[int]] # The group name group_name: str = "weight_update_group" diff --git a/test/registered/rl/test_update_weights_from_distributed.py b/test/registered/rl/test_update_weights_from_distributed.py index fe7e593a3c47..9738416c391b 100644 --- a/test/registered/rl/test_update_weights_from_distributed.py +++ b/test/registered/rl/test_update_weights_from_distributed.py @@ -411,7 +411,7 @@ def run_decode(max_new_tokens=32): # Get weights from the training engine and update the inference engine. names = [parameter_name for parameter_name in update_parameters] - dtypes = [torch.bfloat16 if backend == "Engine" else "bfloat16"] * len(names) + dtypes = ["bfloat16"] * len(names) shapes = [state_dict_key_to_shape[parameter_name] for parameter_name in names] if pause_generation_mode in ["in_place", "retract"]: From 00d35548c30a58ac5925d3b59d01c634120d8d98 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Thu, 11 Jun 2026 21:16:54 -0700 Subject: [PATCH 26/45] address comments --- .../sglang/srt/arg_groups/speculative_hook.py | 2 +- python/sglang/srt/debug_utils/dumper.py | 9 +- .../srt/disaggregation/encode_grpc_server.py | 3 +- .../srt/disaggregation/encode_server.py | 4 +- python/sglang/srt/entrypoints/api_specs.py | 340 ++++++++++++ python/sglang/srt/entrypoints/engine.py | 11 +- python/sglang/srt/entrypoints/http_server.py | 505 ++++-------------- python/sglang/srt/managers/communicator.py | 18 +- python/sglang/srt/managers/io_struct.py | 118 ++-- python/sglang/srt/sampling/sampling_params.py | 141 ++--- .../managers/test_profile_merger_http_api.py | 28 +- .../unit/utils/test_profile_merger.py | 9 +- 12 files changed, 581 insertions(+), 607 deletions(-) create mode 100644 python/sglang/srt/entrypoints/api_specs.py diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 61caa3b75de4..33241131969d 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -372,7 +372,7 @@ def _handle_eagle_family(server_args: "ServerArgs") -> None: # pass + per-branch expand pass with prefix-tail dup). Only these backends implement # it; flashmla / trtllm_mla / cutlass_mla can't express the per-branch tree, so reject. # TODO: @rainj-me, Temp allow fa4 to support the two-pass - _PAGE_TREE_SPEC_BACKENDS = ("flashinfer", "fa3", "triton", "fa4") + _PAGE_TREE_SPEC_BACKENDS = ("flashinfer", "fa3", "triton") if ( server_args.speculative_eagle_topk > 1 and server_args.page_size > 1 diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 21a73cabbc7e..f3f48c24981d 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -20,6 +20,9 @@ import torch import torch.distributed as dist +import zmq + +from sglang.srt.managers.io_struct import sock_recv, sock_send # -------------------------------------- config base ------------------------------------------ @@ -1419,10 +1422,6 @@ def _create_zmq_rpc_broadcast( handler, timeout_seconds: int = 60 ) -> Optional["_ZmqRpcBroadcast"]: """A general-purpose minimal RPC to support broadcasting executions to multi processes""" - import zmq - - from sglang.srt.managers.io_struct import sock_recv, sock_send - rank = _get_rank() world_size = dist.get_world_size() if dist.is_initialized() else 1 @@ -1478,8 +1477,6 @@ def __init__(self, socket, debug_name: str): def __getattr__(self, method_name: str): def call(*args, **kwargs): - from sglang.srt.managers.io_struct import sock_recv, sock_send - sock_send( self._socket, { diff --git a/python/sglang/srt/disaggregation/encode_grpc_server.py b/python/sglang/srt/disaggregation/encode_grpc_server.py index 2e43ad80eff4..71c3d9729a94 100644 --- a/python/sglang/srt/disaggregation/encode_grpc_server.py +++ b/python/sglang/srt/disaggregation/encode_grpc_server.py @@ -26,6 +26,7 @@ handle_scheduler_receive_url_request, launch_encoder, ) +from sglang.srt.managers.io_struct import async_sock_send from sglang.srt.managers.schedule_batch import Modality from sglang.srt.server_args import PortArgs, ServerArgs from sglang.srt.utils import random_uuid @@ -87,8 +88,6 @@ def __init__( async def Encode( self, request: sglang_encoder_pb2.EncodeRequest, context ) -> sglang_encoder_pb2.EncodeResponse: - from sglang.srt.managers.io_struct import async_sock_send - try: request_dict = { "mm_items": list(request.mm_items), diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index 9c71edcd85c6..3c358ee699e4 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -41,11 +41,11 @@ init_distributed_environment, initialize_model_parallel, ) +from sglang.srt.entrypoints.api_specs import ProfileReqInputSpec from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import initialize_dp_attention from sglang.srt.managers.io_struct import ( ProfileReq, - ProfileReqInput, ProfileReqType, async_sock_recv, sock_send, @@ -3740,7 +3740,7 @@ async def health_generate(): @app.api_route("/start_profile", methods=["GET", "POST"]) -async def start_profile_async(obj: Optional[ProfileReqInput] = None): +async def start_profile_async(obj: Optional[ProfileReqInputSpec] = None): if dp_dispatcher is not None: profile_req = None if obj is not None: diff --git a/python/sglang/srt/entrypoints/api_specs.py b/python/sglang/srt/entrypoints/api_specs.py new file mode 100644 index 000000000000..7737d400af9f --- /dev/null +++ b/python/sglang/srt/entrypoints/api_specs.py @@ -0,0 +1,340 @@ +from dataclasses import dataclass, field +from typing import ( + Any, + Dict, + List, + Literal, + Optional, + Union, +) + +from sglang.srt.entrypoints.openai.protocol import Tool + +# All spec classes used for APIs are defined here. +# They are extracted from the io_struct.py. + + +@dataclass +class AttachHiCacheStorageReqInputSpec: + """Dynamically attach (enable) HiCache storage backend at runtime. + + Note: `hicache_storage_backend_extra_config_json` is a JSON string. It may contain both: + - backend-specific configs (e.g., mooncake master address) + - prefetch-related knobs (prefetch_threshold, prefetch_timeout_*, hicache_storage_pass_prefix_keys) + """ + + hicache_storage_backend: str + hicache_storage_backend_extra_config_json: Optional[str] = None + hicache_storage_prefetch_policy: Optional[str] = None + hicache_write_policy: Optional[str] = None + + def __post_init__(self): + if self.hicache_storage_prefetch_policy is None: + pass + else: + allowed = ["best_effort", "wait_complete", "timeout"] + if self.hicache_storage_prefetch_policy not in allowed: + raise ValueError( + f"Invalid hicache_storage_prefetch_policy: {self.hicache_storage_prefetch_policy!r}. " + f"Expected one of {allowed}." + ) + + if self.hicache_write_policy is None: + return + allowed = ["write_back", "write_through", "write_through_selective"] + if self.hicache_write_policy not in allowed: + raise ValueError( + f"Invalid hicache_write_policy: {self.hicache_write_policy!r}. " + f"Expected one of {allowed}." + ) + + +@dataclass +class UpdateWeightFromDiskReqInputSpec: + # The model path with the new weights + model_path: str + # The format to load the weights + load_format: Optional[str] = None + # Whether to abort all requests before updating weights + abort_all_requests: bool = False + # Optional: Update weight version along with weights + weight_version: Optional[str] = None + # Whether to update weights asynchronously + is_async: bool = False + # Whether to call torch.cuda.empty_cache() during flush + torch_empty_cache: bool = False + # Whether to keep the scheduler paused after weight update + keep_pause: bool = False + # Whether to recapture cuda graph after weight update + recapture_cuda_graph: bool = False + # The trainer step id. Used to know which step's weights are used for sampling. + token_step: int = 0 + # Whether to flush the cache after updating weights + flush_cache: bool = True + # Tensor metadata + manifest: Optional[Dict[str, Any]] = None + + +@dataclass +class InitWeightsSendGroupForRemoteInstanceReqInputSpec: + # The master address + master_address: str + # The ports for each rank's communication group + ports: str + # The rank in the communication group + group_rank: int + # The world size + world_size: int + # The group name + group_name: str = "weight_send_group" + # The backend + backend: str = "nccl" + + +@dataclass +class InitWeightsUpdateGroupReqInputSpec: + # The master address + master_address: str + # The master port + master_port: int + # The rank offset + rank_offset: int + # The world size + world_size: int + # The group name + group_name: str = "weight_update_group" + # The backend + backend: str = "nccl" + + +@dataclass +class SendWeightsToRemoteInstanceReqInputSpec: + # The master address + master_address: str + # The ports for each rank's communication group + ports: str + # The group name + group_name: str = "weight_send_group" + + +@dataclass +class DestroyWeightsUpdateGroupReqInputSpec: + group_name: str = "weight_update_group" + + +@dataclass +class UpdateWeightsFromTensorReqInputSpec: + # List[Union[str, bytes]] + serialized_named_tensors: List[Union[str, bytes]] + # Optional format specification for loading + load_format: Optional[str] = None + # Whether to flush the cache after updating weights + flush_cache: bool = True + # Whether to abort all requests before updating weights + abort_all_requests: bool = False + # Optional: Update weight version along with weights + weight_version: Optional[str] = None + # Optional: Determine whether to disable updating the draft model + disable_draft_model: Optional[bool] = None + # Whether to call torch.cuda.empty_cache() during flush + torch_empty_cache: bool = False + + +@dataclass +class UpdateWeightsFromDistributedReqInputSpec: + names: List[str] + dtypes: List[str] + shapes: List[List[int]] + # The group name + group_name: str = "weight_update_group" + # Whether to flush the cache after updating weights + flush_cache: bool = True + # Whether to abort all requests before updating weights + abort_all_requests: bool = False + # Optional: Update weight version along with weights + weight_version: Optional[str] = None + # Optional format specification for loading + load_format: Optional[str] = None + # Whether to call torch.cuda.empty_cache() during flush + torch_empty_cache: bool = False + + +@dataclass +class UpdateWeightsFromIPCReqInputSpec: + # ZMQ socket paths for each device UUID + zmq_handles: Dict[str, str] + # Whether to flush cache after weight update + flush_cache: bool = True + # Optional: Update weight version along with weights + weight_version: Optional[str] = None + # Whether to call torch.cuda.empty_cache() during flush + torch_empty_cache: bool = False + + +@dataclass +class UpdateWeightVersionReqInputSpec: + # The new weight version + new_version: str + # Whether to abort all running requests before updating + abort_all_requests: bool = True + + +@dataclass +class GetWeightsByNameReqInputSpec: + name: str + truncate_size: int = 100 + + +@dataclass +class SetInternalStateReqSpec: + server_args: Dict[str, Any] + + +@dataclass +class ReleaseMemoryOccupationReqInputSpec: + tags: Optional[List[str]] = None + + +@dataclass +class ResumeMemoryOccupationReqInputSpec: + tags: Optional[List[str]] = None + + +@dataclass +class CheckWeightsReqInputSpec: + action: str = "checksum" + + +@dataclass +class SlowDownReqInputSpec: + forward_sleep_time: Optional[float] + + +@dataclass +class LoadLoRAAdapterReqInputSpec: + # The name of the lora module to newly loaded. + lora_name: str + # The path of loading. + lora_path: str + # Whether to pin the LoRA adapter in memory. + pinned: bool = False + # The unique identifier for the LoRA adapter, which automatically generated in the `TokenizerManager`. + lora_id: Optional[str] = None + + +@dataclass +class LoadLoRAAdapterFromTensorsReqInputSpec: + lora_name: str + config_dict: Dict[str, Any] + serialized_tensors: str + pinned: bool = False + added_tokens_config: Optional[Dict[str, Any]] = None + lora_id: Optional[str] = None + load_format: Optional[str] = None + + +@dataclass +class UnloadLoRAAdapterReqInputSpec: + # The name of lora module to unload. + lora_name: str + # The unique identifier for the LoRA adapter, which automatically generated in the `TokenizerManager`. + lora_id: Optional[str] = None + + +@dataclass +class OpenSessionReqInputSpec: + capacity_of_str_len: int + session_id: Optional[str] = None + streaming: Optional[bool] = None + timeout: Optional[float] = None + + +@dataclass +class CloseSessionReqInputSpec: + session_id: str + + +@dataclass +class ConfigureLoggingReqSpec: + log_requests: Optional[bool] = None + log_requests_level: Optional[int] = None + log_requests_format: Optional[str] = None + log_level: Optional[str] = None + dump_requests_folder: Optional[str] = None + dump_requests_threshold: Optional[int] = None + crash_dump_folder: Optional[str] = None + dump_requests_exclude_meta_keys: Optional[List[str]] = None + + +@dataclass +class AbortReqSpec: + rid: Optional[str] = None + # Whether to abort all requests + abort_all: bool = False + + +@dataclass +class PauseGenerationReqInputSpec: + mode: Literal["abort", "retract", "in_place"] = "abort" + + def __post_init__(self): + allowed = ["abort", "retract", "in_place"] + if self.mode not in allowed: + raise ValueError( + f"Invalid mode: {self.mode!r}. " f"Expected one of {allowed}." + ) + + +@dataclass +class ContinueGenerationReqInputSpec: + torch_empty_cache: bool = True + + +@dataclass +class SeparateReasoningReqInputSpec: + 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 VertexGenerateReqInputSpec: + instances: List[dict] + parameters: Optional[dict] = None + + +@dataclass +class ParseFunctionCallReqSpec: + text: str # The text to parse. + tools: List[Tool] = field( + default_factory=list + ) # A list of available function tools (name, parameters, etc.). + tool_call_parser: Optional[str] = ( + None # Specify the parser type, e.g. 'llama3', 'qwen25', or 'mistral'. If not specified, tries all. + ) + + +@dataclass +class ProfileReqInputSpec: + # The output directory + output_dir: Optional[str] = None + # Specify the steps to start the profiling + start_step: Optional[int] = None + # If set, it profile as many as this number of steps. + # If it is set, profiling is automatically stopped after this step, and + # the caller doesn't need to run stop_profile. + num_steps: Optional[int] = None + # The activities to record. The choices are ["CPU", "GPU", "MEM", "RPD"] + activities: Optional[List[str]] = None + # Whether profile by stages (e.g., prefill and decode) separately + profile_by_stage: bool = False + # Whether to record source information (file and line number) for the ops. + with_stack: Optional[bool] = None + # Whether to save information about operator’s input shapes. + record_shapes: Optional[bool] = None + # Merge profiles from all ranks into a single trace + merge_profiles: bool = False + # The prefix of the profile filenames + profile_prefix: Optional[str] = None + # Only profile these stages and ignore others + profile_stages: Optional[List[str]] = None diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index c49e8137ffb8..e3d31f636395 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -863,14 +863,15 @@ def _launch_subprocesses( for p in detoken_procs: scheduler_init_result.all_child_pids.append(p.pid) - if server_args.tokenizer_worker_num > 1 or server_args.enable_http2: - # Launch multi-tokenizer router - tokenizer_manager = MultiTokenizerRouter(server_args, port_args) - template_manager = None - else: + if server_args.tokenizer_worker_num == 1 and not server_args.enable_http2: + # Launch a single tokenizer manager without the router tokenizer_manager, template_manager = init_tokenizer_manager_func( server_args, port_args ) + else: + # Launch multi-tokenizer router + tokenizer_manager = MultiTokenizerRouter(server_args, port_args) + template_manager = None # Wait for the model to finish loading scheduler_init_result.wait_for_ready() diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 16da66a6cea7..57a803118f78 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -33,7 +33,6 @@ Callable, Dict, List, - Literal, Optional, Union, ) @@ -69,6 +68,37 @@ AnthropicMessagesRequest, ) from sglang.srt.entrypoints.anthropic.serving import AnthropicServing +from sglang.srt.entrypoints.api_specs import ( + AbortReqSpec, + AttachHiCacheStorageReqInputSpec, + CheckWeightsReqInputSpec, + CloseSessionReqInputSpec, + ConfigureLoggingReqSpec, + ContinueGenerationReqInputSpec, + DestroyWeightsUpdateGroupReqInputSpec, + GetWeightsByNameReqInputSpec, + InitWeightsSendGroupForRemoteInstanceReqInputSpec, + InitWeightsUpdateGroupReqInputSpec, + LoadLoRAAdapterFromTensorsReqInputSpec, + LoadLoRAAdapterReqInputSpec, + OpenSessionReqInputSpec, + ParseFunctionCallReqSpec, + PauseGenerationReqInputSpec, + ProfileReqInputSpec, + ReleaseMemoryOccupationReqInputSpec, + ResumeMemoryOccupationReqInputSpec, + SendWeightsToRemoteInstanceReqInputSpec, + SeparateReasoningReqInputSpec, + SetInternalStateReqSpec, + SlowDownReqInputSpec, + UnloadLoRAAdapterReqInputSpec, + UpdateWeightFromDiskReqInputSpec, + UpdateWeightsFromDistributedReqInputSpec, + UpdateWeightsFromIPCReqInputSpec, + UpdateWeightsFromTensorReqInputSpec, + UpdateWeightVersionReqInputSpec, + VertexGenerateReqInputSpec, +) from sglang.srt.entrypoints.engine import ( Engine, init_tokenizer_manager, @@ -111,80 +141,31 @@ from sglang.srt.environ import envs from sglang.srt.function_call.function_call_parser import FunctionCallParser from sglang.srt.managers.io_struct import ( - CheckWeightsReqInput as CheckWeightsReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - CloseSessionReqInput as CloseSessionReqInputIpc, -) -from sglang.srt.managers.io_struct import ConfigureLoggingReq as ConfigureLoggingReqIpc -from sglang.srt.managers.io_struct import ( - ContinueGenerationReqInput as ContinueGenerationReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - DestroyWeightsUpdateGroupReqInput as DestroyWeightsUpdateGroupReqInputIpc, -) -from sglang.srt.managers.io_struct import ( + CheckWeightsReqInput, + CloseSessionReqInput, + ConfigureLoggingReq, + ContinueGenerationReqInput, + DestroyWeightsUpdateGroupReqInput, DumperControlReqInput, EmbeddingReqInput, GenerateReqInput, -) -from sglang.srt.managers.io_struct import ( - GetWeightsByNameReqInput as GetWeightsByNameReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - InitWeightsSendGroupForRemoteInstanceReqInput as InitWeightsSendGroupForRemoteInstanceReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - InitWeightsUpdateGroupReqInput as InitWeightsUpdateGroupReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - LoadLoRAAdapterFromTensorsReqInput as LoadLoRAAdapterFromTensorsReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - LoadLoRAAdapterReqInput as LoadLoRAAdapterReqInputIpc, -) -from sglang.srt.managers.io_struct import OpenSessionReqInput as OpenSessionReqInputIpc -from sglang.srt.managers.io_struct import ( - ParseFunctionCallReq, -) -from sglang.srt.managers.io_struct import ( - PauseGenerationReqInput as PauseGenerationReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - ProfileReqInput, -) -from sglang.srt.managers.io_struct import ( - ReleaseMemoryOccupationReqInput as ReleaseMemoryOccupationReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - ResumeMemoryOccupationReqInput as ResumeMemoryOccupationReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - SendWeightsToRemoteInstanceReqInput as SendWeightsToRemoteInstanceReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - SeparateReasoningReqInput, -) -from sglang.srt.managers.io_struct import SetInternalStateReq as SetInternalStateReqIpc -from sglang.srt.managers.io_struct import SlowDownReqInput as SlowDownReqInputIpc -from sglang.srt.managers.io_struct import ( - UnloadLoRAAdapterReqInput as UnloadLoRAAdapterReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - UpdateWeightFromDiskReqInput as UpdateWeightFromDiskReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - UpdateWeightsFromDistributedReqInput as UpdateWeightsFromDistributedReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - UpdateWeightsFromIPCReqInput as UpdateWeightsFromIPCReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - UpdateWeightsFromTensorReqInput as UpdateWeightsFromTensorReqInputIpc, -) -from sglang.srt.managers.io_struct import ( - UpdateWeightVersionReqInput, - VertexGenerateReqInput, + GetWeightsByNameReqInput, + InitWeightsSendGroupForRemoteInstanceReqInput, + InitWeightsUpdateGroupReqInput, + LoadLoRAAdapterFromTensorsReqInput, + LoadLoRAAdapterReqInput, + OpenSessionReqInput, + PauseGenerationReqInput, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, + SendWeightsToRemoteInstanceReqInput, + SetInternalStateReq, + SlowDownReqInput, + UnloadLoRAAdapterReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromIPCReqInput, + UpdateWeightsFromTensorReqInput, ) from sglang.srt.managers.multi_tokenizer_mixin import ( MultiTokenizerRouter, @@ -728,18 +709,13 @@ async def get_load(): ] -@dataclasses.dataclass -class SetInternalStateReq: - server_args: Dict[str, Any] - - # example usage: # curl -s -X POST http://localhost:30000/set_internal_state -H "Content-Type: application/json" -d '{"server_args": {"pp_max_micro_batch_size": 8}}' @app.api_route("/set_internal_state", methods=["POST", "PUT"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def set_internal_state(obj: SetInternalStateReq, request: Request): +async def set_internal_state(obj: SetInternalStateReqSpec, request: Request): res = await _global_state.tokenizer_manager.set_internal_state( - SetInternalStateReqIpc(**obj.__dict__) + SetInternalStateReq(**obj.__dict__) ) return res @@ -939,41 +915,6 @@ async def clear_hicache_storage_backend(): ) -@dataclasses.dataclass -class AttachHiCacheStorageReqInput: - """Dynamically attach (enable) HiCache storage backend at runtime. - - Note: `hicache_storage_backend_extra_config_json` is a JSON string. It may contain both: - - backend-specific configs (e.g., mooncake master address) - - prefetch-related knobs (prefetch_threshold, prefetch_timeout_*, hicache_storage_pass_prefix_keys) - """ - - hicache_storage_backend: str - hicache_storage_backend_extra_config_json: Optional[str] = None - hicache_storage_prefetch_policy: Optional[str] = None - hicache_write_policy: Optional[str] = None - - def __post_init__(self): - if self.hicache_storage_prefetch_policy is None: - pass - else: - allowed = ["best_effort", "wait_complete", "timeout"] - if self.hicache_storage_prefetch_policy not in allowed: - raise ValueError( - f"Invalid hicache_storage_prefetch_policy: {self.hicache_storage_prefetch_policy!r}. " - f"Expected one of {allowed}." - ) - - if self.hicache_write_policy is None: - return - allowed = ["write_back", "write_through", "write_through_selective"] - if self.hicache_write_policy not in allowed: - raise ValueError( - f"Invalid hicache_write_policy: {self.hicache_write_policy!r}. " - f"Expected one of {allowed}." - ) - - # example usage: # curl -s -X PUT http://127.0.0.1:30000/hicache/storage-backend \ # -H 'Content-Type: application/json' \ @@ -985,7 +926,7 @@ def __post_init__(self): # }' @app.api_route("/hicache/storage-backend", methods=["PUT"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def attach_hicache_storage_backend(obj: AttachHiCacheStorageReqInput): +async def attach_hicache_storage_backend(obj: AttachHiCacheStorageReqInputSpec): """Attach (enable) HiCache storage backend at runtime. Only allowed when there are NO running / queued requests. @@ -1059,10 +1000,10 @@ async def hicache_storage_backend_status(): @app.api_route("/start_profile", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def start_profile_async(obj: Optional[ProfileReqInput] = None): +async def start_profile_async(obj: Optional[ProfileReqInputSpec] = None): """Start profiling.""" if obj is None: - obj = ProfileReqInput() + obj = ProfileReqInputSpec() await _global_state.tokenizer_manager.start_profile( output_dir=obj.output_dir, @@ -1149,39 +1090,15 @@ async def dump_expert_distribution_record_async(): ) -@dataclasses.dataclass -class UpdateWeightFromDiskReqInput: - # The model path with the new weights - model_path: str - # The format to load the weights - load_format: Optional[str] = None - # Whether to abort all requests before updating weights - abort_all_requests: bool = False - # Optional: Update weight version along with weights - weight_version: Optional[str] = None - # Whether to update weights asynchronously - is_async: bool = False - # Whether to call torch.cuda.empty_cache() during flush - torch_empty_cache: bool = False - # Whether to keep the scheduler paused after weight update - keep_pause: bool = False - # Whether to recapture cuda graph after weight update - recapture_cuda_graph: bool = False - # The trainer step id. Used to know which step's weights are used for sampling. - token_step: int = 0 - # Whether to flush the cache after updating weights - flush_cache: bool = True - # Tensor metadata - manifest: Optional[Dict[str, Any]] = None - - @app.post("/update_weights_from_disk") @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def update_weights_from_disk(obj: UpdateWeightFromDiskReqInput, request: Request): +async def update_weights_from_disk( + obj: UpdateWeightFromDiskReqInputSpec, request: Request +): """Update the weights from disk inplace without re-launching the server.""" success, message, num_paused_requests = ( await _global_state.tokenizer_manager.update_weights_from_disk( - UpdateWeightFromDiskReqInputIpc(**obj.__dict__), + UpdateWeightFromDiskReqInput(**obj.__dict__), request, ) ) @@ -1203,30 +1120,14 @@ async def update_weights_from_disk(obj: UpdateWeightFromDiskReqInput, request: R ) -@dataclasses.dataclass -class InitWeightsSendGroupForRemoteInstanceReqInput: - # The master address - master_address: str - # The ports for each rank's communication group - ports: str - # The rank in the communication group - group_rank: int - # The world size - world_size: int - # The group name - group_name: str = "weight_send_group" - # The backend - backend: str = "nccl" - - @app.post("/init_weights_send_group_for_remote_instance") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def init_weights_send_group_for_remote_instance( - obj: InitWeightsSendGroupForRemoteInstanceReqInput, request: Request + obj: InitWeightsSendGroupForRemoteInstanceReqInputSpec, request: Request ): success, message = ( await _global_state.tokenizer_manager.init_weights_send_group_for_remote_instance( - InitWeightsSendGroupForRemoteInstanceReqInputIpc(**obj.__dict__), + InitWeightsSendGroupForRemoteInstanceReqInput(**obj.__dict__), request, ) ) @@ -1237,24 +1138,14 @@ async def init_weights_send_group_for_remote_instance( return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) -@dataclasses.dataclass -class SendWeightsToRemoteInstanceReqInput: - # The master address - master_address: str - # The ports for each rank's communication group - ports: str - # The group name - group_name: str = "weight_send_group" - - @app.post("/send_weights_to_remote_instance") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def send_weights_to_remote_instance( - obj: SendWeightsToRemoteInstanceReqInput, request: Request + obj: SendWeightsToRemoteInstanceReqInputSpec, request: Request ): success, message = ( await _global_state.tokenizer_manager.send_weights_to_remote_instance( - SendWeightsToRemoteInstanceReqInputIpc(**obj.__dict__), + SendWeightsToRemoteInstanceReqInput(**obj.__dict__), request, ) ) @@ -1303,30 +1194,14 @@ async def remote_instance_transfer_engine_info(rank: int = None): ) -@dataclasses.dataclass -class InitWeightsUpdateGroupReqInput: - # The master address - master_address: str - # The master port - master_port: int - # The rank offset - rank_offset: int - # The world size - world_size: int - # The group name - group_name: str = "weight_update_group" - # The backend - backend: str = "nccl" - - @app.post("/init_weights_update_group") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def init_weights_update_group( - obj: InitWeightsUpdateGroupReqInput, request: Request + obj: InitWeightsUpdateGroupReqInputSpec, request: Request ): """Initialize the parameter update group.""" success, message = await _global_state.tokenizer_manager.init_weights_update_group( - InitWeightsUpdateGroupReqInputIpc(**obj.__dict__), + InitWeightsUpdateGroupReqInput(**obj.__dict__), request, ) content = {"success": success, "message": message} @@ -1336,20 +1211,15 @@ async def init_weights_update_group( return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) -@dataclasses.dataclass -class DestroyWeightsUpdateGroupReqInput: - group_name: str = "weight_update_group" - - @app.post("/destroy_weights_update_group") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def destroy_weights_update_group( - obj: DestroyWeightsUpdateGroupReqInput, request: Request + obj: DestroyWeightsUpdateGroupReqInputSpec, request: Request ): """Destroy the parameter update group.""" success, message = ( await _global_state.tokenizer_manager.destroy_weights_update_group( - DestroyWeightsUpdateGroupReqInputIpc(**obj.__dict__), + DestroyWeightsUpdateGroupReqInput(**obj.__dict__), request, ) ) @@ -1359,28 +1229,10 @@ async def destroy_weights_update_group( ) -@dataclasses.dataclass -class UpdateWeightsFromTensorReqInput: - # List[Union[str, bytes]] - serialized_named_tensors: List[Union[str, bytes]] - # Optional format specification for loading - load_format: Optional[str] = None - # Whether to flush the cache after updating weights - flush_cache: bool = True - # Whether to abort all requests before updating weights - abort_all_requests: bool = False - # Optional: Update weight version along with weights - weight_version: Optional[str] = None - # Optional: Determine whether to disable updating the draft model - disable_draft_model: Optional[bool] = None - # Whether to call torch.cuda.empty_cache() during flush - torch_empty_cache: bool = False - - @app.post("/update_weights_from_tensor") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weights_from_tensor( - obj: UpdateWeightsFromTensorReqInput, request: Request + obj: UpdateWeightsFromTensorReqInputSpec, request: Request ): """Update the weights from tensor inplace without re-launching the server. Notes: @@ -1390,7 +1242,7 @@ async def update_weights_from_tensor( """ success, message = await _global_state.tokenizer_manager.update_weights_from_tensor( - UpdateWeightsFromTensorReqInputIpc(**obj.__dict__), request + UpdateWeightsFromTensorReqInput(**obj.__dict__), request ) content = {"success": success, "message": message} @@ -1399,34 +1251,15 @@ async def update_weights_from_tensor( ) -@dataclasses.dataclass -class UpdateWeightsFromDistributedReqInput: - names: List[str] - dtypes: List[str] - shapes: List[List[int]] - # The group name - group_name: str = "weight_update_group" - # Whether to flush the cache after updating weights - flush_cache: bool = True - # Whether to abort all requests before updating weights - abort_all_requests: bool = False - # Optional: Update weight version along with weights - weight_version: Optional[str] = None - # Optional format specification for loading - load_format: Optional[str] = None - # Whether to call torch.cuda.empty_cache() during flush - torch_empty_cache: bool = False - - @app.post("/update_weights_from_distributed") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weights_from_distributed( - obj: UpdateWeightsFromDistributedReqInput, request: Request + obj: UpdateWeightsFromDistributedReqInputSpec, request: Request ): """Update model parameter from distributed online.""" success, message = ( await _global_state.tokenizer_manager.update_weights_from_distributed( - UpdateWeightsFromDistributedReqInputIpc(**obj.__dict__), + UpdateWeightsFromDistributedReqInput(**obj.__dict__), request, ) ) @@ -1438,24 +1271,14 @@ async def update_weights_from_distributed( return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) -@dataclasses.dataclass -class UpdateWeightsFromIPCReqInput: - # ZMQ socket paths for each device UUID - zmq_handles: Dict[str, str] - # Whether to flush cache after weight update - flush_cache: bool = True - # Optional: Update weight version along with weights - weight_version: Optional[str] = None - # Whether to call torch.cuda.empty_cache() during flush - torch_empty_cache: bool = False - - @app.post("/update_weights_from_ipc") @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Request): +async def update_weights_from_ipc( + obj: UpdateWeightsFromIPCReqInputSpec, request: Request +): """Update the weights from IPC (Inter-Process Communication) for checkpoint-engine integration.""" success, message = await _global_state.tokenizer_manager.update_weights_from_ipc( - UpdateWeightsFromIPCReqInputIpc(**obj.__dict__), + UpdateWeightsFromIPCReqInput(**obj.__dict__), request, ) @@ -1470,7 +1293,7 @@ async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Re @app.post("/update_weight_version") @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def update_weight_version(obj: UpdateWeightVersionReqInput, request: Request): +async def update_weight_version(obj: UpdateWeightVersionReqInputSpec, request: Request): """Update the weight version. This operation requires no active requests.""" if obj.abort_all_requests: _global_state.tokenizer_manager.abort_request(abort_all=True) @@ -1499,19 +1322,13 @@ async def update_weight_version(obj: UpdateWeightVersionReqInput, request: Reque ) -@dataclasses.dataclass -class GetWeightsByNameReqInput: - name: str - truncate_size: int = 100 - - @app.api_route("/get_weights_by_name", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def get_weights_by_name(obj: GetWeightsByNameReqInput, request: Request): +async def get_weights_by_name(obj: GetWeightsByNameReqInputSpec, request: Request): """Get model parameter by name.""" try: ret = await _global_state.tokenizer_manager.get_weights_by_name( - GetWeightsByNameReqInputIpc(**obj.__dict__), + GetWeightsByNameReqInput(**obj.__dict__), request, ) if ret is None: @@ -1522,61 +1339,46 @@ async def get_weights_by_name(obj: GetWeightsByNameReqInput, request: Request): return _create_error_response(e) -@dataclasses.dataclass -class ReleaseMemoryOccupationReqInput: - tags: Optional[List[str]] = None - - @app.api_route("/release_memory_occupation", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) async def release_memory_occupation( - obj: ReleaseMemoryOccupationReqInput, request: Request + obj: ReleaseMemoryOccupationReqInputSpec, request: Request ): """Release GPU memory occupation temporarily.""" try: await _global_state.tokenizer_manager.release_memory_occupation( - ReleaseMemoryOccupationReqInputIpc(**obj.__dict__), + ReleaseMemoryOccupationReqInput(**obj.__dict__), request, ) except Exception as e: return _create_error_response(e) -@dataclasses.dataclass -class ResumeMemoryOccupationReqInput: - tags: Optional[List[str]] = None - - @app.api_route("/resume_memory_occupation", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) async def resume_memory_occupation( - obj: ResumeMemoryOccupationReqInput, request: Request + obj: ResumeMemoryOccupationReqInputSpec, request: Request ): """Resume GPU memory occupation.""" try: await _global_state.tokenizer_manager.resume_memory_occupation( - ResumeMemoryOccupationReqInputIpc(**obj.__dict__), + ResumeMemoryOccupationReqInput(**obj.__dict__), request, ) except Exception as e: return _create_error_response(e) -@dataclasses.dataclass -class CheckWeightsReqInput: - action: str = "checksum" - - @app.api_route("/weights_checker", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) async def check_weights( - obj: Optional[CheckWeightsReqInput] = None, request: Request = None + obj: Optional[CheckWeightsReqInputSpec] = None, request: Request = None ): if obj is None: - obj = CheckWeightsReqInput() + obj = CheckWeightsReqInputSpec() success, message, ranks, per_engine_checksum = ( await _global_state.tokenizer_manager.check_weights( - CheckWeightsReqInputIpc(**obj.__dict__), + CheckWeightsReqInput(**obj.__dict__), request, ) ) @@ -1588,14 +1390,9 @@ async def check_weights( return ORJSONResponse(body, status_code=200 if success else HTTPStatus.BAD_REQUEST) -@dataclasses.dataclass -class SlowDownReqInput: - forward_sleep_time: Optional[float] - - @app.api_route("/slow_down", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def slow_down(obj: SlowDownReqInput, request: Request): +async def slow_down(obj: SlowDownReqInputSpec, request: Request): """Slow down the system deliberately. Only for testing. Example scenario: when we want to test performance of D in large-scale PD disaggregation and have no enough nodes for P, we can use this to slow down D to let it have enough running sequences, and then disable slowdown @@ -1603,31 +1400,19 @@ async def slow_down(obj: SlowDownReqInput, request: Request): """ try: await _global_state.tokenizer_manager.slow_down( - SlowDownReqInputIpc(**obj.__dict__), + SlowDownReqInput(**obj.__dict__), request, ) except Exception as e: return _create_error_response(e) -@dataclasses.dataclass -class LoadLoRAAdapterReqInput: - # The name of the lora module to newly loaded. - lora_name: str - # The path of loading. - lora_path: str - # Whether to pin the LoRA adapter in memory. - pinned: bool = False - # The unique identifier for the LoRA adapter, which automatically generated in the `TokenizerManager`. - lora_id: Optional[str] = None - - @app.api_route("/load_lora_adapter", methods=["POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def load_lora_adapter(obj: LoadLoRAAdapterReqInput, request: Request): +async def load_lora_adapter(obj: LoadLoRAAdapterReqInputSpec, request: Request): """Load a new LoRA adapter without re-launching the server.""" result = await _global_state.tokenizer_manager.load_lora_adapter( - LoadLoRAAdapterReqInputIpc(**obj.__dict__), + LoadLoRAAdapterReqInput(**obj.__dict__), request, ) result = msgspec.to_builtins(result) @@ -1644,24 +1429,13 @@ async def load_lora_adapter(obj: LoadLoRAAdapterReqInput, request: Request): ) -@dataclasses.dataclass -class LoadLoRAAdapterFromTensorsReqInput: - lora_name: str - config_dict: Dict[str, Any] - serialized_tensors: str - pinned: bool = False - added_tokens_config: Optional[Dict[str, Any]] = None - lora_id: Optional[str] = None - load_format: Optional[str] = None - - @app.api_route("/load_lora_adapter_from_tensors", methods=["POST"]) async def load_lora_adapter_from_tensors( - obj: LoadLoRAAdapterFromTensorsReqInput, request: Request + obj: LoadLoRAAdapterFromTensorsReqInputSpec, request: Request ): """Load a new LoRA adapter from tensors without re-launching the server.""" result = await _global_state.tokenizer_manager.load_lora_adapter_from_tensors( - LoadLoRAAdapterFromTensorsReqInputIpc(**obj.__dict__), + LoadLoRAAdapterFromTensorsReqInput(**obj.__dict__), request, ) result = msgspec.to_builtins(result) @@ -1672,20 +1446,12 @@ async def load_lora_adapter_from_tensors( return ORJSONResponse(result, status_code=HTTPStatus.BAD_REQUEST) -@dataclasses.dataclass -class UnloadLoRAAdapterReqInput: - # The name of lora module to unload. - lora_name: str - # The unique identifier for the LoRA adapter, which automatically generated in the `TokenizerManager`. - lora_id: Optional[str] = None - - @app.api_route("/unload_lora_adapter", methods=["POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def unload_lora_adapter(obj: UnloadLoRAAdapterReqInput, request: Request): +async def unload_lora_adapter(obj: UnloadLoRAAdapterReqInputSpec, request: Request): """Load a new LoRA adapter without re-launching the server.""" result = await _global_state.tokenizer_manager.unload_lora_adapter( - UnloadLoRAAdapterReqInputIpc(**obj.__dict__), + UnloadLoRAAdapterReqInput(**obj.__dict__), request, ) result = msgspec.to_builtins(result) @@ -1702,20 +1468,12 @@ async def unload_lora_adapter(obj: UnloadLoRAAdapterReqInput, request: Request): ) -@dataclasses.dataclass -class OpenSessionReqInput: - capacity_of_str_len: int - session_id: Optional[str] = None - streaming: Optional[bool] = None - timeout: Optional[float] = None - - @app.api_route("/open_session", methods=["GET", "POST"]) -async def open_session(obj: OpenSessionReqInput, request: Request): +async def open_session(obj: OpenSessionReqInputSpec, request: Request): """Open a session, and return its unique session id.""" try: session_id = await _global_state.tokenizer_manager.open_session( - OpenSessionReqInputIpc(**obj.__dict__), + OpenSessionReqInput(**obj.__dict__), request, ) if session_id is None: @@ -1727,17 +1485,12 @@ async def open_session(obj: OpenSessionReqInput, request: Request): return _create_error_response(e) -@dataclasses.dataclass -class CloseSessionReqInput: - session_id: str - - @app.api_route("/close_session", methods=["GET", "POST"]) -async def close_session(obj: CloseSessionReqInput, request: Request): +async def close_session(obj: CloseSessionReqInputSpec, request: Request): """Close the session.""" try: await _global_state.tokenizer_manager.close_session( - CloseSessionReqInputIpc(**obj.__dict__), + CloseSessionReqInput(**obj.__dict__), request, ) return Response(status_code=200) @@ -1745,38 +1498,19 @@ async def close_session(obj: CloseSessionReqInput, request: Request): return _create_error_response(e) -@dataclasses.dataclass -class ConfigureLoggingReq: - log_requests: Optional[bool] = None - log_requests_level: Optional[int] = None - log_requests_format: Optional[str] = None - log_level: Optional[str] = None - dump_requests_folder: Optional[str] = None - dump_requests_threshold: Optional[int] = None - crash_dump_folder: Optional[str] = None - dump_requests_exclude_meta_keys: Optional[List[str]] = None - - @app.api_route("/configure_logging", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def configure_logging(obj: ConfigureLoggingReq, request: Request): +async def configure_logging(obj: ConfigureLoggingReqSpec, request: Request): """Configure the request logging options.""" _global_state.tokenizer_manager.configure_logging( - ConfigureLoggingReqIpc(**obj.__dict__) + ConfigureLoggingReq(**obj.__dict__) ) return Response(status_code=200) -@dataclasses.dataclass -class AbortReq: - rid: Optional[str] = None - # Whether to abort all requests - abort_all: bool = False - - @app.post("/abort_request") @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def abort_request(obj: AbortReq, request: Request): +async def abort_request(obj: AbortReqSpec, request: Request): """Abort a request.""" try: _global_state.tokenizer_manager.abort_request( @@ -1788,7 +1522,7 @@ async def abort_request(obj: AbortReq, request: Request): @app.post("/parse_function_call") -async def parse_function_call_request(obj: ParseFunctionCallReq, request: Request): +async def parse_function_call_request(obj: ParseFunctionCallReqSpec, request: Request): """ A native API endpoint to parse function calls from a text. """ @@ -1810,7 +1544,9 @@ async def parse_function_call_request(obj: ParseFunctionCallReq, request: Reques @app.post("/separate_reasoning") -async def separate_reasoning_request(obj: SeparateReasoningReqInput, request: Request): +async def separate_reasoning_request( + obj: SeparateReasoningReqInputSpec, request: Request +): """ A native API endpoint to separate reasoning from a text. """ @@ -1840,24 +1576,12 @@ async def separate_reasoning_request(obj: SeparateReasoningReqInput, request: Re return ORJSONResponse(content=response_data, status_code=200) -@dataclasses.dataclass -class PauseGenerationReqInput: - mode: Literal["abort", "retract", "in_place"] = "abort" - - def __post_init__(self): - allowed = ["abort", "retract", "in_place"] - if self.mode not in allowed: - raise ValueError( - f"Invalid mode: {self.mode!r}. " f"Expected one of {allowed}." - ) - - @app.post("/pause_generation") @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def pause_generation(obj: PauseGenerationReqInput, request: Request): +async def pause_generation(obj: PauseGenerationReqInputSpec, request: Request): """Pause generation.""" await _global_state.tokenizer_manager.pause_generation( - PauseGenerationReqInputIpc(**obj.__dict__) + PauseGenerationReqInput(**obj.__dict__) ) return ORJSONResponse( content={"message": "Generation paused successfully.", "status": "ok"}, @@ -1865,17 +1589,12 @@ async def pause_generation(obj: PauseGenerationReqInput, request: Request): ) -@dataclasses.dataclass -class ContinueGenerationReqInput: - torch_empty_cache: bool = True - - @app.post("/continue_generation") @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def continue_generation(obj: ContinueGenerationReqInput, request: Request): +async def continue_generation(obj: ContinueGenerationReqInputSpec, request: Request): """Continue generation.""" await _global_state.tokenizer_manager.continue_generation( - ContinueGenerationReqInputIpc(**obj.__dict__) + ContinueGenerationReqInput(**obj.__dict__) ) return ORJSONResponse( content={"message": "Generation continued successfully.", "status": "ok"}, @@ -2214,7 +1933,7 @@ async def sagemaker_chat_completions( ## Vertex AI API @app.post(os.environ.get("AIP_PREDICT_ROUTE", "/vertex_generate")) -async def vertex_generate(vertex_req: VertexGenerateReqInput, raw_request: Request): +async def vertex_generate(vertex_req: VertexGenerateReqInputSpec, raw_request: Request): if not vertex_req.instances: return [] inputs = {} diff --git a/python/sglang/srt/managers/communicator.py b/python/sglang/srt/managers/communicator.py index 58dd0eed6d9f..65910865aafc 100644 --- a/python/sglang/srt/managers/communicator.py +++ b/python/sglang/srt/managers/communicator.py @@ -35,8 +35,6 @@ def __init__(self, sender: zmq.Socket, fan_out: int, mode="queueing"): assert mode in ["queueing", "watching"] async def queueing_call(self, obj: T): - from sglang.srt.managers.multi_tokenizer_mixin import SenderWrapper - ready_event = asyncio.Event() if self._result_event is not None or len(self._ready_queue) > 0: self._ready_queue.append(ready_event) @@ -45,10 +43,11 @@ async def queueing_call(self, obj: T): assert self._result_values is None if obj is not None: - if isinstance(self._sender, SenderWrapper): - self._sender.send_obj(obj) - else: + if isinstance(self._sender, zmq.Socket): sock_send(self._sender, obj) + else: + # for no zmq.Socket, it will be SenderWrapper + self._sender.send_obj(obj) self._result_event = asyncio.Event() self._result_values = [] @@ -62,18 +61,17 @@ async def queueing_call(self, obj: T): return result_values async def watching_call(self, obj): - from sglang.srt.managers.multi_tokenizer_mixin import SenderWrapper - if self._result_event is None: assert self._result_values is None self._result_values = [] self._result_event = asyncio.Event() if obj is not None: - if isinstance(self._sender, SenderWrapper): - self._sender.send_obj(obj) - else: + if isinstance(self._sender, zmq.Socket): sock_send(self._sender, obj) + else: + # for no zmq.Socket, it will be SenderWrapper + self._sender.send_obj(obj) # Capture local refs before await -- after event fires, the first # awakened coroutine clears shared state; later awaiters use local refs. diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index e7a1b8fa9db2..7e2a4a9e5326 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -105,28 +105,6 @@ class BaseBatchReqIpc(msgspec.Struct, tag=True, kw_only=True): http_worker_ipcs: Optional[List[Optional[str]]] = None -class SpeculativeDecodingMetricsMixinReq(BaseBatchReqIpc, kw_only=True): - """ - Mixin class containing speculative decoding metrics. - - This class consolidates speculative decoding metrics that are shared across - batch output types that support speculative decoding to avoid code duplication. - """ - - # Verify count: number of verification forward passes - spec_verify_ct: Optional[List[int]] - - # Accepted drafts: Number of accepted draft tokens during speculative decoding - # (strict drafts-only count, excludes the bonus token). - spec_num_correct_drafts: Optional[List[int]] - - # Acceptance histogram: List of lists, where each inner list represents histogram counts. - # List index = number of accepted tokens in a step, List value = count of steps with that many accepted tokens. - # Example: histogram[0] = 5 means 5 steps with 0 accepted tokens, histogram[3] = 10 means 10 steps with 3 accepted tokens. - # Empty list [] when speculative decoding is disabled. - spec_correct_drafts_histogram: Optional[List[List[int]]] - - # Parameters for a session @dataclass class SessionParams: @@ -1122,7 +1100,7 @@ def __iter__(self): return iter(self.batch) -class BatchTokenIDOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): +class BatchTokenIDOutput(BaseBatchReqIpc, kw_only=True): # The finish reason finished_reasons: List[Optional[Dict[str, Any]]] # List[BaseFinishReason] # For incremental decoding @@ -1228,8 +1206,22 @@ class BatchTokenIDOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): # For observability time_stats: Optional[List[SchedulerReqTimeStats]] = None + # From the SpeculativeDecodingMetricsMixinReq class + # Verify count: number of verification forward passes + spec_verify_ct: Optional[List[int]] -class BatchStrOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): + # Accepted drafts: Number of accepted draft tokens during speculative decoding + # (strict drafts-only count, excludes the bonus token). + spec_num_correct_drafts: Optional[List[int]] + + # Acceptance histogram: List of lists, where each inner list represents histogram counts. + # List index = number of accepted tokens in a step, List value = count of steps with that many accepted tokens. + # Example: histogram[0] = 5 means 5 steps with 0 accepted tokens, histogram[3] = 10 means 10 steps with 3 accepted tokens. + # Empty list [] when speculative decoding is disabled. + spec_correct_drafts_histogram: Optional[List[List[int]]] + + +class BatchStrOutput(BaseBatchReqIpc, kw_only=True): # The finish reason finished_reasons: List[Optional[Dict[str, Any]]] # List[dict] # The output decoded strings @@ -1329,6 +1321,20 @@ class BatchStrOutput(SpeculativeDecodingMetricsMixinReq, kw_only=True): # For observability time_stats: Optional[List[SchedulerReqTimeStats]] = None + # From the SpeculativeDecodingMetricsMixinReq class + # Verify count: number of verification forward passes + spec_verify_ct: Optional[List[int]] + + # Accepted drafts: Number of accepted draft tokens during speculative decoding + # (strict drafts-only count, excludes the bonus token). + spec_num_correct_drafts: Optional[List[int]] + + # Acceptance histogram: List of lists, where each inner list represents histogram counts. + # List index = number of accepted tokens in a step, List value = count of steps with that many accepted tokens. + # Example: histogram[0] = 5 means 5 steps with 0 accepted tokens, histogram[3] = 10 means 10 steps with 3 accepted tokens. + # Empty list [] when speculative decoding is disabled. + spec_correct_drafts_histogram: Optional[List[List[int]]] + class BatchEmbeddingOutput(BaseBatchReqIpc, kw_only=True): # The finish reason @@ -1663,14 +1669,6 @@ class DestroyWeightsUpdateGroupReqOutput(BaseReqIpc, kw_only=True): message: str -@dataclass -class UpdateWeightVersionReqInput(BaseReq): - # The new weight version - new_version: str - # Whether to abort all running requests before updating - abort_all_requests: bool = True - - class GetWeightsByNameReqInput(BaseReqIpc, kw_only=True): name: str truncate_size: int = 100 @@ -1761,32 +1759,6 @@ class SetInternalStateReqOutput(BaseReqIpc, kw_only=True): server_args: Dict[str, object] -@dataclass -class ProfileReqInput(BaseReq): - # The output directory - output_dir: Optional[str] = None - # Specify the steps to start the profiling - start_step: Optional[int] = None - # If set, it profile as many as this number of steps. - # If it is set, profiling is automatically stopped after this step, and - # the caller doesn't need to run stop_profile. - num_steps: Optional[int] = None - # The activities to record. The choices are ["CPU", "GPU", "MEM", "RPD"] - activities: Optional[List[str]] = None - # Whether profile by stages (e.g., prefill and decode) separately - profile_by_stage: bool = False - # Whether to record source information (file and line number) for the ops. - with_stack: Optional[bool] = None - # Whether to save information about operator’s input shapes. - record_shapes: Optional[bool] = None - # Merge profiles from all ranks into a single trace - merge_profiles: bool = False - # The prefix of the profile filenames - profile_prefix: Optional[str] = None - # Only profile these stages and ignore others - profile_stages: Optional[List[str]] = None - - class ProfileReqType(Enum): START_PROFILE = 1 STOP_PROFILE = 2 @@ -1874,30 +1846,6 @@ class Tool: type: Optional[str] = "function" -@dataclass -class ParseFunctionCallReq(BaseReq): - text: str # The text to parse. - tools: List[Tool] = field( - default_factory=list - ) # A list of available function tools (name, parameters, etc.). - tool_call_parser: Optional[str] = ( - None # Specify the parser type, e.g. 'llama3', 'qwen25', or 'mistral'. If not specified, tries all. - ) - - -@dataclass -class SeparateReasoningReqInput(BaseReq): - 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): - instances: List[dict] - parameters: Optional[dict] = None - - class RpcReqInput(BaseReqIpc, kw_only=True): method: str parameters: Optional[Dict] = None @@ -2172,7 +2120,8 @@ def _check_all_req_types(): is_base_req = ( issubclass(class_type[1], BaseReq) or issubclass(class_type[1], BaseBatchReq) - or issubclass(class_type[1], msgspec.Struct) # allow IPC objects + or issubclass(class_type[1], BaseReqIpc) + or issubclass(class_type[1], BaseBatchReqIpc) ) if is_io_struct and not is_base_req: raise ValueError(f"{name} is not a subclass of BaseReq or BaseBatchReq.") @@ -2194,7 +2143,7 @@ def enc_hook(obj: Any) -> Any: if isinstance(obj, torch.Tensor): # encode torch tensor as Tuple(shape, dtype, data) tensor_dtype = str(obj.dtype).split(".")[-1] # e.g., "float32" - raw_data = obj.flatten().cpu().contiguous().view(torch.uint8).numpy().data + raw_data = obj.numpy().tobytes() return (obj.shape, tensor_dtype, raw_data) elif isinstance(obj, array): return (obj.typecode, obj.tobytes()) @@ -2243,7 +2192,6 @@ def dec_hook(tp: Type, obj: Any) -> Any: cls for cls in BaseReqIpc.__subclasses__() + BaseBatchReqIpc.__subclasses__() - + SpeculativeDecodingMetricsMixinReq.__subclasses__() + [PickleWrapperIpc] ) # str primitive types should encode to bytes diff --git a/python/sglang/srt/sampling/sampling_params.py b/python/sglang/srt/sampling/sampling_params.py index 3b0c0ea7fba2..5cba35ffe324 100644 --- a/python/sglang/srt/sampling/sampling_params.py +++ b/python/sglang/srt/sampling/sampling_params.py @@ -14,9 +14,10 @@ """Sampling parameters for text generation.""" import logging -from dataclasses import dataclass from typing import Any, Dict, List, Optional, Union +import msgspec + # sre_parse is deprecated in Python 3.11+, use re._parser instead try: import re._parser as sre_parse @@ -29,8 +30,7 @@ logger = logging.getLogger(__name__) -@dataclass -class SamplingParams: +class SamplingParams(msgspec.Struct, kw_only=True): """ The sampling parameters. @@ -39,105 +39,76 @@ class SamplingParams: for the documentation. """ - max_new_tokens: Optional[int] - stop_strs: Optional[Union[str, List[str]]] - stop_token_ids: Optional[List[int]] - stop_regex_strs: Optional[Union[str, List[str]]] - stop_str_max_len: int - stop_regex_max_len: int - temperature: float - top_p: float - top_k: int - min_p: float - frequency_penalty: float - presence_penalty: float - repetition_penalty: float - min_new_tokens: int - regex: Optional[str] - n: int - json_schema: Optional[str] - ebnf: Optional[str] - structural_tag: Optional[str] - ignore_eos: bool - skip_special_tokens: bool - spaces_between_special_tokens: bool - no_stop_trim: bool - custom_params: Optional[Dict[str, Any]] - stream_interval: Optional[int] - logit_bias: Optional[Dict[str, float]] - sampling_seed: Optional[int] + max_new_tokens: Optional[int] = 128 + stop: Optional[Union[str, List[str]]] = None + stop_strs: Optional[Union[str, List[str]]] = None + stop_token_ids: Optional[List[int]] = None + stop_regex: Optional[Union[str, List[str]]] = None + stop_regex_strs: Optional[Union[str, List[str]]] = None + stop_str_max_len: int = 0 + stop_regex_max_len: int = 0 + temperature: float = 1.0 + top_p: float = 1.0 + top_k: int = -1 + min_p: float = 0.0 + frequency_penalty: float = 0.0 + presence_penalty: float = 0.0 + repetition_penalty: float = 1.0 + min_new_tokens: int = 0 + regex: Optional[str] = None + n: int = 1 + json_schema: Optional[str] = None + ebnf: Optional[str] = None + structural_tag: Optional[str] = None + ignore_eos: bool = False + skip_special_tokens: bool = True + spaces_between_special_tokens: bool = True + no_stop_trim: bool = False + custom_params: Optional[Dict[str, Any]] = None + stream_interval: Optional[int] = None + logit_bias: Optional[Dict[str, float]] = None + sampling_seed: Optional[int] = None - def __init__( - self, - max_new_tokens: int = 128, - stop: Optional[Union[str, List[str]]] = None, - stop_token_ids: Optional[List[int]] = None, - stop_regex: Optional[Union[str, List[str]]] = None, - temperature: float = 1.0, - top_p: float = 1.0, - top_k: int = -1, - min_p: float = 0.0, - frequency_penalty: float = 0.0, - presence_penalty: float = 0.0, - repetition_penalty: float = 1.0, - min_new_tokens: int = 0, - n: int = 1, - json_schema: Optional[str] = None, - regex: Optional[str] = None, - ebnf: Optional[str] = None, - structural_tag: Optional[str] = None, - ignore_eos: bool = False, - skip_special_tokens: bool = True, - spaces_between_special_tokens: bool = True, - no_stop_trim: bool = False, - custom_params: Optional[Dict[str, Any]] = None, - stream_interval: Optional[int] = None, - logit_bias: Optional[Dict[str, float]] = None, - sampling_seed: Optional[int] = None, - ) -> None: + def __post_init__(self): # For non-optional params, treat None as "use default" so that callers # (e.g. /generate) can pass null without crashing verify(). - self.max_new_tokens = max_new_tokens - self.stop_strs = stop - if stop_token_ids: - filtered = {int(t) for t in stop_token_ids if t is not None} + self.stop_strs = self.stop or [] + if self.stop_token_ids: + filtered = {int(t) for t in self.stop_token_ids if t is not None} self.stop_token_ids = filtered or None else: self.stop_token_ids = None - self.stop_regex_strs = stop_regex - self.temperature = temperature if temperature is not None else 1.0 - self.top_p = top_p if top_p is not None else 1.0 - self.top_k = top_k if top_k is not None else -1 - self.min_p = min_p if min_p is not None else 0.0 + self.stop_regex_strs = self.stop_regex or [] + self.temperature = self.temperature if self.temperature is not None else 1.0 + self.top_p = self.top_p if self.top_p is not None else 1.0 + self.top_k = self.top_k if self.top_k is not None else -1 + self.min_p = self.min_p if self.min_p is not None else 0.0 self.frequency_penalty = ( - frequency_penalty if frequency_penalty is not None else 0.0 + self.frequency_penalty if self.frequency_penalty is not None else 0.0 ) self.presence_penalty = ( - presence_penalty if presence_penalty is not None else 0.0 + self.presence_penalty if self.presence_penalty is not None else 0.0 ) self.repetition_penalty = ( - repetition_penalty if repetition_penalty is not None else 1.0 + self.repetition_penalty if self.repetition_penalty is not None else 1.0 + ) + self.min_new_tokens = ( + self.min_new_tokens if self.min_new_tokens is not None else 0 ) - self.min_new_tokens = min_new_tokens if min_new_tokens is not None else 0 - self.regex = regex - self.n = n if n is not None else 1 - self.json_schema = json_schema - self.ebnf = ebnf - self.structural_tag = structural_tag - self.ignore_eos = ignore_eos if ignore_eos is not None else False + self.regex = self.regex + self.n = self.n if self.n is not None else 1 + self.ignore_eos = self.ignore_eos if self.ignore_eos is not None else False self.skip_special_tokens = ( - skip_special_tokens if skip_special_tokens is not None else True + self.skip_special_tokens if self.skip_special_tokens is not None else True ) self.spaces_between_special_tokens = ( - spaces_between_special_tokens - if spaces_between_special_tokens is not None + self.spaces_between_special_tokens + if self.spaces_between_special_tokens is not None else True ) - self.no_stop_trim = no_stop_trim if no_stop_trim is not None else False - self.custom_params = custom_params - self.stream_interval = stream_interval - self.logit_bias = logit_bias - self.sampling_seed = sampling_seed + self.no_stop_trim = ( + self.no_stop_trim if self.no_stop_trim is not None else False + ) # Process some special cases if 0 <= self.temperature < _SAMPLING_EPS: diff --git a/test/registered/unit/managers/test_profile_merger_http_api.py b/test/registered/unit/managers/test_profile_merger_http_api.py index 92aebef3eb21..ae95c4190847 100644 --- a/test/registered/unit/managers/test_profile_merger_http_api.py +++ b/test/registered/unit/managers/test_profile_merger_http_api.py @@ -1,7 +1,7 @@ import json import unittest -from sglang.srt.managers.io_struct import ProfileReqInput +from sglang.srt.entrypoints.api_specs import ProfileReqInputSpec from sglang.test.ci.ci_register import ( register_amd_ci, register_cpu_ci, @@ -16,7 +16,7 @@ class TestProfileMergerHTTPAPI(unittest.TestCase): def test_profile_req_input_merge_profiles_json_serialization(self): # Test with merge_profiles=True - req_input = ProfileReqInput( + req_input = ProfileReqInputSpec( output_dir="/tmp/test", num_steps=5, activities=["CPU", "GPU"], @@ -53,8 +53,8 @@ def test_profile_req_input_merge_profiles_json_deserialization(self): "merge_profiles": True, } - # Create ProfileReqInput from dict (as HTTP server would do) - req_input = ProfileReqInput(**json_data) + # Create ProfileReqInputSpec from dict (as HTTP server would do) + req_input = ProfileReqInputSpec(**json_data) self.assertTrue(req_input.merge_profiles) self.assertEqual(req_input.output_dir, "/tmp/test") @@ -66,13 +66,13 @@ def test_profile_req_input_merge_profiles_default_value(self): # Test with minimal data json_data = {"output_dir": "/tmp/test"} - req_input = ProfileReqInput(**json_data) + req_input = ProfileReqInputSpec(**json_data) self.assertFalse(req_input.merge_profiles) def test_profile_req_input_merge_profiles_explicit_false(self): json_data = {"output_dir": "/tmp/test", "merge_profiles": False} - req_input = ProfileReqInput(**json_data) + req_input = ProfileReqInputSpec(**json_data) self.assertFalse(req_input.merge_profiles) def test_http_api_parameter_flow(self): @@ -85,8 +85,8 @@ def test_http_api_parameter_flow(self): "merge_profiles": True, } - # Create ProfileReqInput as HTTP server would - obj = ProfileReqInput(**request_data) + # Create ProfileReqInputSpec as HTTP server would + obj = ProfileReqInputSpec(**request_data) # Verify the parameter is set correctly self.assertTrue(obj.merge_profiles) @@ -98,23 +98,23 @@ def test_http_api_parameter_flow(self): def test_http_api_parameter_validation(self): # Test with True json_data = {"merge_profiles": True} - req_input = ProfileReqInput(**json_data) + req_input = ProfileReqInputSpec(**json_data) self.assertTrue(req_input.merge_profiles) # Test with False json_data = {"merge_profiles": False} - req_input = ProfileReqInput(**json_data) + req_input = ProfileReqInputSpec(**json_data) self.assertFalse(req_input.merge_profiles) # Test with string "true" (should be converted by JSON parser) json_data = {"merge_profiles": "true"} - req_input = ProfileReqInput(**json_data) + req_input = ProfileReqInputSpec(**json_data) self.assertEqual(req_input.merge_profiles, "true") # String, not boolean def test_http_api_backward_compatibility(self): # Test minimal request (no merge_profiles) json_data = {} - req_input = ProfileReqInput(**json_data) + req_input = ProfileReqInputSpec(**json_data) self.assertFalse(req_input.merge_profiles) # Should default to False # Test with other parameters but no merge_profiles @@ -123,7 +123,7 @@ def test_http_api_backward_compatibility(self): "num_steps": 5, "activities": ["CPU", "GPU"], } - req_input = ProfileReqInput(**json_data) + req_input = ProfileReqInputSpec(**json_data) self.assertFalse(req_input.merge_profiles) # Should default to False def test_http_api_parameter_combinations(self): @@ -163,7 +163,7 @@ def test_http_api_parameter_combinations(self): for test_case in test_cases: with self.subTest(test_case["name"]): - req_input = ProfileReqInput(**test_case["data"]) + req_input = ProfileReqInputSpec(**test_case["data"]) self.assertEqual(req_input.merge_profiles, test_case["expected_merge"]) diff --git a/test/registered/unit/utils/test_profile_merger.py b/test/registered/unit/utils/test_profile_merger.py index e716f85b11a8..d6d39705085b 100644 --- a/test/registered/unit/utils/test_profile_merger.py +++ b/test/registered/unit/utils/test_profile_merger.py @@ -13,7 +13,8 @@ import tempfile import unittest -from sglang.srt.managers.io_struct import ProfileReq, ProfileReqInput, ProfileReqType +from sglang.srt.entryponts.api_specs import ProfileReqInputSpec +from sglang.srt.managers.io_struct import ProfileReq, ProfileReqType from sglang.srt.utils.profile_merger import ProfileMerger from sglang.test.ci.ci_register import ( register_amd_ci, @@ -208,11 +209,11 @@ def test_merge_chrome_traces(self): class TestProfileMergerIntegration(unittest.TestCase): def test_data_structures_merge_profiles(self): - # Test ProfileReqInput - req_input = ProfileReqInput() + # Test ProfileReqInputSpec + req_input = ProfileReqInputSpec() self.assertFalse(req_input.merge_profiles) - req_input = ProfileReqInput(merge_profiles=True) + req_input = ProfileReqInputSpec(merge_profiles=True) self.assertTrue(req_input.merge_profiles) # Test ProfileReq From 5d78f74c21515d927f249030dfd089ae0d3aaa98 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Fri, 12 Jun 2026 10:40:37 -0700 Subject: [PATCH 27/45] fix ci test --- python/sglang/srt/managers/io_struct.py | 2 +- test/registered/unit/utils/test_profile_merger.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 7e2a4a9e5326..d7483b7bbfd0 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -2143,7 +2143,7 @@ def enc_hook(obj: Any) -> Any: if isinstance(obj, torch.Tensor): # encode torch tensor as Tuple(shape, dtype, data) tensor_dtype = str(obj.dtype).split(".")[-1] # e.g., "float32" - raw_data = obj.numpy().tobytes() + raw_data = obj.flatten().cpu().contiguous().view(torch.uint8).numpy().data return (obj.shape, tensor_dtype, raw_data) elif isinstance(obj, array): return (obj.typecode, obj.tobytes()) diff --git a/test/registered/unit/utils/test_profile_merger.py b/test/registered/unit/utils/test_profile_merger.py index d6d39705085b..8521e547ee90 100644 --- a/test/registered/unit/utils/test_profile_merger.py +++ b/test/registered/unit/utils/test_profile_merger.py @@ -13,7 +13,7 @@ import tempfile import unittest -from sglang.srt.entryponts.api_specs import ProfileReqInputSpec +from sglang.srt.entrypoints.api_specs import ProfileReqInputSpec from sglang.srt.managers.io_struct import ProfileReq, ProfileReqType from sglang.srt.utils.profile_merger import ProfileMerger from sglang.test.ci.ci_register import ( From e32ec71f491807e57cb07e489b2caffea57197dc Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Fri, 12 Jun 2026 14:01:25 -0700 Subject: [PATCH 28/45] fix samplingparam re-post-init after normalize --- python/sglang/srt/sampling/sampling_params.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/sampling/sampling_params.py b/python/sglang/srt/sampling/sampling_params.py index 5cba35ffe324..8b93d7926c33 100644 --- a/python/sglang/srt/sampling/sampling_params.py +++ b/python/sglang/srt/sampling/sampling_params.py @@ -68,17 +68,25 @@ class SamplingParams(msgspec.Struct, kw_only=True): stream_interval: Optional[int] = None logit_bias: Optional[Dict[str, float]] = None sampling_seed: Optional[int] = None + is_normalized: bool = False def __post_init__(self): # For non-optional params, treat None as "use default" so that callers # (e.g. /generate) can pass null without crashing verify(). - self.stop_strs = self.stop or [] + + # Since the __post_init__ is called after msgpack deserialization, + # we try to avoid init again to reset the values. + if self.is_normalized: + return + + self.stop_strs = self.stop + self.stop_regex_strs = self.stop_regex + if self.stop_token_ids: filtered = {int(t) for t in self.stop_token_ids if t is not None} self.stop_token_ids = filtered or None else: self.stop_token_ids = None - self.stop_regex_strs = self.stop_regex or [] self.temperature = self.temperature if self.temperature is not None else 1.0 self.top_p = self.top_p if self.top_p is not None else 1.0 self.top_k = self.top_k if self.top_k is not None else -1 @@ -210,6 +218,8 @@ def normalize(self, tokenizer): self.stop_regex_max_len = stop_regex_max_len + self.is_normalized = True + # This function gets a strict upperbound on the maximum number of tokens that would need # to be buffered to match the input regex string From 87b52ac898de6d0d3a504e56dc2631c124a76547 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Fri, 12 Jun 2026 18:37:09 -0700 Subject: [PATCH 29/45] support metrics obj searialized with msgpack --- python/sglang/srt/managers/io_struct.py | 13 ++-- .../srt/observability/req_time_stats.py | 47 ++++++++++---- python/sglang/srt/observability/trace.py | 65 ++++++++----------- 3 files changed, 68 insertions(+), 57 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index d7483b7bbfd0..866a9ba1aa29 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -42,7 +42,9 @@ 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 ( - ReqTimeStatsBase, + APIServerReqTimeStats, + DPControllerReqTimeStats, + MetricsCollectorWrapper, SchedulerReqTimeStats, ) from sglang.srt.sampling.sampling_params import SamplingParams @@ -824,8 +826,8 @@ class TokenizedGenerateReqInput(BaseReqIpc, kw_only=True): multi_item_delimiter_indices: Optional[List[int]] = None # For observability - # Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] - time_stats: Optional[ReqTimeStatsBase] = None + # Optional[] + time_stats: Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] = None class BatchTokenizedGenerateReqInput(BaseBatchReqIpc, kw_only=True): @@ -1079,8 +1081,7 @@ class TokenizedEmbeddingReqInput(BaseReqIpc, kw_only=True): # Pre-computed delimiter indices for multi-item scoring multi_item_delimiter_indices: Optional[List[int]] = None # For observability - # Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] - time_stats: Optional[ReqTimeStatsBase] = None + time_stats: Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] = None # Whether to return pooled hidden states (pre-head transformer output) return_pooled_hidden_states: bool = False @@ -2149,6 +2150,8 @@ def enc_hook(obj: Any) -> Any: return (obj.typecode, obj.tobytes()) elif isinstance(obj, np.floating): return float(obj) + elif isinstance(obj, MetricsCollectorWrapper): + return None else: if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): logger.info(f"Object of type {type(obj)} is encoding with pickle.") diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index 6a9a86a1ad49..2c5c8b2a9a45 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/python/sglang/srt/observability/req_time_stats.py @@ -18,9 +18,10 @@ import logging import time import uuid -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +import msgspec from typing_extensions import Self from sglang.srt.disaggregation.utils import DisaggregationMode @@ -219,17 +220,38 @@ class RequestStage: ANONYMOUS = RequestStageConfig("") -@dataclass -class ReqTimeStatsBase: +# The metrics collector is not serializable, +# and it is only used for observing metrics in the current process. +class MetricsCollectorWrapper: # the non-serializable wrapper + __slots__ = ("inner",) + + def __init__(self, inner=None): + self.inner = inner + + def __getattr__(self, name): # delegate to internal object + if name.startswith("__") and name.endswith("__"): + raise AttributeError(name) + return getattr(self.inner, name) + + +class ReqTimeStatsBase(msgspec.Struct, kw_only=True, tag=True): enable_metrics: bool = False - # use object to avoid msgpack decode build issue, not use Any due to the - # custom object serialized by pickle. - # Optional[Union[SchedulerMetricsCollector, TokenizerMetricsCollector]] - metrics_collector: Optional[object] = None - # Union[TraceReqContext, TraceNullContext] - trace_ctx: object = field(default_factory=TraceNullContext) + trace_ctx: Union[TraceReqContext, TraceNullContext] = msgspec.field( + default_factory=TraceNullContext + ) + metrics_collector: Optional[MetricsCollectorWrapper] = None disagg_mode: Optional[DisaggregationMode] = DisaggregationMode.NULL diff_realtime_monotonic: float = 0.0 + is_initialized: bool = False + + def __post_init__(self): + if self.is_initialized: + return + + if not isinstance(self.metrics_collector, MetricsCollectorWrapper): + self.metrics_collector = MetricsCollectorWrapper(self.metrics_collector) + + self.is_initialized = True @classmethod def new_from_obj(cls, obj: Optional[ReqTimeStatsBase], *args, **kwargs) -> Self: @@ -237,7 +259,7 @@ def new_from_obj(cls, obj: Optional[ReqTimeStatsBase], *args, **kwargs) -> Self: new_obj = cls(*args, **kwargs) if obj is None: return new_obj - for key, value in obj.__dict__.items(): + for key, value in msgspec.structs.asdict(obj).items(): if hasattr(new_obj, key): setattr(new_obj, key, value) @@ -261,7 +283,7 @@ def set_metrics_collector( ): if collector: self.enable_metrics = True - self.metrics_collector = collector + self.metrics_collector = MetricsCollectorWrapper(collector) def observe_per_stage_req_latency(self, stage: RequestStageConfig, latency: float): if self.enable_metrics and stage.metrics_is_observed: @@ -324,7 +346,6 @@ def __setstate__(self, state: object): self.__dict__.update(state) -@dataclass class APIServerReqTimeStats(ReqTimeStatsBase): # get by time.perf_counter() created_time: float = 0.0 @@ -475,7 +496,6 @@ def convert_to_gen_ai_span_attrs(self): return span_attrs -@dataclass class DPControllerReqTimeStats(ReqTimeStatsBase): # propagated from tokenizer/grpc_server, get by time.perf_counter() created_time: float = 0.0 @@ -521,7 +541,6 @@ def set_dp_dispatch_finish_time(self, ts=None): ) -@dataclass class SchedulerReqTimeStats(ReqTimeStatsBase): """ Store the timestamps for each stage of a request. diff --git a/python/sglang/srt/observability/trace.py b/python/sglang/srt/observability/trace.py index 657c3c3d3a7a..e5d1dfb9eb62 100644 --- a/python/sglang/srt/observability/trace.py +++ b/python/sglang/srt/observability/trace.py @@ -24,6 +24,8 @@ from dataclasses import dataclass from typing import Any, Dict, List, Mapping, Optional +import msgspec + from sglang.srt.utils import get_int_env_var logger = logging.getLogger(__name__) @@ -258,49 +260,39 @@ def trace_set_thread_info( ) -class TraceReqContext: - def __init__( - self, - rid, - bootstrap_room=None, - role="unified", - module_name="", - external_trace_header: Optional[Dict[str, str]] = None, - ): - self.rid: str = str(rid) +class TraceReqContext(msgspec.Struct, tag=True): + rid: str + trace_level: int + pid: int + tracing_enable: bool = False + start_time_ns: Optional[int] = None + thread_context: Optional[TraceThreadContext] = None + bootstrap_room: Optional[int] = None + role: str = "unified" + module_name: str = "" + is_copy: bool = False + root_span: Optional[trace.span.Span] = None + root_span_context: Optional[context.Context] = None + last_span_context: Optional[trace.span.SpanContext] = None + external_trace_header: Optional[Dict[str, str]] = None + events_cache: List[TraceEvent] = [] + is_initialized: bool = False + + def __post_init__(self): + self.rid = str(self.rid) self.trace_level = global_trace_level - self.tracing_enable: bool = opentelemetry_initialized and self.trace_level > 0 + self.tracing_enable = opentelemetry_initialized and self.trace_level > 0 # Filter by --trace-modules only for explicitly named modules; contexts # created with the default empty module_name are always traced. if ( - module_name + self.module_name and global_trace_modules is not None - and module_name not in global_trace_modules + and self.module_name not in global_trace_modules ): self.tracing_enable = False - if not self.tracing_enable: - return - - self.start_time_ns: Optional[int] = None - self.thread_context: Optional[TraceThreadContext] = None - self.bootstrap_room: Optional[int] = bootstrap_room - self.role: str = role - self.module_name = module_name - - # Indicates whether this instance is a replica from the main process. - # When True, root_span is None and only root_span_context is preserved. - self.is_copy: bool = False - self.root_span: Optional[trace.span.Span] = None - self.root_span_context: Optional[context.Context] = None - # Record the most recently completed span as the previous span for the next span to be created. - self.last_span_context: Optional[trace.span.SpanContext] = None - self.external_trace_header: Optional[Dict[str, str]] = external_trace_header - - self.events_cache: List[TraceEvent] = [] - - self.pid: int = threading.get_native_id() + self.pid = threading.get_native_id() def is_tracing_enabled(self) -> bool: return self.tracing_enable @@ -768,12 +760,9 @@ def __del__(self): # Temporary remove the @dataclass decorator to make msgpack treat # TraceReqContext as a custom class and serialize with pickle like TraceReqContext. # When the TraceReqContext fixed with msgpack, we can consider to add @dataclass back. -class TraceNullContext: +class TraceNullContext(msgspec.Struct, tag=True): tracing_enable: bool = False - def __init__(self, tracing_enable=False): - self.tracing_enable = tracing_enable - def __getattr__(self, name): # fix the .__dataclass_fields__ is not a dict error if name.startswith("__") and name.endswith("__"): From 73d5008db587759eb71734f086d532de266278c9 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Fri, 12 Jun 2026 19:22:08 -0700 Subject: [PATCH 30/45] lint --- python/sglang/srt/entrypoints/http_server.py | 40 ++++++++------------ 1 file changed, 15 insertions(+), 25 deletions(-) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index d9e35f989b52..5bad78d58bcc 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -1174,11 +1174,9 @@ async def update_weights_from_disk( success, message, num_paused_requests, - ) = ( - await _global_state.tokenizer_manager.update_weights_from_disk( - UpdateWeightFromDiskReqInput(**obj.__dict__), - request, - ) + ) = await _global_state.tokenizer_manager.update_weights_from_disk( + UpdateWeightFromDiskReqInput(**obj.__dict__), + request, ) content = { @@ -1206,11 +1204,9 @@ async def init_weights_send_group_for_remote_instance( ( success, message, - ) = ( - await _global_state.tokenizer_manager.init_weights_send_group_for_remote_instance( - InitWeightsSendGroupForRemoteInstanceReqInput(**obj.__dict__), - request, - ) + ) = await _global_state.tokenizer_manager.init_weights_send_group_for_remote_instance( + InitWeightsSendGroupForRemoteInstanceReqInput(**obj.__dict__), + request, ) content = {"success": success, "message": message} if success: @@ -1227,11 +1223,9 @@ async def send_weights_to_remote_instance( ( success, message, - ) = ( - await _global_state.tokenizer_manager.send_weights_to_remote_instance( - SendWeightsToRemoteInstanceReqInput(**obj.__dict__), - request, - ) + ) = await _global_state.tokenizer_manager.send_weights_to_remote_instance( + SendWeightsToRemoteInstanceReqInput(**obj.__dict__), + request, ) content = {"success": success, "message": message} if success: @@ -1304,11 +1298,9 @@ async def destroy_weights_update_group( ( success, message, - ) = ( - await _global_state.tokenizer_manager.destroy_weights_update_group( - DestroyWeightsUpdateGroupReqInput(**obj.__dict__), - request, - ) + ) = await _global_state.tokenizer_manager.destroy_weights_update_group( + DestroyWeightsUpdateGroupReqInput(**obj.__dict__), + request, ) content = {"success": success, "message": message} return ORJSONResponse( @@ -1347,11 +1339,9 @@ async def update_weights_from_distributed( ( success, message, - ) = ( - await _global_state.tokenizer_manager.update_weights_from_distributed( - UpdateWeightsFromDistributedReqInput(**obj.__dict__), - request, - ) + ) = await _global_state.tokenizer_manager.update_weights_from_distributed( + UpdateWeightsFromDistributedReqInput(**obj.__dict__), + request, ) content = {"success": success, "message": message} From bab26bc4dad1f1b15dd9399ef6cbb4cdbfb7bf57 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Fri, 12 Jun 2026 19:38:07 -0700 Subject: [PATCH 31/45] fix ci test --- python/sglang/srt/observability/req_time_stats.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index bf31aabe99d8..ee173a0e7893 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/python/sglang/srt/observability/req_time_stats.py @@ -1204,7 +1204,6 @@ def format_wallclock(perf_counter_time: float) -> str: return f"{convert_time_to_realtime(perf_counter_time):.3f}" -@dataclass class EncoderReqTimeStats(ReqTimeStatsBase): mm_encode_start_time: float = 0.0 mm_encode_end_time: float = 0.0 From e0328e68ad015d602a8ee10c67d2e90ce70fa4e5 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Sat, 13 Jun 2026 11:47:21 -0700 Subject: [PATCH 32/45] fix ci test --- python/sglang/srt/managers/io_struct.py | 16 ++++++- python/sglang/srt/managers/scheduler.py | 8 ++-- python/sglang/srt/observability/trace.py | 56 ++++++++++++++---------- 3 files changed, 52 insertions(+), 28 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 866a9ba1aa29..fe5b45dfb177 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -47,6 +47,10 @@ MetricsCollectorWrapper, SchedulerReqTimeStats, ) +from sglang.srt.observability.trace import ( + TraceContext, + TraceSpanContext, +) from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.server_args import ServerArgs from sglang.srt.utils import ImageData, VideoData @@ -826,7 +830,6 @@ class TokenizedGenerateReqInput(BaseReqIpc, kw_only=True): multi_item_delimiter_indices: Optional[List[int]] = None # For observability - # Optional[] time_stats: Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] = None @@ -2176,6 +2179,17 @@ def dec_hook(tp: Type, obj: Any) -> Any: res = array(typecode) res.frombytes(raw_data) return res + elif tp is TraceContext: + if isinstance(obj, dict): + obj = {k: pickle.loads(v) if v is not None else v for k, v in obj.items()} + return TraceContext(obj) + elif tp is TraceSpanContext: + trace_id, span_id, is_remote, trace_flags, trace_state, _ = obj + if trace_flags is not None: + trace_flags = pickle.loads(trace_flags) + if trace_state is not None: + trace_state = pickle.loads(trace_state) + return TraceSpanContext(trace_id, span_id, is_remote, trace_flags, trace_state) else: if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): logger.info( diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index b331593687b9..43a46d25c39a 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3681,7 +3681,7 @@ def handle_rpc_request(self, recv_req: RpcReqInput): logger.error(f"Failed to call rpc {recv_req.method}: {str(e)}") barrier() - return RpcReqOutput(success, "" if not exec else str(exec)) + return RpcReqOutput(success=success, message="" if not exec else str(exec)) def abort_request(self, recv_req: AbortReq): # todo hisparse, release resources for abort requests in hisparse coordinator @@ -3899,14 +3899,16 @@ def init_weights_send_group_for_remote_instance( success, message = self.tp_worker.init_weights_send_group_for_remote_instance( recv_req ) - return InitWeightsSendGroupForRemoteInstanceReqOutput(success, message) + return InitWeightsSendGroupForRemoteInstanceReqOutput( + success=success, message=message + ) def send_weights_to_remote_instance( self, recv_req: SendWeightsToRemoteInstanceReqInput ): """Send the seed instance weights to the destination instance.""" success, message = self.tp_worker.send_weights_to_remote_instance(recv_req) - return SendWeightsToRemoteInstanceReqOutput(success, message) + return SendWeightsToRemoteInstanceReqOutput(success=success, message=message) def slow_down(self, recv_req: SlowDownReqInput): t = recv_req.forward_sleep_time diff --git a/python/sglang/srt/observability/trace.py b/python/sglang/srt/observability/trace.py index e5d1dfb9eb62..67faedd898f9 100644 --- a/python/sglang/srt/observability/trace.py +++ b/python/sglang/srt/observability/trace.py @@ -21,7 +21,6 @@ import threading import time import uuid -from dataclasses import dataclass from typing import Any, Dict, List, Mapping, Optional import msgspec @@ -62,6 +61,10 @@ _trace_context_propagator = TraceContextTextMapPropagator() + TraceSpan = trace.span.Span + TraceContext = context.Context + TraceSpanContext = trace.span.SpanContext + opentelemetry_imported = True except ImportError: @@ -69,6 +72,13 @@ class id_generator: class IdGenerator: pass + # when the opentelemetry package is not installed, + # msgspec decode still need to recognize the types. + class TracePlaceHolder: + pass + + TraceSpan = TraceContext = TraceSpanContext = TracePlaceHolder + logger.debug("opentelemetry package is not installed, tracing disabled") @@ -81,39 +91,35 @@ def set_global_trace_level(level: int): global_trace_level = level -@dataclass -class TraceThreadInfo: +class TraceThreadInfo(msgspec.Struct): host_id: str pid: int thread_label: str - tp_rank: int - dp_rank: int - pp_rank: int + tp_rank: Optional[int] + dp_rank: Optional[int] + pp_rank: Optional[int] -@dataclass -class TraceEvent: +class TraceEvent(msgspec.Struct): event_name: str ts: int attrs: Dict[str, Any] -@dataclass -class TraceSliceContext: +class TraceSliceContext(msgspec.Struct): slice_name: str start_time_ns: int end_time_ns: Optional[int] = None - span: Optional[trace.span.Span] = None + span: Optional[TraceSpan] = None level: int = 1 attrs: Optional[Dict[str, Any]] = None events: Optional[List[TraceEvent]] = None -@dataclass -class TraceThreadContext: +class TraceThreadContext(msgspec.Struct): thread_info: TraceThreadInfo cur_slice_stack: Optional[List[TraceSliceContext]] = None - thread_span: Optional[trace.span.Span] = None + thread_span: Optional[TraceSpan] = None class TraceCustomIdGenerator(id_generator.IdGenerator): @@ -262,8 +268,8 @@ def trace_set_thread_info( class TraceReqContext(msgspec.Struct, tag=True): rid: str - trace_level: int - pid: int + trace_level: int = msgspec.field(default=global_trace_level) + pid: int = msgspec.field(default_factory=lambda: threading.get_native_id()) tracing_enable: bool = False start_time_ns: Optional[int] = None thread_context: Optional[TraceThreadContext] = None @@ -271,16 +277,19 @@ class TraceReqContext(msgspec.Struct, tag=True): role: str = "unified" module_name: str = "" is_copy: bool = False - root_span: Optional[trace.span.Span] = None - root_span_context: Optional[context.Context] = None - last_span_context: Optional[trace.span.SpanContext] = None + # Optional[trace.span.Span] + root_span: Optional[TraceSpan] = None + # Optional[context.Context] + root_span_context: Optional[TraceContext] = None + # Optional[trace.span.SpanContext] + last_span_context: Optional[TraceSpanContext] = None external_trace_header: Optional[Dict[str, str]] = None events_cache: List[TraceEvent] = [] - is_initialized: bool = False def __post_init__(self): + self.rid = str(self.rid) - self.trace_level = global_trace_level + self.tracing_enable = opentelemetry_initialized and self.trace_level > 0 # Filter by --trace-modules only for explicitly named modules; contexts @@ -291,7 +300,7 @@ def __post_init__(self): and self.module_name not in global_trace_modules ): self.tracing_enable = False - + # pid has to be override after transporting the object. self.pid = threading.get_native_id() def is_tracing_enabled(self) -> bool: @@ -427,9 +436,8 @@ def copy_for_thread(self) -> TraceReqContext: prev_span_context = cur_slice.span.get_span_context() # Create new instance with shared state - copied = TraceReqContext.__new__(TraceReqContext) + copied = TraceReqContext(self.rid) copied.tracing_enable = self.tracing_enable - copied.rid = self.rid copied.bootstrap_room = self.bootstrap_room copied.start_time_ns = self.start_time_ns copied.role = self.role From 25db6a7a459b4e5f41c3b01500606337d54fceca Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Sat, 13 Jun 2026 17:09:13 -0700 Subject: [PATCH 33/45] fix ci test --- python/sglang/srt/server_args.py | 9 +++++++++ test/registered/unit/observability/test_trace.py | 2 +- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 84fdda0f0d4b..84f6cae6042c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -922,10 +922,16 @@ class ServerArgs: # For msProbe msprobe_dump_config: Optional[str] = None + # Whether the post-init processing has been completed + is_initialized: bool = False + def __post_init__(self): """ Orchestrates the handling of various server arguments, ensuring proper configuration and validation. """ + # Avoid re-initialization if already done + if self.is_initialized: + return self._maybe_download_model_for_runai() @@ -1084,6 +1090,9 @@ def __post_init__(self): # Handle any other necessary validations. self._handle_other_validations() + # Finalize initialization. + self.is_initialized = True + def _maybe_download_model_for_runai(self): if is_runai_obj_uri(self.model_path): ObjectStorageModel.download_and_get_path(self.model_path) diff --git a/test/registered/unit/observability/test_trace.py b/test/registered/unit/observability/test_trace.py index 7667bce92a21..14414f1ef039 100644 --- a/test/registered/unit/observability/test_trace.py +++ b/test/registered/unit/observability/test_trace.py @@ -227,7 +227,7 @@ def test_getstate_disabled(self): self.assertEqual(state, {"tracing_enable": False}) def test_setstate_disabled(self): - ctx = TraceReqContext.__new__(TraceReqContext) + ctx = TraceReqContext(rid="req-1") ctx.__setstate__({"tracing_enable": True, "is_copy": False}) # opentelemetry_initialized is False → tracing forced off self.assertFalse(ctx.tracing_enable) From 48e8f822efb6b68bfd00f0e81e5e7b75b4a46305 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Sun, 14 Jun 2026 18:53:07 -0700 Subject: [PATCH 34/45] refactor BaseReqIpc and BaseBatchReqIpc, fix ci tests --- python/sglang/srt/managers/io_struct.py | 302 +++++++++--------- .../srt/managers/multi_tokenizer_mixin.py | 25 +- .../scheduler_components/output_sender.py | 8 +- .../sglang/srt/managers/tokenizer_manager.py | 5 +- python/sglang/srt/observability/trace.py | 6 +- python/sglang/srt/server_args.py | 3 + test/registered/core/test_srt_engine.py | 2 + test/registered/observability/test_tracing.py | 17 +- .../unit/observability/test_trace.py | 4 +- 9 files changed, 193 insertions(+), 179 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index fe5b45dfb177..d5b0b3227f13 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -22,7 +22,6 @@ import logging import pickle import uuid -from abc import ABC from array import array from collections import Counter from dataclasses import dataclass, field @@ -65,48 +64,14 @@ logger = logging.getLogger(__name__) -@dataclass -class BaseReq(ABC): - rid: Optional[Union[str, List[str]]] = field(default=None, kw_only=True) - http_worker_ipc: Optional[str] = field(default=None, kw_only=True) - - def regenerate_rid(self): - """Generate a new request ID and return it.""" - if isinstance(self.rid, list): - self.rid = [uuid.uuid4().hex for _ in range(len(self.rid))] - else: - self.rid = uuid.uuid4().hex - return self.rid - - def _validate_rid_uniqueness(self): - """Validate that request IDs within a batch are unique.""" - if isinstance(self.rid, list) and len(set(self.rid)) != len(self.rid): - counts = Counter(self.rid) - duplicates = [rid for rid, count in counts.items() if count > 1] - raise ValueError( - f"Duplicate request IDs detected within the request: {duplicates}" - ) - - # The BaseReq IPC class for IPC object -class BaseReqIpc(msgspec.Struct, tag=True, kw_only=True): +class BaseReq(msgspec.Struct, tag=True, kw_only=True): rid: Optional[str] = None http_worker_ipc: Optional[str] = None -@dataclass -class BaseBatchReq(ABC): - rids: Optional[List[str]] = field(default=None, kw_only=True) - http_worker_ipcs: Optional[List[str]] = field(default=None, kw_only=True) - - 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 - - # The BaseBatchReq IPC class for IPC object -class BaseBatchReqIpc(msgspec.Struct, tag=True, kw_only=True): +class BaseBatchReq(msgspec.Struct, tag=True, kw_only=True): rids: Optional[List[str]] = None http_worker_ipcs: Optional[List[Optional[str]]] = None @@ -139,7 +104,11 @@ class SessionParams: @dataclass -class GenerateReqInput(BaseReq): +class GenerateReqInput: + # From BaseReq dataclass + rid: Optional[Union[str, List[str]]] = field(default=None, kw_only=True) + http_worker_ipc: Optional[str] = field(default=None, kw_only=True) + # The input prompt. It can be a single prompt or a batch of prompts. text: Optional[Union[List[str], str]] = None # The token ids for text. @@ -288,6 +257,23 @@ class GenerateReqInput(BaseReq): # Batch-level: List[List[int]] (one per request). After __getitem__: List[int]. multi_item_delimiter_indices: Optional[Union[List[List[int]], List[int]]] = None + def regenerate_rid(self): + """Generate a new request ID and return it.""" + if isinstance(self.rid, list): + self.rid = [uuid.uuid4().hex for _ in range(len(self.rid))] + else: + self.rid = uuid.uuid4().hex + return self.rid + + def _validate_rid_uniqueness(self): + """Validate that request IDs within a batch are unique.""" + if isinstance(self.rid, list) and len(set(self.rid)) != len(self.rid): + counts = Counter(self.rid) + duplicates = [rid for rid, count in counts.items() if count > 1] + raise ValueError( + f"Duplicate request IDs detected within the request: {duplicates}" + ) + def contains_mm_input(self) -> bool: return ( has_valid_data(self.image_data) @@ -736,7 +722,7 @@ def __getitem__(self, i): return sub -class TokenizedGenerateReqInput(BaseReqIpc, kw_only=True): +class TokenizedGenerateReqInput(BaseReq, kw_only=True): # The input text input_text: Optional[Union[str, List[Union[str, List[str]]]]] # str # The input token ids @@ -833,7 +819,7 @@ class TokenizedGenerateReqInput(BaseReqIpc, kw_only=True): time_stats: Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] = None -class BatchTokenizedGenerateReqInput(BaseBatchReqIpc, kw_only=True): +class BatchTokenizedGenerateReqInput(BaseBatchReq, kw_only=True): # The batch of tokenized requests batch: List[TokenizedGenerateReqInput] @@ -848,7 +834,11 @@ def __iter__(self): @dataclass -class EmbeddingReqInput(BaseReq): +class EmbeddingReqInput: + # From BaseReq dataclass + rid: Optional[Union[str, List[str]]] = field(default=None, kw_only=True) + http_worker_ipc: Optional[str] = field(default=None, kw_only=True) + # The input prompt. It can be a single prompt or a batch of prompts. text: Optional[Union[List[List[str]], List[str], str]] = None # The image input. It can be an image instance, file name, URL, or base64 encoded string. @@ -916,6 +906,23 @@ class EmbeddingReqInput(BaseReq): # Batch-level: List[List[int]] (one per request). After __getitem__: List[int]. multi_item_delimiter_indices: Optional[Union[List[List[int]], List[int]]] = None + def regenerate_rid(self): + """Generate a new request ID and return it.""" + if isinstance(self.rid, list): + self.rid = [uuid.uuid4().hex for _ in range(len(self.rid))] + else: + self.rid = uuid.uuid4().hex + return self.rid + + def _validate_rid_uniqueness(self): + """Validate that request IDs within a batch are unique.""" + if isinstance(self.rid, list) and len(set(self.rid)) != len(self.rid): + counts = Counter(self.rid) + duplicates = [rid for rid, count in counts.items() if count > 1] + raise ValueError( + f"Duplicate request IDs detected within the request: {duplicates}" + ) + def normalize_batch_and_arguments(self): # at least one of text, input_ids, or image should be provided if self.text is None and self.input_ids is None and self.image_data is None: @@ -1059,7 +1066,7 @@ def __getitem__(self, i): return sub -class TokenizedEmbeddingReqInput(BaseReqIpc, kw_only=True): +class TokenizedEmbeddingReqInput(BaseReq, kw_only=True): # The input text input_text: Optional[Union[str, List[Union[str, List[str]]]]] # str # The input token ids @@ -1090,7 +1097,7 @@ class TokenizedEmbeddingReqInput(BaseReqIpc, kw_only=True): return_pooled_hidden_states: bool = False -class BatchTokenizedEmbeddingReqInput(BaseBatchReqIpc, kw_only=True): +class BatchTokenizedEmbeddingReqInput(BaseBatchReq, kw_only=True): # The batch of tokenized embedding requests batch: List[TokenizedEmbeddingReqInput] @@ -1104,7 +1111,7 @@ def __iter__(self): return iter(self.batch) -class BatchTokenIDOutput(BaseBatchReqIpc, kw_only=True): +class BatchTokenIDOutput(BaseBatchReq, kw_only=True): # The finish reason finished_reasons: List[Optional[Dict[str, Any]]] # List[BaseFinishReason] # For incremental decoding @@ -1225,7 +1232,7 @@ class BatchTokenIDOutput(BaseBatchReqIpc, kw_only=True): spec_correct_drafts_histogram: Optional[List[List[int]]] -class BatchStrOutput(BaseBatchReqIpc, kw_only=True): +class BatchStrOutput(BaseBatchReq, kw_only=True): # The finish reason finished_reasons: List[Optional[Dict[str, Any]]] # List[dict] # The output decoded strings @@ -1340,7 +1347,7 @@ class BatchStrOutput(BaseBatchReqIpc, kw_only=True): spec_correct_drafts_histogram: Optional[List[List[int]]] -class BatchEmbeddingOutput(BaseBatchReqIpc, kw_only=True): +class BatchEmbeddingOutput(BaseBatchReq, kw_only=True): # The finish reason finished_reasons: List[Optional[Dict[str, Any]]] # List[BaseFinishReason] # The output embedding @@ -1367,58 +1374,58 @@ class BatchEmbeddingOutput(BaseBatchReqIpc, kw_only=True): is_pooled_hidden_states_stacked: bool = False -class ClearHiCacheReqInput(BaseReqIpc, kw_only=True): +class ClearHiCacheReqInput(BaseReq, kw_only=True): pass -class ClearHiCacheReqOutput(BaseReqIpc, kw_only=True): +class ClearHiCacheReqOutput(BaseReq, kw_only=True): success: bool -class FlushCacheReqInput(BaseReqIpc, kw_only=True): +class FlushCacheReqInput(BaseReq, kw_only=True): timeout_s: Optional[float] = None -class FlushCacheReqOutput(BaseReqIpc, kw_only=True): +class FlushCacheReqOutput(BaseReq, kw_only=True): success: bool message: str = "" -class AddExternalCorpusReqInput(BaseReqIpc, kw_only=True): +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 -class AddExternalCorpusReqOutput(BaseReqIpc, kw_only=True): +class AddExternalCorpusReqOutput(BaseReq, kw_only=True): success: bool corpus_id: str = "" message: str = "" loaded_token_count: int = 0 -class RemoveExternalCorpusReqInput(BaseReqIpc, kw_only=True): +class RemoveExternalCorpusReqInput(BaseReq, kw_only=True): corpus_id: str -class RemoveExternalCorpusReqOutput(BaseReqIpc, kw_only=True): +class RemoveExternalCorpusReqOutput(BaseReq, kw_only=True): success: bool message: str = "" -class ListExternalCorporaReqInput(BaseReqIpc, kw_only=True): +class ListExternalCorporaReqInput(BaseReq, kw_only=True): pass -class ListExternalCorporaReqOutput(BaseReqIpc, kw_only=True): +class ListExternalCorporaReqOutput(BaseReq, kw_only=True): success: bool corpus_token_counts: Dict[str, int] = field(default_factory=dict) message: str = "" # see also AttachHiCacheStorageReqInput in the http_server.py -class AttachHiCacheStorageReqInput(BaseReqIpc, kw_only=True): +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: @@ -1432,23 +1439,23 @@ class AttachHiCacheStorageReqInput(BaseReqIpc, kw_only=True): hicache_write_policy: Optional[str] = None -class AttachHiCacheStorageReqOutput(BaseReqIpc, kw_only=True): +class AttachHiCacheStorageReqOutput(BaseReq, kw_only=True): success: bool message: str = "" -class DetachHiCacheStorageReqInput(BaseReqIpc, kw_only=True): +class DetachHiCacheStorageReqInput(BaseReq, kw_only=True): """Dynamically detach (disable) HiCache storage backend at runtime.""" pass -class DetachHiCacheStorageReqOutput(BaseReqIpc, kw_only=True): +class DetachHiCacheStorageReqOutput(BaseReq, kw_only=True): success: bool message: str = "" -class PauseGenerationReqInput(BaseReqIpc, kw_only=True): +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. @@ -1477,7 +1484,7 @@ def __post_init__(self): ) -class ContinueGenerationReqInput(BaseReqIpc, kw_only=True): +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 @@ -1486,19 +1493,19 @@ class ContinueGenerationReqInput(BaseReqIpc, kw_only=True): torch_empty_cache: bool = True -class TokenizerWorkerRegistrationReq(BaseReqIpc, kw_only=True): +class TokenizerWorkerRegistrationReq(BaseReq, kw_only=True): """Sent by each TokenizerWorker on startup to register its IPC name with the router.""" worker_ipc_name: str -class PauseContinueBroadcastReq(BaseReqIpc, kw_only=True): +class PauseContinueBroadcastReq(BaseReq, kw_only=True): """Broadcast from router to all workers to set is_pause state.""" is_pause: bool -class UpdateWeightFromDiskReqInput(BaseReqIpc, kw_only=True): +class UpdateWeightFromDiskReqInput(BaseReq, kw_only=True): # The model path with the new weights model_path: str # The format to load the weights @@ -1523,14 +1530,14 @@ class UpdateWeightFromDiskReqInput(BaseReqIpc, kw_only=True): manifest: Optional[Dict[str, Any]] = None -class UpdateWeightFromDiskReqOutput(BaseReqIpc, kw_only=True): +class UpdateWeightFromDiskReqOutput(BaseReq, kw_only=True): success: bool message: str # Number of paused requests during weight sync. num_paused_requests: Optional[int] = 0 -class UpdateWeightsFromDistributedReqInput(BaseReqIpc, kw_only=True): +class UpdateWeightsFromDistributedReqInput(BaseReq, kw_only=True): names: List[str] dtypes: List[str] shapes: List[List[int]] @@ -1548,12 +1555,12 @@ class UpdateWeightsFromDistributedReqInput(BaseReqIpc, kw_only=True): torch_empty_cache: bool = False -class UpdateWeightsFromDistributedReqOutput(BaseReqIpc, kw_only=True): +class UpdateWeightsFromDistributedReqOutput(BaseReq, kw_only=True): success: bool message: str -class UpdateWeightsFromTensorReqInput(BaseReqIpc, kw_only=True): +class UpdateWeightsFromTensorReqInput(BaseReq, kw_only=True): """Update model weights from tensor input. - Tensors are serialized for transmission @@ -1576,12 +1583,12 @@ class UpdateWeightsFromTensorReqInput(BaseReqIpc, kw_only=True): torch_empty_cache: bool = False -class UpdateWeightsFromTensorReqOutput(BaseReqIpc, kw_only=True): +class UpdateWeightsFromTensorReqOutput(BaseReq, kw_only=True): success: bool message: str -class InitWeightsSendGroupForRemoteInstanceReqInput(BaseReqIpc, kw_only=True): +class InitWeightsSendGroupForRemoteInstanceReqInput(BaseReq, kw_only=True): # The master address master_address: str # The ports for each rank's communication group @@ -1598,7 +1605,7 @@ class InitWeightsSendGroupForRemoteInstanceReqInput(BaseReqIpc, kw_only=True): # Now UpdateWeightsFromIPCReqInput and UpdateWeightsFromIPCReqOutput # are only used by Checkpoint Engine (https://github.com/MoonshotAI/checkpoint-engine) -class UpdateWeightsFromIPCReqInput(BaseReqIpc, kw_only=True): +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 @@ -1609,17 +1616,17 @@ class UpdateWeightsFromIPCReqInput(BaseReqIpc, kw_only=True): torch_empty_cache: bool = False -class UpdateWeightsFromIPCReqOutput(BaseReqIpc, kw_only=True): +class UpdateWeightsFromIPCReqOutput(BaseReq, kw_only=True): success: bool message: str -class InitWeightsSendGroupForRemoteInstanceReqOutput(BaseReqIpc, kw_only=True): +class InitWeightsSendGroupForRemoteInstanceReqOutput(BaseReq, kw_only=True): success: bool message: str -class SendWeightsToRemoteInstanceReqInput(BaseReqIpc, kw_only=True): +class SendWeightsToRemoteInstanceReqInput(BaseReq, kw_only=True): # The master address master_address: str # The ports for each rank's communication group @@ -1628,23 +1635,23 @@ class SendWeightsToRemoteInstanceReqInput(BaseReqIpc, kw_only=True): group_name: str = "weight_send_group" -class SendWeightsToRemoteInstanceReqOutput(BaseReqIpc, kw_only=True): +class SendWeightsToRemoteInstanceReqOutput(BaseReq, kw_only=True): success: bool message: str -class UpdateExpertBackupReq(BaseReqIpc, kw_only=True): +class UpdateExpertBackupReq(BaseReq, kw_only=True): pass -class BackupDramReq(BaseReqIpc, kw_only=True): +class BackupDramReq(BaseReq, kw_only=True): rank: int weight_pointer_map: Dict[str, Any] session_id: str buffer_size: int -class InitWeightsUpdateGroupReqInput(BaseReqIpc, kw_only=True): +class InitWeightsUpdateGroupReqInput(BaseReq, kw_only=True): # The master address master_address: str # The master port @@ -1659,68 +1666,68 @@ class InitWeightsUpdateGroupReqInput(BaseReqIpc, kw_only=True): backend: str = "nccl" -class InitWeightsUpdateGroupReqOutput(BaseReqIpc, kw_only=True): +class InitWeightsUpdateGroupReqOutput(BaseReq, kw_only=True): success: bool message: str -class DestroyWeightsUpdateGroupReqInput(BaseReqIpc, kw_only=True): +class DestroyWeightsUpdateGroupReqInput(BaseReq, kw_only=True): group_name: str = "weight_update_group" -class DestroyWeightsUpdateGroupReqOutput(BaseReqIpc, kw_only=True): +class DestroyWeightsUpdateGroupReqOutput(BaseReq, kw_only=True): success: bool message: str -class GetWeightsByNameReqInput(BaseReqIpc, kw_only=True): +class GetWeightsByNameReqInput(BaseReq, kw_only=True): name: str truncate_size: int = 100 -class GetWeightsByNameReqOutput(BaseReqIpc, kw_only=True): +class GetWeightsByNameReqOutput(BaseReq, kw_only=True): parameter: list -class ReleaseMemoryOccupationReqInput(BaseReqIpc, kw_only=True): +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 -class ReleaseMemoryOccupationReqOutput(BaseReqIpc, kw_only=True): +class ReleaseMemoryOccupationReqOutput(BaseReq, kw_only=True): pass -class ResumeMemoryOccupationReqInput(BaseReqIpc, kw_only=True): +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 -class ResumeMemoryOccupationReqOutput(BaseReqIpc, kw_only=True): +class ResumeMemoryOccupationReqOutput(BaseReq, kw_only=True): pass -class CheckWeightsReqInput(BaseReqIpc, kw_only=True): +class CheckWeightsReqInput(BaseReq, kw_only=True): action: str = "checksum" -class CheckWeightsReqOutput(BaseReqIpc, kw_only=True): +class CheckWeightsReqOutput(BaseReq, kw_only=True): success: bool message: str payload: Optional[Dict] = None -class SlowDownReqInput(BaseReqIpc, kw_only=True): +class SlowDownReqInput(BaseReq, kw_only=True): forward_sleep_time: Optional[float] -class SlowDownReqOutput(BaseReqIpc, kw_only=True): +class SlowDownReqOutput(BaseReq, kw_only=True): pass -class AbortReq(BaseReqIpc, kw_only=True): +class AbortReq(BaseReq, kw_only=True): # Whether to abort all requests abort_all: bool = False # The finished reason data @@ -1733,15 +1740,15 @@ def __post_init__(self): self.rid = "" -class ActiveRanksOutput(BaseReqIpc, kw_only=True): +class ActiveRanksOutput(BaseReq, kw_only=True): status: List[bool] -class GetInternalStateReq(BaseReqIpc, kw_only=True): +class GetInternalStateReq(BaseReq, kw_only=True): pass -class GetInternalStateReqOutput(BaseReqIpc, kw_only=True): +class GetInternalStateReqOutput(BaseReq, kw_only=True): server_args: ServerArgs last_gen_throughput: float memory_usage: Dict[str, Union[float, int]] @@ -1752,13 +1759,13 @@ class GetInternalStateReqOutput(BaseReqIpc, kw_only=True): # The ServerArgs has a lot of fields with different types # let's use object to make it serialize/deserialize via pickle for now -class SetInternalStateReq(BaseReqIpc, kw_only=True): +class SetInternalStateReq(BaseReq, kw_only=True): server_args: Dict[str, object] # The ServerArgs has a lot of fields with different types # let's use object to make it serialize/deserialize via pickle for now -class SetInternalStateReqOutput(BaseReqIpc, kw_only=True): +class SetInternalStateReqOutput(BaseReq, kw_only=True): updated: bool server_args: Dict[str, object] @@ -1768,7 +1775,7 @@ class ProfileReqType(Enum): STOP_PROFILE = 2 -class ProfileReq(BaseReqIpc, kw_only=True): +class ProfileReq(BaseReq, kw_only=True): req_type: ProfileReqType output_dir: Optional[str] = None start_step: Optional[int] = None @@ -1783,16 +1790,16 @@ class ProfileReq(BaseReqIpc, kw_only=True): profile_stages: Optional[List[str]] = None -class ProfileReqOutput(BaseReqIpc, kw_only=True): +class ProfileReqOutput(BaseReq, kw_only=True): success: bool message: str -class FreezeGCReq(BaseReqIpc, kw_only=True): +class FreezeGCReq(BaseReq, kw_only=True): pass -class ConfigureLoggingReq(BaseReqIpc, kw_only=True): +class ConfigureLoggingReq(BaseReq, kw_only=True): log_requests: Optional[bool] = None log_requests_level: Optional[int] = None log_requests_format: Optional[str] = None @@ -1803,23 +1810,23 @@ class ConfigureLoggingReq(BaseReqIpc, kw_only=True): dump_requests_exclude_meta_keys: Optional[List[str]] = None -class OpenSessionReqInput(BaseReqIpc, kw_only=True): +class OpenSessionReqInput(BaseReq, kw_only=True): capacity_of_str_len: int session_id: Optional[str] = None streaming: Optional[bool] = None timeout: Optional[float] = None -class CloseSessionReqInput(BaseReqIpc, kw_only=True): +class CloseSessionReqInput(BaseReq, kw_only=True): session_id: str -class OpenSessionReqOutput(BaseReqIpc, kw_only=True): +class OpenSessionReqOutput(BaseReq, kw_only=True): session_id: Optional[str] success: bool -class HealthCheckOutput(BaseReqIpc, kw_only=True): +class HealthCheckOutput(BaseReq, kw_only=True): pass @@ -1829,11 +1836,11 @@ class ExpertDistributionReqType(Enum): DUMP_RECORD = 3 -class ExpertDistributionReq(BaseReqIpc, kw_only=True): +class ExpertDistributionReq(BaseReq, kw_only=True): action: ExpertDistributionReqType -class ExpertDistributionReqOutput(BaseReqIpc, kw_only=True): +class ExpertDistributionReqOutput(BaseReq, kw_only=True): pass @@ -1850,17 +1857,17 @@ class Tool: type: Optional[str] = "function" -class RpcReqInput(BaseReqIpc, kw_only=True): +class RpcReqInput(BaseReq, kw_only=True): method: str parameters: Optional[Dict] = None -class RpcReqOutput(BaseReqIpc, kw_only=True): +class RpcReqOutput(BaseReq, kw_only=True): success: bool message: str -class LoadLoRAAdapterReqInput(BaseReqIpc, kw_only=True): +class LoadLoRAAdapterReqInput(BaseReq, kw_only=True): # The name of the lora module to newly loaded. lora_name: str # The path of loading. @@ -1879,7 +1886,7 @@ def to_ref(self) -> LoRARef: ) -class UnloadLoRAAdapterReqInput(BaseReqIpc, kw_only=True): +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`. @@ -1892,7 +1899,7 @@ def to_ref(self) -> LoRARef: ) -class LoadLoRAAdapterFromTensorsReqInput(BaseReqIpc, kw_only=True): +class LoadLoRAAdapterFromTensorsReqInput(BaseReq, kw_only=True): lora_name: str config_dict: Dict[str, Any] serialized_tensors: str @@ -1910,7 +1917,7 @@ def to_ref(self) -> LoRARef: ) -class LoRAUpdateOutput(BaseReqIpc, kw_only=True): +class LoRAUpdateOutput(BaseReq, kw_only=True): success: bool error_message: Optional[str] = None loaded_adapters: Optional[Dict[str, Union[str, LoRARef]]] = None @@ -1926,7 +1933,7 @@ class BlockReqType(Enum): UNBLOCK = 2 -class BlockReqInput(BaseReqIpc, kw_only=True): +class BlockReqInput(BaseReq, kw_only=True): req_type: BlockReqType @@ -2014,7 +2021,7 @@ class QueueMetrics: retracted: int = field(metadata={"metric": ("gauge", "Retracted requests count")}) -class GetLoadsReqInput(BaseReqIpc, kw_only=True): +class GetLoadsReqInput(BaseReq, kw_only=True): """Request for /v1/loads endpoint.""" VALID_SECTIONS = frozenset( @@ -2035,7 +2042,7 @@ def __post_init__(self): ) -class GetLoadsReqOutput(BaseReqIpc, kw_only=True): +class GetLoadsReqOutput(BaseReq, kw_only=True): """Per-DP-rank load metrics for /v1/loads endpoint.""" dp_rank: Optional[int] @@ -2075,32 +2082,32 @@ class GetLoadsReqOutput(BaseReqIpc, kw_only=True): queues: Optional[QueueMetrics] = None -class WatchLoadUpdateReq(BaseReqIpc, kw_only=True): +class WatchLoadUpdateReq(BaseReq, kw_only=True): loads: List[GetLoadsReqOutput] -class SetInjectDumpMetadataReqInput(BaseReqIpc, kw_only=True): +class SetInjectDumpMetadataReqInput(BaseReq, kw_only=True): dump_metadata: Dict[str, Any] -class SetInjectDumpMetadataReqOutput(BaseReqIpc, kw_only=True): +class SetInjectDumpMetadataReqOutput(BaseReq, kw_only=True): success: bool -class LazyDumpTensorsReqInput(BaseReqIpc, kw_only=True): +class LazyDumpTensorsReqInput(BaseReq, kw_only=True): pass -class LazyDumpTensorsReqOutput(BaseReqIpc, kw_only=True): +class LazyDumpTensorsReqOutput(BaseReq, kw_only=True): success: bool -class DumperControlReqInput(BaseReqIpc, kw_only=True): +class DumperControlReqInput(BaseReq, kw_only=True): method: str body: Dict[str, Any] -class DumperControlReqOutput(BaseReqIpc, kw_only=True): +class DumperControlReqOutput(BaseReq, kw_only=True): success: bool response: List[Dict[str, Any]] error: str = "" @@ -2115,17 +2122,14 @@ def _check_all_req_types(): for class_type in all_classes: # check its name name = class_type[0] + # GenerateReqInput and EmbeddingReqInput are not io object. + if name in ("GenerateReqInput", "EmbeddingReqInput"): + continue is_io_struct = ( - name.endswith("Req") - or name.endswith("Input") - or name.endswith("Output") - or name.endswith("Ipc") # allow IPC-specific struct naming + name.endswith("Req") or name.endswith("Input") or name.endswith("Output") ) - is_base_req = ( - issubclass(class_type[1], BaseReq) - or issubclass(class_type[1], BaseBatchReq) - or issubclass(class_type[1], BaseReqIpc) - or issubclass(class_type[1], BaseBatchReqIpc) + is_base_req = issubclass(class_type[1], BaseReq) or issubclass( + class_type[1], BaseBatchReq ) if is_io_struct and not is_base_req: raise ValueError(f"{name} is not a subclass of BaseReq or BaseBatchReq.") @@ -2139,7 +2143,7 @@ def _check_all_req_types(): # Below are msgpack serialization and ipc utils -class PickleWrapperIpc(msgspec.Struct, tag=True): +class PickleWrapper(msgspec.Struct, tag=True): data: bytes @@ -2207,20 +2211,28 @@ def dec_hook(tp: Type, obj: Any) -> Any: _struct_types = tuple( cls - for cls in BaseReqIpc.__subclasses__() - + BaseBatchReqIpc.__subclasses__() - + [PickleWrapperIpc] + for cls in BaseReq.__subclasses__() + + BaseBatchReq.__subclasses__() + + [PickleWrapper] ) # str primitive types should encode to bytes _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) -def _maybe_wrap_pickle(obj: Any) -> PickleWrapperIpc: +# technically all the msgpack struct type should be define in this file +# in case there are other types has to be defined outside of this file +# and need to be serialized via msgpack +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) -> PickleWrapper: # for msgspec.Struct type, we can directly serialize if isinstance(obj, msgspec.Struct): return obj @@ -2229,19 +2241,17 @@ def _maybe_wrap_pickle(obj: Any) -> PickleWrapperIpc: return obj if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): - logger.info(f"Object of type {type(obj)} is wrapped via PickleWrapperIpc.") - return PickleWrapperIpc(pickle.dumps(obj)) + logger.info(f"Object of type {type(obj)} is wrapped via PickleWrapper.") + return PickleWrapper(pickle.dumps(obj)) def _maybe_unwrap_pickle(obj: Any) -> Any: - # only PickleWrapperIpc type needs to be unwrapped, + # only PickleWrapper type needs to be unwrapped, # other types can be directly used - if isinstance(obj, PickleWrapperIpc): + 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 PickleWrapperIpc." - ) + logger.info(f"Object of type {type(obj)} is unwrapped from PickleWrapper.") return obj return obj diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index 3c8da710a6cc..08c4c70f1c47 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -29,7 +29,7 @@ import threading import zlib from multiprocessing import shared_memory -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional import psutil import setproctitle @@ -40,9 +40,7 @@ from sglang.srt.managers.disagg_service import start_disagg_service from sglang.srt.managers.io_struct import ( BaseBatchReq, - BaseBatchReqIpc, BaseReq, - BaseReqIpc, BatchEmbeddingOutput, BatchStrOutput, BatchTokenIDOutput, @@ -342,13 +340,13 @@ def multi_http_worker_event_loop(self: DetokenizerManager): # Fan out the output back to the originating tokenizer worker(s). # In multi-detokenizer mode the upstream MultiDetokenizerRouter may # forward either batched or single requests, so handle both shapes. - if isinstance(recv_obj, BaseBatchReqIpc): + if isinstance(recv_obj, BaseBatchReq): for i, ipc_name in enumerate(recv_obj.http_worker_ipcs): new_output = _handle_output_by_index(output, i) self.socket_mapping.send_output( ipc_name, new_output, is_tokenizer=True ) - elif isinstance(recv_obj, BaseReqIpc): + elif isinstance(recv_obj, BaseReq): self.socket_mapping.send_output( recv_obj.http_worker_ipc, output, is_tokenizer=True ) @@ -468,9 +466,9 @@ async def handle_loop(self): await self._distribute_result_to_workers(recv_obj) async def _distribute_result_to_workers(self, recv_obj): - if isinstance(recv_obj, BaseReqIpc): + if isinstance(recv_obj, BaseReq): ipc_names = [recv_obj.http_worker_ipc] - elif isinstance(recv_obj, BaseBatchReqIpc): + elif isinstance(recv_obj, BaseBatchReq): ipc_names = recv_obj.http_worker_ipcs else: raise ValueError(f"Unknown recv_obj type: {type(recv_obj)}") @@ -514,7 +512,7 @@ def event_loop(self): continue # Single request: route by its own http_worker_ipc. - if isinstance(recv_obj, BaseReqIpc): + if isinstance(recv_obj, BaseReq): assert ( recv_obj.http_worker_ipc is not None ), f"Single req {recv_obj.rid=} missing http_worker_ipc" @@ -522,7 +520,7 @@ def event_loop(self): continue # Batch request. - if isinstance(recv_obj, BaseBatchReqIpc): + if isinstance(recv_obj, BaseBatchReq): # Idle/no-op batch (rids=[]): broadcast to all detokenizers if not recv_obj.rids: for ipc in self.ipc_name_list: @@ -656,15 +654,6 @@ async def _apply_pause_continue_broadcast(self, obj: PauseContinueBroadcastReq): self._pause_continue_future.set_result(True) self._pause_continue_future = None - def _attach_multi_http_worker_info(self, req: Union[BaseReq, BaseBatchReq]): - - if isinstance(req, BaseReq): - req.http_worker_ipc = self.tokenizer_ipc_name - elif isinstance(req, BaseBatchReq): - req.http_worker_ipcs = [self.tokenizer_ipc_name] * len(req.rids) - else: - raise ValueError(f"Unknown req type: {type(req)}") - async def print_exception_wrapper(func): """ diff --git a/python/sglang/srt/managers/scheduler_components/output_sender.py b/python/sglang/srt/managers/scheduler_components/output_sender.py index 4f63fbb561b7..d4be68081b35 100644 --- a/python/sglang/srt/managers/scheduler_components/output_sender.py +++ b/python/sglang/srt/managers/scheduler_components/output_sender.py @@ -2,7 +2,7 @@ import zmq -from sglang.srt.managers.io_struct import BaseBatchReqIpc, BaseReqIpc, sock_send +from sglang.srt.managers.io_struct import BaseBatchReq, BaseReq, sock_send class SenderWrapper: @@ -11,14 +11,14 @@ def __init__(self, socket: zmq.Socket): def send_output( self, - output: Union[BaseReqIpc, BaseBatchReqIpc], - recv_obj: Optional[Union[BaseReqIpc, BaseBatchReqIpc]] = None, + output: Union[BaseReq, BaseBatchReq], + recv_obj: Optional[Union[BaseReq, BaseBatchReq]] = None, ): if self.socket is None: return if ( - isinstance(recv_obj, BaseReqIpc) + isinstance(recv_obj, BaseReq) and recv_obj.http_worker_ipc is not None and output.http_worker_ipc is None ): diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index d68b43acebae..de5f26cdf4ee 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -600,7 +600,10 @@ async def generate_request( ) if self.server_args.tokenizer_worker_num > 1: - self._attach_multi_http_worker_info(obj) + # here the obj is either GenerateReqInput or EmbeddingReqInput, + # both of which have the http_worker_ipc attribute, + # so we can set it directly without checking the type again. + obj.http_worker_ipc = self.tokenizer_ipc_name self._init_req_state(obj, request) if self.server_args.language_only: self._handle_epd_disaggregation_encode_request(obj) diff --git a/python/sglang/srt/observability/trace.py b/python/sglang/srt/observability/trace.py index 67faedd898f9..4c451eaefa20 100644 --- a/python/sglang/srt/observability/trace.py +++ b/python/sglang/srt/observability/trace.py @@ -287,9 +287,8 @@ class TraceReqContext(msgspec.Struct, tag=True): events_cache: List[TraceEvent] = [] def __post_init__(self): - self.rid = str(self.rid) - + self.trace_level = global_trace_level self.tracing_enable = opentelemetry_initialized and self.trace_level > 0 # Filter by --trace-modules only for explicitly named modules; contexts @@ -391,7 +390,8 @@ def __getstate__(self) -> Optional[Dict[str, Any]]: return state def __setstate__(self, state: Dict[str, Any]): - self.__dict__.update(state) + for k, v in state.items(): + setattr(self, k, v) if not opentelemetry_initialized: self.tracing_enable = False if not self.tracing_enable: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 84f6cae6042c..b36c17cb0e35 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -922,6 +922,9 @@ class ServerArgs: # For msProbe msprobe_dump_config: Optional[str] = None + # For grpc + enable_grpc: bool = False + # Whether the post-init processing has been completed is_initialized: bool = False diff --git a/test/registered/core/test_srt_engine.py b/test/registered/core/test_srt_engine.py index f5c6ecfc8eac..6e0a3c391b31 100644 --- a/test/registered/core/test_srt_engine.py +++ b/test/registered/core/test_srt_engine.py @@ -121,6 +121,8 @@ def test_7_engine_offline_throughput(self): server_args = ServerArgs( model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, ) + # reset servre_args initialization from test cases + server_args.is_initialized = False bench_args = BenchArgs(num_prompts=10) result = throughput_test(server_args=server_args, bench_args=bench_args) self.assertGreater(result["total_throughput"], 3000) diff --git a/test/registered/observability/test_tracing.py b/test/registered/observability/test_tracing.py index 9ac85ca6a584..f76b93275269 100644 --- a/test/registered/observability/test_tracing.py +++ b/test/registered/observability/test_tracing.py @@ -16,13 +16,18 @@ import multiprocessing as mp import time import unittest -from dataclasses import dataclass from typing import List, Optional, Union +import msgspec import requests import zmq from sglang import Engine +from sglang.srt.managers.io_struct import ( + hook_custom_types, + sock_recv, + sock_send, +) from sglang.srt.observability.req_time_stats import RequestStage from sglang.srt.observability.trace import ( TraceReqContext, @@ -104,12 +109,14 @@ def _get_span_names_by_level(level: int) -> List[str]: ] -@dataclass -class Req: +class Req(msgspec.Struct, tag=True): rid: int req_context: Optional[Union[TraceReqContext]] = None +hook_custom_types(Req) + + def _subprocess_worker(): """Worker function for subprocess trace context propagation test. Must be at module level for pickle compatibility with spawn. @@ -121,7 +128,7 @@ def _subprocess_worker(): recv_from_main = get_zmq_socket(context, zmq.PULL, "ipc:///tmp/zmq_test.ipc", True) try: - req = recv_from_main.recv_pyobj() + req = sock_recv(recv_from_main) req.req_context.rebuild_thread_context() req.req_context.trace_slice_start("work", level=1) time.sleep(0.2) @@ -235,7 +242,7 @@ def test_context_propagate(self): req.req_context.trace_req_start() req.req_context.trace_slice_start("dispatch", level=1) time.sleep(0.2) - send_to_subproc.send_pyobj(req) + sock_send(send_to_subproc, req) req.req_context.trace_slice_end("dispatch", level=1) subproc.join() diff --git a/test/registered/unit/observability/test_trace.py b/test/registered/unit/observability/test_trace.py index 14414f1ef039..547beab659f3 100644 --- a/test/registered/unit/observability/test_trace.py +++ b/test/registered/unit/observability/test_trace.py @@ -539,7 +539,7 @@ def test_setstate_enabled(self): state = ctx.__getstate__() ctx.trace_req_finish(ts=2000) - ctx2 = TraceReqContext.__new__(TraceReqContext) + ctx2 = TraceReqContext(rid="req-2") ctx2.__setstate__(state) self.assertTrue(ctx2.tracing_enable) self.assertTrue(ctx2.is_copy) @@ -567,7 +567,7 @@ def test_setstate_with_last_span_context(self): ctx.trace_req_finish(ts=3000) self.assertIsNotNone(state.get("last_span_context")) - ctx2 = TraceReqContext.__new__(TraceReqContext) + ctx2 = TraceReqContext(rid="req-2") ctx2.__setstate__(state) self.assertIsNotNone(ctx2.last_span_context) From 526f74082dc94b112963887604013df59c2ade6a Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Mon, 15 Jun 2026 10:47:31 -0700 Subject: [PATCH 35/45] fix ci test --- python/sglang/srt/managers/io_struct.py | 7 +++---- python/sglang/srt/managers/scheduler.py | 2 +- python/sglang/srt/server_args.py | 9 --------- test/registered/core/test_srt_engine.py | 2 -- 4 files changed, 4 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index d5b0b3227f13..8c6c2e7a721f 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -51,7 +51,6 @@ TraceSpanContext, ) from sglang.srt.sampling.sampling_params import SamplingParams -from sglang.srt.server_args import ServerArgs from sglang.srt.utils import ImageData, VideoData from sglang.srt.utils.field_validators import validate_optional_list_i64_1d_2d @@ -1749,7 +1748,7 @@ class GetInternalStateReq(BaseReq, kw_only=True): class GetInternalStateReqOutput(BaseReq, kw_only=True): - server_args: ServerArgs + server_args: Dict[str, Any] last_gen_throughput: float memory_usage: Dict[str, Union[float, int]] effective_max_running_requests_per_dp: int @@ -1760,14 +1759,14 @@ class GetInternalStateReqOutput(BaseReq, kw_only=True): # The ServerArgs has a lot of fields with different types # let's use object to make it serialize/deserialize via pickle for now class SetInternalStateReq(BaseReq, kw_only=True): - server_args: Dict[str, object] + server_args: Dict[str, Any] # The ServerArgs has a lot of fields with different types # let's use object to make it serialize/deserialize via pickle for now class SetInternalStateReqOutput(BaseReq, kw_only=True): updated: bool - server_args: Dict[str, object] + server_args: Dict[str, Any] class ProfileReqType(Enum): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 43a46d25c39a..ff2a535d0c4d 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3574,7 +3574,7 @@ def flush_cache(self, empty_cache: bool = True): return success def get_internal_state(self, recv_req: GetInternalStateReq): - server_args = get_global_server_args() + server_args = vars(get_global_server_args()) last_gen_throughput = self.metrics_reporter.last_gen_throughput memory_usage = { "weight": round(self.tp_worker.model_runner.weight_load_mem_usage, 2), diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index b36c17cb0e35..c9868e4d482f 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -925,16 +925,10 @@ class ServerArgs: # For grpc enable_grpc: bool = False - # Whether the post-init processing has been completed - is_initialized: bool = False - def __post_init__(self): """ Orchestrates the handling of various server arguments, ensuring proper configuration and validation. """ - # Avoid re-initialization if already done - if self.is_initialized: - return self._maybe_download_model_for_runai() @@ -1093,9 +1087,6 @@ def __post_init__(self): # Handle any other necessary validations. self._handle_other_validations() - # Finalize initialization. - self.is_initialized = True - def _maybe_download_model_for_runai(self): if is_runai_obj_uri(self.model_path): ObjectStorageModel.download_and_get_path(self.model_path) diff --git a/test/registered/core/test_srt_engine.py b/test/registered/core/test_srt_engine.py index 6e0a3c391b31..f5c6ecfc8eac 100644 --- a/test/registered/core/test_srt_engine.py +++ b/test/registered/core/test_srt_engine.py @@ -121,8 +121,6 @@ def test_7_engine_offline_throughput(self): server_args = ServerArgs( model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, ) - # reset servre_args initialization from test cases - server_args.is_initialized = False bench_args = BenchArgs(num_prompts=10) result = throughput_test(server_args=server_args, bench_args=bench_args) self.assertGreater(result["total_throughput"], 3000) From 55a446ab7fe45801e34ea9e13af1b59bb0ae58d6 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Mon, 15 Jun 2026 11:01:56 -0700 Subject: [PATCH 36/45] remove comment --- python/sglang/srt/arg_groups/speculative_hook.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 290dd7da4db8..690b684f9dd7 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -358,7 +358,6 @@ def _handle_eagle_family(server_args: "ServerArgs") -> None: # topk > 1 + page_size > 1 needs the two-pass cascade draft-decode (shared prefix # pass + per-branch expand pass with prefix-tail dup). Only these backends implement # it; flashmla / trtllm_mla / cutlass_mla can't express the per-branch tree, so reject. - # TODO: @rainj-me, Temp allow fa4 to support the two-pass _PAGE_TREE_SPEC_BACKENDS = ("flashinfer", "fa3", "triton") if ( server_args.speculative_eagle_topk > 1 From e8ac7a047366bb4a704de3c5c2de5f9ad8a349aa Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Mon, 15 Jun 2026 14:06:24 -0700 Subject: [PATCH 37/45] support array-like to reduce msg size after serialization --- python/sglang/srt/entrypoints/http_server.py | 6 +++--- python/sglang/srt/managers/io_struct.py | 10 +++------- python/sglang/srt/managers/scheduler.py | 12 ++++++++++-- .../sglang/srt/managers/tokenizer_control_mixin.py | 2 +- 4 files changed, 17 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 3646dbae6f11..45a21ba6a682 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -1498,7 +1498,7 @@ async def load_lora_adapter(obj: LoadLoRAAdapterReqInputSpec, request: Request): LoadLoRAAdapterReqInput(**obj.__dict__), request, ) - result = msgspec.to_builtins(result) + result = msgspec.structs.asdict(result) if result.get("success", False): return ORJSONResponse( @@ -1521,7 +1521,7 @@ async def load_lora_adapter_from_tensors( LoadLoRAAdapterFromTensorsReqInput(**obj.__dict__), request, ) - result = msgspec.to_builtins(result) + result = msgspec.structs.asdict(result) if result.get("success", False): return ORJSONResponse(result, status_code=HTTPStatus.OK) @@ -1537,7 +1537,7 @@ async def unload_lora_adapter(obj: UnloadLoRAAdapterReqInputSpec, request: Reque UnloadLoRAAdapterReqInput(**obj.__dict__), request, ) - result = msgspec.to_builtins(result) + result = msgspec.structs.asdict(result) if result.get("success", False): return ORJSONResponse( diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 4fb183ee2257..5cc41f334335 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -64,13 +64,13 @@ # The BaseReq IPC class for IPC object -class BaseReq(msgspec.Struct, tag=True, kw_only=True): +class BaseReq(msgspec.Struct, tag=True, kw_only=True, array_like=True): rid: Optional[str] = None http_worker_ipc: Optional[str] = None # The BaseBatchReq IPC class for IPC object -class BaseBatchReq(msgspec.Struct, tag=True, kw_only=True): +class BaseBatchReq(msgspec.Struct, tag=True, kw_only=True, array_like=True): rids: Optional[List[str]] = None http_worker_ipcs: Optional[List[Optional[str]]] = None @@ -1782,14 +1782,10 @@ class GetInternalStateReqOutput(BaseReq, kw_only=True): step_time_dict: Optional[Dict[str, float]] = None -# The ServerArgs has a lot of fields with different types -# let's use object to make it serialize/deserialize via pickle for now class SetInternalStateReq(BaseReq, kw_only=True): server_args: Dict[str, Any] -# The ServerArgs has a lot of fields with different types -# let's use object to make it serialize/deserialize via pickle for now class SetInternalStateReqOutput(BaseReq, kw_only=True): updated: bool server_args: Dict[str, Any] @@ -2168,7 +2164,7 @@ def _check_all_req_types(): # Below are msgpack serialization and ipc utils -class PickleWrapper(msgspec.Struct, tag=True): +class PickleWrapper(msgspec.Struct, tag=True, array_like=True): data: bytes diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c1c3df24ef47..dac8b9e9c5d1 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3579,7 +3579,10 @@ def flush_cache(self, empty_cache: bool = True): return success def get_internal_state(self, recv_req: GetInternalStateReq): - server_args = vars(get_global_server_args()) + server_args = dict(vars(get_global_server_args())) + # This field is not serializable. + server_args.pop("model_config", None) + last_gen_throughput = self.metrics_reporter.last_gen_throughput memory_usage = { "weight": round(self.tp_worker.model_runner.weight_load_mem_usage, 2), @@ -3655,9 +3658,14 @@ def set_internal_state(self, recv_req: SetInternalStateReq): for k, v in server_args_dict.items(): setattr(get_global_server_args(), k, v) logger.info(f"Global server args updated! {get_global_server_args()=}") + + server_args = dict(vars(get_global_server_args())) + # This field is not serializable. + server_args.pop("model_config", None) + return SetInternalStateReqOutput( updated=True, - server_args=vars(get_global_server_args()), + server_args=server_args, ) def save_remote_model(self, **kwargs): diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index a57688eba865..dc2c5788b14e 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -801,7 +801,7 @@ async def get_internal_state(self: TokenizerManager) -> List[Dict[Any, Any]]: ) results = [] for response in responses: - res = msgspec.to_builtins(response) + res = msgspec.structs.asdict(response) server_args = res.pop("server_args", None) or {} res = {k: v for k, v in res.items() if v is not None} results.append(server_args | res) From 381c099fee1fb700e2635f1630008b826dc9e794 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Mon, 15 Jun 2026 15:02:33 -0700 Subject: [PATCH 38/45] fix ci test --- test/registered/observability/test_tracing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/registered/observability/test_tracing.py b/test/registered/observability/test_tracing.py index f76b93275269..5b063e17acac 100644 --- a/test/registered/observability/test_tracing.py +++ b/test/registered/observability/test_tracing.py @@ -18,12 +18,12 @@ import unittest from typing import List, Optional, Union -import msgspec import requests import zmq from sglang import Engine from sglang.srt.managers.io_struct import ( + BaseReq, hook_custom_types, sock_recv, sock_send, @@ -109,7 +109,7 @@ def _get_span_names_by_level(level: int) -> List[str]: ] -class Req(msgspec.Struct, tag=True): +class Req(BaseReq, kw_only=True): rid: int req_context: Optional[Union[TraceReqContext]] = None From 9760869a060b50aaf2e141b4494ea584604c5dd0 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Mon, 15 Jun 2026 20:38:02 -0700 Subject: [PATCH 39/45] refactor http2 server --- python/sglang/srt/entrypoints/engine.py | 2 +- python/sglang/srt/entrypoints/http_server.py | 192 +++++++++--------- .../sglang/srt/managers/tokenizer_manager.py | 6 +- python/sglang/srt/server_args.py | 5 +- 4 files changed, 100 insertions(+), 105 deletions(-) diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index e3d31f636395..6a2d229edb59 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -863,7 +863,7 @@ def _launch_subprocesses( for p in detoken_procs: scheduler_init_result.all_child_pids.append(p.pid) - if server_args.tokenizer_worker_num == 1 and not server_args.enable_http2: + if server_args.tokenizer_worker_num == 1: # Launch a single tokenizer manager without the router tokenizer_manager, template_manager = init_tokenizer_manager_func( server_args, port_args diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 19b8decffcbc..7d48c92c3390 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -232,32 +232,6 @@ def get_global_state() -> _GlobalState: return _global_state -async def _init_granian_worker() -> ServerArgs: - main_pid = get_main_process_id() - port_args, server_args, scheduler_info = read_from_shared_memory( - f"multi_tokenizer_args_{main_pid}" - ) - - tokenizer_manager = TokenizerManager(server_args, port_args) - template_manager = TemplateManager() - template_manager.initialize_templates( - tokenizer_manager=tokenizer_manager, - model_path=server_args.model_path, - chat_template=server_args.chat_template, - completion_template=server_args.completion_template, - ) - tokenizer_manager.max_req_input_len = scheduler_info["max_req_input_len"] - - set_global_state( - _GlobalState( - tokenizer_manager=tokenizer_manager, - template_manager=template_manager, - scheduler_info=scheduler_info, - ) - ) - return server_args - - async def init_multi_tokenizer() -> ServerArgs: """ Initialization function for multi-process tokenizer mode. @@ -315,10 +289,6 @@ async def lifespan(fast_api_app: FastAPI): server_args = fast_api_app.server_args warmup_thread_kwargs = fast_api_app.warmup_thread_kwargs thread_label = "Tokenizer" - elif envs.SGLANG_GRANIAN_PARENT_PID.get() is not None: - server_args = await _init_granian_worker() - warmup_thread_kwargs = dict(server_args=server_args) - thread_label = "Tokenizer" else: # Initialize multi-tokenizer support for worker processes server_args = await init_multi_tokenizer() @@ -2312,29 +2282,53 @@ def _close_main_process_sockets(): setattr(tm, attr, None) -def _run_granian_server(server_args: ServerArgs): - """Launch Granian with HTTP/2 support""" - from granian import Granian - from granian.constants import HTTPModes, Interfaces, Loops +async def _run_granian_server( + host, + port, + log_level, + ssl_certfile=None, + ssl_keyfile=None, + ssl_ca_certs=None, + ssl_keyfile_password=None, + ssl_verify=False, +): + """Serve the in-process ASGI app with Granian (embedded mode) over HTTP/2. + + Unlike Granian's default multi-process server, the embedded server runs a + single worker as an asyncio task inside the current process. It therefore + serves the live ``app`` object directly and reuses the already-initialized + global state (tokenizer manager, templates, ...) through the normal + single-tokenizer lifespan path -- no shared memory or worker re-init needed. + The event loop is uvloop. + """ + import signal + + from granian.constants import HTTPModes, Interfaces + from granian.server.embed import Server as GranianEmbeddedServer - granian_kwargs = dict( - target="sglang.srt.entrypoints.http_server:app", - address=server_args.host, - port=server_args.port, + server = GranianEmbeddedServer( + target=app, + address=host, + port=port, interface=Interfaces.ASGI, http=HTTPModes.auto, - loop=Loops.uvloop, - log_level=server_args.log_level_http or server_args.log_level or "info", - workers=1, - ) - - ssl_enabled = server_args.ssl_certfile and server_args.ssl_keyfile - if ssl_enabled: - granian_kwargs["ssl_cert"] = server_args.ssl_certfile - granian_kwargs["ssl_key"] = server_args.ssl_keyfile - - server = Granian(**granian_kwargs) - server.serve() + log_level=log_level, + ssl_cert=ssl_certfile, + ssl_key=ssl_keyfile, + ssl_key_password=ssl_keyfile_password, + ssl_ca=ssl_ca_certs, + ssl_client_verify=ssl_verify, + ) + + # The embedded server does not install its own signal handlers, so wire + # SIGINT/SIGTERM to a graceful stop, mirroring uvicorn's behavior. + loop = asyncio.get_running_loop() + for sig in (signal.SIGINT, signal.SIGTERM): + try: + loop.add_signal_handler(sig, server.stop) + except (NotImplementedError, ValueError): + pass + await server.serve() def _setup_and_run_http_server( @@ -2367,35 +2361,6 @@ def _setup_and_run_http_server( if server_args.enable_metrics: add_prometheus_track_response_middleware(app) - # Use Granian for HTTP/2 server - if server_args.enable_http2: - # Reuse the multi-tokenizer shared memory mechanism to pass - # init args (port_args, server_args, scheduler_info) to - # Granian workers, which are independent processes. - multi_tokenizer_args_shm = write_data_for_multi_tokenizer( - port_args, server_args, scheduler_infos[0] - ) - try: - if server_args.ssl_certfile: - logger.info( - f"SSL enabled: certfile={server_args.ssl_certfile}, " - f"keyfile={server_args.ssl_keyfile}" - ) - logger.info( - f"Starting Granian HTTP/2 server on " - f"{server_args.host}:{server_args.port}" - ) - # Propagate the main process PID via os.environ so Granian - # workers (forked or spawned) can locate the shared memory - # segment created above. - envs.SGLANG_GRANIAN_PARENT_PID.set(os.getpid()) - # _close_main_process_sockets() - _run_granian_server(server_args) - finally: - if multi_tokenizer_args_shm is not None: - multi_tokenizer_args_shm.unlink() - return - # Pass additional arguments to the lifespan function. # They will be used for additional initialization setups. if server_args.tokenizer_worker_num == 1: @@ -2447,7 +2412,25 @@ def _setup_and_run_http_server( # Listen for HTTP requests if server_args.tokenizer_worker_num == 1: - if server_args.enable_ssl_refresh: + if server_args.enable_http2: + logger.info( + f"Starting embedded Granian HTTP/2 server on " + f"{server_args.host}:{server_args.port}" + ) + ssl_verify = server_args.ssl_verify() + uvloop.run( + _run_granian_server( + host=server_args.host, + port=server_args.port, + log_level=server_args.log_level_http or server_args.log_level, + ssl_certfile=server_args.ssl_certfile, + ssl_keyfile=server_args.ssl_keyfile, + ssl_ca_certs=server_args.ssl_ca_certs, + ssl_keyfile_password=server_args.ssl_keyfile_password, + ssl_verify=ssl_verify, + ) + ) + elif server_args.enable_ssl_refresh: # Use Config/Server API for access to the SSLContext. config = uvicorn.Config( app, @@ -2516,21 +2499,40 @@ async def _run_with_ssl_refresh(): "SSL refresh will be disabled." ) - uvicorn.run( - "sglang.srt.entrypoints.http_server:app", - host=server_args.host, - port=server_args.port, - root_path=server_args.fastapi_root_path, - log_level=server_args.log_level_http or server_args.log_level, - timeout_keep_alive=envs.SGLANG_TIMEOUT_KEEP_ALIVE.get(), - timeout_worker_healthcheck=envs.SGLANG_UVICORN_WORKER_HEALTHCHECK_TIMEOUT.get(), - loop="uvloop", - workers=server_args.tokenizer_worker_num, - ssl_keyfile=server_args.ssl_keyfile, - ssl_certfile=server_args.ssl_certfile, - ssl_ca_certs=server_args.ssl_ca_certs, - ssl_keyfile_password=server_args.ssl_keyfile_password, - ) + if server_args.enable_http2: + logger.info( + f"Starting embedded Granian HTTP/2 server on " + f"{server_args.host}:{server_args.port}" + ) + ssl_verify = server_args.ssl_verify() + uvloop.run( + _run_granian_server( + host=server_args.host, + port=server_args.port, + log_level=server_args.log_level_http or server_args.log_level, + ssl_certfile=server_args.ssl_certfile, + ssl_keyfile=server_args.ssl_keyfile, + ssl_ca_certs=server_args.ssl_ca_certs, + ssl_keyfile_password=server_args.ssl_keyfile_password, + ssl_verify=ssl_verify, + ) + ) + else: + uvicorn.run( + "sglang.srt.entrypoints.http_server:app", + host=server_args.host, + port=server_args.port, + root_path=server_args.fastapi_root_path, + log_level=server_args.log_level_http or server_args.log_level, + timeout_keep_alive=envs.SGLANG_TIMEOUT_KEEP_ALIVE.get(), + timeout_worker_healthcheck=envs.SGLANG_UVICORN_WORKER_HEALTHCHECK_TIMEOUT.get(), + loop="uvloop", + workers=server_args.tokenizer_worker_num, + ssl_keyfile=server_args.ssl_keyfile, + ssl_certfile=server_args.ssl_certfile, + ssl_ca_certs=server_args.ssl_ca_certs, + ssl_keyfile_password=server_args.ssl_keyfile_password, + ) finally: if server_args.tokenizer_worker_num > 1: if multi_tokenizer_args_shm is not None: diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 10d1a16726c0..6aea8a95333b 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -380,11 +380,7 @@ def init_ipc_channels(self, port_args: PortArgs): self.recv_from_detokenizer = get_zmq_socket( context, zmq.PULL, port_args.tokenizer_ipc_name, True ) - # for http2, use the connect instead of binding - if ( - self.server_args.tokenizer_worker_num == 1 - and not self.server_args.enable_http2 - ): + if self.server_args.tokenizer_worker_num == 1: send_to_scheduler = get_zmq_socket( context, zmq.PUSH, port_args.scheduler_input_ipc_name, True ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 216c864c5ba2..a73681273d4c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -922,9 +922,6 @@ class ServerArgs: # For msProbe msprobe_dump_config: Optional[str] = None - # For grpc - enable_grpc: bool = False - def __post_init__(self): """ Orchestrates the handling of various server arguments, ensuring proper configuration and validation. @@ -8430,7 +8427,7 @@ def init_new( else: nccl_port = server_args.nccl_port - if server_args.tokenizer_worker_num == 1 and not server_args.enable_http2: + if server_args.tokenizer_worker_num == 1: tokenizer_worker_ipc_name = None else: tokenizer_worker_ipc_name = ( From 0322d27b170ed7f973af20642ba0f45d8967abb5 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Tue, 16 Jun 2026 11:51:03 -0700 Subject: [PATCH 40/45] set the granian backlog and backpressure exactly like uvicorn --- python/sglang/srt/entrypoints/http_server.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 7d48c92c3390..144aa83c4238 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -2291,6 +2291,8 @@ async def _run_granian_server( ssl_ca_certs=None, ssl_keyfile_password=None, ssl_verify=False, + backlog=2048, + backpressure=2048, ): """Serve the in-process ASGI app with Granian (embedded mode) over HTTP/2. @@ -2299,7 +2301,8 @@ async def _run_granian_server( serves the live ``app`` object directly and reuses the already-initialized global state (tokenizer manager, templates, ...) through the normal single-tokenizer lifespan path -- no shared memory or worker re-init needed. - The event loop is uvloop. + The event loop is uvloop. The default backlog and backpressure values are set + exactly like uvicorn's defaults. """ import signal @@ -2318,6 +2321,8 @@ async def _run_granian_server( ssl_key_password=ssl_keyfile_password, ssl_ca=ssl_ca_certs, ssl_client_verify=ssl_verify, + backlog=backlog, + backpressure=backpressure, ) # The embedded server does not install its own signal handlers, so wire From ccc1f050a92cb8d7bb7ac3f9df0580c33146e98e Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Tue, 16 Jun 2026 14:05:54 -0700 Subject: [PATCH 41/45] rename api_specs.py to http_api_specs.py --- .../sglang/srt/disaggregation/encode_server.py | 2 +- .../{api_specs.py => http_api_specs.py} | 0 python/sglang/srt/entrypoints/http_server.py | 14 +++++++------- python/sglang/srt/managers/io_struct.py | 16 ++++++---------- .../managers/test_profile_merger_http_api.py | 2 +- .../registered/unit/utils/test_profile_merger.py | 2 +- 6 files changed, 16 insertions(+), 20 deletions(-) rename python/sglang/srt/entrypoints/{api_specs.py => http_api_specs.py} (100%) diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index d7b6608e5e7e..6fb8f60f067c 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -27,6 +27,7 @@ from fastapi.responses import ORJSONResponse, Response from transformers import AutoProcessor +from python.sglang.srt.entrypoints.http_api_specs import ProfileReqInputSpec from sglang.srt.configs.device_config import DeviceConfig from sglang.srt.configs.load_config import LoadConfig from sglang.srt.configs.model_config import ModelConfig @@ -42,7 +43,6 @@ init_distributed_environment, initialize_model_parallel, ) -from sglang.srt.entrypoints.api_specs import ProfileReqInputSpec from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import initialize_dp_attention from sglang.srt.managers.io_struct import ( diff --git a/python/sglang/srt/entrypoints/api_specs.py b/python/sglang/srt/entrypoints/http_api_specs.py similarity index 100% rename from python/sglang/srt/entrypoints/api_specs.py rename to python/sglang/srt/entrypoints/http_api_specs.py diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 144aa83c4238..e8624923ea10 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -69,7 +69,13 @@ AnthropicMessagesRequest, ) from sglang.srt.entrypoints.anthropic.serving import AnthropicServing -from sglang.srt.entrypoints.api_specs import ( +from sglang.srt.entrypoints.engine import ( + Engine, + init_tokenizer_manager, + run_detokenizer_process, + run_scheduler_process, +) +from sglang.srt.entrypoints.http_api_specs import ( AbortReqSpec, AttachHiCacheStorageReqInputSpec, CheckWeightsReqInputSpec, @@ -100,12 +106,6 @@ UpdateWeightVersionReqInputSpec, VertexGenerateReqInputSpec, ) -from sglang.srt.entrypoints.engine import ( - Engine, - init_tokenizer_manager, - run_detokenizer_process, - run_scheduler_process, -) from sglang.srt.entrypoints.ollama.protocol import ( OllamaChatRequest, OllamaGenerateRequest, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 5cc41f334335..c8122f391d4a 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -2186,11 +2186,9 @@ def enc_hook(obj: Any) -> Any: try: return pickle.dumps(obj) except Exception as e: - if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): - logger.info( - f"Return None since failed to encode object of type {type(obj)} with pickle: {e}" - ) - return None + raise ValueError( + f"Failed to encode object of type {type(obj)} with pickle: {e}." + ) def dec_hook(tp: Type, obj: Any) -> Any: @@ -2223,11 +2221,9 @@ def dec_hook(tp: Type, obj: Any) -> Any: try: return pickle.loads(obj) except Exception as e: - if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): - logger.info( - f"Return None since failed to decode object of type {type(obj)} with pickle: {e}" - ) - return None + raise ValueError( + f"Failed to decode object of type {type(obj)} with pickle: {e}." + ) _struct_types = tuple( diff --git a/test/registered/unit/managers/test_profile_merger_http_api.py b/test/registered/unit/managers/test_profile_merger_http_api.py index ae95c4190847..b6ce1058d347 100644 --- a/test/registered/unit/managers/test_profile_merger_http_api.py +++ b/test/registered/unit/managers/test_profile_merger_http_api.py @@ -1,7 +1,7 @@ import json import unittest -from sglang.srt.entrypoints.api_specs import ProfileReqInputSpec +from sglang.srt.entrypoints.http_api_specs import ProfileReqInputSpec from sglang.test.ci.ci_register import ( register_amd_ci, register_cpu_ci, diff --git a/test/registered/unit/utils/test_profile_merger.py b/test/registered/unit/utils/test_profile_merger.py index 8521e547ee90..43439a80f835 100644 --- a/test/registered/unit/utils/test_profile_merger.py +++ b/test/registered/unit/utils/test_profile_merger.py @@ -13,7 +13,7 @@ import tempfile import unittest -from sglang.srt.entrypoints.api_specs import ProfileReqInputSpec +from sglang.srt.entrypoints.http_api_specs import ProfileReqInputSpec from sglang.srt.managers.io_struct import ProfileReq, ProfileReqType from sglang.srt.utils.profile_merger import ProfileMerger from sglang.test.ci.ci_register import ( From e5933186c965b569db66847c9fdb246366591dd8 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Tue, 16 Jun 2026 20:31:15 -0700 Subject: [PATCH 42/45] use msgpack to serialize mm_inputs and fix ci --- .../srt/disaggregation/encode_server.py | 2 +- python/sglang/srt/managers/io_struct.py | 34 ++++- python/sglang/srt/managers/mm_utils.py | 10 -- python/sglang/srt/managers/schedule_batch.py | 139 +++++++++++++++--- .../sglang/srt/multimodal/evs/evs_module.py | 3 - 5 files changed, 147 insertions(+), 41 deletions(-) diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index 6fb8f60f067c..4be83de4d1e4 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -27,7 +27,6 @@ from fastapi.responses import ORJSONResponse, Response from transformers import AutoProcessor -from python.sglang.srt.entrypoints.http_api_specs import ProfileReqInputSpec from sglang.srt.configs.device_config import DeviceConfig from sglang.srt.configs.load_config import LoadConfig from sglang.srt.configs.model_config import ModelConfig @@ -43,6 +42,7 @@ init_distributed_environment, initialize_model_parallel, ) +from sglang.srt.entrypoints.http_api_specs import ProfileReqInputSpec from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import initialize_dp_attention from sglang.srt.managers.io_struct import ( diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index c8122f391d4a..a13b78b1c1e8 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -38,7 +38,11 @@ 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 +from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalInputs, + MultimodalProcessorOutput, +) from sglang.srt.multimodal.mm_utils import has_valid_data from sglang.srt.observability.req_time_stats import ( APIServerReqTimeStats, @@ -48,6 +52,7 @@ ) from sglang.srt.observability.trace import ( TraceContext, + TraceSpan, TraceSpanContext, ) from sglang.srt.sampling.sampling_params import SamplingParams @@ -743,7 +748,7 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True): # The input token ids input_ids: Optional[array] # Optional[array[int]] # The multimodal inputs - mm_inputs: Optional[object] + mm_inputs: Optional[MultimodalProcessorOutput] # The sampling parameters sampling_params: SamplingParams # Whether to return the logprobs @@ -1087,7 +1092,7 @@ class TokenizedEmbeddingReqInput(BaseReq, kw_only=True): # The input token ids input_ids: Optional[array] # array[int] # The image inputs - image_inputs: Optional[object] + image_inputs: Optional[MultimodalInputs] # The token type ids token_type_ids: Optional[List[int]] # Dummy sampling params for compatibility @@ -2144,7 +2149,12 @@ def _check_all_req_types(): # check its name name = class_type[0] # GenerateReqInput and EmbeddingReqInput are not io object. - if name in ("GenerateReqInput", "EmbeddingReqInput"): + # SKip checking MultimodalProcessorOutput + if name in ( + "GenerateReqInput", + "EmbeddingReqInput", + "MultimodalProcessorOutput", + ): continue is_io_struct = ( name.endswith("Req") or name.endswith("Input") or name.endswith("Output") @@ -2174,11 +2184,19 @@ def enc_hook(obj: Any) -> Any: tensor_dtype = str(obj.dtype).split(".")[-1] # e.g., "float32" raw_data = obj.flatten().cpu().contiguous().view(torch.uint8).numpy().data return (obj.shape, tensor_dtype, raw_data) + if 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, array): return (obj.typecode, obj.tobytes()) elif isinstance(obj, np.floating): return float(obj) elif isinstance(obj, MetricsCollectorWrapper): + # Don't serialize the MetricsCollectorWrapper object + # due to it hold a MetricsCollector which has _thread.RLock. + return None + elif isinstance(obj, TraceSpan): + # Don't serialize the TraceSpan object due to it hold _thread.RLock. return None else: if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): @@ -2193,20 +2211,26 @@ def enc_hook(obj: Any) -> Any: def dec_hook(tp: Type, obj: Any) -> Any: if tp is torch.Tensor: - # Convert ``obj`` (which should be a ``tuple``) to a torch.Tensor shape, dtype, data = obj tensor_dtype = getattr(torch, dtype) return torch.frombuffer(data, dtype=tensor_dtype).reshape(shape) + elif tp is np.ndarray: + shape, dtype, data = obj + return np.frombuffer(data, dtype=np.dtype(dtype)).reshape(shape) elif tp is array: typecode, raw_data = obj res = array(typecode) res.frombytes(raw_data) return res elif tp is TraceContext: + # The TraceContext is represented as a dict after encoding, + # we need to decode it back to TraceContext if isinstance(obj, dict): obj = {k: pickle.loads(v) if v is not None else v for k, v in obj.items()} return TraceContext(obj) elif tp is TraceSpanContext: + # The TraceSpanContext is represented as a tuple after encoding, + # we need to decode it back to TraceSpanContext trace_id, span_id, is_remote, trace_flags, trace_state, _ = obj if trace_flags is not None: trace_flags = pickle.loads(trace_flags) diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 8e24df79d23d..515376ebc6a2 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -1675,16 +1675,6 @@ def wrap_shm_features(obj): item.precomputed_embeddings = _wrap_tensor_or_list( item.precomputed_embeddings, precomputed_hash=item_hash ) - # Use pickle to serialize the mm_inputs with SHM pointers, - # so that msgpack doesn't try to serialize the potentially - # large tensors inside mm_inputs - obj.mm_inputs = pickle.dumps(obj.mm_inputs) - - if hasattr(obj, "image_inputs"): - # Use pickle to serialize the image_inputs - # so that msgpack doesn't try to serialize the potentially - # large tensors inside mm_inputs - obj.image_inputs = pickle.dumps(obj.image_inputs) return obj diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 01cb7afb8938..8ead1f555a56 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -40,6 +40,7 @@ import logging import re from array import array +from collections.abc import MutableMapping from concurrent.futures import Future from enum import Enum, auto from functools import lru_cache @@ -56,6 +57,7 @@ Union, ) +import msgspec import numpy as np import torch @@ -237,8 +239,54 @@ class MultimodalInputFormat(Enum): PRECOMPUTED_EMBEDDING = auto() -@dataclasses.dataclass -class MultimodalDataItem: +class DataType(Enum): + TENSOR = auto() + NDARRAY = auto() + + +class TensorOrNdarray(msgspec.Struct): + """A wrapper that can hold either a torch.Tensor or a numpy.ndarray, with an indicator of which type it holds. + This is used in MultimodalDataItem to allow for flexible storage of features and precomputed embeddings, + which may be returned as either tensors or ndarrays from different multimodal processors. + """ + + data_type: DataType + tensor: Optional[torch.Tensor] = None + ndarray: Optional[np.ndarray] = None + + @property + def value(self): + return self.tensor if self.data_type == DataType.TENSOR else self.ndarray + + +class ModelSpecificData(MutableMapping): + """A custom type wrapper for model-specific data, which cannot be directly serialized via msgpack since + it may contain various types of data, e.g. tensors, lists, dicts, etc. By wrapping it with ModelSpecificData, + it will be automatically serialized via pickle, we need to sort out all known types later. + """ + + data: dict[str, Any] = {} + + def __init__(self, data: dict[str, Any] | None = None): + self.data: dict[str, Any] = dict(data) if data else {} # per-instance + + def __getitem__(self, key): + return self.data[key] + + def __setitem__(self, key, value): + self.data[key] = value + + def __delitem__(self, key): + del self.data[key] + + def __iter__(self): + return iter(self.data) + + def __len__(self): + return len(self.data) + + +class MultimodalDataItem(msgspec.Struct, kw_only=True): """ One MultimodalDataItem represents a single multimodal input (one image, one video, or one audio). For example, if there are 3 images and 1 audio, there will be 4 MultimodalDataItems. @@ -249,35 +297,84 @@ class MultimodalDataItem: """ modality: Modality - hash: int = None - pad_value: int = None - offsets: Optional[list] = None + hash: Optional[int] = None + pad_value: Optional[int] = None + offsets: Optional[List[Tuple[int, int]]] = None format: MultimodalInputFormat = MultimodalInputFormat.NORMAL # the raw features returned by processor, e.g. pixel_values or audio_features - feature: Union[torch.Tensor, np.ndarray] = None + _feature: Optional[Union[TensorOrNdarray, List[TensorOrNdarray]]] = None # the precomputed embeddings, passed as final encoder embeddings # One and only one of the feature and precomputed_embeddings will be empty - precomputed_embeddings: Optional[Union[torch.Tensor, np.ndarray]] = None + _precomputed_embeddings: Optional[TensorOrNdarray] = None # Model-specific data stored in a dictionary - model_specific_data: dict[str, Any] = dataclasses.field(default_factory=dict) + model_specific_data: ModelSpecificData = msgspec.field( + default_factory=ModelSpecificData + ) - def __getattr__(self, name: str): - if ( - "model_specific_data" in self.__dict__ - and name in self.__dict__["model_specific_data"] - ): - return self.__dict__["model_specific_data"][name] + @property + def precomputed_embeddings(self): + return ( + self._precomputed_embeddings.value if self._precomputed_embeddings else None + ) + + @precomputed_embeddings.setter + def precomputed_embeddings(self, value: Optional[Union[torch.Tensor, np.ndarray]]): + if value is None: + self._precomputed_embeddings = None + elif isinstance(value, torch.Tensor): + self._precomputed_embeddings = TensorOrNdarray( + data_type=DataType.TENSOR, tensor=value + ) + elif isinstance(value, np.ndarray): + self._precomputed_embeddings = TensorOrNdarray( + data_type=DataType.NDARRAY, ndarray=value + ) + else: + raise ValueError( + f"precomputed_embeddings must be a torch.Tensor, np.ndarray, or None, got {type(value)}" + ) + + @property + def feature(self): + return self._feature.value if self._feature else None + + @feature.setter + def feature(self, value: Optional[Union[torch.Tensor, np.ndarray]]): + if value is None: + self._feature = None + elif isinstance(value, torch.Tensor): + self._feature = TensorOrNdarray(data_type=DataType.TENSOR, tensor=value) + elif isinstance(value, np.ndarray): + self._feature = TensorOrNdarray(data_type=DataType.NDARRAY, ndarray=value) else: - raise AttributeError( - f"'{self.__class__.__name__}' object has no attribute '{name}'" + raise ValueError( + f"feature must be a torch.Tensor, np.ndarray, or None, got {type(value)}" ) + def __getattr__(self, name: str): + # runs ONLY when normal slot/field/method lookup fails + if name.startswith("__") and name.endswith("__"): + raise AttributeError(name) # never route dunders here + try: + msd = object.__getattribute__(self, "model_specific_data") # no re-entry + except AttributeError: + raise AttributeError(name) + if name in msd: + return msd[name] + raise AttributeError( + f"'{type(self).__name__}' object has no attribute '{name}'" + ) + def __setitem__(self, key: str, value: Any): - if key in self.__dict__: - self.__dict__[key] = value + if key == "feature": + self.feature = value + elif key == "precomputed_embeddings": + self.precomputed_embeddings = value + elif key in self.__struct_fields__: + setattr(self, key, value) else: self.model_specific_data[key] = value @@ -374,8 +471,7 @@ def reconstruct(self, target_device: int): self.model_specific_data[extra_key] = extra_data -@dataclasses.dataclass -class MultimodalProcessorOutput: +class MultimodalProcessorOutput(msgspec.Struct): """Raw output from multimodal processors before scheduler-side preparation (pad, hash). This is the typed replacement for the dict previously returned by @@ -458,8 +554,7 @@ def build_padded_input_ids(input_ids, mm_items: List[MultimodalDataItem]): return padded_input_ids -@dataclasses.dataclass -class MultimodalInputs: +class MultimodalInputs(msgspec.Struct): """The multimodal data related inputs.""" # items of data diff --git a/python/sglang/srt/multimodal/evs/evs_module.py b/python/sglang/srt/multimodal/evs/evs_module.py index a7ab6296e074..debc722bcd4d 100644 --- a/python/sglang/srt/multimodal/evs/evs_module.py +++ b/python/sglang/srt/multimodal/evs/evs_module.py @@ -13,7 +13,6 @@ # ============================================================================== -import dataclasses import typing from abc import ABC, abstractmethod from dataclasses import dataclass @@ -29,12 +28,10 @@ from .evs_core import compute_retention_mask, replace_offsets_with_tokens_per_frame -@dataclasses.dataclass(kw_only=True) class EVSDataItem(MultimodalDataItem): thw_grids: list[tuple[int, int, int]] -@dataclasses.dataclass(kw_only=True) class VideoEVSDataItem(EVSDataItem): pre_chunked_input_ids: torch.Tensor From 7fd3ecad7bce8242bbd3fa6d927f0490478404d8 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Tue, 16 Jun 2026 21:49:57 -0700 Subject: [PATCH 43/45] fix ci issue --- python/sglang/srt/managers/schedule_batch.py | 38 ++++++++++++++++++-- 1 file changed, 36 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 8ead1f555a56..8876ee2b14ec 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -242,6 +242,8 @@ class MultimodalInputFormat(Enum): class DataType(Enum): TENSOR = auto() NDARRAY = auto() + SHM_DATA = auto() + CUDA_IPC_PROXY = auto() class TensorOrNdarray(msgspec.Struct): @@ -253,10 +255,23 @@ class TensorOrNdarray(msgspec.Struct): data_type: DataType tensor: Optional[torch.Tensor] = None ndarray: Optional[np.ndarray] = None + shm_data: Optional[object] = ( + None # use object instead of ShmPointerMMData to avoid circular import + ) + cuda_ipc_proxy: Optional[CudaIpcTensorTransportProxy] = None @property def value(self): - return self.tensor if self.data_type == DataType.TENSOR else self.ndarray + if self.data_type == DataType.TENSOR: + return self.tensor + elif self.data_type == DataType.NDARRAY: + return self.ndarray + elif self.data_type == DataType.SHM_DATA: + return self.shm_data + elif self.data_type == DataType.CUDA_IPC_PROXY: + return self.cuda_ipc_proxy + else: + raise ValueError(f"Invalid data_type: {self.data_type}") class ModelSpecificData(MutableMapping): @@ -342,13 +357,32 @@ def feature(self): return self._feature.value if self._feature else None @feature.setter - def feature(self, value: Optional[Union[torch.Tensor, np.ndarray]]): + def feature(self, value: Optional[Union[TensorOrNdarray, List[TensorOrNdarray]]]): + from sglang.srt.managers.mm_utils import ShmPointerMMData + if value is None: self._feature = None + elif isinstance(value, list): + if len(value) > 0 and isinstance(value[0], ShmPointerMMData): + self._feature = [ + TensorOrNdarray(data_type=DataType.SHM_DATA, shm_data=item) + for item in value + ] + elif len(value) > 0 and isinstance(value[0], torch.Tensor): + self._feature = [ + TensorOrNdarray(data_type=DataType.TENSOR, tensor=item) + for item in value + ] elif isinstance(value, torch.Tensor): self._feature = TensorOrNdarray(data_type=DataType.TENSOR, tensor=value) elif isinstance(value, np.ndarray): self._feature = TensorOrNdarray(data_type=DataType.NDARRAY, ndarray=value) + elif isinstance(value, ShmPointerMMData): + self._feature = TensorOrNdarray(data_type=DataType.SHM_DATA, shm_data=value) + elif isinstance(value, CudaIpcTensorTransportProxy): + self._feature = TensorOrNdarray( + data_type=DataType.CUDA_IPC_PROXY, cuda_ipc_proxy=value + ) else: raise ValueError( f"feature must be a torch.Tensor, np.ndarray, or None, got {type(value)}" From 2ba59f9ab89aa0f578fceb2f5ff6ed2514f83b82 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Wed, 17 Jun 2026 12:57:29 -0700 Subject: [PATCH 44/45] revert the MultimodalDataItem to dataclass since lots of feature still rely on dataclass --- python/sglang/srt/managers/schedule_batch.py | 196 ++++++++++++------- 1 file changed, 127 insertions(+), 69 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 8876ee2b14ec..0dd5e3186e0f 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -246,7 +246,8 @@ class DataType(Enum): CUDA_IPC_PROXY = auto() -class TensorOrNdarray(msgspec.Struct): +@dataclasses.dataclass +class TensorOrNdarray: """A wrapper that can hold either a torch.Tensor or a numpy.ndarray, with an indicator of which type it holds. This is used in MultimodalDataItem to allow for flexible storage of features and precomputed embeddings, which may be returned as either tensors or ndarrays from different multimodal processors. @@ -273,6 +274,41 @@ def value(self): else: raise ValueError(f"Invalid data_type: {self.data_type}") + @classmethod + def from_value(cls, value: Any): + from sglang.srt.managers.mm_utils import ShmPointerMMData + + if value is None: + return None + elif isinstance(value, list): + if len(value) > 0 and isinstance(value[0], ShmPointerMMData): + return [ + TensorOrNdarray(data_type=DataType.SHM_DATA, shm_data=item) + for item in value + ] + elif len(value) > 0 and isinstance(value[0], torch.Tensor): + return [ + TensorOrNdarray(data_type=DataType.TENSOR, tensor=item) + for item in value + ] + elif isinstance(value, torch.Tensor): + return TensorOrNdarray(data_type=DataType.TENSOR, tensor=value) + elif isinstance(value, np.ndarray): + return TensorOrNdarray(data_type=DataType.NDARRAY, ndarray=value) + elif isinstance(value, ShmPointerMMData): + return TensorOrNdarray(data_type=DataType.SHM_DATA, shm_data=value) + elif isinstance(value, CudaIpcTensorTransportProxy): + return TensorOrNdarray( + data_type=DataType.CUDA_IPC_PROXY, cuda_ipc_proxy=value + ) + elif isinstance(value, TensorOrNdarray): + return value + else: + raise ValueError( + f"value must be a torch.Tensor, np.ndarray, ShmPointerMMData, " + f"CudaIpcTensorTransportProxy, List[torch.Tensor], List[ShmPointerMMData] or None, got {type(value)}" + ) + class ModelSpecificData(MutableMapping): """A custom type wrapper for model-specific data, which cannot be directly serialized via msgpack since @@ -300,8 +336,22 @@ def __iter__(self): def __len__(self): return len(self.data) + @classmethod + def from_value(cls, value: Any): + if value is None: + return ModelSpecificData() + elif isinstance(value, dict): + return ModelSpecificData(value) + elif isinstance(value, ModelSpecificData): + return value + else: + raise ValueError( + f"value must be a dict, ModelSpecificData or None, got {type(value)}" + ) + -class MultimodalDataItem(msgspec.Struct, kw_only=True): +@dataclasses.dataclass +class MultimodalDataItem: """ One MultimodalDataItem represents a single multimodal input (one image, one video, or one audio). For example, if there are 3 images and 1 audio, there will be 4 MultimodalDataItems. @@ -309,6 +359,9 @@ class MultimodalDataItem(msgspec.Struct, kw_only=True): Each item has its own hash and pad_value, enabling per-image RadixAttention caching. We put the common fields first and the model-specific fields in model_specific_data. + + Keep this class as dataclass instead of msgpack struct due to heavily usage and most fields + are not directly serialized (e.g. feature, precomputed_embeddings, model_specific_data) """ modality: Modality @@ -325,90 +378,90 @@ class MultimodalDataItem(msgspec.Struct, kw_only=True): _precomputed_embeddings: Optional[TensorOrNdarray] = None # Model-specific data stored in a dictionary - model_specific_data: ModelSpecificData = msgspec.field( + _model_specific_data: ModelSpecificData = msgspec.field( default_factory=ModelSpecificData ) + def __init__( + self, + modality: Modality, + hash: Optional[int] = None, + pad_value: Optional[int] = None, + offsets: Optional[List[Tuple[int, int]]] = None, + format: MultimodalInputFormat = MultimodalInputFormat.NORMAL, + feature: Optional[Union[torch.Tensor, np.ndarray]] = None, + precomputed_embeddings: Optional[Union[torch.Tensor, np.ndarray]] = None, + model_specific_data: Optional[Dict[str, Any]] = None, + ): + self.modality = modality + self.hash = hash + self.pad_value = pad_value + self.offsets = offsets + self.format = format + self._feature = TensorOrNdarray.from_value(feature) + self._precomputed_embeddings = TensorOrNdarray.from_value( + precomputed_embeddings + ) + self._model_specific_data = ModelSpecificData.from_value(model_specific_data) + @property def precomputed_embeddings(self): - return ( - self._precomputed_embeddings.value if self._precomputed_embeddings else None - ) + if self._precomputed_embeddings is None: + return None + elif isinstance(self._precomputed_embeddings, list): + return [t.value for t in self._precomputed_embeddings] + return self._precomputed_embeddings.value @precomputed_embeddings.setter - def precomputed_embeddings(self, value: Optional[Union[torch.Tensor, np.ndarray]]): - if value is None: - self._precomputed_embeddings = None - elif isinstance(value, torch.Tensor): - self._precomputed_embeddings = TensorOrNdarray( - data_type=DataType.TENSOR, tensor=value - ) - elif isinstance(value, np.ndarray): - self._precomputed_embeddings = TensorOrNdarray( - data_type=DataType.NDARRAY, ndarray=value - ) - else: - raise ValueError( - f"precomputed_embeddings must be a torch.Tensor, np.ndarray, or None, got {type(value)}" - ) + def precomputed_embeddings(self, value: Any): + self._precomputed_embeddings = TensorOrNdarray.from_value(value) @property def feature(self): - return self._feature.value if self._feature else None + if self._feature is None: + return None + elif isinstance(self._feature, list): + return [t.value for t in self._feature] + return self._feature.value @feature.setter - def feature(self, value: Optional[Union[TensorOrNdarray, List[TensorOrNdarray]]]): - from sglang.srt.managers.mm_utils import ShmPointerMMData + def feature(self, value: Any): + self._feature = TensorOrNdarray.from_value(value) - if value is None: - self._feature = None - elif isinstance(value, list): - if len(value) > 0 and isinstance(value[0], ShmPointerMMData): - self._feature = [ - TensorOrNdarray(data_type=DataType.SHM_DATA, shm_data=item) - for item in value - ] - elif len(value) > 0 and isinstance(value[0], torch.Tensor): - self._feature = [ - TensorOrNdarray(data_type=DataType.TENSOR, tensor=item) - for item in value - ] - elif isinstance(value, torch.Tensor): - self._feature = TensorOrNdarray(data_type=DataType.TENSOR, tensor=value) - elif isinstance(value, np.ndarray): - self._feature = TensorOrNdarray(data_type=DataType.NDARRAY, ndarray=value) - elif isinstance(value, ShmPointerMMData): - self._feature = TensorOrNdarray(data_type=DataType.SHM_DATA, shm_data=value) - elif isinstance(value, CudaIpcTensorTransportProxy): - self._feature = TensorOrNdarray( - data_type=DataType.CUDA_IPC_PROXY, cuda_ipc_proxy=value - ) + @property + def model_specific_data(self): + return self._model_specific_data.data + + @model_specific_data.setter + def model_specific_data(self, value: Any): + self._model_specific_data = ModelSpecificData.from_value(value) + + def __getattr__(self, name: str): + if ( + "_model_specific_data" in self.__dict__ + and name in self.__dict__["_model_specific_data"].data + ): + return self.__dict__["_model_specific_data"].data[name] else: - raise ValueError( - f"feature must be a torch.Tensor, np.ndarray, or None, got {type(value)}" + raise AttributeError( + f"'{self.__class__.__name__}' object has no attribute '{name}'" ) - def __getattr__(self, name: str): - # runs ONLY when normal slot/field/method lookup fails - if name.startswith("__") and name.endswith("__"): - raise AttributeError(name) # never route dunders here - try: - msd = object.__getattribute__(self, "model_specific_data") # no re-entry - except AttributeError: - raise AttributeError(name) - if name in msd: - return msd[name] - raise AttributeError( - f"'{type(self).__name__}' object has no attribute '{name}'" - ) + def __setattr__(self, name: str, value: Any): + # Route attribute assignment through the same dispatch as item + # assignment / set(), so unknown names land in model_specific_data and + # survive msgpack (de)serialization, mirroring __getattr__. + self.__setitem__(name, value) def __setitem__(self, key: str, value: Any): if key == "feature": - self.feature = value + self._feature = TensorOrNdarray.from_value(value) elif key == "precomputed_embeddings": - self.precomputed_embeddings = value - elif key in self.__struct_fields__: - setattr(self, key, value) + self._precomputed_embeddings = TensorOrNdarray.from_value(value) + elif key == "model_specific_data": + self._model_specific_data = ModelSpecificData.from_value(value) + elif key in self.__dataclass_fields__ or key in self.__dict__: + self.__dict__[key] = value else: self.model_specific_data[key] = value @@ -490,9 +543,11 @@ def has_cuda_ipc_proxy(self): def reconstruct(self, target_device: int): """materialize cuda ipc proxy tensors in-place on target_device""" if isinstance(self.feature, CudaIpcTensorTransportProxy): - self.feature = self.feature.reconstruct_on_target_device(target_device) + self._feature = TensorOrNdarray.from_value( + self.feature.reconstruct_on_target_device(target_device) + ) if isinstance(self.precomputed_embeddings, CudaIpcTensorTransportProxy): - self.precomputed_embeddings = ( + self._precomputed_embeddings = TensorOrNdarray.from_value( self.precomputed_embeddings.reconstruct_on_target_device(target_device) ) for extra_key in self.model_specific_data: @@ -594,7 +649,7 @@ class MultimodalInputs(msgspec.Struct): # items of data mm_items: List[MultimodalDataItem] padded_input_ids: Optional[List[int]] = None - image_pad_len: Optional[list] = None + image_pad_len: Optional[List[int]] = None num_image_tokens: Optional[int] = None # image @@ -622,6 +677,9 @@ class MultimodalInputs(msgspec.Struct): media_nums_per_sample: Optional[List[int]] = None visible_frame_counts: Optional[torch.Tensor] = None + data_offsets: Optional[List[int]] = None # used by nano_nemotron_vl + image_offsets: Optional[List[int]] = None # used by llava and llavavid + def release_features(self): """Release feature tensors to free GPU memory.""" for item in self.mm_items: From 8467177f9eb3ee1bfb5f0c17d1c282a2abe3b273 Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Wed, 17 Jun 2026 15:01:44 -0700 Subject: [PATCH 45/45] fix ci test --- python/sglang/srt/managers/schedule_batch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index a2395a077975..14d3b0090293 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -818,7 +818,7 @@ def merge(self, other: MultimodalInputs): [self.mrope_position_delta, other.mrope_position_delta], dim=0 ) - for key, val in other.__dict__.items(): + for key in other.__struct_fields__: if "_id" in key: # set token_ids if getattr(self, key, None) is None: