Merge pull request deepspeedai#20 from rraminen/PR_Update_Dockerfile

Update dockerfile.rocm for cupy and h5py
jithunnair-amd · Jun 28, 2021 · 0cf3306 · 0cf3306
2 parents 1850f88 + f6c79ae
commit 0cf3306
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm
@@ -124,7 +124,8 @@ RUN pip install psutil \
                 numpy \
                 sklearn \
                 scikit-learn \
-                mpi4py
+                mpi4py \
+                h5py 
 
 ##############################################################################
 ## SSH daemon port inside container cannot conflict with host OS port
@@ -150,6 +151,15 @@ RUN cat /etc/ssh/sshd_config > ${STAGE_DIR}/sshd_config && \
 RUN rm -rf /usr/lib/python3/dist-packages/yaml && \
         rm -rf /usr/lib/python3/dist-packages/PyYAML-*
 
+##############################################################################
+## CuPy installation
+###############################################################################
+RUN git clone https://github.com/ROCmSoftwarePlatform/cupy ${STAGE_DIR}/cupy
+RUN cd ${STAGE_DIR}/cupy && \
+        git submodule update --init && \
+        CUPY_INSTALL_USE_HIP=1 ROCM_HOME=/opt/rocm pip install -e . --no-cache-dir -vvvv
+RUN rm -rf ${STAGE_DIR}/cupy
+
 ##############################################################################
 ## Add deepspeed user
 ###############################################################################