# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import os
import socket

import pytest
import torch
from vllm.model_executor.models.utils import PPMissingLayer, make_empty_intermediate_tensors_factory, make_layers
from vllm.sequence import IntermediateTensors
from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors

import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.parallel_state import (
    destroy_distributed_env,
    get_classifier_free_guidance_rank,
    get_pp_group,
    init_distributed_environment,
    initialize_model_parallel,
)
from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
from vllm_omni.platforms import current_omni_platform

pytestmark = [pytest.mark.parallel]


def _find_free_port() -> str:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("127.0.0.1", 0))
        return str(s.getsockname()[1])


def update_environment_variables(envs_dict: dict[str, str]) -> None:
    for k, v in envs_dict.items():
        os.environ[k] = v


# ---------------------------------------------------------------------------
# Shared stubs used by both unit and distributed tests
# ---------------------------------------------------------------------------


class FakeWork:
    """Drop-in for torch.distributed.Work that records whether wait() was called."""

    def __init__(self):
        self.waited = False

    def wait(self):
        self.waited = True


class SimpleScheduler:
    """Minimal diffusion-step scheduler: latents -= 0.1 * noise_pred."""

    def step(self, noise_pred: torch.Tensor, t, latents: torch.Tensor, return_dict: bool = False):
        return (latents - 0.1 * noise_pred,)


class FakeVAE:
    def __init__(self, distributed_enabled: bool = False):
        self.calls = 0
        self.distributed_enabled = distributed_enabled

    def decode(self, z: torch.Tensor):
        """Original decode docstring."""
        self.calls += 1
        return (z + 1,)

    def is_distributed_enabled(self) -> bool:
        return self.distributed_enabled


class MockPipelineParallel(PipelineParallelMixin, CFGParallelMixin):
    """Minimal pipeline used to exercise PipelineParallelMixin.

    Uses vLLM's ``make_layers`` for layer partitioning — the same utility used
    by real DiT models — so the PP layer-split logic is exercised faithfully.

    Each layer's weights are seeded by ``seed + layer_index`` so that layer ``i``
    is initialized identically on every rank regardless of which ranks are active,
    allowing the distributed output to be compared against the single-GPU baseline.

    Args:
        num_layers: Total number of Linear layers.
        dim:        Input / hidden dimension.
        seed:       Base RNG seed; layer ``i`` uses ``seed + i``.
        device:     Target device for layer weights (default: CPU).
        dtype:      Target dtype for layer weights (default: float32).
    """

    def __init__(
        self,
        num_layers: int = 4,
        dim: int = 64,
        seed: int = 42,
        device: torch.device | None = None,
        dtype: torch.dtype = torch.float32,
    ):
        self.start_layer, self.end_layer, self.layers = make_layers(
            num_layers,
            lambda prefix: torch.nn.Linear(dim, dim, bias=False),
            prefix="layers",
        )

        for i, layer in enumerate(self.layers):
            if not isinstance(layer, PPMissingLayer):
                torch.manual_seed(seed + i)
                torch.nn.init.normal_(layer.weight, mean=0.0, std=0.02)

        self.layers.to(device=device, dtype=dtype)

        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states"], dim)
        self.scheduler = SimpleScheduler()

    def predict_noise(self, x=None, intermediate_tensors=None, **_kwargs) -> torch.Tensor | IntermediateTensors:
        """Layer-split forward pass.

        * First PP rank: uses ``x`` from caller kwargs.
        * Later PP ranks: overrides ``x`` with ``intermediate_tensors["hidden_states"]``
          (which transparently waits for the async receive).
        * Non-last PP ranks return ``IntermediateTensors``; the last rank
          returns the plain noise-prediction tensor.
        """
        if intermediate_tensors is not None:
            x = intermediate_tensors["hidden_states"]

        for i in range(self.start_layer, self.end_layer):
            x = self.layers[i](x)

        pp_group = get_pp_group()
        if not pp_group.is_last_rank:
            return IntermediateTensors({"hidden_states": x})
        return x


# ---------------------------------------------------------------------------
# 1.  AsyncLatents – unit tests (no distributed env required)
# ---------------------------------------------------------------------------


