Blame for scripts/run_long_form.py

"""Run one real long-form Music3 generation in a fresh process.

This is intentionally a small end-to-end harness.  It invokes the normal
phased pipeline with the explicit experimental long-form flag, fixed
flush=32, and the existing chunked Flow path.  It does not run the M9-A
AR-only benchmark or compare multiple modes in one process.
"""

from __future__ import annotations

import argparse
import json
import platform
import subprocess
import sys
import threading
import time
import traceback
import wave
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Sequence

import mlx.core as mx

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from generation.ar_profiler import (  # noqa: E402
    _memory_pressure,
    _physical_memory_footprint_bytes,
    _physical_memory_bytes,
    _process_memory,
    _swap_used_bytes,
)
from generation.phased_pipeline import (  # noqa: E402
    DEFAULT_MODEL,
    PhasedPipeline,
)


ALLOWED_DURATIONS = (120, 180, 240)
DEFAULT_OUTPUT_ROOT = ROOT / "outputs" / "long_form_test"
DEFAULT_CAPTION = (
    "dark melodic techno, 128 BPM, female airy vocal, lonely but not sad, "
    "slow atmospheric intro, heavy bass enters after 16 bars, "
    "long continuous arrangement"
)
DEFAULT_LYRICS = (
    "[verse]\n"
    "Under a dead sky, I count the lights\n"
    "Every quiet window holds a different flame\n"
    "I walk the avenue where all the old trains fade\n"
    "The night keeps turning softly, calling out my name\n"
    "[pre-chorus]\n"
    "The city breathes in neon tides\n"
    "The wires sing above me, restless in the rain\n"
    "I feel the distant thunder moving through the lane\n"
    "One more step toward morning, one more pulse again\n"
    "[chorus]\n"
    "Rise above the noise, carry me through the night\n"
    "Hold on to the rhythm, keep the fire bright\n"
    "When the world is breaking, we can make it right\n"
    "Rise above the noise, carry me through the night\n"
    "[verse]\n"
    "Every shadow turns into a sign\n"
    "Every empty station leaves a trace behind\n"
    "I gather up the echoes that I thought I lost in time\n"
    "The road is wide and open, and the dark is mine\n"
    "[pre-chorus]\n"
    "The city breathes in neon tides\n"
    "The rain is on the rooftops, silver in the light\n"
    "The beat beneath the pavement keeps us synchronized\n"
    "One more step together, one more turn tonight\n"
    "[chorus]\n"
    "Rise above the noise, carry me through the night\n"
    "Hold on to the rhythm, keep the fire bright\n"
    "When the world is breaking, we can make it right\n"
    "Rise above the noise, carry me through the night\n"
    "[bridge]\n"
    "Hold the silence, let the pulse return\n"
    "Leave behind the lessons that were never ours to learn\n"
    "There is still a signal underneath the burn\n"
    "Follow where it leads us, watch the colors turn\n"
    "[chorus]\n"
    "Rise above the noise, carry me through the night\n"
    "Hold on to the rhythm, keep the fire bright\n"
    "When the world is breaking, we can make it right\n"
    "Rise above the noise, carry me through the night\n"
    "[verse]\n"
    "Morning finds us moving through a field of blue\n"
    "All the roads we wandered still come back to you\n"
    "I can hear the open sky beginning to renew\n"
    "Every distant heartbeat says we will get through\n"
    "[pre-chorus]\n"
    "The city breathes in neon tides\n"
    "The last stars disappear but leave the skyline bright\n"
    "We carry all the motion of a thousand nights\n"
    "One more step together, one more turn in flight\n"
    "[chorus]\n"
    "Rise above the noise, carry me through the night\n"
    "Hold on to the rhythm, keep the fire bright\n"
    "When the world is breaking, we can make it right\n"
    "Rise above the noise, carry me through the night\n"
    "[instrumental]\n"
    "[verse]\n"
    "Through the open doorway, warmer air arrives\n"
    "All the hidden colors wake and come alive\n"
    "We were never waiting, we were learning how to drive\n"
    "With the sound around us, we remember why\n"
    "[bridge]\n"
    "Hold the silence, let the pulse return\n"
    "Let the old horizons lose the shape they used to wear\n"
    "Every new direction is a promise in the air\n"
    "We can keep on walking, we can meet it there\n"
    "[chorus]\n"
    "Rise above the noise, carry me through the night\n"
    "Hold on to the rhythm, keep the fire bright\n"
    "When the world is breaking, we can make it right\n"
    "Rise above the noise, carry me through the night\n"
    "[instrumental]\n"
    "[verse]\n"
    "Every shadow turns into a sign\n"
    "Every moving light becomes a thread we can align\n"
    "I can feel the future in the rhythm of the line\n"
    "Keep the signal open, keep your hand in mine\n"
    "[pre-chorus]\n"
    "The city breathes in neon tides\n"
    "The beat is getting stronger as the hours slide\n"
    "Nothing in the distance has to stay outside\n"
    "One more step together, let the current guide\n"
    "[chorus]\n"
    "Rise above the noise, carry me through the night\n"
    "Hold on to the rhythm, keep the fire bright\n"
    "When the world is breaking, we can make it right\n"
    "Rise above the noise, carry me through the night\n"
    "[verse]\n"
    "When the morning opens, there is more to say\n"
    "We will take the long road, we will find a way\n"
    "Every beat remembers everything we made\n"
    "And the song keeps rising with the light of day\n"
    "[chorus]\n"
    "Rise above the noise, carry me through the night\n"
    "Hold on to the rhythm, keep the fire bright\n"
    "When the world is breaking, we can make it right\n"
    "Rise above the noise, carry me through the night\n"
)
DEFAULT_SEED = 0
DEFAULT_STEPS = 30
PRESSURE_ORDER = {None: -1, "normal": 0, "warning": 1, "critical": 2}


