crunch_scripts/run-command

   1 #!/usr/bin/env python
   2
   3 import arvados
   4 import re
   5 import os
   6 import subprocess
   7 import sys
   8 import shutil
   9 import subst
  10
  11 os.umask(0077)
  12
  13 t = arvados.current_task().tmpdir
  14
  15 os.chdir(arvados.current_task().tmpdir)
  16 os.mkdir("tmpdir")
  17 os.mkdir("output")
  18
  19 os.chdir("output")
  20
  21 if len(arvados.current_task()['parameters']) > 0:
  22     p = arvados.current_task()['parameters']
  23 else:
  24     p = arvados.current_job()['script_parameters']
  25
  26 links = []
  27
  28 def sub_link(v):
  29     r = os.path.basename(v)
  30     os.symlink(os.path.join(os.environ['TASK_KEEPMOUNT'], v) , r)
  31     links.append(r)
  32     return r
  33
  34 def sub_tmpdir(v):
  35     return os.path.join(arvados.current_task().tmpdir, 'tmpdir')
  36
  37 subst.default_subs["link "] = sub_link
  38 subst.default_subs["tmpdir"] = sub_tmpdir
  39
  40 rcode = 1
  41
  42 try:
  43     cmd = []
  44     for c in p["command"]:
  45         cmd.append(subst.do_substitution(p, c))
  46
  47     stdoutname = None
  48     stdoutfile = None
  49     if "stdout" in p:
  50         stdoutname = subst.do_substitution(p, p["stdout"])
  51         stdoutfile = open(stdoutname, "wb")
  52
  53     print("Running command: {}{}".format(' '.join(cmd), (" > " + stdoutname) if stdoutname != None else ""))
  54
  55     rcode = subprocess.call(cmd, stdout=stdoutfile)
  56
  57 except Exception as e:
  58     print("Caught exception {}".format(e))
  59
  60 finally:
  61     for l in links:
  62         os.unlink(l)
  63
  64     out = arvados.CollectionWriter()
  65     out.write_directory_tree(".", max_manifest_depth=0)
  66     arvados.current_task().set_output(out.finish())
  67
  68 if rcode == 0:
  69     os.chdir("..")
  70     shutil.rmtree("tmpdir")
  71     shutil.rmtree("output")
  72
  73 sys.exit(rcode)