Skip to content

Add ROCJPEG decoder

Add ROCJPEG decoder #7

Workflow file for this run

name: Linux ROCm x86
on:
pull_request:
push:
branches:
- nightly
- main
- release/*
tags:
- v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
- v[0-9]+.[0-9]+.[0-9]+
workflow_dispatch:
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
cancel-in-progress: true
permissions:
id-token: write
contents: read
defaults:
run:
shell: bash -l -eo pipefail {0}
jobs:
generate-matrix:
uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
with:
package-type: wheel
os: linux
test-infra-repository: pytorch/test-infra
test-infra-ref: main
with-cpu: disable
with-xpu: disable
with-cuda: disable
with-rocm: enable
build-python-only: "disable"
build:
needs: generate-matrix
strategy:
fail-fast: false
name: Build and Upload wheel
uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
with:
repository: meta-pytorch/torchcodec
ref: ""
test-infra-repository: pytorch/test-infra
test-infra-ref: main
build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
pre-script: packaging/pre_build_script.sh
post-script: packaging/post_build_script.sh
smoke-test-script: packaging/fake_smoke_test.py
package-name: torchcodec
trigger-event: ${{ github.event_name }}
build-platform: "python-build-package"
# ENABLE_ROCM=1 compiles the rocJPEG GPU JPEG decoder into
# libtorchcodec_image.so (the ROCm counterpart of nvJPEG). Video decoding
# stays on CPU. The FFmpeg core libs are still built against all FFmpeg.
build-command: "BUILD_AGAINST_ALL_FFMPEG_FROM_S3=1 ENABLE_ROCM=1 python -m build --wheel -vvv --no-isolation"
install-and-test:
# ROCm runners are not on EC2 and authenticate to AWS (ECR) via GitHub OIDC.
# Fork PRs don't get an OIDC id-token from GitHub, so the job can't pull the
# docker image and always fails. Skip it for fork PRs; it still runs on
# same-repo PRs, push, and workflow_dispatch.
if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository
needs: build
strategy:
fail-fast: false
matrix:
python-version: ['3.10']
rocm-version: ['7.1']
ffmpeg-version-for-tests: ['7']
uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
permissions:
id-token: write
contents: read
with:
repository: meta-pytorch/torchcodec
runner: "linux.rocm.gpu.gfx950.1"
gpu-arch-type: "rocm"
gpu-arch-version: ${{ matrix.rocm-version }}
# ROCm self-hosted runners are not on EC2, so they can't pull the default
# ECR-hosted docker image (no AWS credentials). Use a Docker Hub image and
# run without sudo, matching pytorch/vision's ROCm job.
docker-image: pytorch/manylinux2_28-builder:rocm${{ matrix.rocm-version }}
no-sudo: true
timeout: 120
test-infra-ref: main
download-artifact: meta-pytorch_torchcodec__${{ matrix.python-version }}_rocm${{ matrix.rocm-version }}_x86_64
script: |
set -euo pipefail
python -m pip install --upgrade pip
echo '::group::Install PyTorch (ROCm)'
pip install --progress-bar=off --pre torch \
--index-url https://download.pytorch.org/whl/nightly/rocm${{ matrix.rocm-version }}
echo '::endgroup::'
echo '::group::Assert we are on ROCm'
rocm-smi
python -c "
import torch
print('torch:', torch.__version__, '| hip:', torch.version.hip)
assert torch.version.hip is not None, 'torch is not a ROCm/HIP build'
assert torch.cuda.is_available(), 'no GPU visible to torch'
print('device:', torch.cuda.get_device_name(0))
"
echo '::endgroup::'
# Prevent the checked-out src/ tree from shadowing the installed wheel.
bash packaging/remove_src.sh
echo '::group::Install rocJPEG runtime'
# The wheel deliberately does not bundle librocjpeg (see repair_wheel.py),
# so decode_jpeg(device="cuda") needs it present at runtime.
bash packaging/install_rocjpeg.sh
echo '::endgroup::'
echo '::group::Install torchcodec from the wheel'
python -m pip install "${RUNNER_ARTIFACT_DIR}"/*.whl -vvv
echo '::endgroup::'
echo '::group::Install FFmpeg and test dependencies'
bash packaging/install_ffmpeg.sh ${{ matrix.ffmpeg-version-for-tests }}
bash packaging/install_test_dependencies.sh
echo '::endgroup::'
echo '::group::Run Python tests'
# torch.cuda.is_available() is True on ROCm (HIP masquerades as CUDA), so
# conftest.py won't auto-skip needs_cuda tests. This wheel has no GPU
# decode compiled, so we deselect them explicitly. When the rocJPEG
# decoder lands, its GPU tests get wired in here.
FAIL_WITHOUT_JPEG=1 FAIL_WITHOUT_PNG=1 FAIL_WITHOUT_WEBP=1 FAIL_WITHOUT_AVIF=1 \
pytest -m "not needs_cuda" --override-ini="addopts=-v" test --tb=short
echo '::endgroup::'