[CI] separate distributed tests from torch cpu tests (#4313)

* [CI] separate distributed tests from torch cpu tests

* remove TF related env

Jenkinsfile

@@ -62,6 +62,14 @@ def unit_test_linux(backend, dev) {
   }
 }
 
+def unit_distributed_linux(backend, dev) {
+  init_git()
+  unpack_lib("dgl-${dev}-linux", dgl_linux_libs)
+  timeout(time: 30, unit: 'MINUTES') {
+    sh "bash tests/scripts/task_distributed_test.sh ${backend} ${dev}"
+  }
+}
+
 def unit_test_cugraph(backend, dev) {
   init_git()
   unpack_lib("dgl-${dev}-linux", dgl_linux_libs)
@@ -452,6 +460,28 @@ pipeline {
                 }
               }
             }
+            stage('Distributed') {
+              agent {
+                docker {
+                  label "linux-cpu-node"
+                  image "dgllib/dgl-ci-cpu:cu101_v220629"
+                  args "--shm-size=4gb"
+                  alwaysPull true
+                }
+              }
+              stages {
+                stage('Distributed Torch CPU Unit test') {
+                  steps {
+                    unit_distributed_linux('pytorch', 'cpu')
+                  }
+                }
+              }
+              post {
+                always {
+                  cleanWs disableDeferredWipeout: true, deleteDirs: true
+                }
+              }
+            }
             stage('PyTorch Cugraph GPU') {
               agent {
                 docker {

tests/scripts/task_distributed_test.sh

+#!/bin/bash
+
+. /opt/conda/etc/profile.d/conda.sh
+
+function fail {
+    echo FAIL: $@
+    exit -1
+}
+
+function usage {
+    echo "Usage: $0 backend device"
+}
+
+if [ $# -ne 2 ]; then
+    usage
+    fail "Error: must specify backend and device"
+fi
+
+[ $1 == "pytorch" ] || fail "Distrbuted tests run on pytorch backend only."
+[ $2 == "cpu" ] || fail "Distrbuted tests run on cpu only."
+
+export DGLBACKEND=$1
+export DGLTESTDEV=$2
+export DGL_LIBRARY_PATH=${PWD}/build
+export PYTHONPATH=tests:${PWD}/python:$PYTHONPATH
+export DGL_DOWNLOAD_DIR=${PWD}
+unset TORCH_ALLOW_TF32_CUBLAS_OVERRIDE
+
+export CUDA_VISIBLE_DEVICES=-1
+
+conda activate ${DGLBACKEND}-ci
+
+python3 -m pip install pytest psutil pyyaml pydantic pandas rdflib ogb filelock || fail "pip install"
+
+export PYTHONUNBUFFERED=1
+export OMP_NUM_THREADS=1
+export DMLC_LOG_DEBUG=1
+
+python3 -m pytest -v --capture=tee-sys --junitxml=pytest_distributed.xml --durations=100 tests/distributed/*.py || fail "distributed"
+
+PYTHONPATH=tools:$PYTHONPATH python3 -m pytest -v --capture=tee-sys --junitxml=pytest_tools.xml --durations=100 tests/tools/*.py || fail "tools"

tests/scripts/task_unit_test.sh

@@ -36,13 +36,3 @@ conda activate ${DGLBACKEND}-ci
 python3 -m pip install pytest psutil pyyaml pydantic pandas rdflib ogb || fail "pip install"
 python3 -m pytest -v --junitxml=pytest_compute.xml --durations=100 tests/compute || fail "compute"
 python3 -m pytest -v --junitxml=pytest_backend.xml --durations=100 tests/$DGLBACKEND || fail "backend-specific"
-
-
-export PYTHONUNBUFFERED=1
-export OMP_NUM_THREADS=1
-export DMLC_LOG_DEBUG=1
-if [ $2 != "gpu" ] && [ $DGLBACKEND == "pytorch" ]; then
-    python3 -m pip install filelock
-    python3 -m pytest -v --capture=tee-sys --junitxml=pytest_distributed.xml --durations=100 tests/distributed/*.py || fail "distributed"
-    PYTHONPATH=tools:$PYTHONPATH python3 -m pytest -v --capture=tee-sys --junitxml=pytest_tools.xml --durations=100 tests/tools/*.py || fail "tools"
-fi