Skip to content

Commit

Permalink
Add pytorch-notebook image variants with cuda 11 and 12 (x86_64 versi…
Browse files Browse the repository at this point in the history
…ons only) (#2091)

* feat: build cuda variants of pytorch

* feat: build with variant tag

* style: remove unused import

* refactor: rename get_prefix params

(cherry picked from commit 12b50af)

* revert: drop ROOT_CONTAINER addition from Makefile

(cherry picked from commit f423145)

* style: use consistent three empty lines in Makefile

(cherry picked from commit 446b45a)

* refactor: add default value for parent-image

(cherry picked from commit 32955ce)

* revert: use original workflow structure

(cherry picked from commit 68c6744)

* refactor: use single build image step

(cherry picked from commit 5f1ac0a)

* fix: run merge tags regardless of repository owner

(cherry picked from commit 3fce366)

* refactor: build cuda12 instead of cuda tag

(cherry picked from commit 217144e)

* docs: add note about CUDA tags to documentation

* refactor: add default value for variant in build-test-upload

* refactor: swap ordering of cuda11/cuda12 variants

* refactor: remove optional str type in arg parser

* fix: add proper env variables to CUDA Dockerfiles

* fix: remove CUDA build for aarch64

* fix: use latest NVIDIA documentation link

* fix: skip aarch64 tags file for CUDA variants

---------

Co-authored-by: zynaa <7562909-zynaa@users.noreply.gitlab.com>
  • Loading branch information
johanna-reiml-hpi and zynaa authored Feb 24, 2024
1 parent fe21954 commit eccda24
Show file tree
Hide file tree
Showing 15 changed files with 307 additions and 84 deletions.
7 changes: 5 additions & 2 deletions .github/actions/load-image/action.yml
Original file line number Diff line number Diff line change
Expand Up @@ -8,17 +8,20 @@ inputs:
platform:
description: Image platform
required: true
variant:
description: Variant tag prefix
required: true

runs:
using: composite
steps:
- name: Download built image 📥
uses: actions/download-artifact@v4
with:
name: ${{ inputs.image }}-${{ inputs.platform }}
name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}
path: /tmp/jupyter/images/
- name: Load downloaded image to docker 📥
run: |
zstd --uncompress --stdout --rm /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}.tar.zst | docker load
zstd --uncompress --stdout --rm /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}.tar.zst | docker load
docker image ls --all
shell: bash
48 changes: 36 additions & 12 deletions .github/workflows/docker-build-test-upload.yml
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,20 @@ on:
description: Parent image name
required: true
type: string
parent-variant:
description: Parent variant tag prefix
required: false
type: string
default: default
image:
description: Image name
required: true
type: string
variant:
description: Variant tag prefix
required: false
type: string
default: default
platform:
description: Image platform
required: true
Expand All @@ -29,6 +39,19 @@ jobs:
runs-on: ${{ inputs.runs-on }}

steps:
# Image with CUDA needs extra disk space
- name: Free disk space 🧹
if: contains(inputs.variant, 'cuda') && inputs.platform == 'x86_64'
uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be
with:
tool-cache: false
android: true
dotnet: true
haskell: true
large-packages: false
docker-images: false
swap-storage: false

- name: Checkout Repo ⚡️
uses: actions/checkout@v4
- name: Create dev environment 📦
Expand All @@ -52,14 +75,15 @@ jobs:
with:
image: ${{ inputs.parent-image }}
platform: ${{ inputs.platform }}
variant: ${{ inputs.parent-variant }}

- name: Pull ubuntu:22.04 image 📥
if: inputs.parent-image == ''
run: docker pull ubuntu:22.04
shell: bash

- name: Build image 🛠
run: docker build --rm --force-rm --tag ${{ env.REGISTRY }}/${{ env.OWNER }}/${{ inputs.image }} images/${{ inputs.image }}/ --build-arg REGISTRY=${{ env.REGISTRY }} --build-arg OWNER=${{ env.OWNER }}
run: docker build --rm --force-rm --tag ${{ env.REGISTRY }}/${{ env.OWNER }}/${{ inputs.image }} images/${{ inputs.image }}/${{ inputs.variant != 'default' && inputs.variant || '.' }}/ --build-arg REGISTRY=${{ env.REGISTRY }} --build-arg OWNER=${{ env.OWNER }}
env:
DOCKER_BUILDKIT: 1
# Full logs for CI build
Expand All @@ -72,39 +96,39 @@ jobs:

- name: Write tags file 🏷
run: |
python3 -m tagging.write_tags_file --short-image-name ${{ inputs.image }} --tags-dir /tmp/jupyter/tags/ --registry ${{ env.REGISTRY }} --owner ${{ env.OWNER }}
python3 -m tagging.write_tags_file --short-image-name ${{ inputs.image }} --tags-dir /tmp/jupyter/tags/ --registry ${{ env.REGISTRY }} --owner ${{ env.OWNER }} --variant ${{ inputs.variant }}
shell: bash
- name: Upload tags file 💾
uses: actions/upload-artifact@v4
with:
name: ${{ inputs.image }}-${{ inputs.platform }}-tags
path: /tmp/jupyter/tags/${{ inputs.platform }}-${{ inputs.image }}.txt
name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}-tags
path: /tmp/jupyter/tags/${{ inputs.platform }}-${{ inputs.variant }}-${{ inputs.image }}.txt
retention-days: 3

