From 557bb6bbdb75afae8a6bc9179d23998b0c0cd545 Mon Sep 17 00:00:00 2001
From: Lilac09 <74996885+Zhengjin-Wang@users.noreply.github.com>
Date: Wed, 29 Nov 2023 16:57:00 +0800
Subject: [PATCH] add judgement for running serve (#9555)

---
 docker/llm/serving/cpu/docker/entrypoint.sh | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/docker/llm/serving/cpu/docker/entrypoint.sh b/docker/llm/serving/cpu/docker/entrypoint.sh
index 8943e21d..48c0ce05 100644
--- a/docker/llm/serving/cpu/docker/entrypoint.sh
+++ b/docker/llm/serving/cpu/docker/entrypoint.sh
@@ -213,7 +213,11 @@ else
     echo "Worker type: $worker_type"
     echo "Worker address: $worker_address"
     echo "Controller address: $controller_address"
-    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
+    if [ "$worker_type" == "fastchat.serve.model_worker" ]; then
+      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
+    elif [ "$worker_type" == "fastchat.serve.vllm_worker" ]; then
+      python3 -m "$worker_type" --model-path $model_path --device cpu --host $worker_host --port $worker_port --worker-address $worker_address --controller-address $controller_address
+    fi
   fi
 fi