[FAST_BUILD] Merge manifests without pulling images (#2377)

Author: mathbunnyru

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

@@ -0,0 +1,80 @@
+name: Merge single platform tags
+
+env:
+  REGISTRY: quay.io
+  PUSH_TO_REGISTRY: ${{ (github.repository_owner == 'jupyter' || github.repository_owner == 'mathbunnyru') && (github.ref == 'refs/heads/main' || github.event_name == 'schedule') }}
+
+on:
+  workflow_call:
+    inputs:
+      image:
+        description: Image name
+        required: true
+        type: string
+      variant:
+        description: Variant tag prefix
+        required: true
+        type: string
+      timeout-minutes:
+        description: Timeout in minutes
+        default: 5
+        type: number
+    secrets:
+      REGISTRY_USERNAME:
+        required: true
+      REGISTRY_TOKEN:
+        required: true
+
+permissions:
+  contents: read
+
+jobs:
+  tag-merge:
+    runs-on: ubuntu-24.04
+    timeout-minutes: ${{ inputs.timeout-minutes }}
+
+    steps:
+      - name: Checkout Repo ⚡️
+        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+
+      - name: Create dev environment 📦
+        uses: ./.github/actions/create-dev-env
+
+      - name: Download aarch64 tags file 🏷
+        uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 # v6.0.0
+        with:
+          name: ${{ inputs.image }}-aarch64-${{ inputs.variant }}-tags
+          path: /tmp/jupyter/tags/
+        if: ${{ !contains(inputs.variant, 'cuda11') }}
+
+      - name: Download x86_64 tags file 🏷
+        uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53 # v6.0.0
+        with:
+          name: ${{ inputs.image }}-x86_64-${{ inputs.variant }}-tags
+          path: /tmp/jupyter/tags/
+
+      - name: Login to Registry 🔐
+        if: env.PUSH_TO_REGISTRY == 'true'
+        run: |
+          docker login ${{ env.REGISTRY }} \
+            --username ${{ secrets.REGISTRY_USERNAME }} \
+            --password ${{ secrets.REGISTRY_TOKEN }} || \
+          docker login ${{ env.REGISTRY }} \
+            --username ${{ secrets.REGISTRY_USERNAME }} \
+            --password ${{ secrets.REGISTRY_TOKEN }}
+        shell: bash
+        id: login
+
+      - name: Merge tags for the images 🔀
+        run: |
+          python3 -m tagging.apps.merge_tags \
+            --image ${{ inputs.image }} \
+            --variant ${{ inputs.variant }} \
+            --tags-dir /tmp/jupyter/tags/
+        shell: bash
+
+      - name: Logout from Registry 🔐
+        if: always() && env.PUSH_TO_REGISTRY == 'true' && steps.login.outcome == 'success'
+        run: |
+          docker logout ${{ env.REGISTRY }}
+        shell: bash

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

@@ -0,0 +1,41 @@
+name: Download a Docker image and its tags from GitHub artifacts, apply them, and push the image to the Registry; then merge them
+
+on:
+  workflow_call:
+    inputs:
+      image:
+        description: Image name
+        required: true
+        type: string
+      variant:
+        description: Variant tag prefix
+        required: true
+        type: string
+    secrets:
+      REGISTRY_USERNAME:
+        required: true
+      REGISTRY_TOKEN:
+        required: true
+
+permissions:
+  contents: read
+
+jobs:
+  tag-push:
+    uses: ./.github/workflows/docker-tag-push.yml
+    with:
+      image: ${{ inputs.image }}
+      variant: ${{ inputs.variant }}
+    secrets:
+      REGISTRY_USERNAME: ${{ secrets.REGISTRY_USERNAME }}
+      REGISTRY_TOKEN: ${{ secrets.REGISTRY_TOKEN }}
+
+  tag-merge:
+    uses: ./.github/workflows/docker-tag-merge.yml
+    needs: tag-push
+    with:
+      image: ${{ inputs.image }}
+      variant: ${{ inputs.variant }}
+    secrets:
+      REGISTRY_USERNAME: ${{ secrets.REGISTRY_USERNAME }}
+      REGISTRY_TOKEN: ${{ secrets.REGISTRY_TOKEN }}

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

@@ -18,7 +18,7 @@ on:
         type: string
       timeout-minutes:
         description: Timeout in minutes
-        default: 20
+        default: 15
         type: number
     secrets:
       REGISTRY_USERNAME:
@@ -33,29 +33,22 @@ jobs:
   tag-push:
     runs-on: ubuntu-24.04
     timeout-minutes: ${{ inputs.timeout-minutes }}
+    strategy:
+      matrix:
+        platform: [aarch64, x86_64]
 
     steps:
       - name: Checkout Repo ⚡️
         uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
-      - name: Free disk space 🧹
-        uses: ./.github/actions/free-disk-space
       - name: Create dev environment 📦
         uses: ./.github/actions/create-dev-env
 
-      - name: Download aarch64 image tar and apply tags 🏷
+      - name: Download image tar and apply tags 🏷
         uses: ./.github/actions/apply-single-tags
         with:
           image: ${{ inputs.image }}
           variant: ${{ inputs.variant }}
-          platform: aarch64
-        if: ${{ !contains(inputs.variant, 'cuda11') }}
-
-      - 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
+          platform: ${{ matrix.platform }}
 
       - name: Login to Registry 🔐
         if: env.PUSH_TO_REGISTRY == 'true'
@@ -76,14 +69,6 @@ jobs:
           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/
-        shell: bash
-
       - name: Logout from Registry 🔐
         if: always() && env.PUSH_TO_REGISTRY == 'true' && steps.login.outcome == 'success'
         run: |

.github/workflows/docker.yml

@@ -15,7 +15,7 @@ on:
       # We use local reusable workflows to make architecture clean and simple
       # https://docs.github.com/en/actions/sharing-automations/reusing-workflows
       - ".github/workflows/docker-build-test-upload.yml"
-      - ".github/workflows/docker-tag-push.yml"
+      - ".github/workflows/docker-tag-push-merge.yml"
       - ".github/workflows/docker-wiki-update.yml"
 
       # We use local composite actions to combine multiple workflow steps within one action
@@ -39,7 +39,7 @@ on:
     paths:
       - ".github/workflows/docker.yml"
       - ".github/workflows/docker-build-test-upload.yml"
-      - ".github/workflows/docker-tag-push.yml"
+      - ".github/workflows/docker-tag-push-merge.yml"
       - ".github/workflows/docker-wiki-update.yml"
 
       - ".github/actions/apply-single-tags/action.yml"
@@ -376,7 +376,7 @@ jobs:
     needs: [aarch64-base, x86_64-base, aarch64-minimal, x86_64-minimal]
 
   tag-push:
-    uses: ./.github/workflows/docker-tag-push.yml
+    uses: ./.github/workflows/docker-tag-push-merge.yml
     with:
       image: ${{ matrix.image }}
       variant: ${{ matrix.variant }}
@@ -443,7 +443,7 @@ jobs:
     if: ${{ !contains(github.event.pull_request.title, '[FAST_BUILD]') }}
 
   tag-push-fast:
-    uses: ./.github/workflows/docker-tag-push.yml
+    uses: ./.github/workflows/docker-tag-push-merge.yml
     with:
       image: ${{ matrix.image }}
       variant: ${{ matrix.variant }}

docs/using/custom-images.md

@@ -21,7 +21,7 @@ Our repository provides several customization points:
 - `REGISTRY`, `OWNER` (part of `env` in some GitHub workflows) - these allow to properly tag and refer to images during following steps:
   - [`build-test-upload`](https://github.com/jupyter/docker-stacks/blob/main/.github/workflows/docker-build-test-upload.yml)
   - [`contributed-recipes`](https://github.com/jupyter/docker-stacks/blob/main/.github/workflows/contributed-recipes.yml)
-  - [`tag-push`](https://github.com/jupyter/docker-stacks/blob/main/.github/workflows/docker-tag-push.yml)
+  - [`tag-push-merge`](https://github.com/jupyter/docker-stacks/blob/main/.github/workflows/docker-tag-push-merge.yml)
 
 These customization points can't be changed during runtime.
 Read more about [Docker build arguments](https://docs.docker.com/build/building/variables/#arg-usage-example) and [GitHub environment variables for a single workflow](https://docs.github.com/en/actions/writing-workflows/choosing-what-your-workflow-does/store-information-in-variables#defining-environment-variables-for-a-single-workflow).

tagging/apps/merge_tags.py

@@ -16,17 +16,15 @@
 from tagging.apps.config import Config
 from tagging.utils.get_platform import ALL_PLATFORMS
 from tagging.utils.get_prefix import get_file_prefix_for_platform
-from tagging.utils.git_helper import GitHelper
 
 docker = plumbum.local["docker"]
 
 LOGGER = logging.getLogger(__name__)
 
 
-def read_local_tags_from_files(config: Config) -> tuple[list[str], set[str]]:
+def read_local_tags_from_files(config: Config) -> set[str]:
     LOGGER.info(f"Read tags from file(s) for image: {config.image}")
 
-    all_local_tags = []
     merged_local_tags = set()
     for platform in ALL_PLATFORMS:
         LOGGER.info(f"Reading tags for platform: {platform}")
@@ -42,73 +40,67 @@ def read_local_tags_from_files(config: Config) -> tuple[list[str], set[str]]:
 
         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 all_local_tags, merged_local_tags
+    return merged_local_tags
 
 
 @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4))
-def pull_tag(tag: str) -> None:
-    LOGGER.info(f"Pulling tag: {tag}")
-    docker["pull", tag] & plumbum.FG
-    LOGGER.info(f"Tag {tag} pulled successfully")
+def inspect_manifest(tag: str) -> None:
+    LOGGER.info(f"Inspecting manifest for tag: {tag}")
+    docker["buildx", "imagetools", "inspect", tag] & plumbum.FG
+    LOGGER.info(f"Manifest {tag} exists")
 
 
-def pull_missing_tags(merged_tag: str, all_local_tags: list[str]) -> list[str]:
-    existing_platform_tags = []
+def find_platform_tags(merged_tag: str) -> list[str]:
+    platform_tags = []
 
     for platform in ALL_PLATFORMS:
         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")
+        LOGGER.warning(f"Trying to inspect: {platform_tag} in the registry")
         try:
-            pull_tag(platform_tag)
-            existing_platform_tags.append(platform_tag)
-            LOGGER.info(f"Tag {platform_tag} pulled successfully")
+            inspect_manifest(platform_tag)
+            platform_tags.append(platform_tag)
+            LOGGER.info(f"Tag {platform_tag} found successfully")
         except RetryError:
-            LOGGER.warning(f"Pull failed, tag {platform_tag} doesn't exist")
+            LOGGER.warning(f"Manifest for tag {platform_tag} doesn't exist")
 
-    return existing_platform_tags
+    return platform_tags
 
 
-def merge_tags(
-    merged_tag: str, all_local_tags: list[str], push_to_registry: bool
-) -> None:
+def merge_tags(merged_tag: 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)
+    platform_tags = find_platform_tags(merged_tag)
+    if not platform_tags:
+        assert not push_to_registry, (
+            f"No platform tags found for merged tag: {merged_tag}, "
+            "and push to registry is enabled. "
+            "Cannot create a manifest for a non-existing image."
+        )
+        LOGGER.info(
+            f"Not running merge for tag: {merged_tag} as no platform tags found"
+        )
+        return
+
     args = [
         "buildx",
         "imagetools",
         "create",
-        *existing_platform_tags,
+        *platform_tags,
         "--tag",
         merged_tag,
     ]
     if not push_to_registry:
         args.append("--dry-run")
 
-    commit_hash_tag = GitHelper.commit_hash_tag()
-    if not push_to_registry and merged_tag.endswith(commit_hash_tag):
-        LOGGER.info(
-            f"Not running merge for tag: {merged_tag} as it's a commit SHA tag and it wasn't pushed to registry"
-        )
-        return
-
     LOGGER.info(f"Running command: {' '.join(args)}")
     docker[args] & plumbum.FG
     if push_to_registry:
-        LOGGER.info(f"Pushed merged tag: {merged_tag} to registry")
+        LOGGER.info(f"Pushed merged tag: {merged_tag}")
     else:
-        LOGGER.info(f"Skipping push for tag: {merged_tag}")
+        LOGGER.info(f"Skipped push for tag: {merged_tag}")
 
 
 if __name__ == "__main__":
@@ -119,8 +111,8 @@ def merge_tags(
 
     LOGGER.info(f"Merging tags for image: {config.image}")
 
-    all_local_tags, merged_local_tags = read_local_tags_from_files(config)
+    merged_local_tags = read_local_tags_from_files(config)
     for tag in merged_local_tags:
-        merge_tags(tag, all_local_tags, push_to_registry)
+        merge_tags(tag, push_to_registry)
 
     LOGGER.info(f"Successfully merged tags for image: {config.image}")