From 9450048acb141151895a4405c7d5d118165677de Mon Sep 17 00:00:00 2001
From: Qiao Longfei <qiaolongfei@baidu.com>
Date: Thu, 29 Nov 2018 16:48:43 +0800
Subject: [PATCH] add PADDLE_ENABLE_REMOTE_PREFETCH to enable remote prefetch
 test=develop

---
 paddle/fluid/API.spec                                       | 2 +-
 python/paddle/fluid/layers/nn.py                            | 6 ++++--
 python/paddle/fluid/tests/unittests/dist_ctr.py             | 2 ++
 python/paddle/fluid/tests/unittests/test_dist_transpiler.py | 2 ++
 .../fluid/tests/unittests/test_lookup_remote_table_op.py    | 1 +
 5 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/API.spec b/paddle/fluid/API.spec
index 59605a7e16..c40f603341 100644
--- a/paddle/fluid/API.spec
+++ b/paddle/fluid/API.spec
@@ -50,7 +50,7 @@ paddle.fluid.initializer.MSRAInitializer.__init__ ArgSpec(args=['self', 'uniform
 paddle.fluid.initializer.force_init_on_cpu ArgSpec(args=[], varargs=None, keywords=None, defaults=None)
 paddle.fluid.initializer.init_on_cpu ArgSpec(args=[], varargs='args', keywords='kwds', defaults=None)
 paddle.fluid.layers.fc ArgSpec(args=['input', 'size', 'num_flatten_dims', 'param_attr', 'bias_attr', 'act', 'is_test', 'name'], varargs=None, keywords=None, defaults=(1, None, None, None, False, None))
-paddle.fluid.layers.embedding ArgSpec(args=['input', 'size', 'is_sparse', 'is_distributed', 'padding_idx', 'param_attr', 'dtype', 'remote_prefetch'], varargs=None, keywords=None, defaults=(False, False, None, None, 'float32', False))
+paddle.fluid.layers.embedding ArgSpec(args=['input', 'size', 'is_sparse', 'is_distributed', 'padding_idx', 'param_attr', 'dtype'], varargs=None, keywords=None, defaults=(False, False, None, None, 'float32'))
 paddle.fluid.layers.dynamic_lstm ArgSpec(args=['input', 'size', 'h_0', 'c_0', 'param_attr', 'bias_attr', 'use_peepholes', 'is_reverse', 'gate_activation', 'cell_activation', 'candidate_activation', 'dtype', 'name'], varargs=None, keywords=None, defaults=(None, None, None, None, True, False, 'sigmoid', 'tanh', 'tanh', 'float32', None))
 paddle.fluid.layers.dynamic_lstmp ArgSpec(args=['input', 'size', 'proj_size', 'param_attr', 'bias_attr', 'use_peepholes', 'is_reverse', 'gate_activation', 'cell_activation', 'candidate_activation', 'proj_activation', 'dtype', 'name'], varargs=None, keywords=None, defaults=(None, None, True, False, 'sigmoid', 'tanh', 'tanh', 'tanh', 'float32', None))
 paddle.fluid.layers.dynamic_gru ArgSpec(args=['input', 'size', 'param_attr', 'bias_attr', 'is_reverse', 'gate_activation', 'candidate_activation', 'h_0'], varargs=None, keywords=None, defaults=(None, None, False, 'sigmoid', 'tanh', None))
diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py
index 2051a1ea01..29a0de29dc 100644
--- a/python/paddle/fluid/layers/nn.py
+++ b/python/paddle/fluid/layers/nn.py
@@ -287,8 +287,7 @@ def embedding(input,
               is_distributed=False,
               padding_idx=None,
               param_attr=None,
-              dtype='float32',
-              remote_prefetch=False):
+              dtype='float32'):
     """
     **Embedding Layer**
 
@@ -327,6 +326,9 @@ def embedding(input,
     """
 
     helper = LayerHelper('embedding', **locals())
+    remote_prefetch = False
+    if os.environ.get('PADDLE_ENABLE_REMOTE_PREFETCH'):
+        remote_prefetch = True
     if remote_prefetch:
         assert is_sparse is True and is_distributed is False
     w = helper.create_parameter(
diff --git a/python/paddle/fluid/tests/unittests/dist_ctr.py b/python/paddle/fluid/tests/unittests/dist_ctr.py
index 088f16a8ac..d7ddafdb17 100644
--- a/python/paddle/fluid/tests/unittests/dist_ctr.py
+++ b/python/paddle/fluid/tests/unittests/dist_ctr.py
@@ -16,11 +16,13 @@ from __future__ import print_function
 
 import paddle
 import paddle.fluid as fluid
+import os
 
 import dist_ctr_reader
 from test_dist_base import TestDistRunnerBase, runtime_main
 
 IS_SPARSE = True
+os.environ['PADDLE_ENABLE_REMOTE_PREFETCH'] = "1"
 
 # Fix seed for test
 fluid.default_startup_program().random_seed = 1
diff --git a/python/paddle/fluid/tests/unittests/test_dist_transpiler.py b/python/paddle/fluid/tests/unittests/test_dist_transpiler.py
index 3a905d4a02..26c596ba54 100644
--- a/python/paddle/fluid/tests/unittests/test_dist_transpiler.py
+++ b/python/paddle/fluid/tests/unittests/test_dist_transpiler.py
@@ -792,6 +792,8 @@ class TestNCCL2Transpile(TranspilerTest):
 # test for remote prefetch
 class TestRemoteLookupTable(TestDistLookupTableBase):
     def net_conf(self):
+        import os
+        os.environ['PADDLE_ENABLE_REMOTE_PREFETCH'] = "1"
         self.network_with_table(
             is_sparse=True, is_distributed=False, remote_prefetch=True)
 
diff --git a/python/paddle/fluid/tests/unittests/test_lookup_remote_table_op.py b/python/paddle/fluid/tests/unittests/test_lookup_remote_table_op.py
index b46e61d2ef..47830fb56b 100644
--- a/python/paddle/fluid/tests/unittests/test_lookup_remote_table_op.py
+++ b/python/paddle/fluid/tests/unittests/test_lookup_remote_table_op.py
@@ -174,6 +174,7 @@ class TestListenAndServOp(unittest.TestCase):
                     self.assertTrue((result_array[i] == id).all())
 
     def test_lookup_remote_table(self):
+        os.environ['PADDLE_ENABLE_REMOTE_PREFETCH'] = "1"
         # run pserver on CPU in sync mode
         p0 = self._start_pserver(0, False, True, run_pserver)
         self._wait_ps_ready(p0.pid)
-- 
GitLab