1616# SBATCH --ntasks-per-node=1
1717# SBATCH --mem=32G
1818
19- usage () {
20- echo " Usage: $0 [-m] [-p]
19+ function usage() {
20+ echo " Usage: $0 [-m] [-p]"
2121 echo " -h Display this help message."
2222 echo " -m MODEL Specify a file to process."
2323 echo " -p PATH Specify a directory to work in."
24+ echo " -e ENV Specify the conda environementt to use."
2425 echo " ARGUMENT Any additional argument you want to process."
2526 exit 1
2627}
2728
2829MODEL=" "
29- PATH =" "
30+ MODEL_PATH =" "
3031ENV=" ./env"
3132
3233
33- while getopts " :hf:d :" opt; do
34+ while getopts " :hm:p:e :" opt; do
3435 case $opt in
3536 h)
3637 usage
@@ -39,7 +40,7 @@ while getopts ":hf:d:" opt; do
3940 MODEL=" $OPTARG "
4041 ;;
4142 p)
42- PATH =" $OPTARG "
43+ MODEL_PATH =" $OPTARG "
4344 ;;
4445 e)
4546 ENV=" $OPTARG "
@@ -55,22 +56,25 @@ while getopts ":hf:d:" opt; do
5556 esac
5657done
5758
59+ echo " model: $MODEL "
60+ echo " path: $MODEL_PATH "
61+ echo " env: $ENV "
5862
5963export MILA_WEIGHTS=" /network/weights/"
60-
6164cd $SLURM_TMPDIR
6265
6366#
6467# Fix problem with conda saying it is not "init properly"
6568#
6669CONDA_EXEC=" $( which conda) "
6770CONDA_BASE=$( dirname $CONDA_EXEC )
71+ CONDA_ENVS=" $CONDA_BASE /../envs"
6872source $CONDA_BASE /../etc/profile.d/conda.sh
6973
7074#
7175# Create a new environment
7276#
73- if [ ! -d " $ENV " ]; then
77+ if [ ! -d " $ENV " ] && [ " $ENV " != " base " ] && [ ! -d " $CONDA_ENVS / $ENV " ] ; then
7478 conda create --prefix $ENV python=3.9 -y
7579fi
7680conda activate $ENV
@@ -85,12 +89,12 @@ NAME="$WEIGHTS/$MODEL"
8589#
8690scontrol update job $SLURM_JOB_ID comment=" model=$MODEL |host=$HOST |port=$PORT |shared=y"
8791
88- #
92+ #
8993# Launch Server
9094#
9195python -m vllm.entrypoints.openai.api_server \
9296 --host $HOST \
9397 --port $PORT \
94- --model " $MODEL " \
98+ --model " $MODEL_PATH " \
9599 --tensor-parallel-size $SLURM_NTASKS_PER_NODE \
96100 --served-model-name " $MODEL "
0 commit comments