tinyfan 0.2.8__tar.gz → 0.3.0__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.
- {tinyfan-0.2.8 → tinyfan-0.3.0}/PKG-INFO +6 -7
- {tinyfan-0.2.8 → tinyfan-0.3.0}/README.md +1 -1
- {tinyfan-0.2.8 → tinyfan-0.3.0}/pyproject.toml +4 -5
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/argo_typing.py +11 -1
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/codegen.py +79 -40
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/flow.py +11 -2
- tinyfan-0.3.0/src/tinyfan/stores/gcs_pickle.py +149 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/utils/embed.py +2 -2
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/test_codegen.py +33 -0
- tinyfan-0.3.0/tests/test_gcs_pickle_store.py +71 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/uv.lock +423 -326
- tinyfan-0.2.8/src/tinyfan/stores/gcs_pickle.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/.github/workflows/publish.yml +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/.gitignore +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/.python-version +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/examples/README.md +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/examples/configmap.yaml +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/examples/configs.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/examples/getting-started.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/examples/packaged/.python-version +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/examples/packaged/Dockerfile +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/examples/packaged/README.md +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/examples/packaged/pyproject.toml +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/examples/packaged/src/tinyfan_examples_packaged/__init__.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/examples/packaged/src/tinyfan_examples_packaged/asset.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/examples/resource.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/examples/secret.yaml +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/__init__.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/cli.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/config.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/dagviz.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/flowrundata.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/py.typed +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/resources/__init__.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/resources/base.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/stores/__init__.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/stores/base.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/stores/bigquery.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/stores/naive.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/utils/__init__.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/utils/exjson.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/utils/merge.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/utils/typing.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/src/tinyfan/utils/yaml.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/__init__.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/codegen_samples/directory/__init__.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/codegen_samples/directory/asset.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/codegen_samples/directory/asset2.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/codegen_samples/singlefile.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/conftest.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/helpers.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/openapi2jsonschema.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/test_config.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/test_decorators.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/utils/__init__.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/utils/embed_samples/directory/__init__.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/utils/embed_samples/directory/dep.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/utils/embed_samples/singlefile.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/utils/test_embed.py +0 -0
- {tinyfan-0.2.8 → tinyfan-0.3.0}/tests/utils/test_exjson.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: tinyfan
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.3.0
|
|
4
4
|
Summary: Tiny Argo Workflows python devkit
|
|
5
5
|
Project-URL: Documentation, https://github.com/eunchuldev/tinyfan#readme
|
|
6
6
|
Project-URL: Issues, https://github.com/eunchuldev/tinyfan/issues
|
|
@@ -9,14 +9,12 @@ Author-email: eunchuldev <eunchulsong@gmail.com>
|
|
|
9
9
|
Keywords: argo-workflow,pipeline,tinyfan,workflow
|
|
10
10
|
Classifier: License :: OSI Approved :: MIT License
|
|
11
11
|
Classifier: Operating System :: OS Independent
|
|
12
|
-
Classifier: Programming Language :: Python :: 3.10
|
|
13
|
-
Classifier: Programming Language :: Python :: 3.11
|
|
14
12
|
Classifier: Programming Language :: Python :: 3.12
|
|
15
13
|
Classifier: Typing :: Typed
|
|
16
14
|
Requires-Python: >=3.12
|
|
17
|
-
Requires-Dist: croniter>=
|
|
18
|
-
Requires-Dist: pyyaml>=6.0.
|
|
19
|
-
Requires-Dist: typer-slim>=0.
|
|
15
|
+
Requires-Dist: croniter>=6
|
|
16
|
+
Requires-Dist: pyyaml>=6.0.3
|
|
17
|
+
Requires-Dist: typer-slim>=0.24.0
|
|
20
18
|
Provides-Extra: all
|
|
21
19
|
Requires-Dist: google-cloud-bigquery[pandas]; extra == 'all'
|
|
22
20
|
Requires-Dist: google-cloud-storage; extra == 'all'
|
|
@@ -40,6 +38,7 @@ Requires-Dist: pytest>=8.3.3; extra == 'dev'
|
|
|
40
38
|
Requires-Dist: ruff; extra == 'dev'
|
|
41
39
|
Requires-Dist: types-croniter; extra == 'dev'
|
|
42
40
|
Requires-Dist: types-jsonschema; extra == 'dev'
|
|
41
|
+
Requires-Dist: types-networkx; extra == 'dev'
|
|
43
42
|
Requires-Dist: types-pyyaml; extra == 'dev'
|
|
44
43
|
Provides-Extra: gcs
|
|
45
44
|
Requires-Dist: google-cloud-storage; extra == 'gcs'
|
|
@@ -99,7 +98,7 @@ def greeting(world: str):
|
|
|
99
98
|
```shell
|
|
100
99
|
# Apply the changes to argo workflow
|
|
101
100
|
|
|
102
|
-
tinyfan
|
|
101
|
+
tinyfan main.py | kubectl apply -f -
|
|
103
102
|
```
|
|
104
103
|
|
|
105
104
|
# Real World Example (still tiny though!)
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[project]
|
|
2
2
|
name = "tinyfan"
|
|
3
|
-
version = "0.
|
|
3
|
+
version = "0.3.0"
|
|
4
4
|
description = 'Tiny Argo Workflows python devkit'
|
|
5
5
|
readme = "README.md"
|
|
6
6
|
requires-python = ">=3.12"
|
|
@@ -9,12 +9,10 @@ authors = [{ name = "eunchuldev", email = "eunchulsong@gmail.com" }]
|
|
|
9
9
|
classifiers = [
|
|
10
10
|
"Operating System :: OS Independent",
|
|
11
11
|
"License :: OSI Approved :: MIT License",
|
|
12
|
-
"Programming Language :: Python :: 3.10",
|
|
13
|
-
"Programming Language :: Python :: 3.11",
|
|
14
12
|
"Programming Language :: Python :: 3.12",
|
|
15
13
|
"Typing :: Typed",
|
|
16
14
|
]
|
|
17
|
-
dependencies = ["croniter>=
|
|
15
|
+
dependencies = ["croniter>=6", "pyyaml>=6.0.3", "typer-slim>=0.24.0"]
|
|
18
16
|
|
|
19
17
|
[project.scripts]
|
|
20
18
|
tinyfan = "tinyfan.cli:main"
|
|
@@ -34,6 +32,7 @@ dev = [
|
|
|
34
32
|
"types-PyYAML",
|
|
35
33
|
"types-jsonschema",
|
|
36
34
|
"types-croniter",
|
|
35
|
+
"types-networkx",
|
|
37
36
|
"tinyfan[all]",
|
|
38
37
|
]
|
|
39
38
|
|
|
@@ -53,7 +52,7 @@ addopts = ["--import-mode=importlib"]
|
|
|
53
52
|
exclude = ['build']
|
|
54
53
|
|
|
55
54
|
[[tool.mypy.overrides]]
|
|
56
|
-
module = ["lark.*", "jsonpatch.*"]
|
|
55
|
+
module = ["lark.*", "jsonpatch.*", "google.*"]
|
|
57
56
|
ignore_missing_imports = true
|
|
58
57
|
|
|
59
58
|
[tool.ruff]
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
from typing import TypedDict, NotRequired
|
|
1
|
+
from typing import Any, TypedDict, NotRequired
|
|
2
2
|
|
|
3
3
|
|
|
4
4
|
class ExecAction(TypedDict):
|
|
@@ -194,6 +194,16 @@ class VolumeMount(TypedDict, total=False):
|
|
|
194
194
|
subPathExpr: str
|
|
195
195
|
|
|
196
196
|
|
|
197
|
+
class TemplateInputs(TypedDict, total=False):
|
|
198
|
+
artifacts: list[dict[str, Any]]
|
|
199
|
+
parameters: list[dict[str, Any]]
|
|
200
|
+
|
|
201
|
+
|
|
202
|
+
class TemplateOutputs(TypedDict, total=False):
|
|
203
|
+
artifacts: list[dict[str, Any]]
|
|
204
|
+
parameters: list[dict[str, Any]]
|
|
205
|
+
|
|
206
|
+
|
|
197
207
|
# Ref: https://argo-workflows.readthedocs.io/en/latest/fields/#scripttemplate
|
|
198
208
|
class ScriptTemplate(TypedDict, total=False):
|
|
199
209
|
args: list[str]
|
|
@@ -1,4 +1,5 @@
|
|
|
1
1
|
from tinyfan.flow import Flow, Asset, FLOW_REGISTER, DEFAULT_IMAGE
|
|
2
|
+
import copy
|
|
2
3
|
import importlib.util
|
|
3
4
|
from typing import Self
|
|
4
5
|
from .utils.yaml import dump as yaml_dump, dump_all as yaml_dump_all
|
|
@@ -11,7 +12,7 @@ import os
|
|
|
11
12
|
import pkgutil
|
|
12
13
|
import datetime
|
|
13
14
|
import croniter
|
|
14
|
-
from .argo_typing import ScriptTemplate
|
|
15
|
+
from .argo_typing import ScriptTemplate, TemplateInputs, TemplateOutputs
|
|
15
16
|
from .utils.embed import embed
|
|
16
17
|
from .utils.merge import deepmerge, dropnone
|
|
17
18
|
from .config import CLI_CONFIG_STATE, Config, ConfigValue, ConfigMapKeyRef, ConfigMapRef, SecretKeyRef, SecretRef
|
|
@@ -23,6 +24,34 @@ VALID_DEPENDS_REGEXP = r"^(?!.*\.[A-Za-z0-9_]+\.[A-Za-z0-9_]+)(?!.*\.(?!Succeede
|
|
|
23
24
|
EXTRACT_DEPENDS_REGEXP = r"(?<!\.)\b([A-Za-z0-9_]+)\b"
|
|
24
25
|
|
|
25
26
|
RUNDATA_FILE_PATH = "/tmp/tinyfan/rundata.json"
|
|
27
|
+
INTERNAL_INPUT_PARAMETERS = {"rundata", "asset_name", "moduledata"}
|
|
28
|
+
INTERNAL_OUTPUT_PARAMETERS = {"rundata"}
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def merge_template_io(template_io: dict, user_io: TemplateInputs | TemplateOutputs | None, reserved_parameter_names: set[str]) -> dict:
|
|
32
|
+
result = copy.deepcopy(template_io)
|
|
33
|
+
if user_io is None:
|
|
34
|
+
return result
|
|
35
|
+
for key, value in user_io.items():
|
|
36
|
+
if key in ("parameters", "artifacts"):
|
|
37
|
+
if not isinstance(value, list):
|
|
38
|
+
raise TypeError(f"asset {key} must be a list")
|
|
39
|
+
if key == "parameters":
|
|
40
|
+
conflicts = [
|
|
41
|
+
param.get("name")
|
|
42
|
+
for param in value
|
|
43
|
+
if isinstance(param, dict) and param.get("name") in reserved_parameter_names
|
|
44
|
+
]
|
|
45
|
+
if conflicts:
|
|
46
|
+
names = ", ".join(f"`{name}`" for name in conflicts)
|
|
47
|
+
raise ValueError(f"asset parameters cannot redefine tinyfan internal parameter(s): {names}")
|
|
48
|
+
result.setdefault(key, [])
|
|
49
|
+
result[key].extend(copy.deepcopy(value))
|
|
50
|
+
elif isinstance(result.get(key), dict) and isinstance(value, dict):
|
|
51
|
+
result[key] = deepmerge(result[key], copy.deepcopy(value))
|
|
52
|
+
else:
|
|
53
|
+
result[key] = copy.deepcopy(value)
|
|
54
|
+
return result
|
|
26
55
|
|
|
27
56
|
|
|
28
57
|
def brackets_balanced(code: str) -> bool:
|
|
@@ -257,7 +286,7 @@ class AssetTree:
|
|
|
257
286
|
"generateName": f"""{self.flow.name}-{encode_schedule_to_k8sname(tz, schedule)}-""",
|
|
258
287
|
},
|
|
259
288
|
"spec": {
|
|
260
|
-
"
|
|
289
|
+
"schedules": [schedule],
|
|
261
290
|
"timezone": tz,
|
|
262
291
|
"workflowSpec": {
|
|
263
292
|
"entrypoint": "flow",
|
|
@@ -364,32 +393,40 @@ class AssetTree:
|
|
|
364
393
|
"synchronization": {
|
|
365
394
|
"mutexes": [{"name": f"{self.flow.name}-{node.asset.name.replace('_', '-')}"}]
|
|
366
395
|
},
|
|
367
|
-
"inputs":
|
|
368
|
-
|
|
369
|
-
|
|
370
|
-
|
|
371
|
-
|
|
372
|
-
|
|
373
|
-
{
|
|
374
|
-
"name": "asset_name",
|
|
375
|
-
"value": "{{= nil }}",
|
|
376
|
-
},
|
|
377
|
-
{
|
|
378
|
-
"name": "moduledata",
|
|
379
|
-
"value": "{{= nil }}",
|
|
380
|
-
},
|
|
381
|
-
]
|
|
382
|
-
},
|
|
383
|
-
"outputs": {
|
|
384
|
-
"parameters": [
|
|
385
|
-
{
|
|
386
|
-
"name": "rundata",
|
|
387
|
-
"valueFrom": {
|
|
388
|
-
"path": RUNDATA_FILE_PATH,
|
|
396
|
+
"inputs": merge_template_io(
|
|
397
|
+
{
|
|
398
|
+
"parameters": [
|
|
399
|
+
{
|
|
400
|
+
"name": "rundata",
|
|
401
|
+
"value": "{{= nil }}",
|
|
389
402
|
},
|
|
390
|
-
|
|
391
|
-
|
|
392
|
-
|
|
403
|
+
{
|
|
404
|
+
"name": "asset_name",
|
|
405
|
+
"value": "{{= nil }}",
|
|
406
|
+
},
|
|
407
|
+
{
|
|
408
|
+
"name": "moduledata",
|
|
409
|
+
"value": "{{= nil }}",
|
|
410
|
+
},
|
|
411
|
+
]
|
|
412
|
+
},
|
|
413
|
+
node.asset.inputs,
|
|
414
|
+
INTERNAL_INPUT_PARAMETERS,
|
|
415
|
+
),
|
|
416
|
+
"outputs": merge_template_io(
|
|
417
|
+
{
|
|
418
|
+
"parameters": [
|
|
419
|
+
{
|
|
420
|
+
"name": "rundata",
|
|
421
|
+
"valueFrom": {
|
|
422
|
+
"path": RUNDATA_FILE_PATH,
|
|
423
|
+
},
|
|
424
|
+
}
|
|
425
|
+
],
|
|
426
|
+
},
|
|
427
|
+
node.asset.outputs,
|
|
428
|
+
INTERNAL_OUTPUT_PARAMETERS,
|
|
429
|
+
),
|
|
393
430
|
}
|
|
394
431
|
for node in relatives
|
|
395
432
|
]
|
|
@@ -461,17 +498,19 @@ def codegen(
|
|
|
461
498
|
serviceAccountName: str | None = None,
|
|
462
499
|
) -> str:
|
|
463
500
|
"""Generate argocd workflow resource as yaml from tinyfan definitions"""
|
|
464
|
-
|
|
465
|
-
|
|
466
|
-
|
|
467
|
-
|
|
468
|
-
|
|
469
|
-
|
|
470
|
-
|
|
471
|
-
|
|
472
|
-
|
|
473
|
-
|
|
474
|
-
|
|
475
|
-
|
|
476
|
-
|
|
501
|
+
try:
|
|
502
|
+
CLI_CONFIG_STATE.codegen = True
|
|
503
|
+
if flow:
|
|
504
|
+
result = AssetTree(flow).compile(embedded, container, serviceAccountName)
|
|
505
|
+
elif location:
|
|
506
|
+
import_all_submodules(location)
|
|
507
|
+
result = "\n---\n".join(
|
|
508
|
+
AssetTree(flow).compile(embedded, container, serviceAccountName) for flow in FLOW_REGISTER.values()
|
|
509
|
+
)
|
|
510
|
+
else:
|
|
511
|
+
result = "\n---\n".join(
|
|
512
|
+
AssetTree(flow).compile(embedded, container, serviceAccountName) for flow in FLOW_REGISTER.values()
|
|
513
|
+
)
|
|
514
|
+
finally:
|
|
515
|
+
CLI_CONFIG_STATE.codegen = False
|
|
477
516
|
return result
|
|
@@ -6,7 +6,7 @@ from .flowrundata import FlowRunData, StoreIdx, UMeta
|
|
|
6
6
|
from dataclasses import dataclass, field
|
|
7
7
|
from typing import Callable, Any, Generic, TypeVar, Awaitable, Coroutine, cast
|
|
8
8
|
from .resources.base import ResourceBase
|
|
9
|
-
from .argo_typing import ScriptTemplate
|
|
9
|
+
from .argo_typing import ScriptTemplate, TemplateInputs, TemplateOutputs
|
|
10
10
|
from .config import Config, ConfigValue
|
|
11
11
|
|
|
12
12
|
FLOW_REGISTER = {}
|
|
@@ -51,6 +51,8 @@ class Asset(Generic[Ret, UMeta, StoreIdx]):
|
|
|
51
51
|
name: str = field(init=False)
|
|
52
52
|
container: ScriptTemplate | None = None
|
|
53
53
|
serviceAccountName: str | None = None
|
|
54
|
+
inputs: TemplateInputs | None = None
|
|
55
|
+
outputs: TemplateOutputs | None = None
|
|
54
56
|
|
|
55
57
|
def __post_init__(self):
|
|
56
58
|
self.name = self.func.__name__
|
|
@@ -81,7 +83,10 @@ class Asset(Generic[Ret, UMeta, StoreIdx]):
|
|
|
81
83
|
for name, prundata in parent_flowrundatas.items():
|
|
82
84
|
parent = self.flow.assets[name]
|
|
83
85
|
index = prundata.get("store_entry_idx", None)
|
|
84
|
-
|
|
86
|
+
if index is not None:
|
|
87
|
+
data = parent.store.retrieve(index, source_rundata=prundata, target_rundata=rundata)
|
|
88
|
+
else:
|
|
89
|
+
data = None
|
|
85
90
|
if name in func_param_names:
|
|
86
91
|
params[name] = data
|
|
87
92
|
ret = self.func(**params)
|
|
@@ -118,6 +123,8 @@ def asset(
|
|
|
118
123
|
metadata: UMeta | None = None,
|
|
119
124
|
container: ScriptTemplate | None = None,
|
|
120
125
|
serviceAccountName: str | None = None,
|
|
126
|
+
inputs: TemplateInputs | None = None,
|
|
127
|
+
outputs: TemplateOutputs | None = None,
|
|
121
128
|
) -> Callable[..., AssetFunc[Ret, UMeta, StoreIdx]]:
|
|
122
129
|
def wrapper(func: Callable[..., Ret]) -> AssetFunc[Ret, UMeta, StoreIdx]:
|
|
123
130
|
asset = Asset[Ret, UMeta, StoreIdx](
|
|
@@ -130,6 +137,8 @@ def asset(
|
|
|
130
137
|
metadata=metadata,
|
|
131
138
|
container=container,
|
|
132
139
|
serviceAccountName=serviceAccountName,
|
|
140
|
+
inputs=inputs,
|
|
141
|
+
outputs=outputs,
|
|
133
142
|
)
|
|
134
143
|
return AssetFunc(func, asset)
|
|
135
144
|
|
|
@@ -0,0 +1,149 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import pickle
|
|
4
|
+
import re
|
|
5
|
+
import uuid
|
|
6
|
+
from collections.abc import Mapping
|
|
7
|
+
from typing import Any, NotRequired, TypedDict
|
|
8
|
+
|
|
9
|
+
from ..config import ConfigValue
|
|
10
|
+
from .base import FlowRunData, StoreBase
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class GcsPickleStoreIndex(TypedDict):
|
|
14
|
+
bucket: str
|
|
15
|
+
name: str
|
|
16
|
+
generation: NotRequired[int]
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class GcsPickleStoreMetadata(TypedDict, total=False):
|
|
20
|
+
blob_name: str
|
|
21
|
+
prefix: str
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
ConfigurableString = str | ConfigValue
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
_GCS_EXTRA_ERROR = "GCS extra not installed. Please install with `pip install tinyfan[gcs]`."
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def _new_storage_client(client_kwargs: Mapping[str, Any]) -> Any:
|
|
31
|
+
try:
|
|
32
|
+
from google.cloud import storage
|
|
33
|
+
except ImportError as exc:
|
|
34
|
+
raise ImportError(_GCS_EXTRA_ERROR) from exc
|
|
35
|
+
|
|
36
|
+
return storage.Client(**client_kwargs)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def _resolve(value: ConfigurableString) -> str:
|
|
40
|
+
if isinstance(value, ConfigValue):
|
|
41
|
+
return value.get_value()
|
|
42
|
+
return value
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def _path_part(value: object, fallback: str) -> str:
|
|
46
|
+
text = str(value if value is not None else fallback).strip()
|
|
47
|
+
text = re.sub(r"[^A-Za-z0-9._=-]+", "-", text)
|
|
48
|
+
return text.strip(".-") or fallback
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def _join_blob_name(*parts: str) -> str:
|
|
52
|
+
return "/".join(part.strip("/") for part in parts if part and part.strip("/"))
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def _run_part(rundata: FlowRunData[Any, Any]) -> str:
|
|
56
|
+
for key in ("data_interval_start", "ts", "ds"):
|
|
57
|
+
value = rundata.get(key)
|
|
58
|
+
if value is None:
|
|
59
|
+
continue
|
|
60
|
+
if hasattr(value, "isoformat"):
|
|
61
|
+
value = value.isoformat()
|
|
62
|
+
return _path_part(value, "manual")
|
|
63
|
+
return "manual"
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
class GcsPickleStore(StoreBase[object, GcsPickleStoreMetadata, GcsPickleStoreIndex]):
|
|
67
|
+
bucket: ConfigurableString
|
|
68
|
+
prefix: ConfigurableString
|
|
69
|
+
|
|
70
|
+
def __init__(
|
|
71
|
+
self,
|
|
72
|
+
bucket: ConfigurableString,
|
|
73
|
+
prefix: ConfigurableString = "tinyfan",
|
|
74
|
+
*,
|
|
75
|
+
client: Any | None = None,
|
|
76
|
+
client_kwargs: Mapping[str, Any] | None = None,
|
|
77
|
+
pickle_protocol: int = pickle.HIGHEST_PROTOCOL,
|
|
78
|
+
) -> None:
|
|
79
|
+
self.bucket = bucket
|
|
80
|
+
self.prefix = prefix
|
|
81
|
+
self._client = client
|
|
82
|
+
self._client_kwargs = dict(client_kwargs or {})
|
|
83
|
+
self.pickle_protocol = pickle_protocol
|
|
84
|
+
|
|
85
|
+
@staticmethod
|
|
86
|
+
def id() -> str:
|
|
87
|
+
return "tinyfan.gcs_pickle"
|
|
88
|
+
|
|
89
|
+
def _get_client(self) -> Any:
|
|
90
|
+
if self._client is None:
|
|
91
|
+
self._client = _new_storage_client(self._client_kwargs)
|
|
92
|
+
return self._client
|
|
93
|
+
|
|
94
|
+
def _bucket_name(self) -> str:
|
|
95
|
+
bucket = _resolve(self.bucket).strip()
|
|
96
|
+
if not bucket:
|
|
97
|
+
raise ValueError("GCS bucket name must not be empty.")
|
|
98
|
+
return bucket
|
|
99
|
+
|
|
100
|
+
def _blob_name(self, rundata: FlowRunData[GcsPickleStoreMetadata, GcsPickleStoreIndex]) -> str:
|
|
101
|
+
metadata = rundata.get("metadata")
|
|
102
|
+
if isinstance(metadata, Mapping):
|
|
103
|
+
blob_name = metadata.get("blob_name")
|
|
104
|
+
if blob_name is not None:
|
|
105
|
+
blob_name = str(blob_name).strip("/")
|
|
106
|
+
if not blob_name:
|
|
107
|
+
raise ValueError("GCS blob_name metadata must not be empty.")
|
|
108
|
+
return blob_name
|
|
109
|
+
|
|
110
|
+
prefix = str(metadata.get("prefix") or _resolve(self.prefix))
|
|
111
|
+
else:
|
|
112
|
+
prefix = _resolve(self.prefix)
|
|
113
|
+
|
|
114
|
+
flow_name = _path_part(rundata.get("flow_name"), "flow")
|
|
115
|
+
asset_name = _path_part(rundata.get("asset_name"), "asset")
|
|
116
|
+
run_id = _run_part(rundata)
|
|
117
|
+
filename = f"{uuid.uuid4().hex}.pickle"
|
|
118
|
+
return _join_blob_name(prefix, flow_name, asset_name, run_id, filename)
|
|
119
|
+
|
|
120
|
+
def store(
|
|
121
|
+
self,
|
|
122
|
+
value: object,
|
|
123
|
+
rundata: FlowRunData[GcsPickleStoreMetadata, GcsPickleStoreIndex],
|
|
124
|
+
) -> GcsPickleStoreIndex:
|
|
125
|
+
bucket_name = self._bucket_name()
|
|
126
|
+
blob = self._get_client().bucket(bucket_name).blob(self._blob_name(rundata))
|
|
127
|
+
data = pickle.dumps(value, protocol=self.pickle_protocol)
|
|
128
|
+
blob.upload_from_string(data, content_type="application/octet-stream")
|
|
129
|
+
|
|
130
|
+
index: GcsPickleStoreIndex = {
|
|
131
|
+
"bucket": bucket_name,
|
|
132
|
+
"name": blob.name,
|
|
133
|
+
}
|
|
134
|
+
generation = getattr(blob, "generation", None)
|
|
135
|
+
if generation is not None:
|
|
136
|
+
index["generation"] = int(generation)
|
|
137
|
+
return index
|
|
138
|
+
|
|
139
|
+
def retrieve(
|
|
140
|
+
self,
|
|
141
|
+
index: GcsPickleStoreIndex,
|
|
142
|
+
source_rundata: FlowRunData[GcsPickleStoreMetadata, GcsPickleStoreIndex],
|
|
143
|
+
target_rundata: FlowRunData,
|
|
144
|
+
) -> object:
|
|
145
|
+
del source_rundata, target_rundata
|
|
146
|
+
|
|
147
|
+
bucket = self._get_client().bucket(index["bucket"])
|
|
148
|
+
blob = bucket.blob(index["name"], generation=index.get("generation"))
|
|
149
|
+
return pickle.loads(blob.download_as_bytes())
|
|
@@ -54,7 +54,7 @@ def zip_directory(
|
|
|
54
54
|
files -= set().union(*map(path.glob, excludes))
|
|
55
55
|
for file_path in map(str, files):
|
|
56
56
|
arcname = os.path.relpath(file_path, start=parent_dir)
|
|
57
|
-
with open(file_path, "r") as f:
|
|
57
|
+
with open(file_path, "r", encoding="utf-8") as f:
|
|
58
58
|
try:
|
|
59
59
|
source = f.read()
|
|
60
60
|
except Exception as e:
|
|
@@ -67,7 +67,7 @@ def zip_directory(
|
|
|
67
67
|
|
|
68
68
|
|
|
69
69
|
def embed_singlefile(path: str) -> str:
|
|
70
|
-
with open(path, "r") as f:
|
|
70
|
+
with open(path, "r", encoding="utf-8") as f:
|
|
71
71
|
sourcecode = f.read()
|
|
72
72
|
return sourcecode
|
|
73
73
|
|
|
@@ -5,6 +5,7 @@ from jsonschema.exceptions import ValidationError
|
|
|
5
5
|
from typing import Generator
|
|
6
6
|
import os
|
|
7
7
|
import pytest
|
|
8
|
+
import yaml
|
|
8
9
|
|
|
9
10
|
|
|
10
11
|
# type: ignore
|
|
@@ -103,6 +104,38 @@ def test_codegen_flow(validate_crds, sample_flow):
|
|
|
103
104
|
assert_code_is_running(actual)
|
|
104
105
|
|
|
105
106
|
|
|
107
|
+
def test_codegen_asset_inputs_outputs(validate_crds):
|
|
108
|
+
flow: Flow = Flow("test-io")
|
|
109
|
+
|
|
110
|
+
@asset(
|
|
111
|
+
flow=flow,
|
|
112
|
+
schedule="@daily",
|
|
113
|
+
inputs={"parameters": [{"name": "message", "value": "hello"}]},
|
|
114
|
+
outputs={"parameters": [{"name": "summary", "valueFrom": {"path": "/tmp/summary.txt"}}]},
|
|
115
|
+
)
|
|
116
|
+
def asset1(): # type: ignore
|
|
117
|
+
return "hello"
|
|
118
|
+
|
|
119
|
+
try:
|
|
120
|
+
actual = codegen(flow=flow)
|
|
121
|
+
validate_crds(actual)
|
|
122
|
+
manifest = next(yaml.load_all(actual, yaml.Loader))
|
|
123
|
+
template = next(
|
|
124
|
+
template
|
|
125
|
+
for template in manifest["spec"]["workflowSpec"]["templates"]
|
|
126
|
+
if template["name"] == "asset1"
|
|
127
|
+
)
|
|
128
|
+
input_params = {param["name"]: param for param in template["inputs"]["parameters"]}
|
|
129
|
+
output_params = {param["name"]: param for param in template["outputs"]["parameters"]}
|
|
130
|
+
|
|
131
|
+
assert input_params["message"] == {"name": "message", "value": "hello"}
|
|
132
|
+
assert {"rundata", "asset_name", "moduledata"}.issubset(input_params.keys())
|
|
133
|
+
assert output_params["summary"] == {"name": "summary", "valueFrom": {"path": "/tmp/summary.txt"}}
|
|
134
|
+
assert output_params["rundata"] == {"name": "rundata", "valueFrom": {"path": "/tmp/tinyfan/rundata.json"}}
|
|
135
|
+
finally:
|
|
136
|
+
FLOW_REGISTER.clear()
|
|
137
|
+
|
|
138
|
+
|
|
106
139
|
def test_codegen_implicit(validate_crds, sample_flow):
|
|
107
140
|
actual = codegen()
|
|
108
141
|
validate_crds(actual)
|
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from datetime import UTC, datetime
|
|
4
|
+
|
|
5
|
+
from tinyfan.stores.gcs_pickle import GcsPickleStore
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class FakeBlob:
|
|
9
|
+
def __init__(self, name: str, generation: int | None = None) -> None:
|
|
10
|
+
self.name = name
|
|
11
|
+
self.generation = generation
|
|
12
|
+
self.content_type: str | None = None
|
|
13
|
+
self.data: bytes | None = None
|
|
14
|
+
|
|
15
|
+
def upload_from_string(self, data: bytes, content_type: str) -> None:
|
|
16
|
+
self.data = data
|
|
17
|
+
self.content_type = content_type
|
|
18
|
+
self.generation = (self.generation or 0) + 1
|
|
19
|
+
|
|
20
|
+
def download_as_bytes(self) -> bytes:
|
|
21
|
+
assert self.data is not None
|
|
22
|
+
return self.data
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class FakeBucket:
|
|
26
|
+
def __init__(self, name: str) -> None:
|
|
27
|
+
self.name = name
|
|
28
|
+
self.blobs: dict[str, FakeBlob] = {}
|
|
29
|
+
|
|
30
|
+
def blob(self, name: str, generation: int | None = None) -> FakeBlob:
|
|
31
|
+
blob = self.blobs.setdefault(name, FakeBlob(name, generation))
|
|
32
|
+
blob.generation = generation or blob.generation
|
|
33
|
+
return blob
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
class FakeClient:
|
|
37
|
+
def __init__(self) -> None:
|
|
38
|
+
self.buckets: dict[str, FakeBucket] = {}
|
|
39
|
+
|
|
40
|
+
def bucket(self, name: str) -> FakeBucket:
|
|
41
|
+
return self.buckets.setdefault(name, FakeBucket(name))
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def test_gcs_pickle_store_roundtrip() -> None:
|
|
45
|
+
client = FakeClient()
|
|
46
|
+
store = GcsPickleStore("test-bucket", prefix="runs", client=client)
|
|
47
|
+
value = {"hello": ["world", 1]}
|
|
48
|
+
rundata = {
|
|
49
|
+
"flow_name": "flow",
|
|
50
|
+
"asset_name": "asset",
|
|
51
|
+
"ds": "2026-05-31",
|
|
52
|
+
"data_interval_start": datetime(2026, 5, 31, 12, 30, tzinfo=UTC),
|
|
53
|
+
}
|
|
54
|
+
|
|
55
|
+
index = store.store(value, rundata)
|
|
56
|
+
|
|
57
|
+
assert index["bucket"] == "test-bucket"
|
|
58
|
+
assert index["name"].startswith("runs/flow/asset/2026-05-31T12-30-00-00-00/")
|
|
59
|
+
assert index["name"].endswith(".pickle")
|
|
60
|
+
assert index["generation"] == 1
|
|
61
|
+
assert client.bucket("test-bucket").blob(index["name"]).content_type == "application/octet-stream"
|
|
62
|
+
assert store.retrieve(index, source_rundata=rundata, target_rundata={}) == value
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def test_gcs_pickle_store_uses_metadata_blob_name() -> None:
|
|
66
|
+
store = GcsPickleStore("test-bucket", client=FakeClient())
|
|
67
|
+
rundata = {"metadata": {"blob_name": "/custom/path/value.pkl"}}
|
|
68
|
+
|
|
69
|
+
index = store.store("value", rundata)
|
|
70
|
+
|
|
71
|
+
assert index["name"] == "custom/path/value.pkl"
|