Skip to content

Commit d7c88aa

Browse files
author
pytorchbot
committed
2026-05-14 nightly release (134250a)
1 parent 5fb36db commit d7c88aa

24 files changed

Lines changed: 1016 additions & 601 deletions
Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
name: Linux CPU aarch64
2+
3+
on:
4+
pull_request:
5+
push:
6+
branches:
7+
- nightly
8+
- main
9+
- release/*
10+
tags:
11+
- v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
12+
workflow_dispatch:
13+
14+
concurrency:
15+
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
16+
cancel-in-progress: true
17+
18+
permissions:
19+
id-token: write
20+
contents: write
21+
22+
defaults:
23+
run:
24+
shell: bash -l -eo pipefail {0}
25+
26+
jobs:
27+
28+
generate-matrix:
29+
uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
30+
with:
31+
package-type: wheel
32+
os: linux-aarch64
33+
test-infra-repository: pytorch/test-infra
34+
test-infra-ref: main
35+
with-xpu: disable
36+
with-rocm: disable
37+
with-cuda: disable
38+
build-python-only: "disable"
39+
40+
build:
41+
needs: generate-matrix
42+
strategy:
43+
fail-fast: false
44+
name: Build and Upload Linux aarch64 wheel
45+
uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
46+
with:
47+
repository: meta-pytorch/torchcodec
48+
ref: ""
49+
test-infra-repository: pytorch/test-infra
50+
test-infra-ref: main
51+
build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
52+
pre-script: packaging/pre_build_script.sh
53+
post-script: packaging/post_build_script.sh
54+
smoke-test-script: packaging/fake_smoke_test.py
55+
package-name: torchcodec
56+
trigger-event: ${{ github.event_name }}
57+
architecture: aarch64
58+
build-platform: "python-build-package"
59+
build-command: "BUILD_AGAINST_ALL_FFMPEG_FROM_S3=1 python -m build --wheel -vvv --no-isolation"
60+
61+
install-and-test:
62+
runs-on: linux.arm64.2xlarge
63+
strategy:
64+
fail-fast: false
65+
matrix:
66+
python-version: ['3.10']
67+
ffmpeg-version-for-tests: ['7.0.1']
68+
needs: build
69+
steps:
70+
- name: Check out repo
71+
uses: actions/checkout@v6
72+
73+
- name: Remove src/ folder
74+
run: bash packaging/remove_src.sh
75+
76+
- name: Setup conda env
77+
uses: conda-incubator/setup-miniconda@v3
78+
with:
79+
auto-update-conda: true
80+
# Using miniforge instead of miniconda ensures that the default
81+
# conda channel is conda-forge instead of main/default. This ensures
82+
# ABI consistency between dependencies:
83+
# https://conda-forge.org/docs/user/transitioning_from_defaults/
84+
miniforge-version: latest
85+
activate-environment: test
86+
python-version: ${{ matrix.python-version }}
87+
88+
- name: Update pip
89+
run: python -m pip install --upgrade pip
90+
91+
- name: Install PyTorch
92+
run: bash packaging/install_pytorch.sh cpu "torch torchvision"
93+
94+
- uses: actions/download-artifact@v4
95+
with:
96+
name: meta-pytorch_torchcodec__${{ matrix.python-version }}_cpu_aarch64
97+
path: dist/
98+
99+
- name: Install torchcodec from the wheel
100+
run: bash packaging/install_torchcodec_wheel.sh
101+
102+
- name: Install ffmpeg, post build
103+
run: bash packaging/install_ffmpeg.sh ${{ matrix.ffmpeg-version-for-tests }}
104+
105+
- name: Install test dependencies
106+
run: bash packaging/install_test_dependencies.sh
107+
108+
- name: Run Python tests
109+
run: |
110+
pytest --override-ini="addopts=-v" test

.github/workflows/linux_cuda_wheel.yaml

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,69 @@ jobs:
132132
run: |
133133
${CONDA_RUN} time python benchmarks/decoders/gpu_benchmark.py --devices=cuda:0,cpu --resize_devices=none
134134
135+
install-and-test-on-cpu-only-machine:
136+
# This job tests that CUDA wheels work fine on CPU-only machines. Note that
137+
# we still install a CUDA-enabled version of torch, and that's by design.
138+
# Essentially, what we want to make sure is that
139+
# `pip install torch torchcodec` works on CPU-only machines, and this
140+
# command should install CUDA-enabled versions of both torch and torchcodec.
141+
# It's critical that this job runs on a CPU-only machine.
142+
runs-on: ubuntu-latest
143+
needs: build
144+
env:
145+
PYTHON_VERSION: '3.10'
146+
CUDA_VERSION: '12.6'
147+
FFMPEG_VERSION: '7'
148+
steps:
149+
- name: Setup env vars
150+
run: |
151+
cuda_version_without_periods=$(echo "${{ env.CUDA_VERSION }}" | sed 's/\.//g')
152+
echo cuda_version_without_periods=${cuda_version_without_periods} >> $GITHUB_ENV
153+
python_version_without_periods=$(echo "${{ env.PYTHON_VERSION }}" | sed 's/\.//g')
154+
echo python_version_without_periods=${python_version_without_periods} >> $GITHUB_ENV
155+
156+
- name: Check out repo
157+
uses: actions/checkout@v6
158+
159+
- name: Remove src/ folder
160+
run: bash packaging/remove_src.sh
161+
162+
- name: Setup conda env
163+
uses: conda-incubator/setup-miniconda@v3
164+
with:
165+
auto-update-conda: true
166+
miniforge-version: latest
167+
activate-environment: test
168+
python-version: ${{ env.PYTHON_VERSION }}
169+
170+
- name: Update pip
171+
run: python -m pip install --upgrade pip
172+
173+
- name: Install PyTorch
174+
run: bash packaging/install_pytorch.sh cu${{ env.cuda_version_without_periods }} "torch torchvision"
175+
176+
- uses: actions/download-artifact@v4
177+
with:
178+
name: meta-pytorch_torchcodec__${{ env.PYTHON_VERSION }}_cu${{ env.cuda_version_without_periods }}_x86_64
179+
path: dist/
180+
181+
- name: Install torchcodec from the wheel
182+
run: bash packaging/install_torchcodec_wheel.sh "*cu${{ env.cuda_version_without_periods }}-cp${{ env.python_version_without_periods }}*.whl"
183+
184+
- name: Install ffmpeg
185+
run: bash packaging/install_ffmpeg.sh ${{ env.FFMPEG_VERSION }}
186+
187+
- name: Install test dependencies
188+
run: bash packaging/install_test_dependencies.sh
189+
190+
- name: Assert CUDA is not available
191+
run: |
192+
python -c "import torch; assert not torch.cuda.is_available()"
193+
194+
- name: Run Python tests
195+
run: |
196+
pytest --override-ini="addopts=-v" test --tb=short
197+
135198
build-docs:
136199
runs-on: linux.g5.4xlarge.nvidia.gpu
137200
env:

packaging/update_ci_for_release.sh

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ done
1818
# This must NOT touch install-and-test-third-party-interface or build-docs.
1919
WHEEL_FILES=(
2020
"${WORKFLOW_DIR}/linux_wheel.yaml"
21+
"${WORKFLOW_DIR}/linux_aarch64_wheel.yaml"
2122
"${WORKFLOW_DIR}/linux_cuda_wheel.yaml"
2223
"${WORKFLOW_DIR}/linux_cuda_aarch64_wheel.yaml"
2324
"${WORKFLOW_DIR}/macos_wheel.yaml"

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
#include "Logging.h"
1717
#include "NVDECCache.h"
1818

19+
#include "NPPRuntimeLoader.h"
1920
#include "NVCUVIDRuntimeLoader.h"
2021
#include "nvcuvid_include/cuviddec.h"
2122
#include "nvcuvid_include/nvcuvid.h"
@@ -69,8 +70,7 @@ static DecoderCapsCache& getDecoderCapsCache() {
6970
return cache;
7071
}
7172

72-
// TODO: rename private variant "default" to "nvdec" to match public name.
73-
static bool g_cuda_default = registerDeviceInterface(
73+
static bool g_cuda_nvdec = registerDeviceInterface(
7474
DeviceInterfaceKey(kStableCUDA, /*variant=*/"default"),
7575
[](const StableDevice& device) {
7676
return new BetaCudaDeviceInterface(device);
@@ -267,12 +267,20 @@ void cudaBufferFreeCallback(void* opaque, [[maybe_unused]] uint8_t* data) {
267267

268268
BetaCudaDeviceInterface::BetaCudaDeviceInterface(const StableDevice& device)
269269
: DeviceInterface(device) {
270-
STD_TORCH_CHECK(
271-
g_cuda_default, "BetaCudaDeviceInterface was not registered!");
270+
STD_TORCH_CHECK(g_cuda_nvdec, "NvdecCudaDeviceInterface was not registered!");
272271
STD_TORCH_CHECK(
273272
device_.type() == kStableCUDA, "Unsupported device: must be CUDA");
274273

275274
initializeCudaContextWithPytorch(device_);
275+
276+
// Note: we could consider *not* erroring when NPP is unavailable, and just
277+
// fallback to the CPU for the color-conversion. This would be similar to what
278+
// we do when NVCUVID is not available (we fallback to the CPU for the
279+
// decoding step).
280+
STD_TORCH_CHECK(
281+
loadNPPLibrary(),
282+
"Failed to load NPP library. NPP is required for CUDA color conversion.");
283+
276284
nppCtx_ = getNppStreamContext(device_);
277285

278286
nvcuvidAvailable_ = loadNVCUVIDLibrary();
@@ -906,7 +914,7 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
906914
// ffmpeg interface does it with maybeConvertAVFrameToNV12OrRGB24().
907915
STD_TORCH_CHECK(
908916
gpuFrame->format == AV_PIX_FMT_CUDA,
909-
"Expected CUDA format frame from BETA CUDA interface");
917+
"Expected CUDA format frame from NVDEC CUDA interface");
910918

911919
cudaStream_t nvdecStream = getCurrentCudaStream(device_.index());
912920

@@ -923,7 +931,7 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
923931
// preAllocatedOutputTensor has post-rotation dimensions, but NV12->RGB
924932
// conversion outputs pre-rotation dimensions, so we can't use it as the
925933
// conversion destination or validate it against the frame shape.
926-
// Once we support native transforms on the default CUDA interface,
934+
// Once we support native transforms on the NVDEC CUDA interface,
927935
// rotation should be handled as part of the transform pipeline instead.
928936
frameOutput.data = convertNV12FrameToRGB(
929937
gpuFrame,
@@ -966,7 +974,7 @@ void BetaCudaDeviceInterface::applyRotation(
966974
}
967975

968976
std::string BetaCudaDeviceInterface::getDetails() {
969-
std::string details = "Beta CUDA Device Interface.";
977+
std::string details = "NVDEC CUDA Device Interface.";
970978
if (cpuFallback_) {
971979
details += " Using CPU fallback.";
972980
if (!nvcuvidAvailable_) {

src/torchcodec/_core/BetaCudaDeviceInterface.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
// This source code is licensed under the BSD-style license found in the
55
// LICENSE file in the root directory of this source tree.
66

7-
// BETA CUDA device interface that provides direct control over NVDEC
7+
// NVDEC CUDA device interface that provides direct control over NVDEC
88
// while keeping FFmpeg for demuxing. A lot of the logic, particularly the use
99
// of a cache for the decoders, is inspired by DALI's implementation which is
1010
// APACHE 2.0:

src/torchcodec/_core/CMakeLists.txt

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -145,20 +145,16 @@ function(make_torchcodec_libraries
145145
)
146146

147147
if(ENABLE_CUDA)
148-
list(APPEND core_sources CudaDeviceInterface.cpp BetaCudaDeviceInterface.cpp NVDECCache.cpp CUDACommon.cpp NVCUVIDRuntimeLoader.cpp)
148+
list(APPEND core_sources CudaDeviceInterface.cpp BetaCudaDeviceInterface.cpp NVDECCache.cpp CUDACommon.cpp NVCUVIDRuntimeLoader.cpp NPPRuntimeLoader.cpp)
149149
endif()
150150

151151
set(core_library_dependencies
152152
${ffmpeg_target}
153153
${TORCH_LIBRARIES}
154154
)
155155

156-
if(ENABLE_CUDA)
157-
list(APPEND core_library_dependencies
158-
${CUDA_nppi_LIBRARY}
159-
${CUDA_nppicc_LIBRARY}
160-
)
161-
endif()
156+
# Note: NPP (nppi, nppicc) is NOT linked here. It is loaded at runtime
157+
# via NPPRuntimeLoader.cpp, following the same pattern as NVCUVID.
162158

163159
make_torchcodec_sublibrary(
164160
"${core_library_name}"

src/torchcodec/_core/CudaDeviceInterface.cpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
#include "Cache.h"
55
#include "CudaDeviceInterface.h"
66
#include "FFMPEGCommon.h"
7+
#include "NPPRuntimeLoader.h"
78
#include "StableABICompat.h"
89
#include "ValidationUtils.h"
910

@@ -99,6 +100,10 @@ CudaDeviceInterface::CudaDeviceInterface(const StableDevice& device)
99100

100101
initializeCudaContextWithPytorch(device_);
101102

103+
STD_TORCH_CHECK(
104+
loadNPPLibrary(),
105+
"Failed to load NPP library. NPP is required for CUDA color conversion.");
106+
102107
hardwareDeviceCtx_ = getHardwareDeviceContext(device_);
103108
nppCtx_ = getNppStreamContext(device_);
104109
}

src/torchcodec/_core/Encoder.cpp

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -784,7 +784,7 @@ void VideoEncoder::initializeEncoder(
784784
auto tensorDevice = frames_.device();
785785
StableDevice stableDevice(
786786
static_cast<StableDeviceType>(tensorDevice.type()), tensorDevice.index());
787-
// The default CUDA interface is decode-only; encoders need the FFmpeg-based
787+
// The NVDEC CUDA interface is decode-only; encoders need the FFmpeg-based
788788
// one.
789789
deviceInterface_ = createDeviceInterface(
790790
stableDevice, stableDevice.type() == kStableCUDA ? "ffmpeg" : "default");
@@ -1055,8 +1055,12 @@ MultiStreamEncoder::~MultiStreamEncoder() {
10551055
close();
10561056
}
10571057

1058-
MultiStreamEncoder::MultiStreamEncoder(std::string_view fileName) {
1058+
MultiStreamEncoder::MultiStreamEncoder() {
10591059
setFFmpegLogLevel();
1060+
}
1061+
1062+
void MultiStreamEncoder::open(std::string_view fileName) {
1063+
STD_TORCH_CHECK(!headerWritten_, "open() was already called.");
10601064

10611065
AVFormatContext* avFormatContext = nullptr;
10621066
int status = avformat_alloc_output_context2(
@@ -1078,13 +1082,17 @@ MultiStreamEncoder::MultiStreamEncoder(std::string_view fileName) {
10781082
fileName,
10791083
", make sure it's a valid path? ",
10801084
getFFMPEGErrorStringFromErrorCode(status));
1085+
1086+
openStreamsAndWriteHeader();
10811087
}
10821088

1083-
MultiStreamEncoder::MultiStreamEncoder(
1089+
void MultiStreamEncoder::open(
10841090
std::string_view formatName,
1085-
std::unique_ptr<AVIOContextHolder> avioContextHolder)
1086-
: avioContextHolder_(std::move(avioContextHolder)) {
1087-
setFFmpegLogLevel();
1091+
std::unique_ptr<AVIOContextHolder> avioContextHolder) {
1092+
STD_TORCH_CHECK(!headerWritten_, "open() was already called.");
1093+
1094+
avioContextHolder_ = std::move(avioContextHolder);
1095+
10881096
// Map mkv -> matroska when used as format name
10891097
formatName = (formatName == "mkv") ? "matroska" : formatName;
10901098
AVFormatContext* avFormatContext = nullptr;
@@ -1101,6 +1109,8 @@ MultiStreamEncoder::MultiStreamEncoder(
11011109
avFormatContext_.reset(avFormatContext);
11021110

11031111
avFormatContext_->pb = avioContextHolder_->getAVIOContext();
1112+
1113+
openStreamsAndWriteHeader();
11041114
}
11051115

11061116
void MultiStreamEncoder::addVideoStream(
@@ -1121,7 +1131,7 @@ void MultiStreamEncoder::addVideoStream(
11211131
STD_TORCH_CHECK(frameRate > 0, "frame_rate must be > 0, got ", frameRate);
11221132
videoStream_ = VideoStream{};
11231133
StableDevice stableDevice(std::move(device));
1124-
// The default CUDA interface is decode-only; encoders need the FFmpeg-based
1134+
// The NVDEC CUDA interface is decode-only; encoders need the FFmpeg-based
11251135
// one.
11261136
videoStream_->deviceInterface = createDeviceInterface(
11271137
stableDevice, stableDevice.type() == kStableCUDA ? "ffmpeg" : "default");
@@ -1383,8 +1393,7 @@ void MultiStreamEncoder::initializeAudioStream() {
13831393
audioStream.avAudioFifo.reset(avAudioFifo);
13841394
}
13851395

1386-
void MultiStreamEncoder::open() {
1387-
STD_TORCH_CHECK(!headerWritten_, "open() was already called.");
1396+
void MultiStreamEncoder::openStreamsAndWriteHeader() {
13881397
STD_TORCH_CHECK(
13891398
videoStream_.has_value() || audioStream_.has_value(),
13901399
"Call addVideoStream() or addAudioStream() before open().");

0 commit comments

Comments
 (0)