# L-13 MCS 507 Wed 20 Sep 2023 : mpi4py_dynamic.py

"""
Simulates dynamic load balancing.
At the command prompt, type
mpiexec -n 3 python3 mpi4py_dynamic.py
where '3' is the number of processes.
Every process runs the script below:
"""

from mpi4py import MPI
from random import randint, seed
from time import sleep

seed(231112117) # for deterministic runs

COMM = MPI.COMM_WORLD
SIZE = COMM.Get_size()
RANK = COMM.Get_rank()

def worker():
    """
    A worker receives a job, sleeps as many
    seconds as the value of the job and
    then sends the job back to the manager.
    """
    job = COMM.recv(source=0, tag=RANK)
    print(RANK, 'receives', job)
    sleep(job)
    print(RANK, 'sends', job)
    COMM.send(job, dest=0, tag=RANK)

def manager():
    """
    Distributes jobs to the workers
    and probes for their return.
    """
    for i in range(1, SIZE):
        job = randint(2, 11)
        print('manager sends job', job, 'to', i)
        COMM.send(job, dest=i, tag=i)
    print('manager waits for jobs to return...')
    count = 0
    while(count < SIZE-1):
        state = MPI.Status()
        okay = COMM.Iprobe(source=MPI.ANY_SOURCE, \
            tag=MPI.ANY_TAG, status=state)
        if(okay):
            node = state.Get_source()
            data = COMM.recv(source=node, tag=node)
            print('received', data, 'from', node)
            count = count + 1
        else:
            sleep(0.5)

if(RANK == 0):
    manager()
else:
    worker()