Processing Zarr files

[orliac]

Hi. I'm seeking advice on what would be the most natural way to process a Zarr file with cupynumeric where I process each chunk with a custom function and perform a reduction to get the final result. In dask-cuda I'm using map_blocks to that end.

[Jacobfaib]

Hello

So Legate provides a (very!) experimental Zarr IO backend that you might be able to use for this. Unfortunately, the actual read/write API is relatively limited:

from legate.core.experimental.io.zarr import read_array

logical_array = read_array("path/to/file.zarr") # returns legate.core.LogicalArray

You could further post-process the array using python tasks:

from legate.core.task import task, InputArray, OutputArray
import cupynumeric as cnp

@task
def post_process(arr_in: InputArray, arr_out: OutputArray) -> None:
  cnp_in = cnp.asarray(arr_in)
  cnp_out = cnp.asarray(arr_out)
  cnp_out[:] = ...

post_process(arr_in=logical_array, arr_out=logical_array)
arr = cnp.asarray(logical_array)

But we currently don't have mechanism to provide callbacks during the actual loading of the array like map_blocks I'm afraid.

[orliac]

Many thanks for your feedback @Jacobfaib and @manopapad.
My zarr file represents a data set with the following dimensions: [100, 130816, 512, 4], complex64 data.
Working on this dummy piece of code:

@task
def post_process(arr_in: InputArray, arr_out: ReductionArray[ReductionOpKind.ADD]) -> None:
    cnp_in = cnp.asarray(arr_in)
    cnp_out = cnp.asarray(arr_out)
    cnp_out += cnp.ones((100, 100))
    print("@task post_process adding ones[100,100] to arr_out", flush=True)
    print(f"Shape: {cnp_in.shape}, Size: {cnp_in.size}")
  
logical_out_array = cnp.zeros((100, 100))
post_process(arr_in=logical_vis_array, arr_out=logical_out_array)
arr = cnp.asarray(logical_out_array)
print(arr, arr.shape)

Fails with:

Unhandled Python exception: 'maximum recursion depth exceeded', aborting!
Signal 6 received by node 1, process 233473 (thread 7f41f47a0740) - obtaining backtrace
Signal 6 received by process 233473 (thread 7f41f47a0740) at:
stack trace: 12 frames
  [0] = /lib64/libc.so.6(+0x59db0) [0x7f6e71578db0]
  [1] = /lib64/libc.so.6(+0xa642c) [0x7f6e715c542c]
  [2] = /lib64/libc.so.6(raise+0x16) [0x7f6e71578d06]
  [3] = /lib64/libc.so.6(abort+0xd3) [0x7f6e7154b7d3]
  [4] = /work/scitas-ge/orliac/miniconda3/envs/CONV_LEGATE/lib/python3.12/site-packages/legate/core/_lib/task/task_info.cpython-312-x86_64-linux-gnu.so(+0x14aae) [0x7f6dcd602aae]
  [5] = /work/scitas-ge/orliac/miniconda3/envs/CONV_LEGATE/lib/python3.12/site-packages/legate/core/_lib/data/../../../../../../liblegate.so.25.01.00(legate::detail::legion_task_body(void (*)(legate::TaskContext), legate::VariantCode, std::optional<std::basic_string_view<char, std::char_traits<char> > >, void const*, unsigned long, Realm::Processor)+0x879) [0x7f6decb269c9]
  [6] = /work/scitas-ge/orliac/miniconda3/envs/CONV_LEGATE/lib/python3.12/site-packages/legate/core/_lib/task/task_info.cpython-312-x86_64-linux-gnu.so(void legate::detail::task_wrapper_dyn_name<&(_py_variant(legate::TaskContext)), (legate::VariantCode)1>(void const*, unsigned long, void const*, unsigned long, Realm::Processor)+0x34) [0x7f6dcd6144f4]
  [7] = /work/scitas-ge/orliac/miniconda3/envs/CONV_LEGATE/lib/python3.12/site-packages/legate/core/_lib/data/../../../../../.././librealm-legate.so.1(+0x7fdab1) [0x7f6dd0d29ab1]
  [8] = /work/scitas-ge/orliac/miniconda3/envs/CONV_LEGATE/lib/python3.12/site-packages/legate/core/_lib/data/../../../../../.././librealm-legate.so.1(+0x7fdb46) [0x7f6dd0d29b46]
  [9] = /work/scitas-ge/orliac/miniconda3/envs/CONV_LEGATE/lib/python3.12/site-packages/legate/core/_lib/data/../../../../../.././librealm-legate.so.1(+0x7fc08a) [0x7f6dd0d2808a]
  [10] = /work/scitas-ge/orliac/miniconda3/envs/CONV_LEGATE/lib/python3.12/site-packages/legate/core/_lib/data/../../../../../.././librealm-legate.so.1(+0x802c67) [0x7f6dd0d2ec67]
  [11] = /lib64/libc.so.6(+0x2e330) [0x7f6e7154d330]
