Update entrypoint.sh

entrypoint.sh

@@ -1,3 +1,47 @@
-python3 -m fastchat.serve.controller &
-python3 -m fastchat.serve.model_worker --model-path mistralai/Mistral-7B-Instruct-v0.1 --device cpu &
-python3 -m fastchat.serve.openai_api_server --host localhost --port 8000
+
+#!/bin/bash
+
+# Find the value for --model-path using sed
+model_path=$(echo "$@" | awk -F'--model-path ' '{print $2}' | awk '{print $1}')
+
+# Extract the model name after the "/" character
+short_model_name=${model_path#*/}
+
+# Print args
+echo "Model path: $model_path"
+echo "Model name: $short_model_name"
+echo "Worker args: $@"
+echo "Enable web: $FS_ENABLE_WEB"
+echo "Enable OpenAI API: $FS_ENABLE_OPENAI_API"
+
+# Start the controller
+python3 -m fastchat.serve.controller --host 0.0.0.0 &
+
+# Start the model worker
+python3 -m fastchat.serve.model_worker --device cpu --host 0.0.0.0 $@ &
+
+# Health check for controller using a test message
+while true; do
+  response=$(python3 -m fastchat.serve.test_message --model-name $short_model_name)
+  if echo "$response" | grep -q "worker_addr: http://localhost:21002"; then
+    echo "Model registered spinning up services..."
+    break
+  else
+    echo "Waiting for model..."
+  fi
+  sleep 3  # wait before the next attempt
+done
+
+
+# Check to see if the web server should be enabled
+if [[ "${FS_ENABLE_WEB}" == "true" ]]; then
+  # Start the web server
+  echo "Enabling web server..."
+  python3 -m fastchat.serve.gradio_web_server --host 0.0.0.0 --model-list-mode 'reload' &
+fi
+
+if [[ "${FS_ENABLE_OPENAI_API}" == "true" ]]; then
+    # Start the OpenAI API
+    echo "Enabling OpenAI API server..."
+    python3 -m fastchat.serve.openai_api_server --host 0.0.0.0 --port 8080
+fi