mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
Merge ab66b12125 into 3b8ad37296
This commit is contained in:
yi-coding
GitHub
commit
6fcbac3dac
3 changed files with 541 additions and 4 deletions
|
|
@ -1,15 +1,29 @@
|
|||
import asyncio
|
||||
import enum
|
||||
import heapq
|
||||
import logging as _logging
|
||||
import random as _random
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from functools import partial as _partial
|
||||
from typing import Final, TypeAlias
|
||||
from typing import Protocol as _Protocol
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm import print_verbose
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_router_logger as _verbose_router_logger
|
||||
from litellm.caching.caching import DualCache, RedisCache
|
||||
from litellm.caching.redis_cache import (
|
||||
RedisCircuitBreakerOpenError as _RedisCircuitBreakerOpenError,
|
||||
)
|
||||
from litellm.caching.redis_cache import (
|
||||
_is_redis_health_failure, # pyright: ignore[reportPrivateUsage] # RedisCache's connectivity classifier has no public counterpart
|
||||
)
|
||||
from litellm.caching.redis_cache import (
|
||||
log_redis_failure as _log_redis_failure,
|
||||
)
|
||||
from litellm.constants import DEFAULT_IN_MEMORY_TTL, DEFAULT_POLLING_INTERVAL
|
||||
from litellm.exceptions import Timeout
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
|
@ -17,6 +31,25 @@ from litellm.types.llms.base import LiteLLMBaseModel
|
|||
SCHEDULER_QUEUE_TARGET: Final = "scheduler_queue"
|
||||
QueueEntry: TypeAlias = tuple[int, str]
|
||||
_QUEUE_ENTRIES: Final = TypeAdapter(list[QueueEntry])
|
||||
_Queue: TypeAlias = tuple[QueueEntry, ...]
|
||||
_QueueResult: TypeAlias = tuple[bool, _Queue]
|
||||
_QUEUE_SNAPSHOT: Final = TypeAdapter(_Queue)
|
||||
_RAW_QUEUE: Final = TypeAdapter[str | bytes | None](str | bytes | None)
|
||||
_CAS_RESULT: Final = TypeAdapter(bool)
|
||||
_QUEUE_READ: Final = "return redis.call('GET', KEYS[1])"
|
||||
_QUEUE_COMPARE_SET: Final = """
|
||||
local current = redis.call('GET', KEYS[1])
|
||||
if (ARGV[1] == '0' and current ~= false) or
|
||||
(ARGV[1] == '1' and current ~= ARGV[2]) then
|
||||
return 0
|
||||
end
|
||||
if ARGV[4] == '' then
|
||||
redis.call('SET', KEYS[1], ARGV[3])
|
||||
else
|
||||
redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[4])
|
||||
end
|
||||
return 1
|
||||
"""
|
||||
|
||||
|
||||
class SchedulerCacheKeys(enum.Enum):
|
||||
|
|
@ -30,6 +63,98 @@ class FlowItem(LiteLLMBaseModel):
|
|||
model_name: str
|
||||
|
||||
|
||||
class _LocalQueueCache(_Protocol):
|
||||
@property
|
||||
def default_in_memory_ttl(self) -> float | None: ...
|
||||
|
||||
async def async_get_cache(self, key: str, *, local_only: bool) -> object: ...
|
||||
|
||||
async def async_set_cache(self, key: str, value: Sequence[QueueEntry], *, local_only: bool) -> None: ...
|
||||
|
||||
|
||||
async def _read_local_queue(cache: _LocalQueueCache, key: str) -> _Queue:
|
||||
local: Final = await cache.async_get_cache(key=key, local_only=True)
|
||||
return _QUEUE_SNAPSHOT.validate_python(local or ())
|
||||
|
||||
|
||||
async def _write_local_queue(cache: _LocalQueueCache, key: str, queue: _Queue) -> None:
|
||||
await cache.async_set_cache(key=key, value=list(queue), local_only=True)
|
||||
|
||||
|
||||
def _enqueue(entry: QueueEntry, queue: _Queue) -> _QueueResult:
|
||||
return False, queue if entry in queue else tuple(sorted((*queue, entry)))
|
||||
|
||||
|
||||
def _remove(request_id: str, queue: _Queue) -> _QueueResult:
|
||||
return False, tuple(sorted(item for item in queue if item[1] != request_id))
|
||||
|
||||
|
||||
def _poll(entry: QueueEntry, queue: _Queue) -> _QueueResult:
|
||||
pending: Final = tuple(sorted(queue if entry in queue else (*queue, entry)))
|
||||
return (True, pending[1:]) if pending[0] == entry else (False, pending)
|
||||
|
||||
|
||||
def _redis_unavailable(error: Exception) -> bool:
|
||||
return isinstance(error, _RedisCircuitBreakerOpenError) or _is_redis_health_failure(error)
|
||||
|
||||
|
||||
async def _read_redis_queue(cache: RedisCache, key: str) -> tuple[str | bytes | None, _Queue]:
|
||||
raw: Final = _RAW_QUEUE.validate_python(await cache.async_register_script(_QUEUE_READ)(keys=(key,), args=()))
|
||||
return raw, () if raw is None else _QUEUE_SNAPSHOT.validate_json(raw)
|
||||
|
||||
|
||||
async def _mutate_redis_queue(
|
||||
cache: RedisCache, key: str, model_name: str, ttl: int | None, change: Callable[[_Queue], _QueueResult]
|
||||
) -> _QueueResult:
|
||||
script: Final = cache.async_register_script(_QUEUE_COMPARE_SET)
|
||||
for attempt in range(100):
|
||||
raw, queue = await _read_redis_queue(cache, key)
|
||||
admitted, updated = change(queue)
|
||||
if updated == queue:
|
||||
return admitted, updated
|
||||
if _CAS_RESULT.validate_python(
|
||||
await script(
|
||||
keys=(key,),
|
||||
args=(
|
||||
"0" if raw is None else "1",
|
||||
raw or "",
|
||||
_QUEUE_SNAPSHOT.dump_json(updated),
|
||||
"" if ttl is None else str(ttl),
|
||||
),
|
||||
)
|
||||
):
|
||||
return admitted, updated
|
||||
if attempt < 99:
|
||||
await asyncio.sleep(_random.uniform(0, min(0.05, 0.001 * 2 ** min(attempt, 6))))
|
||||
raise Timeout(
|
||||
message="Scheduler queue remained busy after concurrent updates", model=model_name, llm_provider="openai"
|
||||
)
|
||||
|
||||
|
||||
async def _shared_queue_result(
|
||||
cache: _LocalQueueCache,
|
||||
redis_cache: RedisCache,
|
||||
key: str,
|
||||
model_name: str,
|
||||
change: Callable[[_Queue], _QueueResult],
|
||||
) -> _QueueResult:
|
||||
try:
|
||||
return await _mutate_redis_queue(
|
||||
redis_cache,
|
||||
key,
|
||||
model_name,
|
||||
redis_cache.default_ttl if cache.default_in_memory_ttl is None else int(cache.default_in_memory_ttl),
|
||||
change,
|
||||
)
|
||||
except Exception as error:
|
||||
if not _redis_unavailable(error):
|
||||
raise
|
||||
_log_redis_failure(
|
||||
_verbose_router_logger, _logging.WARNING, "Scheduler using local queue while Redis is unavailable", error
|
||||
)
|
||||
return change(await _read_local_queue(cache, key))
|
||||
|
||||
|
||||
class Scheduler:
|
||||
cache: DualCache
|
||||
|
||||
|
|
@ -49,7 +174,22 @@ class Scheduler:
|
|||
self.cache = DualCache(redis_cache=redis_cache, default_in_memory_ttl=default_in_memory_ttl)
|
||||
self.polling_interval = polling_interval or DEFAULT_POLLING_INTERVAL # default to 3ms
|
||||
|
||||
@with_service_target(SCHEDULER_QUEUE_TARGET)
|
||||
async def _update_shared_queue(
|
||||
self, redis_cache: RedisCache, model_name: str, change: Callable[[_Queue], _QueueResult]
|
||||
) -> bool:
|
||||
key: Final = f"{SchedulerCacheKeys.queue.value}:{model_name}"
|
||||
result: Final = await _shared_queue_result(self.cache, redis_cache, key, model_name, change)
|
||||
await _write_local_queue(self.cache, key, result[1])
|
||||
return result[0]
|
||||
|
||||
async def add_request(self, request: FlowItem):
|
||||
redis_cache: Final = self.cache.redis_cache
|
||||
if isinstance(redis_cache, RedisCache):
|
||||
await self._update_shared_queue(
|
||||
redis_cache, request.model_name, _partial(_enqueue, (request.priority, request.request_id))
|
||||
)
|
||||
return
|
||||
# We use the priority directly, as lower values indicate higher priority
|
||||
# get the queue
|
||||
queue: Final = await self.get_queue(model_name=request.model_name)
|
||||
|
|
@ -78,6 +218,11 @@ class Scheduler:
|
|||
if len(health_deployments) > 0:
|
||||
return True
|
||||
|
||||
redis_cache: Final = self.cache.redis_cache
|
||||
if isinstance(redis_cache, RedisCache):
|
||||
return await self._update_shared_queue(
|
||||
redis_cache, request.model_name, _partial(_poll, (request.priority, request.request_id))
|
||||
)
|
||||
queue: Final = await self.get_queue(model_name=request.model_name)
|
||||
entry: Final = (request.priority, request.request_id)
|
||||
print_verbose(f"queue: {queue}, seeking {entry}")
|
||||
|
|
@ -109,14 +254,24 @@ class Scheduler:
|
|||
return
|
||||
await asyncio.sleep(self.polling_interval)
|
||||
finally:
|
||||
await asyncio.shield(self.remove_request(request_id=request.request_id, model_name=request.model_name))
|
||||
await asyncio.shield(self._cleanup_request(request))
|
||||
raise Timeout(message="Request timed out while polling queue", model=request.model_name, llm_provider="openai")
|
||||
|
||||
async def _cleanup_request(self, request: FlowItem) -> None:
|
||||
try:
|
||||
await self.remove_request(request_id=request.request_id, model_name=request.model_name)
|
||||
except Exception as error: # noqa: BLE001 # Cleanup must preserve every request outcome, including cancellation
|
||||
_log_redis_failure(_verbose_router_logger, _logging.WARNING, "Scheduler queue cleanup failed", error)
|
||||
|
||||
async def remove_request(self, request_id: str, model_name: str) -> None:
|
||||
"""
|
||||
Remove a specific request from the priority queue for a model.
|
||||
Used when a request times out while waiting in the queue.
|
||||
"""
|
||||
redis_cache: Final = self.cache.redis_cache
|
||||
if isinstance(redis_cache, RedisCache):
|
||||
await self._update_shared_queue(redis_cache, model_name, _partial(_remove, request_id))
|
||||
return
|
||||
queue: Final = await self.get_queue(model_name=model_name)
|
||||
filtered_queue: Final = [item for item in queue if item[1] != request_id]
|
||||
heapq.heapify(filtered_queue) # restore heap invariant after filtering
|
||||
|
|
@ -151,6 +306,17 @@ class Scheduler:
|
|||
Redis hands the queue back as JSON lists, so every entry is validated into the
|
||||
(priority, request_id) tuple the heap operations compare against.
|
||||
"""
|
||||
redis_cache: Final = self.cache.redis_cache
|
||||
if isinstance(redis_cache, RedisCache):
|
||||
key: Final = f"{SchedulerCacheKeys.queue.value}:{model_name}"
|
||||
try:
|
||||
_, snapshot = await _read_redis_queue(redis_cache, key)
|
||||
except Exception as error:
|
||||
if not _redis_unavailable(error):
|
||||
raise
|
||||
return list(await _read_local_queue(self.cache, key))
|
||||
await _write_local_queue(self.cache, key, snapshot)
|
||||
return list(snapshot)
|
||||
if self.cache is not None:
|
||||
_cache_key: Final = f"{SchedulerCacheKeys.queue.value}:{model_name}"
|
||||
response: Final = await self.cache.async_get_cache(key=_cache_key)
|
||||
|
|
|
|||
|
|
@ -1,15 +1,22 @@
|
|||
import asyncio
|
||||
import os
|
||||
import threading
|
||||
import uuid
|
||||
from typing import Final
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from integration._support.async_server import LoopbackServer, close_writer
|
||||
from integration._support.ports import allocate_port
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.caching.affinity_cache import claim_affinity_pin
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.scheduler import FlowItem, Scheduler, SchedulerCacheKeys
|
||||
|
||||
|
||||
def _redis_cache() -> RedisCache:
|
||||
|
|
@ -95,6 +102,212 @@ async def test_sync_and_async_batch_reads_share_one_redis_without_sync_reads_goi
|
|||
assert await dual_cache.async_batch_get_cache(keys=[survivor_key]) == [expected[survivor_key]]
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(loop_scope="function")
|
||||
async def shared_scheduler_cache() -> AsyncIterator[RedisCache]:
|
||||
cache: Final = RedisCache(
|
||||
host=os.environ["REDIS_HOST"],
|
||||
port=int(os.environ["REDIS_PORT"]),
|
||||
namespace="scheduler-atomic-" + uuid.uuid4().hex,
|
||||
)
|
||||
try:
|
||||
yield cache
|
||||
finally:
|
||||
client: Final = cache.init_async_client()
|
||||
await client.delete(cache.check_and_fix_namespace(f"{SchedulerCacheKeys.queue.value}:shared"))
|
||||
await client.aclose()
|
||||
cache.redis_client.close()
|
||||
|
||||
|
||||
async def test_shared_scheduler_preserves_concurrent_admissions(shared_scheduler_cache: RedisCache) -> None:
|
||||
first: Final = Scheduler(redis_cache=shared_scheduler_cache)
|
||||
second: Final = Scheduler(redis_cache=shared_scheduler_cache)
|
||||
await asyncio.gather(first.get_queue("shared"), second.get_queue("shared"))
|
||||
await asyncio.gather(
|
||||
first.add_request(FlowItem(priority=1, request_id="first", model_name="shared")),
|
||||
second.add_request(FlowItem(priority=2, request_id="second", model_name="shared")),
|
||||
)
|
||||
assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [[1, "first"], [2, "second"]]
|
||||
assert await first.get_queue("shared") == [(1, "first"), (2, "second")]
|
||||
assert await second.get_queue("shared") == [(1, "first"), (2, "second")]
|
||||
|
||||
|
||||
async def test_shared_scheduler_does_not_resurrect_peer_removal(shared_scheduler_cache: RedisCache) -> None:
|
||||
await shared_scheduler_cache.async_set_cache("scheduler:queue:shared", [[1, "removed"], [2, "waiting"]])
|
||||
first: Final = Scheduler(redis_cache=shared_scheduler_cache)
|
||||
second: Final = Scheduler(redis_cache=shared_scheduler_cache)
|
||||
await asyncio.gather(first.get_queue("shared"), second.get_queue("shared"))
|
||||
await first.remove_request("removed", "shared")
|
||||
await second.add_request(FlowItem(priority=3, request_id="new", model_name="shared"))
|
||||
assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [[2, "waiting"], [3, "new"]]
|
||||
|
||||
|
||||
async def test_shared_scheduler_observes_peer_pop(shared_scheduler_cache: RedisCache) -> None:
|
||||
await shared_scheduler_cache.async_set_cache("scheduler:queue:shared", [[1, "first"], [2, "second"]])
|
||||
first: Final = Scheduler(redis_cache=shared_scheduler_cache)
|
||||
second: Final = Scheduler(redis_cache=shared_scheduler_cache)
|
||||
await asyncio.gather(first.get_queue("shared"), second.get_queue("shared"))
|
||||
assert await first.poll(FlowItem(priority=1, request_id="first", model_name="shared"), [])
|
||||
assert await second.poll(FlowItem(priority=2, request_id="second", model_name="shared"), [])
|
||||
assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == []
|
||||
|
||||
|
||||
async def test_shared_scheduler_keeps_other_requests_during_concurrent_cleanup(
|
||||
shared_scheduler_cache: RedisCache,
|
||||
) -> None:
|
||||
schedulers: Final = tuple(Scheduler(redis_cache=shared_scheduler_cache) for _ in range(16))
|
||||
await asyncio.gather(
|
||||
*(
|
||||
scheduler.add_request(FlowItem(priority=i, request_id=str(i), model_name="shared"))
|
||||
for i, scheduler in enumerate(schedulers)
|
||||
)
|
||||
)
|
||||
await asyncio.gather(
|
||||
*(scheduler.remove_request(str(i), "shared") for i, scheduler in enumerate(schedulers) if i % 2 == 0)
|
||||
)
|
||||
assert await schedulers[0].get_queue("shared") == [(i, str(i)) for i in range(16) if i % 2]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def interrupted_scheduler_cache(shared_scheduler_cache: RedisCache) -> Iterator[tuple[RedisCache, threading.Event]]:
|
||||
interrupted: Final = threading.Event()
|
||||
|
||||
async def relay(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
||||
upstream_reader, upstream_writer = await asyncio.open_connection(
|
||||
os.environ["REDIS_HOST"], int(os.environ["REDIS_PORT"])
|
||||
)
|
||||
|
||||
async def forward(source: asyncio.StreamReader, destination: asyncio.StreamWriter) -> None:
|
||||
try:
|
||||
while chunk := await source.read(65536):
|
||||
if interrupted.is_set():
|
||||
return
|
||||
destination.write(chunk)
|
||||
await destination.drain()
|
||||
except ConnectionError:
|
||||
return
|
||||
finally:
|
||||
await close_writer(destination)
|
||||
|
||||
await asyncio.gather(forward(reader, upstream_writer), forward(upstream_reader, writer))
|
||||
|
||||
port: Final = allocate_port()
|
||||
server: Final = LoopbackServer(port, relay)
|
||||
server.start()
|
||||
cache: Final = RedisCache(host="127.0.0.1", port=port, namespace=shared_scheduler_cache.namespace)
|
||||
try:
|
||||
yield cache, interrupted
|
||||
finally:
|
||||
cache.redis_client.close()
|
||||
server.stop()
|
||||
|
||||
|
||||
async def test_shared_scheduler_keeps_local_fallback_out_of_recovered_shared_queue(
|
||||
shared_scheduler_cache: RedisCache,
|
||||
interrupted_scheduler_cache: tuple[RedisCache, threading.Event],
|
||||
) -> None:
|
||||
cache, interrupted = interrupted_scheduler_cache
|
||||
scheduler: Final = Scheduler(redis_cache=cache)
|
||||
try:
|
||||
await scheduler.add_request(FlowItem(priority=1, request_id="removed", model_name="shared"))
|
||||
interrupted.set()
|
||||
await scheduler.add_request(FlowItem(priority=2, request_id="local", model_name="shared"))
|
||||
assert await scheduler.get_queue("shared") == [(1, "removed"), (2, "local")]
|
||||
await Scheduler(redis_cache=shared_scheduler_cache).remove_request("removed", "shared")
|
||||
interrupted.clear()
|
||||
await scheduler.add_request(FlowItem(priority=3, request_id="new", model_name="shared"))
|
||||
assert await scheduler.get_queue("shared") == [(3, "new")]
|
||||
assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [[3, "new"]]
|
||||
finally:
|
||||
interrupted.clear()
|
||||
await cache.init_async_client().aclose()
|
||||
|
||||
|
||||
async def test_shared_scheduler_preserves_invalid_queue_for_diagnosis(shared_scheduler_cache: RedisCache) -> None:
|
||||
await shared_scheduler_cache.async_set_cache("scheduler:queue:shared", {"invalid": "queue"})
|
||||
scheduler: Final = Scheduler(redis_cache=shared_scheduler_cache)
|
||||
with pytest.raises(ValidationError):
|
||||
await scheduler.add_request(FlowItem(priority=1, request_id="first", model_name="shared"))
|
||||
assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == {"invalid": "queue"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("priority", [0, -1, 2**80])
|
||||
async def test_shared_scheduler_retains_priority_and_request_id(
|
||||
shared_scheduler_cache: RedisCache, priority: int
|
||||
) -> None:
|
||||
scheduler: Final = Scheduler(redis_cache=shared_scheduler_cache)
|
||||
scheduler.cache.update_cache_ttl(default_in_memory_ttl=17, default_redis_ttl=None)
|
||||
await scheduler.add_request(FlowItem(priority=priority, request_id="请求-é", model_name="shared"))
|
||||
assert await scheduler.get_queue("shared") == [(priority, "请求-é")]
|
||||
assert (
|
||||
0
|
||||
< shared_scheduler_cache.redis_client.ttl(
|
||||
shared_scheduler_cache.check_and_fix_namespace("scheduler:queue:shared")
|
||||
)
|
||||
<= 17
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("outcome", ["admitted", "failed", "cancelled"])
|
||||
async def test_shared_scheduler_cleanup_failure_preserves_request_outcome(
|
||||
shared_scheduler_cache: RedisCache,
|
||||
outcome: Literal["admitted", "failed", "cancelled"],
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
scheduler: Final = Scheduler(redis_cache=shared_scheduler_cache)
|
||||
original: Final = RuntimeError("health lookup failed")
|
||||
|
||||
async def health_after_corrupting_queue() -> tuple[object, ...]:
|
||||
await shared_scheduler_cache.async_set_cache("scheduler:queue:shared", {"invalid": "queue"})
|
||||
match outcome:
|
||||
case "failed":
|
||||
raise original
|
||||
case "cancelled":
|
||||
raise asyncio.CancelledError
|
||||
case "admitted":
|
||||
return ("healthy",)
|
||||
|
||||
async def wait() -> None:
|
||||
await scheduler.wait_for_turn(
|
||||
FlowItem(priority=1, request_id="cleanup", model_name="shared"),
|
||||
timeout=5,
|
||||
get_healthy_deployments=health_after_corrupting_queue,
|
||||
)
|
||||
|
||||
match outcome:
|
||||
case "admitted":
|
||||
await wait()
|
||||
case "failed":
|
||||
with pytest.raises(RuntimeError) as caught:
|
||||
await wait()
|
||||
assert caught.value is original
|
||||
case "cancelled":
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await wait()
|
||||
assert "Scheduler queue cleanup failed" in caplog.text
|
||||
assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == {"invalid": "queue"}
|
||||
|
||||
|
||||
async def test_shared_scheduler_handles_three_hundred_concurrent_enqueues_and_removals(
|
||||
shared_scheduler_cache: RedisCache,
|
||||
) -> None:
|
||||
schedulers: Final = tuple(Scheduler(redis_cache=shared_scheduler_cache) for _ in range(300))
|
||||
for scheduler in schedulers:
|
||||
scheduler.cache.update_cache_ttl(default_in_memory_ttl=60, default_redis_ttl=None)
|
||||
requests: Final = tuple(
|
||||
FlowItem(priority=i % 8, request_id=f"burst-{i:03}", model_name="shared") for i in range(300)
|
||||
)
|
||||
|
||||
await asyncio.gather(*(scheduler.add_request(request) for scheduler, request in zip(schedulers, requests)))
|
||||
|
||||
assert await schedulers[0].get_queue("shared") == sorted(
|
||||
(request.priority, request.request_id) for request in requests
|
||||
)
|
||||
await asyncio.gather(
|
||||
*(scheduler.remove_request(request.request_id, "shared") for scheduler, request in zip(schedulers, requests))
|
||||
)
|
||||
assert await schedulers[0].get_queue("shared") == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_session_repin_claims_keep_the_shared_winner_after_recovery() -> None:
|
||||
redis: Final = _redis_cache()
|
||||
|
|
|
|||
|
|
@ -5,13 +5,18 @@ import asyncio
|
|||
import importlib
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
from typing import Final
|
||||
from collections.abc import Iterator, Sequence
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.exceptions import Timeout
|
||||
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
|
||||
|
|
@ -332,6 +337,159 @@ async def test_wait_for_turn_finishes_removal_when_cancelled_again_during_cleanu
|
|||
assert await scheduler.get_queue("sched-model") == [(0, "head")]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def scheduler_redis_replies() -> Iterator[tuple[RedisCache, AsyncMock, AsyncMock]]:
|
||||
reads: Final = AsyncMock(return_value=None)
|
||||
writes: Final = AsyncMock(return_value=1)
|
||||
|
||||
async def execute(
|
||||
*, keys: Sequence[str], args: Sequence[str | bytes | int | float], client: object = None
|
||||
) -> object:
|
||||
return await (writes if args else reads)(keys=keys, args=args, client=client)
|
||||
|
||||
with (
|
||||
patch("redis.Redis.execute_command", return_value={"redis_version": "unit-fixture"}),
|
||||
patch("redis.asyncio.Redis.execute_command", new_callable=AsyncMock, return_value=True),
|
||||
patch("redis.commands.core.AsyncScript.__call__", side_effect=execute),
|
||||
):
|
||||
cache: Final = RedisCache(host="scheduler-unit.invalid", port=6379)
|
||||
yield cache, reads, writes
|
||||
cache.redis_client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shared_scheduler_retries_conflicts_using_the_new_peer_snapshot(
|
||||
scheduler_redis_replies: tuple[RedisCache, AsyncMock, AsyncMock],
|
||||
) -> None:
|
||||
cache, reads, writes = scheduler_redis_replies
|
||||
scheduler: Final = Scheduler(redis_cache=cache)
|
||||
reads.side_effect = (None, b'[[0,"peer"]]', RedisConnectionError("offline"))
|
||||
writes.side_effect = (0, 1)
|
||||
|
||||
await scheduler.add_request(FlowItem(priority=1, request_id="new", model_name="shared"))
|
||||
|
||||
assert await scheduler.get_queue("shared") == [(0, "peer"), (1, "new")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shared_scheduler_stops_after_conflicts_without_publishing_an_uncommitted_queue(
|
||||
scheduler_redis_replies: tuple[RedisCache, AsyncMock, AsyncMock],
|
||||
) -> None:
|
||||
cache, _reads, writes = scheduler_redis_replies
|
||||
scheduler: Final = Scheduler(redis_cache=cache)
|
||||
writes.return_value = 0
|
||||
|
||||
with patch("random.uniform", return_value=0), pytest.raises(Timeout, match="queue remained busy"):
|
||||
await scheduler.add_request(FlowItem(priority=1, request_id="never-committed", model_name="shared"))
|
||||
|
||||
assert await scheduler.cache.async_get_cache("scheduler:queue:shared", local_only=True) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("operation", ["enqueue", "read"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_shared_scheduler_rejects_malformed_snapshots_without_replacing_local_history(
|
||||
scheduler_redis_replies: tuple[RedisCache, AsyncMock, AsyncMock], operation: str
|
||||
) -> None:
|
||||
cache, reads, _writes = scheduler_redis_replies
|
||||
scheduler: Final = Scheduler(redis_cache=cache)
|
||||
await scheduler.cache.async_set_cache("scheduler:queue:shared", [(0, "last-known")], local_only=True)
|
||||
reads.side_effect = (b'{"invalid":"queue"}',)
|
||||
|
||||
call: Final = (
|
||||
scheduler.add_request(FlowItem(priority=1, request_id="new", model_name="shared"))
|
||||
if operation == "enqueue"
|
||||
else scheduler.get_queue("shared")
|
||||
)
|
||||
with pytest.raises(ValidationError):
|
||||
await call
|
||||
|
||||
assert await scheduler.cache.async_get_cache("scheduler:queue:shared", local_only=True) == [(0, "last-known")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("outcome", ["admitted", "failed", "cancelled"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_shared_scheduler_cleanup_failure_preserves_the_request_outcome(
|
||||
scheduler_redis_replies: tuple[RedisCache, AsyncMock, AsyncMock], outcome: Literal["admitted", "failed", "cancelled"]
|
||||
) -> None:
|
||||
cache, reads, _writes = scheduler_redis_replies
|
||||
scheduler: Final = Scheduler(redis_cache=cache)
|
||||
original: Final = RuntimeError("health lookup failed")
|
||||
reads.side_effect = (None, b'{"invalid":"queue"}')
|
||||
|
||||
async def health() -> Sequence[object]:
|
||||
match outcome:
|
||||
case "failed":
|
||||
raise original
|
||||
case "cancelled":
|
||||
raise asyncio.CancelledError
|
||||
case "admitted":
|
||||
return ("healthy",)
|
||||
|
||||
async def wait() -> None:
|
||||
await scheduler.wait_for_turn(
|
||||
FlowItem(priority=1, request_id="waiting", model_name="shared"), timeout=5, get_healthy_deployments=health
|
||||
)
|
||||
|
||||
match outcome:
|
||||
case "admitted":
|
||||
await wait()
|
||||
case "failed":
|
||||
with pytest.raises(RuntimeError) as caught:
|
||||
await wait()
|
||||
assert caught.value is original
|
||||
case "cancelled":
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await wait()
|
||||
assert await scheduler.cache.async_get_cache("scheduler:queue:shared", local_only=True) == [(1, "waiting")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shared_scheduler_keeps_admission_and_removal_in_its_fallback_snapshot(
|
||||
scheduler_redis_replies: tuple[RedisCache, AsyncMock, AsyncMock],
|
||||
) -> None:
|
||||
cache, reads, _writes = scheduler_redis_replies
|
||||
scheduler: Final = Scheduler(redis_cache=cache)
|
||||
head: Final = FlowItem(priority=0, request_id="head", model_name="shared")
|
||||
later: Final = FlowItem(priority=2, request_id="later", model_name="shared")
|
||||
reads.side_effect = (
|
||||
b'[[0,"head"],[2,"later"]]',
|
||||
b'[[0,"head"],[2,"later"]]',
|
||||
b'[[0,"head"],[2,"later"]]',
|
||||
RedisConnectionError("offline after admission"),
|
||||
b'[[2,"later"],[3,"peer"]]',
|
||||
RedisConnectionError("offline after removal"),
|
||||
b'[[3,"peer"]]',
|
||||
)
|
||||
|
||||
await scheduler.add_request(head)
|
||||
assert not await scheduler.poll(later, [])
|
||||
assert await scheduler.poll(head, [])
|
||||
assert await scheduler.get_queue("shared") == [(2, "later")]
|
||||
await scheduler.remove_request("later", "shared")
|
||||
assert await scheduler.get_queue("shared") == [(3, "peer")]
|
||||
assert await scheduler.get_queue("shared") == [(3, "peer")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shared_scheduler_does_not_upload_outage_only_requests_when_redis_recovers(
|
||||
scheduler_redis_replies: tuple[RedisCache, AsyncMock, AsyncMock],
|
||||
) -> None:
|
||||
cache, reads, writes = scheduler_redis_replies
|
||||
scheduler: Final = Scheduler(redis_cache=cache)
|
||||
reads.side_effect = (
|
||||
RedisConnectionError("offline during enqueue"),
|
||||
RedisConnectionError("still offline"),
|
||||
b'[[0,"peer"]]',
|
||||
RedisConnectionError("offline after recovery"),
|
||||
)
|
||||
|
||||
await scheduler.add_request(FlowItem(priority=1, request_id="local-only", model_name="shared"))
|
||||
assert await scheduler.get_queue("shared") == [(1, "local-only")]
|
||||
await scheduler.add_request(FlowItem(priority=2, request_id="new", model_name="shared"))
|
||||
assert await scheduler.get_queue("shared") == [(0, "peer"), (2, "new")]
|
||||
assert json.loads(writes.call_args.kwargs["args"][2]) == [[0, "peer"], [2, "new"]]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _vcr_outcome_gate(request, vcr):
|
||||
install_live_call_probe(request, vcr)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue