From 6060f8ce0c1f4bc23b3e6e391b1d33ee19e95f3f Mon Sep 17 00:00:00 2001 From: Lucas Armand Date: Fri, 12 Dec 2025 10:04:33 -0800 Subject: [PATCH] updated start_server.sh --- start_server.sh | 217 +++++++++++++++++++++++++++++++----------------- 1 file changed, 139 insertions(+), 78 deletions(-) diff --git a/start_server.sh b/start_server.sh index 3bb80c4..5e8f93f 100755 --- a/start_server.sh +++ b/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 <> "$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 <