Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add async write logic #885

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

Sarafudinov
Copy link

No description provided.

@Sarafudinov Sarafudinov changed the title add async write logic Add async write logic Jun 17, 2024

async def write(self, data):
# Write data directly to S3 object
await self.s3_fs._call_s3(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This call will overwrite the file on each call to .write(). Ideally, we'd want to be able to push to an open HTTP stream, but that doesn't seem possible. However, s3 allows for multi-part uploads, with the condition that each partc is >=5MB. So we also would need buffering of bytes until there are enough.

self.closed = False
self.loc = 0

async def write(self, data):
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Multipart upload can be supported (I worked at Amazon in the past)

Let's keep the part size as 5.5MB just to be safe.

part_size = 5.5 * 1024 * 1024

start the multipart upload session

response = await self.s3_fs._call_s3(
            "create_multipart_upload",
            Bucket=self.bucket,
            Key=self.key,
            {'PartSize': part_size}
        )

The response key contains UploadId which is the unique identifier to associate with each subsequent operation

upload_id = response['UploadId']

We need to start using IO steam to batch up the data

data_stream = io.BytesIO(data)

now keep chunking data

chunk_num = 0
uploaded = []
while True:
    chunk_data = data_stream.read(part_size)
    if not chunk_data:
        break
    response = await self.s3_fs._call_s3(
                "upload_part",
                Bucket=self.bucket,
                Key=self.key,
                UploadId=upload_id,
                PartNumber=chunk_num,
                Body=chunk_data
            )
    chunk_num += 1
    uploaded.append({'PartNumber': chunk_num, 'ETag': response['ETag']})

Now we just call complete multipart upload

await self.s3_fs._call_s3(
            "complete_multipart_upload",
            Bucket=self.bucket,
            Key=self.key,
            UploadId=upload_id,
            MultipartUpload={'Parts': uploaded}
        )

make sure to have this in try-catch block so that in catch you can invoke abort_multipart_upload else your S3 account will keep getting charged

https://docs.aws.amazon.com/AmazonS3/latest/API/API_CreateMultipartUpload.html

After you initiate a multipart upload and upload one or more parts, to stop being charged for storing the uploaded parts, you must either complete or abort the multipart upload.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We already have exactly this model for files, where the user API is sync. The idea was, that we should be able to have a true async stream ready for writing, rather than pausing for bigger writes every time we have enough for a new Part. It seems that is not possible.

Implementing what you say - essentially making write() async but keeping the same logic as the standard file - may still be worth while for the sake of fsspec.generic.rsync . GCS and ab2 both support exactly the same pattern.

One possible problem: is it required that a previous Part is finished before the next one can begin?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(also note the method s3fs.S3FileSystem.clear_multipart_uploads for dealing with MPUs that might have failed to complete or abort).

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One possible problem: is it required that a previous Part is finished before the next one can begin?

Actually it's allowed https://aws.amazon.com/blogs/compute/uploading-large-objects-to-amazon-s3-using-multipart-upload-and-transfer-acceleration/

A multipart upload allows an application to upload a large object as a set of smaller parts uploaded in parallel

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants