From 0ea67f633b6272c809244925c9aa05f16adf2401 Mon Sep 17 00:00:00 2001 From: Ran Chen Date: Tue, 16 Jun 2020 15:17:40 -0700 Subject: [PATCH] Try to deduce job, replica and task from config.list_logical_devices() again We used to assume that job is localhost, which may not be true in a multi worker environment. list_logical_devices() isn't always safe since it may return remote devices as well, but we don't have list_local_devices() yet, and we're already using list_logical_devices() in MirorredStrategy (when user doesn't explicit pass devices), so this shouldn't make things worse. An alternative is to re-consider whether we need device_util.canonicalize at all, but there are multiple usage of it. This requires further effort. PiperOrigin-RevId: 316768135 Change-Id: Idb2653497e0c4e746a61c44b4d452ab1d052d014 --- tensorflow/python/distribute/BUILD | 8 ++++ tensorflow/python/distribute/device_util.py | 13 ++++- .../python/distribute/device_util_test.py | 48 +++++++++++++++---- 3 files changed, 59 insertions(+), 10 deletions(-) diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD index a3655699669..e39631d634f 100644 --- a/tensorflow/python/distribute/BUILD +++ b/tensorflow/python/distribute/BUILD @@ -121,9 +121,17 @@ cuda_py_test( name = "device_util_test", srcs = ["device_util_test.py"], deps = [ + ":combinations", ":device_util", + ":multi_worker_test_base", + ":multi_worker_util", + "//tensorflow/core:protos_all_py", "//tensorflow/python:client_testlib", + "//tensorflow/python:extra_py_tests_deps", "//tensorflow/python:framework_ops", + "//tensorflow/python:training_server_lib", + "//tensorflow/python/eager:context", + "@absl_py//absl/testing:parameterized", ], ) diff --git a/tensorflow/python/distribute/device_util.py b/tensorflow/python/distribute/device_util.py index 7f32ed39aed..0aee34e33ae 100644 --- a/tensorflow/python/distribute/device_util.py +++ b/tensorflow/python/distribute/device_util.py @@ -19,6 +19,7 @@ from __future__ import division from __future__ import print_function from tensorflow.python.eager import context +from tensorflow.python.framework import config from tensorflow.python.framework import device as tf_device from tensorflow.python.framework import ops @@ -55,8 +56,16 @@ def canonicalize(d, default=None): result = tf_device.DeviceSpec( replica=0, task=0, device_type="CPU", device_index=0) if ops.executing_eagerly_outside_functions(): - # The default job is localhost if eager execution is enabled - result = result.replace(job="localhost") + # Try to deduce job, replica and task in case it's in a multi worker setup. + # TODO(b/151452748): Using list_logical_devices is not always safe since it + # may return remote devices as well, but we're already doing this elsewhere. + host_cpu = tf_device.DeviceSpec.from_string( + config.list_logical_devices("CPU")[0].name) + if host_cpu.job: + result = result.make_merged_spec(host_cpu) + else: + # The default job is localhost if eager execution is enabled + result = result.replace(job="localhost") if default: # Overrides any defaults with values from the default device if given. result = result.make_merged_spec( diff --git a/tensorflow/python/distribute/device_util_test.py b/tensorflow/python/distribute/device_util_test.py index 2f0d7ed3b31..df53fe0288a 100644 --- a/tensorflow/python/distribute/device_util_test.py +++ b/tensorflow/python/distribute/device_util_test.py @@ -18,16 +18,27 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from absl.testing import parameterized + +from tensorflow.core.protobuf import tensorflow_server_pb2 +from tensorflow.python.distribute import combinations from tensorflow.python.distribute import device_util +from tensorflow.python.distribute import multi_worker_test_base from tensorflow.python.eager import context from tensorflow.python.framework import ops -from tensorflow.python.framework import test_util from tensorflow.python.platform import test +from tensorflow.python.training import server_lib -class DeviceUtilTest(test.TestCase): +class DeviceUtilTest(test.TestCase, parameterized.TestCase): - @test_util.run_deprecated_v1 + def setUp(self): + super(DeviceUtilTest, self).setUp() + context._reset_context() # pylint: disable=protected-access + + @combinations.generate( + combinations.combine(mode="graph") + ) def testCurrentDeviceWithGlobalGraph(self): with ops.device("/cpu:0"): self.assertEqual(device_util.current(), "/device:CPU:0") @@ -51,11 +62,16 @@ class DeviceUtilTest(test.TestCase): self.assertEqual(device_util.current(), "/job:localhost/replica:0/task:0/device:CPU:0") - @test_util.run_deprecated_v1 - def testCanonicalizeWithoutDefaultDevice(self): - self.assertEqual( - device_util.canonicalize("/cpu:0"), - "/replica:0/task:0/device:CPU:0") + @combinations.generate(combinations.combine(mode=["graph", "eager"])) + def testCanonicalizeWithoutDefaultDevice(self, mode): + if mode == "graph": + self.assertEqual( + device_util.canonicalize("/cpu:0"), + "/replica:0/task:0/device:CPU:0") + else: + self.assertEqual( + device_util.canonicalize("/cpu:0"), + "/job:localhost/replica:0/task:0/device:CPU:0") self.assertEqual( device_util.canonicalize("/job:worker/cpu:0"), "/job:worker/replica:0/task:0/device:CPU:0") @@ -63,6 +79,22 @@ class DeviceUtilTest(test.TestCase): device_util.canonicalize("/job:worker/task:1/cpu:0"), "/job:worker/replica:0/task:1/device:CPU:0") + @combinations.generate(combinations.combine(mode=["eager"])) + def testCanonicalizeWithoutDefaultDeviceCollectiveEnabled(self): + cluster_spec = server_lib.ClusterSpec( + multi_worker_test_base.create_cluster_spec( + has_chief=False, num_workers=1, num_ps=0, has_eval=False)) + server_def = tensorflow_server_pb2.ServerDef( + cluster=cluster_spec.as_cluster_def(), + job_name="worker", + task_index=0, + protocol="grpc", + port=0) + context.context().enable_collective_ops(server_def) + self.assertEqual( + device_util.canonicalize("/cpu:0"), + "/job:worker/replica:0/task:0/device:CPU:0") + def testCanonicalizeWithDefaultDevice(self): self.assertEqual( device_util.canonicalize("/job:worker/task:1/cpu:0", default="/gpu:0"),