def _generate_task_from_execution( metadata_handler: metadata.Metadata, pipeline: pipeline_pb2.Pipeline, node: pipeline_pb2.PipelineNode, execution: metadata_store_pb2.Execution) -> task_lib.Task: """Generates `ExecNodeTask` given execution.""" contexts = metadata_handler.store.get_contexts_by_execution(execution.id) exec_properties = _extract_properties(execution) input_artifacts = execution_lib.get_artifacts_dict( metadata_handler, execution.id, metadata_store_pb2.Event.INPUT) outputs_resolver = outputs_utils.OutputsResolver(node, pipeline.pipeline_info, pipeline.runtime_spec, pipeline.execution_mode) return task_lib.ExecNodeTask( node_uid=task_lib.NodeUid.from_pipeline_node(pipeline, node), execution=execution, contexts=contexts, exec_properties=exec_properties, input_artifacts=input_artifacts, output_artifacts=outputs_resolver.generate_output_artifacts( execution.id), executor_output_uri=outputs_resolver.get_executor_output_uri( execution.id), stateful_working_dir=outputs_resolver.get_stateful_working_directory( execution.id))
def _generate_task( self, node: pipeline_pb2.PipelineNode, node_executions: Sequence[metadata_store_pb2.Execution] ) -> task_lib.Task: """Generates a node execution task. If node execution is not feasible, `None` is returned. Args: node: The pipeline node for which to generate a task. node_executions: Node executions fetched from MLMD. Returns: Returns an `ExecNodeTask` if node can be executed. If an error occurs, a `FinalizePipelineTask` is returned to abort the pipeline execution. """ result = task_gen_utils.generate_task_from_active_execution( self._mlmd_handle, self._pipeline, node, node_executions) if result: return result node_uid = task_lib.NodeUid.from_pipeline_node(self._pipeline, node) resolved_info = task_gen_utils.generate_resolved_info( self._mlmd_handle, node) if resolved_info.input_artifacts is None: 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 failure to resolve ' f'inputs; problematic node uid: {node_uid}'))) execution = execution_publish_utils.register_execution( metadata_handler=self._mlmd_handle, execution_type=node.node_info.type, contexts=resolved_info.contexts, input_artifacts=resolved_info.input_artifacts, exec_properties=resolved_info.exec_properties) outputs_resolver = outputs_utils.OutputsResolver( node, self._pipeline.pipeline_info, self._pipeline.runtime_spec, self._pipeline.execution_mode) return task_lib.ExecNodeTask( node_uid=node_uid, execution=execution, contexts=resolved_info.contexts, input_artifacts=resolved_info.input_artifacts, exec_properties=resolved_info.exec_properties, output_artifacts=outputs_resolver.generate_output_artifacts( execution.id), executor_output_uri=outputs_resolver.get_executor_output_uri( execution.id), stateful_working_dir=outputs_resolver. get_stateful_working_directory(execution.id), pipeline=self._pipeline)
def setUp(self): super().setUp() pipeline_runtime_spec = pipeline_pb2.PipelineRuntimeSpec() pipeline_runtime_spec.pipeline_root.field_value.string_value = self.tmp_dir pipeline_runtime_spec.pipeline_run_id.field_value.string_value = ( 'test_run_0') self._output_resolver = outputs_utils.OutputsResolver( pipeline_node=_PIPLINE_NODE, pipeline_info=_PIPELING_INFO, pipeline_runtime_spec=pipeline_runtime_spec)
def _generate_task( self, metadata_handler: metadata.Metadata, node: pipeline_pb2.PipelineNode) -> Optional[task_lib.Task]: """Generates a node execution task. If node execution is not feasible, `None` is returned. Args: metadata_handler: A handler to access MLMD db. node: The pipeline node for which to generate a task. Returns: Returns a `Task` or `None` if task generation is deemed infeasible. """ if not task_gen_utils.is_feasible_node(node): return None executions = task_gen_utils.get_executions(metadata_handler, node) result = task_gen_utils.generate_task_from_active_execution( metadata_handler, self._pipeline, node, executions) if result: return result resolved_info = task_gen_utils.generate_resolved_info( metadata_handler, node) if resolved_info.input_artifacts is None: # TODO(goutham): If the pipeline can't make progress, there should be a # standard mechanism to surface it to the user. logging.warning( 'Task cannot be generated for node %s since no input artifacts ' 'are resolved.', node.node_info.id) return None execution = execution_publish_utils.register_execution( metadata_handler=metadata_handler, execution_type=node.node_info.type, contexts=resolved_info.contexts, input_artifacts=resolved_info.input_artifacts, exec_properties=resolved_info.exec_properties) outputs_resolver = outputs_utils.OutputsResolver( node, self._pipeline.pipeline_info, self._pipeline.runtime_spec, self._pipeline.execution_mode) return task_lib.ExecNodeTask( node_uid=task_lib.NodeUid.from_pipeline_node(self._pipeline, node), execution=execution, contexts=resolved_info.contexts, input_artifacts=resolved_info.input_artifacts, exec_properties=resolved_info.exec_properties, output_artifacts=outputs_resolver.generate_output_artifacts( execution.id), executor_output_uri=outputs_resolver.get_executor_output_uri( execution.id), stateful_working_dir=outputs_resolver. get_stateful_working_directory(execution.id))
def __init__(self, pipeline_node: pipeline_pb2.PipelineNode, mlmd_connection: metadata.Metadata, pipeline_info: pipeline_pb2.PipelineInfo, pipeline_runtime_spec: pipeline_pb2.PipelineRuntimeSpec, executor_spec: Optional[message.Message] = None, custom_driver_spec: Optional[message.Message] = None, platform_spec: Optional[message.Message] = None, custom_executor_operators: Optional[Dict[ Any, Type[ExecutorOperator]]] = None): """Initializes a Launcher. Args: pipeline_node: The specification of the node that this launcher lauches. mlmd_connection: ML metadata connection. pipeline_info: The information of the pipeline that this node runs in. pipeline_runtime_spec: The runtime information of the pipeline that this node runs in. executor_spec: Specification for the executor of the node. This is expected for all components nodes. This will be used to determine the specific ExecutorOperator class to be used to execute and will be passed into ExecutorOperator. custom_driver_spec: Specification for custom driver. This is expected only for advanced use cases. platform_spec: Platform config that will be used as auxiliary info of the node execution. This will be passed to ExecutorOperator along with the `executor_spec`. custom_executor_operators: a map of ExecutorSpec to its ExecutorOperation implementation. Raises: ValueError: when component and component_config are not launchable by the launcher. """ del custom_driver_spec self._pipeline_node = pipeline_node self._mlmd_connection = mlmd_connection self._pipeline_info = pipeline_info self._pipeline_runtime_spec = pipeline_runtime_spec self._executor_operators = {} self._executor_operators.update(DEFAULT_EXECUTOR_OPERATORS) self._executor_operators.update(custom_executor_operators or {}) self._executor_operator = self._executor_operators[type( executor_spec)](executor_spec, platform_spec) self._output_resolver = outputs_utils.OutputsResolver( pipeline_node=self._pipeline_node, pipeline_info=self._pipeline_info, pipeline_runtime_spec=self._pipeline_runtime_spec)
def __init__(self, pipeline_node: pipeline_pb2.PipelineNode, mlmd_connection: metadata.Metadata, pipeline_info: pipeline_pb2.PipelineInfo, pipeline_runtime_spec: pipeline_pb2.PipelineRuntimeSpec, custom_executor_operators: Optional[Dict[ Any, Type[ExecutorOperator]]] = None): """Initializes a Launcher. Args: pipeline_node: The specification of the node that this launcher lauches. mlmd_connection: ML metadata connection. The connection is expected to not be opened before launcher is initiated. pipeline_info: The information of the pipeline that this node runs in. pipeline_runtime_spec: The runtime information of the pipeline that this node runs in. custom_executor_operators: a map of ExcutorSpec to its ExecutorOperation implementation. Raises: ValueError: when component and component_config are not launchable by the launcher. """ self._pipeline_node = pipeline_node self._mlmd_connection = mlmd_connection self._pipeline_info = pipeline_info self._pipeline_runtime_spec = pipeline_runtime_spec self._executor_operators = {} self._executor_operators.update(DEFAULT_EXECUTOR_OPERATORS) self._executor_operators.update(custom_executor_operators or {}) executor_spec_name = self._pipeline_node.executor.WhichOneof('spec') self._executor_spec = getattr(self._pipeline_node.executor, executor_spec_name) self._executor_operator = self._executor_operators[type( self._executor_spec)](self._executor_spec) self._output_resolver = outputs_utils.OutputsResolver( pipeline_node=self._pipeline_node, pipeline_info=self._pipeline_info, pipeline_runtime_spec=self._pipeline_runtime_spec)
def _output_resolver(self, execution_mode=pipeline_pb2.Pipeline.SYNC): return outputs_utils.OutputsResolver( pipeline_node=_PIPELINE_NODE, pipeline_info=_PIPELINE_INFO, pipeline_runtime_spec=self._pipeline_runtime_spec, execution_mode=execution_mode)
def _resolve_inputs_and_generate_tasks_for_node( self, node: pipeline_pb2.PipelineNode, ) -> List[task_lib.Task]: """Generates tasks for a node by freshly resolving inputs.""" result = [] node_uid = task_lib.NodeUid.from_pipeline_node(self._pipeline, node) resolved_info = task_gen_utils.generate_resolved_info( self._mlmd_handle, node) if resolved_info is None: result.append( task_lib.UpdateNodeStateTask(node_uid=node_uid, state=pstate.NodeState.SKIPPED)) return result if not resolved_info.input_artifacts: error_msg = f'failure to resolve inputs; node uid: {node_uid}' result.append( task_lib.UpdateNodeStateTask(node_uid=node_uid, state=pstate.NodeState.FAILED, status=status_lib.Status( code=status_lib.Code.ABORTED, message=error_msg))) return result # TODO(b/207038460): Update sync pipeline to support ForEach. input_artifacts = resolved_info.input_artifacts[0] execution = execution_publish_utils.register_execution( metadata_handler=self._mlmd_handle, execution_type=node.node_info.type, contexts=resolved_info.contexts, input_artifacts=input_artifacts, exec_properties=resolved_info.exec_properties) outputs_resolver = outputs_utils.OutputsResolver( node, self._pipeline.pipeline_info, self._pipeline.runtime_spec, self._pipeline.execution_mode) output_artifacts = outputs_resolver.generate_output_artifacts( execution.id) # For mixed service nodes, we ensure node services and check service # status; pipeline is aborted if the service jobs have failed. service_status = self._ensure_node_services_if_mixed(node.node_info.id) if service_status == service_jobs.ServiceStatus.FAILED: error_msg = f'associated service job failed; node uid: {node_uid}' result.append( task_lib.UpdateNodeStateTask(node_uid=node_uid, state=pstate.NodeState.FAILED, status=status_lib.Status( code=status_lib.Code.ABORTED, message=error_msg))) return result outputs_utils.make_output_dirs(output_artifacts) result.append( task_lib.UpdateNodeStateTask(node_uid=node_uid, state=pstate.NodeState.RUNNING)) result.append( task_lib.ExecNodeTask( node_uid=node_uid, execution_id=execution.id, contexts=resolved_info.contexts, input_artifacts=input_artifacts, exec_properties=resolved_info.exec_properties, output_artifacts=output_artifacts, executor_output_uri=outputs_resolver.get_executor_output_uri( execution.id), stateful_working_dir=outputs_resolver. get_stateful_working_directory(execution.id), tmp_dir=outputs_resolver.make_tmp_dir(execution.id), pipeline=self._pipeline)) return result
def _maybe_generate_task( self, node: pipeline_pb2.PipelineNode, node_executions: Sequence[metadata_store_pb2.Execution], successful_node_ids: MutableSet[str], ) -> Optional[task_lib.Task]: """Generates a task to execute or `None` if no action is required. If node is executable, `ExecNodeTask` is returned. If node execution is infeasible due to unsatisfied preconditions such as missing inputs or service job failure, task to abort the pipeline is returned. If cache is enabled and previously computed outputs are found, those outputs are used to finish the execution. Since node execution can be elided, `None` is returned after adding the node_id to `successful_node_ids` set. Args: node: The pipeline node for which to generate a task. node_executions: Node executions fetched from MLMD. successful_node_ids: Set that tracks successful node ids. Returns: Returns an `ExecNodeTask` if node can be executed. If an error occurs, a `FinalizePipelineTask` is returned to abort the pipeline execution. """ result = task_gen_utils.generate_task_from_active_execution( self._mlmd_handle, self._pipeline, node, node_executions) if result: return result node_uid = task_lib.NodeUid.from_pipeline_node(self._pipeline, node) resolved_info = task_gen_utils.generate_resolved_info( self._mlmd_handle, node) if resolved_info.input_artifacts is None: return self._abort_task( f'failure to resolve inputs; node uid {node_uid}') execution = execution_publish_utils.register_execution( metadata_handler=self._mlmd_handle, execution_type=node.node_info.type, contexts=resolved_info.contexts, input_artifacts=resolved_info.input_artifacts, exec_properties=resolved_info.exec_properties) outputs_resolver = outputs_utils.OutputsResolver( node, self._pipeline.pipeline_info, self._pipeline.runtime_spec, self._pipeline.execution_mode) output_artifacts = outputs_resolver.generate_output_artifacts( execution.id) # Check if we can elide node execution by reusing previously computed # outputs for the node. cache_context = cache_utils.get_cache_context( self._mlmd_handle, pipeline_node=node, pipeline_info=self._pipeline.pipeline_info, executor_spec=_get_executor_spec(self._pipeline, node.node_info.id), input_artifacts=resolved_info.input_artifacts, output_artifacts=output_artifacts, parameters=resolved_info.exec_properties) contexts = resolved_info.contexts + [cache_context] if node.execution_options.caching_options.enable_cache: cached_outputs = cache_utils.get_cached_outputs( self._mlmd_handle, cache_context=cache_context) if cached_outputs is not None: logging.info( 'Eliding node execution, using cached outputs; node uid: %s', node_uid) execution_publish_utils.publish_cached_execution( self._mlmd_handle, contexts=contexts, execution_id=execution.id, output_artifacts=cached_outputs) successful_node_ids.add(node.node_info.id) pstate.record_state_change_time() return None # For mixed service nodes, we ensure node services and check service # status; pipeline is aborted if the service jobs have failed. service_status = self._ensure_node_services_if_mixed(node.node_info.id) if service_status == service_jobs.ServiceStatus.FAILED: return self._abort_task( f'associated service job failed; node uid: {node_uid}') return task_lib.ExecNodeTask( node_uid=node_uid, execution_id=execution.id, contexts=contexts, input_artifacts=resolved_info.input_artifacts, exec_properties=resolved_info.exec_properties, output_artifacts=output_artifacts, executor_output_uri=outputs_resolver.get_executor_output_uri( execution.id), stateful_working_dir=outputs_resolver. get_stateful_working_directory(execution.id), pipeline=self._pipeline)
def _generate_task( self, metadata_handler: metadata.Metadata, node: pipeline_pb2.PipelineNode) -> Optional[task_lib.Task]: """Generates a node execution task. If a node execution is not feasible, `None` is returned. Args: metadata_handler: A handler to access MLMD db. node: The pipeline node for which to generate a task. Returns: Returns a `Task` or `None` if task generation is deemed infeasible. """ executions = task_gen_utils.get_executions(metadata_handler, node) result = task_gen_utils.generate_task_from_active_execution( metadata_handler, self._pipeline, node, executions) if result: return result resolved_info = task_gen_utils.generate_resolved_info( metadata_handler, node) if resolved_info.input_artifacts is None or not any( resolved_info.input_artifacts.values()): logging.info( 'Task cannot be generated for node %s since no input artifacts ' 'are resolved.', node.node_info.id) return None # If the latest successful execution had the same resolved input artifacts, # the component should not be triggered, so task is not generated. # TODO(b/170231077): This logic should be handled by the resolver when it's # implemented. Also, currently only the artifact ids of previous execution # are checked to decide if a new execution is warranted but it may also be # necessary to factor in the difference of execution properties. latest_exec = task_gen_utils.get_latest_successful_execution( executions) if latest_exec: artifact_ids_by_event_type = ( execution_lib.get_artifact_ids_by_event_type_for_execution_id( metadata_handler, latest_exec.id)) latest_exec_input_artifact_ids = artifact_ids_by_event_type.get( metadata_store_pb2.Event.INPUT, set()) current_exec_input_artifact_ids = set( a.id for a in itertools.chain( *resolved_info.input_artifacts.values())) if latest_exec_input_artifact_ids == current_exec_input_artifact_ids: return None node_uid = task_lib.NodeUid.from_pipeline_node(self._pipeline, node) execution = execution_publish_utils.register_execution( metadata_handler=metadata_handler, execution_type=node.node_info.type, contexts=resolved_info.contexts, input_artifacts=resolved_info.input_artifacts, exec_properties=resolved_info.exec_properties) outputs_resolver = outputs_utils.OutputsResolver( node, self._pipeline.pipeline_info, self._pipeline.runtime_spec, self._pipeline.execution_mode) # For mixed service nodes, we ensure node services and check service # status; the node is aborted if its service jobs have failed. service_status = self._ensure_node_services_if_mixed(node.node_info.id) if service_status is not None: if service_status != service_jobs.ServiceStatus.RUNNING: return self._abort_node_task(node_uid) return task_lib.ExecNodeTask( node_uid=node_uid, execution=execution, contexts=resolved_info.contexts, input_artifacts=resolved_info.input_artifacts, exec_properties=resolved_info.exec_properties, output_artifacts=outputs_resolver.generate_output_artifacts( execution.id), executor_output_uri=outputs_resolver.get_executor_output_uri( execution.id), stateful_working_dir=outputs_resolver. get_stateful_working_directory(execution.id), pipeline=self._pipeline)
def _generate_tasks_for_node( self, metadata_handler: metadata.Metadata, node: pipeline_pb2.PipelineNode) -> List[task_lib.Task]: """Generates a node execution task. If a node execution is not feasible, `None` is returned. Args: metadata_handler: A handler to access MLMD db. node: The pipeline node for which to generate a task. Returns: Returns a `Task` or `None` if task generation is deemed infeasible. """ result = [] node_uid = task_lib.NodeUid.from_pipeline_node(self._pipeline, node) executions = task_gen_utils.get_executions(metadata_handler, node) exec_node_task = task_gen_utils.generate_task_from_active_execution( metadata_handler, self._pipeline, node, executions) if exec_node_task: result.append( task_lib.UpdateNodeStateTask(node_uid=node_uid, state=pstate.NodeState.RUNNING)) result.append(exec_node_task) return result resolved_info = task_gen_utils.generate_resolved_info( metadata_handler, node) # TODO(b/207038460): Update async pipeline to support ForEach. if (resolved_info is None or not resolved_info.input_artifacts or resolved_info.input_artifacts[0] is None or not any(resolved_info.input_artifacts[0].values())): logging.info( 'Task cannot be generated for node %s since no input artifacts ' 'are resolved.', node.node_info.id) return result input_artifact = resolved_info.input_artifacts[0] executor_spec_fingerprint = hashlib.sha256() executor_spec = task_gen_utils.get_executor_spec( self._pipeline_state.pipeline, node.node_info.id) if executor_spec is not None: executor_spec_fingerprint.update( executor_spec.SerializeToString(deterministic=True)) resolved_info.exec_properties[ constants. EXECUTOR_SPEC_FINGERPRINT_KEY] = executor_spec_fingerprint.hexdigest( ) # If the latest execution had the same resolved input artifacts, execution # properties and executor specs, we should not trigger a new execution. latest_exec = task_gen_utils.get_latest_execution(executions) if latest_exec: artifact_ids_by_event_type = ( execution_lib.get_artifact_ids_by_event_type_for_execution_id( metadata_handler, latest_exec.id)) latest_exec_input_artifact_ids = artifact_ids_by_event_type.get( metadata_store_pb2.Event.INPUT, set()) current_exec_input_artifact_ids = set( a.id for a in itertools.chain(*input_artifact.values())) latest_exec_properties = task_gen_utils.extract_properties( latest_exec) current_exec_properties = resolved_info.exec_properties latest_exec_executor_spec_fp = latest_exec_properties[ constants.EXECUTOR_SPEC_FINGERPRINT_KEY] current_exec_executor_spec_fp = resolved_info.exec_properties[ constants.EXECUTOR_SPEC_FINGERPRINT_KEY] if (latest_exec_input_artifact_ids == current_exec_input_artifact_ids and _exec_properties_match(latest_exec_properties, current_exec_properties) and latest_exec_executor_spec_fp == current_exec_executor_spec_fp): result.append( task_lib.UpdateNodeStateTask( node_uid=node_uid, state=pstate.NodeState.STARTED)) return result execution = execution_publish_utils.register_execution( metadata_handler=metadata_handler, execution_type=node.node_info.type, contexts=resolved_info.contexts, input_artifacts=input_artifact, exec_properties=resolved_info.exec_properties) outputs_resolver = outputs_utils.OutputsResolver( node, self._pipeline.pipeline_info, self._pipeline.runtime_spec, self._pipeline.execution_mode) # For mixed service nodes, we ensure node services and check service # status; the node is aborted if its service jobs have failed. service_status = self._ensure_node_services_if_mixed(node.node_info.id) if service_status is not None: if service_status != service_jobs.ServiceStatus.RUNNING: error_msg = f'associated service job failed; node uid: {node_uid}' result.append( task_lib.UpdateNodeStateTask( node_uid=node_uid, state=pstate.NodeState.FAILED, status=status_lib.Status(code=status_lib.Code.ABORTED, message=error_msg))) return result output_artifacts = outputs_resolver.generate_output_artifacts( execution.id) outputs_utils.make_output_dirs(output_artifacts) result.append( task_lib.UpdateNodeStateTask(node_uid=node_uid, state=pstate.NodeState.RUNNING)) result.append( task_lib.ExecNodeTask( node_uid=node_uid, execution_id=execution.id, contexts=resolved_info.contexts, input_artifacts=input_artifact, exec_properties=resolved_info.exec_properties, output_artifacts=output_artifacts, executor_output_uri=outputs_resolver.get_executor_output_uri( execution.id), stateful_working_dir=outputs_resolver. get_stateful_working_directory(execution.id), tmp_dir=outputs_resolver.make_tmp_dir(execution.id), pipeline=self._pipeline)) return result