Merge 8fdc28005071d2319e9455eeba2b875536a64c42 into 553741557a57df44a46fb3ab72b6aafeaf945337

csrc/engine/scheduler.hpp

@@ -143,21 +143,25 @@ struct Scheduler {
 
     std::vector<SequencePtr> abort_sequence(std::vector<std::string> &request_ids) {
         std::vector<SequencePtr> 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;

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)

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)
 

tests/engine/conftest.py

@@ -0,0 +1,146 @@
+"""CPU fixtures for the native scheduler; no model weights or device execution."""
+
+from collections.abc import Callable, Iterator
+from itertools import count
+from types import SimpleNamespace
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pytest
+
+if TYPE_CHECKING:
+    from torch.distributed import ProcessGroup
+
+    from deepinfer._C.engine import BlockManager, Scheduler, Sequence, SequenceMetadataPool
+
+
+@pytest.fixture(scope="module")
+def scheduler_process_group() -> Iterator["ProcessGroup"]:
+    import torch.distributed as dist
+
+    owns_group = not dist.is_initialized()
+    if owns_group:
+        dist.init_process_group("gloo", store=dist.HashStore(), rank=0, world_size=1)
+    elif dist.get_world_size() != 1 or dist.get_backend() != "gloo":
+        pytest.skip("Native scheduler unit tests require a single-rank CPU Gloo group")
+    try:
+        yield dist.group.WORLD
+    finally:
+        if owns_group:
+            dist.destroy_process_group()
+
+
+@pytest.fixture
+def scheduler_factory(
+    scheduler_process_group: "ProcessGroup", monkeypatch: pytest.MonkeyPatch
+) -> Callable[..., tuple["Scheduler", "BlockManager", "BlockManager | None"]]:
+    from deepinfer._C.engine import BlockManager, Scheduler
+
+    def make(
+        num_blocks: int = 33,
+        *,
+        num_states: int | None = None,
+        stop_preempted_request: bool = True,
+        nextn: int = 0,
+        max_num_batched_tokens: int = 64,
+        max_num_seqs: int = 4,
+        max_model_len: int = 256,
+    ) -> tuple["Scheduler", "BlockManager", "BlockManager | None"]:
+        monkeypatch.setenv("DEEPINFER_USE_NEXTN_DECODE", str(nextn))
+        blocks = BlockManager(num_blocks)
+        states = BlockManager(num_states) if num_states is not None else None
+        config = SimpleNamespace(
+            max_model_len=max_model_len,
+            runtime_max_model_len=max_model_len + 64,
+            max_num_batched_tokens=max_num_batched_tokens,
+            max_num_seqs=max_num_seqs,
+            stop_preempted_request=stop_preempted_request,
+            chunked_prefill=True,
+            engine_type=0,
+            prefill_attempt_interval=2,
+            prefill_any_attempt_interval=100,
+        )
+        scheduler = Scheduler(
+            block_size=16,
+            block_manager=blocks,
+            state_manager=states,
+            eos_token_id=2,
+            dp_group=scheduler_process_group,
+            scheduler_config=config,
+            truncate_len=0,
+            lookahead_tokens=0,
+            prefill_max_sum_seqlen_k=4096,
+            flops_estimator=None,
+            engram_lookup_table=[],
+            cp_size=1,
+            vocab_size=100,
+            log_stats=False,
+        )
+        return scheduler, blocks, states
+
+    return make
+
+
+@pytest.fixture
+def sequence_factory() -> Callable[..., "Sequence"]:
+    from deepinfer._C.engine import Sequence
+    from deepinfer.sampling_params import SamplingParams
+
+    ids = count()
+
+    def make(length: int = 128, *, request_id: str = "request", block_size: int = 16) -> "Sequence":
+        tokens = np.ones(length, dtype=np.int32)
+        return Sequence(
+            seq_id=next(ids),
+            prompt_token_ids=tokens,
+            compressed_prompt_token_ids=tokens,
+            multi_modal_data=None,
+            block_size=block_size,
+            request_id=request_id,
+            user_id="test",
+            sampling_params=SamplingParams(max_tokens=4, ignore_eos=True, eos_id=2, seq_seed=0),
+            arrival_time=0.0,
+            golden_token_ids=None,
+        )
+
+    return make
+
+
+@pytest.fixture
+def metadata_pool() -> "SequenceMetadataPool":
+    from deepinfer._C.engine import SequenceMetadataPool
+
+    return SequenceMetadataPool()
+
+
+@pytest.fixture
+def unscheduled_outputs() -> Callable:
+    import torch
+
+    from deepinfer._C.engine import RequestOutput, RequestOutputPool, build_request_outputs
+
+    def build(outputs: object) -> list[object]:
+        pool = RequestOutputPool()
+        pointers = build_request_outputs(
+            outputs,
+            64,
+            torch.empty(0, dtype=torch.int32),
+            torch.empty(0, dtype=torch.float32),
+            top_logprobs_=None,
+            top_token_ids_=None,
+            sum_smaller_probs_=None,
+            accepted_token_ids_=None,
+            accepted_token_logprobs_=None,
+            accepted_top_logprobs_=None,
+            accepted_top_token_ids_=None,
+            accepted_sum_smaller_probs_=None,
+            accepted_token_last_logprobs_=None,
+            pushed_tokens_before_finish=np.empty(0, dtype=np.int32),
+            prompt_logprobs=None,
+            request_output_pool=pool,
+        )
+        result = [RequestOutput.clone(int(pointer)) for pointer in pointers]
+        pool.release(pointers)
+        return result
+
+    return build

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)

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

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