Files
unshackle/scripts/live_segmentbase_ab.py
T
imSp4rky 7eee34014e perf(dash): collapse single-URL byte-range segments into one direct download
Also throttle hedge-median recompute to 0.5s, cache replace_fields format
patterns and short-circuit when no $ token, and stream HLS merge with
copyfileobj. Adds test_dash_collapse and live A/B tooling.
2026-07-16 12:42:16 -06:00

131 lines
4.8 KiB
Python

"""Live A/B test for SegmentBase (byte-range) DASH downloads.
Drives the real DASH pipeline against public SegmentBase test manifests and reports
wall time, size, throughput and output validity (ffprobe). On branches that have the
single-URL collapse (DASH._collapsible_single_url) it also reports whether the
collapse fired; on checkouts without it (e.g. dev) the same script runs the classic
per-segment + merge path, so running it from both checkouts gives an A/B comparison:
uv run python scripts/live_segmentbase_ab.py --track largest
/path/to/dev/.venv/bin/python scripts/live_segmentbase_ab.py --track largest
Throwaway-ish operational tool: it makes real network requests to public test assets.
"""
from __future__ import annotations
import shutil
import subprocess
import time
from pathlib import Path
from tempfile import mkdtemp
from typing import Any, Optional
import click
from requests import Session
from unshackle.core.manifests import DASH
from unshackle.core.tracks.track import DownloadContext
# (label, SegmentBase MPD). netflix-1a is collapse-eligible (Initialization range from
# byte 0); sony-1a has no Initialization range so every checkout takes the classic path.
STREAMS: list[tuple[str, str]] = [
("netflix-1a", "https://dash.akamaized.net/dash264/TestCases/1a/netflix/exMPD_BIP_TC1.mpd"),
("sony-1a", "https://dash.akamaized.net/dash264/TestCases/1a/sony/SNE_DASH_SD_CASE1A_REVISED.mpd"),
]
guard_results: list[bool] = []
def install_guard_spy() -> bool:
"""Record collapse decisions when this checkout has the guard; no-op otherwise."""
orig = getattr(DASH, "_collapsible_single_url", None)
if orig is None:
return False
def spy(is_subtitle: bool, segments: Any, init_len: Optional[int]) -> bool:
result = orig(is_subtitle, segments, init_len)
guard_results.append(result)
return result
DASH._collapsible_single_url = staticmethod(spy) # type: ignore[method-assign]
return True
def ffprobe(path: Path) -> str:
if not path.exists() or not shutil.which("ffprobe"):
return "n/a"
r = subprocess.run(
["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "csv=p=0", str(path)],
capture_output=True,
text=True,
)
return f"{float(r.stdout.strip()):.0f}s" if r.stdout.strip() else "INVALID"
def run_one(label: str, url: str, which: str, workers: int, has_guard: bool) -> str:
guard_results.clear()
session = Session()
try:
tracks = DASH.from_url(url=url, session=session).to_tracks(language="en")
except Exception as e: # noqa: BLE001
return f"{label:12} PARSE ERROR {type(e).__name__}: {e}"
vids = list(tracks.videos)
if not vids:
return f"{label:12} NO VIDEO TRACKS"
key = lambda t: t.bitrate or 0 # noqa: E731
track = min(vids, key=key) if which == "smallest" else max(vids, key=key)
work = Path(mkdtemp(prefix="segbase_ab_"))
save_path = work / "out.mp4"
# adaptive_workers only exists on the downloader-speed branch; skip elsewhere
extra = {"adaptive_workers": True} if "adaptive_workers" in DownloadContext.__dataclass_fields__ else {}
ctx = DownloadContext(
save_path=save_path,
save_dir=work / "segs",
progress=lambda **ev: None,
session=session,
proxy=None,
max_workers=workers,
**extra,
)
t0 = time.time()
err = ""
try:
DASH.download_track(track=track, ctx=ctx)
except Exception as e: # noqa: BLE001
err = f" ERROR {type(e).__name__}: {e}"
dt = time.time() - t0
size = save_path.stat().st_size if save_path.exists() else 0
mbps = (size / 1e6 / dt) if dt > 0 and size else 0.0
probe = ffprobe(save_path)
shutil.rmtree(work, ignore_errors=True)
if has_guard:
collapse = "yes" if guard_results and guard_results[0] else "no"
else:
collapse = "unsupported"
res = f"{getattr(track, 'width', '?')}x{getattr(track, 'height', '?')}"
return (
f"{label:12} {dt:6.1f}s {size / 1e6:8.1f}MB {mbps:6.1f}MB/s [{res}] "
f"collapse: {collapse:11} duration: {probe}{err}"
)
@click.command()
@click.option("--workers", default=16, show_default=True, help="max workers (adaptive cap where supported)")
@click.option("--track", "which", type=click.Choice(["smallest", "largest"]), default="smallest", show_default=True)
@click.option("--only", default="", help="substring filter on stream labels")
def main(workers: int, which: str, only: str) -> None:
has_guard = install_guard_spy()
print(f"track: {which}, max_workers: {workers}, collapse guard present: {has_guard}\n")
for label, url in STREAMS:
if only and only.lower() not in label.lower():
continue
print(run_one(label, url, which, workers, has_guard))
if __name__ == "__main__":
main()