diff --git a/docs/design/disaggregated_diffusion.md b/docs/design/disaggregated_diffusion.md deleted file mode 100644 index 8703a9e0d9d0..000000000000 --- a/docs/design/disaggregated_diffusion.md +++ /dev/null @@ -1,203 +0,0 @@ -# Design Doc: Disaggregated Diffusion Inference (Diff-Disagg) in Dynamo - -## 1. Motivation - -Current diffusion inference in Dynamo (e.g., SGLang backend) is monolithic: a single worker loads all model components (Text Encoder, Transformer/UNet, VAE). This design faces several challenges with modern large-scale diffusion models (Flux, SD3, Video Models): - -1. **Huge Text Encoders**: Models like Flux use T5-XXL, consuming 10-20GB VRAM just for the encoder, which is only used once at the beginning. -2. **Resource Inefficiency**: During the long denoising loop (20-50 steps), the text encoder weights occupy precious H100/H200 memory without being used. -3. **Heterogeneous Compute**: - * **Text Encoding**: Compute-intensive, single pass. Can run on lower-end GPUs or even CPU. - * **Denoising**: Memory-bandwidth and compute-intensive, iterative. Requires high-end GPUs (H100/H200). - * **VAE Decoding**: Memory-intensive, single pass. Can be offloaded. - -**Goal**: Decompose the diffusion pipeline into flexible, independent stages (Encoder, Denoiser, VAE) that can be deployed on different hardware and scaled independently. - -## 2. Architecture: Router-Orchestrated Multi-Stage Pipeline - -We adopt a **Router-Centric** architecture similar to Dynamo's existing EPD (Encoder-Prefill-Decode) flow for LLMs. The Frontend/Router acts as the orchestrator, chaining calls between specialized workers. - -### 2.1 Component Roles - -1. **Frontend / Global Router**: - * Acts as the conductor. - * Receives the initial user request. - * Orchestrates the sequence: `Encoder -> Denoiser -> VAE`. - * Maintains the request state and context. - -2. **Stage Workers**: - * **Encoder Worker**: Loads CLIP/T5. Input: Text. Output: Embeddings (Tensor Bytes). - * **Denoiser Worker**: Loads Transformer/UNet. Input: Embeddings + Latents (optional). Output: Denoised Latents. - * **VAE Worker**: Loads VAE. Input: Denoised Latents. Output: Image/Video (Pixel Data). - -### 2.2 Data Flow (Request/Response) - -```mermaid -sequenceDiagram - participant Client - participant Frontend - participant Router - participant EncoderWorker - participant DenoiserWorker - participant VAEWorker - - Client->>Frontend: POST /v1/images/generations (Prompt) - - Note over Frontend: Step 1: Text Encoding - Frontend->>Router: Route(ModelType.Encoder, Prompt) - Router->>EncoderWorker: GenerateEmbeddings(Prompt) - EncoderWorker-->>Frontend: Return {prompt_embeds, pooled_embeds} - - Note over Frontend: Step 2: Denoising - Frontend->>Router: Route(ModelType.Denoiser, Embeddings) - Router->>DenoiserWorker: Denoise(Embeddings, Latents=None) - DenoiserWorker-->>Frontend: Return {denoised_latents} - - Note over Frontend: Step 3: VAE Decoding - Frontend->>Router: Route(ModelType.VAE, Latents) - Router->>VAEWorker: Decode(Latents) - VAEWorker-->>Frontend: Return {image_data / url} - - Frontend-->>Client: Final Response -``` - -*Optimization Note*: For large data (like Video Latents), we can use **Reference Passing** via a shared storage (Object Store / Shared Memory) instead of passing raw bytes through the Router/Frontend. - -## 3. Detailed Design - -### 3.1 Protocol Extensions (`dynamo/common/protocols`) - -We need to define data structures for intermediate results. - -**New Protocol Types:** - -```python -class DiffusionEmbeddingData(BaseModel): - """Output from Encoder Stage""" - prompt_embeds: bytes # Serialized Tensor (e.g., safetensors/numpy) - pooled_prompt_embeds: Optional[bytes] = None - negative_prompt_embeds: Optional[bytes] = None - negative_pooled_prompt_embeds: Optional[bytes] = None - -class DiffusionLatentData(BaseModel): - """Output from Denoiser Stage""" - latents: bytes # Serialized Tensor - shape: List[int] - dtype: str - -class StageRequest(BaseModel): - """Generic Request for a specific stage""" - stage: str # "encoder", "denoiser", "vae" - input_data: Union[str, DiffusionEmbeddingData, DiffusionLatentData] - params: Dict[str, Any] # Generation params (steps, cfg, etc.) -``` - -### 3.2 ModelType Expansion (`dynamo/llm/src/model.rs` & Python Enums) - -Extend `ModelType` to support fine-grained stages. - -```python -class ModelType(IntFlag): - # Existing - Tokens = auto() - # ... - - # New Diffusion Stages - DiffusionEncoder = auto() # Text Encoder only - DiffusionDenoiser = auto() # Transformer/UNet only - DiffusionVAE = auto() # VAE only -``` - -### 3.3 SGLang Backend Adaptation (`components/src/dynamo/sglang`) - -We need to modify `init_diffusion.py` and handlers to support partial loading. - -**Configuration:** -Add `--diffusion-stage` argument to `sglang` worker. - -* `--diffusion-stage full` (Default): Loads everything (current behavior). -* `--diffusion-stage encoder`: Loads only Text Encoders. -* `--diffusion-stage denoiser`: Loads only Transformer, accepts Embeddings. -* `--diffusion-stage vae`: Loads only VAE, accepts Latents. - -**Handler Implementation:** - -1. **`EncoderHandler`**: - * Uses `pipe.encode_prompt()`. - * Returns serialized embeddings. - -2. **`DenoiserHandler`**: - * Initializes pipeline with `text_encoder=None`, `vae=None`. - * Implements `generate(prompt_embeds=...)`. - * Returns latents (skips VAE decode). - -3. **`VAEHandler`**: - * Initializes pipeline with `text_encoder=None`, `transformer=None`. - * Implements `decode(latents=...)`. - -### 3.4 Router Logic (`components/src/dynamo/global_router`) - -The Global Router needs to be aware of these new `ModelType`s. - -* **Registration**: Workers register with specific types (e.g., `ModelType.DiffusionEncoder`). -* **Routing**: - * `handle_encoder_request`: Routes to Encoder Pool. - * `handle_denoiser_request`: Routes to Denoiser Pool. - * `handle_vae_request`: Routes to VAE Pool. - -## 4. Implementation Plan - -### POC (Proof of Concept) - -Full POC code lives in [`examples/disagg_diffusion/`](../../examples/disagg_diffusion/). - -#### Phase 0: Offline Validation (`phase0_validate/`) - -Single-GPU script that proves diffusers supports split execution. -Runs Encoder → Denoiser → VAE as three separate stages with serialized -intermediate tensors, then compares the result against a monolithic run. - -**Key risk validated**: Can diffusers pipelines accept `prompt_embeds` and -return `output_type="latent"` to bypass text encoder / VAE respectively? - -#### Phase 1: Dynamo Stage Workers (`phase1_workers/`) - -Three independent Dynamo workers, each loading only its model component: - -| Worker | Loads | Endpoint | Input | Output | -|--------|-------|----------|-------|--------| -| `encoder_worker.py` | CLIP + T5 | `disagg_diffusion.encoder.encode` | text prompt | embeddings (b64) | -| `denoiser_worker.py` | Transformer | `disagg_diffusion.denoiser.denoise` | embeddings + params | latents (b64) | -| `vae_worker.py` | VAE | `disagg_diffusion.vae.decode` | latents | image (b64 PNG) | - -Intermediate data is serialized as base64-encoded `torch.save` bytes. -Protocol types are defined in `protocol.py`. - -#### Phase 2: Orchestrator Client (`phase2_orchestrator/`) - -A lightweight Python client that connects to the Dynamo runtime, calls the -three stage endpoints in sequence, and produces the final image. - -### Production Roadmap - -#### Phase 3: Protocol & Types -1. Define `DiffusionEmbeddingData` and `DiffusionLatentData` in `dynamo/common/protocols`. -2. Add `DiffusionEncoder`, `DiffusionDenoiser`, `DiffusionVAE` to `ModelType` enum (Rust + Python). - -#### Phase 4: SGLang Modularization -1. Modify `init_diffusion.py` to accept `--diffusion-stage {full,encoder,denoiser,vae}`. -2. Implement `EncoderHandler`, `DenoiserHandler`, `VAEHandler` in `sglang/request_handlers/disagg_diffusion/`. -3. Add logic to conditionally load model components based on stage. - -#### Phase 5: Frontend / Router Orchestration -1. Update `Frontend` to detect if the backend is disaggregated diffusion. -2. Implement the multi-step orchestration logic (`Encoder → Denoiser → VAE`) - in `Frontend` or a specialized `DiffusionOrchestrator` component. -3. Support configurable stage DAGs (e.g., skip VAE for latent-only output). - -#### Phase 6: Optimization -1. **Reference Passing**: Replace base64 tensor payloads with shared-memory - handles or object-store URLs (e.g., `media_output_fs_url`). -2. **Pipelining**: Overlap encoding of request N+1 with denoising of request N. -3. **Encoder Caching**: LRU cache for repeated prompts at the Encoder stage. -4. **Video Extension**: Extend to video models (larger latents, more frames). diff --git a/examples/disagg_diffusion/DESIGN.md b/examples/disagg_diffusion/DESIGN.md new file mode 100644 index 000000000000..5bdefa6a0c9e --- /dev/null +++ b/examples/disagg_diffusion/DESIGN.md @@ -0,0 +1,475 @@ +# Disaggregated Diffusion Pipeline — Design Document + +## 1. Overview + +### Motivation + +Modern diffusion pipelines (text-to-video, text-to-image, omni-modal) are +composed of heterogeneous stages — text encoding, iterative denoising, VAE +decoding — each with fundamentally different compute profiles: + +- **Different optimization strategies per stage.** Encoders are + memory-bound single-pass transforms; denoisers are compute-bound + multi-step loops that benefit from tensor parallelism; VAE decoders + are memory-intensive but run only once. Forcing all three into a + single process prevents stage-specific tuning (parallelism, batching, + memory management, quantization). + +- **Shifting compute balance.** As diffusion models mature, the DiT no + longer dominates the entire pipeline — faster denoisers (fewer steps, + distilled models) shift the bottleneck to encoding and decoding. + Multi-task pipelines are emerging (e.g. OneVideo: encode + denoise + + decode + audio in one model) where each task demands independent + scaling. + +- **Omni-modal future.** Models that jointly produce video, audio, + image, and text require stage-level separation so each modality's + compute can scale independently without wasting GPU resources. + +### Solution + +Decompose the pipeline into **N independent stages**, each running as a +Dynamo RPC worker on dedicated GPU(s). An orchestrator chains stages +together, using NIXL RDMA for GPU-direct tensor transfer between them. + +| Goal | Mechanism | +|---|---| +| Stage-level scaling | N workers per stage, auto-discovered via etcd | +| Pipeline parallelism | Semaphore admission, independent worker pools | +| GPU-direct transfer | NIXL RDMA — only ~1.5 KB metadata over RPC | +| Loose coupling | Workers are independent processes, no shared state | +| Dynamic scaling | Add/remove workers at runtime, no restart | +| Auto routing | Idle-queue dispatch, backpressure, retry on failure | + + +## 2. Architecture + +### 2.1 Architecture Diagram + +``` + ┌───────────┐ + │ etcd │ + │ registry │ + └─────┬─────┘ + register │ discover + ┌───────────────────────────────┼───────────────────────────────┐ + │ Orchestrator │ + │ │ + │ HTTP ──► handle_generate() ──► dispatch_with_retry() │ + │ PipelineTracker per-stage WorkerManager │ + │ Semaphore(depth) acquire → direct() → release │ + │ │ + └──────┬──────────────────────┬──────────────────────┬──────────┘ + │ Dynamo RPC │ Dynamo RPC │ Dynamo RPC + │ (JSON ~1 KB) │ (NIXL meta ~1.5 KB) │ (NIXL meta ~1.5 KB) + ▼ ▼ ▼ + ┌──────────┐ ┌──────────┐ ┌──────────┐ + │Encoder-0 │ │Denoiser-0│ │ VAE-0 │ + │ GPU 0 │ │ GPU 1,2 │ │ GPU 7 │ + └────┬─────┘ └──┬───┬───┘ └────┬─────┘ + │ NIXL RDMA pull │ │ NIXL RDMA pull │ + │◄──────────────────┘ └────────────────────►│ + │ (embeddings) (latents) │ + ┌──────────┐ ┌──────────┐ ┌──────────┐ + │Encoder-1 │ │Denoiser-1│ │ VAE-1 │ + │ GPU 3 │ │ GPU 3,4 │ │ ... │ + └──────────┘ └──────────┘ └──────────┘ + ...↕ ┌──────────┐ ...↕ + dynamic │Denoiser-2│ dynamic + add/remove │ GPU 5,6 │ add/remove + └──────────┘ + ...↕ + dynamic + add/remove +``` + +**Key points:** +- Workers self-register with etcd; the orchestrator discovers them + automatically. A settle period (`WORKER_SETTLE_S`, default 30s) + waits for slow-starting workers after the first instance appears. +- Dynamo RPC carries only small JSON payloads and NIXL metadata + (~1.5 KB). Actual tensors (embeddings, latents) transfer GPU-to-GPU + via NIXL RDMA, never touching the CPU or RPC channel. +- **RDMA is receiver-initiated (one-sided pull):** the denoiser pulls + embeddings from the encoder, the VAE pulls latents from the denoiser. + After the pull completes, the sender receives a completion + notification and releases the GPU buffer. +- Each stage can have a different number of workers and GPU count. + The denoiser typically uses TP > 1 (multi-GPU), while encoder and VAE + each use a single GPU. + +### 2.2 Process Model + +Each worker is an independent OS process with no shared state: + +``` +Dynamo Worker Process (e.g. denoiser_worker.py) +├── @dynamo_worker ← Dynamo runtime bootstrap +│ ├── serve_endpoint("generate") ← Dynamo RPC from orchestrator +│ │ └── handle_generate() +│ │ └── StageClient.forward() ← ZMQ to local backend subprocess +│ └── serve_endpoint("health") +│ +└── Backend subprocess(es) (spawned at startup, TP=2 → 2 processes) + ├── Rank 0 (master, handles ZMQ + NIXL) + │ ├── NixlReceiveStage ← RDMA-pull from previous stage + │ ├── ComputeStage(s) ← model inference (TP all-reduce) + │ └── NixlSendStage ← register output as RDMA-readable + └── Rank 1 (slave, TP compute only) + ├── NixlReceiveStage ← receives tensors via TP broadcast + ├── ComputeStage(s) ← model inference (TP all-reduce) + └── NixlSendStage ← no-op (only rank 0 sends) +``` + +- The Dynamo worker process handles RPC and control-plane logic. +- The backend subprocess runs the actual model inference. +- Communication between the two is via ZMQ REQ/REP (same-host IPC). +- With TP > 1: rank 0 does the NIXL pull then broadcasts tensors to + other ranks via `torch.distributed.broadcast`. This is transparent + to the compute stages — TP is internal to the worker. + +### 2.3 Request Flow + +A single video generation request follows this path: + +``` +Client Orchestrator Encoder-k Denoiser-j VAE-i + │ │ │ │ │ + │─POST /v1/videos/──────►│ │ │ │ + │ generations │ │ │ │ + │ [acquire semaphore] │ │ │ + │ │ │ │ │ + │ │──EncoderRequest──────►│ │ │ + │ │ (prompt, cfg) [encode text] │ │ + │ │ [register readable] │ │ + │ │◄──NIXL metadata──────│ │ │ + │ │ (~1.5 KB) [hold buffer] │ │ + │ │ │ │ │ + │ │──DenoiserRequest───────────────────►│ │ + │ │ (NIXL meta, params) │ │ + │ │ │ [RDMA pull embeddings] │ + │ │ │─────────────►│ │ + │ │ [completion notif] │ │ + │ │ [release buffer] │ │ + │ │ │ [denoise N steps] │ + │ │ │ [register readable] │ + │ │◄──NIXL metadata────────────────────│ │ + │ │ │ [hold buffer] │ + │ │ │ │ │ + │ │──VAEDecodeRequest──────────────────────────────►│ + │ │ (NIXL meta, req_id) │ │ + │ │ │ │ [RDMA pull latents] + │ │ │ │───────────►│ + │ │ │ [completion notif] │ + │ │ │ [release buffer] [decode] + │ │◄──{video_path}──────────────────────────────────│ + │ [release semaphore] │ │ │ + │◄──{url, timings}───────│ │ │ │ +``` + +Only ~1.5 KB of NIXL metadata travels over Dynamo RPC between stages. +The actual tensor data (embeddings: ~tens of MB, latents: ~hundreds of +MB) transfers GPU-to-GPU via NIXL RDMA without CPU involvement. + +**Buffer lifecycle:** The sender holds the GPU buffer (via `readable_op` +in `_active_readables`) until the receiver completes the RDMA pull and +the sender receives a completion notification. `_poll_completed()` +checks and releases finished buffers at the start of each new request. + +### 2.4 Pipeline Parallelism + +An `asyncio.Semaphore(pipeline_depth)` gates admission. Each stage has +its own independent worker pool, so multiple requests overlap: + +``` +Time ──────────────────────────────────────────────────► + +Req A: [Enc-0][======Den-0======][VAE-0] +Req B: [Enc-0][======Den-1======][VAE-0] +Req C: [Enc-0][======Den-2======][VAE-0] +Req D: [Enc-0][======Den-0======][VAE-0] +``` + +- `pipeline_depth` defaults to `MAX_PIPELINE_DEPTH` (env, default 4) + or the total number of workers across all stages. +- Each request independently acquires workers from each stage's pool. +- The denoiser is typically the bottleneck (many diffusion steps), so + encoder and VAE workers are freed quickly to serve other requests. + + +## 3. Component Design + +### 3.1 Orchestrator + +`orchestrator/run_disagg.py` — aiohttp HTTP server that chains stages. + +**Endpoints:** + +| Method | Path | Description | +|---|---|---| +| `POST` | `/v1/videos/generations` | Submit generation request | +| `GET` | `/health` | Orchestrator liveness | +| `GET` | `/health/stages` | Per-stage health (queries each worker) | +| `GET` | `/pipeline/status` | Active requests, queue depth, latencies | +| `GET` | `/videos/` | Serve generated video file | + +**Key components:** + +- `PipelineTracker` — Tracks active requests per stage, completed/failed + counts, and rolling average stage latencies. +- `dispatch_with_retry(mgr, request_id, json)` — Acquire worker → + dispatch → on failure retry on a different worker (up to + `STAGE_DISPATCH_RETRIES` attempts). +- `admission = asyncio.Semaphore(pipeline_depth)` — Limits concurrent + in-flight requests across the entire pipeline. + +### 3.2 WorkerManager + +`orchestrator/worker_manager.py` — Per-stage worker pool with +busy/idle tracking. + +- Backed by `asyncio.Queue` (idle pool) — `acquire_worker()` awaits + the queue, `dispatch()` returns the worker to it on completion. +- `_call_direct()` sends to a specific worker via + `client.direct(json, worker_id)`. +- `status()` returns per-worker `{id, status, request_id, completed, + avg_latency_s}` plus stage-level `{queue_depth, completed, failed}`. + +### 3.3 Tensor Transfer — NIXL + +`workers/nixl_transfer.py` — GPU-direct RDMA transfer between stages. + +**Connection management — `PersistentConnector`:** + +Each `NixlTensorSender` and `NixlTensorReceiver` owns a +`PersistentConnector` instance (subclass of `nixl_connect.Connector`) +that overrides `_create_connection()` to reuse a single +`Connection` (= one `nixl_agent` / UCX endpoint) across all operations. + +``` +Per-process NIXL agents (e.g. denoiser with TP=2): + + Rank 0 subprocess: + ├── NixlReceiveStage → NixlTensorReceiver → PersistentConnector → agent A + └── NixlSendStage → NixlTensorSender → PersistentConnector → agent B + + Rank 1 subprocess: + └── (no NIXL agents — receives tensors via TP broadcast from rank 0) +``` + +- Agents are created eagerly at stage construction time, before any + request arrives. This ensures UCX endpoints are fully initialized. +- Each sender/receiver has its own `PersistentConnector` → its own + agent. This avoids UCX state conflicts between read and write roles. +- Agents are reused across all requests — never recreated. + +**UCX transport configuration:** + +For intra-node transfers, set `UCX_TLS=cuda_ipc,tcp,self,cuda_copy,cma` +to force NVLink (cuda_ipc) instead of IB RDMA. The default +`UCX_TLS=all` may select IB which fails across NUMA boundaries on +multi-socket systems. + +**Sender (`NixlTensorSender`):** + +```python +sender = NixlTensorSender() # creates agent eagerly +meta, readable_op = sender.send({"latents": t}) # → (metadata, handle) +# caller holds readable_op until receiver completes RDMA pull +``` + +1. Flatten all tensors into a single contiguous GPU buffer. +2. Create NIXL `Descriptor` and register as readable. +3. Return `(metadata_dict, readable_op)` — caller must hold + `readable_op` to prevent GC of the descriptor and GPU buffer. + +**Receiver (`NixlTensorReceiver`):** + +```python +receiver = NixlTensorReceiver() # creates agent eagerly +tensors = receiver.recv(meta, device="cuda") # → {"latents": tensor} +``` + +1. Allocate flat GPU buffer on target device. +2. `begin_read(rdma_meta, descriptor)` → one-sided RDMA pull. +3. `wait_for_completion()` — blocks until transfer finishes. +4. Slice flat buffer into individual tensors. + +**Buffer lifecycle:** + +``` +Sender (Encoder/Denoiser) Receiver (Denoiser/VAE) +───────────────────────── ─────────────────────── +create_readable(descriptor) + → register GPU buffer + → return (meta, readable_op) + begin_read(meta) + → RDMA pull from sender GPU + wait_for_completion() + → transfer done + ← completion notification +_poll_completed() + → status == COMPLETE + → drop readable_op → GC buffer +``` + +### 3.4 Worker Interface + +All three workers follow an identical pattern: + +```python +@dynamo_worker(enable_nats=False) +async def worker(runtime: DistributedRuntime): + # 1. Launch backend subprocess + processes, client, server_args = launch_stage_server(...) + + # 2. Define Dynamo RPC handlers + async def handle_generate(request, context): + output = await client.forward([build_req(...)]) + yield result_dict # JSON with NIXL metadata or video path + + # 3. Serve endpoints + gen_ep = runtime.endpoint("disagg_diffusion..generate") + await gen_ep.serve_endpoint(handle_generate) +``` + +**Stage IO contract:** + +| Stage | Input | Output | +|---|---|---| +| Encoder | prompt, guidance_scale | `transfer_meta` (embedding metadata) | +| Denoiser | `transfer_meta` + inference params | `transfer_meta` (latent metadata) | +| VAE | `transfer_meta` + request_id | `video_path` | + +Fallback: when NIXL is unavailable, tensors are serialized via +`torch.save()` + base64 in the `tensor_data` field. + + +## 4. Implementation: SGLang Backend + +### 4.1 PartialGPUWorker + +`workers/partial_gpu_worker.py` — Extends SGLang's `GPUWorker` to load +only the modules each stage needs. + +- Overrides only `init_device_and_model()`. All other `GPUWorker` + behavior (forward execution, memory analysis, LoRA) is inherited. +- `build_partial_pipeline()` dynamically loads only the specified + modules (e.g. `["transformer", "scheduler"]` for denoiser). + +**NIXL pipeline stages:** + +- `NixlReceiveStage(PipelineStage)` — Prepended at denoiser/VAE entry. + Rank 0: RDMA-pulls tensors via `NixlTensorReceiver`. + TP > 1: broadcasts pulled tensors to other ranks via + `torch.distributed.broadcast`. + Falls back to `_device_move()` for ZMQ path. + +- `NixlSendStage(PipelineStage)` — Appended at encoder/denoiser exit. + Rank 0 only: registers tensors as NIXL-readable, stores `readable_op` + in `_active_readables`, returns metadata in `OutputBatch`. + Rank > 0: returns empty `OutputBatch` (only rank 0's output is used). + +**Stage builders:** + +| Function | Stages | +|---|---| +| `build_encoder_stages()` | `TextEncodingStage` → `NixlSendStage` | +| `build_denoiser_stages()` | `NixlReceiveStage` → `LatentPrep` → `TimestepPrep` → `DenoisingStage` → `NixlSendStage` | +| `build_vae_stages()` | `NixlReceiveStage` → `DecodingStage` | + +### 4.2 StageClient + +`workers/sglang_utils.py:StageClient` — Async ZMQ REQ/REP client. + +- `asyncio.Lock` serializes concurrent calls (ZMQ REQ requires strict + send/recv alternation). +- Configurable timeout: `STAGE_FORWARD_TIMEOUT_S` (default 120s). +- On timeout: `_reset_socket()` closes and reconnects the ZMQ socket + to recover from the broken REQ state (prevents cascading EFSM errors). + +### 4.3 HunyuanVideo Specifics + +- **Dual encoder:** auto-detects `text_encoder` + `text_encoder_2` + (Llama + CLIP) from `model_index.json`. +- **HunyuanConfig patching:** wraps `HunyuanConfig.__init__` to supply + `task_type=T2V` when omitted. +- **Triton norm workaround:** wraps `norm_infer` to call + `.contiguous()` on non-contiguous tensors from attention reshapes. +- **Component config sync:** reads all component `config.json` files + to ensure correct parameters even for non-loaded components. + + +## 5. Deployment + +### 5.1 Example: 8-GPU HunyuanVideo + +```bash +# 1 encoder (GPU 0) + 3 denoisers TP=2 (GPU 1,2 / 3,4 / 5,6) + 1 VAE (GPU 7) +UCX_TLS=cuda_ipc,tcp,self,cuda_copy,cma \ +GPU_ENC=0 GPU_DEN="1,2;3,4;5,6" GPU_VAE=7 PORT=8091 \ + ./run_all.sh +``` + +### 5.2 Environment Variables + +| Variable | Default | Description | +|---|---|---| +| `GPU_ENC` | `0` | GPU(s) for encoder workers (`;`-separated for multi-worker) | +| `GPU_DEN` | `1,2` | GPU(s) for denoiser workers (`,` = TP within worker, `;` = multiple workers) | +| `GPU_VAE` | `3` | GPU(s) for VAE workers | +| `PORT` | `8080` | HTTP port | +| `UCX_TLS` | `all` | UCX transport list — set to `cuda_ipc,tcp,self,cuda_copy,cma` for intra-node | +| `MAX_PIPELINE_DEPTH` | `4` | Max concurrent requests in pipeline | +| `STAGE_FORWARD_TIMEOUT_S` | `120` | ZMQ timeout per stage request (seconds) | +| `WORKER_SETTLE_S` | `30` | Seconds to wait for slow workers after first discovery | +| `STAGE_DISPATCH_RETRIES` | `2` | Retry count on stage dispatch failure | +| `DISABLE_NIXL` | `false` | Force ZMQ fallback (disable NIXL) | + +### 5.3 Measured Performance (HunyuanVideo, 9 frames, 3 steps) + +| Stage | Workers | Avg Latency | +|---|---|---| +| Encoder | 1 × GPU | 0.33s | +| Denoiser | 3 × TP=2 | 3.7s | +| VAE | 1 × GPU | 3.8s | +| **End-to-end** | | **~7.9s** | +| **Throughput (20 concurrent)** | | **~0.25 req/s** (VAE-bound) | + + +## 6. Roadmap + +- [x] Multi-worker scaling — N workers per stage, etcd auto-discovery +- [x] Pipeline parallelism — overlapping requests across stages +- [x] NIXL GPU-direct transfer — GPU-to-GPU RDMA, only metadata over RPC +- [x] TP support — TP broadcast in NixlReceiveStage, rank-0-only send +- [x] PersistentConnector — stable NIXL connections, no per-request agent churn +- [ ] Runtime scaling — external auto-scaler based on `/pipeline/status` +- [ ] Smart routing — load-aware dispatch, affinity, request priority +- [ ] Fault tolerance — health-check + eviction, graceful degradation +- [ ] Streaming decode output + richer observability +- [ ] Multi-model support — omni-modal pipelines, heterogeneous stage graphs + + +## 7. File Map + +``` +examples/disagg_diffusion/ +├── DESIGN.md ← this document +├── run_all.sh ← launch etcd + all workers + orchestrator +├── stress_test.sh ← concurrent load testing script +│ +├── orchestrator/ +│ ├── run_disagg.py ← HTTP server, pipeline dispatch, admission control +│ ├── worker_manager.py ← WorkerManager: idle-pool, acquire/dispatch/status +│ └── request_validation.py ← Pydantic-based HTTP request validation +│ +└── workers/ + ├── protocol.py ← Pydantic request/response models (Dynamo RPC) + ├── encoder_worker.py ← Encoder Dynamo worker + ├── denoiser_worker.py ← Denoiser Dynamo worker (supports TP) + ├── vae_worker.py ← VAE Dynamo worker + ├── nixl_transfer.py ← NixlTensorSender / NixlTensorReceiver (RDMA) + ├── partial_gpu_worker.py ← PartialGPUWorker, NixlSend/ReceiveStage, launcher + └── sglang_utils.py ← StageClient, build_partial_pipeline, launch_stage_server +``` diff --git a/examples/disagg_diffusion/README.md b/examples/disagg_diffusion/README.md index 72a8793cc506..f756ea006cdb 100644 --- a/examples/disagg_diffusion/README.md +++ b/examples/disagg_diffusion/README.md @@ -1,60 +1,180 @@ -# Disaggregated Diffusion Inference POC +# Disaggregated Diffusion Inference (HunyuanVideo) -Split a monolithic diffusion pipeline (Text Encoder → Transformer → VAE) into -independent stages that can run on separate GPUs and scale independently. +Split a monolithic video diffusion pipeline into independent Dynamo workers on separate GPUs. +Tensor data transfers between stages use **NIXL RDMA** (GPU-direct); only small metadata +travels over Dynamo RPC. -Design doc: [docs/design/disaggregated_diffusion.md](../../docs/design/disaggregated_diffusion.md) +Supports HunyuanVideo (13B, dual Llama+CLIP encoder) and Wan2.2-TI2V models. -## Phases +## Architecture -### Phase 0: Offline Validation (no Dynamo) +``` +┌─────────────────────────────────────────────────────────────────────┐ +│ Orchestrator (HTTP API) │ +│ run_disagg.py │ +└──────┬──────────────────────┬───────────────────────────┬───────────┘ + │ Dynamo RPC │ Dynamo RPC │ Dynamo RPC + │ (metadata) │ (metadata) │ (metadata) + ▼ ▼ ▼ +┌──────────────┐ ┌─────────────────────┐ ┌──────────────────────┐ +│ Encoder │ │ Denoiser │ │ VAE Decoder │ +│ Worker │ │ Worker │ │ Worker │ +│ │ │ │ │ │ +│ GPU 0 │ │ GPU 1,2 (TP=2) │ │ GPU 3 │ +│ ~18 GB VRAM │ │ ~24 GB VRAM/GPU │ │ ~6 GB VRAM │ +│ │ │ │ │ │ +│ Llama 8B │ │ HunyuanVideo DiT │ │ 3D VAE │ +│ + CLIP │ │ 9.5B params │ │ │ +└──────┬───────┘ └──────┬──────────────┘ └──────┬───────────────┘ + │ │ │ + └──── NIXL RDMA ──►└──── NIXL RDMA ─────────►│ + (embeddings, (latents, + GPU-direct) GPU-direct) +``` + +### Request Flow + +``` + 1. User ──POST /v1/videos/generations──► Orchestrator + 2. Orchestrator ──EncoderRequest──► Encoder Worker + 3. Encoder: TextEncoding → NixlSendStage (register embeddings on GPU) + 4. Encoder ──{nixl_metadata}──► Orchestrator + 5. Orchestrator ──DenoiserRequest + nixl_meta──► Denoiser Worker + 6. Denoiser: NixlReceive (RDMA pull embeddings) → LatentPrep → Denoise (N steps) → NixlSend + 7. Denoiser ──{nixl_metadata}──► Orchestrator + 8. Orchestrator ──VAERequest + nixl_meta──► VAE Worker + 9. VAE: NixlReceive (RDMA pull latents) → Decode → Save MP4 +10. VAE ──{video_path}──► Orchestrator ──► User +``` + +### Worker Internal Architecture + +Each worker wraps an SGLang Scheduler subprocess: + +``` +Dynamo Worker Process (e.g. encoder_worker.py) +├── @dynamo_worker +│ └── serve_endpoint("generate") ← Dynamo RPC from orchestrator +│ └── StageClient.forward() ← ZMQ to local Scheduler +│ +└── SGLang Scheduler subprocess ← spawned by launch_partial_server() + └── PartialGPUWorker + ├── TextEncodingStage ← model inference + └── NixlSendStage ← register tensors for RDMA +``` + +### Pipeline Parallelism + +Multiple requests overlap across stages: + +``` +Request 1: [ Encoder ] ──► [ Denoiser ~~~~~~~~ ] ──► [ VAE ] +Request 2: [ Encoder ] ──► [ Denoiser ~~~~~~~~ ] ──► [ VAE ] +Request 3: [ Encoder ] ──► [ Denoiser ~~~~~~~~ ] +``` + +### Multi-Worker Scaling + +Each stage supports **multiple workers** — just launch more processes with the +same Dynamo component name. They register via etcd and the orchestrator +round-robins requests automatically. Use `;` in GPU specs to separate workers: -Proves that diffusers supports split execution: encode, denoise, and VAE decode -can run independently with serialized intermediate tensors. +``` + ┌─ Encoder_0 (GPU 0) ─┐ ┌─ Denoiser_0 TP=2 (GPU 1,2) ─┐ ┌─ VAE_0 (GPU 3) ─┐ + Request ─┤ ├──►─┤ ├──►─┤ ├─► Video + └─ Encoder_1 (GPU 4) ─┘ └─ Denoiser_1 TP=2 (GPU 5,6) ─┘ └─ VAE_1 (GPU 7) ─┘ + round-robin round-robin round-robin +``` ```bash -# Single GPU, ~24 GB VRAM for FLUX.1-schnell -python phase0_validate/validate_split.py \ - --model black-forest-labs/FLUX.1-schnell \ - --prompt "A photo of a cat sitting on a windowsill" \ - --output-dir /tmp/disagg_validate +# 8 GPU — 2 workers per stage +GPU_ENC="0;4" GPU_DEN="1,2;5,6" GPU_VAE="3;7" ./run_all.sh --test --quick + +# Asymmetric (1 encoder, 3 denoisers, 1 VAE) +GPU_ENC="0" GPU_DEN="1,2;3,4;5,6" GPU_VAE="7" ./run_all.sh ``` -### Phase 1: Dynamo Stage Workers +Single-worker specs (no `;`) are fully backward compatible. + +No pre-pairing is needed — each `send()` allocates a new NIXL buffer and +the `_keep_alive` coroutine holds it until the receiver's RDMA pull completes, +so any sender→receiver combination is safe. + +## Quick Start -Three independent Dynamo workers, each loading only its model component: +One script launches everything (etcd + 3 workers + orchestrator): ```bash -# Terminal 1: Encoder Worker (loads CLIP + T5, ~12 GB) -python phase1_workers/encoder_worker.py --model black-forest-labs/FLUX.1-schnell +conda activate omni +export HF_HUB_CACHE=/path/to/huggingface/hub -# Terminal 2: Denoiser Worker (loads Transformer, ~24 GB) -python phase1_workers/denoiser_worker.py --model black-forest-labs/FLUX.1-schnell +# Launch all services + send a test request +./run_all.sh --test -# Terminal 3: VAE Worker (loads VAE, ~1 GB) -python phase1_workers/vae_worker.py --model black-forest-labs/FLUX.1-schnell -``` +# Quick smoke test (9 frames, 3 steps, ~30s) +./run_all.sh --test --quick -### Phase 2: Orchestrator +# Just launch services (no test request) +./run_all.sh +``` -Chains the three stage endpoints into an end-to-end generation pipeline: +Or launch each service manually: ```bash -python phase2_orchestrator/run_disagg.py \ - --prompt "A photo of a cat sitting on a windowsill" \ - --output /tmp/disagg_output.png +# Terminal 0: etcd +etcd --data-dir /tmp/etcd_disagg --listen-client-urls http://0.0.0.0:2379 + +# Terminal 1-3: Workers +CUDA_VISIBLE_DEVICES=0 python workers/encoder_worker.py +CUDA_VISIBLE_DEVICES=1,2 python workers/denoiser_worker.py +CUDA_VISIBLE_DEVICES=3 python workers/vae_worker.py + +# Terminal 4: Orchestrator +python orchestrator/run_disagg.py ``` -Or use the all-in-one launch script: +Generate a video (61 frames, 50 steps, 544x960 by default): ```bash -bash launch/run_all.sh black-forest-labs/FLUX.1-schnell "A photo of a cat" +curl -X POST http://localhost:8080/v1/videos/generations \ + -H "Content-Type: application/json" \ + -d '{"prompt": "A golden retriever running on a sunny beach with waves crashing"}' ``` +## Workers + +| Worker | Model Component | VRAM | Dynamo Endpoint | +|--------|----------------|------|-----------------| +| `encoder_worker.py` | Llama 8B + CLIP text encoders | ~18 GB | `disagg_diffusion.encoder.generate` | +| `denoiser_worker.py` | HunyuanVideo DiT (TP=2) | ~24 GB/GPU | `disagg_diffusion.denoiser.generate` | +| `vae_worker.py` | 3D VAE decoder | ~6 GB | `disagg_diffusion.vae.generate` | + +## Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `MODEL_PATH` | `hunyuanvideo-community/HunyuanVideo` | HuggingFace model ID or local path | +| `GPU_ENC` / `GPU_DEN` / `GPU_VAE` | `0` / `1,2` / `3` | GPU assignment per stage (use `;` to launch multiple workers, e.g. `"0;4"`) | +| `TP_SIZE` | auto from `GPU_DEN` | Tensor parallelism for denoiser | +| `PORT` | `8080` | Orchestrator HTTP port | +| `OUTPUT_DIR` | `/tmp/disagg_videos` | Video output directory | + ## Supported Models -Any diffusers pipeline that exposes `encode_prompt()` and supports -`prompt_embeds` / `output_type="latent"`. Tested with: +- **`hunyuanvideo-community/HunyuanVideo`** — 13B, dual encoder (Llama 8B + CLIP), recommended +- `Wan-AI/Wan2.2-TI2V-5B-Diffusers` — 5B, single encoder + +## Roadmap -- `black-forest-labs/FLUX.1-schnell` (recommended, 4 steps) -- `stabilityai/stable-diffusion-3.5-medium` +- [x] **Multi-worker scaling** — launch N workers per stage via `run_all.sh`; Dynamo auto-discovers and round-robins +- [ ] **Dynamic scaling** — auto-scale workers based on queue depth, add/remove denoiser replicas +- [ ] **Streaming output** — stream decoded frames to client as they are produced +- [ ] **Orchestrator improvements** — smarter scheduling, request priority, load balancing across replicas +- [ ] **Metrics & observability** — per-stage latency, GPU utilization, NIXL throughput, Prometheus export +- [ ] **Request cancellation** — cancel in-flight requests, free GPU resources immediately + +## Dependencies + +```bash +pip install ai-dynamo-runtime sglang imageio imageio-ffmpeg pyzmq setproctitle +``` diff --git a/examples/disagg_diffusion/launch/run_all.sh b/examples/disagg_diffusion/launch/run_all.sh deleted file mode 100755 index aa3f131b3285..000000000000 --- a/examples/disagg_diffusion/launch/run_all.sh +++ /dev/null @@ -1,86 +0,0 @@ -#!/usr/bin/env bash -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Disaggregated Diffusion POC — all-in-one launcher -# -# Starts three stage workers (Encoder, Denoiser, VAE) as background -# processes, waits for them to be ready, then runs the orchestrator. -# -# Requirements: -# - 3 GPUs (or 1 large GPU >48 GB — set SINGLE_GPU=1) -# - diffusers, transformers, torch, dynamo runtime, uvloop -# -# Usage: -# bash launch/run_all.sh [MODEL] [PROMPT] -# -# Examples: -# bash launch/run_all.sh -# bash launch/run_all.sh black-forest-labs/FLUX.1-schnell "A sunset over mountains" -# SINGLE_GPU=1 bash launch/run_all.sh # all stages on GPU 0 - -set -euo pipefail - -export MODEL_PATH="${1:-black-forest-labs/FLUX.1-schnell}" -export PROMPT="${2:-A photo of a cat sitting on a windowsill}" -SCRIPT_DIR="$(cd "$(dirname "$0")/.." && pwd)" -WORKERS_DIR="${SCRIPT_DIR}/phase1_workers" -ORCH_DIR="${SCRIPT_DIR}/phase2_orchestrator" - -echo "============================================" -echo " Disaggregated Diffusion POC" -echo " Model: ${MODEL_PATH}" -echo " Prompt: ${PROMPT}" -echo "============================================" -echo "" - -PIDS=() - -cleanup() { - echo "" - echo "Shutting down workers …" - for pid in "${PIDS[@]}"; do - kill "$pid" 2>/dev/null || true - done - wait 2>/dev/null - echo "Done." -} -trap cleanup EXIT - -if [[ "${SINGLE_GPU:-0}" == "1" ]]; then - GPU_ENC=0; GPU_DEN=0; GPU_VAE=0 -else - GPU_ENC=0; GPU_DEN=1; GPU_VAE=2 -fi - -# --- Start workers -------------------------------------------------------- - -echo "[1/3] Starting Encoder Worker (GPU ${GPU_ENC}) …" -CUDA_VISIBLE_DEVICES=${GPU_ENC} python "${WORKERS_DIR}/encoder_worker.py" \ - 2>&1 | sed 's/^/ [encoder] /' & -PIDS+=($!) - -echo "[2/3] Starting Denoiser Worker (GPU ${GPU_DEN}) …" -CUDA_VISIBLE_DEVICES=${GPU_DEN} python "${WORKERS_DIR}/denoiser_worker.py" \ - 2>&1 | sed 's/^/ [denoiser] /' & -PIDS+=($!) - -echo "[3/3] Starting VAE Worker (GPU ${GPU_VAE}) …" -CUDA_VISIBLE_DEVICES=${GPU_VAE} python "${WORKERS_DIR}/vae_worker.py" \ - 2>&1 | sed 's/^/ [vae] /' & -PIDS+=($!) - -# Wait for model loading. In production use Dynamo health checks. -WAIT_SECS="${WAIT_SECS:-90}" -echo "" -echo "Waiting ${WAIT_SECS}s for workers to load models …" -sleep "${WAIT_SECS}" - -# --- Run orchestrator ----------------------------------------------------- - -echo "" -echo "Running orchestrator …" -echo "" - -export OUTPUT="/tmp/disagg_output.png" -python "${ORCH_DIR}/run_disagg.py" diff --git a/examples/disagg_diffusion/orchestrator/request_validation.py b/examples/disagg_diffusion/orchestrator/request_validation.py new file mode 100644 index 000000000000..bcb4dd18d81e --- /dev/null +++ b/examples/disagg_diffusion/orchestrator/request_validation.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Request validation helpers for disaggregated diffusion orchestrator.""" + +from protocol import GenerateRequest + + +def validate_generate_http_body(body: dict) -> dict: + """Validate and normalize POST /v1/videos/generations request body.""" + req = GenerateRequest.model_validate(body) + req_dict = req.model_dump() + resp_format = body.get("response_format", "url") + if resp_format not in {"url", "b64_json"}: + raise ValueError("response_format must be 'url' or 'b64_json'") + req_dict["response_format"] = resp_format + return req_dict diff --git a/examples/disagg_diffusion/orchestrator/run_disagg.py b/examples/disagg_diffusion/orchestrator/run_disagg.py new file mode 100755 index 000000000000..05b777bdaf6f --- /dev/null +++ b/examples/disagg_diffusion/orchestrator/run_disagg.py @@ -0,0 +1,408 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Disaggregated Diffusion Orchestrator — HTTP Server with Pipeline Parallelism + +Persistent server that accepts video generation requests and chains +Encoder → Denoiser → VAE via Dynamo RPC + NIXL RDMA. + +Pipeline parallelism: multiple requests can be in different stages +simultaneously. Per-stage WorkerManagers track busy/idle state for +each GPU worker and dispatch requests to specific instances via +``client.direct()``, enabling backpressure-aware scheduling and +per-worker observability through ``/pipeline/status``. + +Each stage worker wraps an SGLang Scheduler subprocess via +launch_partial_server(), supporting TP for the denoiser and NIXL RDMA +for GPU-direct tensor transfer between stages. + +Usage: + python run_disagg.py [--port 8080] + +API: + POST /v1/videos/generations + GET /health + GET /health/stages + GET /pipeline/status + GET /videos/ +""" + +import asyncio +import json +import logging +import os +import sys +import time +import uuid +from collections import defaultdict +from typing import Dict + +import uvloop + +sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "workers")) + +from protocol import ( # noqa: E402 + DenoiserRequest, EncoderRequest, VAEDecodeRequest, + HealthRequest, +) +from dynamo.runtime import DistributedRuntime, dynamo_worker # noqa: E402 +from worker_manager import WorkerManager # noqa: E402 +from pydantic import ValidationError # noqa: E402 +from request_validation import validate_generate_http_body # noqa: E402 + +logger = logging.getLogger(__name__) + +PORT = int(os.environ.get("PORT", "8080")) +HOST = os.environ.get("HOST", "0.0.0.0") +OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/tmp/disagg_videos") +MAX_PIPELINE_DEPTH = int(os.environ.get("MAX_PIPELINE_DEPTH", "4")) +STAGE_DISPATCH_RETRIES = int(os.environ.get("STAGE_DISPATCH_RETRIES", "2")) + + +async def query_stage_health(health_client, stage_name: str) -> dict: + """Query a stage's StageEngine health via its Dynamo health endpoint.""" + try: + health_req = HealthRequest() + stream = await health_client.generate(health_req.model_dump_json()) + result = None + async for chunk in stream: + data = chunk.data() if hasattr(chunk, "data") else chunk + if isinstance(data, str): + data = json.loads(data) + result = data + return {"stage": stage_name, **result} if result else {"stage": stage_name, "error": "empty response"} + except Exception as e: + return {"stage": stage_name, "error": str(e)} + + +class PipelineTracker: + """Tracks requests flowing through the 3-stage pipeline.""" + + STAGES = ("encoder", "denoiser", "vae") + + def __init__(self): + self._active: Dict[str, str] = {} + self._completed = 0 + self._failed = 0 + self._stage_times: Dict[str, list] = defaultdict(list) + self._lock = asyncio.Lock() + + async def enter(self, request_id: str, stage: str): + async with self._lock: + self._active[request_id] = stage + + async def leave(self, request_id: str, stage: str, elapsed: float): + async with self._lock: + self._stage_times[stage].append(elapsed) + if request_id in self._active and self._active[request_id] == stage: + if stage == "vae": + del self._active[request_id] + + async def mark_done(self, request_id: str): + async with self._lock: + self._active.pop(request_id, None) + self._completed += 1 + + async def mark_failed(self, request_id: str): + async with self._lock: + self._active.pop(request_id, None) + self._failed += 1 + + async def status(self) -> dict: + async with self._lock: + per_stage = defaultdict(list) + for rid, stage in self._active.items(): + per_stage[stage].append(rid) + avg_times = {} + for stage in self.STAGES: + times = self._stage_times[stage] + avg_times[stage] = round(sum(times) / len(times), 3) if times else 0 + return { + "active_requests": dict(per_stage), + "active_count": len(self._active), + "completed": self._completed, + "failed": self._failed, + "avg_stage_seconds": avg_times, + } + + +@dynamo_worker(enable_nats=False) +async def worker(runtime: DistributedRuntime): + encoder_client = await runtime.endpoint("disagg_diffusion.encoder.generate").client() + denoiser_client = await runtime.endpoint("disagg_diffusion.denoiser.generate").client() + vae_client = await runtime.endpoint("disagg_diffusion.vae.generate").client() + + encoder_health_client = await runtime.endpoint("disagg_diffusion.encoder.health").client() + denoiser_health_client = await runtime.endpoint("disagg_diffusion.denoiser.health").client() + vae_health_client = await runtime.endpoint("disagg_diffusion.vae.health").client() + + health_clients = { + "encoder": encoder_health_client, + "denoiser": denoiser_health_client, + "vae": vae_health_client, + } + + logger.info("Waiting for stage workers …") + await encoder_client.wait_for_instances() + await denoiser_client.wait_for_instances() + await vae_client.wait_for_instances() + + # Wait for additional workers that may still be registering. + # wait_for_instances() returns after the first instance; model loading + # times vary, so observe all stages for a short settle window. + WORKER_SETTLE_S = int(os.environ.get("WORKER_SETTLE_S", "30")) + if WORKER_SETTLE_S > 0: + import time as _time + deadline = _time.monotonic() + WORKER_SETTLE_S + prev_counts = (-1, -1, -1) + while _time.monotonic() < deadline: + cur_counts = ( + len(encoder_client.instance_ids()), + len(denoiser_client.instance_ids()), + len(vae_client.instance_ids()), + ) + if cur_counts != prev_counts: + logger.info( + "Discovered workers so far: encoder=%d denoiser=%d vae=%d", + cur_counts[0], cur_counts[1], cur_counts[2], + ) + prev_counts = cur_counts + await asyncio.sleep(2) + logger.info("Worker settle period done (%ds)", WORKER_SETTLE_S) + + # Discover registered worker instances per stage + enc_ids = encoder_client.instance_ids() + den_ids = denoiser_client.instance_ids() + vae_ids = vae_client.instance_ids() + logger.info( + "Workers: encoder=%d %s, denoiser=%d %s, vae=%d %s", + len(enc_ids), enc_ids, len(den_ids), den_ids, len(vae_ids), vae_ids, + ) + + os.makedirs(OUTPUT_DIR, exist_ok=True) + + managers: Dict[str, WorkerManager] = { + "encoder": WorkerManager("encoder", encoder_client, enc_ids), + "denoiser": WorkerManager("denoiser", denoiser_client, den_ids), + "vae": WorkerManager("vae", vae_client, vae_ids), + } + n_workers_total = sum(m.worker_count for m in managers.values()) + pipeline_depth = MAX_PIPELINE_DEPTH if MAX_PIPELINE_DEPTH > 0 else n_workers_total + admission = asyncio.Semaphore(pipeline_depth) + tracker = PipelineTracker() + + async def dispatch_with_retry( + mgr: WorkerManager, request_id: str, request_json: str, + ) -> tuple: + """Dispatch to a stage worker; on failure retry on a different worker.""" + for attempt in range(STAGE_DISPATCH_RETRIES + 1): + wid = await mgr.acquire_worker() + try: + result, elapsed = await mgr.dispatch(wid, request_id, request_json) + if "error" in result and result["error"]: + raise RuntimeError(result["error"]) + return result, elapsed, wid + except Exception as e: + logger.warning( + "[%s] %s worker %d failed (attempt %d/%d): %s", + request_id, mgr.stage_name, wid, attempt + 1, + STAGE_DISPATCH_RETRIES + 1, e, + ) + if attempt >= STAGE_DISPATCH_RETRIES: + raise + raise RuntimeError("unreachable") + + async def handle_generate(request: dict) -> dict: + request_id = str(uuid.uuid4())[:8] + seed = request.get("seed") or int(time.time()) % 1000000 + timings: Dict[str, float] = {} + + async with admission: + try: + # Stage 1: Encoder + enc_req = EncoderRequest( + prompt=request["prompt"], + negative_prompt=request.get("negative_prompt", ""), + guidance_scale=request.get("guidance_scale", 1.0), + ) + await tracker.enter(request_id, "encoder") + enc_resp, timings["encoder_s"], enc_wid = await dispatch_with_retry( + managers["encoder"], request_id, enc_req.model_dump_json(), + ) + await tracker.leave(request_id, "encoder", timings["encoder_s"]) + logger.info("[%s] Encoder (worker %d): %.2fs", request_id, enc_wid, timings["encoder_s"]) + + # Stage 2: Denoiser + den_req = DenoiserRequest( + transfer_meta=enc_resp.get("transfer_meta", {}), + tensor_data=enc_resp.get("tensor_data", {}), + height=request.get("height", 544), + width=request.get("width", 960), + num_frames=request.get("num_frames", 61), + num_inference_steps=request.get("num_inference_steps", 50), + guidance_scale=request.get("guidance_scale", 1.0), + seed=seed, + ) + await tracker.enter(request_id, "denoiser") + den_resp, timings["denoiser_s"], den_wid = await dispatch_with_retry( + managers["denoiser"], request_id, den_req.model_dump_json(), + ) + await tracker.leave(request_id, "denoiser", timings["denoiser_s"]) + logger.info("[%s] Denoiser (worker %d): %.2fs", request_id, den_wid, timings["denoiser_s"]) + + # Stage 3: VAE + vae_req = VAEDecodeRequest( + transfer_meta=den_resp.get("transfer_meta", {}), + tensor_data=den_resp.get("tensor_data", {}), + request_id=request_id, + ) + await tracker.enter(request_id, "vae") + vae_resp, timings["vae_s"], vae_wid = await dispatch_with_retry( + managers["vae"], request_id, vae_req.model_dump_json(), + ) + await tracker.leave(request_id, "vae", timings["vae_s"]) + logger.info("[%s] VAE (worker %d): %.2fs", request_id, vae_wid, timings["vae_s"]) + except Exception: + await tracker.mark_failed(request_id) + raise + + timings["total_s"] = round(sum(timings.values()), 3) + await tracker.mark_done(request_id) + logger.info("[%s] Total: %.2fs", request_id, timings["total_s"]) + + filename = vae_resp["video_path"] + resp_format = request.get("response_format", "url") + + if resp_format == "url": + data = [{"url": f"/videos/{filename}"}] + else: + import base64 + filepath = os.path.join(OUTPUT_DIR, filename) + with open(filepath, "rb") as f: + data = [{"b64_json": base64.b64encode(f.read()).decode("ascii")}] + + return { + "id": f"video-{request_id}", + "created": int(time.time()), + "data": data, + "timings": timings, + } + + from aiohttp import web + + async def handle_post(http_request: web.Request) -> web.Response: + try: + body = await http_request.json() + except Exception: + return web.json_response({"error": "invalid JSON body"}, status=400) + + try: + req_dict = validate_generate_http_body(body) + result = await handle_generate(req_dict) + return web.json_response(result) + except ValidationError as e: + return web.json_response( + {"error": "invalid request", "details": e.errors()}, + status=400, + ) + except ValueError as e: + return web.json_response({"error": str(e)}, status=400) + except Exception as e: + logger.error("Request failed: %s", e, exc_info=True) + return web.json_response({"error": str(e)}, status=500) + + async def handle_health(http_request: web.Request) -> web.Response: + return web.json_response({"status": "ok"}) + + async def handle_stages_health(http_request: web.Request) -> web.Response: + results = await asyncio.gather( + *[ + query_stage_health(client, name) + for name, client in health_clients.items() + ] + ) + return web.json_response({"stages": list(results)}) + + async def handle_pipeline_status(http_request: web.Request) -> web.Response: + pipeline = await tracker.status() + stages = {name: mgr.status() for name, mgr in managers.items()} + pipeline["stages"] = stages + pipeline["pipeline_depth"] = { + "active": pipeline["active_count"], + "max": pipeline_depth, + } + return web.json_response(pipeline) + + async def handle_video(http_request: web.Request) -> web.Response: + filename = http_request.match_info["filename"] + if "/" in filename or "\\" in filename or ".." in filename: + return web.json_response({"error": "invalid filename"}, status=400) + filepath = os.path.join(OUTPUT_DIR, filename) + if not os.path.exists(filepath): + return web.json_response({"error": "not found"}, status=404) + return web.FileResponse(filepath, headers={"Content-Type": "video/mp4"}) + + async def handle_index(http_request: web.Request) -> web.Response: + files = sorted( + [f for f in os.listdir(OUTPUT_DIR) if f.endswith(".mp4")], + key=lambda f: os.path.getmtime(os.path.join(OUTPUT_DIR, f)), + reverse=True, + ) if os.path.isdir(OUTPUT_DIR) else [] + latest = files[0] if files else None + video_tag = ( + f'' + if latest else "

