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 .arvjob import ArvadosJob
   7 from .arvcontainer import ArvadosContainer
   8 from .pathmapper import ArvPathMapper
   9 from .runner import make_builder
  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 def set_cluster_target(tool, arvrunner, builder, runtimeContext):
  21     cluster_target_req = None
  22     for field in ("hints", "requirements"):
  23         if field not in tool:
  24             continue
  25         for item in tool[field]:
  26             if item["class"] == "http://arvados.org/cwl#ClusterTarget":
  27                 cluster_target_req = item
  28
  29     if cluster_target_req is None:
  30         return runtimeContext
  31
  32     with SourceLine(cluster_target_req, None, WorkflowException, runtimeContext.debug):
  33         runtimeContext = runtimeContext.copy()
  34         runtimeContext.submit_runner_cluster = builder.do_eval(cluster_target_req.get("cluster_id")) or runtimeContext.submit_runner_cluster
  35         runtimeContext.project_uuid = builder.do_eval(cluster_target_req.get("project_uuid")) or runtimeContext.project_uuid
  36         validate_cluster_target(arvrunner, runtimeContext)
  37
  38     return runtimeContext
  39
  40
  41 class ArvadosCommandTool(CommandLineTool):
  42     """Wrap cwltool CommandLineTool to override selected methods."""
  43
  44     def __init__(self, arvrunner, toolpath_object, loadingContext):
  45         super(ArvadosCommandTool, self).__init__(toolpath_object, loadingContext)
  46         self.arvrunner = arvrunner
  47
  48     def make_job_runner(self, runtimeContext):
  49         if runtimeContext.work_api == "containers":
  50             return partial(ArvadosContainer, self.arvrunner, runtimeContext)
  51         elif runtimeContext.work_api == "jobs":
  52             return partial(ArvadosJob, self.arvrunner)
  53         else:
  54             raise Exception("Unsupported work_api %s", runtimeContext.work_api)
  55
  56     def make_path_mapper(self, reffiles, stagedir, runtimeContext, separateDirs):
  57         if runtimeContext.work_api == "containers":
  58             return ArvPathMapper(self.arvrunner, reffiles+runtimeContext.extra_reffiles, runtimeContext.basedir,
  59                                  "/keep/%s",
  60                                  "/keep/%s/%s")
  61         elif runtimeContext.work_api == "jobs":
  62             return ArvPathMapper(self.arvrunner, reffiles, runtimeContext.basedir,
  63                                  "$(task.keep)/%s",
  64                                  "$(task.keep)/%s/%s")
  65
  66     def job(self, joborder, output_callback, runtimeContext):
  67         builder = make_builder(joborder, self.hints, self.requirements, runtimeContext)
  68         runtimeContext = set_cluster_target(self.tool, self.arvrunner, builder, runtimeContext)
  69
  70         if runtimeContext.work_api == "containers":
  71             dockerReq, is_req = self.get_requirement("DockerRequirement")
  72             if dockerReq and dockerReq.get("dockerOutputDirectory"):
  73                 runtimeContext.outdir = dockerReq.get("dockerOutputDirectory")
  74                 runtimeContext.docker_outdir = dockerReq.get("dockerOutputDirectory")
  75             else:
  76                 runtimeContext.outdir = "/var/spool/cwl"
  77                 runtimeContext.docker_outdir = "/var/spool/cwl"
  78         elif runtimeContext.work_api == "jobs":
  79             runtimeContext.outdir = "$(task.outdir)"
  80             runtimeContext.docker_outdir = "$(task.outdir)"
  81             runtimeContext.tmpdir = "$(task.tmpdir)"
  82             runtimeContext.docker_tmpdir = "$(task.tmpdir)"
  83         return super(ArvadosCommandTool, self).job(joborder, output_callback, runtimeContext)
  84
  85 class ArvadosExpressionTool(ExpressionTool):
  86     def __init__(self, arvrunner, toolpath_object, loadingContext):
  87         super(ArvadosExpressionTool, self).__init__(toolpath_object, loadingContext)
  88         self.arvrunner = arvrunner
  89
  90     def job(self,
  91             job_order,         # type: Mapping[Text, Text]
  92             output_callback,  # type: Callable[[Any, Any], Any]
  93             runtimeContext     # type: RuntimeContext
  94            ):
  95         return super(ArvadosExpressionTool, self).job(job_order, self.arvrunner.get_wrapped_callback(output_callback), runtimeContext)