From 74579ecea3380a8be2afb1ca81f1c4e7fcad8151 Mon Sep 17 00:00:00 2001 From: Zimin Li Date: Fri, 3 Jul 2026 07:32:25 +0000 Subject: [PATCH 1/2] feat: support node-specific launcher directories --- scripts/backends/mpi_base.py | 25 ++++++ scripts/icclrun_logic.py | 163 ++++++++++++++++++++++++++--------- 2 files changed, 145 insertions(+), 43 deletions(-) diff --git a/scripts/backends/mpi_base.py b/scripts/backends/mpi_base.py index 8078044..3857eba 100644 --- a/scripts/backends/mpi_base.py +++ b/scripts/backends/mpi_base.py @@ -1,4 +1,5 @@ import os +import subprocess class BaseMpiBackend: @@ -40,6 +41,30 @@ def get_launch_command(self, config, executable, user_args, launcher_obj): if not launcher_script: launcher_script = launcher_obj.ensure_launcher_exists() + # If any node-specific `dir` is given, stage the generated wrapper at one identical '/tmp' path on all nodes. + # Node-specific dirs may differ, but 'mpirun' can launch only one script path. + if any("dir" in node for node in config["nodes"]): + remote_launcher = f"/tmp/infiniccl_{os.path.basename(launcher_script)}" + subprocess.run(["cp", launcher_script, remote_launcher], check=True) + os.chmod(remote_launcher, 0o755) + + for node in config["nodes"]: + user = node.get("user", config.get("common_user", "root")) + + if launcher_obj._is_local(node["ip"]): + continue + + subprocess.run( + [ + "scp", + launcher_script, + f"{user}@{node['ip']}:{remote_launcher}", + ], + check=True, + ) + + launcher_script = remote_launcher + # Fetch specific base runner flags (OMPI vs MPICH). cmd = self.get_base_mpi_args(hostfile_path, total_slots) diff --git a/scripts/icclrun_logic.py b/scripts/icclrun_logic.py index 4760de7..c0f0d63 100644 --- a/scripts/icclrun_logic.py +++ b/scripts/icclrun_logic.py @@ -1,4 +1,5 @@ import os +import re import socket import subprocess import sys @@ -6,6 +7,18 @@ import yaml +PATH_LIKE = {"LD_LIBRARY_PATH", "PATH", "CPATH", "LIBRARY_PATH"} + + +def _is_infiniccl_source_dir(path): + cmake_path = Path(path).expanduser() / "CMakeLists.txt" + try: + cmake_text = cmake_path.read_text(errors="ignore") + except OSError: + return False + + return re.search(r"project\s*\(\s*InfiniCCL(?:\s|\))", cmake_text) is not None + class ICCLLauncher: def __init__(self, manual_config_path): @@ -24,13 +37,39 @@ def __init__(self, manual_config_path): break if not self.config: - print("Error: cluster.yaml not found.") + print("Error: `cluster.yaml` not found.") sys.exit(1) - # Detect InfiniCCL Root - self.infiniccl_root = os.environ.get("INFINICCL_ROOT") - if not self.infiniccl_root: - script_dir = os.path.dirname(os.path.realpath(__file__)) + # Detect the InfiniCCL source root. Prefer the configured tree + # checkout before environment variables, which may point at stale installs. + script_dir = os.path.dirname(os.path.realpath(__file__)) + candidates = [] + common_dir = self.config.get("common_dir") + if common_dir: + common_dir = os.path.expanduser(common_dir) + candidates.extend([common_dir]) + candidates.extend( + root + for root in [ + os.environ.get("INFINICCL_ROOT"), + os.environ.get("InfiniCCL_ROOT"), + ] + if root + ) + candidates.extend( + [ + os.path.join(script_dir, ".."), + os.path.join(script_dir, "../.."), + os.path.join(script_dir, "../../.."), + ] + ) + + for candidate in candidates: + root = os.path.abspath(os.path.expanduser(candidate)) + if _is_infiniccl_source_dir(root): + self.infiniccl_root = root + break + else: self.infiniccl_root = os.path.abspath(os.path.join(script_dir, "../../..")) def _is_local(self, ip): @@ -66,41 +105,55 @@ def _is_local(self, ip): return False def orchestrate_build(self): - common_dir = self.config["common_dir"] + global_common_dir = self.config["common_dir"] infiniccl_root = self.infiniccl_root - is_internal = os.path.abspath(common_dir) == os.path.abspath(infiniccl_root) - base_install = self.config.get("install_dir", common_dir) for node in self.config["nodes"]: + node_common_dir = node.get("dir", global_common_dir) + node_is_local = self._is_local(node["ip"]) + base_install = self.config.get("install_dir", node_common_dir) + arch = node["type"] install_path = os.path.join(base_install, "install", arch) user_cmake_flags = node.get( "cmake_flags", self.config.get("cmake_flags", "") ) - # Build the library using YAML flags. + internal_cmd = ( + f"mkdir -p {node_common_dir}/build/{arch} && cd {node_common_dir}/build/{arch} && " + f"cmake -DCMAKE_INSTALL_PREFIX={install_path} {user_cmake_flags} {node_common_dir} && " + f"make -j$(nproc) install" + ) lib_cmd = ( f"mkdir -p {infiniccl_root}/build/{arch} && cd {infiniccl_root}/build/{arch} && " f"cmake -DCMAKE_INSTALL_PREFIX={install_path} {user_cmake_flags} {infiniccl_root} && " f"make -j$(nproc) install" ) + app_cmd = ( + f"export INFINICCL_INSTALL={install_path} && " + f"mkdir -p {node_common_dir}/build/{arch} && cd {node_common_dir}/build/{arch} && " + f"cmake {user_cmake_flags} {node_common_dir} && " + f"make -j$(nproc)" + ) + external_cmd = f"{lib_cmd} && {app_cmd}" - if is_internal: - full_cmd = lib_cmd + if node_is_local: + full_cmd = ( + internal_cmd + if _is_infiniccl_source_dir(node_common_dir) + else external_cmd + ) else: - app_cmd = ( - f"export INFINICCL_INSTALL={install_path} && " - f"mkdir -p {common_dir}/build/{arch} && cd {common_dir}/build/{arch} && " - f"cmake {user_cmake_flags} {common_dir} && " - f"make -j$(nproc)" + full_cmd = ( + f'if grep -Eq "project[[:space:]]*\\([[:space:]]*InfiniCCL([[:space:]]|\\))" {node_common_dir}/CMakeLists.txt 2>/dev/null; ' + f"then {internal_cmd}; else {external_cmd}; fi" ) - full_cmd = f"{lib_cmd} && {app_cmd}" # Execute via SSH or locally. user = node.get("user", self.config.get("common_user", "root")) print(f"[*] Orchestrating {arch} on {node['ip']}...") - if self._is_local(node["ip"]): + if node_is_local: subprocess.run(["bash", "-l", "-c", full_cmd], check=True) else: subprocess.run( @@ -111,22 +164,10 @@ def orchestrate_build(self): return self.ensure_launcher_exists() def ensure_launcher_exists(self): - common_dir = Path(self.config["common_dir"]).expanduser().resolve() - wrapper_path = str(common_dir / "build" / "run_wrapper.sh") + global_common_dir = Path(self.config["common_dir"]).expanduser().resolve() + wrapper_path = str(global_common_dir / "build" / "run_wrapper.sh") os.makedirs(os.path.dirname(wrapper_path), exist_ok=True) - infiniccl_root_dir = Path(self.infiniccl_root).expanduser().resolve() - is_internal = common_dir == infiniccl_root_dir - - base_install = ( - Path(self.config.get("install_dir", self.config["common_dir"])) - .expanduser() - .resolve() - ) - - bin_sub = "examples/$1" if is_internal else "$1" - - PATH_LIKE = {"LD_LIBRARY_PATH", "PATH", "CPATH", "LIBRARY_PATH"} global_env = self.config.get("backend_env", {}) blocks = [] @@ -135,7 +176,23 @@ def ensure_launcher_exists(self): for node in self.config["nodes"]: ip_or_host = node["ip"] n_type = node["type"] - install_lib = base_install / "install" / n_type / "lib" + + node_common_dir_str = node.get("dir", self.config["common_dir"]) + node_is_local = self._is_local(ip_or_host) + node_common_dir = ( + str(Path(node_common_dir_str).expanduser().resolve()) + if node_is_local + else node_common_dir_str + ) + + base_install_str = self.config.get("install_dir", node_common_dir_str) + base_install = ( + str(Path(base_install_str).expanduser().resolve()) + if node_is_local + else base_install_str + ) + + install_lib = os.path.join(base_install, "install", n_type, "lib") resolved_ips = set() @@ -174,15 +231,14 @@ def ensure_launcher_exists(self): ) sys.exit(1) - node_env = node.get("backend_env", {}).copy() - - for concat_var in PATH_LIKE: - if concat_var in global_env and concat_var in node_env: - node_env[concat_var] = ( - f"{node_env[concat_var]}:{global_env[concat_var]}" - ) + node_env = node.get("backend_env", {}) + merged_env = global_env.copy() - merged_env = {**global_env, **node_env} + for env_key, env_value in node_env.items(): + if env_key in PATH_LIKE and env_key in merged_env: + merged_env[env_key] = f"{env_value}:{merged_env[env_key]}" + else: + merged_env[env_key] = env_value ip_conditions = [ f'[[ "$HOST_IPS" == *"{ip}"* ]]' for ip in sorted(resolved_ips) @@ -191,8 +247,9 @@ def ensure_launcher_exists(self): condition = " || ".join(ip_conditions) export_lines = [] + export_lines.append(f'INSTALL_LIB="$(expand_path "{install_lib}")"') export_lines.append( - f'export LD_LIBRARY_PATH="{install_lib}:${{LD_LIBRARY_PATH:-}}"' + 'export LD_LIBRARY_PATH="$INSTALL_LIB:${LD_LIBRARY_PATH:-}"' ) for k, v in merged_env.items(): @@ -202,6 +259,14 @@ def ensure_launcher_exists(self): export_lines.append(f'export {k}="{v}"') export_lines.append(f'ARCH="{n_type}"') + export_lines.append(f'NODE_COMMON_DIR="$(expand_path "{node_common_dir}")"') + export_lines.append( + 'if grep -Eq "project[[:space:]]*\\([[:space:]]*InfiniCCL([[:space:]]|\\))" "$NODE_COMMON_DIR/CMakeLists.txt" 2>/dev/null; then' + ) + export_lines.append(' BIN_SUB="examples/$1"') + export_lines.append("else") + export_lines.append(' BIN_SUB="$1"') + export_lines.append("fi") body = "\n".join(f" {line}" for line in export_lines) @@ -213,6 +278,14 @@ def ensure_launcher_exists(self): script_lines = [ "#!/bin/bash", "", + "expand_path() {", + ' case "$1" in', + ' "~") printf "%s\\n" "$HOME" ;;', + ' "~/"*) printf "%s\\n" "$HOME/${1#~/}" ;;', + ' *) printf "%s\\n" "$1" ;;', + " esac", + "}", + "", 'HOST_IPS="$(hostname -I)"', "", "".join(blocks).rstrip(), @@ -223,7 +296,11 @@ def ensure_launcher_exists(self): " exit 1", "fi", "", - f'EXE="{common_dir}/build/$ARCH/{bin_sub}"', + 'EXE="$NODE_COMMON_DIR/build/$ARCH/$BIN_SUB"', + 'if [[ ! -x "$EXE" ]]; then', + ' echo "[ERROR] Executable not found or not executable: $EXE"', + " exit 1", + "fi", "", "shift", "", From 7e0e468d7002e64ec55561f830fa77f8e375b3ab Mon Sep 17 00:00:00 2001 From: Zimin Li Date: Fri, 3 Jul 2026 08:10:31 +0000 Subject: [PATCH 2/2] docs: update `examples/cluster.yaml` to reflect the node-specific `dir` field --- examples/cluster.yaml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/examples/cluster.yaml b/examples/cluster.yaml index eb85c4e..86378c3 100644 --- a/examples/cluster.yaml +++ b/examples/cluster.yaml @@ -39,13 +39,15 @@ backend_env: nodes: - ip: "localhost" # IP address or hostname. user: "node1_username" # [OPTIONAL] Overrides `common_user` for this node. + dir: "path/to/project" # [OPTIONAL] Node-specific path to the project, overrides `common_dir`. type: "nvidia" # Architecture type. slots: 8 # Number of GPUs/Processes to run on this node. cmake_flags: "-D..." # [OPTIONAL] Node-specific CMake options (e.g., `-DUSE_CUDA=ON`). backend_env: # [OPTIONAL] Node-specific environment variables. NODE1_ENV_VAR1: "value1" + # Add more nodes as needed, following the same structure. - ip: "" - type: "metax" + type: "..." slots: 8 cmake_flags: "-D..."