-
Notifications
You must be signed in to change notification settings - Fork 0
/
launch_pipelines.py
386 lines (313 loc) · 12.2 KB
/
launch_pipelines.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
import argparse
import datetime
import json
import logging
import pydantic
import uuid
import yaml
from pathlib import Path
from seqerakit import seqeraplatform
from seqerakit.helper import parse_launch_block
import yaml
## Globals
# Global UUID for the launch name
workflow_uuid = str(uuid.uuid4()).replace("-", "")[:15]
# Global date for the launch name
date = datetime.datetime.now().strftime("%Y%m%d")
class SeqeraKitError(Exception):
"""Exception for failure to use Tower CLI."""
pass
class Pipeline(pydantic.BaseModel):
"""A pipeline to launch."""
name: str
url: str
latest: bool
profiles: list[str]
revision: str | None = None
class ComputeEnvironment(pydantic.BaseModel):
"""A compute environment to launch a pipeline on."""
ref: str
name: str
workdir: str
workspace: str
profiles: list[str] = []
class LaunchConfig(pydantic.BaseModel):
"""A pipeline and compute environment to launch a pipeline on
TODO: Make this a full-fledged class with methods.
"""
pipeline: "Pipeline"
compute_environment: "ComputeEnvironment"
def __eq__(self, other, strict: bool = False) -> bool:
# Check classes make sense
if other.__class__ is self.__class__:
# If strict, check the entire class matches
if strict:
return self == other
# Checks pipeline name and compute env name match
# Will ignore if other variables are different.
else:
return self.pipeline.model_dump().get(
"name"
) == other.pipeline.model_dump().get(
"name"
) and self.compute_environment.model_dump().get(
"name"
) == other.compute_environment.model_dump().get(
"name"
)
else:
return NotImplemented
def launch_pipeline(
self, seqera: seqeraplatform.SeqeraPlatform, wait: str = "SUBMITTED"
) -> dict[str, str | bool | None]:
"""
Launch a pipeline.
Args:
seqera (seqeraplatform.SeqeraPlatform): A SeqeraPlatform object.
wait (str, optional): The wait status for the pipeline. Defaults to "SUBMITTED".
Raises:
SeqeraKitError: If the pipeline fails to launch.
Returns:
dict[str, str]: The launched pipeline.
"""
# Pre-create some variables to make things easier.
run_name = "_".join(
[self.pipeline.name, self.compute_environment.ref, date, workflow_uuid]
)
# It's never good to create a path with string handling but it's the quickest way here.
workdir = "/".join(
[self.compute_environment.workdir, self.pipeline.name, "work-" + date]
)
outdir = "/".join(
[
self.compute_environment.workdir,
self.pipeline.name,
"results-test-" + date,
]
)
# Add trailing slash to outdir if it doesn't exist
if not outdir.endswith("/"):
outdir += "/"
# Create params dict
params = {"outdir": outdir}
# Launch the pipeline and wait for submission.
logging.info(
f"Launching pipeline {self.pipeline.name} on {self.compute_environment.ref}."
)
args_dict = {
"workspace": self.compute_environment.workspace,
"compute-env": self.compute_environment.name,
"work-dir": workdir,
"name": run_name,
"wait": wait,
"params": params,
"pipeline": self.pipeline.url,
}
if self.pipeline.revision is not None:
args_dict.update({"revision": self.pipeline.revision})
if self.pipeline.profiles != [] or self.compute_environment.profiles != []:
# Create profiles string
args_dict.update(
{
"profile": str(
",".join(
self.pipeline.profiles + self.compute_environment.profiles
)
)
}
)
default_response = {
"workflowId": None,
"workflowUrl": None,
"workspaceId": None,
"workspaceRef": None,
"workflowName": run_name,
"computeEnvironment": self.compute_environment.name,
"launchSuccess": False,
"error": "",
}
try:
# Use seqerakit helper function to construct arguments
args_list = parse_launch_block(args_dict)
launched_pipeline = seqera.launch(*args_list, to_json=True)
# If dryrun, return default response
if seqera.dryrun:
return default_response
# If we fail to add the pipeline for a predictable reason we can log and continue
except seqeraplatform.ResourceCreationError as err:
logging.info(
f"Failed to launch pipeline {run_name}. Logging and proceeding..."
)
message = "\n".join(err.args)
logging.debug(message)
# Raise pipeline launch error here:
default_response.update({"error": message})
return default_response
# If we fail to add the pipeline for an unpredictable reason we log and fail
except json.decoder.JSONDecodeError as err:
logging.error(f"Failed to launch pipeline {run_name}.")
logging.debug(err.doc)
# Raise pipeline launch error here:
raise SeqeraKitError(err.doc)
# Add pipeline launch info to dict
launched_pipeline.update(
{
"workflowName": run_name,
"computeEnvironment": self.compute_environment.name,
"launchSuccess": True,
"error": "",
}
)
return launched_pipeline
# Need to use update_forward_refs() to resolve circular references in Pydantic.
LaunchConfig.model_rebuild()
def parse_args() -> argparse.Namespace:
"""
Parse command line arguments.
Returns:
argparse.Namespace: The parsed arguments.
"""
parser = argparse.ArgumentParser(
description="Launch a matrix of pipelines and compute environments."
)
parser.add_argument(
"-l",
"--log_level",
default="INFO",
choices=("CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG"),
help="The desired log level (default: INFO).",
type=str.upper,
)
parser.add_argument(
"-o", "--output", type=str, required=True, help="Output filename for JSON file."
)
parser.add_argument(
"-i",
"--inputs",
nargs="+",
required=True,
type=Path,
help="The input yaml files to read. Must contain keys 'include', 'exclude', 'compute-envs' and 'pipelines'.",
)
parser.add_argument(
"-d",
"--dryrun",
action="store_true",
help="Dry run the pipeline launch without actually launching.",
)
return parser.parse_args()
def read_yaml(paths: list[str]) -> list[LaunchConfig]:
"""
Read multiple YAML files of pipeline, compute-env, include and exclude YAML files
then create a list of launch configs from the resulting mix. Assumes keys 'include',
'exclude', 'compute-envs' and 'pipelines' are present in the YAML files.
Args:
paths (list[str]): The paths to the YAML files.
Returns:
list[Pipeline]: A list of pipelines read from YAML.
"""
logging.info("Reading launch details...")
# Pre-populate empty output to fill
# This saves us doing lots of if/or statements for getting the contents
objects: dict = {"pipelines": [], "compute-envs": [], "include": [], "exclude": []}
for path in paths:
with open(path) as pipeline_file:
# Extend existing dictionary values (lists)
# We grab keys from pre-populated dict so we can do some key checking
file_contents = yaml.safe_load(pipeline_file)
for key in file_contents.keys():
if key in objects.keys():
objects[key] = objects[key] + file_contents[key]
else:
raise KeyError(f"Unexpected key in YAML file: {key}")
# Get pipeline details from 'pipelines' key
pipelines = [Pipeline(**pipeline) for pipeline in objects["pipelines"]]
# Get compute env details from 'compute-envs' key
compute_envs = [
ComputeEnvironment(**compute_env) for compute_env in objects["compute-envs"]
]
# Get include and exclude details from 'include' and 'exclude' keys
include = [LaunchConfig(**include) for include in objects["include"]]
exclude = [LaunchConfig(**exclude) for exclude in objects["exclude"]]
# Create matrix of pipeline * compute-envs to LaunchConfigs
launch_configs = create_launch_config(pipelines, compute_envs)
# Add any included LaunchConfigs and remove excluded LaunchConfigs
complete_launch_configs = filter_launch_configs(launch_configs, include, exclude)
return complete_launch_configs
def create_launch_config(
pipelines: list[Pipeline],
compute_envs: list[ComputeEnvironment],
include: list[LaunchConfig] = [],
) -> list[LaunchConfig]:
"""
Create a list of launch configs from a list of pipelines and compute environments.
Args:
pipelines (list[Pipeline]): A list of pipelines.
compute_envs (list[ComputeEnvironment]): A list of compute environments.
include (list[LaunchConfig], optional): A list of launch configs to include in addition to pipelines * compute envs. Defaults to [].
Returns:
list[LaunchConfig]: A list of launch configs.
"""
launch_configs = []
# Might be able to do this cleaner with itertools.combinations()
for pipeline in pipelines:
for compute_env in compute_envs:
launch_config = LaunchConfig(
pipeline=pipeline, compute_environment=compute_env
)
launch_configs.append(launch_config)
for launch_config in include:
launch_configs.append(launch_config)
return launch_configs
def filter_launch_configs(
launch_configs: list[LaunchConfig],
include: list[LaunchConfig] = [],
exclude: list[LaunchConfig] = [],
) -> list[LaunchConfig]:
"""
Filter a list of launch configs by include and exclude lists.
Args:
launch_configs (list[LaunchConfig]): A list of initial launch configs.
include (list[LaunchConfig]): A list of launch configs to include.
exclude (list[LaunchConfig]): A list of launch configs to exclude.
Returns:
list[LaunchConfig]: A list of filtered launch configs.
"""
logging.info("Adding include launch configs to full set...")
full_launch_configs = launch_configs + include
logging.info("Removing exclude launch configs from full set...")
filtered_launch_configs = [
launch_config
for launch_config in full_launch_configs
if launch_config not in exclude
]
return filtered_launch_configs
def launch_pipelines(
seqera: seqeraplatform.SeqeraPlatform,
launch_configs: list[LaunchConfig],
) -> list[dict[str, str | bool | None]]:
"""
Launch a list of pipelines.
Args:
seqera (seqeraplatform.SeqeraPlatform): A SeqeraPlatform object.
launch_configs (list[LaunchConfig]): A list of launch configs.
Returns:
list[dict[str, str]]: A list of launched pipelines.
"""
logging.info("Launching pipelines.")
launched_pipelines = [
launch_config.launch_pipeline(seqera=seqera) for launch_config in launch_configs
]
logging.info("Pipelines launched.")
return launched_pipelines
def main() -> None:
args = parse_args()
logging.basicConfig(level=logging.INFO)
seqera = seqeraplatform.SeqeraPlatform(dryrun=args.dryrun)
complete_launch_configs = read_yaml(args.inputs)
launched_pipelines = launch_pipelines(seqera, complete_launch_configs)
logging.info(f"Writing launches to JSON file {args.output}")
with open(args.output, "w") as output_file:
json.dump(launched_pipelines, output_file, indent=4)
if __name__ == "__main__":
main()