Migrate 4-GPU/8-GPU workflow jobs to stage-c and add CI registry decorators (#17299)

2026-01-31 22:37:22 -08:00
parent 95180484e9
commit a0bae4c343
39 changed files with 189 additions and 284 deletions
@@ -0,0 +1,109 @@
+import os
+import shutil
+import tempfile
+import unittest
+from pathlib import Path
+
+import requests
+
+from sglang.srt.environ import envs
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import 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=120,
+    suite="stage-b-test-small-1-gpu",
+    disabled="Temporarily disabled",
+)
+
+
+class TestStartProfile(CustomTestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.output_dir = tempfile.mkdtemp()
+        envs.SGLANG_TORCH_PROFILER_DIR.set(cls.output_dir)
+        envs.SGLANG_PROFILE_V2.set(True)
+        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_profile_by_stage(self):
+        self._start_profile(
+            profile_by_stage=True,
+            num_steps=10,
+        )
+
+        self._post_request()
+
+        self._check_profile_output(pattern="*-prefill*", expect_existence=True)
+        self._check_profile_output(pattern="*-decode*", expect_existence=True)
+
+    def test_decode_only(self):
+        self._start_profile(
+            profile_by_stage=True,
+            profile_stages=["decode"],
+            num_steps=10,
+        )
+
+        self._post_request()
+
+        self._check_profile_output(pattern="*-prefill*", expect_existence=False)  # NOTE
+        self._check_profile_output(pattern="*-decode*", expect_existence=True)
+
+    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(self.output_dir):
+            shutil.rmtree(self.output_dir)
+
+    def _check_profile_output(self, pattern: str, expect_existence: bool):
+        self.assertTrue(
+            os.path.isdir(self.output_dir), "Output directory does not exist."
+        )
+        self.assertEqual(
+            len(list(Path(self.output_dir).glob(pattern))) > 0,
+            expect_existence,
+            f"Does not find {pattern=} ({list(Path(self.output_dir).glob('**/*'))=})",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()