KnowHow

技術的なメモを中心にまとめます。
検索にて調べることができます。

SLURMのgenerate_pbs_nodefileからhostsファイルを調整するスクリプト

登録日 :2024/11/28 05:03
カテゴリ :SLURM

#!/bin/bash

#
# created by Me
# 2024.11.28
# version 1.0
#

# Uses Slurm environment variables to produce a $PBS_NODEFILE -style
#   output file.  The output goes in a temporary file and the name of
#   the file is printed on hosts file.


function echo_usage() {
    echo "Usage  : $0 [Positive number]"
    echo "Example: $0 40"

}

function echo_messages(){
    echo "------------------------------------------------------------------"
    echo "Description:"
    echo " Please enter the number of cores you want to use per node."
    echo " Uses Slurm environment variables to produce a PBS_NODEFILE -style,"
    echo " output hosts file."
    echo ""
    echo_usage
    echo "------------------------------------------------------------------"
    echo ""
}


function repeat_nodes() {
    local REPEAT_COUNT=$1
    local OUTPUT_FILE=$2
    local NODEFILE="nod1 nod2"

    # initialize file
    > "$OUTPUT_FILE"

    for node in $NODEFILE; do
        for ((i=1; i<=$REPEAT_COUNT; i++)); do
            echo "$node" >> "$OUTPUT_FILE"
        done
    done

    #echo "Nodes repeated $REPEAT_COUNT times and written to $OUTPUT_FILE"
    echo "Create hosts file"
}


######################################
# main
# example
# repeat_nodes 10 "hosts"
#####################################

# check an argument

if [ $# -eq 0 ]; then
    while true; do
        read -p ">> Enter the repeat count (positive integer): " REPEAT_COUNT
        if [[ $REPEAT_COUNT =~ ^[0-9]+$ ]]; then
            break
        else
            echo_messages
            #echo "Error: Please enter a positive integer"
        fi
    done
elif [ $# -eq 1 ]; then
    REPEAT_COUNT=$1
    if ! [[ $REPEAT_COUNT =~ ^[0-9]+$ ]]; then
        echo "Error: Argument must be a positive integer"
        exit 1
    fi
else
    echo "Usage: $0 [repeat_count]"
    exit 1
fi

# Execute command
repeat_nodes "$REPEAT_COUNT" hosts