Merge tags and create multi-platform images (#1789)

* Merge tags and create multi-platform images

* Add new files :)

* Provide images

* Fix

* Fix

* Use the same naming style as python for the images

* Fix docker image rmi command

* Use plumbum.FG

* Fix

* Add missing path

* Few merging fixes

* Add pushing manifests

* Add missing path

* Add arch to apply tags

* Use platform word

* Remove unused step

* Merge tags for tensorflow as well

* Do not use print

* Make merge_tags work when image doesn't exist

* Add logs

* Add icon

* Add tensorflow-notebook workaround

* Unify platform usage

* Remove unused function

* Fix

* Update docs

tagging/apply_tags.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+import argparse
+import logging
+from pathlib import Path
+
+import plumbum
+
+docker = plumbum.local["docker"]
+
+LOGGER = logging.getLogger(__name__)
+
+
+def apply_tags(
+    short_image_name: str,
+    owner: str,
+    tags_dir: Path,
+    platform: str,
+) -> None:
+    """
+    Tags <owner>/<short_image_name>:latest with the tags
+    reported by all taggers for the given image.
+    """
+    LOGGER.info(f"Tagging image: {short_image_name}")
+
+    image = f"{owner}/{short_image_name}:latest"
+    filename = f"{platform}-{short_image_name}.txt"
+    tags = (tags_dir / filename).read_text().splitlines()
+
+    for tag in tags:
+        LOGGER.info(f"Applying tag: {tag}")
+        docker["tag", image, tag] & plumbum.FG
+
+    LOGGER.info("Removing latest tag from the image")
+    docker["image", "rmi", image] & plumbum.FG
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+
+    arg_parser = argparse.ArgumentParser()
+    arg_parser.add_argument(
+        "--short-image-name",
+        required=True,
+        help="Short image name to apply tags for",
+    )
+    arg_parser.add_argument(
+        "--tags-dir",
+        required=True,
+        type=Path,
+        help="Directory with saved tags file",
+    )
+    arg_parser.add_argument(
+        "--platform",
+        required=True,
+        type=str,
+        choices=["x86_64", "aarch64"],
+        help="Image platform",
+    )
+    arg_parser.add_argument(
+        "--owner",
+        required=True,
+        help="Owner of the image",
+    )
+    args = arg_parser.parse_args()
+
+    apply_tags(args.short_image_name, args.owner, args.tags_dir, args.platform)

tagging/get_platform.py

@@ -0,0 +1,14 @@
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+import platform
+
+ALL_PLATFORMS = {"x86_64", "aarch64"}
+
+
+def get_platform() -> str:
+    machine = platform.machine()
+    return {
+        "aarch64": "aarch64",
+        "arm64": "aarch64",  # To support local building on aarch64 Macs
+        "x86_64": "x86_64",
+    }[machine]

tagging/get_tags_prefix.py

@@ -1,8 +0,0 @@
-# Copyright (c) Jupyter Development Team.
-# Distributed under the terms of the Modified BSD License.
-import platform
-
-
-def get_tags_prefix() -> str:
-    machine = platform.machine()
-    return "" if machine == "x86_64" else f"{machine}-"

tagging/merge_tags.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+import argparse
+import logging
+from pathlib import Path
+
+import plumbum
+
+from tagging.get_platform import ALL_PLATFORMS
+
+docker = plumbum.local["docker"]
+
+LOGGER = logging.getLogger(__name__)
+
+
+def merge_tags(
+    short_image_name: str,
+    tags_dir: Path,
+) -> None:
+    """
+    Merge tags for x86_64 and aarch64 images when possible.
+    """
+    LOGGER.info(f"Merging tags for image: {short_image_name}")
+
+    all_tags: set[str] = set()
+
+    for platform in ALL_PLATFORMS:
+        filename = f"{platform}-{short_image_name}.txt"
+        tags = (tags_dir / filename).read_text().splitlines()
+        all_tags.update(tag.replace(platform + "-", "") for tag in tags)
+
+    LOGGER.info(f"Got tags: {all_tags}")
+
+    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
+        )
+        docker["manifest", "push", tag] & plumbum.FG
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+
+    arg_parser = argparse.ArgumentParser()
+    arg_parser.add_argument(
+        "--short-image-name",
+        required=True,
+        help="Short image name to apply tags for",
+    )
+    arg_parser.add_argument(
+        "--tags-dir",
+        required=True,
+        type=Path,
+        help="Directory with saved tags file",
+    )
+    args = arg_parser.parse_args()
+
+    merge_tags(args.short_image_name, args.tags_dir)

tagging/write_manifest.py

