[Refactor]: clean up diffusion scheduler interfaces (#5216)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: Zeyu Huang | 黃澤宇 <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py

@@ -20,7 +20,7 @@ from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
 )
 from vllm_omni.diffusion.worker.input_batch import InputBatch
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-from vllm_omni.diffusion.worker.utils import DiffusionRequestState
+from vllm_omni.diffusion.worker.utils import StepRequestState
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
@@ -38,8 +38,8 @@ def _pipeline():
     return pipeline
 
 
-def _state(request_id: str, step_index: int) -> DiffusionRequestState:
-    state = DiffusionRequestState(
+def _state(request_id: str, step_index: int) -> StepRequestState:
+    state = StepRequestState(
         request_id=request_id,
         sampling=SimpleNamespace(),
         prompt="prompt",
@@ -133,7 +133,7 @@ def test_prepare_encode_preserves_normal_hunyuan_bot_task_semantics(
 
     monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt)
     pipeline.prepare_model_inputs = fake_prepare_model_inputs
-    state = DiffusionRequestState(
+    state = StepRequestState(
         request_id="req-bot-task",
         sampling=sampling,
         prompt=prompt_item,

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -36,10 +36,10 @@ from vllm_omni.diffusion.sched.interface import RequestBatchSamplingParamsKey
 from vllm_omni.diffusion.worker.utils import RunnerOutput
 
 # Default values for every batch-key field, so SimpleNamespace-based
-# sampling_params satisfy ``get_sampling_params_key``'s attribute lookups.
+# sampling_params satisfy ``RequestScheduler._build_sampling_params_key``'s attribute lookups.
 _SAMPLING_KEY_DEFAULTS = {f.name: f.default for f in _dc_fields(RequestBatchSamplingParamsKey)}
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.cpu]
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 
 
 # ───────────────────────────────────────── helpers ─────────────────────────

tests/diffusion/test_diffusion_scheduler.py

@@ -14,10 +14,10 @@ from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedErr
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
+    BaseScheduler,
     DiffusionRequestStatus,
     RequestScheduler,
     Scheduler,
-    SchedulerInterface,
     StepScheduler,
 )
 from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
@@ -85,7 +85,7 @@ def _cached_ids(sched_output) -> list[str]:
     return list(sched_output.scheduled_cached_reqs.request_ids)
 
 
-class _StubScheduler(SchedulerInterface):
+class _StubScheduler:
     def __init__(self, request: OmniDiffusionRequest, output) -> None:
         self._request = request
         self._output = output
@@ -154,8 +154,14 @@ class _StubScheduler(SchedulerInterface):
         return None
 
 
-class TestGetSamplingParamsKey:
-    """Pure-function tests for the batch-compatibility key builder."""
+class _ConcreteScheduler(BaseScheduler):
+    def update_from_output(self, sched_output, output) -> set[str]:
+        del sched_output, output
+        return set()
+
+
+class TestGetStepBatchSamplingParamsKey:
+    """Tests for the step-batch compatibility key builder on BaseScheduler."""
 
     @staticmethod
     def _make(lora_int_id: int | None = None, lora_scale: float = 1.0) -> OmniDiffusionRequest:
@@ -176,34 +182,32 @@ class TestGetSamplingParamsKey:
         )
 
     def test_distinguishes_lora_id(self) -> None:
-        from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
-
-        assert get_sampling_params_key(self._make(lora_int_id=1)) != get_sampling_params_key(self._make(lora_int_id=2))
+        scheduler = _ConcreteScheduler()
+        assert scheduler._build_sampling_params_key(self._make(lora_int_id=1)) != scheduler._build_sampling_params_key(
+            self._make(lora_int_id=2)
+        )
 
     def test_distinguishes_lora_scale(self) -> None:
-        from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
-
-        assert get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5)) != get_sampling_params_key(
-            self._make(lora_int_id=1, lora_scale=1.0)
-        )
+        scheduler = _ConcreteScheduler()
+        assert scheduler._build_sampling_params_key(
+            self._make(lora_int_id=1, lora_scale=0.5)
+        ) != scheduler._build_sampling_params_key(self._make(lora_int_id=1, lora_scale=1.0))
 
     def test_treats_no_lora_as_distinct_bucket(self) -> None:
-        from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
-
-        assert get_sampling_params_key(self._make(lora_int_id=None)) != get_sampling_params_key(
-            self._make(lora_int_id=1)
-        )
+        scheduler = _ConcreteScheduler()
+        assert scheduler._build_sampling_params_key(
+            self._make(lora_int_id=None)
+        ) != scheduler._build_sampling_params_key(self._make(lora_int_id=1))
 
     def test_equal_for_same_lora_identity(self) -> None:
-        from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
-
-        a = get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
-        b = get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
+        scheduler = _ConcreteScheduler()
+        a = scheduler._build_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
+        b = scheduler._build_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
         assert a == b
 
 
 class TestGetRequestBatchSamplingParamsKey:
-    """Pure-function tests for the request-batch compatibility key builder."""
+    """Tests for the request-batch compatibility key builder on RequestScheduler."""
 
     @staticmethod
     def _make(
@@ -220,21 +224,19 @@ class TestGetRequestBatchSamplingParamsKey:
         return OmniDiffusionRequest(prompt="prompt", sampling_params=sp, request_id=f"req-{num_inference_steps}")
 
     def test_distinguishes_num_inference_steps(self) -> None:
-        from vllm_omni.diffusion.sched.base_scheduler import get_request_batch_sampling_params_key
-
-        assert get_request_batch_sampling_params_key(
+        scheduler = RequestScheduler()
+        assert scheduler._build_sampling_params_key(
             self._make(num_inference_steps=2)
-        ) != get_request_batch_sampling_params_key(self._make(num_inference_steps=4))
+        ) != scheduler._build_sampling_params_key(self._make(num_inference_steps=4))
 
     def test_ignores_seed_and_generator(self) -> None:
-        from vllm_omni.diffusion.sched.base_scheduler import get_request_batch_sampling_params_key
-
+        scheduler = RequestScheduler()
         gen_a = torch.Generator(device="cpu").manual_seed(1)
         gen_b = torch.Generator(device="cpu").manual_seed(2)
 
-        assert get_request_batch_sampling_params_key(
+        assert scheduler._build_sampling_params_key(
             self._make(seed=1, generator=gen_a)
-        ) == get_request_batch_sampling_params_key(self._make(seed=2, generator=gen_b))
+        ) == scheduler._build_sampling_params_key(self._make(seed=2, generator=gen_b))
 
 
 class TestRequestScheduler:

tests/diffusion/test_diffusion_step_pipeline.py

@@ -43,7 +43,7 @@ from vllm_omni.diffusion.sched.interface import (
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 from vllm_omni.diffusion.worker.input_batch import InputBatch
-from vllm_omni.diffusion.worker.utils import DiffusionRequestState, RunnerOutput
+from vllm_omni.diffusion.worker.utils import RunnerOutput, StepRequestState
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -369,8 +369,8 @@ def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
     )
 
 
-def _make_input_batch_state(request_id: str, latent_value: float) -> DiffusionRequestState:
-    state = DiffusionRequestState(
+def _make_input_batch_state(request_id: str, latent_value: float) -> StepRequestState:
+    state = StepRequestState(
         request_id=request_id,
         sampling=SimpleNamespace(),
         prompt=None,

tests/diffusion/test_input_batch_txt_seq_lens.py

@@ -2,14 +2,14 @@ import pytest
 import torch
 
 from vllm_omni.diffusion.worker.input_batch import _prepare_request_prompt_field
-from vllm_omni.diffusion.worker.utils import DiffusionRequestState
+from vllm_omni.diffusion.worker.utils import StepRequestState
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-pytestmark = [pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
 def test_prepare_request_prompt_field_refreshes_txt_seq_lens_after_padding():
-    state = DiffusionRequestState(
+    state = StepRequestState(
         request_id="req-0",
         sampling=OmniDiffusionSamplingParams(),
         prompt_embeds=torch.randn(1, 10, 4),

tests/diffusion/test_worker_wrapper_base.py

@@ -184,7 +184,7 @@ class TestWorkerWrapperBaseDelegation:
         mock_reqs = [mocker.Mock()]
         result = wrapper.execute_model(mock_reqs, mock_od_config)
 
-        wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config, kv_prefetch_jobs=None)
+        wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config, kv_prefetch_job=None)
         assert result == mock_output
 
     def test_load_weights_delegation(self, mocker: MockerFixture, mock_od_config):