docker: fix setting tags (#20504)

* docker: fix setting tags

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* pip install

* cmd

* latest

* print

* cuda 12.1.1

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.actions/assistant.py

@@ -481,6 +481,25 @@ def convert_version2nightly(ver_file: str = "src/version.info") -> None:
         with open(ver_file, "w") as fo:
             fo.write(version + os.linesep)
 
+    @staticmethod
+    def generate_docker_tags(
+        release_version: str,
+        python_version: str,
+        torch_version: str,
+        cuda_version: str,
+        docker_project: str = "pytorchlightning/pytorch_lightning",
+        add_latest: bool = False,
+    ) -> None:
+        """Generate docker tags for the given versions."""
+        tags = [f"latest-py{python_version}-torch{torch_version}-cuda{cuda_version}"]
+        if release_version:
+            tags += [f"{release_version}-py{python_version}-torch{torch_version}-cuda{cuda_version}"]
+        if add_latest:
+            tags += ["latest"]
+
+        tags = [f"{docker_project}:{tag}" for tag in tags]
+        print(",".join(tags))
+
 
 if __name__ == "__main__":
     import sys

.github/checkgroup.yml

@@ -145,17 +145,17 @@ subprojects:
       - "!*.md"
       - "!**/*.md"
     checks:
-      - "build-cuda (3.10, 2.1.2, 12.1.0)"
-      - "build-cuda (3.11, 2.2.2, 12.1.0)"
-      - "build-cuda (3.11, 2.3.1, 12.1.0)"
-      - "build-cuda (3.11, 2.4.1, 12.1.0)"
-      - "build-cuda (3.12, 2.5.1, 12.1.0)"
+      - "build-cuda (3.10, 2.1.2, 12.1.1)"
+      - "build-cuda (3.11, 2.2.2, 12.1.1)"
+      - "build-cuda (3.11, 2.3.1, 12.1.1)"
+      - "build-cuda (3.11, 2.4.1, 12.1.1)"
+      - "build-cuda (3.12, 2.5.1, 12.1.1)"
       #- "build-NGC"
-      - "build-pl (3.10, 2.1, 12.1.0)"
-      - "build-pl (3.11, 2.2, 12.1.0)"
-      - "build-pl (3.11, 2.3, 12.1.0)"
-      - "build-pl (3.11, 2.4, 12.1.0)"
-      - "build-pl (3.12, 2.5, 12.1.0)"
+      - "build-pl (3.10, 2.1, 12.1.1)"
+      - "build-pl (3.11, 2.2, 12.1.1)"
+      - "build-pl (3.11, 2.3, 12.1.1)"
+      - "build-pl (3.11, 2.4, 12.1.1)"
+      - "build-pl (3.12, 2.5, 12.1.1, true)"
 
   # SECTION: lightning_fabric
 

.github/workflows/docker-build.yml

@@ -43,11 +43,11 @@ jobs:
         include:
           # We only release one docker image per PyTorch version.
           # Make sure the matrix here matches the one below.
-          - { python_version: "3.10", pytorch_version: "2.1", cuda_version: "12.1.0" }
-          - { python_version: "3.11", pytorch_version: "2.2", cuda_version: "12.1.0" }
-          - { python_version: "3.11", pytorch_version: "2.3", cuda_version: "12.1.0" }
-          - { python_version: "3.11", pytorch_version: "2.4", cuda_version: "12.1.0" }
-          - { python_version: "3.12", pytorch_version: "2.5", cuda_version: "12.1.0" }
+          - { python_version: "3.10", pytorch_version: "2.1", cuda_version: "12.1.1" }
+          - { python_version: "3.11", pytorch_version: "2.2", cuda_version: "12.1.1" }
+          - { python_version: "3.11", pytorch_version: "2.3", cuda_version: "12.1.1" }
+          - { python_version: "3.11", pytorch_version: "2.4", cuda_version: "12.1.1" }
+          - { python_version: "3.12", pytorch_version: "2.5", cuda_version: "12.1.1", latest: "true" }
     steps:
       - uses: actions/checkout@v4
         with:
@@ -65,23 +65,14 @@ jobs:
         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.11' and pt_ver == '2.3' and cuda_ver == '12.1.0':
-            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
+          pip install -q -r .actions/requirements.txt
+          tags=$(python .actions/assistant.py generate_docker_tags \
+            --release_version="${{ env.RELEASE_VERSION }}" \
+            --python_version="${{ matrix.python_version }}" \
+            --torch_version="${{ matrix.pytorch_version }}" \
+            --cuda_version="${{ matrix.cuda_version }}" \
+            --add_latest="${{ matrix.latest || 'false' }}")
+          echo "DOCKER_TAGS=$tags" >> $GITHUB_ENV
 
       - uses: docker/build-push-action@v6
         with:
@@ -104,11 +95,11 @@ jobs:
         include:
           # These are the base images for PL release docker images.
           # Make sure the matrix here matches the one above.
-          - { python_version: "3.10", pytorch_version: "2.1.2", cuda_version: "12.1.0" }
-          - { python_version: "3.11", pytorch_version: "2.2.2", cuda_version: "12.1.0" }
-          - { python_version: "3.11", pytorch_version: "2.3.1", cuda_version: "12.1.0" }
-          - { python_version: "3.11", pytorch_version: "2.4.1", cuda_version: "12.1.0" }
-          - { python_version: "3.12", pytorch_version: "2.5.1", cuda_version: "12.1.0" }
+          - { python_version: "3.10", pytorch_version: "2.1.2", cuda_version: "12.1.1" }
+          - { python_version: "3.11", pytorch_version: "2.2.2", cuda_version: "12.1.1" }
+          - { python_version: "3.11", pytorch_version: "2.3.1", cuda_version: "12.1.1" }
+          - { python_version: "3.11", pytorch_version: "2.4.1", cuda_version: "12.1.1" }
+          - { python_version: "3.12", pytorch_version: "2.5.1", cuda_version: "12.1.1" }
     steps:
       - uses: actions/checkout@v4
       - uses: docker/setup-buildx-action@v3