sdk/cwl/arvados_cwl/arvtool.py

   1 # Copyright (C) The Arvados Authors. All rights reserved.
   2 #
   3 # SPDX-License-Identifier: Apache-2.0
   4
   5 from cwltool.command_line_tool import CommandLineTool, ExpressionTool
   6 from .arvcontainer import ArvadosContainer
   7 from .pathmapper import ArvPathMapper
   8 from .runner import make_builder
   9 from ._version import __version__
  10 from functools import partial
  11 from schema_salad.sourceline import SourceLine
  12 from cwltool.errors import WorkflowException
  13
  14 def validate_cluster_target(arvrunner, runtimeContext):
  15     if (runtimeContext.submit_runner_cluster and
  16         runtimeContext.submit_runner_cluster not in arvrunner.api._rootDesc["remoteHosts"] and
  17         runtimeContext.submit_runner_cluster != arvrunner.api._rootDesc["uuidPrefix"]):
  18         raise WorkflowException("Unknown or invalid cluster id '%s' known remote clusters are %s" % (runtimeContext.submit_runner_cluster,
  19                                                                                                      ", ".join(list(arvrunner.api._rootDesc["remoteHosts"].keys()))))
  20     if runtimeContext.project_uuid:
  21         cluster_target = runtimeContext.submit_runner_cluster or arvrunner.api._rootDesc["uuidPrefix"]
  22         if not runtimeContext.project_uuid.startswith(cluster_target):
  23             raise WorkflowException("Project uuid '%s' should start with id of target cluster '%s'" % (runtimeContext.project_uuid, cluster_target))
  24
  25         try:
  26             if runtimeContext.project_uuid[5:12] == '-tpzed-':
  27                 arvrunner.api.users().get(uuid=runtimeContext.project_uuid).execute()
  28             else:
  29                 proj = arvrunner.api.groups().get(uuid=runtimeContext.project_uuid).execute()
  30                 if proj["group_class"] != "project":
  31                     raise Exception("not a project, group_class is '%s'" % (proj["group_class"]))
  32         except Exception as e:
  33             raise WorkflowException("Invalid project uuid '%s': %s" % (runtimeContext.project_uuid, e))
  34
  35 def set_cluster_target(tool, arvrunner, builder, runtimeContext):
  36     cluster_target_req = None
  37     for field in ("hints", "requirements"):
  38         if field not in tool:
  39             continue
  40         for item in tool[field]:
  41             if item["class"] == "http://arvados.org/cwl#ClusterTarget":
  42                 cluster_target_req = item
  43
  44     if cluster_target_req is None:
  45         return runtimeContext
  46
  47     with SourceLine(cluster_target_req, None, WorkflowException, runtimeContext.debug):
  48         runtimeContext = runtimeContext.copy()
  49         runtimeContext.submit_runner_cluster = builder.do_eval(cluster_target_req.get("cluster_id")) or runtimeContext.submit_runner_cluster
  50         runtimeContext.project_uuid = builder.do_eval(cluster_target_req.get("project_uuid")) or runtimeContext.project_uuid
  51         validate_cluster_target(arvrunner, runtimeContext)
  52
  53     return runtimeContext
  54
  55
  56 class ArvadosCommandTool(CommandLineTool):
  57     """Wrap cwltool CommandLineTool to override selected methods."""
  58
  59     def __init__(self, arvrunner, toolpath_object, loadingContext):
  60         super(ArvadosCommandTool, self).__init__(toolpath_object, loadingContext)
  61
  62         (docker_req, docker_is_req) = self.get_requirement("DockerRequirement")
  63         if not docker_req:
  64             self.hints.append({"class": "DockerRequirement",
  65                                "dockerPull": "arvados/jobs:"+__version__})
  66
  67         self.arvrunner = arvrunner
  68
  69     def make_job_runner(self, runtimeContext):
  70         if runtimeContext.work_api == "containers":
  71             return partial(ArvadosContainer, self.arvrunner, runtimeContext)
  72         else:
  73             raise Exception("Unsupported work_api %s", runtimeContext.work_api)
  74
  75     def make_path_mapper(self, reffiles, stagedir, runtimeContext, separateDirs):
  76         if runtimeContext.work_api == "containers":
  77             return ArvPathMapper(self.arvrunner, reffiles+runtimeContext.extra_reffiles, runtimeContext.basedir,
  78                                  "/keep/%s",
  79                                  "/keep/%s/%s")
  80
  81     def job(self, joborder, output_callback, runtimeContext):
  82         builder = make_builder(joborder, self.hints, self.requirements, runtimeContext, self.metadata)
  83         runtimeContext = set_cluster_target(self.tool, self.arvrunner, builder, runtimeContext)
  84
  85         if runtimeContext.work_api == "containers":
  86             dockerReq, is_req = self.get_requirement("DockerRequirement")
  87             if dockerReq and dockerReq.get("dockerOutputDirectory"):
  88                 runtimeContext.outdir = dockerReq.get("dockerOutputDirectory")
  89                 runtimeContext.docker_outdir = dockerReq.get("dockerOutputDirectory")
  90             else:
  91                 runtimeContext.outdir = "/var/spool/cwl"
  92                 runtimeContext.docker_outdir = "/var/spool/cwl"
  93         return super(ArvadosCommandTool, self).job(joborder, output_callback, runtimeContext)
  94
  95 class ArvadosExpressionTool(ExpressionTool):
  96     def __init__(self, arvrunner, toolpath_object, loadingContext):
  97         super(ArvadosExpressionTool, self).__init__(toolpath_object, loadingContext)
  98         self.arvrunner = arvrunner
  99
 100     def job(self,
 101             job_order,         # type: Mapping[Text, Text]
 102             output_callback,  # type: Callable[[Any, Any], Any]
 103             runtimeContext     # type: RuntimeContext
 104            ):
 105         return super(ArvadosExpressionTool, self).job(job_order, self.arvrunner.get_wrapped_callback(output_callback), runtimeContext)