import os

import subprocess

import socket

from IPython.parallel.error import UnmetDependency

# Determine if this engine (host) can accept the job.

if 'host' in job and socket.getfqdn() not in job['host']:

# UnmetDependency exception forces the scheduler to resubmit the task to a different engine.

raise UnmetDependency

# Make sure that 'cmd' field exists in the job.

if 'cmd' not in job:

return dict(error='"cmd" field does not exist in the job')

# Change the working directory if specified.

if 'cd' in job:

os.chdir(job['cd'])

# Set stdout and stderr file objects if specified.

stdout = open(job['out'], 'w') if 'out' in job else None

stderr = open(job['err'], 'w') if 'err' in job else None

# Run the command.

try:

# We force to use '/bin/bash -o pipefail' so that the return value presents the value of the rightmost command to exit with a non-zero status code.

returncode = subprocess.call(['/bin/bash', '-c', 'set -o pipefail; ' + job['cmd']], stdout=stdout, stderr=stderr)

return dict(code=returncode)

except OSError, e:

"""A "bakapara" client for IPython cluster.

Args: identical to IPython.parallel.Client

"""

def __init__(self, **args):

self.rc = Client(**args)

self.lview = None

self.ar = None

self.jobs = None

self.indices = None

self.finished = set()

# Obtain the host names and PIDs of engines.

self.pids = self.rc[:].apply(os.getpid).get_dict()

self.hosts = self.rc[:].apply(socket.getfqdn).get_dict()

# Change the working directory of each engine.

self.rc[:].apply(os.chdir, os.getcwd())

def run(self, jobs, targets=None):

"""Runs the jobs on the cluster.

Args:

jobs (list): list of dictionary objects describing the jobs.

targets (int, list of ints, 'all', or None): the engine(s) on which the jobs will run.

Returns:

bool: True if successful, False otherwise (e.g., jobs are running).

"""

if self.ar is not None and not self.ar.ready():

return False

self.lview = self.rc.load_balanced_view(targets)

self.ar = self.lview.map_async(runjob, jobs)

self.jobs = jobs

self.indices = dict([(k, v) for v, k in enumerate(self.ar.msg_ids)])

self.finished = set()

return True

def wait(self, timeout=1e-3):

"""Waits for the jobs to complete and writes job results.

Args:

timeout (float): a time in seconds, after which to give up.

"""

if self.ar is None:

return

# Find finished msg_ids.

pending = set(self.ar.msg_ids)

try:

self.rc.wait(pending, timeout)

except TimeoutError:

pass

finished = pending.difference(self.rc.outstanding)

# Overwrite the results in the job array.

for msg_id in finished:

i = self.indices[msg_id]

if i in self.finished:

continue

job = self.jobs[i]

meta = self.rc.metadata[msg_id]

result = self.rc.results[msg_id][0]

for key in ('submitted', 'started', 'completed', 'received'):

result[key] = meta[key].isoformat()

for key in ('engine_id', 'pyerr', 'pyout', 'status', 'msg_id'):

result[key] = meta[key]

result['elapsed'] = str(meta['completed'] - meta['started'])

result['host'] = self.hosts[meta['engine_id']]

job['result'] = result

self.finished.add(i)

def ready():

"""Returns whether the jobs have completed."""

return self.ar is not None or self.ar.ready()

def successful():

"""Returns whether the jobs completed without raising an exception.

Raises:

AssertionError: the result is not ready.

"""

return self.ar is not None and self.ar.successful()

def abort(self, **args):

"""Aborts jobs.

Args: identical to IPython.parallel.client.view.LoadBalancedView.abort()

"""

if self.lview is not None:

self.lview.abort()

def interrupt(self):

"""Sends SIGINT signal to engines (experimental).

http://mail.scipy.org/pipermail/ipython-dev/2014-March/013426.html

"""

self.abort()

for i in self.rc.ids:

host = self.hosts[i]

pid = self.pids[i]

if host == socket.getfqdn():

os.kill(pid, signal.SIGINT)

else:

os.system('ssh {} kill -INT {}'.format(host, pid))

def shutdown(self, **args):

"""Terminates one or more engine processes, optionally including the hub.

Args: identical to IPython.parallel.Client.shutdown

"""

if self.lview is None:

return False

self.lview.shutdown(**args)

def status(self, interval=1., timeout=-1, fo=sys.stdout):

"""Waits for the jobs, printing progress at regular intervals

Args:

interval (float): a time in seconds, after which to print the progress.

timeout (float): a time in seconds, after which to give up waiting.

fo (file): a file object to which the progress is printed.

"""

if self.ar is None:

return

if timeout is None:

timeout = -1

# Make sure to write the job results into the job objects.

self.wait(1e-3)

tic = time.time()

while not self.ar.ready() and (timeout < 0 or time.time() - tic <= timeout):

self.wait(interval)

clear_output(wait=True)

dt = datetime.timedelta(seconds=self.ar.elapsed)

fo.write('{}/{} tasks finished after {}'.format(self.ar.progress, len(self.ar), str(dt)))

fo.flush()

else:

fo.write('\n')

dt = datetime.timedelta(seconds=self.ar.elapsed)

clear_output(wait=True)

fo.write('{} tasks completed in {}\n'.format(len(self.ar), str(dt)))

fo.flush()

def __len__(self):

"""Returns the number of engines."""

import argparse

import json

import logging as lg

import operator

# Parse the command-line arguments.

parser = argparse.ArgumentParser(

description="Run jobs on the IPython cluster environment in 'bakapara' manner."

)

parser.add_argument(

'-p', '--profile', default='default',

help='specify the name of the Cluster profile'

)

parser.add_argument(

'-i', '--input',

help='specify the filename for reading the jobs (default: STDIN)'

)

parser.add_argument(

'-o', '--output',

help='specify the filename for storing the results (default: STDOUT)'

)

parser.add_argument(

'-l', '--logging',

help='specify the filename for storing the logs (default: STDERR)'

)

args = parser.parse_args(args)

# Determine the stream for reading the jobs.

fi = open(args.input) if args.input else sys.stdin

# Determine the stream for storing the results.

fo = open(args.output, 'w') if args.output else sys.stdout

# Determine the stream for storing the logs.

fe = open(args.logging, 'w') if args.logging else sys.stderr

# Initialize the default logger.

lg.basicConfig(level=lg.INFO, format='%(asctime)s %(message)s', stream=fe)

# Read the jobs.

jobs = map(lambda line: json.loads(line), fi)

# Initialize Bakapara client.

bp = Bakapara(profile=args.profile)

# Run the jobs.

lg.info('total {} jobs on {} engines'.format(len(jobs), len(bp)))

bp.run(jobs)

# Loop until all jobs are done.

finished = set()

while len(finished) < len(jobs):

bp.wait(1.)

just_finished = bp.finished - finished

if just_finished:

finished_jobs = [jobs[i] for i in just_finished]

finished_jobs.sort(key=lambda x: x['result']['completed'])

i = len(finished) + 1

for job in finished_jobs:

fo.write('{}\n'.format(json.dumps(job)))

fo.flush()

lg.info('[{i}/{n}] returned {code} in {elapsed} sec on {host}: {cmd}'.format(i=i, n=len(jobs), code=job['result']['code'], elapsed=job['result']['elapsed'], host=job['result']['host'], cmd=job['cmd']))

i += 1

finished |= bp.finished

lg.info('completed')

fo.close()