优化训练任务启动命令

在 AIStudio 中，训练任务运行在容器内，其退出码（exit code）对平台调度器至关重要。平台依据退出码判断训练任务是成功、失败还是异常退出。

Exit Code	平台状态
0	运行成功。平台认为模型训练成功。
非 0	运行失败。平台认为模型训练失败或异常终止。

因此，用户编写的启动脚本必须确保：

• 错误不会被隐藏（如被 tee 吞掉）
• 脚本失败时必须返回非 0
• 日志足够让用户排障
• 整个流程可观测、可追踪

下面我们通过重构一个训练任务启动命令示例，展示如何将其改造成平台级可靠脚本。

原始启动命令

该训练任务的启动命令存在以下问题：

1. tee 导致异常被吃掉，没有 pipefail。在 cmd | tee log 中，因为 tee 总是成功，管道 exit code 会变成 0。
2. 无错误捕获。失败后脚本仍继续，容器 exit code 不是训练真正 exit code。缺少错误处理机制，无法在训练任务失败时采取措施。
3. 无阶段日志，无法跟踪每个阶段的开始和结束时间以及执行状态，排障时很难知道卡在哪一步。
4. GPU 时间被浪费，将大文件下载、解压和数据准备放在训练任务启动时执行，导致 GPU 节点空转。

wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json​
wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt​
mkdir -p checkpoints/gpt2_345m​
cd checkpoints/gpt2_345m​
wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/megatron_l​
m_345m/versions/v0.0/zip -O megatron_lm_345m_v0.0.zip​
unzip megatron_lm_345m_v0.0.zip​
rm megatron_lm_345m_v0.0.zip​
cd ../..​
git clone https://github.com/EastInsure/Megatron-DeepSpeed.git​
cd Megatron-DeepSpeed​
bash train/7B1_test.sh $MASTER_ADDR $MASTER_PORT $WORLD_SIZE $RANK | tee $(pwd)/M​egatron-DeepSpeed.log

训练脚本/命令必须确保退出码不被 tee 吃掉

在 Shell 中，默认情况下管道命令（如 cmd1 | cmd2）的退出码取决于管道中最后一个命令。如果用户使用 tee 保存日志（python train.py | tee log.txt），即使训练脚本失败，只要 tee 成功，整个命令就会返回 0，导致平台在任务失败时仍展示运行成功。

解决方案：在启动命令头部（包括引用的脚本头部），设置 set -o pipefail，保证管道中的任何错误都自动传播，导致整个管道失败。同时需确保管道错误可被捕获，任何失败立即导致任务显性失败。

错误捕获

Bash 默认是非常宽容的。如果脚本中某一行命令失败了（返回非 0 状态码），Bash 会忽略这个错误（静默失败，Silent Failure），继续执行下一行。在 AI 训练流程中，这意味着任务实际上已经出错了，但容器依然在空转，浪费宝贵的 GPU 算力。

场景演示：假设您的脚本逻辑是：进入数据目录，解压数据，然后开始训练。

• 不加 set -e。cd 命令失败不被捕获，结果删除了不该删的文件，且浪费了时间排查为什么训练找不到数据。

  cd /data/datasets/coco  # 假设这个目录不存在，命令失败
  # Bash 忽略了上面的错误，继续执行
  rm -rf * # 灾难发生！它删除了当前工作目录（可能是根目录）下的所有文件
  python train.py         # 找不到数据，报错退出

• 加上 set -e (安全)。如果 cd 目录不存在，任务立即标记为失败，您能立刻看到平台任务报错。

  set -e
  cd /data/datasets/coco  # 目录不存在，命令失败，返回非 0
  # 脚本立即在此处终止！后续命令不会执行。

解决方案：在实际的启动命令中，我们推荐将 set -eu 与我们之前提到的 pipefail 组合使用（严格模式），在启动命令及所有 shell 脚本的开头（第一行 #!/bin/bash 之后）添加以下指令，确保任何错误立即暴露

set -euo pipefail

作用：

• -e：只要有一个命令失败就退出
• -u：未定义变量视为错误
• -o pipefail：管道错误不被隐藏

注意

• -e/-u/-o pipefail 可分开使用。
• set -euo pipefail 设置的是当前 Shell 进程的选项。如果在启动命令里执行 bash train.sh 时，会启动了一个全新的子进程（Child Process）。默认情况下，子进程不会继承父 Shell 的这些选项设置。请在引用的脚本头部添加 set -euo pipefail，启用严格模式。
• 即使使用了 set -e，在某些特定场景（如需要判断某个命令是否失败并执行特定逻辑）下，可以使用 || 操作符或 if 语句来暂时豁免错误检查。

# 允许 mkdir 失败（例如目录已存在），脚本继续执行
mkdir /existing/dir || true

推荐使用 Process Substitution

在任务启动命令中，常见的日志写法是：

bash train.sh | tee train.log

虽然这种写法简单，但并不适用于深度学习训练任务。在分布式训练、多 GPU 训练以及高负载场景下，将训练脚本放入 Subshell（如使用管道）会导致信号传递中断，造成日志延迟、信号无法传递（导致僵尸进程）等问题。

在 Shell 层面，为了同时将日志输出到控制台和文件，Bash 提供了一种特殊的语法 > >(...)，它允许我们将输出重定向到一个进程，但不将主命令放入 Subshell。

## 在外部使用 > >(tee ...) 处理 Shell 和 Supervisor 日志
bash train-with-torchrun.sh > >(tee execution_monitor.log) 2>&1

记录关键阶段标记

