[Refactor] Refactor Diffusion Scheduler/Executor Boundaries and Request State Flow (#1625)

Signed-off-by: jader <yjader@foxmail.com>
Signed-off-by: JiangJie Zhang <76905040+yJader@users.noreply.github.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>

tests/diffusion/test_diffusion_scheduler.py

@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import threading
+from unittest.mock import Mock, patch
+
+import pytest
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.sched import (
+    DiffusionRequestStatus,
+    RequestScheduler,
+    Scheduler,
+    SchedulerInterface,
+)
+from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _make_request(req_id: str) -> OmniDiffusionRequest:
+    return OmniDiffusionRequest(
+        prompts=[f"prompt_{req_id}"],
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        request_ids=[req_id],
+    )
+
+
+def _make_request_output(req_id: str, *, error: str | None = None) -> DiffusionOutput:
+    del req_id
+    return DiffusionOutput(output=None, error=error)
+
+
+def _new_ids(sched_output) -> list[str]:
+    return [req.sched_req_id for req in sched_output.scheduled_new_reqs]
+
+
+def _cached_ids(sched_output) -> list[str]:
+    return list(sched_output.scheduled_cached_reqs.sched_req_ids)
+
+
+class _StubScheduler(SchedulerInterface):
+    def __init__(self, request: OmniDiffusionRequest, output: DiffusionOutput) -> None:
+        self._request = request
+        self._output = output
+        self.initialized_with = None
+        self._sched_req_id = request.request_ids[0]
+        self._state = None
+        self._scheduled = False
+
+    def initialize(self, od_config) -> None:
+        self.initialized_with = od_config
+
+    def add_request(self, request: OmniDiffusionRequest) -> str:
+        assert request is self._request
+        self._state = Mock(sched_req_id=self._sched_req_id, req=request)
+        return self._sched_req_id
+
+    def schedule(self):
+        if self._scheduled or self._state is None:
+            return Mock(
+                scheduled_new_reqs=[],
+                scheduled_cached_reqs=CachedRequestData.make_empty(),
+                scheduled_req_ids=[],
+                is_empty=True,
+            )
+        self._scheduled = True
+        return Mock(
+            scheduled_new_reqs=[NewRequestData.from_state(self._state)],
+            scheduled_cached_reqs=CachedRequestData.make_empty(),
+            scheduled_req_ids=[self._state.sched_req_id],
+            is_empty=False,
+        )
+
+    def update_from_output(self, sched_output, output: DiffusionOutput) -> set[str]:
+        del sched_output
+        assert output is self._output
+        return {self._sched_req_id}
+
+    def has_requests(self) -> bool:
+        return not self._scheduled
+
+    def get_request_state(self, sched_req_id: str):
+        del sched_req_id
+        return self._state
+
+    def get_sched_req_id(self, request_id: str) -> str | None:
+        if request_id in self._request.request_ids:
+            return self._sched_req_id
+        return None
+
+    def pop_request_state(self, sched_req_id: str):
+        del sched_req_id
+        return self._state
+
+    def preempt_request(self, sched_req_id: str) -> bool:
+        del sched_req_id
+        return False
+
+    def finish_requests(self, sched_req_ids, status) -> None:
+        del sched_req_ids, status
+        return None
+
+    def close(self) -> None:
+        return None
+
+
+class TestRequestScheduler:
+    def setup_method(self) -> None:
+        self.scheduler: RequestScheduler = RequestScheduler()
+        self.scheduler.initialize(Mock())
+
+    def test_single_request_success_lifecycle(self) -> None:
+        req_id = self.scheduler.add_request(_make_request("a"))
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.WAITING
+
+        sched_output = self.scheduler.schedule()
+        assert _new_ids(sched_output) == [req_id]
+        assert _cached_ids(sched_output) == []
+        assert sched_output.num_running_reqs == 1
+        assert sched_output.num_waiting_reqs == 0
+
+        finished = self.scheduler.update_from_output(sched_output, _make_request_output(req_id))
+        assert finished == {req_id}
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
+        assert self.scheduler.has_requests() is False
+
+    def test_error_output_marks_finished_error(self) -> None:
+        req_id = self.scheduler.add_request(_make_request("err"))
+
+        sched_output = self.scheduler.schedule()
+        finished = self.scheduler.update_from_output(
+            sched_output,
+            _make_request_output(req_id, error="worker failed"),
+        )
+
+        assert finished == {req_id}
+        state = self.scheduler.get_request_state(req_id)
+        assert state.status == DiffusionRequestStatus.FINISHED_ERROR
+        assert state.error == "worker failed"
+
+    def test_empty_output_without_error_marks_completed(self) -> None:
+        req_id = self.scheduler.add_request(_make_request("empty"))
+
+        sched_output = self.scheduler.schedule()
+        finished = self.scheduler.update_from_output(sched_output, _make_request_output(req_id))
+
+        assert finished == {req_id}
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
+
+    def test_fifo_single_request_scheduling(self) -> None:
+        req_id_a = self.scheduler.add_request(_make_request("a"))
+        req_id_b = self.scheduler.add_request(_make_request("b"))
+
+        first = self.scheduler.schedule()
+        assert _new_ids(first) == [req_id_a]
+        assert _cached_ids(first) == []
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 1
+
+        # Request A is still running; scheduling again should not pull B.
+        second = self.scheduler.schedule()
+        assert _new_ids(second) == []
+        assert _cached_ids(second) == [req_id_a]
+        assert second.num_running_reqs == 1
+        assert second.num_waiting_reqs == 1
+
+        self.scheduler.update_from_output(first, _make_request_output(req_id_a))
+
+        third = self.scheduler.schedule()
+        assert _new_ids(third) == [req_id_b]
+        assert _cached_ids(third) == []
+        assert third.num_running_reqs == 1
+        assert third.num_waiting_reqs == 0
+
+    def test_abort_request_for_waiting_and_running(self) -> None:
+        req_id_a = self.scheduler.add_request(_make_request("a"))
+        req_id_b = self.scheduler.add_request(_make_request("b"))
+
+        # Abort waiting request.
+        self.scheduler.finish_requests(req_id_b, DiffusionRequestStatus.FINISHED_ABORTED)
+        state_b = self.scheduler.get_request_state(req_id_b)
+        assert state_b.status == DiffusionRequestStatus.FINISHED_ABORTED
+
+        # A should still run normally.
+        output_a = self.scheduler.schedule()
+        assert _new_ids(output_a) == [req_id_a]
+
+        # Abort running request.
+        self.scheduler.finish_requests(req_id_a, DiffusionRequestStatus.FINISHED_ABORTED)
+        state_a = self.scheduler.get_request_state(req_id_a)
+        assert state_a.status == DiffusionRequestStatus.FINISHED_ABORTED
+
+        assert self.scheduler.has_requests() is False
+        assert self.scheduler.schedule().scheduled_req_ids == []
+
+    def test_has_requests_state_transition(self) -> None:
+        assert self.scheduler.has_requests() is False
+
+        req_id = self.scheduler.add_request(_make_request("has"))
+        assert self.scheduler.has_requests() is True
+
+        sched_output = self.scheduler.schedule()
+        assert self.scheduler.has_requests() is True
+
+        self.scheduler.update_from_output(sched_output, _make_request_output(req_id))
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
+        assert self.scheduler.has_requests() is False
+
+    def test_request_id_mapping_lifecycle(self) -> None:
+        request = OmniDiffusionRequest(
+            prompts=["prompt_map_a", "prompt_map_b"],
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+            request_ids=["map-a", "map-b"],
+        )
+
+        sched_req_id = self.scheduler.add_request(request)
+
+        assert self.scheduler.get_sched_req_id("map-a") == sched_req_id
+        assert self.scheduler.get_sched_req_id("map-b") == sched_req_id
+
+        self.scheduler.pop_request_state(sched_req_id)
+
+        assert self.scheduler.get_sched_req_id("map-a") is None
+        assert self.scheduler.get_sched_req_id("map-b") is None
+
+
+class TestDiffusionEngine:
+    def test_add_req_and_wait_for_response_single_path(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = RequestScheduler()
+        engine.scheduler.initialize(Mock())
+        engine.executor = Mock()
+        engine._rpc_lock = threading.Lock()
+
+        request = _make_request("engine")
+        expected = DiffusionOutput(output=None)
+        engine.executor.add_req.return_value = expected
+
+        output = engine.add_req_and_wait_for_response(request)
+
+        assert output is expected
+        engine.executor.add_req.assert_called_once_with(request)
+
+    def test_supports_scheduler_interface_injection(self) -> None:
+        request = _make_request("engine_iface")
+        expected = DiffusionOutput(output=None)
+        scheduler = _StubScheduler(request, expected)
+
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = scheduler
+        engine.executor = Mock()
+        engine.executor.add_req = Mock(return_value=expected)
+        engine._rpc_lock = threading.Lock()
+
+        output = engine.add_req_and_wait_for_response(request)
+
+        assert output is expected
+        engine.executor.add_req.assert_called_once_with(request)
+
+    def test_initializes_injected_scheduler(self) -> None:
+        request = _make_request("init")
+        scheduler = _StubScheduler(request, DiffusionOutput(output=None))
+        od_config = Mock(model_class_name="mock_model")
+        fake_executor_cls = Mock(return_value=Mock())
+
+        with (
+            patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func", return_value=None),
+            patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func", return_value=None),
+            patch("vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class", return_value=fake_executor_cls),
+            patch.object(DiffusionEngine, "_dummy_run", return_value=None),
+        ):
+            DiffusionEngine(od_config, scheduler=scheduler)
+
+        assert scheduler.initialized_with is od_config
+        fake_executor_cls.assert_called_once_with(od_config)
+
+    def test_scheduler_alias_keeps_default_request_scheduler(self) -> None:
+        scheduler = Scheduler()
+        scheduler.initialize(Mock())
+
+        req_id = scheduler.add_request(_make_request("alias"))
+        sched_output = scheduler.schedule()
+        finished = scheduler.update_from_output(sched_output, _make_request_output(req_id))
+
+        assert req_id in finished
+        assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
+
+    def test_dummy_run_raises_on_output_error(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.od_config = Mock(model_class_name="mock_model")
+        engine.pre_process_func = None
+        engine.add_req_and_wait_for_response = Mock(return_value=DiffusionOutput(error="boom"))
+
+        with pytest.raises(RuntimeError, match="Dummy run failed: boom"):
+            engine._dummy_run()

tests/diffusion/test_diffusion_step_pipeline.py

@@ -25,10 +25,9 @@ from vllm_omni.diffusion.ipc import (
     unpack_diffusion_output_shm,
 )
 from vllm_omni.diffusion.sched.interface import (
-    DiffusionRequestState as SchedulerRequestState,
-)
-from vllm_omni.diffusion.sched.interface import (
+    CachedRequestData,
     DiffusionSchedulerOutput,
+    NewRequestData,
 )
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
@@ -224,7 +223,19 @@ def _make_distributed_runner(mode: str, device: torch.device):
 def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_ids=None):
     return DiffusionSchedulerOutput(
         step_id=step_id,
-        req_states=[SchedulerRequestState(sched_req_id=sched_req_id, req=req)],
+        scheduled_new_reqs=[NewRequestData(sched_req_id=sched_req_id, req=req)],
+        scheduled_cached_reqs=CachedRequestData.make_empty(),
+        finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
+        num_running_reqs=1,
+        num_waiting_reqs=0,
+    )
+
+
+def _make_cached_scheduler_output(sched_req_id="req-1", step_id=1, finished_req_ids=None):
+    return DiffusionSchedulerOutput(
+        step_id=step_id,
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=CachedRequestData(sched_req_ids=[sched_req_id]),
         finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
         num_running_reqs=1,
         num_waiting_reqs=0,
@@ -297,7 +308,7 @@ class TestRunner:
         assert first.result is None
         assert "req-1" in runner.state_cache
 
-        second = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=1))
+        second = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
         assert second.req_id == "req-1"
         assert second.step_index == 2
         assert second.finished is True
@@ -369,7 +380,7 @@ class TestWorker:
         worker = object.__new__(DiffusionWorker)
         expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None)
         scheduler_output = SimpleNamespace(
-            req_states=[
+            scheduled_new_reqs=[
                 SimpleNamespace(
                     req=SimpleNamespace(
                         sampling_params=SimpleNamespace(lora_request=None),
@@ -389,7 +400,7 @@ class TestWorker:
     def test_clears_active_lora_before_stepwise_execution(self):
         worker = object.__new__(DiffusionWorker)
         scheduler_output = SimpleNamespace(
-            req_states=[
+            scheduled_new_reqs=[
                 SimpleNamespace(
                     req=SimpleNamespace(
                         sampling_params=SimpleNamespace(lora_request=None),
@@ -413,7 +424,7 @@ class TestWorker:
     def test_rejects_lora_requests_in_step_mode(self):
         worker = object.__new__(DiffusionWorker)
         scheduler_output = SimpleNamespace(
-            req_states=[
+            scheduled_new_reqs=[
                 SimpleNamespace(
                     req=SimpleNamespace(
                         sampling_params=SimpleNamespace(lora_request=object()),

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -9,8 +9,9 @@ import pytest
 import torch
 
 from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
-from vllm_omni.diffusion.scheduler import Scheduler
+from vllm_omni.diffusion.sched import RequestScheduler
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
 
@@ -30,42 +31,43 @@ def _mock_request(tag: str) -> Mock:
     return req
 
 
-def _make_scheduler():
-    """Create a ``Scheduler`` whose *mq* / *result_mq* are backed by
-    plain ``queue.Queue`` objects (thread-safe, no real IPC).
+def _make_executor(num_gpus: int = 1):
+    """Create a ``MultiprocDiffusionExecutor`` without launching workers.
 
-    Returns ``(scheduler, request_queue, result_queue)``.
+    Returns ``(executor, request_queue, result_queue)``.
     """
-    sched = Scheduler()
-    sched.num_workers = 1
-    sched._lock = threading.Lock()
+    od_cfg = Mock()
+    od_cfg.num_gpus = num_gpus
+
+    with patch.object(MultiprocDiffusionExecutor, "_init_executor"):
+        executor = MultiprocDiffusionExecutor(od_cfg)
 
     req_q: queue.Queue = queue.Queue()
     res_q: queue.Queue = queue.Queue()
 
-    mock_mq = Mock()
-    mock_mq.enqueue = req_q.put
+    mock_broadcast_mq = Mock()
+    mock_broadcast_mq.enqueue = req_q.put
 
     mock_rmq = Mock()
-    mock_rmq.dequeue = lambda timeout=None: res_q.get(timeout=timeout if timeout else 10)
+    mock_rmq.dequeue = lambda timeout=None: res_q.get(timeout=timeout if timeout is not None else 10)
 
-    sched.mq = mock_mq
-    sched.result_mq = mock_rmq
-    return sched, req_q, res_q
-
-
-def _make_executor(scheduler):
-    """Create a ``MultiprocDiffusionExecutor`` wired to *scheduler*
-    without launching real worker processes.
-    """
-    od_cfg = Mock()
-    od_cfg.num_gpus = 1
-    with patch.object(MultiprocDiffusionExecutor, "_init_executor"):
-        executor = MultiprocDiffusionExecutor(od_cfg)
-    executor.scheduler = scheduler
+    executor._broadcast_mq = mock_broadcast_mq
+    executor._result_mq = mock_rmq
     executor._closed = False
     executor._processes = []
-    return executor
+    return executor, req_q, res_q
+
+
+def _make_engine(num_gpus: int = 1):
+    """Create a lightweight ``DiffusionEngine`` wired to mocked executor."""
+    executor, req_q, res_q = _make_executor(num_gpus)
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    sched = RequestScheduler()
+    sched.initialize(Mock())
+    engine.scheduler = sched
+    engine.executor = executor
+    engine._rpc_lock = threading.Lock()
+    return engine, executor, req_q, res_q
 
 
 def _start_worker(req_q, res_q, count=2):
@@ -91,8 +93,8 @@ def _start_worker(req_q, res_q, count=2):
     return t
 
 
-def _inject_interleave(scheduler):
-    """Monkey-patch ``scheduler.mq.enqueue`` so that:
+def _inject_interleave(executor):
+    """Monkey-patch ``executor._broadcast_mq.enqueue`` so that:
 
     * The thread named **thread_a** *blocks* after its enqueue until the
       thread named **thread_b** has finished entirely.
@@ -102,7 +104,7 @@ def _inject_interleave(scheduler):
     """
     a_enqueued = threading.Event()
     b_complete = threading.Event()
-    orig_enqueue = scheduler.mq.enqueue  # points to req_q.put
+    orig_enqueue = executor._broadcast_mq.enqueue  # points to req_q.put
 
     def _controlled(item):
         orig_enqueue(item)
@@ -110,7 +112,7 @@ def _inject_interleave(scheduler):
             a_enqueued.set()  # tell B: "A has enqueued"
             b_complete.wait(5)  # block A until B finishes
 
-    scheduler.mq.enqueue = _controlled
+    executor._broadcast_mq.enqueue = _controlled
     return a_enqueued, b_complete
 
 
@@ -118,21 +120,21 @@ def _inject_interleave(scheduler):
 
 
 class TestConcurrentAddReqBug:
-    """Two concurrent ``Scheduler.add_req()`` calls swap results."""
+    """Two concurrent ``add_req_and_wait_for_response()`` calls swap results."""
 
     def test_results_are_correctly_routed(self):
-        sched, req_q, res_q = _make_scheduler()
-        a_enqueued, b_complete = _inject_interleave(sched)
+        engine, executor, req_q, res_q = _make_engine()
+        a_enqueued, b_complete = _inject_interleave(executor)
         wt = _start_worker(req_q, res_q, count=2)
 
         results: dict[str, DiffusionOutput] = {}
 
         def _a():
-            results["A"] = sched.add_req(_mock_request("A"))
+            results["A"] = engine.add_req_and_wait_for_response(_mock_request("A"))
 
         def _b():
             a_enqueued.wait(5)  # wait for A to enqueue
-            results["B"] = sched.add_req(_mock_request("B"))
+            results["B"] = engine.add_req_and_wait_for_response(_mock_request("B"))
             b_complete.set()  # release A
 
         ta = threading.Thread(target=_a, name="thread_a")
@@ -156,15 +158,14 @@ class TestConcurrentCollectiveRpcBug:
     """Two concurrent ``collective_rpc()`` calls swap results."""
 
     def test_results_are_correctly_routed(self):
-        sched, req_q, res_q = _make_scheduler()
-        executor = _make_executor(sched)
-        a_enqueued, b_complete = _inject_interleave(sched)
+        engine, executor, req_q, res_q = _make_engine()
+        a_enqueued, b_complete = _inject_interleave(executor)
         wt = _start_worker(req_q, res_q, count=2)
 
         results: dict[str, object] = {}
 
         def _a():
-            results["A"] = executor.collective_rpc(
+            results["A"] = engine.collective_rpc(
                 "ping",
                 args=("call_A",),
                 unique_reply_rank=0,
@@ -172,7 +173,7 @@ class TestConcurrentCollectiveRpcBug:
 
         def _b():
             a_enqueued.wait(5)
-            results["B"] = executor.collective_rpc(
+            results["B"] = engine.collective_rpc(
                 "ping",
                 args=("call_B",),
                 unique_reply_rank=0,
@@ -198,19 +199,18 @@ class TestConcurrentAddReqVsCollectiveRpcBug:
     """``add_req`` and ``collective_rpc`` running concurrently swap results."""
 
     def test_results_are_correctly_routed(self):
-        sched, req_q, res_q = _make_scheduler()
-        executor = _make_executor(sched)
-        a_enqueued, b_complete = _inject_interleave(sched)
+        engine, executor, req_q, res_q = _make_engine()
+        a_enqueued, b_complete = _inject_interleave(executor)
         wt = _start_worker(req_q, res_q, count=2)
 
         results: dict[str, object] = {}
 
         def _a():  # add_req path
-            results["A"] = sched.add_req(_mock_request("A"))
+            results["A"] = engine.add_req_and_wait_for_response(_mock_request("A"))
 
         def _b():  # collective_rpc path
             a_enqueued.wait(5)
-            results["B"] = executor.collective_rpc(
+            results["B"] = engine.collective_rpc(
                 "ping",
                 args=("call_B",),
                 unique_reply_rank=0,
@@ -241,31 +241,30 @@ class TestSerialOperations:
     """
 
     def test_serial_add_req_returns_correct_result(self):
-        sched, req_q, res_q = _make_scheduler()
+        engine, _, req_q, res_q = _make_engine()
         wt = _start_worker(req_q, res_q, count=1)
 
-        result = sched.add_req(_mock_request("X"))
+        result = engine.add_req_and_wait_for_response(_mock_request("X"))
         wt.join(5)
 
         assert isinstance(result, DiffusionOutput)
         assert result.error == "result_for_X"
 
     def test_serial_add_req_multiple_sequential(self):
-        sched, req_q, res_q = _make_scheduler()
+        engine, _, req_q, res_q = _make_engine()
         wt = _start_worker(req_q, res_q, count=3)
 
         for tag in ("one", "two", "three"):
-            out = sched.add_req(_mock_request(tag))
+            out = engine.add_req_and_wait_for_response(_mock_request(tag))
             assert out.error == f"result_for_{tag}"
 
         wt.join(5)
 
     def test_serial_collective_rpc_single_rank(self):
-        sched, req_q, res_q = _make_scheduler()
-        executor = _make_executor(sched)
+        engine, _, req_q, res_q = _make_engine()
         wt = _start_worker(req_q, res_q, count=1)
 
-        result = executor.collective_rpc(
+        result = engine.collective_rpc(
             "ping",
             args=("Y",),
             unique_reply_rank=0,
@@ -278,27 +277,24 @@ class TestSerialOperations:
         """``collective_rpc`` without *unique_reply_rank* collects
         ``num_gpus`` responses.
         """
-        sched, req_q, res_q = _make_scheduler()
-        executor = _make_executor(sched)
-        executor.od_config.num_gpus = 2
+        engine, _, _, res_q = _make_engine(num_gpus=2)
 
         # Pre-populate two results (simulating two workers replying)
         res_q.put(_tagged_output("rank0"))
         res_q.put(_tagged_output("rank1"))
 
-        results = executor.collective_rpc("ping", args=("multi",))
+        results = engine.collective_rpc("ping", args=("multi",))
 
         assert len(results) == 2
         assert results[0].error == "rank0"
         assert results[1].error == "rank1"
 
     def test_serial_add_req_then_collective_rpc(self):
-        sched, req_q, res_q = _make_scheduler()
-        executor = _make_executor(sched)
+        engine, _, req_q, res_q = _make_engine()
         wt = _start_worker(req_q, res_q, count=2)
 
-        gen_out = sched.add_req(_mock_request("gen"))
-        rpc_out = executor.collective_rpc(
+        gen_out = engine.add_req_and_wait_for_response(_mock_request("gen"))
+        rpc_out = engine.collective_rpc(
             "ping",
             args=("rpc",),
             unique_reply_rank=0,
@@ -310,29 +306,30 @@ class TestSerialOperations:
 
     def test_serial_add_req_error_propagation(self):
         """``add_req`` should raise when the worker reports an error."""
-        sched, _, res_q = _make_scheduler()
+        engine, _, _, res_q = _make_engine()
         # Put an error response directly
         res_q.put({"status": "error", "error": "boom"})
 
-        with pytest.raises(RuntimeError, match="worker error"):
-            sched.add_req(_mock_request("fail"))
+        out = engine.add_req_and_wait_for_response(_mock_request("fail"))
+
+        assert isinstance(out, DiffusionOutput)
+        assert out.error is not None
+        assert "boom" in out.error
 
     def test_serial_collective_rpc_error_propagation(self):
         """``collective_rpc`` should raise when the worker reports an error."""
-        sched, _, res_q = _make_scheduler()
-        executor = _make_executor(sched)
+        engine, _, _, res_q = _make_engine()
         res_q.put({"status": "error", "error": "kaboom"})
 
         with pytest.raises(RuntimeError, match="kaboom"):
-            executor.collective_rpc("bad", unique_reply_rank=0)
+            engine.collective_rpc("bad", unique_reply_rank=0)
 
     def test_collective_rpc_closed_executor_raises(self):
-        sched, _, _ = _make_scheduler()
-        executor = _make_executor(sched)
+        engine, executor, _, _ = _make_engine()
         executor._closed = True
 
         with pytest.raises(RuntimeError, match="closed"):
-            executor.collective_rpc("anything")
+            engine.collective_rpc("anything")
 
 
 # ─────────── timeout regression: RPC must not block on a stalled lock ─────
@@ -340,23 +337,22 @@ class TestSerialOperations:
 
 class TestCollectiveRpcTimeoutWhileLockHeld:
     """``collective_rpc(timeout=...)`` must honour its timeout even when
-    another thread holds ``scheduler._lock`` indefinitely (e.g. a stalled
+    another thread holds ``engine._rpc_lock`` indefinitely (e.g. a stalled
     ``add_req`` waiting on an unresponsive worker).
     """
 
     def test_rpc_times_out_when_lock_held_directly(self):
         """Simplest case: lock is manually held by another thread."""
-        sched, req_q, res_q = _make_scheduler()
-        executor = _make_executor(sched)
+        engine, _, _, _ = _make_engine()
 
         stall_started = threading.Event()
 
         def _hold_lock():
-            sched._lock.acquire()
+            engine._rpc_lock.acquire()
             stall_started.set()
             # Hold the lock far longer than the RPC timeout.
             threading.Event().wait(30)
-            sched._lock.release()
+            engine._rpc_lock.release()
 
         stall_thread = threading.Thread(target=_hold_lock, daemon=True)
         stall_thread.start()
@@ -364,24 +360,23 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
 
         # collective_rpc should raise TimeoutError, NOT block forever.
         with pytest.raises(TimeoutError):
-            executor.collective_rpc("health", timeout=0.5)
+            engine.collective_rpc("health", timeout=0.5)
 
     def test_rpc_times_out_when_add_req_stalled_on_worker(self):
         """Real-world scenario the bot flagged:
 
-        ``add_req`` holds ``_lock`` while blocked on ``result_mq.dequeue()``
-        because the worker never replies.  A concurrent
-        ``collective_rpc(timeout=...)`` must still time out instead of
-        hanging forever waiting for the lock.
+        ``add_req`` holds ``_rpc_lock`` while blocked on
+        ``executor._result_mq.dequeue()`` because the worker never replies.
+        A concurrent ``collective_rpc(timeout=...)`` must still time out
+        instead of hanging forever waiting for the lock.
         """
-        sched, req_q, res_q = _make_scheduler()
-        executor = _make_executor(sched)
+        engine, executor, _, _ = _make_engine()
 
         add_req_blocked = threading.Event()
 
         # Patch dequeue: signal once entered, then block indefinitely
         # (simulates a worker that never sends a result).
-        orig_dequeue = sched.result_mq.dequeue
+        orig_dequeue = executor._result_mq.dequeue
 
         def _hanging_dequeue(timeout=None):
             add_req_blocked.set()
@@ -389,13 +384,13 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
             threading.Event().wait(30)
             return orig_dequeue(timeout=timeout)
 
-        sched.result_mq.dequeue = _hanging_dequeue
+        executor._result_mq.dequeue = _hanging_dequeue
 
         # Thread running add_req — acquires the lock, enqueues, then
         # blocks on dequeue forever (worker hang).
         def _stalled_add_req():
             try:
-                sched.add_req(_mock_request("stalled"))
+                engine.add_req_and_wait_for_response(_mock_request("stalled"))
             except Exception:
                 pass
 
@@ -407,23 +402,19 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
 
         # collective_rpc should time out at lock acquisition, not hang.
         with pytest.raises(TimeoutError):
-            executor.collective_rpc("health_check", timeout=0.5)
+            engine.collective_rpc("health_check", timeout=0.5)
 
     def test_rpc_without_timeout_still_waits_for_lock(self):
         """When no timeout is given, ``collective_rpc`` should still wait
         for the lock (blocking) — existing behaviour preserved.
         """
-        sched, req_q, res_q = _make_scheduler()
-        executor = _make_executor(sched)
-
-        lock_released = threading.Event()
+        engine, _, _, res_q = _make_engine()
 
         def _hold_and_release():
-            sched._lock.acquire()
+            engine._rpc_lock.acquire()
             # Hold for a short time then release.
             threading.Event().wait(0.3)
-            sched._lock.release()
-            lock_released.set()
+            engine._rpc_lock.release()
 
         # Pre-populate a result so collective_rpc succeeds after lock.
         res_q.put(_tagged_output("ok"))
@@ -431,8 +422,8 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
         t = threading.Thread(target=_hold_and_release, daemon=True)
         t.start()
 
-        # No timeout → should block until lock is released, then succeed.
-        result = executor.collective_rpc(
+        # No timeout -> should block until lock is released, then succeed.
+        result = engine.collective_rpc(
             "ping",
             args=("wait",),
             unique_reply_rank=0,