: oss: now we have rpath remove LD_LIBRARY_PATH

python/tests/test_rust_bindings_load.py

@@ -0,0 +1,87 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-ignore-all-errors
+
+import importlib
+import os
+import pathlib
+import re
+import shutil
+import subprocess
+import sys
+import sysconfig
+
+import pytest
+
+
+@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="Linux-only")
+def test_import_is_hermetic_without_loader_env():
+    """
+    Import the extension in a clean loader environment so RUNPATH must do the work.
+    This catches cases where the wheel works locally only because LD_LIBRARY_PATH is set.
+    """
+    # Keep PYTHONPATH (Buck's link-tree). Only scrub loader paths.
+    env = {
+        k: v
+        for k, v in os.environ.items()
+        if k not in ("LD_LIBRARY_PATH", "DYLD_LIBRARY_PATH")
+    }
+    env["PYTHONNOUSERSITE"] = "1"
+    # Ensure child's sys.path matches parent explicitly (works in Buck & pip)
+    env["PYTHONPATH"] = os.pathsep.join(sys.path)
+    subprocess.run(
+        [sys.executable, "-c", "import monarch._rust_bindings"], env=env, check=True
+    )
+
+
+@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="Linux-only")
+def test_runpath_and_needed_dependencies_resolve():
+    if not shutil.which("readelf") or not shutil.which("ldd"):
+        pytest.skip("readelf/ldd not available")
+
+    # 1) Import (normal env) and locate the .so
+    mod = importlib.import_module("monarch._rust_bindings")
+    so_file = getattr(mod, "__file__", None)
+    if not so_file or so_file == "static-extension" or not os.path.isabs(so_file):
+        pytest.skip("extension is statically linked here; no shared object to inspect")
+
+    so_path = pathlib.Path(so_file)
+
+    # 2) Inspect dynamic section
+    dyn = subprocess.run(
+        ["readelf", "-d", str(so_path)], capture_output=True, text=True, check=True
+    ).stdout
+
+    # RUNPATH sanity: we expect our relocation entries to be present
+    # ($ORIGIN for sidecars; parent dirs; and jump to the env's lib)
+    assert "$ORIGIN" in dyn, f"RUNPATH missing $ORIGIN:\n{dyn}"
+    assert "$ORIGIN/.." in dyn, f"RUNPATH missing $ORIGIN/..:\n{dyn}"
+    assert "$ORIGIN/../../.." in dyn, f"RUNPATH missing $ORIGIN/../../..:\n{dyn}"
+
+    # 3) Parse NEEDED
+    needed = re.findall(r"\(NEEDED\)\s+Shared library: \[(.+?)\]", dyn)
+
+    # 4) ldd resolution check (conditional)
+    ldd = subprocess.run(
+        ["ldd", str(so_path)], capture_output=True, text=True, check=True
+    ).stdout
+
+    # libpython may be absent if built as abi3; only assert when present
+    if any("libpython" in n for n in needed):
+        py_hits = [ln for ln in ldd.splitlines() if "libpython" in ln]
+        assert py_hits, f"ldd did not list libpython:\n{ldd}"
+        assert all(
+            "not found" not in ln for ln in py_hits
+        ), f"libpython unresolved:\n{ldd}"
+
+    # CUDA runtime: if you linked against cudart, ensure it resolves
+    if any("cudart" in n for n in needed):
+        cu_hits = [ln for ln in ldd.splitlines() if "cudart" in ln]
+        assert cu_hits, f"ldd did not list libcudart:\n{ldd}"
+        assert all(
+            "not found" not in ln for ln in cu_hits
+        ), f"libcudart unresolved:\n{ldd}"

scripts/common-setup.sh

@@ -15,7 +15,6 @@ setup_conda_environment() {
     echo "Setting up conda environment with Python ${python_version}..."
     conda create -n venv python="${python_version}" -y
     conda activate venv
-    export LD_LIBRARY_PATH="${CONDA_PREFIX}/lib:$LD_LIBRARY_PATH"
     export PATH=/opt/rh/devtoolset-10/root/usr/bin/:$PATH
     python -m pip install --upgrade pip
 }

setup.py

@@ -8,6 +8,7 @@
 
 import shutil
 import subprocess
+import sys
 import sysconfig
 
 import torch
@@ -119,6 +120,45 @@ def run(self):
 with open("README.md", encoding="utf8") as f:
     readme = f.read()
 
+if sys.platform.startswith("linux"):
+    # Always include the active env's lib (Conda-safe)
+    conda_lib = os.path.join(sys.prefix, "lib")
+
+    # Only use LIBDIR if it actually contains the current libpython
+    import os, sysconfig
+
+    ldlib = sysconfig.get_config_var("LDLIBRARY") or ""
+    libdir = sysconfig.get_config_var("LIBDIR") or ""
+    py_lib = ""
+    if libdir and ldlib:
+        cand = os.path.join(libdir, ldlib)
+        if os.path.exists(cand) and os.path.realpath(libdir) != os.path.realpath(
+            conda_lib
+        ):
+            py_lib = libdir
+
+    # Prefer sidecar .so next to the extension; then the conda env;
+    # then (optionally) py_lib
+    flags = [
+        "-C",
+        "link-arg=-Wl,--enable-new-dtags",
+        "-C",
+        "link-arg=-Wl,-z,origin",
+        "-C",
+        "link-arg=-Wl,-rpath,$ORIGIN",
+        "-C",
+        "link-arg=-Wl,-rpath,$ORIGIN/..",
+        "-C",
+        "link-arg=-Wl,-rpath,$ORIGIN/../../..",
+        "-C",
+        "link-arg=-Wl,-rpath," + conda_lib,
+    ]
+    if py_lib:
+        flags += ["-C", "link-arg=-Wl,-rpath," + py_lib]
+
+    cur = os.environ.get("RUSTFLAGS", "")
+    os.environ["RUSTFLAGS"] = (cur + " " + " ".join(flags)).strip()
+
 rust_extensions = [
     RustBin(
         target="process_allocator",