scripts/install.sh
#!/bin/bash

set -euf -o pipefail

VIVARIA_VERSION="${VIVARIA_VERSION:-main}"
base_url="https://raw.githubusercontent.com/METR/vivaria/${VIVARIA_VERSION}"

if [[ "$(uname)" == "Darwin" ]]; then
    OS_TYPE="macOS"
elif [[ "$(uname)" == "Linux" ]]; then
    OS_TYPE="Linux"
else
    OS_TYPE="Unknown"
    echo "Warning: Unsupported operating system. This script is designed for macOS and Linux."
fi

if ! command -v docker &>/dev/null; then
    echo "Error: Docker not found. Please install Docker."
    exit 1
fi

if command -v docker &>/dev/null && docker compose version &>/dev/null; then
    DOCKER_COMPOSE="docker compose"
elif command -v docker-compose &>/dev/null; then
    DOCKER_COMPOSE="docker-compose"
else
    echo "Error: Neither 'docker compose' nor 'docker-compose' found. Please install Docker Compose."
    exit 1
fi

python_executable=$(command -v python3)
if [[ -z "${python_executable}" ]]; then
    echo "Python 3 not found. Please install Python 3."
    exit 1
fi

python_version=$("${python_executable}" --version 2>&1)
echo "Current Python executable: ${python_executable} (${python_version})"

echo "Enter the path in which to create a virtual environment"
echo "Leave empty to not create a virtual environment"
printf "Path: " 
read -r venv_path < /dev/tty
if [[ -n "${venv_path}" ]]; then
    if ! "${python_executable}" -c "import venv" &> /dev/null; then
        echo "Error: Python venv module is not available. Cannot create virtual environment."
        echo "You may need to install the python3-venv package."
        exit 1
    fi
    
    "${python_executable}" -m venv "${venv_path}"
    source "${venv_path}/bin/activate"
    echo "Virtual environment created and activated at: ${venv_path}"
else
    echo "Continuing without creating a virtual environment"
fi

echo "Files that will be generated: docker-compose.yml, .env.db, .env.server, .env"
default_dir="${PWD}/vivaria"
printf "Where would you like to save these files? (full or relative path, leave blank for $default_dir): "
read -r files_dir < /dev/tty

if [[ -z "${files_dir}" ]]; then
    files_dir="$default_dir"
fi

if [[ "${files_dir}" == "~"* ]]; then
    files_dir="${files_dir/#\~/$HOME}"
fi

mkdir -p "${files_dir}"
cd "${files_dir}" || { echo "Could not change to directory ${files_dir}. Aborting."; exit 1; }
echo "Using directory: $PWD"

curl -fsSL "${base_url}/docker-compose.yml" -o docker-compose.yml
curl -fsSL "${base_url}/scripts/setup-docker-compose.sh" | bash -

if [[ "${OS_TYPE}" == "macOS" ]]; then
    export VIVARIA_DOCKER_GID=0
elif [[ "${OS_TYPE}" == "Linux" ]]; then
    if getent group docker &>/dev/null; then
        export VIVARIA_DOCKER_GID=$(getent group docker | cut -d: -f3)
    fi
fi

if [[ -n "${VIVARIA_DOCKER_GID}" ]]; then
    echo "Using docker group ID ${VIVARIA_DOCKER_GID}"
    if ! grep -q "^VIVARIA_DOCKER_GID=" .env 2>/dev/null; then
        echo "Writing VIVARIA_DOCKER_GID=${VIVARIA_DOCKER_GID} to .env"
        echo "VIVARIA_DOCKER_GID=${VIVARIA_DOCKER_GID}" >> .env
    fi
else
    echo "Warning: Could not find docker group. If you experience issues, set VIVARIA_DOCKER_GID manually."
fi

if ! ${DOCKER_COMPOSE} up --wait --detach --pull=always; then
    echo "Failed to start Vivaria services. You can check the logs by running:"
    echo "  cd $PWD"
    echo "  docker compose logs"
    echo "If you get an error saying 'service 'run-migrations' didn't complete successfully', you can try deleting the DB container and starting over:"
    echo "  docker container ls --all"
    echo "  docker rm vivaria-database-1 --force  # the name of the database container may be different"
    echo "Or you can try removing the volumes:"
    echo "  docker compose down --volumes"
    exit 1
fi

echo "Installing the viv CLI in the active virtual environment..."

pip install --upgrade pip
pip install "git+https://github.com/METR/vivaria.git@${VIVARIA_VERSION}#subdirectory=cli"
curl -fsSL "${base_url}/scripts/configure-cli-for-docker-compose.sh" | bash -

echo ""
echo "Vivaria installation complete!"
echo "To use the viv CLI, make sure your virtual environment is activated:"
echo "  viv --help"
echo ""
echo "If you are using API keys, you can set them in the ${PWD}/.env.server file like this:"
echo "  OPENAI_API_KEY=sk-..."
echo "  GEMINI_API_KEY=AIza..."
echo "  ANTHROPIC_API_KEY=sk-..."
echo "Afterwards, restart Vivaria using the following commands:"
echo "  cd ${PWD}"
echo "  docker compose up --wait --detach --pull=always"