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