From eda96d073f575ebac70e29c99088bd048636a5be Mon Sep 17 00:00:00 2001
From: Akshay Modi <nareshmodi@google.com>
Date: Fri, 20 Sep 2019 17:12:46 -0700
Subject: [PATCH] Use the SendTensor op by default
PiperOrigin-RevId: 270381552
---
tensorflow/core/common_runtime/eager/context.cc | 2 +-
.../python/data/experimental/kernel_tests/replicate_test.py | 2 --
tensorflow/python/eager/memory_tests/remote_memory_test.py | 3 ---
tensorflow/python/eager/remote.py | 3 ---
tensorflow/python/eager/remote_benchmarks_test.py | 2 --
5 files changed, 1 insertion(+), 11 deletions(-)
diff --git a/tensorflow/core/common_runtime/eager/context.cc b/tensorflow/core/common_runtime/eager/context.cc
index ef639d33396..3bb1954bfb4 100644
--- a/tensorflow/core/common_runtime/eager/context.cc
+++ b/tensorflow/core/common_runtime/eager/context.cc
@@ -680,7 +680,7 @@ Status EagerContext::InitializeRemoteMaster(
remote_contexts_ = remote_contexts;
use_send_tensor_rpc_ =
- ReadBoolFromEnvVar("TF_EAGER_REMOTE_USE_SEND_TENSOR_RPC", false);
+ ReadBoolFromEnvVar("TF_EAGER_REMOTE_USE_SEND_TENSOR_RPC", true);
local_unowned_device_manager_ = local_device_mgr;
local_device_manager_ = nullptr;
diff --git a/tensorflow/python/data/experimental/kernel_tests/replicate_test.py b/tensorflow/python/data/experimental/kernel_tests/replicate_test.py
index 46a9ac4df7e..a831655a2ae 100644
--- a/tensorflow/python/data/experimental/kernel_tests/replicate_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/replicate_test.py
@@ -17,7 +17,6 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
-import os
from absl.testing import parameterized
from tensorflow.core.protobuf import cluster_pb2
@@ -140,7 +139,6 @@ class RemoteReplicateTest(test_base.DatasetTestBase, parameterized.TestCase):
super(RemoteReplicateTest, self).__init__(methodName)
self._cached_server1 = server_lib.Server.create_local_server()
self._cached_server2 = server_lib.Server.create_local_server()
- os.environ["TF_EAGER_REMOTE_USE_SEND_TENSOR_RPC"] = "1"
self._cached_server1_target = self._cached_server1.target[len("grpc://"):]
self._cached_server2_target = self._cached_server2.target[len("grpc://"):]
self._device0 = "/job:%s/replica:0/task:0/device:CPU:0" % JOB_NAME
diff --git a/tensorflow/python/eager/memory_tests/remote_memory_test.py b/tensorflow/python/eager/memory_tests/remote_memory_test.py
index 0ef4d1f8065..71389fd35e5 100644
--- a/tensorflow/python/eager/memory_tests/remote_memory_test.py
+++ b/tensorflow/python/eager/memory_tests/remote_memory_test.py
@@ -18,8 +18,6 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
-import os
-
from tensorflow.python.eager import def_function
from tensorflow.python.eager import remote
from tensorflow.python.eager import test
@@ -36,7 +34,6 @@ class RemoteWorkerMemoryTest(test.TestCase):
super(RemoteWorkerMemoryTest, self).__init__(method)
# used for remote worker tests
- os.environ["TF_EAGER_REMOTE_USE_SEND_TENSOR_RPC"] = "1"
self._cached_server = server_lib.Server.create_local_server()
self._cached_server_target = self._cached_server.target[len("grpc://"):]
diff --git a/tensorflow/python/eager/remote.py b/tensorflow/python/eager/remote.py
index 71fefd56160..ce89a8a83b3 100644
--- a/tensorflow/python/eager/remote.py
+++ b/tensorflow/python/eager/remote.py
@@ -18,7 +18,6 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
-import os
from absl import logging
from tensorflow.core.protobuf.tensorflow_server_pb2 import ServerDef
@@ -128,8 +127,6 @@ def connect_to_cluster(cluster_spec_or_resolver,
cluster=cluster_def, job_name=job_name, task_index=task_index,
protocol=protocol)
- # TODO(nareshmodi): Make this default since it works in more situations.
- os.environ["TF_EAGER_REMOTE_USE_SEND_TENSOR_RPC"] = "1"
context.set_server_def(server_def)
if make_master_device_default and isinstance(
diff --git a/tensorflow/python/eager/remote_benchmarks_test.py b/tensorflow/python/eager/remote_benchmarks_test.py
index c3b76f49193..458a8b60667 100644
--- a/tensorflow/python/eager/remote_benchmarks_test.py
+++ b/tensorflow/python/eager/remote_benchmarks_test.py
@@ -27,7 +27,6 @@ from __future__ import division
from __future__ import print_function
import gc
-import os
import time
from six.moves import xrange # pylint: disable=redefined-builtin
@@ -64,7 +63,6 @@ class RemoteWorkerMicroBenchmarks(test.Benchmark):
def __init__(self):
# used for remote benchmarks
- os.environ["TF_EAGER_REMOTE_USE_SEND_TENSOR_RPC"] = "1"
self._cached_server1 = server_lib.Server.create_local_server()
self._cached_server_target1 = self._cached_server1.target[len("grpc://"):]
self._cached_server2 = server_lib.Server.create_local_server()