From 08c88f75271ca86b347f7a0ca8f4659837685cd1 Mon Sep 17 00:00:00 2001
From: Rob Ballantyne <rob@robballantyne.co.uk>
Date: Wed, 20 Aug 2025 09:34:09 +0100
Subject: [PATCH] Improve testability

---
 start_server.sh | 29 +++++++++++++++++++++--------
 1 file changed, 21 insertions(+), 8 deletions(-)

diff --git a/start_server.sh b/start_server.sh
index 8ef61a7..2d53525 100755
--- a/start_server.sh
+++ b/start_server.sh
@@ -41,20 +41,33 @@ echo_var DEBUG_LOG
 echo_var PYWORKER_LOG
 echo_var MODEL_LOG
 
-env | grep _ >> /etc/environment;
-
+# 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#*=}
+            printf '%s="%s"\n' "$name" "$value"
+        done > /etc/environment
+fi
 
 if [ ! -d "$ENV_PATH" ]
 then
     echo "setting up venv"
-    curl -LsSf https://astral.sh/uv/install.sh | sh
-    source ~/.local/bin/env
-    git clone https://github.com/vast-ai/pyworker "$SERVER_DIR"
+    if ! which uv; then
+        curl -LsSf https://astral.sh/uv/install.sh | sh
+        source ~/.local/bin/env
+    fi
 
-    uv venv --managed-python "$WORKSPACE_DIR/worker-env" -p 3.10
-    source "$WORKSPACE_DIR/worker-env/bin/activate"
+    # 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")
+    fi
 
-    uv pip install -r vast-pyworker/requirements.txt
+    uv venv --managed-python "$ENV_PATH" -p 3.10
+    source "$ENV_PATH/bin/activate"
+
+    uv pip install -r "${SERVER_DIR}/requirements.txt"
 
     touch ~/.no_auto_tmux
 else