class TestAsyncLatents:
    """Verifies the lazy-resolution behaviour of AsyncLatents without a real process group."""

    pytestmark = [pytest.mark.cpu]

    def _make(self, tensor: torch.Tensor, handles: list | None = None, postproc: list | None = None) -> AsyncLatents:
        return AsyncLatents({"latents": tensor}, handles or [], postproc or [])

    def test_resolve_returns_wrapped_tensor(self):
        t = torch.randn(2, 4)
        al = self._make(t)
        assert al._resolve() is t

    def test_attribute_access_resolves(self):
        t = torch.randn(2, 4)
        al = self._make(t)
        assert al.shape == t.shape
        assert al.dtype == t.dtype

    def test_torch_function_protocol(self):
        """torch ops that receive an AsyncLatents should see the underlying tensor."""
        t = torch.randn(2, 4)
        al = self._make(t)
        mask = torch.ones_like(t)
        result = mask * al  # triggers __torch_function__
        torch.testing.assert_close(result, mask * t)

    def test_torch_function_with_list_arg(self):
        """__torch_function__ must unwrap AsyncLatents inside list/tuple args."""
        t = torch.randn(2, 4)
        al = self._make(t)
        result = torch.cat([al, al], dim=0)
        torch.testing.assert_close(result, torch.cat([t, t], dim=0))

    def test_torch_tensor_conversion(self):
        """torch.as_tensor on an AsyncLatents must share storage with the underlying tensor (no copy)."""
        t = torch.randn(2, 4)
        al = self._make(t)
        result = torch.as_tensor(al)
        assert result.data_ptr() == t.data_ptr(), "torch.as_tensor copied the data instead of sharing storage"

    def test_handles_are_waited_before_resolve(self):
        t = torch.randn(2, 4)
        h1, h2 = FakeWork(), FakeWork()
        al = self._make(t, handles=[h1, h2])
        _ = al.shape  # trigger resolution
        assert h1.waited and h2.waited, "Not all handles were waited on"

    def test_postproc_callbacks_invoked_on_resolve(self):
        t = torch.randn(2, 4)
        log: list[int] = []
        al = self._make(t, postproc=[lambda: log.append(1), lambda: log.append(2)])
        _ = al.shape
        assert log == [1, 2], f"postproc not called in order: {log}"

    def test_idempotent_resolve(self):
        """handle.wait() must not be called twice if _resolve() is called twice."""
        t = torch.randn(2, 4)
        h = FakeWork()
        al = self._make(t, handles=[h])
        _ = al.shape  # first resolve
        h.waited = False  # reset sentinel
        _ = al.dtype  # second resolve
        assert not h.waited, "handle.wait() was called a second time"


# ---------------------------------------------------------------------------
# 2.  _sync_pp_send / diffuse wrapper – unit tests (no distributed env required)
# ---------------------------------------------------------------------------


class TestSyncPPSend:
    """Verifies PipelineParallelMixin's internal PP-send flush."""

    pytestmark = [pytest.mark.cpu]

    @staticmethod
    def _make_pipeline() -> PipelineParallelMixin:
        # Instantiate a bare mixin — no layers, no distributed env needed.
        # _sync_pp_send only touches _pp_send_work, so this is sufficient.
        class _BarePP(PipelineParallelMixin, CFGParallelMixin):
            pass

        return _BarePP()

    def test_noop_when_work_list_empty(self):
        pipeline = self._make_pipeline()
        pipeline._sync_pp_send()
        assert pipeline._pp_send_work == []

    def test_waits_all_pending_handles(self):
        pipeline = self._make_pipeline()
        works = [FakeWork(), FakeWork(), FakeWork()]
        pipeline._pp_send_work = works
        pipeline._sync_pp_send()
        assert all(w.waited for w in works), "Some handles were not waited on"

    def test_clears_work_list_after_sync(self):
        pipeline = self._make_pipeline()
        pipeline._pp_send_work = [FakeWork()]
        pipeline._sync_pp_send()
        assert pipeline._pp_send_work == []


