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.

Sign up for GitHub

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

Jump to bottom

Add async write logic #885

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

Add async write logic #885

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 28 additions & 0 deletions s3fs/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2086,8 +2086,15 @@ async def _invalidate_region_cache(self):
async def open_async(self, path, mode="rb", **kwargs):
if "b" not in mode or kwargs.get("compression"):
raise ValueError
if "w" in mode or "a" in mode:
return await self._open_for_writing(path, mode, **kwargs)
return S3AsyncStreamedFile(self, path, mode)

async def _open_for_writing(self, path, mode, **kwargs):
# Parse the path to get bucket and key
bucket, key, _ = self.split_path(path)
return S3AsyncStreamWriter(self, bucket, key)


class S3File(AbstractBufferedFile):
"""
Expand Down Expand Up @@ -2429,6 +2436,26 @@ def _abort_mpu(self):
self.mpu = None


# Define a new class to represent the file-like object for writing
class S3AsyncStreamWriter:
def __init__(self, s3_fs, bucket, key):
self.s3_fs = s3_fs
self.bucket = bucket
self.key = key
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

# 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.

"put_object",
Bucket=self.bucket,
Key=self.key,
Body=data,
)
self.loc += len(data)


class S3AsyncStreamedFile(AbstractAsyncStreamedFile):
def __init__(self, fs, path, mode):
self.fs = fs
Expand Down Expand Up @@ -2481,3 +2508,4 @@ async def _call_and_read():
resp["Body"].close()

return await _error_wrapper(_call_and_read, retries=fs.retries)

12 changes: 12 additions & 0 deletions s3fs/tests/test_s3fs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2696,6 +2696,18 @@ async def read_stream():
break
out.append(got)

async def write_stream():
fs = S3FileSystem(
anon=False,
client_kwargs={"endpoint_url": endpoint_uri},
skip_instance_cache=True,
)
await fs._mkdir(test_bucket_name)
f = await fs.open_async(fn, mode="wb")
await f.write(data)
await f.close()

asyncio.run(write_stream())
asyncio.run(read_stream())
assert b"".join(out) == data

Expand Down