experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Take device locality into account during prioritization.

Browse Source 
After this CL, if multiple devices with identical device type are viable for a placement of an op, the local device (if available) will be selected. (Prior to this change, the device whose job name comes first alphabetically would be selected.)

PiperOrigin-RevId: 312716604
Change-Id: I484c00cf0d34acc23c32ab8dd1cc5c394d32f0f3
This commit is contained in:
Jiri Simsa 2020-05-21 12:18:17 -07:00 committed by TensorFlower Gardener
parent eaa710abca
commit 37ab9c3bfc
2 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorflow/core/common_runtime/device_set.cc
View File

@ -116,12 +116,15 @@ void DeviceSet::SortPrioritizedDeviceVector(PrioritizedDeviceVector* vector) {
if (a_type_name != b_type_name) {
auto a_priority = DeviceFactory::DevicePriority(a_type_name);
auto b_priority = DeviceFactory::DevicePriority(b_type_name);
// First sort by prioritized device type (higher is preferred) and
// then by device name (lexicographically).
if (a_priority != b_priority) {
return a_priority > b_priority;
}
}
if (a.first->IsLocal() != b.first->IsLocal()) {
return a.first->IsLocal();
}
return StringPiece(a.first->name()) < StringPiece(b.first->name());
};
std::sort(vector->begin(), vector->end(), device_sort);

4
tensorflow/core/common_runtime/device_set.h
View File

@ -90,8 +90,8 @@ class DeviceSet {
//
// After a call to this function, the argument vector will be sorted by
// explicit priority (the second element in the `std::pair<DeviceType,
// int32>`), then by `DeviceTypeOrder` of the device type, and lastly
// by device name.
// int32>`), then by `DeviceTypeOrder` of the device type, then by device
// locality, and lastly by device name.
static void SortPrioritizedDeviceVector(PrioritizedDeviceVector* vector);
// Sorts a PrioritizedDeviceTypeVector according to types and explicit