[BEAM-6027] Fix slow downloads when reading from GCS

[fabito]

Overrides io.RawIOBase.readall in filesystemio.DownloaderStream as proposed in BEAM-6027.
It improves download time in ~40x.

[udim]

Thank you for your contribution!
Very cool

[udim]

cc: @chamikaramj
Can you run some benchmarks and include your results here?

[fabito]

Hi @udim ,

Using this snippet:

import tempfile
import timeit

from apache_beam.io.filesystems import FileSystems
from apache_beam.io.gcp import gcsio
from apache_beam.io.filesystemio import DownloaderStream


# https://issues.apache.org/jira/browse/BEAM-6027
def downloader_stream_readall(self):
    res = []
    while True:
        data = self.read(gcsio.DEFAULT_READ_BUFFER_SIZE)
        if not data:
            break
        res.append(data)
    return b''.join(res)


original_read_all = DownloaderStream.readall


if __name__ == '__main__':
    test_file = 'gs://cloud-samples-tests/vision/saigon.mp4'
    num_executions = 1

    def test_original():
        DownloaderStream.readall = original_read_all
        with FileSystems.open(test_file) as audio_file:
            with tempfile.NamedTemporaryFile(mode='w+b') as temp:
                temp.write(audio_file.read())

    def test_refactored():
        DownloaderStream.readall = downloader_stream_readall
        with FileSystems.open(test_file) as audio_file:
            with tempfile.NamedTemporaryFile(mode='w+b') as temp:
                temp.write(audio_file.read())

    print(timeit.timeit("test_original()", setup="from __main__ import test_original", number=num_executions))
    print(timeit.timeit("test_refactored()", setup="from __main__ import test_refactored", number=num_executions))

I got the following output:

120.99772120200214
1.0684915780002484

Hope that helps

[chamikaramj]

Thanks for the update. This looks great.

Seems like the file you used for your microbenchmark is about 4MB which will be within the first chunk for the new buffer size. Can you try running a Beam pipeline with a larger input (say 10GB) with Dataflow to confirm that there's no regression at large scale ?

[fabito]

Yes I can. Any advice on how this pipeline would be and how can we measure the reading performance?
Maybe something like this:

    with beam.Pipeline(options=pipeline_options) as pipeline:
        _ = (
            pipeline
            | 'Read 10Gb file' >> beam.io.ReadAllFromText('gs://bucket/10Gb.txt')
            | 'Write file' >> beam.io.WriteToText('gs://bucket/10Gb_copy*.txt')
        )

[chamikaramj]

Yeah, that pipeline looks good. End-to-end execution time and Total vCPU time show in Dataflow console should be good metrics to compare.

[fabito]

I ran the same pipeline, first with the original code and after with the new implementation of read_all. I used a text file with ~8.81Gb. Apparently performance wasn't affected. Check below the console snapshots from both executions:

Before the change:

After the change:

[chamikaramj]

Thanks. LGTM.

Thanks for fixing this.

Added one comment. Also please fixup your commits into a single commit for merging.

[chamikaramj]

Where is this function used ?

Prob. remove if unused.

[chamikaramj]

Ah actually seems like you are overriding the function here: https://docs.python.org/3/library/io.html#io.IOBase

[chamikaramj]

Sorry, still have a question.

Does Beam call readlll() function anywhere ? I couldn't find a usage. Beam textio for example, invokes read() not readall().
https://github.com/apache/beam/blob/master/sdks/python/apache_beam/io/textio.py#L272

If it does, I'm not sure what will prevent us from reading a huge amount of data into memory and running into OOMs.

[udim]

I only found this usage in ReadableFile (relatively new) where we don't specify the size:

beam/sdks/python/apache_beam/io/fileio.py

Lines 150 to 154 in 1382505

	def open(self, mime_type='text/plain'):
	return filesystems.FileSystems.open(self.metadata.path)
	def read(self):
	return self.open().read()

[chamikaramj]

That makes sense. I think ReadableFile is intended for small files. But probably we should add a readall() method there as well and update read() to take a buffer (not in this PR).

cc: @pabloem

[chamikaramj]

Thanks. I'll squash and merge.

[udim]

Adding note to not forget to "run python postcommit" before merging