make warmup timeout configurable

Author: billishyahao

python/sglang/srt/entrypoints/http_server.py

@@ -1529,11 +1529,16 @@ def _execute_server_warmup(
 
     try:
         if server_args.disaggregation_mode == "null":
+            logger.info(f"Start of co-locate warmup ...")
             res = requests.post(
                 url + request_name,
                 json=json_data,
                 headers=headers,
-                timeout=600,
+                timeout=(
+                    server_args.warmup_timeout
+                    if server_args.warmup_timeout is not None
+                    else 300
+                ),
             )
             assert res.status_code == 200, f"{res}"
             _global_state.tokenizer_manager.server_status = ServerStatus.Up
@@ -1559,7 +1564,11 @@ def _execute_server_warmup(
                 url + request_name,
                 json=json_data,
                 headers=headers,
-                timeout=1800,  # because of deep gemm precache is very long if not precache.
+                timeout=(
+                    server_args.warmup_timeout
+                    if server_args.warmup_timeout is not None
+                    else 1800
+                ),  # because of deep gemm precache is very long if not precache.
             )
             if res.status_code == 200:
                 logger.info(

python/sglang/srt/server_args.py

@@ -292,6 +292,7 @@ class ServerArgs:
     constrained_json_disable_any_whitespace: bool = False
     watchdog_timeout: float = 300
     dist_timeout: Optional[int] = None  # timeout for torch.distributed
+    warmup_timeout: Optional[float] = None  # timeout for warmup request
     download_dir: Optional[str] = None
     base_gpu_id: int = 0
     gpu_id_step: int = 1
@@ -2287,6 +2288,12 @@ def add_cli_args(parser: argparse.ArgumentParser):
             default=ServerArgs.dist_timeout,
             help="Set timeout for torch.distributed initialization.",
         )
+        parser.add_argument(
+            "--warmup-timeout",
+            type=float,
+            default=ServerArgs.warmup_timeout,
+            help="Set warmup timeout in seconds. If a warmup forward batch takes longer than this, the server will crash to prevent hanging. Recommend to increase warmup timeout to 1800 to accommodate some kernel JIT precache e.g. deep gemm",
+        )
         parser.add_argument(
             "--download-dir",
             type=str,