PyTorch 1.8 (Python 3.8 & CUDA 11.1)

Author: hephaex

python-pytorch/Dockerfile.1.8-py38-cuda11.1

@@ -1,16 +1,20 @@
 FROM lablup/common-base:py38-cuda11.1
-# Install PyTorch and MXNet
-ENV PYTORCH_VERSION=1.7.1
-ENV TORCHVISION_VERSION=0.8.2
-ENV TORCHAUDIO_VERSION=0.7.2
-ENV TORCHTEXT_VERSION=0.8.1
+
+# Install PyTorch
+ENV PYTORCH_VERSION=1.8.1
+ENV TORCHVISION_VERSION=0.9.1
+ENV TORCHAUDIO_VERSION=0.8.1
+ENV TORCHTEXT_VERSION=0.9.1
 ENV TENSORBOARDX_VERSION=2.1
+ENV MXNET_VERSION=1.7.0
 
 RUN python3 -m pip uninstall -y torch && \
     python3 -m pip install --no-cache-dir \
-    	https://download.pytorch.org/whl/cu110/torch-${PYTORCH_VERSION}%2Bcu110-cp38-cp38-linux_x86_64.whl \
-	https://download.pytorch.org/whl/cu110/torchvision-${TORCHVISION_VERSION}%2Bcu110-cp38-cp38-linux_x86_64.whl \
+        https://download.pytorch.org/whl/cu111/torch-${PYTORCH_VERSION}%2Bcu111-cp38-cp38-linux_x86_64.whl \
+        https://download.pytorch.org/whl/cu111/torchvision-${TORCHVISION_VERSION}%2Bcu111-cp38-cp38-linux_x86_64.whl \
         https://download.pytorch.org/whl/torchaudio-${TORCHAUDIO_VERSION}-cp38-cp38-linux_x86_64.whl \
+	https://download.pytorch.org/whl/cu111/torchcsprng-0.2.1%2Bcu111-cp38-cp38-linux_x86_64.whl \
+	https://download.pytorch.org/whl/torchserve-0.3.0-py2.py3-none-any.whl \
         https://download.pytorch.org/whl/torchtext-${TORCHTEXT_VERSION}-cp38-cp38-linux_x86_64.whl && \
     python3 -m pip install --no-cache-dir tensorboardX==${TENSORBOARDX_VERSION}
 
@@ -28,23 +32,25 @@ RUN python3 -m pip install --extra-index-url \
 RUN ldconfig /usr/local/cuda/targets/x86_64-linux/lib/stubs && \
     HOROVOD_GPU_ALLREDUCE=NCCL HOROVOD_GPU_BROADCAST=NCCL \
     HOROVOD_WITH_TENSORFLOW=1 HOROVOD_WITH_PYTORCH=1 HOROVOD_WITHOUT_MXNET=1 \
-    pip install --no-cache-dir horovod==0.21.1 && \
+    pip install --no-cache-dir horovod==0.21.3 && \
     ldconfig
 
 RUN python3 -m pip install --no-cache-dir \
     	mpi4py==3.0.3 \
-        mlflow==1.12.1 \
-	nni==1.9 \
+        mlflow==1.15.0 \
+	nni==2.1 \
 	scikit-nni==0.2.1
 
 RUN apt autoclean && \
     rm -rf /var/lib/apt/lists/* && \	
     rm -rf /root/.cache && \
     rm -rf /tmp/*
+
 COPY ./service-defs /etc/backend.ai/service-defs
+COPY ./runner-scripts/bootstrap.sh runner-scripts/setup_multinode.py /opt/container/
 
 # Install ipython kernelspec
-Run python3 -m ipykernel install --display-name "PyTorch 1.7.1 on Python 3.8 (CUDA 11.1)" && \
+Run python3 -m ipykernel install --display-name "PyTorch 1.8.1 on Python 3.8 (CUDA 11.1)" && \
     cat /usr/local/share/jupyter/kernels/python3/kernel.json
 
 # Backend.AI specifics

python-pytorch/runner-scripts/bootstrap.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+BAI_MULTINODE_CONFIG_TF=$(/opt/backend.ai/bin/python /opt/container/setup_multinode.py)
+if [ -z "$BAI_MULTINODE_CONFIG_TF" ];
+then
+	echo "";
+else
+	echo ${BAI_MULTINODE_CONFIG_TF}
+	export TF_CONFIG="${BAI_MULTINODE_CONFIG_TF}"
+fi

python-pytorch/runner-scripts/setup_multinode.py

@@ -0,0 +1,19 @@
+import os
+import json
+
+if 'BACKENDAI_CLUSTER_HOST' in os.environ: # Start mutli-instance setup.
+    env = {}
+    env['cluster'] = {}
+    env['cluster']['worker'] = []
+    for container in os.environ['BACKENDAI_CLUSTER_HOSTS'].split(","):
+        env['cluster']['worker'].append(container + ":2220")
+    env['task'] = {}
+    if os.environ['BACKENDAI_CLUSTER_ROLE'] == 'main':
+        env['task']['type'] = "worker" # Was chief. but recent TF choose first worker as chief.
+        env['task']["index"] = str(int(os.environ['BACKENDAI_CLUSTER_IDX']) - 1) # Index starts from 0
+    else:
+        env['task']['type'] = "worker"
+        env['task']["index"] = os.environ['BACKENDAI_CLUSTER_IDX'] 
+    print(json.dumps(env))
+else:
+    print("")