# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression tests for the busy-loop-based RPC routing in DiffusionEngine.

These tests guard the invariants established when ``collective_rpc`` was
moved off ``_rpc_lock`` mutual exclusion and onto a queue drained by the
engine's busy-loop thread:

* Only the busy-loop thread ever calls ``executor.collective_rpc`` after
  the loop starts.
* Concurrent ``engine.collective_rpc`` and per-request ``execute_fn``
  invocations never overlap on the executor.
* Results are routed back to the correct caller (sync, async, and
  per-request paths).
* Pending RPCs are failed cleanly on shutdown.
* The bootstrap path (busy loop not yet started) calls the executor
  directly so ``_dummy_run`` keeps working.
"""

from __future__ import annotations

import asyncio
import queue
import threading
import time
from dataclasses import fields as _dc_fields
from types import SimpleNamespace
from typing import Any

import pytest

from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode, _RpcTask
from vllm_omni.diffusion.sched import RequestScheduler
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 ``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, pytest.mark.core_model]


# ───────────────────────────────────────── helpers ─────────────────────────


class _ConcurrencyTrackingExecutor:
    """Fake executor that records every ``collective_rpc`` invocation and
    flags any overlap. Used to assert structural serialization.
    """

    def __init__(self, rpc_delay: float = 0.0):
        self._active = 0
        self._lock = threading.Lock()
        self.max_concurrent = 0
        self.calls: list[dict[str, Any]] = []
        self.thread_ids: set[int] = set()
        self.rpc_delay = rpc_delay
        self.is_failed = False
        self._closed = False
        self.od_config = SimpleNamespace()
        # Pause tests: a call whose gate is registered (by request id, or by
        # method name for ``synchronize_device``) blocks inside the executor
        # until the gate is set; ``version`` tags each call so a test can tell
        # which "weights" a request ran on.
        self.gates: dict[str, threading.Event] = {}
        self.version = "v1"

    def collective_rpc(
        self,
        method: str,
        timeout: float | None = None,
        args: tuple = (),
        kwargs: dict | None = None,
        unique_reply_rank: int | None = None,
        exec_all_ranks: bool = False,
    ) -> Any:
        with self._lock:
            self._active += 1
            self.max_concurrent = max(self.max_concurrent, self._active)
            self.thread_ids.add(threading.get_ident())
            self.calls.append(
                {
                    "method": method,
                    "args": args,
                    "kwargs": kwargs,
                    "unique_reply_rank": unique_reply_rank,
                    "thread": threading.get_ident(),
                    "version": self.version,
                }
            )
        try:
            if self.rpc_delay:
                time.sleep(self.rpc_delay)
            if method == "synchronize_device":
                gate = self.gates.get(method)
                if gate is not None:
                    gate.wait(timeout=5.0)
                return [None]
            # Distinguish per-request execution from raw RPC by method name.
            if method in {"execute_model", "execute_stepwise", "generate"}:
                # args[0] is the request-like object for execute_model.
                req = args[0] if args else None
                tag = req.request_id if req is not None and hasattr(req, "request_id") else "unknown"
                gate = self.gates.get(tag)
                if gate is not None:
                    gate.wait(timeout=5.0)
                return DiffusionOutput(error=f"result_for_{tag}", finished=True)
            tag = args[0] if args else method
            return DiffusionOutput(error=f"rpc_result_for_{tag}")
        finally:
            with self._lock:
                self._active -= 1

    def execute_request(self, scheduler_output) -> RunnerOutput:
        # Mimic the real MultiprocDiffusionExecutor.execute_request: it
        # forwards a single request through collective_rpc.
        new_req = scheduler_output.scheduled_new_reqs[0]
        req = new_req.req
        result = self.collective_rpc(
            "execute_model",
            args=(req, self.od_config),
            unique_reply_rank=0,
            exec_all_ranks=True,
        )
        return RunnerOutput(
            request_id=req.request_id,
            step_index=None,
            finished=True,
            result=result,
        )

    def shutdown(self) -> None:
        self._closed = True


def _make_request(tag: str):
    sampling_params = dict(_SAMPLING_KEY_DEFAULTS)
    sampling_params["num_inference_steps"] = 1
    sampling_params["extra_args"] = {}
    return SimpleNamespace(
        request_id=tag,
        prompt=f"prompt_{tag}",
        sampling_params=SimpleNamespace(**sampling_params),
        diffusion_kv_requests=None,
    )


def _make_engine_with_loop(
    loop: asyncio.AbstractEventLoop | None,
    rpc_delay: float = 0.0,
    *,
    scheduler: RequestScheduler | None = None,
    start_loop: bool = True,
):
    """Construct a ``DiffusionEngine`` skeleton with a real busy loop.

    The engine is wired with a fake executor that asserts no concurrent
    calls and a real ``RequestScheduler``. ``start_loop=False`` leaves the
    engine in its bootstrap state (no busy-loop thread).
    """
    engine = DiffusionEngine.__new__(DiffusionEngine)
    engine._closed = False
    engine.od_config = SimpleNamespace(streaming_output=False)
    engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay)

    sched = scheduler
    if sched is None:
        sched = RequestScheduler()
        sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0))
    engine.scheduler = sched
    engine.step_execution = False
    engine.supports_request_batch = False
    engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH
    engine.execute_fn = engine.executor.execute_request

    engine._rpc_lock = threading.RLock()
    engine._cv = threading.Condition(engine._rpc_lock)
    engine._out_streams = {}
    engine._closed = False
    engine.abort_queue = queue.Queue()
    engine._rpc_queue = queue.Queue()
    engine._scheduling_paused = False

    engine.main_loop = loop
    engine.stop_event = threading.Event()
    if not start_loop:
        engine._loop_started = False
        engine.worker_thread = None
        return engine
    engine._loop_started = True
    engine.worker_thread = threading.Thread(target=engine._busy_loop, daemon=True)
    engine.worker_thread.start()
    return engine


def _stop_engine(engine: DiffusionEngine) -> None:
    if engine.worker_thread is None:
        return
    for gate in engine.executor.gates.values():
        gate.set()
    with engine._cv:
        engine.stop_event.set()
        engine._cv.notify_all()
    # Race-proof shutdown for the test: drain any RPCs still queued and
    # fail them with the documented shutdown error before the busy loop
    # has a chance to pick them up after its current in-flight call
    # returns. The engine's own ``_fail_pending_rpcs`` then has nothing
    # left to do.
    while True:
        try:
            task = engine._rpc_queue.get_nowait()
        except queue.Empty:
            break
        if not task.future.done():
            task.future.set_exception(RuntimeError("DiffusionEngine is shutting down."))
    engine.worker_thread.join(timeout=5)
    assert not engine.worker_thread.is_alive(), "Busy loop thread did not stop"


async def _consume_final_output(generator):
    final_output = None
    async for output in generator:
        final_output = output
    if final_output is None:
        raise RuntimeError("Diffusion execution finished without output.")
    return final_output


# ─────────────────────── single-thread invariant ───────────────────────────


@pytest.mark.asyncio
async def test_executor_only_called_from_busy_loop_thread():
    """All executor calls — both per-request and raw RPC — must come from
    the busy-loop thread, never from a caller's thread."""
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop)
    busy_tid = engine.worker_thread.ident
    try:
        # Per-request path
        await _consume_final_output(engine.async_add_req_and_stream_response(_make_request("req1")))
        # Raw RPC path (sync from a worker thread)
        result = await asyncio.to_thread(engine.collective_rpc, "ping", args=("a",), unique_reply_rank=0)
        assert result.error == "rpc_result_for_a"
        # Raw RPC path (async)
        result_async = await engine.async_collective_rpc("ping", args=("b",), unique_reply_rank=0)
        assert result_async.error == "rpc_result_for_b"
    finally:
        _stop_engine(engine)

    # Every recorded call ran on the busy-loop thread.
    assert engine.executor.thread_ids == {busy_tid}, (
        f"executor.collective_rpc must run only on the busy-loop thread "
        f"(expected {{{busy_tid}}}, got {engine.executor.thread_ids})"
    )


@pytest.mark.asyncio
async def test_executor_calls_never_overlap_under_load():
    """Stress: many concurrent ``collective_rpc`` callers + a request must
    never produce overlapping executor calls."""
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop, rpc_delay=0.005)
    try:

        async def _rpc(i: int):
            return await engine.async_collective_rpc("ping", args=(f"x{i}",), unique_reply_rank=0)

        async def _request(i: int):
            return await _consume_final_output(engine.async_add_req_and_stream_response(_make_request(f"r{i}")))

        tasks = [_rpc(i) for i in range(20)] + [_request(i) for i in range(5)]
        results = await asyncio.gather(*tasks)
    finally:
        _stop_engine(engine)

    assert engine.executor.max_concurrent == 1, (
        f"Detected concurrent executor calls (max_concurrent={engine.executor.max_concurrent})"
    )
    # All 25 calls should have a result.
    assert len(results) == 25
    assert all(r is not None for r in results)


# ─────────────────────────── result routing ────────────────────────────────


@pytest.mark.asyncio
async def test_collective_rpc_results_routed_to_correct_caller():
    """Under concurrent calls with distinct args, each caller must get its
    own result back — not another caller's."""
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop, rpc_delay=0.01)
    try:

        async def _call(tag: str):
            res = await engine.async_collective_rpc("ping", args=(tag,), unique_reply_rank=0)
            return tag, res

        tags = [f"t{i}" for i in range(15)]
        results = await asyncio.gather(*[_call(t) for t in tags])
    finally:
        _stop_engine(engine)

    for tag, res in results:
        assert res.error == f"rpc_result_for_{tag}", f"caller for {tag!r} received {res.error!r}"


