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) {
|
||||
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, "/",
|
||||
|
@ -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.
|
||||
|
Loading…
Reference in New Issue
Block a user