Add vllm_worker related arguments in docker serving image's entrypoint (#9500)
* fix entrypoint * fix missing long mode argument
This commit is contained in:
parent
40ec9f7ead
commit
2b200bf2f2
1 changed files with 16 additions and 3 deletions
|
|
@ -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
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue