From b8eea223427479a111127583e1a7769dae0900c2 Mon Sep 17 00:00:00 2001
From: Tianqi Chen <tqchen@users.noreply.github.com>
Date: Wed, 13 Nov 2019 20:49:34 -0800
Subject: [PATCH] [CI][DOCKER] Add ONNX runtime dep (#4314)

* [DOCKER] Add ONNX runtime dep

* Improve ci script
---
 Jenkinsfile                              | 22 ++++++++++++----------
 docker/install/ubuntu_install_onnx.sh    |  1 +
 tests/scripts/task_python_frontend.sh    |  1 +
 tests/scripts/task_python_integration.sh |  2 ++
 4 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/Jenkinsfile b/Jenkinsfile
index 17f9a5669d039..e016e11ad1da9 100644
--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -64,6 +64,8 @@ docker_run = 'docker/bash.sh'
 // timeout in minutes
 max_time = 120
 
+workspace = 'workspace/exec_${EXECUTOR_NUMBER}'
+
 // initialize source codes
 def init_git() {
   checkout scm
@@ -86,7 +88,7 @@ def init_git_win() {
 stage("Sanity Check") {
   timeout(time: max_time, unit: 'MINUTES') {
     node('CPU') {
-      ws('workspace/tvm/sanity') {
+      ws('${workspace}/tvm/sanity') {
         init_git()
         sh "${docker_run} ${ci_lint}  ./tests/scripts/task_lint.sh"
       }
@@ -134,7 +136,7 @@ def unpack_lib(name, libs) {
 stage('Build') {
   parallel 'BUILD: GPU': {
     node('GPUBUILD') {
-      ws('workspace/tvm/build-gpu') {
+      ws('${workspace}/tvm/build-gpu') {
         init_git()
         sh """
            mkdir -p build
@@ -182,7 +184,7 @@ stage('Build') {
   },
   'BUILD: CPU': {
     node('CPU') {
-      ws('workspace/tvm/build-cpu') {
+      ws('${workspace}/tvm/build-cpu') {
         init_git()
         sh """
            mkdir -p build
@@ -213,7 +215,7 @@ stage('Build') {
   },
   'BUILD : i386': {
     node('CPU') {
-      ws('workspace/tvm/build-i386') {
+      ws('${workspace}/tvm/build-i386') {
         init_git()
         sh """
            mkdir -p build
@@ -238,7 +240,7 @@ stage('Build') {
 stage('Unit Test') {
   parallel 'python3: GPU': {
     node('TensorCore') {
-      ws('workspace/tvm/ut-python-gpu') {
+      ws('${workspace}/tvm/ut-python-gpu') {
         init_git()
         unpack_lib('gpu', tvm_multilib)
         timeout(time: max_time, unit: 'MINUTES') {
@@ -250,7 +252,7 @@ stage('Unit Test') {
   },
   'python3: i386': {
     node('CPU') {
-      ws('workspace/tvm/ut-python-i386') {
+      ws('${workspace}/tvm/ut-python-i386') {
         init_git()
         unpack_lib('i386', tvm_multilib)
         timeout(time: max_time, unit: 'MINUTES') {
@@ -263,7 +265,7 @@ stage('Unit Test') {
   },
   'java: GPU': {
     node('GPU') {
-      ws('workspace/tvm/ut-java') {
+      ws('${workspace}/tvm/ut-java') {
         init_git()
         unpack_lib('gpu', tvm_multilib)
         timeout(time: max_time, unit: 'MINUTES') {
@@ -277,7 +279,7 @@ stage('Unit Test') {
 stage('Integration Test') {
   parallel 'topi: GPU': {
     node('GPU') {
-      ws('workspace/tvm/topi-python-gpu') {
+      ws('${workspace}/tvm/topi-python-gpu') {
         init_git()
         unpack_lib('gpu', tvm_multilib)
         timeout(time: max_time, unit: 'MINUTES') {
@@ -288,7 +290,7 @@ stage('Integration Test') {
   },
   'frontend: GPU': {
     node('GPU') {
-      ws('workspace/tvm/frontend-python-gpu') {
+      ws('${workspace}/tvm/frontend-python-gpu') {
         init_git()
         unpack_lib('gpu', tvm_multilib)
         timeout(time: max_time, unit: 'MINUTES') {
@@ -299,7 +301,7 @@ stage('Integration Test') {
   },
   'legacy: GPU': {
     node('GPU') {
-      ws('workspace/tvm/legacy-python-gpu') {
+      ws('${workspace}/tvm/legacy-python-gpu') {
         init_git()
         unpack_lib('gpu', tvm_multilib)
         timeout(time: max_time, unit: 'MINUTES') {
diff --git a/docker/install/ubuntu_install_onnx.sh b/docker/install/ubuntu_install_onnx.sh
index 54210b83f4d69..a915ca02c05a1 100755
--- a/docker/install/ubuntu_install_onnx.sh
+++ b/docker/install/ubuntu_install_onnx.sh
@@ -22,6 +22,7 @@ set -o pipefail
 
 # fix to certain version for now
 pip3 install onnx==1.5.0
+pip3 install onnxruntime==1.0.0
 
 # torch depends on a number of other packages, but unhelpfully, does
 # not expose that in the wheel!!!
diff --git a/tests/scripts/task_python_frontend.sh b/tests/scripts/task_python_frontend.sh
index 78ad8e8ca73a1..7a7bcacf61408 100755
--- a/tests/scripts/task_python_frontend.sh
+++ b/tests/scripts/task_python_frontend.sh
@@ -21,6 +21,7 @@ set -u
 
 export PYTHONPATH=nnvm/python:python:topi/python
 # to avoid openblas threading error
+export TVM_BIND_THREADS=0
 export OMP_NUM_THREADS=1
 
 # Rebuild cython
diff --git a/tests/scripts/task_python_integration.sh b/tests/scripts/task_python_integration.sh
index ab0f3a8523079..ebbcf2106617b 100755
--- a/tests/scripts/task_python_integration.sh
+++ b/tests/scripts/task_python_integration.sh
@@ -21,6 +21,8 @@ set -u
 
 export PYTHONPATH=python:topi/python:apps/extension/python
 export LD_LIBRARY_PATH="build:${LD_LIBRARY_PATH:-}"
+export TVM_BIND_THREADS=0
+export TVM_NUM_THREADS=2
 
 rm -rf python/tvm/*.pyc python/tvm/*/*.pyc python/tvm/*/*/*.pyc