# Copyright 2025 - Oumi
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import functools
import pathlib
from contextlib import contextmanager
from typing import Optional
import torch
from oumi.core.configs.params.profiler_params import ProfilerParams
from oumi.core.distributed import DeviceRankInfo, get_device_rank_info
from oumi.utils.logging import logger
_PROFILER_LOG_PREFIX = "PROF:"
_PROFILER_DEFAULT_SUB_DIR = "profiler"
def _configure_torch_profile_save_dir(
params: ProfilerParams, training_output_dir: Optional[str]
) -> ProfilerParams:
"""Auto-generates ProfilerParams.saved_dir if not specified explicitly."""
if not params.save_dir and training_output_dir:
params.save_dir = str(
pathlib.Path(training_output_dir) / _PROFILER_DEFAULT_SUB_DIR
)
return params
def _on_trace_ready(
prof,
*,
out_prefix: str,
enable_cpu_profiling: bool,
enable_cuda_profiling: bool,
params: ProfilerParams,
save_dir_path: Optional[pathlib.Path],
) -> None:
logger.info(f"{_PROFILER_LOG_PREFIX} on_trace_ready(out_prefix={out_prefix})")
sort_by = []
if enable_cpu_profiling:
sort_by.extend(
[
"cpu_time_total",
"self_cpu_time_total",
]
)
if params.profile_memory:
sort_by.extend(
[
"cpu_memory_usage",
"self_cpu_memory_usage",
]
)
if enable_cuda_profiling:
sort_by.extend(
[
"cuda_time_total",
"self_cuda_time_total",
]
)
if params.profile_memory:
sort_by.extend(
[
"cuda_memory_usage",
"self_cuda_memory_usage",
]
)
# if `params.record_shapes` is True, then also generate reports with breakdowns
# by tensor shapes. Otherwise (the default), only produce profiling reports
# without shape breakdowns (less verbose).
for group_by_input_shape in [False] + ([True] if params.record_shapes else []):
group_by_shape_tag = "_by_shape" if group_by_input_shape else ""
prof_avgs = prof.key_averages(group_by_input_shape=group_by_input_shape)
for sort_key in sort_by:
prof_table = prof_avgs.table(sort_by=sort_key, row_limit=params.row_limit)
logger.info(
f"{_PROFILER_LOG_PREFIX} {sort_key} "
f"[group_by_shape={group_by_input_shape}]"
f"\n{prof_table}\n"
)
if save_dir_path:
file_path: pathlib.Path = (
save_dir_path / f"{out_prefix}_{sort_key}{group_by_shape_tag}.txt"
)
with file_path.open("w") as f:
f.write(prof_table)
if save_dir_path:
# Save gzip-ed trace (JSON traces compress extremely well).
# Open traces on https://ui.perfetto.dev/ or `chrome://tracing`
file_name: pathlib.Path = save_dir_path / f"{out_prefix}_pt_trace.json.gz"
logger.info(f"Exporting profiler Chrome trace to {file_name} ...")
prof.export_chrome_trace(str(file_name))
[docs]
@contextmanager
def torch_profile(
params: ProfilerParams,
*,
training_output_dir: Optional[str],
record_function_name: str = "oumi.train",
):
"""Creates PyTorch Profiler context manager.
Args:
params: Profiler config.
training_output_dir: If `ProfilerParams.save_dir` is not specified, then
a "profiler" sub-directory will be created under `training_output_dir`,
and used to save profiler traces.
record_function_name: The name to use with `torch.profiler.record_function()`
for top-level `train()` operation.
Yields:
torch.profiler.profile or None: The newly-created Profiler object if profiling
is enabled, or `None` otherwise.
Example:
To profile a training loop::
with torch_profile(params, record_function_name="oumi.train") as prof:
for i in range(n):
training_step()
if prof is not None:
prof.step()
"""
params = _configure_torch_profile_save_dir(params, training_output_dir)
device_rank_info: DeviceRankInfo = get_device_rank_info()
out_prefix: str = (
f"prof_{device_rank_info.rank:03}_local_{device_rank_info.local_rank:02}"
)
profile_activities = []
enable_cpu_profiling = params.enable_cpu_profiling
if enable_cpu_profiling:
profile_activities.append(torch.profiler.ProfilerActivity.CPU)
enable_cuda_profiling = False
if params.enable_cuda_profiling:
enable_cuda_profiling = torch.cuda.is_available()
if enable_cuda_profiling:
profile_activities.append(torch.profiler.ProfilerActivity.CUDA)
else:
logger.warning(
f"{_PROFILER_LOG_PREFIX} CUDA profiling is requested "
"while CUDA is not available!"
)
if not profile_activities:
# Nothing to profile. Return noop/null context.
logger.info(f"{_PROFILER_LOG_PREFIX} Torch Profiler disabled!")
yield None
return
logger.info(f"{_PROFILER_LOG_PREFIX} Starting profiling...")
logger.info(f"{_PROFILER_LOG_PREFIX} Save dir: {params.save_dir}")
save_dir_path: Optional[pathlib.Path] = (
pathlib.Path(params.save_dir) if params.save_dir else None
)
if save_dir_path:
save_dir_path.mkdir(parents=True, exist_ok=True)
logger.info(f"{_PROFILER_LOG_PREFIX} Save dir: {save_dir_path}")
logger.info(f"{_PROFILER_LOG_PREFIX} Output prefix: {out_prefix}")
logger.info(f"{_PROFILER_LOG_PREFIX} Function: {record_function_name}")
logger.info(f"{_PROFILER_LOG_PREFIX} Params: {params}")
# See also torch.profiler.tensorboard_trace_handler
trace_handler = functools.partial(
_on_trace_ready,
out_prefix=out_prefix,
enable_cpu_profiling=enable_cpu_profiling,
enable_cuda_profiling=enable_cuda_profiling,
params=params,
save_dir_path=save_dir_path,
)
schedule = None
if params.schedule.enable_schedule:
schedule = torch.profiler.schedule(
skip_first=params.schedule.skip_first,
wait=params.schedule.wait,
warmup=params.schedule.warmup,
active=params.schedule.active,
repeat=params.schedule.repeat,
)
profiler = torch.profiler.profile(
activities=profile_activities,
on_trace_ready=trace_handler,
record_shapes=params.record_shapes,
profile_memory=params.profile_memory,
with_stack=params.with_stack,
with_flops=params.with_flops,
with_modules=params.with_modules,
schedule=schedule,
)
with profiler:
try:
with torch.profiler.record_function(record_function_name):
yield profiler
except Exception as e:
# The inner function raised an error
import traceback
logger.error(
f"{_PROFILER_LOG_PREFIX}"
+ "".join(traceback.format_exception(None, e, e.__traceback__))
)
raise
logger.info(f"{_PROFILER_LOG_PREFIX} Finished post-processing!")
return
