diff --git a/docs/scheduler.md b/docs/scheduler.md index 6f1c99e1..cb9c2a57 100644 --- a/docs/scheduler.md +++ b/docs/scheduler.md @@ -294,12 +294,12 @@ own capacity checks. With no selected work, no batch is returned. | Restore | If allocation succeeds, copy KV and optional hidden/token tensors to device, synchronize the current CUDA stream, delete the CPU snapshot and enqueue the saved phase. If allocation fails, retain the snapshot and defer. | | Abort a suspended request | Discard its snapshot, release runner resources/RNG and abort it through the scheduler. | -Deferred IDs are put back at the front in their existing order. Fresh admission -has `max_admission_bypasses` accounting to prevent unlimited bypass of blocked -requests. Successful restoration takes an early continuation: it does not use -the common bypass-counter update. Blocked restoration also takes an early -continuation rather than checking the fresh-allocation bypass limit. These are -current branch behaviors, not a guarantee of starvation-free restoration. +Deferred IDs are put back at the front in their existing order. Both fresh +admission and restoration count as bypasses of earlier blocked requests. A +failed fresh admission or blocked restoration stops the scan when its own +counter reaches `max_admission_bypasses`; a successful admission or restoration +increments the deferred counters and stops when any of them reaches the limit. +This bounds repeated bypasses; it is not a general starvation-freedom guarantee. #### Implemented recovery versus possible extensions diff --git a/tests/test_prefill_livelock.py b/tests/test_prefill_livelock.py new file mode 100644 index 00000000..9a0b9a71 --- /dev/null +++ b/tests/test_prefill_livelock.py @@ -0,0 +1,211 @@ +"""Regression coverage for priority admission under staged prefill pressure. + +A waiting request can be bypassed while other requests are admitted or +restored. The admission scan must count both kinds of progress against +``max_admission_bypasses``; otherwise a blocked resumption can be retried +forever while the requests already running never get a chance to drain. +""" + +from collections import defaultdict + +import pytest + +from tests.helpers import tiny_ouro_model +from vllm_rlt import ( + CacheConfig, + ExecutionConfig, + ExitConfig, + SamplingParams, + SchedulerConfig, +) +from vllm_rlt.core.kv_cache_manager import KVCacheManager +from vllm_rlt.core.scheduler import Scheduler +from vllm_rlt.engine.stage_engine import StageEngine +from vllm_rlt.request import Request, RequestPhase + +REQUESTS = ( + (0, "r0", [24, 25, 9, 13], 1, 1), + (2, "r1", [6], 2, 2), + (2, "r2", [16, 52, 33], 1, 1), + (3, "r3", [26, 42, 2, 30], 5, 1), + (5, "r4", [30, 25, 32, 37, 13, 58, 54, 26], 3, 1), + (5, "r5", [32, 15], 5, 0), + (6, "r6", [45], 3, 2), + (8, "r7", [34, 27, 31, 58, 59], 2, 1), + (10, "r8", [47, 8, 43, 17, 7, 53, 5], 5, 2), + (11, "r9", [40, 54, 25, 7, 43, 4, 22], 1, 1), +) + +REVIEWER_REQUESTS = ( + (0, "r0", [6], 4, 1), + (4, "r1", [24, 54], 4, 0), + (5, "r2", [48, 52, 43], 6, 1), + (6, "r3", [17, 39, 14, 39, 3], 3, 1), + (8, "r4", [44, 11, 28, 41, 26, 52, 47, 56, 33, 24], 4, 1), + (8, "r5", [29, 33, 18, 58, 3, 56, 2, 24, 30], 6, 2), + (9, "r6", [59, 25, 28, 58, 57, 34], 2, 1), + (14, "r7", [36, 12, 16], 3, 2), + (18, "r8", [2, 12, 21, 12], 6, 0), + (18, "r9", [33, 33, 24], 3, 2), +) + + +def run_workload( + requests, + *, + mode="refill", + preemption=True, + wavefront=True, + asynchronous=False, + delayed=False, +): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(40, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + mode=mode, + policy="priority", + enable_preemption=preemption, + max_num_seqs=4, + max_num_batched_tokens=6, + prefill_chunk_size=4, + wavefront_chunked_prefill=wavefront, + ), + exit_config=ExitConfig("ouro_delayed" if delayed else "ouro"), + execution_config=ExecutionConfig(async_scheduling=asynchronous), + ) + arrivals = defaultdict(list) + for step, request_id, prompt, max_tokens, priority in requests: + arrivals[step].append((request_id, prompt, max_tokens, priority)) + finished = {} + engine_steps = 0 + try: + for step in range(500): + for request_id, prompt, max_tokens, priority in arrivals[step]: + engine.add_request( + request_id, + prompt, + SamplingParams( + max_tokens=max_tokens, + ignore_eos=True, + priority=priority, + max_loops=4, + exit_threshold=0.6, + ), + ) + if not engine.has_unfinished_requests(): + if step > requests[-1][0]: + break + continue + outputs = engine.step() + engine_steps += 1 + for output in outputs: + if output.finished: + finished[output.request_id] = ( + output.token_ids, + output.exit_depths, + output.finish_reason, + ) + assert len(finished) == len(requests), ( + f"only {len(finished)}/{len(requests)} requests finished in 500 steps; " + f"preemptions={engine.preemption.preemptions if engine.preemption else 0}" + ) + assert set(finished) == {request_id for _, request_id, *_ in requests} + assert not engine.has_unfinished_requests() + assert engine.cache_manager.num_used_blocks == 0 + if engine.preemption is not None: + assert not engine.preemption.snapshots + return ( + finished, + engine.preemption.preemptions if engine.preemption else 0, + engine_steps, + ) + finally: + engine.close() + + +@pytest.mark.parametrize("requests", [REQUESTS, REVIEWER_REQUESTS], ids=["seed-8", "reviewer"]) +@pytest.mark.parametrize("mode,wavefront", [("refill", True), ("mixed", True), ("mixed", False)]) +def test_priority_staged_prefill_drains_under_cache_pressure(requests, mode, wavefront): + reference, _, _ = run_workload(requests, preemption=False) + actual, preemptions, engine_steps = run_workload(requests, mode=mode, wavefront=wavefront) + assert preemptions > 0 + assert engine_steps < 500 + assert actual == reference + + +def test_priority_wavefront_prefill_drains_with_async_scheduling(): + reference, _, _ = run_workload(REQUESTS, preemption=False, delayed=True) + actual, preemptions, engine_steps = run_workload(REQUESTS, asynchronous=True, delayed=True) + assert preemptions > 0 + assert engine_steps < 500 + assert actual == reference + + +def test_resumption_counts_as_an_admission_bypass(): + engine = StageEngine( + tiny_ouro_model(), + attention_backend="torch", + cache_config=CacheConfig(24, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + policy="fcfs", + enable_preemption=True, + max_num_seqs=3, + max_admission_bypasses=1, + prefill_chunk_size=4, + ), + ) + try: + engine.add_request("resumed", [2], SamplingParams(max_tokens=1, priority=10)) + engine.add_request("running-1", [1] * 4, SamplingParams(max_tokens=1)) + engine.add_request("running-2", [1] * 4, SamplingParams(max_tokens=1)) + engine.scheduler._admit() + engine.add_request("long", [1] * 9, SamplingParams(max_tokens=1)) + assert engine.scheduler.preempt(engine.scheduler.requests["long"]) + + engine.scheduler._admit() + assert engine.scheduler.requests["long"].phase == RequestPhase.WAITING + assert engine.scheduler.requests["long"].admission_bypasses == 1 + assert engine.scheduler.requests["resumed"].phase == RequestPhase.PREFILL + + engine.abort_request("running-1") + engine.scheduler._admit() + assert engine.scheduler.requests["long"].phase == RequestPhase.WAITING + engine.abort_request("running-2") + engine.scheduler._admit() + assert engine.scheduler.requests["long"].phase == RequestPhase.PREFILL + finally: + engine.close() + + +def test_blocked_restoration_respects_the_bypass_limit(): + cache = KVCacheManager( + num_layers=1, + num_kv_heads=1, + head_dim=4, + num_blocks=16, + block_size=2, + max_loops=2, + incremental_allocation=True, + ) + scheduler = Scheduler( + SchedulerConfig( + policy="fcfs", + max_num_seqs=4, + max_admission_bypasses=1, + prefill_chunk_size=4, + ), + cache, + ) + scheduler.add_request(Request("blocked", [1, 2, 3, 4], SamplingParams(max_tokens=1))) + scheduler.add_request(Request("fresh", [5], SamplingParams(max_tokens=1))) + scheduler.requests["blocked"].admission_bypasses = 1 + scheduler.resume_callback = lambda request: False if request.request_id == "blocked" else None + + scheduler._admit() + + assert scheduler.requests["blocked"].phase == RequestPhase.WAITING + assert scheduler.requests["blocked"].admission_bypasses == 1 + assert scheduler.requests["fresh"].phase == RequestPhase.WAITING + assert cache.num_used_blocks == 0 diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index 637cb0f1..14d95950 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -475,8 +475,9 @@ def _admit(self): """Admit waiting requests: order -> slot -> restore/allocate -> defer/enter. A blocked long request may let a short one pass. Each successful fresh - admission increments earlier blocked requests' bypass counts; reaching - the limit stops further bypasses so active requests can drain. + admission or resumption increments earlier blocked requests' bypass + counts; fresh admission and blocked resumption stop the scan at the + limit so active requests can drain. """ if not self.queues[RequestPhase.WAITING]: return @@ -515,25 +516,27 @@ def _admit(self): self._enter_prefill(request) else: resume_result = self._try_resume(request) - if resume_result == ResumeResult.RESTORED: - active_count += 1 - continue if resume_result == ResumeResult.BLOCKED: deferred.append(request_id) + if request.admission_bypasses >= self.config.max_admission_bypasses: + break continue - plan = self._plan_admission(request, active_count) - admitted = self._commit_admission(request, plan) - while not admitted and self.preempt(request, waiting_only=True): - # Recompute prefix claims and growth reservations after - # releasing a resident allocation; the previous plan is stale. + # RESTORED skips fresh allocation but shares the accounting + # below, just like an admitted fresh request. + if resume_result == ResumeResult.NOT_PREEMPTED: plan = self._plan_admission(request, active_count) admitted = self._commit_admission(request, plan) - if not admitted: - deferred.append(request_id) - if request.admission_bypasses >= self.config.max_admission_bypasses: - break - continue + while not admitted and self.preempt(request, waiting_only=True): + # Recompute prefix claims and growth reservations after + # releasing a resident allocation; the previous plan is stale. + plan = self._plan_admission(request, active_count) + admitted = self._commit_admission(request, plan) + if not admitted: + deferred.append(request_id) + if request.admission_bypasses >= self.config.max_admission_bypasses: + break + continue active_count += 1 for blocked_id in deferred: