minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Extract profile request cleanups (#29098)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-06-24 11:22:58 -07:00
committed by GitHub
parent d6aacd2801
commit d5c566e59b
16 changed files with 179 additions and 202 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/encode_server.py
+9 -45
View File
@@ -46,7 +46,6 @@ from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import initialize_dp_attention from sglang.srt.layers.dp_attention import initialize_dp_attention
from sglang.srt.managers.io_struct import ( from sglang.srt.managers.io_struct import (
ProfileReq, ProfileReq,
ProfileReqInput,
ProfileReqType, ProfileReqType,
async_sock_recv, async_sock_recv,
async_sock_send, async_sock_send,
@@ -2981,13 +2980,8 @@ async def _dp_worker_handle_profile(
) -> dict: ) -> dict:
prefix = f"dp_rank={dp_rank}: " prefix = f"dp_rank={dp_rank}: "
if dp_type == "start_profile": if dp_type == "start_profile":
obj = request.get("profile_req") req = request.get("profile_req") or ProfileReq()
# `is None` (not `if not obj`) so empty dict still raises. req.req_type = ProfileReqType.START_PROFILE
req = (
ProfileReq(**obj)
if obj is not None
else ProfileReq(ProfileReqType.START_PROFILE)
)
if enc.profiler is None: if enc.profiler is None:
enc.profiler = EncoderProfiler(dp_rank) enc.profiler = EncoderProfiler(dp_rank)
ok, msg = enc.profiler.start(req) ok, msg = enc.profiler.start(req)
@@ -3207,7 +3201,7 @@ async def run_encoder(
while True: while True:
request = await async_sock_recv(encoder.schedule_socket) request = await async_sock_recv(encoder.schedule_socket)
if isinstance(request, ProfileReq): if isinstance(request, ProfileReq):
if request.type == ProfileReqType.START_PROFILE: if request.req_type == ProfileReqType.START_PROFILE:
if encoder.profiler is None: if encoder.profiler is None:
encoder.profiler = EncoderProfiler(encoder.rank) encoder.profiler = EncoderProfiler(encoder.rank)
encoder.profiler.start(request) encoder.profiler.start(request)
@@ -3862,51 +3856,21 @@ async def health_generate():
@app.api_route("/start_profile", methods=["GET", "POST"]) @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: if dp_dispatcher is not None:
profile_req = None
if obj is not None: if obj is not None:
profile_req = { obj.req_type = ProfileReqType.START_PROFILE
"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,
}
try: try:
results = await dp_dispatcher.broadcast( results = await dp_dispatcher.broadcast(
{"_dp_type": "start_profile", "profile_req": profile_req} {"_dp_type": "start_profile", "profile_req": obj}
) )
except MMError as e: except MMError as e:
return Response(content=f"{e}\n", status_code=int(e.code)) return Response(content=f"{e}\n", status_code=int(e.code))
return _summarise_dp_broadcast(results) return _summarise_dp_broadcast(results)
if encoder is None: if encoder is None:
return Response(content="encoder not ready\n", status_code=503) return Response(content="encoder not ready\n", status_code=503)
req = None req = obj or ProfileReq()
if obj is None: req.req_type = ProfileReqType.START_PROFILE
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,
)
for socket in send_sockets: for socket in send_sockets:
sock_send(socket, req) sock_send(socket, req)
if encoder.profiler is None: if encoder.profiler is None:
@@ -3937,7 +3901,7 @@ async def stop_profile_async():
return Response( return Response(
content="profiling not initialized\n", status_code=HTTPStatus.BAD_REQUEST 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: for socket in send_sockets:
sock_send(socket, req) sock_send(socket, req)
ok, msg = encoder.profiler.stop() ok, msg = encoder.profiler.stop()
python/sglang/srt/entrypoints/engine.py
+14 -3
View File
@@ -40,6 +40,7 @@ from typing import (
Optional, Optional,
Tuple, Tuple,
Union, Union,
cast,
) )
import torch import torch
@@ -68,6 +69,8 @@ from sglang.srt.managers.io_struct import (
LoadLoRAAdapterReqInput, LoadLoRAAdapterReqInput,
MultimodalDataInputFormat, MultimodalDataInputFormat,
OpenSessionReqInput, OpenSessionReqInput,
ProfileReq,
ProfileReqType,
ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqInput,
ResumeMemoryOccupationReqInput, ResumeMemoryOccupationReqInput,
RpcReqInput, RpcReqInput,
@@ -93,6 +96,7 @@ from sglang.srt.plugins import load_plugins
from sglang.srt.server_args import PortArgs, ServerArgs from sglang.srt.server_args import PortArgs, ServerArgs
from sglang.srt.utils import ( from sglang.srt.utils import (
MultiprocessingSerializer, MultiprocessingSerializer,
SerializedTensorPayload,
assert_pkg_version, assert_pkg_version,
configure_logger, configure_logger,
get_bool_env_var, get_bool_env_var,
@@ -100,6 +104,7 @@ from sglang.srt.utils import (
kill_process_tree, kill_process_tree,
launch_dummy_health_check_server, launch_dummy_health_check_server,
maybe_reindex_device_id, maybe_reindex_device_id,
normalize_serialized_named_tensor_payloads,
numa_utils, numa_utils,
set_prometheus_multiproc_dir, set_prometheus_multiproc_dir,
set_ulimit, set_ulimit,
@@ -964,7 +969,8 @@ class Engine(EngineScoreMixin, EngineBase):
self.loop.run_until_complete(self.tokenizer_manager.close_session(obj, None)) self.loop.run_until_complete(self.tokenizer_manager.close_session(obj, None))
def start_profile(self, **kwargs): 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): def stop_profile(self):
self.loop.run_until_complete(self.tokenizer_manager.stop_profile()) self.loop.run_until_complete(self.tokenizer_manager.stop_profile())
@@ -1053,14 +1059,19 @@ class Engine(EngineScoreMixin, EngineBase):
def update_weights_from_tensor( def update_weights_from_tensor(
self, self,
named_tensors: List[Tuple[str, torch.Tensor]], named_tensors: Union[
List[Tuple[str, torch.Tensor]],
List[SerializedTensorPayload],
],
load_format: Optional[str] = None, load_format: Optional[str] = None,
flush_cache: bool = True, flush_cache: bool = True,
): ):
"""Update weights from distributed source. If there are going to be more updates, set `flush_cache` to be false """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.""" to avoid duplicated cache cleaning operation."""
if load_format == "flattened_bucket": if load_format == "flattened_bucket":
serialized_named_tensors = named_tensors serialized_named_tensors = normalize_serialized_named_tensor_payloads(
cast(List[SerializedTensorPayload], named_tensors)
)
else: else:
serialized_named_tensors = [ serialized_named_tensors = [
MultiprocessingSerializer.serialize(named_tensors) MultiprocessingSerializer.serialize(named_tensors)
python/sglang/srt/entrypoints/grpc_bridge.py
+4 -2
View File
@@ -479,8 +479,10 @@ class RuntimeHandle:
def start_profile(self, output_dir: Optional[str], chunk_callback) -> None: def start_profile(self, output_dir: Optional[str], chunk_callback) -> None:
async def _payload(): async def _payload():
kwargs = {"output_dir": output_dir} if output_dir else {} from sglang.srt.managers.io_struct import ProfileReq
await self.tokenizer_manager.start_profile(**kwargs)
req = ProfileReq(output_dir=output_dir) if output_dir else ProfileReq()
await self.tokenizer_manager.start_profile(req)
return {"message": "Profiling started."} return {"message": "Profiling started."}
self._submit_json_unary("start_profile", _payload, chunk_callback) self._submit_json_unary("start_profile", _payload, chunk_callback)
python/sglang/srt/entrypoints/grpc_server.py
+2 -2
View File
@@ -105,7 +105,7 @@ def _add_admin_routes(app, request_manager):
record_shapes = (record_shapes is not False) and env_record_shapes record_shapes = (record_shapes is not False) and env_record_shapes
req = ProfileReq( req = ProfileReq(
type=ProfileReqType.START_PROFILE, req_type=ProfileReqType.START_PROFILE,
output_dir=body.get("output_dir"), output_dir=body.get("output_dir"),
start_step=body.get("start_step"), start_step=body.get("start_step"),
num_steps=body.get("num_steps"), num_steps=body.get("num_steps"),
@@ -134,7 +134,7 @@ def _add_admin_routes(app, request_manager):
async def stop_profile_handler(request): async def stop_profile_handler(request):
try: try:
req = ProfileReq(type=ProfileReqType.STOP_PROFILE) req = ProfileReq(req_type=ProfileReqType.STOP_PROFILE)
results = await request_manager.send_communicator_req( results = await request_manager.send_communicator_req(
req, "profile_communicator", timeout=600.0 req, "profile_communicator", timeout=600.0
) )
python/sglang/srt/entrypoints/http_server.py
+3 -17
View File
@@ -124,7 +124,7 @@ from sglang.srt.managers.io_struct import (
OpenSessionReqInput, OpenSessionReqInput,
ParseFunctionCallReq, ParseFunctionCallReq,
PauseGenerationReqInput, PauseGenerationReqInput,
ProfileReqInput, ProfileReq,
ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqInput,
ResumeMemoryOccupationReqInput, ResumeMemoryOccupationReqInput,
SendWeightsToRemoteInstanceReqInput, SendWeightsToRemoteInstanceReqInput,
@@ -1024,23 +1024,9 @@ async def hicache_storage_backend_status():
@app.api_route("/start_profile", methods=["GET", "POST"]) @app.api_route("/start_profile", methods=["GET", "POST"])
@auth_level(AuthLevel.ADMIN_OPTIONAL) @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.""" """Start profiling."""
if obj is None: await _global_state.tokenizer_manager.start_profile(obj or ProfileReq())
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,
)
return Response( return Response(
content="Start profiling.\n", content="Start profiling.\n",
status_code=200, status_code=200,
python/sglang/srt/managers/io_struct.py
+8 -22
View File
@@ -1769,8 +1769,14 @@ class SetInternalStateReqOutput(BaseReq):
server_args: Dict[str, Any] server_args: Dict[str, Any]
class ProfileReqType(Enum):
START_PROFILE = 1
STOP_PROFILE = 2
@dataclass @dataclass
class ProfileReqInput(BaseReq): class ProfileReq(BaseReq):
req_type: ProfileReqType = ProfileReqType.START_PROFILE
# The output directory # The output directory
output_dir: Optional[str] = None output_dir: Optional[str] = None
# Specify the steps to start the profiling # Specify the steps to start the profiling
@@ -1787,6 +1793,7 @@ class ProfileReqInput(BaseReq):
with_stack: Optional[bool] = None with_stack: Optional[bool] = None
# Whether to save information about operator’s input shapes. # Whether to save information about operator’s input shapes.
record_shapes: Optional[bool] = None record_shapes: Optional[bool] = None
profile_id: Optional[str] = None
# Merge profiles from all ranks into a single trace # Merge profiles from all ranks into a single trace
merge_profiles: bool = False merge_profiles: bool = False
# The prefix of the profile filenames # The prefix of the profile filenames
@@ -1795,27 +1802,6 @@ class ProfileReqInput(BaseReq):
profile_stages: Optional[List[str]] = None 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 @dataclass
class ProfileReqOutput(BaseReq): class ProfileReqOutput(BaseReq):
success: bool success: bool
python/sglang/srt/managers/scheduler.py
+8 -4
View File
@@ -3810,7 +3810,7 @@ class Scheduler(
logger.error(f"Failed to call rpc {recv_req.method}: {str(e)}") logger.error(f"Failed to call rpc {recv_req.method}: {str(e)}")
barrier() 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): def abort_request(self, recv_req: AbortReq):
if (chunked_req := self.chunked_req) is not None: 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( success, message = self.tp_worker.init_weights_send_group_for_remote_instance(
recv_req recv_req
) )
return InitWeightsSendGroupForRemoteInstanceReqOutput(success, message) return InitWeightsSendGroupForRemoteInstanceReqOutput(
success=success, message=message
)
def send_weights_to_remote_instance( def send_weights_to_remote_instance(
self, recv_req: SendWeightsToRemoteInstanceReqInput self, recv_req: SendWeightsToRemoteInstanceReqInput
): ):
"""Send the seed instance weights to the destination instance.""" """Send the seed instance weights to the destination instance."""
success, message = self.tp_worker.send_weights_to_remote_instance(recv_req) 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): def slow_down(self, recv_req: SlowDownReqInput):
t = recv_req.forward_sleep_time t = recv_req.forward_sleep_time
@@ -4065,7 +4067,9 @@ class Scheduler(
# Radix-native: open is implicit; explicit open only permits id reuse. # Radix-native: open is implicit; explicit open only permits id reuse.
session_id = recv_req.session_id session_id = recv_req.session_id
self.tree_cache.register_session(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: else:
output = self.session_controller.open(recv_req) 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: if self.ps.pp_rank == 0 and self.ps.tp_rank == 0 and self.ps.attn_cp_rank == 0:
python/sglang/srt/managers/scheduler_components/profiler_manager.py
+1 -1
View File
@@ -412,7 +412,7 @@ class SchedulerProfilerManager:
self._start_profile() self._start_profile()
def _profile(self, recv_req: ProfileReq): 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: if recv_req.profile_by_stage or recv_req.start_step:
return self._init_profile( return self._init_profile(
recv_req.output_dir, recv_req.output_dir,
python/sglang/srt/managers/scheduler_components/weight_updater.py
+11 -7
View File
@@ -116,12 +116,14 @@ class SchedulerWeightUpdaterManager:
self.flush_cache_after_weight_update(recv_req) self.flush_cache_after_weight_update(recv_req)
if not success: if not success:
logger.error(message) 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): def init_weights_update_group(self, recv_req: InitWeightsUpdateGroupReqInput):
"""Initialize the online model parameter update group.""" """Initialize the online model parameter update group."""
success, message = self.tp_worker.init_weights_update_group(recv_req) 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( def destroy_weights_update_group(
self, self,
@@ -129,7 +131,7 @@ class SchedulerWeightUpdaterManager:
): ):
"""Destroy the online model parameter update group.""" """Destroy the online model parameter update group."""
success, message = self.tp_worker.destroy_weights_update_group(recv_req) 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( def update_weights_from_distributed(
self, self,
@@ -142,7 +144,9 @@ class SchedulerWeightUpdaterManager:
self.flush_cache_after_weight_update(recv_req) self.flush_cache_after_weight_update(recv_req)
else: else:
logger.error(message) logger.error(message)
return UpdateWeightsFromDistributedReqOutput(success, message) return UpdateWeightsFromDistributedReqOutput(
success=success, message=message
)
def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput): def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
"""Update the online model parameter from tensors.""" """Update the online model parameter from tensors."""
@@ -157,7 +161,7 @@ class SchedulerWeightUpdaterManager:
else: else:
logger.error(message) logger.error(message)
torch.distributed.barrier(group=self.tp_cpu_group) 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): def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput):
"""Update the online model parameter from IPC for checkpoint-engine integration.""" """Update the online model parameter from IPC for checkpoint-engine integration."""
@@ -171,11 +175,11 @@ class SchedulerWeightUpdaterManager:
if not success: if not success:
logger.error(message) logger.error(message)
torch.distributed.barrier(group=self.tp_cpu_group) 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): def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput):
parameter = self.tp_worker.get_weights_by_name(recv_req) 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): def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput):
assert ( assert (
python/sglang/srt/managers/tokenizer_control_mixin.py
+17 -28
View File
@@ -76,7 +76,10 @@ from sglang.srt.managers.io_struct import (
) )
from sglang.srt.managers.load_snapshot import LoadSnapshot from sglang.srt.managers.load_snapshot import LoadSnapshot
from sglang.srt.server_args import LoRARef, ServerArgs 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 from sglang.utils import TypeBasedDispatcher
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -319,43 +322,25 @@ class TokenizerControlMixin:
async def start_profile( async def start_profile(
self: TokenizerManager, self: TokenizerManager,
output_dir: Optional[str] = None, req: Optional[ProfileReq] = 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,
): ):
self.auto_create_handle_loop() 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") 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( env_record_shapes: bool = get_bool_env_var(
"SGLANG_PROFILE_RECORD_SHAPES", "true" "SGLANG_PROFILE_RECORD_SHAPES", "true"
) )
record_shapes = (record_shapes is not False) and env_record_shapes req.record_shapes = (req.record_shapes is not False) and env_record_shapes
req = ProfileReq( req.profile_id = req.profile_id or str(time.time())
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,
)
return await self._execute_profile(req) return await self._execute_profile(req)
async def stop_profile(self: TokenizerManager): async def stop_profile(self: TokenizerManager):
self.auto_create_handle_loop() self.auto_create_handle_loop()
req = ProfileReq(type=ProfileReqType.STOP_PROFILE) req = ProfileReq(req_type=ProfileReqType.STOP_PROFILE)
return await self._execute_profile(req) return await self._execute_profile(req)
async def _execute_profile(self: TokenizerManager, req: ProfileReq): async def _execute_profile(self: TokenizerManager, req: ProfileReq):
@@ -476,6 +461,10 @@ class TokenizerControlMixin:
if obj.abort_all_requests: if obj.abort_all_requests:
self.abort_request(abort_all=True) 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: async with self.is_pause_cond:
is_paused = self.is_pause is_paused = self.is_pause
if is_paused: if is_paused:
python/sglang/srt/managers/tokenizer_manager.py
+14 -14
View File
@@ -1141,15 +1141,15 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
self.fake_bootstrap_room_counter += 1 self.fake_bootstrap_room_counter += 1
tokenized_obj = TokenizedGenerateReqInput( tokenized_obj = TokenizedGenerateReqInput(
input_text, input_text=input_text,
input_ids_arr, input_ids=input_ids_arr,
mm_inputs, mm_inputs=mm_inputs,
sampling_params, sampling_params=sampling_params,
obj.return_logprob, return_logprob=obj.return_logprob,
obj.logprob_start_len, logprob_start_len=obj.logprob_start_len,
obj.top_logprobs_num, top_logprobs_num=obj.top_logprobs_num,
obj.token_ids_logprob, token_ids_logprob=obj.token_ids_logprob,
obj.stream, stream=obj.stream,
rid=obj.rid, rid=obj.rid,
http_worker_ipc=obj.http_worker_ipc, http_worker_ipc=obj.http_worker_ipc,
bootstrap_host=obj.bootstrap_host, bootstrap_host=obj.bootstrap_host,
@@ -1190,11 +1190,11 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
) )
tokenized_obj = TokenizedEmbeddingReqInput( tokenized_obj = TokenizedEmbeddingReqInput(
input_text, input_text=input_text,
input_ids_arr, input_ids=input_ids_arr,
mm_inputs, image_inputs=mm_inputs,
token_type_ids, token_type_ids=token_type_ids,
sampling_params, sampling_params=sampling_params,
positional_embed_overrides=positional_embed_overrides, positional_embed_overrides=positional_embed_overrides,
rid=obj.rid, rid=obj.rid,
priority=obj.priority, priority=obj.priority,
python/sglang/srt/session/session_controller.py
+3 -3
View File
@@ -365,10 +365,10 @@ class SessionController:
session_id = recv_req.session_id session_id = recv_req.session_id
if session_id in self.sessions: if session_id in self.sessions:
logger.warning(f"session id {session_id} already exist, cannot open.") 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: elif session_id is None:
logger.warning("session id is None, cannot open.") logger.warning("session id is None, cannot open.")
return OpenSessionReqOutput(session_id, False) return OpenSessionReqOutput(session_id=session_id, success=False)
else: else:
self.sessions[session_id] = Session( self.sessions[session_id] = Session(
recv_req.capacity_of_str_len, recv_req.capacity_of_str_len,
@@ -379,7 +379,7 @@ class SessionController:
log_info_on_rank0( log_info_on_rank0(
logger, f"Session opened: {session_id} (active={len(self.sessions)})" 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): def close(self, recv_req: CloseSessionReqInput):
session_id = recv_req.session_id session_id = recv_req.session_id
python/sglang/srt/utils/common.py
+34
View File
@@ -17,6 +17,7 @@ from __future__ import annotations
import argparse import argparse
import asyncio import asyncio
import binascii
import builtins import builtins
import ctypes import ctypes
import functools import functools
@@ -2388,6 +2389,39 @@ class MultiprocessingSerializer:
return SafeUnpickler(io.BytesIO(data)).load() 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): class SafeUnpickler(pickle.Unpickler):
ALLOWED_MODULE_PREFIXES = { ALLOWED_MODULE_PREFIXES = {
# --- Python types --- # --- Python types ---
test/registered/unit/managers/test_profile_merger_http_api.py
+38 -37
View File
@@ -1,22 +1,23 @@
import json import json
import unittest 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 ( from sglang.test.ci.ci_register import (
register_amd_ci, register_amd_ci,
register_cpu_ci, register_cpu_ci,
register_cuda_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_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_amd_ci(est_time=9, suite="stage-b-test-1-gpu-small-amd")
register_cpu_ci(est_time=8, suite="base-c-test-cpu") register_cpu_ci(est_time=8, suite="base-c-test-cpu")
class TestProfileMergerHTTPAPI(unittest.TestCase): class TestProfileMergerHTTPAPI(CustomTestCase):
def test_profile_req_input_merge_profiles_json_serialization(self): def test_profile_req_merge_profiles_json_serialization(self):
# Test with merge_profiles=True # Test with merge_profiles=True
req_input = ProfileReqInput( req = ProfileReq(
output_dir="/tmp/test", output_dir="/tmp/test",
num_steps=5, num_steps=5,
activities=["CPU", "GPU"], activities=["CPU", "GPU"],
@@ -26,11 +27,11 @@ class TestProfileMergerHTTPAPI(unittest.TestCase):
# Convert to dict (as would happen in HTTP request) # Convert to dict (as would happen in HTTP request)
req_dict = { req_dict = {
"output_dir": req_input.output_dir, "output_dir": req.output_dir,
"num_steps": req_input.num_steps, "num_steps": req.num_steps,
"activities": req_input.activities, "activities": req.activities,
"profile_by_stage": req_input.profile_by_stage, "profile_by_stage": req.profile_by_stage,
"merge_profiles": req_input.merge_profiles, "merge_profiles": req.merge_profiles,
} }
# Test JSON serialization # Test JSON serialization
@@ -43,7 +44,7 @@ class TestProfileMergerHTTPAPI(unittest.TestCase):
self.assertEqual(parsed_data["activities"], ["CPU", "GPU"]) self.assertEqual(parsed_data["activities"], ["CPU", "GPU"])
self.assertTrue(parsed_data["profile_by_stage"]) 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 # Test JSON data as would come from HTTP request
json_data = { json_data = {
"output_dir": "/tmp/test", "output_dir": "/tmp/test",
@@ -53,27 +54,27 @@ class TestProfileMergerHTTPAPI(unittest.TestCase):
"merge_profiles": True, "merge_profiles": True,
} }
# Create ProfileReqInput from dict (as HTTP server would do) # Create ProfileReq from dict (as HTTP server would do)
req_input = ProfileReqInput(**json_data) req = ProfileReq(**json_data)
self.assertTrue(req_input.merge_profiles) self.assertTrue(req.merge_profiles)
self.assertEqual(req_input.output_dir, "/tmp/test") self.assertEqual(req.output_dir, "/tmp/test")
self.assertEqual(req_input.num_steps, 10) self.assertEqual(req.num_steps, 10)
self.assertEqual(req_input.activities, ["CPU", "GPU", "MEM"]) self.assertEqual(req.activities, ["CPU", "GPU", "MEM"])
self.assertFalse(req_input.profile_by_stage) 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 # Test with minimal data
json_data = {"output_dir": "/tmp/test"} json_data = {"output_dir": "/tmp/test"}
req_input = ProfileReqInput(**json_data) req = ProfileReq(**json_data)
self.assertFalse(req_input.merge_profiles) 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} json_data = {"output_dir": "/tmp/test", "merge_profiles": False}
req_input = ProfileReqInput(**json_data) req = ProfileReq(**json_data)
self.assertFalse(req_input.merge_profiles) self.assertFalse(req.merge_profiles)
def test_http_api_parameter_flow(self): def test_http_api_parameter_flow(self):
# Simulate HTTP request data # Simulate HTTP request data
@@ -85,8 +86,8 @@ class TestProfileMergerHTTPAPI(unittest.TestCase):
"merge_profiles": True, "merge_profiles": True,
} }
# Create ProfileReqInput as HTTP server would # Create ProfileReq as HTTP server would
obj = ProfileReqInput(**request_data) obj = ProfileReq(**request_data)
# Verify the parameter is set correctly # Verify the parameter is set correctly
self.assertTrue(obj.merge_profiles) self.assertTrue(obj.merge_profiles)
@@ -98,24 +99,24 @@ class TestProfileMergerHTTPAPI(unittest.TestCase):
def test_http_api_parameter_validation(self): def test_http_api_parameter_validation(self):
# Test with True # Test with True
json_data = {"merge_profiles": True} json_data = {"merge_profiles": True}
req_input = ProfileReqInput(**json_data) req = ProfileReq(**json_data)
self.assertTrue(req_input.merge_profiles) self.assertTrue(req.merge_profiles)
# Test with False # Test with False
json_data = {"merge_profiles": False} json_data = {"merge_profiles": False}
req_input = ProfileReqInput(**json_data) req = ProfileReq(**json_data)
self.assertFalse(req_input.merge_profiles) self.assertFalse(req.merge_profiles)
# Test with string "true" (should be converted by JSON parser) # Test with string "true" (should be converted by JSON parser)
json_data = {"merge_profiles": "true"} json_data = {"merge_profiles": "true"}
req_input = ProfileReqInput(**json_data) req = ProfileReq(**json_data)
self.assertEqual(req_input.merge_profiles, "true") # String, not boolean self.assertEqual(req.merge_profiles, "true") # String, not boolean
def test_http_api_backward_compatibility(self): def test_http_api_backward_compatibility(self):
# Test minimal request (no merge_profiles) # Test minimal request (no merge_profiles)
json_data = {} json_data = {}
req_input = ProfileReqInput(**json_data) req = ProfileReq(**json_data)
self.assertFalse(req_input.merge_profiles) # Should default to False self.assertFalse(req.merge_profiles) # Should default to False
# Test with other parameters but no merge_profiles # Test with other parameters but no merge_profiles
json_data = { json_data = {
@@ -123,8 +124,8 @@ class TestProfileMergerHTTPAPI(unittest.TestCase):
"num_steps": 5, "num_steps": 5,
"activities": ["CPU", "GPU"], "activities": ["CPU", "GPU"],
} }
req_input = ProfileReqInput(**json_data) req = ProfileReq(**json_data)
self.assertFalse(req_input.merge_profiles) # Should default to False self.assertFalse(req.merge_profiles) # Should default to False
def test_http_api_parameter_combinations(self): def test_http_api_parameter_combinations(self):
test_cases = [ test_cases = [
@@ -163,8 +164,8 @@ class TestProfileMergerHTTPAPI(unittest.TestCase):
for test_case in test_cases: for test_case in test_cases:
with self.subTest(test_case["name"]): with self.subTest(test_case["name"]):
req_input = ProfileReqInput(**test_case["data"]) req = ProfileReq(**test_case["data"])
self.assertEqual(req_input.merge_profiles, test_case["expected_merge"]) self.assertEqual(req.merge_profiles, test_case["expected_merge"])
if __name__ == "__main__": if __name__ == "__main__":
test/registered/unit/observability/test_trace.py
+3 -3
View File
@@ -227,7 +227,7 @@ class TestTraceReqContextDisabled(unittest.TestCase):
self.assertEqual(state, {"tracing_enable": False}) self.assertEqual(state, {"tracing_enable": False})
def test_setstate_disabled(self): def test_setstate_disabled(self):
ctx = TraceReqContext.__new__(TraceReqContext) ctx = TraceReqContext(rid="req-1")
ctx.__setstate__({"tracing_enable": True, "is_copy": False}) ctx.__setstate__({"tracing_enable": True, "is_copy": False})
# opentelemetry_initialized is False → tracing forced off # opentelemetry_initialized is False → tracing forced off
self.assertFalse(ctx.tracing_enable) self.assertFalse(ctx.tracing_enable)
@@ -539,7 +539,7 @@ class TestTraceReqContextEnabled(unittest.TestCase):
state = ctx.__getstate__() state = ctx.__getstate__()
ctx.trace_req_finish(ts=2000) ctx.trace_req_finish(ts=2000)
ctx2 = TraceReqContext.__new__(TraceReqContext) ctx2 = TraceReqContext(rid="req-2")
ctx2.__setstate__(state) ctx2.__setstate__(state)
self.assertTrue(ctx2.tracing_enable) self.assertTrue(ctx2.tracing_enable)
self.assertTrue(ctx2.is_copy) self.assertTrue(ctx2.is_copy)
@@ -567,7 +567,7 @@ class TestTraceReqContextEnabled(unittest.TestCase):
ctx.trace_req_finish(ts=3000) ctx.trace_req_finish(ts=3000)
self.assertIsNotNone(state.get("last_span_context")) self.assertIsNotNone(state.get("last_span_context"))
ctx2 = TraceReqContext.__new__(TraceReqContext) ctx2 = TraceReqContext(rid="req-2")
ctx2.__setstate__(state) ctx2.__setstate__(state)
self.assertIsNotNone(ctx2.last_span_context) self.assertIsNotNone(ctx2.last_span_context)
test/registered/unit/utils/test_profile_merger.py
+10 -14
View File
@@ -13,20 +13,21 @@ import shutil
import tempfile import tempfile
import unittest 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.srt.utils.profile_merger import ProfileMerger
from sglang.test.ci.ci_register import ( from sglang.test.ci.ci_register import (
register_amd_ci, register_amd_ci,
register_cpu_ci, register_cpu_ci,
register_cuda_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_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_amd_ci(est_time=8, suite="stage-b-test-1-gpu-small-amd")
register_cpu_ci(est_time=8, suite="base-c-test-cpu") register_cpu_ci(est_time=8, suite="base-c-test-cpu")
class TestProfileMerger(unittest.TestCase): class TestProfileMerger(CustomTestCase):
def setUp(self): def setUp(self):
self.temp_dir = tempfile.mkdtemp() self.temp_dir = tempfile.mkdtemp()
self.profile_id = "test_profile_123" self.profile_id = "test_profile_123"
@@ -205,21 +206,15 @@ class TestProfileMerger(unittest.TestCase):
empty_merger.merge_chrome_traces() empty_merger.merge_chrome_traces()
class TestProfileMergerIntegration(unittest.TestCase): class TestProfileMergerIntegration(CustomTestCase):
def test_data_structures_merge_profiles(self): 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 # Test ProfileReq
req = ProfileReq(type=ProfileReqType.START_PROFILE) req = ProfileReq()
self.assertFalse(req.merge_profiles) 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) self.assertTrue(req.merge_profiles)
def test_integration_parameters(self): def test_integration_parameters(self):
@@ -231,7 +226,8 @@ class TestProfileMergerIntegration(unittest.TestCase):
) )
sig = inspect.signature(TokenizerControlMixin.start_profile) 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 # Test SchedulerProfilerMixin
from sglang.srt.managers.scheduler_components.profiler_manager import ( from sglang.srt.managers.scheduler_components.profiler_manager import (
@@ -248,7 +244,7 @@ class TestProfileMergerIntegration(unittest.TestCase):
self.assertIn("merge_profiles", sig.parameters) self.assertIn("merge_profiles", sig.parameters)
class TestProfileMergerEdgeCases(unittest.TestCase): class TestProfileMergerEdgeCases(CustomTestCase):
def setUp(self): def setUp(self):
self.temp_dir = tempfile.mkdtemp() self.temp_dir = tempfile.mkdtemp()
self.profile_id = "test_edge_cases" self.profile_id = "test_edge_cases"