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.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()
python/sglang/srt/entrypoints/engine.py
+14 -3
View File
@@ -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)
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:
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)
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
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
)
python/sglang/srt/entrypoints/http_server.py
+3 -17
View File
@@ -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,
python/sglang/srt/managers/io_struct.py
+8 -22
View File
@@ -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
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)}")
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:
python/sglang/srt/managers/scheduler_components/profiler_manager.py
+1 -1
View File
@@ -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,
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)
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 (
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.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:
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
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,
python/sglang/srt/session/session_controller.py
+3 -3
View File
@@ -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
python/sglang/srt/utils/common.py
+34
View File
@@ -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 ---