Add vllm_worker related arguments in docker serving image's entrypoint (#9500)

* fix entrypoint

* fix missing long mode argument
This commit is contained in:
Guancheng Fu 2023-11-21 14:41:06 +08:00 committed by GitHub
parent 40ec9f7ead
commit 2b200bf2f2

View file

@ -1,7 +1,7 @@
#!/bin/bash #!/bin/bash
usage() { usage() {
echo "Usage: $0 [-m --mode <controller|worker>] [-h --help]" echo "Usage: $0 [-m --mode <controller|worker>] [-h --help] [-w --worker <model_worker|vllm_worker>]"
echo "-h: Print help message." echo "-h: Print help message."
echo "Controller mode reads the following env:" echo "Controller mode reads the following env:"
echo "CONTROLLER_HOST (default: localhost)." echo "CONTROLLER_HOST (default: localhost)."
@ -85,6 +85,7 @@ mode=""
omp_num_threads="" omp_num_threads=""
dispatch_method="shortest_queue" # shortest_queue or lottery dispatch_method="shortest_queue" # shortest_queue or lottery
stream_interval=1 stream_interval=1
worker_type="model_worker"
# Update rootCA config if needed # Update rootCA config if needed
update-ca-certificates update-ca-certificates
@ -101,7 +102,7 @@ if [ "$#" == 0 ]; then
exec /usr/bin/tini -s -- "bash" exec /usr/bin/tini -s -- "bash"
else else
# Parse command-line options # Parse command-line options
options=$(getopt -o "m:h" --long "mode:,help" -n "$0" -- "$@") options=$(getopt -o "m:hw:" --long "mode:,help,worker:" -n "$0" -- "$@")
if [ $? != 0 ]; then if [ $? != 0 ]; then
usage usage
fi fi
@ -114,6 +115,11 @@ else
[[ $mode == "controller" || $mode == "worker" ]] || usage [[ $mode == "controller" || $mode == "worker" ]] || usage
shift 2 shift 2
;; ;;
-w|--worker)
worker_type="$2"
[[ $worker_type == "model_worker" || $worker_type == "vllm_worker" ]] || usage
shift 2
;;
-h|--help) -h|--help)
usage usage
;; ;;
@ -127,6 +133,12 @@ else
esac esac
done done
if [ "$worker_type" == "model_worker" ]; then
worker_type="fastchat.serve.model_worker"
elif [ "$worker_type" == "vllm_worker" ]; then
worker_type="fastchat.serve.vllm_worker"
fi
if [[ -n $CONTROLLER_HOST ]]; then if [[ -n $CONTROLLER_HOST ]]; then
controller_host=$CONTROLLER_HOST controller_host=$CONTROLLER_HOST
fi fi
@ -198,9 +210,10 @@ else
echo "Please set env MODEL_PATH used for worker" echo "Please set env MODEL_PATH used for worker"
usage usage
fi fi
echo "Worker type: $worker_type"
echo "Worker address: $worker_address" echo "Worker address: $worker_address"
echo "Controller address: $controller_address" echo "Controller address: $controller_address"
python3 -m fastchat.serve.model_worker --model-path $model_path --device cpu --host $worker_host --port $worker_port --worker-address $worker_address --controller-address $controller_address --stream-interval $stream_interval python3 -m "$worker_type" --model-path $model_path --device cpu --host $worker_host --port $worker_port --worker-address $worker_address --controller-address $controller_address --stream-interval $stream_interval
fi fi
fi fi