[RFC]: Exit on OOM

Motivation

This RFC documents a requirement for vLLM-Omni to handle out-of-memory (OOM) errors in a manner consistent with vLLM: errors should propagate through all layers and trigger graceful termination of the entire inference process group.

Problem Statement (Current State Before #1163 )

When diffusion model stages encountered CUDA OOM errors during generation:

1. The error was caught locally in AsyncOmniDiffusion and wrapped in a RuntimeError
2. The stage worker continued running in a potentially corrupted state with partial allocations
3. Subsequent requests could fail unpredictably or hang indefinitely
4. No explicit cleanup was performed on the failed stage
5. The service appeared healthy but could not successfully process requests

Requirements

The implementation must satisfy the following requirements:

1. Consistent with vLLM behavior: OOM errors must propagate through all layers (engine → stage → entrypoint) and trigger graceful process termination, matching vLLM's established error handling pattern
2. Complete process group termination: The entire inference process group (including all stage workers and the main API process) must exit cleanly when any stage encounters OOM
3. Graceful resource cleanup: GPU resources must be properly released through standard cleanup mechanisms rather than forceful termination
4. Signal failure to orchestrator: The process exit must be detectable by external container orchestrators (Kubernetes, Docker) for automatic restart
5. Cover all stage types: Both diffusion and autoregressive (AR) stages must follow the same OOM propagation and termination path

Current Implementation in #1163

The current implementation is a partial step toward the requirement. Below is the current state and gaps from the target behavior.

Current Changes

vllm_omni/entrypoints/async_omni.py

Added handle_oom() method (marked as temporary workaround):

def handle_oom(self, result: dict[str, Any], stage: OmniStage, stage_id: int) -> None:
    if getattr(stage, "stage_type", None) == "diffusion":
        error_text = str(result.get("error", ""))
        error_lower = error_text.lower()
        is_oom = "out of memory" in error_lower or "cuda oom" in error_lower
        if is_oom:
            logger.critical(
                f"[{self._name}] Diffusion stage {stage_id} reported OOM; stopping stage worker.",
            )
            try:
                stage.stop_stage_worker()
            except Exception as e:
                logger.warning(
                    f"[{self._name}] Failed to stop diffusion stage {stage_id}: {e}",
                )
            logger.critical(
                f"[{self._name}] Exiting API process due to diffusion stage {stage_id} OOM.",
            )
            import os as _os
            _os._exit(1)

vllm_omni/entrypoints/async_omni_diffusion.py

Removed local try-except to let exceptions propagate:

# Before:
try:
    result = await loop.run_in_executor(self._executor, self.engine.step, request)
    result = result[0]
except Exception as e:
    logger.error("Generation failed for request %s: %s", request_id, e)
    raise RuntimeError(f"Diffusion generation failed: {e}") from e

# After:
result = await loop.run_in_executor(self._executor, self.engine.step, request)
result = result[0]

Gaps from Target Behavior

Aspect	Target (vLLM-like)	Current Implementation
Error propagation	Layer-by-layer propagation through engine → stage → entrypoint	Error detected at entrypoint level via string matching
OOM detection	Exception type based (e.g., torch.cuda.OutOfMemoryError)	String substring matching in error messages
Termination scope	Entire process group	Only API process via os._exit()
Cleanup mechanism	Graceful cleanup with proper resource release	Attempts stop_stage_worker() but force exits immediately
Stage coverage	All stage types (AR + diffusion)	Diffusion stages only
exit method	sys.exit() with cleanup handlers	os._exit(1) bypassing all cleanup

Open Issues

The following issues are not addressed by the current implementation:

1. Incomplete error propagation: OOM errors are detected at the entrypoint level rather than being propagated layer-by-layer from the engine through each abstraction level
2. Fragile error detection: String-based matching ("out of memory" in error_lower) is fragile and may break with PyTorch or driver updates
3. Inconsistent stage coverage: Only diffusion stages have special handling; AR stages use different code paths
4. Non-graceful termination: os._exit(1) bypasses Python cleanup handlers and does not ensure proper resource release across the entire process group
5. No alignment with vLLM patterns: The current implementation does not follow vLLM's established patterns for OOM propagation and process termination

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.