Build, compress, cache, copy and remove conda environments on slurm using micromamba.

env-manage.sh

#!/bin/bash

#SBATCH --cpus-per-task=32
#SBATCH --mem=32000
#SBATCH --partition=short
#SBATCH --output=env-manage.log
#SBATCH --mail-type=END,FAIL
#SBATCH --mail-user=

set -e

# Default config variables
ENV_STORAGE_ROOT=           # The directory where the environment will be stored
ENV_DEST_DIR=               # The directory where the environment will be created and copied to during usage
ENV_YML=                    # The environment.yml file
MICROMAMBA_EXEC=            # The micromamba executable


# Function to create environment
build_env() {
    echo Building envrionment using \$ENV_YML=$ENV_YML on host $(cat /etc/hostname).

    rm -rf $ENV_DEST_DIR
    mkdir -p $ENV_DEST_DIR
    mkdir -p $ENV_STORAGE_DIR
    cd $ENV_DEST_DIR

    trap 'echo "Error occurred. Removing destination directory and exiting..."; rm -rf $ENV_DEST_DIR; exit 1' ERR

    eval "$($MICROMAMBA_EXEC shell hook --shell bash)" # To allow micromamba activation

    micromamba create -p ./env -r ./rootenv -y -f "$ENV_YML"

    TIMESTAMP=$(date -Iminutes)
    micromamba activate "$ENV_DEST_DIR/env"
    micromamba env export > "$ENV_STORAGE_DIR/environment-$TIMESTAMP.yml"
    micromamba deactivate

    if [[ $pv_available == true ]]; then
        tar cf - ./env -P | pv -s $(du -sb ./env | awk '{print $1}') > env.tar
    else
        tar cf env.tar ./env
    fi
    zstd -T32 -9 --rsyncable env.tar
    rsync -avP env.tar.zst "$ENV_STORAGE_DIR/env-$TIMESTAMP.tar.zst"
    ln -s -f -r "$ENV_STORAGE_DIR/env-$TIMESTAMP.tar.zst" "$ENV_STORAGE_DIR/env-latest.tar.zst"
    rm -rf "$ENV_DEST_DIR"
    exit 0
}

# Function to copy environment
copy_env() {
    echo Copying envrionment using \$TMPDIR=$TMPDIR on host $(cat /etc/hostname).

    if [[ $TMPDIR == "" ]]; then
        echo "\$TMPDIR is not set."
        exit 1
    fi

    if [ -L "$ENV_DEST_DIR/env" ] && [ -e "$ENV_DEST_DIR/env" ]; then
        echo "Environment already exists."
        exit 0
    fi

    mkdir -p "$ENV_DEST_DIR"
    ln -s -r -f "$TMPDIR/env" "$ENV_DEST_DIR"

    trap 'echo -n "Copying environment was terminated, removing the destination directory and exiting..."; no_ask_remove=true; remove_env; exit 1' SIGINT SIGTERM ERR

    cd "$TMPDIR"
    rsync -avP "$(realpath $ENV_STORAGE_DIR/env-latest.tar.zst)" "$ENV_DEST_DIR/env.tar.zst"
    zstd -d -T16 "$ENV_DEST_DIR/env.tar.zst" -o "$TMPDIR/env.tar"
    if [[ $pv_available == true ]]; then
        pv "$TMPDIR/env.tar" | tar -xf -
    else
        tar -xf "$TMPDIR/env.tar"
    fi
    exit 0
}

# Function to remove environment
remove_env() {
    echo Removing environment $ENV_DEST_DIR on host $(cat /etc/hostname).

    if [[ $force_remove ]]; then
        rm -rf "$ENV_DEST_DIR"
        echo "Environment destination directory has been forcefully deleted."
        exit 0
    fi

    if [[ $(readlink -f -- "$ENV_DEST_DIR/env") == "$TMPDIR/env" ]]; then
        if [[ $no_ask_remove == true ]]; then
            REPLY=Y
        else
            echo "Do you want to delete /tmp/mmicik_env? (y/N) "
            read REPLY || true
        fi
        if [[ $REPLY == [Yy]* ]]; then
            rm -rf "$ENV_DEST_DIR"
            echo "Environment destination directory has been deleted."
        fi
        exit 0
    fi

    if [[ -L "$ENV_DEST_DIR/env" || -d "$ENV_DEST_DIR/env" ]]; then
        echo "Environment destination directory does not exist."
    else
        echo "Environment destination directory does not belong to this job, not removing."
    fi
    exit 0
}

