refactor(BA-4186): move out of convention function into valid directory (#8509)

Author: kwonkwonn

changes/8509.enhance.md

@@ -0,0 +1 @@
+Resolve cyclic import error from AppProxy with moving non-relevant function out

src/ai/backend/appproxy/coordinator/api/types.py

@@ -25,6 +25,15 @@ class StubResponseModel(BaseModel):
     success: Annotated[bool, Field(default=True)]
 
 
+class AppProxyStatusResponse(BaseModel):
+    """Response from AppProxy /status endpoint."""
+
+    api_version: str = Field(description="AppProxy API version (e.g., 'v1', 'v2')")
+    advertise_address: str | None = Field(
+        default=None, description="Advertised address for AppProxy"
+    )
+
+
 class CircuitListResponseModel(BaseModel):
     circuits: list[SerializableCircuit]
 

src/ai/backend/appproxy/coordinator/server.py

@@ -71,6 +71,7 @@
     mime_match,
     ping_redis_connection,
 )
+from ai.backend.appproxy.coordinator.api.types import AppProxyStatusResponse
 from ai.backend.appproxy.coordinator.models.worker import WorkerStatus
 from ai.backend.common import redis_helper
 from ai.backend.common.clients.valkey_client.valkey_leader.client import ValkeyLeaderClient
@@ -845,10 +846,11 @@ async def status(request: web.Request) -> web.Response:
     root_ctx: RootContext = request.app["_root.context"]
     request["do_not_print_access_log"] = True
     advertised_addr = root_ctx.local_config.proxy_coordinator.advertise_base_url
