Merge pull request #26921 from feihugis:TEST_PrefetchDatasetOp

PiperOrigin-RevId: 239471462
tensorflow · Mar 20, 2019 · 2749bed · 2749bed
2 parents 31606f7 + 7a4186b
commit 2749bed
Show file tree

Hide file tree

Showing 3 changed files with 657 additions and 1 deletion.
diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
@@ -393,6 +393,26 @@ tf_kernel_library(
     ],
 )
 
+tf_cc_test(
+    name = "prefetch_dataset_op_test",
+    size = "small",
+    srcs = ["prefetch_dataset_op_test.cc"],
+    deps = [
+        ":dataset_test_base",
+        ":dataset_utils",
+        ":iterator_ops",
+        ":prefetch_dataset_op",
+        ":tensor_slice_dataset_op",
+        "//tensorflow/core:core_cpu_internal",
+        "//tensorflow/core:dataset_ops_op_lib",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib_internal",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+    ],
+)
+
 tf_kernel_library(
     name = "repeat_dataset_op",
     srcs = ["repeat_dataset_op.cc"],

diff --git a/tensorflow/core/kernels/data/prefetch_dataset_op.cc b/tensorflow/core/kernels/data/prefetch_dataset_op.cc
@@ -386,7 +386,10 @@ void PrefetchDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
                  ParseScalarArgument<int64>(ctx, "buffer_size", &buffer_size));
   OP_REQUIRES(ctx,
               buffer_size >= 0 || buffer_size == PrefetchAutotuner::kAutoTune,
-              errors::InvalidArgument("buffer_size must be >= 0"));
+              errors::InvalidArgument("buffer_size must be >= 0 or set "
+                                      "buffer_size to be ",
+                                      PrefetchAutotuner::kAutoTune,
+                                      " for auto-tuning"));
 
   if (buffer_size == PrefetchAutotuner::kAutoTune) {
     metrics::RecordTFDataAutotune(kDatasetName);