greene_startup.sh

#!/bin/bash

# Step 0: Ensure .cache and .local link are symlinks
if [ ! -L ~/.cache ]; then
    # create the directory if it doesn't exist
    mkdir -p /vast/${USER}/.cache
    ln -s /vast/${USER}/.cache ~/.cache
fi

if [ ! -L ~/.local ]; then
    # create the directory if it doesn't exist
    mkdir -p /vast/${USER}/.local
    ln -s /vast/${USER}/.local ~/.local
fi

# Step 1: Check for overlay file and create if not found
overlay_name=${1:-default_overlay}
mkdir -p /scratch/${USER}/overlays
overlay_path="/scratch/${USER}/overlays/${overlay_name}.img"

if [ ! -f "$overlay_path" ]; then
    echo "Overlay file not found. Creating..."
    singularity overlay create --size 15000 $overlay_path
fi

# Step 2: Ask user for singularity image
read -p "Enter the singularity image path (default: docker://yanwk/comfyui-boot): " singularity_image
singularity_image=${singularity_image:-docker://yanwk/comfyui-boot}

export SINGULARITY_CACHEDIR=/vast/${USER}/.singularity

# Check if the image is a Docker image
if [[ $singularity_image == docker://* ]]; then
    # Extract image name and tag, default tag to latest if not specified
    image_name=$(echo $singularity_image | sed 's|docker://||' | cut -d '/' -f2-)
    if [[ $image_name != *":"* ]]; then
        image_name="${image_name}:latest"
    fi
    # Replace '/' and ':' with '_' to create a valid filename
    cache_filename="${image_name//\//_}"
    cache_filename="${cache_filename//:/_}.sif"
    cache_path="/scratch/${USER}/images/${cache_filename}"

    # Check if the image exists in the cache
    if [ -f "$cache_path" ]; then
        echo "Image already cached at $cache_path"
    else
        echo "Pulling image to cache..."
        mkdir -p "/scratch/${USER}/images"
        singularity pull --name "$cache_path" "$singularity_image"
    fi
    singularity_image=$cache_path
fi

# Step 4: Ask user for resources
read -p "Enter the number of GPUs to request (default 1): " gpu_num
gpu_num=${gpu_num:-1}
read -p "Enter the amount of memory in GB to request (default 32): " mem
mem=${mem:-32}
# Request resources and launch bash

# if running comfyui download install comfyui script
if [[ $singularity_image == *comfyui* ]]; then
    # download script from gist https://gist.githubusercontent.com/venetanji/9c000f77ba85512f9cc2e29c686d0f0d/raw/385df524f7eee4d57472423b49342a5a72d4e14b/run_comfy.sh
    curl -s "https://gist.githubusercontent.com/venetanji/9c000f77ba85512f9cc2e29c686d0f0d/raw/run_comfy.sh?q=$RANDOM" > /scratch/${USER}/run_comfy.sh
    chmod +x /scratch/${USER}/run_comfy.sh
    singularity_command=/scratch/${USER}/run_comfy.sh
    
else
    singularity_command=bin/bash
fi

mem=$(($mem*1024))
srun --pty -c 2 --mem=$mem --gres=gpu:$gpu_num singularity exec --nv --overlay $overlay_path:rw $singularity_image $singularity_command

# Step 5: Return hostname and exit

exit