From 4531336f4952e22d507060a0127c9aa1444fa74e Mon Sep 17 00:00:00 2001
From: Rob Ballantyne <rob@dynamedia.uk>
Date: Mon, 12 Feb 2024 12:02:18 +0000
Subject: [PATCH] Serverless worker startup fix

---
 .../opt/ai-dock/bin/supervisor-comfyui-rp-api.sh          | 5 ++++-
 build/COPY_ROOT/opt/ai-dock/bin/supervisor-comfyui.sh     | 4 +++-
 build/COPY_ROOT/opt/ai-dock/bin/supervisor-serverless.sh  | 8 ++++----
 3 files changed, 11 insertions(+), 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
index 1a6fe476..7a595dcf 100755
--- 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,9 +7,12 @@ 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
@@ -17,7 +20,7 @@ function start() {
 
     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
index 570c2dc9..cbb5769d 100755
--- 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
index b100ebb3..f4b89f51 100755
--- 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"