Disable rocJPEG for ROCm wheel builds

[atalman]

Problem

The ROCm 7.14 torchvision build produces a wheel whose image extension cannot be
imported at all. Smoke test
(test-infra run):

torchvision: 0.29.0.dev20260806+rocm7.14
AttributeError: '_OpNamespace' 'image' object has no attribute '_jpeg_version'

torch.ops.image is empty because torchvision/image_stable.so fails to
dlopen, and torchvision swallows the load error.

Cause

setup.py:449 enables rocJPEG whenever the header is present:

if USE_ROCJPEG and IS_ROCM and (torch.cuda.is_available() or FORCE_CUDA):
    rocjpeg_found = ROCM_HOME is not None and (Path(ROCM_HOME) / "include/rocjpeg/rocjpeg.h").exists()

so image_stable.so links -lrocjpeg. The repair step then drops it:

Relocating image_stable.so
  libpng16.so.16 / libjpeg.so.8 / libwebp.so.7   -> bundled
  librocjpeg.so.1     -> Omitting
  libamdhip64.so.7 / libc10.so / libtorch.so / libtorch_hip.so -> Omitting

Omitting the torch libs is fine — they resolve through torchvision's RPATH into
torch/lib. librocjpeg is different: it 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 the extension is unloadable, and all torchvision.io image ops vanish —
CPU jpeg/png/webp included, not just the GPU path.

Why 7.2 is green today

By accident, not by design. Both builds print USE_ROCJPEG = True; the
difference is only whether the header exists in the builder image:

rocm7.2  : Building torchvision without ROCJPEG support     <- rocm/dev-almalinux-8 has no header
rocm7.14 : Building torchvision with ROCJPEG image support  <- TheRock image ships it in _rocm_sdk_devel

Nothing sets TORCHVISION_USE_ROCJPEG anywhere in this repo or in test-infra
today. This PR makes the existing 7.2 behaviour explicit and applies it to 7.14
too, so the two stay consistent as ROCm 7.1/7.2 → 7.2/7.14 lands
(pytorch/test-infra#8451).

Note on the mechanism

The export is appended to BUILD_ENV_FILE rather than just exported: the
pre-script is invoked as a subprocess (${CONDA_RUN} bash ${SCRIPT}) in its own
workflow step, so a plain export would not reach the build. Both
"Build the wheel" steps in build_wheels_linux.yml do source "${BUILD_ENV_FILE}",
which is how FORCE_CUDA and friends already propagate. The guard on
${BUILD_ENV_FILE:-} keeps the script usable when run by hand.

Test plan

Gate exercised directly:

CU_VERSION	appended to BUILD_ENV_FILE
rocm7.2	export TORCHVISION_USE_ROCJPEG=0
rocm7.14	export TORCHVISION_USE_ROCJPEG=0
cu126 / cpu / xpu	(nothing)

• bash -n and shellcheck -S warning clean.
• With BUILD_ENV_FILE unset the block is a no-op and does not error under
  set -u.
• TORCHVISION_USE_ROCJPEG=0 makes USE_ROCJPEG False, so the whole rocJPEG
  block in setup.py is skipped: no -lrocjpeg, no ROCJPEG_FOUND, and the
  image extension goes back to the CPU-only link line that ROCm 7.2 ships today.

Follow-up

If GPU JPEG on ROCm is wanted, the fix is to make librocjpeg.so.1 resolvable —
bundle it into torchvision.libs/, or add an RPATH from torchvision/ to the
ROCm SDK wheels mirroring repair_wheel.py::rocm_rpaths() — rather than
re-enabling the flag as-is.

cc @jeffdaily @jithunnair-amd

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9586

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[NicolasHug]

SGTM, I am going to remove the entire ROCJPEG code from TorchVision soon anyway. As discussed on slack, I'm migrating it to TorchCodec in meta-pytorch/torchcodec#1554 (would love some help with the CI there!)