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