mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
1587faeeb9
* test(unit): wait on the work itself instead of fixed sleeps in 21 flaky files Replace sleeps that raced background logging, cache writes, refreshes and thread work with either awaiting that work directly (logging worker flush, the batch logger's flush_queue, the tasks the call spawned, thread joins) or an Event the test's own fake sets when it runs. Each wait keeps a 5s hang guard that never decides pass or fail. Tests that need a fake clock, real sockets, or measure latency are left unchanged, and no production code moves * test(unit): await the tasks each completion spawns instead of yielding once The logging tests yielded once with asyncio.sleep(0) so the task that enqueues the success log could run before the worker flush. That relied on the enqueue happening before the task's first await. Snapshot the running tasks around each acompletion call and await the ones it spawned, so the guarantee holds however many steps sit before the enqueue
274 lines
10 KiB
Python
274 lines
10 KiB
Python
"""The /metrics app must render off the event loop, coalesce concurrent scrapes and stream chunks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import threading
|
|
import time
|
|
from collections.abc import Iterator, Mapping, Sequence
|
|
from typing import Final
|
|
|
|
import httpx
|
|
import pytest
|
|
from prometheus_client import CollectorRegistry, Gauge
|
|
from prometheus_client.metrics_core import GaugeMetricFamily
|
|
from prometheus_client.registry import Collector
|
|
|
|
from litellm.integrations.prometheus_metrics_endpoint import (
|
|
RESPONSE_CHUNK_SIZE_BYTES,
|
|
make_metrics_asgi_app,
|
|
)
|
|
|
|
_GATE_TIMEOUT_SECONDS: Final = 10.0
|
|
_SECOND_SCRAPE_SETTLE_SECONDS: Final = 0.2
|
|
|
|
|
|
class _SlowCollector(Collector):
|
|
"""Blocking collector standing in for a large registry render."""
|
|
|
|
def __init__(self, block_seconds: float, sample_count: int = 1) -> None:
|
|
self.block_seconds = block_seconds
|
|
self.sample_count = sample_count
|
|
self.collect_calls = 0
|
|
|
|
def collect(self) -> Iterator[GaugeMetricFamily]:
|
|
self.collect_calls += 1
|
|
time.sleep(self.block_seconds)
|
|
family: Final = GaugeMetricFamily("slow_metric", "slow", labels=("idx",))
|
|
for idx in range(self.sample_count):
|
|
family.add_metric((str(idx),), 1.0)
|
|
yield family
|
|
|
|
|
|
def _client(registry: CollectorRegistry) -> httpx.AsyncClient:
|
|
return httpx.AsyncClient(
|
|
transport=httpx.ASGITransport(app=make_metrics_asgi_app(registry)),
|
|
base_url="http://metrics.test",
|
|
)
|
|
|
|
|
|
def _registry_with(collector: Collector) -> CollectorRegistry:
|
|
registry: Final = CollectorRegistry()
|
|
registry.register(collector)
|
|
return registry
|
|
|
|
|
|
async def _scrape(client: httpx.AsyncClient, headers: Mapping[str, str] | None = None) -> httpx.Response:
|
|
return await client.get("/metrics", headers=headers)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_render_does_not_block_the_event_loop():
|
|
collector: Final = _GatedCollector()
|
|
async with _client(_registry_with(collector)) as client:
|
|
try:
|
|
scrape: Final = asyncio.create_task(_scrape(client))
|
|
assert await asyncio.to_thread(collector.started.wait, _GATE_TIMEOUT_SECONDS), "render never started"
|
|
assert not collector.finished.is_set(), "event loop was blocked while the registry was rendered"
|
|
collector.release.set()
|
|
response: Final = await scrape
|
|
finally:
|
|
collector.release.set()
|
|
|
|
assert b"gated_metric" in response.content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_identical_scrapes_share_one_render():
|
|
collector: Final = _SlowCollector(block_seconds=0.2)
|
|
async with _client(_registry_with(collector)) as client:
|
|
responses: Final[Sequence[httpx.Response]] = await asyncio.gather(*(_scrape(client) for _ in range(5)))
|
|
|
|
assert collector.collect_calls == 1
|
|
for response in responses:
|
|
assert b"slow_metric" in response.content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sequential_scrapes_are_rendered_fresh():
|
|
collector: Final = _SlowCollector(block_seconds=0.0)
|
|
async with _client(_registry_with(collector)) as client:
|
|
await _scrape(client)
|
|
await _scrape(client)
|
|
|
|
assert collector.collect_calls == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_gzip_is_used_when_the_scraper_accepts_it():
|
|
registry: Final = CollectorRegistry()
|
|
Gauge("plain_metric", "plain", registry=registry).set(1)
|
|
|
|
async with _client(registry) as client:
|
|
compressed: Final = await _scrape(client, headers={"accept-encoding": "gzip"})
|
|
plain: Final = await _scrape(client, headers={"accept-encoding": "identity"})
|
|
|
|
assert compressed.headers["content-encoding"] == "gzip"
|
|
assert "content-encoding" not in plain.headers
|
|
assert compressed.content == plain.content
|
|
assert b"plain_metric" in plain.content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_name_filter_restricts_the_rendered_registry():
|
|
registry: Final = CollectorRegistry()
|
|
Gauge("wanted_metric", "wanted", registry=registry).set(1)
|
|
Gauge("other_metric", "other", registry=registry).set(1)
|
|
|
|
async with _client(registry) as client:
|
|
response: Final = await client.get("/metrics", params={"name[]": "wanted_metric"})
|
|
|
|
assert b"wanted_metric" in response.content
|
|
assert b"other_metric" not in response.content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_large_payload_is_streamed_in_chunks():
|
|
registry: Final = _registry_with(_SlowCollector(block_seconds=0.0, sample_count=5000))
|
|
chunk_sizes: Final[list[int]] = [] # mutable-ok: records the ASGI body parts the app emitted
|
|
|
|
async def send(message: Mapping[str, object]) -> None:
|
|
if message["type"] == "http.response.body":
|
|
body = message["body"]
|
|
assert isinstance(body, bytes)
|
|
chunk_sizes.append(len(body))
|
|
|
|
incoming: Final = iter(({"type": "http.request", "body": b"", "more_body": False},))
|
|
|
|
async def receive() -> Mapping[str, object]:
|
|
request: Final = next(incoming, None)
|
|
if request is not None:
|
|
return request
|
|
await asyncio.Event().wait()
|
|
return {"type": "http.disconnect"}
|
|
|
|
app: Final = make_metrics_asgi_app(registry)
|
|
await app(
|
|
{
|
|
"type": "http",
|
|
"method": "GET",
|
|
"path": "/metrics",
|
|
"headers": (),
|
|
"query_string": b"",
|
|
},
|
|
receive,
|
|
send,
|
|
)
|
|
|
|
assert sum(chunk_sizes) > RESPONSE_CHUNK_SIZE_BYTES
|
|
assert len(chunk_sizes) > 2
|
|
assert max(chunk_sizes) <= RESPONSE_CHUNK_SIZE_BYTES
|
|
|
|
|
|
class _GatedCollector(Collector):
|
|
"""Blocking collector that parks in the worker thread until the test releases it."""
|
|
|
|
def __init__(self) -> None:
|
|
self.started = threading.Event()
|
|
self.release = threading.Event()
|
|
self.finished = threading.Event()
|
|
self._lock = threading.Lock()
|
|
self.collect_calls = 0
|
|
|
|
def collect(self) -> Iterator[GaugeMetricFamily]:
|
|
with self._lock:
|
|
self.collect_calls += 1
|
|
self.started.set()
|
|
self.release.wait(timeout=_GATE_TIMEOUT_SECONDS)
|
|
self.finished.set()
|
|
family: Final = GaugeMetricFamily("gated_metric", "gated")
|
|
family.add_metric((), 1.0)
|
|
yield family
|
|
|
|
|
|
async def _scrape_pair_concurrently(
|
|
registry: CollectorRegistry, collector: _GatedCollector, headers: Sequence[Mapping[str, str]]
|
|
) -> Sequence[httpx.Response]:
|
|
"""Issue the second scrape only once the first one's render is parked inside the worker thread."""
|
|
async with _client(registry) as client:
|
|
try:
|
|
first: Final = asyncio.create_task(_scrape(client, headers=headers[0]))
|
|
assert await asyncio.to_thread(collector.started.wait, _GATE_TIMEOUT_SECONDS), "first render never started"
|
|
second: Final = asyncio.create_task(_scrape(client, headers=headers[1]))
|
|
await asyncio.sleep(_SECOND_SCRAPE_SETTLE_SECONDS)
|
|
collector.release.set()
|
|
return await asyncio.gather(first, second)
|
|
finally:
|
|
collector.release.set()
|
|
|
|
|
|
@pytest.mark.parametrize("reverse", (False, True), ids=("as-listed", "reversed"))
|
|
@pytest.mark.parametrize(
|
|
"spellings",
|
|
(
|
|
({"accept-encoding": "gzip"}, {"accept-encoding": "gzip, deflate"}),
|
|
({"accept": "*/*"}, {"accept": "text/plain;version=0.0.4;q=0.5,*/*;q=0.1"}),
|
|
),
|
|
ids=("accept-encoding", "accept"),
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_header_spellings_with_the_same_output_share_one_render(
|
|
spellings: Sequence[Mapping[str, str]], reverse: bool
|
|
):
|
|
collector: Final = _GatedCollector()
|
|
ordered: Final = tuple(reversed(spellings)) if reverse else spellings
|
|
|
|
responses: Final = await _scrape_pair_concurrently(_registry_with(collector), collector, ordered)
|
|
|
|
assert collector.collect_calls == 1, "the second scrape rendered the registry again instead of joining the first"
|
|
for response in responses:
|
|
assert b"gated_metric" in response.content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_different_output_formats_are_rendered_separately():
|
|
collector: Final = _GatedCollector()
|
|
|
|
responses: Final = await _scrape_pair_concurrently(
|
|
_registry_with(collector),
|
|
collector,
|
|
({"accept": "text/plain"}, {"accept": "application/openmetrics-text"}),
|
|
)
|
|
|
|
assert collector.collect_calls == 2, "scrapes wanting different exposition formats must not share a render"
|
|
assert responses[0].headers["content-type"] != responses[1].headers["content-type"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_gzip_and_plain_scrapes_each_get_their_own_encoding():
|
|
collector: Final = _GatedCollector()
|
|
|
|
responses: Final = await _scrape_pair_concurrently(
|
|
_registry_with(collector),
|
|
collector,
|
|
({"accept-encoding": "gzip"}, {"accept-encoding": "identity"}),
|
|
)
|
|
|
|
assert collector.collect_calls == 2, "scrapes wanting different content encodings must not share a render"
|
|
assert responses[0].headers["content-encoding"] == "gzip"
|
|
assert "content-encoding" not in responses[1].headers
|
|
for response in responses:
|
|
assert b"gated_metric" in response.content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_finishing_render_does_not_evict_another_that_is_still_in_flight():
|
|
collector: Final = _GatedCollector()
|
|
async with _client(_registry_with(collector)) as client:
|
|
try:
|
|
parked: Final = asyncio.create_task(_scrape(client))
|
|
assert await asyncio.to_thread(collector.started.wait, _GATE_TIMEOUT_SECONDS), "first render never started"
|
|
|
|
unrelated: Final = await client.get("/metrics", params={"name[]": "no_such_metric"})
|
|
assert unrelated.status_code == 200
|
|
|
|
joiner: Final = asyncio.create_task(_scrape(client))
|
|
await asyncio.sleep(_SECOND_SCRAPE_SETTLE_SECONDS)
|
|
collector.release.set()
|
|
responses: Final = await asyncio.gather(parked, joiner)
|
|
finally:
|
|
collector.release.set()
|
|
|
|
assert collector.collect_calls == 1, "an unrelated render finishing evicted the render still in flight"
|
|
for response in responses:
|
|
assert b"gated_metric" in response.content
|