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