feat(sdk): deprecate .add_node_selector_constraint in favor of .set_accelerator_type

sdk/python/kfp/components/pipeline_task.py

@@ -275,7 +275,7 @@ def set_cpu_limit(self, cpu: str) -> 'PipelineTask':
 
     def set_accelerator_limit(self, limit: int) -> 'PipelineTask':
         """Sets accelerator limit (maximum) for the task. Only applies if
-        accelerator type is also set via .add_node_selector_constraint().
+        accelerator type is also set via .set_accelerator_type().
 
         Args:
             limit: Maximum number of accelerators allowed.
@@ -302,7 +302,7 @@ def set_accelerator_limit(self, limit: int) -> 'PipelineTask':
 
     def set_gpu_limit(self, gpu: str) -> 'PipelineTask':
         """Sets GPU limit (maximum) for the task. Only applies if accelerator
-        type is also set via .add_node_selector_constraint().
+        type is also set via .add_accelerator_type().
 
         Args:
             gpu: The maximum GPU reuqests allowed. This string should be a positive integer number of GPUs.
@@ -401,8 +401,21 @@ def add_node_selector_constraint(self, accelerator: str) -> 'PipelineTask':
         """Sets accelerator type to use when executing this task.
 
         Args:
-            value: The name of the accelerator. Available values include
-                ``'NVIDIA_TESLA_K80'`` and ``'TPU_V3'``.
+            accelerator: The name of the accelerator, such as ``'NVIDIA_TESLA_K80'``, ``'TPU_V3'``, ``'nvidia.com/gpu'`` or ``'cloud-tpus.google.com/v3'``.
+
+        Returns:
+            Self return to allow chained setting calls.
+        """
+        warnings.warn(
+            f'{self.add_node_selector_constraint.__name__!r} is deprecated. Please use {self.set_accelerator_type.__name__!r} instead.',
+            category=DeprecationWarning)
+        return self.set_accelerator_type(accelerator)
+
+    def set_accelerator_type(self, accelerator: str) -> 'PipelineTask':
+        """Sets accelerator type to use when executing this task.
+
+        Args:
+            accelerator: The name of the accelerator, such as ``'NVIDIA_TESLA_K80'``, ``'TPU_V3'``, ``'nvidia.com/gpu'`` or ``'cloud-tpus.google.com/v3'``.
 
         Returns:
             Self return to allow chained setting calls.

sdk/python/kfp/components/pipeline_task_test.py

@@ -191,6 +191,20 @@ def test_set_valid_gpu_limit(self, gpu_limit: str,
         self.assertEqual(expected_gpu_number,
                          task.container_spec.resources.accelerator_count)
 
+    def test_add_valid_node_selector_constraint(self):
+        task = pipeline_task.PipelineTask(
+            component_spec=structures.ComponentSpec.from_yaml_documents(
+                V2_YAML),
+            args={'input1': 'value'},
+        )
+        with self.assertWarnsRegex(
+                DeprecationWarning,
+                "'add_node_selector_constraint' is deprecated. Please use 'set_accelerator_type' instead."
+        ):
+            task.add_node_selector_constraint('TPU_V3')
+        self.assertEqual(task.container_spec.resources.accelerator_type,
+                         'TPU_V3')
+
     @parameterized.parameters(
         {
             'limit': '123',
@@ -276,25 +290,25 @@ def test_set_memory_limit(self, memory: str, expected_memory_number: int):
         self.assertEqual(expected_memory_number,
                          task.container_spec.resources.memory_limit)
 
-    def test_add_node_selector_constraint_type_only(self):
+    def test_set_accelerator_type_with_type_only(self):
         task = pipeline_task.PipelineTask(
             component_spec=structures.ComponentSpec.from_yaml_documents(
                 V2_YAML),
             args={'input1': 'value'},
         )
-        task.add_node_selector_constraint('NVIDIA_TESLA_K80')
+        task.set_accelerator_type('NVIDIA_TESLA_K80')
         self.assertEqual(
             structures.ResourceSpec(
                 accelerator_type='NVIDIA_TESLA_K80', accelerator_count=1),
             task.container_spec.resources)
 
-    def test_add_node_selector_constraint_accelerator_count(self):
+    def test_set_accelerator_type_with_accelerator_count(self):
         task = pipeline_task.PipelineTask(
             component_spec=structures.ComponentSpec.from_yaml_documents(
                 V2_YAML),
             args={'input1': 'value'},
         )
-        task.set_gpu_limit('5').add_node_selector_constraint('TPU_V3')
+        task.set_accelerator_limit('5').set_accelerator_type('TPU_V3')
         self.assertEqual(
             structures.ResourceSpec(
                 accelerator_type='TPU_V3', accelerator_count=5),

sdk/python/test_data/pipelines/pipeline_with_resource_spec.py

@@ -39,8 +39,8 @@ def my_pipeline(input_location: str = 'gs://test-bucket/pipeline_root',
         training_op(
             examples=ingestor.outputs['examples'],
             optimizer=optimizer,
-            n_epochs=n_epochs).set_cpu_limit('4').set_memory_limit('14Gi')
-        .add_node_selector_constraint('tpu-v3').set_accelerator_limit(1))
+            n_epochs=n_epochs).set_cpu_limit('4').set_memory_limit(
+                '14Gi').set_accelerator_type('tpu-v3').set_accelerator_limit(1))
 
 
 if __name__ == '__main__':