Add MPI/rccl-tests installation support

cijohnson · cijohnson · commit e14c7283b5e6 · 2026-05-17T18:08:10.000-04:00
- Add install_mpi() function to rccl_lib.py for UCX and OpenMPI installation
- Add install_rccl_tests() function with MPI support and soft link creation
- Add install_rccl_test.py for automated RCCL tests installation

Made-with: Cursor
Signed-off-by: Ignatious Johnson &lt;ichristo@amd.com&gt;
diff --git a/cvs/input/config_file/rccl/install_rccl_config.json b/cvs/input/config_file/rccl/install_rccl_config.json
@@ -0,0 +1,28 @@
+{
+    "_comment": "Minimal RCCL JSON for cvs/tests/rccl/install_rccl_test.py only. Adjust paths for your cluster shared storage and ROCm layout.",
+    "rccl": {
+        "_comment_rccl_tests_dir": "Where sparse checkout is built and *_perf symlinks are created (must be visible on all nodes if shared filesystem).",
+        "rccl_tests_dir": "/opt/rccl-tests-install",
+
+        "_comment_rccl_tests_repo": "Mono-repo containing projects/rccl-tests (sparse clone).",
+        "rccl_tests_repo": "https://github.com/ROCm/rocm-systems.git",
+
+        "_comment_rccl_tests_branch": "Branch passed to git clone -b.",
+        "rccl_tests_branch": "develop",
+
+        "_comment_rocm_path_var": "ROCm root for hipcc/rocm_agent_enumerator and UCX --with-rocm. install.sh --rocm_home.",
+        "rocm_path_var": "/opt/rocm",
+
+        "_comment_rccl_path_var": "RCCL install prefix for install.sh --rccl_home. Often same as rocm_path_var when RCCL is under ROCm.",
+        "rccl_path_var": "/opt/rocm",
+
+        "_comment_mpi_path_var": "Open MPI prefix (directory containing bin/mpirun). Required for multi-node runs (num_nodes>1). If mpirun is missing here, UCX/OpenMPI are built into this path. Single-node: may be omitted.",
+        "mpi_path_var": "/usr",
+
+        "_comment_ucx_version": "Used only when MPI is auto-installed (multi-node and mpirun missing at mpi_path_var).",
+        "ucx_version": "1.18.0",
+
+        "_comment_ompi_version": "Used only when MPI is auto-installed.",
+        "ompi_version": "5.0.7"
+    }
+}
diff --git a/cvs/lib/rccl_lib.py b/cvs/lib/rccl_lib.py
@@ -1041,3 +1041,275 @@ def rccl_perf(
         check_lat_dip(test_name, results_for_verification, test_exp_dict)
 
     return all_raw_results
+
+
+def install_mpi(shdl, install_path, rocm_path, ucx_version="1.18.0", ompi_version="5.0.7"):
+    """Install UCX and OpenMPI for RCCL tests with MPI support.
+
+    Args:
+        shdl: Single SSH handle for executing commands on head node (shared storage)
+        install_path (str): Directory where UCX and OpenMPI should be installed
+        rocm_path (str): Path to ROCm installation
+        ucx_version (str, optional): UCX version to build. Defaults to "1.18.0"
+        ompi_version (str, optional): OpenMPI version to build. Defaults to "5.0.7"
+
+    Returns:
+        dict: Paths to UCX and OpenMPI installations
+
+    Raises:
+        Exception: If download, build, or installation fails
+    """
+    ucx_dir = f"ucx-{ucx_version}"
+    ompi_dir = f"ompi-{ompi_version}"
+    ompi_base_ver = ".".join(ompi_version.split(".")[:-1])  # e.g., "5.0" from "5.0.7"
+
+    ucx_install_dir = f"{install_path}/{ucx_dir}/install"
+    ompi_install_dir = f"{install_path}/{ompi_dir}/install"
+
+    # Require real artifacts, not only prefix directories (partial installs / timeouts can leave empty dirs).
+    check_ucx = shdl.exec(
+        f"test -x {ucx_install_dir}/bin/ucx_info && echo EXISTS",
+        timeout=10,
+        print_console=False,
+    )
+    check_ompi = shdl.exec(
+        f"test -x {ompi_install_dir}/bin/mpirun && echo EXISTS",
+        timeout=10,
+        print_console=False,
+    )
+
+    ucx_installed = any("EXISTS" in v for v in check_ucx.values())
+    ompi_installed = any("EXISTS" in v for v in check_ompi.values())
+
+    if ucx_installed and ompi_installed:
+        log.info("UCX and OpenMPI already installed in shared storage")
+        # Check if directory links exist, create them if missing
+        check_links = shdl.exec(
+            f"test -L {install_path}/bin && test -L {install_path}/lib && test -L {install_path}/include && echo LINKS_EXIST",
+            timeout=10,
+            print_console=False,
+        )
+        links_exist = any("LINKS_EXIST" in v for v in check_links.values())
+
+        if not links_exist:
+            log.info("[MPI] Directory links missing, creating them...")
+            link_cmd = f"""
+            cd {install_path} &&
+            ln -sf {ompi_install_dir}/bin bin &&
+            ln -sf {ompi_install_dir}/lib lib &&
+            ln -sf {ompi_install_dir}/include include
+            """
+            try:
+                shdl.exec(link_cmd, timeout=30, print_console=False)
+                log.info("[MPI] Directory soft links created successfully")
+            except Exception as e:
+                raise Exception(f"Failed to create MPI directory links: {e}")
+        else:
+            log.info("[MPI] Directory links already exist")
+
+        return {"ucx_path": ucx_install_dir, "ompi_path": install_path}
+
+    # Install UCX
+    if not ucx_installed:
+        log.info(f"[MPI] Installing UCX {ucx_version}...")
+        ucx_cmd = f"""
+        mkdir -p {install_path} &&
+        cd {install_path} &&
+        rm -rf {ucx_dir} ucx-{ucx_version}.tar.gz &&
+        wget https://github.com/openucx/ucx/releases/download/v{ucx_version}/ucx-{ucx_version}.tar.gz &&
+        mkdir -p {ucx_dir} &&
+        tar -zxf ucx-{ucx_version}.tar.gz -C {ucx_dir} --strip-components=1 &&
+        cd {ucx_dir} &&
+        mkdir -p build &&
+        cd build &&
+        ../configure --prefix={ucx_install_dir} --with-rocm={rocm_path} &&
+        make -j$(nproc) &&
+        make install
+        """
+        try:
+            shdl.exec(ucx_cmd, timeout=600, print_console=False)
+        except Exception as e:
+            raise Exception(f"Failed to install UCX: {e}")
+
+    # Install OpenMPI
+    if not ompi_installed:
+        log.info(f"[MPI] Installing OpenMPI {ompi_version}...")
+        ompi_cmd = f"""
+        mkdir -p {install_path} &&
+        cd {install_path} &&
+        rm -rf {ompi_dir} openmpi-{ompi_version}.tar.gz &&
+        wget https://download.open-mpi.org/release/open-mpi/v{ompi_base_ver}/openmpi-{ompi_version}.tar.gz &&
+        mkdir -p {ompi_dir} &&
+        tar -zxf openmpi-{ompi_version}.tar.gz -C {ompi_dir} --strip-components=1 &&
+        cd {ompi_dir} &&
+        mkdir -p build &&
+        cd build &&
+        ../configure --prefix={ompi_install_dir} --with-ucx={ucx_install_dir} --disable-oshmem --disable-mpi-fortran &&
+        make -j$(nproc) &&
+        make install
+        """
+        # Open MPI: configure + make -j + make install often exceeds 15 minutes on shared hosts.
+        try:
+            shdl.exec(ompi_cmd, timeout=1200, print_console=False)
+        except Exception as e:
+            raise Exception(f"Failed to install OpenMPI: {e}")
+
+    # Create directory-level soft links for MPI
+    # Structure: /mnt/scratch1/amd/ichristo/mpi/bin -> ompi-5.0.7/install/bin
+    link_cmd = f"""
+    cd {install_path} &&
+    ln -sf {ompi_install_dir}/bin bin &&
+    ln -sf {ompi_install_dir}/lib lib &&
+    ln -sf {ompi_install_dir}/include include
+    """
+    log.info("[MPI] Creating directory soft links for MPI bin, lib, and include...")
+    try:
+        shdl.exec(link_cmd, timeout=30, print_console=False)
+    except Exception as e:
+        raise Exception(f"Failed to create MPI soft links: {e}")
+
+    # Verify installations (same criteria as skip-if-installed checks above).
+    verify_ucx = shdl.exec(f"test -x {ucx_install_dir}/bin/ucx_info && echo EXISTS", timeout=10, print_console=False)
+    verify_ompi = shdl.exec(f"test -x {ompi_install_dir}/bin/mpirun && echo EXISTS", timeout=10, print_console=False)
+
+    ucx_verified = any("EXISTS" in v for v in verify_ucx.values())
+    ompi_verified = any("EXISTS" in v for v in verify_ompi.values())
+
+    if not ucx_verified:
+        raise Exception("UCX installation verification failed")
+    if not ompi_verified:
+        raise Exception("OpenMPI installation verification failed")
+
+    log.info("[MPI] UCX and OpenMPI installation complete on all nodes")
+    return {"ucx_path": ucx_install_dir, "ompi_path": install_path}
+
+
+def install_rccl_tests(
+    shdl,
+    install_path,
+    repo_url="https://github.com/ROCm/rocm-systems.git",
+    repo_branch="develop",
+    with_mpi=False,
+    mpi_home=None,
+    rccl_home=None,
+    rocm_home=None,
+):
+    """Install rccl-tests from ROCm/rocm-systems monorepo.
+
+    Args:
+        shdl: Single SSH handle for executing commands on head node (shared storage)
+        install_path (str): Directory where rccl-tests should be installed
+        repo_url (str, optional): Git repository URL. Defaults to ROCm/rocm-systems
+        repo_branch (str, optional): Git branch to clone. Defaults to 'develop'
+        with_mpi (bool, optional): Build with MPI support. Defaults to False
+        mpi_home (str, optional): Path to MPI installation. Required if with_mpi=True
+        rccl_home (str, optional): Path to RCCL library install (e.g. /opt/rocm)
+        rocm_home (str, optional): Path to ROCm root for hipcc/amdclang (e.g. /opt/rocm).
+            Defaults to rccl_home, then /opt/rocm. Must not be the rccl-tests work dir.
+
+    Returns:
+        str: Path to the install directory containing all *_perf binaries
+
+    Raises:
+        Exception: If clone, build, or verification fails
+    """
+    rocm = (rocm_home or "/opt/rocm").rstrip("/")
+    if rccl_home is None:
+        rccl = rocm
+    else:
+        rccl = str(rccl_home).rstrip("/")
+    if with_mpi and not mpi_home:
+        raise ValueError("mpi_home is required when with_mpi=True")
+    sparse_dir = "projects/rccl-tests"
+    clone_root = f"{install_path}/rocm-systems"
+
+    # Check if installation already exists (check for any *_perf binary)
+    check = shdl.exec(
+        f"ls {install_path}/*_perf 2>/dev/null | head -1 | grep -q _perf && echo EXISTS",
+        timeout=10,
+        print_console=False,
+    )
+    if any("EXISTS" in v for v in check.values()):
+        log.info(f"rccl-tests binaries already present in shared storage: {install_path}")
+        return install_path
+
+    # Clone rocm-systems repository
+    clone_cmd = (
+        f"mkdir -p {install_path} && "
+        f"rm -rf {clone_root} && "
+        f"git clone --depth 1 --filter=blob:none --sparse "
+        f"-b {repo_branch} {repo_url} {clone_root} && "
+        f"cd {clone_root} && git sparse-checkout set {sparse_dir}"
+    )
+    log.info(f"[rccl-tests] Cloning {sparse_dir} from rocm-systems to shared storage...")
+    try:
+        shdl.exec(clone_cmd, timeout=180, print_console=False)
+    except Exception as e:
+        raise Exception(f"Failed to clone rccl-tests: {e}")
+
+    # Upstream rccl-tests Makefile sets rpath to $(NCCL_HOME) (install root), but
+    # shared libraries live under $(NCCL_HOME)/lib. Patch rpath so binaries resolve
+    # against the configured ROCm install instead of system /opt/rocm.
+    patch_makefile_cmd = (
+        f"cd {clone_root}/{sparse_dir}/src && "
+        f"sed -i 's|-Wl,-rpath,$(NCCL_HOME)|-Wl,-rpath,$(NCCL_HOME)/lib -Wl,-rpath,$(ROCM_PATH)/lib|g' Makefile"
+    )
+    log.info("[rccl-tests] Patching Makefile rpath to prefer configured ROCm lib directories...")
+    try:
+        shdl.exec(patch_makefile_cmd, timeout=30, print_console=False)
+    except Exception as e:
+        raise Exception(f"Failed to patch rccl-tests Makefile rpath: {e}")
+
+    # Build rccl-tests using install.sh script
+    gpu_target = f"$({rocm}/bin/rocm_agent_enumerator | grep -v gfx000 | head -1)"
+
+    hip_compiler = f"{rocm}/llvm/bin/amdclang++"
+
+    if with_mpi:
+        build_cmd = (
+            f"cd {clone_root}/{sparse_dir} && "
+            f"ROCM_PATH={rocm} "
+            f"./install.sh --mpi --mpi_home {mpi_home} --rccl_home {rccl} --rocm_home {rocm} "
+            f"--hip_compiler {hip_compiler} --gpu_targets {gpu_target}"
+        )
+        log.info("[rccl-tests] Building rccl-tests with MPI support using install.sh in shared storage...")
+    else:
+        build_cmd = (
+            f"cd {clone_root}/{sparse_dir} && "
+            f"ROCM_PATH={rocm} "
+            f"./install.sh --rccl_home {rccl} --rocm_home {rocm} "
+            f"--hip_compiler {hip_compiler} --gpu_targets {gpu_target}"
+        )
+        log.info("[rccl-tests] Building rccl-tests without MPI using install.sh in shared storage...")
+
+    try:
+        shdl.exec(build_cmd, timeout=300, print_console=False)
+    except Exception as e:
+        raise Exception(f"Failed to build rccl-tests: {e}")
+
+    # Create soft links for all *_perf binaries
+    link_cmd = (
+        f"cd {install_path} && "
+        f"for perf_binary in {clone_root}/{sparse_dir}/build/*_perf; do "
+        f"  if [ -f \"$perf_binary\" ]; then "
+        f"    ln -sf \"$perf_binary\" \"$(basename \"$perf_binary\")\"; "
+        f"  fi; "
+        f"done"
+    )
+    log.info("[rccl-tests] Creating soft links for *_perf binaries...")
+    try:
+        shdl.exec(link_cmd, timeout=30, print_console=False)
+    except Exception as e:
+        raise Exception(f"Failed to create soft links: {e}")
+
+    # Verify installation (check for any *_perf binary)
+    verify = shdl.exec(
+        f"ls {install_path}/*_perf 2>/dev/null | head -1 | grep -q _perf && echo EXISTS",
+        timeout=10,
+        print_console=False,
+    )
+    if not any("EXISTS" in v for v in verify.values()):
+        raise Exception("rccl-tests build completed but no *_perf binaries found in shared storage")
+
+    log.info(f"[rccl-tests] Build complete in shared storage: {install_path}")
+    return install_path
diff --git a/cvs/tests/rccl/install_rccl_test.py b/cvs/tests/rccl/install_rccl_test.py
