interact

#!/bin/bash -l
# Script to start interactive job on ARC systems

# Define color
RED='\033[0;31m'
YELLOW='\033[0;33m'
NC='\033[0m' # No Color

# User-provided arguments
original_args="$@"
# We'll add these
extra_args="--job-name='INTERACT'"

# System-specific defaults
case "${SYSNAME}" in
    tinkercliffs)
        def_partition="normal_q"
        def_gres=""
        ;;
    owl)
        def_partition="normal_q"
        def_gres=""
        ;;
    falcon)
        def_partition="v100_normal_q,t4_normal_q"
        def_gres="gpu:1"
        ;;
    *)
        echo "SYSNAME '${SYSNAME}' not recognized. Exiting."
        exit 1
        ;;
esac

# Print help if requested or if no argument is provided
if [[ " $original_args" == *" --help"* || $# -eq 0 ]]; then
  echo "Request an interactive job"
  echo "Usage: interact [OPTIONS]"
  echo "where OPTIONS are any valid Slurm (sbatch/salloc/srun) options, or:"
  echo "  --verbose     Print the Slurm command being run before starting job (for debugging purposes)"
  echo "  --help        Print this message"
  echo
  echo "To run, an account should be provided with -A [account] or --account=[account]. See 'quota' for a list of valid accounts and utilization metrics."
  exit
fi

# Add verbosity if requested
[[ " $original_args" == *" --verbose"* ]] && verbose=1 || verbose=0

# If account not specified by user
if [[ ! " $original_args" == *" -A"* && ! " $original_args" == *" --account"* ]]; then
  echo "An account must be specified with -A or --account. A list of valid allocation accounts (and remaining balances) can be found with the command 'quota'."
  exit 1
fi

# If partition not specified by user
if [[ ! " $original_args" == *" -p"* && ! " $original_args" == *" --partition"* ]]; then
  echo -e "${YELLOW} --- No partition specified:${NC} setting --partition=$def_partition"
  extra_args="--partition=$def_partition $extra_args"
fi

# If gres not specified by user (and def_gres not empty)
if [[ $original_args != *"--gres=gpu"* && \
    $original_args != *"--gres gpu"* && \
    $original_args != *"--gpus"* && \
    ! -z $def_gres ]]; then
    echo -e "${YELLOW} --- No gres specified:${NC} setting --gres=$def_gres"
    extra_args="--gres=$def_gres $extra_args"
fi

# Check if a GPU is requested on a non-GPU partition
if [[ $original_args =~ "--gres=gpu" || $original_args =~ "--gres gpu" || $original_args =~ "--gpus" ]]; then
    # Check if a non-GPU partition is specified using regex to catch -p, --partition, spaces, and equals signs
    if [[ $original_args =~ (-p|--partition)[[:space:]=]*(normal_q|preemptable_q) ]]; then
        echo -e "${RED} --- Invalid Request:${NC} You cannot request GPUs on a non-GPU partition (normal_q or preemptable_q)."
        exit 1
    fi
fi

# Resolve the effective partition for QoS mapping & if the user passed -p or --partition, extract it from their args.
if [[ ! " $original_args" == *" -p"* && ! " $original_args" == *" --partition"* ]]; then
    effective_partition="$def_partition"
else
    effective_partition="$(echo "$original_args" | grep -oP '(?<=-p |--partition=|--partition )\S+')"
fi

# If QoS not specified by user, map the effective partition to its interactive QoS and inject it. 
# Todo: (add mapping for new nodes & partitions)
if [[ ! " $original_args" == *" --qos"* ]]; then
    case "${SYSNAME}:${effective_partition}" in
        tinkercliffs:normal_q)            def_qos="tc_normal_int" ;;
        tinkercliffs:a100_normal_q)       def_qos="tc_a100_normal_int" ;;
        tinkercliffs:h200_normal_q)       def_qos="tc_h200_normal_int" ;;
        owl:normal_q)                     def_qos="owl_normal_int" ;;
        falcon:v100_normal_q,t4_normal_q) def_qos="fal_v100_normal_int" ;;
        falcon:v100_normal_q)             def_qos="fal_v100_normal_int" ;;
        falcon:t4_normal_q)               def_qos="fal_t4_normal_int" ;;
        falcon:a30_normal_q)              def_qos="fal_a30_normal_int" ;;
        falcon:l40s_normal_q)             def_qos="fal_l40s_normal_int" ;;
        *)
            echo -e "${YELLOW} --- Warning:${NC} No interactive QoS mapping for ${SYSNAME}:${effective_partition}. Proceeding without --qos."
            def_qos=""
            ;;
    esac

    if [[ -n "$def_qos" ]]; then
        #uncomment to debug QoS mapping
        # echo -e "${YELLOW} --- No QoS specified:${NC} setting --qos=$def_qos"
        extra_args="--qos=$def_qos $extra_args"
    fi
fi

# Command to be run
# Later arguments appear to be given priority so add user-provided
# arguments second on the off chance of conflicts
cmd="srun $extra_args $original_args --pty $SHELL"
# Print the command being run if user requested verbose mode
[[ $verbose -gt 0 ]] && echo "Running: $cmd"

# Print the warning message in red
echo -e "${RED} --- Warning:${NC}"
echo "     Your session consumes resources (CPUs, memory, and GPUs) while it remains open."
echo "     Close your session whenever you finish your work."
echo "     Other users cannot use the resources allocated to your job until you close your session."
echo "     Consider the use of batch jobs to optimize resources allocation."

# Run
$cmd