ci/run_n_monitor: Add --include-stage argument

Add basic filtering by stages, starting with stage inclusion.

Signed-off-by: Daniel Stone <daniels@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30784>

bin/ci/ci_run_n_monitor.py

@@ -114,6 +114,7 @@ def monitor_pipeline(
     project: gitlab.v4.objects.Project,
     pipeline: gitlab.v4.objects.ProjectPipeline,
     target_jobs_regex: re.Pattern,
+    include_stage_regex: re.Pattern,
     dependencies: set[str],
     force_manual: bool,
     stress: int,
@@ -137,7 +138,9 @@ def monitor_pipeline(
     if stress:
         # When stress test, it is necessary to collect this information before start.
         for job in pipeline.jobs.list(all=True, include_retried=True):
-            if target_jobs_regex.fullmatch(job.name) and job.status in COMPLETED_STATUSES:
+            if target_jobs_regex.fullmatch(job.name) and \
+               include_stage_regex.fullmatch(job.stage) and \
+               job.status in COMPLETED_STATUSES:
                 stress_status_counter[job.name][job.status] += 1
                 execution_times[job.name][job.id] = (job_duration(job), job.status, job.web_url)
 
@@ -158,7 +161,8 @@ def monitor_pipeline(
         to_cancel = []
         jobs_waiting.clear()
         for job in sorted(pipeline.jobs.list(all=True), key=lambda j: j.name):
-            if target_jobs_regex.fullmatch(job.name):
+            if target_jobs_regex.fullmatch(job.name) and \
+               include_stage_regex.fullmatch(job.stage):
                 target_id = job.id
                 target_status = job.status
 
@@ -360,10 +364,20 @@ def parse_args() -> argparse.Namespace:
         "--target",
         metavar="target-job",
         help="Target job regex. For multiple targets, pass multiple values, "
-             "eg. `--target foo bar`.",
+             "eg. `--target foo bar`. Only jobs in the target stage(s) "
+             "supplied, and their dependencies, will be considered.",
         required=True,
         nargs=argparse.ONE_OR_MORE,
     )
+    parser.add_argument(
+        "--include-stage",
+        metavar="include-stage",
+        help="Job stages to include when searching for target jobs. "
+             "For multiple targets, pass multiple values, eg. "
+             "`--include-stage foo bar`.",
+        default=[".*"],
+        nargs=argparse.ONE_OR_MORE,
+    )
     parser.add_argument(
         "--token",
         metavar="token",
@@ -438,6 +452,7 @@ def print_detected_jobs(
 def find_dependencies(
     token: str | None,
     target_jobs_regex: re.Pattern,
+    include_stage_regex: re.Pattern,
     project_path: str,
     iid: int
 ) -> set[str]:
@@ -466,7 +481,7 @@ def find_dependencies(
         gql_instance, {"projectPath": project_path.path_with_namespace, "iid": iid}
     )
 
-    target_dep_dag = filter_dag(dag, target_jobs_regex)
+    target_dep_dag = filter_dag(dag, target_jobs_regex, include_stage_regex)
     if not target_dep_dag:
         print(Fore.RED + "The job(s) were not found in the pipeline." + Fore.RESET)
         sys.exit(1)
@@ -584,14 +599,22 @@ def main() -> None:
 
         target_jobs_regex = re.compile(target)
 
+        include_stage = '|'.join(args.include_stage)
+        include_stage = include_stage.strip()
+
+        print("🞋 target from stages: " + Fore.BLUE + include_stage + Style.RESET_ALL)  # U+1F78B Round target
+
+        include_stage_regex = re.compile(include_stage)
+
         deps = find_dependencies(
             token=token,
             target_jobs_regex=target_jobs_regex,
+            include_stage_regex=include_stage_regex,
             iid=pipe.iid,
             project_path=cur_project
         )
         target_job_id, ret, exec_t = monitor_pipeline(
-            cur_project, pipe, target_jobs_regex, deps, args.force_manual, args.stress
+            cur_project, pipe, target_jobs_regex, include_stage_regex, deps, args.force_manual, args.stress
         )
 
         if target_job_id:

bin/ci/gitlab_gql.py

@@ -325,9 +325,15 @@ def create_job_needs_dag(gl_gql: GitlabGQL, params, disable_cache: bool = True)
     return final_dag
 
 
-def filter_dag(dag: Dag, regex: Pattern) -> Dag:
-    jobs_with_regex: set[str] = {job for job in dag if regex.fullmatch(job)}
-    return Dag({job: data for job, data in dag.items() if job in sorted(jobs_with_regex)})
+def filter_dag(dag: Dag, job_name_regex: Pattern, include_stage_regex: Pattern) -> Dag:
+    filtered_jobs: Dag = Dag({})
+    for (job, data) in dag.items():
+        if not job_name_regex.fullmatch(job):
+            continue
+        if not include_stage_regex.fullmatch(data["stage"]):
+            continue
+        filtered_jobs[job] = data
+    return filtered_jobs
 
 
 def print_dag(dag: Dag) -> None:
@@ -472,8 +478,16 @@ def parse_args() -> Namespace:
         "--regex",
         type=str,
         required=False,
+        default=".*",
         help="Regex pattern for the job name to be considered",
     )
+    parser.add_argument(
+        "--include-stage",
+        type=str,
+        required=False,
+        default=".*",
+        help="Regex pattern for the stage name to be considered",
+    )
     mutex_group_print = parser.add_mutually_exclusive_group()
     mutex_group_print.add_argument(
         "--print-dag",
@@ -515,8 +529,7 @@ def main():
             gl_gql, {"projectPath": args.project_path, "iid": iid}, disable_cache=True
         )
 
-        if args.regex:
-            dag = filter_dag(dag, re.compile(args.regex))
+        dag = filter_dag(dag, re.compile(args.regex), re.compile(args.include_stage))
 
         print_dag(dag)