runnable 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (71) hide show
  1. runnable/__init__.py +34 -0
  2. runnable/catalog.py +141 -0
  3. runnable/cli.py +272 -0
  4. runnable/context.py +34 -0
  5. runnable/datastore.py +686 -0
  6. runnable/defaults.py +179 -0
  7. runnable/entrypoints.py +484 -0
  8. runnable/exceptions.py +94 -0
  9. runnable/executor.py +431 -0
  10. runnable/experiment_tracker.py +139 -0
  11. runnable/extensions/catalog/__init__.py +21 -0
  12. runnable/extensions/catalog/file_system/__init__.py +0 -0
  13. runnable/extensions/catalog/file_system/implementation.py +226 -0
  14. runnable/extensions/catalog/k8s_pvc/__init__.py +0 -0
  15. runnable/extensions/catalog/k8s_pvc/implementation.py +16 -0
  16. runnable/extensions/catalog/k8s_pvc/integration.py +59 -0
  17. runnable/extensions/executor/__init__.py +714 -0
  18. runnable/extensions/executor/argo/__init__.py +0 -0
  19. runnable/extensions/executor/argo/implementation.py +1182 -0
  20. runnable/extensions/executor/argo/specification.yaml +51 -0
  21. runnable/extensions/executor/k8s_job/__init__.py +0 -0
  22. runnable/extensions/executor/k8s_job/implementation_FF.py +259 -0
  23. runnable/extensions/executor/k8s_job/integration_FF.py +69 -0
  24. runnable/extensions/executor/local/__init__.py +0 -0
  25. runnable/extensions/executor/local/implementation.py +69 -0
  26. runnable/extensions/executor/local_container/__init__.py +0 -0
  27. runnable/extensions/executor/local_container/implementation.py +367 -0
  28. runnable/extensions/executor/mocked/__init__.py +0 -0
  29. runnable/extensions/executor/mocked/implementation.py +220 -0
  30. runnable/extensions/experiment_tracker/__init__.py +0 -0
  31. runnable/extensions/experiment_tracker/mlflow/__init__.py +0 -0
  32. runnable/extensions/experiment_tracker/mlflow/implementation.py +94 -0
  33. runnable/extensions/nodes.py +675 -0
  34. runnable/extensions/run_log_store/__init__.py +0 -0
  35. runnable/extensions/run_log_store/chunked_file_system/__init__.py +0 -0
  36. runnable/extensions/run_log_store/chunked_file_system/implementation.py +106 -0
  37. runnable/extensions/run_log_store/chunked_k8s_pvc/__init__.py +0 -0
  38. runnable/extensions/run_log_store/chunked_k8s_pvc/implementation.py +21 -0
  39. runnable/extensions/run_log_store/chunked_k8s_pvc/integration.py +61 -0
  40. runnable/extensions/run_log_store/db/implementation_FF.py +157 -0
  41. runnable/extensions/run_log_store/db/integration_FF.py +0 -0
  42. runnable/extensions/run_log_store/file_system/__init__.py +0 -0
  43. runnable/extensions/run_log_store/file_system/implementation.py +136 -0
  44. runnable/extensions/run_log_store/generic_chunked.py +541 -0
  45. runnable/extensions/run_log_store/k8s_pvc/__init__.py +0 -0
  46. runnable/extensions/run_log_store/k8s_pvc/implementation.py +21 -0
  47. runnable/extensions/run_log_store/k8s_pvc/integration.py +56 -0
  48. runnable/extensions/secrets/__init__.py +0 -0
  49. runnable/extensions/secrets/dotenv/__init__.py +0 -0
  50. runnable/extensions/secrets/dotenv/implementation.py +100 -0
  51. runnable/extensions/secrets/env_secrets/__init__.py +0 -0
  52. runnable/extensions/secrets/env_secrets/implementation.py +42 -0
  53. runnable/graph.py +464 -0
  54. runnable/integration.py +205 -0
  55. runnable/interaction.py +399 -0
  56. runnable/names.py +546 -0
  57. runnable/nodes.py +489 -0
  58. runnable/parameters.py +183 -0
  59. runnable/pickler.py +102 -0
  60. runnable/sdk.py +470 -0
  61. runnable/secrets.py +95 -0
  62. runnable/tasks.py +392 -0
  63. runnable/utils.py +630 -0
  64. runnable-0.2.0.dist-info/METADATA +437 -0
  65. runnable-0.2.0.dist-info/RECORD +69 -0
  66. runnable-0.2.0.dist-info/entry_points.txt +44 -0
  67. runnable-0.1.0.dist-info/METADATA +0 -16
  68. runnable-0.1.0.dist-info/RECORD +0 -6
  69. /runnable/{.gitkeep → extensions/__init__.py} +0 -0
  70. {runnable-0.1.0.dist-info → runnable-0.2.0.dist-info}/LICENSE +0 -0
  71. {runnable-0.1.0.dist-info → runnable-0.2.0.dist-info}/WHEEL +0 -0
