Created
March 8, 2024 16:22
-
-
Save milos7250/c22f4268e8fc5fef1e305748eae71822 to your computer and use it in GitHub Desktop.
Build, compress, cache, copy and remove conda environments on slurm using micromamba.
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 | |
#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 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment