Standardizing device names to the newer /device:<type>:<index> format by making all the device factories produce the new device names.

The python API would still support the legacy /<type>:<index> format so the C++ layer would accept both legacy and standardized names but the C++ layer would produce only new device names now.

PiperOrigin-RevId: 170758313

1 files changed, 37 insertions, 5 deletions

diff --git a/tensorflow/core/util/device_name_utils.cc b/tensorflow/core/util/device_name_utils.cc
index e667791c89..2d797c855a 100644
--- a/tensorflow/core/util/device_name_utils.cc
+++ b/tensorflow/core/util/device_name_utils.cc
@@ -104,11 +104,12 @@ string DeviceNameUtils::FullName(const string& job, int replica, int task,
   return DeviceName(job, replica, task, "/device:", type, id);
 }
 
-/* static */
-string DeviceNameUtils::LegacyName(const string& job, int replica, int task,
-                                   const string& type, int id) {
+namespace {
+string LegacyName(const string& job, int replica, int task, const string& type,
+                  int id) {
   return DeviceName(job, replica, task, "/", str_util::Lowercase(type), id);
 }
+}  // anonymous namespace
 
 bool DeviceNameUtils::ParseFullName(StringPiece fullname, ParsedName* p) {
   p->Clear();
@@ -185,6 +186,18 @@ bool DeviceNameUtils::ParseFullName(StringPiece fullname, ParsedName* p) {
 }
 
 /* static */
+string DeviceNameUtils::CanonicalizeDeviceName(StringPiece fullname) {
+  ParsedName parsed_name;
+  if (ParseLocalName(fullname, &parsed_name)) {
+    return ParsedNameToString(parsed_name);
+  }
+  if (ParseFullName(fullname, &parsed_name)) {
+    return ParsedNameToString(parsed_name);
+  }
+  return "";
+}
+
+/* static */
 string DeviceNameUtils::ParsedNameToString(const ParsedName& pn) {
   string buf;
   if (pn.has_job) strings::StrAppend(&buf, "/job:", pn.job);
@@ -338,8 +351,16 @@ bool DeviceNameUtils::IsSameAddressSpace(StringPiece src, StringPiece dst) {
 
 /* static */
 string DeviceNameUtils::LocalName(StringPiece type, int id) {
+  return strings::StrCat("/device:", type, ":", id);
+}
+
+namespace {
+// Returns the legacy local device name given its "type" and "id" (which is
+// '/device:type:id').
+string LegacyLocalName(StringPiece type, int id) {
   return strings::StrCat(type, ":", id);
 }
+}  // anonymous namespace
 
 /* static */
 string DeviceNameUtils::LocalName(StringPiece fullname) {
@@ -353,12 +374,14 @@ bool DeviceNameUtils::ParseLocalName(StringPiece name, ParsedName* p) {
   if (!ConsumeDeviceType(&name, &p->type)) {
     return false;
   }
+  p->has_type = true;
   if (!str_util::ConsumePrefix(&name, ":")) {
     return false;
   }
   if (!ConsumeNumber(&name, &p->id)) {
     return false;
   }
+  p->has_id = true;
   return name.empty();
 }
 
@@ -393,8 +416,17 @@ std::vector<string> DeviceNameUtils::GetNamesForDeviceMappings(
   if (pn.has_job && pn.has_replica && pn.has_task && pn.has_type && pn.has_id) {
     return {
         DeviceNameUtils::FullName(pn.job, pn.replica, pn.task, pn.type, pn.id),
-        DeviceNameUtils::LegacyName(pn.job, pn.replica, pn.task, pn.type,
-                                    pn.id)};
+        LegacyName(pn.job, pn.replica, pn.task, pn.type, pn.id)};
+  } else {
+    return {};
+  }
+}
+
+std::vector<string> DeviceNameUtils::GetLocalNamesForDeviceMappings(
+    const ParsedName& pn) {
+  if (pn.has_type && pn.has_id) {
+    return {DeviceNameUtils::LocalName(pn.type, pn.id),
+            LegacyLocalName(pn.type, pn.id)};
   } else {
     return {};
   }