#!/bin/bash
# Detect resources of a list of hosts and print OAR commands to create them in OAR database.

usage() {
        cat <<EOF
Usage:
    ${0##*/} [options] <hostfile>

Generate a script for the creation of the OAR resources of a list of hosts, by
connecting to each host and gathering system information.
Hosts are read one per line from hostfile (must be readable by user oar) or
STDIN (if hostfile is -).

Options:
  -t          create the thread hierarchy level (HyperThreading/SMT)
  -n          create the numa hierarchy level and numamem properties
  -g          create the gpu hierarchy level and gpudevice properties (but not populated)
  -o <file>   script filename
  -x          execute the script at the end and remove it
  -y          assume "yes" as answer to all prompts and run non-interactively
  -v          print the generated script
  -e <path>   path to the ssh command
  -s <path>   path to the oarnodesetting command
  -p <path>   path to the oarproperty command
  -c <path>   path to the oar.conf file
  -D          enable debug mode
  -h          display this message

EOF

}

parse_opts() {
    while getopts "c:e:p:s:o:xytngvDh" opt; do
        case $opt in
        c)
            OARCONFFILE=$OPTARG
            ;;
        e)
            OPENSSH_CMD=$OPTARG
            ;;
        s)
            OARNODESETTING_CMD=$OPTARG
            ;;
        p)
            OARPROPERTY_CMD=$OPTARG
            ;;
        o)
            SCRIPT_FILE=$OPTARG
            ;;
        x)
            EXECUTE=1
            ;;
        v)
            PRINT=1
            ;;
        y)
            NON_INTERACTIVE=1
            ;;
        t)
            WITH_THREAD=1
            ;;
        n)
            WITH_NUMA=1
            ;;
        g)
            WITH_GPU=1
            ;;
        D)
            DEBUG=1
            ;;
        h)
            usage
            exit 0
            ;;
        \?)
            echo -e "Error: invalid option: -$OPTARG\n\n" >&2
            usage 1>&2
            exit 1
          ;;
      esac
    done
}

OARCONFFILE="/etc/oar/oar.conf"
OPENSSH_CMD="/usr/bin/ssh"
OARPROPERTY_CMD="oarproperty"
OARNODESETTING_CMD="oarnodesetting"
unset WITH_THREAD
unset WITH_NUMA
unset WITH_GPU

parse_opts "$@"
shift $((OPTIND-1))
HOSTFILE=$1
shift
OPTIND=1
parse_opts "$@"

if ! [ -r "$OARCONFFILE" ]; then
    echo -e "Error: cannot read OAR configuration file $OARCONFFILE\n\n" 1>&2
    usage 1>&2
    exit 1
fi

# check host file
if [ -z "$HOSTFILE" -o \( "$HOSTFILE" != "-" -a ! -r "$HOSTFILE" \) ]; then
    echo -e "Error: cannot read host file (warning: must be readable by the oar user)\n" 1>&2
    usage 1>&2
    exit 1
fi
. "$OARCONFFILE" || exit 20

if [ "$HOSTFILE" == "-" ]; then
    if [ -z "$NON_INTERACTIVE" ]; then
        echo -e "Error: reading nodes from STDIN requires the script to run non-interactive (-y option).\n" 1>&2
        usage 1>&2
        exit 1
    fi
    echo "# Reading nodes from STDIN..."
    exec 4<&0
    exec 0<&-
else
    echo "# Reading nodes from ${HOSTFILE}..."
    exec 4<$HOSTFILE
fi

if [ -z "$NON_INTERACTIVE" ]; then
    read -s -n 1 -p "# Is the OAR SSH key configured on all hosts? (press y to confirm)"
    echo
    [ "$REPLY" != "y" ] && exit 2
fi

# Check output file
if [ -n "$SCRIPT_FILE" ]; then
    if [ -e "$SCRIPT_FILE" ]; then
        if [ -z "$NON_INTERACTIVE" ]; then
            read -s -n 1 -p "# Script file $SCRIPT_FILE already exists, overwrite it? (press y to confirm)"
            echo
            [ "$REPLY" != "y" ] && exit 2
        fi
        echo -n > $SCRIPT_FILE || exit 4
    fi
