• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 #include "tensorflow/core/util/device_name_utils.h"
17 
18 #include "tensorflow/core/lib/core/errors.h"
19 #include "tensorflow/core/lib/strings/str_util.h"
20 #include "tensorflow/core/lib/strings/strcat.h"
21 #include "tensorflow/core/platform/logging.h"
22 
23 namespace tensorflow {
24 
IsAlpha(char c)25 static bool IsAlpha(char c) {
26   return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z');
27 }
28 
IsAlphaNum(char c)29 static bool IsAlphaNum(char c) { return IsAlpha(c) || (c >= '0' && c <= '9'); }
30 
31 // Returns true iff "in" is a valid job name.
IsJobName(StringPiece in)32 static bool IsJobName(StringPiece in) {
33   if (in.empty()) return false;
34   if (!IsAlpha(in[0])) return false;
35   for (size_t i = 1; i < in.size(); ++i) {
36     if (!(IsAlphaNum(in[i]) || in[i] == '_')) return false;
37   }
38   return true;
39 }
40 
41 // Returns true and fills in "*job" iff "*in" starts with a job name.
ConsumeJobName(StringPiece * in,string * job)42 static bool ConsumeJobName(StringPiece* in, string* job) {
43   if (in->empty()) return false;
44   if (!IsAlpha((*in)[0])) return false;
45   size_t i = 1;
46   for (; i < in->size(); ++i) {
47     const char c = (*in)[i];
48     if (c == '/') break;
49     if (!(IsAlphaNum(c) || c == '_')) {
50       return false;
51     }
52   }
53   job->assign(in->data(), i);
54   in->remove_prefix(i);
55   return true;
56 }
57 
58 // Returns true and fills in "*device_type" iff "*in" starts with a device type
59 // name.
ConsumeDeviceType(StringPiece * in,string * device_type)60 static bool ConsumeDeviceType(StringPiece* in, string* device_type) {
61   if (in->empty()) return false;
62   if (!IsAlpha((*in)[0])) return false;
63   size_t i = 1;
64   for (; i < in->size(); ++i) {
65     const char c = (*in)[i];
66     if (c == '/' || c == ':') break;
67     if (!(IsAlphaNum(c) || c == '_')) {
68       return false;
69     }
70   }
71   device_type->assign(in->data(), i);
72   in->remove_prefix(i);
73   return true;
74 }
75 
76 // Returns true and fills in "*val" iff "*in" starts with a decimal
77 // number.
ConsumeNumber(StringPiece * in,int * val)78 static bool ConsumeNumber(StringPiece* in, int* val) {
79   uint64 tmp;
80   if (str_util::ConsumeLeadingDigits(in, &tmp)) {
81     *val = tmp;
82     return true;
83   } else {
84     return false;
85   }
86 }
87 
88 // Returns a fully qualified device name given the parameters.
DeviceName(const string & job,int replica,int task,const string & device_prefix,const string & device_type,int id)89 static string DeviceName(const string& job, int replica, int task,
90                          const string& device_prefix, const string& device_type,
91                          int id) {
92   CHECK(IsJobName(job)) << job;
93   CHECK_LE(0, replica);
94   CHECK_LE(0, task);
95   CHECK(!device_type.empty());
96   CHECK_LE(0, id);
97   return strings::StrCat("/job:", job, "/replica:", replica, "/task:", task,
98                          device_prefix, device_type, ":", id);
99 }
100 
101 /* static */
FullName(const string & job,int replica,int task,const string & type,int id)102 string DeviceNameUtils::FullName(const string& job, int replica, int task,
103                                  const string& type, int id) {
104   return DeviceName(job, replica, task, "/device:", type, id);
105 }
106 
107 namespace {
LegacyName(const string & job,int replica,int task,const string & type,int id)108 string LegacyName(const string& job, int replica, int task, const string& type,
109                   int id) {
110   return DeviceName(job, replica, task, "/", str_util::Lowercase(type), id);
111 }
112 }  // anonymous namespace
113 
ParseFullName(StringPiece fullname,ParsedName * p)114 bool DeviceNameUtils::ParseFullName(StringPiece fullname, ParsedName* p) {
115   p->Clear();
116   if (fullname == "/") {
117     return true;
118   }
119   while (!fullname.empty()) {
120     bool progress = false;
121     if (str_util::ConsumePrefix(&fullname, "/job:")) {
122       p->has_job = !str_util::ConsumePrefix(&fullname, "*");
123       if (p->has_job && !ConsumeJobName(&fullname, &p->job)) {
124         return false;
125       }
126       progress = true;
127     }
128     if (str_util::ConsumePrefix(&fullname, "/replica:")) {
129       p->has_replica = !str_util::ConsumePrefix(&fullname, "*");
130       if (p->has_replica && !ConsumeNumber(&fullname, &p->replica)) {
131         return false;
132       }
133       progress = true;
134     }
135     if (str_util::ConsumePrefix(&fullname, "/task:")) {
136       p->has_task = !str_util::ConsumePrefix(&fullname, "*");
137       if (p->has_task && !ConsumeNumber(&fullname, &p->task)) {
138         return false;
139       }
140       progress = true;
141     }
142     if (str_util::ConsumePrefix(&fullname, "/device:")) {
143       p->has_type = !str_util::ConsumePrefix(&fullname, "*");
144       if (p->has_type && !ConsumeDeviceType(&fullname, &p->type)) {
145         return false;
146       }
147       if (!str_util::ConsumePrefix(&fullname, ":")) {
148         p->has_id = false;
149       } else {
150         p->has_id = !str_util::ConsumePrefix(&fullname, "*");
151         if (p->has_id && !ConsumeNumber(&fullname, &p->id)) {
152           return false;
153         }
154       }
155       progress = true;
156     }
157 
158     // Handle legacy naming convention for cpu and gpu.
159     if (str_util::ConsumePrefix(&fullname, "/cpu:") ||
160         str_util::ConsumePrefix(&fullname, "/CPU:")) {
161       p->has_type = true;
162       p->type = "CPU";  // Treat '/cpu:..' as uppercase '/device:CPU:...'
163       p->has_id = !str_util::ConsumePrefix(&fullname, "*");
164       if (p->has_id && !ConsumeNumber(&fullname, &p->id)) {
165         return false;
166       }
167       progress = true;
168     }
169     if (str_util::ConsumePrefix(&fullname, "/gpu:") ||
170         str_util::ConsumePrefix(&fullname, "/GPU:")) {
171       p->has_type = true;
172       p->type = "GPU";  // Treat '/gpu:..' as uppercase '/device:GPU:...'
173       p->has_id = !str_util::ConsumePrefix(&fullname, "*");
174       if (p->has_id && !ConsumeNumber(&fullname, &p->id)) {
175         return false;
176       }
177       progress = true;
178     }
179 
180     if (!progress) {
181       return false;
182     }
183   }
184   return true;
185 }
186 
187 namespace {
188 
CompleteName(const DeviceNameUtils::ParsedName & parsed_basename,DeviceNameUtils::ParsedName * parsed_name)189 void CompleteName(const DeviceNameUtils::ParsedName& parsed_basename,
190                   DeviceNameUtils::ParsedName* parsed_name) {
191   if (!parsed_name->has_job) {
192     parsed_name->job = parsed_basename.job;
193     parsed_name->has_job = true;
194   }
195   if (!parsed_name->has_replica) {
196     parsed_name->replica = parsed_basename.replica;
197     parsed_name->has_replica = true;
198   }
199   if (!parsed_name->has_task) {
200     parsed_name->task = parsed_basename.task;
201     parsed_name->has_task = true;
202   }
203   if (!parsed_name->has_type) {
204     parsed_name->type = parsed_basename.type;
205     parsed_name->has_type = true;
206   }
207   if (!parsed_name->has_id) {
208     parsed_name->id = parsed_basename.id;
209     parsed_name->has_id = true;
210   }
211 }
212 
213 }  // namespace
214 
215 /* static */
CanonicalizeDeviceName(StringPiece fullname,StringPiece basename,string * canonical_name)216 Status DeviceNameUtils::CanonicalizeDeviceName(StringPiece fullname,
217                                                StringPiece basename,
218                                                string* canonical_name) {
219   *canonical_name = "";
220   ParsedName parsed_basename;
221   if (!ParseFullName(basename, &parsed_basename)) {
222     return errors::InvalidArgument("Could not parse basename: ", basename,
223                                    " into a device specification.");
224   }
225   if (!(parsed_basename.has_job && parsed_basename.has_replica &&
226         parsed_basename.has_task && parsed_basename.has_type &&
227         parsed_basename.has_id)) {
228     return errors::InvalidArgument("Basename: ", basename,
229                                    " should be fully "
230                                    "specified.");
231   }
232   ParsedName parsed_name;
233   if (ParseLocalName(fullname, &parsed_name)) {
234     CompleteName(parsed_basename, &parsed_name);
235     *canonical_name = ParsedNameToString(parsed_name);
236     return Status::OK();
237   }
238   if (ParseFullName(fullname, &parsed_name)) {
239     CompleteName(parsed_basename, &parsed_name);
240     *canonical_name = ParsedNameToString(parsed_name);
241     return Status::OK();
242   }
243   return errors::InvalidArgument("Could not parse ", fullname,
244                                  " into a device "
245                                  "specification.");
246 }
247 
248 /* static */
ParsedNameToString(const ParsedName & pn)249 string DeviceNameUtils::ParsedNameToString(const ParsedName& pn) {
250   string buf;
251   if (pn.has_job) strings::StrAppend(&buf, "/job:", pn.job);
252   if (pn.has_replica) strings::StrAppend(&buf, "/replica:", pn.replica);
253   if (pn.has_task) strings::StrAppend(&buf, "/task:", pn.task);
254   if (pn.has_type) {
255     strings::StrAppend(&buf, "/device:", pn.type, ":");
256     if (pn.has_id) {
257       strings::StrAppend(&buf, pn.id);
258     } else {
259       strings::StrAppend(&buf, "*");
260     }
261   }
262   return buf;
263 }
264 
265 /* static */
IsSpecification(const ParsedName & less_specific,const ParsedName & more_specific)266 bool DeviceNameUtils::IsSpecification(const ParsedName& less_specific,
267                                       const ParsedName& more_specific) {
268   if (less_specific.has_job &&
269       (!more_specific.has_job || (less_specific.job != more_specific.job))) {
270     return false;
271   }
272   if (less_specific.has_replica &&
273       (!more_specific.has_replica ||
274        (less_specific.replica != more_specific.replica))) {
275     return false;
276   }
277   if (less_specific.has_task &&
278       (!more_specific.has_task || (less_specific.task != more_specific.task))) {
279     return false;
280   }
281   if (less_specific.has_type &&
282       (!more_specific.has_type || (less_specific.type != more_specific.type))) {
283     return false;
284   }
285   if (less_specific.has_id &&
286       (!more_specific.has_id || (less_specific.id != more_specific.id))) {
287     return false;
288   }
289   return true;
290 }
291 
EnsureSpecification(ParsedName * more_specific,const ParsedName & less_specific)292 void DeviceNameUtils::EnsureSpecification(ParsedName* more_specific,
293                                           const ParsedName& less_specific) {
294   if (less_specific.has_job) {
295     more_specific->has_job = true;
296     more_specific->job = less_specific.job;
297   }
298   if (less_specific.has_replica) {
299     more_specific->has_replica = true;
300     more_specific->replica = less_specific.replica;
301   }
302   if (less_specific.has_task) {
303     more_specific->has_task = true;
304     more_specific->task = less_specific.task;
305   }
306   if (less_specific.has_type) {
307     more_specific->has_type = true;
308     more_specific->type = less_specific.type;
309   }
310   if (less_specific.has_id) {
311     more_specific->has_id = true;
312     more_specific->id = less_specific.id;
313   }
314 }
315 
316 /* static */
IsCompleteSpecification(const ParsedName & pattern,const ParsedName & name)317 bool DeviceNameUtils::IsCompleteSpecification(const ParsedName& pattern,
318                                               const ParsedName& name) {
319   CHECK(name.has_job && name.has_replica && name.has_task && name.has_type &&
320         name.has_id);
321 
322   if (pattern.has_job && (pattern.job != name.job)) return false;
323   if (pattern.has_replica && (pattern.replica != name.replica)) return false;
324   if (pattern.has_task && (pattern.task != name.task)) return false;
325   if (pattern.has_type && (pattern.type != name.type)) return false;
326   if (pattern.has_id && (pattern.id != name.id)) return false;
327   return true;
328 }
329 
330 /* static */
MergeDevNames(ParsedName * target,const ParsedName & other,bool allow_soft_placement)331 Status DeviceNameUtils::MergeDevNames(ParsedName* target,
332                                       const ParsedName& other,
333                                       bool allow_soft_placement) {
334   if (other.has_job) {
335     if (target->has_job && target->job != other.job) {
336       return errors::InvalidArgument(
337           "Cannot merge devices with incompatible jobs: '",
338           ParsedNameToString(*target), "' and '", ParsedNameToString(other),
339           "'");
340     } else {
341       target->has_job = other.has_job;
342       target->job = other.job;
343     }
344   }
345 
346   if (other.has_replica) {
347     if (target->has_replica && target->replica != other.replica) {
348       return errors::InvalidArgument(
349           "Cannot merge devices with incompatible replicas: '",
350           ParsedNameToString(*target), "' and '", ParsedNameToString(other),
351           "'");
352     } else {
353       target->has_replica = other.has_replica;
354       target->replica = other.replica;
355     }
356   }
357 
358   if (other.has_task) {
359     if (target->has_task && target->task != other.task) {
360       return errors::InvalidArgument(
361           "Cannot merge devices with incompatible tasks: '",
362           ParsedNameToString(*target), "' and '", ParsedNameToString(other),
363           "'");
364     } else {
365       target->has_task = other.has_task;
366       target->task = other.task;
367     }
368   }
369 
370   if (other.has_type) {
371     if (target->has_type && target->type != other.type) {
372       if (!allow_soft_placement) {
373         return errors::InvalidArgument(
374             "Cannot merge devices with incompatible types: '",
375             ParsedNameToString(*target), "' and '", ParsedNameToString(other),
376             "'");
377       } else {
378         target->has_id = false;
379         target->has_type = false;
380         return Status::OK();
381       }
382     } else {
383       target->has_type = other.has_type;
384       target->type = other.type;
385     }
386   }
387 
388   if (other.has_id) {
389     if (target->has_id && target->id != other.id) {
390       if (!allow_soft_placement) {
391         return errors::InvalidArgument(
392             "Cannot merge devices with incompatible ids: '",
393             ParsedNameToString(*target), "' and '", ParsedNameToString(other),
394             "'");
395       } else {
396         target->has_id = false;
397         return Status::OK();
398       }
399     } else {
400       target->has_id = other.has_id;
401       target->id = other.id;
402     }
403   }
404 
405   return Status::OK();
406 }
407 
408 /* static */
IsSameAddressSpace(const ParsedName & a,const ParsedName & b)409 bool DeviceNameUtils::IsSameAddressSpace(const ParsedName& a,
410                                          const ParsedName& b) {
411   return (a.has_job && b.has_job && (a.job == b.job)) &&
412          (a.has_replica && b.has_replica && (a.replica == b.replica)) &&
413          (a.has_task && b.has_task && (a.task == b.task));
414 }
415 
416 /* static */
IsSameAddressSpace(StringPiece src,StringPiece dst)417 bool DeviceNameUtils::IsSameAddressSpace(StringPiece src, StringPiece dst) {
418   ParsedName x;
419   ParsedName y;
420   return ParseFullName(src, &x) && ParseFullName(dst, &y) &&
421          IsSameAddressSpace(x, y);
422 }
423 
424 /* static */
LocalName(StringPiece type,int id)425 string DeviceNameUtils::LocalName(StringPiece type, int id) {
426   return strings::StrCat("/device:", type, ":", id);
427 }
428 
429 namespace {
430 // Returns the legacy local device name given its "type" and "id" (which is
431 // '/device:type:id').
LegacyLocalName(StringPiece type,int id)432 string LegacyLocalName(StringPiece type, int id) {
433   return strings::StrCat(type, ":", id);
434 }
435 }  // anonymous namespace
436 
437 /* static */
LocalName(StringPiece fullname)438 string DeviceNameUtils::LocalName(StringPiece fullname) {
439   ParsedName x;
440   CHECK(ParseFullName(fullname, &x)) << fullname;
441   return LocalName(x.type, x.id);
442 }
443 
444 /* static */
ParseLocalName(StringPiece name,ParsedName * p)445 bool DeviceNameUtils::ParseLocalName(StringPiece name, ParsedName* p) {
446   if (!ConsumeDeviceType(&name, &p->type)) {
447     return false;
448   }
449   p->has_type = true;
450   if (!str_util::ConsumePrefix(&name, ":")) {
451     return false;
452   }
453   if (!ConsumeNumber(&name, &p->id)) {
454     return false;
455   }
456   p->has_id = true;
457   return name.empty();
458 }
459 
460 /* static */
SplitDeviceName(StringPiece name,string * task,string * device)461 bool DeviceNameUtils::SplitDeviceName(StringPiece name, string* task,
462                                       string* device) {
463   ParsedName pn;
464   if (ParseFullName(name, &pn) && pn.has_type && pn.has_id) {
465     task->clear();
466     task->reserve(
467         (pn.has_job ? (5 + pn.job.size()) : 0) +
468         (pn.has_replica ? (9 + 4 /*estimated UB for # replica digits*/) : 0) +
469         (pn.has_task ? (6 + 4 /*estimated UB for # task digits*/) : 0));
470     if (pn.has_job) {
471       strings::StrAppend(task, "/job:", pn.job);
472     }
473     if (pn.has_replica) {
474       strings::StrAppend(task, "/replica:", pn.replica);
475     }
476     if (pn.has_task) {
477       strings::StrAppend(task, "/task:", pn.task);
478     }
479     device->clear();
480     strings::StrAppend(device, pn.type, ":", pn.id);
481     return true;
482   }
483   return false;
484 }
485 
GetNamesForDeviceMappings(const ParsedName & pn)486 std::vector<string> DeviceNameUtils::GetNamesForDeviceMappings(
487     const ParsedName& pn) {
488   if (pn.has_job && pn.has_replica && pn.has_task && pn.has_type && pn.has_id) {
489     return {
490         DeviceNameUtils::FullName(pn.job, pn.replica, pn.task, pn.type, pn.id),
491         LegacyName(pn.job, pn.replica, pn.task, pn.type, pn.id)};
492   } else {
493     return {};
494   }
495 }
496 
GetLocalNamesForDeviceMappings(const ParsedName & pn)497 std::vector<string> DeviceNameUtils::GetLocalNamesForDeviceMappings(
498     const ParsedName& pn) {
499   if (pn.has_type && pn.has_id) {
500     return {DeviceNameUtils::LocalName(pn.type, pn.id),
501             LegacyLocalName(pn.type, pn.id)};
502   } else {
503     return {};
504   }
505 }
506 
DeviceNameToCpuDeviceName(const string & device_name,string * host_device_name)507 /*static*/ Status DeviceNameUtils::DeviceNameToCpuDeviceName(
508     const string& device_name, string* host_device_name) {
509   DeviceNameUtils::ParsedName device;
510   if (!DeviceNameUtils::ParseFullName(device_name, &device)) {
511     return errors::Internal("Could not parse device name ", device_name);
512   }
513   device.type = "CPU";
514   device.id = 0;
515   *host_device_name = DeviceNameUtils::ParsedNameToString(device);
516   return Status::OK();
517 }
518 
519 }  // namespace tensorflow
520