diff --git a/mooncake-integration/ollama/.gitignore b/mooncake-integration/ollama/.gitignore new file mode 100644 index 00000000..8dd38df0 --- /dev/null +++ b/mooncake-integration/ollama/.gitignore @@ -0,0 +1,12 @@ +# Generated state lives OUTSIDE this tree (see scripts/env.sh -> $OMB_STATE). +# These patterns guard against accidental in-tree artifacts only. +__pycache__/ +*.pyc +*.pyo + +# the Stage-2 microbench binary, if built in-place +/bridge/cbridge/omb_kvbench + +# editor / OS +.DS_Store +*.swp diff --git a/mooncake-integration/ollama/README.md b/mooncake-integration/ollama/README.md new file mode 100644 index 00000000..6d79d350 --- /dev/null +++ b/mooncake-integration/ollama/README.md @@ -0,0 +1,173 @@ +# Ollama × Mooncake — global KV-cache reuse for local agent swarms + +> **6 coding agents. One shared 30k-token repo context. 1 prefill. 5 RDMA restores. 0 redundant prefill.** +> +> `100% of redundant prefill eliminated` · `prefill work stays flat as the swarm grows` · `−69% TTFT for an uncontended cross-GPU restore, −34% swarm mean at 7B` · measured on 8×H200 + +Let a swarm of local coding agents share one long-prefix KV cache. The first +agent prefills a large repo/context once and publishes its KV to the +[Mooncake](https://github.com/kvcache-ai/Mooncake) Store; every other agent — in +a different process, on a different GPU, or on a different node — **restores that +KV instead of re-prefilling it**. + +![Redundant prefill work vs swarm size](docs/figures/scaling_agents.png) + +*The whole point in one chart: without sharing, every agent re-prefills the same +context, so prefill work grows linearly with the swarm. With Mooncake, the context +is prefilled **once** and restored by everyone else — the work stays flat no matter +how many agents join. Real measurement, 7B model, 30k-token context, 1→8 agents.* + +Ollama has no KV-cache sharing of its own: a `grep` of its tree for +`slot-save-path` / `cache-reuse` / `state_seq` is empty, so prefix reuse is +confined to a single process's RAM and dies with the request. Mooncake already +pools and shares KV for vLLM / SGLang / TRT-LLM in the data centre; this +integration extends that reuse down to the workstation/edge multi-agent +scenario, with a sidecar design that needs almost no change to Ollama itself. + +Built and measured on real hardware: 8× NVIDIA H200, Mellanox RDMA (mlx5), the +real `mooncake-transfer-engine` distributed store over RDMA, a CUDA build of +llama.cpp, and Qwen2.5-Coder 1.5B / 7B GGUF models. + +## Results (measured) + +* **Redundant prefill is eliminated, and the win compounds with the swarm.** At 8 + agents sharing a 30k context, a no-sharing baseline re-prefills **231k tokens**; + Mooncake prefills **29k once** and restores the rest — prefill work stays flat + no matter how many agents join (only the pioneer prefills; everyone else + restores). This is a property of the mechanism, not of timing, and holds in + every run. +* **Throughput / TTFT track store bandwidth.** A single uncontended cross-GPU + restore cuts **time-to-first-token by ~69%** (5.99 s → 1.83 s); across a + 6-agent swarm the **mean TTFT drops ~34%** at 7B/30k (and ~33% at 7B/16k, 1.25× + throughput). The gain grows with model size and context length, and the single + store client + single NIC serialize concurrent restores — so the swarm mean + trails the single-reuser number until the store is scaled out. +* A **restore-vs-recompute arbiter** that learns store bandwidth and prefill + rate online and **never sustains a loss**: it restores when it helps and falls + back to recompute when the store is too slow to beat a local prefill. +* **35–44 GB/s** striped zero-copy KV transfer over RDMA (uncontended). +* Stage-2 in-process state export is **4–5× faster** than the file path, and the + `ON_DEVICE` flag keeps KV on the GPU for GPUDirect RDMA (zero host copy). + +Full methodology and figures: **[docs/REPORT.md](docs/REPORT.md)**. + +## What's new here + +| Contribution | What it is | Why it matters | +|---|---|---| +| **Cross-process KV bus** | a KV cache shared across processes, GPUs and nodes via the Mooncake Store | brings data-centre KV pooling to local agent swarms; Ollama had none | +| **Content-addressed chained keys** | model fingerprint + forward-chained per-block hash | exact longest-prefix matching; KV is never mixed across models/tokenizers | +| **Radix prefix index** | cross-process analogue of RadixAttention over remote KV snapshots | O(blocks) longest-prefix lookup + per-prefix hotness for replication | +| **Restore-vs-recompute arbiter** | learns store bandwidth and prefill rate online, picks the cheaper | makes sharing **loss-free** — safe to enable unconditionally | +| **Striped RDMA + Stage-2 ON_DEVICE** | pre-registered staging pool; cgo `state_seq_*_ext` keeps KV on the GPU | sustains 35–44 GB/s; targets GPUDirect with zero host copy | + +## Architecture (3 decoupled processes) + +![Architecture](docs/figures/architecture.png) + +``` +agents / patched Ollama ─► bridge (Go sidecar) ─► store-proxy (Python) ─► Mooncake Store + key + radix index + arbiter 1 warm store client, RDMA / GPUDirect + 3-stage Lookup/Prepare/Commit striped batch RDMA, DRAM pool, replicas + │ pre-registered staging + llama.cpp /slots (Stage 1) + cgo _ext ON_DEVICE (Stage 2) +``` + +* **`bridge/`** (Go) — the brain. Pure Go, no heavy deps. Serves a gRPC + `KVCacheBus` (Lookup/Prepare/Commit) **and** an HTTP/JSON gateway with a + Prometheus `/metrics` endpoint. +* **`store-proxy/`** (Python) — owns the single warm + `mooncake.store.MooncakeDistributedStore` handle (the official client) and + exposes it over gRPC. KV blobs are passed as **file paths**, so a multi-GiB + blob is copied at most once and never streams through Go. +* **llama.cpp server** — unmodified for Stage 1; driven via its `/slots` + save/restore endpoints. + +## Repo layout + +``` +bridge/ Go sidecar (the core deliverable) + internal/cachekey/ GGUF parser + chained block-hash cache keys + internal/prefixindex/ cross-process radix tree of KV prefixes + internal/arbiter/ restore-vs-recompute cost arbiter (online learning) + internal/orchestrator/ 3-stage Lookup → Prepare(restore) → Commit(save) + internal/store/ pluggable backend: mooncake (proxy) | local file + internal/llamabridge/ Stage-1 /slots HTTP client + internal/seqstate/ Stage-2 cgo binding (llama_state_seq_*_ext, ON_DEVICE) + internal/server/ gRPC KVCacheBus + HTTP/JSON + /metrics + api/bridge.proto agent ↔ sidecar contract + cbridge/omb_kvbench.cpp Stage-2 libllama microbenchmark +store-proxy/ Python gRPC wrapper over MooncakeDistributedStore +proto/storeproxy.proto sidecar ↔ store-proxy contract +ollama-patches/ diffs to ollama/ollama (+ README, integration notes) +benchmarks/ agent_swarm.py, run_matrix.py, plot.py, micro-benchmarks +deploy/ docker-compose, prometheus.yml, grafana dashboard +scripts/ env.sh + setup/start/stop/demo scripts (state kept off-tree) +docs/ REPORT.md, figures/, result JSONs +``` + +`scripts/env.sh` keeps **all** build/runtime state (Go toolchain, venv, llama.cpp, +models, logs) under `$OMB_STATE` (default: a `.omb-state/` dir *outside* the repo), +so the source tree stays clean and nothing is written under `/`. + +## Quickstart + +```bash +# 1. one-time setup (local Go, venv + mooncake, llama.cpp CUDA build, models) +bash scripts/setup_go.sh +bash scripts/setup_py.sh +bash scripts/setup_llama.sh +bash scripts/setup_models.sh --with-7b + +# 2. run the multi-agent demo (brings up master + store-proxy + 2 llama + sidecar) +bash scripts/demo.sh 7b 6 + +# 3. (optional) the full benchmark matrix + figures +python benchmarks/run_matrix.py --config benchmarks/workloads/matrix.json --out "$OMB_STATE/run/matrix_results.json" +python benchmarks/plot.py --matrix "$OMB_STATE/run/matrix_results.json" \ + --per-agent "$OMB_STATE/run/demo_result.json" --scaling benchmarks/scaling.json + +# 4. (optional) Stage-2 KV-state microbenchmark +bash bridge/cbridge/build_kvbench.sh +CUDA_VISIBLE_DEVICES=0 "$OMB_STATE/run/omb_kvbench" "$OMB_STATE/models/qwen2.5-coder-7b-instruct-q4_k_m.gguf" 16000 + +# tear everything down (only this integration's processes) +bash scripts/stack_down.sh +``` + +The demo ends on a one-glance scoreboard (real run, 6 agents / 30k ctx, concurrency 2): + +``` +================================================================ + SCOREBOARD +================================================================ + agents sharing context : 6 + prefill tokens, no sharing : 173,653 + prefill tokens, Mooncake : 28,998 (144,655 saved) + total prefill work avoided : 83% + redundant prefill killed : 100% (only the pioneer must prefill) + mean TTFT : 5868 ms -> 3898 ms + swarm throughput : 1.01x + cache hits : 5/6 +================================================================ +``` + +*Two honest numbers: 83% of **all** swarm prefill work is avoided, and 100% of the +**redundant** prefill is gone — the one pioneer prefill is unavoidable. A single +uncontended cross-GPU restore cuts TTFT ~69% (5.99 s → 1.83 s); the swarm mean +trails it because one store client + one NIC serialize concurrent restores.* + +Observability: the sidecar exposes Prometheus at `http://127.0.0.1:52052/metrics` +(`mooncake_bridge_saved_prefill_tokens_total` is the primary series); a Grafana +dashboard is in `deploy/grafana_dashboard.json`. + +## Design notes + +* **Correctness first:** any difference in model / tokenizer / RoPE / KV dtype / + layout lands in a different key space — KV is never mixed across models. The + chained block hash makes longest-prefix matching exact. +* **No disk in the hot path:** slot files live on `/dev/shm`; KV moves + file↔store at most once via RDMA into pre-registered, pinned buffers. +* **Scoped runtime:** every cache/build/model lives under this directory, and + all process management is scoped to this workspace. diff --git a/mooncake-integration/ollama/benchmarks/agent_swarm.py b/mooncake-integration/ollama/benchmarks/agent_swarm.py new file mode 100644 index 00000000..aaa5ec4f --- /dev/null +++ b/mooncake-integration/ollama/benchmarks/agent_swarm.py @@ -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() diff --git a/mooncake-integration/ollama/benchmarks/microbench_store.py b/mooncake-integration/ollama/benchmarks/microbench_store.py new file mode 100644 index 00000000..3274503b --- /dev/null +++ b/mooncake-integration/ollama/benchmarks/microbench_store.py @@ -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() diff --git a/mooncake-integration/ollama/benchmarks/microbench_striped.py b/mooncake-integration/ollama/benchmarks/microbench_striped.py new file mode 100644 index 00000000..acc74d5b --- /dev/null +++ b/mooncake-integration/ollama/benchmarks/microbench_striped.py @@ -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() diff --git a/mooncake-integration/ollama/benchmarks/plot.py b/mooncake-integration/ollama/benchmarks/plot.py new file mode 100644 index 00000000..eea3d424 --- /dev/null +++ b/mooncake-integration/ollama/benchmarks/plot.py @@ -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() diff --git a/mooncake-integration/ollama/benchmarks/run_matrix.py b/mooncake-integration/ollama/benchmarks/run_matrix.py new file mode 100644 index 00000000..66675538 --- /dev/null +++ b/mooncake-integration/ollama/benchmarks/run_matrix.py @@ -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() diff --git a/mooncake-integration/ollama/benchmarks/scaling.json b/mooncake-integration/ollama/benchmarks/scaling.json new file mode 100644 index 00000000..f07991ab --- /dev/null +++ b/mooncake-integration/ollama/benchmarks/scaling.json @@ -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 + } +] \ No newline at end of file diff --git a/mooncake-integration/ollama/benchmarks/smoke_e2e.py b/mooncake-integration/ollama/benchmarks/smoke_e2e.py new file mode 100644 index 00000000..8dd7b7dd --- /dev/null +++ b/mooncake-integration/ollama/benchmarks/smoke_e2e.py @@ -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() diff --git a/mooncake-integration/ollama/benchmarks/workloads/matrix.json b/mooncake-integration/ollama/benchmarks/workloads/matrix.json new file mode 100644 index 00000000..9960dcaf --- /dev/null +++ b/mooncake-integration/ollama/benchmarks/workloads/matrix.json @@ -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" + } + ] +} \ No newline at end of file diff --git a/mooncake-integration/ollama/bridge/api/bridge.proto b/mooncake-integration/ollama/bridge/api/bridge.proto new file mode 100644 index 00000000..26d55f08 --- /dev/null +++ b/mooncake-integration/ollama/bridge/api/bridge.proto @@ -0,0 +1,150 @@ +syntax = "proto3"; + +// KVCacheBus is the contract between a caller (a coding Agent, the bench +// harness, or patched Ollama) and the Go sidecar. It implements the three-stage +// reuse flow: Lookup (longest-prefix match) -> Prepare (Load/restore the matched +// KV into a llama.cpp slot) -> Commit (Save the produced KV back to the Mooncake +// Store). +package bridge.v1; + +option go_package = "github.com/mooncake-ai/ollama-mooncake-bridge/internal/bridge/pb;bridgepb"; + +service KVCacheBus { + rpc Health(HealthRequest) returns (HealthReply); + // Lookup is read-only: longest block-prefix match + arbiter decision, no I/O + // into llama.cpp. Useful for inspection / benchmarking. + rpc Lookup(LookupRequest) returns (LookupReply); + // Prepare = Lookup + (if a hit and the cost arbiter approves) restore the + // matched KV prefix into target.slot so llama.cpp only prefills the tail. + rpc Prepare(PrepareRequest) returns (PrepareReply); + // Commit = Save the KV currently held in target.slot and Put it to the store. + rpc Commit(CommitRequest) returns (CommitReply); + rpc Stats(StatsRequest) returns (StatsReply); +} + +// ModelFingerprint: every field that changes KV bytes. A mismatch => miss. +message ModelFingerprint { + string model_digest = 1; + string arch = 2; + string tokenizer_hash = 3; + string rope_hash = 4; + string kv_type = 5; // f16 | q8_0 | q4_0 | bf16 + string kv_layout = 6; // dense | swa + int32 n_ctx_train = 7; + int32 block_size = 8; + // Optional path to the GGUF file. If set, the sidecar parses it once to fill + // any empty arch/tokenizer/rope/ctx fields (and a content-addressed digest), + // so callers only need to know the model file. + string model_path = 9; +} + +// CachePolicy mirrors the user-visible options.mooncake.* request extension. +message CachePolicy { + bool enable = 1; + string namespace = 2; // e.g. "repo:my-org/my-repo@main" + bool read = 3; + bool write = 4; + // granularity and transport are accepted for forward compatibility. Block + // granularity is the implemented mode; the store transport (tcp|rdma) is + // chosen once by the store-proxy at startup rather than per request. + string granularity = 5; // block (implemented) | sequence (reserved) + int32 block_size = 6; // overrides fingerprint block_size if > 0 + string transport = 7; // advisory; effective transport is set on the proxy + uint32 replica_num = 8; // hot-prefix replication + bool soft_pin = 9; // keep hot prefixes resident + int32 min_prefix_blocks = 10; // admission: don't store prefixes shorter than this +} + +// LlamaTarget identifies a running llama.cpp server slot to restore into / read +// from. base_url like "http://127.0.0.1:52070"; slot_id is the /slots id. +message LlamaTarget { + string base_url = 1; + int32 slot_id = 2; +} + +message LookupRequest { + ModelFingerprint fp = 1; + CachePolicy policy = 2; + repeated int32 tokens = 3; +} +message LookupReply { + bool hit = 1; + int32 matched_blocks = 2; + int32 matched_tokens = 3; + int32 total_blocks = 4; + int32 total_tokens = 5; + string decision = 6; // restore | recompute | miss + string matched_key = 7; + string reason = 8; // arbiter explanation + string error = 9; +} + +message PrepareRequest { + ModelFingerprint fp = 1; + CachePolicy policy = 2; + repeated int32 tokens = 3; + LlamaTarget target = 4; // if unset, behaves like Lookup (plan only) +} +message PrepareReply { + bool hit = 1; + int32 matched_blocks = 2; + int32 matched_tokens = 3; + int32 total_blocks = 4; + int32 total_tokens = 5; + string decision = 6; + string matched_key = 7; + bool restored = 8; + int32 restored_tokens = 9; + double restore_ms = 10; // llama /slots restore time + double store_get_ms = 11; // store GetFile time + uint64 bytes = 12; + string reason = 13; + string error = 14; +} + +message CommitRequest { + ModelFingerprint fp = 1; + CachePolicy policy = 2; + repeated int32 tokens = 3; // tokens resident in the slot (the prefill prefix) + LlamaTarget target = 4; + // Optional: the caller's observed prefill so the arbiter learns the live + // prefill rate (tokens/s) for restore-vs-recompute decisions. + int32 prefill_n = 5; + double prefill_ms = 6; +} +message CommitReply { + bool ok = 1; + bool stored = 2; // false if skipped (already present / below threshold) + int32 stored_blocks = 3; + int32 stored_tokens = 4; + uint64 bytes = 5; + double save_ms = 6; // llama /slots save time + double store_put_ms = 7; // store PutFile time + string key = 8; + string reason = 9; + string error = 10; +} + +message HealthRequest {} +message HealthReply { + bool ok = 1; + string version = 2; + string store_backend = 3; + bool store_ok = 4; + string protocol = 5; + string detail = 6; +} + +message StatsRequest {} +message StatsReply { + uint64 prepare_total = 1; + uint64 hits = 2; + uint64 misses = 3; + uint64 hit_blocks = 4; + uint64 miss_blocks = 5; + uint64 saved_prefill_tokens = 6; // primary effectiveness metric + uint64 restore_count = 7; + uint64 commit_count = 8; + uint64 bytes_get = 9; + uint64 bytes_put = 10; +} diff --git a/mooncake-integration/ollama/bridge/cbridge/build_kvbench.sh b/mooncake-integration/ollama/bridge/cbridge/build_kvbench.sh new file mode 100644 index 00000000..33e3f98a --- /dev/null +++ b/mooncake-integration/ollama/bridge/cbridge/build_kvbench.sh @@ -0,0 +1,16 @@ +#!/usr/bin/env bash +# Build the Stage-2 KV-state microbenchmark against the locally-built libllama. +set -euo pipefail +source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/../../scripts/env.sh" + +SRC="$WS/bridge/cbridge/omb_kvbench.cpp" +OUT="$OMB_RUN/omb_kvbench" +INC="$LLAMA_DIR/include" +GGML_INC="$LLAMA_DIR/ggml/include" +LIBDIR="$LLAMA_BUILD/bin" + +g++ -std=c++17 -O2 -o "$OUT" "$SRC" \ + -I"$INC" -I"$GGML_INC" \ + -L"$LIBDIR" -lllama -lggml -lggml-base \ + -Wl,-rpath,"$LIBDIR" +echo "built: $OUT" diff --git a/mooncake-integration/ollama/bridge/cbridge/omb_kvbench.cpp b/mooncake-integration/ollama/bridge/cbridge/omb_kvbench.cpp new file mode 100644 index 00000000..0a5c91df --- /dev/null +++ b/mooncake-integration/ollama/bridge/cbridge/omb_kvbench.cpp @@ -0,0 +1,147 @@ +// omb_kvbench — Stage-2 KV-state microbenchmark (libllama, in-process). +// +// Demonstrates, in-process against libllama, the three ways to get a sequence's +// KV out of a running model and why the Stage-2 path matters: +// +// (A) llama_state_seq_save_file(...) -- the Stage-1 path the sidecar +// uses today via /slots: GPU +// -> host -> serialize -> file +// (B) llama_state_seq_get_data_ext(..., NONE) -- raw host export: GPU -> host +// (C) llama_state_seq_get_size_ext(..., ON_DEVICE) -- the Stage-2 target: the +// KV stays in device buffers, +// ready for Mooncake Transfer +// Engine GPUDirect RDMA with +// NO host copy (avoids the +// double-copy of llama.cpp +// issue #8915). +// +// It also verifies correctness: export seq 0, import into seq 1, decode one +// token from each and confirm the KV round-trips. +// +// Build: see cbridge/build_kvbench.sh. Run: omb_kvbench [n_prompt] [n_ctx] [ngl] +#include "llama.h" + +#include +#include +#include +#include +#include +#include +#include + +using clk = std::chrono::high_resolution_clock; +static double ms(clk::time_point a, clk::time_point b) { + return std::chrono::duration(b - a).count(); +} + +int main(int argc, char ** argv) { + if (argc < 2) { fprintf(stderr, "usage: %s [n_prompt=4000] [n_ctx] [ngl=99]\n", argv[0]); return 2; } + const char * model_path = argv[1]; + int n_prompt = argc > 2 ? atoi(argv[2]) : 4000; + // KV cells are shared across sequences; we use two seqs (export/import), so + // size the context for two full copies of the prompt. + int n_ctx = argc > 3 ? atoi(argv[3]) : (n_prompt + 512) * 2; + int ngl = argc > 4 ? atoi(argv[4]) : 99; + + llama_backend_init(); + llama_model_params mparams = llama_model_default_params(); + mparams.n_gpu_layers = ngl; + llama_model * model = llama_model_load_from_file(model_path, mparams); + if (!model) { fprintf(stderr, "model load failed\n"); return 1; } + + llama_context_params cparams = llama_context_default_params(); + cparams.n_ctx = n_ctx; + cparams.n_seq_max = 2; + cparams.n_batch = 2048; + llama_context * ctx = llama_init_from_model(model, cparams); + if (!ctx) { fprintf(stderr, "ctx init failed\n"); return 1; } + const llama_vocab * vocab = llama_model_get_vocab(model); + + // Synthetic code-like prompt, tokenized. + std::string text; + while ((int) text.size() < n_prompt * 5) + text += "func process(ctx Context, d []Record) (Result, error) { return agg(transform(validate(d))) }\n"; + std::vector toks(text.size() + 16); + int n = llama_tokenize(vocab, text.c_str(), (int) text.size(), toks.data(), (int) toks.size(), true, false); + if (n <= 0) { fprintf(stderr, "tokenize failed: %d\n", n); return 1; } + if (n > n_prompt) n = n_prompt; + toks.resize(n); + + // Prefill seq 0 in <= n_batch chunks. + int n_batch_sz = (int) cparams.n_batch; + auto t0 = clk::now(); + for (int start = 0; start < n; start += n_batch_sz) { + int cnt = (n - start < n_batch_sz) ? (n - start) : n_batch_sz; + llama_batch b = llama_batch_init(cnt, 0, 1); + for (int i = 0; i < cnt; i++) { + b.token[i] = toks[start + i]; b.pos[i] = start + i; + b.n_seq_id[i] = 1; b.seq_id[i][0] = 0; + b.logits[i] = (start + i == n - 1); + } + b.n_tokens = cnt; + if (llama_decode(ctx, b) != 0) { fprintf(stderr, "decode failed at %d\n", start); return 1; } + llama_batch_free(b); + } + double prefill_ms = ms(t0, clk::now()); + + // KV byte/token (host state size / tokens). + size_t sz_host = llama_state_seq_get_size_ext(ctx, 0, LLAMA_STATE_SEQ_FLAGS_NONE); + size_t sz_dev = llama_state_seq_get_size_ext(ctx, 0, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); + + // (B) raw host export GPU->host + std::vector buf(sz_host); + t0 = clk::now(); + size_t got = llama_state_seq_get_data_ext(ctx, buf.data(), buf.size(), 0, LLAMA_STATE_SEQ_FLAGS_NONE); + double host_get_ms = ms(t0, clk::now()); + + // (A) Stage-1 file path (what /slots save does): write to tmpfs + const char * fpath = "/dev/shm/omb_kvbench_seq0.bin"; + t0 = clk::now(); + size_t fsaved = llama_state_seq_save_file(ctx, fpath, 0, toks.data(), toks.size()); + double file_save_ms = ms(t0, clk::now()); + + // import (host) into seq 1 + t0 = clk::now(); + size_t set = llama_state_seq_set_data_ext(ctx, buf.data(), got, 1, LLAMA_STATE_SEQ_FLAGS_NONE); + double host_set_ms = ms(t0, clk::now()); + + // correctness: decode one token after the prefix in BOTH seqs, compare argmax. + auto next_logits = [&](int seq) -> const float * { + llama_batch b = llama_batch_init(1, 0, 1); + b.token[0] = toks.back(); b.pos[0] = n; b.n_seq_id[0] = 1; b.seq_id[0][0] = seq; b.logits[0] = 1; b.n_tokens = 1; + llama_decode(ctx, b); + const float * lg = llama_get_logits_ith(ctx, 0); + llama_batch_free(b); + return lg; + }; + int n_vocab = llama_vocab_n_tokens(vocab); + const float * l0 = next_logits(0); + std::vector l0c(l0, l0 + n_vocab); + const float * l1 = next_logits(1); + auto argmax = [&](const float * l) { int a = 0; for (int i = 1; i < n_vocab; i++) if (l[i] > l[a]) a = i; return a; }; + int a0 = argmax(l0c.data()), a1 = argmax(l1); + bool ok = (set > 0) && (a0 == a1); + + double bpt = (double) sz_host / n; + printf("\n==== omb_kvbench: %s ====\n", model_path); + printf("prompt tokens : %d (prefill %.1f ms, %.0f tok/s)\n", n, prefill_ms, n / (prefill_ms / 1e3)); + printf("KV state size (host) : %.1f MiB (%.0f bytes/token)\n", sz_host / 1048576.0, bpt); + printf("KV state size (ondev) : %.1f MiB\n", sz_dev / 1048576.0); + printf("\n--- export paths (the Stage-1 vs Stage-2 comparison) ---\n"); + printf("(A) /slots file save : %8.1f ms (%.2f GB/s) GPU->host->serialize->tmpfs\n", + file_save_ms, fsaved / (file_save_ms / 1e3) / 1e9); + printf("(B) host get_data_ext : %8.1f ms (%.2f GB/s) GPU->host (one copy)\n", + host_get_ms, got / (host_get_ms / 1e3) / 1e9); + printf("(C) ON_DEVICE export : (stays on device; hand the device buffer to\n"); + printf(" Mooncake TE registerLocalMemory for GPUDirect\n"); + printf(" RDMA -- zero host copy, the Stage-2 target)\n"); + printf("import set_data_ext : %8.1f ms\n", host_set_ms); + printf("\nfile-save overhead vs raw host copy : %.2fx slower\n", file_save_ms / host_get_ms); + printf("KV round-trip correctness (seq0==seq1 argmax): %s\n", ok ? "PASS" : "FAIL"); + + remove(fpath); + llama_free(ctx); + llama_model_free(model); + llama_backend_free(); + return ok ? 0 : 1; +} diff --git a/mooncake-integration/ollama/bridge/cmd/bridged/main.go b/mooncake-integration/ollama/bridge/cmd/bridged/main.go new file mode 100644 index 00000000..055bf7ad --- /dev/null +++ b/mooncake-integration/ollama/bridge/cmd/bridged/main.go @@ -0,0 +1,151 @@ +// Command bridged is the ollama-mooncake-bridge sidecar daemon. It serves the +// KVCacheBus over gRPC (TCP + optional Unix socket) and an HTTP/JSON gateway +// (with Prometheus /metrics), backed by either the real Mooncake Store (via the +// Python store proxy) or a local filesystem store. +package main + +import ( + "context" + "flag" + "log" + "net" + "net/http" + "os" + "os/signal" + "syscall" + "time" + + bridgepb "github.com/mooncake-ai/ollama-mooncake-bridge/internal/bridge/pb" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/arbiter" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/metrics" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/orchestrator" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/server" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/store" + "github.com/prometheus/client_golang/prometheus" + "google.golang.org/grpc" +) + +func main() { + var ( + grpcAddr = flag.String("grpc-addr", "127.0.0.1:52051", "gRPC listen address") + grpcUnix = flag.String("grpc-unix", "", "optional gRPC unix socket path") + httpAddr = flag.String("http-addr", "127.0.0.1:52052", "HTTP/JSON + /metrics listen address") + backend = flag.String("store-backend", "mooncake", "store backend: mooncake | local") + proxyAddr = flag.String("store-proxy-addr", "127.0.0.1:52060", "mooncake store proxy gRPC address") + localDir = flag.String("local-store-dir", "./run/store-local", "local backend directory") + slotSavePath = flag.String("slot-save-path", "./run/slots/", "llama.cpp --slot-save-path (shared FS, trailing slash)") + blockSize = flag.Int("block-size", 256, "default cache block size (tokens)") + replicaNum = flag.Uint("replica-num", 1, "default store replica count") + minPrefix = flag.Int("min-prefix-blocks", 1, "do not cache prefixes shorter than this many blocks") + cleanup = flag.Bool("cleanup-files", true, "delete local slot files after store transfer") + priorGBps = flag.Float64("prior-get-gbps", 8.0, "arbiter prior: store read GB/s") + priorTPS = flag.Float64("prior-prefill-tps", 8000, "arbiter prior: prefill tokens/s") + minTokens = flag.Int("arb-min-tokens", 256, "arbiter: minimum matched tokens to bother restoring") + safety = flag.Float64("arb-safety", 0.9, "arbiter: restore only if est_restore < safety*est_recompute") + ) + flag.Parse() + log.SetFlags(log.LstdFlags | log.Lmicroseconds) + + // Store backend. + var be store.Backend + var err error + switch *backend { + case "local": + be, err = store.NewLocalBackend(*localDir) + case "mooncake": + be, err = store.DialMooncake(*proxyAddr) + default: + log.Fatalf("unknown store backend %q", *backend) + } + if err != nil { + log.Fatalf("store backend init: %v", err) + } + defer be.Close() + + // Probe store health (non-fatal: log and continue so /healthz reflects it). + { + ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second) + if hi, herr := be.Health(ctx); herr != nil { + log.Printf("WARNING store health: %v", herr) + } else { + log.Printf("store backend=%s protocol=%s master=%s device=%s", hi.Backend, hi.Protocol, hi.Master, hi.Device) + } + cancel() + } + + reg := prometheus.NewRegistry() + mx := metrics.New(reg) + arb := arbiter.New(arbiter.Config{ + PriorGetGBps: *priorGBps, PriorPrefillToksS: *priorTPS, PriorRestoreMsFix: 8, + MinTokens: *minTokens, SafetyMargin: *safety, Alpha: 0.3, + }) + orch := orchestrator.New(be, arb, mx, orchestrator.Config{ + SlotSavePath: *slotSavePath, DefaultBlockSize: *blockSize, + DefaultReplicaNum: uint32(*replicaNum), MinPrefixBlocks: *minPrefix, CleanupFiles: *cleanup, + }) + + // gRPC server. + gs := grpc.NewServer( + grpc.MaxRecvMsgSize(256<<20), + grpc.MaxSendMsgSize(256<<20), + ) + bridgepb.RegisterKVCacheBusServer(gs, server.NewGRPC(orch, be, mx)) + + var listeners []net.Listener + tl, err := net.Listen("tcp", *grpcAddr) + if err != nil { + log.Fatalf("grpc tcp listen: %v", err) + } + listeners = append(listeners, tl) + if *grpcUnix != "" { + os.Remove(*grpcUnix) + ul, uerr := net.Listen("unix", *grpcUnix) + if uerr != nil { + log.Fatalf("grpc unix listen: %v", uerr) + } + listeners = append(listeners, ul) + } + for _, l := range listeners { + go func(l net.Listener) { + log.Printf("gRPC KVCacheBus on %s", l.Addr()) + if serr := gs.Serve(l); serr != nil { + log.Printf("grpc serve(%s) stopped: %v", l.Addr(), serr) + } + }(l) + } + + // HTTP gateway + metrics. + hs := server.NewHTTP(orch, be, mx, reg) + httpSrv := &http.Server{Addr: *httpAddr, Handler: hs.Mux()} + go func() { + log.Printf("HTTP/JSON + /metrics on %s", *httpAddr) + if herr := httpSrv.ListenAndServe(); herr != nil && herr != http.ErrServerClosed { + log.Printf("http serve stopped: %v", herr) + } + }() + + // periodic gauge refresh so /metrics shows learned rates even when idle + stop := make(chan struct{}) + go func() { + t := time.NewTicker(2 * time.Second) + defer t.Stop() + for { + select { + case <-stop: + return + case <-t.C: + orch.RefreshGauges() + } + } + }() + + sig := make(chan os.Signal, 1) + signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM) + <-sig + log.Printf("shutting down...") + close(stop) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + _ = httpSrv.Shutdown(ctx) + gs.GracefulStop() +} diff --git a/mooncake-integration/ollama/bridge/go.mod b/mooncake-integration/ollama/bridge/go.mod new file mode 100644 index 00000000..e62898cb --- /dev/null +++ b/mooncake-integration/ollama/bridge/go.mod @@ -0,0 +1,21 @@ +module github.com/mooncake-ai/ollama-mooncake-bridge + +go 1.23 + +require ( + github.com/prometheus/client_golang v1.19.1 + github.com/prometheus/client_model v0.5.0 + google.golang.org/grpc v1.64.1 + google.golang.org/protobuf v1.34.2 +) + +require ( + github.com/beorn7/perks v1.0.1 // indirect + github.com/cespare/xxhash/v2 v2.2.0 // indirect + github.com/prometheus/common v0.48.0 // indirect + github.com/prometheus/procfs v0.12.0 // indirect + golang.org/x/net v0.26.0 // indirect + golang.org/x/sys v0.21.0 // indirect + golang.org/x/text v0.16.0 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20240318140521-94a12d6c2237 // indirect +) diff --git a/mooncake-integration/ollama/bridge/go.sum b/mooncake-integration/ollama/bridge/go.sum new file mode 100644 index 00000000..a2c5717e --- /dev/null +++ b/mooncake-integration/ollama/bridge/go.sum @@ -0,0 +1,24 @@ +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= +github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= +github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= +github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/prometheus/client_golang v1.19.1 h1:wZWJDwK+NameRJuPGDhlnFgx8e8HN3XHQeLaYJFJBOE= +github.com/prometheus/client_golang v1.19.1/go.mod h1:mP78NwGzrVks5S2H6ab8+ZZGJLZUq1hoULYBAYBw1Ho= +github.com/prometheus/client_model v0.5.0 h1:VQw1hfvPvk3Uv6Qf29VrPF32JB6rtbgI6cYPYQjL0Qw= +github.com/prometheus/client_model v0.5.0/go.mod h1:dTiFglRmd66nLR9Pv9f0mZi7B7fk5Pm3gvsjB5tr+kI= +github.com/prometheus/common v0.48.0 h1:QO8U2CdOzSn1BBsmXJXduaaW+dY/5QLjfB8svtSzKKE= +github.com/prometheus/common v0.48.0/go.mod h1:0/KsvlIEfPQCQ5I2iNSAWKPZziNCvRs5EC6ILDTlAPc= +github.com/prometheus/procfs v0.12.0 h1:jluTpSng7V9hY0O2R9DzzJHYb2xULk9VTR1V1R/k6Bo= +github.com/prometheus/procfs v0.12.0/go.mod h1:pcuDEFsWDnvcgNzo4EEweacyhjeA9Zk3cnaOZAZEfOo= +golang.org/x/net v0.26.0 h1:soB7SVo0PWrY4vPW/+ay0jKDNScG2X9wFeYlXIvJsOQ= +golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE= +golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws= +golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4= +golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI= +google.golang.org/genproto/googleapis/rpc v0.0.0-20240318140521-94a12d6c2237 h1:NnYq6UN9ReLM9/Y01KWNOWyI5xQ9kbIms5GGJVwS/Yc= +google.golang.org/genproto/googleapis/rpc v0.0.0-20240318140521-94a12d6c2237/go.mod h1:WtryC6hu0hhx87FDGxWCDptyssuo68sk10vYjF+T9fY= +google.golang.org/grpc v1.64.1 h1:LKtvyfbX3UGVPFcGqJ9ItpVWW6oN/2XqTxfAnwRRXiA= +google.golang.org/grpc v1.64.1/go.mod h1:hiQF4LFZelK2WKaP6W0L92zGHtiQdZxk8CrSdvyjeP0= +google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg= +google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw= diff --git a/mooncake-integration/ollama/bridge/internal/arbiter/arbiter.go b/mooncake-integration/ollama/bridge/internal/arbiter/arbiter.go new file mode 100644 index 00000000..74f26528 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/arbiter/arbiter.go @@ -0,0 +1,213 @@ +// Package arbiter decides, per request, whether reusing a cached KV prefix is +// actually worth it — i.e. whether *fetching* the KV from the store is cheaper +// than *recomputing* it with a prefill. +// +// Why this matters: a KV cache only helps if the store can deliver KV bytes +// faster than the GPU regenerates them during prefill. The break-even bandwidth +// is +// +// B* = (KV bytes per token) * (prefill tokens/sec) +// +// For the models we test that is ~0.6–0.85 GB/s. Over slow TCP loopback we +// measured ~0.5 GB/s (so recompute wins); over RDMA ~30–40 GB/s (so restore +// wins by ~50x). A naive "always restore on hit" policy would *lose* on slow +// transports. The arbiter instead estimates both times from quantities it +// learns online and picks the cheaper — making the system robust on any fabric. +// +// All rates are tracked with EWMAs fed by real observations from the +// orchestrator, so the estimates self-calibrate to the live hardware. +package arbiter + +import ( + "fmt" + "sync" +) + +// ewma is a thread-safe exponentially-weighted moving average. The FIRST real +// sample *replaces* the prior (seed-on-first) so the estimate adapts to live +// hardware within a single observation instead of slowly decaying from a guess. +type ewma struct { + mu sync.Mutex + alpha float64 + val float64 + n uint64 +} + +func newEWMA(alpha, prior float64) *ewma { return &ewma{alpha: alpha, val: prior} } + +func (e *ewma) update(sample float64) { + if sample <= 0 { + return + } + e.mu.Lock() + if e.n == 0 { + e.val = sample // seed on first real observation + } else { + e.val = e.alpha*sample + (1-e.alpha)*e.val + } + e.n++ + e.mu.Unlock() +} + +func (e *ewma) get() float64 { + e.mu.Lock() + defer e.mu.Unlock() + return e.val +} + +// Config holds priors and policy knobs. +type Config struct { + // Priors used before any observation. + PriorGetGBps float64 // store read bandwidth, GB/s + PriorPrefillToksS float64 // prefill throughput, tokens/s + PriorRestoreMsFix float64 // fixed restore overhead (file write + llama load), ms + // Policy. + MinTokens int // never bother below this many matched tokens + SafetyMargin float64 // restore only if est_restore < margin * est_recompute (e.g. 0.85) + Alpha float64 // EWMA smoothing +} + +func DefaultConfig() Config { + return Config{ + PriorGetGBps: 2.0, // effective end-to-end restore bandwidth (store fetch + GPU load) + PriorPrefillToksS: 8000., // typical small-model prefill on a modern GPU + PriorRestoreMsFix: 8.0, + MinTokens: 256, + SafetyMargin: 0.9, + Alpha: 0.3, + } +} + +// modelState holds the learned rates for one (model, kv-type). Keeping them +// per-model means a fast tiny model and a slow large model don't pollute each +// other's restore-vs-recompute decision. +type modelState struct { + getGBps *ewma // effective end-to-end restore bandwidth (store fetch + GPU load) + prefillTPS *ewma // prefill throughput + restoreFix *ewma // fixed restore overhead, ms +} + +// Arbiter is safe for concurrent use. +type Arbiter struct { + cfg Config + mu sync.Mutex + states map[string]*modelState +} + +func New(cfg Config) *Arbiter { + if cfg.Alpha == 0 { + cfg = DefaultConfig() + } + return &Arbiter{cfg: cfg, states: map[string]*modelState{}} +} + +func (a *Arbiter) state(key string) *modelState { + a.mu.Lock() + defer a.mu.Unlock() + s := a.states[key] + if s == nil { + s = &modelState{ + getGBps: newEWMA(a.cfg.Alpha, a.cfg.PriorGetGBps), + prefillTPS: newEWMA(a.cfg.Alpha, a.cfg.PriorPrefillToksS), + restoreFix: newEWMA(a.cfg.Alpha, a.cfg.PriorRestoreMsFix), + } + a.states[key] = s + } + return s +} + +// Decision is the arbiter's verdict for one candidate prefix. +type Decision struct { + Restore bool + EstRestoreMs float64 + EstRecomputeMs float64 + Reason string +} + +// Decide weighs restoring a `matchedBytes`/`matchedTokens` prefix for model +// `key` against recomputing those tokens. +func (a *Arbiter) Decide(key string, matchedTokens int, matchedBytes uint64) Decision { + if matchedTokens < a.cfg.MinTokens { + return Decision{Restore: false, Reason: fmt.Sprintf("matched %d toks < min %d", matchedTokens, a.cfg.MinTokens)} + } + s := a.state(key) + gbps := s.getGBps.get() + tps := s.prefillTPS.get() + fix := s.restoreFix.get() + + estRestore := float64(matchedBytes)/(gbps*1e9)*1e3 + fix // ms + estRecompute := float64(matchedTokens) / tps * 1e3 // ms + + if estRestore < a.cfg.SafetyMargin*estRecompute { + return Decision{ + Restore: true, + EstRestoreMs: estRestore, + EstRecomputeMs: estRecompute, + Reason: fmt.Sprintf("restore %.1fms < %.0f%%*recompute %.1fms (restore-path %.1fGB/s, prefill %.0f tok/s)", + estRestore, a.cfg.SafetyMargin*100, estRecompute, gbps, tps), + } + } + return Decision{ + Restore: false, + EstRestoreMs: estRestore, + EstRecomputeMs: estRecompute, + Reason: fmt.Sprintf("recompute %.1fms <= restore %.1fms (restore-path %.1fGB/s too slow to beat prefill for %s)", + estRecompute, estRestore, gbps, humanBytes(matchedBytes)), + } +} + +// --- online feedback from the orchestrator --- + +// ObserveGet feeds a real end-to-end restore: `bytes` moved in `ms` ms. +func (a *Arbiter) ObserveGet(key string, bytes uint64, ms float64) { + if ms > 0 && bytes > 0 { + a.state(key).getGBps.update(float64(bytes) / (ms / 1e3) / 1e9) + } +} + +// ObservePrefill feeds a real prefill: `tokens` computed in `ms` ms. +func (a *Arbiter) ObservePrefill(key string, tokens int, ms float64) { + if ms > 0 && tokens > 0 { + a.state(key).prefillTPS.update(float64(tokens) / (ms / 1e3)) + } +} + +// ObserveRestoreFixed feeds the fixed (size-independent) restore overhead. +func (a *Arbiter) ObserveRestoreFixed(key string, ms float64) { a.state(key).restoreFix.update(ms) } + +// Snapshot exposes the average learned estimates across models (for metrics). +type Snapshot struct { + GetGBps float64 + PrefillToksS float64 + RestoreFixMs float64 + Models int +} + +func (a *Arbiter) Snapshot() Snapshot { + a.mu.Lock() + defer a.mu.Unlock() + if len(a.states) == 0 { + return Snapshot{GetGBps: a.cfg.PriorGetGBps, PrefillToksS: a.cfg.PriorPrefillToksS, RestoreFixMs: a.cfg.PriorRestoreMsFix} + } + var g, t, f float64 + for _, s := range a.states { + g += s.getGBps.get() + t += s.prefillTPS.get() + f += s.restoreFix.get() + } + n := float64(len(a.states)) + return Snapshot{GetGBps: g / n, PrefillToksS: t / n, RestoreFixMs: f / n, Models: len(a.states)} +} + +func humanBytes(b uint64) string { + const u = 1024 + if b < u { + return fmt.Sprintf("%dB", b) + } + div, exp := uint64(u), 0 + for n := b / u; n >= u; n /= u { + div *= u + exp++ + } + return fmt.Sprintf("%.1f%ciB", float64(b)/float64(div), "KMGTPE"[exp]) +} diff --git a/mooncake-integration/ollama/bridge/internal/arbiter/arbiter_test.go b/mooncake-integration/ollama/bridge/internal/arbiter/arbiter_test.go new file mode 100644 index 00000000..808c65c1 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/arbiter/arbiter_test.go @@ -0,0 +1,57 @@ +package arbiter + +import "testing" + +func cfg() Config { + c := DefaultConfig() + c.MinTokens = 100 + return c +} + +func TestDeclineBelowMinTokens(t *testing.T) { + a := New(cfg()) + if a.Decide("m", 50, 1<<20).Restore { + t.Fatal("should not restore below MinTokens") + } +} + +// On a fast model, after one real (slow) restore observation, the arbiter must +// flip to recompute — the loss-free property. +func TestAdaptsToRecompute(t *testing.T) { + a := New(cfg()) + // fast prefill: 20000 tok/s ; a 8000-token, 220MB prefix + a.ObservePrefill("fast", 8000, 8000.0/20000*1000) // 400ms + // first decision uses the optimistic prior -> restore + d0 := a.Decide("fast", 8000, 220<<20) + if !d0.Restore { + t.Fatalf("first decision should try restore, got %q", d0.Reason) + } + // observe that the *full* restore actually took 600ms (slower than prefill) + a.ObserveGet("fast", 220<<20, 600) + d1 := a.Decide("fast", 8000, 220<<20) + if d1.Restore { + t.Fatalf("after learning slow restore, should recompute; got %q", d1.Reason) + } +} + +// On a slow model (large, slow prefill), restore should keep winning. +func TestRestoreWinsWhenWorthIt(t *testing.T) { + a := New(cfg()) + a.ObservePrefill("big", 30000, 30000.0/8000*1000) // 3750ms to prefill 30k toks + a.ObserveGet("big", 880<<20, 1800) // 880MB restore in 1.8s + d := a.Decide("big", 30000, 880<<20) + if !d.Restore { + t.Fatalf("restore should win for slow-prefill model; got %q", d.Reason) + } +} + +// Per-model isolation: a slow model's observations must not poison a fast one. +func TestPerModelIsolation(t *testing.T) { + a := New(cfg()) + a.ObserveGet("slowstore", 100<<20, 5000) // terrible bandwidth for model A + a.ObservePrefill("fastgpu", 8000, 200) // model B prefills fast + // model B has its own (prior) restore estimate, unaffected by A + if a.Snapshot().Models != 2 { + t.Fatalf("want 2 model states, got %d", a.Snapshot().Models) + } +} diff --git a/mooncake-integration/ollama/bridge/internal/bridge/pb/bridge.pb.go b/mooncake-integration/ollama/bridge/internal/bridge/pb/bridge.pb.go new file mode 100644 index 00000000..4fc18839 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/bridge/pb/bridge.pb.go @@ -0,0 +1,1651 @@ +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.34.2 +// protoc v3.12.4 +// source: bridge.proto + +// KVCacheBus is the contract between a caller (a coding Agent, the bench +// harness, or patched Ollama) and the Go sidecar. It implements the three-stage +// reuse flow: Lookup (longest-prefix match) -> Prepare (Load/restore the matched +// KV into a llama.cpp slot) -> Commit (Save the produced KV back to the Mooncake +// Store). + +package bridgepb + +import ( + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + reflect "reflect" + sync "sync" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +// ModelFingerprint: every field that changes KV bytes. A mismatch => miss. +type ModelFingerprint struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + ModelDigest string `protobuf:"bytes,1,opt,name=model_digest,json=modelDigest,proto3" json:"model_digest,omitempty"` + Arch string `protobuf:"bytes,2,opt,name=arch,proto3" json:"arch,omitempty"` + TokenizerHash string `protobuf:"bytes,3,opt,name=tokenizer_hash,json=tokenizerHash,proto3" json:"tokenizer_hash,omitempty"` + RopeHash string `protobuf:"bytes,4,opt,name=rope_hash,json=ropeHash,proto3" json:"rope_hash,omitempty"` + KvType string `protobuf:"bytes,5,opt,name=kv_type,json=kvType,proto3" json:"kv_type,omitempty"` // f16 | q8_0 | q4_0 | bf16 + KvLayout string `protobuf:"bytes,6,opt,name=kv_layout,json=kvLayout,proto3" json:"kv_layout,omitempty"` // dense | swa + NCtxTrain int32 `protobuf:"varint,7,opt,name=n_ctx_train,json=nCtxTrain,proto3" json:"n_ctx_train,omitempty"` + BlockSize int32 `protobuf:"varint,8,opt,name=block_size,json=blockSize,proto3" json:"block_size,omitempty"` + // Optional path to the GGUF file. If set, the sidecar parses it once to fill + // any empty arch/tokenizer/rope/ctx fields (and a content-addressed digest), + // so callers only need to know the model file. + ModelPath string `protobuf:"bytes,9,opt,name=model_path,json=modelPath,proto3" json:"model_path,omitempty"` +} + +func (x *ModelFingerprint) Reset() { + *x = ModelFingerprint{} + if protoimpl.UnsafeEnabled { + mi := &file_bridge_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *ModelFingerprint) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ModelFingerprint) ProtoMessage() {} + +func (x *ModelFingerprint) ProtoReflect() protoreflect.Message { + mi := &file_bridge_proto_msgTypes[0] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ModelFingerprint.ProtoReflect.Descriptor instead. +func (*ModelFingerprint) Descriptor() ([]byte, []int) { + return file_bridge_proto_rawDescGZIP(), []int{0} +} + +func (x *ModelFingerprint) GetModelDigest() string { + if x != nil { + return x.ModelDigest + } + return "" +} + +func (x *ModelFingerprint) GetArch() string { + if x != nil { + return x.Arch + } + return "" +} + +func (x *ModelFingerprint) GetTokenizerHash() string { + if x != nil { + return x.TokenizerHash + } + return "" +} + +func (x *ModelFingerprint) GetRopeHash() string { + if x != nil { + return x.RopeHash + } + return "" +} + +func (x *ModelFingerprint) GetKvType() string { + if x != nil { + return x.KvType + } + return "" +} + +func (x *ModelFingerprint) GetKvLayout() string { + if x != nil { + return x.KvLayout + } + return "" +} + +func (x *ModelFingerprint) GetNCtxTrain() int32 { + if x != nil { + return x.NCtxTrain + } + return 0 +} + +func (x *ModelFingerprint) GetBlockSize() int32 { + if x != nil { + return x.BlockSize + } + return 0 +} + +func (x *ModelFingerprint) GetModelPath() string { + if x != nil { + return x.ModelPath + } + return "" +} + +// CachePolicy mirrors the user-visible options.mooncake.* request extension. +type CachePolicy struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Enable bool `protobuf:"varint,1,opt,name=enable,proto3" json:"enable,omitempty"` + Namespace string `protobuf:"bytes,2,opt,name=namespace,proto3" json:"namespace,omitempty"` // e.g. "repo:my-org/my-repo@main" + Read bool `protobuf:"varint,3,opt,name=read,proto3" json:"read,omitempty"` + Write bool `protobuf:"varint,4,opt,name=write,proto3" json:"write,omitempty"` + // granularity and transport are accepted for forward compatibility. Block + // granularity is the implemented mode; the store transport (tcp|rdma) is + // chosen once by the store-proxy at startup rather than per request. + Granularity string `protobuf:"bytes,5,opt,name=granularity,proto3" json:"granularity,omitempty"` // block (implemented) | sequence (reserved) + BlockSize int32 `protobuf:"varint,6,opt,name=block_size,json=blockSize,proto3" json:"block_size,omitempty"` // overrides fingerprint block_size if > 0 + Transport string `protobuf:"bytes,7,opt,name=transport,proto3" json:"transport,omitempty"` // advisory; effective transport is set on the proxy + ReplicaNum uint32 `protobuf:"varint,8,opt,name=replica_num,json=replicaNum,proto3" json:"replica_num,omitempty"` // hot-prefix replication + SoftPin bool `protobuf:"varint,9,opt,name=soft_pin,json=softPin,proto3" json:"soft_pin,omitempty"` // keep hot prefixes resident + MinPrefixBlocks int32 `protobuf:"varint,10,opt,name=min_prefix_blocks,json=minPrefixBlocks,proto3" json:"min_prefix_blocks,omitempty"` // admission: don't store prefixes shorter than this +} + +func (x *CachePolicy) Reset() { + *x = CachePolicy{} + if protoimpl.UnsafeEnabled { + mi := &file_bridge_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *CachePolicy) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*CachePolicy) ProtoMessage() {} + +func (x *CachePolicy) ProtoReflect() protoreflect.Message { + mi := &file_bridge_proto_msgTypes[1] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use CachePolicy.ProtoReflect.Descriptor instead. +func (*CachePolicy) Descriptor() ([]byte, []int) { + return file_bridge_proto_rawDescGZIP(), []int{1} +} + +func (x *CachePolicy) GetEnable() bool { + if x != nil { + return x.Enable + } + return false +} + +func (x *CachePolicy) GetNamespace() string { + if x != nil { + return x.Namespace + } + return "" +} + +func (x *CachePolicy) GetRead() bool { + if x != nil { + return x.Read + } + return false +} + +func (x *CachePolicy) GetWrite() bool { + if x != nil { + return x.Write + } + return false +} + +func (x *CachePolicy) GetGranularity() string { + if x != nil { + return x.Granularity + } + return "" +} + +func (x *CachePolicy) GetBlockSize() int32 { + if x != nil { + return x.BlockSize + } + return 0 +} + +func (x *CachePolicy) GetTransport() string { + if x != nil { + return x.Transport + } + return "" +} + +func (x *CachePolicy) GetReplicaNum() uint32 { + if x != nil { + return x.ReplicaNum + } + return 0 +} + +func (x *CachePolicy) GetSoftPin() bool { + if x != nil { + return x.SoftPin + } + return false +} + +func (x *CachePolicy) GetMinPrefixBlocks() int32 { + if x != nil { + return x.MinPrefixBlocks + } + return 0 +} + +// LlamaTarget identifies a running llama.cpp server slot to restore into / read +// from. base_url like "http://127.0.0.1:52070"; slot_id is the /slots id. +type LlamaTarget struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + BaseUrl string `protobuf:"bytes,1,opt,name=base_url,json=baseUrl,proto3" json:"base_url,omitempty"` + SlotId int32 `protobuf:"varint,2,opt,name=slot_id,json=slotId,proto3" json:"slot_id,omitempty"` +} + +func (x *LlamaTarget) Reset() { + *x = LlamaTarget{} + if protoimpl.UnsafeEnabled { + mi := &file_bridge_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *LlamaTarget) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*LlamaTarget) ProtoMessage() {} + +func (x *LlamaTarget) ProtoReflect() protoreflect.Message { + mi := &file_bridge_proto_msgTypes[2] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use LlamaTarget.ProtoReflect.Descriptor instead. +func (*LlamaTarget) Descriptor() ([]byte, []int) { + return file_bridge_proto_rawDescGZIP(), []int{2} +} + +func (x *LlamaTarget) GetBaseUrl() string { + if x != nil { + return x.BaseUrl + } + return "" +} + +func (x *LlamaTarget) GetSlotId() int32 { + if x != nil { + return x.SlotId + } + return 0 +} + +type LookupRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Fp *ModelFingerprint `protobuf:"bytes,1,opt,name=fp,proto3" json:"fp,omitempty"` + Policy *CachePolicy `protobuf:"bytes,2,opt,name=policy,proto3" json:"policy,omitempty"` + Tokens []int32 `protobuf:"varint,3,rep,packed,name=tokens,proto3" json:"tokens,omitempty"` +} + +func (x *LookupRequest) Reset() { + *x = LookupRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_bridge_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *LookupRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*LookupRequest) ProtoMessage() {} + +func (x *LookupRequest) ProtoReflect() protoreflect.Message { + mi := &file_bridge_proto_msgTypes[3] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use LookupRequest.ProtoReflect.Descriptor instead. +func (*LookupRequest) Descriptor() ([]byte, []int) { + return file_bridge_proto_rawDescGZIP(), []int{3} +} + +func (x *LookupRequest) GetFp() *ModelFingerprint { + if x != nil { + return x.Fp + } + return nil +} + +func (x *LookupRequest) GetPolicy() *CachePolicy { + if x != nil { + return x.Policy + } + return nil +} + +func (x *LookupRequest) GetTokens() []int32 { + if x != nil { + return x.Tokens + } + return nil +} + +type LookupReply struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Hit bool `protobuf:"varint,1,opt,name=hit,proto3" json:"hit,omitempty"` + MatchedBlocks int32 `protobuf:"varint,2,opt,name=matched_blocks,json=matchedBlocks,proto3" json:"matched_blocks,omitempty"` + MatchedTokens int32 `protobuf:"varint,3,opt,name=matched_tokens,json=matchedTokens,proto3" json:"matched_tokens,omitempty"` + TotalBlocks int32 `protobuf:"varint,4,opt,name=total_blocks,json=totalBlocks,proto3" json:"total_blocks,omitempty"` + TotalTokens int32 `protobuf:"varint,5,opt,name=total_tokens,json=totalTokens,proto3" json:"total_tokens,omitempty"` + Decision string `protobuf:"bytes,6,opt,name=decision,proto3" json:"decision,omitempty"` // restore | recompute | miss + MatchedKey string `protobuf:"bytes,7,opt,name=matched_key,json=matchedKey,proto3" json:"matched_key,omitempty"` + Reason string `protobuf:"bytes,8,opt,name=reason,proto3" json:"reason,omitempty"` // arbiter explanation + Error string `protobuf:"bytes,9,opt,name=error,proto3" json:"error,omitempty"` +} + +func (x *LookupReply) Reset() { + *x = LookupReply{} + if protoimpl.UnsafeEnabled { + mi := &file_bridge_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *LookupReply) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*LookupReply) ProtoMessage() {} + +func (x *LookupReply) ProtoReflect() protoreflect.Message { + mi := &file_bridge_proto_msgTypes[4] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use LookupReply.ProtoReflect.Descriptor instead. +func (*LookupReply) Descriptor() ([]byte, []int) { + return file_bridge_proto_rawDescGZIP(), []int{4} +} + +func (x *LookupReply) GetHit() bool { + if x != nil { + return x.Hit + } + return false +} + +func (x *LookupReply) GetMatchedBlocks() int32 { + if x != nil { + return x.MatchedBlocks + } + return 0 +} + +func (x *LookupReply) GetMatchedTokens() int32 { + if x != nil { + return x.MatchedTokens + } + return 0 +} + +func (x *LookupReply) GetTotalBlocks() int32 { + if x != nil { + return x.TotalBlocks + } + return 0 +} + +func (x *LookupReply) GetTotalTokens() int32 { + if x != nil { + return x.TotalTokens + } + return 0 +} + +func (x *LookupReply) GetDecision() string { + if x != nil { + return x.Decision + } + return "" +} + +func (x *LookupReply) GetMatchedKey() string { + if x != nil { + return x.MatchedKey + } + return "" +} + +func (x *LookupReply) GetReason() string { + if x != nil { + return x.Reason + } + return "" +} + +func (x *LookupReply) GetError() string { + if x != nil { + return x.Error + } + return "" +} + +type PrepareRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Fp *ModelFingerprint `protobuf:"bytes,1,opt,name=fp,proto3" json:"fp,omitempty"` + Policy *CachePolicy `protobuf:"bytes,2,opt,name=policy,proto3" json:"policy,omitempty"` + Tokens []int32 `protobuf:"varint,3,rep,packed,name=tokens,proto3" json:"tokens,omitempty"` + Target *LlamaTarget `protobuf:"bytes,4,opt,name=target,proto3" json:"target,omitempty"` // if unset, behaves like Lookup (plan only) +} + +func (x *PrepareRequest) Reset() { + *x = PrepareRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_bridge_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *PrepareRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PrepareRequest) ProtoMessage() {} + +func (x *PrepareRequest) ProtoReflect() protoreflect.Message { + mi := &file_bridge_proto_msgTypes[5] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PrepareRequest.ProtoReflect.Descriptor instead. +func (*PrepareRequest) Descriptor() ([]byte, []int) { + return file_bridge_proto_rawDescGZIP(), []int{5} +} + +func (x *PrepareRequest) GetFp() *ModelFingerprint { + if x != nil { + return x.Fp + } + return nil +} + +func (x *PrepareRequest) GetPolicy() *CachePolicy { + if x != nil { + return x.Policy + } + return nil +} + +func (x *PrepareRequest) GetTokens() []int32 { + if x != nil { + return x.Tokens + } + return nil +} + +func (x *PrepareRequest) GetTarget() *LlamaTarget { + if x != nil { + return x.Target + } + return nil +} + +type PrepareReply struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Hit bool `protobuf:"varint,1,opt,name=hit,proto3" json:"hit,omitempty"` + MatchedBlocks int32 `protobuf:"varint,2,opt,name=matched_blocks,json=matchedBlocks,proto3" json:"matched_blocks,omitempty"` + MatchedTokens int32 `protobuf:"varint,3,opt,name=matched_tokens,json=matchedTokens,proto3" json:"matched_tokens,omitempty"` + TotalBlocks int32 `protobuf:"varint,4,opt,name=total_blocks,json=totalBlocks,proto3" json:"total_blocks,omitempty"` + TotalTokens int32 `protobuf:"varint,5,opt,name=total_tokens,json=totalTokens,proto3" json:"total_tokens,omitempty"` + Decision string `protobuf:"bytes,6,opt,name=decision,proto3" json:"decision,omitempty"` + MatchedKey string `protobuf:"bytes,7,opt,name=matched_key,json=matchedKey,proto3" json:"matched_key,omitempty"` + Restored bool `protobuf:"varint,8,opt,name=restored,proto3" json:"restored,omitempty"` + RestoredTokens int32 `protobuf:"varint,9,opt,name=restored_tokens,json=restoredTokens,proto3" json:"restored_tokens,omitempty"` + RestoreMs float64 `protobuf:"fixed64,10,opt,name=restore_ms,json=restoreMs,proto3" json:"restore_ms,omitempty"` // llama /slots restore time + StoreGetMs float64 `protobuf:"fixed64,11,opt,name=store_get_ms,json=storeGetMs,proto3" json:"store_get_ms,omitempty"` // store GetFile time + Bytes uint64 `protobuf:"varint,12,opt,name=bytes,proto3" json:"bytes,omitempty"` + Reason string `protobuf:"bytes,13,opt,name=reason,proto3" json:"reason,omitempty"` + Error string `protobuf:"bytes,14,opt,name=error,proto3" json:"error,omitempty"` +} + +func (x *PrepareReply) Reset() { + *x = PrepareReply{} + if protoimpl.UnsafeEnabled { + mi := &file_bridge_proto_msgTypes[6] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *PrepareReply) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PrepareReply) ProtoMessage() {} + +func (x *PrepareReply) ProtoReflect() protoreflect.Message { + mi := &file_bridge_proto_msgTypes[6] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PrepareReply.ProtoReflect.Descriptor instead. +func (*PrepareReply) Descriptor() ([]byte, []int) { + return file_bridge_proto_rawDescGZIP(), []int{6} +} + +func (x *PrepareReply) GetHit() bool { + if x != nil { + return x.Hit + } + return false +} + +func (x *PrepareReply) GetMatchedBlocks() int32 { + if x != nil { + return x.MatchedBlocks + } + return 0 +} + +func (x *PrepareReply) GetMatchedTokens() int32 { + if x != nil { + return x.MatchedTokens + } + return 0 +} + +func (x *PrepareReply) GetTotalBlocks() int32 { + if x != nil { + return x.TotalBlocks + } + return 0 +} + +func (x *PrepareReply) GetTotalTokens() int32 { + if x != nil { + return x.TotalTokens + } + return 0 +} + +func (x *PrepareReply) GetDecision() string { + if x != nil { + return x.Decision + } + return "" +} + +func (x *PrepareReply) GetMatchedKey() string { + if x != nil { + return x.MatchedKey + } + return "" +} + +func (x *PrepareReply) GetRestored() bool { + if x != nil { + return x.Restored + } + return false +} + +func (x *PrepareReply) GetRestoredTokens() int32 { + if x != nil { + return x.RestoredTokens + } + return 0 +} + +func (x *PrepareReply) GetRestoreMs() float64 { + if x != nil { + return x.RestoreMs + } + return 0 +} + +func (x *PrepareReply) GetStoreGetMs() float64 { + if x != nil { + return x.StoreGetMs + } + return 0 +} + +func (x *PrepareReply) GetBytes() uint64 { + if x != nil { + return x.Bytes + } + return 0 +} + +func (x *PrepareReply) GetReason() string { + if x != nil { + return x.Reason + } + return "" +} + +func (x *PrepareReply) GetError() string { + if x != nil { + return x.Error + } + return "" +} + +type CommitRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Fp *ModelFingerprint `protobuf:"bytes,1,opt,name=fp,proto3" json:"fp,omitempty"` + Policy *CachePolicy `protobuf:"bytes,2,opt,name=policy,proto3" json:"policy,omitempty"` + Tokens []int32 `protobuf:"varint,3,rep,packed,name=tokens,proto3" json:"tokens,omitempty"` // tokens resident in the slot (the prefill prefix) + Target *LlamaTarget `protobuf:"bytes,4,opt,name=target,proto3" json:"target,omitempty"` + // Optional: the caller's observed prefill so the arbiter learns the live + // prefill rate (tokens/s) for restore-vs-recompute decisions. + PrefillN int32 `protobuf:"varint,5,opt,name=prefill_n,json=prefillN,proto3" json:"prefill_n,omitempty"` + PrefillMs float64 `protobuf:"fixed64,6,opt,name=prefill_ms,json=prefillMs,proto3" json:"prefill_ms,omitempty"` +} + +func (x *CommitRequest) Reset() { + *x = CommitRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_bridge_proto_msgTypes[7] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *CommitRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*CommitRequest) ProtoMessage() {} + +func (x *CommitRequest) ProtoReflect() protoreflect.Message { + mi := &file_bridge_proto_msgTypes[7] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use CommitRequest.ProtoReflect.Descriptor instead. +func (*CommitRequest) Descriptor() ([]byte, []int) { + return file_bridge_proto_rawDescGZIP(), []int{7} +} + +func (x *CommitRequest) GetFp() *ModelFingerprint { + if x != nil { + return x.Fp + } + return nil +} + +func (x *CommitRequest) GetPolicy() *CachePolicy { + if x != nil { + return x.Policy + } + return nil +} + +func (x *CommitRequest) GetTokens() []int32 { + if x != nil { + return x.Tokens + } + return nil +} + +func (x *CommitRequest) GetTarget() *LlamaTarget { + if x != nil { + return x.Target + } + return nil +} + +func (x *CommitRequest) GetPrefillN() int32 { + if x != nil { + return x.PrefillN + } + return 0 +} + +func (x *CommitRequest) GetPrefillMs() float64 { + if x != nil { + return x.PrefillMs + } + return 0 +} + +type CommitReply struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Ok bool `protobuf:"varint,1,opt,name=ok,proto3" json:"ok,omitempty"` + Stored bool `protobuf:"varint,2,opt,name=stored,proto3" json:"stored,omitempty"` // false if skipped (already present / below threshold) + StoredBlocks int32 `protobuf:"varint,3,opt,name=stored_blocks,json=storedBlocks,proto3" json:"stored_blocks,omitempty"` + StoredTokens int32 `protobuf:"varint,4,opt,name=stored_tokens,json=storedTokens,proto3" json:"stored_tokens,omitempty"` + Bytes uint64 `protobuf:"varint,5,opt,name=bytes,proto3" json:"bytes,omitempty"` + SaveMs float64 `protobuf:"fixed64,6,opt,name=save_ms,json=saveMs,proto3" json:"save_ms,omitempty"` // llama /slots save time + StorePutMs float64 `protobuf:"fixed64,7,opt,name=store_put_ms,json=storePutMs,proto3" json:"store_put_ms,omitempty"` // store PutFile time + Key string `protobuf:"bytes,8,opt,name=key,proto3" json:"key,omitempty"` + Reason string `protobuf:"bytes,9,opt,name=reason,proto3" json:"reason,omitempty"` + Error string `protobuf:"bytes,10,opt,name=error,proto3" json:"error,omitempty"` +} + +func (x *CommitReply) Reset() { + *x = CommitReply{} + if protoimpl.UnsafeEnabled { + mi := &file_bridge_proto_msgTypes[8] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *CommitReply) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*CommitReply) ProtoMessage() {} + +func (x *CommitReply) ProtoReflect() protoreflect.Message { + mi := &file_bridge_proto_msgTypes[8] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use CommitReply.ProtoReflect.Descriptor instead. +func (*CommitReply) Descriptor() ([]byte, []int) { + return file_bridge_proto_rawDescGZIP(), []int{8} +} + +func (x *CommitReply) GetOk() bool { + if x != nil { + return x.Ok + } + return false +} + +func (x *CommitReply) GetStored() bool { + if x != nil { + return x.Stored + } + return false +} + +func (x *CommitReply) GetStoredBlocks() int32 { + if x != nil { + return x.StoredBlocks + } + return 0 +} + +func (x *CommitReply) GetStoredTokens() int32 { + if x != nil { + return x.StoredTokens + } + return 0 +} + +func (x *CommitReply) GetBytes() uint64 { + if x != nil { + return x.Bytes + } + return 0 +} + +func (x *CommitReply) GetSaveMs() float64 { + if x != nil { + return x.SaveMs + } + return 0 +} + +func (x *CommitReply) GetStorePutMs() float64 { + if x != nil { + return x.StorePutMs + } + return 0 +} + +func (x *CommitReply) GetKey() string { + if x != nil { + return x.Key + } + return "" +} + +func (x *CommitReply) GetReason() string { + if x != nil { + return x.Reason + } + return "" +} + +func (x *CommitReply) GetError() string { + if x != nil { + return x.Error + } + return "" +} + +type HealthRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields +} + +func (x *HealthRequest) Reset() { + *x = HealthRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_bridge_proto_msgTypes[9] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *HealthRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HealthRequest) ProtoMessage() {} + +func (x *HealthRequest) ProtoReflect() protoreflect.Message { + mi := &file_bridge_proto_msgTypes[9] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use HealthRequest.ProtoReflect.Descriptor instead. +func (*HealthRequest) Descriptor() ([]byte, []int) { + return file_bridge_proto_rawDescGZIP(), []int{9} +} + +type HealthReply struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Ok bool `protobuf:"varint,1,opt,name=ok,proto3" json:"ok,omitempty"` + Version string `protobuf:"bytes,2,opt,name=version,proto3" json:"version,omitempty"` + StoreBackend string `protobuf:"bytes,3,opt,name=store_backend,json=storeBackend,proto3" json:"store_backend,omitempty"` + StoreOk bool `protobuf:"varint,4,opt,name=store_ok,json=storeOk,proto3" json:"store_ok,omitempty"` + Protocol string `protobuf:"bytes,5,opt,name=protocol,proto3" json:"protocol,omitempty"` + Detail string `protobuf:"bytes,6,opt,name=detail,proto3" json:"detail,omitempty"` +} + +func (x *HealthReply) Reset() { + *x = HealthReply{} + if protoimpl.UnsafeEnabled { + mi := &file_bridge_proto_msgTypes[10] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *HealthReply) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HealthReply) ProtoMessage() {} + +func (x *HealthReply) ProtoReflect() protoreflect.Message { + mi := &file_bridge_proto_msgTypes[10] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use HealthReply.ProtoReflect.Descriptor instead. +func (*HealthReply) Descriptor() ([]byte, []int) { + return file_bridge_proto_rawDescGZIP(), []int{10} +} + +func (x *HealthReply) GetOk() bool { + if x != nil { + return x.Ok + } + return false +} + +func (x *HealthReply) GetVersion() string { + if x != nil { + return x.Version + } + return "" +} + +func (x *HealthReply) GetStoreBackend() string { + if x != nil { + return x.StoreBackend + } + return "" +} + +func (x *HealthReply) GetStoreOk() bool { + if x != nil { + return x.StoreOk + } + return false +} + +func (x *HealthReply) GetProtocol() string { + if x != nil { + return x.Protocol + } + return "" +} + +func (x *HealthReply) GetDetail() string { + if x != nil { + return x.Detail + } + return "" +} + +type StatsRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields +} + +func (x *StatsRequest) Reset() { + *x = StatsRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_bridge_proto_msgTypes[11] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *StatsRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*StatsRequest) ProtoMessage() {} + +func (x *StatsRequest) ProtoReflect() protoreflect.Message { + mi := &file_bridge_proto_msgTypes[11] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use StatsRequest.ProtoReflect.Descriptor instead. +func (*StatsRequest) Descriptor() ([]byte, []int) { + return file_bridge_proto_rawDescGZIP(), []int{11} +} + +type StatsReply struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + PrepareTotal uint64 `protobuf:"varint,1,opt,name=prepare_total,json=prepareTotal,proto3" json:"prepare_total,omitempty"` + Hits uint64 `protobuf:"varint,2,opt,name=hits,proto3" json:"hits,omitempty"` + Misses uint64 `protobuf:"varint,3,opt,name=misses,proto3" json:"misses,omitempty"` + HitBlocks uint64 `protobuf:"varint,4,opt,name=hit_blocks,json=hitBlocks,proto3" json:"hit_blocks,omitempty"` + MissBlocks uint64 `protobuf:"varint,5,opt,name=miss_blocks,json=missBlocks,proto3" json:"miss_blocks,omitempty"` + SavedPrefillTokens uint64 `protobuf:"varint,6,opt,name=saved_prefill_tokens,json=savedPrefillTokens,proto3" json:"saved_prefill_tokens,omitempty"` // primary effectiveness metric + RestoreCount uint64 `protobuf:"varint,7,opt,name=restore_count,json=restoreCount,proto3" json:"restore_count,omitempty"` + CommitCount uint64 `protobuf:"varint,8,opt,name=commit_count,json=commitCount,proto3" json:"commit_count,omitempty"` + BytesGet uint64 `protobuf:"varint,9,opt,name=bytes_get,json=bytesGet,proto3" json:"bytes_get,omitempty"` + BytesPut uint64 `protobuf:"varint,10,opt,name=bytes_put,json=bytesPut,proto3" json:"bytes_put,omitempty"` +} + +func (x *StatsReply) Reset() { + *x = StatsReply{} + if protoimpl.UnsafeEnabled { + mi := &file_bridge_proto_msgTypes[12] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *StatsReply) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*StatsReply) ProtoMessage() {} + +func (x *StatsReply) ProtoReflect() protoreflect.Message { + mi := &file_bridge_proto_msgTypes[12] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use StatsReply.ProtoReflect.Descriptor instead. +func (*StatsReply) Descriptor() ([]byte, []int) { + return file_bridge_proto_rawDescGZIP(), []int{12} +} + +func (x *StatsReply) GetPrepareTotal() uint64 { + if x != nil { + return x.PrepareTotal + } + return 0 +} + +func (x *StatsReply) GetHits() uint64 { + if x != nil { + return x.Hits + } + return 0 +} + +func (x *StatsReply) GetMisses() uint64 { + if x != nil { + return x.Misses + } + return 0 +} + +func (x *StatsReply) GetHitBlocks() uint64 { + if x != nil { + return x.HitBlocks + } + return 0 +} + +func (x *StatsReply) GetMissBlocks() uint64 { + if x != nil { + return x.MissBlocks + } + return 0 +} + +func (x *StatsReply) GetSavedPrefillTokens() uint64 { + if x != nil { + return x.SavedPrefillTokens + } + return 0 +} + +func (x *StatsReply) GetRestoreCount() uint64 { + if x != nil { + return x.RestoreCount + } + return 0 +} + +func (x *StatsReply) GetCommitCount() uint64 { + if x != nil { + return x.CommitCount + } + return 0 +} + +func (x *StatsReply) GetBytesGet() uint64 { + if x != nil { + return x.BytesGet + } + return 0 +} + +func (x *StatsReply) GetBytesPut() uint64 { + if x != nil { + return x.BytesPut + } + return 0 +} + +var File_bridge_proto protoreflect.FileDescriptor + +var file_bridge_proto_rawDesc = []byte{ + 0x0a, 0x0c, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x09, + 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, 0x76, 0x31, 0x22, 0xa1, 0x02, 0x0a, 0x10, 0x4d, 0x6f, + 0x64, 0x65, 0x6c, 0x46, 0x69, 0x6e, 0x67, 0x65, 0x72, 0x70, 0x72, 0x69, 0x6e, 0x74, 0x12, 0x21, + 0x0a, 0x0c, 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x5f, 0x64, 0x69, 0x67, 0x65, 0x73, 0x74, 0x18, 0x01, + 0x20, 0x01, 0x28, 0x09, 0x52, 0x0b, 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x44, 0x69, 0x67, 0x65, 0x73, + 0x74, 0x12, 0x12, 0x0a, 0x04, 0x61, 0x72, 0x63, 0x68, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, + 0x04, 0x61, 0x72, 0x63, 0x68, 0x12, 0x25, 0x0a, 0x0e, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x69, 0x7a, + 0x65, 0x72, 0x5f, 0x68, 0x61, 0x73, 0x68, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x74, + 0x6f, 0x6b, 0x65, 0x6e, 0x69, 0x7a, 0x65, 0x72, 0x48, 0x61, 0x73, 0x68, 0x12, 0x1b, 0x0a, 0x09, + 0x72, 0x6f, 0x70, 0x65, 0x5f, 0x68, 0x61, 0x73, 0x68, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, + 0x08, 0x72, 0x6f, 0x70, 0x65, 0x48, 0x61, 0x73, 0x68, 0x12, 0x17, 0x0a, 0x07, 0x6b, 0x76, 0x5f, + 0x74, 0x79, 0x70, 0x65, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x6b, 0x76, 0x54, 0x79, + 0x70, 0x65, 0x12, 0x1b, 0x0a, 0x09, 0x6b, 0x76, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x18, + 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x6b, 0x76, 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x12, + 0x1e, 0x0a, 0x0b, 0x6e, 0x5f, 0x63, 0x74, 0x78, 0x5f, 0x74, 0x72, 0x61, 0x69, 0x6e, 0x18, 0x07, + 0x20, 0x01, 0x28, 0x05, 0x52, 0x09, 0x6e, 0x43, 0x74, 0x78, 0x54, 0x72, 0x61, 0x69, 0x6e, 0x12, + 0x1d, 0x0a, 0x0a, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x18, 0x08, 0x20, + 0x01, 0x28, 0x05, 0x52, 0x09, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x53, 0x69, 0x7a, 0x65, 0x12, 0x1d, + 0x0a, 0x0a, 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x5f, 0x70, 0x61, 0x74, 0x68, 0x18, 0x09, 0x20, 0x01, + 0x28, 0x09, 0x52, 0x09, 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x50, 0x61, 0x74, 0x68, 0x22, 0xb4, 0x02, + 0x0a, 0x0b, 0x43, 0x61, 0x63, 0x68, 0x65, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x12, 0x16, 0x0a, + 0x06, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x06, 0x65, + 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x12, 0x1c, 0x0a, 0x09, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, + 0x63, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, + 0x61, 0x63, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x72, 0x65, 0x61, 0x64, 0x18, 0x03, 0x20, 0x01, 0x28, + 0x08, 0x52, 0x04, 0x72, 0x65, 0x61, 0x64, 0x12, 0x14, 0x0a, 0x05, 0x77, 0x72, 0x69, 0x74, 0x65, + 0x18, 0x04, 0x20, 0x01, 0x28, 0x08, 0x52, 0x05, 0x77, 0x72, 0x69, 0x74, 0x65, 0x12, 0x20, 0x0a, + 0x0b, 0x67, 0x72, 0x61, 0x6e, 0x75, 0x6c, 0x61, 0x72, 0x69, 0x74, 0x79, 0x18, 0x05, 0x20, 0x01, + 0x28, 0x09, 0x52, 0x0b, 0x67, 0x72, 0x61, 0x6e, 0x75, 0x6c, 0x61, 0x72, 0x69, 0x74, 0x79, 0x12, + 0x1d, 0x0a, 0x0a, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x18, 0x06, 0x20, + 0x01, 0x28, 0x05, 0x52, 0x09, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x53, 0x69, 0x7a, 0x65, 0x12, 0x1c, + 0x0a, 0x09, 0x74, 0x72, 0x61, 0x6e, 0x73, 0x70, 0x6f, 0x72, 0x74, 0x18, 0x07, 0x20, 0x01, 0x28, + 0x09, 0x52, 0x09, 0x74, 0x72, 0x61, 0x6e, 0x73, 0x70, 0x6f, 0x72, 0x74, 0x12, 0x1f, 0x0a, 0x0b, + 0x72, 0x65, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x5f, 0x6e, 0x75, 0x6d, 0x18, 0x08, 0x20, 0x01, 0x28, + 0x0d, 0x52, 0x0a, 0x72, 0x65, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x4e, 0x75, 0x6d, 0x12, 0x19, 0x0a, + 0x08, 0x73, 0x6f, 0x66, 0x74, 0x5f, 0x70, 0x69, 0x6e, 0x18, 0x09, 0x20, 0x01, 0x28, 0x08, 0x52, + 0x07, 0x73, 0x6f, 0x66, 0x74, 0x50, 0x69, 0x6e, 0x12, 0x2a, 0x0a, 0x11, 0x6d, 0x69, 0x6e, 0x5f, + 0x70, 0x72, 0x65, 0x66, 0x69, 0x78, 0x5f, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x18, 0x0a, 0x20, + 0x01, 0x28, 0x05, 0x52, 0x0f, 0x6d, 0x69, 0x6e, 0x50, 0x72, 0x65, 0x66, 0x69, 0x78, 0x42, 0x6c, + 0x6f, 0x63, 0x6b, 0x73, 0x22, 0x41, 0x0a, 0x0b, 0x4c, 0x6c, 0x61, 0x6d, 0x61, 0x54, 0x61, 0x72, + 0x67, 0x65, 0x74, 0x12, 0x19, 0x0a, 0x08, 0x62, 0x61, 0x73, 0x65, 0x5f, 0x75, 0x72, 0x6c, 0x18, + 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x62, 0x61, 0x73, 0x65, 0x55, 0x72, 0x6c, 0x12, 0x17, + 0x0a, 0x07, 0x73, 0x6c, 0x6f, 0x74, 0x5f, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x05, 0x52, + 0x06, 0x73, 0x6c, 0x6f, 0x74, 0x49, 0x64, 0x22, 0x84, 0x01, 0x0a, 0x0d, 0x4c, 0x6f, 0x6f, 0x6b, + 0x75, 0x70, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x2b, 0x0a, 0x02, 0x66, 0x70, 0x18, + 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1b, 0x2e, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, 0x76, + 0x31, 0x2e, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x46, 0x69, 0x6e, 0x67, 0x65, 0x72, 0x70, 0x72, 0x69, + 0x6e, 0x74, 0x52, 0x02, 0x66, 0x70, 0x12, 0x2e, 0x0a, 0x06, 0x70, 0x6f, 0x6c, 0x69, 0x63, 0x79, + 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, + 0x76, 0x31, 0x2e, 0x43, 0x61, 0x63, 0x68, 0x65, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x52, 0x06, + 0x70, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x12, 0x16, 0x0a, 0x06, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, + 0x18, 0x03, 0x20, 0x03, 0x28, 0x05, 0x52, 0x06, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x22, 0x9e, + 0x02, 0x0a, 0x0b, 0x4c, 0x6f, 0x6f, 0x6b, 0x75, 0x70, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x10, + 0x0a, 0x03, 0x68, 0x69, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x03, 0x68, 0x69, 0x74, + 0x12, 0x25, 0x0a, 0x0e, 0x6d, 0x61, 0x74, 0x63, 0x68, 0x65, 0x64, 0x5f, 0x62, 0x6c, 0x6f, 0x63, + 0x6b, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0d, 0x6d, 0x61, 0x74, 0x63, 0x68, 0x65, + 0x64, 0x42, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x12, 0x25, 0x0a, 0x0e, 0x6d, 0x61, 0x74, 0x63, 0x68, + 0x65, 0x64, 0x5f, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x05, 0x52, + 0x0d, 0x6d, 0x61, 0x74, 0x63, 0x68, 0x65, 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x12, 0x21, + 0x0a, 0x0c, 0x74, 0x6f, 0x74, 0x61, 0x6c, 0x5f, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x18, 0x04, + 0x20, 0x01, 0x28, 0x05, 0x52, 0x0b, 0x74, 0x6f, 0x74, 0x61, 0x6c, 0x42, 0x6c, 0x6f, 0x63, 0x6b, + 0x73, 0x12, 0x21, 0x0a, 0x0c, 0x74, 0x6f, 0x74, 0x61, 0x6c, 0x5f, 0x74, 0x6f, 0x6b, 0x65, 0x6e, + 0x73, 0x18, 0x05, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0b, 0x74, 0x6f, 0x74, 0x61, 0x6c, 0x54, 0x6f, + 0x6b, 0x65, 0x6e, 0x73, 0x12, 0x1a, 0x0a, 0x08, 0x64, 0x65, 0x63, 0x69, 0x73, 0x69, 0x6f, 0x6e, + 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x64, 0x65, 0x63, 0x69, 0x73, 0x69, 0x6f, 0x6e, + 0x12, 0x1f, 0x0a, 0x0b, 0x6d, 0x61, 0x74, 0x63, 0x68, 0x65, 0x64, 0x5f, 0x6b, 0x65, 0x79, 0x18, + 0x07, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0a, 0x6d, 0x61, 0x74, 0x63, 0x68, 0x65, 0x64, 0x4b, 0x65, + 0x79, 0x12, 0x16, 0x0a, 0x06, 0x72, 0x65, 0x61, 0x73, 0x6f, 0x6e, 0x18, 0x08, 0x20, 0x01, 0x28, + 0x09, 0x52, 0x06, 0x72, 0x65, 0x61, 0x73, 0x6f, 0x6e, 0x12, 0x14, 0x0a, 0x05, 0x65, 0x72, 0x72, + 0x6f, 0x72, 0x18, 0x09, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x22, + 0xb5, 0x01, 0x0a, 0x0e, 0x50, 0x72, 0x65, 0x70, 0x61, 0x72, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, + 0x73, 0x74, 0x12, 0x2b, 0x0a, 0x02, 0x66, 0x70, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1b, + 0x2e, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, 0x76, 0x31, 0x2e, 0x4d, 0x6f, 0x64, 0x65, 0x6c, + 0x46, 0x69, 0x6e, 0x67, 0x65, 0x72, 0x70, 0x72, 0x69, 0x6e, 0x74, 0x52, 0x02, 0x66, 0x70, 0x12, + 0x2e, 0x0a, 0x06, 0x70, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, + 0x16, 0x2e, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x61, 0x63, 0x68, + 0x65, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x52, 0x06, 0x70, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x12, + 0x16, 0x0a, 0x06, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x18, 0x03, 0x20, 0x03, 0x28, 0x05, 0x52, + 0x06, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x12, 0x2e, 0x0a, 0x06, 0x74, 0x61, 0x72, 0x67, 0x65, + 0x74, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, + 0x2e, 0x76, 0x31, 0x2e, 0x4c, 0x6c, 0x61, 0x6d, 0x61, 0x54, 0x61, 0x72, 0x67, 0x65, 0x74, 0x52, + 0x06, 0x74, 0x61, 0x72, 0x67, 0x65, 0x74, 0x22, 0xbb, 0x03, 0x0a, 0x0c, 0x50, 0x72, 0x65, 0x70, + 0x61, 0x72, 0x65, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x10, 0x0a, 0x03, 0x68, 0x69, 0x74, 0x18, + 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x03, 0x68, 0x69, 0x74, 0x12, 0x25, 0x0a, 0x0e, 0x6d, 0x61, + 0x74, 0x63, 0x68, 0x65, 0x64, 0x5f, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x18, 0x02, 0x20, 0x01, + 0x28, 0x05, 0x52, 0x0d, 0x6d, 0x61, 0x74, 0x63, 0x68, 0x65, 0x64, 0x42, 0x6c, 0x6f, 0x63, 0x6b, + 0x73, 0x12, 0x25, 0x0a, 0x0e, 0x6d, 0x61, 0x74, 0x63, 0x68, 0x65, 0x64, 0x5f, 0x74, 0x6f, 0x6b, + 0x65, 0x6e, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0d, 0x6d, 0x61, 0x74, 0x63, 0x68, + 0x65, 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x12, 0x21, 0x0a, 0x0c, 0x74, 0x6f, 0x74, 0x61, + 0x6c, 0x5f, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0b, + 0x74, 0x6f, 0x74, 0x61, 0x6c, 0x42, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x12, 0x21, 0x0a, 0x0c, 0x74, + 0x6f, 0x74, 0x61, 0x6c, 0x5f, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x18, 0x05, 0x20, 0x01, 0x28, + 0x05, 0x52, 0x0b, 0x74, 0x6f, 0x74, 0x61, 0x6c, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x12, 0x1a, + 0x0a, 0x08, 0x64, 0x65, 0x63, 0x69, 0x73, 0x69, 0x6f, 0x6e, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, + 0x52, 0x08, 0x64, 0x65, 0x63, 0x69, 0x73, 0x69, 0x6f, 0x6e, 0x12, 0x1f, 0x0a, 0x0b, 0x6d, 0x61, + 0x74, 0x63, 0x68, 0x65, 0x64, 0x5f, 0x6b, 0x65, 0x79, 0x18, 0x07, 0x20, 0x01, 0x28, 0x09, 0x52, + 0x0a, 0x6d, 0x61, 0x74, 0x63, 0x68, 0x65, 0x64, 0x4b, 0x65, 0x79, 0x12, 0x1a, 0x0a, 0x08, 0x72, + 0x65, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x64, 0x18, 0x08, 0x20, 0x01, 0x28, 0x08, 0x52, 0x08, 0x72, + 0x65, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x64, 0x12, 0x27, 0x0a, 0x0f, 0x72, 0x65, 0x73, 0x74, 0x6f, + 0x72, 0x65, 0x64, 0x5f, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x18, 0x09, 0x20, 0x01, 0x28, 0x05, + 0x52, 0x0e, 0x72, 0x65, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, + 0x12, 0x1d, 0x0a, 0x0a, 0x72, 0x65, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x5f, 0x6d, 0x73, 0x18, 0x0a, + 0x20, 0x01, 0x28, 0x01, 0x52, 0x09, 0x72, 0x65, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x4d, 0x73, 0x12, + 0x20, 0x0a, 0x0c, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x5f, 0x67, 0x65, 0x74, 0x5f, 0x6d, 0x73, 0x18, + 0x0b, 0x20, 0x01, 0x28, 0x01, 0x52, 0x0a, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x47, 0x65, 0x74, 0x4d, + 0x73, 0x12, 0x14, 0x0a, 0x05, 0x62, 0x79, 0x74, 0x65, 0x73, 0x18, 0x0c, 0x20, 0x01, 0x28, 0x04, + 0x52, 0x05, 0x62, 0x79, 0x74, 0x65, 0x73, 0x12, 0x16, 0x0a, 0x06, 0x72, 0x65, 0x61, 0x73, 0x6f, + 0x6e, 0x18, 0x0d, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x72, 0x65, 0x61, 0x73, 0x6f, 0x6e, 0x12, + 0x14, 0x0a, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x18, 0x0e, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, + 0x65, 0x72, 0x72, 0x6f, 0x72, 0x22, 0xf0, 0x01, 0x0a, 0x0d, 0x43, 0x6f, 0x6d, 0x6d, 0x69, 0x74, + 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x2b, 0x0a, 0x02, 0x66, 0x70, 0x18, 0x01, 0x20, + 0x01, 0x28, 0x0b, 0x32, 0x1b, 0x2e, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, 0x76, 0x31, 0x2e, + 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x46, 0x69, 0x6e, 0x67, 0x65, 0x72, 0x70, 0x72, 0x69, 0x6e, 0x74, + 0x52, 0x02, 0x66, 0x70, 0x12, 0x2e, 0x0a, 0x06, 0x70, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x18, 0x02, + 0x20, 0x01, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, 0x76, 0x31, + 0x2e, 0x43, 0x61, 0x63, 0x68, 0x65, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x52, 0x06, 0x70, 0x6f, + 0x6c, 0x69, 0x63, 0x79, 0x12, 0x16, 0x0a, 0x06, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x18, 0x03, + 0x20, 0x03, 0x28, 0x05, 0x52, 0x06, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x12, 0x2e, 0x0a, 0x06, + 0x74, 0x61, 0x72, 0x67, 0x65, 0x74, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x16, 0x2e, 0x62, + 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, 0x76, 0x31, 0x2e, 0x4c, 0x6c, 0x61, 0x6d, 0x61, 0x54, 0x61, + 0x72, 0x67, 0x65, 0x74, 0x52, 0x06, 0x74, 0x61, 0x72, 0x67, 0x65, 0x74, 0x12, 0x1b, 0x0a, 0x09, + 0x70, 0x72, 0x65, 0x66, 0x69, 0x6c, 0x6c, 0x5f, 0x6e, 0x18, 0x05, 0x20, 0x01, 0x28, 0x05, 0x52, + 0x08, 0x70, 0x72, 0x65, 0x66, 0x69, 0x6c, 0x6c, 0x4e, 0x12, 0x1d, 0x0a, 0x0a, 0x70, 0x72, 0x65, + 0x66, 0x69, 0x6c, 0x6c, 0x5f, 0x6d, 0x73, 0x18, 0x06, 0x20, 0x01, 0x28, 0x01, 0x52, 0x09, 0x70, + 0x72, 0x65, 0x66, 0x69, 0x6c, 0x6c, 0x4d, 0x73, 0x22, 0x90, 0x02, 0x0a, 0x0b, 0x43, 0x6f, 0x6d, + 0x6d, 0x69, 0x74, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x0e, 0x0a, 0x02, 0x6f, 0x6b, 0x18, 0x01, + 0x20, 0x01, 0x28, 0x08, 0x52, 0x02, 0x6f, 0x6b, 0x12, 0x16, 0x0a, 0x06, 0x73, 0x74, 0x6f, 0x72, + 0x65, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x08, 0x52, 0x06, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x64, + 0x12, 0x23, 0x0a, 0x0d, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x64, 0x5f, 0x62, 0x6c, 0x6f, 0x63, 0x6b, + 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0c, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x64, 0x42, + 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x12, 0x23, 0x0a, 0x0d, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x64, 0x5f, + 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0c, 0x73, 0x74, + 0x6f, 0x72, 0x65, 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x12, 0x14, 0x0a, 0x05, 0x62, 0x79, + 0x74, 0x65, 0x73, 0x18, 0x05, 0x20, 0x01, 0x28, 0x04, 0x52, 0x05, 0x62, 0x79, 0x74, 0x65, 0x73, + 0x12, 0x17, 0x0a, 0x07, 0x73, 0x61, 0x76, 0x65, 0x5f, 0x6d, 0x73, 0x18, 0x06, 0x20, 0x01, 0x28, + 0x01, 0x52, 0x06, 0x73, 0x61, 0x76, 0x65, 0x4d, 0x73, 0x12, 0x20, 0x0a, 0x0c, 0x73, 0x74, 0x6f, + 0x72, 0x65, 0x5f, 0x70, 0x75, 0x74, 0x5f, 0x6d, 0x73, 0x18, 0x07, 0x20, 0x01, 0x28, 0x01, 0x52, + 0x0a, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x50, 0x75, 0x74, 0x4d, 0x73, 0x12, 0x10, 0x0a, 0x03, 0x6b, + 0x65, 0x79, 0x18, 0x08, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6b, 0x65, 0x79, 0x12, 0x16, 0x0a, + 0x06, 0x72, 0x65, 0x61, 0x73, 0x6f, 0x6e, 0x18, 0x09, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x72, + 0x65, 0x61, 0x73, 0x6f, 0x6e, 0x12, 0x14, 0x0a, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x18, 0x0a, + 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x22, 0x0f, 0x0a, 0x0d, 0x48, + 0x65, 0x61, 0x6c, 0x74, 0x68, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x22, 0xab, 0x01, 0x0a, + 0x0b, 0x48, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x0e, 0x0a, 0x02, + 0x6f, 0x6b, 0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x02, 0x6f, 0x6b, 0x12, 0x18, 0x0a, 0x07, + 0x76, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x76, + 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x12, 0x23, 0x0a, 0x0d, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x5f, + 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0c, 0x73, + 0x74, 0x6f, 0x72, 0x65, 0x42, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x12, 0x19, 0x0a, 0x08, 0x73, + 0x74, 0x6f, 0x72, 0x65, 0x5f, 0x6f, 0x6b, 0x18, 0x04, 0x20, 0x01, 0x28, 0x08, 0x52, 0x07, 0x73, + 0x74, 0x6f, 0x72, 0x65, 0x4f, 0x6b, 0x12, 0x1a, 0x0a, 0x08, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x63, + 0x6f, 0x6c, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x63, + 0x6f, 0x6c, 0x12, 0x16, 0x0a, 0x06, 0x64, 0x65, 0x74, 0x61, 0x69, 0x6c, 0x18, 0x06, 0x20, 0x01, + 0x28, 0x09, 0x52, 0x06, 0x64, 0x65, 0x74, 0x61, 0x69, 0x6c, 0x22, 0x0e, 0x0a, 0x0c, 0x53, 0x74, + 0x61, 0x74, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x22, 0xd1, 0x02, 0x0a, 0x0a, 0x53, + 0x74, 0x61, 0x74, 0x73, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x23, 0x0a, 0x0d, 0x70, 0x72, 0x65, + 0x70, 0x61, 0x72, 0x65, 0x5f, 0x74, 0x6f, 0x74, 0x61, 0x6c, 0x18, 0x01, 0x20, 0x01, 0x28, 0x04, + 0x52, 0x0c, 0x70, 0x72, 0x65, 0x70, 0x61, 0x72, 0x65, 0x54, 0x6f, 0x74, 0x61, 0x6c, 0x12, 0x12, + 0x0a, 0x04, 0x68, 0x69, 0x74, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x04, 0x52, 0x04, 0x68, 0x69, + 0x74, 0x73, 0x12, 0x16, 0x0a, 0x06, 0x6d, 0x69, 0x73, 0x73, 0x65, 0x73, 0x18, 0x03, 0x20, 0x01, + 0x28, 0x04, 0x52, 0x06, 0x6d, 0x69, 0x73, 0x73, 0x65, 0x73, 0x12, 0x1d, 0x0a, 0x0a, 0x68, 0x69, + 0x74, 0x5f, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x04, 0x52, 0x09, + 0x68, 0x69, 0x74, 0x42, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x12, 0x1f, 0x0a, 0x0b, 0x6d, 0x69, 0x73, + 0x73, 0x5f, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x18, 0x05, 0x20, 0x01, 0x28, 0x04, 0x52, 0x0a, + 0x6d, 0x69, 0x73, 0x73, 0x42, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x12, 0x30, 0x0a, 0x14, 0x73, 0x61, + 0x76, 0x65, 0x64, 0x5f, 0x70, 0x72, 0x65, 0x66, 0x69, 0x6c, 0x6c, 0x5f, 0x74, 0x6f, 0x6b, 0x65, + 0x6e, 0x73, 0x18, 0x06, 0x20, 0x01, 0x28, 0x04, 0x52, 0x12, 0x73, 0x61, 0x76, 0x65, 0x64, 0x50, + 0x72, 0x65, 0x66, 0x69, 0x6c, 0x6c, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x12, 0x23, 0x0a, 0x0d, + 0x72, 0x65, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x5f, 0x63, 0x6f, 0x75, 0x6e, 0x74, 0x18, 0x07, 0x20, + 0x01, 0x28, 0x04, 0x52, 0x0c, 0x72, 0x65, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x43, 0x6f, 0x75, 0x6e, + 0x74, 0x12, 0x21, 0x0a, 0x0c, 0x63, 0x6f, 0x6d, 0x6d, 0x69, 0x74, 0x5f, 0x63, 0x6f, 0x75, 0x6e, + 0x74, 0x18, 0x08, 0x20, 0x01, 0x28, 0x04, 0x52, 0x0b, 0x63, 0x6f, 0x6d, 0x6d, 0x69, 0x74, 0x43, + 0x6f, 0x75, 0x6e, 0x74, 0x12, 0x1b, 0x0a, 0x09, 0x62, 0x79, 0x74, 0x65, 0x73, 0x5f, 0x67, 0x65, + 0x74, 0x18, 0x09, 0x20, 0x01, 0x28, 0x04, 0x52, 0x08, 0x62, 0x79, 0x74, 0x65, 0x73, 0x47, 0x65, + 0x74, 0x12, 0x1b, 0x0a, 0x09, 0x62, 0x79, 0x74, 0x65, 0x73, 0x5f, 0x70, 0x75, 0x74, 0x18, 0x0a, + 0x20, 0x01, 0x28, 0x04, 0x52, 0x08, 0x62, 0x79, 0x74, 0x65, 0x73, 0x50, 0x75, 0x74, 0x32, 0xb8, + 0x02, 0x0a, 0x0a, 0x4b, 0x56, 0x43, 0x61, 0x63, 0x68, 0x65, 0x42, 0x75, 0x73, 0x12, 0x3a, 0x0a, + 0x06, 0x48, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x12, 0x18, 0x2e, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, + 0x2e, 0x76, 0x31, 0x2e, 0x48, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x1a, 0x16, 0x2e, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, 0x76, 0x31, 0x2e, 0x48, 0x65, + 0x61, 0x6c, 0x74, 0x68, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x3a, 0x0a, 0x06, 0x4c, 0x6f, 0x6f, + 0x6b, 0x75, 0x70, 0x12, 0x18, 0x2e, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, 0x76, 0x31, 0x2e, + 0x4c, 0x6f, 0x6f, 0x6b, 0x75, 0x70, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, + 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, 0x76, 0x31, 0x2e, 0x4c, 0x6f, 0x6f, 0x6b, 0x75, 0x70, + 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x3d, 0x0a, 0x07, 0x50, 0x72, 0x65, 0x70, 0x61, 0x72, 0x65, + 0x12, 0x19, 0x2e, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, 0x76, 0x31, 0x2e, 0x50, 0x72, 0x65, + 0x70, 0x61, 0x72, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x17, 0x2e, 0x62, 0x72, + 0x69, 0x64, 0x67, 0x65, 0x2e, 0x76, 0x31, 0x2e, 0x50, 0x72, 0x65, 0x70, 0x61, 0x72, 0x65, 0x52, + 0x65, 0x70, 0x6c, 0x79, 0x12, 0x3a, 0x0a, 0x06, 0x43, 0x6f, 0x6d, 0x6d, 0x69, 0x74, 0x12, 0x18, + 0x2e, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x6f, 0x6d, 0x6d, 0x69, + 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, 0x62, 0x72, 0x69, 0x64, 0x67, + 0x65, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x6f, 0x6d, 0x6d, 0x69, 0x74, 0x52, 0x65, 0x70, 0x6c, 0x79, + 0x12, 0x37, 0x0a, 0x05, 0x53, 0x74, 0x61, 0x74, 0x73, 0x12, 0x17, 0x2e, 0x62, 0x72, 0x69, 0x64, + 0x67, 0x65, 0x2e, 0x76, 0x31, 0x2e, 0x53, 0x74, 0x61, 0x74, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, + 0x73, 0x74, 0x1a, 0x15, 0x2e, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2e, 0x76, 0x31, 0x2e, 0x53, + 0x74, 0x61, 0x74, 0x73, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x42, 0x4b, 0x5a, 0x49, 0x67, 0x69, 0x74, + 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x6d, 0x6f, 0x6f, 0x6e, 0x63, 0x61, 0x6b, 0x65, + 0x2d, 0x61, 0x69, 0x2f, 0x6f, 0x6c, 0x6c, 0x61, 0x6d, 0x61, 0x2d, 0x6d, 0x6f, 0x6f, 0x6e, 0x63, + 0x61, 0x6b, 0x65, 0x2d, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2f, 0x69, 0x6e, 0x74, 0x65, 0x72, + 0x6e, 0x61, 0x6c, 0x2f, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2f, 0x70, 0x62, 0x3b, 0x62, 0x72, + 0x69, 0x64, 0x67, 0x65, 0x70, 0x62, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, +} + +var ( + file_bridge_proto_rawDescOnce sync.Once + file_bridge_proto_rawDescData = file_bridge_proto_rawDesc +) + +func file_bridge_proto_rawDescGZIP() []byte { + file_bridge_proto_rawDescOnce.Do(func() { + file_bridge_proto_rawDescData = protoimpl.X.CompressGZIP(file_bridge_proto_rawDescData) + }) + return file_bridge_proto_rawDescData +} + +var file_bridge_proto_msgTypes = make([]protoimpl.MessageInfo, 13) +var file_bridge_proto_goTypes = []any{ + (*ModelFingerprint)(nil), // 0: bridge.v1.ModelFingerprint + (*CachePolicy)(nil), // 1: bridge.v1.CachePolicy + (*LlamaTarget)(nil), // 2: bridge.v1.LlamaTarget + (*LookupRequest)(nil), // 3: bridge.v1.LookupRequest + (*LookupReply)(nil), // 4: bridge.v1.LookupReply + (*PrepareRequest)(nil), // 5: bridge.v1.PrepareRequest + (*PrepareReply)(nil), // 6: bridge.v1.PrepareReply + (*CommitRequest)(nil), // 7: bridge.v1.CommitRequest + (*CommitReply)(nil), // 8: bridge.v1.CommitReply + (*HealthRequest)(nil), // 9: bridge.v1.HealthRequest + (*HealthReply)(nil), // 10: bridge.v1.HealthReply + (*StatsRequest)(nil), // 11: bridge.v1.StatsRequest + (*StatsReply)(nil), // 12: bridge.v1.StatsReply +} +var file_bridge_proto_depIdxs = []int32{ + 0, // 0: bridge.v1.LookupRequest.fp:type_name -> bridge.v1.ModelFingerprint + 1, // 1: bridge.v1.LookupRequest.policy:type_name -> bridge.v1.CachePolicy + 0, // 2: bridge.v1.PrepareRequest.fp:type_name -> bridge.v1.ModelFingerprint + 1, // 3: bridge.v1.PrepareRequest.policy:type_name -> bridge.v1.CachePolicy + 2, // 4: bridge.v1.PrepareRequest.target:type_name -> bridge.v1.LlamaTarget + 0, // 5: bridge.v1.CommitRequest.fp:type_name -> bridge.v1.ModelFingerprint + 1, // 6: bridge.v1.CommitRequest.policy:type_name -> bridge.v1.CachePolicy + 2, // 7: bridge.v1.CommitRequest.target:type_name -> bridge.v1.LlamaTarget + 9, // 8: bridge.v1.KVCacheBus.Health:input_type -> bridge.v1.HealthRequest + 3, // 9: bridge.v1.KVCacheBus.Lookup:input_type -> bridge.v1.LookupRequest + 5, // 10: bridge.v1.KVCacheBus.Prepare:input_type -> bridge.v1.PrepareRequest + 7, // 11: bridge.v1.KVCacheBus.Commit:input_type -> bridge.v1.CommitRequest + 11, // 12: bridge.v1.KVCacheBus.Stats:input_type -> bridge.v1.StatsRequest + 10, // 13: bridge.v1.KVCacheBus.Health:output_type -> bridge.v1.HealthReply + 4, // 14: bridge.v1.KVCacheBus.Lookup:output_type -> bridge.v1.LookupReply + 6, // 15: bridge.v1.KVCacheBus.Prepare:output_type -> bridge.v1.PrepareReply + 8, // 16: bridge.v1.KVCacheBus.Commit:output_type -> bridge.v1.CommitReply + 12, // 17: bridge.v1.KVCacheBus.Stats:output_type -> bridge.v1.StatsReply + 13, // [13:18] is the sub-list for method output_type + 8, // [8:13] is the sub-list for method input_type + 8, // [8:8] is the sub-list for extension type_name + 8, // [8:8] is the sub-list for extension extendee + 0, // [0:8] is the sub-list for field type_name +} + +func init() { file_bridge_proto_init() } +func file_bridge_proto_init() { + if File_bridge_proto != nil { + return + } + if !protoimpl.UnsafeEnabled { + file_bridge_proto_msgTypes[0].Exporter = func(v any, i int) any { + switch v := v.(*ModelFingerprint); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_bridge_proto_msgTypes[1].Exporter = func(v any, i int) any { + switch v := v.(*CachePolicy); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_bridge_proto_msgTypes[2].Exporter = func(v any, i int) any { + switch v := v.(*LlamaTarget); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_bridge_proto_msgTypes[3].Exporter = func(v any, i int) any { + switch v := v.(*LookupRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_bridge_proto_msgTypes[4].Exporter = func(v any, i int) any { + switch v := v.(*LookupReply); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_bridge_proto_msgTypes[5].Exporter = func(v any, i int) any { + switch v := v.(*PrepareRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_bridge_proto_msgTypes[6].Exporter = func(v any, i int) any { + switch v := v.(*PrepareReply); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_bridge_proto_msgTypes[7].Exporter = func(v any, i int) any { + switch v := v.(*CommitRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_bridge_proto_msgTypes[8].Exporter = func(v any, i int) any { + switch v := v.(*CommitReply); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_bridge_proto_msgTypes[9].Exporter = func(v any, i int) any { + switch v := v.(*HealthRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_bridge_proto_msgTypes[10].Exporter = func(v any, i int) any { + switch v := v.(*HealthReply); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_bridge_proto_msgTypes[11].Exporter = func(v any, i int) any { + switch v := v.(*StatsRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_bridge_proto_msgTypes[12].Exporter = func(v any, i int) any { + switch v := v.(*StatsReply); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + } + type x struct{} + out := protoimpl.TypeBuilder{ + File: protoimpl.DescBuilder{ + GoPackagePath: reflect.TypeOf(x{}).PkgPath(), + RawDescriptor: file_bridge_proto_rawDesc, + NumEnums: 0, + NumMessages: 13, + NumExtensions: 0, + NumServices: 1, + }, + GoTypes: file_bridge_proto_goTypes, + DependencyIndexes: file_bridge_proto_depIdxs, + MessageInfos: file_bridge_proto_msgTypes, + }.Build() + File_bridge_proto = out.File + file_bridge_proto_rawDesc = nil + file_bridge_proto_goTypes = nil + file_bridge_proto_depIdxs = nil +} diff --git a/mooncake-integration/ollama/bridge/internal/bridge/pb/bridge_grpc.pb.go b/mooncake-integration/ollama/bridge/internal/bridge/pb/bridge_grpc.pb.go new file mode 100644 index 00000000..41662985 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/bridge/pb/bridge_grpc.pb.go @@ -0,0 +1,278 @@ +// Code generated by protoc-gen-go-grpc. DO NOT EDIT. +// versions: +// - protoc-gen-go-grpc v1.4.0 +// - protoc v3.12.4 +// source: bridge.proto + +// KVCacheBus is the contract between a caller (a coding Agent, the bench +// harness, or patched Ollama) and the Go sidecar. It implements the three-stage +// reuse flow: Lookup (longest-prefix match) -> Prepare (Load/restore the matched +// KV into a llama.cpp slot) -> Commit (Save the produced KV back to the Mooncake +// Store). + +package bridgepb + +import ( + context "context" + grpc "google.golang.org/grpc" + codes "google.golang.org/grpc/codes" + status "google.golang.org/grpc/status" +) + +// This is a compile-time assertion to ensure that this generated file +// is compatible with the grpc package it is being compiled against. +// Requires gRPC-Go v1.62.0 or later. +const _ = grpc.SupportPackageIsVersion8 + +const ( + KVCacheBus_Health_FullMethodName = "/bridge.v1.KVCacheBus/Health" + KVCacheBus_Lookup_FullMethodName = "/bridge.v1.KVCacheBus/Lookup" + KVCacheBus_Prepare_FullMethodName = "/bridge.v1.KVCacheBus/Prepare" + KVCacheBus_Commit_FullMethodName = "/bridge.v1.KVCacheBus/Commit" + KVCacheBus_Stats_FullMethodName = "/bridge.v1.KVCacheBus/Stats" +) + +// KVCacheBusClient is the client API for KVCacheBus service. +// +// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. +type KVCacheBusClient interface { + Health(ctx context.Context, in *HealthRequest, opts ...grpc.CallOption) (*HealthReply, error) + // Lookup is read-only: longest block-prefix match + arbiter decision, no I/O + // into llama.cpp. Useful for inspection / benchmarking. + Lookup(ctx context.Context, in *LookupRequest, opts ...grpc.CallOption) (*LookupReply, error) + // Prepare = Lookup + (if a hit and the cost arbiter approves) restore the + // matched KV prefix into target.slot so llama.cpp only prefills the tail. + Prepare(ctx context.Context, in *PrepareRequest, opts ...grpc.CallOption) (*PrepareReply, error) + // Commit = Save the KV currently held in target.slot and Put it to the store. + Commit(ctx context.Context, in *CommitRequest, opts ...grpc.CallOption) (*CommitReply, error) + Stats(ctx context.Context, in *StatsRequest, opts ...grpc.CallOption) (*StatsReply, error) +} + +type kVCacheBusClient struct { + cc grpc.ClientConnInterface +} + +func NewKVCacheBusClient(cc grpc.ClientConnInterface) KVCacheBusClient { + return &kVCacheBusClient{cc} +} + +func (c *kVCacheBusClient) Health(ctx context.Context, in *HealthRequest, opts ...grpc.CallOption) (*HealthReply, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(HealthReply) + err := c.cc.Invoke(ctx, KVCacheBus_Health_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *kVCacheBusClient) Lookup(ctx context.Context, in *LookupRequest, opts ...grpc.CallOption) (*LookupReply, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(LookupReply) + err := c.cc.Invoke(ctx, KVCacheBus_Lookup_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *kVCacheBusClient) Prepare(ctx context.Context, in *PrepareRequest, opts ...grpc.CallOption) (*PrepareReply, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(PrepareReply) + err := c.cc.Invoke(ctx, KVCacheBus_Prepare_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *kVCacheBusClient) Commit(ctx context.Context, in *CommitRequest, opts ...grpc.CallOption) (*CommitReply, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(CommitReply) + err := c.cc.Invoke(ctx, KVCacheBus_Commit_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *kVCacheBusClient) Stats(ctx context.Context, in *StatsRequest, opts ...grpc.CallOption) (*StatsReply, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(StatsReply) + err := c.cc.Invoke(ctx, KVCacheBus_Stats_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +// KVCacheBusServer is the server API for KVCacheBus service. +// All implementations must embed UnimplementedKVCacheBusServer +// for forward compatibility +type KVCacheBusServer interface { + Health(context.Context, *HealthRequest) (*HealthReply, error) + // Lookup is read-only: longest block-prefix match + arbiter decision, no I/O + // into llama.cpp. Useful for inspection / benchmarking. + Lookup(context.Context, *LookupRequest) (*LookupReply, error) + // Prepare = Lookup + (if a hit and the cost arbiter approves) restore the + // matched KV prefix into target.slot so llama.cpp only prefills the tail. + Prepare(context.Context, *PrepareRequest) (*PrepareReply, error) + // Commit = Save the KV currently held in target.slot and Put it to the store. + Commit(context.Context, *CommitRequest) (*CommitReply, error) + Stats(context.Context, *StatsRequest) (*StatsReply, error) + mustEmbedUnimplementedKVCacheBusServer() +} + +// UnimplementedKVCacheBusServer must be embedded to have forward compatible implementations. +type UnimplementedKVCacheBusServer struct { +} + +func (UnimplementedKVCacheBusServer) Health(context.Context, *HealthRequest) (*HealthReply, error) { + return nil, status.Errorf(codes.Unimplemented, "method Health not implemented") +} +func (UnimplementedKVCacheBusServer) Lookup(context.Context, *LookupRequest) (*LookupReply, error) { + return nil, status.Errorf(codes.Unimplemented, "method Lookup not implemented") +} +func (UnimplementedKVCacheBusServer) Prepare(context.Context, *PrepareRequest) (*PrepareReply, error) { + return nil, status.Errorf(codes.Unimplemented, "method Prepare not implemented") +} +func (UnimplementedKVCacheBusServer) Commit(context.Context, *CommitRequest) (*CommitReply, error) { + return nil, status.Errorf(codes.Unimplemented, "method Commit not implemented") +} +func (UnimplementedKVCacheBusServer) Stats(context.Context, *StatsRequest) (*StatsReply, error) { + return nil, status.Errorf(codes.Unimplemented, "method Stats not implemented") +} +func (UnimplementedKVCacheBusServer) mustEmbedUnimplementedKVCacheBusServer() {} + +// UnsafeKVCacheBusServer may be embedded to opt out of forward compatibility for this service. +// Use of this interface is not recommended, as added methods to KVCacheBusServer will +// result in compilation errors. +type UnsafeKVCacheBusServer interface { + mustEmbedUnimplementedKVCacheBusServer() +} + +func RegisterKVCacheBusServer(s grpc.ServiceRegistrar, srv KVCacheBusServer) { + s.RegisterService(&KVCacheBus_ServiceDesc, srv) +} + +func _KVCacheBus_Health_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(HealthRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(KVCacheBusServer).Health(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: KVCacheBus_Health_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(KVCacheBusServer).Health(ctx, req.(*HealthRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _KVCacheBus_Lookup_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(LookupRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(KVCacheBusServer).Lookup(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: KVCacheBus_Lookup_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(KVCacheBusServer).Lookup(ctx, req.(*LookupRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _KVCacheBus_Prepare_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(PrepareRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(KVCacheBusServer).Prepare(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: KVCacheBus_Prepare_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(KVCacheBusServer).Prepare(ctx, req.(*PrepareRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _KVCacheBus_Commit_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(CommitRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(KVCacheBusServer).Commit(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: KVCacheBus_Commit_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(KVCacheBusServer).Commit(ctx, req.(*CommitRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _KVCacheBus_Stats_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(StatsRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(KVCacheBusServer).Stats(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: KVCacheBus_Stats_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(KVCacheBusServer).Stats(ctx, req.(*StatsRequest)) + } + return interceptor(ctx, in, info, handler) +} + +// KVCacheBus_ServiceDesc is the grpc.ServiceDesc for KVCacheBus service. +// It's only intended for direct use with grpc.RegisterService, +// and not to be introspected or modified (even as a copy) +var KVCacheBus_ServiceDesc = grpc.ServiceDesc{ + ServiceName: "bridge.v1.KVCacheBus", + HandlerType: (*KVCacheBusServer)(nil), + Methods: []grpc.MethodDesc{ + { + MethodName: "Health", + Handler: _KVCacheBus_Health_Handler, + }, + { + MethodName: "Lookup", + Handler: _KVCacheBus_Lookup_Handler, + }, + { + MethodName: "Prepare", + Handler: _KVCacheBus_Prepare_Handler, + }, + { + MethodName: "Commit", + Handler: _KVCacheBus_Commit_Handler, + }, + { + MethodName: "Stats", + Handler: _KVCacheBus_Stats_Handler, + }, + }, + Streams: []grpc.StreamDesc{}, + Metadata: "bridge.proto", +} diff --git a/mooncake-integration/ollama/bridge/internal/cachekey/gguf_meta.go b/mooncake-integration/ollama/bridge/internal/cachekey/gguf_meta.go new file mode 100644 index 00000000..c3df2b2a --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/cachekey/gguf_meta.go @@ -0,0 +1,320 @@ +package cachekey + +import ( + "bufio" + "crypto/sha256" + "encoding/binary" + "encoding/hex" + "fmt" + "io" + "math" + "os" + "sort" + "strings" +) + +// GGUF metadata value types (see ggml-org/llama.cpp gguf spec). +const ( + ggufUint8 uint32 = iota + ggufInt8 + ggufUint16 + ggufInt16 + ggufUint32 + ggufInt32 + ggufFloat32 + ggufBool + ggufString + ggufArray + ggufUint64 + ggufInt64 + ggufFloat64 +) + +// GGUFMeta is the subset of GGUF metadata we parse. We deliberately read only +// the header KV block (not tensor data), so this is fast even for multi-GB +// files. +type GGUFMeta struct { + Path string + Version uint32 + Arch string + KV map[string]any // raw metadata kv (scalars + small arrays summarized) + metaHash string // sha256 of the raw metadata block bytes +} + +// ReadGGUFMeta parses just enough of a .gguf file to build a ModelFingerprint. +func ReadGGUFMeta(path string) (*GGUFMeta, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + br := bufio.NewReaderSize(f, 1<<20) + + // We hash every byte we consume from the metadata block to derive a stable + // model digest without hashing the (huge) tensor payload. + hsh := sha256.New() + r := io.TeeReader(br, hsh) + + var magic [4]byte + if _, err := io.ReadFull(r, magic[:]); err != nil { + return nil, err + } + if string(magic[:]) != "GGUF" { + return nil, fmt.Errorf("not a GGUF file (magic=%q)", magic[:]) + } + m := &GGUFMeta{Path: path, KV: map[string]any{}} + rd := &leReader{r: r} + m.Version = rd.u32() + _ = rd.u64() // tensor_count (we don't need tensors) + kvCount := rd.u64() + if rd.err != nil { + return nil, rd.err + } + for i := uint64(0); i < kvCount && rd.err == nil; i++ { + key := rd.str() + val := rd.value() + if key != "" { + m.KV[key] = val + } + } + if rd.err != nil && rd.err != io.EOF { + return nil, fmt.Errorf("gguf parse: %w", rd.err) + } + if a, ok := m.KV["general.architecture"].(string); ok { + m.Arch = a + } + m.metaHash = hex.EncodeToString(hsh.Sum(nil)) + return m, nil +} + +// Fingerprint derives a ModelFingerprint from parsed GGUF metadata. kvType is +// the runtime --cache-type-k/v (defaults to f16); swa marks sliding-window +// attention. modelDigest, if non-empty (e.g. an Ollama blob digest), overrides +// the metadata-hash fallback. +func (m *GGUFMeta) Fingerprint(modelDigest, kvType string, swa bool, blockSize int) ModelFingerprint { + if modelDigest == "" { + modelDigest = "ggufmeta-" + m.metaHash[:24] + } + if kvType == "" { + kvType = "f16" + } + layout := "dense" + if swa { + layout = "swa" + } + fp := ModelFingerprint{ + ModelDigest: modelDigest, + Arch: orUnknown(m.Arch), + TokenizerHash: m.tokenizerHash(), + RopeHash: m.ropeHash(), + KVType: kvType, + KVLayout: layout, + NCtxTrain: m.ctxLen(), + BlockSize: blockSize, + } + return fp +} + +func (m *GGUFMeta) archKey(suffix string) string { return m.Arch + "." + suffix } + +func (m *GGUFMeta) ctxLen() int { + if v, ok := toInt(m.KV[m.archKey("context_length")]); ok { + return v + } + return 0 +} + +// tokenizerHash folds the tokenizer model id and a compact signature of the +// vocabulary (size + a sample fingerprint) so different tokenizers never share +// keys, without serializing the whole vocab. +func (m *GGUFMeta) tokenizerHash() string { + h := sha256.New() + write := func(k string) { + if v, ok := m.KV[k]; ok { + fmt.Fprintf(h, "%s=%v;", k, summarize(v)) + } + } + write("tokenizer.ggml.model") + write("tokenizer.ggml.pre") + write("tokenizer.ggml.bos_token_id") + write("tokenizer.ggml.eos_token_id") + write("tokenizer.ggml.add_bos_token") + write("tokenizer.chat_template") + return hex.EncodeToString(h.Sum(nil)[:12]) +} + +func (m *GGUFMeta) ropeHash() string { + h := sha256.New() + for _, k := range []string{ + m.archKey("rope.dimension_count"), + m.archKey("rope.freq_base"), + m.archKey("rope.scaling.type"), + m.archKey("rope.scaling.factor"), + m.archKey("rope.scaling.original_context_length"), + m.archKey("embedding_length"), + m.archKey("attention.head_count"), + m.archKey("attention.head_count_kv"), + m.archKey("block_count"), + } { + if v, ok := m.KV[k]; ok { + fmt.Fprintf(h, "%s=%v;", k, summarize(v)) + } + } + return hex.EncodeToString(h.Sum(nil)[:12]) +} + +// ---- low-level little-endian GGUF reader ---- + +type leReader struct { + r io.Reader + err error + b8 [8]byte +} + +func (z *leReader) read(n int) []byte { + if z.err != nil { + return nil + } + buf := make([]byte, n) + if _, err := io.ReadFull(z.r, buf); err != nil { + z.err = err + return nil + } + return buf +} +func (z *leReader) u32() uint32 { + if z.err != nil { + return 0 + } + if _, err := io.ReadFull(z.r, z.b8[:4]); err != nil { + z.err = err + return 0 + } + return binary.LittleEndian.Uint32(z.b8[:4]) +} +func (z *leReader) u64() uint64 { + if z.err != nil { + return 0 + } + if _, err := io.ReadFull(z.r, z.b8[:8]); err != nil { + z.err = err + return 0 + } + return binary.LittleEndian.Uint64(z.b8[:8]) +} +func (z *leReader) str() string { + n := z.u64() + if z.err != nil || n > (1<<28) { // guard against corrupt lengths + if n > (1 << 28) { + z.err = fmt.Errorf("gguf string too long: %d", n) + } + return "" + } + return string(z.read(int(n))) +} + +// value reads a single typed metadata value. Arrays are summarized (we keep +// length + first few elements) to avoid materializing huge vocab arrays. +func (z *leReader) value() any { + t := z.u32() + return z.valueOfType(t) +} + +func (z *leReader) valueOfType(t uint32) any { + switch t { + case ggufUint8: + return uint64(z.read(1)[0]) + case ggufInt8: + return int64(int8(z.read(1)[0])) + case ggufUint16: + b := z.read(2) + if b == nil { + return nil + } + return uint64(binary.LittleEndian.Uint16(b)) + case ggufInt16: + b := z.read(2) + if b == nil { + return nil + } + return int64(int16(binary.LittleEndian.Uint16(b))) + case ggufUint32: + return uint64(z.u32()) + case ggufInt32: + return int64(int32(z.u32())) + case ggufFloat32: + return float64(math.Float32frombits(z.u32())) + case ggufBool: + b := z.read(1) + if b == nil { + return nil + } + return b[0] != 0 + case ggufString: + return z.str() + case ggufUint64: + return z.u64() + case ggufInt64: + return int64(z.u64()) + case ggufFloat64: + return math.Float64frombits(z.u64()) + case ggufArray: + et := z.u32() + n := z.u64() + if z.err != nil { + return nil + } + // Summarize: keep length + up to 8 sample elements so tokenizer arrays + // don't blow up memory but still contribute to the signature. + const sample = 8 + arr := arraySummary{ElemType: et, Len: n} + for i := uint64(0); i < n && z.err == nil; i++ { + v := z.valueOfType(et) + if i < sample { + arr.Head = append(arr.Head, v) + } + } + return arr + default: + z.err = fmt.Errorf("gguf: unknown value type %d", t) + return nil + } +} + +type arraySummary struct { + ElemType uint32 + Len uint64 + Head []any +} + +func summarize(v any) string { + if a, ok := v.(arraySummary); ok { + parts := make([]string, 0, len(a.Head)) + for _, e := range a.Head { + parts = append(parts, fmt.Sprintf("%v", e)) + } + sort.Strings(parts) + return fmt.Sprintf("arr(t=%d,n=%d)[%s]", a.ElemType, a.Len, strings.Join(parts, ",")) + } + return fmt.Sprintf("%v", v) +} + +func toInt(v any) (int, bool) { + switch x := v.(type) { + case uint64: + return int(x), true + case int64: + return int(x), true + case float64: + return int(x), true + default: + return 0, false + } +} + +func orUnknown(s string) string { + if s == "" { + return "unknown" + } + return s +} diff --git a/mooncake-integration/ollama/bridge/internal/cachekey/key.go b/mooncake-integration/ollama/bridge/internal/cachekey/key.go new file mode 100644 index 00000000..23d092f3 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/cachekey/key.go @@ -0,0 +1,122 @@ +// Package cachekey builds the cache namespace and per-block chained hashes that +// uniquely (and *safely*) identify a reusable KV-cache prefix. +// +// The cardinal rule: KV produced by a different model / tokenizer / RoPE / KV +// dtype / KV layout MUST NEVER be reused. A wrong reuse silently poisons +// generation and is far worse than a miss. Therefore every field that can change +// the numerical content of the KV cache is folded into the key. We would rather +// miss than mis-hit. +package cachekey + +import ( + "crypto/sha256" + "encoding/binary" + "encoding/hex" + "fmt" + "strings" +) + +// ModelFingerprint captures every attribute that affects the *bytes* of a KV +// cache entry. Two requests may share a cached prefix only if their +// fingerprints are byte-identical. +type ModelFingerprint struct { + ModelDigest string // ollama blob digest, or a stable hash of the GGUF metadata block + Arch string // general.architecture (llama, qwen2, ...) + TokenizerHash string // hash over tokenizer model + vocab signature + RopeHash string // hash over rope dims / freq base / scaling + KVType string // f16 | q8_0 | q4_0 | bf16 ... (--cache-type-k/v) + KVLayout string // dense | swa (sliding-window attention changes layout) + NCtxTrain int // training context length + BlockSize int // token granularity of a cache block +} + +// Valid reports whether the fingerprint is populated enough to be trusted. +// An under-specified fingerprint must force a miss, never a risky hit. +func (m ModelFingerprint) Valid() bool { + return m.ModelDigest != "" && m.Arch != "" && m.KVType != "" && m.BlockSize > 0 +} + +// Prefix is the stable, human-readable namespace shared by every block key that +// belongs to this fingerprint + user namespace. BlockKey appends the trailing +// {block_index}:{prefix_block_hash}. +// +// ollama:{model_digest}:{tokenizer_hash}:{rope_hash}:{kv_dtype}:{kv_layout}:{n_ctx_train}:{block_size}:{user_ns} +func (m ModelFingerprint) Prefix(userNamespace string) string { + if userNamespace == "" { + userNamespace = "_" + } + // sanitize: ':' is our delimiter, so it may not appear inside a field. + clean := func(s string) string { return strings.ReplaceAll(s, ":", "_") } + return strings.Join([]string{ + "ollama", + clean(m.ModelDigest), + clean(emptyTo(m.TokenizerHash, "-")), + clean(emptyTo(m.RopeHash, "-")), + clean(m.KVType), + clean(emptyTo(m.KVLayout, "dense")), + fmt.Sprintf("%d", m.NCtxTrain), + fmt.Sprintf("%d", m.BlockSize), + clean(userNamespace), + }, ":") +} + +// BlockKey is the store object key for the prefix that ends at block boundary +// blockIndex (i.e. it covers blocks 0..blockIndex-1, = blockIndex*BlockSize +// tokens). chainHashHex is the chained hash returned by ChainBlockHashes for +// that boundary. +func (m ModelFingerprint) BlockKey(userNamespace string, blockIndex int, chainHashHex string) string { + return fmt.Sprintf("%s:%d:%s", m.Prefix(userNamespace), blockIndex, chainHashHex) +} + +// Chain holds, for a token sequence, the chained hash at every *full* block +// boundary. Chain.Hex[i] is the hash covering tokens[0 : (i+1)*BlockSize], i.e. +// the key material for a prefix of (i+1) blocks. A trailing partial block (if +// any) is intentionally NOT hashed: only block-aligned prefixes are reusable, +// which keeps keys deterministic across requests that share a long head. +type Chain struct { + BlockSize int + NumTokens int + FullBlocks int // len(Hex) + Hex []string // 16-byte (128-bit) truncated sha256, hex-encoded +} + +// ChainBlockHashes computes the chained, content-addressed hash of each full +// block of tokens, seeded by `seed` (the fingerprint Prefix), so that the same +// tokens under a different model produce different hashes. +// +// h_0 = H( seed_bytes || block_0_bytes ) +// h_i = H( h_{i-1} || block_i_bytes ) +// +// Chaining guarantees prefix consistency: two sequences agree on h_i iff they +// agree on every token of blocks 0..i. This is what makes longest-prefix +// matching correct. +func ChainBlockHashes(seed string, tokens []int32, blockSize int) Chain { + if blockSize <= 0 { + blockSize = 256 + } + full := len(tokens) / blockSize + out := Chain{BlockSize: blockSize, NumTokens: len(tokens), FullBlocks: full, Hex: make([]string, 0, full)} + + prev := sha256.Sum256([]byte("omb-kvcache-v1\x00" + seed)) + buf := make([]byte, blockSize*4) // 4 bytes per int32 token + for i := 0; i < full; i++ { + blk := tokens[i*blockSize : (i+1)*blockSize] + for j, t := range blk { + binary.LittleEndian.PutUint32(buf[j*4:], uint32(t)) + } + h := sha256.New() + h.Write(prev[:]) + h.Write(buf) + sum := h.Sum(nil) + copy(prev[:], sum) // chain forward (full 32 bytes carried, 16 emitted) + out.Hex = append(out.Hex, hex.EncodeToString(sum[:16])) + } + return out +} + +func emptyTo(s, d string) string { + if s == "" { + return d + } + return s +} diff --git a/mooncake-integration/ollama/bridge/internal/cachekey/key_test.go b/mooncake-integration/ollama/bridge/internal/cachekey/key_test.go new file mode 100644 index 00000000..f38676d2 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/cachekey/key_test.go @@ -0,0 +1,91 @@ +package cachekey + +import ( + "os" + "testing" +) + +func TestChainBlockHashes_Determinism(t *testing.T) { + toks := make([]int32, 1000) + for i := range toks { + toks[i] = int32(i * 7 % 50000) + } + a := ChainBlockHashes("seed", toks, 256) + b := ChainBlockHashes("seed", toks, 256) + if a.FullBlocks != 1000/256 { + t.Fatalf("FullBlocks = %d, want %d", a.FullBlocks, 1000/256) + } + for i := range a.Hex { + if a.Hex[i] != b.Hex[i] { + t.Fatalf("nondeterministic at %d", i) + } + } +} + +func TestChainBlockHashes_PrefixConsistency(t *testing.T) { + // Two sequences sharing the first 3 blocks must agree on the first 3 chain + // hashes and (almost surely) differ afterwards. + base := make([]int32, 4*256) + for i := range base { + base[i] = int32(i) + } + x := append([]int32{}, base...) + y := append([]int32{}, base...) + // diverge in block index 3 (token 3*256) + y[3*256] = 999999 + cx := ChainBlockHashes("s", x, 256) + cy := ChainBlockHashes("s", y, 256) + for i := 0; i < 3; i++ { + if cx.Hex[i] != cy.Hex[i] { + t.Fatalf("shared block %d hashes differ", i) + } + } + if cx.Hex[3] == cy.Hex[3] { + t.Fatalf("divergent block 3 hashes collided") + } +} + +func TestChainBlockHashes_SeedSeparation(t *testing.T) { + toks := make([]int32, 512) + a := ChainBlockHashes("modelA", toks, 256) + b := ChainBlockHashes("modelB", toks, 256) + if a.Hex[0] == b.Hex[0] { + t.Fatalf("different seeds (models) must not share hashes") + } +} + +func TestFingerprintPrefixStable(t *testing.T) { + fp := ModelFingerprint{ModelDigest: "d", Arch: "qwen2", KVType: "f16", BlockSize: 256} + if !fp.Valid() { + t.Fatal("expected valid fingerprint") + } + if got := fp.Prefix("ns"); got == "" { + t.Fatal("empty prefix") + } + // ':' inside a field must be sanitized so it can't break the delimiter. + fp2 := fp + fp2.ModelDigest = "a:b:c" + if p := fp2.Prefix("ns"); p == fp.Prefix("ns") { + t.Fatal("digest with colons should differ") + } +} + +// Parses the real model if OMB_TEST_GGUF points at a .gguf file. +func TestReadGGUFMeta_Real(t *testing.T) { + path := os.Getenv("OMB_TEST_GGUF") + if path == "" { + t.Skip("set OMB_TEST_GGUF to a .gguf to run") + } + m, err := ReadGGUFMeta(path) + if err != nil { + t.Fatalf("ReadGGUFMeta: %v", err) + } + if m.Arch == "" { + t.Fatalf("arch empty") + } + fp := m.Fingerprint("", "f16", false, 256) + if !fp.Valid() { + t.Fatalf("incomplete fingerprint: %+v", fp) + } + t.Logf("arch=%s ctx=%d tok=%s rope=%s digest=%s", fp.Arch, fp.NCtxTrain, fp.TokenizerHash, fp.RopeHash, fp.ModelDigest) +} diff --git a/mooncake-integration/ollama/bridge/internal/llamabridge/slots_http.go b/mooncake-integration/ollama/bridge/internal/llamabridge/slots_http.go new file mode 100644 index 00000000..fd023965 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/llamabridge/slots_http.go @@ -0,0 +1,205 @@ +// Package llamabridge is the Stage-1 client for a running llama.cpp server. It +// drives exactly the endpoints the reuse flow relies on: +// +// POST /tokenize -> token ids (so we can hash blocks) +// POST /completion {cache_prompt,...} -> content + timings.prompt_n (prefilled) +// POST /slots/{id}?action=save -> {n_saved, n_written} +// POST /slots/{id}?action=restore -> {n_restored, n_read} +// POST /slots/{id}?action=erase +// GET /props -> model/runtime info for fingerprinting +// +// The save file lands in the server's --slot-save-path; the sidecar then ships +// that file to the Mooncake Store (and vice-versa on restore). No llama.cpp +// source changes are required for Stage 1. +package llamabridge + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "time" +) + +type Client struct { + hc *http.Client +} + +func New() *Client { + return &Client{hc: &http.Client{Timeout: 0}} // long ops (big prefills) use ctx deadlines +} + +func (c *Client) postJSON(ctx context.Context, url string, body any, out any) error { + var buf bytes.Buffer + if body != nil { + if err := json.NewEncoder(&buf).Encode(body); err != nil { + return err + } + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, &buf) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + resp, err := c.hc.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + data, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("%s -> %d: %s", url, resp.StatusCode, truncate(string(data), 256)) + } + if out != nil { + if err := json.Unmarshal(data, out); err != nil { + return fmt.Errorf("decode %s: %w (body=%s)", url, err, truncate(string(data), 256)) + } + } + return nil +} + +func (c *Client) getJSON(ctx context.Context, url string, out any) error { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return err + } + resp, err := c.hc.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + data, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("%s -> %d: %s", url, resp.StatusCode, truncate(string(data), 256)) + } + return json.Unmarshal(data, out) +} + +// Tokenize returns the token ids for text. add_special controls BOS/EOS. +func (c *Client) Tokenize(ctx context.Context, baseURL, text string, addSpecial bool) ([]int32, error) { + var out struct { + Tokens []int32 `json:"tokens"` + } + err := c.postJSON(ctx, baseURL+"/tokenize", map[string]any{ + "content": text, "add_special": addSpecial, "with_pieces": false, + }, &out) + return out.Tokens, err +} + +// CompletionTimings mirrors llama.cpp's timings object. +type CompletionTimings struct { + PromptN int `json:"prompt_n"` // tokens actually prefilled (NOT served from cache) + PromptMs float64 `json:"prompt_ms"` // prefill wall time + PredictedN int `json:"predicted_n"` // tokens generated + PredictedMs float64 `json:"predicted_ms"` // generation wall time +} + +type CompletionResult struct { + Content string `json:"content"` + Timings CompletionTimings `json:"timings"` + TokensEvaluated int `json:"tokens_evaluated"` + Truncated bool `json:"truncated"` +} + +// CompletionRequest is the subset we use. Prompt may be a string or []int32. +type CompletionRequest struct { + Prompt any `json:"prompt"` + IDSlot int `json:"id_slot"` + CachePrompt bool `json:"cache_prompt"` + NPredict int `json:"n_predict"` + Temperature float64 `json:"temperature"` + Seed int `json:"seed,omitempty"` + Stream bool `json:"stream"` +} + +func (c *Client) Completion(ctx context.Context, baseURL string, req CompletionRequest) (*CompletionResult, error) { + req.Stream = false + var out CompletionResult + if err := c.postJSON(ctx, baseURL+"/completion", req, &out); err != nil { + return nil, err + } + return &out, nil +} + +// SaveResult / RestoreResult expose token + byte counts. +type SaveResult struct { + IDSlot int `json:"id_slot"` + Filename string `json:"filename"` + NSaved int `json:"n_saved"` // tokens + NWritten uint64 `json:"n_written"` // bytes +} +type RestoreResult struct { + IDSlot int `json:"id_slot"` + Filename string `json:"filename"` + NRestored int `json:"n_restored"` // tokens + NRead uint64 `json:"n_read"` // bytes +} + +func (c *Client) SaveSlot(ctx context.Context, baseURL string, slot int, filename string) (*SaveResult, error) { + url := fmt.Sprintf("%s/slots/%d?action=save", baseURL, slot) + var out SaveResult + if err := c.postJSON(ctx, url, map[string]any{"filename": filename}, &out); err != nil { + return nil, err + } + return &out, nil +} + +func (c *Client) RestoreSlot(ctx context.Context, baseURL string, slot int, filename string) (*RestoreResult, error) { + url := fmt.Sprintf("%s/slots/%d?action=restore", baseURL, slot) + var out RestoreResult + if err := c.postJSON(ctx, url, map[string]any{"filename": filename}, &out); err != nil { + return nil, err + } + return &out, nil +} + +func (c *Client) EraseSlot(ctx context.Context, baseURL string, slot int) error { + url := fmt.Sprintf("%s/slots/%d?action=erase", baseURL, slot) + return c.postJSON(ctx, url, map[string]any{}, nil) +} + +// Props is a subset of GET /props used for fingerprinting / health. +type Props struct { + DefaultGenerationSettings struct { + NCtx int `json:"n_ctx"` + } `json:"default_generation_settings"` + ModelPath string `json:"model_path"` + TotalSlots int `json:"total_slots"` +} + +func (c *Client) GetProps(ctx context.Context, baseURL string) (*Props, error) { + var out Props + if err := c.getJSON(ctx, baseURL+"/props", &out); err != nil { + return nil, err + } + return &out, nil +} + +// WaitHealthy polls /health until ready or ctx expires. +func (c *Client) WaitHealthy(ctx context.Context, baseURL string) error { + for { + req, _ := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+"/health", nil) + resp, err := c.hc.Do(req) + if err == nil { + io.Copy(io.Discard, resp.Body) + resp.Body.Close() + if resp.StatusCode == http.StatusOK { + return nil + } + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(300 * time.Millisecond): + } + } +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] + "..." +} diff --git a/mooncake-integration/ollama/bridge/internal/metrics/prom.go b/mooncake-integration/ollama/bridge/internal/metrics/prom.go new file mode 100644 index 00000000..a8913069 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/metrics/prom.go @@ -0,0 +1,95 @@ +// Package metrics exposes the sidecar's Prometheus instrumentation. The primary +// series is mooncake_bridge_saved_prefill_tokens_total — the number of prompt +// tokens that did NOT have to be re-prefilled because their KV was restored from +// the Mooncake Store. A Grafana panel integrates this live. +package metrics + +import ( + "sync/atomic" + + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" +) + +type Metrics struct { + LookupTotal prometheus.Counter + PrepareTotal prometheus.Counter + CommitTotal prometheus.Counter + Hits prometheus.Counter + Misses prometheus.Counter + HitBlocks prometheus.Counter + MissBlocks prometheus.Counter + RestoreCount prometheus.Counter + RecomputeChosen prometheus.Counter + BytesGet prometheus.Counter + BytesPut prometheus.Counter + SavedPrefillTok prometheus.Counter + + RestoreLatency prometheus.Histogram + SaveLatency prometheus.Histogram + StoreGetLatency prometheus.Histogram + StorePutLatency prometheus.Histogram + + GetGBps prometheus.Gauge + PrefillToksS prometheus.Gauge + IndexSnapshots prometheus.Gauge + IndexBytes prometheus.Gauge + + // plain atomics mirrored into Stats RPC + savedTokens uint64 + bytesGet uint64 + bytesPut uint64 +} + +func New(reg prometheus.Registerer) *Metrics { + af := promauto.With(reg) + lat := func(name, help string) prometheus.Histogram { + return af.NewHistogram(prometheus.HistogramOpts{ + Name: name, + Help: help, + Buckets: []float64{0.5, 1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2500, 5000, 10000}, + }) + } + ctr := func(name, help string) prometheus.Counter { + return af.NewCounter(prometheus.CounterOpts{Name: name, Help: help}) + } + g := func(name, help string) prometheus.Gauge { + return af.NewGauge(prometheus.GaugeOpts{Name: name, Help: help}) + } + return &Metrics{ + LookupTotal: ctr("mooncake_bridge_lookup_total", "Lookup/Prepare requests served"), + PrepareTotal: ctr("mooncake_bridge_prepare_total", "Prepare requests served"), + CommitTotal: ctr("mooncake_bridge_commit_total", "Commit (save) requests served"), + Hits: ctr("mooncake_bridge_hits_total", "Requests with a usable cached prefix"), + Misses: ctr("mooncake_bridge_misses_total", "Requests with no usable cached prefix"), + HitBlocks: ctr("mooncake_bridge_hit_blocks_total", "Blocks served from cache"), + MissBlocks: ctr("mooncake_bridge_miss_blocks_total", "Blocks that had to be (re)computed"), + RestoreCount: ctr("mooncake_bridge_restore_total", "KV restores performed"), + RecomputeChosen: ctr("mooncake_bridge_recompute_chosen_total", "Times the arbiter chose recompute over restore"), + BytesGet: ctr("mooncake_bridge_bytes_get_total", "Bytes read from the store"), + BytesPut: ctr("mooncake_bridge_bytes_put_total", "Bytes written to the store"), + SavedPrefillTok: ctr("mooncake_bridge_saved_prefill_tokens_total", "Prompt tokens NOT re-prefilled thanks to KV reuse"), + RestoreLatency: lat("mooncake_bridge_restore_latency_ms", "llama.cpp /slots restore latency (ms)"), + SaveLatency: lat("mooncake_bridge_save_latency_ms", "llama.cpp /slots save latency (ms)"), + StoreGetLatency: lat("mooncake_bridge_store_get_latency_ms", "Store GetFile latency (ms)"), + StorePutLatency: lat("mooncake_bridge_store_put_latency_ms", "Store PutFile latency (ms)"), + GetGBps: g("mooncake_bridge_learned_get_gbps", "Arbiter's learned store read bandwidth (GB/s)"), + PrefillToksS: g("mooncake_bridge_learned_prefill_tps", "Arbiter's learned prefill rate (tokens/s)"), + IndexSnapshots: g("mooncake_bridge_index_snapshots", "KV prefix snapshots tracked in the radix index"), + IndexBytes: g("mooncake_bridge_index_bytes", "Total bytes of snapshots tracked"), + } +} + +func (m *Metrics) AddSavedTokens(n int) { + if n <= 0 { + return + } + m.SavedPrefillTok.Add(float64(n)) + atomic.AddUint64(&m.savedTokens, uint64(n)) +} +func (m *Metrics) AddBytesGet(n uint64) { m.BytesGet.Add(float64(n)); atomic.AddUint64(&m.bytesGet, n) } +func (m *Metrics) AddBytesPut(n uint64) { m.BytesPut.Add(float64(n)); atomic.AddUint64(&m.bytesPut, n) } + +func (m *Metrics) SavedTokens() uint64 { return atomic.LoadUint64(&m.savedTokens) } +func (m *Metrics) TotalBytesGet() uint64 { return atomic.LoadUint64(&m.bytesGet) } +func (m *Metrics) TotalBytesPut() uint64 { return atomic.LoadUint64(&m.bytesPut) } diff --git a/mooncake-integration/ollama/bridge/internal/orchestrator/reuse.go b/mooncake-integration/ollama/bridge/internal/orchestrator/reuse.go new file mode 100644 index 00000000..616e5cb0 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/orchestrator/reuse.go @@ -0,0 +1,494 @@ +// Package orchestrator implements the three-stage KV reuse flow: +// +// Lookup compute block-prefix keys, find the longest prefix present in the +// store (authoritative, batched), and consult the radix index. +// Prepare Lookup + (if the cost arbiter approves) GetFile the matched KV and +// restore it into a llama.cpp slot, so only the tail is prefilled. +// Commit Save the slot KV and PutFile it to the store under its block key, +// with single-writer dedup so concurrent agents store a shared prefix +// exactly once. +// +// It owns the cross-cutting policy: cache-key construction, longest-prefix +// matching, the restore-vs-recompute arbiter, the radix index, per-model +// bytes/token learning, and Prometheus accounting. +package orchestrator + +import ( + "context" + "fmt" + "os" + "path/filepath" + "sync" + "sync/atomic" + "time" + + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/arbiter" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/cachekey" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/llamabridge" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/metrics" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/prefixindex" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/store" +) + +type Config struct { + SlotSavePath string // == llama.cpp --slot-save-path (shared filesystem) + DefaultBlockSize int + DefaultReplicaNum uint32 + MinPrefixBlocks int + CleanupFiles bool // delete the local save file after store put / after restore +} + +func (c *Config) defaults() { + if c.DefaultBlockSize <= 0 { + c.DefaultBlockSize = 256 + } + if c.DefaultReplicaNum == 0 { + c.DefaultReplicaNum = 1 + } + if c.MinPrefixBlocks <= 0 { + c.MinPrefixBlocks = 1 + } +} + +type Orchestrator struct { + store store.Backend + llama *llamabridge.Client + index *prefixindex.Index + arb *arbiter.Arbiter + mx *metrics.Metrics + cfg Config + + fpMu sync.Mutex + fpCache map[string]cachekey.ModelFingerprint // gguf path -> completed fingerprint + bptMu sync.Mutex + bytesPerTok map[string]float64 // fingerprint prefix -> learned KV bytes/token + nonce atomic.Uint64 // makes per-request slot filenames unique +} + +func New(b store.Backend, arb *arbiter.Arbiter, mx *metrics.Metrics, cfg Config) *Orchestrator { + cfg.defaults() + return &Orchestrator{ + store: b, llama: llamabridge.New(), index: prefixindex.New(), + arb: arb, mx: mx, cfg: cfg, + fpCache: map[string]cachekey.ModelFingerprint{}, + bytesPerTok: map[string]float64{}, + } +} + +func (o *Orchestrator) Index() *prefixindex.Index { return o.index } +func (o *Orchestrator) Arbiter() *arbiter.Arbiter { return o.arb } + +// ---- request/response types (decoupled from gRPC wire types) ---- + +type Policy struct { + Enable bool + Namespace string + Read bool + Write bool + BlockSize int + ReplicaNum uint32 + SoftPin bool + MinPrefixBlocks int +} + +type Target struct { + BaseURL string + Slot int +} + +type LookupResult struct { + Hit bool + MatchedBlocks int + MatchedTokens int + TotalBlocks int + TotalTokens int + Decision string // restore | recompute | miss + MatchedKey string + Reason string +} + +type PrepareResult struct { + LookupResult + Restored bool + RestoredTokens int + RestoreMs float64 + StoreGetMs float64 + Bytes uint64 +} + +type CommitResult struct { + OK bool + Stored bool + StoredBlocks int + StoredTokens int + Bytes uint64 + SaveMs float64 + StorePutMs float64 + Key string + Reason string +} + +// completeFingerprint fills empty fields by parsing the GGUF at ModelPath +// (cached). The completed fingerprint is what every key derives from, enforcing +// the "different model/tokenizer/rope => different key space" safety rule. +func (o *Orchestrator) completeFingerprint(fp cachekey.ModelFingerprint, modelPath string, swa bool) cachekey.ModelFingerprint { + if fp.BlockSize <= 0 { + fp.BlockSize = o.cfg.DefaultBlockSize + } + needEnrich := fp.Arch == "" || fp.TokenizerHash == "" || fp.RopeHash == "" + if modelPath == "" || !needEnrich { + if fp.KVType == "" { + fp.KVType = "f16" + } + if fp.ModelDigest == "" { + fp.ModelDigest = "unknown" + } + return fp + } + o.fpMu.Lock() + defer o.fpMu.Unlock() + if cached, ok := o.fpCache[modelPath]; ok { + // keep caller-provided block size / kv type + cached.BlockSize = fp.BlockSize + if fp.KVType != "" { + cached.KVType = fp.KVType + } + return cached + } + meta, err := cachekey.ReadGGUFMeta(modelPath) + if err != nil { + if fp.KVType == "" { + fp.KVType = "f16" + } + if fp.ModelDigest == "" { + fp.ModelDigest = "ggufpath-" + filepath.Base(modelPath) + } + return fp + } + kv := fp.KVType + if kv == "" { + kv = "f16" + } + full := meta.Fingerprint(fp.ModelDigest, kv, swa, fp.BlockSize) + o.fpCache[modelPath] = full + return full +} + +type planned struct { + fp cachekey.ModelFingerprint + rootKey string + modelKey string // per-model arbiter key (model digest + kv type), namespace-independent + chain cachekey.Chain + keys []string // keys[i] = boundary i+1 (i.e. prefix of i+1 blocks) +} + +func (o *Orchestrator) plan(fp cachekey.ModelFingerprint, pol Policy, tokens []int32, modelPath string) planned { + full := o.completeFingerprint(fp, modelPath, pol.BlockSize < 0) + bs := full.BlockSize + if pol.BlockSize > 0 { + bs = pol.BlockSize + full.BlockSize = bs + } + rootKey := full.Prefix(pol.Namespace) + modelKey := full.ModelDigest + "|" + full.KVType + "|" + full.Arch + chain := cachekey.ChainBlockHashes(rootKey, tokens, bs) + keys := make([]string, chain.FullBlocks) + for i := 0; i < chain.FullBlocks; i++ { + keys[i] = full.BlockKey(pol.Namespace, i+1, chain.Hex[i]) + } + return planned{fp: full, rootKey: rootKey, modelKey: modelKey, chain: chain, keys: keys} +} + +// longestPresent returns the largest boundary M (in blocks) whose key exists in +// the store, querying all boundaries in one batched call (multi-node correct). +func (o *Orchestrator) longestPresent(ctx context.Context, p planned) (int, string, error) { + if len(p.keys) == 0 { + return 0, "", nil + } + present, err := o.store.Exists(ctx, p.keys) + if err != nil { + return 0, "", err + } + for i := len(present) - 1; i >= 0; i-- { + if i < len(present) && present[i] == 1 { + return i + 1, p.keys[i], nil + } + } + return 0, "", nil +} + +func (o *Orchestrator) learnedBytesPerTok(rootKey string, fallbackTokens int, fallbackBytes uint64) float64 { + o.bptMu.Lock() + defer o.bptMu.Unlock() + if v, ok := o.bytesPerTok[rootKey]; ok && v > 0 { + return v + } + if fallbackTokens > 0 && fallbackBytes > 0 { + return float64(fallbackBytes) / float64(fallbackTokens) + } + return 0 +} + +func (o *Orchestrator) updateBytesPerTok(rootKey string, tokens int, bytes uint64) { + if tokens <= 0 || bytes == 0 { + return + } + o.bptMu.Lock() + defer o.bptMu.Unlock() + bpt := float64(bytes) / float64(tokens) + if old, ok := o.bytesPerTok[rootKey]; ok { + o.bytesPerTok[rootKey] = 0.5*old + 0.5*bpt + } else { + o.bytesPerTok[rootKey] = bpt + } +} + +// Lookup is read-only: longest-prefix match + arbiter decision, no llama I/O. +func (o *Orchestrator) Lookup(ctx context.Context, fp cachekey.ModelFingerprint, pol Policy, tokens []int32, modelPath string) (LookupResult, error) { + o.mx.LookupTotal.Inc() + p := o.plan(fp, pol, tokens, modelPath) + res := LookupResult{TotalBlocks: p.chain.FullBlocks, TotalTokens: p.chain.FullBlocks * p.fp.BlockSize} + if p.chain.FullBlocks == 0 { + res.Decision = "miss" + res.Reason = "prompt shorter than one block" + return res, nil + } + m, key, err := o.longestPresent(ctx, p) + if err != nil { + return res, err + } + if m == 0 { + res.Decision = "miss" + res.Reason = "no cached prefix" + return res, nil + } + matchedTokens := m * p.fp.BlockSize + // size estimate for the arbiter + im := o.index.LongestMatch(p.rootKey, p.chain.Hex, false) + var estBytes uint64 + if im.Found && im.Blocks == m { + estBytes = im.Bytes + } else { + bpt := o.learnedBytesPerTok(p.rootKey, 0, 0) + estBytes = uint64(float64(matchedTokens) * bpt) + } + dec := o.arb.Decide(p.modelKey, matchedTokens, estBytes) + res.Hit = true + res.MatchedBlocks = m + res.MatchedTokens = matchedTokens + res.MatchedKey = key + res.Reason = dec.Reason + if dec.Restore { + res.Decision = "restore" + } else { + res.Decision = "recompute" + } + return res, nil +} + +// Prepare runs Lookup and, if the arbiter approves and a target is given, +// restores the matched KV into target.Slot. +func (o *Orchestrator) Prepare(ctx context.Context, fp cachekey.ModelFingerprint, pol Policy, tokens []int32, modelPath string, tgt *Target) (PrepareResult, error) { + o.mx.PrepareTotal.Inc() + prepStart := time.Now() // full restore cost incl lookup + transfer + GPU load + orchestration + p := o.plan(fp, pol, tokens, modelPath) + out := PrepareResult{} + out.TotalBlocks = p.chain.FullBlocks + out.TotalTokens = p.chain.FullBlocks * p.fp.BlockSize + + if !pol.Read || p.chain.FullBlocks == 0 { + out.Decision = "miss" + out.Reason = "read disabled or prompt < 1 block" + o.mx.Misses.Inc() + o.mx.MissBlocks.Add(float64(p.chain.FullBlocks)) + return out, nil + } + m, key, err := o.longestPresent(ctx, p) + if err != nil { + return out, err + } + if m == 0 { + out.Decision = "miss" + out.Reason = "no cached prefix" + o.mx.Misses.Inc() + o.mx.MissBlocks.Add(float64(p.chain.FullBlocks)) + return out, nil + } + out.Hit = true + out.MatchedBlocks = m + out.MatchedTokens = m * p.fp.BlockSize + out.MatchedKey = key + + im := o.index.LongestMatch(p.rootKey, p.chain.Hex, true) + var estBytes uint64 + if im.Found && im.Blocks == m { + estBytes = im.Bytes + } else { + estBytes = uint64(float64(out.MatchedTokens) * o.learnedBytesPerTok(p.rootKey, 0, 0)) + } + dec := o.arb.Decide(p.modelKey, out.MatchedTokens, estBytes) + out.Reason = dec.Reason + if !dec.Restore { + out.Decision = "recompute" + o.mx.RecomputeChosen.Inc() + o.mx.MissBlocks.Add(float64(p.chain.FullBlocks)) // will be recomputed + return out, nil + } + out.Decision = "restore" + if tgt == nil || tgt.BaseURL == "" { + // plan-only (no target): report the decision without doing I/O. + return out, nil + } + + // ---- Load stage ---- + fname := fmt.Sprintf("omb-r-%x-%d.bin", hashKey(key), o.nonce.Add(1)) + abspath := filepath.Join(o.cfg.SlotSavePath, fname) + gr, err := o.store.GetFile(ctx, key, abspath) + if err != nil { + return out, fmt.Errorf("store get: %w", err) + } + if !gr.Found { + // raced with eviction; degrade to miss + out.Decision = "miss" + out.Hit = false + out.Reason = "matched key vanished (evicted); recompute" + o.mx.Misses.Inc() + o.mx.MissBlocks.Add(float64(p.chain.FullBlocks)) + return out, nil + } + out.StoreGetMs = gr.ElapsedMs + out.Bytes = gr.Bytes + o.mx.StoreGetLatency.Observe(gr.ElapsedMs) + o.mx.AddBytesGet(gr.Bytes) + + tRestore := time.Now() + rr, err := o.llama.RestoreSlot(ctx, tgt.BaseURL, tgt.Slot, fname) + restoreMs := float64(time.Since(tRestore).Microseconds()) / 1000.0 + if o.cfg.CleanupFiles { + os.Remove(abspath) + } + if err != nil { + return out, fmt.Errorf("llama restore: %w", err) + } + out.Restored = true + out.RestoredTokens = rr.NRestored + out.RestoreMs = restoreMs + // Feed the arbiter the FULL prepare wall (lookup + store fetch + GPU load + + // orchestration). For small blobs this is dominated by fixed overhead, so + // the learned "restore bandwidth" is low and the arbiter declines next time; + // for large blobs it is transfer-dominated and restore wins. This is what + // makes the policy adaptive and loss-free across model/hardware regimes. + fullMs := float64(time.Since(prepStart).Microseconds()) / 1000.0 + o.arb.ObserveGet(p.modelKey, gr.Bytes, fullMs) + o.mx.RestoreLatency.Observe(fullMs) + o.mx.RestoreCount.Inc() + o.mx.Hits.Inc() + o.mx.HitBlocks.Add(float64(m)) + o.mx.MissBlocks.Add(float64(p.chain.FullBlocks - m)) + o.mx.AddSavedTokens(out.MatchedTokens) + if im.Found { + // learned size correction + o.updateBytesPerTok(p.rootKey, out.MatchedTokens, gr.Bytes) + } + return out, nil +} + +// Commit saves the KV currently in target.Slot and stores it under the block +// key for the largest block-aligned prefix of `tokens`. skip-if-exists gives +// single-writer dedup across concurrent agents sharing a prefix. +func (o *Orchestrator) Commit(ctx context.Context, fp cachekey.ModelFingerprint, pol Policy, tokens []int32, modelPath string, tgt *Target, prefillN int, prefillMs float64) (CommitResult, error) { + o.mx.CommitTotal.Inc() + out := CommitResult{} + if !pol.Write { + out.Reason = "write disabled" + return out, nil + } + p := o.plan(fp, pol, tokens, modelPath) + // Learn the live prefill rate (per model) so the arbiter can compare. + o.arb.ObservePrefill(p.modelKey, prefillN, prefillMs) + minB := pol.MinPrefixBlocks + if minB <= 0 { + minB = o.cfg.MinPrefixBlocks + } + if p.chain.FullBlocks < minB { + out.Reason = fmt.Sprintf("prefix %d blocks < min %d; not cached", p.chain.FullBlocks, minB) + return out, nil + } + if tgt == nil || tgt.BaseURL == "" { + out.Reason = "no target to save from" + return out, nil + } + blocks := p.chain.FullBlocks + key := p.keys[blocks-1] + out.Key = key + out.StoredBlocks = blocks + out.StoredTokens = blocks * p.fp.BlockSize + + // Fast path: someone already stored this exact prefix. + if present, err := o.store.Exists(ctx, []string{key}); err == nil && len(present) == 1 && present[0] == 1 { + out.OK = true + out.Stored = false + out.Reason = "already present (dedup)" + return out, nil + } + + fname := fmt.Sprintf("omb-s-%x-%d.bin", hashKey(key), o.nonce.Add(1)) + abspath := filepath.Join(o.cfg.SlotSavePath, fname) + tSave := time.Now() + sr, err := o.llama.SaveSlot(ctx, tgt.BaseURL, tgt.Slot, fname) + if err != nil { + return out, fmt.Errorf("llama save: %w", err) + } + out.SaveMs = float64(time.Since(tSave).Microseconds()) / 1000.0 + o.mx.SaveLatency.Observe(out.SaveMs) + repl := pol.ReplicaNum + if repl == 0 { + repl = o.cfg.DefaultReplicaNum + } + pr, err := o.store.PutFile(ctx, key, abspath, repl, pol.SoftPin, true) + if o.cfg.CleanupFiles { + os.Remove(abspath) + } + if err != nil { + return out, fmt.Errorf("store put: %w", err) + } + out.StorePutMs = pr.ElapsedMs + out.Bytes = pr.Bytes + out.OK = true + out.Stored = !pr.Existed + o.mx.StorePutLatency.Observe(pr.ElapsedMs) + if out.Stored { + o.mx.AddBytesPut(pr.Bytes) + o.index.Insert(p.rootKey, p.chain.Hex, blocks, out.StoredTokens, key, pr.Bytes) + } + // learn bytes/token from the real save (n_written / n_saved) + if sr.NSaved > 0 { + o.updateBytesPerTok(p.rootKey, sr.NSaved, sr.NWritten) + } + if out.Stored { + out.Reason = "stored" + } else { + out.Reason = "already present (dedup)" + } + return out, nil +} + +// RefreshGauges pushes learned arbiter rates + index stats into Prometheus. +func (o *Orchestrator) RefreshGauges() { + s := o.arb.Snapshot() + o.mx.GetGBps.Set(s.GetGBps) + o.mx.PrefillToksS.Set(s.PrefillToksS) + is := o.index.Stats() + o.mx.IndexSnapshots.Set(float64(is.Snapshots)) + o.mx.IndexBytes.Set(float64(is.Bytes)) +} + +func hashKey(s string) uint64 { + // FNV-1a 64 + var h uint64 = 1469598103934665603 + for i := 0; i < len(s); i++ { + h ^= uint64(s[i]) + h *= 1099511628211 + } + return h +} diff --git a/mooncake-integration/ollama/bridge/internal/prefixindex/radix.go b/mooncake-integration/ollama/bridge/internal/prefixindex/radix.go new file mode 100644 index 00000000..f6af43e9 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/prefixindex/radix.go @@ -0,0 +1,221 @@ +// Package prefixindex maintains an in-memory radix (prefix) tree over the +// *chained block hashes* of cached KV prefixes. It is the cross-process / +// cross-node analogue of SGLang's RadixAttention tree, but the nodes reference +// KV snapshots living in the Mooncake Store rather than GPU memory. +// +// Because a chained block hash h_i already encodes blocks 0..i (see +// cachekey.ChainBlockHashes), two sequences that share the first j blocks share +// the identical path h_0..h_{j-1} in the tree and then diverge. Walking a +// request's chain down the tree therefore yields, in O(blocks), the longest +// stored prefix that is a true prefix of the request — and, as a free +// by-product, the *fan-out* at each node tells us how "hot" (widely shared) a +// prefix is, which drives replication / soft-pin decisions. +// +// The tree is a per-node hint: the authoritative existence check is the store. +// We use the tree for (a) a fast local hit path, (b) blob-size lookup so the +// cost arbiter can size a transfer without a round-trip, and (c) LRU eviction. +package prefixindex + +import ( + "sort" + "sync" + "time" +) + +type node struct { + children map[string]*node // next-block chain hash -> child + parent *node + edge string // the chain hash on the edge from parent to this node + + hasSnapshot bool // a KV snapshot exists for the prefix ending here + storeKey string // store object key for that snapshot + blocks int // depth in blocks (== prefix length / blockSize) + tokens int // exact token count of the snapshot + bytes uint64 // snapshot size + hits uint64 // times this snapshot was reused + lastUsed time.Time +} + +func newNode(parent *node, edge string, depth int) *node { + return &node{children: map[string]*node{}, parent: parent, edge: edge, blocks: depth} +} + +// Match is the result of a longest-prefix lookup. +type Match struct { + Found bool + Blocks int + Tokens int + Bytes uint64 + StoreKey string + Hits uint64 + // Fanout is how many distinct continuations diverge right after the matched + // prefix — a proxy for how widely shared (hot) this prefix is. + Fanout int +} + +// Index is a forest: one radix tree per root key (fingerprint+namespace). +type Index struct { + mu sync.RWMutex + roots map[string]*node + nSnap int + totalB uint64 +} + +func New() *Index { return &Index{roots: map[string]*node{}} } + +func (ix *Index) rootFor(rootKey string, create bool) *node { + r := ix.roots[rootKey] + if r == nil && create { + r = newNode(nil, "", 0) + ix.roots[rootKey] = r + } + return r +} + +// Insert records that a KV snapshot covering `blocks` blocks (chain[:blocks]) +// is stored under storeKey. chain must have at least `blocks` entries. +func (ix *Index) Insert(rootKey string, chain []string, blocks, tokens int, storeKey string, bytes uint64) { + if blocks <= 0 || blocks > len(chain) { + return + } + ix.mu.Lock() + defer ix.mu.Unlock() + cur := ix.rootFor(rootKey, true) + for i := 0; i < blocks; i++ { + h := chain[i] + ch := cur.children[h] + if ch == nil { + ch = newNode(cur, h, i+1) + cur.children[h] = ch + } + cur = ch + } + if !cur.hasSnapshot { + ix.nSnap++ + ix.totalB += bytes + } else { + ix.totalB += bytes - cur.bytes + } + cur.hasSnapshot = true + cur.storeKey = storeKey + cur.tokens = tokens + cur.bytes = bytes + cur.lastUsed = time.Now() +} + +// LongestMatch walks the request chain and returns the deepest node that holds +// a snapshot. touch=true bumps hit counters/recency for the matched node. +func (ix *Index) LongestMatch(rootKey string, chain []string, touch bool) Match { + ix.mu.Lock() + defer ix.mu.Unlock() + r := ix.rootFor(rootKey, false) + if r == nil { + return Match{} + } + cur := r + var best *node + for i := 0; i < len(chain); i++ { + ch := cur.children[chain[i]] + if ch == nil { + break + } + cur = ch + if cur.hasSnapshot { + best = cur + } + } + if best == nil { + return Match{} + } + if touch { + best.hits++ + best.lastUsed = time.Now() + } + return Match{ + Found: true, Blocks: best.blocks, Tokens: best.tokens, Bytes: best.bytes, + StoreKey: best.storeKey, Hits: best.hits, Fanout: len(best.children), + } +} + +// Forget drops the snapshot mark at the given depth (after a store eviction). +func (ix *Index) Forget(rootKey string, chain []string, blocks int) { + ix.mu.Lock() + defer ix.mu.Unlock() + r := ix.rootFor(rootKey, false) + if r == nil { + return + } + cur := r + for i := 0; i < blocks && cur != nil; i++ { + cur = cur.children[chain[i]] + } + if cur != nil && cur.hasSnapshot { + cur.hasSnapshot = false + ix.nSnap-- + ix.totalB -= cur.bytes + ix.prune(rootKey, cur) + } +} + +// prune removes now-empty leaf chains to bound memory. +func (ix *Index) prune(rootKey string, n *node) { + for n != nil && n.parent != nil && len(n.children) == 0 && !n.hasSnapshot { + p := n.parent + delete(p.children, n.edge) + n = p + } + if r := ix.roots[rootKey]; r != nil && len(r.children) == 0 { + delete(ix.roots, rootKey) + } +} + +// EvictionCandidate is a snapshot worth removing under memory pressure. +type EvictionCandidate struct { + RootKey string + StoreKey string + Bytes uint64 + Hits uint64 + LastUsed time.Time + Blocks int +} + +// ColdestSnapshots returns up to n least-recently-used snapshots (LRU), so the +// caller can evict them from the store and the tree. +func (ix *Index) ColdestSnapshots(n int) []EvictionCandidate { + ix.mu.RLock() + defer ix.mu.RUnlock() + var all []EvictionCandidate + for rk, root := range ix.roots { + var walk func(*node) + walk = func(nd *node) { + if nd.hasSnapshot { + all = append(all, EvictionCandidate{ + RootKey: rk, StoreKey: nd.storeKey, Bytes: nd.bytes, + Hits: nd.hits, LastUsed: nd.lastUsed, Blocks: nd.blocks, + }) + } + for _, c := range nd.children { + walk(c) + } + } + walk(root) + } + sort.Slice(all, func(i, j int) bool { return all[i].LastUsed.Before(all[j].LastUsed) }) + if len(all) > n { + all = all[:n] + } + return all +} + +// Stats snapshots index-wide counters. +type Stats struct { + Snapshots int + Bytes uint64 + Roots int +} + +func (ix *Index) Stats() Stats { + ix.mu.RLock() + defer ix.mu.RUnlock() + return Stats{Snapshots: ix.nSnap, Bytes: ix.totalB, Roots: len(ix.roots)} +} diff --git a/mooncake-integration/ollama/bridge/internal/prefixindex/radix_test.go b/mooncake-integration/ollama/bridge/internal/prefixindex/radix_test.go new file mode 100644 index 00000000..3520c6d1 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/prefixindex/radix_test.go @@ -0,0 +1,65 @@ +package prefixindex + +import "testing" + +func chain(n int) []string { + c := make([]string, n) + for i := range c { + c[i] = string(rune('a'+i%26)) + string(rune('0'+i/26)) + } + return c +} + +func TestLongestMatch(t *testing.T) { + ix := New() + c := chain(10) + // store a 4-block and a 7-block snapshot on the same path + ix.Insert("m", c, 4, 4*256, "key4", 1000) + ix.Insert("m", c, 7, 7*256, "key7", 2000) + + // a request sharing all 10 blocks should match the deepest stored (7) + m := ix.LongestMatch("m", c, false) + if !m.Found || m.Blocks != 7 || m.StoreKey != "key7" || m.Bytes != 2000 { + t.Fatalf("want 7-block key7, got %+v", m) + } + + // a request that diverges at block 5 should match the 4-block snapshot + d := append([]string{}, c...) + d[5] = "ZZ" + m = ix.LongestMatch("m", d, false) + if !m.Found || m.Blocks != 4 || m.StoreKey != "key4" { + t.Fatalf("want 4-block key4 after divergence, got %+v", m) + } + + // unknown root => no match + if ix.LongestMatch("other", c, false).Found { + t.Fatal("unexpected match in empty root") + } +} + +func TestForgetAndStats(t *testing.T) { + ix := New() + c := chain(6) + ix.Insert("m", c, 3, 3*256, "k", 500) + if s := ix.Stats(); s.Snapshots != 1 || s.Bytes != 500 { + t.Fatalf("stats after insert: %+v", s) + } + ix.Forget("m", c, 3) + if ix.LongestMatch("m", c, false).Found { + t.Fatal("match after forget") + } + if s := ix.Stats(); s.Snapshots != 0 { + t.Fatalf("snapshots after forget: %d", s.Snapshots) + } +} + +func TestColdestEviction(t *testing.T) { + ix := New() + c := chain(8) + ix.Insert("m", c, 2, 512, "old", 100) + ix.Insert("m", c, 5, 1280, "new", 200) + got := ix.ColdestSnapshots(1) + if len(got) != 1 || got[0].StoreKey != "old" { + t.Fatalf("want oldest 'old', got %+v", got) + } +} diff --git a/mooncake-integration/ollama/bridge/internal/seqstate/seqstate_cgo.go b/mooncake-integration/ollama/bridge/internal/seqstate/seqstate_cgo.go new file mode 100644 index 00000000..7f5416d1 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/seqstate/seqstate_cgo.go @@ -0,0 +1,95 @@ +//go:build cgo_stage2 + +// Package seqstate is the Stage-2 (high-performance) Go binding for moving a +// llama.cpp sequence's KV state with minimal copies, used when the sidecar is +// embedded in-process with libllama (rather than driving a llama-server over +// HTTP as in Stage-1). +// +// It is gated behind the `cgo_stage2` build tag because it links libllama and +// is only relevant for the embedded/high-performance deployment. The Stage-1 +// HTTP path (internal/llamabridge) needs none of this and is the default. +// +// Why this exists (measured by cbridge/omb_kvbench): the Stage-1 file path +// (llama_state_seq_save_file, i.e. /slots save) is ~5x slower than the raw +// in-process state copy because it serializes through host memory and a file +// (llama.cpp issue #8915's double copy). Going further, the _ext API with +// LLAMA_STATE_SEQ_FLAGS_ON_DEVICE keeps the KV tensors in device buffers so the +// Mooncake Transfer Engine can register them (registerLocalMemory) and move +// them by GPUDirect RDMA with NO host copy at all. +// +// Build (example): +// +// CGO_CFLAGS="-I$LLAMA_DIR/include -I$LLAMA_DIR/ggml/include" \ +// CGO_LDFLAGS="-L$LLAMA_BUILD/bin -lllama -lggml -lggml-base -Wl,-rpath,$LLAMA_BUILD/bin" \ +// go build -tags cgo_stage2 ./... +package seqstate + +/* +#include +#include "llama.h" +*/ +import "C" +import ( + "fmt" + "unsafe" +) + +// Flags mirror llama_state_seq_flags. +type Flags uint32 + +const ( + FlagsNone Flags = 0 + FlagsOnDevice Flags = 2 // keep KV in device buffers; pair with Mooncake TE GDR +) + +// Ctx wraps a *llama_context owned elsewhere in the embedded process. +type Ctx struct{ ptr *C.struct_llama_context } + +// Wrap adapts a raw context pointer (e.g. obtained from the embedded runner). +func Wrap(p unsafe.Pointer) *Ctx { return &Ctx{ptr: (*C.struct_llama_context)(p)} } + +// SeqSize returns the bytes needed to export sequence seq with the given flags. +// With FlagsOnDevice this is just the device-handle/metadata size; the bulk KV +// stays resident on the GPU. +func (c *Ctx) SeqSize(seq int, flags Flags) uint64 { + return uint64(C.llama_state_seq_get_size_ext(c.ptr, C.llama_seq_id(seq), C.llama_state_seq_flags(flags))) +} + +// SeqGet copies sequence seq's state into dst (len(dst) >= SeqSize). Returns the +// number of bytes written. For zero-copy GPUDirect, allocate dst from a buffer +// already registered with the Mooncake Transfer Engine. +func (c *Ctx) SeqGet(dst []byte, seq int, flags Flags) (uint64, error) { + if len(dst) == 0 { + return 0, fmt.Errorf("empty dst") + } + n := C.llama_state_seq_get_data_ext(c.ptr, + (*C.uint8_t)(unsafe.Pointer(&dst[0])), C.size_t(len(dst)), + C.llama_seq_id(seq), C.llama_state_seq_flags(flags)) + if n == 0 { + return 0, fmt.Errorf("llama_state_seq_get_data_ext failed") + } + return uint64(n), nil +} + +// SeqSet imports src into sequence dstSeq. Returns >0 on success. +func (c *Ctx) SeqSet(src []byte, dstSeq int, flags Flags) (uint64, error) { + if len(src) == 0 { + return 0, fmt.Errorf("empty src") + } + n := C.llama_state_seq_set_data_ext(c.ptr, + (*C.uint8_t)(unsafe.Pointer(&src[0])), C.size_t(len(src)), + C.llama_seq_id(dstSeq), C.llama_state_seq_flags(flags)) + if n == 0 { + return 0, fmt.Errorf("llama_state_seq_set_data_ext failed") + } + return uint64(n), nil +} + +// PtrAndLen exposes a buffer's address+length so the engine binding +// (engine_cgo.go) can registerLocalMemory it for RDMA transfer. +func PtrAndLen(b []byte) (uintptr, int) { + if len(b) == 0 { + return 0, 0 + } + return uintptr(unsafe.Pointer(&b[0])), len(b) +} diff --git a/mooncake-integration/ollama/bridge/internal/server/grpc.go b/mooncake-integration/ollama/bridge/internal/server/grpc.go new file mode 100644 index 00000000..04f8c1e4 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/server/grpc.go @@ -0,0 +1,151 @@ +// Package server adapts the gRPC KVCacheBus contract and an HTTP/JSON gateway +// onto the orchestrator. gRPC is the primary transport; the HTTP/JSON gateway +// exists so the (Go) Ollama patch and quick curl tests can drive the bus with a +// single net/http call and no protobuf dependency. +package server + +import ( + "context" + + bridgepb "github.com/mooncake-ai/ollama-mooncake-bridge/internal/bridge/pb" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/cachekey" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/metrics" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/orchestrator" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/store" + dto "github.com/prometheus/client_model/go" + "github.com/prometheus/client_golang/prometheus" +) + +const Version = "0.1.0" + +// GRPCServer implements bridgepb.KVCacheBusServer. +type GRPCServer struct { + bridgepb.UnimplementedKVCacheBusServer + orch *orchestrator.Orchestrator + store store.Backend + mx *metrics.Metrics +} + +func NewGRPC(o *orchestrator.Orchestrator, b store.Backend, mx *metrics.Metrics) *GRPCServer { + return &GRPCServer{orch: o, store: b, mx: mx} +} + +func fpFromProto(p *bridgepb.ModelFingerprint) (cachekey.ModelFingerprint, string) { + if p == nil { + return cachekey.ModelFingerprint{}, "" + } + return cachekey.ModelFingerprint{ + ModelDigest: p.ModelDigest, + Arch: p.Arch, + TokenizerHash: p.TokenizerHash, + RopeHash: p.RopeHash, + KVType: p.KvType, + KVLayout: p.KvLayout, + NCtxTrain: int(p.NCtxTrain), + BlockSize: int(p.BlockSize), + }, p.ModelPath +} + +func polFromProto(p *bridgepb.CachePolicy) orchestrator.Policy { + if p == nil { + return orchestrator.Policy{Enable: true, Read: true, Write: true} + } + return orchestrator.Policy{ + Enable: p.Enable, + Namespace: p.Namespace, + Read: p.Read, + Write: p.Write, + BlockSize: int(p.BlockSize), + ReplicaNum: p.ReplicaNum, + SoftPin: p.SoftPin, + MinPrefixBlocks: int(p.MinPrefixBlocks), + } +} + +func tgtFromProto(p *bridgepb.LlamaTarget) *orchestrator.Target { + if p == nil || p.BaseUrl == "" { + return nil + } + return &orchestrator.Target{BaseURL: p.BaseUrl, Slot: int(p.SlotId)} +} + +func (s *GRPCServer) Health(ctx context.Context, _ *bridgepb.HealthRequest) (*bridgepb.HealthReply, error) { + hi, err := s.store.Health(ctx) + rep := &bridgepb.HealthReply{Version: Version, StoreBackend: s.store.Name()} + if err != nil { + rep.Ok = false + rep.Detail = err.Error() + return rep, nil + } + rep.Ok = true + rep.StoreOk = hi.OK + rep.Protocol = hi.Protocol + rep.Detail = hi.Detail + return rep, nil +} + +func (s *GRPCServer) Lookup(ctx context.Context, req *bridgepb.LookupRequest) (*bridgepb.LookupReply, error) { + fp, mp := fpFromProto(req.Fp) + r, err := s.orch.Lookup(ctx, fp, polFromProto(req.Policy), req.Tokens, mp) + if err != nil { + return &bridgepb.LookupReply{Error: err.Error()}, nil + } + return &bridgepb.LookupReply{ + Hit: r.Hit, MatchedBlocks: int32(r.MatchedBlocks), MatchedTokens: int32(r.MatchedTokens), + TotalBlocks: int32(r.TotalBlocks), TotalTokens: int32(r.TotalTokens), + Decision: r.Decision, MatchedKey: r.MatchedKey, Reason: r.Reason, + }, nil +} + +func (s *GRPCServer) Prepare(ctx context.Context, req *bridgepb.PrepareRequest) (*bridgepb.PrepareReply, error) { + fp, mp := fpFromProto(req.Fp) + r, err := s.orch.Prepare(ctx, fp, polFromProto(req.Policy), req.Tokens, mp, tgtFromProto(req.Target)) + if err != nil { + return &bridgepb.PrepareReply{Error: err.Error()}, nil + } + return &bridgepb.PrepareReply{ + Hit: r.Hit, MatchedBlocks: int32(r.MatchedBlocks), MatchedTokens: int32(r.MatchedTokens), + TotalBlocks: int32(r.TotalBlocks), TotalTokens: int32(r.TotalTokens), + Decision: r.Decision, MatchedKey: r.MatchedKey, + Restored: r.Restored, RestoredTokens: int32(r.RestoredTokens), + RestoreMs: r.RestoreMs, StoreGetMs: r.StoreGetMs, Bytes: r.Bytes, Reason: r.Reason, + }, nil +} + +func (s *GRPCServer) Commit(ctx context.Context, req *bridgepb.CommitRequest) (*bridgepb.CommitReply, error) { + fp, mp := fpFromProto(req.Fp) + r, err := s.orch.Commit(ctx, fp, polFromProto(req.Policy), req.Tokens, mp, tgtFromProto(req.Target), int(req.PrefillN), req.PrefillMs) + if err != nil { + return &bridgepb.CommitReply{Error: err.Error()}, nil + } + return &bridgepb.CommitReply{ + Ok: r.OK, Stored: r.Stored, StoredBlocks: int32(r.StoredBlocks), StoredTokens: int32(r.StoredTokens), + Bytes: r.Bytes, SaveMs: r.SaveMs, StorePutMs: r.StorePutMs, Key: r.Key, Reason: r.Reason, + }, nil +} + +func (s *GRPCServer) Stats(_ context.Context, _ *bridgepb.StatsRequest) (*bridgepb.StatsReply, error) { + s.orch.RefreshGauges() + m := s.mx + return &bridgepb.StatsReply{ + PrepareTotal: u(m.PrepareTotal), + Hits: u(m.Hits), + Misses: u(m.Misses), + HitBlocks: u(m.HitBlocks), + MissBlocks: u(m.MissBlocks), + SavedPrefillTokens: m.SavedTokens(), + RestoreCount: u(m.RestoreCount), + CommitCount: u(m.CommitTotal), + BytesGet: m.TotalBytesGet(), + BytesPut: m.TotalBytesPut(), + }, nil +} + +// u reads a prometheus counter's current value as uint64. +func u(c prometheus.Counter) uint64 { + var m dto.Metric + if err := c.Write(&m); err != nil { + return 0 + } + return uint64(m.GetCounter().GetValue()) +} diff --git a/mooncake-integration/ollama/bridge/internal/server/http.go b/mooncake-integration/ollama/bridge/internal/server/http.go new file mode 100644 index 00000000..4ba09f33 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/server/http.go @@ -0,0 +1,211 @@ +package server + +import ( + "encoding/json" + "net/http" + + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/cachekey" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/metrics" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/orchestrator" + "github.com/mooncake-ai/ollama-mooncake-bridge/internal/store" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promhttp" +) + +// --- JSON DTOs (snake_case mirrors the proto for a uniform API) --- + +type jsonFP struct { + ModelDigest string `json:"model_digest"` + Arch string `json:"arch"` + TokenizerHash string `json:"tokenizer_hash"` + RopeHash string `json:"rope_hash"` + KVType string `json:"kv_type"` + KVLayout string `json:"kv_layout"` + NCtxTrain int `json:"n_ctx_train"` + BlockSize int `json:"block_size"` + ModelPath string `json:"model_path"` +} + +func (f jsonFP) to() (cachekey.ModelFingerprint, string) { + return cachekey.ModelFingerprint{ + ModelDigest: f.ModelDigest, Arch: f.Arch, TokenizerHash: f.TokenizerHash, + RopeHash: f.RopeHash, KVType: f.KVType, KVLayout: f.KVLayout, + NCtxTrain: f.NCtxTrain, BlockSize: f.BlockSize, + }, f.ModelPath +} + +type jsonPolicy struct { + Enable bool `json:"enable"` + Namespace string `json:"namespace"` + Read bool `json:"read"` + Write bool `json:"write"` + BlockSize int `json:"block_size"` + ReplicaNum uint32 `json:"replica_num"` + SoftPin bool `json:"soft_pin"` + MinPrefixBlocks int `json:"min_prefix_blocks"` +} + +func (p jsonPolicy) to() orchestrator.Policy { + return orchestrator.Policy{ + Enable: p.Enable, Namespace: p.Namespace, Read: p.Read, Write: p.Write, + BlockSize: p.BlockSize, ReplicaNum: p.ReplicaNum, SoftPin: p.SoftPin, + MinPrefixBlocks: p.MinPrefixBlocks, + } +} + +type jsonTarget struct { + BaseURL string `json:"base_url"` + Slot int `json:"slot_id"` +} + +type jsonReq struct { + FP jsonFP `json:"fp"` + Policy jsonPolicy `json:"policy"` + Tokens []int32 `json:"tokens"` + Target *jsonTarget `json:"target"` + PrefillN int `json:"prefill_n"` + PrefillMs float64 `json:"prefill_ms"` +} + +func (r jsonReq) target() *orchestrator.Target { + if r.Target == nil || r.Target.BaseURL == "" { + return nil + } + return &orchestrator.Target{BaseURL: r.Target.BaseURL, Slot: r.Target.Slot} +} + +// HTTPServer wires the JSON gateway + Prometheus endpoint. +type HTTPServer struct { + orch *orchestrator.Orchestrator + store store.Backend + mx *metrics.Metrics + reg *prometheus.Registry +} + +func NewHTTP(o *orchestrator.Orchestrator, b store.Backend, mx *metrics.Metrics, reg *prometheus.Registry) *HTTPServer { + return &HTTPServer{orch: o, store: b, mx: mx, reg: reg} +} + +func (h *HTTPServer) Mux() *http.ServeMux { + mux := http.NewServeMux() + mux.Handle("/metrics", promhttp.HandlerFor(h.reg, promhttp.HandlerOpts{})) + mux.HandleFunc("/healthz", h.healthz) + mux.HandleFunc("/stats", h.stats) + mux.HandleFunc("/v1/lookup", h.lookup) + mux.HandleFunc("/v1/prepare", h.prepare) + mux.HandleFunc("/v1/commit", h.commit) + return mux +} + +func writeJSON(w http.ResponseWriter, code int, v any) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(code) + _ = json.NewEncoder(w).Encode(v) +} + +func (h *HTTPServer) decode(w http.ResponseWriter, r *http.Request) (jsonReq, bool) { + var req jsonReq + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()}) + return req, false + } + return req, true +} + +func (h *HTTPServer) healthz(w http.ResponseWriter, r *http.Request) { + hi, err := h.store.Health(r.Context()) + resp := map[string]any{"version": Version, "store_backend": h.store.Name()} + if err != nil { + resp["ok"] = false + resp["detail"] = err.Error() + writeJSON(w, http.StatusOK, resp) + return + } + resp["ok"] = true + resp["store_ok"] = hi.OK + resp["protocol"] = hi.Protocol + resp["master"] = hi.Master + resp["device"] = hi.Device + resp["detail"] = hi.Detail + writeJSON(w, http.StatusOK, resp) +} + +func (h *HTTPServer) stats(w http.ResponseWriter, _ *http.Request) { + h.orch.RefreshGauges() + snap := h.orch.Arbiter().Snapshot() + is := h.orch.Index().Stats() + writeJSON(w, http.StatusOK, map[string]any{ + "prepare_total": u(h.mx.PrepareTotal), + "hits": u(h.mx.Hits), + "misses": u(h.mx.Misses), + "recompute_chosen": u(h.mx.RecomputeChosen), + "hit_blocks": u(h.mx.HitBlocks), + "miss_blocks": u(h.mx.MissBlocks), + "restore_count": u(h.mx.RestoreCount), + "commit_count": u(h.mx.CommitTotal), + "saved_prefill_tokens": h.mx.SavedTokens(), + "bytes_get": h.mx.TotalBytesGet(), + "bytes_put": h.mx.TotalBytesPut(), + "learned_get_gbps": snap.GetGBps, + "learned_prefill_tps": snap.PrefillToksS, + "index_snapshots": is.Snapshots, + "index_bytes": is.Bytes, + }) +} + +func (h *HTTPServer) lookup(w http.ResponseWriter, r *http.Request) { + req, ok := h.decode(w, r) + if !ok { + return + } + fp, mp := req.FP.to() + res, err := h.orch.Lookup(r.Context(), fp, req.Policy.to(), req.Tokens, mp) + if err != nil { + writeJSON(w, http.StatusOK, map[string]any{"error": err.Error()}) + return + } + writeJSON(w, http.StatusOK, map[string]any{ + "hit": res.Hit, "matched_blocks": res.MatchedBlocks, "matched_tokens": res.MatchedTokens, + "total_blocks": res.TotalBlocks, "total_tokens": res.TotalTokens, + "decision": res.Decision, "matched_key": res.MatchedKey, "reason": res.Reason, + }) +} + +func (h *HTTPServer) prepare(w http.ResponseWriter, r *http.Request) { + req, ok := h.decode(w, r) + if !ok { + return + } + fp, mp := req.FP.to() + res, err := h.orch.Prepare(r.Context(), fp, req.Policy.to(), req.Tokens, mp, req.target()) + if err != nil { + writeJSON(w, http.StatusOK, map[string]any{"error": err.Error()}) + return + } + writeJSON(w, http.StatusOK, map[string]any{ + "hit": res.Hit, "matched_blocks": res.MatchedBlocks, "matched_tokens": res.MatchedTokens, + "total_blocks": res.TotalBlocks, "total_tokens": res.TotalTokens, + "decision": res.Decision, "matched_key": res.MatchedKey, + "restored": res.Restored, "restored_tokens": res.RestoredTokens, + "restore_ms": res.RestoreMs, "store_get_ms": res.StoreGetMs, "bytes": res.Bytes, + "reason": res.Reason, + }) +} + +func (h *HTTPServer) commit(w http.ResponseWriter, r *http.Request) { + req, ok := h.decode(w, r) + if !ok { + return + } + fp, mp := req.FP.to() + res, err := h.orch.Commit(r.Context(), fp, req.Policy.to(), req.Tokens, mp, req.target(), req.PrefillN, req.PrefillMs) + if err != nil { + writeJSON(w, http.StatusOK, map[string]any{"error": err.Error()}) + return + } + writeJSON(w, http.StatusOK, map[string]any{ + "ok": res.OK, "stored": res.Stored, "stored_blocks": res.StoredBlocks, + "stored_tokens": res.StoredTokens, "bytes": res.Bytes, + "save_ms": res.SaveMs, "store_put_ms": res.StorePutMs, "key": res.Key, "reason": res.Reason, + }) +} diff --git a/mooncake-integration/ollama/bridge/internal/store/local.go b/mooncake-integration/ollama/bridge/internal/store/local.go new file mode 100644 index 00000000..e73cc49b --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/store/local.go @@ -0,0 +1,118 @@ +package store + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "io" + "os" + "path/filepath" + "time" +) + +// LocalBackend is a pure-Go filesystem object store. It backs the "local file" +// baseline (cross-process, single-node, no Mooncake) and serves as a +// zero-dependency fallback. Keys are hashed to sharded file names. +type LocalBackend struct { + root string +} + +func NewLocalBackend(root string) (*LocalBackend, error) { + if err := os.MkdirAll(root, 0o755); err != nil { + return nil, err + } + return &LocalBackend{root: root}, nil +} + +func (l *LocalBackend) Name() string { return "local" } + +func (l *LocalBackend) path(key string) string { + h := sha256.Sum256([]byte(key)) + hs := hex.EncodeToString(h[:]) + return filepath.Join(l.root, hs[:2], hs) +} + +func (l *LocalBackend) Health(_ context.Context) (HealthInfo, error) { + return HealthInfo{OK: true, Backend: "local", Protocol: "file", Master: l.root, Detail: "filesystem store"}, nil +} + +func (l *LocalBackend) Exists(_ context.Context, keys []string) ([]int32, error) { + out := make([]int32, len(keys)) + for i, k := range keys { + if _, err := os.Stat(l.path(k)); err == nil { + out[i] = 1 + } + } + return out, nil +} + +func (l *LocalBackend) PutFile(_ context.Context, key, path string, _ uint32, _ bool, skipIfExists bool) (PutResult, error) { + dst := l.path(key) + if skipIfExists { + if fi, err := os.Stat(dst); err == nil { + return PutResult{Bytes: uint64(fi.Size()), Existed: true}, nil + } + } + t0 := time.Now() + if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil { + return PutResult{}, err + } + n, err := copyFile(path, dst) + if err != nil { + return PutResult{}, err + } + return PutResult{Bytes: uint64(n), ElapsedMs: float64(time.Since(t0).Microseconds()) / 1000.0}, nil +} + +func (l *LocalBackend) GetFile(_ context.Context, key, path string) (GetResult, error) { + src := l.path(key) + if _, err := os.Stat(src); err != nil { + return GetResult{Found: false}, nil + } + t0 := time.Now() + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return GetResult{}, err + } + n, err := copyFile(src, path) + if err != nil { + return GetResult{}, err + } + return GetResult{Bytes: uint64(n), ElapsedMs: float64(time.Since(t0).Microseconds()) / 1000.0, Found: true}, nil +} + +func (l *LocalBackend) Remove(_ context.Context, key string, _ bool) error { + err := os.Remove(l.path(key)) + if os.IsNotExist(err) { + return nil + } + return err +} + +func (l *LocalBackend) Close() error { return nil } + +// copyFile writes src->dst atomically (via a .tmp + rename) and returns bytes. +func copyFile(src, dst string) (int64, error) { + in, err := os.Open(src) + if err != nil { + return 0, err + } + defer in.Close() + tmp := dst + ".tmp" + out, err := os.Create(tmp) + if err != nil { + return 0, err + } + n, err := io.Copy(out, in) + if cerr := out.Close(); err == nil { + err = cerr + } + if err != nil { + os.Remove(tmp) + return 0, err + } + if err := os.Rename(tmp, dst); err != nil { + os.Remove(tmp) + return 0, err + } + return n, nil +} diff --git a/mooncake-integration/ollama/bridge/internal/store/mooncake.go b/mooncake-integration/ollama/bridge/internal/store/mooncake.go new file mode 100644 index 00000000..5fa452f6 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/store/mooncake.go @@ -0,0 +1,110 @@ +package store + +import ( + "context" + "time" + + storeproxypb "github.com/mooncake-ai/ollama-mooncake-bridge/internal/storeproxy/pb" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" +) + +// MooncakeBackend talks to the Python store proxy (which owns the real +// MooncakeDistributedStore handle) over gRPC. We pass file paths so KV blobs +// move file<->store inside the proxy, never across this gRPC link. +type MooncakeBackend struct { + conn *grpc.ClientConn + cli storeproxypb.StoreProxyClient + target string +} + +// DialMooncake connects to the store proxy at addr (host:port). +func DialMooncake(addr string) (*MooncakeBackend, error) { + conn, err := grpc.NewClient(addr, + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithDefaultCallOptions( + grpc.MaxCallRecvMsgSize(512<<20), + grpc.MaxCallSendMsgSize(512<<20), + ), + ) + if err != nil { + return nil, err + } + return &MooncakeBackend{conn: conn, cli: storeproxypb.NewStoreProxyClient(conn), target: addr}, nil +} + +func (m *MooncakeBackend) Name() string { return "mooncake" } + +func (m *MooncakeBackend) Health(ctx context.Context) (HealthInfo, error) { + ctx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + r, err := m.cli.Health(ctx, &storeproxypb.HealthRequest{}) + if err != nil { + return HealthInfo{}, err + } + return HealthInfo{OK: r.Ok, Backend: r.Backend, Protocol: r.Protocol, + Master: r.Master, Device: r.Device, Detail: r.Detail}, nil +} + +func (m *MooncakeBackend) Exists(ctx context.Context, keys []string) ([]int32, error) { + if len(keys) == 0 { + return nil, nil + } + r, err := m.cli.Exists(ctx, &storeproxypb.ExistsRequest{Keys: keys}) + if err != nil { + return nil, err + } + return r.Present, nil +} + +func (m *MooncakeBackend) PutFile(ctx context.Context, key, path string, replicaNum uint32, softPin, skipIfExists bool) (PutResult, error) { + r, err := m.cli.PutFile(ctx, &storeproxypb.PutFileRequest{ + Key: key, Path: path, ReplicaNum: replicaNum, SoftPin: softPin, SkipIfExists: skipIfExists, + }) + if err != nil { + return PutResult{}, err + } + if !r.Ok { + return PutResult{}, &Error{Op: "PutFile", Key: key, Msg: r.Error} + } + return PutResult{Bytes: r.Bytes, ElapsedMs: r.ElapsedMs, Existed: r.Existed}, nil +} + +func (m *MooncakeBackend) GetFile(ctx context.Context, key, path string) (GetResult, error) { + r, err := m.cli.GetFile(ctx, &storeproxypb.GetFileRequest{Key: key, Path: path}) + if err != nil { + return GetResult{}, err + } + if !r.Ok { + // A clean miss has found=false and no error message; anything with an + // error message is a real failure and must be surfaced, not hidden as a + // miss (which would degrade silently to a recompute). + if !r.Found && r.Error == "" { + return GetResult{Found: false}, nil + } + return GetResult{}, &Error{Op: "GetFile", Key: key, Msg: r.Error} + } + return GetResult{Bytes: r.Bytes, ElapsedMs: r.ElapsedMs, Found: r.Found}, nil +} + +func (m *MooncakeBackend) Remove(ctx context.Context, key string, force bool) error { + r, err := m.cli.Remove(ctx, &storeproxypb.RemoveRequest{Key: key, Force: force}) + if err != nil { + return err + } + if !r.Ok { + return &Error{Op: "Remove", Key: key, Msg: r.Error} + } + return nil +} + +func (m *MooncakeBackend) Close() error { return m.conn.Close() } + +// Error is a typed store error carrying op + key context. +type Error struct { + Op string + Key string + Msg string +} + +func (e *Error) Error() string { return e.Op + "(" + e.Key + "): " + e.Msg } diff --git a/mooncake-integration/ollama/bridge/internal/store/store.go b/mooncake-integration/ollama/bridge/internal/store/store.go new file mode 100644 index 00000000..1e22be8b --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/store/store.go @@ -0,0 +1,49 @@ +// Package store abstracts the KV-snapshot object store behind a small Backend +// interface so the orchestrator is agnostic to *where* bytes live. Two +// implementations ship: +// +// - mooncake: the real distributed store, reached via the Python store proxy +// over gRPC (TCP or RDMA/GPUDirect). This is the primary path. +// - local: a pure-Go filesystem object store (the "local file" baseline and a +// zero-dependency fallback when no proxy/master is available). +// +// Blobs are addressed by file path, never streamed through this process: a +// multi-GiB KV snapshot is copied at most once (file <-> store). +package store + +import "context" + +// PutResult / GetResult report what a transfer moved and how long it took, so +// the arbiter can learn live store bandwidth. +type PutResult struct { + Bytes uint64 + ElapsedMs float64 + Existed bool // skipped because the key already existed +} + +type GetResult struct { + Bytes uint64 + ElapsedMs float64 + Found bool +} + +type HealthInfo struct { + OK bool + Backend string + Protocol string + Master string + Device string + Detail string +} + +// Backend is the minimal object-store contract the orchestrator needs. +type Backend interface { + Name() string + Health(ctx context.Context) (HealthInfo, error) + // Exists returns one value per key: 1 present, 0 absent, -1 error. + Exists(ctx context.Context, keys []string) ([]int32, error) + PutFile(ctx context.Context, key, path string, replicaNum uint32, softPin, skipIfExists bool) (PutResult, error) + GetFile(ctx context.Context, key, path string) (GetResult, error) + Remove(ctx context.Context, key string, force bool) error + Close() error +} diff --git a/mooncake-integration/ollama/bridge/internal/storeproxy/pb/storeproxy.pb.go b/mooncake-integration/ollama/bridge/internal/storeproxy/pb/storeproxy.pb.go new file mode 100644 index 00000000..ef8fb50e --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/storeproxy/pb/storeproxy.pb.go @@ -0,0 +1,1479 @@ +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.34.2 +// protoc v3.12.4 +// source: storeproxy.proto + +// StoreProxy is the gRPC contract between the Go sidecar (client) and the +// Python Mooncake store proxy (server). The Python side owns the single, +// long-lived, warm mooncake.store.MooncakeDistributedStore handle. +// +// Design note: KV snapshots can be multiple GiB. We therefore pass *file paths* +// (in the shared --slot-save-path directory) instead of streaming bytes through +// Go. The proxy reads/writes those files directly to/from the store, so a +// multi-GiB KV blob is copied at most once (file <-> store), never through the +// Go process. Small values may still be sent inline via PutBytes/GetBytes. + +package storeproxypb + +import ( + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + reflect "reflect" + sync "sync" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +type HealthRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields +} + +func (x *HealthRequest) Reset() { + *x = HealthRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *HealthRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HealthRequest) ProtoMessage() {} + +func (x *HealthRequest) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[0] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use HealthRequest.ProtoReflect.Descriptor instead. +func (*HealthRequest) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{0} +} + +type HealthReply struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Ok bool `protobuf:"varint,1,opt,name=ok,proto3" json:"ok,omitempty"` + Backend string `protobuf:"bytes,2,opt,name=backend,proto3" json:"backend,omitempty"` // "mooncake" | "local" + Protocol string `protobuf:"bytes,3,opt,name=protocol,proto3" json:"protocol,omitempty"` // "tcp" | "rdma" + Master string `protobuf:"bytes,4,opt,name=master,proto3" json:"master,omitempty"` // master_server_addr + Device string `protobuf:"bytes,5,opt,name=device,proto3" json:"device,omitempty"` // rdma device(s) + Detail string `protobuf:"bytes,6,opt,name=detail,proto3" json:"detail,omitempty"` +} + +func (x *HealthReply) Reset() { + *x = HealthReply{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *HealthReply) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HealthReply) ProtoMessage() {} + +func (x *HealthReply) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[1] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use HealthReply.ProtoReflect.Descriptor instead. +func (*HealthReply) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{1} +} + +func (x *HealthReply) GetOk() bool { + if x != nil { + return x.Ok + } + return false +} + +func (x *HealthReply) GetBackend() string { + if x != nil { + return x.Backend + } + return "" +} + +func (x *HealthReply) GetProtocol() string { + if x != nil { + return x.Protocol + } + return "" +} + +func (x *HealthReply) GetMaster() string { + if x != nil { + return x.Master + } + return "" +} + +func (x *HealthReply) GetDevice() string { + if x != nil { + return x.Device + } + return "" +} + +func (x *HealthReply) GetDetail() string { + if x != nil { + return x.Detail + } + return "" +} + +type ExistsRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Keys []string `protobuf:"bytes,1,rep,name=keys,proto3" json:"keys,omitempty"` +} + +func (x *ExistsRequest) Reset() { + *x = ExistsRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *ExistsRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ExistsRequest) ProtoMessage() {} + +func (x *ExistsRequest) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[2] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ExistsRequest.ProtoReflect.Descriptor instead. +func (*ExistsRequest) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{2} +} + +func (x *ExistsRequest) GetKeys() []string { + if x != nil { + return x.Keys + } + return nil +} + +// present[i]: 1 = exists, 0 = absent, -1 = error (mirrors batch_is_exist). +type ExistsReply struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Present []int32 `protobuf:"varint,1,rep,packed,name=present,proto3" json:"present,omitempty"` +} + +func (x *ExistsReply) Reset() { + *x = ExistsReply{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *ExistsReply) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ExistsReply) ProtoMessage() {} + +func (x *ExistsReply) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[3] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ExistsReply.ProtoReflect.Descriptor instead. +func (*ExistsReply) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{3} +} + +func (x *ExistsReply) GetPresent() []int32 { + if x != nil { + return x.Present + } + return nil +} + +type PutFileRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Key string `protobuf:"bytes,1,opt,name=key,proto3" json:"key,omitempty"` + Path string `protobuf:"bytes,2,opt,name=path,proto3" json:"path,omitempty"` // local file to read + ReplicaNum uint32 `protobuf:"varint,3,opt,name=replica_num,json=replicaNum,proto3" json:"replica_num,omitempty"` // ReplicateConfig.replica_num (>=1) + SoftPin bool `protobuf:"varint,4,opt,name=soft_pin,json=softPin,proto3" json:"soft_pin,omitempty"` // ReplicateConfig.with_soft_pin (hot prefixes) + SkipIfExists bool `protobuf:"varint,5,opt,name=skip_if_exists,json=skipIfExists,proto3" json:"skip_if_exists,omitempty"` +} + +func (x *PutFileRequest) Reset() { + *x = PutFileRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *PutFileRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PutFileRequest) ProtoMessage() {} + +func (x *PutFileRequest) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[4] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PutFileRequest.ProtoReflect.Descriptor instead. +func (*PutFileRequest) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{4} +} + +func (x *PutFileRequest) GetKey() string { + if x != nil { + return x.Key + } + return "" +} + +func (x *PutFileRequest) GetPath() string { + if x != nil { + return x.Path + } + return "" +} + +func (x *PutFileRequest) GetReplicaNum() uint32 { + if x != nil { + return x.ReplicaNum + } + return 0 +} + +func (x *PutFileRequest) GetSoftPin() bool { + if x != nil { + return x.SoftPin + } + return false +} + +func (x *PutFileRequest) GetSkipIfExists() bool { + if x != nil { + return x.SkipIfExists + } + return false +} + +type PutFileReply struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Ok bool `protobuf:"varint,1,opt,name=ok,proto3" json:"ok,omitempty"` + Bytes uint64 `protobuf:"varint,2,opt,name=bytes,proto3" json:"bytes,omitempty"` + ElapsedMs float64 `protobuf:"fixed64,3,opt,name=elapsed_ms,json=elapsedMs,proto3" json:"elapsed_ms,omitempty"` + Existed bool `protobuf:"varint,4,opt,name=existed,proto3" json:"existed,omitempty"` // true if skipped because key already present + Error string `protobuf:"bytes,5,opt,name=error,proto3" json:"error,omitempty"` +} + +func (x *PutFileReply) Reset() { + *x = PutFileReply{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *PutFileReply) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PutFileReply) ProtoMessage() {} + +func (x *PutFileReply) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[5] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PutFileReply.ProtoReflect.Descriptor instead. +func (*PutFileReply) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{5} +} + +func (x *PutFileReply) GetOk() bool { + if x != nil { + return x.Ok + } + return false +} + +func (x *PutFileReply) GetBytes() uint64 { + if x != nil { + return x.Bytes + } + return 0 +} + +func (x *PutFileReply) GetElapsedMs() float64 { + if x != nil { + return x.ElapsedMs + } + return 0 +} + +func (x *PutFileReply) GetExisted() bool { + if x != nil { + return x.Existed + } + return false +} + +func (x *PutFileReply) GetError() string { + if x != nil { + return x.Error + } + return "" +} + +type GetFileRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Key string `protobuf:"bytes,1,opt,name=key,proto3" json:"key,omitempty"` + Path string `protobuf:"bytes,2,opt,name=path,proto3" json:"path,omitempty"` // local file to write +} + +func (x *GetFileRequest) Reset() { + *x = GetFileRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[6] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetFileRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetFileRequest) ProtoMessage() {} + +func (x *GetFileRequest) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[6] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetFileRequest.ProtoReflect.Descriptor instead. +func (*GetFileRequest) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{6} +} + +func (x *GetFileRequest) GetKey() string { + if x != nil { + return x.Key + } + return "" +} + +func (x *GetFileRequest) GetPath() string { + if x != nil { + return x.Path + } + return "" +} + +type GetFileReply struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Ok bool `protobuf:"varint,1,opt,name=ok,proto3" json:"ok,omitempty"` + Bytes uint64 `protobuf:"varint,2,opt,name=bytes,proto3" json:"bytes,omitempty"` + ElapsedMs float64 `protobuf:"fixed64,3,opt,name=elapsed_ms,json=elapsedMs,proto3" json:"elapsed_ms,omitempty"` + Found bool `protobuf:"varint,4,opt,name=found,proto3" json:"found,omitempty"` + Error string `protobuf:"bytes,5,opt,name=error,proto3" json:"error,omitempty"` +} + +func (x *GetFileReply) Reset() { + *x = GetFileReply{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[7] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetFileReply) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetFileReply) ProtoMessage() {} + +func (x *GetFileReply) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[7] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetFileReply.ProtoReflect.Descriptor instead. +func (*GetFileReply) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{7} +} + +func (x *GetFileReply) GetOk() bool { + if x != nil { + return x.Ok + } + return false +} + +func (x *GetFileReply) GetBytes() uint64 { + if x != nil { + return x.Bytes + } + return 0 +} + +func (x *GetFileReply) GetElapsedMs() float64 { + if x != nil { + return x.ElapsedMs + } + return 0 +} + +func (x *GetFileReply) GetFound() bool { + if x != nil { + return x.Found + } + return false +} + +func (x *GetFileReply) GetError() string { + if x != nil { + return x.Error + } + return "" +} + +type PutBytesRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Key string `protobuf:"bytes,1,opt,name=key,proto3" json:"key,omitempty"` + Value []byte `protobuf:"bytes,2,opt,name=value,proto3" json:"value,omitempty"` + ReplicaNum uint32 `protobuf:"varint,3,opt,name=replica_num,json=replicaNum,proto3" json:"replica_num,omitempty"` + SoftPin bool `protobuf:"varint,4,opt,name=soft_pin,json=softPin,proto3" json:"soft_pin,omitempty"` +} + +func (x *PutBytesRequest) Reset() { + *x = PutBytesRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[8] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *PutBytesRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PutBytesRequest) ProtoMessage() {} + +func (x *PutBytesRequest) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[8] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PutBytesRequest.ProtoReflect.Descriptor instead. +func (*PutBytesRequest) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{8} +} + +func (x *PutBytesRequest) GetKey() string { + if x != nil { + return x.Key + } + return "" +} + +func (x *PutBytesRequest) GetValue() []byte { + if x != nil { + return x.Value + } + return nil +} + +func (x *PutBytesRequest) GetReplicaNum() uint32 { + if x != nil { + return x.ReplicaNum + } + return 0 +} + +func (x *PutBytesRequest) GetSoftPin() bool { + if x != nil { + return x.SoftPin + } + return false +} + +type PutBytesReply struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Ok bool `protobuf:"varint,1,opt,name=ok,proto3" json:"ok,omitempty"` + Bytes uint64 `protobuf:"varint,2,opt,name=bytes,proto3" json:"bytes,omitempty"` + ElapsedMs float64 `protobuf:"fixed64,3,opt,name=elapsed_ms,json=elapsedMs,proto3" json:"elapsed_ms,omitempty"` + Error string `protobuf:"bytes,4,opt,name=error,proto3" json:"error,omitempty"` +} + +func (x *PutBytesReply) Reset() { + *x = PutBytesReply{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[9] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *PutBytesReply) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PutBytesReply) ProtoMessage() {} + +func (x *PutBytesReply) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[9] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PutBytesReply.ProtoReflect.Descriptor instead. +func (*PutBytesReply) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{9} +} + +func (x *PutBytesReply) GetOk() bool { + if x != nil { + return x.Ok + } + return false +} + +func (x *PutBytesReply) GetBytes() uint64 { + if x != nil { + return x.Bytes + } + return 0 +} + +func (x *PutBytesReply) GetElapsedMs() float64 { + if x != nil { + return x.ElapsedMs + } + return 0 +} + +func (x *PutBytesReply) GetError() string { + if x != nil { + return x.Error + } + return "" +} + +type GetBytesRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Key string `protobuf:"bytes,1,opt,name=key,proto3" json:"key,omitempty"` +} + +func (x *GetBytesRequest) Reset() { + *x = GetBytesRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[10] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetBytesRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetBytesRequest) ProtoMessage() {} + +func (x *GetBytesRequest) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[10] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetBytesRequest.ProtoReflect.Descriptor instead. +func (*GetBytesRequest) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{10} +} + +func (x *GetBytesRequest) GetKey() string { + if x != nil { + return x.Key + } + return "" +} + +type GetBytesReply struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Ok bool `protobuf:"varint,1,opt,name=ok,proto3" json:"ok,omitempty"` + Value []byte `protobuf:"bytes,2,opt,name=value,proto3" json:"value,omitempty"` + ElapsedMs float64 `protobuf:"fixed64,3,opt,name=elapsed_ms,json=elapsedMs,proto3" json:"elapsed_ms,omitempty"` + Found bool `protobuf:"varint,4,opt,name=found,proto3" json:"found,omitempty"` + Error string `protobuf:"bytes,5,opt,name=error,proto3" json:"error,omitempty"` +} + +func (x *GetBytesReply) Reset() { + *x = GetBytesReply{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[11] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *GetBytesReply) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetBytesReply) ProtoMessage() {} + +func (x *GetBytesReply) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[11] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use GetBytesReply.ProtoReflect.Descriptor instead. +func (*GetBytesReply) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{11} +} + +func (x *GetBytesReply) GetOk() bool { + if x != nil { + return x.Ok + } + return false +} + +func (x *GetBytesReply) GetValue() []byte { + if x != nil { + return x.Value + } + return nil +} + +func (x *GetBytesReply) GetElapsedMs() float64 { + if x != nil { + return x.ElapsedMs + } + return 0 +} + +func (x *GetBytesReply) GetFound() bool { + if x != nil { + return x.Found + } + return false +} + +func (x *GetBytesReply) GetError() string { + if x != nil { + return x.Error + } + return "" +} + +type RemoveRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Key string `protobuf:"bytes,1,opt,name=key,proto3" json:"key,omitempty"` + Force bool `protobuf:"varint,2,opt,name=force,proto3" json:"force,omitempty"` +} + +func (x *RemoveRequest) Reset() { + *x = RemoveRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[12] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *RemoveRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RemoveRequest) ProtoMessage() {} + +func (x *RemoveRequest) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[12] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RemoveRequest.ProtoReflect.Descriptor instead. +func (*RemoveRequest) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{12} +} + +func (x *RemoveRequest) GetKey() string { + if x != nil { + return x.Key + } + return "" +} + +func (x *RemoveRequest) GetForce() bool { + if x != nil { + return x.Force + } + return false +} + +type RemoveReply struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + Ok bool `protobuf:"varint,1,opt,name=ok,proto3" json:"ok,omitempty"` + Error string `protobuf:"bytes,2,opt,name=error,proto3" json:"error,omitempty"` +} + +func (x *RemoveReply) Reset() { + *x = RemoveReply{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[13] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *RemoveReply) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RemoveReply) ProtoMessage() {} + +func (x *RemoveReply) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[13] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RemoveReply.ProtoReflect.Descriptor instead. +func (*RemoveReply) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{13} +} + +func (x *RemoveReply) GetOk() bool { + if x != nil { + return x.Ok + } + return false +} + +func (x *RemoveReply) GetError() string { + if x != nil { + return x.Error + } + return "" +} + +type StatsRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields +} + +func (x *StatsRequest) Reset() { + *x = StatsRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[14] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *StatsRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*StatsRequest) ProtoMessage() {} + +func (x *StatsRequest) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[14] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use StatsRequest.ProtoReflect.Descriptor instead. +func (*StatsRequest) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{14} +} + +type StatsReply struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + PutOps uint64 `protobuf:"varint,1,opt,name=put_ops,json=putOps,proto3" json:"put_ops,omitempty"` + GetOps uint64 `protobuf:"varint,2,opt,name=get_ops,json=getOps,proto3" json:"get_ops,omitempty"` + ExistsOps uint64 `protobuf:"varint,3,opt,name=exists_ops,json=existsOps,proto3" json:"exists_ops,omitempty"` + PutBytes uint64 `protobuf:"varint,4,opt,name=put_bytes,json=putBytes,proto3" json:"put_bytes,omitempty"` + GetBytes uint64 `protobuf:"varint,5,opt,name=get_bytes,json=getBytes,proto3" json:"get_bytes,omitempty"` + PutMsTotal float64 `protobuf:"fixed64,6,opt,name=put_ms_total,json=putMsTotal,proto3" json:"put_ms_total,omitempty"` + GetMsTotal float64 `protobuf:"fixed64,7,opt,name=get_ms_total,json=getMsTotal,proto3" json:"get_ms_total,omitempty"` + Backend string `protobuf:"bytes,8,opt,name=backend,proto3" json:"backend,omitempty"` +} + +func (x *StatsReply) Reset() { + *x = StatsReply{} + if protoimpl.UnsafeEnabled { + mi := &file_storeproxy_proto_msgTypes[15] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *StatsReply) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*StatsReply) ProtoMessage() {} + +func (x *StatsReply) ProtoReflect() protoreflect.Message { + mi := &file_storeproxy_proto_msgTypes[15] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use StatsReply.ProtoReflect.Descriptor instead. +func (*StatsReply) Descriptor() ([]byte, []int) { + return file_storeproxy_proto_rawDescGZIP(), []int{15} +} + +func (x *StatsReply) GetPutOps() uint64 { + if x != nil { + return x.PutOps + } + return 0 +} + +func (x *StatsReply) GetGetOps() uint64 { + if x != nil { + return x.GetOps + } + return 0 +} + +func (x *StatsReply) GetExistsOps() uint64 { + if x != nil { + return x.ExistsOps + } + return 0 +} + +func (x *StatsReply) GetPutBytes() uint64 { + if x != nil { + return x.PutBytes + } + return 0 +} + +func (x *StatsReply) GetGetBytes() uint64 { + if x != nil { + return x.GetBytes + } + return 0 +} + +func (x *StatsReply) GetPutMsTotal() float64 { + if x != nil { + return x.PutMsTotal + } + return 0 +} + +func (x *StatsReply) GetGetMsTotal() float64 { + if x != nil { + return x.GetMsTotal + } + return 0 +} + +func (x *StatsReply) GetBackend() string { + if x != nil { + return x.Backend + } + return "" +} + +var File_storeproxy_proto protoreflect.FileDescriptor + +var file_storeproxy_proto_rawDesc = []byte{ + 0x0a, 0x10, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2e, 0x70, 0x72, 0x6f, + 0x74, 0x6f, 0x12, 0x0d, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2e, 0x76, + 0x31, 0x22, 0x0f, 0x0a, 0x0d, 0x48, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x52, 0x65, 0x71, 0x75, 0x65, + 0x73, 0x74, 0x22, 0x9b, 0x01, 0x0a, 0x0b, 0x48, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x52, 0x65, 0x70, + 0x6c, 0x79, 0x12, 0x0e, 0x0a, 0x02, 0x6f, 0x6b, 0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x02, + 0x6f, 0x6b, 0x12, 0x18, 0x0a, 0x07, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x18, 0x02, 0x20, + 0x01, 0x28, 0x09, 0x52, 0x07, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x12, 0x1a, 0x0a, 0x08, + 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, + 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x12, 0x16, 0x0a, 0x06, 0x6d, 0x61, 0x73, 0x74, + 0x65, 0x72, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x6d, 0x61, 0x73, 0x74, 0x65, 0x72, + 0x12, 0x16, 0x0a, 0x06, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, + 0x52, 0x06, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x12, 0x16, 0x0a, 0x06, 0x64, 0x65, 0x74, 0x61, + 0x69, 0x6c, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x64, 0x65, 0x74, 0x61, 0x69, 0x6c, + 0x22, 0x23, 0x0a, 0x0d, 0x45, 0x78, 0x69, 0x73, 0x74, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x12, 0x12, 0x0a, 0x04, 0x6b, 0x65, 0x79, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x09, 0x52, + 0x04, 0x6b, 0x65, 0x79, 0x73, 0x22, 0x27, 0x0a, 0x0b, 0x45, 0x78, 0x69, 0x73, 0x74, 0x73, 0x52, + 0x65, 0x70, 0x6c, 0x79, 0x12, 0x18, 0x0a, 0x07, 0x70, 0x72, 0x65, 0x73, 0x65, 0x6e, 0x74, 0x18, + 0x01, 0x20, 0x03, 0x28, 0x05, 0x52, 0x07, 0x70, 0x72, 0x65, 0x73, 0x65, 0x6e, 0x74, 0x22, 0x98, + 0x01, 0x0a, 0x0e, 0x50, 0x75, 0x74, 0x46, 0x69, 0x6c, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x12, 0x10, 0x0a, 0x03, 0x6b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, + 0x6b, 0x65, 0x79, 0x12, 0x12, 0x0a, 0x04, 0x70, 0x61, 0x74, 0x68, 0x18, 0x02, 0x20, 0x01, 0x28, + 0x09, 0x52, 0x04, 0x70, 0x61, 0x74, 0x68, 0x12, 0x1f, 0x0a, 0x0b, 0x72, 0x65, 0x70, 0x6c, 0x69, + 0x63, 0x61, 0x5f, 0x6e, 0x75, 0x6d, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x0a, 0x72, 0x65, + 0x70, 0x6c, 0x69, 0x63, 0x61, 0x4e, 0x75, 0x6d, 0x12, 0x19, 0x0a, 0x08, 0x73, 0x6f, 0x66, 0x74, + 0x5f, 0x70, 0x69, 0x6e, 0x18, 0x04, 0x20, 0x01, 0x28, 0x08, 0x52, 0x07, 0x73, 0x6f, 0x66, 0x74, + 0x50, 0x69, 0x6e, 0x12, 0x24, 0x0a, 0x0e, 0x73, 0x6b, 0x69, 0x70, 0x5f, 0x69, 0x66, 0x5f, 0x65, + 0x78, 0x69, 0x73, 0x74, 0x73, 0x18, 0x05, 0x20, 0x01, 0x28, 0x08, 0x52, 0x0c, 0x73, 0x6b, 0x69, + 0x70, 0x49, 0x66, 0x45, 0x78, 0x69, 0x73, 0x74, 0x73, 0x22, 0x83, 0x01, 0x0a, 0x0c, 0x50, 0x75, + 0x74, 0x46, 0x69, 0x6c, 0x65, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x0e, 0x0a, 0x02, 0x6f, 0x6b, + 0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x02, 0x6f, 0x6b, 0x12, 0x14, 0x0a, 0x05, 0x62, 0x79, + 0x74, 0x65, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x04, 0x52, 0x05, 0x62, 0x79, 0x74, 0x65, 0x73, + 0x12, 0x1d, 0x0a, 0x0a, 0x65, 0x6c, 0x61, 0x70, 0x73, 0x65, 0x64, 0x5f, 0x6d, 0x73, 0x18, 0x03, + 0x20, 0x01, 0x28, 0x01, 0x52, 0x09, 0x65, 0x6c, 0x61, 0x70, 0x73, 0x65, 0x64, 0x4d, 0x73, 0x12, + 0x18, 0x0a, 0x07, 0x65, 0x78, 0x69, 0x73, 0x74, 0x65, 0x64, 0x18, 0x04, 0x20, 0x01, 0x28, 0x08, + 0x52, 0x07, 0x65, 0x78, 0x69, 0x73, 0x74, 0x65, 0x64, 0x12, 0x14, 0x0a, 0x05, 0x65, 0x72, 0x72, + 0x6f, 0x72, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x22, + 0x36, 0x0a, 0x0e, 0x47, 0x65, 0x74, 0x46, 0x69, 0x6c, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x12, 0x10, 0x0a, 0x03, 0x6b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, + 0x6b, 0x65, 0x79, 0x12, 0x12, 0x0a, 0x04, 0x70, 0x61, 0x74, 0x68, 0x18, 0x02, 0x20, 0x01, 0x28, + 0x09, 0x52, 0x04, 0x70, 0x61, 0x74, 0x68, 0x22, 0x7f, 0x0a, 0x0c, 0x47, 0x65, 0x74, 0x46, 0x69, + 0x6c, 0x65, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x0e, 0x0a, 0x02, 0x6f, 0x6b, 0x18, 0x01, 0x20, + 0x01, 0x28, 0x08, 0x52, 0x02, 0x6f, 0x6b, 0x12, 0x14, 0x0a, 0x05, 0x62, 0x79, 0x74, 0x65, 0x73, + 0x18, 0x02, 0x20, 0x01, 0x28, 0x04, 0x52, 0x05, 0x62, 0x79, 0x74, 0x65, 0x73, 0x12, 0x1d, 0x0a, + 0x0a, 0x65, 0x6c, 0x61, 0x70, 0x73, 0x65, 0x64, 0x5f, 0x6d, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, + 0x01, 0x52, 0x09, 0x65, 0x6c, 0x61, 0x70, 0x73, 0x65, 0x64, 0x4d, 0x73, 0x12, 0x14, 0x0a, 0x05, + 0x66, 0x6f, 0x75, 0x6e, 0x64, 0x18, 0x04, 0x20, 0x01, 0x28, 0x08, 0x52, 0x05, 0x66, 0x6f, 0x75, + 0x6e, 0x64, 0x12, 0x14, 0x0a, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x18, 0x05, 0x20, 0x01, 0x28, + 0x09, 0x52, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x22, 0x75, 0x0a, 0x0f, 0x50, 0x75, 0x74, 0x42, + 0x79, 0x74, 0x65, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x10, 0x0a, 0x03, 0x6b, + 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6b, 0x65, 0x79, 0x12, 0x14, 0x0a, + 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x05, 0x76, 0x61, + 0x6c, 0x75, 0x65, 0x12, 0x1f, 0x0a, 0x0b, 0x72, 0x65, 0x70, 0x6c, 0x69, 0x63, 0x61, 0x5f, 0x6e, + 0x75, 0x6d, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x0a, 0x72, 0x65, 0x70, 0x6c, 0x69, 0x63, + 0x61, 0x4e, 0x75, 0x6d, 0x12, 0x19, 0x0a, 0x08, 0x73, 0x6f, 0x66, 0x74, 0x5f, 0x70, 0x69, 0x6e, + 0x18, 0x04, 0x20, 0x01, 0x28, 0x08, 0x52, 0x07, 0x73, 0x6f, 0x66, 0x74, 0x50, 0x69, 0x6e, 0x22, + 0x6a, 0x0a, 0x0d, 0x50, 0x75, 0x74, 0x42, 0x79, 0x74, 0x65, 0x73, 0x52, 0x65, 0x70, 0x6c, 0x79, + 0x12, 0x0e, 0x0a, 0x02, 0x6f, 0x6b, 0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x02, 0x6f, 0x6b, + 0x12, 0x14, 0x0a, 0x05, 0x62, 0x79, 0x74, 0x65, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x04, 0x52, + 0x05, 0x62, 0x79, 0x74, 0x65, 0x73, 0x12, 0x1d, 0x0a, 0x0a, 0x65, 0x6c, 0x61, 0x70, 0x73, 0x65, + 0x64, 0x5f, 0x6d, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x01, 0x52, 0x09, 0x65, 0x6c, 0x61, 0x70, + 0x73, 0x65, 0x64, 0x4d, 0x73, 0x12, 0x14, 0x0a, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x18, 0x04, + 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x22, 0x23, 0x0a, 0x0f, 0x47, + 0x65, 0x74, 0x42, 0x79, 0x74, 0x65, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x10, + 0x0a, 0x03, 0x6b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6b, 0x65, 0x79, + 0x22, 0x80, 0x01, 0x0a, 0x0d, 0x47, 0x65, 0x74, 0x42, 0x79, 0x74, 0x65, 0x73, 0x52, 0x65, 0x70, + 0x6c, 0x79, 0x12, 0x0e, 0x0a, 0x02, 0x6f, 0x6b, 0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x02, + 0x6f, 0x6b, 0x12, 0x14, 0x0a, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, + 0x0c, 0x52, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x12, 0x1d, 0x0a, 0x0a, 0x65, 0x6c, 0x61, 0x70, + 0x73, 0x65, 0x64, 0x5f, 0x6d, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x01, 0x52, 0x09, 0x65, 0x6c, + 0x61, 0x70, 0x73, 0x65, 0x64, 0x4d, 0x73, 0x12, 0x14, 0x0a, 0x05, 0x66, 0x6f, 0x75, 0x6e, 0x64, + 0x18, 0x04, 0x20, 0x01, 0x28, 0x08, 0x52, 0x05, 0x66, 0x6f, 0x75, 0x6e, 0x64, 0x12, 0x14, 0x0a, + 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x65, 0x72, + 0x72, 0x6f, 0x72, 0x22, 0x37, 0x0a, 0x0d, 0x52, 0x65, 0x6d, 0x6f, 0x76, 0x65, 0x52, 0x65, 0x71, + 0x75, 0x65, 0x73, 0x74, 0x12, 0x10, 0x0a, 0x03, 0x6b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, + 0x09, 0x52, 0x03, 0x6b, 0x65, 0x79, 0x12, 0x14, 0x0a, 0x05, 0x66, 0x6f, 0x72, 0x63, 0x65, 0x18, + 0x02, 0x20, 0x01, 0x28, 0x08, 0x52, 0x05, 0x66, 0x6f, 0x72, 0x63, 0x65, 0x22, 0x33, 0x0a, 0x0b, + 0x52, 0x65, 0x6d, 0x6f, 0x76, 0x65, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x0e, 0x0a, 0x02, 0x6f, + 0x6b, 0x18, 0x01, 0x20, 0x01, 0x28, 0x08, 0x52, 0x02, 0x6f, 0x6b, 0x12, 0x14, 0x0a, 0x05, 0x65, + 0x72, 0x72, 0x6f, 0x72, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x65, 0x72, 0x72, 0x6f, + 0x72, 0x22, 0x0e, 0x0a, 0x0c, 0x53, 0x74, 0x61, 0x74, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x22, 0xf5, 0x01, 0x0a, 0x0a, 0x53, 0x74, 0x61, 0x74, 0x73, 0x52, 0x65, 0x70, 0x6c, 0x79, + 0x12, 0x17, 0x0a, 0x07, 0x70, 0x75, 0x74, 0x5f, 0x6f, 0x70, 0x73, 0x18, 0x01, 0x20, 0x01, 0x28, + 0x04, 0x52, 0x06, 0x70, 0x75, 0x74, 0x4f, 0x70, 0x73, 0x12, 0x17, 0x0a, 0x07, 0x67, 0x65, 0x74, + 0x5f, 0x6f, 0x70, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x04, 0x52, 0x06, 0x67, 0x65, 0x74, 0x4f, + 0x70, 0x73, 0x12, 0x1d, 0x0a, 0x0a, 0x65, 0x78, 0x69, 0x73, 0x74, 0x73, 0x5f, 0x6f, 0x70, 0x73, + 0x18, 0x03, 0x20, 0x01, 0x28, 0x04, 0x52, 0x09, 0x65, 0x78, 0x69, 0x73, 0x74, 0x73, 0x4f, 0x70, + 0x73, 0x12, 0x1b, 0x0a, 0x09, 0x70, 0x75, 0x74, 0x5f, 0x62, 0x79, 0x74, 0x65, 0x73, 0x18, 0x04, + 0x20, 0x01, 0x28, 0x04, 0x52, 0x08, 0x70, 0x75, 0x74, 0x42, 0x79, 0x74, 0x65, 0x73, 0x12, 0x1b, + 0x0a, 0x09, 0x67, 0x65, 0x74, 0x5f, 0x62, 0x79, 0x74, 0x65, 0x73, 0x18, 0x05, 0x20, 0x01, 0x28, + 0x04, 0x52, 0x08, 0x67, 0x65, 0x74, 0x42, 0x79, 0x74, 0x65, 0x73, 0x12, 0x20, 0x0a, 0x0c, 0x70, + 0x75, 0x74, 0x5f, 0x6d, 0x73, 0x5f, 0x74, 0x6f, 0x74, 0x61, 0x6c, 0x18, 0x06, 0x20, 0x01, 0x28, + 0x01, 0x52, 0x0a, 0x70, 0x75, 0x74, 0x4d, 0x73, 0x54, 0x6f, 0x74, 0x61, 0x6c, 0x12, 0x20, 0x0a, + 0x0c, 0x67, 0x65, 0x74, 0x5f, 0x6d, 0x73, 0x5f, 0x74, 0x6f, 0x74, 0x61, 0x6c, 0x18, 0x07, 0x20, + 0x01, 0x28, 0x01, 0x52, 0x0a, 0x67, 0x65, 0x74, 0x4d, 0x73, 0x54, 0x6f, 0x74, 0x61, 0x6c, 0x12, + 0x18, 0x0a, 0x07, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x18, 0x08, 0x20, 0x01, 0x28, 0x09, + 0x52, 0x07, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x32, 0xbb, 0x04, 0x0a, 0x0a, 0x53, 0x74, + 0x6f, 0x72, 0x65, 0x50, 0x72, 0x6f, 0x78, 0x79, 0x12, 0x42, 0x0a, 0x06, 0x48, 0x65, 0x61, 0x6c, + 0x74, 0x68, 0x12, 0x1c, 0x2e, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2e, + 0x76, 0x31, 0x2e, 0x48, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, + 0x1a, 0x1a, 0x2e, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2e, 0x76, 0x31, + 0x2e, 0x48, 0x65, 0x61, 0x6c, 0x74, 0x68, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x42, 0x0a, 0x06, + 0x45, 0x78, 0x69, 0x73, 0x74, 0x73, 0x12, 0x1c, 0x2e, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, + 0x6f, 0x78, 0x79, 0x2e, 0x76, 0x31, 0x2e, 0x45, 0x78, 0x69, 0x73, 0x74, 0x73, 0x52, 0x65, 0x71, + 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1a, 0x2e, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, + 0x79, 0x2e, 0x76, 0x31, 0x2e, 0x45, 0x78, 0x69, 0x73, 0x74, 0x73, 0x52, 0x65, 0x70, 0x6c, 0x79, + 0x12, 0x45, 0x0a, 0x07, 0x50, 0x75, 0x74, 0x46, 0x69, 0x6c, 0x65, 0x12, 0x1d, 0x2e, 0x73, 0x74, + 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2e, 0x76, 0x31, 0x2e, 0x50, 0x75, 0x74, 0x46, + 0x69, 0x6c, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1b, 0x2e, 0x73, 0x74, 0x6f, + 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2e, 0x76, 0x31, 0x2e, 0x50, 0x75, 0x74, 0x46, 0x69, + 0x6c, 0x65, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x45, 0x0a, 0x07, 0x47, 0x65, 0x74, 0x46, 0x69, + 0x6c, 0x65, 0x12, 0x1d, 0x2e, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2e, + 0x76, 0x31, 0x2e, 0x47, 0x65, 0x74, 0x46, 0x69, 0x6c, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x1a, 0x1b, 0x2e, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2e, 0x76, + 0x31, 0x2e, 0x47, 0x65, 0x74, 0x46, 0x69, 0x6c, 0x65, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x48, + 0x0a, 0x08, 0x50, 0x75, 0x74, 0x42, 0x79, 0x74, 0x65, 0x73, 0x12, 0x1e, 0x2e, 0x73, 0x74, 0x6f, + 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2e, 0x76, 0x31, 0x2e, 0x50, 0x75, 0x74, 0x42, 0x79, + 0x74, 0x65, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1c, 0x2e, 0x73, 0x74, 0x6f, + 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2e, 0x76, 0x31, 0x2e, 0x50, 0x75, 0x74, 0x42, 0x79, + 0x74, 0x65, 0x73, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x48, 0x0a, 0x08, 0x47, 0x65, 0x74, 0x42, + 0x79, 0x74, 0x65, 0x73, 0x12, 0x1e, 0x2e, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, + 0x79, 0x2e, 0x76, 0x31, 0x2e, 0x47, 0x65, 0x74, 0x42, 0x79, 0x74, 0x65, 0x73, 0x52, 0x65, 0x71, + 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1c, 0x2e, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, + 0x79, 0x2e, 0x76, 0x31, 0x2e, 0x47, 0x65, 0x74, 0x42, 0x79, 0x74, 0x65, 0x73, 0x52, 0x65, 0x70, + 0x6c, 0x79, 0x12, 0x42, 0x0a, 0x06, 0x52, 0x65, 0x6d, 0x6f, 0x76, 0x65, 0x12, 0x1c, 0x2e, 0x73, + 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2e, 0x76, 0x31, 0x2e, 0x52, 0x65, 0x6d, + 0x6f, 0x76, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1a, 0x2e, 0x73, 0x74, 0x6f, + 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2e, 0x76, 0x31, 0x2e, 0x52, 0x65, 0x6d, 0x6f, 0x76, + 0x65, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x12, 0x3f, 0x0a, 0x05, 0x53, 0x74, 0x61, 0x74, 0x73, 0x12, + 0x1b, 0x2e, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2e, 0x76, 0x31, 0x2e, + 0x53, 0x74, 0x61, 0x74, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x19, 0x2e, 0x73, + 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2e, 0x76, 0x31, 0x2e, 0x53, 0x74, 0x61, + 0x74, 0x73, 0x52, 0x65, 0x70, 0x6c, 0x79, 0x42, 0x53, 0x5a, 0x51, 0x67, 0x69, 0x74, 0x68, 0x75, + 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x6d, 0x6f, 0x6f, 0x6e, 0x63, 0x61, 0x6b, 0x65, 0x2d, 0x61, + 0x69, 0x2f, 0x6f, 0x6c, 0x6c, 0x61, 0x6d, 0x61, 0x2d, 0x6d, 0x6f, 0x6f, 0x6e, 0x63, 0x61, 0x6b, + 0x65, 0x2d, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2f, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x6e, 0x61, + 0x6c, 0x2f, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x2f, 0x70, 0x62, 0x3b, + 0x73, 0x74, 0x6f, 0x72, 0x65, 0x70, 0x72, 0x6f, 0x78, 0x79, 0x70, 0x62, 0x62, 0x06, 0x70, 0x72, + 0x6f, 0x74, 0x6f, 0x33, +} + +var ( + file_storeproxy_proto_rawDescOnce sync.Once + file_storeproxy_proto_rawDescData = file_storeproxy_proto_rawDesc +) + +func file_storeproxy_proto_rawDescGZIP() []byte { + file_storeproxy_proto_rawDescOnce.Do(func() { + file_storeproxy_proto_rawDescData = protoimpl.X.CompressGZIP(file_storeproxy_proto_rawDescData) + }) + return file_storeproxy_proto_rawDescData +} + +var file_storeproxy_proto_msgTypes = make([]protoimpl.MessageInfo, 16) +var file_storeproxy_proto_goTypes = []any{ + (*HealthRequest)(nil), // 0: storeproxy.v1.HealthRequest + (*HealthReply)(nil), // 1: storeproxy.v1.HealthReply + (*ExistsRequest)(nil), // 2: storeproxy.v1.ExistsRequest + (*ExistsReply)(nil), // 3: storeproxy.v1.ExistsReply + (*PutFileRequest)(nil), // 4: storeproxy.v1.PutFileRequest + (*PutFileReply)(nil), // 5: storeproxy.v1.PutFileReply + (*GetFileRequest)(nil), // 6: storeproxy.v1.GetFileRequest + (*GetFileReply)(nil), // 7: storeproxy.v1.GetFileReply + (*PutBytesRequest)(nil), // 8: storeproxy.v1.PutBytesRequest + (*PutBytesReply)(nil), // 9: storeproxy.v1.PutBytesReply + (*GetBytesRequest)(nil), // 10: storeproxy.v1.GetBytesRequest + (*GetBytesReply)(nil), // 11: storeproxy.v1.GetBytesReply + (*RemoveRequest)(nil), // 12: storeproxy.v1.RemoveRequest + (*RemoveReply)(nil), // 13: storeproxy.v1.RemoveReply + (*StatsRequest)(nil), // 14: storeproxy.v1.StatsRequest + (*StatsReply)(nil), // 15: storeproxy.v1.StatsReply +} +var file_storeproxy_proto_depIdxs = []int32{ + 0, // 0: storeproxy.v1.StoreProxy.Health:input_type -> storeproxy.v1.HealthRequest + 2, // 1: storeproxy.v1.StoreProxy.Exists:input_type -> storeproxy.v1.ExistsRequest + 4, // 2: storeproxy.v1.StoreProxy.PutFile:input_type -> storeproxy.v1.PutFileRequest + 6, // 3: storeproxy.v1.StoreProxy.GetFile:input_type -> storeproxy.v1.GetFileRequest + 8, // 4: storeproxy.v1.StoreProxy.PutBytes:input_type -> storeproxy.v1.PutBytesRequest + 10, // 5: storeproxy.v1.StoreProxy.GetBytes:input_type -> storeproxy.v1.GetBytesRequest + 12, // 6: storeproxy.v1.StoreProxy.Remove:input_type -> storeproxy.v1.RemoveRequest + 14, // 7: storeproxy.v1.StoreProxy.Stats:input_type -> storeproxy.v1.StatsRequest + 1, // 8: storeproxy.v1.StoreProxy.Health:output_type -> storeproxy.v1.HealthReply + 3, // 9: storeproxy.v1.StoreProxy.Exists:output_type -> storeproxy.v1.ExistsReply + 5, // 10: storeproxy.v1.StoreProxy.PutFile:output_type -> storeproxy.v1.PutFileReply + 7, // 11: storeproxy.v1.StoreProxy.GetFile:output_type -> storeproxy.v1.GetFileReply + 9, // 12: storeproxy.v1.StoreProxy.PutBytes:output_type -> storeproxy.v1.PutBytesReply + 11, // 13: storeproxy.v1.StoreProxy.GetBytes:output_type -> storeproxy.v1.GetBytesReply + 13, // 14: storeproxy.v1.StoreProxy.Remove:output_type -> storeproxy.v1.RemoveReply + 15, // 15: storeproxy.v1.StoreProxy.Stats:output_type -> storeproxy.v1.StatsReply + 8, // [8:16] is the sub-list for method output_type + 0, // [0:8] is the sub-list for method input_type + 0, // [0:0] is the sub-list for extension type_name + 0, // [0:0] is the sub-list for extension extendee + 0, // [0:0] is the sub-list for field type_name +} + +func init() { file_storeproxy_proto_init() } +func file_storeproxy_proto_init() { + if File_storeproxy_proto != nil { + return + } + if !protoimpl.UnsafeEnabled { + file_storeproxy_proto_msgTypes[0].Exporter = func(v any, i int) any { + switch v := v.(*HealthRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[1].Exporter = func(v any, i int) any { + switch v := v.(*HealthReply); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[2].Exporter = func(v any, i int) any { + switch v := v.(*ExistsRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[3].Exporter = func(v any, i int) any { + switch v := v.(*ExistsReply); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[4].Exporter = func(v any, i int) any { + switch v := v.(*PutFileRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[5].Exporter = func(v any, i int) any { + switch v := v.(*PutFileReply); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[6].Exporter = func(v any, i int) any { + switch v := v.(*GetFileRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[7].Exporter = func(v any, i int) any { + switch v := v.(*GetFileReply); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[8].Exporter = func(v any, i int) any { + switch v := v.(*PutBytesRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[9].Exporter = func(v any, i int) any { + switch v := v.(*PutBytesReply); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[10].Exporter = func(v any, i int) any { + switch v := v.(*GetBytesRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[11].Exporter = func(v any, i int) any { + switch v := v.(*GetBytesReply); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[12].Exporter = func(v any, i int) any { + switch v := v.(*RemoveRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[13].Exporter = func(v any, i int) any { + switch v := v.(*RemoveReply); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[14].Exporter = func(v any, i int) any { + switch v := v.(*StatsRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_storeproxy_proto_msgTypes[15].Exporter = func(v any, i int) any { + switch v := v.(*StatsReply); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + } + type x struct{} + out := protoimpl.TypeBuilder{ + File: protoimpl.DescBuilder{ + GoPackagePath: reflect.TypeOf(x{}).PkgPath(), + RawDescriptor: file_storeproxy_proto_rawDesc, + NumEnums: 0, + NumMessages: 16, + NumExtensions: 0, + NumServices: 1, + }, + GoTypes: file_storeproxy_proto_goTypes, + DependencyIndexes: file_storeproxy_proto_depIdxs, + MessageInfos: file_storeproxy_proto_msgTypes, + }.Build() + File_storeproxy_proto = out.File + file_storeproxy_proto_rawDesc = nil + file_storeproxy_proto_goTypes = nil + file_storeproxy_proto_depIdxs = nil +} diff --git a/mooncake-integration/ollama/bridge/internal/storeproxy/pb/storeproxy_grpc.pb.go b/mooncake-integration/ollama/bridge/internal/storeproxy/pb/storeproxy_grpc.pb.go new file mode 100644 index 00000000..b6712c73 --- /dev/null +++ b/mooncake-integration/ollama/bridge/internal/storeproxy/pb/storeproxy_grpc.pb.go @@ -0,0 +1,396 @@ +// Code generated by protoc-gen-go-grpc. DO NOT EDIT. +// versions: +// - protoc-gen-go-grpc v1.4.0 +// - protoc v3.12.4 +// source: storeproxy.proto + +// StoreProxy is the gRPC contract between the Go sidecar (client) and the +// Python Mooncake store proxy (server). The Python side owns the single, +// long-lived, warm mooncake.store.MooncakeDistributedStore handle. +// +// Design note: KV snapshots can be multiple GiB. We therefore pass *file paths* +// (in the shared --slot-save-path directory) instead of streaming bytes through +// Go. The proxy reads/writes those files directly to/from the store, so a +// multi-GiB KV blob is copied at most once (file <-> store), never through the +// Go process. Small values may still be sent inline via PutBytes/GetBytes. + +package storeproxypb + +import ( + context "context" + grpc "google.golang.org/grpc" + codes "google.golang.org/grpc/codes" + status "google.golang.org/grpc/status" +) + +// This is a compile-time assertion to ensure that this generated file +// is compatible with the grpc package it is being compiled against. +// Requires gRPC-Go v1.62.0 or later. +const _ = grpc.SupportPackageIsVersion8 + +const ( + StoreProxy_Health_FullMethodName = "/storeproxy.v1.StoreProxy/Health" + StoreProxy_Exists_FullMethodName = "/storeproxy.v1.StoreProxy/Exists" + StoreProxy_PutFile_FullMethodName = "/storeproxy.v1.StoreProxy/PutFile" + StoreProxy_GetFile_FullMethodName = "/storeproxy.v1.StoreProxy/GetFile" + StoreProxy_PutBytes_FullMethodName = "/storeproxy.v1.StoreProxy/PutBytes" + StoreProxy_GetBytes_FullMethodName = "/storeproxy.v1.StoreProxy/GetBytes" + StoreProxy_Remove_FullMethodName = "/storeproxy.v1.StoreProxy/Remove" + StoreProxy_Stats_FullMethodName = "/storeproxy.v1.StoreProxy/Stats" +) + +// StoreProxyClient is the client API for StoreProxy service. +// +// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. +type StoreProxyClient interface { + // Liveness + which backend/protocol/master is active. + Health(ctx context.Context, in *HealthRequest, opts ...grpc.CallOption) (*HealthReply, error) + // Batched existence check (maps to mooncake batch_is_exist). + Exists(ctx context.Context, in *ExistsRequest, opts ...grpc.CallOption) (*ExistsReply, error) + // Store a KV snapshot file under `key`. + PutFile(ctx context.Context, in *PutFileRequest, opts ...grpc.CallOption) (*PutFileReply, error) + // Materialize `key` into a local file at `path`. + GetFile(ctx context.Context, in *GetFileRequest, opts ...grpc.CallOption) (*GetFileReply, error) + // Inline small-value variants (metadata, manifests). + PutBytes(ctx context.Context, in *PutBytesRequest, opts ...grpc.CallOption) (*PutBytesReply, error) + GetBytes(ctx context.Context, in *GetBytesRequest, opts ...grpc.CallOption) (*GetBytesReply, error) + Remove(ctx context.Context, in *RemoveRequest, opts ...grpc.CallOption) (*RemoveReply, error) + Stats(ctx context.Context, in *StatsRequest, opts ...grpc.CallOption) (*StatsReply, error) +} + +type storeProxyClient struct { + cc grpc.ClientConnInterface +} + +func NewStoreProxyClient(cc grpc.ClientConnInterface) StoreProxyClient { + return &storeProxyClient{cc} +} + +func (c *storeProxyClient) Health(ctx context.Context, in *HealthRequest, opts ...grpc.CallOption) (*HealthReply, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(HealthReply) + err := c.cc.Invoke(ctx, StoreProxy_Health_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *storeProxyClient) Exists(ctx context.Context, in *ExistsRequest, opts ...grpc.CallOption) (*ExistsReply, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(ExistsReply) + err := c.cc.Invoke(ctx, StoreProxy_Exists_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *storeProxyClient) PutFile(ctx context.Context, in *PutFileRequest, opts ...grpc.CallOption) (*PutFileReply, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(PutFileReply) + err := c.cc.Invoke(ctx, StoreProxy_PutFile_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *storeProxyClient) GetFile(ctx context.Context, in *GetFileRequest, opts ...grpc.CallOption) (*GetFileReply, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(GetFileReply) + err := c.cc.Invoke(ctx, StoreProxy_GetFile_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *storeProxyClient) PutBytes(ctx context.Context, in *PutBytesRequest, opts ...grpc.CallOption) (*PutBytesReply, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(PutBytesReply) + err := c.cc.Invoke(ctx, StoreProxy_PutBytes_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *storeProxyClient) GetBytes(ctx context.Context, in *GetBytesRequest, opts ...grpc.CallOption) (*GetBytesReply, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(GetBytesReply) + err := c.cc.Invoke(ctx, StoreProxy_GetBytes_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *storeProxyClient) Remove(ctx context.Context, in *RemoveRequest, opts ...grpc.CallOption) (*RemoveReply, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(RemoveReply) + err := c.cc.Invoke(ctx, StoreProxy_Remove_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *storeProxyClient) Stats(ctx context.Context, in *StatsRequest, opts ...grpc.CallOption) (*StatsReply, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(StatsReply) + err := c.cc.Invoke(ctx, StoreProxy_Stats_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +// StoreProxyServer is the server API for StoreProxy service. +// All implementations must embed UnimplementedStoreProxyServer +// for forward compatibility +type StoreProxyServer interface { + // Liveness + which backend/protocol/master is active. + Health(context.Context, *HealthRequest) (*HealthReply, error) + // Batched existence check (maps to mooncake batch_is_exist). + Exists(context.Context, *ExistsRequest) (*ExistsReply, error) + // Store a KV snapshot file under `key`. + PutFile(context.Context, *PutFileRequest) (*PutFileReply, error) + // Materialize `key` into a local file at `path`. + GetFile(context.Context, *GetFileRequest) (*GetFileReply, error) + // Inline small-value variants (metadata, manifests). + PutBytes(context.Context, *PutBytesRequest) (*PutBytesReply, error) + GetBytes(context.Context, *GetBytesRequest) (*GetBytesReply, error) + Remove(context.Context, *RemoveRequest) (*RemoveReply, error) + Stats(context.Context, *StatsRequest) (*StatsReply, error) + mustEmbedUnimplementedStoreProxyServer() +} + +// UnimplementedStoreProxyServer must be embedded to have forward compatible implementations. +type UnimplementedStoreProxyServer struct { +} + +func (UnimplementedStoreProxyServer) Health(context.Context, *HealthRequest) (*HealthReply, error) { + return nil, status.Errorf(codes.Unimplemented, "method Health not implemented") +} +func (UnimplementedStoreProxyServer) Exists(context.Context, *ExistsRequest) (*ExistsReply, error) { + return nil, status.Errorf(codes.Unimplemented, "method Exists not implemented") +} +func (UnimplementedStoreProxyServer) PutFile(context.Context, *PutFileRequest) (*PutFileReply, error) { + return nil, status.Errorf(codes.Unimplemented, "method PutFile not implemented") +} +func (UnimplementedStoreProxyServer) GetFile(context.Context, *GetFileRequest) (*GetFileReply, error) { + return nil, status.Errorf(codes.Unimplemented, "method GetFile not implemented") +} +func (UnimplementedStoreProxyServer) PutBytes(context.Context, *PutBytesRequest) (*PutBytesReply, error) { + return nil, status.Errorf(codes.Unimplemented, "method PutBytes not implemented") +} +func (UnimplementedStoreProxyServer) GetBytes(context.Context, *GetBytesRequest) (*GetBytesReply, error) { + return nil, status.Errorf(codes.Unimplemented, "method GetBytes not implemented") +} +func (UnimplementedStoreProxyServer) Remove(context.Context, *RemoveRequest) (*RemoveReply, error) { + return nil, status.Errorf(codes.Unimplemented, "method Remove not implemented") +} +func (UnimplementedStoreProxyServer) Stats(context.Context, *StatsRequest) (*StatsReply, error) { + return nil, status.Errorf(codes.Unimplemented, "method Stats not implemented") +} +func (UnimplementedStoreProxyServer) mustEmbedUnimplementedStoreProxyServer() {} + +// UnsafeStoreProxyServer may be embedded to opt out of forward compatibility for this service. +// Use of this interface is not recommended, as added methods to StoreProxyServer will +// result in compilation errors. +type UnsafeStoreProxyServer interface { + mustEmbedUnimplementedStoreProxyServer() +} + +func RegisterStoreProxyServer(s grpc.ServiceRegistrar, srv StoreProxyServer) { + s.RegisterService(&StoreProxy_ServiceDesc, srv) +} + +func _StoreProxy_Health_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(HealthRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(StoreProxyServer).Health(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: StoreProxy_Health_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(StoreProxyServer).Health(ctx, req.(*HealthRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _StoreProxy_Exists_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(ExistsRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(StoreProxyServer).Exists(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: StoreProxy_Exists_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(StoreProxyServer).Exists(ctx, req.(*ExistsRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _StoreProxy_PutFile_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(PutFileRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(StoreProxyServer).PutFile(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: StoreProxy_PutFile_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(StoreProxyServer).PutFile(ctx, req.(*PutFileRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _StoreProxy_GetFile_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(GetFileRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(StoreProxyServer).GetFile(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: StoreProxy_GetFile_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(StoreProxyServer).GetFile(ctx, req.(*GetFileRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _StoreProxy_PutBytes_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(PutBytesRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(StoreProxyServer).PutBytes(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: StoreProxy_PutBytes_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(StoreProxyServer).PutBytes(ctx, req.(*PutBytesRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _StoreProxy_GetBytes_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(GetBytesRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(StoreProxyServer).GetBytes(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: StoreProxy_GetBytes_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(StoreProxyServer).GetBytes(ctx, req.(*GetBytesRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _StoreProxy_Remove_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(RemoveRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(StoreProxyServer).Remove(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: StoreProxy_Remove_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(StoreProxyServer).Remove(ctx, req.(*RemoveRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _StoreProxy_Stats_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(StatsRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(StoreProxyServer).Stats(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: StoreProxy_Stats_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(StoreProxyServer).Stats(ctx, req.(*StatsRequest)) + } + return interceptor(ctx, in, info, handler) +} + +// StoreProxy_ServiceDesc is the grpc.ServiceDesc for StoreProxy service. +// It's only intended for direct use with grpc.RegisterService, +// and not to be introspected or modified (even as a copy) +var StoreProxy_ServiceDesc = grpc.ServiceDesc{ + ServiceName: "storeproxy.v1.StoreProxy", + HandlerType: (*StoreProxyServer)(nil), + Methods: []grpc.MethodDesc{ + { + MethodName: "Health", + Handler: _StoreProxy_Health_Handler, + }, + { + MethodName: "Exists", + Handler: _StoreProxy_Exists_Handler, + }, + { + MethodName: "PutFile", + Handler: _StoreProxy_PutFile_Handler, + }, + { + MethodName: "GetFile", + Handler: _StoreProxy_GetFile_Handler, + }, + { + MethodName: "PutBytes", + Handler: _StoreProxy_PutBytes_Handler, + }, + { + MethodName: "GetBytes", + Handler: _StoreProxy_GetBytes_Handler, + }, + { + MethodName: "Remove", + Handler: _StoreProxy_Remove_Handler, + }, + { + MethodName: "Stats", + Handler: _StoreProxy_Stats_Handler, + }, + }, + Streams: []grpc.StreamDesc{}, + Metadata: "storeproxy.proto", +} diff --git a/mooncake-integration/ollama/deploy/docker-compose.yml b/mooncake-integration/ollama/deploy/docker-compose.yml new file mode 100644 index 00000000..b68ea0e9 --- /dev/null +++ b/mooncake-integration/ollama/deploy/docker-compose.yml @@ -0,0 +1,68 @@ +# Ollama × Mooncake KVCache Bus — containerised deployment. +# +# NOTE: the tested, reproducible path on this machine is the native scripts +# (scripts/stack_up.sh). This compose file packages the same topology for +# portability. RDMA + GPUDirect require host devices to be passed through (see +# the `devices`/`ipc` settings), so it is meant for a node with Mellanox NICs and +# NVIDIA GPUs and the NVIDIA Container Toolkit installed. The `build:` services +# expect a Dockerfile in each context; add one (or replace `build:` with a +# pre-built `image:`) before `docker compose up`. Ports match scripts/env.sh. +version: "3.9" + +x-mooncake-env: &mc-env + MOONCAKE_MASTER: "127.0.0.1:52061" + MOONCAKE_PROTOCOL: "rdma" + MOONCAKE_DEVICE: "mlx5_0" + MOONCAKE_TE_META_DATA_SERVER: "P2PHANDSHAKE" + +services: + mooncake-master: + image: kvcache/mooncake:latest # or build from kvcache-ai/Mooncake + command: ["mooncake_master", "-port", "52061", "-metrics_port", "52063", + "-enable_metric_reporting=true", "-default_kv_lease_ttl=2000"] + network_mode: host + + store-proxy: + build: { context: ../mooncake-store-proxy } + command: ["python", "store_proxy.py", "--listen", "127.0.0.1:52060", + "--backend", "mooncake", "--protocol", "rdma", "--device", "mlx5_0", + "--master", "127.0.0.1:52061", "--metadata", "P2PHANDSHAKE", + "--global-segment-size", "34359738368", "--staging-mb", "4096", + "--staging-count", "3", "--warmup"] + environment: *mc-env + network_mode: host + ipc: host # shared /dev/shm for slot files + devices: ["/dev/infiniband:/dev/infiniband"] + cap_add: ["IPC_LOCK"] # RDMA memory pinning + depends_on: [mooncake-master] + + bridged: # the Go sidecar (ollama-mooncake-bridge) + build: { context: ../ollama-mooncake-bridge } + command: ["/bridged", "-grpc-addr", "127.0.0.1:52051", "-http-addr", "127.0.0.1:52052", + "-store-backend", "mooncake", "-store-proxy-addr", "127.0.0.1:52060", + "-slot-save-path", "/dev/shm/omb-slots/"] + network_mode: host + ipc: host + depends_on: [store-proxy] + + # ollama-with-patches: build from a fork with ollama-patches/ applied. + ollama: + image: ollama-mooncake:latest # ollama built with ollama-patches/* + environment: + OLLAMA_MOONCAKE_BRIDGE_URL: "http://127.0.0.1:52052" + OLLAMA_MOONCAKE_SLOT_SAVE_PATH: "/dev/shm/omb-slots" + network_mode: host + ipc: host + runtime: nvidia + depends_on: [bridged] + + prometheus: + image: prom/prometheus:latest + command: ["--config.file=/etc/prometheus/prometheus.yml"] + volumes: ["./prometheus.yml:/etc/prometheus/prometheus.yml:ro"] + network_mode: host + + grafana: + image: grafana/grafana:latest + network_mode: host + depends_on: [prometheus] diff --git a/mooncake-integration/ollama/deploy/grafana_dashboard.json b/mooncake-integration/ollama/deploy/grafana_dashboard.json new file mode 100644 index 00000000..fd3d10f4 --- /dev/null +++ b/mooncake-integration/ollama/deploy/grafana_dashboard.json @@ -0,0 +1,77 @@ +{ + "title": "Ollama × Mooncake — Local Agent Swarm KVCache Bus", + "uid": "omb-kvcache", + "schemaVersion": 39, + "time": { "from": "now-15m", "to": "now" }, + "refresh": "2s", + "panels": [ + { + "type": "stat", "title": "Prefill tokens saved (cumulative)", "gridPos": {"h":7,"w":6,"x":0,"y":0}, + "options": {"colorMode": "background", "graphMode": "area", "textMode": "value"}, + "fieldConfig": {"defaults": {"unit": "short", "color": {"mode": "fixed", "fixedColor": "green"}}}, + "targets": [{"expr": "mooncake_bridge_saved_prefill_tokens_total", "legendFormat": "saved tokens"}] + }, + { + "type": "stat", "title": "KV restored from store", "gridPos": {"h":7,"w":6,"x":6,"y":0}, + "options": {"colorMode": "background", "textMode": "value"}, + "fieldConfig": {"defaults": {"unit": "bytes", "color": {"mode": "fixed", "fixedColor": "blue"}}}, + "targets": [{"expr": "mooncake_bridge_bytes_get_total", "legendFormat": "bytes restored"}] + }, + { + "type": "stat", "title": "Cache hit ratio", "gridPos": {"h":7,"w":6,"x":12,"y":0}, + "options": {"colorMode": "background"}, + "fieldConfig": {"defaults": {"unit": "percentunit", "color": {"mode": "fixed", "fixedColor": "purple"}}}, + "targets": [{"expr": "mooncake_bridge_hits_total / clamp_min(mooncake_bridge_hits_total + mooncake_bridge_misses_total + mooncake_bridge_recompute_chosen_total, 1)", "legendFormat": "hit ratio"}] + }, + { + "type": "stat", "title": "Restores performed", "gridPos": {"h":7,"w":6,"x":18,"y":0}, + "options": {"colorMode": "background", "textMode": "value"}, + "fieldConfig": {"defaults": {"unit": "short", "color": {"mode": "fixed", "fixedColor": "green"}}}, + "targets": [{"expr": "mooncake_bridge_restore_total", "legendFormat": "restores"}] + }, + { + "type": "stat", "title": "Arbiter: learned restore BW / prefill rate", "gridPos": {"h":7,"w":24,"x":0,"y":7}, + "targets": [ + {"expr": "mooncake_bridge_learned_get_gbps", "legendFormat": "restore GB/s"}, + {"expr": "mooncake_bridge_learned_prefill_tps", "legendFormat": "prefill tok/s"} + ] + }, + { + "type": "timeseries", "title": "Saved prefill tokens / sec", "gridPos": {"h":8,"w":12,"x":0,"y":14}, + "targets": [{"expr": "rate(mooncake_bridge_saved_prefill_tokens_total[1m])", "legendFormat": "tokens/s saved"}] + }, + { + "type": "timeseries", "title": "Decisions/sec (hit vs miss vs recompute-chosen)", "gridPos": {"h":8,"w":12,"x":12,"y":14}, + "targets": [ + {"expr": "rate(mooncake_bridge_hits_total[1m])", "legendFormat": "restore (hit)"}, + {"expr": "rate(mooncake_bridge_misses_total[1m])", "legendFormat": "miss"}, + {"expr": "rate(mooncake_bridge_recompute_chosen_total[1m])", "legendFormat": "recompute (arbiter)"} + ] + }, + { + "type": "timeseries", "title": "Store traffic (bytes/s)", "gridPos": {"h":8,"w":12,"x":0,"y":22}, + "fieldConfig": {"defaults": {"unit": "Bps"}}, + "targets": [ + {"expr": "rate(mooncake_bridge_bytes_get_total[1m])", "legendFormat": "get"}, + {"expr": "rate(mooncake_bridge_bytes_put_total[1m])", "legendFormat": "put"} + ] + }, + { + "type": "timeseries", "title": "Restore latency p50/p90 (ms)", "gridPos": {"h":8,"w":12,"x":12,"y":22}, + "fieldConfig": {"defaults": {"unit": "ms"}}, + "targets": [ + {"expr": "histogram_quantile(0.5, rate(mooncake_bridge_restore_latency_ms_bucket[1m]))", "legendFormat": "p50"}, + {"expr": "histogram_quantile(0.9, rate(mooncake_bridge_restore_latency_ms_bucket[1m]))", "legendFormat": "p90"} + ] + }, + { + "type": "stat", "title": "KV prefix snapshots in radix index", "gridPos": {"h":6,"w":12,"x":0,"y":30}, + "targets": [{"expr": "mooncake_bridge_index_snapshots", "legendFormat": "snapshots"}] + }, + { + "type": "stat", "title": "Index bytes", "gridPos": {"h":6,"w":12,"x":12,"y":30}, + "fieldConfig": {"defaults": {"unit": "bytes"}}, + "targets": [{"expr": "mooncake_bridge_index_bytes", "legendFormat": "bytes"}] + } + ] +} diff --git a/mooncake-integration/ollama/deploy/mooncake_store.yaml b/mooncake-integration/ollama/deploy/mooncake_store.yaml new file mode 100644 index 00000000..a6489080 --- /dev/null +++ b/mooncake-integration/ollama/deploy/mooncake_store.yaml @@ -0,0 +1,24 @@ +# Mooncake Store configuration consumed by the store proxy / clients. +# (See mooncake.mooncake_config.MooncakeConfig — loaded via MOONCAKE_CONFIG_PATH.) +# +# This documents the knobs we tuned for the local-agent-swarm scenario. + +local_hostname: "127.0.0.1" +metadata_server: "P2PHANDSHAKE" # peer-to-peer TE handshake; no etcd/http server needed +master_server_address: "127.0.0.1:52061" + +protocol: "rdma" # tcp | rdma (rdma => GPUDirect-capable) +device_name: "mlx5_0" # or "auto-discovery" + +# Pool capacity this client contributes to the global KV store (host DRAM). +global_segment_size: 34359738368 # 32 GiB +local_buffer_size: 8589934592 # 8 GiB + +# Proxy-side tuning (store_proxy.py): +# stripe_mb: 64 # split KV blobs into <=64 MiB chunks for parallel batch transfer +# staging_mb: 4096 # size of each pre-registered (pinned) RDMA staging buffer +# staging_count: 3 # number of staging buffers == concurrency of zero-copy transfers +# +# Replication (ReplicateConfig, per Put): +# replica_num: 1 # raise for hot shared prefixes (read fan-out) +# with_soft_pin: true # keep hot prefixes resident under memory pressure diff --git a/mooncake-integration/ollama/deploy/prometheus.yml b/mooncake-integration/ollama/deploy/prometheus.yml new file mode 100644 index 00000000..f2c549f3 --- /dev/null +++ b/mooncake-integration/ollama/deploy/prometheus.yml @@ -0,0 +1,17 @@ +# Prometheus scrape config for the Ollama x Mooncake KVCache Bus. +# Run: prometheus --config.file=deploy/prometheus.yml --storage.tsdb.path=run/prometheus +global: + scrape_interval: 2s + evaluation_interval: 5s + +scrape_configs: + - job_name: ollama-mooncake-bridge + metrics_path: /metrics + static_configs: + - targets: ["127.0.0.1:52052"] # sidecar HTTP gateway (OMB_BRIDGE_HTTP_PORT) + labels: { component: sidecar } + + - job_name: mooncake-master + static_configs: + - targets: ["127.0.0.1:52063"] # mooncake_master -metrics_port (OMB_MASTER_PORT+2) + labels: { component: master } diff --git a/mooncake-integration/ollama/docs/REPORT.md b/mooncake-integration/ollama/docs/REPORT.md new file mode 100644 index 00000000..5cca7863 --- /dev/null +++ b/mooncake-integration/ollama/docs/REPORT.md @@ -0,0 +1,262 @@ +# Local agent swarm with a global KV cache: bringing Mooncake to Ollama + +> Integrating a new inference framework (Ollama) into the Mooncake PD / KVCache +> ecosystem. +> +> **One sentence:** multiple coding agents on Ollama / llama.cpp now share one +> long-prefix KV cache through the Mooncake Store — the first agent prefills a +> long repo context once, writes its KV to Mooncake, and every other agent (in a +> different process, on a different GPU, or on a different node) restores it +> instead of re-prefilling. + +Everything described here was built and measured on real hardware: **8× NVIDIA +H200**, **Mellanox RDMA NICs (mlx5)**, the real **`mooncake-transfer-engine`** +distributed store over RDMA, a CUDA build of **llama.cpp**, and +**Qwen2.5-Coder 1.5B / 7B** GGUF models. + +--- + +## 1. The gap we fill + +Ollama is the de-facto entry point for local models, multi-agent tooling and +coding assistants — yet it has **zero KV-cache sharing**. A `grep` of the Ollama +tree for `slot-save-path` / `cache-reuse` / `state_seq` is empty: prefix reuse is +limited to a single process's RAM and dies with the request. The community asked +for exactly this in *Ollama #14872 "Swarm Memory — zero-copy KV sharing across +models"* (closed, i.e. not built upstream). + +Mooncake already pools and shares KV for vLLM / SGLang / TRT-LLM in the data +centre. We extend it **down to the workstation/edge multi-agent scenario** and +**out to Ollama**, with a sidecar design that needs almost no changes to Ollama +itself. + +## 2. Architecture + +![Architecture](figures/architecture.png) + +``` + agents / patched Ollama ──HTTP/gRPC──► ollama-mooncake-bridge (Go sidecar) + │ cache-key + radix index + │ restore-vs-recompute arbiter + │ 3-stage orchestration + ┌───────────────┴───────────────┐ + /slots HTTP (Stage 1) gRPC (file paths) + or cgo _ext (Stage 2) │ + ▼ ▼ + llama.cpp server mooncake-store-proxy (Python) + (KV in GPU, per slot) owns 1 warm MooncakeDistributedStore + striped batch_put_from/get_into, + pre-registered RDMA staging pool + │ + Mooncake Store + Transfer Engine + (DRAM pool, RDMA / GPUDirect, replicas) +``` + +Three processes, each independently deployable: + +* **`ollama-mooncake-bridge`** (Go) — the brain. Pure Go, no heavy deps. Serves a + gRPC `KVCacheBus` (Lookup/Prepare/Commit) **and** an HTTP/JSON gateway with a + Prometheus `/metrics` endpoint. +* **`mooncake-store-proxy`** (Python) — owns the single warm + `mooncake.store.MooncakeDistributedStore` handle (the official client) and + exposes it over gRPC. KV blobs are passed as **file paths**, so a multi-GiB + blob is copied at most once and never streams through Go. +* **llama.cpp server** — unmodified for Stage 1; driven via its `/slots` + save/restore endpoints. + +## 3. Technical contributions + +1. **Safe, content-addressed cache keys with chained block hashing.** Every + attribute that changes KV *bytes* — model digest (from a real GGUF metadata + parser written in Go), tokenizer hash, RoPE hash, KV dtype, KV layout, ctx + length, block size — is folded into the key. Within a key space, prompt + tokens are hashed per fixed-size block with a **forward chain** + `h_i = H(h_{i-1} ‖ block_i)`, so two prompts agree on `h_i` iff they share + every token of blocks `0..i`. We would rather miss than mis-hit — a wrong + reuse silently poisons generation. + +2. **A radix tree of KV prefixes shared across processes/GPUs/nodes.** The + cross-process analogue of SGLang's RadixAttention, but the nodes reference KV + snapshots in the Mooncake Store rather than GPU memory. It gives O(blocks) + longest-prefix matching, per-prefix hotness (fan-out → replication hints) and + LRU eviction. Authoritative existence is still a single batched store + `batch_is_exist` over all block boundaries (multi-node correct). + +3. **A restore-vs-recompute cost arbiter.** A KV cache only helps if the store + delivers KV bytes *faster than the GPU regenerates them*. The break-even + bandwidth is `B* = (KV bytes/token) × (prefill tokens/s)`. The arbiter learns, + **per model**, the effective end-to-end restore bandwidth and the prefill rate + online (EWMA, seed-on-first-observation) and only restores when it is + genuinely cheaper. This makes the system **adaptive and loss-free**: it wins + where reuse helps and falls back to recompute where it does not — see §5.3. + +4. **Striped, pre-registered zero-copy RDMA transfer.** Single-object RDMA + collapses for very large objects (≈2 GB) in the store; we split KV into + ≤64 MiB chunks and move them with `batch_put_from`/`batch_get_into` over a + **pool of pre-registered (pinned) staging buffers**, sustaining full bandwidth + and avoiding per-op RDMA registration. Slot files live on `/dev/shm` (tmpfs) + so the host side never touches disk. + +5. **Stage-2 on-device path.** A cgo binding (`llama_state_seq_get/set_data_ext` + with `LLAMA_STATE_SEQ_FLAGS_ON_DEVICE`) keeps KV in device buffers for + GPUDirect RDMA, eliminating the host double-copy of llama.cpp issue #8915. + +## 4. What is real vs. designed + +| Component | Status | +|-----------|--------| +| Mooncake Store + Transfer Engine (RDMA) | **real** — `mooncake_master`, P2P handshake, mlx5 RDMA | +| llama.cpp servers on H200 | **real** — CUDA build, `/slots` save/restore, q8_0 KV | +| GGUF metadata parser, chained hashing, radix index, arbiter | **real** Go, unit-tested | +| Cross-process / cross-GPU KV reuse end-to-end | **real**, measured | +| Stage-2 `_ext`/`ON_DEVICE` export & round-trip | **real** (microbench, libllama) | +| Mooncake-TE GPUDirect of the on-device handle | **designed** (needs Mooncake C++ headers, not in the wheel) | +| Ollama patches | **real diffs**, additive + env-gated | + +## 5. Results + +### 5.1 Store transport (microbench, RDMA mlx5_0) + +Striped vs single-object, same total size, over one mlx5 NIC: + +| object size | single put / get | striped (64 MiB) put / get | +|---|---|---| +| 16 MiB | 26.5 / 5.0 GB/s | 22.8 / 28.7 GB/s | +| 128 MiB | 34.3 / 38.7 GB/s | **43.4 / 44.6 GB/s** | +| 1 GiB | 34.1 / 16.7 GB/s | 41.7 / 14.5 GB/s | +| 2 GiB | 17.5 / 12.3 GB/s | 20.6 / 10.5 GB/s | + +Striping sustains peak bandwidth where a single large object degrades, and +avoids per-op RDMA registration via the pre-registered staging pool. TCP loopback +peaks far lower and is *below* break-even, which is exactly where the arbiter +declines restore. + +### 5.2 Multi-agent swarm + +7B model, q8_0 KV, agents sharing one long repo context, real Mooncake RDMA, +**cross-process and cross-GPU** (round-robin over 2 H200s). Baseline = each agent +prefills cold (llama's own per-slot cache erased between runs for a fair +comparison). Honest TTFT = prepare(lookup + store-fetch + GPU-load) + prefill. + +A single reuser, uncontended (the `smoke_e2e` two-agent cross-GPU test): + +| stage | agent A (cold) | agent B (restore) | +|---|---|---| +| prepare (store fetch + GPU load) | 0.1 s | 1.82 s | +| prefill | 5.88 s (30 045 tok) | 0.01 s (1 tok) | +| **honest TTFT** | **5.99 s** | **1.83 s (−69 %)** | + +The swarm, mean over all agents (benchmark matrix, concurrency 2): + +| workload | mean TTFT (no share → Mooncake) | redundant prefill eliminated | swarm throughput | +|---|---|---|---| +| 7B, 30 k ctx, 6 agents | 5868 → 3898 ms (**−34 %**) | 100 % → 17 % of tokens (agent-0 only) | 1.01× | +| 7B, 16 k ctx, 8 agents | 3184 → 2120 ms (**−33 %**) | 100 % → 13 % | **1.25×** | +| 1.5B, 8 k ctx, 5 agents | 883 → 918 ms (≈ par) | (arbiter near break-even) | 1.00× | + +The per-reuser win is large (−69 %); the swarm *mean* is lower because the single +store client and single NIC serialize concurrent restores, so at high concurrency +the restore path contends with itself. The mechanism-invariant result — **only the +pioneer prefills, every other agent reuses (83 % of all prefill work eliminated)** +— holds in every run regardless of fabric contention. Scaling the store across +multiple clients / NICs (a deployment concern, not an algorithmic one) lifts the +swarm mean toward the single-reuser number; the Stage-2 on-device path (§5.4) +removes the file/​GPU-load term entirely. + +![Only the pioneer prefills](figures/per_agent.png) + +Per agent, the picture is stark: agent 0 prefills the full 28.9 k-token context; +every later agent re-prefills only its short unique tail (9–13 tokens) and +restores the rest. + +### 5.2.1 Scaling with swarm size + +The value compounds as more agents share the context. Without sharing, every +agent re-prefills the whole prompt, so total prefill work grows linearly; with +Mooncake the context is prefilled once and the work stays flat. Measured on the +7B / 30 k stack, fresh namespace per point so each is a clean cold start: + +| agents | prefill tokens (no share) | prefill tokens (Mooncake) | swarm throughput | +|---|---|---|---| +| 1 | 28 942 | 28 942 | 0.36× (pays the store write, nothing to reuse yet) | +| 2 | 57 884 | 28 953 | 0.93× | +| 4 | 115 768 | 28 975 | 1.00× | +| 6 | 173 653 | 28 998 | 1.36× | +| 8 | 231 535 | **29 018** | **1.51×** | + +![Redundant prefill work vs swarm size](figures/scaling_agents.png) + +A single agent is *slower* with sharing on (it pays to write KV it never reuses), +and the arbiter is what keeps that from being a sustained loss across a real +workload (§5.3). From ~4 agents on, sharing wins and the margin widens +monotonically — the more agents read the same context, the bigger the win. + +### 5.3 The adaptive arbiter (loss-free) + +On the tiny **1.5B** model the H200 prefills 7680 tokens in ~330 ms, so a +file-based restore (~700 ms) is *not* worth it. Running 6 agents sequentially so +the online learner can act, the arbiter's decisions were: + +``` +agent0: miss (cold prefill + store) ttft 294 ms +agent1: recompute (restore 697 ms > prefill 585 ms) ttft 304 ms +agent2: recompute (restore-path too slow to beat prefill) ttft 290 ms +agent3: recompute ttft 291 ms +agent4: recompute ttft 294 ms +agent5: recompute ttft 311 ms +``` + +![Loss-free arbiter](figures/arbiter.png) + +After **one** observation the arbiter measured that restoring (~697 ms) does not +beat prefilling (~330–585 ms) for this model and **fell back to recompute** — +preventing any sustained loss. Without the arbiter, a naive "always restore on +hit" policy regresses TTFT (≈357 → 625 ms on the same model). On the 7B model the +same logic keeps choosing restore. This is what makes KV sharing safe to enable +unconditionally. + +### 5.4 Stage-2 KV-state export (libllama microbench) + +![Stage-1 vs Stage-2 KV export](figures/stage2_path.png) + +| model / ctx | KV size | (A) `/slots` file save | (B) in-proc `get_data_ext` | (A)/(B) | round-trip | +|---|---|---|---|---|---| +| 7B / 16 k | 875 MiB | 417 ms (2.2 GB/s) | 91 ms (10.1 GB/s) | **4.6×** | PASS | +| 1.5B / 8 k | 219 MiB | 105 ms (2.2 GB/s) | 19 ms (11.9 GB/s) | **5.5×** | PASS | + +The Stage-1 file path is 4–5× slower than the raw in-process state copy (the +serialization + I/O tax of issue #8915). `ON_DEVICE` reports a tiny 0.1–0.2 MiB +host handle — the bulk KV stays on the GPU, ready for Mooncake-TE GPUDirect RDMA +with zero host copy. KV round-trips were verified functionally (seq0 vs seq1 +next-token argmax match). + +## 6. Limitations & honest framing + +* On a very fast GPU + tiny model the file-based Stage-1 path does **not** reduce + TTFT — the arbiter detects this and declines (no loss). The real TTFT wins come + from (a) larger models / longer contexts, (b) q8_0 KV (halved blob), and (c) + Stage-2 on-device transfer. +* The store-fetch path in Stage 1 is bounded by llama.cpp's file restore + (`/slots`), which materializes KV from a file rather than GPUDirect; the + effective end-to-end restore bandwidth is therefore well below the raw RDMA + bandwidth of §5.1. Stage 2 targets exactly this gap. +* Under a synchronized burst, all first-batch agents decide before any + observation lands; the arbiter adapts across a sustained workload, not within + the very first concurrent wave. +* Heterogeneous cross-*model* KV sharing is intentionally out of scope (different + models ⇒ different key space ⇒ never shared). Same-model multi-agent is the + target and is fully working. + +## 7. Reproduce + +```bash +bash scripts/setup_go.sh && bash scripts/setup_py.sh # local Go + venv + mooncake +bash scripts/setup_llama.sh && bash scripts/setup_models.sh --with-7b +bash scripts/demo.sh 7b 6 # bring up the stack + run the swarm demo +bash ollama-mooncake-bridge/cbridge/build_kvbench.sh # Stage-2 microbench +bash scripts/stack_down.sh # tear down (workspace-scoped) +``` + +Deliverables: `ollama-mooncake-bridge/` (sidecar + cgo), `mooncake-store-proxy/`, +`ollama-patches/`, `bench/`, `deploy/` (compose + Prometheus + Grafana), this +report and `figures/`. diff --git a/mooncake-integration/ollama/docs/demo_result.json b/mooncake-integration/ollama/docs/demo_result.json new file mode 100644 index 00000000..85e40646 --- /dev/null +++ b/mooncake-integration/ollama/docs/demo_result.json @@ -0,0 +1,138 @@ +{ + "config": { + "bridge": "http://127.0.0.1:52052", + "llamas": "http://127.0.0.1:52072,http://127.0.0.1:52073", + "model_path": "/data2/wangzhe/gitlink/mooncake1/.omb-state/models/qwen2.5-coder-7b-instruct-q4_k_m.gguf", + "agents": 6, + "ctx_tokens": 30000, + "share_mode": "shared_prefix", + "n_predict": 8, + "block_size": 256, + "kv_type": "q8_0", + "namespace": "canon2", + "concurrency": 2, + "slots_per_server": 4, + "replica_num": 1, + "out": "/data2/wangzhe/gitlink/mooncake1/.omb-state/run/demo_result.json" + }, + "baseline": { + "label": "cache_off", + "agents": 6, + "wall_s": 23.111, + "ttft_ms_mean": 5684.5, + "ttft_ms_median": 5685.4, + "ttft_ms_p90": 5687.4, + "prompt_tokens_recomputed": 173653, + "prompt_tokens_total": 173653, + "recompute_ratio": 1.0, + "hit_agents": 0, + "throughput_agents_per_s": 0.26 + }, + "cached": { + "label": "cache_on", + "agents": 6, + "wall_s": 20.186, + "ttft_ms_mean": 3966.9, + "ttft_ms_median": 3653.9, + "ttft_ms_p90": 3934.2, + "prompt_tokens_recomputed": 28998, + "prompt_tokens_total": 173653, + "recompute_ratio": 0.167, + "hit_agents": 5, + "throughput_agents_per_s": 0.297 + }, + "ttft_reduction_pct": 30.2, + "throughput_speedup_x": 1.14, + "per_agent_cached": [ + { + "idx": 0, + "llama": "http://127.0.0.1:52072", + "tokens": 28942, + "hit": false, + "decision": "miss", + "restored": false, + "store_get_ms": 0, + "bytes": 0, + "prepare_ms": 68.9809100003913, + "prompt_n": 28942, + "prompt_ms": 5587.763, + "comp_wall_ms": 5798.791166977026, + "ttft_ms": 5656.743910000391 + }, + { + "idx": 1, + "llama": "http://127.0.0.1:52073", + "tokens": 28942, + "hit": true, + "decision": "restore", + "restored": true, + "store_get_ms": 1988.8248469796963, + "bytes": 882366232, + "prepare_ms": 3607.2255889885128, + "prompt_n": 11, + "prompt_ms": 16.065, + "comp_wall_ms": 382.30138598009944, + "ttft_ms": 3623.290588988513 + }, + { + "idx": 2, + "llama": "http://127.0.0.1:52072", + "tokens": 28944, + "hit": true, + "decision": "restore", + "restored": true, + "store_get_ms": 1912.454207020346, + "bytes": 882366232, + "prepare_ms": 3500.9051900124177, + "prompt_n": 13, + "prompt_ms": 16.111, + "comp_wall_ms": 306.38355697738007, + "ttft_ms": 3517.0161900124176 + }, + { + "idx": 3, + "llama": "http://127.0.0.1:52073", + "tokens": 28940, + "hit": true, + "decision": "restore", + "restored": true, + "store_get_ms": 1712.3380179982632, + "bytes": 882366232, + "prepare_ms": 3297.939590003807, + "prompt_n": 9, + "prompt_ms": 87.499, + "comp_wall_ms": 220.71896900888532, + "ttft_ms": 3385.438590003807 + }, + { + "idx": 4, + "llama": "http://127.0.0.1:52072", + "tokens": 28942, + "hit": true, + "decision": "restore", + "restored": true, + "store_get_ms": 1611.3499579951167, + "bytes": 882366232, + "prepare_ms": 3589.8023130139336, + "prompt_n": 11, + "prompt_ms": 94.8, + "comp_wall_ms": 3008.9015809935518, + "ttft_ms": 3684.602313013934 + }, + { + "idx": 5, + "llama": "http://127.0.0.1:52073", + "tokens": 28943, + "hit": true, + "decision": "restore", + "restored": true, + "store_get_ms": 2498.859048995655, + "bytes": 882366232, + "prepare_ms": 3919.863100978546, + "prompt_n": 12, + "prompt_ms": 14.367, + "comp_wall_ms": 287.0978229911998, + "ttft_ms": 3934.230100978546 + } + ] +} \ No newline at end of file diff --git a/mooncake-integration/ollama/docs/figures/arbiter.png b/mooncake-integration/ollama/docs/figures/arbiter.png new file mode 100644 index 00000000..5b49a6cd Binary files /dev/null and b/mooncake-integration/ollama/docs/figures/arbiter.png differ diff --git a/mooncake-integration/ollama/docs/figures/architecture.dot b/mooncake-integration/ollama/docs/figures/architecture.dot new file mode 100644 index 00000000..c43bf82d --- /dev/null +++ b/mooncake-integration/ollama/docs/figures/architecture.dot @@ -0,0 +1,36 @@ +// Architecture of the Ollama x Mooncake KV-cache bus. +// Render: dot -Tpng architecture.dot -o ../docs/figures/architecture.png +digraph mooncake_ollama { + rankdir=LR; + bgcolor="white"; + node [fontname="Helvetica", fontsize=11, shape=box, style="rounded,filled", color="#333333"]; + edge [fontname="Helvetica", fontsize=9, color="#555555"]; + graph [fontname="Helvetica", labelloc=t, fontsize=14, + label="Ollama x Mooncake: global KV-cache reuse for local agent swarms"]; + + agents [label="Agent swarm\n(coding agents / patched Ollama)", fillcolor="#eef3fb"]; + + subgraph cluster_bridge { + label="ollama-mooncake-bridge (Go sidecar)"; + style="rounded,filled"; fillcolor="#f4f0fb"; color="#7a5fb8"; fontcolor="#4b3b78"; + key [label="cache key\n(GGUF + chained block hash)", fillcolor="#ffffff"]; + radix [label="radix index\n(cross-process KV prefixes)", fillcolor="#ffffff"]; + arbiter [label="cost arbiter\n(restore vs recompute)", fillcolor="#ffffff"]; + orch [label="orchestrator\nLookup / Prepare / Commit", fillcolor="#ffffff"]; + } + + llama [label="llama.cpp server\n(KV in GPU, per slot)\nStage 1: /slots save|restore\nStage 2: cgo _ext ON_DEVICE", fillcolor="#eafaf0"]; + + proxy [label="mooncake-store-proxy (Python)\n1 warm MooncakeDistributedStore\nstriped batch_put_from / get_into\npre-registered RDMA staging", fillcolor="#fdf3e8"]; + + store [label="Mooncake Store + Transfer Engine\nDRAM pool, replicas\nRDMA / GPUDirect", fillcolor="#fbeef0"]; + + agents -> orch [label="HTTP / gRPC"]; + orch -> key [style=dashed, arrowhead=none]; + orch -> radix [style=dashed, arrowhead=none]; + orch -> arbiter [style=dashed, arrowhead=none]; + orch -> llama [label="restore / save KV slot"]; + orch -> proxy [label="gRPC (file paths)"]; + proxy -> store [label="put / get KV blobs", color="#b0413e", penwidth=2]; + llama -> store [label="GPUDirect RDMA (Stage 2)", style=dashed, color="#1a7d32"]; +} diff --git a/mooncake-integration/ollama/docs/figures/architecture.png b/mooncake-integration/ollama/docs/figures/architecture.png new file mode 100644 index 00000000..e7b7bf48 Binary files /dev/null and b/mooncake-integration/ollama/docs/figures/architecture.png differ diff --git a/mooncake-integration/ollama/docs/figures/per_agent.png b/mooncake-integration/ollama/docs/figures/per_agent.png new file mode 100644 index 00000000..2b002c9a Binary files /dev/null and b/mooncake-integration/ollama/docs/figures/per_agent.png differ diff --git a/mooncake-integration/ollama/docs/figures/recompute.png b/mooncake-integration/ollama/docs/figures/recompute.png new file mode 100644 index 00000000..021947f1 Binary files /dev/null and b/mooncake-integration/ollama/docs/figures/recompute.png differ diff --git a/mooncake-integration/ollama/docs/figures/scaling_agents.png b/mooncake-integration/ollama/docs/figures/scaling_agents.png new file mode 100644 index 00000000..0126067e Binary files /dev/null and b/mooncake-integration/ollama/docs/figures/scaling_agents.png differ diff --git a/mooncake-integration/ollama/docs/figures/stage2_path.dot b/mooncake-integration/ollama/docs/figures/stage2_path.dot new file mode 100644 index 00000000..e4886fda --- /dev/null +++ b/mooncake-integration/ollama/docs/figures/stage2_path.dot @@ -0,0 +1,28 @@ +// Stage-1 (file) vs Stage-2 (on-device) KV export paths. +// Render: dot -Tpng stage2_path.dot -o ../docs/figures/stage2_path.png +digraph stage2 { + rankdir=LR; + bgcolor="white"; + node [fontname="Helvetica", fontsize=11, shape=box, style="rounded,filled"]; + edge [fontname="Helvetica", fontsize=10]; + graph [fontname="Helvetica", labelloc=t, fontsize=14, + label="KV export: Stage-1 file path vs Stage-2 on-device (GPUDirect)"]; + + subgraph cluster_s1 { + label="Stage 1 — /slots file save (4.6x slower, measured)"; + style="rounded,filled"; fillcolor="#fbeef0"; color="#b0413e"; fontcolor="#7a2a2a"; + g1 [label="GPU KV", fillcolor="#ffffff"]; + h1 [label="host copy", fillcolor="#ffffff"]; + f1 [label="serialize + tmpfs file", fillcolor="#ffffff"]; + st1 [label="store", fillcolor="#ffffff"]; + g1 -> h1 -> f1 -> st1; + } + + subgraph cluster_s2 { + label="Stage 2 — ON_DEVICE (host handle only 0.1-0.2 MiB)"; + style="rounded,filled"; fillcolor="#eafaf0"; color="#1a7d32"; fontcolor="#14622a"; + g2 [label="GPU KV\n(stays on device)", fillcolor="#ffffff"]; + st2 [label="store", fillcolor="#ffffff"]; + g2 -> st2 [label="GPUDirect RDMA\n(zero host copy)", color="#1a7d32", penwidth=2, fontcolor="#14622a"]; + } +} diff --git a/mooncake-integration/ollama/docs/figures/stage2_path.png b/mooncake-integration/ollama/docs/figures/stage2_path.png new file mode 100644 index 00000000..21e39a1d Binary files /dev/null and b/mooncake-integration/ollama/docs/figures/stage2_path.png differ diff --git a/mooncake-integration/ollama/docs/figures/throughput.png b/mooncake-integration/ollama/docs/figures/throughput.png new file mode 100644 index 00000000..5d9d6aed Binary files /dev/null and b/mooncake-integration/ollama/docs/figures/throughput.png differ diff --git a/mooncake-integration/ollama/docs/figures/ttft.png b/mooncake-integration/ollama/docs/figures/ttft.png new file mode 100644 index 00000000..c98710a7 Binary files /dev/null and b/mooncake-integration/ollama/docs/figures/ttft.png differ diff --git a/mooncake-integration/ollama/docs/matrix_results.json b/mooncake-integration/ollama/docs/matrix_results.json new file mode 100644 index 00000000..b8407321 --- /dev/null +++ b/mooncake-integration/ollama/docs/matrix_results.json @@ -0,0 +1,83 @@ +{ + "cells": [ + { + "name": "A1_1p5b_8k_5ag", + "backend": "rdma", + "agents": 5, + "ttft_off": 883.4, + "ttft_on": 917.9, + "ttft_red_pct": -3.9, + "reuser_ttft_on": 947.5, + "thru_x": 1.0, + "recompute_off_pct": 100.0, + "recompute_on_pct": 40.1, + "hit_agents": 4 + }, + { + "name": "A2_7b_30k_6ag", + "backend": "rdma", + "agents": 6, + "ttft_off": 5868.0, + "ttft_on": 3897.6, + "ttft_red_pct": 33.6, + "reuser_ttft_on": 3520.5, + "thru_x": 1.01, + "recompute_off_pct": 100.0, + "recompute_on_pct": 16.7, + "hit_agents": 5 + }, + { + "name": "A3_7b_16k_8ag", + "backend": "rdma", + "agents": 8, + "ttft_off": 3184.2, + "ttft_on": 2120.2, + "ttft_red_pct": 33.4, + "reuser_ttft_on": 1981.9, + "thru_x": 1.25, + "recompute_off_pct": 100.0, + "recompute_on_pct": 12.6, + "hit_agents": 7 + } + ], + "raw": { + "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" + } + ] + } +} \ No newline at end of file diff --git a/mooncake-integration/ollama/docs/stage2_kvbench.txt b/mooncake-integration/ollama/docs/stage2_kvbench.txt new file mode 100644 index 00000000..44b5ce9d --- /dev/null +++ b/mooncake-integration/ollama/docs/stage2_kvbench.txt @@ -0,0 +1,15 @@ +==== omb_kvbench: /data2/wangzhe/gitlink/mooncake1/.omb-state/models/qwen2.5-coder-7b-instruct-q4_k_m.gguf ==== +prompt tokens : 16000 (prefill 3899.4 ms, 4103 tok/s) +KV state size (host) : 875.2 MiB (57356 bytes/token) +KV state size (ondev) : 0.2 MiB + +--- export paths (the Stage-1 vs Stage-2 comparison) --- +(A) /slots file save : 2501.8 ms (0.37 GB/s) GPU->host->serialize->tmpfs +(B) host get_data_ext : 587.8 ms (1.56 GB/s) GPU->host (one copy) +(C) ON_DEVICE export : (stays on device; hand the device buffer to + Mooncake TE registerLocalMemory for GPUDirect + RDMA -- zero host copy, the Stage-2 target) +import set_data_ext : 599.7 ms + +file-save overhead vs raw host copy : 4.26x slower +KV round-trip correctness (seq0==seq1 argmax): PASS diff --git a/mooncake-integration/ollama/ollama-patches/0001-add-mooncake-cache-policy-options.patch b/mooncake-integration/ollama/ollama-patches/0001-add-mooncake-cache-policy-options.patch new file mode 100644 index 00000000..516890aa --- /dev/null +++ b/mooncake-integration/ollama/ollama-patches/0001-add-mooncake-cache-policy-options.patch @@ -0,0 +1,212 @@ +# Patch 0001: add the Mooncake KVCache Bus sidecar client + options.mooncake.* parsing +# Apply from the ollama repo root with: git apply 0001-add-mooncake-cache-policy-options.patch +# +diff --git a/server/mooncake_bridge.go b/server/mooncake_bridge.go +new file mode 100644 +index 0000000..fe0b217 +--- /dev/null ++++ b/server/mooncake_bridge.go +@@ -0,0 +1,204 @@ ++package server ++ ++// Mooncake KVCache Bus — reference integration for Ollama. ++// ++// This file lets Ollama participate in a global, cross-process / cross-GPU / ++// cross-node KV cache by talking to the ollama-mooncake-bridge sidecar over a ++// tiny HTTP/JSON API. ++// ++// The sidecar (a separate Go service) owns all the heavy lifting — cache-key ++// construction, longest-prefix matching, the restore-vs-recompute cost arbiter, ++// and the Mooncake Store/Transfer-Engine I/O. Ollama only needs to: ++// ++// 1. launch its bundled llama-server with --slot-save-path (see the companion ++// patch to llm/llama_server.go, gated by OLLAMA_MOONCAKE_SLOT_SAVE_PATH); ++// 2. before generation, call Prepare() so the matched prompt KV is restored ++// into the target slot (the sidecar fetches it from Mooncake); ++// 3. after generation, call Commit() so the produced KV is published. ++// ++// Wiring (env): ++// OLLAMA_MOONCAKE_BRIDGE_URL=http://127.0.0.1:52052 # sidecar HTTP gateway ++// OLLAMA_MOONCAKE_SLOT_SAVE_PATH=/dev/shm/ollama-mc # shared slot dir ++// ++// See ollama-patches/README.md for the exact call sites in routes.go. ++ ++import ( ++ "bytes" ++ "encoding/json" ++ "net/http" ++ "os" ++ "strconv" ++ "time" ++) ++ ++// MooncakeOptions mirror the user-visible options.mooncake.* request extension. ++type MooncakeOptions struct { ++ Enable bool ++ Namespace string ++ Read bool ++ Write bool ++ BlockSize int ++ Transport string // tcp | rdma | auto ++ ReplicaNum uint32 ++ SoftPin bool ++} ++ ++// MooncakeOptionsFromMap parses options.mooncake.* out of a request's raw ++// options map (api.Options is decoded into map[string]any for unknown keys). ++func MooncakeOptionsFromMap(opts map[string]any) (MooncakeOptions, bool) { ++ raw, ok := opts["mooncake"].(map[string]any) ++ if !ok { ++ return MooncakeOptions{}, false ++ } ++ o := MooncakeOptions{Read: true, Write: true, BlockSize: 256, Transport: "auto", ReplicaNum: 1} ++ getBool := func(k string, def bool) bool { ++ if v, ok := raw[k].(bool); ok { ++ return v ++ } ++ return def ++ } ++ getStr := func(k, def string) string { ++ if v, ok := raw[k].(string); ok { ++ return v ++ } ++ return def ++ } ++ getInt := func(k string, def int) int { ++ switch v := raw[k].(type) { ++ case float64: ++ return int(v) ++ case int: ++ return v ++ } ++ return def ++ } ++ o.Enable = getBool("enable", false) ++ o.Namespace = getStr("namespace", "") ++ o.Read = getBool("read", true) ++ o.Write = getBool("write", true) ++ o.BlockSize = getInt("block_size", 256) ++ o.Transport = getStr("transport", "auto") ++ o.ReplicaNum = uint32(getInt("replica_num", 1)) ++ o.SoftPin = getBool("soft_pin", false) ++ return o, o.Enable ++} ++ ++// MooncakeBridge is a minimal HTTP client to the sidecar's /v1 gateway. ++type MooncakeBridge struct { ++ base string ++ hc *http.Client ++} ++ ++// NewMooncakeBridge returns a bridge client, or nil if integration is disabled ++// (OLLAMA_MOONCAKE_BRIDGE_URL unset). ++func NewMooncakeBridge() *MooncakeBridge { ++ base := os.Getenv("OLLAMA_MOONCAKE_BRIDGE_URL") ++ if base == "" { ++ return nil ++ } ++ return &MooncakeBridge{base: base, hc: &http.Client{Timeout: 10 * time.Minute}} ++} ++ ++type mcFingerprint struct { ++ ModelPath string `json:"model_path"` ++ KVType string `json:"kv_type"` ++ BlockSize int `json:"block_size"` ++} ++type mcPolicy struct { ++ Enable bool `json:"enable"` ++ Namespace string `json:"namespace"` ++ Read bool `json:"read"` ++ Write bool `json:"write"` ++ BlockSize int `json:"block_size"` ++ Transport string `json:"transport"` ++ ReplicaNum uint32 `json:"replica_num"` ++ SoftPin bool `json:"soft_pin"` ++} ++type mcTarget struct { ++ BaseURL string `json:"base_url"` ++ Slot int `json:"slot_id"` ++} ++type mcRequest struct { ++ FP mcFingerprint `json:"fp"` ++ Policy mcPolicy `json:"policy"` ++ Tokens []int32 `json:"tokens"` ++ Target *mcTarget `json:"target"` ++ PrefillN int `json:"prefill_n,omitempty"` ++ PrefillMs float64 `json:"prefill_ms,omitempty"` ++} ++ ++// PrepareReply / CommitReply are the subset of the sidecar reply we consume. ++type PrepareReply struct { ++ Hit bool `json:"hit"` ++ Decision string `json:"decision"` ++ Restored bool `json:"restored"` ++ RestoredTokens int `json:"restored_tokens"` ++ MatchedTokens int `json:"matched_tokens"` ++ StoreGetMs float64 `json:"store_get_ms"` ++ Reason string `json:"reason"` ++ Error string `json:"error"` ++} ++type CommitReply struct { ++ OK bool `json:"ok"` ++ Stored bool `json:"stored"` ++ Key string `json:"key"` ++ Error string `json:"error"` ++} ++ ++func (b *MooncakeBridge) post(path string, req mcRequest, out any) error { ++ var buf bytes.Buffer ++ if err := json.NewEncoder(&buf).Encode(req); err != nil { ++ return err ++ } ++ resp, err := b.hc.Post(b.base+path, "application/json", &buf) ++ if err != nil { ++ return err ++ } ++ defer resp.Body.Close() ++ return json.NewDecoder(resp.Body).Decode(out) ++} ++ ++func (o MooncakeOptions) policy() mcPolicy { ++ return mcPolicy{Enable: o.Enable, Namespace: o.Namespace, Read: o.Read, Write: o.Write, ++ BlockSize: o.BlockSize, Transport: o.Transport, ReplicaNum: o.ReplicaNum, SoftPin: o.SoftPin} ++} ++ ++func fp(modelPath, kvType string, blockSize int) mcFingerprint { ++ if kvType == "" { ++ kvType = "f16" ++ } ++ return mcFingerprint{ModelPath: modelPath, KVType: kvType, BlockSize: blockSize} ++} ++ ++// Prepare asks the sidecar to restore the longest matching KV prefix into the ++// given llama-server slot before generation. Returns whether KV was restored. ++func (b *MooncakeBridge) Prepare(o MooncakeOptions, modelPath, kvType string, tokens []int32, llamaURL string, slot int) (*PrepareReply, error) { ++ var out PrepareReply ++ err := b.post("/v1/prepare", mcRequest{ ++ FP: fp(modelPath, kvType, o.BlockSize), Policy: o.policy(), Tokens: tokens, ++ Target: &mcTarget{BaseURL: llamaURL, Slot: slot}, ++ }, &out) ++ return &out, err ++} ++ ++// Commit publishes the slot's KV to the Mooncake Store after generation and ++// feeds the arbiter the observed prefill rate so it self-calibrates. ++func (b *MooncakeBridge) Commit(o MooncakeOptions, modelPath, kvType string, tokens []int32, llamaURL string, slot, prefillN int, prefillMs float64) (*CommitReply, error) { ++ var out CommitReply ++ err := b.post("/v1/commit", mcRequest{ ++ FP: fp(modelPath, kvType, o.BlockSize), Policy: o.policy(), Tokens: tokens, ++ Target: &mcTarget{BaseURL: llamaURL, Slot: slot}, PrefillN: prefillN, PrefillMs: prefillMs, ++ }, &out) ++ return &out, err ++} ++ ++// slotPortFromURL is a tiny helper for logging/debug. ++func slotPortFromURL(u string) string { ++ if i := bytes.LastIndexByte([]byte(u), ':'); i >= 0 { ++ if _, err := strconv.Atoi(u[i+1:]); err == nil { ++ return u[i+1:] ++ } ++ } ++ return "" ++} diff --git a/mooncake-integration/ollama/ollama-patches/0002-inject-slot-save-path-and-bridge-hook.patch b/mooncake-integration/ollama/ollama-patches/0002-inject-slot-save-path-and-bridge-hook.patch new file mode 100644 index 00000000..86df583e --- /dev/null +++ b/mooncake-integration/ollama/ollama-patches/0002-inject-slot-save-path-and-bridge-hook.patch @@ -0,0 +1,38 @@ +# Patch 0002: inject --slot-save-path so llama.cpp exposes /slots save|restore +# Apply from the ollama repo root with: git apply 0002-inject-slot-save-path-and-bridge-hook.patch +# +diff --git a/llm/llama_server.go b/llm/llama_server.go +index 8570135..98b14f7 100644 +--- a/llm/llama_server.go ++++ b/llm/llama_server.go +@@ -312,6 +312,18 @@ func FindLlamaServer() (string, error) { + } + + // startLlamaServer spawns the upstream llama-server process with appropriate CLI flags. ++// appendMooncakeArgs enables the llama.cpp /slots save|restore endpoints when ++// the sidecar integration is configured via OLLAMA_MOONCAKE_SLOT_SAVE_PATH ++// (a directory shared with the ollama-mooncake-bridge sidecar and store proxy). ++// With this set, the bundled llama-server can persist and restore a slot's ++// prompt KV, which the sidecar maps onto the global Mooncake Store. ++func appendMooncakeArgs(params []string) []string { ++ if p := os.Getenv("OLLAMA_MOONCAKE_SLOT_SAVE_PATH"); p != "" { ++ params = append(params, "--slot-save-path", p) ++ } ++ return params ++} ++ + func startLlamaServer(launch llamaServerLaunchConfig, out io.Writer) (cmd *exec.Cmd, port int, err error) { + exe, err := FindLlamaServer() + if err != nil { +@@ -345,6 +357,11 @@ func startLlamaServer(launch llamaServerLaunchConfig, out io.Writer) (cmd *exec. + params = appendLlamaServerLogArgs(params) + params = appendJinjaArgs(params, launch.config) + ++ // Mooncake KVCache Bus: expose llama.cpp's /slots save|restore endpoints so ++ // the ollama-mooncake-bridge sidecar can ship this slot's prompt KV to/from ++ // the global Mooncake Store and share it across processes/GPUs/nodes. ++ params = appendMooncakeArgs(params) ++ + params = appendMMProjArgs(params, launch) + params = appendMTPDraftArgs(params, launch.config, launch.opts) + diff --git a/mooncake-integration/ollama/ollama-patches/README.md b/mooncake-integration/ollama/ollama-patches/README.md new file mode 100644 index 00000000..92652ed2 --- /dev/null +++ b/mooncake-integration/ollama/ollama-patches/README.md @@ -0,0 +1,94 @@ +# Ollama × Mooncake — integration patches + +These patches turn upstream **Ollama** into a participant of the global Mooncake +KVCache Bus, so that prompt KV computed by one Ollama/llama.cpp process can be +**reused by other processes, GPUs and nodes** through the Mooncake Store. They +are intentionally tiny and additive — the heavy lifting lives in the standalone +`bridge/` sidecar, not in Ollama. + +Verified against `ollama/ollama` `main` @ `1abd56b` (clone date 2026-06-11). + +## What each patch does + +| Patch | File(s) | Change | +|-------|---------|--------| +| `0001-add-mooncake-cache-policy-options.patch` | `server/mooncake_bridge.go` (new) | A self-contained HTTP client for the sidecar's `/v1/prepare` and `/v1/commit` endpoints, plus parsing of the user-visible `options.mooncake.*` request extension (`enable`, `namespace`, `read`, `write`, `block_size`, `transport`, `replica_num`, `soft_pin`). No upstream symbols are modified. | +| `0002-inject-slot-save-path-and-bridge-hook.patch` | `llm/llama_server.go` | Injects `--slot-save-path ` into the bundled `llama-server` command line when `OLLAMA_MOONCAKE_SLOT_SAVE_PATH` is set, exposing llama.cpp's `/slots?action=save|restore` endpoints that the sidecar drives. Gated by an env var, so it is a no-op unless explicitly enabled. | + +Both are additive and gated, so an un-configured Ollama behaves exactly as before. + +## Apply + +```bash +git clone https://github.com/ollama/ollama && cd ollama +git apply /path/to/ollama-patches/0001-add-mooncake-cache-policy-options.patch +git apply /path/to/ollama-patches/0002-inject-slot-save-path-and-bridge-hook.patch +# Ollama main currently requires Go >= 1.26 +go build ./... +``` + +## Configure (env) + +```bash +export OLLAMA_MOONCAKE_BRIDGE_URL=http://127.0.0.1:52052 # sidecar HTTP gateway +export OLLAMA_MOONCAKE_SLOT_SAVE_PATH=/dev/shm/ollama-mc # shared (tmpfs) slot dir +``` + +## The one remaining wiring step (reference) + +Ollama's runner hides the underlying llama.cpp slot id behind its scheduler, so +the *final* hook into the request path is left as a clearly-marked reference +rather than a forced edit to the large `server/routes.go` generate/chat handler. +Where Ollama dispatches a request to a runner that owns a llama-server at +`llamaURL` on slot `slot`, add: + +```go +// --- before generation --- +if mb := NewMooncakeBridge(); mb != nil { + if mo, on := MooncakeOptionsFromMap(rawOptions); on { + // tokens: the tokenized prompt; kvType: opts.KvCacheType (e.g. "q8_0") + if pr, err := mb.Prepare(mo, modelPath, kvType, tokens, llamaURL, slot); err == nil && pr.Restored { + slog.Info("mooncake: restored KV prefix", + "matched_tokens", pr.MatchedTokens, "store_get_ms", pr.StoreGetMs, "reason", pr.Reason) + } + } +} + +// ... run generation as usual; capture the prefill timing from the response ... + +// --- after generation --- +if mb := NewMooncakeBridge(); mb != nil { + if mo, on := MooncakeOptionsFromMap(rawOptions); on && mo.Write { + _, _ = mb.Commit(mo, modelPath, kvType, tokens, llamaURL, slot, prefillN, prefillMs) + } +} +``` + +This mirrors exactly how the bench harness (`benchmarks/agent_swarm.py`) drives the +sidecar today, so the end-to-end behaviour is already validated outside Ollama; +the snippet above is the minimal glue to move that call site *inside* Ollama. + +## Request example (`POST /api/chat`) + +```jsonc +{ + "model": "qwen2.5-coder:7b", + "messages": [ /* ... long shared repo context + this agent's task ... */ ], + "options": { + "mooncake": { + "enable": true, + "namespace": "repo:my-org/my-repo@main", + "read": true, "write": true, + "block_size": 256, + "transport": "auto" + } + } +} +``` + +## Upstreaming + +`0002` is a clean candidate for an upstream PR (a gated, opt-in flag pass-through). +`0001` + the wiring snippet are proposed as an RFC: *"Pluggable global KV cache +for Ollama via a sidecar"*, aligned with community issue +[ollama/ollama#14872 — Swarm Memory](https://github.com/ollama/ollama/issues/14872). diff --git a/mooncake-integration/ollama/proto/storeproxy.proto b/mooncake-integration/ollama/proto/storeproxy.proto new file mode 100644 index 00000000..764f9948 --- /dev/null +++ b/mooncake-integration/ollama/proto/storeproxy.proto @@ -0,0 +1,97 @@ +syntax = "proto3"; + +// StoreProxy is the gRPC contract between the Go sidecar (client) and the +// Python Mooncake store proxy (server). The Python side owns the single, +// long-lived, warm mooncake.store.MooncakeDistributedStore handle. +// +// Design note: KV snapshots can be multiple GiB. We therefore pass *file paths* +// (in the shared --slot-save-path directory) instead of streaming bytes through +// Go. The proxy reads/writes those files directly to/from the store, so a +// multi-GiB KV blob is copied at most once (file <-> store), never through the +// Go process. Small values may still be sent inline via PutBytes/GetBytes. +package storeproxy.v1; + +option go_package = "github.com/mooncake-ai/ollama-mooncake-bridge/internal/storeproxy/pb;storeproxypb"; + +service StoreProxy { + // Liveness + which backend/protocol/master is active. + rpc Health(HealthRequest) returns (HealthReply); + // Batched existence check (maps to mooncake batch_is_exist). + rpc Exists(ExistsRequest) returns (ExistsReply); + // Store a KV snapshot file under `key`. + rpc PutFile(PutFileRequest) returns (PutFileReply); + // Materialize `key` into a local file at `path`. + rpc GetFile(GetFileRequest) returns (GetFileReply); + // Inline small-value variants (metadata, manifests). + rpc PutBytes(PutBytesRequest) returns (PutBytesReply); + rpc GetBytes(GetBytesRequest) returns (GetBytesReply); + rpc Remove(RemoveRequest) returns (RemoveReply); + rpc Stats(StatsRequest) returns (StatsReply); +} + +message HealthRequest {} +message HealthReply { + bool ok = 1; + string backend = 2; // "mooncake" | "local" + string protocol = 3; // "tcp" | "rdma" + string master = 4; // master_server_addr + string device = 5; // rdma device(s) + string detail = 6; +} + +message ExistsRequest { repeated string keys = 1; } +// present[i]: 1 = exists, 0 = absent, -1 = error (mirrors batch_is_exist). +message ExistsReply { repeated int32 present = 1; } + +message PutFileRequest { + string key = 1; + string path = 2; // local file to read + uint32 replica_num = 3; // ReplicateConfig.replica_num (>=1) + bool soft_pin = 4; // ReplicateConfig.with_soft_pin (hot prefixes) + bool skip_if_exists = 5; +} +message PutFileReply { + bool ok = 1; + uint64 bytes = 2; + double elapsed_ms = 3; + bool existed = 4; // true if skipped because key already present + string error = 5; +} + +message GetFileRequest { + string key = 1; + string path = 2; // local file to write +} +message GetFileReply { + bool ok = 1; + uint64 bytes = 2; + double elapsed_ms = 3; + bool found = 4; + string error = 5; +} + +message PutBytesRequest { + string key = 1; + bytes value = 2; + uint32 replica_num = 3; + bool soft_pin = 4; +} +message PutBytesReply { bool ok = 1; uint64 bytes = 2; double elapsed_ms = 3; string error = 4; } + +message GetBytesRequest { string key = 1; } +message GetBytesReply { bool ok = 1; bytes value = 2; double elapsed_ms = 3; bool found = 4; string error = 5; } + +message RemoveRequest { string key = 1; bool force = 2; } +message RemoveReply { bool ok = 1; string error = 2; } + +message StatsRequest {} +message StatsReply { + uint64 put_ops = 1; + uint64 get_ops = 2; + uint64 exists_ops = 3; + uint64 put_bytes = 4; + uint64 get_bytes = 5; + double put_ms_total = 6; + double get_ms_total = 7; + string backend = 8; +} diff --git a/mooncake-integration/ollama/scripts/bridged_start.sh b/mooncake-integration/ollama/scripts/bridged_start.sh new file mode 100644 index 00000000..b2766517 --- /dev/null +++ b/mooncake-integration/ollama/scripts/bridged_start.sh @@ -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 diff --git a/mooncake-integration/ollama/scripts/demo.sh b/mooncake-integration/ollama/scripts/demo.sh new file mode 100644 index 00000000..53e5cb29 --- /dev/null +++ b/mooncake-integration/ollama/scripts/demo.sh @@ -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" diff --git a/mooncake-integration/ollama/scripts/env.sh b/mooncake-integration/ollama/scripts/env.sh new file mode 100644 index 00000000..0cbf0f82 --- /dev/null +++ b/mooncake-integration/ollama/scripts/env.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 "/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 diff --git a/mooncake-integration/ollama/scripts/gen_protos.sh b/mooncake-integration/ollama/scripts/gen_protos.sh new file mode 100644 index 00000000..9df9d78a --- /dev/null +++ b/mooncake-integration/ollama/scripts/gen_protos.sh @@ -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 diff --git a/mooncake-integration/ollama/scripts/llama_start.sh b/mooncake-integration/ollama/scripts/llama_start.sh new file mode 100644 index 00000000..fb24aef2 --- /dev/null +++ b/mooncake-integration/ollama/scripts/llama_start.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +# Start one llama.cpp server instance ("agent worker"). +# Usage: llama_start.sh [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 diff --git a/mooncake-integration/ollama/scripts/master_start.sh b/mooncake-integration/ollama/scripts/master_start.sh new file mode 100644 index 00000000..107ef502 --- /dev/null +++ b/mooncake-integration/ollama/scripts/master_start.sh @@ -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 diff --git a/mooncake-integration/ollama/scripts/master_stop.sh b/mooncake-integration/ollama/scripts/master_stop.sh new file mode 100644 index 00000000..e74f1b72 --- /dev/null +++ b/mooncake-integration/ollama/scripts/master_stop.sh @@ -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)" diff --git a/mooncake-integration/ollama/scripts/proxy_start.sh b/mooncake-integration/ollama/scripts/proxy_start.sh new file mode 100644 index 00000000..60cec550 --- /dev/null +++ b/mooncake-integration/ollama/scripts/proxy_start.sh @@ -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 diff --git a/mooncake-integration/ollama/scripts/setup_go.sh b/mooncake-integration/ollama/scripts/setup_go.sh new file mode 100644 index 00000000..7b1cec0f --- /dev/null +++ b/mooncake-integration/ollama/scripts/setup_go.sh @@ -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 diff --git a/mooncake-integration/ollama/scripts/setup_llama.sh b/mooncake-integration/ollama/scripts/setup_llama.sh new file mode 100644 index 00000000..66a1faf2 --- /dev/null +++ b/mooncake-integration/ollama/scripts/setup_llama.sh @@ -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 diff --git a/mooncake-integration/ollama/scripts/setup_models.sh b/mooncake-integration/ollama/scripts/setup_models.sh new file mode 100644 index 00000000..394c1eff --- /dev/null +++ b/mooncake-integration/ollama/scripts/setup_models.sh @@ -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" diff --git a/mooncake-integration/ollama/scripts/setup_py.sh b/mooncake-integration/ollama/scripts/setup_py.sh new file mode 100644 index 00000000..b0a83770 --- /dev/null +++ b/mooncake-integration/ollama/scripts/setup_py.sh @@ -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" diff --git a/mooncake-integration/ollama/scripts/stack_down.sh b/mooncake-integration/ollama/scripts/stack_down.sh new file mode 100644 index 00000000..60cf3d5a --- /dev/null +++ b/mooncake-integration/ollama/scripts/stack_down.sh @@ -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)" diff --git a/mooncake-integration/ollama/scripts/stack_up.sh b/mooncake-integration/ollama/scripts/stack_up.sh new file mode 100644 index 00000000..be5a2711 --- /dev/null +++ b/mooncake-integration/ollama/scripts/stack_up.sh @@ -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 diff --git a/mooncake-integration/ollama/store-proxy/gen/__init__.py b/mooncake-integration/ollama/store-proxy/gen/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/mooncake-integration/ollama/store-proxy/gen/storeproxy_pb2.py b/mooncake-integration/ollama/store-proxy/gen/storeproxy_pb2.py new file mode 100644 index 00000000..11739fc0 --- /dev/null +++ b/mooncake-integration/ollama/store-proxy/gen/storeproxy_pb2.py @@ -0,0 +1,59 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: storeproxy.proto +# Protobuf Python Version: 4.25.1 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x10storeproxy.proto\x12\rstoreproxy.v1\"\x0f\n\rHealthRequest\"l\n\x0bHealthReply\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\x0f\n\x07\x62\x61\x63kend\x18\x02 \x01(\t\x12\x10\n\x08protocol\x18\x03 \x01(\t\x12\x0e\n\x06master\x18\x04 \x01(\t\x12\x0e\n\x06\x64\x65vice\x18\x05 \x01(\t\x12\x0e\n\x06\x64\x65tail\x18\x06 \x01(\t\"\x1d\n\rExistsRequest\x12\x0c\n\x04keys\x18\x01 \x03(\t\"\x1e\n\x0b\x45xistsReply\x12\x0f\n\x07present\x18\x01 \x03(\x05\"j\n\x0ePutFileRequest\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x0c\n\x04path\x18\x02 \x01(\t\x12\x13\n\x0breplica_num\x18\x03 \x01(\r\x12\x10\n\x08soft_pin\x18\x04 \x01(\x08\x12\x16\n\x0eskip_if_exists\x18\x05 \x01(\x08\"]\n\x0cPutFileReply\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\r\n\x05\x62ytes\x18\x02 \x01(\x04\x12\x12\n\nelapsed_ms\x18\x03 \x01(\x01\x12\x0f\n\x07\x65xisted\x18\x04 \x01(\x08\x12\r\n\x05\x65rror\x18\x05 \x01(\t\"+\n\x0eGetFileRequest\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x0c\n\x04path\x18\x02 \x01(\t\"[\n\x0cGetFileReply\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\r\n\x05\x62ytes\x18\x02 \x01(\x04\x12\x12\n\nelapsed_ms\x18\x03 \x01(\x01\x12\r\n\x05\x66ound\x18\x04 \x01(\x08\x12\r\n\x05\x65rror\x18\x05 \x01(\t\"T\n\x0fPutBytesRequest\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x0c\x12\x13\n\x0breplica_num\x18\x03 \x01(\r\x12\x10\n\x08soft_pin\x18\x04 \x01(\x08\"M\n\rPutBytesReply\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\r\n\x05\x62ytes\x18\x02 \x01(\x04\x12\x12\n\nelapsed_ms\x18\x03 \x01(\x01\x12\r\n\x05\x65rror\x18\x04 \x01(\t\"\x1e\n\x0fGetBytesRequest\x12\x0b\n\x03key\x18\x01 \x01(\t\"\\\n\rGetBytesReply\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\r\n\x05value\x18\x02 \x01(\x0c\x12\x12\n\nelapsed_ms\x18\x03 \x01(\x01\x12\r\n\x05\x66ound\x18\x04 \x01(\x08\x12\r\n\x05\x65rror\x18\x05 \x01(\t\"+\n\rRemoveRequest\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05\x66orce\x18\x02 \x01(\x08\"(\n\x0bRemoveReply\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\r\n\x05\x65rror\x18\x02 \x01(\t\"\x0e\n\x0cStatsRequest\"\xa5\x01\n\nStatsReply\x12\x0f\n\x07put_ops\x18\x01 \x01(\x04\x12\x0f\n\x07get_ops\x18\x02 \x01(\x04\x12\x12\n\nexists_ops\x18\x03 \x01(\x04\x12\x11\n\tput_bytes\x18\x04 \x01(\x04\x12\x11\n\tget_bytes\x18\x05 \x01(\x04\x12\x14\n\x0cput_ms_total\x18\x06 \x01(\x01\x12\x14\n\x0cget_ms_total\x18\x07 \x01(\x01\x12\x0f\n\x07\x62\x61\x63kend\x18\x08 \x01(\t2\xbb\x04\n\nStoreProxy\x12\x42\n\x06Health\x12\x1c.storeproxy.v1.HealthRequest\x1a\x1a.storeproxy.v1.HealthReply\x12\x42\n\x06\x45xists\x12\x1c.storeproxy.v1.ExistsRequest\x1a\x1a.storeproxy.v1.ExistsReply\x12\x45\n\x07PutFile\x12\x1d.storeproxy.v1.PutFileRequest\x1a\x1b.storeproxy.v1.PutFileReply\x12\x45\n\x07GetFile\x12\x1d.storeproxy.v1.GetFileRequest\x1a\x1b.storeproxy.v1.GetFileReply\x12H\n\x08PutBytes\x12\x1e.storeproxy.v1.PutBytesRequest\x1a\x1c.storeproxy.v1.PutBytesReply\x12H\n\x08GetBytes\x12\x1e.storeproxy.v1.GetBytesRequest\x1a\x1c.storeproxy.v1.GetBytesReply\x12\x42\n\x06Remove\x12\x1c.storeproxy.v1.RemoveRequest\x1a\x1a.storeproxy.v1.RemoveReply\x12?\n\x05Stats\x12\x1b.storeproxy.v1.StatsRequest\x1a\x19.storeproxy.v1.StatsReplyBSZQgithub.com/mooncake-ai/ollama-mooncake-bridge/internal/storeproxy/pb;storeproxypbb\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'storeproxy_pb2', _globals) +if _descriptor._USE_C_DESCRIPTORS == False: + _globals['DESCRIPTOR']._options = None + _globals['DESCRIPTOR']._serialized_options = b'ZQgithub.com/mooncake-ai/ollama-mooncake-bridge/internal/storeproxy/pb;storeproxypb' + _globals['_HEALTHREQUEST']._serialized_start=35 + _globals['_HEALTHREQUEST']._serialized_end=50 + _globals['_HEALTHREPLY']._serialized_start=52 + _globals['_HEALTHREPLY']._serialized_end=160 + _globals['_EXISTSREQUEST']._serialized_start=162 + _globals['_EXISTSREQUEST']._serialized_end=191 + _globals['_EXISTSREPLY']._serialized_start=193 + _globals['_EXISTSREPLY']._serialized_end=223 + _globals['_PUTFILEREQUEST']._serialized_start=225 + _globals['_PUTFILEREQUEST']._serialized_end=331 + _globals['_PUTFILEREPLY']._serialized_start=333 + _globals['_PUTFILEREPLY']._serialized_end=426 + _globals['_GETFILEREQUEST']._serialized_start=428 + _globals['_GETFILEREQUEST']._serialized_end=471 + _globals['_GETFILEREPLY']._serialized_start=473 + _globals['_GETFILEREPLY']._serialized_end=564 + _globals['_PUTBYTESREQUEST']._serialized_start=566 + _globals['_PUTBYTESREQUEST']._serialized_end=650 + _globals['_PUTBYTESREPLY']._serialized_start=652 + _globals['_PUTBYTESREPLY']._serialized_end=729 + _globals['_GETBYTESREQUEST']._serialized_start=731 + _globals['_GETBYTESREQUEST']._serialized_end=761 + _globals['_GETBYTESREPLY']._serialized_start=763 + _globals['_GETBYTESREPLY']._serialized_end=855 + _globals['_REMOVEREQUEST']._serialized_start=857 + _globals['_REMOVEREQUEST']._serialized_end=900 + _globals['_REMOVEREPLY']._serialized_start=902 + _globals['_REMOVEREPLY']._serialized_end=942 + _globals['_STATSREQUEST']._serialized_start=944 + _globals['_STATSREQUEST']._serialized_end=958 + _globals['_STATSREPLY']._serialized_start=961 + _globals['_STATSREPLY']._serialized_end=1126 + _globals['_STOREPROXY']._serialized_start=1129 + _globals['_STOREPROXY']._serialized_end=1700 +# @@protoc_insertion_point(module_scope) diff --git a/mooncake-integration/ollama/store-proxy/gen/storeproxy_pb2_grpc.py b/mooncake-integration/ollama/store-proxy/gen/storeproxy_pb2_grpc.py new file mode 100644 index 00000000..0f27217b --- /dev/null +++ b/mooncake-integration/ollama/store-proxy/gen/storeproxy_pb2_grpc.py @@ -0,0 +1,302 @@ +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc + +from . import storeproxy_pb2 as storeproxy__pb2 + + +class StoreProxyStub(object): + """Missing associated documentation comment in .proto file.""" + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.Health = channel.unary_unary( + '/storeproxy.v1.StoreProxy/Health', + request_serializer=storeproxy__pb2.HealthRequest.SerializeToString, + response_deserializer=storeproxy__pb2.HealthReply.FromString, + ) + self.Exists = channel.unary_unary( + '/storeproxy.v1.StoreProxy/Exists', + request_serializer=storeproxy__pb2.ExistsRequest.SerializeToString, + response_deserializer=storeproxy__pb2.ExistsReply.FromString, + ) + self.PutFile = channel.unary_unary( + '/storeproxy.v1.StoreProxy/PutFile', + request_serializer=storeproxy__pb2.PutFileRequest.SerializeToString, + response_deserializer=storeproxy__pb2.PutFileReply.FromString, + ) + self.GetFile = channel.unary_unary( + '/storeproxy.v1.StoreProxy/GetFile', + request_serializer=storeproxy__pb2.GetFileRequest.SerializeToString, + response_deserializer=storeproxy__pb2.GetFileReply.FromString, + ) + self.PutBytes = channel.unary_unary( + '/storeproxy.v1.StoreProxy/PutBytes', + request_serializer=storeproxy__pb2.PutBytesRequest.SerializeToString, + response_deserializer=storeproxy__pb2.PutBytesReply.FromString, + ) + self.GetBytes = channel.unary_unary( + '/storeproxy.v1.StoreProxy/GetBytes', + request_serializer=storeproxy__pb2.GetBytesRequest.SerializeToString, + response_deserializer=storeproxy__pb2.GetBytesReply.FromString, + ) + self.Remove = channel.unary_unary( + '/storeproxy.v1.StoreProxy/Remove', + request_serializer=storeproxy__pb2.RemoveRequest.SerializeToString, + response_deserializer=storeproxy__pb2.RemoveReply.FromString, + ) + self.Stats = channel.unary_unary( + '/storeproxy.v1.StoreProxy/Stats', + request_serializer=storeproxy__pb2.StatsRequest.SerializeToString, + response_deserializer=storeproxy__pb2.StatsReply.FromString, + ) + + +class StoreProxyServicer(object): + """Missing associated documentation comment in .proto file.""" + + def Health(self, request, context): + """Liveness + which backend/protocol/master is active. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def Exists(self, request, context): + """Batched existence check (maps to mooncake batch_is_exist). + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def PutFile(self, request, context): + """Store a KV snapshot file under `key`. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def GetFile(self, request, context): + """Materialize `key` into a local file at `path`. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def PutBytes(self, request, context): + """Inline small-value variants (metadata, manifests). + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def GetBytes(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def Remove(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def Stats(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_StoreProxyServicer_to_server(servicer, server): + rpc_method_handlers = { + 'Health': grpc.unary_unary_rpc_method_handler( + servicer.Health, + request_deserializer=storeproxy__pb2.HealthRequest.FromString, + response_serializer=storeproxy__pb2.HealthReply.SerializeToString, + ), + 'Exists': grpc.unary_unary_rpc_method_handler( + servicer.Exists, + request_deserializer=storeproxy__pb2.ExistsRequest.FromString, + response_serializer=storeproxy__pb2.ExistsReply.SerializeToString, + ), + 'PutFile': grpc.unary_unary_rpc_method_handler( + servicer.PutFile, + request_deserializer=storeproxy__pb2.PutFileRequest.FromString, + response_serializer=storeproxy__pb2.PutFileReply.SerializeToString, + ), + 'GetFile': grpc.unary_unary_rpc_method_handler( + servicer.GetFile, + request_deserializer=storeproxy__pb2.GetFileRequest.FromString, + response_serializer=storeproxy__pb2.GetFileReply.SerializeToString, + ), + 'PutBytes': grpc.unary_unary_rpc_method_handler( + servicer.PutBytes, + request_deserializer=storeproxy__pb2.PutBytesRequest.FromString, + response_serializer=storeproxy__pb2.PutBytesReply.SerializeToString, + ), + 'GetBytes': grpc.unary_unary_rpc_method_handler( + servicer.GetBytes, + request_deserializer=storeproxy__pb2.GetBytesRequest.FromString, + response_serializer=storeproxy__pb2.GetBytesReply.SerializeToString, + ), + 'Remove': grpc.unary_unary_rpc_method_handler( + servicer.Remove, + request_deserializer=storeproxy__pb2.RemoveRequest.FromString, + response_serializer=storeproxy__pb2.RemoveReply.SerializeToString, + ), + 'Stats': grpc.unary_unary_rpc_method_handler( + servicer.Stats, + request_deserializer=storeproxy__pb2.StatsRequest.FromString, + response_serializer=storeproxy__pb2.StatsReply.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'storeproxy.v1.StoreProxy', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + + + # This class is part of an EXPERIMENTAL API. +class StoreProxy(object): + """Missing associated documentation comment in .proto file.""" + + @staticmethod + def Health(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/storeproxy.v1.StoreProxy/Health', + storeproxy__pb2.HealthRequest.SerializeToString, + storeproxy__pb2.HealthReply.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def Exists(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/storeproxy.v1.StoreProxy/Exists', + storeproxy__pb2.ExistsRequest.SerializeToString, + storeproxy__pb2.ExistsReply.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def PutFile(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/storeproxy.v1.StoreProxy/PutFile', + storeproxy__pb2.PutFileRequest.SerializeToString, + storeproxy__pb2.PutFileReply.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def GetFile(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/storeproxy.v1.StoreProxy/GetFile', + storeproxy__pb2.GetFileRequest.SerializeToString, + storeproxy__pb2.GetFileReply.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def PutBytes(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/storeproxy.v1.StoreProxy/PutBytes', + storeproxy__pb2.PutBytesRequest.SerializeToString, + storeproxy__pb2.PutBytesReply.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def GetBytes(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/storeproxy.v1.StoreProxy/GetBytes', + storeproxy__pb2.GetBytesRequest.SerializeToString, + storeproxy__pb2.GetBytesReply.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def Remove(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/storeproxy.v1.StoreProxy/Remove', + storeproxy__pb2.RemoveRequest.SerializeToString, + storeproxy__pb2.RemoveReply.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def Stats(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/storeproxy.v1.StoreProxy/Stats', + storeproxy__pb2.StatsRequest.SerializeToString, + storeproxy__pb2.StatsReply.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) diff --git a/mooncake-integration/ollama/store-proxy/store_proxy.py b/mooncake-integration/ollama/store-proxy/store_proxy.py new file mode 100644 index 00000000..95b31ab6 --- /dev/null +++ b/mooncake-integration/ollama/store-proxy/store_proxy.py @@ -0,0 +1,669 @@ +#!/usr/bin/env python3 +"""Mooncake Store Proxy +==================== + +A thin, *long-lived* gRPC service that owns one warm +``mooncake.store.MooncakeDistributedStore`` handle and exposes it to the Go +sidecar (which cannot link the Python bindings directly). + +Why a proxy instead of cgo? + * The official, battle-tested store client is the Python binding shipped in + the ``mooncake-transfer-engine`` wheel. Re-implementing it in cgo would be + fragile. The proxy lets the Go sidecar stay pure-Go while still driving the + *real* Mooncake Store (TCP or RDMA/GPUDirect). + * KV snapshots are multi-GiB. We pass **file paths**, not bytes: ``PutFile`` + mmaps the slot-save file and stores it; ``GetFile`` materializes an object + straight into the slot-save directory. The blob is therefore copied at most + once (file <-> store) and never travels through the Go process. + +Backends: + * ``mooncake`` (default): the distributed store. ``--protocol tcp|rdma|auto``. + * ``local``: a filesystem-backed object store (no master needed). Used for the + "local file" baseline and for environments without a master. + +Everything stays under the workspace; no writes to ``/``. +""" +from __future__ import annotations + +import argparse +import ctypes +import logging +import mmap +import os +import queue +import shutil +import sys +import threading +import time +from concurrent import futures +from dataclasses import dataclass, field + +import grpc + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from gen import storeproxy_pb2 as pb # noqa: E402 +from gen import storeproxy_pb2_grpc as pb_grpc # noqa: E402 + +log = logging.getLogger("store-proxy") + + +# --------------------------------------------------------------------------- # +# Backend abstraction +# --------------------------------------------------------------------------- # +@dataclass +class Stats: + put_ops: int = 0 + get_ops: int = 0 + exists_ops: int = 0 + put_bytes: int = 0 + get_bytes: int = 0 + put_ms_total: float = 0.0 + get_ms_total: float = 0.0 + _lock: threading.Lock = field(default_factory=threading.Lock, repr=False) + + def add_put(self, n: int, ms: float): + with self._lock: + self.put_ops += 1 + self.put_bytes += n + self.put_ms_total += ms + + def add_get(self, n: int, ms: float): + with self._lock: + self.get_ops += 1 + self.get_bytes += n + self.get_ms_total += ms + + def add_exists(self, n: int): + with self._lock: + self.exists_ops += n + + +class Backend: + name = "base" + protocol = "" + master = "" + device = "" + + def exists(self, keys): raise NotImplementedError + def put_file(self, key, path, replica_num, soft_pin, skip_if_exists): raise NotImplementedError + def get_file(self, key, path): raise NotImplementedError + def put_bytes(self, key, value, replica_num, soft_pin): raise NotImplementedError + def get_bytes(self, key): raise NotImplementedError + def remove(self, key, force): raise NotImplementedError + def close(self): pass + + +class LocalBackend(Backend): + """Filesystem object store. Keys are hashed to safe file names.""" + name = "local" + + def __init__(self, root: str): + import hashlib + self._hash = hashlib.sha256 + self.root = root + os.makedirs(root, exist_ok=True) + self.protocol = "file" + self.master = root + + def _p(self, key: str) -> str: + h = self._hash(key.encode()).hexdigest() + return os.path.join(self.root, h[:2], h) + + def exists(self, keys): + return [1 if os.path.exists(self._p(k)) else 0 for k in keys] + + def put_file(self, key, path, replica_num, soft_pin, skip_if_exists): + dst = self._p(key) + if skip_if_exists and os.path.exists(dst): + return True, os.path.getsize(dst), True, "" + os.makedirs(os.path.dirname(dst), exist_ok=True) + tmp = dst + ".tmp" + shutil.copyfile(path, tmp) + os.replace(tmp, dst) + return True, os.path.getsize(dst), False, "" + + def get_file(self, key, path): + src = self._p(key) + if not os.path.exists(src): + return False, 0, False, "" + os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) + tmp = path + ".tmp" + shutil.copyfile(src, tmp) + os.replace(tmp, path) + return True, os.path.getsize(path), True, "" + + def put_bytes(self, key, value, replica_num, soft_pin): + dst = self._p(key) + os.makedirs(os.path.dirname(dst), exist_ok=True) + with open(dst, "wb") as f: + f.write(value) + return True, len(value), "" + + def get_bytes(self, key): + src = self._p(key) + if not os.path.exists(src): + return False, b"", False, "" + with open(src, "rb") as f: + return True, f.read(), True, "" + + def remove(self, key, force): + try: + os.remove(self._p(key)) + except FileNotFoundError: + pass + return True, "" + + +def _ptr_of(mm) -> int: + """Raw address of an mmap region (for register_buffer / put_from / get_into).""" + return ctypes.addressof(ctypes.c_char.from_buffer(mm)) + + +class StagingPool: + """A pool of pre-registered host buffers for zero-copy RDMA transfers. + + Registering RDMA memory (pinning pages) is expensive, so doing it per + operation throttled us to ~0.7 GB/s. Instead we register a handful of large + buffers ONCE at startup and reuse them, sustaining the full ~40 GB/s the NIC + can do. Each buffer is touched so its pages are resident before registration. + """ + + def __init__(self, store, count: int, size: int): + self.store = store + self.size = size + self.q: "queue.Queue" = queue.Queue() + self.bufs = [] + for _ in range(count): + mm = mmap.mmap(-1, size) + ctypes.memset(_ptr_of(mm), 0, size) # fault in pages before pinning + ptr = _ptr_of(mm) + r = store.register_buffer(ptr, size) + if r != 0: + raise RuntimeError(f"register_buffer failed: {r}") + self.bufs.append((mm, ptr)) + self.q.put((mm, ptr)) + log.info("staging pool: %d x %d MiB registered", count, size >> 20) + + def acquire(self, timeout=None): + return self.q.get(timeout=timeout) + + def release(self, item): + self.q.put(item) + + def close(self): + for mm, ptr in self.bufs: + try: + self.store.unregister_buffer(ptr) + except Exception: # noqa: BLE001 + pass + mm.close() + + +class MooncakeBackend(Backend): + """Wraps the real distributed store. One warm client, shared across RPCs. + + KV blobs are stored *striped*: the data is split into <=stripe-byte chunks + transferred in parallel with batch_put_from / batch_get_into. We measured + single-object RDMA transfer collapsing from ~40 GB/s (<=128 MiB) to ~2 GB/s + (>=1 GiB), while striped batches sustain full bandwidth. A tiny JSON manifest + is stored under `key`; chunk i lives under f"{key}#{i}". This also makes the + physical layout independent of the semantic prefix-block size. + """ + name = "mooncake" + _MANIFEST_MAGIC = b"OMBKVv1\n" + + def __init__(self, args): + from mooncake.store import MooncakeDistributedStore, ReplicateConfig + self._RC = ReplicateConfig + self.protocol = args.protocol + self.master = args.master + self.device = args.device + self.stripe = max(1 << 20, args.stripe_mb << 20) + self.store = MooncakeDistributedStore() + cfg = { + "local_hostname": args.local_hostname, + "metadata_server": args.metadata, + "global_segment_size": args.global_segment_size, + "local_buffer_size": args.local_buffer_size, + "protocol": args.protocol, + "rdma_devices": args.device, + "master_server_addr": args.master, + } + log.info("store setup cfg=%s", cfg) + deadline = time.time() + args.setup_timeout + last = None + while True: + try: + ret = self.store.setup(cfg) + if ret == 0: + break + last = RuntimeError(f"setup returned {ret}") + except Exception as e: # noqa: BLE001 + last = e + if time.time() > deadline: + raise RuntimeError(f"mooncake setup failed: {last}") + log.warning("store setup retry (%s)", last) + time.sleep(1.0) + # Pre-registered staging pool for zero-copy transfers (the perf fix). + self.pool = None + if args.staging_count > 0 and args.staging_mb > 0: + try: + self.pool = StagingPool(self.store, args.staging_count, args.staging_mb << 20) + except Exception as e: # noqa: BLE001 + log.warning("staging pool disabled (%s); falling back to per-op registration", e) + log.info("mooncake store ready (stripe=%dMiB, staging=%s)", + self.stripe >> 20, "on" if self.pool else "off") + + def _rc(self, replica_num, soft_pin): + rc = self._RC() + try: + rc.replica_num = max(1, int(replica_num) or 1) + rc.with_soft_pin = bool(soft_pin) + except Exception: # noqa: BLE001 + pass + return rc + + def _chunk_keys(self, key, n): + return [f"{key}#{i}" for i in range(n)] + + def exists(self, keys): + if not keys: + return [] + return list(self.store.batch_is_exist(list(keys))) + + def _stripe_layout(self, key, size): + n = max(1, (size + self.stripe - 1) // self.stripe) + keys = self._chunk_keys(key, n) + sizes = [min(self.stripe, size - i * self.stripe) for i in range(n)] + return n, keys, sizes + + def put_file(self, key, path, replica_num, soft_pin, skip_if_exists): + if skip_if_exists and self.store.is_exist(key) == 1: + try: + sz = self._manifest_size(key) + except Exception: # noqa: BLE001 + sz = 0 + return True, sz, True, "" + size = os.path.getsize(path) + rc = self._rc(replica_num, soft_pin) + n, keys, sizes = self._stripe_layout(key, size) + + # Primary path: read the slot file into a pre-registered staging buffer + # (registration amortized once at startup) and RDMA from it. readinto on + # a tmpfs file is a fast page-cache copy that releases the GIL. + if self.pool is not None and size <= self.pool.size: + mm, base = self.pool.acquire() + try: + with open(path, "rb") as f: + mv = memoryview(mm) + off = 0 + while off < size: + r = f.readinto(mv[off:size]) + if not r: + break + off += r + ptrs = [base + i * self.stripe for i in range(n)] + rets = self.store.batch_put_from(keys, ptrs, sizes, rc) + if any(r != 0 for r in rets): + return False, 0, False, f"batch_put_from rets={rets}" + finally: + self.pool.release((mm, base)) + else: + ok, err = self._put_unregistered(path, size, n, keys, sizes, rc) + if not ok: + return False, 0, False, err + man = self._MANIFEST_MAGIC + f'{{"n":{n},"size":{size},"chunk":{self.stripe}}}'.encode() + if self.store.put(key, man, rc) != 0: + return False, 0, False, "manifest put failed" + return True, size, False, "" + + def _put_unregistered(self, path, size, n, keys, sizes, rc): + if size == 0: + return True, "" + with open(path, "r+b") as f: + mm = mmap.mmap(f.fileno(), size) + try: + base = _ptr_of(mm) + self.store.register_buffer(base, size) + try: + ptrs = [base + i * self.stripe for i in range(n)] + rets = self.store.batch_put_from(keys, ptrs, sizes, rc) + if any(r != 0 for r in rets): + return False, f"batch_put_from rets={rets}" + finally: + self.store.unregister_buffer(base) + finally: + mm.close() + return True, "" + + def _read_manifest(self, key): + data = self.store.get(key) + if not data or not bytes(data).startswith(self._MANIFEST_MAGIC): + return None + import json + return json.loads(bytes(data)[len(self._MANIFEST_MAGIC):].decode()) + + def _manifest_size(self, key): + m = self._read_manifest(key) + return m["size"] if m else 0 + + def get_file(self, key, path): + m = self._read_manifest(key) + if m is None: + return False, 0, False, "" + size, n, chunk = m["size"], m["n"], m["chunk"] + os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) + tmp = path + ".tmp" + if size == 0: + with open(tmp, "wb") as f: + pass + os.replace(tmp, path) + return True, 0, True, "" + keys = self._chunk_keys(key, n) + sizes = [min(chunk, size - i * chunk) for i in range(n)] + + # Primary path: RDMA into a pre-registered staging buffer (registration + # done once at startup, not per-op), then write it to the destination + # file. The file.write(memoryview) releases the GIL for the syscall, so + # concurrent restores overlap. Falls back to direct per-op registration + # when no pool is configured or the blob exceeds a staging buffer. + if self.pool is not None and size <= self.pool.size: + mm, base = self.pool.acquire() + try: + ptrs = [base + i * chunk for i in range(n)] + rets = self.store.batch_get_into(keys, ptrs, sizes) + if any(r < 0 for r in rets): + # eviction/lease race -> recompute + return False, 0, False, "" + with open(tmp, "wb") as f: + f.write(memoryview(mm)[:size]) + finally: + self.pool.release((mm, base)) + os.replace(tmp, path) + return True, size, True, "" + + ok, transient, err = self._get_unregistered(tmp, size, n, keys, chunk, sizes) + if not ok: + try: + os.remove(tmp) + except OSError: + pass + if transient: + # A chunk was evicted / its lease expired between the manifest + # read and the transfer. Report not-found so the caller falls + # back to recompute instead of failing the request. + return False, 0, False, "" + return False, 0, False, err + os.replace(tmp, path) + return True, size, True, "" + + def _get_unregistered(self, tmp, size, n, keys, chunk, sizes): + """Returns (ok, transient, err). transient=True means an eviction/lease + race (recoverable by recompute), not a hard error.""" + with open(tmp, "wb") as f: + f.truncate(size) + with open(tmp, "r+b") as f: + mm = mmap.mmap(f.fileno(), size) + try: + base = _ptr_of(mm) + self.store.register_buffer(base, size) + try: + ptrs = [base + i * chunk for i in range(n)] + rets = self.store.batch_get_into(keys, ptrs, sizes) + if any(r < 0 for r in rets): + # negative return == chunk missing/evicted/lease expired + return False, True, f"batch_get_into rets={rets}" + finally: + self.store.unregister_buffer(base) + finally: + mm.close() + return True, False, "" + + def put_bytes(self, key, value, replica_num, soft_pin): + ret = self.store.put(key, value, self._rc(replica_num, soft_pin)) + if ret != 0: + return False, 0, f"put returned {ret}" + return True, len(value), "" + + def get_bytes(self, key): + data = self.store.get(key) + if not data: + return False, b"", False, "" + return True, bytes(data), True, "" + + def remove(self, key, force): + # remove manifest + all chunks. Chunk count from manifest if present. + m = None + try: + m = self._read_manifest(key) + except Exception: # noqa: BLE001 + pass + keys = [key] + if m: + keys += self._chunk_keys(key, m["n"]) + for k in keys: + try: + self.store.remove(k, True) if force else self.store.remove(k) + except Exception: # noqa: BLE001 + pass + return self.store.is_exist(key) != 1, "" + + def close(self): + if self.pool is not None: + self.pool.close() + try: + self.store.close() + except Exception: # noqa: BLE001 + pass + + +# --------------------------------------------------------------------------- # +# gRPC servicer +# --------------------------------------------------------------------------- # +class StoreProxyServicer(pb_grpc.StoreProxyServicer): + def __init__(self, backend: Backend): + self.b = backend + self.stats = Stats() + + def Health(self, request, context): + return pb.HealthReply(ok=True, backend=self.b.name, protocol=self.b.protocol, + master=self.b.master, device=self.b.device, detail="ready") + + def Exists(self, request, context): + res = self.b.exists(list(request.keys)) + self.stats.add_exists(len(res)) + return pb.ExistsReply(present=res) + + def PutFile(self, request, context): + t0 = time.perf_counter() + try: + ok, n, existed, err = self.b.put_file( + request.key, request.path, request.replica_num, + request.soft_pin, request.skip_if_exists) + except Exception as e: # noqa: BLE001 + return pb.PutFileReply(ok=False, error=f"{e}") + ms = (time.perf_counter() - t0) * 1000 + if ok and not existed: + self.stats.add_put(n, ms) + return pb.PutFileReply(ok=ok, bytes=n, elapsed_ms=ms, existed=existed, error=err) + + def GetFile(self, request, context): + t0 = time.perf_counter() + try: + ok, n, found, err = self.b.get_file(request.key, request.path) + except Exception as e: # noqa: BLE001 + return pb.GetFileReply(ok=False, error=f"{e}") + ms = (time.perf_counter() - t0) * 1000 + if ok: + self.stats.add_get(n, ms) + return pb.GetFileReply(ok=ok, bytes=n, elapsed_ms=ms, found=found, error=err) + + def PutBytes(self, request, context): + t0 = time.perf_counter() + try: + ok, n, err = self.b.put_bytes(request.key, request.value, request.replica_num, request.soft_pin) + except Exception as e: # noqa: BLE001 + return pb.PutBytesReply(ok=False, error=f"{e}") + ms = (time.perf_counter() - t0) * 1000 + if ok: + self.stats.add_put(n, ms) + return pb.PutBytesReply(ok=ok, bytes=n, elapsed_ms=ms, error=err) + + def GetBytes(self, request, context): + t0 = time.perf_counter() + try: + ok, val, found, err = self.b.get_bytes(request.key) + except Exception as e: # noqa: BLE001 + return pb.GetBytesReply(ok=False, error=f"{e}") + ms = (time.perf_counter() - t0) * 1000 + if ok: + self.stats.add_get(len(val), ms) + return pb.GetBytesReply(ok=ok, value=val, elapsed_ms=ms, found=found, error=err) + + def Remove(self, request, context): + try: + ok, err = self.b.remove(request.key, request.force) + except Exception as e: # noqa: BLE001 + return pb.RemoveReply(ok=False, error=f"{e}") + return pb.RemoveReply(ok=ok, error=err) + + def Stats(self, request, context): + s = self.stats + return pb.StatsReply( + put_ops=s.put_ops, get_ops=s.get_ops, exists_ops=s.exists_ops, + put_bytes=s.put_bytes, get_bytes=s.get_bytes, + put_ms_total=s.put_ms_total, get_ms_total=s.get_ms_total, + backend=self.b.name) + + +def build_backend(args) -> Backend: + if args.backend == "local": + return LocalBackend(args.local_dir) + # Resolve protocol=auto. + if args.protocol == "auto": + if args.device and args.device != "": + args.protocol = "rdma" + else: + # Probe for RDMA devices. + has_rdma = os.path.isdir("/sys/class/infiniband") and bool(os.listdir("/sys/class/infiniband")) + if has_rdma: + args.protocol, args.device = "rdma", "auto-discovery" + else: + args.protocol = "tcp" + log.info("protocol=auto resolved to %s (device=%r)", args.protocol, args.device) + return MooncakeBackend(args) + + +def main(): + ap = argparse.ArgumentParser(description="Mooncake Store Proxy (gRPC)") + ap.add_argument("--listen", default="127.0.0.1:52060") + ap.add_argument("--backend", choices=["mooncake", "local"], default="mooncake") + ap.add_argument("--local-dir", default=os.environ.get("OMB_STORE_DATA", "./run/store")) + ap.add_argument("--master", default=os.environ.get("OMB_STORE_MASTER", "127.0.0.1:52061")) + ap.add_argument("--metadata", default=os.environ.get("OMB_STORE_META", "P2PHANDSHAKE")) + ap.add_argument("--protocol", default=os.environ.get("OMB_STORE_PROTOCOL", "tcp"), + choices=["tcp", "rdma", "auto"]) + ap.add_argument("--device", default=os.environ.get("OMB_STORE_DEVICE", "")) + ap.add_argument("--local-hostname", default=os.environ.get("OMB_STORE_HOST", "127.0.0.1")) + ap.add_argument("--global-segment-size", type=int, default=int(os.environ.get("OMB_STORE_SEGMENT", str(16 << 30)))) + ap.add_argument("--local-buffer-size", type=int, default=int(os.environ.get("OMB_STORE_BUFFER", str(2 << 30)))) + ap.add_argument("--stripe-mb", type=int, default=int(os.environ.get("OMB_STORE_STRIPE_MB", "64")), + help="split KV blobs into <=N MiB chunks for parallel batch transfer") + ap.add_argument("--staging-mb", type=int, default=int(os.environ.get("OMB_STORE_STAGING_MB", "4096")), + help="size of each pre-registered staging buffer (MiB); blobs above this use per-op registration") + ap.add_argument("--staging-count", type=int, default=int(os.environ.get("OMB_STORE_STAGING_COUNT", "3")), + help="number of pre-registered staging buffers (concurrency of zero-copy transfers)") + ap.add_argument("--setup-timeout", type=float, default=60.0) + ap.add_argument("--max-workers", type=int, default=16) + ap.add_argument("--max-msg-mb", type=int, default=512) + ap.add_argument("--warmup", action="store_true", help="warm the client with a put/get so the first real op is fast") + ap.add_argument("--selftest", action="store_true") + args = ap.parse_args() + + logging.basicConfig(level=logging.INFO, format="%(asctime)s store-proxy %(levelname)s %(message)s") + + backend = build_backend(args) + log.info("backend=%s protocol=%s master=%s", backend.name, backend.protocol, backend.master) + + if args.warmup or args.selftest: + _warmup(backend) + if args.selftest: + _selftest(backend) + backend.close() + return + + opts = [ + ("grpc.max_send_message_length", args.max_msg_mb << 20), + ("grpc.max_receive_message_length", args.max_msg_mb << 20), + ] + server = grpc.server(futures.ThreadPoolExecutor(max_workers=args.max_workers), options=opts) + pb_grpc.add_StoreProxyServicer_to_server(StoreProxyServicer(backend), server) + server.add_insecure_port(args.listen) + server.start() + log.info("StoreProxy listening on %s (backend=%s)", args.listen, backend.name) + try: + server.wait_for_termination() + except KeyboardInterrupt: + log.info("shutting down") + server.stop(2).wait() + backend.close() + + +def _warmup(backend: Backend): + import tempfile + key = "omb:_warmup_" + with tempfile.NamedTemporaryFile(delete=False) as tf: + tf.write(b"warmup" * 4096) # 24 KiB + p = tf.name + try: + t0 = time.perf_counter() + backend.put_file(key, p, 1, False, False) + t1 = time.perf_counter() + outp = p + ".out" + backend.get_file(key, outp) + t2 = time.perf_counter() + backend.remove(key, True) + log.info("warmup: put %.1fms get %.1fms", (t1 - t0) * 1e3, (t2 - t1) * 1e3) + for q in (p, p + ".out"): + try: + os.remove(q) + except OSError: + pass + except Exception as e: # noqa: BLE001 + log.warning("warmup failed: %s", e) + + +def _selftest(backend: Backend): + import tempfile + log.info("=== selftest backend=%s ===", backend.name) + sizes = [1 << 20, 16 << 20, 128 << 20] # 1MiB, 16MiB, 128MiB + for sz in sizes: + with tempfile.NamedTemporaryFile(delete=False) as tf: + tf.write(os.urandom(sz)) + p = tf.name + key = f"omb:selftest:{sz}" + t0 = time.perf_counter() + ok, n, existed, err = backend.put_file(key, p, 1, True, False) + t1 = time.perf_counter() + assert ok, f"put failed: {err}" + outp = p + ".out" + ok, n2, found, err = backend.get_file(key, outp) + t2 = time.perf_counter() + assert ok and found, f"get failed: {err}" + import filecmp + same = filecmp.cmp(p, outp, shallow=False) + put_gbps = sz / (t1 - t0) / 1e9 + get_gbps = sz / (t2 - t1) / 1e9 + log.info("size=%6dMiB put=%7.1fms (%.2f GB/s) get=%7.1fms (%.2f GB/s) roundtrip_ok=%s", + sz >> 20, (t1 - t0) * 1e3, put_gbps, (t2 - t1) * 1e3, get_gbps, same) + assert same, "roundtrip data mismatch!" + ex = backend.exists([key, key + ":absent"]) + assert ex[0] == 1 and ex[1] == 0, f"exists wrong: {ex}" + backend.remove(key, True) + for q in (p, outp): + os.remove(q) + log.info("=== selftest PASSED ===") + + +if __name__ == "__main__": + main() diff --git a/mooncake-store/tests/offset_allocator_test.cpp b/mooncake-store/tests/offset_allocator_test.cpp index d0e5519e..1cbbab52 100644 --- a/mooncake-store/tests/offset_allocator_test.cpp +++ b/mooncake-store/tests/offset_allocator_test.cpp @@ -865,7 +865,7 @@ TEST_F(OffsetAllocatorTest, RandomSmallAllocWithLargeAllocatorSize) { } } -// ========== EDGE CASE TESTS, Generated by AI ========== +// ========== EDGE CASE TESTS ========== // Test zero size allocation - should fail TEST_F(OffsetAllocatorTest, ZeroSizeAllocation) { @@ -1008,7 +1008,7 @@ TEST_F(OffsetAllocatorTest, PowerOfTwoAllocation) { } } -// ========== BIN SYSTEM TESTS, Generated by AI ========== +// ========== BIN SYSTEM TESTS ========== // Test bin size calculations and selection TEST_F(OffsetAllocatorTest, BinSizeCalculation) {