[ci][docker] Fall back to tlcpackstaging if images don't exist

See #11768

.gitignore

@@ -265,3 +265,6 @@ gallery/how_to/work_with_microtvm/micro_tvmc.py
 
 # Test sample data files
 !tests/python/ci/sample_prs/*.json
+
+# Used in CI to communicate between Python and Jenkins
+.docker-image-names/

jenkins/Deploy.groovy.j2

@@ -86,6 +86,21 @@ def deploy() {
         }
       }
     }
+    if (env.BRANCH_NAME == 'main' && env.RETAG_STAGING_IMAGES == 'yes') {
+      for (tlcpackstaging_image in used_tlcpackstaging_images) {
+        def tlcpack_image = tlcpackstaging_image.replace('tlcpackstaging', 'tlcpack')
+        def tag = tlcpackstaging_image.split(':')[1]
+        sh(
+          script: """
+            set -eux
+            CONTENT_TYPE="application/vnd.docker.distribution.manifest.v2+json"
+            MANIFEST=$(curl -H "Accept: \${CONTENT_TYPE}" "https://hub.docker.com/v2/${tlcpackstaging_image}/manifests/${tag}")
+            curl -X PUT -H "Content-Type: \${CONTENT_TYPE}" -d "\${MANIFEST}" "https://hub.docker.com/v2/${tlcpack_image}/manifests/${tag}"
+          """,
+          label: 'Tag tlcpackstaging image to tlcpack',
+        )
+      }
+    }
     if (env.BRANCH_NAME == 'main' && env.DEPLOY_DOCKER_IMAGES == 'yes' && rebuild_docker_images && upstream_revision != null) {
       node('CPU') {
         ws({{ m.per_exec_ws('tvm/deploy-docker') }}) {

jenkins/Jenkinsfile.j2

@@ -82,6 +82,7 @@ docker_build = 'docker/build.sh'
 // timeout in minutes
 max_time = 180
 rebuild_docker_images = false
+used_tlcpackstaging_images = []
 
 // skips builds from branch indexing; sourced from https://www.jvt.me/posts/2020/02/23/jenkins-multibranch-skip-branch-index/
 // execute this before anything else, including requesting any time on an agent

jenkins/Prepare.groovy.j2

@@ -141,6 +141,23 @@ def prepare() {
     node('CPU-SMALL') {
       ws("workspace/exec_${env.EXECUTOR_NUMBER}/tvm/prepare") {
         init_git()
+
+        sh(
+          script: "./tests/scripts/determine_docker_images.py {% for image in images %}{{ image.name }}={% raw %}${{% endraw %}{{ image.name }}{% raw %}}{% endraw %} {% endfor %}",
+          label: 'Decide whether to use tlcpack or tlcpackstaging for Docker images',
+        )
+        // Pull image names from the results of should_rebuild_docker.py
+        {% for image in images %}
+        {{ image.name }} = sh(
+          script: "cat .docker-image.names/{{ image.name }}",
+          label: "Find docker image name for {{ image.name }}",
+          returnStdout: true,
+        ).trim()
+        if ({{ image.name }}.contains("tlcpackstaging")) {
+          used_tlcpackstaging_images.add({{ image.name }})
+        }
+        {% endfor %}
+
         {% for image in images %}
         {{ image.name }} = params.{{ image.name }}_param ?: {{ image.name }}
         {% endfor %}

tests/python/ci/test_ci.py

@@ -788,6 +788,59 @@ def run(type, data, check):
     )
 
 
+@pytest.mark.parametrize(
+    "images,expected",
+    [
+        (
+            ["ci_arm=tlcpack/ci-arm:abc-abc-123", "ci_lint=tlcpack/ci-lint:abc-abc-234"],
+            {
+                "ci_arm": "tlcpack/ci-arm:abc-abc-123",
+                "ci_lint": "tlcpack/ci-lint:abc-abc-234",
+            },
+        ),
+        (
+            ["ci_arm2=tlcpack/ci-arm2:abc-abc-123"],
+            {
+                "ci_arm2": "tlcpackstaging/ci-arm2:abc-abc-123",
+            },
+        ),
+    ],
+)
+def test_determine_docker_images(tmpdir_factory, images, expected):
+    tag_script = REPO_ROOT / "tests" / "scripts" / "determine_docker_images.py"
+
+    dir = tmpdir_factory.mktemp("tmp_git_dir")
+
+    docker_data = {
+        "repositories/tlcpack/ci-arm/tags/abc-abc-123": {},
+        "repositories/tlcpack/ci-lint/tags/abc-abc-234": {},
+    }
+
+    proc = subprocess.run(
+        [
+            str(tag_script),
+            "--testing-docker-data",
+            json.dumps(docker_data),
+            "--base-dir",
+            dir,
+        ]
+        + images,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        encoding="utf-8",
+        cwd=dir,
+        check=False,
+    )
+    if proc.returncode != 0:
+        raise RuntimeError(f"Failed to run script:\n{proc.stdout}")
+
+    for expected_filename, expected_image in expected.items():
+        with open(Path(dir) / expected_filename) as f:
+            actual_image = f.read()
+
+        assert actual_image == expected_image
+
+
 @pytest.mark.parametrize(
     "changed_files,name,check,expected_code",
     [

tests/scripts/determine_docker_images.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import argparse
+import datetime
+import json
+import logging
+import urllib.error
+from pathlib import Path
+
+from typing import Dict, Any
+
+
+from http_utils import get
+from cmd_utils import init_log, REPO_ROOT
+
+
+DOCKER_API_BASE = "https://hub.docker.com/v2/"
+PAGE_SIZE = 25
+TEST_DATA = None
+
+
+def docker_api(url: str, use_pagination: bool = False) -> Dict[str, Any]:
+    """
+    Run a paginated fetch from the public Docker Hub API
+    """
+    if TEST_DATA is not None:
+        if url not in TEST_DATA:
+            raise urllib.error.HTTPError(url, 404, "Not found", {}, None)
+        return TEST_DATA[url]
+    pagination = ""
+    if use_pagination:
+        pagination = f"?page_size={PAGE_SIZE}&page=1"
+    url = DOCKER_API_BASE + url + pagination
+    r, headers = get(url)
+    reset = headers.get("x-ratelimit-reset")
+    if reset is not None:
+        reset = datetime.datetime.fromtimestamp(int(reset))
+        reset = reset.isoformat()
+    logging.info(
+        f"Docker API Rate Limit: {headers.get('x-ratelimit-remaining')} / {headers.get('x-ratelimit-limit')} (reset at {reset})"
+    )
+    return r
+
+
+def image_exists(spec: str) -> bool:
+    name, tag = spec.split(":")
+    try:
+        r = docker_api(f"repositories/{name}/tags/{tag}")
+        logging.info(f"Image exists, got response: {json.dumps(r, indent=2)}")
+        return True
+    except urllib.error.HTTPError as e:
+        # Image was not found
+        logging.exception(e)
+        return False
+
+
+if __name__ == "__main__":
+    init_log()
+    parser = argparse.ArgumentParser(
+        description="Writes out Docker images names to be used to .docker-image-names/"
+    )
+    parser.add_argument(
+        "--testing-docker-data",
+        help="(testing only) JSON data to mock response from Docker Hub API",
+    )
+    parser.add_argument(
+        "--base-dir",
+        default=".docker-image-names",
+        required=True,
+        help="(testing only) Folder to write image names to",
+    )
+    args, other = parser.parse_known_args()
+    name_dir = Path(args.base_dir)
+
+    images = {}
+    for item in other:
+        name, tag = item.split("=")
+        images[name] = tag
+
+    if args.testing_docker_data is not None:
+        TEST_DATA = json.loads(args.testing_docker_data)
+
+    logging.info(f"Checking if these images exist in tlcpack: {images}")
+
+    name_dir.mkdir(exist_ok=True)
+    images_to_use = {}
+    for filename, tag in images.items():
+        if image_exists(tag):
+            logging.info(f"{tag} found in tlcpack")
+            images_to_use[filename] = tag
+        else:
+            logging.info(f"{tag} not found in tlcpack, using tlcpackstaging")
+            images_to_use[filename] = tag.replace("tlcpack", "tlcpackstaging")
+
+    for filename, image in images_to_use.items():
+        logging.info(f"Writing image {image} to {name_dir / filename}")
+        with open(name_dir / filename, "w") as f:
+            f.write(image)