| 142 | } |
| 143 | |
| 144 | Status GrpcServer::Init(const GrpcServerOptions& opts) { |
| 145 | mutex_lock l(mu_); |
| 146 | CHECK_EQ(state_, NEW); |
| 147 | master_env_.env = env_; |
| 148 | worker_env_.env = env_; |
| 149 | |
| 150 | // Check parameters before DeviceFactory::AddDevices, |
| 151 | // otherwise if 'task_index=-1' the program will abort. |
| 152 | |
| 153 | int requested_port; |
| 154 | TF_RETURN_IF_ERROR(GetPort(&requested_port)); |
| 155 | |
| 156 | SessionOptions sess_opts; |
| 157 | ConfigProto config = server_def_.default_session_config(); |
| 158 | sess_opts.config = config; |
| 159 | |
| 160 | // Configure shared devices between master and worker. |
| 161 | string name_prefix = |
| 162 | strings::StrCat("/job:", server_def_.job_name(), "/replica:0", |
| 163 | "/task:", server_def_.task_index()); |
| 164 | std::vector<std::unique_ptr<Device>> devices; |
| 165 | TF_RETURN_IF_ERROR( |
| 166 | DeviceFactory::AddDevices(sess_opts, name_prefix, &devices)); |
| 167 | worker_env_.device_mgr = new DeviceMgr(std::move(devices)); |
| 168 | master_env_.local_devices = worker_env_.device_mgr->ListDevices(); |
| 169 | worker_env_.local_devices = worker_env_.device_mgr->ListDevices(); |
| 170 | worker_env_.rendezvous_mgr = opts.rendezvous_mgr_func == nullptr |
| 171 | ? new RpcRendezvousMgr(&worker_env_) |
| 172 | : opts.rendezvous_mgr_func(&worker_env_); |
| 173 | string unused; |
| 174 | string default_worker_name; |
| 175 | if (!DeviceNameUtils::SplitDeviceName(master_env_.local_devices[0]->name(), |
| 176 | &default_worker_name, &unused)) { |
| 177 | return errors::Internal("Could not parse worker name."); |
| 178 | } |
| 179 | |
| 180 | // N.B. The order of initialization here is intricate, because we |
| 181 | // wish to allow `requested_port == 0` (for choosing any port, |
| 182 | // mostly for testing). Therefore, the construction of the channel |
| 183 | // and worker caches depends on `bound_port_`, which is not set |
| 184 | // until we call `builder.BuildAndStart()`. We must create the |
| 185 | // service objects before calling `builder.BuildAndStart()`, but |
| 186 | // `master_env_` and `worker_env_` are only partially |
| 187 | // configured. However, this is not dangerous, because we do not |
| 188 | // start serving requests until `this->Start()` is called, which |
| 189 | // happens after this method returns. |
| 190 | // |
| 191 | // TODO(mrry): Provide a general mechanism for dynamically setting |
| 192 | // the identities of tasks in the worker pool after the service is |
| 193 | // running. |
| 194 | ::grpc::ServerBuilder builder; |
| 195 | builder.AddListeningPort(strings::StrCat("0.0.0.0:", requested_port), |
| 196 | GetServerCredentials(server_def_), &bound_port_); |
| 197 | builder.SetMaxMessageSize(std::numeric_limits<int32>::max()); |
| 198 | |
| 199 | builder.SetOption( |
| 200 | std::unique_ptr<::grpc::ServerBuilderOption>(new NoReusePortOption)); |
| 201 | // Allow subclasses to specify more args to pass to the gRPC server. |
no test coverage detected