"""Push protocol benchmark suite — Phase 7. Measures wall-clock time, bytes sent, and throughput for the scenarios defined in push-v2.md Phase 7: 1. tiny — 1 commit, 0 objects 2. incremental — 10 commits, 5 objects each; second push sends only new objects 3. cold_large — 1 commit, 500 × 4KB objects (~2 MB total) 4. repush_noop — same 500-object push repeated (all objects already in store) 5. many_small — 1 commit, 2000 × 256-byte objects 6. few_large — 1 commit, 5 × 500KB objects Run: python3 tests/bench_push.py python3 tests/bench_push.py --runs 5 Each scenario is repeated RUNS times (default 3). Prints a table with p50, p95 wall-clock (ms), total bytes sent, and throughput (MB/s). """ from __future__ import annotations import asyncio import os import statistics import sys import tempfile import time from pathlib import Path from typing import Any os.environ.setdefault("MUSE_ENV", "test") sys.path.insert(0, str(Path(__file__).parent.parent)) import msgpack from httpx import AsyncClient, ASGITransport from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.pool import NullPool from muse.core.types import blob_id, now_utc_iso from musehub.types.json_types import JSONObject, JSONValue from muse.core.mpack import WIRE_CONTENT_TYPE, MuseWireFrameWriter from musehub.db.database import Base, get_db import musehub.db.database as _database import musehub.db.muse_cli_models # noqa: F401 — register all ORM models from musehub.main import app from musehub.auth.request_signing import MSignContext, require_signed_request, optional_signed_request from musehub.models.wire import SFRAME_HEADER, SFRAME_COMMIT_PACK, SFRAME_END, SFRAME_OBJECT _fw = MuseWireFrameWriter() _DB_URL = os.environ.get( "TEST_DATABASE_URL", "postgresql+asyncpg://musehub:musehub@localhost:5434/musehub_test", ) _ENGINE = create_async_engine(_DB_URL, poolclass=NullPool) _SESSION_FACTORY = async_sessionmaker(bind=_ENGINE, expire_on_commit=False) _AUTH_CTX = MSignContext( handle="bench-user", identity_id="bench-user-id", is_agent=False, is_admin=False, ) # ── frame helpers ────────────────────────────────────────────────────────────── def _pack(obj: JSONValue) -> bytes: return msgpack.packb(obj, use_bin_type=True) def _wrap(ft: str, data: JSONValue) -> bytes: return _fw.wrap(frame_type=ft, payload=_pack(data)) _commit_counter = 0 def _make_commit(snapshot_id: str, parent_id: str | None = None) -> JSONObject: global _commit_counter _commit_counter += 1 cid = blob_id(f"bench-commit-{_commit_counter}-{os.urandom(4).hex()}".encode()) return { "commit_id": cid, "parent_ids": [parent_id] if parent_id else [], "parent_commit_id": parent_id, "parent2_commit_id": None, "snapshot_id": snapshot_id, "branch": "main", "message": "bench commit", "author": "bench", "committed_at": now_utc_iso(), "signature": "", "signer_key_id": "", "agent_id": "", "model_id": "", "metadata": {}, } def _make_snapshot(snap_id: str, manifest: JSONObject | None = None) -> JSONObject: return {"snapshot_id": snap_id, "manifest": manifest or {}} def _h_frame(n_objects: int, n_commits: int, force: bool = False) -> bytes: return _wrap(SFRAME_HEADER, { "t": SFRAME_HEADER, "branch": "main", "force": force, "have": [], "head": "", "n_objects": n_objects, "n_commits": n_commits, }) def _o_frame(oid: str, raw: bytes) -> bytes: return _wrap(SFRAME_OBJECT, { "t": SFRAME_OBJECT, "id": oid, "content": raw, "enc": "raw", "path": "bench.bin", "sz": len(raw), }) def _c_frame(commits: list[dict], snapshots: list[dict]) -> bytes: return _wrap(SFRAME_COMMIT_PACK, { "t": SFRAME_COMMIT_PACK, "commits": commits, "snapshots": snapshots, }) def _e_frame(n_objects: int, n_commits: int) -> bytes: return _wrap(SFRAME_END, { "t": SFRAME_END, "n_objects": n_objects, "n_commits": n_commits, }) def _push_body( objects: list[tuple[str, bytes]], commits: list[dict], snapshots: list[dict], force: bool = False, ) -> bytes: parts = [_h_frame(len(objects), len(commits), force=force)] for oid, raw in objects: parts.append(_o_frame(oid, raw)) parts.append(_c_frame(commits, snapshots)) parts.append(_e_frame(len(objects), len(commits))) return b"".join(parts) # ── bench harness ────────────────────────────────────────────────────────────── async def _create_repo_api(client: AsyncClient, name: str) -> tuple[str, str]: """Create a repo via the API and return (owner, slug).""" import musehub.db.database as _db_mod from sqlalchemy.ext.asyncio import AsyncSession async with _SESSION_FACTORY() as session: from tests.factories import create_repo repo = await create_repo(session, owner="bench-user", name=name) await session.commit() return repo.owner, repo.slug async def _do_push(client: AsyncClient, owner: str, slug: str, body: bytes) -> JSONObject: resp = await client.post( f"/{owner}/{slug}/push/stream", content=body, headers={"Content-Type": WIRE_CONTENT_TYPE}, ) unpacker = msgpack.Unpacker(raw=False) unpacker.feed(resp.content) last: JSONObject = {} for frame in unpacker: last = frame return last class BenchResult: def __init__(self, name: str) -> None: self.name = name self.times_ms: list[float] = [] self.bytes_sent: int = 0 def record(self, elapsed_s: float, body_bytes: int) -> None: self.times_ms.append(elapsed_s * 1000) self.bytes_sent = body_bytes def p50(self) -> float: return statistics.median(self.times_ms) def p95(self) -> float: n = len(self.times_ms) if n < 2: return self.times_ms[0] return sorted(self.times_ms)[max(0, int(n * 0.95) - 1)] def throughput_mbs(self) -> float: p50_s = self.p50() / 1000 if p50_s <= 0: return 0.0 return (self.bytes_sent / (1024 * 1024)) / p50_s # ── scenarios ───────────────────────────────────────────────────────────────── async def bench_tiny(client: AsyncClient, runs: int) -> BenchResult: """1 commit, 0 objects — measures pure protocol + DB overhead.""" result = BenchResult("tiny (1 commit, 0 obj)") owner, slug = await _create_repo_api(client, f"bench-tiny-{os.urandom(4).hex()}") prev_cid: str | None = None for i in range(runs): snap_id = blob_id(f"bench-tiny-snap-{i}-{os.urandom(4).hex()}".encode()) snap = _make_snapshot(snap_id) commit = _make_commit(snap_id, parent_id=prev_cid) body = _push_body([], [commit], [snap]) t0 = time.perf_counter() r = await _do_push(client, owner, slug, body) elapsed = time.perf_counter() - t0 assert r.get("ok") is True, f"tiny push failed: {r}" prev_cid = commit["commit_id"] result.record(elapsed, len(body)) return result async def bench_cold_large(client: AsyncClient, runs: int) -> BenchResult: """1 commit, 500 × 4KB objects (~2 MB) — cold upload throughput.""" result = BenchResult("cold_large (500 × 4KB, ~2MB)") n_obj = 500 obj_size = 4096 for run_i in range(runs): owner, slug = await _create_repo_api(client, f"bench-cold-{run_i}-{os.urandom(4).hex()}") objects = [] for i in range(n_obj): raw = os.urandom(obj_size - 8) + i.to_bytes(4, "big") + run_i.to_bytes(4, "big") oid = blob_id(raw) objects.append((oid, raw)) manifest = {f"f{i}.bin": oid for i, (oid, _) in enumerate(objects)} snap_id = blob_id(f"bench-cold-snap-{run_i}-{os.urandom(4).hex()}".encode()) snap = _make_snapshot(snap_id, manifest) commit = _make_commit(snap_id) body = _push_body(objects, [commit], [snap]) t0 = time.perf_counter() r = await _do_push(client, owner, slug, body) elapsed = time.perf_counter() - t0 assert r.get("ok") is True, f"cold_large push failed: {r}" result.record(elapsed, len(body)) return result async def bench_repush_noop(client: AsyncClient, runs: int) -> BenchResult: """Push same 500 objects twice — second push dedup-skips all objects.""" result = BenchResult("repush_noop (500 obj dedup-skip)") n_obj = 500 obj_size = 4096 owner, slug = await _create_repo_api(client, f"bench-repush-{os.urandom(4).hex()}") objects = [] for i in range(n_obj): raw = os.urandom(obj_size - 4) + i.to_bytes(4, "big") oid = blob_id(raw) objects.append((oid, raw)) manifest_base = {f"f{i}.bin": oid for i, (oid, _) in enumerate(objects)} snap_id_1 = blob_id(b"bench-repush-snap-prime-" + os.urandom(4)) snap1 = _make_snapshot(snap_id_1, manifest_base) commit1 = _make_commit(snap_id_1) prime_body = _push_body(objects, [commit1], [snap1]) r0 = await _do_push(client, owner, slug, prime_body) assert r0.get("ok") is True, f"repush prime failed: {r0}" prev_cid = commit1["commit_id"] for run_i in range(runs): snap_id_n = blob_id(f"bench-repush-snap-{run_i}-{os.urandom(4).hex()}".encode()) snap_n = _make_snapshot(snap_id_n, manifest_base) commit_n = _make_commit(snap_id_n, parent_id=prev_cid) body = _push_body(objects, [commit_n], [snap_n]) t0 = time.perf_counter() r = await _do_push(client, owner, slug, body) elapsed = time.perf_counter() - t0 assert r.get("ok") is True, f"repush_noop failed: {r}" prev_cid = commit_n["commit_id"] result.record(elapsed, len(body)) return result async def bench_many_small(client: AsyncClient, runs: int) -> BenchResult: """1 commit, 2000 × 256-byte objects — many-small-objects throughput.""" result = BenchResult("many_small (2000 × 256B)") n_obj = 2000 obj_size = 256 for run_i in range(runs): owner, slug = await _create_repo_api(client, f"bench-small-{run_i}-{os.urandom(4).hex()}") objects = [] for i in range(n_obj): raw = (i.to_bytes(4, "big") + run_i.to_bytes(4, "big")).ljust(obj_size, b"\xab") oid = blob_id(raw) objects.append((oid, raw)) manifest = {f"f{i}.bin": oid for i, (oid, _) in enumerate(objects)} snap_id = blob_id(f"bench-small-snap-{run_i}-{os.urandom(4).hex()}".encode()) snap = _make_snapshot(snap_id, manifest) commit = _make_commit(snap_id) body = _push_body(objects, [commit], [snap]) t0 = time.perf_counter() r = await _do_push(client, owner, slug, body) elapsed = time.perf_counter() - t0 assert r.get("ok") is True, f"many_small push failed: {r}" result.record(elapsed, len(body)) return result async def bench_few_large(client: AsyncClient, runs: int) -> BenchResult: """1 commit, 5 × 500KB objects (~2.5 MB) — large-object server CPU.""" result = BenchResult("few_large (5 × 500KB, ~2.5MB)") n_obj = 5 obj_size = 512 * 1024 for run_i in range(runs): owner, slug = await _create_repo_api(client, f"bench-large-{run_i}-{os.urandom(4).hex()}") objects = [] for i in range(n_obj): raw = os.urandom(obj_size - 8) + i.to_bytes(4, "big") + run_i.to_bytes(4, "big") oid = blob_id(raw) objects.append((oid, raw)) manifest = {f"track{i}.wav": oid for i, (oid, _) in enumerate(objects)} snap_id = blob_id(f"bench-large-snap-{run_i}-{os.urandom(4).hex()}".encode()) snap = _make_snapshot(snap_id, manifest) commit = _make_commit(snap_id) body = _push_body(objects, [commit], [snap]) t0 = time.perf_counter() r = await _do_push(client, owner, slug, body) elapsed = time.perf_counter() - t0 assert r.get("ok") is True, f"few_large push failed: {r}" result.record(elapsed, len(body)) return result async def bench_incremental(client: AsyncClient, runs: int) -> BenchResult: """10-commit seed push; then measure incremental push of 5 new objects.""" result = BenchResult("incremental (seed 50 obj; push +5 obj)") n_commits_seed = 10 obj_per_commit = 5 obj_size = 8192 owner, slug = await _create_repo_api(client, f"bench-inc-{os.urandom(4).hex()}") # Seed push: 10 commits × 5 objects each all_objects: list[tuple[str, bytes]] = [] for i in range(n_commits_seed * obj_per_commit): raw = os.urandom(obj_size - 4) + i.to_bytes(4, "big") oid = blob_id(raw) all_objects.append((oid, raw)) commits_seed = [] snaps_seed = [] prev_cid: str | None = None for ci in range(n_commits_seed): chunk = all_objects[ci * obj_per_commit:(ci + 1) * obj_per_commit] manifest = {f"f{j}.bin": oid for j, (oid, _) in enumerate(chunk)} snap_id = blob_id(f"bench-inc-seed-{ci}-{os.urandom(4).hex()}".encode()) snaps_seed.append(_make_snapshot(snap_id, manifest)) c = _make_commit(snap_id, parent_id=prev_cid) commits_seed.append(c) prev_cid = c["commit_id"] seed_body = _push_body(all_objects, commits_seed, snaps_seed) r_seed = await _do_push(client, owner, slug, seed_body) assert r_seed.get("ok") is True, f"incremental seed failed: {r_seed}" # Incremental push: 5 new objects only for run_i in range(runs): new_objects: list[tuple[str, bytes]] = [] for i in range(obj_per_commit): raw = os.urandom(obj_size - 8) + run_i.to_bytes(4, "big") + i.to_bytes(4, "big") oid = blob_id(raw) new_objects.append((oid, raw)) manifest2 = {f"new{i}.bin": oid for i, (oid, _) in enumerate(new_objects)} snap_id2 = blob_id(f"bench-inc-push2-{run_i}-{os.urandom(4).hex()}".encode()) snap2 = _make_snapshot(snap_id2, manifest2) commit2 = _make_commit(snap_id2, parent_id=prev_cid) body2 = _push_body(new_objects, [commit2], [snap2]) t0 = time.perf_counter() r = await _do_push(client, owner, slug, body2) elapsed = time.perf_counter() - t0 assert r.get("ok") is True, f"incremental push2 failed: {r}" prev_cid = commit2["commit_id"] result.record(elapsed, len(body2)) return result # ── table printer ───────────────────────────────────────────────────────────── def _print_results(results: list[BenchResult]) -> None: col_w = [42, 12, 12, 14, 12] sep = " " header = sep.join(s.ljust(w) for s, w in zip( ["Scenario", "p50 (ms)", "p95 (ms)", "bytes sent", "MB/s"], col_w, )) rule = sep.join("-" * w for w in col_w) print() print("Push benchmark results") print("=" * (sum(col_w) + len(sep) * (len(col_w) - 1))) print(header) print(rule) for r in results: mb = r.bytes_sent / (1024 * 1024) row = [ r.name, f"{r.p50():.1f}", f"{r.p95():.1f}", f"{mb:.3f} MB", f"{r.throughput_mbs():.2f}", ] print(sep.join(s.ljust(w) for s, w in zip(row, col_w))) print() # ── main ────────────────────────────────────────────────────────────────────── async def main(runs: int = 3) -> None: # Schema setup async with _ENGINE.begin() as conn: await conn.run_sync(Base.metadata.drop_all) await conn.run_sync(Base.metadata.create_all) # Temp object storage — same approach as conftest._tmp_objects_dir _tmp = tempfile.mkdtemp(prefix="bench_objects_") import musehub.storage.backends as _backends import musehub.services.musehub_wire as _wire_svc import musehub.api.routes.wire as _wire_route from musehub.config import settings _test_backend = _backends.LocalBackend(objects_dir=_tmp) _wire_svc.get_backend = lambda: _test_backend # type: ignore[method-assign] _wire_route.get_backend = lambda: _test_backend # type: ignore[method-assign] settings.musehub_objects_dir = _tmp # Stub background jobs import musehub.services.musehub_jobs as _jobs async def _noop() -> None: pass _jobs.enqueue_push_intel = _noop # type: ignore[method-assign] _jobs.enqueue_profile_snapshot = _noop # type: ignore[method-assign] # Wire the app's get_db to use our test engine (same pattern as conftest). from typing import AsyncGenerator from sqlalchemy.ext.asyncio import AsyncSession _database._engine = _ENGINE _database._async_session_factory = _SESSION_FACTORY async def _override_get_db() -> AsyncGenerator[AsyncSession, None]: async with _SESSION_FACTORY() as req_session: yield req_session # Inject auth app.dependency_overrides[get_db] = _override_get_db app.dependency_overrides[require_signed_request] = lambda: _AUTH_CTX app.dependency_overrides[optional_signed_request] = lambda: _AUTH_CTX try: async with AsyncClient( transport=ASGITransport(app=app), base_url="https://localhost:1337", ) as client: print(f"\nRunning {runs} repetition(s) per scenario…\n") results = [] for bench_fn in [ bench_tiny, bench_incremental, bench_cold_large, bench_repush_noop, bench_many_small, bench_few_large, ]: print(f" {bench_fn.__name__}…", end="", flush=True) r = await bench_fn(client, runs) results.append(r) print(f" done p50={r.p50():.0f}ms throughput={r.throughput_mbs():.2f} MB/s") _print_results(results) finally: app.dependency_overrides.clear() await _ENGINE.dispose() if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Push protocol benchmark suite — Phase 7") parser.add_argument("--runs", type=int, default=3, help="Repetitions per scenario (default 3)") args = parser.parse_args() asyncio.run(main(runs=args.runs))