crunch_scripts/run-command

   1 #!/usr/bin/env python
   2
   3 import logging
   4
   5 logger = logging.getLogger('run-command')
   6 log_handler = logging.StreamHandler()
   7 log_handler.setFormatter(logging.Formatter("run-command: %(message)s"))
   8 logger.addHandler(log_handler)
   9 logger.setLevel(logging.INFO)
  10
  11 import arvados
  12 import re
  13 import os
  14 import subprocess
  15 import sys
  16 import shutil
  17 import crunchutil.subst as subst
  18 import time
  19 import arvados.commands.put as put
  20 import signal
  21 import stat
  22 import copy
  23 import traceback
  24 import pprint
  25 import multiprocessing
  26 import crunchutil.robust_put as robust_put
  27 import crunchutil.vwd as vwd
  28 import argparse
  29 import json
  30 import tempfile
  31 import errno
  32
  33 parser = argparse.ArgumentParser()
  34 parser.add_argument('--dry-run', action='store_true')
  35 parser.add_argument('--script-parameters', type=str, default="{}")
  36 args = parser.parse_args()
  37
  38 os.umask(0077)
  39
  40 if not args.dry_run:
  41     api = arvados.api('v1')
  42     t = arvados.current_task().tmpdir
  43     os.chdir(arvados.current_task().tmpdir)
  44     os.mkdir("tmpdir")
  45     os.mkdir("output")
  46
  47     os.chdir("output")
  48
  49     outdir = os.getcwd()
  50
  51     taskp = None
  52     jobp = arvados.current_job()['script_parameters']
  53     if len(arvados.current_task()['parameters']) > 0:
  54         taskp = arvados.current_task()['parameters']
  55 else:
  56     outdir = "/tmp"
  57     jobp = json.loads(args.script_parameters)
  58     os.environ['JOB_UUID'] = 'zzzzz-8i9sb-1234567890abcde'
  59     os.environ['TASK_UUID'] = 'zzzzz-ot0gb-1234567890abcde'
  60     os.environ['CRUNCH_SRC'] = '/tmp/crunche-src'
  61     if 'TASK_KEEPMOUNT' not in os.environ:
  62         os.environ['TASK_KEEPMOUNT'] = '/keep'
  63
  64 links = []
  65
  66 def sub_tmpdir(v):
  67     return os.path.join(arvados.current_task().tmpdir, 'tmpdir')
  68
  69 def sub_outdir(v):
  70     return outdir
  71
  72 def sub_cores(v):
  73      return str(multiprocessing.cpu_count())
  74
  75 def sub_jobid(v):
  76      return os.environ['JOB_UUID']
  77
  78 def sub_taskid(v):
  79      return os.environ['TASK_UUID']
  80
  81 def sub_jobsrc(v):
  82      return os.environ['CRUNCH_SRC']
  83
  84 subst.default_subs["task.tmpdir"] = sub_tmpdir
  85 subst.default_subs["task.outdir"] = sub_outdir
  86 subst.default_subs["job.srcdir"] = sub_jobsrc
  87 subst.default_subs["node.cores"] = sub_cores
  88 subst.default_subs["job.uuid"] = sub_jobid
  89 subst.default_subs["task.uuid"] = sub_taskid
  90
  91 class SigHandler(object):
  92     def __init__(self):
  93         self.sig = None
  94
  95     def send_signal(self, subprocesses, signum):
  96         for sp in subprocesses:
  97             sp.send_signal(signum)
  98         self.sig = signum
  99
 100 # http://rightfootin.blogspot.com/2006/09/more-on-python-flatten.html
 101 def flatten(l, ltypes=(list, tuple)):
 102     ltype = type(l)
 103     l = list(l)
 104     i = 0
 105     while i < len(l):
 106         while isinstance(l[i], ltypes):
 107             if not l[i]:
 108                 l.pop(i)
 109                 i -= 1
 110                 break
 111             else:
 112                 l[i:i + 1] = l[i]
 113         i += 1
 114     return ltype(l)
 115
 116 def add_to_group(gr, match):
 117     m = match.groups()
 118     if m not in gr:
 119         gr[m] = []
 120     gr[m].append(match.group(0))
 121
 122 def var_items(p, c, key):
 123     if "var" in c:
 124         # Var specifies
 125         return (c["var"], get_items(p, c[key]))
 126     else:
 127         if isinstance(c[key], list):
 128             return (None, get_items(p, c[key]))
 129         m = re.match("^\$\((.*)\)$", c[key])
 130         if m and m.group(1) in p:
 131             return (m.group(1), get_items(p, c[key]))
 132         else:
 133             # backwards compatible, foreach specifies bare parameter name to use
 134             return (c[key], get_items(p, p[c[key]]))
 135
 136 def expand_item(p, c):
 137     if isinstance(c, dict):
 138         if "foreach" in c and "command" in c:
 139             var, items = var_items(p, c, "foreach")
 140             r = []
 141             for i in items:
 142                 params = copy.copy(p)
 143                 params[var] = i
 144                 r.append(expand_item(params, c["command"]))
 145             return r
 146         if "list" in c and "index" in c and "command" in c:
 147             var, items = var_items(p, c, "list")
 148             params = copy.copy(p)
 149             params[var] = items[int(c["index"])]
 150             return expand_item(params, c["command"])
 151         if "regex" in c:
 152             pattern = re.compile(c["regex"])
 153             if "filter" in c:
 154                 var, items = var_items(p, c, "filter")
 155                 return [i for i in items if pattern.match(i)]
 156             elif "group" in c:
 157                 var, items = var_items(p, c, "group")
 158                 groups = {}
 159                 for i in items:
 160                     match = pattern.match(i)
 161                     if match:
 162                         add_to_group(groups, match)
 163                 return [groups[k] for k in groups]
 164             elif "extract" in c:
 165                 var, items = var_items(p, c, "extract")
 166                 r = []
 167                 for i in items:
 168                     match = pattern.match(i)
 169                     if match:
 170                         r.append(list(match.groups()))
 171                 return r
 172         if "batch" in c and "size" in c:
 173             var, items = var_items(p, c, "batch")
 174             sz = int(c["size"])
 175             r = []
 176             for j in xrange(0, len(items), sz):
 177                 r.append(items[j:j+sz])
 178             return r
 179     elif isinstance(c, list):
 180         return [expand_item(p, arg) for arg in c]
 181     elif isinstance(c, basestring):
 182         m = re.match("^\$\((.*)\)$", c)
 183         if m and m.group(1) in p:
 184             return expand_item(p, p[m.group(1)])
 185         else:
 186             return subst.do_substitution(p, c)
 187
 188     return []
 189
 190 def get_items(p, value):
 191     value = expand_item(p, value)
 192     if isinstance(value, list):
 193         return value
 194     elif isinstance(value, basestring):
 195         mode = os.stat(value).st_mode
 196         prefix = value[len(os.environ['TASK_KEEPMOUNT'])+1:]
 197         if mode is not None:
 198             if stat.S_ISDIR(mode):
 199                 items = [os.path.join(value, l) for l in os.listdir(value)]
 200             elif stat.S_ISREG(mode):
 201                 with open(value) as f:
 202                     items = [line.rstrip("\r\n") for line in f]
 203             return items
 204     raise Exception("get_items did not yield a list")
 205
 206 stdoutname = None
 207 stdoutfile = None
 208 stdinname = None
 209 stdinfile = None
 210
 211 def recursive_foreach(params, fvars):
 212     var = fvars[0]
 213     fvars = fvars[1:]
 214     items = get_items(params, params[var])
 215     logger.info("parallelizing on %s with items %s" % (var, items))
 216     if items is not None:
 217         for i in items:
 218             params = copy.copy(params)
 219             params[var] = i
 220             if len(fvars) > 0:
 221                 recursive_foreach(params, fvars)
 222             else:
 223                 if not args.dry_run:
 224                     arvados.api().job_tasks().create(body={
 225                         'job_uuid': arvados.current_job()['uuid'],
 226                         'created_by_job_task_uuid': arvados.current_task()['uuid'],
 227                         'sequence': 1,
 228                         'parameters': params
 229                     }).execute()
 230                 else:
 231                     if isinstance(params["command"][0], list):
 232                         for c in params["command"]:
 233                             logger.info(flatten(expand_item(params, c)))
 234                     else:
 235                         logger.info(flatten(expand_item(params, params["command"])))
 236     else:
 237         logger.error("parameter %s with value %s in task.foreach yielded no items" % (var, params[var]))
 238         sys.exit(1)
 239
 240 try:
 241     if "task.foreach" in jobp:
 242         if args.dry_run or arvados.current_task()['sequence'] == 0:
 243             # This is the first task to start the other tasks and exit
 244             fvars = jobp["task.foreach"]
 245             if isinstance(fvars, basestring):
 246                 fvars = [fvars]
 247             if not isinstance(fvars, list) or len(fvars) == 0:
 248                 logger.error("value of task.foreach must be a string or non-empty list")
 249                 sys.exit(1)
 250             recursive_foreach(jobp, jobp["task.foreach"])
 251             if not args.dry_run:
 252                 if "task.vwd" in jobp:
 253                     # Set output of the first task to the base vwd collection so it
 254                     # will be merged with output fragments from the other tasks by
 255                     # crunch.
 256                     arvados.current_task().set_output(subst.do_substitution(jobp, jobp["task.vwd"]))
 257                 else:
 258                     arvados.current_task().set_output(None)
 259             sys.exit(0)
 260     else:
 261         # This is the only task so taskp/jobp are the same
 262         taskp = jobp
 263 except Exception as e:
 264     logger.exception("caught exception")
 265     logger.error("job parameters were:")
 266     logger.error(pprint.pformat(jobp))
 267     sys.exit(1)
 268
 269 try:
 270     if not args.dry_run:
 271         if "task.vwd" in taskp:
 272             # Populate output directory with symlinks to files in collection
 273             vwd.checkout(subst.do_substitution(taskp, taskp["task.vwd"]), outdir)
 274
 275         if "task.cwd" in taskp:
 276             os.chdir(subst.do_substitution(taskp, taskp["task.cwd"]))
 277
 278     cmd = []
 279     if isinstance(taskp["command"][0], list):
 280         for c in taskp["command"]:
 281             cmd.append(flatten(expand_item(taskp, c)))
 282     else:
 283         cmd.append(flatten(expand_item(taskp, taskp["command"])))
 284
 285     if "task.stdin" in taskp:
 286         stdinname = subst.do_substitution(taskp, taskp["task.stdin"])
 287         if not args.dry_run:
 288             stdinfile = open(stdinname, "rb")
 289
 290     if "task.stdout" in taskp:
 291         stdoutname = subst.do_substitution(taskp, taskp["task.stdout"])
 292         if not args.dry_run:
 293             stdoutfile = open(stdoutname, "wb")
 294
 295     logger.info("{}{}{}".format(' | '.join([' '.join(c) for c in cmd]), (" < " + stdinname) if stdinname is not None else "", (" > " + stdoutname) if stdoutname is not None else ""))
 296
 297     if args.dry_run:
 298         sys.exit(0)
 299 except subst.SubstitutionError as e:
 300     logger.error(str(e))
 301     logger.error("task parameters were:")
 302     logger.error(pprint.pformat(taskp))
 303     sys.exit(1)
 304 except Exception as e:
 305     logger.exception("caught exception")
 306     logger.error("task parameters were:")
 307     logger.error(pprint.pformat(taskp))
 308     sys.exit(1)
 309
 310 try:
 311     subprocesses = []
 312     close_streams = []
 313     if stdinfile:
 314         close_streams.append(stdinfile)
 315     next_stdin = stdinfile
 316
 317     for i in xrange(len(cmd)):
 318         if i == len(cmd)-1:
 319             # this is the last command in the pipeline, so its stdout should go to stdoutfile
 320             next_stdout = stdoutfile
 321         else:
 322             # this is an intermediate command in the pipeline, so its stdout should go to a pipe
 323             next_stdout = subprocess.PIPE
 324
 325         sp = subprocess.Popen(cmd[i], shell=False, stdin=next_stdin, stdout=next_stdout)
 326
 327         # Need to close the FDs on our side so that subcommands will get SIGPIPE if the
 328         # consuming process ends prematurely.
 329         if sp.stdout:
 330             close_streams.append(sp.stdout)
 331
 332         # Send this processes's stdout to to the next process's stdin
 333         next_stdin = sp.stdout
 334
 335         subprocesses.append(sp)
 336
 337     # File descriptors have been handed off to the subprocesses, so close them here.
 338     for s in close_streams:
 339         s.close()
 340
 341     # Set up signal handling
 342     sig = SigHandler()
 343
 344     # Forward terminate signals to the subprocesses.
 345     signal.signal(signal.SIGINT, lambda signum, frame: sig.send_signal(subprocesses, signum))
 346     signal.signal(signal.SIGTERM, lambda signum, frame: sig.send_signal(subprocesses, signum))
 347     signal.signal(signal.SIGQUIT, lambda signum, frame: sig.send_signal(subprocesses, signum))
 348
 349     active = 1
 350     pids = set([s.pid for s in subprocesses])
 351     rcode = {}
 352     while len(pids) > 0:
 353         (pid, status) = os.wait()
 354         pids.discard(pid)
 355         rcode[pid] = (status >> 8)
 356
 357     if sig.sig is not None:
 358         logger.critical("terminating on signal %s" % sig.sig)
 359         sys.exit(2)
 360     else:
 361         for i in xrange(len(cmd)):
 362             r = rcode[subprocesses[i].pid]
 363             logger.info("%s completed with exit code %i (%s)" % (cmd[i][0], r, "success" if r == 0 else "failed"))
 364
 365 except Exception as e:
 366     logger.exception("caught exception")
 367
 368 # restore default signal handlers.
 369 signal.signal(signal.SIGINT, signal.SIG_DFL)
 370 signal.signal(signal.SIGTERM, signal.SIG_DFL)
 371 signal.signal(signal.SIGQUIT, signal.SIG_DFL)
 372
 373 for l in links:
 374     os.unlink(l)
 375
 376 logger.info("the following output files will be saved to keep:")
 377
 378 subprocess.call(["find", ".", "-type", "f", "-printf", "run-command: %12.12s %h/%f\\n"], stdout=sys.stderr)
 379
 380 logger.info("start writing output to keep")
 381
 382 if "task.vwd" in taskp:
 383     if "task.foreach" in jobp:
 384         # This is a subtask, so don't merge with the original collection, that will happen at the end
 385         outcollection = vwd.checkin(subst.do_substitution(taskp, taskp["task.vwd"]), outdir, merge=False).manifest_text()
 386     else:
 387         # Just a single task, so do merge with the original collection
 388         outcollection = vwd.checkin(subst.do_substitution(taskp, taskp["task.vwd"]), outdir, merge=True).manifest_text()
 389 else:
 390     outcollection = robust_put.upload(outdir, logger)
 391
 392 success = reduce(lambda x, y: x & (y == 0), [True]+rcode.values())
 393
 394 api.job_tasks().update(uuid=arvados.current_task()['uuid'],
 395                                      body={
 396                                          'output': outcollection,
 397                                          'success': success,
 398                                          'progress':1.0
 399                                      }).execute()
 400
 401 sys.exit(rcode)