log

on rank: 3
on rank: 4
[0.6762652  0.36602538 0.66778483 ... 0.81427016 0.02274506 0.5112974 ]
on rank: 3
[0.87459513 0.16242899 0.82267847 ... 0.45909029 0.2555669  0.14410488]
on rank: 4
[0.53281283 0.69011684 0.96558067 ... 0.2611109  0.48600367 0.92250252]
on rank: 3
[0.955936   0.9542669  0.33336958 ... 0.4231765  0.68975261 0.02536984]
on rank: 4
on rank: 2
[0.46576689 0.40748544 0.80455681 ... 0.42885748 0.81368281 0.38118024]
on rank: 3
[0.88111086 0.99480552 0.16140886 ... 0.90521507 0.77372323 0.13990896]
on rank: 4
[0.82990596 0.48924338 0.41863186 ... 0.59769472 0.51270926 0.81132706]
on rank: 3
[0.45036641 0.42551552 0.52030069 ... 0.77405095 0.84382401 0.89169458]
on rank: 4
[0.31825974 0.62104552 0.75929952 ... 0.37520872 0.16978756 0.8760555 ]
on rank: 3
[0.96237011 0.44356861 0.78833879 ... 0.3674196  0.86802084 0.71364039]
[0.97582544 0.80295255 0.25654445 ... 0.84160663 0.12412553 0.59906167]
on rank: 4
on rank: 2
[0.1688846  0.58157204 0.16451889 ... 0.34421949 0.49695815 0.92643581]
on rank: 3
[0.98214734 0.98315998 0.63971929 ... 0.14208908 0.75490514 0.53435884]
on rank: 4
[0.57835946 0.14757299 0.84228858 ... 0.64323714 0.4602186  0.79279534]
on rank: 3
[0.68370391 0.2832432  0.6213475  ... 0.61192222 0.22898187 0.2974134 ]
[0.66759712 0.65493042 0.81066896 ... 0.48954846 0.72284914 0.76437382]
on rank: 3
on rank: 4
[0.32983885 0.2539114  0.79152594 ... 0.02218682 0.86351997 0.47719963]
on rank: 2
[0.30644285 0.73747115 0.85986612 ... 0.73666464 0.10342705 0.52379661]
on rank: 3
[0.76916236 0.77598458 0.28500003 ... 0.90905628 0.23665702 0.57075437]
on rank: 4
[0.52920866 0.26291031 0.89066068 ... 0.75433109 0.9493736  0.95311504]
on rank: 3
[0.98245076 0.68116538 0.20551943 ... 0.19077821 0.97035904 0.55776512]
on rank: 2
[0.08761166 0.51342471 0.78168433 ... 0.22695929 0.04876412 0.13969128]
on rank: 3
[0.80238936 0.84823687 0.91772985 ... 0.21073148 0.15191539 0.2641806 ]
on rank: 4
[0.51222459 0.30124403 0.91731018 ... 0.65364405 0.3013653  0.72670106]
on rank: 2
[0.50083549 0.64954766 0.05770564 ... 0.27849506 0.07379375 0.46671673]
on rank: 3
[0.24594152 0.13169863 0.52008051 ... 0.49674937 0.40600912 0.50324279]
on rank: 4
[0.54202305 0.85822515 0.41030172 ... 0.64581797 0.70166402 0.30382231]
on rank: 3
[0.70767039 0.90951326 0.31133026 ... 0.49314863 0.06560531 0.16518159]
on rank: 2
[0.90398164 0.60180442 0.9443617  ... 0.14031918 0.02232488 0.48822571]
[0.05478206 0.81664613 0.35910074 ... 0.96420982 0.32538952 0.59723687]
on rank: 3
on rank: 4
[0.41548813 0.88569158 0.29297808 ... 0.54687345 0.16020542 0.09430924]
on rank: 2
[0.84546033 0.12713573 0.86581912 ... 0.76424964 0.98041288 0.3324159 ]
on rank: 3
[0.71561284 0.38432743 0.18767541 ... 0.12406488 0.11866775 0.87406871]
on rank: 2
[0.93821379 0.12543426 0.34341433 ... 0.57030132 0.9727206  0.70089124]
on rank: 4
[0.12517972 0.53384516 0.71692212 ... 0.21513075 0.00232064 0.76206317]
on rank: 3
[0.23141619 0.49854579 0.88881073 ... 0.81373299 0.4846015  0.3575222 ]
[0.92257104 0.48420128 0.44848163 ... 0.99039567 0.79186595 0.03491569]
on rank: 3
on rank: 2
[0.13133474 0.55528742 0.63292403 ... 0.84422123 0.50021923 0.22490392]
on rank: 4
[0.87358901 0.1544822  0.749749   ... 0.31119323 0.5052893  0.19191903]
on rank: 3
[0.56583095 0.02431988 0.71280272 ... 0.39936182 0.27877391 0.22340064]
on rank: 2
[0.62742026 0.91286224 0.05855439 ... 0.51768182 0.31498707 0.85661797]
on rank: 4
[0.78144947 0.61066993 0.1755132  ... 0.29480791 0.87427025 0.61911644]
on rank: 3
[0.92764739 0.43795626 0.45164963 ... 0.13309332 0.15574926 0.18489999]
on rank: 2
[0.92782642 0.36555016 0.48214365 ... 0.63561091 0.87621868 0.84461551]
on rank: 3
[0.81782931 0.91192807 0.70472449 ... 0.72206016 0.94951062 0.0270532 ]
on rank: 4
[0.7637982  0.12859063 0.00491193 ... 0.07275175 0.02714237 0.03680531]
on rank: 3
[0.82922057 0.72431252 0.45917991 ... 0.3639784  0.46698012 0.0553549 ]
[0.79833215 0.37720596 0.47819629 ... 0.97509617 0.83669869 0.68710309]
on rank: 3
on rank: 4
[0.36659509 0.61956985 0.8437419  ... 0.93418933 0.91230952 0.22763575]
on rank: 2
[0.54514914 0.79057549 0.20953065 ... 0.06259719 0.01732728 0.22779676]
on rank: 3
[0.50680852 0.74784516 0.3917731  ... 0.56205127 0.14644047 0.55942508]
on rank: 4
[0.16010723 0.89391557 0.00727873 ... 0.94842165 0.3289769  0.04128895]
on rank: 3
[0.80243112 0.28737219 0.78483333 ... 0.99765049 0.5003993  0.30659048]
on rank: 2
[0.98190252 0.05404361 0.29106037 ... 0.63044645 0.83272067 0.35665344]
on rank: 3
[0.43604235 0.3627482  0.74784993 ... 0.83750898 0.22771437 0.65857736]
on rank: 4
[0.86145385 0.97696374 0.76890572 ... 0.58332048 0.7052309  0.78694877]
on rank: 3
[0.19728786 0.90364256 0.14541513 ... 0.44514958 0.71732876 0.89461189]
on rank: 2
[0.98778021 0.29503345 0.18655413 ... 0.69174964 0.32790176 0.08562617]
on rank: 4
[0.72826404 0.88720794 0.42520134 ... 0.65761562 0.86800046 0.01064764]
on rank: 3
[0.91831824 0.11894736 0.55936059 ... 0.01655764 0.50530541 0.60582922]
on rank: 2
[0.55202718 0.66102594 0.71290269 ... 0.94206327 0.06368831 0.77837592]
[0.78273617 0.71191591 0.27537167 ... 0.40160733 0.75219183 0.52767671]
on rank: 4
[0.18369987 0.16518824 0.98569219 ... 0.83331517 0.14594768 0.52862712]
on rank: 2
[0.03203684 0.90570367 0.32990894 ... 0.96342936 0.62205253 0.15890437]
on rank: 4
[0.07906099 0.43661611 0.00932421 ... 0.36108141 0.19721442 0.45146561]
on rank: 2
[0.10959625 0.5699617  0.58444372 ... 0.98312544 0.34038455 0.12052567]
on rank: 4
[0.11351032 0.70019545 0.85659508 ... 0.33768587 0.4381742  0.11141565]
on rank: 2
[0.11467384 0.95662757 0.27713965 ... 0.52380709 0.25276928 0.45763576]
on rank: 4
[0.98849967 0.18133257 0.25786528 ... 0.85140901 0.60453729 0.70112326]
on rank: 2
[0.12650615 0.51441426 0.54008579 ... 0.33433543 0.70892265 0.05704064]
on rank: 4
[0.06179748 0.72890846 0.90119964 ... 0.86525527 0.46014486 0.44775929]
on rank: 2
[0.37237507 0.89626974 0.48030807 ... 0.17981154 0.62025363 0.93082311]
on rank: 4
[0.4199309  0.60691144 0.83223096 ... 0.68900948 0.13088822 0.83614849]
on rank: 2
[0.0333568  0.94207194 0.84041287 ... 0.4965042  0.80985039 0.27780345]
on rank: 4
[0.82649037 0.39312487 0.38275763 ... 0.66184616 0.55448179 0.22289818]
on rank: 2
[0.59972367 0.23660624 0.61304144 ... 0.8432855  0.57396537 0.66696322]
on rank: 4
[0.73638404 0.32416858 0.10096265 ... 0.91413991 0.68288748 0.56960323]
on rank: 2
[0.26676647 0.58787766 0.64258518 ... 0.59977315 0.8691493  0.81694284]
on rank: 4
[0.36080418 0.77036432 0.26019506 ... 0.98416404 0.64319957 0.09826414]
on rank: 2
[0.73296521 0.97174026 0.78322377 ... 0.59036903 0.22357104 0.09051722]
on rank: 4
[0.64763093 0.04049762 0.23235353 ... 0.34726947 0.86169054 0.62781441]
on rank: 2
[0.92169614 0.9898451  0.40096562 ... 0.8222193  0.64446272 0.18012898]
[0.35295633 0.8059343  0.41423439 ... 0.33008426 0.81416168 0.82325237]
on rank: 2
[0.21761518 0.75705988 0.69187121 ... 0.28424747 0.32016067 0.427896  ]
on rank: 2
[0.27539021 0.67511978 0.98492016 ... 0.66133562 0.75453656 0.23543746]
on rank: 2
[0.60006236 0.10945516 0.37048505 ... 0.73805396 0.776438   0.97534687]
on rank: 2
[0.90321743 0.70065045 0.04812071 ... 0.24920127 0.63280322 0.13230854]
on rank: 2
[0.3033475  0.23233823 0.49161606 ... 0.44838765 0.65622138 0.82372118]
on rank: 2
[0.12800169 0.02475751 0.03915679 ... 0.38747201 0.43100064 0.12585467]
on rank: 2
[0.61887436 0.16392423 0.20214295 ... 0.55868692 0.43958721 0.97016008]
on rank: 2
[0.25641897 0.34256872 0.51111756 ... 0.69751212 0.10591505 0.29269547]
on rank: 2
[0.84608528 0.54703337 0.59115594 ... 0.47893436 0.17921986 0.3440337 ]
on rank: 2
[0.65593448 0.2484794  0.496993   ... 0.62272806 0.21674368 0.73041811]
on rank: 2
[0.9218968  0.44637063 0.49744715 ... 0.30897724 0.86883827 0.44634603]
on rank: 2
[0.8837767  0.36708401 0.49818932 ... 0.05460973 0.70123494 0.98856517]
on rank: 2
[0.26752397 0.75503283 0.38552768 ... 0.75014506 0.55600362 0.49910974]
on rank: 2
[0.0130105  0.21319671 0.52268074 ... 0.59308191 0.17284722 0.10492754]
on rank: 2
[0.07367413 0.97387486 0.40212462 ... 0.00357962 0.41816324 0.05680377]
on rank: 2
[0.54517198 0.56434908 0.95470926 ... 0.10839728 0.48502013 0.36642402]
on rank: 2
[0.07569338 0.53976229 0.5742005  ... 0.59393042 0.301477   0.62426507]

