New Dask + ITK blogpost

[GenevieveBuckley]

Now that [release v5.3rc03 of ITK is available (which should include this PR), it would be good to do a follow up blogpost to this one about using Dask + ITK together.

The purpose of this would be:

1. To tell other people that ITK images can now be serialized & used with Dask, in the hopes that they experiment with this on their own, and/or
2. in the event that we find there are still problems that need fixing, to document the current state of work.

The first step is re-running the code from the earlier blogpost with ITK v5.3rc03 or above and seeing whether that works or not. Then we write up whatever we find.

Here are the comments specifically discussing what should be included in a followup blogpost:

• ENH: pickle serialization for itk images (issue 1091) InsightSoftwareConsortium/ITK#2829 (comment)
• ENH: pickle serialization for itk images (issue 1091) InsightSoftwareConsortium/ITK#2829 (comment)
• ENH: pickle serialization for itk images (issue 1091) InsightSoftwareConsortium/ITK#2829 (comment)

Links:

• Blogpost: https://blog.dask.org/2019/08/09/image-itk
• PR discussing blogpost: Blogpost using Dask Array with ITK #27
• PR enabling pickle serialization for ITK images: ENH: pickle serialization for itk images (issue 1091) InsightSoftwareConsortium/ITK#2829

[GenevieveBuckley]

Link to the data: https://drive.google.com/drive/folders/13mpIfqspKTIINkfoWbFsVtFF8D7jbTqJ (linked to from this earlier blogpost about image loading)

Link to the PSF: https://drive.google.com/drive/folders/13udO-h9epItG5MNWBp0VxBkKCllYBLQF (discussed here)

Details from the blogpost:

We will use image data generously provided by Gokul Upadhyayula at the Advanced Bioimaging Center at UC Berkeley and discussed in this paper (preprint), though the workloads presented here should work for any kind of imaging data, or array data generally.

[GenevieveBuckley]

Nope, it doesn't work. Previously, we had hoped that the simpler application of itk's richardson lucy deconvolution would now be possible with map_blocks. Unfortunately, it looks like there is still some kind of problem blocking it from happening.

deconvolved = da.map_blocks(itk.richardson_lucy_deconvolution_image_filter, imgs, kernel=psf, iterations=1, dtype=np.float32)

I wish I had a good error message or something to add here. However, whatever the problem is only happens during the compute call, so I'm not getting anything useful out to help with debugging.

I'm pretty sure the problem is not related to the leading dimensions, because I can use np.expand_dims on regular numpy array input to itk.richardson_lucy_deconvolution and that runs ok.

I don't think I have more time to play with this, sadly.

cc @thewtex

[thewtex]

@GenevieveBuckley yes, absolutely, great plan!

Will be you at SciPy US this year to hash this out? 🚶‍♀️ 🤝

@PranjalSahu built on your work, made a few fixes:

https://www.youtube.com/watch?v=CMmoa8pP_eo

These are now published in itk-5.3rc4.post1 (published yesterday). TBD if the issue you mentioned still occurs.

[GenevieveBuckley]

@GenevieveBuckley yes, absolutely, great plan!

Well, I don't know if there is a plan anymore, since the current situation for me is "it doesn't work and I can't summarize why". I guess having this issue thread discussion is better than nothing, though.

Will be you at SciPy US this year to hash this out? 🚶‍♀️ 🤝

Nope! (I was recently in the US, but didn't realize at the time that would have meant a better opportunity to videochat while in the same timezone. Oh well, we could still videochat, it's just a lot more inconvenient)

@PranjalSahu built on your work, made a few fixes:

https://www.youtube.com/watch?v=CMmoa8pP_eo

Very cool! I will have to check the talk out 😄

These are now published in itk-5.3rc4.post1 (published yesterday). TBD if the issue you mentioned still occurs.

We'll have to check, this could be promising!

[GenevieveBuckley]

These are now published in itk-5.3rc4.post1 (published yesterday). TBD if the issue you mentioned still occurs.

We'll have to check, this could be promising!

Update: Unfortunately it looks like there are still problems. I see this error when I try to compute the deconvolution with itk

import itk
import numpy as np
from dask.array.image import imread

imgs = imread('raw-488nm/*.tif').astype(np.float32)
psf = imread('psfs_z0p1/PSF_488nm_dz100nm.tif').astype(np.float32)

deconvolved = da.map_blocks(itk.richardson_lucy_deconvolution_image_filter,
                            imgs,
                            kernel_image=psf,
                            number_of_iterations=1,
                            dtype=np.float32)

output = deconvolved[100, ...].compute()

Traceback:

2022-07-12 17:55:14,883 - distributed.nanny - WARNING - Restarting worker
2022-07-12 17:55:16,364 - distributed.nanny - WARNING - Restarting worker
2022-07-12 17:55:17,415 - distributed.nanny - WARNING - Restarting worker
2022-07-12 17:55:20,493 - distributed.nanny - WARNING - Restarting worker

---------------------------------------------------------------------------
KilledWorker                              Traceback (most recent call last)
Input In [16], in <cell line: 1>()
----> 1 output = deconvolved[100, ...].compute()
      2 print(output.shape)

File ~\.conda\envs\itk-dask\lib\site-packages\dask\base.py:314, in DaskMethodsMixin.compute(self, **kwargs)
    290 def compute(self, **kwargs):
    291     """Compute this dask collection
    292 
    293     This turns a lazy Dask collection into its in-memory equivalent.
   (...)
    312     dask.base.compute
    313     """
--> 314     (result,) = compute(self, traverse=False, **kwargs)
    315     return result

File ~\.conda\envs\itk-dask\lib\site-packages\dask\base.py:602, in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
    599     keys.append(x.__dask_keys__())
    600     postcomputes.append(x.__dask_postcompute__())
--> 602 results = schedule(dsk, keys, **kwargs)
    603 return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])

