feat(BA-992): Offload health check capability to AppProxy (#5134)

Co-authored-by: HyeockJinKim <[email protected]>

Author: 2 people

changes/5134.breaking.md

@@ -0,0 +1,3 @@
+- Health check capability temporarily broken on OSS AppProxy due to architectural changes
+- Users must disable health check feature in `model-definition.yaml` to use model services on Open Source Backend.AI
+- OSS AppProxy support will be restored in future releases

changes/5134.feature.md

@@ -0,0 +1 @@
+- Offload model service health check architecture to AppProxy with Redis-based route management for improved scalability and real-time endpoint monitoring

src/ai/backend/agent/agent.py

@@ -2934,15 +2934,11 @@ async def start_and_monitor_model_service_health(
             # if everything went well then krunner itself will report the status via zmq
             await self.anycast_and_broadcast_event(
                 ModelServiceStatusAnycastEvent(
-                    kernel_obj.kernel_id,
                     kernel_obj.session_id,
-                    model["name"],
                     ModelServiceStatus.UNHEALTHY,
                 ),
                 ModelServiceStatusBroadcastEvent(
-                    kernel_obj.kernel_id,
                     kernel_obj.session_id,
-                    model["name"],
                     ModelServiceStatus.UNHEALTHY,
                 ),
             )

src/ai/backend/agent/kernel.py

@@ -50,15 +50,11 @@
 from ai.backend.common.events.event_types.kernel.types import (
     KernelLifecycleEventReason,
 )
-from ai.backend.common.events.event_types.model_serving.anycast import (
-    ModelServiceStatusAnycastEvent,
-)
 from ai.backend.common.json import load_json
 from ai.backend.common.types import (
     AgentId,
     CommitStatus,
     KernelId,
-    ModelServiceStatus,
     ServicePort,
     SessionId,
     SessionTypes,
@@ -1129,18 +1125,8 @@ async def read_output(self) -> None:
                         case b"model-service-result":
                             await self.model_service_queue.put(msg_data)
                         case b"model-service-status":
-                            response = load_json(msg_data)
-                            event = ModelServiceStatusAnycastEvent(
-                                self.kernel_id,
-                                self.session_id,
-                                response["model_name"],
-                                (
-                                    ModelServiceStatus.HEALTHY
-                                    if response["is_healthy"]
-                                    else ModelServiceStatus.UNHEALTHY
-                                ),
-                            )
-                            await self.event_producer.anycast_event(event)
+                            # no-op
+                            pass
                         case b"apps-result":
                             await self.service_apps_info_queue.put(msg_data)
                         case b"stdout":

src/ai/backend/common/data/config/types.py

@@ -1,5 +1,7 @@
 from dataclasses import dataclass
-from typing import Optional
+from typing import Annotated, Optional
+
+from pydantic import BaseModel, Field
 
 from ai.backend.common.typed_validators import HostPortPair
 
@@ -10,3 +12,15 @@ class EtcdConfigData:
     addr: HostPortPair
     user: Optional[str]
     password: Optional[str]
+
+
+class HealthCheckConfig(BaseModel):
+    """
+    Health check configuration matching model-definition.yaml schema
+    """
+
+    interval: Annotated[float, Field(default=10.0, ge=0)] = 10.0
+    path: str
+    max_retries: Annotated[int, Field(default=10, ge=1)] = 10
+    max_wait_time: Annotated[float, Field(default=15.0, ge=0)] = 15.0
+    expected_status_code: Annotated[int, Field(default=200, ge=100, le=599)] = 200

src/ai/backend/common/events/event_types/model_serving/__init__.py

@@ -0,0 +1,39 @@
+import uuid
+from dataclasses import dataclass
+from typing import Optional, override
+
+from ai.backend.common.events.types import AbstractEvent, EventDomain
+from ai.backend.common.events.user_event.user_event import UserEvent
+from ai.backend.common.types import ModelServiceStatus, SessionId
+
+
+@dataclass
+class ModelServiceStatusEventArgs(AbstractEvent):
+    session_id: SessionId
+    new_status: ModelServiceStatus
+
+    def serialize(self) -> tuple:
+        return (
+            str(self.session_id),
+            self.new_status.value,
+        )
+
+    @classmethod
+    def deserialize(cls, value: tuple):
+        return cls(
+            session_id=SessionId(uuid.UUID(value[0])),
+            new_status=ModelServiceStatus(value[1]),
+        )
+
+    @classmethod
+    @override
+    def event_domain(cls) -> EventDomain:
+        return EventDomain.MODEL_SERVING
+
+    @override
+    def domain_id(self) -> Optional[str]:
+        return None
+
+    @override
+    def user_event(self) -> Optional[UserEvent]:
+        return None

src/ai/backend/common/events/event_types/model_serving/anycast.py

@@ -4,81 +4,81 @@
 
 from ai.backend.common.events.types import AbstractAnycastEvent, EventDomain
 from ai.backend.common.events.user_event.user_event import UserEvent
-from ai.backend.common.types import KernelId, ModelServiceStatus, SessionId
+
+from . import ModelServiceStatusEventArgs
+
+
+class ModelServiceStatusAnycastEvent(ModelServiceStatusEventArgs, AbstractAnycastEvent):
+    @classmethod
+    @override
+    def event_name(cls) -> str:
+        return "model_service_status_updated"
 
 
 @dataclass
-class ModelServiceStatusEventArgs(AbstractAnycastEvent):
-    kernel_id: KernelId
-    session_id: SessionId
-    model_name: str
-    new_status: ModelServiceStatus
+class RouteCreationEvent(AbstractAnycastEvent):
+    route_id: uuid.UUID
 
     def serialize(self) -> tuple:
-        return (
-            str(self.kernel_id),
-            str(self.session_id),
-            self.model_name,
-            self.new_status.value,
-        )
+        return (str(self.route_id),)
 
     @classmethod
     def deserialize(cls, value: tuple):
-        return cls(
-            kernel_id=KernelId(uuid.UUID(value[0])),
-            session_id=SessionId(uuid.UUID(value[1])),
-            model_name=value[2],
-            new_status=ModelServiceStatus(value[3]),
-        )
+        return cls(uuid.UUID(value[0]))
 
     @classmethod
     @override
     def event_domain(cls) -> EventDomain:
-        return EventDomain.MODEL_SERVING
+        return EventDomain.MODEL_ROUTE
 
     @override
     def domain_id(self) -> Optional[str]:
-        return None
+        return str(self.route_id)
 
     @override
     def user_event(self) -> Optional[UserEvent]:
         return None
 
 
-class ModelServiceStatusAnycastEvent(ModelServiceStatusEventArgs):
+class RouteCreatedAnycastEvent(RouteCreationEvent):
     @classmethod
     @override
     def event_name(cls) -> str:
-        return "model_service_status_updated"
+        return "route_created"
+
+
+class RouteTerminatingEvent(RouteCreationEvent):
+    @classmethod
+    @override
+    def event_name(cls) -> str:
+        return "route_terminating"
 
 
 @dataclass
-class RouteCreationEvent(AbstractAnycastEvent):
-    route_id: uuid.UUID
+class EndpointRouteListUpdatedEvent(AbstractAnycastEvent):
+    endpoint_id: uuid.UUID
 
     def serialize(self) -> tuple:
-        return (str(self.route_id),)
+        return (str(self.endpoint_id),)
 
     @classmethod
     def deserialize(cls, value: tuple):
         return cls(uuid.UUID(value[0]))
 
+    @classmethod
+    @override
+    def event_name(cls) -> str:
+        return "endpoint_route_list_updated"
+
     @classmethod
     @override
     def event_domain(cls) -> EventDomain:
         return EventDomain.MODEL_ROUTE
 
     @override
     def domain_id(self) -> Optional[str]:
-        return str(self.route_id)
+        return str(self.endpoint_id)
 
     @override
     def user_event(self) -> Optional[UserEvent]:
         return None
-
-
-class RouteCreatedAnycastEvent(RouteCreationEvent):
-    @classmethod
-    @override
-    def event_name(cls) -> str:
-        return "route_created"

src/ai/backend/common/events/event_types/model_serving/broadcast.py

@@ -1,51 +1,11 @@
-import uuid
-from dataclasses import dataclass
-from typing import Optional, override
+from typing import override
 
-from ai.backend.common.events.types import AbstractBroadcastEvent, EventDomain
-from ai.backend.common.events.user_event.user_event import UserEvent
-from ai.backend.common.types import KernelId, ModelServiceStatus, SessionId
+from ai.backend.common.events.types import AbstractBroadcastEvent
 
+from . import ModelServiceStatusEventArgs
 
-@dataclass
-class ModelServiceStatusEventArgs(AbstractBroadcastEvent):
-    kernel_id: KernelId
-    session_id: SessionId
-    model_name: str
-    new_status: ModelServiceStatus
 
-    def serialize(self) -> tuple:
-        return (
-            str(self.kernel_id),
-            str(self.session_id),
-            self.model_name,
-            self.new_status.value,
-        )
-
-    @classmethod
-    def deserialize(cls, value: tuple):
-        return cls(
-            kernel_id=KernelId(uuid.UUID(value[0])),
-            session_id=SessionId(uuid.UUID(value[1])),
-            model_name=value[2],
-            new_status=ModelServiceStatus(value[3]),
-        )
-
-    @classmethod
-    @override
-    def event_domain(cls) -> EventDomain:
-        return EventDomain.MODEL_SERVING
-
-    @override
-    def domain_id(self) -> Optional[str]:
-        return None
-
-    @override
-    def user_event(self) -> Optional[UserEvent]:
-        return None
-
-
-class ModelServiceStatusBroadcastEvent(ModelServiceStatusEventArgs):
+class ModelServiceStatusBroadcastEvent(ModelServiceStatusEventArgs, AbstractBroadcastEvent):
     @classmethod
     @override
     def event_name(cls) -> str:

src/ai/backend/manager/api/service.py

@@ -30,6 +30,7 @@
     AccessKey,
     ClusterMode,
     RuntimeVariant,
+    VFolderID,
     VFolderMount,
     VFolderUsageMode,
 )
@@ -74,7 +75,7 @@
 from ..errors.api import InvalidAPIParameters
 from ..errors.storage import VFolderNotFound
 from ..models import (
-    ModelServicePredicateChecker,
+    ModelServiceHelper,
     UserRole,
     UserRow,
     query_accessible_vfolders,
@@ -521,7 +522,7 @@ async def _validate(request: web.Request, params: NewServiceRequestModel) -> Val
         raise InvalidAPIParameters(f"Cannot spawn more than {_m} sessions for a single service")
 
     async with root_ctx.db.begin_readonly() as conn:
-        checked_scaling_group = await ModelServicePredicateChecker.check_scaling_group(
+        checked_scaling_group = await ModelServiceHelper.check_scaling_group(
             conn,
             params.config.scaling_group,
             owner_access_key,
@@ -579,7 +580,7 @@ async def _validate(request: web.Request, params: NewServiceRequestModel) -> Val
 
         model_id = folder_row["id"]
 
-        vfolder_mounts = await ModelServicePredicateChecker.check_extra_mounts(
+        vfolder_mounts = await ModelServiceHelper.check_extra_mounts(
             conn,
             root_ctx.config_provider.legacy_etcd_config_loader,
             root_ctx.storage_manager,
@@ -596,11 +597,19 @@ async def _validate(request: web.Request, params: NewServiceRequestModel) -> Val
         )
 
     if params.runtime_variant == RuntimeVariant.CUSTOM:
-        yaml_path = await ModelServicePredicateChecker.validate_model_definition(
+        vfid = VFolderID(folder_row["quota_scope_id"], folder_row["id"])
+        yaml_path = await ModelServiceHelper.validate_model_definition_file_exists(
             root_ctx.storage_manager,
-            folder_row,
+            folder_row["host"],
+            vfid,
             params.config.model_definition_path,
         )
+        await ModelServiceHelper.validate_model_definition(
+            root_ctx.storage_manager,
+            folder_row["host"],
+            vfid,
+            yaml_path,
+        )
     else:
         if (
             params.runtime_variant != RuntimeVariant.CMD

src/ai/backend/manager/clients/wsproxy/client.py

@@ -1,4 +1,5 @@
 from typing import Any, Mapping
+from uuid import UUID
 
 import aiohttp
 
@@ -10,7 +11,7 @@ def __init__(self, address: str, token: str) -> None:
 
     async def create_endpoint(
         self,
-        endpoint_id: str,
+        endpoint_id: UUID,
         body: Mapping[str, Any],
     ) -> dict[str, Any]:
         async with aiohttp.ClientSession() as session:
@@ -26,7 +27,7 @@ async def create_endpoint(
 
     async def delete_endpoint(
         self,
-        endpoint_id: str,
+        endpoint_id: UUID,
     ) -> None:
         async with aiohttp.ClientSession() as session:
             async with session.delete(