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 __init__(self,
roll_site_session_id,
rp_ctx: RollPairContext,
options: dict = None):
if options is None:
options = {}
self.roll_site_session_id = roll_site_session_id
self.rp_ctx = rp_ctx
self.role = options["self_role"]
self.party_id = str(options["self_party_id"])
self.proxy_endpoint = options["proxy_endpoint"]
self.is_standalone = self.rp_ctx.get_session().get_option(
SessionConfKeys.CONFKEY_SESSION_DEPLOY_MODE) == "standalone"
if self.is_standalone:
self.stub = None
else:
channel = self.grpc_channel_factory.create_channel(
self.proxy_endpoint)
self.stub = proxy_pb2_grpc.DataTransferServiceStub(channel)
self.init_job_session_pair(self.roll_site_session_id,
self.rp_ctx.session_id)
self.pushing_task_count = 0
self.rp_ctx.get_session().add_exit_task(self.push_complete)
L.info(f"inited RollSiteContext: {self.__dict__}")
def test_get_route_table(self, point='localhost:9470', party='10001', key="shiqili"):
channel = grpc.insecure_channel(point)
stub = proxy_pb2_grpc.DataTransferServiceStub(channel)
salt = self._get_salt()
md5hash = hashlib.md5((salt + key).encode("utf-8")).hexdigest()
topic_dst = proxy_pb2.Topic(partyId=party)
metadata = proxy_pb2.Metadata(dst=topic_dst, operator="get_route_table")
data = proxy_pb2.Data(key=md5hash, value=salt.encode('utf-8'))
packet = proxy_pb2.Packet(header=metadata, body=data)
ret_packet = stub.unaryCall(packet)
print(ret_packet.body.value.decode('utf8'))
return ret_packet.body.value.decode('utf8')
def __init__(self,
roll_site_session_id,
rp_ctx: RollPairContext,
options: dict = None):
if options is None:
options = {}
self.roll_site_session_id = roll_site_session_id
self.rp_ctx = rp_ctx
self.role = options["self_role"]
self.party_id = str(options["self_party_id"])
self._options = options
endpoint = options["proxy_endpoint"]
if isinstance(endpoint, str):
splitted = endpoint.split(':')
self.proxy_endpoint = ErEndpoint(host=splitted[0].strip(),
port=int(splitted[1].strip()))
elif isinstance(endpoint, ErEndpoint):
self.proxy_endpoint = endpoint
else:
raise ValueError("endpoint only support str and ErEndpoint type")
self.is_standalone = RollSiteConfKeys.EGGROLL_ROLLSITE_DEPLOY_MODE.get_with(
options) == "standalone"
if self.is_standalone:
self.stub = None
else:
channel = self.grpc_channel_factory.create_channel(
self.proxy_endpoint)
self.stub = proxy_pb2_grpc.DataTransferServiceStub(channel)
self.pushing_latch = CountDownLatch(0)
self.rp_ctx.get_session().add_exit_task(self._wait_push_complete)
self._wait_push_exit_timeout = int(
RollSiteConfKeys.EGGROLL_ROLLSITE_PUSH_OVERALL_TIMEOUT_SEC.
get_with(options))
L.info(f"inited RollSiteContext: {self.__dict__}")
def _push_partition(ertask):
rs_header._partition_id = ertask._inputs[0]._id
L.trace(
f"pushing rollpair partition. rs_key={rs_key}, partition_id={rs_header._partition_id}, rs_header={rs_header}"
)
from eggroll.core.grpc.factory import GrpcChannelFactory
from eggroll.core.proto import proxy_pb2_grpc
grpc_channel_factory = GrpcChannelFactory()
with create_adapter(ertask._inputs[0]) as db, db.iteritems() as rb:
# NOTICE AGAIN: all modifications to rs_header are limited in bs_helper.
# rs_header is shared by bs_helper and here. any modification in bs_helper affects this header.
# Remind that python's object references are passed by value,
# meaning the 'pointer' is copied, while the contents are modificable
bs_helper = _BatchStreamHelper(rs_header)
bin_batch_streams = bs_helper._generate_batch_streams(
pair_iter=rb,
batches_per_stream=batches_per_stream,
body_bytes=body_bytes)
channel = grpc_channel_factory.create_channel(endpoint)
stub = proxy_pb2_grpc.DataTransferServiceStub(channel)
for batch_stream in bin_batch_streams:
batch_stream_data = list(batch_stream)
cur_retry = 0
exception = None
while cur_retry < max_retry_cnt:
L.trace(
f'pushing rollpair partition stream. rs_key={rs_key}, partition_id={rs_header._partition_id}, rs_header={rs_header}, cur_retry={cur_retry}'
)
try:
stub.push(bs_helper.generate_packet(
batch_stream_data, cur_retry),
timeout=per_stream_timeout)
exception = None
break
except Exception as e:
if cur_retry < max_retry_cnt - long_retry_cnt:
retry_interval = round(
min(2 * cur_retry, 20) +
random.random() * 10, 3)
else:
retry_interval = round(
300 + random.random() * 10, 3)
L.warn(
f"push rp partition error. rs_key={rs_key}, partition_id={rs_header._partition_id}, rs_header={rs_header}, max_retry_cnt={max_retry_cnt}, cur_retry={cur_retry}, retry_interval={retry_interval}",
exc_info=e)
time.sleep(retry_interval)
if isinstance(e, RpcError) and e.code(
).name == 'UNAVAILABLE':
channel = grpc_channel_factory.create_channel(
endpoint, refresh=True)
stub = proxy_pb2_grpc.DataTransferServiceStub(
channel)
exception = e
finally:
cur_retry += 1
if exception is not None:
L.exception(
f"push partition failed. rs_key={rs_key}, partition_id={rs_header._partition_id}, rs_header={rs_header}, cur_retry={cur_retry}",
exc_info=exception)
raise exception
L.trace(
f'pushed rollpair partition stream. rs_key={rs_key}, partition_id={rs_header._partition_id}, rs_header={rs_header}, retry count={cur_retry - 1}'
)
L.trace(
f"pushed rollpair partition. rs_key={rs_key}, partition_id={rs_header._partition_id}, rs_header={rs_header}"
)
def _push_bytes(self,
obj,
rs_header: ErRollSiteHeader,
options: dict = None):
if options is None:
options = {}
start_time = time.time()
rs_key = rs_header.get_rs_key()
int_size = 4
if L.isEnabledFor(logging.DEBUG):
L.debug(f"pushing object: rs_key={rs_key}, rs_header={rs_header}")
def _generate_obj_bytes(py_obj, body_bytes):
key_id = 0
obj_bytes = pickle.dumps(py_obj)
obj_bytes_len = len(obj_bytes)
cur_pos = 0
while cur_pos <= obj_bytes_len:
yield key_id.to_bytes(
int_size, "big"), obj_bytes[cur_pos:cur_pos + body_bytes]
key_id += 1
cur_pos += body_bytes
rs_header._partition_id = 0
rs_header._total_partitions = 1
serdes = options.get('serdes', None)
if serdes is not None:
rs_header._options['serdes'] = serdes
wrapee_cls = options.get('wrapee_cls', None)
if serdes is not None:
rs_header._options['wrapee_cls'] = wrapee_cls
# NOTICE: all modifications to rs_header are limited in bs_helper.
# rs_header is shared by bs_helper and here. any modification in bs_helper affects this header.
# Remind that python's object references are passed by value,
# meaning the 'pointer' is copied, while the contents are modificable
bs_helper = _BatchStreamHelper(rs_header=rs_header)
bin_batch_streams = bs_helper._generate_batch_streams(
pair_iter=_generate_obj_bytes(obj, self.batch_body_bytes),
batches_per_stream=self.push_batches_per_stream,
body_bytes=self.batch_body_bytes)
grpc_channel_factory = GrpcChannelFactory()
channel = grpc_channel_factory.create_channel(self.ctx.proxy_endpoint)
stub = proxy_pb2_grpc.DataTransferServiceStub(channel)
max_retry_cnt = self.push_max_retry
long_retry_cnt = self.push_long_retry
per_stream_timeout = self.push_per_stream_timeout
# if use stub.push.future here, retry mechanism is a problem to solve
for batch_stream in bin_batch_streams:
cur_retry = 0
batch_stream_data = list(batch_stream)
exception = None
while cur_retry < max_retry_cnt:
L.trace(
f'pushing object stream. rs_key={rs_key}, rs_header={rs_header}, cur_retry={cur_retry}'
)
try:
stub.push(bs_helper.generate_packet(
batch_stream_data, cur_retry),
timeout=per_stream_timeout)
exception = None
break
except Exception as e:
if cur_retry <= max_retry_cnt - long_retry_cnt:
retry_interval = round(
min(2 * cur_retry, 20) + random.random() * 10, 3)
else:
retry_interval = round(300 + random.random() * 10, 3)
L.warn(
f"push object error. rs_key={rs_key}, partition_id={rs_header._partition_id}, rs_header={rs_header}, max_retry_cnt={max_retry_cnt}, cur_retry={cur_retry}, retry_interval={retry_interval}",
exc_info=e)
time.sleep(retry_interval)
if isinstance(e,
RpcError) and e.code().name == 'UNAVAILABLE':
channel = grpc_channel_factory.create_channel(
self.ctx.proxy_endpoint, refresh=True)
stub = proxy_pb2_grpc.DataTransferServiceStub(channel)
exception = e
finally:
cur_retry += 1
if exception is not None:
L.exception(
f"push object failed. rs_key={rs_key}, partition_id={rs_header._partition_id}, rs_header={rs_header}, cur_retry={cur_retry}",
exc_info=exception)
raise exception
L.trace(
f'pushed object stream. rs_key={rs_key}, rs_header={rs_header}, cur_retry={cur_retry - 1}'
)
L.debug(
f"pushed object: rs_key={rs_key}, rs_header={rs_header}, is_none={obj is None}, elapsed={time.time() - start_time}"
)
self.ctx.pushing_latch.count_down()
def get_proxy_data_channel():
channel = grpc.insecure_channel('{}:{}'.format(PROXY_HOST, PROXY_PORT))
stub = proxy_pb2_grpc.DataTransferServiceStub(channel)
return channel, stub
