[FAST_BUILD] Apply and merge tags in the same place (#2274)

* Apply and merge tags in the same place

* Multiple fixes

* Refactor merge_tags.py

* Small fixes

* Download aarch64 first

* Revert a change back

* Show docker images

.github/actions/apply-single-tags/action.yml

@@ -0,0 +1,49 @@
+name: Apply single platform tags
+description: Download the image tar, load it to Docker and apply tags to it
+
+inputs:
+  image:
+    description: Image name
+    required: true
+  platform:
+    description: Image platform
+    required: true
+  variant:
+    description: Variant tag prefix
+    required: true
+
+runs:
+  using: composite
+  steps:
+    - name: Load image to Docker 📥
+      uses: ./.github/actions/load-image
+      with:
+        image: ${{ inputs.image }}
+        platform: ${{ inputs.platform }}
+        variant: ${{ inputs.variant }}
+
+    - name: Download tags file 📥
+      uses: actions/download-artifact@v4
+      with:
+        name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}-tags
+        path: /tmp/jupyter/tags/
+
+    - name: Apply tags to the loaded image 🏷
+      run: |
+        python3 -m tagging.apps.apply_tags \
+          --registry ${{ env.REGISTRY }} \
+          --owner ${{ env.OWNER }} \
+          --image ${{ inputs.image }} \
+          --variant ${{ inputs.variant }} \
+          --platform ${{ inputs.platform }} \
+          --tags-dir /tmp/jupyter/tags/
+      shell: bash
+
+    # 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
+      shell: bash
+
+    - name: Show Docker images 📦
+      run: docker image ls --all
+      shell: bash

.github/actions/load-image/action.yml

@@ -20,6 +20,7 @@ runs:
       with:
         name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}
         path: /tmp/jupyter/images/
+
     - name: Load downloaded image to docker 📥
       run: |
         zstd \
@@ -28,5 +29,8 @@ runs:
           --rm \
           /tmp/jupyter/images/${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}.tar.zst \
           | docker load
-        docker image ls --all
+      shell: bash
+
+    - name: Show Docker images 📦
+      run: docker image ls --all
       shell: bash

.github/dependabot.yml

@@ -9,6 +9,10 @@ updates:
     directory: /
     schedule:
       interval: weekly
+  - package-ecosystem: github-actions
+    directory: .github/actions/apply-single-tags/
+    schedule:
+      interval: weekly
   - package-ecosystem: github-actions
     directory: .github/actions/create-dev-env/
     schedule:

.github/workflows/docker-merge-tags.yml

@@ -1,69 +0,0 @@
-name: Download all tags from GitHub artifacts and create multi-platform manifests
-
-env:
-  OWNER: ${{ github.repository_owner }}
-  PUSH_TO_REGISTRY: ${{ (github.repository_owner == 'jupyter' || github.repository_owner == 'mathbunnyru') && (github.ref == 'refs/heads/main' || github.event_name == 'schedule') }}
-
-on:
-  workflow_call:
-    inputs:
-      variant:
-        description: Variant tag prefix
-        required: true
-        type: string
-      image:
-        description: Image name
-        required: true
-        type: string
-      timeout-minutes:
-        description: Timeout in minutes
-        type: number
-        default: 10
-    secrets:
-      REGISTRY_USERNAME:
-        required: true
-      REGISTRY_TOKEN:
-        required: true
-
-jobs:
-  merge-tags:
-    runs-on: ubuntu-24.04
-    timeout-minutes: ${{ inputs.timeout-minutes }}
-
-    steps:
-      - name: Checkout Repo ⚡️
-        uses: actions/checkout@v4
-      - name: Create dev environment 📦
-        uses: ./.github/actions/create-dev-env
-
-      - name: Download x86_64 tags file 📥
-        uses: actions/download-artifact@v4
-        with:
-          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-${{ inputs.variant }}-tags
-          path: /tmp/jupyter/tags/
-        if: ${{ !contains(inputs.variant, 'cuda') }}
-
-      - name: Login to Registry 🔐
-        if: env.PUSH_TO_REGISTRY == 'true'
-        uses: docker/login-action@74a5d142397b4f367a81961eba4e8cd7edddf772 # v3.4.0
-        with:
-          registry: quay.io
-          username: ${{ secrets.REGISTRY_USERNAME }}
-          password: ${{ secrets.REGISTRY_TOKEN }}
-
-      - name: Merge tags for the images 🔀
-        run: |
-          python3 -m tagging.apps.merge_tags \
-            --image ${{ inputs.image }} \
-            --variant ${{ inputs.variant }} \
-            --tags-dir /tmp/jupyter/tags/ || \
-          python3 -m tagging.apps.merge_tags \
-            --image ${{ inputs.image }} \
-            --variant ${{ inputs.variant }} \
-            --tags-dir /tmp/jupyter/tags/
-        shell: bash

.github/workflows/docker-tag-push.yml

@@ -12,10 +12,6 @@ on:
         description: Image name
         required: true
         type: string
-      platform:
-        description: Image platform
-        required: true
-        type: string
       variant:
         description: Variant tag prefix
         required: true
@@ -40,12 +36,21 @@ jobs:
         uses: actions/checkout@v4
       - name: Create dev environment 📦
         uses: ./.github/actions/create-dev-env
-      - name: Load image to Docker 📥
-        uses: ./.github/actions/load-image
+
+      - name: Download aarch64 image tar and apply tags 🏷
+        uses: ./.github/actions/apply-single-tags
         with:
           image: ${{ inputs.image }}
-          platform: ${{ inputs.platform }}
           variant: ${{ inputs.variant }}
+          platform: aarch64
+        if: ${{ !contains(inputs.variant, 'cuda') }}
+
+      - name: Download x86_64 image tar and apply tags 🏷
+        uses: ./.github/actions/apply-single-tags
+        with:
+          image: ${{ inputs.image }}
+          variant: ${{ inputs.variant }}
+          platform: x86_64
 
       - name: Login to Registry 🔐
         if: env.PUSH_TO_REGISTRY == 'true'
@@ -55,27 +60,21 @@ jobs:
           username: ${{ secrets.REGISTRY_USERNAME }}
           password: ${{ secrets.REGISTRY_TOKEN }}
 
-      - name: Download tags file 📥
-        uses: actions/download-artifact@v4
-        with:
-          name: ${{ inputs.image }}-${{ inputs.platform }}-${{ inputs.variant }}-tags
-          path: /tmp/jupyter/tags/
-      - name: Apply tags to the loaded image 🏷
-        run: |
-          python3 -m tagging.apps.apply_tags \
-            --registry ${{ env.REGISTRY }} \
-            --owner ${{ env.OWNER }} \
-            --image ${{ inputs.image }} \
-            --variant ${{ inputs.variant }} \
-            --platform ${{ inputs.platform }} \
-            --tags-dir /tmp/jupyter/tags/
-      # 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
-
-      - name: Push Images to Registry 📤
+      - name: Push single platform images to Registry 📤
         if: env.PUSH_TO_REGISTRY == 'true'
         run: |
           docker push --all-tags ${{ env.REGISTRY }}/${{ env.OWNER }}/${{ inputs.image }} || \
           docker push --all-tags ${{ env.REGISTRY }}/${{ env.OWNER }}/${{ inputs.image }}
         shell: bash
+
+      - name: Merge tags for the images 🔀
+        run: |
+          python3 -m tagging.apps.merge_tags \
+            --image ${{ inputs.image }} \
+            --variant ${{ inputs.variant }} \
+            --tags-dir /tmp/jupyter/tags/ || \
+          python3 -m tagging.apps.merge_tags \
+            --image ${{ inputs.image }} \
+            --variant ${{ inputs.variant }} \
+            --tags-dir /tmp/jupyter/tags/
+        shell: bash

.github/workflows/docker.yml

@@ -21,6 +21,7 @@ on:
 
       # We use local composite actions to combine multiple workflow steps within one action
       # https://docs.github.com/en/actions/sharing-automations/creating-actions/about-custom-actions#composite-actions
+      - ".github/actions/apply-single-tags/action.yml"
       - ".github/actions/create-dev-env/action.yml"
       - ".github/actions/load-image/action.yml"
 
@@ -42,6 +43,7 @@ on:
       - ".github/workflows/docker-tag-push.yml"
       - ".github/workflows/docker-wiki-update.yml"
 
+      - ".github/actions/apply-single-tags/action.yml"
       - ".github/actions/create-dev-env/action.yml"
       - ".github/actions/load-image/action.yml"
 
@@ -336,68 +338,9 @@ jobs:
     needs: [x86_64-pyspark]
     if: ${{ !contains(github.event.pull_request.title, '[FAST_BUILD]') }}
 
-  aarch64-images-tag-push:
+  images-tag-push:
     uses: ./.github/workflows/docker-tag-push.yml
     with:
-      platform: aarch64
-      image: ${{ matrix.image }}
-      variant: ${{ matrix.variant }}
-    secrets:
-      REGISTRY_USERNAME: ${{ secrets.QUAY_USERNAME }}
-      REGISTRY_TOKEN: ${{ secrets.QUAY_ROBOT_TOKEN }}
-    strategy:
-      matrix:
-        image:
-          [
-            docker-stacks-foundation,
-            base-notebook,
-            minimal-notebook,
-            scipy-notebook,
-            r-notebook,
-            julia-notebook,
-            tensorflow-notebook,
-            pytorch-notebook,
-            datascience-notebook,
-            pyspark-notebook,
-            all-spark-notebook,
-          ]
-        variant: [default]
-    needs:
-      [
-        aarch64-foundation,
-        aarch64-base,
-        aarch64-minimal,
-        aarch64-scipy,
-        aarch64-r,
-        aarch64-julia,
-        aarch64-tensorflow,
-        aarch64-pytorch,
-        aarch64-datascience,
-        aarch64-pyspark,
-        aarch64-all-spark,
-      ]
-    if: ${{ !contains(github.event.pull_request.title, '[FAST_BUILD]') }}
-
-  aarch64-images-tag-push-fast:
-    uses: ./.github/workflows/docker-tag-push.yml
-    with:
-      platform: aarch64
-      image: ${{ matrix.image }}
-      variant: ${{ matrix.variant }}
-    secrets:
-      REGISTRY_USERNAME: ${{ secrets.QUAY_USERNAME }}
-      REGISTRY_TOKEN: ${{ secrets.QUAY_ROBOT_TOKEN }}
-    strategy:
-      matrix:
-        image: [docker-stacks-foundation, base-notebook]
-        variant: [default]
-    needs: [aarch64-foundation, aarch64-base]
-    if: contains(github.event.pull_request.title, '[FAST_BUILD]')
-
-  x86_64-images-tag-push:
-    uses: ./.github/workflows/docker-tag-push.yml
-    with:
-      platform: x86_64
       image: ${{ matrix.image }}
       variant: ${{ matrix.variant }}
     secrets:
@@ -429,6 +372,18 @@ jobs:
             variant: cuda12
     needs:
       [
+        aarch64-foundation,
+        aarch64-base,
+        aarch64-minimal,
+        aarch64-scipy,
+        aarch64-r,
+        aarch64-julia,
+        aarch64-tensorflow,
+        aarch64-pytorch,
+        aarch64-datascience,
+        aarch64-pyspark,
+        aarch64-all-spark,
+
         x86_64-foundation,
         x86_64-base,
         x86_64-minimal,
@@ -446,10 +401,9 @@ jobs:
       ]
     if: ${{ !contains(github.event.pull_request.title, '[FAST_BUILD]') }}
 
-  x86_64-images-tag-push-fast:
+  images-tag-push-fast:
     uses: ./.github/workflows/docker-tag-push.yml
     with:
-      platform: x86_64
       image: ${{ matrix.image }}
       variant: ${{ matrix.variant }}
     secrets:
@@ -459,69 +413,19 @@ jobs:
       matrix:
         image: [docker-stacks-foundation, base-notebook]
         variant: [default]
