A bash script used to select one or multiple free CUDA device(s).

select-cuda.sh

#!/bin/bash
# Usage:
# source ./select-cuda.sh [-c <cnt>] [-f]
#   -c cnt: count of GPUs that you want to use
#   -f    : select a GPU even it's being used
# 
# Source this file everytime before running your ML scripts.

count=1
force=0

OPTIND=1

while getopts c:f flags
do
    case "${flags}" in
        c) count=${OPTARG};;
        f) force=1;;
    esac
done

echo "Requsting $count device"'(s)' >&2
if [ "$force" -eq 1 ]; then
    echo "Force mode enabled" >&2
else
    echo "Force mode disabled" >&2
fi

device_info=$(nvidia-smi --format=csv,noheader,nounits --query-gpu=index,memory.free,memory.used | awk '/([0-9]+),\s*([0-9]+),\s*([0-9]+)\s*/{print $1,$2,$3=="0"}' | sort -nrk2 | head -n $count)

free_or_not=$(echo "$device_info" | awk '{print $3}')
free_count=$(echo "$free_or_not" | grep -c "1")

skip=0

if [ "$free_count" -lt "$count" ]; then
    echo 'Allocated '$count' device(s), but only '$free_count' device(s) are free.' >&2

    if [ "$force" -eq 0 ]; then
        echo "Add -f to allocate busy device(s)" >&2
        echo "Failed. Exiting..." >&2
        skip=1
    fi
fi

if [ "$skip" -eq 0 ]; then
    assigned_gpus=$(echo "$device_info" | awk '{print $1}' | tr -d '\n' | sed 's/,$//')

    echo "Selecting the following CUDA devices:"
    echo "$assigned_gpus"
    export CUDA_VISIBLE_DEVICES=$assigned_gpus
fi