else
    SCRIPT_FILE=$(mktemp -p "" ${0##*/}.XXXXX)
fi

# Add properties
echo "$OARPROPERTY_CMD -a cpu 2> /dev/null || true" > $SCRIPT_FILE
echo "$OARPROPERTY_CMD -a core 2> /dev/null || true" >> $SCRIPT_FILE
[ -n "$WITH_THREAD" ] && echo "$OARPROPERTY_CMD -a thread 2> /dev/null || true" >> $SCRIPT_FILE
echo "$OARPROPERTY_CMD -c -a host 2> /dev/null || true" >> $SCRIPT_FILE
echo "$OARPROPERTY_CMD -a mem 2> /dev/null || true" >> $SCRIPT_FILE
[ -n "$WITH_NUMA" ] && echo "$OARPROPERTY_CMD -a numa 2> /dev/null || true" >> $SCRIPT_FILE
[ -n "$WITH_NUMA" ] && echo "$OARPROPERTY_CMD -a numamem 2> /dev/null || true" >> $SCRIPT_FILE
[ -n "$WITH_GPU" ] && echo "$OARPROPERTY_CMD -a gpu 2> /dev/null || true" >> $SCRIPT_FILE
[ -n "$WITH_GPU" ] && echo "$OARPROPERTY_CMD -c -a gpudevice 2> /dev/null || true" >> $SCRIPT_FILE

# Init CPU id
CPU=$($OARNODESETTING_CMD --last-property-value cpu 2> /dev/null)
[ -n "$CPU" ] || CPU=0
# Init NUMA id
if [ -n "$WITH_NUMA" ]; then
    NUMA=$($OARNODESETTING_CMD --last-property-value numa 2> /dev/null)
fi
[ -n "$NUMA" ] || NUMA=0
# Init CORE id
CORE=$($OARNODESETTING_CMD --last-property-value core 2> /dev/null)
[ -n "$CORE" ] || CORE=0
# Init THREAD id
if [ -n "$WITH_THREAD" ]; then
    THREAD=$($OARNODESETTING_CMD --last-property-value thread 2> /dev/null)
fi
[ -n "$THREAD" ] || THREAD=0

write_oarnodesetting() {
    if [ -n "${CPUSET[*]}" ]; then
    [ -n "$DEBUG" ] && ( IFS=, ; echo "DEBUG|found: host=$host cpu=$CPU${WITH_NUMA:+ numa=$NUMA} core=$CORE${WITH_THREAD:+ thread=$THREAD} cpuset=${CPUSET[*]} mem=$MEM${WITH_NUMA:+ numamem=$NUMAMEM}" )
        ( IFS=, ; echo "$OARNODESETTING_CMD -a -h $host -p host=$host -p cpu=$CPU${WITH_NUMA:+ -p numa=$NUMA} -p core=$CORE${WITH_THREAD:+ -p thread=$THREAD} -p cpuset=${CPUSET[*]} -p mem=$MEM${WITH_NUMA:+ -p numamem=$NUMAMEM}${WITH_GPU:+ -p gpudevice=''}" >> "$SCRIPT_FILE" )
        unset CPUSET
    fi
}
while read -u 4 host; do
    echo "# Fetching info from $host... "
    while IFS= read -r line; do
        [ -n "$DEBUG" ] && echo "DEBUG|read: $line"
        if [[ $line =~ ^[[:space:]]+PU[[:space:]]P#([[:digit:]]+)$ ]]; then
            ((THREAD++))
            CPUSET+=( "${BASH_REMATCH[1]}" )
            [ -n "$WITH_THREAD" ] && write_oarnodesetting # with threads: write as soon as a PU is found
        elif [[ $line =~ ^[[:space:]]+Core[[:space:]]P#[[:digit:]]+[[:space:]]\+[[:space:]]PU[[:space:]]P#([[:digit:]]+)$ ]]; then
            ((CORE++))
            ((THREAD++))
            CPUSET+=( "${BASH_REMATCH[1]}" )
            [ -n "$WITH_THREAD" ] && write_oarnodesetting # with threads: write as soon as a PU is found
        else
            [ -z "$WITH_THREAD" ] && write_oarnodesetting # without threads: write all PUs after a new core/cpu/package/... is found
            if [[ $line =~ ^Machine[[:space:]]\(([[:digit:]]+)([^[:space:]]+).*\)$ ]]; then
                MEM=${BASH_REMATCH[1]}
                if [ "${BASH_REMATCH[2]}" != "GB" ]; then
                    echo "Warning: host memory unit is ${BASH_REMATCH[2]}, not GB"
                fi
            elif [[ $line =~ ^[[:space:]]+Package[[:space:]]P#[[:digit:]]+$ ]]; then
                ((CPU++))
            elif [[ $line =~ ^Machine[[:space:]]\(([[:digit:]]+)([^[:space:]]+).*\)[[:space:]]\+[[:space:]]Package[[:space:]]P#[[:digit:]]+$ ]]; then
                MEM=${BASH_REMATCH[1]}
                if [ "${BASH_REMATCH[2]}" != "GB" ]; then
                    echo "Warning: host memory unit is ${BASH_REMATCH[2]}, not GB"
                fi
                ((CPU++))
            elif [[ $line =~ ^[[:space:]]+NUMANode[[:space:]]P#[[:digit:]]+[[:space:]]\(([[:digit:]]+)(.+)\)$ ]]; then
                NUMAMEM=${BASH_REMATCH[1]}
                if [ "${BASH_REMATCH[2]}" != "GB" ]; then
                    echo "Warning: CPU memory unit is ${BASH_REMATCH[2]}, not GB"
                fi
                ((NUMA++))
            elif [[ $line =~ ^[[:space:]]+Core[[:space:]]P#[[:digit:]]+$ ]]; then
                ((CORE++))
            fi
        fi
    done < <($OPENSSH_CMD $host hwloc-ls -p --no-io --no-caches)
    [ -z "$WITH_THREAD" ] && write_oarnodesetting # without threads: latest found PUs may still have to be written (if no other line after the last PU line)
done
echo -ne "# Done\n\n"

if [ -n "$PRINT" ]; then
    echo -ne "# Generated script content:\n\n"
    cat $SCRIPT_FILE
    echo
fi
if [ -n "$EXECUTE" ]; then
    echo -ne "# Executing script...\n\n"
    . $SCRIPT_FILE
    echo -ne "\n# Removing script file.\n"
    rm $SCRIPT_FILE
else
    cat <<EOF
# Please please find the generated script here: $SCRIPT_FILE
# It everything looks correct to you, source it to create the resources, by running:
# . $SCRIPT_FILE

EOF
fi
