Last active
July 4, 2024 12:07
-
-
Save venetanji/05b9e503273b3845c4e27a4879be3ded to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/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