File ~\.conda\envs\itk-dask\lib\site-packages\distributed\client.py:3000, in Client.get(self, dsk, keys, workers, allow_other_workers, resources, sync, asynchronous, direct, retries, priority, fifo_timeout, actors, **kwargs)
   2998         should_rejoin = False
   2999 try:
-> 3000     results = self.gather(packed, asynchronous=asynchronous, direct=direct)
   3001 finally:
   3002     for f in futures.values():

File ~\.conda\envs\itk-dask\lib\site-packages\distributed\client.py:2174, in Client.gather(self, futures, errors, direct, asynchronous)
   2172 else:
   2173     local_worker = None
-> 2174 return self.sync(
   2175     self._gather,
   2176     futures,
   2177     errors=errors,
   2178     direct=direct,
   2179     local_worker=local_worker,
   2180     asynchronous=asynchronous,
   2181 )

File ~\.conda\envs\itk-dask\lib\site-packages\distributed\utils.py:336, in SyncMethodMixin.sync(self, func, asynchronous, callback_timeout, *args, **kwargs)
    334     return future
    335 else:
--> 336     return sync(
    337         self.loop, func, *args, callback_timeout=callback_timeout, **kwargs
    338     )

File ~\.conda\envs\itk-dask\lib\site-packages\distributed\utils.py:403, in sync(loop, func, callback_timeout, *args, **kwargs)
    401 if error:
    402     typ, exc, tb = error
--> 403     raise exc.with_traceback(tb)
    404 else:
    405     return result

File ~\.conda\envs\itk-dask\lib\site-packages\distributed\utils.py:376, in sync.<locals>.f()
    374         future = asyncio.wait_for(future, callback_timeout)
    375     future = asyncio.ensure_future(future)
--> 376     result = yield future
    377 except Exception:
    378     error = sys.exc_info()

File ~\.conda\envs\itk-dask\lib\site-packages\tornado\gen.py:769, in Runner.run(self)
    766 exc_info = None
    768 try:
