message: Override buffer size when running on cloud tpu.
  A small buffer size adds a performance slowdown. Here we override the buffer size to a minimum recommended buffer size.
diff: @@ -38,6 +38,15 @@ namespace data {
  constexpr char kCurrentFileIndex[] = "current_file_index";
  constexpr char kOffset[] = "offset";
  +constexpr char kGcsFsPrefix[] = "gs://";
  +constexpr int64 kCloudTpuBlockSize = 127LL << 20; // 127MB.
  +
  +bool is_cloud_tpu_gcs_fs() {
  +#if defined(PLATFORM_CLOUD_TPU) && defined(TPU_GCS_FS)
  + return true;
  +#endif
  + return false;
  +}
  class TFRecordDatasetOp::Dataset : public DatasetBase {
  public:
  @@ -224,11 +233,13 @@ void TFRecordDatasetOp::MakeDataset(OpKernelContext* ctx,
  ctx, filenames_tensor->dims() <= 1,
  errors::InvalidArgument("`filenames` must be a scalar or a vector."));
  + bool is_gcs_fs = true;
  std::vector<string> filenames;
  filenames.reserve(filenames_tensor->NumElements());
  for (int i = 0; i < filenames_tensor->NumElements(); ++i) {
  VLOG(2) << "Reading file: " << filenames_tensor->flat<tstring>()(i);
  filenames.push_back(filenames_tensor->flat<tstring>()(i));
  + is_gcs_fs &= absl::StartsWith(filenames[i], kGcsFsPrefix);
  }
  tstring compression_type;
  @@ -242,6 +253,14 @@ void TFRecordDatasetOp::MakeDataset(OpKernelContext* ctx,
  errors::InvalidArgument(
  "`buffer_size` must be >= 0 (0 == no buffering)"));
  + if (is_gcs_fs && is_cloud_tpu_gcs_fs() && buffer_size < kCloudTpuBlockSize) {
  + LOG(WARNING) << "User buffer size is too small for reading Cloud TPU "
  + << "TFRecords stored in GCS. Overriding " << buffer_size
  + << " to the minimum recommended buffer_size = "
  + << kCloudTpuBlockSize;
  + buffer_size = kCloudTpuBlockSize;
  + }
  +
  *output =
  new Dataset(ctx, std::move(filenames), compression_type, buffer_size);
  }
  