Use a fixed training container and runtime state

2026-07-23 19:39:07 +08:00
parent 98c159679b
commit 54ac00152a
12 changed files with 152 additions and 33 deletions
@@ -5,6 +5,7 @@ checkpoints/
 outputs/
 wandb/
 tensorboard/
+.runtime/
 *.log
 *.pt
 *.ckpt
@@ -252,14 +252,16 @@ bash scripts/resume_pretrain_8192_8gpu_mbs14.sh
 Gracefully stop training and save a checkpoint at the next iteration boundary:

 ```bash
-CHECKPOINT_DIR=/mnt/beegfs/yi/laoyao-2b-pretraining/runs/<run>/checkpoints \
-bash scripts/graceful_stop_megatron_training.sh <container-name>
+bash scripts/graceful_stop_megatron_training.sh
 ```

 The Python training entrypoint listens for `SIGINT` by default through
 Megatron-Bridge's distributed signal handler. The helper signals direct torchrun
 workers and intentionally leaves PID 1 alive while the checkpoint is written.
-Do not use `docker kill --signal=SIGINT` for a routine stop.
+All long-running launchers use the fixed default container name
+`laoyao-2b-pretraining`. They also atomically record the active checkpoint
+directory in `.runtime/active_training.txt`, so the stop command needs no
+arguments. Do not use `docker kill --signal=SIGINT` for a routine stop.

 Before stopping, check the latest checkpoint:

@@ -16,6 +16,8 @@
 - `stop_laoyao_megatron_server.sh`: 停止 Megatron inference server 容器。
 - `graceful_stop_megatron_training.sh`: 向 torchrun 的训练 worker 发送信号，在当前
  iteration 完成后临时保存 checkpoint 并安全退出。
+- `record_active_training.sh`: launcher 启动容器后原子记录当前 container、run 和
+  checkpoint 路径，供状态检查与无参数 graceful stop 使用。

 ## g0050 下载与部署

@@ -97,13 +99,33 @@ handler，并监听 `SIGINT`。收到信号后，各 rank 会：
 写完前终止 rank。应运行：

 ```bash
-CHECKPOINT_DIR=/mnt/beegfs/yi/laoyao-2b-pretraining/runs/<run>/checkpoints \
-bash scripts/graceful_stop_megatron_training.sh <container-name>
+bash scripts/graceful_stop_megatron_training.sh
 ```

 脚本通过 `docker top` 只选择 torchrun 的直接 worker 子进程，不会误发给
 dataloader worker，也不会先终止 torchrun。默认等待 900 秒且不会在超时后自动