def validate_long_form_seconds(seconds: int) -> int:
    if seconds not in ALLOWED_DURATIONS:
        raise ValueError(f"seconds must be one of {ALLOWED_DURATIONS}")
    return seconds


def long_form_output_dir(root: str | Path, seconds: int) -> Path:
    validate_long_form_seconds(seconds)
    return Path(root) / f"{seconds}s"


def read_wav_metadata(path: Path) -> dict[str, Any]:
    with wave.open(str(path), "rb") as handle:
        frames = handle.getnframes()
        sample_rate = handle.getframerate()
        channels = handle.getnchannels()
    return {
        "frames": frames,
        "sample_rate": sample_rate,
        "channels": channels,
        "duration_seconds": frames / sample_rate if sample_rate else 0.0,
    }


def _event_time(events: Sequence[dict[str, Any]], name: str) -> float | None:
    for event in events:
        if event.get("name") == name:
            value = event.get("t_s")
            return float(value) if value is not None else None
    return None


def _duration_between(
    events: Sequence[dict[str, Any]], start_name: str, end_name: str
) -> float | None:
    start = _event_time(events, start_name)
    end = _event_time(events, end_name)
    if start is None or end is None or end < start:
        return None
    return round(end - start, 3)


def stage_timings(events: Sequence[dict[str, Any]]) -> dict[str, float | None]:
    return {
        "ar_stage_seconds": _duration_between(events, "stage_a_start", "stage_a_done"),
        "flow_vocoder_seconds": _duration_between(
            events, "stage_b_weights", "stage_b_done"
        ),
        "flow_stage_seconds": _duration_between(
            events, "stage_b_start", "stage_b_done"
        ),
        "wav_write_seconds": _duration_between(events, "stage_b_done", "write_done"),
    }


def _run(command: list[str]) -> str | None:
    try:
        result = subprocess.run(command, capture_output=True, text=True, check=False)
    except OSError:
        return None
    return result.stdout.strip() if result.returncode == 0 else None


def _environment(model: str) -> dict[str, Any]:
    import importlib.metadata

    def version(name: str) -> str | None:
        try:
            return importlib.metadata.version(name)
        except importlib.metadata.PackageNotFoundError:
            return None

    return {
        "commit_hash": _run(["git", "rev-parse", "HEAD"]),
        "mac_model": _run(["sysctl", "-n", "hw.model"]),
        "chip": "Apple M4" if _run(["sysctl", "-n", "hw.model"]) == "Mac16,10" else None,
        "apple_silicon": platform.machine(),
        "ram_bytes": _physical_memory_bytes(),
        "macos": platform.mac_ver()[0],
        "python": sys.version.split()[0],
        "mlx": version("mlx"),
        "mlx_audio": version("mlx-audio"),
        "model": model,
        "quantization": "4bit" if "4bit" in model.lower() else "unknown",
    }


