自動再投入シェルスクリプト

はじめに

卒業研究中に作成した「自動再投入シェルスクリプト」の備忘録です．

研究内容

現在の研究テーマは、Detection Transformer (DETR) 系検出器を用いた小型物体検出の精度向上です．

• 実行環境は大学のGPGPUサーバ（SLURMを利用）
• このサーバではジョブの最大実行時間が24時間に制限
• 長時間学習が必要なDETRを実行する際，途中から再開できる仕組みが必要

スクリプトの詳細

• detr_checkpoint.sh: SLURM を使って DETR の学習を1回実行 するスクリプト
• auto_continue_detr.sh: detr_checkpoint.sh を 自動で繰り返し実行・監視 し、失敗時もチェックポイントから再開するスクリプト

detr_checkpoint.sh

#!/bin/bash
#SBATCH --job-name=detr-checkpoint
#SBATCH --partition="パーティション名"
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=16
#SBATCH --gpus-per-task=4
#SBATCH --mem=512G
#SBATCH --time=24:00:00
#SBATCH --output=slurm-%j.out
#SBATCH --container="コンテナ名"
#SBATCH --mail-type=END,FAIL
#SBATCH --mail-user="メールアドレス"

cd /detr/detr

# 出力ディレクトリの設定
OUTPUT_DIR="/root/detr/output/CheckPoints$(date +%Y%m%d_%H%M%S)"
mkdir -p $OUTPUT_DIR

echo "=== DETR Training Started ==="
echo "Job ID: $SLURM_JOB_ID"
echo "Output Directory: $OUTPUT_DIR"
echo "Start Time: $(date)"

# 最新のチェックポイントを確認
RESUME_OPTION=""
# 過去の出力ディレクトリから最新のチェックポイントを検索
LATEST_CHECKPOINT=$(find /root/detr/output/CheckPoints* -name "checkpoint.pth" 2>/dev/null | sort -r | head -1)

if [ -n "$LATEST_CHECKPOINT" ]; then
    RESUME_OPTION="--resume $LATEST_CHECKPOINT"
    echo "最新のチェックポイントから再開: $LATEST_CHECKPOINT"
else
    echo "新規トレーニングを開始"
fi

echo "=== Training Parameters ==="
echo "Epochs: 500"
echo "Batch Size: 8"
echo "Learning Rate: 1e-4"
echo "Resume Option: $RESUME_OPTION"
echo "=============================="

# バックグラウンドでDETRトレーニングを実行
python -m torch.distributed.launch --nproc_per_node=4 --use_env main.py \
  --coco_path /root/data/coco \
  --output_dir $OUTPUT_DIR \
  --epochs 500 \
  --batch_size 8 \
  --lr 1e-4 \
  --lr_backbone 1e-5 \
  --weight_decay 1e-4 \
  --num_workers 2 \
  $RESUME_OPTION &

TRAIN_PID=$!

# SIGTERM/SIGINTを受信した時の処理
cleanup() {
    echo "=== 終了信号を受信 ==="
    echo "トレーニングプロセス (PID: $TRAIN_PID) に終了信号を送信中..."
    kill -TERM $TRAIN_PID 2>/dev/null
    
    # プロセスが終了するまで待機（最大60秒）
    local count=0
    while kill -0 $TRAIN_PID 2>/dev/null && [ $count -lt 60 ]; do
        sleep 1
        count=$((count + 1))
    done
    
    # まだ動いている場合は強制終了
    if kill -0 $TRAIN_PID 2>/dev/null; then
        echo "強制終了中..."
        kill -KILL $TRAIN_PID 2>/dev/null
    fi
    
    echo "=== チェックポイント確認 ==="
    if [ -f "$OUTPUT_DIR/checkpoint.pth" ]; then
        echo "チェックポイントが保存されています: $OUTPUT_DIR/checkpoint.pth"
        ls -la $OUTPUT_DIR/checkpoint.pth
    else
        echo "警告: チェックポイントファイルが見つかりません"
    fi

    echo "=== 終了処理完了 ==="
    exit 0
}

# シグナルハンドラーを設定
trap cleanup SIGTERM SIGINT

# トレーニングプロセスの完了を待機
wait $TRAIN_PID
EXIT_CODE=$?

echo "=== Training Completed ==="
echo "Exit Code: $EXIT_CODE"
echo "End Time: $(date)"

# 実行結果の確認
if [ $EXIT_CODE -eq 0 ]; then
    echo "トレーニングが正常に完了しました"
else
    echo "トレーニングでエラーが発生しました (Exit Code: $EXIT_CODE)"
fi

# 最終的なチェックポイントの確認
echo "=== Final Checkpoint Status ==="
if [ -f "$OUTPUT_DIR/checkpoint.pth" ]; then
    echo "最終チェックポイント: $OUTPUT_DIR/checkpoint.pth"
    ls -la $OUTPUT_DIR/checkpoint.pth
else
    echo "チェックポイントファイルが見つかりません"
fi

# ディスク使用量の確認
echo "=== Disk Usage ==="
du -sh $OUTPUT_DIR
echo "=================="

auto_continue_detr.sh

#!/bin/bash
# auto_continue_detr.sh

SCRIPT_NAME="detr_checkpoint.sh"
MAX_JOBS=10
CURRENT_JOB=1

echo "=== DETR継続実行スクリプト ==="
echo "最大ジョブ数: $MAX_JOBS"
echo "開始時刻: $(date)"

