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