prefer wheel-provided libraries, use RTLD_LOCAL

.pre-commit-config.yaml

@@ -2,7 +2,7 @@
 
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v5.0.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
@@ -11,12 +11,12 @@ repos:
     hooks:
       - id: isort
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.4.4
+    rev: v0.7.3
     hooks:
       - id: ruff
       - id: ruff-format
   - repo: https://github.com/rapidsai/pre-commit-hooks
-    rev: v0.0.3
+    rev: v0.4.0
     hooks:
       - id: verify-copyright
 

VERSION

@@ -1 +1 @@
-1.17.0
+1.15.0.post2

python/libucx/libucx/load.py

@@ -16,36 +16,61 @@
 import ctypes
 import os
 
-# IMPORTANT: The load order here matters! libucm.so depends on symbols in libucs.so, but
-# it does not express this via a DT_NEEDED entry, presumably because libucs.so also has
-# a dependency on libucm.so and the libraries are attempting to avoid a circular
-# dependency. Moreover, it seems like if libucs.so is not loaded before libuct.so and
-# libucp.so something is set up incorrectly, perhaps with the atexit handlers, because
-# on library close there is a double free issue. Therefore, libucs.so must be loaded
-# first. The other libraries may then be loaded in any order. The libraries themselves
-# all have $ORIGIN RPATHs to find each other.
 UCX_LIBRARIES = [
     "libucs.so",
     "libucs_signal.so",
-    "libucm.so",
     "libuct.so",
     "libucp.so",
 ]
 
 
+# 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
+# a library to be loaded later and successfully satisfy that dependency
+# without polluting the global symbol table with symbols from
+# that library 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__), "lib", soname)):
+        return ctypes.CDLL(lib, PREFERRED_LOAD_FLAG)
+    return None
+
+
 def load_library():
-    # Dynamically load libucx.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.
+    """Dynamically load UCX libraries"""
+    prefer_system_installation = (
+        os.getenv("RAPIDS_LIBUCX_PREFER_SYSTEM_LIBRARY", "false").lower() != "false"
+    )
+
     libraries = []
     for lib in UCX_LIBRARIES:
-        try:
-            libucx_lib = ctypes.CDLL(lib, ctypes.RTLD_GLOBAL)
-        except OSError:
-            libucx_lib = ctypes.CDLL(
-                os.path.join(os.path.dirname(__file__), "lib", lib),
-                ctypes.RTLD_GLOBAL,
-            )
+        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:
+                libucx_lib = _load_system_installation(lib)
+            except OSError:
+                libucx_lib = _load_wheel_installation(lib)
+        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.
+            libucx_lib = _load_wheel_installation(lib)
+
         libraries.append(libucx_lib)
 
     return libraries

python/libucx/setup.py

@@ -67,6 +67,10 @@ def run(self):
                     subprocess.run(["make", "install"])
                     # The config file built into UCX is not relocatable. We need to fix
                     # that so that we can package up UCX and distribute it in a wheel.
+                    # This was fixed in
+                    # https://github.com/openucx/ucx/commit/c3437851f4abb5b7da966109d3762966ac30c476
+                    # and merged into version 1.17, so we can remove this patch once we
+                    # no longer need to support older versions.
                     subprocess.run(
                         [
                             "sed",