Azure ci pipeline for distributed environment tests (#5881)

orttraining/orttraining/test/python/orttraining_distributed_tests.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import sys
+import argparse
+
+from _test_commons import run_subprocess
+
+import logging
+
+logging.basicConfig(
+    format="%(asctime)s %(name)s [%(levelname)s] - %(message)s",
+    level=logging.DEBUG)
+log = logging.getLogger("DistributedTests")
+
+def parse_arguments():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--cwd", help="Path to the current working directory")
+    return parser.parse_args()
+
+def main():
+    import torch
+    ngpus = torch.cuda.device_count()
+
+    if ngpus < 2:
+        raise RuntimeError("Cannot run distributed tests with less than 2 gpus.")
+
+    args = parse_arguments()
+    cwd = args.cwd
+
+    log.info("Running distributed tests pipeline")
+
+    # TODO: Add distributed test suite here.
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tools/ci_build/github/azure-pipelines/orttraining-linux-gpu-distributed-test-ci-pipeline.yml

@@ -0,0 +1,48 @@
+trigger: none
+
+jobs:
+- job: Onnxruntime_Linux_GPU_Distributed_Test
+
+  timeoutInMinutes: 120
+  pool: 'Linux-Multi-GPU-V100'
+
+  steps:
+  - checkout: self
+    clean: true
+    submodules: recursive
+
+  - template: templates/run-docker-build-steps.yml
+    parameters:
+      RunDockerBuildArgs: |
+        -o ubuntu16.04 -d gpu -r $(Build.BinariesDirectory) \
+        -t onnxruntime_distributed_tests_image \
+        -x " \
+          --config RelWithDebInfo \
+          --enable_training \
+          --update --build \
+          --build_wheel \
+          " \
+        -m
+      DisplayName: 'Build'
+
+    # all distributed tests
+  - script: |
+      docker run \
+        --gpus all \
+        --shm-size=1024m \
+        --rm \
+        --volume $(Build.SourcesDirectory):/onnxruntime_src \
+        --volume $(Build.BinariesDirectory):/build \
+        onnxruntime_distributed_tests_image \
+          /build/RelWithDebInfo/launch_test.py \
+            --cmd_line_with_args "python orttraining_distributed_tests.py" \
+            --cwd /build/RelWithDebInfo \
+    displayName: 'Run orttraining_distributed_tests.py'
+    condition: succeededOrFailed()
+    timeoutInMinutes: 30
+  
+  - template: templates/component-governance-component-detection-steps.yml
+    parameters:
+      condition: 'succeeded'
+        
+  - template: templates/clean-agent-build-directory-step.yml

tools/ci_build/github/linux/docker/Dockerfile.ubuntu_gpu

@@ -2,10 +2,11 @@ FROM nvidia/cuda:10.1-cudnn7-devel-ubuntu16.04
 
 ARG PYTHON_VERSION=3.6
 ARG BUILD_EXTR_PAR
+ARG INSTALL_DEPS_EXTRA_ARGS
 
 ADD scripts /tmp/scripts
 RUN /tmp/scripts/install_ubuntu.sh -p $PYTHON_VERSION && \
-    /tmp/scripts/install_deps.sh -p $PYTHON_VERSION -d gpu -x "$BUILD_EXTR_PAR" && \
+    /tmp/scripts/install_deps.sh -p $PYTHON_VERSION -d gpu -x "$BUILD_EXTR_PAR" $INSTALL_DEPS_EXTRA_ARGS && \
     rm -rf /tmp/scripts
 
 WORKDIR /root

tools/ci_build/github/linux/docker/scripts/install_deps.sh

@@ -2,13 +2,15 @@
 set -e -x
 
 SCRIPT_DIR="$( dirname "${BASH_SOURCE[0]}" )"
+INSTALL_DEPS_DISTRIBUTED_SETUP=false
 
-while getopts p:d:x: parameter_Option
+while getopts p:d:x:m parameter_Option
 do case "${parameter_Option}"
 in
 p) PYTHON_VER=${OPTARG};;
 d) DEVICE_TYPE=${OPTARG};;
 x) BUILD_EXTR_PAR=${OPTARG};;
