feat(BA-8642): add the coordinator client that pushes ingresses to api-delivery workers

changes/15915.feature.md

@@ -0,0 +1 @@
+Add the AppProxy coordinator client that pushes ingresses to workers with the `api` delivery mode

src/ai/backend/appproxy/coordinator/clients/worker/client.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import aiohttp
+
+from ai.backend.appproxy.common.config import SecretConfig
+from ai.backend.appproxy.coordinator.errors.ingress import (
+    ApiDeliveryWorkerRejectedError,
+    ApiDeliveryWorkerServerError,
+    ApiDeliveryWorkerUnreachableError,
+)
+from ai.backend.common.api_handlers import BaseRequestModel
+from ai.backend.common.dto.appproxy.v2.ingress.request import (
+    IngressRemoveRequest,
+    IngressUpsertRequest,
+)
+from ai.backend.common.dto.appproxy.v2.ingress.types import Ingress, IngressKey
+
+_TOKEN_HEADER = "X-BackendAI-Token"
+
+
+class AppProxyWorkerClientPool:
+    """Holds one AppProxyWorkerClient per worker address."""
+
+    _secrets: SecretConfig
+    _clients: dict[str, AppProxyWorkerClient]
+
+    def __init__(self, secrets: SecretConfig) -> None:
+        self._secrets = secrets
+        self._clients = {}
+
+    def load_client(self, addr: str) -> AppProxyWorkerClient:
+        client = self._clients.get(addr)
+        if client is None:
+            session = aiohttp.ClientSession(
+                base_url=f"http://{addr}",
+                timeout=aiohttp.ClientTimeout(total=30.0, connect=10.0),
+            )
+            client = AppProxyWorkerClient(session, addr, self._secrets)
+            self._clients[addr] = client
+        return client
+
+    async def close(self) -> None:
+        for client in self._clients.values():
+            await client.close()
+        self._clients.clear()
+
+
+class AppProxyWorkerClient:
+    """Sends ingress pushes to the API port of one worker node."""
+
+    _session: aiohttp.ClientSession
+    _addr: str
+    _secrets: SecretConfig
+
+    def __init__(self, session: aiohttp.ClientSession, addr: str, secrets: SecretConfig) -> None:
+        self._session = session
+        self._addr = addr
+        self._secrets = secrets
+
+    async def close(self) -> None:
+        await self._session.close()
+
+    async def upsert(self, ingresses: Sequence[Ingress]) -> None:
+        await self._post("/v2/ingress/upsert", IngressUpsertRequest(ingresses=list(ingresses)))
+
+    async def remove(self, keys: Sequence[IngressKey]) -> None:
+        await self._post("/v2/ingress/remove", IngressRemoveRequest(keys=list(keys)))
+
+    async def _post(self, path: str, request: BaseRequestModel) -> None:
+        try:
+            async with self._session.post(
+                path,
+                json=request.model_dump(mode="json"),
+                headers={_TOKEN_HEADER: self._secrets.api_secret},
+            ) as resp:
+                body = await resp.read()
+                extra_msg = f"worker {self._addr} answered {resp.status}"
+                extra_data = {
+                    "worker_addr": self._addr,
+                    "status": resp.status,
+                    "body": body.decode(errors="replace"),
+                }
+                match resp.status // 100:
+                    case 4:
+                        raise ApiDeliveryWorkerRejectedError(
+                            extra_msg=extra_msg, extra_data=extra_data
+                        )
+                    case 5:
+                        raise ApiDeliveryWorkerServerError(
+                            extra_msg=extra_msg, extra_data=extra_data
+                        )
+        except aiohttp.ClientError as e:
+            raise ApiDeliveryWorkerUnreachableError(
+                extra_msg=f"worker {self._addr} is unreachable",
+                extra_data={"worker_addr": self._addr},
+            ) from e

src/ai/backend/appproxy/coordinator/errors/ingress.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from typing import override
+
+from aiohttp import web
+
+from ai.backend.common.exception import (
+    BackendAIError,
+    ErrorCode,
+    ErrorDetail,
+    ErrorDomain,
+    ErrorOperation,
+)
+
+
+class ApiDeliveryWorkerUnreachableError(BackendAIError, web.HTTPBadGateway):
+    """Raised when an api-delivery worker cannot be reached."""
+
+    error_type = "https://api.backend.ai/probs/appproxy-coordinator/api-delivery-worker-unreachable"
+    error_title = "An api-delivery worker cannot be reached."
+
+    @override
+    def error_code(self) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.APPPROXY,
+            operation=ErrorOperation.REQUEST,
+            error_detail=ErrorDetail.UNREACHABLE,
+        )
+
+
+class ApiDeliveryWorkerRejectedError(BackendAIError, web.HTTPBadRequest):
+    """Raised when an api-delivery worker answers a request with a 4xx."""
+
+    error_type = "https://api.backend.ai/probs/appproxy-coordinator/api-delivery-worker-rejected"
+    error_title = "An api-delivery worker rejected the request."
+
+    @override
+    def error_code(self) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.APPPROXY,
+            operation=ErrorOperation.REQUEST,
+            error_detail=ErrorDetail.BAD_REQUEST,
+        )
+
+
+class ApiDeliveryWorkerServerError(BackendAIError, web.HTTPBadGateway):
+    """Raised when an api-delivery worker answers a request with a 5xx."""
+
+    error_type = (
+        "https://api.backend.ai/probs/appproxy-coordinator/api-delivery-worker-server-error"
+    )
+    error_title = "An api-delivery worker failed to handle the request."
+
+    @override
+    def error_code(self) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.APPPROXY,
+            operation=ErrorOperation.REQUEST,
+            error_detail=ErrorDetail.INTERNAL_ERROR,
+        )

