mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
b020971e78
Pull Request resolved: https://github.com/pytorch/pytorch/pull/156083 Approved by: https://github.com/jingsh ghstack dependencies: #156079, #156082
56 lines
1.6 KiB
Bash
56 lines
1.6 KiB
Bash
#!/bin/bash
|
|
|
|
# This script makes it easy to parallelize data collection across multiple GPUs
|
|
|
|
# Check if tmux is installed
|
|
if ! command -v tmux &> /dev/null; then
|
|
echo "tmux is not installed. Please install it and try again."
|
|
exit 1
|
|
fi
|
|
|
|
# Check if the correct number of arguments is provided
|
|
if [ "$#" -ne 5 ]; then
|
|
echo "Usage: $0 \"<python_command>\" <comma_separated_device_numbers> <num_samples to generate> <CONDA_ENV> <OUTPUT_DIR>"
|
|
echo "Example: $0 \"python run.py --a b --b c\" 1,4,5,3 1000 pytorch-3.10 a100"
|
|
exit 1
|
|
fi
|
|
|
|
PYTHON_COMMAND=$1
|
|
DEVICE_NUMBERS=$2
|
|
NUM_SAMPLES=$3
|
|
CONDA_ENV=$4
|
|
OUTPUT_DIR=$5
|
|
|
|
# Create a new tmux session
|
|
SESSION_NAME="parallel_run_$(date +%s)"
|
|
tmux new-session -d -s "$SESSION_NAME"
|
|
|
|
# Split the device numbers
|
|
IFS=',' read -ra DEVICES <<< "$DEVICE_NUMBERS"
|
|
|
|
NUM_GPUS=${#DEVICES[@]}
|
|
NUM_SAMPLES_PER_GPU=$((NUM_SAMPLES / NUM_GPUS))
|
|
echo "AutoHeuristic will collect ${NUM_SAMPLES} samples split across ${NUM_GPUS} GPUs"
|
|
echo "Each GPU will collect ${NUM_SAMPLES_PER_GPU}"
|
|
|
|
# Function to create a new pane and run the script
|
|
create_pane() {
|
|
local device=$1
|
|
tmux split-window -t "$SESSION_NAME"
|
|
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
|
|
}
|
|
|
|
# Create panes for each device number
|
|
for device in "${DEVICES[@]}"; do
|
|
create_pane ${device}
|
|
done
|
|
|
|
# Remove the first pane (empty one)
|
|
tmux kill-pane -t "$SESSION_NAME.0"
|
|
|
|
# Arrange panes in a tiled layout
|
|
tmux select-layout -t "$SESSION_NAME" tiled
|
|
|
|
# Attach to the tmux session
|
|
tmux attach-session -t "$SESSION_NAME"
|