e9c341afe8
Transferred EAGLE metadata buffers are reusable slot views. The previous clone-based mitigation protected correctness but added copies on the transfer hot path and hid the actual lifetime contract. This change makes the transferred request own the metadata slot while it sits in the decode waiting queue, then releases it immediately after process_prebuilt has consumed top-k and hidden state into the prebuilt batch. Abort paths also release any held decode metadata slot. Constraint: Decode disaggregation metadata buffers are reusable slot views consumed later by process_prebuilt. Rejected: Clone transferred EAGLE tensors at commit time | correct but less efficient and masks the ownership contract. Rejected: Release in process_batch_result_prebuilt | holds slots across forward longer than needed. Confidence: medium Scope-risk: moderate Directive: Do not free successful EAGLE transfer metadata in pop_transferred unless process_prebuilt consumption is also moved earlier. Tested: Remote py_compile for decode.py, scheduler.py, scheduler_output_processor_mixin.py, and test_decode_queue_compaction.py. Tested: Remote focused lifecycle tests passed: 3 passed. Tested: Remote full test_decode_queue_compaction.py passed: 11 passed, 5 warnings. Not-tested: Fresh ETE runtime validation of EAGLE accept-length recovery after the C48 sync.
730 lines
26 KiB
Python
730 lines
26 KiB
Python
"""Unit tests for decode queue one-pass compaction."""
|
|
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from typing import Any, cast
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
|
|
from sglang.srt.disaggregation.base import KVPoll
|
|
from sglang.srt.disaggregation.decode import (
|
|
DecodePreallocQueue,
|
|
DecodeRequest,
|
|
DecodeTransferQueue,
|
|
SchedulerDisaggregationDecodeMixin,
|
|
)
|
|
from sglang.srt.disaggregation.utils import ReqToMetadataIdxAllocator
|
|
from sglang.srt.managers.scheduler_output_processor_mixin import (
|
|
SchedulerOutputProcessorMixin,
|
|
)
|
|
from sglang.srt.managers.schedule_batch import FINISH_ABORT
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cpu_ci(est_time=8, suite="stage-a-test-cpu")
|
|
|
|
|
|
class FakeReq:
|
|
def __init__(self, rid, bootstrap_room):
|
|
self.rid = rid
|
|
self.bootstrap_room = bootstrap_room
|
|
self.bootstrap_host = "host"
|
|
self.return_logprob = False
|
|
self.latencies = []
|
|
self.origin_input_ids = []
|
|
self.output_ids = []
|
|
self.is_retracted = True
|
|
self.load_calls = []
|
|
self.prealloc_done = False
|
|
self.finished_reason = cast(Any, None)
|
|
self.cached_tokens = 0
|
|
self.init_next_round_calls = []
|
|
self.sampling_params = SimpleNamespace(max_new_tokens=0)
|
|
|
|
class FakeTimeStats:
|
|
def __init__(self):
|
|
self.forward_entry_time = None
|
|
|
|
def set_bootstrap_done_time(self):
|
|
return None
|
|
|
|
def set_decode_transfer_queue_entry_time(self):
|
|
return None
|
|
|
|
def set_wait_queue_entry_time(self):
|
|
return None
|
|
|
|
def set_forward_entry_time(self, ts):
|
|
self.forward_entry_time = ts
|
|
|
|
def set_decode_prebuilt_finish_time(self):
|
|
return None
|
|
|
|
def set_quick_finish_time(self):
|
|
return None
|
|
|
|
self.time_stats = FakeTimeStats()
|
|
|
|
def add_latency(self, stage):
|
|
self.latencies.append(stage)
|
|
|
|
def load_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
|
|
self.load_calls.append((req_to_token_pool, token_to_kv_pool_allocator))
|
|
|
|
def init_next_round_input(self, tree_cache):
|
|
self.init_next_round_calls.append(tree_cache)
|
|
|
|
def check_finished(self):
|
|
return None
|
|
|
|
def finished(self):
|
|
return self.finished_reason is not None
|
|
|
|
|
|
class FakeReceiver:
|
|
def __init__(self, should_fail=False):
|
|
self.should_fail = should_fail
|
|
self.init_calls = []
|
|
self.clear_calls = 0
|
|
|
|
def init(self, *args):
|
|
self.init_calls.append(args)
|
|
|
|
def failure_exception(self):
|
|
if self.should_fail:
|
|
raise RuntimeError("boom")
|
|
|
|
def clear(self):
|
|
self.clear_calls += 1
|
|
return None
|
|
|
|
|
|
class FakeBatch:
|
|
def __init__(self):
|
|
self.prepared = False
|
|
self.processed = []
|
|
|
|
def prepare_for_prebuilt(self):
|
|
self.prepared = True
|
|
|
|
def process_prebuilt(self, server_args, future_map):
|
|
self.processed.append((server_args, future_map))
|
|
|
|
|
|
class FakeItem:
|
|
def __init__(self, value):
|
|
self.value = value
|
|
|
|
def item(self):
|
|
return self.value
|
|
|
|
|
|
class FakeTensor:
|
|
def __init__(self, values):
|
|
self.values = values
|
|
|
|
def __getitem__(self, item):
|
|
if isinstance(item, tuple):
|
|
row, col = item
|
|
return FakeTensor(self.values[row][col])
|
|
return FakeTensor(self.values[item])
|
|
|
|
def cpu(self):
|
|
return self
|
|
|
|
def numpy(self):
|
|
return self.values
|
|
|
|
|
|
class FakeAllocator(ReqToMetadataIdxAllocator):
|
|
def __init__(self):
|
|
super().__init__(size=0)
|
|
self.freed = []
|
|
|
|
def free(self, free_index):
|
|
self.freed.append(free_index)
|
|
|
|
|
|
class TestDecodeQueueCompaction(CustomTestCase):
|
|
def test_decode_transfer_queue_compacts_in_one_pass(self):
|
|
streamed = []
|
|
released = []
|
|
committed = []
|
|
allocator = FakeAllocator()
|
|
|
|
queue = DecodeTransferQueue.__new__(DecodeTransferQueue)
|
|
queue.gloo_group = None
|
|
queue.req_to_metadata_buffer_idx_allocator = allocator
|
|
queue.tp_rank = 0
|
|
queue.metadata_buffers = cast(Any, object())
|
|
queue.tree_cache = cast(Any, object())
|
|
queue.spec_algorithm = cast(Any, SimpleNamespace(is_none=lambda: True))
|
|
queue.scheduler = cast(
|
|
Any,
|
|
SimpleNamespace(
|
|
stream_output=lambda reqs, return_logprob: streamed.extend(
|
|
req.rid for req in reqs
|
|
),
|
|
enable_metrics=False,
|
|
token_to_kv_pool_allocator=SimpleNamespace(
|
|
get_kvcache=lambda: SimpleNamespace()
|
|
),
|
|
),
|
|
)
|
|
|
|
keep = DecodeRequest(
|
|
req=cast(Any, FakeReq("keep", 1)),
|
|
kv_receiver=cast(Any, FakeReceiver()),
|
|
metadata_buffer_index=10,
|
|
)
|
|
success = DecodeRequest(
|
|
req=cast(Any, FakeReq("success", 2)),
|
|
kv_receiver=cast(Any, FakeReceiver()),
|
|
metadata_buffer_index=11,
|
|
)
|
|
failed = DecodeRequest(
|
|
req=cast(Any, FakeReq("failed", 3)),
|
|
kv_receiver=cast(Any, FakeReceiver(should_fail=True)),
|
|
metadata_buffer_index=12,
|
|
)
|
|
skipped = DecodeRequest(
|
|
req=cast(Any, FakeReq("skip", 4)),
|
|
kv_receiver=cast(Any, FakeReceiver()),
|
|
metadata_buffer_index=13,
|
|
)
|
|
queue.queue = [keep, success, failed, skipped]
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.disaggregation.decode.poll_and_all_reduce",
|
|
return_value=[
|
|
KVPoll.Transferring,
|
|
KVPoll.Success,
|
|
KVPoll.Failed,
|
|
KVPoll.Success,
|
|
],
|
|
),
|
|
patch(
|
|
"sglang.srt.disaggregation.decode.release_kv_cache",
|
|
lambda req, tree_cache, is_insert=False: released.append(
|
|
(req.rid, is_insert)
|
|
),
|
|
),
|
|
):
|
|
queue._commit_transfer_to_req = lambda decode_req: (
|
|
committed.append(decode_req.req.rid) or True
|
|
)
|
|
transferred = queue.pop_transferred(
|
|
rids_to_check=["keep", "success", "failed"]
|
|
)
|
|
|
|
self.assertEqual([req.rid for req in transferred], ["success"])
|
|
self.assertEqual(committed, ["success"])
|
|
self.assertEqual(streamed, ["failed"])
|
|
self.assertEqual(released, [("failed", False)])
|
|
self.assertEqual(allocator.freed, [11, 12])
|
|
self.assertEqual(queue.queue, [keep, skipped])
|
|
|
|
def test_decode_transfer_queue_keeps_metadata_waiters(self):
|
|
allocator = FakeAllocator()
|
|
|
|
queue = DecodeTransferQueue.__new__(DecodeTransferQueue)
|
|
queue.gloo_group = None
|
|
queue.req_to_metadata_buffer_idx_allocator = allocator
|
|
queue.tp_rank = 0
|
|
queue.metadata_buffers = cast(Any, object())
|
|
queue.tree_cache = cast(Any, object())
|
|
queue.spec_algorithm = cast(Any, SimpleNamespace(is_none=lambda: True))
|
|
queue.scheduler = cast(
|
|
Any,
|
|
SimpleNamespace(
|
|
stream_output=lambda reqs, return_logprob: None,
|
|
enable_metrics=False,
|
|
token_to_kv_pool_allocator=SimpleNamespace(
|
|
get_kvcache=lambda: SimpleNamespace()
|
|
),
|
|
),
|
|
)
|
|
|
|
waiting = DecodeRequest(
|
|
req=cast(Any, FakeReq("waiting", 1)),
|
|
kv_receiver=cast(Any, FakeReceiver()),
|
|
metadata_buffer_index=10,
|
|
)
|
|
queue.queue = [waiting]
|
|
|
|
with patch(
|
|
"sglang.srt.disaggregation.decode.poll_and_all_reduce",
|
|
return_value=[KVPoll.Success],
|
|
):
|
|
queue._commit_transfer_to_req = lambda decode_req: False
|
|
transferred = queue.pop_transferred()
|
|
|
|
self.assertEqual(transferred, [])
|
|
self.assertEqual(allocator.freed, [])
|
|
self.assertEqual(queue.queue, [waiting])
|
|
|
|
def test_commit_transfer_to_req_waits_for_real_metadata(self):
|
|
queue = DecodeTransferQueue.__new__(DecodeTransferQueue)
|
|
queue.metadata_buffers = cast(
|
|
Any,
|
|
SimpleNamespace(
|
|
get_buf=lambda idx: (
|
|
[FakeItem(0)],
|
|
[FakeItem(0)],
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
[FakeItem(0)],
|
|
)
|
|
),
|
|
)
|
|
queue.scheduler = cast(
|
|
Any,
|
|
SimpleNamespace(
|
|
server_args=SimpleNamespace(disaggregation_transfer_backend="mooncake")
|
|
),
|
|
)
|
|
queue.spec_algorithm = cast(Any, SimpleNamespace(is_none=lambda: True))
|
|
|
|
receiver = FakeReceiver()
|
|
decode_req = DecodeRequest(
|
|
req=cast(Any, FakeReq("waiting", 3)),
|
|
kv_receiver=cast(Any, receiver),
|
|
metadata_buffer_index=9,
|
|
)
|
|
|
|
should_remove = queue._commit_transfer_to_req(decode_req)
|
|
|
|
self.assertIs(should_remove, False)
|
|
self.assertIs(decode_req.kv_receiver, receiver)
|
|
self.assertEqual(receiver.clear_calls, 0)
|
|
self.assertEqual(decode_req.req.output_ids, [])
|
|
|
|
def test_commit_transfer_to_req_aborts_on_room_mismatch(self):
|
|
queue = DecodeTransferQueue.__new__(DecodeTransferQueue)
|
|
queue.metadata_buffers = cast(
|
|
Any,
|
|
SimpleNamespace(
|
|
get_buf=lambda idx: (
|
|
[FakeItem(0)],
|
|
[FakeItem(0)],
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
[FakeItem(99)],
|
|
)
|
|
),
|
|
)
|
|
queue.scheduler = cast(
|
|
Any,
|
|
SimpleNamespace(
|
|
server_args=SimpleNamespace(disaggregation_transfer_backend="mooncake")
|
|
),
|
|
)
|
|
queue.spec_algorithm = cast(Any, SimpleNamespace(is_none=lambda: True))
|
|
|
|
receiver = FakeReceiver()
|
|
decode_req = DecodeRequest(
|
|
req=cast(Any, FakeReq("corrupt", 3)),
|
|
kv_receiver=cast(Any, receiver),
|
|
metadata_buffer_index=9,
|
|
)
|
|
|
|
aborted = []
|
|
with patch(
|
|
"sglang.srt.disaggregation.decode.prepare_abort",
|
|
lambda req, message, status_code: aborted.append(
|
|
(req.rid, message, status_code)
|
|
),
|
|
):
|
|
should_remove = queue._commit_transfer_to_req(decode_req)
|
|
|
|
self.assertIs(should_remove, True)
|
|
self.assertEqual(receiver.clear_calls, 1)
|
|
self.assertIsNone(decode_req.kv_receiver)
|
|
self.assertEqual(len(aborted), 1)
|
|
self.assertEqual(aborted[0][0], "corrupt")
|
|
|
|
def test_pop_transferred_holds_eagle_metadata_slot_until_prebuilt_consumes(self):
|
|
queue = DecodeTransferQueue.__new__(DecodeTransferQueue)
|
|
allocator = FakeAllocator()
|
|
output_topk_p = torch.arange(16, dtype=torch.float32)
|
|
output_topk_index = torch.arange(16, dtype=torch.int64)
|
|
output_hidden_states = torch.arange(8, dtype=torch.float32)
|
|
queue.gloo_group = None
|
|
queue.req_to_metadata_buffer_idx_allocator = allocator
|
|
queue.tp_rank = 0
|
|
queue.metadata_buffers = cast(
|
|
Any,
|
|
SimpleNamespace(
|
|
get_buf=lambda idx: (
|
|
torch.tensor([7], dtype=torch.int32),
|
|
torch.tensor([320], dtype=torch.int32),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
output_topk_p,
|
|
output_topk_index,
|
|
output_hidden_states,
|
|
torch.tensor([3], dtype=torch.int64),
|
|
)
|
|
),
|
|
)
|
|
queue.tree_cache = cast(Any, object())
|
|
queue.scheduler = cast(
|
|
Any,
|
|
SimpleNamespace(
|
|
server_args=SimpleNamespace(disaggregation_transfer_backend="mooncake"),
|
|
stream_output=lambda reqs, return_logprob: None,
|
|
enable_metrics=False,
|
|
),
|
|
)
|
|
queue.spec_algorithm = cast(Any, SimpleNamespace(is_none=lambda: False))
|
|
|
|
receiver = FakeReceiver()
|
|
decode_req = DecodeRequest(
|
|
req=cast(Any, FakeReq("eagle", 3)),
|
|
kv_receiver=cast(Any, receiver),
|
|
metadata_buffer_index=9,
|
|
)
|
|
queue.queue = [decode_req]
|
|
|
|
with patch(
|
|
"sglang.srt.disaggregation.decode.poll_and_all_reduce",
|
|
return_value=[KVPoll.Success],
|
|
):
|
|
transferred = queue.pop_transferred()
|
|
|
|
self.assertEqual(transferred, [decode_req.req])
|
|
self.assertEqual(queue.queue, [])
|
|
self.assertEqual(allocator.freed, [])
|
|
|
|
self.assertEqual(decode_req.req.output_ids, [7])
|
|
self.assertEqual(decode_req.req.cached_tokens, 320)
|
|
self.assertEqual(decode_req.req.metadata_buffer_index, 9)
|
|
self.assertEqual(
|
|
decode_req.req.output_topk_p.data_ptr(), output_topk_p.data_ptr()
|
|
)
|
|
self.assertEqual(
|
|
decode_req.req.output_topk_index.data_ptr(), output_topk_index.data_ptr()
|
|
)
|
|
self.assertEqual(
|
|
decode_req.req.hidden_states_tensor.data_ptr(),
|
|
output_hidden_states.data_ptr(),
|
|
)
|
|
|
|
def test_free_decode_metadata_index_if_held_releases_once(self):
|
|
allocator = FakeAllocator()
|
|
req = FakeReq("eagle", 3)
|
|
req.metadata_buffer_index = 9
|
|
|
|
scheduler = SchedulerOutputProcessorMixin.__new__(SchedulerOutputProcessorMixin)
|
|
scheduler.req_to_metadata_buffer_idx_allocator = allocator
|
|
|
|
scheduler._free_decode_metadata_index_if_held(req)
|
|
scheduler._free_decode_metadata_index_if_held(req)
|
|
|
|
self.assertEqual(allocator.freed, [9])
|
|
self.assertEqual(req.metadata_buffer_index, -1)
|
|
|
|
def test_resume_retracted_reqs_compacts_queue_in_one_pass(self):
|
|
prealloc_queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
prealloc_queue.req_to_token_pool = cast(
|
|
Any, SimpleNamespace(available_size=lambda: 1)
|
|
)
|
|
prealloc_queue.token_to_kv_pool_allocator = cast(Any, object())
|
|
prealloc_queue.num_reserved_decode_tokens = 1
|
|
prealloc_queue._allocatable_tokens = (
|
|
lambda retractable_tokens=None, count_retracted=False: 8
|
|
)
|
|
prealloc_queue._pre_alloc = lambda req: setattr(req, "prealloc_done", True)
|
|
|
|
first = FakeReq("resume", 1)
|
|
first.origin_input_ids = [1, 2]
|
|
first.output_ids = [3]
|
|
skipped = FakeReq("skip", 2)
|
|
skipped.origin_input_ids = [4]
|
|
blocked = FakeReq("blocked", 3)
|
|
blocked.origin_input_ids = [5, 6, 7, 8, 9, 10, 11, 12]
|
|
|
|
prealloc_queue.retracted_queue = cast(Any, [first, skipped, blocked])
|
|
|
|
resumed = prealloc_queue.resume_retracted_reqs(
|
|
rids_to_check=["resume", "blocked"]
|
|
)
|
|
|
|
self.assertEqual(resumed, [first])
|
|
self.assertIs(first.is_retracted, False)
|
|
self.assertIs(first.prealloc_done, True)
|
|
self.assertEqual(len(first.load_calls), 1)
|
|
self.assertEqual(prealloc_queue.retracted_queue, [skipped, blocked])
|
|
|
|
def test_pop_preallocated_still_removes_failed_reqs_after_block(self):
|
|
streamed = []
|
|
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
queue._resolve_pending_reqs = lambda: None
|
|
queue._update_handshake_waiters = lambda rids_to_check=None: None
|
|
queue.req_to_token_pool = cast(
|
|
Any,
|
|
SimpleNamespace(
|
|
available_size=lambda: 1,
|
|
req_to_token=FakeTensor([[1, 2, 3, 4, 5, 6, 7, 8]]),
|
|
write=lambda *args, **kwargs: None,
|
|
),
|
|
)
|
|
queue.req_to_metadata_buffer_idx_allocator = cast(
|
|
Any, SimpleNamespace(available_size=lambda: 1, alloc=lambda: 7)
|
|
)
|
|
queue.token_to_kv_pool_allocator = cast(Any, SimpleNamespace(page_size=1))
|
|
queue.token_to_kv_pool = cast(Any, object())
|
|
queue.draft_token_to_kv_pool = None
|
|
queue.num_reserved_decode_tokens = 1
|
|
queue.scheduler = cast(
|
|
Any,
|
|
SimpleNamespace(
|
|
stream_output=lambda reqs, return_logprob: streamed.extend(
|
|
req.rid for req in reqs
|
|
),
|
|
enable_metrics=False,
|
|
running_batch=SimpleNamespace(reqs=[]),
|
|
sliding_window_size=4,
|
|
),
|
|
)
|
|
queue._allocatable_tokens = (
|
|
lambda retractable_tokens=None, count_retracted=True: 2
|
|
)
|
|
queue._pre_alloc = lambda req, **kwargs: (
|
|
setattr(req, "req_pool_idx", 0),
|
|
torch.arange(len(req.origin_input_ids), dtype=torch.int64),
|
|
)[1]
|
|
|
|
blocked_req = FakeReq("blocked", 1)
|
|
blocked_req.origin_input_ids = [1, 2, 3]
|
|
blocked_req.sampling_params = SimpleNamespace(max_new_tokens=8)
|
|
blocked = DecodeRequest(
|
|
req=cast(Any, blocked_req),
|
|
kv_receiver=cast(Any, FakeReceiver()),
|
|
waiting_for_input=True,
|
|
)
|
|
failed_req = FakeReq("failed", 2)
|
|
failed_req.finished_reason = FINISH_ABORT("boom")
|
|
failed = DecodeRequest(
|
|
req=cast(Any, failed_req),
|
|
kv_receiver=cast(Any, FakeReceiver()),
|
|
waiting_for_input=True,
|
|
)
|
|
tail = DecodeRequest(
|
|
req=cast(Any, FakeReq("tail", 3)),
|
|
kv_receiver=cast(Any, FakeReceiver()),
|
|
waiting_for_input=True,
|
|
)
|
|
queue.queue = cast(Any, [blocked, failed, tail])
|
|
|
|
preallocated, failed_reqs = queue.pop_preallocated(
|
|
rids_to_check=["blocked", "failed", "tail"]
|
|
)
|
|
|
|
self.assertEqual(preallocated, [])
|
|
self.assertEqual([req.req.rid for req in failed_reqs], ["failed"])
|
|
self.assertEqual(streamed, ["failed"])
|
|
self.assertEqual(queue.queue, [blocked, tail])
|
|
|
|
def test_pop_preallocated_compacts_queue_in_one_pass(self):
|
|
streamed = []
|
|
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
queue._resolve_pending_reqs = lambda: None
|
|
queue._update_handshake_waiters = lambda rids_to_check=None: None
|
|
queue.req_to_token_pool = cast(
|
|
Any,
|
|
SimpleNamespace(
|
|
available_size=lambda: 1,
|
|
req_to_token=FakeTensor([[1, 2, 3, 4, 5, 6, 7, 8]]),
|
|
write=lambda *args, **kwargs: None,
|
|
),
|
|
)
|
|
queue.req_to_metadata_buffer_idx_allocator = cast(
|
|
Any, SimpleNamespace(available_size=lambda: 1, alloc=lambda: 7)
|
|
)
|
|
queue.token_to_kv_pool_allocator = cast(Any, SimpleNamespace(page_size=1))
|
|
queue.token_to_kv_pool = cast(Any, object())
|
|
queue.draft_token_to_kv_pool = None
|
|
queue.num_reserved_decode_tokens = 1
|
|
queue.scheduler = cast(
|
|
Any,
|
|
SimpleNamespace(
|
|
stream_output=lambda reqs, return_logprob: streamed.extend(
|
|
req.rid for req in reqs
|
|
),
|
|
enable_metrics=False,
|
|
running_batch=SimpleNamespace(reqs=[]),
|
|
sliding_window_size=4,
|
|
),
|
|
)
|
|
queue._allocatable_tokens = (
|
|
lambda retractable_tokens=None, count_retracted=True: 6
|
|
)
|
|
queue._pre_alloc = lambda req, **kwargs: (
|
|
setattr(req, "req_pool_idx", 0),
|
|
torch.arange(len(req.origin_input_ids), dtype=torch.int64),
|
|
)[1]
|
|
|
|
skipped = DecodeRequest(
|
|
req=cast(Any, FakeReq("skip", 1)),
|
|
kv_receiver=cast(Any, FakeReceiver()),
|
|
waiting_for_input=True,
|
|
)
|
|
failed_req = FakeReq("failed", 2)
|
|
failed_req.finished_reason = FINISH_ABORT("boom")
|
|
failed = DecodeRequest(
|
|
req=cast(Any, failed_req),
|
|
kv_receiver=cast(Any, FakeReceiver()),
|
|
waiting_for_input=True,
|
|
)
|
|
waiting = DecodeRequest(
|
|
req=cast(Any, FakeReq("wait", 3)),
|
|
kv_receiver=cast(Any, FakeReceiver()),
|
|
waiting_for_input=False,
|
|
)
|
|
success_req = FakeReq("success", 4)
|
|
success_req.origin_input_ids = [1, 2]
|
|
success_req.sampling_params = SimpleNamespace(max_new_tokens=2)
|
|
success = DecodeRequest(
|
|
req=cast(Any, success_req),
|
|
kv_receiver=cast(Any, FakeReceiver()),
|
|
waiting_for_input=True,
|
|
)
|
|
blocked_req = FakeReq("blocked", 5)
|
|
blocked_req.origin_input_ids = [1, 2, 3, 4, 5, 6]
|
|
blocked_req.sampling_params = SimpleNamespace(max_new_tokens=8)
|
|
blocked = DecodeRequest(
|
|
req=cast(Any, blocked_req),
|
|
kv_receiver=cast(Any, FakeReceiver()),
|
|
waiting_for_input=True,
|
|
)
|
|
tail = DecodeRequest(
|
|
req=cast(Any, FakeReq("tail", 6)),
|
|
kv_receiver=cast(Any, FakeReceiver()),
|
|
waiting_for_input=True,
|
|
)
|
|
|
|
queue.queue = cast(Any, [skipped, failed, waiting, success, blocked, tail])
|
|
|
|
preallocated, failed_reqs = queue.pop_preallocated(
|
|
rids_to_check=["failed", "wait", "success", "blocked"]
|
|
)
|
|
|
|
self.assertEqual([req.req.rid for req in preallocated], ["success"])
|
|
self.assertEqual([req.req.rid for req in failed_reqs], ["failed"])
|
|
self.assertEqual(streamed, ["failed"])
|
|
self.assertEqual(success.metadata_buffer_index, 7)
|
|
self.assertEqual(queue.queue, [skipped, waiting, blocked, tail])
|
|
|
|
def test_get_new_prebuilt_batch_slices_waiting_queue_prefix(self):
|
|
allocator = FakeAllocator()
|
|
scheduler = cast(Any, SimpleNamespace())
|
|
scheduler.grammar_manager = SimpleNamespace(
|
|
has_waiting_grammars=lambda: False,
|
|
get_ready_grammar_requests=lambda: [],
|
|
)
|
|
scheduler._add_request_to_queue = lambda req: None
|
|
scheduler.waiting_queue = [FakeReq(f"req-{i}", i) for i in range(5)]
|
|
scheduler.running_batch = SimpleNamespace(batch_size=lambda: 1)
|
|
scheduler.req_to_token_pool = SimpleNamespace(size=8)
|
|
scheduler.max_running_requests = 4
|
|
scheduler.tree_cache = object()
|
|
scheduler.token_to_kv_pool_allocator = object()
|
|
scheduler.model_config = object()
|
|
scheduler.enable_overlap = False
|
|
scheduler.spec_algorithm = object()
|
|
scheduler.server_args = object()
|
|
scheduler.future_map = object()
|
|
scheduler.req_to_metadata_buffer_idx_allocator = allocator
|
|
scheduler._free_decode_metadata_index_if_held = (
|
|
SchedulerOutputProcessorMixin._free_decode_metadata_index_if_held.__get__(
|
|
scheduler
|
|
)
|
|
)
|
|
for i, req in enumerate(scheduler.waiting_queue[:3]):
|
|
req.metadata_buffer_index = 20 + i
|
|
|
|
captured = {}
|
|
|
|
def fake_init_new(reqs, *args, **kwargs):
|
|
captured["reqs"] = list(reqs)
|
|
batch = FakeBatch()
|
|
captured["batch"] = batch
|
|
return batch
|
|
|
|
with patch(
|
|
"sglang.srt.disaggregation.decode.ScheduleBatch.init_new", fake_init_new
|
|
):
|
|
batch = SchedulerDisaggregationDecodeMixin.get_new_prebuilt_batch(scheduler)
|
|
|
|
self.assertIs(batch, captured["batch"])
|
|
self.assertEqual(
|
|
[req.rid for req in captured["reqs"]], ["req-0", "req-1", "req-2"]
|
|
)
|
|
self.assertEqual(
|
|
[req.rid for req in scheduler.waiting_queue], ["req-3", "req-4"]
|
|
)
|
|
for req in captured["reqs"]:
|
|
self.assertEqual(req.init_next_round_calls, [scheduler.tree_cache])
|
|
self.assertIsNotNone(req.time_stats.forward_entry_time)
|
|
self.assertTrue(captured["batch"].prepared)
|
|
self.assertEqual(
|
|
captured["batch"].processed,
|
|
[(scheduler.server_args, scheduler.future_map)],
|
|
)
|
|
self.assertEqual(allocator.freed, [20, 21, 22])
|
|
for req in captured["reqs"]:
|
|
self.assertEqual(req.metadata_buffer_index, -1)
|
|
|
|
def test_get_new_prebuilt_batch_keeps_waiting_queue_when_no_capacity(self):
|
|
scheduler = cast(Any, SimpleNamespace())
|
|
scheduler.grammar_manager = SimpleNamespace(
|
|
has_waiting_grammars=lambda: False,
|
|
get_ready_grammar_requests=lambda: [],
|
|
)
|
|
scheduler._add_request_to_queue = lambda req: None
|
|
scheduler.waiting_queue = [FakeReq(f"req-{i}", i) for i in range(3)]
|
|
scheduler.running_batch = SimpleNamespace(batch_size=lambda: 4)
|
|
scheduler.req_to_token_pool = SimpleNamespace(size=8)
|
|
scheduler.max_running_requests = 4
|
|
scheduler.tree_cache = object()
|
|
scheduler.token_to_kv_pool_allocator = object()
|
|
scheduler.model_config = object()
|
|
scheduler.enable_overlap = False
|
|
scheduler.spec_algorithm = object()
|
|
scheduler.server_args = object()
|
|
scheduler.future_map = object()
|
|
|
|
with patch(
|
|
"sglang.srt.disaggregation.decode.ScheduleBatch.init_new",
|
|
side_effect=AssertionError(
|
|
"init_new should not be called without capacity"
|
|
),
|
|
):
|
|
batch = SchedulerDisaggregationDecodeMixin.get_new_prebuilt_batch(scheduler)
|
|
|
|
self.assertIsNone(batch)
|
|
self.assertEqual(
|
|
[req.rid for req in scheduler.waiting_queue], ["req-0", "req-1", "req-2"]
|
|
)
|
|
for req in scheduler.waiting_queue:
|
|
self.assertEqual(req.init_next_round_calls, [])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|