Skip to content

Commit 090b7e6

Browse files
tetrontetron
committed
Add test coverage for CUDA checks
1 parent a30b00b commit 090b7e6

File tree

6 files changed

+195
-5
lines changed

6 files changed

+195
-5
lines changed

‎cwltool/extensions-v1.1.yml‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,12 @@ $graph:
110110
cudaDeviceCountMin:
111111
type: ['null', int, cwl:Expression]
112112
default: 1
113-
doc: Minimum number of GPU devices to request, default 1.
113+
doc: |
114+
Minimum number of GPU devices to request. If not specified,
115+
same as `cudaDeviceCountMax`. If neither are specified,
116+
default 1.
114117
cudaDeviceCountMax:
115118
type: ['null', int, cwl:Expression]
116-
doc: Maximum number of GPU devices to request. If not specified, same as `deviceCountMin`.
119+
doc: |
120+
Maximum number of GPU devices to request. If not specified,
121+
same as `cudaDeviceCountMin`.

‎cwltool/extensions.yml‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -220,7 +220,12 @@ $graph:
220220
cudaDeviceCountMin:
221221
type: ['null', int, cwl:Expression]
222222
default: 1
223-
doc: Minimum number of GPU devices to request, default 1.
223+
doc: |
224+
Minimum number of GPU devices to request. If not specified,
225+
same as `cudaDeviceCountMax`. If neither are specified,
226+
default 1.
224227
cudaDeviceCountMax:
225228
type: ['null', int, cwl:Expression]
226-
doc: Maximum number of GPU devices to request. If not specified, same as `deviceCountMin`.
229+
doc: |
230+
Maximum number of GPU devices to request. If not specified,
231+
same as `cudaDeviceCountMin`.

‎tests/test_cuda.py‎

Lines changed: 135 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,24 @@
1+
import mock
12
import pytest
3+
from schema_salad.avro import schema
24

5+
from cwltool.builder import Builder
6+
from cwltool.context import LoadingContext, RuntimeContext
37
from cwltool.cuda import cuda_version_and_device_count
8+
from cwltool.errors import WorkflowException
9+
from cwltool.job import CommandLineJob
10+
from cwltool.load_tool import load_tool
411
from cwltool.main import main
12+
from cwltool.pathmapper import MapperEnt, PathMapper
13+
from cwltool.process import use_custom_schema, use_standard_schema
14+
from cwltool.stdfsaccess import StdFsAccess
15+
from cwltool.update import INTERNAL_VERSION, ORIGINAL_CWLVERSION
16+
from cwltool.utils import CWLObjectType
517

618
from .util import get_data, needs_docker, needs_singularity_3_or_newer
719

20+
from unittest.mock import MagicMock
21+
822
cuda_version = cuda_version_and_device_count()
923

1024

