Add the functionality to read chunk by chunk

🚀 The feature

This issue is a continuation of the discussion from here.

We are proposing to add some functionality to allow reading chunk by chunk from a stream. Ideally, this should be done without loading the entire stream/response into memory first, but this may not be possible in all cases.

Motivation, pitch

From @pmeier

Having something like a ChunkReader datapipe can be useful in general not just for HTTP requests, right?

from torchdata.datapipes.iter import IterDataPipe


class ChunkReader(IterDataPipe):
    def __init__(self, datapipe, *, chunk_size=32 * 1024 * 1024):
        self.datapipe = datapipe
        self.chunk_size = chunk_size

    def __iter__(self):
        for path, stream in self.datapipe:
            for chunk in iter(lambda: stream.read(self.chunk_size), b""):
                # filter out keep-alive new chunks from HTTP streams
                if not chunk:
                    continue

                yield path, chunk

A utility like that is needed anyway for Saver anyway. AFAIK, there is currently no builtin functionality to read data from a stream, correct? I'm guessing that is why the Saver example uses bytes as input.

This can be useful for reading from file stream in chunks (e.g. after FileOpener), writing to files in chunks (e.g. Saver, and etc (e.g. HttpReader).

Alternatives

Instead of adding a new DataPipe, we can consider modifying specific existing DataPipes to add the functionality of reading by chunk. This may be useful in avoiding the need to read the entire stream/response into memory (as I believe it is the case in the current implementation of HttpReader, please correct me if this is wrong).

We can also do both.

Additional context

No response

cc: @ejguan @VitalyFedyunin

Start by reading the existing FileOpener, Saver, and HttpReader DataPipes, especially the Saver example at torchdata/datapipes/iter/util/saver.py. Review the linked discussion to determine whether a new ChunkReader or changes to existing DataPipes are expected. Done means streams can be consumed chunk by chunk without unnecessarily loading the entire response into memory, with the behavior covered by appropriate tests.