ci: Enhance nightly/mr/weekly error reporting

.gitlab/scripts/build.sh

@@ -48,6 +48,11 @@ if [[ -n "$TE_GIT_REF" ]]; then
     ADDITIONAL_PARAMS+=("--build-arg TE_COMMIT=${TE_GIT_REF}")
 fi
 
+if [[ "$FILE" == "Dockerfile.linting" ]]; then
+    ADDITIONAL_PARAMS+=("--build-arg CI_SERVER_URL=${CI_SERVER_URL}")
+    ADDITIONAL_PARAMS+=("--secret id=NEMO_CI_TRIAGE_TOKEN,env=PAT")
+fi
+
 echo $(git rev-parse HEAD)
 
 JET_API_VERSION=$(curl -s -u "$ARTIFACTORY_USER:$ARTIFACTORY_TOKEN" "https://sc-hw-artf.nvidia.com/artifactory/api/pypi/hw-joc-pypi/simple/jet-api/" | grep -o 'href="../../jet-api/[0-9.]*/' | sed 's|href="../../jet-api/||;s|/||' | sort -V -r | head -n1)

.gitlab/stages/02.test.yml

@@ -82,6 +82,7 @@ test:unit_tests_configure:
         "--dependent-job test:unit_tests_configure"
         "--slurm-account ${CI_SLURM_ACCOUNT}"
         "--no-enable-warmup"
+        "--enable-error-extraction"
       )
     - |
       export PYTHONPATH=$(pwd)

.gitlab/stages/03.integration-tests.yml

@@ -56,6 +56,7 @@ integration:configure:
         "--no-enable-warmup"
         "--dependent-job integration:configure"
         "--enable-lightweight-mode"
+        "--enable-error-extraction"
       )
     - |
       export PYTHONPATH=$(pwd)

.gitlab/stages/04.functional-tests.yml

@@ -87,6 +87,7 @@ functional:configure:
         "--record-checkpoints ${RECORD_CHECKPOINTS}"
         "--slurm-account ${CI_SLURM_ACCOUNT}"
         "--no-enable-warmup"
+        "--enable-error-extraction"
       )
     - |
       SMOKE_ARGS=(
@@ -101,6 +102,7 @@ functional:configure:
         "--record-checkpoints false"
         "--slurm-account ${CI_SLURM_ACCOUNT}"
         "--no-enable-warmup"
+        "--enable-error-extraction"
       )
     - |
       export PYTHONPATH=$(pwd)

docker/Dockerfile.linting

@@ -21,3 +21,13 @@ ARG JET_API_VERSION
 RUN --mount=type=secret,id=JET_INDEX_URLS \
       JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS) && \ 
       uv pip install --no-cache-dir "jet-client~=2.0" --upgrade $JET_INDEX_URLS
+
+# Keep this in the internal-only stage so public CI has no internal service dependency.
+ARG CI_SERVER_URL
+ARG NEMO_CI_TRIAGE_COMMIT=8e65fa4ae20b58578d0e0f20ebea37ee7d92c8ea
+RUN --mount=type=secret,id=NEMO_CI_TRIAGE_TOKEN \
+      GIT_CONFIG_COUNT=1 \
+      GIT_CONFIG_KEY_0=http.extraHeader \
+      GIT_CONFIG_VALUE_0="Authorization: Basic $(printf 'oauth2:%s' "$(cat /run/secrets/NEMO_CI_TRIAGE_TOKEN)" | base64 -w0)" \
+      uv pip install --no-cache-dir \
+        "nemo-ci-triage @ git+${CI_SERVER_URL}/dl/nemo/nemo-ci-triage.git@${NEMO_CI_TRIAGE_COMMIT}"

tests/test_utils/python_scripts/generate_jet_trigger_job.py

@@ -9,6 +9,26 @@
 from tests.test_utils.python_scripts import recipe_parser
 
 BASE_PATH = pathlib.Path(__file__).parent.resolve()
+TRIAGE_LOG_PATH = "jet_workload.log"
+TRIAGE_REPORT_PATH = "error_report.json"
+
+
+def build_test_script(command: str) -> str:
+    """Wrap a workload command with non-blocking error extraction."""
+    return "\n".join(
+        [
+            "set +e",
+            "set -o pipefail",
+            f"{command} 2>&1 | tee {TRIAGE_LOG_PATH}",
+            'exit_code=${PIPESTATUS[0]}',
+            "set -e",
+            (
+                f"extract-errors {TRIAGE_LOG_PATH} --output {TRIAGE_REPORT_PATH} "
+                '--exit-code "$exit_code" || true'
+            ),
+            'exit "$exit_code"',
+        ]
+    )
 
 
 @click.command()
@@ -70,6 +90,11 @@
         "Empty/unset disables the cadence filter."
     ),
 )
