[Backend] Turn to official dlpack for Tensorflow (dmlc#1511)

* Turn to official dlpack * fix * fix
zeta1999 · May 11, 2020 · 6d04885 · 6d04885
1 parent 0285021
commit 6d04885
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 8 deletions.
diff --git a/Jenkinsfile b/Jenkinsfile
@@ -101,7 +101,8 @@ pipeline {
       agent { 
         docker {
           label "linux-c52x-node"
-          image "dgllib/dgl-ci-lint" 
+          image "dgllib/dgl-ci-lint"  
+          alwaysPull true
         }
       }
       steps {
@@ -121,6 +122,7 @@ pipeline {
             docker {
               label "linux-c52x-node"
               image "dgllib/dgl-ci-cpu:conda" 
+              alwaysPull true
             }
           }
           steps {
@@ -138,6 +140,7 @@ pipeline {
               label "linux-c52x-node"
               image "dgllib/dgl-ci-gpu:conda"
               args "-u root"
+              alwaysPull true
             }
           }
           steps {
@@ -173,6 +176,7 @@ pipeline {
             docker { 
               label "linux-c52x-node"
               image "dgllib/dgl-ci-cpu:conda"
+              alwaysPull true
             }
           }
           steps {
@@ -200,6 +204,7 @@ pipeline {
             docker {
               label "linux-c52x-node"
               image "dgllib/dgl-ci-cpu:conda" 
+              alwaysPull true
             }
           }
           stages {
@@ -221,6 +226,7 @@ pipeline {
               label "linux-gpu-node"
               image "dgllib/dgl-ci-gpu:conda" 
               args "--runtime nvidia"
+              alwaysPull true
             }
           }
           stages {
@@ -241,6 +247,7 @@ pipeline {
             docker {
               label "linux-c52x-node"
               image "dgllib/dgl-ci-cpu:conda" 
+              alwaysPull true
             }
           }
           stages {
@@ -292,6 +299,7 @@ pipeline {
               label "linux-gpu-node"
               image "dgllib/dgl-ci-gpu:conda"
               args "--runtime nvidia"
+              alwaysPull true
             }
           }
           stages {
@@ -318,6 +326,7 @@ pipeline {
             docker {
               label "linux-c52x-node"
               image "dgllib/dgl-ci-cpu:conda" 
+              alwaysPull true
             }
           }
           stages {
@@ -344,6 +353,7 @@ pipeline {
               label "linux-gpu-node" 
               image "dgllib/dgl-ci-gpu:conda"
               args "--runtime nvidia"
+              alwaysPull true
             }
           }
           stages {

diff --git a/docker/install/conda_env/tensorflow_cpu.yml b/docker/install/conda_env/tensorflow_cpu.yml
@@ -3,8 +3,7 @@ dependencies:
   - python=3.6.9
   - pip
   - pip:
-    - tensorflow==2.2.0rc1
-    # - tf-nightly==2.2.0.dev20200327
+    - tensorflow==2.2.0
     - tfdlpack
     - pytest
     - nose

diff --git a/docker/install/conda_env/tensorflow_gpu.yml b/docker/install/conda_env/tensorflow_gpu.yml
@@ -3,8 +3,7 @@ dependencies:
   - python=3.6.9
   - pip
   - pip:
-    - tensorflow==2.2.0rc1
-    # - tf-nightly==2.2.0.dev20200327
+    - tensorflow==2.2.0
     - tfdlpack-gpu
     - pytest
     - nose

diff --git a/python/dgl/backend/tensorflow/tensor.py b/python/dgl/backend/tensorflow/tensor.py
@@ -13,15 +13,15 @@
 from ... import kernel as K
 from ...function.base import TargetCode
 
-if os.getenv("USE_OFFICIAL_TFDLPACK", False):
+if not os.getenv("USE_TFDLPACK", False):
     if LooseVersion(tf.__version__) < LooseVersion("2.2.0"):
         raise RuntimeError("DGL requires tensorflow>=2.2.0 for the official DLPack support.")
 
     def zerocopy_to_dlpack(input):
         return tf.experimental.dlpack.to_dlpack(input)
 
     def zerocopy_from_dlpack(dlpack_tensor):
-        # TODO(Jinjing): Tensorflow requires memory to be 64-bit aligned. We check the
+        # TODO(Jinjing): Tensorflow requires memory to be 64-bytes aligned. We check the
         #   alignment and make a copy if needed. The functionality is better in TF's main repo.
         aligned = nd.from_dlpack(dlpack_tensor).to_dlpack(64)
         return tf.experimental.dlpack.from_dlpack(aligned)

diff --git a/tests/scripts/build_dgl.sh b/tests/scripts/build_dgl.sh
@@ -22,7 +22,7 @@ rm -rf _download
 
 pushd build
 cmake $CMAKE_VARS ..
-make -j4
+make -j8
 popd
 
 pushd python