From bd3e0032a1c4df91d9f0e3e9b52918b52ccc6041 Mon Sep 17 00:00:00 2001
From: LucasArmandVast <lucas@vast.ai>
Date: Thu, 18 Dec 2025 00:01:52 -0500
Subject: [PATCH] Add SDK version checking (#76)

---
 start_server.sh | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/start_server.sh b/start_server.sh
index 84656c2..4d12a86 100755
--- a/start_server.sh
+++ b/start_server.sh
@@ -46,6 +46,21 @@ JSON
     exit 1
 }
 
+function install_vastai_sdk() {
+    if [ -n "${SDK_VERSION:-}" ]; then
+        echo "Installing vastai-sdk version ${SDK_VERSION}"
+        if ! uv pip install "vastai-sdk==${SDK_VERSION}"; then
+            report_error_and_exit "Failed to install vastai-sdk==${SDK_VERSION}"
+        fi
+    else
+        echo "Installing default vastai-sdk"
+        if ! uv pip install vastai-sdk; then
+            report_error_and_exit "Failed to install vastai-sdk"
+        fi
+    fi
+}
+
+
 [ -n "$BACKEND" ] && [ -z "$HF_TOKEN" ] && report_error_and_exit "HF_TOKEN must be set when BACKEND is 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!"
@@ -123,6 +138,8 @@ then
         report_error_and_exit "Failed to install Python requirements"
     fi
 
+    install_vastai_sdk
+
     if ! touch ~/.no_auto_tmux; then
         report_error_and_exit "Failed to create ~/.no_auto_tmux"
     fi