No videos generated yet.

" + ) + history = "".join( + f'
  • {f}
    • ' for f in files[:20] + ) + html = ( + "Disagg Diffusion" + "" + f"

    Latest Video

    {video_tag}" + f"

    History

      {history}
    " + ) + return web.Response(text=html, content_type="text/html") + + http_app = web.Application() + http_app.router.add_get("/", handle_index) + http_app.router.add_post("/v1/videos/generations", handle_post) + http_app.router.add_get("/health", handle_health) + http_app.router.add_get("/health/stages", handle_stages_health) + http_app.router.add_get("/pipeline/status", handle_pipeline_status) + http_app.router.add_get("/videos/{filename}", handle_video) + + runner = web.AppRunner(http_app) + await runner.setup() + + bound_port = PORT + for attempt in range(10): + try: + site = web.TCPSite(runner, HOST, bound_port, reuse_address=True) + await site.start() + break + except OSError as e: + if e.errno == 98 and attempt < 9: + logger.warning("Port %d in use, trying %d", bound_port, bound_port + 1) + bound_port += 1 + else: + raise + + logger.info("Server listening on http://%s:%d (pipeline depth=%d)", HOST, bound_port, pipeline_depth) + logger.info(" GET / <- latest video preview") + logger.info(" POST /v1/videos/generations") + logger.info(" GET /health") + logger.info(" GET /health/stages") + logger.info(" GET /pipeline/status") + logger.info(" GET /videos/.mp4") + + await asyncio.Event().wait() + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(name)s] %(levelname)s %(message)s") + uvloop.install() + asyncio.run(worker()) diff --git a/examples/disagg_diffusion/orchestrator/worker_manager.py b/examples/disagg_diffusion/orchestrator/worker_manager.py new file mode 100644 index 000000000000..14ca2527a1ed --- /dev/null +++ b/examples/disagg_diffusion/orchestrator/worker_manager.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Per-stage worker pool manager with busy/idle tracking and targeted dispatch. + +Each pipeline stage (encoder, denoiser, vae) gets its own WorkerManager +that tracks which Dynamo worker instances are idle or busy and dispatches +requests to specific workers via ``client.direct()``. This replaces the +previous ``asyncio.Semaphore``-based flow control with explicit state that +is queryable via ``/pipeline/status``. +""" + +import asyncio +import json +import logging +import time +from dataclasses import dataclass +from typing import Any, Dict, List, Literal, Optional, Tuple + +logger = logging.getLogger(__name__) + + +@dataclass +class WorkerState: + """Runtime state for a single Dynamo worker instance.""" + + worker_id: int + status: Literal["idle", "busy"] = "idle" + current_request_id: Optional[str] = None + completed_count: int = 0 + total_latency_s: float = 0.0 + + +class WorkerManager: + """Manages the worker pool for one pipeline stage. + + Provides acquire/dispatch semantics: callers first ``acquire_worker()`` + (blocks until one is idle), then ``dispatch()`` to send the request to + that specific worker. On completion (or failure) the worker is + automatically returned to the idle pool. + """ + + def __init__(self, stage_name: str, client: Any, worker_ids: List[int]): + self._stage_name = stage_name + self._client = client + self._workers: Dict[int, WorkerState] = { + wid: WorkerState(worker_id=wid) for wid in worker_ids + } + self._idle_queue: asyncio.Queue[int] = asyncio.Queue() + for wid in worker_ids: + self._idle_queue.put_nowait(wid) + + # Observability counters + self._total_completed = 0 + self._total_failed = 0 + self._waiting_count = 0 # requests blocked waiting for an idle worker + + @property + def stage_name(self) -> str: + return self._stage_name + + @property + def worker_count(self) -> int: + return len(self._workers) + + async def acquire_worker(self) -> int: + """Block until a worker becomes idle. Returns worker_id.""" + self._waiting_count += 1 + try: + wid = await self._idle_queue.get() + finally: + self._waiting_count -= 1 + return wid + + async def dispatch( + self, worker_id: int, request_id: str, request_json: str + ) -> Tuple[dict, float]: + """Send request to a specific worker via ``client.direct()`` and track state. + + Returns (result_dict, elapsed_seconds). + """ + ws = self._workers[worker_id] + ws.status = "busy" + ws.current_request_id = request_id + t0 = time.monotonic() + try: + result = await self._call_direct(worker_id, request_json) + elapsed = time.monotonic() - t0 + ws.completed_count += 1 + ws.total_latency_s += elapsed + self._total_completed += 1 + return result, elapsed + except Exception: + self._total_failed += 1 + raise + finally: + ws.status = "idle" + ws.current_request_id = None + self._idle_queue.put_nowait(worker_id) + + async def _call_direct(self, worker_id: int, request_json: str) -> dict: + """Issue a Dynamo RPC to a specific worker instance.""" + result = None + stream = await self._client.direct(request_json, worker_id) + async for chunk in stream: + data = chunk.data() if hasattr(chunk, "data") else chunk + if isinstance(data, str): + data = json.loads(data) + result = data + if result is None: + raise RuntimeError( + f"Empty response from {self._stage_name} worker {worker_id}" + ) + return result + + def status(self) -> dict: + """Snapshot of per-worker state for ``/pipeline/status``.""" + workers = [] + for ws in self._workers.values(): + avg = ( + round(ws.total_latency_s / ws.completed_count, 3) + if ws.completed_count + else 0 + ) + workers.append( + { + "id": ws.worker_id, + "status": ws.status, + "request_id": ws.current_request_id, + "completed": ws.completed_count, + "avg_latency_s": avg, + } + ) + return { + "stage": self._stage_name, + "workers": workers, + "queue_depth": self._waiting_count, + "completed": self._total_completed, + "failed": self._total_failed, + } diff --git a/examples/disagg_diffusion/phase1_workers/denoiser_worker.py b/examples/disagg_diffusion/phase1_workers/denoiser_worker.py deleted file mode 100755 index ebc29696ba20..000000000000 --- a/examples/disagg_diffusion/phase1_workers/denoiser_worker.py +++ /dev/null @@ -1,142 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Disaggregated Diffusion — Denoiser Worker - -Loads only the Transformer (DiT) and scheduler. Accepts pre-computed -embeddings and returns denoised latents (skips VAE decode). - -Usage: - python denoiser_worker.py --model black-forest-labs/FLUX.1-schnell -""" - -import asyncio -import logging -import os -import sys - -import torch -import uvloop - -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - -from protocol import ( # noqa: E402 - DenoiserRequest, - DenoiserResponse, - b64_to_tensors, - tensors_to_b64, -) - -from dynamo.runtime import DistributedRuntime, dynamo_endpoint, dynamo_worker # noqa: E402 - -logger = logging.getLogger(__name__) - -MODEL_PATH = os.environ.get("MODEL_PATH", "black-forest-labs/FLUX.1-schnell") -DEVICE = os.environ.get("DEVICE", "cuda") - - -class DenoiserStage: - """Denoiser stage: embeddings → latents (no text encoder, no VAE).""" - - def __init__(self): - self.pipe = None - self.vae_scaling_factor = 1.0 - self.vae_shift_factor = None - - def load_model(self): - from diffusers import FluxPipeline - - logger.info("Loading transformer from %s …", MODEL_PATH) - self.pipe = FluxPipeline.from_pretrained( - MODEL_PATH, torch_dtype=torch.bfloat16 - ) - self.pipe.to(DEVICE) - - # Capture VAE config before discarding - self.vae_scaling_factor = self.pipe.vae.config.scaling_factor - self.vae_shift_factor = getattr(self.pipe.vae.config, "shift_factor", None) - - # Free text encoders + VAE - self.pipe.text_encoder = None - self.pipe.text_encoder_2 = None - self.pipe.tokenizer = None - self.pipe.tokenizer_2 = None - self.pipe.vae = None - torch.cuda.empty_cache() - - vram = torch.cuda.memory_allocated() / 1e6 - logger.info("Denoiser ready — VRAM: %.0f MB (transformer only)", vram) - - @dynamo_endpoint(DenoiserRequest, DenoiserResponse) - async def generate(self, request: DenoiserRequest): - logger.info( - "Denoising %dx%d, %d steps, seed=%d", - request.width, request.height, - request.num_inference_steps, request.seed, - ) - - embeddings = b64_to_tensors(request.embeddings_b64, DEVICE) - generator = torch.Generator(device=DEVICE).manual_seed(request.seed) - - loop = asyncio.get_event_loop() - result = await loop.run_in_executor( - None, - lambda: self.pipe( - prompt_embeds=embeddings["prompt_embeds"], - pooled_prompt_embeds=embeddings["pooled_prompt_embeds"], - num_inference_steps=request.num_inference_steps, - guidance_scale=request.guidance_scale, - generator=generator, - height=request.height, - width=request.width, - output_type="latent", - ), - ) - latents = result.images - - latent_payload = { - "latents": latents, - "scaling_factor": torch.tensor(self.vae_scaling_factor), - } - if self.vae_shift_factor is not None: - latent_payload["shift_factor"] = torch.tensor(self.vae_shift_factor) - - response = DenoiserResponse( - latents_b64=tensors_to_b64(latent_payload), - shape=list(latents.shape), - ) - logger.info("Denoised — latents %s", list(latents.shape)) - yield response.model_dump() - - -@dynamo_worker() -async def worker(runtime: DistributedRuntime): - endpoint = runtime.endpoint("disagg_diffusion.denoiser.generate") - - stage = DenoiserStage() - loop = asyncio.get_event_loop() - await loop.run_in_executor(None, stage.load_model) - - logger.info("Serving denoiser endpoint: disagg_diffusion.denoiser.generate") - await endpoint.serve_endpoint(stage.generate) - - -if __name__ == "__main__": - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s [%(name)s] %(levelname)s %(message)s", - ) - uvloop.install() - asyncio.run(worker()) diff --git a/examples/disagg_diffusion/phase1_workers/encoder_worker.py b/examples/disagg_diffusion/phase1_workers/encoder_worker.py deleted file mode 100755 index f52c2fabfd14..000000000000 --- a/examples/disagg_diffusion/phase1_workers/encoder_worker.py +++ /dev/null @@ -1,111 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Disaggregated Diffusion — Encoder Worker - -Loads only the text encoders (CLIP + T5) and serves an endpoint that -converts text prompts into serialized embeddings. - -Usage: - python encoder_worker.py --model black-forest-labs/FLUX.1-schnell -""" - -import asyncio -import logging -import os -import sys - -import torch -import uvloop - -# Allow importing protocol.py from the same directory -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - -from protocol import EncoderRequest, EncoderResponse, tensors_to_b64 # noqa: E402 - -from dynamo.runtime import DistributedRuntime, dynamo_endpoint, dynamo_worker # noqa: E402 - -logger = logging.getLogger(__name__) - -MODEL_PATH = os.environ.get("MODEL_PATH", "black-forest-labs/FLUX.1-schnell") -DEVICE = os.environ.get("DEVICE", "cuda") - - -class EncoderStage: - """Text encoder stage: CLIP + T5 → embeddings.""" - - def __init__(self): - self.pipe = None - - def load_model(self): - from diffusers import FluxPipeline - - logger.info("Loading text encoders from %s …", MODEL_PATH) - self.pipe = FluxPipeline.from_pretrained( - MODEL_PATH, torch_dtype=torch.bfloat16 - ) - self.pipe.to(DEVICE) - - # Free transformer + VAE — we only need text encoders - self.pipe.transformer = None - self.pipe.vae = None - torch.cuda.empty_cache() - - vram = torch.cuda.memory_allocated() / 1e6 - logger.info("Encoder ready — VRAM: %.0f MB (text encoders only)", vram) - - @dynamo_endpoint(EncoderRequest, EncoderResponse) - async def generate(self, request: EncoderRequest): - logger.info("Encoding prompt: %.80s…", request.prompt) - - loop = asyncio.get_event_loop() - prompt_embeds, pooled_prompt_embeds, text_ids = await loop.run_in_executor( - None, - lambda: self.pipe.encode_prompt(prompt=request.prompt, prompt_2=None), - ) - - embeddings = { - "prompt_embeds": prompt_embeds, - "pooled_prompt_embeds": pooled_prompt_embeds, - "text_ids": text_ids, - } - - response = EncoderResponse( - embeddings_b64=tensors_to_b64(embeddings), - shapes={k: list(v.shape) for k, v in embeddings.items()}, - ) - logger.info("Encoded — prompt_embeds %s", list(prompt_embeds.shape)) - yield response.model_dump() - - -@dynamo_worker() -async def worker(runtime: DistributedRuntime): - endpoint = runtime.endpoint("disagg_diffusion.encoder.generate") - - stage = EncoderStage() - loop = asyncio.get_event_loop() - await loop.run_in_executor(None, stage.load_model) - - logger.info("Serving encoder endpoint: disagg_diffusion.encoder.generate") - await endpoint.serve_endpoint(stage.generate) - - -if __name__ == "__main__": - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s [%(name)s] %(levelname)s %(message)s", - ) - uvloop.install() - asyncio.run(worker()) diff --git a/examples/disagg_diffusion/phase1_workers/protocol.py b/examples/disagg_diffusion/phase1_workers/protocol.py deleted file mode 100644 index b4aa6c1a19c6..000000000000 --- a/examples/disagg_diffusion/phase1_workers/protocol.py +++ /dev/null @@ -1,127 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Protocol types for disaggregated diffusion stages. - -These Pydantic models define the request/response contracts between -Encoder, Denoiser, and VAE workers. Tensor payloads are serialized -as base64-encoded torch.save bytes. For production use, replace with -shared-memory or object-store references. -""" - -import base64 -import io -from typing import Any, Dict, List, Optional - -import torch -from pydantic import BaseModel - - -# --------------------------------------------------------------------------- -# Tensor serialization helpers -# --------------------------------------------------------------------------- - -def tensor_to_b64(t: torch.Tensor) -> str: - buf = io.BytesIO() - torch.save(t.cpu(), buf) - return base64.b64encode(buf.getvalue()).decode("ascii") - - -def b64_to_tensor(s: str, device: str = "cuda") -> torch.Tensor: - raw = base64.b64decode(s) - return torch.load(io.BytesIO(raw), weights_only=True).to(device) - - -def tensors_to_b64(tensors: Dict[str, torch.Tensor]) -> str: - buf = io.BytesIO() - cpu_tensors = {k: v.cpu() for k, v in tensors.items()} - torch.save(cpu_tensors, buf) - return base64.b64encode(buf.getvalue()).decode("ascii") - - -def b64_to_tensors(s: str, device: str = "cuda") -> Dict[str, torch.Tensor]: - raw = base64.b64decode(s) - data = torch.load(io.BytesIO(raw), weights_only=True) - return {k: v.to(device) for k, v in data.items()} - - -# --------------------------------------------------------------------------- -# Stage 1: Encoder -# --------------------------------------------------------------------------- - -class EncoderRequest(BaseModel): - prompt: str - model: str = "black-forest-labs/FLUX.1-schnell" - - -class EncoderResponse(BaseModel): - """Serialized text embeddings.""" - embeddings_b64: str # base64(torch.save({prompt_embeds, pooled_prompt_embeds, text_ids})) - shapes: Dict[str, List[int]] - - -# --------------------------------------------------------------------------- -# Stage 2: Denoiser -# --------------------------------------------------------------------------- - -class DenoiserRequest(BaseModel): - embeddings_b64: str - model: str = "black-forest-labs/FLUX.1-schnell" - height: int = 512 - width: int = 512 - num_inference_steps: int = 4 - guidance_scale: float = 0.0 - seed: int = 42 - - -class DenoiserResponse(BaseModel): - """Serialized denoised latents.""" - latents_b64: str # base64(torch.save({latents, scaling_factor, shift_factor})) - shape: List[int] - - -# --------------------------------------------------------------------------- -# Stage 3: VAE Decoder -# --------------------------------------------------------------------------- - -class VAEDecodeRequest(BaseModel): - latents_b64: str - model: str = "black-forest-labs/FLUX.1-schnell" - - -class VAEDecodeResponse(BaseModel): - """Final output image.""" - image_b64: Optional[str] = None # base64 PNG - url: Optional[str] = None - - -# --------------------------------------------------------------------------- -# End-to-end (for orchestrator convenience) -# --------------------------------------------------------------------------- - -class GenerateRequest(BaseModel): - prompt: str - model: str = "black-forest-labs/FLUX.1-schnell" - height: int = 512 - width: int = 512 - num_inference_steps: int = 4 - guidance_scale: float = 0.0 - seed: int = 42 - response_format: str = "b64_json" - - -class GenerateResponse(BaseModel): - image_b64: Optional[str] = None - url: Optional[str] = None - timings: Optional[Dict[str, float]] = None diff --git a/examples/disagg_diffusion/phase1_workers/vae_worker.py b/examples/disagg_diffusion/phase1_workers/vae_worker.py deleted file mode 100755 index c1c8bc5c8c58..000000000000 --- a/examples/disagg_diffusion/phase1_workers/vae_worker.py +++ /dev/null @@ -1,122 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Disaggregated Diffusion — VAE Worker - -Loads only the VAE decoder. Accepts denoised latents and returns -the final image as base64-encoded PNG. - -Usage: - python vae_worker.py --model black-forest-labs/FLUX.1-schnell -""" - -import asyncio -import base64 -import io -import logging -import os -import sys - -import numpy as np -import torch -import uvloop -from PIL import Image - -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - -from protocol import VAEDecodeRequest, VAEDecodeResponse, b64_to_tensors # noqa: E402 - -from dynamo.runtime import DistributedRuntime, dynamo_endpoint, dynamo_worker # noqa: E402 - -logger = logging.getLogger(__name__) - -MODEL_PATH = os.environ.get("MODEL_PATH", "black-forest-labs/FLUX.1-schnell") -DEVICE = os.environ.get("DEVICE", "cuda") - - -class VAEStage: - """VAE decode stage: latents → image pixels.""" - - def __init__(self): - self.vae = None - - def load_model(self): - from diffusers import AutoencoderKL - - logger.info("Loading VAE from %s …", MODEL_PATH) - self.vae = AutoencoderKL.from_pretrained( - MODEL_PATH, subfolder="vae", torch_dtype=torch.bfloat16 - ) - self.vae.to(DEVICE) - - vram = torch.cuda.memory_allocated() / 1e6 - logger.info("VAE ready — VRAM: %.0f MB", vram) - - @dynamo_endpoint(VAEDecodeRequest, VAEDecodeResponse) - async def generate(self, request: VAEDecodeRequest): - logger.info("Decoding latents …") - - data = b64_to_tensors(request.latents_b64, DEVICE) - latents = data["latents"] - scaling_factor = data["scaling_factor"].item() - shift_factor = data.get("shift_factor") - if shift_factor is not None: - shift_factor = shift_factor.item() - - # Undo pipeline's latent scaling - if shift_factor is not None: - latents = latents / scaling_factor + shift_factor - else: - latents = latents / scaling_factor - - loop = asyncio.get_event_loop() - - def _decode(): - with torch.no_grad(): - decoded = self.vae.decode(latents, return_dict=False)[0] - decoded = (decoded / 2 + 0.5).clamp(0, 1) - return decoded.cpu().permute(0, 2, 3, 1).float().numpy() - - pixels = await loop.run_in_executor(None, _decode) - - img = Image.fromarray((pixels[0] * 255).round().astype(np.uint8)) - buf = io.BytesIO() - img.save(buf, format="PNG") - image_b64 = base64.b64encode(buf.getvalue()).decode("ascii") - - response = VAEDecodeResponse(image_b64=image_b64) - logger.info("Decoded — image %dx%d", img.width, img.height) - yield response.model_dump() - - -@dynamo_worker() -async def worker(runtime: DistributedRuntime): - endpoint = runtime.endpoint("disagg_diffusion.vae.generate") - - stage = VAEStage() - loop = asyncio.get_event_loop() - await loop.run_in_executor(None, stage.load_model) - - logger.info("Serving VAE endpoint: disagg_diffusion.vae.generate") - await endpoint.serve_endpoint(stage.generate) - - -if __name__ == "__main__": - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s [%(name)s] %(levelname)s %(message)s", - ) - uvloop.install() - asyncio.run(worker()) diff --git a/examples/disagg_diffusion/phase2_orchestrator/run_disagg.py b/examples/disagg_diffusion/phase2_orchestrator/run_disagg.py deleted file mode 100755 index 615ccff5f401..000000000000 --- a/examples/disagg_diffusion/phase2_orchestrator/run_disagg.py +++ /dev/null @@ -1,164 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Phase 2: Disaggregated Diffusion Orchestrator - -Connects to the three stage workers (Encoder, Denoiser, VAE) via Dynamo -RPC and chains them into an end-to-end image generation pipeline. - -This plays the same role as the Frontend/Global Router in the EPD -architecture, but implemented as a lightweight client for the POC. - -Usage: - # Ensure the three workers are already running (see launch/run_all.sh) - python run_disagg.py \\ - --prompt "A photo of a cat sitting on a windowsill" \\ - --output /tmp/disagg_output.png -""" - -import asyncio -import base64 -import json -import logging -import os -import sys -import time - -import uvloop - -sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "phase1_workers")) - -from protocol import ( # noqa: E402 - DenoiserRequest, - EncoderRequest, - VAEDecodeRequest, -) - -from dynamo.runtime import DistributedRuntime, dynamo_worker # noqa: E402 - -logger = logging.getLogger(__name__) - -PROMPT = os.environ.get("PROMPT", "A photo of a cat sitting on a windowsill") -MODEL = os.environ.get("MODEL_PATH", "black-forest-labs/FLUX.1-schnell") -OUTPUT = os.environ.get("OUTPUT", "/tmp/disagg_output.png") -HEIGHT = int(os.environ.get("HEIGHT", "512")) -WIDTH = int(os.environ.get("WIDTH", "512")) -NUM_STEPS = int(os.environ.get("NUM_STEPS", "4")) -SEED = int(os.environ.get("SEED", "42")) - - -async def call_stage(client, request_json: str) -> dict: - """Call a Dynamo endpoint, collect the streamed response. - - Each stage worker yields exactly one response dict. - """ - result = None - stream = await client.generate(request_json) - async for chunk in stream: - # Dynamo streams return objects with a .data() method or raw dicts - data = chunk.data() if hasattr(chunk, "data") else chunk - if isinstance(data, str): - data = json.loads(data) - result = data - - if result is None: - raise RuntimeError("Empty response from stage") - return result - - -@dynamo_worker() -async def worker(runtime: DistributedRuntime): - """Orchestrator: chain Encoder → Denoiser → VAE.""" - - # Create clients to the three stage endpoints - encoder_client = await runtime.endpoint( - "disagg_diffusion.encoder.generate" - ).client() - denoiser_client = await runtime.endpoint( - "disagg_diffusion.denoiser.generate" - ).client() - vae_client = await runtime.endpoint( - "disagg_diffusion.vae.generate" - ).client() - - logger.info("Connected to all three stage endpoints") - timings = {} - - # ── Stage 1: Encoder ───────────────────────────────────────────── - logger.info("[1/3] Encoding prompt …") - t0 = time.monotonic() - - encoder_req = EncoderRequest(prompt=PROMPT, model=MODEL) - encoder_resp = await call_stage(encoder_client, encoder_req.model_dump_json()) - - timings["encoder_s"] = time.monotonic() - t0 - logger.info(" Done in %.2fs — shapes: %s", timings["encoder_s"], encoder_resp.get("shapes")) - - # ── Stage 2: Denoiser ──────────────────────────────────────────── - logger.info("[2/3] Denoising %dx%d, %d steps …", WIDTH, HEIGHT, NUM_STEPS) - t0 = time.monotonic() - - denoiser_req = DenoiserRequest( - embeddings_b64=encoder_resp["embeddings_b64"], - model=MODEL, - height=HEIGHT, - width=WIDTH, - num_inference_steps=NUM_STEPS, - guidance_scale=0.0, - seed=SEED, - ) - denoiser_resp = await call_stage(denoiser_client, denoiser_req.model_dump_json()) - - timings["denoiser_s"] = time.monotonic() - t0 - logger.info(" Done in %.2fs — latent shape: %s", timings["denoiser_s"], denoiser_resp.get("shape")) - - # ── Stage 3: VAE Decode ────────────────────────────────────────── - logger.info("[3/3] VAE decoding …") - t0 = time.monotonic() - - vae_req = VAEDecodeRequest( - latents_b64=denoiser_resp["latents_b64"], - model=MODEL, - ) - vae_resp = await call_stage(vae_client, vae_req.model_dump_json()) - - timings["vae_s"] = time.monotonic() - t0 - logger.info(" Done in %.2fs", timings["vae_s"]) - - # ── Save output ────────────────────────────────────────────────── - image_bytes = base64.b64decode(vae_resp["image_b64"]) - with open(OUTPUT, "wb") as f: - f.write(image_bytes) - - timings["total_s"] = sum(timings.values()) - - logger.info("") - logger.info("=" * 50) - logger.info("Pipeline complete!") - logger.info(" Encoder: %.2fs", timings["encoder_s"]) - logger.info(" Denoiser: %.2fs", timings["denoiser_s"]) - logger.info(" VAE: %.2fs", timings["vae_s"]) - logger.info(" Total: %.2fs", timings["total_s"]) - logger.info(" Output: %s", OUTPUT) - logger.info("=" * 50) - - -if __name__ == "__main__": - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s [%(name)s] %(levelname)s %(message)s", - ) - uvloop.install() - asyncio.run(worker()) diff --git a/examples/disagg_diffusion/run_all.sh b/examples/disagg_diffusion/run_all.sh new file mode 100755 index 000000000000..328eff6d942b --- /dev/null +++ b/examples/disagg_diffusion/run_all.sh @@ -0,0 +1,171 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Launch all disaggregated diffusion services (etcd + workers + orchestrator) +# and optionally send a test request. +# +# Each stage supports multiple workers: use ';' to separate workers in GPU +# specs. Each worker is an independent top-level process; Dynamo discovers +# them via etcd and the orchestrator dispatches to specific workers via client.direct(). +# +# Usage: +# ./run_all.sh # launch all services (1 worker/stage) +# ./run_all.sh --test # launch + send a test request +# ./run_all.sh --test --quick # launch + quick smoke test (9 frames, 3 steps) +# +# # Multi-worker (8 GPU): +# GPU_ENC="0;4" GPU_DEN="1,2;5,6" GPU_VAE="3;7" ./run_all.sh --test --quick +# +# Environment variables: +# MODEL_PATH HuggingFace model (default: hunyuanvideo-community/HunyuanVideo) +# GPU_ENC GPU(s) for encoder (default: 0; use "0;4" for 2 workers) +# GPU_DEN GPU(s) for denoiser (default: 1,2; use "1,2;5,6" for 2 TP=2 workers) +# GPU_VAE GPU(s) for VAE (default: 3; use "3;7" for 2 workers) +# PORT HTTP port (default: 8080) + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +WORKERS_DIR="$SCRIPT_DIR/workers" +ORCH_DIR="$SCRIPT_DIR/orchestrator" +LOG_DIR="/tmp/disagg_logs" + +GPU_ENC="${GPU_ENC:-0}" +GPU_DEN="${GPU_DEN:-1,2}" +GPU_VAE="${GPU_VAE:-3}" +PORT="${PORT:-8080}" + +DO_TEST=false +QUICK=false +for arg in "$@"; do + case "$arg" in + --test) DO_TEST=true ;; + --quick) QUICK=true ;; + esac +done + +mkdir -p "$LOG_DIR" +PIDS=() + +cleanup() { + echo "" + echo "Shutting down all services..." + for pid in "${PIDS[@]}"; do + kill "$pid" 2>/dev/null || true + done + # Kill child processes (sglang schedulers) + for pid in "${PIDS[@]}"; do + pkill -P "$pid" 2>/dev/null || true + done + wait 2>/dev/null + echo "All services stopped." +} +trap cleanup EXIT INT TERM + +echo "==========================================" +echo " Disaggregated Diffusion — Launch All" +echo "==========================================" +echo " Encoder: GPU $GPU_ENC" +echo " Denoiser: GPU $GPU_DEN" +echo " VAE: GPU $GPU_VAE" +echo " HTTP port: $PORT" +echo " Logs: $LOG_DIR/" +echo "==========================================" + +# 1. etcd +if ! pgrep -x etcd > /dev/null 2>&1; then + echo "[1/5] Starting etcd..." + etcd --data-dir /tmp/etcd_disagg \ + --listen-client-urls http://0.0.0.0:2379 \ + --advertise-client-urls http://127.0.0.1:2379 \ + > "$LOG_DIR/etcd.log" 2>&1 & + PIDS+=($!) + sleep 2 +else + echo "[1/5] etcd already running, skipping." +fi + +# 2. Encoder Worker(s) +STEP=2 +IFS=';' read -ra ENC_GPUS <<< "$GPU_ENC" +for i in "${!ENC_GPUS[@]}"; do + port=$((15600 + i * 10)) + echo "[$STEP] Starting Encoder Worker $i (GPU ${ENC_GPUS[$i]}, port $port)..." + CUDA_VISIBLE_DEVICES="${ENC_GPUS[$i]}" SCHEDULER_PORT=$port \ + python "$WORKERS_DIR/encoder_worker.py" > "$LOG_DIR/encoder_$i.log" 2>&1 & + PIDS+=($!) + STEP=$((STEP + 1)) +done + +# 3. Denoiser Worker(s) +IFS=';' read -ra DEN_GPUS <<< "$GPU_DEN" +for i in "${!DEN_GPUS[@]}"; do + port=$((15700 + i * 10)) + echo "[$STEP] Starting Denoiser Worker $i (GPU ${DEN_GPUS[$i]}, port $port)..." + CUDA_VISIBLE_DEVICES="${DEN_GPUS[$i]}" SCHEDULER_PORT=$port \ + python "$WORKERS_DIR/denoiser_worker.py" > "$LOG_DIR/denoiser_$i.log" 2>&1 & + PIDS+=($!) + STEP=$((STEP + 1)) +done + +# 4. VAE Worker(s) +IFS=';' read -ra VAE_GPUS <<< "$GPU_VAE" +for i in "${!VAE_GPUS[@]}"; do + port=$((15800 + i * 10)) + echo "[$STEP] Starting VAE Worker $i (GPU ${VAE_GPUS[$i]}, port $port)..." + CUDA_VISIBLE_DEVICES="${VAE_GPUS[$i]}" SCHEDULER_PORT=$port \ + python "$WORKERS_DIR/vae_worker.py" > "$LOG_DIR/vae_$i.log" 2>&1 & + PIDS+=($!) + STEP=$((STEP + 1)) +done + +N_WORKERS=$(( ${#ENC_GPUS[@]} + ${#DEN_GPUS[@]} + ${#VAE_GPUS[@]} )) +echo "" +echo "Launched $N_WORKERS worker(s): ${#ENC_GPUS[@]} encoder, ${#DEN_GPUS[@]} denoiser, ${#VAE_GPUS[@]} vae" + +# 5. Orchestrator +echo "[$STEP] Starting Orchestrator (port $PORT)..." +PORT="$PORT" python "$ORCH_DIR/run_disagg.py" \ + > "$LOG_DIR/orchestrator.log" 2>&1 & +PIDS+=($!) + +echo "" +echo "All services launching. Waiting for workers to be ready..." +echo " tail -f $LOG_DIR/encoder_*.log # monitor encoder(s)" +echo " tail -f $LOG_DIR/denoiser_*.log # monitor denoiser(s)" +echo " tail -f $LOG_DIR/vae_*.log # monitor vae(s)" +echo " tail -f $LOG_DIR/orchestrator.log" +echo "" + +# Wait for orchestrator HTTP to be ready +for i in $(seq 1 120); do + if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then + echo "Orchestrator ready at http://localhost:$PORT" + break + fi + if [ "$i" -eq 120 ]; then + echo "ERROR: Orchestrator not ready after 120s. Check logs in $LOG_DIR/" + exit 1 + fi + sleep 1 +done + +# Test request +if [ "$DO_TEST" = true ]; then + echo "" + echo "Sending test request..." + if [ "$QUICK" = true ]; then + curl -s -X POST "http://localhost:$PORT/v1/videos/generations" \ + -H "Content-Type: application/json" \ + -d '{"prompt": "A cat walking on green grass", "num_frames": 9, "num_inference_steps": 3}' | python -m json.tool + else + curl -s -X POST "http://localhost:$PORT/v1/videos/generations" \ + -H "Content-Type: application/json" \ + -d '{"prompt": "A golden retriever running on a sunny beach with waves crashing in the background"}' | python -m json.tool + fi +fi + +echo "" +echo "Services running. Press Ctrl+C to stop all." +wait diff --git a/examples/disagg_diffusion/stress_test.sh b/examples/disagg_diffusion/stress_test.sh new file mode 100755 index 000000000000..6b3e01ad912b --- /dev/null +++ b/examples/disagg_diffusion/stress_test.sh @@ -0,0 +1,251 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Concurrent stress test for disaggregated diffusion pipeline. +# +# Sends N_REQUESTS concurrent requests, then validates: +# 1. All requests return HTTP 200 with valid JSON +# 2. All produced .mp4 files are non-zero size +# 3. /pipeline/status shows completed=N, failed=0 +# 4. All workers per stage were utilized +# +# Prerequisites: services must already be running (via run_all.sh). +# +# Usage: +# ./stress_test.sh # 20 requests, default port +# N_REQUESTS=50 ./stress_test.sh # 50 requests +# PORT=8081 ./stress_test.sh # custom port + +set -euo pipefail + +PORT="${PORT:-8080}" +N_REQUESTS="${N_REQUESTS:-20}" +BASE_URL="http://localhost:$PORT" +RESULT_DIR="/tmp/stress_test_$$" + +mkdir -p "$RESULT_DIR" + +echo "==========================================" +echo " Disagg Diffusion — Concurrent Stress Test" +echo "==========================================" +echo " Target: $BASE_URL" +echo " Requests: $N_REQUESTS" +echo " Results: $RESULT_DIR/" +echo "==========================================" + +# Verify the server is up +if ! curl -sf "$BASE_URL/health" > /dev/null 2>&1; then + echo "ERROR: Server not reachable at $BASE_URL/health" + echo " Start services first with run_all.sh" + exit 1 +fi + +# Record pre-test status +PRE_STATUS=$(curl -sf "$BASE_URL/pipeline/status") +PRE_COMPLETED=$(echo "$PRE_STATUS" | python3 -c "import sys,json; print(json.load(sys.stdin)['completed'])") +echo "Pre-test completed count: $PRE_COMPLETED" +echo "" + +PROMPTS=( + "A cat walking on green grass" + "A dog running on a sandy beach" + "A bird flying over a mountain lake" + "A fish swimming in a coral reef" + "A horse galloping through a meadow" + "A butterfly landing on a red flower" + "A wolf howling at the full moon" + "A dolphin jumping out of the ocean" + "A panda eating bamboo in the forest" + "A fox running through autumn leaves" + "A penguin sliding on ice in Antarctica" + "An eagle soaring above snowy peaks" + "A tiger walking through a jungle" + "A rabbit hopping through a garden" + "A deer drinking from a stream" + "An owl perched on a branch at night" + "A lion resting under an acacia tree" + "A whale breaching in the deep ocean" + "A parrot perched on a tropical branch" + "A turtle walking slowly on the sand" +) + +# Launch all requests concurrently +echo "Sending $N_REQUESTS concurrent requests..." +T0=$(date +%s) +PIDS=() + +for i in $(seq 0 $((N_REQUESTS - 1))); do + idx=$((i % ${#PROMPTS[@]})) + prompt="${PROMPTS[$idx]}" + ( + curl -sf -X POST "$BASE_URL/v1/videos/generations" \ + -H "Content-Type: application/json" \ + -d "{\"prompt\": \"$prompt\", \"num_frames\": 9, \"num_inference_steps\": 3}" \ + -o "$RESULT_DIR/resp_$i.json" \ + -w "%{http_code}" \ + > "$RESULT_DIR/status_$i.txt" 2>/dev/null + ) & + PIDS+=($!) +done + +echo "All $N_REQUESTS requests launched (PIDs: ${#PIDS[@]}). Waiting..." + +# Wait for all requests to finish +FAILURES=0 +for i in "${!PIDS[@]}"; do + if ! wait "${PIDS[$i]}"; then + FAILURES=$((FAILURES + 1)) + echo " Request $i: curl failed (PID ${PIDS[$i]})" + fi +done + +T1=$(date +%s) +ELAPSED=$((T1 - T0)) +echo "" +echo "All requests completed in ${ELAPSED}s" +echo "" + +# ── Validation ── + +PASS=0 +FAIL=0 +VIDEOS=() + +echo "── Validation ──" + +# Check 1: HTTP status codes +echo "" +echo "1) HTTP status codes:" +for i in $(seq 0 $((N_REQUESTS - 1))); do + status_file="$RESULT_DIR/status_$i.txt" + if [ -f "$status_file" ]; then + code=$(cat "$status_file") + if [ "$code" = "200" ]; then + PASS=$((PASS + 1)) + else + FAIL=$((FAIL + 1)) + echo " FAIL: request $i returned HTTP $code" + fi + else + FAIL=$((FAIL + 1)) + echo " FAIL: request $i — no status file (curl crashed)" + fi +done +echo " $PASS/$N_REQUESTS returned HTTP 200" + +# Check 2: Valid JSON with video URL +echo "" +echo "2) Response JSON + video files:" +VIDEO_PASS=0 +VIDEO_FAIL=0 +for i in $(seq 0 $((N_REQUESTS - 1))); do + resp="$RESULT_DIR/resp_$i.json" + if [ ! -f "$resp" ] || [ ! -s "$resp" ]; then + VIDEO_FAIL=$((VIDEO_FAIL + 1)) + continue + fi + url=$(python3 -c " +import json, sys +try: + d = json.load(open('$resp')) + print(d['data'][0].get('url', '')) +except Exception: + print('') +" 2>/dev/null) + if [ -n "$url" ]; then + # Extract filename from /videos/ + fname=$(basename "$url") + VIDEOS+=("$fname") + VIDEO_PASS=$((VIDEO_PASS + 1)) + else + VIDEO_FAIL=$((VIDEO_FAIL + 1)) + echo " FAIL: request $i — no video URL in response" + fi +done +echo " $VIDEO_PASS/$N_REQUESTS have valid video URLs" + +# Check 3: Video files exist and are non-zero +echo "" +echo "3) Video file integrity:" +VIDEO_DIR="/tmp/disagg_videos" +FILE_PASS=0 +FILE_FAIL=0 +for fname in "${VIDEOS[@]}"; do + fpath="$VIDEO_DIR/$fname" + if [ -f "$fpath" ] && [ -s "$fpath" ]; then + FILE_PASS=$((FILE_PASS + 1)) + else + FILE_FAIL=$((FILE_FAIL + 1)) + echo " FAIL: $fname missing or empty" + fi +done +echo " $FILE_PASS/${#VIDEOS[@]} video files valid (non-zero size)" + +# Check 4: Pipeline status — completed count and worker utilization +echo "" +echo "4) Pipeline status:" +POST_STATUS=$(curl -sf "$BASE_URL/pipeline/status") +POST_COMPLETED=$(echo "$POST_STATUS" | python3 -c "import sys,json; print(json.load(sys.stdin)['completed'])") +POST_FAILED=$(echo "$POST_STATUS" | python3 -c "import sys,json; print(json.load(sys.stdin)['failed'])") +NEW_COMPLETED=$((POST_COMPLETED - PRE_COMPLETED)) +echo " Completed: $NEW_COMPLETED (expected: $N_REQUESTS)" +echo " Failed: $POST_FAILED" + +# Check worker utilization +echo "" +echo "5) Worker utilization:" +echo "$POST_STATUS" | python3 -c " +import json, sys +status = json.load(sys.stdin) +all_used = True +for stage_name, stage in status.get('stages', {}).items(): + workers = stage.get('workers', []) + used = [w for w in workers if w['completed'] > 0] + total = len(workers) + print(f' {stage_name}: {len(used)}/{total} workers used', end='') + for w in workers: + print(f\" [id={w['id']} completed={w['completed']} avg={w['avg_latency_s']:.2f}s]\", end='') + print() + if len(used) < total: + all_used = False +if all_used: + print(' All workers utilized.') +else: + print(' WARNING: Not all workers were utilized.') +" + +# ── Summary ── +echo "" +echo "==========================================" +TOTAL_CHECKS=0 +TOTAL_PASS=0 + +# HTTP check +TOTAL_CHECKS=$((TOTAL_CHECKS + 1)) +if [ "$PASS" -eq "$N_REQUESTS" ]; then TOTAL_PASS=$((TOTAL_PASS + 1)); fi + +# Video URL check +TOTAL_CHECKS=$((TOTAL_CHECKS + 1)) +if [ "$VIDEO_PASS" -eq "$N_REQUESTS" ]; then TOTAL_PASS=$((TOTAL_PASS + 1)); fi + +# File integrity check +TOTAL_CHECKS=$((TOTAL_CHECKS + 1)) +if [ "$FILE_PASS" -eq "${#VIDEOS[@]}" ] && [ "${#VIDEOS[@]}" -gt 0 ]; then TOTAL_PASS=$((TOTAL_PASS + 1)); fi + +# Pipeline completed count check +TOTAL_CHECKS=$((TOTAL_CHECKS + 1)) +if [ "$NEW_COMPLETED" -eq "$N_REQUESTS" ] && [ "$POST_FAILED" -eq "0" ]; then TOTAL_PASS=$((TOTAL_PASS + 1)); fi + +if [ "$TOTAL_PASS" -eq "$TOTAL_CHECKS" ]; then + echo " RESULT: ALL PASSED ($TOTAL_PASS/$TOTAL_CHECKS checks)" + echo " $N_REQUESTS requests, ${ELAPSED}s total, $(echo "scale=1; $ELAPSED / $N_REQUESTS" | bc)s avg" +else + echo " RESULT: $TOTAL_PASS/$TOTAL_CHECKS checks passed" +fi +echo "==========================================" + +# Cleanup temp dir +rm -rf "$RESULT_DIR" + +[ "$TOTAL_PASS" -eq "$TOTAL_CHECKS" ] && exit 0 || exit 1 diff --git a/examples/disagg_diffusion/tests/test_protocol_models.py b/examples/disagg_diffusion/tests/test_protocol_models.py new file mode 100644 index 000000000000..993f45fab85a --- /dev/null +++ b/examples/disagg_diffusion/tests/test_protocol_models.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 + +from pathlib import Path +import sys + + +ROOT = Path(__file__).resolve().parents[1] +WORKERS_DIR = ROOT / "workers" +sys.path.insert(0, str(WORKERS_DIR)) + +from protocol import DenoiserResponse, EncoderResponse, GenerateRequest, VAEDecodeRequest # noqa: E402 + + +def test_protocol_mutable_defaults_are_isolated(): + a = EncoderResponse() + b = EncoderResponse() + a.transfer_meta["x"] = 1 + a.shapes["latents"] = [1, 2, 3] + + assert b.transfer_meta == {} + assert b.shapes == {} + + +def test_protocol_other_dict_defaults_are_isolated(): + a = DenoiserResponse() + b = DenoiserResponse() + a.transfer_meta["k"] = "v" + a.shape.append(42) + + assert b.transfer_meta == {} + assert b.shape == [] + + x = VAEDecodeRequest(transfer_meta={}) + y = VAEDecodeRequest(transfer_meta={}) + x.tensor_data["t"] = "blob" + assert y.tensor_data == {} + + +def test_generate_request_seed_is_optional(): + req = GenerateRequest(prompt="hello") + assert req.seed is None diff --git a/examples/disagg_diffusion/tests/test_request_validation.py b/examples/disagg_diffusion/tests/test_request_validation.py new file mode 100644 index 000000000000..084232fbf0ae --- /dev/null +++ b/examples/disagg_diffusion/tests/test_request_validation.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python3 + +from pathlib import Path +import sys + +import pytest +from pydantic import ValidationError + + +ROOT = Path(__file__).resolve().parents[1] +WORKERS_DIR = ROOT / "workers" +ORCH_DIR = ROOT / "orchestrator" +sys.path.insert(0, str(WORKERS_DIR)) +sys.path.insert(0, str(ORCH_DIR)) + +from request_validation import validate_generate_http_body # noqa: E402 + + +def test_validate_generate_http_body_accepts_minimal_payload(): + out = validate_generate_http_body({"prompt": "a cat"}) + assert out["prompt"] == "a cat" + assert out["response_format"] == "url" + assert out["seed"] is None + + +def test_validate_generate_http_body_rejects_invalid_response_format(): + with pytest.raises(ValueError, match="response_format"): + validate_generate_http_body( + { + "prompt": "a cat", + "response_format": "invalid", + } + ) + + +def test_validate_generate_http_body_rejects_invalid_schema(): + with pytest.raises(ValidationError): + validate_generate_http_body({"height": 544}) diff --git a/examples/disagg_diffusion/phase0_validate/validate_split.py b/examples/disagg_diffusion/validate/validate_split.py similarity index 100% rename from examples/disagg_diffusion/phase0_validate/validate_split.py rename to examples/disagg_diffusion/validate/validate_split.py diff --git a/examples/disagg_diffusion/phase1_workers/__init__.py b/examples/disagg_diffusion/workers/__init__.py similarity index 100% rename from examples/disagg_diffusion/phase1_workers/__init__.py rename to examples/disagg_diffusion/workers/__init__.py diff --git a/examples/disagg_diffusion/workers/denoiser_worker.py b/examples/disagg_diffusion/workers/denoiser_worker.py new file mode 100755 index 000000000000..35774e461006 --- /dev/null +++ b/examples/disagg_diffusion/workers/denoiser_worker.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Disaggregated Diffusion — Denoiser Worker (Dynamo RPC) + +Wraps SGLang Scheduler subprocess(es) running NixlReceiveStage + Denoising + +NixlSendStage. Supports TP via launch_partial_server(tp_size=N). + +Process architecture:: + + Dynamo Worker Process (this file) + |-- @dynamo_worker + | |-- serve_endpoint("generate") <-- Dynamo RPC from orchestrator + | | +-- handle_generate() + | | +-- StageClient.forward() <-- ZMQ to local Scheduler + | +-- serve_endpoint("health") + | + +-- SGLang Scheduler subprocess(es) (spawned by launch_partial_server) + +-- PartialGPUWorker (TP=N) + |-- NixlReceiveStage <-- RDMA-pull embeddings from encoder + |-- LatentPreparationStage + |-- TimestepPreparationStage + |-- DenoisingStage + +-- NixlSendStage <-- register latents as NIXL-readable +""" + +import asyncio +import json +import logging +import multiprocessing as mp +import os +import sys + +import uvloop + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from dynamo.runtime import DistributedRuntime, dynamo_worker # noqa: E402 + +logger = logging.getLogger(__name__) + +MODEL_PATH = os.environ.get("MODEL_PATH", "hunyuanvideo-community/HunyuanVideo") +SCHEDULER_PORT = int(os.environ.get("SCHEDULER_PORT", "15700")) + + +@dynamo_worker(enable_nats=False) +async def worker(runtime: DistributedRuntime): + from sglang_utils import launch_stage_server, build_req + from partial_gpu_worker import build_denoiser_stages + + # Auto-detect GPU count from CUDA_VISIBLE_DEVICES + num_gpus = len(os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")) + tp_size = int(os.environ.get("TP_SIZE", str(num_gpus))) + + logger.info("Launching denoiser Scheduler: num_gpus=%d, tp=%d, port=%d", num_gpus, tp_size, SCHEDULER_PORT) + processes, client, server_args = launch_stage_server( + MODEL_PATH, ["transformer", "scheduler"], build_denoiser_stages, + SCHEDULER_PORT, tp_size=tp_size, num_gpus=num_gpus, + client_name="denoiser", + ) + + # ── Dynamo RPC handlers ────────────────────────────────────────── + + async def handle_generate(request, context): + try: + if isinstance(request, str): + request = json.loads(request) + + req = build_req( + prompt="(embeddings via NIXL)", + negative_prompt="", + height=request.get("height", 544), + width=request.get("width", 960), + num_frames=request.get("num_frames", 61), + num_inference_steps=request.get("num_inference_steps", 50), + guidance_scale=request.get("guidance_scale", 1.0), + seed=request.get("seed", 42), + ) + req.do_classifier_free_guidance = (req.guidance_scale > 1.0) + + # Pass NIXL metadata or raw tensor data for NixlReceiveStage + transfer_meta = request.get("transfer_meta", {}) + tensor_data = request.get("tensor_data", {}) + if transfer_meta: + req._nixl_transfer_meta = transfer_meta + elif tensor_data: + # ZMQ fallback: deserialize tensors and inject onto Req + import torch, base64, io + from sglang_utils import inject_tensors_to_req + tensors = {} + for k, v in tensor_data.items(): + if isinstance(v, list): + tensors[k] = [torch.load(io.BytesIO(base64.b64decode(b)), weights_only=True) for b in v] + else: + tensors[k] = torch.load(io.BytesIO(base64.b64decode(v)), weights_only=True) + inject_tensors_to_req(req, tensors) + logger.info("Injected %d tensor fields via ZMQ fallback", len(tensors)) + + output = await client.forward([req]) + if output.error: + yield {"error": str(output.error), "transfer_meta": {}, "shape": []} + return + + result = output.output + transfer_meta_out = result.get("_nixl_transfer_meta", {}) + if transfer_meta_out: + logger.info("Denoised — NIXL latent metadata ready") + yield {"transfer_meta": transfer_meta_out, "shape": []} + else: + # ZMQ fallback for latents + import torch, base64, io + td = {} + for k, v in result.items(): + if isinstance(v, torch.Tensor): + buf = io.BytesIO() + torch.save(v.cpu(), buf) + td[k] = base64.b64encode(buf.getvalue()).decode() + logger.info("Denoised — ZMQ fallback (%d tensor fields)", len(td)) + yield {"transfer_meta": {}, "tensor_data": td, "shape": []} + + except Exception as e: + logger.error("Denoiser generate failed: %s", e, exc_info=True) + yield {"error": str(e), "transfer_meta": {}, "shape": []} + + async def handle_health(request, context): + yield { + "status": "ok", "stage": "denoiser", + "model": MODEL_PATH, "tp_size": tp_size, + } + + # ── Serve Dynamo endpoints ─────────────────────────────────────── + + gen_ep = runtime.endpoint("disagg_diffusion.denoiser.generate") + health_ep = runtime.endpoint("disagg_diffusion.denoiser.health") + + logger.info("Serving: disagg_diffusion.denoiser.generate + health") + try: + await asyncio.gather( + gen_ep.serve_endpoint(handle_generate), + health_ep.serve_endpoint(handle_health), + ) + finally: + client.close() + for p in processes: + p.terminate() + for p in processes: + p.join(timeout=10) + + +if __name__ == "__main__": + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(name)s] %(levelname)s %(message)s", + ) + mp.set_start_method("spawn", force=True) + uvloop.install() + asyncio.run(worker()) diff --git a/examples/disagg_diffusion/workers/encoder_worker.py b/examples/disagg_diffusion/workers/encoder_worker.py new file mode 100755 index 000000000000..b4743b0fe7b3 --- /dev/null +++ b/examples/disagg_diffusion/workers/encoder_worker.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Disaggregated Diffusion — Encoder Worker (Dynamo RPC) + +Wraps an SGLang Scheduler subprocess running TextEncodingStage + NixlSendStage. +Exposes a Dynamo RPC endpoint. The Scheduler handles model loading, GPU +execution, and NIXL tensor registration. This worker bridges Dynamo RPC <-> ZMQ. + +Process architecture:: + + Dynamo Worker Process (this file) + |-- @dynamo_worker + | |-- serve_endpoint("generate") <-- Dynamo RPC from orchestrator + | | +-- handle_generate() + | | +-- StageClient.forward() <-- ZMQ to local Scheduler + | +-- serve_endpoint("health") + | + +-- SGLang Scheduler subprocess (spawned by launch_partial_server) + +-- PartialGPUWorker + |-- TextEncodingStage + +-- NixlSendStage <-- registers embeddings as NIXL-readable +""" + +import asyncio +import json +import logging +import multiprocessing as mp +import os +import sys + +import uvloop + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from dynamo.runtime import DistributedRuntime, dynamo_worker # noqa: E402 + +logger = logging.getLogger(__name__) + +MODEL_PATH = os.environ.get("MODEL_PATH", "hunyuanvideo-community/HunyuanVideo") +SCHEDULER_PORT = int(os.environ.get("SCHEDULER_PORT", "15600")) + + +@dynamo_worker(enable_nats=False) +async def worker(runtime: DistributedRuntime): + from sglang_utils import launch_stage_server, detect_encoder_modules, build_req + from partial_gpu_worker import build_encoder_stages + + enc_modules = detect_encoder_modules(MODEL_PATH) + logger.info("Launching encoder Scheduler: modules=%s, port=%d", enc_modules, SCHEDULER_PORT) + processes, client, server_args = launch_stage_server( + MODEL_PATH, enc_modules, build_encoder_stages, + SCHEDULER_PORT, client_name="encoder", + ) + + # ── Dynamo RPC handlers ────────────────────────────────────────── + + async def handle_generate(request, context): + try: + if isinstance(request, str): + request = json.loads(request) + + req = build_req( + prompt=request.get("prompt", ""), + negative_prompt=request.get("negative_prompt", ""), + guidance_scale=request.get("guidance_scale", 1.0), + ) + + output = await client.forward([req]) + if output.error: + yield {"error": str(output.error), "transfer_meta": {}, "shapes": {}} + return + + result = output.output + transfer_meta = result.get("_nixl_transfer_meta", {}) + if transfer_meta: + logger.info("Encoded prompt — NIXL metadata ready") + yield {"transfer_meta": transfer_meta, "shapes": {}} + else: + # ZMQ fallback: forward raw tensors (when NIXL is disabled) + import torch, base64, io + tensor_data = {} + for k, v in result.items(): + if isinstance(v, torch.Tensor): + buf = io.BytesIO() + torch.save(v.cpu(), buf) + tensor_data[k] = base64.b64encode(buf.getvalue()).decode() + elif isinstance(v, list) and v and isinstance(v[0], torch.Tensor): + tensor_data[k] = [] + for t in v: + buf = io.BytesIO() + torch.save(t.cpu(), buf) + tensor_data[k].append(base64.b64encode(buf.getvalue()).decode()) + logger.info("Encoded prompt — ZMQ fallback (%d tensor fields)", len(tensor_data)) + yield {"transfer_meta": {}, "tensor_data": tensor_data, "shapes": {}} + + except Exception as e: + logger.error("Encoder generate failed: %s", e, exc_info=True) + yield {"error": str(e), "transfer_meta": {}, "shapes": {}} + + async def handle_health(request, context): + yield {"status": "ok", "stage": "encoder", "model": MODEL_PATH} + + # ── Serve Dynamo endpoints ─────────────────────────────────────── + + gen_ep = runtime.endpoint("disagg_diffusion.encoder.generate") + health_ep = runtime.endpoint("disagg_diffusion.encoder.health") + + logger.info("Serving: disagg_diffusion.encoder.generate + health") + try: + await asyncio.gather( + gen_ep.serve_endpoint(handle_generate), + health_ep.serve_endpoint(handle_health), + ) + finally: + client.close() + for p in processes: + p.terminate() + for p in processes: + p.join(timeout=10) + + +if __name__ == "__main__": + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(name)s] %(levelname)s %(message)s", + ) + mp.set_start_method("spawn", force=True) + uvloop.install() + asyncio.run(worker()) diff --git a/examples/disagg_diffusion/workers/nixl_transfer.py b/examples/disagg_diffusion/workers/nixl_transfer.py new file mode 100644 index 000000000000..ce284363273b --- /dev/null +++ b/examples/disagg_diffusion/workers/nixl_transfer.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""NIXL RDMA tensor transfer for disaggregated diffusion stages. + +Provides GPU-direct tensor transfer between stage workers. Only small +metadata (shapes, dtypes, NIXL descriptor ~1.5 KB) travels over the ZMQ +control plane; actual tensor data (embeddings, latents) transfers +GPU->GPU via NIXL RDMA. + +Follows the PersistentConnector pattern from embedding_transfer.py: +- Each Sender/Receiver owns its PersistentConnector (one Connection/agent). +- Remote._release is nooped to keep agent pairs alive. +- The sender returns a ``readable_op`` handle that the caller must hold + until the receiver completes the RDMA pull. + +Usage inside PipelineStage.forward() (synchronous context):: + + sender = NixlTensorSender() # creates agent eagerly + meta, readable_op = sender.send({"latents": tensor}) + # ... pass meta via ZMQ, hold readable_op until COMPLETE ... + + receiver = NixlTensorReceiver() # creates agent eagerly + tensors = receiver.recv(meta, device="cuda") # RDMA pull +""" + +from __future__ import annotations + +import asyncio +import logging +import os +from typing import Dict, Tuple + +import torch + +logger = logging.getLogger(__name__) + +try: + import dynamo.nixl_connect as nixl_connect + + NIXL_AVAILABLE = not os.environ.get("DISABLE_NIXL", "").lower() in ( + "1", + "true", + "yes", + ) + if not NIXL_AVAILABLE: + logger.info("NIXL disabled via DISABLE_NIXL env var") +except ImportError: + NIXL_AVAILABLE = False + logger.info("NIXL not available — falling back to ZMQ tensor transfer") + + +# --------------------------------------------------------------------------- +# PersistentConnector + Remote._release noop +# --------------------------------------------------------------------------- +# Exact pattern from components/src/dynamo/common/multimodal/embedding_transfer.py + +if NIXL_AVAILABLE: + + class PersistentConnector(nixl_connect.Connector): + """Connector that reuses a single Connection for all operations.""" + + def __init__(self): + super().__init__() + self._connection = None + + async def _create_connection(self) -> nixl_connect.Connection: + if self._connection is None: + self._connection = nixl_connect.Connection(self, 1) + await self._connection.initialize() + return self._connection + + # NOTE: We do NOT noop Remote._release here. Our recv() is synchronous + # (awaits wait_for_completion before returning), so the transfer is + # always complete before Remote is GC'd. Keeping the remote agent + # registered across requests causes NIXL_ERR_NOT_ALLOWED on the + # second add_remote_agent() call. + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_event_loop = None + +def _run_coro(coro): + """Run a coroutine from synchronous context (sglang scheduler thread).""" + global _event_loop + if _event_loop is None or _event_loop.is_closed(): + _event_loop = asyncio.new_event_loop() + return _event_loop.run_until_complete(coro) + + +# --------------------------------------------------------------------------- +# NixlTensorSender +# --------------------------------------------------------------------------- + + +class NixlTensorSender: + """Register GPU tensors as NIXL-readable. Returns (metadata, readable_op). + + Each instance owns a PersistentConnector whose Connection (nixl_agent) + is created eagerly in ``__init__`` so UCX is fully initialized before + any transfer occurs. + + The caller **must** hold ``readable_op`` until the receiver completes + the RDMA pull. + """ + + def __init__(self): + self.connector = PersistentConnector() + # Eagerly create the Connection / nixl_agent so UCX is ready + _run_coro(self.connector._create_connection()) + + def send(self, tensors: Dict[str, torch.Tensor]) -> Tuple[dict, object]: + """Register tensors and return (metadata_dict, readable_op).""" + return _run_coro(self._async_send(tensors)) + + async def _async_send( + self, tensors: Dict[str, torch.Tensor] + ) -> Tuple[dict, object]: + import time as _time + t0 = _time.monotonic() + + # Flatten all tensors into a single contiguous buffer + vals = list(tensors.values()) + if len(vals) == 1: + flat = vals[0].contiguous().view(-1) + else: + flat = torch.cat([t.contiguous().view(-1) for t in vals]) + t1 = _time.monotonic() + + descriptor = nixl_connect.Descriptor(flat) + t2 = _time.monotonic() + + readable = await self.connector.create_readable(descriptor) + t3 = _time.monotonic() + + raw_meta = readable.metadata() + t4 = _time.monotonic() + + logger.info( + "NIXL send breakdown: flatten=%.1fms descriptor=%.1fms " + "create_readable=%.1fms metadata=%.1fms total=%.1fms " + "(buf=%s %.2fMB)", + (t1-t0)*1000, (t2-t1)*1000, (t3-t2)*1000, (t4-t3)*1000, + (t4-t0)*1000, list(flat.shape), flat.nbytes/1e6, + ) + + meta = { + "tensor_keys": list(tensors.keys()), + "shapes": {k: list(t.shape) for k, t in tensors.items()}, + "dtypes": { + k: str(t.dtype).removeprefix("torch.") for k, t in tensors.items() + }, + "nixl_metadata": raw_meta.model_dump() + if hasattr(raw_meta, "model_dump") + else raw_meta, + } + + # Return both — caller holds readable to prevent GC + return meta, readable + + +# --------------------------------------------------------------------------- +# NixlTensorReceiver +# --------------------------------------------------------------------------- + + +class NixlTensorReceiver: + """Pull tensors from a remote sender via NIXL RDMA. + + Each instance owns a PersistentConnector whose Connection (nixl_agent) + is created eagerly in ``__init__``. + """ + + def __init__(self): + self.connector = PersistentConnector() + # Eagerly create the Connection / nixl_agent so UCX is ready + _run_coro(self.connector._create_connection()) + + def recv(self, meta: dict, device: str = "cuda") -> Dict[str, torch.Tensor]: + """Pull tensors described by metadata. Returns {name: tensor}.""" + return _run_coro(self._async_recv(meta, device)) + + async def _async_recv(self, meta: dict, device: str) -> Dict[str, torch.Tensor]: + # Calculate total size and per-tensor specs + specs = [] + total_bytes = 0 + for key in meta["tensor_keys"]: + shape = meta["shapes"][key] + dtype = getattr(torch, meta["dtypes"][key]) + numel = 1 + for s in shape: + numel *= s + size = numel * dtype.itemsize + specs.append((key, shape, dtype, size)) + total_bytes += size + + # Allocate receive buffer directly on target device (GPU-direct) + flat = torch.empty(total_bytes, dtype=torch.uint8, device=device) + descriptor = nixl_connect.Descriptor(flat) + + rdma_meta = nixl_connect.RdmaMetadata.model_validate(meta["nixl_metadata"]) + read_op = await self.connector.begin_read(rdma_meta, descriptor) + await read_op.wait_for_completion() + + # Slice the flat buffer into individual tensors + result = {} + offset = 0 + for key, shape, dtype, size in specs: + result[key] = flat[offset : offset + size].view(dtype=dtype).reshape(shape) + offset += size + + return result diff --git a/examples/disagg_diffusion/workers/partial_gpu_worker.py b/examples/disagg_diffusion/workers/partial_gpu_worker.py new file mode 100644 index 000000000000..0966b6499fc0 --- /dev/null +++ b/examples/disagg_diffusion/workers/partial_gpu_worker.py @@ -0,0 +1,721 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""PartialGPUWorker, NixlSendStage, NixlReceiveStage, and subprocess launcher. + +This module is the single source of truth for disaggregated diffusion's +integration with sglang. Everything that runs **inside the Scheduler +subprocess** lives here — no Dynamo imports, no heavy dependencies beyond +sglang + torch. + +Architecture:: + + Main process Subprocess (spawned) + ─────────── ──────────────────── + launch_partial_server() ──► _run_partial_scheduler_process() + ├─ monkey-patch GPUWorker → PartialGPUWorker + └─ run_scheduler_process() + └─ Scheduler() + └─ PartialGPUWorker() + └─ pipeline.forward() + SchedulerClient ◄────ZMQ────► Scheduler.event_loop() + +``NixlSendStage`` is appended as the last pipeline stage for +encoder / denoiser. It registers ``Req`` tensors as NIXL-readable and +returns an ``OutputBatch`` with NIXL metadata. + +``NixlReceiveStage`` is prepended at the start of denoiser / VAE +pipelines. It RDMA-pulls tensors from the previous stage. + +``DecodingStage`` (VAE) already returns ``OutputBatch``, so no extra +send stage is needed for the VAE worker. +""" + +from __future__ import annotations + +import logging +import multiprocessing as mp +import os +import sys +from typing import Callable, Dict, List, Optional + +import torch +from setproctitle import setproctitle + +from sglang.multimodal_gen.runtime.distributed import ( + get_tp_rank, + get_tp_world_size, + maybe_init_distributed_environment_and_model_parallel, + model_parallel_is_initialized, +) +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_classifier_free_guidance_rank, + get_classifier_free_guidance_world_size, + get_ring_parallel_rank, + get_ring_parallel_world_size, + get_ulysses_parallel_rank, + get_ulysses_parallel_world_size, +) +from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req, OutputBatch, RequestMetrics +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage +from sglang.multimodal_gen.runtime.server_args import ServerArgs + +# layerwise_offload may not exist in all sglang versions — guard import +try: + from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin +except ImportError: + OffloadableDiTMixin = None # type: ignore[misc,assignment] + +logger = logging.getLogger(__name__) + + +# ═══════════════════════════════════════════════════════════════════════ +# NIXL Pipeline Stages (run inside Scheduler subprocess) +# ═══════════════════════════════════════════════════════════════════════ + + +class NixlReceiveStage(PipelineStage): + """Pull tensor fields onto the current GPU via NIXL RDMA. + + Prepend at the start of denoiser / VAE pipelines. The ``Req`` + carries NIXL metadata (set by the orchestrator); this stage uses + it to pull the actual tensor data directly from the sender's GPU + without a CPU round-trip. + """ + + def __init__(self, tensor_fields: List[str]): + super().__init__() + self._tensor_fields = tensor_fields + # Eagerly create receiver so its NIXL agent/UCX endpoint is ready + # before any transfer occurs in this process. + # Only rank 0 uses NIXL — TP slaves skip to avoid UCX conflicts. + from nixl_transfer import NIXL_AVAILABLE + if NIXL_AVAILABLE and get_tp_rank() == 0: + from nixl_transfer import NixlTensorReceiver + self._receiver = NixlTensorReceiver() + else: + self._receiver = None + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + tp_world = get_tp_world_size() + tp_rank = get_tp_rank() + nixl_meta = getattr(batch, "_nixl_transfer_meta", None) + + if tp_world <= 1: + # No TP — simple path + if nixl_meta is not None and self._receiver is not None: + return self._nixl_pull(batch, nixl_meta) + return self._device_move(batch) + + # TP mode: rank 0 decides NIXL vs ZMQ, broadcast decision to all ranks + # (torch.distributed.broadcast requires ALL ranks to participate) + import torch.distributed as dist + flag = torch.zeros(1, dtype=torch.int32, device="cuda") + if tp_rank == 0 and nixl_meta is not None and self._receiver is not None: + flag[0] = 1 + dist.broadcast(flag, src=0) + + if flag.item(): + # NIXL path: rank 0 pulls, then broadcasts tensors to all TP ranks + if tp_rank == 0: + batch = self._nixl_pull(batch, nixl_meta) + self._tp_broadcast_fields(batch, src_rank=0) + return batch + + return self._device_move(batch) + + _NIXL_PULL_MAX_RETRIES = int(os.environ.get("NIXL_PULL_MAX_RETRIES", "5")) + _NIXL_PULL_BACKOFF_S = float(os.environ.get("NIXL_PULL_BACKOFF_S", "0.5")) + + def _nixl_pull(self, batch: Req, meta: dict) -> Req: + # Retry on REMOTE_DISCONNECT — NIXL 1:N fan-out can be flaky on + # initial peer connection when multiple receivers target one sender. + last_err = None + for attempt in range(self._NIXL_PULL_MAX_RETRIES + 1): + try: + tensors = self._receiver.recv(meta, device="cuda") + break + except Exception as e: + if "REMOTE_DISCONNECT" in str(e) and attempt < self._NIXL_PULL_MAX_RETRIES: + import time + wait = self._NIXL_PULL_BACKOFF_S * (attempt + 1) + logger.warning( + "NIXL pull attempt %d/%d failed: %s, retrying in %.1fs", + attempt + 1, self._NIXL_PULL_MAX_RETRIES, e, wait, + ) + # Reuse the same receiver — PersistentConnector handles + # agent reuse. Creating a fresh receiver was itself a + # source of REMOTE_DISCONNECT (new agent overwhelms UCX). + time.sleep(wait) + last_err = e + continue + raise + else: + raise RuntimeError(f"NIXL pull failed after {self._NIXL_PULL_MAX_RETRIES} retries") from last_err + # Reconstruct indexed fields (e.g. prompt_embeds_0, prompt_embeds_1 + # → prompt_embeds list) + reconstructed = {} + indexed = {} # base_name → {idx: tensor} + for k, v in tensors.items(): + parts = k.rsplit("_", 1) + if len(parts) == 2 and parts[1].isdigit(): + indexed.setdefault(parts[0], {})[int(parts[1])] = v + else: + reconstructed[k] = v + for base, idx_map in indexed.items(): + reconstructed[base] = [idx_map[i] for i in sorted(idx_map)] + from sglang_utils import inject_tensors_to_req + inject_tensors_to_req(batch, reconstructed) + return batch + + def _tp_broadcast_fields(self, batch: Req, src_rank: int = 0) -> None: + """Broadcast tensor fields from src_rank to all TP ranks. + + Uses broadcast_object_list for shape/dtype metadata, then + point-to-point broadcast for GPU tensor data. + """ + import torch.distributed as dist + tp_rank = get_tp_rank() + device = torch.device("cuda") + + # Collect field data on rank 0, None on others + field_data = [None] # single-element list for broadcast_object_list + if tp_rank == src_rank: + meta = {} + for field in self._tensor_fields: + val = getattr(batch, field, None) + if val is None: + continue + if isinstance(val, list): + meta[field] = [{"shape": list(t.shape), "dtype": str(t.dtype)} for t in val if isinstance(t, torch.Tensor)] + elif isinstance(val, torch.Tensor): + meta[field] = {"shape": list(val.shape), "dtype": str(val.dtype)} + field_data = [meta] + + dist.broadcast_object_list(field_data, src=src_rank) + meta = field_data[0] + if meta is None: + return + + # Allocate tensors on non-src ranks, then broadcast data + for field, info in meta.items(): + if isinstance(info, list): + # List of tensors + if tp_rank == src_rank: + tensors = getattr(batch, field) + else: + tensors = [torch.empty(m["shape"], dtype=getattr(torch, m["dtype"].removeprefix("torch.")), device=device) for m in info] + setattr(batch, field, tensors) + for t in tensors: + dist.broadcast(t, src=src_rank) + elif isinstance(info, dict): + # Single tensor + if tp_rank != src_rank: + t = torch.empty(info["shape"], dtype=getattr(torch, info["dtype"].removeprefix("torch.")), device=device) + setattr(batch, field, t) + else: + t = getattr(batch, field) + dist.broadcast(t, src=src_rank) + + def _device_move(self, batch: Req) -> Req: + """Fallback: move CPU tensors to GPU (ZMQ pickle path).""" + device = torch.device("cuda") + for field in self._tensor_fields: + val = getattr(batch, field, None) + if val is None: + continue + if isinstance(val, list): + setattr(batch, field, [ + t.to(device).contiguous() if isinstance(t, torch.Tensor) else t + for t in val + ]) + elif isinstance(val, torch.Tensor): + setattr(batch, field, val.to(device).contiguous()) + return batch + + +class NixlSendStage(PipelineStage): + """Register ``Req`` tensors as NIXL-readable and return metadata. + + Append as the **last** stage in encoder / denoiser partial pipelines. + Returns an ``OutputBatch`` whose ``output`` dict contains: + - ``_nixl_transfer_meta``: NIXL metadata for the receiver (~1.5 KB) + - tensor shapes/dtypes for logging + + Actual tensor data stays on GPU — only metadata travels over ZMQ. + Falls back to sending raw tensors when NIXL is unavailable. + + Holds ``readable_op`` handles in ``_active_readables`` until the + receiver completes the RDMA pull — prevents premature GC of the + Connection/Descriptor/GPU buffer that causes REMOTE_DISCONNECT. + """ + + def __init__(self, output_fields: List[str]): + super().__init__() + self._output_fields = output_fields + # Eagerly create sender so its NIXL agent/UCX endpoint is ready + # before any transfer occurs in this process. + # Only rank 0 uses NIXL — TP slaves skip to avoid UCX conflicts. + from nixl_transfer import NIXL_AVAILABLE + if NIXL_AVAILABLE and get_tp_rank() == 0: + from nixl_transfer import NixlTensorSender + self._sender = NixlTensorSender() + else: + self._sender = None + # (readable_op, flat_buffer_ref) — kept alive until receiver completes + self._active_readables: list[tuple[object, object]] = [] + + def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: + # TP rank > 0: only rank 0's output is returned to the orchestrator + if get_tp_rank() != 0: + return OutputBatch(output={}, metrics=RequestMetrics(request_id="nixl")) + + self._poll_completed() # release finished readables first + + tensors = self._extract_tensors(batch) + if not tensors: + return OutputBatch(output={}, metrics=RequestMetrics(request_id="nixl")) + + from nixl_transfer import NIXL_AVAILABLE + if NIXL_AVAILABLE and self._sender is not None: + return self._nixl_send(tensors) + return self._fallback_send(tensors) + + def _poll_completed(self): + """Release readable_ops whose RDMA pull has completed.""" + still_active = [] + for readable, buf_ref in self._active_readables: + try: + status = readable.status + name = getattr(status, "name", str(status)) + if "COMPLETE" in name: + logger.debug("NIXL readable completed, releasing buffer") + continue # drop — GC releases descriptor + buffer + still_active.append((readable, buf_ref)) + except Exception: + logger.warning("NIXL readable status check failed, keeping alive to be safe") + still_active.append((readable, buf_ref)) + self._active_readables = still_active + + def _extract_tensors(self, batch: Req) -> Dict[str, torch.Tensor]: + """Flatten list-valued fields into individual tensors.""" + result: Dict[str, torch.Tensor] = {} + for field in self._output_fields: + val = getattr(batch, field, None) + if val is None: + continue + if isinstance(val, list): + if len(val) == 0: + continue + if len(val) == 1: + result[field] = val[0] + else: + # Dual-encoder: store each element separately for NIXL + for i, t in enumerate(val): + result[f"{field}_{i}"] = t + elif isinstance(val, torch.Tensor): + result[field] = val + return result + + def _nixl_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: + meta, readable_op = self._sender.send(tensors) + # Hold readable_op (+ implicit buffer ref) until receiver completes pull + self._active_readables.append((readable_op, None)) + return OutputBatch(output={"_nixl_transfer_meta": meta}, metrics=RequestMetrics(request_id="nixl")) + + def _fallback_send(self, tensors: Dict[str, torch.Tensor]) -> OutputBatch: + """Fallback: send raw tensors via ZMQ pickle.""" + # Reconstruct list-valued fields for backward compat + output: Dict[str, object] = {} + for k, v in tensors.items(): + if "_" in k and k.rsplit("_", 1)[1].isdigit(): + base, idx = k.rsplit("_", 1) + output.setdefault(f"_list_{base}", {})[int(idx)] = v + else: + output[k] = v + # Reassemble lists + for base, idx_map in list(output.items()): + if base.startswith("_list_"): + real_key = base[6:] + output[real_key] = [idx_map[i] for i in sorted(idx_map)] + del output[base] + return OutputBatch(output=output, metrics=RequestMetrics(request_id="nixl")) + + +# ═══════════════════════════════════════════════════════════════════════ +# Stage builder functions (picklable — used as subprocess args) +# ═══════════════════════════════════════════════════════════════════════ + + +def build_encoder_stages(pipeline, server_args): + """TextEncodingStage → NixlSendStage(prompt_embeds, …). + + Automatically detects all loaded text encoders/tokenizers so that + both single-encoder (Wan) and dual-encoder (HunyuanVideo) models work. + """ + from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import ( + TextEncodingStage, + ) + from sglang_utils import get_component_backend + + # Collect all available text encoders and tokenizers + text_encoders = [] + tokenizers = [] + for name in ["text_encoder", "text_encoder_2", "text_encoder_3"]: + enc = pipeline.get_module(name) + if enc is not None: + text_encoders.append(enc) + logger.info("%s backend: %s", name, get_component_backend(enc)) + for name in ["tokenizer", "tokenizer_2", "tokenizer_3"]: + tok = pipeline.get_module(name) + if tok is not None: + tokenizers.append(tok) + + assert len(text_encoders) > 0, "No text encoders found in pipeline" + assert len(text_encoders) == len(tokenizers), ( + f"Encoder/tokenizer count mismatch: {len(text_encoders)} vs {len(tokenizers)}" + ) + + # Always list both fields; NixlSendStage skips None values, + # so negative_prompt_embeds is simply omitted when CFG is disabled. + return [ + TextEncodingStage(text_encoders=text_encoders, tokenizers=tokenizers), + NixlSendStage(["prompt_embeds", "negative_prompt_embeds"]), + ] + + +def build_denoiser_stages(pipeline, server_args): + """NixlReceive → LatentPrep → TimestepPrep → Denoising → NixlSend.""" + from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import ( + LatentPreparationStage, + ) + from sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation import ( + TimestepPreparationStage, + ) + from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import ( + DenoisingStage, + ) + from sglang_utils import get_component_backend + + transformer = pipeline.get_module("transformer") + scheduler = pipeline.get_module("scheduler") + logger.info("transformer backend: %s", get_component_backend(transformer)) + + return [ + NixlReceiveStage(["prompt_embeds", "negative_prompt_embeds"]), + LatentPreparationStage(scheduler=scheduler, transformer=transformer), + TimestepPreparationStage(scheduler=scheduler), + DenoisingStage(transformer=transformer, scheduler=scheduler), + NixlSendStage(["latents"]), + ] + + +def build_vae_stages(pipeline, server_args): + """NixlReceive → DecodingStage.""" + from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import ( + DecodingStage, + ) + from sglang_utils import get_component_backend + + vae = pipeline.get_module("vae") + logger.info("vae backend: %s", get_component_backend(vae)) + return [ + NixlReceiveStage(["latents"]), + DecodingStage(vae=vae, pipeline=pipeline), + ] + + +# ═══════════════════════════════════════════════════════════════════════ +# PartialGPUWorker +# ═══════════════════════════════════════════════════════════════════════ + + +class PartialGPUWorker(GPUWorker): + """GPUWorker that loads only a subset of pipeline modules. + + Only ``init_device_and_model`` is overridden. Everything else + (``execute_forward``, ``do_mem_analysis``, LoRA, etc.) is inherited. + """ + + def __init__( + self, + required_modules: List[str], + custom_stages_fn: Optional[Callable] = None, + **kwargs, + ): + self._required_modules = required_modules + self._custom_stages_fn = custom_stages_fn + super().__init__(**kwargs) + + def init_device_and_model(self) -> None: + torch.get_device_module().set_device(self.local_rank) + + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(self.master_port) + os.environ["LOCAL_RANK"] = str(self.local_rank) + os.environ["RANK"] = str(self.rank) + os.environ["WORLD_SIZE"] = str(self.server_args.num_gpus) + + dist_kwargs = dict( + tp_size=self.server_args.tp_size, + enable_cfg_parallel=self.server_args.enable_cfg_parallel, + ulysses_degree=getattr(self.server_args, "ulysses_degree", 1), + ring_degree=getattr(self.server_args, "ring_degree", 1), + sp_size=self.server_args.sp_degree, + dp_size=self.server_args.dp_size, + distributed_init_method=f"tcp://127.0.0.1:{self.master_port}", + ) + # dist_timeout only in newer sglang versions + import inspect + sig = inspect.signature(maybe_init_distributed_environment_and_model_parallel) + if "dist_timeout" in sig.parameters: + dist_kwargs["dist_timeout"] = self.server_args.dist_timeout + maybe_init_distributed_environment_and_model_parallel(**dist_kwargs) + + if model_parallel_is_initialized(): + suffix = "" + if get_tp_world_size() != 1: + suffix += f"_TP{get_tp_rank()}" + if get_ulysses_parallel_world_size() != 1: + suffix += f"_U{get_ulysses_parallel_rank()}" + if get_ring_parallel_world_size() != 1: + suffix += f"_R{get_ring_parallel_rank()}" + if get_classifier_free_guidance_world_size() != 1: + suffix += f"_C{get_classifier_free_guidance_rank()}" + setproctitle(f"sgl_diffusion::partial{suffix}") + else: + setproctitle(f"sgl_diffusion::partial_{self.local_rank}") + + from sglang_utils import build_partial_pipeline + + self.pipeline = build_partial_pipeline( + self.server_args, + required_modules=self._required_modules, + ) + + if self._custom_stages_fn is not None: + stages = self._custom_stages_fn(self.pipeline, self.server_args) + for stage in stages: + name = type(stage).__name__ + self.pipeline.add_stage(stage, name) + + if getattr(self.server_args, "dit_layerwise_offload", False) and OffloadableDiTMixin is not None: + for module_name in [ + "transformer", "transformer_2", + "video_dit", "video_dit_2", "audio_dit", + ]: + dit = self.pipeline.get_module(module_name) + if dit is not None: + if isinstance(dit, OffloadableDiTMixin): + dit.configure_layerwise_offload(self.server_args) + else: + logger.info( + "Module %s does not support layerwise offload.", + type(dit).__name__, + ) + + logger.info( + "PartialGPUWorker %d: modules=%s stages=%s", + self.rank, + self._required_modules, + list(self.pipeline._stage_name_mapping.keys()), + ) + + +# ═══════════════════════════════════════════════════════════════════════ +# Subprocess entry point + launcher +# ═══════════════════════════════════════════════════════════════════════ + + +def _patch_triton_norm_contiguous(): + """Patch SGLang's triton norm_infer to handle non-contiguous tensors. + + HunyuanVideo's transformer produces non-contiguous intermediate tensors + (e.g. from attention reshapes) which trip the triton kernel assertion + ``assert x.stride(-1) == 1``. + + Must patch both ``triton_ops`` and ``layernorm`` modules because + layernorm uses ``from triton_ops import norm_infer`` (direct binding). + """ + try: + import sglang.multimodal_gen.runtime.layers.triton_ops as triton_ops + import sglang.multimodal_gen.runtime.layers.layernorm as layernorm_mod + + _orig = triton_ops.norm_infer + + def _safe_norm_infer(x, *args, **kwargs): + if not x.is_contiguous(): + x = x.contiguous() + return _orig(x, *args, **kwargs) + + triton_ops.norm_infer = _safe_norm_infer + layernorm_mod.norm_infer = _safe_norm_infer + except (ImportError, AttributeError): + pass + + +def _run_partial_scheduler_process( + required_modules: List[str], + custom_stages_fn: Callable, + local_rank: int, + rank: int, + master_port: int, + server_args: ServerArgs, + pipe_writer, + task_pipe_r, + result_pipe_w, + task_pipes_to_slaves: list, + result_pipes_from_slaves: list, +) -> None: + """Subprocess entry point: patch GPUWorker then delegate to sglang.""" + # Ensure sglang_utils etc. are importable in the subprocess + workers_dir = os.path.dirname(os.path.abspath(__file__)) + if workers_dir not in sys.path: + sys.path.insert(0, workers_dir) + + _patch_triton_norm_contiguous() + + import sglang.multimodal_gen.runtime.managers.scheduler as sched_mod + + _rm, _csf = required_modules, custom_stages_fn + + class _PatchedGPUWorker(PartialGPUWorker): + def __init__(self, **kwargs): + super().__init__(required_modules=_rm, custom_stages_fn=_csf, **kwargs) + + sched_mod.GPUWorker = _PatchedGPUWorker + + from sglang.multimodal_gen.runtime.managers.gpu_worker import ( + run_scheduler_process, + ) + + run_scheduler_process( + local_rank, rank, master_port, server_args, pipe_writer, + task_pipe_r, result_pipe_w, + task_pipes_to_slaves, result_pipes_from_slaves, + ) + + +def launch_partial_server( + server_args: ServerArgs, + required_modules: List[str], + custom_stages_fn: Callable, +) -> list[mp.Process]: + """Spawn Scheduler subprocess(es) with ``PartialGPUWorker``. + + Mirrors sglang's ``launch_server()`` logic: spawns ``num_gpus`` + processes with master/slave pipe wiring so that TP, SP, and CFG + parallelism all work out of the box. + + * ``num_gpus == 1`` → single process, no pipes. + * ``num_gpus > 1`` → rank 0 is master (has ZMQ receiver + pipes to + slaves), ranks 1..N are slaves (coordinate via ``torch.distributed`` + broadcast). + + Returns the process list (for later shutdown). + + After this returns, call ``async_scheduler_client.initialize(server_args)`` + to connect from the main process. + """ + from sglang.multimodal_gen.runtime.utils.logging_utils import configure_logger + + configure_logger(server_args) + + num_gpus = server_args.num_gpus + master_port = server_args.master_port + processes: list[mp.Process] = [] + + # --- pipes for master ↔ slave coordination (same as launch_server) --- + task_pipes_to_slaves_w: list = [] + task_pipes_to_slaves_r: list = [] + result_pipes_from_slaves_w: list = [] + result_pipes_from_slaves_r: list = [] + + for _ in range(num_gpus - 1): + r, w = mp.Pipe(duplex=False) + task_pipes_to_slaves_r.append(r) + task_pipes_to_slaves_w.append(w) + + for _ in range(num_gpus - 1): + r, w = mp.Pipe(duplex=False) + result_pipes_from_slaves_r.append(r) + result_pipes_from_slaves_w.append(w) + + # --- spawn one process per GPU --- + readiness_readers: list = [] + + for i in range(num_gpus): + reader, writer = mp.Pipe(duplex=False) + readiness_readers.append(reader) + + if i == 0: + # Rank 0 (master): owns ZMQ receiver + write-ends to slaves + args = ( + required_modules, + custom_stages_fn, + i, # local_rank + i, # rank + master_port, + server_args, + writer, # pipe_writer (readiness signal) + None, # task_pipe_r (master doesn't receive tasks) + None, # result_pipe_w (master doesn't send results) + task_pipes_to_slaves_w, + result_pipes_from_slaves_r, + ) + else: + # Rank > 0 (slave): coordinates via torch.distributed broadcast + args = ( + required_modules, + custom_stages_fn, + i, # local_rank + i, # rank + master_port, + server_args, + writer, # pipe_writer (readiness signal) + None, # task_pipe_r + None, # result_pipe_w + task_pipes_to_slaves_r[i - 1], + result_pipes_from_slaves_w[i - 1], + ) + + process = mp.Process( + target=_run_partial_scheduler_process, + args=args, + name=f"sglang-partial-worker-{i}", + daemon=True, + ) + process.start() + writer.close() + processes.append(process) + + # --- close unused pipe ends in the parent (same as launch_server) --- + for p in task_pipes_to_slaves_w: + p.close() + for p in task_pipes_to_slaves_r: + p.close() + for p in result_pipes_from_slaves_w: + p.close() + for p in result_pipes_from_slaves_r: + p.close() + + # --- wait for all workers to report ready --- + for i, reader in enumerate(readiness_readers): + try: + data = reader.recv() + except EOFError: + processes[i].join() + raise RuntimeError( + f"Partial scheduler rank {i} died (exit code {processes[i].exitcode}). " + "Check logs above for errors." + ) + reader.close() + if data.get("status") != "ready": + raise RuntimeError(f"Partial scheduler rank {i} failed to initialize.") + + pids = [p.pid for p in processes] + logger.info( + "Partial scheduler ready: %d GPU(s), pids=%s", num_gpus, pids, + ) + return processes diff --git a/examples/disagg_diffusion/workers/protocol.py b/examples/disagg_diffusion/workers/protocol.py new file mode 100644 index 000000000000..c86082f2af34 --- /dev/null +++ b/examples/disagg_diffusion/workers/protocol.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Protocol types for disaggregated diffusion stages. + +Uses NIXL RDMA for GPU-direct tensor transfer between stage workers. +Only small metadata (shapes, dtypes, NIXL descriptor) travels over Dynamo RPC. +Tensor transfer is handled by NixlSendStage/NixlReceiveStage in +partial_gpu_worker.py and nixl_transfer.py. +""" + +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, Field + + +# --------------------------------------------------------------------------- +# Stage 1: Encoder +# --------------------------------------------------------------------------- + +class EncoderRequest(BaseModel): + prompt: str + negative_prompt: str = "" + guidance_scale: float = 1.0 + + +class EncoderResponse(BaseModel): + transfer_meta: Dict[str, Any] = Field(default_factory=dict) + shapes: Dict[str, List[int]] = Field(default_factory=dict) + + +# --------------------------------------------------------------------------- +# Stage 2: Denoiser +# --------------------------------------------------------------------------- + +class DenoiserRequest(BaseModel): + transfer_meta: Dict[str, Any] + tensor_data: Dict[str, Any] = Field(default_factory=dict) + height: int = 544 + width: int = 960 + num_frames: int = 61 + num_inference_steps: int = 50 + guidance_scale: float = 1.0 + seed: int = 42 + + +class DenoiserResponse(BaseModel): + transfer_meta: Dict[str, Any] = Field(default_factory=dict) + shape: List[int] = Field(default_factory=list) + + +# --------------------------------------------------------------------------- +# Stage 3: VAE Decoder +# --------------------------------------------------------------------------- + +class VAEDecodeRequest(BaseModel): + transfer_meta: Dict[str, Any] + tensor_data: Dict[str, Any] = Field(default_factory=dict) + request_id: str = "" + + +class VAEDecodeResponse(BaseModel): + video_path: str = "" + num_frames: int = 0 + + +# --------------------------------------------------------------------------- +# End-to-end (orchestrator convenience) +# --------------------------------------------------------------------------- + +class GenerateRequest(BaseModel): + prompt: str + negative_prompt: str = "" + height: int = 544 + width: int = 960 + num_frames: int = 61 + num_inference_steps: int = 50 + guidance_scale: float = 1.0 + seed: Optional[int] = None + + +# --------------------------------------------------------------------------- +# Health (per-stage) +# --------------------------------------------------------------------------- + +class HealthRequest(BaseModel): + pass + + +class HealthResponse(BaseModel): + status: str = "ok" + stage: str = "" + model: str = "" diff --git a/examples/disagg_diffusion/workers/sglang_utils.py b/examples/disagg_diffusion/workers/sglang_utils.py new file mode 100644 index 000000000000..47e181bb121c --- /dev/null +++ b/examples/disagg_diffusion/workers/sglang_utils.py @@ -0,0 +1,386 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""SGLang PipelineStage utilities for disaggregated diffusion workers. + +Provides helpers to load partial pipelines (only the modules each worker +needs), launch stage servers, and convert between Dynamo protocol types +and SGLang's Req dataclass. + +Also contains shared utilities (StageClient, model detection, compatibility +patches) used by Dynamo workers. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +from typing import Dict, List, Optional + +import torch + +logger = logging.getLogger(__name__) + + +def build_partial_pipeline( + server_args, + required_modules: List[str], +): + """Load a pipeline with only *required_modules* populated. + + Auto-detects pipeline class from model_index.json, suppresses automatic + stage creation, and syncs all component configs (even unloaded ones). + """ + from sglang.multimodal_gen.runtime.pipelines_core import get_model_info + + model_info = get_model_info(server_args.model_path) + base_pipeline_cls = model_info.pipeline_cls + + # Build a partial pipeline class that: + # 1. Suppresses automatic stage creation + # 2. Skips LoRA init (which assumes 'transformer' always exists) + def _noop_create_stages(self, server_args): + return None + + def _safe_init(self, **kwargs): + # Call ComposedPipelineBase.__init__ directly, skipping LoRAPipeline + # which tries to access self.modules['transformer'] + from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, + ) + ComposedPipelineBase.__init__(self, **kwargs) + + partial_cls = type( + f"_Partial{base_pipeline_cls.__name__}", + (base_pipeline_cls,), + { + "create_pipeline_stages": _noop_create_stages, + "__init__": _safe_init, + }, + ) + + pipeline = partial_cls( + model_path=server_args.model_path, + server_args=server_args, + required_config_modules=required_modules, + ) + + _sync_all_component_configs(server_args, pipeline) + + return pipeline + + +def _sync_all_component_configs(server_args, pipeline): + """Read config.json for every component in model_index.json and update + the corresponding arch_config in ``server_args.pipeline_config``, ensuring + correct parameters even for components whose weights are not loaded. + """ + import json + + CONFIG_ATTR_MAP = { + "vae": ("vae_config", "update_model_arch"), + "video_vae": ("vae_config", "update_model_arch"), + "transformer": ("dit_config", "update_model_arch"), + "video_dit": ("dit_config", "update_model_arch"), + "audio_dit": ("audio_dit_config", "update_model_arch"), + "audio_vae": ("audio_vae_config", "update_model_arch"), + } + + # pipeline.model_path is resolved to a local path by _load_config() + # (may differ from server_args.model_path if the original was a hub ID). + model_path = pipeline.model_path + model_index_path = os.path.join(model_path, "model_index.json") + if not os.path.isfile(model_index_path): + return + + with open(model_index_path, "r") as f: + model_index = json.load(f) + + pipeline_config = server_args.pipeline_config + for component_name, mapping in CONFIG_ATTR_MAP.items(): + config_attr, update_method_name = mapping + cfg = getattr(pipeline_config, config_attr, None) + if cfg is None: + continue + + if component_name not in model_index: + continue + + config_json_path = os.path.join(model_path, component_name, "config.json") + if not os.path.isfile(config_json_path): + continue + + with open(config_json_path, "r") as f: + hf_config = json.load(f) + + hf_config.pop("_class_name", None) + hf_config.pop("_diffusers_version", None) + + update_fn = getattr(cfg, update_method_name, None) + if update_fn is not None: + update_fn(hf_config) + logger.info("Synced %s config from %s (e.g. z_dim=%s)", + config_attr, + config_json_path, + getattr(getattr(cfg, "arch_config", cfg), "z_dim", "N/A")) + + if hasattr(cfg, "post_init"): + cfg.post_init() + + +def get_component_backend(module) -> str: + """Return a human-readable string indicating which backend loaded *module*.""" + mod = type(module).__module__ or "" + cls = type(module).__qualname__ + if mod.startswith("sglang."): + return f"sglang-optimized ({cls})" + if mod.startswith("diffusers."): + return f"native-diffusers ({cls})" + if mod.startswith("transformers."): + return f"native-transformers ({cls})" + return f"unknown ({mod}.{cls})" + + +def build_req( + prompt: str, + negative_prompt: Optional[str] = "", + height: int = 544, + width: int = 960, + num_frames: int = 61, + num_inference_steps: int = 50, + guidance_scale: float = 1.0, + seed: int = 42, + device: str = "cuda", + **extra_fields, +) -> "Req": + """Construct a minimal SGLang ``Req`` for running pipeline stages.""" + from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req + from sglang.multimodal_gen.configs.sample.sampling_params import DataType + + req = Req( + data_type=DataType.VIDEO, + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + seed=seed, + generator=torch.Generator(device="cpu").manual_seed(seed), + do_classifier_free_guidance=(guidance_scale > 1.0), + save_output=False, + return_file_paths_only=False, + ) + + for k, v in extra_fields.items(): + setattr(req, k, v) + + return req + + +def inject_tensors_to_req( + req, + tensors: Dict[str, object], + list_fields: Optional[List[str]] = None, +): + """Inject received tensors back into a ``Req``. + + Values may be bare tensors (single-encoder outputs) or lists of tensors + (dual-encoder outputs like HunyuanVideo's Llama + CLIP embeddings). + ``list_fields`` specifies Req attributes that expect a list value; + bare tensors are auto-wrapped in a single-element list. + """ + list_fields = set(list_fields or [ + "prompt_embeds", + "negative_prompt_embeds", + "pooled_embeds", + "neg_pooled_embeds", + "image_embeds", + ]) + for key, value in tensors.items(): + if key in list_fields: + if isinstance(value, list): + # Multi-encoder: already a list of tensors + setattr(req, key, value) + else: + # Single encoder: wrap in list + setattr(req, key, [value]) + else: + setattr(req, key, value) + return req + + +# ═══════════════════════════════════════════════════════════════════════ +# Shared utilities — used by Dynamo workers and the standalone E2E script +# ═══════════════════════════════════════════════════════════════════════ + + +class StageClient: + """Async ZMQ REQ client that talks to a SGLang Scheduler subprocess.""" + + FORWARD_TIMEOUT_S = float(os.environ.get("STAGE_FORWARD_TIMEOUT_S", "120")) + + def __init__(self, endpoint: str, name: str = ""): + import zmq, zmq.asyncio + self._zmq = zmq + self._name = name + self._endpoint = endpoint + self._ctx = zmq.asyncio.Context() + self._sock = self._ctx.socket(zmq.REQ) + self._sock.connect(endpoint) + self._lock = asyncio.Lock() + logger.info("StageClient(%s) connected to %s", name, endpoint) + + def _reset_socket(self): + """Recreate ZMQ socket after a timeout leaves it in a broken state. + + ZMQ REQ sockets enforce strict send/recv alternation. If recv + times out after send, the socket is stuck waiting for a reply and + all subsequent sends raise ``EFSM``. The only recovery is to + close and reconnect. + """ + self._sock.close(linger=0) + self._sock = self._ctx.socket(self._zmq.REQ) + self._sock.connect(self._endpoint) + logger.warning("StageClient(%s) reset ZMQ socket after timeout", self._name) + + async def forward(self, reqs): + """Send request(s) and receive response (with timeout).""" + async with self._lock: + await self._sock.send_pyobj(reqs) + try: + return await asyncio.wait_for( + self._sock.recv_pyobj(), + timeout=self.FORWARD_TIMEOUT_S, + ) + except (asyncio.TimeoutError, TimeoutError): + self._reset_socket() + raise + + def close(self): + self._sock.close() + self._ctx.term() + + +def patch_hunyuan_config(): + """HunyuanConfig inherits ``task_type`` from PipelineConfig without a + default value, so ``HunyuanConfig()`` crashes. Wrap __init__ to supply + ``task_type=T2V`` when omitted. Idempotent. + """ + from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType + try: + from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import ( + HunyuanConfig, FastHunyuanConfig, + ) + except ImportError: + return + + for cls in (HunyuanConfig, FastHunyuanConfig): + if getattr(cls, "_task_type_patched", False): + continue + orig = cls.__init__ + + def _patched(self, *a, task_type=ModelTaskType.T2V, _orig=orig, **kw): + _orig(self, *a, task_type=task_type, **kw) + + cls.__init__ = _patched + cls._task_type_patched = True + + +def detect_encoder_modules(model_path: str) -> List[str]: + """Return the required_modules list for the encoder stage. + + Auto-detects dual-encoder models (e.g. HunyuanVideo with Llama + CLIP). + """ + try: + from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( + maybe_download_model_index, verify_model_config_and_directory, + ) + config = (verify_model_config_and_directory(model_path) + if os.path.exists(model_path) + else maybe_download_model_index(model_path)) + modules = ["text_encoder", "tokenizer"] + if "text_encoder_2" in config: + modules += ["text_encoder_2", "tokenizer_2"] + modules.append("scheduler") + return modules + except Exception: + pass + # Fallback: include dual encoders for known models + if "hunyuan" in model_path.lower(): + return ["text_encoder", "text_encoder_2", + "tokenizer", "tokenizer_2", "scheduler"] + return ["text_encoder", "tokenizer", "scheduler"] + + +def save_video(frames_tensor, output_path: str, fps: int = 24): + """Save decoded video tensor [B,C,T,H,W] as mp4. + + Returns (filepath, num_frames). + """ + import numpy as np + + if isinstance(frames_tensor, dict): + for v in frames_tensor.values(): + if hasattr(v, "shape"): + frames_tensor = v + break + + if hasattr(frames_tensor, "cpu"): + frames_tensor = frames_tensor.cpu().float().numpy() + + # [B, C, T, H, W] -> [T, H, W, C] + frames = (frames_tensor[0].transpose(1, 2, 3, 0) * 255).clip(0, 255).astype(np.uint8) + + os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True) + try: + import imageio + imageio.mimwrite(output_path, frames, fps=fps, codec="libx264") + except Exception as e: + logger.warning("mp4 export failed (%s), saving first frame as PNG", e) + from PIL import Image + output_path = output_path.rsplit(".", 1)[0] + ".png" + Image.fromarray(frames[0]).save(output_path) + + return output_path, len(frames) + + +def launch_stage_server(model_path, required_modules, custom_stages_fn, + scheduler_port, tp_size=1, num_gpus=None, + client_name=""): + """Patch configs, create ServerArgs, launch Scheduler, return (processes, client, server_args). + + Consolidates the boilerplate shared by encoder, denoiser, and VAE workers: + patch_hunyuan_config → ServerArgs → set_global → launch_partial_server → StageClient. + """ + from sglang.multimodal_gen.runtime.server_args import ( + ServerArgs, set_global_server_args, + ) + from partial_gpu_worker import launch_partial_server + + patch_hunyuan_config() + + if num_gpus is None: + num_gpus = tp_size + + server_args = ServerArgs.from_kwargs( + model_path=model_path, + num_gpus=num_gpus, + tp_size=tp_size, + scheduler_port=scheduler_port, + ) + set_global_server_args(server_args) + + processes = launch_partial_server( + server_args, + required_modules=required_modules, + custom_stages_fn=custom_stages_fn, + ) + + endpoint = server_args.scheduler_endpoint + if callable(endpoint): + endpoint = endpoint() + client = StageClient(endpoint, client_name) + return processes, client, server_args diff --git a/examples/disagg_diffusion/workers/vae_worker.py b/examples/disagg_diffusion/workers/vae_worker.py new file mode 100755 index 000000000000..95f37d1eeb22 --- /dev/null +++ b/examples/disagg_diffusion/workers/vae_worker.py @@ -0,0 +1,145 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Disaggregated Diffusion — VAE Worker (Dynamo RPC) + +Wraps an SGLang Scheduler subprocess running NixlReceiveStage + DecodingStage. +Receives latents via NIXL RDMA, decodes to video frames, saves as mp4. + +Process architecture:: + + Dynamo Worker Process (this file) + |-- @dynamo_worker + | |-- serve_endpoint("generate") <-- Dynamo RPC from orchestrator + | | +-- handle_generate() + | | +-- StageClient.forward() <-- ZMQ to local Scheduler + | +-- serve_endpoint("health") + | + +-- SGLang Scheduler subprocess (spawned by launch_partial_server) + +-- PartialGPUWorker + |-- NixlReceiveStage <-- RDMA-pull latents from denoiser + +-- DecodingStage <-- VAE decode to video frames +""" + +import asyncio +import json +import logging +import multiprocessing as mp +import os +import sys +import uuid + +import uvloop + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from dynamo.runtime import DistributedRuntime, dynamo_worker # noqa: E402 + +logger = logging.getLogger(__name__) + +MODEL_PATH = os.environ.get("MODEL_PATH", "hunyuanvideo-community/HunyuanVideo") +SCHEDULER_PORT = int(os.environ.get("SCHEDULER_PORT", "15800")) +OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/tmp/disagg_videos") + + +@dynamo_worker(enable_nats=False) +async def worker(runtime: DistributedRuntime): + from sglang_utils import launch_stage_server, build_req, save_video + from partial_gpu_worker import build_vae_stages + + os.makedirs(OUTPUT_DIR, exist_ok=True) + + logger.info("Launching VAE Scheduler: port=%d", SCHEDULER_PORT) + processes, client, server_args = launch_stage_server( + MODEL_PATH, ["vae", "scheduler"], build_vae_stages, + SCHEDULER_PORT, client_name="vae", + ) + + # ── Dynamo RPC handlers ────────────────────────────────────────── + + async def handle_generate(request, context): + try: + if isinstance(request, str): + request = json.loads(request) + + req = build_req( + prompt="", + height=request.get("height", 544), + width=request.get("width", 960), + num_frames=request.get("num_frames", 61), + num_inference_steps=1, + guidance_scale=0.0, + seed=request.get("seed", 42), + ) + + # Pass NIXL metadata or raw tensor data for NixlReceiveStage + transfer_meta = request.get("transfer_meta", {}) + tensor_data = request.get("tensor_data", {}) + if transfer_meta: + req._nixl_transfer_meta = transfer_meta + elif tensor_data: + # ZMQ fallback: deserialize latents and inject onto Req + import torch, base64, io + from sglang_utils import inject_tensors_to_req + tensors = {} + for k, v in tensor_data.items(): + if isinstance(v, list): + tensors[k] = [torch.load(io.BytesIO(base64.b64decode(b)), weights_only=True) for b in v] + else: + tensors[k] = torch.load(io.BytesIO(base64.b64decode(v)), weights_only=True) + inject_tensors_to_req(req, tensors) + logger.info("Injected %d tensor fields via ZMQ fallback", len(tensors)) + + output = await client.forward([req]) + if output.error: + yield {"error": str(output.error), "video_path": "", "num_frames": 0} + return + + # Extract frames and save video + frames_tensor = output.output + request_id = request.get("request_id") or str(uuid.uuid4())[:8] + out_path = os.path.join(OUTPUT_DIR, f"{request_id}.mp4") + + loop = asyncio.get_event_loop() + filepath, n_frames = await loop.run_in_executor( + None, save_video, frames_tensor, out_path, + ) + filename = os.path.basename(filepath) + logger.info("Decoded — %d frames -> %s", n_frames, filename) + yield {"video_path": filename, "num_frames": n_frames} + + except Exception as e: + logger.error("VAE generate failed: %s", e, exc_info=True) + yield {"error": str(e), "video_path": "", "num_frames": 0} + + async def handle_health(request, context): + yield {"status": "ok", "stage": "vae", "model": MODEL_PATH} + + # ── Serve Dynamo endpoints ─────────────────────────────────────── + + gen_ep = runtime.endpoint("disagg_diffusion.vae.generate") + health_ep = runtime.endpoint("disagg_diffusion.vae.health") + + logger.info("Serving: disagg_diffusion.vae.generate + health") + try: + await asyncio.gather( + gen_ep.serve_endpoint(handle_generate), + health_ep.serve_endpoint(handle_health), + ) + finally: + client.close() + for p in processes: + p.terminate() + for p in processes: + p.join(timeout=10) + + +if __name__ == "__main__": + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(name)s] %(levelname)s %(message)s", + ) + mp.set_start_method("spawn", force=True) + uvloop.install() + asyncio.run(worker())