From c702edae5f287bdba837e936c8134dade2ffa6a2 Mon Sep 17 00:00:00 2001
From: Dan Saunders <danjsaund@gmail.com>
Date: Fri, 26 Sep 2025 20:19:14 -0400
Subject: [PATCH] use container venv

---
 docker/Dockerfile          |  8 +++++---
 docker/Dockerfile-uv-base  | 11 ++++++-----
 scripts/unsloth_install.py | 14 ++++++++------
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index b0f7d8a88..0680c0420 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -8,10 +8,12 @@ ARG CUDA="118"
 ARG PYTORCH_VERSION="2.1.2"
 ARG GIT_REF="refs/heads/main"
 ARG GIT_SHA="HEAD"
+ARG VENV_PYTHON="/workspace/axolotl-venv/bin/python"
 
 ENV PYTORCH_VERSION=$PYTORCH_VERSION
 ENV GIT_REF=$GIT_REF
 ENV GIT_SHA=$GIT_SHA
+ENV VENV_PYTHON=$VENV_PYTHON
 
 RUN apt-get update && \
     apt-get install -y --allow-change-held-packages vim curl nano libnccl2 libnccl-dev rsync s3fs && \
@@ -29,13 +31,13 @@ RUN git fetch origin "$GIT_REF" && git checkout "$GIT_SHA"
 
 # If AXOLOTL_EXTRAS is set, append it in brackets
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
-        uv pip install --system --no-build-isolation -e .[ring-flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
+        uv pip install --python "$VENV_PYTHON" --no-build-isolation -e .[ring-flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
     else \
-        uv pip install --system --no-build-isolation -e .[ring-flash-attn,optimizers,ray] $AXOLOTL_ARGS; \
+        uv pip install --python "$VENV_PYTHON" --no-build-isolation -e .[ring-flash-attn,optimizers,ray] $AXOLOTL_ARGS; \
     fi && \
     python scripts/unsloth_install.py | sh && \
     python scripts/cutcrossentropy_install.py | sh && \
-    uv pip install --system pytest
+    uv pip install --python "$VENV_PYTHON" pytest
 
 # fix so that git fetch/pull from remote works with shallow clone
 RUN git config remote.origin.fetch "+refs/heads/*:refs/remotes/origin/*" && \
diff --git a/docker/Dockerfile-uv-base b/docker/Dockerfile-uv-base
index 5c25df120..ff72ec50b 100644
--- a/docker/Dockerfile-uv-base
+++ b/docker/Dockerfile-uv-base
@@ -13,6 +13,7 @@ ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 9.0+PTX"
 ENV PYTHON_VERSION=$PYTHON_VERSION
 ENV TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST
 ENV UV_TORCH_BACKEND="cu${CUDA}"
+ENV VENV_PYTHON=/workspace/axolotl-venv/bin/python
 
 RUN apt-get update \
     && apt-get install -y wget git build-essential ninja-build git-lfs libaio-dev pkg-config curl && rm -rf /var/lib/apt/lists/* \
@@ -29,8 +30,8 @@ RUN uv venv --no-project --relocatable axolotl-venv
 
 ENV PATH="/workspace/axolotl-venv/bin:${PATH}"
 
-RUN uv pip install --system packaging setuptools wheel psutil \
-    && uv pip install --system torch==${PYTORCH_VERSION} \
-    && uv pip install --system --no-build-isolation "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" \
-    && uv pip install --system "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main" \
-    && uv pip install --system awscli pydantic
+RUN uv pip install --python "$VENV_PYTHON" packaging setuptools wheel psutil \
+    && uv pip install --python "$VENV_PYTHON" torch==${PYTORCH_VERSION} \
+    && uv pip install --python "$VENV_PYTHON" --no-build-isolation "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" \
+    && uv pip install --python "$VENV_PYTHON" "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main" \
+    && uv pip install --python "$VENV_PYTHON" awscli pydantic
diff --git a/scripts/unsloth_install.py b/scripts/unsloth_install.py
index 717d0cbe2..0bfc1af21 100644
--- a/scripts/unsloth_install.py
+++ b/scripts/unsloth_install.py
@@ -2,6 +2,9 @@
 
 from __future__ import annotations
 
+import sys
+from shlex import quote
+
 try:
     import torch
 except ImportError as exc:  # pragma: no cover
@@ -9,19 +12,18 @@ except ImportError as exc:  # pragma: no cover
 
 from packaging.version import Version as V
 
-torch_version = V(torch.__version__.split("+")[0])
-
-# Unsloth supports torch >= 2.6.0 via the 2025.9 builds.
 MIN_TORCH = V("2.6.0")
 
-if torch_version < MIN_TORCH:
+python_version = V(torch.__version__.split("+")[0])
+if python_version < MIN_TORCH:
     raise RuntimeError(
         f"Torch {torch.__version__} detected, but Unsloth requires >= {MIN_TORCH}."
     )
 
+python_path = quote(sys.executable)
 commands = (
-    "uv pip install --system --no-deps unsloth-zoo==2025.9.12 && "
-    'uv pip install --system --no-deps "unsloth[huggingface]==2025.9.9"'
+    f"uv pip install --python {python_path} --no-deps unsloth-zoo==2025.9.12 && "
+    f'uv pip install --python {python_path} --no-deps "unsloth[huggingface]==2025.9.9"'
 )
 
 print(commands)