Profile Cubed task memory usage with Memray

[tomwhite]

I've been using Memray to profile the memory usage of Cubed tasks running locally using a local Lithops executor (since it runs tasks in a separate process).

This give great insight into how memory is allocated for blockwise operations. Here is an example timeline for an array add operation, where you can see Zarr chunks being loaded and stored:

More details in this presentation.

This issue is to add the instrumentation code to make it easy to profile any computation.

[tomwhite]

The code I used to produce the above plot is in this branch, in particular this commit.

It would be be good to have more control over i) the file that the profile info is written to, and ii) which tasks to run it for (it currently just runs it for the first task - i.e. with block ID composed of zeros).

Then to generate the visualization I ran:

rm -fr memray/*
pytest -vs 'cubed/tests/test_mem_utilization.py::test_add' --runslow
for f in $(ls memray); do echo $f; python -m memray flamegraph --temporal -f -o memray/$f.html memray/$f; done
# view html files in memray directory