Repository navigation
Expand file tree
/
Copy pathtest_utils.py
More file actions
137 lines (123 loc) · 4.64 KB
/
Copy pathtest_utils.py
File metadata and controls
137 lines (123 loc) · 4.64 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0
import gc
import logging
import os
import re
import subprocess
import sys
from cuda.core import Device, system
from cuda.core._utils.cuda_utils import CUDAError
try:
import cupy as cp
except ImportError:
cp = None
try:
import matplotlib
except ImportError:
matplotlib = None
else:
# disable plot windows from popping out when testing locally
matplotlib.use("Agg")
import pytest
try:
DEVICE_COUNT = system.get_num_devices()
except CUDAError:
DEVICE_COUNT = 0
try:
cc = tuple(Device().compute_capability)
except RuntimeError:
cc = (0, 0)
class SampleTestError(Exception):
pass
def parse_python_script(filepath):
if filepath.endswith(".py"):
with open(filepath, encoding="utf-8") as f:
script = f.read()
else:
raise ValueError(f"{filepath} not supported")
return script
def run_sample(samples_path, filename, env=None, use_subprocess=False, use_mpi=False):
requires_mgpu = filename.endswith("_mgpu.py")
if DEVICE_COUNT == 0 and "cpu_execution" not in filename:
pytest.skip(f"Sample ({filename}) skipped because 0 CUDA devices were found.")
if requires_mgpu and DEVICE_COUNT == 1:
pytest.skip(f"Sample ({filename}) skipped due to limited device counts : ({DEVICE_COUNT})")
fullpath = os.path.join(samples_path, filename)
script = parse_python_script(fullpath)
if env is None:
env = {}
# Set __file__ so that example scripts can resolve relative paths
# (e.g., to the common directory)
env["__file__"] = fullpath
try:
old_argv = sys.argv
sys.argv = [fullpath]
SYS_PATH_BACKUP = sys.path.copy()
sys.path.append(samples_path)
ROOT_LOG_LEVEL_BACKUP = logging.getLogger().level
if use_mpi:
assert use_subprocess
# Check if the filename indicates with how many processes to run, for example:
# `example_something_4p.py` is to be run with 4 processes.
m = re.search(r".*_(\d+)p.py$", filename)
uses_nccl = (
"distributed/linalg/advanced/matmul" in fullpath
or "distributed/linalg/generic/direct_solver" in fullpath
or "nccl" in filename
)
if m:
num_procs = m.group(1)
if uses_nccl and int(num_procs) > DEVICE_COUNT:
pytest.skip(
f"This test requires {num_procs} processes but NCCL only allows one "
f"process per GPU and there are {DEVICE_COUNT} GPUs"
)
elif uses_nccl:
# For NCCL we can't have more processes than GPUs. And don't use more
# than 4 processes since the distributed examples aren't designed
# to run with more.
num_procs = str(min(4, DEVICE_COUNT))
else:
# Run with 2 processes by default.
num_procs = "2"
cmd = ["mpiexec", "-n", num_procs, sys.executable, fullpath]
else:
cmd = [sys.executable, fullpath]
if use_subprocess:
result = subprocess.run(cmd, capture_output=True, text=True, env=os.environ)
if result.returncode != 0:
if "ModuleNotFoundError" in result.stderr:
raise ModuleNotFoundError(result.stderr)
else:
raise RuntimeError(f"Subprocess failed: {result.stderr}")
else:
exec(script, env)
except ImportError as e:
# for samples requiring any of optional dependencies
for m in ("torch", "cupy"):
if f"No module named '{m}'" in str(e):
pytest.skip(f"{m} uninstalled, skipping related tests")
break
else:
raise
except Exception as e:
if str(e) == "libcudadevrt.a not found":
pytest.skip(f"Skipping test {filename} since libcudadevrt.a is not found.")
elif "requires libmathdx 0.3.2" in str(e):
pytest.skip(f"Skipping test {filename}: {e}")
else:
msg = "\n"
msg += f"Got error ({filename}):\n"
msg += str(e)
raise SampleTestError(msg) from e
finally:
sys.path = SYS_PATH_BACKUP
sys.argv = old_argv
logging.getLogger().setLevel(ROOT_LOG_LEVEL_BACKUP)
# further reduce the memory watermark
gc.collect()
if cp is not None:
cp.get_default_memory_pool().free_all_blocks()
cp.get_default_pinned_memory_pool().free_all_blocks()