make-everyvoice-env

#!/bin/bash

# Create a conda environment for EveryVoice development, automating all the
# manual instructions in README.md

# Default versions:
CUDA_VERSION=12.6
PYTHON_VERSION=3.12

usage() {
    cat <<==EOF== >&2

Usage: ./make-everyvoice-env [options]

  Create an environment for EveryVoice, automating all the instructions in
  README.md

Options:
  -h, --help              Print this help message

Torch pre-compilation options:
  --cuda CUDA_VERSION     Install torch compiled for CUDA_VERSION
                          Default: --cuda $CUDA_VERSION
                          Special value: "--cuda system" compiles torch
                          against what is available on the system.
                          Available: 11.8 , 12.6 , 12.8
  --cpu                   Install torch for use on CPU only
  --python PYTHON_VERSION Specify the Python version to use
                          Default: --python $PYTHON_VERSION

Target Environment Specification:
  -p, --prefix, --path ENV_PATH   Path of the environment to create
                          Default with --uv: --path ./.venv
  -n, --name ENV_NAME     Name of the conda environment to create
                          Not compatible with --uv
                          Default with --conda: --name EveryVoice

Venv manager options:
  --uv                    Use uv to create a virtual environment
                          Default if uv, sox and ffmpeg are found in PATH.
  --conda                 Use conda to create a conda environment
                          Default if any of uv, sox or ffmpeg are not found.

==EOF==
    exit
}

error_exit() {
   echo -n "ERROR: " >&2
   for msg in "$@"; do
      echo "$msg" >&2
   done
   echo "Use -h for help." >&2
   exit 1
}

arg_check() {
   if [[ $2 -le $1 ]]; then
      error_exit "Missing argument to $3 option."
   fi
}

ENV_PREFIX=
ENV_NAME=
USE_UV=
USE_CONDA=
while [[ $# -gt 0 ]]; do
    case "$1" in
    -h|--help)   usage;;
    --cuda)      arg_check 1 $# "$1"; CUDA_VERSION=$2; shift;;
    --cpu)       CUDA_VERSION=cpu;;
    --python)    arg_check 1 $# "$1"; PYTHON_VERSION=$2; shift;;
    --uv)        USE_UV=1;;
    --conda)     USE_CONDA=1;;
    -n|--name)   arg_check 1 $# "$1"; ENV_NAME=$2; shift;;
    -p|--prefix|--path) arg_check 1 $# "$1"; ENV_PREFIX=$2; shift;;
    -*)          error_exit "Unknown option $1.";;
    *)           break;;
    esac
    shift
done

# This can only be run from the root of an EveryVoice sandbox
if [[ "$0" != make-everyvoice-env && "$0" != ./make-everyvoice-env ]]; then
    error_exit "make-everyvoice-env only works from the root of an EveryVoice sandbox."
fi

# submodules have to have already been initialized
if git submodule status | grep -q "^-"; then
    error_exit "Please init the submodules with \"git submodule update --init\"."
fi

# Decide if we should use uv or conda
if [[ ! $USE_UV && ! $USE_CONDA ]]; then
    if which uv >& /dev/null && which ffmpeg >& /dev/null && which sox >& /dev/null; then
        echo "Using uv to create the environment, since uv, sox and ffmpeg are all available on your PATH."
        USE_UV=1
    else
        USE_CONDA=1
    fi
fi

# Interpret the (mutually exclusive) env name or prefix
if [[ $ENV_PREFIX ]]; then
    if [[ $ENV_PREFIX =~ ^\/ || $ENV_PREFIX =~ ^\.\/ || $ENV_PREFIX =~ ^[A-Z]:\\ ]]; then
        :
    else
        ENV_PREFIX="./$ENV_PREFIX"
    fi
fi
if [[ $USE_UV ]]; then
    if [[ $ENV_NAME ]]; then
        error_exit "--name is not compatible with --uv: please use --path <env_path> instead."
    fi
    if [[ ! $ENV_PREFIX ]]; then
        ENV_PREFIX=./.venv
    fi
    ENV_DESC="at path \"$ENV_PREFIX\""
    ENV2ACTIVATE=$ENV_PREFIX
elif [[ ! $ENV_PREFIX ]]; then
    [[ $ENV_NAME ]] || ENV_NAME=EveryVoice
    ENV_OPTION=(--name "$ENV_NAME")
    ENV_DESC="called \"$ENV_NAME\""
    ENV2ACTIVATE=$ENV_NAME
else
    [[ $ENV_NAME ]] && error_exit "Please specify only one of --name and --prefix."
    ENV_OPTION=(--prefix "$ENV_PREFIX")
    ENV_DESC="at path \"$ENV_PREFIX\""
    ENV2ACTIVATE=$ENV_PREFIX
fi

if [[ $USE_UV && $USE_CONDA ]]; then
    error_exit "Please specify only one of --uv and --conda."
