3609: Inherit --retries from _util. Be more specific about error being caught. ...
[arvados.git] / crunch_scripts / run-command
index 34419b4de42fb7d908f9a4726fb301a4d3439af9..28adb749cb63cc5df3f239ea380c3af04fbb88a5 100755 (executable)
@@ -28,6 +28,7 @@ import crunchutil.vwd as vwd
 import argparse
 import json
 import tempfile
+import errno
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--dry-run', action='store_true')
@@ -96,36 +97,74 @@ class SigHandler(object):
             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 = c["list"]
-            items = get_items(p, p[var])
+            var, items = var_items(p, c, "list")
             params = copy.copy(p)
             params[var] = items[int(c["index"])]
-            return expand_list(params, c["command"])
+            return expand_item(params, c["command"])
         if "regex" in c:
             pattern = re.compile(c["regex"])
             if "filter" in c:
-                items = get_items(p, p[c["filter"]])
+                var, items = var_items(p, c, "filter")
                 return [i for i in items if pattern.match(i)]
             elif "group" in c:
-                items = get_items(p, p[c["group"]])
+                var, items = var_items(p, c, "group")
                 groups = {}
                 for i in items:
                     match = pattern.match(i)
@@ -133,52 +172,58 @@ def expand_item(p, c):
                         add_to_group(groups, match)
                 return [groups[k] for k in groups]
             elif "extract" in c:
-                items = get_items(p, p[c["extract"]])
+                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)
+        return [expand_item(p, arg) for arg in c]
     elif isinstance(c, basestring):
-        return [subst.do_substitution(p, c)]
-
-    return []
+        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)
 
-def expand_list(p, l):
-    if isinstance(l, basestring):
-        return expand_item(p, l)
-    else:
-        return [exp for arg in l for exp in expand_item(p, arg)]
+    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):
-    if isinstance(value, dict):
-        return expand_item(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 = [os.path.join(fn, l) for l in os.listdir(fn)]
-        elif stat.S_ISREG(mode):
-            with open(fn) as f:
-                items = [line.rstrip("\r\n") 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
 
+# Construct the cross product of all values of each variable listed in fvars
 def recursive_foreach(params, fvars):
     var = fvars[0]
     fvars = fvars[1:]
@@ -199,7 +244,11 @@ def recursive_foreach(params, fvars):
                         'parameters': params
                     }).execute()
                 else:
-                    logger.info(expand_list(params, params["command"]))
+                    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)
@@ -227,7 +276,13 @@ try:
     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)
 
+try:
     if not args.dry_run:
         if "task.vwd" in taskp:
             # Populate output directory with symlinks to files in collection
@@ -239,9 +294,9 @@ try:
     cmd = []
     if isinstance(taskp["command"][0], list):
         for c in taskp["command"]:
-            cmd.append(expand_list(taskp, c))
+            cmd.append(flatten(expand_item(taskp, c)))
     else:
-        cmd.append(expand_list(taskp, taskp["command"]))
+        cmd.append(flatten(expand_item(taskp, taskp["command"])))
 
     if "task.stdin" in taskp:
         stdinname = subst.do_substitution(taskp, taskp["task.stdin"])
@@ -308,18 +363,23 @@ try:
     signal.signal(signal.SIGQUIT, lambda signum, frame: sig.send_signal(subprocesses, signum))
 
     active = 1
-    while active > 0:
-        os.waitpid(0, 0)
-        active = sum([1 if s.poll() is None else 0 for s in subprocesses])
-
-    # wait for process to complete.
-    rcode = subprocesses[len(subprocesses)-1].returncode
+    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:
         logger.critical("terminating on signal %s" % sig.sig)
         sys.exit(2)
     else:
-        logger.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:
     logger.exception("caught exception")
@@ -348,10 +408,13 @@ if "task.vwd" in taskp:
 else:
     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()