tests/unit/appproxy/coordinator/clients/worker/BUILD

@@ -0,0 +1,3 @@
+python_tests(
+    name="tests",
+)

tests/unit/appproxy/coordinator/clients/worker/test_client.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from http import HTTPStatus
+from unittest.mock import AsyncMock, MagicMock
+
+import aiohttp
+import pytest
+
+from ai.backend.appproxy.common.config import SecretConfig
+from ai.backend.appproxy.coordinator.clients.worker.client import (
+    AppProxyWorkerClient,
+    AppProxyWorkerClientPool,
+)
+from ai.backend.appproxy.coordinator.errors.ingress import (
+    ApiDeliveryWorkerRejectedError,
+    ApiDeliveryWorkerServerError,
+    ApiDeliveryWorkerUnreachableError,
+)
+from ai.backend.common.dto.appproxy.v2.ingress.types import (
+    Ingress,
+    IngressAccess,
+    IngressAccessMode,
+    IngressAuthority,
+    IngressBackend,
+    IngressBackendID,
+    IngressFrontend,
+    IngressFrontendKind,
+    IngressID,
+    IngressProtocol,
+    IngressTargetGroup,
+    IngressTargetGroupID,
+    PortFrontend,
+)
+from ai.backend.common.exception import BackendAIError
+
+API_SECRET = "api-secret"
+WORKER_ADDR = "10.0.0.1:10201"
+
+
+@dataclass(frozen=True)
+class _ErrorAnswerCase:
+    status: HTTPStatus
+    expected_error: type[BackendAIError]
+
+
+def _answering_session(status: int, body: bytes) -> MagicMock:
+    response = MagicMock()
+    response.status = status
+    response.read = AsyncMock(return_value=body)
+    session = MagicMock()
+    session.post.return_value.__aenter__.return_value = response
+    return session
+
+
+@pytest.fixture
+def secrets() -> SecretConfig:
+    return SecretConfig(api_secret=API_SECRET, jwt_secret="jwt-secret")
+
+
+@pytest.fixture
+def ingress() -> Ingress:
+    return Ingress(
+        id=IngressID("ingress-1"),
+        authority=IngressAuthority("worker-1"),
+        revision=3,
+        frontend=IngressFrontend(
+            kind=IngressFrontendKind.PORT,
+            protocol=IngressProtocol.HTTP,
+            port=PortFrontend(number=10205),
+        ),
+        access=IngressAccess(mode=IngressAccessMode.PUBLIC, allowed_client_ips=[]),
+        backend=IngressBackend(
+            id=IngressBackendID("backend-1"),
+            target_groups=[
+                IngressTargetGroup(id=IngressTargetGroupID("group-1"), weight=1, targets=[])
+            ],
+        ),
+    )
+
+
+@pytest.fixture
+def applying_session() -> MagicMock:
+    return _answering_session(HTTPStatus.OK, b"")
+
+
+@pytest.fixture
+def unreachable_session() -> MagicMock:
+    session = MagicMock()
+    session.post.side_effect = aiohttp.ClientConnectionError("connection refused")
+    return session
+
+
+class TestAppProxyWorkerClient:
+    async def test_push_carries_the_api_secret(
+        self, applying_session: MagicMock, secrets: SecretConfig, ingress: Ingress
+    ) -> None:
+        client = AppProxyWorkerClient(applying_session, WORKER_ADDR, secrets)
+        await client.upsert([ingress])
+        assert applying_session.post.call_args.kwargs["headers"] == {
+            "X-BackendAI-Token": API_SECRET
+        }
+
+    @pytest.mark.parametrize(
+        "case",
+        [
+            _ErrorAnswerCase(
+                status=HTTPStatus.BAD_REQUEST, expected_error=ApiDeliveryWorkerRejectedError
+            ),
+            _ErrorAnswerCase(
+                status=HTTPStatus.UNAUTHORIZED, expected_error=ApiDeliveryWorkerRejectedError
+            ),
+            _ErrorAnswerCase(
+                status=HTTPStatus.CONFLICT, expected_error=ApiDeliveryWorkerRejectedError
+            ),
+            _ErrorAnswerCase(
+                status=HTTPStatus.INTERNAL_SERVER_ERROR,
+                expected_error=ApiDeliveryWorkerServerError,
+            ),
+            _ErrorAnswerCase(
+                status=HTTPStatus.NOT_IMPLEMENTED, expected_error=ApiDeliveryWorkerServerError
+            ),
+        ],
+        ids=lambda case: str(case.status.value),
+    )
+    async def test_error_answer_raises_with_the_worker_status(
+        self, case: _ErrorAnswerCase, secrets: SecretConfig, ingress: Ingress
+    ) -> None:
+        session = _answering_session(case.status, b"rejected")
+        client = AppProxyWorkerClient(session, WORKER_ADDR, secrets)
+        with pytest.raises(case.expected_error) as exc_info:
+            await client.upsert([ingress])
+        assert exc_info.value.extra_data == {
+            "worker_addr": WORKER_ADDR,
+            "status": case.status,
+            "body": "rejected",
+        }
+
+    async def test_unreachable_worker_raises(
+        self, unreachable_session: MagicMock, secrets: SecretConfig, ingress: Ingress
+    ) -> None:
+        client = AppProxyWorkerClient(unreachable_session, WORKER_ADDR, secrets)
+        with pytest.raises(ApiDeliveryWorkerUnreachableError):
+            await client.upsert([ingress])
+
+
+class TestAppProxyWorkerClientPool:
+    async def test_client_of_an_address_is_reused(self, secrets: SecretConfig) -> None:
+        pool = AppProxyWorkerClientPool(secrets)
+        try:
+            assert pool.load_client(WORKER_ADDR) is pool.load_client(WORKER_ADDR)
+        finally:
+            await pool.close()