My Portfolio

Epilogでクリーンアップする方法まとめ３

登録日	：2025/03/30 12:21
カテゴリ	：SLURM

SLURM 21.8.8環境でBatchFlagを使用してジョブタイプを判別し、適切なクリーンアップ処理を行うEpilogスクリプトは以下のように実装できます。

バッチ処理のロジックについて、SLURMジョブIDに紐ずくプロセスだけを削除するような仕組みを追加しました。ジョブIDに紐ずくPIDを親プロセスから子プロセスまでを収集する外部スクリプトを作成して組み込んでいます。

1. Epilogスクリプト

Epilogで実行するスクリプトを以下のように作成

#!/bin/bash

#--------------------------------------
# SLURM remained process killing
# Created by Me
#Ver1.0 2025.03.15
# Ver1.1 2025.03.20
# Ver2.0 2025.03.30
#-------------------------------------

#
# settings
#

# Dry-run mode setting (0=run mode, 1=log only)
#DRY_RUN=${DRY_RUN:-0}
DRY_RUN=1

# Log settings
BASE_DIR="/home/settings/Epilog/log"
LOG_DIR=$BASE_DIR/${SLURM_JOB_ID}
mkdir -p $LOG_DIR
LOG_FILE="$LOG_DIR/cleanup_${SLURM_JOB_ID}_$(hostname)_$(date +%Y%m%d%H%M%S).log"

# external scripts
SCRIPT="/home/settings/Epilog/script/fetch_pids"

#
# declare function
#

# BatchFlag get function
function get_batch_flag() {
    local retries=3
    for ((i=1; i<=retries; i++)); do
        JOB_INFO=$(sudo scontrol show job $SLURM_JOB_ID 2>/dev/null) && {
            BATCH_FLAG=$(echo "$JOB_INFO" | grep -oP 'BatchFlag=\K\d+')
            [ -n "$BATCH_FLAG" ] && return 0
        }
        sleep 1
    done
    return 1
}

# Process deletion/logging functions
function delete_processes() {
    local pids=$1
    local job_type=$2

    if [ $DRY_RUN -eq 1 ]; then
        echo "[DRY-RUN] ${job_type} processes to be killed:"
        echo "$pids"
    else
        echo "[ACTION] Killing ${job_type} processes:"
        echo "$pids"
        sudo kill -9 $pids 2>/dev/null || true
    fi
}

# Batch job processing
function handle_batch_job() {
    local pid=$(ps -u $SLURM_JOB_USER -o pid,cmd | grep -E "[s]lurm" | grep job | grep $SLURM_JOB_ID | awk '{print $1}')
    echo "SLURM PID: $pid"
    echo "SLURM_TASK_PID: $SLURM_TASK_PID"
    echo $SCRIPT" : start"
    #product
    local pids=$($SCRIPT $pid $SLURM_JOB_USER)
    #test------------------------
    #remark root
    #local pids=$($SCRIPT 1 root)
    #----------------------------
    echo $SCRIPT" : end"
    delete_processes "$pids" "Batch"
}

# Interactive job processing
function handle_interactive_job() {
    local pids=$(pgrep -u $SLURM_JOB_USER | grep -v "^$$\$")
    delete_processes "$pids" "Interactive"
}

# Fail-safe processing
function handle_failsafe() {
    local pids=$(ps -u $SLURM_JOB_USER -o pid,cmd | grep -E "[s]lurm" | grep job | grep $SLURM_JOB_ID | awk '{print $1}')
    delete_processes "$pids" "Failsafe"
}

#
# main
#
{
sleep 2
EXIT_CODE=$(sacct -j $SLURM_JOB_ID --format=State --noheader | tr -d ' ' | head -n 1)
echo "=== Epilog Start [$(date)] ==="
echo "Job ID: $SLURM_JOB_ID"
echo "User: $SLURM_JOB_USER"
echo "STATUS: $EXIT_CODE"
echo "Node: $(hostname)"
echo "Dry-run Mode: $DRY_RUN"

# Get BatchFlag
if get_batch_flag; then
    case $BATCH_FLAG in
        "1")
            echo "[Batch Job] Process cleanup started"
            handle_batch_job
            ;;
        "0")
            echo "[Interactive Job] Full cleanup started"
            handle_interactive_job
            ;;
        *)
            echo "[WARNING] Invalid BatchFlag value: $BATCH_FLAG"
            handle_failsafe
            ;;
    esac
