rocm: don't bundle librocjpeg; RPATH libtorchcodec_image.so instead

Bundling librocjpeg moves it out of _rocm_sdk_core/lib, breaking AMD's
own relative RPATH inside it that points to librocm_sysdeps_* transitive
deps. We then have to re-patch those paths ourselves, which is fragile
and was causing the missing librocm_sysdeps_va.so.2 / libamdhip64.so.7
failures at runtime.

The cleaner approach: exclude librocjpeg from auditwheel bundling and
instead append _rocm_sdk_core/lib (ROCm >= 7.14) and /opt/rocm/lib
(ROCm <= 7.2) to the RPATH of libtorchcodec_image.so itself. At runtime
the dynamic linker finds librocjpeg in its original location; AMD's own
RPATH on that library then correctly resolves all transitive deps without
any intervention from us.

_find_rocjpeg_lib() is kept: auditwheel still needs librocjpeg in
LD_LIBRARY_PATH to resolve the DT_NEEDED during its analysis pass, even
though we exclude it from bundling.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: akashveramd

.github/workflows/linux_rocm.yaml

@@ -139,10 +139,10 @@ jobs:
         bash packaging/remove_src.sh
 
         echo '::group::Install rocJPEG runtime'
-        # librocjpeg is bundled into the wheel by repair_wheel.py (ROCm 7.14
-        # pip-wheel build) or must be present on the system (ROCm <= 7.2 system
-        # install).  install_rocjpeg.sh ensures the runtime side-deps (libva)
-        # are present regardless of how librocjpeg itself was obtained.
+        # librocjpeg is NOT bundled into the wheel.  Instead, libtorchcodec_image.so
+        # has an RPATH entry pointing to _rocm_sdk_core/lib (ROCm 7.14 pip-wheel) and
+        # /opt/rocm/lib (ROCm <= 7.2 system install).  install_rocjpeg.sh ensures the
+        # runtime side-deps (libva) are present regardless of layout.
         bash packaging/install_rocjpeg.sh
         echo '::endgroup::'
 
@@ -162,36 +162,24 @@ jobs:
         echo '::endgroup::'
 
         echo '::group::Run FFmpeg-free image decoder tests (incl. rocJPEG GPU)'