-强杀容器。训练入口也支持显式切换信号：
+强杀容器。状态默认读取：
+
+```text
+.runtime/active_training.txt
+```
+
+所有长期训练 launcher 的默认容器名固定为：
+
+```text
+laoyao-2b-pretraining
+```
+
+训练 launcher 确认容器启动后还会原子更新状态文件，记录 checkpoint 目录、run
+目录和日志路径。即使状态文件缺失，停止脚本也会操作上述固定容器名。多实验并存时
+仍可显式覆盖：
+
+```bash
+CHECKPOINT_DIR=/path/to/checkpoints \
+bash scripts/graceful_stop_megatron_training.sh <container-name>
+```
+
+训练入口也支持显式切换信号：

 ```text
 --graceful-exit-signal SIGINT   # 默认
@@ -1,13 +1,35 @@
 #!/usr/bin/env bash
 set -euo pipefail

-CONTAINER_NAME="${1:-${CONTAINER_NAME:-laoyao-stage2-1-plainqa-cpt}}"
+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 ;;
@@ -22,25 +44,37 @@ if [[ "$(docker inspect -f '{{.State.Running}}' "$CONTAINER_NAME" 2>/dev/null ||
  exit 3
 fi

-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"
-)"
-if [[ -z "$torchrun_pid" ]]; then
-  echo "ERROR: torchrun process not found in container: $CONTAINER_NAME" >&2
-  exit 4
-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

-mapfile -t worker_pids < <(
-  awk -v parent="$torchrun_pid" \
-    '$2 == parent && /laoyao_2b_moe_pretrain[.]py/ {print $1}' \
-    <<<"$container_top"
-)
-if [[ "${#worker_pids[@]}" -ne "$EXPECTED_WORKERS" ]]; 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
+  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"
@@ -49,6 +83,7 @@ if [[ -n "$tracker" && -f "$tracker" ]]; then
 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}"
@@ -83,3 +118,8 @@ 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
@@ -0,0 +1,40 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [[ "$#" -ne 4 ]]; then
+  echo "usage: $0 <container-name> <checkpoint-dir> <run-dir> <log-file>" >&2
+  exit 2
+fi
+
+CONTAINER_NAME="$1"
+CHECKPOINT_DIR="$2"
+RUN_DIR="$3"
+LOG_FILE="$4"
+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}"
+STARTUP_WAIT_SECONDS="${STARTUP_WAIT_SECONDS:-30}"
+
+deadline=$((SECONDS + STARTUP_WAIT_SECONDS))
+while [[ "$(docker inspect -f '{{.State.Running}}' "$CONTAINER_NAME" 2>/dev/null || true)" != "true" ]]; do
+  if (( SECONDS >= deadline )); then
+    echo "ERROR: container did not enter running state within ${STARTUP_WAIT_SECONDS}s: $CONTAINER_NAME" >&2
+    exit 3
+  fi
+  sleep 1
+done
+
+mkdir -p "$(dirname "$TRAINING_STATE_FILE")"
+temporary_file="$(mktemp "${TRAINING_STATE_FILE}.tmp.XXXXXX")"
+trap 'rm -f "$temporary_file"' EXIT
+{
+  printf 'container_name=%s\n' "$CONTAINER_NAME"
+  printf 'checkpoint_dir=%s\n' "$CHECKPOINT_DIR"
+  printf 'run_dir=%s\n' "$RUN_DIR"
+  printf 'log_file=%s\n' "$LOG_FILE"
+  printf 'recorded_at=%s\n' "$(date --iso-8601=seconds)"
+} >"$temporary_file"
+mv "$temporary_file" "$TRAINING_STATE_FILE"
+trap - EXIT
+
+echo "active_training_state=$TRAINING_STATE_FILE"
@@ -1,6 +1,7 @@
 #!/usr/bin/env bash
 set -euo pipefail

+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 REPO_ROOT="${REPO_ROOT:-/ssd/workspace/yi/laoyao_2b_moe}"
 IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}"
 RUN_NAME="${RUN_NAME:-pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000}"
@@ -18,7 +19,7 @@ EVAL_INTERVAL="${EVAL_INTERVAL:-15000}"
 EVAL_ITERS="${EVAL_ITERS:-10}"
 DATASET_WORKERS="${DATASET_WORKERS:-4}"
 LOG_FILE="${LOG_FILE:-/tmp/laoyao_pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000_resume.log}"
-CONTAINER_NAME="${CONTAINER_NAME:-laoyao_pretrain_resume}"
+CONTAINER_NAME="${CONTAINER_NAME:-laoyao-2b-pretraining}"

 if [[ ! -d "${REPO_ROOT}" ]]; then
  echo "ERROR: repo dir not found: ${REPO_ROOT}" >&2
@@ -91,6 +92,7 @@ nohup docker run --rm --name "${CONTAINER_NAME}" \
      --recompute-num-layers 1" \
  > "${LOG_FILE}" 2>&1 &

+bash "${SCRIPT_DIR}/record_active_training.sh" "$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE"
 echo "Launched. Check with:"
 echo "  tail -f ${LOG_FILE}"
 echo "  docker ps"
@@ -1,6 +1,7 @@
 #!/usr/bin/env bash
 set -euo pipefail

+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 # Host path layout on g0050 after /ssd was replaced by /data.
 HOST_WORKSPACE="${HOST_WORKSPACE:-/data/workspace/yi}"
 CONTAINER_WORKSPACE="${CONTAINER_WORKSPACE:-/ssd/workspace/yi}"
@@ -27,7 +28,7 @@ LR="${LR:-3.0e-4}"
 MIN_LR="${MIN_LR:-5.0e-6}"
 WARMUP_FRACTION="${WARMUP_FRACTION:-0.01}"
 LOG_FILE="${LOG_FILE:-/tmp/laoyao_stage1_2_8192_8gpu_mbs14.log}"
-CONTAINER_NAME="${CONTAINER_NAME:-laoyao_stage1_2_pretrain}"
+CONTAINER_NAME="${CONTAINER_NAME:-laoyao-2b-pretraining}"

 HOST_DATA_MANIFEST="${DATA_MANIFEST/${CONTAINER_WORKSPACE}/${HOST_WORKSPACE}}"
 HOST_LOAD_DIR="${LOAD_DIR/${CONTAINER_WORKSPACE}/${HOST_WORKSPACE}}"
@@ -115,6 +116,8 @@ nohup docker run --rm --name "${CONTAINER_NAME}" \
      --recompute-num-layers 1" \
  > "${LOG_FILE}" 2>&1 &

+bash "${SCRIPT_DIR}/record_active_training.sh" \
+  "$CONTAINER_NAME" "$HOST_CKPT_DIR" "$(dirname "$HOST_CKPT_DIR")" "$LOG_FILE"
 echo "Launched. Check with:"
 echo "  tail -f ${LOG_FILE}"
 echo "  docker ps"
@@ -1,6 +1,7 @@
 #!/usr/bin/env bash
 set -euo pipefail

+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao-2b-pretraining/code/laoyao_2b_moe}"
 IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}"
 DATA_MANIFEST="${DATA_MANIFEST:-/data/yi/laoyao-2b-pretraining/megatron/stage2_1_8192_glm52_plainqa_cpt_v4_20260722/manifest.json}"
@@ -13,7 +14,7 @@ RUN_DIR="${RUN_DIR:-/mnt/beegfs/yi/laoyao-2b-pretraining/runs/${RUN_NAME}}"
 CKPT_DIR="${CKPT_DIR:-${RUN_DIR}/checkpoints}"
 TENSORBOARD_DIR="${TENSORBOARD_DIR:-${RUN_DIR}/tensorboard}"
 LOG_FILE="${LOG_FILE:-${RUN_DIR}/train.log}"
-CONTAINER_NAME="${CONTAINER_NAME:-laoyao-stage2-1-plainqa-cpt}"
+CONTAINER_NAME="${CONTAINER_NAME:-laoyao-2b-pretraining}"

 SEQ_LENGTH="${SEQ_LENGTH:-8192}"
 MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-14}"
@@ -153,4 +154,5 @@ nohup docker run --rm --name "$CONTAINER_NAME" \
    --recompute-granularity full --recompute-method uniform --recompute-num-layers 1 \
    $STAGE_TRANSITION_ARG" >"$LOG_FILE" 2>&1 &

+bash "${SCRIPT_DIR}/record_active_training.sh" "$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE"
 echo "Launched container=$CONTAINER_NAME log=$LOG_FILE"
@@ -1,6 +1,7 @@
 #!/usr/bin/env bash
 set -euo pipefail

+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao-2b-pretraining/code/laoyao_2b_moe}"
 IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}"
 DATA_MANIFEST="${DATA_MANIFEST:-/data/yi/laoyao-2b-pretraining/megatron/stage2_1_8192_glm52_zh10_v3_20260719/manifest.json}"
@@ -13,7 +14,7 @@ RUN_DIR="${RUN_DIR:-/mnt/beegfs/yi/laoyao-2b-pretraining/runs/${RUN_NAME}}"
 CKPT_DIR="${CKPT_DIR:-${RUN_DIR}/checkpoints}"
 TENSORBOARD_DIR="${TENSORBOARD_DIR:-${RUN_DIR}/tensorboard}"
 LOG_FILE="${LOG_FILE:-${RUN_DIR}/train.log}"
-CONTAINER_NAME="${CONTAINER_NAME:-laoyao-stage2-1-pretrain}"
+CONTAINER_NAME="${CONTAINER_NAME:-laoyao-2b-pretraining}"

 SEQ_LENGTH="${SEQ_LENGTH:-8192}"
 MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-14}"
@@ -140,4 +141,5 @@ nohup docker run --rm --name "$CONTAINER_NAME" \
    --recompute-granularity full --recompute-method uniform --recompute-num-layers 1 \
    $STAGE_TRANSITION_ARG" >"$LOG_FILE" 2>&1 &

+bash "${SCRIPT_DIR}/record_active_training.sh" "$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE"
 echo "Launched container=$CONTAINER_NAME log=$LOG_FILE"
@@ -1,6 +1,7 @@
 #!/usr/bin/env bash
 set -euo pipefail

+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao-2b-pretraining/code/laoyao_2b_moe}"
 IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}"
 DATA_MANIFEST="${DATA_MANIFEST:-/data/yi/laoyao-2b-pretraining/megatron/stage2_2_8192_glm52_plainqa_cpt_v4_20260722/manifest.json}"
@@ -13,7 +14,7 @@ RUN_DIR="${RUN_DIR:-/mnt/beegfs/yi/laoyao-2b-pretraining/runs/${RUN_NAME}}"
 CKPT_DIR="${CKPT_DIR:-${RUN_DIR}/checkpoints}"
 TENSORBOARD_DIR="${TENSORBOARD_DIR:-${RUN_DIR}/tensorboard}"
 LOG_FILE="${LOG_FILE:-${RUN_DIR}/train.log}"
-CONTAINER_NAME="${CONTAINER_NAME:-laoyao-stage2-2-plainqa-cpt}"
+CONTAINER_NAME="${CONTAINER_NAME:-laoyao-2b-pretraining}"

 SEQ_LENGTH="${SEQ_LENGTH:-8192}"
 MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-14}"
@@ -146,4 +147,5 @@ nohup docker run --rm --name "$CONTAINER_NAME" \
    --recompute-granularity full --recompute-method uniform --recompute-num-layers 1 \
    $STAGE_TRANSITION_ARG" >"$LOG_FILE" 2>&1 &

+bash "${SCRIPT_DIR}/record_active_training.sh" "$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE"
 echo "Launched container=$CONTAINER_NAME log=$LOG_FILE"
@@ -1,6 +1,7 @@
 #!/usr/bin/env bash
 set -euo pipefail

+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 REPO_ROOT="${REPO_ROOT:-/mnt/beegfs/yi/laoyao-2b-pretraining/code/laoyao_2b_moe}"
 IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}"
 DATA_MANIFEST="${DATA_MANIFEST:-/data/yi/laoyao-2b-pretraining/megatron/stage2_2_8192_glm52_zh10_v3_20260719/manifest.json}"
@@ -13,7 +14,7 @@ RUN_DIR="${RUN_DIR:-/mnt/beegfs/yi/laoyao-2b-pretraining/runs/${RUN_NAME}}"
 CKPT_DIR="${CKPT_DIR:-${RUN_DIR}/checkpoints}"
 TENSORBOARD_DIR="${TENSORBOARD_DIR:-${RUN_DIR}/tensorboard}"
 LOG_FILE="${LOG_FILE:-${RUN_DIR}/train.log}"
-CONTAINER_NAME="${CONTAINER_NAME:-laoyao-stage2-2-pretrain}"
+CONTAINER_NAME="${CONTAINER_NAME:-laoyao-2b-pretraining}"

 SEQ_LENGTH="${SEQ_LENGTH:-8192}"
 MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-14}"
@@ -146,4 +147,5 @@ nohup docker run --rm --name "$CONTAINER_NAME" \
    --recompute-granularity full --recompute-method uniform --recompute-num-layers 1 \
    $STAGE_TRANSITION_ARG" >"$LOG_FILE" 2>&1 &

+bash "${SCRIPT_DIR}/record_active_training.sh" "$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE"
 echo "Launched container=$CONTAINER_NAME log=$LOG_FILE"
@@ -74,5 +74,6 @@ cfg.train.exit_signal = signal.SIGINT
 只记录退出请求；标准训练循环会在 iteration 边界执行 checkpoint，因此不会从
 Python signal callback 中直接运行 CUDA、NCCL 或文件写入。

-操作时使用 `scripts/graceful_stop_megatron_training.sh`，不要把 SIGINT 直接发给
-torchrun/container PID 1。
+训练 launcher 会将活动训练信息写入 `.runtime/active_training.txt`。操作时直接运行
+`scripts/graceful_stop_megatron_training.sh`，无需查询 container name；不要把
+SIGINT 直接发给 torchrun/container PID 1。