gabriel / musehub public
bench_push.py python
494 lines 18.5 KB
Raw
sha256:a34090cc4a394a78bd72cbbe34b08cc59525141e19135b6c0ab154f10611b9ef debug(push/stream): instrument O-frame decode path with INF… Sonnet 4.6 patch 123 days ago
1 """Push protocol benchmark suite — Phase 7.
2
3 Measures wall-clock time, bytes sent, and throughput for the scenarios
4 defined in push-v2.md Phase 7:
5
6 1. tiny — 1 commit, 0 objects
7 2. incremental — 10 commits, 5 objects each; second push sends only new objects
8 3. cold_large — 1 commit, 500 × 4KB objects (~2 MB total)
9 4. repush_noop — same 500-object push repeated (all objects already in store)
10 5. many_small — 1 commit, 2000 × 256-byte objects
11 6. few_large — 1 commit, 5 × 500KB objects
12
13 Run:
14 python3 tests/bench_push.py
15 python3 tests/bench_push.py --runs 5
16
17 Each scenario is repeated RUNS times (default 3). Prints a table with
18 p50, p95 wall-clock (ms), total bytes sent, and throughput (MB/s).
19 """
20 from __future__ import annotations
21
22 import asyncio
23 import os
24 import statistics
25 import sys
26 import tempfile
27 import time
28 from pathlib import Path
29 from typing import Any
30
31 os.environ.setdefault("MUSE_ENV", "test")
32 sys.path.insert(0, str(Path(__file__).parent.parent))
33
34 import msgpack
35 from httpx import AsyncClient, ASGITransport
36 from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
37 from sqlalchemy.pool import NullPool
38
39 from muse.core.types import blob_id, now_utc_iso
40 from musehub.types.json_types import JSONObject, JSONValue
41 from muse.core.mpack import WIRE_CONTENT_TYPE, MuseWireFrameWriter
42 from musehub.db.database import Base, get_db
43 import musehub.db.database as _database
44 import musehub.db.muse_cli_models # noqa: F401 — register all ORM models
45 from musehub.main import app
46 from musehub.auth.request_signing import MSignContext, require_signed_request, optional_signed_request
47 from musehub.models.wire import SFRAME_HEADER, SFRAME_COMMIT_PACK, SFRAME_END, SFRAME_OBJECT
48
49 _fw = MuseWireFrameWriter()
50
51 _DB_URL = os.environ.get(
52 "TEST_DATABASE_URL",
53 "postgresql+asyncpg://musehub:musehub@localhost:5434/musehub_test",
54 )
55 _ENGINE = create_async_engine(_DB_URL, poolclass=NullPool)
56 _SESSION_FACTORY = async_sessionmaker(bind=_ENGINE, expire_on_commit=False)
57
58 _AUTH_CTX = MSignContext(
59 handle="bench-user",
60 identity_id="bench-user-id",
61 is_agent=False,
62 is_admin=False,
63 )
64
65 # ── frame helpers ──────────────────────────────────────────────────────────────
66
67 def _pack(obj: JSONValue) -> bytes:
68 return msgpack.packb(obj, use_bin_type=True)
69
70 def _wrap(ft: str, data: JSONValue) -> bytes:
71 return _fw.wrap(frame_type=ft, payload=_pack(data))
72
73 _commit_counter = 0
74
75 def _make_commit(snapshot_id: str, parent_id: str | None = None) -> JSONObject:
76 global _commit_counter
77 _commit_counter += 1
78 cid = blob_id(f"bench-commit-{_commit_counter}-{os.urandom(4).hex()}".encode())
79 return {
80 "commit_id": cid,
81 "parent_ids": [parent_id] if parent_id else [],
82 "parent_commit_id": parent_id,
83 "parent2_commit_id": None,
84 "snapshot_id": snapshot_id,
85 "branch": "main",
86 "message": "bench commit",
87 "author": "bench",
88 "committed_at": now_utc_iso(),
89 "signature": "",
90 "signer_key_id": "",
91 "agent_id": "",
92 "model_id": "",
93 "metadata": {},
94 }
95
96 def _make_snapshot(snap_id: str, manifest: JSONObject | None = None) -> JSONObject:
97 return {"snapshot_id": snap_id, "manifest": manifest or {}}
98
99 def _h_frame(n_objects: int, n_commits: int, force: bool = False) -> bytes:
100 return _wrap(SFRAME_HEADER, {
101 "t": SFRAME_HEADER, "branch": "main", "force": force,
102 "have": [], "head": "", "n_objects": n_objects, "n_commits": n_commits,
103 })
104
105 def _o_frame(oid: str, raw: bytes) -> bytes:
106 return _wrap(SFRAME_OBJECT, {
107 "t": SFRAME_OBJECT, "id": oid, "content": raw,
108 "enc": "raw", "path": "bench.bin", "sz": len(raw),
109 })
110
111 def _c_frame(commits: list[dict], snapshots: list[dict]) -> bytes:
112 return _wrap(SFRAME_COMMIT_PACK, {
113 "t": SFRAME_COMMIT_PACK, "commits": commits, "snapshots": snapshots,
114 })
115
116 def _e_frame(n_objects: int, n_commits: int) -> bytes:
117 return _wrap(SFRAME_END, {
118 "t": SFRAME_END, "n_objects": n_objects, "n_commits": n_commits,
119 })
120
121 def _push_body(
122 objects: list[tuple[str, bytes]],
123 commits: list[dict],
124 snapshots: list[dict],
125 force: bool = False,
126 ) -> bytes:
127 parts = [_h_frame(len(objects), len(commits), force=force)]
128 for oid, raw in objects:
129 parts.append(_o_frame(oid, raw))
130 parts.append(_c_frame(commits, snapshots))
131 parts.append(_e_frame(len(objects), len(commits)))
132 return b"".join(parts)
133
134
135 # ── bench harness ──────────────────────────────────────────────────────────────
136
137 async def _create_repo_api(client: AsyncClient, name: str) -> tuple[str, str]:
138 """Create a repo via the API and return (owner, slug)."""
139 import musehub.db.database as _db_mod
140 from sqlalchemy.ext.asyncio import AsyncSession
141 async with _SESSION_FACTORY() as session:
142 from tests.factories import create_repo
143 repo = await create_repo(session, owner="bench-user", name=name)
144 await session.commit()
145 return repo.owner, repo.slug
146
147 async def _do_push(client: AsyncClient, owner: str, slug: str, body: bytes) -> JSONObject:
148 resp = await client.post(
149 f"/{owner}/{slug}/push/stream",
150 content=body,
151 headers={"Content-Type": WIRE_CONTENT_TYPE},
152 )
153 unpacker = msgpack.Unpacker(raw=False)
154 unpacker.feed(resp.content)
155 last: JSONObject = {}
156 for frame in unpacker:
157 last = frame
158 return last
159
160
161 class BenchResult:
162 def __init__(self, name: str) -> None:
163 self.name = name
164 self.times_ms: list[float] = []
165 self.bytes_sent: int = 0
166
167 def record(self, elapsed_s: float, body_bytes: int) -> None:
168 self.times_ms.append(elapsed_s * 1000)
169 self.bytes_sent = body_bytes
170
171 def p50(self) -> float:
172 return statistics.median(self.times_ms)
173
174 def p95(self) -> float:
175 n = len(self.times_ms)
176 if n < 2:
177 return self.times_ms[0]
178 return sorted(self.times_ms)[max(0, int(n * 0.95) - 1)]
179
180 def throughput_mbs(self) -> float:
181 p50_s = self.p50() / 1000
182 if p50_s <= 0:
183 return 0.0
184 return (self.bytes_sent / (1024 * 1024)) / p50_s
185
186
187 # ── scenarios ─────────────────────────────────────────────────────────────────
188
189 async def bench_tiny(client: AsyncClient, runs: int) -> BenchResult:
190 """1 commit, 0 objects — measures pure protocol + DB overhead."""
191 result = BenchResult("tiny (1 commit, 0 obj)")
192 owner, slug = await _create_repo_api(client, f"bench-tiny-{os.urandom(4).hex()}")
193
194 prev_cid: str | None = None
195 for i in range(runs):
196 snap_id = blob_id(f"bench-tiny-snap-{i}-{os.urandom(4).hex()}".encode())
197 snap = _make_snapshot(snap_id)
198 commit = _make_commit(snap_id, parent_id=prev_cid)
199 body = _push_body([], [commit], [snap])
200
201 t0 = time.perf_counter()
202 r = await _do_push(client, owner, slug, body)
203 elapsed = time.perf_counter() - t0
204 assert r.get("ok") is True, f"tiny push failed: {r}"
205 prev_cid = commit["commit_id"]
206 result.record(elapsed, len(body))
207
208 return result
209
210
211 async def bench_cold_large(client: AsyncClient, runs: int) -> BenchResult:
212 """1 commit, 500 × 4KB objects (~2 MB) — cold upload throughput."""
213 result = BenchResult("cold_large (500 × 4KB, ~2MB)")
214 n_obj = 500
215 obj_size = 4096
216
217 for run_i in range(runs):
218 owner, slug = await _create_repo_api(client, f"bench-cold-{run_i}-{os.urandom(4).hex()}")
219 objects = []
220 for i in range(n_obj):
221 raw = os.urandom(obj_size - 8) + i.to_bytes(4, "big") + run_i.to_bytes(4, "big")
222 oid = blob_id(raw)
223 objects.append((oid, raw))
224
225 manifest = {f"f{i}.bin": oid for i, (oid, _) in enumerate(objects)}
226 snap_id = blob_id(f"bench-cold-snap-{run_i}-{os.urandom(4).hex()}".encode())
227 snap = _make_snapshot(snap_id, manifest)
228 commit = _make_commit(snap_id)
229 body = _push_body(objects, [commit], [snap])
230
231 t0 = time.perf_counter()
232 r = await _do_push(client, owner, slug, body)
233 elapsed = time.perf_counter() - t0
234 assert r.get("ok") is True, f"cold_large push failed: {r}"
235 result.record(elapsed, len(body))
236
237 return result
238
239
240 async def bench_repush_noop(client: AsyncClient, runs: int) -> BenchResult:
241 """Push same 500 objects twice — second push dedup-skips all objects."""
242 result = BenchResult("repush_noop (500 obj dedup-skip)")
243 n_obj = 500
244 obj_size = 4096
245 owner, slug = await _create_repo_api(client, f"bench-repush-{os.urandom(4).hex()}")
246
247 objects = []
248 for i in range(n_obj):
249 raw = os.urandom(obj_size - 4) + i.to_bytes(4, "big")
250 oid = blob_id(raw)
251 objects.append((oid, raw))
252
253 manifest_base = {f"f{i}.bin": oid for i, (oid, _) in enumerate(objects)}
254 snap_id_1 = blob_id(b"bench-repush-snap-prime-" + os.urandom(4))
255 snap1 = _make_snapshot(snap_id_1, manifest_base)
256 commit1 = _make_commit(snap_id_1)
257 prime_body = _push_body(objects, [commit1], [snap1])
258 r0 = await _do_push(client, owner, slug, prime_body)
259 assert r0.get("ok") is True, f"repush prime failed: {r0}"
260 prev_cid = commit1["commit_id"]
261
262 for run_i in range(runs):
263 snap_id_n = blob_id(f"bench-repush-snap-{run_i}-{os.urandom(4).hex()}".encode())
264 snap_n = _make_snapshot(snap_id_n, manifest_base)
265 commit_n = _make_commit(snap_id_n, parent_id=prev_cid)
266 body = _push_body(objects, [commit_n], [snap_n])
267
268 t0 = time.perf_counter()
269 r = await _do_push(client, owner, slug, body)
270 elapsed = time.perf_counter() - t0
271 assert r.get("ok") is True, f"repush_noop failed: {r}"
272 prev_cid = commit_n["commit_id"]
273 result.record(elapsed, len(body))
274
275 return result
276
277
278 async def bench_many_small(client: AsyncClient, runs: int) -> BenchResult:
279 """1 commit, 2000 × 256-byte objects — many-small-objects throughput."""
280 result = BenchResult("many_small (2000 × 256B)")
281 n_obj = 2000
282 obj_size = 256
283
284 for run_i in range(runs):
285 owner, slug = await _create_repo_api(client, f"bench-small-{run_i}-{os.urandom(4).hex()}")
286 objects = []
287 for i in range(n_obj):
288 raw = (i.to_bytes(4, "big") + run_i.to_bytes(4, "big")).ljust(obj_size, b"\xab")
289 oid = blob_id(raw)
290 objects.append((oid, raw))
291
292 manifest = {f"f{i}.bin": oid for i, (oid, _) in enumerate(objects)}
293 snap_id = blob_id(f"bench-small-snap-{run_i}-{os.urandom(4).hex()}".encode())
294 snap = _make_snapshot(snap_id, manifest)
295 commit = _make_commit(snap_id)
296 body = _push_body(objects, [commit], [snap])
297
298 t0 = time.perf_counter()
299 r = await _do_push(client, owner, slug, body)
300 elapsed = time.perf_counter() - t0
301 assert r.get("ok") is True, f"many_small push failed: {r}"
302 result.record(elapsed, len(body))
303
304 return result
305
306
307 async def bench_few_large(client: AsyncClient, runs: int) -> BenchResult:
308 """1 commit, 5 × 500KB objects (~2.5 MB) — large-object server CPU."""
309 result = BenchResult("few_large (5 × 500KB, ~2.5MB)")
310 n_obj = 5
311 obj_size = 512 * 1024
312
313 for run_i in range(runs):
314 owner, slug = await _create_repo_api(client, f"bench-large-{run_i}-{os.urandom(4).hex()}")
315 objects = []
316 for i in range(n_obj):
317 raw = os.urandom(obj_size - 8) + i.to_bytes(4, "big") + run_i.to_bytes(4, "big")
318 oid = blob_id(raw)
319 objects.append((oid, raw))
320
321 manifest = {f"track{i}.wav": oid for i, (oid, _) in enumerate(objects)}
322 snap_id = blob_id(f"bench-large-snap-{run_i}-{os.urandom(4).hex()}".encode())
323 snap = _make_snapshot(snap_id, manifest)
324 commit = _make_commit(snap_id)
325 body = _push_body(objects, [commit], [snap])
326
327 t0 = time.perf_counter()
328 r = await _do_push(client, owner, slug, body)
329 elapsed = time.perf_counter() - t0
330 assert r.get("ok") is True, f"few_large push failed: {r}"
331 result.record(elapsed, len(body))
332
333 return result
334
335
336 async def bench_incremental(client: AsyncClient, runs: int) -> BenchResult:
337 """10-commit seed push; then measure incremental push of 5 new objects."""
338 result = BenchResult("incremental (seed 50 obj; push +5 obj)")
339 n_commits_seed = 10
340 obj_per_commit = 5
341 obj_size = 8192
342
343 owner, slug = await _create_repo_api(client, f"bench-inc-{os.urandom(4).hex()}")
344
345 # Seed push: 10 commits × 5 objects each
346 all_objects: list[tuple[str, bytes]] = []
347 for i in range(n_commits_seed * obj_per_commit):
348 raw = os.urandom(obj_size - 4) + i.to_bytes(4, "big")
349 oid = blob_id(raw)
350 all_objects.append((oid, raw))
351
352 commits_seed = []
353 snaps_seed = []
354 prev_cid: str | None = None
355 for ci in range(n_commits_seed):
356 chunk = all_objects[ci * obj_per_commit:(ci + 1) * obj_per_commit]
357 manifest = {f"f{j}.bin": oid for j, (oid, _) in enumerate(chunk)}
358 snap_id = blob_id(f"bench-inc-seed-{ci}-{os.urandom(4).hex()}".encode())
359 snaps_seed.append(_make_snapshot(snap_id, manifest))
360 c = _make_commit(snap_id, parent_id=prev_cid)
361 commits_seed.append(c)
362 prev_cid = c["commit_id"]
363
364 seed_body = _push_body(all_objects, commits_seed, snaps_seed)
365 r_seed = await _do_push(client, owner, slug, seed_body)
366 assert r_seed.get("ok") is True, f"incremental seed failed: {r_seed}"
367
368 # Incremental push: 5 new objects only
369 for run_i in range(runs):
370 new_objects: list[tuple[str, bytes]] = []
371 for i in range(obj_per_commit):
372 raw = os.urandom(obj_size - 8) + run_i.to_bytes(4, "big") + i.to_bytes(4, "big")
373 oid = blob_id(raw)
374 new_objects.append((oid, raw))
375
376 manifest2 = {f"new{i}.bin": oid for i, (oid, _) in enumerate(new_objects)}
377 snap_id2 = blob_id(f"bench-inc-push2-{run_i}-{os.urandom(4).hex()}".encode())
378 snap2 = _make_snapshot(snap_id2, manifest2)
379 commit2 = _make_commit(snap_id2, parent_id=prev_cid)
380 body2 = _push_body(new_objects, [commit2], [snap2])
381
382 t0 = time.perf_counter()
383 r = await _do_push(client, owner, slug, body2)
384 elapsed = time.perf_counter() - t0
385 assert r.get("ok") is True, f"incremental push2 failed: {r}"
386 prev_cid = commit2["commit_id"]
387 result.record(elapsed, len(body2))
388
389 return result
390
391
392 # ── table printer ─────────────────────────────────────────────────────────────
393
394 def _print_results(results: list[BenchResult]) -> None:
395 col_w = [42, 12, 12, 14, 12]
396 sep = " "
397 header = sep.join(s.ljust(w) for s, w in zip(
398 ["Scenario", "p50 (ms)", "p95 (ms)", "bytes sent", "MB/s"],
399 col_w,
400 ))
401 rule = sep.join("-" * w for w in col_w)
402 print()
403 print("Push benchmark results")
404 print("=" * (sum(col_w) + len(sep) * (len(col_w) - 1)))
405 print(header)
406 print(rule)
407 for r in results:
408 mb = r.bytes_sent / (1024 * 1024)
409 row = [
410 r.name,
411 f"{r.p50():.1f}",
412 f"{r.p95():.1f}",
413 f"{mb:.3f} MB",
414 f"{r.throughput_mbs():.2f}",
415 ]
416 print(sep.join(s.ljust(w) for s, w in zip(row, col_w)))
417 print()
418
419
420 # ── main ──────────────────────────────────────────────────────────────────────
421
422 async def main(runs: int = 3) -> None:
423 # Schema setup
424 async with _ENGINE.begin() as conn:
425 await conn.run_sync(Base.metadata.drop_all)
426 await conn.run_sync(Base.metadata.create_all)
427
428 # Temp object storage — same approach as conftest._tmp_objects_dir
429 _tmp = tempfile.mkdtemp(prefix="bench_objects_")
430 import musehub.storage.backends as _backends
431 import musehub.services.musehub_wire as _wire_svc
432 import musehub.api.routes.wire as _wire_route
433 from musehub.config import settings
434
435 _test_backend = _backends.LocalBackend(objects_dir=_tmp)
436 _wire_svc.get_backend = lambda: _test_backend # type: ignore[method-assign]
437 _wire_route.get_backend = lambda: _test_backend # type: ignore[method-assign]
438 settings.musehub_objects_dir = _tmp
439
440 # Stub background jobs
441 import musehub.services.musehub_jobs as _jobs
442 async def _noop() -> None:
443 pass
444 _jobs.enqueue_push_intel = _noop # type: ignore[method-assign]
445 _jobs.enqueue_profile_snapshot = _noop # type: ignore[method-assign]
446
447 # Wire the app's get_db to use our test engine (same pattern as conftest).
448 from typing import AsyncGenerator
449 from sqlalchemy.ext.asyncio import AsyncSession
450 _database._engine = _ENGINE
451 _database._async_session_factory = _SESSION_FACTORY
452
453 async def _override_get_db() -> AsyncGenerator[AsyncSession, None]:
454 async with _SESSION_FACTORY() as req_session:
455 yield req_session
456
457 # Inject auth
458 app.dependency_overrides[get_db] = _override_get_db
459 app.dependency_overrides[require_signed_request] = lambda: _AUTH_CTX
460 app.dependency_overrides[optional_signed_request] = lambda: _AUTH_CTX
461
462 try:
463 async with AsyncClient(
464 transport=ASGITransport(app=app),
465 base_url="https://localhost:1337",
466 ) as client:
467 print(f"\nRunning {runs} repetition(s) per scenario…\n")
468 results = []
469 for bench_fn in [
470 bench_tiny,
471 bench_incremental,
472 bench_cold_large,
473 bench_repush_noop,
474 bench_many_small,
475 bench_few_large,
476 ]:
477 print(f" {bench_fn.__name__}…", end="", flush=True)
478 r = await bench_fn(client, runs)
479 results.append(r)
480 print(f" done p50={r.p50():.0f}ms throughput={r.throughput_mbs():.2f} MB/s")
481
482 _print_results(results)
483 finally:
484 app.dependency_overrides.clear()
485
486 await _ENGINE.dispose()
487
488
489 if __name__ == "__main__":
490 import argparse
491 parser = argparse.ArgumentParser(description="Push protocol benchmark suite — Phase 7")
492 parser.add_argument("--runs", type=int, default=3, help="Repetitions per scenario (default 3)")
493 args = parser.parse_args()
494 asyncio.run(main(runs=args.runs))
File History 1 commit
sha256:a34090cc4a394a78bd72cbbe34b08cc59525141e19135b6c0ab154f10611b9ef debug(push/stream): instrument O-frame decode path with INF… Sonnet 4.6 patch 123 days ago