litellm/litellm/scheduler.py

import asyncio
import enum
import heapq
import time
from collections.abc import Awaitable, Callable, Sequence
from typing import Final, TypeAlias

from pydantic import TypeAdapter

from litellm import print_verbose
from litellm._internal_context import with_service_target
from litellm.caching.caching import DualCache, RedisCache
from litellm.constants import DEFAULT_IN_MEMORY_TTL, DEFAULT_POLLING_INTERVAL
from litellm.exceptions import Timeout
from litellm.types.llms.base import LiteLLMBaseModel

SCHEDULER_QUEUE_TARGET: Final = "scheduler_queue"
QueueEntry: TypeAlias = tuple[int, str]
_QUEUE_ENTRIES: Final = TypeAdapter(list[QueueEntry])


class SchedulerCacheKeys(enum.Enum):
    queue = "scheduler:queue"
    default_in_memory_ttl = DEFAULT_IN_MEMORY_TTL  # cache queue in-memory for 5s when redis cache available


class FlowItem(LiteLLMBaseModel):
    priority: int  # Priority between 0 and 255
    request_id: str
    model_name: str


class Scheduler:
    cache: DualCache

    def __init__(
        self,
        polling_interval: float | None = None,
        redis_cache: RedisCache | None = None,
    ):
        """
        polling_interval: float or null - frequency of polling queue. Default is 3ms.
        """
        self.queue: list[QueueEntry] = []
        default_in_memory_ttl: float | None = None
        if redis_cache is not None:
            # if redis-cache available frequently poll that instead of using in-memory.
            default_in_memory_ttl = SchedulerCacheKeys.default_in_memory_ttl.value
        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

    async def add_request(self, request: FlowItem):
        # 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)
        # update the queue
        heapq.heappush(queue, (request.priority, request.request_id))

        # save the queue
        await self.save_queue(queue=queue, model_name=request.model_name)

    async def poll(self, request: FlowItem, health_deployments: Sequence[object]) -> bool:
        """
        Return if request can be processed.

        Returns:
        - True:
            * If healthy deployments are available
            * OR If request at the top of queue
        - False:
            * If no healthy deployments available
            * AND request not at the top of queue

        A request the queue no longer holds (its cache key expired or a concurrent writer erased the entry)
        is put back at its priority so it keeps its place in the order instead of failing or jumping ahead
        """
        print_verbose(f"len(health_deployments): {len(health_deployments)}")
        if len(health_deployments) > 0:
            return True

        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}")
        if entry not in queue:
            print_verbose(f"queue no longer holds {entry}, re-enqueueing it")
            heapq.heappush(queue, entry)
            if queue[0] != entry:
                await self.save_queue(queue=queue, model_name=request.model_name)
                return False
        elif queue[0] != entry:
            return False

        heapq.heappop(queue)
        await self.save_queue(queue=queue, model_name=request.model_name)
        print_verbose(f"Popped id: {request.request_id}")
        return True

    async def wait_for_turn(
        self,
        request: FlowItem,
        timeout: float,
        get_healthy_deployments: Callable[[], Awaitable[Sequence[object]]],
    ) -> None:
        try:
            await self.add_request(request=request)
            end_time: Final = time.monotonic() + timeout
            while time.monotonic() < end_time:
                if await self.poll(request=request, health_deployments=await get_healthy_deployments()):
                    return
                await asyncio.sleep(self.polling_interval)
        finally:
            await asyncio.shield(self.remove_request(request_id=request.request_id, model_name=request.model_name))
        raise Timeout(message="Request timed out while polling queue", model=request.model_name, llm_provider="openai")

    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.
        """
        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
        await self.save_queue(queue=filtered_queue, model_name=model_name)
        print_verbose(f"Removed request_id: {request_id} from queue for model: {model_name}")

    async def peek(self, id: str, model_name: str, health_deployments: list) -> bool:
        """Return if the id is at the top of the queue. Don't pop the value from heap."""
        queue: Final = await self.get_queue(model_name=model_name)
        if not queue:
            raise Exception(f"Incorrectly setup. Queue is invalid. Queue={queue}")

        # ------------
        # Setup values
        # ------------

        # Check if the id is at the top of the heap
        if queue[0][1] == id:
            return True

        return False

    def get_queue_status(self):
        """Get the status of items in the queue"""
        return self.queue

    @with_service_target(SCHEDULER_QUEUE_TARGET)
    async def get_queue(self, model_name: str) -> list[QueueEntry]:
        """
        Return a queue for that specific model group.

        Redis hands the queue back as JSON lists, so every entry is validated into the
        (priority, request_id) tuple the heap operations compare against.
        """
        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)
            if not isinstance(response, list):
                return []
            return _QUEUE_ENTRIES.validate_python(response)
        return self.queue

    @with_service_target(SCHEDULER_QUEUE_TARGET)
    async def save_queue(self, queue: list[QueueEntry], model_name: str) -> None:
        """
        Save the updated queue of the model group
        """
        if self.cache is not None:
            _cache_key: Final = f"{SchedulerCacheKeys.queue.value}:{model_name}"
            await self.cache.async_set_cache(key=_cache_key, value=queue)