From ef94af1f7c911f04c7602a5ea27909bb8ea1fd66 Mon Sep 17 00:00:00 2001 From: Yuefeng Zhou Date: Mon, 9 Nov 2020 17:03:37 -0800 Subject: [PATCH] Add the unit test for a fixed big. PiperOrigin-RevId: 341515086 Change-Id: Ia6ba3e17a44d8d619748006f414bb217bc703ba7 --- .../coordinator/cluster_coordinator_test.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tensorflow/python/distribute/coordinator/cluster_coordinator_test.py b/tensorflow/python/distribute/coordinator/cluster_coordinator_test.py index 53f26d8736a..b38688ebaf5 100644 --- a/tensorflow/python/distribute/coordinator/cluster_coordinator_test.py +++ b/tensorflow/python/distribute/coordinator/cluster_coordinator_test.py @@ -658,6 +658,26 @@ class ClusterCoordinatorTest(TestCaseWithErrorReportingThread): 'only accepts a `tf.function` or a concrete function.'): self.coordinator.schedule(func, args=(1,)) + def testDatasetPartiallyCreatedOnCoordinator(self): + dataset = dataset_ops.DatasetV2.range(1, 10) + + @def_function.function + def input_fn(): + return dataset.shuffle(9) + + @def_function.function + def worker_fn(iterator): + x = next(iterator) + return x + + per_worker_dataset = self.coordinator.create_per_worker_dataset(input_fn) + self.coordinator.schedule(worker_fn, args=(iter(per_worker_dataset),)) + + with self.assertRaisesRegexp( + coordinator_lib.InputError, + 'error message is Failed copying input tensor from'): + self.coordinator.join() + class LimitedClosureQueueSizeBasicTest(ClusterCoordinatorTest): """Test basic functionality works with explicit maximum closure queue size.