From b924c02b0c5a80f4152c0351d8876a6c011647b5 Mon Sep 17 00:00:00 2001
From: FightingZhen <295632982@qq.com>
Date: Sun, 28 Dec 2025 10:42:44 +0800
Subject: [PATCH] Set Megatron related environment variable with ENV in Ascend
 dockerfile

---
 docker/ascend/Dockerfile.ascend_8.2.rc1_a2 | 3 ++-
 docker/ascend/Dockerfile.ascend_8.2.rc1_a3 | 3 ++-
 docker/ascend/Dockerfile.ascend_8.3.rc1_a2 | 3 ++-
 docker/ascend/Dockerfile.ascend_8.3.rc1_a3 | 3 ++-
 4 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/docker/ascend/Dockerfile.ascend_8.2.rc1_a2 b/docker/ascend/Dockerfile.ascend_8.2.rc1_a2
index ceefe23cb06..1c4a5fec4ef 100644
--- a/docker/ascend/Dockerfile.ascend_8.2.rc1_a2
+++ b/docker/ascend/Dockerfile.ascend_8.2.rc1_a2
@@ -41,11 +41,12 @@ RUN ARCH=$(uname -m) && \
     cd vllm-ascend && pip install -v -e . && cd .. && \
     # Install MindSpeed & Megatron
     pip install -e MindSpeed && \
-    echo "export PYTHONPATH=\$PYTHONPATH:/Megatron-LM" >> ~/.bashrc && \
     # Clear extra files
     rm -rf /tmp/* /var/tmp/* && \
     pip cache purge
 
+ENV PYTHONPATH="/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}"
+
 # Prepare and install verl (update frequently)
 RUN git clone --depth 1 https://github.com/volcengine/verl.git && \
     cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \
diff --git a/docker/ascend/Dockerfile.ascend_8.2.rc1_a3 b/docker/ascend/Dockerfile.ascend_8.2.rc1_a3
index 2b16e075f3a..c8d7febc884 100644
--- a/docker/ascend/Dockerfile.ascend_8.2.rc1_a3
+++ b/docker/ascend/Dockerfile.ascend_8.2.rc1_a3
@@ -41,11 +41,12 @@ RUN ARCH=$(uname -m) && \
     cd vllm-ascend && pip install -v -e . && cd .. && \
     # Install MindSpeed & Megatron
     pip install -e MindSpeed && \
-    echo "export PYTHONPATH=\$PYTHONPATH:/Megatron-LM" >> ~/.bashrc && \
     # Clear extra files
     rm -rf /tmp/* /var/tmp/* && \
     pip cache purge
 
+ENV PYTHONPATH="/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}"
+
 # Prepare and install verl (update frequently)
 RUN git clone --depth 1 https://github.com/volcengine/verl.git && \
     cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \
diff --git a/docker/ascend/Dockerfile.ascend_8.3.rc1_a2 b/docker/ascend/Dockerfile.ascend_8.3.rc1_a2
index 200e7a05f35..946504224b4 100644
--- a/docker/ascend/Dockerfile.ascend_8.3.rc1_a2
+++ b/docker/ascend/Dockerfile.ascend_8.3.rc1_a2
@@ -41,7 +41,6 @@ RUN ARCH=$(uname -m) && \
     cd vllm-ascend && pip install -v -e . && cd .. && \
     # Install MindSpeed & Megatron
     pip install -e MindSpeed && \
-    echo "export PYTHONPATH=\$PYTHONPATH:/Megatron-LM" >> ~/.bashrc && \
     # Remove existing triton or triton-ascend installed by some third-party packages
     pip uninstall -y triton triton-ascend && \
     # Install mbridge
@@ -50,6 +49,8 @@ RUN ARCH=$(uname -m) && \
     rm -rf /tmp/* /var/tmp/* && \
     pip cache purge
 
+ENV PYTHONPATH="/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}"
+
 # Prepare and install verl (update frequently)
 RUN git clone --depth 1 https://github.com/volcengine/verl.git && \
     cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \
diff --git a/docker/ascend/Dockerfile.ascend_8.3.rc1_a3 b/docker/ascend/Dockerfile.ascend_8.3.rc1_a3
index bbf7de87bbe..4e7fec47e5b 100644
--- a/docker/ascend/Dockerfile.ascend_8.3.rc1_a3
+++ b/docker/ascend/Dockerfile.ascend_8.3.rc1_a3
@@ -41,7 +41,6 @@ RUN ARCH=$(uname -m) && \
     cd vllm-ascend && pip install -v -e . && cd .. && \
     # Install MindSpeed & Megatron
     pip install -e MindSpeed && \
-    echo "export PYTHONPATH=\$PYTHONPATH:/Megatron-LM" >> ~/.bashrc && \
     # Remove existing triton or triton-ascend installed by some third-party packages
     pip uninstall -y triton triton-ascend && \
     # Install mbridge
@@ -50,6 +49,8 @@ RUN ARCH=$(uname -m) && \
     rm -rf /tmp/* /var/tmp/* && \
     pip cache purge
 
+ENV PYTHONPATH="/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}"
+
 # Prepare and install verl (update frequently)
 RUN git clone --depth 1 https://github.com/volcengine/verl.git && \
     cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \