Skip to content

Commit

Permalink
rename method
Browse files Browse the repository at this point in the history
  • Loading branch information
connor-mccarthy committed Mar 13, 2023
1 parent 10c4db9 commit dc9a368
Showing 3 changed files with 37 additions and 10 deletions.
21 changes: 17 additions & 4 deletions sdk/python/kfp/components/pipeline_task.py
Original file line number Diff line number Diff line change
@@ -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'``.
value: The name of the accelerator, such as ``'NVIDIA_TESLA_K80'`` and ``'TPU_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:
value: The name of the accelerator, such as ``'NVIDIA_TESLA_K80'`` and ``'TPU_V3'``.
Returns:
Self return to allow chained setting calls.
22 changes: 18 additions & 4 deletions sdk/python/kfp/components/pipeline_task_test.py
Original file line number Diff line number Diff line change
@@ -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),
4 changes: 2 additions & 2 deletions sdk/python/test_data/pipelines/pipeline_with_resource_spec.py
Original file line number Diff line number Diff line change
@@ -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__':

0 comments on commit dc9a368

Please sign in to comment.