126 lines
4.1 KiB
Bash
Executable File
126 lines
4.1 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
|
|
TRAINING_STATE_FILE="${TRAINING_STATE_FILE:-${REPO_ROOT}/.runtime/active_training.txt}"
|
|
CONTAINER_NAME="${1:-${CONTAINER_NAME:-laoyao-2b-pretraining}}"
|
|
CHECKPOINT_DIR="${CHECKPOINT_DIR:-}"
|
|
GRACEFUL_SIGNAL="${GRACEFUL_SIGNAL:-INT}"
|
|
EXPECTED_WORKERS="${EXPECTED_WORKERS:-8}"
|
|
WORKER_DISCOVERY_WAIT_SECONDS="${WORKER_DISCOVERY_WAIT_SECONDS:-120}"
|
|
WAIT_SECONDS="${WAIT_SECONDS:-900}"
|
|
POLL_SECONDS="${POLL_SECONDS:-5}"
|
|
|
|
read_state_value() {
|
|
local key="$1"
|
|
awk -v key="$key" '
|
|
index($0, key "=") == 1 {
|
|
sub(/^[^=]*=/, "")
|
|
print
|
|
exit
|
|
}
|
|
' "$TRAINING_STATE_FILE"
|
|
}
|
|
|
|
if [[ -f "$TRAINING_STATE_FILE" ]]; then
|
|
state_container_name="$(read_state_value container_name)"
|
|
if [[ -z "$CHECKPOINT_DIR" && "$state_container_name" == "$CONTAINER_NAME" ]]; then
|
|
CHECKPOINT_DIR="$(read_state_value checkpoint_dir)"
|
|
fi
|
|
fi
|
|
|
|
case "$GRACEFUL_SIGNAL" in
|
|
INT|SIGINT) GRACEFUL_SIGNAL=INT ;;
|
|
TERM|SIGTERM) GRACEFUL_SIGNAL=TERM ;;
|
|
*)
|
|
echo "ERROR: GRACEFUL_SIGNAL must be INT or TERM, got: $GRACEFUL_SIGNAL" >&2
|
|
exit 2
|
|
;;
|
|
esac
|
|
|
|
if [[ "$(docker inspect -f '{{.State.Running}}' "$CONTAINER_NAME" 2>/dev/null || true)" != "true" ]]; then
|
|
echo "ERROR: container is not running: $CONTAINER_NAME" >&2
|
|
exit 3
|
|
fi
|
|
|
|
discovery_deadline=$((SECONDS + WORKER_DISCOVERY_WAIT_SECONDS))
|
|
torchrun_pid=""
|
|
worker_pids=()
|
|
while true; do
|
|
container_top="$(docker top "$CONTAINER_NAME" -eo pid,ppid,args)"
|
|
torchrun_pid="$(
|
|
awk '/[/]usr[/]local[/]bin[/]torchrun|torchrun --nproc_per_node/ {print $1; exit}' <<<"$container_top"
|
|
)"
|
|
worker_pids=()
|
|
if [[ -n "$torchrun_pid" ]]; then
|
|
mapfile -t worker_pids < <(
|
|
awk -v parent="$torchrun_pid" \
|
|
'$2 == parent && /laoyao_2b_moe_pretrain[.]py/ {print $1}' \
|
|
<<<"$container_top"
|
|
)
|
|
fi
|
|
|
|
if [[ "${#worker_pids[@]}" -eq "$EXPECTED_WORKERS" ]]; then
|
|
break
|
|
fi
|
|
if [[ "$(docker inspect -f '{{.State.Running}}' "$CONTAINER_NAME" 2>/dev/null || true)" != "true" ]]; then
|
|
echo "ERROR: container exited before training workers became ready: $CONTAINER_NAME" >&2
|
|
exit 4
|
|
fi
|
|
if (( SECONDS >= discovery_deadline )); then
|
|
echo "ERROR: expected $EXPECTED_WORKERS direct torchrun workers, found ${#worker_pids[@]}" >&2
|
|
printf 'candidate process tree:\n%s\n' "$container_top" >&2
|
|
exit 5
|
|
fi
|
|
sleep 1
|
|
done
|
|
|
|
tracker="${CHECKPOINT_DIR:+${CHECKPOINT_DIR}/latest_checkpointed_iteration.txt}"
|
|
before_checkpoint="unknown"
|
|
if [[ -n "$tracker" && -f "$tracker" ]]; then
|
|
before_checkpoint="$(tr -d '[:space:]' < "$tracker")"
|
|
fi
|
|
|
|
echo "container=$CONTAINER_NAME"
|
|
echo "training_state=$TRAINING_STATE_FILE"
|
|
echo "torchrun_pid=$torchrun_pid"
|
|
echo "worker_pids=${worker_pids[*]}"
|
|
echo "signal=SIG${GRACEFUL_SIGNAL}"
|
|
echo "checkpoint_before=$before_checkpoint"
|
|
echo "Sending the signal to rank workers only; torchrun PID 1 remains alive while checkpointing."
|
|
kill "-${GRACEFUL_SIGNAL}" "${worker_pids[@]}"
|
|
|
|
deadline=$((SECONDS + WAIT_SECONDS))
|
|
while [[ "$(docker inspect -f '{{.State.Running}}' "$CONTAINER_NAME" 2>/dev/null || true)" == "true" ]]; do
|
|
if (( SECONDS >= deadline )); then
|
|
echo "ERROR: graceful shutdown did not finish within ${WAIT_SECONDS}s." >&2
|
|
echo "The script will not force-kill the container while a checkpoint may be in progress." >&2
|
|
exit 6
|
|
fi
|
|
|
|
current_checkpoint="$before_checkpoint"
|
|
if [[ -n "$tracker" && -f "$tracker" ]]; then
|
|
current_checkpoint="$(tr -d '[:space:]' < "$tracker")"
|
|
fi
|
|
echo "waiting: elapsed=${SECONDS}s checkpoint=${current_checkpoint}"
|
|
sleep "$POLL_SECONDS"
|
|
done
|
|
|
|
after_checkpoint="unknown"
|
|
if [[ -n "$tracker" && -f "$tracker" ]]; then
|
|
after_checkpoint="$(tr -d '[:space:]' < "$tracker")"
|
|
fi
|
|
|
|
echo "container_stopped=true"
|
|
echo "checkpoint_after=$after_checkpoint"
|
|
if [[ -n "$tracker" && "$after_checkpoint" == "$before_checkpoint" ]]; then
|
|
echo "ERROR: container exited without advancing the checkpoint tracker." >&2
|
|
exit 7
|
|
fi
|
|
|
|
if [[ -f "$TRAINING_STATE_FILE" && "$(read_state_value container_name)" == "$CONTAINER_NAME" ]]; then
|
|
rm -f "$TRAINING_STATE_FILE"
|
|
echo "active_training_state_removed=$TRAINING_STATE_FILE"
|
|
fi
|