#!/usr/bin/env python
from __future__ import print_function
from dask.distributed import Client
import subprocess
import os
import socket

import sys
import argparse

def execall(client, args):
    cmd = ' '.join(args.cmd)
    full_cmd = 'cd {}; {}'.format(os.getcwd(), cmd)
    def myfunc(cmd):
        result = subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True)
        return result
    result = client.run(myfunc, full_cmd)
    for key in result.keys():
        print('Worker {}:'.format(key))
        print(result[key].decode('utf-8'))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Execute a command on all workers')
    parser.add_argument('cmd', nargs='*', help='The command to execute')
    args = parser.parse_args()
    ip = socket.gethostbyname(socket.gethostname())
    dask_scheduler = '{}:8786'.format(ip)
    try:
        client = Client(dask_scheduler, timeout=2)
    except IOError:
        print('Warning: using local dask client')
        client = Client()
    execall(client, args)
    client.close()
