do a shallow git clone of the pyworker

lib/metrics.py

@@ -45,7 +45,6 @@ class Metrics:
         self.model_metrics.workload_received += workload
         self.model_metrics.requests_recieved.add(reqnum)
         self.model_metrics.requests_working.add(reqnum)
-        self.update_pending = True
 
     def _request_end(self, workload: float, reqnum: int) -> None:
         """
@@ -79,10 +78,10 @@ class Metrics:
             elapsed = time.time() - self.last_metric_update
             if self.system_metrics.model_is_loaded is False and elapsed >= 10:
                 log.debug(f"sending loading model metrics after {int(elapsed)}s wait")
-                self.__send_metrics_and_reset()
+                self.__send_metrics_and_reset(elapsed)
             elif self.update_pending or elapsed > 10:
                 log.debug(f"sending loaded model metrics after {int(elapsed)}s wait")
-                self.__send_metrics_and_reset()
+                self.__send_metrics_and_reset(elapsed)
 
     def _model_loaded(self, max_throughput: float) -> None:
         self.system_metrics.model_loading_time = (
@@ -97,13 +96,13 @@ class Metrics:
 
     #######################################Private#######################################
 
-    def __send_metrics_and_reset(self):
+    def __send_metrics_and_reset(self, elapsed):
 
         def compute_autoscaler_data() -> AutoScalaerData:
             return AutoScalaerData(
                 id=self.id,
                 loadtime=(self.system_metrics.model_loading_time or 0.0),
-                cur_load=(self.model_metrics.workload_processing),
+                cur_load=(self.model_metrics.workload_processing / elapsed),
                 max_perf=self.model_metrics.max_throughput,
                 cur_perf=self.model_metrics.cur_perf,
                 error_msg=self.model_metrics.error_msg or "",

start_server.sh

@@ -3,7 +3,8 @@
 set -e -o pipefail
 
 WORKSPACE_DIR="${WORKSPACE_DIR:-/workspace}"
-SERVER_DIR="$WORKSPACE_DIR/worker"
+
+SERVER_DIR="$WORKSPACE_DIR/vast-pyworker"
 ENV_PATH="$WORKSPACE_DIR/worker-env"
 DEBUG_LOG="$WORKSPACE_DIR/debug.log"
 PYWORKER_LOG="$WORKSPACE_DIR/pyworker.log"
@@ -21,23 +22,24 @@ function echo_var(){
     echo "$1: ${!1}"
 }
 
-# Updated validation - BACKEND no longer required, but MODEL_LOG still is
+[ -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
 
-echo "start_server.sh - SDK Worker Version"
+
+echo "start_server.sh"
 date
 
+echo_var BACKEND
 echo_var REPORT_ADDR
 echo_var WORKER_PORT
 echo_var WORKSPACE_DIR
+echo_var SERVER_DIR
 echo_var ENV_PATH
 echo_var DEBUG_LOG
 echo_var PYWORKER_LOG
 echo_var MODEL_LOG
-echo_var MODEL_SERVER_URL
-echo_var PYWORKER_REPO
-echo_var PYWORKER_REF
 
 # Populate /etc/environment with quoted values
 if ! grep -q "VAST" /etc/environment; then
@@ -57,12 +59,10 @@ then
     fi
 
     if [[ ! -d $SERVER_DIR ]]; then
-        echo "Cloning worker repository..."
         git clone --depth=1 "${PYWORKER_REPO:-https://github.com/vast-ai/pyworker}" "$SERVER_DIR"
     fi
 
     if [[ -n ${PYWORKER_REF:-} ]]; then
-        echo "Checking out ref: $PYWORKER_REF"
         (
             cd "$SERVER_DIR"
             git fetch --depth=1 origin "$PYWORKER_REF"
@@ -73,15 +73,7 @@ then
     uv venv --python-preference only-managed "$ENV_PATH" -p 3.10
     source "$ENV_PATH/bin/activate"
 
-    # Install vast-sdk from server-side-sdk branch
-    echo "Installing vast-sdk from GitHub (server-side-sdk branch)..."
-    uv pip install "git+https://github.com/vast-ai/vast-sdk.git@server-side-sdk"
-
-    # Install requirements from worker repo if they exist
-    if [ -f "${SERVER_DIR}/requirements.txt" ]; then
-        echo "Installing additional dependencies from requirements.txt..."
-        uv pip install -r "${SERVER_DIR}/requirements.txt"
-    fi
+    uv pip install -r "${SERVER_DIR}/requirements.txt"
 
     touch ~/.no_auto_tmux
 else
@@ -91,12 +83,7 @@ else
     echo "venv: $VIRTUAL_ENV"
 fi
 
-# Check that worker.py exists
-if [ ! -f "$SERVER_DIR/worker.py" ]; then
-    echo "ERROR: worker.py not found in $SERVER_DIR"
-    echo "Please ensure your PYWORKER_REPO contains a worker.py file"
-    exit 1
-fi
+[ ! -d "$SERVER_DIR/workers/$BACKEND" ] && echo "$BACKEND not supported!" && exit 1
 
 if [ "$USE_SSL" = true ]; then
 
@@ -134,6 +121,9 @@ EOF
         POST "https://console.vast.ai/api/v0/sign_cert/?instance_id=$CONTAINER_ID" > /etc/instance.crt;
 fi
 
+
+
+
 export REPORT_ADDR WORKER_PORT USE_SSL UNSECURED
 
 cd "$SERVER_DIR"
@@ -144,6 +134,5 @@ echo "launching PyWorker server"
 # from the run prior to reboot. past logs are saved in $MODEL_LOG.old for debugging only
 [ -e "$MODEL_LOG" ] && cat "$MODEL_LOG" >> "$MODEL_LOG.old" && : > "$MODEL_LOG"
 
-# Launch the SDK-based worker instead of the old backend system
-(python3 worker.py |& tee -a "$PYWORKER_LOG") &
-echo "launching PyWorker server done"
+(python3 -m "workers.$BACKEND.server" |& tee -a "$PYWORKER_LOG") &
+echo "launching PyWorker server done"