slurmstepd: error:  mpi/pmix_v5: _errhandler: kh038 [3]: pmixp_client_v2.c:211: Error handler invoked: status = -61, source = [slurm.pmix.350807.0:3]
slurmstepd: error: *** STEP 350807.0 ON kh003 CANCELLED AT 2025-03-06T18:29:00 ***
srun: Job step aborted: Waiting up to 32 seconds for job step to finish.
slurmstepd: error:  mpi/pmix_v5: _errhandler: kh027 [2]: pmixp_client_v2.c:211: Error handler invoked: status = -61, source = [slurm.pmix.350807.0:2]
slurmstepd: error:  mpi/pmix_v5: _errhandler: kh019 [1]: pmixp_client_v2.c:211: Error handler invoked: status = -61, source = [slurm.pmix.350807.0:1]
srun: error: kh027: task 2: Aborted (core dumped)
srun: Terminating StepId=350807.0
srun: error: kh038: task 3: Aborted (core dumped)
srun: error: kh019: task 1: Aborted (core dumped)
slurmstepd: error:  mpi/pmix_v5: _errhandler: kh003 [0]: pmixp_client_v2.c:211: Error handler invoked: status = -61, source = [slurm.pmix.350807.0:0]
srun: error: kh003: task 0: Killed

while I'm just trying to access the shape and size of the data to be processed by each task.

It works fine though with smaller dummy arrays like:

@task
def my_task(arr: InputArray) -> None:
    arr_cnp = cnp.asarray(arr)  # Convert to cuNumeric array
    print(f"Shape: {arr_cnp.shape}, Size: {arr_cnp.size}", flush=True)

arr = cnp.ones((100, 200))  # Create a Legate array
my_task(arr)
arr = cnp.ones((10000, 20000))  # Create a Legate array
my_task(arr)

which gives:

Shape: (100, 200), Size: 20000
Shape: (10000, 5000), Size: 50000000
Shape: (10000, 5000), Size: 50000000
Shape: (10000, 5000), Size: 50000000
Shape: (10000, 5000), Size: 50000000

while working with 4 tasks (via Slurm)

I need to be able to process each zarr chunk assigned to each task individually. What is the way to retrieve the 4-dimensional indices of the chunks to be processed by each task?

[Jacobfaib]

I need to be able to process each zarr chunk assigned to each task individually. What is the way to retrieve the 4-dimensional indices of the chunks to be processed by each task?

So this is going to be tricky. We have an upcoming release where you will be able to request the TaskContext be passed into your python task as well:

from legate.core import TaskContext
from legate.core.task import task

@task
def task_with_ctx(ctx: TaskContext, other_args...) -> None:
    ...

# note, no task context in the task call, it is automagically inserted later
task_with_ctx(other_args...)

That TaskContext has methods get_task_index() and get_launch_domain() which get you the task-local N-dimensional index and the total task domain just like in the C++ API.

HOWEVER, those additional methods on the Python TaskContext did not quite make it into the release :(. However:

1. We will be open-sourcing Legate as part of this release (so you could build the latest)...
2. Or you could use nightly builds once the PR extending TaskContext is merged.

So all this is to say: you can't do this quite yet, but you will be able to Very Soon(TM). Sorry about that...

---

As for your other segfault error, it's hard to say without a symbolized backtrace... maybe @manopapad

[Jacobfaib]

Note to pull from the nightlies channel when downloading via conda pass -c legate/label/experimental to your conda install cupynumeric command.

I'll post on here when the TaskContext feature has landed there

[orliac]

Many thanks for the quick feedback @Jacobfaib! I'll wait for the next release then.

[Jacobfaib]

@orliac sorry for the delay. The TaskContext feature has now landed. You should be able to use the above experimental channel to pull the latest version. The properties TaskContext.task_index and TaskContext.launch_domain will contain the necessary information. The former returns an N-dimensional tuple with each tasks index, while the second returns a pair of tuples specifying the lower and upper bounds of the task domain (think of these like the bottom left and top right corners of a rectangle, if the launch domain is 2D).

Let me know if this helps