[Refactor] Unify torch profiler for omni and diffusion models (#2099)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Samit <285365963@qq.com>
Co-authored-by: lishunyang lishunyang12@163.com

vllm_omni/diffusion/diffusion_engine.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import os
 import threading
 import time
 from collections.abc import Iterable
@@ -309,124 +308,31 @@ class DiffusionEngine:
                     self.scheduler.pop_request_state(target_sched_req_id)
                     return output
 
-    def start_profile(self, trace_filename: str | None = None) -> None:
-        """
-        Start torch profiling on all diffusion workers.
-
-        Creates a directory (if needed) and sets up a base filename template
-        for per-rank profiler traces (typically saved as <template>_rank<N>.json).
+    def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
+        """Start or stop torch profiling on all diffusion workers.
 
         Args:
-            trace_filename: Optional base filename (without extension or rank suffix).
-                            If None, generates one using current timestamp.
+            is_start: True to start profiling, False to stop.
+            profile_prefix: Optional prefix for trace filename (vLLM compat).
+
+        Note:
+            Matches vLLM's worker.profile() signature for consistency.
+            Traces are saved automatically via on_trace_ready callback.
         """
-        if trace_filename is None:
-            trace_filename = f"stage_0_diffusion_{int(time.time())}_rank"
-
-        trace_dir = os.environ.get("VLLM_TORCH_PROFILER_DIR", "./profiles")
-
-        # Expand ~ and ~user, then make absolute (robust against cwd changes)
-        trace_dir = os.path.expanduser(trace_dir)
-        trace_dir = os.path.abspath(trace_dir)
-
-        try:
-            os.makedirs(trace_dir, exist_ok=True)
-        except OSError as exc:
-            logger.error(f"Failed to create profiler directory {trace_dir}: {exc}")
-            raise
-
-        # Build final template path (without rank or extension — torch.profiler appends those)
-        full_template = os.path.join(trace_dir, trace_filename)
-
-        expected_pattern = f"{full_template}*.json"
-        logger.info(f"Starting diffusion profiling → {expected_pattern}")
-
-        # Also log the absolute directory once (useful in multi-node or containers)
-        logger.debug(f"Profiler output directory: {trace_dir}")
-
-        # Propagate to all workers
-        try:
-            self.collective_rpc(method="start_profile", args=(full_template,))
-        except Exception as e:
-            logger.error("Failed to start profiling on workers", exc_info=True)
-            raise RuntimeError(f"Could not start profiler: {e}") from e
-
-    def stop_profile(self) -> dict:
-        """
-        Stop profiling on all workers and collect the final trace/table paths.
-
-        The worker (torch_profiler.py) now handles trace export, compression to .gz,
-        and deletion of the original .json file. This method only collects and
-        reports the paths returned by the workers.
-
-        Returns:
-            dict with keys:
-            - "traces": list of final trace file paths (usually .json.gz)
-            - "tables": list of table strings (one per rank)
-        """
-        logger.info("Stopping diffusion profiling and collecting results...")
-
-        try:
-            # Give worker enough time — export + compression + table can be slow
-            results = self.collective_rpc(method="stop_profile", timeout=600)
-        except Exception:
-            logger.error("Failed to stop profiling on workers", exc_info=True)
-            return {"traces": [], "tables": []}
-
-        output_files = {"traces": [], "tables": []}
-        successful_traces = 0
-
-        if not results:
-            logger.warning("No profiling results returned from any rank")
-            return output_files
-
-        for rank, res in enumerate(results):
-            if not isinstance(res, dict):
-                logger.warning(f"Rank {rank}: invalid result format (got {type(res)})")
-                continue
-
-            # 1. Trace file — should be .json.gz if compression succeeded
-            trace_path = res.get("trace")
-            if trace_path:
-                # We trust the worker — it created/compressed the file
-                logger.info(f"[Rank {rank}] Final trace: {trace_path}")
-                output_files["traces"].append(trace_path)
-                successful_traces += 1
-
-                # Optional: warn if path looks suspicious (e.g. still .json)
-                if not trace_path.endswith((".json.gz", ".json")):
-                    logger.warning(f"Rank {rank}: unusual trace path extension: {trace_path}")
-
-            # 2. Table file — plain text
-            table = res.get("table")
-            if table:
-                output_files["tables"].append(table)
-
-        # Final summary logging
-        num_ranks = len(results)
-        if successful_traces > 0:
-            final_paths_str = ", ".join(output_files["traces"][:3])
-            if len(output_files["traces"]) > 3:
-                final_paths_str += f" ... (+{len(output_files['traces']) - 3} more)"
-
-            logger.info(
-                f"Profiling stopped. Collected {successful_traces} trace file(s) "
-                f"from {num_ranks} rank(s). "
-                f"Final trace paths: {final_paths_str}"
-            )
-        elif output_files["traces"]:
-            logger.info(
-                f"Profiling stopped but no traces were successfully collected. "
-                f"Reported paths: {', '.join(output_files['traces'][:3])}"
-                f"{' ...' if len(output_files['traces']) > 3 else ''}"
-            )
+        if is_start:
+            if profile_prefix is None:
+                profile_prefix = f"diffusion_{int(time.time())}"
+            logger.info(f"Starting diffusion profiling with prefix: {profile_prefix}")
         else:
-            logger.info("Profiling stopped — no trace files were collected from any rank.")
+            logger.info("Stopping diffusion profiling...")
 
-        if output_files["tables"]:
-            logger.debug(f"Collected {len(output_files['tables'])} profiling table(s)")
-
-        return output_files
+        try:
+            self.collective_rpc(method="profile", args=(is_start, profile_prefix))
+        except Exception as e:
+            action = "start" if is_start else "stop"
+            logger.error(f"Failed to {action} profiling on workers", exc_info=True)
+            if is_start:
+                raise RuntimeError(f"Could not {action} profiler: {e}") from e
 
     def _dummy_run(self):
         """A dummy run to warm up the model."""