feat: add forward timeout (#17831)

Co-authored-by: qiuxuan.lzw <qiuxuan.lzw@alibaba-inc.com>

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
 

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")

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)

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()