feat(sdk): rename kfp artifact attributes (support custom artifact types pt. 1)

sdk/python/kfp/compiler/pipeline_spec_builder.py

@@ -915,8 +915,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)
 

sdk/python/kfp/compiler/test_data/pipelines/pipeline_with_gcpc_types.py

@@ -20,7 +20,7 @@
 
 
 class VertexModel(dsl.Artifact):
-    TYPE_NAME = 'google.VertexModel'
+    schema_title = 'google.VertexModel'
 
 
 producer_op = components.load_component_from_text("""

sdk/python/kfp/compiler/test_data/pipelines/pipeline_with_importer_and_gcpc_types.py

@@ -21,7 +21,7 @@
 
 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("""

sdk/python/kfp/components/component_factory.py

@@ -142,24 +142,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]):

sdk/python/kfp/components/importer_node.py

@@ -55,24 +55,24 @@ def pipeline_with_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',
         implementation=structures.Implementation(
             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(

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 @@ def my_pipeline():
 
     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 @@ def __init__(self,
         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,

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 @@ def __exit__(self, *unused_args) -> None:
         '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',
 }
 

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'})