#!/usr/bin/env python

import os
import time
import sys
import subprocess

def emit_output(p):
    (sout, serr) = p.communicate(None)
    if len(sout): 
        print >> sys.stdout, sout,
    sys.stdout.flush()

def wait(n):
    global n_completed
    while (len(procs) >= n):
        t = time.time()
        for p in procs:
            if p.poll() is not None:
                print >> sys.stderr, '%.2f -- %s is done, status %d' % (t-times[p], cmds[p].rstrip(), p.returncode)
                emit_output(p)
                n_completed=n_completed+1
                busy[prochost[p]] = 0
                del prochost[p]
                del cmds[p]
                del times[p]
                procs.remove(p)
            elif t-times[p] > timeout:
                print >> sys.stderr, '%s timed out!' % cmds[p].rstrip()
                p.terminate()
                emit_output(p)
                busy[prochost[p]] = 0
                del prochost[p]
                del cmds[p]
                del times[p]
                procs.remove(p)
        time.sleep(1.0)
    time.sleep(0.05)

def free_host():
    for host in hosts:
        if busy[host] != 1:
            return host
    return None

task_file = sys.argv[1]
hosts_file = sys.argv[2]
nprocs = int(sys.argv[3])  # Oversubscription parameter: number of processors to devote to a task

hosts = []
tasks = []
nname = -1  # prefix of host specification
procs = []  # process, defined as task on a host
cmds = {}   # command assocated with process
prochost = {} # host associate with process
times = {}
busy = {}
timeout = 2*3600
n_active = 0 # number of active jobs we expect to have running should equal Nhosts * Ncore_per_host / nprocs


n_completed=0
start_time = time.time()

with open(task_file) as f:
    for task in f:
        if task[0] != '#':
            tasks.append(task.rstrip())

with open(hosts_file) as f:
    for host in f:
        if host[:2] == 'mu': 
            cores = int(12 / nprocs)  # Use nprocs to reduce the core subscription
            nname = 6
        elif host[:2] == 'ml': 
            cores = 16
            nname = 5
        elif host[:2] == 'mp': 
            cores = 8
            nname = 5
        else:
            cores = 1
            nname = -1
        for core in range(cores):
            h = "%s-%d" % (host[:nname], core)
            hosts.append(h)
            busy[h] = 0
            n_active = n_active + 1

print 'running %d commands at a time from %d tasks in %s' \
    % (n_active, len(tasks), task_file)

for task in tasks:
    h = free_host()
    busy[h] = 1
    command = '/usr/bin/ssh -n -q %s \"cd /scratch3/watersc1/;  %s\"' % (h[:nname], task)
    print 'executing ', command
    sys.stdout.flush()
    try:
        p = subprocess.Popen(command, bufsize=-1, shell=True, stdin=None, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    except:
        print >> sys.stderr, "Exception! %s" % command
    procs.append(p)
    prochost[p] = h
    cmds[p] = command
    times[p] = time.time()
    wait(n_active)

wait(1)

t = time.time()
print >> sys.stderr, 'all done. %d tasks completed in %2.f seconds (%.2f per second)' % (n_completed, t-start_time, n_completed/(t-start_time))
