setcuda.sh

#!/bin/bash

# Check if NVIDIA driver is available
if ! command -v nvidia-smi &> /dev/null; then
  echo "[setcuda] Error: NVIDIA driver not found"
  return 1
fi

# Check if the `setcuda` command is already defined
if declare -f setcuda > /dev/null; then
  if [[ -z "$SETCUDA_NUM_GPU" ]]; then
    echo "[setcuda] Error: `setcuda` command is already defined by another script."
    return 1
  fi

  # Assume that the script is already loaded
  return
fi

# Number of available GPUs
SETCUDA_NUM_GPU=$(nvidia-smi --list-gpus | wc -l)

# Function to set or clear CUDA_VISIBLE_DEVICES
function setcuda() {
  # Show help if no argument is provided
  if [ -z "$1" ]; then
    echo "Usage: setcuda <none|any|gpu_ids>"
    echo ""
    echo "Examples:"
    echo "  setcuda any     # Enables all GPUs"
    echo "  setcuda none    # Disables all GPUs"
    echo "  setcuda 0       # Enables specified GPU 0"
    echo "  setcuda 0,2     # Enables specified GPUs 0 and 2"
    return
  fi

  case "$1" in
    any)
      unset CUDA_VISIBLE_DEVICES
      ;;

    none)
      export CUDA_VISIBLE_DEVICES=""
      ;;

    *)
      # Check each GPU ID in the comma-separated list
      local IFS=','
      local -a GPU_IDS
      read -ra GPU_IDS <<< "$1"

      local -A seen_gpus
      local gpu_id
      for gpu_id in "${GPU_IDS[@]}"; do
        # Check if GPU ID is valid
        if ! [[ "$gpu_id" =~ ^[0-9]+$ ]] || (( gpu_id >= SETCUDA_NUM_GPU )); then
          echo "Error: Invalid GPU $gpu_id."
          return 1
        fi

        # Check for duplicate GPU IDs
        if [[ -n "${seen_gpus[$gpu_id]}" ]]; then
          echo "Error: Duplicate GPU $gpu_id."
          return 1
        fi

        seen_gpus[$gpu_id]=1
      done

      export CUDA_VISIBLE_DEVICES="$1"
      ;;
  esac

  setcuda_update_ps1
}

# Function to update the PS1 prompt
function setcuda_update_ps1() {
  local devices
  if [ -z "${CUDA_VISIBLE_DEVICES+x}" ]; then
    devices="any"
  elif [ -z "$CUDA_VISIBLE_DEVICES" ]; then
    devices="none"
  else
    devices="$CUDA_VISIBLE_DEVICES"
  fi


  PS1=$(echo "$PS1" | sed -E 's/\[cuda:[^]]*\] //')
  PS1="[cuda:$devices] $PS1"
}

# Add the update function to PROMPT_COMMAND
PROMPT_COMMAND="setcuda_update_ps1; $PROMPT_COMMAND"