Support gated launch to defer startup memory allocation (#35927)

python/sglang/srt/distributed/bootstrap.py

@@ -20,6 +20,7 @@ from sglang.srt.distributed import (
     set_mscclpp_all_reduce,
     set_torch_symm_mem_all_reduce,
 )
+from sglang.srt.distributed.gated_launch import maybe_wait_for_gated_launch
 from sglang.srt.distributed.parallel_state import (
     _tag_groups_for_flashinfer_allreduce_only,
 )
@@ -132,6 +133,10 @@ def init_torch_distributed(
         ):
             _prewarm_tp_lm_head_all_to_all()
 
+    maybe_wait_for_gated_launch(
+        host=server_args.host, port=server_args.gated_launch_port
+    )
+
     pre_model_load_memory = get_available_gpu_memory(
         device,
         ps.gpu_id,

python/sglang/srt/distributed/gated_launch.py

@@ -0,0 +1,96 @@
+import logging
+import threading
+import time
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+import uvicorn
+from fastapi import FastAPI
+from fastapi.responses import PlainTextResponse
+
+from sglang.srt.distributed import get_world_group
+
+logger = logging.getLogger(__name__)
+
+POLL_INTERVAL_SECONDS = 1.0
+LOG_INTERVAL_SECONDS = 10.0
+
+_instance: Optional["_GatedLaunchServer"] = None
+
+
+def maybe_wait_for_gated_launch(*, host: str, port: Optional[int]) -> None:
+    global _instance
+
+    if port is None or _instance is not None:
+        return
+
+    world_group = get_world_group()
+
+    _instance = _GatedLaunchServer()
+    if world_group.rank_in_group == 0:
+        _instance.serve(host=host, port=port)
+
+    logger.info(f"Gated launch waiting for activation. rank={world_group.rank}")
+    tic = time.perf_counter()
+    _wait_until_activated(world_group=world_group, server=_instance)
+    logger.info(f"Gated launch activated. elapsed={time.perf_counter() - tic:.2f} s")
+
+
+def _wait_until_activated(*, world_group, server: "_GatedLaunchServer") -> None:
+    activated = torch.zeros(1, dtype=torch.int32)
+    started_at = time.perf_counter()
+    next_log_at = started_at + LOG_INTERVAL_SECONDS
+
+    while True:
+        activated[0] = int(server.activated)
+
+        if world_group.world_size > 1:
+            dist.broadcast(
+                activated,
+                src=world_group.ranks[0],
+                group=world_group.cpu_group,
+            )
+
+        if bool(activated[0]):
+            return
+
+        if (now := time.perf_counter()) >= next_log_at:
+            logger.info(
+                f"Gated launch still waiting for activation. "
+                f"rank={world_group.rank} elapsed={now - started_at:.0f} s"
+            )
+            next_log_at = now + LOG_INTERVAL_SECONDS
+
+        time.sleep(POLL_INTERVAL_SECONDS)
+
+
+class _GatedLaunchServer:
+    def __init__(self):
+        self.activated = False
+        self._server: Optional[uvicorn.Server] = None
+        self._thread: Optional[threading.Thread] = None
+
+    def serve(self, *, host: str, port: int) -> None:
+        config = uvicorn.Config(
+            _build_app(self), host=host, port=port, log_level="warning"
+        )
+        self._server = uvicorn.Server(config)
+        self._thread = threading.Thread(target=self._server.run, daemon=True)
+        self._thread.start()
+        logger.info(f"Gated launch control server started on {host}:{port}")
+
+
+def _build_app(server: _GatedLaunchServer) -> FastAPI:
+    app = FastAPI()
+
+    @app.get("/health")
+    def health():
+        return PlainTextResponse("OK")
+
+    @app.post("/gate/activate")
+    def activate():
+        server.activated = True
+        return PlainTextResponse("OK")
+
+    return app

python/sglang/srt/server_args.py

@@ -1046,6 +1046,11 @@ class ServerArgs:
         ),
         NS("parallel"),
     ] = None
+    gated_launch_port: A[
+        Optional[int],
+        "The port of the gated launch control server. When set, every rank blocks right after the distributed environment is initialized, before any sizable GPU allocation, until `POST /gate/activate` is sent to this port on the host of the first rank. This lets an external orchestrator defer the memory hungry part of startup to a safe window. Defaults to None, which disables the gate.",
+        NS("parallel"),
+    ] = None
     nnodes: A[int, "The number of nodes.", NS("parallel")] = 1
     node_rank: A[int, "The node rank.", NS("parallel")] = 0
     tp_size: A[

test/registered/core/test_gated_launch.py

@@ -0,0 +1,124 @@
+import os
+import subprocess
+import time
+import unittest
+
+import psutil
+import requests
+
+from sglang.srt.utils.common import kill_process_tree
+from sglang.srt.utils.network import get_open_port
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+)
+
+register_cuda_ci(
+    est_time=180, stage="nightly", runner_config="1-gpu-large", nightly=True
+)
+
+MEM_FRACTION_STATIC = 0.6
+GATED_MEMORY_CEILING_MB = 8 * 1024
+SERVING_MEMORY_FLOOR_MB = 8 * 1024
+
+
+class TestGatedLaunch(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        _, host, port = cls.base_url.split(":")
+        cls.gate_port = get_open_port()
+        cls.gate_url = f"http:{host}:{cls.gate_port}"
+
+        command = [
+            "sglang",
+            "serve",
+            "--model-path",
+            DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+            "--host",
+            host[2:],
+            "--port",
+            port,
+            "--gated-launch-port",
+            str(cls.gate_port),
+            "--mem-fraction-static",
+            str(MEM_FRACTION_STATIC),
+        ]
+        cls.process = subprocess.Popen(command, env=os.environ.copy())
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process") and cls.process:
+            kill_process_tree(cls.process.pid)
+
+    def test_gated_launch_defers_startup_until_activated(self):
+        """The engine holds off every sizable allocation until it is activated."""
+        self._wait_for_health(self.gate_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH)
+
+        with self.assertRaises(requests.exceptions.RequestException):
+            requests.get(f"{self.base_url}/health", timeout=5)
+
+        gated_memory_mb = self._device_memory_mb()
+        self.assertLess(gated_memory_mb, GATED_MEMORY_CEILING_MB)
+
+        for _ in range(2):
+            response = requests.post(f"{self.gate_url}/gate/activate", timeout=5)
+            self.assertEqual(response.status_code, 200)
+
+        self._wait_for_health(self.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH)
+
+        response = requests.post(
+            f"{self.base_url}/generate",
+            json={
+                "text": "The capital of France is",
+                "sampling_params": {"max_new_tokens": 8, "temperature": 0},
+            },
+            timeout=60,
+        )
+        self.assertEqual(response.status_code, 200)
+        self.assertTrue(response.json()["text"])
+
+        self.assertGreater(self._device_memory_mb(), SERVING_MEMORY_FLOOR_MB)
+
+    def _wait_for_health(self, url: str, timeout: float) -> None:
+        deadline = time.perf_counter() + timeout
+        while time.perf_counter() < deadline:
+            self.assertIsNone(
+                self.process.poll(), msg=f"server died while waiting for {url}"
+            )
+            try:
+                if requests.get(f"{url}/health", timeout=5).status_code == 200:
+                    return
+            except requests.exceptions.RequestException:
+                pass
+            time.sleep(1)
+        self.fail(f"{url} did not become healthy within {timeout}s")
+
+    def _device_memory_mb(self) -> int:
+        parent = psutil.Process(self.process.pid)
+        pids = {parent.pid} | {child.pid for child in parent.children(recursive=True)}
+
+        output = subprocess.check_output(
+            [
+                "nvidia-smi",
+                "--query-compute-apps=pid,used_gpu_memory",
+                "--format=csv,noheader,nounits",
+            ],
+            text=True,
+        )
+
+        total_mb = 0
+        for line in output.splitlines():
+            if not line.strip():
+                continue
+            pid, used_mb = (field.strip() for field in line.split(","))
+            if int(pid) in pids:
+                total_mb += int(used_mb)
+        return total_mb
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/distributed/test_gated_launch.py

@@ -0,0 +1,197 @@
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=30, suite="base-a-test-cpu")
+
+import multiprocessing
+import threading
+import time
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import requests
+import torch.distributed as dist
+
+from sglang.srt.distributed import gated_launch
+from sglang.srt.distributed.gated_launch import (
+    POLL_INTERVAL_SECONDS,
+    _GatedLaunchServer,
+    _wait_until_activated,
+    maybe_wait_for_gated_launch,
+)
+from sglang.srt.utils.network import get_open_port
+from sglang.test.test_utils import CustomTestCase
+
+_JOIN_TIMEOUT_SECONDS = 120
+
+
+def _fake_world_group(*, rank: int, world_size: int, cpu_group=None):
+    return SimpleNamespace(
+        rank=rank,
+        rank_in_group=rank,
+        ranks=list(range(world_size)),
+        world_size=world_size,
+        cpu_group=cpu_group,
+    )
+
+
+def _activate_over_http(base_url: str, delay: float):
+    time.sleep(delay)
+    deadline = time.perf_counter() + 30
+    while time.perf_counter() < deadline:
+        try:
+            if requests.post(f"{base_url}/gate/activate", timeout=1).status_code == 200:
+                return
+        except requests.exceptions.RequestException:
+            pass
+        time.sleep(0.1)
+
+
+def _gated_launch_worker(
+    rank: int, dist_port: int, gate_port: int, activate_after: float, out
+):
+    dist.init_process_group(
+        backend="gloo",
+        init_method=f"tcp://127.0.0.1:{dist_port}",
+        rank=rank,
+        world_size=2,
+    )
+    try:
+        if rank == 0:
+            threading.Thread(
+                target=_activate_over_http,
+                args=(f"http://127.0.0.1:{gate_port}", activate_after),
+                daemon=True,
+            ).start()
+
+        world_group = _fake_world_group(
+            rank=rank, world_size=2, cpu_group=dist.group.WORLD
+        )
+        started_at = time.perf_counter()
+        with patch.object(gated_launch, "get_world_group", return_value=world_group):
+            maybe_wait_for_gated_launch(host="127.0.0.1", port=gate_port)
+        out.put((rank, time.perf_counter() - started_at))
+    finally:
+        dist.destroy_process_group()
+
+
+class TestGatedLaunchServer(CustomTestCase):
+    def setUp(self):
+        self.server = _GatedLaunchServer()
+        self.port = get_open_port()
+        self.server.serve(host="127.0.0.1", port=self.port)
+        self.base_url = f"http://127.0.0.1:{self.port}"
+        self._wait_until_listening()
+
+    def _wait_until_listening(self):
+        deadline = time.perf_counter() + 30
+        while time.perf_counter() < deadline:
+            try:
+                requests.get(f"{self.base_url}/health", timeout=1)
+                return
+            except requests.exceptions.RequestException:
+                time.sleep(0.1)
+        self.fail(f"control server did not start listening on port {self.port}")
+
+    def test_health_is_served_while_the_gate_is_still_closed(self):
+        """The control port answers before activation so a caller can find it."""
+        response = requests.get(f"{self.base_url}/health", timeout=5)
+
+        self.assertEqual(response.status_code, 200)
+        self.assertFalse(self.server.activated)
+
+    def test_activate_flips_the_flag_and_stays_successful_when_repeated(self):
+        """A retried activation still succeeds instead of erroring or toggling back."""
+        first = requests.post(f"{self.base_url}/gate/activate", timeout=5)
+        self.assertEqual(first.status_code, 200)
+        self.assertTrue(self.server.activated)
+
+        second = requests.post(f"{self.base_url}/gate/activate", timeout=5)
+        self.assertEqual(second.status_code, 200)
+        self.assertTrue(self.server.activated)
+
+    def test_activating_one_server_leaves_another_one_closed(self):
+        """The route acts on its own server instead of process wide state."""
+        other = _GatedLaunchServer()
+        other_port = get_open_port()
+        other.serve(host="127.0.0.1", port=other_port)
+
+        requests.post(f"{self.base_url}/gate/activate", timeout=5)
+
+        self.assertTrue(self.server.activated)
+        self.assertFalse(other.activated)
+
+
+class TestWaitUntilActivated(CustomTestCase):
+    def test_single_rank_keeps_polling_until_the_flag_is_set(self):
+        """A lone rank leaves the gate only after its own flag flips."""
+        server = _GatedLaunchServer()
+        activate_after = 2 * POLL_INTERVAL_SECONDS
+        threading.Timer(
+            activate_after, lambda: setattr(server, "activated", True)
+        ).start()
+
+        started_at = time.perf_counter()
+        _wait_until_activated(
+            world_group=_fake_world_group(rank=0, world_size=1), server=server
+        )
+        elapsed = time.perf_counter() - started_at
+
+        self.assertGreaterEqual(elapsed, activate_after)
+
+    def test_second_rank_learns_about_activation_through_the_cpu_group(self):
+        """Driven through maybe_wait_for_gated_launch: the rank without the control server is released by the gloo broadcast."""
+        context = multiprocessing.get_context("spawn")
+        out = context.Queue()
+        dist_port = get_open_port()
+        gate_port = get_open_port()
+        activate_after = 2 * POLL_INTERVAL_SECONDS
+
+        processes = [
+            context.Process(
+                target=_gated_launch_worker,
+                args=(rank, dist_port, gate_port, activate_after, out),
+            )
+            for rank in range(2)
+        ]
+        for process in processes:
+            process.start()
+
+        elapsed_by_rank = {}
+        for _ in processes:
+            rank, elapsed = out.get(timeout=_JOIN_TIMEOUT_SECONDS)
+            elapsed_by_rank[rank] = elapsed
+
+        for process in processes:
+            process.join(timeout=_JOIN_TIMEOUT_SECONDS)
+        for process in processes:
+            self.assertEqual(process.exitcode, 0)
+        self.assertEqual(sorted(elapsed_by_rank), [0, 1])
+        self.assertGreaterEqual(elapsed_by_rank[1], activate_after)
+
+
+class TestMaybeWaitForGatedLaunch(CustomTestCase):
+    def setUp(self):
+        self.addCleanup(setattr, gated_launch, "_instance", None)
+        gated_launch._instance = None
+
+    def test_an_unset_port_leaves_the_startup_path_untouched(self):
+        """Without the flag the gate never reaches the distributed environment."""
+        with patch.object(gated_launch, "get_world_group") as get_world_group:
+            maybe_wait_for_gated_launch(host="127.0.0.1", port=None)
+
+        get_world_group.assert_not_called()
+        self.assertIsNone(gated_launch._instance)
+
+    def test_a_second_call_in_the_same_process_does_not_gate_again(self):
+        """A draft worker re-entering the init path must not wait a second time."""
+        gated_launch._instance = _GatedLaunchServer()
+
+        with patch.object(gated_launch, "get_world_group") as get_world_group:
+            maybe_wait_for_gated_launch(host="127.0.0.1", port=get_open_port())
+
+        get_world_group.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()