diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 8383ceb1f..4a9690d11 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -17,7 +17,8 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_HEALTH_CHECK_TIMEOUT` | Timeout for health check in seconds | `20` | | `SGLANG_EPLB_HEATMAP_COLLECTION_INTERVAL` | The interval of passes to collect the metric of selected count of physical experts on each layer and GPU rank. 0 means disabled. | `0` | | `SGLANG_FORWARD_UNKNOWN_TOOLS` | Forward unknown tool calls to clients instead of dropping them | `false` (drop unknown tools) | -| `SGLANG_QUEUED_TIMEOUT_MS` | Timeout (in ms) for requests in the waiting queue | `-1` | +| `SGLANG_QUEUED_TIMEOUT_MS` | Timeout (in ms) for requests in the waiting queue | `-1` | +| `SGLANG_FORWARD_TIMEOUT_MS` | Timeout (in ms) for requests in the forward batch | `-1` | ## Performance Tuning diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 702e065dd..1d33cd411 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -250,6 +250,7 @@ class Envs: SGLANG_DATA_PARALLEL_BUDGET_INTERVAL = EnvInt(1) SGLANG_QUEUED_TIMEOUT_MS = EnvInt(-1) SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH = EnvBool(False) + SGLANG_FORWARD_TIMEOUT_MS = EnvInt(-1) # Test: pd-disaggregation SGLANG_TEST_PD_DISAGG_BACKEND = EnvStr("mooncake") diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index c4728b714..2c19dd3e8 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -2,6 +2,7 @@ from __future__ import annotations import logging import time +from http import HTTPStatus from typing import TYPE_CHECKING, List, Optional, Tuple, Union import torch @@ -16,6 +17,7 @@ from sglang.srt.managers.io_struct import ( BatchTokenIDOutput, ) from sglang.srt.managers.schedule_batch import ( + FINISH_ABORT, BaseFinishReason, Req, RequestStage, @@ -122,7 +124,19 @@ class SchedulerOutputProcessorMixin: # Check finish conditions logprob_pt = 0 + deadline = -1 + if (timeout_ms := envs.SGLANG_FORWARD_TIMEOUT_MS.get()) > 0: + deadline = time.perf_counter() - timeout_ms / 1000.0 + for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)): + if ( + not req.finished() + and 0 < req.time_stats.forward_entry_time < deadline + ): + req.to_finish = FINISH_ABORT( + "Forward timeout.", HTTPStatus.SERVICE_UNAVAILABLE + ) + if req.finished() or req.is_retracted: # decode req in mixed batch or retracted req continue @@ -388,10 +402,20 @@ class SchedulerOutputProcessorMixin: # NOTE: in any case, we should check finish here # if finished, also clean up committed kv cache and over-allocated kv cache here + deadline = -1 + if (timeout_ms := envs.SGLANG_FORWARD_TIMEOUT_MS.get()) > 0: + deadline = time.perf_counter() - timeout_ms / 1000.0 + # Check finish condition for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)): req: Req + if not req.finished() and 0 < req.time_stats.forward_entry_time < deadline: + # req.set_finish_with_abort() + req.to_finish = FINISH_ABORT( + "Forward timeout.", HTTPStatus.SERVICE_UNAVAILABLE + ) + if self.enable_overlap and (req.finished() or req.is_retracted): # NOTE: This (req.finished() or req.is_retracted) should only happen when overlap scheduling is enabled. # (currently not, e.g. Eagle V1 still check finish during forward) diff --git a/test/registered/scheduler/test_abort.py b/test/registered/scheduler/test_abort.py index 27ba6406e..e57119efc 100644 --- a/test/registered/scheduler/test_abort.py +++ b/test/registered/scheduler/test_abort.py @@ -283,5 +283,41 @@ class TestAbortWithQueueTimeout(CustomTestCase): self.assertEqual(error_count, 1) +class TestAbortWithForwardTimeout(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST + cls.base_url = DEFAULT_URL_FOR_TEST + with envs.SGLANG_FORWARD_TIMEOUT_MS.override( + 1 + ), envs.SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION.override(False): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=["--skip-server-warmup"], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_forward_timeout(self): + response = requests.post( + self.base_url + "/generate", + json={ + "text": "Today is ", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 512, + "ignore_eos": True, + }, + }, + ) + result = response.json() + self.assertEqual(result["object"], "error") + self.assertEqual(result["code"], 503) + + if __name__ == "__main__": unittest.main()