Using dask.distributed for distributing tomography reconstruction computations at ESRF

1. Tomography processing
2. Why we use dask.distributed
3. How we use dask.distributed
4. Benefits and limitations

A tomography scan acquires a collection of images : radios

In parallel-beam geometry, line k of each radios is used to build sinogram k

Radios are chunked, so that each "compute worker" will handle a series of sinograms

Worker k handles chunk k

Project management reasons:

• Focus developments on our core expertise: tomography processing
• Do not re-invent a new computations distribution system: additional liability and maintenance burden
• Use an off-the-shelf solution with a community focused on scientific computing
• In the code, decouple the computations distribution from actual processing (more difficult to achieve with "invasive" solutions like MPI)

Technical reasons:

• Parallel-beam geometry --> "obvious" way of distributing work
• No need of shared memory (in our case)
• "Move computing resources to data, not the other way around"
• Many nice features of distributed: resilience, smart scheduling, dashboard, Actors/RPC, ...
• dask_jobqueue offers an abstraction on top of SLURM/OAR/SGE... ESRF uses two batch schedulers !

We do not use dask itself, but dask.distributed and dask_jobqueue.

The new tomography processing software can do the computations in two modes:

• "Local": computations are distributed on the local machine
• "jobqueue": computations are distributed over SLURM or OAR.

In both cases, the process is the same. A "master process" will

1. Ingest the processing tasks and options (via a config file or dictionary)
2. Spawn a "cluster" of workers
3. Query the workers for their actual computing resources
4. Build a list of tasks (basically: "reconstruct this chunk")
5. Distribute the tasks

Benefits

• Spawning workers and submiting tasks is roughly the same with "local" and "jobqueue" modes
• Client/scheduler model (steps 2, 3, 4 above allow flexibility)
• Asynchronous calls (submit() returns a future) and callbacks

Limitations

• Handling of cuda contexts (see next slide)
• RPC: "Actors" are not as full-featured as the state-less submit() (workaround next slide)
• Heterogeneous computations
  • Alleviated in "local mode" by using SpecCluster instead of LocalCluster. But not dask_jobqueue equivalent (?)
  • Workaround for "jobqueue mode": spawn two (client, cluster) couples, and one "meta-client" ?

In the context of (off-line) tomography reconstruction, the critical resource is GPU memory (more than computing power).

The standard way of using distributed is to submit state-less jobs. However for some reason cuda context are not properly handled between Nanny and Worker:

• Have to manually push() and pop() the Cuda contexts
• cufft context does not seem to be cleaned up after job completion

Solutions:

• Destroy completely the cuda context and manually release GPU memory after each chunk processing (not elegant and error-prone)
• Use stateful computations: Actors