Add profiling script for GPU bottleneck analysis in EasyCom extraction. Introduce command-line options for GPU ID, log directory, and sampling period. Implement stage profiling to measure extraction performance and synchronize CUDA operations. Enhance CausalSourceAV class with new parameters for decoder threads and open readers, optimizing video decoding process.
This commit is contained in:
Xin Wang
547
profile_streaming_bottleneck_gpu7.py
Normal file
547
profile_streaming_bottleneck_gpu7.py
Normal file
@@ -0,0 +1,547 @@
|
||||
#!/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 _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")
|
||||
|
||||
# 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(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
|
||||
)
|
||||
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_tower/video_tower/thinker_forward/cache_clone."
|
||||
)
|
||||
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",
|
||||
"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": "cache processor-ready AV tensors for every causal prefix",
|
||||
"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()
|
||||
Reference in New Issue
Block a user