[tangbinh]

Purpose

Add a Ray-based distributed executor for diffusion models, enabling multi-node GPU execution. Resolves #796.

The existing MultiprocDiffusionExecutor (#865) uses local multiprocessing, which limits execution to a single node. This PR adds RayDiffusionExecutor as an alternative backend that uses Ray actors for worker management, supporting both single-node multi-GPU and multi-node execution across a Ray cluster.

Design follows the same patterns as vLLM's RayDistributedExecutor (vllm/v1/executor/ray_executor.py), adapted for the diffusion executor interface.

Usage:

omni = Omni(
    model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
    distributed_executor_backend="ray",
    num_gpus=2,
    parallel_config=DiffusionParallelConfig(
        sequence_parallel_size=2,
        ulysses_degree=2,
    ),
)

Test Plan

• Unit tests (no GPU needed)

pytest tests/diffusion/executor/test_ray_executor.py -v

• E2E tests (requires GPUs)

# Single GPU
CUDA_VISIBLE_DEVICES=0 pytest tests/e2e/offline_inference/test_ray_executor_t2v.py -v -k single

# Multi-GPU with Ulysses sequence parallelism
CUDA_VISIBLE_DEVICES=0,1 pytest tests/e2e/offline_inference/test_ray_executor_t2v.py -v -k parallel

Test Result

• Unit tests

tests/diffusion/executor/test_ray_executor.py::TestRayDiffusionWorkerWrapper::test_init_worker PASSED
tests/diffusion/executor/test_ray_executor.py::TestRayDiffusionWorkerWrapper::test_execute_raises_when_uninitialized PASSED
tests/diffusion/executor/test_ray_executor.py::TestExecutorFactory::test_get_class_returns_ray_executor PASSED

======================== 3 passed in 1.07s ========================

• E2E: Single GPU (Wan2.1-T2V-1.3B)

tests/e2e/offline_inference/test_ray_executor_t2v.py::test_ray_executor_single_gpu PASSED [35.22s]

• E2E: 2-GPU Ulysses SP (Wan2.1-T2V-1.3B)

tests/e2e/offline_inference/test_ray_executor_t2v.py::test_ray_executor_parallel_ulysses PASSED [60.38s]

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please provide the test scripts & test commands.
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update.
• (Optional) Release notes update.

---

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 760c3512cc

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[yenuo26]

please add @pytest.mark.advanced_model to differentiate test case levels

[tangbinh]

I've added the marker in the latest commit.

[yenuo26]

maybe you can use omni_runner and omni_runner_handler in conftest.py, like e2e/offline_inference/test_qwen3_omni.py

[tangbinh]

I tried to follow the setup but it looks like adding test fixtures for omni_runner and omni_runner_handler is not very necessary for these simple tests. I did extract the duplicated assertion logic into a shared helper to reduce code duplication though.

[hsliuustc0106]

Review Summary

This PR adds a well-structured Ray executor backend for distributed diffusion model execution. The implementation follows vLLM's RayDistributedExecutor patterns and includes good test coverage for the main inference path.

✅ Validated

• Design follows vLLM patterns appropriately
• Unit tests cover worker initialization and error handling
• E2E tests demonstrate single-GPU and multi-GPU Ulysses SP working
• Proper placement group management and worker lifecycle

🔧 Changes Needed

1. Merge conflicts - PR needs rebase onto main

2. Missing test for collective_rpc() - This is public API that downstream code may use for distributed coordination. A unit test would ensure it behaves correctly.

📝 Suggestions (non-blocking)

• The to_cpu() method in DiffusionOutput mutates in-place. Consider adding a docstring noting this behavior for clarity.

---

Thanks for the detailed test evidence in the PR description!

[hsliuustc0106]

The collective_rpc() method is public API but has no test coverage. Consider adding a unit test to verify:

1. It correctly calls execute_method on all workers
2. The unique_reply_rank parameter returns only the specified rank's response
3. Timeout handling raises TimeoutError

Example:

def test_collective_rpc_returns_all_responses(self, mocker, mock_od_config):
    """collective_rpc should aggregate responses from all workers."""
    # Setup mock executor with workers
    # Call collective_rpc with a method name
    # Verify it returns responses from all workers

def test_collective_rpc_unique_reply_rank(self, mocker, mock_od_config):
    """collective_rpc with unique_reply_rank should return single response."""
    # Verify only the specified rank's response is returned

[tangbinh]

Thank you for the suggestions! I've included these tests in the latest commit.

[lishunyang12]

Two things on shutdown/cleanup:

1. shutdown() uses ray.kill() without graceful shutdown (ray_executor.py) — the mp executor sends SHUTDOWN_MESSAGE and does proc.join(30) before proc.terminate(). Here it's straight to SIGKILL. RayDiffusionWorkerWrapper.shutdown() exists but is never called. Consider calling meta.worker.shutdown.remote() with a timeout first.

2. __del__ as cleanup fallback is unreliable — GC may never call it, or call it after Ray has shut down. The mp executor uses weakref.finalize() for this.

[tangbinh]

This PR was accidentally closed during rebasing and cannot be reopened (my apologies). Let's continue the review in PR #2106.