diff --git a/.gitlab/scripts/build.sh b/.gitlab/scripts/build.sh index 15c926ed51f..c72bd38a581 100644 --- a/.gitlab/scripts/build.sh +++ b/.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) diff --git a/.gitlab/stages/02.test.yml b/.gitlab/stages/02.test.yml index a324ce037fb..d81e1be4857 100644 --- a/.gitlab/stages/02.test.yml +++ b/.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) diff --git a/.gitlab/stages/03.integration-tests.yml b/.gitlab/stages/03.integration-tests.yml index 70fa345e513..603c6e09b52 100644 --- a/.gitlab/stages/03.integration-tests.yml +++ b/.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) diff --git a/.gitlab/stages/04.functional-tests.yml b/.gitlab/stages/04.functional-tests.yml index 515aa3e7f7f..f83a2de4563 100644 --- a/.gitlab/stages/04.functional-tests.yml +++ b/.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) diff --git a/docker/Dockerfile.linting b/docker/Dockerfile.linting index bf27b768374..15cea7b6b57 100644 --- a/docker/Dockerfile.linting +++ b/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}" diff --git a/tests/test_utils/python_scripts/generate_jet_trigger_job.py b/tests/test_utils/python_scripts/generate_jet_trigger_job.py index aca23ad9fee..e3b471a629d 100644 --- a/tests/test_utils/python_scripts/generate_jet_trigger_job.py +++ b/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": { diff --git a/tests/test_utils/python_scripts/notify.py b/tests/test_utils/python_scripts/notify.py index 103badc6ce5..db875271d89 100644 --- a/tests/test_utils/python_scripts/notify.py +++ b/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: : All {total_num_jobs} passed." - ) - continue - - unsuccessful_jobs = [job for job in bridges[bridge_name] if job.status != "success"] - messages.append( - f":doctorge: : {len(unsuccessful_jobs)} of {total_num_jobs} failed." - ) - if TAG_TEAM: - messages.append( - f"cc {TEAM_SLUG} <@U09TX0DHZ97>: Critical event, please react as soon as possible." - ) - - for job in unsuccessful_jobs: - messages.append( - f"\tJob: " - ) - - 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} <@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__": diff --git a/tests/test_utils/test_ci_triage.py b/tests/test_utils/test_ci_triage.py new file mode 100644 index 00000000000..74ca1bc354a --- /dev/null +++ b/tests/test_utils/test_ci_triage.py @@ -0,0 +1,147 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import yaml +from click.testing import CliRunner + +from tests.test_utils.python_scripts import generate_jet_trigger_job, recipe_parser + + +@pytest.fixture +def notify_module(): + pytest.importorskip("nemo_ci_triage.slack_notification") + from tests.test_utils.python_scripts import notify + + return notify + + +def test_build_test_script_preserves_workload_exit_code(): + script = generate_jet_trigger_job.build_test_script("python workload.py") + + assert "set +e" in script + assert "set -o pipefail" in script + assert "python workload.py 2>&1 | tee jet_workload.log" in script + assert "exit_code=${PIPESTATUS[0]}" in script + assert "set -e" in script + assert "extract-errors jet_workload.log" in script + assert "--output error_report.json" in script + assert 'exit "$exit_code"' in script + + +@pytest.mark.parametrize("enable_error_extraction", [False, True]) +def test_error_extraction_is_opt_in_for_generated_jobs( + monkeypatch, tmp_path, enable_error_extraction +): + workload = recipe_parser.dotdict( + type="basic", + spec=recipe_parser.dotdict(model="gpt", environment="dev", test_case="triage-test"), + ) + monkeypatch.setattr( + generate_jet_trigger_job.recipe_parser, "load_workloads", lambda **_kwargs: [workload] + ) + output_path = tmp_path / "pipeline.yml" + args = [ + "--scope", + "mr", + "--environment", + "dev", + "--n-repeat", + "1", + "--time-limit", + "60", + "--test-cases", + "all", + "--platform", + "dgx_h100", + "--cluster", + "ghci", + "--output-path", + str(output_path), + "--container-image", + "utility", + "--container-tag", + "test", + "--dependent-job", + "functional:configure", + "--record-checkpoints", + "false", + "--slurm-account", + "mcore", + "--no-enable-warmup", + ] + if enable_error_extraction: + args.append("--enable-error-extraction") + + result = CliRunner().invoke(generate_jet_trigger_job.main, args) + + assert result.exit_code == 0, result.output + job = yaml.safe_load(output_path.read_text())["triage-test"] + if enable_error_extraction: + assert "extract-errors jet_workload.log" in job["script"][0] + assert job["artifacts"]["paths"] == ["results/", "jet_workload.log", "error_report.json"] + else: + assert "extract-errors" not in job["script"][0] + assert job["artifacts"]["paths"] == ["results/"] + + +def test_get_pipeline_jobs_uses_triage_collector(monkeypatch, notify_module): + notify = notify_module + bridge = SimpleNamespace( + name="functional:run_dev_dgx_h100", attributes={"downstream_pipeline": {"id": 101}} + ) + root_pipeline = Mock() + root_pipeline.bridges.list.return_value = [bridge] + project = Mock() + project.pipelines.get.return_value = root_pipeline + handle = Mock() + handle.projects.get.return_value = project + jobs = [{"status": "failed", "gpu": "Unknown"}] + + monkeypatch.setattr(notify, "get_gitlab_handle", lambda: handle) + collector = Mock(return_value=jobs) + monkeypatch.setattr(notify.notification, "get_jobs_from_pipeline", collector) + + assert notify.get_pipeline_jobs(123, "functional:run_") == [ + ("functional:run_dev_dgx_h100", 101, [{"status": "failed", "gpu": "H100"}]) + ] + collector.assert_called_once_with(project, 101) + + +def test_notification_delegates_to_triage_package(monkeypatch, notify_module): + notify = notify_module + pipeline_jobs = [("functional:run_dev_dgx_h100", 101, [{"status": "failed"}])] + sender = Mock() + + monkeypatch.setattr(notify, "WEBHOOK_URL", "https://slack.invalid/webhook") + monkeypatch.setattr(notify, "PROJECT_URL", "https://ci.example.com/ADLR/megatron-lm") + monkeypatch.setattr(notify, "get_pipeline_jobs", lambda *_args: pipeline_jobs) + monkeypatch.setattr(notify.notification, "send_slack_notification", sender) + + result = CliRunner().invoke( + notify.main, + [ + "--pipeline-id", + "123", + "--check-for", + "functional-tests", + "--pipeline-context", + "mr", + "--pipeline-created-at", + "2026-07-12T00:00:00Z", + ], + ) + + assert result.exit_code == 0, result.output + assert ( + notify.notification.JOB_URL_TEMPLATE == "https://ci.example.com/ADLR/megatron-lm/-/jobs/{}" + ) + assert ( + notify.notification.PIPELINE_URL_TEMPLATE + == "https://ci.example.com/ADLR/megatron-lm/-/pipelines/{}" + ) + sender.assert_called_once_with( + "megatron-lm", "mr", pipeline_jobs, None, webhook_url="https://slack.invalid/webhook" + )