Improve setup to allow embedding the optix headers for bdist_wheel, install and develop commands

README.md

@@ -2,37 +2,58 @@
 
 Python wrapper for the OptiX 7 raytracing engine.
 
-Python-OptiX wraps the OptiX C++ API using Cython and provides a simplified 
-interface to the original C-like API using the 
-[CuPy](https://cupy.dev) package.
+Python-OptiX wraps the original OptiX C-like API using Cython while aiming to provide a more
+pythonic, object-oriented interface using the [CuPy](https://cupy.dev) package.
 
 ### Supported Platforms
 
 Only Linux is officially supported at the moment. Experimental windows support is available.
 
 ### OptiX Versions
 
-Python-OptiX currently supports the OptiX releases 7.3.0, 7.4.0 and 7.5.0
+Python-OptiX always supports the most recent version of the OptiX SDK. 
+The current version therefore supports OptiX 7.6.0
 
 ## Installation
 
 ### Dependencies
 
 Install a recent version of the [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads)
-and the [OptiX 7.5.0 SDK](https://developer.nvidia.com/optix/downloads/7.5.0/linux64-x86_64)
+and the [OptiX 7.6.0 SDK](https://developer.nvidia.com/optix/downloads/7.6.0/linux64-x86_64)
 
-Make sure the CUDA header files are installed as well.
+Make sure the CUDA header files are installed as well. 
 
-Add the locations of CUDA and OptiX to the system `PATH` variable if necessary.
+Note, that for some variants of the CUDA Toolkit, 
+like the one installed by the `conda` package manager, these are not installed by default. 
+`conda`-environments require the additional `cudatoolkit-dev` package.
+
+### Environment
+
+`python-optix` requires both the OptiX as well as the CUDA include path during setup as well as runtime 
+to compile the CUDA kernels. Therefore, it is necessary to either add both locations to the system `PATH`
+or set the `CUDA_PATH` and `OPTIX_PATH` variables to the respective locations.
+
+The setup additionally has the option to embed the OptiX header files into the `python-optix` installation. 
+If the variable `OPTIX_EMBED_HEADERS` is set to `1`, the setup will copy the headers from the 
+OptiX SDK directory into the generated wheel.
+
+If this option was chosen during setup, setting the `OPTIX_PATH` is no longer required as the 
+embedded headers will be utilized then.
 
 ### Using pip
 ```
-pip install python-optix
+export OPTIX_PATH=/path/to/optix
+export CUDA_PATH=/path/to/cuda_toolkit
+export OPTIX_EMBED_HEADERS=1 # embed the optix headers into the package
+python -m pip install python-optix
 ```
 
 ### From source
 ```
 git clone https://github.com/mortacious/python-optix.git
 cd python-optix
-python setup.py install
+export OPTIX_PATH=/path/to/optix
+export CUDA_PATH=/path/to/cuda_toolkit
+export OPTIX_EMBED_HEADERS=1 # embed the optix headers into the package
+python -m pip install [-e] .
 ```

examples/hello.py

@@ -10,6 +10,7 @@
 script_dir = os.path.dirname(__file__)
 cuda_src = os.path.join(script_dir, "cuda", "hello.cu")
 
+
 def create_module(ctx, pipeline_opts):
     compile_opts = ox.ModuleCompileOptions(debug_level=ox.CompileDebugLevel.FULL, opt_level=ox.CompileOptimizationLevel.LEVEL_0)
     module = ox.Module(ctx, cuda_src, compile_opts, pipeline_opts)

optix/context.pyx

@@ -144,13 +144,13 @@ cdef class DeviceContext(OptixObject):
         """
         The callback function for logging
         """
-        return self._log_callback_function
+        return self._log_callback
 
     @log_callback.setter
     def log_callback(self, object log_callback_function):
-        self._log_callback_function = log_callback_function
-        if self._log_callback_function is not None:
-            optix_check_return(optixDeviceContextSetLogCallback(self.c_context, context_log_cb, <void*>self._log_callback_function, self._log_callback_level))
+        self._log_callback = log_callback_function
+        if self._log_callback is not None:
+            optix_check_return(optixDeviceContextSetLogCallback(self.c_context, context_log_cb, <void*>self._log_callback, self._log_callback_level))
 
     @property
     def log_callback_level(self):

optix/module.pyx

@@ -326,7 +326,7 @@ cdef class Module(OptixContextObject):
         cdef unsigned int pipeline_payload_values, i
         #cls._check_payload_values(module_compile_options, pipeline_compile_options)
 
-        ptx = cls.compile_cuda_ptx(src, compile_flags, name=program_name)
+        ptx = module.compile_cuda_ptx(src, compile_flags, name=program_name)
         c_ptx = ptx
 
         cdef Task task = Task(module)
@@ -381,8 +381,8 @@ cdef class Module(OptixContextObject):
     def get_default_nvrtc_compile_flags(std=None, rdc=False):
         return get_default_nvrtc_compile_flags(std, rdc)
 
-    @staticmethod
-    def compile_cuda_ptx(src, compile_flags=_nvrtc_compile_flags_default, name=None, **kwargs):
+
+    def compile_cuda_ptx(self, src, compile_flags=_nvrtc_compile_flags_default, name=None, **kwargs):
         """
         Compiles a valid source module into the ptx format. Accepts files containing either source code, ptx, or
         optix-ir code, compiles the source code if necessary and returns valid ptx or optix-ir modules.
@@ -424,9 +424,14 @@ cdef class Module(OptixContextObject):
             cuda_include_path = get_cuda_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)
+                # attempt to load the global path if the local path is not available
                 optix_include_path = get_optix_include_path()
+            if optix_include_path is None:
+                raise ValueError("Unable to locate the optix headers. Make sure that either the OPTIX_PATH environement variable is set"
+                                 "correctly or the optix headers are embedded into this package.")
+            if <Module>self.context.log_callback is not None:
+                # hook into the logging system for this output
+                <Module>self.context.log_callback(4, "build", f"Using optix include path: {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

@@ -120,6 +120,7 @@ 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:

pyproject.toml

@@ -1,3 +1,3 @@
 [build-system]
-requires = ["setuptools", "wheel", "Cython>=0.29.22,<3"]
+requires = ["setuptools", "wheel", "Cython>=0.29.22,<3", "numpy"]
 build-backend = "setuptools.build_meta"

setup.py

@@ -30,14 +30,16 @@ 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()
+cuda_include_path = util.get_cuda_include_path(environment_variable='CUDA_PATH')
+optix_include_path = util.get_optix_include_path(environement_variable='OPTIX_PATH')
 print("Found cuda includes at", cuda_include_path)
 print("Found optix includes at", optix_include_path)
 if cuda_include_path is None:
-    raise RuntimeError("CUDA not found in the system, but is required to build this package.")
+    raise RuntimeError("CUDA not found in the system, but is required to build this package. Consider setting"
+                       "CUDA_PATH to the location of the local cuda toolkit installation.")
 if optix_include_path is None:
-    raise RuntimeError("OptiX not found in the system, but is required to build this package.")
+    raise RuntimeError("OptiX not found in the system, but is required to build this package. Consider setting "
+                       "OPTIX_PATH to the location of the optix SDK.")
 
 optix_version_re = re.compile(r'.*OPTIX_VERSION +(\d{5})')  # get the optix version from the header
 with open(Path(optix_include_path) / "optix.h", 'r') as f:
@@ -75,33 +77,59 @@ def import_module_from_path(path):
 
 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)]
+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)]
 
+from setuptools.command.install import install as _install
+from setuptools.command.develop import develop as _develop
 
-    class custom_bdist_wheel(_bdist_wheel):
-        def finalize_options(self):
-            _bdist_wheel.finalize_options(self)
+class EmbeddHeadersCommandMixin:
+    def update_package_data(self):
+        self.distribution.package_data.update({
+            'optix': [*glob_fix('optix', 'include/**/*')]
+        })
+        print("embedding optix headers into package data",
+              self.distribution.package_data)
 
+    def run(self):
+        embedd = os.getenv("OPTIX_EMBED_HEADERS")
+        if embedd:
             # create the path for the internal headers
-            # due to optix license restrictions those 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/**/*')]
-            })
+                shutil.copytree(optix_include_path, 'optix/include')
+
+            self.update_package_data()
+
+        super().run()
+
+
+class CustomInstallCommand(EmbeddHeadersCommandMixin, _install):
+    pass
+
+
+class CustomDevelopCommand(EmbeddHeadersCommandMixin, _develop):
+    pass
+
+
+cmd_classes = {'install': CustomInstallCommand,
+               'develop': CustomDevelopCommand}
+
+try:
+    from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
 
+    class CustomBdistWheelCommand(EmbeddHeadersCommandMixin, _bdist_wheel):
+        pass
+    cmd_classes['bdist_wheel'] = CustomBdistWheelCommand
 except ImportError:
-    custom_bdist_wheel = None
+    CustomBdistWheel = None
 
 setup(
     name="python-optix",
@@ -142,5 +170,5 @@ def finalize_options(self):
     python_requires=">=3.8",
     package_data=package_data,
     zip_safe=False,
-    cmdclass={'bdist_wheel': custom_bdist_wheel}
+    cmdclass=cmd_classes
 )