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

@@ -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",
     ],
 )
 

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(

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"),