class TestDiffuseWrapper:
    """Verifies that PipelineParallelMixin flushes pending sends when diffuse() exits."""

    pytestmark = [pytest.mark.cpu]

    def test_diffuse_flushes_pending_sends_on_success(self):
        work = FakeWork()

        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
            def diffuse(self):
                self._pp_send_work = [work]
                return "done"

        pipeline = _DiffusePP()

        assert pipeline.diffuse() == "done"
        assert work.waited
        assert pipeline._pp_send_work == []

    def test_diffuse_flushes_pending_sends_on_exception(self):
        work = FakeWork()

        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
            def diffuse(self):
                self._pp_send_work = [work]
                raise RuntimeError("boom")

        pipeline = _DiffusePP()

        with pytest.raises(RuntimeError, match="boom"):
            pipeline.diffuse()
        assert work.waited
        assert pipeline._pp_send_work == []

    def test_diffuse_wrapper_preserves_metadata(self):
        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
            def diffuse(self):
                """Original diffuse docstring."""
                return "done"

        assert _DiffusePP.diffuse.__name__ == "diffuse"
        assert _DiffusePP.diffuse.__doc__ == "Original diffuse docstring."


class TestVaeDecodeGuard:
    pytestmark = [pytest.mark.cpu]

    @staticmethod
    def _make_pipeline(distributed_enabled: bool = False) -> PipelineParallelMixin:
        class _DecodePP(PipelineParallelMixin, CFGParallelMixin):
            def __init__(self):
                self.vae = FakeVAE(distributed_enabled=distributed_enabled)

        return _DecodePP()

    @staticmethod
    def _set_rank(monkeypatch, world_size: int, first_stage: bool) -> None:
        monkeypatch.setattr(pp_module, "get_pipeline_parallel_world_size", lambda: world_size)
        monkeypatch.setattr(pp_module, "is_pipeline_first_stage", lambda: first_stage)

    def test_calls_original_decode_when_pp_disabled(self, monkeypatch):
        self._set_rank(monkeypatch, world_size=1, first_stage=True)
        pipeline = self._make_pipeline()
        z = torch.ones(2, 3)

        output = pipeline.vae.decode(z)[0]

        assert pipeline.vae.calls == 1
        torch.testing.assert_close(output, z + 1)

    def test_calls_original_decode_on_first_stage(self, monkeypatch):
        self._set_rank(monkeypatch, world_size=2, first_stage=True)
        pipeline = self._make_pipeline()
        z = torch.ones(2, 3)

        output = pipeline.vae.decode(z)[0]

        assert pipeline.vae.calls == 1
        torch.testing.assert_close(output, z + 1)

    def test_skips_decode_on_non_first_stage(self, monkeypatch):
        self._set_rank(monkeypatch, world_size=2, first_stage=False)
        pipeline = self._make_pipeline()
        z = torch.ones(2, 3)

        output = pipeline.vae.decode(z)

        assert pipeline.vae.calls == 0
        assert output == (None,)

    def test_calls_original_decode_when_distributed_vae_enabled(self, monkeypatch):
        self._set_rank(monkeypatch, world_size=2, first_stage=False)
        pipeline = self._make_pipeline(distributed_enabled=True)
        z = torch.ones(2, 3)

        output = pipeline.vae.decode(z)[0]

        assert pipeline.vae.calls == 1
        torch.testing.assert_close(output, z + 1)

    def test_decode_wrapper_preserves_metadata(self):
        pipeline = self._make_pipeline()

        assert pipeline.vae.decode.__name__ == "decode"
        assert pipeline.vae.decode.__doc__ == "Original decode docstring."


@pytest.mark.cpu
def test_pipeline_parallel_requires_cfg_mixin():
    with pytest.raises(TypeError, match="inherits PipelineParallelMixin but not CFGParallelMixin"):

        class _MissingCFG(PipelineParallelMixin):
            pass


@pytest.mark.cpu
def test_pipeline_parallel_requires_mro_before_cfg_mixin():
    with pytest.raises(TypeError, match="must inherit PipelineParallelMixin before CFGParallelMixin"):

        class _WrongOrder(CFGParallelMixin, PipelineParallelMixin):
            pass