-    return web.json_response({
-        "api_version": "v2",
-        "advertise_address": advertised_addr,
-    })
+    response = AppProxyStatusResponse(
+        api_version="v2",
+        advertise_address=advertised_addr,
+    )
+    return web.json_response(response.model_dump(mode="json"))
 
 
 def handle_loop_error(

src/ai/backend/common/exception.py

@@ -173,6 +173,7 @@ class ErrorDomain(enum.StrEnum):
     ROLE = "role"
     METRIC = "metric"
     STORAGE_PROXY = "storage-proxy"
+    APPPROXY = "appproxy"
     MESSAGE_QUEUE = "message-queue"
     NOTIFICATION = "notification"
     HEALTH_CHECK = "health-check"

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

@@ -40,6 +40,7 @@
     from ai.backend.common.types import ValkeyProfileTarget
     from ai.backend.manager.agent_cache import AgentRPCCache
     from ai.backend.manager.clients.agent import AgentClientPool
+    from ai.backend.manager.clients.appproxy.client import AppProxyClientPool
     from ai.backend.manager.config.provider import ManagerConfigProvider
     from ai.backend.manager.idle import IdleCheckerHost
     from ai.backend.manager.models.storage import StorageSessionManager
@@ -97,6 +98,7 @@ class RootContext(BaseContext):
     registry: AgentRegistry
     agent_cache: AgentRPCCache
     agent_client_pool: AgentClientPool
+    appproxy_client_pool: AppProxyClientPool
     sokovan_orchestrator: SokovanOrchestrator
     scheduling_controller: SchedulingController
     deployment_controller: DeploymentController

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

@@ -1,23 +1,16 @@
-import json
 import logging
 from collections.abc import Iterable
 from dataclasses import dataclass, field
 from http import HTTPStatus
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any
 
-import aiohttp
 import aiohttp_cors
-import aiotools
 import trafaret as t
 from aiohttp import web
 
 from ai.backend.common import validators as tx
 from ai.backend.logging import BraceStyleAdapter
-from ai.backend.manager.errors.common import (
-    InternalServerError,
-    ObjectNotFound,
-    ServerMisconfiguredError,
-)
+from ai.backend.manager.errors.common import ObjectNotFound
 from ai.backend.manager.models.scaling_group import query_allowed_sgroups
 from ai.backend.manager.models.utils import ExtendedAsyncSAEngine
 
@@ -37,27 +30,6 @@ class WSProxyVersionQueryParams:
     db_ctx: ExtendedAsyncSAEngine = field(hash=False)
 
 
-@aiotools.lru_cache(expire_after=30)  # expire after 30 seconds
-async def query_wsproxy_status(
-    wsproxy_addr: str,
-) -> dict[str, Any]:
-    async with (
-        aiohttp.ClientSession() as session,
-        session.get(
-            wsproxy_addr + "/status",
-            headers={"Accept": "application/json"},
-        ) as resp,
-    ):
-        try:
-            result = await resp.json()
-        except (aiohttp.ContentTypeError, json.JSONDecodeError) as e:
-            log.error("Failed to parse wsproxy status response from {}: {}", wsproxy_addr, e)
-            raise InternalServerError(
-                "Got invalid response from wsproxy when querying status"
-            ) from e
-        return cast(dict[str, Any], result)
-
-
 @auth_required
 @server_status_required(READ_ALLOWED)
 @check_api_params(
@@ -100,29 +72,27 @@ async def get_wsproxy_version(request: web.Request, params: Any) -> web.Response
     domain_name = request["user"]["domain_name"]
     group_id_or_name = params["group"]
     log.info("SGROUPS.LIST(ak:{}, g:{}, d:{})", access_key, group_id_or_name, domain_name)
+    # remove appproxy client pool from root_ctx when db query migrated to service layer.
     async with root_ctx.db.begin_readonly() as conn:
         sgroups = await query_allowed_sgroups(conn, domain_name, group_id_or_name or "", access_key)
-        for sgroup in sgroups:
-            if sgroup.name == scaling_group_name:
-                wsproxy_addr = sgroup.wsproxy_addr
-                if not wsproxy_addr:
-                    wsproxy_version = "v1"
-                else:
-                    try:
-                        wsproxy_status = await query_wsproxy_status(wsproxy_addr)
-                        wsproxy_version = wsproxy_status["api_version"]
-                    except aiohttp.ClientConnectorError:
-                        log.error(
-                            "Failed to query the wsproxy {1} configured for sg:{0}",
-                            scaling_group_name,
-                            wsproxy_addr,
-                        )
-                        return ServerMisconfiguredError()
-                return web.json_response({
-                    "wsproxy_version": wsproxy_version,
-                })
-        else:
+        sgroup_filtered = [sg for sg in sgroups if sg.name == scaling_group_name]
+        if not sgroup_filtered:
             raise ObjectNotFound(object_name="scaling group")
+        sgroup = sgroup_filtered[0]
+
+        if not sgroup.wsproxy_addr:
+            # if wsproxy_addr is not set, raise not found error(migrating from v1 behavior)
+            # It should be either 404 or 500 before wsproxy_addr is mandatory field.
+            raise ObjectNotFound(object_name="AppProxy address")
+        client = root_ctx.appproxy_client_pool.load_client(
+            sgroup.wsproxy_addr, sgroup.wsproxy_api_token or ""
+        )
+        status = await client.fetch_status()
+        wsproxy_version = status.api_version
+
+        return web.json_response({
+            "wsproxy_version": wsproxy_version,
+        })
 
 
 async def init(app: web.Application) -> None:

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

@@ -1,16 +1,30 @@
+from __future__ import annotations
+
+import json
+import logging
 from typing import Any
 from uuid import UUID
 
 import aiohttp
 
+from ai.backend.appproxy.coordinator.api.types import AppProxyStatusResponse
+from ai.backend.common.clients.http_client.client_pool import (
+    ClientKey,
+    ClientPool,
+    tcp_client_session_factory,
+)
 from ai.backend.common.exception import BackendAIError
 from ai.backend.common.metrics.metric import DomainType, LayerType
 from ai.backend.common.resilience.policies.metrics import MetricArgs, MetricPolicy
 from ai.backend.common.resilience.policies.retry import BackoffStrategy, RetryArgs, RetryPolicy
 from ai.backend.common.resilience.resilience import Resilience
+from ai.backend.logging import BraceStyleAdapter
+from ai.backend.manager.errors.appproxy import AppProxyConnectionError, AppProxyResponseError
 
 from .types import CreateEndpointRequestBody
 
+log: BraceStyleAdapter = BraceStyleAdapter(logging.getLogger(__spec__.name))
+
 appproxy_client_resilience = Resilience(
     policies=[
         MetricPolicy(MetricArgs(domain=DomainType.CLIENT, layer=LayerType.WSPROXY_CLIENT)),
@@ -26,6 +40,25 @@
 )
 
 
+class AppProxyClientPool:
+    _client_pool: ClientPool
+
+    def __init__(self) -> None:
+        self._client_pool = ClientPool(tcp_client_session_factory)
+
+    def load_client(self, address: str, token: str) -> AppProxyClient:
+        client_session = self._client_pool.load_client_session(
+            ClientKey(
+                endpoint=address,
+                domain="appproxy",
+            )
+        )
+        return AppProxyClient(client_session, address, token)
+
+    async def close(self) -> None:
+        await self._client_pool.close()
+
+
 class AppProxyClient:
     _client_session: aiohttp.ClientSession
     _address: str
@@ -36,6 +69,26 @@ def __init__(self, client_session: aiohttp.ClientSession, address: str, token: s
         self._address = address
         self._token = token
 
+    @appproxy_client_resilience.apply()
+    async def fetch_status(self) -> AppProxyStatusResponse:
+        try:
+            async with self._client_session.get(
+                "/status",
+                headers={"Accept": "application/json"},
+            ) as resp:
+                data = await resp.json()
+                return AppProxyStatusResponse.model_validate(data)
+        except aiohttp.ClientConnectorError as e:
+            log.error("Failed to connect to app-proxy at {}: {}", self._address, e)
+            raise AppProxyConnectionError(
+                extra_msg=f"Failed to connect to AppProxy at {self._address}"
+            ) from e
+        except (aiohttp.ContentTypeError, json.JSONDecodeError) as e:
+            log.error("Failed to parse app-proxy status response from {}: {}", self._address, e)
+            raise AppProxyResponseError(
+                extra_msg=f"Invalid response from AppProxy at {self._address}"
+            ) from e
+
     @appproxy_client_resilience.apply()
     async def create_endpoint(
         self,

src/ai/backend/manager/errors/appproxy.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from aiohttp import web
+
+from ai.backend.common.exception import (
+    BackendAIError,
+    ErrorCode,
+    ErrorDetail,
+    ErrorDomain,
+    ErrorOperation,
+)
+
+
+class AppProxyConnectionError(BackendAIError, web.HTTPServiceUnavailable):
+    """Raised when connection to AppProxy fails."""
+
+    error_type = "https://api.backend.ai/probs/appproxy-connection-error"
+    error_title = "Failed to connect to AppProxy."
+
+    def error_code(self) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.APPPROXY,
+            operation=ErrorOperation.REQUEST,
+            error_detail=ErrorDetail.UNAVAILABLE,
+        )
+
+
+class AppProxyResponseError(BackendAIError, web.HTTPInternalServerError):
+    """Raised when AppProxy returns an invalid response."""
+
+    error_type = "https://api.backend.ai/probs/appproxy-response-error"
+    error_title = "Invalid response from AppProxy."
+
+    def error_code(self) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.APPPROXY,
+            operation=ErrorOperation.REQUEST,
+            error_detail=ErrorDetail.INVALID_DATA_FORMAT,
+        )

src/ai/backend/manager/registry.py

@@ -204,6 +204,8 @@ class AgentRegistry:
     pending_waits: set[asyncio.Task[None]]
     database_ptask_group: aiotools.PersistentTaskGroup
     webhook_ptask_group: aiotools.PersistentTaskGroup
+    # TODO: Migrate to use root_ctx.appproxy_client_pool instead.
+    # After migration, remove _client_pool and _load_app_proxy_client.
     _client_pool: ClientPool
     _agent_client_pool: AgentClientPool
 
@@ -268,6 +270,7 @@ async def shutdown(self) -> None:
         await self.database_ptask_group.shutdown()
         await self.webhook_ptask_group.shutdown()
 
+    # TODO: Migrate callers to use root_ctx.appproxy_client_pool.load_client() instead.
     def _load_app_proxy_client(self, address: str, token: str) -> AppProxyClient:
         client_session = self._client_pool.load_client_session(
             ClientKey(

src/ai/backend/manager/server.py

@@ -146,6 +146,7 @@
 from .api import ManagerStatus
 from .api.context import RootContext
 from .clients.agent import AgentClientPool, AgentPoolSpec
+from .clients.appproxy.client import AppProxyClientPool
 from .config.bootstrap import BootstrapConfig
 from .config.loader.config_overrider import ConfigOverrider
 from .config.loader.etcd_loader import (
@@ -762,6 +763,7 @@ async def processors_ctx(root_ctx: RootContext) -> AsyncIterator[None]:
                 event_producer=root_ctx.event_producer,
                 agent_cache=root_ctx.agent_cache,
                 notification_center=root_ctx.notification_center,
+                appproxy_client_pool=root_ctx.appproxy_client_pool,
             ),
         ),
         [reporter_monitor, prometheus_monitor, audit_log_monitor],
@@ -1103,6 +1105,7 @@ async def agent_registry_ctx(root_ctx: RootContext) -> AsyncIterator[None]:
             recovery_timeout=60.0,
         ),
     )
+    root_ctx.appproxy_client_pool = AppProxyClientPool()
     root_ctx.registry = AgentRegistry(
         root_ctx.config_provider,
         root_ctx.db,
@@ -1126,6 +1129,7 @@ async def agent_registry_ctx(root_ctx: RootContext) -> AsyncIterator[None]:
         yield
     finally:
         await root_ctx.agent_client_pool.close()
+        await root_ctx.appproxy_client_pool.close()
         await root_ctx.registry.shutdown()