Writing a Zarr from multiple processes or threads

[alexgleith]

Hey folks

I'm just getting started writing Zarr files, and I have a question about how to handle creating a Zarr from many processes at once.

I have a script over here that takes a bunch of SST NetCDF files and combines them into a single Zarr. This fails on some files, because a thread is writing one of the index files when another wants to, I think. I'll copy the important parts of the code in below.

The basic structure is informed by how I'd like to run this operationally, i.e., if we have a new daily file, we should be able to "process_one_file" and add it to the Zarr. But there might be more than one file, and the initial work is to combine 10 years or so, so 3-4000 files.

I know we can use other methods, like open_mfdataset to do this, but I want to exhaust my options on this idempotent method first! Is there a way we can utilise a shared lock using Xarray's to_zarr maybe? Or am I being naive here?!

def process_one_file(arguments):
    path, store_path, create = arguments

    print(f"Working on {path}")

    try:
        data = read_clean_data(path)
        data = data.chunk({"time": 1, "lon": 500, "lat": 500})
        if create:
            data.to_zarr(store_path, mode="w", consolidated=True)
        else:
            data.to_zarr(store_path, mode="a", consolidated=True, append_dim="time")
    except Exception as e:
        print(f"Failed {path} with exception {e}")
        return False

    print(f"Finished {path}")
    return True


def main(source_path, store_path):
    paths = get_paths(source_path)

    # TODO: check if the zarr exists first, and only then create it
    path = paths.pop(0)
    print("Creating template file...")
    process_one_file((path, store_path, True))

    paths = paths[0:13]

    # Process each file in the list on a separate thread
    print(f"Processing {len(paths)} files")
    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
        results = executor.map(process_one_file, [(p, store_path, False) for p in paths])

    failed = not all(results)
    if failed:
        print("Failed to write all files...")