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("run-command: {}{}".format(' '.join(cmd), (" > " + stdoutname) if stdoutname != None else ""))
  54
  55     rcode = subprocess.call(cmd, stdout=stdoutfile)
  56
  57     print("run-command: completed with exit code %i" % rcode)
  58
  59 except Exception as e:
  60     print("run-command: caught exception: {}".format(e))
  61
  62 finally:
  63     for l in links:
  64         os.unlink(l)
  65
  66     print("run-command: the follow output files will be saved to keep:")
  67
  68     subprocess.call(["find", ".", "-type", "f", "-printf", "run-command: %12.12s %h/%f\\n"])
  69
  70     print("run-command: start writing output to keep")
  71
  72     out = arvados.CollectionWriter()
  73     out.write_directory_tree(".", max_manifest_depth=0)
  74     outuuid = out.finish()
  75     arvados.api('v1').job_tasks().update(uuid=arvados.current_task()['uuid'],
  76                                          body={
  77                                              'output':outuuid,
  78                                              'success': (rcode == 0),
  79                                              'progress':1.0
  80                                          }).execute()
  81
  82 if rcode == 0:
  83     os.chdir("..")
  84     shutil.rmtree("tmpdir")
  85     shutil.rmtree("output")
  86     print("run-command: success")
  87 else:
  88     print("run-command: failed")
  89
  90 sys.exit(rcode)