kazeflow 0.1.0a0__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.
- kazeflow-0.1.0a0/PKG-INFO +116 -0
- kazeflow-0.1.0a0/README.md +106 -0
- kazeflow-0.1.0a0/pyproject.toml +28 -0
- kazeflow-0.1.0a0/src/kazeflow/__init__.py +0 -0
- kazeflow-0.1.0a0/src/kazeflow/assets.py +48 -0
- kazeflow-0.1.0a0/src/kazeflow/context.py +10 -0
- kazeflow-0.1.0a0/src/kazeflow/flow.py +156 -0
- kazeflow-0.1.0a0/src/kazeflow/graph.py +29 -0
- kazeflow-0.1.0a0/src/kazeflow/logger.py +23 -0
- kazeflow-0.1.0a0/src/kazeflow/play.py +89 -0
- kazeflow-0.1.0a0/src/kazeflow/py.typed +0 -0
- kazeflow-0.1.0a0/src/kazeflow/result.py +14 -0
- kazeflow-0.1.0a0/src/kazeflow/tui.py +116 -0
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
|
+
Name: kazeflow
|
|
3
|
+
Version: 0.1.0a0
|
|
4
|
+
Summary: A lightweight, asset-based task flow engine
|
|
5
|
+
Author: Kj
|
|
6
|
+
Author-email: Kj <46800811+kj-9@users.noreply.github.com>
|
|
7
|
+
Requires-Dist: rich>=14.1.0
|
|
8
|
+
Requires-Python: >=3.10
|
|
9
|
+
Description-Content-Type: text/markdown
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
# kazeflow
|
|
13
|
+
|
|
14
|
+
`kazeflow` is a lightweight, asset-based task flow engine inspired by Dagster. It is designed to be simple, flexible, and easy to use.
|
|
15
|
+
|
|
16
|
+
## Example
|
|
17
|
+
|
|
18
|
+
When you run this script, `kazeflow` will execute the assets in the correct order, handle the failure of `failing_asset` gracefully, and provide a rich terminal UI to visualize the progress.
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
Here is a simple example of how to use `kazeflow` to define and execute a data flow with dependencies, inputs/outputs, and logging.
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
example.py:
|
|
25
|
+
```python
|
|
26
|
+
import asyncio
|
|
27
|
+
|
|
28
|
+
from kazeflow.assets import asset
|
|
29
|
+
from kazeflow.context import AssetContext
|
|
30
|
+
from kazeflow.flow import Flow
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
# A simple asset with no dependencies
|
|
34
|
+
@asset()
|
|
35
|
+
async def users() -> list[str]:
|
|
36
|
+
"""This asset returns a list of user names."""
|
|
37
|
+
return ["Alice", "Bob", "Charlie"]
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
# This asset depends on the `users` asset.
|
|
41
|
+
# The output of `users` is automatically passed as an argument.
|
|
42
|
+
@asset(deps=["users"])
|
|
43
|
+
async def greetings(users: list[str], context: AssetContext) -> list[str]:
|
|
44
|
+
"""This asset receives the list of users and a context object.
|
|
45
|
+
|
|
46
|
+
It uses the context to get a logger and log a message.
|
|
47
|
+
"""
|
|
48
|
+
context.logger.info(f"Generating greetings for {len(users)} users.")
|
|
49
|
+
return [f"Hello, {user}!" for user in users]
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
# This asset fails intentionally to demonstrate error handling.
|
|
53
|
+
@asset(deps=["users"])
|
|
54
|
+
async def failing_asset(users: list[str]):
|
|
55
|
+
"""This asset always fails."""
|
|
56
|
+
raise ValueError("This asset is designed to fail.")
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
if __name__ == "__main__":
|
|
60
|
+
# Define a flow that includes the final assets we want to generate.
|
|
61
|
+
# kazeflow automatically includes all upstream dependencies.
|
|
62
|
+
flow = Flow(asset_names=["greetings", "failing_asset"])
|
|
63
|
+
|
|
64
|
+
# Run the flow asynchronously.
|
|
65
|
+
# You can limit the number of concurrent assets with `max_concurrency`.
|
|
66
|
+
asyncio.run(flow.run_async(max_concurrency=2))
|
|
67
|
+
|
|
68
|
+
```
|
|
69
|
+
```bash
|
|
70
|
+
❯ uv run example.py
|
|
71
|
+
Task Flow (Execution Order)
|
|
72
|
+
└── users
|
|
73
|
+
├── failing_asset
|
|
74
|
+
└── greetings
|
|
75
|
+
|
|
76
|
+
Execution Logs
|
|
77
|
+
INFO Executing asset: users
|
|
78
|
+
INFO Finished executing asset: users in 0.00s
|
|
79
|
+
INFO Executing asset: greetings
|
|
80
|
+
INFO Generating greetings for 3 users.
|
|
81
|
+
INFO Finished executing asset: greetings in 0.00s
|
|
82
|
+
INFO Executing asset: failing_asset
|
|
83
|
+
ERROR Error executing asset failing_asset: This asset is designed to
|
|
84
|
+
fail.
|
|
85
|
+
╭───────────── Traceback (most recent call last) ─────────────╮
|
|
86
|
+
│ /Users/kh03/work/repos/myflow/src/kazeflow/flow.py:82 in │
|
|
87
|
+
│ _execute_asset │
|
|
88
|
+
│ │
|
|
89
|
+
│ 79 │ │ │ │ input_kwargs["context"] = context │
|
|
90
|
+
│ 80 │ │ │ │
|
|
91
|
+
│ 81 │ │ │ if asyncio.iscoroutinefunction(asset_func │
|
|
92
|
+
│ ❱ 82 │ │ │ │ output = await asset_func(**input_kwa │
|
|
93
|
+
│ 83 │ │ │ else: │
|
|
94
|
+
│ 84 │ │ │ │ # Run sync function in a thread pool │
|
|
95
|
+
│ 85 │ │ │ │ loop = asyncio.get_running_loop() │
|
|
96
|
+
│ │
|
|
97
|
+
│ /Users/kh03/work/repos/myflow/example.py:31 in │
|
|
98
|
+
│ failing_asset │
|
|
99
|
+
│ │
|
|
100
|
+
│ 28 @asset(deps=["users"]) │
|
|
101
|
+
│ 29 async def failing_asset(users: list[str]): │
|
|
102
|
+
│ 30 │ """This asset always fails.""" │
|
|
103
|
+
│ ❱ 31 │ raise ValueError("This asset is designed to fail." │
|
|
104
|
+
│ 32 │
|
|
105
|
+
│ 33 │
|
|
106
|
+
│ 34 if __name__ == "__main__": │
|
|
107
|
+
╰─────────────────────────────────────────────────────────────╯
|
|
108
|
+
ValueError: This asset is designed to fail.
|
|
109
|
+
╭─────────────────────────────── Assets ───────────────────────────────╮
|
|
110
|
+
│ ✓ users (0.00s) │
|
|
111
|
+
│ ✓ greetings (0.00s) │
|
|
112
|
+
│ ✗ failing_asset (0.04s) │
|
|
113
|
+
╰──────────────────────────────────────────────────────────────────────╯
|
|
114
|
+
Overall Progress ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 3/3 0:00:00
|
|
115
|
+
|
|
116
|
+
```
|
|
@@ -0,0 +1,106 @@
|
|
|
1
|
+
|
|
2
|
+
# kazeflow
|
|
3
|
+
|
|
4
|
+
`kazeflow` is a lightweight, asset-based task flow engine inspired by Dagster. It is designed to be simple, flexible, and easy to use.
|
|
5
|
+
|
|
6
|
+
## Example
|
|
7
|
+
|
|
8
|
+
When you run this script, `kazeflow` will execute the assets in the correct order, handle the failure of `failing_asset` gracefully, and provide a rich terminal UI to visualize the progress.
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
Here is a simple example of how to use `kazeflow` to define and execute a data flow with dependencies, inputs/outputs, and logging.
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
example.py:
|
|
15
|
+
```python
|
|
16
|
+
import asyncio
|
|
17
|
+
|
|
18
|
+
from kazeflow.assets import asset
|
|
19
|
+
from kazeflow.context import AssetContext
|
|
20
|
+
from kazeflow.flow import Flow
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
# A simple asset with no dependencies
|
|
24
|
+
@asset()
|
|
25
|
+
async def users() -> list[str]:
|
|
26
|
+
"""This asset returns a list of user names."""
|
|
27
|
+
return ["Alice", "Bob", "Charlie"]
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
# This asset depends on the `users` asset.
|
|
31
|
+
# The output of `users` is automatically passed as an argument.
|
|
32
|
+
@asset(deps=["users"])
|
|
33
|
+
async def greetings(users: list[str], context: AssetContext) -> list[str]:
|
|
34
|
+
"""This asset receives the list of users and a context object.
|
|
35
|
+
|
|
36
|
+
It uses the context to get a logger and log a message.
|
|
37
|
+
"""
|
|
38
|
+
context.logger.info(f"Generating greetings for {len(users)} users.")
|
|
39
|
+
return [f"Hello, {user}!" for user in users]
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
# This asset fails intentionally to demonstrate error handling.
|
|
43
|
+
@asset(deps=["users"])
|
|
44
|
+
async def failing_asset(users: list[str]):
|
|
45
|
+
"""This asset always fails."""
|
|
46
|
+
raise ValueError("This asset is designed to fail.")
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
if __name__ == "__main__":
|
|
50
|
+
# Define a flow that includes the final assets we want to generate.
|
|
51
|
+
# kazeflow automatically includes all upstream dependencies.
|
|
52
|
+
flow = Flow(asset_names=["greetings", "failing_asset"])
|
|
53
|
+
|
|
54
|
+
# Run the flow asynchronously.
|
|
55
|
+
# You can limit the number of concurrent assets with `max_concurrency`.
|
|
56
|
+
asyncio.run(flow.run_async(max_concurrency=2))
|
|
57
|
+
|
|
58
|
+
```
|
|
59
|
+
```bash
|
|
60
|
+
❯ uv run example.py
|
|
61
|
+
Task Flow (Execution Order)
|
|
62
|
+
└── users
|
|
63
|
+
├── failing_asset
|
|
64
|
+
└── greetings
|
|
65
|
+
|
|
66
|
+
Execution Logs
|
|
67
|
+
INFO Executing asset: users
|
|
68
|
+
INFO Finished executing asset: users in 0.00s
|
|
69
|
+
INFO Executing asset: greetings
|
|
70
|
+
INFO Generating greetings for 3 users.
|
|
71
|
+
INFO Finished executing asset: greetings in 0.00s
|
|
72
|
+
INFO Executing asset: failing_asset
|
|
73
|
+
ERROR Error executing asset failing_asset: This asset is designed to
|
|
74
|
+
fail.
|
|
75
|
+
╭───────────── Traceback (most recent call last) ─────────────╮
|
|
76
|
+
│ /Users/kh03/work/repos/myflow/src/kazeflow/flow.py:82 in │
|
|
77
|
+
│ _execute_asset │
|
|
78
|
+
│ │
|
|
79
|
+
│ 79 │ │ │ │ input_kwargs["context"] = context │
|
|
80
|
+
│ 80 │ │ │ │
|
|
81
|
+
│ 81 │ │ │ if asyncio.iscoroutinefunction(asset_func │
|
|
82
|
+
│ ❱ 82 │ │ │ │ output = await asset_func(**input_kwa │
|
|
83
|
+
│ 83 │ │ │ else: │
|
|
84
|
+
│ 84 │ │ │ │ # Run sync function in a thread pool │
|
|
85
|
+
│ 85 │ │ │ │ loop = asyncio.get_running_loop() │
|
|
86
|
+
│ │
|
|
87
|
+
│ /Users/kh03/work/repos/myflow/example.py:31 in │
|
|
88
|
+
│ failing_asset │
|
|
89
|
+
│ │
|
|
90
|
+
│ 28 @asset(deps=["users"]) │
|
|
91
|
+
│ 29 async def failing_asset(users: list[str]): │
|
|
92
|
+
│ 30 │ """This asset always fails.""" │
|
|
93
|
+
│ ❱ 31 │ raise ValueError("This asset is designed to fail." │
|
|
94
|
+
│ 32 │
|
|
95
|
+
│ 33 │
|
|
96
|
+
│ 34 if __name__ == "__main__": │
|
|
97
|
+
╰─────────────────────────────────────────────────────────────╯
|
|
98
|
+
ValueError: This asset is designed to fail.
|
|
99
|
+
╭─────────────────────────────── Assets ───────────────────────────────╮
|
|
100
|
+
│ ✓ users (0.00s) │
|
|
101
|
+
│ ✓ greetings (0.00s) │
|
|
102
|
+
│ ✗ failing_asset (0.04s) │
|
|
103
|
+
╰──────────────────────────────────────────────────────────────────────╯
|
|
104
|
+
Overall Progress ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 3/3 0:00:00
|
|
105
|
+
|
|
106
|
+
```
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "kazeflow"
|
|
3
|
+
version = "0.1.0a0"
|
|
4
|
+
description = "A lightweight, asset-based task flow engine"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
authors = [
|
|
7
|
+
{ name = "Kj", email = "46800811+kj-9@users.noreply.github.com" }
|
|
8
|
+
]
|
|
9
|
+
requires-python = ">=3.10"
|
|
10
|
+
dependencies = [
|
|
11
|
+
"rich>=14.1.0",
|
|
12
|
+
]
|
|
13
|
+
|
|
14
|
+
[build-system]
|
|
15
|
+
requires = ["uv_build>=0.8.17,<0.9.0"]
|
|
16
|
+
build-backend = "uv_build"
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
[tool.ruff.format]
|
|
20
|
+
quote-style = "double"
|
|
21
|
+
|
|
22
|
+
[dependency-groups]
|
|
23
|
+
dev = [
|
|
24
|
+
"ruff>=0.13.0",
|
|
25
|
+
"pytest>=8.4.2",
|
|
26
|
+
"ty>=0.0.1a20",
|
|
27
|
+
"pytest-asyncio>=0.24.0",
|
|
28
|
+
]
|
|
File without changes
|
|
@@ -0,0 +1,48 @@
|
|
|
1
|
+
from typing import Any, Optional, Protocol, TypedDict
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
class NamedCallable(Protocol):
|
|
5
|
+
__name__: str
|
|
6
|
+
|
|
7
|
+
def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class AssetMeta(TypedDict):
|
|
11
|
+
func: NamedCallable
|
|
12
|
+
deps: list[str]
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
_assets: dict[str, AssetMeta] = {}
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def asset(
|
|
19
|
+
deps: Optional[list[str]] = None,
|
|
20
|
+
):
|
|
21
|
+
"""
|
|
22
|
+
A decorator to define an asset, its dependencies, and its configuration schema.
|
|
23
|
+
"""
|
|
24
|
+
|
|
25
|
+
def decorator(func: NamedCallable) -> NamedCallable:
|
|
26
|
+
_assets[func.__name__] = {
|
|
27
|
+
"func": func,
|
|
28
|
+
"deps": deps or [],
|
|
29
|
+
}
|
|
30
|
+
return func
|
|
31
|
+
|
|
32
|
+
return decorator
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def get_asset(name: str) -> AssetMeta:
|
|
36
|
+
"""Retrieves an asset's metadata including its function, dependencies,
|
|
37
|
+
and config schema."""
|
|
38
|
+
if name not in _assets:
|
|
39
|
+
raise ValueError(f"Asset '{name}' not found.")
|
|
40
|
+
return _assets[name]
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def clear_assets() -> None:
|
|
44
|
+
"""Clears all registered assets.
|
|
45
|
+
|
|
46
|
+
This is useful for testing purposes.
|
|
47
|
+
"""
|
|
48
|
+
_assets.clear()
|
|
@@ -0,0 +1,156 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
import inspect
|
|
3
|
+
import logging
|
|
4
|
+
import time
|
|
5
|
+
from graphlib import TopologicalSorter
|
|
6
|
+
from typing import Any, Optional
|
|
7
|
+
|
|
8
|
+
from .assets import get_asset
|
|
9
|
+
from .context import AssetContext
|
|
10
|
+
from .result import AssetResult
|
|
11
|
+
from .tui import FlowTUIRenderer, show_flow_tree
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class Flow:
|
|
15
|
+
"""A class representing a workflow of assets."""
|
|
16
|
+
|
|
17
|
+
def __init__(self, asset_names: list[str]):
|
|
18
|
+
self.asset_names = asset_names
|
|
19
|
+
self.asset_outputs: dict[str, Any] = {}
|
|
20
|
+
|
|
21
|
+
self.graph = self._get_graph()
|
|
22
|
+
|
|
23
|
+
ts = self._get_ts()
|
|
24
|
+
self.static_order = list(ts.static_order())
|
|
25
|
+
|
|
26
|
+
def _get_graph(self) -> dict[str, set[str]]:
|
|
27
|
+
"""Sets up the graph based on asset dependencies."""
|
|
28
|
+
graph: dict[str, set[str]] = {}
|
|
29
|
+
|
|
30
|
+
queue = list(self.asset_names)
|
|
31
|
+
visited = set()
|
|
32
|
+
|
|
33
|
+
while queue:
|
|
34
|
+
asset_name = queue.pop(0)
|
|
35
|
+
if asset_name in visited:
|
|
36
|
+
continue
|
|
37
|
+
visited.add(asset_name)
|
|
38
|
+
|
|
39
|
+
asset = get_asset(asset_name)
|
|
40
|
+
deps = set(asset["deps"])
|
|
41
|
+
graph[asset_name] = deps
|
|
42
|
+
|
|
43
|
+
for dep in deps:
|
|
44
|
+
queue.append(dep)
|
|
45
|
+
|
|
46
|
+
return graph
|
|
47
|
+
|
|
48
|
+
def _get_ts(self) -> TopologicalSorter:
|
|
49
|
+
"""Sets up the topological sorter based on asset dependencies."""
|
|
50
|
+
|
|
51
|
+
return TopologicalSorter(self.graph)
|
|
52
|
+
|
|
53
|
+
async def _execute_asset(
|
|
54
|
+
self, asset_name: str, logger: logging.Logger
|
|
55
|
+
) -> AssetResult:
|
|
56
|
+
"""Executes an asset and returns a result object."""
|
|
57
|
+
start_time = time.monotonic()
|
|
58
|
+
output = None
|
|
59
|
+
exception = None
|
|
60
|
+
success = False
|
|
61
|
+
try:
|
|
62
|
+
logger.info(f"Executing asset: {asset_name}")
|
|
63
|
+
asset = get_asset(asset_name)
|
|
64
|
+
|
|
65
|
+
asset_func = asset["func"]
|
|
66
|
+
deps = asset["deps"]
|
|
67
|
+
|
|
68
|
+
# Only pass outputs that are actual parameters of the asset function
|
|
69
|
+
sig = inspect.signature(asset_func)
|
|
70
|
+
params = sig.parameters
|
|
71
|
+
input_kwargs = {
|
|
72
|
+
dep: self.asset_outputs[dep]
|
|
73
|
+
for dep in deps
|
|
74
|
+
if dep in self.asset_outputs and dep in params
|
|
75
|
+
}
|
|
76
|
+
|
|
77
|
+
if "context" in params:
|
|
78
|
+
context = AssetContext(logger=logger, asset_name=asset_name)
|
|
79
|
+
input_kwargs["context"] = context
|
|
80
|
+
|
|
81
|
+
if asyncio.iscoroutinefunction(asset_func):
|
|
82
|
+
output = await asset_func(**input_kwargs)
|
|
83
|
+
else:
|
|
84
|
+
# Run sync function in a thread pool executor
|
|
85
|
+
loop = asyncio.get_running_loop()
|
|
86
|
+
import functools
|
|
87
|
+
|
|
88
|
+
p = functools.partial(asset_func, **input_kwargs)
|
|
89
|
+
output = await loop.run_in_executor(None, p)
|
|
90
|
+
success = True
|
|
91
|
+
|
|
92
|
+
except Exception as e:
|
|
93
|
+
exception = e
|
|
94
|
+
logger.exception(f"Error executing asset {asset_name}: {e}")
|
|
95
|
+
|
|
96
|
+
duration = time.monotonic() - start_time
|
|
97
|
+
if success:
|
|
98
|
+
logger.info(f"Finished executing asset: {asset_name} in {duration:.2f}s")
|
|
99
|
+
|
|
100
|
+
return AssetResult(
|
|
101
|
+
name=asset_name,
|
|
102
|
+
success=success,
|
|
103
|
+
duration=duration,
|
|
104
|
+
start_time=start_time,
|
|
105
|
+
output=output,
|
|
106
|
+
exception=exception,
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
async def run_async(
|
|
110
|
+
self,
|
|
111
|
+
config: Optional[dict[str, Any]] = None,
|
|
112
|
+
max_concurrency: Optional[int] = None,
|
|
113
|
+
) -> None:
|
|
114
|
+
"""Executes the assets in the flow asynchronously with a concurrency limit."""
|
|
115
|
+
if max_concurrency is not None and max_concurrency <= 0:
|
|
116
|
+
raise ValueError("max_concurrency must be a positive integer or None.")
|
|
117
|
+
|
|
118
|
+
show_flow_tree(self.graph)
|
|
119
|
+
ts = self._get_ts()
|
|
120
|
+
ts.prepare()
|
|
121
|
+
|
|
122
|
+
tui = FlowTUIRenderer(total_assets=len(self.static_order))
|
|
123
|
+
running_tasks_map: dict[asyncio.Task, tuple[str, int]] = {}
|
|
124
|
+
|
|
125
|
+
with tui:
|
|
126
|
+
while ts.is_active():
|
|
127
|
+
ready_to_run = list(ts.get_ready())
|
|
128
|
+
|
|
129
|
+
limit = max_concurrency if max_concurrency is not None else float("inf")
|
|
130
|
+
|
|
131
|
+
# Start new tasks if we have capacity and there are tasks ready to run
|
|
132
|
+
while len(running_tasks_map) < limit and ready_to_run:
|
|
133
|
+
asset_name = ready_to_run.pop(0)
|
|
134
|
+
progress_task_id = tui.add_running_task(asset_name)
|
|
135
|
+
async_task = asyncio.create_task(
|
|
136
|
+
self._execute_asset(asset_name, tui.logger)
|
|
137
|
+
)
|
|
138
|
+
running_tasks_map[async_task] = (asset_name, progress_task_id)
|
|
139
|
+
|
|
140
|
+
if not running_tasks_map:
|
|
141
|
+
break # Nothing running, nothing new to start
|
|
142
|
+
|
|
143
|
+
done, _ = await asyncio.wait(
|
|
144
|
+
running_tasks_map.keys(), return_when=asyncio.FIRST_COMPLETED
|
|
145
|
+
)
|
|
146
|
+
|
|
147
|
+
for task in done:
|
|
148
|
+
asset_name, progress_task_id = running_tasks_map.pop(task)
|
|
149
|
+
|
|
150
|
+
asset_result: AssetResult = task.result()
|
|
151
|
+
|
|
152
|
+
tui.complete_running_task(progress_task_id, asset_result)
|
|
153
|
+
|
|
154
|
+
if asset_result.success:
|
|
155
|
+
self.asset_outputs[asset_name] = asset_result.output
|
|
156
|
+
ts.done(asset_name)
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
from rich.console import Console
|
|
2
|
+
from rich.panel import Panel
|
|
3
|
+
|
|
4
|
+
console = Console()
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
def draw_simple_dag():
|
|
8
|
+
# ノード
|
|
9
|
+
a = Panel("A", style="green", expand=False)
|
|
10
|
+
b = Panel("B", style="yellow", expand=False)
|
|
11
|
+
c = Panel("C", style="yellow", expand=False)
|
|
12
|
+
d = Panel("D", style="grey50", expand=False)
|
|
13
|
+
|
|
14
|
+
lines = [
|
|
15
|
+
f" {a}",
|
|
16
|
+
" │",
|
|
17
|
+
" ▼",
|
|
18
|
+
f"{b} {c}",
|
|
19
|
+
" │ │",
|
|
20
|
+
" └──┬──┘",
|
|
21
|
+
" ▼",
|
|
22
|
+
f" {d}",
|
|
23
|
+
]
|
|
24
|
+
|
|
25
|
+
for line in lines:
|
|
26
|
+
console.print(line)
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
draw_simple_dag()
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
from typing import Optional
|
|
3
|
+
|
|
4
|
+
from rich.console import Console
|
|
5
|
+
from rich.logging import RichHandler
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def get_logger(
|
|
9
|
+
name: str = __name__, console: Optional[Console] = None
|
|
10
|
+
) -> logging.Logger:
|
|
11
|
+
"""Get a logger that uses rich for pretty printing."""
|
|
12
|
+
logger = logging.getLogger(name)
|
|
13
|
+
# Clear existing handlers to allow for reconfiguration
|
|
14
|
+
if logger.hasHandlers():
|
|
15
|
+
logger.handlers.clear()
|
|
16
|
+
|
|
17
|
+
logger.setLevel(logging.INFO)
|
|
18
|
+
handler = RichHandler(
|
|
19
|
+
show_path=False, rich_tracebacks=True, console=console, show_time=False
|
|
20
|
+
)
|
|
21
|
+
logger.addHandler(handler)
|
|
22
|
+
logger.propagate = False
|
|
23
|
+
return logger
|
|
@@ -0,0 +1,89 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
import time
|
|
3
|
+
from graphlib import TopologicalSorter
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
async def task(name, delay=1, records: list = []):
|
|
7
|
+
start = time.time()
|
|
8
|
+
print(f"start {name}")
|
|
9
|
+
await asyncio.sleep(delay)
|
|
10
|
+
end = time.time()
|
|
11
|
+
print(f"done {name}")
|
|
12
|
+
|
|
13
|
+
records.append((name, start, end))
|
|
14
|
+
return name
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
dag = {
|
|
18
|
+
"A": [],
|
|
19
|
+
"B": ["A"],
|
|
20
|
+
"C": ["A"],
|
|
21
|
+
"D": ["B", "C"],
|
|
22
|
+
"E": ["C"],
|
|
23
|
+
"F": ["C"],
|
|
24
|
+
}
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
async def run_tasks():
|
|
28
|
+
ts = TopologicalSorter(dag)
|
|
29
|
+
ts.prepare()
|
|
30
|
+
|
|
31
|
+
records = []
|
|
32
|
+
running = {
|
|
33
|
+
asyncio.create_task(task(n, delay=1 + len(n), records=records)): n
|
|
34
|
+
for n in ts.get_ready()
|
|
35
|
+
}
|
|
36
|
+
|
|
37
|
+
while running:
|
|
38
|
+
done, _ = await asyncio.wait(
|
|
39
|
+
running.keys(), return_when=asyncio.FIRST_COMPLETED
|
|
40
|
+
)
|
|
41
|
+
|
|
42
|
+
for d in done:
|
|
43
|
+
name = running.pop(d)
|
|
44
|
+
ts.done(name)
|
|
45
|
+
for new in ts.get_ready():
|
|
46
|
+
running[
|
|
47
|
+
asyncio.create_task(task(new, delay=1 + len(new), records=records))
|
|
48
|
+
] = new
|
|
49
|
+
|
|
50
|
+
return records
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
def text_gantt(records):
|
|
54
|
+
base = min(start for _, start, _ in records)
|
|
55
|
+
end_time = max(end for _, _, end in records)
|
|
56
|
+
length = int(end_time - base) + 1 # 秒単位
|
|
57
|
+
|
|
58
|
+
tasks = sorted(set(name for name, _, _ in records))
|
|
59
|
+
|
|
60
|
+
lines = []
|
|
61
|
+
for name in tasks:
|
|
62
|
+
line = [" "] * length
|
|
63
|
+
for n, start, end in records:
|
|
64
|
+
if n == name:
|
|
65
|
+
s = int(start - base)
|
|
66
|
+
e = int(end - base)
|
|
67
|
+
for i in range(s, e):
|
|
68
|
+
line[i] = "#"
|
|
69
|
+
lines.append(f"{name:>3} |{''.join(line)}|")
|
|
70
|
+
return "\n".join(lines)
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def text_graph(dag):
|
|
74
|
+
lines = []
|
|
75
|
+
for node, deps in dag.items():
|
|
76
|
+
if deps:
|
|
77
|
+
for dep in deps:
|
|
78
|
+
lines.append(f"{dep} --> {node}")
|
|
79
|
+
else:
|
|
80
|
+
lines.append(f"{node}")
|
|
81
|
+
return "\n".join(lines)
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
if __name__ == "__main__":
|
|
85
|
+
print("\nDependency Graph (text-based):")
|
|
86
|
+
print(text_graph(dag))
|
|
87
|
+
records = asyncio.run(run_tasks())
|
|
88
|
+
print("\nGantt Chart (text-based):")
|
|
89
|
+
print(text_gantt(records))
|
|
File without changes
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
from dataclasses import dataclass
|
|
2
|
+
from typing import Any, Optional
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
@dataclass
|
|
6
|
+
class AssetResult:
|
|
7
|
+
"""Holds the result of a single asset's execution."""
|
|
8
|
+
|
|
9
|
+
name: str
|
|
10
|
+
success: bool
|
|
11
|
+
duration: float
|
|
12
|
+
start_time: float
|
|
13
|
+
output: Optional[Any] = None
|
|
14
|
+
exception: Optional[Exception] = None
|
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
from typing import Any, Dict, Set
|
|
2
|
+
|
|
3
|
+
from rich.console import Console, Group
|
|
4
|
+
from rich.live import Live
|
|
5
|
+
from rich.panel import Panel
|
|
6
|
+
from rich.progress import (
|
|
7
|
+
BarColumn,
|
|
8
|
+
MofNCompleteColumn,
|
|
9
|
+
Progress,
|
|
10
|
+
SpinnerColumn,
|
|
11
|
+
TextColumn,
|
|
12
|
+
TimeElapsedColumn,
|
|
13
|
+
)
|
|
14
|
+
from rich.tree import Tree
|
|
15
|
+
|
|
16
|
+
from .logger import get_logger
|
|
17
|
+
from .result import AssetResult
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def show_flow_tree(graph: Dict[str, Set[str]]) -> None:
|
|
21
|
+
"""Displays the task flow as a rich tree, in execution order."""
|
|
22
|
+
|
|
23
|
+
# Reverse the graph to show data flow from dependencies to dependents
|
|
24
|
+
reversed_graph = {node: set() for node in graph}
|
|
25
|
+
for node, deps in graph.items():
|
|
26
|
+
for dep in deps:
|
|
27
|
+
if dep in reversed_graph:
|
|
28
|
+
reversed_graph[dep].add(node)
|
|
29
|
+
|
|
30
|
+
# Find root nodes (assets with no dependencies)
|
|
31
|
+
root_nodes = [node for node, deps in graph.items() if not deps]
|
|
32
|
+
|
|
33
|
+
tree = Tree("[bold green]Task Flow (Execution Order)[/bold green]")
|
|
34
|
+
added_nodes = set()
|
|
35
|
+
|
|
36
|
+
def add_to_tree(parent_tree: Tree, node_name: str) -> None:
|
|
37
|
+
if node_name in added_nodes:
|
|
38
|
+
return
|
|
39
|
+
added_nodes.add(node_name)
|
|
40
|
+
node_tree = parent_tree.add(node_name)
|
|
41
|
+
# Use reversed_graph to find nodes that depend on the current one
|
|
42
|
+
for dependent_node in sorted(list(reversed_graph.get(node_name, []))):
|
|
43
|
+
add_to_tree(node_tree, dependent_node)
|
|
44
|
+
|
|
45
|
+
for root in sorted(root_nodes):
|
|
46
|
+
add_to_tree(tree, root)
|
|
47
|
+
|
|
48
|
+
Console().print(tree)
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
class FlowTUIRenderer:
|
|
52
|
+
"""Manages the Text-based User Interface for flow execution using rich."""
|
|
53
|
+
|
|
54
|
+
def __init__(self, total_assets: int):
|
|
55
|
+
self.completed_progress = Progress(TextColumn("✓ [green]{task.description}"))
|
|
56
|
+
self.failed_progress = Progress(TextColumn("✗ [red]{task.description}"))
|
|
57
|
+
self.running_progress = Progress(
|
|
58
|
+
TextColumn(" [purple]Running: {task.description}"),
|
|
59
|
+
SpinnerColumn("simpleDots"),
|
|
60
|
+
TimeElapsedColumn(),
|
|
61
|
+
)
|
|
62
|
+
self.overall_progress = Progress(
|
|
63
|
+
TextColumn("[bold blue]Overall Progress"),
|
|
64
|
+
BarColumn(),
|
|
65
|
+
MofNCompleteColumn(),
|
|
66
|
+
TimeElapsedColumn(),
|
|
67
|
+
)
|
|
68
|
+
self.progress_group = Group(
|
|
69
|
+
Panel(
|
|
70
|
+
Group(
|
|
71
|
+
self.completed_progress,
|
|
72
|
+
self.failed_progress,
|
|
73
|
+
self.running_progress,
|
|
74
|
+
),
|
|
75
|
+
title="Assets",
|
|
76
|
+
),
|
|
77
|
+
self.overall_progress,
|
|
78
|
+
)
|
|
79
|
+
self.overall_task_id = self.overall_progress.add_task(
|
|
80
|
+
"Assets", total=total_assets
|
|
81
|
+
)
|
|
82
|
+
self.live = Live(self.progress_group)
|
|
83
|
+
self.logger = get_logger(__name__, console=self.live.console)
|
|
84
|
+
self.results: list[AssetResult] = []
|
|
85
|
+
|
|
86
|
+
def __enter__(self) -> Live:
|
|
87
|
+
Console().print("\n[bold underline green]Execution Logs[/bold underline green]")
|
|
88
|
+
return self.live.__enter__()
|
|
89
|
+
|
|
90
|
+
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
91
|
+
# Sort results by start time before final render
|
|
92
|
+
self.results.sort(key=lambda r: r.start_time)
|
|
93
|
+
|
|
94
|
+
# Now, populate the final progress bars
|
|
95
|
+
for result in self.results:
|
|
96
|
+
description = f"{result.name:<30} ({result.duration:.2f}s)"
|
|
97
|
+
if result.success:
|
|
98
|
+
self.completed_progress.add_task(description)
|
|
99
|
+
else:
|
|
100
|
+
self.failed_progress.add_task(description)
|
|
101
|
+
|
|
102
|
+
return self.live.__exit__(exc_type, exc_val, exc_tb)
|
|
103
|
+
|
|
104
|
+
def add_running_task(self, name: str) -> int:
|
|
105
|
+
"""Adds a task to the running progress bar."""
|
|
106
|
+
return self.running_progress.add_task(name, total=1)
|
|
107
|
+
|
|
108
|
+
def complete_running_task(self, task_id: int, result: AssetResult) -> None:
|
|
109
|
+
"""Moves a task from running to a temporary list for later sorting."""
|
|
110
|
+
self.running_progress.stop_task(task_id)
|
|
111
|
+
self.running_progress.update(task_id, visible=False)
|
|
112
|
+
|
|
113
|
+
# Store the result for final sorting
|
|
114
|
+
self.results.append(result)
|
|
115
|
+
|
|
116
|
+
self.overall_progress.update(self.overall_task_id, advance=1)
|