@pytest.mark.asyncio
async def test_sync_collective_rpc_from_worker_thread():
    """Sync ``collective_rpc`` from non-event-loop threads still receives
    the correct result via the busy loop."""
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop)
    try:
        results: list[Any] = [None] * 8

        def _worker(idx: int):
            results[idx] = engine.collective_rpc("ping", args=(f"s{idx}",), unique_reply_rank=0)

        threads = [threading.Thread(target=_worker, args=(i,)) for i in range(8)]
        for t in threads:
            t.start()
        for t in threads:
            t.join(5)
            assert not t.is_alive()
    finally:
        _stop_engine(engine)

    for i, r in enumerate(results):
        assert r is not None
        assert r.error == f"rpc_result_for_s{i}"


# ─────────────────────────── timeout behaviour ─────────────────────────────


@pytest.mark.asyncio
async def test_collective_rpc_times_out_when_busy_loop_busy():
    """If the busy loop is occupied, ``collective_rpc(timeout=...)`` must
    raise ``TimeoutError`` rather than block indefinitely."""
    loop = asyncio.get_running_loop()
    # 2-second per-RPC delay simulates a busy worker.
    engine = _make_engine_with_loop(loop, rpc_delay=2.0)
    try:
        # Kick off one slow RPC to occupy the busy loop.
        slow = asyncio.create_task(engine.async_collective_rpc("slow", args=("s",), unique_reply_rank=0))
        # Yield so the slow task is enqueued.
        await asyncio.sleep(0.05)

        with pytest.raises(TimeoutError):
            await asyncio.to_thread(
                engine.collective_rpc,
                "ping",
                args=("x",),
                unique_reply_rank=0,
                timeout=0.2,
            )

        # The slow call should still complete normally.
        result = await slow
        assert result.error == "rpc_result_for_s"
    finally:
        _stop_engine(engine)


