update fedml dockerfile and jenkinsfile.

tony92151 · Oct 25, 2022 · 37c544b · 37c544b
1 parent dfa5319
commit 37c544b
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 deletions.
diff --git a/devops/scripts/build-fedml-docker.sh b/devops/scripts/build-fedml-docker.sh
@@ -14,8 +14,13 @@ OUTPUT_IMAGE=fedml/fedml:latest-torch1.12.1-cuda11.3-cudnn8-devel
 NVIDIA_BASE_IMAGE=nvidia/cuda:11.3.1-cudnn8-devel-ubuntu18.04
 PYTORCH_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cu113
 PYTORCH_GEOMETRIC_URL=https://data.pyg.org/whl/torch-1.12.0+cu113.html
+export FEDML_VERSION=`cat python/setup.py |grep version= |awk -F'=' '{print $2}' |awk -F',' '{print $1}'|awk -F'"' '{print $2}'`
+CURRENT_IMAGE=fedml/fedml:${FEDML_VERSION}-torch1.12.1-cuda11.3-cudnn8-devel
+
 cd ./docker
 bash build-docker.sh $ARCH $OS $DISTRO $PYTHON_VERSION $PYTORCH_VERSION $NCCL_VERSION $CUDA_VERSION \
      $OUTPUT_IMAGE $NVIDIA_BASE_IMAGE $PYTORCH_EXTRA_INDEX_URL $PYTORCH_GEOMETRIC_URL
 
+docker tag $OUTPUT_IMAGE $CURRENT_IMAGE
+
 cd $pwd
diff --git a/devops/scripts/push-fedml-docker.sh b/devops/scripts/push-fedml-docker.sh
@@ -1,4 +1,5 @@
 #!/bin/bash
 
 export FEDML_VERSION=`cat python/setup.py |grep version= |awk -F'=' '{print $2}' |awk -F',' '{print $1}'|awk -F'"' '{print $2}'`
-docker push fedml/fedml:latest-torch1.12.1-cuda11.3-cudnn8-devel
+docker push fedml/fedml:latest-torch1.12.1-cuda11.3-cudnn8-devel
+docker push fedml/fedml:${FEDML_VERSION}-torch1.12.1-cuda11.3-cudnn8-devel
diff --git a/doc/en/starter/installation.md b/doc/en/starter/installation.md
@@ -79,7 +79,7 @@ Please refer to the following commands and remember to change `WORKSPACE` to you
 
 **(1) Pull the Docker image and prepare the docker environment**
 ```
-FEDML_DOCKER_IMAGE=fedml/fedml:cuda-11.4.0-devel-ubuntu20.04
+FEDML_DOCKER_IMAGE=fedml/fedml:latest-torch1.12.1-cuda11.3-cudnn8-devel
 docker pull FEDML_DOCKER_IMAGE
 
 # if you want to use GPUs in your host OS, please follow this link: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#docker
@@ -92,7 +92,7 @@ sudo chmod 777 /var/run/docker.sock
 **(2) Run Docker with interactive mode**
 
 ```
-FEDML_DOCKER_IMAGE=fedml/fedml:cuda-11.4.0-devel-ubuntu20.04
+FEDML_DOCKER_IMAGE=fedml/fedml:latest-torch1.12.1-cuda11.3-cudnn8-devel
 WORKSPACE=/home/chaoyanghe/sourcecode/FedML_startup/FedML
 
 docker run -t -i -v $WORKSPACE:$WORKSPACE --shm-size=64g --ulimit nofile=65535 --ulimit memlock=-1 --privileged \