--> 769     value = future.result()
    770 except Exception:
    771     exc_info = sys.exc_info()

File ~\.conda\envs\itk-dask\lib\site-packages\distributed\client.py:2037, in Client._gather(self, futures, errors, direct, local_worker)
   2035         exc = CancelledError(key)
   2036     else:
-> 2037         raise exception.with_traceback(traceback)
   2038     raise exc
   2039 if errors == "skip":

KilledWorker: ("('richardson_lucy_deconvolution_image_filter-getitem-21573412bfdc1b613ca3e9d5e79d94bb', 0, 0, 0)", <WorkerState 'tcp://127.0.0.1:58059', name: 1, status: closed, memory: 0, processing: 1>)

[PranjalSahu]

I will also test it on my system.

[PranjalSahu]

I tested it on Linux and the computation could finish for this code:
output = deconvolved[5, ...].compute()

I tested it locally. Did you perform computation on some remote server?

Ok I can reproduce the error when using LocalCluster with Dask.
It is the same issue of import.

ImportError: PyCapsule_Import could not import module "_ITKCommonPython" 2022-07-12 14:18:39,986 - distributed.core - ERROR - PyCapsule_Import could not import module "_ITKCommonPython"

[PranjalSahu]

@GenevieveBuckley Ok I got it to work. I still don't know the root cause but I did something similar before.
Please try doing this.

def mymethod(*args, **kwargs):
    import itk
    output = itk.richardson_lucy_deconvolution_image_filter(*args, **kwargs)
    return output

deconvolved = da.map_blocks(mymethod,
                            imgs,
                            kernel_image=psf,
                            number_of_iterations=1,
                            dtype=np.float32)
output = deconvolved[10, ...].compute()
print(output.shape)

[GenevieveBuckley]

Oh, that's weird - if you import itk within the function it works for you, but not if itk is imported outside of it? That's very odd. I'll give that a go (tomorrow, probably)

[PranjalSahu]

I have tested it on entire dataset using 2 workers each with 32 GB limit.
The workers are restarting when I perform compute on all 50 images (probably memory leak or in-sufficient memory).
I am able to run it with 25 images like this

output = deconvolved[10:35, ...].compute()

Also I see some "ConnectionResetError: [Errno 104] Connection reset by peer" errors. But no computation fails.

This is the dask dashboard for this computation.

[GenevieveBuckley]

@GenevieveBuckley Ok I got it to work. I still don't know the root cause but I did something similar before. Please try doing this.

def mymethod(*args, **kwargs):
    import itk
    output = itk.richardson_lucy_deconvolution_image_filter(*args, **kwargs)
    return output

deconvolved = da.map_blocks(mymethod,
                            imgs,
                            kernel_image=psf,
                            number_of_iterations=1,
                            dtype=np.float32)
output = deconvolved[10, ...].compute()
print(output.shape)

Confirming this, I also see it working if itk is imported inside the function passed to dask's map_blocks, but not for the same thing with the import located outside the function.

[GenevieveBuckley]

This PR by Matt McCormick could fix the problem: InsightSoftwareConsortium/ITK#3494

[GenevieveBuckley]

Status: ITK PR #3494 has fixed the problem! We will wait for the next ITK release candidate (higher than the current v5.3rc03) to become available, and then write the blogpost.

It'll be much better to release the blogpost when readers can install the ITK pre-release and try things themselves (plus it'll be easier to write then).

Ideally, the blogpost should ideally have two examples:

1. The image deconvolution example from the original dask + itk blogpost, and
2. This ITK mesh example from Pranjal Sahu would also be good to include:

import dask
import itk
from dask import delayed

if __name__== "__main__":
    from dask.distributed import Client, LocalCluster
    cluster = LocalCluster()
    client = Client(cluster)

    @delayed
    def performsum(m):
        return m.GetNumberOfPoints()

    m = itk.Mesh[itk.F, 3].New()
    a = performsum(m).compute()
    print(a)