# ---------------------------------------------------------------------------
# Distributed test helpers
# ---------------------------------------------------------------------------


def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.device:
    """Initialise the distributed environment for a spawned worker."""
    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
    current_omni_platform.set_device(device)
    update_environment_variables(
        {
            "RANK": str(local_rank),
            "LOCAL_RANK": str(local_rank),
            "WORLD_SIZE": str(world_size),
            "MASTER_ADDR": "localhost",
            "MASTER_PORT": master_port,
        }
    )
    init_distributed_environment()
    return device


def make_pipeline_and_inputs(
    test_config: dict, dtype: torch.dtype, device: torch.device, do_true_cfg: bool = False
) -> tuple["MockPipelineParallel", dict, dict | None]:
    """Create a MockPipelineParallel and seeded inputs from a test_config dict.

    Must be called after ``initialize_model_parallel`` so that ``make_layers``
    can read the PP group to determine this rank's layer slice.

    Returns ``(pipeline, positive_kwargs, negative_kwargs)``.
    ``negative_kwargs`` is ``None`` when ``do_true_cfg=False``.
    """
    pipeline = MockPipelineParallel(
        num_layers=test_config["num_layers"],
        dim=test_config["dim"],
        seed=test_config["model_seed"],
        device=device,
        dtype=dtype,
    )

    torch.manual_seed(test_config["input_seed"])
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(test_config["input_seed"])
    pos_x = {"x": torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)}

    negative_kwargs = None
    if do_true_cfg:
        torch.manual_seed(test_config["input_seed"] + 1)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
        neg_x = torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)
        negative_kwargs = {"x": neg_x}

    return pipeline, pos_x, negative_kwargs


# ---------------------------------------------------------------------------
# 3.  isend_tensor_dict / irecv_tensor_dict  (2 GPUs)
# ---------------------------------------------------------------------------


def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, result_queue):
    device = init_dist(local_rank, world_size, master_port)
    initialize_model_parallel(pipeline_parallel_size=world_size)
    pp_group = get_pp_group()

    if pp_group.is_first_rank:
        torch.manual_seed(77)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(77)
        tensor = torch.randn(3, 5, dtype=torch.float32, device=device)
        handles = pp_group.isend_tensor_dict({"t": tensor})
        for h in handles:
            h.wait()
        result_queue.put(("sent", tensor.cpu()))
    else:
        received = AsyncIntermediateTensors(*pp_group.irecv_tensor_dict())
        result_queue.put(("received", received["t"].cpu()))

    if torch.distributed.is_initialized():
        torch.distributed.barrier()
    destroy_distributed_env()


