DeepInfer/tests/engine/test_scheduler_resources.py

"""KV blocks and state slots must always find their way back to the managers, and the requests holding them must always be able to finish.

Drives the real C++ scheduler through the part of LLMEngine.step() that touches it: schedule, free the finished
sequences, feed the (fake) model outputs back. No GPU or model is needed, only the compiled extension and a CPU
process group.
"""

from __future__ import annotations

from collections.abc import Generator
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import pytest
import torch  # noqa: F401  Load torch libraries before the C++ extension.
import torch.distributed as dist

from deepinfer._C.engine import BlockManager, Scheduler, Sequence, SequenceMetadataPool, process_MTPInfo
from deepinfer.sampling_params import SamplingParams

pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]

EOS = 1
VOCAB_SIZE = 1000
GENERATED_TOKEN = 5  # never EOS, so requests end by max_tokens


@pytest.fixture(scope="module")
def cpu_group(tmp_path_factory: pytest.TempPathFactory) -> Generator[dist.ProcessGroup]:
    initialized_here = not dist.is_initialized()
    if initialized_here:
        rendezvous = Path(tmp_path_factory.mktemp("rendezvous")) / "store"
        dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=0, world_size=1)
    yield dist.new_group(ranks=[0], backend="gloo")
    if initialized_here:
        dist.destroy_process_group()


class _Engine:
    """schedule -> free finished sequences -> process outputs, without speculative decoding or zero-overhead scheduling."""

    def __init__(
        self,
        group: dist.ProcessGroup,
        monkeypatch: pytest.MonkeyPatch,
        *,
        num_blocks: int,
        block_size: int,
        max_num_batched_tokens: int,
        num_states: int = 0,
        max_num_seqs: int = 2,
        prefill_attempt_interval: int = 2,
    ) -> None:
        monkeypatch.setenv("DEEPINFER_USE_NEXTN_DECODE", "0")  # read by the C++ scheduler when it is built
        self.num_blocks = num_blocks
        self.num_states = num_states
        self.block_size = block_size
        self.block_manager = BlockManager(num_blocks)
        self.state_manager = BlockManager(num_states) if num_states else None
        max_model_len = 64
        config = SimpleNamespace(
            max_model_len=max_model_len,
            runtime_max_model_len=max_model_len + 16,
            max_num_batched_tokens=max_num_batched_tokens,
            max_num_seqs=max_num_seqs,
            stop_preempted_request=False,
            chunked_prefill=True,
            engine_type=0,  # EngineType.NORMAL
            prefill_attempt_interval=prefill_attempt_interval,
            prefill_any_attempt_interval=prefill_attempt_interval + 1,
        )
        self.scheduler = Scheduler(
            block_size, self.block_manager, self.state_manager, EOS, group, config, 0, 0, 2**30, None, [], 1, VOCAB_SIZE, False
        )
        self.pool = SequenceMetadataPool()
        self.sequences: dict[str, Sequence] = {}

    def add_request(
        self, prompt_len: int, max_tokens: int, *, multi_modal_data: dict | None = None, allocate_blocks: bool = False
    ) -> Sequence:
        request_id = str(len(self.sequences))
        prompt = np.full(prompt_len, 7, dtype=np.int32)
        seq = Sequence(
            seq_id=len(self.sequences), prompt_token_ids=prompt, compressed_prompt_token_ids=prompt, multi_modal_data=multi_modal_data,
            block_size=self.block_size, request_id=request_id, user_id="user",
            sampling_params=SamplingParams(max_tokens=max_tokens, eos_id=EOS), arrival_time=0.0, golden_token_ids=None,
        )  # fmt: skip
        self.sequences[request_id] = seq
        if allocate_blocks:
            assert self.scheduler.try_allocate_seq(seq)
        self.scheduler.add_sequence(seq)
        return seq

    def step(self, *, decode_only: bool = False) -> None:
        if decode_only:
            # In a data-parallel engine every rank has to agree to a prompt step, so a rank can run decode steps between two chunks
            # of a prompt it is prefilling. A decode_step that is no attempt point gives the same schedule on a single rank.
            self.scheduler.decode_step = 1
        seq_metadata_list, outputs = self.scheduler.schedule(self.pool)
        for seq in outputs.finished_seqs:
            self.scheduler.free_seq(seq)
        n = len(outputs.scheduled_seqs)
        process_MTPInfo(outputs, False, torch.zeros((n, 0), dtype=torch.int32), None)
        self.scheduler.process_outputs(
            outputs, torch.empty(0, dtype=torch.int32), torch.full((n,), GENERATED_TOKEN, dtype=torch.int32), None
        )
        self.pool.release(seq_metadata_list)

    def run(self, max_steps: int = 300) -> None:
        for _ in range(max_steps):
            if not self.scheduler.has_unfinished_seqs():
                return
            self.step()
        raise AssertionError(f"no progress after {max_steps} steps: {self.scheduler.get_num_unfinished_seqs()} requests unfinished")

    def assert_all_resources_returned(self) -> None:
        # Block 0 (and state slot 0) is reserved by the managers and never handed out.
        assert self.scheduler.get_num_free_gpu_blocks() == self.num_blocks - 1
        if self.state_manager is not None:
            assert self.state_manager.get_num_free_blocks() == self.num_states - 1


