AWS S3 DataPipe doesn't terminate when used with .load_from_tar(mode="r|")

[NivekT]

🐛 Describe the bug

Note this issue only occurs when mode="r|", not when it is mode="r:".

Sample code:

dp = IterableWrapper(["s3://BUCKET/images0.tar"])
# dp = dp.load_files_by_s3(region="us-east-1").load_from_tar(mode="r:")  # This is fast and works as expected
dp = dp.load_files_by_s3(region="us-east-1").load_from_tar(mode="r|")  # This is slow and doesn't terminate, it repeatedly yields certain files
start = time.time()
for x in dp:
    print(x)
print(time.time() - start)

Not sure if mode="r|" is incompatible with S3 DataPipe/SDK or there is something off with our implementation.

Versions

main branch

cc: @ejguan

[ejguan]

Here are what I have found regarding S3Handler:

Originally I was expecting the returned stream from S3handler is non-seekable stream. But, it turns out that the whole archive/files will be dumped into memory based on the implementation here:

• C++ side:

  data/torchdata/csrc/pybind/pybind.cpp

  Line 28 in a435c7f

  return py::bytes(result);

• Python side:

  data/torchdata/datapipes/iter/load/s3io.py

  Line 135 in a435c7f

  yield url, StreamWrapper(BytesIO(self.handler.s3_read(url)))

  
  And, that is the reason that the performance seems on parity with or without BytesIO in this issue.

In order to make it streaming, we need to have a way to pybind C++ stream IO to python, which is non-trivial. See a code reference: https://github.com/CadQuery/OCP/blob/master/pystreambuf.h

However, it doesn't give us an answer about why it's not working with mode r| properly. I did test those two modes on a local tar archive, and both modes are working properly.

Then, it definitely a bug regarding S3Loader.

Edit: Tested S3FileLoader, the performance is on-par for both modes.....

[NivekT]

The main branch is fine. The root cause is here.

Sorry for the false alarm :/

One discovery is that S3 DataPipe downloads the whole archive into memory instead of chunks. We would prefer it to stream. I'll look into fsspec shortly to examine its behavior.