From 16d78792b86173fd31c5eeb8e383c5e2e89c98cc Mon Sep 17 00:00:00 2001 From: Akash Verma Date: Mon, 10 Aug 2026 22:10:49 +0000 Subject: [PATCH 1/5] Bundle librocjpeg into ROCm wheels and set RPATH for zero-config runtime resolution. --- packaging/repair_wheel.py | 231 +++++++++++++++++++++++++++++++++++--- 1 file changed, 218 insertions(+), 13 deletions(-) diff --git a/packaging/repair_wheel.py b/packaging/repair_wheel.py index 01b97ea9b..183d0c806 100644 --- a/packaging/repair_wheel.py +++ b/packaging/repair_wheel.py @@ -41,6 +41,11 @@ def _is_cuda_wheel(wheel): return re.search(r"[+_]cu\d", Path(wheel).name) is not None +def _is_rocm_wheel(wheel): + # Detect a ROCm wheel from its local-version tag (e.g. "+rocm7.14") in the filename. + return re.search(r"[+_]rocm", Path(wheel).name) is not None + + def run(cmd, **kwargs): cmd = [str(c) for c in cmd] print("+ " + " ".join(cmd), flush=True) @@ -127,6 +132,155 @@ def _find_nvjpeg_license(): return None +def _find_rocjpeg_license(): + """Find rocjpeg's LICENSE file to ship alongside the bundled binary.""" + search_roots = [] + for var in ("ROCM_HOME", "ROCM_PATH"): + if v := os.environ.get(var): + search_roots.append(Path(v)) + try: + result = subprocess.run( + [ + sys.executable, + "-c", + "from torch.utils.cpp_extension import ROCM_HOME; print(ROCM_HOME or '')", + ], + capture_output=True, + text=True, + check=False, + ) + if result.returncode == 0 and result.stdout.strip(): + search_roots.append(Path(result.stdout.strip())) + except Exception: + pass + search_roots.append(Path("/opt/rocm")) + + for root in search_roots: + candidate = root / "share" / "doc" / "rocjpeg" / "LICENSE" + if candidate.is_file(): + return candidate + return None + + +def _find_rocjpeg_lib(): + """Find librocjpeg.so at wheel repair time so auditwheel can bundle it. + + Searches ROCM_HOME / ROCM_PATH env vars, torch's ROCM_HOME, and the + standard /opt/rocm fallback. + """ + search_roots = [] + for var in ("ROCM_HOME", "ROCM_PATH"): + if v := os.environ.get(var): + search_roots.append(Path(v)) + # Ask torch where it found ROCm at its own build time. + try: + result = subprocess.run( + [ + sys.executable, + "-c", + "from torch.utils.cpp_extension import ROCM_HOME; print(ROCM_HOME or '')", + ], + capture_output=True, + text=True, + check=False, + ) + if result.returncode == 0 and result.stdout.strip(): + search_roots.append(Path(result.stdout.strip())) + except Exception: + pass + search_roots.append(Path("/opt/rocm")) + + for root in search_roots: + for lib_dir in (root / "lib", root / "lib64"): + candidate = lib_dir / "librocjpeg.so.1" + if not candidate.exists(): + # Try unversioned symlink + candidate = lib_dir / "librocjpeg.so" + if candidate.exists(): + return lib_dir + 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. + + Two layouts are covered: + - ROCm >= 7.14 (TheRock / rocm-sdk-* Python wheels): + libamdhip64 etc. live in /_rocm_sdk_core/lib/. + $ORIGIN/../../_rocm_sdk_core/lib reaches that dir from + /torchcodec.libs/librocjpeg-HASH.so.1. + - 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. + """ + patchelf = shutil.which("patchelf") + if not patchelf: + raise RuntimeError( + "patchelf not found; install it (pip install patchelf) before " + "repairing ROCm wheels." + ) + + rpath = ":".join([ + "$ORIGIN", + # ROCm >= 7.14: rocm-sdk-core Python wheel layout + "$ORIGIN/../../_rocm_sdk_core/lib", + "$ORIGIN/../../_rocm_sdk_core/lib/rocm_sysdeps/lib", + # ROCm <= 7.2: standard system install path (AMD installer always + # creates /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) + 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) + + # Update the RECORD file so pip's integrity check passes. + # RECORD format: path,sha256=,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} + for record_file in record_files: + lines = record_file.read_text(encoding="utf-8").splitlines() + new_lines = [] + for line in lines: + parts = line.split(",") + if len(parts) >= 3 and parts[0] in patched_rel_names: + data = (tmp_path / parts[0]).read_bytes() + h = base64.urlsafe_b64encode( + hashlib.sha256(data).digest() + ).rstrip(b"=").decode() + new_lines.append(f"{parts[0]},sha256={h},{len(data)}") + else: + new_lines.append(line) + record_file.write_text("\n".join(new_lines) + "\n", encoding="utf-8") + + # Repack wheel preserving zip metadata. + patched_path = wheel_path.with_suffix(".patched.whl") + with zipfile.ZipFile(wheel_path, "r") as src_zf, \ + zipfile.ZipFile(patched_path, "w", compression=zipfile.ZIP_DEFLATED) as dst_zf: + for item in src_zf.infolist(): + patched_file = tmp_path / item.filename + if patched_file.is_file(): + dst_zf.write(patched_file, item.filename) + else: + # Directory entries or missing files: copy as-is + dst_zf.writestr(item, src_zf.read(item.filename)) + wheel_path.unlink() + patched_path.rename(wheel_path) + + def repair_linux(wheels): run([sys.executable, "-m", "pip", "install", "--upgrade", "auditwheel"]) run(["auditwheel", "--version"]) @@ -134,11 +288,22 @@ def repair_linux(wheels): # for auditwheel to graft libs, it must be able to find them, so we set # LD_LIBRARY_PATH: jpeg/png/webp are from conda, libavif is from the S3 # build dir, and (for CUDA wheels) libnvjpeg is from the CUDA toolkit. + # For ROCm wheels, librocjpeg comes from the ROCm install. lib_dirs = [str(_avif_lib_dir())] if conda_prefix := env.get("CONDA_PREFIX"): lib_dirs.append(str(Path(conda_prefix) / "lib")) if any(_is_cuda_wheel(w) for w in wheels): lib_dirs.extend(sorted({str(f.parent) for f in _find_nvjpeg_libs()})) + if any(_is_rocm_wheel(w) for w in wheels): + if rocjpeg_lib_dir := _find_rocjpeg_lib(): + lib_dirs.append(str(rocjpeg_lib_dir)) + print(f"Found librocjpeg in {rocjpeg_lib_dir}", flush=True) + else: + print( + "WARNING: librocjpeg not found; rocJPEG will not be bundled. " + "Set ROCM_HOME or ROCM_PATH if ROCm is in a non-standard location.", + flush=True, + ) env["LD_LIBRARY_PATH"] = os.pathsep.join( [*lib_dirs, env.get("LD_LIBRARY_PATH", "")] ) @@ -168,17 +333,11 @@ def repair_linux(wheels): "libnvshmem*", "libnvfatbin*", "libnvcuvid*", - # rocJPEG: the GPU JPEG decoder our image lib links on ROCm. Unlike - # nvJPEG (which we bundle), rocJPEG is not shipped by the torch-ROCm - # wheel, and bundling it would drag in torch's ROCm libs under mismatched - # (hashed) sonames. So we treat it as a runtime dependency provided by the - # ROCm install, like FFmpeg. decode_jpeg(device='cuda') therefore needs - # ROCm (with rocJPEG) present at runtime. - # TODO_ROCM: Should we still try to ship librocjpeg? - "librocjpeg*", # ROCm/HIP runtime and its system deps: provided by the torch-ROCm wheel - # (torch/lib/) at runtime, exactly like the CUDA libs above. Never bundle - # them, they'd duplicate torch's copies and bloat the 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*", @@ -204,6 +363,13 @@ def repair_linux(wheels): "libelf*", "libbz2*", "liblzma*", + # rocm_sysdeps_* are vendored system libs bundled inside rocm-sdk-core; + # librocm_kpack, libLLVM, libclang-cpp are pulled in transitively by + # libamd_comgr. All resolved at runtime via _rocm_sdk_core or /opt/rocm. + "librocm_sysdeps_*", + "librocm_kpack*", + "libLLVM*", + "libclang-cpp*", ): excludes += ["--exclude", pattern] for wheel in wheels: @@ -212,6 +378,13 @@ 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. + 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) + def repair_macos(wheels): run([sys.executable, "-m", "pip", "install", "--upgrade", "delocate"]) @@ -429,6 +602,16 @@ def _resolve_avif_licenses(): licenses["LICENSE.libnvjpeg-NVIDIA-CUDA-EULA.txt"] = nvjpeg_license print(f" LICENSE.libnvjpeg-NVIDIA-CUDA-EULA.txt <- {nvjpeg_license}") + if _is_rocm_wheel(wheel): + if (rocjpeg_license := _find_rocjpeg_license()) is None: + raise RuntimeError( + f"{wheel.name} bundles librocjpeg but the rocjpeg LICENSE " + "could not be located to ship alongside it. " + "Set ROCM_HOME or ROCM_PATH to the ROCm install root." + ) + licenses["LICENSE.librocjpeg-MIT.txt"] = rocjpeg_license + print(f" LICENSE.librocjpeg-MIT.txt <- {rocjpeg_license}") + unpack_dir = scratch / "unpack" if unpack_dir.is_dir(): shutil.rmtree(unpack_dir) @@ -515,6 +698,9 @@ def _is_avif(lib): stem.startswith("avif") and stem.endswith(".dll") ) + def _is_rocjpeg(lib): + return lib.startswith("librocjpeg") + def _is_nvjpeg(lib): return lib.startswith("libnvjpeg") or ( lib.startswith("nvjpeg") and lib.endswith(".dll") @@ -550,6 +736,7 @@ def _is_allowed(lib): or _is_webp(lib) or _is_avif(lib) or _is_nvjpeg(lib) + or _is_rocjpeg(lib) ): return True if platform.system() == "Darwin" and lib.startswith(("libc++", "libpython")): @@ -631,7 +818,7 @@ def _assert_linux_libjpeg_is_turbo(zf): "found at build time." ) - def _assert_third_party_licenses(zf, is_cuda): + def _assert_third_party_licenses(zf, is_cuda, is_rocm): """Every bundled third-party lib must ship its license text under .dist-info/licenses/third_party/ (see bundle_third_party_licenses).""" license_files = [ @@ -641,9 +828,12 @@ def _assert_third_party_licenses(zf, is_cuda): ] # keyword each bundled lib's license file must be identifiable by. CUDA # wheels also bundle libnvjpeg, whose NVIDIA CUDA EULA must ship too. + # ROCm wheels bundle librocjpeg, whose MIT license must ship too. keywords = ["jpeg", "png", "zlib", "webp", "avif", "dav1d", "yuv"] if is_cuda: keywords.append("nvjpeg") + if is_rocm: + keywords.append("rocjpeg") for keyword in keywords: if not any(keyword in n.lower() for n in license_files): raise RuntimeError( @@ -654,7 +844,9 @@ def _assert_third_party_licenses(zf, is_cuda): for wheel in DIST_DIR.glob("*.whl"): print(f"Checking bundled libraries in {wheel.name}") with zipfile.ZipFile(wheel) as zf: - _assert_third_party_licenses(zf, _is_cuda_wheel(wheel)) + is_cuda = _is_cuda_wheel(wheel) + is_rocm = _is_rocm_wheel(wheel) + _assert_third_party_licenses(zf, is_cuda, is_rocm) names = zf.namelist() libs = sorted({n.rsplit("/", 1)[-1] for n in names if _is_shared_lib(n)}) if unexpected := [lib for lib in libs if not _is_allowed(lib)]: @@ -676,7 +868,9 @@ def _assert_third_party_licenses(zf, is_cuda): "animated webp decoding)." ) is_cuda = _is_cuda_wheel(wheel) + is_rocm = _is_rocm_wheel(wheel) bundles_nvjpeg = any(_is_nvjpeg(lib) for lib in libs) + bundles_rocjpeg = any(_is_rocjpeg(lib) for lib in libs) if is_cuda and not bundles_nvjpeg: raise RuntimeError( f"{wheel.name} is a CUDA wheel but does not bundle libnvjpeg. " @@ -688,6 +882,17 @@ def _assert_third_party_licenses(zf, is_cuda): raise RuntimeError( f"{wheel.name} is not a CUDA wheel but bundles libnvjpeg." ) + if is_rocm and not bundles_rocjpeg: + raise RuntimeError( + f"{wheel.name} is a ROCm wheel but does not bundle librocjpeg. " + "GPU JPEG decoding (decode_jpeg(..., device='cuda')) needs it. " + "Check that librocjpeg is findable at repair time " + "(set ROCM_HOME or ROCM_PATH) and not excluded." + ) + if not is_rocm and bundles_rocjpeg: + raise RuntimeError( + f"{wheel.name} is not a ROCm wheel but bundles librocjpeg." + ) if encoders := [lib for lib in libs if _is_avif_encoder(lib)]: raise RuntimeError( f"{wheel.name} bundles AV1 codec libraries that must not " @@ -700,7 +905,7 @@ def _assert_third_party_licenses(zf, is_cuda): "ship (libheif is a user-supplied runtime dependency, like " "FFmpeg): " + " ".join(lgpl) ) - MAX_WHEEL_BYTES = (14 if is_cuda else 6) * 1024 * 1024 + MAX_WHEEL_BYTES = (14 if is_cuda else 7 if is_rocm else 6) * 1024 * 1024 wheel_bytes = wheel.stat().st_size if wheel_bytes > MAX_WHEEL_BYTES: raise RuntimeError( From d99b6d58c36ab9d93789441c8d30d53831387c9b Mon Sep 17 00:00:00 2001 From: Akash Verma Date: Mon, 17 Aug 2026 19:19:30 +0000 Subject: [PATCH 2/5] Add VA-API and DRM diagnostics to ROCm CI to debug rocJPEG init failure Co-authored-by: Cursor --- .github/workflows/linux_rocm.yaml | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/.github/workflows/linux_rocm.yaml b/.github/workflows/linux_rocm.yaml index 1659434b2..31d881e84 100644 --- a/.github/workflows/linux_rocm.yaml +++ b/.github/workflows/linux_rocm.yaml @@ -144,6 +144,32 @@ jobs: bash packaging/install_rocjpeg.sh echo '::endgroup::' + echo '::group::VA-API and DRM diagnostics' + # libva libraries present in the container + ldconfig -p | grep -E "libva|libdrm" || true + # DRM render nodes - permissions and group ownership + ls -la /dev/dri/ || true + # Current user and groups (container user must be in video/render group) + id + # Check whether mesa-amdgpu-va-drivers was actually installed + # (install_rocjpeg.sh silently falls back and skips it if AMD repo is absent) + rpm -q mesa-amdgpu-va-drivers libva-amdgpu rocjpeg 2>&1 || true + # Find VA driver .so files - if amdgpu_drv_video.so is missing, vaInit will fail + find /usr -name "*drv_video*" 2>/dev/null || true + # Try vainfo on every render node (no-sudo: true so we can't install vainfo) + if command -v vainfo &>/dev/null; then + for node in /dev/dri/renderD*; do + echo "--- vainfo on ${node} ---" + vainfo --display drm --device "${node}" 2>&1 || true + done + else + echo "vainfo not installed, skipping VA-API smoke test" + fi + # Show which VA driver is selected + echo "LIBVA_DRIVER_NAME=${LIBVA_DRIVER_NAME:-}" + echo "LIBVA_DRIVERS_PATH=${LIBVA_DRIVERS_PATH:-}" + echo '::endgroup::' + echo '::group::Install torchcodec from the wheel' python -m pip install "${RUNNER_ARTIFACT_DIR}"/*.whl -vvv echo '::endgroup::' From 456e2748020b7f0d9ef89246f9f9e64423e3d85d Mon Sep 17 00:00:00 2001 From: Akash Verma Date: Mon, 17 Aug 2026 19:21:48 +0000 Subject: [PATCH 3/5] Switch ROCm CI test job from 7.1 to 7.2 Co-authored-by: Cursor --- .github/workflows/linux_rocm.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/linux_rocm.yaml b/.github/workflows/linux_rocm.yaml index 31d881e84..9940b7ed0 100644 --- a/.github/workflows/linux_rocm.yaml +++ b/.github/workflows/linux_rocm.yaml @@ -72,7 +72,7 @@ jobs: fail-fast: false matrix: python-version: ['3.10'] - rocm-version: ['7.1'] + rocm-version: ['7.2'] uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write From 6811de8d2899adb2623255a34356a2456df9c4a3 Mon Sep 17 00:00:00 2001 From: Akash Verma Date: Tue, 18 Aug 2026 16:59:40 +0000 Subject: [PATCH 4/5] Exclude libva* from auditwheel bundling to fix ROCm wheel repair librocjpeg links against libva.so.2 and libva-drm.so.2 (VA-API, used by rocJPEG's HYBRID GPU-JPEG backend). auditwheel was pulling both into the wheel because they weren't excluded, and repair_wheel.py's check_bundling() then rejected them as unexpected. libva/libva-drm are system-provided display-stack libraries, present on any ROCm install alongside libdrm (which is already excluded). Add "libva*" to the same exclude list so auditwheel leaves them on the system rather than bundling them. Co-authored-by: Cursor --- packaging/repair_wheel.py | 1 + 1 file changed, 1 insertion(+) diff --git a/packaging/repair_wheel.py b/packaging/repair_wheel.py index 183d0c806..dc58bf9c7 100644 --- a/packaging/repair_wheel.py +++ b/packaging/repair_wheel.py @@ -360,6 +360,7 @@ def repair_linux(wheels): "librccl*", "libnuma*", "libdrm*", + "libva*", # VA-API libs pulled in by librocjpeg's HYBRID backend; system-provided alongside libdrm "libelf*", "libbz2*", "liblzma*", From 7f42e27bf60d7d265691ade5fd6237df7e9a04c9 Mon Sep 17 00:00:00 2001 From: Akash Verma Date: Tue, 18 Aug 2026 19:24:33 +0000 Subject: [PATCH 5/5] Skip dnf rocjpeg install when already present via ROCm 7.14 pip wheels ROCm >= 7.14 distributes the full ROCm stack (including rocJPEG) as pip wheels (_rocm_sdk_core / _rocm_sdk_devel site-packages) rather than system RPMs. The rocjpeg-devel, libva-amdgpu and mesa-amdgpu-va-drivers DNF packages therefore don't exist on the ROCm 7.14 builder image and the install was failing with "No package rocjpeg available". Check for librocjpeg.so under /opt/conda (pip-wheel install path) and /opt/rocm (classic RPM install path) before attempting dnf install. If already present, skip the install entirely. The existing dnf path is preserved for ROCm 7.2 and earlier where RPMs are the only source. Co-authored-by: Cursor --- packaging/install_rocjpeg.sh | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/packaging/install_rocjpeg.sh b/packaging/install_rocjpeg.sh index 0d99837e4..b795bc299 100755 --- a/packaging/install_rocjpeg.sh +++ b/packaging/install_rocjpeg.sh @@ -20,6 +20,11 @@ set -euo pipefail # test runners). On the plain build image that repo is absent; there we only need # to *compile* against rocjpeg.h/librocjpeg.so, so fall back to installing just # those with --nodeps (base libva provides the libva.so.2 soname we link). +# +# ROCm >= 7.14 distributes the full ROCm stack (including rocJPEG) as pip wheels +# (_rocm_sdk_core / _rocm_sdk_devel site-packages). In that case librocjpeg.so +# and rocjpeg.h are already present and the dnf packages don't exist, so we +# skip the install entirely. install_rocjpeg_build_only() { dnf install -y --refresh libva dnf install -y "dnf-command(download)" >/dev/null 2>&1 || dnf install -y dnf-plugins-core @@ -28,5 +33,16 @@ install_rocjpeg_build_only() { rpm -Uvh --nodeps "${rpm_dir}"/rocjpeg*.rpm } -dnf install -y --refresh rocjpeg-devel libva-amdgpu mesa-amdgpu-va-drivers \ - || install_rocjpeg_build_only +# Check if librocjpeg is already available (e.g. via ROCm 7.14+ pip wheels). +if python3 -c " +import glob, sys +# _rocm_sdk_core and _rocm_sdk_devel are the pip-wheel-based ROCm installs +hits = (glob.glob('/opt/conda/**/librocjpeg.so*', recursive=True) + + glob.glob('/opt/rocm/lib/librocjpeg.so*')) +sys.exit(0 if hits else 1) +" 2>/dev/null; then + echo "librocjpeg already present (ROCm pip-wheel install); skipping dnf install." +else + dnf install -y --refresh rocjpeg-devel libva-amdgpu mesa-amdgpu-va-drivers \ + || install_rocjpeg_build_only +fi