[Bugfix] raise error in diffusion engine and fix offload test (#933)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Jiangyun ZhuandHongsheng Liu authored and GitHub committed
1 parent
7f0825832a
commit
578ed1967e
4 files changed
+136
-133
No files matched your search
@@ -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, (
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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.")
|
||||
|
||||
@@ -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."""
|
||||
|
||||
Reference in new issue
Block a user