prefer wheel-provided libcuspatial.so in load_library(), use RTLD_LOCAL

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