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/core/status_test_util.h"
20 #include "tensorflow/core/lib/strings/str_util.h"
21 #include "tensorflow/core/platform/test.h"
22 #include "tensorflow/core/platform/test_benchmark.h"
23
24 namespace tensorflow {
25
26 namespace {
27
RoundTripParsedName(const string & original,const string & expected)28 bool RoundTripParsedName(const string& original, const string& expected) {
29 DeviceNameUtils::ParsedName p;
30 if (!DeviceNameUtils::ParseFullName(original, &p)) {
31 return false;
32 }
33 string round_tripped = DeviceNameUtils::ParsedNameToString(p);
34 return (round_tripped == expected);
35 }
36
37 enum NamePart { kJob = 0x01, kReplica = 0x02, kTask = 0x04, kDevice = 0x08 };
38
RoundTripPartialName(int parts_to_test,const std::vector<string> & parts,bool explicitDevice)39 bool RoundTripPartialName(int parts_to_test, const std::vector<string>& parts,
40 bool explicitDevice) {
41 string original, expected;
42 if (parts_to_test & kJob) {
43 strings::StrAppend(&original, "/job:", parts[0]);
44 strings::StrAppend(&expected, "/job:", parts[0]);
45 }
46 if (parts_to_test & kReplica) {
47 strings::StrAppend(&original, "/replica:", parts[1]);
48 strings::StrAppend(&expected, "/replica:", parts[1]);
49 }
50 if (parts_to_test & kTask) {
51 strings::StrAppend(&original, "/task:", parts[2]);
52 strings::StrAppend(&expected, "/task:", parts[2]);
53 }
54 if (parts_to_test & kDevice) {
55 if (explicitDevice) {
56 strings::StrAppend(&original, "/device:", parts[3]);
57 strings::StrAppend(&expected, "/device:", parts[3]);
58 } else {
59 strings::StrAppend(&original, "/", parts[3]);
60 strings::StrAppend(&expected,
61 "/device:", absl::AsciiStrToUpper(parts[3]));
62 }
63 }
64 return RoundTripParsedName(original, expected);
65 }
66
67 } // namespace
68
TEST(DeviceNameUtilsTest,Basic)69 TEST(DeviceNameUtilsTest, Basic) {
70 EXPECT_EQ(DeviceNameUtils::FullName("hello", 1, 2, "CPU", 3),
71 "/job:hello/replica:1/task:2/device:CPU:3");
72
73 {
74 DeviceNameUtils::ParsedName p;
75 EXPECT_FALSE(DeviceNameUtils::ParseFullName("foobar", &p));
76 EXPECT_FALSE(DeviceNameUtils::ParseFullName(
77 "/job:123/replica:1/task:2/device:GPU:3", &p));
78 EXPECT_FALSE(
79 DeviceNameUtils::ParseFullName("/job:123/replica:1/task:2/gpu:", &p));
80 EXPECT_FALSE(DeviceNameUtils::ParseFullName(
81 "/job:123/replica:1/task:2/device:gpu:", &p));
82 EXPECT_FALSE(DeviceNameUtils::ParseFullName(
83 "/job:foo/replica:-1/task:2/device:GPU:3", &p));
84 EXPECT_FALSE(DeviceNameUtils::ParseFullName(
85 "/job:foo/replica:1/task:-2/device:GPU:3", &p));
86 EXPECT_FALSE(
87 DeviceNameUtils::ParseFullName("/job:foo/replica:1/task:2/bar:3", &p));
88 EXPECT_FALSE(DeviceNameUtils::ParseFullName(
89 "/job:foo/replica:1/task:2/device:GPU:3/extra", &p));
90 EXPECT_TRUE(DeviceNameUtils::ParseFullName(
91 "/job:foo/replica:1/task:2/device:GPU:3", &p));
92 EXPECT_TRUE(p.has_job);
93 EXPECT_TRUE(p.has_replica);
94 EXPECT_TRUE(p.has_task);
95 EXPECT_TRUE(p.has_type);
96 EXPECT_TRUE(p.has_id);
97 EXPECT_EQ(p.job, "foo");
98 EXPECT_EQ(p.replica, 1);
99 EXPECT_EQ(p.task, 2);
100 EXPECT_EQ(p.type, "GPU");
101 EXPECT_EQ(p.id, 3);
102 }
103 {
104 // Allow _ in job names.
105 DeviceNameUtils::ParsedName p;
106 EXPECT_TRUE(DeviceNameUtils::ParseFullName(
107 "/job:foo_bar/replica:1/task:2/device:GPU:3", &p));
108 EXPECT_TRUE(DeviceNameUtils::ParseFullOrLocalName(
109 "/job:foo_bar/replica:1/task:2/device:GPU:3", &p));
110 EXPECT_TRUE(p.has_job);
111 EXPECT_TRUE(p.has_replica);
112 EXPECT_TRUE(p.has_task);
113 EXPECT_TRUE(p.has_type);
114 EXPECT_TRUE(p.has_id);
115 EXPECT_EQ(p.job, "foo_bar");
116 EXPECT_EQ(p.replica, 1);
117 EXPECT_EQ(p.task, 2);
118 EXPECT_EQ(p.type, "GPU");
119 EXPECT_EQ(p.id, 3);
120 }
121 {
122 // Allow _ in job names.
123 DeviceNameUtils::ParsedName p;
124 EXPECT_TRUE(DeviceNameUtils::ParseFullName(
125 "/job:foo_bar/replica:1/task:2/device:GPU:3", &p));
126 EXPECT_TRUE(p.has_job);
127 EXPECT_TRUE(p.has_replica);
128 EXPECT_TRUE(p.has_task);
129 EXPECT_TRUE(p.has_type);
130 EXPECT_TRUE(p.has_id);
131 EXPECT_EQ(p.job, "foo_bar");
132 EXPECT_EQ(p.replica, 1);
133 EXPECT_EQ(p.task, 2);
134 EXPECT_EQ(p.type, "GPU");
135 EXPECT_EQ(p.id, 3);
136 }
137 {
138 DeviceNameUtils::ParsedName p;
139 EXPECT_TRUE(DeviceNameUtils::ParseFullName("/job:*/replica:4/gpu:*", &p));
140 EXPECT_FALSE(p.has_job);
141 EXPECT_TRUE(p.has_replica);
142 EXPECT_FALSE(p.has_task);
143 EXPECT_TRUE(p.has_type);
144 EXPECT_FALSE(p.has_id);
145 EXPECT_EQ(p.replica, 4);
146 EXPECT_EQ(p.type, "GPU");
147 }
148 {
149 DeviceNameUtils::ParsedName p;
150 EXPECT_TRUE(
151 DeviceNameUtils::ParseFullName("/job:*/replica:4/device:GPU:*", &p));
152 EXPECT_FALSE(p.has_job);
153 EXPECT_TRUE(p.has_replica);
154 EXPECT_FALSE(p.has_task);
155 EXPECT_TRUE(p.has_type);
156 EXPECT_FALSE(p.has_id);
157 EXPECT_EQ(p.replica, 4);
158 EXPECT_EQ(p.type, "GPU");
159 }
160 {
161 DeviceNameUtils::ParsedName p;
162 EXPECT_TRUE(
163 DeviceNameUtils::ParseFullName("/job:*/device:GPU/replica:4", &p));
164 EXPECT_FALSE(p.has_job);
165 EXPECT_TRUE(p.has_replica);
166 EXPECT_FALSE(p.has_task);
167 EXPECT_TRUE(p.has_type);
168 EXPECT_FALSE(p.has_id);
169 EXPECT_EQ(p.replica, 4);
170 EXPECT_EQ(p.type, "GPU");
171 }
172 {
173 DeviceNameUtils::ParsedName p;
174 EXPECT_TRUE(DeviceNameUtils::ParseFullName(
175 "/job:*/replica:4/device:myspecialdevice:13", &p));
176 EXPECT_FALSE(p.has_job);
177 EXPECT_TRUE(p.has_replica);
178 EXPECT_FALSE(p.has_task);
179 EXPECT_TRUE(p.has_type);
180 EXPECT_TRUE(p.has_id);
181 EXPECT_EQ(p.replica, 4);
182 EXPECT_EQ(p.type, "myspecialdevice");
183 EXPECT_EQ(p.id, 13);
184 }
185 {
186 DeviceNameUtils::ParsedName p;
187 EXPECT_TRUE(DeviceNameUtils::ParseFullName("/", &p));
188 EXPECT_FALSE(p.has_job);
189 EXPECT_FALSE(p.has_replica);
190 EXPECT_FALSE(p.has_task);
191 EXPECT_FALSE(p.has_type);
192 EXPECT_FALSE(p.has_id);
193 }
194 {
195 DeviceNameUtils::ParsedName p;
196 EXPECT_TRUE(
197 DeviceNameUtils::ParseFullName("/job:*/replica:4/device:GPU:5", &p));
198 EXPECT_FALSE(p.has_job);
199 EXPECT_TRUE(p.has_replica);
200 EXPECT_FALSE(p.has_task);
201 EXPECT_TRUE(p.has_type);
202 EXPECT_TRUE(p.has_id);
203 EXPECT_EQ(p.replica, 4);
204 EXPECT_EQ(p.type, "GPU");
205 EXPECT_EQ(p.id, 5);
206 }
207 { // Same result if we reorder the components
208 DeviceNameUtils::ParsedName p;
209 EXPECT_TRUE(DeviceNameUtils::ParseFullName("/gpu:*/job:*/replica:4", &p));
210 EXPECT_FALSE(p.has_job);
211 EXPECT_TRUE(p.has_replica);
212 EXPECT_FALSE(p.has_task);
213 EXPECT_TRUE(p.has_type);
214 EXPECT_FALSE(p.has_id);
215 EXPECT_EQ(p.replica, 4);
216 EXPECT_EQ(p.type, "GPU");
217 }
218
219 EXPECT_TRUE(DeviceNameUtils::IsSameAddressSpace(
220 "/job:foo/replica:1/task:2/cpu:3",
221 "/job:foo/replica:1/task:2/device:GPU:4"));
222 EXPECT_FALSE(DeviceNameUtils::IsSameAddressSpace(
223 "/job:foo/replica:1/task:2/cpu:3",
224 "/job:foo/replica:1/task:3/device:GPU:4"));
225 EXPECT_FALSE(DeviceNameUtils::IsSameAddressSpace(
226 "/job:foo/replica:1/task:2/cpu:3",
227 "/job:foo/replica:10/task:2/device:GPU:4"));
228 EXPECT_FALSE(DeviceNameUtils::IsSameAddressSpace(
229 "/job:foo/replica:1/task:2/cpu:3",
230 "/job:bar/replica:1/task:2/device:GPU:4"));
231
232 EXPECT_EQ(DeviceNameUtils::LocalName("CPU", 1), "/device:CPU:1");
233 EXPECT_EQ(DeviceNameUtils::LocalName("GPU", 2), "/device:GPU:2");
234 EXPECT_EQ(DeviceNameUtils::LocalName("MySpecialDevice", 13),
235 "/device:MySpecialDevice:13");
236
237 EXPECT_EQ(
238 DeviceNameUtils::LocalName("/job:foo/replica:1/task:2/device:CPU:3"),
239 "/device:CPU:3");
240
241 EXPECT_EQ(DeviceNameUtils::LocalName("/job:foo/replica:1/task:2/cpu:3"),
242 "/device:CPU:3");
243
244 EXPECT_EQ(
245 DeviceNameUtils::LocalName("/job:foo/replica:1/task:2/device:abc:73"),
246 "/device:abc:73");
247
248 {
249 DeviceNameUtils::ParsedName p;
250 EXPECT_TRUE(DeviceNameUtils::ParseLocalName("CPU:10", &p));
251 EXPECT_TRUE(DeviceNameUtils::ParseFullOrLocalName("CPU:10", &p));
252 EXPECT_EQ(p.type, "CPU");
253 EXPECT_EQ(p.id, 10);
254 EXPECT_FALSE(DeviceNameUtils::ParseLocalName("cpu:abc", &p));
255 EXPECT_FALSE(DeviceNameUtils::ParseLocalName("abc:", &p));
256 EXPECT_FALSE(DeviceNameUtils::ParseLocalName("abc", &p));
257 EXPECT_FALSE(DeviceNameUtils::ParseLocalName("myspecialdevice", &p));
258 EXPECT_FALSE(DeviceNameUtils::ParseFullOrLocalName("myspecialdevice", &p));
259 }
260
261 // Test that all parts are round-tripped correctly.
262 {
263 for (int i = 0; i < 0x10; ++i) {
264 EXPECT_TRUE(RoundTripPartialName(i, {"foo", "3", "2", "CPU:3"},
265 /*explicitDevice=*/false));
266 EXPECT_TRUE(RoundTripPartialName(i, {"foo", "3", "2", "GPU:3"},
267 /*explicitDevice=*/false));
268 EXPECT_TRUE(RoundTripPartialName(i, {"foo", "3", "2", "cpu:3"},
269 /*explicitDevice=*/false));
270 EXPECT_TRUE(RoundTripPartialName(i, {"foo", "3", "2", "gpu:3"},
271 /*explicitDevice=*/false));
272 EXPECT_TRUE(RoundTripPartialName(i, {"foo", "3", "2", "CPU:3"},
273 /*explicitDevice=*/true));
274 EXPECT_TRUE(RoundTripPartialName(i, {"foo", "3", "2", "GPU:3"},
275 /*explicitDevice=*/true));
276 EXPECT_TRUE(RoundTripPartialName(i, {"foo", "3", "2", "cpu:3"},
277 /*explicitDevice=*/true));
278 EXPECT_TRUE(RoundTripPartialName(i, {"foo", "3", "2", "gpu:3"},
279 /*explicitDevice=*/true));
280 EXPECT_TRUE(RoundTripPartialName(i, {"foo", "3", "2", "someDevice:3"},
281 /*explicitDevice=*/true));
282 }
283 }
284 {
285 DeviceNameUtils::ParsedName x, y;
286 DeviceNameUtils::ParseFullName("/job:work/replica:1/task:3/device:GPU:*",
287 &x);
288 DeviceNameUtils::ParseFullName("/device:CPU:*", &y);
289 EXPECT_FALSE(DeviceNameUtils::AreCompatibleDevNames(x, y));
290 }
291 {
292 DeviceNameUtils::ParsedName x, y;
293 DeviceNameUtils::ParseFullName("/job:work/replica:1/task:3", &x);
294 DeviceNameUtils::ParseFullName("/device:CPU:*", &y);
295 EXPECT_TRUE(DeviceNameUtils::AreCompatibleDevNames(x, y));
296 }
297 }
298
IsCSHelper(StringPiece pattern,StringPiece actual)299 static bool IsCSHelper(StringPiece pattern, StringPiece actual) {
300 DeviceNameUtils::ParsedName p, a;
301 EXPECT_TRUE(DeviceNameUtils::ParseFullName(pattern, &p));
302 EXPECT_TRUE(DeviceNameUtils::ParseFullName(actual, &a));
303 return DeviceNameUtils::IsCompleteSpecification(p, a);
304 }
305
TEST(DeviceNameUtilsTest,IsCompleteSpecification)306 TEST(DeviceNameUtilsTest, IsCompleteSpecification) {
307 EXPECT_TRUE(IsCSHelper("/job:*", "/job:work/replica:1/task:2/device:GPU:3"));
308 EXPECT_TRUE(IsCSHelper("/job:*/replica:*",
309 "/job:work/replica:1/task:2/device:GPU:3"));
310 EXPECT_TRUE(
311 IsCSHelper("/job:*/task:*", "/job:work/replica:1/task:2/device:GPU:3"));
312 EXPECT_TRUE(IsCSHelper("/job:*/replica:*/task:*",
313 "/job:work/replica:1/task:2/device:GPU:3"));
314 EXPECT_TRUE(IsCSHelper("/job:*/replica:*/gpu:*",
315 "/job:work/replica:1/task:2/device:GPU:3"));
316 EXPECT_FALSE(
317 IsCSHelper("/cpu:*", "/job:worker/replica:1/task:2/device:GPU:3"));
318 EXPECT_FALSE(
319 IsCSHelper("/device:GPU:2", "/job:worker/replica:1/task:2/device:GPU:1"));
320 EXPECT_TRUE(
321 IsCSHelper("/gpu:*", "/job:worker/replica:1/task:2/device:GPU:3"));
322 }
323
IsSpecHelper(StringPiece pattern,StringPiece actual)324 static bool IsSpecHelper(StringPiece pattern, StringPiece actual) {
325 DeviceNameUtils::ParsedName p, a;
326 EXPECT_TRUE(DeviceNameUtils::ParseFullName(pattern, &p));
327 EXPECT_TRUE(DeviceNameUtils::ParseFullName(actual, &a));
328 return DeviceNameUtils::IsSpecification(p, a);
329 }
330
TEST(DeviceNameUtilsTest,IsSpecification)331 TEST(DeviceNameUtilsTest, IsSpecification) {
332 EXPECT_TRUE(
333 IsSpecHelper("/job:*", "/job:work/replica:1/task:2/device:GPU:3"));
334 EXPECT_TRUE(IsSpecHelper("/job:*", "/job:work/replica:1/device:GPU:3"));
335 EXPECT_TRUE(IsSpecHelper("/job:*", "/job:work/replica:1"));
336 EXPECT_TRUE(IsSpecHelper("/job:*", "/replica:1"));
337 EXPECT_TRUE(IsSpecHelper("/job:*", "/job:work"));
338 EXPECT_TRUE(IsSpecHelper("/job:*/replica:*",
339 "/job:work/replica:1/task:2/device:GPU:3"));
340 EXPECT_TRUE(IsSpecHelper("/job:work/replica:1/gpu:*",
341 "/job:work/replica:1/task:2/device:GPU:3"));
342 EXPECT_TRUE(IsSpecHelper("/job:work/replica:1/device:GPU:3",
343 "/job:work/replica:1/task:2/device:GPU:3"));
344 EXPECT_TRUE(IsSpecHelper("/job:work/replica:1/task:2",
345 "/job:work/replica:1/task:2/device:GPU:3"));
346 EXPECT_TRUE(IsSpecHelper("/job:work/replica:*/task:2",
347 "/job:work/replica:1/task:2/device:GPU:3"));
348 EXPECT_TRUE(IsSpecHelper("/task:*", "/job:*/replica:1/task:2/device:GPU:3"));
349 EXPECT_TRUE(IsSpecHelper("/task:2", "/job:*/replica:1/task:2/device:GPU:3"));
350 EXPECT_TRUE(IsSpecHelper("/cpu:*", "/job:*/replica:1/task:2/cpu:1"));
351 EXPECT_TRUE(IsSpecHelper("/cpu:0", "/cpu:0"));
352 EXPECT_TRUE(
353 IsSpecHelper("/gpu:*", "/job:worker/replica:1/task:2/device:GPU:3"));
354
355 EXPECT_FALSE(
356 IsSpecHelper("/job:worker/replica:1/task:2/device:GPU:3", "/gpu:*"));
357 EXPECT_FALSE(IsSpecHelper("/cpu:*", "/job:*/replica:1/task:2"));
358 EXPECT_FALSE(IsSpecHelper("/cpu:*", "/job:*/replica:1/task:2/device:GPU:1"));
359 EXPECT_FALSE(
360 IsSpecHelper("/cpu:*", "/job:worker/replica:1/task:2/device:GPU:3"));
361 EXPECT_FALSE(IsSpecHelper("/device:GPU:2",
362 "/job:worker/replica:1/task:2/device:GPU:1"));
363 EXPECT_FALSE(IsSpecHelper("/job:work/replica:*/task:0",
364 "/job:work/replica:1/task:2/device:GPU:3"));
365 EXPECT_FALSE(IsSpecHelper("/job:work/replica:0/task:2",
366 "/job:work/replica:*/task:2/device:GPU:3"));
367 }
368
TEST(DeviceNameUtilsTest,SplitDeviceName)369 TEST(DeviceNameUtilsTest, SplitDeviceName) {
370 string task;
371 string device;
372 EXPECT_TRUE(DeviceNameUtils::SplitDeviceName(
373 "/job:foo/replica:1/task:2/cpu:1", &task, &device));
374 EXPECT_EQ("/job:foo/replica:1/task:2", task);
375 EXPECT_EQ("CPU:1", device);
376 EXPECT_TRUE(DeviceNameUtils::SplitDeviceName(
377 "/job:foo/cpu:1/task:2/replica:1", &task, &device));
378 EXPECT_EQ("/job:foo/replica:1/task:2", task);
379 EXPECT_EQ("CPU:1", device);
380 EXPECT_TRUE(
381 DeviceNameUtils::SplitDeviceName("/device:GPU:3", &task, &device));
382 EXPECT_EQ("", task);
383 EXPECT_EQ("GPU:3", device);
384 EXPECT_FALSE(DeviceNameUtils::SplitDeviceName("gpu:3", &task, &device));
385 EXPECT_FALSE(DeviceNameUtils::SplitDeviceName("/job:foo/task:2/replica:1",
386 &task, &device));
387 EXPECT_TRUE(DeviceNameUtils::SplitDeviceName("/device:myspecialdevice:3",
388 &task, &device));
389 EXPECT_EQ("", task);
390 EXPECT_EQ("myspecialdevice:3", device);
391 }
392
Name(const string & str)393 static DeviceNameUtils::ParsedName Name(const string& str) {
394 DeviceNameUtils::ParsedName ret;
395 CHECK(DeviceNameUtils::ParseFullName(str, &ret)) << "Invalid name: " << str;
396 return ret;
397 }
398
MergeDevNamesHelperImpl(const string & name_a,const string & name_b,const string & expected_merge_name,bool allow_soft_placement)399 static void MergeDevNamesHelperImpl(const string& name_a, const string& name_b,
400 const string& expected_merge_name,
401 bool allow_soft_placement) {
402 DeviceNameUtils::ParsedName target_a = Name(name_a);
403 TF_EXPECT_OK(DeviceNameUtils::MergeDevNames(&target_a, Name(name_b),
404 allow_soft_placement));
405 DeviceNameUtils::ParsedName target_b = Name(name_b);
406 TF_EXPECT_OK(DeviceNameUtils::MergeDevNames(&target_b, Name(name_a),
407 allow_soft_placement));
408 EXPECT_EQ(target_a, target_b);
409 EXPECT_EQ(target_a, Name(expected_merge_name));
410 EXPECT_EQ(target_b, Name(expected_merge_name));
411 }
412
MergeDevNamesHelper(const string & name_a,const string & name_b,const string & expected_merge_name)413 static void MergeDevNamesHelper(const string& name_a, const string& name_b,
414 const string& expected_merge_name) {
415 MergeDevNamesHelperImpl(name_a, name_b, expected_merge_name, false);
416 }
417
MergeDevNamesHelperAllowSoftPlacement(const string & name_a,const string & name_b,const string & expected_merge_name)418 static void MergeDevNamesHelperAllowSoftPlacement(
419 const string& name_a, const string& name_b,
420 const string& expected_merge_name) {
421 MergeDevNamesHelperImpl(name_a, name_b, expected_merge_name, true);
422 }
423
MergeDevNamesError(const string & name_a,const string & name_b,const string & expected_error_substr)424 static void MergeDevNamesError(const string& name_a, const string& name_b,
425 const string& expected_error_substr) {
426 DeviceNameUtils::ParsedName target_a = Name(name_a);
427 Status s = DeviceNameUtils::MergeDevNames(&target_a, Name(name_b));
428 EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
429 EXPECT_TRUE(absl::StrContains(s.error_message(), expected_error_substr)) << s;
430 }
431
MergeOverrideHelper(const string & target,const string & name,const string & expected_merge_name)432 static void MergeOverrideHelper(const string& target, const string& name,
433 const string& expected_merge_name) {
434 DeviceNameUtils::ParsedName parsed_target = Name(target);
435 TF_EXPECT_OK(
436 DeviceNameUtils::MergeOverrideDevNames(&parsed_target, Name(name)));
437 DeviceNameUtils::ParsedName parsed_expected = Name(expected_merge_name);
438
439 EXPECT_EQ(parsed_target, parsed_expected)
440 << "parsed_target: " << DeviceNameUtils::ParsedNameToString(parsed_target)
441 << " expected_name: "
442 << DeviceNameUtils::ParsedNameToString(parsed_expected);
443 }
444
MergeUnsetDevNamesHelper(const string & name_a,const string & name_b,const string & expected_merge_name_ab,const string & expected_merge_name_ba)445 static void MergeUnsetDevNamesHelper(const string& name_a, const string& name_b,
446 const string& expected_merge_name_ab,
447 const string& expected_merge_name_ba) {
448 DeviceNameUtils::ParsedName target_a = Name(name_a);
449 DeviceNameUtils::MergeUnsetDevNames(&target_a, Name(name_b));
450 EXPECT_EQ(target_a, Name(expected_merge_name_ab));
451 DeviceNameUtils::ParsedName target_b = Name(name_b);
452 DeviceNameUtils::MergeUnsetDevNames(&target_b, Name(name_a));
453 EXPECT_EQ(target_b, Name(expected_merge_name_ba));
454 }
455
TEST(DeviceNameUtilsTest,MergeDevNames)456 TEST(DeviceNameUtilsTest, MergeDevNames) {
457 // Idempotence tests.
458 MergeDevNamesHelper("", "", "");
459 MergeDevNamesHelper("/job:foo/replica:1/task:2/cpu:1",
460 "/job:foo/replica:1/task:2/cpu:1",
461 "/job:foo/replica:1/task:2/cpu:1");
462
463 // Merging with empty device has no effect.
464 MergeDevNamesHelper("", "/job:foo", "/job:foo");
465 MergeDevNamesHelper("", "/replica:2", "/replica:2");
466 MergeDevNamesHelper("", "/task:7", "/task:7");
467 MergeDevNamesHelper("", "/device:GPU:1", "/device:GPU:1");
468
469 // Combining disjoint names.
470 MergeDevNamesHelper("/job:foo", "/task:7", "/job:foo/task:7");
471 MergeDevNamesHelper("/job:foo", "/device:GPU:1", "/job:foo/device:GPU:1");
472
473 // Combining overlapping names.
474 MergeDevNamesHelper("/job:foo/replica:0", "/replica:0/task:1",
475 "/job:foo/replica:0/task:1");
476
477 // Wildcard tests.
478 MergeDevNamesHelper("", "/gpu:*", "/gpu:*");
479 MergeDevNamesHelper("/gpu:*", "/gpu:*", "/gpu:*");
480 MergeDevNamesHelper("/device:GPU:1", "/gpu:*", "/device:GPU:1");
481
482 // Incompatible components.
483 MergeDevNamesError("/job:foo", "/job:bar", "incompatible jobs");
484 MergeDevNamesError("/replica:0", "/replica:1", "incompatible replicas");
485 MergeDevNamesError("/task:0", "/task:1", "incompatible tasks");
486 MergeDevNamesError("/gpu:*", "/cpu:*", "incompatible types");
487 MergeDevNamesError("/device:GPU:0", "/device:GPU:1", "incompatible ids");
488 }
489
TEST(DeviceNameUtilsTest,MergeDevNamesAllowSoftPlacement)490 TEST(DeviceNameUtilsTest, MergeDevNamesAllowSoftPlacement) {
491 // Incompatible components with allow_soft_placement.
492 MergeDevNamesHelperAllowSoftPlacement("/gpu:*", "/cpu:1", "");
493 MergeDevNamesHelperAllowSoftPlacement("/cpu:*", "/device:GPU:1", "");
494 MergeDevNamesHelperAllowSoftPlacement("/device:GPU:1", "/device:GPU:2",
495 "/device:GPU:*");
496 }
497
TEST(DeviceNameUtilsTest,MergeOverrideDevNames)498 TEST(DeviceNameUtilsTest, MergeOverrideDevNames) {
499 // Idempotence tests.
500 MergeOverrideHelper("", "", "");
501 MergeOverrideHelper("/job:foo/replica:1/task:2/cpu:1",
502 "/job:foo/replica:1/task:2/cpu:1",
503 "/job:foo/replica:1/task:2/cpu:1");
504
505 // Merging with empty device has no effect.
506 MergeOverrideHelper("", "/job:foo", "/job:foo");
507 MergeOverrideHelper("", "/replica:2", "/replica:2");
508 MergeOverrideHelper("", "/task:7", "/task:7");
509 MergeOverrideHelper("", "/device:GPU:1", "/device:GPU:1");
510
511 // Combining disjoint names.
512 MergeOverrideHelper("/job:foo", "/task:7", "/job:foo/task:7");
513 MergeOverrideHelper("/job:foo", "/device:GPU:1", "/job:foo/device:GPU:1");
514
515 // Combining overlapping names.
516 MergeOverrideHelper("/job:foo/replica:0", "/replica:0/task:1",
517 "/job:foo/replica:0/task:1");
518
519 // Wildcard tests.
520 MergeOverrideHelper("", "/gpu:*", "/gpu:*");
521 MergeOverrideHelper("/gpu:*", "/gpu:*", "/gpu:*");
522 MergeOverrideHelper("/device:GPU:1", "/gpu:*", "/device:GPU:1");
523
524 // Testing actual override functionality
525 MergeOverrideHelper("/gpu:0", "/cpu:1", "/cpu:1");
526 MergeOverrideHelper("/gpu:*", "/cpu:1", "/cpu:1");
527 MergeOverrideHelper("/cpu:*", "/device:GPU:1", "/gpu:1");
528 MergeOverrideHelper("/device:GPU:1", "/device:GPU:2", "/device:GPU:2");
529
530 // Override with regular merging
531 MergeOverrideHelper("/job:foo/CPU:*", "/device:GPU:1", "/job:foo/GPU:1");
532 MergeOverrideHelper("/cpu:*", "/job:foo/device:GPU:1", "/job:foo/GPU:1");
533 MergeOverrideHelper("/task:0/cpu:*", "/device:GPU:1", "/task:0/GPU:1");
534 MergeOverrideHelper("/cpu:*", "/task:0/device:GPU:1", "/task:0/GPU:1");
535 }
536
TEST(DeviceNameUtilsTest,MergeUnsetDevNames)537 TEST(DeviceNameUtilsTest, MergeUnsetDevNames) {
538 // Idempotence tests.
539 MergeUnsetDevNamesHelper("", "", "", "");
540 MergeUnsetDevNamesHelper(
541 "/job:foo/replica:1/task:2/cpu:1", "/job:foo/replica:1/task:2/cpu:1",
542 "/job:foo/replica:1/task:2/cpu:1", "/job:foo/replica:1/task:2/cpu:1");
543
544 // Merging with empty device has no effect.
545 MergeUnsetDevNamesHelper("", "/job:foo", "/job:foo", "/job:foo");
546 MergeUnsetDevNamesHelper("", "/replica:2", "/replica:2", "/replica:2");
547 MergeUnsetDevNamesHelper("", "/task:7", "/task:7", "/task:7");
548 MergeUnsetDevNamesHelper("", "/device:GPU:1", "/device:GPU:1",
549 "/device:GPU:1");
550
551 // Combining disjoint names.
552 MergeUnsetDevNamesHelper("/job:foo", "/task:7", "/job:foo/task:7",
553 "/job:foo/task:7");
554 MergeUnsetDevNamesHelper("/job:foo", "/device:GPU:1", "/job:foo/device:GPU:1",
555 "/job:foo/device:GPU:1");
556
557 // Combining overlapping names.
558 MergeUnsetDevNamesHelper("/job:foo/replica:0", "/replica:0/task:1",
559 "/job:foo/replica:0/task:1",
560 "/job:foo/replica:0/task:1");
561
562 // Wildcard tests.
563 MergeUnsetDevNamesHelper("", "/gpu:*", "/gpu:*", "/gpu:*");
564 MergeUnsetDevNamesHelper("/gpu:*", "/gpu:*", "/gpu:*", "/gpu:*");
565 MergeUnsetDevNamesHelper("/device:GPU:1", "/gpu:*", "/device:GPU:1",
566 "/device:GPU:1");
567
568 // Incompatible components.
569 MergeUnsetDevNamesHelper("/job:foo", "/job:bar", "/job:foo", "/job:bar");
570 MergeUnsetDevNamesHelper("/replica:0", "/replica:1", "/replica:0",
571 "/replica:1");
572 MergeUnsetDevNamesHelper("/task:0", "/task:1", "/task:0", "/task:1");
573 MergeUnsetDevNamesHelper("/gpu:*", "/cpu:*", "/gpu:*", "/cpu:*");
574 MergeUnsetDevNamesHelper("/device:GPU:0", "/device:GPU:1", "/device:GPU:0",
575 "/device:GPU:1");
576 MergeUnsetDevNamesHelper("/job:foo/device:GPU", "/job:bar",
577 "/job:foo/device:GPU", "/job:bar/device:GPU");
578 }
579
TEST(DeviceNameUtilsTest,GetNamesForDeviceMappings)580 TEST(DeviceNameUtilsTest, GetNamesForDeviceMappings) {
581 DeviceNameUtils::ParsedName p =
582 Name("/job:foo/replica:10/task:0/device:GPU:1");
583 EXPECT_EQ(absl::StrJoin(DeviceNameUtils::GetNamesForDeviceMappings(p), ","),
584 "/job:foo/replica:10/task:0/device:GPU:1,"
585 "/job:foo/replica:10/task:0/gpu:1");
586 p.has_task = false;
587 EXPECT_EQ(absl::StrJoin(DeviceNameUtils::GetNamesForDeviceMappings(p), ","),
588 "");
589 }
590
TEST(DeviceNameUtilsTest,CanonicalizeDeviceName)591 TEST(DeviceNameUtilsTest, CanonicalizeDeviceName) {
592 string canonical_name;
593 {
594 // Good basename.
595 string basename = "/job:foo/replica:10/task:0/device:CPU:0";
596 TF_EXPECT_OK(DeviceNameUtils::CanonicalizeDeviceName(
597 "/job:foo/replica:10/task:0/device:CPU:1", basename, &canonical_name));
598 EXPECT_EQ("/job:foo/replica:10/task:0/device:CPU:1", canonical_name);
599 TF_EXPECT_OK(DeviceNameUtils::CanonicalizeDeviceName(
600 "/job:foo/task:0/replica:10/device:CPU:1", basename, &canonical_name));
601 EXPECT_EQ("/job:foo/replica:10/task:0/device:CPU:1", canonical_name);
602 TF_EXPECT_OK(DeviceNameUtils::CanonicalizeDeviceName(
603 "/job:foo/task:0/replica:10/cpu:1", basename, &canonical_name));
604 EXPECT_EQ("/job:foo/replica:10/task:0/device:CPU:1", canonical_name);
605 TF_EXPECT_OK(DeviceNameUtils::CanonicalizeDeviceName("CPU:0", basename,
606 &canonical_name));
607 EXPECT_EQ("/job:foo/replica:10/task:0/device:CPU:0", canonical_name);
608 Status s = DeviceNameUtils::CanonicalizeDeviceName(
609 "/job:foo/task:0/replica/cpu:1", basename, &canonical_name);
610 EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
611 EXPECT_EQ("", canonical_name);
612 }
613
614 {
615 // Try out malformed basenames.
616 string fullname = "/device:CPU:0";
617
618 Status s = DeviceNameUtils::CanonicalizeDeviceName(
619 fullname, "/device:CPU:0", &canonical_name);
620 EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
621 EXPECT_EQ("", canonical_name);
622 s = DeviceNameUtils::CanonicalizeDeviceName(
623 fullname, "/job:foo/task:0/replica/cpu:1", &canonical_name);
624 EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
625 EXPECT_EQ("", canonical_name);
626 }
627 }
628
BM_ParseFullName(::testing::benchmark::State & state)629 static void BM_ParseFullName(::testing::benchmark::State& state) {
630 DeviceNameUtils::ParsedName p;
631 for (auto s : state) {
632 DeviceNameUtils::ParseFullName("/job:worker/replica:3/task:0/cpu:0", &p);
633 }
634 }
635 BENCHMARK(BM_ParseFullName);
636
637 } // namespace tensorflow
638