def publish_execution( self, component_info: data_types.ComponentInfo, output_artifacts: Optional[Dict[Text, List[types.Artifact]]] = None, exec_properties: Optional[Dict[Text, Any]] = None): """Publishes a component execution to metadata. This function will do two things: 1. update the execution that was previously registered before execution to complete or skipped state, depending on whether cached results are used. 2. for each input and output artifact, publish an event that associate the artifact to the execution, with type INPUT or OUTPUT respectively Args: component_info: the information of the component output_artifacts: optional key -> Artifacts to be published as outputs of the execution exec_properties: optional execution properties to be published for the execution Returns: A dict containing output artifacts. """ outputs_utils.tag_output_artifacts_with_version(output_artifacts) logging.debug('Outputs: %s', output_artifacts) logging.debug('Execution properties: %s', exec_properties) self._metadata_handler.publish_execution( component_info=component_info, output_artifacts=output_artifacts, exec_properties=exec_properties)
def launch(self) -> Optional[metadata_store_pb2.Execution]: """Executes the component, includes driver, executor and publisher. Returns: The metadata of this execution that is registered in MLMD. It can be None if the driver decides not to run the execution. Raises: Exception: If the executor fails. """ logging.info('Running launcher for %s', self._pipeline_node) if self._system_node_handler: # If this is a system node, runs it and directly return. return self._system_node_handler.run(self._mlmd_connection, self._pipeline_node, self._pipeline_info, self._pipeline_runtime_spec) # Runs as a normal node. prepare_execution_result = self._prepare_execution() (execution_info, contexts, is_execution_needed) = (prepare_execution_result.execution_info, prepare_execution_result.contexts, prepare_execution_result.is_execution_needed) if is_execution_needed: try: executor_output = self._run_executor(execution_info) except Exception as e: # pylint: disable=broad-except execution_output = (e.executor_output if isinstance( e, _ExecutionFailedError) else None) self._publish_failed_execution(execution_info.execution_id, contexts, execution_output) logging.error('Execution %d failed.', execution_info.execution_id) raise finally: self._clean_up_stateless_execution_info(execution_info) logging.info('Execution %d succeeded.', execution_info.execution_id) self._clean_up_stateful_execution_info(execution_info) # TODO(b/182316162): Unify publisher handing so that post-execution # artifact logic is more cleanly handled. # Note that currently both the ExecutionInfo and ExecutorOutput are # consulted in `execution_publish_utils.publish_succeeded_execution()`. outputs_utils.tag_executor_output_with_version(executor_output) outputs_utils.tag_output_artifacts_with_version( execution_info.output_dict) logging.info('Publishing output artifacts %s for execution %s', execution_info.output_dict, execution_info.execution_id) self._publish_successful_execution(execution_info.execution_id, contexts, execution_info.output_dict, executor_output) return prepare_execution_result.execution_metadata
def _publish_execution_results(mlmd_handle: metadata.Metadata, task: task_lib.ExecNodeTask, result: ts.TaskSchedulerResult) -> None: """Publishes execution results to MLMD.""" def _update_state(status: status_lib.Status) -> None: assert status.code != status_lib.Code.OK if status.code == status_lib.Code.CANCELLED: logging.info('Cancelling execution (id: %s); task id: %s; status: %s', task.execution_id, task.task_id, status) execution_state = metadata_store_pb2.Execution.CANCELED else: logging.info( 'Aborting execution (id: %s) due to error (code: %s); task id: %s', task.execution_id, status.code, task.task_id) execution_state = metadata_store_pb2.Execution.FAILED _update_execution_state_in_mlmd(mlmd_handle, task.execution_id, execution_state, status.message) pipeline_state.record_state_change_time() if result.status.code != status_lib.Code.OK: _update_state(result.status) return # TODO(b/182316162): Unify publisher handing so that post-execution artifact # logic is more cleanly handled. outputs_utils.tag_output_artifacts_with_version(task.output_artifacts) publish_params = dict(output_artifacts=task.output_artifacts) if result.output_artifacts is not None: # TODO(b/182316162): Unify publisher handing so that post-execution artifact # logic is more cleanly handled. outputs_utils.tag_output_artifacts_with_version(result.output_artifacts) publish_params['output_artifacts'] = result.output_artifacts elif result.executor_output is not None: if result.executor_output.execution_result.code != status_lib.Code.OK: _update_state( status_lib.Status( code=result.executor_output.execution_result.code, message=result.executor_output.execution_result.result_message)) return # TODO(b/182316162): Unify publisher handing so that post-execution artifact # logic is more cleanly handled. outputs_utils.tag_executor_output_with_version(result.executor_output) publish_params['executor_output'] = result.executor_output execution_publish_utils.publish_succeeded_execution(mlmd_handle, task.execution_id, task.contexts, **publish_params) pipeline_state.record_state_change_time()
def _publish_execution_results(mlmd_handle: metadata.Metadata, task: task_lib.ExecNodeTask, result: ts.TaskSchedulerResult) -> None: """Publishes execution results to MLMD.""" def _update_state(status: status_lib.Status) -> None: assert status.code != status_lib.Code.OK _remove_output_dirs(task, result) _remove_task_dirs(task) if status.code == status_lib.Code.CANCELLED: logging.info( 'Cancelling execution (id: %s); task id: %s; status: %s', task.execution_id, task.task_id, status) execution_state = metadata_store_pb2.Execution.CANCELED else: logging.info( 'Aborting execution (id: %s) due to error (code: %s); task id: %s', task.execution_id, status.code, task.task_id) execution_state = metadata_store_pb2.Execution.FAILED _update_execution_state_in_mlmd(mlmd_handle, task.execution_id, execution_state, status.message) pipeline_state.record_state_change_time() if result.status.code != status_lib.Code.OK: _update_state(result.status) return # TODO(b/182316162): Unify publisher handing so that post-execution artifact # logic is more cleanly handled. outputs_utils.tag_output_artifacts_with_version(task.output_artifacts) if isinstance(result.output, ts.ExecutorNodeOutput): executor_output = result.output.executor_output if executor_output is not None: if executor_output.execution_result.code != status_lib.Code.OK: _update_state( status_lib.Status( code=executor_output.execution_result.code, message=executor_output.execution_result.result_message )) return # TODO(b/182316162): Unify publisher handing so that post-execution # artifact logic is more cleanly handled. outputs_utils.tag_executor_output_with_version(executor_output) _remove_task_dirs(task) execution_publish_utils.publish_succeeded_execution( mlmd_handle, execution_id=task.execution_id, contexts=task.contexts, output_artifacts=task.output_artifacts, executor_output=executor_output) elif isinstance(result.output, ts.ImporterNodeOutput): output_artifacts = result.output.output_artifacts # TODO(b/182316162): Unify publisher handing so that post-execution artifact # logic is more cleanly handled. outputs_utils.tag_output_artifacts_with_version(output_artifacts) _remove_task_dirs(task) execution_publish_utils.publish_succeeded_execution( mlmd_handle, execution_id=task.execution_id, contexts=task.contexts, output_artifacts=output_artifacts) elif isinstance(result.output, ts.ResolverNodeOutput): resolved_input_artifacts = result.output.resolved_input_artifacts execution_publish_utils.publish_internal_execution( mlmd_handle, execution_id=task.execution_id, contexts=task.contexts, output_artifacts=resolved_input_artifacts) else: raise TypeError(f'Unable to process task scheduler result: {result}') pipeline_state.record_state_change_time()
def launch(self) -> Optional[data_types.ExecutionInfo]: """Executes the component, includes driver, executor and publisher. Returns: The metadata of this execution that is registered in MLMD. It can be None if the driver decides not to run the execution. Raises: Exception: If the executor fails. """ logging.info('Running launcher for %s', self._pipeline_node) if self._system_node_handler: # If this is a system node, runs it and directly return. return self._system_node_handler.run(self._mlmd_connection, self._pipeline_node, self._pipeline_info, self._pipeline_runtime_spec) # Runs as a normal node. execution_preparation_result = self._prepare_execution() (execution_info, contexts, is_execution_needed) = ( execution_preparation_result.execution_info, execution_preparation_result.contexts, execution_preparation_result.is_execution_needed) if is_execution_needed: try: executor_watcher = None if self._executor_operator: # Create an execution watcher and save an in memory copy of the # Execution object to execution to it. Launcher calls executor # operator in process, thus there won't be race condition between the # execution watcher and the launcher to write to MLMD. executor_watcher = execution_watcher.ExecutionWatcher( port=portpicker.pick_unused_port(), mlmd_connection=self._mlmd_connection, execution=execution_preparation_result. execution_metadata, creds=grpc.local_server_credentials()) self._executor_operator.with_execution_watcher( executor_watcher.address) executor_watcher.start() executor_output = self._run_executor(execution_info) except Exception as e: # pylint: disable=broad-except execution_output = (e.executor_output if isinstance( e, _ExecutionFailedError) else None) self._publish_failed_execution(execution_info.execution_id, contexts, execution_output) logging.error('Execution %d failed.', execution_info.execution_id) raise finally: self._clean_up_stateless_execution_info(execution_info) if executor_watcher: executor_watcher.stop() logging.info('Execution %d succeeded.', execution_info.execution_id) self._clean_up_stateful_execution_info(execution_info) # TODO(b/182316162): Unify publisher handing so that post-execution # artifact logic is more cleanly handled. # Note that currently both the ExecutionInfo and ExecutorOutput are # consulted in `execution_publish_utils.publish_succeeded_execution()`. outputs_utils.tag_executor_output_with_version(executor_output) outputs_utils.tag_output_artifacts_with_version( execution_info.output_dict) logging.info('Publishing output artifacts %s for execution %s', execution_info.output_dict, execution_info.execution_id) self._publish_successful_execution(execution_info.execution_id, contexts, execution_info.output_dict, executor_output) return execution_info
def _run_executor(args: argparse.Namespace, beam_args: List[str]) -> None: """Selects a particular executor and run it based on name. Args: args: --executor_class_path: The import path of the executor class. --json_serialized_invocation_args: Full JSON-serialized parameters for this execution. beam_args: Optional parameter that maps to the optional_pipeline_args parameter in the pipeline, which provides additional configuration options for apache-beam and tensorflow.logging. For more about the beam arguments please refer to: https://cloud.google.com/dataflow/docs/guides/specifying-exec-params """ logging.set_verbosity(logging.INFO) # Rehydrate inputs/outputs/exec_properties from the serialized metadata. executor_input = pipeline_spec_pb2.ExecutorInput() json_format.Parse( args.json_serialized_invocation_args, executor_input, ignore_unknown_fields=True) inputs_dict = executor_input.inputs.artifacts outputs_dict = executor_input.outputs.artifacts inputs_parameter = executor_input.inputs.parameters if fileio.exists(executor_input.outputs.output_file): # It has a driver that outputs the updated exec_properties in this file. with fileio.open(executor_input.outputs.output_file, 'rb') as output_meta_json: output_metadata = pipeline_spec_pb2.ExecutorOutput() json_format.Parse( output_meta_json.read(), output_metadata, ignore_unknown_fields=True) # Append/Overwrite exec_propertise. for k, v in output_metadata.parameters.items(): inputs_parameter[k].CopyFrom(v) name_from_id = {} inputs = kubeflow_v2_entrypoint_utils.parse_raw_artifact_dict( inputs_dict, name_from_id) outputs = kubeflow_v2_entrypoint_utils.parse_raw_artifact_dict( outputs_dict, name_from_id) exec_properties = kubeflow_v2_entrypoint_utils.parse_execution_properties( inputs_parameter) logging.info('Executor %s do: inputs: %s, outputs: %s, exec_properties: %s', args.executor_class_path, inputs, outputs, exec_properties) executor_cls = import_utils.import_class_by_path(args.executor_class_path) if issubclass(executor_cls, base_beam_executor.BaseBeamExecutor): executor_context = base_beam_executor.BaseBeamExecutor.Context( beam_pipeline_args=beam_args, unique_id='', tmp_dir='/tmp') else: executor_context = base_executor.BaseExecutor.Context( extra_flags=beam_args, unique_id='', tmp_dir='/tmp') executor = executor_cls(executor_context) logging.info('Starting executor') executor.Do(inputs, outputs, exec_properties) # TODO(b/182316162): Unify publisher handling so that post-execution artifact # logic is more cleanly handled. outputs_utils.tag_output_artifacts_with_version(outputs) # pylint: disable=protected-access # TODO(b/169583143): Remove this workaround when TFX migrates to use str-typed # id/name to identify artifacts. # Convert ModelBlessing artifact to use managed MLMD resource name. if (issubclass(executor_cls, evaluator_executor.Executor) and standard_component_specs.BLESSING_KEY in outputs): # Parse the parent prefix for managed MLMD resource name. kubeflow_v2_entrypoint_utils.refactor_model_blessing( artifact_utils.get_single_instance( outputs[standard_component_specs.BLESSING_KEY]), name_from_id) # Log the output metadata to a file. So that it can be picked up by MP. metadata_uri = executor_input.outputs.output_file executor_output = pipeline_spec_pb2.ExecutorOutput() for k, v in kubeflow_v2_entrypoint_utils.translate_executor_output( outputs, name_from_id).items(): executor_output.artifacts[k].CopyFrom(v) fileio.makedirs(os.path.dirname(metadata_uri)) with fileio.open(metadata_uri, 'wb') as f: f.write(json_format.MessageToJson(executor_output))
def run( self, mlmd_connection: metadata.Metadata, pipeline_node: pipeline_pb2.PipelineNode, pipeline_info: pipeline_pb2.PipelineInfo, pipeline_runtime_spec: pipeline_pb2.PipelineRuntimeSpec ) -> data_types.ExecutionInfo: """Runs Importer specific logic. Args: mlmd_connection: ML metadata connection. pipeline_node: The specification of the node that this launcher lauches. 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. Returns: The execution of the run. """ logging.info('Running as an importer node.') with mlmd_connection as m: # 1.Prepares all contexts. contexts = context_lib.prepare_contexts( metadata_handler=m, node_contexts=pipeline_node.contexts) # 2. Resolves execution properties, please note that importers has no # input. exec_properties = data_types_utils.build_parsed_value_dict( inputs_utils.resolve_parameters_with_schema( node_parameters=pipeline_node.parameters)) # 3. Registers execution in metadata. execution = execution_publish_utils.register_execution( metadata_handler=m, execution_type=pipeline_node.node_info.type, contexts=contexts, exec_properties=exec_properties) # 4. Generate output artifacts to represent the imported artifacts. output_spec = pipeline_node.outputs.outputs[ importer.IMPORT_RESULT_KEY] properties = self._extract_proto_map( output_spec.artifact_spec.additional_properties) custom_properties = self._extract_proto_map( output_spec.artifact_spec.additional_custom_properties) output_artifact_class = types.Artifact( output_spec.artifact_spec.type).type output_artifacts = importer.generate_output_dict( metadata_handler=m, uri=str(exec_properties[importer.SOURCE_URI_KEY]), properties=properties, custom_properties=custom_properties, reimport=bool(exec_properties[importer.REIMPORT_OPTION_KEY]), output_artifact_class=output_artifact_class, mlmd_artifact_type=output_spec.artifact_spec.type) result = data_types.ExecutionInfo(execution_id=execution.id, input_dict={}, output_dict=output_artifacts, exec_properties=exec_properties, pipeline_node=pipeline_node, pipeline_info=pipeline_info) # TODO(b/182316162): consider let the launcher level do the publish # for system nodes. So that the version taging logic doesn't need to be # handled per system node. outputs_utils.tag_output_artifacts_with_version(result.output_dict) # 5. Publish the output artifacts. If artifacts are reimported, the # execution is published as CACHED. Otherwise it is published as COMPLETE. if _is_artifact_reimported(output_artifacts): execution_publish_utils.publish_cached_execution( metadata_handler=m, contexts=contexts, execution_id=execution.id, output_artifacts=output_artifacts) else: execution_publish_utils.publish_succeeded_execution( metadata_handler=m, execution_id=execution.id, contexts=contexts, output_artifacts=output_artifacts) return result