feat: Add retry mechanism to file access methods

[pingsutw]

Tracking issue

NA

Why are the changes needed?

Failed to upload a file to s3

      File "/usr/local/lib/python3.11/site-packages/s3fs/core.py", line 719, in _lsdir
        raise translate_boto_error(e)

Message:

    OSError: [Errno 16] Please reduce your request rate.

User error.

What changes were proposed in this pull request?

Catch the error Please reduce your request rate, and retry the request.

How was this patch tested?

Remote cluster

Setup process

import asyncio
import os
from typing import List, Dict, Any, Coroutine
import s3fs

from flytekit import ImageSpec, task, workflow, Resources, FlyteContextManager

# Constants
NUMBER_OF_REQUESTS: int = 100000
BUCKET_NAME: str = 's3://xxx/load-test/'
OBJECT_KEY: str = 'your_object_key'
DATA: str = 'abcsfsdfsdfsfdsdf'
s3fs = s3fs.S3FileSystem()


def write_data(source_file: str, object_key: str):
    ctx = FlyteContextManager.current_context()
    ctx.file_access.put(source_file, f"{BUCKET_NAME}{object_key}")


def make_file_if_not_exists(file_path: str, data: str) -> None:
    if not os.path.exists(file_path):
        with open(file_path, "w") as f:
            f.write(data)


async def make_s3_request(object_key: str, data: str, idx: int) -> Dict[str, Any]:
    if idx % 10 == 0:
        print(f"Making request for object_key: {idx}")
    source_file = f"/root/source_{idx % 10000}.txt"
    make_file_if_not_exists(source_file, data)
    response: Dict[str, Any] = await asyncio.to_thread(write_data, source_file, object_key)
    if idx % 10 == 0:
        print(f"Done with request {idx}")
    # response = await s3fs._put("/tmp/test/file.txt", f"s3://union-cloud-oc-staging-dogfood/load-test/{object_key}")
    return response


async def send_requests() -> List[Coroutine]:
    tasks: List[Coroutine] = [make_s3_request(f"{OBJECT_KEY}_{i}", DATA, i) for i in range(NUMBER_OF_REQUESTS)]
    print("Sending requests")
    return await asyncio.gather(*tasks)


async def main() -> None:
    responses: List[Dict[str, Any]] = await send_requests()
    # Do something with the responses


new_flytekit = "git+https://github.com/flyteorg/flytekit@0db9079b6e11767aac61c5dafc1c025ba04ef31e"
image = ImageSpec(registry="pingsutw", packages=[new_flytekit], apt_packages=["git"])


@task(
    requests=Resources(cpu="1", mem="3000Mi"),
    limits=Resources(cpu="1", mem="5000Mi"),
    interruptible=False,
    # container_image=image,
)
def s3_load_task():
    asyncio.run(main())


@workflow
def wf():
    for i in range(100):
        s3_load_task()


if __name__ == "__main__":
    asyncio.run(main())

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

NA

Docs link

NA

[codecov]

Codecov Report

Attention: Patch coverage is 64.70588% with 6 lines in your changes are missing coverage. Please review.

Project coverage is 83.60%. Comparing base (f45dc39) to head (8cf1bbd).
Report is 7 commits behind head on master.

Files	Patch %	Lines
flytekit/core/data_persistence.py	64.70%	3 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2287      +/-   ##
==========================================
+ Coverage   83.31%   83.60%   +0.28%     
==========================================
  Files         309      311       +2     
  Lines       24064    23966      -98     
  Branches     3496     3503       +7     
==========================================
- Hits        20050    20037      -13     
+ Misses       3386     3298      -88     
- Partials      628      631       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[pingsutw]

s3fs PR. fsspec/s3fs#865

[eapolinario]

exponential backoff + jitter is usually implemented like:

sleep(random.randint(0, min(2 ** retry, 32)))

https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/ is the gold standard describing a few techniques for implementing exp. backoff + jitter.