Skip to content

Commit 773569e

Browse files
vmoenscursoragent
andcommitted
[CI] Add CPU-only build option and fix release workflow
- Add with-cuda, with-rocm, with-cpu inputs to build-wheels-linux.yml - Add with-cuda, with-cpu inputs to build-wheels-windows.yml - Configure release workflow to build CPU-only wheels (TorchRL is pure Python) - Fix test-infra-ref to use 'main' (not pytorch_release which is wrong repo) CPU-only builds are much faster and produce the same wheel since TorchRL has no CUDA-specific compiled code. Co-authored-by: Cursor <[email protected]>
1 parent 3dcbfa7 commit 773569e

File tree

4 files changed

+282
-0
lines changed

4 files changed

+282
-0
lines changed
Lines changed: 245 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,245 @@
1+
name: Build Docker Images
2+
3+
on:
4+
# Manual trigger
5+
workflow_dispatch:
6+
inputs:
7+
image_type:
8+
description: 'Image type to build'
9+
required: true
10+
type: choice
11+
options:
12+
- all
13+
- base
14+
- nightly
15+
- stable
16+
- habitat
17+
force_rebuild:
18+
description: 'Force rebuild even if unchanged'
19+
required: false
20+
type: boolean
21+
default: false
22+
23+
# Weekly schedule for nightly images (Sundays at 2 AM UTC)
24+
schedule:
25+
- cron: '0 2 * * 0'
26+
27+
# Trigger on changes to Docker files
28+
push:
29+
branches:
30+
- main
31+
- docker-images
32+
paths:
33+
- '.github/docker/**'
34+
- '.github/workflows/build-docker-images.yml'
35+
36+
permissions:
37+
contents: read
38+
packages: write
39+
40+
env:
41+
REGISTRY: ghcr.io
42+
IMAGE_NAME: pytorch/torchrl-ci
43+
44+
jobs:
45+
# Matrix strategy for building multiple images
46+
build-base:
47+
if: |
48+
github.event_name == 'schedule' ||
49+
github.event_name == 'push' ||
50+
(github.event_name == 'workflow_dispatch' &&
51+
(github.event.inputs.image_type == 'all' || github.event.inputs.image_type == 'base'))
52+
runs-on: ubuntu-latest
53+
strategy:
54+
matrix:
55+
cuda_version: ['12.4.0', '11.8.0']
56+
python_version: ['3.11', '3.9']
57+
include:
58+
- cuda_version: '12.4.0'
59+
cuda_short: 'cuda12.4'
60+
- cuda_version: '11.8.0'
61+
cuda_short: 'cuda11.8'
62+
steps:
63+
- name: Checkout repository
64+
uses: actions/checkout@v4
65+
66+
- name: Set up Docker Buildx
67+
uses: docker/setup-buildx-action@v3
68+
69+
- name: Log in to GitHub Container Registry
70+
uses: docker/login-action@v3
71+
with:
72+
registry: ${{ env.REGISTRY }}
73+
username: ${{ github.actor }}
74+
password: ${{ secrets.GITHUB_TOKEN }}
75+
76+
- name: Extract metadata
77+
id: meta
78+
uses: docker/metadata-action@v5
79+
with:
80+
images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
81+
tags: |
82+
type=raw,value=base-${{ matrix.cuda_short }}-py${{ matrix.python_version }}
83+
84+
- name: Build and push base image
85+
uses: docker/build-push-action@v5
86+
with:
87+
context: .github/docker
88+
file: .github/docker/Dockerfile.base
89+
push: true
90+
tags: ${{ steps.meta.outputs.tags }}
91+
labels: ${{ steps.meta.outputs.labels }}
92+
build-args: |
93+
CUDA_VERSION=${{ matrix.cuda_version }}
94+
PYTHON_VERSION=${{ matrix.python_version }}
95+
cache-from: type=registry,ref=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:base-${{ matrix.cuda_short }}-py${{ matrix.python_version }}-cache
96+
cache-to: type=registry,ref=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:base-${{ matrix.cuda_short }}-py${{ matrix.python_version }}-cache,mode=max
97+
98+
build-nightly:
99+
needs: build-base
100+
if: |
101+
always() &&
102+
(github.event_name == 'schedule' ||
103+
github.event_name == 'push' ||
104+
(github.event_name == 'workflow_dispatch' &&
105+
(github.event.inputs.image_type == 'all' || github.event.inputs.image_type == 'nightly')))
106+
runs-on: ubuntu-latest
107+
strategy:
108+
matrix:
109+
config:
110+
# Main testing configuration
111+
- cuda_version: '12.4.0'
112+
cuda_short: 'cuda12.4'
113+
cu_version: 'cu124'
114+
python_version: '3.11'
115+
# CPU testing
116+
- cuda_version: '12.2.0'
117+
cuda_short: 'cuda12.2'
118+
cu_version: 'cpu'
119+
python_version: '3.9'
120+
# Old deps testing
121+
- cuda_version: '11.8.0'
122+
cuda_short: 'cuda11.8'
123+
cu_version: 'cu118'
124+
python_version: '3.9'
125+
steps:
126+
- name: Checkout repository
127+
uses: actions/checkout@v4
128+
129+
- name: Set up Docker Buildx
130+
uses: docker/setup-buildx-action@v3
131+
132+
- name: Log in to GitHub Container Registry
133+
uses: docker/login-action@v3
134+
with:
135+
registry: ${{ env.REGISTRY }}
136+
username: ${{ github.actor }}
137+
password: ${{ secrets.GITHUB_TOKEN }}
138+
139+
- name: Get current date
140+
id: date
141+
run: echo "DATE=$(date +'%Y%m%d')" >> $GITHUB_OUTPUT
142+
143+
- name: Extract metadata
144+
id: meta
145+
uses: docker/metadata-action@v5
146+
with:
147+
images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
148+
tags: |
149+
type=raw,value=nightly-${{ matrix.config.cuda_short }}-py${{ matrix.config.python_version }}-${{ steps.date.outputs.DATE }}
150+
type=raw,value=nightly-${{ matrix.config.cuda_short }}-py${{ matrix.config.python_version }}-latest
151+
152+
- name: Build and push nightly image
153+
uses: docker/build-push-action@v5
154+
with:
155+
context: .github/docker
156+
file: .github/docker/Dockerfile.nightly
157+
push: true
158+
tags: ${{ steps.meta.outputs.tags }}
159+
labels: ${{ steps.meta.outputs.labels }}
160+
build-args: |
161+
CUDA_VERSION=${{ matrix.config.cuda_version }}
162+
PYTHON_VERSION=${{ matrix.config.python_version }}
163+
CU_VERSION=${{ matrix.config.cu_version }}
164+
BASE_TAG=base-${{ matrix.config.cuda_short }}-py${{ matrix.config.python_version }}
165+
BUILD_DATE=${{ steps.date.outputs.DATE }}
166+
cache-from: type=registry,ref=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:nightly-${{ matrix.config.cuda_short }}-py${{ matrix.config.python_version }}-cache
167+
cache-to: type=registry,ref=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:nightly-${{ matrix.config.cuda_short }}-py${{ matrix.config.python_version }}-cache,mode=max
168+
169+
build-stable:
170+
needs: build-base
171+
if: |
172+
always() &&
173+
(github.event_name == 'push' ||
174+
(github.event_name == 'workflow_dispatch' &&
175+
(github.event.inputs.image_type == 'all' || github.event.inputs.image_type == 'stable')))
176+
runs-on: ubuntu-latest
177+
strategy:
178+
matrix:
179+
config:
180+
- cuda_version: '12.4.0'
181+
cuda_short: 'cuda12.4'
182+
cu_version: 'cu124'
183+
python_version: '3.10'
184+
- cuda_version: '11.8.0'
185+
cuda_short: 'cuda11.8'
186+
cu_version: 'cu118'
187+
python_version: '3.10'
188+
steps:
189+
- name: Checkout repository
190+
uses: actions/checkout@v4
191+
192+
- name: Set up Docker Buildx
193+
uses: docker/setup-buildx-action@v3
194+
195+
- name: Log in to GitHub Container Registry
196+
uses: docker/login-action@v3
197+
with:
198+
registry: ${{ env.REGISTRY }}
199+
username: ${{ github.actor }}
200+
password: ${{ secrets.GITHUB_TOKEN }}
201+
202+
- name: Get current date
203+
id: date
204+
run: echo "DATE=$(date +'%Y%m%d')" >> $GITHUB_OUTPUT
205+
206+
- name: Extract metadata
207+
id: meta
208+
uses: docker/metadata-action@v5
209+
with:
210+
images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
211+
tags: |
212+
type=raw,value=stable-${{ matrix.config.cuda_short }}-py${{ matrix.config.python_version }}
213+
214+
- name: Build and push stable image
215+
uses: docker/build-push-action@v5
216+
with:
217+
context: .github/docker
218+
file: .github/docker/Dockerfile.stable
219+
push: true
220+
tags: ${{ steps.meta.outputs.tags }}
221+
labels: ${{ steps.meta.outputs.labels }}
222+
build-args: |
223+
CUDA_VERSION=${{ matrix.config.cuda_version }}
224+
PYTHON_VERSION=${{ matrix.config.python_version }}
225+
CU_VERSION=${{ matrix.config.cu_version }}
226+
BASE_TAG=base-${{ matrix.config.cuda_short }}-py${{ matrix.config.python_version }}
227+
BUILD_DATE=${{ steps.date.outputs.DATE }}
228+
cache-from: type=registry,ref=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:stable-${{ matrix.config.cuda_short }}-py${{ matrix.config.python_version }}-cache
229+
cache-to: type=registry,ref=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:stable-${{ matrix.config.cuda_short }}-py${{ matrix.config.python_version }}-cache,mode=max
230+
231+
cleanup-old-images:
232+
needs: [build-nightly]
233+
if: github.event_name == 'schedule'
234+
runs-on: ubuntu-latest
235+
steps:
236+
- name: Delete old nightly images
237+
uses: actions/github-script@v7
238+
with:
239+
script: |
240+
const package_name = 'torchrl-ci';
241+
const keep_count = 3; // Keep last 3 nightly builds per config
242+
243+
// This is a simplified version - you might want to use a dedicated action
244+
// like snok/container-retention-policy for production use
245+
console.log('Image cleanup would run here - implement based on your retention policy');

