diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index 6f407295259..b8ea11f3d2f 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -4,6 +4,7 @@ import os import time import uuid from collections.abc import Callable, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack, contextmanager from dataclasses import dataclass from hashlib import sha256 @@ -15,6 +16,7 @@ from pydantic import JsonValue, TypeAdapter from tests.integration._support.database import read_rows JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +DEPLOYMENT_IDS: Final = TypeAdapter(dict[str, list[str]]) GATEWAY_LIMITS: Final = httpx.Limits(keepalive_expiry=2) T = TypeVar("T") @@ -96,6 +98,7 @@ class Gateway: key: str upstream_url: str database_url: str | None = None + workers: int = 1 def rows(self, query: str, parameters: tuple[str, ...]) -> list[dict[str, JsonValue]]: return read_rows(query, parameters, database_url=self.database_url) @@ -294,8 +297,36 @@ class Scenario: ) identity: Final = string_value(object_value(created["model_info"])["id"]) self.cleanups.callback(self.delete_model, identity) + if self.gateway.workers > 1: + self._await_every_worker(name, identity) return name + def _await_every_worker(self, name: str, identity: str) -> None: + from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS + + eventually( + lambda: self._deployment_ids_by_worker(name), + lambda answers: len({pid for pid, _ in answers}) >= self.gateway.workers + and all(identity in model_ids for _, model_ids in answers), + seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + ) + + def _deployment_ids_by_worker(self, name: str) -> tuple[tuple[str, frozenset[str]], ...]: + with ThreadPoolExecutor(max_workers=8) as pool: + return tuple(pool.map(lambda _: self._deployment_ids_on_a_worker(name), range(16))) + + def _deployment_ids_on_a_worker(self, name: str) -> tuple[str, frozenset[str]]: + response: Final = self.gateway.request( + "GET", + "/health/integration-deployment-ids", + params={"model_name": name}, + headers={"connection": "close"}, + ) + assert response.status_code == 200, f"Deployment ids on a worker: {response.status_code} {response.text}" + return response.headers["x-integration-worker-pid"], frozenset( + DEPLOYMENT_IDS.validate_json(response.content)["model_ids"] + ) + @contextmanager def gateway_from_environment() -> Iterator[Gateway]: diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index b08037e80db..052c52982b6 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -893,7 +893,7 @@ def owned_proxy_process( limits=GATEWAY_LIMITS, ) as client: yield OwnedProxy( - Gateway(client, gateway.key, gateway.upstream_url, owned.database_url), + Gateway(client, gateway.key, gateway.upstream_url, owned.database_url, workers), process, launch.log, healthcheck, diff --git a/tests/unit/integration_support/test_client.py b/tests/unit/integration_support/test_client.py new file mode 100644 index 00000000000..909f69f4f3d --- /dev/null +++ b/tests/unit/integration_support/test_client.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import itertools +import threading +import time +from collections.abc import Iterator +from contextlib import ExitStack +from dataclasses import dataclass, field +from typing import Final + +import httpx + +from tests.integration._support.client import Gateway, Scenario + +DEPLOYMENT: Final = "deployment-under-test" +CREATING_WORKER: Final = "101" +LAGGING_WORKER: Final = "202" + + +@dataclass(frozen=True, slots=True) +class LaggingProxy: + reachable_at: float + synced_at: float + arrivals: Iterator[int] = field(default_factory=itertools.count) + created: threading.Event = field(default_factory=threading.Event) + + def _worker(self) -> str: + arrival: Final = next(self.arrivals) + reachable: Final = time.monotonic() >= self.reachable_at + return LAGGING_WORKER if reachable and arrival % 2 else CREATING_WORKER + + def respond(self, request: httpx.Request) -> httpx.Response: + if request.url.path == "/model/new": + self.created.set() + return httpx.Response(200, json={"model_info": {"id": DEPLOYMENT}}) + worker: Final = self._worker() + serves: Final = self.created.is_set() and (worker == CREATING_WORKER or time.monotonic() >= self.synced_at) + return httpx.Response( + 200, + json={"model_ids": [DEPLOYMENT] if serves else []}, + headers={"x-integration-worker-pid": worker}, + ) + + +def _created_model(proxy: LaggingProxy, workers: int) -> str: + with ( + httpx.Client(base_url="http://127.0.0.1:4000", transport=httpx.MockTransport(proxy.respond)) as client, + ExitStack() as cleanups, + ): + name: Final = Scenario(Gateway(client, "sk-test", "http://127.0.0.1:4001", workers=workers), cleanups).model() + cleanups.pop_all() + return name + + +def test_model_on_multi_worker_proxy_returns_only_once_every_worker_serves_it() -> None: + started: Final = time.monotonic() + proxy: Final = LaggingProxy(reachable_at=started + 0.3, synced_at=started + 0.8) + + _created_model(proxy, workers=2) + + assert time.monotonic() >= proxy.synced_at + + +def test_model_on_single_worker_proxy_returns_without_probing_workers() -> None: + started: Final = time.monotonic() + proxy: Final = LaggingProxy(reachable_at=started, synced_at=started + 60) + + _created_model(proxy, workers=1) + + assert time.monotonic() < proxy.synced_at