[Bugfix] Keep the multi-port DP supervisor serving through fault-tolerant recovery

tests/entrypoints/launchers/test_dp_supervisor.py

@@ -283,10 +283,14 @@ def get_process_timeout(request_timeout, manager_timeout):
 
 
 @pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("fault_tolerant", "was_ready"), [(False, False), (True, False), (False, True)]
+)
 async def test_handles_child_exit(
-    monkeypatch: pytest.MonkeyPatch,
+    monkeypatch: pytest.MonkeyPatch, fault_tolerant: bool, was_ready: bool
 ):
-    supervisor = DPSupervisor(_make_unit_args())
+    supervisor = DPSupervisor(_make_unit_args(enable_fault_tolerance=fault_tolerant))
+    supervisor._was_ready = was_ready
     supervisor._processes = [
         SimpleNamespace(
             name="APIServer_DPRank_4", exitcode=None, is_alive=lambda: True
@@ -320,6 +324,111 @@ async def fake_probe(*_args, **_kwargs) -> bool:
     assert supervisor._is_ready is False
 
 
+@pytest.mark.asyncio
+async def test_fault_tolerant_keeps_survivors_on_child_exit(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """An exited rank with live peers is a scale-down, not a group failure."""
+    supervisor = DPSupervisor(
+        _make_unit_args(enable_fault_tolerance=True, dp_supervisor_probe_interval_s=0.0)
+    )
+    supervisor._was_ready = True
+    survivor_alive = iter([True, True, False])
+    supervisor._processes = [
+        SimpleNamespace(
+            name="APIServer_DPRank_4", is_alive=lambda: next(survivor_alive)
+        ),
+        SimpleNamespace(name="APIServer_DPRank_5", is_alive=lambda: False),
+    ]
+
+    async def fake_child_state(_session, _port) -> str:
+        return "ok"
+
+    monkeypatch.setattr(supervisor, "_child_state", fake_child_state)
+
+    await supervisor._monitor_children()
+    assert supervisor._removed_ports == {8001}
+    assert not supervisor._shutdown_event.is_set()
+    assert next(survivor_alive, None) is None
+
+
+@pytest.mark.asyncio
+async def test_fault_tolerant_waits_while_recovering(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """A recovering rank is not torn down; a gone one only after the threshold."""
+    supervisor = DPSupervisor(
+        _make_unit_args(enable_fault_tolerance=True, dp_supervisor_probe_interval_s=0.0)
+    )
+    supervisor.child_ports = [8000]
+    states = iter(["gone", "recovering", "gone", "gone", "gone"])
+    ready_seen: list[bool] = []
+
+    async def fake_probe(*_args, **_kwargs) -> bool:
+        return True
+
+    async def fake_child_state(_session, _port) -> str:
+        ready_seen.append(supervisor.is_ready)
+        return next(states)
+
+    monkeypatch.setattr(dp_sup, "_probe_endpoint", fake_probe)
+    monkeypatch.setattr(supervisor, "_child_state", fake_child_state)
+
+    await supervisor._probe_all_children()
+    assert ready_seen == [True, False, False, False, False]
+    assert supervisor.is_alive is False
+
+
+@pytest.mark.asyncio
+async def test_fault_tolerant_ready_with_one_serving_rank(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """One healthy rank keeps /ready up; removed ranks are not probed."""
+    supervisor = DPSupervisor(_make_unit_args(enable_fault_tolerance=True))
+    supervisor.child_ports = [8000, 8001, 8002]
+    supervisor._removed_ports = {8002}
+    probed: list[int] = []
+
+    async def fake_child_state(_session, port) -> str:
+        probed.append(port)
+        return "ok" if port == 8000 else "gone"
+
+    monkeypatch.setattr(supervisor, "_child_state", fake_child_state)
+
+    assert await supervisor._probe_fault_tolerant(None) is True
+    assert sorted(probed) == [8000, 8001]
+    assert supervisor.is_ready is True
+
+
+@pytest.mark.asyncio
+async def test_child_state_maps_sentinel_status():
+    supervisor = DPSupervisor(_make_unit_args())
+
+    def session_for(payload):
+        class FakeResponse:
+            status = 200
+
+            async def json(self):
+                return payload
+
+            async def __aenter__(self):
+                return self
+
+            async def __aexit__(self, *_exc):
+                return False
+
+        return SimpleNamespace(get=lambda *_args, **_kwargs: FakeResponse())
+
+    for status, expected in [
+        ("healthy", "ok"),
+        ("unhealthy", "recovering"),
+        ("dead", "gone"),
+    ]:
+        session = session_for({"engines": [{"status": status}]})
+        assert await supervisor._child_state(session, 8000) == expected
+    assert await supervisor._child_state(session_for({}), 8000) == "gone"
+
+
 @pytest.mark.asyncio
 async def test_shutdown_if_supervisor_server_error_on_startup(
     monkeypatch: pytest.MonkeyPatch,
@@ -339,8 +448,10 @@ def __init__(self, _config):
         async def serve(self):
             raise ValueError("supervisor boom")
 
+    alive_seen: list[bool] = []
+
     async def fake_shutdown_children(self):
-        return None
+        alive_seen.append(self.is_alive)
 
     def fake_start_children(self):
         return None
@@ -361,6 +472,7 @@ async def fake_monitor_children(self):
 
     with pytest.raises(ValueError, match="supervisor boom"):
         await supervisor.run()
+    assert alive_seen == [False]
 
 
 # ---------------------------------------------------------------------------

vllm/entrypoints/launchers/dp_supervisor.py

@@ -228,7 +228,7 @@ def _status_response(ok: bool) -> Response:
 
     @app.get("/health", include_in_schema=False)
     async def health() -> Response:
-        return _status_response(app.state.supervisor.is_ready)
+        return _status_response(app.state.supervisor.is_alive)
 
     @app.get("/ready", include_in_schema=False)
     @app.get("/readyz", include_in_schema=False)
@@ -277,6 +277,12 @@ def __init__(self, args: argparse.Namespace):
             for local_rank in range(args.data_parallel_size_local)
         ]
         self._is_ready = False
+        self._was_ready = False
+        self._fault_tolerant = getattr(args, "enable_fault_tolerance", False) and not (
+            envs.VLLM_USE_RUST_FRONTEND and envs.VLLM_RUST_FRONTEND_PATH
+        )
+        self._removed_ports: set[int] = set()
+        self._probe_failures = 0
         self._processes: list[BaseProcess] = []
         self._shutdown_event = asyncio.Event()
         self._shutdown_signal = signal.SIGTERM
@@ -285,6 +291,48 @@ def __init__(self, args: argparse.Namespace):
     def is_ready(self) -> bool:
         return self._is_ready and not self._shutdown_event.is_set()
 
+    @property
+    def is_alive(self) -> bool:
+        return not self._shutdown_event.is_set()
+
+    async def _child_state(self, session: aiohttp.ClientSession, port: int) -> str:
+        """The rank's fault-tolerance state: "ok", "recovering" or "gone"."""
+        probe_ssl = False if self.args.ssl_keyfile and self.args.ssl_certfile else None
+        try:
+            async with session.get(
+                _child_base_url(self.args, port) + "/fault_tolerance/status",
+                ssl=probe_ssl,
+            ) as response:
+                if response.status != HTTPStatus.OK:
+                    return "gone"
+                engines = (await response.json()).get("engines", [])
+        except (aiohttp.ClientError, asyncio.TimeoutError, ValueError):
+            return "gone"
+        statuses = {engine.get("status") for engine in engines}
+        if "healthy" in statuses:
+            return "ok"
+        return "recovering" if "unhealthy" in statuses else "gone"
+
+    async def _probe_fault_tolerant(self, session: aiohttp.ClientSession) -> bool:
+        """Poll the ranks' sentinels; False once no rank is serving or recovering."""
+        states = await asyncio.gather(
+            *(
+                self._child_state(session, port)
+                for port in self.child_ports
+                if port not in self._removed_ports
+            )
+        )
+        self._is_ready = "ok" in states
+        if self._is_ready or "recovering" in states:
+            self._probe_failures = 0
+            return True
+        self._probe_failures += 1
+        if self._probe_failures < self.args.dp_supervisor_probe_failure_threshold:
+            return True
+        logger.info("DPSupervisor found no serving DP Servers.")
+        self._shutdown_event.set()
+        return False
+
     async def run(self) -> None:
         loop = asyncio.get_running_loop()
         decorate_logs("DPSupervisor")
@@ -311,6 +359,7 @@ async def run(self) -> None:
 
             await monitor_task
         finally:
+            self._shutdown_event.set()
             self._is_ready = False
             await self._shutdown_children()
 
@@ -425,43 +474,48 @@ async def _probe_all_children(self) -> None:
         timeout = aiohttp.ClientTimeout(total=self.args.dp_supervisor_probe_timeout_s)
         async with aiohttp.ClientSession(timeout=timeout) as session:
             while not self._shutdown_event.is_set():
-                threshold = (
-                    self.args.dp_supervisor_probe_failure_threshold
-                    if self._is_ready
-                    else 1
-                )
-                results = await asyncio.gather(
-                    *(
-                        _probe_endpoint(
-                            session,
-                            self.args,
-                            port,
-                            "/health",
-                            conn_err_failure_threshold=threshold,
-                            conn_err_retry_delay=self.args.dp_supervisor_probe_interval_s,
-                        )
-                        for port in self.child_ports
-                    ),
-                    return_exceptions=True,
-                )
-                all_healthy = all(r is True for r in results)
-
-                if all_healthy:
-                    # If all healthy, we are ready to receive requests.
-                    # This conditional avoids a potential race condition
-                    # where shutdown is set, THEN the probe returns true.
-                    if not self._shutdown_event.is_set():
-                        self._is_ready = True
-                elif self._is_ready:
-                    # Once ready, any failure in the probe means vLLM is dead.
-                    num_unhealthy = sum(1 for r in results if r is not True)
-                    logger.info(
-                        "DPSupervisor probe found %s unhealthy DP Servers.",
-                        num_unhealthy,
+                if self._fault_tolerant and self._was_ready:
+                    if not await self._probe_fault_tolerant(session):
+                        return
+                else:
+                    threshold = (
+                        self.args.dp_supervisor_probe_failure_threshold
+                        if self._is_ready
+                        else 1
                     )
-                    self._is_ready = False
-                    self._shutdown_event.set()
-                    return
+                    results = await asyncio.gather(
+                        *(
+                            _probe_endpoint(
+                                session,
+                                self.args,
+                                port,
+                                "/health",
+                                conn_err_failure_threshold=threshold,
+                                conn_err_retry_delay=self.args.dp_supervisor_probe_interval_s,
+                            )
+                            for port in self.child_ports
+                        ),
+                        return_exceptions=True,
+                    )
+                    all_healthy = all(r is True for r in results)
+
+                    if all_healthy:
+                        # If all healthy, we are ready to receive requests.
+                        # This conditional avoids a potential race condition
+                        # where shutdown is set, THEN the probe returns true.
+                        if not self._shutdown_event.is_set():
+                            self._is_ready = True
+                            self._was_ready = True
+                    elif self._is_ready:
+                        # Once ready, any failure in the probe means vLLM is dead.
+                        num_unhealthy = sum(1 for r in results if r is not True)
+                        logger.info(
+                            "DPSupervisor probe found %s unhealthy DP Servers.",
+                            num_unhealthy,
+                        )
+                        self._is_ready = False
+                        self._shutdown_event.set()
+                        return
 
                 with contextlib.suppress(asyncio.TimeoutError):
                     logger.debug(
@@ -492,10 +546,22 @@ async def _monitor_children(self) -> None:
         try:
             while not self._shutdown_event.is_set():
                 # 1. Check for dead processes
-                n_failed = len([p for p in self._processes if not p.is_alive()])
-                if n_failed > 0:
-                    logger.info("DPSupervisor found %s exited DP Servers.", n_failed)
+                dead = [p for p in self._processes if not p.is_alive()]
+                if dead and not (
+                    self._fault_tolerant
+                    and self._was_ready
+                    and len(dead) < len(self._processes)
+                ):
+                    logger.info("DPSupervisor found %s exited DP Servers.", len(dead))
                     break
+                for process in dead:
+                    # An exited rank with live peers was removed by the engine.
+                    port = self.child_ports[self._processes.index(process)]
+                    if port not in self._removed_ports:
+                        self._removed_ports.add(port)
+                        logger.info(
+                            "DPSupervisor keeps serving; %s exited.", process.name
+                        )
 
                 # 2. Check if the probe background task crashed or failed.
                 if probe_task.done():