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