VirtualPlacer use default device localhost if device is empty.

PiperOrigin-RevId: 161153235
This commit is contained in:
A. Unique TensorFlower 2017-07-06 17:28:37 -07:00 committed by TensorFlower Gardener
parent eb8bebd670
commit b031107cdf
2 changed files with 26 additions and 0 deletions

View File

@ -77,6 +77,9 @@ string VirtualPlacer::get_canonical_device_name(const NodeDef& node) const {
if (!parsed) {
return get_default_device_name();
} else {
if (parsed_name.job.empty()) {
parsed_name.job = "localhost";
}
device = strings::StrCat(
"/job:", parsed_name.job, "/replica:", parsed_name.replica,
"/task:", parsed_name.task, "/",

View File

@ -51,6 +51,29 @@ TEST(VirtualPlacerTest, LocalDevices) {
placer.get_canonical_device_name(node));
}
TEST(VirtualPlacerTest, EmptyJobBecomesLocalhost) {
// Virtual placer should use "localhost" if device is empty.
// First create a cluster with only localhost devices.
std::unordered_map<string, DeviceProperties> devices;
DeviceProperties cpu_device;
cpu_device.set_type("CPU");
devices["/job:localhost/replica:0/task:0/cpu:0"] = cpu_device;
DeviceProperties gpu_device;
gpu_device.set_type("GPU");
devices["/job:localhost/replica:0/task:0/gpu:0"] = gpu_device;
VirtualCluster cluster(devices);
VirtualPlacer placer(&cluster);
NodeDef node;
node.set_op("Conv2D");
node.set_device("/device:CPU:0");
EXPECT_EQ("/job:localhost/replica:0/task:0/cpu:0",
placer.get_canonical_device_name(node));
node.set_device("/device:GPU:0");
EXPECT_EQ("/job:localhost/replica:0/task:0/gpu:0",
placer.get_canonical_device_name(node));
}
TEST(VirtualPlacerTest, FallBackUnknown) {
// Virtual placer falls back to "UNKNOWN" only if there are no devices in the
// cluster.