File size: 2,009 Bytes
4008bf9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
#!/bin/bash
# This script is adapted from https://github.com/peregilk/ttconnect#ttconnect
zone="us-central2-b" # TPU v4's always are in us-central2-b. Update if using TPU v2/v3's
name=$1
echo "Connecting to $name";
## Some basic checks if the input is valid
output=$(gcloud compute tpus describe $name --zone $zone 2>/dev/null)
if [ $? != 0 ]; then
echo "Could not find a tpu-v4 with this name in the zone $zone. Exiting."
exit 1
fi
tputype=$(echo $output | awk '{print $2}')
tpusize=$(echo $tputype| cut -c4-)
size="$(($tpusize / 8))"
if (( $size < 1 )); then
echo "This is reported as a $tputype with $size tpu(s). This is not a valid tpu-v4 resource. Exiting."
exit 1
fi
# Check if the session exists, if not create it
# If there already is a session with this name, it will just attach
tmux has-session -t $name 2>/dev/null
if [ $? != 0 ]; then
tmux new-session -d -s $name
tmux select-layout main-vertical
for i in $(seq $(($size-1))); do
tmux split-window -v -d -t $name
# Making sure there is space to split
tmux select-layout main-horizontal
done
for i in $(seq $(($size))); do
worker=$(($i -1))
command="gcloud alpha compute tpus tpu-vm ssh $name --zone $zone --worker $worker"
tmux select-pane -t $name:0.$worker
tmux send-keys -t $name "$command" Enter
done
# Select the final layout
if ((size >= 16));then
tmux select-layout tiled
else
tmux select-layout tiled
tmux select-layout main-vertical
fi
# Enable mouse control - for changing pane size
# Disabled for now since it makes copying more difficult
# tmux set-mouse on
# Move cursor to worker 0
tmux select-pane -t $name:0.0
# Resize the left window
tmux resize-pane -L 50
# Set pane synchronization
tmux set-window-option -t $name:0 synchronize-panes on
# Set pane border format
tmux set-option -t $name pane-border-status top
tmux set-option -t $name pane-border-format "worker #{pane_index} "
fi
# Attach to the session
tmux attach -t $name
|