Add TPUHardwareFeature in the tpu topology proto and populate it based on the tpu version (variant) during configuration.

pineapplejuice233 · tensorflower-gardener · commit cfb560ab1a68 · 2022-03-03T14:40:33.000-08:00
PiperOrigin-RevId: 432286170
diff --git a/tensorflow/core/protobuf/tpu/topology.proto b/tensorflow/core/protobuf/tpu/topology.proto
@@ -1,8 +1,26 @@
 syntax = "proto3";
 
+package tensorflow.tpu;
+
 option cc_enable_arenas = true;
 
-package tensorflow.tpu;
+// Describes features of a tpu.
+message TPUHardwareFeature {
+  // Embedding feature of a tpu.
+  enum EmbeddingFeature {
+    // No embedding lookup accelerator available on the tpu.
+    UNSUPPORTED = 0;
+    // Embedding lookup accelerator V1. The embedding lookup operation can only
+    // be placed at the beginning of computation. Only one instance of embedding
+    // lookup layer is allowed.
+    V1 = 1;
+    // Embedding lookup accelerator V2. The embedding lookup operation can be
+    // placed anywhere of the computation. Multiple instances of embedding
+    // lookup layer is allowed.
+    V2 = 2;
+  }
+  EmbeddingFeature embedding_feature = 1;
+}
 
 // Describes the geometry of a TPU mesh.
 message TopologyProto {
@@ -24,4 +42,7 @@ message TopologyProto {
   // in the TPU mesh topology. Each entry [task, device, axis] gives the
   // `axis`-th coordinate in the topology of a task/device pair.
   repeated int32 device_coordinates = 4;
+
+  // TPU supported features.
+  TPUHardwareFeature tpu_hardware_feature = 5;
 }
diff --git a/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py b/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py
@@ -17,6 +17,7 @@
 import collections
 import re
 
+from tensorflow.core.protobuf.tpu import topology_pb2
 from tensorflow.python.distribute.cluster_resolver import cluster_resolver
 from tensorflow.python.framework import config as framework_config
 from tensorflow.python.framework import errors
@@ -217,6 +218,8 @@ def __init__(self,
     else:
       self._coordinator_address = coordinator_address
 
+    self._tpu_topology = None
+
   def __enter__(self):
     self._cloud_tpu_client.enter()
 
@@ -393,6 +396,18 @@ def num_accelerators(self,
       }
     return {'TPU': 0}
 
+  def set_tpu_topology(self, serialized_tpu_topology):
+    """Sets the tpu topology info stored in this resolver."""
+    self._tpu_topology = topology_pb2.TopologyProto()
+    self._tpu_topology.ParseFromString(serialized_tpu_topology)
+
+  @property
+  def tpu_hardware_feature(self):
+    """Returns the tpu topology info stored."""
+    if self._tpu_topology is None:
+      return self._tpu_topology
+    return self._tpu_topology.tpu_hardware_feature
+
   @property
   def environment(self):
     """Returns the current environment which TensorFlow is running in."""
diff --git a/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver_test.py b/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver_test.py
@@ -19,6 +19,7 @@
 import six
 from six.moves.urllib.error import URLError
 
+from tensorflow.core.protobuf.tpu import topology_pb2
 from tensorflow.python import framework
 from tensorflow.python.client import session
 from tensorflow.python.distribute.cluster_resolver.tpu import tpu_cluster_resolver as resolver
@@ -29,6 +30,7 @@
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.training import server_lib
 from tensorflow.python.util import compat
+
 mock = test.mock
 
 try:
@@ -706,6 +708,17 @@ def testLocalTpuResolver(self):
     cr = resolver.TPUClusterResolver(tpu='local')
     self.assertEqual(cr.get_master(), '')
 
+  def testTpuTopology(self):
+    cluster_resolver = resolver.TPUClusterResolver(tpu='local')
+    self.assertIsNone(cluster_resolver._tpu_topology)
+
+    # Test set with tpu topology proto.
+    cluster_resolver.set_tpu_topology(
+        serialized_tpu_topology=topology_pb2.TopologyProto(
+            mesh_shape=[1, 1, 1, 1]).SerializeToString())
+    self.assertIsInstance(cluster_resolver.tpu_hardware_feature,
+                          topology_pb2.TPUHardwareFeature)
+
 
 if __name__ == '__main__':
   test.main()
diff --git a/tensorflow/python/tpu/tpu_strategy_util.py b/tensorflow/python/tpu/tpu_strategy_util.py
@@ -141,6 +141,7 @@ def _tpu_init_fn():
 
   logging.info("Finished initializing TPU system.")
   tpu_topology = topology.Topology(serialized=serialized_topology)
+  cluster_resolver.set_tpu_topology(serialized_topology)
   _INITIALIZED_TPU_SYSTEMS[tpu_name] = tpu_topology
 
   return tpu_topology
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.cluster_resolver.-t-p-u-cluster-resolver.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.cluster_resolver.-t-p-u-cluster-resolver.pbtxt
@@ -15,6 +15,10 @@ tf_class {
     name: "task_type"
     mtype: "<type \'property\'>"
   }
+  member {
+    name: "tpu_hardware_feature"
+    mtype: "<type \'property\'>"
+  }
   member_method {
     name: "__init__"
     argspec: "args=[\'self\', \'tpu\', \'zone\', \'project\', \'job_name\', \'coordinator_name\', \'coordinator_address\', \'credentials\', \'service\', \'discovery_url\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'worker\', \'None\', \'None\', \'default\', \'None\', \'None\'], "
@@ -47,4 +51,8 @@ tf_class {
     name: "num_accelerators"
     argspec: "args=[\'self\', \'task_type\', \'task_id\', \'config_proto\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
   }
+  member_method {
+    name: "set_tpu_topology"
+    argspec: "args=[\'self\', \'serialized_tpu_topology\'], varargs=None, keywords=None, defaults=None"
+  }
 }
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.cluster_resolver.-t-p-u-cluster-resolver.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.cluster_resolver.-t-p-u-cluster-resolver.pbtxt
@@ -15,6 +15,10 @@ tf_class {
     name: "task_type"
     mtype: "<type \'property\'>"
   }
+  member {
+    name: "tpu_hardware_feature"
+    mtype: "<type \'property\'>"
+  }
   member_method {
     name: "__init__"
     argspec: "args=[\'self\', \'tpu\', \'zone\', \'project\', \'job_name\', \'coordinator_name\', \'coordinator_address\', \'credentials\', \'service\', \'discovery_url\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'worker\', \'None\', \'None\', \'default\', \'None\', \'None\'], "
@@ -47,4 +51,8 @@ tf_class {
     name: "num_accelerators"
     argspec: "args=[\'self\', \'task_type\', \'task_id\', \'config_proto\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
   }
+  member_method {
+    name: "set_tpu_topology"
+    argspec: "args=[\'self\', \'serialized_tpu_topology\'], varargs=None, keywords=None, defaults=None"
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -15,6 +15,10 @@ tf_class {`
`15`	`15`	`name: "task_type"`
`16`	`16`	`mtype: "<type \'property\'>"`
`17`	`17`	`}`
	`18`	`+ member {`
	`19`	`+ name: "tpu_hardware_feature"`
	`20`	`+ mtype: "<type \'property\'>"`
	`21`	`+ }`
`18`	`22`	`member_method {`
`19`	`23`	`name: "__init__"`
`20`	`24`	`argspec: "args=[\'self\', \'tpu\', \'zone\', \'project\', \'job_name\', \'coordinator_name\', \'coordinator_address\', \'credentials\', \'service\', \'discovery_url\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'worker\', \'None\', \'None\', \'default\', \'None\', \'None\'], "`
`@@ -47,4 +51,8 @@ tf_class {`
`47`	`51`	`name: "num_accelerators"`
`48`	`52`	`argspec: "args=[\'self\', \'task_type\', \'task_id\', \'config_proto\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "`
`49`	`53`	`}`
	`54`	`+ member_method {`
	`55`	`+ name: "set_tpu_topology"`
	`56`	`+ argspec: "args=[\'self\', \'serialized_tpu_topology\'], varargs=None, keywords=None, defaults=None"`
	`57`	`+ }`
`50`	`58`	`}`