diff --git a/CMakeLists.txt b/CMakeLists.txt index 88911e84a..0c0a2c3f7 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 d9f8ce6be..299b71dc4 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 000000000..83a6a2952 --- /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 000000000..f30206806 --- /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 a790f9861..97076aa7d 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 d5df75df7..b4e595896 100644 --- a/infini_train/src/core/ccl/ccl_utils.cc +++ b/infini_train/src/core/ccl/ccl_utils.cc @@ -1,81 +1,135 @@ #include "infini_train/include/core/ccl/ccl_utils.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/ccl_common.h" +#include "infini_train/include/core/ccl/rendezvous_protocol.h" + +#include namespace infini_train::core { namespace { -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"; +int GetEnvInt(const char *name, int fallback) { + const char *value = std::getenv(name); + return value == nullptr ? fallback : std::stoi(value); } -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"; +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; } -} // namespace -void WriteUniqueIdFile(const CclUniqueId &unique_id, - const std::string &pg_name) { - 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)); - 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; +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; } -void ReadUniqueIdFile(CclUniqueId *unique_id, const std::string &pg_name) { - CHECK_NOTNULL(unique_id); - const std::string file_path = UniqueIdFileName(pg_name); - - while (!std::filesystem::exists(file_path)) { - std::this_thread::sleep_for(std::chrono::microseconds(1000)); +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::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(); +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 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; + 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; +} - CHECK_EQ(bytes.size(), unique_id->Size()) - << "Mismatched unique_id size in file. expected=" << unique_id->Size() - << ", got=" << bytes.size(); - unique_id->Load(bytes.data(), bytes.size()); +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; } -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); - } +} // namespace - const std::string tmp_path = UniqueIdTmpFileName(pg_name); - if (std::filesystem::exists(tmp_path)) { - std::filesystem::remove(tmp_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 ReadUniqueId(CclUniqueId *unique_id, const std::string &pg_name) { + CHECK_NOTNULL(unique_id); + 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()); +} } // 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 000000000..2b4241ac0 --- /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 fc35f5060..431f76b09 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/launch_multinode_ssh.sh b/scripts/multinode_ssh/launch_multinode_ssh.sh new file mode 100755 index 000000000..51a14a383 --- /dev/null +++ b/scripts/multinode_ssh/launch_multinode_ssh.sh @@ -0,0 +1,194 @@ +#!/usr/bin/env bash +set -euo pipefail + +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) 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" +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 + [[ -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" +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 -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 + [[ -z $ssh_user ]] || destination="$ssh_user@$host" + 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 +} + +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" +} + +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= + +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 + +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 +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/tests/distributed/CMakeLists.txt b/tests/distributed/CMakeLists.txt index b8ed49700..87507bc37 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/tests/distributed/test_launch_multinode_ssh.sh b/tests/distributed/test_launch_multinode_ssh.sh new file mode 100755 index 000000000..a3b628092 --- /dev/null +++ b/tests/distributed/test_launch_multinode_ssh.sh @@ -0,0 +1,135 @@ +#!/usr/bin/env bash +set -euo pipefail + +script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +repo=$(cd "$script_dir/../.." && pwd) +launcher="$repo/scripts/multinode_ssh/launch_multinode_ssh.sh" +tmp=$(mktemp -d) +trap 'rm -rf "$tmp"' EXIT + +cat >"$tmp/fake_infini_run" <<'RUN' +#!/usr/bin/env bash +set -euo pipefail +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" + +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" + +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" + +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" + +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 + +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 --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 + +echo "launch_multinode_ssh tests passed" diff --git a/tools/infini_run/CMakeLists.txt b/tools/infini_run/CMakeLists.txt index edc01babb..01c97e0a3 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 51ad1e0c2..caa4e7237 100644 --- a/tools/infini_run/infini_run.cc +++ b/tools/infini_run/infini_run.cc @@ -1,22 +1,28 @@ +// IWYU pragma: no_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"); @@ -24,13 +30,17 @@ 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 { +volatile std::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) { @@ -67,9 +77,10 @@ void SetEnvInt(const char *name, int value) { setenv(name, value_str.c_str(), 1); } -void TerminateChildren(const std::unordered_set &child_pids) { +void TerminateChildren(const std::unordered_set &child_pids, + int signal = SIGTERM) { for (pid_t child_pid : child_pids) { - kill(child_pid, SIGTERM); + kill(child_pid, signal); } } @@ -83,25 +94,6 @@ int ExitCodeFromStatus(int status) { 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; - } - 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(); - } - } - } -} - std::string GenerateLocalRunId() { const auto now = std::chrono::steady_clock::now().time_since_epoch().count(); return std::to_string(getpid()) + "_" + std::to_string(now); @@ -123,6 +115,7 @@ int main(int argc, char **argv) { 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()) @@ -139,12 +132,39 @@ int main(int argc, char **argv) { 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 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); + 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() ? 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 {}; + action.sa_handler = HandleSignal; + sigemptyset(&action.sa_mask); + sigaction(SIGINT, &action, nullptr); + sigaction(SIGTERM, &action, nullptr); std::unordered_set running_children; int exit_code = 0; @@ -153,10 +173,11 @@ int main(int argc, char **argv) { ++local_proc_rank) { pid_t pid = fork(); if (pid < 0) { - perror("fork failed"); - exit_code = 1; + PLOG(ERROR) << "fork failed after starting " << running_children.size() + << " rank processes"; TerminateChildren(running_children); - break; + for (pid_t child : running_children) waitpid(child, nullptr, 0); + return 1; } if (pid == 0) { int global_proc_rank = @@ -164,24 +185,63 @@ int main(int argc, char **argv) { SetEnvInt("LOCAL_WORLD_SIZE", FLAGS_nproc_per_node); 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(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(), + 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, global_proc_rank, local_proc_rank, proc_world_size, + local_proc_rank, run_id.c_str()); + fflush(stderr); execvp(train_program.c_str(), train_argv.data()); perror("exec failed"); - exit(1); + _exit(127); } running_children.insert(pid); } + bool term_sent = false; + auto term_time = std::chrono::steady_clock::now(); while (!running_children.empty()) { + if (g_signal && !term_sent) { + exit_code = 128 + g_signal; + 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)) + TerminateChildren(running_children, SIGKILL); int status; - pid_t child = wait(&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; @@ -199,11 +259,10 @@ int main(int argc, char **argv) { 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) { - CleanupRunUniqueIdFiles(run_id); - } return exit_code; }