@pytest.mark.asyncio
async def test_async_collective_rpc_times_out_when_busy_loop_busy():
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop, rpc_delay=2.0)
    try:
        slow = asyncio.create_task(engine.async_collective_rpc("slow", args=("s",), unique_reply_rank=0))
        await asyncio.sleep(0.05)

        with pytest.raises(TimeoutError):
            await engine.async_collective_rpc("ping", args=("x",), unique_reply_rank=0, timeout=0.2)

        result = await slow
        assert result.error == "rpc_result_for_s"
    finally:
        _stop_engine(engine)


# ─────────────────────────── shutdown handling ─────────────────────────────


@pytest.mark.asyncio
async def test_pending_rpcs_failed_on_shutdown():
    """Shutdown must fail any RPCs still queued so callers don't hang."""
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop, rpc_delay=0.5)
    try:
        # Start a slow RPC that occupies the busy loop, then queue more.
        in_flight = asyncio.create_task(engine.async_collective_rpc("slow", args=("s",), unique_reply_rank=0))
        await asyncio.sleep(0.05)

        pending = [
            asyncio.create_task(engine.async_collective_rpc("ping", args=(f"p{i}",), unique_reply_rank=0))
            for i in range(3)
        ]
        # Give them a moment to enqueue.
        await asyncio.sleep(0.05)
    finally:
        _stop_engine(engine)

    # The in-flight one finishes normally; pending ones should fail.
    assert (await in_flight).error == "rpc_result_for_s"
    for t in pending:
        with pytest.raises(RuntimeError, match="shutting down"):
            await t


# ─────────────────────────── bootstrap path ────────────────────────────────


def test_collective_rpc_before_loop_starts_calls_executor_directly():
    """When ``_loop_started`` is False (e.g. inside ``_dummy_run`` during
    ``__init__``), ``collective_rpc`` must call the executor synchronously
    on the caller's thread without enqueueing.
    """
    engine = DiffusionEngine.__new__(DiffusionEngine)
    engine._closed = False
    engine._loop_started = False
    engine._rpc_lock = threading.RLock()
    engine._cv = threading.Condition(engine._rpc_lock)
    engine.executor = _ConcurrencyTrackingExecutor()

    caller_tid = threading.get_ident()
    result = engine.collective_rpc("ping", args=("boot",), unique_reply_rank=0)

    assert result.error == "rpc_result_for_boot"
    # Critically: ran on the caller's thread, not a busy-loop thread.
    assert engine.executor.thread_ids == {caller_tid}


