Configuring GCS 429 retry handling in TensorFlow tfds.load

Short description
When performing intensive parallel model training, Twitter occasionally exceeds GCS bucket throughput limitations. If all 10 retry attempts fail, then their code correspondingly fails as well:

tensorflow.python.framework.errors_impl.AbortedError: All 10 retry attempts failed. The last failure: Unavailable: Error executing an HTTP request: HTTP response code 429 with body

429 GCS Egress Too High This workload is drawing too much egress bandwidth from Google Cloud Storage and has triggered the Cloud Storage Egress Bandwidth Cap. Reduce the rate of request or contact GCP customer support if you want to increase the bandwidth cap.

In Twitter’s experience with the GCS Connector for Hadoop, the connector has been tuned to handle 429 gracefully (link). For these jobs with the TensorFlow library, they are curious about similar possible retry tunables, and are separately requesting additional throughput cap and considering tuning their parallelism as well.

From the error message, it appears that the library retries the data retrieval 10 times. Is it possible to configure the retry attempts performed by the library when attempting to load data from GCS? (I could not find an obvious parameter to do so.)

Making the retry behavior configurable, such as number of retries and time between retries (avoiding the exponential backoff pathological behavior of getting to "trying again tomorrow"), might provide a mechanism to choose the retry vs. failure tolerance for this type of bursty workloads -- a way to "smooth out the bumps" and avoid job failure during short periods of high throughput utilization.

Stack Trace

 File "/opt/ee/python/3.8/lib/python3.8/site-packages/large_embeddings/projects/recap/run_native.py", line 149, in <module>
    app.run(main)
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/absl/app.py", line 300, in run
    _run_main(main, args)
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/absl/app.py", line 251, in _run_main
    sys.exit(main(argv))
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/large_embeddings/projects/recap/run_native.py", line 141, in main
    maybe_run_training(
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/twitter/ml/torch/experimental/distributed/training.py", line 39, in maybe_run_training
    train_fn(**training_kwargs)
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/large_embeddings/projects/recap/run_native.py", line 115, in train
    train_and_evaluate(
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/large_embeddings/custom_training_loop.py", line 124, in train_and_evaluate
    outputs = train_step_fn()
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/large_embeddings/custom_training_loop.py", line 25, in step_fn
    outputs = pipeline.progress(data_iterator)
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/torchrec/distributed/train_pipeline.py", line 490, in progress
    batch_ip2 = next(dataloader_iter)
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 652, in __next__
    data = self._next_data()
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 692, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 39, in fetch
    data = next(self.dataset_iter)
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/twitter/ml/torch/experimental/datasets/dataset_lib.py", line 106, in _inner
    for elem in self.tf_dataset:
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 800, in __next__
    return self._next_internal()
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 783, in _next_internal
    ret = gen_dataset_ops.iterator_get_next(
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/tensorflow/python/ops/gen_dataset_ops.py", line 2845, in iterator_get_next
    _ops.raise_from_not_ok_status(e, name)
  File "/opt/ee/python/3.8/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 7107, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.AbortedError: All 10 retry attempts failed. The last failure: Error executing an HTTP request: HTTP response code 429 with body '<html><head><meta http-equiv="content-type" content="text/html; charset=utf-8"/><title>Sorry...</title><style> body { font-family: verdana, arial, sans-serif; background-color: #fff; color: #000; }</style></head><body><div><table><tr><td><b><font face=sans-serif size=10><font color=#4285f4>G</font><font color=#ea4335>o</font><font color=#fbbc05>o</font><font color=#4285f4>g</font><font color=#34a853>l</font><font color=#ea4335>e</font></font></b></td><td style="text-align: left; vertical-align: '

**Expected behavior**
Currently when the egress limit is hit and 429 error occurs, the retry mechanism tries 10 times and gives up. We want the ability to configure the retry mechanism and use exponential back off.

** Tensor Flow Version: 2.7 **

Start at the tfds.load entry point and trace how GCS data retrieval handles the 429 response and its current ten retries. Identify the existing retry tests or add coverage for configurable retry count and backoff behavior; done means short GCS throughput bursts no longer fail solely because the fixed retry limit is reached.