From 20abaaa5f5efa9a476fc875b5de5e94d72bfe394 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Thu, 14 Nov 2024 16:02:17 -0600
Subject: [PATCH] prefer wheel-provided libcuspatial.so in load_library(), use
 RTLD_LOCAL (#1483)

Contributes to https://github.com/rapidsai/build-planning/issues/118

Modifies `libcuspatial.load_library()` in the following ways:

* prefer wheel-provided `libcuspatial.so` to system installation
* expose environment variable `RAPIDS_LIBCUSPATIAL_PREFER_SYSTEM_LIBRARY` for switching that preference
* load `libcuspatial.so` with `RTLD_LOCAL`, to prevent adding symbols to the global namespace ([dlopen docs](https://linux.die.net/man/3/dlopen))

## Notes for Reviewers

### How I tested this

See "How I tested this" in https://github.com/rapidsai/cudf/pull/17316

Also opened this PR originally pulling in built packages from https://github.com/rapidsai/cudf/pull/17316

#

Authors:
  - James Lamb (https://github.com/jameslamb)

Approvers:
  - Vyas Ramasubramani (https://github.com/vyasr)
  - Bradley Dice (https://github.com/bdice)

URL: https://github.com/rapidsai/cuspatial/pull/1483
---
 python/libcuspatial/libcuspatial/load.py | 75 ++++++++++++++++--------
 1 file changed, 51 insertions(+), 24 deletions(-)

diff --git a/python/libcuspatial/libcuspatial/load.py b/python/libcuspatial/libcuspatial/load.py
index 7842de9a0..78a8590ef 100644
--- a/python/libcuspatial/libcuspatial/load.py
+++ b/python/libcuspatial/libcuspatial/load.py
@@ -16,8 +16,37 @@
 import ctypes
 import os
 
+# Loading with RTLD_LOCAL adds the library itself to the loader's
+# loaded library cache without loading any symbols into the global
+# namespace. This allows libraries that express a dependency on
+# this library to be loaded later and successfully satisfy this dependency
+# without polluting the global symbol table with symbols from
+# libcuspatial that could conflict with symbols from other 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.
+    """
+    if os.path.isfile(
+        lib := os.path.join(os.path.dirname(__file__), "lib64", soname)
+    ):
+        return ctypes.CDLL(lib, PREFERRED_LOAD_FLAG)
+    return None
+
 
 def load_library():
+    """Dynamically load libcuspatial.so and its dependencies"""
     try:
         # libcudf must be loaded before libcuspatial because libcuspatial
         # references its symbols
@@ -33,32 +62,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