EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[feature] Add start step profile argument in /start_profile (#7608)

Browse Source
This commit is contained in:
kyleliang-nv
2025-07-09 18:42:15 -07:00
committed by GitHub
parent 7590f5224b
commit dd445a41f5
6 changed files with 142 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/srt/entrypoints/http_server.py
View File

@@ -418,6 +418,7 @@ async def start_profile_async(obj: Optional[ProfileReqInput] = None):
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,

2
python/sglang/srt/managers/io_struct.py
View File

@@ -905,6 +905,7 @@ class ProfileReqInput:
# If set, it profile as many as this number of steps.
# If it is set, profiling is automatically stopped after this step, and
# the caller doesn't need to run stop_profile.
start_step: Optional[int] = None
num_steps: Optional[int] = None
activities: Optional[List[str]] = None
profile_by_stage: bool = False
@@ -932,6 +933,7 @@ class ExpertDistributionReqOutput:
class ProfileReq:
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

22
python/sglang/srt/managers/scheduler.py
View File

@@ -485,6 +485,8 @@ class Scheduler(
enable=server_args.enable_memory_saver
)
self.init_profier()
# Init metrics stats
self.init_metrics()
self.init_kv_events(server_args.kv_events_config)
@@ -628,6 +630,7 @@ class Scheduler(
self.torch_profiler_output_dir: Optional[str] = None
self.profiler_activities: Optional[List[str]] = None
self.profile_id: Optional[str] = None
self.profiler_start_forward_ct: Optional[int] = None
self.profiler_target_forward_ct: Optional[int] = None
self.profiler_target_prefill_ct: Optional[int] = None
self.profiler_target_decode_ct: Optional[int] = None
@@ -2389,9 +2392,10 @@ class Scheduler(
def profile(self, recv_req: ProfileReq):
if recv_req.type == ProfileReqType.START_PROFILE:
if recv_req.profile_by_stage:
if recv_req.profile_by_stage or recv_req.start_step:
return self.init_profile(
recv_req.output_dir,
recv_req.start_step,
recv_req.num_steps,
recv_req.activities,
recv_req.with_stack,
@@ -2402,6 +2406,7 @@ class Scheduler(
else:
self.init_profile(
recv_req.output_dir,
recv_req.start_step,
recv_req.num_steps,
recv_req.activities,
recv_req.with_stack,
@@ -2416,6 +2421,7 @@ class Scheduler(
def init_profile(
self,
output_dir: Optional[str],
start_step: Optional[int],
num_steps: Optional[int],
activities: Optional[List[str]],
with_stack: Optional[bool],
@@ -2442,6 +2448,9 @@ class Scheduler(
self.profiler_activities = activities
self.profile_id = profile_id
if start_step:
self.profiler_start_forward_ct = max(start_step, self.forward_ct + 1)
if num_steps:
self.profile_steps = num_steps
if self.profile_by_stage:
@@ -2449,6 +2458,10 @@ class Scheduler(
self.profiler_target_decode_ct = num_steps
self.profiler_prefill_ct = 0
self.profiler_decode_ct = 0
elif start_step:
self.profiler_target_forward_ct = (
self.profiler_start_forward_ct + num_steps
)
else:
self.profiler_target_forward_ct = self.forward_ct + num_steps
# The caller will be notified when reaching profiler_target_forward_ct
@@ -2521,6 +2534,7 @@ class Scheduler(
if "CUDA_PROFILER" in activities:
torch.cuda.cudart().cudaProfilerStart()
self.profile_in_progress = True
return ProfileReqOutput(success=True, message="Succeeded")
@@ -2584,6 +2598,7 @@ class Scheduler(
)
self.torch_profiler = None
self.profile_in_progress = False
self.profiler_start_forward_ct = None
return ProfileReqOutput(success=True, message="Succeeded.")
@@ -2617,6 +2632,11 @@ class Scheduler(
and self.profiler_target_forward_ct <= self.forward_ct
):
self.stop_profile()
if (
self.profiler_start_forward_ct
and self.profiler_start_forward_ct == self.forward_ct
):
self.start_profile()
def expert_distribution_handle(self, recv_req: ExpertDistributionReq):
if recv_req == ExpertDistributionReq.START_RECORD:

2
python/sglang/srt/managers/tokenizer_manager.py
View File

@@ -863,6 +863,7 @@ class TokenizerManager:
async def start_profile(
self,
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,
@@ -875,6 +876,7 @@ class TokenizerManager:
req = ProfileReq(
type=ProfileReqType.START_PROFILE,
output_dir=output_dir,
start_step=start_step,
num_steps=num_steps,
activities=activities,
with_stack=with_stack,

1
test/srt/run_suite.py
View File

@@ -87,6 +87,7 @@ suites = {
TestFile("test_skip_tokenizer_init.py", 117),
TestFile("test_srt_engine.py", 261),
TestFile("test_srt_endpoint.py", 130),
TestFile("test_start_profile.py", 60),
TestFile("test_torch_compile.py", 76),
TestFile("test_torch_compile_moe.py", 172),
TestFile("test_torch_native_attention_backend.py", 123),

115
test/srt/test_start_profile.py Normal file
View File

@@ -0,0 +1,115 @@
"""
Usage:
python3 -m unittest test_srt_engine.TestSRTEngine.test_4_sync_async_stream_combination
"""
import os
import shutil
import unittest
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
OUTPUT_DIR = "./profiler_dir"
class TestStartProfile(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def setUp(self):
self._clear_profile_dir()
def test_start_profile_1(self):
"""Test /start_profile with start_step and num_steps argument. This have to be the first test for start_step to work"""
response = self._start_profile(start_step="15", num_steps=5)
self._post_request()
self._check_non_empty_profile_dir()
def test_start_profile_2(self):
"""Test /start_profile with no argument"""
response = self._start_profile()
self._post_request()
# Before /stop_profile, the profile directory should be empty
self._check_empty_profile_dir()
# Post /stop_profile and check the profile directory is non-empty
response = requests.post(
f"{DEFAULT_URL_FOR_TEST}/stop_profile",
)
self._check_non_empty_profile_dir()
def test_start_profile_3(self):
"""Test /start_profile with num_steps argument"""
response = self._start_profile(num_steps=5)
self._post_request()
self._check_non_empty_profile_dir()
def _start_profile(self, **kwargs):
"""Start profiling with optional parameters."""
response = requests.post(
f"{DEFAULT_URL_FOR_TEST}/start_profile",
json=kwargs if kwargs else None,
)
self.assertEqual(response.status_code, 200)
def _post_request(self):
response = requests.post(
f"{DEFAULT_URL_FOR_TEST}/generate",
json={
"text": "The capital of France is",
"sampling_params": {
"temperature": 0,
"max_new_tokens": 32,
},
},
)
self.assertEqual(response.status_code, 200)
def _clear_profile_dir(self):
if os.path.isdir(OUTPUT_DIR):
# Remove the directory and all its contents
shutil.rmtree(OUTPUT_DIR)
def _check_non_empty_profile_dir(self):
self.assertTrue(os.path.isdir(OUTPUT_DIR), "Output directory does not exist.")
self.assertNotEqual(
len(os.listdir(OUTPUT_DIR)), 0, "Output directory is empty!"
)
def _check_empty_profile_dir(self):
if os.path.isdir(OUTPUT_DIR):
self.assertEqual(
len(os.listdir(OUTPUT_DIR)), 0, "Output directory is non-empty!"
)
if __name__ == "__main__":
os.environ["SGLANG_TORCH_PROFILER_DIR"] = OUTPUT_DIR
unittest.main()