X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/740d71d2b7a797bb2dd2e1e427c12e84c86b9ce6..0812bc1c717e5fed57d420b177f6ca9d41e81032:/crunch_scripts/run-command

diff --git a/crunch_scripts/run-command b/crunch_scripts/run-command
index c624e3cadf..28adb749cb 100755
--- a/crunch_scripts/run-command
+++ b/crunch_scripts/run-command
@@ -1,7 +1,12 @@
 #!/usr/bin/env python
 
 import logging
-logging.basicConfig(level=logging.INFO, format="run-command: %(message)s")
+
+logger = logging.getLogger('run-command')
+log_handler = logging.StreamHandler()
+log_handler.setFormatter(logging.Formatter("run-command: %(message)s"))
+logger.addHandler(log_handler)
+logger.setLevel(logging.INFO)
 
 import arvados
 import re
@@ -20,25 +25,41 @@ import pprint
 import multiprocessing
 import crunchutil.robust_put as robust_put
 import crunchutil.vwd as vwd
+import argparse
+import json
+import tempfile
+import errno
 
-os.umask(0077)
-
-t = arvados.current_task().tmpdir
+parser = argparse.ArgumentParser()
+parser.add_argument('--dry-run', action='store_true')
+parser.add_argument('--script-parameters', type=str, default="{}")
+args = parser.parse_args()
 
-api = arvados.api('v1')
+os.umask(0077)
 
-os.chdir(arvados.current_task().tmpdir)
-os.mkdir("tmpdir")
-os.mkdir("output")
+if not args.dry_run:
+    api = arvados.api('v1')
+    t = arvados.current_task().tmpdir
+    os.chdir(arvados.current_task().tmpdir)
+    os.mkdir("tmpdir")
+    os.mkdir("output")
 
-os.chdir("output")
+    os.chdir("output")
 
-outdir = os.getcwd()
+    outdir = os.getcwd()
 
-taskp = None
-jobp = arvados.current_job()['script_parameters']
-if len(arvados.current_task()['parameters']) > 0:
-    taskp = arvados.current_task()['parameters']
+    taskp = None
+    jobp = arvados.current_job()['script_parameters']
+    if len(arvados.current_task()['parameters']) > 0:
+        taskp = arvados.current_task()['parameters']
+else:
+    outdir = "/tmp"
+    jobp = json.loads(args.script_parameters)
+    os.environ['JOB_UUID'] = 'zzzzz-8i9sb-1234567890abcde'
+    os.environ['TASK_UUID'] = 'zzzzz-ot0gb-1234567890abcde'
+    os.environ['CRUNCH_SRC'] = '/tmp/crunche-src'
+    if 'TASK_KEEPMOUNT' not in os.environ:
+        os.environ['TASK_KEEPMOUNT'] = '/keep'
 
 links = []
 
@@ -71,128 +92,297 @@ class SigHandler(object):
     def __init__(self):
         self.sig = None
 
-    def send_signal(self, sp, signum):
-        sp.send_signal(signum)
+    def send_signal(self, subprocesses, signum):
+        for sp in subprocesses:
+            sp.send_signal(signum)
         self.sig = signum
 
+# http://rightfootin.blogspot.com/2006/09/more-on-python-flatten.html
+def flatten(l, ltypes=(list, tuple)):
+    ltype = type(l)
+    l = list(l)
+    i = 0
+    while i < len(l):
+        while isinstance(l[i], ltypes):
+            if not l[i]:
+                l.pop(i)
+                i -= 1
+                break
+            else:
+                l[i:i + 1] = l[i]
+        i += 1
+    return ltype(l)
+
+def add_to_group(gr, match):
+    m = match.groups()
+    if m not in gr:
+        gr[m] = []
+    gr[m].append(match.group(0))
+
+# Return the name of variable ('var') that will take on each value in 'items'
+# when performing an inner substitution
+def var_items(p, c, key):
+    if "var" in c:
+        # Var specifies the variable name for inner parameter substitution
+        return (c["var"], get_items(p, c[key]))
+    else:
+        # The component function ('key') value is a list, so return the list
+        # directly with no parameter substition.
+        if isinstance(c[key], list):
+            return (None, get_items(p, c[key]))
+
+        # check if c[key] is a string that looks like a parameter
+        m = re.match("^\$\((.*)\)$", c[key])
+        if m and m.group(1) in p:
+            return (m.group(1), get_items(p, c[key]))
+        else:
+            # backwards compatible, foreach specifies bare parameter name to use
+            return (c[key], get_items(p, p[c[key]]))
+
+# "p" is the parameter scope, "c" is the item to be expanded.
+# If "c" is a dict, apply function expansion.
+# If "c" is a list, recursively expand each item and return a new list.
+# If "c" is a string, apply parameter substitution
 def expand_item(p, c):
     if isinstance(c, dict):
         if "foreach" in c and "command" in c:
