653 lines
23 KiB
Python
653 lines
23 KiB
Python
#!/usr/bin/env python3
|
|
"""Profile direct EasyCom Qwen2.5-Omni extraction on one physical GPU.
|
|
|
|
This is a diagnostic entry point for ``streaming_turn_model_directly.py``.
|
|
It accepts the same extraction arguments plus three profiler-only options:
|
|
|
|
``--gpu-id``
|
|
Physical host GPU ID. It defaults to 7 and is exposed to the extraction
|
|
process as ``cuda:0``.
|
|
``--profile-log-dir``
|
|
Directory for the console log, GPU telemetry CSV, and JSON summary.
|
|
``--gpu-sample-seconds``
|
|
``nvidia-smi`` sampling period in seconds.
|
|
|
|
The profiler synchronizes CUDA around individual model stages so their times
|
|
are attributable. That synchronization perturbs throughput; use this script
|
|
to find bottlenecks, not to produce the final runtime measurement.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import csv
|
|
import functools
|
|
import json
|
|
import os
|
|
import re
|
|
import resource
|
|
import statistics
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import types
|
|
from collections import defaultdict
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any, Callable, TextIO, TypeVar
|
|
|
|
|
|
def _pop_cli_option(name: str, default: str) -> str:
|
|
"""Remove one profiler-only option before the normal parser sees it."""
|
|
|
|
for index, argument in enumerate(sys.argv[1:], start=1):
|
|
if argument == name:
|
|
if index + 1 >= len(sys.argv):
|
|
raise SystemExit(f"{name} requires a value")
|
|
value = sys.argv[index + 1]
|
|
del sys.argv[index : index + 2]
|
|
return value
|
|
prefix = f"{name}="
|
|
if argument.startswith(prefix):
|
|
value = argument[len(prefix) :]
|
|
del sys.argv[index]
|
|
return value
|
|
return default
|
|
|
|
|
|
# CUDA visibility must be fixed before importing torch or the Qwen modules.
|
|
PHYSICAL_GPU_ID = int(_pop_cli_option("--gpu-id", "7"))
|
|
PROFILE_LOG_DIR = Path(
|
|
_pop_cli_option("--profile-log-dir", "logs/bottleneck_gpu7")
|
|
).expanduser()
|
|
GPU_SAMPLE_SECONDS = float(_pop_cli_option("--gpu-sample-seconds", "1"))
|
|
if PHYSICAL_GPU_ID < 0:
|
|
raise SystemExit("--gpu-id must be >= 0")
|
|
if GPU_SAMPLE_SECONDS <= 0:
|
|
raise SystemExit("--gpu-sample-seconds must be > 0")
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = str(PHYSICAL_GPU_ID)
|
|
|
|
import torch # noqa: E402 (must follow CUDA_VISIBLE_DEVICES)
|
|
|
|
import streaming_turn_model as core # noqa: E402
|
|
import streaming_turn_model_directly as direct # noqa: E402
|
|
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
class TeeStream:
|
|
"""Write stdout/stderr to both the terminal and a durable log."""
|
|
|
|
def __init__(self, terminal: TextIO, log_file: TextIO) -> None:
|
|
self.terminal = terminal
|
|
self.log_file = log_file
|
|
self.encoding = getattr(terminal, "encoding", "utf-8")
|
|
|
|
def write(self, data: str) -> int:
|
|
terminal_count = self.terminal.write(data)
|
|
self.log_file.write(data)
|
|
self.flush()
|
|
return terminal_count
|
|
|
|
def flush(self) -> None:
|
|
self.terminal.flush()
|
|
self.log_file.flush()
|
|
|
|
def isatty(self) -> bool:
|
|
return self.terminal.isatty()
|
|
|
|
|
|
def _percentile(values: list[float], fraction: float) -> float:
|
|
if not values:
|
|
return 0.0
|
|
ordered = sorted(values)
|
|
index = round((len(ordered) - 1) * fraction)
|
|
return float(ordered[index])
|
|
|
|
|
|
def _sync_cuda() -> None:
|
|
if torch.cuda.is_available() and torch.cuda.is_initialized():
|
|
torch.cuda.synchronize()
|
|
|
|
|
|
class StageProfiler:
|
|
"""Accumulate inclusive wall times for named extraction stages."""
|
|
|
|
def __init__(self) -> None:
|
|
self.durations: dict[str, list[float]] = defaultdict(list)
|
|
self.metadata: dict[str, Any] = {}
|
|
|
|
def measure(
|
|
self,
|
|
stage: str,
|
|
function: Callable[..., T],
|
|
*args: Any,
|
|
cuda_sync: bool = False,
|
|
**kwargs: Any,
|
|
) -> T:
|
|
if cuda_sync:
|
|
_sync_cuda()
|
|
started = time.perf_counter()
|
|
try:
|
|
return function(*args, **kwargs)
|
|
finally:
|
|
if cuda_sync:
|
|
_sync_cuda()
|
|
self.durations[stage].append(time.perf_counter() - started)
|
|
|
|
def rows(self) -> list[dict[str, Any]]:
|
|
extract_time = sum(self.durations.get("extract_total", []))
|
|
rows = []
|
|
for stage, values in self.durations.items():
|
|
total = sum(values)
|
|
rows.append(
|
|
{
|
|
"stage": stage,
|
|
"calls": len(values),
|
|
"total_seconds": total,
|
|
"average_ms": total / len(values) * 1000.0,
|
|
"p50_ms": statistics.median(values) * 1000.0,
|
|
"p95_ms": _percentile(values, 0.95) * 1000.0,
|
|
"maximum_ms": max(values) * 1000.0,
|
|
"percent_of_extract": (
|
|
total / extract_time * 100.0 if extract_time else None
|
|
),
|
|
}
|
|
)
|
|
return sorted(rows, key=lambda row: row["total_seconds"], reverse=True)
|
|
|
|
|
|
PROFILER = StageProfiler()
|
|
|
|
|
|
def _patch_function(
|
|
owner: Any,
|
|
name: str,
|
|
stage: str,
|
|
*,
|
|
cuda_sync: bool = False,
|
|
) -> Callable[..., Any]:
|
|
original = getattr(owner, name)
|
|
|
|
@functools.wraps(original)
|
|
def wrapped(*args: Any, **kwargs: Any) -> Any:
|
|
return PROFILER.measure(
|
|
stage,
|
|
original,
|
|
*args,
|
|
cuda_sync=cuda_sync,
|
|
**kwargs,
|
|
)
|
|
|
|
setattr(owner, name, wrapped)
|
|
return original
|
|
|
|
|
|
def _patch_bound_method(
|
|
owner: Any,
|
|
name: str,
|
|
stage: str,
|
|
*,
|
|
cuda_sync: bool,
|
|
) -> None:
|
|
original = getattr(owner, name)
|
|
|
|
@functools.wraps(original)
|
|
def wrapped(_self: Any, *args: Any, **kwargs: Any) -> Any:
|
|
return PROFILER.measure(
|
|
stage,
|
|
original,
|
|
*args,
|
|
cuda_sync=cuda_sync,
|
|
**kwargs,
|
|
)
|
|
|
|
setattr(owner, name, types.MethodType(wrapped, owner))
|
|
|
|
|
|
def _patch_component_call(component: Any, stage: str) -> None:
|
|
"""Time ``component(...)`` despite Python's class-level special lookup."""
|
|
|
|
component_class = type(component)
|
|
original = component_class.__call__
|
|
|
|
@functools.wraps(original)
|
|
def wrapped(instance: Any, *args: Any, **kwargs: Any) -> Any:
|
|
if instance is component:
|
|
return PROFILER.measure(stage, original, instance, *args, **kwargs)
|
|
return original(instance, *args, **kwargs)
|
|
|
|
# Qwen's processor, feature extractor, video processor, and tokenizer use
|
|
# distinct classes in Transformers 4.52.3. Guard anyway so an unexpected
|
|
# shared class cannot silently replace an earlier component timer.
|
|
marker = "_omni_bottleneck_profile_stage"
|
|
previous_stage = getattr(component_class, marker, None)
|
|
if previous_stage is not None:
|
|
raise RuntimeError(
|
|
f"Cannot separately profile {stage}: {component_class.__name__} "
|
|
f"is already patched as {previous_stage}"
|
|
)
|
|
setattr(component_class, marker, stage)
|
|
setattr(component_class, "__call__", wrapped)
|
|
|
|
|
|
def _install_profiling_hooks() -> None:
|
|
# One-time input setup.
|
|
_patch_function(direct, "_decode_audio_exact", "audio_decode")
|
|
_patch_function(direct, "_timestamp_array", "video_timestamp_index")
|
|
_patch_function(direct.CausalSourceAV, "__init__", "media_reader_init")
|
|
|
|
# Per-state CPU media and processor work.
|
|
_patch_function(direct.CausalSourceAV, "frames", "media_frames")
|
|
_patch_function(direct.CausalSourceAV, "audio_slice", "media_audio_slice")
|
|
_patch_function(core, "_prepare_conversation", "processor")
|
|
_patch_function(core, "process_mm_info", "processor_mm_info")
|
|
_patch_function(
|
|
core,
|
|
"_processor_batch_convert_to_tensors",
|
|
"processor_tensorize",
|
|
)
|
|
_patch_function(
|
|
core,
|
|
"_processor_batch_to_device",
|
|
"processor_h2d",
|
|
cuda_sync=True,
|
|
)
|
|
_patch_function(
|
|
core,
|
|
"_processor_batch_to_dtype",
|
|
"processor_dtype_cast",
|
|
cuda_sync=True,
|
|
)
|
|
|
|
# Per-state GPU/cache work. Leaf GPU calls synchronize for attribution.
|
|
_patch_function(core, "clone_cache", "cache_clone", cuda_sync=True)
|
|
_patch_function(core, "_multimodal_embeddings", "multimodal_merge_total")
|
|
_patch_function(core, "_append_embeddings", "append_total")
|
|
_patch_function(core, "_read_open_state", "readout_total")
|
|
_patch_function(
|
|
core.OpenStreamReplayExtractor,
|
|
"_prepare_pieces",
|
|
"prepare_piece_total",
|
|
)
|
|
_patch_function(
|
|
core.OpenStreamReplayExtractor,
|
|
"_cache_from_intervals",
|
|
"rolling_cache_rebuild",
|
|
)
|
|
_patch_function(
|
|
core.CompletedBlockPieceCache,
|
|
"put",
|
|
"completed_block_cache_store",
|
|
)
|
|
_patch_function(
|
|
core.CompletedBlockPieceCache,
|
|
"get_many",
|
|
"completed_block_cache_lookup",
|
|
)
|
|
_patch_function(torch, "save", "output_save")
|
|
|
|
original_extract = core.OpenStreamReplayExtractor.extract
|
|
|
|
@functools.wraps(original_extract)
|
|
def extract_wrapped(self: Any, *args: Any, **kwargs: Any) -> Any:
|
|
result = PROFILER.measure(
|
|
"extract_total",
|
|
original_extract,
|
|
self,
|
|
*args,
|
|
cuda_sync=True,
|
|
**kwargs,
|
|
)
|
|
PROFILER.metadata["num_states"] = int(len(result.times))
|
|
return result
|
|
|
|
core.OpenStreamReplayExtractor.extract = extract_wrapped
|
|
|
|
original_load = core.load_qwen_omni
|
|
|
|
@functools.wraps(original_load)
|
|
def load_wrapped(*args: Any, **kwargs: Any) -> Any:
|
|
model, processor = PROFILER.measure(
|
|
"model_load",
|
|
original_load,
|
|
*args,
|
|
cuda_sync=True,
|
|
**kwargs,
|
|
)
|
|
|
|
# These are the mutually exclusive heavy GPU leaves inside one state.
|
|
_patch_bound_method(
|
|
model, "get_audio_features", "audio_tower", cuda_sync=True
|
|
)
|
|
_patch_bound_method(
|
|
model, "get_video_features", "video_tower", cuda_sync=True
|
|
)
|
|
_patch_bound_method(model, "get_rope_index", "tmrope", cuda_sync=True)
|
|
_patch_bound_method(
|
|
model.model, "forward", "thinker_forward", cuda_sync=True
|
|
)
|
|
processor_components = {
|
|
"processor_qwen_call": processor,
|
|
"processor_audio": getattr(processor, "feature_extractor", None),
|
|
"processor_video": getattr(processor, "video_processor", None),
|
|
"processor_tokenizer": getattr(processor, "tokenizer", None),
|
|
}
|
|
missing_components = [
|
|
stage
|
|
for stage, component in processor_components.items()
|
|
if component is None
|
|
]
|
|
if missing_components:
|
|
raise RuntimeError(
|
|
"Installed Qwen processor is missing profiling components: "
|
|
+ ", ".join(missing_components)
|
|
)
|
|
for stage, component in processor_components.items():
|
|
_patch_component_call(component, stage)
|
|
_patch_bound_method(
|
|
processor,
|
|
"apply_chat_template",
|
|
"processor_chat_template",
|
|
cuda_sync=False,
|
|
)
|
|
_patch_bound_method(
|
|
processor,
|
|
"replace_multimodal_special_tokens",
|
|
"processor_placeholder_expansion",
|
|
cuda_sync=False,
|
|
)
|
|
PROFILER.metadata["processor_component_classes"] = {
|
|
stage: type(component).__name__
|
|
for stage, component in processor_components.items()
|
|
}
|
|
torch.cuda.reset_peak_memory_stats()
|
|
return model, processor
|
|
|
|
core.load_qwen_omni = load_wrapped
|
|
|
|
|
|
def _start_gpu_monitor(path: Path) -> tuple[subprocess.Popen[str] | None, TextIO | None]:
|
|
command = [
|
|
"nvidia-smi",
|
|
"-i",
|
|
str(PHYSICAL_GPU_ID),
|
|
"--query-gpu=timestamp,index,name,utilization.gpu,utilization.memory,"
|
|
"memory.used,memory.total,power.draw,temperature.gpu",
|
|
"--format=csv,nounits",
|
|
"-l",
|
|
f"{GPU_SAMPLE_SECONDS:g}",
|
|
]
|
|
try:
|
|
output = path.open("w", encoding="utf-8", buffering=1)
|
|
process = subprocess.Popen(
|
|
command,
|
|
stdout=output,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
)
|
|
return process, output
|
|
except (FileNotFoundError, OSError) as exc:
|
|
print(f"WARNING: could not start nvidia-smi monitor: {exc}", flush=True)
|
|
return None, None
|
|
|
|
|
|
def _stop_gpu_monitor(
|
|
process: subprocess.Popen[str] | None,
|
|
output: TextIO | None,
|
|
) -> None:
|
|
if process is not None and process.poll() is None:
|
|
process.terminate()
|
|
try:
|
|
process.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
process.kill()
|
|
process.wait(timeout=5)
|
|
if output is not None:
|
|
output.close()
|
|
|
|
|
|
def _numeric(value: str | None) -> float | None:
|
|
if not value:
|
|
return None
|
|
match = re.search(r"[-+]?[0-9]*\.?[0-9]+", value)
|
|
return float(match.group(0)) if match else None
|
|
|
|
|
|
def _gpu_telemetry_summary(path: Path) -> dict[str, Any]:
|
|
if not path.is_file():
|
|
return {"samples": 0}
|
|
metrics: dict[str, list[float]] = defaultdict(list)
|
|
with path.open("r", encoding="utf-8", errors="replace") as handle:
|
|
reader = csv.DictReader(handle)
|
|
for row in reader:
|
|
for key, value in row.items():
|
|
normalized = key.strip().lower()
|
|
number = _numeric(value)
|
|
if number is None:
|
|
continue
|
|
if normalized.startswith("utilization.gpu"):
|
|
metrics["gpu_util_percent"].append(number)
|
|
elif normalized.startswith("utilization.memory"):
|
|
metrics["memory_util_percent"].append(number)
|
|
elif normalized.startswith("memory.used"):
|
|
metrics["memory_used_mib"].append(number)
|
|
elif normalized.startswith("power.draw"):
|
|
metrics["power_watts"].append(number)
|
|
elif normalized.startswith("temperature.gpu"):
|
|
metrics["temperature_c"].append(number)
|
|
result: dict[str, Any] = {
|
|
"samples": max((len(values) for values in metrics.values()), default=0)
|
|
}
|
|
for name, values in metrics.items():
|
|
result[name] = {
|
|
"mean": statistics.fmean(values),
|
|
"p50": statistics.median(values),
|
|
"p95": _percentile(values, 0.95),
|
|
"max": max(values),
|
|
}
|
|
return result
|
|
|
|
|
|
def _print_summary(
|
|
*,
|
|
rows: list[dict[str, Any]],
|
|
gpu_summary: dict[str, Any],
|
|
overall_seconds: float,
|
|
) -> None:
|
|
extract_seconds = sum(PROFILER.durations.get("extract_total", []))
|
|
num_states = int(PROFILER.metadata.get("num_states", 0))
|
|
print("\n=== Qwen Omni bottleneck profile (inclusive timings) ===")
|
|
print(
|
|
"Nested rows intentionally overlap; compare leaf rows such as "
|
|
"processor_audio/processor_video/processor_tokenizer/processor_h2d "
|
|
"and the model-tower/Thinker rows."
|
|
)
|
|
print(
|
|
f"overall={overall_seconds:.3f}s extract={extract_seconds:.3f}s "
|
|
f"states={num_states} "
|
|
f"throughput={(num_states / extract_seconds if extract_seconds else 0):.3f} state/s"
|
|
)
|
|
print(
|
|
f"{'stage':<27} {'calls':>7} {'total_s':>10} {'avg_ms':>10} "
|
|
f"{'p95_ms':>10} {'extract%':>10}"
|
|
)
|
|
for row in rows:
|
|
percent = row["percent_of_extract"]
|
|
percent_text = "-" if percent is None else f"{percent:.1f}"
|
|
print(
|
|
f"{row['stage']:<27} {row['calls']:>7d} "
|
|
f"{row['total_seconds']:>10.3f} {row['average_ms']:>10.3f} "
|
|
f"{row['p95_ms']:>10.3f} {percent_text:>10}"
|
|
)
|
|
|
|
gpu_util = gpu_summary.get("gpu_util_percent")
|
|
memory_used = gpu_summary.get("memory_used_mib")
|
|
if gpu_util:
|
|
print(
|
|
"GPU telemetry: "
|
|
f"samples={gpu_summary['samples']} "
|
|
f"util mean/p95/max={gpu_util['mean']:.1f}/"
|
|
f"{gpu_util['p95']:.1f}/{gpu_util['max']:.1f}%"
|
|
)
|
|
if memory_used:
|
|
print(
|
|
"GPU memory MiB mean/p95/max="
|
|
f"{memory_used['mean']:.0f}/{memory_used['p95']:.0f}/"
|
|
f"{memory_used['max']:.0f}"
|
|
)
|
|
|
|
leaf_stages = (
|
|
"media_frames",
|
|
"processor_mm_info",
|
|
"processor_chat_template",
|
|
"processor_audio",
|
|
"processor_video",
|
|
"processor_placeholder_expansion",
|
|
"processor_tokenizer",
|
|
"processor_tensorize",
|
|
"processor_h2d",
|
|
"processor_dtype_cast",
|
|
"audio_tower",
|
|
"video_tower",
|
|
"tmrope",
|
|
"thinker_forward",
|
|
"cache_clone",
|
|
)
|
|
totals = {
|
|
stage: sum(PROFILER.durations.get(stage, [])) for stage in leaf_stages
|
|
}
|
|
bottleneck = max(totals, key=totals.get)
|
|
advice = {
|
|
"media_frames": "bulk-decode/cache the selected 2-FPS resized frames",
|
|
"processor_mm_info": "cache the frame-list/media packaging metadata",
|
|
"processor_chat_template": "cache the invariant rendered chat template",
|
|
"processor_audio": "precompute or incrementally slice exact audio features",
|
|
"processor_video": "cache transformed 2-FPS frames or batch equal grids",
|
|
"processor_placeholder_expansion": "cache token-layout templates by prefix shape",
|
|
"processor_tokenizer": "cache token IDs by multimodal prefix shape",
|
|
"processor_tensorize": "return PyTorch tensors directly from the processor",
|
|
"processor_h2d": "use pinned buffers and non-blocking host-to-device copies",
|
|
"processor_dtype_cast": "produce the target floating dtype before transfer",
|
|
"audio_tower": "microbatch the ten 0.2-s prefixes in each 2-s block",
|
|
"video_tower": "cache unique 2-FPS visual inputs and batch equal shapes",
|
|
"tmrope": "cache position IDs for the ten repeated prefix shapes",
|
|
"thinker_forward": "microbatch independent streams/prefix branches",
|
|
"cache_clone": "replace full KV clones with append-and-crop branch reuse",
|
|
}
|
|
print(
|
|
f"Primary measured leaf bottleneck: {bottleneck} "
|
|
f"({totals[bottleneck]:.3f}s). Recommended next optimization: "
|
|
f"{advice[bottleneck]}."
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
PROFILE_LOG_DIR.mkdir(parents=True, exist_ok=True)
|
|
stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
console_log = PROFILE_LOG_DIR / f"profile_gpu{PHYSICAL_GPU_ID}_{stamp}.log"
|
|
gpu_log = PROFILE_LOG_DIR / f"gpu_gpu{PHYSICAL_GPU_ID}_{stamp}.csv"
|
|
json_log = PROFILE_LOG_DIR / f"profile_gpu{PHYSICAL_GPU_ID}_{stamp}.json"
|
|
|
|
original_stdout = sys.stdout
|
|
original_stderr = sys.stderr
|
|
status = "ok"
|
|
error_text: str | None = None
|
|
started = time.perf_counter()
|
|
usage_started = resource.getrusage(resource.RUSAGE_SELF)
|
|
|
|
with console_log.open("w", encoding="utf-8", buffering=1) as log_handle:
|
|
sys.stdout = TeeStream(original_stdout, log_handle)
|
|
sys.stderr = TeeStream(original_stderr, log_handle)
|
|
monitor: subprocess.Popen[str] | None = None
|
|
monitor_output: TextIO | None = None
|
|
try:
|
|
print(f"Physical GPU: {PHYSICAL_GPU_ID} -> process cuda:0")
|
|
print(f"CUDA_VISIBLE_DEVICES={os.environ['CUDA_VISIBLE_DEVICES']}")
|
|
print(f"Console log: {console_log.resolve()}")
|
|
print(f"GPU log: {gpu_log.resolve()}")
|
|
print(f"JSON summary: {json_log.resolve()}")
|
|
print("Extraction argv:", " ".join(sys.argv[1:]))
|
|
print(
|
|
"NOTE: CUDA synchronization is enabled for attribution; "
|
|
"this is not a production throughput benchmark."
|
|
)
|
|
if torch.cuda.is_available():
|
|
print(f"Visible GPU: {torch.cuda.get_device_name(0)}")
|
|
print(f"PyTorch: {torch.__version__}, CUDA: {torch.version.cuda}")
|
|
else:
|
|
raise RuntimeError("CUDA is not available after GPU selection")
|
|
|
|
monitor, monitor_output = _start_gpu_monitor(gpu_log)
|
|
_install_profiling_hooks()
|
|
direct.main()
|
|
except BaseException as exc:
|
|
status = "error"
|
|
error_text = f"{type(exc).__name__}: {exc}"
|
|
raise
|
|
finally:
|
|
_sync_cuda()
|
|
_stop_gpu_monitor(monitor, monitor_output)
|
|
overall_seconds = time.perf_counter() - started
|
|
usage_finished = resource.getrusage(resource.RUSAGE_SELF)
|
|
gpu_summary = _gpu_telemetry_summary(gpu_log)
|
|
rows = PROFILER.rows()
|
|
PROFILER.metadata.update(
|
|
{
|
|
"status": status,
|
|
"error": error_text,
|
|
"physical_gpu_id": PHYSICAL_GPU_ID,
|
|
"visible_cuda_device": 0,
|
|
"gpu_name": torch.cuda.get_device_name(0),
|
|
"torch_version": torch.__version__,
|
|
"torch_cuda_version": torch.version.cuda,
|
|
"overall_seconds": overall_seconds,
|
|
"cpu_user_seconds": (
|
|
usage_finished.ru_utime - usage_started.ru_utime
|
|
),
|
|
"cpu_system_seconds": (
|
|
usage_finished.ru_stime - usage_started.ru_stime
|
|
),
|
|
"max_rss_kib": usage_finished.ru_maxrss,
|
|
"peak_cuda_allocated_mib": (
|
|
torch.cuda.max_memory_allocated() / 1024**2
|
|
),
|
|
"peak_cuda_reserved_mib": (
|
|
torch.cuda.max_memory_reserved() / 1024**2
|
|
),
|
|
"console_log": str(console_log.resolve()),
|
|
"gpu_log": str(gpu_log.resolve()),
|
|
}
|
|
)
|
|
_print_summary(
|
|
rows=rows,
|
|
gpu_summary=gpu_summary,
|
|
overall_seconds=overall_seconds,
|
|
)
|
|
print(
|
|
"Process resources: "
|
|
f"user_cpu={PROFILER.metadata['cpu_user_seconds']:.2f}s "
|
|
f"system_cpu={PROFILER.metadata['cpu_system_seconds']:.2f}s "
|
|
f"max_rss={PROFILER.metadata['max_rss_kib'] / 1024:.1f}MiB "
|
|
f"peak_cuda_allocated="
|
|
f"{PROFILER.metadata['peak_cuda_allocated_mib']:.1f}MiB"
|
|
)
|
|
with json_log.open("w", encoding="utf-8") as handle:
|
|
json.dump(
|
|
{
|
|
"metadata": PROFILER.metadata,
|
|
"gpu_telemetry": gpu_summary,
|
|
"stages": rows,
|
|
"argv": sys.argv[1:],
|
|
},
|
|
handle,
|
|
indent=2,
|
|
)
|
|
print(f"Saved profile summary to {json_log.resolve()}")
|
|
sys.stdout = original_stdout
|
|
sys.stderr = original_stderr
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|