Improved setup by packaging the optix headers only at wheel creation to enable easier installation in virtual/conda environments

Author: Felix Igelbrink

.gitignore

@@ -10,3 +10,4 @@ imgui.ini
 *.cpp
 *.c
 .*
+include/

optix/_version.py

@@ -1 +1 @@
-__version__ = "0.1.2"
+__version__ = "0.1.3"

optix/denoiser.pyx

@@ -319,7 +319,7 @@ cdef class Denoiser(OptixContextObject):
                 denoise_alpha = DenoiserAlphaMode.COPY
 
             assert isinstance(denoise_alpha, DenoiserAlphaMode), "Optix >7.5 changed this from a boolean variable into an enum"
-            params.denoiseAlpha = denoise_alpha.value
+            params.denoiseAlpha = <OptixDenoiserAlphaMode>denoise_alpha.value
         ELSE:
             params.denoiseAlpha = 1 if denoise_alpha else 0
 

optix/module.pyx

@@ -2,7 +2,8 @@
 
 from enum import IntEnum, IntFlag
 import os
-from .path_utility import get_cuda_include_path, get_optix_include_path
+import warnings
+from .path_utility import get_cuda_include_path, get_optix_include_path, get_local_optix_include_path
 from .common cimport optix_check_return, optix_init
 from .context cimport DeviceContext
 from .pipeline cimport PipelineCompileOptions
@@ -408,9 +409,11 @@ cdef class Module(OptixContextObject):
         flags = list(compile_flags)
         # get cuda and optix_include_paths
         cuda_include_path = get_cuda_include_path()
-        print("cuda path", cuda_include_path)
-        optix_include_path = get_optix_include_path()
-
+        optix_include_path = get_local_optix_include_path()
+        if not os.path.exists(optix_include_path):
+            warnings.warn("Local optix not found. This usually indicates some installation issue. Attempting"
+                          " to load the global optix includes instead.", RuntimeWarning)
+            optix_include_path = get_optix_include_path()
         flags.extend([f'-I{cuda_include_path}', f'-I{optix_include_path}'])
         ptx, _ = prog.compile(flags)
         return ptx

optix/path_utility.py

@@ -25,6 +25,7 @@
 
 import os
 from itertools import chain
+import pathlib
 
 _cuda_path_cache = 'NOT_INITIALIZED'
 _optix_path_cache = 'NOT_INITIALIZED'
@@ -98,7 +99,7 @@ def get_optix_path(path_hint=None, environment_variable=None):
                                                                                                  None else None)
         if optix_header_path is None:
             # search on the default path
-            optix_header_path = search_on_path(('../optix/include/optix.h',), keys=('PATH',))
+            optix_header_path = search_on_path(('../optix/include/optix.h',), keys=('PATH', 'OPTIX_PATH'))
 
         if optix_header_path is not None:
             optix_header_path = os.path.normpath(os.path.join(os.path.dirname(optix_header_path), '..'))
@@ -115,6 +116,10 @@ def get_optix_path(path_hint=None, environment_variable=None):
     return _optix_path_cache
 
 
+def get_local_optix_include_path():
+    local_include_path = pathlib.Path(__file__).parent / "include"
+    return str(local_include_path) if local_include_path.exists() else None
+
 def get_optix_include_path(environment_variable=None):
     optix_path = get_optix_path(environment_variable=environment_variable)
     if optix_path is None:
@@ -124,3 +129,4 @@ def get_optix_include_path(environment_variable=None):
         return optix_include_path
     else:
         return None
+

setup.py

@@ -1,8 +1,10 @@
-from setuptools import setup, Extension, find_packages
+from struct import pack
+from setuptools import setup, Extension, find_packages, find_namespace_packages
 from Cython.Build import cythonize
 import re
 import os
 from pathlib import Path
+import shutil
 
 
 # standalone import of a module (https://stackoverflow.com/a/58423785)
@@ -29,6 +31,8 @@ def import_module_from_path(path):
 util = import_module_from_path('optix/path_utility.py')
 cuda_include_path = util.get_cuda_include_path()
 optix_include_path = util.get_optix_include_path()
+print("Found cuda includes at", cuda_include_path)
+print("Found optix includes at", optix_include_path)
 if cuda_include_path is None or optix_include_path is None:
     raise RuntimeError("Cuda or optix not found in the system")
 
@@ -65,6 +69,35 @@ def import_module_from_path(path):
 
 version = import_module_from_path('optix/_version.py').__version__
 
+package_data = {}
+
+try:
+    from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
+
+    def glob_fix(package_name, glob):
+        # this assumes setup.py lives in the folder that contains the package
+        package_path = Path(f'./{package_name}').resolve()
+        return [str(path.relative_to(package_path)) 
+                for path in package_path.glob(glob)]
+
+    class custom_bdist_wheel(_bdist_wheel):
+        def finalize_options(self):
+            _bdist_wheel.finalize_options(self)
+
+            # create the path for the internal headers
+            # due to optix license restrictions those headers 
+            # cannot be distributed on pypi directly so we will add this headers dynamically
+            # upon wheel construction to install them alongside the package
+
+            if not os.path.exists('optix/include/optix.h'):
+                shutil.copytree(optix_include_path, 'optix/include')    
+            self.distribution.package_data.update({
+                'optix': [*glob_fix('optix', 'include/**/*')]
+            })
+
+except ImportError:
+    custom_bdist_wheel = None
+
 setup(
     name="python-optix",
     version=version,
@@ -87,6 +120,7 @@ def import_module_from_path(path):
     classifiers=[
         "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
         "License :: OSI Approved :: MIT License",
         "Operating System :: POSIX :: Linux",
         "Operating System :: Microsoft :: Windows",
@@ -101,5 +135,7 @@ def import_module_from_path(path):
         'examples': ["pillow", "pyopengl", "pyglfw", "pyimgui"]
     },
     python_requires=">=3.8",
-    zip_safe=False
+    package_data=package_data,
+    zip_safe=False,
+    cmdclass={'bdist_wheel': custom_bdist_wheel}
 )