| 173 | } |
| 174 | |
| 175 | tensorflow::Status UpdateTFE_ContextWithServerDef( |
| 176 | int keep_alive_secs, const tensorflow::ServerDef& server_def, |
| 177 | TFE_Context* ctx) { |
| 178 | // We don't use the TF_RETURN_IF_ERROR macro directly since that destroys the |
| 179 | // server object (which currently CHECK-fails) and we miss the error, instead, |
| 180 | // we log the error, and then return to allow the user to see the error |
| 181 | // message. |
| 182 | #define LOG_AND_RETURN_IF_ERROR(...) \ |
| 183 | do { \ |
| 184 | const ::tensorflow::Status _status = (__VA_ARGS__); \ |
| 185 | if (TF_PREDICT_FALSE(!_status.ok())) { \ |
| 186 | LOG(ERROR) << _status.error_message(); \ |
| 187 | return _status; \ |
| 188 | } \ |
| 189 | } while (0); |
| 190 | |
| 191 | string worker_name = |
| 192 | tensorflow::strings::StrCat("/job:", server_def.job_name(), |
| 193 | "/replica:0/task:", server_def.task_index()); |
| 194 | |
| 195 | std::unique_ptr<tensorflow::ServerInterface> server; |
| 196 | LOG_AND_RETURN_IF_ERROR(tensorflow::NewServer(server_def, &server)); |
| 197 | |
| 198 | tensorflow::GrpcServer* grpc_server = |
| 199 | dynamic_cast<tensorflow::GrpcServer*>(server.get()); |
| 200 | if (grpc_server == nullptr) { |
| 201 | LOG_AND_RETURN_IF_ERROR(tensorflow::errors::Internal( |
| 202 | "Currently, TFE_NewContext only supports tensorflow::GrpcServer.")); |
| 203 | } |
| 204 | |
| 205 | tensorflow::uint64 context_id = tensorflow::EagerContext::NewContextId(); |
| 206 | // Make master eager context accessible by local eager service, which might |
| 207 | // receive send tensor requests from remote workers. |
| 208 | LOG_AND_RETURN_IF_ERROR(grpc_server->AddMasterEagerContextToEagerService( |
| 209 | context_id, ctx->context)); |
| 210 | |
| 211 | std::vector<string> remote_workers; |
| 212 | grpc_server->master_env()->worker_cache->ListWorkers(&remote_workers); |
| 213 | remote_workers.erase( |
| 214 | std::remove(remote_workers.begin(), remote_workers.end(), worker_name), |
| 215 | remote_workers.end()); |
| 216 | |
| 217 | std::unique_ptr<tensorflow::DeviceMgr> remote_device_mgr; |
| 218 | LOG_AND_RETURN_IF_ERROR(GetAllRemoteDevices( |
| 219 | remote_workers, grpc_server->master_env()->worker_cache, |
| 220 | &remote_device_mgr)); |
| 221 | |
| 222 | std::vector<tensorflow::DeviceAttributes> cluster_device_attributes; |
| 223 | remote_device_mgr->ListDeviceAttributes(&cluster_device_attributes); |
| 224 | |
| 225 | std::vector<tensorflow::DeviceAttributes> local_device_attributes; |
| 226 | grpc_server->worker_env()->device_mgr->ListDeviceAttributes( |
| 227 | &local_device_attributes); |
| 228 | |
| 229 | // This request make sure that we can create Rendevzous properly between |
| 230 | // Local and Remote context. |
| 231 | tensorflow::eager::CreateContextRequest base_request; |
| 232 | for (const auto& da : cluster_device_attributes) { |
no test coverage detected