gidd_tpu_setup.sh

#set -Eeuo pipefail

printf "n\n" | ssh-keygen -f ~/.ssh/google_compute_engine -t rsa -P "" > /dev/null

printf "$TPU_NAME\n\n\n\n\n\n\n\n" | script -qec 'bash <(curl -sL https://raw.githubusercontent.com/dvruette/EasyDeL/refs/heads/main/tpu_setup.sh)' /dev/null

source /home/dimitri/.bashrc
source ~/easy-venv/bin/activate

git clone https://github.com/dvruette/gidd-easydel.git
cd gidd-easydel
git pull

ZONE=$(curl -s "http://metadata.google.internal/computeMetadata/v1/instance/zone" -H "Metadata-Flavor: Google" | cut -d/ -f4 || true)
TPU_TYPE=$(gcloud compute tpus tpu-vm describe "$TPU_NAME" --zone="$ZONE" --format="value(acceleratorType)" 2>/dev/null | awk -F'/' '{print $NF}')

UV_PATH=$HOME/.local/bin/uv
PY_PATH=$(which python)
eopod run "$UV_PATH pip uninstall --python $PY_PATH easydel && $UV_PATH pip install --python $PY_PATH -U git+https://github.com/dvruette/easydel.git[tpu,torch] && $UV_PATH pip install --python $PY_PATH $(cat requirements.txt | tr '\n' ' ')"

eopod auto-config-ray --self-job --python-path $(which python)

# Install GNU parallel
wget http://ftp.jaist.ac.jp/pub/GNU/parallel/parallel-latest.tar.bz2
tar -xjf parallel-latest.tar.bz2
pushd $(echo parallel-* | awk '{print $1}')
sudo ./configure && make
sudo make install
popd
rm -rf parallel-*

echo "export HF_TOKEN_FOR_EASYDEL=$HF_TOKEN_FOR_EASYDEL" >> ~/.bashrc
echo "export WANDB_API_KEY_FOR_EASYDEL=$WANDB_API_KEY_FOR_EASYDEL" >> ~/.bashrc
echo "export DATA_FILES=$DATA_FILES" >> ~/.bashrc
echo "export SAVE_DIRECTORY=$SAVE_DIRECTORY" >> ~/.bashrc
echo "export TPU_VERSION=$TPU_TYPE" >> ~/.bashrc
echo "export TPU_NAME=$TPU_NAME" >> ~/.bashrc
echo "export TPU_ZONE=$ZONE" >> ~/.bashrc
echo "source $HOME/easy-venv/bin/activate" >> ~/.bashrc
echo "alias s='cd gidd-easydel && tmux'" >> ~/.bashrc
echo "alias a='tmux a'" >> ~/.bashrc
echo "alias reinstall_easydel='eopod run \"\$(which uv) pip uninstall --python \$(which python) easydel && \$(which uv) pip install --python \$(which python) -U git+https://github.com/dvruette/easydel.git[tpu,torch]\"'" >> ~/.bashrc
echo "alias restart_ray='eopod auto-config-ray --self-job --python-path \$(which python)'" >> ~/.bashrc