# Function to print configuration
print_config() {
    echo "ENV_DEST_DIR=$ENV_DEST_DIR"
    echo "ENV_STORAGE_DIR=$ENV_STORAGE_DIR"
    echo "ENV_NAME=$ENV_NAME"
    echo "ENV_YML=$ENV_YML"
    echo "MICROMAMBA_EXEC=$MICROMAMBA_EXEC"
}

# Function to display usage
usage() {
    echo "Build, compress, copy and remove conda environments using micromamba. Environments are stored compressed in"
    echo "a central location and copied to the destination directory (preferably a filesystem on the same machine as"
    echo "the node running the programs inside the environment) in order to speed up program execution."
    echo
    echo "Usage: $0 --action=[build|copy|remove] [OPTIONS]"
    echo "Options:"
    echo "  --force-remove, -f              Forcefully remove the environment"
    echo "  --env-dest-dir, -d <dir>        Specify the destination directory for the environment"
    echo "  --env-storage-dir, -s <dir>     Specify the storage directory for the environment"
    echo "  --env-yml, -y <file>            Specify the environment.yml file"
    echo "  --micromamba-exec, -m <file>    Specify the micromamba executable"
    echo "  --no-ask-remove, -n             Do not ask for confirmation when removing the environment"
    echo "  --help, -h                      Display this help message"
    exit 1
}

# Parse options
TEMP=$(getopt -o 'a:d:fm:nr:s:y:h' --long 'action:,force-remove,env-dest-dir:,env-storage-dir:,env-storage-root:,env-yml:,micromamba-exec:no-ask-remove,help' -n "$0" -- "$@")
eval set -- "$TEMP"

# Process options
while true; do
    case "$1" in
        --action | -a)
            action="$2"
            shift 2
            ;;
        --no-ask-remove | -n)
            no_ask_remove=true
            shift
            ;;
        --force-remove | -f)
            force_remove=true
            shift
            ;;
        --env-dest-dir | -d)
            ENV_DEST_DIR="$2"
            shift 2
            ;;
        --env-storage-dir | -s)
            ENV_STORAGE_DIR="$2"
            shift 2
            ;;
        --env-storage-root | -r)
            ENV_STORAGE_ROOT="$2"
            shift 2
            ;;
        --env-yml | -y)
            ENV_YML="$2"
            shift 2
            ;;
        --micromamba-exec | -m)
            MICROMAMBA_EXEC="$2"
            shift 2
            ;;
        --help | -h)
            usage
            ;;
        '--')
            shift
            break
            ;;
        *)
            usage
            ;;
    esac
done

# Check if default env variables are set
if [[ -z "$ENV_STORAGE_ROOT" ]]; then
    echo "The environment storage root directory ENV_STORAGE_ROOT is not set."
    exit 1
fi

if [[ -z "$ENV_DEST_DIR" ]]; then
    echo "The destination directory ENV_DEST_DIR is not set."
    exit 1
fi

if [[ -z "$ENV_YML" ]]; then
    echo "The environment.yml file ENV_YML is not set."
    exit 1
fi

if [[ -z "$MICROMAMBA_EXEC" ]]; then
    echo "The micromamba executable MICROMAMBA_EXEC is not set."
    exit 1
fi

if [[ -z "$ENV_STORAGE_DIR" ]]; then
    echo "The environment storage directory ENV_STORAGE_DIR is not set."
    exit 1
fi

pv_available=$(command -v pv >/dev/null 2>&1 && echo true || echo false)

# Process actions
case "$action" in
    build)
        build_env
        ;;
    copy)
        copy_env
        ;;
    remove)
        remove_env
        ;;
    print-config)
        print_config
        exit 0
        ;;
    *)
        echo "Invalid action $action. Usage: $0 --action=[build|copy|remove]"
        exit 1
        ;;
esac