test: add CPU fixtures for native scheduler ownership checks

tests/engine/conftest.py

@@ -0,0 +1,113 @@
+"""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()