}

TF_CAPI_EXPORT extern void TFE_AbortCollectiveOps(TFE_Context* ctx,
                                                  TF_Status* status) {
  tensorflow::EagerContext* context =
      tensorflow::ContextFromInterface(tensorflow::unwrap(ctx));
  auto collective_executor_handle = context->GetCollectiveExecutorHandle();
  collective_executor_handle->get()->StartAbort(status->status);
}

#define EXPECT_TF_OK(x) EXPECT_EQ(TF_OK, TF_GetCode(x)) << TF_Message(x)

static const char* content = "abcdefghijklmnopqrstuvwxyz1234567890";
// We will work with content_view instead of content.
static const absl::string_view content_view = content;

namespace gcs = google::cloud::storage;

static std::string InitializeTmpDir() {
  // This env should be something like `gs://bucket/path`

    }
    std::vector<MaybeParallelTensorOwned> result_content;
    result_content.reserve(1);
    result_content.push_back(ParallelTensor::FromTensorHandles(
        parallel_device, std::move(components), status));
    if (TF_GetCode(status) != TF_OK) return result;
    result.emplace(std::move(result_content));
    return result;
  } else if (operation_name == std::string("TPUReplicatedOutput")) {

      (TF_GetCode(status) != TF_OUT_OF_RANGE)) {
    return -1;
  }
  int64_t read;
  auto content_length = stream.headers().find("content-length");
  if (content_length == stream.headers().end()) {
    // When we read a file with offset that is bigger than the actual file size.
    // GCS will return an empty header (e.g no `content-length` header). In this
    // case, we will set read to `0` and continue.
    read = 0;

    auto fcmp = block_map_.upper_bound(fmax);
    if (fcmp != block_map_.begin() && key < (--fcmp)->first) {
      return TF_SetStatus(status, TF_INTERNAL,
                          "Block cache contents are inconsistent.");
    }
  }

  Trim();

  return TF_SetStatus(status, TF_OK, "");
}

void RamFileBlockCache::MaybeFetch(const Key& key,

          // Expect a partial read.
          EXPECT_EQ(got.size(), size - offset)
              << "block size = " << block_size << ", offset = " << offset
              << ", n = " << n;
        }
        // Verify the contents of the read.
        std::vector<char>::const_iterator begin = buf.begin() + offset;
        std::vector<char>::const_iterator end =
            offset + n > buf.size() ? buf.end() : begin + n;

                                    TF_Buffer* buf) {
  string content;
  cell->cell.value().SerializeToString(&content);
  void* data = tensorflow::port::Malloc(content.length());
  content.copy(static_cast<char*>(data), content.length(), 0);
  buf->data = data;
  buf->length = content.length();
  buf->data_deallocator = [](void* data, size_t length) {
    tensorflow::port::Free(data);
  };

}

// Set server def with timeout.
TF_CAPI_EXPORT extern void TFE_ContextSetServerDefWithTimeout(
    TFE_Context* ctx, int keep_alive_secs, const void* proto, size_t proto_len,
    int64_t init_timeout_in_ms, TF_Status* status,
    bool clear_existing_contexts) {
  TFE_ContextSetServerDefWithTimeoutAndRetries(
      ctx, keep_alive_secs, proto, proto_len, init_timeout_in_ms,

                                     init_timeout_in_ms, status,
                                     /*clear_existing_contexts=*/false);
  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
  TFE_DeleteContextOptions(opts);
  TF_DeleteStatus(status);
  return ctx;
}

tensorflow::ServerDef ReplaceTaskInServerDef(
    const tensorflow::ServerDef& server_def, int task_index) {

                  .ok());
  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,