def count_batch(self, rs_header: ErRollSiteHeader, batch_pairs):
L.trace(f'count batch. rs_key={rs_header.get_rs_key()}, rs_header={rs_header}, batch_pairs={batch_pairs}')
batch_seq_id = rs_header._batch_seq
stream_seq_id = rs_header._stream_seq
if self._rs_header is None:
self._rs_header = rs_header
self._rs_key = rs_header.get_rs_key()
L.debug(f"header arrived. rs_key={rs_header.get_rs_key()}, rs_header={rs_header}")
self._header_arrive_event.set()
self._batch_seq_to_pair_counter[batch_seq_id] = batch_pairs
self._stream_seq_to_pair_counter[stream_seq_id] += batch_pairs
self._stream_seq_to_batch_seq[stream_seq_id] = batch_seq_id
def _pull_one(self, rs_header: ErRollSiteHeader, options: dict = None):
if options is None:
options = {}
start_time = time.time()
rs_key = rp_name = rs_header.get_rs_key()
rp_namespace = self.roll_site_session_id
transfer_tag_prefix = "putBatch-" + rs_header.get_rs_key() + "-"
last_total_batches = None
last_cur_pairs = -1
pull_attempts = 0
data_type = None
L.debug(f'pulling rs_key={rs_key}')
try:
# make sure rollpair already created
pull_header_interval = self.pull_header_interval
pull_header_timeout = self.pull_header_timeout # skips pickling self
pull_interval = self.pull_interval # skips pickling self
def get_partition_status(task):
put_batch_task = PutBatchTask(
transfer_tag_prefix + str(task._inputs[0]._id), None)
return put_batch_task.get_status(pull_interval)
def get_status(roll_site):
pull_status = {}
total_pairs = 0
total_batches = 0
all_finished = True
final_options = options.copy()
final_options['create_if_missing'] = False
store = roll_site.ctx.rp_ctx.load(name=rp_name,
namespace=rp_namespace,
options=final_options)
if store is None:
raise ValueError(
f'illegal state for rp_name={rp_name}, rp_namespace={rp_namespace}'
)
all_status = store.with_stores(
get_partition_status,
options={"__op": "get_partition_status"})
for part_id, part_status in all_status:
if not part_status.is_finished:
all_finished = False
pull_status[part_id] = part_status
total_batches += part_status.total_batches
total_pairs += part_status.total_pairs
return pull_status, all_finished, total_batches, total_pairs
def clear_status(task):
return PutBatchTask(transfer_tag_prefix +
str(task._inputs[0]._id)).clear_status()
wait_time = 0
header_response = None
while wait_time < pull_header_timeout and \
(header_response is None or not isinstance(header_response[0][1], ErRollSiteHeader)):
final_options = options.copy()
final_options['create_if_missing'] = True
final_options['total_partitions'] = 1
header_response = self.ctx.rp_ctx.load(name=STATUS_TABLE_NAME,
namespace=rp_namespace,
options=final_options) \
.with_stores(lambda x: PutBatchTask(transfer_tag_prefix + "0").get_header(pull_header_interval), options={"__op": "pull_header"})
wait_time += pull_header_interval
#pull_status, all_finished, total_batches, total_pairs = stat_all_status(self)
L.debug(
f"roll site get header_response: rs_key={rs_key}, rs_header={rs_header}, wait_time={wait_time}"
)
if header_response is None or not isinstance(
header_response[0][1], ErRollSiteHeader):
raise IOError(
f"roll site pull header failed: rs_key={rs_key}, rs_header={rs_header}, timeout={self.pull_header_timeout}"
)
else:
header: ErRollSiteHeader = header_response[0][1]
# TODO:0: push bytes has only one partition, that means it has finished, need not get_status
data_type = header._data_type
L.debug(
f"roll site pull header successful: rs_key={rs_key}, rs_header={header}"
)
pull_status = {}
for cur_retry in range(self.pull_max_retry):
pull_attempts = cur_retry
pull_status, all_finished, total_batches, cur_pairs = get_status(
self)
if not all_finished:
L.debug(
f'getting status NOT finished for rs_key={rs_key}, '
f'rs_header={rs_header}, '
f'cur_status={pull_status}, '
f'attempts={pull_attempts}, '
f'cur_pairs={cur_pairs}, '
f'last_cur_pairs={last_cur_pairs}, '
f'total_batches={total_batches}, '
f'last_total_batches={last_total_batches}, '
f'elapsed={time.time() - start_time}')
if last_cur_pairs == cur_pairs and cur_pairs > 0:
raise IOError(
f"roll site pull waiting failed because there is no updated progress: rs_key={rs_key}, "
f"rs_header={rs_header}, pull_status={pull_status}, last_cur_pairs={last_cur_pairs}, cur_pairs={cur_pairs}"
)
else:
L.debug(
f"getting status DO finished for rs_key={rs_key}, rs_header={rs_header}, pull_status={pull_status}, cur_pairs={cur_pairs}, total_batches={total_batches}"
)
rp = self.ctx.rp_ctx.load(name=rp_name,
namespace=rp_namespace)
clear_future = self._receive_executor_pool.submit(
rp.with_stores,
clear_status,
options={"__op": "clear_status"})
if data_type == "object":
result = pickle.loads(b''.join(
map(
lambda t: t[1],
sorted(rp.get_all(),
key=lambda x: int.from_bytes(
x[0], "big")))))
rp.destroy()
L.debug(
f"pulled object: rs_key={rs_key}, rs_header={rs_header}, is_none={result is None}, "
f"elapsed={time.time() - start_time}")
else:
result = rp
if L.isEnabledFor(logging.DEBUG):
L.debug(
f"pulled roll_pair: rs_key={rs_key}, rs_header={rs_header}, rp.count={rp.count()}, "
f"elapsed={time.time() - start_time}")
clear_future.result()
return result
last_total_batches = total_batches
last_cur_pairs = cur_pairs
raise IOError(
f"roll site pull failed. max try exceeded: {self.pull_max_retry}, rs_key={rs_key}, "
f"rs_header={rs_header}, pull_status={pull_status}")
except Exception as e:
L.exception(
f"fatal error: when pulling rs_key={rs_key}, rs_header={rs_header}, attempts={pull_attempts}"
)
raise e
def _push_rollpair(self,
rp: RollPair,
rs_header: ErRollSiteHeader,
options: dict = None):
if options is None:
options = {}
rs_key = rs_header.get_rs_key()
if L.isEnabledFor(logging.DEBUG):
L.debug(
f"pushing rollpair: rs_key={rs_key}, rs_header={rs_header}, rp.count={rp.count()}"
)
start_time = time.time()
rs_header._total_partitions = rp.get_partitions()
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
batches_per_stream = self.push_batches_per_stream
body_bytes = self.batch_body_bytes
endpoint = self.ctx.proxy_endpoint
max_retry_cnt = self.push_max_retry
long_retry_cnt = self.push_long_retry
per_stream_timeout = self.push_per_stream_timeout
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}"
)
rp.with_stores(_push_partition, options={"__op": "push_partition"})
if L.isEnabledFor(logging.DEBUG):
L.debug(
f"pushed rollpair: rs_key={rs_key}, rs_header={rs_header}, count={rp.count()}, elapsed={time.time() - start_time}"
)
self.ctx.pushing_latch.count_down()
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()