Source code for kfp.dsl.importer_node

# Copyright 2020-2022 The Kubeflow Authors
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.
"""Utility function for building Importer Node spec."""

from typing import Any, Dict, Mapping, Optional, Type, Union

from kfp.dsl import importer_component
from kfp.dsl import pipeline_channel
from kfp.dsl import pipeline_task
from kfp.dsl import placeholders
from kfp.dsl import structures
from kfp.dsl import utils
from kfp.dsl.types import artifact_types
from kfp.dsl.types import type_utils

URI_KEY = 'uri'
OUTPUT_KEY = 'artifact'
METADATA_KEY = 'metadata'

[docs]def importer( artifact_uri: Union[pipeline_channel.PipelineParameterChannel, str], artifact_class: Type[artifact_types.Artifact], reimport: bool = False, metadata: Optional[Mapping[str, Any]] = None, ) -> pipeline_task.PipelineTask: """Imports an existing artifact for use in a downstream component. Args: artifact_uri: The URI of the artifact to import. artifact_class: The artifact class being imported. reimport: Whether to reimport the artifact. metadata: Properties of the artifact. Returns: A task with the artifact accessible via its ``.output`` attribute. Examples:: @dsl.pipeline(name='pipeline-with-importer') def pipeline_with_importer(): importer1 = importer( artifact_uri='gs://ml-pipeline-playground/shakespeare1.txt', artifact_class=Dataset, reimport=False) train(dataset=importer1.output) """ component_inputs: Dict[str, structures.InputSpec] = {} call_inputs: Dict[str, Any] = {} def traverse_dict_and_create_metadata_inputs(d: Any) -> Any: if isinstance(d, pipeline_channel.PipelineParameterChannel): reversed_call_inputs = { pipeline_param_chan: name for name, pipeline_param_chan in call_inputs.items() } # minimizes importer spec interface by not creating new # inputspec/parameters if the same input is used multiple places # in metadata unique_name = reversed_call_inputs.get( d, utils.make_name_unique_by_adding_index( METADATA_KEY, list(call_inputs), '-', ), ) call_inputs[unique_name] = d component_inputs[unique_name] = structures.InputSpec( type=d.channel_type) return placeholders.InputValuePlaceholder( input_name=unique_name)._to_string() elif isinstance(d, dict): # use this instead of list comprehension to ensure compiles are identical across Python versions res = {} for k, v in d.items(): new_k = traverse_dict_and_create_metadata_inputs(k) new_v = traverse_dict_and_create_metadata_inputs(v) res[new_k] = new_v return res elif isinstance(d, list): return [traverse_dict_and_create_metadata_inputs(el) for el in d] elif isinstance(d, str): # extract pipeline channels from f-strings, if any pipeline_channels = pipeline_channel.extract_pipeline_channels_from_any( d) # pass the channel back into the recursive function to create the placeholder, component inputs, and call inputs, then replace the channel with the placeholder for channel in pipeline_channels: input_placeholder = traverse_dict_and_create_metadata_inputs( channel) d = d.replace(channel.pattern, input_placeholder) return d else: return d metadata_with_placeholders = traverse_dict_and_create_metadata_inputs( metadata) component_spec = structures.ComponentSpec( name='importer', implementation=structures.Implementation( importer=structures.ImporterSpec( artifact_uri=placeholders.InputValuePlaceholder( URI_KEY)._to_string(), schema_title=type_utils.create_bundled_artifact_type( artifact_class.schema_title, artifact_class.schema_version), schema_version=artifact_class.schema_version, reimport=reimport, metadata=metadata_with_placeholders)), inputs={ URI_KEY: structures.InputSpec(type='String'), **component_inputs }, outputs={ OUTPUT_KEY: structures.OutputSpec( type=type_utils.create_bundled_artifact_type( artifact_class.schema_title, artifact_class.schema_version)) }, ) importer = importer_component.ImporterComponent( component_spec=component_spec) return importer(uri=artifact_uri, **call_inputs)