+m) INSTALL_DEPS_DISTRIBUTED_SETUP=true
 esac
 done
 
@@ -113,7 +115,7 @@ if [ $DEVICE_TYPE = "gpu" ]; then
     if [[ $BUILD_EXTR_PAR = *--enable_training* ]]; then
       ${PYTHON_EXE} -m pip install -r ${0/%install_deps.sh/training\/requirements.txt}
 
-      if [[ $BUILD_EXTR_PAR = *--enable_training_python_frontend_e2e_tests* || $BUILD_EXTR_PAR = *enable_training_pipeline_e2e_tests* ]]; then
+      if [[ $BUILD_EXTR_PAR = *--enable_training_python_frontend_e2e_tests* || $BUILD_EXTR_PAR = *enable_training_pipeline_e2e_tests* || $INSTALL_DEPS_DISTRIBUTED_SETUP = true ]]; then
         source ${0/%install_deps.sh/install_openmpi.sh}
       fi
     fi

tools/ci_build/github/linux/run_dockerbuild.sh

@@ -5,10 +5,11 @@ SCRIPT_DIR="$( dirname "${BASH_SOURCE[0]}" )"
 SOURCE_ROOT=$(realpath $SCRIPT_DIR/../../../../)
 CUDA_VER=cuda10.1-cudnn7.6
 YOCTO_VERSION="4.19"
+INSTALL_DEPS_DISTRIBUTED_SETUP=false
 ALLOW_RELEASED_ONNX_OPSET_ONLY_ENV="ALLOW_RELEASED_ONNX_OPSET_ONLY="$ALLOW_RELEASED_ONNX_OPSET_ONLY
 echo "ALLOW_RELEASED_ONNX_OPSET_ONLY environment variable is set as "$ALLOW_RELEASED_ONNX_OPSET_ONLY_ENV
 
-while getopts c:o:d:r:p:x:a:v:y:t: parameter_Option
+while getopts c:o:d:r:p:x:a:v:y:t:m parameter_Option
 do case "${parameter_Option}"
 in
 #android, ubuntu16.04, manylinux2010, ubuntu18.04, CentOS7
@@ -31,6 +32,8 @@ y) YOCTO_VERSION=${OPTARG};;
 # an additional name for the resulting docker image (created with "docker tag")
 # this is useful for referencing the image outside of this script
 t) EXTRA_IMAGE_TAG=${OPTARG};;
+# install distributed setup dependencies
+m) INSTALL_DEPS_DISTRIBUTED_SETUP=true;;
 esac
 done
 
@@ -86,8 +89,9 @@ else
         if [ $CUDA_VER = "cuda9.1-cudnn7.1" ]; then
             DOCKER_FILE=Dockerfile.ubuntu_gpu_cuda9
         fi
+        [[ $INSTALL_DEPS_DISTRIBUTED_SETUP = true ]] && INSTALL_DEPS_EXTRA_ARGS="-m" || INSTALL_DEPS_EXTRA_ARGS=""
         $GET_DOCKER_IMAGE_CMD --repository "onnxruntime-$IMAGE" \
-            --docker-build-args="--build-arg BUILD_USER=onnxruntimedev --build-arg BUILD_UID=$(id -u) --build-arg PYTHON_VERSION=${PYTHON_VER} --build-arg BUILD_EXTR_PAR=\"${BUILD_EXTR_PAR}\"" \
+            --docker-build-args="--build-arg BUILD_USER=onnxruntimedev --build-arg BUILD_UID=$(id -u) --build-arg PYTHON_VERSION=${PYTHON_VER} --build-arg BUILD_EXTR_PAR=\"${BUILD_EXTR_PAR}\" --build-arg INSTALL_DEPS_EXTRA_ARGS=${INSTALL_DEPS_EXTRA_ARGS}" \
             --dockerfile $DOCKER_FILE --context .
     elif [ $BUILD_DEVICE = "tensorrt" ]; then
         # TensorRT container release 20.07