else
    echo "[WARNING] Failed to detect job type. Performing safe cleanup"
    handle_failsafe
fi

echo "=== Epilog End [$(date)] ==="
} >> "$LOG_FILE" 2>&1

# env | grep SLURM >> $LOG_FILE

exit 0

バッチジョブ時 (BatchFlag=1)
１　psコマンドでユーザーの全プロセスをリストアップ
２　grepで以下のパターンをフィルタリング:
　　・slurm関連プロセス
　　・コマンドラインにジョブIDを含むプロセス
３　該当PIDを強制終了

インタラクティブジョブ時 (BatchFlag=0)
１　pgrepでユーザーの全プロセスを取得
２　現在のシェルプロセス(PID: $$)を除外
３　残りすべてのPIDを強制終了

リトライロジック:
　・scontrolコマンドが一時的に失敗した場合に3回まで再試行
　・ネットワーク遅延や高負荷状態に対応

セキュリティ:
　・ログファイルのパーミッションを適切に設定（例: chmod 640 /var/log/slurm/epilog/*）

※このスクリプトはSLURM 21.8.8環境で動作検証済みです。本番環境導入前には必ずテスト環境での動作確認をお願いします。
(追記)テスト環境を構築して、Epilogの設定は問題なく構築できることを確認した。
※プロセスの削除が適切に行えたかまでは、テスト環境の規模が小さく確認できていない。
Epilogスクリプトを以下のように修正
・共有フォルダ（NFSマウント）にEpilogスクリプトを配置して参照する
・ログの出力先も共有フォルダ（NFSマウント）として、ジョブIDごとにフォルダを作成する。
・テスト時の切り分けのため、kill実施有無をコメントアウトで切替
スクリプトの注意点として、実行権限をつけるのを忘れないこと。
※実行権限がないと、Epilogスクリプト実行エラーでノードがDrainとなる。

External script

バッチ処理でSLURMジョブIDの親プロセスから子プロセスを収集する外部スクリプト「fetch_pids」を以下のように作成

#!/bin/bash

#
# Recursively collect PIDs associated with the input PID
# Created by Me
# ver1.0 2025.03.30
#

#
# settings input
#

TARGET_PID=$1  # The PID specified in the argument
TARGET_USER=${2:-$SLURM_JOB_USER}  # If the second argument is not specified, $SLURM_JOB_USER is used.

# argument check
if [[ -z "$TARGET_USER" || -z "$TARGET_PID" ]]; then
    echo "Usage: $0 <PID> [USER]"
    # echo "Note: If USER is not specified, \$SLURM_JOB_USER is used."
    exit 1
fi

#
# declair function
#

# A function to check if a PID exists
function pid_exists {
    local pid=$1
    ps -ef | awk -v user="$TARGET_USER" '$1 == user {print $2}' | grep -wq "$pid"
    # echo $?
    return $?  # Returns 0 if it exists, 1 if it doesn't.
}

# A function to recursively get child processes associated with a PID
function get_descendants {
    local current_pid=$1

    # If the PID does not exist, the process exits without doing anything.
    if ! pid_exists "$current_pid"; then
        return
    fi

    echo "$current_pid"  # Print the current PID

    # Get the child process of the current PID
    local children=$(ps -ef | awk -v user="$TARGET_USER" '$1 == user {print $2, $3}' | awk -v ppid="$current_pid" '$2 == ppid {print $1}')
    for child in $children; do
        # Recursively search for child processes
        get_descendants "$child"
    done
}

#
# main
#

# Start with the initial PID and print the results
get_descendants "$TARGET_PID"

2. 設定手順

2.1 スクリプトを配置と権限設定

※各クライアントの計算ノードすべてで実施する

sudo mkdir -p /opt/slurm/etc/scripts/epilog.d
sudo nano /opt/slurm/etc/scripts/epilog.d/20_clean_processes.sh

※SLURMの実行ユーザがslurmの場合。rootで実行している場合はchwon不要

sudo chmod 755 /opt/slurm/etc/scripts/epilog.d/20_clean_processes.sh
sudo chown slurm:slurm /opt/slurm/etc/scripts/epilog.d/20_clean_processes.sh

※SLURMの実行ユーザがslurmの場合。rootで実行している場合はsudousersの設定不要

# /etc/sudoers.d/slurm_epilog
slurm ALL=(ALL) NOPASSWD: /usr/bin/scontrol
slurm ALL=(ALL) NOPASSWD: /usr/bin/pgrep
slurm ALL=(ALL) NOPASSWD: /usr/bin/kill

2.2 ログローテーションの設定

クリーンアップの実行ログを出力するので、ローテーションしてログを定期的にクリアする
※各クライアントの計算ノードすべてで実施する

# /etc/logrotate.d/slurm_epilog
/var/log/slurm/epilog/*.log {
    daily
    missingok
    rotate 30
    compress
    delaycompress
    notifempty
}

2.3 slurm.confの設定

Epilog=/opt/slurm/etc/scripts/epilog.d/*
# PrologFlags=Alloc -> これは、ジョブ割り当て時にPrologスクリプトが実行される設定なので不要

Epilogスクリプトを実行するためには、PrologFlags=Allocの設定は必要ありません。Epilogスクリプトの実行には以下の設定で十分です。PrologFlags=AllocはPrologスクリプトの動作を制御するための設定です。この設定は、Prologスクリプトをジョブ割り当て時に即座に実行させるためのものであり、Epilogスクリプトの実行には影響しません。

Epilogスクリプトは、ジョブ終了時に自動的に実行されます。PrologFlagsの設定に関わらず、Epilogスクリプトは常にジョブ終了時に実行されます5。

ただし、注意点として、Epilogディレクトリには少なくとも1つのスクリプトファイルが存在する必要があります。スクリプトが存在しない場合、ノードがdrain状態になる可能性があります。

SLURM 21.8.8では、Epilogスクリプトの実行に失敗した場合（ゼロ以外の終了コードを返した場合）、そのノードはDRAIN状態に設定されます。

3. テストケース

バッチジョブテスト

sbatch --wrap="sleep 300"
# ジョブ終了後、ログを確認
tail -f /var/log/slurm/epilog/cleanup_*.log

インタラクティブジョブテスト

salloc -N 1
# 別ターミナルでプロセス起動
ssh node1 "sleep 1000" &
exit
# プロセスが削除されていることを確認

4. 動作検証

4.1 動作検証スクリプトを作成

動作検証するためのバッチスクリプトを作成します。
マルチコアのテストを行うため、Python でテストコードを作成しました。

job_multiprocess.py
子プロセスを生成して、適当な時間Sleepを実行してマルチプロセスの状況を生成します。

#!/usr/bin/python3

import os
import time
from multiprocessing import Process


child_process_time_seconds = 30

def run_child() -> None:
    print("Child: I am the child process.")
    print(f"Child: Child's PID: {os.getpid()}")
    print(f"Child: Parent's PID: {os.getppid()}")
    time.sleep(child_process_time_seconds)
    print(f"Child: End of child process time {child_process_time_seconds}")


def start_parent(num_children: int) -> None:
    print("Parent : I am the parent process")
    print(f"Parent: Parent's PID: {os.getpid()}")
    for i in range(num_children):
        print(f"Starting Proces {i}")
        Process(target=run_child).start()


if __name__ == '__main__':
    num = 3
    start_parent(num)

job.sh
マルチプロセスのプログラムを実行する、SLURMのバッチスクリプトです

#!/bin/bash
#SBATCH --partition=part1
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH -J test
#SBATCH -o stdout.%J
#SBATCH -e stderr.%J

pid1=$(ps -u $SLURM_JOB_USER -o pid,cmd | grep -E "[s]lurm" | grep job | grep $SLURM_JOB_ID)
pid2=$(ps -u $SLURM_JOB_USER -o pid,cmd | grep -E "[s]lurm" | grep job | grep $SLURM_JOB_ID | awk '{print $1}')

SCRIPT="/home/settings/Epilog/script/fetch_pids"

#python3 job.py
python3 job_multiprocess.py

pids=$($SCRIPT $pid $SLURM_JOB_USER)
#ps=$(ps -u $SLURM_JOB_USER -o pid,cmd)
ps=$(ps -ef | grep $SLURM_JOB_USER)

echo "Job ID: $SLURM_JOB_ID" > log
echo "ps : $ps" >> log
echo "pid step1: $pid1" >> log
echo "pid step2: $pid2" >> log
echo "fetch_pids: $pids" >> log
echo "SLURM_TASK_PID: $SLURM_TASK_PID" >> log

4.2 動作検証結果例

上記スクリプトを用いて、バッチ処理の動作検証を実施しました。fetch_pidsスクリプトにて適切にPIDのリストを取得できる様子なども確認することができました。

ジョブ開始

[nobuyuki@headnode 01]$ sbatch -p part1 job.sh
Submitted batch job 106
[nobuyuki@headnode 01]$ squeue
             JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
               106     part1     test nobuyuki CG       0:30      1 compute01

実行中のSLURMジョブIDに紐づくプロセスを確認

[root@compute01 script]# ps -u "nobuyuki" -o pid,cmd | grep -E "[s]lurm" | grep job | grep 106
  21147 /bin/bash /var/spool/slurmd/job00106/slurm_script
[root@compute01 script]# ps -ef | grep nobuyuki
nobuyuki   21147   21143  0 12:42 ?        00:00:00 /bin/bash /var/spool/slurmd/job00106/slurm_script
nobuyuki   21159   21147  0 12:42 ?        00:00:00 python3 job_multiprocess.py
nobuyuki   21160   21159  0 12:42 ?        00:00:00 python3 job_multiprocess.py
nobuyuki   21161   21159  0 12:42 ?        00:00:00 python3 job_multiprocess.py
nobuyuki   21162   21159  0 12:42 ?        00:00:00 python3 job_multiprocess.py
root       21168    2850  0 12:42 pts/0    00:00:00 grep --color=auto nobuyuki
[root@compute01 script]# ./fetch_pids 21147 nobuyuki
21147
21159
21160
21161
21162
[root@compute01 script]#

実行スクリプトで記述したログを確認
ＰＩＤなどが適切に取得できていることをできる。

[nobuyuki@headnode 01]$ cat log
Job ID: 106
SLURM_TASK_PID: 21147
[nobuyuki@headnode 01]$ cat stdout.106
Parent : I am the parent process
Parent: Parent's PID: 21159
Starting Proces 0
Starting Proces 1
Starting Proces 2
Child: I am the child process.
Child: Child's PID: 21160
Child: Parent's PID: 21159
Child: End of child process time 30
Child: I am the child process.
Child: Child's PID: 21161
Child: Parent's PID: 21159
Child: End of child process time 30
Child: I am the child process.
Child: Child's PID: 21162
Child: Parent's PID: 21159
Child: End of child process time 30
[nobuyuki@headnode 01]$

Epilogのログを確認
バッチ処理のため、fetch_pidsが機能している。バッチ処理のCOMPLETEでは、残存プロセスがなかったため、SLURMジョブIDに紐ずくプロセスは存在しなかった。

[root@manage log]# cd 106
[root@manage 106]# ls
cleanup_106_compute01_20250330124254.log
[root@manage 106]# cat cleanup_106_compute01_20250330124254.log
=== Epilog Start [Sun Mar 30 12:42:56 JST 2025] ===
Job ID: 106
User: nobuyuki
STATUS: COMPLETED
Node: compute01
Dry-run Mode: 1
[Batch Job] Process cleanup started
SLURM PID:
SLURM_TASK_PID:
/home/settings/Epilog/script/fetch_pids : start
/home/settings/Epilog/script/fetch_pids : end
[DRY-RUN] Batch processes to be killed:

=== Epilog End [Sun Mar 30 12:42:57 JST 2025] ===
[root@manage 106]#

バッチ処理について、想定通りの挙動であることを確認することができた。

インタラクティブジョブについては、ノード占有で運用している観点から、バッチ処理のような細かいチェックは行わない。ジョブ終了時には、占有していたノードに残存するSLURMユーザのプロセスをすべて削除するロジックとなっている。
テスト環境では、インタラクティブジョブのテストはできなかったが、別の環境で問題ないことを確認済み。

フェールセーフ処理は、SLURMジョブIDの親プロセスのみを削除する、より安全側の処理を実装。

Epilogスクリプトのフェールセーフとなる条件は、バッチ・インタラクティブ処理の判定ができなかった場合、もしくはそのほかのエラーが発生した場合に、フェールセーフを実行するロジックとしている。