Function MakeGRPCCluster

tensorflow/core/distributed_runtime/rpcbench_test.cc:44–100

static thread::ThreadPool* worker_threads;

void MakeGRPCCluster(const SessionOptions& options, int n,
                     std::vector<string>* workers,
                     std::vector<DeviceAttributes>* devices) {
  CHECK_GE(n, 1);

  workers->clear();
  std::vector<int> port(n);
  for (int i = 0; i < n; ++i) {
    port[i] = testing::PickUnusedPortOrDie();
    workers->push_back(strings::StrCat("grpc://localhost:", port[i]));
  }

  int num_cpus = 1;
  int num_gpus = 0;
  auto iter = options.config.device_count().find("CPU");
  if (iter != options.config.device_count().end()) {
    num_cpus = iter->second;
  }
  iter = options.config.device_count().find("GPU");
  if (iter != options.config.device_count().end()) {
    num_gpus = iter->second;
  }

  worker_threads = new thread::ThreadPool(Env::Default(), "worker_threads", n);
  for (int worker_idx = 0; worker_idx < n; ++worker_idx) {
    worker_threads->Schedule([worker_idx, n, num_cpus, num_gpus, &port] {
      ServerDef server;
      server.set_protocol("grpc");
      server.set_job_name("localhost");
      server.set_task_index(worker_idx);

      auto job_def = server.mutable_cluster()->add_job();
      job_def->set_name("localhost");
      for (int i = 0; i < n; i++) {
        (*(job_def->mutable_tasks()))[i] =
            strings::StrCat("localhost:", port[i]);
      }

      auto config = server.mutable_default_session_config();
      (*config->mutable_device_count())["CPU"] = num_cpus;
      (*config->mutable_device_count())["GPU"] = num_gpus;

      std::unique_ptr<ServerInterface> svr;
      TF_CHECK_OK(NewServer(server, &svr));
      TF_CHECK_OK(svr->Start());
      TF_CHECK_OK(svr->Join());
    });
  }

  // Get attributes for all devices.
  LOG(ERROR) << "W '" << (*workers)[0] << "'";
  SessionOptions options_copy(options);
  options_copy.target = (*workers)[0];
  std::unique_ptr<GrpcSession> session;
  TF_CHECK_OK(GrpcSession::Create(options_copy, &session));
  TF_CHECK_OK(session->ListDevices(devices));
}