Skip to content

Commit 1ebe689

Browse files
committed
CuPy gets installed by Nox
1 parent a498702 commit 1ebe689

File tree

1 file changed

+14
-12
lines changed

1 file changed

+14
-12
lines changed

noxfile.py

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -235,25 +235,27 @@ def tests(session: Session) -> None:
235235
install_additional_warning = """
236236
Installation failed, this likely means that the required hardware or drivers are missing.
237237
Run without `-- additional-tests` to avoid this."""
238+
# install cuda-python
238239
try:
239240
session.install("cuda-python")
240241
except Exception as error:
241242
session.log(error)
242243
session.warn(install_additional_warning)
244+
# install cupy
245+
try:
243246
try:
244-
try:
245-
# based on the CUDA version, try installing the exact prebuilt cupy version
246-
cuda_cupy_version = f"cupy-cuda{''.join(cuda_version.split('.'))}"
247-
session.install(cuda_cupy_version)
248-
except Exception:
249-
# if the exact prebuilt is not available, try the more general prebuilt
250-
cuda_cupy_version_x = f"cupy-cuda{cuda_version.split('.')[0]}x"
251-
session.warn(f"CuPy exact prebuilt not available for {cuda_version}, trying {cuda_cupy_version_x}")
252-
session.install(cuda_cupy_version_x)
247+
# based on the CUDA version, try installing the exact prebuilt cupy version
248+
cuda_cupy_version = f"cupy-cuda{''.join(cuda_version.split('.'))}"
249+
session.install(cuda_cupy_version)
253250
except Exception:
254-
# if no compatible prebuilt wheel is found, try building CuPy ourselves
255-
session.warn(f"No prebuilt CuPy found for CUDA {cuda_version}, building from source...")
256-
session.install("cupy")
251+
# if the exact prebuilt is not available, try the more general prebuilt
252+
cuda_cupy_version_x = f"cupy-cuda{cuda_version.split('.')[0]}x"
253+
session.warn(f"CuPy exact prebuilt not available for {cuda_version}, trying {cuda_cupy_version_x}")
254+
session.install(cuda_cupy_version_x)
255+
except Exception:
256+
# if no compatible prebuilt wheel is found, try building CuPy ourselves
257+
session.warn(f"No prebuilt CuPy found for CUDA {cuda_version}, building from source...")
258+
session.install("cupy")
257259

258260
# for the last Python version session if all optional dependencies are enabled:
259261
if session.python == python_versions_to_test[-1] and full_install:

0 commit comments

Comments
 (0)