Writing tiles to numpy in parallel

[edwinytgoh]

I am trying to improve the performance of my script that takes an xarray dataset containing a variable with shape (~2e6, 256, 256) that corresponds to the coordinates index, i_local, and j_local, respectively. My goal is to cast the entire variable to float16 and write each index to individual 256x256 PNG files.

This dataset is a subset of a 5PB ocean simulation (LLC4320) that is lazily loaded using xarray, so I'm essentially streaming from one location on disk into numpy files in another location on disk. However, my current implementation is very slow, taking ~8 seconds per tile. I have looked into using map_blocks and GroupBy.apply, but those seem to require helper functions that return arrays, whereas my save_to_numpy function doesn't return anything.

Is there any way to parallelize this operation using xarray and dask? Are there any best practices for this type of operation? Any help would be greatly appreciated.

---

Just in case, here's my save_to_numpy function:

import io
import tarfile
from tqdm.auto import tqdm

def save_to_numpy(ds, out_folder):
    # convert to single precision float16 to save space
    ds = ds.astype("float16").chunk({"index": 2048})
    
    # write numpy tiles to tar file using io.BytesIO()    
    tar_path = os.path.join(out_folder, "output.tar")
    with tarfile.open(tar_path, "a") as tar:
        for idx in tqdm(ds.index.values):  # each iteration takes 8 seconds!
                tile = get_numpy_tile(ds["Theta"].isel(index=idx))
                filename = f"{idx}.npy"
                with io.BytesIO() as buf:
                    np.save(buf, tile)
                    buf.seek(0)
                    info = tarfile.TarInfo(filename)
                    info.size = len(buf.getvalue())
                    tar.addfile(tarinfo=info, fileobj=buf)        

def get_numpy_tile(tile):
    tile = tile.drop_vars(["index", "time", "tile_num"])
    tile = tile.squeeze()
    tile = tile.transpose("j_local", "i_local")
    tile = tile.load()
    tile = tile.values
    tile = tile.astype(np.float16)
    return tile

[TomNicholas]

You could always just use dask.delayed like we did here. Basically create a function which when called pulls out one chunk of your array and writes just that chunk to a file. Then wrap that function in dask.delayed and I think dask should be able to parallelize that.

[edwinytgoh]