@@ -39,7 +53,127 @@ def test_cuda_singularity() -> None:
3953
def test_cuda_no_container() -> None:
4054
params = [
4155
"--enable-ext",
42-
"--singularity",
4356
get_data("tests/wf/nvidia-smi.cwl"),
4457
]
4558
assert main(params) == 0
59+
60+
61+
@pytest.mark.skipif(
62+
cuda_version[0] == "", reason="nvidia-smi required for CUDA not detected"
63+
)
64+
def test_cuda_cc_list() -> None:
65+
params = [
66+
"--enable-ext",
67+
get_data("tests/wf/nvidia-smi-cc.cwl"),
68+
]
69+
assert main(params) == 0
70+
71+
72+
def _makebuilder(cudaReq: CWLObjectType) -> Builder:
73+
return Builder(
74+
{},
75+
[],
76+
[],
77+
{},
78+
schema.Names(),
79+
[cudaReq],
80+
[],
81+
{"cudaDeviceCount": 1},
82+
None,
83+
None,
84+
StdFsAccess,
85+
StdFsAccess(""),
86+
None,
87+
0.1,
88+
False,
89+
False,
90+
False,
91+
"",
92+
"",
93+
"",
94+
"",
95+
INTERNAL_VERSION,
96+
"docker",
97+
)
98+
99+
100+
@mock.patch("subprocess.check_output")
101+
@mock.patch("os.makedirs")
102+
def test_cuda_job_setup_check(makedirs: MagicMock, check_output: MagicMock) -> None:
103+
104+
runtime_context = RuntimeContext({})
105+
106+
cudaReq: CWLObjectType = {
107+
"class": "http://commonwl.org/cwltool#CUDARequirement",
108+
"cudaVersionMin": "1.0",
109+
"cudaComputeCapability": "1.0",
110+
}
111+
builder = _makebuilder(cudaReq)
112+
113+
check_output.return_value = """
114+
<nvidia>
115+
<attached_gpus>1</attached_gpus>
116+
<cuda_version>1.0</cuda_version>
117+
</nvidia>
118+
"""
119+
120+
jb = CommandLineJob(builder, {}, PathMapper, [], [], "")
121+
jb._setup(runtime_context)
122+
123+
124+
@mock.patch("subprocess.check_output")
125+
@mock.patch("os.makedirs")
126+
def test_cuda_job_setup_check_err(makedirs: MagicMock, check_output: MagicMock) -> None:
127+
128+
runtime_context = RuntimeContext({})
129+
130+
cudaReq: CWLObjectType = {
131+
"class": "http://commonwl.org/cwltool#CUDARequirement",
132+
"cudaVersionMin": "2.0",
133+
"cudaComputeCapability": "1.0",
134+
}
135+
builder = _makebuilder(cudaReq)
136+
137+
check_output.return_value = """
138+
<nvidia>
139+
<attached_gpus>1</attached_gpus>
140+
<cuda_version>1.0</cuda_version>
141+
</nvidia>
142+
"""
143+
jb = CommandLineJob(builder, {}, PathMapper, [], [], "")
144+
with pytest.raises(WorkflowException):
145+
jb._setup(runtime_context)
146+
147+
148+
def test_cuda_eval_resource_range() -> None:
149+
with open(get_data("cwltool/extensions-v1.1.yml")) as res:
150+
use_custom_schema("v1.2", "http://commonwl.org/cwltool", res.read())
151+
152+
joborder = {} # type: CWLObjectType
153+
loadingContext = LoadingContext({"do_update": True})
154+
runtime_context = RuntimeContext({})
155+
156+
tool = load_tool(get_data("tests/wf/nvidia-smi-range.cwl"), loadingContext)
157+
builder = _makebuilder(tool.requirements[0])
158+
builder.job = joborder
159+
160+
resources = tool.evalResources(builder, runtime_context)
161+
162+
assert resources["cudaDeviceCount"] == 2
163+
164+
165+
def test_cuda_eval_resource_max() -> None:
166+
with open(get_data("cwltool/extensions-v1.1.yml")) as res:
167+
use_custom_schema("v1.2", "http://commonwl.org/cwltool", res.read())
168+
169+
joborder = {} # type: CWLObjectType
170+
loadingContext = LoadingContext({"do_update": True})
171+
runtime_context = RuntimeContext({})
172+
173+
tool = load_tool(get_data("tests/wf/nvidia-smi-max.cwl"), loadingContext)
174+
builder = _makebuilder(tool.requirements[0])
175+
builder.job = joborder
176+
177+
resources = tool.evalResources(builder, runtime_context)
178+
179+
assert resources["cudaDeviceCount"] == 4

‎tests/wf/nvidia-smi-cc.cwl‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
cwlVersion: v1.2
2+
class: CommandLineTool
3+
$namespaces:
4+
cwltool: "http://commonwl.org/cwltool#"
5+
requirements:
6+
cwltool:CUDARequirement:
7+
cudaVersionMin: "1.0"
8+
cudaComputeCapability: ["1.0", "2.0", "3.0"]
9+
cudaDeviceCountMin: $(inputs.gpus)
10+
inputs:
11+
gpus:
12+
type: int
13+
default: 1
14+
outputs: []
15+
# Assume this will exit non-zero (resulting in a failing test case) if
16+
# nvidia-smi doesn't detect any devices.
17+
baseCommand: "nvidia-smi"

‎tests/wf/nvidia-smi-max.cwl‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
cwlVersion: v1.2
2+
class: CommandLineTool
3+
$namespaces:
4+
cwltool: "http://commonwl.org/cwltool#"
5+
requirements:
6+
cwltool:CUDARequirement:
7+
cudaVersionMin: "1.0"
8+
cudaComputeCapability: "1.0"
9+
cudaDeviceCountMax: 4
10+
inputs: []
11+
outputs: []
12+
# Assume this will exit non-zero (resulting in a failing test case) if
13+
# nvidia-smi doesn't detect any devices.
14+
baseCommand: "nvidia-smi"

‎tests/wf/nvidia-smi-range.cwl‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
cwlVersion: v1.2
2+
class: CommandLineTool
3+
$namespaces:
4+
cwltool: "http://commonwl.org/cwltool#"
5+
requirements:
6+
cwltool:CUDARequirement:
7+
cudaVersionMin: "1.0"
8+
cudaComputeCapability: "1.0"
9+
cudaDeviceCountMin: 2
10+
cudaDeviceCountMax: 4
11+
inputs: []
12+
outputs: []
13+
# Assume this will exit non-zero (resulting in a failing test case) if
14+
# nvidia-smi doesn't detect any devices.
15+
baseCommand: "nvidia-smi"

0 commit comments

Comments
 (0)