From 539fa570068305ab3c10d020f6ce740b9dd2ea02 Mon Sep 17 00:00:00 2001 From: "Will.hou" <1205157517@qq.com> Date: Mon, 27 Apr 2026 21:40:31 +0800 Subject: [PATCH] fix(diffusion): correct metric keys, remove duplication, minor cleanup (#2692) Signed-off-by: willamhou Co-authored-by: Claude Co-authored-by: Happy --- .../test_diffusion_engine_metrics.py | 105 ++++++++++++++++++ vllm_omni/diffusion/diffusion_engine.py | 15 +-- 2 files changed, 111 insertions(+), 9 deletions(-) create mode 100644 tests/diffusion/test_diffusion_engine_metrics.py diff --git a/tests/diffusion/test_diffusion_engine_metrics.py b/tests/diffusion/test_diffusion_engine_metrics.py new file mode 100644 index 000000000..dd8db9fdf --- /dev/null +++ b/tests/diffusion/test_diffusion_engine_metrics.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Source-level regression tests for DiffusionEngine. + +These tests verify naming conventions and patterns by inspecting source code +at the function level using AST. They are intentionally coupled to the source +layout and should be updated whenever the metrics construction code is +refactored. +""" + +from __future__ import annotations + +import ast +import os + +_ENGINE_PATH = os.path.normpath( + os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + "vllm_omni", + "diffusion", + "diffusion_engine.py", + ) +) + + +def _read_engine_source() -> str: + with open(_ENGINE_PATH) as f: + return f.read() + + +def _get_function_source(source: str, class_name: str | None, func_name: str) -> str: + """Extract the source of a specific function/method using AST. + + Args: + source: Full file source code. + class_name: Enclosing class name, or None for module-level functions. + func_name: Function/method name. + + Returns: + Source code of the function body. + """ + tree = ast.parse(source) + if class_name is not None: + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and node.name == class_name: + for item in node.body: + if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and item.name == func_name: + result = ast.get_source_segment(source, item) + assert result is not None, f"{class_name}.{func_name} source not found" + return result + else: + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name: + result = ast.get_source_segment(source, node) + assert result is not None, f"{func_name} source not found" + return result + raise AssertionError(f"Function {class_name + '.' if class_name else ''}{func_name} not found in source") + + +class TestMetricKeys: + """Verify metric naming conventions in DiffusionEngine.step().""" + + def test_no_duplicate_preprocess_key(self) -> None: + """step() should not contain 'preprocessing_time_ms' (duplicate of 'preprocess_time_ms').""" + source = _read_engine_source() + step_source = _get_function_source(source, "DiffusionEngine", "step") + assert "preprocessing_time_ms" not in step_source, ( + "Found duplicate key 'preprocessing_time_ms' in step() — should only use 'preprocess_time_ms'" + ) + + def test_metric_key_naming_consistency(self) -> None: + """In step(): exec_time should use exec_total_time, total_time should use step_total_ms.""" + source = _read_engine_source() + step_source = _get_function_source(source, "DiffusionEngine", "step") + lines = step_source.split("\n") + + found_exec = False + found_total = False + for line in lines: + if '"diffusion_engine_exec_time_ms"' in line: + found_exec = True + assert "exec_total_time" in line, ( + "diffusion_engine_exec_time_ms should measure executor time only (exec_total_time)" + ) + if '"diffusion_engine_total_time_ms"' in line: + found_total = True + assert "step_total_ms" in line, ( + "diffusion_engine_total_time_ms should measure full step time (step_total_ms)" + ) + assert found_exec, "diffusion_engine_exec_time_ms key not found in step()" + assert found_total, "diffusion_engine_total_time_ms key not found in step()" + + +class TestDummyRunAllocation: + """Verify _dummy_run generates exact-sized audio arrays.""" + + def test_no_oversized_allocation(self) -> None: + """_dummy_run should not allocate more audio than needed.""" + source = _read_engine_source() + dummy_source = _get_function_source(source, "DiffusionEngine", "_dummy_run") + assert "audio_sr * audio_duration_sec" not in dummy_source, ( + "_dummy_run should generate exact-sized audio, not allocate and slice" + ) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index abaf59895..80ef0df66 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -190,22 +190,21 @@ class DiffusionEngine: metrics = { "preprocess_time_ms": preprocess_time * 1000, - "diffusion_engine_exec_time_ms": (time.perf_counter() - diffusion_engine_start_time) * 1000, - "diffusion_engine_total_time_ms": exec_total_time * 1000, + "diffusion_engine_exec_time_ms": exec_total_time * 1000, + "diffusion_engine_total_time_ms": step_total_ms, "image_num": int(request.sampling_params.num_outputs_per_prompt), "resolution": int(request.sampling_params.resolution), "postprocess_time_ms": postprocess_time * 1000, } - if self.pre_process_func is not None: - metrics["preprocessing_time_ms"] = preprocess_time * 1000 # Handle single request or multiple requests + is_audio_output = supports_audio_output(self.od_config.model_class_name) if len(request.prompts) == 1: # Single request: return single OmniRequestOutput prompt = request.prompts[0] request_id = request.request_ids[0] if request.request_ids else "" - if supports_audio_output(self.od_config.model_class_name): + if is_audio_output: request_audio_payload = outputs[0] if len(outputs) == 1 else outputs return [ OmniRequestOutput.from_diffusion( @@ -265,7 +264,7 @@ class DiffusionEngine: request_outputs = outputs[start_idx:end_idx] if output_idx < len(outputs) else [] output_idx = end_idx - if supports_audio_output(self.od_config.model_class_name): + if is_audio_output: request_audio_payload = request_outputs[0] if len(request_outputs) == 1 else request_outputs results.append( OmniRequestOutput.from_diffusion( @@ -416,9 +415,7 @@ class DiffusionEngine: if supports_audio_input(self.od_config.model_class_name): audio_sr = 16000 - audio_duration_sec = 4 - audio_array = np.random.randn(audio_sr * audio_duration_sec).astype(np.float32) - dummy_audio = audio_array[audio_sr * 1 : audio_sr * 3] + dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32) else: dummy_audio = None