tools/profile_analysis.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import asyncio
import json
import resource
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from openmusic_analysis.application import build_service


async def main() -> None:
    parser = argparse.ArgumentParser(description="Measure model load and per-representation latency")
    parser.add_argument("audio", type=Path)
    parser.add_argument("--lyrics", type=Path)
    args = parser.parse_args()
    service = build_service()
    lyrics = args.lyrics.read_text(encoding="utf-8") if args.lyrics else None
    measurements = {}
    started = time.perf_counter()
    await service.load_models()
    measurements["model_startup_seconds"] = time.perf_counter() - started
    representations = ["audio.global", "audio.temporal"]
    if lyrics is not None:
        representations.append("lyrics.global")
    for representation in representations:
        started = time.perf_counter()
        await service.analyze(
            args.audio,
            lyrics=lyrics,
            requested_representations=[representation],
            track_id=None,
            content_identity=None,
        )
        measurements[f"{representation}_seconds"] = time.perf_counter() - started
    measurements["peak_process_rss_platform_units"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    try:
        import torch

        if torch.cuda.is_available():
            measurements["peak_cuda_memory_bytes"] = torch.cuda.max_memory_allocated()
    except ImportError:
        pass
    print(json.dumps(measurements, indent=2))


if __name__ == "__main__":
    asyncio.run(main())