Skip to content

Instantly share code, notes, and snippets.

@venetanji
Last active July 4, 2024 12:07
Show Gist options
  • Save venetanji/05b9e503273b3845c4e27a4879be3ded to your computer and use it in GitHub Desktop.
Save venetanji/05b9e503273b3845c4e27a4879be3ded to your computer and use it in GitHub Desktop.
#!/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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment