updated start_server.sh

start_server.sh

@@ -15,18 +15,42 @@ WORKER_PORT="${WORKER_PORT:-3000}"
 mkdir -p "$WORKSPACE_DIR"
 cd "$WORKSPACE_DIR"
 
-# make all output go to $DEBUG_LOG and stdout without having to add `... | tee -a $DEBUG_LOG` to every command
 exec &> >(tee -a "$DEBUG_LOG")
 
 function echo_var(){
     echo "$1: ${!1}"
 }
 
-[ -z "$BACKEND" ] && echo "BACKEND must be set!" && exit 1
-[ -z "$MODEL_LOG" ] && echo "MODEL_LOG must be set!" && exit 1
-[ -z "$HF_TOKEN" ] && echo "HF_TOKEN must be set!" && exit 1
-[ "$BACKEND" = "comfyui" ] && [ -z "$COMFY_MODEL" ] && echo "For comfyui backends, COMFY_MODEL must be set!" && exit 1
+function report_error_and_exit(){
+    local error_msg="$1"
+    echo "ERROR: $error_msg"
 
+    MTOKEN="${MASTER_TOKEN:-}"
+    VERSION="${PYWORKER_VERSION:-0}"
+
+    IFS=',' read -r -a REPORT_ADDRS <<< "${REPORT_ADDR}"
+    for addr in "${REPORT_ADDRS[@]}"; do
+        curl -sS -X POST -H 'Content-Type: application/json' \
+            -d "$(cat <<JSON
+{
+  "id": ${CONTAINER_ID:-0},
+  "mtoken": "${MTOKEN}",
+  "version": "${VERSION}",
+  "error_msg": "${error_msg}",
+  "url": "${URL:-}"
+}
+JSON
+)" "${addr%/}/worker_status/" || true
+    done
+
+    exit 1
+}
+
+[ -z "$BACKEND" ] && report_error_and_exit "BACKEND must be set!"
+[ -z "$MODEL_LOG" ] && report_error_and_exit "MODEL_LOG must be set!"
+[ -z "$HF_TOKEN" ] && report_error_and_exit "HF_TOKEN must be set!"
+[ -z "$CONTAINER_ID" ] && report_error_and_exit "CONTAINER_ID must be set!"
+[ "$BACKEND" = "comfyui" ] && [ -z "$COMFY_MODEL" ] && report_error_and_exit "For comfyui backends, COMFY_MODEL must be set!"
 
 echo "start_server.sh"
 date
@@ -41,55 +65,105 @@ echo_var DEBUG_LOG
 echo_var PYWORKER_LOG
 echo_var MODEL_LOG
 
-# if instance is rebooted, we want to clear out the log file so pyworker doesn't read lines
-# from the run prior to reboot. past logs are saved in $MODEL_LOG.old for debugging only
 if [ -e "$MODEL_LOG" ]; then
     echo "Rotating model log at $MODEL_LOG to $MODEL_LOG.old"
-    cat "$MODEL_LOG" >> "$MODEL_LOG.old" 
-    : > "$MODEL_LOG"
+    if ! cat "$MODEL_LOG" >> "$MODEL_LOG.old"; then
+        report_error_and_exit "Failed to rotate model log"
+    fi
+    if ! : > "$MODEL_LOG"; then
+        report_error_and_exit "Failed to truncate model log"
+    fi
 fi
 
-# Populate /etc/environment with quoted values
-if ! grep -q "VAST" /etc/environment; then
-    env -0 | grep -zEv "^(HOME=|SHLVL=)|CONDA" | while IFS= read -r -d '' line; do
-            name=${line%%=*}
-            value=${line#*=}
+ENV_SENTINEL="VAST_ENV_INITIALIZED=1"
+
+if ! grep -qxF "$ENV_SENTINEL" /etc/environment 2>/dev/null; then
+    {
+        echo "$ENV_SENTINEL"
+
+        # Serialize current environment
+        env -0 | while IFS= read -r -d '' kv; do
+            name=${kv%%=*}
+            value=${kv#*=}
+
+            # Skip volatile or shell-internal variables
+            case "$name" in
+                HOME|SHLVL|PWD|OLDPWD|_ ) continue ;;
+            esac
+
+            # Escape backslashes, quotes, and newlines
+            value=${value//\\/\\\\}
+            value=${value//\"/\\\"}
+            value=${value//$'\n'/\\n}
+
             printf '%s="%s"\n' "$name" "$value"
-        done > /etc/environment
+        done
+    } > /etc/environment.tmp
+
+    if ! mv /etc/environment.tmp /etc/environment; then
+        echo "WARNING: Failed to update /etc/environment, continuing anyway"
+        rm -f /etc/environment.tmp
+    fi
 fi
 
 if [ ! -d "$ENV_PATH" ]
 then
     echo "setting up venv"
     if ! which uv; then
-        curl -LsSf https://astral.sh/uv/install.sh | sh
-        source ~/.local/bin/env
+        if ! curl -LsSf https://astral.sh/uv/install.sh | sh; then
+            report_error_and_exit "Failed to install uv package manager"
+        fi
+        if [[ -f ~/.local/bin/env ]]; then
+            if ! source ~/.local/bin/env; then
+                report_error_and_exit "Failed to source uv environment"
+            fi
+        else
+            echo "WARNING: ~/.local/bin/env not found after uv installation"
+        fi
     fi
 
-    # Fork testing
-    [[ ! -d $SERVER_DIR ]] && git clone "${PYWORKER_REPO:-https://github.com/vast-ai/pyworker}" "$SERVER_DIR"
+    if [[ ! -d $SERVER_DIR ]]; then
+        if ! git clone "${PYWORKER_REPO:-https://github.com/vast-ai/pyworker}" "$SERVER_DIR"; then
+            report_error_and_exit "Failed to clone pyworker repository"
+        fi
+    fi
     if [[ -n ${PYWORKER_REF:-} ]]; then
-        (cd "$SERVER_DIR" && git checkout "$PYWORKER_REF")
+        if ! (cd "$SERVER_DIR" && git checkout "$PYWORKER_REF"); then
+            report_error_and_exit "Failed to checkout pyworker reference: $PYWORKER_REF"
+        fi
     fi
 
-    uv venv --python-preference only-managed "$ENV_PATH" -p 3.10
-    source "$ENV_PATH/bin/activate"
+    if ! uv venv --python-preference only-managed "$ENV_PATH" -p 3.10; then
+        report_error_and_exit "Failed to create virtual environment"
+    fi
 
-    uv pip install -r "${SERVER_DIR}/requirements.txt"
+    if ! source "$ENV_PATH/bin/activate"; then
+        report_error_and_exit "Failed to activate virtual environment"
+    fi
 
-    touch ~/.no_auto_tmux
+    if ! uv pip install -r "${SERVER_DIR}/requirements.txt"; then
+        report_error_and_exit "Failed to install Python requirements"
+    fi
+
+    if ! touch ~/.no_auto_tmux; then
+        report_error_and_exit "Failed to create ~/.no_auto_tmux"
+    fi
 else
-    [[ -f ~/.local/bin/env ]] && source ~/.local/bin/env
-    source "$WORKSPACE_DIR/worker-env/bin/activate"
+    if [[ -f ~/.local/bin/env ]]; then
+        if ! source ~/.local/bin/env; then
+            report_error_and_exit "Failed to source uv environment"
+        fi
+    fi
+    if ! source "$WORKSPACE_DIR/worker-env/bin/activate"; then
+        report_error_and_exit "Failed to activate existing virtual environment"
+    fi
     echo "environment activated"
     echo "venv: $VIRTUAL_ENV"
 fi
 
-[ ! -d "$SERVER_DIR/workers/$BACKEND" ] && echo "$BACKEND not supported!" && exit 1
-
 if [ "$USE_SSL" = true ]; then
 
-    cat << EOF > /etc/openssl-san.cnf
+    if ! cat << EOF > /etc/openssl-san.cnf
     [req]
     default_bits       = 2048
     distinguished_name = req_distinguished_name
@@ -109,77 +183,64 @@ if [ "$USE_SSL" = true ]; then
     [alt_names]
     IP.1   = 0.0.0.0
 EOF
+    then
+        report_error_and_exit "Failed to write OpenSSL config"
+    fi
 
-    openssl req -newkey rsa:2048 -subj "/C=US/ST=CA/CN=pyworker.vast.ai/" \
+    if ! openssl req -newkey rsa:2048 -subj "/C=US/ST=CA/CN=pyworker.vast.ai/" \
         -nodes \
         -sha256 \
         -keyout /etc/instance.key \
         -out /etc/instance.csr \
-        -config /etc/openssl-san.cnf
+        -config /etc/openssl-san.cnf; then
+        report_error_and_exit "Failed to generate SSL certificate request"
+    fi
 
-    curl --header 'Content-Type: application/octet-stream' \
-        --data-binary @//etc/instance.csr \
+    if ! curl --header 'Content-Type: application/octet-stream' \
+        --data-binary @/etc/instance.csr \
         -X \
-        POST "https://console.vast.ai/api/v0/sign_cert/?instance_id=$CONTAINER_ID" > /etc/instance.crt;
+        POST "https://console.vast.ai/api/v0/sign_cert/?instance_id=$CONTAINER_ID" > /etc/instance.crt; then
+        report_error_and_exit "Failed to sign SSL certificate"
+    fi
 fi
 
-
-
-
 export REPORT_ADDR WORKER_PORT USE_SSL UNSECURED
 
-cd "$SERVER_DIR"
+if ! cd "$SERVER_DIR"; then
+    report_error_and_exit "Failed to cd into SERVER_DIR: $SERVER_DIR"
+fi
 
 echo "launching PyWorker server"
 
 set +e
 
-# Try worker entrypoint first
-echo "trying workers.${BACKEND}.worker"
-python3 -m "workers.${BACKEND}.worker" |& tee -a "$PYWORKER_LOG"
-PY_STATUS=${PIPESTATUS[0]}
+PY_STATUS=1
 
-# If that fails, fall back to server
-if [ "${PY_STATUS}" -ne 0 ]; then
-  echo "workers.${BACKEND}.worker failed with status ${PY_STATUS}, trying workers.${BACKEND}.server"
-  python3 -m "workers.${BACKEND}.server" |& tee -a "$PYWORKER_LOG"
-  PY_STATUS=${PIPESTATUS[0]}
+if [ -f "$SERVER_DIR/worker.py" ]; then
+    echo "trying worker.py"
+    python3 -m "worker" |& tee -a "$PYWORKER_LOG"
+    PY_STATUS=${PIPESTATUS[0]}
+fi
+
+if [ "${PY_STATUS}" -ne 0 ] && [ -f "$SERVER_DIR/workers/$BACKEND/worker.py" ]; then
+    echo "trying workers.${BACKEND}.worker"
+    python3 -m "workers.${BACKEND}.worker" |& tee -a "$PYWORKER_LOG"
+    PY_STATUS=${PIPESTATUS[0]}
+fi
+
+if [ "${PY_STATUS}" -ne 0 ] && [ -f "$SERVER_DIR/workers/$BACKEND/server.py" ]; then
+    echo "trying workers.${BACKEND}.server"
+    python3 -m "workers.${BACKEND}.server" |& tee -a "$PYWORKER_LOG"
+    PY_STATUS=${PIPESTATUS[0]}
 fi
 
 set -e
 
 if [ "${PY_STATUS}" -ne 0 ]; then
-  echo "PyWorker exited with status ${PY_STATUS}; notifying autoscaler..."
-  ERROR_MSG="PyWorker exited: code ${PY_STATUS}"
-  MTOKEN="${MASTER_TOKEN:-}"
-  VERSION="${PYWORKER_VERSION:-0}"
-
-  IFS=',' read -r -a REPORT_ADDRS <<< "${REPORT_ADDR}"
-  for addr in "${REPORT_ADDRS[@]}"; do
-    curl -sS -X POST -H 'Content-Type: application/json' \
-      -d "$(cat <<JSON
-{
-  "id": ${CONTAINER_ID:-0},
-  "mtoken": "${MTOKEN}",
-  "version": "${VERSION}",
-  "loadtime": 0,
-  "new_load": 0,
-  "cur_load": 0,
-  "rej_load": 0,
-  "max_perf": 0,
-  "cur_perf": 0,
-  "error_msg": "${ERROR_MSG}",
-  "num_requests_working": 0,
-  "num_requests_recieved": 0,
-  "additional_disk_usage": 0,
-  "working_request_idxs": [],
-  "cur_capacity": 0,
-  "max_capacity": 0,
-  "url": "${URL}"
-}
-JSON
-)" "${addr%/}/worker_status/" || true
-  done
+    if [ ! -f "$SERVER_DIR/worker.py" ] && [ ! -f "$SERVER_DIR/workers/$BACKEND/worker.py" ] && [ ! -f "$SERVER_DIR/workers/$BACKEND/server.py" ]; then
+        report_error_and_exit "Failed to find PyWorker"
+    fi
+    report_error_and_exit "PyWorker exited with status ${PY_STATUS}"
 fi
 
 echo "launching PyWorker server done"