def set_gpu_limit(self, gpu: str) -> 'PipelineTask':
"""Set gpu limit (maximum) for this operator.
Args:
gpu(str): Positive number required for number of GPUs.
Returns:
Self return to allow chained setting calls.
"""
if re.match(r'[1-9]\d*$', gpu) is None:
raise ValueError('GPU must be positive integer.')
gpu = int(gpu)
if self.container_spec is None:
raise ValueError(
'There is no container specified in implementation')
if self.container_spec.resources is not None:
self.container_spec.resources.accelerator_count = gpu
else:
self.container_spec.resources = structures.ResourceSpec(
accelerator_count=gpu)
return self
def set_cpu_limit(self, cpu: str) -> 'PipelineTask':
"""Set cpu limit (maximum) for this operator.
Args:
cpu(str): A string which can be a
number or a number followed by "m", whichmeans 1/1000.
Returns:
Self return to allow chained setting calls.
"""
if re.match(r'([0-9]*[.])?[0-9]+m?$', cpu) is None:
raise ValueError(
'Invalid cpu string. Should be float or integer, or integer'
' followed by "m".')
if cpu.endswith('m'):
cpu = float(cpu[:-1]) / 1000
else:
cpu = float(cpu)
if self.container_spec is None:
raise ValueError(
'There is no container specified in implementation')
if self.container_spec.resources is not None:
self.container_spec.resources.cpu_limit = cpu
else:
self.container_spec.resources = structures.ResourceSpec(
cpu_limit=cpu)
return self
def set_memory_limit(self, memory: str) -> 'PipelineTask':
"""Set memory limit (maximum) for this operator.
Args:
memory(str): a string which can be a number or a number followed by
one of "E", "Ei", "P", "Pi", "T", "Ti", "G", "Gi", "M", "Mi",
"K", "Ki".
Returns:
Self return to allow chained setting calls.
"""
if re.match(r'^[0-9]+(E|Ei|P|Pi|T|Ti|G|Gi|M|Mi|K|Ki){0,1}$',
memory) is None:
raise ValueError(
'Invalid memory string. Should be a number or a number '
'followed by one of "E", "Ei", "P", "Pi", "T", "Ti", "G", '
'"Gi", "M", "Mi", "K", "Ki".')
if memory.endswith('E'):
memory = float(memory[:-1]) * constants._E / constants._G
elif memory.endswith('Ei'):
memory = float(memory[:-2]) * constants._EI / constants._G
elif memory.endswith('P'):
memory = float(memory[:-1]) * constants._P / constants._G
elif memory.endswith('Pi'):
memory = float(memory[:-2]) * constants._PI / constants._G
elif memory.endswith('T'):
memory = float(memory[:-1]) * constants._T / constants._G
elif memory.endswith('Ti'):
memory = float(memory[:-2]) * constants._TI / constants._G
elif memory.endswith('G'):
memory = float(memory[:-1])
elif memory.endswith('Gi'):
memory = float(memory[:-2]) * constants._GI / constants._G
elif memory.endswith('M'):
memory = float(memory[:-1]) * constants._M / constants._G
elif memory.endswith('Mi'):
memory = float(memory[:-2]) * constants._MI / constants._G
elif memory.endswith('K'):
memory = float(memory[:-1]) * constants._K / constants._G
elif memory.endswith('Ki'):
memory = float(memory[:-2]) * constants._KI / constants._G
else:
# By default interpret as a plain integer, in the unit of Bytes.
memory = float(memory) / constants._G
if self.container_spec is None:
raise ValueError(
'There is no container specified in implementation')
if self.container_spec.resources is not None:
self.container_spec.resources.memory_limit = memory
else:
self.container_spec.resources = structures.ResourceSpec(
memory_limit=memory)
return self
def test_add_node_selector_constraint_accelerator_count(self):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
arguments={'input1': 'value'},
)
task.set_gpu_limit('5').add_node_selector_constraint('TPU_V3')
self.assertEqual(
structures.ResourceSpec(accelerator_type='TPU_V3',
accelerator_count=5),
task.container_spec.resources)
def test_add_node_selector_constraint_type_only(self):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
arguments={'input1': 'value'},
)
task.add_node_selector_constraint('NVIDIA_TESLA_K80')
self.assertEqual(
structures.ResourceSpec(accelerator_type='NVIDIA_TESLA_K80',
accelerator_count=1),
task.container_spec.resources)
def add_node_selector_constraint(self, accelerator: str) -> 'PipelineTask':
"""Sets accelerator type requirement for this task.
Args:
value(str): The name of the accelerator. Available values include
'NVIDIA_TESLA_K80', 'TPU_V3'.
Returns:
Self return to allow chained setting calls.
"""
if self.container_spec is None:
raise ValueError(
'There is no container specified in implementation')
if self.container_spec.resources is not None:
self.container_spec.resources.accelerator_type = accelerator
if self.container_spec.resources.accelerator_count is None:
self.container_spec.resources.accelerator_count = 1
else:
self.container_spec.resources = structures.ResourceSpec(
accelerator_count=1, accelerator_type=accelerator)
return self