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