From d5c566e59b7418f547ab6d1342b71a6e5eb5a7db Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Wed, 24 Jun 2026 11:22:58 -0700 Subject: [PATCH] Extract profile request cleanups (#29098) --- .../srt/disaggregation/encode_server.py | 54 +++---------- python/sglang/srt/entrypoints/engine.py | 17 ++++- python/sglang/srt/entrypoints/grpc_bridge.py | 6 +- python/sglang/srt/entrypoints/grpc_server.py | 4 +- python/sglang/srt/entrypoints/http_server.py | 20 +---- python/sglang/srt/managers/io_struct.py | 30 ++------ python/sglang/srt/managers/scheduler.py | 12 ++- .../scheduler_components/profiler_manager.py | 2 +- .../scheduler_components/weight_updater.py | 18 +++-- .../srt/managers/tokenizer_control_mixin.py | 45 +++++------ .../sglang/srt/managers/tokenizer_manager.py | 28 +++---- .../sglang/srt/session/session_controller.py | 6 +- python/sglang/srt/utils/common.py | 34 +++++++++ .../managers/test_profile_merger_http_api.py | 75 ++++++++++--------- .../unit/observability/test_trace.py | 6 +- .../unit/utils/test_profile_merger.py | 24 +++--- 16 files changed, 179 insertions(+), 202 deletions(-) diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index 60cc809f9..5a14c804c 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -46,7 +46,6 @@ from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import initialize_dp_attention from sglang.srt.managers.io_struct import ( ProfileReq, - ProfileReqInput, ProfileReqType, async_sock_recv, async_sock_send, @@ -2981,13 +2980,8 @@ async def _dp_worker_handle_profile( ) -> dict: prefix = f"dp_rank={dp_rank}: " if dp_type == "start_profile": - obj = request.get("profile_req") - # `is None` (not `if not obj`) so empty dict still raises. - req = ( - ProfileReq(**obj) - if obj is not None - else ProfileReq(ProfileReqType.START_PROFILE) - ) + req = request.get("profile_req") or ProfileReq() + req.req_type = ProfileReqType.START_PROFILE if enc.profiler is None: enc.profiler = EncoderProfiler(dp_rank) ok, msg = enc.profiler.start(req) @@ -3207,7 +3201,7 @@ async def run_encoder( while True: 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) @@ -3862,51 +3856,21 @@ 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[ProfileReq] = None): if dp_dispatcher is not None: - profile_req = None if obj is not None: - profile_req = { - "type": ProfileReqType.START_PROFILE, - "output_dir": obj.output_dir, - "start_step": obj.start_step, - "num_steps": obj.num_steps, - "activities": obj.activities, - "with_stack": obj.with_stack, - "record_shapes": obj.record_shapes, - "profile_by_stage": obj.profile_by_stage, - "profile_id": str(time.time()), - "merge_profiles": obj.merge_profiles, - "profile_prefix": obj.profile_prefix, - "profile_stages": obj.profile_stages, - } + obj.req_type = ProfileReqType.START_PROFILE try: results = await dp_dispatcher.broadcast( - {"_dp_type": "start_profile", "profile_req": profile_req} + {"_dp_type": "start_profile", "profile_req": obj} ) except MMError as e: return Response(content=f"{e}\n", status_code=int(e.code)) return _summarise_dp_broadcast(results) if encoder is None: return Response(content="encoder not ready\n", status_code=503) - req = None - if obj is None: - req = ProfileReq(ProfileReqType.START_PROFILE) - else: - req = ProfileReq( - type=ProfileReqType.START_PROFILE, - output_dir=obj.output_dir, - start_step=obj.start_step, - num_steps=obj.num_steps, - activities=obj.activities, - with_stack=obj.with_stack, - record_shapes=obj.record_shapes, - profile_by_stage=obj.profile_by_stage, - profile_id=str(time.time()), - merge_profiles=obj.merge_profiles, - profile_prefix=obj.profile_prefix, - profile_stages=obj.profile_stages, - ) + req = obj or ProfileReq() + req.req_type = ProfileReqType.START_PROFILE for socket in send_sockets: sock_send(socket, req) if encoder.profiler is None: @@ -3937,7 +3901,7 @@ async def stop_profile_async(): return Response( content="profiling not initialized\n", status_code=HTTPStatus.BAD_REQUEST ) - req = ProfileReq(ProfileReqType.STOP_PROFILE) + req = ProfileReq(req_type=ProfileReqType.STOP_PROFILE) for socket in send_sockets: sock_send(socket, req) ok, msg = encoder.profiler.stop() diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 93edb7a83..9998640d3 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -40,6 +40,7 @@ from typing import ( Optional, Tuple, Union, + cast, ) import torch @@ -68,6 +69,8 @@ from sglang.srt.managers.io_struct import ( LoadLoRAAdapterReqInput, MultimodalDataInputFormat, OpenSessionReqInput, + ProfileReq, + ProfileReqType, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, RpcReqInput, @@ -93,6 +96,7 @@ from sglang.srt.plugins import load_plugins from sglang.srt.server_args import PortArgs, ServerArgs from sglang.srt.utils import ( MultiprocessingSerializer, + SerializedTensorPayload, assert_pkg_version, configure_logger, get_bool_env_var, @@ -100,6 +104,7 @@ from sglang.srt.utils import ( kill_process_tree, launch_dummy_health_check_server, maybe_reindex_device_id, + normalize_serialized_named_tensor_payloads, numa_utils, set_prometheus_multiproc_dir, set_ulimit, @@ -964,7 +969,8 @@ class Engine(EngineScoreMixin, EngineBase): self.loop.run_until_complete(self.tokenizer_manager.close_session(obj, None)) def start_profile(self, **kwargs): - self.loop.run_until_complete(self.tokenizer_manager.start_profile(**kwargs)) + req = ProfileReq(req_type=ProfileReqType.START_PROFILE, **kwargs) + self.loop.run_until_complete(self.tokenizer_manager.start_profile(req)) def stop_profile(self): self.loop.run_until_complete(self.tokenizer_manager.stop_profile()) @@ -1053,14 +1059,19 @@ class Engine(EngineScoreMixin, EngineBase): def update_weights_from_tensor( self, - named_tensors: List[Tuple[str, torch.Tensor]], + named_tensors: Union[ + List[Tuple[str, torch.Tensor]], + List[SerializedTensorPayload], + ], load_format: Optional[str] = None, flush_cache: bool = True, ): """Update weights from distributed source. If there are going to be more updates, set `flush_cache` to be false to avoid duplicated cache cleaning operation.""" if load_format == "flattened_bucket": - serialized_named_tensors = named_tensors + serialized_named_tensors = normalize_serialized_named_tensor_payloads( + cast(List[SerializedTensorPayload], named_tensors) + ) else: serialized_named_tensors = [ MultiprocessingSerializer.serialize(named_tensors) diff --git a/python/sglang/srt/entrypoints/grpc_bridge.py b/python/sglang/srt/entrypoints/grpc_bridge.py index 5fa4125a2..5c56edc03 100644 --- a/python/sglang/srt/entrypoints/grpc_bridge.py +++ b/python/sglang/srt/entrypoints/grpc_bridge.py @@ -479,8 +479,10 @@ class RuntimeHandle: def start_profile(self, output_dir: Optional[str], chunk_callback) -> None: async def _payload(): - kwargs = {"output_dir": output_dir} if output_dir else {} - await self.tokenizer_manager.start_profile(**kwargs) + from sglang.srt.managers.io_struct import ProfileReq + + req = ProfileReq(output_dir=output_dir) if output_dir else ProfileReq() + await self.tokenizer_manager.start_profile(req) return {"message": "Profiling started."} self._submit_json_unary("start_profile", _payload, chunk_callback) diff --git a/python/sglang/srt/entrypoints/grpc_server.py b/python/sglang/srt/entrypoints/grpc_server.py index 00944c276..320e04a94 100644 --- a/python/sglang/srt/entrypoints/grpc_server.py +++ b/python/sglang/srt/entrypoints/grpc_server.py @@ -105,7 +105,7 @@ def _add_admin_routes(app, request_manager): 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 @@ def _add_admin_routes(app, request_manager): 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_server.py b/python/sglang/srt/entrypoints/http_server.py index 8f7722311..8b5d9e73d 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -124,7 +124,7 @@ from sglang.srt.managers.io_struct import ( OpenSessionReqInput, ParseFunctionCallReq, PauseGenerationReqInput, - ProfileReqInput, + ProfileReq, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, SendWeightsToRemoteInstanceReqInput, @@ -1024,23 +1024,9 @@ 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[ProfileReq] = None): """Start profiling.""" - if obj is None: - obj = ProfileReqInput() - - await _global_state.tokenizer_manager.start_profile( - output_dir=obj.output_dir, - start_step=obj.start_step, - num_steps=obj.num_steps, - activities=obj.activities, - with_stack=obj.with_stack, - record_shapes=obj.record_shapes, - profile_by_stage=obj.profile_by_stage, - merge_profiles=obj.merge_profiles, - profile_prefix=obj.profile_prefix, - profile_stages=obj.profile_stages, - ) + await _global_state.tokenizer_manager.start_profile(obj or ProfileReq()) return Response( content="Start profiling.\n", status_code=200, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 32bbf12f1..6f87cc667 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1769,8 +1769,14 @@ class SetInternalStateReqOutput(BaseReq): server_args: Dict[str, Any] +class ProfileReqType(Enum): + START_PROFILE = 1 + STOP_PROFILE = 2 + + @dataclass -class ProfileReqInput(BaseReq): +class ProfileReq(BaseReq): + req_type: ProfileReqType = ProfileReqType.START_PROFILE # The output directory output_dir: Optional[str] = None # Specify the steps to start the profiling @@ -1787,6 +1793,7 @@ class ProfileReqInput(BaseReq): with_stack: Optional[bool] = None # Whether to save information about operator’s input shapes. record_shapes: Optional[bool] = None + profile_id: Optional[str] = None # Merge profiles from all ranks into a single trace merge_profiles: bool = False # The prefix of the profile filenames @@ -1795,27 +1802,6 @@ class ProfileReqInput(BaseReq): profile_stages: Optional[List[str]] = None -class ProfileReqType(Enum): - START_PROFILE = 1 - STOP_PROFILE = 2 - - -@dataclass -class ProfileReq(BaseReq): - type: ProfileReqType - output_dir: Optional[str] = None - start_step: Optional[int] = None - num_steps: Optional[int] = None - activities: Optional[List[str]] = None - profile_by_stage: bool = False - with_stack: Optional[bool] = None - record_shapes: Optional[bool] = None - profile_id: Optional[str] = None - merge_profiles: bool = False - profile_prefix: Optional[str] = None - profile_stages: Optional[List[str]] = None - - @dataclass class ProfileReqOutput(BaseReq): success: bool diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f928042d3..92b5a35e8 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3810,7 +3810,7 @@ class Scheduler( 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): if (chunked_req := self.chunked_req) is not None: @@ -4032,14 +4032,16 @@ class Scheduler( 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 @@ -4065,7 +4067,9 @@ class Scheduler( # Radix-native: open is implicit; explicit open only permits id reuse. session_id = recv_req.session_id self.tree_cache.register_session(session_id) - output = OpenSessionReqOutput(session_id, session_id is not None) + output = OpenSessionReqOutput( + session_id=session_id, success=session_id is not None + ) else: output = self.session_controller.open(recv_req) if self.ps.pp_rank == 0 and self.ps.tp_rank == 0 and self.ps.attn_cp_rank == 0: diff --git a/python/sglang/srt/managers/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py index 42b47b598..90945b82d 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 @@ class SchedulerProfilerManager: 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/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index c2dc1c1d8..1abfeef97 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 @@ class SchedulerWeightUpdaterManager: 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 @@ class SchedulerWeightUpdaterManager: ): """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 @@ class SchedulerWeightUpdaterManager: 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 @@ class SchedulerWeightUpdaterManager: 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 @@ class SchedulerWeightUpdaterManager: 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/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index dac49eac2..5bbfe8dbf 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -76,7 +76,10 @@ from sglang.srt.managers.io_struct import ( ) from sglang.srt.managers.load_snapshot import LoadSnapshot from sglang.srt.server_args import LoRARef, ServerArgs -from sglang.srt.utils import get_bool_env_var +from sglang.srt.utils import ( + get_bool_env_var, + normalize_serialized_named_tensor_payloads, +) from sglang.utils import TypeBasedDispatcher if TYPE_CHECKING: @@ -319,43 +322,25 @@ class TokenizerControlMixin: async def start_profile( self: TokenizerManager, - output_dir: Optional[str] = None, - start_step: Optional[int] = None, - num_steps: Optional[int] = None, - activities: Optional[List[str]] = None, - with_stack: Optional[bool] = None, - record_shapes: Optional[bool] = None, - profile_by_stage: bool = False, - merge_profiles: bool = False, - profile_prefix: Optional[str] = None, - profile_stages: Optional[List[str]] = None, + req: Optional[ProfileReq] = None, ): self.auto_create_handle_loop() + req = req or ProfileReq() + req.req_type = ProfileReqType.START_PROFILE env_with_stack: bool = get_bool_env_var("SGLANG_PROFILE_WITH_STACK", "true") - with_stack = False if with_stack is False or env_with_stack is False else True + req.with_stack = ( + False if req.with_stack is False or env_with_stack is False else True + ) env_record_shapes: bool = get_bool_env_var( "SGLANG_PROFILE_RECORD_SHAPES", "true" ) - record_shapes = (record_shapes is not False) and env_record_shapes - req = ProfileReq( - type=ProfileReqType.START_PROFILE, - output_dir=output_dir, - start_step=start_step, - num_steps=num_steps, - activities=activities, - with_stack=with_stack, - record_shapes=record_shapes, - profile_by_stage=profile_by_stage, - profile_id=str(time.time()), - merge_profiles=merge_profiles, - profile_prefix=profile_prefix, - profile_stages=profile_stages, - ) + req.record_shapes = (req.record_shapes is not False) and env_record_shapes + req.profile_id = req.profile_id or str(time.time()) return await self._execute_profile(req) 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): @@ -476,6 +461,10 @@ class TokenizerControlMixin: if obj.abort_all_requests: self.abort_request(abort_all=True) + obj.serialized_named_tensors = normalize_serialized_named_tensor_payloads( + obj.serialized_named_tensors + ) + async with self.is_pause_cond: is_paused = self.is_pause if is_paused: diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 6362e1dd7..532aa7790 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1141,15 +1141,15 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): 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, @@ -1190,11 +1190,11 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): ) 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, diff --git a/python/sglang/srt/session/session_controller.py b/python/sglang/srt/session/session_controller.py index 874de4fdb..f0dd9c7f2 100644 --- a/python/sglang/srt/session/session_controller.py +++ b/python/sglang/srt/session/session_controller.py @@ -365,10 +365,10 @@ class SessionController: 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 @@ class SessionController: 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/srt/utils/common.py b/python/sglang/srt/utils/common.py index ea9e9b252..427ca42b8 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -17,6 +17,7 @@ from __future__ import annotations import argparse import asyncio +import binascii import builtins import ctypes import functools @@ -2388,6 +2389,39 @@ class MultiprocessingSerializer: return SafeUnpickler(io.BytesIO(data)).load() +SerializedTensorPayload = Union[str, bytes, bytearray, memoryview] + + +def _looks_like_pickle_payload(data: bytes) -> bool: + return len(data) >= 2 and data[0] == 0x80 and data[1] <= pickle.HIGHEST_PROTOCOL + + +def normalize_serialized_named_tensor_payload(data: SerializedTensorPayload) -> bytes: + """Normalize a serialized tensor payload to raw MultiprocessingSerializer bytes.""" + if isinstance(data, str): + return pybase64.b64decode(data, validate=True) + + if isinstance(data, (bytes, bytearray, memoryview)): + data = bytes(data) + if _looks_like_pickle_payload(data): + return data + try: + return pybase64.b64decode(data, validate=True) + except (binascii.Error, ValueError): + return data + + raise TypeError( + "serialized_named_tensors entries must be base64 strings or bytes-like " + f"payloads, got {type(data).__name__}" + ) + + +def normalize_serialized_named_tensor_payloads( + payloads: List[SerializedTensorPayload], +) -> List[bytes]: + return [normalize_serialized_named_tensor_payload(data) for data in payloads] + + class SafeUnpickler(pickle.Unpickler): ALLOWED_MODULE_PREFIXES = { # --- Python types --- 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 1c9dc09e6..ab5eae098 100644 --- a/test/registered/unit/managers/test_profile_merger_http_api.py +++ b/test/registered/unit/managers/test_profile_merger_http_api.py @@ -1,22 +1,23 @@ import json import unittest -from sglang.srt.managers.io_struct import ProfileReqInput +from sglang.srt.managers.io_struct import ProfileReq from sglang.test.ci.ci_register import ( register_amd_ci, register_cpu_ci, register_cuda_ci, ) +from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=8, stage="base-b", runner_config="1-gpu-small") register_amd_ci(est_time=9, suite="stage-b-test-1-gpu-small-amd") register_cpu_ci(est_time=8, suite="base-c-test-cpu") -class TestProfileMergerHTTPAPI(unittest.TestCase): - def test_profile_req_input_merge_profiles_json_serialization(self): +class TestProfileMergerHTTPAPI(CustomTestCase): + def test_profile_req_merge_profiles_json_serialization(self): # Test with merge_profiles=True - req_input = ProfileReqInput( + req = ProfileReq( output_dir="/tmp/test", num_steps=5, activities=["CPU", "GPU"], @@ -26,11 +27,11 @@ class TestProfileMergerHTTPAPI(unittest.TestCase): # Convert to dict (as would happen in HTTP request) req_dict = { - "output_dir": req_input.output_dir, - "num_steps": req_input.num_steps, - "activities": req_input.activities, - "profile_by_stage": req_input.profile_by_stage, - "merge_profiles": req_input.merge_profiles, + "output_dir": req.output_dir, + "num_steps": req.num_steps, + "activities": req.activities, + "profile_by_stage": req.profile_by_stage, + "merge_profiles": req.merge_profiles, } # Test JSON serialization @@ -43,7 +44,7 @@ class TestProfileMergerHTTPAPI(unittest.TestCase): self.assertEqual(parsed_data["activities"], ["CPU", "GPU"]) self.assertTrue(parsed_data["profile_by_stage"]) - def test_profile_req_input_merge_profiles_json_deserialization(self): + def test_profile_req_merge_profiles_json_deserialization(self): # Test JSON data as would come from HTTP request json_data = { "output_dir": "/tmp/test", @@ -53,27 +54,27 @@ class TestProfileMergerHTTPAPI(unittest.TestCase): "merge_profiles": True, } - # Create ProfileReqInput from dict (as HTTP server would do) - req_input = ProfileReqInput(**json_data) + # Create ProfileReq from dict (as HTTP server would do) + req = ProfileReq(**json_data) - self.assertTrue(req_input.merge_profiles) - self.assertEqual(req_input.output_dir, "/tmp/test") - self.assertEqual(req_input.num_steps, 10) - self.assertEqual(req_input.activities, ["CPU", "GPU", "MEM"]) - self.assertFalse(req_input.profile_by_stage) + self.assertTrue(req.merge_profiles) + self.assertEqual(req.output_dir, "/tmp/test") + self.assertEqual(req.num_steps, 10) + self.assertEqual(req.activities, ["CPU", "GPU", "MEM"]) + self.assertFalse(req.profile_by_stage) - def test_profile_req_input_merge_profiles_default_value(self): + def test_profile_req_merge_profiles_default_value(self): # Test with minimal data json_data = {"output_dir": "/tmp/test"} - req_input = ProfileReqInput(**json_data) - self.assertFalse(req_input.merge_profiles) + req = ProfileReq(**json_data) + self.assertFalse(req.merge_profiles) - def test_profile_req_input_merge_profiles_explicit_false(self): + def test_profile_req_merge_profiles_explicit_false(self): json_data = {"output_dir": "/tmp/test", "merge_profiles": False} - req_input = ProfileReqInput(**json_data) - self.assertFalse(req_input.merge_profiles) + req = ProfileReq(**json_data) + self.assertFalse(req.merge_profiles) def test_http_api_parameter_flow(self): # Simulate HTTP request data @@ -85,8 +86,8 @@ class TestProfileMergerHTTPAPI(unittest.TestCase): "merge_profiles": True, } - # Create ProfileReqInput as HTTP server would - obj = ProfileReqInput(**request_data) + # Create ProfileReq as HTTP server would + obj = ProfileReq(**request_data) # Verify the parameter is set correctly self.assertTrue(obj.merge_profiles) @@ -98,24 +99,24 @@ class TestProfileMergerHTTPAPI(unittest.TestCase): def test_http_api_parameter_validation(self): # Test with True json_data = {"merge_profiles": True} - req_input = ProfileReqInput(**json_data) - self.assertTrue(req_input.merge_profiles) + req = ProfileReq(**json_data) + self.assertTrue(req.merge_profiles) # Test with False json_data = {"merge_profiles": False} - req_input = ProfileReqInput(**json_data) - self.assertFalse(req_input.merge_profiles) + req = ProfileReq(**json_data) + self.assertFalse(req.merge_profiles) # Test with string "true" (should be converted by JSON parser) json_data = {"merge_profiles": "true"} - req_input = ProfileReqInput(**json_data) - self.assertEqual(req_input.merge_profiles, "true") # String, not boolean + req = ProfileReq(**json_data) + self.assertEqual(req.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) - self.assertFalse(req_input.merge_profiles) # Should default to False + req = ProfileReq(**json_data) + self.assertFalse(req.merge_profiles) # Should default to False # Test with other parameters but no merge_profiles json_data = { @@ -123,8 +124,8 @@ class TestProfileMergerHTTPAPI(unittest.TestCase): "num_steps": 5, "activities": ["CPU", "GPU"], } - req_input = ProfileReqInput(**json_data) - self.assertFalse(req_input.merge_profiles) # Should default to False + req = ProfileReq(**json_data) + self.assertFalse(req.merge_profiles) # Should default to False def test_http_api_parameter_combinations(self): test_cases = [ @@ -163,8 +164,8 @@ class TestProfileMergerHTTPAPI(unittest.TestCase): for test_case in test_cases: with self.subTest(test_case["name"]): - req_input = ProfileReqInput(**test_case["data"]) - self.assertEqual(req_input.merge_profiles, test_case["expected_merge"]) + req = ProfileReq(**test_case["data"]) + self.assertEqual(req.merge_profiles, test_case["expected_merge"]) if __name__ == "__main__": diff --git a/test/registered/unit/observability/test_trace.py b/test/registered/unit/observability/test_trace.py index 7667bce92..547beab65 100644 --- a/test/registered/unit/observability/test_trace.py +++ b/test/registered/unit/observability/test_trace.py @@ -227,7 +227,7 @@ class TestTraceReqContextDisabled(unittest.TestCase): 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 @@ class TestTraceReqContextEnabled(unittest.TestCase): 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 @@ class TestTraceReqContextEnabled(unittest.TestCase): 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/utils/test_profile_merger.py b/test/registered/unit/utils/test_profile_merger.py index fd4d9b1b0..75631e71b 100644 --- a/test/registered/unit/utils/test_profile_merger.py +++ b/test/registered/unit/utils/test_profile_merger.py @@ -13,20 +13,21 @@ import shutil import tempfile import unittest -from sglang.srt.managers.io_struct import ProfileReq, ProfileReqInput, ProfileReqType +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, register_cpu_ci, register_cuda_ci, ) +from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small") register_amd_ci(est_time=8, suite="stage-b-test-1-gpu-small-amd") register_cpu_ci(est_time=8, suite="base-c-test-cpu") -class TestProfileMerger(unittest.TestCase): +class TestProfileMerger(CustomTestCase): def setUp(self): self.temp_dir = tempfile.mkdtemp() self.profile_id = "test_profile_123" @@ -205,21 +206,15 @@ class TestProfileMerger(unittest.TestCase): empty_merger.merge_chrome_traces() -class TestProfileMergerIntegration(unittest.TestCase): +class TestProfileMergerIntegration(CustomTestCase): def test_data_structures_merge_profiles(self): - # Test ProfileReqInput - req_input = ProfileReqInput() - self.assertFalse(req_input.merge_profiles) - - req_input = ProfileReqInput(merge_profiles=True) - self.assertTrue(req_input.merge_profiles) - # Test ProfileReq - req = ProfileReq(type=ProfileReqType.START_PROFILE) + req = ProfileReq() self.assertFalse(req.merge_profiles) + self.assertEqual(req.req_type, ProfileReqType.START_PROFILE) - req = ProfileReq(type=ProfileReqType.START_PROFILE, merge_profiles=True) + req = ProfileReq(merge_profiles=True) self.assertTrue(req.merge_profiles) def test_integration_parameters(self): @@ -231,7 +226,8 @@ class TestProfileMergerIntegration(unittest.TestCase): ) sig = inspect.signature(TokenizerControlMixin.start_profile) - self.assertIn("merge_profiles", sig.parameters) + self.assertIn("req", sig.parameters) + self.assertNotIn("merge_profiles", sig.parameters) # Test SchedulerProfilerMixin from sglang.srt.managers.scheduler_components.profiler_manager import ( @@ -248,7 +244,7 @@ class TestProfileMergerIntegration(unittest.TestCase): self.assertIn("merge_profiles", sig.parameters) -class TestProfileMergerEdgeCases(unittest.TestCase): +class TestProfileMergerEdgeCases(CustomTestCase): def setUp(self): self.temp_dir = tempfile.mkdtemp() self.profile_id = "test_edge_cases"