sdk/python/arvados/commands/run.py

   1 #!/usr/bin/env python
   2
   3 import arvados
   4 import argparse
   5 import json
   6 import re
   7 import os
   8 import stat
   9 import put
  10 import arvados.events
  11 import time
  12
  13 arvrun_parser = argparse.ArgumentParser()
  14 arvrun_parser.add_argument('--dry-run', action="store_true")
  15 arvrun_parser.add_argument('--docker-image', type=str, default="arvados/jobs")
  16 arvrun_parser.add_argument('command')
  17 arvrun_parser.add_argument('args', nargs=argparse.REMAINDER)
  18
  19 needupload_files = []
  20
  21 class ArvFile(object):
  22     def __init__(self, prefix, fn):
  23         self.prefix = prefix
  24         self.fn = fn
  25
  26 def statfile(prefix, fn, pattern):
  27     absfn = os.path.abspath(fn)
  28     if os.path.exists(absfn):
  29         fn = os.path.abspath(fn)
  30         st = os.stat(fn)
  31         if stat.S_ISREG(st.st_mode):
  32             mount = os.path.dirname(fn)+"/.arvados#collection"
  33             if os.path.exists(mount):
  34                 with file(mount, 'r') as f:
  35                     c = json.load(f)
  36                 return prefix+"$(file "+c["portable_data_hash"]+"/" + os.path.basename(fn) + ")"
  37             else:
  38                 needupload_files.append(fn)
  39             return ArvFile(prefix, fn[1:])
  40     return prefix+fn
  41
  42 def main(arguments=None):
  43     args = arvrun_parser.parse_args(arguments)
  44
  45     patterns = [re.compile("(--[^=]+=)(.*)"),
  46                 re.compile("(-[^=]+=)(.*)"),
  47                 re.compile("(-.)(.+)")]
  48
  49     commandargs = []
  50
  51     for a in args.args:
  52         if a[0] == '-':
  53             matched = False
  54             for p in patterns:
  55                 m = p.match(a)
  56                 if m:
  57                     commandargs.append(statfile(m.group(1), m.group(2), p))
  58                     matched = True
  59                     break
  60             if not matched:
  61                 commandargs.append(a)
  62         else:
  63             commandargs.append(statfile('', a, None))
  64
  65     n = True
  66     pathprefix = "/"
  67     files = [c for c in commandargs if isinstance(c, ArvFile)]
  68     if len(files) > 0:
  69         while n:
  70             pathstep = None
  71             for c in files:
  72                 if pathstep is None:
  73                     sp = c.fn.split('/')
  74                     if len(sp) < 2:
  75                         n = False
  76                         break
  77                     pathstep = sp[0] + "/"
  78                 else:
  79                     if not c.fn.startswith(pathstep):
  80                         n = False
  81                         break
  82             if n:
  83                 pathprefix += pathstep
  84                 for c in files:
  85                     c.fn = c.fn[len(pathstep):]
  86
  87         os.chdir(pathprefix)
  88
  89         if args.dry_run:
  90             print("cd %s" % pathprefix)
  91             print("arv-put \"%s\"" % '" "'.join([c.fn for c in files]))
  92             pdh = "$(input)"
  93         else:
  94             pdh = put.main(["--portable-data-hash"]+[c.fn for c in files])
  95
  96     commandargs = [("%s$(file %s/%s)" % (c.prefix, pdh, c.fn)) if isinstance(c, ArvFile) else c for c in commandargs]
  97
  98     cut = None
  99     i = -1
 100     stdio = [None, None]
 101     for j in xrange(0, len(commandargs)):
 102         c = commandargs[j]
 103         if c == '<':
 104             stdio[0] = []
 105             i = 0
 106             cut = j if cut is None else cut
 107         elif c == '>':
 108             stdio[1] = []
 109             i = 1
 110             cut = j if cut is None else cut
 111         elif i > -1:
 112             stdio[i].append(c)
 113
 114     if cut is not None:
 115         commandargs = commandargs[:cut]
 116
 117     component = {
 118         "script": "run-command",
 119         "script_version": "bf243e064a7a2ee4e69a87dc3ba46e949a545150",
 120         "repository": "arvados",
 121         "script_parameters": {
 122             "command": [args.command]+commandargs
 123         },
 124         "runtime_constraints": {
 125             "docker_image": args.docker_image
 126         }
 127     }
 128
 129     if stdio[0]:
 130         component["script_parameters"]["task.stdin"] = stdio[0][0]
 131     if stdio[1]:
 132         component["script_parameters"]["task.stdout"] = stdio[1][0]
 133
 134     pipeline = {
 135         "name": "",
 136         "components": {
 137             args.command: component
 138         },
 139         "state":"RunningOnServer"
 140     }
 141
 142     if args.dry_run:
 143         print(json.dumps(pipeline, indent=4))
 144     else:
 145         api = arvados.api('v1')
 146         pi = api.pipeline_instances().create(body=pipeline).execute()
 147         ws = None
 148         def report(x):
 149             if "event_type" in x:
 150                 print "\n"
 151                 print x
 152                 if x["event_type"] == "stderr":
 153                     print x["properties"]["text"]
 154                 elif x["event_type"] == "update" and x["properties"]["new_attributes"]["state"] in ["Complete", "Failed"]:
 155                     ws.close_connection()
 156
 157         ws =  arvados.events.subscribe(api, [["object_uuid", "=", pi["uuid"]], ["event_type", "in", ["stderr", "update"]]], report)
 158         ws.run_forever()
 159
 160 if __name__ == '__main__':
 161     main()