update container (#69)

* Update Dockerfile * Update docker_rdma_latest.yml * Update Dockerfile * Update docker_rdma_latest.yml * Update Dockerfile * code_clean * new naming
dptech-corp · Feb 9, 2025 · 47704e7 · 47704e7
1 parent b7a1ac5
commit 47704e7
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 77 deletions.
diff --git a/.github/workflows/docker_rdma.yml b/.github/workflows/docker_rdma.yml
diff --git a/.github/workflows/docker_rdma_latest.yml b/.github/workflows/docker_rdma_latest.yml
@@ -4,7 +4,6 @@ on:
   push:
     branches:
     - main
-    - cuda12
 
 jobs:
   docker:
@@ -60,5 +59,5 @@ jobs:
           context: ./docker/rdma/
           push: true
           tags: |
-            dptechnology/unicore:latest-pytorch2.3.0-cuda12.4-rdma
-            dp-ve-registry-cn-beijing.cr.volces.com/dplc/unicore:latest-pytorch2.3.0-cuda12.4-rdma
+            dptechnology/unicore:latest-2407-pytorch2.4.0-cuda12.5-rdma
+            dp-ve-registry-cn-beijing.cr.volces.com/dplc/unicore:latest-2407-pytorch2.4.0-cuda12.5-rdma
diff --git a/docker/rdma/Dockerfile b/docker/rdma/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvcr.io/nvidia/pytorch:24.03-py3
+FROM nvcr.io/nvidia/pytorch:24.07-py3
 
 RUN APT_INSTALL="apt-get install -y --no-install-recommends" && \
     rm -rf /var/lib/apt/lists/* \
@@ -55,23 +55,15 @@ RUN APT_INSTALL="apt-get install -y --no-install-recommends" && \
 # # pytorch
 # # ------------------------------------------------------------------
 ENV TORCH_CUDA_ARCH_LIST "7.0;7.5;8.0;9.0"
-
-RUN pip3 install --no-cache-dir --upgrade sentry-sdk requests ninja typing packaging && rm -rf ~/.cache/pip
+RUN pip3 install --no-cache-dir --upgrade sentry-sdk requests ninja typing packaging wandb rdkit tokenizers lmdb ml-collections tensorboardX && rm -rf ~/.cache/pip
 
 RUN cd /tmp && \
     git clone https://github.com/dptech-corp/Uni-Core && \
     cd Uni-Core && \
     python setup.py install --enable-cuda-ext && \
     rm -rf  /tmp/* && rm -rf ~/.cache/pip
 
-RUN pip3 install --no-cache-dir tokenizers lmdb biopython ml-collections timeout-decorator urllib3 tree dm-tree && rm -rf ~/.cache/pip 
-
-RUN cd /tmp && \
-    git clone https://github.com/Dao-AILab/flash-attention.git && \
-    cd flash-attention && \
-    git checkout v2.6.1 && \
-    MAX_JOBS=4 python setup.py install && \
-    rm -rf  /tmp/* && rm -rf ~/.cache/pip
+RUN pip3 install --no-cache-dir  biopython timeout-decorator urllib3 tree dm-tree && rm -rf ~/.cache/pip 
 
 RUN ldconfig && \
     apt-get clean && \