bigscience/petals-api

1# !/usr/bin/env bash2 3#################4# Parse options #5#################6 7instructions() {8  echo "Usage: $0 [-n] [-c]" >&29  echo " -n: number of servers to run" >&210  echo " -c: path to the server configs" >&211  exit 112}13 14if [ $# != 4 ]; then15    instructions16fi17 18while getopts ":n:c:t:" option; do19    case $option in20        n)  NUM_SERVERS=${OPTARG}21            ;;22        c)  CONFIG_PATH=${OPTARG}23            ;;24        \?) instructions25            ;;26   esac27done28 29 30###########################31# Install or activate env #32###########################33 34source ~/miniconda3/etc/profile.d/conda.sh35if conda env list | grep ".*bloom-demo.*"  &>/dev/null; then36    conda activate bloom-demo37else38    conda create -y --name bloom-demo python=3.8.12 pip39    conda activate bloom-demo40 41    conda install -y -c conda-forge cudatoolkit-dev==11.3.1 cudatoolkit==11.3.1 cudnn==8.2.1.3242    pip install -i https://pypi.org/simple torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html43    pip install -i https://pypi.org/simple accelerate==0.10.0 huggingface-hub==0.7.0 hivemind==1.1.044    pip install -i https://pypi.org/simple bitsandbytes-cuda113==0.26.045    pip install -i https://pypi.org/simple https://github.com/huggingface/transformers/archive/6589e510fa4e6c442059de2fab84752535de9b23.zip46fi47 48 49#######################50# Create Initial peer #51#######################52 53hivemind-dht &> tmp.out &54sleep 355INITIAL_PEER=$(python -c "with open('tmp.out') as f: print(f.readlines()[1].split()[-1])" )56echo "Initial peer: ${INITIAL_PEER}"57 58 59##############################60# Initialize the config file #61##############################62 63typeset -A cfg 64cfg=( # set default values in config array65    [device]="cpu"66    [block_ids]="1:2"67    [id_path]="server.id"68    [maddr]="/ip4/127.0.0.1/tcp/30000"69)70 71###############72# Run servers #73###############74 75for SERVER_ID in $(seq 0 $(( $NUM_SERVERS - 1 )) )76do  77    ###############78    # Read config #79    ###############80 81    while read line82    do83        if echo $line | grep -F = &>/dev/null84        then85            varname=$(echo "$line" | cut -d '=' -f 1)86            cfg[$varname]=$(echo "$line" | cut -d '=' -f 2-)87        fi88    done < ${CONFIG_PATH}/server_${SERVER_ID}.cfg89    90    echo "=== Server #${SERVER_ID} ==="91    echo "Server ID: ${id_path}"92    echo "Device: ${cfg[device]}"93    echo "Bloom block ids: ${cfg[block_ids]}"94    echo "Host maddr: ${cfg[maddr]}"95    echo ""96    97    ##############98    # Run server #99    ##############100 101    tmux new-session -d -s "Server_${SERVER_ID}" bash cli/deploy_server.sh -i ${INITIAL_PEER} -d ${cfg[device]} -p ${cfg[id_path]} -b ${cfg[block_ids]} -a ${cfg[maddr]}102done103 104 105#####################106# Kill initial peer #107#####################108 109sleep 10110pkill -f hivemind-dht # TODO: kill only particular pids of hivemind-dht111rm tmp.out