@@ -4,14 +4,13 @@
 import argparse
 import datetime
 import logging
-import platform
 from pathlib import Path
 
 from docker.models.containers import Container
 
 from tagging.docker_runner import DockerRunner
+from tagging.get_platform import get_platform
 from tagging.get_taggers_and_manifests import get_taggers_and_manifests
-from tagging.get_tags_prefix import get_tags_prefix
 from tagging.git_helper import GitHelper
 from tagging.manifests import ManifestHeader, ManifestInterface
 
@@ -68,11 +67,6 @@ def write_manifest_file(
     (manifest_dir / f"{filename}.md").write_text(markdown_content)
 
 
-def get_file_prefix() -> str:
-    machine = platform.machine()
-    return "amd64" if machine == "x86_64" else "aarch64"
-
-
 def write_manifest(
     short_image_name: str,
     owner: str,
@@ -84,13 +78,15 @@ def write_manifest(
 
     image = f"{owner}/{short_image_name}:latest"
 
-    file_prefix = get_file_prefix()
+    file_prefix = get_platform()
     commit_hash_tag = GitHelper.commit_hash_tag()
     filename = f"{file_prefix}-{short_image_name}-{commit_hash_tag}"
 
     with DockerRunner(image) as container:
-        tags_prefix = get_tags_prefix()
-        all_tags = [tags_prefix + tagger.tag_value(container) for tagger in taggers]
+        tags_prefix = get_platform()
+        all_tags = [
+            tags_prefix + "-" + tagger.tag_value(container) for tagger in taggers
+        ]
         write_build_history_line(
             short_image_name, owner, hist_line_dir, filename, all_tags
         )

tagging/write_tags_file.py

@@ -3,42 +3,41 @@
 # Distributed under the terms of the Modified BSD License.
 import argparse
 import logging
-
-import plumbum
+from pathlib import Path
 
 from tagging.docker_runner import DockerRunner
+from tagging.get_platform import get_platform
 from tagging.get_taggers_and_manifests import get_taggers_and_manifests
-from tagging.get_tags_prefix import get_tags_prefix
-
-docker = plumbum.local["docker"]
 
 LOGGER = logging.getLogger(__name__)
 
 
-def tag_image(short_image_name: str, owner: str) -> None:
+def write_tags_file(
+    short_image_name: str,
+    owner: str,
+    tags_dir: Path,
+) -> None:
     """
-    Tags <owner>/<short_image_name>:latest with the tags reported by all taggers
-    for the given image.
+    Writes tags file for the image <owner>/<short_image_name>:latest
     """
     LOGGER.info(f"Tagging image: {short_image_name}")
     taggers, _ = get_taggers_and_manifests(short_image_name)
 
     image = f"{owner}/{short_image_name}:latest"
-    tags_prefix = get_tags_prefix()
+    tags_prefix = get_platform()
+    filename = f"{tags_prefix}-{short_image_name}.txt"
 
+    tags = [f"{owner}/{short_image_name}:{tags_prefix}-latest"]
     with DockerRunner(image) as container:
         for tagger in taggers:
             tagger_name = tagger.__class__.__name__
             tag_value = tagger.tag_value(container)
             LOGGER.info(
-                f"Applying tag, tagger_name: {tagger_name} tag_value: {tag_value}"
+                f"Calculated tag, tagger_name: {tagger_name} tag_value: {tag_value}"
             )
-            docker[
-                "tag", image, f"{owner}/{short_image_name}:{tags_prefix}{tag_value}"
-            ]()
-    if tags_prefix != "":
-        LOGGER.info(f"Adding {tags_prefix}latest tag")
-        docker["tag", image, f"{owner}/{short_image_name}:{tags_prefix}latest"]()
+            tags.append(f"{owner}/{short_image_name}:{tags_prefix}-{tag_value}")
+    tags_dir.mkdir(parents=True, exist_ok=True)
+    (tags_dir / filename).write_text("\n".join(tags))
 
 
 if __name__ == "__main__":
@@ -48,7 +47,13 @@ if __name__ == "__main__":
     arg_parser.add_argument(
         "--short-image-name",
         required=True,
-        help="Short image name to apply tags for",
+        help="Short image name to write tags for",
+    )
+    arg_parser.add_argument(
+        "--tags-dir",
+        required=True,
+        type=Path,
+        help="Directory to save tags file",
     )
     arg_parser.add_argument(
         "--owner",
@@ -57,4 +62,4 @@ if __name__ == "__main__":
     )
     args = arg_parser.parse_args()
 
-    tag_image(args.short_image_name, args.owner)
+    write_tags_file(args.short_image_name, args.owner, args.tags_dir)