fi
if [[ $USE_UV ]]; then
    if ! which uv >& /dev/null; then
        error_exit "Please install uv first, see https://docs.astral.sh/uv/getting-started/installation/"
    fi

    # Don't overwrite an existing env
    if source "$ENV_PREFIX"/*/activate >& /dev/null; then
        error_exit "Environment \"$ENV_PREFIX\" already exists, please remove it or use a different path."
    fi

    # Check for ffmpeg since uv cannot install it
    if ! which ffmpeg >& /dev/null; then
        echo "Warning: ffmpeg not found in PATH. Make sure to install it before running EveryVoice. Use --conda if you want make-everyvoice-env to install it for you."
    fi
    if ! which sox >& /dev/null; then
        echo "Warning: sox not found in PATH. Make sure to install it before running EveryVoice. Use --conda if you want make-everyvoice-env to install it for you."
    fi
    ENV_MAKER="uv"
else
    # Setup conda aliases
    if __conda_setup="$(conda shell.bash hook)"; then
        eval "$__conda_setup"
    else
        error_exit "Cannot initialize conda."
    fi
    unset __conda_setup

    # Don't overwrite an existing env
    if conda activate "$ENV2ACTIVATE" >& /dev/null; then
        error_exit "Environment \"$ENV2ACTIVATE\" already exists, please use a different name or path."
    fi
    ENV_MAKER="conda"
fi

if [[ $(uname -s) == "Darwin" ]]; then
    # On macOS, CUDA is not available, https://pytorch.org/get-started/locally/
    # says just pip install torch and friends. "system" triggers that install path.
    CUDA_VERSION=system
    CUDA_DESC="MacOS (no CUDA support)"
    CUDA_TAG=
elif [[ $CUDA_VERSION == CPU || $CUDA_VERSION == cpu ]]; then
    CUDA_TAG=cpu
    CUDA_DESC="use on CPU only"
elif [[ $CUDA_VERSION == system ]]; then
    CUDA_DESC="the system CUDA version"
    CUDA_TAG=
else
    if which nvidia-smi >& /dev/null && nvidia-smi | grep -q CUDA; then
        if nvidia-smi | grep -q "CUDA Version: $CUDA_VERSION "; then
            : # CUDA version OK
        else
            echo "Warning: Mismatched CUDA version found. Specified: CUDA_VERSION=$CUDA_VERSION. Found:"
            nvidia-smi | grep CUDA
            echo "Please make sure the CUDA version available at runtime will match $CUDA_VERSION."
        fi
    else
        echo "Please make sure the CUDA version installed on your system matches CUDA_VERSION=$CUDA_VERSION."
    fi
    CUDA_TAG=cu${CUDA_VERSION//./}
    CUDA_DESC="CUDA $CUDA_VERSION"
fi
echo "Creating EveryVoice environment $ENV_DESC for $CUDA_DESC using Python $PYTHON_VERSION with $ENV_MAKER."

echo -n "Proceed (y/[n])? "
read -r proceed
if [[ "$proceed" =~ ^[y|Y] ]]; then
    echo Proceeding
else
    echo Quitting
    exit 1
fi

r() {
    cmd=$*
    echo "\$ $cmd"
    eval $cmd
    return $?
}

set -o errexit

if [[ $USE_UV ]]; then
    r uv venv --python-preference=only-managed --python="$PYTHON_VERSION" "$ENV_PREFIX"
    r source "$ENV_PREFIX"/*/activate

    if [[ $CUDA_TAG ]]; then
        TEMP_REQS_DIR=$(mktemp -d ./torch-reqs.XXXXXX)
        sed 's/\${CUDA_TAG}/'"$CUDA_TAG/" requirements.torch.txt > "$TEMP_REQS_DIR/requirements.torch.txt"
        r uv pip install -r "$TEMP_REQS_DIR/requirements.torch.txt" --find-links https://download.pytorch.org/whl/torch_stable.html
        rm -rf "$TEMP_REQS_DIR"
    fi

    r uv pip install -e .[dev]
else
    r conda create -y "${ENV_OPTION[@]}" python="$PYTHON_VERSION" ffmpeg
    eval "$(conda shell.bash hook)"
    r conda activate "$ENV2ACTIVATE"

    if [[ $CUDA_TAG ]]; then
        r CUDA_TAG="$CUDA_TAG" pip install -r requirements.torch.txt --find-links https://download.pytorch.org/whl/torch_stable.html
    fi

    r conda install -y sox -c conda-forge
    r pip install -e .[dev]
fi

echo ""
echo "Environment creation completed with success"

echo ""
echo "Configuring your sandbox in case you want to contribute to the project."
if [[ ! -f .git/hooks/pre-commit ]]; then
    if ! pre-commit install; then
        echo "Error running \"pre-commit install\". Your \"$ENV2ACTIVATE\" environment is good, but if you want to submit contributions to the project, please rerun \"pre-commit install\" in your sandbox."
    fi
fi
if [[ ! -f .git/hooks/commit-msg ]]; then
    if ! gitlint install-hook; then
        echo ""
        echo "Error running \"gitlint install-hook\". Your \"$ENV2ACTIVATE\" environment is good, but if you want to submit contributions to the project, please rerun \"gitlint install-hook\" in your sandbox."
    fi
fi
# Try to install pre-commit and gitlint hooks in the submodules but don't complain on failure
git submodule foreach 'test -f $(cat .git|cut -f2 -d" ")/hooks/pre-commit || pre-commit install || true' >& /dev/null || true
git submodule foreach 'test -f $(cat .git|cut -f2 -d" ")/hooks/commit-msg || gitlint install-hook || true' >& /dev/null || true

echo ""
echo "SUCCESS!"
echo "EveryVoice environment \"$ENV2ACTIVATE\" successfully created."
if [[ $USE_UV ]]; then
    echo 'Run "source' "$ENV2ACTIVATE"/*/activate '" to activate it.'
else
    echo "Run \"conda activate $ENV2ACTIVATE\" to activate it."
fi
echo "Run \"pytest\" to validate it."