[ci][docker] Automatically update docker images

.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,7 +86,7 @@ def deploy() {
         }
       }
     }
-    if (env.BRANCH_NAME == 'main' && env.DEPLOY_DOCKER_IMAGES == 'yes' && rebuild_docker_images && upstream_revision != null) {
+    if (env.BRANCH_NAME == 'main' && env.PUSH_DOCKER_IMAGES == 'yes' && rebuild_docker_images && upstream_revision != null) {
       node('CPU') {
         ws({{ m.per_exec_ws('tvm/deploy-docker') }}) {
           try {

jenkins/Prepare.groovy.j2

@@ -145,13 +145,6 @@ def prepare() {
         {{ image.name }} = params.{{ image.name }}_param ?: {{ image.name }}
         {% endfor %}
 
-        sh (script: """
-          echo "Docker images being used in this build:"
-          {% for image in images %}
-          echo " {{ image.name }} = ${ {{- image.name -}} }"
-          {% endfor %}
-        """, label: 'Docker image names')
-
         is_docs_only_build = sh (
           returnStatus: true,
           script: './tests/scripts/git_change_docs.sh',
@@ -161,13 +154,32 @@ def prepare() {
         skip_slow_tests = should_skip_slow_tests(env.CHANGE_ID)
         rebuild_docker_images = sh (
           returnStatus: true,
-          script: './tests/scripts/git_change_docker.sh',
+          script: './tests/scripts/should_rebuild_docker.py',
           label: 'Check for any docker changes',
         )
         if (skip_ci) {
           // Don't rebuild when skipping CI
           rebuild_docker_images = false
         }
+        if (!rebuild_docker_images) {
+          // Pull image names from the results of should_rebuild_docker.py
+          if (env.USE_AUTOUPDATED_DOCKER_IMAGES == 'yes') {
+            {% 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 %}
+          }
+        }
+
+        sh (script: """
+          echo "Docker images being used in this build:"
+          {% for image in images %}
+          echo " {{ image.name }} = ${ {{- image.name -}} }"
+          {% endfor %}
+        """, label: 'Docker image names')
       }
     }
   }

tests/scripts/should_rebuild_docker.py

@@ -25,10 +25,11 @@
 
 
 from http_utils import get
-from cmd_utils import Sh, init_log
+from cmd_utils import Sh, init_log, REPO_ROOT
 
 
 DOCKER_API_BASE = "https://hub.docker.com/v2/"
+DOCKER_USER = "tlcpack"
 PAGE_SIZE = 25
 TEST_DATA = None
 
@@ -111,7 +112,9 @@ def find_commit_in_repo(tags: List[Dict[str, Any]]):
 
 def main():
     # Fetch all tlcpack images
-    images = docker_api("repositories/tlcpack")
+    images = docker_api(f"repositories/{DOCKER_USER}")
+    NAMES_DIR = REPO_ROOT / ".docker-image-names"
+    NAMES_DIR.mkdir(exist_ok=True)
 
     # Ignore all non-ci images
     relevant_images = [image for image in images["results"] if image["name"].startswith("ci-")]
@@ -120,7 +123,7 @@ def main():
 
     for image in relevant_images:
         # Check the tags for the image
-        tags = docker_api(f"repositories/tlcpack/{image['name']}/tags")
+        tags = docker_api(f"repositories/{DOCKER_USER}/{image['name']}/tags")
 
         # Find the hash of the most recent tag
         shorthash, tag = find_commit_in_repo(tags)
@@ -132,6 +135,11 @@ def main():
             logging.info(f"Found docker changes from {shorthash} when checking {name}")
             logging.info(diff)
             exit(2)
+        else:
+            fullname = f"{DOCKER_USER}/{image['name']}:{name}"
+            logging.info(f"Using image {fullname} for {image}")
+            with open(NAMES_DIR / image["name"], "w") as f:
+                f.write(fullname)
 
     logging.info("Did not find changes, no rebuild necessary")
     exit(0)
@@ -140,7 +148,8 @@ def main():
 if __name__ == "__main__":
     init_log()
     parser = argparse.ArgumentParser(
-        description="Exits 0 if Docker images don't need to be rebuilt, 1 otherwise"
+        description="Exits 0 if Docker images don't need to be rebuilt, 1 otherwise. "
+        "If 0, image names will be written to a .docker-image-names folder"
     )
     parser.add_argument(
         "--testing-docker-data",