Unrolled build for rust-lang#134898

Rollup merge of rust-lang#134898 - Kobzol:ci-python-script, r=MarcoIeni

Make it easier to run CI jobs locally

This PR extends the Python CI script to perform a poor man's CI-like execution of a given CI job locally. It's not perfect, but it's better than nothing.

r? `@jieyouxu`

.github/workflows/ci.yml

@@ -2,7 +2,7 @@
 # and also on pushes to special branches (auto, try).
 #
 # The actual definition of the executed jobs is calculated by a Python
-# script located at src/ci/github-actions/calculate-job-matrix.py, which
+# script located at src/ci/github-actions/ci.py, which
 # uses job definition data from src/ci/github-actions/jobs.yml.
 # You should primarily modify the `jobs.yml` file if you want to modify
 # what jobs are executed in CI.
@@ -56,7 +56,7 @@ jobs:
       - name: Calculate the CI job matrix
         env:
           COMMIT_MESSAGE: ${{ github.event.head_commit.message }}
-        run: python3 src/ci/github-actions/calculate-job-matrix.py >> $GITHUB_OUTPUT
+        run: python3 src/ci/github-actions/ci.py calculate-job-matrix >> $GITHUB_OUTPUT
         id: jobs
   job:
     name: ${{ matrix.name }}

src/ci/docker/README.md

@@ -1,29 +1,30 @@
 # Docker images for CI
 
 This folder contains a bunch of docker images used by the continuous integration
-(CI) of Rust. An script is accompanied (`run.sh`) with these images to actually
-execute them. To test out an image execute:
+(CI) of Rust. A script is accompanied (`run.sh`) with these images to actually
+execute them.
 
-```
-./src/ci/docker/run.sh $image_name
-```
+Note that a single Docker image can be used by multiple CI jobs, so the job name
+is the important thing that you should know. You can examine the existing CI jobs in
+the [`jobs.yml`](../github-actions/jobs.yml) file.
 
-for example:
+To run a specific CI job locally, you can use the following script:
 
 ```
-./src/ci/docker/run.sh x86_64-gnu
+python3 ./src/ci/github-actions/ci.py run-local <job-name>
 ```
 
-Images will output artifacts in an `obj/$image_name` dir at the root of a repository. Note
-that the script will overwrite the contents of this directory.
-
-To match conditions in rusts CI, also set the environment variable `DEPLOY=1`, e.g.:
+For example, to run the `x86_64-gnu-llvm-18-1` job:
 ```
-DEPLOY=1 ./src/ci/docker/run.sh x86_64-gnu
+python3 ./src/ci/github-actions/ci.py run-local x86_64-gnu-llvm-18-1
 ```
 
+The job will output artifacts in an `obj/<image-name>` dir at the root of a repository. Note
+that the script will overwrite the contents of this directory. `<image-name>` is set based on the
+Docker image executed in the given CI job.
+
 **NOTE**: In CI, the script outputs the artifacts to the `obj` directory,
-while locally, to the `obj/$image_name` directory. This is primarily to prevent
+while locally, to the `obj/<image-name>` directory. This is primarily to prevent
 strange linker errors when using multiple Docker images.
 
 For some Linux workflows (for example `x86_64-gnu-llvm-18-N`), the process is more involved. You will need to see which script is executed for the given workflow inside the [`jobs.yml`](../github-actions/jobs.yml) file and pass it through the `DOCKER_SCRIPT` environment variable. For example, to reproduce the `x86_64-gnu-llvm-18-3` workflow, you can run the following script:

src/ci/github-actions/calculate-job-matrix.py → src/ci/github-actions/ci.py

@@ -1,18 +1,20 @@
 #!/usr/bin/env python3
 
 """
-This script serves for generating a matrix of jobs that should
-be executed on CI.
+This script contains CI functionality.
+It can be used to generate a matrix of jobs that should
+be executed on CI, or run a specific CI job locally.
 
-It reads job definitions from `src/ci/github-actions/jobs.yml`
-and filters them based on the event that happened on CI.
+It reads job definitions from `src/ci/github-actions/jobs.yml`.
 """
 
+import argparse
 import dataclasses
 import json
 import logging
 import os
 import re
+import subprocess
 import typing
 from pathlib import Path
 from typing import List, Dict, Any, Optional
@@ -25,25 +27,35 @@
 Job = Dict[str, Any]
 
 
-def name_jobs(jobs: List[Dict], prefix: str) -> List[Job]:
+def add_job_properties(jobs: List[Dict], prefix: str) -> List[Job]:
     """
-    Add a `name` attribute to each job, based on its image and the given `prefix`.
+    Modify the `name` attribute of each job, based on its base name and the given `prefix`.
+    Add an `image` attribute to each job, based on its image.
     """
+    modified_jobs = []
     for job in jobs:
-        job["name"] = f"{prefix} - {job['image']}"
-    return jobs
+        # Create a copy of the `job` dictionary to avoid modifying `jobs`
+        new_job = dict(job)
+        new_job["image"] = get_job_image(new_job)
+        new_job["name"] = f"{prefix} - {new_job['name']}"
+        modified_jobs.append(new_job)
+    return modified_jobs
 
 
 def add_base_env(jobs: List[Job], environment: Dict[str, str]) -> List[Job]:
     """
     Prepends `environment` to the `env` attribute of each job.
     The `env` of each job has higher precedence than `environment`.
     """
+    modified_jobs = []
     for job in jobs:
         env = environment.copy()
         env.update(job.get("env", {}))
-        job["env"] = env
-    return jobs
+
+        new_job = dict(job)
+        new_job["env"] = env
+        modified_jobs.append(new_job)
+    return modified_jobs
 
 
 @dataclasses.dataclass
@@ -116,7 +128,9 @@ def find_run_type(ctx: GitHubCtx) -> Optional[WorkflowRunType]:
 
 def calculate_jobs(run_type: WorkflowRunType, job_data: Dict[str, Any]) -> List[Job]:
     if isinstance(run_type, PRRunType):
-        return add_base_env(name_jobs(job_data["pr"], "PR"), job_data["envs"]["pr"])
+        return add_base_env(
+            add_job_properties(job_data["pr"], "PR"), job_data["envs"]["pr"]
+        )
     elif isinstance(run_type, TryRunType):
         jobs = job_data["try"]
         custom_jobs = run_type.custom_jobs
