[RFC/Bug] Improve diffusion worker/engine control-plane reliability and output contracts

[hsliuustc0106]

Motivation

The current diffusion worker/engine implementation has made good progress: the engine/worker split is clear, diffusion lifecycle cleanup has improved, and several engine cleanup PRs are already tracking readability and telemetry work. However, a review of the current main branch and PR #4282 surfaced a few remaining reliability and scalability gaps that are not fully covered by existing issues.

The most important gap is distributed control-plane correctness: some all-rank diffusion RPCs can report success based only on rank 0 even when non-rank-0 workers fail. This can silently leave workers in inconsistent state after control operations such as LoRA add/remove, sleep, wake, or future extension RPCs.

Current status

What looks healthy today:

• DiffusionEngine owns request lifecycle, scheduling, postprocessing, and response shaping.
• DiffusionWorker owns device/model execution, with DiffusionModelRunner holding most model-specific execution logic.
• Multiprocess diffusion execution uses a broadcast/control channel and a rank-0 result path.
• Sleep/wake support and cleanup lifecycle have improved through prior work.

Remaining problems:

1. All-rank RPC failure visibility

   • MultiprocDiffusionExecutor.collective_rpc() executes some RPCs on all workers but expects only one response, normally from rank 0.
   • Non-rank workers catch/log exceptions locally, but failures are not propagated to the caller.
   • The caller can receive success even if rank 1+ failed, leaving distributed worker state divergent.
   • This is especially risky for LoRA mutation, sleep/wake, and any future all-rank stateful control RPC.

2. Implicit diffusion output contract

   • DiffusionEngine.step() currently handles loose dict outputs with keys such as video, audio, actions, custom_output, fps, and audio_sample_rate.
   • Multi-prompt output splitting and batch-scoped metadata handling live in the same broad method.
   • This works for current models but becomes brittle as image/video/audio/action outputs and model-specific postprocess paths grow.

3. Residual shutdown edge

   • PR Fix diffusion engine cleanup lifecycle #3494 fixed several important cleanup lifecycle bugs and is already merged.
   • A remaining edge is that close() can return after the background thread join timeout while scheduler/executor shutdown is deferred. That may be intentional, but the engine state and resource ownership should be explicit so callers know whether the engine is fully closed, failed, or pending retry cleanup.

4. Worker API contract cleanup

   • Some worker APIs have unclear contracts, for example sleep() is annotated like a boolean operation while returning byte counts.
   • LoRA public methods assume a LoRA manager exists; models without one can fail via AttributeError rather than returning a clear unsupported-operation result.

Related issues and PRs

• [RFC]: diffusion engine clean up #2335 tracks broader DiffusionEngine cleanup: concurrency, telemetry, CPU utilization, and step readability.
• [BugFix][Core] Address diffusion engine issues in #2335 #2409 is an open broad refactor for queue-backed engine loop, request/RPC ownership, metrics, and step materialization.
• fix(diffusion): correct metrics, recursive CPU offload, batched audio slicing, metadata scope #2703 is an open focused cleanup for metrics, recursive CPU offload, batched audio slicing, and metadata scope.
• fix(diffusion): recursive CPU offload, batched audio slicing, metadata scope #2694 is a similar/open phase-2 cleanup around recursive CPU offload and batched output splitting.
• Fix diffusion engine cleanup lifecycle #3494 is merged and fixed several cleanup lifecycle defects in DiffusionEngine.
• [Feat][sleepmode] add omni sleepmode and ack protocol #2022 is merged and introduced Omni sleep mode and ACK protocol.
• [Frontend] Enable sleep/wake HTTP API without VLLM_SERVER_DEV_MODE #2224 is open for sleep/wake HTTP API exposure.
• [Refactor] Omni Stage Runtime and Distributed Replica Control Plane #3855 is open and refactors stage runtime and distributed replica control-plane ownership.
• Support DROID policy server for Cosmos3 OpenPI #4282 adds/adjusts diffusion engine/worker paths for Cosmos3/OpenPI and highlighted the need for clearer action/custom output ownership.

Proposed fix options

Option A: Minimal correctness fix for all-rank RPCs

Add per-rank success/error aggregation for all-rank collective_rpc() calls while preserving the current request protocol.

Possible shape:

• Each worker sends a small RPC status envelope for all-rank control calls.
• Rank 0 or the executor aggregates all statuses.
• The caller receives success only if every rank succeeded.
• Any rank failure includes rank id, method name, and traceback/error string.

Pros:

• Lowest-risk path.
• Directly fixes the highest-priority correctness issue.
• Can be covered with unit tests by forcing a non-rank-0 RPC failure.

Cons:

• Keeps the existing rank-0-centric result path and may add special cases.

Option B: Typed control-plane result envelope

Introduce a structured result envelope for diffusion RPCs, for example:

@dataclass
class DiffusionRPCResult:
    ok: bool
    method: str
    rank_results: dict[int, RankRPCResult]
    value: Any | None = None

Pros:

• Makes all-rank vs single-rank semantics explicit.
• Gives future control APIs a safer contract.
• Easier to extend for sleep/wake/LoRA/health-check status.

Cons:

• Moderate migration cost across executor, worker, and stage proc layers.

Option C: Align with broader stage-runtime control-plane refactor

Fold the fix into the direction of #3855, where stage runtime and distributed replica ownership are being cleaned up.

Pros:

• Better long-term architecture.
• Can unify diffusion and LLM control-plane behavior.
• Avoids adding another short-lived protocol layer.

Cons:

• Larger scope.
• Does not immediately protect current all-rank diffusion RPCs unless prioritized inside that PR.

