diff --git a/.github/workflows/therock-rccl-ci-linux.yml b/.github/workflows/therock-rccl-ci-linux.yml index 635e0f11944..4faebf311f9 100644 --- a/.github/workflows/therock-rccl-ci-linux.yml +++ b/.github/workflows/therock-rccl-ci-linux.yml @@ -28,6 +28,10 @@ on: notify_webhook: type: string default: '' + madengine_nodes: + description: 'Number of SLURM nodes for MADEngine workloads' + type: string + default: '2' permissions: contents: read @@ -208,3 +212,23 @@ jobs: artifact_run_id: ${{ inputs.artifact_run_id || github.run_id }} notify_email: ${{ inputs.notify_email }} notify_webhook: ${{ inputs.notify_webhook }} + + therock-test-madengine: + name: "Test MADEngine workloads (scheduled)" + if: ${{ !cancelled() && inputs.amdgpu_families == 'gfx950-dcgpu' && (inputs.event_name == 'schedule' || inputs.event_name == 'workflow_dispatch') }} + needs: [therock-build-linux] + permissions: + contents: read + id-token: write + packages: write + uses: ./.github/workflows/therock-rccl-test-madengine.yml + secrets: inherit + with: + amdgpu_families: ${{ inputs.amdgpu_families }} + artifact_group: ${{ inputs.artifact_group }} + test_runs_on: ruby-linux-slurm-scale-runner + artifact_run_id: ${{ inputs.artifact_run_id || github.run_id }} + workload: llama-3.1-70b-training + nodes: ${{ inputs.madengine_nodes || '2' }} + notify_email: ${{ inputs.notify_email }} + teams_webhook: ${{ inputs.notify_webhook }} diff --git a/.github/workflows/therock-rccl-ci.yml b/.github/workflows/therock-rccl-ci.yml index ed6a34f6697..179419b3433 100644 --- a/.github/workflows/therock-rccl-ci.yml +++ b/.github/workflows/therock-rccl-ci.yml @@ -2,7 +2,7 @@ name: TheRock CI for rccl on: schedule: - - cron: '17 6 * * 0' # Weekly, Sunday 06:17 UTC + - cron: '17 6 * * *' # Nightly, 06:17 UTC workflow_dispatch: inputs: artifact_run_id: @@ -21,6 +21,10 @@ on: description: 'Teams webhook URL for test summary notifications' type: string default: '' + madengine_nodes: + description: 'Number of SLURM nodes for MADEngine workloads' + type: string + default: '2' permissions: contents: read @@ -89,6 +93,7 @@ jobs: test_scope: ${{ inputs.test_scope || 'smoke' }} notify_email: ${{ inputs.notify_email || 'collectives-ci@amd.com' }} notify_webhook: ${{ inputs.notify_webhook }} + madengine_nodes: ${{ inputs.madengine_nodes || '2' }} cmake_options: > -DTHEROCK_ENABLE_ALL=OFF -DTHEROCK_BUILD_TESTING=ON diff --git a/.github/workflows/therock-rccl-test-madengine.yml b/.github/workflows/therock-rccl-test-madengine.yml new file mode 100644 index 00000000000..9b158bc79aa --- /dev/null +++ b/.github/workflows/therock-rccl-test-madengine.yml @@ -0,0 +1,187 @@ +name: TheRock Test MADEngine Workloads for rccl + +on: + workflow_call: + inputs: + amdgpu_families: + type: string + artifact_group: + type: string + test_runs_on: + type: string + artifact_run_id: + type: string + workload: + type: string + default: 'llama-3.1-70b-training' + nodes: + type: string + default: '2' + notify_email: + type: string + default: '' + teams_webhook: + type: string + default: '' + workflow_dispatch: + inputs: + amdgpu_families: + type: string + artifact_group: + type: string + test_runs_on: + type: string + default: 'ruby-linux-slurm-scale-runner' + artifact_run_id: + description: 'Reuse artifacts from a previous build run ID' + type: string + workload: + description: 'MADEngine workload to run' + type: string + default: 'llama-3.1-70b-training' + nodes: + description: 'Number of SLURM nodes to allocate' + type: string + default: '2' + notify_email: + description: 'Email address to send test summary report' + type: string + default: '' + teams_webhook: + description: 'Teams webhook URL for notifications' + type: string + default: '' + +permissions: + contents: read + packages: write + +jobs: + test_madengine: + name: 'Test MADEngine ${{ inputs.workload }} (${{ inputs.nodes }}N)' + runs-on: ${{ inputs.test_runs_on }} + defaults: + run: + shell: bash + env: + ARTIFACT_RUN_ID: "${{ inputs.artifact_run_id }}" + RESULTS_DIR: "/apps/rccl-ci/perf" + WORK_DIR: "/apps/rccl-ci/workdir/${{ github.run_id }}" + steps: + - name: Checkout TheRock repository + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + repository: "ROCm/TheRock" + ref: 2e7c190b99caa226a8644eda6eca720b3db102d5 # 2026-07-22 + + - name: Checkout rocm-systems scripts + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + sparse-checkout: projects/rccl/ci/scripts + path: rocm-systems + + - name: Set arch-specific paths + run: | + case "${{ inputs.amdgpu_families }}" in + gfx950-dcgpu) + echo "OUTPUT_ARTIFACTS_DIR=/apps/cvs_tests/dist_new/dist/rocm" >> $GITHUB_ENV + echo "AWS_SHARED_CREDENTIALS_FILE=/apps/cvs_tests/awsconfig/credentials.ini" >> $GITHUB_ENV + echo "CLUSTER=ruby" >> $GITHUB_ENV + # Use /apps NFS for workdir — /home has root_squash which blocks Docker container writes + echo "WORK_DIR=/apps/rccl-ci/workdir/${{ github.run_id }}" >> $GITHUB_ENV + ;; + *) + echo "Unsupported arch for MADEngine: ${{ inputs.amdgpu_families }}" + exit 1 + ;; + esac + + - name: Resolve TheRock artifact run + if: ${{ inputs.artifact_run_id == '' }} + env: + GH_TOKEN: ${{ github.token }} + GITHUB_TOKEN: ${{ github.token }} + run: | + pip3 install -r requirements.txt + run_id="$(python3 -c " + import sys, contextlib + sys.path.insert(0, 'build_tools') + from find_latest_artifacts import find_latest_artifacts + with contextlib.redirect_stdout(sys.stderr): + results = find_latest_artifacts( + artifact_groups=['${{ inputs.artifact_group }}'], + verbose=True, + ) + if not results: + print('No artifacts found for ${{ inputs.artifact_group }}', file=sys.stderr) + sys.exit(1) + print(results[0].workflow_run_id) + ")" + echo "Resolved TheRock run_id=${run_id}" + echo "ARTIFACT_RUN_ID=${run_id}" >> $GITHUB_ENV + + - name: Fetch RCCL artifacts + run: | + pip3 install -r requirements.txt + python3 build_tools/install_rocm_from_artifacts.py \ + --run-id="${ARTIFACT_RUN_ID}" \ + --artifact-group="${{ inputs.artifact_group }}" \ + --output-dir="${OUTPUT_ARTIFACTS_DIR}" \ + --rccl + + - name: Setup Python environment + run: | + mkdir -p "${WORK_DIR}" + python3 -m venv "${WORK_DIR}/venv" + source "${WORK_DIR}/venv/bin/activate" + pip install --upgrade pip + + - name: Login to container registry + run: | + echo "${{ github.token }}" | docker login ghcr.io -u "${{ github.actor }}" --password-stdin + + - name: Run MADEngine workload + timeout-minutes: 210 + env: + NCCL_DEBUG: WARN + MAD_SECRETS_HFTOKEN: ${{ secrets.HF_TOKEN || '' }} + run: | + source "${WORK_DIR}/venv/bin/activate" + export PYTHONPATH="rocm-systems/projects/rccl/ci/scripts:${PYTHONPATH}" + + NOTIFY_ARGS="" + if [ -n "${{ inputs.notify_email }}" ]; then + NOTIFY_ARGS="$NOTIFY_ARGS --notify-email ${{ inputs.notify_email }}" + fi + if [ -n "${{ inputs.teams_webhook }}" ]; then + NOTIFY_ARGS="$NOTIFY_ARGS --teams-webhook ${{ inputs.teams_webhook }}" + fi + + python3 rocm-systems/projects/rccl/ci/scripts/test_madengine.py \ + --artifact-dir "${OUTPUT_ARTIFACTS_DIR}" \ + --workload "${{ inputs.workload }}" \ + --cluster "${CLUSTER}" \ + --nodes "${{ inputs.nodes }}" \ + --results-dir "${RESULTS_DIR}" \ + --work-dir "${WORK_DIR}" \ + --registry "ghcr.io/rocm/rocm-systems" \ + $NOTIFY_ARGS + + - name: Upload results + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: madengine-results-${{ inputs.workload }}-${{ inputs.nodes }}N-${{ inputs.amdgpu_families }} + path: | + ${{ env.WORK_DIR }}/perf.csv + ${{ env.WORK_DIR }}/perf_entry_super.csv + ${{ env.WORK_DIR }}/perf_entry_super.json + ${{ env.WORK_DIR }}/perf_super.csv + ${{ env.WORK_DIR }}/madengine_summary.txt + ${{ env.WORK_DIR }}/manifest.json + retention-days: 30 + + - name: Cleanup work directory + if: always() + run: | + rm -rf "${WORK_DIR}/madengine" "${WORK_DIR}/MAD" "${WORK_DIR}/venv" "${WORK_DIR}/rccl_libs" diff --git a/projects/rccl/ci/scripts/rccl_ci_utils.py b/projects/rccl/ci/scripts/rccl_ci_utils.py new file mode 100644 index 00000000000..f9e7f5bc6d2 --- /dev/null +++ b/projects/rccl/ci/scripts/rccl_ci_utils.py @@ -0,0 +1,186 @@ +"""Shared utilities for RCCL CI test scripts (JAX, PyTorch, MADEngine).""" + +import json +import logging +import os +import smtplib +import sys +import urllib.request +import xml.etree.ElementTree as ET +from email.mime.text import MIMEText +from pathlib import Path + +log = logging.getLogger(__name__) + +SMTP_SERVERS = ["smtp.amd.com", "aussmtp.amd.com", "mail.amd.com", "localhost"] + + +def find_rccl_library(artifact_dir: Path) -> Path: + """Find librccl.so in the artifact directory tree.""" + matches = list(artifact_dir.rglob("librccl.so")) + if not matches: + so_files = list(artifact_dir.rglob("*.so"))[:20] + log.error("librccl.so not found in %s", artifact_dir) + log.error("Shared libraries found: %s", [str(f) for f in so_files]) + sys.exit(1) + lib_path = matches[0].resolve() + log.info("Found librccl.so at: %s", lib_path) + return lib_path + + +def verify_rccl_override(rccl_lib_dir: Path) -> None: + """Verify that the CI-built librccl.so exists on disk.""" + ci_rccl = rccl_lib_dir.resolve() / "librccl.so" + if not ci_rccl.exists(): + log.error("CI-built librccl.so not found at %s", ci_rccl) + sys.exit(1) + log.info("CI-built RCCL: %s (%d bytes)", ci_rccl, ci_rccl.stat().st_size) + + +def parse_junit_xml(xml_path: Path) -> dict: + """Parse JUnit XML and return structured results.""" + tree = ET.parse(xml_path) + root = tree.getroot() + + passed_tests = [] + failed_tests = [] + error_details = [] + tests_run = 0 + failures = 0 + errors = 0 + + for suite in root.iter("testsuite"): + tests_run = int(suite.get("tests", 0)) + failures = int(suite.get("failures", 0)) + errors = int(suite.get("errors", 0)) + + for tc in root.iter("testcase"): + name = tc.get("name", "") + time_s = tc.get("time", "") + duration = f"{float(time_s):.2f}s" if time_s else "" + + failure = tc.find("failure") + error = tc.find("error") + if failure is not None: + failed_tests.append(name) + error_details.append( + f"FAILED: {name}\n {failure.get('message', '')}" + ) + elif error is not None: + failed_tests.append(name) + error_details.append( + f"ERROR: {name}\n {error.get('message', '')}" + ) + else: + passed_tests.append((name, duration)) + + return { + "passed": passed_tests, + "failed": failed_tests, + "error_details": error_details, + "tests_run": tests_run, + "failures": failures, + "errors": errors, + } + + +def write_github_summary(report: str) -> None: + """Write report to GITHUB_STEP_SUMMARY if available.""" + summary_file = os.environ.get("GITHUB_STEP_SUMMARY") + if summary_file: + with open(summary_file, "a") as f: + f.write("```\n") + f.write(report) + f.write("\n```\n") + log.info("Summary written to GITHUB_STEP_SUMMARY") + + +def set_github_output(key: str, value: str) -> None: + """Write a key=value pair to GITHUB_OUTPUT if available.""" + output_file = os.environ.get("GITHUB_OUTPUT") + if output_file: + with open(output_file, "a") as f: + f.write(f"{key}={value}\n") + + +def send_email_report( + report: str, recipient: str, status: str, subject_prefix: str +) -> None: + """Send the summary report via email.""" + subject = f"{subject_prefix}: {status}" + msg = MIMEText(report) + msg["Subject"] = subject + msg["From"] = "rccl-ci@amd.com" + msg["To"] = recipient + + for server in SMTP_SERVERS: + try: + with smtplib.SMTP(server, timeout=10) as s: + s.sendmail(msg["From"], [recipient], msg.as_string()) + log.info("Email sent to %s via %s", recipient, server) + return + except Exception as e: + log.debug("SMTP %s failed: %s", server, e) + continue + log.warning( + "Could not send email to %s (tried: %s)", recipient, ", ".join(SMTP_SERVERS) + ) + + +def send_teams_webhook( + report: str, webhook_url: str, status: str, subject_prefix: str +) -> None: + """Send the summary report to a Microsoft Teams channel via webhook.""" + color = "Good" if status == "PASSED" else "Attention" + run_url = os.environ.get("GITHUB_SERVER_URL", "") + repo = os.environ.get("GITHUB_REPOSITORY", "") + run_id = os.environ.get("GITHUB_RUN_ID", "") + actions_url = f"{run_url}/{repo}/actions/runs/{run_id}" if run_url else "" + + facts = [{"title": "Status", "value": status}] + if actions_url: + facts.append({"title": "Run", "value": f"[View]({actions_url})"}) + + body = [ + { + "type": "TextBlock", + "text": f"{subject_prefix}: {status}", + "weight": "Bolder", + "size": "Medium", + "color": color, + }, + {"type": "FactSet", "facts": facts}, + { + "type": "TextBlock", + "text": report, + "wrap": True, + "fontType": "Monospace", + "size": "Small", + }, + ] + + payload = { + "type": "message", + "attachments": [ + { + "contentType": "application/vnd.microsoft.card.adaptive", + "content": { + "$schema": "http://adaptivecards.io/schemas/adaptive-card.json", + "type": "AdaptiveCard", + "version": "1.4", + "body": body, + }, + } + ], + } + + try: + req = urllib.request.Request( + webhook_url, + data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(req, timeout=15) as resp: + log.info("Teams webhook sent (HTTP %d)", resp.status) + except Exception as e: + log.warning("Failed to send Teams webhook: %s", e) diff --git a/projects/rccl/ci/scripts/test_jax_collective.py b/projects/rccl/ci/scripts/test_jax_collective.py new file mode 100644 index 00000000000..cdade4c6951 --- /dev/null +++ b/projects/rccl/ci/scripts/test_jax_collective.py @@ -0,0 +1,412 @@ +#!/usr/bin/env python3 +"""Run JAX collective smoke tests against CI-built RCCL. + +This script handles: + 1. Discovering the CI-built librccl.so in the artifact directory + 2. Prepending its directory to LD_LIBRARY_PATH so JAX loads it + 3. Cloning the matching JAX test sources (sparse checkout) + 4. Running pytest on pmap_test.py and shard_map_test.py + +Usage from GitHub Actions: + python .github/scripts/test_jax_collective.py \ + --artifact-dir ./build \ + --jax-src ./jax-src \ + --results-log ./jax_collective_results.log +""" + +import argparse +import logging +import os +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path + +from rccl_ci_utils import ( + find_rccl_library, + parse_junit_xml, + send_email_report, + set_github_output, + verify_rccl_override, + write_github_summary, +) + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +log = logging.getLogger(__name__) + +JAX_REPO = "https://github.com/ROCm/jax.git" + +SMOKE_TEST_FILES = [ + "tests/pmap_test.py", + "tests/shard_map_test.py", +] + +SMOKE_TEST_KEYWORDS = [ + "PythonPmapTest and testBasic", + "PythonPmapTest and testGather and not testGatherBool and not testGatherNeg and not testGatherTiled and not testGatherReplica", + "PythonPmapTest and testReduceScatter and not Tiled and not Replica", + "PythonPmapTest and testCollectivePermute and not Grad and not Cyclic", + "PythonPmapTest and testAllToAll and not Replica and not Vmap and not Grad", + "ShardMapTest and test_all_gather and not invariant and not axis_index", + "ShardMapTest and test_matmul_reduce_scatter", + "ShardMapTest and test_collective_permute and not multiple", + "ShardMapTest and test_axis_index and not basic and not twoaxes and not eager", + "ShardMapTest and test_all_to_all and not axis_index and not grad", +] + +XLA_ENV = { + "XLA_PYTHON_CLIENT_ALLOCATOR": "default", + "XLA_PYTHON_CLIENT_PREALLOCATE": "false", + "XLA_FLAGS": ( + "--xla_gpu_force_compilation_parallelism=1 " + "--xla_gpu_enable_nccl_comm_splitting=false " + # Empty value disables all command buffer types (HIP graphs) on ROCm. + "--xla_gpu_enable_command_buffer= " + "--xla_gpu_enable_cublaslt=false" + ), +} + + +def find_lib_dirs(artifact_dir: Path) -> list[Path]: + """Find all directories containing .so files in the artifact tree.""" + lib_dirs: set[Path] = set() + for so_file in artifact_dir.rglob("*.so"): + lib_dirs.add(so_file.parent.resolve()) + for so_file in artifact_dir.rglob("*.so.*"): + lib_dirs.add(so_file.parent.resolve()) + sorted_dirs = sorted(lib_dirs) + for d in sorted_dirs: + count = sum(1 for f in d.iterdir() if ".so" in f.name) + log.info("Found lib dir: %s (%d libs)", d, count) + return sorted_dirs + + +def populate_rocm_lib_dir(lib_dirs: list[Path]) -> None: + """Populate /opt/rocm/lib with symlinks to artifact libraries. + + JAX's xla_rocm_plugin.so has RUNPATH including /opt/rocm/lib. With ELF + RUNPATH semantics, transitive dependencies are resolved via RUNPATH + rather than LD_LIBRARY_PATH. In a container with no system ROCm, we + populate /opt/rocm/lib so the loader can find them. + """ + rocm_lib = Path("/opt/rocm/lib") + try: + rocm_lib.mkdir(parents=True, exist_ok=True) + except PermissionError: + log.warning("Cannot create %s — not running as root", rocm_lib) + return + + count = 0 + for d in lib_dirs: + for so_file in d.iterdir(): + if ".so" not in so_file.name: + continue + target = rocm_lib / so_file.name + if target.exists() or target.is_symlink(): + continue + target.symlink_to(so_file.resolve()) + count += 1 + log.info("Created %d symlinks in %s", count, rocm_lib) + + +def setup_ld_library_path(lib_dirs: list[Path]) -> str: + """Prepend all artifact lib dirs to LD_LIBRARY_PATH.""" + parts = [str(d) for d in lib_dirs] + rocm_lib = Path("/opt/rocm/lib") + if rocm_lib.is_dir(): + parts.insert(0, str(rocm_lib)) + existing = os.environ.get("LD_LIBRARY_PATH", "") + if existing: + parts.append(existing) + new_path = ":".join(parts) + os.environ["LD_LIBRARY_PATH"] = new_path + log.info("LD_LIBRARY_PATH=%s", new_path) + return new_path + + +def setup_xla_environment() -> None: + """Set XLA environment variables required for JAX on ROCm.""" + for key, value in XLA_ENV.items(): + os.environ[key] = value + log.info("Set %s=%s", key, value) + + +def clone_jax_test_sources(jax_src: Path) -> None: + """Sparse-clone ROCm/jax to get test sources matching the installed version. + + Tries to find a tag matching the installed JAX version (e.g. jax-v0.5.3). + Falls back to the default branch HEAD for nightly/dev builds. + """ + if jax_src.exists() and (jax_src / "tests" / "pmap_test.py").exists(): + log.info("JAX test sources already present at %s, skipping clone", jax_src) + return + + import jax + jax_version = jax.__version__ + base_version = jax_version.split(".dev")[0].split("+")[0] + git_ref = f"jax-v{base_version}" + log.info("JAX version: %s, trying tag: %s", jax_version, git_ref) + + result = subprocess.run( + ["git", "ls-remote", "--tags", JAX_REPO, git_ref], + capture_output=True, text=True, + ) + if not result.stdout.strip(): + log.info("Tag %s not found, using default branch HEAD", git_ref) + git_ref = None + + clone_cmd = [ + "git", "clone", + "--depth=1", + "--filter=blob:none", + "--sparse", + JAX_REPO, + str(jax_src), + ] + if git_ref: + clone_cmd.insert(-1, f"--branch={git_ref}") + + log.info("Cloning ROCm/jax (ref=%s, sparse) into %s", git_ref or "HEAD", jax_src) + subprocess.run(clone_cmd, check=True) + subprocess.run( + ["git", "sparse-checkout", "set", "tests/", "build/", "jax/"], + cwd=jax_src, + check=True, + ) + + test_file = jax_src / "tests" / "pmap_test.py" + if not test_file.exists(): + log.error("pmap_test.py not found after clone") + sys.exit(1) + log.info("JAX test sources ready: %s", jax_src) + + +def print_environment_info() -> None: + """Print GPU and environment details for CI logs.""" + log.info("--- Environment Info ---") + try: + import jax + log.info("JAX version: %s", jax.__version__) + devices = jax.devices() + log.info("Devices: %s", devices) + log.info("Device count: %d", jax.device_count()) + log.info("Local device count: %d", jax.local_device_count()) + gpu_devices = [d for d in devices if d.platform != "cpu"] + if not gpu_devices: + log.error("No GPU devices found — JAX fell back to CPU only") + log.error("Check that ROCm libraries are on LD_LIBRARY_PATH") + sys.exit(1) + except Exception as e: + log.error("Failed to query JAX devices: %s", e) + sys.exit(1) + + log.info("LD_LIBRARY_PATH: %s", os.environ.get("LD_LIBRARY_PATH", "")) + for key in sorted(XLA_ENV): + log.info("%s=%s", key, os.environ.get(key, "")) + log.info("--- End Environment Info ---") + + +def run_tests(jax_src: Path, results_log: Path) -> tuple[int, dict]: + """Run pytest on the 10 collective smoke tests and return (exit_code, summary).""" + junit_xml = results_log.parent / "jax_collective_results.xml" + + k_expr = " or ".join(f"({kw})" for kw in SMOKE_TEST_KEYWORDS) + cmd = [ + sys.executable, "-m", "pytest", + "-sv", + "--timeout=120", + "--tb=short", + f"--junitxml={junit_xml}", + "-k", k_expr, + ] + SMOKE_TEST_FILES + + log.info("Running: %s", " ".join(cmd)) + + with open(results_log, "w") as log_file: + proc = subprocess.Popen( + cmd, + cwd=jax_src, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + for line in proc.stdout: + sys.stdout.write(line) + sys.stdout.flush() + log_file.write(line) + proc.wait() + + log.info("Test exit code: %d", proc.returncode) + log.info("Results written to: %s", results_log) + + exit_code = proc.returncode + passed_tests = [] + failed_tests = [] + tests_run = 0 + summary_line = "" + + if junit_xml.exists(): + log.info("Parsing JUnit XML: %s", junit_xml) + junit = parse_junit_xml(junit_xml) + passed_tests = junit["passed"] + failed_tests = junit["failed"] + tests_run = junit["tests_run"] + parts = [] + if passed_tests: + parts.append(f"{len(passed_tests)} passed") + if failed_tests: + parts.append(f"{len(failed_tests)} failed") + summary_line = ", ".join(parts) + + if junit["error_details"]: + log.info("Failure/error details from JUnit XML:") + for detail in junit["error_details"]: + log.info(" %s", detail) + else: + log.warning("JUnit XML not found at %s, falling back to exit code only", junit_xml) + + if tests_run < len(SMOKE_TEST_KEYWORDS): + log.error( + "Expected at least %d smoke tests but only %d were collected — " + "tests may have been skipped or deselected", + len(SMOKE_TEST_KEYWORDS), + tests_run, + ) + exit_code = 1 + + summary = { + "exit_code": exit_code, + "passed": passed_tests, + "failed": failed_tests, + "summary_line": summary_line, + "tests_run": tests_run, + "expected_tests": len(SMOKE_TEST_KEYWORDS), + } + return exit_code, summary + + +def generate_summary_report(summary: dict, rccl_lib: Path) -> str: + """Generate a plain-text summary report.""" + import jax + + status = "PASSED" if summary["exit_code"] == 0 else "FAILED" + devices = jax.devices() + gpu_name = str(devices[0]) if devices else "unknown" + + lines = [ + "RCCL JAX Collective Test Report", + "=" * 40, + f"Status: {status}", + f"Date: {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S UTC')}", + "", + f"JAX: {jax.__version__}", + f"RCCL: {rccl_lib}", + f"GPUs: {len(devices)}x {gpu_name}", + "", + f"Results: {summary['summary_line']}", + ] + + if summary.get("expected_tests") is not None: + lines.append(f"Collected: {summary['tests_run']}/{summary['expected_tests']} expected smoke tests") + lines.append("") + + if summary["failed"]: + lines.append(f"FAILED tests ({len(summary['failed'])}):") + for name in summary["failed"]: + lines.append(f" FAIL {name}") + lines.append("") + + if summary["passed"]: + lines.append(f"PASSED tests ({len(summary['passed'])}):") + for name, duration in summary["passed"]: + lines.append(f" OK {name:60s} {duration}") + lines.append("") + + run_url = os.environ.get("GITHUB_SERVER_URL", "") + repo = os.environ.get("GITHUB_REPOSITORY", "") + run_id = os.environ.get("GITHUB_RUN_ID", "") + if run_url and repo and run_id: + lines.append(f"CI run: {run_url}/{repo}/actions/runs/{run_id}") + + return "\n".join(lines) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--artifact-dir", + type=Path, + required=True, + help="Directory containing CI-built artifacts", + ) + parser.add_argument( + "--jax-src", + type=Path, + required=True, + help="Directory to clone JAX test sources into", + ) + parser.add_argument( + "--results-log", + type=Path, + default=Path("jax_collective_results.log"), + help="Path for test results log file", + ) + parser.add_argument( + "--notify-email", + type=str, + default="", + help="Send summary report to this email address", + ) + parser.add_argument( + "--discover-only", + action="store_true", + help="Only discover library paths and set GITHUB_OUTPUT, then exit", + ) + + args = parser.parse_args() + + # Step 1: Discover RCCL library and all lib dirs in artifacts + rccl_lib = find_rccl_library(args.artifact_dir) + rccl_lib_dir = rccl_lib.parent + lib_dirs = find_lib_dirs(args.artifact_dir) + + set_github_output("RCCL_LIB_DIR", str(rccl_lib_dir)) + + if args.discover_only: + return + + # Step 2: Set up library paths and verify override + populate_rocm_lib_dir(lib_dirs) + setup_ld_library_path(lib_dirs) + verify_rccl_override(rccl_lib_dir) + + # Step 3: Set XLA environment variables + setup_xla_environment() + + # Step 4: Clone JAX test sources + clone_jax_test_sources(args.jax_src) + + # Step 5: Print environment info and run tests + print_environment_info() + exit_code, summary = run_tests(args.jax_src, args.results_log) + + # Step 6: Generate and distribute summary report + report = generate_summary_report(summary, rccl_lib) + log.info("\n%s", report) + write_github_summary(report) + + summary_path = args.results_log.parent / "jax_collective_summary.txt" + summary_path.write_text(report) + log.info("Summary written to: %s", summary_path) + + if args.notify_email: + status = "PASSED" if exit_code == 0 else "FAILED" + send_email_report(report, args.notify_email, status, + subject_prefix="RCCL JAX Collective Test") + + sys.exit(exit_code) + + +if __name__ == "__main__": + main() diff --git a/projects/rccl/ci/scripts/test_madengine.py b/projects/rccl/ci/scripts/test_madengine.py new file mode 100644 index 00000000000..15b766a7ce1 --- /dev/null +++ b/projects/rccl/ci/scripts/test_madengine.py @@ -0,0 +1,1274 @@ +#!/usr/bin/env python3 +"""Run MADEngine AI workloads against CI-built RCCL and track performance. + +This script handles: + 1. Installing madengine from source into the CI venv + 2. Building a Docker overlay image with the CI-built RCCL + 3. Generating a manifest.json for the requested workload + 4. Running the workload via `madengine run` + 5. Parsing perf.csv results and checking for regressions + 6. Appending results to a JSONL datastore for trend analysis + +Usage from GitHub Actions (on ruby-linux-slurm-scale-runner): + python projects/rccl/ci/scripts/test_madengine.py \ + --artifact-dir /apps/cvs_tests/dist_new/dist/rocm \ + --workload llama-3.1-70b-training \ + --cluster ruby \ + --nodes 2 \ + --results-dir /apps/rccl-ci/perf +""" +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import logging +import os +import re +import shutil +import subprocess +import sys +import tempfile +from datetime import datetime, timezone +from pathlib import Path + +from rccl_ci_utils import ( + find_rccl_library, + send_email_report, + send_teams_webhook, + set_github_output, + write_github_summary, +) + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +log = logging.getLogger(__name__) + +PERF_DATASTORE = "madengine_results.jsonl" +REGRESSION_WINDOW = 5 +REGRESSION_THRESHOLD_TRAINING = 0.02 # 2% +REGRESSION_THRESHOLD_INFERENCE = 0.05 # 5% + +MADENGINE_REPO = "https://github.com/ROCm/madengine.git" +MADENGINE_REF = "ec4de0b58c49f05d89dd33e38cc3e81e0fb3d992" +MAD_REPO = "https://github.com/ROCm/MAD.git" +MAD_REF = "688828bd9d4ad5be4196bc6161cabb6d7048ee65" +MAD_BRANCH = "mad-rccl" + +WORKLOAD_CONFIGS = { + "llama-3.1-70b-training": { + "type": "training", + "model_repo": "primus_pyt_megatron_lm_train_llama-3.1-70b", + "model_repo_aliases": [ + "primus_pyt_megatron_lm_train_llama-3.1-70b_overlay", + "primus_pyt_megatron_lm_train_llama-3.1-70b_scaleout", + ], + "base_image": "rocm/primus:v26.4", + "gpu_target": "gfx950", + "metric_key": "tokens_per_second_per_gpu", + "multiple_results": "perf_primus-megatron-Llama-3.1-70B.csv", + "reference_values": { + "2N": 1685, + "4N": 1600, + "16N": 1685, + "32N": 1485, + "44N": 1432, + }, + "slurm_partition": "meta64", + "gpus_per_node": 8, + "time_limit": "03:00:00", + "docker_mounts": {"/dev/infiniband": "/dev/infiniband"}, + "docker_run_options": "--privileged --group-add render --shm-size 64G " + "--device=/dev/infiniband --cap-add IPC_LOCK " + "--ulimit memlock=-1 -v /sys:/sys:ro -v /run/udev:/run/udev:ro", + }, + "llama-4-scout-training": { + "type": "training", + "model_repo": "primus_pyt_megatron_lm_train_llama-4-scout-17b-16e", + "model_repo_aliases": [ + "primus_pyt_megatron_lm_train_llama-4-scout-17b-16e_overlay", + "primus_pyt_megatron_lm_train_llama-4-scout-17b-16e_scaleout", + ], + "base_image": "rocm/primus:v26.4", + "gpu_target": "gfx950", + "metric_key": "tokens_per_second_per_gpu", + "multiple_results": "perf_primus-megatron-Llama-3.1-70B.csv", + "reference_values": { + "2N": 2734, + "4N": 2337, + }, + "slurm_partition": "meta64", + "gpus_per_node": 8, + "time_limit": "02:00:00", + "docker_mounts": {"/dev/infiniband": "/dev/infiniband"}, + "docker_run_options": "--privileged --group-add render --shm-size 64G " + "--device=/dev/infiniband --cap-add IPC_LOCK " + "--ulimit memlock=-1 -v /sys:/sys:ro -v /run/udev:/run/udev:ro", + }, +} + +CLUSTER_CONFIGS = { + "ruby": { + "gpu_target": "gfx950", + "slurm_partition": "meta64", + "slurm_qos": "vip_prio", + "slurm_no_gres": True, # Ruby SLURM has no GPU GRES configured + "mount_host_ib_libs": True, + "nccl_env": { + "NCCL_NET": "IB", + "NCCL_IB_DISABLE": "0", + "NCCL_IB_HCA": "bnxt_re0:1,bnxt_re1:1,bnxt_re2:1,bnxt_re3:1,bnxt_re4:1,bnxt_re5:1,bnxt_re6:1,bnxt_re7:1", + "NCCL_IB_GID_INDEX": "3", + "NCCL_IB_TC": "104", + "NCCL_IB_QPS_PER_CONNECTION": "4", + "NCCL_SOCKET_IFNAME": "fenic0,enp49s0f0np0", + "NCCL_DEBUG": "WARN", + }, + "results_base": "/apps/rccl-ci/perf", + }, +} + + +def install_madengine(work_dir: Path) -> Path: + """Clone and install madengine into the current Python environment.""" + madengine_dir = work_dir / "madengine" + mad_dir = work_dir / "MAD" + + if not madengine_dir.exists(): + log.info("Cloning madengine at %s...", MADENGINE_REF[:12]) + subprocess.run( + ["git", "clone", "--depth=1", MADENGINE_REPO, str(madengine_dir)], + check=True, + ) + subprocess.run( + ["git", "-C", str(madengine_dir), "fetch", "--depth=1", "origin", MADENGINE_REF], + check=True, + ) + subprocess.run( + ["git", "-C", str(madengine_dir), "checkout", MADENGINE_REF], + check=True, + ) + + if not mad_dir.exists(): + log.info("Cloning MAD (%s) at %s...", MAD_BRANCH, MAD_REF[:12]) + subprocess.run( + ["git", "clone", "--depth=1", "--branch", MAD_BRANCH, MAD_REPO, str(mad_dir)], + check=True, + ) + # Verify the clone landed on the expected SHA. The --branch clone + # gives us the branch tip; fetch+checkout is only needed if the + # pinned ref differs from the tip (e.g. after the branch moves). + head = subprocess.run( + ["git", "-C", str(mad_dir), "rev-parse", "HEAD"], + capture_output=True, text=True, check=True, + ).stdout.strip() + if not head.startswith(MAD_REF[:12]): + log.info("MAD HEAD %s != pinned %s, fetching...", head[:12], MAD_REF[:12]) + subprocess.run( + ["git", "-C", str(mad_dir), "fetch", "--depth=1", "origin", MAD_REF], + check=True, + ) + subprocess.run( + ["git", "-C", str(mad_dir), "checkout", MAD_REF], + check=True, + ) + else: + log.info("MAD HEAD matches pinned ref: %s", head[:12]) + + log.info("Installing madengine...") + subprocess.run( + [sys.executable, "-m", "pip", "install", "-e", str(madengine_dir)], + check=True, + ) + + log.info("madengine installed to: %s", madengine_dir) + + scripts_src = mad_dir / "scripts" / "primus_megatron-lm" + if not scripts_src.is_dir(): + scripts_src = mad_dir / "scripts" / "primus" / "megatron-lm" + scripts_dst = work_dir / "scripts" / "primus_megatron-lm" + if scripts_src.is_dir() and not scripts_dst.exists(): + scripts_dst.parent.mkdir(parents=True, exist_ok=True) + subprocess.run(["cp", "-r", str(scripts_src), str(scripts_dst)], check=True) + log.info("Copied MAD primus scripts to %s", scripts_dst) + elif not scripts_src.is_dir(): + log.error("MAD primus scripts not found — expected at %s", scripts_src) + + result = subprocess.run( + ["madengine", "--version"], + capture_output=True, text=True, + ) + if result.returncode == 0: + log.info("madengine version: %s", result.stdout.strip()) + else: + log.warning("madengine --version failed, but install may still be usable") + + return madengine_dir + + +def patch_madengine_for_cluster( + madengine_dir: Path, + no_gres: bool = False, +) -> None: + """Patch madengine source for cluster-specific compatibility.""" + src = madengine_dir / "src" / "madengine" + + if no_gres: + template = src / "deployment" / "templates" / "slurm" / "job.sh.j2" + if not template.exists(): + log.warning("SLURM template not found at %s", template) + else: + content = template.read_text() + patched = content.replace( + "#SBATCH --gpus-per-node={{ gpus_per_node }}\n", "" + ) + if patched != content: + template.write_text(patched) + log.info("Patched SLURM template: removed --gpus-per-node directive") + else: + log.info("SLURM template already patched (no --gpus-per-node)") + + template = src / "deployment" / "templates" / "slurm" / "job.sh.j2" + if template.exists(): + content = template.read_text() + marker = "# Load required modules" + if marker in content and "$HOME/.local/bin" not in content: + patched = content.replace( + marker, + 'export PATH="$HOME/.local/bin:$PATH"\n\n' + marker, + ) + template.write_text(patched) + log.info( + "Patched SLURM template: added $HOME/.local/bin to PATH " + "(SLURM jobs do not inherit user shell PATH)" + ) + else: + log.info("SLURM template PATH patch already present or marker not found") + + slurm_py = src / "deployment" / "slurm.py" + if slurm_py.exists(): + content = slurm_py.read_text() + patched = content.replace( + '["madengine", "--version"],\n' + " capture_output=True,\n" + " text=True,\n" + " timeout=5,", + '["madengine", "--version"],\n' + " capture_output=True,\n" + " text=True,\n" + " timeout=120,", + ) + if patched != content: + slurm_py.write_text(patched) + log.info("Patched slurm.py: increased CLI validation timeout to 120s") + else: + log.info("slurm.py already patched or timeout string not found") + + template = src / "deployment" / "templates" / "slurm" / "job.sh.j2" + if template and template.exists(): + content = template.read_text() + # Patch the MULTI-NODE verification block (inside TASK_SCRIPT_EOF + # heredoc) to install madengine per-node when the head node's venv + # is incompatible (Python 3.10 vs 3.9). The single-node block + # runs on the head node where the venv works — leave it alone. + # + # Find the multi-node block by searching for the verification + # string AFTER the TASK_SCRIPT_EOF heredoc marker. + heredoc_marker = "TASK_SCRIPT_EOF" + heredoc_idx = content.find(heredoc_marker) + if heredoc_idx != -1: + verify_str = 'echo "Verifying madengine availability..."' + mn_verify_idx = content.find(verify_str, heredoc_idx) + if mn_verify_idx == -1: + mn_verify_idx = content.find(verify_str) + if mn_verify_idx != -1: + mn_end_str = "# Create local execution manifest" + mn_end_idx = content.find(mn_end_str, mn_verify_idx) + if mn_end_idx != -1: + replacement = ( + 'echo "Verifying madengine availability..."\n' + 'MAD_CLI_COMMAND=""\n' + 'if command -v madengine >/dev/null 2>&1 && ' + 'madengine --help >/dev/null 2>&1; then\n' + ' MAD_CLI_COMMAND="madengine"\n' + ' echo " ✓ madengine available: ' + '$(madengine --version 2>&1 | head -1)"\n' + 'fi\n' + 'if [ -z "$MAD_CLI_COMMAND" ]; then\n' + ' echo " ⚠ madengine not functional — ' + 'installing for this node\'s Python ($(python3 --version))"\n' + ' SUBMISSION_DIR={{ manifest_file | dirname }}\n' + ' MADENGINE_SRC="$SUBMISSION_DIR/madengine"\n' + ' if [ -d "$MADENGINE_SRC" ] && [ -f "$MADENGINE_SRC/pyproject.toml" ]; then\n' + ' python3 -m venv "$WORKSPACE/node_venv"\n' + ' source "$WORKSPACE/node_venv/bin/activate"\n' + ' pip install --upgrade pip setuptools wheel 2>&1 | tail -3\n' + ' pip install "$MADENGINE_SRC" 2>&1 | tail -20\n' + ' if madengine --version >/dev/null 2>&1; then\n' + ' MAD_CLI_COMMAND="madengine"\n' + ' echo " ✓ madengine installed: ' + '$(madengine --version 2>&1 | head -1)"\n' + ' else\n' + ' echo " ✗ madengine install failed"\n' + ' exit 1\n' + ' fi\n' + ' else\n' + ' echo " ✗ madengine source not found at $MADENGINE_SRC"\n' + ' exit 1\n' + ' fi\n' + 'fi\n' + 'echo ""\n\n' + ) + content = content[:mn_verify_idx] + replacement + content[mn_end_idx:] + template.write_text(content) + log.info("Patched SLURM template: added per-node madengine install (multi-node)") + else: + log.warning("Could not find end of multi-node verification block") + else: + log.warning("Could not find multi-node verification block in template") + else: + log.warning("TASK_SCRIPT_EOF not found — template may not have multi-node support") + + template = src / "deployment" / "templates" / "slurm" / "job.sh.j2" + if template and template.exists(): + content = template.read_text() + old_nfs_pattern = r"\bnfs\b" + new_nfs_pattern = r"\bnfs[0-9]*\b" + if old_nfs_pattern in content and new_nfs_pattern not in content: + content = content.replace(old_nfs_pattern, new_nfs_pattern) + template.write_text(content) + log.info("Patched SLURM template: NFS detection now matches nfs4") + + run_orch = src / "orchestration" / "run_orchestrator.py" + if run_orch.exists(): + content = run_orch.read_text() + patched = content.replace( + 'print(self.console.sh("yum info rocm-libs", canFail=True))', + 'print(self.console.sh("rpm -qi rocm-libs 2>/dev/null ' + '|| echo rocm-libs not installed as RPM", canFail=True))', + ) + if patched != content: + run_orch.write_text(patched) + log.info( + "Patched run_orchestrator.py: replaced 'yum info' with 'rpm -qi' " + "to avoid interactive GPG prompt hang" + ) + else: + log.info("run_orchestrator.py already patched or yum string not found") + + +def get_rccl_commit(rccl_lib: Path | None = None) -> str: + """Derive a unique identifier for the RCCL build. + + Checks, in order: RCCL_COMMIT_HASH env, GITHUB_RUN_ID env, sha256 of + the librccl.so binary. Does NOT fall back to ``git rev-parse HEAD`` + because in CI the checkout is TheRock (not RCCL), which would produce + a constant tag and cause stale cache hits on persistent runners. + """ + commit = os.environ.get("RCCL_COMMIT_HASH", "") + if commit: + return commit[:12] + + run_id = os.environ.get("GITHUB_RUN_ID", "") + if run_id: + return f"run{run_id}" + + if rccl_lib and rccl_lib.exists(): + h = hashlib.sha256(rccl_lib.read_bytes()).hexdigest() + return h[:12] + + return "unknown" + + +def _rccl_uses_kpack(rccl_lib: Path) -> bool: + """Check if librccl.so uses kpack (GPU kernels in separate .kpack files). + + TheRock builds with kpack produce a small .so (~4MB) with a + .rocm_kpack_ref section and an empty (NOBITS) .hip_fatbin section. + These libraries crash on base images whose HIP runtime pre-dates + kpack support. + """ + try: + result = subprocess.run( + ["readelf", "-S", str(rccl_lib)], + capture_output=True, text=True, timeout=10, + ) + return ".rocm_kpack_ref" in result.stdout + except (subprocess.TimeoutExpired, FileNotFoundError): + return False + + +def build_rccl_overlay_image( + rccl_lib: Path, + base_image: str, + gpu_target: str, + work_dir: Path, + registry: str = "", +) -> str: + """Build a Docker overlay image with the CI-built RCCL and push to registry. + + When a registry is provided, the image is tagged and pushed so that + SLURM compute nodes can pull it automatically. Returns the final + image tag (registry-qualified if pushed). + """ + rccl_commit = get_rccl_commit(rccl_lib) + tag = f"{base_image}-rccl-{gpu_target}-{rccl_commit}" + + result = subprocess.run( + ["docker", "image", "inspect", tag], + capture_output=True, + ) + if result.returncode == 0: + log.info("Overlay image already exists on head node: %s", tag) + else: + dockerfile = work_dir / "Dockerfile.rccl-overlay" + rccl_lib_dir = rccl_lib.parent + uses_kpack = _rccl_uses_kpack(rccl_lib) + + staging_dir = work_dir / "rccl_libs" + staging_dir.mkdir(exist_ok=True) + for so_file in rccl_lib_dir.glob("librccl*"): + dest = staging_dir / so_file.name + if not dest.exists(): + subprocess.run(["cp", "-L", str(so_file), str(dest)], check=True) + + if uses_kpack: + kpack_files = list(rccl_lib_dir.rglob("*.kpack")) + if not kpack_files: + kpack_files = list(rccl_lib.parent.parent.rglob("rccl*.kpack")) + has_kpack_files = len(kpack_files) > 0 + if has_kpack_files: + kpack_staging = staging_dir / ".kpack" + kpack_staging.mkdir(exist_ok=True) + for kp in kpack_files: + dest = kpack_staging / kp.name + if not dest.exists(): + subprocess.run(["cp", "-L", str(kp), str(dest)], check=True) + log.info("Found %d kpack file(s): %s", + len(kpack_files), [f.name for f in kpack_files]) + else: + log.warning("RCCL .so has kpack references but no .kpack files found in artifacts") + log.info( + "CI-built librccl.so uses kpack (%.1f MB .so). " + "Building overlay with SDK venv layout for %s.", + rccl_lib.stat().st_size / 1e6, + base_image, + ) + dockerfile.write_text(f"""\ +FROM {base_image} +COPY rccl_libs/ /tmp/rccl_ci/ +RUN set -e; \\ + SDK_LIB="/opt/venv/lib/python3.12/site-packages/_rocm_sdk_libraries/lib"; \\ + SDK_DEV="/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib"; \\ + SDK_KPACK="/opt/venv/lib/python3.12/site-packages/_rocm_sdk_libraries/.kpack"; \\ + cp /tmp/rccl_ci/librccl.so "$SDK_LIB/librccl.so.1"; \\ + cp /tmp/rccl_ci/librccl.so "$SDK_LIB/librccl.so.1.0" 2>/dev/null || true; \\ + cp /tmp/rccl_ci/librccl.so "$SDK_DEV/librccl.so.1"; \\ + cp /tmp/rccl_ci/librccl.so "$SDK_DEV/librccl.so.1.0"; \\ + if [ -d /tmp/rccl_ci/.kpack ] && ls /tmp/rccl_ci/.kpack/*.kpack >/dev/null 2>&1; then \\ + mkdir -p "$SDK_KPACK"; \\ + cp /tmp/rccl_ci/.kpack/*.kpack "$SDK_KPACK/"; \\ + fi; \\ + rm -rf /tmp/rccl_ci +ENV NCCL_DEBUG=WARN +""") + else: + log.info( + "CI-built librccl.so has embedded GPU kernels (%.1f MB)", + rccl_lib.stat().st_size / 1e6, + ) + dockerfile.write_text(f"""\ +FROM {base_image} +COPY rccl_libs/ /tmp/rccl_ci/ +RUN set -e; \\ + RCCL_REAL=$(readlink -f /opt/rocm/lib/librccl.so 2>/dev/null || \\ + find /opt/rocm*/lib -name 'librccl.so.*.*' -not -type l 2>/dev/null | head -1); \\ + cp /tmp/rccl_ci/librccl.so "$RCCL_REAL"; \\ + rm -rf /tmp/rccl_ci +ENV NCCL_DEBUG=WARN +""") + + log.info("Building overlay image: %s", tag) + subprocess.run( + ["docker", "build", "--network=none", "-t", tag, + "-f", str(dockerfile), str(work_dir)], + check=True, + ) + log.info("Overlay image built: %s", tag) + + if registry: + safe_base = base_image.replace("/", "-").replace(":", "-") + push_tag = f"{registry}/rccl-ci:{safe_base}-{rccl_commit}" + log.info("Tagging overlay for registry: %s -> %s", tag, push_tag) + subprocess.run(["docker", "tag", tag, push_tag], check=True) + log.info("Pushing overlay image to registry: %s", push_tag) + subprocess.run(["docker", "push", push_tag], check=True) + log.info("Overlay image pushed: %s", push_tag) + return push_tag + + return tag + + +def generate_manifest( + workload_name: str, + workload_config: dict, + cluster_config: dict, + overlay_image: str, + nodes: int, + work_dir: Path, + nodelist: str = "", + registry: str = "", +) -> Path: + """Generate a madengine manifest.json for the workload. + + Structure follows the reference template from the mad-rccl branch: + deployment config under ``deployment_config``, env vars inside both + ``context.docker_env_vars`` and ``deployment_config.env_vars``, mounts + in ``context.docker_mounts``. + """ + gpus_per_node = workload_config["gpus_per_node"] + + nccl_env = dict(cluster_config.get("nccl_env", {})) + if nodes == 1: + nccl_env.pop("NCCL_NET", None) + ifname = nccl_env.get("NCCL_SOCKET_IFNAME", "") + if "," in ifname: + nccl_env["NCCL_SOCKET_IFNAME"] = ifname.split(",")[0] + + socket_ifname = nccl_env.get("NCCL_SOCKET_IFNAME", "") + + # HF_TOKEN is passed via MAD_SECRETS_HFTOKEN in the process environment + # (set in the workflow). Do NOT write it into the manifest — the manifest + # is uploaded as a CI artifact and would leak the credential. + + model_repo = workload_config["model_repo"] + scripts_dir = work_dir / "scripts" / "primus_megatron-lm" + if not scripts_dir.is_dir(): + scripts_dir = work_dir / "scripts" / "primus" / "megatron-lm" + + image_key = "overlay" + gpu_indices = ",".join(str(i) for i in range(gpus_per_node)) + render_ds = [128 + i for i in range(gpus_per_node)] + + docker_env_vars = { + **nccl_env, + "NCCL_DEBUG": "WARN", + "NCCL_IB_DISABLE": "0", + "NCCL_TIMEOUT": "900", + "IBV_SHOW_WARNINGS": "1", + } + if socket_ifname: + docker_env_vars["GLOO_SOCKET_IFNAME"] = socket_ifname + + docker_mounts = dict(workload_config.get("docker_mounts", {})) + docker_run_opts = workload_config.get("docker_run_options", "") + + slurm_config = { + "partition": cluster_config.get("slurm_partition", workload_config["slurm_partition"]), + "qos": cluster_config.get("slurm_qos", ""), + "nodes": nodes, + "gpus_per_node": gpus_per_node, + "time": workload_config["time_limit"], + "output_dir": "./slurm_output", + "exclusive": True, + "enable_node_check": False, + "network_interface": socket_ifname, + **({"nodelist": nodelist} if nodelist else {}), + } + + manifest = { + "built_images": { + image_key: { + "docker_image": overlay_image, + "local_image": not bool(registry), + "registry_image": overlay_image if registry else None, + "registry": registry or None, + "base_docker": workload_config["base_image"], + "build_status": "SKIPPED", + "build_duration": 0, + "gpu_vendor": "AMD", + }, + }, + "built_models": { + image_key: { + "name": model_repo, + "tags": workload_config.get("tags", ["pyt", "pretrain", "training"]), + "dockerfile": "N/A (overlay image)", + "scripts": f"scripts/{scripts_dir.name}/run.sh", + "n_gpus": "-1", + "owner": "", + "training_precision": "", + "multiple_results": workload_config.get("multiple_results", ""), + "args": f"--model_repo {model_repo}", + "additional_docker_run_options": docker_run_opts, + "data": "", + "cred": "", + "timeout": None, + }, + }, + "context": { + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "docker_gpus": gpu_indices, + "gpu_renderDs": render_ds, + "docker_env_vars": docker_env_vars, + "docker_mounts": docker_mounts, + "docker_build_arg": {}, + }, + "deployment_config": { + "target": "slurm", + "slurm": slurm_config, + "distributed": { + "launcher": "primus", + "backend": "nccl", + "port": 29500, + "nnodes": nodes, + "nproc_per_node": gpus_per_node, + }, + "env_vars": { + **docker_env_vars, + "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1", + "TORCH_NCCL_HIGH_PRIORITY": "1", + "OMP_NUM_THREADS": "8", + "MIOPEN_FIND_MODE": "1", + }, + "debug": False, + "docker_gpus": gpu_indices, + }, + } + + manifest_path = work_dir / "manifest.json" + manifest_path.write_text(json.dumps(manifest, indent=2)) + log.info("Manifest written to: %s", manifest_path) + return manifest_path + + +def run_madengine( + manifest_path: Path, + output_csv: Path, + work_dir: Path, + timeout_minutes: int = 120, +) -> int: + """Run madengine with the given manifest and return the exit code. + + The manifest already contains deployment_config with slurm, distributed, + and env_vars sections. madengine merges deployment_config into + additional_context automatically (run_orchestrator.py:225-234), so we + do not need to duplicate those here. + """ + cmd = [ + "madengine", "run", + "-m", str(manifest_path), + "-o", str(output_csv), + "--live-output", + "--verbose", + ] + + log.info("Running: %s", " ".join(cmd)) + log.info("Timeout: %d minutes", timeout_minutes) + + # Pre-warm: madengine's SLURM deployment validates CLI availability by + # running `madengine --version` with a 5s timeout. Cold import of + # madengine's heavy dependencies (kubernetes, aiohttp, paramiko) can + # exceed 5s. Running it once beforehand populates the bytecode cache. + try: + subprocess.run(["madengine", "--version"], capture_output=True, timeout=120) + except subprocess.TimeoutExpired: + log.warning("madengine --version pre-warm timed out (non-fatal)") + except Exception: + pass + + env = os.environ.copy() + docker_builds_dir = work_dir / "docker_builds" + docker_builds_dir.mkdir(exist_ok=True) + env["MAD_DOCKER_BUILDS"] = str(docker_builds_dir) + + try: + proc = subprocess.run( + cmd, + cwd=work_dir, + env=env, + timeout=timeout_minutes * 60, + ) + log.info("madengine exit code: %d", proc.returncode) + return proc.returncode + except subprocess.TimeoutExpired: + log.error("madengine timed out after %d minutes", timeout_minutes) + return 124 + + +def parse_perf_results(work_dir: Path) -> list[dict]: + """Parse madengine performance results. + + Prefers ``perf_entry_super.json`` (31 fixed columns, per-precision rows + with ``multi_results``). Falls back to ``perf.csv`` (variable-width, + long-format). Returns a list of result dicts — one per row. + """ + super_json = work_dir / "perf_entry_super.json" + if super_json.exists(): + try: + entries = json.loads(super_json.read_text()) + if entries: + log.info("Parsed %d result(s) from perf_entry_super.json", len(entries)) + for e in entries: + log.info(" model=%s perf=%s metric=%s status=%s precision=%s", + e.get("model"), e.get("performance"), + e.get("metric"), e.get("status"), + e.get("training_precision")) + return entries + except (json.JSONDecodeError, TypeError) as exc: + log.warning("Could not parse %s: %s", super_json, exc) + + csv_path = work_dir / "perf.csv" + if csv_path.exists(): + rows = [] + with open(csv_path) as f: + for row in csv.DictReader(f): + log.info("perf.csv row: %s", dict(row)) + rows.append(dict(row)) + if rows: + log.info("Parsed %d row(s) from perf.csv", len(rows)) + return rows + + log.warning("No perf results found in %s", work_dir) + return [] + + +_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") +_ITER_RE = re.compile( + r"iteration\s+(?P\d+)/\s*(?P\d+)" + r".*throughput per GPU \(TFLOP/s/GPU\):\s*[\d.]+/(?P[\d.]+)" + r".*tokens per GPU \(tokens/s/GPU\):\s*[\d.]+/(?P[\d.]+)" +) +_RUN_HEADER_RE = re.compile(r"Running:\s+(.+)\s+-\s+(\w+)\s+-\s+(\w+)\s*$") + + +def parse_live_log_metrics(work_dir: Path) -> list[dict]: + """Parse madengine live logs for training metrics. + + Detects multiple runs within a single log (e.g. BF16 then FP8) by + watching for "Running: Model - Precision - Mode" header lines. + + Returns a list of run dicts with keys: + model, precision, mode, iter, total, tflops_avg, + tokens_per_second_per_gpu, completed, log_file + """ + logs = sorted(work_dir.glob("*.run.live.log")) + if not logs: + return [] + + runs = [] + for log_path in logs: + current = None + with open(log_path) as f: + for raw_line in f: + line = _ANSI_RE.sub("", raw_line) + + hdr = _RUN_HEADER_RE.search(line) + if hdr: + if current: + current.setdefault("iter", 0) + current.setdefault("total", 0) + current["completed"] = ( + current["iter"] > 0 + and current["iter"] == current["total"] + ) + runs.append(current) + current = { + "model": hdr.group(1).strip(), + "precision": hdr.group(2), + "mode": hdr.group(3), + "log_file": str(log_path), + } + continue + + m = _ITER_RE.search(line) + if m: + if current is None: + current = {"log_file": str(log_path)} + current.update({ + "iter": int(m.group("iter")), + "total": int(m.group("total")), + "tflops_avg": float(m.group("tflops_avg")), + "tokens_per_second_per_gpu": float(m.group("tps_avg")), + }) + + if current: + current.setdefault("iter", 0) + current.setdefault("total", 0) + current["completed"] = ( + current["iter"] > 0 + and current["iter"] == current["total"] + ) + runs.append(current) + + return runs + + +def check_regression( + results_dir: Path, + workload_name: str, + scale: str, + current_value: float, + workload_type: str, + precision: str | None = None, +) -> tuple[bool, str]: + """Check if current metric is a regression vs rolling average. + + Returns (is_regression, message). + """ + datastore = results_dir / PERF_DATASTORE + if not datastore.exists(): + return False, "No historical data yet — skipping regression check" + + threshold = ( + REGRESSION_THRESHOLD_TRAINING + if workload_type == "training" + else REGRESSION_THRESHOLD_INFERENCE + ) + + historical = [] + with open(datastore) as f: + for line in f: + line = line.strip() + if not line: + continue + try: + entry = json.loads(line) + if ( + entry.get("workload") == workload_name + and entry.get("scale") == scale + and entry.get("precision") == precision + and entry.get("status") == "pass" + and entry.get("metric_value") is not None + ): + historical.append(entry["metric_value"]) + except json.JSONDecodeError: + continue + + if len(historical) < 3: + return False, f"Only {len(historical)} historical data points — need at least 3 for regression check" + + window = historical[-REGRESSION_WINDOW:] + rolling_avg = sum(window) / len(window) + if rolling_avg == 0: + return False, "Rolling average is 0 — skipping regression check" + pct_change = (current_value - rolling_avg) / rolling_avg + + msg = ( + f"Current: {current_value:.1f}, " + f"Rolling avg ({len(window)} runs): {rolling_avg:.1f}, " + f"Change: {pct_change:+.1%}, " + f"Threshold: -{threshold:.0%}" + ) + + if pct_change < -threshold: + return True, f"REGRESSION DETECTED — {msg}" + + return False, f"No regression — {msg}" + + +def append_result( + results_dir: Path, + workload_name: str, + scale: str, + metric_value: float | None, + status: str, + rccl_commit: str, + extra: dict | None = None, + precision: str | None = None, + tflops: float | None = None, + tokens_per_sec: float | None = None, +) -> None: + """Append a result entry to the JSONL datastore.""" + results_dir.mkdir(parents=True, exist_ok=True) + datastore = results_dir / PERF_DATASTORE + + entry = { + "run_id": os.environ.get("GITHUB_RUN_ID", "local"), + "timestamp": datetime.now(timezone.utc).isoformat(), + "commit": rccl_commit, + "workload": workload_name, + "scale": scale, + "precision": precision, + "tflops_per_gpu": tflops, + "tokens_per_sec_per_gpu": tokens_per_sec, + "metric_value": metric_value, + "status": status, + } + if extra: + entry.update(extra) + + with open(datastore, "a") as f: + f.write(json.dumps(entry) + "\n") + log.info("Result appended to %s", datastore) + + run_id = os.environ.get("GITHUB_RUN_ID", "local") + run_dir = results_dir / "runs" / run_id + run_dir.mkdir(parents=True, exist_ok=True) + + +def generate_summary_report( + workload_name: str, + scale: str, + exit_code: int, + metric_value: float | None, + regression_msg: str, + rccl_commit: str, + cluster: str, + precision_results: list[dict] | None = None, +) -> str: + """Generate a plain-text summary report.""" + status = "PASSED" if exit_code == 0 else "FAILED" + lines = [ + "RCCL MADEngine Workload Test Report", + "=" * 40, + f"Status: {status}", + f"Date: {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S UTC')}", + "", + f"Workload: {workload_name}", + f"Scale: {scale}", + f"Cluster: {cluster}", + f"RCCL: {rccl_commit}", + "", + ] + + if precision_results: + for r in precision_results: + prec = r.get("precision", "?") + val = r.get("metric_value") + tflops = r.get("tflops_avg") + val_s = f"{val:.1f}" if val else "N/A" + tflops_s = f"{tflops:.1f}" if tflops else "" + suffix = f" ({tflops_s} TFLOP/s/GPU)" if tflops_s else "" + lines.append(f"{prec:>4}: {val_s} tok/s/GPU{suffix} [{r['status']}]") + elif metric_value is not None: + lines.append(f"Throughput: {metric_value:.1f} tok/s/GPU") + else: + lines.append("Throughput: N/A (workload did not produce metrics)") + + lines.append("") + lines.append(f"Regression: {regression_msg}") + lines.append("") + + run_url = os.environ.get("GITHUB_SERVER_URL", "") + repo = os.environ.get("GITHUB_REPOSITORY", "") + run_id = os.environ.get("GITHUB_RUN_ID", "") + if run_url and repo and run_id: + lines.append(f"CI run: {run_url}/{repo}/actions/runs/{run_id}") + + return "\n".join(lines) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + "--artifact-dir", + type=Path, + required=True, + help="Directory containing CI-built RCCL artifacts", + ) + parser.add_argument( + "--workload", + type=str, + required=True, + choices=list(WORKLOAD_CONFIGS.keys()), + help="Workload to run", + ) + parser.add_argument( + "--cluster", + type=str, + required=True, + choices=list(CLUSTER_CONFIGS.keys()), + help="Target cluster", + ) + parser.add_argument( + "--nodes", + type=int, + default=2, + help="Number of nodes to allocate (default: 2)", + ) + parser.add_argument( + "--results-dir", + type=Path, + default=None, + help="Directory for JSONL datastore and run artifacts (default: cluster-specific path)", + ) + parser.add_argument( + "--work-dir", + type=Path, + default=None, + help="Working directory for madengine install, overlay build, etc.", + ) + parser.add_argument( + "--timeout-minutes", + type=int, + default=190, + help="Timeout for madengine run in minutes (default: 190)", + ) + parser.add_argument( + "--notify-email", + type=str, + default="", + help="Send summary report to this email address", + ) + parser.add_argument( + "--teams-webhook", + type=str, + default="", + help="Send summary report to this Teams webhook URL", + ) + parser.add_argument( + "--registry", + type=str, + default="", + help="Container registry to push overlay image to (e.g. ghcr.io/rocm/rocm-systems)", + ) + parser.add_argument( + "--skip-overlay-build", + action="store_true", + help="Skip Docker overlay build (use pre-built image specified via --overlay-image)", + ) + parser.add_argument( + "--overlay-image", + type=str, + default="", + help="Pre-built overlay image to use (requires --skip-overlay-build)", + ) + + args = parser.parse_args() + + workload_config = WORKLOAD_CONFIGS[args.workload] + cluster_config = CLUSTER_CONFIGS[args.cluster] + scale = f"{args.nodes}N/{args.nodes * workload_config['gpus_per_node']}GPU" + + results_dir = args.results_dir or Path(cluster_config["results_base"]) + work_dir = args.work_dir or Path(tempfile.mkdtemp(prefix="madengine_ci_")) + log.info("Work directory: %s", work_dir) + log.info("Results directory: %s", results_dir) + + # Step 1: Find RCCL library + rccl_lib = find_rccl_library(args.artifact_dir) + log.info("RCCL library: %s", rccl_lib) + rccl_commit = get_rccl_commit(rccl_lib) + log.info("RCCL commit/tag: %s", rccl_commit) + + # Step 2: Install madengine + madengine_dir = install_madengine(work_dir) + + patch_madengine_for_cluster( + madengine_dir, + no_gres=cluster_config.get("slurm_no_gres", False), + ) + + # Step 3: Build overlay image (or use pre-built) + if args.skip_overlay_build: + if not args.overlay_image: + log.error("--skip-overlay-build requires --overlay-image") + sys.exit(1) + overlay_image = args.overlay_image + else: + overlay_image = build_rccl_overlay_image( + rccl_lib, + workload_config["base_image"], + cluster_config["gpu_target"], + work_dir, + registry=args.registry, + ) + + # Step 4: Generate manifest + # When no registry is configured, the overlay image only exists on the + # node that built it. Pin the SLURM job to that node so madengine can + # find the image locally. + nodelist = "" + if not args.registry: + nodelist = os.environ.get("SLURM_NODELIST", "") + if not nodelist: + hostname = subprocess.run( + ["hostname", "-s"], capture_output=True, text=True, + ).stdout.strip() + if hostname: + nodelist = hostname + if nodelist: + log.info("No registry — pinning SLURM job to build node: %s", nodelist) + + manifest_path = generate_manifest( + args.workload, + workload_config, + cluster_config, + overlay_image, + args.nodes, + work_dir, + nodelist=nodelist, + registry=args.registry, + ) + + # Step 5: Run the workload + output_csv = work_dir / "perf.csv" + exit_code = run_madengine( + manifest_path, output_csv, work_dir, args.timeout_minutes, + ) + + # Step 6: Parse results — structured output first, live log fallback + perf_results = parse_perf_results(work_dir) + live_log_runs = parse_live_log_metrics(work_dir) + + # Save run artifacts + run_id = os.environ.get("GITHUB_RUN_ID", "local") + run_artifacts = results_dir / "runs" / run_id + try: + run_artifacts.mkdir(parents=True, exist_ok=True) + for f in ["perf.csv", "perf_entry_super.csv", "perf_entry_super.json"]: + src = work_dir / f + if src.exists(): + shutil.copy2(str(src), str(run_artifacts / f)) + except OSError as exc: + log.warning("Could not save run artifacts to %s: %s", run_artifacts, exc) + + # Build per-precision results from structured output (primary) or + # live-log scraping (fallback). Each entry carries precision, metric + # value, and a pass/fail status so downstream regression checks and + # datastore writes are driven from one list. + metric_key = workload_config["metric_key"] + precision_results: list[dict] = [] + + if perf_results: + # perf_entry_super.json rows are long-format: metric name is a + # value in the ``metric`` column, performance in ``performance``. + # Filter to the configured metric_key and key by precision. + for row in perf_results: + if row.get("metric") != metric_key: + continue + perf_val = row.get("performance", "") + precision = row.get("training_precision", "") + row_status = row.get("status", "") + if not perf_val: + continue + try: + val = float(perf_val) + except (ValueError, TypeError): + continue + precision_results.append({ + "precision": precision, + "metric_value": val, + "status": "pass" if row_status in ("", "pass", "PASS") else "fail", + "source": "structured", + }) + log.info("Structured result: %s %s = %.1f (status=%s)", + precision, metric_key, val, row_status) + + if not precision_results and live_log_runs: + for run in live_log_runs: + val = run.get("tokens_per_second_per_gpu") + if val is None: + continue + precision_results.append({ + "precision": run.get("precision"), + "metric_value": val, + "tflops_avg": run.get("tflops_avg"), + "status": "pass" if run.get("completed", False) else "fail", + "source": "live_log", + "iter": run.get("iter", 0), + "total": run.get("total", 0), + "log_file": run.get("log_file"), + }) + log.info("Live-log result: %s = %.1f (completed=%s)", + run.get("precision"), val, run.get("completed")) + + metric_value = precision_results[-1]["metric_value"] if precision_results else None + + # Override exit_code if training actually completed successfully. + # madengine can report failure (exit code 3) when its perf collector + # can't parse the output format, even though training ran to completion. + if exit_code != 0 and precision_results: + all_pass = all(r["status"] == "pass" for r in precision_results) + has_metric = all(r.get("metric_value") is not None for r in precision_results) + if all_pass and has_metric: + log.info( + "Overriding madengine exit code %d → 0: all %d precision run(s) " + "passed with metrics", + exit_code, len(precision_results), + ) + for r in precision_results: + log.info(" %s: %.1f %s", r.get("precision", "?"), + r["metric_value"], metric_key) + exit_code = 0 + + # Step 7: Per-precision regression check + regression_msg = "N/A" + is_regression = False + if precision_results: + regression_msgs = [] + for pr in precision_results: + if pr["metric_value"] is not None: + reg, msg = check_regression( + results_dir, args.workload, scale, pr["metric_value"], + workload_config["type"], precision=pr.get("precision"), + ) + regression_msgs.append(f"[{pr.get('precision', '?')}] {msg}") + if reg: + is_regression = True + log.warning(msg) + else: + log.info(msg) + regression_msg = "; ".join(regression_msgs) if regression_msgs else "N/A" + if is_regression: + exit_code = max(exit_code, 1) + + # Step 8: Append result to datastore (one record per precision run) + extra = {"cluster": args.cluster, "overlay_image": overlay_image} + if precision_results: + for pr in precision_results: + append_result( + results_dir, + args.workload, + scale, + pr["metric_value"], + pr["status"], + rccl_commit, + extra=extra, + precision=pr.get("precision"), + tflops=pr.get("tflops_avg"), + tokens_per_sec=pr["metric_value"], + ) + else: + append_result( + results_dir, + args.workload, + scale, + None, + "fail", + rccl_commit, + extra=extra, + ) + + # Step 9: Generate and distribute report + status = "pass" if exit_code == 0 else "fail" + report = generate_summary_report( + args.workload, scale, exit_code, metric_value, + regression_msg, rccl_commit, args.cluster, + precision_results=precision_results if precision_results else None, + ) + log.info("\n%s", report) + write_github_summary(report) + set_github_output("madengine_status", status) + if metric_value is not None: + set_github_output("madengine_metric", f"{metric_value:.1f}") + + summary_path = work_dir / "madengine_summary.txt" + summary_path.write_text(report) + + report_status = "PASSED" if exit_code == 0 else "FAILED" + if args.notify_email: + send_email_report(report, args.notify_email, report_status, + subject_prefix=f"RCCL MADEngine {args.workload}") + if args.teams_webhook: + send_teams_webhook(report, args.teams_webhook, report_status, + subject_prefix=f"RCCL MADEngine {args.workload}") + + sys.exit(exit_code) + + +if __name__ == "__main__": + main() diff --git a/projects/rccl/ci/scripts/test_pytorch_c10d.py b/projects/rccl/ci/scripts/test_pytorch_c10d.py new file mode 100644 index 00000000000..4139ccab0ca --- /dev/null +++ b/projects/rccl/ci/scripts/test_pytorch_c10d.py @@ -0,0 +1,450 @@ +#!/usr/bin/env python3 +"""Run PyTorch c10d NCCL distributed tests against CI-built RCCL. + +This script handles: + 1. Discovering the CI-built librccl.so in the artifact directory + 2. Verifying that LD_LIBRARY_PATH overrides PyTorch's bundled RCCL + 3. Cloning the matching PyTorch test sources (sparse checkout) + 4. Running pytest on test_c10d_nccl.py + +Usage from GitHub Actions: + python .github/scripts/test_pytorch_c10d.py \ + --artifact-dir ./build \ + --pytorch-src ./pytorch-src \ + --results-log ./pytorch_c10d_results.log +""" + +import argparse +import logging +import os +import subprocess +import sys +import tempfile +from datetime import datetime, timezone +from pathlib import Path + +from rccl_ci_utils import ( + find_rccl_library, + parse_junit_xml, + send_email_report, + set_github_output, + verify_rccl_override, + write_github_summary, +) + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +log = logging.getLogger(__name__) + +SMOKE_TESTS = [ + "test_all_reduce_coalesced_nccl", + "test_all_reduce_coalesced_manager_nccl", + "test_allgather_base", + "test_all_gather_into_tensor_coalesced_manager_nccl", + "test_broadcast_coalesced_nccl", + "test_broadcast_subgroup", + "test_reduce_scatter_base_k", + "test_reduce_scatter_tensor_coalesced", + "test_non_blocking_p2p", + "test_send_recv_subgroup", + "test_nccl_barrier_device_ids", + "test_reduce_subgroup", + "test_scatter_subgroup", + "test_gather_subgroup", + "test_all_to_all_single", + "test_init_wo_backend_str", + "test_new_group", + "test_pass_nccl_options_high_priority_stream", + "test_set_process_group_desc", + "test_tensor_dtype_complex", + "test_batch_send_recv_subgroup", + "test_collectives", +] + + +def find_rocm_lib_dir(artifact_dir: Path) -> Path | None: + """Find the dist/rocm/lib directory in artifacts.""" + for d in artifact_dir.rglob("dist/rocm/lib"): + if d.is_dir(): + log.info("Found ROCm lib dir: %s", d) + return d + return None + + +def setup_ld_library_path(rccl_lib_dir: Path, rocm_lib_dir: Path | None) -> str: + """Prepend RCCL and ROCm lib dirs to LD_LIBRARY_PATH.""" + parts = [str(rccl_lib_dir.resolve())] + if rocm_lib_dir: + parts.append(str(rocm_lib_dir.resolve())) + existing = os.environ.get("LD_LIBRARY_PATH", "") + if existing: + parts.append(existing) + new_path = ":".join(parts) + os.environ["LD_LIBRARY_PATH"] = new_path + log.info("LD_LIBRARY_PATH=%s", new_path) + return new_path + + +def clone_pytorch_test_sources(pytorch_src: Path) -> None: + """Sparse-clone PyTorch test sources matching the installed torch version. + + For release builds (e.g. 2.5.0), clones at the matching tag. + For nightly builds (e.g. 2.14.0a0+rocm7.15.0a20260712), clones using + --shallow-since to get commits around the build date, then checks out the + commit closest to that date so test sources match the installed wheel. + """ + import re + from datetime import datetime, timedelta + + import torch + + torch_version = torch.__version__ + base_version = torch_version.split("+")[0] + log.info("PyTorch version: %s", torch_version) + + git_ref = f"v{base_version}" + result = subprocess.run( + ["git", "ls-remote", "--tags", "https://github.com/pytorch/pytorch.git", git_ref], + capture_output=True, + text=True, + ) + + date_match = re.search(r"(\d{8})", torch_version) + use_date_pinning = False + + if result.stdout.strip(): + log.info("Found tag %s", git_ref) + elif date_match: + build_date = date_match.group(1) + dt = datetime.strptime(build_date, "%Y%m%d") + shallow_since = (dt - timedelta(days=2)).strftime("%Y-%m-%d") + log.info("Tag %s not found; nightly build date %s", git_ref, build_date) + git_ref = "nightly" + use_date_pinning = True + else: + log.info("Tag %s not found, using nightly branch HEAD", git_ref) + git_ref = "nightly" + + if use_date_pinning: + log.info("Cloning PyTorch (ref=%s, shallow-since=%s, sparse) into %s", + git_ref, shallow_since, pytorch_src) + subprocess.run( + [ + "git", "clone", + f"--branch={git_ref}", + f"--shallow-since={shallow_since}", + "--filter=blob:none", + "--sparse", + "https://github.com/pytorch/pytorch.git", + str(pytorch_src), + ], + check=True, + ) + else: + log.info("Cloning PyTorch (ref=%s, depth=1, sparse) into %s", git_ref, pytorch_src) + subprocess.run( + [ + "git", "clone", + "--depth=1", + f"--branch={git_ref}", + "--filter=blob:none", + "--sparse", + "https://github.com/pytorch/pytorch.git", + str(pytorch_src), + ], + check=True, + ) + + subprocess.run( + ["git", "sparse-checkout", "set", "test/"], + cwd=pytorch_src, + check=True, + ) + + if use_date_pinning: + before = (dt + timedelta(days=1)).strftime("%Y-%m-%dT00:00:00") + result = subprocess.run( + ["git", "log", f"--before={before}", "--format=%H", "-1"], + cwd=pytorch_src, + capture_output=True, + text=True, + ) + if result.returncode == 0 and result.stdout.strip(): + commit = result.stdout.strip() + log.info("Checking out commit %s (latest before %s)", commit[:12], before) + subprocess.run( + ["git", "checkout", commit], + cwd=pytorch_src, + check=True, + ) + else: + log.warning("Could not find commit before %s, using HEAD of nightly", before) + + test_file = pytorch_src / "test" / "distributed" / "test_c10d_nccl.py" + if not test_file.exists(): + log.error("test_c10d_nccl.py not found after clone") + sys.exit(1) + log.info("Test sources ready: %s", test_file) + + +def patch_missing_torch_modules() -> None: + """Create stubs for internal torch modules missing from nightly wheels.""" + import torch + + torch_dir = Path(torch.__file__).parent + strobelight_dir = torch_dir / "_strobelight" + profiler_file = strobelight_dir / "compile_time_profiler.py" + if not profiler_file.exists(): + log.info("Creating stub for torch._strobelight (missing from nightly wheel)") + strobelight_dir.mkdir(parents=True, exist_ok=True) + (strobelight_dir / "__init__.py").write_text("") + profiler_file.write_text( + "class StrobelightCompileTimeProfiler:\n" + " def __enter__(self): return self\n" + " def __exit__(self, *a): pass\n" + ) + + +def print_environment_info() -> None: + """Print GPU and environment details for CI logs.""" + import torch + + log.info("PyTorch: %s", torch.__version__) + log.info("CUDA/HIP available: %s", torch.cuda.is_available()) + log.info("GPU count: %s", torch.cuda.device_count()) + for i in range(torch.cuda.device_count()): + log.info(" GPU %d: %s", i, torch.cuda.get_device_name(i)) + log.info("LD_LIBRARY_PATH: %s", os.environ.get("LD_LIBRARY_PATH", "")) + + +def run_tests(pytorch_src: Path, results_log: Path, test_scope: str = "smoke") -> tuple[int, dict]: + """Run pytest on test_c10d_nccl.py and return (exit_code, summary_dict).""" + miopen_cache = tempfile.mkdtemp(prefix="miopen_cache_") + os.environ["MIOPEN_USER_DB_PATH"] = miopen_cache + + env = os.environ.copy() + env["PYTHONPATH"] = str(pytorch_src / "test") + ":" + env.get("PYTHONPATH", "") + + junit_xml = results_log.parent / "pytorch_c10d_results.xml" + + if test_scope == "smoke": + k_expr = " or ".join(SMOKE_TESTS) + timeout = "60" + log.info("Running smoke tests (%d tests)", len(SMOKE_TESTS)) + else: + k_expr = "not NCCLTraceTestDumpOnTimeout" + timeout = "600" + log.info("Running all tests (excluding NCCLTraceTestDumpOnTimeout)") + + cmd = [ + sys.executable, + "-m", + "pytest", + str(pytorch_src / "test" / "distributed" / "test_c10d_nccl.py"), + "-v", + f"--timeout={timeout}", + "--tb=short", + f"--junitxml={junit_xml}", + "-k", + k_expr, + ] + log.info("Running: %s", " ".join(cmd)) + + with open(results_log, "w") as log_file: + proc = subprocess.Popen( + cmd, + cwd=pytorch_src, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + for line in proc.stdout: + sys.stdout.write(line) + sys.stdout.flush() + log_file.write(line) + proc.wait() + + log.info("Test exit code: %d", proc.returncode) + log.info("Results written to: %s", results_log) + + exit_code = proc.returncode + passed_tests = [] + failed_tests = [] + error_details = [] + tests_run = 0 + summary_line = "" + + if junit_xml.exists(): + log.info("Parsing JUnit XML: %s", junit_xml) + junit = parse_junit_xml(junit_xml) + passed_tests = junit["passed"] + failed_tests = junit["failed"] + error_details = junit["error_details"] + tests_run = junit["tests_run"] + parts = [] + if passed_tests: + parts.append(f"{len(passed_tests)} passed") + if failed_tests: + parts.append(f"{len(failed_tests)} failed") + summary_line = ", ".join(parts) + + if error_details: + log.info("Failure/error details from JUnit XML:") + for detail in error_details: + log.info(" %s", detail) + else: + log.warning("JUnit XML not found at %s, falling back to exit code only", junit_xml) + + if test_scope == "smoke" and tests_run < len(SMOKE_TESTS): + log.error( + "Expected %d smoke tests but only %d were collected — " + "test names may have changed in the nightly", + len(SMOKE_TESTS), + tests_run, + ) + exit_code = 1 + + summary = { + "exit_code": exit_code, + "test_scope": test_scope, + "passed": passed_tests, + "failed": failed_tests, + "summary_line": summary_line, + "tests_run": tests_run, + "expected_tests": len(SMOKE_TESTS) if test_scope == "smoke" else None, + } + return exit_code, summary + + +def generate_summary_report(summary: dict, rccl_lib: Path) -> str: + """Generate a plain-text summary report.""" + import torch + + status = "PASSED" if summary["exit_code"] == 0 else "FAILED" + gpu_info = [] + for i in range(torch.cuda.device_count()): + gpu_info.append(f" GPU {i}: {torch.cuda.get_device_name(i)}") + + lines = [ + f"RCCL PyTorch c10d Test Report", + f"{'=' * 40}", + f"Status: {status}", + f"Test scope: {summary['test_scope']}", + f"Date: {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S UTC')}", + f"", + f"PyTorch: {torch.__version__}", + f"RCCL: {rccl_lib}", + f"GPUs: {torch.cuda.device_count()}x {torch.cuda.get_device_name(0)}", + f"", + f"Results: {summary['summary_line']}", + ] + + if summary.get("expected_tests") is not None: + lines.append(f"Collected: {summary['tests_run']}/{summary['expected_tests']} expected smoke tests") + lines.append("") + + if summary["failed"]: + lines.append(f"FAILED tests ({len(summary['failed'])}):") + for name in summary["failed"]: + lines.append(f" FAIL {name}") + lines.append("") + + if summary["passed"]: + lines.append(f"PASSED tests ({len(summary['passed'])}):") + for name, duration in summary["passed"]: + lines.append(f" OK {name:60s} {duration}") + lines.append("") + + run_url = os.environ.get("GITHUB_SERVER_URL", "") + repo = os.environ.get("GITHUB_REPOSITORY", "") + run_id = os.environ.get("GITHUB_RUN_ID", "") + if run_url and repo and run_id: + lines.append(f"CI run: {run_url}/{repo}/actions/runs/{run_id}") + + return "\n".join(lines) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--artifact-dir", + type=Path, + required=True, + help="Directory containing CI-built artifacts", + ) + parser.add_argument( + "--pytorch-src", + type=Path, + required=True, + help="Directory to clone PyTorch test sources into", + ) + parser.add_argument( + "--results-log", + type=Path, + default=Path("pytorch_c10d_results.log"), + help="Path for test results log file", + ) + parser.add_argument( + "--test-scope", + choices=["smoke", "all"], + default="smoke", + help="Run smoke tests (22 curated tests, ~3min) or all tests (default: smoke)", + ) + parser.add_argument( + "--notify-email", + type=str, + default="", + help="Send summary report to this email address", + ) + parser.add_argument( + "--discover-only", + action="store_true", + help="Only discover library paths and set GITHUB_OUTPUT, then exit", + ) + + args = parser.parse_args() + + # Step 1: Discover RCCL library path + rccl_lib = find_rccl_library(args.artifact_dir) + rccl_lib_dir = rccl_lib.parent + rocm_lib_dir = find_rocm_lib_dir(args.artifact_dir) + + set_github_output("RCCL_LIB_DIR", str(rccl_lib_dir)) + if rocm_lib_dir: + set_github_output("ROCM_LIB_DIR", str(rocm_lib_dir)) + + if args.discover_only: + return + + # Step 2: Set up LD_LIBRARY_PATH and verify override + setup_ld_library_path(rccl_lib_dir, rocm_lib_dir) + verify_rccl_override(rccl_lib_dir) + + # Step 3: Clone PyTorch test sources + clone_pytorch_test_sources(args.pytorch_src) + + # Step 4: Patch missing modules, print environment info, and run tests + patch_missing_torch_modules() + print_environment_info() + exit_code, summary = run_tests(args.pytorch_src, args.results_log, args.test_scope) + + # Step 5: Generate and distribute summary report + report = generate_summary_report(summary, rccl_lib) + log.info("\n%s", report) + write_github_summary(report) + + summary_path = args.results_log.parent / "pytorch_c10d_summary.txt" + summary_path.write_text(report) + log.info("Summary written to: %s", summary_path) + + if args.notify_email: + status = "PASSED" if exit_code == 0 else "FAILED" + send_email_report(report, args.notify_email, status, + subject_prefix="RCCL PyTorch c10d Test") + + sys.exit(exit_code) + + +if __name__ == "__main__": + main()