Add an ability to specify registry when using docker images (#2008)

* Add an ability to specify registry when using docker images

* Fix typo

* [TMP] Speedup workflow

* Revert "[TMP] Speedup workflow"

This reverts commit 3af0055ccf.

tagging/apply_tags.py

@@ -16,6 +16,7 @@ LOGGER = logging.getLogger(__name__)
 
 def apply_tags(
     short_image_name: str,
+    registry: str,
     owner: str,
     tags_dir: Path,
     platform: str,
@@ -26,7 +27,7 @@ def apply_tags(
     """
     LOGGER.info(f"Tagging image: {short_image_name}")
 
-    image = f"{owner}/{short_image_name}:latest"
+    image = f"{registry}/{owner}/{short_image_name}:latest"
     filename = f"{platform}-{short_image_name}.txt"
     tags = (tags_dir / filename).read_text().splitlines()
 
@@ -60,6 +61,13 @@ if __name__ == "__main__":
         choices=["x86_64", "aarch64", "arm64"],
         help="Image platform",
     )
+    arg_parser.add_argument(
+        "--registry",
+        required=True,
+        type=str,
+        choices=["docker.io", "quay.io"],
+        help="Image registry",
+    )
     arg_parser.add_argument(
         "--owner",
         required=True,
@@ -68,4 +76,6 @@ if __name__ == "__main__":
     args = arg_parser.parse_args()
     args.platform = unify_aarch64(args.platform)
 
-    apply_tags(args.short_image_name, args.owner, args.tags_dir, args.platform)
+    apply_tags(
+        args.short_image_name, args.registry, args.owner, args.tags_dir, args.platform
+    )

tagging/manifests.py

@@ -23,13 +23,18 @@ class ManifestHeader:
     """ManifestHeader doesn't fall under common interface, and we run it separately"""
 
     @staticmethod
-    def create_header(short_image_name: str, owner: str, build_timestamp: str) -> str:
+    def create_header(
+        short_image_name: str, registry: str, owner: str, build_timestamp: str
+    ) -> str:
         commit_hash = GitHelper.commit_hash()
         commit_hash_tag = GitHelper.commit_hash_tag()
         commit_message = GitHelper.commit_message()
 
         image_size = docker[
-            "images", f"{owner}/{short_image_name}:latest", "--format", "{{.Size}}"
+            "images",
+            f"{registry}/{owner}/{short_image_name}:latest",
+            "--format",
+            "{{.Size}}",
         ]().rstrip()
 
         return "\n".join(
@@ -39,7 +44,7 @@ class ManifestHeader:
                 "## Build Info",
                 "",
                 f"* Build datetime: {build_timestamp}",
-                f"* Docker image: {owner}/{short_image_name}:{commit_hash_tag}",
+                f"* Docker image: {registry}/{owner}/{short_image_name}:{commit_hash_tag}",
                 f"* Docker image size: {image_size}",
                 f"* Git commit SHA: [{commit_hash}](https://github.com/jupyter/docker-stacks/commit/{commit_hash})",
                 "* Git commit message:",

tagging/write_manifest.py

@@ -23,6 +23,7 @@ MARKDOWN_LINE_BREAK = "<br />"
 
 def write_build_history_line(
     short_image_name: str,
+    registry: str,
     owner: str,
     hist_line_dir: Path,
     filename: str,
@@ -32,7 +33,7 @@ def write_build_history_line(
 
     date_column = f"`{BUILD_TIMESTAMP}`"
     image_column = MARKDOWN_LINE_BREAK.join(
-        f"`{owner}/{short_image_name}:{tag_value}`" for tag_value in all_tags
+        f"`{registry}/{owner}/{short_image_name}:{tag_value}`" for tag_value in all_tags
     )
     commit_hash = GitHelper.commit_hash()
     links_column = MARKDOWN_LINE_BREAK.join(
@@ -49,6 +50,7 @@ def write_build_history_line(
 
 def write_manifest_file(
     short_image_name: str,
+    registry: str,
     owner: str,
     manifest_dir: Path,
     filename: str,
@@ -59,7 +61,7 @@ def write_manifest_file(
     LOGGER.info(f"Using manifests: {manifest_names}")
 
     markdown_pieces = [
-        ManifestHeader.create_header(short_image_name, owner, BUILD_TIMESTAMP)
+        ManifestHeader.create_header(short_image_name, registry, owner, BUILD_TIMESTAMP)
     ] + [manifest.markdown_piece(container) for manifest in manifests]
     markdown_content = "\n\n".join(markdown_pieces) + "\n"
 
@@ -69,6 +71,7 @@ def write_manifest_file(
 
 def write_manifest(
     short_image_name: str,
+    registry: str,
     owner: str,
     hist_line_dir: Path,
     manifest_dir: Path,
@@ -76,7 +79,7 @@ def write_manifest(
     LOGGER.info(f"Creating manifests for image: {short_image_name}")
     taggers, manifests = get_taggers_and_manifests(short_image_name)
 
-    image = f"{owner}/{short_image_name}:latest"
+    image = f"{registry}/{owner}/{short_image_name}:latest"
 
     file_prefix = get_platform()
     commit_hash_tag = GitHelper.commit_hash_tag()
@@ -88,10 +91,16 @@ def write_manifest(
             tags_prefix + "-" + tagger.tag_value(container) for tagger in taggers
         ]
         write_build_history_line(
-            short_image_name, owner, hist_line_dir, filename, all_tags
+            short_image_name, registry, owner, hist_line_dir, filename, all_tags
         )
         write_manifest_file(
-            short_image_name, owner, manifest_dir, filename, manifests, container
+            short_image_name,
+            registry,
+            owner,
+            manifest_dir,
+            filename,
+            manifests,
+            container,
         )
 
 
@@ -116,6 +125,13 @@ if __name__ == "__main__":
         type=Path,
         help="Directory to save manifest file",
     )
+    arg_parser.add_argument(
+        "--registry",
+        required=True,
+        type=str,
+        choices=["docker.io", "quay.io"],
+        help="Image registry",
+    )
     arg_parser.add_argument(
         "--owner",
         required=True,
@@ -126,5 +142,9 @@ if __name__ == "__main__":
     LOGGER.info(f"Current build timestamp: {BUILD_TIMESTAMP}")
 
     write_manifest(
-        args.short_image_name, args.owner, args.hist_line_dir, args.manifest_dir
+        args.short_image_name,
+        args.registry,
+        args.owner,
+        args.hist_line_dir,
+        args.manifest_dir,
     )

tagging/write_tags_file.py

@@ -14,6 +14,7 @@ LOGGER = logging.getLogger(__name__)
 
 def write_tags_file(
     short_image_name: str,
+    registry: str,
     owner: str,
     tags_dir: Path,
 ) -> None:
@@ -23,11 +24,11 @@ def write_tags_file(
     LOGGER.info(f"Tagging image: {short_image_name}")
     taggers, _ = get_taggers_and_manifests(short_image_name)
 
-    image = f"{owner}/{short_image_name}:latest"
+    image = f"{registry}/{owner}/{short_image_name}:latest"
     tags_prefix = get_platform()
     filename = f"{tags_prefix}-{short_image_name}.txt"
 
-    tags = [f"{owner}/{short_image_name}:{tags_prefix}-latest"]
+    tags = [f"{registry}/{owner}/{short_image_name}:{tags_prefix}-latest"]
     with DockerRunner(image) as container:
         for tagger in taggers:
             tagger_name = tagger.__class__.__name__
@@ -55,6 +56,13 @@ if __name__ == "__main__":
         type=Path,
         help="Directory to save tags file",
     )
+    arg_parser.add_argument(
+        "--registry",
+        required=True,
+        type=str,
+        choices=["docker.io", "quay.io"],
+        help="Image registry",
+    )
     arg_parser.add_argument(
         "--owner",
         required=True,
@@ -62,4 +70,4 @@ if __name__ == "__main__":
     )
     args = arg_parser.parse_args()
 
-    write_tags_file(args.short_image_name, args.owner, args.tags_dir)
+    write_tags_file(args.short_image_name, args.registry, args.owner, args.tags_dir)