@@ -130,7 +144,7 @@ def calculate_jobs(run_type: WorkflowRunType, job_data: Dict[str, Any]) -> List[
             jobs = []
             unknown_jobs = []
             for custom_job in custom_jobs:
-                job = [j for j in job_data["auto"] if j["image"] == custom_job]
+                job = [j for j in job_data["auto"] if j["name"] == custom_job]
                 if not job:
                     unknown_jobs.append(custom_job)
                     continue
@@ -140,10 +154,10 @@ def calculate_jobs(run_type: WorkflowRunType, job_data: Dict[str, Any]) -> List[
                     f"Custom job(s) `{unknown_jobs}` not found in auto jobs"
                 )
 
-        return add_base_env(name_jobs(jobs, "try"), job_data["envs"]["try"])
+        return add_base_env(add_job_properties(jobs, "try"), job_data["envs"]["try"])
     elif isinstance(run_type, AutoRunType):
         return add_base_env(
-            name_jobs(job_data["auto"], "auto"), job_data["envs"]["auto"]
+            add_job_properties(job_data["auto"], "auto"), job_data["envs"]["auto"]
         )
 
     return []
@@ -181,12 +195,64 @@ def format_run_type(run_type: WorkflowRunType) -> str:
         raise AssertionError()
 
 
-if __name__ == "__main__":
-    logging.basicConfig(level=logging.INFO)
+def get_job_image(job: Job) -> str:
+    """
+    By default, the Docker image of a job is based on its name.
+    However, it can be overridden by its IMAGE environment variable.
+    """
+    env = job.get("env", {})
+    # Return the IMAGE environment variable if it exists, otherwise return the job name
+    return env.get("IMAGE", job["name"])
 
-    with open(JOBS_YAML_PATH) as f:
-        data = yaml.safe_load(f)
 
+def is_linux_job(job: Job) -> bool:
+    return "ubuntu" in job["os"]
+
+
+def find_linux_job(job_data: Dict[str, Any], job_name: str, pr_jobs: bool) -> Job:
+    candidates = job_data["pr"] if pr_jobs else job_data["auto"]
+    jobs = [job for job in candidates if job.get("name") == job_name]
+    if len(jobs) == 0:
+        available_jobs = "\n".join(
+            sorted(job["name"] for job in candidates if is_linux_job(job))
+        )
+        raise Exception(f"""Job `{job_name}` not found in {'pr' if pr_jobs else 'auto'} jobs.
+The following jobs are available:
+{available_jobs}""")
+    assert len(jobs) == 1
+
+    job = jobs[0]
+    if not is_linux_job(job):
+        raise Exception("Only Linux jobs can be executed locally")
+    return job
+
+
+def run_workflow_locally(job_data: Dict[str, Any], job_name: str, pr_jobs: bool):
+    DOCKER_DIR = Path(__file__).absolute().parent.parent / "docker"
+
+    job = find_linux_job(job_data, job_name=job_name, pr_jobs=pr_jobs)
+
+    custom_env = {}
+    # Replicate src/ci/scripts/setup-environment.sh
+    # Adds custom environment variables to the job
+    if job_name.startswith("dist-"):
+        if job_name.endswith("-alt"):
+            custom_env["DEPLOY_ALT"] = "1"
+        else:
+            custom_env["DEPLOY"] = "1"
+    custom_env.update({k: str(v) for (k, v) in job.get("env", {}).items()})
+
+    args = [str(DOCKER_DIR / "run.sh"), get_job_image(job)]
+    env_formatted = [f"{k}={v}" for (k, v) in sorted(custom_env.items())]
+    print(f"Executing `{' '.join(env_formatted)} {' '.join(args)}`")
+
+    env = os.environ.copy()
+    env.update(custom_env)
+
+    subprocess.run(args, env=env)
+
+
+def calculate_job_matrix(job_data: Dict[str, Any]):
     github_ctx = get_github_ctx()
 
     run_type = find_run_type(github_ctx)
@@ -197,7 +263,7 @@ def format_run_type(run_type: WorkflowRunType) -> str:
 
     jobs = []
     if run_type is not None:
-        jobs = calculate_jobs(run_type, data)
+        jobs = calculate_jobs(run_type, job_data)
     jobs = skip_jobs(jobs, channel)
 
     if not jobs:
@@ -208,3 +274,45 @@ def format_run_type(run_type: WorkflowRunType) -> str:
     logging.info(f"Output:\n{yaml.dump(dict(jobs=jobs, run_type=run_type), indent=4)}")
     print(f"jobs={json.dumps(jobs)}")
     print(f"run_type={run_type}")
+
+
+def create_cli_parser():
+    parser = argparse.ArgumentParser(
+        prog="ci.py", description="Generate or run CI workflows"
+    )
+    subparsers = parser.add_subparsers(
+        help="Command to execute", dest="command", required=True
+    )
+    subparsers.add_parser(
+        "calculate-job-matrix",
+        help="Generate a matrix of jobs that should be executed in CI",
+    )
+    run_parser = subparsers.add_parser(
+        "run-local", help="Run a CI jobs locally (on Linux)"
+    )
+    run_parser.add_argument(
+        "job_name",
+        help="CI job that should be executed. By default, a merge (auto) "
+        "job with the given name will be executed",
+    )
+    run_parser.add_argument(
+        "--pr", action="store_true", help="Run a PR job instead of an auto job"
+    )
+    return parser
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+
+    with open(JOBS_YAML_PATH) as f:
+        data = yaml.safe_load(f)
+
+    parser = create_cli_parser()
+    args = parser.parse_args()
+
+    if args.command == "calculate-job-matrix":
+        calculate_job_matrix(data)
+    elif args.command == "run-local":
+        run_workflow_locally(data, args.job_name, args.pr)
+    else:
+        raise Exception(f"Unknown command {args.command}")