1#!/bin/bash 2 3# this script makes it easy parallize collecting data across using multiple GPUs 4 5# Check if tmux is installed 6if ! command -v tmux &> /dev/null; then 7 echo "tmux is not installed. Please install it and try again." 8 exit 1 9fi 10 11# Check if the correct number of arguments is provided 12if [ "$#" -ne 5 ]; then 13 echo "Usage: $0 \"<python_command>\" <comma_separated_device_numbers> <num_samples to generate> <CONDA_ENV> <OUTPUT_DIR>" 14 echo "Example: $0 \"python run.py --a b --b c\" 1,4,5,3 1000 pytorch-3.10 a100" 15 exit 1 16fi 17 18PYTHON_COMMAND=$1 19DEVICE_NUMBERS=$2 20NUM_SAMPLES=$3 21CONDA_ENV=$4 22OUTPUT_DIR=$5 23 24# Create a new tmux session 25SESSION_NAME="parallel_run_$(date +%s)" 26tmux new-session -d -s "$SESSION_NAME" 27 28# Split the device numbers 29IFS=',' read -ra DEVICES <<< "$DEVICE_NUMBERS" 30 31NUM_GPUS=${#DEVICES[@]} 32NUM_SAMPLES_PER_GPU=$((NUM_SAMPLES / NUM_GPUS)) 33echo "AutoHeuristic will collect ${NUM_SAMPLES} samples split across ${NUM_GPUS} GPUs" 34echo "Each GPU will collect ${NUM_SAMPLES_PER_GPU}" 35 36# Function to create a new pane and run the script 37create_pane() { 38 local device=$1 39 tmux split-window -t "$SESSION_NAME" 40 tmux send-keys -t "$SESSION_NAME" "conda activate ${CONDA_ENV} && $PYTHON_COMMAND --device $device -o ${OUTPUT_DIR}/data_${device}.txt --num-samples ${NUM_SAMPLES_PER_GPU}" C-m 41} 42 43# Create panes for each device number 44for device in "${DEVICES[@]}"; do 45 create_pane ${device} 46done 47 48# Remove the first pane (empty one) 49tmux kill-pane -t "$SESSION_NAME.0" 50 51# Arrange panes in a tiled layout 52tmux select-layout -t "$SESSION_NAME" tiled 53 54# Attach to the tmux session 55tmux attach-session -t "$SESSION_NAME" 56