// At the bottom custom metadata of instance templates
// Key: startup-script

Copy and paste the value below and update the docker value paths and any env variables
// Value:

#!/bin/bash

# Ensure Python3 is installed
if ! command -v python3 &> /dev/null; then
    echo "Installing Python3..."
    apt-get update
    apt-get install -y python3
    echo "Python3 installed."
else
    echo "Python3 is already installed."
fi

# Ensure curl is installed
if ! command -v curl &> /dev/null; then
    echo "Installing curl..."
    apt-get update
    apt-get install -y curl
    echo "curl installed."
else
    echo "curl is already installed."
fi

# Install NVIDIA drivers using Google's recommended method
if test -f /opt/google/cuda-installer; then
    echo "NVIDIA drivers are already installed."
else
    echo "Installing NVIDIA drivers using Google's recommended installer..."
    mkdir -p /opt/google/cuda-installer
    cd /opt/google/cuda-installer/ || exit
    curl -fSsL -O https://github.com/GoogleCloudPlatform/compute-gpu-installation/releases/download/cuda-installer-v1.1.0/cuda_installer.pyz
    python3 cuda_installer.pyz install_cuda
    echo "NVIDIA drivers installed successfully."
fi

# Check if Docker is installed
if ! command -v docker &> /dev/null; then
    echo "Installing Docker..."
    apt-get update
    apt-get install -y docker.io
    systemctl enable docker
    systemctl start docker
    echo "Docker installed and started."
else
    echo "Docker is already installed."
fi

# Install NVIDIA Container Toolkit if not already installed
if ! command -v nvidia-ctk &> /dev/null; then
    echo "Installing NVIDIA Container Toolkit..."
    curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
    curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
        sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
        tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
    apt-get update
    apt-get install -y nvidia-container-toolkit
    nvidia-ctk runtime configure --runtime=docker
    systemctl restart docker
    echo "NVIDIA Container Toolkit installed and configured."
else
    echo "NVIDIA Container Toolkit is already installed."
fi

# Authenticate with Google Container Registry if necessary
if ! docker images | grep -q 'us.gcr.io/your/path/here'; then
    echo "Pulling Docker image from Google Container Registry..."
    if ! command -v gcloud &> /dev/null; then
        echo "Installing Google Cloud SDK..."
        apt-get install -y google-cloud-sdk
    fi
    gcloud auth configure-docker us.gcr.io
    docker pull us.gcr.io/your/path/here:version-1.0-snapshot
    echo "Docker image pulled."
fi

# Run the Docker container with environment variables
echo "Starting Docker container..."
docker run --gpus all -p 3000:3000 -d \
  -e EXAMPLE_ENV='test1' \
  -e EXAMPLE_ENV_2='server' \
  us.gcr.io/your/path/here:version-1.0-snapshot
echo "Docker container started."