Improve merge_tags: runnable multiple times, run in PRs, simpler function

Author: mathbunnyru

.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 }} \

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}")
-
-    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: {existing_images}")
-        docker["manifest", "create", tag][existing_images] & plumbum.FG
+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")
+
+    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"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}")