def test_collective_rpc_before_loop_starts_serializes_concurrent_callers():
    """Regression: between ``__init__`` returning and the first async
    request starting the busy loop, multiple threads may call
    ``collective_rpc`` concurrently. The pre-loop fast-path must
    serialize them so they cannot race on the shared executor MQ pair.
    """
    engine = DiffusionEngine.__new__(DiffusionEngine)
    engine._closed = False
    engine._loop_started = False
    engine._rpc_lock = threading.RLock()
    engine._cv = threading.Condition(engine._rpc_lock)
    # Non-trivial delay forces overlap if the lock is missing.
    engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=0.02)

    n = 8
    results: list[Any] = [None] * n

    def _call(i: int) -> None:
        results[i] = engine.collective_rpc("ping", args=(f"x{i}",), unique_reply_rank=0)

    threads = [threading.Thread(target=_call, args=(i,)) for i in range(n)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    assert engine.executor.max_concurrent == 1, "Pre-loop collective_rpc must serialize concurrent callers"
    for i, r in enumerate(results):
        assert r is not None and r.error == f"rpc_result_for_x{i}"


@pytest.mark.asyncio
async def test_cancelled_future_does_not_kill_busy_loop():
    """Regression: if a queued RPC future is cancelled while the executor
    call is in flight, ``_process_rpc_queue`` must not raise
    ``InvalidStateError`` when trying to set the result/exception. Doing
    so would crash the busy-loop thread and stall all later requests.
    """
    loop = asyncio.get_running_loop()
    # Slow executor so we can cancel mid-flight.
    engine = _make_engine_with_loop(loop, rpc_delay=0.3)
    try:
        # Submit an RPC and immediately cancel its future to simulate
        # a sync timeout / asyncio cancellation racing the worker.
        with pytest.raises(TimeoutError):
            await asyncio.to_thread(
                engine.collective_rpc,
                "slow",
                args=("cancelme",),
                unique_reply_rank=0,
                timeout=0.05,
            )
        # Give the busy loop time to finish the in-flight slow call and
        # attempt to set state on the cancelled future.
        await asyncio.sleep(0.5)

        # If the busy loop crashed, this follow-up RPC would hang /
        # never complete. Bound it with a timeout to fail fast.
        result = await asyncio.wait_for(
            engine.async_collective_rpc("ping", args=("after",), unique_reply_rank=0),
            timeout=3.0,
        )
        assert result.error == "rpc_result_for_after"
        assert engine.worker_thread.is_alive()
    finally:
        _stop_engine(engine)


@pytest.mark.asyncio
@pytest.mark.parametrize("error_type", [ValueError, RuntimeError])
async def test_native_kv_reservation_error_wakes_stream_without_killing_busy_loop(error_type):
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop)

    def fail_reservation(*args, **kwargs):
        raise error_type("native allocation bug")

    engine.scheduler._diffusion_kv_manager = SimpleNamespace(
        has_request=lambda request_id: False,
        reserve_request=fail_reservation,
        free_request=lambda request_id: None,
    )
    request = _make_request("kv-error")
    request.diffusion_kv_requests = (object(),)
    try:
        response_stream = engine.async_add_req_and_stream_response(request)
        output_queue = engine._out_streams[request.request_id]
        for _ in range(300):
            if not output_queue.empty():
                break
            await asyncio.sleep(0.01)
        assert not output_queue.empty(), "terminal KV allocation error was not delivered to the request stream"

        output = await anext(response_stream)
        assert output.error == "native allocation bug"
        assert output.finished
        assert engine.worker_thread.is_alive()
    finally:
        _stop_engine(engine)


# ─────────────────────────── _RpcTask basics ───────────────────────────────


def test_rpc_task_default_future_is_unique_per_instance():
    """Regression: ``_RpcTask.future`` must default to a *new* Future per
    instance, not a shared one (would cross-resolve all callers).
    """
    a = _RpcTask(method="m", args=(), kwargs=None, deadline=None, unique_reply_rank=0)
    b = _RpcTask(method="m", args=(), kwargs=None, deadline=None, unique_reply_rank=0)
    assert a.future is not b.future
    a.future.set_result("a")
    assert not b.future.done()


# ──────────────── busy-loop drains RPCs without scheduler work ─────────────


@pytest.mark.asyncio
async def test_busy_loop_handles_rpc_without_pending_requests():
    """The wait predicate must wake on RPC submissions even when the
    scheduler queue is empty (regression: previously the loop only woke
    on ``has_requests()``).
    """
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop)
    try:
        # No request has been added; only RPC.
        result = await asyncio.wait_for(
            engine.async_collective_rpc("ping", args=("alone",), unique_reply_rank=0),
            timeout=3.0,
        )
        assert result.error == "rpc_result_for_alone"
    finally:
        _stop_engine(engine)


# ──────────────── interleaved RPC + request integrity ──────────────────────


@pytest.mark.asyncio
async def test_rpc_and_request_results_do_not_swap():
    """A concurrent RPC and request execution must each receive their own
    result (regression for the original race the refactor fixed).
    """
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop, rpc_delay=0.02)
    try:
        rpc_task = asyncio.create_task(engine.async_collective_rpc("ping", args=("rpc1",), unique_reply_rank=0))
        req_task = asyncio.create_task(
            _consume_final_output(engine.async_add_req_and_stream_response(_make_request("req1")))
        )
        rpc_res, req_res = await asyncio.gather(rpc_task, req_task)
    finally:
        _stop_engine(engine)

    assert rpc_res.error == "rpc_result_for_rpc1"
    assert req_res.error == "result_for_req1"
    # Belt and braces: never overlapped on the executor.
    assert engine.executor.max_concurrent == 1


