litellm/tests/unit/test_scheduler.py

# 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.litellm_core_utils.model_cost_table import publish
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)
    publish(litellm.model_cost)


_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