diff --git a/README.md b/README.md index 9a8b365f0..7cba09743 100644 --- a/README.md +++ b/README.md @@ -169,6 +169,24 @@ > 注意：`--cache-dir` 应指向包含 `ceval___ceval-exam` 和 `cais___mmlu` 等数据集子目录的父目录，而不是直接指向这些子目录 - 试验中功能 + - 请求优先级调度 + + OpenAI 兼容接口接受可选的整数参数 `priority`，范围为 0 到 10，数值越大越优先。默认值为 0；相同有效优先级的请求保持 FIFO。该功能只调整等待请求进入 Prefill 的顺序，不抢占正在运行的请求。 + + ```json + { + "model": "model-name", + "messages": [{"role": "user", "content": "Hello"}], + "priority": 8 + } + ``` + + 为避免低优先级请求饥饿，等待请求默认每 5 秒提升一级有效优先级。可在启动服务时调整 Aging 间隔： + + ```bash + python python/infinilm/server/inference_server.py --model= --priority-aging-interval=10 + ``` + - Warm Up ```bash python examples/bench.py --device nvidia --model= --warmup diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 4ae7665c0..13e70f5a1 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -5,6 +5,7 @@ import shutil import warnings +from infinilm.config.engine_config import DEFAULT_PRIORITY_AGING_INTERVAL from infinilm.moe_config import MOE_EP_BACKEND_HELP @@ -81,6 +82,7 @@ def __init__(self): self.num_blocks = self.args.num_blocks self.block_size = self.args.block_size self.max_cache_len = self.args.max_cache_len + self.priority_aging_interval = self.args.priority_aging_interval self.kv_cache_dtype = self.args.kv_cache_dtype self.skip_load = self.args.skip_load self.weight_load_mode = self.args.weight_load_mode @@ -283,6 +285,12 @@ def _add_common_args(self): self.parser.add_argument( "--max-cache-len", type=int, default=4096, help="maximum cache length" ) + self.parser.add_argument( + "--priority-aging-interval", + type=float, + default=DEFAULT_PRIORITY_AGING_INTERVAL, + help="seconds before a waiting request gains one priority level", + ) self.parser.add_argument( "--kv-cache-dtype", type=str, diff --git a/python/infinilm/config/engine_config.py b/python/infinilm/config/engine_config.py index bccb7e758..78e5e92bf 100644 --- a/python/infinilm/config/engine_config.py +++ b/python/infinilm/config/engine_config.py @@ -1,8 +1,25 @@ +import math from dataclasses import dataclass from typing import Optional from infinilm.config.kv_transfer import KVTransferConfig +DEFAULT_PRIORITY_AGING_INTERVAL = 5.0 + + +def validate_priority_aging_interval(interval: float) -> float: + """Validate and normalize the priority aging interval.""" + if ( + isinstance(interval, bool) + or not isinstance(interval, (int, float)) + or not math.isfinite(interval) + or interval <= 0 + ): + raise ValueError( + "`priority_aging_interval` must be a finite number greater than zero." + ) + return float(interval) + @dataclass class EngineConfig: @@ -27,6 +44,7 @@ class EngineConfig: num_blocks: Number of KV cache blocks (only for paged cache). block_size: Size of each KV cache block (only for paged cache). max_cache_len: Maximum sequence length (only for static cache). + priority_aging_interval: Seconds before a waiting request gains one priority level. enable_prefix_caching: Whether to reuse KV cache across requests. temperature: Default sampling temperature. top_p: Default top-p sampling parameter. @@ -69,6 +87,7 @@ class EngineConfig: use_legacy_moe: bool = False kv_transfer_config: Optional[KVTransferConfig] = None enable_prefix_caching: bool = True + priority_aging_interval: float = DEFAULT_PRIORITY_AGING_INTERVAL def __post_init__(self) -> None: if self.num_draft_tokens < 1: @@ -84,6 +103,9 @@ def __post_init__(self) -> None: if self.weight_load_mode not in {"async", "sync"}: raise ValueError("weight_load_mode must be either 'async' or 'sync'") + self.priority_aging_interval = validate_priority_aging_interval( + self.priority_aging_interval + ) if ( self.kv_transfer_config is not None diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59d5a1eca..556772082 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -16,7 +16,10 @@ import janus -from infinilm.config.engine_config import EngineConfig +from infinilm.config.engine_config import ( + DEFAULT_PRIORITY_AGING_INTERVAL, + EngineConfig, +) from infinilm.config.kv_transfer import KVTransferConfig from infinilm.infer_engine import model_uses_mamba_cache, read_hf_config from infinilm.kv_connector import KVConnectorFactory, KVConnectorRole @@ -70,6 +73,7 @@ def __init__(self, config: EngineConfig): self.scheduler = StaticScheduler( max_cache_len=config.max_cache_len, enable_prefix_caching=config.enable_prefix_caching, + priority_aging_interval=config.priority_aging_interval, ) logger.info( f"Using Static KV Cache with max_cache_len={config.max_cache_len}" @@ -110,6 +114,7 @@ def __init__(self, config: EngineConfig): has_mamba_cache=has_mamba_cache, num_mamba_cache_blocks=num_mamba_cache_blocks, enable_prefix_caching=config.enable_prefix_caching, + priority_aging_interval=config.priority_aging_interval, ) logger.info(f"Using Paged KV Cache with num_blocks={config.num_blocks}") if has_mamba_cache: @@ -366,6 +371,7 @@ def __init__( skip_load: bool = False, use_legacy_moe: bool = False, enable_prefix_caching: bool = True, + priority_aging_interval: float = DEFAULT_PRIORITY_AGING_INTERVAL, ): """Initialize LLM. @@ -387,6 +393,8 @@ def __init__( attn_backend: Attention backend to use ('default', 'flash-attn'). use_mla: Whether to use DeepSeek V2 MLA attention when supported. weight_load_mode: Weight loading mode across tensor-parallel workers. + priority_aging_interval: Seconds before a waiting request gains one + priority level. """ config = EngineConfig( model_path=model_path, @@ -407,6 +415,7 @@ def __init__( num_blocks=num_blocks, block_size=block_size, max_cache_len=max_cache_len, + priority_aging_interval=priority_aging_interval, temperature=temperature, top_p=top_p, top_k=top_k, @@ -594,6 +603,7 @@ def __init__( weight_load_mode: str = "async", use_legacy_moe: bool = False, enable_prefix_caching: bool = True, + priority_aging_interval: float = DEFAULT_PRIORITY_AGING_INTERVAL, ): """Initialize AsyncLLMEngine. @@ -619,6 +629,8 @@ def __init__( use_mla: Whether to use DeepSeek V2 MLA attention when supported. skip_load: Whether to skip loading model weights. weight_load_mode: Weight loading mode across tensor-parallel workers. + priority_aging_interval: Seconds before a waiting request gains one + priority level. """ config = EngineConfig( model_path=model_path, @@ -639,6 +651,7 @@ def __init__( num_blocks=num_blocks, block_size=block_size, max_cache_len=max_cache_len, + priority_aging_interval=priority_aging_interval, temperature=temperature, top_p=top_p, top_k=top_k, @@ -784,6 +797,7 @@ def add_request( request_id: Optional[str] = None, # For server use request_data: Optional[dict] = None, + priority: int = 0, ) -> InferenceRequest: """Add a request to the engine. @@ -813,6 +827,7 @@ def add_request( sampling_params: Sampling parameters. request_id: Optional request ID. request_data: Optional request data dict (for server use). + priority: Request admission priority from 0 to 10. Returns: The created InferenceRequest object. @@ -878,6 +893,7 @@ def add_request( sampling_params=sampling_params, eos_token_ids=self.engine.eos_token_ids, request_data=request_data, + priority=priority, ) if request_data and "kv_transfer_params" in request_data: @@ -897,6 +913,7 @@ def add_chat_request( request_id: Optional[str] = None, request_data: Optional[dict] = None, add_generation_prompt: bool = True, + priority: int = 0, **kwargs, ) -> InferenceRequest: """Add a chat request to the engine. @@ -906,6 +923,7 @@ def add_chat_request( sampling_params: Sampling parameters. request_id: Optional request ID. request_data: Optional request data dict. + priority: Request admission priority from 0 to 10. Returns: The created InferenceRequest object. @@ -918,6 +936,7 @@ def add_chat_request( sampling_params=sampling_params, request_id=request_id, request_data=request_data, + priority=priority, ) async def stream_request( diff --git a/python/infinilm/llm/priority_scheduling.py b/python/infinilm/llm/priority_scheduling.py new file mode 100644 index 000000000..410039b0d --- /dev/null +++ b/python/infinilm/llm/priority_scheduling.py @@ -0,0 +1,113 @@ +"""Priority ordering helpers for inference request admission.""" + +import queue +import time +from collections import Counter +from threading import Lock +from typing import Callable + +from infinilm.llm.request import InferenceRequest + + +class PrioritySchedulingStats: + """Track cumulative priority admission statistics.""" + + def __init__(self) -> None: + self.admitted_requests = 0 + self.aged_admissions = 0 + self.total_wait_time_seconds = 0.0 + self.max_wait_time_seconds = 0.0 + self.admitted_by_priority: Counter[int] = Counter() + self._lock = Lock() + + def record_admission( + self, + request: InferenceRequest, + now: float, + aging_interval: float, + ) -> tuple[float, int, bool]: + """Record one admission and return its scheduling details.""" + wait_time = request_wait_time(request, now) + admission_priority = effective_priority(request, now, aging_interval) + aged = admission_priority > request.priority + + with self._lock: + self.admitted_requests += 1 + self.aged_admissions += int(aged) + self.total_wait_time_seconds += wait_time + self.max_wait_time_seconds = max(self.max_wait_time_seconds, wait_time) + self.admitted_by_priority[request.priority] += 1 + return wait_time, admission_priority, aged + + def snapshot(self) -> dict: + """Return a JSON-serializable statistics snapshot.""" + with self._lock: + average_wait_time = ( + self.total_wait_time_seconds / self.admitted_requests + if self.admitted_requests + else 0.0 + ) + return { + "admitted_requests": self.admitted_requests, + "aged_admissions": self.aged_admissions, + "average_wait_time_seconds": average_wait_time, + "max_wait_time_seconds": self.max_wait_time_seconds, + "admitted_by_priority": dict(sorted(self.admitted_by_priority.items())), + } + + +def mark_request_enqueued( + request: InferenceRequest, + sequence: int, + clock: Callable[[], float] = time.monotonic, +) -> None: + """Record stable scheduling metadata on a request's first enqueue.""" + if request.scheduling_enqueue_time is None: + request.scheduling_enqueue_time = clock() + if request.scheduling_sequence is None: + request.scheduling_sequence = sequence + + +def effective_priority( + request: InferenceRequest, + now: float, + aging_interval: float, +) -> int: + """Return the request priority after applying wait-time aging.""" + if aging_interval <= 0: + raise ValueError("`aging_interval` must be greater than zero.") + wait_time = request_wait_time(request, now) + return request.priority + int(wait_time // aging_interval) + + +def request_wait_time(request: InferenceRequest, now: float) -> float: + """Return how long a request has waited for admission.""" + enqueue_time = request.scheduling_enqueue_time + if enqueue_time is None: + enqueue_time = now + return max(0.0, now - enqueue_time) + + +def drain_priority_ordered( + sync_queue, + aging_interval: float, + clock: Callable[[], float] = time.monotonic, +) -> list[InferenceRequest]: + """Drain a finite queue snapshot and return requests in admission order.""" + requests = [] + for _ in range(sync_queue.qsize()): + try: + requests.append(sync_queue.get_nowait()) + except queue.Empty: + break + + now = clock() + requests.sort( + key=lambda request: ( + -effective_priority(request, now, aging_interval), + request.scheduling_sequence + if request.scheduling_sequence is not None + else float("inf"), + ) + ) + return requests diff --git a/python/infinilm/llm/request.py b/python/infinilm/llm/request.py index f4612b816..14cf392cf 100644 --- a/python/infinilm/llm/request.py +++ b/python/infinilm/llm/request.py @@ -21,6 +21,22 @@ logger = logging.getLogger(__name__) +MIN_REQUEST_PRIORITY = 0 +MAX_REQUEST_PRIORITY = 10 + + +def validate_request_priority(priority: int) -> int: + """Validate and return a request priority.""" + error = ( + f"`priority` must be an integer between {MIN_REQUEST_PRIORITY} " + f"and {MAX_REQUEST_PRIORITY}." + ) + if isinstance(priority, bool) or not isinstance(priority, int): + raise ValueError(error) + if not MIN_REQUEST_PRIORITY <= priority <= MAX_REQUEST_PRIORITY: + raise ValueError(error) + return priority + class _SequenceView(Sequence): """Live read-only view over an internally mutable list.""" @@ -158,6 +174,7 @@ def __init__( request_data: Optional[dict] = None, *, has_multimodal_inputs: bool = False, + priority: int = MIN_REQUEST_PRIORITY, ): self.arrival_time: float = arrival_time or time.time() self.finished_time: Optional[float] = None @@ -174,7 +191,9 @@ def __init__( self.has_multimodal_inputs: bool = has_multimodal_inputs or bool( mm_token_index_mappings ) - self.priority: int = 0 + self.priority = validate_request_priority(priority) + self.scheduling_enqueue_time: Optional[float] = None + self.scheduling_sequence: Optional[int] = None # Sampling & stopping criteria self.sampling_params: SamplingParams = sampling_params or SamplingParams() diff --git a/python/infinilm/llm/scheduler.py b/python/infinilm/llm/scheduler.py index c10b55f2f..d51f4d27e 100644 --- a/python/infinilm/llm/scheduler.py +++ b/python/infinilm/llm/scheduler.py @@ -4,11 +4,23 @@ import logging import queue -from typing import List, Optional +import time +from collections import deque +from itertools import count +from typing import Callable, List, Optional import janus +from infinilm.config.engine_config import ( + DEFAULT_PRIORITY_AGING_INTERVAL, + validate_priority_aging_interval, +) from infinilm.llm.cache_manager import BlockManager, MambaCacheManager +from infinilm.llm.priority_scheduling import ( + PrioritySchedulingStats, + drain_priority_ordered, + mark_request_enqueued, +) from infinilm.llm.request import InferenceRequest, RequestStatus logger = logging.getLogger(__name__) @@ -71,10 +83,18 @@ def __init__( has_mamba_cache: bool = False, num_mamba_cache_blocks: int | None = None, enable_prefix_caching: bool = True, + priority_aging_interval: float = DEFAULT_PRIORITY_AGING_INTERVAL, + clock: Callable[[], float] = time.monotonic, ): + self.priority_aging_interval = validate_priority_aging_interval( + priority_aging_interval + ) self.waiting_queue = janus.Queue() self.running_queue = janus.Queue() self.max_batch_size = max_batch_size + self._clock = clock + self._enqueue_sequence = count() + self.priority_scheduling_stats = PrioritySchedulingStats() self.finished_receiving_kv_req_ids: set[str] = set() self.failed_receiving_kv_req_ids: set[str] = set() @@ -106,6 +126,7 @@ def add_request(self, request: InferenceRequest): and not request.has_multimodal_inputs, ) request.status = RequestStatus.WAITING + mark_request_enqueued(request, next(self._enqueue_sequence), self._clock) self.waiting_queue.sync_q.put(request) def _exceeds_token_budget( @@ -133,16 +154,21 @@ def schedule(self) -> Optional[SchedulerOutput]: is_prefill = False current_num_batched_tokens = 0 current_prefill_extra_blocks = 0 + waiting_requests = deque( + drain_priority_ordered( + self.waiting_queue.sync_q, + self.priority_aging_interval, + self._clock, + ) + ) # Process Waiting queue (prefill phase) while ( - len(scheduled_requests) < self.max_batch_size + waiting_requests + and len(scheduled_requests) < self.max_batch_size and current_num_batched_tokens < self.max_num_batched_tokens ): - try: - req = self.waiting_queue.sync_q.get_nowait() - except queue.Empty: - break + req = waiting_requests.popleft() # Skip requests that were already finished (e.g., timed out/canceled while waiting) if req.is_finished(): self.complete_requests([req]) @@ -272,6 +298,23 @@ def schedule(self) -> Optional[SchedulerOutput]: break self.commit_computed_tokens(req, req.num_computed_tokens) + now = self._clock() + wait_time, admission_priority, aged = ( + self.priority_scheduling_stats.record_admission( + req, now, self.priority_aging_interval + ) + ) + logger.debug( + "Admitting request %s with priority=%d, effective_priority=%d, " + "wait_time=%.3fs, waiting_queue_size=%d, aged=%s.", + req.request_id[:8], + req.priority, + admission_priority, + wait_time, + len(waiting_requests) + self.waiting_queue.sync_q.qsize(), + aged, + ) + if load_kv_async: req.status = RequestStatus.WAITING_FOR_REMOTE_KVS self.remote_kv_requests[req.request_id] = req @@ -291,6 +334,8 @@ def schedule(self) -> Optional[SchedulerOutput]: if deferred_requests: for req in deferred_requests: self.waiting_queue.sync_q.put(req) + for req in waiting_requests: + self.waiting_queue.sync_q.put(req) # Return prefill batch if any waiting requests were scheduled if scheduled_requests: @@ -584,6 +629,9 @@ def get_cache_stats(self) -> dict: "num_free_blocks": self.cache_manager.get_num_free_blocks(), "usable_blocks": self.cache_manager.get_total_usable_blocks(), "num_used_blocks": len(self.cache_manager.used_block_ids), + "waiting_queue_size": self.waiting_queue.sync_q.qsize(), + "running_queue_size": self.running_queue.sync_q.qsize(), + "priority_scheduling": self.priority_scheduling_stats.snapshot(), } if self.mamba_cache_manager is not None: stats.update( diff --git a/python/infinilm/llm/static_scheduler.py b/python/infinilm/llm/static_scheduler.py index 876c1e8c6..32811910a 100644 --- a/python/infinilm/llm/static_scheduler.py +++ b/python/infinilm/llm/static_scheduler.py @@ -3,12 +3,23 @@ """ import logging -import queue -from typing import List, Optional +import time +from collections import deque +from itertools import count +from typing import Callable, List, Optional import janus +from infinilm.config.engine_config import ( + DEFAULT_PRIORITY_AGING_INTERVAL, + validate_priority_aging_interval, +) from infinilm.llm.prefix_cache import BlockHash +from infinilm.llm.priority_scheduling import ( + PrioritySchedulingStats, + drain_priority_ordered, + mark_request_enqueued, +) from infinilm.llm.request import ( FinishReason, InferenceRequest, @@ -51,12 +62,20 @@ def __init__( self, max_cache_len: int = 4096, enable_prefix_caching: bool = True, + priority_aging_interval: float = DEFAULT_PRIORITY_AGING_INTERVAL, + clock: Callable[[], float] = time.monotonic, ): + self.priority_aging_interval = validate_priority_aging_interval( + priority_aging_interval + ) self.waiting_queue = janus.Queue() self.running_request: Optional[InferenceRequest] = None self.max_cache_len = max_cache_len self.enable_prefix_caching = enable_prefix_caching self.cached_block_hashes: List[BlockHash] = [] + self._clock = clock + self._enqueue_sequence = count() + self.priority_scheduling_stats = PrioritySchedulingStats() def add_request(self, request: InferenceRequest): if request is not None: @@ -67,10 +86,12 @@ def add_request(self, request: InferenceRequest): self.enable_prefix_caching and not request.has_multimodal_inputs, ) request.status = RequestStatus.WAITING + mark_request_enqueued(request, next(self._enqueue_sequence), self._clock) self.waiting_queue.sync_q.put(request) def schedule(self) -> Optional[StaticSchedulerOutput]: """Schedule and return single request to execute.""" + waiting_requests = None while True: # Case 1: Continue running request (decode phase) if self.running_request is not None: @@ -107,10 +128,17 @@ def schedule(self) -> Optional[StaticSchedulerOutput]: return StaticSchedulerOutput(scheduled_requests=[req], is_prefill=False) # Case 2: Get new request from waiting queue (prefill phase) - try: - req = self.waiting_queue.sync_q.get_nowait() - except queue.Empty: + if waiting_requests is None: + waiting_requests = deque( + drain_priority_ordered( + self.waiting_queue.sync_q, + self.priority_aging_interval, + self._clock, + ) + ) + if not waiting_requests: return None + req = waiting_requests.popleft() if req.is_finished(): continue @@ -166,6 +194,24 @@ def schedule(self) -> Optional[StaticSchedulerOutput]: req.status = RequestStatus.RUNNING self.running_request = req + now = self._clock() + wait_time, admission_priority, aged = ( + self.priority_scheduling_stats.record_admission( + req, now, self.priority_aging_interval + ) + ) + logger.debug( + "Admitting request %s with priority=%d, effective_priority=%d, " + "wait_time=%.3fs, waiting_queue_size=%d, aged=%s.", + req.request_id[:8], + req.priority, + admission_priority, + wait_time, + len(waiting_requests) + self.waiting_queue.sync_q.qsize(), + aged, + ) + for waiting_req in waiting_requests: + self.waiting_queue.sync_q.put(waiting_req) return StaticSchedulerOutput( scheduled_requests=[req], is_prefill=True, prefix_hit_len=prefix_hit_len ) @@ -212,4 +258,5 @@ def get_cache_stats(self) -> dict: self.running_request.request_id if self.running_request else None ), "waiting_queue_size": self.waiting_queue.sync_q.qsize(), + "priority_scheduling": self.priority_scheduling_stats.snapshot(), } diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py index 462c31084..4cb5aa3cc 100644 --- a/python/infinilm/server/inference_server.py +++ b/python/infinilm/server/inference_server.py @@ -17,7 +17,9 @@ from fastapi.responses import JSONResponse, StreamingResponse from infinilm.base_config import BaseConfig from infinilm.config import KVTransferConfig +from infinilm.config.engine_config import DEFAULT_PRIORITY_AGING_INTERVAL from infinilm.llm import AsyncLLMEngine, FinishReason, SamplingParams +from infinilm.llm.request import validate_request_priority from infinilm.moe_config import configure_moe_ep_backend logger = logging.getLogger(__name__) @@ -125,6 +127,7 @@ def __init__( kv_transfer_config: Optional[KVTransferConfig] = None, enable_prefix_caching: bool = True, pre_transpose: bool = False, + priority_aging_interval: float = DEFAULT_PRIORITY_AGING_INTERVAL, ): """Initialize inference server. @@ -154,6 +157,8 @@ def __init__( weight_load_mode: Weight loading mode across tensor-parallel workers. ignore_eos: Whether to ignore EOS tokens during generation. kv_transfer_config: Optional configuration for the KV transfer mechanism. + priority_aging_interval: Seconds before a waiting request gains one + priority level. """ self.model_path = model_path # vLLM-like served model id: directory name of model_path @@ -174,6 +179,7 @@ def __init__( self.num_blocks = num_blocks self.block_size = block_size self.max_cache_len = max_cache_len + self.priority_aging_interval = priority_aging_interval self.temperature = temperature self.top_p = top_p self.top_k = top_k @@ -221,6 +227,7 @@ async def lifespan(app: FastAPI): num_blocks=self.num_blocks, block_size=self.block_size, max_cache_len=self.max_cache_len, + priority_aging_interval=self.priority_aging_interval, temperature=self.temperature, top_p=self.top_p, top_k=self.top_k, @@ -269,6 +276,11 @@ async def chat_completions(request: Request): # Normalize messages to handle multimodal content (list format) data["messages"] = data.get("messages", []) + try: + data["priority"] = validate_request_priority(data.get("priority", 0)) + except ValueError as error: + return JSONResponse(content={"error": str(error)}, status_code=400) + stream = data.get("stream", False) request_id = f"cmpl-{uuid.uuid4().hex}" @@ -399,6 +411,7 @@ async def _stream_chat(self, request_id: str, data: dict, http_request: Request) sampling_params=sampling_params, request_id=request_id, request_data=data, + priority=data["priority"], add_generation_prompt=bool(data.get("add_generation_prompt", True)), chat_template_kwargs=data.get("chat_template_kwargs") or {}, ) @@ -504,6 +517,7 @@ async def _chat(self, request_id: str, data: dict, http_request: Request): sampling_params=sampling_params, request_id=request_id, request_data=data, + priority=data["priority"], add_generation_prompt=bool(data.get("add_generation_prompt", True)), chat_template_kwargs=data.get("chat_template_kwargs") or {}, ) @@ -653,6 +667,7 @@ def main(): num_blocks=cfg.num_blocks, block_size=cfg.block_size, max_cache_len=cfg.max_cache_len, + priority_aging_interval=cfg.priority_aging_interval, temperature=cfg.temperature, top_p=cfg.top_p, top_k=cfg.top_k, diff --git a/test/llm/test_priority_scheduling.py b/test/llm/test_priority_scheduling.py new file mode 100644 index 000000000..8166645c0 --- /dev/null +++ b/test/llm/test_priority_scheduling.py @@ -0,0 +1,260 @@ +import unittest +from unittest.mock import patch + +import janus +from infinilm.llm.priority_scheduling import ( + drain_priority_ordered, + mark_request_enqueued, +) +from infinilm.llm.request import InferenceRequest, RequestStatus +from infinilm.llm.sampling_params import SamplingParams +from infinilm.llm.scheduler import Scheduler +from infinilm.llm.static_scheduler import StaticScheduler + + +class FakeClock: + def __init__(self, now: float = 0.0): + self.now = now + + def __call__(self) -> float: + return self.now + + +def make_request(request_id: str, priority: int = 0) -> InferenceRequest: + return InferenceRequest( + request_id=request_id, + prompt_token_ids=[1, 2], + sampling_params=SamplingParams(max_tokens=1), + priority=priority, + ) + + +class FakeConnector: + def get_num_new_matched_tokens(self, request, num_local_computed_tokens): + return 1, True + + def update_state_after_alloc( + self, request, block_table, num_external_computed_tokens, block_size + ): + pass + + def build_connector_meta(self): + return None + + +class PriorityOrderingTest(unittest.TestCase): + def test_default_priority_preserves_fifo_order(self): + clock = FakeClock() + requests = [make_request("first"), make_request("second")] + request_queue = janus.Queue() + self.addCleanup(request_queue.close) + + for sequence, request in enumerate(requests): + mark_request_enqueued(request, sequence, clock) + request_queue.sync_q.put(request) + + ordered = drain_priority_ordered(request_queue.sync_q, 5.0, clock) + + self.assertEqual( + [request.request_id for request in ordered], ["first", "second"] + ) + + def test_higher_priority_is_admitted_first(self): + clock = FakeClock() + low = make_request("low", priority=1) + high = make_request("high", priority=8) + request_queue = janus.Queue() + self.addCleanup(request_queue.close) + + mark_request_enqueued(low, 0, clock) + request_queue.sync_q.put(low) + clock.now = 1.0 + mark_request_enqueued(high, 1, clock) + request_queue.sync_q.put(high) + + ordered = drain_priority_ordered(request_queue.sync_q, 5.0, clock) + + self.assertEqual([request.request_id for request in ordered], ["high", "low"]) + + def test_aging_prevents_starvation(self): + clock = FakeClock() + old_low = make_request("old-low", priority=0) + request_queue = janus.Queue() + self.addCleanup(request_queue.close) + + mark_request_enqueued(old_low, 0, clock) + request_queue.sync_q.put(old_low) + clock.now = 55.0 + new_high = make_request("new-high", priority=10) + mark_request_enqueued(new_high, 1, clock) + request_queue.sync_q.put(new_high) + + ordered = drain_priority_ordered(request_queue.sync_q, 5.0, clock) + + self.assertEqual( + [request.request_id for request in ordered], ["old-low", "new-high"] + ) + + def test_requeue_preserves_original_age_and_fifo_position(self): + clock = FakeClock() + request = make_request("request") + + mark_request_enqueued(request, 3, clock) + clock.now = 10.0 + mark_request_enqueued(request, 9, clock) + + self.assertEqual(request.scheduling_enqueue_time, 0.0) + self.assertEqual(request.scheduling_sequence, 3) + + def test_invalid_priorities_are_rejected(self): + for priority in (-1, 11, True, 1.5, "1"): + with self.subTest(priority=priority), self.assertRaises(ValueError): + make_request("invalid", priority=priority) + + +class SchedulerPriorityTest(unittest.TestCase): + def test_invalid_aging_intervals_are_rejected(self): + for interval in (0, -1, float("nan"), float("inf"), True, "5"): + with self.subTest(interval=interval), self.assertRaises(ValueError): + StaticScheduler(priority_aging_interval=interval) + + def test_static_scheduler_selects_high_priority_request(self): + scheduler = StaticScheduler(enable_prefix_caching=False) + self.addCleanup(scheduler.waiting_queue.close) + scheduler.add_request(make_request("low", priority=1)) + scheduler.add_request(make_request("high", priority=8)) + + output = scheduler.schedule() + + self.assertEqual(output.scheduled_requests[0].request_id, "high") + + def test_canceled_high_priority_request_does_not_block_queue(self): + scheduler = StaticScheduler(enable_prefix_caching=False) + self.addCleanup(scheduler.waiting_queue.close) + scheduler.add_request(make_request("low", priority=1)) + canceled = make_request("canceled", priority=10) + scheduler.add_request(canceled) + canceled.mark_canceled() + + output = scheduler.schedule() + + self.assertEqual(output.scheduled_requests[0].request_id, "low") + + def test_timed_out_high_priority_request_does_not_block_queue(self): + scheduler = StaticScheduler(enable_prefix_caching=False) + self.addCleanup(scheduler.waiting_queue.close) + scheduler.add_request(make_request("low", priority=1)) + timed_out = make_request("timed-out", priority=10) + scheduler.add_request(timed_out) + timed_out.mark_timeout() + + output = scheduler.schedule() + + self.assertEqual(output.scheduled_requests[0].request_id, "low") + + def test_admission_stats_track_wait_time_and_aging(self): + clock = FakeClock() + scheduler = StaticScheduler( + enable_prefix_caching=False, + priority_aging_interval=5.0, + clock=clock, + ) + self.addCleanup(scheduler.waiting_queue.close) + scheduler.add_request(make_request("aged", priority=1)) + clock.now = 12.0 + + scheduler.schedule() + stats = scheduler.get_cache_stats()["priority_scheduling"] + + self.assertEqual(stats["admitted_requests"], 1) + self.assertEqual(stats["aged_admissions"], 1) + self.assertEqual(stats["admitted_by_priority"], {1: 1}) + self.assertEqual(stats["average_wait_time_seconds"], 12.0) + self.assertEqual(stats["max_wait_time_seconds"], 12.0) + + def test_paged_scheduler_selects_high_priority_request(self): + scheduler = Scheduler( + max_batch_size=1, + num_blocks=32, + block_size=4, + enable_prefix_caching=False, + ) + self.addCleanup(scheduler.waiting_queue.close) + self.addCleanup(scheduler.running_queue.close) + scheduler.add_request(make_request("low", priority=1)) + scheduler.add_request(make_request("high", priority=8)) + + output = scheduler.schedule() + + self.assertEqual(output.scheduled_requests[0].request_id, "high") + self.assertEqual(scheduler.waiting_queue.sync_q.qsize(), 1) + + def test_paged_scheduler_prioritizes_limited_batch(self): + scheduler = Scheduler( + max_batch_size=2, + num_blocks=64, + block_size=4, + enable_prefix_caching=False, + ) + self.addCleanup(scheduler.waiting_queue.close) + self.addCleanup(scheduler.running_queue.close) + scheduler.add_request(make_request("low", priority=1)) + scheduler.add_request(make_request("high", priority=8)) + scheduler.add_request(make_request("medium", priority=4)) + + output = scheduler.schedule() + stats = scheduler.get_cache_stats() + + self.assertEqual( + [request.request_id for request in output.scheduled_requests], + ["high", "medium"], + ) + self.assertEqual(stats["waiting_queue_size"], 1) + self.assertEqual( + stats["priority_scheduling"]["admitted_by_priority"], {4: 1, 8: 1} + ) + + def test_deferred_request_keeps_priority(self): + scheduler = Scheduler( + max_batch_size=1, + num_blocks=32, + block_size=4, + enable_prefix_caching=False, + ) + self.addCleanup(scheduler.waiting_queue.close) + self.addCleanup(scheduler.running_queue.close) + scheduler.add_request(make_request("low", priority=1)) + scheduler.add_request(make_request("high", priority=8)) + with patch.object(scheduler, "can_accept_request", return_value=False): + self.assertIsNone(scheduler.schedule()) + + scheduler.add_request(make_request("medium", priority=4)) + with patch.object(scheduler, "can_accept_request", return_value=True): + output = scheduler.schedule() + + self.assertEqual(output.scheduled_requests[0].request_id, "high") + + def test_remote_kv_request_is_counted_when_admitted(self): + scheduler = Scheduler( + max_batch_size=1, + num_blocks=32, + block_size=4, + connector=FakeConnector(), + enable_prefix_caching=False, + ) + self.addCleanup(scheduler.waiting_queue.close) + self.addCleanup(scheduler.running_queue.close) + request = make_request("remote", priority=7) + scheduler.add_request(request) + + output = scheduler.schedule() + stats = scheduler.get_cache_stats()["priority_scheduling"] + + self.assertEqual(output.scheduled_requests, []) + self.assertEqual(request.status, RequestStatus.WAITING_FOR_REMOTE_KVS) + self.assertEqual(stats["admitted_requests"], 1) + self.assertEqual(stats["admitted_by_priority"], {7: 1}) + + +if __name__ == "__main__": + unittest.main()