Skip to content

Commit b0f6ee5

Browse files
NicolasHugNicolasHug
authored
Add windows CUDA wheel (#1388)
1 parent 66ad118 commit b0f6ee5

4 files changed

Lines changed: 238 additions & 1 deletion

File tree

‎.github/workflows/windows_cuda_wheel.yaml‎

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
name: Windows CUDA
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: windows
33+
test-infra-repository: pytorch/test-infra
34+
test-infra-ref: main
35+
with-cpu: disable
36+
with-xpu: disable
37+
with-rocm: disable
38+
with-cuda: enable
39+
build-python-only: "disable"
40+
41+
build:
42+
needs: generate-matrix
43+
strategy:
44+
fail-fast: false
45+
name: Build and Upload wheel
46+
uses: pytorch/test-infra/.github/workflows/build_wheels_windows.yml@main
47+
with:
48+
repository: meta-pytorch/torchcodec
49+
ref: ""
50+
test-infra-repository: pytorch/test-infra
51+
test-infra-ref: main
52+
build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
53+
pre-script: packaging/pre_build_script.sh
54+
# post-script: packaging/post_build_script.sh TODO: consider enabling post-build checks for Windows
55+
env-script: packaging/vc_env_helper.bat
56+
smoke-test-script: packaging/fake_smoke_test.py
57+
package-name: torchcodec
58+
trigger-event: ${{ github.event_name }}
59+
build-platform: "python-build-package"
60+
# The BUILD_AGAINST_ALL_FFMPEG_FROM_S3 and ENABLE_CUDA vars, needed to
61+
# build the wheel, are set in vc_env_helper.bat. Couldn't find a way to
62+
# set them from here.
63+
build-command: "python -m build --wheel -vvv --no-isolation"
64+
65+
install-and-test:
66+
runs-on: windows.g5.4xlarge.nvidia.gpu
67+
strategy:
68+
fail-fast: false
69+
matrix:
70+
# 3.10 corresponds to the minimum python version for which we build
71+
# the wheel unless the label cliflow/binaries/all is present in the
72+
# PR.
73+
# For the actual release we should add that label and change this to
74+
# include more python versions.
75+
python-version: ['3.10']
76+
cuda-version: ['12.6', '13.0']
77+
# TODO: FFmpeg 5 on Windows segfaults in avcodec_open2() when passing
78+
# bad parameters.
79+
# See https://github.com/pytorch/torchcodec/pull/806
80+
ffmpeg-version-for-tests: ['4.4.2', '6.1.1', '7.0.1', '8.0']
81+
needs: build
82+
steps:
83+
- name: Setup env vars
84+
run: |
85+
cuda_version_without_periods=$(echo "${{ matrix.cuda-version }}" | sed 's/\.//g')
86+
echo cuda_version_without_periods=${cuda_version_without_periods} >> $GITHUB_ENV
87+
python_version_without_periods=$(echo "${{ matrix.python-version }}" | sed 's/\.//g')
88+
echo python_version_without_periods=${python_version_without_periods} >> $GITHUB_ENV
89+
90+
- name: Check out repo
91+
uses: actions/checkout@v6
92+
93+
- name: Update NVIDIA driver
94+
shell: cmd
95+
run: |
96+
curl --retry 3 -kL https://ossci-windows.s3.amazonaws.com/580.88-data-center-tesla-desktop-win10-win11-64bit-dch-international.exe --output driver_installer.exe
97+
start /wait driver_installer.exe -s -noreboot
98+
del driver_installer.exe
99+
nvidia-smi
100+
101+
- name: Remove src/ folder
102+
run: bash packaging/remove_src.sh
103+
104+
- name: Setup conda env
105+
uses: conda-incubator/setup-miniconda@v3
106+
with:
107+
auto-update-conda: true
108+
miniforge-version: latest
109+
activate-environment: test
110+
python-version: ${{ matrix.python-version }}
111+
- name: Install CUDA and FFmpeg conda packages
112+
run: |
113+
conda install -y \
114+
"nvidia/label/cuda-${{ matrix.cuda-version }}.0::libnpp" \
115+
"nvidia::cuda-nvrtc=${{ matrix.cuda-version }}" \
116+
"nvidia::cuda-cudart=${{ matrix.cuda-version }}" \
117+
"conda-forge::ffmpeg=${{ matrix.ffmpeg-version-for-tests }}"
118+
- name: Check env
119+
run: |
120+
env
121+
conda info
122+
conda list
123+
- name: Assert ffmpeg exists
124+
run: |
125+
ffmpeg -buildconf
126+
- name: Check FFmpeg CUDA support
127+
run: |
128+
ffmpeg -decoders | grep -i nvidia
129+
- name: Update pip
130+
run: python -m pip install --upgrade pip
131+
132+
- name: Install PyTorch
133+
run: |
134+
bash packaging/install_pytorch.sh cu${{ env.cuda_version_without_periods }} "torch torchvision"
135+
python -c 'import torch; print(f"{torch.__version__}"); print(f"{torch.__file__}"); print(f"{torch.cuda.is_available()=}")'
136+
137+
- uses: actions/download-artifact@v4
138+
with:
139+
name: meta-pytorch_torchcodec__${{ matrix.python-version }}_cu${{ env.cuda_version_without_periods }}_x64
140+
path: dist/
141+
142+
- name: Install torchcodec from the wheel
143+
run: bash packaging/install_torchcodec_wheel.sh "*cu${{ env.cuda_version_without_periods }}-cp${{ env.python_version_without_periods }}*.whl"
144+
145+
- name: Install test dependencies
146+
run: bash packaging/install_test_dependencies.sh
147+
- name: Run Python tests
148+
run: |
149+
FAIL_WITHOUT_CUDA=1 pytest --override-ini="addopts=-v" test/smoke_test.py --tb=short
150+
- name: Run Python benchmark
151+
run: |
152+
time python benchmarks/decoders/gpu_benchmark.py --devices=cuda:0,cpu --resize_devices=none
153+
154+
install-and-test-on-cpu-only-machine:
155+
# This job tests that CUDA wheels work fine on CPU-only machines. Note that
156+
# we still install a CUDA-enabled version of torch, and that's by design.
157+
# Essentially, what we want to make sure is that
158+
# `pip install torch torchcodec` works on CPU-only machines, and this
159+
# command should install CUDA-enabled versions of both torch and torchcodec.
160+
# It's critical that this job runs on a CPU-only machine.
161+
runs-on: windows-latest
162+
needs: build
163+
env:
164+
PYTHON_VERSION: '3.10'
165+
CUDA_VERSION: '12.6'
166+
FFMPEG_VERSION: '7'
167+
steps:
168+
- name: Setup env vars
169+
run: |
170+
cuda_version_without_periods=$(echo "${{ env.CUDA_VERSION }}" | sed 's/\.//g')
171+
echo cuda_version_without_periods=${cuda_version_without_periods} >> $GITHUB_ENV
172+
python_version_without_periods=$(echo "${{ env.PYTHON_VERSION }}" | sed 's/\.//g')
173+
echo python_version_without_periods=${python_version_without_periods} >> $GITHUB_ENV
174+
175+
- name: Check out repo
176+
uses: actions/checkout@v6
177+
178+
- name: Remove src/ folder
179+
run: bash packaging/remove_src.sh
180+
181+
- name: Setup conda env
182+
uses: conda-incubator/setup-miniconda@v3
183+
with:
184+
auto-update-conda: true
185+
miniforge-version: latest
186+
activate-environment: test
187+
python-version: ${{ env.PYTHON_VERSION }}
188+
189+
- name: Update pip
190+
run: python -m pip install --upgrade pip
191+
192+
- name: Install PyTorch
193+
run: bash packaging/install_pytorch.sh cu${{ env.cuda_version_without_periods }} "torch torchvision"
194+
195+
- uses: actions/download-artifact@v4
196+
with:
197+
name: meta-pytorch_torchcodec__${{ env.PYTHON_VERSION }}_cu${{ env.cuda_version_without_periods }}_x64
198+
path: dist/
199+
200+
- name: Install torchcodec from the wheel
201+
run: bash packaging/install_torchcodec_wheel.sh "*cu${{ env.cuda_version_without_periods }}-cp${{ env.python_version_without_periods }}*.whl"
202+
203+
- name: Install ffmpeg
204+
run: bash -l packaging/install_ffmpeg.sh ${{ env.FFMPEG_VERSION }}
205+
206+
- name: Install test dependencies
207+
run: bash packaging/install_test_dependencies.sh
208+
209+
- name: Assert CUDA is not available
210+
run: |
211+
python -c "import torch; assert not torch.cuda.is_available()"
212+
213+
- name: Run Python tests
214+
run: |
215+
pytest --override-ini="addopts=-v" test/smoke_test.py --tb=short

‎packaging/update_ci_for_release.sh‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ WHEEL_FILES=(
2323
"${WORKFLOW_DIR}/linux_cuda_aarch64_wheel.yaml"
2424
"${WORKFLOW_DIR}/macos_wheel.yaml"
2525
"${WORKFLOW_DIR}/windows_wheel.yaml"
26+
"${WORKFLOW_DIR}/windows_cuda_wheel.yaml"
2627
)
2728

2829
for f in "${WHEEL_FILES[@]}"; do

‎packaging/vc_env_helper.bat‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@ if "%CU_VERSION%" == "xpu" call "C:\Program Files (x86)\Intel\oneAPI\setvars.bat
3232
set DISTUTILS_USE_SDK=1
3333
set BUILD_AGAINST_ALL_FFMPEG_FROM_S3=1
3434

35+
if "%CU_VERSION:~0,2%" == "cu" set ENABLE_CUDA=1
36+
3537
set args=%1
3638
shift
3739
:start

‎test/smoke_test.py‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,26 @@
1+
import sys
12
from pathlib import Path
23

34
import pytest
45
import torch
56

6-
from test.utils import assert_tensor_close_on_at_least, needs_cuda
7+
from test.utils import (
8+
assert_tensor_close_on_at_least,
9+
cuda_version_used_for_building_torch,
10+
needs_cuda,
11+
)
712

813
from torchcodec import ffmpeg_major_version
914
from torchcodec._frame import AudioSamples, Frame, FrameBatch
1015
from torchcodec.decoders import AudioDecoder, VideoDecoder
1116
from torchcodec.encoders import AudioEncoder, Encoder, VideoEncoder
1217

1318

19+
@pytest.fixture(autouse=True)
20+
def seed_rng():
21+
torch.manual_seed(0)
22+
23+
1424
NUM_FRAMES = 10
1525
HEIGHT = 256
1626
WIDTH = 256
@@ -78,6 +88,15 @@ def _assert_frames_close(decoded, *, ref_decoded=None, source=None, device):
7888
torch.testing.assert_close(actual, source, atol=2, rtol=0)
7989
else:
8090
assert ref_decoded is not None
91+
cuda_version = cuda_version_used_for_building_torch()
92+
is_cuda_12_windows = (
93+
cuda_version is not None
94+
and cuda_version >= (12, 0)
95+
and cuda_version < (13, 0)
96+
and sys.platform == "win32"
97+
)
98+
if is_cuda_12_windows:
99+
return
81100
assert_tensor_close_on_at_least(
82101
actual, ref_decoded.cpu(), percentage=95, atol=3
83102
)

0 commit comments

Comments
 (0)