diff --git a/bioengine/apps/proxy_deployment.py b/bioengine/apps/proxy_deployment.py index 7935f79f..0ca13266 100644 --- a/bioengine/apps/proxy_deployment.py +++ b/bioengine/apps/proxy_deployment.py @@ -6,7 +6,7 @@ import os import time import uuid -from typing import Any, Callable, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional, Tuple import ray from pydantic import Field @@ -74,6 +74,20 @@ class _PermanentRegistrationError(RuntimeError): """Registration failed for a reason no amount of retrying will fix.""" +# Serve's own verdict on a deployment, for deployments sitting at zero replicas. +# Zero replicas is the correct idle state under ``min_replicas: 0``, so a replica +# count alone cannot tell a crash from a deliberate scale-down. UPSCALING matters +# as much as HEALTHY: it is the state a scaled-to-zero deployment enters when a +# request wakes it, and deregistering there would remove the service during the +# very wake-up it is meant to allow. +_SERVICEABLE_AT_ZERO_REPLICAS = ("HEALTHY", "UPSCALING", "DOWNSCALING") + + +def _is_serviceable(running: int, status: str) -> bool: + """Whether a sibling deployment can serve a request, now or after an upscale.""" + return running > 0 or status in _SERVICEABLE_AT_ZERO_REPLICAS + + # ``ray_actor_options`` is intentionally minimal here. The proxy needs a # ``runtime_env.pip`` list (aiortc, httpx, hypha-rpc, pydantic + their pins # from ``[worker]``), but computing it requires ``importlib.metadata`` to @@ -1285,10 +1299,10 @@ def _ensure_maintenance_task(self) -> None: # ===== Ray Serve Health Check ===== # Implements periodic health checks for Ray Serve. - async def _sibling_running_counts(self) -> Optional[Dict[str, int]]: - """RUNNING replica count of every sibling deployment (all deployments in - this app but the proxy itself), read out-of-band from the Serve - controller. + async def _sibling_states(self) -> Optional[Dict[str, Tuple[int, str]]]: + """``(RUNNING replica count, deployment status)`` for every sibling + deployment (all deployments in this app but the proxy itself), read + out-of-band from the Serve controller. The controller already health-checks every replica; this reads that collected state and never issues an in-band request, so a saturated app @@ -1296,6 +1310,9 @@ async def _sibling_running_counts(self) -> Optional[Dict[str, int]]: Returns ``None`` when the status can't be determined (controller mid-restart, app not yet in the view) — the caller treats ``None`` as "unknown" and never deregisters on it. + + The status is carried because the replica count alone cannot separate a + deployment idling at ``min_replicas: 0`` from one that crashed to zero. """ from ray import serve as _serve @@ -1305,22 +1322,24 @@ async def _sibling_running_counts(self) -> Optional[Dict[str, int]]: except Exception: self._own_deployment_name = None - def _query() -> Optional[Dict[str, int]]: + def _query() -> Optional[Dict[str, Tuple[int, str]]]: app = _serve.status().applications.get(self.application_id) if app is None: return None - counts: Dict[str, int] = {} + states: Dict[str, Tuple[int, str]] = {} for name, deployment in app.deployments.items(): if name == self._own_deployment_name: continue - counts[name] = sum( + running = sum( count for state, count in deployment.replica_states.items() if str(getattr(state, "value", state)) == "RUNNING" ) - return counts + status = str(getattr(deployment.status, "value", deployment.status)) + states[name] = (running, status) + return states - def _read() -> Optional[Dict[str, int]]: + def _read() -> Optional[Dict[str, Tuple[int, str]]]: try: return _query() except Exception: @@ -1380,20 +1399,20 @@ async def check_health(self): # against any deployment's ``max_ongoing_requests``. Ray's own controller # already health-checks each replica; a sibling that crashes or whose # ``health_check`` raises drops out of the RUNNING count here. - counts = await self._sibling_running_counts() + states = await self._sibling_states() if not self.entry_deployment_ready: - if counts and all(running > 0 for running in counts.values()): + if states and all(_is_serviceable(*s) for s in states.values()): self.entry_deployment_ready = True - for dep in counts: + for dep in states: self._dep_seen_ready[dep] = True logger.info( - f"✅ All deployments of app '{self.application_id}' are RUNNING." + f"✅ All deployments of app '{self.application_id}' are serviceable." ) else: pending = ( - [dep for dep, running in counts.items() if running == 0] - if counts + [dep for dep, s in states.items() if not _is_serviceable(*s)] + if states else "unknown" ) logger.info( @@ -1401,19 +1420,23 @@ async def check_health(self): f"(pending: {pending})." ) return - elif counts is not None: - # A sibling that came up and then dropped to zero means the app can - # no longer serve: deregister so the service disappears from Hypha, - # and re-gate. The outage stays visible in the app status via the - # down deployment itself, so the proxy need not fail its own health. - for dep, running in counts.items(): - if running > 0: + elif states is not None: + # A sibling that came up and then stopped being serviceable means the + # app can no longer serve: deregister so the service disappears from + # Hypha, and re-gate. The outage stays visible in the app status via + # the down deployment itself, so the proxy need not fail its own + # health. Serviceability is not a replica count: a deployment idling + # at ``min_replicas: 0`` is at zero replicas on purpose, and + # deregistering it would be unrecoverable — waking it needs a request, + # and a request needs the registration this would remove. + for dep, (running, status) in states.items(): + if _is_serviceable(running, status): self._dep_seen_ready[dep] = True - elif running == 0 and self._dep_seen_ready.get(dep): + elif self._dep_seen_ready.get(dep): logger.error( f"❌ Deployment '{dep}' of app '{self.application_id}' has " - f"no RUNNING replica. Deregistering Hypha service until it " - f"recovers." + f"no RUNNING replica and is not serviceable (status: " + f"{status}). Deregistering Hypha service until it recovers." ) await self._deregister_services() return diff --git a/tests/apps/test_proxy_entry_saturation_tolerance.py b/tests/apps/test_proxy_entry_saturation_tolerance.py index be234644..d43f9ea5 100644 --- a/tests/apps/test_proxy_entry_saturation_tolerance.py +++ b/tests/apps/test_proxy_entry_saturation_tolerance.py @@ -4,19 +4,24 @@ that call was routed through the entry's router and counted against its ``max_ongoing_requests``, so a saturated-but-healthy entry head-of-line-blocked the probe. Instead the proxy reads every sibling deployment's RUNNING replica -count out-of-band from the Serve controller (``serve.status()``), which a -saturated app can neither block nor fail. +count *and Serve's own status for it* out-of-band from the Serve controller +(``serve.status()``), which a saturated app can neither block nor fail. The gate: -* registers only once every sibling deployment has a RUNNING replica; +* registers only once every sibling deployment is serviceable; * tolerates an "unknown" reading (``None``) — a controller mid-restart never deregisters a healthy app; -* deregisters when a sibling that was seen RUNNING drops to zero (the outage - stays visible in the app status via the down deployment itself, so the proxy - does not fail its own health); +* deregisters when a sibling that was seen serviceable stops being so (the + outage stays visible in the app status via the down deployment itself, so the + proxy does not fail its own health); * saturation is a non-event: a busy replica is still ``RUNNING`` in ``serve.status()``, so the count is unaffected. + +**Serviceable is not "has a replica".** Under ``min_replicas: 0`` zero replicas +is the correct idle state, so the replica count alone cannot separate a crash +from a deliberate scale-down — Serve's status can, and the tests below pin both +directions. """ from __future__ import annotations @@ -80,11 +85,13 @@ def _bare_proxy(**attrs): return inst -def _stub_counts(value): - async def _counts(self=None): +def _stub_states(value): + """``value`` is ``None`` or ``{deployment: (running, status)}``.""" + + async def _states(self=None): return value - return _counts + return _states async def _run(inst): @@ -100,7 +107,7 @@ async def test_saturation_does_not_deregister() -> None: server=_Server(), websocket_service_id="ws", ) - inst._sibling_running_counts = _stub_counts({"EntryDeployment": 1}) + inst._sibling_states = _stub_states({"EntryDeployment": (1, "HEALTHY")}) await _run(inst) assert inst._deregister_services.called is False @@ -114,28 +121,28 @@ async def test_unknown_status_does_not_deregister() -> None: server=_Server(), websocket_service_id="ws", ) - inst._sibling_running_counts = _stub_counts(None) + inst._sibling_states = _stub_states(None) await _run(inst) assert inst._deregister_services.called is False @pytest.mark.asyncio -async def test_sibling_drop_deregisters() -> None: - """A sibling seen RUNNING that drops to zero deregisters the service.""" +async def test_sibling_crash_deregisters() -> None: + """A sibling that was up and is now UNHEALTHY at zero deregisters.""" inst = _bare_proxy( entry_deployment_ready=True, _dep_seen_ready={"RuntimeDeployment": True}, server=_Server(), websocket_service_id="ws", ) - inst._sibling_running_counts = _stub_counts({"RuntimeDeployment": 0}) + inst._sibling_states = _stub_states({"RuntimeDeployment": (0, "UNHEALTHY")}) await _run(inst) assert inst._deregister_services.called is True @pytest.mark.asyncio -async def test_initial_gate_waits_for_all_running() -> None: - """The service registers only once every sibling is RUNNING. +async def test_initial_gate_waits_for_all_siblings() -> None: + """The service registers only once every sibling is serviceable. check_health owns the gate; the background maintenance task does the registering, and must stay idle until the gate opens. @@ -144,16 +151,16 @@ async def test_initial_gate_waits_for_all_running() -> None: inst = _bare_proxy() inst._register_services = registered - inst._sibling_running_counts = _stub_counts( - {"EntryDeployment": 1, "RuntimeDeployment": 0} + inst._sibling_states = _stub_states( + {"EntryDeployment": (1, "HEALTHY"), "RuntimeDeployment": (0, "UPDATING")} ) await _run(inst) assert inst.entry_deployment_ready is False await inst._maintenance_tick() assert registered.called is False - inst._sibling_running_counts = _stub_counts( - {"EntryDeployment": 1, "RuntimeDeployment": 1} + inst._sibling_states = _stub_states( + {"EntryDeployment": (1, "HEALTHY"), "RuntimeDeployment": (1, "HEALTHY")} ) await _run(inst) assert inst.entry_deployment_ready is True @@ -161,6 +168,101 @@ async def test_initial_gate_waits_for_all_running() -> None: assert registered.called is True +# ===== svamp #0059: min_replicas 0 must not be mistaken for an outage ===== + + +@pytest.mark.asyncio +async def test_scaled_to_zero_does_not_deregister() -> None: + """The #0059 regression: idling at ``min_replicas: 0`` is not an outage. + + Deregistering here was unrecoverable by construction — waking the + deployment needs a request, and a request needs the registration that the + deregistration removes. + """ + inst = _bare_proxy( + entry_deployment_ready=True, + _dep_seen_ready={"EntryDeployment": True}, + server=_Server(), + websocket_service_id="ws", + ) + inst._sibling_states = _stub_states({"EntryDeployment": (0, "HEALTHY")}) + await _run(inst) + assert inst._deregister_services.called is False + + +@pytest.mark.asyncio +async def test_wake_up_window_does_not_deregister() -> None: + """A request arriving at a scaled-to-zero deployment must not lose the service. + + UPSCALING with zero replicas is exactly the wake-up this fix exists to + allow; deregistering there would remove the service mid-wake. + """ + inst = _bare_proxy( + entry_deployment_ready=True, + _dep_seen_ready={"EntryDeployment": True}, + server=_Server(), + websocket_service_id="ws", + ) + inst._sibling_states = _stub_states({"EntryDeployment": (0, "UPSCALING")}) + await _run(inst) + assert inst._deregister_services.called is False + + +@pytest.mark.asyncio +async def test_downscaling_to_zero_does_not_deregister() -> None: + """The transient between the last replica going and the status settling.""" + inst = _bare_proxy( + entry_deployment_ready=True, + _dep_seen_ready={"EntryDeployment": True}, + server=_Server(), + websocket_service_id="ws", + ) + inst._sibling_states = _stub_states({"EntryDeployment": (0, "DOWNSCALING")}) + await _run(inst) + assert inst._deregister_services.called is False + + +@pytest.mark.asyncio +async def test_cold_start_at_min_replicas_zero_registers() -> None: + """The half the KTH repro could not reach. + + ``initial_replicas`` defaults to ``None``, and the autoscaler's lower bound + is then ``min_replicas`` — so an app *deployed* at ``min_replicas: 0`` + starts with zero replicas and never had a replica to be "seen ready". The + old gate required ``all(running > 0)`` and so never opened at all, leaving + the app permanently unregistered rather than merely unwakeable. + """ + registered = _Recorder() + inst = _bare_proxy() + inst._register_services = registered + + # Still coming up: not serviceable, gate stays shut. + inst._sibling_states = _stub_states({"EntryDeployment": (0, "UPDATING")}) + await _run(inst) + assert inst.entry_deployment_ready is False + await inst._maintenance_tick() + assert registered.called is False + + # Settled at zero replicas because that is what was asked for. + inst._sibling_states = _stub_states({"EntryDeployment": (0, "HEALTHY")}) + await _run(inst) + assert inst.entry_deployment_ready is True + await inst._maintenance_tick() + assert registered.called is True + + +def test_serviceability_separates_idle_from_crashed() -> None: + """The discriminator itself, stated as a table so it cannot drift silently.""" + assert pd_module._is_serviceable(1, "UNHEALTHY") is True, ( + "a deployment with a RUNNING replica can serve regardless of what Serve " + "thinks of it overall" + ) + for status in ("HEALTHY", "UPSCALING", "DOWNSCALING"): + assert pd_module._is_serviceable(0, status) is True, status + for status in ("UNHEALTHY", "UPDATING"): + assert pd_module._is_serviceable(0, status) is False, status + + def test_probe_is_off_the_data_plane() -> None: """check_health must read sibling status out-of-band, never in-band.""" src = inspect.getsource(_ProxyCls.check_health) @@ -168,4 +270,4 @@ def test_probe_is_off_the_data_plane() -> None: "the entry health probe must not be issued in-band — it counts against " "the entry's max_ongoing_requests and blocks under load." ) - assert "_sibling_running_counts" in src + assert "_sibling_states" in src