13306: Changes to arvados-cwl-runner code after running futurize --stage2
[arvados.git] / 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 cwltool.builder import Builder
7 from .arvjob import ArvadosJob
8 from .arvcontainer import ArvadosContainer
9 from .pathmapper import ArvPathMapper
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 def make_builder(joborder, hints, requirements, runtimeContext):
41     return Builder(
42                  job=joborder,
43                  files=[],               # type: List[Dict[Text, Text]]
44                  bindings=[],            # type: List[Dict[Text, Any]]
45                  schemaDefs={},          # type: Dict[Text, Dict[Text, Any]]
46                  names=None,               # type: Names
47                  requirements=requirements,        # type: List[Dict[Text, Any]]
48                  hints=hints,               # type: List[Dict[Text, Any]]
49                  resources={},           # type: Dict[str, int]
50                  mutation_manager=None,    # type: Optional[MutationManager]
51                  formatgraph=None,         # type: Optional[Graph]
52                  make_fs_access=None,      # type: Type[StdFsAccess]
53                  fs_access=None,           # type: StdFsAccess
54                  job_script_provider=runtimeContext.job_script_provider, # type: Optional[Any]
55                  timeout=runtimeContext.eval_timeout,             # type: float
56                  debug=runtimeContext.debug,               # type: bool
57                  js_console=runtimeContext.js_console,          # type: bool
58                  force_docker_pull=runtimeContext.force_docker_pull,   # type: bool
59                  loadListing="",         # type: Text
60                  outdir="",              # type: Text
61                  tmpdir="",              # type: Text
62                  stagedir="",            # type: Text
63                 )
64
65 class ArvadosCommandTool(CommandLineTool):
66     """Wrap cwltool CommandLineTool to override selected methods."""
67
68     def __init__(self, arvrunner, toolpath_object, loadingContext):
69         super(ArvadosCommandTool, self).__init__(toolpath_object, loadingContext)
70         self.arvrunner = arvrunner
71
72     def make_job_runner(self, runtimeContext):
73         if runtimeContext.work_api == "containers":
74             return partial(ArvadosContainer, self.arvrunner, runtimeContext)
75         elif runtimeContext.work_api == "jobs":
76             return partial(ArvadosJob, self.arvrunner)
77         else:
78             raise Exception("Unsupported work_api %s", runtimeContext.work_api)
79
80     def make_path_mapper(self, reffiles, stagedir, runtimeContext, separateDirs):
81         if runtimeContext.work_api == "containers":
82             return ArvPathMapper(self.arvrunner, reffiles+runtimeContext.extra_reffiles, runtimeContext.basedir,
83                                  "/keep/%s",
84                                  "/keep/%s/%s")
85         elif runtimeContext.work_api == "jobs":
86             return ArvPathMapper(self.arvrunner, reffiles, runtimeContext.basedir,
87                                  "$(task.keep)/%s",
88                                  "$(task.keep)/%s/%s")
89
90     def job(self, joborder, output_callback, runtimeContext):
91         builder = make_builder(joborder, self.hints, self.requirements, runtimeContext)
92         runtimeContext = set_cluster_target(self.tool, self.arvrunner, builder, runtimeContext)
93
94         if runtimeContext.work_api == "containers":
95             dockerReq, is_req = self.get_requirement("DockerRequirement")
96             if dockerReq and dockerReq.get("dockerOutputDirectory"):
97                 runtimeContext.outdir = dockerReq.get("dockerOutputDirectory")
98                 runtimeContext.docker_outdir = dockerReq.get("dockerOutputDirectory")
99             else:
100                 runtimeContext.outdir = "/var/spool/cwl"
101                 runtimeContext.docker_outdir = "/var/spool/cwl"
102         elif runtimeContext.work_api == "jobs":
103             runtimeContext.outdir = "$(task.outdir)"
104             runtimeContext.docker_outdir = "$(task.outdir)"
105             runtimeContext.tmpdir = "$(task.tmpdir)"
106             runtimeContext.docker_tmpdir = "$(task.tmpdir)"
107         return super(ArvadosCommandTool, self).job(joborder, output_callback, runtimeContext)
108
109 class ArvadosExpressionTool(ExpressionTool):
110     def __init__(self, arvrunner, toolpath_object, loadingContext):
111         super(ArvadosExpressionTool, self).__init__(toolpath_object, loadingContext)
112         self.arvrunner = arvrunner
113
114     def job(self,
115             job_order,         # type: Mapping[Text, Text]
116             output_callback,  # type: Callable[[Any, Any], Any]
117             runtimeContext     # type: RuntimeContext
118            ):
119         return super(ArvadosExpressionTool, self).job(job_order, self.arvrunner.get_wrapped_callback(output_callback), runtimeContext)