diff --git a/changes/15915.feature.md b/changes/15915.feature.md new file mode 100644 index 00000000000..2c74f6432e8 --- /dev/null +++ b/changes/15915.feature.md @@ -0,0 +1 @@ +Add the AppProxy coordinator client that pushes ingresses to workers with the `api` delivery mode diff --git a/src/ai/backend/appproxy/coordinator/clients/__init__.py b/src/ai/backend/appproxy/coordinator/clients/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/src/ai/backend/appproxy/coordinator/clients/worker/__init__.py b/src/ai/backend/appproxy/coordinator/clients/worker/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/src/ai/backend/appproxy/coordinator/clients/worker/client.py b/src/ai/backend/appproxy/coordinator/clients/worker/client.py new file mode 100644 index 00000000000..ee23f41f553 --- /dev/null +++ b/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 diff --git a/src/ai/backend/appproxy/coordinator/errors/ingress.py b/src/ai/backend/appproxy/coordinator/errors/ingress.py new file mode 100644 index 00000000000..8073d0370db --- /dev/null +++ b/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, + ) diff --git a/tests/unit/appproxy/coordinator/clients/worker/BUILD b/tests/unit/appproxy/coordinator/clients/worker/BUILD new file mode 100644 index 00000000000..57341b1358b --- /dev/null +++ b/tests/unit/appproxy/coordinator/clients/worker/BUILD @@ -0,0 +1,3 @@ +python_tests( + name="tests", +) diff --git a/tests/unit/appproxy/coordinator/clients/worker/test_client.py b/tests/unit/appproxy/coordinator/clients/worker/test_client.py new file mode 100644 index 00000000000..a062ace4f32 --- /dev/null +++ b/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()