From 3103a71f5a6ee8f2d0f279488fbc8fa8c609c591 Mon Sep 17 00:00:00 2001 From: Andrey Talman Date: Thu, 6 Aug 2026 08:38:51 -0700 Subject: [PATCH] Disable rocJPEG for ROCm wheel builds The ROCm 7.14 build of torchvision produces a wheel whose image extension cannot be imported at all: AttributeError: '_OpNamespace' 'image' object has no attribute '_jpeg_version' setup.py enables rocJPEG whenever $ROCM_HOME/include/rocjpeg/rocjpeg.h is present, which links torchvision/image_stable.so against librocjpeg.so.1. The wheel repair step then omits librocjpeg ("Omitting librocjpeg.so.1"), as it does for libtorch/libc10/libamdhip64 -- but unlike those, librocjpeg is not in the torch wheel either: pytorch's repair_wheel.py has no rocjpeg in ROCM_SO_FILES, and on the ROCm 7.14 TheRock layout it bundles no ROCm libs at all (they live in the separate _rocm_sdk_* wheels, which torchvision has no RPATH to). So image_stable.so fails to dlopen and every torchvision.io image op disappears -- CPU jpeg/png/webp included, not just the GPU path. ROCm 7.2 is unaffected only by accident: rocm/dev-almalinux-8 does not ship the rocjpeg header, so that build already prints "Building torchvision without ROCJPEG support". Setting the flag for all ROCm builds makes that intentional and keeps 7.2 and 7.14 consistent. The export is written to BUILD_ENV_FILE because the pre-script runs as a subprocess (${CONDA_RUN} bash ${SCRIPT}); the build steps source that file. If GPU JPEG on ROCm is wanted later, the fix is to make librocjpeg.so.1 resolvable -- bundle it into torchvision.libs or add an RPATH to the ROCm SDK wheels -- rather than to re-enable the flag as-is. --- packaging/pre_build_script.sh | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/packaging/pre_build_script.sh b/packaging/pre_build_script.sh index 7167d813399..f1337233103 100644 --- a/packaging/pre_build_script.sh +++ b/packaging/pre_build_script.sh @@ -38,5 +38,21 @@ else pip install "auditwheel<6.3.0" fi +# Build the image extension without rocJPEG on ROCm. setup.py enables rocJPEG +# whenever $ROCM_HOME/include/rocjpeg/rocjpeg.h exists, which links +# torchvision/image_stable.so against librocjpeg.so.1. That library is neither +# bundled into the torchvision wheel (the repair step omits it) nor shipped in +# the torch wheel, so the extension fails to dlopen at import and every +# torchvision.io image op silently disappears -- CPU jpeg/png/webp included. +# Written to BUILD_ENV_FILE because this script runs as a subprocess; the build +# steps source that file. +if [[ "${CU_VERSION:-}" == rocm* ]]; then + echo "Disabling rocJPEG support for ${CU_VERSION}" + export TORCHVISION_USE_ROCJPEG=0 + if [[ -n "${BUILD_ENV_FILE:-}" ]]; then + echo "export TORCHVISION_USE_ROCJPEG=0" >> "${BUILD_ENV_FILE}" + fi +fi + pip install numpy pyyaml future ninja pip install --upgrade setuptools==72.1.0