[Bugfix] Fix Parallel State Initialization For MoE + Diffusion (#7676)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/diffusion/distributed/test_parallel_state_sp_groups.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for SP subgroup construction in parallel_state.py."""
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit tests for SP subgroup construction for diffusion parallel state."""
 
 from __future__ import annotations
 
@@ -8,12 +8,55 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+import vllm.distributed.parallel_state as vllm_parallel_state
 
+from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.diffusion.distributed import parallel_state as omni_parallel_state
 from vllm_omni.diffusion.distributed.parallel_state import RankGenerator, set_seq_parallel_pg
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.core_model, pytest.mark.cpu]
 
 
+def ensure_parallel_state_initialized():
+    """Ensure torch distributed is initialized and that all world / count vars are set."""
+    assert torch.distributed.is_initialized()
+    assert omni_parallel_state._WORLD is not None
+    assert vllm_parallel_state._WORLD is not None
+    assert vllm_parallel_state._NODE_COUNT is not None
+
+
+def ensure_parallel_state_not_initialized():
+    """Ensure torch distributed is not initialized and that all world / count vars are unset."""
+    assert not torch.distributed.is_initialized()
+    assert omni_parallel_state._WORLD is None
+    assert vllm_parallel_state._WORLD is None
+    assert vllm_parallel_state._NODE_COUNT is None
+
+
+def test_omni_manages_vllm_distributed_state(monkeypatch):
+    """Verify Omni initializes and tears down vLLM's distributed state properly.
+
+    This is a regression test for ensuring vLLM Omni also sets up vLLM's vars correctly
+    to avoid potential misalignment in cases where Omni uses vLLM's native coordinator,
+    e.g., MoE + diffusion.
+    """
+    ensure_parallel_state_not_initialized()
+    monkeypatch.setattr(omni_parallel_state.current_omni_platform, "get_device_count", lambda: 1)
+    monkeypatch.setattr(omni_parallel_state.current_omni_platform, "set_device", lambda _device: None)
+    omni_parallel_state.init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method=get_distributed_init_method(),
+        backend="gloo",
+    )
+    try:
+        ensure_parallel_state_initialized()
+    finally:
+        omni_parallel_state.destroy_distributed_env()
+    ensure_parallel_state_not_initialized()
+
+
 def _fake_new_group_factory(created_groups: list[SimpleNamespace]):
     def _fake_new_group(ranks, *args, **kwargs):
         group = SimpleNamespace(ranks=list(ranks))