Extract pause_resume_in_place kit; rename test_abort to test_scheduler_control (#22647)

python/sglang/test/kits/pause_generation_kit.py

@@ -0,0 +1,111 @@
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+import requests
+
+_REQUEST_TIMEOUT = 180
+
+
+class PauseResumeInPlaceMixin:
+    """Test pause/resume with in_place mode.
+
+    Sends concurrent requests, pauses mid-generation, verifies no progress
+    during the pause window, then resumes and verifies all requests complete.
+
+    Subclass must set:
+      - pause_generate_url: URL to send /generate requests (or falls back to self.base_url)
+      - pause_target_urls: list of URLs to send /pause_generation and /continue_generation
+    """
+
+    pause_num_requests: int = 32
+    pause_max_new_tokens: int = 512
+    pause_duration: float = 5
+    pause_generate_url: str = ""
+    pause_target_urls: list = []
+
+    def test_pause_resume_in_place(self):
+        generate_url = self.pause_generate_url or self.base_url
+        target_urls = self.pause_target_urls or [self.base_url]
+        num_requests = self.pause_num_requests
+
+        def _generate(prompt_id):
+            return requests.post(
+                generate_url + "/generate",
+                json={
+                    "text": f"Question {prompt_id}: Write a short essay about the number {prompt_id}.",
+                    "sampling_params": {
+                        "temperature": 0.8,
+                        "max_new_tokens": self.pause_max_new_tokens,
+                    },
+                },
+                timeout=_REQUEST_TIMEOUT,
+            )
+
+        with ThreadPoolExecutor(max_workers=num_requests) as executor:
+            futures = {executor.submit(_generate, i): i for i in range(num_requests)}
+
+            time.sleep(1)
+
+            # Pause all targets
+            for url in target_urls:
+                requests.post(
+                    url + "/pause_generation",
+                    json={"mode": "in_place"},
+                    timeout=30,
+                ).raise_for_status()
+
+            time.sleep(0.5)
+            done_before = sum(1 for f in futures if f.done())
+
+            time.sleep(self.pause_duration)
+            done_after = sum(1 for f in futures if f.done())
+
+            self.assertLess(
+                done_before,
+                num_requests,
+                "All requests completed before pause took effect -- "
+                "increase pause_max_new_tokens to make the test meaningful.",
+            )
+
+            self.assertEqual(
+                done_after - done_before,
+                0,
+                f"{done_after - done_before} requests completed during pause "
+                f"({done_before} before, {done_after} after) -- "
+                f"pause_generation was not respected by the scheduler.",
+            )
+
+            # Resume all targets (reverse order to unblock downstream first)
+            for url in reversed(target_urls):
+                requests.post(
+                    url + "/continue_generation",
+                    json={},
+                    timeout=30,
+                ).raise_for_status()
+
+            completed = 0
+            errors = []
+            for future in as_completed(futures, timeout=_REQUEST_TIMEOUT):
+                prompt_id = futures[future]
+                try:
+                    resp = future.result()
+                    if resp.status_code == 200:
+                        body = resp.json()
+                        self.assertIn("text", body)
+                        self.assertGreater(len(body["text"]), 0)
+                        completed += 1
+                    else:
+                        errors.append(f"Request {prompt_id}: status={resp.status_code}")
+                except Exception as e:
+                    errors.append(f"Request {prompt_id}: exception={e}")
+
+        self.assertEqual(
+            completed + len(errors),
+            num_requests,
+            "Some requests did not resolve within the timeout -- likely hung during pause.",
+        )
+        self.assertEqual(
+            completed,
+            num_requests,
+            f"Some requests failed: {completed}/{num_requests} succeeded. Errors: {errors}",
+        )

test/registered/disaggregation/test_disaggregation_basic.py

@@ -1,9 +1,7 @@
 import asyncio
 import json
 import os
-import time
 import unittest
-from concurrent.futures import ThreadPoolExecutor, as_completed
 from types import SimpleNamespace
 
 import aiohttp
@@ -12,6 +10,7 @@ import requests
 from transformers import AutoTokenizer
 
 from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.pause_generation_kit import PauseResumeInPlaceMixin
 from sglang.test.run_eval import run_eval
 from sglang.test.server_fixtures.disaggregation_fixture import (
     PDDisaggregationServerBase,
@@ -25,11 +24,13 @@ from sglang.test.test_utils import (
 register_cuda_ci(est_time=394, suite="stage-b-test-2-gpu-large")
 
 
-class TestDisaggregationAccuracy(PDDisaggregationServerBase):
+class TestDisaggregationAccuracy(PauseResumeInPlaceMixin, PDDisaggregationServerBase):
     @classmethod
     def setUpClass(cls):
         super().setUpClass()
         cls.model = DEFAULT_MODEL_NAME_FOR_TEST
+        cls.pause_generate_url = cls.lb_url
+        cls.pause_target_urls = [cls.prefill_url, cls.decode_url]
         cls.launch_all()
 
     def test_gsm8k(self):
@@ -99,101 +100,6 @@ class TestDisaggregationAccuracy(PDDisaggregationServerBase):
         # ensure the output is a valid JSON
         json.loads(output)
 
-    def test_pause_resume_in_place(self):
-        """Send requests, pause mid-generation, verify no progress during pause, resume."""
-        NUM_REQUESTS = 32
-        MAX_NEW_TOKENS = 512
-        REQUEST_TIMEOUT = 180
-        PAUSE_DURATION = 5
-
-        def _generate(prompt_id):
-            return requests.post(
-                self.lb_url + "/generate",
-                json={
-                    "text": f"Question {prompt_id}: Write a short essay about the number {prompt_id}.",
-                    "sampling_params": {
-                        "temperature": 0.8,
-                        "max_new_tokens": MAX_NEW_TOKENS,
-                    },
-                },
-                timeout=REQUEST_TIMEOUT,
-            )
-
-        with ThreadPoolExecutor(max_workers=NUM_REQUESTS) as executor:
-            futures = {executor.submit(_generate, i): i for i in range(NUM_REQUESTS)}
-
-            time.sleep(1)
-
-            requests.post(
-                self.prefill_url + "/pause_generation",
-                json={"mode": "in_place"},
-                timeout=30,
-            ).raise_for_status()
-            requests.post(
-                self.decode_url + "/pause_generation",
-                json={"mode": "in_place"},
-                timeout=30,
-            ).raise_for_status()
-
-            time.sleep(0.5)
-            done_before = sum(1 for f in futures if f.done())
-
-            time.sleep(PAUSE_DURATION)
-            done_after = sum(1 for f in futures if f.done())
-
-            self.assertLess(
-                done_before,
-                NUM_REQUESTS,
-                "All requests completed before pause took effect — "
-                "increase MAX_NEW_TOKENS to make the test meaningful.",
-            )
-
-            self.assertEqual(
-                done_after - done_before,
-                0,
-                f"{done_after - done_before} requests completed during pause "
-                f"({done_before} before, {done_after} after) — "
-                f"pause_generation was not respected by the disagg scheduler.",
-            )
-
-            requests.post(
-                self.decode_url + "/continue_generation",
-                json={},
-                timeout=30,
-            ).raise_for_status()
-            requests.post(
-                self.prefill_url + "/continue_generation",
-                json={},
-                timeout=30,
-            ).raise_for_status()
-
-            completed = 0
-            errors = []
-            for future in as_completed(futures, timeout=REQUEST_TIMEOUT):
-                prompt_id = futures[future]
-                try:
-                    resp = future.result()
-                    if resp.status_code == 200:
-                        body = resp.json()
-                        self.assertIn("text", body)
-                        self.assertGreater(len(body["text"]), 0)
-                        completed += 1
-                    else:
-                        errors.append(f"Request {prompt_id}: status={resp.status_code}")
-                except Exception as e:
-                    errors.append(f"Request {prompt_id}: exception={e}")
-
-        self.assertEqual(
-            completed + len(errors),
-            NUM_REQUESTS,
-            "Some requests did not resolve within the timeout — likely hung during pause.",
-        )
-        self.assertEqual(
-            completed,
-            NUM_REQUESTS,
-            f"Some requests failed: {completed}/{NUM_REQUESTS} succeeded. Errors: {errors}",
-        )
-
     def test_first_token_finish(self):
         client = openai.Client(api_key="empty", base_url=f"{self.lb_url}/v1")
         tokenizer = AutoTokenizer.from_pretrained(self.model)

test/registered/scheduler/test_scheduler_control.py

@@ -10,6 +10,7 @@ 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.kits.abort_timeout_kit import AbortAllMixin, WaitingTimeoutMixin
+from sglang.test.kits.pause_generation_kit import PauseResumeInPlaceMixin
 from sglang.test.test_utils import (
     DEFAULT_MODEL_NAME_FOR_TEST,
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -19,7 +20,7 @@ from sglang.test.test_utils import (
     run_and_check_memory_leak,
 )
 
-register_cuda_ci(est_time=328, suite="stage-b-test-1-gpu-small")
+register_cuda_ci(est_time=360, suite="stage-b-test-1-gpu-small")
 register_amd_ci(est_time=300, suite="stage-b-test-1-gpu-small-amd")
 
 
@@ -110,7 +111,7 @@ class TestAbortWithApiKey(CustomTestCase):
         )
 
 
-class TestAbortAll(AbortAllMixin, CustomTestCase):
+class TestSchedulerControl(AbortAllMixin, PauseResumeInPlaceMixin, CustomTestCase):
     @classmethod
     def setUpClass(cls):
         cls.model = DEFAULT_MODEL_NAME_FOR_TEST