#!/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.
"""

from IPython.parallel import Client, Reference


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

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

view = rc[:]

# run bintree.py script defining bintree functions, etc.
execfile('bintree.py')

# 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)

# gather the connection information into a dict
ar = view.apply_async(lambda : com.info)
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, 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 = range(1000)
view.scatter('data', data)

# perform cumulative sum via allreduce
view.execute("data_sum = com.allreduce(add, data, flat=False)")
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']