@dataclass
class SystemMonitor:
    interval_seconds: float = 60.0

    def __post_init__(self) -> None:
        self.initial_swap_bytes = _swap_used_bytes()
        self.latest_swap_bytes = self.initial_swap_bytes
        self.max_swap_delta_bytes = 0
        self.max_rss_bytes = 0
        self.max_footprint_bytes = 0
        self.max_mlx_active_bytes = 0
        self.max_mlx_peak_bytes = 0
        self.highest_pressure: str | None = None
        self.sample_count = 0
        self.critical_event = threading.Event()
        self._stop_event = threading.Event()
        self._thread: threading.Thread | None = None

    def sample(self) -> None:
        rss, _vm = _process_memory()
        footprint = _physical_memory_footprint_bytes() or 0
        swap = _swap_used_bytes()
        pressure = _memory_pressure()
        self.max_rss_bytes = max(self.max_rss_bytes, rss)
        self.max_footprint_bytes = max(self.max_footprint_bytes, footprint)
        self.max_mlx_active_bytes = max(self.max_mlx_active_bytes, int(mx.get_active_memory()))
        self.max_mlx_peak_bytes = max(self.max_mlx_peak_bytes, int(mx.get_peak_memory()))
        self.latest_swap_bytes = swap
        if swap is not None and self.initial_swap_bytes is not None:
            self.max_swap_delta_bytes = max(
                self.max_swap_delta_bytes, swap - self.initial_swap_bytes
            )
        if PRESSURE_ORDER.get(pressure, -1) > PRESSURE_ORDER.get(self.highest_pressure, -1):
            self.highest_pressure = pressure
        if pressure == "critical":
            self.critical_event.set()
        self.sample_count += 1

    def _run(self) -> None:
        while not self._stop_event.wait(self.interval_seconds):
            self.sample()

    def start(self) -> None:
        self.sample()
        self._thread = threading.Thread(target=self._run, name="long-form-monitor")
        self._thread.daemon = True
        self._thread.start()

    def stop(self) -> None:
        self._stop_event.set()
        if self._thread is not None:
            self._thread.join(timeout=max(5.0, self.interval_seconds + 1.0))
        self.sample()


def _gb(value: int | None) -> float | None:
    return round(value / 1024**3, 3) if value is not None else None


