def test_check_gpu(self): with self.assertRaises(TypeError): utils.gpu_requested("cpu=1") self.assertFalse(utils.gpu_requested(None)) self.assertFalse(utils.gpu_requested({})) self.assertFalse(utils.gpu_requested({"cpu": 1})) self.assertFalse(utils.gpu_requested({"cpu": 1, "memory": 2})) self.assertTrue(utils.gpu_requested({"gpu": 1})) self.assertTrue(utils.gpu_requested({" gpu ": 1})) self.assertTrue(utils.gpu_requested({"GPU": 1})) self.assertTrue(utils.gpu_requested({"cpu": 1, "memory": 2, "gpu": 1}))
def to_dict(self): template = Template.to_dict(self) if not utils.gpu_requested(self.resources): if self.env is None: self.env = {} self.env.update(OVERWRITE_GPU_ENVS) template["script"] = self.script_dict() return template
def to_dict(self): template = Template.to_dict(self) if (not utils.gpu_requested(self.resources) and states._overwrite_nvidia_gpu_envs): if self.env is None: self.env = {} self.env.update(OVERWRITE_GPU_ENVS) template["script"] = self.script_dict() return template
def to_dict(self): template = Container.to_dict(self) if (not utils.gpu_requested(self.resources) and states._overwrite_nvidia_gpu_envs): if self.env is None: self.env = {} self.env.update(OVERWRITE_GPU_ENVS) if "container" in template: template["script"] = template.pop("container") template["script"].update(self.script_dict()) return template
def to_dict(self): template = Template.to_dict(self) # Inputs parameters = [] if self.args is not None: i = 0 for arg in self.args: if not isinstance(self.args[i], OutputArtifact): if isinstance(arg, OutputJob): for _ in range(3): parameters.append({ "name": utils.input_parameter_name(self.name, i) }) i += 1 else: para_name = utils.input_parameter_name(self.name, i) parameters.append({"name": para_name}) i += 1 # Input # Case 1: add the input parameter if len(parameters) > 0: template["inputs"] = OrderedDict() template["inputs"]["parameters"] = parameters # Case 2: add the input artifact if self.input is not None: _input_list = [] for o in self.input: if isinstance(o, TypedArtifact): _input_list.append(o.to_yaml()) if isinstance(o, OutputArtifact): name = o.artifact["name"] if not any(name == x["name"] for x in _input_list): _input_list.append(o.artifact) if len(_input_list) > 0: if "inputs" not in template: template["inputs"] = OrderedDict() template["inputs"]["artifacts"] = _input_list # Node selector if self.node_selector is not None: # TODO: Support inferring node selector values from Argo parameters template["nodeSelector"] = self.node_selector # Container if (not utils.gpu_requested(self.resources) and states._overwrite_nvidia_gpu_envs): if self.env is None: self.env = {} self.env.update(OVERWRITE_GPU_ENVS) template["container"] = self.container_dict() # Output if self.output is not None: _output_list = [] for o in self.output: _output_list.append(o.to_yaml()) if isinstance(o, TypedArtifact): # Require only one kind of output type template["outputs"] = {"artifacts": _output_list} else: template["outputs"] = {"parameters": _output_list} # Volume if self.volume_mounts is not None: if self.volumes is None: template["volumes"] = [] for volume_mount in self.volume_mounts: template["volumes"].append({ "name": volume_mount.name, "emptyDir": {} }) return template
def to_dict(self): template = Template.to_dict(self) # Inputs parameters = [] if self.args is not None: i = 0 for arg in self.args: if not isinstance(self.args[i], OutputArtifact): if isinstance(arg, OutputJob): for _ in range(3): parameters.append({ "name": utils.input_parameter_name(self.name, i) }) i += 1 else: para_name = utils.input_parameter_name(self.name, i) parameters.append({"name": para_name}) i += 1 # Input # Case 1: add the input parameter if len(parameters) > 0: template["inputs"] = OrderedDict() template["inputs"]["parameters"] = parameters # Case 2: add the input artifact if self.input is not None: _input_list = [] for o in self.input: if isinstance(o, TypedArtifact): _input_list.append(o.to_yaml()) if isinstance(o, OutputArtifact): _input_list.append(o.artifact) if len(_input_list) > 0: if "inputs" not in template: template["inputs"] = OrderedDict() template["inputs"]["artifacts"] = _input_list # Container if not utils.gpu_requested(self.resources): if self.env is None: self.env = {} self.env.update(OVERWRITE_GPU_ENVS) template["container"] = self.container_dict() # Output if self.output is not None: _output_list = [] for o in self.output: _output_list.append(o.to_yaml()) if isinstance(o, TypedArtifact): # Require only one kind of output type template["outputs"] = {"artifacts": _output_list} else: template["outputs"] = {"parameters": _output_list} return template