mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-02 07:13:02 +00:00
ee39de7107
* fix(scheduler): observe shared queue expiry before polling
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(scheduler): defer optional Redis exception import
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* perf(scheduler): share one Redis queue read per polling interval across waiters
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* revert(scheduler): drop shared Redis queue read
This reverts commit 5fd9b44a85
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>
491 lines
18 KiB
Python
491 lines
18 KiB
Python
# What is this?
|
|
## Unit tests for the Scheduler.py (workload prioritization scheduler)
|
|
|
|
import asyncio
|
|
import importlib
|
|
import json
|
|
import os
|
|
import time
|
|
from collections.abc import Sequence
|
|
from typing import Final
|
|
|
|
import pytest
|
|
import redis
|
|
import redis.asyncio
|
|
from fakeredis import FakeRedis, FakeServer
|
|
from fakeredis.aioredis import FakeConnection
|
|
from redis.asyncio import BlockingConnectionPool
|
|
|
|
import litellm
|
|
from litellm.caching.in_memory_cache import InMemoryCache
|
|
from litellm.caching.redis_cache import RedisCache
|
|
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
|
from litellm.scheduler import FlowItem, Scheduler, SchedulerCacheKeys
|
|
from litellm.utils import _invalidate_model_cost_lowercase_map
|
|
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scheduler_diff_model_names():
|
|
"""
|
|
Assert 2 requests to 2 diff model groups are top of their respective queue's
|
|
"""
|
|
scheduler = Scheduler()
|
|
|
|
item1 = FlowItem(priority=0, request_id="10", model_name="gpt-3.5-turbo")
|
|
item2 = FlowItem(priority=0, request_id="11", model_name="gpt-4")
|
|
await scheduler.add_request(item1)
|
|
await scheduler.add_request(item2)
|
|
|
|
assert await scheduler.poll(request=item1, health_deployments=[{"key": "value"}]) == True
|
|
assert await scheduler.poll(request=item2, health_deployments=[{"key": "value"}]) == True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scheduler_poll_persists_queue_to_cache():
|
|
class StubRedisCache:
|
|
def __init__(self):
|
|
self.store = {}
|
|
|
|
async def async_get_cache(self, key, **kwargs):
|
|
return self.store.get(key)
|
|
|
|
async def async_set_cache(self, key, value, **kwargs):
|
|
self.store[key] = value
|
|
|
|
redis_cache = StubRedisCache()
|
|
scheduler = Scheduler(redis_cache=redis_cache)
|
|
|
|
item1 = FlowItem(priority=0, request_id="10", model_name="gpt-3.5-turbo")
|
|
item2 = FlowItem(priority=0, request_id="11", model_name="gpt-3.5-turbo")
|
|
await scheduler.add_request(item1)
|
|
await scheduler.add_request(item2)
|
|
|
|
await scheduler.poll(request=item1, health_deployments=[])
|
|
|
|
queue_key = f"{SchedulerCacheKeys.queue.value}:{item1.model_name}"
|
|
updated_queue = redis_cache.store[queue_key]
|
|
assert updated_queue[0][1] == "11"
|
|
|
|
|
|
@pytest.mark.parametrize("p0, p1", [(0, 0), (0, 1), (1, 0)])
|
|
@pytest.mark.parametrize("healthy_deployments", [[{"key": "value"}], []])
|
|
@pytest.mark.asyncio
|
|
async def test_scheduler_prioritized_requests(p0, p1, healthy_deployments):
|
|
"""
|
|
2 requests for same model group
|
|
"""
|
|
scheduler = Scheduler()
|
|
|
|
item1 = FlowItem(priority=p0, request_id="10", model_name="gpt-3.5-turbo")
|
|
item2 = FlowItem(priority=p1, request_id="11", model_name="gpt-3.5-turbo")
|
|
await scheduler.add_request(item1)
|
|
await scheduler.add_request(item2)
|
|
|
|
if p0 == 0:
|
|
assert (
|
|
await scheduler.peek(
|
|
id="10",
|
|
model_name="gpt-3.5-turbo",
|
|
health_deployments=healthy_deployments,
|
|
)
|
|
== True
|
|
), "queue={}".format(await scheduler.get_queue(model_name="gpt-3.5-turbo"))
|
|
assert (
|
|
await scheduler.peek(
|
|
id="11",
|
|
model_name="gpt-3.5-turbo",
|
|
health_deployments=healthy_deployments,
|
|
)
|
|
== False
|
|
)
|
|
else:
|
|
assert (
|
|
await scheduler.peek(
|
|
id="11",
|
|
model_name="gpt-3.5-turbo",
|
|
health_deployments=healthy_deployments,
|
|
)
|
|
== True
|
|
)
|
|
assert (
|
|
await scheduler.peek(
|
|
id="10",
|
|
model_name="gpt-3.5-turbo",
|
|
health_deployments=healthy_deployments,
|
|
)
|
|
== False
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scheduler_queue_cleanup_on_timeout():
|
|
"""
|
|
Test that a timed-out request is properly removed from the queue.
|
|
This prevents memory leaks from accumulating timed-out requests.
|
|
"""
|
|
scheduler = Scheduler()
|
|
|
|
# Add multiple requests with different priorities
|
|
item1 = FlowItem(priority=0, request_id="req-0", model_name="gpt-3.5-turbo")
|
|
item2 = FlowItem(priority=1, request_id="req-1", model_name="gpt-3.5-turbo")
|
|
item3 = FlowItem(priority=2, request_id="req-2", model_name="gpt-3.5-turbo")
|
|
|
|
await scheduler.add_request(item1)
|
|
await scheduler.add_request(item2)
|
|
await scheduler.add_request(item3)
|
|
|
|
# Verify initial queue size
|
|
queue_before = await scheduler.get_queue(model_name="gpt-3.5-turbo")
|
|
assert len(queue_before) == 3, f"Expected 3 items in queue, got {len(queue_before)}"
|
|
|
|
# Simulate timeout cleanup - remove a non-front request (item2)
|
|
await scheduler.remove_request(request_id="req-1", model_name="gpt-3.5-turbo")
|
|
|
|
# Verify queue was cleaned up
|
|
queue_after = await scheduler.get_queue(model_name="gpt-3.5-turbo")
|
|
assert len(queue_after) == 2, f"Expected 2 items after cleanup, got {len(queue_after)}"
|
|
|
|
# Verify the correct request was removed
|
|
remaining_ids = [item[1] for item in queue_after]
|
|
assert "req-1" not in remaining_ids, "Expected req-1 to be removed"
|
|
assert "req-0" in remaining_ids, "Expected req-0 to remain"
|
|
assert "req-2" in remaining_ids, "Expected req-2 to remain"
|
|
|
|
# Verify remaining items are in correct priority order (0 should be first)
|
|
assert queue_after[0][1] == "req-0", "Expected req-0 (priority 0) to be at front"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poll_admits_request_missing_from_queue_while_a_deployment_is_healthy():
|
|
scheduler: Final = Scheduler()
|
|
|
|
assert await scheduler.poll(
|
|
request=FlowItem(priority=1, request_id="erased-by-concurrent-write", model_name="sched-model"),
|
|
health_deployments=[{"model_info": {"id": "a"}}],
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poll_during_cooldown_admits_only_the_head_of_the_queue():
|
|
scheduler: Final = Scheduler()
|
|
later: Final = FlowItem(priority=2, request_id="later", model_name="sched-model")
|
|
head: Final = FlowItem(priority=1, request_id="head", model_name="sched-model")
|
|
await scheduler.add_request(later)
|
|
await scheduler.add_request(head)
|
|
|
|
assert not await scheduler.poll(request=later, health_deployments=[])
|
|
assert await scheduler.poll(request=head, health_deployments=[])
|
|
assert await scheduler.get_queue("sched-model") == [(2, "later")]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poll_during_cooldown_re_enqueues_a_request_a_concurrent_writer_erased_behind_the_head():
|
|
scheduler: Final = Scheduler()
|
|
still_queued: Final = FlowItem(priority=0, request_id="still-queued", model_name="sched-model")
|
|
erased: Final = FlowItem(priority=1, request_id="erased-by-concurrent-write", model_name="sched-model")
|
|
await scheduler.add_request(still_queued)
|
|
|
|
assert not await scheduler.poll(request=erased, health_deployments=[])
|
|
assert await scheduler.get_queue("sched-model") == [(0, "still-queued"), (1, "erased-by-concurrent-write")]
|
|
assert await scheduler.poll(request=still_queued, health_deployments=[])
|
|
assert await scheduler.poll(request=erased, health_deployments=[])
|
|
assert await scheduler.get_queue("sched-model") == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poll_during_cooldown_admits_an_erased_request_that_outranks_the_queue():
|
|
scheduler: Final = Scheduler()
|
|
await scheduler.add_request(FlowItem(priority=2, request_id="still-queued", model_name="sched-model"))
|
|
|
|
assert await scheduler.poll(
|
|
request=FlowItem(priority=0, request_id="erased-urgent", model_name="sched-model"), health_deployments=[]
|
|
)
|
|
assert await scheduler.get_queue("sched-model") == [(2, "still-queued")]
|
|
|
|
|
|
class _ExpiringCache:
|
|
def __init__(self) -> None:
|
|
self.store: dict[str, object] = {}
|
|
|
|
async def async_get_cache(self, key: str, **kwargs: object) -> object:
|
|
return self.store.get(key)
|
|
|
|
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
|
|
self.store[key] = value
|
|
|
|
|
|
class _SchedulerClock:
|
|
now: float = 1000.0
|
|
|
|
def __call__(self) -> float:
|
|
return self.now
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_waiter_observes_shared_expiry_after_reading_just_before_it(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
clock: Final = _SchedulerClock()
|
|
server: Final = FakeServer()
|
|
|
|
class Pool(BlockingConnectionPool):
|
|
def __init__(self, **kwargs: object) -> None:
|
|
super().__init__(connection_class=FakeConnection, server=server, **kwargs)
|
|
|
|
monkeypatch.setattr(time, "time", clock)
|
|
monkeypatch.setattr(redis, "Redis", FakeRedis)
|
|
monkeypatch.setattr(redis.asyncio, "BlockingConnectionPool", Pool)
|
|
cache: Final = RedisCache(host="scheduler-expiry")
|
|
scheduler: Final = Scheduler(redis_cache=cache)
|
|
scheduler.cache.in_memory_cache = InMemoryCache(clock=clock)
|
|
waiter: Final = FlowItem(priority=2, request_id="waiter", model_name="expiry")
|
|
key: Final = f"{SchedulerCacheKeys.queue.value}:{waiter.model_name}"
|
|
try:
|
|
await cache.async_set_cache(key, [(1, "dead-replica")])
|
|
assert await scheduler.get_queue(waiter.model_name) == [(1, "dead-replica")]
|
|
clock.now = 1000.1
|
|
await scheduler.add_request(waiter)
|
|
clock.now = 1005.05
|
|
assert not await scheduler.poll(waiter, [])
|
|
clock.now = 1005.2
|
|
server.connected = False
|
|
assert not await scheduler.poll(waiter, [])
|
|
server.connected = True
|
|
assert await cache.async_get_cache(key) is None
|
|
assert await scheduler.get_queue(waiter.model_name) == []
|
|
server.connected = False
|
|
assert await scheduler.get_queue(waiter.model_name) == []
|
|
server.connected = True
|
|
assert await scheduler.poll(waiter, [])
|
|
assert await scheduler.get_queue(waiter.model_name) == []
|
|
finally:
|
|
server.connected = True
|
|
await cache.disconnect()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wait_for_turn_during_cooldown_survives_the_queue_key_expiring():
|
|
cache: Final = _ExpiringCache()
|
|
scheduler: Final = Scheduler(redis_cache=cache)
|
|
scheduler.cache.in_memory_cache.cache_dict.clear()
|
|
|
|
async def no_healthy_deployments_after_the_key_expired() -> Sequence[object]:
|
|
cache.store.clear()
|
|
scheduler.cache.in_memory_cache.cache_dict.clear()
|
|
return ()
|
|
|
|
await scheduler.wait_for_turn(
|
|
request=FlowItem(priority=1, request_id="sole-waiter", model_name="sched-model"),
|
|
timeout=5,
|
|
get_healthy_deployments=no_healthy_deployments_after_the_key_expired,
|
|
)
|
|
|
|
assert await scheduler.get_queue("sched-model") == []
|
|
|
|
|
|
|
|
class _JsonRoundTripRedisCache:
|
|
def __init__(self) -> None:
|
|
self.store: dict[str, str] = {}
|
|
|
|
async def async_get_cache(self, key: str, **kwargs: object) -> object:
|
|
raw: Final = self.store.get(key)
|
|
return None if raw is None else json.loads(raw)
|
|
|
|
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
|
|
self.store[key] = json.dumps(value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_second_replica_enqueues_behind_a_queue_decoded_from_redis():
|
|
redis_cache: Final = _JsonRoundTripRedisCache()
|
|
replica_a: Final = Scheduler(redis_cache=redis_cache)
|
|
replica_b: Final = Scheduler(redis_cache=redis_cache)
|
|
waiting_on_a: Final = FlowItem(priority=1, request_id="waiting-on-a", model_name="sched-model")
|
|
urgent_on_b: Final = FlowItem(priority=0, request_id="urgent-on-b", model_name="sched-model")
|
|
await replica_a.add_request(waiting_on_a)
|
|
|
|
await replica_b.add_request(urgent_on_b)
|
|
|
|
assert await replica_b.get_queue("sched-model") == [(0, "urgent-on-b"), (1, "waiting-on-a")]
|
|
assert await replica_b.poll(request=urgent_on_b, health_deployments=[])
|
|
assert await replica_b.poll(request=waiting_on_a, health_deployments=[])
|
|
await replica_b.remove_request(request_id="waiting-on-a", model_name="sched-model")
|
|
assert await replica_b.get_queue("sched-model") == []
|
|
|
|
class _PausesAfterEnqueueScheduler(Scheduler):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.enqueued: Final = asyncio.Event()
|
|
|
|
async def add_request(self, request: FlowItem) -> None:
|
|
await super().add_request(request)
|
|
self.enqueued.set()
|
|
await asyncio.Event().wait()
|
|
|
|
|
|
async def _no_healthy_deployments() -> Sequence[object]:
|
|
return ()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wait_for_turn_removes_entry_when_cancelled_mid_enqueue():
|
|
scheduler: Final = _PausesAfterEnqueueScheduler()
|
|
waiting: Final = asyncio.create_task(
|
|
scheduler.wait_for_turn(
|
|
request=FlowItem(priority=1, request_id="cancelled", model_name="sched-model"),
|
|
timeout=5,
|
|
get_healthy_deployments=_no_healthy_deployments,
|
|
)
|
|
)
|
|
await scheduler.enqueued.wait()
|
|
assert await scheduler.get_queue("sched-model") == [(1, "cancelled")]
|
|
|
|
waiting.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await waiting
|
|
|
|
assert await scheduler.get_queue("sched-model") == []
|
|
|
|
|
|
class _HeldRemovalScheduler(Scheduler):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.removing: Final = asyncio.Event()
|
|
self.finish_removal: Final = asyncio.Event()
|
|
|
|
async def remove_request(self, request_id: str, model_name: str) -> None:
|
|
self.removing.set()
|
|
await self.finish_removal.wait()
|
|
await super().remove_request(request_id=request_id, model_name=model_name)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wait_for_turn_finishes_removal_when_cancelled_again_during_cleanup():
|
|
scheduler: Final = _HeldRemovalScheduler()
|
|
await scheduler.add_request(FlowItem(priority=0, request_id="head", model_name="sched-model"))
|
|
polling: Final = asyncio.Event()
|
|
|
|
async def no_healthy_deployments() -> Sequence[object]:
|
|
polling.set()
|
|
return ()
|
|
|
|
waiting: Final = asyncio.create_task(
|
|
scheduler.wait_for_turn(
|
|
request=FlowItem(priority=1, request_id="cancelled", model_name="sched-model"),
|
|
timeout=5,
|
|
get_healthy_deployments=no_healthy_deployments,
|
|
)
|
|
)
|
|
await polling.wait()
|
|
waiting.cancel()
|
|
await scheduler.removing.wait()
|
|
waiting.cancel()
|
|
scheduler.finish_removal.set()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await waiting
|
|
|
|
assert await scheduler.get_queue("sched-model") == [(0, "head")]
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _vcr_outcome_gate(request, vcr):
|
|
install_live_call_probe(request, vcr)
|
|
yield
|
|
record_vcr_outcome(request, vcr)
|
|
|
|
|
|
@pytest.fixture(scope="function", autouse=True)
|
|
def isolate_litellm_state():
|
|
"""
|
|
Per-function isolation fixture.
|
|
|
|
Resets litellm globals to their true defaults before each test and
|
|
restores them afterward, so tests don't leak side effects.
|
|
Works safely under pytest-xdist parallel execution.
|
|
"""
|
|
original_state = {}
|
|
for attr in (
|
|
"callbacks",
|
|
"success_callback",
|
|
"failure_callback",
|
|
"_async_success_callback",
|
|
"_async_failure_callback",
|
|
):
|
|
if hasattr(litellm, attr):
|
|
val = getattr(litellm, attr)
|
|
original_state[attr] = val.copy() if val else []
|
|
for attr in ("pre_call_rules", "post_call_rules"):
|
|
if hasattr(litellm, attr):
|
|
val = getattr(litellm, attr)
|
|
original_state[attr] = val.copy() if val else []
|
|
for attr in _SCALAR_DEFAULTS:
|
|
if hasattr(litellm, attr):
|
|
original_state[attr] = getattr(litellm, attr)
|
|
if hasattr(litellm, "in_memory_llm_clients_cache"):
|
|
litellm.in_memory_llm_clients_cache.flush_cache()
|
|
for attr in (
|
|
"callbacks",
|
|
"success_callback",
|
|
"failure_callback",
|
|
"_async_success_callback",
|
|
"_async_failure_callback",
|
|
"pre_call_rules",
|
|
"post_call_rules",
|
|
):
|
|
if hasattr(litellm, attr):
|
|
setattr(litellm, attr, [])
|
|
for attr, default_val in _SCALAR_DEFAULTS.items():
|
|
if hasattr(litellm, attr):
|
|
setattr(litellm, attr, default_val)
|
|
yield
|
|
asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue())
|
|
if hasattr(litellm, "in_memory_llm_clients_cache"):
|
|
litellm.in_memory_llm_clients_cache.flush_cache()
|
|
for attr, original_value in original_state.items():
|
|
if hasattr(litellm, attr):
|
|
setattr(litellm, attr, original_value)
|
|
_invalidate_model_cost_lowercase_map()
|
|
|
|
|
|
_SCALAR_DEFAULTS = {
|
|
"num_retries": getattr(litellm, "num_retries", None),
|
|
"num_retries_per_request": getattr(litellm, "num_retries_per_request", None),
|
|
"request_timeout": getattr(litellm, "request_timeout", None),
|
|
"set_verbose": getattr(litellm, "set_verbose", False),
|
|
"cache": getattr(litellm, "cache", None),
|
|
"allowed_fails": getattr(litellm, "allowed_fails", 3),
|
|
"default_fallbacks": getattr(litellm, "default_fallbacks", None),
|
|
"enable_azure_ad_token_refresh": getattr(litellm, "enable_azure_ad_token_refresh", None),
|
|
"tag_budget_config": getattr(litellm, "tag_budget_config", None),
|
|
"model_cost": getattr(litellm, "model_cost", None),
|
|
"token_counter": getattr(litellm, "token_counter", None),
|
|
"disable_aiohttp_transport": getattr(litellm, "disable_aiohttp_transport", False),
|
|
"force_ipv4": getattr(litellm, "force_ipv4", False),
|
|
"drop_params": getattr(litellm, "drop_params", None),
|
|
"modify_params": getattr(litellm, "modify_params", False),
|
|
"api_base": getattr(litellm, "api_base", None),
|
|
"api_key": getattr(litellm, "api_key", None),
|
|
}
|
|
|
|
|
|
@pytest.fixture(scope="module", autouse=True)
|
|
def setup_and_teardown():
|
|
"""
|
|
Module-scoped setup. Reloads litellm only in single-process mode
|
|
(skipped under xdist to avoid cross-worker interference).
|
|
"""
|
|
import litellm
|
|
|
|
worker_id = os.environ.get("PYTEST_XDIST_WORKER", None)
|
|
if worker_id is None:
|
|
importlib.reload(litellm)
|
|
try:
|
|
if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"):
|
|
import litellm.proxy.proxy_server
|
|
|
|
importlib.reload(litellm.proxy.proxy_server)
|
|
except Exception:
|
|
pass
|
|
if hasattr(litellm, "in_memory_llm_clients_cache"):
|
|
litellm.in_memory_llm_clients_cache.flush_cache()
|
|
yield
|