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