SYLAR
a1c37cc6e3
Signed-off-by: lishunyang <lishunyang12@163.com> Signed-off-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
127 lines
3.9 KiB
Python
127 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import os
|
|
import subprocess
|
|
from contextlib import nullcontext
|
|
|
|
import torch
|
|
from torch.profiler import ProfilerActivity, profile
|
|
from vllm.logger import init_logger
|
|
|
|
from .base import ProfilerBase
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class TorchProfiler(ProfilerBase):
|
|
"""
|
|
Torch-based profiler configured for End-to-End continuous recording.
|
|
Uses 'on_trace_ready' to handle Trace export.
|
|
Compression is offloaded to a background subprocess to avoid blocking the worker loop.
|
|
"""
|
|
|
|
_profiler: profile | None = None
|
|
_trace_template: str = ""
|
|
|
|
@classmethod
|
|
def start(cls, trace_path_template: str) -> str:
|
|
"""
|
|
Start the profiler with the given trace path template.
|
|
"""
|
|
# 1. Cleanup any existing profiler
|
|
if cls._profiler is not None:
|
|
logger.warning("[Rank %s] Stopping existing Torch profiler", cls._get_rank())
|
|
cls._profiler.stop()
|
|
cls._profiler = None
|
|
|
|
rank = cls._get_rank()
|
|
|
|
# 2. Make path absolute
|
|
trace_path_template = os.path.abspath(trace_path_template)
|
|
cls._trace_template = trace_path_template
|
|
|
|
# Expected paths
|
|
json_file = f"{trace_path_template}_rank{rank}.json"
|
|
|
|
os.makedirs(os.path.dirname(json_file), exist_ok=True)
|
|
|
|
logger.info(f"[Rank {rank}] Starting End-to-End Torch profiler")
|
|
|
|
# 3. Define the on_trace_ready handler
|
|
def trace_handler(p):
|
|
nonlocal json_file
|
|
|
|
# A. Export JSON Trace
|
|
try:
|
|
p.export_chrome_trace(json_file)
|
|
logger.info(f"[Rank {rank}] Trace exported to {json_file}")
|
|
|
|
try:
|
|
subprocess.Popen(["gzip", "-f", json_file])
|
|
logger.info(f"[Rank {rank}] Triggered background compression for {json_file}")
|
|
# Update variable to point to the eventual file
|
|
json_file = f"{json_file}.gz"
|
|
except Exception as compress_err:
|
|
logger.warning(f"[Rank {rank}] Background gzip failed to start: {compress_err}")
|
|
|
|
except Exception as e:
|
|
logger.warning(f"[Rank {rank}] Failed to export trace: {e}")
|
|
|
|
# 4. Initialize profiler with long active period
|
|
cls._profiler = profile(
|
|
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
|
|
schedule=torch.profiler.schedule(
|
|
wait=0,
|
|
warmup=0,
|
|
active=100000, # long capture window
|
|
),
|
|
on_trace_ready=trace_handler,
|
|
record_shapes=True,
|
|
profile_memory=True,
|
|
with_stack=True,
|
|
with_flops=True,
|
|
)
|
|
|
|
# 5. Start profiling
|
|
cls._profiler.start()
|
|
|
|
# Return the expected final path
|
|
return f"{trace_path_template}_rank{rank}.json.gz"
|
|
|
|
@classmethod
|
|
def stop(cls) -> dict | None:
|
|
if cls._profiler is None:
|
|
return None
|
|
|
|
rank = cls._get_rank()
|
|
|
|
# Determine expected paths
|
|
base_path = f"{cls._trace_template}_rank{rank}"
|
|
gz_path = f"{base_path}.json.gz"
|
|
|
|
try:
|
|
# This triggers trace_handler synchronously
|
|
# Since we removed table generation and backgrounded compression, this returns fast.
|
|
cls._profiler.stop()
|
|
except Exception as e:
|
|
logger.warning(f"[Rank {rank}] Profiler stop failed: {e}")
|
|
|
|
cls._profiler = None
|
|
|
|
# We return the .gz path assuming background compression will succeed.
|
|
return {"trace": gz_path, "table": None}
|
|
|
|
@classmethod
|
|
def step(cls):
|
|
if cls._profiler is not None:
|
|
cls._profiler.step()
|
|
|
|
@classmethod
|
|
def is_active(cls) -> bool:
|
|
return cls._profiler is not None
|
|
|
|
@classmethod
|
|
def get_step_context(cls):
|
|
return nullcontext()
|