From bc92c78a60bc829b320d89fc79e0e228cabcccfd Mon Sep 17 00:00:00 2001 From: brxs Date: Sat, 8 Aug 2026 13:11:24 -0700 Subject: [PATCH] spike: add PyTorch MRT2 validation harness --- docs/issue-109-hardware-checklist.md | 83 ++ docs/spike-mrt2-pytorch.md | 134 +++ spike/mrt2_pytorch/README.md | 65 ++ spike/mrt2_pytorch/__init__.py | 1 + spike/mrt2_pytorch/harness.py | 863 ++++++++++++++++++ spike/mrt2_pytorch/provenance.json | 49 + spike/mrt2_pytorch/requirements-candidate.txt | 10 + spike/mrt2_pytorch/tests/test_harness.py | 87 ++ 8 files changed, 1292 insertions(+) create mode 100644 docs/issue-109-hardware-checklist.md create mode 100644 docs/spike-mrt2-pytorch.md create mode 100644 spike/mrt2_pytorch/README.md create mode 100644 spike/mrt2_pytorch/__init__.py create mode 100644 spike/mrt2_pytorch/harness.py create mode 100644 spike/mrt2_pytorch/provenance.json create mode 100644 spike/mrt2_pytorch/requirements-candidate.txt create mode 100644 spike/mrt2_pytorch/tests/test_harness.py diff --git a/docs/issue-109-hardware-checklist.md b/docs/issue-109-hardware-checklist.md new file mode 100644 index 0000000..d1dafb6 --- /dev/null +++ b/docs/issue-109-hardware-checklist.md @@ -0,0 +1,83 @@ +# Issue #109 Linux/Windows NVIDIA qualification checklist + +Run every item once on representative Ubuntu 22.04+ and Windows 11 systems. +Attach the JSON and logs to #109; do not summarize a failing/missing run as a +pass. + +## Host record + +- [ ] Record OS/build, CPU, physical RAM, GPU model, VRAM, NVIDIA driver, power + mode, and whether the GPU drives a display. +- [ ] Save `nvidia-smi` output before and after each run. +- [ ] Record Python, PyTorch, Transformers, CUDA runtime, cuDNN, and exact + `provenance.json` revisions. +- [ ] Confirm `torch.cuda.is_available()` and the reported compute capability. +- [ ] Start from a clean runtime with no system Python, Git, shell, CUDA toolkit, + or compiler dependency in the packaged execution path. + +## Acquisition/offline proof + +- [ ] Acquire the small model at `7037d99551c84ac5c6afb7f1a5e58c65e7233dbb`. +- [ ] Acquire MusicCoCa at `236c488e38aa98643805514996934d705668298b`. +- [ ] Verify hashes, disconnect network (or set `HF_HUB_OFFLINE=1`), and confirm + the harness starts. A cache miss must fail clearly without a download. +- [ ] Repeat for base `92087988d05d0fe38b11f021f0b0d00a75afb86b` + only if base is a proposed supported model. + +## Required matrix + +For every proposed acceleration mode, run: + +```text +shared-worker × 25 frames × 600 seconds +shared-worker × 5 frames × 600 seconds +per-deck × 25 frames × 600 seconds +per-deck × 5 frames × 600 seconds +``` + +- [ ] Run the dry adapter first and retain it separately as synthetic evidence. +- [ ] Run `python -m spike.mrt2_pytorch.harness --backend upstream ...` for the + full matrix with parity guidance (the default). +- [ ] Repeat with a live app build feeding the native engine and capture its + engine-reported underrun counter; the harness proxy is not a substitute. +- [ ] Confirm both decks prime, generate/play continuously, and report zero + native engine underruns. +- [ ] Confirm p50/p95/p99/max latency, generated-audio/wall ratio, RSS peak, + PyTorch CUDA peak, process VRAM peak, driver, temperature, P-state, and + power rows are present. +- [ ] Inspect thermal clocks/temperature across the full ten minutes; record any + throttling or laptop power-mode dependency. + +## Controls and continuity + +- [ ] At the scheduled change, verify weighted prompt, temperature, top-k, + prompt CFG, note CFG, drum CFG, MIDI onset, and onset-to-sustain take effect + without resetting continuation state. +- [ ] Listen/inspect boundaries in both chunk modes for gaps, repeats, channel + swaps, clipping, or a sample-rate mismatch (required: float32, stereo, + 48 kHz). +- [ ] Run twice from a fresh state with the same seed and record determinism; + confirm changing seed with retained state does not falsely claim reseeding. +- [ ] Exercise text and captured-audio style inputs. +- [ ] Switch small/base models and verify reset/readiness behavior. + +## Supervision/topology + +- [ ] Record cold start, warm-up, readiness, and clean shutdown time. +- [ ] Kill deck A in the per-deck topology; deck B must continue and the parent + must report the failure. +- [ ] Kill the shared worker; both decks must stop/report failure coherently. +- [ ] Close the parent/app during generation and confirm no Python, helper, + compiler, or GPU process remains. +- [ ] Force one malformed control payload and one generation exception; no + silent hang or stale-playing state is allowed. + +## Release gate + +- [ ] Name the lowest GPU/VRAM/driver configuration that passed every required + run on both operating systems. +- [ ] Confirm packaged execution requires no runtime compiler; otherwise mark the + acceleration path no-go. +- [ ] #108 confirms code, derived-weight, MusicCoCa, and original-weight notices. +- [ ] Publish the topology recommendation and raw evidence on #109 before #110 + starts. diff --git a/docs/spike-mrt2-pytorch.md b/docs/spike-mrt2-pytorch.md new file mode 100644 index 0000000..77187c6 --- /dev/null +++ b/docs/spike-mrt2-pytorch.md @@ -0,0 +1,134 @@ +# Issue #109 — PyTorch MRT2 portability spike + +Audit date: 2026-08-08. This note covers software/API validation. No Linux or +Windows NVIDIA machine was available in this workspace, so it contains no +fabricated performance results. + +## Outcome + +The port is consumable without an LSDJ fork: use the Transformers model snapshot +at an exact Hugging Face revision, keep the adapter thin, prefetch all assets, +and run with `local_files_only=True`. The repository is public and Apache-2.0. +The earlier Space named in discovery material is now authentication-gated, but +it is not needed by the proposed path. + +The production decision is **conditional no-go pending hardware qualification**. +Do not begin #110 until both target operating systems pass the two-deck, +ten-minute matrix and #108 resolves weight/processor licensing. This is a +maturity and evidence gate, not a rejection of the implementation. + +## Immutable dependency path + +| Component | Immutable reference | Finding | +| --- | --- | --- | +| PyTorch source/API | `multimodalart/magenta-realtime-torch@6d076baa3df3b10448876c400521a015a5137c59` | Public Apache-2.0 source; no PyTorch-specific release/tag | +| Base model + remote code | `magenta-community/magenta-realtime-2@92087988d05d0fe38b11f021f0b0d00a75afb86b` | Transformers custom model, ~2.46B reported parameters | +| Small model + remote code | `magenta-community/magenta-realtime-2-small@7037d99551c84ac5c6afb7f1a5e58c65e7233dbb` | Transformers custom model, ~282M reported parameters | +| MusicCoCa processor | `magenta-community/magenta-rt-musiccoca-torch@236c488e38aa98643805514996934d705668298b` | Text/audio encoder artifacts; exact revision must be supplied by the adapter | +| Original Google assets | `google/magenta-realtime-2@010aa0dcb0dfd27b24f0ad07b4dad63e8f9521cc` | Declared base model/weight provenance | + +The runnable fixture is `spike/mrt2_pytorch/harness.py`; the complete +machine-readable inventory is `spike/mrt2_pytorch/provenance.json`. + +The audited GitHub `pyproject.toml` is still the upstream JAX/MLX package: it +does not declare a PyTorch extra or pin `torch`/`transformers`. The Transformers +snapshot is therefore the cleaner dependency boundary. LSDJ must own a +target-specific lock for Python 3.12, PyTorch, Transformers, and CUDA wheels; +the spike's exact direct pins are candidates, not a production lock. + +## Control and state parity + +| LSDJ MLX behavior | PyTorch port | Disposition | +| --- | --- | --- | +| Weighted text prompt embeddings | `MusicCoCaProcessor.layer()` returns 12 style tokens | Thin adapter; cache tokens on changes, not per chunk | +| Text negative prompt | Not exposed by the MRT2 deck today | No mapping needed; upstream CFG negatives are masked conditioning, not negative text | +| Temperature and top-k per chunk | `generate(temperature=, top_k=)` | Direct | +| Prompt/note CFG matching `.mlxfn` | `generate(..., guidance=True)` | Direct but more expensive than upstream's default token-CFG path; benchmark parity mode | +| Drum adherence token | `cfg_drums` | Direct; true-CFG mode still treats drums through the learned token | +| Note states `-1/0/1/2/3` and drum `-1/0/1` | Raw `notes`/`drums` arrays | Direct; LSDJ retains onset-to-sustain decay | +| Small/base model selection | Separate pinned model repositories | Direct; switching requires a worker/model restart | +| Per-deck continuation | Returned state contains decoder, RNG, and codec state | One model can safely own two state objects; must be sustained-tested | +| Seed | Seed creates the RNG only when state is new | Adapter documents reset-to-reseed; current LSDJ UI has no MRT2 seed control | +| 25-frame and 5-frame chunks | Arbitrary `frames`; 40 ms/frame, 48 kHz stereo | Direct; output length/continuity must be measured on hardware | +| Audio style sampling | Processor source implements audio embedding/resampling | API match; golden/hardware parity is not present in upstream CI | +| Warm-up/reset | No stable high-level readiness API | Adapter performs a throwaway generate then clears state; readiness contract belongs in #110 | + +The upstream workflow at the audited revision runs macOS MLX tests only. It +does not run the PyTorch port on Linux, Windows, or CUDA, and checkpoint-heavy +parity tests are skipped. Claims in model cards are useful provenance, not a +substitute for LSDJ qualification. + +## Harness topology and ring model + +`shared-worker` loads one model in one process and keeps independent deck A/B +state, scheduling generation round-robin. `per-deck` starts two processes and +therefore two model instances. Both use the same control-change sequence and +the same 1.5-second prebuffer gate. + +The harness reports starvation duration/transitions as `underrun_proxy_*`. +LSDJ's Rust engine counts individual audio callback blocks after the ring is +primed, so the proxy is deliberately not named an engine underrun. Final +qualification must capture both the harness JSON and the app's native telemetry. + +Start with the shared-worker/two-state topology: it avoids duplicating a large +model and has the smallest support-floor risk. Its failure domain covers both +decks and its serialized inference load may miss real time. Move to per-deck +workers only if concurrent GPU execution materially improves the 5-frame +two-deck result on supported hardware and the measured VRAM floor is acceptable. + +## Packaging feasibility + +An installer can ship without user-installed Python, Git, a shell, or a CUDA +toolkit by bundling an embedded Python runtime, exact binary wheels, and +prefetched snapshots. Users still need a compatible NVIDIA driver. The official +PyTorch release matrix publishes the same CUDA wheels for Linux and Windows. + +Risks that must be closed before release: + +- `trust_remote_code=True` executes snapshot code; only the audited commit may + be acquired, hash-checked, and promoted atomically. +- `model.load_processor()` defaults to a mutable repository reference. The LSDJ + adapter must resolve the exact processor revision locally, as the harness does. +- `torch.compile` can require a compiler on Windows. Production must not trigger + an unbundled MSVC/toolchain install. AOTInductor artifacts are GPU-architecture + specific, so one artifact cannot establish a broad GPU support floor. +- Eager, `torch.compile`, CUDA graph, and AOT behavior have not been compared on + the target systems. Upstream exposes the fast CUDA graph path through a stream + surface rather than the simple resumable `generate` call, so a stable chunk API + may require an upstream contribution. +- Produce platform-specific, hash-locked wheels only after choosing the CUDA + runtime and minimum driver from the hardware results. + +## Licensing/provenance escalation + +The fork's code is Apache-2.0. The derived Transformers model cards say +Apache-2.0, while they declare `google/magenta-realtime-2` as their base and the +Google weights are CC-BY-4.0. The MusicCoCa artifact card is also CC-BY-4.0. +Issue #108 must decide the effective redistribution/notice obligations; this +spike makes no legal conclusion. + +## Decision gates for #110 + +Go only when all are true: + +1. Linux and Windows each sustain both decks for ten minutes, at 25 and 5 + frames, with zero native engine underruns on the proposed minimum GPU. +2. The parity-guidance path meets the budget, including a live prompt and note + onset/sustain change; token-CFG results cannot stand in for it. +3. Exact wheel and snapshot locks install offline in a clean bundled runtime. +4. Startup/readiness, one-deck crash behavior, whole-tree shutdown, and device + recovery are demonstrated on both platforms. +5. #108 approves the notices, acknowledgement, and redistribution path. + +Until those gates pass, #109's software deliverables are complete but the +production recommendation remains conditional no-go. + +## Primary evidence + +- Source/API: +- Source license: +- Base model: +- Small model: +- Google model card/weights: +- PyTorch binary matrix: +- Windows compiler caveat: diff --git a/spike/mrt2_pytorch/README.md b/spike/mrt2_pytorch/README.md new file mode 100644 index 0000000..29dc0ee --- /dev/null +++ b/spike/mrt2_pytorch/README.md @@ -0,0 +1,65 @@ +# PyTorch MRT2 two-deck spike harness + +This directory is an isolated issue #109 fixture. It does not import or change +the production backend. It benchmarks the immutable upstream Transformers +snapshots in two process topologies, at 25 frames (~1 second) and 5 frames +(~200 ms), with a 1.5 second playback-prebuffer simulation. + +The harness records cold start, warm-up, per-deck p50/p95/p99 generation +latency, output duration, an underrun proxy, RSS, PyTorch CUDA allocation, and +`nvidia-smi` VRAM/driver/temperature rows. JSON marks dry runs as `synthetic`; +they are never qualification evidence. + +## CI/dry run + +From the repository root: + +```sh +python3 -m unittest discover -s spike/mrt2_pytorch/tests -v +python3 -m spike.mrt2_pytorch.harness \ + --backend dry-run --duration-seconds 2 \ + --output /tmp/mrt2-dry-run.json +``` + +## NVIDIA run + +Use Python 3.12 and install the direct candidate pins: + +```sh +python -m venv .venv-mrt2-spike +.venv-mrt2-spike/bin/python -m pip install -r spike/mrt2_pytorch/requirements-candidate.txt +``` + +On Windows, use `.venv-mrt2-spike\Scripts\python.exe` for the same commands. +Prefetch the immutable snapshots while online: + +```sh +hf download magenta-community/magenta-realtime-2-small \ + --revision 7037d99551c84ac5c6afb7f1a5e58c65e7233dbb +hf download magenta-community/magenta-rt-musiccoca-torch \ + --revision 236c488e38aa98643805514996934d705668298b +``` + +Then disconnect or set `HF_HUB_OFFLINE=1` and run the complete matrix: + +```sh +python -m spike.mrt2_pytorch.harness \ + --backend upstream \ + --model mrt2_small \ + --topologies shared-worker,per-deck \ + --frames 25,5 \ + --duration-seconds 600 \ + --prompt-change-seconds 60 \ + --output mrt2-small-eager.json +``` + +Repeat with `--acceleration torch-compile`; do not treat a runtime compiler as +a distributable solution until Windows packaging proves it needs no developer +toolchain. The default enables classifier-free guidance because that matches +LSDJ's `.mlxfn` path. `--token-cfg` intentionally measures upstream's cheaper, +non-parity conditioning-token path. + +The model adapter always uses exact revisions and `local_files_only=True`. +Any cache miss therefore fails clearly instead of silently downloading a +different asset. See `provenance.json` and +`docs/issue-109-hardware-checklist.md` before running qualification. diff --git a/spike/mrt2_pytorch/__init__.py b/spike/mrt2_pytorch/__init__.py new file mode 100644 index 0000000..8c581c9 --- /dev/null +++ b/spike/mrt2_pytorch/__init__.py @@ -0,0 +1 @@ +"""Issue #109's isolated PyTorch MRT2 validation harness.""" diff --git a/spike/mrt2_pytorch/harness.py b/spike/mrt2_pytorch/harness.py new file mode 100644 index 0000000..962562d --- /dev/null +++ b/spike/mrt2_pytorch/harness.py @@ -0,0 +1,863 @@ +"""Two-deck MRT2 benchmark harness for issue #109. + +This module deliberately has no LSDJ production imports. The real adapter loads +the pinned Hugging Face snapshot with ``local_files_only=True``; the dry adapter +lets CI exercise scheduling, topology, ring accounting, and result schemas with +no model, GPU, or network. +""" + +from __future__ import annotations + +import argparse +import dataclasses +import json +import math +import multiprocessing +import os +import pathlib +import platform +import queue +import statistics +import subprocess +import sys +import time +import traceback +from collections.abc import Iterable +from typing import Any + +SAMPLE_RATE = 48_000 +CHANNELS = 2 +FRAME_SECONDS = 0.04 +DECKS = (0, 1) +DEFAULT_PREBUFFER_SECONDS = 1.5 + +SOURCE_REPOSITORY = "https://github.com/multimodalart/magenta-realtime-torch.git" +SOURCE_REVISION = "6d076baa3df3b10448876c400521a015a5137c59" +MODEL_REVISIONS = { + "mrt2_base": ( + "magenta-community/magenta-realtime-2", + "92087988d05d0fe38b11f021f0b0d00a75afb86b", + ), + "mrt2_small": ( + "magenta-community/magenta-realtime-2-small", + "7037d99551c84ac5c6afb7f1a5e58c65e7233dbb", + ), +} +PROCESSOR_REPOSITORY = "magenta-community/magenta-rt-musiccoca-torch" +PROCESSOR_REVISION = "236c488e38aa98643805514996934d705668298b" + + +@dataclasses.dataclass +class RingBudget: + """Event-level proxy for LSDJ's per-deck 1.5 second prebuffer gate. + + The production engine counts callback blocks that find a primed ring short. + The harness cannot observe the device callback, so it records starvation + intervals and transitions. Both are labelled proxies in the JSON output. + """ + + prebuffer_seconds: float + fill_seconds: float = 0.0 + primed: bool = False + last_time: float | None = None + playback_started_at: float | None = None + underrun_events: int = 0 + underrun_seconds: float = 0.0 + fill_min_seconds: float | None = None + fill_max_seconds: float = 0.0 + + def advance(self, now: float) -> None: + if self.last_time is None: + self.last_time = now + return + elapsed = max(0.0, now - self.last_time) + self.last_time = now + if not self.primed: + return + before = self.fill_seconds + self.fill_seconds = max(0.0, before - elapsed) + self.fill_min_seconds = ( + self.fill_seconds + if self.fill_min_seconds is None + else min(self.fill_min_seconds, self.fill_seconds) + ) + if elapsed > before: + self.underrun_seconds += elapsed - before + if before > 0.0 or self.underrun_events == 0: + self.underrun_events += 1 + + def push(self, audio_seconds: float, now: float) -> None: + self.advance(now) + self.fill_seconds += max(0.0, audio_seconds) + self.fill_max_seconds = max(self.fill_max_seconds, self.fill_seconds) + if not self.primed and self.fill_seconds >= self.prebuffer_seconds: + self.primed = True + self.playback_started_at = now + self.fill_min_seconds = self.fill_seconds + + def ready_for_generation(self, target_seconds: float, chunk_seconds: float) -> bool: + if not self.primed: + return True + return self.fill_seconds <= max(0.0, target_seconds - chunk_seconds) + + def result(self, origin: float) -> dict[str, Any]: + return { + "prebuffer_seconds": self.prebuffer_seconds, + "primed": self.primed, + "time_to_prime_seconds": ( + None + if self.playback_started_at is None + else round(self.playback_started_at - origin, 6) + ), + "fill_final_seconds": round(self.fill_seconds, 6), + "fill_min_seconds": ( + None + if self.fill_min_seconds is None + else round(self.fill_min_seconds, 6) + ), + "fill_max_seconds": round(self.fill_max_seconds, 6), + "underrun_proxy_events": self.underrun_events, + "underrun_proxy_seconds": round(self.underrun_seconds, 6), + } + + +def percentile(values: Iterable[float], percent: float) -> float | None: + ordered = sorted(values) + if not ordered: + return None + if len(ordered) == 1: + return ordered[0] + rank = (len(ordered) - 1) * percent / 100.0 + low = math.floor(rank) + high = math.ceil(rank) + if low == high: + return ordered[low] + return ordered[low] + (ordered[high] - ordered[low]) * (rank - low) + + +def latency_summary(values: list[float]) -> dict[str, Any]: + return { + "count": len(values), + "mean_ms": None if not values else round(statistics.fmean(values) * 1_000, 3), + "p50_ms": _milliseconds(percentile(values, 50)), + "p95_ms": _milliseconds(percentile(values, 95)), + "p99_ms": _milliseconds(percentile(values, 99)), + "max_ms": _milliseconds(max(values) if values else None), + } + + +def _milliseconds(value: float | None) -> float | None: + return None if value is None else round(value * 1_000, 3) + + +def _rss_bytes() -> int | None: + """Current worker RSS without adding a benchmark-only dependency.""" + + if sys.platform.startswith("linux"): + try: + for line in pathlib.Path("/proc/self/status").read_text().splitlines(): + if line.startswith("VmRSS:"): + return int(line.split()[1]) * 1_024 + except (OSError, ValueError, IndexError): + return None + if sys.platform == "win32": + try: + import ctypes + from ctypes import wintypes + + class ProcessMemoryCounters(ctypes.Structure): + _fields_ = [ + ("cb", wintypes.DWORD), + ("PageFaultCount", wintypes.DWORD), + ("PeakWorkingSetSize", ctypes.c_size_t), + ("WorkingSetSize", ctypes.c_size_t), + ("QuotaPeakPagedPoolUsage", ctypes.c_size_t), + ("QuotaPagedPoolUsage", ctypes.c_size_t), + ("QuotaPeakNonPagedPoolUsage", ctypes.c_size_t), + ("QuotaNonPagedPoolUsage", ctypes.c_size_t), + ("PagefileUsage", ctypes.c_size_t), + ("PeakPagefileUsage", ctypes.c_size_t), + ] + + counters = ProcessMemoryCounters() + counters.cb = ctypes.sizeof(counters) + ok = ctypes.windll.psapi.GetProcessMemoryInfo( + ctypes.windll.kernel32.GetCurrentProcess(), + ctypes.byref(counters), + counters.cb, + ) + return int(counters.WorkingSetSize) if ok else None + except (AttributeError, OSError): + return None + try: + import resource + + value = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return int(value if sys.platform == "darwin" else value * 1_024) + except (ImportError, OSError): + return None + + +def _gpu_snapshot(worker_pids: set[int]) -> dict[str, Any] | None: + """Best-effort NVIDIA metadata. Raw rows survive driver schema differences.""" + + try: + gpu = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=index,name,uuid,driver_version,memory.total,memory.used,temperature.gpu,pstate,power.draw", + "--format=csv,noheader,nounits", + ], + check=True, + capture_output=True, + text=True, + timeout=10, + ) + apps = subprocess.run( + [ + "nvidia-smi", + "--query-compute-apps=pid,used_memory", + "--format=csv,noheader,nounits", + ], + check=True, + capture_output=True, + text=True, + timeout=10, + ) + except (FileNotFoundError, subprocess.SubprocessError): + return None + worker_memory_mib = 0.0 + matching_rows: list[str] = [] + for row in apps.stdout.splitlines(): + cells = [cell.strip() for cell in row.split(",")] + try: + pid = int(cells[0]) + used = float(cells[1]) + except (ValueError, IndexError): + continue + if pid in worker_pids: + matching_rows.append(row) + worker_memory_mib += used + return { + "gpu_rows": gpu.stdout.splitlines(), + "matching_compute_rows": matching_rows, + "worker_vram_mib": worker_memory_mib, + } + + +@dataclasses.dataclass(frozen=True) +class RunConfig: + backend: str + topology: str + frames: int + duration_seconds: float + prebuffer_seconds: float + target_ahead_seconds: float + model: str + acceleration: str + guidance: bool + dry_latency_ms: float + startup_timeout_seconds: float + worker_timeout_seconds: float + seed: int + prompt_change_seconds: float + + @property + def chunk_seconds(self) -> float: + return self.frames * FRAME_SECONDS + + +class DryAdapter: + def __init__(self, config: RunConfig): + self.delay = config.dry_latency_ms / 1_000.0 + self.states: dict[int, int] = {} + + def metadata(self) -> dict[str, Any]: + return {"adapter": "dry-run", "accelerator": "simulated"} + + def reset(self) -> None: + self.states.clear() + + def generate(self, deck: int, controls: dict[str, Any], frames: int) -> int: + time.sleep(self.delay) + self.states[deck] = self.states.get(deck, 0) + frames + return frames * round(SAMPLE_RATE * FRAME_SECONDS) + + def device_memory(self) -> dict[str, int | None]: + return { + "cuda_allocated_bytes": None, + "cuda_reserved_bytes": None, + "cuda_peak_allocated_bytes": None, + } + + +class UpstreamAdapter: + """Thin, snapshot-pinned adapter around upstream's Transformers API.""" + + def __init__(self, config: RunConfig): + import torch + from huggingface_hub import snapshot_download + from transformers import AutoModel + + if not torch.cuda.is_available(): + raise RuntimeError("PyTorch reports no CUDA accelerator") + model_repo, model_revision = MODEL_REVISIONS[config.model] + # Offline-only is intentional: installer/acquisition is a separate concern. + model_path = snapshot_download( + repo_id=model_repo, + revision=model_revision, + local_files_only=True, + ) + processor_path = snapshot_download( + repo_id=PROCESSOR_REPOSITORY, + revision=PROCESSOR_REVISION, + local_files_only=True, + ) + self.torch = torch + self.model = ( + AutoModel.from_pretrained( + model_path, + trust_remote_code=True, + dtype=torch.bfloat16, + local_files_only=True, + ) + .to("cuda") + .eval() + ) + self.model.load_processor(processor_path, device="cuda") + if config.acceleration == "torch-compile": + self.model.compile_steps() + elif config.acceleration != "eager": + raise ValueError(f"unsupported acceleration mode {config.acceleration!r}") + self.guidance = config.guidance + self.states: dict[int, Any] = {} + self.style_tokens: dict[ + tuple[tuple[str, ...], tuple[float, ...]], list[int] + ] = {} + self.torch.cuda.reset_peak_memory_stats() + + def metadata(self) -> dict[str, Any]: + torch = self.torch + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + return { + "adapter": "transformers-remote-code", + "torch_version": torch.__version__, + "torch_cuda_runtime": torch.version.cuda, + "cudnn_version": torch.backends.cudnn.version(), + "cuda_device": props.name, + "cuda_capability": list(torch.cuda.get_device_capability()), + "cuda_total_memory_bytes": props.total_memory, + } + + def reset(self) -> None: + self.states.clear() + + def _tokens(self, controls: dict[str, Any]) -> list[int]: + prompts = tuple(controls["prompts"]) + weights = tuple(float(value) for value in controls["weights"]) + key = prompts, weights + if key not in self.style_tokens: + self.style_tokens[key] = self.model.processor.layer(prompts, weights) + return self.style_tokens[key] + + def generate(self, deck: int, controls: dict[str, Any], frames: int) -> int: + torch = self.torch + torch.cuda.synchronize() + audio, state = self.model.generate( + style=self._tokens(controls), + notes=controls["notes"], + drums=controls["drums"], + cfg_musiccoca=controls["cfg_musiccoca"], + cfg_notes=controls["cfg_notes"], + cfg_drums=controls["cfg_drums"], + temperature=controls["temperature"], + top_k=controls["top_k"], + frames=frames, + seed=controls["seed"], + state=self.states.get(deck), + guidance=self.guidance, + ) + torch.cuda.synchronize() + if getattr(audio, "ndim", None) != 2 or audio.shape[1] != CHANNELS: + raise RuntimeError(f"upstream returned invalid audio shape {audio.shape!r}") + self.states[deck] = state + return int(audio.shape[0]) + + def device_memory(self) -> dict[str, int]: + torch = self.torch + return { + "cuda_allocated_bytes": torch.cuda.memory_allocated(), + "cuda_reserved_bytes": torch.cuda.memory_reserved(), + "cuda_peak_allocated_bytes": torch.cuda.max_memory_allocated(), + } + + +def _controls(deck: int, changed: bool, onset: bool, seed: int) -> dict[str, Any]: + notes = [-1] * 128 + if changed: + notes[60 + deck * 7] = 2 if onset else 1 + return { + "prompts": ( + ["warm disco funk", "analog synth bass"] + if not changed + else ["broken beat percussion", "ambient pads"] + ), + "weights": [0.7, 0.3] if not changed else [0.55, 0.45], + "temperature": 1.1 if not changed else 0.95, + "top_k": 50 if not changed else 64, + "cfg_musiccoca": 1.6 if not changed else 2.0, + "cfg_notes": 2.4, + "cfg_drums": 4.0, + "notes": notes, + "drums": [-1] if deck == 0 else [0], + "seed": seed + deck, + } + + +def _worker_main( + worker_id: int, + request_queue: Any, + result_queue: Any, + config: RunConfig, +) -> None: + started = time.perf_counter() + try: + adapter = ( + DryAdapter(config) + if config.backend == "dry-run" + else UpstreamAdapter(config) + ) + result_queue.put( + { + "type": "ready", + "worker": worker_id, + "pid": os.getpid(), + "startup_seconds": time.perf_counter() - started, + "rss_bytes": _rss_bytes(), + "metadata": adapter.metadata(), + "device_memory": adapter.device_memory(), + } + ) + while True: + request = request_queue.get() + action = request["action"] + if action == "shutdown": + result_queue.put( + { + "type": "stopped", + "worker": worker_id, + "rss_bytes": _rss_bytes(), + "device_memory": adapter.device_memory(), + } + ) + return + if action == "reset": + adapter.reset() + result_queue.put({"type": "reset", "worker": worker_id}) + continue + if action != "generate": + raise ValueError(f"unknown worker action {action!r}") + generated_at = time.perf_counter() + sample_frames = adapter.generate( + request["deck"], request["controls"], request["frames"] + ) + result_queue.put( + { + "type": "chunk", + "worker": worker_id, + "deck": request["deck"], + "sequence": request["sequence"], + "control_change": request["control_change"], + "latency_seconds": time.perf_counter() - generated_at, + "sample_frames": sample_frames, + "rss_bytes": _rss_bytes(), + "device_memory": adapter.device_memory(), + } + ) + except BaseException as error: + result_queue.put( + { + "type": "error", + "worker": worker_id, + "pid": os.getpid(), + "error": f"{type(error).__name__}: {error}", + "traceback": traceback.format_exc(), + } + ) + + +def _wait_for_messages( + result_queue: Any, + expected_type: str, + count: int, + timeout: float, +) -> list[dict[str, Any]]: + deadline = time.monotonic() + timeout + messages = [] + while len(messages) < count: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + f"timed out waiting for {count} {expected_type!r} messages" + ) + message = result_queue.get(timeout=remaining) + if message["type"] == "error": + raise RuntimeError( + f"worker failed: {message['error']}\n{message['traceback']}" + ) + if message["type"] != expected_type: + raise RuntimeError( + f"expected worker message {expected_type!r}, got {message['type']!r}" + ) + messages.append(message) + return messages + + +def run_benchmark(config: RunConfig) -> dict[str, Any]: + if config.topology not in {"shared-worker", "per-deck"}: + raise ValueError(f"unknown topology {config.topology!r}") + if config.frames not in {5, 25}: + raise ValueError("issue #109 requires 5-frame or 25-frame runs") + if config.target_ahead_seconds < config.prebuffer_seconds: + raise ValueError("target ahead must be at least the prebuffer threshold") + + context = multiprocessing.get_context("spawn") + result_queue = context.Queue() + worker_ids = [0] if config.topology == "shared-worker" else [0, 1] + request_queues = {worker: context.Queue() for worker in worker_ids} + processes = { + worker: context.Process( + target=_worker_main, + args=(worker, request_queues[worker], result_queue, config), + name=f"mrt2-bench-{worker}", + ) + for worker in worker_ids + } + launch_started = time.perf_counter() + for process in processes.values(): + process.start() + + worker_for_deck = ( + {0: 0, 1: 0} if config.topology == "shared-worker" else {0: 0, 1: 1} + ) + startup: list[dict[str, Any]] = [] + shutdown: list[dict[str, Any]] = [] + try: + startup = _wait_for_messages( + result_queue, + "ready", + len(worker_ids), + config.startup_timeout_seconds, + ) + ready_completed = time.perf_counter() + + # Warm each model process, then clear continuation state before measurement. + for worker in worker_ids: + deck = worker if config.topology == "per-deck" else 0 + request_queues[worker].put( + { + "action": "generate", + "deck": deck, + "sequence": -1, + "frames": config.frames, + "controls": _controls(deck, False, False, config.seed), + "control_change": False, + } + ) + warmup = _wait_for_messages( + result_queue, + "chunk", + len(worker_ids), + config.worker_timeout_seconds, + ) + for worker in worker_ids: + request_queues[worker].put({"action": "reset"}) + _wait_for_messages( + result_queue, + "reset", + len(worker_ids), + config.worker_timeout_seconds, + ) + + origin = time.perf_counter() + deadline = origin + config.duration_seconds + rings = { + deck: RingBudget(config.prebuffer_seconds, last_time=origin) + for deck in DECKS + } + latencies = {deck: [] for deck in DECKS} + change_latencies = {deck: [] for deck in DECKS} + sample_frames = {deck: 0 for deck in DECKS} + sequences = {deck: 0 for deck in DECKS} + changed = {deck: False for deck in DECKS} + onset_pending = {deck: False for deck in DECKS} + inflight_workers: set[int] = set() + rss_peak = {message["worker"]: message.get("rss_bytes") for message in startup} + cuda_peaks = {worker: 0 for worker in worker_ids} + gpu_samples: list[dict[str, Any]] = [] + next_gpu_sample = origin + round_robin = 0 + + while True: + now = time.perf_counter() + bounded_now = min(now, deadline) + for ring in rings.values(): + ring.advance(bounded_now) + + elapsed = bounded_now - origin + for deck in DECKS: + if not changed[deck] and elapsed >= config.prompt_change_seconds: + changed[deck] = True + onset_pending[deck] = True + + if now < deadline: + candidates = list(DECKS) + if config.topology == "shared-worker": + candidates = [round_robin, 1 - round_robin] + for deck in candidates: + worker = worker_for_deck[deck] + if worker in inflight_workers: + continue + if not rings[deck].ready_for_generation( + config.target_ahead_seconds, config.chunk_seconds + ): + continue + request_queues[worker].put( + { + "action": "generate", + "deck": deck, + "sequence": sequences[deck], + "frames": config.frames, + "controls": _controls( + deck, + changed[deck], + onset_pending[deck], + config.seed, + ), + "control_change": onset_pending[deck], + } + ) + sequences[deck] += 1 + inflight_workers.add(worker) + if config.topology == "shared-worker": + round_robin = 1 - deck + + worker_pids = {process.pid for process in processes.values() if process.pid} + if now >= next_gpu_sample: + sample = _gpu_snapshot(worker_pids) + if sample is not None: + sample["elapsed_seconds"] = round(now - origin, 3) + gpu_samples.append(sample) + next_gpu_sample = now + 1.0 + + if now >= deadline and not inflight_workers: + break + try: + message = result_queue.get(timeout=0.05) + except queue.Empty: + continue + if message["type"] == "error": + raise RuntimeError( + f"worker failed: {message['error']}\n{message['traceback']}" + ) + if message["type"] != "chunk": + raise RuntimeError(f"unexpected worker message {message['type']!r}") + worker = message["worker"] + deck = message["deck"] + inflight_workers.discard(worker) + completed = time.perf_counter() + if completed <= deadline: + audio_seconds = message["sample_frames"] / SAMPLE_RATE + rings[deck].push(audio_seconds, completed) + sample_frames[deck] += message["sample_frames"] + latency = message["latency_seconds"] + latencies[deck].append(latency) + if message["control_change"]: + change_latencies[deck].append(latency) + onset_pending[deck] = False + rss = message.get("rss_bytes") + if rss is not None: + rss_peak[worker] = max(rss_peak[worker] or 0, rss) + allocated = message.get("device_memory", {}).get( + "cuda_peak_allocated_bytes" + ) + if allocated is not None: + cuda_peaks[worker] = max(cuda_peaks[worker], allocated) + + ended = time.perf_counter() + for ring in rings.values(): + ring.advance(deadline) + + stop_started = time.perf_counter() + for worker in worker_ids: + request_queues[worker].put({"action": "shutdown"}) + shutdown = _wait_for_messages( + result_queue, + "stopped", + len(worker_ids), + config.worker_timeout_seconds, + ) + for process in processes.values(): + process.join(timeout=config.worker_timeout_seconds) + shutdown_seconds = time.perf_counter() - stop_started + + vram_samples = [sample["worker_vram_mib"] for sample in gpu_samples] + return { + "schema_version": 1, + "qualification": "synthetic" if config.backend == "dry-run" else "hardware", + "config": dataclasses.asdict(config), + "pins": { + "source_repository": SOURCE_REPOSITORY, + "source_revision": SOURCE_REVISION, + "model": { + "repository": MODEL_REVISIONS[config.model][0], + "revision": MODEL_REVISIONS[config.model][1], + }, + "processor": { + "repository": PROCESSOR_REPOSITORY, + "revision": PROCESSOR_REVISION, + }, + }, + "host": { + "platform": platform.platform(), + "system": platform.system(), + "release": platform.release(), + "machine": platform.machine(), + "python": platform.python_version(), + }, + "workers": startup, + "cold_start_wall_seconds": round(ready_completed - launch_started, 6), + "warmup": warmup, + "measurement_wall_seconds": round(ended - origin, 6), + "shutdown_seconds": round(shutdown_seconds, 6), + "shutdown": shutdown, + "failure_domain": ( + "both decks share one process" + if config.topology == "shared-worker" + else "one process failure is isolated to one deck" + ), + "decks": { + str(deck): { + "latency": latency_summary(latencies[deck]), + "control_change_latency": latency_summary(change_latencies[deck]), + "generated_audio_seconds": round( + sample_frames[deck] / SAMPLE_RATE, 6 + ), + "generated_audio_to_wall_ratio": round( + sample_frames[deck] / SAMPLE_RATE / config.duration_seconds, 6 + ), + "ring": rings[deck].result(origin), + } + for deck in DECKS + }, + "memory": { + "worker_rss_peak_bytes": rss_peak, + "worker_cuda_peak_allocated_bytes": cuda_peaks, + "nvidia_worker_vram_peak_mib": max(vram_samples) + if vram_samples + else None, + "nvidia_samples": gpu_samples, + }, + "notes": [ + "underrun_proxy_* is event-level 1.5 s ring simulation, not Rust engine telemetry", + "a hardware qualification must also record the app's engine-reported underrun counter", + ], + } + finally: + for worker, process in processes.items(): + if process.is_alive(): + try: + request_queues[worker].put({"action": "shutdown"}) + process.join(timeout=2) + except (OSError, ValueError): + pass + if process.is_alive(): + process.terminate() + process.join(timeout=2) + for request_queue in request_queues.values(): + request_queue.close() + result_queue.close() + + +def _parse_csv_ints(value: str) -> list[int]: + return [int(item.strip()) for item in value.split(",") if item.strip()] + + +def _parse_csv_strings(value: str) -> list[str]: + return [item.strip() for item in value.split(",") if item.strip()] + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--backend", choices=("dry-run", "upstream"), default="dry-run") + parser.add_argument( + "--topologies", default="shared-worker,per-deck", help="comma-separated matrix" + ) + parser.add_argument("--frames", default="25,5", help="comma-separated matrix") + parser.add_argument("--duration-seconds", type=float, default=600.0) + parser.add_argument( + "--prebuffer-seconds", type=float, default=DEFAULT_PREBUFFER_SECONDS + ) + parser.add_argument( + "--target-ahead-seconds", type=float, default=DEFAULT_PREBUFFER_SECONDS + ) + parser.add_argument("--model", choices=tuple(MODEL_REVISIONS), default="mrt2_small") + parser.add_argument( + "--acceleration", choices=("eager", "torch-compile"), default="eager" + ) + parser.add_argument( + "--token-cfg", + action="store_true", + help="use upstream token CFG instead of MLX-parity classifier-free guidance", + ) + parser.add_argument("--dry-latency-ms", type=float, default=10.0) + parser.add_argument("--startup-timeout-seconds", type=float, default=900.0) + parser.add_argument("--worker-timeout-seconds", type=float, default=300.0) + parser.add_argument("--seed", type=int, default=109) + parser.add_argument("--prompt-change-seconds", type=float, default=30.0) + parser.add_argument("--output", type=pathlib.Path) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + topologies = _parse_csv_strings(args.topologies) + frames = _parse_csv_ints(args.frames) + results = [] + for topology in topologies: + for frame_count in frames: + config = RunConfig( + backend=args.backend, + topology=topology, + frames=frame_count, + duration_seconds=args.duration_seconds, + prebuffer_seconds=args.prebuffer_seconds, + target_ahead_seconds=args.target_ahead_seconds, + model=args.model, + acceleration=args.acceleration, + guidance=not args.token_cfg, + dry_latency_ms=args.dry_latency_ms, + startup_timeout_seconds=args.startup_timeout_seconds, + worker_timeout_seconds=args.worker_timeout_seconds, + seed=args.seed, + prompt_change_seconds=min( + args.prompt_change_seconds, args.duration_seconds / 2 + ), + ) + results.append(run_benchmark(config)) + document = { + "schema_version": 1, + "created_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "results": results, + } + rendered = json.dumps(document, indent=2, sort_keys=True) + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered + "\n", encoding="utf-8") + else: + print(rendered) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/spike/mrt2_pytorch/provenance.json b/spike/mrt2_pytorch/provenance.json new file mode 100644 index 0000000..8a47c12 --- /dev/null +++ b/spike/mrt2_pytorch/provenance.json @@ -0,0 +1,49 @@ +{ + "schema_version": 1, + "audited_at": "2026-08-08", + "code": { + "repository": "https://github.com/multimodalart/magenta-realtime-torch.git", + "revision": "6d076baa3df3b10448876c400521a015a5137c59", + "revision_date": "2026-06-23T23:34:02Z", + "license_file": "https://github.com/multimodalart/magenta-realtime-torch/blob/6d076baa3df3b10448876c400521a015a5137c59/LICENSE", + "license_spdx": "Apache-2.0", + "release_status": "No PyTorch-specific tag or package extra at the audited revision; pin the commit/model snapshot." + }, + "models": { + "mrt2_base": { + "repository": "magenta-community/magenta-realtime-2", + "revision": "92087988d05d0fe38b11f021f0b0d00a75afb86b", + "reported_parameters": 2459164696, + "card_license": "Apache-2.0" + }, + "mrt2_small": { + "repository": "magenta-community/magenta-realtime-2-small", + "revision": "7037d99551c84ac5c6afb7f1a5e58c65e7233dbb", + "reported_parameters": 282195480, + "card_license": "Apache-2.0" + }, + "musiccoca_processor": { + "repository": "magenta-community/magenta-rt-musiccoca-torch", + "revision": "236c488e38aa98643805514996934d705668298b", + "card_license": "CC-BY-4.0" + }, + "original_google_weights": { + "repository": "google/magenta-realtime-2", + "revision": "010aa0dcb0dfd27b24f0ad07b4dad63e8f9521cc", + "card_license": "CC-BY-4.0" + } + }, + "runtime_candidate": { + "python": "3.12", + "torch": "2.12.1", + "transformers": "5.8.0", + "cuda_wheel_candidate": "cu130", + "status": "Registry-verified direct pins; not yet qualified on target hardware." + }, + "known_mutable_defaults_avoided_by_harness": [ + "AutoModel.from_pretrained without revision", + "model.load_processor() default repository without revision", + "hf_hub_download without revision" + ], + "license_review_required": "The derived Transformers cards say Apache-2.0 while the declared base Google weights and MusicCoCa artifacts say CC-BY-4.0. Issue #108 must determine redistribution and attribution obligations." +} diff --git a/spike/mrt2_pytorch/requirements-candidate.txt b/spike/mrt2_pytorch/requirements-candidate.txt new file mode 100644 index 0000000..223ec2a --- /dev/null +++ b/spike/mrt2_pytorch/requirements-candidate.txt @@ -0,0 +1,10 @@ +# Issue #109 benchmark candidate, registry-verified 2026-08-08. +# Direct requirements are exact; produce target-specific hash locks only after +# the Linux/Windows CUDA candidate is selected from real-hardware results. +torch==2.12.1 +transformers==5.8.0 +huggingface-hub==1.1.5 +numpy==2.3.5 +safetensors==0.7.0 +sentencepiece==0.2.1 +resampy==0.4.3 diff --git a/spike/mrt2_pytorch/tests/test_harness.py b/spike/mrt2_pytorch/tests/test_harness.py new file mode 100644 index 0000000..bc1b261 --- /dev/null +++ b/spike/mrt2_pytorch/tests/test_harness.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import unittest + +from spike.mrt2_pytorch.harness import ( + PROCESSOR_REVISION, + SOURCE_REVISION, + RingBudget, + RunConfig, + latency_summary, + percentile, + run_benchmark, +) + + +class RingBudgetTests(unittest.TestCase): + def test_prebuffer_shortfall_is_not_an_underrun(self) -> None: + ring = RingBudget(1.5, last_time=0.0) + + ring.push(1.0, 0.5) + ring.advance(4.0) + + self.assertFalse(ring.primed) + self.assertEqual(ring.underrun_events, 0) + self.assertEqual(ring.underrun_seconds, 0.0) + + def test_primed_ring_records_starvation_once_until_refilled(self) -> None: + ring = RingBudget(1.5, last_time=0.0) + ring.push(1.5, 0.0) + + ring.advance(2.0) + ring.advance(3.0) + + self.assertTrue(ring.primed) + self.assertEqual(ring.underrun_events, 1) + self.assertAlmostEqual(ring.underrun_seconds, 1.5) + + +class SummaryTests(unittest.TestCase): + def test_percentiles_interpolate_deterministically(self) -> None: + self.assertEqual(percentile([1.0, 2.0, 3.0, 4.0], 50), 2.5) + self.assertEqual(percentile([], 99), None) + + def test_latency_schema_reports_milliseconds(self) -> None: + summary = latency_summary([0.001, 0.003]) + + self.assertEqual(summary["count"], 2) + self.assertEqual(summary["p50_ms"], 2.0) + + +class DryRunTests(unittest.TestCase): + def _config(self, topology: str, frames: int) -> RunConfig: + return RunConfig( + backend="dry-run", + topology=topology, + frames=frames, + duration_seconds=0.12, + prebuffer_seconds=0.04, + target_ahead_seconds=max(0.04, frames * 0.04), + model="mrt2_small", + acceleration="eager", + guidance=True, + dry_latency_ms=1.0, + startup_timeout_seconds=10.0, + worker_timeout_seconds=10.0, + seed=109, + prompt_change_seconds=0.02, + ) + + def test_shared_worker_exercises_two_independent_deck_states(self) -> None: + result = run_benchmark(self._config("shared-worker", 5)) + + self.assertEqual(len(result["workers"]), 1) + self.assertGreater(result["decks"]["0"]["latency"]["count"], 0) + self.assertGreater(result["decks"]["1"]["latency"]["count"], 0) + self.assertTrue(result["decks"]["0"]["ring"]["primed"]) + + def test_per_deck_topology_reports_two_processes_and_exact_pins(self) -> None: + result = run_benchmark(self._config("per-deck", 25)) + + self.assertEqual(len(result["workers"]), 2) + self.assertEqual(result["pins"]["source_revision"], SOURCE_REVISION) + self.assertEqual(result["pins"]["processor"]["revision"], PROCESSOR_REVISION) + + +if __name__ == "__main__": + unittest.main()