@pytest.mark.gpu
@pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs")
@pytest.mark.parametrize("pp_size", [2])
def test_isend_irecv_tensor_dict(pp_size: int):
    """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
    mp_context = torch.multiprocessing.get_context("spawn")
    manager = mp_context.Manager()
    q = manager.Queue()

    port = _find_free_port()
    torch.multiprocessing.spawn(isend_irecv_worker, args=(pp_size, port, q), nprocs=pp_size)

    results = {label: tensor for label, tensor in [q.get(), q.get()]}
    torch.testing.assert_close(
        results["received"], results["sent"], rtol=0, atol=0, msg="isend/irecv transferred tensor incorrectly"
    )


# ---------------------------------------------------------------------------
# 4.  predict_noise_maybe_with_cfg
# ---------------------------------------------------------------------------

_baseline_cache: dict[tuple, torch.Tensor] = {}


def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_cfg: bool) -> torch.Tensor:
    """Compute expected single-GPU output using the same MockPipelineParallel.

    Initializes a trivial distributed env (world_size=1) so that ``make_layers`` and the PP/CFG mixins work normally.
    Results are cached so identical configs are only computed once.
    """
    key = (
        test_config["num_layers"],
        test_config["dim"],
        test_config["batch_size"],
        test_config["model_seed"],
        test_config["input_seed"],
        test_config["cfg_scale"],
        dtype,
        do_true_cfg,
    )
    if key in _baseline_cache:
        return _baseline_cache[key]

    device = init_dist(0, 1, _find_free_port())
    initialize_model_parallel(pipeline_parallel_size=1)

    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
        test_config, dtype, device, do_true_cfg=do_true_cfg
    )

    with torch.inference_mode():
        noise_pred = pipeline.predict_noise_maybe_with_cfg(
            do_true_cfg=do_true_cfg,
            true_cfg_scale=test_config["cfg_scale"],
            positive_kwargs=positive_kwargs,
            negative_kwargs=negative_kwargs,
            cfg_normalize=False,
        )

    destroy_distributed_env()

    _baseline_cache[key] = noise_pred.cpu()
    return _baseline_cache[key]


def predict_noise_worker(
    local_rank: int,
    world_size: int,
    master_port: str,
    pp_size: int,
    cfg_size: int,
    do_true_cfg: bool,
    dtype: torch.dtype,
    test_config: dict,
    result_queue,
):
    """Generic predict-noise worker parameterized by PP and CFG topology."""
    device = init_dist(local_rank, world_size, master_port)
    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)

    pp_group = get_pp_group()
    cfg_rank = get_classifier_free_guidance_rank()

    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
        test_config, dtype, device, do_true_cfg=do_true_cfg
    )

    with torch.inference_mode():
        noise_pred = pipeline.predict_noise_maybe_with_cfg(
            do_true_cfg=do_true_cfg,
            true_cfg_scale=test_config["cfg_scale"],
            positive_kwargs=positive_kwargs,
            negative_kwargs=negative_kwargs,
            cfg_normalize=False,
        )
    # This worker exercises predict_noise_maybe_with_cfg directly, bypassing diffuse().
    # Flush the non-last PP rank's async send before barrier / process teardown.
    pipeline._sync_pp_send()

    if pp_group.is_last_rank:
        assert noise_pred is not None
        if cfg_rank == 0:
            result_queue.put(noise_pred.cpu())
    else:
        assert noise_pred is None

    if torch.distributed.is_initialized():
        torch.distributed.barrier()
    destroy_distributed_env()


@pytest.mark.gpu
@pytest.mark.parametrize(
    "pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol",
    [
        pytest.param(
            2,
            1,
            False,
            torch.float32,
            4,
            100,
            1e-5,
            1e-5,
            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
            id="pp2-no_cfg-float32",
        ),
        pytest.param(
            2,
            1,
            False,
            torch.bfloat16,
            4,
            100,
            1e-2,
            1e-2,
            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
            id="pp2-no_cfg-bfloat16",
        ),
        pytest.param(
            2,
            1,
            True,
            torch.bfloat16,
            4,
            100,
            1e-2,
            1e-2,
            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
            id="pp2-seq_cfg-bfloat16",
        ),
        pytest.param(
            2,
            2,
            True,
            torch.bfloat16,
            4,
            100,
            1e-2,
            1e-2,
            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
            id="pp2-cfg2-bfloat16",
        ),
        pytest.param(
            3,
            1,
            False,
            torch.bfloat16,
            6,
            100,
            1e-2,
            1e-2,
            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 3, reason="Need at least 3 GPUs"),
            id="pp3-no_cfg-bfloat16",
        ),
    ],
)
def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol):
    """predict_noise_maybe_with_cfg output matches the single-GPU baseline across PP / CFG topologies."""
    test_config = {
        "num_layers": num_layers,
        "dim": 64,
        "batch_size": 2,
        "cfg_scale": 7.5,
        "model_seed": 42,
        "input_seed": input_seed,
    }

    baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg)

    mp_context = torch.multiprocessing.get_context("spawn")
    manager = mp_context.Manager()
    pp_q = manager.Queue()

    world_size = pp_size * cfg_size
    port = _find_free_port()
    torch.multiprocessing.spawn(
        predict_noise_worker,
        args=(world_size, port, pp_size, cfg_size, do_true_cfg, dtype, test_config, pp_q),
        nprocs=world_size,
    )

    pp_out = pp_q.get()

    assert baseline_out.shape == pp_out.shape
    torch.testing.assert_close(
        pp_out,
        baseline_out,
        rtol=rtol,
        atol=atol,
        msg=f"PP={pp_size} cfg={cfg_size} {'with' if do_true_cfg else 'no'} CFG output differs from baseline ({dtype=})",
    )


# ---------------------------------------------------------------------------
# 5.  scheduler_step_maybe_with_cfg
# ---------------------------------------------------------------------------


def compute_scheduler_step_baseline(test_config: dict, do_true_cfg: bool) -> torch.Tensor:
    """Single-GPU reference: predict_noise + scheduler_step."""
    device = init_dist(0, 1, _find_free_port())
    initialize_model_parallel(pipeline_parallel_size=1)

    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
        test_config, torch.float32, device, do_true_cfg=do_true_cfg
    )
    latents = positive_kwargs["x"]
    t = torch.tensor(500, device=device)

    with torch.inference_mode():
        noise_pred = pipeline.predict_noise_maybe_with_cfg(
            do_true_cfg=do_true_cfg,
            true_cfg_scale=test_config["cfg_scale"],
            positive_kwargs=positive_kwargs,
            negative_kwargs=negative_kwargs,
            cfg_normalize=False,
        )
        result = pipeline.scheduler_step_maybe_with_cfg(
            noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
        )

    destroy_distributed_env()
    return result.cpu()


def scheduler_step_worker(
    local_rank: int,
    world_size: int,
    master_port: str,
    pp_size: int,
    cfg_size: int,
    do_true_cfg: bool,
    test_config: dict,
    result_queue,
):
    device = init_dist(local_rank, world_size, master_port)
    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)

    pp_group = get_pp_group()
    cfg_rank = get_classifier_free_guidance_rank()

    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
        test_config, torch.float32, device, do_true_cfg=do_true_cfg
    )
    latents = positive_kwargs["x"]
    t = torch.tensor(500, device=device)

    with torch.inference_mode():
        noise_pred = pipeline.predict_noise_maybe_with_cfg(
            do_true_cfg=do_true_cfg,
            true_cfg_scale=test_config["cfg_scale"],
            positive_kwargs=positive_kwargs,
            negative_kwargs=negative_kwargs,
            cfg_normalize=False,
        )
        latents = pipeline.scheduler_step_maybe_with_cfg(
            noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
        )
    # This worker exercises scheduler_step_maybe_with_cfg directly, bypassing diffuse().
    # Flush the last PP rank's async latent send before barrier / process teardown.
    pipeline._sync_pp_send()

    if pp_group.is_first_rank and cfg_rank == 0:
        resolved = latents.contiguous()
        result_queue.put(resolved.cpu())

    if torch.distributed.is_initialized():
        torch.distributed.barrier()
    destroy_distributed_env()


@pytest.mark.gpu
@pytest.mark.parametrize(
    "pp_size, cfg_size, do_true_cfg, input_seed",
    [
        pytest.param(
            2,
            1,
            False,
            300,
            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
            id="pp2-no_cfg",
        ),
        pytest.param(
            2,
            2,
            True,
            600,
            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
            id="pp2-cfg2-true_cfg",
        ),
    ],
)
def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
    """Rank 0 latents after scheduler_step match the single-GPU baseline across PP / CFG topologies."""
    test_config = {
        "num_layers": 4,
        "dim": 64,
        "batch_size": 2,
        "cfg_scale": 7.5,
        "model_seed": 42,
        "input_seed": input_seed,
    }

    baseline = compute_scheduler_step_baseline(test_config, do_true_cfg)

    mp_context = torch.multiprocessing.get_context("spawn")
    manager = mp_context.Manager()
    q = manager.Queue()

    port = _find_free_port()
    world_size = pp_size * cfg_size
    torch.multiprocessing.spawn(
        scheduler_step_worker,
        args=(world_size, port, pp_size, cfg_size, do_true_cfg, test_config, q),
        nprocs=world_size,
    )

    result = q.get()
    torch.testing.assert_close(
        result,
        baseline,
        rtol=0,
        atol=0,
        msg=f"PP={pp_size} CFG={cfg_size} scheduler step latents on rank 0 do not match single-GPU baseline",
    )