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