Support adding patch_cmds in pip.override

[Ryang20718]

🚀 feature request

Description

Currently, the following exists for patching targets. However, for certain targets like pytorch, we have custom commands that patch the rpath of shared object libraries in our current workspace. However, this isn't doable with pip.override as we can't quite generate a patch for rpath differences. We can't patch the wheel beforehand because we're using pip torch relies on an rpath that is variable to change depending on the sha256sum of the wheel so we have to patch after the wheel is built as we can't know the sha256sum of the wheel before it's built

pip.override(
    file = "<your python wheel>.whl",
    patch_strip = 1,
    patches = [
        "<patch target>",
    ],
)

Describe the solution you'd like

ideally if pip.override has a patch_cmds optional argument that would be great

pip.override(
    file = "<your python wheel>.whl",
    patch_strip = 1,
    patches = [
        "<patch target>",
    ],
    patch_cmds = [
     ]
)

Describe alternatives you've considered

We've tried patching before the wheel is built, but like the description above says, it's currently a chicken and egg problem

[lalten]

We had the exact same issue. For a while we just manually rpath-patched the .so files, repacked the wheel, and mirrored it on a fileserver that requirements.in points at.
Recently I switched this to patching the preloading logic to just try and find the files in the runfiles:

Diff for Pytorch 2.6.0:

--- a/torch/__init__.py
+++ b/torch/__init__.py
@@ -17,6 +17,7 @@ import importlib
 import inspect
 import math
 import os
+import pathlib
 import platform
 import sys
 import textwrap
@@ -280,27 +281,8 @@ def _preload_cuda_deps(lib_folder: str, lib_name: str) -> None:
     # Should only be called on Linux if default path resolution have failed
     assert platform.system() == "Linux", "Should only be called on Linux"
 
-    lib_path = None
-    for path in sys.path:
-        nvidia_path = os.path.join(path, "nvidia")
-        if not os.path.exists(nvidia_path):
-            continue
-        candidate_lib_paths = glob.glob(
-            os.path.join(nvidia_path, lib_folder, "lib", lib_name)
-        )
-        # if path/nvidia/lib_folder/ is not found look in path/lib_folder/
-        if not candidate_lib_paths:
-            candidate_lib_paths = glob.glob(
-                os.path.join(path, lib_folder, "lib", lib_name)
-            )
-
-        if candidate_lib_paths and not lib_path:
-            lib_path = candidate_lib_paths[0]
-        if lib_path:
-            break
-    if not lib_path:
-        raise ValueError(f"{lib_name} not found in the system path {sys.path}")
-    ctypes.CDLL(lib_path)
+    for lib_path in pathlib.Path.cwd().parent.rglob(lib_name):
+        ctypes.CDLL(lib_path.as_posix())
 
 
 # See Note [Global dependencies]
@@ -314,28 +296,7 @@ def _load_global_deps() -> None:
     here = os.path.abspath(__file__)
     global_deps_lib_path = os.path.join(os.path.dirname(here), "lib", lib_name)
 
-    try:
-        ctypes.CDLL(global_deps_lib_path, mode=ctypes.RTLD_GLOBAL)
-        # Workaround slim-wheel CUDA dependency bugs in cusparse and cudnn by preloading nvjitlink
-        # and nvrtc. In CUDA-12.4+ cusparse depends on nvjitlink, but does not have rpath when
-        # shipped as wheel, which results in OS picking wrong/older version of nvjitlink library
-        # if `LD_LIBRARY_PATH` is defined, see https://github.com/pytorch/pytorch/issues/138460
-        # Similar issue exist in cudnn that dynamically loads nvrtc, unaware of its relative path.
-        # See https://github.com/pytorch/pytorch/issues/145580
-        try:
-            with open("/proc/self/maps") as f:
-                _maps = f.read()
-            # libtorch_global_deps.so always depends in cudart, check if its installed via wheel
-            if "nvidia/cuda_runtime/lib/libcudart.so" not in _maps:
-                return
-            # If all above-mentioned conditions are met, preload nvrtc and nvjitlink
-            # Please note that order are important for CUDA-11.8 , as nvjitlink does not exist there
-            _preload_cuda_deps("cuda_nvrtc", "libnvrtc.so.*[0-9]")
-            _preload_cuda_deps("nvjitlink", "libnvJitLink.so.*[0-9]")
-        except Exception:
-            pass
-
-    except OSError as err:
+    if True:
         # Can only happen for wheel with cuda libs as PYPI deps
         # As PyTorch is not purelib, but nvidia-*-cu12 is
         cuda_libs: _Dict[str, str] = {
@@ -353,11 +314,6 @@ def _load_global_deps() -> None:
             "nccl": "libnccl.so.*[0-9]",
             "nvtx": "libnvToolsExt.so.*[0-9]",
         }
-        is_cuda_lib_err = [
-            lib for lib in cuda_libs.values() if lib.split(".")[0] in err.args[0]
-        ]
-        if not is_cuda_lib_err:
-            raise err
         for lib_folder, lib_name in cuda_libs.items():
             _preload_cuda_deps(lib_folder, lib_name)
         ctypes.CDLL(global_deps_lib_path, mode=ctypes.RTLD_GLOBAL)

However because of torch's RECORD structure you get 20k+ lines of output during the build 🙈, see #3515

Hope that helps :)

[rickeylev]

is this rpath patching still necessary if venv_site_packages is enabled? That should create a normal venv layout, with the files in the right relative locations