.github/workflows/build-wheels-linux.yml

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,21 @@ on:
2929
required: false
3030
type: string
3131
default: ''
32+
with-cuda:
33+
description: 'Build with CUDA? (enable/disable)'
34+
required: false
35+
type: string
36+
default: 'enable'
37+
with-rocm:
38+
description: 'Build with ROCm? (enable/disable)'
39+
required: false
40+
type: string
41+
default: 'enable'
42+
with-cpu:
43+
description: 'Build with CPU? (enable/disable)'
44+
required: false
45+
type: string
46+
default: 'enable'
3247

3348
permissions:
3449
id-token: write
@@ -50,6 +65,9 @@ jobs:
5065
test-infra-ref: ${{ inputs.test-infra-ref || 'main' }}
5166
channel: ${{ inputs.channel || '' }}
5267
use-only-dl-pytorch-org: ${{ inputs.channel == 'release' && 'true' || 'false' }}
68+
with-cuda: ${{ inputs.with-cuda || 'enable' }}
69+
with-rocm: ${{ inputs.with-rocm || 'enable' }}
70+
with-cpu: ${{ inputs.with-cpu || 'enable' }}
5371
build:
5472
needs: generate-matrix
5573
strategy:

.github/workflows/build-wheels-windows.yml

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,16 @@ on:
2929
required: false
3030
type: string
3131
default: ''
32+
with-cuda:
33+
description: 'Build with CUDA? (enable/disable)'
34+
required: false
35+
type: string
36+
default: 'enable'
37+
with-cpu:
38+
description: 'Build with CPU? (enable/disable)'
39+
required: false
40+
type: string
41+
default: 'enable'
3242

