Fix Slurm job Python environment isolation

Author: alec-flowers

configs/prefetch-ai-dynamo-wheel.sh

@@ -29,13 +29,31 @@ runtime_wheel_path() {
     find "${wheel_dir}" -maxdepth 1 -type f -name "${DYNAMO_RUNTIME_WHEEL_PATTERN}" -print -quit
 }
 
+python_with_pip() {
+    if python3 -m pip --version >/dev/null 2>&1; then
+        command -v python3
+        return
+    fi
+
+    if ! command -v uv >/dev/null 2>&1; then
+        echo "ERROR: python3 does not provide pip, and uv is unavailable to create a pip-seeded prefetch venv" >&2
+        return 1
+    fi
+
+    local prefetch_venv="${DYNAMO_PREFETCH_VENV:-${wheel_dir}/.prefetch-venv}"
+    uv venv --seed "${prefetch_venv}" >/dev/null
+    echo "${prefetch_venv}/bin/python"
+}
+
 if [ -f "${wheel_path}" ] && [ -n "$(runtime_wheel_path)" ]; then
     echo "ai-dynamo wheels already staged: ${wheel_dir}"
     exit 0
 fi
 
 download_wheels() {
-    python3 -m pip download \
+    local python_bin
+    python_bin="$(python_with_pip)"
+    "${python_bin}" -m pip download \
         --no-deps \
         --pre \
         --only-binary=:all: \

src/srtctl/templates/job_script_minimal.j2

@@ -82,12 +82,37 @@ export {{ key }}={{ value }}
 echo ""
 echo "Preparing srtctl environment..."
 
-# Install uv if not present (single binary, no dependencies)
-if ! command -v uv &> /dev/null; then
+# SLURM inherits the submitter environment by default. If srtctl was submitted
+# from an activated virtualenv, that venv can point at a pipless or wrong-arch
+# Python on the compute node. Drop it before selecting Python/uv for the job.
+if [ -n "${VIRTUAL_ENV:-}" ]; then
+    echo "Ignoring inherited virtualenv: ${VIRTUAL_ENV}"
+    CLEAN_PATH=""
+    OLD_IFS="${IFS}"
+    IFS=":"
+    for PATH_ENTRY in ${PATH}; do
+        if [ "${PATH_ENTRY}" != "${VIRTUAL_ENV}/bin" ]; then
+            if [ -z "${CLEAN_PATH}" ]; then
+                CLEAN_PATH="${PATH_ENTRY}"
+            else
+                CLEAN_PATH="${CLEAN_PATH}:${PATH_ENTRY}"
+            fi
+        fi
+    done
+    IFS="${OLD_IFS}"
+    export PATH="${CLEAN_PATH}"
+    unset VIRTUAL_ENV
+fi
+
+# Install a job-local uv so inherited submitter binaries cannot shadow the
+# compute-node architecture.
+UV_BIN_DIR="${OUTPUT_DIR}/uv-bin"
+mkdir -p "${UV_BIN_DIR}"
+if ! "${UV_BIN_DIR}/uv" --version >/dev/null 2>&1; then
     echo "Installing uv package manager..."
-    curl -LsSf https://astral.sh/uv/install.sh | sh
-    export PATH="$HOME/.local/bin:$PATH"
+    curl -LsSf https://astral.sh/uv/install.sh | env XDG_BIN_HOME="${UV_BIN_DIR}" INSTALLER_NO_MODIFY_PATH=1 sh
 fi
+export PATH="${UV_BIN_DIR}:$PATH"
 
 echo "Using uv with Python 3.12..."
 

tests/test_configs.py

@@ -596,6 +596,9 @@ def test_sbatch_template_prefetches_dynamo_wheel(self):
 
         assert "export DYNAMO_VERSION=1.2.0.dev20260426" in script
         assert "export DYNAMO_WHEEL_NAME=ai_dynamo-1.2.0.dev20260426-py3-none-any.whl" in script
+        assert "Ignoring inherited virtualenv" in script
+        assert 'unset VIRTUAL_ENV' in script
+        assert 'UV_BIN_DIR="${OUTPUT_DIR}/uv-bin"' in script
         assert "configs/prefetch-ai-dynamo-wheel.sh" in script
 
     def test_setup_script_env_var_override(self, monkeypatch):