spelling fix

PyWorker fixes for cur_load and acks bug

lib/metrics.py

@@ -45,6 +45,7 @@ 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:
         """
@@ -78,10 +79,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(elapsed)
+                self.__send_metrics_and_reset()
             elif self.update_pending or elapsed > 10:
                 log.debug(f"sending loaded model metrics after {int(elapsed)}s wait")
-                self.__send_metrics_and_reset(elapsed)
+                self.__send_metrics_and_reset()
 
     def _model_loaded(self, max_throughput: float) -> None:
         self.system_metrics.model_loading_time = (
@@ -96,13 +97,13 @@ class Metrics:
 
     #######################################Private#######################################
 
-    def __send_metrics_and_reset(self, elapsed):
+    def __send_metrics_and_reset(self):
 
         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 / elapsed),
+                cur_load=(self.model_metrics.workload_processing),
                 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,8 +3,7 @@
 set -e -o pipefail
 
 WORKSPACE_DIR="${WORKSPACE_DIR:-/workspace}"
-
-SERVER_DIR="$WORKSPACE_DIR/vast-pyworker"
+SERVER_DIR="$WORKSPACE_DIR/worker"
 ENV_PATH="$WORKSPACE_DIR/worker-env"
 DEBUG_LOG="$WORKSPACE_DIR/debug.log"
 PYWORKER_LOG="$WORKSPACE_DIR/pyworker.log"
@@ -22,24 +21,23 @@ function echo_var(){
     echo "$1: ${!1}"
 }
 
-[ -z "$BACKEND" ] && echo "BACKEND must be set!" && exit 1
+# Updated validation - BACKEND no longer required, but MODEL_LOG still is
 [ -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"
+echo "start_server.sh - SDK Worker Version"
 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
@@ -58,16 +56,32 @@ then
         source ~/.local/bin/env
     fi
 
-    # Fork testing
-    git clone "${PYWORKER_REPO:-https://github.com/vast-ai/pyworker}" "$SERVER_DIR"
-    if [[ -n ${PYWORKER_REF:-} ]]; then
-        (cd "$SERVER_DIR" && git checkout "$PYWORKER_REF")
+    if [[ ! -d $SERVER_DIR ]]; then
+        echo "Cloning worker repository..."
+        git clone --depth=1 "${PYWORKER_REPO:-https://github.com/vast-ai/pyworker}" "$SERVER_DIR"
     fi
 
-    uv venv --managed-python "$ENV_PATH" -p 3.10
+    if [[ -n ${PYWORKER_REF:-} ]]; then
+        echo "Checking out ref: $PYWORKER_REF"
+        (
+            cd "$SERVER_DIR"
+            git fetch --depth=1 origin "$PYWORKER_REF"
+            git checkout "$PYWORKER_REF"
+        )
+    fi
+
+    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
 
     touch ~/.no_auto_tmux
 else
@@ -77,7 +91,12 @@ else
     echo "venv: $VIRTUAL_ENV"
 fi
 
-[ ! -d "$SERVER_DIR/workers/$BACKEND" ] && echo "$BACKEND not supported!" && exit 1
+# 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
 
 if [ "$USE_SSL" = true ]; then
 
@@ -115,9 +134,6 @@ 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"
@@ -128,5 +144,6 @@ 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"
 
-(python3 -m "workers.$BACKEND.server" |& tee -a "$PYWORKER_LOG") &
+# Launch the SDK-based worker instead of the old backend system
+(python3 worker.py |& tee -a "$PYWORKER_LOG") &
 echo "launching PyWorker server done"

workers/comfyui-json/server.py

@@ -70,7 +70,7 @@ class ComfyWorkflowHandler(EndpointHandler[ComfyWorkflowData]):
 
     @property
     def healthcheck_endpoint(self) -> Optional[str]:
-        return "/health"
+        return f"{MODEL_SERVER_URL}/health"
 
     @classmethod
     def payload_cls(cls) -> Type[ComfyWorkflowData]: