[Refactor] [0/N] Scheduler clean dead code (#5312)
Signed-off-by: R2-Y <ruiruyang2@gmail.com>
This commit is contained in:
rein yang authored and GitHub committed
1 parent
633ebbd8a4
commit
1984346bc7
7 files changed
+93
-664
No files matched your search
@@ -62,7 +62,6 @@ Engine classes for offline and online inference.
|
||||
|
||||
Core scheduling and caching components.
|
||||
|
||||
- [vllm_omni.core.sched.omni_ar_scheduler.KVCacheTransferData][]
|
||||
- [vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler][]
|
||||
- [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][]
|
||||
- [vllm_omni.core.sched.output.OmniCachedRequestData][]
|
||||
|
||||
@@ -4,6 +4,9 @@
|
||||
|
||||
These tests use mock request objects and mock queues. They do not require
|
||||
GPU, vLLM runtime, or any connector.
|
||||
|
||||
Chunk waiting (WAITING_FOR_CHUNK / process_pending_chunks) lives on
|
||||
OmniChunkTransferAdapter — see tests/distributed/omni_connectors/.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -11,6 +14,7 @@ from __future__ import annotations
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod
|
||||
@@ -19,6 +23,8 @@ from vllm_omni.core.sched.omni_scheduling_coordinator import (
|
||||
uses_full_payload_input_coordinator,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Mock helpers
|
||||
# ------------------------------------------------------------------ #
|
||||
@@ -27,7 +33,6 @@ from vllm_omni.core.sched.omni_scheduling_coordinator import (
|
||||
class _RequestStatus:
|
||||
WAITING = "waiting"
|
||||
RUNNING = "running"
|
||||
WAITING_FOR_CHUNK = "waiting_for_chunk"
|
||||
WAITING_FOR_INPUT = "waiting_for_input"
|
||||
FINISHED_STOPPED = "finished_stopped"
|
||||
|
||||
@@ -185,105 +190,12 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase):
|
||||
)
|
||||
|
||||
|
||||
class TestChunkCoordinatorStateTransition(unittest.TestCase):
|
||||
"""Test 5: process_pending_chunks transitions WAITING_FOR_CHUNK → target."""
|
||||
|
||||
def test_ready_request_transitions_to_waiting(self):
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
|
||||
|
||||
req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
|
||||
waiting = MockQueue([req])
|
||||
running: list = []
|
||||
|
||||
coord.process_pending_chunks(
|
||||
waiting,
|
||||
running,
|
||||
chunk_ready_req_ids={"r1"},
|
||||
chunk_finished_req_ids=set(),
|
||||
)
|
||||
|
||||
self.assertEqual(req.status, RequestStatus.WAITING)
|
||||
self.assertIn("r1", coord.requests_with_ready_chunks)
|
||||
|
||||
def test_non_ready_stays_waiting_for_chunk(self):
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
|
||||
|
||||
req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
|
||||
waiting = MockQueue([req])
|
||||
running: list = []
|
||||
|
||||
coord.process_pending_chunks(
|
||||
waiting,
|
||||
running,
|
||||
chunk_ready_req_ids=set(),
|
||||
chunk_finished_req_ids=set(),
|
||||
)
|
||||
|
||||
self.assertEqual(req.status, RequestStatus.WAITING_FOR_CHUNK)
|
||||
|
||||
def test_stage_0_is_noop(self):
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=0)
|
||||
req = _make_request("r1")
|
||||
waiting = MockQueue([req])
|
||||
running: list = []
|
||||
|
||||
coord.process_pending_chunks(
|
||||
waiting,
|
||||
running,
|
||||
chunk_ready_req_ids={"r1"},
|
||||
chunk_finished_req_ids=set(),
|
||||
)
|
||||
self.assertNotEqual(req.status, RequestStatus.WAITING_FOR_CHUNK)
|
||||
|
||||
|
||||
class TestChunkCoordinatorRestoreQueues(unittest.TestCase):
|
||||
"""Test 6: restore_queues returns waiting-for-chunk requests."""
|
||||
|
||||
def test_restore(self):
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
|
||||
|
||||
r1 = _make_request("r1")
|
||||
r2 = _make_request("r2")
|
||||
coord._waiting_for_chunk_waiting.append(r1)
|
||||
coord._waiting_for_chunk_running.append(r2)
|
||||
|
||||
waiting = MockQueue()
|
||||
running: list = []
|
||||
|
||||
coord.restore_queues(waiting, running)
|
||||
|
||||
self.assertIn(r1, waiting)
|
||||
self.assertIn(r2, running)
|
||||
self.assertEqual(len(coord._waiting_for_chunk_waiting), 0)
|
||||
self.assertEqual(len(coord._waiting_for_chunk_running), 0)
|
||||
|
||||
|
||||
class TestChunkCoordinatorFinishedSignal(unittest.TestCase):
|
||||
"""Test 8: chunk_finished_req_ids → finished_requests."""
|
||||
|
||||
def test_finished_signal(self):
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
|
||||
|
||||
req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
|
||||
waiting = MockQueue([req])
|
||||
running: list = []
|
||||
|
||||
coord.process_pending_chunks(
|
||||
waiting,
|
||||
running,
|
||||
chunk_ready_req_ids={"r1"},
|
||||
chunk_finished_req_ids={"r1"},
|
||||
)
|
||||
|
||||
self.assertIn("r1", coord.finished_requests)
|
||||
|
||||
|
||||
class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
|
||||
class TestCoordinatorUpdateRequestMetadata(unittest.TestCase):
|
||||
"""Test update_request_metadata applies scheduling metadata to requests."""
|
||||
|
||||
def test_ar_mode_no_longer_sets_additional_information(self):
|
||||
"""AR mode only processes scheduling metadata, not full payloads."""
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
|
||||
req = _make_request("r1")
|
||||
requests = {"r1": req}
|
||||
@@ -300,7 +212,7 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
|
||||
self.assertIsNone(getattr(req, "additional_information", None))
|
||||
|
||||
def test_generation_mode(self):
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
|
||||
req = _make_request("r1")
|
||||
req.prompt_token_ids = [0, 0, 0]
|
||||
@@ -313,7 +225,6 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
|
||||
request_metadata = {
|
||||
"r1": {
|
||||
"code_predictor_codes": [10, 20, 30],
|
||||
"left_context_size": 25,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -325,10 +236,9 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
|
||||
self.assertEqual(req._all_token_ids, [10, 20, 30])
|
||||
self.assertEqual(req._output_token_ids, [])
|
||||
self.assertIsNone(req.additional_information)
|
||||
self.assertEqual(req._omni_initial_model_buffer, {"meta": {"left_context_size": 25}})
|
||||
|
||||
def test_generation_mode_flattens_tensor_code_predictor_codes(self):
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
|
||||
req = _make_request("r1")
|
||||
req.prompt_token_ids = [9]
|
||||
@@ -349,7 +259,7 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
|
||||
self.assertEqual(req._output_token_ids, [])
|
||||
|
||||
def test_generation_mode_flattens_nested_code_predictor_codes(self):
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
|
||||
req = _make_request("r1")
|
||||
req.prompt_token_ids = [9]
|
||||
@@ -370,53 +280,30 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
|
||||
self.assertEqual(req._output_token_ids, [])
|
||||
|
||||
|
||||
class TestChunkCoordinatorPostprocess(unittest.TestCase):
|
||||
"""Test postprocess_scheduler_output clears ready chunks."""
|
||||
|
||||
def test_clear_ready(self):
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
|
||||
coord.requests_with_ready_chunks = {"r1", "r2"}
|
||||
|
||||
new_req = SimpleNamespace(req_id="r1")
|
||||
cached_reqs = SimpleNamespace(req_ids=["r2"])
|
||||
scheduler_output = SimpleNamespace(
|
||||
scheduled_new_reqs=[new_req],
|
||||
scheduled_cached_reqs=cached_reqs,
|
||||
)
|
||||
|
||||
coord.postprocess_scheduler_output(scheduler_output)
|
||||
|
||||
self.assertEqual(coord.requests_with_ready_chunks, set())
|
||||
|
||||
|
||||
class TestWaitingForInputTransition(unittest.TestCase):
|
||||
"""Test B8: process_pending_full_payload_inputs transitions WAITING_FOR_INPUT."""
|
||||
"""Test process_pending_full_payload_inputs transitions WAITING_FOR_INPUT."""
|
||||
|
||||
def test_transition_on_recv(self):
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
|
||||
req = _make_request("r1", status=RequestStatus.WAITING_FOR_INPUT)
|
||||
waiting = MockQueue([req])
|
||||
running: list = []
|
||||
|
||||
coord.process_pending_full_payload_inputs(
|
||||
waiting,
|
||||
running,
|
||||
stage_recv_req_ids={"r1"},
|
||||
)
|
||||
|
||||
self.assertEqual(req.status, RequestStatus.WAITING)
|
||||
|
||||
def test_stays_waiting_for_input_if_not_received(self):
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
|
||||
req = _make_request("r1", status=RequestStatus.WAITING_FOR_INPUT)
|
||||
waiting = MockQueue([req])
|
||||
running: list = []
|
||||
|
||||
coord.process_pending_full_payload_inputs(
|
||||
waiting,
|
||||
running,
|
||||
stage_recv_req_ids=set(),
|
||||
)
|
||||
|
||||
@@ -424,49 +311,39 @@ class TestWaitingForInputTransition(unittest.TestCase):
|
||||
self.assertEqual(len(coord._waiting_for_input), 1)
|
||||
|
||||
def test_stage_0_is_noop(self):
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=0)
|
||||
coord = OmniSchedulingCoordinator(stage_id=0)
|
||||
|
||||
req = _make_request("r1", status=RequestStatus.WAITING_FOR_INPUT)
|
||||
waiting = MockQueue([req])
|
||||
running: list = []
|
||||
|
||||
coord.process_pending_full_payload_inputs(
|
||||
waiting,
|
||||
running,
|
||||
stage_recv_req_ids={"r1"},
|
||||
)
|
||||
self.assertEqual(req.status, RequestStatus.WAITING_FOR_INPUT)
|
||||
|
||||
def test_restore_queues_includes_waiting_for_input(self):
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
|
||||
r1 = _make_request("r1")
|
||||
coord._waiting_for_input.append(r1)
|
||||
|
||||
waiting = MockQueue()
|
||||
running: list = []
|
||||
|
||||
coord.restore_queues(waiting, running)
|
||||
coord.restore_queues(waiting)
|
||||
|
||||
self.assertIn(r1, waiting)
|
||||
self.assertEqual(len(coord._waiting_for_input), 0)
|
||||
|
||||
def test_full_payload_mode_auto_transitions_waiting_to_waiting_for_input(self):
|
||||
"""In full_payload_mode (async_chunk=False), fresh WAITING requests on
|
||||
non-Stage-0 should be transitioned to WAITING_FOR_INPUT."""
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=10,
|
||||
stage_id=1,
|
||||
async_chunk=False,
|
||||
)
|
||||
"""Fresh downstream WAITING requests enter WAITING_FOR_INPUT."""
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
|
||||
req = _make_request("r1", status=RequestStatus.WAITING)
|
||||
waiting = MockQueue([req])
|
||||
running: list = []
|
||||
|
||||
coord.process_pending_full_payload_inputs(
|
||||
waiting,
|
||||
running,
|
||||
stage_recv_req_ids=set(),
|
||||
)
|
||||
|
||||
@@ -474,37 +351,14 @@ class TestWaitingForInputTransition(unittest.TestCase):
|
||||
self.assertEqual(len(coord._waiting_for_input), 1)
|
||||
self.assertEqual(len(coord.pending_input_registrations), 1)
|
||||
|
||||
def test_async_chunk_mode_does_not_auto_transition(self):
|
||||
"""In async_chunk mode, fresh WAITING requests should NOT be
|
||||
transitioned to WAITING_FOR_INPUT."""
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=10,
|
||||
stage_id=1,
|
||||
async_chunk=True,
|
||||
)
|
||||
|
||||
req = _make_request("r1", status=RequestStatus.WAITING)
|
||||
waiting = MockQueue([req])
|
||||
running: list = []
|
||||
|
||||
coord.process_pending_full_payload_inputs(
|
||||
waiting,
|
||||
running,
|
||||
stage_recv_req_ids=set(),
|
||||
)
|
||||
|
||||
self.assertEqual(req.status, RequestStatus.WAITING)
|
||||
|
||||
def test_pending_input_registrations(self):
|
||||
coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
|
||||
req = _make_request("r1", status=RequestStatus.WAITING_FOR_INPUT)
|
||||
waiting = MockQueue([req])
|
||||
running: list = []
|
||||
|
||||
coord.process_pending_full_payload_inputs(
|
||||
waiting,
|
||||
running,
|
||||
stage_recv_req_ids=set(),
|
||||
)
|
||||
|
||||
@@ -512,18 +366,13 @@ class TestWaitingForInputTransition(unittest.TestCase):
|
||||
self.assertEqual(coord.pending_input_registrations[0].request_id, "r1")
|
||||
|
||||
def test_idle_cycles_retain_received_marker_before_request_appears(self):
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=10,
|
||||
stage_id=1,
|
||||
async_chunk=False,
|
||||
)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
coord._full_payload_input_received.add("late")
|
||||
coord.finished_requests.add("late")
|
||||
|
||||
waiting = MockQueue()
|
||||
running: list = []
|
||||
|
||||
coord.process_pending_full_payload_inputs(waiting, running, stage_recv_req_ids=set())
|
||||
coord.process_pending_full_payload_inputs(waiting, stage_recv_req_ids=set())
|
||||
|
||||
self.assertIn("late", coord._full_payload_input_received)
|
||||
self.assertIn("late", coord.finished_requests)
|
||||
@@ -531,7 +380,7 @@ class TestWaitingForInputTransition(unittest.TestCase):
|
||||
late_req = _make_request("late", status=RequestStatus.WAITING)
|
||||
waiting.add_request(late_req)
|
||||
|
||||
coord.process_pending_full_payload_inputs(waiting, running, stage_recv_req_ids=set())
|
||||
coord.process_pending_full_payload_inputs(waiting, stage_recv_req_ids=set())
|
||||
|
||||
self.assertEqual(late_req.status, RequestStatus.WAITING)
|
||||
self.assertEqual(coord.pending_input_registrations, [])
|
||||
@@ -542,65 +391,18 @@ class TestWaitingForInputTransition(unittest.TestCase):
|
||||
class TestTimeoutDetection(unittest.TestCase):
|
||||
"""Regression tests for orphaned pending-recv timeout detection.
|
||||
|
||||
Covers the full lifecycle:
|
||||
1. Request enters WAITING_FOR_CHUNK from either waiting or running queue
|
||||
2. restore_queues() moves it back to the scheduler queue
|
||||
3. Timeout fires via collect_timed_out_request_ids()
|
||||
4. Scheduler removes from both queues and calls _free_request()
|
||||
Covers WAITING_FOR_INPUT lifecycle timeouts. Chunk waiting timeouts are
|
||||
covered by OmniChunkTransferAdapter tests.
|
||||
"""
|
||||
|
||||
def test_waiting_since_recorded_on_chunk_wait(self):
|
||||
"""_waiting_since is set when a request enters WAITING_FOR_CHUNK."""
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=10,
|
||||
stage_id=1,
|
||||
async_chunk=True,
|
||||
)
|
||||
req = _make_request("r1", status=RequestStatus.WAITING)
|
||||
waiting = MockQueue([req])
|
||||
|
||||
coord.process_pending_chunks(
|
||||
waiting,
|
||||
[],
|
||||
chunk_ready_req_ids=set(),
|
||||
chunk_finished_req_ids=set(),
|
||||
)
|
||||
|
||||
self.assertIn("r1", coord._waiting_since)
|
||||
self.assertEqual(req.status, RequestStatus.WAITING_FOR_CHUNK)
|
||||
|
||||
def test_waiting_since_cleared_on_chunk_arrival(self):
|
||||
"""_waiting_since is cleared when a chunk arrives."""
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=10,
|
||||
stage_id=1,
|
||||
async_chunk=True,
|
||||
)
|
||||
req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
|
||||
waiting = MockQueue([req])
|
||||
|
||||
coord.process_pending_chunks(
|
||||
waiting,
|
||||
[],
|
||||
chunk_ready_req_ids={"r1"},
|
||||
chunk_finished_req_ids=set(),
|
||||
)
|
||||
|
||||
self.assertNotIn("r1", coord._waiting_since)
|
||||
|
||||
def test_waiting_since_recorded_on_input_wait(self):
|
||||
"""_waiting_since is set when a request enters WAITING_FOR_INPUT."""
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=10,
|
||||
stage_id=1,
|
||||
async_chunk=False,
|
||||
)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
req = _make_request("r1", status=RequestStatus.WAITING)
|
||||
waiting = MockQueue([req])
|
||||
|
||||
coord.process_pending_full_payload_inputs(
|
||||
waiting,
|
||||
[],
|
||||
stage_recv_req_ids=set(),
|
||||
)
|
||||
|
||||
@@ -608,11 +410,7 @@ class TestTimeoutDetection(unittest.TestCase):
|
||||
|
||||
def test_waiting_since_cleared_on_input_arrival(self):
|
||||
"""_waiting_since is cleared when input data arrives."""
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=10,
|
||||
stage_id=1,
|
||||
async_chunk=False,
|
||||
)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
req = _make_request("r1", status=RequestStatus.WAITING_FOR_INPUT)
|
||||
coord._waiting_for_input.append(req)
|
||||
coord._waiting_since["r1"] = 0.0
|
||||
@@ -620,7 +418,6 @@ class TestTimeoutDetection(unittest.TestCase):
|
||||
waiting = MockQueue()
|
||||
coord.process_pending_full_payload_inputs(
|
||||
waiting,
|
||||
[],
|
||||
stage_recv_req_ids={"r1"},
|
||||
)
|
||||
|
||||
@@ -629,10 +426,7 @@ class TestTimeoutDetection(unittest.TestCase):
|
||||
|
||||
def test_collect_timed_out_request_ids_no_timeout(self):
|
||||
"""No IDs returned when nothing has timed out."""
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=10,
|
||||
stage_id=1,
|
||||
)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
import time
|
||||
|
||||
coord._waiting_since["r1"] = time.monotonic()
|
||||
@@ -642,10 +436,7 @@ class TestTimeoutDetection(unittest.TestCase):
|
||||
|
||||
def test_collect_timed_out_request_ids_expired(self):
|
||||
"""Timed-out IDs are returned and _waiting_since is cleared."""
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=10,
|
||||
stage_id=1,
|
||||
)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
coord._waiting_since["r1"] = 0.0 # epoch → definitely expired
|
||||
coord._waiting_since["r2"] = 0.0
|
||||
|
||||
@@ -662,29 +453,19 @@ class TestTimeoutDetection(unittest.TestCase):
|
||||
|
||||
def test_collect_removes_from_coordinator_queues(self):
|
||||
"""Timed-out requests are defensively removed from internal queues."""
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=10,
|
||||
stage_id=1,
|
||||
)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
r1 = _make_request("r1")
|
||||
r2 = _make_request("r2")
|
||||
coord._waiting_for_chunk_waiting.append(r1)
|
||||
coord._waiting_for_input.append(r2)
|
||||
coord._waiting_for_input.append(r1)
|
||||
coord._waiting_since["r1"] = 0.0
|
||||
coord._waiting_since["r2"] = 0.0
|
||||
|
||||
result = coord.collect_timed_out_request_ids(timeout_s=1.0)
|
||||
|
||||
self.assertEqual(result, {"r1", "r2"})
|
||||
self.assertEqual(len(coord._waiting_for_chunk_waiting), 0)
|
||||
self.assertEqual(result, {"r1"})
|
||||
self.assertEqual(len(coord._waiting_for_input), 0)
|
||||
|
||||
def test_free_finished_request_clears_waiting_since(self):
|
||||
"""free_finished_request clears coordinator lifecycle markers."""
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=10,
|
||||
stage_id=1,
|
||||
)
|
||||
coord = OmniSchedulingCoordinator(stage_id=1)
|
||||
coord._waiting_since["r1"] = 0.0
|
||||
coord._full_payload_input_received.add("r1")
|
||||
coord.finished_requests.add("r1")
|
||||
@@ -693,174 +474,6 @@ class TestTimeoutDetection(unittest.TestCase):
|
||||
self.assertNotIn("r1", coord._full_payload_input_received)
|
||||
self.assertNotIn("r1", coord.finished_requests)
|
||||
|
||||
def test_timeout_from_running_queue_full_lifecycle(self):
|
||||
"""End-to-end: request from running → WAITING_FOR_CHUNK → restore →
|
||||
timeout → removed from running list.
|
||||
|
||||
This is the critical regression case: WAITING_FOR_CHUNK requests
|
||||
that originated from self.running are placed back into self.running
|
||||
by restore_queues(), but their status remains WAITING_FOR_CHUNK.
|
||||
The scheduler must remove from BOTH queues unconditionally.
|
||||
"""
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=10,
|
||||
stage_id=1,
|
||||
async_chunk=True,
|
||||
)
|
||||
|
||||
# 1) Request starts in running queue with WAITING status
|
||||
req = _make_request("r1", status=RequestStatus.WAITING)
|
||||
running = [req]
|
||||
waiting = MockQueue()
|
||||
|
||||
# 2) process_pending_chunks: moves to WAITING_FOR_CHUNK
|
||||
coord.process_pending_chunks(
|
||||
waiting,
|
||||
running,
|
||||
chunk_ready_req_ids=set(),
|
||||
chunk_finished_req_ids=set(),
|
||||
)
|
||||
self.assertEqual(req.status, RequestStatus.WAITING_FOR_CHUNK)
|
||||
self.assertIn("r1", coord._waiting_since)
|
||||
self.assertEqual(len(coord._waiting_for_chunk_running), 1)
|
||||
|
||||
# 3) restore_queues: back to running (status stays WAITING_FOR_CHUNK)
|
||||
coord.restore_queues(waiting, running)
|
||||
self.assertIn(req, running)
|
||||
self.assertEqual(len(coord._waiting_for_chunk_running), 0)
|
||||
self.assertEqual(req.status, RequestStatus.WAITING_FOR_CHUNK)
|
||||
|
||||
# 4) Force timeout by setting _waiting_since to epoch
|
||||
coord._waiting_since["r1"] = 0.0
|
||||
|
||||
timed_out_ids = coord.collect_timed_out_request_ids(timeout_s=1.0)
|
||||
self.assertEqual(timed_out_ids, {"r1"})
|
||||
|
||||
# 5) Scheduler removes from both queues (simulating the scheduler path)
|
||||
timed_out_id_set = {id(req)}
|
||||
running = [r for r in running if id(r) not in timed_out_id_set]
|
||||
waiting.remove_requests([req])
|
||||
|
||||
self.assertNotIn(req, running)
|
||||
self.assertEqual(len(waiting), 0)
|
||||
|
||||
def test_timeout_from_waiting_queue_full_lifecycle(self):
|
||||
"""End-to-end: request from waiting → WAITING_FOR_CHUNK → restore →
|
||||
timeout → removed from waiting queue."""
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=10,
|
||||
stage_id=1,
|
||||
async_chunk=True,
|
||||
)
|
||||
|
||||
req = _make_request("r1", status=RequestStatus.WAITING)
|
||||
waiting = MockQueue([req])
|
||||
running: list = []
|
||||
|
||||
coord.process_pending_chunks(
|
||||
waiting,
|
||||
running,
|
||||
chunk_ready_req_ids=set(),
|
||||
chunk_finished_req_ids=set(),
|
||||
)
|
||||
self.assertEqual(len(coord._waiting_for_chunk_waiting), 1)
|
||||
|
||||
coord.restore_queues(waiting, running)
|
||||
self.assertIn(req, waiting)
|
||||
|
||||
coord._waiting_since["r1"] = 0.0
|
||||
timed_out_ids = coord.collect_timed_out_request_ids(timeout_s=1.0)
|
||||
self.assertEqual(timed_out_ids, {"r1"})
|
||||
|
||||
waiting.remove_requests([req])
|
||||
self.assertEqual(len(waiting), 0)
|
||||
|
||||
|
||||
class TestOverflowPreemption(unittest.TestCase):
|
||||
"""Tests for P1-1: overflow requests must get WAITING status.
|
||||
|
||||
Overflow happens when multiple WAITING_FOR_CHUNK requests in
|
||||
``_waiting_for_chunk_running`` receive their chunk in the same cycle.
|
||||
``_process_chunk_queue`` restores them to RUNNING (``continue``
|
||||
path) while RUNNING requests without chunks are moved out. If the
|
||||
net result exceeds ``scheduler_max_num_seqs``, the tail is pushed
|
||||
to ``waiting_queue`` and must have status == WAITING.
|
||||
"""
|
||||
|
||||
def test_overflow_sets_waiting_status(self):
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=1,
|
||||
stage_id=1,
|
||||
async_chunk=True,
|
||||
)
|
||||
|
||||
# r1 is currently RUNNING in the queue.
|
||||
# r2, r3 were previously moved to _waiting_for_chunk_running.
|
||||
r1 = _make_request("r1", status=RequestStatus.RUNNING)
|
||||
r2 = _make_request("r2", status=RequestStatus.WAITING_FOR_CHUNK)
|
||||
r3 = _make_request("r3", status=RequestStatus.WAITING_FOR_CHUNK)
|
||||
|
||||
running = [r1]
|
||||
waiting = MockQueue([])
|
||||
coord._waiting_for_chunk_running.extend([r2, r3])
|
||||
|
||||
# restore_queues puts r2, r3 back into running
|
||||
coord.restore_queues(waiting, running)
|
||||
self.assertEqual(len(running), 3)
|
||||
|
||||
# Now process_pending_chunks with r2, r3 chunks ready:
|
||||
# _process_chunk_queue will:
|
||||
# r1 (RUNNING) → no chunk → move to _waiting_for_chunk_running
|
||||
# r2 (WAITING_FOR_CHUNK, chunk ready) → set RUNNING, stay in running
|
||||
# r3 (WAITING_FOR_CHUNK, chunk ready) → set RUNNING, stay in running
|
||||
# running = [r2, r3], len=2 > max=1 → overflow
|
||||
coord.process_pending_chunks(
|
||||
waiting,
|
||||
running,
|
||||
chunk_ready_req_ids={"r2", "r3"},
|
||||
chunk_finished_req_ids=set(),
|
||||
)
|
||||
|
||||
self.assertEqual(len(running), 1)
|
||||
self.assertEqual(len(waiting), 1)
|
||||
overflow_req = list(waiting)[0]
|
||||
self.assertEqual(
|
||||
overflow_req.status,
|
||||
RequestStatus.WAITING,
|
||||
f"Overflowed request should have WAITING status, got {overflow_req.status}",
|
||||
)
|
||||
|
||||
def test_overflow_does_not_strand_request(self):
|
||||
"""Without the fix, the overflowed request would keep its
|
||||
RUNNING status in the waiting queue and never be re-scheduled."""
|
||||
coord = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=1,
|
||||
stage_id=1,
|
||||
async_chunk=True,
|
||||
)
|
||||
|
||||
r1 = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
|
||||
r2 = _make_request("r2", status=RequestStatus.WAITING_FOR_CHUNK)
|
||||
coord._waiting_for_chunk_running.extend([r1, r2])
|
||||
|
||||
running: list = []
|
||||
waiting = MockQueue([])
|
||||
|
||||
coord.restore_queues(waiting, running)
|
||||
self.assertEqual(len(running), 2)
|
||||
|
||||
coord.process_pending_chunks(
|
||||
waiting,
|
||||
running,
|
||||
chunk_ready_req_ids={"r1", "r2"},
|
||||
chunk_finished_req_ids=set(),
|
||||
)
|
||||
|
||||
self.assertEqual(len(running), 1)
|
||||
self.assertEqual(len(waiting), 1)
|
||||
for req in waiting:
|
||||
self.assertNotEqual(req.status, RequestStatus.RUNNING, "Overflowed request must not keep RUNNING status")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,7 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import defaultdict
|
||||
from dataclasses import asdict, dataclass
|
||||
from time import time
|
||||
from typing import Any
|
||||
|
||||
@@ -36,17 +35,6 @@ from vllm_omni.outputs import OmniConnectorOutput
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class KVCacheTransferData:
|
||||
request_id: str
|
||||
layer_blocks: dict[str, Any]
|
||||
block_ids: list[int]
|
||||
metadata: dict[str, Any]
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return asdict(self)
|
||||
|
||||
|
||||
class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
"""Synchronous AutoRegressive scheduler for vLLM-Omni. This class is also
|
||||
used as a base class for the OmniARAsyncScheduler and holds most of the
|
||||
@@ -87,9 +75,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
self.input_coordinator: OmniSchedulingCoordinator | None = None
|
||||
if uses_full_payload_input_coordinator(model_config):
|
||||
self.input_coordinator = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=self.vllm_config.scheduler_config.max_num_seqs,
|
||||
stage_id=getattr(model_config, "stage_id", 0),
|
||||
async_chunk=False,
|
||||
)
|
||||
self._latest_omni_connector_output: OmniConnectorOutput | None = None
|
||||
# Snapshot prompt length for each streaming input update
|
||||
@@ -175,6 +161,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
if criteria_type == "prefill_finished":
|
||||
if confirmed_computed >= request.num_prompt_tokens:
|
||||
self.transfer_triggered_requests.add(request.request_id)
|
||||
|
||||
self._mark_request_for_kv_transfer(request.request_id, confirmed_computed)
|
||||
actually_queued = request.request_id in self.requests_needing_kv_transfer
|
||||
|
||||
@@ -220,7 +207,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
queue.remove(req)
|
||||
self._consume_pending_connector_output(model_mode="ar")
|
||||
self._process_pending_input_timeouts()
|
||||
|
||||
if self.chunk_transfer_adapter:
|
||||
self.chunk_transfer_adapter.process_pending_chunks(
|
||||
self.waiting, self.running, scheduler_requests=self.requests
|
||||
@@ -247,7 +233,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
scheduler_requests=self.requests,
|
||||
)
|
||||
if self.input_coordinator:
|
||||
self.input_coordinator.restore_queues(self.waiting, self.running)
|
||||
self.input_coordinator.restore_queues(self.waiting)
|
||||
try:
|
||||
# Late import to avoid circulars in some launch modes
|
||||
from .output import OmniNewRequestData
|
||||
@@ -655,7 +641,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
Tuple of (req_id, client_index) for requests that were aborted. Will not
|
||||
include any that were already finished.
|
||||
"""
|
||||
|
||||
# TODO(yrr): chunk transfer adapter & input_coordinator unified to one
|
||||
if self.chunk_transfer_adapter:
|
||||
self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
|
||||
|
||||
@@ -823,11 +809,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
finally:
|
||||
self._free_input_coordinator_request(request_id)
|
||||
|
||||
def _free_blocks(self, request: Request):
|
||||
# Helper to match base class structure if not directly available
|
||||
# VLLMScheduler has _free_blocks
|
||||
super()._free_blocks(request)
|
||||
|
||||
def _mark_request_for_kv_transfer(self, req_id: str, seq_len: int) -> None:
|
||||
"""Mark a request as needing KV cache transfer when it finishes."""
|
||||
# Avoid duplicate marking (if already pending in queue)
|
||||
|
||||
@@ -50,14 +50,12 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
self.input_coordinator: OmniSchedulingCoordinator | None = None
|
||||
if uses_full_payload_input_coordinator(model_config):
|
||||
self.input_coordinator = OmniSchedulingCoordinator(
|
||||
scheduler_max_num_seqs=self.vllm_config.scheduler_config.max_num_seqs,
|
||||
stage_id=getattr(model_config, "stage_id", 0),
|
||||
async_chunk=False,
|
||||
)
|
||||
self._latest_omni_connector_output: OmniConnectorOutput | None = None
|
||||
|
||||
def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
|
||||
"""Diffusion fast path:
|
||||
"""One-shot generation fast path:
|
||||
- Feed all input tokens of the request at once
|
||||
(if 0, allocate 1 placeholder token).
|
||||
- If the token budget cannot be satisfied at once, fall back to the
|
||||
@@ -81,7 +79,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
cached_prompt_token_ids: dict[str, list[int]] = {}
|
||||
cached_additional_information: dict[str, dict | None] = {}
|
||||
|
||||
# Temporary queue: preserve waiting order, do not disturb non-diffusion requests
|
||||
# Temporary queue: preserve waiting order while requests await input.
|
||||
skipped_waiting_requests = create_request_queue(self.policy)
|
||||
req_index = 0
|
||||
self._consume_pending_connector_output(model_mode="generation")
|
||||
@@ -144,8 +142,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
if already_finished_reqs:
|
||||
self.running = remove_all(self.running, already_finished_reqs)
|
||||
|
||||
# Fast path selection and scheduling (treat all as diffusion requests,
|
||||
# independent of pooling_params)
|
||||
# Fast path selection and scheduling for one-shot generation requests,
|
||||
# independent of pooling_params.
|
||||
while (
|
||||
self.waiting
|
||||
and token_budget > 0
|
||||
@@ -172,9 +170,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
skipped_waiting_requests.prepend_request(request)
|
||||
continue
|
||||
|
||||
# Uniformly treat as diffusion. A feature flag can be added later
|
||||
# via config or request tag.
|
||||
|
||||
# Allocate all input tokens for the request in one shot
|
||||
# (allocate 1 placeholder if zero)
|
||||
required_tokens = max(len(request.prompt_token_ids), 1)
|
||||
@@ -219,7 +214,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
else:
|
||||
res = super().schedule(throttle_prefills)
|
||||
if self.input_coordinator:
|
||||
self.input_coordinator.restore_queues(self.waiting, self.running)
|
||||
self.input_coordinator.restore_queues(self.waiting)
|
||||
return self._wrap_omni_scheduler_output(res)
|
||||
|
||||
# Compute common prefix blocks (aligned with v1)
|
||||
@@ -346,7 +341,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
scheduler_requests=self.requests,
|
||||
)
|
||||
if self.input_coordinator:
|
||||
self.input_coordinator.restore_queues(self.waiting, self.running)
|
||||
self.input_coordinator.restore_queues(self.waiting)
|
||||
|
||||
return self._wrap_omni_scheduler_output(scheduler_output)
|
||||
|
||||
@@ -397,24 +392,12 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
finally:
|
||||
self._free_input_coordinator_request(request.request_id)
|
||||
|
||||
"""
|
||||
Scheduler for the diffusion model.
|
||||
This scheduler is modified to stop the request immediately for the diffusion model.
|
||||
This is because the diffusion model can generate the final image/audio in one step.
|
||||
Note: This is just a minimal modification to the original scheduler,
|
||||
and there should be some further efforts to optimize the scheduler.
|
||||
The original scheduler is still used for the AR model.
|
||||
"""
|
||||
|
||||
def update_from_output(
|
||||
self,
|
||||
scheduler_output: SchedulerOutput,
|
||||
model_runner_output: OmniModelRunnerOutput,
|
||||
) -> dict[int, EngineCoreOutputs]:
|
||||
"""Update the scheduler state based on the model runner output.
|
||||
|
||||
This method is modified to stop the request immediately for the diffusion model.
|
||||
"""
|
||||
"""Update scheduler state and finish completed one-shot work."""
|
||||
sampled_token_ids = model_runner_output.sampled_token_ids
|
||||
logprobs = model_runner_output.logprobs
|
||||
prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
|
||||
@@ -496,7 +479,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
finish_reason = None
|
||||
routed_experts = None
|
||||
|
||||
# Diffusion request: completes in one step; mark finished and free resources
|
||||
# One-shot generation request: finish after its current input unit
|
||||
# has been fully processed.
|
||||
if (
|
||||
request.status == RequestStatus.FINISHED_STOPPED
|
||||
or (self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens)
|
||||
@@ -509,7 +493,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
request.status = RequestStatus.FINISHED_STOPPED
|
||||
# Optional: set a stop_reason for front-end clarity
|
||||
# (does not affect protocol)
|
||||
request.stop_reason = request.stop_reason # or "generation_done"
|
||||
stopped = True
|
||||
|
||||
if stopped:
|
||||
|
||||
@@ -68,7 +68,6 @@ class OmniSchedulerMixin:
|
||||
)
|
||||
input_coordinator.process_pending_full_payload_inputs(
|
||||
self.waiting,
|
||||
self.running,
|
||||
connector_output.stage_recv_req_ids if connector_output else set(),
|
||||
)
|
||||
|
||||
|
||||
@@ -1,13 +1,11 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Scheduling-side coordination for chunk and full_payload input waiting.
|
||||
"""Scheduling-side coordination for full_payload input waiting.
|
||||
|
||||
Manages WAITING_FOR_CHUNK and WAITING_FOR_INPUT state transitions
|
||||
based on readiness signals from OmniConnectorOutput, without ever
|
||||
calling connector.put()/get().
|
||||
Manages WAITING_FOR_INPUT state transitions based on readiness signals
|
||||
from OmniConnectorOutput, without ever calling connector.put()/get().
|
||||
|
||||
This replaces the scheduling half of OmniChunkTransferAdapter; the
|
||||
transport half lives in OmniConnectorModelRunnerMixin.
|
||||
Chunk waiting (WAITING_FOR_CHUNK) lives on OmniChunkTransferAdapter.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -64,7 +62,7 @@ _FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
|
||||
|
||||
|
||||
def uses_full_payload_input_coordinator(model_config: Any) -> bool:
|
||||
"""Returns True iff this stage parks pending requests in
|
||||
"""Returns True if this stage parks pending requests in
|
||||
WAITING_FOR_INPUT awaiting a full_payload delivery on the worker connector.
|
||||
|
||||
Gated by (model_arch, model_stage) — see _FULL_PAYLOAD_INPUT_STAGES for the
|
||||
@@ -83,31 +81,20 @@ def uses_full_payload_input_coordinator(model_config: Any) -> bool:
|
||||
|
||||
|
||||
class OmniSchedulingCoordinator:
|
||||
"""Pure-scheduling coordinator for chunk and full_payload input waiting.
|
||||
"""Pure-scheduling coordinator for full_payload input waiting.
|
||||
|
||||
The Scheduler owns an instance of this class. It consumes readiness
|
||||
signals produced by the Model Runner's ``OmniConnectorModelRunnerMixin``
|
||||
(via ``OmniConnectorOutput``) and manages ``WAITING_FOR_CHUNK`` and
|
||||
``WAITING_FOR_INPUT`` state transitions accordingly.
|
||||
(via ``OmniConnectorOutput``) and manages ``WAITING_FOR_INPUT`` state
|
||||
transitions accordingly.
|
||||
"""
|
||||
|
||||
def __init__(self, scheduler_max_num_seqs: int, stage_id: int = 0, async_chunk: bool = False):
|
||||
def __init__(self, stage_id: int = 0):
|
||||
self._stage_id = stage_id
|
||||
self._scheduler_max_num_seqs = scheduler_max_num_seqs
|
||||
self._async_chunk = async_chunk
|
||||
|
||||
self.finished_requests: set[str] = set()
|
||||
self.requests_with_ready_chunks: set[str] = set()
|
||||
self._full_payload_input_received: set[str] = set()
|
||||
|
||||
self._waiting_for_chunk_waiting: deque[Any] = deque()
|
||||
self._waiting_for_chunk_running: deque[Any] = deque()
|
||||
|
||||
# Request IDs that were newly registered for chunk recv this cycle.
|
||||
# The engine/Model Runner should call register_chunk_recv() for these
|
||||
# so the bg thread starts polling.
|
||||
self.pending_chunk_registrations: list[Any] = []
|
||||
|
||||
# Requests waiting for full_payload stage input (WAITING_FOR_INPUT).
|
||||
self._waiting_for_input: deque[Any] = deque()
|
||||
# Per-cycle list of minimal handles to ship to the model runner so it
|
||||
@@ -117,67 +104,22 @@ class OmniSchedulingCoordinator:
|
||||
self.pending_input_registrations: list[OmniChunkRecvHandle] = []
|
||||
|
||||
# Monotonic timestamp recording when each request first entered
|
||||
# WAITING_FOR_CHUNK or WAITING_FOR_INPUT. Used by
|
||||
# collect_timed_out_request_ids() to detect orphaned waits.
|
||||
# WAITING_FOR_INPUT. Used by collect_timed_out_request_ids() to
|
||||
# detect orphaned waits.
|
||||
self._waiting_since: dict[str, float] = {}
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Core scheduling methods
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def process_pending_chunks(
|
||||
self,
|
||||
waiting_queue: Any,
|
||||
running_queue: list[Request],
|
||||
chunk_ready_req_ids: set[str],
|
||||
chunk_finished_req_ids: set[str],
|
||||
) -> None:
|
||||
"""Transition requests whose chunks have arrived.
|
||||
|
||||
Args:
|
||||
waiting_queue: Scheduler's waiting request queue.
|
||||
running_queue: Scheduler's running request list.
|
||||
chunk_ready_req_ids: IDs with a newly arrived chunk this cycle.
|
||||
chunk_finished_req_ids: IDs whose final chunk has arrived.
|
||||
"""
|
||||
if self._stage_id == 0 or not self._async_chunk:
|
||||
return
|
||||
|
||||
terminal_ready_req_ids = chunk_ready_req_ids.intersection(chunk_finished_req_ids)
|
||||
self.finished_requests.update(chunk_finished_req_ids - terminal_ready_req_ids)
|
||||
self.pending_chunk_registrations = []
|
||||
|
||||
self._process_chunk_queue(
|
||||
waiting_queue,
|
||||
self._waiting_for_chunk_waiting,
|
||||
RequestStatus.WAITING,
|
||||
chunk_ready_req_ids,
|
||||
)
|
||||
self._process_chunk_queue(
|
||||
running_queue,
|
||||
self._waiting_for_chunk_running,
|
||||
RequestStatus.RUNNING,
|
||||
chunk_ready_req_ids,
|
||||
)
|
||||
self.finished_requests.update(terminal_ready_req_ids)
|
||||
|
||||
while len(running_queue) > self._scheduler_max_num_seqs:
|
||||
request = running_queue.pop()
|
||||
# Must reset status to WAITING so the scheduler treats it as
|
||||
# schedulable work. KV blocks are NOT freed here (unlike a
|
||||
# real preemption), so PREEMPTED would be incorrect.
|
||||
request.status = RequestStatus.WAITING
|
||||
waiting_queue.prepend_requests([request])
|
||||
|
||||
def process_pending_full_payload_inputs(
|
||||
self,
|
||||
waiting_queue: Any,
|
||||
running_queue: list[Request],
|
||||
stage_recv_req_ids: set[str],
|
||||
) -> None:
|
||||
"""Manage WAITING_FOR_INPUT lifecycle for full_payload_mode.
|
||||
|
||||
For non-Stage-0 stages in full_payload_mode (``async_chunk=False``):
|
||||
For non-Stage-0 stages in full_payload mode:
|
||||
1. Fresh WAITING requests are transitioned to WAITING_FOR_INPUT
|
||||
and registered for bg-thread polling.
|
||||
2. WAITING_FOR_INPUT requests whose data has arrived (in
|
||||
@@ -187,7 +129,7 @@ class OmniSchedulingCoordinator:
|
||||
return
|
||||
|
||||
self._full_payload_input_received.update(stage_recv_req_ids)
|
||||
if not self._async_chunk and stage_recv_req_ids:
|
||||
if stage_recv_req_ids:
|
||||
self.finished_requests.update(stage_recv_req_ids)
|
||||
logger.debug(
|
||||
"[Coordinator stage-%s] full_payload recv -> finished_requests: %s",
|
||||
@@ -206,19 +148,29 @@ class OmniSchedulingCoordinator:
|
||||
remaining.append(request)
|
||||
self._waiting_for_input = remaining
|
||||
|
||||
if not self._async_chunk:
|
||||
to_remove: list[Any] = []
|
||||
queue_snapshot = list(waiting_queue)
|
||||
for request in queue_snapshot:
|
||||
if request.status == RequestStatus.WAITING:
|
||||
if request.request_id in self._full_payload_input_received:
|
||||
continue
|
||||
if request.request_id in self.requests_with_ready_chunks:
|
||||
continue
|
||||
if request.request_id in self.finished_requests:
|
||||
continue
|
||||
request.status = RequestStatus.WAITING_FOR_INPUT
|
||||
self._waiting_since.setdefault(request.request_id, time.monotonic())
|
||||
to_remove: list[Any] = []
|
||||
queue_snapshot = list(waiting_queue)
|
||||
for request in queue_snapshot:
|
||||
if request.status == RequestStatus.WAITING:
|
||||
if request.request_id in self._full_payload_input_received:
|
||||
continue
|
||||
if request.request_id in self.finished_requests:
|
||||
continue
|
||||
request.status = RequestStatus.WAITING_FOR_INPUT
|
||||
self._waiting_since.setdefault(request.request_id, time.monotonic())
|
||||
to_remove.append(request)
|
||||
self._waiting_for_input.append(request)
|
||||
self.pending_input_registrations.append(
|
||||
OmniChunkRecvHandle(
|
||||
request_id=request.request_id,
|
||||
external_req_id=getattr(request, "external_req_id", None),
|
||||
)
|
||||
)
|
||||
elif request.status == RequestStatus.WAITING_FOR_INPUT:
|
||||
if request.request_id in stage_recv_req_ids:
|
||||
request.status = RequestStatus.WAITING
|
||||
self._waiting_since.pop(request.request_id, None)
|
||||
else:
|
||||
to_remove.append(request)
|
||||
self._waiting_for_input.append(request)
|
||||
self.pending_input_registrations.append(
|
||||
@@ -227,38 +179,15 @@ class OmniSchedulingCoordinator:
|
||||
external_req_id=getattr(request, "external_req_id", None),
|
||||
)
|
||||
)
|
||||
elif request.status == RequestStatus.WAITING_FOR_INPUT:
|
||||
if request.request_id in stage_recv_req_ids:
|
||||
request.status = RequestStatus.WAITING
|
||||
self._waiting_since.pop(request.request_id, None)
|
||||
else:
|
||||
to_remove.append(request)
|
||||
self._waiting_for_input.append(request)
|
||||
self.pending_input_registrations.append(
|
||||
OmniChunkRecvHandle(
|
||||
request_id=request.request_id,
|
||||
external_req_id=getattr(request, "external_req_id", None),
|
||||
)
|
||||
)
|
||||
if to_remove:
|
||||
# Use the bulk-remove helper: one O(N) sweep instead of N
|
||||
# repeated O(N) removes from a list-backed queue.
|
||||
waiting_queue.remove_requests(to_remove)
|
||||
|
||||
def process_pending_full_payload_inputs_legacy(
|
||||
self,
|
||||
waiting_queue: Any,
|
||||
running_queue: list[Request],
|
||||
stage_recv_req_ids: set[str],
|
||||
) -> None:
|
||||
"""Compatibility wrapper for ``process_pending_full_payload_inputs``."""
|
||||
self.process_pending_full_payload_inputs(waiting_queue, running_queue, stage_recv_req_ids)
|
||||
if to_remove:
|
||||
# Use the bulk-remove helper: one O(N) sweep instead of N
|
||||
# repeated O(N) removes from a list-backed queue.
|
||||
waiting_queue.remove_requests(to_remove)
|
||||
|
||||
def free_finished_request(self, request_id: str) -> None:
|
||||
"""Prune internal tracking sets for a freed request to prevent unbounded growth."""
|
||||
self._full_payload_input_received.discard(request_id)
|
||||
self.finished_requests.discard(request_id)
|
||||
self.requests_with_ready_chunks.discard(request_id)
|
||||
self._waiting_since.pop(request_id, None)
|
||||
|
||||
def collect_timed_out_request_ids(
|
||||
@@ -290,22 +219,16 @@ class OmniSchedulingCoordinator:
|
||||
|
||||
# Defensively remove from coordinator internal queues (may already
|
||||
# be empty if restore_queues() has run).
|
||||
for queue_attr in (
|
||||
"_waiting_for_chunk_waiting",
|
||||
"_waiting_for_chunk_running",
|
||||
"_waiting_for_input",
|
||||
):
|
||||
queue = getattr(self, queue_attr)
|
||||
remaining: deque[Any] = deque()
|
||||
for request in queue:
|
||||
if request.request_id not in timed_out_ids:
|
||||
remaining.append(request)
|
||||
setattr(self, queue_attr, remaining)
|
||||
remaining: deque[Any] = deque()
|
||||
for request in self._waiting_for_input:
|
||||
if request.request_id not in timed_out_ids:
|
||||
remaining.append(request)
|
||||
self._waiting_for_input = remaining
|
||||
|
||||
for req_id in timed_out_ids:
|
||||
self._waiting_since.pop(req_id, None)
|
||||
logger.warning(
|
||||
"[Coordinator stage-%s] Request %s timed out waiting for chunk/input (waited > %.0fs)",
|
||||
"[Coordinator stage-%s] Request %s timed out waiting for input (waited > %.0fs)",
|
||||
self._stage_id,
|
||||
req_id,
|
||||
timeout_s,
|
||||
@@ -316,17 +239,8 @@ class OmniSchedulingCoordinator:
|
||||
def restore_queues(
|
||||
self,
|
||||
waiting_queue: Any,
|
||||
running_queue: list[Request],
|
||||
) -> None:
|
||||
"""Return waiting-for-chunk/input requests to scheduling queues."""
|
||||
for request in self._waiting_for_chunk_waiting:
|
||||
waiting_queue.add_request(request)
|
||||
self._waiting_for_chunk_waiting = deque()
|
||||
|
||||
if self._waiting_for_chunk_running:
|
||||
running_queue.extend(self._waiting_for_chunk_running)
|
||||
self._waiting_for_chunk_running = deque()
|
||||
|
||||
"""Return waiting-for-input requests to the waiting queue."""
|
||||
for request in self._waiting_for_input:
|
||||
waiting_queue.add_request(request)
|
||||
self._waiting_for_input = deque()
|
||||
@@ -411,12 +325,6 @@ class OmniSchedulingCoordinator:
|
||||
|
||||
if model_mode != "ar":
|
||||
new_ids = self._flatten_prompt_token_ids(metadata.get("code_predictor_codes"))
|
||||
runtime_seed = None
|
||||
if "left_context_size" in metadata:
|
||||
runtime_seed = {
|
||||
"meta": {"left_context_size": metadata["left_context_size"]},
|
||||
}
|
||||
request._omni_initial_model_buffer = runtime_seed
|
||||
if new_ids:
|
||||
request.prompt_token_ids = new_ids
|
||||
request.num_prompt_tokens = len(new_ids)
|
||||
@@ -424,57 +332,3 @@ class OmniSchedulingCoordinator:
|
||||
request._all_token_ids.extend(new_ids)
|
||||
request._output_token_ids.clear()
|
||||
request.num_computed_tokens = 0
|
||||
|
||||
def postprocess_scheduler_output(
|
||||
self,
|
||||
scheduler_output: Any,
|
||||
requests: dict[str, Request] | None = None,
|
||||
) -> None:
|
||||
"""Clear per-cycle ready state after scheduler output is materialized."""
|
||||
self._clear_chunk_ready(scheduler_output)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Internal helpers
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def _process_chunk_queue(
|
||||
self,
|
||||
queue: Any,
|
||||
waiting_for_chunk_list: deque[Any],
|
||||
target_status: RequestStatus,
|
||||
chunk_ready_req_ids: set[str],
|
||||
) -> None:
|
||||
queue_snapshot = list(queue)
|
||||
for request in queue_snapshot:
|
||||
if request.status != RequestStatus.WAITING_FOR_CHUNK:
|
||||
if request.request_id in self.requests_with_ready_chunks:
|
||||
continue
|
||||
if request.request_id in self.finished_requests:
|
||||
continue
|
||||
if request.status == RequestStatus.WAITING_FOR_INPUT:
|
||||
continue
|
||||
if request.request_id in chunk_ready_req_ids:
|
||||
self.requests_with_ready_chunks.add(request.request_id)
|
||||
continue
|
||||
self.pending_chunk_registrations.append(request)
|
||||
request.status = RequestStatus.WAITING_FOR_CHUNK
|
||||
self._waiting_since.setdefault(request.request_id, time.monotonic())
|
||||
else:
|
||||
if request.request_id in chunk_ready_req_ids:
|
||||
request.status = target_status
|
||||
self.requests_with_ready_chunks.add(request.request_id)
|
||||
self._waiting_since.pop(request.request_id, None)
|
||||
continue
|
||||
queue.remove(request)
|
||||
waiting_for_chunk_list.append(request)
|
||||
|
||||
def _clear_chunk_ready(self, scheduler_output: Any) -> None:
|
||||
if scheduler_output.scheduled_new_reqs:
|
||||
for req_data in scheduler_output.scheduled_new_reqs:
|
||||
self.requests_with_ready_chunks.discard(
|
||||
getattr(req_data, "req_id", None),
|
||||
)
|
||||
|
||||
if scheduler_output.scheduled_cached_reqs:
|
||||
for req_id in scheduler_output.scheduled_cached_reqs.req_ids:
|
||||
self.requests_with_ready_chunks.discard(req_id)
|
||||
@@ -73,7 +73,7 @@ class OmniCachedRequestData(CachedRequestData):
|
||||
|
||||
@dataclass
|
||||
class OmniChunkRecvHandle:
|
||||
"""Minimal identifier carried from scheduler to runner for chunk-recv
|
||||
"""Minimal identifier carried from scheduler to runner for input-receive
|
||||
registration.
|
||||
|
||||
The runner's ``register_chunk_recv`` only consumes ``request_id`` and
|
||||
|
||||
Reference in new issue
Block a user