From 27d3af01fb1e714a8ea64c5bebd743a02370f461 Mon Sep 17 00:00:00 2001
From: Jinjing Zhou <VoVAllen@users.noreply.github.com>
Date: Tue, 8 Mar 2022 19:27:00 +0800
Subject: [PATCH] Fix benchmark problem (#3762)

* fix

* fix bench

* fix

* Update bench_rgcn.py

* try fix

* fix wrong fix

* Update bench_rgcn.py
---
 Jenkinsfile                                     | 7 ++++---
 benchmarks/benchmarks/model_speed/bench_rgcn.py | 6 +++++-
 benchmarks/benchmarks/utils.py                  | 2 +-
 benchmarks/scripts/publish.sh                   | 4 ++--
 4 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/Jenkinsfile b/Jenkinsfile
index af516ab41e45..092bb1f84474 100644
--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -107,9 +107,10 @@ pipeline {
   stages {
     stage('Regression Test Trigger') {
       agent {
-        kubernetes {
-          yamlFile 'docker/pods/ci-lint.yaml'
-          defaultContainer 'dgl-ci-lint'
+        docker {
+            label 'linux-benchmark-node'
+            image 'dgllib/dgl-ci-lint'
+            alwaysPull true
         }
       }
       when { triggeredBy 'IssueCommentCause' }
diff --git a/benchmarks/benchmarks/model_speed/bench_rgcn.py b/benchmarks/benchmarks/model_speed/bench_rgcn.py
index b01d60b07f22..61dca9f2208c 100644
--- a/benchmarks/benchmarks/model_speed/bench_rgcn.py
+++ b/benchmarks/benchmarks/model_speed/bench_rgcn.py
@@ -41,7 +41,11 @@ def forward(self, g, h, r, norm):
 def track_time(data, use_type_count):
     # args
     if data == 'aifb':
-        num_bases = -1
+        if dgl.__version__.startswith("0.8"):
+            num_bases = None
+        else:
+            num_bases = -1
+            
         l2norm = 0.
     elif data == 'am':
         num_bases = 40
diff --git a/benchmarks/benchmarks/utils.py b/benchmarks/benchmarks/utils.py
index 6cdb7afc0e7b..1258a4f71dea 100644
--- a/benchmarks/benchmarks/utils.py
+++ b/benchmarks/benchmarks/utils.py
@@ -251,7 +251,7 @@ def __getitem__(self, idx):
 
 
 def load_nowplaying_rs():
-    import torchtext
+    import torchtext.legacy as torchtext
     # follow examples/pytorch/pinsage/README to create nowplaying_rs.pkl
     name = 'nowplaying_rs.pkl'
     dataset_dir = os.path.join(os.getcwd(), 'dataset')
diff --git a/benchmarks/scripts/publish.sh b/benchmarks/scripts/publish.sh
index 7845513d3ada..1cf7d2002843 100644
--- a/benchmarks/scripts/publish.sh
+++ b/benchmarks/scripts/publish.sh
@@ -55,14 +55,14 @@ if [[ $DEVICE == "cpu" ]]; then
         --rm \
         $DOCKER_MOUNT_OPT \
         $DOCKER_ENV_OPT \
-        --shm-size="4g" \
+        --shm-size="16g" \
         --hostname=$MACHINE -dit public.ecr.aws/s1o7b3d9/benchmakrk_pyg_dgl:cu111_torch181_pyg170 /bin/bash
 else
     docker run --name dgl-reg \
         --rm --gpus all \
         $DOCKER_MOUNT_OPT \
         $DOCKER_ENV_OPT \
-        --shm-size="4g" \
+        --shm-size="16g" \
         --hostname=$MACHINE -dit public.ecr.aws/s1o7b3d9/benchmakrk_pyg_dgl:cu111_torch181_pyg170 /bin/bash
 fi