TF 2.3 S3 client having permission issues with `tf.data.TFRecordDataset()`

[shaowei-su]

Please make sure that this is a bug. As per our
GitHub Policy,
we only address code/doc bugs, performance issues, feature requests and
build/installation issues on GitHub. tag:bug_template

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): no
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:
• TensorFlow installed from (source or binary):
• TensorFlow version (use command below): 2.3
• Python version: 3.6
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version:N/A
• GPU model and memory:

You can collect some of this information using our environment capture
script
You can also obtain the TensorFlow version with:

1. TF 1.0: python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)"
2. TF 2.0: python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"

Describe the current behavior
TF 2.3 tf.data.TFRecordDataset() failed to load files from S3.
Side note:
tf_file_io is working properly for the same setup.

>> from tensorflow.python.lib.io import file_io as tf_file_io
>> tf_file_io.file_exists('SOME_S3_FILE')

True

See error log below.
Describe the expected behavior
Same API and environments works well with TF 2.1 and TF 1.15.

Standalone code to reproduce the issue
Provide a reproducible test case that is the bare minimum necessary to generate
the problem. If possible, please share a link to Colab/Jupyter/any notebook.

dataset = tf.data.TFRecordDataset('SOME_S3_PATH')
for raw_record in dataset.take(10):
  print(repr(raw_record))

Other info / logs Include any logs or source code that would be helpful to
diagnose the problem. If including tracebacks, please include the full
traceback. Large logs and files should be attached.

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/default_user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 736, in __next__
    return self.next()
  File "/home/default_user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 772, in next
    return self._next_internal()
  File "/home/default_user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 764, in _next_internal
    return structure.from_compatible_tensor_list(self._element_spec, ret)
  File "/home/default_user/.conda/envs/user/lib/python3.6/contextlib.py", line 99, in __exit__
    self.gen.throw(type, value, traceback)
  File "/home/default_user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/eager/context.py", line 2105, in execution_mode
    executor_new.wait()
  File "/home/default_user/.conda/envs/user/lib/python3.6/site-packages/tensorflow/python/eager/executor.py", line 67, in wait
    pywrap_tfe.TFE_ExecutorWaitForAllPendingNodes(self._handle)
tensorflow.python.framework.errors_impl.FailedPreconditionError: AWS Credentials have not been set properly. Unable to access the specified S3 location

[Saduf2019]

@shaowei-su
Please provide with complete stand alone code for us to replicate the issue faced, or if possible share a colab gist with error reported.

[shaowei-su]

Thank you @Saduf2019 taking look into this.
To reproduce the error it requires:

• S3 bucket access
• AWS credential setup for the given S3 bucket

(Unfortunately, my current setup is not public accessible/sharable)
Once its setup then you can run the following tests:

1. Smoke test (should pass)

from tensorflow.python.lib.io import file_io as tf_file_io
tf_file_io.file_exists('SOME_S3_TFRECORD_FILE')

2. TFRecordDataset load (should fail)

dataset = tf.data.TFRecordDataset('SOME_S3_TFRECORD_FILE')
for raw_record in dataset.take(10):
  print(repr(raw_record))

[shaowei-su]

Quick updates on my end:
This issue seems related to the S3 multi part download that's introduced in TF2.2.
By disable the functionality with os.environ['S3_DISABLE_MULTI_PART_DOWNLOAD'] = '1', tf.data.TFRecordDataset() can load data from S3 properly.