diff --git a/sdk/python/kfp/compiler/pipeline_spec_builder.py b/sdk/python/kfp/compiler/pipeline_spec_builder.py index 334be084c6..5087662aba 100644 --- a/sdk/python/kfp/compiler/pipeline_spec_builder.py +++ b/sdk/python/kfp/compiler/pipeline_spec_builder.py @@ -900,8 +900,8 @@ def populate_metrics_in_dag_outputs( if artifact_spec.artifact_type.WhichOneof( 'kind' ) == 'schema_title' and artifact_spec.artifact_type.schema_title in [ - artifact_types.Metrics.TYPE_NAME, - artifact_types.ClassificationMetrics.TYPE_NAME, + artifact_types.Metrics.schema_title, + artifact_types.ClassificationMetrics.schema_title, ]: unique_output_name = '{}-{}'.format(task.name, output_name) diff --git a/sdk/python/kfp/compiler/test_data/pipelines/pipeline_with_gcpc_types.py b/sdk/python/kfp/compiler/test_data/pipelines/pipeline_with_gcpc_types.py index f2e17e19de..aeb1aef88a 100644 --- a/sdk/python/kfp/compiler/test_data/pipelines/pipeline_with_gcpc_types.py +++ b/sdk/python/kfp/compiler/test_data/pipelines/pipeline_with_gcpc_types.py @@ -20,7 +20,7 @@ from kfp.dsl import Input class VertexModel(dsl.Artifact): - TYPE_NAME = 'google.VertexModel' + schema_title = 'google.VertexModel' producer_op = components.load_component_from_text(""" diff --git a/sdk/python/kfp/compiler/test_data/pipelines/pipeline_with_importer_and_gcpc_types.py b/sdk/python/kfp/compiler/test_data/pipelines/pipeline_with_importer_and_gcpc_types.py index c3b1a18ccd..98431d362c 100644 --- a/sdk/python/kfp/compiler/test_data/pipelines/pipeline_with_importer_and_gcpc_types.py +++ b/sdk/python/kfp/compiler/test_data/pipelines/pipeline_with_importer_and_gcpc_types.py @@ -21,7 +21,7 @@ from kfp.dsl import importer class VertexDataset(dsl.Artifact): """An artifact representing a GCPC Vertex Dataset.""" - TYPE_NAME = 'google.VertexDataset' + schema_title = 'google.VertexDataset' consumer_op = components.load_component_from_text(""" diff --git a/sdk/python/kfp/components/component_factory.py b/sdk/python/kfp/components/component_factory.py index ab5056f319..d021628122 100644 --- a/sdk/python/kfp/components/component_factory.py +++ b/sdk/python/kfp/components/component_factory.py @@ -143,24 +143,24 @@ def _annotation_to_type_struct(annotation): return type_struct if issubclass(annotation, artifact_types.Artifact - ) and not annotation.TYPE_NAME.startswith('system.'): + ) and not annotation.schema_title.startswith('system.'): # For artifact classes not under the `system` namespace, - # use its TYPE_NAME as-is. - type_name = annotation.TYPE_NAME + # use its schema_title as-is. + schema_title = annotation.schema_title else: - type_name = str(annotation.__name__) + schema_title = str(annotation.__name__) elif hasattr(annotation, '__forward_arg__'): # Handling typing.ForwardRef('Type_name') - type_name = str(annotation.__forward_arg__) + schema_title = str(annotation.__forward_arg__) else: - type_name = str(annotation) + schema_title = str(annotation) # It's also possible to get the converter by type name - type_struct = type_utils.get_canonical_type_name_for_type(type_name) + type_struct = type_utils.get_canonical_type_name_for_type(schema_title) if type_struct: return type_struct - return type_name + return schema_title def _maybe_make_unique(name: str, names: List[str]): diff --git a/sdk/python/kfp/components/importer_node.py b/sdk/python/kfp/components/importer_node.py index 93a119e722..89c1f9c90e 100644 --- a/sdk/python/kfp/components/importer_node.py +++ b/sdk/python/kfp/components/importer_node.py @@ -55,12 +55,12 @@ def importer( train(dataset=importer1.output) """ if issubclass(artifact_class, artifact_types.Artifact - ) and not artifact_class.TYPE_NAME.startswith('system.'): + ) and not artifact_class.schema_title.startswith('system.'): # For artifact classes not under the `system` namespace, - # use its TYPE_NAME as-is. - type_name = artifact_class.TYPE_NAME + # use its schema_title as-is. + schema_title = artifact_class.schema_title else: - type_name = artifact_class.__name__ + schema_title = artifact_class.__name__ component_spec = structures.ComponentSpec( name='importer', @@ -68,11 +68,11 @@ def importer( importer=structures.ImporterSpec( artifact_uri=placeholders.InputValuePlaceholder( INPUT_KEY).to_placeholder_string(), - type_schema=artifact_class.TYPE_NAME, + type_schema=artifact_class.schema_title, reimport=reimport, metadata=metadata)), inputs={INPUT_KEY: structures.InputSpec(type='String')}, - outputs={OUTPUT_KEY: structures.OutputSpec(type=type_name)}, + outputs={OUTPUT_KEY: structures.OutputSpec(type=schema_title)}, ) importer = importer_component.ImporterComponent( diff --git a/sdk/python/kfp/components/types/artifact_types.py b/sdk/python/kfp/components/types/artifact_types.py index 0820a61037..7b783aa659 100644 --- a/sdk/python/kfp/components/types/artifact_types.py +++ b/sdk/python/kfp/components/types/artifact_types.py @@ -16,7 +16,7 @@ These are only compatible with v2 Pipelines. """ -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Type _GCS_LOCAL_MOUNT_PREFIX = '/gcs/' _MINIO_LOCAL_MOUNT_PREFIX = '/minio/' @@ -63,8 +63,8 @@ class Artifact: Note: Other artifacts are used similarly to the usage of ``Artifact`` in the example above (within ``Input[]`` and ``Output[]``). """ - TYPE_NAME = 'system.Artifact' - VERSION = '0.0.1' + schema_title = 'system.Artifact' + schema_version = '0.0.1' def __init__(self, name: Optional[str] = None, @@ -110,7 +110,7 @@ class Model(Artifact): uri: The model's location on disk or cloud storage. metadata: Arbitrary key-value pairs about the model. """ - TYPE_NAME = 'system.Model' + schema_title = 'system.Model' def __init__(self, name: Optional[str] = None, @@ -141,7 +141,7 @@ class Dataset(Artifact): uri: The dataset's location on disk or cloud storage. metadata: Arbitrary key-value pairs about the dataset. """ - TYPE_NAME = 'system.Dataset' + schema_title = 'system.Dataset' def __init__(self, name: Optional[str] = None, @@ -158,7 +158,7 @@ class Metrics(Artifact): uri: The metrics artifact's location on disk or cloud storage. metadata: Key-value scalar metrics. """ - TYPE_NAME = 'system.Metrics' + schema_title = 'system.Metrics' def __init__(self, name: Optional[str] = None, @@ -184,7 +184,7 @@ class ClassificationMetrics(Artifact): uri: The metrics artifact's location on disk or cloud storage. metadata: The key-value scalar metrics. """ - TYPE_NAME = 'system.ClassificationMetrics' + schema_title = 'system.ClassificationMetrics' def __init__(self, name: Optional[str] = None, @@ -347,7 +347,7 @@ class SlicedClassificationMetrics(Artifact): metadata: Arbitrary key-value pairs about the metrics artifact. """ - TYPE_NAME = 'system.SlicedClassificationMetrics' + schema_title = 'system.SlicedClassificationMetrics' def __init__(self, name: Optional[str] = None, @@ -472,7 +472,7 @@ class HTML(Artifact): uri: The HTML file's location on disk or cloud storage. metadata: Arbitrary key-value pairs about the HTML file. """ - TYPE_NAME = 'system.HTML' + schema_title = 'system.HTML' def __init__(self, name: Optional[str] = None, @@ -489,7 +489,7 @@ class Markdown(Artifact): uri: The markdown file's location on disk or cloud storage. metadata: Arbitrary key-value pairs about the markdown file. """ - TYPE_NAME = 'system.Markdown' + schema_title = 'system.Markdown' def __init__(self, name: Optional[str] = None, @@ -498,8 +498,8 @@ class Markdown(Artifact): super().__init__(uri=uri, name=name, metadata=metadata) -_SCHEMA_TITLE_TO_TYPE: Dict[str, Artifact] = { - x.TYPE_NAME: x for x in [ +_SCHEMA_TITLE_TO_TYPE: Dict[str, Type[Artifact]] = { + x.schema_title: x for x in [ Artifact, Model, Dataset, diff --git a/sdk/python/kfp/components/types/type_utils.py b/sdk/python/kfp/components/types/type_utils.py index a2bcc63aa7..35c7d153c2 100644 --- a/sdk/python/kfp/components/types/type_utils.py +++ b/sdk/python/kfp/components/types/type_utils.py @@ -124,8 +124,8 @@ def get_artifact_type_schema( artifact_class = artifact_class_or_type_name return pipeline_spec_pb2.ArtifactTypeSchema( - schema_title=artifact_class.TYPE_NAME, - schema_version=artifact_class.VERSION) + schema_title=artifact_class.schema_title, + schema_version=artifact_class.schema_version) def get_parameter_type( @@ -357,21 +357,21 @@ IR_TYPE_TO_IN_MEMORY_SPEC_TYPE = { 'Dict', 'BOOLEAN': 'Boolean', - artifact_types.Artifact.TYPE_NAME: + artifact_types.Artifact.schema_title: 'Artifact', - artifact_types.Model.TYPE_NAME: + artifact_types.Model.schema_title: 'Model', - artifact_types.Dataset.TYPE_NAME: + artifact_types.Dataset.schema_title: 'Dataset', - artifact_types.Metrics.TYPE_NAME: + artifact_types.Metrics.schema_title: 'Metrics', - artifact_types.ClassificationMetrics.TYPE_NAME: + artifact_types.ClassificationMetrics.schema_title: 'ClassificationMetrics', - artifact_types.SlicedClassificationMetrics.TYPE_NAME: + artifact_types.SlicedClassificationMetrics.schema_title: 'SlicedClassificationMetrics', - artifact_types.HTML.TYPE_NAME: + artifact_types.HTML.schema_title: 'HTML', - artifact_types.Markdown.TYPE_NAME: + artifact_types.Markdown.schema_title: 'Markdown', } diff --git a/sdk/python/kfp/components/types/type_utils_test.py b/sdk/python/kfp/components/types/type_utils_test.py index d24e8f35b4..24d16bc207 100644 --- a/sdk/python/kfp/components/types/type_utils_test.py +++ b/sdk/python/kfp/components/types/type_utils_test.py @@ -62,8 +62,8 @@ class _ArbitraryClass: class _VertexDummy(artifact_types.Artifact): - TYPE_NAME = 'google.VertexDummy' - VERSION = '0.0.2' + schema_title = 'google.VertexDummy' + schema_version = '0.0.2' def __init__(self): super().__init__(uri='uri', name='name', metadata={'dummy': '123'})