-        # ROCm 7.14 pip-wheel layout: the ROCm runtime and its vendored system
-        # deps (librocm_sysdeps_va.so.2 etc.) live inside _rocm_sdk_core/lib/
-        # and _rocm_sdk_core/lib/rocm_sysdeps/lib/.  The RPATH on the bundled
-        # librocjpeg-HASH.so points there via $ORIGIN, but the dynamic linker
-        # may not expand $ORIGIN correctly without execution permission.  Set
-        # LD_LIBRARY_PATH explicitly as belt-and-suspenders.
-        rocm_core_lib=$(python -c "
-import importlib.util, pathlib, os
-spec = importlib.util.find_spec('_rocm_sdk_core')
-if spec:
-    p = pathlib.Path(spec.submodule_search_locations[0])
-    dirs = [str(p/'lib'), str(p/'lib'/'rocm_sysdeps'/'lib')]
-    print(os.pathsep.join(d for d in dirs if pathlib.Path(d).is_dir()))
-" 2>/dev/null || true)
+        # ROCm 7.14 pip-wheel layout: librocjpeg lives in _rocm_sdk_core/lib and
+        # its transitive deps (librocm_sysdeps_va.so.2 etc.) are in
+        # _rocm_sdk_core/lib/rocm_sysdeps/lib. The RPATH on libtorchcodec_image.so
+        # points to _rocm_sdk_core/lib, and AMD's own RPATH on librocjpeg handles
+        # the sysdeps. Set LD_LIBRARY_PATH as belt-and-suspenders for the HIP
+        # runtime (libamdhip64) that librocjpeg needs.
+        rocm_core_lib=$(python -c "import importlib.util, pathlib, os; spec = importlib.util.find_spec('_rocm_sdk_core'); p = pathlib.Path(spec.submodule_search_locations[0]) if spec else None; dirs = [str(p/'lib'), str(p/'lib'/'rocm_sysdeps'/'lib')] if p else []; print(os.pathsep.join(d for d in dirs if pathlib.Path(d).is_dir()))" 2>/dev/null || true)
         if [ -n "${rocm_core_lib}" ]; then
           export LD_LIBRARY_PATH="${rocm_core_lib}:${LD_LIBRARY_PATH:-}"
           echo "LD_LIBRARY_PATH (rocm_sdk_core): ${rocm_core_lib}"
         fi
-        # Diagnostics: dump ldd and RPATH so any remaining missing dep is visible.
+        # Diagnostics: dump ldd on the image .so so any missing dep is visible.
         image_so=$(python -c "import importlib.util, pathlib; spec = importlib.util.find_spec('torchcodec'); print(pathlib.Path(spec.origin).parent / 'libtorchcodec_image.so')" 2>/dev/null || true)
         if [ -n "${image_so}" ] && [ -f "${image_so}" ]; then
+          echo "RPATH of ${image_so}:"
+          objdump -p "${image_so}" 2>/dev/null | grep -E "RPATH|RUNPATH" || true
           echo "ldd ${image_so}:"
           ldd "${image_so}" || true
-          rocjpeg_so=$(find "$(dirname "${image_so}")/../torchcodec.libs" -name 'librocjpeg*.so*' 2>/dev/null | head -1 || true)
-          if [ -n "${rocjpeg_so}" ]; then
-            echo "ldd ${rocjpeg_so}:"
-            ldd "${rocjpeg_so}" || true
-            echo "RPATH of ${rocjpeg_so}:"
-            objdump -p "${rocjpeg_so}" 2>/dev/null | grep -E "RPATH|RUNPATH" || true
-          fi
         fi
         # torch.cuda.is_available() is True on ROCm (HIP masquerades as CUDA), so
         # the needs_cuda image tests run here and exercise the GPU JPEG decoder,

packaging/repair_wheel.py

@@ -163,7 +163,12 @@ def _find_rocjpeg_license():
 
 
 def _find_rocjpeg_lib():
-    """Find librocjpeg.so at wheel repair time so auditwheel can bundle it.
+    """Find librocjpeg.so at wheel repair time so auditwheel can resolve it.
+
+    auditwheel needs librocjpeg to be resolvable in LD_LIBRARY_PATH during
+    `auditwheel repair` even though we exclude it from bundling (so it stays
+    in the user's ROCm install). This function returns the directory to add to
+    LD_LIBRARY_PATH before calling auditwheel.
 
     Searches ROCM_HOME / ROCM_PATH env vars, torch's ROCM_HOME, the standard
     /opt/rocm fallback, and (for ROCm >= 7.14) the _rocm_sdk_* pip-wheel
@@ -230,16 +235,21 @@ def _find_rocjpeg_lib():
     return None
 
 
-def _patch_rocjpeg_rpath_in_wheel(wheel_path: Path) -> None:
-    """Set RPATH on the bundled librocjpeg-*.so.* inside a repaired wheel so
-    it can find its ROCm runtime deps at the user's install time without
-    LD_LIBRARY_PATH.
+def _patch_image_so_rpath_in_wheel(wheel_path: Path) -> None:
+    """Append ROCm library search paths to libtorchcodec_image.so's RPATH.
+
+    librocjpeg is NOT bundled in the wheel (it is excluded from auditwheel so
+    it stays in the user's ROCm install). At runtime the dynamic linker must
+    find librocjpeg via RPATH on libtorchcodec_image.so itself.
 
     Two layouts are covered:
       - ROCm >= 7.14 (TheRock / rocm-sdk-* Python wheels):
-          libamdhip64 etc. live in <site-packages>/_rocm_sdk_core/lib/.
+          librocjpeg lives in <site-packages>/_rocm_sdk_core/lib/.
           $ORIGIN/../../_rocm_sdk_core/lib reaches that dir from
-          <site-packages>/torchcodec.libs/librocjpeg-HASH.so.1.
+          <site-packages>/torchcodec/libtorchcodec_image.so.
+          AMD already set the correct RPATH inside _rocm_sdk_core's librocjpeg
+          to find librocm_sysdeps_* transitive deps, so we do not need to touch
+          those at all.
       - ROCm <= 7.2 (system install):
           /opt/rocm/lib is the standard path; the AMD installer always
           creates the /opt/rocm symlink even for versioned installs.
@@ -251,39 +261,44 @@ def _patch_rocjpeg_rpath_in_wheel(wheel_path: Path) -> None:
             "repairing ROCm wheels."
         )
 
-    rpath = ":".join([
-        "$ORIGIN",
-        # ROCm >= 7.14 layout A: torch bundles ROCm runtime inside torch/lib/
-        # (same layout PyTorch uses for CUDA runtime on CUDA builds).
-        "$ORIGIN/../../torch/lib",
-        # ROCm >= 7.14 layout B: rocm-sdk-core ships as a separate pip wheel
-        # and ROCm runtime lives in _rocm_sdk_core/lib.
-        "$ORIGIN/../../_rocm_sdk_core/lib",
-        "$ORIGIN/../../_rocm_sdk_core/lib/rocm_sysdeps/lib",
-        # ROCm <= 7.2: standard system install path (AMD installer always
-        # creates the /opt/rocm symlink even for versioned installs like 7.2.0).
-        "/opt/rocm/lib",
-        "/opt/rocm/lib/rocm_sysdeps/lib",
-    ])
-
     import hashlib, base64, tempfile
     with tempfile.TemporaryDirectory() as tmp:
         tmp_path = Path(tmp)
         with zipfile.ZipFile(wheel_path, "r") as zf:
             zf.extractall(tmp_path)
 
-        rocjpeg_libs = list(tmp_path.rglob("librocjpeg*.so*"))
-        if not rocjpeg_libs:
-            print(f"No librocjpeg found in {wheel_path.name}; skipping RPATH patch.", flush=True)
+        image_libs = list(tmp_path.rglob("libtorchcodec_image*.so*"))
+        if not image_libs:
+            print(
+                f"No libtorchcodec_image found in {wheel_path.name}; "
+                "skipping ROCm RPATH patch.",
+                flush=True,
+            )
             return
-        for lib in rocjpeg_libs:
-            print(f"Patching RPATH on {lib.name}: {rpath}", flush=True)
-            subprocess.run([patchelf, "--set-rpath", rpath, str(lib)], check=True)
+
+        for lib in image_libs:
+            # Read the RPATH auditwheel already set (e.g. $ORIGIN/../torchcodec.libs)
+            # and append the ROCm search dirs without clobbering them.
+            result = subprocess.run(
+                [patchelf, "--print-rpath", str(lib)],
+                capture_output=True, text=True, check=True,
+            )
+            existing = result.stdout.strip()
+            extra = ":".join([
+                # ROCm >= 7.14: librocjpeg lives in _rocm_sdk_core/lib alongside
+                # the other ROCm libraries shipped as a pip wheel.
+                "$ORIGIN/../../_rocm_sdk_core/lib",
+                # ROCm <= 7.2: standard system install (always symlinked to /opt/rocm).
+                "/opt/rocm/lib",
+            ])
+            new_rpath = f"{existing}:{extra}" if existing else extra
+            print(f"Setting RPATH on {lib.name}: {new_rpath}", flush=True)
+            subprocess.run([patchelf, "--set-rpath", new_rpath, str(lib)], check=True)
 
         # Update the RECORD file so pip's integrity check passes.
         # RECORD format: path,sha256=<base64url>,size  (or ",," for RECORD itself)
         record_files = list(tmp_path.rglob("RECORD"))
-        patched_rel_names = {lib.relative_to(tmp_path).as_posix() for lib in rocjpeg_libs}
+        patched_rel_names = {lib.relative_to(tmp_path).as_posix() for lib in image_libs}
         for record_file in record_files:
             lines = record_file.read_text(encoding="utf-8").splitlines()
             new_lines = []
@@ -333,7 +348,10 @@ def repair_linux(wheels):
             print(f"Found librocjpeg in {rocjpeg_lib_dir}", flush=True)
         else:
             print(
-                "WARNING: librocjpeg not found; rocJPEG will not be bundled. "
+                "WARNING: librocjpeg not found; auditwheel cannot resolve the "
+                "DT_NEEDED entry for librocjpeg. The wheel will still be built "
+                "but ROCm JPEG decoding will fail at runtime unless librocjpeg "
+                "is reachable via /opt/rocm/lib or _rocm_sdk_core/lib. "
                 "Set ROCM_HOME or ROCM_PATH if ROCm is in a non-standard location.",
                 flush=True,
             )
@@ -366,11 +384,18 @@ def repair_linux(wheels):
         "libnvshmem*",
         "libnvfatbin*",
         "libnvcuvid*",
+        # librocjpeg is NOT bundled. Instead, libtorchcodec_image.so gets an RPATH
+        # entry pointing to _rocm_sdk_core/lib (ROCm >= 7.14) and /opt/rocm/lib
+        # (ROCm <= 7.2), so the linker finds librocjpeg in its original location.
+        # This is intentional: AMD already set correct RPATHs inside their
+        # librocjpeg to find librocm_sysdeps_* and other transitive deps relative
+        # to _rocm_sdk_core/lib. Moving it (bundling) breaks those relative paths
+        # and requires us to re-patch them, which is fragile.
+        "librocjpeg*",
         # ROCm/HIP runtime and its system deps: provided by the torch-ROCm wheel
         # or the system ROCm install at runtime. Never bundle them — they would
         # duplicate torch's copies and bloat the wheel significantly (libLLVM
         # alone is ~200 MB).
-        # librocjpeg itself IS bundled (not listed here); only its deps are excluded.
         "libamdhip64*",
         "libamd_comgr*",
         "libhsa-runtime64*",
@@ -412,12 +437,14 @@ def repair_linux(wheels):
             env=env,
         )
 
-    # After auditwheel bundles librocjpeg-HASH.so.*, patch its RPATH so the
-    # bundled copy can find its ROCm runtime deps (libamdhip64 etc.) at the
-    # user's install time without requiring LD_LIBRARY_PATH.
+    # After auditwheel repair, patch libtorchcodec_image.so's RPATH to include
+    # _rocm_sdk_core/lib (ROCm >= 7.14) and /opt/rocm/lib (ROCm <= 7.2) so the
+    # dynamic linker can find librocjpeg at runtime without LD_LIBRARY_PATH.
+    # librocjpeg itself is NOT bundled; it stays in the ROCm install so AMD's
+    # own RPATH inside it correctly resolves all transitive deps.
     if any(_is_rocm_wheel(w) for w in wheels):
         for repaired_whl in REPAIRED_DIR.glob("*.whl"):
-            _patch_rocjpeg_rpath_in_wheel(repaired_whl)
+            _patch_image_so_rpath_in_wheel(repaired_whl)
 
 
 def repair_macos(wheels):