diff --git a/python/libucxx/libucxx/load.py b/python/libucxx/libucxx/load.py index 4d49d2cd..be514139 100644 --- a/python/libucxx/libucxx/load.py +++ b/python/libucxx/libucxx/load.py @@ -15,8 +15,35 @@ 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 +# libucxx 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 libucxx.so and its dependencies""" # If libucx was installed as a wheel, we must request it to load the library # symbols. Otherwise, we assume that the library was installed in a system path # that ld can find. @@ -28,25 +55,33 @@ def load_library(): libucx.load_library() del libucx - # Dynamically load libucxx.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_LIBUCXX_PREFER_SYSTEM_LIBRARY", "false").lower() != "false" + ) + + soname = "libucxx.so" libucxx_lib = None - try: - libucxx_lib = ctypes.CDLL("libucxx.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 libucxx Python package was a no-op. Note that this approach - # won't work for real editable installs of the libucxx package, but that's not a - # use case I think we need to support. 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, "libucxx.so") - ): - libucxx_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: + libucxx_lib = _load_system_installation(soname) + except OSError: + libucxx_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. + try: + libucxx_lib = _load_wheel_installation(soname) + if libucxx_lib is None: + libucxx_lib = _load_system_installation(soname) + except OSError: + # If none of the searches above succeed, just silently return None + # and rely on other mechanisms (like RPATHs on other DSOs) to + # help the loader find the library. + pass # 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