Add Ollama integration: cross-process KV-cache reuse via the Mooncake Store #5
|
|
@ -0,0 +1,247 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Multi-agent KV-reuse workload driver.
|
||||
|
||||
Models the core scenario: several coding agents share a long repo context and
|
||||
differ only in a short instruction. The first agent prefills + stores its KV;
|
||||
the rest restore it from Mooncake and prefill only their unique tail.
|
||||
|
||||
Run a workload twice (cache OFF = baseline, cache ON) and compare honest
|
||||
end-to-end TTFT, recomputed prefill tokens, hit rate, and aggregate throughput.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import argparse, json, statistics, sys, threading, time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import requests
|
||||
|
||||
SESSION = requests.Session()
|
||||
|
||||
|
||||
def tokenize(llama, text, add_special=True):
|
||||
r = SESSION.post(f"{llama}/tokenize", json={"content": text, "add_special": add_special}, timeout=120)
|
||||
r.raise_for_status()
|
||||
return r.json()["tokens"]
|
||||
|
||||
|
||||
def completion(llama, tokens, slot, n_predict):
|
||||
r = SESSION.post(f"{llama}/completion", json={
|
||||
"prompt": tokens, "id_slot": slot, "cache_prompt": True,
|
||||
"n_predict": n_predict, "temperature": 0.0,
|
||||
}, timeout=600)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
|
||||
def bridge_call(bridge, ep, fp, policy, tokens, target, extra=None):
|
||||
body = {"fp": fp, "policy": policy, "tokens": tokens, "target": target}
|
||||
if extra:
|
||||
body.update(extra)
|
||||
r = SESSION.post(f"{bridge}/v1/{ep}", json=body, timeout=600)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
|
||||
def erase_all_slots(llamas, slots_per_server):
|
||||
"""Wipe llama.cpp's *own* per-slot prefix cache so a 'baseline' run truly
|
||||
prefills cold and the cached run's only reuse path is Mooncake. Without this,
|
||||
llama's single-process cache contaminates the cross-process comparison."""
|
||||
for u in llamas:
|
||||
for s in range(slots_per_server):
|
||||
try:
|
||||
SESSION.post(f"{u}/slots/{s}?action=erase", timeout=30)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
|
||||
|
||||
def build_repo_context(llama, ctx_tokens):
|
||||
"""A synthetic but realistic multi-file Go repo context of ~ctx_tokens."""
|
||||
header = ("You are an autonomous coding agent reviewing a Go service repository.\n"
|
||||
"Below is the full source tree. Read it carefully.\n\n")
|
||||
unit = ("// ===== file internal/pkg{i}/service.go =====\n"
|
||||
"package pkg{i}\n"
|
||||
"import (\"context\"; \"fmt\"; \"time\")\n"
|
||||
"// Service{i} coordinates the {i}-th processing stage of the pipeline.\n"
|
||||
"type Service{i} struct {{ store *Store; clock time.Clock; retries int }}\n"
|
||||
"func (s *Service{i}) Handle(ctx context.Context, r Record) (Result, error) {{\n"
|
||||
" if err := r.Validate(); err != nil {{ return Result{{}}, fmt.Errorf(\"pkg{i}: %w\", err) }}\n"
|
||||
" out := s.store.Transform(ctx, r)\n"
|
||||
" return s.store.Aggregate(ctx, out), nil\n}}\n\n")
|
||||
text = header
|
||||
i = 0
|
||||
# grow until token target reached (tokenize incrementally but cheaply)
|
||||
while True:
|
||||
chunk = "".join(unit.format(i=j) for j in range(i, i + 12))
|
||||
if len(tokenize(llama, text + chunk)) > ctx_tokens:
|
||||
break
|
||||
text += chunk
|
||||
i += 12
|
||||
return text
|
||||
|
||||
|
||||
QUERIES = [
|
||||
"\n\nTASK: Find and explain any nil-pointer risks in Service handlers.",
|
||||
"\n\nTASK: Suggest where to add structured logging across the pipeline.",
|
||||
"\n\nTASK: Write a unit test for the Aggregate step of stage 3.",
|
||||
"\n\nTASK: Identify duplicated validation logic and propose a refactor.",
|
||||
"\n\nTASK: Where would a context cancellation be dropped? Fix it.",
|
||||
"\n\nTASK: Propose a retry/backoff policy for transient store errors.",
|
||||
"\n\nTASK: Add metrics counters; list the call sites to instrument.",
|
||||
"\n\nTASK: Review error wrapping for consistency; cite files.",
|
||||
"\n\nTASK: Find concurrency hazards in shared Store access.",
|
||||
"\n\nTASK: Summarize the end-to-end data flow in 5 bullet points.",
|
||||
]
|
||||
|
||||
|
||||
def make_agents(llama_tok, n, ctx_tokens, share_mode, block_size):
|
||||
"""Return n token sequences that share an EXACT, block-aligned token prefix.
|
||||
|
||||
We tokenize the shared repo context once and block-align it, then append
|
||||
each agent's distinct query as separately-tokenized ids. Concatenating token
|
||||
*ids* (not text) guarantees the shared prefix is byte-identical across agents
|
||||
so the chained block hashes match and the cache actually hits.
|
||||
"""
|
||||
base = build_repo_context(llama_tok, ctx_tokens)
|
||||
base_tokens = tokenize(llama_tok, base)
|
||||
aligned = (len(base_tokens) // block_size) * block_size
|
||||
base_tokens = base_tokens[:aligned]
|
||||
agents = []
|
||||
for i in range(n):
|
||||
q = QUERIES[0] if share_mode == "identical" else QUERIES[i % len(QUERIES)]
|
||||
qt = tokenize(llama_tok, q, add_special=False)
|
||||
agents.append(list(base_tokens) + list(qt))
|
||||
return agents, aligned
|
||||
|
||||
|
||||
def run_agent(idx, bridge, llama, slot, tokens, fp, policy, n_predict, cache_on):
|
||||
res = {"idx": idx, "llama": llama, "tokens": len(tokens)}
|
||||
prep_ms = 0.0
|
||||
if cache_on:
|
||||
t = time.perf_counter()
|
||||
p = bridge_call(bridge, "prepare", fp, policy, tokens, {"base_url": llama, "slot_id": slot})
|
||||
prep_ms = (time.perf_counter() - t) * 1e3
|
||||
if p.get("error"):
|
||||
# surface bridge errors instead of crashing the whole workload
|
||||
raise RuntimeError(f"prepare error (agent {idx}, slot {slot}): {p['error']}")
|
||||
res.update(hit=p.get("hit", False), decision=p.get("decision", "miss"),
|
||||
restored=p.get("restored", False),
|
||||
store_get_ms=p.get("store_get_ms", 0.0), bytes=p.get("bytes", 0))
|
||||
t = time.perf_counter()
|
||||
c = completion(llama, tokens, slot, n_predict)
|
||||
comp_ms = (time.perf_counter() - t) * 1e3
|
||||
if cache_on and policy.get("write"):
|
||||
bridge_call(bridge, "commit", fp, policy, tokens, {"base_url": llama, "slot_id": slot},
|
||||
extra={"prefill_n": c["timings"]["prompt_n"], "prefill_ms": c["timings"]["prompt_ms"]})
|
||||
res.update(prepare_ms=prep_ms, prompt_n=c["timings"]["prompt_n"],
|
||||
prompt_ms=c["timings"]["prompt_ms"], comp_wall_ms=comp_ms,
|
||||
ttft_ms=prep_ms + c["timings"]["prompt_ms"])
|
||||
return res
|
||||
|
||||
|
||||
def run_workload(bridge, llamas, slots_per_server, agents, fp, policy, n_predict, cache_on, concurrency):
|
||||
# Each agent gets a UNIQUE (server, slot) so concurrent agents never share a
|
||||
# llama slot. Agent 0 runs alone first (cold: prefill + store); the rest run
|
||||
# concurrently and reuse. With cache off, all run concurrently.
|
||||
results = []
|
||||
|
||||
def place(i):
|
||||
server = llamas[i % len(llamas)]
|
||||
slot = (i // len(llamas)) % slots_per_server
|
||||
return server, slot
|
||||
|
||||
t0 = time.perf_counter()
|
||||
if cache_on and len(agents) > 1:
|
||||
s, sl = place(0)
|
||||
results.append(run_agent(0, bridge, s, sl, agents[0], fp, policy, n_predict, True))
|
||||
rest = list(range(1, len(agents)))
|
||||
else:
|
||||
rest = list(range(len(agents)))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=concurrency) as ex:
|
||||
futs = []
|
||||
for i in rest:
|
||||
s, sl = place(i)
|
||||
futs.append(ex.submit(run_agent, i, bridge, s, sl, agents[i], fp, policy, n_predict, cache_on))
|
||||
for f in futs:
|
||||
results.append(f.result())
|
||||
wall = time.perf_counter() - t0
|
||||
results.sort(key=lambda r: r["idx"])
|
||||
return results, wall
|
||||
|
||||
|
||||
def summarize(results, wall, label):
|
||||
ttfts = [r["ttft_ms"] for r in results]
|
||||
recomputed = sum(r["prompt_n"] for r in results)
|
||||
total_prompt = sum(r["tokens"] for r in results)
|
||||
hits = sum(1 for r in results if r.get("hit"))
|
||||
return {
|
||||
"label": label,
|
||||
"agents": len(results),
|
||||
"wall_s": round(wall, 3),
|
||||
"ttft_ms_mean": round(statistics.mean(ttfts), 1),
|
||||
"ttft_ms_median": round(statistics.median(ttfts), 1),
|
||||
"ttft_ms_p90": round(sorted(ttfts)[int(0.9 * (len(ttfts) - 1))], 1),
|
||||
"prompt_tokens_recomputed": recomputed,
|
||||
"prompt_tokens_total": total_prompt,
|
||||
"recompute_ratio": round(recomputed / max(total_prompt, 1), 4),
|
||||
"hit_agents": hits,
|
||||
"throughput_agents_per_s": round(len(results) / max(wall, 1e-9), 3),
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--bridge", default="http://127.0.0.1:52052")
|
||||
ap.add_argument("--llamas", required=True, help="comma-separated llama server URLs")
|
||||
ap.add_argument("--model-path", required=True)
|
||||
ap.add_argument("--agents", type=int, default=5)
|
||||
ap.add_argument("--ctx-tokens", type=int, default=16000)
|
||||
ap.add_argument("--share-mode", choices=["identical", "shared_prefix"], default="shared_prefix")
|
||||
ap.add_argument("--n-predict", type=int, default=8)
|
||||
ap.add_argument("--block-size", type=int, default=256)
|
||||
ap.add_argument("--kv-type", default="f16", help="must match the llama server --cache-type (f16|q8_0|q4_0)")
|
||||
ap.add_argument("--namespace", default="bench")
|
||||
ap.add_argument("--concurrency", type=int, default=8)
|
||||
ap.add_argument("--slots-per-server", type=int, default=4)
|
||||
ap.add_argument("--replica-num", type=int, default=1)
|
||||
ap.add_argument("--out", default="")
|
||||
args = ap.parse_args()
|
||||
|
||||
llamas = args.llamas.split(",")
|
||||
fp = {"model_path": args.model_path, "kv_type": args.kv_type, "block_size": args.block_size}
|
||||
pol_on = {"enable": True, "namespace": f"{args.namespace}-{args.ctx_tokens}", "read": True, "write": True,
|
||||
"block_size": args.block_size, "replica_num": args.replica_num, "min_prefix_blocks": 1}
|
||||
pol_off = {**pol_on, "enable": False, "read": False, "write": False}
|
||||
|
||||
print(f"building workload: {args.agents} agents, ctx~{args.ctx_tokens}, share={args.share_mode}")
|
||||
agents, aligned = make_agents(llamas[0], args.agents, args.ctx_tokens, args.share_mode, args.block_size)
|
||||
print(f" shared block-aligned prefix: {aligned} tokens; prompts: {[len(a) for a in agents]}")
|
||||
|
||||
# Baseline: cache OFF (erase llama's own slot cache first => true cold prefill)
|
||||
erase_all_slots(llamas, args.slots_per_server)
|
||||
base_res, base_wall = run_workload(args.bridge, llamas, args.slots_per_server, agents, fp, pol_off, args.n_predict, False, args.concurrency)
|
||||
base = summarize(base_res, base_wall, "cache_off")
|
||||
# Cached: cache ON. Erase slots again so the ONLY reuse is via Mooncake.
|
||||
erase_all_slots(llamas, args.slots_per_server)
|
||||
cached_res, cached_wall = run_workload(args.bridge, llamas, args.slots_per_server, agents, fp, pol_on, args.n_predict, True, args.concurrency)
|
||||
cached = summarize(cached_res, cached_wall, "cache_on")
|
||||
|
||||
ttft_red = 100 * (1 - cached["ttft_ms_mean"] / max(base["ttft_ms_mean"], 1e-9))
|
||||
speedup = base["wall_s"] / max(cached["wall_s"], 1e-9)
|
||||
out = {
|
||||
"config": vars(args), "baseline": base, "cached": cached,
|
||||
"ttft_reduction_pct": round(ttft_red, 1), "throughput_speedup_x": round(speedup, 2),
|
||||
"per_agent_cached": cached_res,
|
||||
}
|
||||
print("\n--- RESULT ---")
|
||||
print(f" mean TTFT: {base['ttft_ms_mean']:.0f}ms -> {cached['ttft_ms_mean']:.0f}ms ({ttft_red:.1f}% lower)")
|
||||
print(f" recompute: {base['recompute_ratio']*100:.0f}% -> {cached['recompute_ratio']*100:.0f}% of prompt tokens")
|
||||
print(f" wall/throughput: {base['wall_s']:.2f}s -> {cached['wall_s']:.2f}s ({speedup:.2f}x)")
|
||||
print(f" hit agents: {cached['hit_agents']}/{cached['agents']}")
|
||||
if args.out:
|
||||
with open(args.out, "w") as f:
|
||||
json.dump(out, f, indent=2)
|
||||
print(f" saved {args.out}")
|
||||
return out
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,81 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Microbenchmark the Mooncake Store transfer path (zero-copy put_from/get_into).
|
||||
|
||||
Measures steady-state put/get bandwidth for a registered host buffer over a
|
||||
given protocol (tcp or rdma). This tells the cost-arbiter how fast the store
|
||||
can deliver KV bytes, which decides restore-vs-recompute.
|
||||
|
||||
Usage:
|
||||
python benchmarks/microbench_store.py --protocol tcp --master 127.0.0.1:52061
|
||||
python benchmarks/microbench_store.py --protocol rdma --device mlx5_0 --master 127.0.0.1:52061
|
||||
"""
|
||||
import argparse, ctypes, os, time, mmap as _mmap
|
||||
from mooncake.store import MooncakeDistributedStore, ReplicateConfig
|
||||
|
||||
|
||||
def ptr_of(buf):
|
||||
return ctypes.addressof(ctypes.c_char.from_buffer(buf))
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--master", default=os.environ.get("OMB_STORE_MASTER", "127.0.0.1:52061"))
|
||||
ap.add_argument("--metadata", default="P2PHANDSHAKE")
|
||||
ap.add_argument("--protocol", default="tcp")
|
||||
ap.add_argument("--device", default="")
|
||||
ap.add_argument("--sizes-mb", default="16,128,512")
|
||||
ap.add_argument("--iters", type=int, default=3)
|
||||
ap.add_argument("--segment-gb", type=int, default=16)
|
||||
ap.add_argument("--register", action="store_true", default=True)
|
||||
args = ap.parse_args()
|
||||
|
||||
s = MooncakeDistributedStore()
|
||||
cfg = {
|
||||
"local_hostname": "127.0.0.1",
|
||||
"metadata_server": args.metadata,
|
||||
"global_segment_size": args.segment_gb << 30,
|
||||
"local_buffer_size": 4 << 30,
|
||||
"protocol": args.protocol,
|
||||
"rdma_devices": args.device,
|
||||
"master_server_addr": args.master,
|
||||
}
|
||||
assert s.setup(cfg) == 0, "setup failed"
|
||||
rc = ReplicateConfig(); rc.replica_num = 1
|
||||
|
||||
sizes = [int(x) << 20 for x in args.sizes_mb.split(",")]
|
||||
maxsz = max(sizes)
|
||||
buf = _mmap.mmap(-1, maxsz)
|
||||
src = ctypes.create_string_buffer(b"", maxsz)
|
||||
# fill with a pattern
|
||||
for i in range(0, maxsz, 1 << 20):
|
||||
buf[i:i+16] = os.urandom(16)
|
||||
p = ptr_of(buf)
|
||||
if args.register:
|
||||
r = s.register_buffer(p, maxsz)
|
||||
print(f"register_buffer -> {r}")
|
||||
|
||||
# warm (absorb one-time ~20s first-put cost)
|
||||
t0 = time.perf_counter()
|
||||
s.put_from("omb:bench:warm", p, 1 << 20, rc)
|
||||
print(f"warmup put: {(time.perf_counter()-t0)*1e3:.0f} ms")
|
||||
s.get_into("omb:bench:warm", p, 1 << 20)
|
||||
s.remove("omb:bench:warm", True)
|
||||
|
||||
print(f"\nprotocol={args.protocol} device={args.device!r}")
|
||||
print(f"{'size':>8} {'put_ms':>9} {'put_GBps':>9} {'get_ms':>9} {'get_GBps':>9}")
|
||||
for sz in sizes:
|
||||
put_ms = []; get_ms = []
|
||||
for it in range(args.iters):
|
||||
k = f"omb:bench:{sz}:{it}"
|
||||
t0 = time.perf_counter(); s.put_from(k, p, sz, rc); put_ms.append((time.perf_counter()-t0)*1e3)
|
||||
t0 = time.perf_counter(); s.get_into(k, p, sz); get_ms.append((time.perf_counter()-t0)*1e3)
|
||||
s.remove(k, True)
|
||||
pm = min(put_ms); gm = min(get_ms)
|
||||
print(f"{sz>>20:>6}MB {pm:>9.1f} {sz/(pm/1e3)/1e9:>9.2f} {gm:>9.1f} {sz/(gm/1e3)/1e9:>9.2f}")
|
||||
if args.register:
|
||||
s.unregister_buffer(p)
|
||||
s.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,57 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Verify striped (parallel, batched) RDMA transfer sustains full bandwidth for
|
||||
large KV blobs by splitting into <=chunk slices and using batch_put_from /
|
||||
batch_get_into. Compares single-object vs striped for the same total size."""
|
||||
import argparse, ctypes, os, time, mmap as _mmap
|
||||
from mooncake.store import MooncakeDistributedStore, ReplicateConfig
|
||||
|
||||
|
||||
def ptr_of(buf): return ctypes.addressof(ctypes.c_char.from_buffer(buf))
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--master", default=os.environ.get("OMB_STORE_MASTER", "127.0.0.1:52061"))
|
||||
ap.add_argument("--protocol", default="rdma")
|
||||
ap.add_argument("--device", default="mlx5_0")
|
||||
ap.add_argument("--total-mb", type=int, default=2048)
|
||||
ap.add_argument("--chunk-mb", type=int, default=64)
|
||||
args = ap.parse_args()
|
||||
|
||||
s = MooncakeDistributedStore()
|
||||
assert s.setup({
|
||||
"local_hostname": "127.0.0.1", "metadata_server": "P2PHANDSHAKE",
|
||||
"global_segment_size": 32 << 30, "local_buffer_size": 8 << 30,
|
||||
"protocol": args.protocol, "rdma_devices": args.device,
|
||||
"master_server_addr": args.master,
|
||||
}) == 0
|
||||
rc = ReplicateConfig(); rc.replica_num = 1
|
||||
|
||||
total = args.total_mb << 20
|
||||
chunk = args.chunk_mb << 20
|
||||
buf = _mmap.mmap(-1, total)
|
||||
buf[0:16] = os.urandom(16)
|
||||
p = ptr_of(buf)
|
||||
s.register_buffer(p, total)
|
||||
s.put_from("warm", p, 1 << 20, rc); s.get_into("warm", p, 1 << 20); s.remove("warm", True)
|
||||
|
||||
# single object
|
||||
t0 = time.perf_counter(); s.put_from("big", p, total, rc); put1 = time.perf_counter() - t0
|
||||
t0 = time.perf_counter(); s.get_into("big", p, total); get1 = time.perf_counter() - t0
|
||||
s.remove("big", True)
|
||||
print(f"single {args.total_mb}MB: put {put1*1e3:7.0f}ms {total/put1/1e9:5.2f}GB/s get {get1*1e3:7.0f}ms {total/get1/1e9:5.2f}GB/s")
|
||||
|
||||
# striped
|
||||
n = (total + chunk - 1) // chunk
|
||||
keys = [f"stripe:{i}" for i in range(n)]
|
||||
ptrs = [p + i * chunk for i in range(n)]
|
||||
sizes = [min(chunk, total - i * chunk) for i in range(n)]
|
||||
t0 = time.perf_counter(); s.batch_put_from(keys, ptrs, sizes, rc); putN = time.perf_counter() - t0
|
||||
t0 = time.perf_counter(); s.batch_get_into(keys, ptrs, sizes); getN = time.perf_counter() - t0
|
||||
s.batch_remove(keys) if hasattr(s, "batch_remove") else [s.remove(k, True) for k in keys]
|
||||
print(f"striped {args.total_mb}MB/{args.chunk_mb}MB x{n}: put {putN*1e3:7.0f}ms {total/putN/1e9:5.2f}GB/s get {getN*1e3:7.0f}ms {total/getN/1e9:5.2f}GB/s")
|
||||
s.unregister_buffer(p); s.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,200 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Plot the benchmark matrix into publication-style figures (PNG)."""
|
||||
import argparse, json, os, re
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
# Cohesive figure style shared by every chart.
|
||||
PALETTE = {"baseline": "#c0392b", "mooncake": "#2e6fdb", "reuser": "#27ae60",
|
||||
"accent": "#7a5fb8", "muted": "#9aa0a6"}
|
||||
plt.rcParams.update({
|
||||
"figure.dpi": 140,
|
||||
"font.size": 12,
|
||||
"axes.titlesize": 15,
|
||||
"axes.titleweight": "bold",
|
||||
"axes.labelsize": 12,
|
||||
"axes.spines.top": False,
|
||||
"axes.spines.right": False,
|
||||
"axes.grid": True,
|
||||
"grid.color": "#e6e6e6",
|
||||
"grid.linewidth": 0.8,
|
||||
"axes.axisbelow": True,
|
||||
"legend.frameon": False,
|
||||
"figure.autolayout": True,
|
||||
})
|
||||
|
||||
|
||||
def _human(name):
|
||||
"""A2_7b_30k_6ag -> '7B / 30k / 6 agents'."""
|
||||
m = re.search(r"(\d+\.?\d*)b_(\d+)k?_(\d+)ag", name, re.I)
|
||||
if not m:
|
||||
return name
|
||||
size, ctx, ag = m.groups()
|
||||
return f"{size}B / {ctx}k / {ag} agents"
|
||||
|
||||
|
||||
def fig_ttft(cells, out):
|
||||
names = [_human(c["name"]) for c in cells]
|
||||
off = [c["ttft_off"] for c in cells]
|
||||
on = [c["ttft_on"] for c in cells]
|
||||
reuser = [c.get("reuser_ttft_on", c["ttft_on"]) for c in cells]
|
||||
x = np.arange(len(names)); w = 0.27
|
||||
fig, ax = plt.subplots(figsize=(max(7, 2.0 * len(names)), 4.4))
|
||||
ax.bar(x - w, off, w, label="no sharing (baseline)", color=PALETTE["baseline"])
|
||||
ax.bar(x, on, w, label="Mooncake (all agents)", color=PALETTE["mooncake"])
|
||||
ax.bar(x + w, reuser, w, label="Mooncake (reusing agents only)", color=PALETTE["reuser"])
|
||||
for i, (a, b) in enumerate(zip(off, on)):
|
||||
ax.text(i, max(a, b) * 1.02, f"-{100*(1-b/a):.0f}%", ha="center", fontsize=10, fontweight="bold", color="#222")
|
||||
ax.set_ylabel("mean TTFT (ms)")
|
||||
ax.set_title("Time-to-first-token: KV reuse via Mooncake")
|
||||
ax.set_xticks(x); ax.set_xticklabels(names); ax.legend()
|
||||
fig.savefig(out); print("wrote", out); plt.close(fig)
|
||||
|
||||
|
||||
def fig_recompute(cells, out):
|
||||
names = [_human(c["name"]) for c in cells]
|
||||
off = [c["recompute_off_pct"] for c in cells]
|
||||
on = [c["recompute_on_pct"] for c in cells]
|
||||
x = np.arange(len(names)); w = 0.38
|
||||
fig, ax = plt.subplots(figsize=(max(7, 2.0 * len(names)), 4.4))
|
||||
ax.bar(x - w/2, off, w, label="baseline", color=PALETTE["baseline"])
|
||||
bars = ax.bar(x + w/2, on, w, label="Mooncake", color=PALETTE["mooncake"])
|
||||
for b, v in zip(bars, on):
|
||||
ax.text(b.get_x() + b.get_width()/2, v + 2, f"{v:.0f}%", ha="center", fontsize=10, color="#222")
|
||||
ax.set_ylabel("% of prompt tokens re-prefilled"); ax.set_ylim(0, 108)
|
||||
ax.set_title("Redundant prefill eliminated"); ax.set_xticks(x)
|
||||
ax.set_xticklabels(names); ax.legend()
|
||||
fig.savefig(out); print("wrote", out); plt.close(fig)
|
||||
|
||||
|
||||
def fig_throughput(cells, out):
|
||||
names = [_human(c["name"]) for c in cells]
|
||||
thru = [c["thru_x"] for c in cells]
|
||||
x = np.arange(len(names))
|
||||
fig, ax = plt.subplots(figsize=(max(7, 2.0 * len(names)), 4.2))
|
||||
bars = ax.bar(x, thru, 0.5, color=PALETTE["accent"])
|
||||
ax.axhline(1.0, ls="--", color=PALETTE["muted"], label="parity")
|
||||
for b, t in zip(bars, thru):
|
||||
ax.text(b.get_x() + b.get_width()/2, t + 0.02, f"{t:.2f}x", ha="center", fontsize=11, fontweight="bold")
|
||||
ax.set_ylabel("throughput speedup (x)"); ax.set_title("Aggregate swarm throughput")
|
||||
ax.set_xticks(x); ax.set_xticklabels(names); ax.legend()
|
||||
fig.savefig(out); print("wrote", out); plt.close(fig)
|
||||
|
||||
|
||||
def fig_per_agent(swarm_json, out):
|
||||
"""Prompt tokens recomputed per agent: the pioneer pays the full prefill,
|
||||
every later agent restores and recomputes almost nothing. This is the
|
||||
'only agent 0 prefills' story, shown directly rather than via TTFT (which
|
||||
on the Stage-1 file path can make a reuser look slower than the pioneer)."""
|
||||
d = json.load(open(swarm_json))
|
||||
pa = d["per_agent_cached"]
|
||||
idx = [r["idx"] for r in pa]
|
||||
recomputed = [r["prompt_n"] for r in pa]
|
||||
total = [r["tokens"] for r in pa]
|
||||
colors = [PALETTE["baseline"]] + [PALETTE["reuser"]] * (len(pa) - 1)
|
||||
fig, ax = plt.subplots(figsize=(max(7.5, 1.0 * len(pa)), 4.6))
|
||||
# faint full-prompt reference so the saving is obvious
|
||||
ax.bar(idx, total, color="#ececec", width=0.74, label="prompt size (tokens)")
|
||||
bars = ax.bar(idx, recomputed, color=colors, width=0.74, label="tokens re-prefilled")
|
||||
for i, (b, rc) in enumerate(zip(bars, recomputed)):
|
||||
tag = f"{rc}" if rc < 1000 else f"{rc/1000:.1f}k"
|
||||
ax.text(b.get_x() + b.get_width() / 2, rc + max(total) * 0.02, tag,
|
||||
ha="center", fontsize=10, fontweight="bold", color="#222")
|
||||
ax.annotate("pioneer pays\nthe full prefill", xy=(0, total[0]), xytext=(0.6, total[0] * 0.82),
|
||||
fontsize=10, color=PALETTE["baseline"],
|
||||
arrowprops=dict(arrowstyle="->", color=PALETTE["baseline"]))
|
||||
saved = sum(total) - sum(recomputed)
|
||||
ax.set_xlabel("agent index (0 = pioneer that prefills + stores)")
|
||||
ax.set_ylabel("prompt tokens")
|
||||
ax.set_xticks(idx)
|
||||
ax.set_title(f"Only the pioneer prefills ({saved//1000}k tokens saved)")
|
||||
ax.legend(loc="center right")
|
||||
fig.savefig(out); print("wrote", out); plt.close(fig)
|
||||
|
||||
|
||||
def fig_scaling(scaling_json, out):
|
||||
"""Baseline work grows linearly with the swarm; Mooncake stays almost flat.
|
||||
scaling_json is a list of {agents, baseline_recompute_tokens,
|
||||
cached_recompute_tokens, baseline_wall_s, cached_wall_s}."""
|
||||
rows = sorted(json.load(open(scaling_json)), key=lambda r: r["agents"])
|
||||
n = [r["agents"] for r in rows]
|
||||
base_tok = [r["baseline_recompute_tokens"] / 1e3 for r in rows]
|
||||
cache_tok = [r["cached_recompute_tokens"] / 1e3 for r in rows]
|
||||
fig, ax = plt.subplots(figsize=(7.5, 4.6))
|
||||
ax.plot(n, base_tok, "o-", color=PALETTE["baseline"], lw=2.4, ms=7,
|
||||
label="no sharing (re-prefill every agent)")
|
||||
ax.plot(n, cache_tok, "o-", color=PALETTE["mooncake"], lw=2.4, ms=7,
|
||||
label="Mooncake (prefill once, restore the rest)")
|
||||
ax.fill_between(n, cache_tok, base_tok, color=PALETTE["mooncake"], alpha=0.10)
|
||||
# endpoint callouts at the widest swarm
|
||||
ax.annotate(f"{base_tok[-1]:.0f}k", xy=(n[-1], base_tok[-1]), xytext=(-4, 6),
|
||||
textcoords="offset points", ha="right", fontsize=11,
|
||||
fontweight="bold", color=PALETTE["baseline"])
|
||||
ax.annotate(f"{cache_tok[-1]:.0f}k (flat)", xy=(n[-1], cache_tok[-1]), xytext=(-4, 8),
|
||||
textcoords="offset points", ha="right", fontsize=11,
|
||||
fontweight="bold", color=PALETTE["mooncake"])
|
||||
ax.text(n[len(n)//2], (base_tok[-1] + cache_tok[0]) / 2,
|
||||
f"{base_tok[-1]-cache_tok[-1]:.0f}k tokens\nnever re-prefilled",
|
||||
ha="center", va="center", fontsize=10, color="#555", style="italic")
|
||||
ax.set_xlabel("agents sharing the context")
|
||||
ax.set_ylabel("prompt tokens prefilled (thousands)")
|
||||
ax.set_title("Redundant prefill work vs swarm size (7B, 30k ctx)")
|
||||
ax.set_xticks(n); ax.set_ylim(bottom=0); ax.legend(loc="upper left")
|
||||
fig.savefig(out); print("wrote", out); plt.close(fig)
|
||||
|
||||
|
||||
def fig_arbiter(out, points=None):
|
||||
"""Restore cost vs recompute cost per (model, context); the arbiter picks the
|
||||
cheaper. Demonstrates the loss-free property: it restores only when restoring
|
||||
is actually faster than recomputing. Values are the arbiter's own estimates
|
||||
from measured online rates.
|
||||
|
||||
points: list of {label, restore_ms, recompute_ms, decision}. Defaults to the
|
||||
two measured regimes (7B restores, 1.5B declines)."""
|
||||
if points is None:
|
||||
points = [
|
||||
{"label": "7B / 30k ctx", "restore_ms": 2160, "recompute_ms": 4034, "decision": "RESTORE"},
|
||||
{"label": "1.5B / 8k ctx", "restore_ms": 697, "recompute_ms": 311, "decision": "RECOMPUTE"},
|
||||
]
|
||||
labels = [p["label"] for p in points]
|
||||
restore = [p["restore_ms"] for p in points]
|
||||
recompute = [p["recompute_ms"] for p in points]
|
||||
x = np.arange(len(labels)); w = 0.36
|
||||
fig, ax = plt.subplots(figsize=(max(6.5, 2.6 * len(labels)), 4.4))
|
||||
ax.bar(x - w / 2, restore, w, label="restore from store", color=PALETTE["mooncake"])
|
||||
ax.bar(x + w / 2, recompute, w, label="recompute (prefill)", color=PALETTE["baseline"])
|
||||
for i, p in enumerate(points):
|
||||
chosen = min(p["restore_ms"], p["recompute_ms"])
|
||||
ax.annotate(f"chooses\n{p['decision']}",
|
||||
xy=(i, chosen), xytext=(i, chosen + max(restore + recompute) * 0.08),
|
||||
ha="center", fontsize=10, fontweight="bold", color="#1a7d32")
|
||||
ax.set_ylabel("estimated cost (ms)")
|
||||
ax.set_title("Loss-free arbiter: restore only when it's cheaper")
|
||||
ax.set_xticks(x); ax.set_xticklabels(labels); ax.legend(loc="upper left")
|
||||
fig.savefig(out); print("wrote", out); plt.close(fig)
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--matrix", default="run/matrix_results.json")
|
||||
ap.add_argument("--per-agent", default="", help="a swarm cell json for the per-agent figure")
|
||||
ap.add_argument("--scaling", default="", help="a scaling json for the agents-scaling figure")
|
||||
ap.add_argument("--outdir", default="docs/figures")
|
||||
args = ap.parse_args()
|
||||
os.makedirs(args.outdir, exist_ok=True)
|
||||
m = json.load(open(args.matrix))
|
||||
cells = m["cells"]
|
||||
fig_ttft(cells, os.path.join(args.outdir, "ttft.png"))
|
||||
fig_recompute(cells, os.path.join(args.outdir, "recompute.png"))
|
||||
fig_throughput(cells, os.path.join(args.outdir, "throughput.png"))
|
||||
if args.per_agent and os.path.exists(args.per_agent):
|
||||
fig_per_agent(args.per_agent, os.path.join(args.outdir, "per_agent.png"))
|
||||
if args.scaling and os.path.exists(args.scaling):
|
||||
fig_scaling(args.scaling, os.path.join(args.outdir, "scaling_agents.png"))
|
||||
fig_arbiter(os.path.join(args.outdir, "arbiter.png"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,104 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Run the benchmark matrix and aggregate results.
|
||||
|
||||
Each cell shells out to agent_swarm.py (cache OFF baseline vs cache ON) for a
|
||||
given (model, context, #agents, kv_type) on the currently-running stack, then
|
||||
collects the JSON into a single matrix file consumed by plot.py.
|
||||
|
||||
The store-backend axis (RDMA / TCP / local-file) is swept by restarting the
|
||||
store layer between groups (re-run proxy_start.sh with a different protocol).
|
||||
"""
|
||||
import argparse, json, os, subprocess, sys, time
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
|
||||
def run_cell(py, cell, bridge, out_dir):
|
||||
out = os.path.join(out_dir, f"cell_{cell['name']}.json")
|
||||
# model_path may be a bare filename; resolve it against $MODELS_DIR so the
|
||||
# checked-in config is portable across machines/checkouts.
|
||||
model_path = cell["model_path"]
|
||||
if not os.path.isabs(model_path) and not os.path.exists(model_path):
|
||||
models_dir = os.environ.get("MODELS_DIR", "")
|
||||
if models_dir:
|
||||
cand = os.path.join(models_dir, os.path.basename(model_path))
|
||||
if os.path.exists(cand):
|
||||
model_path = cand
|
||||
cmd = [py, os.path.join(HERE, "agent_swarm.py"),
|
||||
"--bridge", bridge,
|
||||
"--llamas", cell["llamas"],
|
||||
"--model-path", model_path,
|
||||
"--kv-type", cell.get("kv_type", "f16"),
|
||||
"--agents", str(cell["agents"]),
|
||||
"--ctx-tokens", str(cell["ctx_tokens"]),
|
||||
"--share-mode", cell.get("share_mode", "shared_prefix"),
|
||||
"--slots-per-server", str(cell.get("slots_per_server", 4)),
|
||||
"--concurrency", str(cell.get("concurrency", 4)),
|
||||
"--n-predict", str(cell.get("n_predict", 8)),
|
||||
"--namespace", cell["name"],
|
||||
"--out", out]
|
||||
print(f"\n===== cell {cell['name']} =====")
|
||||
print(" ", " ".join(cmd))
|
||||
t0 = time.time()
|
||||
r = subprocess.run(cmd, capture_output=True, text=True)
|
||||
print(r.stdout[-1200:])
|
||||
if r.returncode != 0:
|
||||
print("STDERR:", r.stderr[-1500:])
|
||||
return None
|
||||
with open(out) as f:
|
||||
data = json.load(f)
|
||||
data["wall_total_s"] = round(time.time() - t0, 1)
|
||||
data["store_backend"] = cell.get("store_backend", "rdma")
|
||||
return data
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--config", required=True, help="JSON file with a list of cells")
|
||||
ap.add_argument("--bridge", default="http://127.0.0.1:52052")
|
||||
ap.add_argument("--out", default="run/matrix_results.json")
|
||||
args = ap.parse_args()
|
||||
|
||||
with open(args.config) as f:
|
||||
cfg = json.load(f)
|
||||
out_dir = os.path.dirname(os.path.abspath(args.out))
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
|
||||
results = []
|
||||
for cell in cfg["cells"]:
|
||||
d = run_cell(sys.executable, cell, args.bridge, out_dir)
|
||||
if d:
|
||||
results.append({"name": cell["name"], **summ(d)})
|
||||
matrix = {"cells": results, "raw": cfg}
|
||||
with open(args.out, "w") as f:
|
||||
json.dump(matrix, f, indent=2)
|
||||
|
||||
print("\n================ MATRIX SUMMARY ================")
|
||||
hdr = f"{'cell':<22}{'backend':<8}{'agents':>7}{'TTFT off':>10}{'TTFT on':>10}{'red%':>7}{'thru x':>8}{'recompute%':>12}"
|
||||
print(hdr); print("-" * len(hdr))
|
||||
for r in results:
|
||||
print(f"{r['name']:<22}{r['backend']:<8}{r['agents']:>7}{r['ttft_off']:>10.0f}{r['ttft_on']:>10.0f}"
|
||||
f"{r['ttft_red_pct']:>7.0f}{r['thru_x']:>8.2f}{r['recompute_on_pct']:>12.0f}")
|
||||
print(f"\nsaved {args.out}")
|
||||
|
||||
|
||||
def summ(d):
|
||||
b, c = d["baseline"], d["cached"]
|
||||
reuse = d["per_agent_cached"][1:] if len(d["per_agent_cached"]) > 1 else d["per_agent_cached"]
|
||||
reuse_ttft = sum(r["ttft_ms"] for r in reuse) / max(len(reuse), 1)
|
||||
return {
|
||||
"backend": d.get("store_backend", "rdma"),
|
||||
"agents": c["agents"],
|
||||
"ttft_off": b["ttft_ms_mean"],
|
||||
"ttft_on": c["ttft_ms_mean"],
|
||||
"ttft_red_pct": d["ttft_reduction_pct"],
|
||||
"reuser_ttft_on": round(reuse_ttft, 1),
|
||||
"thru_x": d["throughput_speedup_x"],
|
||||
"recompute_off_pct": round(b["recompute_ratio"] * 100, 1),
|
||||
"recompute_on_pct": round(c["recompute_ratio"] * 100, 1),
|
||||
"hit_agents": c["hit_agents"],
|
||||
}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,42 @@
|
|||
[
|
||||
{
|
||||
"agents": 1,
|
||||
"baseline_recompute_tokens": 28942,
|
||||
"cached_recompute_tokens": 28942,
|
||||
"baseline_wall_s": 8.09,
|
||||
"cached_wall_s": 22.694,
|
||||
"throughput_x": 0.36
|
||||
},
|
||||
{
|
||||
"agents": 2,
|
||||
"baseline_recompute_tokens": 57884,
|
||||
"cached_recompute_tokens": 28953,
|
||||
"baseline_wall_s": 5.849,
|
||||
"cached_wall_s": 6.269,
|
||||
"throughput_x": 0.93
|
||||
},
|
||||
{
|
||||
"agents": 4,
|
||||
"baseline_recompute_tokens": 115768,
|
||||
"cached_recompute_tokens": 28975,
|
||||
"baseline_wall_s": 8.558,
|
||||
"cached_wall_s": 8.594,
|
||||
"throughput_x": 1.0
|
||||
},
|
||||
{
|
||||
"agents": 6,
|
||||
"baseline_recompute_tokens": 173653,
|
||||
"cached_recompute_tokens": 28998,
|
||||
"baseline_wall_s": 12.961,
|
||||
"cached_wall_s": 9.519,
|
||||
"throughput_x": 1.36
|
||||
},
|
||||
{
|
||||
"agents": 8,
|
||||
"baseline_recompute_tokens": 231535,
|
||||
"cached_recompute_tokens": 29018,
|
||||
"baseline_wall_s": 16.469,
|
||||
"cached_wall_s": 10.878,
|
||||
"throughput_x": 1.51
|
||||
}
|
||||
]
|
||||
|
|
@ -0,0 +1,128 @@
|
|||
#!/usr/bin/env python3
|
||||
"""End-to-end smoke test: cross-process / cross-GPU KV reuse via Mooncake.
|
||||
|
||||
Agent A (on llama server #0, GPU X): prepare -> miss; /completion fully
|
||||
prefills a long prompt; commit -> store the KV to Mooncake.
|
||||
|
||||
Agent B (on llama server #1, GPU Y): prepare -> the sidecar fetches the KV from
|
||||
Mooncake and restores it into B's slot; /completion now prefills ~0 tokens.
|
||||
|
||||
We print the prompt-token count actually re-computed (timings.prompt_n) and the
|
||||
prefill wall time for A vs B. B should be dramatically faster, proving the KV
|
||||
crossed processes/GPUs through the store.
|
||||
"""
|
||||
import argparse, json, time, sys
|
||||
import requests
|
||||
|
||||
|
||||
def tokenize(llama, text, add_special=True):
|
||||
r = requests.post(f"{llama}/tokenize", json={"content": text, "add_special": add_special})
|
||||
r.raise_for_status()
|
||||
return r.json()["tokens"]
|
||||
|
||||
|
||||
def completion(llama, tokens, slot, n_predict=8):
|
||||
r = requests.post(f"{llama}/completion", json={
|
||||
"prompt": tokens, "id_slot": slot, "cache_prompt": True,
|
||||
"n_predict": n_predict, "temperature": 0.0,
|
||||
})
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
|
||||
def bridge_call(bridge, ep, fp, policy, tokens, target):
|
||||
r = requests.post(f"{bridge}/v1/{ep}", json={
|
||||
"fp": fp, "policy": policy, "tokens": tokens, "target": target,
|
||||
})
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--bridge", default="http://127.0.0.1:52052")
|
||||
ap.add_argument("--llama-a", default="http://127.0.0.1:52070")
|
||||
ap.add_argument("--llama-b", default="http://127.0.0.1:52071")
|
||||
ap.add_argument("--model-path", required=True)
|
||||
ap.add_argument("--ctx-tokens", type=int, default=8000)
|
||||
ap.add_argument("--namespace", default="smoke")
|
||||
ap.add_argument("--block-size", type=int, default=256)
|
||||
args = ap.parse_args()
|
||||
|
||||
fp = {"model_path": args.model_path, "kv_type": "f16", "block_size": args.block_size}
|
||||
policy = {"enable": True, "namespace": args.namespace, "read": True, "write": True,
|
||||
"block_size": args.block_size, "replica_num": 1, "min_prefix_blocks": 1}
|
||||
|
||||
# Build a long, code-like context prompt of ~ctx-tokens tokens.
|
||||
unit = ("// module {i}: utility helpers for the data pipeline\n"
|
||||
"func process_{i}(records []Record) (Result, error) {{\n"
|
||||
" // validate, transform, and aggregate the {i}-th shard\n"
|
||||
" return aggregate(transform(validate(records))), nil\n}}\n\n")
|
||||
text = "You are reviewing a large Go repository. Here is the source:\n\n"
|
||||
i = 0
|
||||
while len(tokenize(args.llama_a, text)) < args.ctx_tokens:
|
||||
text += unit.format(i=i)
|
||||
i += 1
|
||||
text += "\nSummarize the overall architecture of this repository."
|
||||
tokens = tokenize(args.llama_a, text)
|
||||
print(f"[setup] prompt = {len(tokens)} tokens, block_size={args.block_size} "
|
||||
f"=> {len(tokens)//args.block_size} full blocks\n")
|
||||
|
||||
# ---- Agent A: cold on server #0 ----
|
||||
ta = time.perf_counter()
|
||||
pa = bridge_call(args.bridge, "prepare", fp, policy, tokens,
|
||||
{"base_url": args.llama_a, "slot_id": 0})
|
||||
pa_wall = (time.perf_counter() - ta) * 1e3
|
||||
t0 = time.perf_counter()
|
||||
ca = completion(args.llama_a, tokens, slot=0)
|
||||
a_wall = (time.perf_counter() - t0) * 1e3
|
||||
co = bridge_call(args.bridge, "commit", fp, policy, tokens,
|
||||
{"base_url": args.llama_a, "slot_id": 0})
|
||||
a_ttft = pa_wall + ca['timings']['prompt_ms']
|
||||
print("AGENT A (cold, server #0):")
|
||||
print(f" prepare : hit={pa['hit']} decision={pa['decision']} wall={pa_wall:.1f}ms ({pa.get('reason','')})")
|
||||
print(f" prefill : prompt_n={ca['timings']['prompt_n']} toks, "
|
||||
f"prompt_ms={ca['timings']['prompt_ms']:.1f}, e2e_wall={a_wall:.1f}ms")
|
||||
print(f" commit : stored={co['stored']} blocks={co['stored_blocks']} "
|
||||
f"bytes={co['bytes']} put_ms={co['store_put_ms']:.1f} key=...{co['key'][-24:]}")
|
||||
print(f" >> honest TTFT(A) = prepare {pa_wall:.1f} + prefill {ca['timings']['prompt_ms']:.1f} = {a_ttft:.1f}ms\n")
|
||||
|
||||
# ---- Agent B: warm on server #1 (must restore from the store) ----
|
||||
tb = time.perf_counter()
|
||||
pb = bridge_call(args.bridge, "prepare", fp, policy, tokens,
|
||||
{"base_url": args.llama_b, "slot_id": 0})
|
||||
pb_wall = (time.perf_counter() - tb) * 1e3
|
||||
t0 = time.perf_counter()
|
||||
cb = completion(args.llama_b, tokens, slot=0)
|
||||
b_wall = (time.perf_counter() - t0) * 1e3
|
||||
b_ttft = pb_wall + cb['timings']['prompt_ms']
|
||||
print("AGENT B (warm, server #1 -- different process & GPU):")
|
||||
print(f" prepare : hit={pb['hit']} decision={pb['decision']} restored={pb['restored']} "
|
||||
f"restored_tokens={pb['restored_tokens']} wall={pb_wall:.1f}ms")
|
||||
print(f" store_get_ms={pb['store_get_ms']:.1f} bytes={pb['bytes']} ({pb.get('reason','')})")
|
||||
print(f" prefill : prompt_n={cb['timings']['prompt_n']} toks, "
|
||||
f"prompt_ms={cb['timings']['prompt_ms']:.1f}, e2e_wall={b_wall:.1f}ms")
|
||||
print(f" >> honest TTFT(B) = prepare {pb_wall:.1f} + prefill {cb['timings']['prompt_ms']:.1f} = {b_ttft:.1f}ms\n")
|
||||
|
||||
# ---- verdict ----
|
||||
saved = ca['timings']['prompt_n'] - cb['timings']['prompt_n']
|
||||
ttft_red = 100.0 * (1 - b_ttft / max(a_ttft, 1e-9))
|
||||
print("=" * 70)
|
||||
print(f" prefill tokens recomputed: A={ca['timings']['prompt_n']} -> B={cb['timings']['prompt_n']} "
|
||||
f"(saved {saved} tokens, {100.0*saved/max(ca['timings']['prompt_n'],1):.1f}%)")
|
||||
print(f" HONEST end-to-end TTFT: A={a_ttft:.1f}ms -> B={b_ttft:.1f}ms "
|
||||
f"(reduced {ttft_red:.1f}%)")
|
||||
print(f" (B breakdown: store_get {pb['store_get_ms']:.1f}ms + restore/overhead "
|
||||
f"{pb_wall - pb['store_get_ms']:.1f}ms + prefill {cb['timings']['prompt_ms']:.1f}ms)")
|
||||
print("=" * 70)
|
||||
stats = requests.get(f"{args.bridge}/stats").json()
|
||||
print(f"sidecar stats: hits={stats['hits']} misses={stats['misses']} "
|
||||
f"saved_prefill_tokens={stats['saved_prefill_tokens']} "
|
||||
f"learned_get_gbps={stats['learned_get_gbps']:.1f} learned_prefill_tps={stats['learned_prefill_tps']:.0f}")
|
||||
ok = cb['timings']['prompt_n'] < ca['timings']['prompt_n'] * 0.3 and pb['restored']
|
||||
print("\nRESULT:", "PASS — KV reused across processes/GPUs via Mooncake" if ok else "FAIL")
|
||||
sys.exit(0 if ok else 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,40 @@
|
|||
{
|
||||
"cells": [
|
||||
{
|
||||
"name": "A1_1p5b_8k_5ag",
|
||||
"llamas": "http://127.0.0.1:52070,http://127.0.0.1:52071",
|
||||
"model_path": "qwen2.5-coder-1.5b-instruct-q4_k_m.gguf",
|
||||
"kv_type": "f16",
|
||||
"agents": 5,
|
||||
"ctx_tokens": 8000,
|
||||
"share_mode": "shared_prefix",
|
||||
"slots_per_server": 4,
|
||||
"concurrency": 2,
|
||||
"store_backend": "rdma"
|
||||
},
|
||||
{
|
||||
"name": "A2_7b_30k_6ag",
|
||||
"llamas": "http://127.0.0.1:52072,http://127.0.0.1:52073",
|
||||
"model_path": "qwen2.5-coder-7b-instruct-q4_k_m.gguf",
|
||||
"kv_type": "q8_0",
|
||||
"agents": 6,
|
||||
"ctx_tokens": 30000,
|
||||
"share_mode": "shared_prefix",
|
||||
"slots_per_server": 4,
|
||||
"concurrency": 2,
|
||||
"store_backend": "rdma"
|
||||
},
|
||||
{
|
||||
"name": "A3_7b_16k_8ag",
|
||||
"llamas": "http://127.0.0.1:52072,http://127.0.0.1:52073",
|
||||
"model_path": "qwen2.5-coder-7b-instruct-q4_k_m.gguf",
|
||||
"kv_type": "q8_0",
|
||||
"agents": 8,
|
||||
"ctx_tokens": 16000,
|
||||
"share_mode": "shared_prefix",
|
||||
"slots_per_server": 4,
|
||||
"concurrency": 2,
|
||||
"store_backend": "rdma"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,39 @@
|
|||
#!/usr/bin/env bash
|
||||
# Build (if needed) and start the Go sidecar (bridged).
|
||||
# Usage: bridged_start.sh [mooncake|local]
|
||||
set -euo pipefail
|
||||
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/env.sh"
|
||||
BACKEND="${1:-mooncake}"
|
||||
|
||||
if [ ! -x "$OMB_RUN/bridged" ] || [ "${OMB_REBUILD:-0}" = "1" ]; then
|
||||
echo "[bridged] building"
|
||||
(cd "$WS/bridge" && go build -o "$OMB_RUN/bridged" ./cmd/bridged)
|
||||
fi
|
||||
|
||||
if ss -ltn 2>/dev/null | grep -qE "[:.]$OMB_BRIDGE_GRPC_PORT\b"; then
|
||||
echo "[bridged] already on $OMB_BRIDGE_GRPC_PORT"; exit 0
|
||||
fi
|
||||
|
||||
nohup "$OMB_RUN/bridged" \
|
||||
-grpc-addr "127.0.0.1:$OMB_BRIDGE_GRPC_PORT" \
|
||||
-grpc-unix "$OMB_SOCK/bridged.sock" \
|
||||
-http-addr "127.0.0.1:$OMB_BRIDGE_HTTP_PORT" \
|
||||
-store-backend "$BACKEND" \
|
||||
-store-proxy-addr "127.0.0.1:$OMB_STORE_PROXY_PORT" \
|
||||
-local-store-dir "$OMB_RUN/store-local" \
|
||||
-slot-save-path "$OMB_SLOTS/" \
|
||||
-block-size "${OMB_BLOCK_SIZE:-256}" \
|
||||
-replica-num "${OMB_REPLICA_NUM:-1}" \
|
||||
-min-prefix-blocks "${OMB_MIN_PREFIX_BLOCKS:-1}" \
|
||||
-prior-get-gbps "${OMB_PRIOR_GBPS:-8}" \
|
||||
-prior-prefill-tps "${OMB_PRIOR_TPS:-8000}" \
|
||||
-arb-min-tokens "${OMB_ARB_MIN_TOKENS:-256}" \
|
||||
-arb-safety "${OMB_ARB_SAFETY:-0.9}" \
|
||||
> "$OMB_LOGS/bridged.log" 2>&1 &
|
||||
echo $! > "$OMB_RUN/bridged.pid"
|
||||
for _ in $(seq 1 40); do
|
||||
curl -fsS "http://127.0.0.1:$OMB_BRIDGE_HTTP_PORT/healthz" >/dev/null 2>&1 && {
|
||||
echo "[bridged] up (grpc $OMB_BRIDGE_GRPC_PORT, http $OMB_BRIDGE_HTTP_PORT, backend $BACKEND)"; exit 0; }
|
||||
sleep 0.25
|
||||
done
|
||||
echo "[bridged] FAILED"; tail -25 "$OMB_LOGS/bridged.log"; exit 1
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
#!/usr/bin/env bash
|
||||
# 10-minute demo: bring up the stack and show a swarm of coding agents sharing
|
||||
# one repo's KV cache through Mooncake. Usage: demo.sh [7b|1.5b] [agents]
|
||||
set -euo pipefail
|
||||
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
source "$HERE/env.sh"
|
||||
|
||||
WHICH="${1:-7b}"; AGENTS="${2:-6}"
|
||||
if [ "$WHICH" = "7b" ]; then
|
||||
MODEL="$MODELS_DIR/qwen2.5-coder-7b-instruct-q4_k_m.gguf"; KV=q8_0; CTX=30000; GPUS=(4 5); PORTS=(52072 52073)
|
||||
else
|
||||
MODEL="$MODELS_DIR/qwen2.5-coder-1.5b-instruct-q4_k_m.gguf"; KV=f16; CTX=16000; GPUS=(2 3); PORTS=(52070 52071)
|
||||
fi
|
||||
|
||||
# llama.cpp's -c is the TOTAL KV context shared across the -np parallel slots, so
|
||||
# each slot must be sized for the full shared prompt plus a generation margin.
|
||||
NP=4
|
||||
PERSLOT=$(( CTX + 2048 ))
|
||||
TOTAL_CTX=$(( PERSLOT * NP ))
|
||||
|
||||
echo "==> master"; bash "$HERE/master_start.sh"
|
||||
echo "==> store proxy"; OMB_STORE_SEGMENT_GB=24 OMB_STORE_STAGING_MB="${OMB_STORE_STAGING_MB:-1280}" OMB_STORE_STAGING_COUNT="${OMB_STORE_STAGING_COUNT:-4}" bash "$HERE/proxy_start.sh" rdma
|
||||
echo "==> 2x llama ($WHICH, kv=$KV, ${PERSLOT}tok/slot x$NP)";
|
||||
OMB_LLAMA_KVTYPE="$KV" bash "$HERE/llama_start.sh" d0 "${GPUS[0]}" "${PORTS[0]}" "$MODEL" "$TOTAL_CTX" "$NP" | tail -1
|
||||
OMB_LLAMA_KVTYPE="$KV" bash "$HERE/llama_start.sh" d1 "${GPUS[1]}" "${PORTS[1]}" "$MODEL" "$TOTAL_CTX" "$NP" | tail -1
|
||||
echo "==> sidecar"; OMB_PRIOR_GBPS=2 OMB_PRIOR_TPS=8000 bash "$HERE/bridged_start.sh" mooncake
|
||||
|
||||
echo
|
||||
echo "######################################################################"
|
||||
echo "# $AGENTS coding agents analyse the SAME ~${CTX}-token repo context."
|
||||
echo "# Agent 0 prefills it once and writes its KV to Mooncake; the rest"
|
||||
echo "# restore that KV across processes & GPUs and skip the prefill."
|
||||
echo "######################################################################"
|
||||
echo
|
||||
python "$WS/benchmarks/agent_swarm.py" \
|
||||
--llamas "http://127.0.0.1:${PORTS[0]},http://127.0.0.1:${PORTS[1]}" \
|
||||
--model-path "$MODEL" --kv-type "$KV" \
|
||||
--agents "$AGENTS" --ctx-tokens "$CTX" --share-mode shared_prefix \
|
||||
--slots-per-server 4 --concurrency "${DEMO_CONCURRENCY:-2}" --namespace demo \
|
||||
--out "$OMB_RUN/demo_result.json"
|
||||
|
||||
# Scoreboard: turn the result JSON into a one-glance before/after summary.
|
||||
python - "$OMB_RUN/demo_result.json" <<'PY'
|
||||
import json, sys
|
||||
d = json.load(open(sys.argv[1]))
|
||||
b, c = d["baseline"], d["cached"]
|
||||
base, cached = b["prompt_tokens_recomputed"], c["prompt_tokens_recomputed"]
|
||||
saved = base - cached
|
||||
# Two distinct metrics: total prefill work avoided across the swarm, vs the
|
||||
# fraction of *redundant* prefill removed (one pioneer prefill is unavoidable).
|
||||
pioneer = base // c["agents"] if c["agents"] else base
|
||||
redundant_base = max(base - pioneer, 1)
|
||||
redundant_left = max(cached - pioneer, 0)
|
||||
total_avoided = 100 * (1 - cached / max(base, 1))
|
||||
redundant_killed = 100 * (1 - redundant_left / redundant_base)
|
||||
line = "=" * 64
|
||||
print("\n" + line); print(" SCOREBOARD"); print(line)
|
||||
print(f" agents sharing context : {c['agents']}")
|
||||
print(f" prefill tokens, no sharing : {base:>9,}")
|
||||
print(f" prefill tokens, Mooncake : {cached:>9,} ({saved:,} saved)")
|
||||
print(f" total prefill work avoided : {total_avoided:.0f}%")
|
||||
print(f" redundant prefill killed : {redundant_killed:.0f}% (only the pioneer must prefill)")
|
||||
print(f" mean TTFT : {b['ttft_ms_mean']:.0f} ms -> {c['ttft_ms_mean']:.0f} ms")
|
||||
print(f" swarm throughput : {d['throughput_speedup_x']:.2f}x")
|
||||
print(f" cache hits : {c['hit_agents']}/{c['agents']}")
|
||||
print(line + "\n")
|
||||
PY
|
||||
|
||||
if [ "${DEMO_VERBOSE:-0}" = "1" ]; then
|
||||
echo "sidecar live stats:"
|
||||
curl -fsS "http://127.0.0.1:$OMB_BRIDGE_HTTP_PORT/stats" | python -m json.tool
|
||||
fi
|
||||
echo "Prometheus metrics: http://127.0.0.1:$OMB_BRIDGE_HTTP_PORT/metrics (set DEMO_VERBOSE=1 for raw stats)"
|
||||
echo "Tear down with: bash scripts/stack_down.sh"
|
||||
|
|
@ -0,0 +1,90 @@
|
|||
# shellcheck shell=bash
|
||||
# Central environment for the Ollama x Mooncake KVCache Bus integration.
|
||||
#
|
||||
# Layout philosophy:
|
||||
# * WS -> the *source* tree (this integration dir inside the Mooncake
|
||||
# repo). Resolved from this file's location, so the checkout
|
||||
# can live at any path without edits. Never written to by the
|
||||
# build/run (keeps the git tree clean for a PR).
|
||||
# * OMB_STATE -> all *generated* state: toolchains, venv, caches, llama.cpp /
|
||||
# ollama checkouts, model weights, logs, run dirs. Defaults to
|
||||
# a sibling ".omb-state" OUTSIDE the repo so heavy artifacts
|
||||
# never pollute the tree and never land on '/'. Override with
|
||||
# OMB_STATE_DIR to point at a large data volume.
|
||||
#
|
||||
# Source at the top of every script: source "<dir>/env.sh"
|
||||
|
||||
# --- Source tree root (this integration dir) ---
|
||||
export WS="$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")/.." && pwd)"
|
||||
|
||||
# --- Generated-state root (kept OUTSIDE the source tree) ---
|
||||
# Default: a ".omb-state" directory next to the top of the git repo if we can
|
||||
# find it, else next to WS. Override via OMB_STATE_DIR for a big data volume.
|
||||
if [ -z "${OMB_STATE_DIR:-}" ]; then
|
||||
_repo_root="$(cd "$WS" && git rev-parse --show-toplevel 2>/dev/null || true)"
|
||||
if [ -n "${_repo_root:-}" ]; then
|
||||
OMB_STATE_DIR="$(dirname "$_repo_root")/.omb-state"
|
||||
else
|
||||
OMB_STATE_DIR="$(dirname "$WS")/.omb-state"
|
||||
fi
|
||||
fi
|
||||
export OMB_STATE="$OMB_STATE_DIR"
|
||||
|
||||
# --- Keep ALL caches off '/' and out of the repo ---
|
||||
export TMPDIR="$OMB_STATE/tmp"
|
||||
export PIP_CACHE_DIR="$OMB_STATE/.cache/pip"
|
||||
export XDG_CACHE_HOME="$OMB_STATE/.cache/xdg"
|
||||
export HF_HOME="$OMB_STATE/.cache/hf"
|
||||
export HUGGINGFACE_HUB_CACHE="$OMB_STATE/.cache/hf/hub"
|
||||
|
||||
# --- Go toolchain (installed locally, never system-wide) ---
|
||||
export GOROOT="$OMB_STATE/go"
|
||||
export GOPATH="$OMB_STATE/.cache/gopath"
|
||||
export GOMODCACHE="$OMB_STATE/.cache/gomod"
|
||||
export GOCACHE="$OMB_STATE/.cache/go-build"
|
||||
export GOBIN="$OMB_STATE/.cache/gopath/bin"
|
||||
export GOFLAGS="-mod=mod"
|
||||
export GOTOOLCHAIN="local" # never auto-download a different toolchain
|
||||
|
||||
# --- Python venv ---
|
||||
export VENV="$OMB_STATE/.venv"
|
||||
|
||||
# --- CUDA / GPU. Restrict our processes to a subset of devices. ---
|
||||
export CUDA_HOME="${CUDA_HOME:-/usr/local/cuda-12.4}"
|
||||
export OMB_GPUS="${OMB_GPUS:-2,3,4,5,6,7}"
|
||||
|
||||
# --- llama.cpp build/run ---
|
||||
export LLAMA_DIR="$OMB_STATE/llama.cpp"
|
||||
export LLAMA_BUILD="$LLAMA_DIR/build"
|
||||
export MODELS_DIR="$OMB_STATE/models"
|
||||
export OLLAMA_DIR="$OMB_STATE/ollama"
|
||||
|
||||
# --- Runtime dirs ---
|
||||
export OMB_RUN="$OMB_STATE/run"
|
||||
export OMB_LOGS="$OMB_STATE/run/logs"
|
||||
# Slot save files live on tmpfs (RAM) so the GPU<->host<->store path never
|
||||
# touches a spinning/NVMe disk. Scoped + cleaned after use.
|
||||
export OMB_SLOTS="${OMB_SLOTS_OVERRIDE:-/dev/shm/omb_slots}"
|
||||
export OMB_SOCK="$OMB_STATE/run/sockets"
|
||||
export OMB_STORE_DATA="$OMB_STATE/run/store" # mooncake local store spill
|
||||
|
||||
# --- PATH ---
|
||||
export PATH="$GOROOT/bin:$GOBIN:$VENV/bin:$CUDA_HOME/bin:$PATH"
|
||||
export LD_LIBRARY_PATH="$CUDA_HOME/lib64:$LLAMA_BUILD/bin:${LD_LIBRARY_PATH:-}"
|
||||
|
||||
# --- Default service ports (loopback only, high range to avoid clashes) ---
|
||||
export OMB_BRIDGE_GRPC_PORT="${OMB_BRIDGE_GRPC_PORT:-52051}" # Go sidecar gRPC
|
||||
export OMB_BRIDGE_HTTP_PORT="${OMB_BRIDGE_HTTP_PORT:-52052}" # Go sidecar metrics/HTTP
|
||||
export OMB_STORE_PROXY_PORT="${OMB_STORE_PROXY_PORT:-52060}" # Python store proxy gRPC
|
||||
export OMB_MASTER_PORT="${OMB_MASTER_PORT:-52061}" # mooncake master
|
||||
export OMB_META_PORT="${OMB_META_PORT:-52062}" # mooncake metadata (http)
|
||||
export OMB_LLAMA_PORT="${OMB_LLAMA_PORT:-52070}" # llama-server base port
|
||||
|
||||
mkdir -p "$TMPDIR" "$OMB_LOGS" "$OMB_SLOTS" "$OMB_SOCK" "$OMB_STORE_DATA" \
|
||||
"$PIP_CACHE_DIR" "$GOPATH" "$GOMODCACHE" "$GOCACHE" "$HF_HOME" "$XDG_CACHE_HOME" 2>/dev/null
|
||||
|
||||
# Helper: activate python venv if present
|
||||
if [ -f "$VENV/bin/activate" ]; then
|
||||
# shellcheck disable=SC1091
|
||||
source "$VENV/bin/activate"
|
||||
fi
|
||||
|
|
@ -0,0 +1,39 @@
|
|||
#!/usr/bin/env bash
|
||||
# Generate gRPC stubs for both protos, into the Go module and the Python proxy.
|
||||
# The generated stubs are checked in, so this is only needed after editing a
|
||||
# .proto. Requires protoc + the Go protoc plugins (installed by setup_go.sh).
|
||||
set -euo pipefail
|
||||
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/env.sh"
|
||||
cd "$WS"
|
||||
|
||||
GO_MOD="github.com/mooncake-ai/ollama-mooncake-bridge"
|
||||
GO_OUT="bridge"
|
||||
|
||||
# Wait for Go protoc plugins if a background install is still finishing.
|
||||
for _ in $(seq 1 60); do
|
||||
[ -x "$GOBIN/protoc-gen-go" ] && [ -x "$GOBIN/protoc-gen-go-grpc" ] && break
|
||||
sleep 2
|
||||
done
|
||||
|
||||
echo "[proto] Go stubs (storeproxy + bridge)"
|
||||
protoc -I proto -I bridge/api \
|
||||
--plugin=protoc-gen-go="$GOBIN/protoc-gen-go" \
|
||||
--plugin=protoc-gen-go-grpc="$GOBIN/protoc-gen-go-grpc" \
|
||||
--go_out="$GO_OUT" --go_opt=module="$GO_MOD" \
|
||||
--go-grpc_out="$GO_OUT" --go-grpc_opt=module="$GO_MOD" \
|
||||
proto/storeproxy.proto bridge/api/bridge.proto
|
||||
|
||||
echo "[proto] Python stubs (storeproxy) -> store-proxy/gen"
|
||||
mkdir -p store-proxy/gen
|
||||
python -m grpc_tools.protoc -I proto \
|
||||
--python_out=store-proxy/gen \
|
||||
--grpc_python_out=store-proxy/gen \
|
||||
proto/storeproxy.proto
|
||||
# make 'gen' a package and fix the absolute import grpc_tools emits
|
||||
touch store-proxy/gen/__init__.py
|
||||
sed -i 's/^import storeproxy_pb2 as/from . import storeproxy_pb2 as/' \
|
||||
store-proxy/gen/storeproxy_pb2_grpc.py 2>/dev/null || true
|
||||
|
||||
echo "[proto] done"
|
||||
find "$GO_OUT/internal" -name '*.pb.go' 2>/dev/null
|
||||
ls store-proxy/gen
|
||||
|
|
@ -0,0 +1,55 @@
|
|||
#!/usr/bin/env bash
|
||||
# Start one llama.cpp server instance ("agent worker").
|
||||
# Usage: llama_start.sh <name> <gpu|auto> <port> [model_gguf] [ctx] [np]
|
||||
# auto-pick only scans the devices listed in $OMB_GPUS.
|
||||
set -euo pipefail
|
||||
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/env.sh"
|
||||
|
||||
NAME="${1:?name}"; GPU="${2:-auto}"; PORT="${3:?port}"
|
||||
MODEL="${4:-$MODELS_DIR/qwen2.5-coder-1.5b-instruct-q4_k_m.gguf}"
|
||||
CTX="${5:-40960}"; NP="${6:-1}"
|
||||
|
||||
pick_free_gpu() {
|
||||
# pick a GPU from $OMB_GPUS that is currently near-idle (< 2 GiB used)
|
||||
local used line idx mem
|
||||
for idx in ${OMB_GPUS//,/ }; do
|
||||
mem=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i "$idx" 2>/dev/null | tr -d ' ')
|
||||
if [ -n "$mem" ] && [ "$mem" -lt 2048 ]; then echo "$idx"; return 0; fi
|
||||
done
|
||||
echo "${OMB_GPUS%%,*}" # fallback: first allowed gpu
|
||||
}
|
||||
|
||||
if [ "$GPU" = "auto" ]; then GPU="$(pick_free_gpu)"; fi
|
||||
|
||||
if ss -ltn 2>/dev/null | grep -q "127.0.0.1:$PORT "; then
|
||||
echo "[llama:$NAME] already on $PORT"; exit 0
|
||||
fi
|
||||
|
||||
LOG="$OMB_LOGS/llama_${NAME}.log"
|
||||
KVTYPE="${OMB_LLAMA_KVTYPE:-f16}" # f16 | q8_0 | q4_0 (smaller KV => faster store transfer)
|
||||
EXTRA=()
|
||||
if [ "$KVTYPE" != "f16" ]; then
|
||||
# quantized KV cache requires flash attention in llama.cpp
|
||||
EXTRA+=(--cache-type-k "$KVTYPE" --cache-type-v "$KVTYPE" -fa on)
|
||||
elif [ "${OMB_LLAMA_FA:-0}" = "1" ]; then
|
||||
EXTRA+=(-fa on)
|
||||
fi
|
||||
echo "[llama:$NAME] gpu=$GPU port=$PORT ctx=$CTX np=$NP kv=$KVTYPE model=$(basename "$MODEL")"
|
||||
CUDA_VISIBLE_DEVICES="$GPU" nohup "$LLAMA_BUILD/bin/llama-server" \
|
||||
-m "$MODEL" \
|
||||
--host 127.0.0.1 --port "$PORT" \
|
||||
-ngl 99 -c "$CTX" -np "$NP" -t 8 \
|
||||
--slot-save-path "$OMB_SLOTS/" \
|
||||
--no-webui \
|
||||
"${EXTRA[@]}" \
|
||||
> "$LOG" 2>&1 &
|
||||
echo $! > "$OMB_RUN/llama_${NAME}.pid"
|
||||
|
||||
# wait for health
|
||||
for _ in $(seq 1 240); do
|
||||
if curl -fsS "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
|
||||
echo "[llama:$NAME] healthy on $PORT (gpu $GPU)"; exit 0
|
||||
fi
|
||||
sleep 0.5
|
||||
done
|
||||
echo "[llama:$NAME] FAILED to become healthy"; tail -30 "$LOG"; exit 1
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
#!/usr/bin/env bash
|
||||
# Start the Mooncake master (idempotent). Logs + glog under run/logs.
|
||||
set -euo pipefail
|
||||
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/env.sh"
|
||||
|
||||
if ss -ltn 2>/dev/null | grep -q ":$OMB_MASTER_PORT "; then
|
||||
echo "[master] already listening on $OMB_MASTER_PORT"
|
||||
exit 0
|
||||
fi
|
||||
nohup "$VENV/bin/mooncake_master" \
|
||||
-port "$OMB_MASTER_PORT" \
|
||||
-metrics_port "$((OMB_MASTER_PORT+2))" \
|
||||
-enable_metric_reporting=true \
|
||||
-default_kv_lease_ttl="${OMB_KV_LEASE_TTL:-30000}" \
|
||||
-default_kv_soft_pin_ttl="${OMB_KV_SOFT_PIN_TTL:-600000}" \
|
||||
-log_dir="$OMB_LOGS" \
|
||||
> "$OMB_LOGS/master.log" 2>&1 &
|
||||
echo $! > "$OMB_RUN/master.pid"
|
||||
for _ in $(seq 1 30); do
|
||||
ss -ltn 2>/dev/null | grep -q ":$OMB_MASTER_PORT " && { echo "[master] up on $OMB_MASTER_PORT (pid $(cat "$OMB_RUN/master.pid"))"; exit 0; }
|
||||
sleep 0.3
|
||||
done
|
||||
echo "[master] FAILED to start"; tail -20 "$OMB_LOGS/master.log"; exit 1
|
||||
|
|
@ -0,0 +1,11 @@
|
|||
#!/usr/bin/env bash
|
||||
# Stop ONLY this workspace's processes. We scope every match to $WS so we never
|
||||
# touch any co-located workspace sharing this machine.
|
||||
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/env.sh"
|
||||
# master is bound to our unique port; safe to match on that.
|
||||
pkill -f "mooncake_master -port $OMB_MASTER_PORT" 2>/dev/null || true
|
||||
# everything else: match only command lines that mention our source path.
|
||||
pkill -f "$WS/benchmarks/microbench" 2>/dev/null || true
|
||||
pkill -f "$WS/store-proxy" 2>/dev/null || true
|
||||
rm -f "$OMB_RUN/master.pid"
|
||||
echo "[master] stopped (workspace-scoped)"
|
||||
|
|
@ -0,0 +1,34 @@
|
|||
#!/usr/bin/env bash
|
||||
# Start the Mooncake store proxy (gRPC) wrapping the real distributed store.
|
||||
# Usage: proxy_start.sh [tcp|rdma] (default: $OMB_STORE_PROTOCOL or rdma)
|
||||
set -euo pipefail
|
||||
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/env.sh"
|
||||
|
||||
PROTO="${1:-${OMB_STORE_PROTOCOL:-rdma}}"
|
||||
DEV="${OMB_STORE_DEVICE:-}"
|
||||
if [ "$PROTO" = "rdma" ] && [ -z "$DEV" ]; then DEV="mlx5_0"; fi
|
||||
|
||||
if ss -ltn 2>/dev/null | grep -qE "[:.]$OMB_STORE_PROXY_PORT\b"; then
|
||||
echo "[proxy] already listening on $OMB_STORE_PROXY_PORT"; exit 0
|
||||
fi
|
||||
|
||||
nohup python "$WS/store-proxy/store_proxy.py" \
|
||||
--listen "127.0.0.1:$OMB_STORE_PROXY_PORT" \
|
||||
--backend mooncake \
|
||||
--protocol "$PROTO" --device "$DEV" \
|
||||
--master "127.0.0.1:$OMB_MASTER_PORT" \
|
||||
--metadata "P2PHANDSHAKE" \
|
||||
--global-segment-size "$(( ${OMB_STORE_SEGMENT_GB:-32} << 30 ))" \
|
||||
--local-buffer-size "$(( ${OMB_STORE_BUFFER_GB:-8} << 30 ))" \
|
||||
--stripe-mb "${OMB_STORE_STRIPE_MB:-64}" \
|
||||
--warmup \
|
||||
> "$OMB_LOGS/store_proxy.log" 2>&1 &
|
||||
echo $! > "$OMB_RUN/store_proxy.pid"
|
||||
echo "[proxy] starting (proto=$PROTO dev=$DEV) pid $(cat "$OMB_RUN/store_proxy.pid")"
|
||||
# RDMA memory registration of the segment+buffer can take ~1-2 min the first time.
|
||||
for _ in $(seq 1 360); do
|
||||
ss -ltn 2>/dev/null | grep -qE "[:.]$OMB_STORE_PROXY_PORT\b" && { echo "[proxy] up on $OMB_STORE_PROXY_PORT"; exit 0; }
|
||||
kill -0 "$(cat "$OMB_RUN/store_proxy.pid")" 2>/dev/null || { echo "[proxy] process died"; tail -25 "$OMB_LOGS/store_proxy.log"; exit 1; }
|
||||
sleep 0.5
|
||||
done
|
||||
echo "[proxy] FAILED"; tail -25 "$OMB_LOGS/store_proxy.log"; exit 1
|
||||
|
|
@ -0,0 +1,30 @@
|
|||
#!/usr/bin/env bash
|
||||
# Install a local Go toolchain under $OMB_STATE/go (never touches the system or
|
||||
# the source tree).
|
||||
set -euo pipefail
|
||||
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/env.sh"
|
||||
|
||||
if [ -x "$GOROOT/bin/go" ]; then
|
||||
echo "[go] already installed: $("$GOROOT/bin/go" version)"
|
||||
else
|
||||
mkdir -p "$OMB_STATE"
|
||||
VER="$(curl -fsSL https://go.dev/VERSION?m=text 2>/dev/null | head -1 || echo go1.23.4)"
|
||||
[ -z "$VER" ] && VER="go1.23.4"
|
||||
TARBALL="${VER}.linux-amd64.tar.gz"
|
||||
echo "[go] downloading $TARBALL"
|
||||
curl -fL --retry 3 -o "$TMPDIR/$TARBALL" "https://go.dev/dl/${TARBALL}"
|
||||
rm -rf "$GOROOT"
|
||||
# tarball unpacks to a 'go/' dir; GOROOT is exactly $OMB_STATE/go
|
||||
tar -C "$OMB_STATE" -xzf "$TMPDIR/$TARBALL"
|
||||
rm -f "$TMPDIR/$TARBALL"
|
||||
"$GOROOT/bin/go" version
|
||||
echo "[go] done"
|
||||
fi
|
||||
|
||||
# protoc plugins for regenerating gRPC stubs (the checked-in stubs already build;
|
||||
# these are only needed if you edit a .proto and run scripts/gen_protos.sh).
|
||||
if [ ! -x "$GOBIN/protoc-gen-go" ] || [ ! -x "$GOBIN/protoc-gen-go-grpc" ]; then
|
||||
echo "[go] installing protoc-gen-go / protoc-gen-go-grpc"
|
||||
"$GOROOT/bin/go" install google.golang.org/protobuf/cmd/protoc-gen-go@v1.34.2 || true
|
||||
"$GOROOT/bin/go" install google.golang.org/grpc/cmd/protoc-gen-go-grpc@v1.4.0 || true
|
||||
fi
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
#!/usr/bin/env bash
|
||||
# Clone + build llama.cpp with CUDA (H200 = sm_90) and SHARED libs (needed for cgo Stage-2).
|
||||
set -euo pipefail
|
||||
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/env.sh"
|
||||
|
||||
if [ ! -d "$LLAMA_DIR/.git" ]; then
|
||||
echo "[llama] cloning"
|
||||
git clone --depth 1 https://github.com/ggml-org/llama.cpp "$LLAMA_DIR"
|
||||
fi
|
||||
cd "$LLAMA_DIR"
|
||||
echo "[llama] HEAD: $(git rev-parse --short HEAD)"
|
||||
|
||||
cmake -B build -G Ninja \
|
||||
-DCMAKE_BUILD_TYPE=Release \
|
||||
-DGGML_CUDA=ON \
|
||||
-DCMAKE_CUDA_ARCHITECTURES=90 \
|
||||
-DBUILD_SHARED_LIBS=ON \
|
||||
-DLLAMA_CURL=OFF \
|
||||
-DLLAMA_BUILD_TESTS=OFF \
|
||||
-DLLAMA_BUILD_EXAMPLES=ON \
|
||||
-DLLAMA_BUILD_SERVER=ON
|
||||
|
||||
# Limit build parallelism to a fixed job count rather than all available cores.
|
||||
cmake --build build --target llama-server llama-cli -j 24
|
||||
echo "[llama] server: $LLAMA_BUILD/bin/llama-server"
|
||||
ls -la "$LLAMA_BUILD/bin" | head
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
#!/usr/bin/env bash
|
||||
# Download small GGUF coder models (weights only; KV is computed at runtime).
|
||||
set -euo pipefail
|
||||
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/env.sh"
|
||||
cd "$MODELS_DIR"
|
||||
|
||||
dl() { # name url
|
||||
local out="$1" url="$2"
|
||||
if [ -f "$out" ] && [ "$(stat -c%s "$out")" -gt 100000000 ]; then
|
||||
echo "[model] have $out ($(du -h "$out" | cut -f1))"; return 0
|
||||
fi
|
||||
echo "[model] downloading $out"
|
||||
curl -fL --retry 4 --retry-delay 3 -o "$out.part" "$url"
|
||||
mv "$out.part" "$out"
|
||||
echo "[model] done $out ($(du -h "$out" | cut -f1))"
|
||||
}
|
||||
|
||||
# 1.5B for fast iteration, 7B for the main results (bigger KV => bigger transfer win).
|
||||
dl qwen2.5-coder-1.5b-instruct-q4_k_m.gguf \
|
||||
"https://huggingface.co/Qwen/Qwen2.5-Coder-1.5B-Instruct-GGUF/resolve/main/qwen2.5-coder-1.5b-instruct-q4_k_m.gguf"
|
||||
|
||||
if [ "${1:-}" = "--with-7b" ]; then
|
||||
dl qwen2.5-coder-7b-instruct-q4_k_m.gguf \
|
||||
"https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct-GGUF/resolve/main/qwen2.5-coder-7b-instruct-q4_k_m.gguf"
|
||||
fi
|
||||
ls -la "$MODELS_DIR"
|
||||
|
|
@ -0,0 +1,44 @@
|
|||
#!/usr/bin/env bash
|
||||
# Create a venv on /data2 and install Python deps (mooncake + grpc + bench tooling).
|
||||
set -euo pipefail
|
||||
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/env.sh"
|
||||
|
||||
if [ ! -d "$VENV" ]; then
|
||||
echo "[py] creating venv at $VENV"
|
||||
python3 -m venv "$VENV"
|
||||
fi
|
||||
# shellcheck disable=SC1091
|
||||
source "$VENV/bin/activate"
|
||||
|
||||
python -m pip install --upgrade pip wheel setuptools >/dev/null
|
||||
|
||||
echo "[py] installing core deps"
|
||||
pip install \
|
||||
"mooncake-transfer-engine==0.3.11.post1" \
|
||||
"grpcio==1.62.3" "grpcio-tools==1.62.3" "protobuf>=4.25,<5" \
|
||||
"numpy" "requests" "aiohttp" "rich" "matplotlib" "pyyaml" "prometheus-client"
|
||||
|
||||
echo "[py] verifying mooncake imports"
|
||||
python - <<'PY'
|
||||
import importlib, sys
|
||||
ok = True
|
||||
for m in ("mooncake.engine", "mooncake.store"):
|
||||
try:
|
||||
mod = importlib.import_module(m)
|
||||
print(f" OK {m} -> {getattr(mod,'__file__','?')}")
|
||||
except Exception as e:
|
||||
ok = False
|
||||
print(f" ERR {m}: {e}")
|
||||
try:
|
||||
from mooncake.store import MooncakeDistributedStore, ReplicateConfig # noqa
|
||||
print(" OK MooncakeDistributedStore + ReplicateConfig")
|
||||
except Exception as e:
|
||||
print(f" WARN MooncakeDistributedStore/ReplicateConfig: {e}")
|
||||
try:
|
||||
from mooncake.engine import TransferEngine
|
||||
print(" OK TransferEngine")
|
||||
except Exception as e:
|
||||
print(f" WARN TransferEngine: {e}")
|
||||
sys.exit(0 if ok else 1)
|
||||
PY
|
||||
echo "[py] done"
|
||||
|
|
@ -0,0 +1,20 @@
|
|||
#!/usr/bin/env bash
|
||||
# Tear down ONLY this workspace's services. Workspace-scoped so any co-located
|
||||
# environment on this machine is never touched.
|
||||
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/env.sh"
|
||||
|
||||
# kill by recorded PIDs first
|
||||
for pf in "$OMB_RUN"/llama_*.pid "$OMB_RUN"/bridged.pid "$OMB_RUN"/store_proxy.pid; do
|
||||
[ -f "$pf" ] || continue
|
||||
pid=$(cat "$pf" 2>/dev/null || true)
|
||||
[ -n "${pid:-}" ] && kill "$pid" 2>/dev/null || true
|
||||
rm -f "$pf"
|
||||
done
|
||||
# belt-and-suspenders: match only command lines mentioning our state/source paths
|
||||
pkill -f "$OMB_RUN/bridged" 2>/dev/null || true
|
||||
pkill -f "$WS/store-proxy" 2>/dev/null || true
|
||||
pkill -f "slot-save-path $OMB_SLOTS" 2>/dev/null || true
|
||||
# master is on our unique port
|
||||
pkill -f "mooncake_master -port $OMB_MASTER_PORT" 2>/dev/null || true
|
||||
rm -f "$OMB_RUN/master.pid"
|
||||
echo "[stack] down (workspace-scoped)"
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
#!/usr/bin/env bash
|
||||
# Bring up the full stack: master + store-proxy + N llama servers + sidecar.
|
||||
# Usage: stack_up.sh [num_servers=2] [tcp|rdma=rdma] [model_gguf] [ctx=40960]
|
||||
set -euo pipefail
|
||||
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
source "$HERE/env.sh"
|
||||
|
||||
N="${1:-2}"; PROTO="${2:-rdma}"
|
||||
MODEL="${3:-$MODELS_DIR/qwen2.5-coder-1.5b-instruct-q4_k_m.gguf}"
|
||||
CTX="${4:-40960}"
|
||||
|
||||
echo "==> master"; bash "$HERE/master_start.sh"
|
||||
echo "==> store proxy"; OMB_STORE_PROTOCOL="$PROTO" bash "$HERE/proxy_start.sh" "$PROTO"
|
||||
echo "==> $N llama servers";
|
||||
for i in $(seq 0 $((N-1))); do
|
||||
bash "$HERE/llama_start.sh" "a$i" auto "$((OMB_LLAMA_PORT + i))" "$MODEL" "$CTX" 1
|
||||
done
|
||||
echo "==> sidecar"; bash "$HERE/bridged_start.sh" mooncake
|
||||
|
||||
echo
|
||||
echo "stack is up:"
|
||||
echo " master 127.0.0.1:$OMB_MASTER_PORT"
|
||||
echo " store proxy 127.0.0.1:$OMB_STORE_PROXY_PORT (proto=$PROTO)"
|
||||
for i in $(seq 0 $((N-1))); do echo " llama a$i 127.0.0.1:$((OMB_LLAMA_PORT + i))"; done
|
||||
echo " sidecar gRPC 127.0.0.1:$OMB_BRIDGE_GRPC_PORT"
|
||||
echo " sidecar HTTP 127.0.0.1:$OMB_BRIDGE_HTTP_PORT (/metrics /stats /healthz)"
|
||||
echo
|
||||
curl -fsS "http://127.0.0.1:$OMB_BRIDGE_HTTP_PORT/healthz" && echo
|
||||
Loading…
Reference in New Issue