while [ $CURRENT_JOB -le $MAX_JOBS ]; do
    echo ""
    echo "=== ジョブ $CURRENT_JOB/$MAX_JOBS を投入 ==="

    # ジョブを投入し、ジョブIDのみを抽出
    SBATCH_OUTPUT=$(sbatch --parsable $SCRIPT_NAME 2>&1)
    
    if [ $? -eq 0 ]; then
        # 出力から数字のみを抽出（ジョブID）
        JOB_ID=$(echo "$SBATCH_OUTPUT" | grep -o '[0-9][0-9]*' | tail -1)
        
        if [ -n "$JOB_ID" ]; then
            echo "ジョブ ID: $JOB_ID が投入されました"
            echo "sbatch出力: $SBATCH_OUTPUT"
        else
            echo "エラー: ジョブIDの抽出に失敗しました"
            echo "sbatch出力: $SBATCH_OUTPUT"
            exit 1
        fi
    else
        echo "エラー: ジョブの投入に失敗しました"
        echo "sbatch出力: $SBATCH_OUTPUT"
        exit 1
    fi

    # ジョブの状態を監視
    echo "ジョブ $JOB_ID の実行を監視中..."

    # 初期状態確認
    while true; do
        JOB_EXISTS=$(squeue -j $JOB_ID 2>/dev/null | grep "$JOB_ID" | wc -l)
        if [ $JOB_EXISTS -eq 0 ]; then
            echo "ジョブ $JOB_ID がキューから消えました"
            break
        fi
        
        # 5分間隔でステータスを表示
        JOB_STATUS=$(squeue -j $JOB_ID --format="%T" --noheader 2>/dev/null | tr -d ' ')
        if [ -n "$JOB_STATUS" ]; then
            echo "$(date): ジョブ $JOB_ID - ステータス: $JOB_STATUS"
        fi
        sleep 300
    done

    # ジョブの最終状態を確認
    echo "ジョブ $JOB_ID が終了しました。最終状態を確認中..."
    sleep 10  # saccctの情報が更新されるまで少し待機

    # sacctでジョブの状態を確認（数値のジョブIDを使用）
    JOB_STATE=$(sacct -j $JOB_ID --format=State --noheader 2>/dev/null | tail -1 | tr -d ' ')
    echo "最終状態: $JOB_STATE"

    # チェックポイント確認関数
    find_latest_checkpoint() {
        find /root/detr/output/CheckPoints* -type f -name "checkpoint.pth" 2>/dev/null | sort -r | head -1
    }

    case $JOB_STATE in
        "COMPLETED")
            echo "🎉 トレーニングが正常に完了しました！"

            LATEST_OUTPUT=$(ls -td /root/detr/output/CheckPoints* 2>/dev/null | head -1)
            if [ -n "$LATEST_OUTPUT" ]; then
                echo "最終出力ディレクトリ: $LATEST_OUTPUT"
                echo "最終結果:"
                ls -la "$LATEST_OUTPUT/"
            fi
            exit 0
            ;;

        "TIMEOUT" | "CANCELLED" | "FAILED" | "NODE_FAIL" | "OUT_OF_MEMORY")
            echo "⚠️ ジョブが $JOB_STATE しました。チェックポイントから継続実行を試みます..."

            LATEST_CHECKPOINT=$(find_latest_checkpoint)
            if [ -n "$LATEST_CHECKPOINT" ]; then
                echo "チェックポイントが見つかりました: $LATEST_CHECKPOINT"
            else
                echo "チェックポイントが見つかりません。実行を中止します。"
                exit 1
            fi
            ;;

        "")
            echo "❓ ジョブ状態が取得できませんでした。手動で確認してください。"
            echo "ジョブ ID: $JOB_ID"
            
            # チェックポイントの存在確認
            LATEST_CHECKPOINT=$(find_latest_checkpoint)
            if [ -n "$LATEST_CHECKPOINT" ]; then
                echo "チェックポイントが見つかりました: $LATEST_CHECKPOINT"
                echo "継続実行を試みます..."
            else
                echo "チェックポイントが見つかりません。実行を中止します。"
                exit 1
            fi
            ;;

        *)
            echo "❓ 不明な状態: $JOB_STATE"
            echo "チェックポイントの存在を確認中..."

            LATEST_CHECKPOINT=$(find_latest_checkpoint)
            if [ -n "$LATEST_CHECKPOINT" ]; then
                echo "チェックポイントが見つかりました: $LATEST_CHECKPOINT"
            else
                echo "チェックポイントが見つかりません。実行を中止します。"
                exit 1
            fi
            ;;
    esac

    # 次のジョブまで少し待機
    echo "30秒後に次のジョブを投入します..."
    sleep 30

    CURRENT_JOB=$((CURRENT_JOB + 1))
done

echo "⚠️  最大ジョブ数に達しました。トレーニングが完了していない可能性があります。"
echo "最新のチェックポイントを確認してください:"
find /root/detr/output/CheckPoints* -name "checkpoint.pth" 2>/dev/null || echo "チェックポイントが見つかりません"

detr_checkpoint.shとauto_continue_detr.shを作成し，以下の通り実行権限を付与する．

bash

chmod +x detr_checkpoint.sh
chmod +x auto_continue_detr.sh

実行権限を付けたら，以下のコマンドでバックグラウンド実行します．

bash

nohup ./auto_continue_detr.sh > auto_continue.log 2>&1 &

実行状況のログは auto_continue.log に保存されるので，進捗やエラーを確認できます．