# ───────────────────────── scheduler pause (keep) ─────────────────────────


class _GatedAdmissionScheduler(RequestScheduler):
    """Batching wait that only ends when the test opens admission."""

    def __init__(self) -> None:
        super().__init__()
        self.admission_open = threading.Event()

    def get_admission_wait_decision(self, *, now: float, dp_concurrent: bool = False):
        decision = super().get_admission_wait_decision(now=now, dp_concurrent=dp_concurrent)
        return decision.__class__(should_wait=True, deadline=None, max_batch=self.max_num_running_reqs)

    def should_end_admission_wait(self, decision, *, now: float, stable_since: float) -> bool:
        return self.admission_open.is_set()


def _pause(engine: DiffusionEngine, timeout: float | None = None):
    return engine.async_collective_rpc("pause_scheduler", timeout=timeout, kwargs={"mode": "keep"})


def _resume(engine: DiffusionEngine):
    return engine.async_collective_rpc("resume_scheduler")


def _methods(engine: DiffusionEngine) -> list[str]:
    return [call["method"] for call in engine.executor.calls]


def _executed(engine: DiffusionEngine) -> list[tuple[str, str]]:
    return [
        (call["args"][0].request_id, call["version"])
        for call in engine.executor.calls
        if call["method"] == "execute_model"
    ]


def _submit(engine: DiffusionEngine, tag: str) -> asyncio.Task:
    return asyncio.create_task(_consume_final_output(engine.async_add_req_and_stream_response(_make_request(tag))))


async def _wait_until(predicate, timeout: float = 3.0) -> None:
    deadline = time.monotonic() + timeout
    while not predicate():
        if time.monotonic() > deadline:
            raise AssertionError("condition not reached in time")
        await asyncio.sleep(0.005)


@pytest.mark.asyncio
async def test_keep_pause_retains_waiting_request_and_serves_control_rpcs():
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop)
    try:
        assert await _pause(engine) is None
        queued = _submit(engine, "queued")
        result = await engine.async_collective_rpc("list_loras", unique_reply_rank=0)
        assert result.error == "rpc_result_for_list_loras"
        await asyncio.sleep(0.05)
        assert engine.scheduler.num_waiting_requests() == 1
        assert _executed(engine) == []
        assert not queued.done()

        await _resume(engine)
        assert (await asyncio.wait_for(queued, 3.0)).error == "result_for_queued"
        assert _methods(engine) == ["synchronize_device", "list_loras", "execute_model"]
    finally:
        _stop_engine(engine)


@pytest.mark.asyncio
async def test_keep_pause_acks_after_running_batch_and_freezes_queue():
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop)
    gate = engine.executor.gates["A"] = threading.Event()
    try:
        a = _submit(engine, "A")
        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
        b = _submit(engine, "B")
        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)

        pause = asyncio.create_task(_pause(engine))
        await asyncio.sleep(0.1)
        assert not pause.done()

        gate.set()
        await asyncio.wait_for(pause, 3.0)
        assert _methods(engine) == ["execute_model", "synchronize_device"]
        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
        assert engine.scheduler.num_waiting_requests() == 1
        assert not b.done()

        engine.executor.version = "v2"
        await _resume(engine)
        assert (await asyncio.wait_for(b, 3.0)).error == "result_for_B"
        assert _executed(engine) == [("A", "v1"), ("B", "v2")]
    finally:
        _stop_engine(engine)