- name: Write manifest and build history file 🏷
run: python3 -m tagging.write_manifest --short-image-name ${{ inputs.image }} --hist-lines-dir /tmp/jupyter/hist_lines/ --manifests-dir /tmp/jupyter/manifests/ --registry ${{ env.REGISTRY }} --owner ${{ env.OWNER }}
run: python3 -m tagging.write_manifest --short-image-name ${{ inputs.image }} --hist-lines-dir /tmp/jupyter/hist_lines/ --manifests-dir /tmp/jupyter/manifests/ --registry ${{ env.REGISTRY }} --owner ${{ env.OWNER }} --variant ${{ inputs.variant }}
shell: bash
- name: Upload manifest file 💾
uses: actions/upload-artifact@v4
with:
name: ${{ inputs.image }}-${{ inputs.platform }}-manifest
path: /tmp/jupyter/manifests/${{ inputs.platform }}-${{ inputs.image }}-*.md
name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}-manifest
path: /tmp/jupyter/manifests/${{ inputs.platform }}-${{ inputs.variant }}-${{ inputs.image }}-*.md
retention-days: 3
- name: Upload build history line 💾
uses: actions/upload-artifact@v4
with:
name: ${{ inputs.image }}-${{ inputs.platform }}-history_line
path: /tmp/jupyter/hist_lines/${{ inputs.platform }}-${{ inputs.image }}-*.txt
name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}-history_line
path: /tmp/jupyter/hist_lines/${{ inputs.platform }}-${{ inputs.variant }}-${{ inputs.image }}-*.txt
retention-days: 3

- name: Save image as a tar for later use 💾
run: |
mkdir -p /tmp/jupyter/images/
docker save ${{ env.REGISTRY }}/${{ env.OWNER }}/${{ inputs.image }} | zstd > /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}.tar.zst
docker save ${{ env.REGISTRY }}/${{ env.OWNER }}/${{ inputs.image }} | zstd > /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}.tar.zst
shell: bash
- name: Upload image as artifact 💾
uses: actions/upload-artifact@v4
with:
name: ${{ inputs.image }}-${{ inputs.platform }}
path: /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}.tar.zst
name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}
path: /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}.tar.zst
retention-days: 3
11 changes: 8 additions & 3 deletions .github/workflows/docker-merge-tags.yml
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,10 @@ env:
on:
workflow_call:
inputs:
variant:
description: Variant tag prefix
required: true
type: string
image:
description: Image name
required: true
Expand All @@ -30,13 +34,14 @@ jobs:
- name: Download x86_64 tags file 📥
uses: actions/download-artifact@v4
with:
name: ${{ inputs.image }}-x86_64-tags
name: ${{ inputs.image }}-x86_64-${{ inputs.variant }}-tags
path: /tmp/jupyter/tags/
- name: Download aarch64 tags file 📥
uses: actions/download-artifact@v4
with:
name: ${{ inputs.image }}-aarch64-tags
name: ${{ inputs.image }}-aarch64-${{ inputs.variant }}-tags
path: /tmp/jupyter/tags/
if: github.repository_owner == 'jupyter' && !contains(inputs.variant, 'cuda')

# Docker might be stuck when pulling images
# https://github.com/docker/for-mac/issues/2083
Expand All @@ -57,5 +62,5 @@ jobs:

- name: Merge tags for the images 🔀
if: env.PUSH_TO_REGISTRY == 'true'
run: python3 -m tagging.merge_tags --short-image-name ${{ inputs.image }} --tags-dir /tmp/jupyter/tags/
run: python3 -m tagging.merge_tags --short-image-name ${{ inputs.image }} --tags-dir /tmp/jupyter/tags/ --variant ${{ inputs.variant }}
shell: bash
22 changes: 20 additions & 2 deletions .github/workflows/docker-tag-push.yml
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,10 @@ on:
description: Image platform
required: true
type: string
variant:
description: Variant tag prefix
required: true
type: string
secrets:
REGISTRY_USERNAME:
required: true
Expand All @@ -27,6 +31,19 @@ jobs:
runs-on: ubuntu-latest

steps:
# Image with CUDA needs extra disk space
- name: Free disk space 🧹
if: contains(inputs.variant, 'cuda') && inputs.platform == 'x86_64'
uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be
with:
tool-cache: false
android: true
dotnet: true
haskell: true
large-packages: false
docker-images: false
swap-storage: false

- name: Checkout Repo ⚡️
uses: actions/checkout@v4
- name: Create dev environment 📦
Expand All @@ -36,6 +53,7 @@ jobs:
with:
image: ${{ inputs.image }}
platform: ${{ inputs.platform }}
variant: ${{ inputs.variant }}

- name: Login to Registry 🔐
if: env.PUSH_TO_REGISTRY == 'true'
Expand All @@ -48,10 +66,10 @@ jobs:
- name: Download tags file 📥
uses: actions/download-artifact@v4
with:
name: ${{ inputs.image }}-${{ inputs.platform }}-tags
name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}-tags
path: /tmp/jupyter/tags/
- name: Apply tags to the loaded image 🏷
run: python3 -m tagging.apply_tags --short-image-name ${{ inputs.image }} --tags-dir /tmp/jupyter/tags/ --platform ${{ inputs.platform }} --registry ${{ env.REGISTRY }} --owner ${{ env.OWNER }}
run: python3 -m tagging.apply_tags --short-image-name ${{ inputs.image }} --tags-dir /tmp/jupyter/tags/ --platform ${{ inputs.platform }} --variant ${{ inputs.variant }} --registry ${{ env.REGISTRY }} --owner ${{ env.OWNER }}
# This step is needed to prevent pushing non-multi-arch "latest" tag
- name: Remove the "latest" tag from the image 🗑️
run: docker image rmi ${{ env.REGISTRY }}/${{ env.OWNER }}/${{ inputs.image }}:latest
Expand Down
Loading

0 comments on commit eccda24

Please sign in to comment.