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

test/srt/run_suite.py

@@ -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),

test/srt/test_start_profile.py

@@ -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()