.ok());
  ASSERT_TRUE(worker_server2->Start().ok());

  TFE_Context* ctx_0 = CreateContext(serialized_server_def_0,
                                     /*isolate_session_state=*/false,
                                     /*init_timeout_in_ms=*/0);
  TFE_Context* ctx_1 = CreateContext(serialized_server_def_1,
                                     /*isolate_session_state=*/false,

  ASSERT_TRUE(worker_server2->Start().ok());

  // Create two contexts.
  int32_t init_timeout_in_ms = 300000;
  TFE_Context* ctx_0 =
      CreateContext(serialized_server_def_0,
                    /*isolate_session_state=*/false, init_timeout_in_ms);
  TFE_Context* ctx_1 =
      CreateContext(serialized_server_def_1,
                    /*isolate_session_state=*/false, init_timeout_in_ms);

  // Remote device on `worker2`.

TFE_TensorHandle* CreateVariable(TFE_Context* ctx, float value,
                                 const tensorflow::string& device_name,
                                 const tensorflow::string& variable_name);

TFE_Context* CreateContext(const std::string& serialized_server_def,
                           bool isolate_session_state,
                           int64_t init_timeout_in_ms);

tensorflow::ServerDef ReplaceTaskInServerDef(

  TFE_DeleteOp(op);
  if (TF_GetCode(status) != TF_OK) return nullptr;
  CHECK_EQ(0, num_retvals);
  TF_DeleteStatus(status);
  return var_handle;
}

TFE_Context* CreateContext(const std::string& serialized_server_def,
                           bool isolate_session_state,
                           int64_t init_timeout_in_ms) {
  TF_Status* status = TF_NewStatus();

// TODO(b/291142876) Simplify TFE_ContextSetServerDefWithTimeoutAndRetries and
// TFE_ContextUpdateServerDefWithTimeout to be simple wrappers around the same
// C++ function.
// Retries are used for CreateContext calls, which is used in
// ParameterServerStrategy initialization to be robust to worker preemption.
TF_CAPI_EXPORT extern void TFE_ContextSetServerDefWithTimeoutAndRetries(