diff --git a/csrc/engine/scheduler.hpp b/csrc/engine/scheduler.hpp index ed6e897..0650e55 100644 --- a/csrc/engine/scheduler.hpp +++ b/csrc/engine/scheduler.hpp @@ -143,21 +143,25 @@ struct Scheduler { std::vector abort_sequence(std::vector &request_ids) { std::vector abort_seqs; - // 多模态处理失败的 seq 会带着 schedule.schedulable=false 滞留在 waiting,两个队列都要扫。 - auto abort_in_queue = [&](auto &queue, auto &request_id) { - for (auto &seq : queue) { - if (seq->request.request_id == request_id && !seq->is_finished()) { - // NOTE: We do NOT abort in PE, so we don't need to maintain - // store ids of aborted seqs to release in scheduler. - seq->schedule.status = SequenceStatus::FINISHED_ABORTED; - num_aborted += 1; - abort_seqs.push_back(seq); - } + auto abort_seq = [&](auto &seq, auto &request_id) { + if (seq->request.request_id == request_id && !seq->is_finished()) { + // NOTE: We do NOT abort in PE, so we don't need to maintain + // store ids of aborted seqs to release in scheduler. + seq->schedule.status = SequenceStatus::FINISHED_ABORTED; + num_aborted += 1; + abort_seqs.push_back(seq); } }; for (auto &request_id : request_ids) { - abort_in_queue(running, request_id); - abort_in_queue(waiting, request_id); + for (auto &seq : running) { + abort_seq(seq, request_id); + } + for (auto &seq : waiting) { + abort_seq(seq, request_id); + } + for (auto &[seq, forked] : prompted) { + abort_seq(seq, request_id); + } } return abort_seqs; } @@ -576,6 +580,21 @@ struct Scheduler { } } + // Canceled forks must reach finalization even when running is full or + // an earlier fork is waiting for clone resources. + for (auto it = prompted.begin(); it != prompted.end();) { + auto &seq = std::get<0>(*it); + if (seq->is_finished()) { + // These forks have never produced a model output: notify the + // caller via ignored_seqs and release ownership via finished_seqs. + ignored_seqs.push_back(seq); + finished_seqs.push_back(std::move(seq)); + it = prompted.erase(it); + } else { + ++it; + } + } + const size_t max_seqs = max_num_seqs; // prompted 队列头部的 seq 依次转入 running(或已结束的转入 finished), @@ -594,11 +613,6 @@ struct Scheduler { while (!prompted.empty() && running.size() < max_seqs) { auto &&[seq, forked] = prompted.front(); - if (seq->schedule.status == SequenceStatus::FINISHED_ABORTED) { // 进 CE 之前就 abort 了 - finished_seqs.push_back(std::move(seq)); - prompted.pop_front(); - continue; - } if (forked) { if ((state_manager && !state_manager->can_allocate_single_nolock()) || !block_manager->can_allocate_single_nolock()) { break; diff --git a/deepinfer/engine/llm_engine.py b/deepinfer/engine/llm_engine.py index cb87941..489d0a1 100644 --- a/deepinfer/engine/llm_engine.py +++ b/deepinfer/engine/llm_engine.py @@ -648,6 +648,11 @@ class LLMEngine: assert isinstance(request_id, (str, dict, set, list, tuple)), f"invalid request_id type: {type(request_id)}" request_id = [request_id] if isinstance(request_id, str) else list(request_id) + if self._so_matcher_builder is not None: + for seq in self._so_matcher_builder.abort(request_id): + # Use the scheduler's normal abort/output path without invoking + # the ready callback, which would start multimodal processing. + self.scheduler.add_sequence(seq) seqs = self.scheduler.abort_sequence(request_id) if self.multi_modal_processor is not None: self.multi_modal_processor.abort(request_id) diff --git a/deepinfer/structure_output/so_matcher_builder.py b/deepinfer/structure_output/so_matcher_builder.py index 030d2dc..aaa4403 100644 --- a/deepinfer/structure_output/so_matcher_builder.py +++ b/deepinfer/structure_output/so_matcher_builder.py @@ -120,6 +120,20 @@ class SOMatcherBuilder: on_ready(seq) self._pending = still_pending + def abort(self, request_ids: list[str]) -> list[object]: + """Remove canceled requests; the caller schedules their terminal output.""" + request_id_set = set(request_ids) + aborted = [] + still_pending = [] + for seq, future, on_ready in self._pending: + if seq.request_id in request_id_set: + future.release() + aborted.append(seq) + else: + still_pending.append((seq, future, on_ready)) + self._pending = still_pending + return aborted + def has_pending(self) -> bool: return bool(self._pending) diff --git a/tests/engine/test_request_cancellation.py b/tests/engine/test_request_cancellation.py new file mode 100644 index 0000000..0a9abc8 --- /dev/null +++ b/tests/engine/test_request_cancellation.py @@ -0,0 +1,67 @@ +from collections.abc import Callable +from concurrent.futures import Future +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch # noqa: F401 Load torch libraries before the C++ extension. + +from deepinfer._C.engine import SequenceMetadataPool, SequenceStatus +from deepinfer.engine.llm_engine import LLMEngine +from deepinfer.structure_output.matcher import XGrammarMatcher, XGrammarMatcherFuture +from deepinfer.structure_output.so_matcher_builder import SOMatcherBuilder + + +@pytest.mark.parametrize("completed", [False, True]) +def test_abort_pending_grammar_uses_scheduler_terminal_path( + scheduler_factory: Callable, + sequence_factory: Callable, + metadata_pool: SequenceMetadataPool, + unscheduled_outputs: Callable, + monkeypatch: pytest.MonkeyPatch, + completed: bool, +) -> None: + scheduler, blocks, _ = scheduler_factory() + engine = object.__new__(LLMEngine) + engine.closed = True # No workers or compiler pool were started. + engine.scheduler = scheduler + engine.multi_modal_processor = SimpleNamespace(abort=Mock(), add=Mock()) + builder = object.__new__(SOMatcherBuilder) + engine._so_matcher_builder = builder + seq = sequence_factory(length=16, request_id="grammar") + task = Future() + if completed: + task.set_result(XGrammarMatcher(None, None, None, [])) + future = XGrammarMatcherFuture(task) + on_ready = Mock() + builder._pending = [(seq, future, on_ready)] + release_handles = Mock() + monkeypatch.setattr("deepinfer.engine.llm_engine.free_multi_modal_pool_handles", release_handles) + free_blocks = blocks.get_num_free_blocks() + + assert engine.abort_request(["grammar", "grammar"]) == (1, 16) + assert engine.abort_request("grammar") == (0, 0) + assert scheduler.get_num_aborted_seqs() == 1 + assert seq.status == SequenceStatus.FINISHED_ABORTED + assert not builder.has_pending() + assert future.task is None + release_handles.assert_called_once_with(seq.multi_modal_data) + on_ready.assert_not_called() + engine.multi_modal_processor.add.assert_not_called() + + builder.pump() + scheduler.decode_step = 1 + metadata, outputs = scheduler.schedule(metadata_pool) + assert not outputs.scheduled_seqs + assert scheduler.get_num_unfinished_seqs() == 0 + assert blocks.get_num_free_blocks() == free_blocks + terminal = unscheduled_outputs(outputs) + assert len(terminal) == 1 + assert terminal[0].request_id == "grammar" + assert terminal[0].finished + assert terminal[0].finish_reason == "abort" + metadata_pool.release(metadata) + scheduler.decode_step = 1 + metadata, outputs = scheduler.schedule(metadata_pool) + assert unscheduled_outputs(outputs) == [] + metadata_pool.release(metadata) diff --git a/tests/engine/test_scheduler_cancellation.py b/tests/engine/test_scheduler_cancellation.py new file mode 100644 index 0000000..70b8afd --- /dev/null +++ b/tests/engine/test_scheduler_cancellation.py @@ -0,0 +1,100 @@ +from collections.abc import Callable + +import pytest +import torch # noqa: F401 Load torch libraries before the C++ extension. + +from deepinfer._C.engine import SequenceMetadataPool, SequenceStatus + + +@pytest.mark.parametrize("queue", ["waiting", "running"]) +def test_abort_existing_queue_is_idempotent(scheduler_factory: Callable, sequence_factory: Callable, queue: str) -> None: + scheduler, _, _ = scheduler_factory() + seq = sequence_factory(request_id="cancel") + other = sequence_factory(request_id="keep") + if queue == "waiting": + scheduler.add_sequence(seq) + scheduler.add_sequence(other) + else: + seq.status = SequenceStatus.RUNNING + other.status = SequenceStatus.RUNNING + scheduler.running = [seq, other] + + assert scheduler.abort_sequence(["missing", "cancel", "cancel"]) == [seq] + assert scheduler.abort_sequence(["cancel"]) == [] + assert scheduler.get_num_aborted_seqs() == 1 + assert seq.status == SequenceStatus.FINISHED_ABORTED + assert other.status in (SequenceStatus.WAITING, SequenceStatus.RUNNING) + + +@pytest.mark.parametrize("blocked_by", ["running_capacity", "earlier_fork"]) +def test_aborted_prompted_fork_finishes_without_waiting_for_admission( + scheduler_factory: Callable, + sequence_factory: Callable, + metadata_pool: SequenceMetadataPool, + unscheduled_outputs: Callable, + blocked_by: str, +) -> None: + scheduler, blocks, states = scheduler_factory(max_num_seqs=1, num_states=5) + initial_blocks = blocks.get_num_free_blocks() + initial_states = states.get_num_free_blocks() + source = sequence_factory(length=32, request_id="source") + assert scheduler.try_allocate_seq(source) + source.num_cached_tokens = source.get_len() - 1 + + def fork(request_id: str, seq_id: int): + child = scheduler.fork_seq(source, None, request_id, source.sampling_params, None, 2, -1, seq_id, 0.0) + child.request_id = request_id + return child + + held_blocks = [] + if blocked_by == "earlier_fork": + other = fork("keep", 2) + # The head fork cannot clone a block, even though running has space. + held_blocks = blocks.allocate_nolock(blocks.get_num_free_blocks()) + else: + other = sequence_factory(length=17, request_id="keep") + assert scheduler.try_allocate_seq(other) + other.num_cached_tokens = 16 + other.status = SequenceStatus.RUNNING + scheduler.running = [other] + canceled = fork("cancel", 3) + before_blocks = blocks.get_num_free_blocks() + before_states = states.get_num_free_blocks() + + assert scheduler.abort_sequence(["cancel", "cancel"]) == [canceled] + assert scheduler.abort_sequence(["cancel"]) == [] + scheduler.decode_step = 1 + metadata, outputs = scheduler.schedule(metadata_pool) + assert outputs.finished_seqs == [canceled] + assert canceled not in outputs.scheduled_seqs + assert outputs.clone_blocks.size == 0 + assert outputs.clone_states.size == 0 + assert blocks.get_num_free_blocks() == before_blocks + assert states.get_num_free_blocks() == before_states + scheduler.free_seq(canceled) + terminal = unscheduled_outputs(outputs) + assert len(terminal) == 1 + assert terminal[0].request_id == "cancel" + assert terminal[0].finished + assert terminal[0].finish_reason == "abort" + metadata_pool.release(metadata) + + # The canceled fork was removed, so it cannot be finalized twice. + scheduler.decode_step = 1 + metadata, outputs = scheduler.schedule(metadata_pool) + assert not outputs.finished_seqs + assert unscheduled_outputs(outputs) == [] + metadata_pool.release(metadata) + + assert scheduler.abort_sequence(["keep"]) == [other] + scheduler.decode_step = 1 + metadata, outputs = scheduler.schedule(metadata_pool) + assert outputs.finished_seqs == [other] + scheduler.free_seq(other) + metadata_pool.release(metadata) + scheduler.free_seq(source) + blocks.unpin_blocks(held_blocks) + assert scheduler.get_num_unfinished_seqs() == 0 + assert scheduler.get_num_aborted_seqs() == 2 + assert blocks.get_num_free_blocks() == initial_blocks + assert states.get_num_free_blocks() == initial_states diff --git a/tests/engine/test_so_matcher_cancellation.py b/tests/engine/test_so_matcher_cancellation.py new file mode 100644 index 0000000..24c3d70 --- /dev/null +++ b/tests/engine/test_so_matcher_cancellation.py @@ -0,0 +1,70 @@ +from concurrent.futures import Future +from types import SimpleNamespace + +import pytest + +from deepinfer.structure_output.matcher import XGrammarMatcher, XGrammarMatcherFuture +from deepinfer.structure_output.so_matcher_builder import SOMatcherBuilder + + +@pytest.mark.parametrize("task_state", ["queued", "running", "completed"]) +def test_abort_removes_pending_request_before_compilation_is_pumped(task_state: str) -> None: + builder = object.__new__(SOMatcherBuilder) + canceled = SimpleNamespace(request_id="cancel", so_matcher=None) + retained = SimpleNamespace(request_id="keep", so_matcher=None) + task = Future() + matcher = XGrammarMatcher(None, None, None, []) + # Stand in for compiled resources so completed-future release is observable. + matcher.ctx = object() + if task_state == "running": + assert task.set_running_or_notify_cancel() + elif task_state == "completed": + task.set_result(matcher) + future = XGrammarMatcherFuture(task) + retained_task = Future() + ready = [] + builder._pending = [ + (canceled, future, ready.append), + (retained, XGrammarMatcherFuture(retained_task), ready.append), + ] + + assert builder.abort(["missing", "cancel", "cancel"]) == [canceled] + assert builder.abort(["cancel"]) == [] + assert builder.get_num_pending() == 1 + assert future.task is None + assert future.res is None + assert canceled.so_matcher is None + assert ready == [] + if task_state == "queued": + assert task.cancelled() + elif task_state == "running": + # A running compile cannot be interrupted; its eventual result must not + # put the canceled sequence back into the admission path. + task.set_result(matcher) + else: + assert matcher.ctx is None + + retained_matcher = XGrammarMatcher(None, None, None, []) + retained_task.set_result(retained_matcher) + builder.pump() + builder.pump() + assert ready == [retained] + assert retained.so_matcher is retained_matcher + assert not builder.has_pending() + assert builder.abort(["keep", "cancel"]) == [] + + +def test_abort_after_pump_leaves_admitted_request_to_scheduler() -> None: + builder = object.__new__(SOMatcherBuilder) + seq = SimpleNamespace(request_id="already-admitted", so_matcher=None) + matcher = XGrammarMatcher(None, None, None, []) + task = Future() + task.set_result(matcher) + ready = [] + builder._pending = [(seq, XGrammarMatcherFuture(task), ready.append)] + + builder.pump() + + assert builder.abort(["already-admitted"]) == [] + assert ready == [seq] + assert seq.so_matcher is matcher