From 8065fed126e7877fbb3568089481318cc5528bb8 Mon Sep 17 00:00:00 2001
From: Salman Mohammadi <salman.mohammadi@outlook.com>
Date: Wed, 2 Jul 2025 15:27:42 +0100
Subject: [PATCH] adding venv to prompt

---
 docker/Dockerfile-base         | 8 ++++++--
 docker/Dockerfile-base-next    | 8 ++++++--
 docker/Dockerfile-base-nightly | 8 ++++++--
 docker/Dockerfile-uv-base      | 3 +++
 4 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base
index 52201f276..a6ff8de69 100644
--- a/docker/Dockerfile-base
+++ b/docker/Dockerfile-base
@@ -22,9 +22,9 @@ RUN apt-get update \
     && mkdir /root/.conda \
     && bash Miniconda3-latest-Linux-x86_64.sh -b \
     && rm -f Miniconda3-latest-Linux-x86_64.sh \
-    && conda create -n "py${PYTHON_VERSION}" python="${PYTHON_VERSION}"
+    && conda create -n "axolotl-py${PYTHON_VERSION}" python="${PYTHON_VERSION}"
 
-ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"
+ENV PATH="/root/miniconda3/envs/axolotl-py${PYTHON_VERSION}/bin:${PATH}"
 
 WORKDIR /workspace
 
@@ -37,3 +37,7 @@ RUN git lfs install --skip-repo && \
     pip3 install awscli && \
     # The base image ships with `pydantic==1.8.2` which is not working
     pip3 install -U --no-cache-dir pydantic==1.10.10
+
+# Initialize conda for bash shell to show environment in prompt
+RUN conda init bash && \
+    echo "conda activate axolotl-py${PYTHON_VERSION}" >> ~/.bashrc
diff --git a/docker/Dockerfile-base-next b/docker/Dockerfile-base-next
index 85bac2516..cd5accdb6 100644
--- a/docker/Dockerfile-base-next
+++ b/docker/Dockerfile-base-next
@@ -22,9 +22,9 @@ RUN apt-get update \
     && mkdir /root/.conda \
     && bash Miniconda3-latest-Linux-x86_64.sh -b \
     && rm -f Miniconda3-latest-Linux-x86_64.sh \
-    && conda create -n "py${PYTHON_VERSION}" python="${PYTHON_VERSION}"
+    && conda create -n "axolotl-py${PYTHON_VERSION}" python="${PYTHON_VERSION}"
 
-ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"
+ENV PATH="/root/miniconda3/envs/axolotl-py${PYTHON_VERSION}/bin:${PATH}"
 
 WORKDIR /workspace
 
@@ -36,3 +36,7 @@ RUN python3 -m pip install --upgrade pip && pip3 install packaging && \
 RUN git lfs install --skip-repo && \
     pip3 install awscli && \
     pip3 install -U --no-cache-dir pydantic==2.10.6
+
+# Initialize conda for bash shell to show environment in prompt
+RUN conda init bash && \
+    echo "conda activate axolotl-py${PYTHON_VERSION}" >> ~/.bashrc
diff --git a/docker/Dockerfile-base-nightly b/docker/Dockerfile-base-nightly
index 85805ea41..e0acca766 100644
--- a/docker/Dockerfile-base-nightly
+++ b/docker/Dockerfile-base-nightly
@@ -22,9 +22,9 @@ RUN apt-get update \
     && mkdir /root/.conda \
     && bash Miniconda3-latest-Linux-x86_64.sh -b \
     && rm -f Miniconda3-latest-Linux-x86_64.sh \
-    && conda create -n "py${PYTHON_VERSION}" python="${PYTHON_VERSION}"
+    && conda create -n "axolotl-py${PYTHON_VERSION}" python="${PYTHON_VERSION}"
 
-ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}"
+ENV PATH="/root/miniconda3/envs/axolotl-py${PYTHON_VERSION}/bin:${PATH}"
 
 WORKDIR /workspace
 
@@ -37,3 +37,7 @@ RUN git lfs install --skip-repo && \
     pip3 install awscli && \
     # The base image ships with `pydantic==1.8.2` which is not working
     pip3 install -U --no-cache-dir pydantic==1.10.10
+
+# Initialize conda for bash shell to show environment in prompt
+RUN conda init bash && \
+    echo "conda activate axolotl-py${PYTHON_VERSION}" >> ~/.bashrc
diff --git a/docker/Dockerfile-uv-base b/docker/Dockerfile-uv-base
index 4b08e55f8..063882599 100644
--- a/docker/Dockerfile-uv-base
+++ b/docker/Dockerfile-uv-base
@@ -34,3 +34,6 @@ RUN uv pip install packaging setuptools wheel psutil \
     && uv pip install --no-build-isolation "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" \
     && uv pip install "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main" \
     && uv pip install awscli pydantic
+
+# Configure shell to display current virtual environment
+RUN echo 'export PS1="(axolotl-venv) \u@\h:\w\$ "' >> ~/.bashrc