Improve engine health monitoring (#4645)

* Improve engine health monitoring and wakeup scheduling

* remove default value from health_probe()

* revert to sync wakeup

Author: lvhan028

lmdeploy/pytorch/engine/engine_loop.py

@@ -446,7 +446,6 @@ async def __no_running_warning():
                 running=next_running,
                 forward_inputs=forward_inputs,
             )
-            scheduler.tick()
             self.inputs_maker.deactivate_evict_seqs()
             has_runable_event.set()
 

lmdeploy/pytorch/engine/inputs_maker.py

@@ -796,6 +796,7 @@ async def _send_next_inputs_impl(self, prefill: bool = None, enable_empty: bool
             session_ids = [seq.session_id for seq in next_running]
             logger.debug(f'Forward session_ids: {session_ids}')
         await self.executor.forward_async(forward_inputs)
+        self.scheduler.tick()
         self.forward_inputs = forward_inputs
         return forward_inputs, next_running
 

lmdeploy/pytorch/paging/scheduler.py

@@ -67,7 +67,7 @@ def __init__(
         self.scheduler_tick = 0
 
     def tick(self):
-        """Mark one scheduler progress step."""
+        """Mark one scheduler progress step (once per forward dispatch)."""
         self.scheduler_tick += 1
 
     @staticmethod

lmdeploy/serve/core/async_engine.py

@@ -280,7 +280,7 @@ def _validate_scheduler_progress(self, metrics, scheduler_stall_timeout: float)
     def _make_health_result(status: str, message: str) -> dict:
         return dict(status=status, message=message)
 
-    async def health_probe(self, timeout: float = 2.0, scheduler_stall_timeout: float = 15.0) -> dict:
+    async def health_probe(self, timeout: float, scheduler_stall_timeout: float) -> dict:
         """Probe backend health with a bounded, non-overlapping call."""
         if self.is_sleeping:
             return self._make_health_result(
@@ -291,7 +291,7 @@ async def health_probe(self, timeout: float = 2.0, scheduler_stall_timeout: floa
         if self._health_probe_task is not None:
             if not self._health_probe_task.done():
                 return self._make_health_result(
-                    status='unhealthy',
+                    status='pending',
                     message='Previous backend health probe is still pending.',
                 )
             try:

lmdeploy/serve/core/health.py

@@ -2,25 +2,73 @@
 from __future__ import annotations
 
 import asyncio
+import os
 import time
 from typing import TYPE_CHECKING
 
+from lmdeploy.utils import get_logger
+
 if TYPE_CHECKING:
     from .async_engine import AsyncEngine
 
+logger = get_logger('lmdeploy')
+
+HEALTH_POLL_INTERVAL = 'LMDEPLOY_HEALTH_POLL_INTERVAL'
+HEALTH_PROBE_TIMEOUT = 'LMDEPLOY_HEALTH_PROBE_TIMEOUT'
+HEALTH_UNHEALTHY_AFTER = 'LMDEPLOY_HEALTH_UNHEALTHY_AFTER'
+
+DEFAULT_PROBE_TIMEOUT = 10.0
+DEFAULT_POLL_INTERVAL = 12.0
+DEFAULT_UNHEALTHY_AFTER = 90.0
+
+
+def _env_override_float(env_var: str, value: float) -> float:
+    """Return ``value`` unless ``env_var`` is set, then parse and return it."""
+    env_value = os.getenv(env_var)
+    if env_value is None:
+        return value
+    try:
+        return float(env_value)
+    except ValueError:
+        return value
+
 
 class EngineHealthMonitor:
     """Background engine health monitor."""
 
     def __init__(self,
                  async_engine: AsyncEngine | None,
-                 poll_interval: float = 5.0,
-                 probe_timeout: float = 2.0,
-                 unhealthy_after: float = 15.0):
+                 poll_interval: float = DEFAULT_POLL_INTERVAL,
+                 probe_timeout: float = DEFAULT_PROBE_TIMEOUT,
+                 unhealthy_after: float = DEFAULT_UNHEALTHY_AFTER):
+        """Initialize the background health monitor.
+
+        Args:
+            async_engine: Engine instance to probe; ``None`` marks the service
+                unhealthy until an engine is attached.
+            poll_interval: Seconds between consecutive ``probe_once()`` calls in
+                the background loop (default 12.0). Should be greater than
+                ``probe_timeout`` to reduce overlapping probes. Overridden by
+                ``LMDEPLOY_HEALTH_POLL_INTERVAL`` when that variable is set.
+            probe_timeout: Maximum seconds to wait for a single
+                ``health_probe()`` (backend ``get_health_status()``) before
+                reporting that probe as unhealthy (default 10.0). Overridden by
+                ``LMDEPLOY_HEALTH_PROBE_TIMEOUT`` when that variable is set.
+            unhealthy_after: Seconds without a successful probe or scheduler
+                progress before the service is considered unhealthy (default
+                90.0). Overridden by ``LMDEPLOY_HEALTH_UNHEALTHY_AFTER`` when
+                that variable is set. Passed to ``health_probe()`` as
+                ``scheduler_stall_timeout``, and also used in ``snapshot()``
+                to expire stale healthy status or stuck ``initializing`` state.
+        """
         self.async_engine = async_engine
-        self.poll_interval = poll_interval
-        self.probe_timeout = probe_timeout
-        self.unhealthy_after = unhealthy_after
+        self.poll_interval = _env_override_float(HEALTH_POLL_INTERVAL, poll_interval)
+        self.probe_timeout = _env_override_float(HEALTH_PROBE_TIMEOUT, probe_timeout)
+        self.unhealthy_after = _env_override_float(HEALTH_UNHEALTHY_AFTER, unhealthy_after)
+        if self.poll_interval <= self.probe_timeout:
+            logger.warning('Engine health poll_interval (%.1fs) should be greater than probe_timeout (%.1fs) '
+                             'to avoid overlapping probes.',
+                             self.poll_interval, self.probe_timeout)
         self._task: asyncio.Task | None = None
         self._started_time = time.monotonic()
         self._last_success_time: float | None = None
@@ -61,14 +109,17 @@ async def probe_once(self):
                               message=f'Engine health probe failed: {e}')
 
         status = result['status']
+        if status == 'pending':
+            logger.info('Engine health probe skipped: previous backend health probe is still pending.')
+            return
         if status in ('healthy', 'sleeping'):
             self._last_success_time = probe_time
         self._snapshot = dict(status=status, message=result['message'])
 
     def snapshot(self) -> dict:
         snapshot = dict(self._snapshot)
         now = time.monotonic()
-        if snapshot['status'] in ('healthy', 'sleeping') and self._last_success_time is not None:
+        if snapshot['status'] == 'healthy' and self._last_success_time is not None:
             if now - self._last_success_time > self.unhealthy_after:
                 snapshot['status'] = 'unhealthy'
                 snapshot['message'] = f'No successful health probe for {now - self._last_success_time:.1f}s.'