organize dockerfiles

Author: pan-x-c

.github/workflows/docker.yaml

@@ -47,11 +47,11 @@ jobs:
         with:
           context: .
           push: true
-          file: scripts/docker_for_megatron/Dockerfile
+          file: scripts/docker/Dockerfile.uv
           shm-size: 64g
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
-      
+
       # This step generates an artifact attestation for the image, which is an unforgeable statement about where and how it was built. It increases supply chain security for people who consume the image. For more information, see [Using artifact attestations to establish provenance for builds](/actions/security-guides/using-artifact-attestations-to-establish-provenance-for-builds).
       - name: Generate artifact attestation
         uses: actions/attest-build-provenance@v3

docs/sphinx_doc/source/tutorial/example_megatron.md

@@ -38,10 +38,10 @@ We provide a Docker setup to simplify environment management.
 #### Build the Docker Image
 
 
-Trinity-RFT provides a dedicated Dockerfile for Megatron-LM located at `scripts/docker_for_megatron/Dockerfile`. You can build the image using the following command:
+Trinity-RFT provides a dedicated Dockerfile for Megatron-LM located at `scripts/docker/Dockerfile.megatron`. You can build the image using the following command:
 
 ```bash
-docker build -f scripts/docker_for_megatron/Dockerfile -t trinity-rft-megatron:latest .
+docker build -f scripts/docker/Dockerfile.megatron -t trinity-rft-megatron:latest .
 ```
 
 > 💡 You can customize the Dockerfile before building — for example, to add pip mirrors or set API keys.

docs/sphinx_doc/source_zh/tutorial/example_megatron.md

@@ -43,11 +43,11 @@ pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \
 
 #### 构建 Docker 镜像
 
-Trinity-RFT 提供了专门用于 Megatron-LM 的 Dockerfile，位于 `scripts/docker_for_megatron/Dockerfile`。
+Trinity-RFT 提供了专门用于 Megatron-LM 的 Dockerfile，位于 `scripts/docker/Dockerfile.megatron`。
 可以使用以下命令构建镜像：
 
 ```bash
-docker build -f scripts/docker_for_megatron/Dockerfile -t trinity-rft-megatron:latest .
+docker build -f scripts/docker/Dockerfile.megatron -t trinity-rft-megatron:latest .
 ```
 
 > 💡 你可以在构建前自定义 Dockerfile —— 例如添加 pip 镜像源或设置 API 密钥。

scripts/docker_for_megatron/Dockerfile renamed to scripts/docker/Dockerfile.megatron

@@ -1,7 +1,7 @@
 # Build and run the docker image with the following command:
 #
 # cd <Trinity-RFT root dir>
-# docker build -f scripts/docker_for_megatron/Dockerfile -t trinity-rft-megatron:latest .
+# docker build -f scripts/docker/Dockerfile.megatron -t trinity-rft-megatron:latest .
 # docker run -it --gpus all --shm-size="64g" --rm -v $PWD:/workspace -v <root_path_of_data_and_checkpoints>:/data trinity-rft-megatron:latest
 
 

scripts/docker/Dockerfile.uv

@@ -0,0 +1,57 @@
+# Build and run the docker image with the following command:
+#
+# cd <Trinity-RFT root dir>
+# docker build -f scripts/docker/Dockerfile.uv -t trinity-rft:latest .
+# docker run -it --gpus all --shm-size="64g" --rm -v $PWD:/workspace -v <root_path_of_data_and_checkpoints>:/data trinity-rft:latest
+
+
+FROM nvcr.io/nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04
+
+WORKDIR /workspace
+
+RUN chmod 1777 /tmp && apt update && apt install -y \
+    build-essential \
+    curl git wget vim tmux net-tools \
+    python3 python3-pip python3-dev python3-packaging python3-venv \
+    libomp-dev infiniband-diags libibverbs-dev librdmacm-dev rdma-core perftest \
+    && rm -rf /var/lib/apt/lists/* \
+    && ln -sf /usr/bin/python3 /usr/bin/python \
+    && ln -sf /usr/bin/pip3 /usr/bin/pip
+
+# For Aliyun users: update pip mirror to aliyun to speed up pip install
+# RUN pip config set global.index-url http://mirrors.cloud.aliyuncs.com/pypi/simple/ \
+#    && pip config set install.trusted-host mirrors.cloud.aliyuncs.com
+
+ENV VIRTUAL_ENV=/opt/venv
+
+# copy the Trinity-RFT dir into the workspace
+COPY . .
+
+# Install uv
+RUN pip install uv && uv venv /opt/venv --python=python3.12
+
+# Install minimal Trinity-RFT
+RUN . /opt/venv/bin/activate && \
+    uv pip install -e.[mm,dev]
+
+# Install flash_attn and Megatron
+RUN . /opt/venv/bin/activate && \
+    uv pip install -e.[flash_attn] && \
+    uv pip install -e .[megatron] && \
+    uv pip install -v --no-build-isolation \
+    --config-settings="--build-option=--cpp_ext" \
+    --config-settings="--build-option=--cuda_ext" \
+    git+https://github.com/NVIDIA/apex.git
+
+# Set Env variables
+
+# WANDB
+# ENV WANDB_API_KEY=
+# ENV WANDB_BASE_URL=
+
+# LLM API
+# ENV OPENAI_API_KEY=
+# ENV DASH_API_KEY=
+
+ENTRYPOINT ["/bin/bash", "-c", "source /opt/venv/bin/activate && exec \"$@\"", "--"]
+CMD ["bash"]