| 97 | Placer::~Placer() {} |
| 98 | |
| 99 | Status Placer::Run() { |
| 100 | if (devices_->devices().empty()) { |
| 101 | return errors::FailedPrecondition("No devices are registered"); |
| 102 | } |
| 103 | |
| 104 | if (VLOG_IS_ON(3)) { |
| 105 | DumpGraphToFile("placer_input", *graph_, nullptr); |
| 106 | } |
| 107 | if (VLOG_IS_ON(5)) { |
| 108 | for (const Node* node : graph_->op_nodes()) { |
| 109 | VLOG(5) << " " << node->name() << ": requested: '" |
| 110 | << node->requested_device() << "' assigned: '" |
| 111 | << node->assigned_device_name() << "'"; |
| 112 | } |
| 113 | } |
| 114 | |
| 115 | std::string cpu_name, gpu_name, all_names; |
| 116 | for (auto d : devices_->devices()) { |
| 117 | all_names += d->name(); |
| 118 | all_names += ";\n"; |
| 119 | if (d->name().find("device:CPU:") != std::string::npos) { |
| 120 | cpu_name = d->name(); |
| 121 | } else if (d->name().find("device:GPU:") != std::string::npos) { |
| 122 | gpu_name = d->name(); |
| 123 | } |
| 124 | } |
| 125 | |
| 126 | bool use_session_group = false; |
| 127 | TF_CHECK_OK(tensorflow::ReadBoolFromEnvVar( |
| 128 | "USE_SESSION_GROUP", false, &use_session_group)); |
| 129 | if (use_session_group) { |
| 130 | for (Node* node : graph_->op_nodes()) { |
| 131 | const std::string dname = node->requested_device(); |
| 132 | if (!dname.empty()) { |
| 133 | if (dname.find("device:CPU:") != std::string::npos) { |
| 134 | node->set_requested_device(cpu_name); |
| 135 | } else if (dname.find("device:GPU:") != std::string::npos) { |
| 136 | node->set_requested_device(gpu_name); |
| 137 | } else { |
| 138 | LOG(ERROR) << "Can not find requested device in current devices set" |
| 139 | << ", node requested device: " << dname |
| 140 | << ", current devices set: " << all_names; |
| 141 | } |
| 142 | } |
| 143 | // we don't continue here, cause there are |
| 144 | // nodes should be assigned device below. |
| 145 | } |
| 146 | } |
| 147 | |
| 148 | FunctionStack stack(function_name_); |
| 149 | ColocationGraph colocation_graph(graph_, stack, flib_def_, devices_, |
| 150 | default_local_device_, allow_soft_placement_, |
| 151 | log_device_placement_); |
| 152 | |
| 153 | TF_RETURN_IF_ERROR(colocation_graph.Initialize()); |
| 154 | |
| 155 | // For each node, assign a device based on the constraints in the disjoint |
| 156 | // node set. |