Download multiple GCS files in parallel (into memory) using Python

Question

I have a storage bucket with a lot of large files (500mb each). At times I need to load multiple files, referenced by name. I have been using the blob.download_as_string() function to download the files one-by-one, but it's extremely slow so I would like to try and download them in parallel instead.

I found the gcloud-aio-storage package, however the documentation is a bit sparse, especially for the download function.

I would prefer to download / store the files in memory instead of downloading to local machine then upload to script.

This is what I've pieced together, though I can't seem to get this to work. I keep getting a timeout error. What am I doing wrong?

Note: Using python 3.7, and latest of all other packages.

test_download.py

from gcloud.aio.storage import Storage
import aiohttp 
import asyncio

async def gcs_download(session, bucket_name, file, storage):
    async with session: 
        bucket = storage.get_bucket(bucket_name)
        blob = await bucket.get_blob(file)
        return  await blob.download()
    

async def get_gcsfiles_async(bucket_name, gcs_files):

    async with aiohttp.ClientSession() as session:
        storage = Storage(session=session)
        coros = (gcs_download(session, bucket_name, file, storage) for file in gcs_files)
        return await asyncio.gather(*coros)
        

Then the way I'm calling / passing in values are as follows:

import test_download as test
import asyncio

bucket_name = 'my_bucket_name'
project_name = 'my_project_name'  ### Where do I reference this???

gcs_files = ['bucket_folder/some-file-2020-10-06.txt', 
            'bucket_folder/some-file-2020-10-07.txt',
            'bucket_folder/some-file-2020-10-08.txt']

result = asyncio.run(test.get_gcsfiles_async(bucket_name, gcs_files))

Any help would be appreciated!

Here is related question, although there are two things to note: Google Storage python api download in parallel

1. When I run the code from the approved answer it ends up getting stuck and never downloads
2. It's from before the gcloud-aio-storage package was released and might not be leveraging the "best" current methods.

Answer 1

It looks like the documentation for that library is lacking, but I could get something running, and it is working on my tests. Something I found out by looking at the code is that you don’t need to use blob.download(), since it calls storage.download() anyways. I based the script below on the usage section, which deals with uploads, but can be rewritten for downloading. storage.download() does not write to a file, since that is done by storage.download_to_filename(). You can check the available download methods here.

async_download.py

import asyncio
from gcloud.aio.auth import Token
from gcloud.aio.storage import Storage

# Used a token from a service account for authentication
sa_token = Token(service_file="../resources/gcs-test-service-account.json", scopes=["https://www.googleapis.com/auth/devstorage.full_control"])

async def async_download(bucket, obj_names):
    async with Storage(token=sa_token) as client:
        tasks = (client.download(bucket, file) for file in obj_names) # Used the built in download method, with required args
        res = await asyncio.gather(*tasks)

    await sa_token.close()
    return res

main.py

import async_download as dl_test
import asyncio

bucket_name = "my-bucket-name"
obj_names = [
    "text1.txt",
    "text2.txt",
    "text3.txt"
]

res = asyncio.run(dl_test.async_download(bucket_name, obj_names))

print(res)

If you want to use a Service Account Token instead, you can follow this guide and use the relevant auth scopes. Since Service Accounts are project-wise, specifying a project is not needed, but I did not see any project name references for a Session either. While the GCP Python library for GCS does not yet support parallel downloads, there is a feature request open for this. There is no ETA for a release of this yet.