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 push(self, obj, parties: list = None, options: dict = None):
if options is None:
options = {}
futures = []
for role_party_id in parties:
self.ctx.pushing_latch.count_up()
dst_role = role_party_id[0]
dst_party_id = str(role_party_id[1])
data_type = 'rollpair' if isinstance(obj, RollPair) else 'object'
rs_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=dst_role,
dst_party_id=dst_party_id,
data_type=data_type)
if isinstance(obj, RollPair):
future = self._run_thread(self._impl_instance._push_rollpair,
obj, rs_header, options)
else:
future = self._run_thread(self._impl_instance._push_bytes, obj,
rs_header, options)
futures.append(future)
return futures
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 pull(self, parties: list = None):
futures = []
for src_role, src_party_id in parties:
src_party_id = str(src_party_id)
rs_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)
futures.append(
self._receive_executor_pool.submit(self._pull_one, rs_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