[tf.data service] Fix flaky GcUnusedJob test.

The dataset needs more than 10 elements because by default we prefetch 16, creating a race condition where the job may finish before being garbage collected, resulting in self.assertEqual(worker._num_tasks(), 1) failing because there are 0 tasks.

PiperOrigin-RevId: 331203813
Change-Id: Iaaff7ac1e1c199c3606a3e51910dc6a704aa1664
This commit is contained in:
Andrew Audibert 2020-09-11 12:48:29 -07:00 committed by TensorFlower Gardener
parent fd1f787727
commit 0f8656631d

View File

@ -547,12 +547,12 @@ class DataServiceOpsTest(test_base.DatasetTestBase, parameterized.TestCase):
dispatcher = self.start_dispatch_server(
job_gc_check_interval_ms=50, job_gc_timeout_ms=20)
worker = self.start_worker_server(dispatcher) # pylint: disable=unused-variable
num_elements = 10
num_elements = 100
ds = _make_distributed_range_dataset(
num_elements, dispatcher, job_name=job_name)
it = iter(ds)
self.assertEqual(0, next(it).numpy())
self.assertEqual(1, worker._num_tasks())
self.assertEqual(next(it).numpy(), 0)
self.assertEqual(worker._num_tasks(), 1)
del it
while worker._num_tasks() > 0:
time.sleep(0.1)