sdk/cwl/arvados_cwl/__init__.py

   1 #!/usr/bin/env python
   2
   3 import argparse
   4 import arvados
   5 import arvados.events
   6 import arvados.commands.keepdocker
   7 import arvados.commands.run
   8 import cwltool.draft2tool
   9 import cwltool.workflow
  10 import cwltool.main
  11 import threading
  12 import cwltool.docker
  13 import fnmatch
  14 import logging
  15 import re
  16 import os
  17 from cwltool.process import get_feature
  18
  19 logger = logging.getLogger('arvados.cwl-runner')
  20 logger.setLevel(logging.INFO)
  21
  22 def arv_docker_get_image(api_client, dockerRequirement, pull_image):
  23     if "dockerImageId" not in dockerRequirement and "dockerPull" in dockerRequirement:
  24         dockerRequirement["dockerImageId"] = dockerRequirement["dockerPull"]
  25
  26     sp = dockerRequirement["dockerImageId"].split(":")
  27     image_name = sp[0]
  28     image_tag = sp[1] if len(sp) > 1 else None
  29
  30     images = arvados.commands.keepdocker.list_images_in_arv(api_client, 3,
  31                                                             image_name=image_name,
  32                                                             image_tag=image_tag)
  33
  34     if not images:
  35         imageId = cwltool.docker.get_image(dockerRequirement, pull_image)
  36         args = [image_name]
  37         if image_tag:
  38             args.append(image_tag)
  39         arvados.commands.keepdocker.main(args)
  40
  41     return dockerRequirement["dockerImageId"]
  42
  43 class CollectionFsAccess(cwltool.draft2tool.StdFsAccess):
  44     def __init__(self, basedir):
  45         self.collections = {}
  46         self.basedir = basedir
  47
  48     def get_collection(self, path):
  49         p = path.split("/")
  50         if arvados.util.keep_locator_pattern.match(p[0]):
  51             if p[0] not in self.collections:
  52                 self.collections[p[0]] = arvados.collection.CollectionReader(p[0])
  53             return (self.collections[p[0]], "/".join(p[1:]))
  54         else:
  55             return (None, path)
  56
  57     def _match(self, collection, patternsegments, parent):
  58         ret = []
  59         for filename in collection:
  60             if fnmatch.fnmatch(filename, patternsegments[0]):
  61                 cur = os.path.join(parent, filename)
  62                 if len(patternsegments) == 1:
  63                     ret.append(cur)
  64                 else:
  65                     ret.extend(self._match(collection[filename], patternsegments[1:], cur))
  66         return ret
  67
  68     def glob(self, pattern):
  69         collection, rest = self.get_collection(pattern)
  70         patternsegments = rest.split("/")
  71         return self._match(collection, patternsegments, collection.manifest_locator())
  72
  73     def open(self, fn, mode):
  74         collection, rest = self.get_collection(fn)
  75         if collection:
  76             return collection.open(rest, mode)
  77         else:
  78             return open(self._abs(fn), mode)
  79
  80     def exists(self, fn):
  81         collection, rest = self.get_collection(fn)
  82         if collection:
  83             return collection.exists(rest)
  84         else:
  85             return os.path.exists(self._abs(fn))
  86
  87 class ArvadosJob(object):
  88     def __init__(self, runner):
  89         self.arvrunner = runner
  90         self.running = False
  91
  92     def run(self, dry_run=False, pull_image=True, **kwargs):
  93         script_parameters = {
  94             "command": self.command_line
  95         }
  96         runtime_constraints = {}
  97
  98         if self.generatefiles:
  99             vwd = arvados.collection.Collection()
 100             for t in self.generatefiles:
 101                 if isinstance(self.generatefiles[t], dict):
 102                     src, rest = self.arvrunner.fs_access.get_collection(self.generatefiles[t]["path"][6:])
 103                     vwd.copy(rest, t, source_collection=src)
 104                 else:
 105                     with vwd.open(t, "w") as f:
 106                         f.write(self.generatefiles[t])
 107             vwd.save_new()
 108             script_parameters["task.vwd"] = vwd.portable_data_hash()
 109
 110         script_parameters["task.env"] = {"TMPDIR": "$(task.tmpdir)"}
 111         if self.environment:
 112             script_parameters["task.env"].update(self.environment)
 113
 114         if self.stdin:
 115             script_parameters["task.stdin"] = self.pathmapper.mapper(self.stdin)[1]
 116
 117         if self.stdout:
 118             script_parameters["task.stdout"] = self.stdout
 119
 120         (docker_req, docker_is_req) = get_feature(self, "DockerRequirement")
 121         if docker_req and kwargs.get("use_container") is not False:
 122             runtime_constraints["docker_image"] = arv_docker_get_image(self.arvrunner.api, docker_req, pull_image)
 123             runtime_constraints["arvados_sdk_version"] = "master"
 124
 125         response = self.arvrunner.api.jobs().create(body={
 126             "script": "run-command",
 127             "repository": "arvados",
 128             "script_version": "master",
 129             "script_parameters": script_parameters,
 130             "runtime_constraints": runtime_constraints
 131         }, find_or_create=kwargs.get("enable_reuse", True)).execute()
 132
 133         self.arvrunner.jobs[response["uuid"]] = self
 134
 135         logger.info("Job %s is %s", response["uuid"], response["state"])
 136
 137         if response["state"] in ("Complete", "Failed", "Cancelled"):
 138             self.done(response)
 139
 140     def done(self, record):
 141         try:
 142             if record["state"] == "Complete":
 143                 processStatus = "success"
 144             else:
 145                 processStatus = "permanentFail"
 146
 147             try:
 148                 outputs = {}
 149                 outputs = self.collect_outputs(record["output"])
 150             except Exception as e:
 151                 logger.warn(str(e))
 152                 processStatus = "permanentFail"
 153
 154             self.output_callback(outputs, processStatus)
 155         finally:
 156             del self.arvrunner.jobs[record["uuid"]]
 157
 158 class ArvPathMapper(cwltool.pathmapper.PathMapper):
 159     def __init__(self, arvrunner, referenced_files, basedir, **kwargs):
 160         self._pathmap = {}
 161         uploadfiles = []
 162
 163         pdh_path = re.compile(r'^[0-9a-f]{32}\+\d+/.+')
 164
 165         for src in referenced_files:
 166             if isinstance(src, basestring) and pdh_path.match(src):
 167                 self._pathmap[src] = (src, "/keep/%s" % src)
 168             else:
 169                 ab = src if os.path.isabs(src) else os.path.join(basedir, src)
 170                 st = arvados.commands.run.statfile("", ab)
 171                 if kwargs.get("conformance_test"):
 172                     self._pathmap[src] = (src, ab)
 173                 elif isinstance(st, arvados.commands.run.UploadFile):
 174                     uploadfiles.append((src, ab, st))
 175                 elif isinstance(st, arvados.commands.run.ArvFile):
 176                     self._pathmap[src] = (ab, st.fn)
 177                 else:
 178                     raise cwltool.workflow.WorkflowException("Input file path '%s' is invalid" % st)
 179
 180         if uploadfiles:
 181             arvados.commands.run.uploadfiles([u[2] for u in uploadfiles], arvrunner.api, dry_run=kwargs.get("dry_run"), num_retries=3)
 182
 183         for src, ab, st in uploadfiles:
 184             self._pathmap[src] = (ab, st.fn)
 185
 186
 187
 188 class ArvadosCommandTool(cwltool.draft2tool.CommandLineTool):
 189     def __init__(self, arvrunner, toolpath_object, **kwargs):
 190         super(ArvadosCommandTool, self).__init__(toolpath_object, **kwargs)
 191         self.arvrunner = arvrunner
 192
 193     def makeJobRunner(self):
 194         return ArvadosJob(self.arvrunner)
 195
 196     def makePathMapper(self, reffiles, input_basedir, **kwargs):
 197         return ArvPathMapper(self.arvrunner, reffiles, input_basedir, **kwargs)
 198
 199
 200 class ArvCwlRunner(object):
 201     def __init__(self, api_client):
 202         self.api = api_client
 203         self.jobs = {}
 204         self.lock = threading.Lock()
 205         self.cond = threading.Condition(self.lock)
 206         self.final_output = None
 207
 208     def arvMakeTool(self, toolpath_object, **kwargs):
 209         if "class" in toolpath_object and toolpath_object["class"] == "CommandLineTool":
 210             return ArvadosCommandTool(self, toolpath_object, **kwargs)
 211         else:
 212             return cwltool.workflow.defaultMakeTool(toolpath_object, **kwargs)
 213
 214     def output_callback(self, out, processStatus):
 215         if processStatus == "success":
 216             logger.info("Overall job status is %s", processStatus)
 217         else:
 218             logger.warn("Overall job status is %s", processStatus)
 219         self.final_output = out
 220
 221     def on_message(self, event):
 222         if "object_uuid" in event:
 223                 if event["object_uuid"] in self.jobs and event["event_type"] == "update":
 224                     if event["properties"]["new_attributes"]["state"] == "Running" and self.jobs[event["object_uuid"]].running is False:
 225                         logger.info("Job %s is Running", event["object_uuid"])
 226                         with self.lock:
 227                             self.jobs[event["object_uuid"]].running = True
 228                     elif event["properties"]["new_attributes"]["state"] in ("Complete", "Failed", "Cancelled"):
 229                         logger.info("Job %s is %s", event["object_uuid"], event["properties"]["new_attributes"]["state"])
 230                         try:
 231                             self.cond.acquire()
 232                             self.jobs[event["object_uuid"]].done(event["properties"]["new_attributes"])
 233                             self.cond.notify()
 234                         finally:
 235                             self.cond.release()
 236
 237     def arvExecutor(self, tool, job_order, input_basedir, args, **kwargs):
 238         events = arvados.events.subscribe(arvados.api('v1'), [["object_uuid", "is_a", "arvados#job"]], self.on_message)
 239
 240         self.fs_access = CollectionFsAccess(input_basedir)
 241
 242         kwargs["fs_access"] = self.fs_access
 243         kwargs["enable_reuse"] = args.enable_reuse
 244
 245         if kwargs.get("conformance_test"):
 246             return cwltool.main.single_job_executor(tool, job_order, input_basedir, args, **kwargs)
 247         else:
 248             jobiter = tool.job(job_order,
 249                             input_basedir,
 250                             self.output_callback,
 251                             **kwargs)
 252
 253             for runnable in jobiter:
 254                 if runnable:
 255                     with self.lock:
 256                         runnable.run(**kwargs)
 257                 else:
 258                     if self.jobs:
 259                         try:
 260                             self.cond.acquire()
 261                             self.cond.wait()
 262                         finally:
 263                             self.cond.release()
 264                     else:
 265                         logger.error("Workflow cannot make any more progress.")
 266                         break
 267
 268             while self.jobs:
 269                 try:
 270                     self.cond.acquire()
 271                     self.cond.wait()
 272                 finally:
 273                     self.cond.release()
 274
 275             events.close()
 276
 277             if self.final_output is None:
 278                 raise cwltool.workflow.WorkflowException("Workflow did not return a result.")
 279
 280             return self.final_output
 281
 282
 283 def main(args, stdout, stderr, api_client=None):
 284     runner = ArvCwlRunner(api_client=arvados.api('v1'))
 285     args.append("--leave-outputs")
 286     parser = cwltool.main.arg_parser()
 287     exgroup = parser.add_mutually_exclusive_group()
 288     exgroup.add_argument("--enable-reuse", action="store_true",
 289                         default=False, dest="enable_reuse",
 290                         help="")
 291     exgroup.add_argument("--disable-reuse", action="store_false",
 292                         default=False, dest="enable_reuse",
 293                         help="")
 294
 295     return cwltool.main.main(args, executor=runner.arvExecutor, makeTool=runner.arvMakeTool, parser=parser)