diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 9e11cd3b8570..2e6cfc4f248b 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,8 +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 - rank = _get_rank() world_size = dist.get_world_size() if dist.is_initialized() else 1 @@ -1433,13 +1434,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() @@ -1476,14 +1477,15 @@ def __init__(self, socket, debug_name: str): def __getattr__(self, method_name: str): def call(*args, **kwargs): - self._socket.send_pyobj( + sock_send( + self._socket, { "method": method_name, "args": args, "kwargs": kwargs, - } + }, ) - response = self._socket.recv_pyobj() + 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..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 @@ -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 c97b7257be07..4be83de4d1e4 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -42,9 +42,15 @@ 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 ProfileReq, ProfileReqInput, ProfileReqType +from sglang.srt.managers.io_struct import ( + ProfileReq, + ProfileReqType, + async_sock_recv, + sock_send, +) 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 @@ -2390,13 +2396,14 @@ 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, "requests": requests, "enter_time": start, - } + }, ) logger.info(f"Dispatching batch of {len(group)} {modality.name} requests") @@ -2442,7 +2449,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) @@ -2978,7 +2985,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) @@ -3197,9 +3204,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) @@ -3601,7 +3608,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) ) @@ -3822,7 +3829,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( @@ -3854,12 +3861,12 @@ 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: 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, @@ -3883,10 +3890,10 @@ 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( - type=ProfileReqType.START_PROFILE, + req_type=ProfileReqType.START_PROFILE, output_dir=obj.output_dir, start_step=obj.start_step, num_steps=obj.num_steps, @@ -3900,7 +3907,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) @@ -3929,9 +3936,9 @@ 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: - 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..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 +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 @@ -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 25d8231fe91f..ab3287961c0e 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, @@ -861,8 +863,8 @@ 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: + # Launch a single tokenizer manager without the router tokenizer_manager, template_manager = init_tokenizer_manager_func( server_args, port_args ) @@ -1223,8 +1225,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..320e04a94d8b 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(req_type=ProfileReqType.STOP_PROFILE) results = await request_manager.send_communicator_req( req, "profile_communicator", timeout=600.0 ) diff --git a/python/sglang/srt/entrypoints/http_api_specs.py b/python/sglang/srt/entrypoints/http_api_specs.py new file mode 100644 index 000000000000..7737d400af9f --- /dev/null +++ b/python/sglang/srt/entrypoints/http_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/http_server.py b/python/sglang/srt/entrypoints/http_server.py index e31cec6cbfdf..e8624923ea10 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 @@ -74,6 +75,37 @@ run_detokenizer_process, run_scheduler_process, ) +from sglang.srt.entrypoints.http_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.ollama.protocol import ( OllamaChatRequest, OllamaGenerateRequest, @@ -110,8 +142,6 @@ 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, @@ -126,13 +156,10 @@ LoadLoRAAdapterFromTensorsReqInput, LoadLoRAAdapterReqInput, OpenSessionReqInput, - ParseFunctionCallReq, PauseGenerationReqInput, - ProfileReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, SendWeightsToRemoteInstanceReqInput, - SeparateReasoningReqInput, SetInternalStateReq, SlowDownReqInput, UnloadLoRAAdapterReqInput, @@ -140,8 +167,6 @@ UpdateWeightsFromDistributedReqInput, UpdateWeightsFromIPCReqInput, UpdateWeightsFromTensorReqInput, - UpdateWeightVersionReqInput, - VertexGenerateReqInput, ) from sglang.srt.managers.multi_tokenizer_mixin import ( MultiTokenizerRouter, @@ -207,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. @@ -290,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() @@ -773,8 +768,10 @@ async def get_load(): # 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) +async def set_internal_state(obj: SetInternalStateReqSpec, request: Request): + res = await _global_state.tokenizer_manager.set_internal_state( + SetInternalStateReq(**obj.__dict__) + ) return res @@ -984,7 +981,7 @@ async def clear_hicache_storage_backend(): # }' @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. @@ -1058,10 +1055,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, @@ -1150,13 +1147,18 @@ async def dump_expert_distribution_record_async(): @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(obj, request) + ) = await _global_state.tokenizer_manager.update_weights_from_disk( + UpdateWeightFromDiskReqInput(**obj.__dict__), + request, + ) content = { "success": success, @@ -1178,13 +1180,14 @@ async def update_weights_from_disk(obj: UpdateWeightFromDiskReqInput, request: R @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( - obj, request + InitWeightsSendGroupForRemoteInstanceReqInput(**obj.__dict__), + request, ) content = {"success": success, "message": message} if success: @@ -1196,13 +1199,14 @@ async def init_weights_send_group_for_remote_instance( @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( - obj, request + SendWeightsToRemoteInstanceReqInput(**obj.__dict__), + request, ) content = {"success": success, "message": message} if success: @@ -1252,11 +1256,12 @@ async def remote_instance_transfer_engine_info(rank: int = None): @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( - obj, request + InitWeightsUpdateGroupReqInput(**obj.__dict__), + request, ) content = {"success": success, "message": message} if success: @@ -1268,13 +1273,16 @@ async def init_weights_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(obj, request) + ) = await _global_state.tokenizer_manager.destroy_weights_update_group( + DestroyWeightsUpdateGroupReqInput(**obj.__dict__), + request, + ) content = {"success": success, "message": message} return ORJSONResponse( content, status_code=200 if success else HTTPStatus.BAD_REQUEST @@ -1284,7 +1292,7 @@ async def destroy_weights_update_group( @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: @@ -1294,7 +1302,7 @@ async def update_weights_from_tensor( """ success, message = await _global_state.tokenizer_manager.update_weights_from_tensor( - obj, request + UpdateWeightsFromTensorReqInput(**obj.__dict__), request ) content = {"success": success, "message": message} @@ -1306,14 +1314,15 @@ async def update_weights_from_tensor( @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( - obj, request + UpdateWeightsFromDistributedReqInput(**obj.__dict__), + request, ) content = {"success": success, "message": message} @@ -1325,10 +1334,13 @@ async def update_weights_from_distributed( @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( - obj, request + UpdateWeightsFromIPCReqInput(**obj.__dict__), + request, ) content = {"success": success, "message": message} @@ -1342,7 +1354,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) @@ -1373,10 +1385,13 @@ async def update_weight_version(obj: UpdateWeightVersionReqInput, request: Reque @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(obj, request) + ret = await _global_state.tokenizer_manager.get_weights_by_name( + GetWeightsByNameReqInput(**obj.__dict__), + request, + ) if ret is None: return _create_error_response("Get parameter by name failed") else: @@ -1388,11 +1403,14 @@ async def get_weights_by_name(obj: GetWeightsByNameReqInput, request: Request): @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(obj, request) + await _global_state.tokenizer_manager.release_memory_occupation( + ReleaseMemoryOccupationReqInput(**obj.__dict__), + request, + ) except Exception as e: return _create_error_response(e) @@ -1400,11 +1418,14 @@ async def release_memory_occupation( @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(obj, request) + await _global_state.tokenizer_manager.resume_memory_occupation( + ResumeMemoryOccupationReqInput(**obj.__dict__), + request, + ) except Exception as e: return _create_error_response(e) @@ -1412,12 +1433,15 @@ async def resume_memory_occupation( @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(obj, request) + await _global_state.tokenizer_manager.check_weights( + CheckWeightsReqInput(**obj.__dict__), + request, + ) ) body = {"success": success, "message": message} if ranks is not None: @@ -1429,25 +1453,32 @@ async def check_weights( @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 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( + SlowDownReqInput(**obj.__dict__), + request, + ) except Exception as e: return _create_error_response(e) @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(obj, request) + result = await _global_state.tokenizer_manager.load_lora_adapter( + LoadLoRAAdapterReqInput(**obj.__dict__), + request, + ) + result = msgspec.structs.asdict(result) - if result.success: + if result.get("success", False): return ORJSONResponse( result, status_code=HTTPStatus.OK, @@ -1461,14 +1492,16 @@ async def load_lora_adapter(obj: LoadLoRAAdapterReqInput, request: Request): @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( - obj, request + LoadLoRAAdapterFromTensorsReqInput(**obj.__dict__), + request, ) + result = msgspec.structs.asdict(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) @@ -1476,11 +1509,15 @@ async def load_lora_adapter_from_tensors( @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(obj, request) + result = await _global_state.tokenizer_manager.unload_lora_adapter( + UnloadLoRAAdapterReqInput(**obj.__dict__), + request, + ) + result = msgspec.structs.asdict(result) - if result.success: + if result.get("success", False): return ORJSONResponse( result, status_code=HTTPStatus.OK, @@ -1493,10 +1530,13 @@ async def unload_lora_adapter(obj: UnloadLoRAAdapterReqInput, request: Request): @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(obj, request) + session_id = await _global_state.tokenizer_manager.open_session( + OpenSessionReqInput(**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." @@ -1507,10 +1547,13 @@ async def open_session(obj: OpenSessionReqInput, request: Request): @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(obj, request) + await _global_state.tokenizer_manager.close_session( + CloseSessionReqInput(**obj.__dict__), + request, + ) return Response(status_code=200) except Exception as e: return _create_error_response(e) @@ -1518,15 +1561,17 @@ async def close_session(obj: CloseSessionReqInput, request: Request): @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(obj) + _global_state.tokenizer_manager.configure_logging( + ConfigureLoggingReq(**obj.__dict__) + ) return Response(status_code=200) @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( @@ -1538,7 +1583,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. """ @@ -1560,7 +1605,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. """ @@ -1592,9 +1639,11 @@ async def separate_reasoning_request(obj: SeparateReasoningReqInput, request: Re @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(obj) + await _global_state.tokenizer_manager.pause_generation( + PauseGenerationReqInput(**obj.__dict__) + ) return ORJSONResponse( content={"message": "Generation paused successfully.", "status": "ok"}, status_code=200, @@ -1603,9 +1652,11 @@ async def pause_generation(obj: PauseGenerationReqInput, request: Request): @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(obj) + await _global_state.tokenizer_manager.continue_generation( + ContinueGenerationReqInput(**obj.__dict__) + ) return ORJSONResponse( content={"message": "Generation continued successfully.", "status": "ok"}, status_code=200, @@ -1942,7 +1993,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 = {} @@ -2231,29 +2282,58 @@ 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, + backlog=2048, + backpressure=2048, +): + """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. The default backlog and backpressure values are set + exactly like uvicorn's defaults. + """ + 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, + backlog=backlog, + backpressure=backpressure, + ) + + # 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( @@ -2286,35 +2366,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: @@ -2366,7 +2417,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, @@ -2435,21 +2504,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/environ.py b/python/sglang/srt/environ.py index f67f4da2bced..7b123eff00d7 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -216,6 +216,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..65910865aafc 100644 --- a/python/sglang/srt/managers/communicator.py +++ b/python/sglang/srt/managers/communicator.py @@ -7,6 +7,8 @@ import zmq +from sglang.srt.managers.io_struct import sock_send + T = TypeVar("T") @@ -41,7 +43,11 @@ async def queueing_call(self, obj: T): assert self._result_values is None if obj is not None: - self._sender.send_pyobj(obj) + 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 = [] @@ -61,7 +67,11 @@ async def watching_call(self, obj): self._result_event = asyncio.Event() if obj is not None: - self._sender.send_pyobj(obj) + 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/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index 30560c0d0653..368c43c0a487 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_recv, + sock_send, ) 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 @@ -380,11 +382,11 @@ 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 - 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 b05334deaa0a..fa9a1d262f90 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 @@ -160,10 +162,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..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 @@ -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 074ab43a270e..85a62b794b4f 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -19,26 +19,42 @@ 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 -from typing import TYPE_CHECKING, Annotated, Any, Dict, List, Literal, Optional, Union +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 +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, + MultimodalInputs, + MultimodalProcessorOutput, +) from sglang.srt.multimodal.mm_utils import has_valid_data from sglang.srt.observability.req_time_stats import ( APIServerReqTimeStats, DPControllerReqTimeStats, + MetricsCollectorWrapper, SchedulerReqTimeStats, ) +from sglang.srt.observability.trace import ( + TraceContext, + TraceSpan, + TraceSpanContext, +) from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.utils import ImageData, VideoData from sglang.srt.utils.field_validators import validate_optional_list_i64_1d_2d @@ -49,62 +65,19 @@ else: Image = Any +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}" - ) - - -@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 - - -@dataclass -class SpeculativeDecodingMetricsMixin: - """ - Mixin class containing speculative decoding metrics. +# The BaseReq IPC class for IPC object +class BaseReq(msgspec.Struct, tag=True, kw_only=True, array_like=True): + rid: Optional[str] = None + http_worker_ipc: Optional[str] = None - 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: 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] - - # 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]] +# The BaseBatchReq IPC class for IPC object +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 # Parameters for a session @@ -135,7 +108,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. @@ -284,6 +261,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) @@ -748,14 +742,13 @@ def __getitem__(self, i): return sub -@dataclass -class TokenizedGenerateReqInput(BaseReq): +class TokenizedGenerateReqInput(BaseReq, kw_only=True): # The input text - input_text: str + input_text: Optional[Union[str, List[Union[str, List[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[MultimodalProcessorOutput] # The sampling parameters sampling_params: SamplingParams # Whether to return the logprobs @@ -765,7 +758,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 @@ -780,7 +773,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 @@ -831,7 +824,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 # Encoder URL snapshot frozen at tokenizer-side dispatch time so that @@ -846,8 +839,7 @@ class TokenizedGenerateReqInput(BaseReq): time_stats: Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] = None -@dataclass -class BatchTokenizedGenerateReqInput(BaseBatchReq): +class BatchTokenizedGenerateReqInput(BaseBatchReq, kw_only=True): # The batch of tokenized requests batch: List[TokenizedGenerateReqInput] @@ -862,7 +854,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. @@ -930,6 +926,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: @@ -1073,16 +1086,15 @@ def __getitem__(self, i): return sub -@dataclass -class TokenizedEmbeddingReqInput(BaseReq): +class TokenizedEmbeddingReqInput(BaseReq, kw_only=True): # The input text - input_text: str + input_text: Optional[Union[str, List[Union[str, List[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[MultimodalInputs] # 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. @@ -1105,8 +1117,7 @@ class TokenizedEmbeddingReqInput(BaseReq): return_pooled_hidden_states: bool = False -@dataclass -class BatchTokenizedEmbeddingReqInput(BaseBatchReq): +class BatchTokenizedEmbeddingReqInput(BaseBatchReq, kw_only=True): # The batch of tokenized embedding requests batch: List[TokenizedEmbeddingReqInput] @@ -1120,16 +1131,15 @@ def __iter__(self): return iter(self.batch) -@dataclass -class BatchTokenIDOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin): +class BatchTokenIDOutput(BaseBatchReq, 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] @@ -1142,49 +1152,85 @@ 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[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 - output_hidden_states: List[List[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 # 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 # 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 @@ -1194,15 +1240,28 @@ class BatchTokenIDOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin): audio_tokens: Optional[List[int]] = None video_tokens: Optional[List[int]] = None + # From the SpeculativeDecodingMetricsMixinReq class + # Verify count: number of verification forward passes + spec_verify_ct: Optional[List[int]] -@dataclass -class BatchStrOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin): + # 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(BaseBatchReq, 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] @@ -1211,48 +1270,84 @@ 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[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 - output_hidden_states: List[List[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); # 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 # 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 @@ -1262,22 +1357,35 @@ class BatchStrOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin): audio_tokens: Optional[List[int]] = None video_tokens: Optional[List[int]] = None + # From the SpeculativeDecodingMetricsMixinReq class + # Verify count: number of verification forward passes + spec_verify_ct: Optional[List[int]] -@dataclass -class BatchEmbeddingOutput(BaseBatchReq): + # 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(BaseBatchReq, 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[Union[float, List[float]]], Dict[int, float], 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 @@ -1286,73 +1394,63 @@ 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 + is_pooled_hidden_states_stacked: bool = False -@dataclass -class ClearHiCacheReqInput(BaseReq): +class ClearHiCacheReqInput(BaseReq, kw_only=True): pass -@dataclass -class ClearHiCacheReqOutput(BaseReq): +class ClearHiCacheReqOutput(BaseReq, kw_only=True): success: bool -@dataclass -class FlushCacheReqInput(BaseReq): +class FlushCacheReqInput(BaseReq, kw_only=True): timeout_s: Optional[float] = None -@dataclass -class FlushCacheReqOutput(BaseReq): +class FlushCacheReqOutput(BaseReq, kw_only=True): success: bool message: str = "" -@dataclass -class AddExternalCorpusReqInput(BaseReq): +class AddExternalCorpusReqInput(BaseReq, kw_only=True): corpus_id: Optional[str] = None file_path: Optional[str] = None documents: Optional[List[str]] = None token_chunks: Optional[List[List[int]]] = None -@dataclass -class AddExternalCorpusReqOutput(BaseReq): +class AddExternalCorpusReqOutput(BaseReq, kw_only=True): success: bool corpus_id: str = "" message: str = "" loaded_token_count: int = 0 -@dataclass -class RemoveExternalCorpusReqInput(BaseReq): +class RemoveExternalCorpusReqInput(BaseReq, kw_only=True): corpus_id: str -@dataclass -class RemoveExternalCorpusReqOutput(BaseReq): +class RemoveExternalCorpusReqOutput(BaseReq, kw_only=True): success: bool message: str = "" -@dataclass -class ListExternalCorporaReqInput(BaseReq): +class ListExternalCorporaReqInput(BaseReq, kw_only=True): pass -@dataclass -class ListExternalCorporaReqOutput(BaseReq): +class ListExternalCorporaReqOutput(BaseReq, kw_only=True): success: bool corpus_token_counts: Dict[str, int] = field(default_factory=dict) message: str = "" -@dataclass -class AttachHiCacheStorageReqInput(BaseReq): +# see also AttachHiCacheStorageReqInput in the http_server.py +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: @@ -1365,48 +1463,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(BaseReq, kw_only=True): success: bool message: str = "" -@dataclass -class DetachHiCacheStorageReqInput(BaseReq): +class DetachHiCacheStorageReqInput(BaseReq, kw_only=True): """Dynamically detach (disable) HiCache storage backend at runtime.""" pass -@dataclass -class DetachHiCacheStorageReqOutput(BaseReq): +class DetachHiCacheStorageReqOutput(BaseReq, kw_only=True): success: bool message: str = "" -@dataclass -class PauseGenerationReqInput(BaseReq): +class PauseGenerationReqInput(BaseReq, kw_only=True): """ Note that the PauseGenerationRequests is only supported in SGLang Server. abort: Abort and return all requests currently being processed. @@ -1425,7 +1499,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"] @@ -1435,8 +1509,7 @@ def __post_init__(self): ) -@dataclass -class ContinueGenerationReqInput(BaseReq): +class ContinueGenerationReqInput(BaseReq, kw_only=True): # Call torch.cuda.empty_cache() before un-pausing. Returns blocks # cached by the PyTorch allocator (left over from transient allocs # during post-weight-update processing) back to the driver before @@ -1445,22 +1518,19 @@ class ContinueGenerationReqInput(BaseReq): torch_empty_cache: bool = True -@dataclass -class TokenizerWorkerRegistration: +class TokenizerWorkerRegistrationReq(BaseReq, kw_only=True): """Sent by each TokenizerWorker on startup to register its IPC name with the router.""" worker_ipc_name: str -@dataclass -class PauseContinueBroadcast: +class PauseContinueBroadcastReq(BaseReq, kw_only=True): """Broadcast from router to all workers to set is_pause state.""" is_pause: bool -@dataclass -class UpdateWeightFromDiskReqInput(BaseReq): +class UpdateWeightFromDiskReqInput(BaseReq, kw_only=True): # The model path with the new weights model_path: str # The format to load the weights @@ -1485,16 +1555,14 @@ class UpdateWeightFromDiskReqInput(BaseReq): manifest: Optional[Dict[str, Any]] = None -@dataclass -class UpdateWeightFromDiskReqOutput(BaseReq): +class UpdateWeightFromDiskReqOutput(BaseReq, kw_only=True): success: bool message: str # Number of paused requests during weight sync. num_paused_requests: Optional[int] = 0 -@dataclass -class UpdateWeightsFromDistributedReqInput(BaseReq): +class UpdateWeightsFromDistributedReqInput(BaseReq, kw_only=True): names: List[str] dtypes: List[str] shapes: List[List[int]] @@ -1512,21 +1580,20 @@ class UpdateWeightsFromDistributedReqInput(BaseReq): torch_empty_cache: bool = False -@dataclass -class UpdateWeightsFromDistributedReqOutput(BaseReq): +class UpdateWeightsFromDistributedReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class UpdateWeightsFromTensorReqInput(BaseReq): +class UpdateWeightsFromTensorReqInput(BaseReq, 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 @@ -1541,14 +1608,12 @@ class UpdateWeightsFromTensorReqInput(BaseReq): torch_empty_cache: bool = False -@dataclass -class UpdateWeightsFromTensorReqOutput(BaseReq): +class UpdateWeightsFromTensorReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class InitWeightsSendGroupForRemoteInstanceReqInput(BaseReq): +class InitWeightsSendGroupForRemoteInstanceReqInput(BaseReq, kw_only=True): # The master address master_address: str # The ports for each rank's communication group @@ -1565,8 +1630,7 @@ class InitWeightsSendGroupForRemoteInstanceReqInput(BaseReq): # Now UpdateWeightsFromIPCReqInput and UpdateWeightsFromIPCReqOutput # are only used by Checkpoint Engine (https://github.com/MoonshotAI/checkpoint-engine) -@dataclass -class UpdateWeightsFromIPCReqInput(BaseReq): +class UpdateWeightsFromIPCReqInput(BaseReq, kw_only=True): # ZMQ socket paths for each device UUID zmq_handles: Dict[str, str] # Whether to flush cache after weight update @@ -1577,20 +1641,17 @@ class UpdateWeightsFromIPCReqInput(BaseReq): torch_empty_cache: bool = False -@dataclass -class UpdateWeightsFromIPCReqOutput(BaseReq): +class UpdateWeightsFromIPCReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class InitWeightsSendGroupForRemoteInstanceReqOutput(BaseReq): +class InitWeightsSendGroupForRemoteInstanceReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class SendWeightsToRemoteInstanceReqInput(BaseReq): +class SendWeightsToRemoteInstanceReqInput(BaseReq, kw_only=True): # The master address master_address: str # The ports for each rank's communication group @@ -1599,31 +1660,27 @@ class SendWeightsToRemoteInstanceReqInput(BaseReq): group_name: str = "weight_send_group" -@dataclass -class SendWeightsToRemoteInstanceReqOutput(BaseReq): +class SendWeightsToRemoteInstanceReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class UpdateExpertBackupReq(BaseReq): +class UpdateExpertBackupReq(BaseReq, kw_only=True): pass -@dataclass -class BackupDramReq(BaseReq): +class BackupDramReq(BaseReq, kw_only=True): rank: int weight_pointer_map: Dict[str, Any] session_id: str buffer_size: int -@dataclass -class InitWeightsUpdateGroupReqInput(BaseReq): +class InitWeightsUpdateGroupReqInput(BaseReq, kw_only=True): # The master address master_address: str # The master port - master_port: int + master_port: Union[int, str] # The rank offset rank_offset: int # The world size @@ -1634,90 +1691,68 @@ class InitWeightsUpdateGroupReqInput(BaseReq): backend: str = "nccl" -@dataclass -class InitWeightsUpdateGroupReqOutput(BaseReq): +class InitWeightsUpdateGroupReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class DestroyWeightsUpdateGroupReqInput(BaseReq): +class DestroyWeightsUpdateGroupReqInput(BaseReq, kw_only=True): group_name: str = "weight_update_group" -@dataclass -class DestroyWeightsUpdateGroupReqOutput(BaseReq): +class DestroyWeightsUpdateGroupReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class UpdateWeightVersionReqInput(BaseReq): - # The new weight version - new_version: str - # Whether to abort all running requests before updating - abort_all_requests: bool = True - - -@dataclass -class GetWeightsByNameReqInput(BaseReq): +class GetWeightsByNameReqInput(BaseReq, kw_only=True): name: str truncate_size: int = 100 -@dataclass -class GetWeightsByNameReqOutput(BaseReq): +class GetWeightsByNameReqOutput(BaseReq, kw_only=True): parameter: list -@dataclass -class ReleaseMemoryOccupationReqInput(BaseReq): +class ReleaseMemoryOccupationReqInput(BaseReq, kw_only=True): # Optional tags to identify the memory region, which is primarily used for RL # Currently we only support `weights` and `kv_cache` tags: Optional[List[str]] = None -@dataclass -class ReleaseMemoryOccupationReqOutput(BaseReq): +class ReleaseMemoryOccupationReqOutput(BaseReq, kw_only=True): pass -@dataclass -class ResumeMemoryOccupationReqInput(BaseReq): +class ResumeMemoryOccupationReqInput(BaseReq, kw_only=True): # Optional tags to identify the memory region, which is primarily used for RL # Currently we only support `weights` and `kv_cache` tags: Optional[List[str]] = None -@dataclass -class ResumeMemoryOccupationReqOutput(BaseReq): +class ResumeMemoryOccupationReqOutput(BaseReq, kw_only=True): pass -@dataclass -class CheckWeightsReqInput(BaseReq): +class CheckWeightsReqInput(BaseReq, kw_only=True): action: str = "checksum" -@dataclass -class CheckWeightsReqOutput(BaseReq): +class CheckWeightsReqOutput(BaseReq, kw_only=True): success: bool message: str payload: Optional[Dict] = None -@dataclass -class SlowDownReqInput(BaseReq): +class SlowDownReqInput(BaseReq, kw_only=True): forward_sleep_time: Optional[float] -@dataclass -class SlowDownReqOutput(BaseReq): +class SlowDownReqOutput(BaseReq, kw_only=True): pass -@dataclass -class AbortReq(BaseReq): +class AbortReq(BaseReq, kw_only=True): # Whether to abort all requests abort_all: bool = False # The finished reason data @@ -1730,66 +1765,39 @@ def __post_init__(self): self.rid = "" -@dataclass -class ActiveRanksOutput(BaseReq): +class ActiveRanksOutput(BaseReq, kw_only=True): status: List[bool] -@dataclass -class GetInternalStateReq(BaseReq): +class GetInternalStateReq(BaseReq, kw_only=True): pass -@dataclass -class GetInternalStateReqOutput(BaseReq): - internal_state: Dict[Any, Any] +class GetInternalStateReqOutput(BaseReq, kw_only=True): + server_args: Dict[str, Any] + 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 -@dataclass -class SetInternalStateReq(BaseReq): +class SetInternalStateReq(BaseReq, kw_only=True): server_args: Dict[str, Any] -@dataclass -class SetInternalStateReqOutput(BaseReq): +class SetInternalStateReqOutput(BaseReq, kw_only=True): updated: bool server_args: Dict[str, Any] -@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 -@dataclass -class ProfileReq(BaseReq): - type: ProfileReqType +class ProfileReq(BaseReq, kw_only=True): + req_type: ProfileReqType output_dir: Optional[str] = None start_step: Optional[int] = None num_steps: Optional[int] = None @@ -1803,19 +1811,16 @@ class ProfileReq(BaseReq): profile_stages: Optional[List[str]] = None -@dataclass -class ProfileReqOutput(BaseReq): +class ProfileReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class FreezeGCReq(BaseReq): +class FreezeGCReq(BaseReq, kw_only=True): pass -@dataclass -class ConfigureLoggingReq(BaseReq): +class ConfigureLoggingReq(BaseReq, kw_only=True): log_requests: Optional[bool] = None log_requests_level: Optional[int] = None log_requests_format: Optional[str] = None @@ -1826,27 +1831,23 @@ class ConfigureLoggingReq(BaseReq): dump_requests_exclude_meta_keys: Optional[List[str]] = None -@dataclass -class OpenSessionReqInput(BaseReq): +class OpenSessionReqInput(BaseReq, kw_only=True): capacity_of_str_len: int session_id: Optional[str] = None streaming: Optional[bool] = None timeout: Optional[float] = None -@dataclass -class CloseSessionReqInput(BaseReq): +class CloseSessionReqInput(BaseReq, kw_only=True): session_id: str -@dataclass -class OpenSessionReqOutput(BaseReq): +class OpenSessionReqOutput(BaseReq, kw_only=True): session_id: Optional[str] success: bool -@dataclass -class HealthCheckOutput(BaseReq): +class HealthCheckOutput(BaseReq, kw_only=True): pass @@ -1856,13 +1857,11 @@ class ExpertDistributionReqType(Enum): DUMP_RECORD = 3 -@dataclass -class ExpertDistributionReq(BaseReq): +class ExpertDistributionReq(BaseReq, kw_only=True): action: ExpertDistributionReqType -@dataclass -class ExpertDistributionReqOutput(BaseReq): +class ExpertDistributionReqOutput(BaseReq, kw_only=True): pass @@ -1879,44 +1878,17 @@ 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 - - -@dataclass -class RpcReqInput(BaseReq): +class RpcReqInput(BaseReq, kw_only=True): method: str parameters: Optional[Dict] = None -@dataclass -class RpcReqOutput(BaseReq): +class RpcReqOutput(BaseReq, kw_only=True): success: bool message: str -@dataclass -class LoadLoRAAdapterReqInput(BaseReq): +class LoadLoRAAdapterReqInput(BaseReq, kw_only=True): # The name of the lora module to newly loaded. lora_name: str # The path of loading. @@ -1935,8 +1907,7 @@ def to_ref(self) -> LoRARef: ) -@dataclass -class UnloadLoRAAdapterReqInput(BaseReq): +class UnloadLoRAAdapterReqInput(BaseReq, kw_only=True): # The name of lora module to unload. lora_name: str # The unique identifier for the LoRA adapter, which automatically generated in the `TokenizerManager`. @@ -1949,8 +1920,7 @@ def to_ref(self) -> LoRARef: ) -@dataclass -class LoadLoRAAdapterFromTensorsReqInput(BaseReq): +class LoadLoRAAdapterFromTensorsReqInput(BaseReq, kw_only=True): lora_name: str config_dict: Dict[str, Any] serialized_tensors: str @@ -1968,11 +1938,10 @@ def to_ref(self) -> LoRARef: ) -@dataclass -class LoRAUpdateOutput(BaseReq): +class LoRAUpdateOutput(BaseReq, 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 = ( @@ -1985,9 +1954,8 @@ class BlockReqType(Enum): UNBLOCK = 2 -@dataclass -class BlockReqInput(BaseReq): - type: BlockReqType +class BlockReqInput(BaseReq, kw_only=True): + req_type: BlockReqType @dataclass @@ -2074,15 +2042,14 @@ class QueueMetrics: retracted: int = field(metadata={"metric": ("gauge", "Retracted requests count")}) -@dataclass -class GetLoadsReqInput(BaseReq): +class GetLoadsReqInput(BaseReq, kw_only=True): """Request for /v1/loads endpoint.""" VALID_SECTIONS = frozenset( {"core", "memory", "spec", "lora", "disagg", "queues", "all"} ) - include: List[str] = field(default_factory=lambda: ["all"]) + include: List[str] = msgspec.field(default_factory=lambda: ["all"]) dp_rank: Optional[int] = None def __post_init__(self): @@ -2096,53 +2063,38 @@ def __post_init__(self): ) -@dataclass -class GetLoadsReqOutput(BaseReq): +class GetLoadsReqOutput(BaseReq, kw_only=True): """Per-DP-rank load metrics for /v1/loads endpoint.""" - dp_rank: int - timestamp: 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_waiting_uncached_tokens: int = field( - metadata={ - "metric": ( - "gauge", - "Number of uncached input tokens waiting for prefill compute", - ) - } - ) - num_used_tokens: int = field( - metadata={"metric": ("gauge", "Number of tokens in use")} - ) + dp_rank: Optional[int] + timestamp: Optional[float] + + # 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 uncached input tokens waiting for prefill compute")} + num_waiting_uncached_tokens: Optional[int] + # metadata={"metric": ("gauge", "Number of tokens in use")} + num_used_tokens: Optional[int] # num_used_tokens plus pending tokens not already allocated in the KV pool. # Used for DP balance. - num_total_tokens: int = field( - metadata={"metric": ("gauge", "Used tokens plus pending unallocated tokens")} - ) - max_total_num_tokens: int = field( - metadata={"metric": ("gauge", "Maximum token capacity")} - ) + # metadata={"metric": ("gauge", "Used tokens plus pending unallocated 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 @@ -2151,34 +2103,28 @@ class GetLoadsReqOutput(BaseReq): queues: Optional[QueueMetrics] = None -@dataclass -class SetInjectDumpMetadataReqInput(BaseReq): +class SetInjectDumpMetadataReqInput(BaseReq, kw_only=True): dump_metadata: Dict[str, Any] -@dataclass -class SetInjectDumpMetadataReqOutput(BaseReq): +class SetInjectDumpMetadataReqOutput(BaseReq, kw_only=True): success: bool -@dataclass -class LazyDumpTensorsReqInput(BaseReq): +class LazyDumpTensorsReqInput(BaseReq, kw_only=True): pass -@dataclass -class LazyDumpTensorsReqOutput(BaseReq): +class LazyDumpTensorsReqOutput(BaseReq, kw_only=True): success: bool -@dataclass -class DumperControlReqInput(BaseReq): +class DumperControlReqInput(BaseReq, kw_only=True): method: str body: Dict[str, Any] -@dataclass -class DumperControlReqOutput(BaseReq): +class DumperControlReqOutput(BaseReq, kw_only=True): success: bool response: List[Dict[str, Any]] error: str = "" @@ -2193,6 +2139,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. + # SKip checking MultimodalProcessorOutput + if name in ( + "GenerateReqInput", + "EmbeddingReqInput", + "MultimodalProcessorOutput", + ): + continue is_io_struct = ( name.endswith("Req") or name.endswith("Input") or name.endswith("Output") ) @@ -2208,3 +2162,164 @@ def _check_all_req_types(): _check_all_req_types() + + +# Below are msgpack serialization and ipc utils +class PickleWrapper(msgspec.Struct, tag=True, array_like=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) + 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(): + logger.info(f"Object of type {type(obj)} is encoding with pickle.") + try: + return pickle.dumps(obj) + except Exception as e: + raise ValueError( + f"Failed to encode object of type {type(obj)} with pickle: {e}." + ) + + +def dec_hook(tp: Type, obj: Any) -> Any: + if tp is 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) + 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( + f"Object of type {type(obj)} and {str(tp)} is decoding with pickle." + ) + try: + return pickle.loads(obj) + except Exception as e: + raise ValueError( + f"Failed to decode object of type {type(obj)} with pickle: {e}." + ) + + +_struct_types = tuple( + cls + 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) + + +# 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 + # 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 PickleWrapper.") + return PickleWrapper(pickle.dumps(obj)) + + +def _maybe_unwrap_pickle(obj: Any) -> Any: + # only PickleWrapper type needs to be unwrapped, + # other types can be directly used + if isinstance(obj, PickleWrapper): + obj = pickle.loads(obj.data) + if envs.SGLANG_LOG_PICKLE_IPC_OBJECTS.get(): + logger.info(f"Object of type {type(obj)} is unwrapped from PickleWrapper.") + return obj + + return obj + + +def msgpack_encode(obj: Any) -> bytes: + return _msgpack_encoder.encode(_maybe_wrap_pickle(obj)) + + +def msgpack_decode(data: bytes) -> Any: + return _maybe_unwrap_pickle(_msgpack_decoder.decode(data)) + + +def sock_send(socket: Socket, obj: Any, flags=0): + socket.send(msgpack_encode(obj), flags=flags) + + +def sock_recv(socket: Socket, flags=0) -> Any: + data = socket.recv(flags=flags) + 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): + await socket.send(msgpack_encode(obj), flags=flags) + + +async def async_sock_recv(socket: AsyncSocket, flags=0) -> Any: + data = await socket.recv(flags=flags) + 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/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 22e43d6abf0e..bcc87e135a66 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -1747,6 +1747,7 @@ def wrap_shm_features(obj): item.precomputed_embeddings = _wrap_tensor_or_list( item.precomputed_embeddings, precomputed_hash=item_hash ) + return obj diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index d00d77138d4e..56851fa88312 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 @@ -46,9 +46,13 @@ BatchTokenIDOutput, ContinueGenerationReqInput, FreezeGCReq, - PauseContinueBroadcast, + PauseContinueBroadcastReq, PauseGenerationReqInput, - TokenizerWorkerRegistration, + TokenizerWorkerRegistrationReq, + async_sock_recv, + async_sock_send, + sock_recv, + sock_send, ) from sglang.srt.managers.load_snapshot import ( create_load_snapshot_reader, @@ -97,7 +101,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( @@ -334,7 +338,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 @@ -431,9 +435,9 @@ def _register_load_snapshot_reader(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( @@ -447,7 +451,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 @@ -456,15 +460,15 @@ 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): @@ -505,7 +509,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): @@ -599,18 +603,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): @@ -618,7 +622,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": @@ -633,15 +638,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: @@ -655,15 +660,6 @@ async def _apply_pause_continue_broadcast(self, obj: PauseContinueBroadcast): 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): """ @@ -744,11 +740,20 @@ 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, + ): self.port_args = port_args self.send_to_scheduler = send_to_scheduler - def send_pyobj(self, obj): - if isinstance(obj, BaseReq): + def send_obj(self, obj): + 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 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/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 02eef3460038..14d3b0090293 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,6 +239,117 @@ class MultimodalInputFormat(Enum): PRECOMPUTED_EMBEDDING = auto() +class DataType(Enum): + TENSOR = auto() + NDARRAY = auto() + SHM_DATA = auto() + CUDA_IPC_PROXY = auto() + + +@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. + """ + + 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): + 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}") + + @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 + 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) + + @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)}" + ) + + @dataclasses.dataclass class MultimodalDataItem: """ @@ -246,37 +359,108 @@ class MultimodalDataItem: 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 - 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 __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): + 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: Any): + self._precomputed_embeddings = TensorOrNdarray.from_value(value) + + @property + def feature(self): + 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: Any): + self._feature = TensorOrNdarray.from_value(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"] + "_model_specific_data" in self.__dict__ + and name in self.__dict__["_model_specific_data"].data ): - return self.__dict__["model_specific_data"][name] + return self.__dict__["_model_specific_data"].data[name] else: raise AttributeError( f"'{self.__class__.__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 in self.__dict__: + if key == "feature": + self._feature = TensorOrNdarray.from_value(value) + elif key == "precomputed_embeddings": + 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 @@ -359,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: @@ -374,8 +560,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,14 +643,13 @@ 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 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 @@ -493,6 +677,9 @@ class MultimodalInputs: 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: @@ -631,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: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index df9deb1f5f9a..7169b790877a 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 @@ -1607,7 +1608,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) self.flush_wrapper.check_pending() if self.external_corpus_manager is not None: @@ -3578,9 +3579,12 @@ 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 = 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), "kvcache": round( self.token_to_kv_pool_allocator.get_kvcache().mem_usage, 2 @@ -3588,24 +3592,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 - - # This field is not serializable. - ret.pop("model_config", None) - - return GetInternalStateReqOutput(internal_state=ret) + step_time_dict = self.metrics_reporter.step_time_dict + + 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 @@ -3648,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): @@ -3679,7 +3694,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 @@ -3897,14 +3912,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/managers/scheduler_components/kv_events_publisher.py b/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py index b3301cd1d185..86e530ab935d 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..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 BaseBatchReq, BaseReq +from sglang.srt.managers.io_struct import BaseBatchReq, BaseReq, sock_send class SenderWrapper: @@ -25,4 +25,4 @@ def send_output( # 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/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 8d26f7934063..1d7b998d76d0 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -208,11 +208,15 @@ 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: 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)] + is_pooled_hidden_states_stacked = True else: stacked_phs = phs_list else: @@ -232,6 +236,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/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py index 42b47b598032..90945b82d63a 100644 --- a/python/sglang/srt/managers/scheduler_components/profiler_manager.py +++ b/python/sglang/srt/managers/scheduler_components/profiler_manager.py @@ -412,7 +412,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 4c0f6fa917e9..60ad1b4b4753 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, @@ -103,7 +104,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) @@ -112,7 +113,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_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index c2dc1c1d8c15..1abfeef97559 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py @@ -116,12 +116,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, @@ -129,7 +131,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, @@ -142,7 +144,9 @@ 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.""" @@ -157,7 +161,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.""" @@ -171,11 +175,11 @@ 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) - return GetWeightsByNameReqOutput(parameter) + return GetWeightsByNameReqOutput(parameter=parameter) def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput): assert ( 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..dc2c5788b14e 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -8,6 +8,8 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import fastapi +import msgspec +import pybase64 from sglang.srt.managers.communicator import FanOutCommunicator from sglang.srt.managers.io_struct import ( @@ -336,7 +338,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 +355,7 @@ async def start_profile( async def stop_profile(self: TokenizerManager): self.auto_create_handle_loop() - req = ProfileReq(type=ProfileReqType.STOP_PROFILE) + req = ProfileReq(req_type=ProfileReqType.STOP_PROFILE) return await self._execute_profile(req) async def _execute_profile(self: TokenizerManager, req: ProfileReq): @@ -474,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: @@ -791,8 +799,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.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) # Many DP ranks - return [res.internal_state for res in responses] + return results async def set_internal_state( self: TokenizerManager, obj: SetInternalStateReq @@ -857,7 +871,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 +883,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 5312e9c3cedf..6c70859ae859 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -74,6 +74,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 @@ -373,22 +374,22 @@ 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) @@ -595,7 +596,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) @@ -1125,15 +1129,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, @@ -1174,11 +1178,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, @@ -1310,7 +1314,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( @@ -1326,7 +1330,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( @@ -1651,7 +1655,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( @@ -1662,7 +1666,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: @@ -1676,7 +1680,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( @@ -1721,7 +1725,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 @@ -1761,12 +1765,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 @@ -1813,7 +1817,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), @@ -2026,11 +2030,14 @@ 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] + + pooled_hidden_states = recv_obj.pooled_hidden_states + if recv_obj.pooled_hidden_states is not None: + 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. @@ -2696,7 +2703,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/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 diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index 0c654795d8af..24c48bca7291 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 @@ -220,17 +221,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 - metrics_collector: Optional[ - Union[SchedulerMetricsCollector, TokenizerMetricsCollector] - ] = None - trace_ctx: Union[TraceReqContext, TraceNullContext] = field( + trace_ctx: Union[TraceReqContext, TraceNullContext] = msgspec.field( default_factory=TraceNullContext ) - disagg_mode: DisaggregationMode = DisaggregationMode.NULL + 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: @@ -238,7 +260,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) @@ -262,7 +284,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: @@ -351,7 +373,6 @@ def decode_json(self, state: Dict[str, Any]): self.__setstate__(state) -@dataclass class APIServerReqTimeStats(ReqTimeStatsBase): # get by time.perf_counter() created_time: float = 0.0 @@ -514,7 +535,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 @@ -560,7 +580,6 @@ def set_dp_dispatch_finish_time(self, ts=None): ) -@dataclass class SchedulerReqTimeStats(ReqTimeStatsBase): """ Store the timestamps for each stage of a request. @@ -1168,7 +1187,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 diff --git a/python/sglang/srt/observability/trace.py b/python/sglang/srt/observability/trace.py index fb74fc4bf8f3..4c451eaefa20 100644 --- a/python/sglang/srt/observability/trace.py +++ b/python/sglang/srt/observability/trace.py @@ -21,9 +21,10 @@ import threading import time import uuid -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__) @@ -60,6 +61,10 @@ _trace_context_propagator = TraceContextTextMapPropagator() + TraceSpan = trace.span.Span + TraceContext = context.Context + TraceSpanContext = trace.span.SpanContext + opentelemetry_imported = True except ImportError: @@ -67,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") @@ -79,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): @@ -258,49 +266,41 @@ 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 = 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 + bootstrap_room: Optional[int] = None + role: str = "unified" + module_name: str = "" + is_copy: bool = False + # 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] = [] + + 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() + # pid has to be override after transporting the object. + self.pid = threading.get_native_id() def is_tracing_enabled(self) -> bool: return self.tracing_enable @@ -390,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: @@ -435,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 @@ -765,11 +765,17 @@ def __del__(self): self.abort(abort_info={"reason": "have unclosed span, auto closed"}) -@dataclass -class TraceNullContext: +# 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(msgspec.Struct, tag=True): tracing_enable: bool = False 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/python/sglang/srt/sampling/sampling_params.py b/python/sglang/srt/sampling/sampling_params.py index b5a2ce262e27..d75b31913603 100644 --- a/python/sglang/srt/sampling/sampling_params.py +++ b/python/sglang/srt/sampling/sampling_params.py @@ -17,6 +17,8 @@ import math 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 @@ -59,7 +61,7 @@ def raise_if_tokenizer_required( ) -class SamplingParams: +class SamplingParams(msgspec.Struct, kw_only=True): """ The sampling parameters. @@ -68,77 +70,84 @@ class SamplingParams: for the documentation. """ - 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: + 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 + 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.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} + + # 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 = 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.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: @@ -245,6 +254,8 @@ def normalize(self, tokenizer): tokenizer, self.stop_strs, self.stop_regex_strs, self.min_new_tokens ) + 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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 4a0eb44215a5..12b7fc73b9af 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -571,10 +571,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/srt/session/session_controller.py b/python/sglang/srt/session/session_controller.py index 874de4fdb2f5..f0dd9c7f2f6b 100644 --- a/python/sglang/srt/session/session_controller.py +++ b/python/sglang/srt/session/session_controller.py @@ -365,10 +365,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, @@ -379,7 +379,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 diff --git a/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py b/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py index 5f727e127221..744e1019d26a 100644 --- a/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py +++ b/python/sglang/test/scripted_runtime/tokenizer_recv_proxy.py @@ -11,6 +11,7 @@ BatchTokenizedGenerateReqInput, TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, + sock_recv, ) _WORK_REQ_TYPES = ( @@ -28,7 +29,7 @@ def __init__(self, *, underlying: zmq.Socket) -> None: self._buffer: deque = deque() self.work_reqs_seen: int = 0 - def recv_pyobj(self, flags: int = 0) -> Any: + def recv(self, flags: int = 0) -> Any: self._drain_underlying() if self._buffer: @@ -64,7 +65,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 if isinstance(req, _WORK_REQ_TYPES): diff --git a/test/registered/observability/test_tracing.py b/test/registered/observability/test_tracing.py index 9ac85ca6a584..5b063e17acac 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 requests import zmq from sglang import Engine +from sglang.srt.managers.io_struct import ( + BaseReq, + 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(BaseReq, kw_only=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/rl/test_update_weights_from_distributed.py b/test/registered/rl/test_update_weights_from_distributed.py index 0aa9f82a6683..9321dc967db7 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"]: 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__": 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..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.managers.io_struct import ProfileReqInput +from sglang.srt.entrypoints.http_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/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": diff --git a/test/registered/unit/observability/test_trace.py b/test/registered/unit/observability/test_trace.py index 7667bce92a21..547beab659f3 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) @@ -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) 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() diff --git a/test/registered/unit/utils/test_profile_merger.py b/test/registered/unit/utils/test_profile_merger.py index eb4cb9bcb906..43439a80f835 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.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 ( register_amd_ci, @@ -208,18 +209,18 @@ 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 - req = ProfileReq(type=ProfileReqType.START_PROFILE) + req = ProfileReq(req_type=ProfileReqType.START_PROFILE) self.assertFalse(req.merge_profiles) - req = ProfileReq(type=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):