diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py index cd22b34e0..5b43400e8 100644 --- a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py +++ b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py @@ -3,6 +3,7 @@ from pathlib import Path import pytest import torch +from vllm.distributed.parallel_state import cleanup_dist_env_and_memory from tests.utils import GPUMemoryMonitor from vllm_omni.utils.platform_utils import is_npu, is_rocm @@ -17,34 +18,39 @@ from vllm_omni import Omni models = ["riverclouds/qwen_image_random"] +def inference(model_name: str, offload: bool = True): + torch.cuda.empty_cache() + device_index = torch.cuda.current_device() + monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02) + monitor.start() + m = Omni(model=model_name, enable_cpu_offload=offload) + torch.cuda.reset_peak_memory_stats(device=device_index) + height = 256 + width = 256 + + m.generate( + "a photo of a cat sitting on a laptop keyboard", + height=height, + width=width, + num_inference_steps=9, + guidance_scale=0.0, + generator=torch.Generator("cuda").manual_seed(42), + ) + peak = monitor.peak_used_mb + monitor.stop() + + return peak + + @pytest.mark.skipif(is_npu() or is_rocm(), reason="Hardware not supported") @pytest.mark.parametrize("model_name", models) def test_cpu_offload_diffusion_model(model_name: str): - def inference(offload: bool = True): - torch.cuda.empty_cache() - device_index = torch.cuda.current_device() - monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02) - monitor.start() - m = Omni(model=model_name, enable_cpu_offload=offload) - torch.cuda.reset_peak_memory_stats(device=device_index) - height = 256 - width = 256 - - m.generate( - "a photo of a cat sitting on a laptop keyboard", - height=height, - width=width, - num_inference_steps=9, - guidance_scale=0.0, - generator=torch.Generator("cuda").manual_seed(42), - ) - peak = monitor.peak_used_mb - monitor.stop() - - return peak - - offload_peak_memory = inference(offload=True) - no_offload_peak_memory = inference(offload=False) + try: + no_offload_peak_memory = inference(model_name, offload=False) + cleanup_dist_env_and_memory() + offload_peak_memory = inference(model_name, offload=True) + except Exception: + pytest.fail("Inference failed") print(f"Offload peak memory: {offload_peak_memory} MB") print(f"No offload peak memory: {no_offload_peak_memory} MB") assert offload_peak_memory + 2500 < no_offload_peak_memory, ( diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index e85ae0d76..2e00f41b9 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -61,130 +61,124 @@ class DiffusionEngine: raise e def step(self, requests: list[OmniDiffusionRequest]): - try: - # Apply pre-processing if available - if self.pre_process_func is not None: - preprocess_start_time = time.time() - requests = self.pre_process_func(requests) - preprocess_time = time.time() - preprocess_start_time - logger.info(f"Pre-processing completed in {preprocess_time:.4f} seconds") + # Apply pre-processing if available + if self.pre_process_func is not None: + preprocess_start_time = time.time() + requests = self.pre_process_func(requests) + preprocess_time = time.time() - preprocess_start_time + logger.info(f"Pre-processing completed in {preprocess_time:.4f} seconds") - output = self.add_req_and_wait_for_response(requests) - if output.error: - raise Exception(f"{output.error}") - logger.info("Generation completed successfully.") + output = self.add_req_and_wait_for_response(requests) + if output.error: + raise Exception(f"{output.error}") + logger.info("Generation completed successfully.") - if output.output is None: - logger.warning("Output is None, returning empty OmniRequestOutput") - # Return empty output for the first request - if len(requests) > 0: - request = requests[0] - request_id = request.request_id or "" - prompt = request.prompt - if isinstance(prompt, list): - prompt = prompt[0] if prompt else None - return OmniRequestOutput.from_diffusion( - request_id=request_id, - images=[], - prompt=prompt, - metrics={}, - latents=None, - ) - return None - - postprocess_start_time = time.time() - outputs = self.post_process_func(output.output) if self.post_process_func is not None else output.output - postprocess_time = time.time() - postprocess_start_time - logger.info(f"Post-processing completed in {postprocess_time:.4f} seconds") - - # Convert to OmniRequestOutput format - # Ensure outputs is a list - if not isinstance(outputs, list): - outputs = [outputs] if outputs is not None else [] - - # Handle single request or multiple requests - if len(requests) == 1: - # Single request: return single OmniRequestOutput + if output.output is None: + logger.warning("Output is None, returning empty OmniRequestOutput") + # Return empty output for the first request + if len(requests) > 0: request = requests[0] request_id = request.request_id or "" prompt = request.prompt if isinstance(prompt, list): prompt = prompt[0] if prompt else None + return OmniRequestOutput.from_diffusion( + request_id=request_id, + images=[], + prompt=prompt, + metrics={}, + latents=None, + ) + return None + + postprocess_start_time = time.time() + outputs = self.post_process_func(output.output) if self.post_process_func is not None else output.output + postprocess_time = time.time() - postprocess_start_time + logger.info(f"Post-processing completed in {postprocess_time:.4f} seconds") + + # Convert to OmniRequestOutput format + # Ensure outputs is a list + if not isinstance(outputs, list): + outputs = [outputs] if outputs is not None else [] + + # Handle single request or multiple requests + if len(requests) == 1: + # Single request: return single OmniRequestOutput + request = requests[0] + request_id = request.request_id or "" + prompt = request.prompt + if isinstance(prompt, list): + prompt = prompt[0] if prompt else None + + metrics = {} + if output.trajectory_timesteps is not None: + metrics["trajectory_timesteps"] = output.trajectory_timesteps + + if supports_audio_output(self.od_config.model_class_name): + audio_payload = outputs[0] if len(outputs) == 1 else outputs + return OmniRequestOutput.from_diffusion( + request_id=request_id, + images=[], + prompt=prompt, + metrics=metrics, + latents=output.trajectory_latents, + multimodal_output={"audio": audio_payload}, + final_output_type="audio", + ) + else: + return OmniRequestOutput.from_diffusion( + request_id=request_id, + images=outputs, + prompt=prompt, + metrics=metrics, + latents=output.trajectory_latents, + ) + else: + # Multiple requests: return list of OmniRequestOutput + # Split images based on num_outputs_per_prompt for each request + results = [] + output_idx = 0 + + for request in requests: + request_id = request.request_id or "" + prompt = request.prompt + if isinstance(prompt, list): + prompt = prompt[0] if prompt else None + + # Get images for this request + num_outputs = request.num_outputs_per_prompt + request_outputs = outputs[output_idx : output_idx + num_outputs] if output_idx < len(outputs) else [] + output_idx += num_outputs metrics = {} if output.trajectory_timesteps is not None: metrics["trajectory_timesteps"] = output.trajectory_timesteps if supports_audio_output(self.od_config.model_class_name): - audio_payload = outputs[0] if len(outputs) == 1 else outputs - return OmniRequestOutput.from_diffusion( - request_id=request_id, - images=[], - prompt=prompt, - metrics=metrics, - latents=output.trajectory_latents, - multimodal_output={"audio": audio_payload}, - final_output_type="audio", + audio_payload = request_outputs[0] if len(request_outputs) == 1 else request_outputs + results.append( + OmniRequestOutput.from_diffusion( + request_id=request_id, + images=[], + prompt=prompt, + metrics=metrics, + latents=output.trajectory_latents, + multimodal_output={"audio": audio_payload}, + final_output_type="audio", + ) ) else: - return OmniRequestOutput.from_diffusion( - request_id=request_id, - images=outputs, - prompt=prompt, - metrics=metrics, - latents=output.trajectory_latents, - ) - else: - # Multiple requests: return list of OmniRequestOutput - # Split images based on num_outputs_per_prompt for each request - results = [] - output_idx = 0 - - for request in requests: - request_id = request.request_id or "" - prompt = request.prompt - if isinstance(prompt, list): - prompt = prompt[0] if prompt else None - - # Get images for this request - num_outputs = request.num_outputs_per_prompt - request_outputs = ( - outputs[output_idx : output_idx + num_outputs] if output_idx < len(outputs) else [] - ) - output_idx += num_outputs - - metrics = {} - if output.trajectory_timesteps is not None: - metrics["trajectory_timesteps"] = output.trajectory_timesteps - - if supports_audio_output(self.od_config.model_class_name): - audio_payload = request_outputs[0] if len(request_outputs) == 1 else request_outputs - results.append( - OmniRequestOutput.from_diffusion( - request_id=request_id, - images=[], - prompt=prompt, - metrics=metrics, - latents=output.trajectory_latents, - multimodal_output={"audio": audio_payload}, - final_output_type="audio", - ) - ) - else: - results.append( - OmniRequestOutput.from_diffusion( - request_id=request_id, - images=request_outputs, - prompt=prompt, - metrics=metrics, - latents=output.trajectory_latents, - ) + results.append( + OmniRequestOutput.from_diffusion( + request_id=request_id, + images=request_outputs, + prompt=prompt, + metrics=metrics, + latents=output.trajectory_latents, ) + ) - return results - except Exception as e: - logger.error(f"Generation failed: {e}") - return None + return results @staticmethod def make_engine(config: OmniDiffusionConfig) -> "DiffusionEngine": diff --git a/vllm_omni/diffusion/scheduler.py b/vllm_omni/diffusion/scheduler.py index 386218a61..21b0d4914 100644 --- a/vllm_omni/diffusion/scheduler.py +++ b/vllm_omni/diffusion/scheduler.py @@ -62,6 +62,9 @@ class Scheduler: raise RuntimeError("Result queue not initialized") output = self.result_mq.dequeue() + # {"status": "error", "error": str(e)} + if isinstance(output, dict) and output.get("status") == "error": + raise RuntimeError("worker error") return output except zmq.error.Again: logger.error("Timeout waiting for response from scheduler.") diff --git a/vllm_omni/diffusion/worker/gpu_diffusion_worker.py b/vllm_omni/diffusion/worker/gpu_diffusion_worker.py index 5a4d4d016..90d169a0e 100644 --- a/vllm_omni/diffusion/worker/gpu_diffusion_worker.py +++ b/vllm_omni/diffusion/worker/gpu_diffusion_worker.py @@ -279,7 +279,7 @@ class WorkerProc: return result, should_reply except Exception as e: logger.error(f"Error executing RPC: {e}", exc_info=True) - return {"status": "error", "error": str(e)}, should_reply + raise e def worker_busy_loop(self) -> None: """Main busy loop for Multiprocessing Workers."""