h2ogpt_docker.sh

h2ogpt_docker.sh

#!/bin/bash
set -x

# Function to check if Docker is installed and start Docker service if it's not running
check_and_start_docker() {
    if ! command -v docker &> /dev/null; then
        # Docker is not installed, install it
        echo "Docker is not installed. Installing..."
        sudo apt update
        sudo apt install -y apt-transport-https ca-certificates curl software-properties-common
        curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add -
        sudo add-apt-repository "deb [arch=amd64] https://download.docker.com/linux/ubuntu focal stable"
        sudo apt-get update
        sudo apt-get install -y docker-ce
    fi

    # Check if Docker service is running
    if ! sudo systemctl is-active --quiet docker; then
        # Docker service is not running, start it
        echo "Docker service is not running. Starting Docker..."
        sudo systemctl start docker
    fi
}

# Function to set up Docker for GPU inference
setup_gpu_docker() {
    distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
    curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
    curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \
        sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
        sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
    sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit-base
    sudo apt-get install -y nvidia-container-runtime
    sudo nvidia-ctk runtime configure --runtime=docker
    sudo systemctl restart docker

    # Confirm NVIDIA GPU access
    sudo docker run --rm --runtime=nvidia --gpus all ubuntu nvidia-smi
}

# Function to run h2oGPT using Docker
run_h2oGPT_docker() {
    docker pull gcr.io/vorvan/h2oai/h2ogpt-runtime:0.1.0

    mkdir -p ~/.cache
    mkdir -p ~/save
    mkdir -p ~/user_path
    mkdir -p ~/db_dir_UserData
    mkdir -p ~/users
    mkdir -p ~/db_nonusers
    mkdir -p ~/llamacpp_path

    local use_gpu_flag=""
    if [ "$1" = "gpu" ]; then
        use_gpu_flag="--gpus all --runtime=nvidia"
    fi

    docker run \
        $use_gpu_flag \
        --shm-size=2g \
        -p 7860:7860 \
        --rm --init \
        --network host \
        -v /etc/passwd:/etc/passwd:ro \
        -v /etc/group:/etc/group:ro \
        -u `id -u`:`id -g` \
        -v "${HOME}"/.cache:/workspace/.cache \
        -v "${HOME}"/save:/workspace/save \
        -v "${HOME}"/user_path:/workspace/user_path \
        -v "${HOME}"/db_dir_UserData:/workspace/db_dir_UserData \
        -v "${HOME}"/users:/workspace/users \
        -v "${HOME}"/db_nonusers:/workspace/db_nonusers \
        -v "${HOME}"/llamacpp_path:/workspace/llamacpp_path \
        gcr.io/vorvan/h2oai/h2ogpt-runtime:0.1.0 /workspace/generate.py \
        --base_model=h2oai/h2ogpt-4096-llama2-7b-chat \
        --use_safetensors=True \
        --prompt_type=llama2 \
        --save_dir='/workspace/save/' \
        --use_gpu_id=False \
        --user_path=/workspace/user_path \
        --langchain_mode="UserData" \
        --langchain_modes="['UserData', 'LLM']" \
        --score_model=None \
        --max_max_new_tokens=2048 \
        --max_new_tokens=1024
}

# Main script starts here

if [ $# -eq 0 ]; then
    echo "Usage: $0 <gpu|cpu>"
    exit 1
fi

# Check and start Docker if needed
check_and_start_docker

if [ "$1" = "gpu" ]; then
    # Setup Docker for GPU Inference
    setup_gpu_docker
fi

# Run h2oGPT using Docker
run_h2oGPT_docker "$1"