def _sync_job_part(self, job_id, partition_id):
if partition_id not in self._job_part_map or \
self._job_part_map[partition_id] is None or \
self._job_part_map[partition_id].job_id != job_id:
etcd_key = common.portal_job_part_etcd_key(self._portal_name,
job_id, partition_id)
data = self._etcd.get_data(etcd_key)
if data is None:
self._job_part_map[partition_id] = dp_pb.PortalJobPart(
job_id=job_id, rank_id=-1, partition_id=partition_id)
else:
self._job_part_map[partition_id] = \
text_format.Parse(data, dp_pb.PortalJobPart())
return self._job_part_map[partition_id]
def _sync_job_part(self, job_id, partition_id):
if partition_id not in self._job_part_map or \
self._job_part_map[partition_id] is None or \
self._job_part_map[partition_id].job_id != job_id:
kvstore_key = common.portal_job_part_kvstore_key(
self._portal_name, job_id, partition_id)
data = self._kvstore.get_data(kvstore_key)
if data is None:
self._job_part_map[partition_id] = dp_pb.PortalJobPart(
job_id=job_id, rank_id=-1, partition_id=partition_id)
else:
self._job_part_map[partition_id] = \
text_format.Parse(data, dp_pb.PortalJobPart(),
allow_unknown_field=True)
return self._job_part_map[partition_id]
def _finish_job_part(self, job_id, partition_id, src_state, target_state):
job_part = self._sync_job_part(job_id, partition_id)
assert job_part is not None and job_part.part_state == src_state
new_job_part = dp_pb.PortalJobPart()
new_job_part.MergeFrom(job_part)
new_job_part.part_state = target_state
new_job_part.rank_id = -1
self._update_job_part(new_job_part)
def _try_to_alloc_part(self, rank_id, src_state, target_state):
alloc_partition_id = None
processing_job = self._sync_processing_job()
assert processing_job is not None
job_id = self._processing_job.job_id
for partition_id in range(self._output_partition_num):
part_job = self._sync_job_part(job_id, partition_id)
if part_job.part_state == src_state and \
alloc_partition_id is None:
alloc_partition_id = partition_id
if part_job.part_state == target_state and \
part_job.rank_id == rank_id:
alloc_partition_id = partition_id
break
if alloc_partition_id is None:
return None
part_job = self._job_part_map[alloc_partition_id]
if part_job.part_state == src_state:
new_job_part = dp_pb.PortalJobPart(job_id=job_id,
rank_id=rank_id,
partition_id=alloc_partition_id,
part_state=target_state)
self._update_job_part(new_job_part)
return alloc_partition_id
