diff --git a/python/sglang/srt/distributed/bootstrap.py b/python/sglang/srt/distributed/bootstrap.py index 442d204ca..fc6cd4adc 100644 --- a/python/sglang/srt/distributed/bootstrap.py +++ b/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, diff --git a/python/sglang/srt/distributed/gated_launch.py b/python/sglang/srt/distributed/gated_launch.py new file mode 100644 index 000000000..da36301ae --- /dev/null +++ b/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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 80136a2a7..f50cf13be 100644 --- a/python/sglang/srt/server_args.py +++ b/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[ diff --git a/test/registered/core/test_gated_launch.py b/test/registered/core/test_gated_launch.py new file mode 100644 index 000000000..24f6171b9 --- /dev/null +++ b/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() diff --git a/test/registered/unit/distributed/test_gated_launch.py b/test/registered/unit/distributed/test_gated_launch.py new file mode 100644 index 000000000..e5a2336dd --- /dev/null +++ b/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()