-    needs: [x86_64-foundation, x86_64-base]
-    if: contains(github.event.pull_request.title, '[FAST_BUILD]')
-
-  merge-tags:
-    uses: ./.github/workflows/docker-merge-tags.yml
-    with:
-      image: ${{ matrix.image }}
-      variant: ${{ matrix.variant }}
-    secrets:
-      REGISTRY_USERNAME: ${{ secrets.QUAY_USERNAME }}
-      REGISTRY_TOKEN: ${{ secrets.QUAY_ROBOT_TOKEN }}
-    strategy:
-      matrix:
-        image:
-          [
-            docker-stacks-foundation,
-            base-notebook,
-            minimal-notebook,
-            scipy-notebook,
-            r-notebook,
-            julia-notebook,
-            tensorflow-notebook,
-            pytorch-notebook,
-            datascience-notebook,
-            pyspark-notebook,
-            all-spark-notebook,
-          ]
-        variant: [default]
-        include:
-          - image: tensorflow-notebook
-            variant: cuda
-          - image: pytorch-notebook
-            variant: cuda11
-          - image: pytorch-notebook
-            variant: cuda12
-    needs: [aarch64-images-tag-push, x86_64-images-tag-push]
-    if: ${{ !contains(github.event.pull_request.title, '[FAST_BUILD]') }}
-
-  merge-tags-fast:
-    uses: ./.github/workflows/docker-merge-tags.yml
-    with:
-      image: ${{ matrix.image }}
-      variant: ${{ matrix.variant }}
-    secrets:
-      REGISTRY_USERNAME: ${{ secrets.QUAY_USERNAME }}
-      REGISTRY_TOKEN: ${{ secrets.QUAY_ROBOT_TOKEN }}
-    strategy:
-      matrix:
-        image: [docker-stacks-foundation, base-notebook]
-        variant: [default]
-    needs: [aarch64-images-tag-push-fast, x86_64-images-tag-push-fast]
+    needs: [aarch64-foundation, aarch64-base, x86_64-foundation, x86_64-base]
     if: contains(github.event.pull_request.title, '[FAST_BUILD]')
 
   wiki-update:
     uses: ./.github/workflows/docker-wiki-update.yml
-    needs: [aarch64-images-tag-push, x86_64-images-tag-push]
+    needs: [images-tag-push]
     if: ${{ !contains(github.event.pull_request.title, '[FAST_BUILD]') }}
     permissions:
       contents: write
 
   wiki-update-fast:
     uses: ./.github/workflows/docker-wiki-update.yml
-    needs: [aarch64-images-tag-push-fast, x86_64-images-tag-push-fast]
+    needs: [images-tag-push-fast]
     if: contains(github.event.pull_request.title, '[FAST_BUILD]')
 
   contributed-recipes:

tagging/apps/merge_tags.py

@@ -16,59 +16,79 @@ docker = plumbum.local["docker"]
 LOGGER = logging.getLogger(__name__)
 
 
-def read_tags_from_files(config: Config) -> set[str]:
+def read_local_tags_from_files(config: Config) -> tuple[list[str], set[str]]:
     LOGGER.info(f"Read tags from file(s) for image: {config.image}")
 
-    tags: set[str] = set()
+    all_local_tags = []
+    merged_local_tags = set()
     for platform in ALL_PLATFORMS:
         LOGGER.info(f"Reading tags for platform: {platform}")
 
         file_prefix = get_file_prefix_for_platform(platform, config.variant)
         filename = f"{file_prefix}-{config.image}.txt"
         path = config.tags_dir / filename
-        if path.exists():
-            LOGGER.info(f"Tag file: {path} found")
-            lines = path.read_text().splitlines()
-            tags.update(tag.replace(platform + "-", "") for tag in lines)
-        else:
+        if not path.exists():
             LOGGER.info(f"Tag file: {path} doesn't exist")
+            continue
+
+        LOGGER.info(f"Tag file: {path} found")
+        for tag in path.read_text().splitlines():
+            all_local_tags.append(tag)
+            merged_local_tags.add(tag.replace(platform + "-", ""))
 
     LOGGER.info(f"Tags read for image: {config.image}")
