fix: cap cold MCP schema workers at usable CPUs and deflake wall-clock unit tests (#45791)

* test: inject clocks and wait on conditions in three wall-clock-flaky unit tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test: keep generic API logger tests on 3.10 and cover automatic sending

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test: pin router usage minute in shared rate-limit header test

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test: stabilize three wall-clock-sensitive unit tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(mcp): cap cold schema validation workers at the CPUs the proxy can use

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test: inject the cache clock and wait on lock contention instead of 1s deadlines

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(mcp): wait on listing events instead of 1s deadlines in the discovery drain test

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: mateo <mateo@berri.ai>
This commit is contained in:
devin-ai-integration[bot] 2026-10-15 00:47:59 -07:00 • committed by GitHub
parent de0ed1557e
commit bf56719f09
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 166 additions and 44 deletions

View file

@ -2,6 +2,7 @@ from __future__ import annotations
import hashlib
import json
import os
import re
from collections import deque
from collections.abc import Iterator, Mapping, Sequence
@ -55,6 +56,7 @@ _MAX_VALIDATION_DEPTH: Final = 64
_MAX_VALIDATION_NODES: Final = 10_000
_MAX_VALIDATION_CHARACTERS: Final = 1_048_576
_VALIDATION_TIMEOUT_SECONDS: Final = 30
_MAX_VALIDATION_WORKERS: Final = 4
def _validation_nodes(value: JsonValue | Mapping[str, JsonValue]) -> Iterator[tuple[int, int]]:
@ -101,6 +103,11 @@ def _validate_tool_arguments(schema: Mapping[str, JsonValue], arguments: Mapping
return None
def _schema_validation_worker_count() -> int:
available: Final = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count() or 1
return max(1, min(_MAX_VALIDATION_WORKERS, available))
async def _tool_argument_validation_error(
schema: Mapping[str, JsonValue],
arguments: Mapping[str, JsonValue],
@ -110,7 +117,7 @@ async def _tool_argument_validation_error(
if limit_error is not None:
return limit_error
existing: Final = _SCHEMA_VALIDATION_LIMITER.get(None)
limiter: Final = existing if existing is not None else anyio.CapacityLimiter(4)
limiter: Final = existing if existing is not None else anyio.CapacityLimiter(_schema_validation_worker_count())
if existing is None:
_SCHEMA_VALIDATION_LIMITER.set(limiter)
try:

View file

@ -6,6 +6,7 @@ from __future__ import annotations
import asyncio
import hashlib
import json
from datetime import datetime, timezone
from itertools import chain
from typing import Any, Final
from unittest.mock import AsyncMock, MagicMock
@ -837,7 +838,12 @@ async def test_the_prefetch_reads_cooldowns_again_once_the_read_interval_elapses
@pytest.mark.asyncio
async def test_the_prefetch_mget_carries_only_the_keys_whose_read_is_due(monkeypatch):
first_time: Final = 1_000_000.0
pinned: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc)
monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time)
monkeypatch.setattr(
"litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime",
lambda: pinned,
)
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
router = _router(redis_cache, routing_strategy="usage-based-routing-v2")

View file

@ -443,6 +443,42 @@ async def test_generic_api_callback_sumologic_uses_ndjson():
assert len(mine) == 2, f"Expected this test's 2 calls as NDJSON lines, got {len(mine)} of {len(records)}"
@pytest.mark.asyncio
async def test_generic_api_callback_sends_queued_events_on_its_own(monkeypatch: pytest.MonkeyPatch):
mock_post = AsyncMock()
mock_post.return_value.status_code = 200
mock_post.return_value.text = "OK"
test_endpoint = "https://example.com/api/logs"
test_headers = {"Authorization": "Bearer test_token"}
monkeypatch.setenv("GENERIC_LOGGER_ENDPOINT", test_endpoint)
generic_logger = GenericAPILogger(endpoint=test_endpoint, headers=test_headers, flush_interval=0.05)
generic_logger.async_httpx_client.post = mock_post
monkeypatch.setattr(litellm, "callbacks", [generic_logger])
for i in range(3):
await litellm.acompletion(
model="gpt-5.5",
messages=[{"role": "user", "content": f"Hello, world! {i}"}],
mock_response="hi",
user="test_user",
)
async def _poll() -> None:
while sum(len(json.loads(call.kwargs["data"])) for call in mock_post.call_args_list) < 3:
await asyncio.sleep(0.01)
await asyncio.wait_for(_poll(), timeout=10)
posted_bodies: Final = [json.loads(call.kwargs["data"]) for call in mock_post.call_args_list]
for body in posted_bodies:
assert isinstance(body, list), "Each posted body should be a JSON array"
posted_items: Final = sum(posted_bodies, [])
contents: Final = sorted(item["messages"][0]["content"] for item in posted_items)
assert contents == [f"Hello, world! {i}" for i in range(3)]
@pytest.mark.asyncio
async def test_generic_api_callback_invalid_log_format():
"""

View file

@ -1080,11 +1080,19 @@ def test_auth_check_and_project_id_make_one_round_trip_when_langfuse_is_down(
[(500, 2), (503, 2), (429, 1), (404, 1)],
ids=["http-500", "http-503", "http-429", "http-404"],
)
def test_cold_prompt_miss_never_sleeps_when_langfuse_is_down(status: int, round_trips: int):
def test_cold_prompt_miss_never_sleeps_when_langfuse_is_down(
status: int, round_trips: int, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A cold ``get_prompt`` fetches inline on the event loop; with the generated client's default retries a
429 carrying ``Retry-After: 30`` used to hold the loop for a minute. A 5xx gets the v2 client's one
quick retry, a 429 or 4xx none."""
requests: list[httpx.Request] = []
sleeps: Final[list[float]] = []
monkeypatch.setattr(
langfuse_http_client,
"time",
SimpleNamespace(sleep=sleeps.append, time=time.time),
)
def fail(request: httpx.Request) -> httpx.Response:
requests.append(request)
@ -1097,18 +1105,25 @@ def test_cold_prompt_miss_never_sleeps_when_langfuse_is_down(status: int, round_
httpx_client=httpx.Client(transport=httpx.MockTransport(fail)),
)
started = monotonic()
with pytest.raises(LangfusePromptError) as caught:
client.get_prompt("greeting")
assert len(requests) == round_trips
assert monotonic() - started < 0.5
assert sleeps == [], "a cold prompt miss should not sleep for REST retries"
assert caught.value.status_code == status
@pytest.mark.parametrize("first_failure", [503, "connect-error"], ids=["http-503", "connect-error"])
def test_one_transient_failure_on_a_cold_prompt_miss_does_not_fail_the_call(first_failure: int | str):
def test_one_transient_failure_on_a_cold_prompt_miss_does_not_fail_the_call(
first_failure: int | str, monkeypatch: pytest.MonkeyPatch
) -> None:
"""The v2 client retried a cold fetch once; a single Langfuse blip must not fail the LLM call."""
requests: list[httpx.Request] = []
sleeps: Final[list[float]] = []
monkeypatch.setattr(
langfuse_http_client,
"time",
SimpleNamespace(sleep=sleeps.append, time=time.time),
)
def flaky(request: httpx.Request) -> httpx.Response:
requests.append(request)
@ -1125,10 +1140,9 @@ def test_one_transient_failure_on_a_cold_prompt_miss_does_not_fail_the_call(firs
httpx_client=httpx.Client(transport=httpx.MockTransport(flaky)),
)
started = monotonic()
assert client.get_prompt("greeting").compile() == "hello"
assert len(requests) == 2
assert monotonic() - started < 0.5
assert sleeps == [], "one transient failure should not sleep for REST retries"
assert client.get_prompt("greeting").compile() == "hello", "the retried prompt is cached like any other"
assert len(requests) == 2

View file

@ -4121,6 +4121,18 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(
)
mcp_server._stateful_session_owners[session_id] = mcp_server._owner_fingerprint_for(owner_auth)
class _RecordingLock(asyncio.Lock):
def __init__(self) -> None:
super().__init__()
self.contended = asyncio.Event()
async def acquire(self) -> bool:
if self.locked():
self.contended.set()
return await super().acquire()
session_lock: Final = _RecordingLock()
mcp_server._stateful_session_locks[session_id] = session_lock
gate = asyncio.Event()
request_in_handle = asyncio.Event()
response_handled = asyncio.Event()
@ -4186,12 +4198,19 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(
),
):
req_task = asyncio.create_task(call(request_body))
await asyncio.wait_for(request_in_handle.wait(), timeout=1.0)
await request_in_handle.wait()
resp_task = asyncio.create_task(call(response_body))
# Under a flat substring heuristic the response would acquire the
# lock held by req_task and this wait would time out (deadlock).
await asyncio.wait_for(response_handled.wait(), timeout=1.0)
response_handled_task: Final = asyncio.create_task(response_handled.wait())
contended_task: Final = asyncio.create_task(session_lock.contended.wait())
_, pending = await asyncio.wait(
(response_handled_task, contended_task),
return_when=asyncio.FIRST_COMPLETED,
)
for event_task in pending:
event_task.cancel()
await asyncio.gather(*pending, return_exceptions=True)
assert response_handled.is_set(), "response POST contended for the request's session lock"
await resp_task
assert not req_task.done()
@ -4207,6 +4226,7 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(
finally:
gate.set()
mcp_server._stateful_session_auth_contexts.pop(session_id, None)
mcp_server._stateful_session_auth_context_last_seen.pop(session_id, None)
mcp_server._stateful_session_owners.pop(session_id, None)
mcp_server._stateful_session_locks.pop(session_id, None)
mcp_server._stateful_session_active_request_counts.pop(session_id, None)

View file

@ -1427,7 +1427,12 @@ async def test_mcp_tool_search_leaves_the_listed_tools_slot_empty(monkeypatch: p
@pytest.mark.asyncio
async def test_cold_concurrent_schema_validation_accepts_valid_arguments() -> None:
from litellm.proxy._experimental.mcp_server.tool_search import _tool_argument_validation_error
import psutil
from litellm.proxy._experimental.mcp_server.tool_search import (
_MAX_VALIDATION_WORKERS,
_tool_argument_validation_error,
)
schema: Final = {
"type": "object",
@ -1435,6 +1440,7 @@ async def test_cold_concurrent_schema_validation_accepts_valid_arguments() -> No
"properties": {"amount": {"$ref": "#/$defs/amount"}},
}
original_affinity: Final = os.sched_getaffinity(0) if sys.platform == "linux" else None
existing_children: Final = frozenset(child.pid for child in psutil.Process().children())
try:
if original_affinity is not None:
os.sched_setaffinity(0, {min(original_affinity)})
@ -1443,6 +1449,14 @@ async def test_cold_concurrent_schema_validation_accepts_valid_arguments() -> No
*(_tool_argument_validation_error(schema, {"amount": 0.75}, timeout_seconds=90.0) for _ in range(8))
)
assert results == [None] * 8, "Cold and warm workers must accept valid concurrent tool arguments"
workers: Final = tuple(
child
for child in psutil.Process().children()
if child.pid not in existing_children and "anyio.to_process" in child.cmdline()
)
assert 1 <= len(workers) <= (1 if original_affinity is not None else _MAX_VALIDATION_WORKERS), (
"Validation must not start more cold workers than the CPUs it can run on"
)
finally:
if original_affinity is not None:
os.sched_setaffinity(0, original_affinity)

View file

@ -1287,17 +1287,17 @@ async def test_discovery_concurrent_listings_drain_on_failure_and_cancellation(o
):
task = asyncio.create_task(GatewayOperations().execute(DiscoverRequest(), prepare_context(UserAPIKeyAuth(user_id="scoped"))))
try:
await asyncio.wait_for(asyncio.gather(*(event.wait() for event in ready)), 1)
await asyncio.gather(*(event.wait() for event in ready))
if outcome == "cancel":
task.cancel()
else:
release.set()
if outcome == "success":
result = await asyncio.wait_for(task, 1)
result = await task
assert result.capabilities.model_dump(exclude_none=True) == {}
else:
with pytest.raises(asyncio.CancelledError if outcome == "cancel" else ValueError):
await asyncio.wait_for(task, 1)
await task
assert all(event.is_set() for event in closed)
assert tools.call_args.kwargs["log_list_tools_to_spendlogs"] is False
finally:

View file

@ -499,12 +499,14 @@ async def test_native_count_finishing_after_quarter_worker_budget_keeps_plan_and
from litellm.litellm_core_utils import logging_worker
from litellm.litellm_core_utils.logging_worker import LoggingWorker
release: Final = asyncio.Event()
counted: Final = asyncio.Event()
loop: Final = asyncio.get_running_loop()
real_time: Final = loop.time
monkeypatch.setattr(loop, "time", lambda: real_time() + (2.05 if counted.is_set() else 0.0))
async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int:
if not release.is_set():
asyncio.get_running_loop().call_later(2.05, release.set)
await release.wait()
counted.set()
await asyncio.sleep(0.001)
return await _count(model, api_key, body)
worker: Final = LoggingWorker(timeout=8.0)
@ -519,7 +521,6 @@ async def test_native_count_finishing_after_quarter_worker_budget_keeps_plan_and
assert observed.plan is not None and observed.plan.breakpoints[0].prefix_tokens == 5000
assert payload["response_cost"] is not None and worker._timeout_total == 0
finally:
release.set()
await worker.stop()

View file

@ -48,15 +48,15 @@ async def test_fork_lookups_overlap_with_a_bounded_number_of_requests() -> None:
source: Final = GitLab(ROISettings(source_provider="gitlab"), httpx.MockTransport(respond))
pending: Final = asyncio.create_task(source.pulls("group/repo", date(2026, 9, 1), date(2026, 9, 30)))
try:
first_wave: Final = tuple([await asyncio.wait_for(started.get(), timeout=1) for _ in range(8)])
first_wave: Final = tuple([await started.get() for _ in range(8)])
assert len(set(first_wave)) == 8
assert started.empty()
release[first_wave[0] - 2].set()
next_id: Final = await asyncio.wait_for(started.get(), timeout=1)
next_id: Final = await started.get()
assert next_id not in first_wave
for event in release:
event.set()
pulls: Final = await asyncio.wait_for(pending, timeout=1)
pulls: Final = await pending
assert tuple(pull.head.repo.full_name if pull.head and pull.head.repo else None for pull in pulls) == tuple(
None if source_id == 3 else f"fork-{source_id}/repo" for source_id in source_ids
)

View file

@ -27,6 +27,17 @@ from litellm.proxy.utils import hash_token
from litellm.proxy.db.log_db_metrics import record_db_io
class _FakeClock:
def __init__(self) -> None:
self.now = 1000.0
def advance(self, seconds: float) -> None:
self.now += seconds
def __call__(self) -> float:
return self.now
def _digest_row(digest: str, key_alias: str | None, team_id: str | None, user_id: str | None) -> dict[str, str | None]:
return {"digest": digest, "key_alias": key_alias, "team_id": team_id, "user_id": user_id}
@ -422,34 +433,44 @@ async def test_recover_key_metadata_from_spend_logs_rescans_when_the_window_chan
@pytest.mark.asyncio
async def test_recover_key_metadata_from_spend_logs_retries_a_failed_query_only_after_the_miss_ttl():
digest = hash_token("cli-session-retry")
window = (datetime(2026, 9, 7), datetime(2026, 9, 10))
cache = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL)
mock_prisma = MagicMock()
query_raw = _spend_log_transaction(mock_prisma, AsyncMock(side_effect=PrismaError("statement timeout")))
started = time.time()
assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {}
query_raw = _spend_log_transaction(mock_prisma, _query_raw_spend_logs([_spend_log_row(digest, "back-online", None, None)]))
async def test_recover_key_metadata_from_spend_logs_retries_a_failed_query_only_after_the_miss_ttl() -> None:
digest: Final = hash_token("cli-session-retry")
window: Final = (datetime(2026, 9, 7), datetime(2026, 9, 10))
clock: Final = _FakeClock()
cache: Final = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL, clock=clock)
mock_prisma: Final = MagicMock()
query_raw: Final = _spend_log_transaction(
mock_prisma,
AsyncMock(
side_effect=[
PrismaError("statement timeout"),
[_spend_log_row(digest, "back-online", None, None)],
]
),
)
assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {}
query_raw.assert_not_awaited()
miss_key = next(key for key in cache.ttl_dict if digest in key and not key.endswith(":missed-before"))
assert cache.ttl_dict[miss_key] - started <= SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1
cache.ttl_dict[miss_key] = time.time() - 1
miss_key: Final = next(key for key in cache.ttl_dict if digest in key and not key.endswith(":missed-before"))
assert cache.ttl_dict[miss_key] == clock.now + SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL
result = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache)
assert await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache) == {}
assert query_raw.await_count == 1
clock.advance(SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1)
result: Final = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache)
assert result[digest]["key_alias"] == "back-online"
assert query_raw.await_count == 2
@pytest.mark.asyncio
async def test_recover_key_metadata_from_spend_logs_never_caches_repeated_query_failures_as_long_as_a_hit():
async def test_recover_key_metadata_from_spend_logs_never_caches_repeated_query_failures_as_long_as_a_hit() -> None:
digest: Final = hash_token("cli-session-repeated-failure")
window: Final = (datetime(2026, 9, 7), datetime(2026, 9, 10))
cache: Final = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL)
clock: Final = _FakeClock()
cache: Final = InMemoryCache(default_ttl=SPEND_LOG_KEY_METADATA_CACHE_TTL, clock=clock)
mock_prisma: Final = MagicMock()
_spend_log_transaction(
query_raw: Final = _spend_log_transaction(
mock_prisma,
AsyncMock(
side_effect=[
@ -461,17 +482,19 @@ async def test_recover_key_metadata_from_spend_logs_never_caches_repeated_query_
)
await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache)
miss_key: Final = next(key for key in cache.ttl_dict if digest in key and not key.endswith(":missed-before"))
cache.ttl_dict[miss_key] = time.time() - 1
second_query_started: Final = time.time()
assert cache.ttl_dict[miss_key] == clock.now + SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL
clock.advance(SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1)
await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache)
assert cache.ttl_dict[miss_key] - second_query_started <= SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1
cache.ttl_dict[miss_key] = time.time() - 1
assert query_raw.await_count == 2
assert cache.ttl_dict[miss_key] == clock.now + SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL
clock.advance(SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL + 1)
result: Final = await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=cache)
assert result[digest]["key_alias"] == "back-online"
assert query_raw.await_count == 3
@pytest.mark.asyncio

View file

@ -1071,6 +1071,7 @@ def _shared_redis_stub(store: dict) -> MagicMock:
@pytest.mark.asyncio
@pytest.mark.usefixtures("router_minute_pinned")
async def test_headers_on_fresh_worker_reflect_shared_redis_usage():
store: dict = {}
worker_a = _rpm_tpm_router("lit-3058-workers")