fix(mlnode): scheduler-heartbeat liveness for vLLM runner + damped proxy health probe

mlnode/packages/api/src/api/inference/vllm/runner.py

@@ -1,4 +1,5 @@
 import os
+import re
 import subprocess
 import time
 import requests
@@ -23,6 +24,87 @@
 
 logger = create_logger(__name__)
 
+HANG_CONSECUTIVE = 3
+
+# (metric suffix, cross-series reduce). Sum/max across ALL series: /metrics
+# can briefly expose several registries mid-restart.
+_HEARTBEAT_SERIES = (
+    ("iteration_tokens_total_count", sum),
+    ("num_requests_running", max),
+    ("num_requests_waiting", max),
+)
+
+
+def _int_env(name: str, default: int) -> int:
+    """os.environ[name] as int; unset/empty -> default, garbage -> default + warning."""
+    raw = os.getenv(name)
+    if raw is None or not raw.strip():
+        return default
+    try:
+        return int(raw)
+    except ValueError:
+        logger.warning("%s=%r is not an integer -- using %d", name, raw, default)
+        return default
+
+
+def _scrape_heartbeat(host: str, port: int):
+    """(iters, running, waiting) floats from /metrics; None per unreadable value.
+    Raises requests.ConnectionError when the instance is unreachable (dead)."""
+    try:
+        txt = requests.get(f"http://{host}:{port}/metrics", timeout=3).text
+    except requests.exceptions.ConnectTimeout:
+        # Saturation, not death: the caller's grace window decides.
+        return None, None, None
+    except requests.ConnectionError:
+        raise
+    except Exception:
+        return None, None, None
+    values = []
+    for name, reduce_fn in _HEARTBEAT_SERIES:
+        try:
+            series = re.findall(rf"(?m)^vllm:{name}\S*\s+([0-9eE.+-]+)", txt)
+            values.append(reduce_fn(float(x) for x in series) if series else None)
+        except Exception:
+            values.append(None)
+    return tuple(values)
+
+
+def _heartbeat_verdict(port: int, st: dict, iters, running, waiting, now: float, grace: int) -> bool:
+    """Advance one port's heartbeat state (mutates st); True = instance alive.
+    Branch order is semantic -- reorder only with the unit tests in hand."""
+    read_ok = iters is not None or running is not None or waiting is not None
+    has_work = (running is not None and running > 0) or (waiting is not None and waiting > 0)
+    # Fresh baseline on either: legitimate idle (readable scrape, no work) or a
+    # stepping engine -- ANY counter change counts, including the reset after an
+    # engine restart (a ">" test would false-HUNG on the stale high baseline).
+    # A blind scrape (all None) is neither: refreshing on it would silently
+    # disable hang detection.
+    if (read_ok and not has_work) or (
+        iters is not None and (st["iter"] is None or iters != st["iter"])
+    ):
+        st["ts"] = now
+        st["iter"] = iters
+        st["hung"] = 0
+        return True
+    if iters is None:
+        # Counter unreadable: alive only within grace.
+        return now - st["ts"] <= grace
+    if now - st["ts"] > grace:
+        # Frozen with work past grace: require consecutive verdicts so one
+        # confused scrape cannot start the restart escalation.
+        st["hung"] += 1
+        if st["hung"] >= HANG_CONSECUTIVE:
+            logger.error(
+                "vLLM instance on port %s: scheduler heartbeat frozen >%ss "
+                "over %d consecutive checks with running=%s waiting=%s -- "
+                "reporting unhealthy for restart",
+                port, grace, st["hung"], running, waiting,
+            )
+            return False
+        return True
+    st["hung"] = 0
+    return True
+
 
 class IVLLMRunner(ITrackableTask):
     @abstractmethod
@@ -63,6 +145,7 @@ def __init__(
         self.dtype = dtype
         self.additional_args = additional_args or []
         self.processes: List[subprocess.Popen] = []
+        self._hb = {}  # per-port heartbeat state; outlives engine restarts
 
     def _get_arg_value(self, name: str, default: int = 1) -> int:
         if name in self.additional_args:
@@ -224,15 +307,28 @@ def is_running(self) -> bool:
     def is_available(self) -> bool:
         if not self.is_running():
             return False
-        try:
-            # Check if any backend is available
-            for port in range(self.VLLM_PORT + 1, self.VLLM_PORT + len(self.processes) + 1):
-                resp = requests.get(f"http://{self.VLLM_HOST}:{port}/health", timeout=2)
-                if resp.status_code == 200:
-                    return True
-            return False
-        except (requests.ConnectionError, requests.Timeout):
-            return False
+        # Scheduler-heartbeat liveness. vLLM /health returns 200 even when the
+        # scheduler is deadlocked (check_health only reads the `errored` flag)
+        # and misses its deadline while merely busy. The iteration counter
+        # advances on every engine step and freezes only on a real stall; the
+        # grace window also absorbs the counter's observed pre-stall flatline.
+        grace = _int_env("MLNODE_HANG_GRACE_SEC", 120)
+        now = time.time()
+        healthy = True
+        # Check every instance: a hang must not hide behind a healthy sibling.
+        for port in range(self.VLLM_PORT + 1, self.VLLM_PORT + len(self.processes) + 1):
+            try:
+                iters, running, waiting = _scrape_heartbeat(self.VLLM_HOST, port)
+            except requests.ConnectionError:
+                # Refused/unreachable -> this instance is dead.
+                return False
+            st = self._hb.setdefault(port, {"ts": now, "iter": iters, "hung": 0})
+            if grace <= 0:
+                # Detection disabled: only process death / refusal mark unhealthy.
+                continue
+            ok = _heartbeat_verdict(port, st, iters, running, waiting, now, grace)
+            healthy = healthy and ok
+        return healthy
 
     def get_error_if_exist(self) -> Optional[str]:
         for p in self.processes:

mlnode/packages/api/src/api/proxy.py

@@ -22,6 +22,9 @@
 
 vllm_backend_ports: List[int] = []
 vllm_healthy: Dict[int, bool] = {}
+# /health prober damping -- contract in _apply_health_damping.
+HEALTH_FAIL_STREAK: int = 3
+vllm_health_fail_streak: Dict[int, int] = {}
 vllm_counts: Dict[int, int] = {}
 poc_status_by_port: Dict[int, str] = {}  # PoC status: "IDLE", "GENERATING", "STOPPED", or ""
 pow_generate_rr_index: int = 0
@@ -189,6 +192,19 @@ async def call_backend(port: int, method: str, path: str, json_body: dict = None
         raise ValueError(f"Unsupported method: {method}")
 
 
+def _apply_health_damping(port: int, probe_ok: bool) -> bool:
+    """Damped health verdict: unhealthy only after HEALTH_FAIL_STREAK
+    consecutive probe failures, instant recovery on success; a backend that
+    has never succeeded (startup) stays down regardless of streak."""
+    if probe_ok:
+        vllm_health_fail_streak[port] = 0
+        return True
+    vllm_health_fail_streak[port] = vllm_health_fail_streak.get(port, 0) + 1
+    if vllm_healthy.get(port) and vllm_health_fail_streak[port] < HEALTH_FAIL_STREAK:
+        return True
+    return False
+
+
 async def _health_check_vllm(interval: float = 2.0):
     """Health check for vLLM backends and manage compatibility server."""
     logger.info("Health check loop started, checking every %s seconds", interval)
@@ -205,12 +221,14 @@ async def _health_check_vllm(interval: float = 2.0):
                 if vllm_client is None:
                     logger.warning(f"Health check skipped for port {p}: vllm_client is None")
                     continue
-                r = await vllm_client.get(f"http://{VLLM_HOST}:{p}/health", timeout=2)
+                r = await vllm_client.get(f"http://{VLLM_HOST}:{p}/health", timeout=10)
                 ok = r.status_code == 200
                 logger.debug("Health check for port %d: status=%d, ok=%s", p, r.status_code, ok)
             except Exception as e:
                 logger.debug("Health check for port %d failed: %s", p, e)
 
+            ok = _apply_health_damping(p, ok)
+
             prev = vllm_healthy.get(p)
             if prev != ok:
                 logger.info("%s:%d is %s", VLLM_HOST, p, "UP" if ok else "DOWN")

mlnode/packages/api/tests/unit/test_proxy.py

@@ -442,4 +442,61 @@ async def test_resource_cleanup_verification():
 
     # Verify backend counts reset
     for port in [5001, 5002]:
-        assert proxy_module.vllm_counts[port] == 0 
+        assert proxy_module.vllm_counts[port] == 0 
+
+# ---------------------------------------------------------------------------
+# Health-probe damping (_apply_health_damping)
+# ---------------------------------------------------------------------------
+
+@pytest.fixture
+def clean_health_state():
+    """Reset the module-level health dicts around each damping test."""
+    saved_healthy = dict(proxy_module.vllm_healthy)
+    saved_streak = dict(proxy_module.vllm_health_fail_streak)
+    proxy_module.vllm_healthy.clear()
+    proxy_module.vllm_health_fail_streak.clear()
+    yield
+    proxy_module.vllm_healthy.clear()
+    proxy_module.vllm_healthy.update(saved_healthy)
+    proxy_module.vllm_health_fail_streak.clear()
+    proxy_module.vllm_health_fail_streak.update(saved_streak)
+
+
+def _drive(port, probes, start_healthy):
+    """Feed a probe sequence through the damping and record verdicts, updating
+    vllm_healthy the same way _health_check_vllm does."""
+    proxy_module.vllm_healthy[port] = start_healthy
+    out = []
+    for probe in probes:
+        ok = proxy_module._apply_health_damping(port, probe)
+        proxy_module.vllm_healthy[port] = ok
+        out.append(ok)
+    return out
+
+
+def test_damping_startup_stays_down_until_first_success(clean_health_state):
+    assert _drive(5001, [False, False, True], start_healthy=False) == [False, False, True]
+
+
+def test_damping_two_blips_still_healthy(clean_health_state):
+    """Missed /health deadlines on a busy-but-alive engine (up to
+    HEALTH_FAIL_STREAK-1 in a row) must not stop the compatibility server."""
+    assert _drive(5001, [True, False, False, True], start_healthy=True) == [True, True, True, True]
+
+
+def test_damping_three_consecutive_failures_is_down(clean_health_state):
+    assert _drive(5001, [True, False, False, False], start_healthy=True) == [True, True, True, False]
+
+
+def test_damping_recovers_instantly_on_success(clean_health_state):
+    assert _drive(5001, [False, False, False, True], start_healthy=True) == [True, True, False, True]
+
+
+def test_damping_ports_are_independent(clean_health_state):
+    proxy_module.vllm_healthy.update({5001: True, 5002: True})
+    # 5001 fails 3x -> down; 5002 keeps succeeding -> up
+    for _ in range(3):
+        proxy_module.vllm_healthy[5001] = proxy_module._apply_health_damping(5001, False)
+        proxy_module.vllm_healthy[5002] = proxy_module._apply_health_damping(5002, True)
+    assert proxy_module.vllm_healthy[5001] is False
+    assert proxy_module.vllm_healthy[5002] is True