There are still some inconsistencies:
I don’t see this print on your code snippet.
Also, I see two times:
Which is a bit weird.
Maybe we should try to debug step by step: what happens if you remove these two lines:
and
?
In order to see that, you should try without dask-mpi and just a LocalCluster
.