#!/usr/bin/env python
"""
Script for setting up and using [all]reduce with a binary-tree engine interconnect.

usage: `python bintree_script.py`

This spanning tree strategy ensures that a single node node mailbox will never
receive more that 2 messages at once. This is very important to scale to large
clusters (e.g. 1000 nodes) since if you have many incoming messages of a couple
of megabytes you might saturate the network interface of a single node and
potentially its memory buffers if the messages are not consumed in a streamed
manner.

Note that the AllReduce scheme implemented with the spanning tree strategy
impose the aggregation function to be commutative and distributive. It might
not be the case if you implement the naive gather / reduce / broadcast strategy
where you can reorder the partial data before performing the reduce.
"""

import ipyparallel as ipp

# connect client and create views
rc = ipp.Client()
rc.block = True
ids = rc.ids

root_id = ids[0]
root = rc[root_id]

view = rc[:]

from bintree import bintree, print_bintree

# generate binary tree of parents
btree = bintree(ids)

print("setting up binary tree interconnect:")
print_bintree(btree)

view.run('bintree.py')
view.scatter('id', ids, flatten=True)
view['root_id'] = root_id

# create the Communicator objects on the engines
view.execute('com = BinaryTreeCommunicator(id, root = id==root_id )')
pub_url = root.apply_sync(lambda: com.pub_url)  # noqa: F821

# gather the connection information into a dict
ar = view.apply_async(lambda: com.info)  # noqa: F821
peers = ar.get_dict()
# this is a dict, keyed by engine ID, of the connection info for the EngineCommunicators


# connect the engines to each other:
def connect(com, peers, tree, pub_url, root_id):
    """this function will be called on the engines"""
    com.connect(peers, tree, pub_url, root_id)


view.apply_sync(connect, ipp.Reference('com'), peers, btree, pub_url, root_id)


# functions that can be used for reductions
# max and min builtins can be used as well
def add(a, b):
    """cumulative sum reduction"""
    return a + b


def mul(a, b):
    """cumulative product reduction"""
    return a * b


view['add'] = add
view['mul'] = mul

# scatter some data
data = list(range(1000))
view.scatter('data', data)

# perform cumulative sum via allreduce
print("reducing")
ar = view.execute("data_sum = com.allreduce(add, data, flat=False)", block=False)
ar.wait_interactive()
print("allreduce sum of data on all engines:", view['data_sum'])

# perform cumulative sum *without* final broadcast
# when not broadcasting with allreduce, the final result resides on the root node:
view.execute("ids_sum = com.reduce(add, id, flat=True)")
print("reduce sum of engine ids (not broadcast):", root['ids_sum'])
print("partial result on each engine:", view['ids_sum'])
