From 3a70714219e7d0cae243bfd80cacb9de26e55872 Mon Sep 17 00:00:00 2001
From: Jirka Borovec <6035284+Borda@users.noreply.github.com>
Date: Wed, 30 Aug 2023 19:11:35 +0200
Subject: [PATCH] ci: fix & fuse docker builds (#18440)

---
 .github/checkgroup.yml                        | 12 +---
 .../{ci-dockers.yml => docker-build.yml}      | 42 +++++++++--
 .github/workflows/docs-build.yml              |  2 +-
 .github/workflows/release-docker.yml          | 72 -------------------
 4 files changed, 39 insertions(+), 89 deletions(-)
 rename .github/workflows/{ci-dockers.yml => docker-build.yml} (73%)
 delete mode 100644 .github/workflows/release-docker.yml

diff --git a/.github/checkgroup.yml b/.github/checkgroup.yml
index bc06ee126e..48204db5ff 100644
--- a/.github/checkgroup.yml
+++ b/.github/checkgroup.yml
@@ -134,7 +134,7 @@ subprojects:
   - id: "pytorch_lightning: Docker"
     paths:
       - ".actions/*"
-      - ".github/workflows/ci-dockers.yml"
+      - ".github/workflows/docker-build.yml"
       - "dockers/**"
       - "requirements/pytorch/**"
       - "requirements/fabric/**"
@@ -425,13 +425,3 @@ subprojects:
       - "install-pkg (windows-2022, lightning, 3.11)"
       - "install-pkg (windows-2022, notset, 3.8)"
       - "install-pkg (windows-2022, notset, 3.11)"
-
-  - id: "release dockers"
-    paths:
-      - ".github/workflows/release-docker.yml"
-      - "dockers/release/*"
-    checks:
-      - "publish-docker (3.9, 1.13, 12.0.1)"
-      - "publish-docker (3.9, 1.13, 12.0.1)"
-      - "publish-docker (3.10, 2.0, 11.8.0)"
-      - "publish-docker (3.10, 2.0, 12.0.1)"
diff --git a/.github/workflows/ci-dockers.yml b/.github/workflows/docker-build.yml
similarity index 73%
rename from .github/workflows/ci-dockers.yml
rename to .github/workflows/docker-build.yml
index c6ecb5491f..7dc8704251 100644
--- a/.github/workflows/ci-dockers.yml
+++ b/.github/workflows/docker-build.yml
@@ -1,4 +1,4 @@
-name: Docker
+name: Docker builds
 
 on:
   push:
@@ -8,7 +8,7 @@ on:
     types: [opened, reopened, ready_for_review, synchronize]  # added `ready_for_review` since draft is skipped
     paths:
       - ".actions/*"
-      - ".github/workflows/ci-dockers.yml"
+      - ".github/workflows/docker-build.yml"
       - "dockers/**"
       - "requirements/*.txt"
       - "requirements/pytorch/**"
@@ -28,7 +28,7 @@ concurrency:
 
 env:
   PUSH_NIGHTLY: ${{ github.event_name == 'schedule' }}
-  PUSH_RELEASE: ${{ github.event_name == 'release' }}
+  PUSH_RELEASE: ${{ startsWith(github.ref, 'refs/tags/') || github.event_name == 'release' }}
 
 jobs:
   build-pl:
@@ -52,15 +52,47 @@ jobs:
         with:
           submodules: true
       - uses: docker/setup-buildx-action@v2
+      - uses: docker/login-action@v2
+        if: env.PUSH_RELEASE == 'true'
+        with:
+          username: ${{ secrets.DOCKER_USERNAME }}
+          password: ${{ secrets.DOCKER_PASSWORD }}
+
+      - name: Get release version
+        if: github.event_name == 'release'
+        # For workflows triggered by release, `GITHUB_REF` is the release tag created.
+        run: echo "RELEASE_VERSION=$(echo ${GITHUB_REF##*/})" >> $GITHUB_ENV
+      - name: Set tags
+        run: |
+          import os
+
+          repo = "pytorchlightning/pytorch_lightning"
+          ver = os.getenv('RELEASE_VERSION')
+          py_ver = "${{ matrix.python_version }}"
+          pt_ver = "${{ matrix.pytorch_version }}"
+          cuda_ver = "${{ matrix.cuda_version }}"
+          tags = [f"latest-py{py_ver}-torch{pt_ver}-cuda{cuda_ver}"]
+          if ver:
+            tags += [f"{ver}-py{py_ver}-torch{pt_ver}-cuda{cuda_ver}"]
+          if py_ver == '3.10' and pt_ver == '2.0' and cuda_ver == '12.0.1':
+            tags += ["latest"]
+
+          tags = [f"{repo}:{tag}" for tag in tags]
+          with open(os.getenv('GITHUB_ENV'), "a") as gh_env:
+              gh_env.write("DOCKER_TAGS=" + ",".join(tags))
+        shell: python
+
       - uses: docker/build-push-action@v4
         with:
           build-args: |
             PYTHON_VERSION=${{ matrix.python_version }}
             PYTORCH_VERSION=${{ matrix.pytorch_version }}
             CUDA_VERSION=${{ matrix.cuda_version }}
+            LIGHTNING_VERSION=${{ env.RELEASE_VERSION }}
           file: dockers/release/Dockerfile
           push: ${{ env.PUSH_RELEASE }}  # pushed in release-docker.yml only when PL is released
-        timeout-minutes: 50
+          tags: ${{ env.DOCKER_TAGS }}
+        timeout-minutes: 35
 
   build-cuda:
     if: github.event.pull_request.draft == false
@@ -92,7 +124,7 @@ jobs:
             CUDA_VERSION=${{ matrix.cuda_version }}
           file: dockers/base-cuda/Dockerfile
           push: ${{ env.PUSH_NIGHTLY }}
-          tags: pytorchlightning/pytorch_lightning:base-cuda-py${{ matrix.python_version }}-torch${{ matrix.pytorch_version }}-cuda${{ matrix.cuda_version }}
+          tags: "pytorchlightning/pytorch_lightning:base-cuda-py${{ matrix.python_version }}-torch${{ matrix.pytorch_version }}-cuda${{ matrix.cuda_version }}"
         timeout-minutes: 95
       - uses: ravsamhq/notify-slack-action@v2
         if: failure() && env.PUSH_NIGHTLY == 'true'
diff --git a/.github/workflows/docs-build.yml b/.github/workflows/docs-build.yml
index 5bd3814e91..fd106c9bc7 100644
--- a/.github/workflows/docs-build.yml
+++ b/.github/workflows/docs-build.yml
@@ -1,4 +1,4 @@
-name: Build Docs
+name: Docs builds
 # https://github.com/marketplace/actions/sphinx-build
 
 on:
diff --git a/.github/workflows/release-docker.yml b/.github/workflows/release-docker.yml
deleted file mode 100644
index 57145627d1..0000000000
--- a/.github/workflows/release-docker.yml
+++ /dev/null
@@ -1,72 +0,0 @@
-name: Docker
-
-on:
-  push:
-    branches: [master, "release/*"]
-  pull_request:
-    branches: [master, "release/*"]
-    paths:
-      - ".github/workflows/release-docker.yml"
-      - "dockers/release/*"
-  release:
-    types: [published]
-
-jobs:
-  publish-docker:
-    runs-on: ubuntu-latest
-    strategy:
-      fail-fast: false
-      matrix:
-        include:
-          # We only release one docker image per PyTorch version.
-          - {python_version: "3.9", pytorch_version: "1.13", cuda_version: "11.8.0"}
-          - {python_version: "3.9", pytorch_version: "1.13", cuda_version: "12.0.1"}
-          - {python_version: "3.10", pytorch_version: "2.0", cuda_version: "11.8.0"}
-          - {python_version: "3.10", pytorch_version: "2.0", cuda_version: "12.0.1"}
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v3
-        with:
-          submodules: true
-
-      - name: Get release version
-        id: get_version
-        run: echo "RELEASE_VERSION=$(echo ${GITHUB_REF##*/})" >> $GITHUB_OUTPUT
-
-      - name: Publish Releases to Docker
-        uses: docker/build-push-action@v4
-        with:
-          repository: pytorchlightning/pytorch_lightning
-          username: ${{ secrets.DOCKER_USERNAME }}
-          password: ${{ secrets.DOCKER_PASSWORD }}
-          # only on releases
-          push: ${{ startsWith(github.ref, 'refs/tags/') || github.event_name == 'release' }}
-          file: dockers/release/Dockerfile
-          build_args: |
-            PYTHON_VERSION=${{ matrix.python_version }}
-            PYTORCH_VERSION=${{ matrix.pytorch_version }}
-            CUDA_VERSION=${{ matrix.cuda_version }}
-            LIGHTNING_VERSION=${{ steps.get_version.outputs.RELEASE_VERSION }}
-          tags: |
-            ${{ steps.get_version.outputs.RELEASE_VERSION }}-py${{ matrix.python_version }}-torch${{ matrix.pytorch_version }}-cuda${{ matrix.cuda_version }}
-            latest-py${{ matrix.python_version }}-torch${{ matrix.pytorch_version }}-cuda${{ matrix.cuda_version }}
-        timeout-minutes: 55
-
-      - name: Publish Latest to Docker
-        uses: docker/build-push-action@v4
-        # Only latest Python and PyTorch
-        if: matrix.python_version == '3.10' && matrix.pytorch_version == '2.0'
-        with:
-          repository: pytorchlightning/pytorch_lightning
-          username: ${{ secrets.DOCKER_USERNAME }}
-          password: ${{ secrets.DOCKER_PASSWORD }}
-          # only on releases
-          push: ${{ startsWith(github.ref, 'refs/tags/') || github.event_name == 'release' }}
-          file: dockers/release/Dockerfile
-          build_args: |
-            PYTHON_VERSION=${{ matrix.python_version }}
-            PYTORCH_VERSION=${{ matrix.pytorch_version }}
-            CUDA_VERSION=${{ matrix.cuda_version }}
-            LIGHTNING_VERSION=${{ steps.get_version.outputs.RELEASE_VERSION }}
-          tags: "latest"
-        timeout-minutes: 55