lightweight-dag 0.1.1__tar.gz
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.
- lightweight_dag-0.1.1/PKG-INFO +393 -0
- lightweight_dag-0.1.1/README.md +382 -0
- lightweight_dag-0.1.1/pyproject.toml +36 -0
- lightweight_dag-0.1.1/pyproject.toml.orig +35 -0
- lightweight_dag-0.1.1/src/lightweight_dag/__init__.py +13 -0
- lightweight_dag-0.1.1/src/lightweight_dag/dag.py +306 -0
- lightweight_dag-0.1.1/src/lightweight_dag/node_registry.py +83 -0
- lightweight_dag-0.1.1/src/lightweight_dag/suite.py +80 -0
|
@@ -0,0 +1,393 @@
|
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
|
+
Name: lightweight-dag
|
|
3
|
+
Version: 0.1.1
|
|
4
|
+
Summary: A lightweight Python library for defining and executing DAG workflows
|
|
5
|
+
Author: Lorenzo Gardini
|
|
6
|
+
Author-email: Lorenzo Gardini <l.gardo98@gmail.com>
|
|
7
|
+
Requires-Dist: pydantic>=2.13.5
|
|
8
|
+
Requires-Dist: pyyaml>=6.0.3
|
|
9
|
+
Requires-Python: >=3.11
|
|
10
|
+
Description-Content-Type: text/markdown
|
|
11
|
+
|
|
12
|
+
# lightweight-dag
|
|
13
|
+
|
|
14
|
+
Python library for defining, validating, and executing directed acyclic graphs
|
|
15
|
+
(DAGs). Operations are registered Python functions. DAGs can be loaded from
|
|
16
|
+
YAML or dictionaries, grouped into named suites, and executed for only the
|
|
17
|
+
requested outputs.
|
|
18
|
+
|
|
19
|
+
Requires Python 3.11 or later.
|
|
20
|
+
|
|
21
|
+
## Installation
|
|
22
|
+
|
|
23
|
+
Install the published package from PyPI with uv:
|
|
24
|
+
|
|
25
|
+
```bash
|
|
26
|
+
uv add lightweight-dag
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
To install the development and test dependencies when working from a clone of
|
|
30
|
+
the repository:
|
|
31
|
+
|
|
32
|
+
```bash
|
|
33
|
+
uv sync --group dev
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
The public API is available from the package root:
|
|
37
|
+
|
|
38
|
+
```python
|
|
39
|
+
from lightweight_dag import Dag, DagExecutionResult, DagSuite, Node, dag_node
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
## Registering operations
|
|
43
|
+
|
|
44
|
+
Operations must be synchronous functions defined at module level. Passing one
|
|
45
|
+
argument to the decorator registers a global operation:
|
|
46
|
+
|
|
47
|
+
```python
|
|
48
|
+
from lightweight_dag import dag_node
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
@dag_node("normalize")
|
|
52
|
+
def normalize(values):
|
|
53
|
+
return [value.strip().lower() for value in values]
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
Passing two arguments registers an operation for a specific DAG:
|
|
57
|
+
|
|
58
|
+
```python
|
|
59
|
+
from lightweight_dag import dag_node
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
@dag_node("example_report", "filter_orders")
|
|
63
|
+
def filter_monthly_orders(orders, minimum_total):
|
|
64
|
+
return [order for order in orders if order["total"] >= minimum_total]
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
A DAG without `name` looks up only global operations. A named DAG looks up only
|
|
68
|
+
operations registered for that name; there is no fallback to the global scope.
|
|
69
|
+
The same operation can therefore have different implementations in different
|
|
70
|
+
scopes. Registering an operation twice in the same scope is an error.
|
|
71
|
+
|
|
72
|
+
## Defining and running a suite
|
|
73
|
+
|
|
74
|
+
A YAML suite maps DAG names to DAG definitions. Each `Node.inputs` entry must
|
|
75
|
+
refer to another node in the same DAG. Source nodes have no inputs and produce
|
|
76
|
+
their values through their registered operation. The block-style suite example
|
|
77
|
+
in [examples/dags.yaml](examples/dags.yaml) defines both `monthly_report` and
|
|
78
|
+
`sales_audit`; the following code loads and executes both DAGs.
|
|
79
|
+
|
|
80
|
+
Register each operation in the scope of the DAG that uses it. The same
|
|
81
|
+
operation ID can have different implementations in different DAG scopes:
|
|
82
|
+
|
|
83
|
+
```python
|
|
84
|
+
from pathlib import Path
|
|
85
|
+
|
|
86
|
+
from lightweight_dag import DagSuite, dag_node
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
@dag_node("monthly_report", "load_orders")
|
|
90
|
+
def load_monthly_orders():
|
|
91
|
+
return [{"total": 75}, {"total": 35}]
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
@dag_node("monthly_report", "filter_orders")
|
|
95
|
+
def filter_monthly_orders(orders, minimum_total):
|
|
96
|
+
return [order for order in orders if order["total"] >= minimum_total]
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
@dag_node("monthly_report", "summarize_orders")
|
|
100
|
+
def summarize_monthly_orders(eligible_orders):
|
|
101
|
+
return {"count": len(eligible_orders)}
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
@dag_node("sales_audit", "load_orders")
|
|
105
|
+
def load_audit_orders():
|
|
106
|
+
return [{"total": 75}, {"total": 35}]
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
@dag_node("sales_audit", "calculate_total")
|
|
110
|
+
def calculate_total(orders):
|
|
111
|
+
return sum(order["total"] for order in orders)
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
@dag_node("sales_audit", "classify_total")
|
|
115
|
+
def classify_total(total, review_threshold):
|
|
116
|
+
return "review" if total >= review_threshold else "ok"
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
suite = DagSuite.from_yaml(Path("examples/dags.yaml").read_text())
|
|
120
|
+
monthly_report = suite["monthly_report"]
|
|
121
|
+
sales_audit = suite["sales_audit"]
|
|
122
|
+
|
|
123
|
+
report_result = monthly_report.run(outputs=["summary"])
|
|
124
|
+
audit_result = sales_audit.run(outputs=["audit_result"])
|
|
125
|
+
|
|
126
|
+
assert report_result.results["summary"] == {"count": 1}
|
|
127
|
+
assert audit_result.results["audit_result"] == "review"
|
|
128
|
+
```
|
|
129
|
+
|
|
130
|
+
The operation registry is scoped: `monthly_report` cannot use operations
|
|
131
|
+
registered only for `sales_audit`, and named DAGs do not fall back to global
|
|
132
|
+
operations. The YAML key supplies the DAG name; an explicit `name` field is
|
|
133
|
+
optional, but if present it must match the key. `suite[name]` raises `KeyError`
|
|
134
|
+
for an unknown name.
|
|
135
|
+
|
|
136
|
+
`DagSuite.from_dict(...)` accepts the same name-to-definition mapping. For a
|
|
137
|
+
single DAG, use `Dag.from_yaml(...)` or `Dag.from_dict(...)`. This standalone
|
|
138
|
+
example builds a DAG from a dictionary and executes its output:
|
|
139
|
+
|
|
140
|
+
```python
|
|
141
|
+
from lightweight_dag import Dag, dag_node
|
|
142
|
+
|
|
143
|
+
|
|
144
|
+
@dag_node("make_number")
|
|
145
|
+
def make_number():
|
|
146
|
+
return 21
|
|
147
|
+
|
|
148
|
+
|
|
149
|
+
@dag_node("double_number")
|
|
150
|
+
def double_number(number):
|
|
151
|
+
return number * 2
|
|
152
|
+
|
|
153
|
+
|
|
154
|
+
dag = Dag.from_dict(
|
|
155
|
+
{
|
|
156
|
+
"nodes": [
|
|
157
|
+
{"node_id": "number", "operation": "make_number"},
|
|
158
|
+
{
|
|
159
|
+
"node_id": "answer",
|
|
160
|
+
"operation": "double_number",
|
|
161
|
+
"inputs": ["number"],
|
|
162
|
+
},
|
|
163
|
+
]
|
|
164
|
+
}
|
|
165
|
+
)
|
|
166
|
+
result = dag.run(outputs=["answer"])
|
|
167
|
+
assert result.results["answer"] == 42
|
|
168
|
+
```
|
|
169
|
+
|
|
170
|
+
To load a single DAG from YAML, register its global operations and pass the
|
|
171
|
+
YAML document to `Dag.from_yaml(...)`:
|
|
172
|
+
|
|
173
|
+
```python
|
|
174
|
+
from lightweight_dag import Dag, dag_node
|
|
175
|
+
|
|
176
|
+
|
|
177
|
+
@dag_node("yaml_source")
|
|
178
|
+
def yaml_source(): return 21
|
|
179
|
+
|
|
180
|
+
|
|
181
|
+
@dag_node("yaml_double")
|
|
182
|
+
def yaml_double(number): return number * 2
|
|
183
|
+
|
|
184
|
+
|
|
185
|
+
dag = Dag.from_yaml("""{nodes: [
|
|
186
|
+
{node_id: number, operation: yaml_source},
|
|
187
|
+
{node_id: answer, operation: yaml_double, inputs: [number]}
|
|
188
|
+
]}""")
|
|
189
|
+
assert dag.run(outputs=["answer"]).results["answer"] == 42
|
|
190
|
+
```
|
|
191
|
+
|
|
192
|
+
## Inputs and execution
|
|
193
|
+
|
|
194
|
+
Call `Dag.run()` with no arguments to execute every node. To execute only
|
|
195
|
+
selected outputs and their dependencies, pass their IDs in `outputs`. Each
|
|
196
|
+
operation receives its predecessor node results and the YAML `params` as
|
|
197
|
+
keyword arguments. Argument names correspond to predecessor IDs or parameter
|
|
198
|
+
keys.
|
|
199
|
+
|
|
200
|
+
```python
|
|
201
|
+
all_results = monthly_report.run()
|
|
202
|
+
assert set(all_results.results) == {"orders", "eligible_orders", "summary"}
|
|
203
|
+
|
|
204
|
+
summary_result = monthly_report.run(outputs=["summary"])
|
|
205
|
+
assert summary_result.results["summary"] == {"count": 1}
|
|
206
|
+
|
|
207
|
+
if summary_result.status:
|
|
208
|
+
print(summary_result.results["summary"])
|
|
209
|
+
```
|
|
210
|
+
|
|
211
|
+
An unknown output raises `ValueError`. The result contains `status` and
|
|
212
|
+
`results`; `results` contains the computed node values. With `outputs=[]`, no
|
|
213
|
+
nodes are executed and `results` is empty. Omitting `outputs` is different:
|
|
214
|
+
all nodes are selected, including independent branches.
|
|
215
|
+
|
|
216
|
+
Shared predecessors are computed only once, and nodes outside the requested
|
|
217
|
+
paths are not executed. In this graph, requesting only `d` computes `a`, `b`,
|
|
218
|
+
and `d`; `c` is not executed:
|
|
219
|
+
|
|
220
|
+
```mermaid
|
|
221
|
+
flowchart TD
|
|
222
|
+
a --> b --> c
|
|
223
|
+
b --> d
|
|
224
|
+
```
|
|
225
|
+
|
|
226
|
+
```python
|
|
227
|
+
from lightweight_dag import Dag, dag_node
|
|
228
|
+
|
|
229
|
+
|
|
230
|
+
@dag_node("get_a")
|
|
231
|
+
def get_a():
|
|
232
|
+
return "a"
|
|
233
|
+
|
|
234
|
+
|
|
235
|
+
@dag_node("make_b")
|
|
236
|
+
def make_b(a):
|
|
237
|
+
return f"{a}->b"
|
|
238
|
+
|
|
239
|
+
|
|
240
|
+
@dag_node("make_c")
|
|
241
|
+
def make_c(b):
|
|
242
|
+
return f"{b}->c"
|
|
243
|
+
|
|
244
|
+
|
|
245
|
+
@dag_node("make_d")
|
|
246
|
+
def make_d(b):
|
|
247
|
+
return f"{b}->d"
|
|
248
|
+
|
|
249
|
+
|
|
250
|
+
dag = Dag.from_dict({"nodes": [
|
|
251
|
+
{"node_id": "a", "operation": "get_a"},
|
|
252
|
+
{"node_id": "b", "operation": "make_b", "inputs": ["a"]},
|
|
253
|
+
{"node_id": "c", "operation": "make_c", "inputs": ["b"]},
|
|
254
|
+
{"node_id": "d", "operation": "make_d", "inputs": ["b"]},
|
|
255
|
+
]})
|
|
256
|
+
result = dag.run(outputs=["d"])
|
|
257
|
+
|
|
258
|
+
assert result.results == {
|
|
259
|
+
"a": "a",
|
|
260
|
+
"b": "a->b",
|
|
261
|
+
"d": "a->b->d",
|
|
262
|
+
}
|
|
263
|
+
assert "c" not in result.results
|
|
264
|
+
```
|
|
265
|
+
|
|
266
|
+
### Using an LLM with LangChain
|
|
267
|
+
|
|
268
|
+
This example simulates reading a customer record from a database, then sends
|
|
269
|
+
that context and a user prompt to an LLM through LangChain. Install the optional
|
|
270
|
+
integration and set an API key before running it:
|
|
271
|
+
|
|
272
|
+
```bash
|
|
273
|
+
pip install langchain-openai
|
|
274
|
+
export OPENAI_API_KEY="your-api-key"
|
|
275
|
+
```
|
|
276
|
+
|
|
277
|
+
```python
|
|
278
|
+
from lightweight_dag import Dag, dag_node
|
|
279
|
+
from langchain_openai import ChatOpenAI
|
|
280
|
+
from langchain_core.prompts import ChatPromptTemplate
|
|
281
|
+
|
|
282
|
+
|
|
283
|
+
@dag_node("load_customer")
|
|
284
|
+
def load_customer():
|
|
285
|
+
# Replace this in-memory record with a database query.
|
|
286
|
+
return {"name": "Ada", "plan": "Pro", "last_order": "Notebook"}
|
|
287
|
+
|
|
288
|
+
|
|
289
|
+
@dag_node("answer_user")
|
|
290
|
+
def answer_user(customer_data, user_prompt):
|
|
291
|
+
prompt = ChatPromptTemplate.from_messages([
|
|
292
|
+
("system", "Answer using the customer data provided: {customer_data}"),
|
|
293
|
+
("human", "{user_prompt}"),
|
|
294
|
+
])
|
|
295
|
+
chain = prompt | ChatOpenAI(model="gpt-4o-mini")
|
|
296
|
+
response = chain.invoke({
|
|
297
|
+
"customer_data": customer_data,
|
|
298
|
+
"user_prompt": user_prompt,
|
|
299
|
+
})
|
|
300
|
+
return response.content
|
|
301
|
+
|
|
302
|
+
|
|
303
|
+
support_dag = Dag.from_dict({"nodes": [
|
|
304
|
+
{"node_id": "customer_data", "operation": "load_customer"},
|
|
305
|
+
{
|
|
306
|
+
"node_id": "answer",
|
|
307
|
+
"operation": "answer_user",
|
|
308
|
+
"inputs": ["customer_data"],
|
|
309
|
+
"params": {"user_prompt": "Qual è il mio ultimo ordine?"},
|
|
310
|
+
},
|
|
311
|
+
]})
|
|
312
|
+
result = support_dag.run(outputs=["answer"])
|
|
313
|
+
print(result.results["answer"])
|
|
314
|
+
```
|
|
315
|
+
|
|
316
|
+
```python
|
|
317
|
+
from lightweight_dag import Dag, dag_node
|
|
318
|
+
|
|
319
|
+
|
|
320
|
+
@dag_node("branching_demo", "label")
|
|
321
|
+
def label_node(label, **predecessors): return label
|
|
322
|
+
|
|
323
|
+
|
|
324
|
+
branching_dag = Dag.from_dict({"name": "branching_demo", "nodes": [
|
|
325
|
+
{"node_id": "A", "operation": "label", "params": {"label": "A"}},
|
|
326
|
+
{"node_id": "B", "operation": "label", "inputs": ["A"], "params": {"label": "B"}},
|
|
327
|
+
{"node_id": "E", "operation": "label", "inputs": ["B"], "params": {"label": "E"}},
|
|
328
|
+
{"node_id": "G", "operation": "label", "inputs": ["E"], "params": {"label": "G"}},
|
|
329
|
+
{"node_id": "H", "operation": "label", "inputs": ["E"], "params": {"label": "H"}},
|
|
330
|
+
{"node_id": "F", "operation": "label", "inputs": ["B"], "params": {"label": "F"}},
|
|
331
|
+
{"node_id": "I", "operation": "label", "inputs": ["F"], "params": {"label": "I"}},
|
|
332
|
+
{"node_id": "J", "operation": "label", "inputs": ["F"], "params": {"label": "J"}},
|
|
333
|
+
{"node_id": "K", "operation": "label", "inputs": ["G", "H"], "params": {"label": "K"}},
|
|
334
|
+
{"node_id": "N", "operation": "label", "inputs": ["I", "J", "G"], "params": {"label": "N"}},
|
|
335
|
+
]})
|
|
336
|
+
result = branching_dag.run(outputs=["G", "J", "A"])
|
|
337
|
+
assert set(result.results) == {"A", "B", "E", "F", "G", "J"}
|
|
338
|
+
```
|
|
339
|
+
|
|
340
|
+
If an operation raises an exception, its result and the results of its
|
|
341
|
+
dependent nodes become `None`, and `status` becomes `False`. Requested,
|
|
342
|
+
independent branches continue to execute. This example demonstrates the
|
|
343
|
+
failure result:
|
|
344
|
+
|
|
345
|
+
```python
|
|
346
|
+
from lightweight_dag import Dag, dag_node
|
|
347
|
+
|
|
348
|
+
|
|
349
|
+
@dag_node("always_fail")
|
|
350
|
+
def always_fail(): raise RuntimeError("operation failed")
|
|
351
|
+
|
|
352
|
+
|
|
353
|
+
@dag_node("consume_result")
|
|
354
|
+
def consume_result(value): return value
|
|
355
|
+
|
|
356
|
+
|
|
357
|
+
failure_dag = Dag.from_dict({"nodes": [
|
|
358
|
+
{"node_id": "failed", "operation": "always_fail"},
|
|
359
|
+
{"node_id": "dependent", "operation": "consume_result", "inputs": ["failed"]},
|
|
360
|
+
]})
|
|
361
|
+
failure_result = failure_dag.run(outputs=["dependent"])
|
|
362
|
+
assert failure_result.status is False
|
|
363
|
+
assert failure_result.results["failed"] is None
|
|
364
|
+
assert failure_result.results["dependent"] is None
|
|
365
|
+
```
|
|
366
|
+
|
|
367
|
+
## Validation and security
|
|
368
|
+
|
|
369
|
+
Loading validates that node IDs are unique, every dependency refers to an
|
|
370
|
+
existing node, the graph is acyclic, and operations are registered in the
|
|
371
|
+
correct scope. It also rejects unknown fields, collisions between node inputs
|
|
372
|
+
and params, and duplicate YAML keys. `params` values must be strings or
|
|
373
|
+
numbers. `Node`, `Dag`, and `DagSuite` definitions, including their
|
|
374
|
+
collections, are immutable after validation; `DagExecutionResult` contains
|
|
375
|
+
the results of an individual execution.
|
|
376
|
+
|
|
377
|
+
Topological sorting uses Kahn's algorithm with $O(V+E)$ complexity, where $V$
|
|
378
|
+
is the number of nodes and $E$ is the number of dependencies. Execution sorts
|
|
379
|
+
only the subgraph needed for the requested outputs.
|
|
380
|
+
|
|
381
|
+
YAML definitions can invoke Python functions registered by the application.
|
|
382
|
+
Only load and execute configurations from trusted sources.
|
|
383
|
+
|
|
384
|
+
## Testing and building
|
|
385
|
+
|
|
386
|
+
```bash
|
|
387
|
+
uv run pytest
|
|
388
|
+
uv build
|
|
389
|
+
```
|
|
390
|
+
|
|
391
|
+
`pytest` automatically measures line and branch coverage for the entire
|
|
392
|
+
package. To run an individual test without coverage, use
|
|
393
|
+
`uv run pytest --no-cov ...`.
|
|
@@ -0,0 +1,382 @@
|
|
|
1
|
+
# lightweight-dag
|
|
2
|
+
|
|
3
|
+
Python library for defining, validating, and executing directed acyclic graphs
|
|
4
|
+
(DAGs). Operations are registered Python functions. DAGs can be loaded from
|
|
5
|
+
YAML or dictionaries, grouped into named suites, and executed for only the
|
|
6
|
+
requested outputs.
|
|
7
|
+
|
|
8
|
+
Requires Python 3.11 or later.
|
|
9
|
+
|
|
10
|
+
## Installation
|
|
11
|
+
|
|
12
|
+
Install the published package from PyPI with uv:
|
|
13
|
+
|
|
14
|
+
```bash
|
|
15
|
+
uv add lightweight-dag
|
|
16
|
+
```
|
|
17
|
+
|
|
18
|
+
To install the development and test dependencies when working from a clone of
|
|
19
|
+
the repository:
|
|
20
|
+
|
|
21
|
+
```bash
|
|
22
|
+
uv sync --group dev
|
|
23
|
+
```
|
|
24
|
+
|
|
25
|
+
The public API is available from the package root:
|
|
26
|
+
|
|
27
|
+
```python
|
|
28
|
+
from lightweight_dag import Dag, DagExecutionResult, DagSuite, Node, dag_node
|
|
29
|
+
```
|
|
30
|
+
|
|
31
|
+
## Registering operations
|
|
32
|
+
|
|
33
|
+
Operations must be synchronous functions defined at module level. Passing one
|
|
34
|
+
argument to the decorator registers a global operation:
|
|
35
|
+
|
|
36
|
+
```python
|
|
37
|
+
from lightweight_dag import dag_node
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
@dag_node("normalize")
|
|
41
|
+
def normalize(values):
|
|
42
|
+
return [value.strip().lower() for value in values]
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
Passing two arguments registers an operation for a specific DAG:
|
|
46
|
+
|
|
47
|
+
```python
|
|
48
|
+
from lightweight_dag import dag_node
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
@dag_node("example_report", "filter_orders")
|
|
52
|
+
def filter_monthly_orders(orders, minimum_total):
|
|
53
|
+
return [order for order in orders if order["total"] >= minimum_total]
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
A DAG without `name` looks up only global operations. A named DAG looks up only
|
|
57
|
+
operations registered for that name; there is no fallback to the global scope.
|
|
58
|
+
The same operation can therefore have different implementations in different
|
|
59
|
+
scopes. Registering an operation twice in the same scope is an error.
|
|
60
|
+
|
|
61
|
+
## Defining and running a suite
|
|
62
|
+
|
|
63
|
+
A YAML suite maps DAG names to DAG definitions. Each `Node.inputs` entry must
|
|
64
|
+
refer to another node in the same DAG. Source nodes have no inputs and produce
|
|
65
|
+
their values through their registered operation. The block-style suite example
|
|
66
|
+
in [examples/dags.yaml](examples/dags.yaml) defines both `monthly_report` and
|
|
67
|
+
`sales_audit`; the following code loads and executes both DAGs.
|
|
68
|
+
|
|
69
|
+
Register each operation in the scope of the DAG that uses it. The same
|
|
70
|
+
operation ID can have different implementations in different DAG scopes:
|
|
71
|
+
|
|
72
|
+
```python
|
|
73
|
+
from pathlib import Path
|
|
74
|
+
|
|
75
|
+
from lightweight_dag import DagSuite, dag_node
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
@dag_node("monthly_report", "load_orders")
|
|
79
|
+
def load_monthly_orders():
|
|
80
|
+
return [{"total": 75}, {"total": 35}]
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
@dag_node("monthly_report", "filter_orders")
|
|
84
|
+
def filter_monthly_orders(orders, minimum_total):
|
|
85
|
+
return [order for order in orders if order["total"] >= minimum_total]
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
@dag_node("monthly_report", "summarize_orders")
|
|
89
|
+
def summarize_monthly_orders(eligible_orders):
|
|
90
|
+
return {"count": len(eligible_orders)}
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
@dag_node("sales_audit", "load_orders")
|
|
94
|
+
def load_audit_orders():
|
|
95
|
+
return [{"total": 75}, {"total": 35}]
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
@dag_node("sales_audit", "calculate_total")
|
|
99
|
+
def calculate_total(orders):
|
|
100
|
+
return sum(order["total"] for order in orders)
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
@dag_node("sales_audit", "classify_total")
|
|
104
|
+
def classify_total(total, review_threshold):
|
|
105
|
+
return "review" if total >= review_threshold else "ok"
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
suite = DagSuite.from_yaml(Path("examples/dags.yaml").read_text())
|
|
109
|
+
monthly_report = suite["monthly_report"]
|
|
110
|
+
sales_audit = suite["sales_audit"]
|
|
111
|
+
|
|
112
|
+
report_result = monthly_report.run(outputs=["summary"])
|
|
113
|
+
audit_result = sales_audit.run(outputs=["audit_result"])
|
|
114
|
+
|
|
115
|
+
assert report_result.results["summary"] == {"count": 1}
|
|
116
|
+
assert audit_result.results["audit_result"] == "review"
|
|
117
|
+
```
|
|
118
|
+
|
|
119
|
+
The operation registry is scoped: `monthly_report` cannot use operations
|
|
120
|
+
registered only for `sales_audit`, and named DAGs do not fall back to global
|
|
121
|
+
operations. The YAML key supplies the DAG name; an explicit `name` field is
|
|
122
|
+
optional, but if present it must match the key. `suite[name]` raises `KeyError`
|
|
123
|
+
for an unknown name.
|
|
124
|
+
|
|
125
|
+
`DagSuite.from_dict(...)` accepts the same name-to-definition mapping. For a
|
|
126
|
+
single DAG, use `Dag.from_yaml(...)` or `Dag.from_dict(...)`. This standalone
|
|
127
|
+
example builds a DAG from a dictionary and executes its output:
|
|
128
|
+
|
|
129
|
+
```python
|
|
130
|
+
from lightweight_dag import Dag, dag_node
|
|
131
|
+
|
|
132
|
+
|
|
133
|
+
@dag_node("make_number")
|
|
134
|
+
def make_number():
|
|
135
|
+
return 21
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
@dag_node("double_number")
|
|
139
|
+
def double_number(number):
|
|
140
|
+
return number * 2
|
|
141
|
+
|
|
142
|
+
|
|
143
|
+
dag = Dag.from_dict(
|
|
144
|
+
{
|
|
145
|
+
"nodes": [
|
|
146
|
+
{"node_id": "number", "operation": "make_number"},
|
|
147
|
+
{
|
|
148
|
+
"node_id": "answer",
|
|
149
|
+
"operation": "double_number",
|
|
150
|
+
"inputs": ["number"],
|
|
151
|
+
},
|
|
152
|
+
]
|
|
153
|
+
}
|
|
154
|
+
)
|
|
155
|
+
result = dag.run(outputs=["answer"])
|
|
156
|
+
assert result.results["answer"] == 42
|
|
157
|
+
```
|
|
158
|
+
|
|
159
|
+
To load a single DAG from YAML, register its global operations and pass the
|
|
160
|
+
YAML document to `Dag.from_yaml(...)`:
|
|
161
|
+
|
|
162
|
+
```python
|
|
163
|
+
from lightweight_dag import Dag, dag_node
|
|
164
|
+
|
|
165
|
+
|
|
166
|
+
@dag_node("yaml_source")
|
|
167
|
+
def yaml_source(): return 21
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
@dag_node("yaml_double")
|
|
171
|
+
def yaml_double(number): return number * 2
|
|
172
|
+
|
|
173
|
+
|
|
174
|
+
dag = Dag.from_yaml("""{nodes: [
|
|
175
|
+
{node_id: number, operation: yaml_source},
|
|
176
|
+
{node_id: answer, operation: yaml_double, inputs: [number]}
|
|
177
|
+
]}""")
|
|
178
|
+
assert dag.run(outputs=["answer"]).results["answer"] == 42
|
|
179
|
+
```
|
|
180
|
+
|
|
181
|
+
## Inputs and execution
|
|
182
|
+
|
|
183
|
+
Call `Dag.run()` with no arguments to execute every node. To execute only
|
|
184
|
+
selected outputs and their dependencies, pass their IDs in `outputs`. Each
|
|
185
|
+
operation receives its predecessor node results and the YAML `params` as
|
|
186
|
+
keyword arguments. Argument names correspond to predecessor IDs or parameter
|
|
187
|
+
keys.
|
|
188
|
+
|
|
189
|
+
```python
|
|
190
|
+
all_results = monthly_report.run()
|
|
191
|
+
assert set(all_results.results) == {"orders", "eligible_orders", "summary"}
|
|
192
|
+
|
|
193
|
+
summary_result = monthly_report.run(outputs=["summary"])
|
|
194
|
+
assert summary_result.results["summary"] == {"count": 1}
|
|
195
|
+
|
|
196
|
+
if summary_result.status:
|
|
197
|
+
print(summary_result.results["summary"])
|
|
198
|
+
```
|
|
199
|
+
|
|
200
|
+
An unknown output raises `ValueError`. The result contains `status` and
|
|
201
|
+
`results`; `results` contains the computed node values. With `outputs=[]`, no
|
|
202
|
+
nodes are executed and `results` is empty. Omitting `outputs` is different:
|
|
203
|
+
all nodes are selected, including independent branches.
|
|
204
|
+
|
|
205
|
+
Shared predecessors are computed only once, and nodes outside the requested
|
|
206
|
+
paths are not executed. In this graph, requesting only `d` computes `a`, `b`,
|
|
207
|
+
and `d`; `c` is not executed:
|
|
208
|
+
|
|
209
|
+
```mermaid
|
|
210
|
+
flowchart TD
|
|
211
|
+
a --> b --> c
|
|
212
|
+
b --> d
|
|
213
|
+
```
|
|
214
|
+
|
|
215
|
+
```python
|
|
216
|
+
from lightweight_dag import Dag, dag_node
|
|
217
|
+
|
|
218
|
+
|
|
219
|
+
@dag_node("get_a")
|
|
220
|
+
def get_a():
|
|
221
|
+
return "a"
|
|
222
|
+
|
|
223
|
+
|
|
224
|
+
@dag_node("make_b")
|
|
225
|
+
def make_b(a):
|
|
226
|
+
return f"{a}->b"
|
|
227
|
+
|
|
228
|
+
|
|
229
|
+
@dag_node("make_c")
|
|
230
|
+
def make_c(b):
|
|
231
|
+
return f"{b}->c"
|
|
232
|
+
|
|
233
|
+
|
|
234
|
+
@dag_node("make_d")
|
|
235
|
+
def make_d(b):
|
|
236
|
+
return f"{b}->d"
|
|
237
|
+
|
|
238
|
+
|
|
239
|
+
dag = Dag.from_dict({"nodes": [
|
|
240
|
+
{"node_id": "a", "operation": "get_a"},
|
|
241
|
+
{"node_id": "b", "operation": "make_b", "inputs": ["a"]},
|
|
242
|
+
{"node_id": "c", "operation": "make_c", "inputs": ["b"]},
|
|
243
|
+
{"node_id": "d", "operation": "make_d", "inputs": ["b"]},
|
|
244
|
+
]})
|
|
245
|
+
result = dag.run(outputs=["d"])
|
|
246
|
+
|
|
247
|
+
assert result.results == {
|
|
248
|
+
"a": "a",
|
|
249
|
+
"b": "a->b",
|
|
250
|
+
"d": "a->b->d",
|
|
251
|
+
}
|
|
252
|
+
assert "c" not in result.results
|
|
253
|
+
```
|
|
254
|
+
|
|
255
|
+
### Using an LLM with LangChain
|
|
256
|
+
|
|
257
|
+
This example simulates reading a customer record from a database, then sends
|
|
258
|
+
that context and a user prompt to an LLM through LangChain. Install the optional
|
|
259
|
+
integration and set an API key before running it:
|
|
260
|
+
|
|
261
|
+
```bash
|
|
262
|
+
pip install langchain-openai
|
|
263
|
+
export OPENAI_API_KEY="your-api-key"
|
|
264
|
+
```
|
|
265
|
+
|
|
266
|
+
```python
|
|
267
|
+
from lightweight_dag import Dag, dag_node
|
|
268
|
+
from langchain_openai import ChatOpenAI
|
|
269
|
+
from langchain_core.prompts import ChatPromptTemplate
|
|
270
|
+
|
|
271
|
+
|
|
272
|
+
@dag_node("load_customer")
|
|
273
|
+
def load_customer():
|
|
274
|
+
# Replace this in-memory record with a database query.
|
|
275
|
+
return {"name": "Ada", "plan": "Pro", "last_order": "Notebook"}
|
|
276
|
+
|
|
277
|
+
|
|
278
|
+
@dag_node("answer_user")
|
|
279
|
+
def answer_user(customer_data, user_prompt):
|
|
280
|
+
prompt = ChatPromptTemplate.from_messages([
|
|
281
|
+
("system", "Answer using the customer data provided: {customer_data}"),
|
|
282
|
+
("human", "{user_prompt}"),
|
|
283
|
+
])
|
|
284
|
+
chain = prompt | ChatOpenAI(model="gpt-4o-mini")
|
|
285
|
+
response = chain.invoke({
|
|
286
|
+
"customer_data": customer_data,
|
|
287
|
+
"user_prompt": user_prompt,
|
|
288
|
+
})
|
|
289
|
+
return response.content
|
|
290
|
+
|
|
291
|
+
|
|
292
|
+
support_dag = Dag.from_dict({"nodes": [
|
|
293
|
+
{"node_id": "customer_data", "operation": "load_customer"},
|
|
294
|
+
{
|
|
295
|
+
"node_id": "answer",
|
|
296
|
+
"operation": "answer_user",
|
|
297
|
+
"inputs": ["customer_data"],
|
|
298
|
+
"params": {"user_prompt": "Qual è il mio ultimo ordine?"},
|
|
299
|
+
},
|
|
300
|
+
]})
|
|
301
|
+
result = support_dag.run(outputs=["answer"])
|
|
302
|
+
print(result.results["answer"])
|
|
303
|
+
```
|
|
304
|
+
|
|
305
|
+
```python
|
|
306
|
+
from lightweight_dag import Dag, dag_node
|
|
307
|
+
|
|
308
|
+
|
|
309
|
+
@dag_node("branching_demo", "label")
|
|
310
|
+
def label_node(label, **predecessors): return label
|
|
311
|
+
|
|
312
|
+
|
|
313
|
+
branching_dag = Dag.from_dict({"name": "branching_demo", "nodes": [
|
|
314
|
+
{"node_id": "A", "operation": "label", "params": {"label": "A"}},
|
|
315
|
+
{"node_id": "B", "operation": "label", "inputs": ["A"], "params": {"label": "B"}},
|
|
316
|
+
{"node_id": "E", "operation": "label", "inputs": ["B"], "params": {"label": "E"}},
|
|
317
|
+
{"node_id": "G", "operation": "label", "inputs": ["E"], "params": {"label": "G"}},
|
|
318
|
+
{"node_id": "H", "operation": "label", "inputs": ["E"], "params": {"label": "H"}},
|
|
319
|
+
{"node_id": "F", "operation": "label", "inputs": ["B"], "params": {"label": "F"}},
|
|
320
|
+
{"node_id": "I", "operation": "label", "inputs": ["F"], "params": {"label": "I"}},
|
|
321
|
+
{"node_id": "J", "operation": "label", "inputs": ["F"], "params": {"label": "J"}},
|
|
322
|
+
{"node_id": "K", "operation": "label", "inputs": ["G", "H"], "params": {"label": "K"}},
|
|
323
|
+
{"node_id": "N", "operation": "label", "inputs": ["I", "J", "G"], "params": {"label": "N"}},
|
|
324
|
+
]})
|
|
325
|
+
result = branching_dag.run(outputs=["G", "J", "A"])
|
|
326
|
+
assert set(result.results) == {"A", "B", "E", "F", "G", "J"}
|
|
327
|
+
```
|
|
328
|
+
|
|
329
|
+
If an operation raises an exception, its result and the results of its
|
|
330
|
+
dependent nodes become `None`, and `status` becomes `False`. Requested,
|
|
331
|
+
independent branches continue to execute. This example demonstrates the
|
|
332
|
+
failure result:
|
|
333
|
+
|
|
334
|
+
```python
|
|
335
|
+
from lightweight_dag import Dag, dag_node
|
|
336
|
+
|
|
337
|
+
|
|
338
|
+
@dag_node("always_fail")
|
|
339
|
+
def always_fail(): raise RuntimeError("operation failed")
|
|
340
|
+
|
|
341
|
+
|
|
342
|
+
@dag_node("consume_result")
|
|
343
|
+
def consume_result(value): return value
|
|
344
|
+
|
|
345
|
+
|
|
346
|
+
failure_dag = Dag.from_dict({"nodes": [
|
|
347
|
+
{"node_id": "failed", "operation": "always_fail"},
|
|
348
|
+
{"node_id": "dependent", "operation": "consume_result", "inputs": ["failed"]},
|
|
349
|
+
]})
|
|
350
|
+
failure_result = failure_dag.run(outputs=["dependent"])
|
|
351
|
+
assert failure_result.status is False
|
|
352
|
+
assert failure_result.results["failed"] is None
|
|
353
|
+
assert failure_result.results["dependent"] is None
|
|
354
|
+
```
|
|
355
|
+
|
|
356
|
+
## Validation and security
|
|
357
|
+
|
|
358
|
+
Loading validates that node IDs are unique, every dependency refers to an
|
|
359
|
+
existing node, the graph is acyclic, and operations are registered in the
|
|
360
|
+
correct scope. It also rejects unknown fields, collisions between node inputs
|
|
361
|
+
and params, and duplicate YAML keys. `params` values must be strings or
|
|
362
|
+
numbers. `Node`, `Dag`, and `DagSuite` definitions, including their
|
|
363
|
+
collections, are immutable after validation; `DagExecutionResult` contains
|
|
364
|
+
the results of an individual execution.
|
|
365
|
+
|
|
366
|
+
Topological sorting uses Kahn's algorithm with $O(V+E)$ complexity, where $V$
|
|
367
|
+
is the number of nodes and $E$ is the number of dependencies. Execution sorts
|
|
368
|
+
only the subgraph needed for the requested outputs.
|
|
369
|
+
|
|
370
|
+
YAML definitions can invoke Python functions registered by the application.
|
|
371
|
+
Only load and execute configurations from trusted sources.
|
|
372
|
+
|
|
373
|
+
## Testing and building
|
|
374
|
+
|
|
375
|
+
```bash
|
|
376
|
+
uv run pytest
|
|
377
|
+
uv build
|
|
378
|
+
```
|
|
379
|
+
|
|
380
|
+
`pytest` automatically measures line and branch coverage for the entire
|
|
381
|
+
package. To run an individual test without coverage, use
|
|
382
|
+
`uv run pytest --no-cov ...`.
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "lightweight-dag"
|
|
3
|
+
version = "0.1.1"
|
|
4
|
+
description = "A lightweight Python library for defining and executing DAG workflows"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.11"
|
|
7
|
+
dependencies = [
|
|
8
|
+
"pydantic>=2.13.5",
|
|
9
|
+
"pyyaml>=6.0.3",
|
|
10
|
+
]
|
|
11
|
+
|
|
12
|
+
[[project.authors]]
|
|
13
|
+
name = "Lorenzo Gardini"
|
|
14
|
+
email = "l.gardo98@gmail.com"
|
|
15
|
+
|
|
16
|
+
[build-system]
|
|
17
|
+
requires = ["uv_build>=0.10.0,<0.11.0"]
|
|
18
|
+
build-backend = "uv_build"
|
|
19
|
+
|
|
20
|
+
[dependency-groups]
|
|
21
|
+
dev = [
|
|
22
|
+
"pytest>=9.1.1",
|
|
23
|
+
"pytest-cov>=7.1.0",
|
|
24
|
+
]
|
|
25
|
+
|
|
26
|
+
[tool.pytest.ini_options]
|
|
27
|
+
testpaths = ["tests"]
|
|
28
|
+
addopts = [
|
|
29
|
+
"--cov=lightweight_dag",
|
|
30
|
+
"--cov-branch",
|
|
31
|
+
"--cov-report=term-missing",
|
|
32
|
+
]
|
|
33
|
+
|
|
34
|
+
[tool.coverage.run]
|
|
35
|
+
branch = true
|
|
36
|
+
source = ["lightweight_dag"]
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "lightweight-dag"
|
|
3
|
+
version = "0.1.1"
|
|
4
|
+
description = "A lightweight Python library for defining and executing DAG workflows"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
authors = [
|
|
7
|
+
{ name = "Lorenzo Gardini", email = "l.gardo98@gmail.com" }
|
|
8
|
+
]
|
|
9
|
+
requires-python = ">=3.11"
|
|
10
|
+
dependencies = [
|
|
11
|
+
"pydantic>=2.13.5",
|
|
12
|
+
"pyyaml>=6.0.3",
|
|
13
|
+
]
|
|
14
|
+
|
|
15
|
+
[build-system]
|
|
16
|
+
requires = ["uv_build>=0.10.0,<0.11.0"]
|
|
17
|
+
build-backend = "uv_build"
|
|
18
|
+
|
|
19
|
+
[dependency-groups]
|
|
20
|
+
dev = [
|
|
21
|
+
"pytest>=9.1.1",
|
|
22
|
+
"pytest-cov>=7.1.0",
|
|
23
|
+
]
|
|
24
|
+
|
|
25
|
+
[tool.pytest.ini_options]
|
|
26
|
+
testpaths = ["tests"]
|
|
27
|
+
addopts = [
|
|
28
|
+
"--cov=lightweight_dag",
|
|
29
|
+
"--cov-branch",
|
|
30
|
+
"--cov-report=term-missing",
|
|
31
|
+
]
|
|
32
|
+
|
|
33
|
+
[tool.coverage.run]
|
|
34
|
+
branch = true
|
|
35
|
+
source = ["lightweight_dag"]
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
"""Public API for building and running lightweight DAG workflows."""
|
|
2
|
+
|
|
3
|
+
from .dag import Dag, DagExecutionResult, Node
|
|
4
|
+
from .node_registry import dag_node
|
|
5
|
+
from .suite import DagSuite
|
|
6
|
+
|
|
7
|
+
__all__ = [
|
|
8
|
+
"Dag",
|
|
9
|
+
"DagExecutionResult",
|
|
10
|
+
"DagSuite",
|
|
11
|
+
"Node",
|
|
12
|
+
"dag_node",
|
|
13
|
+
]
|
|
@@ -0,0 +1,306 @@
|
|
|
1
|
+
"""Typed DAG models, graph validation, execution, and named DAG suites."""
|
|
2
|
+
|
|
3
|
+
from collections import deque
|
|
4
|
+
from functools import cached_property
|
|
5
|
+
from typing import Any
|
|
6
|
+
|
|
7
|
+
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
|
8
|
+
|
|
9
|
+
from lightweight_dag.node_registry import get_registered_nodes
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class _FrozenList(list):
|
|
13
|
+
"""A list-compatible container that rejects in-place mutation."""
|
|
14
|
+
|
|
15
|
+
@staticmethod
|
|
16
|
+
def _reject_mutation(*args, **kwargs):
|
|
17
|
+
raise TypeError("DAG model collections are immutable")
|
|
18
|
+
|
|
19
|
+
__delitem__ = __iadd__ = __imul__ = __setitem__ = _reject_mutation
|
|
20
|
+
append = clear = extend = insert = pop = remove = reverse = sort = _reject_mutation
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class _FrozenDict(dict):
|
|
24
|
+
"""A dict-compatible container that rejects in-place mutation."""
|
|
25
|
+
|
|
26
|
+
@staticmethod
|
|
27
|
+
def _reject_mutation(*args, **kwargs):
|
|
28
|
+
raise TypeError("DAG model collections are immutable")
|
|
29
|
+
|
|
30
|
+
__delitem__ = __ior__ = __setitem__ = _reject_mutation
|
|
31
|
+
clear = pop = popitem = setdefault = update = _reject_mutation
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _load_yaml(yaml_str: str) -> Any:
|
|
35
|
+
"""Load safe YAML while rejecting duplicate mapping keys."""
|
|
36
|
+
import yaml
|
|
37
|
+
from yaml.constructor import ConstructorError
|
|
38
|
+
|
|
39
|
+
class UniqueKeyLoader(yaml.SafeLoader):
|
|
40
|
+
pass
|
|
41
|
+
|
|
42
|
+
def construct_unique_mapping(loader, node, deep=False):
|
|
43
|
+
loader.flatten_mapping(node)
|
|
44
|
+
mapping = {}
|
|
45
|
+
for key_node, value_node in node.value:
|
|
46
|
+
key = loader.construct_object(key_node, deep=deep)
|
|
47
|
+
if key in mapping:
|
|
48
|
+
raise ConstructorError(
|
|
49
|
+
"while constructing a mapping",
|
|
50
|
+
node.start_mark,
|
|
51
|
+
f"found duplicate key {key!r}",
|
|
52
|
+
key_node.start_mark,
|
|
53
|
+
)
|
|
54
|
+
mapping[key] = loader.construct_object(value_node, deep=deep)
|
|
55
|
+
return mapping
|
|
56
|
+
|
|
57
|
+
UniqueKeyLoader.add_constructor(
|
|
58
|
+
yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
|
|
59
|
+
construct_unique_mapping,
|
|
60
|
+
)
|
|
61
|
+
return yaml.load(yaml_str, Loader=UniqueKeyLoader)
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
class Node(BaseModel):
|
|
65
|
+
"""One operation in a DAG and the IDs of its predecessor nodes."""
|
|
66
|
+
|
|
67
|
+
node_id: str = Field(..., description="The unique identifier for this node")
|
|
68
|
+
operation: str = Field(..., description="The operation performed by this node")
|
|
69
|
+
inputs: list[str] = Field(
|
|
70
|
+
default_factory=list, description="List of input node IDs"
|
|
71
|
+
)
|
|
72
|
+
params: dict[str, int | float | str] = Field(
|
|
73
|
+
default_factory=dict, description="Dictionary of parameters for this node"
|
|
74
|
+
)
|
|
75
|
+
|
|
76
|
+
model_config = ConfigDict(
|
|
77
|
+
extra="forbid",
|
|
78
|
+
frozen=True,
|
|
79
|
+
)
|
|
80
|
+
|
|
81
|
+
@field_validator("inputs")
|
|
82
|
+
@classmethod
|
|
83
|
+
def freeze_inputs(cls, value: list[str]) -> list[str]:
|
|
84
|
+
return _FrozenList(value)
|
|
85
|
+
|
|
86
|
+
@field_validator("params")
|
|
87
|
+
@classmethod
|
|
88
|
+
def freeze_params(
|
|
89
|
+
cls, value: dict[str, int | float | str]
|
|
90
|
+
) -> dict[str, int | float | str]:
|
|
91
|
+
return _FrozenDict(value)
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
class DagExecutionResult(BaseModel):
|
|
95
|
+
"""Execution status and values produced or retained during a DAG run."""
|
|
96
|
+
|
|
97
|
+
status: bool
|
|
98
|
+
results: dict[str, Any] = Field(
|
|
99
|
+
default_factory=dict, description="Dictionary of results for this DAG execution"
|
|
100
|
+
)
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
class Dag(BaseModel):
|
|
104
|
+
"""A validated directed acyclic graph of registered operations.
|
|
105
|
+
|
|
106
|
+
Unnamed DAGs resolve global operations. Named DAGs resolve operations from
|
|
107
|
+
their matching registry scope only.
|
|
108
|
+
"""
|
|
109
|
+
|
|
110
|
+
nodes: list[Node]
|
|
111
|
+
name: str | None = Field(
|
|
112
|
+
default=None, description="Optional registry scope for this DAG"
|
|
113
|
+
)
|
|
114
|
+
|
|
115
|
+
model_config = ConfigDict(
|
|
116
|
+
extra="forbid",
|
|
117
|
+
frozen=True,
|
|
118
|
+
)
|
|
119
|
+
|
|
120
|
+
@field_validator("name")
|
|
121
|
+
@classmethod
|
|
122
|
+
def check_name(cls, value: str | None) -> str | None:
|
|
123
|
+
"""Reject blank names while allowing unnamed DAGs."""
|
|
124
|
+
if value is not None and not value.strip():
|
|
125
|
+
raise ValueError("DAG name must be non-empty")
|
|
126
|
+
return value
|
|
127
|
+
|
|
128
|
+
@field_validator("nodes")
|
|
129
|
+
@classmethod
|
|
130
|
+
def freeze_collections(cls, value: list) -> list:
|
|
131
|
+
return _FrozenList(value)
|
|
132
|
+
|
|
133
|
+
@cached_property
|
|
134
|
+
def _nodes_by_id(self) -> dict[str, Node]:
|
|
135
|
+
return _FrozenDict({node.node_id: node for node in self.nodes})
|
|
136
|
+
|
|
137
|
+
@model_validator(mode="after")
|
|
138
|
+
def check_duplicate_node_ids(self) -> "Dag":
|
|
139
|
+
"""Reject repeated node IDs."""
|
|
140
|
+
node_ids = [node.node_id for node in self.nodes]
|
|
141
|
+
if len(node_ids) != len(set(node_ids)):
|
|
142
|
+
raise ValueError("Node IDs must be unique")
|
|
143
|
+
return self
|
|
144
|
+
|
|
145
|
+
@model_validator(mode="after")
|
|
146
|
+
def check_inputs_exist(self) -> "Dag":
|
|
147
|
+
"""Ensure node inputs refer to existing nodes."""
|
|
148
|
+
node_ids = {node.node_id for node in self.nodes}
|
|
149
|
+
missing_inputs = {
|
|
150
|
+
input_id
|
|
151
|
+
for node in self.nodes
|
|
152
|
+
for input_id in node.inputs
|
|
153
|
+
if input_id not in node_ids
|
|
154
|
+
}
|
|
155
|
+
if missing_inputs:
|
|
156
|
+
raise ValueError(
|
|
157
|
+
f"Inputs reference non-existent nodes: {sorted(missing_inputs)}"
|
|
158
|
+
)
|
|
159
|
+
return self
|
|
160
|
+
|
|
161
|
+
@model_validator(mode="after")
|
|
162
|
+
def check_input_parameter_conflicts(self) -> "Dag":
|
|
163
|
+
"""Reject keyword names that collide between inputs and params."""
|
|
164
|
+
conflicts = {
|
|
165
|
+
node.node_id: sorted(set(node.inputs) & node.params.keys())
|
|
166
|
+
for node in self.nodes
|
|
167
|
+
if set(node.inputs) & node.params.keys()
|
|
168
|
+
}
|
|
169
|
+
if conflicts:
|
|
170
|
+
raise ValueError(f"Node inputs conflict with parameters: {conflicts}")
|
|
171
|
+
return self
|
|
172
|
+
|
|
173
|
+
@model_validator(mode="after")
|
|
174
|
+
def check_acyclic(self) -> "Dag":
|
|
175
|
+
"""Reject dependency cycles."""
|
|
176
|
+
self.topological_sort(self.nodes)
|
|
177
|
+
return self
|
|
178
|
+
|
|
179
|
+
@model_validator(mode="after")
|
|
180
|
+
def check_operations_are_registered(self) -> "Dag":
|
|
181
|
+
"""Ensure every operation exists in this DAG's registry scope."""
|
|
182
|
+
registered_functions = (
|
|
183
|
+
get_registered_nodes()
|
|
184
|
+
if self.name is None
|
|
185
|
+
else get_registered_nodes(self.name)
|
|
186
|
+
)
|
|
187
|
+
function_not_found = [
|
|
188
|
+
node.node_id
|
|
189
|
+
for node in self.nodes
|
|
190
|
+
if node.operation not in registered_functions
|
|
191
|
+
]
|
|
192
|
+
if function_not_found:
|
|
193
|
+
raise ValueError(
|
|
194
|
+
f"Operations reference non-existent functions: {sorted(function_not_found)}"
|
|
195
|
+
)
|
|
196
|
+
return self
|
|
197
|
+
|
|
198
|
+
@classmethod
|
|
199
|
+
def from_yaml(cls, yaml_str: str) -> "Dag":
|
|
200
|
+
"""Parse and validate one DAG from a YAML document."""
|
|
201
|
+
data = _load_yaml(yaml_str)
|
|
202
|
+
return cls.model_validate(data)
|
|
203
|
+
|
|
204
|
+
@classmethod
|
|
205
|
+
def from_dict(cls, data: dict) -> "Dag":
|
|
206
|
+
"""Validate one DAG from a dictionary."""
|
|
207
|
+
return cls.model_validate(data)
|
|
208
|
+
|
|
209
|
+
@classmethod
|
|
210
|
+
def topological_sort(cls, nodes: list[Node]) -> list[str]:
|
|
211
|
+
"""Return node IDs in dependency order using Kahn's linear-time algorithm."""
|
|
212
|
+
node_ids = [node.node_id for node in nodes]
|
|
213
|
+
if len(node_ids) != len(set(node_ids)):
|
|
214
|
+
raise ValueError("Node IDs must be unique")
|
|
215
|
+
|
|
216
|
+
known_ids = set(node_ids)
|
|
217
|
+
missing_inputs = {
|
|
218
|
+
input_id
|
|
219
|
+
for node in nodes
|
|
220
|
+
for input_id in node.inputs
|
|
221
|
+
if input_id not in known_ids
|
|
222
|
+
}
|
|
223
|
+
if missing_inputs:
|
|
224
|
+
raise ValueError(
|
|
225
|
+
f"Inputs reference non-existent nodes: {sorted(missing_inputs)}"
|
|
226
|
+
)
|
|
227
|
+
|
|
228
|
+
in_degree = {}
|
|
229
|
+
dependents = {node_id: [] for node_id in node_ids}
|
|
230
|
+
for node in nodes:
|
|
231
|
+
dependencies = set(node.inputs) & known_ids
|
|
232
|
+
in_degree[node.node_id] = len(dependencies)
|
|
233
|
+
for dependency in dependencies:
|
|
234
|
+
dependents[dependency].append(node.node_id)
|
|
235
|
+
|
|
236
|
+
ready = deque(node.node_id for node in nodes if in_degree[node.node_id] == 0)
|
|
237
|
+
result = []
|
|
238
|
+
while ready:
|
|
239
|
+
node_id = ready.popleft()
|
|
240
|
+
result.append(node_id)
|
|
241
|
+
for dependent in dependents[node_id]:
|
|
242
|
+
in_degree[dependent] -= 1
|
|
243
|
+
if in_degree[dependent] == 0:
|
|
244
|
+
ready.append(dependent)
|
|
245
|
+
|
|
246
|
+
if len(result) != len(nodes):
|
|
247
|
+
raise ValueError("DAG contains a cycle")
|
|
248
|
+
|
|
249
|
+
return result
|
|
250
|
+
|
|
251
|
+
def run(self, outputs: list[str] | None = None) -> DagExecutionResult:
|
|
252
|
+
"""Execute requested outputs and their transitive dependencies.
|
|
253
|
+
|
|
254
|
+
Each operation receives predecessor results and node params as keyword
|
|
255
|
+
arguments. Failed nodes are skipped along with their descendants, while
|
|
256
|
+
independent requested branches continue to run. If outputs is omitted,
|
|
257
|
+
every node is selected; an empty list selects no nodes.
|
|
258
|
+
"""
|
|
259
|
+
nodes_by_id = self._nodes_by_id
|
|
260
|
+
selected_outputs = list(nodes_by_id) if outputs is None else outputs
|
|
261
|
+
unknown_outputs = set(selected_outputs) - nodes_by_id.keys()
|
|
262
|
+
if unknown_outputs:
|
|
263
|
+
raise ValueError(
|
|
264
|
+
f"Outputs reference non-existent nodes: {sorted(unknown_outputs)}"
|
|
265
|
+
)
|
|
266
|
+
|
|
267
|
+
required_node_ids = set()
|
|
268
|
+
required_nodes = []
|
|
269
|
+
pending_node_ids = list(selected_outputs)
|
|
270
|
+
while pending_node_ids:
|
|
271
|
+
node_id = pending_node_ids.pop()
|
|
272
|
+
if node_id not in required_node_ids:
|
|
273
|
+
required_node_ids.add(node_id)
|
|
274
|
+
node = nodes_by_id[node_id]
|
|
275
|
+
required_nodes.append(node)
|
|
276
|
+
pending_node_ids.extend(node.inputs)
|
|
277
|
+
sorted_node_ids = [node_id for node_id in self.topological_sort(required_nodes)]
|
|
278
|
+
context = {}
|
|
279
|
+
is_run_successful = True
|
|
280
|
+
failed_node_ids = set()
|
|
281
|
+
registered_functions = (
|
|
282
|
+
get_registered_nodes()
|
|
283
|
+
if self.name is None
|
|
284
|
+
else get_registered_nodes(self.name)
|
|
285
|
+
)
|
|
286
|
+
|
|
287
|
+
for node_id in sorted_node_ids:
|
|
288
|
+
node = nodes_by_id[node_id]
|
|
289
|
+
if any(input_id in failed_node_ids for input_id in node.inputs):
|
|
290
|
+
context[node_id] = None
|
|
291
|
+
failed_node_ids.add(node_id)
|
|
292
|
+
is_run_successful = False
|
|
293
|
+
continue
|
|
294
|
+
|
|
295
|
+
node_inputs = {input_id: context[input_id] for input_id in node.inputs}
|
|
296
|
+
node_inputs.update(node.params)
|
|
297
|
+
try:
|
|
298
|
+
result = registered_functions[node.operation](**node_inputs)
|
|
299
|
+
except Exception:
|
|
300
|
+
context[node_id] = None
|
|
301
|
+
failed_node_ids.add(node_id)
|
|
302
|
+
is_run_successful = False
|
|
303
|
+
else:
|
|
304
|
+
context[node_id] = result
|
|
305
|
+
|
|
306
|
+
return DagExecutionResult(status=is_run_successful, results=context)
|
|
@@ -0,0 +1,83 @@
|
|
|
1
|
+
"""Registry for synchronous operations used by DAGs."""
|
|
2
|
+
|
|
3
|
+
import inspect
|
|
4
|
+
from collections.abc import Callable
|
|
5
|
+
from threading import Lock
|
|
6
|
+
from typing import Any, TypeVar, overload
|
|
7
|
+
|
|
8
|
+
F = TypeVar("F", bound=Callable[..., Any])
|
|
9
|
+
|
|
10
|
+
_nodes: dict[tuple[str | None, str], Callable[..., Any]] = {}
|
|
11
|
+
_nodes_lock = Lock()
|
|
12
|
+
_MISSING = object()
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
@overload
|
|
16
|
+
def dag_node(node_id: str) -> Callable[[F], F]: ...
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@overload
|
|
20
|
+
def dag_node(dag_name: str, node_id: str) -> Callable[[F], F]: ...
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def dag_node(
|
|
24
|
+
dag_name_or_node_id: str,
|
|
25
|
+
node_id: str | None | object = _MISSING,
|
|
26
|
+
) -> Callable[[F], F]:
|
|
27
|
+
"""Register a global operation or one scoped to a named DAG.
|
|
28
|
+
|
|
29
|
+
Use ``@dag_node("operation")`` for a global operation or
|
|
30
|
+
``@dag_node("dag_name", "operation")`` for a DAG-specific version.
|
|
31
|
+
Named DAG scopes do not fall back to global registrations.
|
|
32
|
+
"""
|
|
33
|
+
if node_id is _MISSING:
|
|
34
|
+
dag_name = None
|
|
35
|
+
operation_id = dag_name_or_node_id
|
|
36
|
+
else:
|
|
37
|
+
dag_name = dag_name_or_node_id
|
|
38
|
+
operation_id = node_id
|
|
39
|
+
|
|
40
|
+
if not isinstance(operation_id, str) or not operation_id:
|
|
41
|
+
raise ValueError("node_id must be a non-empty string")
|
|
42
|
+
if dag_name is not None and (not isinstance(dag_name, str) or not dag_name):
|
|
43
|
+
raise ValueError("dag_name must be a non-empty string")
|
|
44
|
+
|
|
45
|
+
def register(func: F) -> F:
|
|
46
|
+
if inspect.iscoroutinefunction(func):
|
|
47
|
+
raise TypeError("@dag_node does not support async functions")
|
|
48
|
+
if not inspect.isfunction(func):
|
|
49
|
+
raise TypeError("@dag_node can only decorate functions")
|
|
50
|
+
if func.__qualname__ != func.__name__:
|
|
51
|
+
raise TypeError("@dag_node can only decorate module-level functions")
|
|
52
|
+
key = (dag_name, operation_id)
|
|
53
|
+
with _nodes_lock:
|
|
54
|
+
if key in _nodes:
|
|
55
|
+
raise ValueError(f"Node ID '{operation_id}' is already registered")
|
|
56
|
+
_nodes[key] = func
|
|
57
|
+
return func
|
|
58
|
+
|
|
59
|
+
return register
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def get_registered_nodes(
|
|
63
|
+
dag_name: str | None = None,
|
|
64
|
+
) -> dict[str, Callable[..., Any]]:
|
|
65
|
+
"""Return a copy of the global registry or exactly one named-DAG scope."""
|
|
66
|
+
with _nodes_lock:
|
|
67
|
+
return {
|
|
68
|
+
node_id: func
|
|
69
|
+
for (scope, node_id), func in _nodes.items()
|
|
70
|
+
if scope == dag_name
|
|
71
|
+
}
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
def get_node_ids(dag_name: str | None = None) -> list[str]:
|
|
75
|
+
"""Return operation IDs in the global registry or named-DAG scope."""
|
|
76
|
+
with _nodes_lock:
|
|
77
|
+
return [node_id for scope, node_id in _nodes if scope == dag_name]
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
def get_node(node_id: str, dag_name: str | None = None) -> Callable[..., Any] | None:
|
|
81
|
+
"""Look up an operation by ID in the requested registry scope."""
|
|
82
|
+
with _nodes_lock:
|
|
83
|
+
return _nodes.get((dag_name, node_id))
|
|
@@ -0,0 +1,80 @@
|
|
|
1
|
+
"""Named DAG suite loading and validation."""
|
|
2
|
+
|
|
3
|
+
from collections.abc import Mapping
|
|
4
|
+
from typing import Any
|
|
5
|
+
|
|
6
|
+
from pydantic import BaseModel, ConfigDict, field_validator, model_validator
|
|
7
|
+
|
|
8
|
+
from lightweight_dag.dag import Dag, _FrozenDict, _load_yaml
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class DagSuite(BaseModel):
|
|
12
|
+
"""A name-indexed collection of DAGs, accessible with ``suite[name]``."""
|
|
13
|
+
|
|
14
|
+
dags: dict[str, Dag]
|
|
15
|
+
|
|
16
|
+
model_config = ConfigDict(
|
|
17
|
+
extra="forbid",
|
|
18
|
+
frozen=True,
|
|
19
|
+
)
|
|
20
|
+
|
|
21
|
+
@field_validator("dags")
|
|
22
|
+
@classmethod
|
|
23
|
+
def freeze_dags(cls, value: dict[str, Dag]) -> dict[str, Dag]:
|
|
24
|
+
return _FrozenDict(value)
|
|
25
|
+
|
|
26
|
+
@model_validator(mode="after")
|
|
27
|
+
def check_dag_names_match_keys(self) -> "DagSuite":
|
|
28
|
+
"""Ensure every suite key matches the corresponding DAG name."""
|
|
29
|
+
mismatched_names = {
|
|
30
|
+
suite_name: dag.name
|
|
31
|
+
for suite_name, dag in self.dags.items()
|
|
32
|
+
if suite_name != dag.name
|
|
33
|
+
}
|
|
34
|
+
if mismatched_names:
|
|
35
|
+
raise ValueError(f"Suite keys must match DAG names: {mismatched_names}")
|
|
36
|
+
return self
|
|
37
|
+
|
|
38
|
+
def __getitem__(self, name: str) -> Dag:
|
|
39
|
+
"""Return the DAG registered under ``name``."""
|
|
40
|
+
return self.dags[name]
|
|
41
|
+
|
|
42
|
+
@classmethod
|
|
43
|
+
def from_dict(cls, data: Mapping[str, Mapping[str, Any] | Dag]) -> "DagSuite":
|
|
44
|
+
"""Build a suite from a mapping of names to DAG definitions or models."""
|
|
45
|
+
if not isinstance(data, Mapping):
|
|
46
|
+
raise TypeError("DAG suite must be a mapping of names to DAG definitions")
|
|
47
|
+
|
|
48
|
+
dags = {}
|
|
49
|
+
for name, dag_data in data.items():
|
|
50
|
+
if not isinstance(name, str):
|
|
51
|
+
raise TypeError("DAG suite names must be strings")
|
|
52
|
+
if not name:
|
|
53
|
+
raise ValueError("DAG suite names must be non-empty strings")
|
|
54
|
+
|
|
55
|
+
if isinstance(dag_data, Dag):
|
|
56
|
+
if dag_data.name != name:
|
|
57
|
+
raise ValueError(
|
|
58
|
+
f"DAG name '{dag_data.name}' does not match suite key '{name}'"
|
|
59
|
+
)
|
|
60
|
+
dags[name] = dag_data
|
|
61
|
+
continue
|
|
62
|
+
|
|
63
|
+
if not isinstance(dag_data, Mapping):
|
|
64
|
+
raise TypeError(
|
|
65
|
+
f"DAG definition for '{name}' must be a mapping or a Dag"
|
|
66
|
+
)
|
|
67
|
+
configured_name = dag_data.get("name", name)
|
|
68
|
+
if configured_name != name:
|
|
69
|
+
raise ValueError(
|
|
70
|
+
f"DAG name '{configured_name}' does not match suite key '{name}'"
|
|
71
|
+
)
|
|
72
|
+
dags[name] = Dag.model_validate({**dag_data, "name": name})
|
|
73
|
+
|
|
74
|
+
return cls(dags=dags)
|
|
75
|
+
|
|
76
|
+
@classmethod
|
|
77
|
+
def from_yaml(cls, yaml_str: str) -> "DagSuite":
|
|
78
|
+
"""Parse a name-to-DAG mapping from YAML and validate each DAG."""
|
|
79
|
+
data = _load_yaml(yaml_str)
|
|
80
|
+
return cls.from_dict(data)
|