-            var = c["foreach"]
-            items = get_items(p, p[var])
+            var, items = var_items(p, c, "foreach")
             r = []
             for i in items:
                 params = copy.copy(p)
                 params[var] = i
-                r.extend(expand_list(params, c["command"]))
+                r.append(expand_item(params, c["command"]))
+            return r
+        if "list" in c and "index" in c and "command" in c:
+            var, items = var_items(p, c, "list")
+            params = copy.copy(p)
+            params[var] = items[int(c["index"])]
+            return expand_item(params, c["command"])
+        if "regex" in c:
+            pattern = re.compile(c["regex"])
+            if "filter" in c:
+                var, items = var_items(p, c, "filter")
+                return [i for i in items if pattern.match(i)]
+            elif "group" in c:
+                var, items = var_items(p, c, "group")
+                groups = {}
+                for i in items:
+                    match = pattern.match(i)
+                    if match:
+                        add_to_group(groups, match)
+                return [groups[k] for k in groups]
+            elif "extract" in c:
+                var, items = var_items(p, c, "extract")
+                r = []
+                for i in items:
+                    match = pattern.match(i)
+                    if match:
+                        r.append(list(match.groups()))
+                return r
+        if "batch" in c and "size" in c:
+            var, items = var_items(p, c, "batch")
+            sz = int(c["size"])
+            r = []
+            for j in xrange(0, len(items), sz):
+                r.append(items[j:j+sz])
             return r
     elif isinstance(c, list):
-        return expand_list(p, c)
-    elif isinstance(c, str) or isinstance(c, unicode):
-        return [subst.do_substitution(p, c)]
-
-    return []
-
-def expand_list(p, l):
-    return [exp for arg in l for exp in expand_item(p, arg)]
-
+        return [expand_item(p, arg) for arg in c]
+    elif isinstance(c, basestring):
+        m = re.match("^\$\((.*)\)$", c)
+        if m and m.group(1) in p:
+            return expand_item(p, p[m.group(1)])
+        else:
+            return subst.do_substitution(p, c)
+
+    raise Exception("expand_item() unexpected parameter type %s" % (type(c))
+
+# Evaluate in a list context
+# "p" is the parameter scope, "value" will be evaluated
+# if "value" is a list after expansion, return that
+# if "value" is a path to a directory, return a list consisting of each entry in the directory
+# if "value" is a path to a file, return a list consisting of each line of the file
 def get_items(p, value):
+    value = expand_item(p, value)
     if isinstance(value, list):
-        return expand_list(p, value)
-
-    fn = subst.do_substitution(p, value)
-    mode = os.stat(fn).st_mode
-    prefix = fn[len(os.environ['TASK_KEEPMOUNT'])+1:]
-    if mode is not None:
-        if stat.S_ISDIR(mode):
-            items = ["$(dir %s/%s/)" % (prefix, l) for l in os.listdir(fn)]
-        elif stat.S_ISREG(mode):
-            with open(fn) as f:
-                items = [line for line in f]
-        return items
-    else:
-        return None
+        return value
+    elif isinstance(value, basestring):
+        mode = os.stat(value).st_mode
+        prefix = value[len(os.environ['TASK_KEEPMOUNT'])+1:]
+        if mode is not None:
+            if stat.S_ISDIR(mode):
+                items = [os.path.join(value, l) for l in os.listdir(value)]
+            elif stat.S_ISREG(mode):
+                with open(value) as f:
+                    items = [line.rstrip("\r\n") for line in f]
+            return items
+    raise Exception("get_items did not yield a list")
 
 stdoutname = None
 stdoutfile = None
 stdinname = None
 stdinfile = None
-rcode = 1
 
-try:
-    if "task.foreach" in jobp:
-        if arvados.current_task()['sequence'] == 0:
-            var = jobp["task.foreach"]
-            items = get_items(jobp, jobp[var])
-            logging.info("parallelizing on %s with items %s" % (var, items))
-            if items is not None:
-                for i in items:
-                    params = copy.copy(jobp)
-                    params[var] = i
+# Construct the cross product of all values of each variable listed in fvars
+def recursive_foreach(params, fvars):
+    var = fvars[0]
+    fvars = fvars[1:]
+    items = get_items(params, params[var])
+    logger.info("parallelizing on %s with items %s" % (var, items))
+    if items is not None:
+        for i in items:
+            params = copy.copy(params)
+            params[var] = i
+            if len(fvars) > 0:
+                recursive_foreach(params, fvars)
+            else:
+                if not args.dry_run:
                     arvados.api().job_tasks().create(body={
                         'job_uuid': arvados.current_job()['uuid'],
                         'created_by_job_task_uuid': arvados.current_task()['uuid'],
                         'sequence': 1,
                         'parameters': params
-                        }
-                    ).execute()
+                    }).execute()
+                else:
+                    if isinstance(params["command"][0], list):
+                        for c in params["command"]:
+                            logger.info(flatten(expand_item(params, c)))
+                    else:
+                        logger.info(flatten(expand_item(params, params["command"])))
+    else:
+        logger.error("parameter %s with value %s in task.foreach yielded no items" % (var, params[var]))
+        sys.exit(1)
+
+try:
+    if "task.foreach" in jobp:
+        if args.dry_run or arvados.current_task()['sequence'] == 0:
+            # This is the first task to start the other tasks and exit
+            fvars = jobp["task.foreach"]
+            if isinstance(fvars, basestring):
+                fvars = [fvars]
+            if not isinstance(fvars, list) or len(fvars) == 0:
+                logger.error("value of task.foreach must be a string or non-empty list")
+                sys.exit(1)
+            recursive_foreach(jobp, jobp["task.foreach"])
+            if not args.dry_run:
                 if "task.vwd" in jobp:
-                    # Base vwd collection will be merged with output fragments from
-                    # the other tasks by crunch.
+                    # Set output of the first task to the base vwd collection so it
+                    # will be merged with output fragments from the other tasks by
+                    # crunch.
                     arvados.current_task().set_output(subst.do_substitution(jobp, jobp["task.vwd"]))
                 else:
                     arvados.current_task().set_output(None)
-                sys.exit(0)
-            else:
-                sys.exit(1)
+            sys.exit(0)
     else:
+        # This is the only task so taskp/jobp are the same
         taskp = jobp
+except Exception as e:
+    logger.exception("caught exception")
+    logger.error("job parameters were:")
+    logger.error(pprint.pformat(jobp))
+    sys.exit(1)
 
-    if "task.vwd" in taskp:
-        # Populate output directory with symlinks to files in collection
-        vwd.checkout(subst.do_substitution(taskp, taskp["task.vwd"]), outdir)
-
-    if "task.cwd" in taskp:
-        os.chdir(subst.do_substitution(taskp, taskp["task.cwd"]))
-
-    cmd = expand_list(taskp, taskp["command"])
+try:
+    if not args.dry_run:
+        if "task.vwd" in taskp:
+            # Populate output directory with symlinks to files in collection
+            vwd.checkout(subst.do_substitution(taskp, taskp["task.vwd"]), outdir)
+
+        if "task.cwd" in taskp:
+            os.chdir(subst.do_substitution(taskp, taskp["task.cwd"]))
+
+    cmd = []
+    if isinstance(taskp["command"][0], list):
+        for c in taskp["command"]:
+            cmd.append(flatten(expand_item(taskp, c)))
+    else:
+        cmd.append(flatten(expand_item(taskp, taskp["command"])))
 
     if "task.stdin" in taskp:
         stdinname = subst.do_substitution(taskp, taskp["task.stdin"])
-        stdinfile = open(stdinname, "rb")
+        if not args.dry_run:
+            stdinfile = open(stdinname, "rb")
 
     if "task.stdout" in taskp:
         stdoutname = subst.do_substitution(taskp, taskp["task.stdout"])
-        stdoutfile = open(stdoutname, "wb")
+        if not args.dry_run:
+            stdoutfile = open(stdoutname, "wb")
 
-    logging.info("{}{}{}".format(' '.join(cmd), (" < " + stdinname) if stdinname is not None else "", (" > " + stdoutname) if stdoutname is not None else ""))
+    logger.info("{}{}{}".format(' | '.join([' '.join(c) for c in cmd]), (" < " + stdinname) if stdinname is not None else "", (" > " + stdoutname) if stdoutname is not None else ""))
 
+    if args.dry_run:
+        sys.exit(0)
+except subst.SubstitutionError as e:
+    logger.error(str(e))
+    logger.error("task parameters were:")
+    logger.error(pprint.pformat(taskp))
+    sys.exit(1)
 except Exception as e:
-    logging.exception("caught exception")
-    logging.error("task parameters was:")
-    logging.error(pprint.pformat(taskp))
+    logger.exception("caught exception")
+    logger.error("task parameters were:")
+    logger.error(pprint.pformat(taskp))
     sys.exit(1)
 
 try:
-    sp = subprocess.Popen(cmd, shell=False, stdin=stdinfile, stdout=stdoutfile)
-    sig = SigHandler()
+    subprocesses = []
+    close_streams = []
+    if stdinfile:
+        close_streams.append(stdinfile)
+    next_stdin = stdinfile
+
+    for i in xrange(len(cmd)):
+        if i == len(cmd)-1:
+            # this is the last command in the pipeline, so its stdout should go to stdoutfile
+            next_stdout = stdoutfile
+        else:
+            # this is an intermediate command in the pipeline, so its stdout should go to a pipe
+            next_stdout = subprocess.PIPE
+
+        sp = subprocess.Popen(cmd[i], shell=False, stdin=next_stdin, stdout=next_stdout)
+
+        # Need to close the FDs on our side so that subcommands will get SIGPIPE if the
+        # consuming process ends prematurely.
+        if sp.stdout:
+            close_streams.append(sp.stdout)
 
-    # forward signals to the process.
-    signal.signal(signal.SIGINT, lambda signum, frame: sig.send_signal(sp, signum))
-    signal.signal(signal.SIGTERM, lambda signum, frame: sig.send_signal(sp, signum))
-    signal.signal(signal.SIGQUIT, lambda signum, frame: sig.send_signal(sp, signum))
+        # Send this processes's stdout to to the next process's stdin
+        next_stdin = sp.stdout
 
-    # wait for process to complete.
-    rcode = sp.wait()
+        subprocesses.append(sp)
+
+    # File descriptors have been handed off to the subprocesses, so close them here.
+    for s in close_streams:
+        s.close()
+
+    # Set up signal handling
+    sig = SigHandler()
+
+    # Forward terminate signals to the subprocesses.
+    signal.signal(signal.SIGINT, lambda signum, frame: sig.send_signal(subprocesses, signum))
+    signal.signal(signal.SIGTERM, lambda signum, frame: sig.send_signal(subprocesses, signum))
+    signal.signal(signal.SIGQUIT, lambda signum, frame: sig.send_signal(subprocesses, signum))
+
+    active = 1
+    pids = set([s.pid for s in subprocesses])
+    rcode = {}
+    while len(pids) > 0:
+        (pid, status) = os.wait()
+        pids.discard(pid)
+        if not taskp.get("task.ignore_rcode"):
+            rcode[pid] = (status >> 8)
+        else:
+            rcode[pid] = 0
 
     if sig.sig is not None:
-        logging.critical("terminating on signal %s" % sig.sig)
+        logger.critical("terminating on signal %s" % sig.sig)
         sys.exit(2)
     else:
-        logging.info("completed with exit code %i (%s)" % (rcode, "success" if rcode == 0 else "failed"))
+        for i in xrange(len(cmd)):
+            r = rcode[subprocesses[i].pid]
+            logger.info("%s completed with exit code %i (%s)" % (cmd[i][0], r, "success" if r == 0 else "failed"))
 
 except Exception as e:
-    logging.exception("caught exception")
+    logger.exception("caught exception")
 
 # restore default signal handlers.
 signal.signal(signal.SIGINT, signal.SIG_DFL)
@@ -202,11 +392,11 @@ signal.signal(signal.SIGQUIT, signal.SIG_DFL)
 for l in links:
     os.unlink(l)
 
-logging.info("the following output files will be saved to keep:")
+logger.info("the following output files will be saved to keep:")
 
 subprocess.call(["find", ".", "-type", "f", "-printf", "run-command: %12.12s %h/%f\\n"], stdout=sys.stderr)
 
-logging.info("start writing output to keep")
+logger.info("start writing output to keep")
 
 if "task.vwd" in taskp:
     if "task.foreach" in jobp:
@@ -216,12 +406,15 @@ if "task.vwd" in taskp:
         # Just a single task, so do merge with the original collection
         outcollection = vwd.checkin(subst.do_substitution(taskp, taskp["task.vwd"]), outdir, merge=True).manifest_text()
 else:
-    outcollection = robust_put.upload(outdir)
+    outcollection = robust_put.upload(outdir, logger)
+
+# Success if no non-zero return codes
+success = not any([status != 0 for status in rcode.values()])
 
 api.job_tasks().update(uuid=arvados.current_task()['uuid'],
                                      body={
                                          'output': outcollection,
-                                         'success': (rcode == 0),
+                                         'success': success,
                                          'progress':1.0
                                      }).execute()