Thanks for the reply! I wrote a new function called save_to_numpy_parallel and decorated my save_tile function with @dask.delayed. However, I encountered the following error upon execution:

  File "src/generate_llc_tiles.py", line 111, in save_to_numpy_parallel
    results = dask.compute(*delayed_ops, num_workers=20)
  File "/python3.8/site-packages/dask/base.py", line 573, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/python3.8/site-packages/dask/multiprocessing.py", line 220, in get
    result = get_async(
  File "/python3.8/site-packages/dask/local.py", line 494, in get_async
    fire_tasks(chunksize)
  File "/python3.8/site-packages/dask/local.py", line 476, in fire_tasks
    dumps((dsk[key], data)),
  File "/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 73, in dumps
    cp.dump(obj)
  File "/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 632, in dump
    return Pickler.dump(self, obj)
TypeError: cannot pickle '_io.BufferedWriter' object

Do you have any suggestions on how to address this error?

My modified save_to_numpy_parallel function is as follows:

import dask
dask.config.set(scheduler='processes')
dask.config.set({"distributed.worker.use-file-locking": False})    

def save_to_numpy_parallel(ds, out_folder):
    # convert to single precision float16 to save space
    ds = ds.astype("float16").chunk({"index": 2048})
    # write numpy tiles to tar file using io.BytesIO()    
    tar_path = os.path.join(out_folder, "output.tar")
    with tarfile.open(tar_path, "a") as tar:
        delayed_ops = []  # for use with dask.compute()
        for idx in ds.index.values:
            delayed_ops.append(save_tile(ds, idx, tar))
        # Compute the delayed operations
        results = dask.compute(*delayed_ops, num_workers=20)

@dask.delayed
def save_tile(ds, idx, tar):
    tile = get_numpy_tile(ds["Theta"].isel(index=idx))
    with io.BytesIO() as buf:
        np.save(buf, tile)
        buf.seek(0)
        info = tarfile.TarInfo(f"{idx}.npy")
        info.size = len(buf.getvalue())
        tar.addfile(tarinfo=info, fileobj=buf)
···

[TomNicholas]

I'm afraid I'm not sure 😕 There are some results for that error online though, I would start there.

[TomNicholas]

Worth also asking on the pangeo discourse too, with an explanation of why you're trying to achieve what you're doing.

[dcherian]

Seems like you can't use cloudpickle with that io thing,. Dask needs this to communicate between workers.

I would write out a simple dask.array version of the problem and ask on the dask tracker.

[edwinytgoh]

Worth also asking on the pangeo discourse too, with an explanation of why you're trying to achieve what you're doing.

Thanks for the suggestion, I'll check out the pangeo discourse as well. I actually just responded to @rabernat below about what we're trying to use the dataset for, and why we're thinking about saving each tile as a separate .npy file.

I also attempted to remove the tar writing portion of my save_to_numpy_parallel function so that each dask delayed op directly writes to a single .npy file, but it resulted in an out-of-memory issue even with 20 workers. I'll try to reduce the number of workers and report back.

[rabernat]

Why not just use Zarr? That's basically what Zarr does: split arrays into a bunch of chunks.

[edwinytgoh]

The .chunk({"index": 2048}) call makes use of all cores when running get_batch(), which otherwise only uses 1 core/thread. However, it doesn't seem to meaningfully affect timings.

In terms of the chunk size on disk, here's a sample output from lfs getstripe 2012.zarr:

2012.zarr/Theta
stripe_count:  1 stripe_size:   1048576 pattern:       raid0 stripe_offset: -1

./0.0.306
lmm_stripe_count:  1
lmm_stripe_size:   1048576
lmm_pattern:       raid0
lmm_layout_gen:    0
lmm_stripe_offset: 1
        obdidx           objid           objid           group
             1         5859423       0x59685f                0

./0.0.204
lmm_stripe_count:  1
lmm_stripe_size:   1048576
lmm_pattern:       raid0
lmm_layout_gen:    0
lmm_stripe_offset: 5
        obdidx           objid           objid           group
             5         7938280       0x7920e8                0

./0.0.178
lmm_stripe_count:  1
lmm_stripe_size:   1048576
lmm_pattern:       raid0
lmm_layout_gen:    0
lmm_stripe_offset: 6
        obdidx           objid           objid           group
             6         7977223       0x79b907                0

Thanks for the suggestion to use xbatch! I will give that a shot after I figure out the individual numpy tiles. In the meantime, I'll also try running a benchmark with get_batch = lambda sst: sst[..., np.arange(0, 256)].values to verify whether the issue really lies with the random sampling.

[edwinytgoh]

import numpy as np
import xarray as xr
ds = xr.open_dataset("/scratch/LLC4320/SST/2012.zarr")
sst = ds["Theta"]
# profile the time taken to load the "batch" of first 256 tiles (~32MB per batch)
get_batch = lambda sst: sst[..., np.arange(0, 256)].values
get_batch(sst)  # warm up run
%timeit -n 4 -r 10 get_batch(sst)  #  single-worker analogue to 4 epochs, 10 batches per epoch

Gives: 252 ms ± 5.09 ms per loop (mean ± std. dev. of 10 runs, 4 loops each)

Manual chunking with .chunk()

ds = xr.open_dataset("/scratch/LLC4320/SST/2012.zarr").chunk({"index": 2048}): 495 ms ± 8.06 ms per loop

Different indexing variations:

• 2048 "batch", np.arange(0, 2048): 882 ms ± 8.22 ms per loop
• 2048 "batch" but reading across chunks, np.arange(1024, 3072): 1.13 s ± 12.5 ms per loop
• 4096 "batch", np.arange(0, 4096): 1.8 s ± 13.4 ms per loop

[rabernat]

@edwinytgoh, you're still not understanding my question about the chunks on disk. I apologize for not being clear.

Please run the following code and share the output

import zarr
group = zarr.open("/scratch/LLC4320/SST/2012.zarr/")
print(group["Theta"].info)

[edwinytgoh]

Ah, I apologize for misunderstanding your question - I was under the impression that you were talking about actual hardware-level chunking across the Lustre filesystem. Here's the output from group["Theta"].info:

Name               : /Theta
Type               : zarr.core.Array
Data type          : float16
Shape              : (256, 256, 897920)
Chunk shape        : (256, 256, 2048)
Order              : C
Read-only          : False
Compressor         : Blosc(cname='lz4', clevel=5, shuffle=SHUFFLE, blocksize=0)
Store type         : zarr.storage.DirectoryStore
No. bytes          : 117692170240 (109.6G)
No. bytes stored   : 73543091370 (68.5G)
Storage ratio      : 1.6
Chunks initialized : 439/439

[rabernat]

OK, excellent, so this confirms that you chunks on disk are indeed 2048

Chunk shape        : (256, 256, 2048)

If you really want to shuffle your data over the full index axis, you might want to consider writing a copy with

Chunk shape        : (256, 256, 1)

This should be equivalent to what you're trying to do with numpy tiles: one chunk per image.

Also, experiment with not using Dask at all. Xarray by default uses lazy array access and won't load the data until you need it. Also remember that whenever you call .chunk, everything downstream of that has to load the entire chunk into memory.

[negin513]

You might want to consider using xbatcher which is developed for ML applications. It basically makes it easy to feed Xarray objects to machine learning libraries such as PyTorch or Keras.
Here is a good example of using Zarr + Xbatcher.