From 52afcb6055544bf180f872413264d5a992f6eb78 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Mon, 5 Jan 2026 17:32:20 +0800 Subject: [PATCH] [Feature] Support torch profiler across omni stages (#553) Signed-off-by: gcanlin Signed-off-by: lishunyang Co-authored-by: lishunyang --- .../offline_inference/qwen2_5_omni/end2end.py | 19 ++++++ .../offline_inference/qwen3_omni/end2end.py | 17 +++++ vllm_omni/entrypoints/async_omni.py | 36 ++++++++-- vllm_omni/entrypoints/async_omni_llm.py | 10 +-- vllm_omni/entrypoints/omni.py | 68 ++++++++++++++++++- vllm_omni/entrypoints/omni_stage.py | 67 ++++++++++++++++++ vllm_omni/entrypoints/stage_utils.py | 6 ++ 7 files changed, 214 insertions(+), 9 deletions(-) diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index d134e170c..d1e0264cd 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -6,6 +6,7 @@ with the correct prompt format on Qwen2.5-Omni """ import os +import time from typing import NamedTuple import librosa @@ -376,11 +377,17 @@ def main(args): for i, prompt in enumerate(prompts): prompt["modalities"] = output_modalities + profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) + if profiler_enabled: + omni_llm.start_profile(stages=[0]) omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator) # Determine output directory: prefer --output-dir; fallback to --output-wav output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav os.makedirs(output_dir, exist_ok=True) + + total_requests = len(prompts) + processed_count = 0 for stage_outputs in omni_generator: if stage_outputs.final_output_type == "text": for output in stage_outputs.request_output: @@ -408,6 +415,18 @@ def main(args): sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000) print(f"Request ID: {request_id}, Saved audio to {output_wav}") + processed_count += len(stage_outputs.request_output) + if profiler_enabled and processed_count >= total_requests: + print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...") + # Stop the profiler while workers are still alive + omni_llm.stop_profile() + + print("[Info] Waiting 30s for workers to write massive trace files to disk...") + time.sleep(30) + print("[Info] Trace export wait finished.") + + omni_llm.close() + def parse_args(): parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models") diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index 1441441c4..7280d7d2e 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -6,6 +6,7 @@ with the correct prompt format on Qwen3-Omni (thinker only). """ import os +import time from typing import NamedTuple import librosa @@ -319,11 +320,17 @@ def main(args): for i, prompt in enumerate(prompts): prompt["modalities"] = output_modalities + profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) + if profiler_enabled: + omni_llm.start_profile(stages=[0]) omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator) # Determine output directory: prefer --output-dir; fallback to --output-wav output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav os.makedirs(output_dir, exist_ok=True) + total_requests = len(prompts) + processed_count = 0 + for stage_outputs in omni_generator: if stage_outputs.final_output_type == "text": for output in stage_outputs.request_output: @@ -359,6 +366,16 @@ def main(args): # Save audio file with explicit WAV format sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV") print(f"Request ID: {request_id}, Saved audio to {output_wav}") + + processed_count += len(stage_outputs.request_output) + if profiler_enabled and processed_count >= total_requests: + print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...") + # Stop the profiler while workers are still alive + omni_llm.stop_profile() + + print("[Info] Waiting 30s for workers to write trace files to disk...") + time.sleep(30) + print("[Info] Trace export wait time finished.") omni_llm.close() diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index b9784979b..bb496198a 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -614,11 +614,39 @@ class AsyncOmni(OmniBase): """Generate outputs for a request from a pooling model.""" raise NotImplementedError("encode() is not implemented for AsyncOmni") - async def start_profile(self) -> None: - raise NotImplementedError("start_profile() is not implemented for AsyncOmni") + async def start_profile(self, stages: list[int] | None = None) -> None: + """Start profiling for specified stages. - async def stop_profile(self) -> None: - raise NotImplementedError("stop_profile() is not implemented for AsyncOmni") + Async wrapper around the base implementation for API consistency. + + Args: + stages: List of stage IDs to start profiling. If None, starts + profiling for all stages that have profiling enabled. + + Example: + >>> await async_omni.start_profile() + >>> async for output in async_omni.generate(...): + ... pass + >>> await async_omni.stop_profile() + """ + super().start_profile(stages) + + async def stop_profile(self, stages: list[int] | None = None) -> None: + """Stop profiling for specified stages. + + Async wrapper around the base implementation for API consistency. + + Args: + stages: List of stage IDs to stop profiling. If None, stops + profiling for all stages. + + Example: + >>> await async_omni.start_profile() + >>> async for output in async_omni.generate(...): + ... pass + >>> await async_omni.stop_profile() + """ + super().stop_profile(stages) async def pause_generation( self, diff --git a/vllm_omni/entrypoints/async_omni_llm.py b/vllm_omni/entrypoints/async_omni_llm.py index a1433d4f3..567af0377 100644 --- a/vllm_omni/entrypoints/async_omni_llm.py +++ b/vllm_omni/entrypoints/async_omni_llm.py @@ -165,19 +165,21 @@ class AsyncOmniLLM(AsyncLLM): except RuntimeError: pass - if envs.VLLM_TORCH_PROFILER_DIR: + if envs.VLLM_TORCH_PROFILER_DIR and not envs.VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM: logger.info( - "Torch profiler enabled. AsyncLLM CPU traces will be collected under %s", # noqa: E501 + "Torch profiler enabled. AsyncOmniLLM CPU traces will be collected under %s", envs.VLLM_TORCH_PROFILER_DIR, ) - worker_name = f"{socket.gethostname()}_{os.getpid()}.async_llm" + worker_name = f"{socket.gethostname()}_{os.getpid()}.async_omni_llm" self.profiler = torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, ], with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK, on_trace_ready=torch.profiler.tensorboard_trace_handler( - envs.VLLM_TORCH_PROFILER_DIR, worker_name=worker_name, use_gzip=True + envs.VLLM_TORCH_PROFILER_DIR, + worker_name=worker_name, + use_gzip=envs.VLLM_TORCH_PROFILER_USE_GZIP, ), ) else: diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index d093db21c..168999ac9 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -30,7 +30,7 @@ from vllm_omni.distributed.ray_utils.utils import ( ) from vllm_omni.entrypoints.log_utils import OrchestratorMetrics from vllm_omni.entrypoints.omni_stage import OmniStage -from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK +from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load from vllm_omni.entrypoints.utils import ( get_final_stage_id_for_e2e, @@ -334,6 +334,72 @@ class OmniBase: logger.error(f"[{self._name}] Stage initialization failed. Troubleshooting Steps:\n{formatted_suggestions}") + def start_profile(self, stages: list[int] | None = None) -> None: + """Start profiling for specified stages. + + Sends start_profile command to stage workers. Profiling must be enabled + via VLLM_TORCH_PROFILER_DIR environment variable. + + Args: + stages: List of stage IDs to start profiling. If None, starts + profiling for all stages that have profiling enabled. + + Example: + >>> # Profile all stages + >>> omni.start_profile() + >>> outputs = omni.generate(prompts, sampling_params) + >>> omni.stop_profile() + + >>> # Profile only stage 0 and 2 + >>> omni.start_profile(stages=[0, 2]) + """ + if stages is None: + stages = list(range(len(self.stage_list))) + + for stage_id in stages: + if stage_id < len(self.stage_list): + try: + self.stage_list[stage_id].submit({"type": OmniStageTaskType.PROFILER_START}) + logger.info("[%s] Sent start_profile to stage-%s", self._name, stage_id) + except Exception as e: + logger.warning( + "[%s] Failed to send start_profile to stage-%s: %s", + self._name, + stage_id, + e, + ) + + def stop_profile(self, stages: list[int] | None = None) -> None: + """Stop profiling for specified stages. + + Sends stop_profile command to stage workers to finalize and save traces. + + Args: + stages: List of stage IDs to stop profiling. If None, stops + profiling for all stages. + + Example: + >>> omni.start_profile() + >>> outputs = omni.generate(prompts, sampling_params) + >>> omni.stop_profile() + >>> # Traces saved to VLLM_TORCH_PROFILER_DIR/stage_X_/ + """ + if stages is None: + stages = list(range(len(self.stage_list))) + + for stage_id in stages: + if stage_id < len(self.stage_list): + try: + self.stage_list[stage_id].submit({"type": OmniStageTaskType.PROFILER_STOP}) + logger.info("[%s] Sent stop_profile to stage-%s", self._name, stage_id) + except Exception as e: + logger.warning( + "[%s] Failed to send stop_profile to stage-%s: %s", + self._name, + stage_id, + e, + ) + def close(self) -> None: """Close all stage processes and clean up resources.""" if hasattr(self, "_weak_finalizer"): diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index ff6adf8e9..fa70e1bc2 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -41,6 +41,7 @@ from vllm_omni.entrypoints.stage_utils import ( SHUTDOWN_TASK, OmniStageTaskType, _to_dict, + is_profiler_task, maybe_dump_to_shm, set_stage_devices, ) @@ -606,6 +607,32 @@ def _stage_worker( pass logger.debug("Engine initialized") + # Check if stage engine supports profiling (via vLLM's built-in profiler) + has_profiler = hasattr(stage_engine, "start_profile") and hasattr(stage_engine, "stop_profile") + if has_profiler: + logger.info( + "[Stage-%s] vLLM profiler support detected (model_stage=%s)", + stage_id, + engine_args.get("model_stage", None), + ) + + def handle_profiler_task(task_type: OmniStageTaskType) -> None: + """Handle profiler task.""" + if task_type == OmniStageTaskType.PROFILER_START: + if has_profiler: + try: + stage_engine.start_profile() + logger.info("[Stage-%s] Profiler started via vLLM engine", stage_id) + except Exception as e: + logger.warning("[Stage-%s] Failed to start profiler: %s", stage_id, e) + elif task_type == OmniStageTaskType.PROFILER_STOP: + if has_profiler: + try: + stage_engine.stop_profile() + logger.info("[Stage-%s] Profiler stopped via vLLM engine", stage_id) + except Exception as e: + logger.warning("[Stage-%s] Failed to stop profiler: %s", stage_id, e) + # Initialize OmniConnectors if configured connectors = {} if connectors_config: @@ -636,6 +663,11 @@ def _stage_worker( logger.info("Received shutdown signal") break + # Handle profiler control commands + if is_profiler_task(task_type): + handle_profiler_task(task_type) + continue + batch_tasks: list[dict[str, Any]] = [task] start_time = _time.time() if max_batch_size > 1: @@ -645,6 +677,11 @@ def _stage_worker( if extra == SHUTDOWN_TASK: in_q.put(SHUTDOWN_TASK) break + # Handle profiler commands that arrive during batching + extra_type = extra.get("type") if isinstance(extra, dict) else None + if is_profiler_task(extra_type): + handle_profiler_task(extra_type) + continue # Don't add to batch_tasks batch_tasks.append(extra) end_time = _time.time() duration = end_time - start_time @@ -1076,6 +1113,33 @@ async def _stage_worker_async( if stage_type != "diffusion": await stage_engine.reset_mm_cache() logger.debug("[Stage-%s] Engine initialized", stage_id) + + # Check if stage engine supports profiling (via vLLM's built-in profiler) + has_profiler = hasattr(stage_engine, "start_profile") and hasattr(stage_engine, "stop_profile") + if has_profiler: + logger.info( + "[Stage-%s] vLLM profiler support detected for async stage (model_stage=%s)", + stage_id, + engine_args.get("model_stage", None), + ) + + async def handle_profiler_task_async(task_type: OmniStageTaskType) -> None: + """Handle profiler task asynchronously.""" + if task_type == OmniStageTaskType.PROFILER_START: + if has_profiler: + try: + await stage_engine.start_profile() + logger.info("[Stage-%s] Profiler started via vLLM engine", stage_id) + except Exception as e: + logger.warning("[Stage-%s] Failed to start profiler: %s", stage_id, e) + elif task_type == OmniStageTaskType.PROFILER_STOP: + if has_profiler: + try: + await stage_engine.stop_profile() + logger.info("[Stage-%s] Profiler stopped via vLLM engine", stage_id) + except Exception as e: + logger.warning("[Stage-%s] Failed to stop profiler: %s", stage_id, e) + # Signal readiness to orchestrator and send vllm_config back to main process try: # Send vllm_config back to main process so it can be accessed via @@ -1183,8 +1247,11 @@ async def _stage_worker_async( elif task_type == OmniStageTaskType.ABORT: rid = task["request_id"] asyncio.create_task(stage_engine.abort(rid)) + elif is_profiler_task(task_type): + await handle_profiler_task_async(task_type) else: asyncio.create_task(generation_single_request(task)) + except queue.Empty: await asyncio.sleep(0.001) batch_request_outputs: list[Any] = [] diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 986cdc0b1..0656652c4 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -16,11 +16,17 @@ class OmniStageTaskType(enum.Enum): GENERATE = "generate" ABORT = "abort" SHUTDOWN = "shutdown" + PROFILER_START = "profiler_start" + PROFILER_STOP = "profiler_stop" SHUTDOWN_TASK = {"type": OmniStageTaskType.SHUTDOWN} +def is_profiler_task(task_type: OmniStageTaskType) -> bool: + return task_type in (OmniStageTaskType.PROFILER_START, OmniStageTaskType.PROFILER_STOP) + + def set_stage_devices( stage_id: int, devices: str | int | None,