+@click.option(
+    "--enable-error-extraction/--no-enable-error-extraction",
+    default=False,
+    help="Extract a structured error report from GitLab child-job output.",
+)
 def main(
     scope: str,
     environment: str,
@@ -91,7 +116,8 @@ def main(
     enable_lightweight_mode: bool = False,
     enable_warmup: Optional[bool] = None,
     cadence: Optional[str] = None,
-):
+    enable_error_extraction: bool = False,
+) -> None:
     # Treat empty string as "no cadence filter" so callers can wire shell
     # variables in directly without conditional flag emission.
     cadence_arg = cadence or None
@@ -217,14 +243,20 @@ def main(
                 elif warmup_job != "":
                     needs.append({"job": warmup_job})
 
+            test_script = " ".join(script)
+            artifact_paths = ["results/"]
+            if enable_error_extraction:
+                test_script = build_test_script(test_script)
+                artifact_paths.extend([TRIAGE_LOG_PATH, TRIAGE_REPORT_PATH])
+
             gitlab_pipeline[test_case['spec']['test_case']] = {
                 "stage": f"{test_case['spec']['model']}",
                 "image": f"{container_image}:{container_tag}",
                 "tags": job_tags,
                 "timeout": "7 days",
                 "needs": needs,
-                "script": [" ".join(script)],
-                "artifacts": {"paths": ["results/"], "when": "always"},
+                "script": [test_script],
+                "artifacts": {"paths": artifact_paths, "when": "always"},
                 "allow_failure": test_case["spec"].get("allow_failure", False)
                 or test_case["spec"]["model"] == "gpt-nemo",
                 "retry": {

tests/test_utils/python_scripts/notify.py

@@ -5,50 +5,67 @@
 
 import click
 import gitlab
-import pandas as pd
-import requests
-import slack_sdk
+from nemo_ci_triage.slack_notification import notification
 
 PROJECT_ID = int(os.getenv("CI_PROJECT_ID", 19378))
 WEBHOOK_URL = os.getenv("WEBHOOK_URL", "")
-GITLAB_ENDPOINT = os.getenv('GITLAB_ENDPOINT')
-TAG_TEAM = bool(os.getenv('TAG_TEAM', 0))
-TEAM_SLUG = str(os.getenv('TEAM_SLUG'))
+GITLAB_ENDPOINT = os.getenv("GITLAB_ENDPOINT")
+if not GITLAB_ENDPOINT:
+    raise ValueError("GITLAB_ENDPOINT is required")
+SERVER_URL = f"https://{GITLAB_ENDPOINT}"
+PROJECT_URL = os.getenv("CI_PROJECT_URL", f"{SERVER_URL}/ADLR/megatron-lm")
+TAG_TEAM = os.getenv("TAG_TEAM", "0") == "1"
+TEAM_SLUG = os.getenv("TEAM_SLUG", "")
+
+JOB_PREFIXES = {
+    "unit-tests": "test:unit_tests",
+    "integration-tests": "integration:run_",
+    "functional-tests": "functional:run_",
+    "smoke-tests": "functional:smoke-",
+}
 
 logging.basicConfig()
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
 
-def get_gitlab_handle():
-    return gitlab.Gitlab(f"https://{GITLAB_ENDPOINT}", private_token=os.getenv("RO_API_TOKEN"))
+def get_gitlab_handle() -> gitlab.Gitlab:
+    return gitlab.Gitlab(SERVER_URL, private_token=os.getenv("RO_API_TOKEN"))
 
 
-def get_jobs_per_bridge(pipeline_id: int, type_of_job: str):
-    bridge = {}
-    for pipeline_bridge in (
-        get_gitlab_handle()
-        .projects.get(PROJECT_ID)
-        .pipelines.get(pipeline_id)
-        .bridges.list(get_all=True)
-    ):
-        if (
-            not pipeline_bridge.name.startswith(type_of_job)
-            or pipeline_bridge.attributes['downstream_pipeline'] is None
-        ):
+def _bridge_gpu(bridge_name: str) -> str:
+    for gpu in ("GB200", "H100", "A100"):
+        if gpu.lower() in bridge_name.lower():
+            return gpu
+    return "Unknown"
+
+
+def get_pipeline_jobs(pipeline_id: int, job_prefix: str) -> list[tuple[str, int, list[dict]]]:
+    """Collect Megatron-LM's direct child pipelines using nemo-ci-triage-2."""
+    project = get_gitlab_handle().projects.get(PROJECT_ID)
+    root_pipeline = project.pipelines.get(pipeline_id)
+    pipeline_jobs = []
+
+    for bridge in root_pipeline.bridges.list(get_all=True):
+        downstream = bridge.attributes.get("downstream_pipeline")
+        if not bridge.name.startswith(job_prefix) or downstream is None:
             continue
 
-        if pipeline_bridge.name not in bridge:
-            bridge[pipeline_bridge.name] = []
+        child_pipeline_id = downstream["id"]
+        jobs = notification.get_jobs_from_pipeline(project, child_pipeline_id)
+        bridge_gpu = _bridge_gpu(bridge.name)
+        for job in jobs:
+            if job["gpu"] == "Unknown":
+                job["gpu"] = bridge_gpu
+        pipeline_jobs.append((bridge.name, child_pipeline_id, jobs))
+
+    return pipeline_jobs
 
-        for job in (
-            get_gitlab_handle()
-            .projects.get(PROJECT_ID)
-            .pipelines.get(pipeline_bridge.attributes['downstream_pipeline']['id'])
-            .jobs.list(get_all=True)
-        ):
-            bridge[pipeline_bridge.name].append(job)
-    return bridge
+
+def configure_notification_urls() -> None:
+    """Point nemo-ci-triage-2's notification links at Megatron-LM."""
+    notification.JOB_URL_TEMPLATE = f"{PROJECT_URL}/-/jobs/{{}}"
+    notification.PIPELINE_URL_TEMPLATE = f"{PROJECT_URL}/-/pipelines/{{}}"
 
 
 @click.command()
@@ -59,59 +76,24 @@ def get_jobs_per_bridge(pipeline_id: int, type_of_job: str):
     type=click.Choice(["unit-tests", "integration-tests", "functional-tests", "smoke-tests"]),
 )
 @click.option("--pipeline-context", required=True, type=str)
-@click.option("--pipeline-created-at", required=True, type=str)
-def main(pipeline_id: int, check_for: str, pipeline_context: str, pipeline_created_at: str):
-    if check_for == "unit-tests":
-        bridges = get_jobs_per_bridge(pipeline_id, "test:unit_tests")
-
-    if check_for == "integration-tests":
-        bridges = get_jobs_per_bridge(pipeline_id, "integration:run_")
-
-    if check_for == "functional-tests":
-        bridges = get_jobs_per_bridge(pipeline_id, "functional:run_")
+@click.option("--pipeline-created-at", required=True, type=str, expose_value=False)
+def main(pipeline_id: int, check_for: str, pipeline_context: str) -> None:
+    pipeline_jobs = get_pipeline_jobs(pipeline_id, JOB_PREFIXES[check_for])
 
     if check_for == "smoke-tests":
-        bridges = get_jobs_per_bridge(pipeline_id, "functional:smoke-")
-        if all(job.status == "success" for jobs in bridges.values() for job in jobs):
+        if all(job["status"] == "success" for _, _, jobs in pipeline_jobs for job in jobs):
             logger.info("All smoke tests passed, skipping Slack notification")
             return
 
-    pipeline_created_at_day = pd.Timestamp(pipeline_created_at).strftime("%Y-%m-%d")
-
-    messages = []
-
-    for bridge_name in bridges.keys():
-
-        total_num_jobs = len(bridges[bridge_name])
-        if all(job.status == "success" for job in bridges[bridge_name]):
-            messages.append(
-                f":doge3d: <https://{GITLAB_ENDPOINT}/ADLR/megatron-lm/-/pipelines/{pipeline_id}|Report - {pipeline_created_at_day} - {pipeline_context} - {bridge_name}>: All {total_num_jobs} passed."
-            )
-            continue
-
-        unsuccessful_jobs = [job for job in bridges[bridge_name] if job.status != "success"]
-        messages.append(
-            f":doctorge: <https://{GITLAB_ENDPOINT}/ADLR/megatron-lm/-/pipelines/{pipeline_id}|Report - {pipeline_created_at_day} - {pipeline_context} - {bridge_name}>: {len(unsuccessful_jobs)} of {total_num_jobs} failed."
-        )
-        if TAG_TEAM:
-            messages.append(
-                f"cc {TEAM_SLUG} <!subteam^S0A7B4U1T3P> <@U09TX0DHZ97>: Critical event, please react as soon as possible."
-            )
-
-        for job in unsuccessful_jobs:
-            messages.append(
-                f"\tJob: <https://{GITLAB_ENDPOINT}/ADLR/megatron-lm/-/jobs/{job.id}|{job.name}>"
-            )
-
-    messages.append("===============================================")
-
     if not WEBHOOK_URL:
         logger.info("No webhook URL configured, skipping Slack notification")
         return
 
-    for message in messages:
-        response = slack_sdk.webhook.WebhookClient(WEBHOOK_URL).send(text=message)
-        logger.info(response.status_code)
+    configure_notification_urls()
+    slack_mentions = f"{TEAM_SLUG} <!subteam^S0A7B4U1T3P> <@U09TX0DHZ97>" if TAG_TEAM else None
+    notification.send_slack_notification(
+        "megatron-lm", pipeline_context, pipeline_jobs, slack_mentions, webhook_url=WEBHOOK_URL
+    )
 
 
 if __name__ == "__main__":