From 19bfa74b5fc1e47aeb973edf45b5c69421a2eb09 Mon Sep 17 00:00:00 2001
From: Ayaz Salikhov <mathbunnyru@gmail.com>
Date: Fri, 28 Mar 2025 14:07:28 +0000
Subject: [PATCH] Improve merge_tags: runnable multiple times, run in PRs,
 simpler function

---
 .github/workflows/docker-merge-tags.yml |  1 -
 tagging/apps/merge_tags.py              | 57 +++++++++++++++----------
 2 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/.github/workflows/docker-merge-tags.yml b/.github/workflows/docker-merge-tags.yml
index 19d755d3..aafdf34d 100644
--- a/.github/workflows/docker-merge-tags.yml
+++ b/.github/workflows/docker-merge-tags.yml
@@ -66,7 +66,6 @@ jobs:
           password: ${{ secrets.REGISTRY_TOKEN }}
 
       - name: Merge tags for the images 🔀
-        if: env.PUSH_TO_REGISTRY == 'true'
         run: |
           python3 -m tagging.apps.merge_tags \
             --image ${{ inputs.image }} \
diff --git a/tagging/apps/merge_tags.py b/tagging/apps/merge_tags.py
index dbc90d68..6e5b6457 100755
--- a/tagging/apps/merge_tags.py
+++ b/tagging/apps/merge_tags.py
@@ -2,6 +2,7 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 import logging
+import os
 
 import plumbum
 
@@ -36,36 +37,46 @@ def read_tags_from_files(config: Config) -> set[str]:
     return tags
 
 
-def merge_tags(config: Config) -> None:
-    LOGGER.info(f"Merging tags for image: {config.image}")
+def merge_tags(tag: str, push_to_registry: bool) -> None:
+    LOGGER.info(f"Trying to merge tag: {tag}")
+    all_platform_tags = []
+    for platform in ALL_PLATFORMS:
+        platform_tag = tag.replace(":", f":{platform}-")
+        LOGGER.info(f"Trying to pull: {platform_tag}")
+        try:
+            docker["pull", platform_tag] & plumbum.FG
+            all_platform_tags.append(platform_tag)
+            LOGGER.info("Pull success")
+        except plumbum.ProcessExecutionError:
+            LOGGER.info("Pull failed, image with this tag and platform doesn't exist")
 
-    all_tags = read_tags_from_files(config)
-    for tag in all_tags:
-        LOGGER.info(f"Trying to merge tag: {tag}")
-        existing_images = []
-        for platform in ALL_PLATFORMS:
-            image_with_platform = tag.replace(":", f":{platform}-")
-            LOGGER.info(f"Trying to pull: {image_with_platform}")
-            try:
-                docker["pull", image_with_platform] & plumbum.FG
-                existing_images.append(image_with_platform)
-                LOGGER.info("Pull success")
-            except plumbum.ProcessExecutionError:
-                LOGGER.info(
-                    "Pull failed, image with this tag and platform doesn't exist"
-                )
+    LOGGER.info(f"Found images: {all_platform_tags}")
+    try:
+        docker["manifest", "rm", tag] & plumbum.FG
+        LOGGER.info(f"Manifest {tag} already exists, removing it")
+    except plumbum.ProcessExecutionError:
+        LOGGER.info(f"Manifest {tag} doesn't exist")
 
-        LOGGER.info(f"Found images: {existing_images}")
-        docker["manifest", "create", tag][existing_images] & plumbum.FG
+    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}")
+
+    if push_to_registry:
+        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"All tags merged for image: {config.image}")
+    else:
+        LOGGER.info(f"Skipping push for tag: {tag}")
 
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO)
 
     config = common_arguments_parser(image=True, variant=True, tags_dir=True)
-    merge_tags(config)
+    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)
+    LOGGER.info(f"Successfully merged tags for image: {config.image}")