chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[feature] Add layerwise NVTX support (#11870)

Browse Source
This commit is contained in:
kyleliang-nv
2025-11-15 19:20:56 -08:00
committed by GitHub
parent 1ca205f6da
commit 597d416070
7 changed files with 718 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/srt/run_suite.py
+2 -2
View File
@@ -117,7 +117,7 @@ suites = {
TestFile("test_srt_endpoint.py", 130),
TestFile("test_srt_engine.py", 450),
TestFile("test_standalone_speculative_decoding.py", 150),
TestFile("test_start_profile.py", 60),
TestFile("test_start_profile.py", 180),
TestFile("test_profile_merger.py", 60),
TestFile("test_profile_merger_http_api.py", 15),
TestFile("test_swa_unittest.py", 1),
@@ -213,8 +213,8 @@ suites = {
TestFile("lora/test_lora_openai_api.py", 30),
TestFile("openai_server/features/test_lora_openai_compatible.py", 150),
TestFile("batch_invariant/test_batch_invariant_ops.py", 10),
TestFile("test_deepseek_v3_deterministic.py", 240),
TestFile("test_cpp_radix_cache.py", 60),
TestFile("test_deepseek_v3_deterministic.py", 240),
],
"nightly-4-gpu-b200": [
TestFile("nightly/test_flashinfer_trtllm_gen_moe_backend.py", 300),
test/srt/test_start_profile.py
+227 -1
View File
@@ -1,10 +1,19 @@
"""
Usage:
python3 -m unittest test_srt_engine.TestSRTEngine.test_4_sync_async_stream_combination
# From the test/srt directory:
cd test/srt
python3 -m unittest test_start_profile.TestStartProfile
python3 -m unittest test_start_profile.TestStartProfileWithNsys
# Run specific tests:
python3 -m unittest test_start_profile.TestStartProfile.test_start_profile_1
python3 -m unittest test_start_profile.TestStartProfileWithNsys.test_start_profile_cuda_profiler
"""
import os
import shutil
import subprocess
import time
import unittest
import requests
@@ -22,6 +31,15 @@ from sglang.test.test_utils import (
OUTPUT_DIR = "./profiler_dir"
def _is_nsys_available():
"""Check if nsys (Nsight Systems) is available on the system."""
try:
result = subprocess.run(["nsys", "--version"], capture_output=True, timeout=5)
return result.returncode == 0
except (FileNotFoundError, subprocess.TimeoutExpired):
return False
class TestStartProfile(CustomTestCase):
@classmethod
@@ -112,5 +130,213 @@ class TestStartProfile(CustomTestCase):
)
class TestStartProfileWithNsys(CustomTestCase):
"""Test /start_profile with CUDA_PROFILER (requires nsys wrapper)
Each test starts its own clean server instance with nsys profiling.
"""
@classmethod
def setUpClass(cls):
if not _is_nsys_available():
raise unittest.SkipTest("nsys (Nsight Systems) is not available")
envs.SGLANG_TORCH_PROFILER_DIR.set(OUTPUT_DIR)
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
# Use a different port to avoid conflicts with other tests
cls.base_url = "http://127.0.0.1:21100"
def setUp(self):
"""Start a clean server with nsys for each test"""
# Kill any existing processes on this port
self._kill_existing_server()
# Clean up old profile files for this test
test_name = self.id().split(".")[-1] # Get test method name
self.nsys_output_file = f"nsys_profile_{test_name}"
if os.path.isdir(OUTPUT_DIR):
profile_file = os.path.join(OUTPUT_DIR, f"{self.nsys_output_file}.nsys-rep")
if os.path.exists(profile_file):
try:
os.remove(profile_file)
except OSError:
pass
# Launch server with nsys wrapper
self.process = self._popen_launch_server_with_nsys(
self.model,
self.base_url,
self.nsys_output_file,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
)
def tearDown(self):
"""Kill server and verify profile was created"""
# Kill server first to let nsys finalize the .nsys-rep file
kill_process_tree(self.process.pid)
# Also ensure nsys agent processes are killed
try:
subprocess.run(
["pkill", "-f", "nsys.*--start-agent"],
timeout=5,
stderr=subprocess.DEVNULL,
)
except (subprocess.TimeoutExpired, FileNotFoundError):
pass
# Give nsys a moment to finalize the report
time.sleep(3)
# Verify the .nsys-rep file was created
self._verify_nsys_profile_created()
def _kill_existing_server(self):
"""Kill any existing server process on our port and orphaned nsys agents"""
try:
# Kill server on our port
subprocess.run(["lsof", "-ti", ":21100"], capture_output=True, timeout=5)
subprocess.run(["pkill", "-f", "sglang.launch_server.*21100"], timeout=5)
# Kill any orphaned nsys agent processes
subprocess.run(
["pkill", "-f", "nsys.*--start-agent"],
timeout=5,
stderr=subprocess.DEVNULL, # Suppress "no process found" errors
)
time.sleep(2) # Wait for cleanup
except (subprocess.TimeoutExpired, FileNotFoundError):
pass
def _popen_launch_server_with_nsys(self, model, base_url, output_file, timeout):
"""Launch server wrapped with nsys profile -c cudaProfilerApi
Each test gets its own output file for complete isolation.
"""
_, host, port = base_url.split(":")
host = host[2:]
# Build the server launch command
command = [
"nsys",
"profile",
"-c",
"cudaProfilerApi",
"--capture-range-end",
"stop", # Stop after first cudaProfilerStop()
"-o",
os.path.join(OUTPUT_DIR, output_file),
"python3",
"-m",
"sglang.launch_server",
"--model-path",
model,
"--host",
host,
"--port",
port,
]
# Create output directory if it doesn't exist
os.makedirs(OUTPUT_DIR, exist_ok=True)
# Launch the process - capture output to keep test output clean
process = subprocess.Popen(
command,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
)
# Wait for server to be ready
start_time = time.perf_counter()
elapsed = 0
with requests.Session() as session:
while elapsed < timeout:
elapsed = time.perf_counter() - start_time
return_code = process.poll()
if return_code is not None:
raise Exception(
f"Server process exited with code {return_code}. "
"Check server logs above for errors."
)
try:
response = session.get(f"{base_url}/health_generate", timeout=5)
if response.status_code == 200:
return process
except (requests.RequestException, requests.Timeout):
pass
time.sleep(5)
# Timeout reached
kill_process_tree(process.pid)
raise TimeoutError(
f"Server failed to start within {timeout} seconds. "
f"Check the server logs above for more information."
)
def _verify_nsys_profile_created(self):
"""Verify that the .nsys-rep file was created after server shutdown."""
if not os.path.isdir(OUTPUT_DIR):
raise AssertionError("Output directory does not exist.")
expected_file = f"{self.nsys_output_file}.nsys-rep"
profile_path = os.path.join(OUTPUT_DIR, expected_file)
if not os.path.exists(profile_path):
files = os.listdir(OUTPUT_DIR)
raise AssertionError(
f"Expected profile file '{expected_file}' not found. "
f"Files present: {files}"
)
def test_start_profile_cuda_profiler_with_start_step(self):
"""Test /start_profile with CUDA_PROFILER, start_step, and num_steps"""
# Use start_step to let server warm up before profiling
response = self._start_profile(
activities=["CUDA_PROFILER"], start_step=10, num_steps=3
)
self._post_request()
# Profile verification happens in tearDown()
def test_start_profile_cuda_profiler(self):
"""Test /start_profile with CUDA_PROFILER activity (no start_step)"""
# Simple num_steps test - profiling starts immediately
response = self._start_profile(activities=["CUDA_PROFILER"], num_steps=5)
self._post_request()
# Profile verification happens in tearDown()
def _start_profile(self, **kwargs):
"""Start profiling with optional parameters."""
response = requests.post(
f"{self.base_url}/start_profile",
json=kwargs if kwargs else None,
)
self.assertEqual(response.status_code, 200)
return response
def _post_request(self):
response = requests.post(
f"{self.base_url}/generate",
json={
"text": "The capital of France is",
"sampling_params": {
"temperature": 0,
"max_new_tokens": 32,
},
},
)
self.assertEqual(response.status_code, 200)
if __name__ == "__main__":
unittest.main()