ttconnect

#!/bin/bash

# Set default zone if not provided
zone=${2:-us-central2-b}
name=$1

# Check if the name argument is provided
if [ -z "$name" ]; then
    echo "Usage: $0 NAME [ZONE]"
    exit 1
fi

echo "Connecting to $name in zone $zone";

# Some basic checks if the input is valid
output=$(gcloud compute tpus describe $name --zone $zone 2>&1)
if [ $? -ne 0 ]; then
    echo "Could not find a tpu-v4 with the name $name in the zone $zone. Exiting."
    exit 1
fi

# Extract the accelerator type and numeric size
tputype=$(echo "$output" | grep 'acceleratorType' | awk '{print $2}' | tr -d '"' | tr -d '\r' | tr -d '\n')
tpusize=$(echo "$tputype" | grep -o '[0-9]\+$')

# Initialize divisor based on TPU type
if [[ "$tputype" == *"v5litepod"* ]]; then
    divisor=4
else
    divisor=8
fi

# Calculate the size with integer arithmetic
let size=tpusize/divisor

if (( size < 1 )); then
    echo "This is reported as a $tputype with $size tpu(s). This is not a valid tpu resource. Exiting."
    exit 1
fi

# Function to check if tmux session exists with exact name
session_exists() {
    tmux list-sessions | grep -q "^$1:"
}

# Check if the tmux session exists, if not create it
# This now uses exact matching for session names
session_exists $name
if [ $? != 0 ]; then
    tmux new-session -d -s $name
    tmux select-layout main-horizontal

    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-horizontal
    fi

    # 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, now ensuring exact name match
if session_exists $name; then
    tmux attach -t $name
else
    echo "No tmux session with the exact name '$name' found."
fi