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