670 lines
26 KiB
Python
670 lines
26 KiB
Python
#!/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()
|