From cae91bb5610287a9ba6136994a4ba9834ff071c2 Mon Sep 17 00:00:00 2001 From: Gyubong Date: Tue, 20 Oct 2026 16:41:07 +0900 Subject: [PATCH 1/4] feat(BA-8642): add the coordinator client that pushes ingresses to api-delivery workers ApiDeliveryWorkerClient posts ingress upserts and removals to one worker's /v2/ingress with the API secret, and ApiDeliveryWorkerClientPool holds one client per worker address. Co-Authored-By: Claude Opus 5.5 --- .../appproxy/coordinator/errors/ingress.py | 30 ++++ .../appproxy/coordinator/ingress/__init__.py | 0 .../appproxy/coordinator/ingress/client.py | 89 ++++++++++++ tests/unit/appproxy/coordinator/ingress/BUILD | 3 + .../coordinator/ingress/test_client.py | 129 ++++++++++++++++++ 5 files changed, 251 insertions(+) create mode 100644 src/ai/backend/appproxy/coordinator/errors/ingress.py create mode 100644 src/ai/backend/appproxy/coordinator/ingress/__init__.py create mode 100644 src/ai/backend/appproxy/coordinator/ingress/client.py create mode 100644 tests/unit/appproxy/coordinator/ingress/BUILD create mode 100644 tests/unit/appproxy/coordinator/ingress/test_client.py 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..437c0c58c22 --- /dev/null +++ b/src/ai/backend/appproxy/coordinator/errors/ingress.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from typing import override + +from aiohttp import web + +from ai.backend.common.exception import ( + BackendAIError, + ErrorCode, + ErrorDetail, + ErrorDomain, + ErrorOperation, +) + + +class ApiDeliveryWorkerRequestError(BackendAIError, web.HTTPBadGateway): + """Raised when an api-delivery worker cannot be reached or answers a request with an error.""" + + error_type = ( + "https://api.backend.ai/probs/appproxy-coordinator/api-delivery-worker-request-failed" + ) + error_title = "A request to an api-delivery worker failed." + + @override + def error_code(self) -> ErrorCode: + return ErrorCode( + domain=ErrorDomain.APPPROXY, + operation=ErrorOperation.REQUEST, + error_detail=ErrorDetail.UNREACHABLE, + ) diff --git a/src/ai/backend/appproxy/coordinator/ingress/__init__.py b/src/ai/backend/appproxy/coordinator/ingress/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/src/ai/backend/appproxy/coordinator/ingress/client.py b/src/ai/backend/appproxy/coordinator/ingress/client.py new file mode 100644 index 00000000000..5430871242e --- /dev/null +++ b/src/ai/backend/appproxy/coordinator/ingress/client.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +from collections.abc import Sequence +from http import HTTPStatus + +import aiohttp + +from ai.backend.appproxy.common.config import SecretConfig +from ai.backend.appproxy.coordinator.errors.ingress import ApiDeliveryWorkerRequestError +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 ApiDeliveryWorkerClientPool: + """Holds one ApiDeliveryWorkerClient per api-delivery worker address.""" + + _secrets: SecretConfig + _clients: dict[str, ApiDeliveryWorkerClient] + + def __init__(self, secrets: SecretConfig) -> None: + self._secrets = secrets + self._clients = {} + + def load_client(self, addr: str) -> ApiDeliveryWorkerClient: + 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 = ApiDeliveryWorkerClient(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 ApiDeliveryWorkerClient: + """Sends ingress pushes to the API port of one api-delivery 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() + if resp.status >= HTTPStatus.BAD_REQUEST: + raise ApiDeliveryWorkerRequestError( + extra_msg=f"worker {self._addr} answered {resp.status}", + extra_data={ + "worker_addr": self._addr, + "status": resp.status, + "body": body.decode(errors="replace"), + }, + ) + except aiohttp.ClientError as e: + raise ApiDeliveryWorkerRequestError( + extra_msg=f"worker {self._addr} is unreachable", + extra_data={"worker_addr": self._addr}, + ) from e diff --git a/tests/unit/appproxy/coordinator/ingress/BUILD b/tests/unit/appproxy/coordinator/ingress/BUILD new file mode 100644 index 00000000000..57341b1358b --- /dev/null +++ b/tests/unit/appproxy/coordinator/ingress/BUILD @@ -0,0 +1,3 @@ +python_tests( + name="tests", +) diff --git a/tests/unit/appproxy/coordinator/ingress/test_client.py b/tests/unit/appproxy/coordinator/ingress/test_client.py new file mode 100644 index 00000000000..3a57bd86e00 --- /dev/null +++ b/tests/unit/appproxy/coordinator/ingress/test_client.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +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.errors.ingress import ApiDeliveryWorkerRequestError +from ai.backend.appproxy.coordinator.ingress.client import ( + ApiDeliveryWorkerClient, + ApiDeliveryWorkerClientPool, +) +from ai.backend.common.dto.appproxy.v2.ingress.types import ( + Ingress, + IngressAccess, + IngressAccessMode, + IngressAuthority, + IngressBackend, + IngressBackendID, + IngressFrontend, + IngressFrontendKind, + IngressID, + IngressProtocol, + IngressTargetGroup, + IngressTargetGroupID, + PortFrontend, +) + +API_SECRET = "api-secret" +WORKER_ADDR = "10.0.0.1:10201" + + +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 TestApiDeliveryWorkerClient: + async def test_push_carries_the_api_secret( + self, applying_session: MagicMock, secrets: SecretConfig, ingress: Ingress + ) -> None: + client = ApiDeliveryWorkerClient(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( + "status", + [ + HTTPStatus.BAD_REQUEST, + HTTPStatus.UNAUTHORIZED, + HTTPStatus.CONFLICT, + HTTPStatus.NOT_IMPLEMENTED, + ], + ids=lambda status: str(status.value), + ) + async def test_error_answer_raises_with_the_worker_status( + self, status: HTTPStatus, secrets: SecretConfig, ingress: Ingress + ) -> None: + session = _answering_session(status, b"rejected") + client = ApiDeliveryWorkerClient(session, WORKER_ADDR, secrets) + with pytest.raises(ApiDeliveryWorkerRequestError) as exc_info: + await client.upsert([ingress]) + assert exc_info.value.extra_data == { + "worker_addr": WORKER_ADDR, + "status": status, + "body": "rejected", + } + + async def test_unreachable_worker_raises( + self, unreachable_session: MagicMock, secrets: SecretConfig, ingress: Ingress + ) -> None: + client = ApiDeliveryWorkerClient(unreachable_session, WORKER_ADDR, secrets) + with pytest.raises(ApiDeliveryWorkerRequestError): + await client.upsert([ingress]) + + +class TestApiDeliveryWorkerClientPool: + async def test_client_of_an_address_is_reused(self, secrets: SecretConfig) -> None: + pool = ApiDeliveryWorkerClientPool(secrets) + try: + assert pool.load_client(WORKER_ADDR) is pool.load_client(WORKER_ADDR) + finally: + await pool.close() From 33340263308649c06e7e78202d579d3b9080374c Mon Sep 17 00:00:00 2001 From: Gyubong Date: Tue, 20 Oct 2026 16:41:29 +0900 Subject: [PATCH 2/4] chore: add news fragment for PR #15915 Co-Authored-By: Claude Opus 5.5 --- changes/15915.feature.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/15915.feature.md 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 From 0620d974e5f6dd68da0df0fe5b8fb417eeb92bc4 Mon Sep 17 00:00:00 2001 From: Gyubong Date: Tue, 20 Oct 2026 16:50:17 +0900 Subject: [PATCH 3/4] refactor(BA-8642): tell an unreachable worker from its 4xx and 5xx answers A 4xx answer raises ApiDeliveryWorkerRejectedError (400) so the caller does not retry it, a 5xx raises ApiDeliveryWorkerServerError (502), and a connection failure raises ApiDeliveryWorkerUnreachableError (502). Co-Authored-By: Claude Opus 5.5 --- .../appproxy/coordinator/errors/ingress.py | 40 ++++++++++++++-- .../appproxy/coordinator/ingress/client.py | 33 ++++++++----- .../coordinator/ingress/test_client.py | 48 ++++++++++++++----- 3 files changed, 92 insertions(+), 29 deletions(-) diff --git a/src/ai/backend/appproxy/coordinator/errors/ingress.py b/src/ai/backend/appproxy/coordinator/errors/ingress.py index 437c0c58c22..8073d0370db 100644 --- a/src/ai/backend/appproxy/coordinator/errors/ingress.py +++ b/src/ai/backend/appproxy/coordinator/errors/ingress.py @@ -13,18 +13,48 @@ ) -class ApiDeliveryWorkerRequestError(BackendAIError, web.HTTPBadGateway): - """Raised when an api-delivery worker cannot be reached or answers a request with an error.""" +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-request-failed" + "https://api.backend.ai/probs/appproxy-coordinator/api-delivery-worker-server-error" ) - error_title = "A request to an api-delivery worker failed." + 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.UNREACHABLE, + error_detail=ErrorDetail.INTERNAL_ERROR, ) diff --git a/src/ai/backend/appproxy/coordinator/ingress/client.py b/src/ai/backend/appproxy/coordinator/ingress/client.py index 5430871242e..73b855a4374 100644 --- a/src/ai/backend/appproxy/coordinator/ingress/client.py +++ b/src/ai/backend/appproxy/coordinator/ingress/client.py @@ -1,12 +1,15 @@ from __future__ import annotations from collections.abc import Sequence -from http import HTTPStatus import aiohttp from ai.backend.appproxy.common.config import SecretConfig -from ai.backend.appproxy.coordinator.errors.ingress import ApiDeliveryWorkerRequestError +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, @@ -73,17 +76,23 @@ async def _post(self, path: str, request: BaseRequestModel) -> None: headers={_TOKEN_HEADER: self._secrets.api_secret}, ) as resp: body = await resp.read() - if resp.status >= HTTPStatus.BAD_REQUEST: - raise ApiDeliveryWorkerRequestError( - extra_msg=f"worker {self._addr} answered {resp.status}", - extra_data={ - "worker_addr": self._addr, - "status": resp.status, - "body": body.decode(errors="replace"), - }, - ) + 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 ApiDeliveryWorkerRequestError( + raise ApiDeliveryWorkerUnreachableError( extra_msg=f"worker {self._addr} is unreachable", extra_data={"worker_addr": self._addr}, ) from e diff --git a/tests/unit/appproxy/coordinator/ingress/test_client.py b/tests/unit/appproxy/coordinator/ingress/test_client.py index 3a57bd86e00..e4cc4871a38 100644 --- a/tests/unit/appproxy/coordinator/ingress/test_client.py +++ b/tests/unit/appproxy/coordinator/ingress/test_client.py @@ -1,5 +1,6 @@ from __future__ import annotations +from dataclasses import dataclass from http import HTTPStatus from unittest.mock import AsyncMock, MagicMock @@ -7,7 +8,11 @@ import pytest from ai.backend.appproxy.common.config import SecretConfig -from ai.backend.appproxy.coordinator.errors.ingress import ApiDeliveryWorkerRequestError +from ai.backend.appproxy.coordinator.errors.ingress import ( + ApiDeliveryWorkerRejectedError, + ApiDeliveryWorkerServerError, + ApiDeliveryWorkerUnreachableError, +) from ai.backend.appproxy.coordinator.ingress.client import ( ApiDeliveryWorkerClient, ApiDeliveryWorkerClientPool, @@ -27,11 +32,18 @@ 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 @@ -90,25 +102,37 @@ async def test_push_carries_the_api_secret( } @pytest.mark.parametrize( - "status", + "case", [ - HTTPStatus.BAD_REQUEST, - HTTPStatus.UNAUTHORIZED, - HTTPStatus.CONFLICT, - HTTPStatus.NOT_IMPLEMENTED, + _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 status: str(status.value), + ids=lambda case: str(case.status.value), ) async def test_error_answer_raises_with_the_worker_status( - self, status: HTTPStatus, secrets: SecretConfig, ingress: Ingress + self, case: _ErrorAnswerCase, secrets: SecretConfig, ingress: Ingress ) -> None: - session = _answering_session(status, b"rejected") + session = _answering_session(case.status, b"rejected") client = ApiDeliveryWorkerClient(session, WORKER_ADDR, secrets) - with pytest.raises(ApiDeliveryWorkerRequestError) as exc_info: + with pytest.raises(case.expected_error) as exc_info: await client.upsert([ingress]) assert exc_info.value.extra_data == { "worker_addr": WORKER_ADDR, - "status": status, + "status": case.status, "body": "rejected", } @@ -116,7 +140,7 @@ async def test_unreachable_worker_raises( self, unreachable_session: MagicMock, secrets: SecretConfig, ingress: Ingress ) -> None: client = ApiDeliveryWorkerClient(unreachable_session, WORKER_ADDR, secrets) - with pytest.raises(ApiDeliveryWorkerRequestError): + with pytest.raises(ApiDeliveryWorkerUnreachableError): await client.upsert([ingress]) From fd12686f750293064da50847f1059d4bbf1f44fa Mon Sep 17 00:00:00 2001 From: Gyubong Date: Tue, 20 Oct 2026 16:59:45 +0900 Subject: [PATCH 4/4] refactor(BA-8642): move the worker client to clients/worker as AppProxyWorkerClient Co-Authored-By: Claude Opus 5.5 --- .../{ingress => clients}/__init__.py | 0 .../coordinator/clients/worker/__init__.py | 0 .../{ingress => clients/worker}/client.py | 14 ++++++------- .../{ingress => clients/worker}/BUILD | 0 .../worker}/test_client.py | 20 +++++++++---------- 5 files changed, 17 insertions(+), 17 deletions(-) rename src/ai/backend/appproxy/coordinator/{ingress => clients}/__init__.py (100%) create mode 100644 src/ai/backend/appproxy/coordinator/clients/worker/__init__.py rename src/ai/backend/appproxy/coordinator/{ingress => clients/worker}/client.py (88%) rename tests/unit/appproxy/coordinator/{ingress => clients/worker}/BUILD (100%) rename tests/unit/appproxy/coordinator/{ingress => clients/worker}/test_client.py (90%) diff --git a/src/ai/backend/appproxy/coordinator/ingress/__init__.py b/src/ai/backend/appproxy/coordinator/clients/__init__.py similarity index 100% rename from src/ai/backend/appproxy/coordinator/ingress/__init__.py rename to src/ai/backend/appproxy/coordinator/clients/__init__.py 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/ingress/client.py b/src/ai/backend/appproxy/coordinator/clients/worker/client.py similarity index 88% rename from src/ai/backend/appproxy/coordinator/ingress/client.py rename to src/ai/backend/appproxy/coordinator/clients/worker/client.py index 73b855a4374..ee23f41f553 100644 --- a/src/ai/backend/appproxy/coordinator/ingress/client.py +++ b/src/ai/backend/appproxy/coordinator/clients/worker/client.py @@ -20,24 +20,24 @@ _TOKEN_HEADER = "X-BackendAI-Token" -class ApiDeliveryWorkerClientPool: - """Holds one ApiDeliveryWorkerClient per api-delivery worker address.""" +class AppProxyWorkerClientPool: + """Holds one AppProxyWorkerClient per worker address.""" _secrets: SecretConfig - _clients: dict[str, ApiDeliveryWorkerClient] + _clients: dict[str, AppProxyWorkerClient] def __init__(self, secrets: SecretConfig) -> None: self._secrets = secrets self._clients = {} - def load_client(self, addr: str) -> ApiDeliveryWorkerClient: + 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 = ApiDeliveryWorkerClient(session, addr, self._secrets) + client = AppProxyWorkerClient(session, addr, self._secrets) self._clients[addr] = client return client @@ -47,8 +47,8 @@ async def close(self) -> None: self._clients.clear() -class ApiDeliveryWorkerClient: - """Sends ingress pushes to the API port of one api-delivery worker node.""" +class AppProxyWorkerClient: + """Sends ingress pushes to the API port of one worker node.""" _session: aiohttp.ClientSession _addr: str diff --git a/tests/unit/appproxy/coordinator/ingress/BUILD b/tests/unit/appproxy/coordinator/clients/worker/BUILD similarity index 100% rename from tests/unit/appproxy/coordinator/ingress/BUILD rename to tests/unit/appproxy/coordinator/clients/worker/BUILD diff --git a/tests/unit/appproxy/coordinator/ingress/test_client.py b/tests/unit/appproxy/coordinator/clients/worker/test_client.py similarity index 90% rename from tests/unit/appproxy/coordinator/ingress/test_client.py rename to tests/unit/appproxy/coordinator/clients/worker/test_client.py index e4cc4871a38..a062ace4f32 100644 --- a/tests/unit/appproxy/coordinator/ingress/test_client.py +++ b/tests/unit/appproxy/coordinator/clients/worker/test_client.py @@ -8,15 +8,15 @@ 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.appproxy.coordinator.ingress.client import ( - ApiDeliveryWorkerClient, - ApiDeliveryWorkerClientPool, -) from ai.backend.common.dto.appproxy.v2.ingress.types import ( Ingress, IngressAccess, @@ -91,11 +91,11 @@ def unreachable_session() -> MagicMock: return session -class TestApiDeliveryWorkerClient: +class TestAppProxyWorkerClient: async def test_push_carries_the_api_secret( self, applying_session: MagicMock, secrets: SecretConfig, ingress: Ingress ) -> None: - client = ApiDeliveryWorkerClient(applying_session, WORKER_ADDR, secrets) + client = AppProxyWorkerClient(applying_session, WORKER_ADDR, secrets) await client.upsert([ingress]) assert applying_session.post.call_args.kwargs["headers"] == { "X-BackendAI-Token": API_SECRET @@ -127,7 +127,7 @@ 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 = ApiDeliveryWorkerClient(session, WORKER_ADDR, secrets) + 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 == { @@ -139,14 +139,14 @@ async def test_error_answer_raises_with_the_worker_status( async def test_unreachable_worker_raises( self, unreachable_session: MagicMock, secrets: SecretConfig, ingress: Ingress ) -> None: - client = ApiDeliveryWorkerClient(unreachable_session, WORKER_ADDR, secrets) + client = AppProxyWorkerClient(unreachable_session, WORKER_ADDR, secrets) with pytest.raises(ApiDeliveryWorkerUnreachableError): await client.upsert([ingress]) -class TestApiDeliveryWorkerClientPool: +class TestAppProxyWorkerClientPool: async def test_client_of_an_address_is_reused(self, secrets: SecretConfig) -> None: - pool = ApiDeliveryWorkerClientPool(secrets) + pool = AppProxyWorkerClientPool(secrets) try: assert pool.load_client(WORKER_ADDR) is pool.load_client(WORKER_ADDR) finally: