diff --git a/.gitignore b/.gitignore index 35dc7571e..a631aa05f 100644 --- a/.gitignore +++ b/.gitignore @@ -174,6 +174,7 @@ CLAUDE.md # Codex AGENTS.md +.codex .codex/ # cursor diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index 6c209e565..a7b508302 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -138,9 +138,27 @@ Single-stage diffusion serving with torch profiler: vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers \ --omni \ --port 8091 \ - --profiler-config '{"profiler": "torch", "torch_profiler_dir": "./vllm_profile"}' + --profiler-config '{ + "profiler": "torch", + "torch_profiler_dir": "/tmp/vllm_profile_wan22_i2v", + "torch_profiler_with_stack": true, + "torch_profiler_with_flops": false, + "torch_profiler_use_gzip": true, + "torch_profiler_dump_cuda_time_total": true, + "torch_profiler_record_shapes": true, + "torch_profiler_with_memory": true, + "delay_iterations": 0, + "max_iterations": 0 + }' ``` +Useful optional fields: + +- `torch_profiler_with_stack`: export `by_stack` operator views and stack text files. +- `torch_profiler_record_shapes`: export `by_shape` operator views. +- `torch_profiler_with_memory`: dump `memory_snapshot_rank*.pickle` when the backend supports memory history. +- `torch_profiler_use_gzip`: write the trace as `trace_rank*.json.gz`. + Single-stage diffusion serving with Nsight Systems: ```bash @@ -177,8 +195,11 @@ For mixed-stage pipelines, use explicit `stages` and pass the same stage list to Torch profiler output: -- Chrome/Perfetto traces under `torch_profiler_dir` -- Optional aggregated CUDA-time tables under the same directory +- Chrome/Perfetto trace: `trace_rank*.json` or `trace_rank*.json.gz` +- Excel workbook: `ops_rank*.xlsx` with `summary`, and optional `by_shape` / `by_stack` sheets +- Stack exports: `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt` when stack capture is enabled +- Memory snapshot: `memory_snapshot_rank*.pickle` when memory capture is enabled and supported by the backend +- Optional aggregated CUDA-time tables under the same session directory CUDA profiler / Nsight Systems output: diff --git a/tests/profile/test_omni_torch_profiler.py b/tests/profile/test_omni_torch_profiler.py new file mode 100644 index 000000000..3920078af --- /dev/null +++ b/tests/profile/test_omni_torch_profiler.py @@ -0,0 +1,582 @@ +# tests/test_omni_torch_profiler.py +from __future__ import annotations + +import gzip +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace + +import pytest +from openpyxl import load_workbook + +import vllm_omni.profiler.omni_torch_profiler as profiler_mod +from vllm_omni.profiler.omni_torch_profiler import OmniTorchProfilerWrapper + + +@pytest.fixture(autouse=True) +def patch_worker_profiler_init(monkeypatch): + def fake_init(self, profiler_config): + self.profiler_config = profiler_config + + monkeypatch.setattr( + profiler_mod.WorkerProfiler, + "__init__", + fake_init, + ) + + +@dataclass +class DummyProfilerConfig: + torch_profiler_dir: str + torch_profiler_use_gzip: bool = False + torch_profiler_record_shapes: bool = True + torch_profiler_with_memory: bool = True + torch_profiler_with_stack: bool = True + torch_profiler_with_flops: bool = False + torch_profiler_dump_cuda_time_total: bool = False + + +class FakeEvent: + def __init__( + self, + *, + name: str = "aten::mm", + count: int = 1, + input_shapes=None, + stack=None, + self_cpu_time_total: float = 10.0, + cpu_time_total: float = 12.0, + self_cuda_time_total: float = 20.0, + cuda_time_total: float = 25.0, + self_xpu_time_total: float = 0.0, + xpu_time_total: float = 0.0, + self_cpu_memory_usage: int = 128, + cpu_memory_usage: int = 256, + self_cuda_memory_usage: int = 1024, + cuda_memory_usage: int = 2048, + self_xpu_memory_usage: int = 0, + xpu_memory_usage: int = 0, + device_type: str = "CUDA", + node_id: int = 0, + overload_name: str = "", + is_async: bool = False, + is_legacy: bool = False, + ): + self.key = name + self.name = name + self.count = count + self.input_shapes = input_shapes if input_shapes is not None else [[2, 2], [2, 2]] + self.stack = stack if stack is not None else ["frame_a", "frame_b"] + self.self_cpu_time_total = self_cpu_time_total + self.cpu_time_total = cpu_time_total + self.self_cuda_time_total = self_cuda_time_total + self.cuda_time_total = cuda_time_total + self.self_xpu_time_total = self_xpu_time_total + self.xpu_time_total = xpu_time_total + self.self_cpu_memory_usage = self_cpu_memory_usage + self.cpu_memory_usage = cpu_memory_usage + self.self_cuda_memory_usage = self_cuda_memory_usage + self.cuda_memory_usage = cuda_memory_usage + self.self_xpu_memory_usage = self_xpu_memory_usage + self.xpu_memory_usage = xpu_memory_usage + self.device_type = device_type + self.node_id = node_id + self.overload_name = overload_name + self.is_async = is_async + self.is_legacy = is_legacy + + +class FakeEventList(list): + def table(self, sort_by=None, row_limit=-1): + return f"fake_table(sort_by={sort_by}, row_limit={row_limit}, len={len(self)})" + + +class FakeTorchProfiler: + def __init__(self, on_trace_ready=None): + self.started = False + self.stopped = False + self.on_trace_ready = on_trace_ready + self.exported_traces = [] + self.exported_stacks = [] + + def start(self): + self.started = True + + def stop(self): + self.stopped = True + if self.on_trace_ready is not None: + self.on_trace_ready(self) + + def export_chrome_trace(self, path): + Path(path).write_text('{"traceEvents": []}') + self.exported_traces.append(path) + + def export_stacks(self, path, metric): + Path(path).write_text(f"metric={metric}\nstack_line_1\nstack_line_2\n") + self.exported_stacks.append((path, metric)) + + def key_averages(self, group_by_input_shape=False, group_by_stack_n=0): + if group_by_input_shape: + return FakeEventList( + [ + FakeEvent( + name="aten::bmm", + input_shapes=[[4, 8, 16], [4, 16, 32]], + ) + ] + ) + if group_by_stack_n: + return FakeEventList( + [ + FakeEvent( + name="aten::all_reduce", + stack=["python_a", "python_b", "python_c"], + ) + ] + ) + return FakeEventList( + [ + FakeEvent(name="aten::mm"), + FakeEvent(name="nccl:all_reduce"), + ] + ) + + +@pytest.fixture +def fake_config(tmp_path): + return DummyProfilerConfig(torch_profiler_dir=str(tmp_path)) + + +@pytest.fixture +def fake_profiler_factory(monkeypatch): + created = {} + + def fake_profile(*args, **kwargs): + profiler = FakeTorchProfiler(on_trace_ready=kwargs.get("on_trace_ready")) + created["profiler"] = profiler + created["args"] = args + created["kwargs"] = kwargs + return profiler + + monkeypatch.setattr(profiler_mod.torch.profiler, "profile", fake_profile) + return created + + +@pytest.fixture +def wrapper(fake_config, fake_profiler_factory): + return OmniTorchProfilerWrapper( + profiler_config=fake_config, + worker_name="worker0", + local_rank=0, + activities=["CPU", "CUDA"], + ) + + +def test_set_trace_filename_creates_timestamped_session_dir(wrapper, monkeypatch, tmp_path): + class FixedDatetime: + @classmethod + def now(cls): + class _Now: + def strftime(self, fmt): + return "20260403-034200" + + return _Now() + + monkeypatch.setattr(profiler_mod, "datetime", FixedDatetime) + + wrapper.set_trace_filename("stage_0_llm_1234567890") + + session_dir = Path(wrapper._session_dir) + assert session_dir.exists() + assert session_dir.parent == tmp_path + assert session_dir.name == "20260403-034200_stage_0_llm_1234567890" + + +def test_set_trace_filename_with_full_path_creates_timestamped_leaf(wrapper, monkeypatch, tmp_path): + class FixedDatetime: + @classmethod + def now(cls): + class _Now: + def strftime(self, fmt): + return "20260403-111111" + + return _Now() + + monkeypatch.setattr(profiler_mod, "datetime", FixedDatetime) + + target = tmp_path / "nested" / "stage_x" + wrapper.set_trace_filename(str(target)) + + session_dir = Path(wrapper._session_dir) + assert session_dir.exists() + assert session_dir.parent == target.parent + assert session_dir.name == "20260403-111111_stage_x" + + +def test_on_trace_ready_exports_trace_json(wrapper): + wrapper.set_trace_filename("case_trace") + + wrapper._on_trace_ready(wrapper.profiler) + + trace_path = Path(wrapper._trace_path) + assert trace_path.exists() + assert trace_path.name == "trace_rank0.json" + assert trace_path.read_text() == '{"traceEvents": []}' + + +def test_on_trace_ready_exports_gzip_trace(fake_config, fake_profiler_factory, monkeypatch): + fake_config.torch_profiler_use_gzip = True + + wrapper = OmniTorchProfilerWrapper( + profiler_config=fake_config, + worker_name="worker0", + local_rank=0, + activities=["CPU", "CUDA"], + ) + wrapper.set_trace_filename("case_gzip") + + def fake_popen(cmd): + assert cmd[:2] == ["gzip", "-f"] + src = Path(cmd[2]) + gz_path = src.with_suffix(src.suffix + ".gz") + gz_path.write_bytes(gzip.compress(src.read_bytes())) + src.unlink() + + class DummyProc: + pass + + return DummyProc() + + monkeypatch.setattr(profiler_mod.subprocess, "Popen", fake_popen) + + wrapper._on_trace_ready(wrapper.profiler) + + assert wrapper._trace_path.endswith(".json.gz") + gz_path = Path(wrapper._trace_path) + assert gz_path.exists() + assert gzip.decompress(gz_path.read_bytes()) == b'{"traceEvents": []}' + + +def test_start_enables_memory_history(wrapper, monkeypatch): + calls = [] + + monkeypatch.setattr(profiler_mod.torch.cuda, "is_available", lambda: True) + + def fake_record_memory_history(*args, **kwargs): + calls.append((args, kwargs)) + + monkeypatch.setattr( + profiler_mod.torch.cuda.memory, + "_record_memory_history", + fake_record_memory_history, + ) + + wrapper.set_trace_filename("case_memory_start") + wrapper._start() + + assert wrapper.profiler.started is True + assert wrapper._memory_history_enabled is True + assert len(calls) == 1 + assert calls[0][1]["enabled"] == "all" + assert calls[0][1]["context"] == "all" + assert calls[0][1]["stacks"] == "python" + assert calls[0][1]["max_entries"] == 100000 + assert calls[0][1]["clear_history"] is True + + +def test_start_skips_memory_history_when_memory_disabled(fake_config, fake_profiler_factory, monkeypatch): + fake_config.torch_profiler_with_memory = False + + wrapper = OmniTorchProfilerWrapper( + profiler_config=fake_config, + worker_name="worker0", + local_rank=0, + activities=["CPU", "CUDA"], + ) + + called = {"n": 0} + + monkeypatch.setattr(profiler_mod.torch.cuda, "is_available", lambda: True) + + def fake_record_memory_history(*args, **kwargs): + called["n"] += 1 + + monkeypatch.setattr( + profiler_mod.torch.cuda.memory, + "_record_memory_history", + fake_record_memory_history, + ) + + wrapper.set_trace_filename("case_skip_memory") + wrapper._start() + + assert called["n"] == 0 + assert wrapper._memory_history_enabled is False + + +def test_try_dump_memory_snapshot_writes_pickle(wrapper, monkeypatch): + wrapper.set_trace_filename("case_snapshot") + wrapper._memory_history_enabled = True + wrapper._memory_history_backend = "CUDA" + wrapper._memory_history_module = profiler_mod.torch.cuda.memory + + disable_calls = [] + + def fake_record_memory_history(*args, **kwargs): + disable_calls.append((args, kwargs)) + + def fake_dump_snapshot(path): + Path(path).write_bytes(b"fake pickle bytes") + + monkeypatch.setattr( + profiler_mod.torch.cuda.memory, + "_record_memory_history", + fake_record_memory_history, + ) + monkeypatch.setattr( + profiler_mod.torch.cuda.memory, + "_dump_snapshot", + fake_dump_snapshot, + ) + + wrapper._try_dump_memory_snapshot() + + snapshot = Path(wrapper._artifact_paths["memory_snapshot"]) + assert snapshot.exists() + assert snapshot.name == "memory_snapshot_rank0.pickle" + assert snapshot.read_bytes() == b"fake pickle bytes" + assert wrapper._memory_history_enabled is False + + assert disable_calls[-1][1]["enabled"] is None + + +def test_stop_always_dumps_memory_snapshot_on_success_path(wrapper, monkeypatch): + wrapper.set_trace_filename("case_stop") + + record_calls = [] + dump_calls = [] + + monkeypatch.setattr(profiler_mod.torch.cuda, "is_available", lambda: True) + + def fake_record_memory_history(*args, **kwargs): + record_calls.append((args, kwargs)) + + def fake_dump_snapshot(path): + dump_calls.append(path) + Path(path).write_bytes(b"snapshot-bytes") + + monkeypatch.setattr( + profiler_mod.torch.cuda.memory, + "_record_memory_history", + fake_record_memory_history, + ) + monkeypatch.setattr( + profiler_mod.torch.cuda.memory, + "_dump_snapshot", + fake_dump_snapshot, + ) + + wrapper._start() + wrapper._stop() + + session_dir = Path(wrapper._session_dir) + + assert wrapper.profiler.started is True + assert wrapper.profiler.stopped is True + assert (session_dir / "memory_snapshot_rank0.pickle").exists() + assert len(dump_calls) == 1 + assert record_calls[0][1]["enabled"] == "all" + assert record_calls[-1][1]["enabled"] is None + + +def test_on_stop_hook_generates_stack_and_excel_artifacts(wrapper): + wrapper.set_trace_filename("case_artifacts") + wrapper._on_stop_hook() + + session_dir = Path(wrapper._session_dir) + + assert not (session_dir / "ops_summary_rank0.txt").exists() + assert not (session_dir / "ops_by_shape_rank0.txt").exists() + assert not (session_dir / "ops_by_stack_rank0.txt").exists() + assert (session_dir / "stacks_cpu_rank0.txt").exists() + assert (session_dir / "stacks_cuda_rank0.txt").exists() + assert (session_dir / "ops_rank0.xlsx").exists() + + +def test_excel_contains_expected_sheets(wrapper): + wrapper.set_trace_filename("case_excel") + wrapper._on_stop_hook() + + xlsx_path = Path(wrapper._session_dir) / "ops_rank0.xlsx" + wb = load_workbook(xlsx_path) + + assert "summary" in wb.sheetnames + assert "by_shape" in wb.sheetnames + assert "by_stack" in wb.sheetnames + + +def test_excel_summary_has_expected_columns(wrapper): + wrapper.set_trace_filename("case_excel_columns") + wrapper._on_stop_hook() + + xlsx_path = Path(wrapper._session_dir) / "ops_rank0.xlsx" + wb = load_workbook(xlsx_path) + ws = wb["summary"] + + headers = [cell.value for cell in next(ws.iter_rows(min_row=1, max_row=1))] + assert "name" in headers + assert "count" in headers + assert "self_cpu_time_total_us" in headers + assert "self_cuda_time_total_us" in headers + assert "self_cpu_memory_usage_bytes" in headers + assert "self_cuda_memory_usage_bytes" in headers + assert "input_shapes" in headers + assert "stack" in headers + + +def test_get_results_returns_all_artifact_paths(wrapper, monkeypatch): + wrapper.set_trace_filename("case_results") + + monkeypatch.setattr(profiler_mod.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr( + profiler_mod.torch.cuda.memory, + "_record_memory_history", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + profiler_mod.torch.cuda.memory, + "_dump_snapshot", + lambda path: Path(path).write_bytes(b"snapshot"), + ) + + wrapper._start() + wrapper._stop() + + results = wrapper.get_results() + + assert "trace" in results + assert "table" in results + assert "session_dir" in results + assert "ops" in results + assert "memory_snapshot" in results + assert Path(results["session_dir"]).exists() + assert Path(results["ops"]).exists() + assert Path(results["table"]).exists() + assert Path(results["table"]).name == "ops_rank0.xlsx" + assert Path(results["memory_snapshot"]).exists() + + +def test_start_uses_xpu_memory_history_when_available(wrapper, monkeypatch): + calls = [] + + def fake_record_memory_history(*args, **kwargs): + calls.append((args, kwargs)) + + fake_memory_module = SimpleNamespace( + _record_memory_history=fake_record_memory_history, + ) + monkeypatch.setattr( + wrapper, + "_resolve_memory_history_backend", + lambda: ("XPU", fake_memory_module), + ) + + wrapper.set_trace_filename("case_xpu_memory_start") + wrapper._start() + + assert wrapper._memory_history_enabled is True + assert wrapper._memory_history_backend == "XPU" + assert wrapper._memory_history_module is fake_memory_module + assert calls[0][1]["enabled"] == "all" + + +def test_start_uses_npu_memory_history_when_available(wrapper, monkeypatch): + calls = [] + + def fake_record_memory_history(*args, **kwargs): + calls.append((args, kwargs)) + + fake_memory_module = SimpleNamespace( + _record_memory_history=fake_record_memory_history, + ) + monkeypatch.setattr( + wrapper, + "_resolve_memory_history_backend", + lambda: ("NPU", fake_memory_module), + ) + + wrapper.set_trace_filename("case_npu_memory_start") + wrapper._start() + + assert wrapper._memory_history_enabled is True + assert wrapper._memory_history_backend == "NPU" + assert wrapper._memory_history_module is fake_memory_module + assert calls[0][1]["enabled"] == "all" + + +def test_start_skips_memory_history_when_backend_api_missing(wrapper, monkeypatch): + fake_memory_module = SimpleNamespace() + monkeypatch.setattr( + wrapper, + "_resolve_memory_history_backend", + lambda: ("XPU", fake_memory_module), + ) + + wrapper.set_trace_filename("case_missing_memory_api") + wrapper._start() + + assert wrapper._memory_history_enabled is False + assert wrapper._memory_history_backend is None + assert wrapper._memory_history_module is None + + +def test_try_dump_memory_snapshot_uses_resolved_backend_module(wrapper): + wrapper.set_trace_filename("case_xpu_snapshot") + wrapper._memory_history_enabled = True + wrapper._memory_history_backend = "XPU" + + calls = [] + + def fake_record_memory_history(*args, **kwargs): + calls.append((args, kwargs)) + + def fake_dump_snapshot(path): + Path(path).write_bytes(b"xpu snapshot bytes") + + wrapper._memory_history_module = SimpleNamespace( + _record_memory_history=fake_record_memory_history, + _dump_snapshot=fake_dump_snapshot, + ) + + wrapper._try_dump_memory_snapshot() + + snapshot = Path(wrapper._artifact_paths["memory_snapshot"]) + assert snapshot.exists() + assert snapshot.read_bytes() == b"xpu snapshot bytes" + assert calls[-1][1]["enabled"] is None + assert wrapper._memory_history_enabled is False + assert wrapper._memory_history_backend is None + assert wrapper._memory_history_module is None + + +def test_event_list_to_rows_contains_expected_fields(wrapper): + rows = wrapper._event_list_to_rows( + [ + FakeEvent( + name="aten::linear", + input_shapes=[[8, 16], [16, 32]], + stack=["f1", "f2"], + ) + ] + ) + + assert len(rows) == 1 + row = rows[0] + assert row["name"] == "aten::linear" + assert row["count"] == 1 + assert row["self_cpu_time_total_us"] == 10.0 + assert row["self_cuda_time_total_us"] == 20.0 + assert row["self_cpu_memory_usage_bytes"] == 128 + assert row["self_cuda_memory_usage_bytes"] == 1024 + assert "[[8, 16], [16, 32]]" == row["input_shapes"] + assert row["stack"] == "f1\nf2" diff --git a/vllm_omni/profiler/omni_torch_profiler.py b/vllm_omni/profiler/omni_torch_profiler.py index 2257a2128..023ad5009 100644 --- a/vllm_omni/profiler/omni_torch_profiler.py +++ b/vllm_omni/profiler/omni_torch_profiler.py @@ -5,8 +5,8 @@ from __future__ import annotations import os import subprocess -import time -from typing import Literal +from datetime import datetime +from typing import Any, Literal import torch from typing_extensions import override @@ -62,6 +62,13 @@ class OmniTorchProfilerWrapper(WorkerProfiler): self._trace_path: str | None = None self._table_path: str | None = None + self._activities = activities + self._session_dir: str | None = None + self._artifact_paths: dict[str, str | None] = {} + self._memory_history_enabled = False + self._memory_history_backend: str | None = None + self._memory_history_module = None + if local_rank in (None, 0): logger.info_once( "Omni torch profiling enabled. Traces will be saved to: %s", @@ -72,6 +79,9 @@ class OmniTorchProfilerWrapper(WorkerProfiler): self.dump_cpu_time_total = "CPU" in activities and len(activities) == 1 self.profiler = self._create_profiler(profiler_config, activities) + def _rank(self) -> int: + return 0 if self.local_rank is None else self.local_rank + def _get_default_activities(self) -> list[TorchProfilerActivity]: """Get default activities for this platform. @@ -106,19 +116,58 @@ class OmniTorchProfilerWrapper(WorkerProfiler): Can also be a full path (e.g. from diffusion engine). """ self._trace_filename = filename + self._session_dir = None + self._ensure_session_dir() + + def _ensure_session_dir(self) -> str: + """Create one timestamped directory for this profiling run.""" + if self._session_dir is not None: + os.makedirs(self._session_dir, exist_ok=True) + return self._session_dir + + ts = datetime.now().strftime("%Y%m%d-%H%M%S") + base_name = self._trace_filename or self._worker_name + + if os.path.dirname(base_name): + parent_dir = os.path.dirname(base_name) + leaf_name = os.path.basename(base_name) + session_name = f"{ts}_{leaf_name}" + self._session_dir = os.path.join(parent_dir, session_name) + else: + session_name = f"{ts}_{base_name}" + self._session_dir = os.path.join(self._trace_dir, session_name) + + os.makedirs(self._session_dir, exist_ok=True) + self._artifact_paths["session_dir"] = self._session_dir + return self._session_dir + + def _artifact_path(self, stem: str, suffix: str) -> str: + """Build artifact path under the session directory.""" + return os.path.join( + self._ensure_session_dir(), + f"{stem}_rank{self._rank()}{suffix}", + ) + + def _write_text_artifact(self, name: str, content: str) -> str: + path = self._artifact_path(name, ".txt") + with open(path, "w") as f: + f.write(content) + self._artifact_paths[name] = path + return path + + def _has_cuda_like_activity(self) -> bool: + return any(a in self._activities for a in ("CUDA", "MUSA")) + + def _get_time_sort_key(self) -> str: + if self._has_cuda_like_activity(): + return "self_cuda_time_total" + return "self_cpu_time_total" def _on_trace_ready(self, prof) -> None: """Custom trace handler: export chrome trace with omni naming.""" - rank = self.local_rank - filename = self._trace_filename or f"{self._worker_name}_{int(time.time())}" - # If filename already contains a directory, use as-is (e.g. from - # diffusion engine which builds full path). Otherwise join with trace_dir. - if os.path.dirname(filename): - json_file = f"{filename}_rank{rank}.json" - else: - json_file = os.path.join(self._trace_dir, f"{filename}_rank{rank}.json") + rank = self._rank() - os.makedirs(os.path.dirname(json_file), exist_ok=True) + json_file = self._artifact_path("trace", ".json") try: prof.export_chrome_trace(json_file) @@ -143,18 +192,211 @@ class OmniTorchProfilerWrapper(WorkerProfiler): else: self._trace_path = json_file + self._artifact_paths["trace"] = self._trace_path + except Exception as e: logger.warning("[Rank %s] Failed to export trace: %s", rank, e) + def _try_enable_memory_history(self) -> None: + """Enable backend-specific memory history for snapshot analysis.""" + if not self.profiler_config.torch_profiler_with_memory: + return + + backend_name, memory_module = self._resolve_memory_history_backend() + if backend_name is None or memory_module is None: + return + + record_memory_history = getattr(memory_module, "_record_memory_history", None) + if record_memory_history is None: + logger.info( + "[Rank %s] %s memory history is not supported on this platform", + self._rank(), + backend_name, + ) + return + + try: + record_memory_history( + enabled="all", + context="all", + stacks="python", + max_entries=100000, + clear_history=True, + ) + self._memory_history_enabled = True + self._memory_history_backend = backend_name + self._memory_history_module = memory_module + logger.info("[Rank %s] %s memory history enabled", self._rank(), backend_name) + except Exception as e: + logger.warning( + "[Rank %s] Failed to enable %s memory history: %s", + self._rank(), + backend_name, + e, + ) + + def _try_dump_memory_snapshot(self) -> None: + """Dump a backend-specific memory snapshot into the session directory.""" + if not self._memory_history_enabled: + return + + try: + if self._memory_history_module is None or self._memory_history_backend is None: + return + + dump_snapshot = getattr(self._memory_history_module, "_dump_snapshot", None) + if dump_snapshot is None: + logger.info( + "[Rank %s] %s memory snapshot is not supported on this platform", + self._rank(), + self._memory_history_backend, + ) + return + + snapshot_file = self._artifact_path("memory_snapshot", ".pickle") + dump_snapshot(snapshot_file) + self._artifact_paths["memory_snapshot"] = snapshot_file + logger.info( + "[Rank %s] %s memory snapshot dumped to %s", + self._rank(), + self._memory_history_backend, + snapshot_file, + ) + except Exception as e: + logger.warning( + "[Rank %s] Failed to dump %s memory snapshot: %s", + self._rank(), + self._memory_history_backend, + e, + ) + finally: + try: + if self._memory_history_module is not None: + disable_memory_history = getattr( + self._memory_history_module, + "_record_memory_history", + None, + ) + if disable_memory_history is not None: + disable_memory_history(enabled=None) + except Exception: + pass + self._memory_history_enabled = False + self._memory_history_backend = None + self._memory_history_module = None + + def _resolve_memory_history_backend(self) -> tuple[str | None, Any]: + """Resolve the memory backend that supports history/snapshot APIs.""" + backend_specs = [ + ("CUDA", self._has_cuda_like_activity(), getattr(torch, "cuda", None)), + ("NPU", "NPU" in self._activities, getattr(torch, "npu", None)), + ("XPU", "XPU" in self._activities, getattr(torch, "xpu", None)), + ("MUSA", "MUSA" in self._activities, getattr(torch, "musa", None)), + ] + + for backend_name, enabled, device_module in backend_specs: + if not enabled or device_module is None: + continue + + is_available = getattr(device_module, "is_available", None) + if callable(is_available) and not is_available(): + continue + + memory_module = getattr(device_module, "memory", None) + if memory_module is not None: + return backend_name, memory_module + + return None, None + + def _safe_get(self, obj, name: str, default=None): + return getattr(obj, name, default) + + def _event_list_to_rows(self, event_list) -> list[dict]: + rows = [] + for evt in event_list: + row = { + "name": self._safe_get(evt, "key", None) or self._safe_get(evt, "name", None), + "count": self._safe_get(evt, "count", None), + "device_type": self._safe_get(evt, "device_type", None), + "node_id": self._safe_get(evt, "node_id", None), + "self_cpu_time_total_us": self._safe_get(evt, "self_cpu_time_total", None), + "cpu_time_total_us": self._safe_get(evt, "cpu_time_total", None), + "self_cuda_time_total_us": self._safe_get(evt, "self_cuda_time_total", None), + "cuda_time_total_us": self._safe_get(evt, "cuda_time_total", None), + "self_xpu_time_total_us": self._safe_get(evt, "self_xpu_time_total", None), + "xpu_time_total_us": self._safe_get(evt, "xpu_time_total", None), + "self_cpu_memory_usage_bytes": self._safe_get(evt, "self_cpu_memory_usage", None), + "cpu_memory_usage_bytes": self._safe_get(evt, "cpu_memory_usage", None), + "self_cuda_memory_usage_bytes": self._safe_get(evt, "self_cuda_memory_usage", None), + "cuda_memory_usage_bytes": self._safe_get(evt, "cuda_memory_usage", None), + "self_xpu_memory_usage_bytes": self._safe_get(evt, "self_xpu_memory_usage", None), + "xpu_memory_usage_bytes": self._safe_get(evt, "xpu_memory_usage", None), + "input_shapes": str(self._safe_get(evt, "input_shapes", None)), + "stack": "\n".join(self._safe_get(evt, "stack", []) or []), + "overload_name": self._safe_get(evt, "overload_name", None), + "is_async": self._safe_get(evt, "is_async", None), + "is_legacy": self._safe_get(evt, "is_legacy", None), + } + rows.append(row) + return rows + + def _write_excel_artifact(self, name: str, sheets: dict[str, list[dict]]) -> str: + path = self._artifact_path(name, ".xlsx") + + try: + import pandas as pd + except Exception as e: + logger.warning( + "[Rank %s] pandas not available, skip Excel export: %s", + self._rank(), + e, + ) + return path + + with pd.ExcelWriter(path, engine="openpyxl") as writer: + for sheet_name, rows in sheets.items(): + df = pd.DataFrame(rows) + + safe_sheet_name = sheet_name if sheet_name else "Sheet1" + + df.to_excel( + writer, + sheet_name=safe_sheet_name, + index=False, + freeze_panes=(1, 0), + ) + + ws = writer.sheets[safe_sheet_name] + ws.auto_filter.ref = ws.dimensions + + for col_cells in ws.columns: + max_len = 0 + col_letter = col_cells[0].column_letter + for cell in col_cells[:200]: + try: + val = "" if cell.value is None else str(cell.value) + max_len = max(max_len, len(val)) + except Exception: + pass + ws.column_dimensions[col_letter].width = min(max(max_len + 2, 12), 80) + + self._artifact_paths[name] = path + return path + @override def _start(self) -> None: + self._ensure_session_dir() + self._try_enable_memory_history() self.profiler.start() @override def _stop(self) -> None: """Stop profiler, export trace via on_trace_ready, and dump table.""" self.profiler.stop() - self._on_stop_hook() + try: + self._on_stop_hook() + finally: + self._try_dump_memory_snapshot() def _on_stop_hook(self) -> None: """Hook called after profiler.stop(). @@ -163,6 +405,68 @@ class OmniTorchProfilerWrapper(WorkerProfiler): Base implementation handles CUDA time total dump. """ rank = self.local_rank + sort_key = self._get_time_sort_key() + + excel_sheets: dict[str, list[dict]] = {} + + # 1) Summary op table + summary_events = self.profiler.key_averages() + excel_sheets["summary"] = self._event_list_to_rows(summary_events) + + # 2) Shape-grouped op table + if self.profiler_config.torch_profiler_record_shapes: + try: + shape_events = self.profiler.key_averages( + group_by_input_shape=True, + ) + excel_sheets["by_shape"] = self._event_list_to_rows(shape_events) + except Exception as e: + logger.warning( + "[Rank %s] Failed to export shape-grouped op table: %s", + rank, + e, + ) + + # 3) Stack-grouped op table + if self.profiler_config.torch_profiler_with_stack: + try: + stack_events = self.profiler.key_averages( + group_by_stack_n=8, + ) + excel_sheets["by_stack"] = self._event_list_to_rows(stack_events) + except Exception as e: + logger.warning( + "[Rank %s] Failed to export stack-grouped op table: %s", + rank, + e, + ) + + # 4) Export stack files + try: + cpu_stack_file = self._artifact_path("stacks_cpu", ".txt") + self.profiler.export_stacks( + cpu_stack_file, + metric="self_cpu_time_total", + ) + self._artifact_paths["stacks_cpu"] = cpu_stack_file + except Exception as e: + logger.warning("[Rank %s] export_stacks(cpu) failed: %s", rank, e) + + if self._has_cuda_like_activity(): + try: + cuda_stack_file = self._artifact_path("stacks_cuda", ".txt") + self.profiler.export_stacks( + cuda_stack_file, + metric="self_cuda_time_total", + ) + self._artifact_paths["stacks_cuda"] = cuda_stack_file + except Exception as e: + logger.warning("[Rank %s] export_stacks(cuda) failed: %s", rank, e) + + try: + self._table_path = self._write_excel_artifact("ops", excel_sheets) + except Exception as e: + logger.warning("[Rank %s] Failed to export Excel workbook: %s", rank, e) if self.profiler_config.torch_profiler_dump_cuda_time_total: profiler_dir = self.profiler_config.torch_profiler_dir @@ -190,6 +494,7 @@ class OmniTorchProfilerWrapper(WorkerProfiler): return { "trace": self._trace_path, "table": self._table_path, + **self._artifact_paths, }