use container venv

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/*" && \

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