平台任务日志功能现已自动为每一行任务日志注入时间戳与 Worker 名称，您无需再手动在脚本中打印 date 或 $HOSTNAME。

尽管如此，我们仍强烈建议在脚本的关键节点（如环境准备完毕、开始训练、保存模型等）打印清晰的阶段标记。这能帮助您在成千上万行复杂的框架底层日志中，快速定位任务目前处于哪个逻辑阶段。

推荐做法：

只需打印业务含义，无需关注元数据。平台会自动将其格式化为：[2025-11-28 12:00:00] [worker-0] === [Phase] Launching torchrun ===。

# 仅需关注任务逻辑进度
echo "=== [Phase] Environment Setup Complete ==="
# ... (中间是大量安装日志)
echo "=== [Phase] Launching torchrun ==="
# ... (中间是大量训练日志)
echo "=== [Phase] Training Finished, Starting Upload ==="

错误现场保留

在训练任务失败时，让容器立即退出虽然能让平台快速标记错误，但往往不利于排障。用户通常希望进入容器检查日志、临时文件、GPU 状态等。因此，在启动命令中对训练脚本的退出码进行判断，并在失败时保留现场一段时间，是非常成熟的生产实践。

如果退出码非零，打印清晰错误提示并 sleep 指定时间，保持容器存活供用户调试；保留结束后再用原始退出码退出，确保平台能正确识别失败状态：

#!/bin/bash
set -euo pipefail

# 1. 显式初始化错误码变量
RET=0

# 2. 执行训练，分离日志
# 使用 "||" 捕获真实退出码，防止脚本立即退出
echo "[$(date +'%F %T')] Starting training..."
bash train/7B1_test.sh "$@" \
    > >(tee -i train.log) 2>&1 || RET=$?

# 3. (关键) 等待日志落盘
# 防止 tee 进程的数据滞留在缓冲区
sleep 2

# 4. 结果处理与现场保留逻辑
if [[ $RET -ne 0 ]]; then
    # 打印带时间戳和主机名的错误信息
    echo "[$(date +'%F %T')] [$HOSTNAME] Training crashed with exit code $RET."
    echo "Keeping container alive for 10000s to allow debugging..."
    echo "You can now open a terminal to this container."
    
    # 5. 挂起容器，保留现场
    # 注意：这期间 GPU 仍会被占用
    sleep 10000
else
    echo "[$(date +'%F %T')] Training finished successfully."
fi

# 6. 最终以实际状态码退出
exit "$RET"

这种模式同时兼顾可观测性（错误明确）、可调试性（容器不立即消失） 与 准确性（失败状态正确上报），强烈推荐在所有模型训练与数据处理任务中采用。

优化后启动命令示例

以下是一个经过优化的平台训练任务启动命令示例，做到了让容器退出码可信、训练任务可观测、失败可排障、脚本行为确定且平台调度器可以正确感知异常。

#!/bin/bash
# 1. 开启严格模式：任何命令失败立即退出，管道错误不隐藏，未定义变量报错
set -euo pipefail

# --- 环境变量与路径配置 (根据实际挂载路径修改) ---
# 假设数据和预训练模型已挂载在 /mnt 目录下
PRETRAINED_PATH="/mnt/data/megatron_lm_345m" 
CHECKPOINT_PATH="/mnt/data/checkpoints/gpt2_345m"
CODE_DIR="Megatron-DeepSpeed"

echo "[$(date +'%F %T')] Environment setup started."

# --- 代码准备阶段 ---
# 仅在目录不存在时克隆，避免重复执行报错(注意 git clone 网络限制问题)
if [ ! -d "$CODE_DIR" ]; then
    echo "Cloning repository..."
    git clone https://github.com/EastInsure/Megatron-DeepSpeed.git
else
    echo "Repository already exists. Skipping clone."
fi

cd "$CODE_DIR"

# 确保 Checkpoint 输出目录存在
mkdir -p "$CHECKPOINT_PATH"

# --- 训练执行阶段 (核心优化) ---
echo "[$(date +'%F %T')] Starting training task..."

# 初始化退出码
RET=0

# 执行训练脚本：
# 1. 使用 > >(tee ...) 替代管道 | tee，防止信号丢失和僵尸进程
# 2. 使用 || RET=$? 捕获真实退出码，防止脚本因 set -e 立即退出
# 3. 2>&1 确保标准错误也能被记录到日志
bash train/7B1_test.sh $MASTER_ADDR $MASTER_PORT $WORLD_SIZE $RANK \
    > >(tee -i $(pwd)/Megatron-DeepSpeed.log) 2>&1 || RET=$?

# --- 结果处理与调试阶段 ---

# 等待日志落盘 (防止容器销毁导致日志截断)
sleep 2

if [[ $RET -ne 0 ]]; then
    echo "========================================================"
    echo "[$(date +'%F %T')] [$HOSTNAME] Training CRASHED with exit code: $RET"
    echo "Check execution_monitor.log for details."
    echo "--------------------------------------------------------"
    echo "⚠️  DEBUG MODE ENABLED: Keeping container alive for 10000s."
    echo "   You can now use 'Web Terminal' to debug."
    echo "   To stop manually: Run 'kill 1' or stop job from dashboard."
    echo "========================================================"
    
    # 挂起容器，保留现场供用户排查 (注意：此时 GPU 仍在计费)
    sleep 10000
else
    echo "[$(date +'%F %T')] Training finished successfully."
fi

# 最终以实际状态码退出，通知平台调度器任务结果
exit "$RET"