[ci][docker] Regenerate Jenkinsfile on each run (#11886)

This makes it so the generated PRs pass CI

Co-authored-by: driazati <driazati@users.noreply.github.com>

jenkins/requirements.txt

@@ -1,2 +1 @@
-Jinja2
-
+Jinja2>=3.0.0

tests/python/ci/test_ci.py

@@ -845,13 +845,14 @@ def run(source_type, data, check):
             ]
         },
         expected="Tag names were the same, no update needed",
+        expected_images=[],
     ),
     dict(
         tlcpackstaging_body={
             "results": [
                 {
                     "last_updated": "2022-06-01T00:00:00.123456Z",
-                    "name": "abc-abc-234",
+                    "name": "abc-abc-234-staging",
                 },
             ]
         },
@@ -864,6 +865,9 @@ def run(source_type, data, check):
             ]
         },
         expected="Using tlcpackstaging tag on tlcpack",
+        expected_images=[
+            "ci_arm = 'tlcpack/ci-arm:abc-abc-234-staging'",
+        ],
     ),
     dict(
         tlcpackstaging_body={
@@ -883,9 +887,14 @@ def run(source_type, data, check):
             ]
         },
         expected="Found newer image, using: tlcpack",
+        expected_images=[
+            "ci_arm = 'tlcpack/ci-arm:abc-abc-234'",
+        ],
     ),
 )
-def test_open_docker_update_pr(tmpdir_factory, tlcpackstaging_body, tlcpack_body, expected):
+def test_open_docker_update_pr(
+    tmpdir_factory, tlcpackstaging_body, tlcpack_body, expected, expected_images
+):
     """Test workflow to open a PR to update Docker images"""
     tag_script = REPO_ROOT / "tests" / "scripts" / "open_docker_update_pr.py"
 
@@ -926,6 +935,10 @@ def test_open_docker_update_pr(tmpdir_factory, tlcpackstaging_body, tlcpack_body
         check=False,
     )
 
+    for line in expected_images:
+        if line not in proc.stdout:
+            raise RuntimeError(f"Missing line {line} in output:\n{proc.stdout}")
+
     if proc.returncode != 0:
         raise RuntimeError(f"Process failed:\nstdout:\n{proc.stdout}\n\nstderr:\n{proc.stderr}")
 

tests/scripts/open_docker_update_pr.py

@@ -25,10 +25,12 @@
 from urllib import error
 from typing import List, Dict, Any, Optional, Callable
 from git_utils import git, parse_remote, GitHubRepo
-from cmd_utils import REPO_ROOT, init_log
+from cmd_utils import REPO_ROOT, init_log, Sh
 from should_rebuild_docker import docker_api
 
 JENKINSFILE = REPO_ROOT / "jenkins" / "Jenkinsfile.j2"
+GENERATED_JENKINSFILE = REPO_ROOT / "Jenkinsfile"
+GENERATE_SCRIPT = REPO_ROOT / "jenkins" / "generate.py"
 GITHUB_TOKEN = os.environ["GITHUB_TOKEN"]
 BRANCH = "nightly-docker-update"
 
@@ -124,6 +126,7 @@ def latest_tlcpackstaging_image(source: str) -> Optional[str]:
 
     # Build a new Jenkinsfile with the latest images from tlcpack or tlcpackstaging
     new_content = []
+    replacements = {}
     for line in content:
         m = re.match(r"^(ci_[a-zA-Z0-9]+) = \'(.*)\'", line.strip())
         if m is not None:
@@ -135,7 +138,9 @@ def latest_tlcpackstaging_image(source: str) -> Optional[str]:
                 new_content.append(line)
             else:
                 logging.info(f"Using new image {new_image}")
-                new_content.append(f"{groups[0]} = '{new_image}'\n")
+                new_line = f"{groups[0]} = '{new_image}'\n"
+                new_content.append(new_line)
+                replacements[line] = new_line
         else:
             new_content.append(line)
 
@@ -147,6 +152,20 @@ def latest_tlcpackstaging_image(source: str) -> Optional[str]:
         with open(JENKINSFILE, "w") as f:
             f.write("".join(new_content))
 
+    # Re-generate the Jenkinsfile
+    logging.info(f"Editing {GENERATED_JENKINSFILE}")
+    with open(GENERATED_JENKINSFILE) as f:
+        generated_content = f.read()
+
+    for original_line, new_line in replacements.items():
+        generated_content = generated_content.replace(original_line, new_line)
+
+    if args.dry_run:
+        print(f"Would have written:\n{generated_content}")
+    else:
+        with open(GENERATED_JENKINSFILE, "w") as f:
+            f.write(generated_content)
+
     # Publish the PR
     title = "[ci][docker] Nightly Docker image update"
     body = "This bumps the Docker images to the latest versions from Docker Hub."
@@ -158,6 +177,7 @@ def latest_tlcpackstaging_image(source: str) -> Optional[str]:
         logging.info(f"Creating git commit")
         git(["checkout", "-B", BRANCH])
         git(["add", str(JENKINSFILE.relative_to(REPO_ROOT))])
+        git(["add", str(GENERATED_JENKINSFILE.relative_to(REPO_ROOT))])
         git(["config", "user.name", "tvm-bot"])
         git(["config", "user.email", "95660001+tvm-bot@users.noreply.github.com"])
         git(["commit", "-m", message])