#1806 Batch upload flyte directory Pull request opened by pingsutw TL;DR s3fs will try to read all the files in the directory into memory by default. we can specify the batch size for

s3fsFileSystem.put

, and it will dramatically reduce memory consumption. Before: After (batch size = 100): Type ☑︎ Bug Fix ☐ Feature ☐ Plugin Are all requirements met? ☑︎ Code completed ☑︎ Smoke tested ☐ Unit tests added ☐ Code documentation added ☐ Any pending items have an associated Issue Complete description

from typing_extensions import Annotated

from flytekit import task, workflow
from flytekit.types.directory import FlyteDirectory
import os
from memory_profiler import profile
from flytekit import BatchSize
directory = "/tmp/test"


@task
@profile
def generate_files(num_files: int, filesize_mb: int) -> Annotated[FlyteDirectory, BatchSize(100)]:
    os.mkdir(directory)
    for i in range(num_files):
        file_path = os.path.join(directory, f"file_{i}.txt")
        with open(file_path, "w") as f:
            f.write(str(os.urandom(filesize_mb * 1024 * 1024)))
    return FlyteDirectory(directory)


@workflow
def ton_of_files_wf(num_files: int = 1000, filesize_mb: int = 8) -> FlyteDirectory:
    return generate_files(num_files=num_files, filesize_mb=filesize_mb)


if __name__ == '__main__':
    print(ton_of_files_wf())

Tracking Issue NA Follow-up issue NA flyteorg/flytekit GitHub Actions: build (ubuntu-latest, 3.9) GitHub Actions: build (ubuntu-latest, 3.8) GitHub Actions: docs GitHub Actions: lint GitHub Actions: Docs Warnings ✅ 2 other checks have passed 2/7 successful checks

#1806 Batch upload flyte directory Pull request merged by eapolinario flyteorg/flytekit