refactor: use endpoint_id instead of endpoint name for routing

- Update route_payload to use endpoint_id instead of endpoint name
- Update AuthData to expect endpoint_id (int) instead of endpoint (str)
- Update ClientState to track endpoint_id
- Update comfyui client functions to use endpoint_id
- Fetch endpoint info (id + api_key) instead of just api_key

This aligns with the autoscaler changes in AUTO-848 that switched
to ID-based endpoint lookups for improved security and consistency.

lib/data_types.py

@@ -66,7 +66,7 @@ class AuthData:
     """data used to authenticate requester"""
 
     cost: str
-    endpoint: str
+    endpoint_id: int
     reqnum: int
     request_idx: int
     signature: str

lib/test_utils.py

@@ -75,6 +75,7 @@ def print_truncate_res(res: str):
 @dataclass
 class ClientState:
     endpoint_group_name: str
+    endpoint_id: int
     api_key: str
     server_url: str
     worker_endpoint: str
@@ -95,7 +96,7 @@ class ClientState:
             self.status = ClientStatus.Error
             return
         route_payload = {
-            "endpoint": self.endpoint_group_name,
+            "endpoint_id": self.endpoint_id,
             "api_key": self.api_key,
             "cost": self.payload.count_workload(),
         }
@@ -244,16 +245,19 @@ def run_test(
     print_thread = threading.Thread(target=print_state, args=(clients, num_requests))
     print_thread.daemon = True  # makes threads get killed on program exit
     print_thread.start()
-    endpoint_api_key = Endpoint.get_endpoint_api_key(
+    endpoint_info = Endpoint.get_endpoint_info(
         endpoint_name=endpoint_group_name, account_api_key=api_key, instance=instance
     )
-    if not endpoint_api_key:
+    if not endpoint_info:
         log.debug(f"Endpoint {endpoint_group_name} not found for API key")
         return
+    endpoint_id = endpoint_info["id"]
+    endpoint_api_key = endpoint_info["api_key"]
     try:
         for _ in range(num_requests):
             client = ClientState(
                 endpoint_group_name=endpoint_group_name,
+                endpoint_id=endpoint_id,
                 api_key=endpoint_api_key,
                 server_url=server_url,
                 worker_endpoint=worker_endpoint,

requirements.txt

@@ -1,6 +1,4 @@
-aiohttp==3.10.1
-aiodns~=3.6.0
-pycares~=4.11.0
+aiohttp[speedups]==3.10.1
 anyio~=4.4
 lib~=4.0
 nltk~=3.9
@@ -10,4 +8,4 @@ Requests~=2.32
 transformers~=4.52
 utils==1.0.*
 hf_transfer>=0.1.9
-vastai-sdk>=0.2.0
+vastai-sdk>=0.2.0

start_server.sh

@@ -22,49 +22,10 @@ function echo_var(){
     echo "$1: ${!1}"
 }
 
-function report_error_and_exit(){
-    local error_msg="$1"
-    echo "ERROR: $error_msg"
-    
-    # Report error to autoscaler
-    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
-    
-    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!"
+[ -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"
@@ -84,86 +45,51 @@ echo_var MODEL_LOG
 # 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"
-    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
+    cat "$MODEL_LOG" >> "$MODEL_LOG.old" 
+    : > "$MODEL_LOG"
 fi
 
 # Populate /etc/environment with quoted values
 if ! grep -q "VAST" /etc/environment; then
-    if ! env -0 | grep -zEv "^(HOME=|SHLVL=)|CONDA" | while IFS= read -r -d '' line; do
+    env -0 | grep -zEv "^(HOME=|SHLVL=)|CONDA" | while IFS= read -r -d '' line; do
             name=${line%%=*}
             value=${line#*=}
             printf '%s="%s"\n' "$name" "$value"
-        done > /etc/environment; then
-        echo "WARNING: Failed to populate /etc/environment, continuing anyway"
-    fi
+        done > /etc/environment
 fi
 
 if [ ! -d "$ENV_PATH" ]
 then
     echo "setting up venv"
     if ! which uv; then
-        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
+        curl -LsSf https://astral.sh/uv/install.sh | sh
+        source ~/.local/bin/env
     fi
 
     # Fork testing
-    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
+    [[ ! -d $SERVER_DIR ]] && git clone "${PYWORKER_REPO:-https://github.com/vast-ai/pyworker}" "$SERVER_DIR"
     if [[ -n ${PYWORKER_REF:-} ]]; then
-        if ! (cd "$SERVER_DIR" && git checkout "$PYWORKER_REF"); then
-            report_error_and_exit "Failed to checkout pyworker reference: $PYWORKER_REF"
-        fi
+        (cd "$SERVER_DIR" && git checkout "$PYWORKER_REF")
     fi
 
-    if ! uv venv --python-preference only-managed "$ENV_PATH" -p 3.10; then
-        report_error_and_exit "Failed to create virtual environment"
-    fi
-    
-    if ! source "$ENV_PATH/bin/activate"; then
-        report_error_and_exit "Failed to activate virtual environment"
-    fi
+    uv venv --python-preference only-managed "$ENV_PATH" -p 3.10
+    source "$ENV_PATH/bin/activate"
 
-    if ! uv pip install -r "${SERVER_DIR}/requirements.txt"; then
-        report_error_and_exit "Failed to install Python requirements"
-    fi
+    uv pip install -r "${SERVER_DIR}/requirements.txt"
 
-    if ! touch ~/.no_auto_tmux; then
-        report_error_and_exit "Failed to create ~/.no_auto_tmux"
-    fi
+    touch ~/.no_auto_tmux
 else
-    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
+    [[ -f ~/.local/bin/env ]] && source ~/.local/bin/env
+    source "$WORKSPACE_DIR/worker-env/bin/activate"
     echo "environment activated"
     echo "venv: $VIRTUAL_ENV"
 fi
 
-[ ! -d "$SERVER_DIR/workers/$BACKEND" ] && report_error_and_exit "$BACKEND not supported!"
+[ ! -d "$SERVER_DIR/workers/$BACKEND" ] && echo "$BACKEND not supported!" && exit 1
 
 if [ "$USE_SSL" = true ]; then
 
-    if ! cat << EOF > /etc/openssl-san.cnf
+    cat << EOF > /etc/openssl-san.cnf
     [req]
     default_bits       = 2048
     distinguished_name = req_distinguished_name
@@ -183,25 +109,18 @@ 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
 
-    if ! openssl req -newkey rsa:2048 -subj "/C=US/ST=CA/CN=pyworker.vast.ai/" \
+    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; then
-        report_error_and_exit "Failed to generate SSL certificate request"
-    fi
+        -config /etc/openssl-san.cnf
 
-    if ! curl --header 'Content-Type: application/octet-stream' \
-        --data-binary @/etc/instance.csr \
+    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; then
-        report_error_and_exit "Failed to sign SSL certificate"
-    fi
+        POST "https://console.vast.ai/api/v0/sign_cert/?instance_id=$CONTAINER_ID" > /etc/instance.crt;
 fi
 
 
@@ -209,9 +128,7 @@ fi
 
 export REPORT_ADDR WORKER_PORT USE_SSL UNSECURED
 
-if ! cd "$SERVER_DIR"; then
-    report_error_and_exit "Failed to cd into SERVER_DIR: $SERVER_DIR"
-fi
+cd "$SERVER_DIR"
 
 echo "launching PyWorker server"
 
@@ -221,7 +138,37 @@ PY_STATUS=${PIPESTATUS[0]}
 set -e
 
 if [ "${PY_STATUS}" -ne 0 ]; then
-  report_error_and_exit "PyWorker exited with status ${PY_STATUS}"
+  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
 fi
 
 echo "launching PyWorker server done"

workers/comfyui/client.py

@@ -13,11 +13,11 @@ from vastai import Serverless
 ENDPOINT_NAME = "my-comfyui-endpoint"
 COST = 100 # Use a constant cost for image generation
 
-def call_default_workflow(client: Serverless) -> None:
+def call_default_workflow(endpoint_id: int, api_key: str, server_url: str) -> None:
     WORKER_ENDPOINT = "/prompt"
     COST = 100
     route_payload = {
-        "endpoint": endpoint_group_name,
+        "endpoint_id": endpoint_id,
         "api_key": api_key,
         "cost": COST,
     }
@@ -32,7 +32,7 @@ def call_default_workflow(client: Serverless) -> None:
     auth_data = dict(
         signature=message["signature"],
         cost=message["cost"],
-        endpoint=message["endpoint"],
+        endpoint_id=message["endpoint_id"],
         reqnum=message["reqnum"],
         url=message["url"],
     )
@@ -52,12 +52,12 @@ def call_default_workflow(client: Serverless) -> None:
 
 
 def call_custom_workflow_for_sd3(
-    endpoint_group_name: str, api_key: str, server_url: str
+    endpoint_id: int, api_key: str, server_url: str
 ) -> None:
     WORKER_ENDPOINT = "/custom-workflow"
     COST = 100
     route_payload = {
-        "endpoint": endpoint_group_name,
+        "endpoint_id": endpoint_id,
         "api_key": api_key,
         "cost": COST,
     }
@@ -72,7 +72,7 @@ def call_custom_workflow_for_sd3(
     auth_data = dict(
         signature=message["signature"],
         cost=message["cost"],
-        endpoint=message["endpoint"],
+        endpoint_id=message["endpoint_id"],
         reqnum=message["reqnum"],
         url=message["url"],
         request_idx=message["request_idx"],
@@ -146,25 +146,28 @@ def call_custom_workflow_for_sd3(
 if __name__ == "__main__":
     from lib.test_utils import test_args
 
+    log = logging.getLogger(__name__)
     args = test_args.parse_args()
-    endpoint_api_key = Endpoint.get_endpoint_api_key(
+    endpoint_info = Endpoint.get_endpoint_info(
         endpoint_name=args.endpoint_group_name,
         account_api_key=args.api_key,
         instance=args.instance,
     )
-    if endpoint_api_key:
+    if endpoint_info:
+        endpoint_id = endpoint_info["id"]
+        endpoint_api_key = endpoint_info["api_key"]
         try:
             call_default_workflow(
+                endpoint_id=endpoint_id,
                 api_key=endpoint_api_key,
-                endpoint_group_name=args.endpoint_group_name,
                 server_url=args.server_url,
             )
             call_custom_workflow_for_sd3(
+                endpoint_id=endpoint_id,
                 api_key=endpoint_api_key,
-                endpoint_group_name=args.endpoint_group_name,
                 server_url=args.server_url,
             )
         except Exception as e:
             log.error(f"Error during API call: {e}")
     else:
-        log.error(f"Failed to get API key for endpoint {args.endpoint_group_name} ")
+        log.error(f"Failed to get endpoint info for {args.endpoint_group_name}")

workers/openai/test_load.py

@@ -60,7 +60,7 @@ def do_one(endpoint_name: str,
            worker_session):
     try:
         workload = payload.count_workload()
-        route_payload = {"endpoint": endpoint_name, "api_key": endpoint_api_key, "cost": workload}
+        route_payload = {"endpoint_id": endpoint_id, "api_key": endpoint_api_key, "cost": workload}
         headers = {"Authorization": f"Bearer {endpoint_api_key}"}
         start = time.time()
         r0 = route_session.post(urljoin(server_url, "/route/"), json=route_payload, headers=headers, timeout=4)