@pytest.mark.asyncio
async def test_keep_pause_ack_does_not_release_later_rpcs_before_the_batch_is_emitted():
    """A pause ACK must not open the queue to RPCs that touch device memory
    (sleep) before the batch that just ran has been delivered.
    """
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop)
    barrier_gate = engine.executor.gates["synchronize_device"] = threading.Event()
    methods_at_emit: list[list[str]] = []
    emit_outputs = engine._emit_outputs

    def _recording_emit(*args, **kwargs):
        methods_at_emit.append(_methods(engine))
        return emit_outputs(*args, **kwargs)

    engine._emit_outputs = _recording_emit
    try:
        gate = engine.executor.gates["A"] = threading.Event()
        a = _submit(engine, "A")
        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
        pause = asyncio.create_task(_pause(engine))
        await _wait_until(lambda: engine._scheduling_paused)
        gate.set()
        # The barrier holds the busy loop inside the post-execute RPC drain,
        # so the sleep is queued before the pause ACK is published.
        await _wait_until(lambda: "synchronize_device" in _methods(engine))
        sleep_rpc = asyncio.create_task(engine.async_collective_rpc("handle_sleep_task", args=("t",)))
        await _wait_until(lambda: not engine._rpc_queue.empty())
        barrier_gate.set()

        await asyncio.wait_for(pause, 3.0)
        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
        assert (await asyncio.wait_for(sleep_rpc, 3.0)).error == "rpc_result_for_t"
    finally:
        _stop_engine(engine)

    assert _methods(engine) == ["execute_model", "synchronize_device", "handle_sleep_task"]
    assert methods_at_emit and "handle_sleep_task" not in methods_at_emit[0]


@pytest.mark.asyncio
async def test_timed_out_pause_still_ends_the_drain_before_the_batch_is_emitted():
    """A pause whose caller already gave up is still a drain barrier: its
    cancelled task must not let a later RPC run before the batch is emitted.
    """
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop)
    methods_at_emit: list[list[str]] = []
    emit_outputs = engine._emit_outputs

    def _recording_emit(*args, **kwargs):
        methods_at_emit.append(_methods(engine))
        return emit_outputs(*args, **kwargs)

    engine._emit_outputs = _recording_emit
    try:
        gate = engine.executor.gates["A"] = threading.Event()
        a = _submit(engine, "A")
        await _wait_until(lambda: _executed(engine) == [("A", "v1")])

        # The caller times out while A is still executing: the pause future is
        # cancelled, but its task stays queued ahead of the sleep.
        with pytest.raises(TimeoutError):
            await _pause(engine, timeout=0.1)
        sleep_rpc = asyncio.create_task(engine.async_collective_rpc("handle_sleep_task", args=("t",)))
        await _wait_until(lambda: engine._rpc_queue.qsize() == 2)
        gate.set()

        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
        assert (await asyncio.wait_for(sleep_rpc, 3.0)).error == "rpc_result_for_t"
    finally:
        _stop_engine(engine)

    assert _methods(engine) == ["execute_model", "handle_sleep_task"]
    assert methods_at_emit and "handle_sleep_task" not in methods_at_emit[0]


@pytest.mark.asyncio
async def test_pause_rejects_unsupported_modes_without_closing_the_gate():
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop)
    try:
        for kwargs in ({"mode": "abort"}, {"mode": "wait"}, {}):
            with pytest.raises(ValueError):
                await engine.async_collective_rpc("pause_scheduler", kwargs=kwargs)
        engine.execution_mode = DiffusionExecutionMode.STEP_BATCH
        with pytest.raises(NotImplementedError):
            await _pause(engine)
        engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH

        assert (await asyncio.wait_for(_submit(engine, "still-runs"), 3.0)).error == "result_for_still-runs"
        assert _methods(engine) == ["execute_model"]
    finally:
        _stop_engine(engine)


@pytest.mark.asyncio
async def test_keep_pause_during_batching_wait_holds_the_batch():
    loop = asyncio.get_running_loop()
    sched = _GatedAdmissionScheduler()
    sched.initialize(SimpleNamespace(max_num_seqs=4, request_batch_max_wait_ms=1000.0))
    engine = _make_engine_with_loop(loop, scheduler=sched)
    try:
        waiting = _submit(engine, "W")
        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
        await asyncio.sleep(0.05)

        await asyncio.wait_for(_pause(engine), 3.0)
        assert _methods(engine) == ["synchronize_device"]
        assert engine.scheduler.num_waiting_requests() == 1

        await _resume(engine)
        await asyncio.sleep(0.05)
        assert _executed(engine) == []
        sched.admission_open.set()
        assert (await asyncio.wait_for(waiting, 3.0)).error == "result_for_W"
    finally:
        sched.admission_open.set()
        _stop_engine(engine)


@pytest.mark.asyncio
async def test_pause_timeout_leaves_gate_closed_until_explicit_resume():
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop)
    gate = engine.executor.gates["A"] = threading.Event()
    try:
        a = _submit(engine, "A")
        await _wait_until(lambda: _executed(engine) == [("A", "v1")])
        b = _submit(engine, "B")
        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)

        with pytest.raises(TimeoutError):
            await _pause(engine, timeout=0.1)
        gate.set()
        assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A"
        await asyncio.sleep(0.1)
        assert _executed(engine) == [("A", "v1")]
        assert not b.done()

        await _resume(engine)
        assert (await asyncio.wait_for(b, 3.0)).error == "result_for_B"
        assert "synchronize_device" not in _methods(engine)
    finally:
        _stop_engine(engine)


@pytest.mark.asyncio
async def test_repeated_pause_and_resume_are_idempotent():
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop)
    try:
        await _resume(engine)
        await _pause(engine)
        await _pause(engine)
        queued = _submit(engine, "Q")
        await asyncio.sleep(0.05)
        assert _executed(engine) == []
        await _resume(engine)
        await _resume(engine)
        assert (await asyncio.wait_for(queued, 3.0)).error == "result_for_Q"
        assert _methods(engine) == ["synchronize_device", "synchronize_device", "execute_model"]
    finally:
        _stop_engine(engine)


@pytest.mark.asyncio
async def test_abort_of_waiting_request_while_paused_delivers_terminal():
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop)
    try:
        await _pause(engine)
        b = _submit(engine, "B")
        await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1)
        engine.abort("B")

        result = await asyncio.wait_for(b, 3.0)
        assert result.aborted is True
        assert engine.scheduler.get_request_state("B") is None
        assert _executed(engine) == []

        await _resume(engine)
        assert (await asyncio.wait_for(_submit(engine, "C"), 3.0)).error == "result_for_C"
        assert _executed(engine) == [("C", "v1")]
        assert engine.worker_thread.is_alive()
    finally:
        _stop_engine(engine)


def test_pause_and_resume_before_loop_starts_use_engine_local_path():
    engine = _make_engine_with_loop(None, start_loop=False)

    assert engine.collective_rpc("pause_scheduler", kwargs={"mode": "keep"}) is None
    assert _methods(engine) == ["synchronize_device"]

    assert engine.collective_rpc("resume_scheduler") is None
    assert engine.add_req_and_wait_for_response(_make_request("X")).error == "result_for_X"
    assert _executed(engine) == [("X", "v1")]


@pytest.mark.asyncio
async def test_requests_without_pause_never_trigger_the_barrier():
    loop = asyncio.get_running_loop()
    engine = _make_engine_with_loop(loop)
    try:
        first = await asyncio.wait_for(_submit(engine, "r1"), 3.0)
        await engine.async_collective_rpc("ping", args=("p",), unique_reply_rank=0)
        second = await asyncio.wait_for(_submit(engine, "r2"), 3.0)
    finally:
        _stop_engine(engine)

    assert (first.error, second.error) == ("result_for_r1", "result_for_r2")
    assert _methods(engine) == ["execute_model", "ping", "execute_model"]