diff --git a/python/sglang/test/kits/pause_generation_kit.py b/python/sglang/test/kits/pause_generation_kit.py new file mode 100644 index 000000000..00c1ef4cf --- /dev/null +++ b/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}", + ) diff --git a/test/registered/disaggregation/test_disaggregation_basic.py b/test/registered/disaggregation/test_disaggregation_basic.py index deeda4a8a..e20d59ea8 100644 --- a/test/registered/disaggregation/test_disaggregation_basic.py +++ b/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) diff --git a/test/registered/scheduler/test_abort.py b/test/registered/scheduler/test_scheduler_control.py similarity index 98% rename from test/registered/scheduler/test_abort.py rename to test/registered/scheduler/test_scheduler_control.py index b82a22e9a..2042ff00f 100644 --- a/test/registered/scheduler/test_abort.py +++ b/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