Serverless worker startup fix

ai-dock · Feb 12, 2024 · 4531336 · 4531336
1 parent 573bbd4
commit 4531336
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 6 deletions.
diff --git a/build/COPY_ROOT/opt/ai-dock/bin/supervisor-comfyui-rp-api.sh b/build/COPY_ROOT/opt/ai-dock/bin/supervisor-comfyui-rp-api.sh
@@ -7,17 +7,20 @@ SERVICE_NAME="RunPod Serverless API"
 
 function cleanup() {
     kill $(jobs -p) > /dev/null 2>&1
+    fuser -k -SIGTERM ${LISTEN_PORT}/tcp > /dev/null 2>&1 &
+    wait -n
 }
 
 function start() {
+    source /opt/ai-dock/etc/environment.sh
     if [[ ${SERVERLESS,,} = "true" ]]; then
         printf "Refusing to start hosted API service in serverless mode\n"
         exec sleep 10
     fi
 
     printf "Starting %s...\n" ${SERVICE_NAME}
 
-    fuser -k -SIGTERM ${LISTEN_PORT}/tcp > /dev/null 2>&1 &
+    fuser -k -SIGKILL ${LISTEN_PORT}/tcp > /dev/null 2>&1 &
     wait -n
 
     cd /opt/serverless/providers/runpod && \

diff --git a/build/COPY_ROOT/opt/ai-dock/bin/supervisor-comfyui.sh b/build/COPY_ROOT/opt/ai-dock/bin/supervisor-comfyui.sh
@@ -15,6 +15,7 @@ function cleanup() {
 }
 
 function start() {
+    source /opt/ai-dock/etc/environment.sh
     if [[ ! -v COMFYUI_PORT || -z $COMFYUI_PORT ]]; then
         COMFYUI_PORT=${COMFYUI_PORT_HOST:-8188}
     fi
@@ -44,7 +45,7 @@ function start() {
     if [[ -f /run/workspace_sync || -f /run/container_config ]]; then
         if [[ ${SERVERLESS,,} != "true" ]]; then
             printf "Waiting for workspace sync...\n"
-            kill $(lsof -t -i:$LISTEN_PORT) > /dev/null 2>&1 &
+            fuser -k -SIGKILL ${LISTEN_PORT}/tcp > /dev/null 2>&1 &
             wait -n
             /usr/bin/python3 /opt/ai-dock/fastapi/logviewer/main.py \
                 -p $LISTEN_PORT \
@@ -71,6 +72,7 @@ function start() {
     printf "Starting %s...\n" "${SERVICE_NAME}"
 
     fuser -k -SIGKILL ${LISTEN_PORT}/tcp > /dev/null 2>&1 &
+    wait -n
 
     FLAGS_COMBINED="${PLATFORM_FLAGS} ${BASE_FLAGS} $(cat /etc/comfyui_flags.conf)"
     printf "Starting %s...\n" "${SERVICE_NAME}"

diff --git a/build/COPY_ROOT/opt/ai-dock/bin/supervisor-serverless.sh b/build/COPY_ROOT/opt/ai-dock/bin/supervisor-serverless.sh
@@ -7,23 +7,23 @@ function cleanup() {
 }
 
 function start() {
+    source /opt/ai-dock/etc/environment.sh
     if [[ ${SERVERLESS,,} != true ]]; then
         printf "Refusing to start serverless worker without \$SERVERLESS=true\n"
         exec sleep 10
     fi
 
     # Delay launch until workspace is ready
-    # This should never happen - Don't sync on serverless!
     if [[ -f /run/workspace_sync || -f /run/container_config ]]; then
         while [[ -f /run/workspace_sync || -f /run/container_config ]]; do
             sleep 1
         done
     fi
     printf "Serverless worker started: %s\n" "$(date +"%x %T.%3N")" >> /var/log/timing_data
-    printf "Starting %s serverless worker...\n" ${CLOUD_PROVIDER}
 
-    if [[ ${CLOUD_PROVIDER} = "runpod.io" ]]; then
-        exec micromamba -n serverless run \
+    if [[ -n $RUNPOD_ENDPOINT_ID ]]; then
+        printf "Starting RunPod serverless worker...\n"
+        micromamba -n serverless run \
             python -u /opt/serverless/providers/runpod/worker.py
     else
         printf "No serverless worker available in this environment"