# 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"