@@ -0,0 +1,1182 @@
1
+ import json
2
+ import logging
3
+ import random
4
+ import shlex
5
+ import string
6
+ from abc import ABC, abstractmethod
7
+ from collections import OrderedDict
8
+ from typing import Any, Dict, List, Optional, Union, cast
9
+
10
+ from pydantic import BaseModel, ConfigDict, Field, computed_field, field_serializer, field_validator
11
+ from pydantic.functional_serializers import PlainSerializer
12
+ from ruamel.yaml import YAML
13
+ from typing_extensions import Annotated
14
+
15
+ from runnable import defaults, exceptions, integration, parameters, utils
16
+ from runnable.defaults import TypeMapVariable
17
+ from runnable.extensions.executor import GenericExecutor
18
+ from runnable.extensions.nodes import DagNode, MapNode, ParallelNode
19
+ from runnable.graph import Graph, create_node, search_node_by_internal_name
20
+ from runnable.integration import BaseIntegration
21
+ from runnable.nodes import BaseNode
22
+
23
+ logger = logging.getLogger(defaults.NAME)
24
+
25
+ # TODO: Leave the run log in consistent state.
26
+
27
+ """
28
+ executor:
29
+ type: argo
30
+ config:
31
+ image: # apply to template
32
+ max_workflow_duration: # Apply to spec
33
+ nodeSelector: #Apply to spec
34
+ parallelism: #apply to spec
35
+ resources: # convert to podSpecPath
36
+ limits:
37
+ requests:
38
+ retryStrategy:
39
+ max_step_duration: # apply to templateDefaults
40
+ step_timeout: # apply to templateDefaults
41
+ tolerations: # apply to spec
42
+ imagePullPolicy: # apply to template
43
+
44
+ overrides:
45
+ override:
46
+ tolerations: # template
47
+ image: # container
48
+ max_step_duration: # template
49
+ step_timeout: #template
50
+ nodeSelector: #template
51
+ parallelism: # this need to applied for map
52
+ resources: # container
53
+ imagePullPolicy: #container
54
+ retryStrategy: # template
55
+ """
56
+
57
+
58
+ class SecretEnvVar(BaseModel):
59
+ """
60
+ Renders:
61
+ env:
62
+ - name: MYSECRETPASSWORD
63
+ valueFrom:
64
+ secretKeyRef:
65
+ name: my-secret
66
+ key: mypassword
67
+ """
68
+
69
+ environment_variable: str = Field(serialization_alias="name")
70
+ secret_name: str = Field(exclude=True)
71
+ secret_key: str = Field(exclude=True)
72
+
73
+ @computed_field # type: ignore
74
+ @property
75
+ def valueFrom(self) -> Dict[str, Dict[str, str]]:
76
+ return {
77
+ "secretKeyRef": {
78
+ "name": self.secret_name,
79
+ "key": self.secret_key,
80
+ }
81
+ }
82
+
83
+
84
+ class EnvVar(BaseModel):
85
+ """
86
+ Renders:
87
+ parameters: # in arguments
88
+ - name: x
89
+ value: 3 # This is optional for workflow parameters
90
+
91
+ """
92
+
93
+ name: str
94
+ value: Union[str, int, float] = Field(default="")
95
+
96
+
97
+ class Parameter(BaseModel):
98
+ name: str
99
+ value: Optional[str] = None
100
+
101
+ @field_serializer("name")
102
+ def serialize_name(self, name: str) -> str:
103
+ return f"{str(name)}"
104
+
105
+ @field_serializer("value")
106
+ def serialize_value(self, value: str) -> str:
107
+ return f"{value}"
108
+
109
+
110
+ class OutputParameter(Parameter):
111
+ """
112
+ Renders:
113
+ - name: step-name
114
+ valueFrom:
115
+ path: /tmp/output.txt
116
+ """
117
+
118
+ path: str = Field(default="/tmp/output.txt", exclude=True)
119
+
120
+ @computed_field # type: ignore
121
+ @property
122
+ def valueFrom(self) -> Dict[str, str]:
123
+ return {"path": self.path}
124
+
125
+
126
+ class Argument(BaseModel):
127
+ """
128
+ Templates are called with arguments, which become inputs for the template
129
+ Renders:
130
+ arguments:
131
+ parameters:
132
+ - name: The name of the parameter
133
+ value: The value of the parameter
134
+ """
135
+
136
+ name: str
137
+ value: str
138
+
139
+ @field_serializer("name")
140
+ def serialize_name(self, name: str) -> str:
141
+ return f"{str(name)}"
142
+
143
+ @field_serializer("value")
144
+ def serialize_value(self, value: str) -> str:
145
+ return f"{value}"
146
+
147
+
148
+ class Request(BaseModel):
149
+ """
150
+ The default requests
151
+ """
152
+
153
+ memory: str = "1Gi"
154
+ cpu: str = "250m"
155
+
156
+
157
+ VendorGPU = Annotated[
158
+ Optional[int],
159
+ PlainSerializer(lambda x: str(x), return_type=str, when_used="unless-none"),
160
+ ]
161
+
162
+
163
+ class Limit(Request):
164
+ """
165
+ The default limits
166
+ """
167
+
168
+ gpu: VendorGPU = Field(default=None, serialization_alias="nvidia.com/gpu")
169
+
170
+
171
+ class Resources(BaseModel):
172
+ limits: Limit = Field(default=Limit(), serialization_alias="limits")
173
+ requests: Request = Field(default=Request(), serialization_alias="requests")
174
+
175
+
176
+ class BackOff(BaseModel):
177
+ duration_in_seconds: int = Field(default=2 * 60, serialization_alias="duration")
178
+ factor: float = Field(default=2, serialization_alias="factor")
179
+ max_duration: int = Field(default=60 * 60, serialization_alias="maxDuration")
180
+
181
+ @field_serializer("duration_in_seconds")
182
+ def cast_duration_as_str(self, duration_in_seconds: int, _info) -> str:
183
+ return str(duration_in_seconds)
184
+
185
+ @field_serializer("max_duration")
186
+ def cast_mas_duration_as_str(self, max_duration: int, _info) -> str:
187
+ return str(max_duration)
188
+
189
+
190
+ class Retry(BaseModel):
191
+ limit: int = 0
192
+ retry_policy: str = Field(default="Always", serialization_alias="retryPolicy")
193
+ back_off: BackOff = Field(default=BackOff(), serialization_alias="backoff")
194
+
195
+ @field_serializer("limit")
196
+ def cast_limit_as_str(self, limit: int, _info) -> str:
197
+ return str(limit)
198
+
199
+
200
+ class Toleration(BaseModel):
201
+ effect: str
202
+ key: str
203
+ operator: str
204
+ value: str
205
+
206
+
207
+ class TemplateDefaults(BaseModel):
208
+ max_step_duration: int = Field(
209
+ default=60 * 60 * 2,
210
+ serialization_alias="activeDeadlineSeconds",
211
+ gt=0,
212
+ description="Max run time of a step",
213
+ )
214
+
215
+ @computed_field # type: ignore
216
+ @property
217
+ def timeout(self) -> str:
218
+ return f"{self.max_step_duration + 60*60}s"
219
+
220
+
221
+ ShlexCommand = Annotated[str, PlainSerializer(lambda x: shlex.split(x), return_type=List[str])]
222
+
223
+
224
+ class Container(BaseModel):
225
+ image: str
226
+ command: ShlexCommand
227
+ volume_mounts: Optional[List["ContainerVolume"]] = Field(default=None, serialization_alias="volumeMounts")
228
+ image_pull_policy: str = Field(default="", serialization_alias="imagePullPolicy")
229
+ resources: Optional[Resources] = Field(default=None, serialization_alias="resources")
230
+
231
+ env_vars: List[EnvVar] = Field(default_factory=list, exclude=True)
232
+ secrets_from_k8s: List[SecretEnvVar] = Field(default_factory=list, exclude=True)
233
+
234
+ @computed_field # type: ignore
235
+ @property
236
+ def env(self) -> Optional[List[Union[EnvVar, SecretEnvVar]]]:
237
+ if not self.env_vars and not self.secrets_from_k8s:
238
+ return None
239
+
240
+ return self.env_vars + self.secrets_from_k8s
241
+
242
+
243
+ class DagTaskTemplate(BaseModel):
244
+ """
245
+ dag:
246
+ tasks:
247
+ name: A
248
+ template: nested-diamond
249
+ arguments:
250
+ parameters: [{name: message, value: A}]
251
+ """
252
+
253
+ name: str
254
+ template: str
255
+ depends: List[str] = []
256
+ arguments: Optional[List[Argument]] = Field(default=None)
257
+ with_param: Optional[str] = Field(default=None, serialization_alias="withParam")
258
+
259
+ @field_serializer("depends")
260
+ def transform_depends_as_str(self, depends: List[str]) -> str:
261
+ return " || ".join(depends)
262
+
263
+ @field_serializer("arguments", when_used="unless-none")
264
+ def empty_arguments_to_none(self, arguments: List[Argument]) -> Dict[str, List[Argument]]:
265
+ return {"parameters": arguments}
266
+
267
+
268
+ class ContainerTemplate(BaseModel):
269
+ # These templates are used for actual execution nodes.
270
+ name: str
271
+ active_deadline_seconds: Optional[int] = Field(default=None, serialization_alias="activeDeadlineSeconds", gt=0)
272
+ node_selector: Optional[Dict[str, str]] = Field(default=None, serialization_alias="nodeSelector")
273
+ retry_strategy: Optional[Retry] = Field(default=None, serialization_alias="retryStrategy")
274
+ tolerations: Optional[List[Toleration]] = Field(default=None, serialization_alias="tolerations")
275
+
276
+ container: Container
277
+
278
+ outputs: Optional[List[OutputParameter]] = Field(default=None, serialization_alias="outputs")
279
+ inputs: Optional[List[Parameter]] = Field(default=None, serialization_alias="inputs")
280
+
281
+ def __hash__(self):
282
+ return hash(self.name)
283
+
284
+ @field_serializer("outputs", when_used="unless-none")
285
+ def reshape_outputs(self, outputs: List[OutputParameter]) -> Dict[str, List[OutputParameter]]:
286
+ return {"parameters": outputs}
287
+
288
+ @field_serializer("inputs", when_used="unless-none")
289
+ def reshape_inputs(self, inputs: List[Parameter]) -> Dict[str, List[Parameter]]:
290
+ return {"parameters": inputs}
291
+
292
+
293
+ class DagTemplate(BaseModel):
294
+ # These are used for parallel, map nodes dag definition
295
+ name: str = "magnus-dag"
296
+ tasks: List[DagTaskTemplate] = Field(default=[], exclude=True)
297
+ inputs: Optional[List[Parameter]] = Field(default=None, serialization_alias="inputs")
298
+ parallelism: Optional[int] = None
299
+ fail_fast: bool = Field(default=True, serialization_alias="failFast")
300
+
301
+ @field_validator("parallelism")
302
+ @classmethod
303
+ def validate_parallelism(cls, parallelism: Optional[int]) -> Optional[int]:
304
+ if parallelism is not None and parallelism <= 0:
305
+ raise ValueError("Parallelism must be a positive integer greater than 0")
306
+ return parallelism
307
+
308
+ @computed_field # type: ignore
309
+ @property
310
+ def dag(self) -> Dict[str, List[DagTaskTemplate]]:
311
+ return {"tasks": self.tasks}
312
+
313
+ @field_serializer("inputs", when_used="unless-none")
314
+ def reshape_inputs(self, inputs: List[Parameter], _info) -> Dict[str, List[Parameter]]:
315
+ return {"parameters": inputs}
316
+
317
+
318
+ class Volume(BaseModel):
319
+ """
320
+ spec config requires, name and persistentVolumeClaim
321
+ step requires name and mountPath
322
+ """
323
+
324
+ name: str
325
+ claim: str = Field(exclude=True)
326
+ mount_path: str = Field(serialization_alias="mountPath", exclude=True)
327
+
328
+ @computed_field # type: ignore
329
+ @property
330
+ def persistentVolumeClaim(self) -> Dict[str, str]:
331
+ return {"claimName": self.claim}
332
+
333
+
334
+ class ContainerVolume(BaseModel):
335
+ name: str
336
+ mount_path: str = Field(serialization_alias="mountPath")
337
+
338
+
339
+ class UserVolumeMounts(BaseModel):
340
+ """
341
+ The volume specification as user defines it.
342
+ """
343
+
344
+ name: str # This is the name of the PVC on K8s
345
+ mount_path: str # This is mount path on the container
346
+
347
+
348
+ class NodeRenderer(ABC):
349
+ allowed_node_types: List[str] = []
350
+
351
+ def __init__(self, executor: "ArgoExecutor", node: BaseNode) -> None:
352
+ self.executor = executor
353
+ self.node = node
354
+
355
+ @abstractmethod
356
+ def render(self, list_of_iter_values: Optional[List] = None):
357
+ pass
358
+
359
+
360
+ class ExecutionNode(NodeRenderer):
361
+ allowed_node_types = ["task", "stub", "success", "fail"]
362
+
363
+ def render(self, list_of_iter_values: Optional[List] = None):
364
+ """
365
+ Compose the map variable and create the execution command.
366
+ Create an input to the command.
367
+ create_container_template : creates an argument for the list of iter values
368
+ """
369
+ map_variable = self.executor.compose_map_variable(list_of_iter_values)
370
+ command = utils.get_node_execution_command(
371
+ self.node,
372
+ over_write_run_id=self.executor._run_id_placeholder,
373
+ map_variable=map_variable,
374
+ )
375
+
376
+ inputs = []
377
+ if list_of_iter_values:
378
+ for val in list_of_iter_values:
379
+ inputs.append(Parameter(name=val))
380
+
381
+ # Create the container template
382
+ container_template = self.executor.create_container_template(
383
+ working_on=self.node,
384
+ command=command,
385
+ inputs=inputs,
386
+ )
387
+
388
+ self.executor._container_templates.append(container_template)
389
+
390
+
391
+ class DagNodeRenderer(NodeRenderer):
392
+ allowed_node_types = ["dag"]
393
+
394
+ def render(self, list_of_iter_values: Optional[List] = None):
395
+ self.node = cast(DagNode, self.node)
396
+ task_template_arguments = []
397
+ dag_inputs = []
398
+ if list_of_iter_values:
399
+ for value in list_of_iter_values:
400
+ task_template_arguments.append(Argument(name=value, value="{{inputs.parameters." + value + "}}"))
401
+ dag_inputs.append(Parameter(name=value))
402
+
403
+ clean_name = self.executor.get_clean_name(self.node)
404
+ fan_out_template = self.executor._create_fan_out_template(
405
+ composite_node=self.node, list_of_iter_values=list_of_iter_values
406
+ )
407
+ fan_out_template.arguments = task_template_arguments if task_template_arguments else None
408
+
409
+ fan_in_template = self.executor._create_fan_in_template(
410
+ composite_node=self.node, list_of_iter_values=list_of_iter_values
411
+ )
412
+ fan_in_template.arguments = task_template_arguments if task_template_arguments else None
413
+
414
+ self.executor._gather_task_templates_of_dag(
415
+ self.node.branch,
416
+ dag_name=f"{clean_name}-branch",
417
+ list_of_iter_values=list_of_iter_values,
418
+ )
419
+
420
+ branch_template = DagTaskTemplate(
421
+ name=f"{clean_name}-branch",
422
+ template=f"{clean_name}-branch",
423
+ arguments=task_template_arguments if task_template_arguments else None,
424
+ )
425
+ branch_template.depends.append(f"{clean_name}-fan-out.Succeeded")
426
+ fan_in_template.depends.append(f"{clean_name}-branch.Succeeded")
427
+ fan_in_template.depends.append(f"{clean_name}-branch.Failed")
428
+
429
+ self.executor._dag_templates.append(
430
+ DagTemplate(
431
+ tasks=[fan_out_template, branch_template, fan_in_template],
432
+ name=clean_name,
433
+ inputs=dag_inputs if dag_inputs else None,
434
+ )
435
+ )
436
+
437
+
438
+ class ParallelNodeRender(NodeRenderer):
439
+ allowed_node_types = ["parallel"]
440
+
441
+ def render(self, list_of_iter_values: Optional[List] = None):
442
+ self.node = cast(ParallelNode, self.node)
443
+ task_template_arguments = []
444
+ dag_inputs = []
445
+ if list_of_iter_values:
446
+ for value in list_of_iter_values:
447
+ task_template_arguments.append(Argument(name=value, value="{{inputs.parameters." + value + "}}"))
448
+ dag_inputs.append(Parameter(name=value))
449
+
450
+ clean_name = self.executor.get_clean_name(self.node)
451
+ fan_out_template = self.executor._create_fan_out_template(
452
+ composite_node=self.node, list_of_iter_values=list_of_iter_values
453
+ )
454
+ fan_out_template.arguments = task_template_arguments if task_template_arguments else None
455
+
456
+ fan_in_template = self.executor._create_fan_in_template(
457
+ composite_node=self.node, list_of_iter_values=list_of_iter_values
458
+ )
459
+ fan_in_template.arguments = task_template_arguments if task_template_arguments else None
460
+
461
+ branch_templates = []
462
+ for name, branch in self.node.branches.items():
463
+ branch_name = self.executor.sanitize_name(name)
464
+ self.executor._gather_task_templates_of_dag(
465
+ branch,
466
+ dag_name=f"{clean_name}-{branch_name}",
467
+ list_of_iter_values=list_of_iter_values,
468
+ )
469
+ task_template = DagTaskTemplate(
470
+ name=f"{clean_name}-{branch_name}",
471
+ template=f"{clean_name}-{branch_name}",
472
+ arguments=task_template_arguments if task_template_arguments else None,
473
+ )
474
+ task_template.depends.append(f"{clean_name}-fan-out.Succeeded")
475
+ fan_in_template.depends.append(f"{task_template.name}.Succeeded")
476
+ fan_in_template.depends.append(f"{task_template.name}.Failed")
477
+ branch_templates.append(task_template)
478
+
479
+ executor_config = self.executor._resolve_executor_config(self.node)
480
+
481
+ self.executor._dag_templates.append(
482
+ DagTemplate(
483
+ tasks=[fan_out_template] + branch_templates + [fan_in_template],
484
+ name=clean_name,
485
+ inputs=dag_inputs if dag_inputs else None,
486
+ parallelism=executor_config.get("parallelism", None),
487
+ )
488
+ )
489
+
490
+
491
+ class MapNodeRender(NodeRenderer):
492
+ allowed_node_types = ["map"]
493
+
494
+ def render(self, list_of_iter_values: Optional[List] = None):
495
+ self.node = cast(MapNode, self.node)
496
+ task_template_arguments = []
497
+ dag_inputs = []
498
+ if list_of_iter_values:
499
+ for value in list_of_iter_values:
500
+ task_template_arguments.append(Argument(name=value, value="{{inputs.parameters." + value + "}}"))
501
+ dag_inputs.append(Parameter(name=value))
502
+
503
+ clean_name = self.executor.get_clean_name(self.node)
504
+ fan_out_template = self.executor._create_fan_out_template(
505
+ composite_node=self.node, list_of_iter_values=list_of_iter_values
506
+ )
507
+ fan_out_template.arguments = task_template_arguments if task_template_arguments else None
508
+
509
+ fan_in_template = self.executor._create_fan_in_template(
510
+ composite_node=self.node, list_of_iter_values=list_of_iter_values
511
+ )
512
+ fan_in_template.arguments = task_template_arguments if task_template_arguments else None
513
+
514
+ if not list_of_iter_values:
515
+ list_of_iter_values = []
516
+
517
+ list_of_iter_values.append(self.node.iterate_as)
518
+
519
+ self.executor._gather_task_templates_of_dag(
520
+ self.node.branch,
521
+ dag_name=f"{clean_name}-map",
522
+ list_of_iter_values=list_of_iter_values,
523
+ )
524
+
525
+ task_template = DagTaskTemplate(
526
+ name=f"{clean_name}-map",
527
+ template=f"{clean_name}-map",
528
+ arguments=task_template_arguments if task_template_arguments else None,
529
+ )
530
+ task_template.with_param = "{{tasks." + f"{clean_name}-fan-out" + ".outputs.parameters." + "iterate-on" + "}}"
531
+
532
+ argument = Argument(name=self.node.iterate_as, value="{{item}}")
533
+ if task_template.arguments is None:
534
+ task_template.arguments = []
535
+ task_template.arguments.append(argument)
536
+
537
+ task_template.depends.append(f"{clean_name}-fan-out.Succeeded")
538
+ fan_in_template.depends.append(f"{clean_name}-map.Succeeded")
539
+ fan_in_template.depends.append(f"{clean_name}-map.Failed")
540
+
541
+ executor_config = self.executor._resolve_executor_config(self.node)
542
+
543
+ self.executor._dag_templates.append(
544
+ DagTemplate(
545
+ tasks=[fan_out_template, task_template, fan_in_template],
546
+ name=clean_name,
547
+ inputs=dag_inputs if dag_inputs else None,
548
+ parallelism=executor_config.get("parallelism", None),
549
+ fail_fast=executor_config.get("fail_fast", True),
550
+ )
551
+ )
552
+
553
+
554
+ def get_renderer(node):
555
+ renderers = NodeRenderer.__subclasses__()
556
+
557
+ for renderer in renderers:
558
+ if node.node_type in renderer.allowed_node_types:
559
+ return renderer
560
+ raise Exception("This node type is not render-able")
561
+
562
+
563
+ class MetaData(BaseModel):
564
+ generate_name: str = Field(default="magnus-dag-", serialization_alias="generateName")
565
+ annotations: Optional[Dict[str, str]] = Field(default_factory=dict)
566
+ labels: Optional[Dict[str, str]] = Field(default_factory=dict)
567
+ namespace: Optional[str] = Field(default=None)
568
+
569
+
570
+ class Spec(BaseModel):
571
+ active_deadline_seconds: int = Field(serialization_alias="activeDeadlineSeconds")
572
+ entrypoint: str = Field(default="magnus-dag")
573
+ node_selector: Optional[Dict[str, str]] = Field(default_factory=dict, serialization_alias="nodeSelector")
574
+ tolerations: Optional[List[Toleration]] = Field(default=None, serialization_alias="tolerations")
575
+ parallelism: Optional[int] = Field(default=None, serialization_alias="parallelism")
576
+ # TODO: This has to be user driven
577
+ pod_gc: Dict[str, str] = Field(default={"strategy": "OnPodCompletion"}, serialization_alias="podGC")
578
+
579
+ retry_strategy: Retry = Field(default=Retry(), serialization_alias="retryStrategy")
580
+ service_account_name: Optional[str] = Field(default=None, serialization_alias="serviceAccountName")
581
+
582
+ templates: List[Union[DagTemplate, ContainerTemplate]] = Field(default_factory=list)
583
+ template_defaults: Optional[TemplateDefaults] = Field(default=None, serialization_alias="templateDefaults")
584
+
585
+ arguments: Optional[List[EnvVar]] = Field(default_factory=list)
586
+ persistent_volumes: List[UserVolumeMounts] = Field(default_factory=list, exclude=True)
587
+
588
+ @field_validator("parallelism")
589
+ @classmethod
590
+ def validate_parallelism(cls, parallelism: Optional[int]) -> Optional[int]:
591
+ if parallelism is not None and parallelism <= 0:
592
+ raise ValueError("Parallelism must be a positive integer greater than 0")
593
+ return parallelism
594
+
595
+ @computed_field # type: ignore
596
+ @property
597
+ def volumes(self) -> List[Volume]:
598
+ volumes: List[Volume] = []
599
+ claim_names = {}
600
+ for i, user_volume in enumerate(self.persistent_volumes):
601
+ if user_volume.name in claim_names:
602
+ raise Exception(f"Duplicate claim name {user_volume.name}")
603
+ claim_names[user_volume.name] = user_volume.name
604
+
605
+ volume = Volume(name=f"executor-{i}", claim=user_volume.name, mount_path=user_volume.mount_path)
606
+ volumes.append(volume)
607
+ return volumes
608
+
609
+ @field_serializer("arguments", when_used="unless-none")
610
+ def reshape_arguments(self, arguments: List[EnvVar], _info) -> Dict[str, List[EnvVar]]:
611
+ return {"parameters": arguments}
612
+
613
+
614
+ class Workflow(BaseModel):
615
+ api_version: str = Field(
616
+ default="argoproj.io/v1alpha1",
617
+ serialization_alias="apiVersion",
618
+ )
619
+ kind: str = "Workflow"
620
+ metadata: MetaData = Field(default=MetaData())
621
+ spec: Spec
622
+
623
+
624
+ class Override(BaseModel):
625
+ model_config = ConfigDict(extra="ignore")
626
+
627
+ image: str
628
+ tolerations: Optional[List[Toleration]] = Field(default=None)
629
+
630
+ max_step_duration_in_seconds: int = Field(
631
+ default=2 * 60 * 60, # 2 hours
632
+ gt=0,
633
+ )
634
+
635
+ node_selector: Optional[Dict[str, str]] = Field(
636
+ default=None,
637
+ serialization_alias="nodeSelector",
638
+ )
639
+
640
+ parallelism: Optional[int] = Field(
641
+ default=None,
642
+ serialization_alias="parallelism",
643
+ )
644
+
645
+ resources: Resources = Field(
646
+ default=Resources(),
647
+ serialization_alias="resources",
648
+ )
649
+
650
+ image_pull_policy: str = Field(default="")
651
+
652
+ retry_strategy: Retry = Field(
653
+ default=Retry(),
654
+ serialization_alias="retryStrategy",
655
+ description="Common across all templates",
656
+ )
657
+
658
+ @field_validator("parallelism")
659
+ @classmethod
660
+ def validate_parallelism(cls, parallelism: Optional[int]) -> Optional[int]:
661
+ if parallelism is not None and parallelism <= 0:
662
+ raise ValueError("Parallelism must be a positive integer greater than 0")
663
+ return parallelism
664
+
665
+
666
+ class ArgoExecutor(GenericExecutor):
667
+ service_name: str = "argo"
668
+
669
+ model_config = ConfigDict(extra="forbid")
670
+
671
+ image: str
672
+ expose_parameters_as_inputs: bool = True
673
+ secrets_from_k8s: List[SecretEnvVar] = Field(default_factory=list)
674
+ output_file: str = "argo-pipeline.yaml"
675
+
676
+ # Metadata related fields
677
+ name: str = Field(default="magnus-dag-", description="Used as an identifier for the workflow")
678
+ annotations: Dict[str, str] = Field(default_factory=dict)
679
+ labels: Dict[str, str] = Field(default_factory=dict)
680
+
681
+ max_workflow_duration_in_seconds: int = Field(
682
+ 2 * 24 * 60 * 60, # 2 days
683
+ serialization_alias="activeDeadlineSeconds",
684
+ gt=0,
685
+ )
686
+ node_selector: Optional[Dict[str, str]] = Field(
687
+ default=None,
688
+ serialization_alias="nodeSelector",
689
+ )
690
+ parallelism: Optional[int] = Field(
691
+ default=None,
692
+ serialization_alias="parallelism",
693
+ )
694
+ resources: Resources = Field(
695
+ default=Resources(),
696
+ serialization_alias="resources",
697
+ exclude=True,
698
+ )
699
+ retry_strategy: Retry = Field(
700
+ default=Retry(),
701
+ serialization_alias="retryStrategy",
702
+ description="Common across all templates",
703
+ )
704
+ max_step_duration_in_seconds: int = Field(
705
+ default=2 * 60 * 60, # 2 hours
706
+ gt=0,
707
+ )
708
+ tolerations: Optional[List[Toleration]] = Field(default=None)
709
+ image_pull_policy: str = Field(default="")
710
+ service_account_name: Optional[str] = None
711
+ persistent_volumes: List[UserVolumeMounts] = Field(default_factory=list)
712
+
713
+ _run_id_placeholder: str = "{{workflow.parameters.run_id}}"
714
+ _container_templates: List[ContainerTemplate] = []
715
+ _dag_templates: List[DagTemplate] = []
716
+ _clean_names: Dict[str, str] = {}
717
+ _container_volumes: List[ContainerVolume] = []
718
+
719
+ @field_validator("parallelism")
720
+ @classmethod
721
+ def validate_parallelism(cls, parallelism: Optional[int]) -> Optional[int]:
722
+ if parallelism is not None and parallelism <= 0:
723
+ raise ValueError("Parallelism must be a positive integer greater than 0")
724
+ return parallelism
725
+
726
+ @computed_field # type: ignore
727
+ @property
728
+ def step_timeout(self) -> int:
729
+ """
730
+ Maximum time the step can take to complete, including the pending state.
731
+ """
732
+ return self.max_step_duration_in_seconds + 2 * 60 * 60 # 2 hours + max_step_duration_in_seconds
733
+
734
+ @property
735
+ def metadata(self) -> MetaData:
736
+ return MetaData(
737
+ generate_name=self.name,
738
+ annotations=self.annotations,
739
+ labels=self.labels,
740
+ )
741
+
742
+ @property
743
+ def spec(self) -> Spec:
744
+ return Spec(
745
+ active_deadline_seconds=self.max_workflow_duration_in_seconds,
746
+ node_selector=self.node_selector,
747
+ tolerations=self.tolerations,
748
+ parallelism=self.parallelism,
749
+ retry_strategy=self.retry_strategy,
750
+ service_account_name=self.service_account_name,
751
+ persistent_volumes=self.persistent_volumes,
752
+ template_defaults=TemplateDefaults(max_step_duration=self.max_step_duration_in_seconds),
753
+ )
754
+
755
+ def prepare_for_graph_execution(self):
756
+ """
757
+ This method should be called prior to calling execute_graph.
758
+ Perform any steps required before doing the graph execution.
759
+
760
+ The most common implementation is to prepare a run log for the run if the run uses local interactive compute.
761
+
762
+ But in cases of actual rendering the job specs (eg: AWS step functions, K8's) we check if the services are OK.
763
+ We do not set up a run log as its not relevant.
764
+ """
765
+
766
+ integration.validate(self, self._context.run_log_store)
767
+ integration.configure_for_traversal(self, self._context.run_log_store)
768
+
769
+ integration.validate(self, self._context.catalog_handler)
770
+ integration.configure_for_traversal(self, self._context.catalog_handler)
771
+
772
+ integration.validate(self, self._context.secrets_handler)
773
+ integration.configure_for_traversal(self, self._context.secrets_handler)
774
+
775
+ integration.validate(self, self._context.experiment_tracker)
776
+ integration.configure_for_traversal(self, self._context.experiment_tracker)
777
+
778
+ def prepare_for_node_execution(self):
779
+ """
780
+ Perform any modifications to the services prior to execution of the node.
781
+
782
+ Args:
783
+ node (Node): [description]
784
+ map_variable (dict, optional): [description]. Defaults to None.
785
+ """
786
+
787
+ super().prepare_for_node_execution()
788
+ self._set_up_run_log(exists_ok=True)
789
+
790
+ def execute_node(self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs):
791
+ step_log = self._context.run_log_store.create_step_log(node.name, node._get_step_log_name(map_variable))
792
+
793
+ self.add_code_identities(node=node, step_log=step_log)
794
+
795
+ step_log.step_type = node.node_type
796
+ step_log.status = defaults.PROCESSING
797
+ self._context.run_log_store.add_step_log(step_log, self._context.run_id)
798
+
799
+ super()._execute_node(node, map_variable=map_variable, **kwargs)
800
+
801
+ # Implicit fail
802
+ if self._context.dag:
803
+ # functions and notebooks do not have dags
804
+ _, current_branch = search_node_by_internal_name(dag=self._context.dag, internal_name=node.internal_name)
805
+ _, next_node_name = self._get_status_and_next_node_name(node, current_branch, map_variable=map_variable)
806
+ if next_node_name:
807
+ # Terminal nodes do not have next node name
808
+ next_node = current_branch.get_node_by_name(next_node_name)
809
+
810
+ if next_node.node_type == defaults.FAIL:
811
+ self.execute_node(next_node, map_variable=map_variable)
812
+
813
+ step_log = self._context.run_log_store.get_step_log(node._get_step_log_name(map_variable), self._context.run_id)
814
+ if step_log.status == defaults.FAIL:
815
+ raise Exception(f"Step {node.name} failed")
816
+
817
+ def fan_out(self, node: BaseNode, map_variable: TypeMapVariable = None):
818
+ super().fan_out(node, map_variable)
819
+
820
+ # If its a map node, write the list values to "/tmp/output.txt"
821
+ if node.node_type == "map":
822
+ node = cast(MapNode, node)
823
+ iterate_on = self._context.run_log_store.get_parameters(self._context.run_id)[node.iterate_on]
824
+
825
+ with open("/tmp/output.txt", mode="w", encoding="utf-8") as myfile:
826
+ json.dump(iterate_on, myfile, indent=4)
827
+
828
+ def _get_parameters(self) -> Dict[str, Any]:
829
+ params = {}
830
+ if self._context.parameters_file:
831
+ # Parameters from the parameters file if defined
832
+ params.update(utils.load_yaml(self._context.parameters_file))
833
+ # parameters from environment variables supersede file based
834
+ params.update(parameters.get_user_set_parameters())
835
+
836
+ return params
837
+
838
+ def sanitize_name(self, name):
839
+ return name.replace(" ", "-").replace(".", "-").replace("_", "-")
840
+
841
+ def get_clean_name(self, node: BaseNode):
842
+ # Cache names for the node
843
+ if node.internal_name not in self._clean_names:
844
+ sanitized = self.sanitize_name(node.name)
845
+ tag = "".join(random.choices(string.ascii_lowercase + string.digits, k=6))
846
+ self._clean_names[node.internal_name] = f"{sanitized}-{node.node_type}-{tag}"
847
+
848
+ return self._clean_names[node.internal_name]
849
+
850
+ def compose_map_variable(self, list_of_iter_values: Optional[List] = None) -> TypeMapVariable:
851
+ map_variable = OrderedDict()
852
+
853
+ # If we are inside a map node, compose a map_variable
854
+ # The values of "iterate_as" are sent over as inputs to the container template
855
+ if list_of_iter_values:
856
+ for var in list_of_iter_values:
857
+ map_variable[var] = "{{inputs.parameters." + str(var) + "}}"
858
+
859
+ return map_variable # type: ignore
860
+
861
+ def create_container_template(
862
+ self,
863
+ working_on: BaseNode,
864
+ command: str,
865
+ inputs: Optional[List] = None,
866
+ outputs: Optional[List] = None,
867
+ overwrite_name: str = "",
868
+ ):
869
+ effective_node_config = self._resolve_executor_config(working_on)
870
+
871
+ override: Override = Override(**effective_node_config)
872
+
873
+ container = Container(
874
+ command=command,
875
+ image=override.image,
876
+ volume_mounts=self._container_volumes,
877
+ image_pull_policy=override.image_pull_policy,
878
+ resources=override.resources,
879
+ secrets_from_k8s=self.secrets_from_k8s,
880
+ )
881
+
882
+ if working_on.name == self._context.dag.start_at and self.expose_parameters_as_inputs:
883
+ for key, value in self._get_parameters().items():
884
+ # Get the value from work flow parameters for dynamic behavior
885
+ if isinstance(value, int) or isinstance(value, float) or isinstance(value, str):
886
+ env_var = EnvVar(
887
+ name=defaults.PARAMETER_PREFIX + key,
888
+ value="{{workflow.parameters." + key + "}}",
889
+ )
890
+ container.env_vars.append(env_var)
891
+
892
+ clean_name = self.get_clean_name(working_on)
893
+ if overwrite_name:
894
+ clean_name = overwrite_name
895
+
896
+ container_template = ContainerTemplate(
897
+ name=clean_name,
898
+ active_deadline_seconds=(
899
+ override.max_step_duration_in_seconds
900
+ if self.max_step_duration_in_seconds != override.max_step_duration_in_seconds
901
+ else None
902
+ ),
903
+ container=container,
904
+ retry_strategy=override.retry_strategy if self.retry_strategy != override.retry_strategy else None,
905
+ tolerations=override.tolerations if self.tolerations != override.tolerations else None,
906
+ node_selector=override.node_selector if self.node_selector != override.node_selector else None,
907
+ )
908
+
909
+ # inputs are the "iterate_as" value map variables in the same order as they are observed
910
+ # We need to expose the map variables in the command of the container
911
+ if inputs:
912
+ if not container_template.inputs:
913
+ container_template.inputs = []
914
+ container_template.inputs.extend(inputs)
915
+
916
+ # The map step fan out would create an output that we should propagate via Argo
917
+ if outputs:
918
+ if not container_template.outputs:
919
+ container_template.outputs = []
920
+ container_template.outputs.extend(outputs)
921
+
922
+ return container_template
923
+
924
+ def _create_fan_out_template(self, composite_node, list_of_iter_values: Optional[List] = None):
925
+ clean_name = self.get_clean_name(composite_node)
926
+ inputs = []
927
+ # If we are fanning out already map state, we need to send the map variable inside
928
+ # The container template also should be accepting an input parameter
929
+ map_variable = None
930
+ if list_of_iter_values:
931
+ map_variable = self.compose_map_variable(list_of_iter_values=list_of_iter_values)
932
+
933
+ for val in list_of_iter_values:
934
+ inputs.append(Parameter(name=val))
935
+
936
+ command = utils.get_fan_command(
937
+ mode="out",
938
+ node=composite_node,
939
+ run_id=self._run_id_placeholder,
940
+ map_variable=map_variable,
941
+ )
942
+
943
+ outputs = []
944
+ # If the node is a map node, we have to set the output parameters
945
+ # Output is always the step's internal name + iterate-on
946
+ if composite_node.node_type == "map":
947
+ output_parameter = OutputParameter(name="iterate-on")
948
+ outputs.append(output_parameter)
949
+
950
+ # Create the node now
951
+ step_config = {"command": command, "type": "task", "next": "dummy"}
952
+ node = create_node(name=f"{clean_name}-fan-out", step_config=step_config)
953
+
954
+ container_template = self.create_container_template(
955
+ working_on=node,
956
+ command=command,
957
+ outputs=outputs,
958
+ inputs=inputs,
959
+ overwrite_name=f"{clean_name}-fan-out",
960
+ )
961
+
962
+ self._container_templates.append(container_template)
963
+ return DagTaskTemplate(name=f"{clean_name}-fan-out", template=f"{clean_name}-fan-out")
964
+
965
+ def _create_fan_in_template(self, composite_node, list_of_iter_values: Optional[List] = None):
966
+ clean_name = self.get_clean_name(composite_node)
967
+ inputs = []
968
+ # If we are fanning in already map state, we need to send the map variable inside
969
+ # The container template also should be accepting an input parameter
970
+ map_variable = None
971
+ if list_of_iter_values:
972
+ map_variable = self.compose_map_variable(list_of_iter_values=list_of_iter_values)
973
+
974
+ for val in list_of_iter_values:
975
+ inputs.append(Parameter(name=val))
976
+
977
+ command = utils.get_fan_command(
978
+ mode="in",
979
+ node=composite_node,
980
+ run_id=self._run_id_placeholder,
981
+ map_variable=map_variable,
982
+ )
983
+
984
+ step_config = {"command": command, "type": "task", "next": "dummy"}
985
+ node = create_node(name=f"{clean_name}-fan-in", step_config=step_config)
986
+ container_template = self.create_container_template(
987
+ working_on=node,
988
+ command=command,
989
+ inputs=inputs,
990
+ overwrite_name=f"{clean_name}-fan-in",
991
+ )
992
+ self._container_templates.append(container_template)
993
+ clean_name = self.get_clean_name(composite_node)
994
+ return DagTaskTemplate(name=f"{clean_name}-fan-in", template=f"{clean_name}-fan-in")
995
+
996
+ def _gather_task_templates_of_dag(
997
+ self, dag: Graph, dag_name="magnus-dag", list_of_iter_values: Optional[List] = None
998
+ ):
999
+ current_node = dag.start_at
1000
+ previous_node = None
1001
+ previous_node_template_name = None
1002
+
1003
+ templates: Dict[str, DagTaskTemplate] = {}
1004
+
1005
+ if not list_of_iter_values:
1006
+ list_of_iter_values = []
1007
+
1008
+ while True:
1009
+ working_on = dag.get_node_by_name(current_node)
1010
+ if previous_node == current_node:
1011
+ raise Exception("Potentially running in a infinite loop")
1012
+
1013
+ render_obj = get_renderer(working_on)(executor=self, node=working_on)
1014
+ render_obj.render(list_of_iter_values=list_of_iter_values.copy())
1015
+
1016
+ clean_name = self.get_clean_name(working_on)
1017
+
1018
+ # If a task template for clean name exists, retrieve it (could have been created by on_failure)
1019
+ template = templates.get(clean_name, DagTaskTemplate(name=clean_name, template=clean_name))
1020
+
1021
+ # Link the current node to previous node, if the previous node was successful.
1022
+ if previous_node:
1023
+ template.depends.append(f"{previous_node_template_name}.Succeeded")
1024
+
1025
+ templates[clean_name] = template
1026
+
1027
+ # On failure nodes
1028
+ if working_on.node_type not in ["success", "fail"] and working_on._get_on_failure_node():
1029
+ failure_node = dag.get_node_by_name(working_on._get_on_failure_node())
1030
+
1031
+ failure_template_name = self.get_clean_name(failure_node)
1032
+ # If a task template for clean name exists, retrieve it
1033
+ failure_template = templates.get(
1034
+ failure_template_name,
1035
+ DagTaskTemplate(name=failure_template_name, template=failure_template_name),
1036
+ )
1037
+ failure_template.depends.append(f"{clean_name}.Failed")
1038
+
1039
+ templates[failure_template_name] = failure_template
1040
+
1041
+ # If we are in a map node, we need to add the values as arguments
1042
+ template = templates[clean_name]
1043
+ if list_of_iter_values:
1044
+ if not template.arguments:
1045
+ template.arguments = []
1046
+ for value in list_of_iter_values:
1047
+ template.arguments.append(Argument(name=value, value="{{inputs.parameters." + value + "}}"))
1048
+
1049
+ # Move ahead to the next node
1050
+ previous_node = current_node
1051
+ previous_node_template_name = self.get_clean_name(working_on)
1052
+
1053
+ if working_on.node_type in ["success", "fail"]:
1054
+ break
1055
+
1056
+ current_node = working_on._get_next_node()
1057
+
1058
+ # Add the iteration values as input to dag template
1059
+ dag_template = DagTemplate(tasks=list(templates.values()), name=dag_name)
1060
+ if list_of_iter_values:
1061
+ if not dag_template.inputs:
1062
+ dag_template.inputs = []
1063
+ dag_template.inputs.extend([Parameter(name=val) for val in list_of_iter_values])
1064
+
1065
+ # Add the dag template to the list of templates
1066
+ self._dag_templates.append(dag_template)
1067
+
1068
+ def _get_template_defaults(self) -> TemplateDefaults:
1069
+ user_provided_config = self.model_dump(by_alias=False)
1070
+
1071
+ return TemplateDefaults(**user_provided_config)
1072
+
1073
+ def execute_graph(self, dag: Graph, map_variable: Optional[dict] = None, **kwargs):
1074
+ # TODO: Add metadata
1075
+ arguments = []
1076
+ # Expose "simple" parameters as workflow arguments for dynamic behavior
1077
+ if self.expose_parameters_as_inputs:
1078
+ for key, value in self._get_parameters().items():
1079
+ if isinstance(value, dict) or isinstance(value, list):
1080
+ continue
1081
+ env_var = EnvVar(name=key, value=value)
1082
+ arguments.append(env_var)
1083
+
1084
+ run_id_var = EnvVar(name="run_id", value="{{workflow.uid}}")
1085
+ arguments.append(run_id_var)
1086
+
1087
+ # # TODO: Experimental feature
1088
+
1089
+ # original_run_id_var = EnvVar(name="original_run_id")
1090
+ # arguments.append(original_run_id_var)
1091
+
1092
+ for volume in self.spec.volumes:
1093
+ self._container_volumes.append(ContainerVolume(name=volume.name, mount_path=volume.mount_path))
1094
+
1095
+ # Container specifications are globally collected and added at the end.
1096
+ # Dag specifications are added as part of the dag traversal.
1097
+ templates: List[Union[DagTemplate, ContainerTemplate]] = []
1098
+ self._gather_task_templates_of_dag(dag=dag, list_of_iter_values=[])
1099
+ templates.extend(self._dag_templates)
1100
+ templates.extend(self._container_templates)
1101
+
1102
+ spec = self.spec
1103
+ spec.templates = templates
1104
+ spec.arguments = arguments
1105
+ workflow = Workflow(metadata=self.metadata, spec=spec)
1106
+
1107
+ yaml = YAML()
1108
+ with open(self.output_file, "w") as f:
1109
+ yaml.indent(mapping=2, sequence=4, offset=2)
1110
+
1111
+ yaml.dump(workflow.model_dump(by_alias=True, exclude_none=True), f)
1112
+
1113
+ def execute_job(self, node: BaseNode):
1114
+ """
1115
+ Use K8's job instead
1116
+ """
1117
+ raise NotImplementedError("Use K8's job instead")
1118
+
1119
+ def send_return_code(self, stage="traversal"):
1120
+ """
1121
+ Convenience function used by pipeline to send return code to the caller of the cli
1122
+
1123
+ Raises:
1124
+ Exception: If the pipeline execution failed
1125
+ """
1126
+ if stage != "traversal": # traversal does no actual execution, so return code is pointless
1127
+ run_id = self._context.run_id
1128
+
1129
+ run_log = self._context.run_log_store.get_run_log_by_id(run_id=run_id, full=False)
1130
+ if run_log.status == defaults.FAIL:
1131
+ raise exceptions.ExecutionFailedError(run_id)
1132
+
1133
+
1134
+ class FileSystemRunLogStore(BaseIntegration):
1135
+ """
1136
+ Only local execution mode is possible for Buffered Run Log store
1137
+ """
1138
+
1139
+ executor_type = "argo"
1140
+ service_type = "run_log_store" # One of secret, catalog, datastore
1141
+ service_provider = "file-system" # The actual implementation of the service
1142
+
1143
+ def validate(self, **kwargs):
1144
+ msg = (
1145
+ "Argo cannot run work with file-system run log store. "
1146
+ "Unless you have made a mechanism to use volume mounts."
1147
+ "Using this run log store if the pipeline has concurrent tasks might lead to unexpected results"
1148
+ )
1149
+ logger.warning(msg)
1150
+
1151
+
1152
+ class ChunkedFileSystemRunLogStore(BaseIntegration):
1153
+ """
1154
+ Only local execution mode is possible for Buffered Run Log store
1155
+ """
1156
+
1157
+ executor_type = "argo"
1158
+ service_type = "run_log_store" # One of secret, catalog, datastore
1159
+ service_provider = "chunked-fs" # The actual implementation of the service
1160
+
1161
+ def validate(self, **kwargs):
1162
+ msg = (
1163
+ "Argo cannot run work with chunked file-system run log store. "
1164
+ "Unless you have made a mechanism to use volume mounts"
1165
+ )
1166
+ logger.warning(msg)
1167
+
1168
+
1169
+ class FileSystemCatalog(BaseIntegration):
1170
+ """
1171
+ Only local execution mode is possible for Buffered Run Log store
1172
+ """
1173
+
1174
+ executor_type = "argo"
1175
+ service_type = "catalog" # One of secret, catalog, datastore
1176
+ service_provider = "file-system" # The actual implementation of the service
1177
+
1178
+ def validate(self, **kwargs):
1179
+ msg = (
1180
+ "Argo cannot run work with file-system run log store. Unless you have made a mechanism to use volume mounts"
1181
+ )
1182
+ logger.warning(msg)