3343
permissions:
3444
id-token: write
@@ -50,6 +60,8 @@ jobs:
5060
test-infra-ref: ${{ inputs.test-infra-ref || 'main' }}
5161
channel: ${{ inputs.channel || '' }}
5262
use-only-dl-pytorch-org: ${{ inputs.channel == 'release' && 'true' || 'false' }}
63+
with-cuda: ${{ inputs.with-cuda || 'enable' }}
64+
with-cpu: ${{ inputs.with-cpu || 'enable' }}
5365
build:
5466
needs: generate-matrix
5567
strategy:

.github/workflows/release.yml

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -261,6 +261,10 @@ jobs:
261261
test-infra-ref: main
262262
tensordict-source: ${{ inputs.tensordict_source || 'auto' }}
263263
channel: release
264+
# TorchRL is pure Python - only build CPU wheels (same wheel works with any CUDA)
265+
with-cuda: disable
266+
with-rocm: disable
267+
with-cpu: enable
264268
secrets: inherit
265269

266270
build-windows:
@@ -275,6 +279,9 @@ jobs:
275279
test-infra-ref: main
276280
tensordict-source: ${{ inputs.tensordict_source || 'auto' }}
277281
channel: release
282+
# TorchRL is pure Python - only build CPU wheels (same wheel works with any CUDA)
283+
with-cuda: disable
284+
with-cpu: enable
278285
secrets: inherit
279286

280287
build-macos:

0 commit comments

Comments
 (0)