How do get values as they're available on a map_blocks call

I’m trying to process a large chunk of data using map_blocks but I don’t know how to access the data as it is read while processing in parallel.

out = dask.array.map_blocks( torchit, dtype="float32", chunks = chunks )
print("processing")
for err in out:
    start = time.time()
    val = err.compute()
    print( (time.time() - start), "batch complete" )

If I do it this way, then each chunks gets computed one at a time. I think they can be performed in parallel though. I am not dead set on using map_blocks maybe I am using the wrong method to begin with.

For a solution to this problem, I used pythons futures.concurrent.

ex = concurrent.futures.ThreadPoolExecutor()
L = [ ex.submit( torchit, block_id ) for block id in range( n ) ]
for future in L:
    val = future.result()
ex.shutdown()

I thought it would be nice to stick with dask, and I think I could have managed using dask.distributed with a Client and Future. I didn’t see the advantage over just using concurrent.futures.

I am not exactly sure what you are trying to accomplish with “torchit”, but I recommend looking into .compute() and .persist(). Dask will automatically parallelize based on available compute resources. If you run a debuger and stop it on a line in your function, it will actually bring you to a thread with a specific chunk of data (once for each chunk of data) once you have called compute.

out.compute() will compute and collect the result into one (numpy/cupy) array.

out.persist() will also compute but keep it in the chunked form across threads.

Hello Vincent!

torchit is a function that takes the input block, runs it through a pytorch nn and returns the same sized block. Using out.compute would require too much memory because the whole out array would be in memory at the same time. I could use a proxy method around torchit so that out is a small array.

def proxyit(id):
    val = torchit(id)
    #do something with val
    return 0

It sounds like out.persist is what I need. Where I could use.

futures = out.persist()
for future in futures:
    val = future.result()
    #process val

Would that leave me with a full out array in memory though?

persist() will keep “out” loaded in a distributed manner (over many threads/cores/cpus depending on your configuration). People use persist generally to keep the original dataset in memory when performing many calculations on this one object. Otherwise dask will load the data, do computation 1, unload the data, load the data, do computation 2 etc. Stacking computations will make this even worse. Persists helps you to tell dask that this is important to keep.

It depends on what you want to do after obtaining out. If it is your final result, and you cannot use “compute” then you will probably want to store it in a zarr of hdf5 file. For that, you do not even have to call persist. If it is an image, you can add it as a dask array to napari. There you also do not have to call persist.

It all depends on what you want to do with out :wink:

Originally I was saving out which I was using ngff_zarr to handle. I think this question was more derived from that situation: What if I don’t want to save out, but I do want to process it? Eg. write some features to a file and move on. Without ever loading the full amount of input data or out into memory at once because it is too big.

For example, I have a series of images (10000, 1, 64, 64, 64), and I want to calculate a bunch of morphology features and write values in the same order.

You can simply chain various map_blocks() to define an elaborate pipeline. Dask does not actually do anything until you tell it to.

a = da.arange(5, chunks=2) # a dask array
b = da.arange(5, chunks=2) # a dask array
c = da.map_blocks(lambda x,y: x + y**2, a,b) # also a dask array
d = c+1 # STILL a dask array

ZERO computation has occurred. ZERO data has been loaded.
THEN we tell what we actually want to get out of our pipeline.

### if you want:
# (1) to get the result as a standard numpy array
# (2) something ready to use, but nicely spread across our machine(s)
# (3) to explore/visualize the image (in Napari)
# (4) the result as a saved file

# (1)
out = c.compute()

# all intermediate steps are used, and deleted once completed. d is never computed.
# dask distributes the required computations over the available threads/CPUs/nodes
# The output is transfered to one object on one CPU.
# you will get a memory error because your image size is too large

# (2)
c.persist() # The same calculation as compute
# the chunks stay as numpy arrays on many threads/CPUs/nodes
# "c" is still behaves as a dask array in your script
# Ideal if we want to do many follow up computations on "c". Like:
dask.compute(c.mean(), c.max(), c.min(), d.mean(), d.max())

# (3)
import napari
napari.view_image(d) # napari retrieves the information it needs to render from different threads
# it may be beneficial to use .persist() or even .compute().

# (4)
d.to_zarr() # Once a chunk is computed, it gets stored to a file.
# memory is freed up for another computation
# d still behaves as a dask array.
# getting an item from d requires you to call compute again
# alternatively, you can load it from the saved file, but you have to re-specify!