def test_is_latest_execution_successful(self):
executions = []
self.assertFalse(
task_gen_utils.is_latest_execution_successful(executions))
# A successful execution.
executions.append(
metadata_store_pb2.Execution(
id=1,
type_id=2,
create_time_since_epoch=10,
last_known_state=metadata_store_pb2.Execution.COMPLETE))
self.assertTrue(
task_gen_utils.is_latest_execution_successful(executions))
# An older failed execution should not matter.
executions.append(
metadata_store_pb2.Execution(
id=2,
type_id=2,
create_time_since_epoch=5,
last_known_state=metadata_store_pb2.Execution.FAILED))
self.assertTrue(
task_gen_utils.is_latest_execution_successful(executions))
# A newer failed execution returns False.
executions.append(
metadata_store_pb2.Execution(
id=3,
type_id=2,
create_time_since_epoch=15,
last_known_state=metadata_store_pb2.Execution.FAILED))
self.assertFalse(
task_gen_utils.is_latest_execution_successful(executions))
# Finally, a newer successful execution returns True.
executions.append(
metadata_store_pb2.Execution(
id=4,
type_id=2,
create_time_since_epoch=20,
last_known_state=metadata_store_pb2.Execution.COMPLETE))
self.assertTrue(
task_gen_utils.is_latest_execution_successful(executions))
def generate(self) -> List[task_lib.Task]:
"""Generates tasks for executing the next executable nodes in the pipeline.
The returned tasks must have `exec_task` populated. List may be empty if
no nodes are ready for execution.
Returns:
A `list` of tasks to execute.
"""
layers = topsort.topsorted_layers(
[node.pipeline_node for node in self._pipeline.nodes],
get_node_id_fn=lambda node: node.node_info.id,
get_parent_nodes=(
lambda node: [self._node_map[n] for n in node.upstream_nodes]),
get_child_nodes=(
lambda node:
[self._node_map[n] for n in node.downstream_nodes]))
result = []
for nodes in layers:
# Boolean that's set if there's at least one successfully executed node
# in the current layer.
executed_nodes = False
for node in nodes:
if node.node_info.id in self._ignore_node_ids:
logging.info(
'Ignoring node for task generation: %s',
task_lib.NodeUid.from_pipeline_node(
self._pipeline, node))
continue
# If a task for the node is already tracked by the task queue, it need
# not be considered for generation again.
if self._is_task_id_tracked_fn(
task_lib.exec_node_task_id_from_pipeline_node(
self._pipeline, node)):
continue
executions = task_gen_utils.get_executions(
self._mlmd_handle, node)
if (executions
and task_gen_utils.is_latest_execution_successful(
executions)):
executed_nodes = True
continue
# If all upstream nodes are executed but current node is not executed,
# the node is deemed ready for execution.
if self._upstream_nodes_executed(node):
task = self._generate_task(node)
if task:
result.append(task)
# If there are no executed nodes in the current layer, downstream nodes
# need not be checked.
if not executed_nodes:
break
return result
def _upstream_nodes_executed(self, node: pipeline_pb2.PipelineNode) -> bool:
"""Returns `True` if all the upstream nodes have been successfully executed."""
upstream_nodes = [
node for node_id, node in self._node_map.items()
if node_id in set(node.upstream_nodes)
]
if not upstream_nodes:
return True
for node in upstream_nodes:
upstream_node_executions = task_gen_utils.get_executions(
self._mlmd_handle, node)
if not task_gen_utils.is_latest_execution_successful(
upstream_node_executions):
return False
return True
def generate(self) -> List[task_pb2.Task]:
"""Generates tasks for executing the next executable nodes in the pipeline.
The returned tasks must have `exec_task` populated. List may be empty if
no nodes are ready for execution.
Returns:
A `list` of tasks to execute.
"""
layers = topsort.topsorted_layers(
[node.pipeline_node for node in self._pipeline.nodes],
get_node_id_fn=lambda node: node.node_info.id,
get_parent_nodes=(
lambda node: [self._node_map[n] for n in node.upstream_nodes]),
get_child_nodes=(
lambda node:
[self._node_map[n] for n in node.downstream_nodes]))
tasks = []
with self._mlmd_connection as m:
# TODO(goutham): Cache executions and/or use TaskQueue so that we don't
# have to make MLMD queries for upstream nodes in each iteration.
for nodes in layers:
# Boolean that's set if there's at least one successfully executed node
# in the current layer.
executed_nodes = False
for node in nodes:
executions = task_gen_utils.get_executions(m, node)
if (executions
and task_gen_utils.is_latest_execution_successful(
executions)):
executed_nodes = True
continue
# If all upstream nodes are executed but current node is not executed,
# the node is deemed ready for execution.
if self._upstream_nodes_executed(m, node):
task = self._generate_task(m, node)
if task:
tasks.append(task)
# If there are no executed nodes in the current layer, downstream nodes
# need not be checked.
if not executed_nodes:
break
return tasks
def _upstream_nodes_executed(self,
cur_node: pipeline_pb2.PipelineNode) -> bool:
"""Returns `True` if all the upstream nodes have been successfully executed."""
upstream_nodes = [
node for node_id, node in self._node_map.items()
if node_id in set(cur_node.upstream_nodes)
]
for node in upstream_nodes:
if self._service_job_manager.is_pure_service_node(
self._pipeline_state, node.node_info.id):
service_status = self._service_job_manager.ensure_node_services(
self._pipeline_state, node.node_info.id)
if service_status == service_jobs.ServiceStatus.SUCCESS:
continue
else:
return False
node_executions = task_gen_utils.get_executions(
self._mlmd_handle, node)
if not task_gen_utils.is_latest_execution_successful(
node_executions):
return False
return True
def generate(self) -> List[task_lib.Task]:
"""Generates tasks for executing the next executable nodes in the pipeline.
The returned tasks must have `exec_task` populated. List may be empty if
no nodes are ready for execution.
Returns:
A `list` of tasks to execute.
"""
layers = topsort.topsorted_layers(
[node.pipeline_node for node in self._pipeline.nodes],
get_node_id_fn=lambda node: node.node_info.id,
get_parent_nodes=(
lambda node: [self._node_map[n] for n in node.upstream_nodes]),
get_child_nodes=(
lambda node:
[self._node_map[n] for n in node.downstream_nodes]))
result = []
for layer_num, nodes in enumerate(layers):
# Boolean that's set if there's at least one successfully executed node
# in the current layer.
completed_node_ids = set()
for node in nodes:
node_uid = task_lib.NodeUid.from_pipeline_node(
self._pipeline, node)
node_id = node.node_info.id
if self._service_job_manager.is_pure_service_node(
self._pipeline_state, node.node_info.id):
if not self._upstream_nodes_executed(node):
continue
service_status = self._service_job_manager.ensure_node_services(
self._pipeline_state, node_id)
if service_status == service_jobs.ServiceStatus.SUCCESS:
logging.info('Service node completed successfully: %s',
node_uid)
completed_node_ids.add(node_id)
elif service_status == service_jobs.ServiceStatus.FAILED:
logging.error('Failed service node: %s', node_uid)
return [
task_lib.FinalizePipelineTask(
pipeline_uid=self._pipeline_state.pipeline_uid,
status=status_lib.Status(
code=status_lib.Code.ABORTED,
message=
(f'Aborting pipeline execution due to service '
f'node failure; failed node uid: {node_uid}'
)))
]
else:
logging.info('Pure service node in progress: %s',
node_uid)
continue
# If a task for the node is already tracked by the task queue, it need
# not be considered for generation again.
if self._is_task_id_tracked_fn(
task_lib.exec_node_task_id_from_pipeline_node(
self._pipeline, node)):
continue
executions = task_gen_utils.get_executions(
self._mlmd_handle, node)
if (executions
and task_gen_utils.is_latest_execution_successful(
executions)):
completed_node_ids.add(node_id)
continue
# If all upstream nodes are executed but current node is not executed,
# the node is deemed ready for execution.
if self._upstream_nodes_executed(node):
task = self._generate_task(node)
if task_lib.is_finalize_pipeline_task(task):
return [task]
else:
result.append(task)
# If there are no completed nodes in the current layer, downstream nodes
# need not be checked.
if not completed_node_ids:
break
# If all nodes in the final layer are completed successfully , the
# pipeline can be finalized.
# TODO(goutham): If there are conditional eval nodes, not all nodes may be
# executed in the final layer. Handle this case when conditionals are
# supported.
if layer_num == len(layers) - 1 and completed_node_ids == set(
node.node_info.id for node in nodes):
return [
task_lib.FinalizePipelineTask(
pipeline_uid=self._pipeline_state.pipeline_uid,
status=status_lib.Status(code=status_lib.Code.OK))
]
return result
