tree c12bbbd08b1a85d01ecba7f5a35a90fadbefc6a5
parent 01ce6f2f76917839a738be3571d10bbdae5cc6d6
author Ran Chen <crccw@google.com> 1628164379 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> 1628164703 -0700

Improve cancellation of MultiDeviceIterator::GetNextFromShard

Cancelling GetNextFromShard shouldn't cancel the whole MultiDeviceIterator. Although in practice tf.distribute always fetch from all shards, it's possible that one shard iterator is deleted while the others are still fetching. This didn't surface any problems before because RemoteCall didn't propagate the cancellation.

This change:
1. Propagate cancellation in RemoteCall.
2. Cancelling a GetNextFromShard only cancels itself. It no longer cancels MultiDeviceIterator. The cancellation is not efficient but it should be rare.

PiperOrigin-RevId: 388906528
Change-Id: Iadc4e34d401aecd071a7b0f86ce9f0580abe9ef5