test(integration): wait for every worker in Scenario.model (#47277)

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]:

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,

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