my_dask_script.py

# Setup Dask on MPI
# Run with
#    mpirun python my_dask_scripy.py
#    mpirun -np 5 --oversubscribe python my_dask_script.py  # on a laptop

from dask_mpi import initialize
initialize(nthreads=1, bokeh=False)

from dask.distributed import Client, wait, futures_of
client = Client()


# Use Dask to "load" data (actually just create random data here)
import dask.array as da
x = da.random.random(100000000, chunks=(1000000,))
x = x.persist()
wait(x)

# Find out where data is on each worker
# TODO: This could be improved on the Dask side to reduce boiler plate
from toolz import first
from collections import defaultdict
key_to_part_dict = {str(part.key): part for part in futures_of(x)}
who_has = client.who_has(x)
worker_map = defaultdict(list)
for key, workers in who_has.items():
    worker_map[first(workers)].append(key_to_part_dict[key])


# Call an MPI-enabled function on the list of data present on each worker
from my_mpi_lib import print_data_and_rank

futures = [client.submit(print_data_and_rank, list_of_parts, workers=worker)
           for worker, list_of_parts in worker_map.items()]

wait(futures)

client.close()

my_mpi_lib.py

from mpi4py import MPI

comm = MPI.COMM_WORLD

def print_data_and_rank(chunks: list):
    """ Fake function that mocks out how an MPI function should operate

    -   It takes in a list of chunks of data that are present on this machine
    -   It does whatever it wants to with this data and MPI
        Here for simplicity we just print the data and print the rank
    -   Maybe it returns something
    """
    rank = comm.Get_rank()

    for chunk in chunks:
        print("on rank:", rank)
        print(chunk)

    return sum(chunk.sum() for chunk in chunks)

Currently depends on dask/dask-mpi#16