Option D: Typed diffusion output/postprocess envelope

Introduce a typed output object for diffusion postprocess/model outputs, separate from the RPC fix. For example:

@dataclass
class DiffusionPostprocessResult:
    video: Any | None = None
    audio: Any | None = None
    actions: Any | None = None
    custom_output: dict[str, Any] | None = None
    fps: float | None = None
    audio_sample_rate: int | None = None

Pros:

• Reduces DiffusionEngine.step() branching.
• Makes image/video/audio/action ownership explicit.
• Complements fix(diffusion): correct metrics, recursive CPU offload, batched audio slicing, metadata scope #2703/fix(diffusion): recursive CPU offload, batched audio slicing, metadata scope #2694 output-splitting cleanup.

Cons:

• Needs compatibility shims for existing model postprocess functions.

Recommendation

Prioritize Option A first to close the distributed correctness hole. Then decide whether Option B should be the stable control-plane contract or whether it should be folded into #3855. In parallel, #2703/#2694 can continue improving output splitting, but a typed output envelope should be considered before more model-specific dict keys are added to DiffusionEngine.step().

Discussion period

Please use this issue for design discussion from June 13, 2026 through June 27, 2026. After that, we should decide whether to:

• land a minimal all-rank RPC aggregation fix,
• fold the work into [Refactor] Omni Stage Runtime and Distributed Replica Control Plane #3855,
• or split the effort into separate RPC-contract and output-contract PRs.

[hsliuustc0106]

diffusion_engine.py audit — cleanup opportunities

I did a full audit of diffusion_engine.py (973 lines) to identify code that can be moved out of the core engine. This is directly relevant to point 2 (implicit diffusion output contract) and also touches the broader cleanup goals in #2335 / #2409 / #2703 / #2694.

What can move out

1. Model capability query functions → new diffusion/capabilities.py (~80 lines)

These are pure functions with no dependency on DiffusionEngine state. They query DiffusionModelRegistry for model capabilities and are used both externally and internally:

Function	Lines	External callers
supports_multimodal_input()	56–72	tests, serving_chat.py
image_color_format()	75–77	internal only
supports_audio_output()	80–84	async_omni_engine.py, tests
_move_tensor_tree_to_cpu()	87–96	tests
get_dummy_run_num_frames()	99–105	tests
get_extra_body_params()	108–119	serving_chat.py
get_extra_output_params()	122–133	serving_chat.py

These don't belong in the engine file. They could go into a new diffusion/capabilities.py or be merged into the existing model_metadata.py.

2. Output formatting from step() → standalone function (~175 lines)

step() is ~260 lines, but ~175 lines (lines 299–476) are purely output formatting — converting DiffusionOutput → list[OmniRequestOutput]. This code:

• Determines output type (text / audio / image)
• Builds metrics dicts
• Slices audio/action payloads for multi-prompt requests
• Constructs OmniRequestOutput.from_diffusion(...) calls with heavy branching

This is exactly the "brittle dict-output branching" called out in point 2 of the RFC. After extraction, step() would shrink to ~85 lines (just pre-process → execute → post-process → delegate to formatter).

This extraction also makes Option D from the RFC more tractable — once the formatting is a standalone pure function, introducing a DiffusionPostprocessResult typed envelope is a much cleaner change because the conversion boundary is explicit.

3. _dummy_run() request construction (~35 lines)

The request-building part of _dummy_run() (constructing the warmup request with dummy image/audio) is a pure function that depends only on capability queries. Only the final add_req_and_wait_for_response() call needs to stay in the engine.

What should stay

The core engine concerns are clean and should remain:

• Lifecycle (__init__, close, make_engine)
• Async loop management (_check_and_start_background_loop, _busy_loop)
• Request submission & result waiting (add_request, get_result, async_add_req_and_wait_for_response)
• RPC infrastructure (collective_rpc, async_collective_rpc, _submit_rpc, _process_rpc_queue)
• Abort handling
• Scheduler/executor output handling (_handle_finished_requests, _finalize_finished_request)

Estimated impact

What	Lines moved	Engine size after
Capability functions → capabilities.py	~80	—
Output formatting → standalone function	~175	step(): ~260 → ~85 lines
Dummy run construction → helper	~35	—
Total	~290	~680 → ~390 lines of class code

Proposal

I'd like to send a PR for items 1 and 2 first (capabilities extraction + output formatting extraction), since they're purely structural and don't change behavior. Item 2 in particular sets up a clean boundary for Option D's typed output envelope in a follow-up. Item 3 (dummy run) is a smaller change that can go in the same PR or separately.

Thoughts? Does this align with the cleanup direction envisioned in this RFC?

[hsliuustc0106]

Thanks, this aligns well with the cleanup direction for point 2.

I agree items 1 and 2 are good first steps as a behavior-preserving PR:

• move capability/model metadata helpers out of diffusion_engine.py
• extract output formatting from step() into a standalone conversion boundary

A couple of requests:

1. Please keep compatibility imports/re-exports if external callers currently import these helpers from diffusion_engine.py.
2. Add focused tests proving the extracted formatter preserves current image/audio/video/action/custom_output behavior.
3. Keep this separate from the all-rank RPC reliability fix. The formatter extraction prepares Option D, but Option A is still the highest-priority correctness issue because non-rank worker failures can remain invisible.

After that lands, a follow-up typed DiffusionPostprocessResult would be much easier to review. I would not consider this cleanup PR sufficient to close this RFC by itself; the RPC aggregation/control-plane issue still needs a concrete fix path.

[Gaohan123]

@fhfuih PTAL