mutex_lock session_lock(session->mu);
    const Graph& graph = session->graph->graph;

    const string& mutation_warning = session->graph->sessions[session];
    if (!mutation_warning.empty()) {
      // TODO(b/74949947): turn this back into an error status
      LOG(WARNING) << mutation_warning;
      session->graph->sessions[session].clear();
    }

    const auto num_nodes = graph.num_node_ids();

      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();
  }