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
TEST(DeviceNameUtilsTest,MergeDevNames)445 TEST(DeviceNameUtilsTest, MergeDevNames) {
446 // Idempotence tests.
447 MergeDevNamesHelper("", "", "");
448 MergeDevNamesHelper("/job:foo/replica:1/task:2/cpu:1",
449 "/job:foo/replica:1/task:2/cpu:1",
450 "/job:foo/replica:1/task:2/cpu:1");
451
452 // Merging with empty device has no effect.
453 MergeDevNamesHelper("", "/job:foo", "/job:foo");
454 MergeDevNamesHelper("", "/replica:2", "/replica:2");
455 MergeDevNamesHelper("", "/task:7", "/task:7");
456 MergeDevNamesHelper("", "/device:GPU:1", "/device:GPU:1");
457
458 // Combining disjoint names.
459 MergeDevNamesHelper("/job:foo", "/task:7", "/job:foo/task:7");
460 MergeDevNamesHelper("/job:foo", "/device:GPU:1", "/job:foo/device:GPU:1");
461
462 // Combining overlapping names.
463 MergeDevNamesHelper("/job:foo/replica:0", "/replica:0/task:1",
464 "/job:foo/replica:0/task:1");
465
466 // Wildcard tests.
467 MergeDevNamesHelper("", "/gpu:*", "/gpu:*");
468 MergeDevNamesHelper("/gpu:*", "/gpu:*", "/gpu:*");
469 MergeDevNamesHelper("/device:GPU:1", "/gpu:*", "/device:GPU:1");
470
471 // Incompatible components.
472 MergeDevNamesError("/job:foo", "/job:bar", "incompatible jobs");
473 MergeDevNamesError("/replica:0", "/replica:1", "incompatible replicas");
474 MergeDevNamesError("/task:0", "/task:1", "incompatible tasks");
475 MergeDevNamesError("/gpu:*", "/cpu:*", "incompatible types");
476 MergeDevNamesError("/device:GPU:0", "/device:GPU:1", "incompatible ids");
477 }
478
TEST(DeviceNameUtilsTest,MergeDevNamesAllowSoftPlacement)479 TEST(DeviceNameUtilsTest, MergeDevNamesAllowSoftPlacement) {
480 // Incompatible components with allow_soft_placement.
481 MergeDevNamesHelperAllowSoftPlacement("/gpu:*", "/cpu:1", "");
482 MergeDevNamesHelperAllowSoftPlacement("/cpu:*", "/device:GPU:1", "");
483 MergeDevNamesHelperAllowSoftPlacement("/device:GPU:1", "/device:GPU:2",
484 "/device:GPU:*");
485 }
486
TEST(DeviceNameUtilsTest,MergeOverrideDevNames)487 TEST(DeviceNameUtilsTest, MergeOverrideDevNames) {
488 // Idempotence tests.
489 MergeOverrideHelper("", "", "");
490 MergeOverrideHelper("/job:foo/replica:1/task:2/cpu:1",
491 "/job:foo/replica:1/task:2/cpu:1",
492 "/job:foo/replica:1/task:2/cpu:1");
493
494 // Merging with empty device has no effect.
495 MergeOverrideHelper("", "/job:foo", "/job:foo");
496 MergeOverrideHelper("", "/replica:2", "/replica:2");
497 MergeOverrideHelper("", "/task:7", "/task:7");
498 MergeOverrideHelper("", "/device:GPU:1", "/device:GPU:1");
499
500 // Combining disjoint names.
501 MergeOverrideHelper("/job:foo", "/task:7", "/job:foo/task:7");
502 MergeOverrideHelper("/job:foo", "/device:GPU:1", "/job:foo/device:GPU:1");
503
504 // Combining overlapping names.
505 MergeOverrideHelper("/job:foo/replica:0", "/replica:0/task:1",
506 "/job:foo/replica:0/task:1");
507
508 // Wildcard tests.
509 MergeOverrideHelper("", "/gpu:*", "/gpu:*");
510 MergeOverrideHelper("/gpu:*", "/gpu:*", "/gpu:*");
511 MergeOverrideHelper("/device:GPU:1", "/gpu:*", "/device:GPU:1");
512
513 // Testing actual override functionality
514 MergeOverrideHelper("/gpu:0", "/cpu:1", "/cpu:1");
515 MergeOverrideHelper("/gpu:*", "/cpu:1", "/cpu:1");
516 MergeOverrideHelper("/cpu:*", "/device:GPU:1", "/gpu:1");
517 MergeOverrideHelper("/device:GPU:1", "/device:GPU:2", "/device:GPU:2");
518
519 // Override with regular merging
520 MergeOverrideHelper("/job:foo/CPU:*", "/device:GPU:1", "/job:foo/GPU:1");
521 MergeOverrideHelper("/cpu:*", "/job:foo/device:GPU:1", "/job:foo/GPU:1");
522 MergeOverrideHelper("/task:0/cpu:*", "/device:GPU:1", "/task:0/GPU:1");
523 MergeOverrideHelper("/cpu:*", "/task:0/device:GPU:1", "/task:0/GPU:1");
524 }
525
TEST(DeviceNameUtilsTest,GetNamesForDeviceMappings)526 TEST(DeviceNameUtilsTest, GetNamesForDeviceMappings) {
527 DeviceNameUtils::ParsedName p =
528 Name("/job:foo/replica:10/task:0/device:GPU:1");
529 EXPECT_EQ(absl::StrJoin(DeviceNameUtils::GetNamesForDeviceMappings(p), ","),
530 "/job:foo/replica:10/task:0/device:GPU:1,"
531 "/job:foo/replica:10/task:0/gpu:1");
532 p.has_task = false;
533 EXPECT_EQ(absl::StrJoin(DeviceNameUtils::GetNamesForDeviceMappings(p), ","),
534 "");
535 }
536
TEST(DeviceNameUtilsTest,CanonicalizeDeviceName)537 TEST(DeviceNameUtilsTest, CanonicalizeDeviceName) {
538 string canonical_name;
539 {
540 // Good basename.
541 string basename = "/job:foo/replica:10/task:0/device:CPU:0";
542 TF_EXPECT_OK(DeviceNameUtils::CanonicalizeDeviceName(
543 "/job:foo/replica:10/task:0/device:CPU:1", basename, &canonical_name));
544 EXPECT_EQ("/job:foo/replica:10/task:0/device:CPU:1", canonical_name);
545 TF_EXPECT_OK(DeviceNameUtils::CanonicalizeDeviceName(
546 "/job:foo/task:0/replica:10/device:CPU:1", basename, &canonical_name));
547 EXPECT_EQ("/job:foo/replica:10/task:0/device:CPU:1", canonical_name);
548 TF_EXPECT_OK(DeviceNameUtils::CanonicalizeDeviceName(
549 "/job:foo/task:0/replica:10/cpu:1", basename, &canonical_name));
550 EXPECT_EQ("/job:foo/replica:10/task:0/device:CPU:1", canonical_name);
551 TF_EXPECT_OK(DeviceNameUtils::CanonicalizeDeviceName("CPU:0", basename,
552 &canonical_name));
553 EXPECT_EQ("/job:foo/replica:10/task:0/device:CPU:0", canonical_name);
554 Status s = DeviceNameUtils::CanonicalizeDeviceName(
555 "/job:foo/task:0/replica/cpu:1", basename, &canonical_name);
556 EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
557 EXPECT_EQ("", canonical_name);
558 }
559
560 {
561 // Try out malformed basenames.
562 string fullname = "/device:CPU:0";
563
564 Status s = DeviceNameUtils::CanonicalizeDeviceName(
565 fullname, "/device:CPU:0", &canonical_name);
566 EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
567 EXPECT_EQ("", canonical_name);
568 s = DeviceNameUtils::CanonicalizeDeviceName(
569 fullname, "/job:foo/task:0/replica/cpu:1", &canonical_name);
570 EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
571 EXPECT_EQ("", canonical_name);
572 }
573 }
574
BM_ParseFullName(::testing::benchmark::State & state)575 static void BM_ParseFullName(::testing::benchmark::State& state) {
576 DeviceNameUtils::ParsedName p;
577 for (auto s : state) {
578 DeviceNameUtils::ParseFullName("/job:worker/replica:3/task:0/cpu:0", &p);
579 }
580 }
581 BENCHMARK(BM_ParseFullName);
582
583 } // namespace tensorflow
584