Merge pull request #63 from CCInc/master

Update build flags to support other enviornments

Author: nicolas-chaulet

.github/workflows/deploy.yaml

@@ -17,7 +17,7 @@ jobs:
             - name: Install dependencies
               run: |
                   python -m pip install --upgrade pip
-                  pip install torch numpy scikit-learn flake8 setuptools wheel twine numba
+                  pip install torch "numpy<1.20" scikit-learn flake8 setuptools wheel twine numba
             - name: Build package
               run: |
                   python setup.py build_ext --inplace

.github/workflows/tests.yaml

@@ -22,7 +22,7 @@ jobs:
         - name: Install dependencies
           run: |
               python -m pip install --upgrade pip
-              pip install numpy scikit-learn flake8 setuptools numba==0.49.1
+              pip install "numpy<1.20" scikit-learn flake8 setuptools numba==0.49.1
 
         - name: Install torch windows + linux
           if: ${{matrix.os != 'macos-latest'}}

setup.py

@@ -1,4 +1,6 @@
 from setuptools import setup, find_packages
+import os
+import glob
 
 try:
     import torch
@@ -11,14 +13,16 @@
 except:
     raise ModuleNotFoundError("Please install pytorch >= 1.1 before proceeding.")
 
-import glob
-
-from os import path
-
-this_directory = path.abspath(path.dirname(__file__))
-with open(path.join(this_directory, "README.md"), encoding="utf-8") as f:
-    long_description = f.read()
-
+WITH_CUDA = torch.cuda.is_available() and CUDA_HOME is not None
+WITH_CPU = True
+if os.getenv('FORCE_CUDA', '0') == '1':
+    WITH_CUDA = True
+if os.getenv('FORCE_ONLY_CUDA', '0') == '1':
+    WITH_CUDA = True
+    WITH_CPU = False
+if os.getenv('FORCE_ONLY_CPU', '0') == '1':
+    WITH_CUDA = False
+    WITH_CPU = True
 
 def get_ext_modules():
     TORCH_MAJOR = int(torch.__version__.split(".")[0])
@@ -31,7 +35,7 @@ def get_ext_modules():
     ext_sources = glob.glob("{}/src/*.cpp".format(ext_src_root)) + glob.glob("{}/src/*.cu".format(ext_src_root))
 
     ext_modules = []
-    if CUDA_HOME:
+    if WITH_CUDA:
         ext_modules.append(
             CUDAExtension(
                 name="torch_points_kernels.points_cuda",
@@ -47,16 +51,17 @@ def get_ext_modules():
     cpu_ext_src_root = "cpu"
     cpu_ext_sources = glob.glob("{}/src/*.cpp".format(cpu_ext_src_root))
 
-    ext_modules.append(
-        CppExtension(
-            name="torch_points_kernels.points_cpu",
-            sources=cpu_ext_sources,
-            include_dirs=["{}/include".format(cpu_ext_src_root)],
-            extra_compile_args={
-                "cxx": extra_compile_args,
-            },
+    if WITH_CPU:
+        ext_modules.append(
+            CppExtension(
+                name="torch_points_kernels.points_cpu",
+                sources=cpu_ext_sources,
+                include_dirs=["{}/include".format(cpu_ext_src_root)],
+                extra_compile_args={
+                    "cxx": extra_compile_args,
+                },
+            )
         )
-    )
     return ext_modules
 
 
@@ -68,8 +73,11 @@ def __init__(self, *args, **kwargs):
 def get_cmdclass():
     return {"build_ext": CustomBuildExtension}
 
+this_directory = os.path.abspath(os.path.dirname(__file__))
+with open(os.path.join(this_directory, "README.md"), encoding="utf-8") as f:
+    long_description = f.read()
 
-requirements = ["torch>=1.1.0", "numba", "scikit-learn"]
+requirements = ["torch>=1.1.0", "numba", "numpy<1.20", "scikit-learn"]
 
 url = "https://github.com/nicolas-chaulet/torch-points-kernels"
 __version__ = "0.6.10"