feat: sync up-to-date health check information with AppProxy (#5230)

Author: kyujin-cho

changes/5230.feature.md

@@ -0,0 +1 @@
+Sync model service's health information real-time with AppProxy

src/ai/backend/common/clients/valkey_client/valkey_live/client.py

@@ -1,3 +1,4 @@
+import json
 import logging
 from collections.abc import Mapping
 from typing import (
@@ -7,6 +8,7 @@
     Self,
     cast,
 )
+from uuid import UUID
 
 from glide import (
     Batch,
@@ -22,6 +24,7 @@
     create_layer_aware_valkey_decorator,
     create_valkey_client,
 )
+from ai.backend.common.data.config.types import HealthCheckConfig
 from ai.backend.common.metrics.metric import LayerType
 from ai.backend.common.types import RedisTarget
 from ai.backend.logging.utils import BraceStyleAdapter
@@ -473,6 +476,32 @@ async def hgetall_str(self, key: str) -> dict[str, str]:
 
         return str_result
 
+    @valkey_decorator()
+    async def update_appproxy_redis_info(
+        self,
+        endpoint_id: UUID,
+        connection_info: dict[str, Any],
+        health_check_config: HealthCheckConfig | None,
+    ) -> None:
+        pipe = self._create_batch()
+        pipe.set(
+            f"endpoint.{endpoint_id}.route_connection_info",
+            json.dumps(connection_info),
+            expiry=ExpirySet(ExpiryType.SEC, 3600),
+        )
+        pipe.set(
+            f"endpoint.{endpoint_id}.health_check_enabled",
+            "true" if health_check_config is not None else "false",
+            expiry=ExpirySet(ExpiryType.SEC, 3600),
+        )
+        if health_check_config:
+            pipe.set(
+                f"endpoint.{endpoint_id}.health_check_config",
+                health_check_config.model_dump_json(),
+                expiry=ExpirySet(ExpiryType.SEC, 3600),
+            )
+        await self._client.client.exec(pipe, True)
+
     @valkey_decorator()
     async def delete_key(self, key: str) -> int:
         """

src/ai/backend/manager/models/endpoint.py

@@ -70,6 +70,7 @@
 from .image import ImageRow
 from .routing import RouteStatus
 from .scaling_group import scaling_groups
+from .session import SessionStatus
 from .user import UserRow
 from .vfolder import prepare_vfolder_mounts
 
@@ -543,7 +544,9 @@ async def generate_redis_route_info(
             [
                 r.session
                 for r in active_routes
-                if r.status in RouteStatus.active_route_statuses() and r.session
+                if r.status in RouteStatus.active_route_statuses()
+                and r.session
+                and r.session_row.status == SessionStatus.RUNNING
             ],
         )
         session_id_to_route_map = {r.session: r for r in active_routes}

src/ai/backend/manager/registry.py

@@ -3649,7 +3649,8 @@ async def create_appproxy_endpoint(
 
         model = await VFolderRow.get(db_sess, endpoint.model)
 
-        health_check_information = await self.get_health_check_info(endpoint, model)
+        health_check_config = await self.get_health_check_info(endpoint, model)
+
         request_body = {
             "version": "v2",
             "service_name": endpoint.name,
@@ -3666,8 +3667,8 @@ async def create_appproxy_endpoint(
                 },
             },
             "open_to_public": endpoint.open_to_public,
-            "health_check": health_check_information.model_dump(mode="json")
-            if health_check_information
+            "health_check": health_check_config.model_dump(mode="json")
+            if health_check_config
             else None,
         }
         endpoint_json = await wsproxy_client.create_endpoint(endpoint.id, request_body)
@@ -3692,12 +3693,13 @@ async def notify_endpoint_route_update_to_appproxy(self, endpoint_id: uuid.UUID)
         async with self.db.begin_readonly_session() as db_sess:
             endpoint = await EndpointRow.get(db_sess, endpoint_id)
             connection_info = await endpoint.generate_redis_route_info(db_sess)
+            model = await VFolderRow.get(db_sess, endpoint.model)
 
-        await self.valkey_live.delete_key(f"endpoint.{endpoint.id}.route_connection_info")
-        await self.valkey_live.store_live_data(
-            f"endpoint.{endpoint.id}.route_connection_info",
-            json.dumps(connection_info),
-            ex=3600,
+        health_check_config = await self.get_health_check_info(endpoint, model)
+        await self.valkey_live.update_appproxy_redis_info(
+            endpoint.id,
+            connection_info,
+            health_check_config,
         )
 
         await self.event_producer.anycast_event(EndpointRouteListUpdatedEvent(endpoint.id))