Source

xnet-hacks / xnet / tool.py


import sys
import json
from eventlet.green import urllib2, httplib
from xnet.xjob import XJobSpec, XResult, XJobInfo
from xnet.helpers import json_encode
from xnet.conn import ServerConnection
import argparse
import py

def listxjobs(conn):
    response = conn.GET("/listxjobs")
    content = response.read()
    xjoblist = json.loads(content)
    return xjoblist

def getresultinfo(conn, xjobid):
    response = conn.GET("/getresult/" + str(xjobid))
    content = response.read()
    data = json.loads(content)
    result = XResult.from_json(data['result'])
    xjobspec = XJobSpec.from_json(data['xjobspec'])
    return result, xjobspec

def submitxjob(conn, xjobspec):
    response = conn.POST("/", json_encode(xjobspec))
    d = json.loads(response.read())['xjobid']
    return d

def getstdout(conn, xjobid):
    token = 0
    while True:
        response = conn.GET("/getstdout/" + str(xjobid), headers={'range': token})
        content = response.read()
        try:
            stdout = json.loads(content)['stdout']
        except:
            print content
            break
        token = max(map(int, stdout.keys())) or token
        for line in stdout.values():
            print line

class Submit:
    def addoption(self, parser):
        parser.add_argument("positional", nargs="+",
            help="command with options to be executed")
    
    def run(self, args):
        conn = ServerConnection()
        xjobid = submitxjob(conn, XJobSpec(args.positional))
        print "submitted xjob with id", xjobid

class Result:
    def addoption(self, parser):
        parser.add_argument("xjobid", type=int,
            help="id for which to obtain the result")
    def run(self, args):
        conn = ServerConnection()
        result, xjobspec = getresultinfo(conn, args.xjobid)
        tw = py.io.TerminalWriter()
        #tw.sep("=")
        tw.line("xjob execution: %s" % " ".join(xjobspec.args))
        tw.sep("-", "stdout")
        tw.write(result.out)
        tw.sep("-", "stderr")
        tw.write(result.err)
        tw.sep("-", "return code %s" % result.retcode)

class List:
    def addoption(self, parser):
        pass
    def run(self, args):
        conn = ServerConnection()
        xjoblist = listxjobs(conn)
        print xjoblist

class Attach:
    def addoption(self, parser):
        parser.add_argument("xjobid", type=int,
                            help="id for which to attach to stdout")
    def run(self, args):
        conn = ServerConnection()
        getstdout(conn, args.xjobid)

def prepare_parser():
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers()
    
    addcommand(subparsers, Submit)
    addcommand(subparsers, List)
    addcommand(subparsers, Result)
    addcommand(subparsers, Attach)
    return parser

def addcommand(subparsers, cls):
    parser = subparsers.add_parser(cls.__name__.lower())
    inst = cls()
    inst.addoption(parser)
    parser.set_defaults(run=inst.run)

def main(args=None):
    if args is None:
        args = sys.argv[1:]
    args = prepare_parser().parse_args(args)
    args.run(args)