VirtualPlacer use default device localhost if device is empty.
PiperOrigin-RevId: 161153235
This commit is contained in:
parent
eb8bebd670
commit
b031107cdf
@ -77,6 +77,9 @@ string VirtualPlacer::get_canonical_device_name(const NodeDef& node) const {
|
|||||||
if (!parsed) {
|
if (!parsed) {
|
||||||
return get_default_device_name();
|
return get_default_device_name();
|
||||||
} else {
|
} else {
|
||||||
|
if (parsed_name.job.empty()) {
|
||||||
|
parsed_name.job = "localhost";
|
||||||
|
}
|
||||||
device = strings::StrCat(
|
device = strings::StrCat(
|
||||||
"/job:", parsed_name.job, "/replica:", parsed_name.replica,
|
"/job:", parsed_name.job, "/replica:", parsed_name.replica,
|
||||||
"/task:", parsed_name.task, "/",
|
"/task:", parsed_name.task, "/",
|
||||||
|
@ -51,6 +51,29 @@ TEST(VirtualPlacerTest, LocalDevices) {
|
|||||||
placer.get_canonical_device_name(node));
|
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) {
|
TEST(VirtualPlacerTest, FallBackUnknown) {
|
||||||
// Virtual placer falls back to "UNKNOWN" only if there are no devices in the
|
// Virtual placer falls back to "UNKNOWN" only if there are no devices in the
|
||||||
// cluster.
|
// cluster.
|
||||||
|
Loading…
x
Reference in New Issue
Block a user