def schedule(self) -> task_scheduler.TaskSchedulerResult:
def _as_dict(proto_map) -> Dict[str, types.Property]:
return {
k: data_types_utils.get_value(v)
for k, v in proto_map.items()
}
pipeline_node = self.task.get_pipeline_node()
output_spec = pipeline_node.outputs.outputs[importer.IMPORT_RESULT_KEY]
properties = _as_dict(output_spec.artifact_spec.additional_properties)
custom_properties = _as_dict(
output_spec.artifact_spec.additional_custom_properties)
output_artifacts = importer.generate_output_dict(
metadata_handler=self.mlmd_handle,
uri=str(self.task.exec_properties[importer.SOURCE_URI_KEY]),
properties=properties,
custom_properties=custom_properties,
reimport=bool(
self.task.exec_properties[importer.REIMPORT_OPTION_KEY]),
output_artifact_class=types.Artifact(
output_spec.artifact_spec.type).type,
mlmd_artifact_type=output_spec.artifact_spec.type)
return task_scheduler.TaskSchedulerResult(
status=status_lib.Status(code=status_lib.Code.OK),
output=task_scheduler.ImporterNodeOutput(
output_artifacts=output_artifacts))
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 = inputs_utils.resolve_parameters(
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)
# 5. Publish the output artifacts.
execution_publish_utils.publish_succeeded_execution(
metadata_handler=m,
execution_id=execution.id,
contexts=contexts,
output_artifacts=output_artifacts)
return 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)
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