tree b7e80069b2fb46aca976f53efed6ea2f482b5e3d
parent e687cab61615a95b8aea79120c9f168d4cc30955
author Andrew Audibert <aaudibert@google.com> 1610648404 -0800
committer TensorFlower Gardener <gardener@tensorflow.org> 1610648778 -0800

[tf.data service] Improve cancellation for tf.data service requests.

1. If a DataServiceDataset iterator is cancelled, it will now call TryCancel on its outstanding RPCs.
2. As a result, we can reduce the frequency of returning from blocked round-robin requests to check whether the iterator is cancelled. This may avoid delays in GetNext() that could happen if one consumer reads from a round earlier than others, and needs to perform multiple retries with exponential backoff.
3. Because of (2), server shutdown may take up to 1 minute if a round-robin request is blocked waiting for other consumers. To prevent slow unit tests, certain tests store their servers globally so that they are destroyed immediately at process exit without waiting for their outstanding RPCs to finish.

Running data_service_ops_test.py locally, this CL reduces the time from 27 seconds to 20 seconds

PiperOrigin-RevId: 351825888
Change-Id: Iba20a456bdabf251d03b94f090fe760616d3da4d