-    return tags
+    return all_local_tags, merged_local_tags
 
 
-def merge_tags(tag: str, push_to_registry: bool) -> None:
-    LOGGER.info(f"Trying to merge tag: {tag}")
-    all_platform_tags = []
+def pull_missing_tags(merged_tag: str, all_local_tags: list[str]) -> list[str]:
+    existing_platform_tags = []
+
     for platform in ALL_PLATFORMS:
-        platform_tag = tag.replace(":", f":{platform}-")
-        LOGGER.info(f"Trying to pull: {platform_tag}")
+        platform_tag = merged_tag.replace(":", f":{platform}-")
+        if platform_tag in all_local_tags:
+            LOGGER.info(
+                f"Tag {platform_tag} already exists locally, not pulling it from registry"
+            )
+            existing_platform_tags.append(platform_tag)
+            continue
+
+        LOGGER.warning(f"Trying to pull: {platform_tag} from registry")
         try:
             docker["pull", platform_tag] & plumbum.FG
-            all_platform_tags.append(platform_tag)
-            LOGGER.info("Pull success")
+            existing_platform_tags.append(platform_tag)
+            LOGGER.info(f"Tag {platform_tag} pulled successfully")
         except plumbum.ProcessExecutionError:
-            LOGGER.info("Pull failed, image with this tag and platform doesn't exist")
+            LOGGER.warning(f"Pull failed, tag {platform_tag} doesn't exist")
 
-    LOGGER.info(f"Found images: {all_platform_tags}")
+    return existing_platform_tags
+
+
+def merge_tags(
+    merged_tag: str, all_local_tags: list[str], push_to_registry: bool
+) -> None:
+    LOGGER.info(f"Trying to merge tag: {merged_tag}")
+
+    existing_platform_tags = pull_missing_tags(merged_tag, all_local_tags)
+
+    # This allows to rerun the script without having to remove the manifest manually
     try:
-        docker["manifest", "rm", tag] & plumbum.FG
-        LOGGER.info(f"Manifest {tag} already exists, removing it")
+        docker["manifest", "rm", merged_tag] & plumbum.FG
+        LOGGER.warning(f"Manifest {merged_tag} was present locally, removed it")
     except plumbum.ProcessExecutionError:
-        LOGGER.info(f"Manifest {tag} doesn't exist")
+        pass
 
     if push_to_registry:
-        # We need images to have been already pushed to the registry
-        # before creating the manifest
-        LOGGER.info(f"Creating manifest for tag: {tag}")
-        docker["manifest", "create", tag][all_platform_tags] & plumbum.FG
-        LOGGER.info(f"Successfully created manifest for tag: {tag}")
+        # Unforunately, `docker manifest create` requires images to have been already pushed to the registry
+        # which is not true for new tags in PRs
+        LOGGER.info(f"Creating manifest for tag: {merged_tag}")
+        docker["manifest", "create", merged_tag][existing_platform_tags] & plumbum.FG
+        LOGGER.info(f"Successfully created manifest for tag: {merged_tag}")
 
-        LOGGER.info(f"Pushing manifest for tag: {tag}")
-        docker["manifest", "push", tag] & plumbum.FG
-        LOGGER.info(f"Successfully merged and pushed tag: {tag}")
+        LOGGER.info(f"Pushing manifest for tag: {merged_tag}")
+        docker["manifest", "push", merged_tag] & plumbum.FG
+        LOGGER.info(f"Successfully merged and pushed tag: {merged_tag}")
     else:
-        LOGGER.info(f"Skipping push for tag: {tag}")
+        LOGGER.info(f"Skipping push for tag: {merged_tag}")
 
 
 if __name__ == "__main__":
@@ -78,7 +98,9 @@ if __name__ == "__main__":
     push_to_registry = os.environ.get("PUSH_TO_REGISTRY", "false").lower() == "true"
 
     LOGGER.info(f"Merging tags for image: {config.image}")
-    all_tags = read_tags_from_files(config)
-    for tag in all_tags:
-        merge_tags(tag, push_to_registry)
+
+    all_local_tags, merged_local_tags = read_local_tags_from_files(config)
+    for tag in merged_local_tags:
+        merge_tags(tag, all_local_tags, push_to_registry)
+
     LOGGER.info(f"Successfully merged tags for image: {config.image}")