Add Ollama integration: cross-process KV-cache reuse via the Mooncake Store #5
|
|
@ -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)
|
||||
|
|
@ -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)
|
||||
|
|
@ -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()
|
||||
Loading…
Reference in New Issue