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}",
+        )