What is the best way to use dask to transform a large number of files?

I have a large number of relatively small files in s3. I need to read each file, do some processing, and write them back to Google Cloud Storage. Each file is small enough to fit in memory. It is important to preserve the name and contents of each file.

What would be the best dask abstraction to use to scale this work across multiple cores/threads?

I tried to use dask.bag to process the files and was successful with small batches, but ran into memory issues when trying to process a larger number of files.

I used include_path when reading the files but could not figure out a way for to_textfiles to use the path information without first creating a list of paths, which requires reading all files in memory before the first transformed file is written.

This is what I attempted:

import json
import dask.bag as db

from dask.distributed import performance_report
from dask.distributed import Client

def transform_data(data):
    extracted_at = data.pop("extracted_at")
    return dict(extracted_at=extracted_at, data=data)


def transform_file(file):
    """ We need to separate json lines, manipulate each record, and return a json string
    """
    contents, path = file
    data = [
        json.dumps(transform_data(json.loads(line)))
        for line in contents.replace("}{", "}\n{").splitlines() # fix missing \n between json strings
    ]
    return "\n".join(data), path


def get_paths(path):
    return path[-1].replace(".gz", ".jsonl.gz")


if __name__ == "__main__":
    client = Client(processes=False, n_workers=4)
    b = db.read_text("s3://mybucket/prefix/**/*.gz", include_path=True)

    with performance_report():
        xformed = b.map(transform_file)
        data = xformed.map(lambda x: x[0])  # each item is a tuple of (data, path)
        paths = xformed.map(get_paths).compute()
        data.to_textfiles(paths)

Is it possible to write the files as they are processed to avoid the memory issue? Is bag not the best abstraction for this use-case?

Thanks!