diff --git a/python/libcuspatial/libcuspatial/load.py b/python/libcuspatial/libcuspatial/load.py index 7842de9a0..0a6f3966c 100644 --- a/python/libcuspatial/libcuspatial/load.py +++ b/python/libcuspatial/libcuspatial/load.py @@ -16,8 +16,38 @@ import ctypes import os +# RTLD_LOCAL is here for safety... using it loads symbols into the +# library-specific table maintained by the loader, but not into the +# global namespace where they may conflict with symbols from other +# loaded DSOs. +PREFERRED_LOAD_FLAG = ctypes.RTLD_LOCAL + + +def _load_system_installation(soname: str): + """Try to dlopen() the library indicated by ``soname`` + + Raises ``OSError`` if library cannot be loaded. + """ + return ctypes.CDLL(soname, PREFERRED_LOAD_FLAG) + + +def _load_wheel_installation(soname: str): + """Try to dlopen() the library indicated by ``soname`` + + Returns ``None`` if the library cannot be loaded. + """ + out = None + for lib_dir in ("lib", "lib64"): + if os.path.isfile( + lib := os.path.join(os.path.dirname(__file__), lib_dir, soname) + ): + out = ctypes.CDLL(lib, PREFERRED_LOAD_FLAG) + break + return out + def load_library(): + """Dynamically load libcuspatial.so and its dependencies""" try: # libcudf must be loaded before libcuspatial because libcuspatial # references its symbols @@ -33,32 +63,30 @@ def load_library(): # the loader can find it. pass - # Dynamically load libcuspatial.so. Prefer a system library if one is - # present to avoid clobbering symbols that other packages might expect, - # but if no other library is present use the one in the wheel. + prefer_system_installation = ( + os.getenv("RAPIDS_LIBCUSPATIAL_PREFER_SYSTEM_LIBRARY", "false").lower() + != "false" + ) + + soname = "libcuspatial.so" libcuspatial_lib = None - try: - libcuspatial_lib = ctypes.CDLL("libcuspatial.so", ctypes.RTLD_GLOBAL) - except OSError: - # If neither of these directories contain the library, we assume we are - # in an environment where the C++ library is already installed - # somewhere else and the CMake build of the libcuspatial Python package - # was a no-op. - # - # Note that this approach won't work for real editable installs of the - # libcuspatial package. scikit-build-core has limited support for - # importlib.resources so there isn't a clean way to support that case - # yet. - for lib_dir in ("lib", "lib64"): - if os.path.isfile( - lib := os.path.join( - os.path.dirname(__file__), lib_dir, "libcuspatial.so" - ) - ): - libcuspatial_lib = ctypes.CDLL(lib, ctypes.RTLD_GLOBAL) - break + if prefer_system_installation: + # Prefer a system library if one is present to + # avoid clobbering symbols that other packages might expect, but if no + # other library is present use the one in the wheel. + try: + libcuspatial_lib = _load_system_installation(soname) + except OSError: + libcuspatial_lib = _load_wheel_installation(soname) + else: + # Prefer the libraries bundled in this package. If they aren't found + # (which might be the case in builds where the library was prebuilt + # before packaging the wheel), look for a system installation. + libcuspatial_lib = _load_wheel_installation(soname) + if libcuspatial_lib is None: + libcuspatial_lib = _load_system_installation(soname) # The caller almost never needs to do anything with this library, but no # harm in offering the option since this object at least provides a handle - # to inspect where the library was loaded from. + # to inspect where the libcuspatial was loaded from. return libcuspatial_lib