def run_long_form_test(
    seconds: int,
    *,
    output_root: str | Path = DEFAULT_OUTPUT_ROOT,
) -> dict[str, Any]:
    validate_long_form_seconds(seconds)
    run_dir = long_form_output_dir(output_root, seconds)
    run_dir.mkdir(parents=True, exist_ok=True)
    output_path = run_dir / "output.wav"
    config_path = run_dir / "config.json"
    summary_path = run_dir / "summary.json"
    log_path = run_dir / "benchmark.log"
    environment = _environment(DEFAULT_MODEL)
    config = {
        "target_seconds": seconds,
        "caption": DEFAULT_CAPTION,
        "lyrics": DEFAULT_LYRICS,
        "seed": DEFAULT_SEED,
        "steps": DEFAULT_STEPS,
        "model": DEFAULT_MODEL,
        "quantization": "4bit",
        "experimental_long_form": True,
        "flush_interval": 32,
        "output": str(output_path),
        "environment": environment,
    }
    config_path.write_text(json.dumps(config, indent=2, ensure_ascii=False) + "\n")
    log_lines = [
        f"start target_seconds={seconds} seed={DEFAULT_SEED} steps={DEFAULT_STEPS}",
        "mode=experimental_long_form flush_interval=32 phased_pipeline=true",
        f"output={output_path}",
    ]

    monitor = SystemMonitor()
    pipeline_report: dict[str, Any] | None = None
    error: str | None = None
    started = time.perf_counter()
    monitor.start()
    try:
        pipeline = PhasedPipeline(
            label=f"long_form_test/{seconds}s", model_path=DEFAULT_MODEL
        )
        pipeline_report = pipeline.generate(
            caption=DEFAULT_CAPTION,
            lyrics=DEFAULT_LYRICS,
            duration=seconds,
            steps=DEFAULT_STEPS,
            seed=DEFAULT_SEED,
            output=str(output_path),
            experimental_long_form=True,
            stop_event=monitor.critical_event,
        )
    except BaseException as exc:  # noqa: BLE001 - preserve a useful failure artifact
        error = (
            "KeyboardInterrupt: manually stopped after sustained swap growth"
            if isinstance(exc, KeyboardInterrupt)
            else f"{type(exc).__name__}: {exc}"
        )
        log_lines.append(traceback.format_exc())
    finally:
        monitor.stop()
    elapsed = time.perf_counter() - started

    wav_metadata = None
    if output_path.is_file():
        try:
            wav_metadata = read_wav_metadata(output_path)
        except (OSError, wave.Error) as exc:
            error = error or f"WAVReadError: {exc}"

    events = pipeline_report.get("events", []) if pipeline_report else []
    timings = stage_timings(events)
    actual_seconds = wav_metadata["duration_seconds"] if wav_metadata else None
    duration_matches = (
        actual_seconds is not None
        and abs(actual_seconds - seconds) <= max(1.0, seconds * 0.02)
    )
    if error:
        status = "failed"
    elif monitor.critical_event.is_set():
        status = "safety_stopped"
        error = "critical memory pressure reported by monitor"
    elif not output_path.is_file() or not duration_matches:
        status = "incomplete"
    else:
        status = "completed"

    pipeline_peak = None
    if pipeline_report is not None:
        pipeline_peak = float(pipeline_report.get("max_stage_peak_gb", 0.0))
    monitor_peak = _gb(monitor.max_mlx_peak_bytes) or 0.0
    end_to_end = round(elapsed, 3)
    summary = {
        "status": status,
        "success": status == "completed",
        "error": error,
        "target_seconds": seconds,
        "actual_wav_seconds": round(actual_seconds, 3) if actual_seconds is not None else None,
        "duration_matches_target": duration_matches,
        "ar_stage_seconds": timings["ar_stage_seconds"],
        "flow_vocoder_seconds": timings["flow_vocoder_seconds"],
        "flow_stage_seconds": timings["flow_stage_seconds"],
        "wav_write_seconds": timings["wav_write_seconds"],
        "end_to_end_seconds": end_to_end,
        "rtf": round(elapsed / actual_seconds, 3) if actual_seconds else None,
        "mlx_peak_gb": round(max(pipeline_peak or 0.0, monitor_peak), 3),
        "process_peak_rss_gb": _gb(monitor.max_rss_bytes),
        "process_peak_footprint_gb": _gb(monitor.max_footprint_bytes),
        "system_swap_delta_gb": _gb(monitor.max_swap_delta_bytes),
        "highest_memory_pressure": monitor.highest_pressure,
        "system_responsiveness": "not_directly_measured_by_runner",
        "monitor_sample_count": monitor.sample_count,
        "wav": wav_metadata,
        "pipeline_report": pipeline_report,
        "config": config,
    }
    log_lines.extend(
        [
            f"status={status}",
            f"actual_wav_seconds={summary['actual_wav_seconds']}",
            f"end_to_end_seconds={end_to_end}",
            f"ar_stage_seconds={timings['ar_stage_seconds']}",
            f"flow_vocoder_seconds={timings['flow_vocoder_seconds']}",
            f"mlx_peak_gb={summary['mlx_peak_gb']}",
            f"process_peak_rss_gb={summary['process_peak_rss_gb']}",
            f"process_peak_footprint_gb={summary['process_peak_footprint_gb']}",
            f"system_swap_delta_gb={summary['system_swap_delta_gb']}",
            f"highest_memory_pressure={summary['highest_memory_pressure']}",
        ]
    )
    for event in events:
        log_lines.append(
            f"event={event.get('name')} t_s={event.get('t_s')} note={event.get('note', '')}"
        )
    log_path.write_text("\n".join(log_lines) + "\n")
    summary_path.write_text(json.dumps(summary, indent=2, ensure_ascii=False) + "\n")
    print(json.dumps(summary, indent=2, ensure_ascii=False))
    return summary


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Music3 M4 16GB long-form real-world test")
    parser.add_argument("--seconds", type=int, choices=ALLOWED_DURATIONS, required=True)
    parser.add_argument("--output-root", type=Path, default=DEFAULT_OUTPUT_ROOT)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = _parser().parse_args(argv)
    summary = run_long_form_test(args.seconds, output_root=args.output_root)
    return 0 if summary["status"] == "completed" else 2


if __name__ == "__main__":
    raise SystemExit(main())