docker: fix setting tags

Author: Borda

.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,
+    ) -> str:
+        """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]
+        return ",".join(tags)
+
 
 if __name__ == "__main__":
     import sys

.github/workflows/docker-build.yml

@@ -47,7 +47,7 @@ jobs:
           - { 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.12", pytorch_version: "2.5", cuda_version: "12.1.0" , latest: true}
     steps:
       - uses: actions/checkout@v4
         with:
@@ -65,23 +65,13 @@ 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
+          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 }}")
+          echo "DOCKER_TAGS=$tags" >> $GITHUB_ENV
 
       - uses: docker/build-push-action@v6
         with: