Skip to content

Commit

Permalink
fix docker build
Browse files Browse the repository at this point in the history
  • Loading branch information
casperdcl committed Aug 10, 2023
1 parent 372cae0 commit 3525666
Showing 1 changed file with 13 additions and 5 deletions.
18 changes: 13 additions & 5 deletions setup.py
100755 → 100644
Original file line number Diff line number Diff line change
@@ -1,17 +1,16 @@
#!/usr/bin/env python3
"""
Compile CUDA source code and setup Python 3 package 'nimpa'
for namespace 'niftypet'.
"""
import logging
import re
import sys
from pathlib import Path

from setuptools import find_packages, setup
from setuptools_scm import get_version

from niftypet.ninst import cudasetup as cs
from niftypet.ninst import dinf
from niftypet.ninst import install_tools as tls

__version__ = get_version(root=".", relative_to=__file__)
Expand All @@ -36,6 +35,8 @@
Cnt = resources.get_setup()

build_ver = ".".join(__version__.split('.')[:3]).split(".dev")[0]
# some setup kwargs cannot be in `pyproject.toml` since
# `install_requires` is dynamically set depending on CUDA GPU detection
setup_kwargs = {
"use_scm_version": True, "packages": find_packages(exclude=["tests"]), "package_data": {
"niftypet": [
Expand All @@ -51,11 +52,18 @@

try:
import cuvec as cu
from miutil import cuinfo
from skbuild import setup as sksetup
assert cu.include_path.is_dir()
nvcc_arches = {"{2:d}{3:d}".format(*i) for i in dinf.gpuinfo() if i[2:4] >= (3, 5)}
if nvcc_arches:
cmake_args.append("-DCMAKE_CUDA_ARCHITECTURES=" + ";".join(sorted(nvcc_arches)))
try:
nvcc_arch_raw = map(cuinfo.compute_capability, range(cuinfo.num_devices()))
nvcc_arches = {"%d%d" % i for i in nvcc_arch_raw if i >= (3, 5)}
if nvcc_arches:
cmake_args.append("-DCMAKE_CUDA_ARCHITECTURES=" + ";".join(sorted(nvcc_arches)))
except Exception as exc:
if "sdist" not in sys.argv or any(i in sys.argv for i in ["build", "bdist", "wheel"]):
log.warning("CUDA device detection error:\n%s", exc)
log.warning("Compiling for all architectures")
except Exception as exc:
log.warning("Import or CUDA device detection error:\n%s", exc)
setup(**setup_kwargs)
Expand Down

0 comments on commit 3525666

Please sign in to comment.