diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 428283c296eb..70a460059507 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -29,7 +29,15 @@ from concurrent.futures import Future from dataclasses import dataclass from http import HTTPStatus -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Optional, + Tuple, +) import msgspec import numpy as np @@ -89,6 +97,7 @@ ) from sglang.srt.managers.scheduler_components.request_list import ( ListKind, + RequestList, _RequestList, ) from sglang.srt.managers.utils import ( @@ -407,14 +416,6 @@ class DecodeRequest: hicache_load_consumer_index: int = -1 hicache_restore_status: HiCacheRestoreResult = HiCacheRestoreResult.PENDING - @property - def seqlen(self) -> int: - return self.req.seqlen - - @property - def priority(self) -> Optional[int]: - return self.req.priority - class _PendingPublish(msgspec.Struct, frozen=True): allocated: torch.cuda.Event @@ -492,7 +493,11 @@ def __init__( self.num_reserved_decode_tokens = num_reserved_decode_tokens self.transfer_backend = transfer_backend # Queue for requests pending pre-allocation - self.queue: List[DecodeRequest] = [] + self.queue = _RequestList( + ListKind.PD_PREALLOC, + abort=self._abort_prealloc, + req_of=lambda x: x.req, + ) self.retracted_queue = _RequestList( ListKind.PD_RETRACTED, abort=self._abort_retracted ) @@ -1003,7 +1008,7 @@ def _create_receiver_and_enqueue( decode_req = DecodeRequest( req=req, kv_receiver=kv_receiver, is_rebootstrap=is_rebootstrap ) - self.queue.append(decode_req) + self.queue.join([decode_req]) return decode_req def hold_rebootstrap(self, req: Req) -> None: @@ -1022,7 +1027,7 @@ def hold_rebootstrap(self, req: Req) -> None: def enqueue_held_rebootstrap(self) -> None: """Enqueue all staged rebootstrap requests when generation resumes.""" - held = self.held_rebootstrap_reqs.take(lambda req: True) + held = self.held_rebootstrap_reqs.take() for req in held: self.add(req, is_rebootstrap=True) @@ -1063,6 +1068,14 @@ def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool: return True return False + def _abort_prealloc(self, recv_req: AbortReq) -> None: + for decode_req in self.queue.items(): + if recv_req.matches(decode_req.req.rid): + logger.debug(f"Abort prealloc queue request. {decode_req.req.rid=}") + decode_req.kv_receiver.abort() + if get_parallel().pp_size > 1: + prepare_abort(decode_req.req, "Aborted by AbortReq.") + def _abort_retracted(self, recv_req: AbortReq) -> None: # Retracted requests hold a KV backup. for req in self.retracted_queue.take(lambda req: recv_req.matches(req.rid)): @@ -1084,9 +1097,9 @@ def _abort_held(self, recv_req: AbortReq) -> None: make_abort_req(req, finished_reason=recv_req.finished_reason), req ) - def request_lists(self) -> Dict[ListKind, List[Req]]: + def request_lists(self) -> Dict[ListKind, RequestList]: return { - ListKind.PD_PREALLOC: [decode_req.req for decode_req in self.queue], + ListKind.PD_PREALLOC: self.queue, ListKind.PD_RETRACTED: self.retracted_queue, ListKind.PD_HELD: self.held_rebootstrap_reqs, } @@ -1227,17 +1240,18 @@ def _update_handshake_waiters( # is blocked. if self.pp_size > 1: polls = poll_and_all_reduce_pp( - (decode_req.req.rid for decode_req in self.queue), + (req.rid for req in self.queue), KVPoll.WaitingForInput, pp_good_rids, pp_bad_rids, ) else: polls = poll_and_all_reduce( - [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group + [decode_req.kv_receiver for decode_req in self.queue.items()], + self.gloo_group, ) - for decode_req, poll in zip(self.queue, polls): + for decode_req, poll in zip(self.queue.items(), polls): if poll is None: continue if rids_to_check is not None and decode_req.req.rid not in rids_to_check: @@ -1475,14 +1489,17 @@ def pop_preallocated( full_allocatable_tokens -= reserved_restore_tokens # Sort by priority before any index-based bookkeeping so that both the # abort-scan loop and the preallocation loop operate on the same order. + queue = list(self.queue.items()) if self.scheduler.enable_priority_scheduling: priority_sign = ( 1 if self.scheduler.schedule_low_priority_values_first else -1 ) - self.queue.sort(key=lambda r: r.req.priority * priority_sign) + queue.sort(key=lambda r: r.req.priority * priority_sign) + self.queue.clear() + self.queue.join(queue) # First, remove all failed requests from the queue - for i, decode_req in enumerate(self.queue): + for i, decode_req in enumerate(queue): if rids_to_check is not None and decode_req.req.rid not in rids_to_check: continue if isinstance(decode_req.req.finished_reason, FINISH_ABORT): @@ -1533,11 +1550,11 @@ def pop_preallocated( running_loras = { req.lora_id for batch in running_batches for req in batch.reqs } - running_loras.update(r.req.lora_id for r in self.transfer_queue.queue) + running_loras.update(req.lora_id for req in self.transfer_queue.queue) running_loras.update(req.lora_id for req in self.scheduler.waiting_queue) # Then, preallocate the remaining requests if possible - for i, decode_req in enumerate(self.queue): + for i, decode_req in enumerate(queue): if rids_to_check is not None and decode_req.req.rid not in rids_to_check: continue @@ -1781,9 +1798,7 @@ def pop_preallocated( if decode_req.req.kv.holds_mamba and not decode_req.req.kv.holds_kv: release_kv_cache(decode_req.req, self.tree_cache, checkpoint=False) - self.queue = [ - entry for i, entry in enumerate(self.queue) if i not in indices_to_remove - ] + self.queue.leave([queue[i].req for i in indices_to_remove]) return preallocated_reqs, failed_reqs @@ -2240,7 +2255,10 @@ def _host_destinations_quiescent(self) -> bool: queue, and no failed transfer awaiting its drain.""" return not ( self.has_published_destinations - or any(decode_req.host_staged for decode_req in self.transfer_queue.queue) + or any( + decode_req.host_staged + for decode_req in self.transfer_queue.queue.items() + ) or self.transfer_queue.has_pending_deferred_releases() ) @@ -2258,7 +2276,7 @@ def note_destinations_queued(self, count: int) -> None: def num_tokens_pre_allocated(self): return sum( decode_req.req.extend_end - for decode_req in self.transfer_queue.queue + for decode_req in self.transfer_queue.queue.items() if not decode_req.host_staged ) @@ -2268,7 +2286,7 @@ def pre_allocated_token_counts(self) -> Tuple[int, int]: charged the way the pool stats charge a request's private KV.""" page_size = self.token_to_kv_pool_allocator.page_size full_tokens = swa_tokens = 0 - for decode_req in self.transfer_queue.queue: + for decode_req in self.transfer_queue.queue.items(): if decode_req.host_staged: continue full_tokens += kv_private_tokens(decode_req.req.kv, page_size) @@ -2296,7 +2314,7 @@ def _need_space_for_single_req( def _active_req_count(self, extra_reserved_reqs: int = 0) -> int: return ( len(self.scheduler.running_batch.reqs) - + sum(not r.host_staged for r in self.transfer_queue.queue) + + sum(not r.host_staged for r in self.transfer_queue.queue.items()) + len(self.scheduler.waiting_queue) + extra_reserved_reqs ) @@ -2836,7 +2854,11 @@ def __init__( scheduler: Scheduler, tree_cache: BasePrefixCache, ): - self.queue: List[DecodeRequest] = [] + self.queue = _RequestList( + ListKind.PD_RECEIVING, + abort=self._abort_receiving, + req_of=lambda x: x.req, + ) self.gloo_group = get_parallel().attn_tp_group.cpu_group self.req_to_metadata_buffer_idx_allocator = req_to_metadata_buffer_idx_allocator self.tp_rank = get_parallel().tp_rank @@ -2860,22 +2882,35 @@ def __init__( # Aborted-mid-transfer requests whose KV pages/slot are held until drained # or timed out. Entries: # (decode_req, start_time, deadline, metadata_idx, required_acks). - self._deferred_releases: List[Tuple[DecodeRequest, float, float, int, int]] = [] + self._deferred_releases = _RequestList( + ListKind.PD_RELEASING, req_of=lambda x: x[0].req + ) - def request_lists(self) -> Dict[ListKind, List[Req]]: + def request_lists(self) -> Dict[ListKind, RequestList]: return { - ListKind.PD_RECEIVING: [decode_req.req for decode_req in self.queue], + ListKind.PD_RECEIVING: self.queue, # Aborted, with KV held until the transfer drains. - ListKind.PD_RELEASING: [ - decode_req.req for decode_req, *_ in self._deferred_releases - ], + ListKind.PD_RELEASING: self._deferred_releases, } + def _abort_receiving(self, recv_req: AbortReq) -> None: + for decode_req in self.queue.items(): + if recv_req.matches(decode_req.req.rid): + logger.debug(f"Abort transfer queue request. {decode_req.req.rid=}") + if decode_req.host_staged: + # Keep the host destination alive until prefill stops writing. + prepare_abort(decode_req.req, "Aborted by AbortReq.") + continue + # The receiver arms drain-ack accounting before sending the + # ABORT (see CommonKVReceiver._send_abort_notification), so + # an ack racing back is never dropped. + decode_req.kv_receiver.abort() + def add(self, decode_req: DecodeRequest) -> None: - self.queue.append(decode_req) + self.queue.join([decode_req]) def extend(self, decode_reqs: List[DecodeRequest]) -> None: - self.queue.extend(decode_reqs) + self.queue.join(decode_reqs) # This queue now covers them. prealloc_queue = self.scheduler.disagg_decode_prealloc_queue if prealloc_queue is not None: @@ -3066,20 +3101,20 @@ def _commit_transfer_to_req(self, decode_req: DecodeRequest): def _pollers(self) -> list: if self.scheduler.enable_decode_hicache: - return [HiCacheRestoreGatedKVReceiver(dr) for dr in self.queue] - return [dr.kv_receiver for dr in self.queue] + return [HiCacheRestoreGatedKVReceiver(dr) for dr in self.queue.items()] + return [dr.kv_receiver for dr in self.queue.items()] def _poll_with_metadata_gate(self) -> List[int]: return poll_and_all_reduce( self._pollers(), self.gloo_group, - decode_reqs=self.queue, + decode_reqs=list(self.queue.items()), metadata_buffers=self.metadata_buffers, ) def _poll_with_staging(self) -> list: return poll_and_all_reduce_with_staging( - self.queue, + list(self.queue.items()), self.staging_handler, self.gloo_group, metadata_buffers=self.metadata_buffers, @@ -3127,7 +3162,7 @@ def pop_transferred( def publish_pending_metadata(self, polls: Optional[List[int]] = None) -> None: pending = [] ready = [] - for i, decode_req in enumerate(self.queue): + for i, decode_req in enumerate(self.queue.items()): metadata = decode_req.pending_metadata if metadata is None: continue @@ -3166,7 +3201,7 @@ def pop_transferred_with_outcomes( self._process_hicache_local_restores( [ decode_req - for decode_req in self.queue + for decode_req in self.queue.items() if rids_to_check is None or decode_req.req.rid in rids_to_check ] ) @@ -3184,7 +3219,8 @@ def pop_transferred_with_outcomes( failed: List[Tuple[int, DecodeRequest, bool]] = [] outcomes: Dict[str, bool] = {} device_admission_refused = False - for i, (decode_req, poll) in enumerate(zip(self.queue, polls)): + queue = list(self.queue.items()) + for i, (decode_req, poll) in enumerate(zip(queue, polls)): if rids_to_check is not None: if decode_req.req.rid not in rids_to_check: continue @@ -3336,21 +3372,17 @@ def pop_transferred_with_outcomes( # Held for deferred release; metadata buffer freed at resolve time. continue if self.enable_staging and self.staging_handler.is_staging_room( - self.queue[i].req.bootstrap_room + queue[i].req.bootstrap_room ): - self.staging_handler.unregister_decode_req( - self.queue[i].req.bootstrap_room - ) - idx = self.queue[i].metadata_buffer_index + self.staging_handler.unregister_decode_req(queue[i].req.bootstrap_room) + idx = queue[i].metadata_buffer_index assert idx != -1 # Reset so the next owner sees actual_room == 0 ("not yet written") # instead of the stale value, avoiding a false-positive mismatch. self.metadata_buffers.clear_slot(idx) self.req_to_metadata_buffer_idx_allocator.free(idx) - self.queue = [ - entry for i, entry in enumerate(self.queue) if i not in indices_to_remove - ] + self.queue.leave([queue[i].req for i in indices_to_remove]) polls = [poll for i, poll in enumerate(polls) if i not in indices_to_remove] self.publish_pending_metadata(polls) @@ -3378,14 +3410,16 @@ def _defer_release(self, decode_req: DecodeRequest) -> None: or not decode_req.kv_receiver.abort_notified else len(decode_req.kv_receiver.bootstrap_infos) ) - self._deferred_releases.append( - ( - decode_req, - start_time, - deadline, - decode_req.metadata_buffer_index, - required_acks, - ) + self._deferred_releases.join( + [ + ( + decode_req, + start_time, + deadline, + decode_req.metadata_buffer_index, + required_acks, + ) + ] ) def _do_release(self, decode_req: DecodeRequest, idx: int) -> None: @@ -3423,7 +3457,8 @@ def resolve_deferred_releases(self) -> None: now = time.monotonic() drained = [] ready = [] - for decode_req, _, deadline, _, required_acks in self._deferred_releases: + releases = list(self._deferred_releases.items()) + for decode_req, _, deadline, _, required_acks in releases: if self._is_host_destination(decode_req): # A timeout cannot prove that a remote write has stopped. is_drained = decode_req.kv_receiver.is_abort_release_safe() @@ -3442,9 +3477,8 @@ def resolve_deferred_releases(self) -> None: ready = all_reduce_flags( flags=ready, group=self.gloo_group, op=torch.distributed.ReduceOp.MIN ) - still_held = [] to_release = [] - for entry, is_ready, is_drained in zip(self._deferred_releases, ready, drained): + for entry, is_ready, is_drained in zip(releases, ready, drained): decode_req, start_time, _, idx, _ = entry if is_ready: to_release.append( @@ -3456,11 +3490,9 @@ def resolve_deferred_releases(self) -> None: is_drained, ) ) - else: - still_held.append(entry) - # Commit the survivors before releasing so a _do_release exception can't - # leave a released entry in the list (double-free / None receiver on retry). - self._deferred_releases = still_held + # Drop the released entries before releasing so a _do_release exception + # can't leave one in the list (double-free / None receiver on retry). + self._deferred_releases.leave([decode_req.req for decode_req, *_ in to_release]) for decode_req, start_time, idx, room, drained in to_release: if not drained: logger.warning( diff --git a/python/sglang/srt/disaggregation/decode_hicache_mixin.py b/python/sglang/srt/disaggregation/decode_hicache_mixin.py index 46cbb8a9735e..891e33c6fed8 100644 --- a/python/sglang/srt/disaggregation/decode_hicache_mixin.py +++ b/python/sglang/srt/disaggregation/decode_hicache_mixin.py @@ -169,7 +169,7 @@ def _hicache_pending_restore_tokens(self) -> int: return 0 return sum( dr.prefix_match.restore_token_count - for dr in self.transfer_queue.queue + for dr in self.transfer_queue.queue.items() if dr.prefix_match is not None and dr.hicache_restore_lock is None ) diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 039d6a11dff3..47eda5fd75c1 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -182,7 +182,7 @@ def busy() -> bool: return bool( any( not decode_req.host_staged - for decode_req in scheduler.disagg_decode_transfer_queue.queue + for decode_req in scheduler.disagg_decode_transfer_queue.queue.items() ) or scheduler.disagg_decode_prealloc_queue.has_published_destinations or ( diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 8f7459279dcb..9c9b6efcad68 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -2,7 +2,15 @@ import logging import weakref -from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Tuple, Union +from typing import ( + TYPE_CHECKING, + Dict, + List, + NamedTuple, + Optional, + Tuple, + Union, +) import torch @@ -39,7 +47,11 @@ def load_blocks_to_device_buffer_mha(*args, **kwargs): from sglang.srt.configs.model_config import dsa_layer_skips_topk, is_deepseek_dsa from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import Req -from sglang.srt.managers.scheduler_components.request_list import ListKind +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + RequestList, + _RequestList, +) from sglang.srt.mem_cache.allocator.hisparse import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, HiSparseTokenToKVPoolAllocator, @@ -666,7 +678,9 @@ def __init__( self.write_staging_stream = device_module.Stream() self.decode_backup_stream = device_module.Stream() - self.ack_staging_queue: List[HiSparseAct] = [] + self.ack_staging_queue = _RequestList( + ListKind.HISPARSE_STAGING, req_of=lambda x: x.req + ) self.decode_producer_stream = None self._backup_done_event = device_module.Event() self._has_pending_backup = False @@ -874,7 +888,7 @@ def admit_request_into_staging(self, req: Req) -> None: if device_indices.is_cuda: device_indices.record_stream(self.write_staging_stream) - self.ack_staging_queue.append(HiSparseAct(start_event, finish_event, req)) + self.ack_staging_queue.join([HiSparseAct(start_event, finish_event, req)]) def admit_request_direct(self, req: Req) -> None: """Direct-to-host path: KV data already resides in host pool via RDMA. @@ -1058,19 +1072,17 @@ def _grow_device_buffers( reserved_positions = (seq_lens - 1).clamp(max=self.device_buffer_size) return self.req_to_device_buffer[req_pool_indices, reserved_positions] - def request_lists(self) -> Dict[ListKind, List[Req]]: - return {ListKind.HISPARSE_STAGING: [act.req for act in self.ack_staging_queue]} - - def has_ongoing_staging(self) -> bool: - return len(self.ack_staging_queue) > 0 + def request_lists(self) -> Dict[ListKind, RequestList]: + return {ListKind.HISPARSE_STAGING: self.ack_staging_queue} def collect_ready_reqs(self) -> List[Req]: ready_reqs: List[Req] = [] if len(self.ack_staging_queue) == 0: return ready_reqs + staged = list(self.ack_staging_queue.items()) finish_count = 0 - for _, finish_event, _ in self.ack_staging_queue: + for _, finish_event, _ in staged: if not finish_event.query(): break finish_count += 1 @@ -1083,13 +1095,12 @@ def collect_ready_reqs(self) -> List[Req]: group=self.tp_group, ) finish_count = int(queue_size.item()) - while finish_count > 0: - _, _, req = self.ack_staging_queue.pop(0) + for _, _, req in staged[:finish_count]: + self.ack_staging_queue.leave([req]) # prepare device buffer and update req self.alloc_device_buffer(req) self._skip_first_backup[req.kv.req_pool_idx] = True req.hisparse_staging = False - finish_count -= 1 ready_reqs.append(req) return ready_reqs @@ -1371,9 +1382,7 @@ def abort_staging_request(self, req: Req) -> None: but has not yet completed (i.e. req.hisparse_staging is True). """ # Remove from staging queue - self.ack_staging_queue = [ - act for act in self.ack_staging_queue if act.req is not req - ] + self.ack_staging_queue.leave([req]) # Wait for any in-flight staging DMA to complete before freeing self.write_staging_stream.synchronize() diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 995d612b6fbc..171199a00c7d 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2582,6 +2582,7 @@ def maybe_init_uocr_runtime(self) -> None: decode_graph_io=lambda: resolve_uocr_decode_graph_io( runner.decode_cuda_graph_runner ), + abort_windows=self._abort_uocr_doc_windows, ) self.resolve_forward_inputs = self.uocr_runtime.resolve_forward_inputs @@ -5698,6 +5699,7 @@ def is_fully_idle(self, for_health_check=False, ignore_waiting=False) -> bool: ListKind.PD_HELD, # Counted by uocr blocks_flush below. ListKind.UOCR_PARKED, + ListKind.UOCR_RELEASED, # Finished, and held only on the output rank; all ranks must agree. ListKind.DEFERRED, } @@ -5989,12 +5991,10 @@ def record_weight_version_change(self, new_version: Optional[str]) -> None: *self.collect_inflight_reqs(), *self.waiting_queue, *([self.chunked_req] if self.chunked_req is not None else []), - *self.deferred_responses.held_reqs(), + *self.deferred_responses.held, } if self.hisparse_coordinator is not None: - live_reqs.update( - act.req for act in self.hisparse_coordinator.ack_staging_queue - ) + live_reqs.update(self.hisparse_coordinator.ack_staging_queue) num_recorded = record_weight_version_events(live_reqs, old_version=old_version) logger.info( f"Weight version changed. {old_version=} {new_version=} {num_recorded=}" @@ -6019,7 +6019,6 @@ def _collect_inflight_batches(self) -> list[ScheduleBatch | None]: return inflight_batches def abort_request(self, recv_req: AbortReq): - self.deferred_responses.abort(recv_req) if (chunked_req := self.chunked_req) is not None: if recv_req.matches(chunked_req.rid): self._pending_chunked_abort_req = chunked_req @@ -6035,15 +6034,6 @@ def abort_request(self, recv_req: AbortReq): if recv_req.matches(req.rid): to_del.append(i) - if self.uocr_runtime is not None: - # Parked document windows are off the waiting queue; abort them alike. - self._abort_uocr_doc_windows( - self.uocr_runtime.drop_parked( - "" if recv_req.abort_all else recv_req.rid - ), - finished_reason=recv_req.finished_reason, - ) - # Sort in reverse order to avoid index issues when deleting for i in reversed(to_del): # Abort method 1: directly pop from the queue @@ -6082,29 +6072,6 @@ def abort_request(self, recv_req: AbortReq): if isinstance(request_list, RequestList): request_list.abort(recv_req) - # Delete requests not in the waiting queue when PD disaggregation is enabled - if self.disaggregation_mode == DisaggregationMode.DECODE: - # Abort requests that have not yet finished preallocation - for decode_req in self.disagg_decode_prealloc_queue.queue: - if recv_req.matches(decode_req.req.rid): - logger.debug(f"Abort prealloc queue request. {decode_req.req.rid=}") - decode_req.kv_receiver.abort() - if get_parallel().pp_size > 1: - prepare_abort(decode_req.req, "Aborted by AbortReq.") - - # Abort requests waiting for kvcache to release tree cache - for decode_req in self.disagg_decode_transfer_queue.queue: - if recv_req.matches(decode_req.req.rid): - logger.debug(f"Abort transfer queue request. {decode_req.req.rid=}") - if decode_req.host_staged: - # Keep the host destination alive until prefill stops writing. - prepare_abort(decode_req.req, "Aborted by AbortReq.") - continue - # The receiver arms drain-ack accounting before sending the - # ABORT (see CommonKVReceiver._send_abort_notification), so - # an ack racing back is never dropped. - decode_req.kv_receiver.abort() - # Delete requests in the running batch for req in self.collect_inflight_reqs(): if not req.finished() and recv_req.matches(req.rid): diff --git a/python/sglang/srt/managers/scheduler_components/deferred_responses.py b/python/sglang/srt/managers/scheduler_components/deferred_responses.py index dbcb3f7f1138..501462cc23d3 100644 --- a/python/sglang/srt/managers/scheduler_components/deferred_responses.py +++ b/python/sglang/srt/managers/scheduler_components/deferred_responses.py @@ -1,11 +1,15 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Callable, Dict, Optional from sglang.srt.disaggregation.mode import DisaggregationMode from sglang.srt.managers.req_types import FINISH_ABORT -from sglang.srt.managers.scheduler_components.request_list import ListKind +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + RequestList, + _RequestList, +) if TYPE_CHECKING: from sglang.srt.managers.deferred_output import ResponseCompletion @@ -32,22 +36,22 @@ def __init__( self.disaggregation_mode = disaggregation_mode self.output_streamer = output_streamer self.send_abort = send_abort - self.held: Dict[str, Tuple[Req, ResponseCompletion]] = {} + # (req, completion) per held response. + self.held = _RequestList( + ListKind.DEFERRED, abort=self._abort, req_of=lambda x: x[0] + ) - def held_reqs(self) -> List[Req]: - return [req for req, _ in self.held.values()] - - def request_lists(self) -> Dict[ListKind, List[Req]]: - return {ListKind.DEFERRED: self.held_reqs()} + def request_lists(self) -> Dict[ListKind, RequestList]: + return {ListKind.DEFERRED: self.held} def defer(self, req: Req, completion: ResponseCompletion) -> None: if self.disaggregation_mode != DisaggregationMode.NULL: raise ValueError("deferred responses do not support PD disaggregation") if not req.finished() or req.finished_output: raise ValueError("only finished, unstreamed responses can be deferred") - if req.rid in self.held: + if any(held.rid == req.rid for held in self.held): raise ValueError(f"response {req.rid!r} is already deferred") - self.held[req.rid] = (req, completion) + self.held.join([(req, completion)]) req.defer_output = True self.output_streamer.defer_outputs = True @@ -55,15 +59,15 @@ def stream_released(self) -> None: if not self.held: return ready = [] - for rid, (req, completion) in list(self.held.items()): + for req, completion in self.held.items(): if not completion.done(): continue - del self.held[rid] + self.held.leave([req]) req.defer_output = False try: completion.result() except Exception as exc: - logger.exception("Deferred response failed for %s", rid) + logger.exception("Deferred response failed for %s", req.rid) self._abort_response( req, FINISH_ABORT(str(exc), status_code=500).to_json() ) @@ -84,12 +88,12 @@ def _abort_response( req.finished_output = True self.send_abort(req, finished_reason) - def abort(self, recv_req: AbortReq) -> None: - for rid, (req, completion) in list(self.held.items()): - if recv_req.matches(rid): - del self.held[rid] + def _abort(self, recv_req: AbortReq) -> None: + for req, completion in self.held.items(): + if recv_req.matches(req.rid): + self.held.leave([req]) try: completion.cancel() except Exception: - logger.exception("Failed to cancel deferred work for %s", rid) + logger.exception("Failed to cancel deferred work for %s", req.rid) self._abort_response(req, recv_req.finished_reason) diff --git a/python/sglang/srt/managers/scheduler_components/load_inquirer.py b/python/sglang/srt/managers/scheduler_components/load_inquirer.py index 2ca992438da2..585ca9008470 100644 --- a/python/sglang/srt/managers/scheduler_components/load_inquirer.py +++ b/python/sglang/srt/managers/scheduler_components/load_inquirer.py @@ -199,7 +199,7 @@ def get_loads(self) -> LoadSnapshot: ) ready_reqs = [ decode_req.req - for decode_req in self.get_disagg_decode_prealloc_queue().queue + for decode_req in self.get_disagg_decode_prealloc_queue().queue.items() if decode_req.waiting_for_input ] decode_prealloc_ready = len(ready_reqs) diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 8fa273de5df9..3a442a15dd4a 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -1449,7 +1449,7 @@ def _update_decode_queue_stats(self, priority_enabled: bool) -> None: self.scheduler.disagg_decode_transfer_queue.num_pending_deferred_releases() ) host_reqs = ( - [req for req in transfer_queue if req.host_staged] + [d.req for d in transfer_queue.items() if d.host_staged] if get_disagg().disaggregation_decode_host_receive_threshold < 1 else [] ) diff --git a/python/sglang/srt/managers/scheduler_components/request_list.py b/python/sglang/srt/managers/scheduler_components/request_list.py index d92f4492b748..46ab6004e53f 100644 --- a/python/sglang/srt/managers/scheduler_components/request_list.py +++ b/python/sglang/srt/managers/scheduler_components/request_list.py @@ -4,6 +4,7 @@ from abc import ABC, abstractmethod from typing import ( TYPE_CHECKING, + Any, Callable, Dict, Iterable, @@ -37,6 +38,7 @@ class ListKind(enum.Enum): GRAMMAR = enum.auto() DLLM = enum.auto() UOCR_PARKED = enum.auto() + UOCR_RELEASED = enum.auto() DEFERRED = enum.auto() @@ -57,9 +59,10 @@ def leave(self, reqs: Iterable[Req]) -> None: ... def abort(self, recv_req: AbortReq) -> None: """Abort the requests `recv_req` names.""" - def take(self, pred: Callable[[Req], bool]) -> List[Req]: - """Drop the requests matching `pred` and return them, in list order.""" - taken = [req for req in self if pred(req)] + def take(self, pred: Optional[Callable[[Req], bool]] = None) -> List[Req]: + """Drop the requests matching `pred`, or all of them, and return them in + list order.""" + taken = [req for req in self if pred is None or pred(req)] self.leave(taken) return taken @@ -73,50 +76,63 @@ def __len__(self) -> int: class _RequestList(RequestList): """Requests in join order, held by identity. Each change rebinds a new dict, - so leaving while iterating is safe.""" + so leaving while iterating is safe. An owner that wraps each request joins + the wrappers and passes `req_of` to unwrap them.""" def __init__( self, kind: ListKind, *, abort: Optional[Callable[[AbortReq], None]] = None, + req_of: Callable[[Any], Req] = lambda x: x, ) -> None: self.kind = kind self._abort = abort - self._reqs: Dict[Req, None] = {} - - def join(self, reqs: Iterable[Req]) -> None: - """Add the given requests; ones already in keep their place.""" - joining = [req for req in reqs if req not in self._reqs] + self._req_of = req_of + self._items: Dict[Req, Any] = {} + + def join(self, items: Iterable) -> None: + """Add the given items; ones whose request is already in keep their place.""" + joining = {} + for item in items: + req = self._req_of(item) + if req not in self._items: + joining.setdefault(req, item) if joining: - self._reqs = {**self._reqs, **dict.fromkeys(joining)} + self._items = {**self._items, **joining} def leave(self, reqs: Iterable[Req]) -> None: """Drop the given requests; ones not in are ignored.""" - self._drop({req for req in reqs if req in self._reqs}) + self._drop({req for req in reqs if req in self._items}) def leave_finished(self) -> None: - self._drop({req for req in self._reqs if req.finished()}) + self._drop({req for req in self._items if req.finished()}) def clear(self) -> None: - self._reqs = {} + self._items = {} def abort(self, recv_req: AbortReq) -> None: if self._abort is not None: self._abort(recv_req) + def items(self) -> Iterator: + """The joined items, in list order.""" + return iter(self._items.values()) + def __iter__(self) -> Iterator[Req]: - return iter(self._reqs) + return iter(self._items) def __len__(self) -> int: - return len(self._reqs) + return len(self._items) def __contains__(self, req: Req) -> bool: - return req in self._reqs + return req in self._items def _drop(self, leaving: Set[Req]) -> None: if leaving: - self._reqs = {req: None for req in self._reqs if req not in leaving} + self._items = { + req: item for req, item in self._items.items() if req not in leaving + } def move(reqs: Iterable[Req], src: RequestList, dst: RequestList) -> None: diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index ab80e818ec0f..9deec6545af7 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -1748,7 +1748,7 @@ def _pp_pd_get_prealloc_ids(self: Scheduler): if self.pp_group.is_first_rank: # First rank, pop the preallocated reqs from the prealloc queue good_prealloc_rids, bad_prealloc_rids = self.get_rids( - self.disagg_decode_prealloc_queue.queue, + list(self.disagg_decode_prealloc_queue.queue.items()), False, [KVPoll.WaitingForInput], [KVPoll.Failed], @@ -1758,7 +1758,7 @@ def _pp_pd_get_prealloc_ids(self: Scheduler): prev_prealloc_rids = self._pp_recv_pyobj_from_prev_stage() prev_good_prealloc_rids, prev_bad_prealloc_rids = prev_prealloc_rids curr_good_prealloc_rids, curr_bad_prealloc_rids = self.get_rids( - self.disagg_decode_prealloc_queue.queue, + list(self.disagg_decode_prealloc_queue.queue.items()), False, [KVPoll.WaitingForInput], [KVPoll.Failed], @@ -1771,9 +1771,9 @@ def _pp_pd_get_prealloc_ids(self: Scheduler): ) # Same abort routing as the prefill bootstrap consensus above. aborted_rids = { - decode_req.req.rid - for decode_req in self.disagg_decode_prealloc_queue.queue - if isinstance(decode_req.req.finished_reason, FINISH_ABORT) + req.rid + for req in self.disagg_decode_prealloc_queue.queue + if isinstance(req.finished_reason, FINISH_ABORT) } good_prealloc_rids, bad_prealloc_rids = self._route_aborts_to_bad( good_prealloc_rids, bad_prealloc_rids, aborted_rids @@ -1805,7 +1805,7 @@ def _pp_pd_get_decode_transferred_ids(self: Scheduler) -> Dict[str, bool]: waiting queues stay in step with the others. """ succeeded, failed = self.get_rids( - self.disagg_decode_transfer_queue.queue, + list(self.disagg_decode_transfer_queue.queue.items()), False, [KVPoll.Success], [KVPoll.Failed], diff --git a/python/sglang/srt/models/unlimited_ocr_v2_support/document_prefetch.py b/python/sglang/srt/models/unlimited_ocr_v2_support/document_prefetch.py index 037c9ffed4c0..0287ad846d3a 100644 --- a/python/sglang/srt/models/unlimited_ocr_v2_support/document_prefetch.py +++ b/python/sglang/srt/models/unlimited_ocr_v2_support/document_prefetch.py @@ -22,6 +22,12 @@ PrefetchedEmbeddingFunc, _move_items_to_device, ) +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + RequestList, + _RequestList, + move, +) from sglang.srt.models.unlimited_ocr_v2_support.document_reuse import ( forced_len, req_state, @@ -29,6 +35,7 @@ from sglang.srt.models.unlimited_ocr_v2_support.window_params import KEY_KEEP if TYPE_CHECKING: + from sglang.srt.managers.io_struct import AbortReq from sglang.srt.managers.req_types import MultimodalDataItem from sglang.srt.managers.schedule_batch import Req @@ -179,19 +186,25 @@ def __init__( device: torch.device, forward_stream: Optional[torch.cuda.Stream], tp_size: int, + abort_windows: Callable[..., None], ): + self._abort_windows = abort_windows # None when the model encodes nothing ahead (no early embeddings). self._early = early self._device = device self._forward_stream = forward_stream # Vision runs on every TP rank in lockstep; it is not encoded out of step. self._encode_ahead = early is not None and tp_size == 1 - # Parked follow-ups, keyed by the rid they wait for. - self._parked: Dict[str, Req] = {} + # Follow-ups waiting for their predecessor. + self._parked = _RequestList( + ListKind.UOCR_PARKED, abort=lambda r: self._abort(self._parked, r) + ) # Completion tokens of finished predecessors whose follow-up has not arrived. self._finished: Dict[str, int] = {} # Follow-ups whose predecessor finished, for the scheduler to queue. - self._released: List[Req] = [] + self._released = _RequestList( + ListKind.UOCR_RELEASED, abort=lambda r: self._abort(self._released, r) + ) def park(self, req: Req) -> bool: """Hold a follow-up until its predecessor finishes; True if it was parked.""" @@ -207,7 +220,7 @@ def park(self, req: Req) -> bool: return False if not envs.SGLANG_ENABLE_UOCR_DOC_PREFETCH.get(): return False - self._parked[after] = req + self._parked.join([req]) mm = req.multimodal_inputs if self._encode_ahead and mm is not None and mm.mm_items: self._early.encode_ahead( @@ -220,7 +233,7 @@ def park(self, req: Req) -> bool: return True def on_window_finished(self, req: Req) -> None: - """Mark the follow-up of a finished ``req`` released, budget capped. + """Mark every follow-up of a finished ``req`` released, budget capped. Called from the KV-release path after the window's KV was retained; the scheduler queues released windows with ``take_released``. A retracted @@ -233,30 +246,30 @@ def on_window_finished(self, req: Req) -> None: self.release_early(req.rid) out = len(req.output_ids) if req.finished_len is None else req.finished_len used = max(0, out - forced_len(req)) - follow = self._parked.pop(req.rid, None) - if follow is None: + follows = [r for r in self._parked if _after(r) == req.rid] + if not follows: if _is_predecessor(req): self._finished[req.rid] = used return - _apply_budget(follow, used) - self._released.append(follow) + for follow in follows: + _apply_budget(follow, used) + move(follows, self._parked, self._released) def take_released(self) -> List[Req]: """Follow-ups released since the last call, in release order.""" - out, self._released = self._released, [] - return out - - def drop(self, rid_prefix: str) -> List[Req]: - """Remove the parked or released follow-ups an abort names.""" - hit = [k for k, r in self._parked.items() if r.rid.startswith(rid_prefix)] - out = [self._parked.pop(k) for k in hit] - kept = [] - for r in self._released: - (out if r.rid.startswith(rid_prefix) else kept).append(r) - self._released = kept - for r in out: + return self._released.take() + + def request_lists(self) -> Dict[ListKind, RequestList]: + return { + ListKind.UOCR_PARKED: self._parked, + ListKind.UOCR_RELEASED: self._released, + } + + def _abort(self, windows: RequestList, recv_req: AbortReq) -> None: + aborted = windows.take(lambda r: recv_req.matches(r.rid)) + for r in aborted: self.release_early(r.rid) - return out + self._abort_windows(aborted, finished_reason=recv_req.finished_reason) def forget_doc(self, doc_id: str) -> None: """Drop finished-predecessor records of a document that is done.""" @@ -267,10 +280,6 @@ def held_count(self) -> int: """Follow-up windows held off the waiting queue.""" return len(self._parked) + len(self._released) - def held_reqs(self) -> List[Req]: - """Follow-up windows held off the waiting queue, parked then released.""" - return [*self._parked.values(), *self._released] - def release_early(self, rid: str) -> None: """Free a window's early embeddings once it no longer needs them.""" if self._early is not None: @@ -278,9 +287,9 @@ def release_early(self, rid: str) -> None: def reset(self) -> List[Req]: """Forget everything; returns the held follow-ups for the caller to abort.""" - held = list(self._parked.values()) + self._released + held = [*self._parked, *self._released] self._parked.clear() - self._released = [] + self._released.clear() self._finished.clear() if self._early is not None: self._early.clear() diff --git a/python/sglang/srt/models/unlimited_ocr_v2_support/scheduler_runtime.py b/python/sglang/srt/models/unlimited_ocr_v2_support/scheduler_runtime.py index 71904561c66e..a696ab9c6c77 100644 --- a/python/sglang/srt/models/unlimited_ocr_v2_support/scheduler_runtime.py +++ b/python/sglang/srt/models/unlimited_ocr_v2_support/scheduler_runtime.py @@ -6,7 +6,10 @@ from typing import TYPE_CHECKING, Callable, Dict, List, Optional from sglang.srt.managers.overlap_utils import resolve_forward_inputs -from sglang.srt.managers.scheduler_components.request_list import ListKind +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + RequestList, +) from sglang.srt.models.unlimited_ocr_v2_support.async_shm import ( create_async_shm_materializer, ) @@ -75,6 +78,7 @@ def __init__( decode_graph_io: Callable[[], Optional[UOCRDecodeGraphIO]], crop_transfer, enable_overlap: bool, + abort_windows: Callable[..., None], ): self._crop_transfer = crop_transfer self.overlap_release = DeferredOverlapRelease() if enable_overlap else None @@ -83,6 +87,7 @@ def __init__( device=device, forward_stream=forward_stream, tp_size=tp_size, + abort_windows=abort_windows, ) backend = _v2_backend(attn_backend) self._reuse = DocReuse( @@ -103,8 +108,8 @@ def take_released(self) -> List[Req]: """Parked windows whose predecessor finished, to queue in this order.""" return self._prefetch.take_released() - def request_lists(self) -> Dict[ListKind, List[Req]]: - return {ListKind.UOCR_PARKED: self._prefetch.held_reqs()} + def request_lists(self) -> Dict[ListKind, RequestList]: + return self._prefetch.request_lists() def on_req_dropped(self, req: Req) -> None: """Release reuse and prefetch state of a request aborted before running.""" @@ -115,10 +120,6 @@ def free_for_decode(self, batch: ScheduleBatch) -> None: """Evict retained windows before the decode batch would retract.""" self._reuse.free_for_decode(batch) - def drop_parked(self, rid_prefix: str) -> List[Req]: - """Remove and return the held-off windows an abort names.""" - return self._prefetch.drop(rid_prefix) - def release_doc(self, doc_id: str) -> None: """Free a document's retained window once no follow-up will consume it.""" self._reuse.release_doc(doc_id) diff --git a/python/sglang/test/scheduler_test_utils.py b/python/sglang/test/scheduler_test_utils.py index 74095c260ff0..d73bdbc8dd87 100644 --- a/python/sglang/test/scheduler_test_utils.py +++ b/python/sglang/test/scheduler_test_utils.py @@ -3,7 +3,7 @@ from __future__ import annotations from types import SimpleNamespace -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Iterable from sglang.srt.disaggregation.mode import DisaggregationMode from sglang.srt.managers.scheduler_components.deferred_responses import ( @@ -26,6 +26,14 @@ class FakeReq(SimpleNamespace): __hash__ = object.__hash__ +def decode_queue(kind: ListKind, decode_reqs: Iterable = ()) -> _RequestList: + """A PD decode queue of `kind` holding `decode_reqs`, each a wrapper with + `.req`.""" + queue = _RequestList(kind, req_of=lambda x: x.req) + queue.join(decode_reqs) + return queue + + def empty_request_lists(scheduler: Scheduler) -> None: """Give a scheduler built with `Scheduler.__new__` every request list, empty; the test then fills the ones it needs.""" diff --git a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py index 46d4eae739e7..15790241fa61 100644 --- a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py +++ b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py @@ -20,9 +20,11 @@ ) from sglang.srt.disaggregation.utils import ReqToMetadataIdxAllocator from sglang.srt.managers.schedule_batch import ReqKvInfo +from sglang.srt.managers.scheduler_components.request_list import ListKind from sglang.srt.runtime_context import publish, reset_context from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.scheduler_test_utils import FakeReq, decode_queue from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=13, suite="base-a-test-cpu") @@ -126,11 +128,12 @@ def make_queue(self): ) event = SimpleNamespace(query=MagicMock(return_value=True)) entry = SimpleNamespace( + req=FakeReq(), kv_receiver=receiver, pending_metadata=(event, ("pages", 3, None), {}), is_rebootstrap=False, ) - queue.queue = [entry] + queue.queue = decode_queue(ListKind.PD_RECEIVING, [entry]) return queue, entry, event, receiver def test_waits_for_every_rank_and_sends_once(self): @@ -203,7 +206,7 @@ def failure_exception(self): def _gated_decode_req(poll, restore_status): return SimpleNamespace( - req=SimpleNamespace(rid="r", bootstrap_room=7, return_logprob=False), + req=FakeReq(rid="r", bootstrap_room=7, return_logprob=False), kv_receiver=_GateReceiver(poll), metadata_buffer_index=3, hicache_restore_status=restore_status, @@ -241,7 +244,7 @@ def test_pop_transferred_branches_on_reduced_poll_only( ): def make_queue(decode_req, peer_poll): queue = DecodeTransferQueue.__new__(DecodeTransferQueue) - queue.queue = [decode_req] + queue.queue = decode_queue(ListKind.PD_RECEIVING, [decode_req]) queue.enable_staging = False queue.enable_host_receive = False queue.enable_deferred_kv_release = False @@ -278,7 +281,7 @@ def reduce_with_peer( decode_req = _gated_decode_req(KVPoll.Success, HiCacheRestoreResult.FAILED) queue = make_queue(decode_req, peer_poll=KVPoll.Failed) self.assertEqual(queue.pop_transferred(), []) - self.assertEqual(queue.queue, []) + self.assertEqual(list(queue.queue.items()), []) queue._commit_transfer_to_req.assert_not_called() queue._record_transfer_failure.assert_called_once_with("HiCache restore failed") mock_prepare_abort.assert_called_once() @@ -298,7 +301,7 @@ def reduce_with_peer( decode_req = _gated_decode_req(KVPoll.Transferring, HiCacheRestoreResult.FAILED) queue = make_queue(decode_req, peer_poll=KVPoll.Success) self.assertEqual(queue.pop_transferred(), []) - self.assertEqual(queue.queue, [decode_req]) + self.assertEqual(list(queue.queue.items()), [decode_req]) mock_prepare_abort.assert_not_called() mock_release_kv_cache.assert_not_called() @@ -338,7 +341,7 @@ def _deferring_queue(self, decode_reqs): queue = DecodePreallocQueue.__new__(DecodePreallocQueue) queue.pp_size = 1 queue.gloo_group = object() - queue.queue = list(decode_reqs) + queue.queue = decode_queue(ListKind.PD_PREALLOC, decode_reqs) queue.pending_reqs = [] queue.retracted_queue = [] queue.num_reserved_decode_tokens = 0 @@ -369,7 +372,9 @@ def pre_alloc(req, *args): translate_kv_indices_for_transfer=lambda indices: indices, ) queue.kv_manager = SimpleNamespace(kv_args=SimpleNamespace(state_types=[])) - queue.transfer_queue = SimpleNamespace(queue=[], enable_staging=False) + queue.transfer_queue = SimpleNamespace( + queue=decode_queue(ListKind.PD_RECEIVING), enable_staging=False + ) queue.tree_cache = MagicMock() queue.scheduler = MagicMock( enable_priority_scheduling=False, @@ -382,7 +387,7 @@ def pre_alloc(req, *args): @staticmethod def _decode_req(receiver): - req = SimpleNamespace( + req = FakeReq( rid="deferred", origin_input_ids=[1] * 5, output_ids=[], @@ -490,7 +495,7 @@ def test_recompute_waits_for_held_metadata_to_be_sent(self): queue.pop_preallocated() queue._submit_prefill_recompute.assert_not_called() transfer = DecodeTransferQueue.__new__(DecodeTransferQueue) - transfer.queue = [decode_req] + transfer.queue = decode_queue(ListKind.PD_RECEIVING, [decode_req]) transfer.gloo_group = object() transfer.scheduler = queue.scheduler with patch( diff --git a/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py b/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py index 7ea29e5767ab..2c9c802751de 100644 --- a/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py +++ b/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py @@ -34,7 +34,12 @@ ) from sglang.srt.disaggregation.decode import DecodeTransferQueue from sglang.srt.disaggregation.decode_hicache_mixin import HiCacheRestoreResult +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.scheduler_test_utils import FakeReq from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=12, suite="base-a-test-cpu") @@ -482,8 +487,10 @@ def free(self, idx): def _make_queue(timeout=30.0, enable_metrics=False): q = DecodeTransferQueue.__new__(DecodeTransferQueue) - q._deferred_releases = [] - q.queue = [] + q._deferred_releases = _RequestList( + ListKind.PD_RELEASING, req_of=lambda x: x[0].req + ) + q.queue = _RequestList(ListKind.PD_RECEIVING, req_of=lambda x: x.req) q.gloo_group = dist.group.WORLD q.tp_rank = 0 q.enable_deferred_kv_release = True @@ -523,7 +530,7 @@ def _make_decode_req(room, idx, mgr, n_prefill_ranks=1): clear=lambda: None, ) return SimpleNamespace( - req=SimpleNamespace(bootstrap_room=room, rid=str(room), return_logprob=False), + req=FakeReq(bootstrap_room=room, rid=str(room), return_logprob=False), host_staged=False, hicache_restore_status=HiCacheRestoreResult.READY, prefix_match=None, @@ -556,11 +563,11 @@ def test_unpublished_rank_holds_for_peer_without_waiting_for_acks(self): # ABORT was sent before metadata, so this rank has no ack generation. # A peer that already published still makes every rank retain capacity. queue._defer_release(entry) - self.assertEqual(queue._deferred_releases[0][-1], 0) + self.assertEqual(list(queue._deferred_releases.items())[0][-1], 0) with patch.object(decode_mod, "release_kv_cache") as release: queue.resolve_deferred_releases() release.assert_called_once_with(entry.req, queue.tree_cache, checkpoint=False) - self.assertEqual(queue._deferred_releases, []) + self.assertEqual(list(queue._deferred_releases.items()), []) self.assertEqual(queue.req_to_metadata_buffer_idx_allocator.freed, [7]) def test_host_release_requires_drain_on_every_rank_even_after_timeout(self): @@ -594,7 +601,7 @@ def test_host_release_requires_drain_on_every_rank_even_after_timeout(self): mgr.note_abort_ack(1, 0, generations[1]) q.resolve_deferred_releases() self.assertEqual(discard.call_count, 2) - self.assertEqual(q._deferred_releases, []) + self.assertEqual(list(q._deferred_releases.items()), []) device_release.assert_not_called() def test_holds_until_drained_then_releases(self): @@ -629,7 +636,7 @@ def test_holds_until_drained_then_releases(self): rel.assert_called_once_with(dreq.req, q.tree_cache, checkpoint=False) # Held state fully cleaned up. - self.assertEqual(q._deferred_releases, []) + self.assertEqual(list(q._deferred_releases.items()), []) self.assertEqual(q.num_pending_deferred_releases(), 0) self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [idx]) self.assertEqual(q.metadata_buffers.bootstrap_room[idx], 0) @@ -646,7 +653,7 @@ def test_releases_on_timeout_without_ack(self): q = _make_queue(timeout=30.0, enable_metrics=True) dreq = _make_decode_req(room, idx, mgr, n_prefill_ranks=1) # Force an already-expired deadline (no ack will ever arrive). - q._deferred_releases.append((dreq, 10.0, float("-inf"), idx, 1)) + q._deferred_releases.join([(dreq, 10.0, float("-inf"), idx, 1)]) with ( patch.object(decode_mod.time, "monotonic", return_value=42.0), @@ -655,7 +662,7 @@ def test_releases_on_timeout_without_ack(self): q.resolve_deferred_releases() rel.assert_called_once_with(dreq.req, q.tree_cache, checkpoint=False) - self.assertEqual(q._deferred_releases, []) + self.assertEqual(list(q._deferred_releases.items()), []) self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [idx]) self.assertIsNone(dreq.kv_receiver) q.scheduler.metrics_collector.observe_decode_deferred_kv_release.assert_called_once_with( @@ -671,8 +678,8 @@ def test_failed_release_is_isolated_and_not_retried(self): good = _make_decode_req(700, 1, mgr) bad = _make_decode_req(701, 2, mgr) # Both already past deadline -> both selected for release. - q._deferred_releases.append((bad, 10.0, float("-inf"), 2, 1)) - q._deferred_releases.append((good, 10.0, float("-inf"), 1, 1)) + q._deferred_releases.join([(bad, 10.0, float("-inf"), 2, 1)]) + q._deferred_releases.join([(good, 10.0, float("-inf"), 1, 1)]) calls = [] @@ -689,7 +696,7 @@ def fake_release(req, tree_cache, checkpoint): # The good one still released despite the bad one throwing. self.assertIn(good.req, calls) # Nothing left held, and a second call is a clean no-op (no retry). - self.assertEqual(q._deferred_releases, []) + self.assertEqual(list(q._deferred_releases.items()), []) q.resolve_deferred_releases() q.scheduler.metrics_collector.observe_decode_deferred_kv_release.assert_has_calls( [ @@ -705,7 +712,9 @@ def test_defer_release_records_deadline_and_idx(self): with patch.object(decode_mod.time, "monotonic", return_value=5.0): q._defer_release(dreq) self.assertEqual(len(q._deferred_releases), 1) - held_req, start_time, deadline, held_idx, required = q._deferred_releases[0] + held_req, start_time, deadline, held_idx, required = list( + q._deferred_releases.items() + )[0] self.assertIs(held_req, dreq) self.assertEqual(start_time, 5.0) self.assertEqual(deadline, 17.5) @@ -798,7 +807,7 @@ def _run_release_consensus(rank: int, rendezvous: str) -> None: if not receiver.abort_notified: receiver.bootstrap_infos = None generation = mgr.register_deferred_abort_room(43) - q.queue = [dreq] + q.queue.join([dreq]) with ( patch.object(decode_mod, "prepare_abort"), patch.object(decode_mod, "release_kv_cache"), @@ -866,7 +875,7 @@ def ensure(force_arm=False): receiver.force_arm_calls = [] receiver.ensure_abort_notified = ensure return SimpleNamespace( - req=SimpleNamespace(bootstrap_room=room, rid=f"r{room}", return_logprob=False), + req=FakeReq(bootstrap_room=room, rid=f"r{room}", return_logprob=False), kv_receiver=receiver, metadata_buffer_index=idx, hicache_restore_status=None, @@ -900,7 +909,7 @@ def test_non_abort_failure_notifies_and_defers(self): mgr.enable_deferred_decode_kv_release = True q = _make_pop_queue(mgr) entry = _make_failed_req(mgr, room=7, idx=3, abort_notified=False) - q.queue = [entry] + q.queue.join([entry]) released = self._pop(q) @@ -912,7 +921,7 @@ def test_non_abort_failure_notifies_and_defers(self): self.assertEqual(len(q._deferred_releases), 1) # Metadata slot stays owned by the hold until resolve time. self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, []) - self.assertEqual(q.queue, []) + self.assertEqual(list(q.queue.items()), []) def test_unnotifiable_failure_releases_immediately(self): # Metadata never published: no prefill holds this destination, so no @@ -923,12 +932,12 @@ def test_unnotifiable_failure_releases_immediately(self): entry = _make_failed_req( mgr, room=8, idx=4, abort_notified=False, notifiable=False ) - q.queue = [entry] + q.queue.join([entry]) released = self._pop(q) self.assertEqual(released, [entry]) - self.assertEqual(q._deferred_releases, []) + self.assertEqual(list(q._deferred_releases.items()), []) self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [4]) def test_failure_before_deferred_publish_releases_immediately(self): @@ -939,7 +948,7 @@ def test_failure_before_deferred_publish_releases_immediately(self): q = _make_pop_queue(mgr) entry = _make_failed_req(mgr, room=10, idx=6, abort_notified=False) entry.awaiting_publish = True - q.queue = [entry] + q.queue.join([entry]) # Under prefill_complete, clear() of a failed receiver that decode did not # abort notifies prefill and arms drain tracking no ACK will resolve. entry.kv_receiver.clear = lambda: mgr.register_deferred_abort_room(10) @@ -947,7 +956,7 @@ def test_failure_before_deferred_publish_releases_immediately(self): released = self._pop(q) self.assertEqual(released, [entry]) - self.assertEqual(q._deferred_releases, []) + self.assertEqual(list(q._deferred_releases.items()), []) self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [6]) self.assertNotIn(10, mgr._deferred_abort_ack_tracker) @@ -968,7 +977,7 @@ def test_held_destination_is_sent_only_by_queued_requests(self): for entry in (held, sent, queued): entry.kv_receiver.send_metadata = MagicMock() held_receiver = held.kv_receiver # the release path nulls entry.kv_receiver - q.queue = [held, sent, queued] + q.queue.join([held, sent, queued]) q._poll_with_metadata_gate = lambda: [ KVPoll.Failed, KVPoll.Failed, @@ -980,12 +989,12 @@ def test_held_destination_is_sent_only_by_queued_requests(self): held_receiver.send_metadata.assert_not_called() queued.kv_receiver.send_metadata.assert_called_once_with("queued") self.assertIsNone(queued.pending_metadata) - self.assertEqual(q.queue, [queued]) + self.assertEqual(list(q.queue.items()), [queued]) self.assertEqual(released, [held]) self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [3]) required_acks = { entry.req.bootstrap_room: acks - for entry, _, _, _, acks in q._deferred_releases + for entry, _, _, _, acks in q._deferred_releases.items() } self.assertEqual(required_acks, {8: 2}) @@ -1008,7 +1017,7 @@ def test_failed_held_destination_respects_delayed_conclusion(self): entry.kv_receiver.send_metadata = MagicMock() aborted_receiver = aborted.kv_receiver pending = aborted.pending_metadata - q.queue = [concluded, aborted, queued] + q.queue.join([concluded, aborted, queued]) q._poll_with_metadata_gate = lambda: [ KVPoll.Failed, KVPoll.Failed, @@ -1018,7 +1027,7 @@ def test_failed_held_destination_respects_delayed_conclusion(self): with patch.object(decode_mod, "prepare_abort"): q.pop_transferred_with_outcomes({concluded.req.rid: True}) - self.assertEqual(q.queue, [aborted, queued]) + self.assertEqual(list(q.queue.items()), [aborted, queued]) self.assertEqual(released, [concluded]) aborted_receiver.send_metadata.assert_not_called() self.assertIs(aborted.pending_metadata, pending) @@ -1032,9 +1041,9 @@ def test_failed_held_destination_respects_delayed_conclusion(self): q.pop_transferred_with_outcomes({aborted.req.rid: True}) self.assertEqual(released, [concluded, aborted]) - self.assertEqual(q.queue, [queued]) + self.assertEqual(list(q.queue.items()), [queued]) self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [3, 4]) - self.assertEqual(q._deferred_releases, []) + self.assertEqual(list(q._deferred_releases.items()), []) self.assertEqual(mgr._deferred_abort_ack_tracker, {}) aborted_receiver.send_metadata.assert_not_called() queued.kv_receiver.send_metadata.assert_called_once_with("queued") @@ -1045,13 +1054,13 @@ def test_backend_optout_keeps_immediate_release_and_sends_no_abort(self): q = _make_pop_queue(mgr) entry = _make_failed_req(mgr, room=9, idx=5, abort_notified=False) receiver = entry.kv_receiver # the release path nulls entry.kv_receiver - q.queue = [entry] + q.queue.join([entry]) released = self._pop(q) self.assertFalse(receiver.abort_notified) self.assertEqual(released, [entry]) - self.assertEqual(q._deferred_releases, []) + self.assertEqual(list(q._deferred_releases.items()), []) class TestHiSparseDestinationFollowsTheHold(CustomTestCase): @@ -1089,7 +1098,7 @@ def test_deferred_failure_keeps_the_destination_until_release(self): mgr = _make_manager() log = [] q = self._queue(mgr, log) - q.queue = [_make_failed_req(mgr, room=12, idx=6, abort_notified=False)] + q.queue.join([_make_failed_req(mgr, room=12, idx=6, abort_notified=False)]) at_failure = self._pop_and_resolve(q, mgr, 12, log) @@ -1100,11 +1109,13 @@ def test_immediate_failure_frees_the_destination_once(self): mgr = _make_manager() log = [] q = self._queue(mgr, log) - q.queue = [ - _make_failed_req( - mgr, room=13, idx=7, abort_notified=False, notifiable=False - ) - ] + q.queue.join( + [ + _make_failed_req( + mgr, room=13, idx=7, abort_notified=False, notifiable=False + ) + ] + ) at_failure = self._pop_and_resolve(q, mgr, 13, log) diff --git a/test/registered/unit/disaggregation/test_prefill_complete_integration.py b/test/registered/unit/disaggregation/test_prefill_complete_integration.py index 0865e8159453..980dbfac05f7 100644 --- a/test/registered/unit/disaggregation/test_prefill_complete_integration.py +++ b/test/registered/unit/disaggregation/test_prefill_complete_integration.py @@ -38,7 +38,7 @@ from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.scheduler_test_utils import FakeReq +from sglang.test.scheduler_test_utils import FakeReq, decode_queue from sglang.test.test_utils import CustomTestCase, enter_scope register_cpu_ci(est_time=11, suite="base-a-test-cpu") @@ -279,7 +279,7 @@ def test_mismatched_peer_fails_request_before_destination_registration(self): def test_scheduler_does_not_reserve_request_rows_for_unready_prefill(self): sender, receiver = self.sender(), self.receiver() - req = SimpleNamespace( + req = FakeReq( rid="req", bootstrap_room=1, finished_reason=None, @@ -288,7 +288,8 @@ def test_scheduler_does_not_reserve_request_rows_for_unready_prefill(self): ) entry = DecodeRequest(req=req, kv_receiver=receiver) queue = DecodePreallocQueue.__new__(DecodePreallocQueue) - queue.queue, queue.pending_reqs = [entry], [] + queue.queue = decode_queue(ListKind.PD_PREALLOC, [entry]) + queue.pending_reqs = [] queue.pp_size, queue.tp_rank = 1, 0 queue.gloo_group = object() queue._resolve_pending_reqs = Mock() @@ -324,7 +325,7 @@ def test_scheduler_does_not_reserve_request_rows_for_unready_prefill(self): self.assertTrue(entry.waiting_for_input) queue.req_to_metadata_buffer_idx_allocator.available_size.assert_called_once() queue.req_to_token_pool.available_size.assert_called_once() - self.assertEqual(queue.queue, [entry]) + self.assertEqual(list(queue.queue.items()), [entry]) def test_exhausted_optimistic_attempt_follows_allocation_policy(self): override = get_context().override_server_args( diff --git a/test/registered/unit/disaggregation/test_unified_memory_move_gate.py b/test/registered/unit/disaggregation/test_unified_memory_move_gate.py index 2f6f4dd87645..fc71fbd7b0e8 100644 --- a/test/registered/unit/disaggregation/test_unified_memory_move_gate.py +++ b/test/registered/unit/disaggregation/test_unified_memory_move_gate.py @@ -20,10 +20,12 @@ from sglang.srt.disaggregation.mode import DisaggregationMode from sglang.srt.disaggregation.utils import disagg_transfers_in_flight +from sglang.srt.managers.scheduler_components.request_list import ListKind from sglang.srt.mem_cache.kv.alloc.band import Band from sglang.srt.mem_cache.kv.alloc.chain import ChainAllocator from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.scheduler_test_utils import FakeReq, decode_queue from sglang.test.test_utils import CustomTestCase from sglang.test.unified_allocator_fixtures import ( UnifiedGeometry, @@ -36,7 +38,7 @@ class _FakeTransferQueue: def __init__(self): - self.queue: List[object] = [] + self.queue = decode_queue(ListKind.PD_RECEIVING) class _FakePreallocQueue: @@ -87,8 +89,8 @@ def test_closed_while_destination_published_but_not_queued(self): self.assertTrue(busy()) # Handing the batch to the transfer queue transfers responsibility. - scheduler.disagg_decode_transfer_queue.queue.append( - SimpleNamespace(host_staged=False) + scheduler.disagg_decode_transfer_queue.queue.join( + [SimpleNamespace(req=FakeReq(), host_staged=False)] ) scheduler.disagg_decode_prealloc_queue.note_destinations_queued(1) self.assertTrue(busy(), "transfer queue still holds it") @@ -316,8 +318,8 @@ def check_pins(blocked): scheduler.disagg_decode_transfer_queue = _FakeTransferQueue() scheduler.disagg_decode_prealloc_queue = _FakePreallocQueue() check_pins(False) - scheduler.disagg_decode_transfer_queue.queue.append( - SimpleNamespace(host_staged=False) + scheduler.disagg_decode_transfer_queue.queue.join( + [SimpleNamespace(req=FakeReq(), host_staged=False)] ) check_pins(True) scheduler.disagg_decode_transfer_queue.queue.clear() diff --git a/test/registered/unit/managers/scheduler_components/test_request_list.py b/test/registered/unit/managers/scheduler_components/test_request_list.py index ad2e98977667..28e1f65a6046 100644 --- a/test/registered/unit/managers/scheduler_components/test_request_list.py +++ b/test/registered/unit/managers/scheduler_components/test_request_list.py @@ -18,6 +18,7 @@ from sglang.srt.managers.scheduler_components.deferred_responses import ( DeferredResponses, ) +from sglang.srt.models.unlimited_ocr_v2_support.document_prefetch import DocPrefetch from sglang.srt.models.unlimited_ocr_v2_support.scheduler_runtime import ( UocrSchedulerRuntime, ) @@ -34,6 +35,7 @@ DllmManager, HiSparseCoordinator, UocrSchedulerRuntime, + DocPrefetch, DeferredResponses, ) @@ -41,14 +43,10 @@ # shrinks: a list that migrates leaves it. _NOT_YET_LISTS = { "Scheduler.waiting_queue", - "DecodePreallocQueue.queue", # Shares its entries with `queue`. "DecodePreallocQueue.pending_reqs", - "DecodeTransferQueue.queue", - "DecodeTransferQueue._deferred_releases", # A subset of `waiting_queue`. "DllmManager.staging_queue", - "DeferredResponses.held", } _HOLDS_REQUESTS = re.compile(r"\b(Req|DecodeRequest)\b") diff --git a/test/registered/unit/managers/test_hisparse_unit.py b/test/registered/unit/managers/test_hisparse_unit.py index c7eafbcef0f8..c9a026a420c2 100644 --- a/test/registered/unit/managers/test_hisparse_unit.py +++ b/test/registered/unit/managers/test_hisparse_unit.py @@ -30,6 +30,7 @@ register_cuda_ci, register_xpu_ci, ) +from sglang.test.scheduler_test_utils import FakeReq register_cuda_ci(est_time=13, stage="base-b", runner_config="1-gpu-small") register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd") @@ -60,7 +61,7 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None): origin_input_ids = list(range(64)) if output_ids is None: output_ids = [] - req = SimpleNamespace( + req = FakeReq( rid=rid, origin_input_ids=origin_input_ids, output_ids=output_ids, diff --git a/test/registered/unit/managers/test_load_inquirer.py b/test/registered/unit/managers/test_load_inquirer.py index 8324003c0d9b..a77980b99153 100644 --- a/test/registered/unit/managers/test_load_inquirer.py +++ b/test/registered/unit/managers/test_load_inquirer.py @@ -1,8 +1,9 @@ import unittest from dataclasses import replace from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import Mock, patch +from sglang.srt.disaggregation.decode import DecodeRequest from sglang.srt.disaggregation.mode import DisaggregationMode from sglang.srt.managers.schedule_policy import ( CacheAwarePolicy, @@ -11,8 +12,14 @@ from sglang.srt.managers.scheduler_components.load_inquirer import ( SchedulerLoadInquirer, ) +from sglang.srt.managers.scheduler_components.request_list import ( + ListKind, + _RequestList, +) from sglang.srt.mem_cache.base_prefix_cache import CacheRequestHandle +from sglang.srt.observability.metrics_collector import SchedulerStats from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.scheduler_test_utils import FakeReq, decode_queue register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -109,6 +116,49 @@ def test_storage_exclusion_is_clipped_to_uncached_prompt(self): CacheRequestHandle("waiting", 0) ) + def test_decode_loads_count_prealloc_requests_waiting_for_input(self): + """A preallocated request whose input has arrived is ready, with its tokens.""" + ready, pending = FakeReq(seqlen=30), FakeReq(seqlen=70) + prealloc = decode_queue( + ListKind.PD_PREALLOC, + [ + DecodeRequest(req=ready, kv_receiver=None, waiting_for_input=True), + DecodeRequest(req=pending, kv_receiver=None), + ], + ) + inquirer = replace( + self.make_inquirer(disaggregation_mode=DisaggregationMode.DECODE), + pool_stats_observer=SimpleNamespace( + get_pool_stats=lambda: SimpleNamespace( + get_kv_token_stats=lambda: (0, 0.0) + ) + ), + spec_algorithm=SimpleNamespace(is_none=lambda: True), + get_running_batch=lambda: SimpleNamespace(reqs=[]), + get_stats=SchedulerStats, + get_disagg_decode_prealloc_queue=lambda: SimpleNamespace( + queue=prealloc, retracted_queue=_RequestList(ListKind.PD_RETRACTED) + ), + get_disagg_decode_transfer_queue=lambda: SimpleNamespace( + queue=decode_queue(ListKind.PD_RECEIVING) + ), + get_total_prefill_uncached_tokens=lambda: 0, + get_total_prefill_busy_us=lambda: 0, + get_decode_moment_totals=lambda: (0,), + ) + module = "sglang.srt.managers.scheduler_components.load_inquirer" + with ( + patch( + f"{module}.get_lora", return_value=SimpleNamespace(enable_lora=False) + ), + patch(f"{module}.get_parallel", return_value=SimpleNamespace(dp_rank=0)), + ): + loads = inquirer.get_loads() + + self.assertEqual(loads.disaggregation.decode_prealloc_queue_reqs, 2) + self.assertEqual(loads.queues.prealloc_ready, 1) + self.assertEqual(loads.num_prealloc_ready_tokens, 30) + def test_storage_exclusion_requires_per_request_matching(self): for prefix_matched, expected in ((True, 60), (False, 65)): with self.subTest(prefix_matched=prefix_matched): diff --git a/test/registered/unit/managers/test_overlap_uocr_decode.py b/test/registered/unit/managers/test_overlap_uocr_decode.py index 13f380d02522..6ea7bc772c51 100644 --- a/test/registered/unit/managers/test_overlap_uocr_decode.py +++ b/test/registered/unit/managers/test_overlap_uocr_decode.py @@ -19,6 +19,7 @@ ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.scheduler_test_utils import FakeReq from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=5, suite="base-a-test-cpu") @@ -67,6 +68,7 @@ def _hooks(current_io): decode_graph_io=current_io, crop_transfer=Mock(), enable_overlap=False, + abort_windows=Mock(), ) @@ -300,7 +302,7 @@ def test_held_or_matched_windows_block_a_flush(self): hooks._reuse.sources[1] = object() self.assertTrue(hooks.blocks_flush(ignore_waiting=True)) hooks._reuse.sources.clear() - hooks._prefetch._parked["w0"] = SimpleNamespace(rid="w1") + hooks._prefetch._parked.join([FakeReq(rid="w1")]) self.assertTrue(hooks.blocks_flush(ignore_waiting=False)) self.assertFalse(hooks.blocks_flush(ignore_waiting=True)) @@ -308,8 +310,8 @@ def test_flush_drops_document_state_before_the_pools_and_aborts_held_windows(sel """Held windows must get an answer, and the retained windows must be forgotten before the clear returns their KV to the pools.""" hooks, _ = self._hooks() - held = SimpleNamespace(rid="w1") - hooks._prefetch._parked["w0"] = held + held = FakeReq(rid="w1") + hooks._prefetch._parked.join([held]) hooks._reuse.docs["doc"] = object() seen = [] scheduler = object.__new__(Scheduler) diff --git a/test/registered/unit/managers/test_pp_decode_transfer_outcome.py b/test/registered/unit/managers/test_pp_decode_transfer_outcome.py index 4421e991c1ab..4bbfe7d8baa9 100644 --- a/test/registered/unit/managers/test_pp_decode_transfer_outcome.py +++ b/test/registered/unit/managers/test_pp_decode_transfer_outcome.py @@ -14,8 +14,10 @@ from sglang.srt.disaggregation.base import KVPoll from sglang.srt.disaggregation.common.conn import KVTransferError from sglang.srt.disaggregation.decode import DecodeTransferQueue +from sglang.srt.managers.scheduler_components.request_list import ListKind from sglang.srt.managers.scheduler_pp_mixin import SchedulerPPMixin from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.scheduler_test_utils import FakeReq, decode_queue from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -25,7 +27,9 @@ def _stage(first_rank: bool, succeeded, failed, prev=None): """A PP stage whose own receivers report `succeeded` / `failed` rids.""" return SimpleNamespace( pp_group=SimpleNamespace(is_first_rank=first_rank), - disagg_decode_transfer_queue=SimpleNamespace(queue=[]), + disagg_decode_transfer_queue=SimpleNamespace( + queue=decode_queue(ListKind.PD_RECEIVING) + ), get_rids=lambda queue, is_send, *groups: (list(succeeded), list(failed)), _pp_recv_pyobj_from_prev_stage=lambda: prev, ) @@ -83,7 +87,7 @@ def failure_exception(self): def _decode_req(rid, poll): return SimpleNamespace( - req=SimpleNamespace( + req=FakeReq( rid=rid, bootstrap_room=7, return_logprob=False, finished_reason=None ), kv_receiver=_Receiver(poll), @@ -96,7 +100,7 @@ def _decode_req(rid, poll): def _transfer_queue(decode_reqs): queue = DecodeTransferQueue.__new__(DecodeTransferQueue) - queue.queue = list(decode_reqs) + queue.queue = decode_queue(ListKind.PD_RECEIVING, decode_reqs) queue.enable_staging = False queue.enable_host_receive = False queue.enable_deferred_kv_release = False @@ -149,7 +153,7 @@ def test_consensus_failure_overrides_a_local_success(self): self.assertEqual(reqs, [admitted.req]) self.assertEqual(outcomes, {"s": False, "f": True}) - self.assertEqual(queue.queue, [not_yet_agreed]) + self.assertEqual(list(queue.queue.items()), [not_yet_agreed]) queue._commit_transfer_to_req.assert_called_once_with(admitted) queue._release_request.assert_called_once_with(failed_elsewhere) queue._record_transfer_failure.assert_called_once_with( @@ -164,7 +168,7 @@ def test_local_failure_is_reported_downstream(self): self.assertEqual(reqs, []) self.assertEqual(outcomes, {"f": True}) - self.assertEqual(queue.queue, []) + self.assertEqual(list(queue.queue.items()), []) def test_non_pp_path_keeps_the_request_list(self): admitted = _decode_req("s", KVPoll.Success) diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index 8b3084ab6f1e..161792e9f0ff 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -36,7 +36,11 @@ from sglang.srt.runtime_context import get_context, publish, reset_context # noqa: E402 from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.scheduler_test_utils import FakeReq, empty_request_lists +from sglang.test.scheduler_test_utils import ( + FakeReq, + decode_queue, + empty_request_lists, +) from sglang.test.separate_buffer_allocator_double import ( bind_separate_buffer_capacity, ) @@ -468,7 +472,7 @@ def setUp(self): publish(ServerArgs(model_path="dummy"), role="tokenizer") def _new_decode_req(self, rid: str, priority: int, *, failed: bool = False): - req = SimpleNamespace( + req = FakeReq( rid=rid, priority=priority, origin_input_ids=[1, 2, 3], @@ -491,7 +495,7 @@ def _new_decode_req(self, rid: str, priority: int, *, failed: bool = False): def _new_queue(self, decode_reqs, *, low_priority_values_first: bool = False): queue = DecodePreallocQueue.__new__(DecodePreallocQueue) queue.pp_size = 1 - queue.queue = list(decode_reqs) + queue.queue = decode_queue(ListKind.PD_PREALLOC, decode_reqs) queue.pending_reqs = [] queue.retracted_queue = [] queue.num_reserved_decode_tokens = 0 @@ -529,7 +533,9 @@ def pre_alloc_mock(req, prefix_indices=None, prefix_len=0, total_prefix_len=0): queue.token_to_kv_pool_allocator.page_size = 1 queue.token_to_kv_pool_allocator.available_size.return_value = 1000 queue.token_to_kv_pool = MagicMock() - queue.transfer_queue = SimpleNamespace(queue=[], enable_staging=False) + queue.transfer_queue = SimpleNamespace( + queue=decode_queue(ListKind.PD_RECEIVING), enable_staging=False + ) queue.kv_manager = SimpleNamespace( kv_args=SimpleNamespace(state_types=[]), ) @@ -575,9 +581,13 @@ def test_hisparse_prealloc_reserves_scratch_for_pending_transfers(self): queue.token_to_kv_pool_allocator.page_size = 64 queue.token_to_kv_pool.page_size = 64 queue.draft_token_to_kv_pool = None - queue.transfer_queue.queue = [ - SimpleNamespace(host_staged=False) for _ in range(pending) - ] + queue.transfer_queue.queue = decode_queue( + ListKind.PD_RECEIVING, + [ + SimpleNamespace(req=FakeReq(), host_staged=False) + for _ in range(pending) + ], + ) queue._uses_swa_tail_prealloc = MagicMock(return_value=False) queue._uses_swa_reservation = MagicMock(return_value=False) queue._send_kv_metadata = MagicMock() @@ -586,7 +596,7 @@ def test_hisparse_prealloc_reserves_scratch_for_pending_transfers(self): self.assertEqual(failed, []) self.assertEqual(preallocated, reqs[:admitted]) - self.assertEqual(queue.queue, reqs[admitted:]) + self.assertEqual(list(queue.queue.items()), reqs[admitted:]) def test_prealloc_lora_slots_cover_inflight_microbatches_and_queues(self): """In-flight requests retain their adapter slots.""" @@ -606,7 +616,7 @@ def test_prealloc_lora_slots_cover_inflight_microbatches_and_queues(self): queue = self._new_queue(reqs) queue.pp_size = 2 - queue.transfer_queue.queue = [transferring] + queue.transfer_queue.queue = decode_queue(ListKind.PD_RECEIVING, [transferring]) scheduler = queue.scheduler scheduler.running_batch.reqs = [finished.req] scheduler.running_mbs = [ @@ -632,7 +642,7 @@ def test_prealloc_lora_slots_cover_inflight_microbatches_and_queues(self): self.assertEqual(preallocated, [reqs[0], reqs[2]]) self.assertEqual(failed, []) - self.assertEqual(queue.queue, [reqs[1]]) + self.assertEqual(list(queue.queue.items()), [reqs[1]]) self.assertEqual( [call.args[0] for call in queue._pre_alloc.call_args_list], [reqs[0].req, reqs[2].req], @@ -692,7 +702,7 @@ def test_failed_request_indices_stay_valid_after_priority_sort(self): [decode_req.req.rid for decode_req in preallocated], ["healthy-high"] ) self.assertEqual([decode_req.req.rid for decode_req in failed], ["failed-low"]) - self.assertEqual(queue.queue, []) + self.assertEqual(list(queue.queue.items()), []) queue.scheduler.output_streamer.stream_output.assert_called_once_with( [failed_low.req], failed_low.req.return_logprob ) @@ -718,7 +728,9 @@ def test_admission_without_radix_cache_charges_the_alloc_hook(self): fill_len=3, prefix_len=0 ) self.assertEqual(preallocated, [decode_req] if fits else []) - self.assertEqual(queue.queue, [] if fits else [decode_req]) + self.assertEqual( + list(queue.queue.items()), [] if fits else [decode_req] + ) self.assertEqual(failed, []) diff --git a/test/registered/unit/managers/test_scheduler_weight_version_tracking.py b/test/registered/unit/managers/test_scheduler_weight_version_tracking.py index 0d5ca0720d66..1d36ebcd1221 100644 --- a/test/registered/unit/managers/test_scheduler_weight_version_tracking.py +++ b/test/registered/unit/managers/test_scheduler_weight_version_tracking.py @@ -78,11 +78,7 @@ def _scheduler( waiting_queue=list(waiting), chunked_req=chunked, hisparse_coordinator=( - SimpleNamespace( - ack_staging_queue=[SimpleNamespace(req=req) for req in staging] - ) - if staging - else None + SimpleNamespace(ack_staging_queue=list(staging)) if staging else None ), ) diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index a3a2b1893897..b3942fc42f76 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -43,6 +43,7 @@ DecodePrefixMatch, ) from sglang.srt.managers.schedule_batch import Req, ReqKvInfo +from sglang.srt.managers.scheduler_components.request_list import ListKind from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, InsertParams, @@ -53,6 +54,7 @@ from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey from sglang.srt.mem_cache.unified_cache.component_type import ComponentType from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.test.scheduler_test_utils import decode_queue from sglang.test.separate_buffer_allocator_double import ( bind_separate_buffer_capacity, ) @@ -422,7 +424,7 @@ def test_pop_preallocated_rechecks_budget_after_lock(self): # `not decode_req.is_rebootstrap` gate in pop_preallocated). decode_req.is_rebootstrap = False - queue.queue = [decode_req] + queue.queue = decode_queue(ListKind.PD_PREALLOC, [decode_req]) queue.pending_reqs = [] queue.retracted_queue = [] queue.num_reserved_decode_tokens = 0 diff --git a/test/registered/unit/mem_cache/test_decode_retraction_backup.py b/test/registered/unit/mem_cache/test_decode_retraction_backup.py index 65cee20f13c7..b5ab46ca2fef 100644 --- a/test/registered/unit/mem_cache/test_decode_retraction_backup.py +++ b/test/registered/unit/mem_cache/test_decode_retraction_backup.py @@ -52,7 +52,7 @@ HiCacheDraftPlan, ) from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.scheduler_test_utils import FakeReq +from sglang.test.scheduler_test_utils import FakeReq, decode_queue from sglang.test.test_utils import CustomTestCase from sglang.test.tree_slots import tree_slots_of @@ -373,8 +373,10 @@ def _receive_queue(self, env): transfer_backend=TransferBackend.FAKE, is_mla_backend=False, enable_staging=False, - transfer_queue=SimpleNamespace(queue=[], enable_staging=False), - queue=[], + transfer_queue=SimpleNamespace( + queue=decode_queue(ListKind.PD_RECEIVING), enable_staging=False + ), + queue=decode_queue(ListKind.PD_PREALLOC), pending_reqs=[], retracted_queue=_RequestList(ListKind.PD_RETRACTED), _prefill_dp_rank_queries={}, @@ -555,7 +557,7 @@ def test_host_receive_threshold_controls_device_allocation(self, _world_size): receiver = Mock(supports_host_destination=True) receiver.poll.return_value = KVPoll.WaitingForInput decode_req = DecodeRequest(req=req, kv_receiver=receiver) - queue.queue = [decode_req] + queue.queue = decode_queue(ListKind.PD_PREALLOC, [decode_req]) self.assertEqual(queue.pop_preallocated(), ([decode_req], [])) self.assertEqual(decode_req.host_staged, host_staged) self.assertEqual(len(host_admissions), int(host_staged)) @@ -605,7 +607,7 @@ def test_host_receive_restores_target_and_draft_kv(self, _world_size): receiver = Mock(supports_host_destination=False) receiver.poll.return_value = KVPoll.WaitingForInput decode_req = DecodeRequest(req=req, kv_receiver=receiver) - queue.queue = [decode_req] + queue.queue = decode_queue(ListKind.PD_PREALLOC, [decode_req]) host_free_before = cache.host_pool_group.available_size() blocker = env.allocator.alloc(self.pool_size) self.assertEqual(queue.pop_preallocated(), ([], [])) @@ -675,9 +677,9 @@ def test_host_receive_restores_target_and_draft_kv(self, _world_size): sampling_params=SamplingParams(max_new_tokens=1), ) pending = DecodeRequest(req=pending_req, kv_receiver=receiver) - queue.queue = [pending] + queue.queue = decode_queue(ListKind.PD_PREALLOC, [pending]) self.assertEqual(queue.pop_preallocated(), ([], [])) - self.assertEqual(queue.queue, [pending]) + self.assertEqual(list(queue.queue.items()), [pending]) self.assertIsNone(pending_req.kv.req_pool_idx) self.assertIsNone(pending_req.kv.retraction_backup) self.assertEqual( @@ -828,7 +830,7 @@ def make_req(rid, num_tokens): receiver = Mock(supports_host_destination=True) receiver.poll.return_value = KVPoll.WaitingForInput decode_req = DecodeRequest(req=req, kv_receiver=receiver) - queue.queue = [decode_req] + queue.queue = decode_queue(ListKind.PD_PREALLOC, [decode_req]) host_free_before = cache.host_pool_group.available_size() self.assertEqual(queue.pop_preallocated(), ([decode_req], [])) self.assertIsNone(req.kv.req_pool_idx) diff --git a/test/registered/unit/mem_cache/test_mamba_capacity.py b/test/registered/unit/mem_cache/test_mamba_capacity.py index 290ef99737ba..da0d4e6450fc 100644 --- a/test/registered/unit/mem_cache/test_mamba_capacity.py +++ b/test/registered/unit/mem_cache/test_mamba_capacity.py @@ -18,7 +18,7 @@ from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool from sglang.srt.mem_cache.static_kv_storage import StaticKVStorage from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.scheduler_test_utils import FakeReq +from sglang.test.scheduler_test_utils import FakeReq, decode_queue register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -244,7 +244,7 @@ def test_pd_preallocation_leaves_blocked_request_queued(self): ) queue = DecodePreallocQueue.__new__(DecodePreallocQueue) queue.pp_size = 1 - queue.queue = [decode_req] + queue.queue = decode_queue(ListKind.PD_PREALLOC, [decode_req]) queue.pending_reqs = [] queue.retracted_queue = _RequestList(ListKind.PD_RETRACTED) queue.retracted_queue.join([req]) @@ -276,7 +276,7 @@ def test_pd_preallocation_leaves_blocked_request_queued(self): ), ): self.assertEqual(queue.pop_preallocated(), ([], [])) - self.assertEqual(queue.queue, [decode_req]) + self.assertEqual(list(queue.queue.items()), [decode_req]) queue._match_prefix_and_lock.assert_not_called() queue._pre_alloc.assert_not_called() self.assertIsNone(req.kv.req_pool_idx) diff --git a/test/registered/unit/observability/test_forward_pass_metrics.py b/test/registered/unit/observability/test_forward_pass_metrics.py index edad67588a5c..1ed15fc950ab 100644 --- a/test/registered/unit/observability/test_forward_pass_metrics.py +++ b/test/registered/unit/observability/test_forward_pass_metrics.py @@ -28,6 +28,7 @@ SchedulerMetricsReporter, _CacheHitRateWindow, ) +from sglang.srt.managers.scheduler_components.request_list import ListKind from sglang.srt.managers.scheduler_pp_mixin import PPBatchMetadata, SchedulerPPMixin from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.model_executor.forward_mode import ForwardMode @@ -49,6 +50,7 @@ ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils.device_timer import DeviceTimer, _TimingInterval, device_timer_ctx +from sglang.test.scheduler_test_utils import FakeReq, decode_queue from sglang.test.test_utils import CustomTestCase, enter_scope @@ -1136,11 +1138,12 @@ def test_host_receive_metrics_survive_queue_drain(self): queue=[], pre_allocated_token_counts=lambda: (0, 0) ) host_reqs = [ - types.SimpleNamespace(host_staged=True, priority=priority) + types.SimpleNamespace(req=FakeReq(priority=priority), host_staged=True) for priority in (1, 2) ] + device_req = types.SimpleNamespace(req=FakeReq(priority=1), host_staged=False) self.scheduler.disagg_decode_transfer_queue = types.SimpleNamespace( - queue=[*host_reqs, types.SimpleNamespace(host_staged=False, priority=1)], + queue=decode_queue(ListKind.PD_RECEIVING, [*host_reqs, device_req]), num_pending_deferred_releases=lambda: 0, ) for _ in host_reqs: diff --git a/test/registered/unit/observability/test_prealloc_token_usage.py b/test/registered/unit/observability/test_prealloc_token_usage.py index 5ae199d9aa1c..9b9e8b67bf01 100644 --- a/test/registered/unit/observability/test_prealloc_token_usage.py +++ b/test/registered/unit/observability/test_prealloc_token_usage.py @@ -14,12 +14,14 @@ SchedulerMetricsReporter, ) from sglang.srt.managers.scheduler_components.pool_stats_observer import PoolStats +from sglang.srt.managers.scheduler_components.request_list import ListKind from sglang.srt.mem_cache.unified_cache.component_type import ComponentType from sglang.srt.observability.metrics_collector import ( SchedulerMetricsCollector, SchedulerStats, ) from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.scheduler_test_utils import FakeReq, decode_queue from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=5, suite="base-a-test-cpu") @@ -33,12 +35,14 @@ def _transfer_req(fill_len, *, swa_tail_len=None, host_staged=False): kv = ReqKvInfo(kv_allocated_len=fill_len, kv_committed_len=fill_len) if swa_tail_len is not None: kv.set_evicted_seqlen(ComponentType.SWA, fill_len - swa_tail_len) - return SimpleNamespace(req=SimpleNamespace(kv=kv), host_staged=host_staged) + return SimpleNamespace(req=FakeReq(kv=kv), host_staged=host_staged) def _queue(reqs): queue = object.__new__(DecodePreallocQueue) - queue.transfer_queue = SimpleNamespace(queue=reqs) + queue.transfer_queue = SimpleNamespace( + queue=decode_queue(ListKind.PD_RECEIVING, reqs) + ) queue.token_to_kv_pool_allocator = SimpleNamespace(page_size=PAGE) return queue diff --git a/test/registered/unit/session/test_uocr_doc_prefetch_cuda.py b/test/registered/unit/session/test_uocr_doc_prefetch_cuda.py index 568fabcb4094..1b69dbd4906c 100644 --- a/test/registered/unit/session/test_uocr_doc_prefetch_cuda.py +++ b/test/registered/unit/session/test_uocr_doc_prefetch_cuda.py @@ -9,6 +9,7 @@ from sglang.srt.managers import mm_schedule from sglang.srt.models.unlimited_ocr_v2_support import document_prefetch as p from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.scheduler_test_utils import FakeReq from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") @@ -28,9 +29,10 @@ def test_parked_window_features_move_like_prefill_and_are_held(self): device=device, forward_stream=torch.cuda.Stream(device), tp_size=1, + abort_windows=None, ) item = SimpleNamespace(hash=7, feature=torch.ones(3, 4), offsets=[(2, 4)]) - req = SimpleNamespace( + req = FakeReq( rid="w1", sampling_params=SimpleNamespace( custom_params={p.KEY_AFTER: "w0", p.KEY_BUDGET: 10}, diff --git a/test/registered/unit/session/test_uocr_doc_prefetch_unit.py b/test/registered/unit/session/test_uocr_doc_prefetch_unit.py index 4d3686671bc0..57340b22845a 100644 --- a/test/registered/unit/session/test_uocr_doc_prefetch_unit.py +++ b/test/registered/unit/session/test_uocr_doc_prefetch_unit.py @@ -5,9 +5,11 @@ import torch +from sglang.srt.managers.io_struct import AbortReq from sglang.srt.models.unlimited_ocr_v2_support import document_prefetch as p from sglang.srt.models.unlimited_ocr_v2_support.document_reuse import UocrReqState from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.scheduler_test_utils import FakeReq from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=2, suite="base-a-test-cpu") @@ -19,7 +21,7 @@ def _req(rid, *, after=None, budget=None, keep=False, max_new=1000): if after is not None: cp[p.KEY_AFTER] = after cp[p.KEY_BUDGET] = budget - return SimpleNamespace( + return FakeReq( rid=rid, sampling_params=SimpleNamespace(custom_params=cp, max_new_tokens=max_new), multimodal_inputs=None, @@ -30,12 +32,21 @@ def _req(rid, *, after=None, budget=None, keep=False, max_new=1000): ) -def _prefetch(): +def _prefetch(aborted=None): return p.DocPrefetch( - early=None, device=torch.device("cpu"), forward_stream=None, tp_size=1 + early=None, + device=torch.device("cpu"), + forward_stream=None, + tp_size=1, + abort_windows=lambda reqs, finished_reason: aborted.extend(reqs), ) +def _abort(prefetch, rid): + for windows in prefetch.request_lists().values(): + windows.abort(AbortReq(rid=rid)) + + class _Sched: """Scheduler stand-in: a waiting queue that parks follow-ups first.""" @@ -96,6 +107,20 @@ def test_a_follow_up_arriving_late_is_not_parked(self): # Spliced text is not generation; it does not spend the budget. self.assertEqual(w1.sampling_params.max_new_tokens, 479) + def test_every_follow_up_of_a_window_is_released(self): + """Two follow-ups waiting on one window are both queued when it finishes.""" + prefetch = _prefetch() + sched = _Sched(prefetch) + w0 = _req("w0", keep=True) + w1a = _req("w1a", after="w0", budget=1000) + w1b = _req("w1b", after="w0", budget=1000) + sched.add(w1a) + sched.add(w1b) + _finish(w0, 300) + prefetch.on_window_finished(w0) + sched.ingest() + self.assertEqual(sched.queue, [w1a, w1b]) + def test_a_retracted_window_releases_nothing(self): """Retraction passes the same cache hook unfinished; releasing there would start the follow-up before the window's KV is retained.""" @@ -111,11 +136,13 @@ def test_a_retracted_window_releases_nothing(self): def test_aborting_the_document_drops_its_held_windows(self): """An abort must reach a window still parked off the queue, and one released but not yet queued by the scheduler.""" - prefetch = _prefetch() + aborted = [] + prefetch = _prefetch(aborted) sched = _Sched(prefetch) w1 = _req("doc-w1-ab", after="doc-w0-cd", budget=1000) sched.add(w1) - self.assertEqual(prefetch.drop("doc-w1"), [w1]) + _abort(prefetch, "doc-w1") + self.assertEqual(aborted, [w1]) self.assertEqual(prefetch.held_count(), 0) w0 = _req("doc-w0-ef", keep=True) @@ -123,7 +150,8 @@ def test_aborting_the_document_drops_its_held_windows(self): sched.add(w2) _finish(w0, 10) prefetch.on_window_finished(w0) - self.assertEqual(prefetch.drop("doc-w2"), [w2]) + _abort(prefetch, "doc-w2") + self.assertEqual(aborted, [w1, w2]) sched.ingest() self.assertEqual(sched.queue, []) diff --git a/test/registered/unit/utils/test_weight_versions.py b/test/registered/unit/utils/test_weight_versions.py index 315b167be8bf..37adfe50e95f 100644 --- a/test/registered/unit/utils/test_weight_versions.py +++ b/test/registered/unit/utils/test_weight_versions.py @@ -400,9 +400,7 @@ def test_records_on_all_pp_microbatches(self): def test_records_on_hisparse_staging_requests(self): """Requests parked in the HiSparse staging queue are swept like any live request.""" staging_req = _ReqStub(2) - coordinator = SimpleNamespace( - ack_staging_queue=[SimpleNamespace(req=staging_req)] - ) + coordinator = SimpleNamespace(ack_staging_queue=[staging_req]) scheduler = self._scheduler("v1", [], [], hisparse=coordinator) Scheduler.record_weight_version_change(scheduler, new_version="v2")