run.py

#!/usr/bin/env python3
"""interop_bench orchestrator.

Runs a publisher × subscriber matrix of MoQ client implementations against a
local relay, one pair at a time:

    fixture -> publisher adapter -> moq-relay -> subscriber adapter -> received.media
                     |                                   |
                events_pub.csv                      events_sub.csv

For each pair it manages the relay, spawns the two adapter processes with the
shared env contract (see README), samples their CPU/RSS, and hands the pair
directory to analyze.py. Ends with a run-level report.md.

NOTE: this file is LLM-generated.

Usage:
    ./run.py                      # all available pairs
    ./run.py --pairs membrane:moq-cli,moq-cli:membrane
    ./run.py --list               # probe adapters and show availability
"""

import argparse
import json
import os
import shutil
import signal
import subprocess
import sys
import time
from pathlib import Path

import analyze

BENCH_DIR = Path(__file__).resolve().parent
# Convenience fallback for running inside a membrane_moq_plugin checkout that
# carries a built reference clone; everything else resolves via env or $PATH.
LEGACY_MOQ_REF = BENCH_DIR.parent / ".reference" / "moq"

DEFAULT_FIXTURES = {
    "h264": BENCH_DIR / "fixtures" / "bbb_20s_25fps.h264",
    "h264-bframes": BENCH_DIR / "fixtures" / "bbb_20s_25fps_bframes.h264",
    "h265": BENCH_DIR / "fixtures" / "bbb_20s_25fps_hevc.h265",
    "aac": BENCH_DIR / "fixtures" / "bbb_20s.aac",
}
SUFFIX_MEDIA = {".h264": "h264", ".h265": "h265", ".aac": "aac"}
DEFAULT_FPS = 25
DEFAULT_RELAY_URL = "https://localhost:4443/anon"
WARMUP_S = 2.0  # publisher head start before the subscriber joins


def find_tool(env_var, names, legacy):
    """Resolve an external binary: $env_var override, then $PATH, then the
    reference-clone debug build. None if nowhere."""
    override = os.environ.get(env_var)
    if override:
        return override
    for name in names:
        path = shutil.which(name)
        if path:
            return path
    if legacy.exists():
        return str(legacy)
    return None


# --- process tree resource accounting ----------------------------------------


def read_stats(stats_path: Path):
    """Resource summary written by wrap.py, reshaped for metrics/report."""
    if not stats_path.exists():
        return None
    stats = json.loads(stats_path.read_text())
    if stats["wall_s"] <= 0:
        return None
    return {
        "cpu_mean_pct": round(100 * (stats["user_s"] + stats["sys_s"]) / stats["wall_s"], 1),
        "rss_max_mb": round(stats["max_rss_bytes"] / 2**20, 1),
        "wall_s": stats["wall_s"],
    }


# --- relay management --------------------------------------------------------


