import os
from IPython.kernel import client
import time


if __name__ == "__main__":
    
    os.system('nohup ipcontroller &> con.out &')
    time.sleep(2)
    
    os.system('nohup ipengine &> eng1.out &')
    os.system('nohup ipengine &> eng2.out &')
    os.system('nohup ipengine &> eng3.out &')
    time.sleep(2)
        
    rc = client.MultiEngineClient()
    tc = client.TaskClient()
    
    
    while len(rc) < 3:    
        time.sleep(2)
        
    print "The cluster is setup"
    
    rc.execute('import os,sys')
    pwd = os.getcwd()     
    rc.execute('os.chdir("' + pwd + '")')    
    rc.execute('sys.path.append("' + pwd + '")')    
    rc.execute('import test_cluster2 as M')
        
    cmd = """        
        resp = V
    """
    
    print 'Distributing tasks...'    
    tids=[]
    for i in range(1000):        
        tids.append(tc.run(client.StringTask(cmd, pull=['resp'], push=dict(V=i))))
    
    print len(tids), 'tasks started. Waiting to finish...'
    tc.barrier(tids)
    
    resp = []        
    for tid in range(len(tids)):
        res = tc.get_task_result(tids[tid])
        if None<>res.failure:
            res.failure.printDetailedTraceback()
            res.failure.raiseException()            
        resp.append(res.results['resp'])
    
    print "Sum of responses: " , sum(resp)
    
    rc.kill(controller = True)