@pytest.mark.parametrize("num_states", [0, 8], ids=["no_state", "with_state"])
def test_abort_between_prefill_chunks_releases_the_sequence_resources(
    cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch, num_states: int
) -> None:
    engine = _Engine(cpu_group, monkeypatch, num_blocks=17, block_size=4, max_num_batched_tokens=8, num_states=num_states)
    long_prompt = engine.add_request(prompt_len=30, max_tokens=3)  # 4 chunks of at most 8 tokens
    engine.add_request(prompt_len=6, max_tokens=3)

    engine.step()
    # The first chunk is done: the request waits for its next chunk and already owns the KV blocks of the whole prompt.
    assert long_prompt.status.name == "WAITING"
    assert long_prompt.num_cached_tokens == 8
    assert engine.scheduler.get_num_free_gpu_blocks() < engine.num_blocks - 1

    engine.scheduler.abort_sequence([long_prompt.request_id])
    engine.run()

    assert long_prompt.status.name == "FINISHED_ABORTED"
    engine.assert_all_resources_returned()


@pytest.mark.parametrize("num_states", [0, 8], ids=["no_state", "with_state"])
def test_abort_during_a_decode_step_releases_the_sequence_resources(
    cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch, num_states: int
) -> None:
    engine = _Engine(cpu_group, monkeypatch, num_blocks=17, block_size=4, max_num_batched_tokens=8, num_states=num_states)
    engine.add_request(prompt_len=6, max_tokens=12)
    engine.step()
    long_prompt = engine.add_request(prompt_len=30, max_tokens=3)  # 4 chunks of at most 8 tokens
    engine.step()
    assert long_prompt.status.name == "WAITING"
    assert long_prompt.num_cached_tokens == 8

    # The aborted prompt is dropped from the queue by a decode step this time, not by a prompt step.
    engine.scheduler.abort_sequence([long_prompt.request_id])
    engine.step(decode_only=True)
    assert long_prompt.status.name == "FINISHED_ABORTED"
    assert not any(seq.request_id == long_prompt.request_id for seq in engine.scheduler.waiting)

    engine.run()
    engine.assert_all_resources_returned()


def test_abort_before_the_first_chunk_leaves_nothing_to_release(cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch) -> None:
    engine = _Engine(cpu_group, monkeypatch, num_blocks=17, block_size=4, max_num_batched_tokens=8, num_states=8)
    request = engine.add_request(prompt_len=30, max_tokens=3)
    engine.scheduler.abort_sequence([request.request_id])

    engine.run()

    assert request.status.name == "FINISHED_ABORTED"
    engine.assert_all_resources_returned()


