def _get_remote_obj_store_table(parties, rollsite):
from eggroll.roll_site.utils.roll_site_utils import create_store_name
from eggroll.core.transfer_model import ErRollSiteHeader
tables = []
for role_party_id in parties:
_role = role_party_id[0]
_party_id = str(role_party_id[1])
_options = {}
obj_type = 'object'
roll_site_header = ErRollSiteHeader(
roll_site_session_id=rollsite.roll_site_session_id,
name=rollsite.name,
tag=rollsite.tag,
src_role=rollsite.local_role,
src_party_id=rollsite.party_id,
dst_role=_role,
dst_party_id=_party_id,
data_type=obj_type,
options=_options)
_tagged_key = create_store_name(roll_site_header)
namespace = rollsite.roll_site_session_id
tables.append(rollsite.ctx.rp_ctx.load(namespace, _tagged_key))
return tables
def __init__(self, adapter: RollSiteAdapter, options: dict = None):
if options is None:
options = {}
self.adapter = adapter
self.roll_site_header: ErRollSiteHeader = adapter.roll_site_header
self.namespace = adapter.namespace
self.name = create_store_name(self.roll_site_header)
self.tagged_key = ''
self.obj_type = adapter.obj_type
self.proxy_endpoint = adapter.proxy_endpoint
channel = self.grpc_channel_factory.create_channel(self.proxy_endpoint)
self.stub = proxy_pb2_grpc.DataTransferServiceStub(channel)
static_er_conf = get_static_er_conf()
self.__bin_packet_len = int(options.get(
RollSiteConfKeys.EGGROLL_ROLLSITE_ADAPTER_SENDBUF_SIZE.key,
static_er_conf.get(RollSiteConfKeys.EGGROLL_ROLLSITE_ADAPTER_SENDBUF_SIZE.key,
RollSiteConfKeys.EGGROLL_ROLLSITE_ADAPTER_SENDBUF_SIZE.default_value)))
self.total_written = 0
self.ba = bytearray(self.__bin_packet_len)
self.buffer = ArrayByteBuffer(self.ba)
self.writer = PairBinWriter(pair_buffer=self.buffer)
self.push_batch_cnt = 0
self.push_pair_cnt = 0
self.topic_src = proxy_pb2.Topic(name=self.name, partyId=self.roll_site_header._src_party_id,
role=self.roll_site_header._src_role, callback=None)
self.topic_dst = proxy_pb2.Topic(name=self.name, partyId=self.roll_site_header._dst_party_id,
role=self.roll_site_header._dst_role, callback=None)
def _thread_receive(self, packet, namespace, roll_site_header: ErRollSiteHeader):
try:
table_name = create_store_name(roll_site_header)
is_standalone = self.ctx.rp_ctx.get_session().get_option(
SessionConfKeys.CONFKEY_SESSION_DEPLOY_MODE) == "standalone"
if is_standalone:
status_rp = self.ctx.rp_ctx.load(namespace, STATUS_TABLE_NAME + DELIM + self.ctx.roll_site_session_id)
retry_cnt = 0
# TODO:0: sleep retry count and timeout
while True:
msg = f"retry pull: retry_cnt: {retry_cnt}," + \
f" tagged_key: '{table_name}', packet: {to_one_line_string(packet)}, namespace: {namespace}"
if retry_cnt % 10 == 0:
L.info(msg)
else:
L.debug(msg)
retry_cnt += 1
ret_list = status_rp.get(table_name)
if ret_list:
table_namespace = ret_list[2]
table_name = ret_list[1]
obj_type = ret_list[0]
break
time.sleep(min(0.1 * retry_cnt, 30))
else:
retry_cnt = 0
ret_packet = self.stub.unaryCall(packet)
while ret_packet.header.ack != 123:
msg = f"retry pull: retry_cnt: {retry_cnt}," + \
f" store_name: '{table_name}', packet: {to_one_line_string(packet)}, namespace: {namespace}"
if retry_cnt % 10 == 0:
L.info(msg)
else:
L.debug(msg)
retry_cnt += 1
if ret_packet.header.ack in ERROR_STATES:
raise IOError("receive terminated")
ret_packet = self.stub.unaryCall(packet)
time.sleep(min(0.1 * retry_cnt, 30))
obj_type = ret_packet.body.value
table_namespace = self.roll_site_session_id
L.info(f"pull status done: table_name:{table_name}, packet:{to_one_line_string(packet)}, namespace:{namespace}")
rp = self.ctx.rp_ctx.load(namespace=table_namespace, name=table_name)
if obj_type == b'object':
result = rp.get(table_name)
L.info(f"pull success: {table_name}, type: {obj_type}")
else:
result = rp
L.info(f"pull success: {table_name}, count: {rp.count()}, type: {obj_type}")
return result
except Exception as e:
L.exception(f"pull error:{e}")
raise GrpcCallError("push", self.ctx.proxy_endpoint, e)
finally:
end_wall_time = time.time()
end_cpu_time = time.perf_counter()
P.info(f'{{"metric_type": "func_profile", "qualname": "RollSite.pull", "cpu_time": {end_cpu_time - self._pull_start_cpu_time}, "wall_time": {end_wall_time - self._pull_start_wall_time}}}')
def pull(self, parties: list = None):
self._pull_start_wall_time = time.time()
self._pull_start_cpu_time = time.perf_counter()
futures = []
for src_role, src_party_id in parties:
src_party_id = str(src_party_id)
roll_site_header = ErRollSiteHeader(
roll_site_session_id=self.roll_site_session_id,
name=self.name,
tag=self.tag,
src_role=src_role,
src_party_id=src_party_id,
dst_role=self.local_role,
dst_party_id=self.party_id)
_tagged_key = create_store_name(roll_site_header)
name = _tagged_key
model = proxy_pb2.Model(name=_stringify(roll_site_header))
task_info = proxy_pb2.Task(taskId=name, model=model)
topic_src = proxy_pb2.Topic(name="get_status",
partyId=src_party_id,
role=src_role,
callback=None)
topic_dst = proxy_pb2.Topic(name="get_status",
partyId=self.party_id,
role=self.local_role,
callback=None)
get_status_command = proxy_pb2.Command(name="get_status")
conf_test = proxy_pb2.Conf(overallTimeout=1000,
completionWaitTimeout=1000,
packetIntervalTimeout=1000,
maxRetries=10)
metadata = proxy_pb2.Metadata(task=task_info,
src=topic_src,
dst=topic_dst,
command=get_status_command,
operator="getStatus",
seq=0,
ack=0)
packet = proxy_pb2.Packet(header=metadata)
namespace = self.roll_site_session_id
L.info(
f"pulling prepared tagged_key: {_tagged_key}, packet:{to_one_line_string(packet)}"
)
futures.append(
RollSite.receive_exeutor_pool.submit(RollSite._thread_receive,
self, packet, namespace,
roll_site_header))
return futures
def push(self, obj, parties: list = None):
L.info(
f"pushing: self:{self.__dict__}, obj_type:{type(obj)}, parties:{parties}"
)
self._push_start_wall_time = time.time()
self._push_start_cpu_time = time.perf_counter()
futures = []
for role_party_id in parties:
self.ctx.pushing_task_count += 1
_role = role_party_id[0]
_party_id = str(role_party_id[1])
_options = {}
obj_type = 'rollpair' if isinstance(obj, RollPair) else 'object'
roll_site_header = ErRollSiteHeader(
roll_site_session_id=self.roll_site_session_id,
name=self.name,
tag=self.tag,
src_role=self.local_role,
src_party_id=self.party_id,
dst_role=_role,
dst_party_id=_party_id,
data_type=obj_type,
options=_options)
_tagged_key = create_store_name(roll_site_header)
L.debug(f"pushing start party:{type(obj)}, {_tagged_key}")
namespace = self.roll_site_session_id
if isinstance(obj, RollPair):
rp = obj
else:
rp = self.ctx.rp_ctx.load(namespace, _tagged_key)
rp.put(_tagged_key, obj)
rp.disable_gc()
L.info(f"pushing prepared: {type(obj)}, tag_key:{_tagged_key}")
def map_values(_tagged_key, is_standalone, roll_site_header):
if is_standalone:
dst_name = _tagged_key
store_type = rp.get_store_type()
else:
dst_name = DELIM.join([
_tagged_key, self.dst_host,
str(self.dst_port), obj_type
])
store_type = StoreTypes.ROLLPAIR_ROLLSITE
if is_standalone:
status_rp = self.ctx.rp_ctx.load(
namespace,
STATUS_TABLE_NAME + DELIM + self.roll_site_session_id,
options=_options)
status_rp.disable_gc()
if isinstance(obj, RollPair):
status_rp.put(_tagged_key,
(obj_type.encode("utf-8"), rp.get_name(),
rp.get_namespace()))
else:
status_rp.put(
_tagged_key,
(obj_type.encode("utf-8"), dst_name, namespace))
else:
store = rp.get_store()
store_locator = store._store_locator
new_store_locator = ErStoreLocator(
store_type=store_type,
namespace=namespace,
name=dst_name,
total_partitions=store_locator._total_partitions,
partitioner=store_locator._partitioner,
serdes=store_locator._serdes)
# TODO:0: move options from job to store when database modification finished
options = {
"roll_site_header": roll_site_header,
"proxy_endpoint": self.ctx.proxy_endpoint,
"obj_type": obj_type
}
if isinstance(obj, RollPair):
roll_site_header._options[
'total_partitions'] = obj.get_store(
)._store_locator._total_partitions
L.info(
f"RollSite.push: pushing {roll_site_header}, type: RollPair, count: {obj.count()}"
)
else:
L.info(
f"RollSite.push: pushing {roll_site_header}, type: object"
)
rp.map_values(
lambda v: v,
output=ErStore(store_locator=new_store_locator),
options=options)
L.info(
f"RollSite.push: push {roll_site_header} done. type:{type(obj)}"
)
return _tagged_key
future = RollSite.receive_exeutor_pool.submit(
map_values, _tagged_key, self._is_standalone, roll_site_header)
if not self._is_standalone and (obj_type == 'object'
or obj_type == b'object'):
tmp_rp = rp
else:
tmp_rp = None
future.add_done_callback(
functools.partial(self._push_callback, tmp_rp=tmp_rp))
futures.append(future)
return futures
def push(self, obj, parties: list = None):
L.info(f"pushing: self:{self.__dict__}, obj_type:{type(obj)}, parties:{parties}")
self._push_start_wall_time = time.time()
self._push_start_cpu_time = time.perf_counter()
self.ctx.pushing_task_count += 1
futures = []
for role_party_id in parties:
_role = role_party_id[0]
_party_id = str(role_party_id[1])
_options = {}
obj_type = 'rollpair' if isinstance(obj, RollPair) else 'object'
roll_site_header = ErRollSiteHeader(
roll_site_session_id=self.roll_site_session_id,
name=self.name,
tag=self.tag,
src_role=self.local_role,
src_party_id=self.party_id,
dst_role=_role,
dst_party_id=_party_id,
data_type=obj_type,
options=_options)
_tagged_key = create_store_name(roll_site_header)
L.debug(f"pushing start party:{type(obj)}, {_tagged_key}")
namespace = self.roll_site_session_id
if isinstance(obj, RollPair):
rp = obj
else:
rp = self.ctx.rp_ctx.load(namespace, _tagged_key)
rp.put(_tagged_key, obj)
rp.disable_gc()
L.info(f"pushing prepared: {type(obj)}, tag_key:{_tagged_key}")
def map_values(_tagged_key):
is_standalone = self.ctx.rp_ctx.get_session().get_option(
SessionConfKeys.CONFKEY_SESSION_DEPLOY_MODE) == DeployModes.STANDALONE
if is_standalone:
dst_name = _tagged_key
store_type = rp.get_store_type()
else:
dst_name = DELIM.join([_tagged_key,
self.dst_host,
str(self.dst_port),
obj_type])
store_type = StoreTypes.ROLLPAIR_ROLLSITE
if is_standalone:
status_rp = self.ctx.rp_ctx.load(namespace, STATUS_TABLE_NAME + DELIM + self.roll_site_session_id, options=_options)
status_rp.disable_gc()
if isinstance(obj, RollPair):
status_rp.put(_tagged_key, (obj_type.encode("utf-8"), rp.get_name(), rp.get_namespace()))
else:
status_rp.put(_tagged_key, (obj_type.encode("utf-8"), dst_name, namespace))
else:
store = rp.get_store()
store_locator = store._store_locator
new_store_locator = ErStoreLocator(store_type=store_type,
namespace=namespace,
name=dst_name,
total_partitions=store_locator._total_partitions,
partitioner=store_locator._partitioner,
serdes=store_locator._serdes)
# TODO:0: move options from job to store when database modification finished
options = {"roll_site_header": roll_site_header,
"proxy_endpoint": self.ctx.proxy_endpoint,
"obj_type": obj_type}
if isinstance(obj, RollPair):
roll_site_header._options['total_partitions'] = obj.get_store()._store_locator._total_partitions
L.debug(f"pushing map_values: {dst_name}, count: {obj.count()}, tag_key:{_tagged_key}")
rp.map_values(lambda v: v,
output=ErStore(store_locator=new_store_locator),
options=options)
L.info(f"pushing map_values done:{type(obj)}, tag_key:{_tagged_key}")
return _tagged_key
future = self.process_pool.submit(map_values, _tagged_key)
future.add_done_callback(self._push_callback)
futures.append(future)
self.process_pool.shutdown(wait=False)
return futures