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