347 lines
11 KiB
Python
347 lines
11 KiB
Python
"""
|
|
Usage:
|
|
# 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
|
|
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.utils import kill_process_tree
|
|
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
|
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,
|
|
)
|
|
|
|
register_cuda_ci(est_time=41, suite="stage-b-test-1-gpu-small")
|
|
register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd")
|
|
|
|
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
|
|
def setUpClass(cls):
|
|
envs.SGLANG_TORCH_PROFILER_DIR.set(OUTPUT_DIR)
|
|
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!"
|
|
)
|
|
|
|
|
|
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()
|