class Relay:
    """Uses an already-running relay when one exists, otherwise starts one
    (resolved via $MOQ_RELAY / $PATH / the reference clone) and stops it
    afterwards."""

    def __init__(self, url):
        self.url = url
        self.proc = None

    def __enter__(self):
        if self.url != DEFAULT_RELAY_URL:
            print(f"using external relay ({self.url})")
            return self
        running = subprocess.run(["pgrep", "-x", "moq-relay"], capture_output=True)
        if running.returncode == 0:
            print(f"using already-running moq-relay ({self.url})")
            return self
        relay_bin = find_tool(
            "MOQ_RELAY", ["moq-relay"], LEGACY_MOQ_REF / "target" / "debug" / "moq-relay"
        )
        if relay_bin is None:
            sys.exit(
                "no running relay and no moq-relay binary"
                " (checked $MOQ_RELAY, $PATH, the reference clone);"
                " enter the nix dev shell or set --relay-url to a running relay"
            )
        self.proc = subprocess.Popen(
            [relay_bin, str(BENCH_DIR / "relay.toml")],
            cwd=BENCH_DIR,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        time.sleep(1.5)
        if self.proc.poll() is not None:
            sys.exit("moq-relay exited immediately; run it manually to see why")
        print(f"started moq-relay (pid {self.proc.pid}, {self.url})")
        return self

    def __exit__(self, *exc):
        if self.proc and self.proc.poll() is None:
            self.proc.terminate()
            try:
                self.proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.proc.kill()


# --- adapter execution -------------------------------------------------------


def load_registry():
    return json.loads((BENCH_DIR / "adapters" / "registry.json").read_text())


def base_env():
    env = os.environ.copy()
    # The nix package ships the moq-cli crate's binary as `moq`.
    cli = find_tool("MOQ_CLI", ["moq", "moq-cli"], LEGACY_MOQ_REF / "target" / "debug" / "moq-cli")
    if cli:
        env["MOQ_CLI"] = cli
    return env


def probe(adapter):
    try:
        result = subprocess.run(
            adapter["probe"], cwd=BENCH_DIR, env=base_env(),
            capture_output=True, text=True, timeout=60,
        )
        return result.returncode == 0, result.stderr.strip()
    except (subprocess.TimeoutExpired, FileNotFoundError) as error:
        return False, str(error)


def spawn(cmd, env, log_path, stats_path):
    log = open(log_path, "w")
    wrapped = [sys.executable, str(BENCH_DIR / "wrap.py"), str(stats_path), "--"] + cmd
    proc = subprocess.Popen(
        wrapped, cwd=BENCH_DIR, env=env, stdout=log, stderr=subprocess.STDOUT,
        start_new_session=True,
    )
    proc.log_file = log
    return proc


def stop(proc, grace_s):
    """Wait for a clean exit; escalate SIGINT -> SIGKILL on the process group.

    Returns (exit_code, forced): forced means the process only stopped because
    the harness signalled it (normal for subscribers that follow a broadcast
    indefinitely instead of exiting when it ends).
    """
    forced = False
    try:
        proc.wait(timeout=grace_s)
    except subprocess.TimeoutExpired:
        forced = True
        try:
            os.killpg(proc.pid, signal.SIGINT)
            proc.wait(timeout=5)
        except (subprocess.TimeoutExpired, ProcessLookupError):
            try:
                os.killpg(proc.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
            proc.wait()
    proc.log_file.close()
    return proc.returncode, forced


def run_pair(pub_name, sub_name, registry, args, run_dir, fps, clip_duration_s):
    pub = registry["publishers"][pub_name]
    sub = registry["subscribers"][sub_name]
    pair_dir = run_dir / f"{pub_name}__{sub_name}"
    pair_dir.mkdir(parents=True)
    broadcast = f"bench-{pub_name}-{sub_name}.hang"
    timeout_s = clip_duration_s + 40

    env = base_env()
    env.update(
        MOQ_URL=args.relay_url,
        MOQ_BROADCAST=broadcast,
        BENCH_MEDIA=args.media,
        BENCH_INPUT=str(args.fixture),
        BENCH_INPUT_MP4=str(args.fixture.with_suffix(".mp4")),
        # rounded for adapters doing shell arithmetic (audio: ~43.07 -> 43)
        BENCH_FPS=str(round(fps)),
        BENCH_OUTPUT=str(pair_dir / "received.media"),
        BENCH_TIMEOUT_MS=str(int(timeout_s * 1000)),
    )

    print(f"\n=== {pub_name} -> {sub_name}  (broadcast {broadcast}) ===")

    pub_env = env | pub.get("env", {}) | {"BENCH_EVENTS": str(pair_dir / "events_pub.csv")}
    pub_proc = spawn(pub["cmd"], pub_env, pair_dir / "pub.log", pair_dir / "stats_pub.json")
    pub_spawn_ns = time.time_ns()

    time.sleep(WARMUP_S)
    if pub_proc.poll() is not None:
        print(f"publisher exited during warmup (code {pub_proc.returncode}); see pub.log")

    sub_env = env | sub.get("env", {}) | {"BENCH_EVENTS": str(pair_dir / "events_sub.csv")}
    sub_proc = spawn(sub["cmd"], sub_env, pair_dir / "sub.log", pair_dir / "stats_sub.json")
    sub_spawn_ns = time.time_ns()

    # The publisher ends the broadcast when done; the subscriber should then
    # exit on its own. Grace periods cover implementations that don't.
    pub_exit, pub_forced = stop(pub_proc, timeout_s)
    sub_exit, sub_forced = stop(sub_proc, 15)

    meta = {
        "pub": pub_name,
        "sub": sub_name,
        "pub_exit": pub_exit,
        "sub_exit": sub_exit,
        "pub_forced_stop": pub_forced,
        "sub_forced_stop": sub_forced,
        "pub_resources": read_stats(pair_dir / "stats_pub.json"),
        "sub_resources": read_stats(pair_dir / "stats_sub.json"),
        "pub_spawn_wall_ns": pub_spawn_ns,
        "sub_spawn_wall_ns": sub_spawn_ns,
        "pub_events": pub["events"],
        "sub_events": sub["events"],
        "sub_output_format": sub["output"],
        "media": args.media,
        "fps": fps,
        "fixture": str(args.fixture),
        "broadcast": broadcast,
    }
    (pair_dir / "meta.json").write_text(json.dumps(meta, indent=2) + "\n")

    metrics = analyze.analyze_pair(pair_dir)
    d = metrics["delivery"]
    print(
        f"{metrics['status']}: {d['frames_matched']}/{d['frames_fixture']} frames"
        f" (corrupt {d['corrupt']}), startup {metrics['startup_ms']} ms,"
        f" e2e {metrics['e2e_latency'] and metrics['e2e_latency']['p50_ms']} ms p50"
    )
    return metrics


# --- main ---------------------------------------------------------------------


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--pairs", help="comma-separated pub:sub pairs (default: all available)")
    parser.add_argument("--media", choices=[*DEFAULT_FIXTURES, "video", "audio"],
                        help="fixture format (default: inferred from the fixture suffix;"
                             " video/audio are aliases for h264/aac)")
    parser.add_argument("--fixture", type=Path, default=None)
    parser.add_argument("--fps", type=int, default=DEFAULT_FPS,
                        help="video fixture frame rate (aac derives its own from the ADTS header)")
    parser.add_argument("--relay-url", default=DEFAULT_RELAY_URL)
    parser.add_argument("--results", type=Path, default=BENCH_DIR / "results")
    parser.add_argument("--list", action="store_true", help="probe adapters and exit")
    args = parser.parse_args()

    if args.media:
        args.media = analyze.canonical_media(args.media)
    if args.fixture is None:
        args.fixture = DEFAULT_FIXTURES[args.media or "h264"]
    if args.media is None:
        # A B-frame fixture shares the .h264 suffix; select it via --media.
        args.media = SUFFIX_MEDIA.get(args.fixture.suffix, "h264")

    registry = load_registry()

    def supports_media(adapter):
        return args.media in adapter.get("media", ["h264"])

    available = {"publishers": {}, "subscribers": {}}
    for role in available:
        for name, adapter in registry[role].items():
            if not supports_media(adapter):
                available[role][name] = False
                if args.list:
                    print(f"{role[:-1]} {name}: skipped (no {args.media} adapter)")
                continue
            ok, why = probe(adapter)
            available[role][name] = ok
            if args.list or not ok:
                status = "available" if ok else f"unavailable ({why or 'probe failed'})"
                print(f"{role[:-1]} {name}: {status}")
    if args.list:
        return

    if args.pairs:
        pairs = [tuple(pair.split(":")) for pair in args.pairs.split(",")]
        for pub_name, sub_name in pairs:
            for role, name in (("publishers", pub_name), ("subscribers", sub_name)):
                if not supports_media(registry[role][name]):
                    sys.exit(f"{role[:-1]} {name} has no {args.media} adapter")
    else:
        pairs = [
            (p, s)
            for p, ok_p in available["publishers"].items() if ok_p
            for s, ok_s in available["subscribers"].items() if ok_s
        ]
    if not pairs:
        sys.exit("no runnable pairs")

    if not args.fixture.exists():
        sys.exit(f"fixture missing: {args.fixture} (run make_fixture.sh)")
    fixture_bytes = args.fixture.read_bytes()
    frames = analyze.fixture_frames(args.media, fixture_bytes)
    fps = analyze.adts_frame_rate(fixture_bytes) if args.media == "aac" else args.fps
    clip_duration_s = len(frames) / fps
    print(f"fixture: {args.fixture.name} ({args.media}),"
          f" {len(frames)} frames, {clip_duration_s:.0f}s @ {fps:.1f}fps")

    run_dir = args.results / time.strftime("%Y%m%d_%H%M%S")
    run_dir.mkdir(parents=True)

    with Relay(args.relay_url):
        for pub_name, sub_name in pairs:
            run_pair(pub_name, sub_name, registry, args, run_dir, fps, clip_duration_s)

    report = analyze.write_report(run_dir)
    print(f"\n{report.read_text()}")
    print(f"results: {run_dir}")


if __name__ == "__main__":
    main()