Refactor Heath Checkers

Author: AryanAhadinia

src/swiss_ai_model_launch/cli/display/live.py

@@ -19,7 +19,6 @@
 }
 
 _MODEL_HEALTH_STYLE: dict[ModelHealth, str] = {
-    ModelHealth.WAITING: "[yellow]WAITING[/yellow]",
     ModelHealth.HEALTHY: "[green]HEALTHY[/green]",
     ModelHealth.ERROR: "[orange]ERROR[/orange]",
     ModelHealth.NOT_DEPLOYED: "[dim]NOT DEPLOYED[/dim]",

src/swiss_ai_model_launch/cli/display/state.py

@@ -11,7 +11,7 @@ def __init__(self) -> None:
         self.partition: str | None = None
         self.job_id: int | None = None
         self.job_status: JobStatus | None = None
-        self.model_health: ModelHealth = ModelHealth.WAITING
+        self.model_health: ModelHealth = ModelHealth.NOT_DEPLOYED
         self.served_model_name: str | None = None
         self.out_logs: deque[str] = deque()
         self.err_logs: deque[str] = deque()

src/swiss_ai_model_launch/cli/healthcheck/checker.py

@@ -7,9 +7,7 @@
 _TIMEOUT_SECONDS = 10
 
 
-async def check_model_health(
-    served_model_name: str, api_key: str, ever_healthy: bool = False
-) -> ModelHealth:
+async def check_model_health(model_name: str, api_key: str) -> ModelHealth:
     try:
         async with httpx.AsyncClient() as client:
             response = await client.post(
@@ -19,14 +17,14 @@ async def check_model_health(
                     "Authorization": f"Bearer {api_key}",
                 },
                 json={
-                    "model": served_model_name,
+                    "model": model_name,
                     "messages": [_MESSAGE],
                     "stream": False,
                 },
                 timeout=_TIMEOUT_SECONDS,
             )
         if response.is_success:
             return ModelHealth.HEALTHY
-        return ModelHealth.NOT_RESPONDING if ever_healthy else ModelHealth.NOT_DEPLOYED
+        return ModelHealth.NOT_RESPONDING
     except (httpx.TransportError, httpx.TimeoutException):
         return ModelHealth.ERROR

src/swiss_ai_model_launch/cli/healthcheck/model_health.py

@@ -2,7 +2,6 @@
 
 
 class ModelHealth(Enum):
-    WAITING = "WAITING"
     HEALTHY = "HEALTHY"
     ERROR = "ERROR"
     NOT_DEPLOYED = "NOT_DEPLOYED"

src/swiss_ai_model_launch/cli/main.py

@@ -509,9 +509,9 @@ async def _monitor() -> None:
             job_status = await launcher.get_job_status(job_id)
             state.update(job_status=job_status)
 
-            model_health = await check_model_health(
-                served, cscs_api_key, ever_healthy=ever_healthy
-            )
+            model_health = await check_model_health(served, cscs_api_key)
+            if model_health == ModelHealth.NOT_RESPONDING and not ever_healthy:
+                model_health = ModelHealth.NOT_DEPLOYED
             ever_healthy = ever_healthy or model_health == ModelHealth.HEALTHY
             state.update(model_health=model_health)
 

tests/integration/test_firecrest_launcher.py

@@ -1,15 +1,13 @@
-import asyncio
 import importlib.resources
 import json
 import os
 
 import firecrest as f7t
 import pytest
 
-from swiss_ai_model_launch.cli.healthcheck import ModelHealth, check_model_health
 from swiss_ai_model_launch.launchers.firecrest_launcher import FirecRESTLauncher
 from swiss_ai_model_launch.launchers.launch_request import LaunchRequest
-from swiss_ai_model_launch.launchers.launcher import JobStatus
+from tests.integration.utils import wait_for_job_running, wait_for_model_healthy
 
 _LAUNCH_TIMEOUT = 60
 _HEALTH_TIMEOUT = 120
@@ -74,42 +72,6 @@ def cscs_api_key(env: dict[str, str]) -> str:
     return env["SML_CSCS_API_KEY"]
 
 
-async def _wait_for_job_running(
-    launcher: FirecRESTLauncher,
-    job_id: int,
-    timeout_min: int,
-    poll_interval_seconds: int = 15,
-) -> None:
-    deadline = asyncio.get_event_loop().time() + timeout_min * 60
-    while asyncio.get_event_loop().time() < deadline:
-        await asyncio.sleep(poll_interval_seconds)
-        status = await launcher.get_job_status(job_id)
-        print(f"[job {job_id}] status: {status.value}")
-        if status == JobStatus.RUNNING:
-            return
-        if status == JobStatus.TIMEOUT:
-            pytest.fail(f"Job {job_id} timed out before becoming RUNNING.")
-    pytest.fail(f"Job {job_id} didn't reach RUNNING within {timeout_min} mins.")
-
-
-async def _wait_for_model_healthy(
-    served_model_name: str,
-    api_key: str,
-    timeout_min: int,
-    poll_interval_seconds: int = 30,
-) -> None:
-    deadline = asyncio.get_event_loop().time() + timeout_min * 60
-    while asyncio.get_event_loop().time() < deadline:
-        await asyncio.sleep(poll_interval_seconds)
-        health = await check_model_health(served_model_name, api_key)
-        print(f"[{served_model_name}] health: {health.value}")
-        if health == ModelHealth.HEALTHY:
-            return
-    pytest.fail(
-        f"Model '{served_model_name}' didn't become HEALTHY within {timeout_min} mins."
-    )
-
-
 @pytest.mark.parametrize("launch_request", _LAUNCH_REQUESTS)  # type: ignore[misc]
 async def test_launch_apertus_and_health(
     launcher: FirecRESTLauncher,
@@ -123,7 +85,7 @@ async def test_launch_apertus_and_health(
     assert served_model_name
 
     try:
-        await _wait_for_job_running(launcher, job_id, _LAUNCH_TIMEOUT)
-        await _wait_for_model_healthy(served_model_name, cscs_api_key, _HEALTH_TIMEOUT)
+        await wait_for_job_running(launcher, job_id, _LAUNCH_TIMEOUT)
+        await wait_for_model_healthy(served_model_name, cscs_api_key, _HEALTH_TIMEOUT)
     finally:
         await launcher.cancel_job(job_id)

tests/integration/utils.py

@@ -0,0 +1,41 @@
+import asyncio
+
+import pytest
+
+from swiss_ai_model_launch.cli.healthcheck import ModelHealth, check_model_health
+from swiss_ai_model_launch.launchers.firecrest_launcher import FirecRESTLauncher
+from swiss_ai_model_launch.launchers.launcher import JobStatus
+
+
+async def wait_for_job_running(
+    launcher: FirecRESTLauncher,
+    job_id: int,
+    timeout_min: int,
+    poll_interval_seconds: int = 15,
+) -> None:
+    deadline = asyncio.get_event_loop().time() + timeout_min * 60
+    while asyncio.get_event_loop().time() < deadline:
+        await asyncio.sleep(poll_interval_seconds)
+        status = await launcher.get_job_status(job_id)
+        print(f"[job {job_id}] status: {status.value}")
+        if status == JobStatus.RUNNING:
+            return
+        if status == JobStatus.TIMEOUT:
+            pytest.fail(f"Job {job_id} timed out before becoming RUNNING.")
+    pytest.fail(f"Job {job_id} didn't reach RUNNING within {timeout_min} mins.")
+
+
+async def wait_for_model_healthy(
+    model_name: str,
+    api_key: str,
+    timeout_min: int,
+    poll_interval_seconds: int = 30,
+) -> None:
+    deadline = asyncio.get_event_loop().time() + timeout_min * 60
+    while asyncio.get_event_loop().time() < deadline:
+        await asyncio.sleep(poll_interval_seconds)
+        health = await check_model_health(model_name, api_key)
+        print(f"[{model_name}] health: {health.value}")
+        if health == ModelHealth.HEALTHY:
+            return
+    pytest.fail(f"'{model_name}' didn't become HEALTHY within {timeout_min} mins.")