diff --git a/sdk/python/kfp/components/pipeline_task.py b/sdk/python/kfp/components/pipeline_task.py index 71760984590..be5c79b057b 100644 --- a/sdk/python/kfp/components/pipeline_task.py +++ b/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. diff --git a/sdk/python/kfp/components/pipeline_task_test.py b/sdk/python/kfp/components/pipeline_task_test.py index b8f477d1a55..b56452ed788 100644 --- a/sdk/python/kfp/components/pipeline_task_test.py +++ b/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), diff --git a/sdk/python/test_data/pipelines/pipeline_with_resource_spec.py b/sdk/python/test_data/pipelines/pipeline_with_resource_spec.py index 049fc6f5cd0..e10384493bc 100644 --- a/sdk/python/test_data/pipelines/pipeline_with_resource_spec.py +++ b/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__':