dist_mgr->GetCoordinationServiceAgent();
  if (coord_agent == nullptr) {
    status->status = tensorflow::errors::FailedPrecondition(
        "Coordination service is not enabled.");
    return;
  }
  absl::Duration timeout;
  if (timeout_in_ms > 0) {
    timeout = absl::Milliseconds(timeout_in_ms);
  } else {
    // Block until the key-value is set or the worker shuts down.
    timeout = absl::InfiniteDuration();
  }