diff --git a/.gitignore b/.gitignore index 184ff17ab25e..e9b9743f1359 100644 --- a/.gitignore +++ b/.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/ diff --git a/Jenkinsfile b/Jenkinsfile index d7d261ec9967..3f82ff184013 100755 --- a/Jenkinsfile +++ b/Jenkinsfile @@ -45,7 +45,7 @@ // 'python3 jenkins/generate.py' // Note: This timestamp is here to ensure that updates to the Jenkinsfile are // always rebased on main before merging: -// Generated at 2022-06-20T19:48:32.482249 +// Generated at 2022-06-22T10:07:00.173803 import org.jenkinsci.plugins.pipeline.modeldefinition.Utils // NOTE: these lines are scanned by docker/dev_common.sh. Please update the regex as needed. --> @@ -244,6 +244,55 @@ def prepare() { node('CPU-SMALL') { ws("workspace/exec_${env.EXECUTOR_NUMBER}/tvm/prepare") { init_git() + + if (env.DETERMINE_DOCKER_IMAGES == 'yes') { + sh( + script: "./tests/scripts/determine_docker_images.py ci_arm=${ci_arm} ci_cpu=${ci_cpu} ci_gpu=${ci_gpu} ci_hexagon=${ci_hexagon} ci_i386=${ci_i386} ci_lint=${ci_lint} ci_qemu=${ci_qemu} ci_wasm=${ci_wasm} ", + label: 'Decide whether to use tlcpack or tlcpackstaging for Docker images', + ) + // Pull image names from the results of should_rebuild_docker.py + ci_arm = sh( + script: "cat .docker-image-names/ci_arm", + label: "Find docker image name for ci_arm", + returnStdout: true, + ).trim() + ci_cpu = sh( + script: "cat .docker-image-names/ci_cpu", + label: "Find docker image name for ci_cpu", + returnStdout: true, + ).trim() + ci_gpu = sh( + script: "cat .docker-image-names/ci_gpu", + label: "Find docker image name for ci_gpu", + returnStdout: true, + ).trim() + ci_hexagon = sh( + script: "cat .docker-image-names/ci_hexagon", + label: "Find docker image name for ci_hexagon", + returnStdout: true, + ).trim() + ci_i386 = sh( + script: "cat .docker-image-names/ci_i386", + label: "Find docker image name for ci_i386", + returnStdout: true, + ).trim() + ci_lint = sh( + script: "cat .docker-image-names/ci_lint", + label: "Find docker image name for ci_lint", + returnStdout: true, + ).trim() + ci_qemu = sh( + script: "cat .docker-image-names/ci_qemu", + label: "Find docker image name for ci_qemu", + returnStdout: true, + ).trim() + ci_wasm = sh( + script: "cat .docker-image-names/ci_wasm", + label: "Find docker image name for ci_wasm", + returnStdout: true, + ).trim() + } + ci_arm = params.ci_arm_param ?: ci_arm ci_cpu = params.ci_cpu_param ?: ci_cpu ci_gpu = params.ci_gpu_param ?: ci_gpu diff --git a/jenkins/Prepare.groovy.j2 b/jenkins/Prepare.groovy.j2 index 894ddc72eeb7..d9cfa440c7e2 100644 --- a/jenkins/Prepare.groovy.j2 +++ b/jenkins/Prepare.groovy.j2 @@ -141,6 +141,22 @@ def prepare() { node('CPU-SMALL') { ws("workspace/exec_${env.EXECUTOR_NUMBER}/tvm/prepare") { init_git() + + if (env.DETERMINE_DOCKER_IMAGES == 'yes') { + 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() + {% endfor %} + } + {% for image in images %} {{ image.name }} = params.{{ image.name }}_param ?: {{ image.name }} {% endfor %} diff --git a/tests/python/ci/test_ci.py b/tests/python/ci/test_ci.py index b712b6780cd7..2d3373c54535 100644 --- a/tests/python/ci/test_ci.py +++ b/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", [ diff --git a/tests/scripts/determine_docker_images.py b/tests/scripts/determine_docker_images.py new file mode 100755 index 000000000000..dbcde82cff7a --- /dev/null +++ b/tests/scripts/determine_docker_images.py @@ -0,0 +1,115 @@ +#!/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", + 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, spec in images.items(): + if image_exists(spec): + logging.info(f"{spec} found in tlcpack") + images_to_use[filename] = spec + else: + logging.info(f"{spec} not found in tlcpack, using tlcpackstaging") + part, tag = spec.split(":") + user, repo = part.split("/") + tlcpackstaging_tag = f"tlcpackstaging/{repo.replace('-', '_')}:{tag}" + images_to_use[filename] = tlcpackstaging_tag + + 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)