From 8e3841f28993a845a503dacea22691b2dcb5ffec Mon Sep 17 00:00:00 2001 From: JYMiracle305 <604951424@qq.com> Date: Tue, 27 Oct 2026 09:12:57 +0000 Subject: [PATCH 1/4] feat: add SSH multinode launcher --- scripts/multinode_ssh/README.md | 104 ++++++++ scripts/multinode_ssh/launch_multinode_ssh.sh | 239 ++++++++++++++++++ .../test_launch_multinode_ssh.sh | 73 ++++++ 3 files changed, 416 insertions(+) create mode 100644 scripts/multinode_ssh/README.md create mode 100755 scripts/multinode_ssh/launch_multinode_ssh.sh create mode 100755 scripts/multinode_ssh/test_launch_multinode_ssh.sh diff --git a/scripts/multinode_ssh/README.md b/scripts/multinode_ssh/README.md new file mode 100644 index 00000000..8b1c1a4e --- /dev/null +++ b/scripts/multinode_ssh/README.md @@ -0,0 +1,104 @@ +# SSH Multinode Launcher + +This directory contains: + +- `launch_multinode_ssh.sh`: starts InfiniTrain on all nodes from one control node. +- `test_launch_multinode_ssh.sh`: tests launcher argument handling, rank assignment, + shell quoting, Docker command generation, and node startup delay with a fake + `infini_run`. It does not connect to a real cluster or use GPUs. + +## Prerequisites + +- The repository and training assets are available at the same paths in every + training container. +- The control node can use non-interactive SSH authentication for remote hosts. +- The SSH account can run `docker exec` on the remote host when + `--remote-docker-container` is used. +- `build/infini_run` and the selected training executable have already been built. +- `NCCL_SOCKET_IFNAME` is set to the physical interface that routes to the other + training nodes. Do not allow NCCL to select Docker `veth` interfaces. + +Check the interface with: + +```bash +ip route get 192.168.162.63 +``` + +## Two-node DP example + +The following example starts one GPU process on each node and runs one GPT-2 +training iteration. Update the variables for the target environment. + +```bash +CONTROL_HOST=192.168.162.62 +REMOTE_HOST=192.168.162.63 +REMOTE_CONTAINER=infinitrain +WORKDIR="$HOME/code/InfiniTrain" +SSH_KEY="$HOME/.ssh/id_ed25519" +NCCL_IFNAME=ens22f0 + +scripts/multinode_ssh/launch_multinode_ssh.sh \ + --hosts "$CONTROL_HOST,$REMOTE_HOST" \ + --local-host "$CONTROL_HOST" \ + --remote-docker-container "$REMOTE_CONTAINER" \ + --ssh-option "IdentityFile=$SSH_KEY" \ + --ssh-option IdentitiesOnly=yes \ + --workdir "$WORKDIR" \ + --nproc-per-node 1 \ + --master-addr "$CONTROL_HOST" \ + --node-start-delay 5 \ + --env "NCCL_SOCKET_IFNAME=$NCCL_IFNAME" \ + --env NCCL_IB_DISABLE=1 \ + -- ./build/gpt2 \ + --device cuda \ + --tensor_parallel 1 \ + --pipeline_parallel 1 \ + --batch_size 1 \ + --sequence_length 64 \ + --num_iteration 1 \ + --total_batch_size 128 \ + --input_bin /data1/shared/InfiniTrain-dev/data/llmc/gpt2/tinyshakespeare/tiny_shakespeare_train.bin \ + --tokenizer_bin /data1/shared/InfiniTrain-dev/data/llmc/gpt2/gpt2_tokenizer.bin \ + --llmc_filepath /data1/shared/InfiniTrain-dev/data/llmc/gpt2/gpt2_124M.bin +``` + +For pure data parallelism: + +```text +DP size = number of nodes x processes per node +TP size = 1 +PP size = 1 +``` + +`total_batch_size` must be divisible by: + +```text +batch_size x sequence_length x DP size +``` + +Increase `--nproc-per-node` to use more GPUs on each node. For two nodes with +four processes per node and the same batch and sequence lengths, use +`--total_batch_size 512` for one microbatch per optimizer step. + +The launcher generates a unique rendezvous ID when `--rdzv-id` is omitted. +Per-node logs are written to `logs//` unless `--log-dir` is supplied. + +## Inspect commands without launching + +```bash +scripts/multinode_ssh/launch_multinode_ssh.sh \ + --hosts 192.168.162.62,192.168.162.63 \ + --local-host 192.168.162.62 \ + --workdir "$HOME/code/InfiniTrain" \ + --dry-run \ + -- ./build/gpt2 --device cuda +``` + +## Launcher regression test + +```bash +scripts/multinode_ssh/test_launch_multinode_ssh.sh +``` + +This is a local simulated test. A passing result does not verify SSH, Docker, +NCCL, GPUs, model files, or training on the real cluster. diff --git a/scripts/multinode_ssh/launch_multinode_ssh.sh b/scripts/multinode_ssh/launch_multinode_ssh.sh new file mode 100755 index 00000000..7a67cbf8 --- /dev/null +++ b/scripts/multinode_ssh/launch_multinode_ssh.sh @@ -0,0 +1,239 @@ +#!/usr/bin/env bash + +set -euo pipefail + +usage() { + cat <<'EOF' +Launch InfiniTrain on multiple nodes from one control node. + +Usage: + scripts/multinode_ssh/launch_multinode_ssh.sh [launcher options] -- TRAIN_PROGRAM [ARGS...] + +Launcher options: + --hosts HOSTS Comma-separated hosts in node-rank order. + --hostfile PATH One host per line; blank lines and # comments ignored. + --workdir PATH Repository path on every node (default: current directory). + --nproc-per-node N Processes started by infini_run on each node (default: 1). + --master-addr HOST Rendezvous address (default: first host). + --master-port PORT Rendezvous port (default: 29500). + --rdzv-id ID Shared job ID (default: generated). + --infini-run PATH infini_run path relative to workdir (default: ./build/infini_run). + --ssh-user USER SSH user for remote nodes. + --ssh-port PORT SSH port (default: 22). + --remote-docker-container NAME + Run remote commands through docker exec NAME. + --ssh-option OPTION Additional ssh option; may be repeated. + --log-dir PATH Per-node log directory (default: ./logs/). + --local-host HOST Host entry to execute locally instead of through SSH. + --node-start-delay SEC Seconds to wait between starting adjacent nodes (default: 0). + --env NAME=VALUE Export an environment variable on every node; may be repeated. + --dry-run Print commands without starting them. + -h, --help Show this help. + +Example: + scripts/multinode_ssh/launch_multinode_ssh.sh \ + --hosts 192.168.162.62,192.168.162.63 \ + --local-host 192.168.162.62 \ + --workdir "$HOME/code/InfiniTrain" \ + --nproc-per-node 4 \ + --master-addr 192.168.162.62 \ + --env NCCL_SOCKET_IFNAME=ens22f0 \ + --env NCCL_IB_DISABLE=1 \ + -- ./build/gpt2 --device cuda --tensor_parallel 2 --pipeline_parallel 2 +EOF +} + +die() { + echo "error: $*" >&2 + exit 2 +} + +quote_command() { + local quoted=() + local arg + for arg in "$@"; do + printf -v arg '%q' "$arg" + quoted+=("$arg") + done + local IFS=' ' + printf '%s' "${quoted[*]}" +} + +hosts_csv="" +hostfile="" +workdir="$PWD" +nproc_per_node=1 +master_addr="" +master_port=29500 +rdzv_id="" +infini_run="./build/infini_run" +ssh_user="" +ssh_port=22 +remote_docker_container="" +log_dir="" +local_host="" +node_start_delay=0 +dry_run=0 +declare -a ssh_options=() +declare -a exports=() + +while (($#)); do + case "$1" in + --hosts) [[ $# -ge 2 ]] || die "--hosts requires a value"; hosts_csv=$2; shift 2 ;; + --hostfile) [[ $# -ge 2 ]] || die "--hostfile requires a value"; hostfile=$2; shift 2 ;; + --workdir) [[ $# -ge 2 ]] || die "--workdir requires a value"; workdir=$2; shift 2 ;; + --nproc-per-node) [[ $# -ge 2 ]] || die "--nproc-per-node requires a value"; nproc_per_node=$2; shift 2 ;; + --master-addr) [[ $# -ge 2 ]] || die "--master-addr requires a value"; master_addr=$2; shift 2 ;; + --master-port) [[ $# -ge 2 ]] || die "--master-port requires a value"; master_port=$2; shift 2 ;; + --rdzv-id) [[ $# -ge 2 ]] || die "--rdzv-id requires a value"; rdzv_id=$2; shift 2 ;; + --infini-run) [[ $# -ge 2 ]] || die "--infini-run requires a value"; infini_run=$2; shift 2 ;; + --ssh-user) [[ $# -ge 2 ]] || die "--ssh-user requires a value"; ssh_user=$2; shift 2 ;; + --ssh-port) [[ $# -ge 2 ]] || die "--ssh-port requires a value"; ssh_port=$2; shift 2 ;; + --remote-docker-container) [[ $# -ge 2 ]] || die "--remote-docker-container requires a value"; remote_docker_container=$2; shift 2 ;; + --ssh-option) [[ $# -ge 2 ]] || die "--ssh-option requires a value"; ssh_options+=("$2"); shift 2 ;; + --log-dir) [[ $# -ge 2 ]] || die "--log-dir requires a value"; log_dir=$2; shift 2 ;; + --local-host) [[ $# -ge 2 ]] || die "--local-host requires a value"; local_host=$2; shift 2 ;; + --node-start-delay) [[ $# -ge 2 ]] || die "--node-start-delay requires a value"; node_start_delay=$2; shift 2 ;; + --env) [[ $# -ge 2 ]] || die "--env requires NAME=VALUE"; exports+=("$2"); shift 2 ;; + --dry-run) dry_run=1; shift ;; + -h|--help) usage; exit 0 ;; + --) shift; break ;; + *) die "unknown launcher option: $1" ;; + esac +done + +[[ $# -gt 0 ]] || die "missing training program after --" +[[ -z $hosts_csv || -z $hostfile ]] || die "use only one of --hosts and --hostfile" +[[ -n $hosts_csv || -n $hostfile ]] || die "--hosts or --hostfile is required" +[[ $nproc_per_node =~ ^[1-9][0-9]*$ ]] || die "--nproc-per-node must be positive" +[[ $master_port =~ ^[1-9][0-9]*$ ]] || die "--master-port must be positive" +[[ $ssh_port =~ ^[1-9][0-9]*$ ]] || die "--ssh-port must be positive" +[[ $node_start_delay =~ ^[0-9]+([.][0-9]+)?$ ]] || die "--node-start-delay must be non-negative" + +declare -a hosts=() +if [[ -n $hosts_csv ]]; then + IFS=',' read -r -a hosts <<<"$hosts_csv" +else + [[ -r $hostfile ]] || die "cannot read hostfile: $hostfile" + while IFS= read -r line || [[ -n $line ]]; do + line=${line%%#*} + read -r line <<<"$line" + [[ -n $line ]] && hosts+=("$line") + done <"$hostfile" +fi +[[ ${#hosts[@]} -gt 0 ]] || die "host list is empty" + +for entry in "${exports[@]}"; do + [[ $entry == *=* && $entry != =* ]] || die "invalid --env value: $entry" +done + +master_addr=${master_addr:-${hosts[0]}} +rdzv_id=${rdzv_id:-"infinitrain-$(date +%Y%m%d-%H%M%S)-$$"} +log_dir=${log_dir:-"$PWD/logs/$rdzv_id"} +mkdir -p "$log_dir" + +train_command=("$@") +nnodes=${#hosts[@]} +declare -a child_pids=() +declare -a child_nodes=() +stopping=0 + +stop_children() { + ((stopping == 0)) || return + stopping=1 + local pid + for pid in "${child_pids[@]:-}"; do + kill -TERM -- "-$pid" 2>/dev/null || true + done +} +trap 'stop_children; exit 130' INT TERM + +for node_rank in "${!hosts[@]}"; do + host=${hosts[$node_rank]} + node_command=( + "$infini_run" + "--nnodes=$nnodes" + "--nproc_per_node=$nproc_per_node" + "--node_rank=$node_rank" + "--rdzv_endpoint=$master_addr:$master_port" + "--rdzv_id=$rdzv_id" + "${train_command[@]}" + ) + + printf -v quoted_workdir '%q' "$workdir" + command_text="cd $quoted_workdir && " + for entry in "${exports[@]}"; do + printf -v quoted_export '%q' "$entry" + command_text+="export $quoted_export; " + done + command_text+="exec $(quote_command "${node_command[@]}")" + + log_file="$log_dir/node-${node_rank}-${host//[^a-zA-Z0-9_.-]/_}.log" + if [[ -n $local_host && $host == "$local_host" ]]; then + launch=(bash -lc "$command_text") + else + destination=$host + [[ -z $ssh_user ]] || destination="$ssh_user@$host" + launch=(ssh -o BatchMode=yes -o ConnectTimeout=10 -p "$ssh_port") + for option in "${ssh_options[@]}"; do + launch+=(-o "$option") + done + if [[ -n $remote_docker_container ]]; then + docker_command=(docker exec "$remote_docker_container" bash -lc "$command_text") + launch+=("$destination" "$(quote_command "${docker_command[@]}")") + else + launch+=("$destination" "$command_text") + fi + fi + + launch=(setsid "${launch[@]}") + + printf '[node %d %s] ' "$node_rank" "$host" + quote_command "${launch[@]}" + printf ' > %q 2>&1\n' "$log_file" + if ((dry_run)); then + continue + fi + + "${launch[@]}" >"$log_file" 2>&1 & + child_pids+=("$!") + child_nodes+=("$host") + + if ((node_rank + 1 < nnodes)) && [[ $node_start_delay != 0 ]]; then + sleep "$node_start_delay" + fi +done + +((dry_run)) && exit 0 + +status=0 +declare -A node_by_pid=() +declare -A running_pids=() +for index in "${!child_pids[@]}"; do + pid=${child_pids[$index]} + node_by_pid[$pid]=${child_nodes[$index]} + running_pids[$pid]=1 +done + +while ((${#running_pids[@]} > 0)); do + finished_pid="" + if wait -n -p finished_pid; then + rc=0 + else + rc=$? + fi + [[ -n ${finished_pid:-} ]] || die "failed to identify completed node launcher" + unset 'running_pids[$finished_pid]' + if ((rc != 0)); then + echo "node ${node_by_pid[$finished_pid]} failed (launcher status $rc)" >&2 + status=1 + stop_children + for pid in "${!running_pids[@]}"; do + wait "$pid" 2>/dev/null || true + done + break + fi +done + +trap - INT TERM +exit "$status" diff --git a/scripts/multinode_ssh/test_launch_multinode_ssh.sh b/scripts/multinode_ssh/test_launch_multinode_ssh.sh new file mode 100755 index 00000000..27a8b612 --- /dev/null +++ b/scripts/multinode_ssh/test_launch_multinode_ssh.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash + +set -euo pipefail + +script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +launcher="$script_dir/launch_multinode_ssh.sh" +tmpdir=$(mktemp -d) +trap 'rm -rf "$tmpdir"' EXIT + +cat >"$tmpdir/fake_infini_run" <<'EOF' +#!/usr/bin/env bash +set -euo pipefail +printf '%s\n' "$*" +EOF +chmod +x "$tmpdir/fake_infini_run" + +output=$( + "$launcher" \ + --hosts local.test,remote.test \ + --local-host local.test \ + --workdir "$tmpdir" \ + --infini-run ./fake_infini_run \ + --nproc-per-node 2 \ + --master-addr 10.0.0.1 \ + --rdzv-id test-job \ + --env NCCL_DEBUG=INFO \ + --remote-docker-container test-container \ + --dry-run \ + -- ./train --flag 'value with spaces' +) + +grep -q -- '--node_rank=0' <<<"$output" +grep -q -- '--node_rank=1' <<<"$output" +grep -q -- '--nnodes=2' <<<"$output" +grep -q -- '--nproc_per_node=2' <<<"$output" +grep -q -- '--rdzv_endpoint=10.0.0.1:29500' <<<"$output" +grep -q -- '--rdzv_id=test-job' <<<"$output" +grep -q -- 'local.test' <<<"$output" +grep -q -- 'remote.test' <<<"$output" +grep -q -- 'docker.*exec.*test-container' <<<"$output" + +"$launcher" \ + --hosts local.test \ + --local-host local.test \ + --workdir "$tmpdir" \ + --infini-run ./fake_infini_run \ + --rdzv-id local-test \ + --log-dir "$tmpdir/logs" \ + -- ./train --flag value + +grep -q -- '--node_rank=0' "$tmpdir/logs/node-0-local.test.log" +grep -q -- './train --flag value' "$tmpdir/logs/node-0-local.test.log" + +SECONDS=0 +"$launcher" \ + --hosts local.test,local.test \ + --local-host local.test \ + --workdir "$tmpdir" \ + --infini-run ./fake_infini_run \ + --rdzv-id delayed-local-test \ + --node-start-delay 1 \ + --log-dir "$tmpdir/delayed-logs" \ + -- ./train +((SECONDS >= 1)) + +if "$launcher" --hosts local.test --node-start-delay invalid -- ./train \ + >"$tmpdir/invalid-delay.out" 2>&1; then + echo "invalid node start delay unexpectedly succeeded" >&2 + exit 1 +fi +grep -q -- '--node-start-delay must be non-negative' "$tmpdir/invalid-delay.out" + +echo "launch_multinode_ssh tests passed" From a72d70de55b96e203db97d533ede5511c637c55c Mon Sep 17 00:00:00 2001 From: JYMiracle305 <604951424@qq.com> Date: Mon, 2 Nov 2026 08:49:56 +0000 Subject: [PATCH 2/4] feat: validate NVIDIA SSH multinode training --- infini_train/src/core/ccl/ccl_utils.cc | 156 ++++++- scripts/multinode_ssh/README.md | 146 ++++--- scripts/multinode_ssh/VALIDATION_REPORT.md | 68 ++++ scripts/multinode_ssh/launch_multinode_ssh.sh | 375 ++++++++--------- .../test_launch_multinode_ssh.sh | 172 +++++--- tools/infini_run/infini_run.cc | 385 ++++++++++++------ 6 files changed, 816 insertions(+), 486 deletions(-) create mode 100644 scripts/multinode_ssh/VALIDATION_REPORT.md diff --git a/infini_train/src/core/ccl/ccl_utils.cc b/infini_train/src/core/ccl/ccl_utils.cc index d5df75df..f0d7f666 100644 --- a/infini_train/src/core/ccl/ccl_utils.cc +++ b/infini_train/src/core/ccl/ccl_utils.cc @@ -1,19 +1,128 @@ #include "infini_train/include/core/ccl/ccl_utils.h" +#include +#include +#include + +#include #include #include #include #include #include #include +#include +#include #include #include "glog/logging.h" #include "infini_train/include/core/ccl/ccl_common.h" +#include + namespace infini_train::core { namespace { +constexpr uint32_t kRendezvousMagic = 0x49545244; +struct RendezvousHeader { + uint32_t magic, operation, key_size, value_size; +}; + +int GetEnvInt(const char *name, int fallback) { + const char *value = std::getenv(name); + return value == nullptr ? fallback : std::stoi(value); +} + +bool UseTcpRendezvous() { + return GetEnvInt("WORLD_SIZE", 1) > GetEnvInt("LOCAL_WORLD_SIZE", 1); +} + +bool ReadAll(int fd, void *buffer, size_t size) { + auto *bytes = static_cast(buffer); + while (size > 0) { + const ssize_t count = read(fd, bytes, size); + if (count < 0 && errno == EINTR) continue; + if (count <= 0) return false; + bytes += count; + size -= static_cast(count); + } + return true; +} + +bool WriteAll(int fd, const void *buffer, size_t size) { + const auto *bytes = static_cast(buffer); + while (size > 0) { + const ssize_t count = write(fd, bytes, size); + if (count < 0 && errno == EINTR) continue; + if (count <= 0) return false; + bytes += count; + size -= static_cast(count); + } + return true; +} + +int ConnectRendezvous() { + const char *host = std::getenv("MASTER_ADDR"); + const char *port = std::getenv("MASTER_PORT"); + CHECK(host != nullptr && port != nullptr) + << "MASTER_ADDR and MASTER_PORT are required for multi-node rendezvous"; + addrinfo hints{}; + hints.ai_family = AF_UNSPEC; + hints.ai_socktype = SOCK_STREAM; + addrinfo *addresses = nullptr; + CHECK_EQ(getaddrinfo(host, port, &hints, &addresses), 0) + << "Failed to resolve rendezvous endpoint " << host << ':' << port; + const int timeout = GetEnvInt("INFINI_RDZV_TIMEOUT_SEC", 120); + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(timeout); + int fd = -1; + while (std::chrono::steady_clock::now() < deadline && fd < 0) { + for (addrinfo *address = addresses; address != nullptr; + address = address->ai_next) { + const int candidate = socket(address->ai_family, address->ai_socktype, + address->ai_protocol); + if (candidate < 0) continue; + if (connect(candidate, address->ai_addr, address->ai_addrlen) == 0) { + fd = candidate; + break; + } + close(candidate); + } + if (fd < 0) std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + freeaddrinfo(addresses); + CHECK_GE(fd, 0) << "Timed out connecting to rendezvous " << host << ':' + << port; + return fd; +} + +std::string RendezvousRequest(uint32_t operation, const std::string &key, + std::string_view value) { + CHECK_LE(key.size(), std::numeric_limits::max()); + CHECK_LE(value.size(), std::numeric_limits::max()); + const int fd = ConnectRendezvous(); + const RendezvousHeader header{htonl(kRendezvousMagic), htonl(operation), + htonl(static_cast(key.size())), + htonl(static_cast(value.size()))}; + CHECK(WriteAll(fd, &header, sizeof(header)) && + WriteAll(fd, key.data(), key.size()) && + (value.empty() || WriteAll(fd, value.data(), value.size()))) + << "Failed to send rendezvous request for " << key; + uint32_t response[2] = {}; + CHECK(ReadAll(fd, response, sizeof(response))); + CHECK_EQ(ntohl(response[0]), 0) << "Rendezvous request failed for " << key; + std::string result(ntohl(response[1]), '\0'); + CHECK(result.empty() || ReadAll(fd, result.data(), result.size())); + close(fd); + return result; +} + +std::string RendezvousKey(const std::string &pg_name) { + const char *run_id = std::getenv("INFINI_RUN_ID"); + CHECK(run_id != nullptr) << "INFINI_RUN_ID is required"; + return std::string(run_id) + "/ccl/" + pg_name; +} + std::string UniqueIdFileName(const std::string &pg_name) { const char *run_id = std::getenv("INFINI_RUN_ID"); const std::string prefix = run_id == nullptr ? "" : std::string(run_id) + "_"; @@ -29,16 +138,20 @@ std::string UniqueIdTmpFileName(const std::string &pg_name) { void WriteUniqueIdFile(const CclUniqueId &unique_id, const std::string &pg_name) { + if (UseTcpRendezvous()) { + RendezvousRequest( + 1, RendezvousKey(pg_name), + std::string_view(static_cast(unique_id.Data()), + unique_id.Size())); + return; + } const std::string tmp_path = UniqueIdTmpFileName(pg_name); - std::ofstream ofs(tmp_path, std::ios::binary); CHECK(ofs.good()) << "Failed to open unique_id tmp file for write: " << tmp_path; - const size_t size = unique_id.Size(); ofs.write(reinterpret_cast(unique_id.Data()), - static_cast(size)); + static_cast(unique_id.Size())); ofs.close(); - const std::string file_path = UniqueIdFileName(pg_name); CHECK_EQ(std::rename(tmp_path.c_str(), file_path.c_str()), 0) << "Failed to rename unique_id file from " << tmp_path << " to " @@ -47,35 +160,34 @@ void WriteUniqueIdFile(const CclUniqueId &unique_id, void ReadUniqueIdFile(CclUniqueId *unique_id, const std::string &pg_name) { CHECK_NOTNULL(unique_id); + if (UseTcpRendezvous()) { + const std::string bytes = RendezvousRequest(2, RendezvousKey(pg_name), {}); + CHECK_EQ(bytes.size(), unique_id->Size()); + unique_id->Load(bytes.data(), bytes.size()); + return; + } const std::string file_path = UniqueIdFileName(pg_name); - + const int timeout = GetEnvInt("INFINI_RDZV_TIMEOUT_SEC", 120); + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(timeout); while (!std::filesystem::exists(file_path)) { - std::this_thread::sleep_for(std::chrono::microseconds(1000)); + CHECK(std::chrono::steady_clock::now() < deadline) + << "Timed out waiting for unique-id file " << file_path; + std::this_thread::sleep_for(std::chrono::milliseconds(1)); } - std::ifstream ifs(file_path, std::ios::binary); CHECK(ifs.good()) << "Failed to open unique_id file for read: " << file_path; - std::string bytes((std::istreambuf_iterator(ifs)), std::istreambuf_iterator()); - ifs.close(); - - CHECK_EQ(bytes.size(), unique_id->Size()) - << "Mismatched unique_id size in file. expected=" << unique_id->Size() - << ", got=" << bytes.size(); + CHECK_EQ(bytes.size(), unique_id->Size()); unique_id->Load(bytes.data(), bytes.size()); } void CleanupUniqueIdFile(const std::string &pg_name) { - const std::string file_path = UniqueIdFileName(pg_name); - if (std::filesystem::exists(file_path)) { - std::filesystem::remove(file_path); - } - - const std::string tmp_path = UniqueIdTmpFileName(pg_name); - if (std::filesystem::exists(tmp_path)) { - std::filesystem::remove(tmp_path); + if (UseTcpRendezvous()) return; + for (const auto &path : + {UniqueIdFileName(pg_name), UniqueIdTmpFileName(pg_name)}) { + if (std::filesystem::exists(path)) std::filesystem::remove(path); } } - } // namespace infini_train::core diff --git a/scripts/multinode_ssh/README.md b/scripts/multinode_ssh/README.md index 8b1c1a4e..e9fbb8a6 100644 --- a/scripts/multinode_ssh/README.md +++ b/scripts/multinode_ssh/README.md @@ -1,104 +1,118 @@ -# SSH Multinode Launcher +# NVIDIA SSH Multinode Launcher -This directory contains: +`launch_multinode_ssh.sh` starts one `infini_run` per node from a control host. +A two-node, eight-GPU-per-node job creates 16 rank processes, records separate +rank logs, and returns nonzero if either node fails. -- `launch_multinode_ssh.sh`: starts InfiniTrain on all nodes from one control node. -- `test_launch_multinode_ssh.sh`: tests launcher argument handling, rank assignment, - shell quoting, Docker command generation, and node startup delay with a fake - `infini_run`. It does not connect to a real cluster or use GPUs. +## Preparation -## Prerequisites +- Use the same repository commit and matching binaries on every node. +- Ensure non-interactive SSH and known-host verification work before launch. +- The remote account must be able to run the build, and optionally `docker + exec`, without relying on an interactive shell profile. +- Select an unused TCP port reachable from every node. Node rank 0 hosts the + rendezvous service; multi-node CCL unique-ID exchange does not require NFS. +- Set `NCCL_SOCKET_IFNAME` from the actual route. Confirm RDMA from NCCL logs; + successful communication alone is not proof that RDMA was used. -- The repository and training assets are available at the same paths in every - training container. -- The control node can use non-interactive SSH authentication for remote hosts. -- The SSH account can run `docker exec` on the remote host when - `--remote-docker-container` is used. -- `build/infini_run` and the selected training executable have already been built. -- `NCCL_SOCKET_IFNAME` is set to the physical interface that routes to the other - training nodes. Do not allow NCCL to select Docker `veth` interfaces. +A hostfile is plain text in node-rank order; blank lines and `#` comments are +ignored: -Check the interface with: - -```bash -ip route get 192.168.162.63 +```text +# node rank order +TRAIN_NODE_0 +TRAIN_NODE_1 ``` -## Two-node DP example +For a short list, `--hosts TRAIN_NODE_0,TRAIN_NODE_1` avoids maintaining a +separate file. + +## Two Nodes, Eight GPUs Each -The following example starts one GPU process on each node and runs one GPT-2 -training iteration. Update the variables for the target environment. +Run only on the control node. The second node needs no manual command. Set the +following environment-specific values first. `CONTROL_HOST` must be reachable by +every training node, and `NCCL_IFNAME` must be the interface selected by the route +to the peer. ```bash -CONTROL_HOST=192.168.162.62 -REMOTE_HOST=192.168.162.63 -REMOTE_CONTAINER=infinitrain -WORKDIR="$HOME/code/InfiniTrain" -SSH_KEY="$HOME/.ssh/id_ed25519" -NCCL_IFNAME=ens22f0 +CONTROL_HOST=TRAIN_NODE_0 +REMOTE_HOST=TRAIN_NODE_1 +SSH_USER=TRAIN_USER +SSH_KEY=/path/to/private_key +REMOTE_CONTAINER=TRAINING_CONTAINER +WORKDIR=/path/to/InfiniTrain +NCCL_IFNAME=TRAINING_INTERFACE +LOG_ROOT=/path/to/controller/logs +REMOTE_LOG_ROOT=/path/to/node/logs + scripts/multinode_ssh/launch_multinode_ssh.sh \ --hosts "$CONTROL_HOST,$REMOTE_HOST" \ --local-host "$CONTROL_HOST" \ - --remote-docker-container "$REMOTE_CONTAINER" \ + --ssh-user "$SSH_USER" \ --ssh-option "IdentityFile=$SSH_KEY" \ - --ssh-option IdentitiesOnly=yes \ + --remote-docker-container "$REMOTE_CONTAINER" \ --workdir "$WORKDIR" \ - --nproc-per-node 1 \ + --infini-run ./build/infini_run \ + --nproc-per-node 8 \ --master-addr "$CONTROL_HOST" \ - --node-start-delay 5 \ + --master-port 29500 \ + --rdzv-timeout 180 \ + --job-timeout 1800 \ + --log-root "$LOG_ROOT" \ + --remote-log-root "$REMOTE_LOG_ROOT" \ --env "NCCL_SOCKET_IFNAME=$NCCL_IFNAME" \ - --env NCCL_IB_DISABLE=1 \ + --env NCCL_DEBUG=INFO \ + --follow-output \ -- ./build/gpt2 \ --device cuda \ - --tensor_parallel 1 \ - --pipeline_parallel 1 \ + --tensor_parallel 4 \ + --pipeline_parallel 2 \ --batch_size 1 \ --sequence_length 64 \ - --num_iteration 1 \ - --total_batch_size 128 \ - --input_bin /data1/shared/InfiniTrain-dev/data/llmc/gpt2/tinyshakespeare/tiny_shakespeare_train.bin \ - --tokenizer_bin /data1/shared/InfiniTrain-dev/data/llmc/gpt2/gpt2_tokenizer.bin \ - --llmc_filepath /data1/shared/InfiniTrain-dev/data/llmc/gpt2/gpt2_124M.bin + --num_iteration 20 \ + --total_batch_size 16 \ + --input_bin /path/to/gpt2_train.bin \ + --llmc_filepath /path/to/gpt2_124M.bin ``` -For pure data parallelism: +The tested PP2 command intentionally omits `--tokenizer_bin`: current text +sampling asserts that pipeline parallel size is 1 when generation runs. -```text -DP size = number of nodes x processes per node -TP size = 1 -PP size = 1 -``` +`--remote-docker-container` applies to SSH nodes. A host selected by +`--local-host` runs in the controller's current environment, so launch the +controller from the correct container when node 0 also uses a container. -`total_batch_size` must be divisible by: +## Logs and Failure Handling + +Preflight checks the work directory, executable files, Git SHA, requested data +paths, binary SHA-256, `nvidia-smi`, and visible GPU count. The control log +contains metadata, the complete invocation, preflight output, and one SSH +wrapper log per node. Rank output is stored at: ```text -batch_size x sequence_length x DP size +REMOTE_LOG_ROOT/JOB_ID/node-N/rank-R.stdout.log +REMOTE_LOG_ROOT/JOB_ID/node-N/rank-R.stderr.log ``` -Increase `--nproc-per-node` to use more GPUs on each node. For two nodes with -four processes per node and the same batch and sequence lengths, use -`--total_batch_size 512` for one microbatch per optimizer step. - -The launcher generates a unique rendezvous ID when `--rdzv-id` is omitted. -Per-node logs are written to `logs//` unless `--log-dir` is supplied. +`--follow-output` tails the final global rank's stderr on the controller while +preserving the original file. Training loss is normally emitted on stderr. +Every rank first prints hostname, global rank, local rank, world size, device, +and job ID. -## Inspect commands without launching +The launcher does not restart failed jobs. On node failure, SSH loss, timeout, +SIGINT, or SIGTERM, it reads this job's PID file, verifies the command contains +the job ID, and terminates only that process group. It never uses broad `pkill`. -```bash -scripts/multinode_ssh/launch_multinode_ssh.sh \ - --hosts 192.168.162.62,192.168.162.63 \ - --local-host 192.168.162.62 \ - --workdir "$HOME/code/InfiniTrain" \ - --dry-run \ - -- ./build/gpt2 --device cuda -``` +## Validation -## Launcher regression test +Run the local regression first: ```bash scripts/multinode_ssh/test_launch_multinode_ssh.sh ``` -This is a local simulated test. A passing result does not verify SSH, Docker, -NCCL, GPUs, model files, or training on the real cluster. +Then expand in order: single-node 8 ranks, two-node 1 rank each, two-node 8 +ranks each, short training, numerical comparison, and failure injection. Keep +`NCCL_DEBUG=INFO` for communication-path evidence and retain generated job +folders. Local simulation must not be reported as GPU or multi-node validation. diff --git a/scripts/multinode_ssh/VALIDATION_REPORT.md b/scripts/multinode_ssh/VALIDATION_REPORT.md new file mode 100644 index 00000000..a6f4d1e7 --- /dev/null +++ b/scripts/multinode_ssh/VALIDATION_REPORT.md @@ -0,0 +1,68 @@ +# NVIDIA Two-Node Validation Report + +## Code and Environment + +- Date: 2026-11-02 +- Base: `origin/feat/ssh-multinode-launcher` at + `924d21256f3baf66bccb7458aa77395f0fbf5841` +- Observed master: `5a7be78e9b9d42c2f26efe33984dd3167dc60dc9` +- Development branch: `codex/ssh-multinode-nvidia-validation` +- Nodes: two private training hosts (identifiers intentionally omitted) +- GPUs: 8 x NVIDIA A100-SXM4-80GB per node; driver 535.161.08 +- SSH identity, host addresses, container name, worktree path, and interface: + intentionally omitted; configure them as described in `README.md` + +The branch remains based directly on the requested feature branch. Master had +broad unrelated changes, so it was not merged or rebased. + +## Results + +| Area | Result | Evidence | +|---|---|---| +| Parser, quoting, hostfile, invalid config | Passed | local launcher regression | +| Simulated success, node failure, timeout, SIGTERM cleanup | Passed | exits 0, 17, 124, 130; job PID checked | +| Live controller output | Passed | `--follow-output` observed start and final lines | +| TCP rendezvous PUT/GET | Passed | local protocol test against node-0 server | +| C++ build and rank tests | Passed | full build; rank CTest 2/2 | +| Two-node 1 process per node | Passed | real SSH smoke | +| Two-node 8 processes per node | Passed | 16 unique global ranks; local devices 0-7 | +| GPT2 TP4 x PP2 x DP2, 20 steps | Passed | final loss 4.468029 | +| Qwen3 TP4 x PP2 x DP2, 20 steps, BF16 | Passed | final loss 2.951596 | +| Llama3 TP4 x PP2 x DP2, 20 steps, BF16 | Passed | final loss 3.553573 | +| NCCL communication path | Passed for observed run | NCCL log reported `NET/IBext_v10` and GDRDMA | +| Real-node bad SSH / bad path | Partially verified | preflight/path failures verified; full matrix not retained | +| Real GPU rank failure / initialization timeout / Ctrl-C | Not verified | only simulated failure paths completed | +| Performance baseline | Not verified | no throughput target; correctness was prioritized | + +## Numerical Comparison + +Single-node references used TP4 x PP2 x DP1 with gradient accumulation adjusted +to preserve the global batch. Two-node runs used TP4 x PP2 x DP2. Initial +weights, data order, optimizer, global batch, and 20-step count were held fixed. + +| Model | Single-node final | Two-node final | Maximum absolute loss difference | +|---|---:|---:|---:| +| GPT2 | 4.468028 | 4.468029 | 0.000001 | +| Qwen3 BF16 | 2.950893 | 2.951596 | 0.004184 | +| Llama3 BF16 | 3.553813 | 3.553573 | 0.002242 | + +Maximum relative differences were 0.000022% (GPT2), 0.146558% (Qwen3), and +0.069305% (Llama3). Qwen3 and Llama3 use BF16 and different reduction order / +gradient accumulation between DP1 and DP2; these figures are recorded as +observations, not as a newly invented project tolerance. + +## Logs and Known Limits + +Control and rank log roots are operator-configured; their concrete validation +locations are intentionally omitted. The exact command is retained +in `controller-command.txt`, and runtime/version information is in metadata and +preflight logs. + +- PP2 text generation with `--tokenizer_bin` currently aborts when generation + is triggered because that path requires pipeline parallel size 1. Validated + PP2 commands omit tokenizer-based generation. +- The rendezvous service binds the selected node-0 port on all interfaces and + assumes a trusted training network; it is not an authenticated service. +- GPU preflight checks visibility/count, not exclusivity or free memory. +- More-node scaling still needs rendezvous load tests, larger failure matrices, + and a cluster log-retention policy. diff --git a/scripts/multinode_ssh/launch_multinode_ssh.sh b/scripts/multinode_ssh/launch_multinode_ssh.sh index 7a67cbf8..51a14a38 100755 --- a/scripts/multinode_ssh/launch_multinode_ssh.sh +++ b/scripts/multinode_ssh/launch_multinode_ssh.sh @@ -1,239 +1,194 @@ #!/usr/bin/env bash - set -euo pipefail -usage() { - cat <<'EOF' -Launch InfiniTrain on multiple nodes from one control node. - -Usage: - scripts/multinode_ssh/launch_multinode_ssh.sh [launcher options] -- TRAIN_PROGRAM [ARGS...] - -Launcher options: - --hosts HOSTS Comma-separated hosts in node-rank order. - --hostfile PATH One host per line; blank lines and # comments ignored. - --workdir PATH Repository path on every node (default: current directory). - --nproc-per-node N Processes started by infini_run on each node (default: 1). - --master-addr HOST Rendezvous address (default: first host). - --master-port PORT Rendezvous port (default: 29500). - --rdzv-id ID Shared job ID (default: generated). - --infini-run PATH infini_run path relative to workdir (default: ./build/infini_run). - --ssh-user USER SSH user for remote nodes. - --ssh-port PORT SSH port (default: 22). - --remote-docker-container NAME - Run remote commands through docker exec NAME. - --ssh-option OPTION Additional ssh option; may be repeated. - --log-dir PATH Per-node log directory (default: ./logs/). - --local-host HOST Host entry to execute locally instead of through SSH. - --node-start-delay SEC Seconds to wait between starting adjacent nodes (default: 0). - --env NAME=VALUE Export an environment variable on every node; may be repeated. - --dry-run Print commands without starting them. - -h, --help Show this help. - -Example: - scripts/multinode_ssh/launch_multinode_ssh.sh \ - --hosts 192.168.162.62,192.168.162.63 \ - --local-host 192.168.162.62 \ - --workdir "$HOME/code/InfiniTrain" \ - --nproc-per-node 4 \ - --master-addr 192.168.162.62 \ - --env NCCL_SOCKET_IFNAME=ens22f0 \ - --env NCCL_IB_DISABLE=1 \ - -- ./build/gpt2 --device cuda --tensor_parallel 2 --pipeline_parallel 2 -EOF -} - -die() { - echo "error: $*" >&2 - exit 2 -} - -quote_command() { - local quoted=() - local arg - for arg in "$@"; do - printf -v arg '%q' "$arg" - quoted+=("$arg") - done - local IFS=' ' - printf '%s' "${quoted[*]}" -} - -hosts_csv="" -hostfile="" -workdir="$PWD" -nproc_per_node=1 -master_addr="" -master_port=29500 -rdzv_id="" -infini_run="./build/infini_run" -ssh_user="" -ssh_port=22 -remote_docker_container="" -log_dir="" -local_host="" -node_start_delay=0 -dry_run=0 -declare -a ssh_options=() -declare -a exports=() - +original_argv=("$@") + +usage() { sed -n '2,38p' "$0" | sed 's/^# //'; } +# Launch InfiniTrain on multiple nodes from one control node. +# Usage: launch_multinode_ssh.sh [options] -- TRAIN_PROGRAM [ARGS...] +# --hosts HOSTS | --hostfile PATH +# --workdir PATH Repository path on every node. +# --nproc-per-node N Default: 1. +# --master-addr HOST Default: first host. +# --master-port PORT Default: 29500. +# --rdzv-id ID Default: generated unique ID. +# --rdzv-timeout SEC Default: 120. +# --job-timeout SEC Default: 0 (disabled). +# --infini-run PATH Default: ./build/infini_run. +# --ssh-user USER --ssh-port PORT +# --ssh-option OPTION Repeatable; host checking remains enabled. +# --local-host HOST Execute matching host locally. +# --remote-docker-container NAME +# --log-root PATH Default: ./logs. +# --remote-log-root PATH Default: WORKDIR/logs. +# --env NAME=VALUE Repeatable. +# --check-path PATH Repeatable readable data/weight check. +# --skip-gpu-check For CPU/simulated tests. +# --follow-output Stream the last global rank stderr. +# --dry-run + +die() { echo "error: $*" >&2; exit 2; } +quote_command() { local result= arg; for arg in "$@"; do printf -v arg '%q' "$arg"; result+="${result:+ }$arg"; done; printf '%s' "$result"; } + +hosts_csv= hostfile= workdir=$PWD nproc=1 master_addr= master_port=29500 +rdzv_id= rdzv_timeout=120 job_timeout=0 infini_run=./build/infini_run +ssh_user= ssh_port=22 container= local_host= log_root="$PWD/logs" remote_log_root= +dry_run=0 skip_gpu=0 follow_output=0 +declare -a ssh_options=() exports=() check_paths=() while (($#)); do case "$1" in - --hosts) [[ $# -ge 2 ]] || die "--hosts requires a value"; hosts_csv=$2; shift 2 ;; - --hostfile) [[ $# -ge 2 ]] || die "--hostfile requires a value"; hostfile=$2; shift 2 ;; - --workdir) [[ $# -ge 2 ]] || die "--workdir requires a value"; workdir=$2; shift 2 ;; - --nproc-per-node) [[ $# -ge 2 ]] || die "--nproc-per-node requires a value"; nproc_per_node=$2; shift 2 ;; - --master-addr) [[ $# -ge 2 ]] || die "--master-addr requires a value"; master_addr=$2; shift 2 ;; - --master-port) [[ $# -ge 2 ]] || die "--master-port requires a value"; master_port=$2; shift 2 ;; - --rdzv-id) [[ $# -ge 2 ]] || die "--rdzv-id requires a value"; rdzv_id=$2; shift 2 ;; - --infini-run) [[ $# -ge 2 ]] || die "--infini-run requires a value"; infini_run=$2; shift 2 ;; - --ssh-user) [[ $# -ge 2 ]] || die "--ssh-user requires a value"; ssh_user=$2; shift 2 ;; - --ssh-port) [[ $# -ge 2 ]] || die "--ssh-port requires a value"; ssh_port=$2; shift 2 ;; - --remote-docker-container) [[ $# -ge 2 ]] || die "--remote-docker-container requires a value"; remote_docker_container=$2; shift 2 ;; - --ssh-option) [[ $# -ge 2 ]] || die "--ssh-option requires a value"; ssh_options+=("$2"); shift 2 ;; - --log-dir) [[ $# -ge 2 ]] || die "--log-dir requires a value"; log_dir=$2; shift 2 ;; - --local-host) [[ $# -ge 2 ]] || die "--local-host requires a value"; local_host=$2; shift 2 ;; - --node-start-delay) [[ $# -ge 2 ]] || die "--node-start-delay requires a value"; node_start_delay=$2; shift 2 ;; - --env) [[ $# -ge 2 ]] || die "--env requires NAME=VALUE"; exports+=("$2"); shift 2 ;; - --dry-run) dry_run=1; shift ;; - -h|--help) usage; exit 0 ;; - --) shift; break ;; - *) die "unknown launcher option: $1" ;; + --hosts) hosts_csv=${2:?}; shift 2;; --hostfile) hostfile=${2:?}; shift 2;; + --workdir) workdir=${2:?}; shift 2;; --nproc-per-node) nproc=${2:?}; shift 2;; + --master-addr) master_addr=${2:?}; shift 2;; --master-port) master_port=${2:?}; shift 2;; + --rdzv-id) rdzv_id=${2:?}; shift 2;; --rdzv-timeout) rdzv_timeout=${2:?}; shift 2;; + --job-timeout) job_timeout=${2:?}; shift 2;; --infini-run) infini_run=${2:?}; shift 2;; + --ssh-user) ssh_user=${2:?}; shift 2;; --ssh-port) ssh_port=${2:?}; shift 2;; + --ssh-option) ssh_options+=("${2:?}"); shift 2;; + --remote-docker-container) container=${2:?}; shift 2;; --local-host) local_host=${2:?}; shift 2;; + --log-root) log_root=${2:?}; shift 2;; --remote-log-root) remote_log_root=${2:?}; shift 2;; + --env) exports+=("${2:?}"); shift 2;; --check-path) check_paths+=("${2:?}"); shift 2;; --skip-gpu-check) skip_gpu=1; shift;; + --follow-output) follow_output=1; shift;; --dry-run) dry_run=1; shift;; -h|--help) usage; exit 0;; --) shift; break;; + *) die "unknown launcher option: $1";; esac done - [[ $# -gt 0 ]] || die "missing training program after --" [[ -z $hosts_csv || -z $hostfile ]] || die "use only one of --hosts and --hostfile" [[ -n $hosts_csv || -n $hostfile ]] || die "--hosts or --hostfile is required" -[[ $nproc_per_node =~ ^[1-9][0-9]*$ ]] || die "--nproc-per-node must be positive" -[[ $master_port =~ ^[1-9][0-9]*$ ]] || die "--master-port must be positive" -[[ $ssh_port =~ ^[1-9][0-9]*$ ]] || die "--ssh-port must be positive" -[[ $node_start_delay =~ ^[0-9]+([.][0-9]+)?$ ]] || die "--node-start-delay must be non-negative" +for value in "$nproc" "$master_port" "$ssh_port" "$rdzv_timeout"; do [[ $value =~ ^[1-9][0-9]*$ ]] || die "positive integer required: $value"; done +[[ $job_timeout =~ ^[0-9]+$ ]] || die "--job-timeout must be non-negative" +for entry in "${exports[@]}"; do [[ $entry == [A-Za-z_]*=* ]] || die "invalid --env: $entry"; done declare -a hosts=() -if [[ -n $hosts_csv ]]; then - IFS=',' read -r -a hosts <<<"$hosts_csv" -else +if [[ -n $hosts_csv ]]; then IFS=',' read -r -a hosts <<<"$hosts_csv"; else [[ -r $hostfile ]] || die "cannot read hostfile: $hostfile" - while IFS= read -r line || [[ -n $line ]]; do - line=${line%%#*} - read -r line <<<"$line" - [[ -n $line ]] && hosts+=("$line") - done <"$hostfile" + while IFS= read -r line || [[ -n $line ]]; do line=${line%%#*}; read -r line <<<"$line"; [[ -n $line ]] && hosts+=("$line"); done <"$hostfile" fi [[ ${#hosts[@]} -gt 0 ]] || die "host list is empty" - -for entry in "${exports[@]}"; do - [[ $entry == *=* && $entry != =* ]] || die "invalid --env value: $entry" -done - +declare -A seen=(); for host in "${hosts[@]}"; do [[ -n $host ]] || die "empty host"; [[ -z ${seen[$host]:-} ]] || die "duplicate host: $host"; seen[$host]=1; done master_addr=${master_addr:-${hosts[0]}} -rdzv_id=${rdzv_id:-"infinitrain-$(date +%Y%m%d-%H%M%S)-$$"} -log_dir=${log_dir:-"$PWD/logs/$rdzv_id"} -mkdir -p "$log_dir" - -train_command=("$@") -nnodes=${#hosts[@]} -declare -a child_pids=() -declare -a child_nodes=() -stopping=0 - -stop_children() { - ((stopping == 0)) || return - stopping=1 - local pid - for pid in "${child_pids[@]:-}"; do - kill -TERM -- "-$pid" 2>/dev/null || true - done -} -trap 'stop_children; exit 130' INT TERM - -for node_rank in "${!hosts[@]}"; do - host=${hosts[$node_rank]} - node_command=( - "$infini_run" - "--nnodes=$nnodes" - "--nproc_per_node=$nproc_per_node" - "--node_rank=$node_rank" - "--rdzv_endpoint=$master_addr:$master_port" - "--rdzv_id=$rdzv_id" - "${train_command[@]}" - ) - - printf -v quoted_workdir '%q' "$workdir" - command_text="cd $quoted_workdir && " - for entry in "${exports[@]}"; do - printf -v quoted_export '%q' "$entry" - command_text+="export $quoted_export; " - done - command_text+="exec $(quote_command "${node_command[@]}")" - - log_file="$log_dir/node-${node_rank}-${host//[^a-zA-Z0-9_.-]/_}.log" - if [[ -n $local_host && $host == "$local_host" ]]; then - launch=(bash -lc "$command_text") +rdzv_id=${rdzv_id:-"infinitrain-$(date -u +%Y%m%dT%H%M%SZ)-$$-$RANDOM"} +[[ $rdzv_id =~ ^[A-Za-z0-9_.-]+$ ]] || die "rdzv-id may contain only letters, digits, dot, underscore, dash" +remote_log_root=${remote_log_root:-"$workdir/logs"} +job_dir="$log_root/$rdzv_id"; mkdir -p "$job_dir" +train=("$@"); nnodes=${#hosts[@]}; expected_sha= reference_fingerprint= +if [[ -d $workdir/.git ]]; then expected_sha=$(git -c safe.directory="$workdir" -C "$workdir" rev-parse HEAD 2>/dev/null || true); fi + +ssh_base=(ssh -o BatchMode=yes -o ConnectTimeout=10 -p "$ssh_port") +for option in "${ssh_options[@]}"; do ssh_base+=(-o "$option"); done +node_shell() { + local host=$1 command=$2 destination=$host + if [[ -n $local_host && $host == "$local_host" ]]; then bash -c "$command" else - destination=$host [[ -z $ssh_user ]] || destination="$ssh_user@$host" - launch=(ssh -o BatchMode=yes -o ConnectTimeout=10 -p "$ssh_port") - for option in "${ssh_options[@]}"; do - launch+=(-o "$option") - done - if [[ -n $remote_docker_container ]]; then - docker_command=(docker exec "$remote_docker_container" bash -lc "$command_text") - launch+=("$destination" "$(quote_command "${docker_command[@]}")") - else - launch+=("$destination" "$command_text") - fi + if [[ -n $container ]]; then + local wrapped; wrapped=$(quote_command docker exec "$container" bash -c "$command") + "${ssh_base[@]}" "$destination" "$wrapped" + else "${ssh_base[@]}" "$destination" "$command"; fi fi +} - launch=(setsid "${launch[@]}") - - printf '[node %d %s] ' "$node_rank" "$host" - quote_command "${launch[@]}" - printf ' > %q 2>&1\n' "$log_file" - if ((dry_run)); then - continue +follow_node_shell() { + local host=$1 command=$2 destination=$host + if [[ -n $local_host && $host == "$local_host" ]]; then + exec bash -c "$command" fi + [[ -z $ssh_user ]] || destination="$ssh_user@$host" + if [[ -n $container ]]; then + local wrapped + wrapped=$(quote_command docker exec "$container" bash -c "$command") + exec "${ssh_base[@]}" "$destination" "$wrapped" + fi + exec "${ssh_base[@]}" "$destination" "$command" +} - "${launch[@]}" >"$log_file" 2>&1 & - child_pids+=("$!") - child_nodes+=("$host") - - if ((node_rank + 1 < nnodes)) && [[ $node_start_delay != 0 ]]; then - sleep "$node_start_delay" +remote_job_dir="$remote_log_root/$rdzv_id" +code_sha=${expected_sha:-unknown} +code_dirty=unknown +if [[ -d $workdir/.git ]]; then + if [[ -z $(git -c safe.directory="$workdir" -C "$workdir" status --porcelain --untracked-files=normal 2>/dev/null) ]]; then code_dirty=false; else code_dirty=true; fi +fi +printf 'job_id=%s\ncode_sha=%s\ncode_dirty=%s\nbase_sha=%s\nmaster_sha=%s\nstart_utc=%s\ncontroller=%s\nmaster=%s:%s\nnnodes=%d\nnproc_per_node=%d\n' \ + "$rdzv_id" "$code_sha" "$code_dirty" "924d21256f3baf66bccb7458aa77395f0fbf5841" \ + "5a7be78e9b9d42c2f26efe33984dd3167dc60dc9" "$(date -u +%FT%TZ)" "$(hostname)" \ + "$master_addr" "$master_port" "$nnodes" "$nproc" >"$job_dir/metadata.txt" +printf '%q ' "$0" "${original_argv[@]}" >"$job_dir/controller-command.txt" +printf '\n' >>"$job_dir/controller-command.txt" + +for rank in "${!hosts[@]}"; do + host=${hosts[$rank]} + printf -v preflight 'set -e; cd %q; test -x %q; test -x %q; command -v git >/dev/null; command -v sha256sum >/dev/null; printf "hostname=%%s\nsha=%%s\ninfini_run_sha256=%%s\ntrain_sha256=%%s\n" "$(hostname)" "$(git -c safe.directory=%q rev-parse HEAD)" "$(sha256sum %q)" "$(sha256sum %q)"' "$workdir" "$infini_run" "${train[0]}" "$workdir" "$infini_run" "${train[0]}" + for path in "${check_paths[@]}"; do printf -v qpath "%q" "$path"; preflight+="; test -r $qpath"; done + preflight+="; { command -v nvcc >/dev/null && nvcc --version || true; }; { command -v ldconfig >/dev/null && ldconfig -p 2>/dev/null | grep -i nccl || true; }" + if ((skip_gpu == 0)); then preflight+="; command -v nvidia-smi >/dev/null; nvidia-smi --query-gpu=index,name,driver_version,memory.total --format=csv,noheader; nvidia-smi --query-gpu=index --format=csv,noheader | awk 'END { exit(NR < $nproc) }'"; fi + if ((dry_run)); then echo "$preflight"; continue; fi + if ! node_shell "$host" "$preflight" >"$job_dir/preflight-$rank.log" 2>&1; then + echo "FAILED preflight node=$rank host=$host log=$job_dir/preflight-$rank.log" >&2; exit 1 fi + node_sha=$(sed -n 's/^sha=//p' "$job_dir/preflight-$rank.log" | head -1) + [[ -z $expected_sha || $node_sha == "$expected_sha" ]] || die "code SHA mismatch on $host: $node_sha != $expected_sha" + fingerprint=$(sed -n '/^infini_run_sha256=/p; /^train_sha256=/p' "$job_dir/preflight-$rank.log") + if [[ -z $reference_fingerprint ]]; then reference_fingerprint=$fingerprint; else [[ $fingerprint == "$reference_fingerprint" ]] || die "binary SHA-256 mismatch on $host"; fi done +declare -a pids=() node_for_pid=() +follow_pid= -((dry_run)) && exit 0 - -status=0 -declare -A node_by_pid=() -declare -A running_pids=() -for index in "${!child_pids[@]}"; do - pid=${child_pids[$index]} - node_by_pid[$pid]=${child_nodes[$index]} - running_pids[$pid]=1 +stopping=0 +cleanup_nodes() { + ((stopping == 0)) || return; stopping=1 + local -a cleanup_pids=() + for rank in "${!hosts[@]}"; do + host=${hosts[$rank]}; printf -v qdir '%q' "$remote_job_dir"; printf -v qid '%q' "$rdzv_id" + cleanup=$(quote_command bash -c 'p=$1; id=$2; if test -r "$p"; then pid=$(cat "$p"); case "$pid" in ""|*[!0-9]*) exit 0;; esac; if test -r /proc/$pid/cmdline && tr "\0" " " /dev/null; then kill -TERM -- -$pid 2>/dev/null || true; sleep 1; kill -KILL -- -$pid 2>/dev/null || true; fi; fi' _ "$remote_job_dir/node-$rank/launcher.pid" "$rdzv_id") + node_shell "$host" "$cleanup" >/dev/null 2>&1 & cleanup_pids+=("$!") + done + for pid in "${cleanup_pids[@]}"; do wait "$pid" 2>/dev/null || true; done +} +stop_follower() { if [[ -n ${follow_pid:-} ]]; then kill "$follow_pid" 2>/dev/null || true; wait "$follow_pid" 2>/dev/null || true; follow_pid=; fi; } +watchdog_pid= +abort_job() { local rc=$1 reason=$2; cleanup_nodes; stop_follower; printf 'end_utc=%s\nexit_code=%d\nfirst_error=%s\n' "$(date -u +%FT%TZ)" "$rc" "$reason" >>"$job_dir/metadata.txt"; echo "FAILED job=$rdzv_id exit=$rc error=$reason logs=$job_dir remote_logs=$remote_job_dir" >&2; exit "$rc"; } +if ((job_timeout > 0)); then (sleep "$job_timeout"; kill -USR1 $$) & watchdog_pid=$!; fi +trap 'abort_job 124 "job timeout after ${job_timeout}s"' USR1 +trap 'abort_job 130 "controller interrupted"' INT TERM + +for rank in "${!hosts[@]}"; do + host=${hosts[$rank]}; node_log_dir="$remote_job_dir/node-$rank" + node_cmd=("$infini_run" "--nnodes=$nnodes" "--nproc_per_node=$nproc" "--node_rank=$rank" + "--rdzv_endpoint=$master_addr:$master_port" "--rdzv_id=$rdzv_id" "--rdzv_timeout=$rdzv_timeout" + "--log_dir=$node_log_dir" "${train[@]}") + printf -v qwork '%q' "$workdir"; printf -v qdir '%q' "$remote_job_dir" + command="set -e; mkdir -p $qdir/node-$rank; cd $qwork; " + for entry in "${exports[@]}"; do printf -v qentry '%q' "$entry"; command+="export $qentry; "; done + command+="setsid $(quote_command "${node_cmd[@]}") & pid=\$!; echo \$pid >$qdir/node-$rank/launcher.pid; trap 'kill -TERM -- -\$pid 2>/dev/null || true' INT TERM HUP EXIT; wait \$pid" + echo "launch node=$rank host=$host log=$job_dir/node-$rank-controller.log" + if ((dry_run)); then echo "$command"; continue; fi + node_shell "$host" "$command" >"$job_dir/node-$rank-controller.log" 2>&1 & + pids+=("$!"); node_for_pid+=("$rank:$host") done +if ((follow_output && !dry_run)); then + follow_node_rank=$((nnodes - 1)); follow_rank=$((nnodes * nproc - 1)); follow_host=${hosts[$follow_node_rank]} + follow_file="$remote_job_dir/node-$follow_node_rank/rank-$follow_rank.stderr.log" + printf -v qfollow '%q' "$follow_file" + follow_command="while ! test -f $qfollow; do sleep 0.2; done; exec tail -s 0.1 -n +1 -F $qfollow" + echo "follow global_rank=$follow_rank host=$follow_host file=$follow_file" + follow_node_shell "$follow_host" "$follow_command" & follow_pid=$! +fi +((dry_run)) && exit 0 -while ((${#running_pids[@]} > 0)); do - finished_pid="" - if wait -n -p finished_pid; then - rc=0 - else - rc=$? - fi - [[ -n ${finished_pid:-} ]] || die "failed to identify completed node launcher" - unset 'running_pids[$finished_pid]' - if ((rc != 0)); then - echo "node ${node_by_pid[$finished_pid]} failed (launcher status $rc)" >&2 - status=1 - stop_children - for pid in "${!running_pids[@]}"; do - wait "$pid" 2>/dev/null || true - done - break - fi +status=0 first_error= deadline=0 +((job_timeout > 0)) && deadline=$((SECONDS + job_timeout)) +declare -A running=(); for i in "${!pids[@]}"; do running[${pids[$i]}]=${node_for_pid[$i]}; done +while ((${#running[@]})); do + if ((deadline > 0 && SECONDS >= deadline)); then status=124; first_error="job timeout after ${job_timeout}s"; cleanup_nodes; break; fi + finished=; if wait -n -p finished; then rc=0; else rc=$?; fi + [[ -n ${finished:-} ]] || { sleep 1; continue; } + node=${running[$finished]}; unset 'running[$finished]' + if ((rc != 0)); then status=$rc; first_error="node $node exited with status $rc"; cleanup_nodes; break; fi done - -trap - INT TERM +if ((follow_output)); then sleep 0.3; fi +stop_follower +for pid in "${!running[@]}"; do wait "$pid" 2>/dev/null || true; done +[[ -z $watchdog_pid ]] || kill "$watchdog_pid" 2>/dev/null || true +trap - INT TERM USR1 +printf 'end_utc=%s\nexit_code=%d\nfirst_error=%s\n' "$(date -u +%FT%TZ)" "$status" "$first_error" >>"$job_dir/metadata.txt" +if ((status == 0)); then echo "SUCCESS job=$rdzv_id logs=$job_dir remote_logs=$remote_job_dir"; else echo "FAILED job=$rdzv_id exit=$status error=$first_error logs=$job_dir remote_logs=$remote_job_dir" >&2; fi exit "$status" diff --git a/scripts/multinode_ssh/test_launch_multinode_ssh.sh b/scripts/multinode_ssh/test_launch_multinode_ssh.sh index 27a8b612..6c9ea136 100755 --- a/scripts/multinode_ssh/test_launch_multinode_ssh.sh +++ b/scripts/multinode_ssh/test_launch_multinode_ssh.sh @@ -1,73 +1,135 @@ #!/usr/bin/env bash - set -euo pipefail script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) launcher="$script_dir/launch_multinode_ssh.sh" -tmpdir=$(mktemp -d) -trap 'rm -rf "$tmpdir"' EXIT +repo=$(cd "$script_dir/../.." && pwd) +tmp=$(mktemp -d) +trap 'rm -rf "$tmp"' EXIT -cat >"$tmpdir/fake_infini_run" <<'EOF' +cat >"$tmp/fake_infini_run" <<'RUN' #!/usr/bin/env bash set -euo pipefail -printf '%s\n' "$*" -EOF -chmod +x "$tmpdir/fake_infini_run" +printf '%q ' "$@" +printf '\n' +log_dir= +node_rank=0 +nproc=1 +for arg in "$@"; do + case "$arg" in + --log_dir=*) log_dir=${arg#*=} ;; + --node_rank=*) node_rank=${arg#*=} ;; + --nproc_per_node=*) nproc=${arg#*=} ;; + esac +done +case "${FAKE_MODE:-ok}" in + fail) exit 17 ;; + sleep) + trap 'exit 143' TERM INT + while :; do sleep 1; done + ;; + output) + mkdir -p "$log_dir" + rank=$(((node_rank + 1) * nproc - 1)) + printf 'live-start rank=%d\n' "$rank" >"$log_dir/rank-$rank.stderr.log" + sleep 1 + printf 'live-end rank=%d\n' "$rank" >>"$log_dir/rank-$rank.stderr.log" + ;; +esac +RUN +chmod +x "$tmp/fake_infini_run" + +output=$("$launcher" --hosts alpha,beta --workdir "$repo" \ + --infini-run "$tmp/fake_infini_run" --nproc-per-node 8 \ + --master-addr 10.0.0.1 --rdzv-id dry-run-test \ + --log-root "$tmp/logs" --remote-log-root "$tmp/remote logs" \ + --remote-docker-container test-container --skip-gpu-check --dry-run \ + -- /bin/echo 'value with spaces' 'semi;colon') +grep -q -- 'node_rank=0' <<<"$output" +grep -q -- 'node_rank=1' <<<"$output" +grep -q -- 'nproc_per_node=8' <<<"$output" +grep -q -- '10.0.0.1:29500' <<<"$output" -output=$( - "$launcher" \ - --hosts local.test,remote.test \ - --local-host local.test \ - --workdir "$tmpdir" \ - --infini-run ./fake_infini_run \ - --nproc-per-node 2 \ - --master-addr 10.0.0.1 \ - --rdzv-id test-job \ - --env NCCL_DEBUG=INFO \ - --remote-docker-container test-container \ - --dry-run \ - -- ./train --flag 'value with spaces' -) +printf '# rank order\nlocal.test\n' >"$tmp/hosts" +"$launcher" --hostfile "$tmp/hosts" --local-host local.test --workdir "$repo" \ + --infini-run "$tmp/fake_infini_run" --rdzv-id success-test \ + --log-root "$tmp/logs" --remote-log-root "$tmp/remote" --skip-gpu-check \ + --env 'SAFE_VALUE=value with spaces;still-data' -- /bin/echo 'arg with spaces' +grep -q 'exit_code=0' "$tmp/logs/success-test/metadata.txt" +grep -q 'arg\\ with\\ spaces' "$tmp/logs/success-test/node-0-controller.log" +grep -q -- '--hostfile' "$tmp/logs/success-test/controller-command.txt" -grep -q -- '--node_rank=0' <<<"$output" -grep -q -- '--node_rank=1' <<<"$output" -grep -q -- '--nnodes=2' <<<"$output" -grep -q -- '--nproc_per_node=2' <<<"$output" -grep -q -- '--rdzv_endpoint=10.0.0.1:29500' <<<"$output" -grep -q -- '--rdzv_id=test-job' <<<"$output" -grep -q -- 'local.test' <<<"$output" -grep -q -- 'remote.test' <<<"$output" -grep -q -- 'docker.*exec.*test-container' <<<"$output" +FAKE_MODE=output "$launcher" --hosts local.test --local-host local.test \ + --workdir "$repo" --infini-run "$tmp/fake_infini_run" --nproc-per-node 2 \ + --rdzv-id output-test --log-root "$tmp/logs" \ + --remote-log-root "$tmp/remote" --skip-gpu-check --follow-output \ + -- /bin/true >"$tmp/follow-output.log" +grep -q 'live-start rank=1' "$tmp/follow-output.log" +grep -q 'live-end rank=1' "$tmp/follow-output.log" -"$launcher" \ - --hosts local.test \ - --local-host local.test \ - --workdir "$tmpdir" \ - --infini-run ./fake_infini_run \ - --rdzv-id local-test \ - --log-dir "$tmpdir/logs" \ - -- ./train --flag value +if FAKE_MODE=fail "$launcher" --hosts local.test --local-host local.test \ + --workdir "$repo" --infini-run "$tmp/fake_infini_run" \ + --rdzv-id failure-test --log-root "$tmp/logs" \ + --remote-log-root "$tmp/remote" --skip-gpu-check -- /bin/true; then + echo "non-zero node unexpectedly succeeded" >&2 + exit 1 +fi +grep -q 'exit_code=17' "$tmp/logs/failure-test/metadata.txt" -grep -q -- '--node_rank=0' "$tmpdir/logs/node-0-local.test.log" -grep -q -- './train --flag value' "$tmpdir/logs/node-0-local.test.log" +FAKE_MODE=sleep "$launcher" --hosts local.test --local-host local.test \ + --workdir "$repo" --infini-run "$tmp/fake_infini_run" \ + --rdzv-id signal-test --log-root "$tmp/logs" \ + --remote-log-root "$tmp/remote" --skip-gpu-check -- /bin/true & +launcher_pid=$! +pid_file="$tmp/remote/signal-test/node-0/launcher.pid" +for _ in {1..150}; do + [[ -r $pid_file ]] && break + sleep 0.1 +done +[[ -r $pid_file ]] || { echo "signal test never launched remote job" >&2; exit 1; } +kill -TERM "$launcher_pid" +if wait "$launcher_pid"; then + echo "signalled launcher unexpectedly succeeded" >&2 + exit 1 +fi +sleep 1 +if [[ -r $pid_file ]] && kill -0 "$(cat "$pid_file")" 2>/dev/null; then + echo "remote job process remained after signal" >&2 + exit 1 +fi -SECONDS=0 -"$launcher" \ - --hosts local.test,local.test \ - --local-host local.test \ - --workdir "$tmpdir" \ - --infini-run ./fake_infini_run \ - --rdzv-id delayed-local-test \ - --node-start-delay 1 \ - --log-dir "$tmpdir/delayed-logs" \ - -- ./train -((SECONDS >= 1)) +start=$SECONDS +if FAKE_MODE=sleep "$launcher" --hosts local.test --local-host local.test \ + --workdir "$repo" --infini-run "$tmp/fake_infini_run" \ + --rdzv-id timeout-test --job-timeout 2 --log-root "$tmp/logs" \ + --remote-log-root "$tmp/remote" --skip-gpu-check -- /bin/true; then + echo "timed job unexpectedly succeeded" >&2 + exit 1 +else + rc=$? + [[ $rc == 124 ]] || { echo "timeout returned $rc" >&2; exit 1; } +fi +((SECONDS - start < 20)) +grep -q 'exit_code=124' "$tmp/logs/timeout-test/metadata.txt" +pid_file="$tmp/remote/timeout-test/node-0/launcher.pid" +if [[ -r $pid_file ]] && kill -0 "$(cat "$pid_file")" 2>/dev/null; then + echo "remote job process remained after timeout" >&2 + exit 1 +fi + +if "$launcher" --hosts duplicate,duplicate --skip-gpu-check -- /bin/true \ + >"$tmp/invalid" 2>&1; then + echo "duplicate hosts unexpectedly accepted" >&2 + exit 1 +fi +grep -q 'duplicate host' "$tmp/invalid" -if "$launcher" --hosts local.test --node-start-delay invalid -- ./train \ - >"$tmpdir/invalid-delay.out" 2>&1; then - echo "invalid node start delay unexpectedly succeeded" >&2 +if "$launcher" --hosts local.test --local-host local.test --workdir "$repo" \ + --infini-run "$tmp/fake_infini_run" --rdzv-id missing-path-test \ + --log-root "$tmp/logs" --remote-log-root "$tmp/remote" --skip-gpu-check \ + --check-path "$tmp/does-not-exist" -- /bin/true; then + echo "missing check path unexpectedly accepted" >&2 exit 1 fi -grep -q -- '--node-start-delay must be non-negative' "$tmpdir/invalid-delay.out" echo "launch_multinode_ssh tests passed" diff --git a/tools/infini_run/infini_run.cc b/tools/infini_run/infini_run.cc index 51ad1e0c..64094f84 100644 --- a/tools/infini_run/infini_run.cc +++ b/tools/infini_run/infini_run.cc @@ -1,16 +1,24 @@ +#include +#include +#include +#include #include #include -#include #include #include +#include #include #include #include +#include #include -#include +#include +#include +#include +#include #include -#include +#include #include #include @@ -20,190 +28,301 @@ DEFINE_int32(nnodes, 1, "Total number of nodes"); DEFINE_int32(nproc_per_node, 1, "Number of processes per node"); DEFINE_int32(node_rank, 0, "Rank of this node"); -DEFINE_string(rdzv_endpoint, "127.0.0.1:29500", - "Rendezvous endpoint (host:port)"); - +DEFINE_string(rdzv_endpoint, "127.0.0.1:29500", "Rendezvous endpoint"); DEFINE_string(rdzv_id, "", "Unique job ID shared by all nodes"); +DEFINE_int32(rdzv_timeout, 120, "Rendezvous timeout in seconds"); +DEFINE_string(log_dir, "", "Directory for per-rank stdout/stderr logs"); namespace { +constexpr uint32_t kMagic = 0x49545244; +volatile sig_atomic_t g_signal = 0; +void HandleSignal(int signal) { g_signal = signal; } bool IsLauncherFlag(const std::string &flag) { return flag == "--nnodes" || flag == "--nproc_per_node" || flag == "--node_rank" || flag == "--rdzv_endpoint" || - flag == "--rdzv_id"; + flag == "--rdzv_id" || flag == "--rdzv_timeout" || flag == "--log_dir"; } - int FindTrainProgramIndex(int argc, char **argv) { for (int i = 1; i < argc; ++i) { const std::string arg = argv[i]; - if (arg == "--") { - return i + 1; - } - if (arg.rfind("--", 0) != 0) { - return i; - } - - const size_t eq_pos = arg.find('='); - const std::string flag = arg.substr(0, eq_pos); - if (IsLauncherFlag(flag) && eq_pos == std::string::npos) { - ++i; - } + if (arg == "--") return i + 1; + if (!arg.starts_with("--")) return i; + const size_t eq = arg.find('='); + if (IsLauncherFlag(arg.substr(0, eq)) && eq == std::string::npos) ++i; } return argc; } - -std::vector BuildLauncherArgv(int train_program_index, char **argv) { - std::vector launcher_argv; - launcher_argv.push_back(argv[0]); - for (int i = 1; i < train_program_index; ++i) { - launcher_argv.push_back(argv[i]); - } - launcher_argv.push_back(nullptr); - return launcher_argv; +std::vector BuildLauncherArgv(int end, char **argv) { + std::vector result{argv[0]}; + for (int i = 1; i < end; ++i) result.push_back(argv[i]); + result.push_back(nullptr); + return result; } - void SetEnvInt(const char *name, int value) { - const auto value_str = std::to_string(value); - setenv(name, value_str.c_str(), 1); + const std::string text = std::to_string(value); + setenv(name, text.c_str(), 1); } - -void TerminateChildren(const std::unordered_set &child_pids) { - for (pid_t child_pid : child_pids) { - kill(child_pid, SIGTERM); +bool ReadAll(int fd, void *buffer, size_t size) { + auto *bytes = static_cast(buffer); + while (size > 0) { + ssize_t n = read(fd, bytes, size); + if (n < 0 && errno == EINTR) continue; + if (n <= 0) return false; + bytes += n; + size -= n; } + return true; +} +bool WriteAll(int fd, const void *buffer, size_t size) { + const auto *bytes = static_cast(buffer); + while (size > 0) { + ssize_t n = write(fd, bytes, size); + if (n < 0 && errno == EINTR) continue; + if (n <= 0) return false; + bytes += n; + size -= n; + } + return true; } -int ExitCodeFromStatus(int status) { - if (WIFEXITED(status)) { - return WEXITSTATUS(status); +class RendezvousServer { + public: + RendezvousServer(int port, int timeout) : timeout_(timeout) { + fd_ = socket(AF_INET, SOCK_STREAM, 0); + CHECK_GE(fd_, 0); + int reuse = 1; + setsockopt(fd_, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)); + sockaddr_in address{}; + address.sin_family = AF_INET; + address.sin_addr.s_addr = htonl(INADDR_ANY); + address.sin_port = htons(port); + CHECK_EQ(bind(fd_, reinterpret_cast(&address), sizeof(address)), + 0) + << "Failed to bind rendezvous port " << port << ": " << strerror(errno); + CHECK_EQ(listen(fd_, 64), 0); + thread_ = std::thread([this] { Serve(); }); } - if (WIFSIGNALED(status)) { - return 128 + WTERMSIG(status); + ~RendezvousServer() { + shutdown(fd_, SHUT_RDWR); + close(fd_); + if (thread_.joinable()) thread_.join(); + for (auto &worker : workers_) + if (worker.joinable()) worker.join(); } - return 1; -} -void CleanupRunUniqueIdFiles(const std::string &run_id) { - const std::string prefix = "cclUniqueId_" + run_id + "_"; - for (const auto &entry : - std::filesystem::directory_iterator(std::filesystem::current_path())) { - if (!entry.is_regular_file()) { - continue; + private: + struct Header { + uint32_t magic, operation, key_size, value_size; + }; + void Serve() { + while (true) { + int client = accept(fd_, nullptr, nullptr); + if (client < 0) break; + workers_.emplace_back(&RendezvousServer::Handle, this, client); } - const std::string filename = entry.path().filename().string(); - if (filename.rfind(prefix, 0) == 0) { - std::error_code ec; - std::filesystem::remove(entry.path(), ec); - if (ec) { - LOG(WARNING) << "Failed to remove unique-id file " << entry.path() - << ": " << ec.message(); - } + } + void Handle(int client) { + Header wire{}; + if (!ReadAll(client, &wire, sizeof(wire))) { + close(client); + return; } + uint32_t operation = ntohl(wire.operation), key_size = ntohl(wire.key_size), + value_size = ntohl(wire.value_size); + if (ntohl(wire.magic) != kMagic || key_size == 0 || key_size > 4096 || + value_size > 1024 * 1024) { + close(client); + return; + } + std::string key(key_size, '\0'), value(value_size, '\0'); + if (!ReadAll(client, key.data(), key.size()) || + (!value.empty() && !ReadAll(client, value.data(), value.size()))) { + close(client); + return; + } + uint32_t status = 0; + if (operation == 1) { + { + std::lock_guard lock(mutex_); + values_[key] = std::move(value); + } + condition_.notify_all(); + } else if (operation == 2) { + std::unique_lock lock(mutex_); + if (!condition_.wait_for(lock, std::chrono::seconds(timeout_), + [&] { return values_.contains(key); })) + status = 1; + else + value = values_.at(key); + } else + status = 2; + CHECK_LE(value.size(), std::numeric_limits::max()); + uint32_t response[2] = {htonl(status), + htonl(static_cast(value.size()))}; + WriteAll(client, response, sizeof(response)); + if (status == 0 && operation == 2) + WriteAll(client, value.data(), value.size()); + close(client); } -} + int fd_ = -1, timeout_; + std::thread thread_; + std::vector workers_; + std::mutex mutex_; + std::condition_variable condition_; + std::map values_; +}; -std::string GenerateLocalRunId() { - const auto now = std::chrono::steady_clock::now().time_since_epoch().count(); - return std::to_string(getpid()) + "_" + std::to_string(now); +int ExitCode(int status) { + if (WIFEXITED(status)) return WEXITSTATUS(status); + return WIFSIGNALED(status) ? 128 + WTERMSIG(status) : 1; +} +void SignalChildren(const std::unordered_set &children, int signal) { + for (pid_t pid : children) kill(pid, signal); +} +std::string LocalRunId() { + return std::to_string(getpid()) + "_" + + std::to_string( + std::chrono::steady_clock::now().time_since_epoch().count()); +} +void CleanupFiles(const std::string &run_id) { + const std::string prefix = "cclUniqueId_" + run_id + "_"; + for (const auto &entry : std::filesystem::directory_iterator(".")) { + if (entry.is_regular_file() && + entry.path().filename().string().starts_with(prefix)) { + std::error_code error; + std::filesystem::remove(entry.path(), error); + } + } } - } // namespace int main(int argc, char **argv) { - const int train_program_index = FindTrainProgramIndex(argc, argv); - std::vector launcher_argv = - BuildLauncherArgv(train_program_index, argv); - int launcher_argc = static_cast(launcher_argv.size()) - 1; + const int train_index = FindTrainProgramIndex(argc, argv); + auto launcher_argv = BuildLauncherArgv(train_index, argv); + int launcher_argc = launcher_argv.size() - 1; char **launcher_argv_ptr = launcher_argv.data(); gflags::ParseCommandLineFlags(&launcher_argc, &launcher_argv_ptr, true); google::InitGoogleLogging(argv[0]); + CHECK_GT(FLAGS_nnodes, 0); + CHECK_GT(FLAGS_nproc_per_node, 0); + CHECK_GE(FLAGS_node_rank, 0); + CHECK_LT(FLAGS_node_rank, FLAGS_nnodes); + CHECK_GT(FLAGS_rdzv_timeout, 0); + CHECK_LT(train_index, argc); + const size_t separator = FLAGS_rdzv_endpoint.rfind(':'); + CHECK_NE(separator, std::string::npos) << "rdzv_endpoint must be host:port"; + CHECK(FLAGS_nnodes == 1 || !FLAGS_rdzv_id.empty()); - CHECK_GT(FLAGS_nnodes, 0) << "nnodes must be positive"; - CHECK_GT(FLAGS_nproc_per_node, 0) << "nproc_per_node must be positive"; - CHECK_GE(FLAGS_node_rank, 0) << "node_rank must be non-negative"; - CHECK_LT(FLAGS_node_rank, FLAGS_nnodes) - << "node_rank must be less than nnodes"; - CHECK_NE(FLAGS_rdzv_endpoint.find(':'), std::string::npos) - << "rdzv_endpoint must be host:port"; - CHECK(FLAGS_nnodes == 1 || !FLAGS_rdzv_id.empty()) - << "rdzv_id must be set to the same unique job ID on every node for " - "multi-node training"; - - CHECK_LT(train_program_index, argc) << "No training program specified!"; + const std::string master_addr = FLAGS_rdzv_endpoint.substr(0, separator); + const std::string master_port = FLAGS_rdzv_endpoint.substr(separator + 1); + char *port_end = nullptr; + errno = 0; + const long parsed_port = std::strtol(master_port.c_str(), &port_end, 10); + CHECK(errno == 0 && port_end != master_port.c_str() && *port_end == '\0' && + parsed_port > 0 && parsed_port <= 65535) + << "Invalid rendezvous port: " << master_port; + const std::string run_id = + FLAGS_rdzv_id.empty() ? LocalRunId() : FLAGS_rdzv_id; + setenv("INFINI_RUN_ID", run_id.c_str(), 1); + setenv("INFINI_RDZV_TIMEOUT_SEC", std::to_string(FLAGS_rdzv_timeout).c_str(), + 1); + std::unique_ptr server; + if (FLAGS_nnodes > 1 && FLAGS_node_rank == 0) + server = std::make_unique(static_cast(parsed_port), + FLAGS_rdzv_timeout); + if (!FLAGS_log_dir.empty()) + std::filesystem::create_directories(FLAGS_log_dir); + struct sigaction action {}; + action.sa_handler = HandleSignal; + sigemptyset(&action.sa_mask); + sigaction(SIGINT, &action, nullptr); + sigaction(SIGTERM, &action, nullptr); - std::string train_program = argv[train_program_index]; std::vector train_argv; - for (int i = train_program_index; i < argc; ++i) { - train_argv.push_back(argv[i]); - } + for (int i = train_index; i < argc; ++i) train_argv.push_back(argv[i]); train_argv.push_back(nullptr); - - int proc_world_size = FLAGS_nnodes * FLAGS_nproc_per_node; - std::string master_addr = - FLAGS_rdzv_endpoint.substr(0, FLAGS_rdzv_endpoint.find(':')); - std::string master_port = - FLAGS_rdzv_endpoint.substr(FLAGS_rdzv_endpoint.find(':') + 1); - const std::string run_id = - FLAGS_rdzv_id.empty() ? GenerateLocalRunId() : FLAGS_rdzv_id; - - std::unordered_set running_children; + std::unordered_set children; int exit_code = 0; - - for (int local_proc_rank = 0; local_proc_rank < FLAGS_nproc_per_node; - ++local_proc_rank) { + const int world_size = FLAGS_nnodes * FLAGS_nproc_per_node; + for (int local_rank = 0; local_rank < FLAGS_nproc_per_node; ++local_rank) { pid_t pid = fork(); if (pid < 0) { - perror("fork failed"); - exit_code = 1; - TerminateChildren(running_children); - break; + PLOG(ERROR) << "fork failed after starting " << children.size() + << " rank processes"; + SignalChildren(children, SIGTERM); + for (pid_t child : children) waitpid(child, nullptr, 0); + return 1; } if (pid == 0) { - int global_proc_rank = - FLAGS_node_rank * FLAGS_nproc_per_node + local_proc_rank; + const int rank = FLAGS_node_rank * FLAGS_nproc_per_node + local_rank; SetEnvInt("LOCAL_WORLD_SIZE", FLAGS_nproc_per_node); - + SetEnvInt("WORLD_SIZE", world_size); + SetEnvInt("RANK", rank); + SetEnvInt("LOCAL_RANK", local_rank); setenv("MASTER_ADDR", master_addr.c_str(), 1); setenv("MASTER_PORT", master_port.c_str(), 1); - setenv("INFINI_RUN_ID", run_id.c_str(), 1); - - SetEnvInt("RANK", global_proc_rank); - SetEnvInt("LOCAL_RANK", local_proc_rank); - - SetEnvInt("WORLD_SIZE", proc_world_size); - - execvp(train_program.c_str(), train_argv.data()); + if (!FLAGS_log_dir.empty()) { + const std::string prefix = + FLAGS_log_dir + "/rank-" + std::to_string(rank); + int out = open((prefix + ".stdout.log").c_str(), + O_WRONLY | O_CREAT | O_TRUNC, 0644); + int err = open((prefix + ".stderr.log").c_str(), + O_WRONLY | O_CREAT | O_TRUNC, 0644); + CHECK_GE(out, 0); + CHECK_GE(err, 0); + dup2(out, STDOUT_FILENO); + dup2(err, STDERR_FILENO); + close(out); + close(err); + } + char hostname[256] = {}; + gethostname(hostname, sizeof(hostname) - 1); + fprintf(stderr, + "INFINITRAIN_RANK hostname=%s global_rank=%d local_rank=%d " + "world_size=%d device=%d job_id=%s\n", + hostname, rank, local_rank, world_size, local_rank, + run_id.c_str()); + fflush(stderr); + execvp(train_argv[0], train_argv.data()); perror("exec failed"); - exit(1); + _exit(127); } - running_children.insert(pid); + children.insert(pid); } - while (!running_children.empty()) { - int status; - pid_t child = wait(&status); + bool term_sent = false; + auto term_time = std::chrono::steady_clock::now(); + while (!children.empty()) { + if (g_signal && !term_sent) { + exit_code = 128 + g_signal; + SignalChildren(children, SIGTERM); + term_sent = true; + term_time = std::chrono::steady_clock::now(); + } + if (term_sent && + std::chrono::steady_clock::now() - term_time > std::chrono::seconds(10)) + SignalChildren(children, SIGKILL); + int status = 0; + pid_t child = waitpid(-1, &status, term_sent ? WNOHANG : 0); + if (child == 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + continue; + } if (child < 0) { - if (errno == EINTR) { - continue; - } - perror("wait failed"); - if (exit_code == 0) { - exit_code = 1; - } - TerminateChildren(running_children); + if (errno == EINTR) continue; + exit_code = 1; break; } - - running_children.erase(child); - const int child_exit_code = ExitCodeFromStatus(status); - if (child_exit_code != 0 && exit_code == 0) { - exit_code = child_exit_code; - TerminateChildren(running_children); + children.erase(child); + const int code = ExitCode(status); + if (code != 0 && exit_code == 0) { + exit_code = code; + SignalChildren(children, SIGTERM); + term_sent = true; + term_time = std::chrono::steady_clock::now(); } } - - if (FLAGS_nnodes == 1) { - CleanupRunUniqueIdFiles(run_id); - } + if (FLAGS_nnodes == 1) CleanupFiles(run_id); return exit_code; } From afc48a7f4c68e4c7de037632aab06a9fdfe4100d Mon Sep 17 00:00:00 2001 From: JYMiracle305 <604951424@qq.com> Date: Mon, 2 Nov 2026 09:48:35 +0000 Subject: [PATCH 3/4] fix: satisfy IWYU include checks --- infini_train/src/core/ccl/ccl_utils.cc | 3 +++ tools/infini_run/infini_run.cc | 9 ++++++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/infini_train/src/core/ccl/ccl_utils.cc b/infini_train/src/core/ccl/ccl_utils.cc index f0d7f666..fc34b4f9 100644 --- a/infini_train/src/core/ccl/ccl_utils.cc +++ b/infini_train/src/core/ccl/ccl_utils.cc @@ -1,15 +1,18 @@ #include "infini_train/include/core/ccl/ccl_utils.h" #include +#include #include #include #include #include +#include #include #include #include #include +#include #include #include #include diff --git a/tools/infini_run/infini_run.cc b/tools/infini_run/infini_run.cc index 64094f84..8b634757 100644 --- a/tools/infini_run/infini_run.cc +++ b/tools/infini_run/infini_run.cc @@ -1,12 +1,16 @@ +// IWYU pragma: no_include #include #include #include +#include #include #include #include +#include #include #include +#include #include #include #include @@ -17,9 +21,12 @@ #include #include #include +#include #include +#include #include #include +#include #include #include "gflags/gflags.h" @@ -35,7 +42,7 @@ DEFINE_string(log_dir, "", "Directory for per-rank stdout/stderr logs"); namespace { constexpr uint32_t kMagic = 0x49545244; -volatile sig_atomic_t g_signal = 0; +volatile std::sig_atomic_t g_signal = 0; void HandleSignal(int signal) { g_signal = signal; } bool IsLauncherFlag(const std::string &flag) { From b33bf4318538e6aa2272c0d715ae822a5c84cddc Mon Sep 17 00:00:00 2001 From: JYMiracle305 <604951424@qq.com> Date: Tue, 3 Nov 2026 09:59:33 +0000 Subject: [PATCH 4/4] refactor: unify TCP rendezvous handling --- CMakeLists.txt | 17 + infini_train/include/core/ccl/ccl_utils.h | 7 +- .../include/core/ccl/rendezvous_protocol.h | 26 ++ .../include/core/ccl/rendezvous_server.h | 21 ++ .../include/nn/parallel/process_group.h | 1 - infini_train/src/core/ccl/ccl_utils.cc | 103 ++---- .../src/core/ccl/rendezvous_server.cc | 185 ++++++++++ infini_train/src/nn/parallel/process_group.cc | 9 +- scripts/multinode_ssh/README.md | 118 ------ scripts/multinode_ssh/VALIDATION_REPORT.md | 68 ---- tests/distributed/CMakeLists.txt | 11 + .../distributed}/test_launch_multinode_ssh.sh | 2 +- tools/infini_run/CMakeLists.txt | 6 +- tools/infini_run/infini_run.cc | 337 +++++++----------- 14 files changed, 426 insertions(+), 485 deletions(-) create mode 100644 infini_train/include/core/ccl/rendezvous_protocol.h create mode 100644 infini_train/include/core/ccl/rendezvous_server.h create mode 100644 infini_train/src/core/ccl/rendezvous_server.cc delete mode 100644 scripts/multinode_ssh/README.md delete mode 100644 scripts/multinode_ssh/VALIDATION_REPORT.md rename {scripts/multinode_ssh => tests/distributed}/test_launch_multinode_ssh.sh (98%) diff --git a/CMakeLists.txt b/CMakeLists.txt index 88911e84..0c0a2c3f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -73,6 +73,11 @@ endif() file(GLOB_RECURSE SRC ${PROJECT_SOURCE_DIR}/infini_train/src/*.cc) list(FILTER SRC EXCLUDE REGEX ".*kernels/cpu/.*") +set(RENDEZVOUS_SERVER_SOURCE + ${PROJECT_SOURCE_DIR}/infini_train/src/core/ccl/rendezvous_server.cc +) +list(REMOVE_ITEM SRC ${RENDEZVOUS_SERVER_SOURCE}) + # Exclude backend-specific runtime/ccl translation units when the corresponding # backend is disabled. This keeps each build self-contained and avoids pulling # in headers (e.g. / ) that are not on the @@ -86,6 +91,18 @@ endif() # CPU kernels (*.cc) file(GLOB_RECURSE CPU_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/kernels/cpu/*.cc) +# ------------------------------------------------------------------------------ +# Rendezvous library +# ------------------------------------------------------------------------------ + +add_library(infini_train_rendezvous STATIC ${RENDEZVOUS_SERVER_SOURCE}) +add_library(InfiniTrain::rendezvous ALIAS infini_train_rendezvous) +target_compile_features(infini_train_rendezvous PUBLIC cxx_std_20) +target_include_directories(infini_train_rendezvous PUBLIC + "$" +) +target_link_libraries(infini_train_rendezvous PUBLIC glog) + # ------------------------------------------------------------------------------ # CPU kernels library # ------------------------------------------------------------------------------ diff --git a/infini_train/include/core/ccl/ccl_utils.h b/infini_train/include/core/ccl/ccl_utils.h index d9f8ce6b..299b71dc 100644 --- a/infini_train/include/core/ccl/ccl_utils.h +++ b/infini_train/include/core/ccl/ccl_utils.h @@ -5,11 +5,8 @@ namespace infini_train::core { class CclUniqueId; -void WriteUniqueIdFile(const CclUniqueId &unique_id, - const std::string &pg_name); +void WriteUniqueId(const CclUniqueId &unique_id, const std::string &pg_name); -void ReadUniqueIdFile(CclUniqueId *unique_id, const std::string &pg_name); - -void CleanupUniqueIdFile(const std::string &pg_name); +void ReadUniqueId(CclUniqueId *unique_id, const std::string &pg_name); } // namespace infini_train::core diff --git a/infini_train/include/core/ccl/rendezvous_protocol.h b/infini_train/include/core/ccl/rendezvous_protocol.h new file mode 100644 index 00000000..83a6a295 --- /dev/null +++ b/infini_train/include/core/ccl/rendezvous_protocol.h @@ -0,0 +1,26 @@ +#pragma once + +#include + +namespace infini_train::core::rendezvous { +inline constexpr uint32_t kMagic = 0x49545244; +inline constexpr uint32_t kMaxKeySize = 4096; +inline constexpr uint32_t kMaxValueSize = 1024 * 1024; + +enum class Operation : uint32_t { + kPut = 1, + kGet = 2, +}; + +struct Header { + uint32_t magic; + uint32_t operation; + uint32_t key_size; + uint32_t value_size; +}; + +struct Response { + uint32_t status; + uint32_t value_size; +}; +} // namespace infini_train::core::rendezvous diff --git a/infini_train/include/core/ccl/rendezvous_server.h b/infini_train/include/core/ccl/rendezvous_server.h new file mode 100644 index 00000000..f3020680 --- /dev/null +++ b/infini_train/include/core/ccl/rendezvous_server.h @@ -0,0 +1,21 @@ +#pragma once + +#include +#include + +namespace infini_train::core { +class RendezvousServer { + public: + RendezvousServer(const std::string &host, int port, int timeout); + ~RendezvousServer(); + + RendezvousServer(const RendezvousServer &) = delete; + RendezvousServer &operator=(const RendezvousServer &) = delete; + + int port() const; + + private: + class Impl; + std::unique_ptr impl_; +}; +} // namespace infini_train::core diff --git a/infini_train/include/nn/parallel/process_group.h b/infini_train/include/nn/parallel/process_group.h index a790f986..97076aa7 100644 --- a/infini_train/include/nn/parallel/process_group.h +++ b/infini_train/include/nn/parallel/process_group.h @@ -139,7 +139,6 @@ class ProcessGroup { const std::string name_ = ""; - bool is_main_process_ = false; Device::DeviceType backend_ = Device::DeviceType::kInvalid; // Save impl for convenience diff --git a/infini_train/src/core/ccl/ccl_utils.cc b/infini_train/src/core/ccl/ccl_utils.cc index fc34b4f9..b4e59589 100644 --- a/infini_train/src/core/ccl/ccl_utils.cc +++ b/infini_train/src/core/ccl/ccl_utils.cc @@ -8,38 +8,26 @@ #include #include #include -#include #include -#include -#include -#include -#include #include +#include #include #include #include "glog/logging.h" #include "infini_train/include/core/ccl/ccl_common.h" +#include "infini_train/include/core/ccl/rendezvous_protocol.h" #include namespace infini_train::core { namespace { -constexpr uint32_t kRendezvousMagic = 0x49545244; -struct RendezvousHeader { - uint32_t magic, operation, key_size, value_size; -}; - int GetEnvInt(const char *name, int fallback) { const char *value = std::getenv(name); return value == nullptr ? fallback : std::stoi(value); } -bool UseTcpRendezvous() { - return GetEnvInt("WORLD_SIZE", 1) > GetEnvInt("LOCAL_WORLD_SIZE", 1); -} - bool ReadAll(int fd, void *buffer, size_t size) { auto *bytes = static_cast(buffer); while (size > 0) { @@ -99,22 +87,24 @@ int ConnectRendezvous() { return fd; } -std::string RendezvousRequest(uint32_t operation, const std::string &key, - std::string_view value) { +std::string RendezvousRequest(rendezvous::Operation operation, + const std::string &key, std::string_view value) { CHECK_LE(key.size(), std::numeric_limits::max()); CHECK_LE(value.size(), std::numeric_limits::max()); const int fd = ConnectRendezvous(); - const RendezvousHeader header{htonl(kRendezvousMagic), htonl(operation), - htonl(static_cast(key.size())), - htonl(static_cast(value.size()))}; + const rendezvous::Header header{htonl(rendezvous::kMagic), + htonl(static_cast(operation)), + htonl(static_cast(key.size())), + htonl(static_cast(value.size()))}; CHECK(WriteAll(fd, &header, sizeof(header)) && WriteAll(fd, key.data(), key.size()) && (value.empty() || WriteAll(fd, value.data(), value.size()))) << "Failed to send rendezvous request for " << key; - uint32_t response[2] = {}; - CHECK(ReadAll(fd, response, sizeof(response))); - CHECK_EQ(ntohl(response[0]), 0) << "Rendezvous request failed for " << key; - std::string result(ntohl(response[1]), '\0'); + rendezvous::Response response{}; + CHECK(ReadAll(fd, &response, sizeof(response))); + CHECK_EQ(ntohl(response.status), 0) + << "Rendezvous request failed for " << key; + std::string result(ntohl(response.value_size), '\0'); CHECK(result.empty() || ReadAll(fd, result.data(), result.size())); close(fd); return result; @@ -126,71 +116,20 @@ std::string RendezvousKey(const std::string &pg_name) { return std::string(run_id) + "/ccl/" + pg_name; } -std::string UniqueIdFileName(const std::string &pg_name) { - const char *run_id = std::getenv("INFINI_RUN_ID"); - const std::string prefix = run_id == nullptr ? "" : std::string(run_id) + "_"; - return "cclUniqueId_" + prefix + pg_name + ".bin"; -} - -std::string UniqueIdTmpFileName(const std::string &pg_name) { - const char *run_id = std::getenv("INFINI_RUN_ID"); - const std::string prefix = run_id == nullptr ? "" : std::string(run_id) + "_"; - return "cclUniqueId_" + prefix + pg_name + ".tmp"; -} } // namespace -void WriteUniqueIdFile(const CclUniqueId &unique_id, - const std::string &pg_name) { - if (UseTcpRendezvous()) { - RendezvousRequest( - 1, RendezvousKey(pg_name), - std::string_view(static_cast(unique_id.Data()), - unique_id.Size())); - return; - } - const std::string tmp_path = UniqueIdTmpFileName(pg_name); - std::ofstream ofs(tmp_path, std::ios::binary); - CHECK(ofs.good()) << "Failed to open unique_id tmp file for write: " - << tmp_path; - ofs.write(reinterpret_cast(unique_id.Data()), - static_cast(unique_id.Size())); - ofs.close(); - const std::string file_path = UniqueIdFileName(pg_name); - CHECK_EQ(std::rename(tmp_path.c_str(), file_path.c_str()), 0) - << "Failed to rename unique_id file from " << tmp_path << " to " - << file_path; +void WriteUniqueId(const CclUniqueId &unique_id, const std::string &pg_name) { + RendezvousRequest( + rendezvous::Operation::kPut, RendezvousKey(pg_name), + std::string_view(static_cast(unique_id.Data()), + unique_id.Size())); } -void ReadUniqueIdFile(CclUniqueId *unique_id, const std::string &pg_name) { +void ReadUniqueId(CclUniqueId *unique_id, const std::string &pg_name) { CHECK_NOTNULL(unique_id); - if (UseTcpRendezvous()) { - const std::string bytes = RendezvousRequest(2, RendezvousKey(pg_name), {}); - CHECK_EQ(bytes.size(), unique_id->Size()); - unique_id->Load(bytes.data(), bytes.size()); - return; - } - const std::string file_path = UniqueIdFileName(pg_name); - const int timeout = GetEnvInt("INFINI_RDZV_TIMEOUT_SEC", 120); - const auto deadline = - std::chrono::steady_clock::now() + std::chrono::seconds(timeout); - while (!std::filesystem::exists(file_path)) { - CHECK(std::chrono::steady_clock::now() < deadline) - << "Timed out waiting for unique-id file " << file_path; - std::this_thread::sleep_for(std::chrono::milliseconds(1)); - } - std::ifstream ifs(file_path, std::ios::binary); - CHECK(ifs.good()) << "Failed to open unique_id file for read: " << file_path; - std::string bytes((std::istreambuf_iterator(ifs)), - std::istreambuf_iterator()); + const std::string bytes = RendezvousRequest(rendezvous::Operation::kGet, + RendezvousKey(pg_name), {}); CHECK_EQ(bytes.size(), unique_id->Size()); unique_id->Load(bytes.data(), bytes.size()); } - -void CleanupUniqueIdFile(const std::string &pg_name) { - if (UseTcpRendezvous()) return; - for (const auto &path : - {UniqueIdFileName(pg_name), UniqueIdTmpFileName(pg_name)}) { - if (std::filesystem::exists(path)) std::filesystem::remove(path); - } -} } // namespace infini_train::core diff --git a/infini_train/src/core/ccl/rendezvous_server.cc b/infini_train/src/core/ccl/rendezvous_server.cc new file mode 100644 index 00000000..2b4241ac --- /dev/null +++ b/infini_train/src/core/ccl/rendezvous_server.cc @@ -0,0 +1,185 @@ +#include "infini_train/include/core/ccl/rendezvous_server.h" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/core/ccl/rendezvous_protocol.h" + +#include + +namespace infini_train::core { +namespace { +bool ReadAll(int fd, void *buffer, size_t size) { + auto *bytes = static_cast(buffer); + while (size > 0) { + const ssize_t count = read(fd, bytes, size); + if (count < 0 && errno == EINTR) continue; + if (count <= 0) return false; + bytes += count; + size -= static_cast(count); + } + return true; +} + +bool WriteAll(int fd, const void *buffer, size_t size) { + const auto *bytes = static_cast(buffer); + while (size > 0) { + const ssize_t count = write(fd, bytes, size); + if (count < 0 && errno == EINTR) continue; + if (count <= 0) return false; + bytes += count; + size -= static_cast(count); + } + return true; +} +} // namespace + +class RendezvousServer::Impl { + public: + Impl(const std::string &host, int port, int timeout) : timeout_(timeout) { + addrinfo hints{}; + hints.ai_family = AF_UNSPEC; + hints.ai_socktype = SOCK_STREAM; + addrinfo *addresses = nullptr; + const std::string port_text = std::to_string(port); + CHECK_EQ(getaddrinfo(host.c_str(), port_text.c_str(), &hints, &addresses), + 0) + << "Failed to resolve rendezvous bind address " << host; + for (addrinfo *address = addresses; address != nullptr; + address = address->ai_next) { + fd_ = socket(address->ai_family, address->ai_socktype, + address->ai_protocol); + if (fd_ < 0) continue; + int reuse = 1; + setsockopt(fd_, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)); + if (bind(fd_, address->ai_addr, address->ai_addrlen) == 0) break; + close(fd_); + fd_ = -1; + } + freeaddrinfo(addresses); + CHECK_GE(fd_, 0) << "Failed to bind rendezvous endpoint " << host << ':' + << port << ": " << strerror(errno); + CHECK_EQ(listen(fd_, 64), 0); + + sockaddr_storage bound_address{}; + socklen_t bound_address_size = sizeof(bound_address); + CHECK_EQ(getsockname(fd_, reinterpret_cast(&bound_address), + &bound_address_size), + 0); + if (bound_address.ss_family == AF_INET) { + port_ = ntohs(reinterpret_cast(&bound_address)->sin_port); + } else { + port_ = + ntohs(reinterpret_cast(&bound_address)->sin6_port); + } + thread_ = std::thread([this] { Serve(); }); + } + + ~Impl() { + shutdown(fd_, SHUT_RDWR); + close(fd_); + if (thread_.joinable()) thread_.join(); + for (auto &worker : workers_) { + if (worker.joinable()) worker.join(); + } + } + + int port() const { return port_; } + + private: + void Serve() { + while (true) { + const int client = accept(fd_, nullptr, nullptr); + if (client < 0) break; + workers_.emplace_back(&Impl::Handle, this, client); + } + } + + void Handle(int client) { + rendezvous::Header wire{}; + if (!ReadAll(client, &wire, sizeof(wire))) { + close(client); + return; + } + const auto operation = + static_cast(ntohl(wire.operation)); + const uint32_t key_size = ntohl(wire.key_size); + const uint32_t value_size = ntohl(wire.value_size); + if (ntohl(wire.magic) != rendezvous::kMagic || key_size == 0 || + key_size > rendezvous::kMaxKeySize || + value_size > rendezvous::kMaxValueSize) { + close(client); + return; + } + + std::string key(key_size, '\0'); + std::string value(value_size, '\0'); + if (!ReadAll(client, key.data(), key.size()) || + (!value.empty() && !ReadAll(client, value.data(), value.size()))) { + close(client); + return; + } + + uint32_t status = 0; + if (operation == rendezvous::Operation::kPut) { + { + std::lock_guard lock(mutex_); + values_[key] = std::move(value); + } + condition_.notify_all(); + } else if (operation == rendezvous::Operation::kGet) { + std::unique_lock lock(mutex_); + if (!condition_.wait_for(lock, std::chrono::seconds(timeout_), + [&] { return values_.contains(key); })) { + status = 1; + } else { + value = values_.at(key); + } + } else { + status = 2; + } + + const rendezvous::Response response{ + htonl(status), htonl(static_cast(value.size()))}; + WriteAll(client, &response, sizeof(response)); + if (status == 0 && operation == rendezvous::Operation::kGet) { + WriteAll(client, value.data(), value.size()); + } + close(client); + } + + int fd_ = -1; + int timeout_; + int port_ = -1; + std::thread thread_; + std::vector workers_; + std::mutex mutex_; + std::condition_variable condition_; + std::map values_; +}; + +RendezvousServer::RendezvousServer(const std::string &host, int port, + int timeout) + : impl_(std::make_unique(host, port, timeout)) {} + +RendezvousServer::~RendezvousServer() = default; + +int RendezvousServer::port() const { return impl_->port(); } +} // namespace infini_train::core diff --git a/infini_train/src/nn/parallel/process_group.cc b/infini_train/src/nn/parallel/process_group.cc index fc35f506..431f76b0 100644 --- a/infini_train/src/nn/parallel/process_group.cc +++ b/infini_train/src/nn/parallel/process_group.cc @@ -66,10 +66,6 @@ ProcessGroup::ProcessGroup(Device::DeviceType backend, } ProcessGroup::~ProcessGroup() { - if (is_main_process_) { - core::CleanupUniqueIdFile(name_); - } - // NOTE(zbl): The destruction could happen after GPU runtime is destroyed, // so do not explicitly destroy streams or comms here } @@ -107,10 +103,9 @@ void ProcessGroup::InitMultiProcess(const std::vector &ranks) { int min_rank = std::ranges::min(ranks); if (min_rank < upper_rank && min_rank >= lower_rank) { - is_main_process_ = true; - core::WriteUniqueIdFile(*unique_id, name_); + core::WriteUniqueId(*unique_id, name_); } else { - core::ReadUniqueIdFile(unique_id.get(), name_); + core::ReadUniqueId(unique_id.get(), name_); } const auto local_rank_count = diff --git a/scripts/multinode_ssh/README.md b/scripts/multinode_ssh/README.md deleted file mode 100644 index e9fbb8a6..00000000 --- a/scripts/multinode_ssh/README.md +++ /dev/null @@ -1,118 +0,0 @@ -# NVIDIA SSH Multinode Launcher - -`launch_multinode_ssh.sh` starts one `infini_run` per node from a control host. -A two-node, eight-GPU-per-node job creates 16 rank processes, records separate -rank logs, and returns nonzero if either node fails. - -## Preparation - -- Use the same repository commit and matching binaries on every node. -- Ensure non-interactive SSH and known-host verification work before launch. -- The remote account must be able to run the build, and optionally `docker - exec`, without relying on an interactive shell profile. -- Select an unused TCP port reachable from every node. Node rank 0 hosts the - rendezvous service; multi-node CCL unique-ID exchange does not require NFS. -- Set `NCCL_SOCKET_IFNAME` from the actual route. Confirm RDMA from NCCL logs; - successful communication alone is not proof that RDMA was used. - -A hostfile is plain text in node-rank order; blank lines and `#` comments are -ignored: - -```text -# node rank order -TRAIN_NODE_0 -TRAIN_NODE_1 -``` - -For a short list, `--hosts TRAIN_NODE_0,TRAIN_NODE_1` avoids maintaining a -separate file. - -## Two Nodes, Eight GPUs Each - -Run only on the control node. The second node needs no manual command. Set the -following environment-specific values first. `CONTROL_HOST` must be reachable by -every training node, and `NCCL_IFNAME` must be the interface selected by the route -to the peer. - -```bash -CONTROL_HOST=TRAIN_NODE_0 -REMOTE_HOST=TRAIN_NODE_1 -SSH_USER=TRAIN_USER -SSH_KEY=/path/to/private_key -REMOTE_CONTAINER=TRAINING_CONTAINER -WORKDIR=/path/to/InfiniTrain -NCCL_IFNAME=TRAINING_INTERFACE -LOG_ROOT=/path/to/controller/logs -REMOTE_LOG_ROOT=/path/to/node/logs - - -scripts/multinode_ssh/launch_multinode_ssh.sh \ - --hosts "$CONTROL_HOST,$REMOTE_HOST" \ - --local-host "$CONTROL_HOST" \ - --ssh-user "$SSH_USER" \ - --ssh-option "IdentityFile=$SSH_KEY" \ - --remote-docker-container "$REMOTE_CONTAINER" \ - --workdir "$WORKDIR" \ - --infini-run ./build/infini_run \ - --nproc-per-node 8 \ - --master-addr "$CONTROL_HOST" \ - --master-port 29500 \ - --rdzv-timeout 180 \ - --job-timeout 1800 \ - --log-root "$LOG_ROOT" \ - --remote-log-root "$REMOTE_LOG_ROOT" \ - --env "NCCL_SOCKET_IFNAME=$NCCL_IFNAME" \ - --env NCCL_DEBUG=INFO \ - --follow-output \ - -- ./build/gpt2 \ - --device cuda \ - --tensor_parallel 4 \ - --pipeline_parallel 2 \ - --batch_size 1 \ - --sequence_length 64 \ - --num_iteration 20 \ - --total_batch_size 16 \ - --input_bin /path/to/gpt2_train.bin \ - --llmc_filepath /path/to/gpt2_124M.bin -``` - -The tested PP2 command intentionally omits `--tokenizer_bin`: current text -sampling asserts that pipeline parallel size is 1 when generation runs. - -`--remote-docker-container` applies to SSH nodes. A host selected by -`--local-host` runs in the controller's current environment, so launch the -controller from the correct container when node 0 also uses a container. - -## Logs and Failure Handling - -Preflight checks the work directory, executable files, Git SHA, requested data -paths, binary SHA-256, `nvidia-smi`, and visible GPU count. The control log -contains metadata, the complete invocation, preflight output, and one SSH -wrapper log per node. Rank output is stored at: - -```text -REMOTE_LOG_ROOT/JOB_ID/node-N/rank-R.stdout.log -REMOTE_LOG_ROOT/JOB_ID/node-N/rank-R.stderr.log -``` - -`--follow-output` tails the final global rank's stderr on the controller while -preserving the original file. Training loss is normally emitted on stderr. -Every rank first prints hostname, global rank, local rank, world size, device, -and job ID. - -The launcher does not restart failed jobs. On node failure, SSH loss, timeout, -SIGINT, or SIGTERM, it reads this job's PID file, verifies the command contains -the job ID, and terminates only that process group. It never uses broad `pkill`. - -## Validation - -Run the local regression first: - -```bash -scripts/multinode_ssh/test_launch_multinode_ssh.sh -``` - -Then expand in order: single-node 8 ranks, two-node 1 rank each, two-node 8 -ranks each, short training, numerical comparison, and failure injection. Keep -`NCCL_DEBUG=INFO` for communication-path evidence and retain generated job -folders. Local simulation must not be reported as GPU or multi-node validation. diff --git a/scripts/multinode_ssh/VALIDATION_REPORT.md b/scripts/multinode_ssh/VALIDATION_REPORT.md deleted file mode 100644 index a6f4d1e7..00000000 --- a/scripts/multinode_ssh/VALIDATION_REPORT.md +++ /dev/null @@ -1,68 +0,0 @@ -# NVIDIA Two-Node Validation Report - -## Code and Environment - -- Date: 2026-11-02 -- Base: `origin/feat/ssh-multinode-launcher` at - `924d21256f3baf66bccb7458aa77395f0fbf5841` -- Observed master: `5a7be78e9b9d42c2f26efe33984dd3167dc60dc9` -- Development branch: `codex/ssh-multinode-nvidia-validation` -- Nodes: two private training hosts (identifiers intentionally omitted) -- GPUs: 8 x NVIDIA A100-SXM4-80GB per node; driver 535.161.08 -- SSH identity, host addresses, container name, worktree path, and interface: - intentionally omitted; configure them as described in `README.md` - -The branch remains based directly on the requested feature branch. Master had -broad unrelated changes, so it was not merged or rebased. - -## Results - -| Area | Result | Evidence | -|---|---|---| -| Parser, quoting, hostfile, invalid config | Passed | local launcher regression | -| Simulated success, node failure, timeout, SIGTERM cleanup | Passed | exits 0, 17, 124, 130; job PID checked | -| Live controller output | Passed | `--follow-output` observed start and final lines | -| TCP rendezvous PUT/GET | Passed | local protocol test against node-0 server | -| C++ build and rank tests | Passed | full build; rank CTest 2/2 | -| Two-node 1 process per node | Passed | real SSH smoke | -| Two-node 8 processes per node | Passed | 16 unique global ranks; local devices 0-7 | -| GPT2 TP4 x PP2 x DP2, 20 steps | Passed | final loss 4.468029 | -| Qwen3 TP4 x PP2 x DP2, 20 steps, BF16 | Passed | final loss 2.951596 | -| Llama3 TP4 x PP2 x DP2, 20 steps, BF16 | Passed | final loss 3.553573 | -| NCCL communication path | Passed for observed run | NCCL log reported `NET/IBext_v10` and GDRDMA | -| Real-node bad SSH / bad path | Partially verified | preflight/path failures verified; full matrix not retained | -| Real GPU rank failure / initialization timeout / Ctrl-C | Not verified | only simulated failure paths completed | -| Performance baseline | Not verified | no throughput target; correctness was prioritized | - -## Numerical Comparison - -Single-node references used TP4 x PP2 x DP1 with gradient accumulation adjusted -to preserve the global batch. Two-node runs used TP4 x PP2 x DP2. Initial -weights, data order, optimizer, global batch, and 20-step count were held fixed. - -| Model | Single-node final | Two-node final | Maximum absolute loss difference | -|---|---:|---:|---:| -| GPT2 | 4.468028 | 4.468029 | 0.000001 | -| Qwen3 BF16 | 2.950893 | 2.951596 | 0.004184 | -| Llama3 BF16 | 3.553813 | 3.553573 | 0.002242 | - -Maximum relative differences were 0.000022% (GPT2), 0.146558% (Qwen3), and -0.069305% (Llama3). Qwen3 and Llama3 use BF16 and different reduction order / -gradient accumulation between DP1 and DP2; these figures are recorded as -observations, not as a newly invented project tolerance. - -## Logs and Known Limits - -Control and rank log roots are operator-configured; their concrete validation -locations are intentionally omitted. The exact command is retained -in `controller-command.txt`, and runtime/version information is in metadata and -preflight logs. - -- PP2 text generation with `--tokenizer_bin` currently aborts when generation - is triggered because that path requires pipeline parallel size 1. Validated - PP2 commands omit tokenizer-based generation. -- The rendezvous service binds the selected node-0 port on all interfaces and - assumes a trusted training network; it is not an authenticated service. -- GPU preflight checks visibility/count, not exclusivity or free memory. -- More-node scaling still needs rendezvous load tests, larger failure matrices, - and a cluster log-retention policy. diff --git a/tests/distributed/CMakeLists.txt b/tests/distributed/CMakeLists.txt index b8ed4970..87507bc3 100644 --- a/tests/distributed/CMakeLists.txt +++ b/tests/distributed/CMakeLists.txt @@ -23,3 +23,14 @@ set_tests_properties(RankTest.MultiNodeSingleProcessIsParallel LABELS cpu TIMEOUT 10 ) + +add_test( + NAME MultinodeSshLauncherTest + COMMAND bash ${CMAKE_CURRENT_SOURCE_DIR}/test_launch_multinode_ssh.sh +) + +set_tests_properties(MultinodeSshLauncherTest + PROPERTIES + LABELS cpu + TIMEOUT 60 +) diff --git a/scripts/multinode_ssh/test_launch_multinode_ssh.sh b/tests/distributed/test_launch_multinode_ssh.sh similarity index 98% rename from scripts/multinode_ssh/test_launch_multinode_ssh.sh rename to tests/distributed/test_launch_multinode_ssh.sh index 6c9ea136..a3b62809 100755 --- a/scripts/multinode_ssh/test_launch_multinode_ssh.sh +++ b/tests/distributed/test_launch_multinode_ssh.sh @@ -2,8 +2,8 @@ set -euo pipefail script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) -launcher="$script_dir/launch_multinode_ssh.sh" repo=$(cd "$script_dir/../.." && pwd) +launcher="$repo/scripts/multinode_ssh/launch_multinode_ssh.sh" tmp=$(mktemp -d) trap 'rm -rf "$tmp"' EXIT diff --git a/tools/infini_run/CMakeLists.txt b/tools/infini_run/CMakeLists.txt index edc01bab..01c97e0a 100644 --- a/tools/infini_run/CMakeLists.txt +++ b/tools/infini_run/CMakeLists.txt @@ -1,2 +1,6 @@ add_executable(infini_run infini_run.cc) -target_link_libraries(infini_run PRIVATE gflags glog) +target_link_libraries(infini_run PRIVATE + InfiniTrain::rendezvous + gflags + glog +) diff --git a/tools/infini_run/infini_run.cc b/tools/infini_run/infini_run.cc index 8b634757..caa4e723 100644 --- a/tools/infini_run/infini_run.cc +++ b/tools/infini_run/infini_run.cc @@ -1,9 +1,5 @@ // IWYU pragma: no_include -#include #include -#include -#include -#include #include #include @@ -11,37 +7,33 @@ #include #include #include -#include #include #include #include -#include #include -#include -#include #include -#include #include #include -#include #include #include -#include #include #include "gflags/gflags.h" #include "glog/logging.h" +#include "infini_train/include/core/ccl/rendezvous_server.h" + DEFINE_int32(nnodes, 1, "Total number of nodes"); DEFINE_int32(nproc_per_node, 1, "Number of processes per node"); DEFINE_int32(node_rank, 0, "Rank of this node"); -DEFINE_string(rdzv_endpoint, "127.0.0.1:29500", "Rendezvous endpoint"); +DEFINE_string(rdzv_endpoint, "127.0.0.1:29500", + "Rendezvous endpoint (host:port)"); + DEFINE_string(rdzv_id, "", "Unique job ID shared by all nodes"); DEFINE_int32(rdzv_timeout, 120, "Rendezvous timeout in seconds"); DEFINE_string(log_dir, "", "Directory for per-rank stdout/stderr logs"); namespace { -constexpr uint32_t kMagic = 0x49545244; volatile std::sig_atomic_t g_signal = 0; void HandleSignal(int signal) { g_signal = signal; } @@ -50,179 +42,99 @@ bool IsLauncherFlag(const std::string &flag) { flag == "--node_rank" || flag == "--rdzv_endpoint" || flag == "--rdzv_id" || flag == "--rdzv_timeout" || flag == "--log_dir"; } + int FindTrainProgramIndex(int argc, char **argv) { for (int i = 1; i < argc; ++i) { const std::string arg = argv[i]; - if (arg == "--") return i + 1; - if (!arg.starts_with("--")) return i; - const size_t eq = arg.find('='); - if (IsLauncherFlag(arg.substr(0, eq)) && eq == std::string::npos) ++i; + if (arg == "--") { + return i + 1; + } + if (arg.rfind("--", 0) != 0) { + return i; + } + + const size_t eq_pos = arg.find('='); + const std::string flag = arg.substr(0, eq_pos); + if (IsLauncherFlag(flag) && eq_pos == std::string::npos) { + ++i; + } } return argc; } -std::vector BuildLauncherArgv(int end, char **argv) { - std::vector result{argv[0]}; - for (int i = 1; i < end; ++i) result.push_back(argv[i]); - result.push_back(nullptr); - return result; -} -void SetEnvInt(const char *name, int value) { - const std::string text = std::to_string(value); - setenv(name, text.c_str(), 1); -} -bool ReadAll(int fd, void *buffer, size_t size) { - auto *bytes = static_cast(buffer); - while (size > 0) { - ssize_t n = read(fd, bytes, size); - if (n < 0 && errno == EINTR) continue; - if (n <= 0) return false; - bytes += n; - size -= n; + +std::vector BuildLauncherArgv(int train_program_index, char **argv) { + std::vector launcher_argv; + launcher_argv.push_back(argv[0]); + for (int i = 1; i < train_program_index; ++i) { + launcher_argv.push_back(argv[i]); } - return true; + launcher_argv.push_back(nullptr); + return launcher_argv; } -bool WriteAll(int fd, const void *buffer, size_t size) { - const auto *bytes = static_cast(buffer); - while (size > 0) { - ssize_t n = write(fd, bytes, size); - if (n < 0 && errno == EINTR) continue; - if (n <= 0) return false; - bytes += n; - size -= n; - } - return true; + +void SetEnvInt(const char *name, int value) { + const auto value_str = std::to_string(value); + setenv(name, value_str.c_str(), 1); } -class RendezvousServer { - public: - RendezvousServer(int port, int timeout) : timeout_(timeout) { - fd_ = socket(AF_INET, SOCK_STREAM, 0); - CHECK_GE(fd_, 0); - int reuse = 1; - setsockopt(fd_, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)); - sockaddr_in address{}; - address.sin_family = AF_INET; - address.sin_addr.s_addr = htonl(INADDR_ANY); - address.sin_port = htons(port); - CHECK_EQ(bind(fd_, reinterpret_cast(&address), sizeof(address)), - 0) - << "Failed to bind rendezvous port " << port << ": " << strerror(errno); - CHECK_EQ(listen(fd_, 64), 0); - thread_ = std::thread([this] { Serve(); }); - } - ~RendezvousServer() { - shutdown(fd_, SHUT_RDWR); - close(fd_); - if (thread_.joinable()) thread_.join(); - for (auto &worker : workers_) - if (worker.joinable()) worker.join(); +void TerminateChildren(const std::unordered_set &child_pids, + int signal = SIGTERM) { + for (pid_t child_pid : child_pids) { + kill(child_pid, signal); } +} - private: - struct Header { - uint32_t magic, operation, key_size, value_size; - }; - void Serve() { - while (true) { - int client = accept(fd_, nullptr, nullptr); - if (client < 0) break; - workers_.emplace_back(&RendezvousServer::Handle, this, client); - } +int ExitCodeFromStatus(int status) { + if (WIFEXITED(status)) { + return WEXITSTATUS(status); } - void Handle(int client) { - Header wire{}; - if (!ReadAll(client, &wire, sizeof(wire))) { - close(client); - return; - } - uint32_t operation = ntohl(wire.operation), key_size = ntohl(wire.key_size), - value_size = ntohl(wire.value_size); - if (ntohl(wire.magic) != kMagic || key_size == 0 || key_size > 4096 || - value_size > 1024 * 1024) { - close(client); - return; - } - std::string key(key_size, '\0'), value(value_size, '\0'); - if (!ReadAll(client, key.data(), key.size()) || - (!value.empty() && !ReadAll(client, value.data(), value.size()))) { - close(client); - return; - } - uint32_t status = 0; - if (operation == 1) { - { - std::lock_guard lock(mutex_); - values_[key] = std::move(value); - } - condition_.notify_all(); - } else if (operation == 2) { - std::unique_lock lock(mutex_); - if (!condition_.wait_for(lock, std::chrono::seconds(timeout_), - [&] { return values_.contains(key); })) - status = 1; - else - value = values_.at(key); - } else - status = 2; - CHECK_LE(value.size(), std::numeric_limits::max()); - uint32_t response[2] = {htonl(status), - htonl(static_cast(value.size()))}; - WriteAll(client, response, sizeof(response)); - if (status == 0 && operation == 2) - WriteAll(client, value.data(), value.size()); - close(client); + if (WIFSIGNALED(status)) { + return 128 + WTERMSIG(status); } - int fd_ = -1, timeout_; - std::thread thread_; - std::vector workers_; - std::mutex mutex_; - std::condition_variable condition_; - std::map values_; -}; - -int ExitCode(int status) { - if (WIFEXITED(status)) return WEXITSTATUS(status); - return WIFSIGNALED(status) ? 128 + WTERMSIG(status) : 1; + return 1; } -void SignalChildren(const std::unordered_set &children, int signal) { - for (pid_t pid : children) kill(pid, signal); -} -std::string LocalRunId() { - return std::to_string(getpid()) + "_" + - std::to_string( - std::chrono::steady_clock::now().time_since_epoch().count()); -} -void CleanupFiles(const std::string &run_id) { - const std::string prefix = "cclUniqueId_" + run_id + "_"; - for (const auto &entry : std::filesystem::directory_iterator(".")) { - if (entry.is_regular_file() && - entry.path().filename().string().starts_with(prefix)) { - std::error_code error; - std::filesystem::remove(entry.path(), error); - } - } + +std::string GenerateLocalRunId() { + const auto now = std::chrono::steady_clock::now().time_since_epoch().count(); + return std::to_string(getpid()) + "_" + std::to_string(now); } + } // namespace int main(int argc, char **argv) { - const int train_index = FindTrainProgramIndex(argc, argv); - auto launcher_argv = BuildLauncherArgv(train_index, argv); - int launcher_argc = launcher_argv.size() - 1; + const int train_program_index = FindTrainProgramIndex(argc, argv); + std::vector launcher_argv = + BuildLauncherArgv(train_program_index, argv); + int launcher_argc = static_cast(launcher_argv.size()) - 1; char **launcher_argv_ptr = launcher_argv.data(); gflags::ParseCommandLineFlags(&launcher_argc, &launcher_argv_ptr, true); google::InitGoogleLogging(argv[0]); - CHECK_GT(FLAGS_nnodes, 0); - CHECK_GT(FLAGS_nproc_per_node, 0); - CHECK_GE(FLAGS_node_rank, 0); - CHECK_LT(FLAGS_node_rank, FLAGS_nnodes); - CHECK_GT(FLAGS_rdzv_timeout, 0); - CHECK_LT(train_index, argc); - const size_t separator = FLAGS_rdzv_endpoint.rfind(':'); - CHECK_NE(separator, std::string::npos) << "rdzv_endpoint must be host:port"; - CHECK(FLAGS_nnodes == 1 || !FLAGS_rdzv_id.empty()); - const std::string master_addr = FLAGS_rdzv_endpoint.substr(0, separator); - const std::string master_port = FLAGS_rdzv_endpoint.substr(separator + 1); + CHECK_GT(FLAGS_nnodes, 0) << "nnodes must be positive"; + CHECK_GT(FLAGS_nproc_per_node, 0) << "nproc_per_node must be positive"; + CHECK_GE(FLAGS_node_rank, 0) << "node_rank must be non-negative"; + CHECK_LT(FLAGS_node_rank, FLAGS_nnodes) + << "node_rank must be less than nnodes"; + CHECK_GT(FLAGS_rdzv_timeout, 0) << "rdzv_timeout must be positive"; + CHECK_NE(FLAGS_rdzv_endpoint.find(':'), std::string::npos) + << "rdzv_endpoint must be host:port"; + CHECK(FLAGS_nnodes == 1 || !FLAGS_rdzv_id.empty()) + << "rdzv_id must be set to the same unique job ID on every node for " + "multi-node training"; + + CHECK_LT(train_program_index, argc) << "No training program specified!"; + + std::string train_program = argv[train_program_index]; + std::vector train_argv; + for (int i = train_program_index; i < argc; ++i) { + train_argv.push_back(argv[i]); + } + train_argv.push_back(nullptr); + + int proc_world_size = FLAGS_nnodes * FLAGS_nproc_per_node; + const size_t separator = FLAGS_rdzv_endpoint.rfind(':'); + std::string master_addr = FLAGS_rdzv_endpoint.substr(0, separator); + std::string master_port = FLAGS_rdzv_endpoint.substr(separator + 1); char *port_end = nullptr; errno = 0; const long parsed_port = std::strtol(master_port.c_str(), &port_end, 10); @@ -230,14 +142,22 @@ int main(int argc, char **argv) { parsed_port > 0 && parsed_port <= 65535) << "Invalid rendezvous port: " << master_port; const std::string run_id = - FLAGS_rdzv_id.empty() ? LocalRunId() : FLAGS_rdzv_id; - setenv("INFINI_RUN_ID", run_id.c_str(), 1); - setenv("INFINI_RDZV_TIMEOUT_SEC", std::to_string(FLAGS_rdzv_timeout).c_str(), - 1); - std::unique_ptr server; - if (FLAGS_nnodes > 1 && FLAGS_node_rank == 0) - server = std::make_unique(static_cast(parsed_port), - FLAGS_rdzv_timeout); + FLAGS_rdzv_id.empty() ? GenerateLocalRunId() : FLAGS_rdzv_id; + const bool endpoint_is_default = + gflags::GetCommandLineFlagInfoOrDie("rdzv_endpoint").is_default; + int effective_master_port = static_cast(parsed_port); + std::unique_ptr server; + if (proc_world_size > 1 && FLAGS_node_rank == 0) { + if (FLAGS_nnodes == 1 && endpoint_is_default) { + master_addr = "127.0.0.1"; + effective_master_port = 0; + } + server = std::make_unique( + master_addr, effective_master_port, FLAGS_rdzv_timeout); + effective_master_port = server->port(); + } + const std::string effective_master_port_text = + std::to_string(effective_master_port); if (!FLAGS_log_dir.empty()) std::filesystem::create_directories(FLAGS_log_dir); struct sigaction action {}; @@ -246,32 +166,38 @@ int main(int argc, char **argv) { sigaction(SIGINT, &action, nullptr); sigaction(SIGTERM, &action, nullptr); - std::vector train_argv; - for (int i = train_index; i < argc; ++i) train_argv.push_back(argv[i]); - train_argv.push_back(nullptr); - std::unordered_set children; + std::unordered_set running_children; int exit_code = 0; - const int world_size = FLAGS_nnodes * FLAGS_nproc_per_node; - for (int local_rank = 0; local_rank < FLAGS_nproc_per_node; ++local_rank) { + + for (int local_proc_rank = 0; local_proc_rank < FLAGS_nproc_per_node; + ++local_proc_rank) { pid_t pid = fork(); if (pid < 0) { - PLOG(ERROR) << "fork failed after starting " << children.size() + PLOG(ERROR) << "fork failed after starting " << running_children.size() << " rank processes"; - SignalChildren(children, SIGTERM); - for (pid_t child : children) waitpid(child, nullptr, 0); + TerminateChildren(running_children); + for (pid_t child : running_children) waitpid(child, nullptr, 0); return 1; } if (pid == 0) { - const int rank = FLAGS_node_rank * FLAGS_nproc_per_node + local_rank; + int global_proc_rank = + FLAGS_node_rank * FLAGS_nproc_per_node + local_proc_rank; SetEnvInt("LOCAL_WORLD_SIZE", FLAGS_nproc_per_node); - SetEnvInt("WORLD_SIZE", world_size); - SetEnvInt("RANK", rank); - SetEnvInt("LOCAL_RANK", local_rank); + setenv("MASTER_ADDR", master_addr.c_str(), 1); - setenv("MASTER_PORT", master_port.c_str(), 1); + setenv("MASTER_PORT", effective_master_port_text.c_str(), 1); + setenv("INFINI_RUN_ID", run_id.c_str(), 1); + setenv("INFINI_RDZV_TIMEOUT_SEC", + std::to_string(FLAGS_rdzv_timeout).c_str(), 1); + + SetEnvInt("RANK", global_proc_rank); + SetEnvInt("LOCAL_RANK", local_proc_rank); + + SetEnvInt("WORLD_SIZE", proc_world_size); + if (!FLAGS_log_dir.empty()) { const std::string prefix = - FLAGS_log_dir + "/rank-" + std::to_string(rank); + FLAGS_log_dir + "/rank-" + std::to_string(global_proc_rank); int out = open((prefix + ".stdout.log").c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644); int err = open((prefix + ".stderr.log").c_str(), @@ -288,48 +214,55 @@ int main(int argc, char **argv) { fprintf(stderr, "INFINITRAIN_RANK hostname=%s global_rank=%d local_rank=%d " "world_size=%d device=%d job_id=%s\n", - hostname, rank, local_rank, world_size, local_rank, - run_id.c_str()); + hostname, global_proc_rank, local_proc_rank, proc_world_size, + local_proc_rank, run_id.c_str()); fflush(stderr); - execvp(train_argv[0], train_argv.data()); + execvp(train_program.c_str(), train_argv.data()); perror("exec failed"); _exit(127); } - children.insert(pid); + running_children.insert(pid); } bool term_sent = false; auto term_time = std::chrono::steady_clock::now(); - while (!children.empty()) { + while (!running_children.empty()) { if (g_signal && !term_sent) { exit_code = 128 + g_signal; - SignalChildren(children, SIGTERM); + TerminateChildren(running_children); term_sent = true; term_time = std::chrono::steady_clock::now(); } if (term_sent && std::chrono::steady_clock::now() - term_time > std::chrono::seconds(10)) - SignalChildren(children, SIGKILL); - int status = 0; + TerminateChildren(running_children, SIGKILL); + int status; pid_t child = waitpid(-1, &status, term_sent ? WNOHANG : 0); if (child == 0) { std::this_thread::sleep_for(std::chrono::milliseconds(100)); continue; } if (child < 0) { - if (errno == EINTR) continue; - exit_code = 1; + if (errno == EINTR) { + continue; + } + perror("wait failed"); + if (exit_code == 0) { + exit_code = 1; + } + TerminateChildren(running_children); break; } - children.erase(child); - const int code = ExitCode(status); - if (code != 0 && exit_code == 0) { - exit_code = code; - SignalChildren(children, SIGTERM); + + running_children.erase(child); + const int child_exit_code = ExitCodeFromStatus(status); + if (child_exit_code != 0 && exit_code == 0) { + exit_code = child_exit_code; + TerminateChildren(running_children); term_sent = true; term_time = std::chrono::steady_clock::now(); } } - if (FLAGS_nnodes == 1) CleanupFiles(run_id); + return exit_code; }