def _preempt_during_chunked_prefill(
    group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch, *, num_states: int = 0
) -> tuple[_Engine, list[Sequence], Sequence]:
    """Two requests decode while a long prompt is prefilled in chunks, and the cache is too small for all of them to keep growing.

    The cache has 6 usable blocks of 4 tokens. The prompt (12 tokens, chunks of 4) reserves 3 of them with its first chunk, the
    decoding requests hold one each, and the last block goes to one of them for its next token. The other one needs a new block too,
    is preempted and has to be recomputed, which takes 2 blocks while only the 1 it just released is free.
    """
    engine = _Engine(group, monkeypatch, num_blocks=7, block_size=4, max_num_batched_tokens=4, max_num_seqs=4, num_states=num_states)
    decoding = [engine.add_request(prompt_len=4, max_tokens=12), engine.add_request(prompt_len=4, max_tokens=12)]
    long_prompt = engine.add_request(prompt_len=12, max_tokens=2)
    for _ in range(10):
        if long_prompt.num_cached_tokens:
            break
        engine.step()
    assert long_prompt.status.name == "WAITING"
    assert long_prompt.num_cached_tokens == 4
    assert engine.scheduler.get_num_free_gpu_blocks() == 1

    for _ in range(10):
        if engine.scheduler.get_num_preempted_seqs():
            break
        engine.step(decode_only=True)
    assert engine.scheduler.get_num_preempted_seqs() == 1
    assert sorted(seq.status.name for seq in decoding) == ["RUNNING", "WAITING"]
    assert long_prompt.status.name == "WAITING"
    return engine, decoding, long_prompt


def test_abort_of_a_preempted_request_does_not_free_its_state_slot_twice(
    cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch
) -> None:
    engine, decoding, _ = _preempt_during_chunked_prefill(cpu_group, monkeypatch, num_states=8)
    preempted = next(seq for seq in decoding if seq.status.name == "WAITING")
    # Preemption already gave its state slot back; releasing the aborted sequence again must not release it a second time.
    engine.scheduler.abort_sequence([preempted.request_id])

    engine.run()

    assert preempted.status.name == "FINISHED_ABORTED"
    engine.assert_all_resources_returned()


def test_preempted_request_does_not_block_a_chunked_prefill_that_holds_blocks(
    cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch
) -> None:
    engine, decoding, long_prompt = _preempt_during_chunked_prefill(cpu_group, monkeypatch)

    # The preempted request cannot be recomputed before the long prompt has finished its prefill and released its blocks.
    engine.run()

    assert [seq.status.name for seq in decoding] == ["FINISHED_LENGTH_CAPPED"] * 2
    assert long_prompt.status.name == "FINISHED_LENGTH_CAPPED"
    engine.assert_all_resources_returned()


def test_preempted_request_is_not_queued_behind_a_sequence_that_cannot_be_scheduled(
    cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch
) -> None:
    # 7 usable blocks of 4 tokens. The two decoding requests need 3 blocks each in the end, the multimodal request holds 3.
    engine = _Engine(cpu_group, monkeypatch, num_blocks=8, block_size=4, max_num_batched_tokens=64, max_num_seqs=3)
    decoding = [engine.add_request(prompt_len=4, max_tokens=8), engine.add_request(prompt_len=4, max_tokens=8)]
    # It is not schedulable until its images are processed, but it owns its KV blocks (examples/benchmark.py allocates them up front).
    multimodal = engine.add_request(prompt_len=12, max_tokens=3, multi_modal_data={"image": 1}, allocate_blocks=True)
    assert not multimodal.schedulable

    for _ in range(40):
        if all(seq.status.name == "FINISHED_LENGTH_CAPPED" for seq in decoding):
            break
        engine.step()

    # The cache is too small for both, so one of them was preempted and recomputed. Queued behind the multimodal request, which
    # cannot be scheduled yet, it would wait for it forever.
    assert engine.scheduler.get_num_preempted_seqs() > 0
    assert [seq.status.name for seq in decoding] == ["FINISHED_LENGTH_CAPPED"] * 2

    engine.scheduler.set_schedulable(multimodal.request_id)
    engine.run()
    assert multimodal.status.name == "FINISHED_LENGTH_CAPPED"
    engine.assert_all_resources_returned()