orchestracto-sdk 0.0.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.
- orchestracto_sdk-0.0.1/PKG-INFO +8 -0
- orchestracto_sdk-0.0.1/orc_sdk/__init__.py +4 -0
- orchestracto_sdk-0.0.1/orc_sdk/processor.py +255 -0
- orchestracto_sdk-0.0.1/orc_sdk/run_step.py +61 -0
- orchestracto_sdk-0.0.1/orc_sdk/step.py +146 -0
- orchestracto_sdk-0.0.1/orc_sdk/step_chain.py +58 -0
- orchestracto_sdk-0.0.1/orc_sdk/utils.py +6 -0
- orchestracto_sdk-0.0.1/orc_sdk/workflow.py +82 -0
- orchestracto_sdk-0.0.1/orchestracto_sdk.egg-info/PKG-INFO +8 -0
- orchestracto_sdk-0.0.1/orchestracto_sdk.egg-info/SOURCES.txt +15 -0
- orchestracto_sdk-0.0.1/orchestracto_sdk.egg-info/dependency_links.txt +1 -0
- orchestracto_sdk-0.0.1/orchestracto_sdk.egg-info/entry_points.txt +3 -0
- orchestracto_sdk-0.0.1/orchestracto_sdk.egg-info/requires.txt +3 -0
- orchestracto_sdk-0.0.1/orchestracto_sdk.egg-info/top_level.txt +1 -0
- orchestracto_sdk-0.0.1/pyproject.toml +20 -0
- orchestracto_sdk-0.0.1/setup.cfg +4 -0
- orchestracto_sdk-0.0.1/tests/test_main.py +66 -0
|
@@ -0,0 +1,255 @@
|
|
|
1
|
+
import argparse
|
|
2
|
+
import dataclasses
|
|
3
|
+
import datetime
|
|
4
|
+
import importlib.util
|
|
5
|
+
import inspect
|
|
6
|
+
import os
|
|
7
|
+
import os.path
|
|
8
|
+
import subprocess
|
|
9
|
+
import sys
|
|
10
|
+
import tempfile
|
|
11
|
+
from typing import Any
|
|
12
|
+
|
|
13
|
+
import yaml
|
|
14
|
+
|
|
15
|
+
from orc_client.client import OrcClient
|
|
16
|
+
|
|
17
|
+
from orc_sdk.workflow import WorkflowRuntimeObject
|
|
18
|
+
from orc_sdk.step import FuncStep, RetValWrapper
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def load_module(file_name, module_name):
|
|
22
|
+
spec = importlib.util.spec_from_file_location(module_name, file_name)
|
|
23
|
+
module = importlib.util.module_from_spec(spec)
|
|
24
|
+
sys.modules[module_name] = module
|
|
25
|
+
spec.loader.exec_module(module)
|
|
26
|
+
return module
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def compile_wf_config(wfro: WorkflowRuntimeObject) -> dict[str, Any]:
|
|
30
|
+
wf_config = {
|
|
31
|
+
"triggers": [],
|
|
32
|
+
"steps": [],
|
|
33
|
+
}
|
|
34
|
+
|
|
35
|
+
steps = wfro.get_steps()
|
|
36
|
+
for step_id, sci_info in steps.items():
|
|
37
|
+
if isinstance(sci_info.sci, FuncStep):
|
|
38
|
+
func_signature = inspect.signature(sci_info.sci.func)
|
|
39
|
+
|
|
40
|
+
passed_args = {}
|
|
41
|
+
for i, arg in enumerate(sci_info.sci.func_args):
|
|
42
|
+
arg_name = list(func_signature.parameters.keys())[i] # FIXME??
|
|
43
|
+
passed_args[arg_name] = arg
|
|
44
|
+
|
|
45
|
+
for name, value in sci_info.sci.func_kwargs.items():
|
|
46
|
+
passed_args[name] = value
|
|
47
|
+
|
|
48
|
+
step_args = []
|
|
49
|
+
for i, name in enumerate(func_signature.parameters):
|
|
50
|
+
if name in passed_args:
|
|
51
|
+
if isinstance(passed_args[name], RetValWrapper):
|
|
52
|
+
src_type = "step_output"
|
|
53
|
+
src_ref = f"{passed_args[name].sci.step_id}.{passed_args[name].name}"
|
|
54
|
+
else:
|
|
55
|
+
src_type = "constant"
|
|
56
|
+
src_ref = passed_args[name]
|
|
57
|
+
else:
|
|
58
|
+
src_type = "constant"
|
|
59
|
+
src_ref = func_signature.parameters[name].default
|
|
60
|
+
|
|
61
|
+
step_args.append({
|
|
62
|
+
"name": name,
|
|
63
|
+
"src_type": src_type,
|
|
64
|
+
"src_ref": src_ref,
|
|
65
|
+
})
|
|
66
|
+
|
|
67
|
+
secrets = []
|
|
68
|
+
for secret in sci_info.sci.secrets:
|
|
69
|
+
secrets.append({
|
|
70
|
+
"key": secret.key,
|
|
71
|
+
"value_ref": secret.value_ref,
|
|
72
|
+
"value_src_type": secret.value_src_type,
|
|
73
|
+
})
|
|
74
|
+
|
|
75
|
+
wf_config["steps"].append({
|
|
76
|
+
"step_id": step_id,
|
|
77
|
+
"task_type": "docker",
|
|
78
|
+
"task_params": {
|
|
79
|
+
"docker_image": "TODO",
|
|
80
|
+
"command": sci_info.sci.func.__name__,
|
|
81
|
+
},
|
|
82
|
+
"args": step_args,
|
|
83
|
+
"secrets": secrets,
|
|
84
|
+
"outputs": [{
|
|
85
|
+
"name": name,
|
|
86
|
+
} for name in sci_info.sci.retval_names],
|
|
87
|
+
"depends_on": list(sci_info.depends_on),
|
|
88
|
+
"cache": {
|
|
89
|
+
"enable": sci_info.sci.cache.enable,
|
|
90
|
+
},
|
|
91
|
+
})
|
|
92
|
+
else:
|
|
93
|
+
raise NotImplementedError
|
|
94
|
+
|
|
95
|
+
wf_config["triggers"] = wfro.triggers
|
|
96
|
+
|
|
97
|
+
return wf_config
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
def get_wfro_from_file(filename: str) -> WorkflowRuntimeObject:
|
|
101
|
+
module = load_module(filename, "user_code")
|
|
102
|
+
for key, obj in module.__dict__.items():
|
|
103
|
+
if getattr(obj, "is_workflow", False):
|
|
104
|
+
wfro = obj() # TODO: args?
|
|
105
|
+
return wfro
|
|
106
|
+
else:
|
|
107
|
+
raise Exception("No workflow found")
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
def get_file_module_root(filename: str) -> str:
|
|
111
|
+
dirname = os.path.dirname(os.path.abspath(filename))
|
|
112
|
+
if not os.path.exists(os.path.join(dirname, "__init__.py")):
|
|
113
|
+
return filename
|
|
114
|
+
return get_file_module_root(dirname)
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
BASE_IMAGE = "cr.eu-north1.nebius.cloud/e00faee7vas5hpsh3s/orchestracto/sdk-simple-runtime:24"
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def login_in_registry(registry_url: str, token: str):
|
|
121
|
+
proc = subprocess.Popen(
|
|
122
|
+
["docker", "login", registry_url, "--password-stdin", "-u", "user"],
|
|
123
|
+
stdin=subprocess.PIPE, stdout=sys.stdout, stderr=sys.stderr
|
|
124
|
+
)
|
|
125
|
+
proc.communicate(token.encode())
|
|
126
|
+
proc.wait()
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
def build_docker_with_reqs(
|
|
130
|
+
workflow_path: str,
|
|
131
|
+
image_name: str,
|
|
132
|
+
wf_file_module: str,
|
|
133
|
+
additional_requirements: list[str],
|
|
134
|
+
base_image: str = BASE_IMAGE,
|
|
135
|
+
) -> str:
|
|
136
|
+
py_packages_installation = f"""
|
|
137
|
+
RUN pip install -U {" ".join(additional_requirements)}
|
|
138
|
+
""" if additional_requirements else ""
|
|
139
|
+
|
|
140
|
+
dockerfile = f"""
|
|
141
|
+
FROM {BASE_IMAGE} AS sdk_runtime
|
|
142
|
+
FROM {base_image}
|
|
143
|
+
USER root
|
|
144
|
+
{py_packages_installation}
|
|
145
|
+
RUN mkdir /orc
|
|
146
|
+
COPY {os.path.basename(wf_file_module)} /orc/lib/{os.path.basename(wf_file_module)}
|
|
147
|
+
COPY --from=sdk_runtime /usr/local/lib/python3.12/site-packages/orc_sdk /orc/lib/orc_sdk
|
|
148
|
+
COPY --from=sdk_runtime /usr/local/bin/orc_run_step /usr/local/bin/orc_run_step
|
|
149
|
+
RUN new_shebang='#!/usr/bin/env python3' && sed -i "1s|.*|$new_shebang|" /usr/local/bin/orc_run_step
|
|
150
|
+
RUN chmod +x /usr/local/bin/orc_run_step
|
|
151
|
+
"""
|
|
152
|
+
|
|
153
|
+
registry_url = os.environ.get("REGISTRY_URL") # TODO: get from //sys/@ui_config?..
|
|
154
|
+
if registry_url is None:
|
|
155
|
+
registry_url = "registry." + os.environ["YT_PROXY"].removeprefix("https://").removeprefix("http://")
|
|
156
|
+
|
|
157
|
+
login_in_registry(registry_url, os.environ["YT_TOKEN"])
|
|
158
|
+
|
|
159
|
+
image_rel_path = workflow_path.removeprefix("//") + "/" + image_name
|
|
160
|
+
docker_tag = f"{registry_url}/home/orchestracto/public_registry/{image_rel_path}:latest"
|
|
161
|
+
|
|
162
|
+
with tempfile.NamedTemporaryFile() as tf:
|
|
163
|
+
with open(tf.name, "w") as f:
|
|
164
|
+
f.write(dockerfile)
|
|
165
|
+
subprocess.check_call(["docker", "build", "-t", docker_tag, "-f", tf.name, "--platform", "linux/amd64", os.path.dirname(wf_file_module), "--push"])
|
|
166
|
+
|
|
167
|
+
return docker_tag
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
def get_docker_image_with_sha(docker_image: str) -> str:
|
|
171
|
+
return subprocess.check_output([
|
|
172
|
+
"docker", "inspect", "--format", "{{index .RepoDigests 0}}", docker_image
|
|
173
|
+
]).decode().strip()
|
|
174
|
+
|
|
175
|
+
|
|
176
|
+
@dataclasses.dataclass
|
|
177
|
+
class WorkflowInfo:
|
|
178
|
+
workflow_path: str
|
|
179
|
+
workflow_config: dict[str, Any]
|
|
180
|
+
|
|
181
|
+
|
|
182
|
+
def get_wf_info_from_file(filename: str) -> WorkflowInfo:
|
|
183
|
+
wfro = get_wfro_from_file(filename)
|
|
184
|
+
wf_config = compile_wf_config(wfro)
|
|
185
|
+
|
|
186
|
+
wf_file_module = get_file_module_root(filename)
|
|
187
|
+
|
|
188
|
+
if os.path.abspath(wf_file_module) != os.path.abspath(filename):
|
|
189
|
+
wf_file_module_dir = os.path.dirname(wf_file_module) # FIXME
|
|
190
|
+
rel_file_path = os.path.abspath(filename).removeprefix(wf_file_module_dir + "/")
|
|
191
|
+
else:
|
|
192
|
+
rel_file_path = os.path.basename(filename) # TODO FIXME
|
|
193
|
+
path_in_container = f"/orc/lib/{rel_file_path}"
|
|
194
|
+
|
|
195
|
+
default_docker_tag = build_docker_with_reqs(
|
|
196
|
+
wfro.workflow_path, "default",
|
|
197
|
+
wf_file_module,
|
|
198
|
+
wfro.additional_requirements
|
|
199
|
+
)
|
|
200
|
+
|
|
201
|
+
wfro_steps = wfro.get_steps()
|
|
202
|
+
|
|
203
|
+
for step in wf_config["steps"]:
|
|
204
|
+
step_id = step["step_id"]
|
|
205
|
+
sci = wfro_steps[step_id].sci
|
|
206
|
+
assert isinstance(sci, FuncStep)
|
|
207
|
+
if sci.additional_requirements or sci.base_image:
|
|
208
|
+
docker_tag = build_docker_with_reqs(
|
|
209
|
+
wfro.workflow_path, step_id,
|
|
210
|
+
wf_file_module,
|
|
211
|
+
wfro.additional_requirements + sci.additional_requirements,
|
|
212
|
+
base_image=sci.base_image or BASE_IMAGE,
|
|
213
|
+
)
|
|
214
|
+
step["task_params"]["docker_image"] = docker_tag
|
|
215
|
+
else:
|
|
216
|
+
step["task_params"]["docker_image"] = default_docker_tag
|
|
217
|
+
|
|
218
|
+
step["task_params"]["env"] = {"PYTHONPATH": "/orc/lib", "YT_BASE_LAYER": step["task_params"]["docker_image"]}
|
|
219
|
+
step["task_params"]["command"] = f"orc_run_step {path_in_container} {step_id} >&2" # FIXME: corrupts cache key
|
|
220
|
+
step["task_params"]["func_code_hash"] = sci.func_code_hash
|
|
221
|
+
|
|
222
|
+
return WorkflowInfo(workflow_path=wfro.workflow_path, workflow_config=wf_config)
|
|
223
|
+
|
|
224
|
+
|
|
225
|
+
def process_python_file(filename: str):
|
|
226
|
+
wf_info = get_wf_info_from_file(filename)
|
|
227
|
+
|
|
228
|
+
if "ORC_URL" not in os.environ:
|
|
229
|
+
yt_proxy = os.environ.get("YT_PROXY")
|
|
230
|
+
if yt_proxy is None:
|
|
231
|
+
raise ValueError("Either ORC_URL or YT_PROXY environment variable should be set")
|
|
232
|
+
orc_url = "https://orc." + yt_proxy.removeprefix("https://").removeprefix("http://")
|
|
233
|
+
else:
|
|
234
|
+
orc_url = os.environ["ORC_URL"]
|
|
235
|
+
|
|
236
|
+
orc_client = OrcClient(orc_url=orc_url, yt_token=os.environ["YT_TOKEN"])
|
|
237
|
+
orc_client.update_workflow(wf_info.workflow_path, wf_info.workflow_config)
|
|
238
|
+
|
|
239
|
+
|
|
240
|
+
def configure_arg_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
|
|
241
|
+
subparsers = parser.add_subparsers(dest="command")
|
|
242
|
+
|
|
243
|
+
get_config = subparsers.add_parser("get-config")
|
|
244
|
+
get_config.add_argument("filename", type=str)
|
|
245
|
+
|
|
246
|
+
process_parser = subparsers.add_parser("process")
|
|
247
|
+
process_parser.add_argument("filename", type=str)
|
|
248
|
+
|
|
249
|
+
return parser
|
|
250
|
+
|
|
251
|
+
|
|
252
|
+
def process_args(args: argparse.Namespace):
|
|
253
|
+
match args.command:
|
|
254
|
+
case "process":
|
|
255
|
+
process_python_file(args.filename)
|
|
@@ -0,0 +1,61 @@
|
|
|
1
|
+
#!/usr/bin/env python3
|
|
2
|
+
|
|
3
|
+
import argparse
|
|
4
|
+
import importlib.util
|
|
5
|
+
import json
|
|
6
|
+
import os
|
|
7
|
+
import sys
|
|
8
|
+
|
|
9
|
+
from orc_sdk.step import FuncStep
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def load_module(file_name, module_name): # FIXME: duplicate
|
|
13
|
+
spec = importlib.util.spec_from_file_location(module_name, file_name)
|
|
14
|
+
module = importlib.util.module_from_spec(spec)
|
|
15
|
+
sys.modules[module_name] = module
|
|
16
|
+
spec.loader.exec_module(module)
|
|
17
|
+
return module
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def main():
|
|
21
|
+
parser = argparse.ArgumentParser()
|
|
22
|
+
parser.add_argument("pyfile_with_workflow")
|
|
23
|
+
parser.add_argument("step_id")
|
|
24
|
+
parser.add_argument("--base-layer")
|
|
25
|
+
|
|
26
|
+
args = parser.parse_args()
|
|
27
|
+
|
|
28
|
+
module = load_module(args.pyfile_with_workflow, "workflow")
|
|
29
|
+
|
|
30
|
+
for key, obj in module.__dict__.items():
|
|
31
|
+
if getattr(obj, "is_workflow", False):
|
|
32
|
+
wfro = obj()
|
|
33
|
+
break
|
|
34
|
+
else:
|
|
35
|
+
raise Exception("No workflow found")
|
|
36
|
+
|
|
37
|
+
steps = wfro.get_steps()
|
|
38
|
+
the_step_sci = steps[args.step_id].sci
|
|
39
|
+
|
|
40
|
+
if not isinstance(the_step_sci, FuncStep):
|
|
41
|
+
raise NotImplementedError
|
|
42
|
+
|
|
43
|
+
for key, value in os.environ.items():
|
|
44
|
+
if key.startswith("YT_SECURE_VAULT_"):
|
|
45
|
+
os.environ[key.removeprefix("YT_SECURE_VAULT_")] = value
|
|
46
|
+
|
|
47
|
+
returned_values = the_step_sci.run()
|
|
48
|
+
|
|
49
|
+
if len(the_step_sci.retval_names) > 0:
|
|
50
|
+
if not isinstance(returned_values, tuple):
|
|
51
|
+
returned_values = (returned_values,)
|
|
52
|
+
|
|
53
|
+
ret_dict = {}
|
|
54
|
+
for idx, value in enumerate(returned_values):
|
|
55
|
+
ret_dict[the_step_sci.retval_names[idx]] = value
|
|
56
|
+
|
|
57
|
+
print(json.dumps(ret_dict))
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
if __name__ == "__main__":
|
|
61
|
+
main()
|
|
@@ -0,0 +1,146 @@
|
|
|
1
|
+
import dataclasses
|
|
2
|
+
import hashlib
|
|
3
|
+
import inspect
|
|
4
|
+
import os
|
|
5
|
+
from typing import Callable, Any, Self
|
|
6
|
+
from functools import wraps
|
|
7
|
+
from collections import defaultdict
|
|
8
|
+
|
|
9
|
+
from orc_sdk.step_chain import StepChainItem
|
|
10
|
+
from orc_sdk.utils import random_string
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
@dataclasses.dataclass
|
|
14
|
+
class RetValWrapper:
|
|
15
|
+
value: Any
|
|
16
|
+
sci: StepChainItem
|
|
17
|
+
name: str
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
@dataclasses.dataclass
|
|
21
|
+
class SecretRecord:
|
|
22
|
+
key: str
|
|
23
|
+
value_ref: str
|
|
24
|
+
value_src_type: str
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
@dataclasses.dataclass
|
|
28
|
+
class SecretsMixin:
|
|
29
|
+
secrets: list[SecretRecord] = dataclasses.field(default_factory=list)
|
|
30
|
+
|
|
31
|
+
def with_secret(self, key: str, value_ref: str, value_src_type: str) -> Self:
|
|
32
|
+
self.secrets.append(SecretRecord(key=key, value_ref=value_ref, value_src_type=value_src_type))
|
|
33
|
+
return self
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
@dataclasses.dataclass
|
|
37
|
+
class CacheSettings:
|
|
38
|
+
enable: bool = dataclasses.field(default=False)
|
|
39
|
+
enable_write: bool | None = dataclasses.field(default=None)
|
|
40
|
+
enable_read: bool | None = dataclasses.field(default=None)
|
|
41
|
+
cache_version: str = dataclasses.field(default="v1")
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
@dataclasses.dataclass
|
|
45
|
+
class CacheMixin:
|
|
46
|
+
cache: CacheSettings = dataclasses.field(default_factory=CacheSettings)
|
|
47
|
+
|
|
48
|
+
def with_cache(self, version: str | None = None) -> Self:
|
|
49
|
+
self.cache.enable = True
|
|
50
|
+
if version is not None:
|
|
51
|
+
self.cache.cache_version = version
|
|
52
|
+
return self
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
@dataclasses.dataclass
|
|
56
|
+
class Step(StepChainItem):
|
|
57
|
+
task_type: str
|
|
58
|
+
task_params: dict[str, Any]
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
@dataclasses.dataclass
|
|
62
|
+
class FuncStep(SecretsMixin, CacheMixin, StepChainItem):
|
|
63
|
+
func: Callable[..., Any] = dataclasses.field(default_factory=lambda: None)
|
|
64
|
+
func_args: tuple[Any, ...] = dataclasses.field(default_factory=list)
|
|
65
|
+
func_kwargs: dict = dataclasses.field(default_factory=dict)
|
|
66
|
+
retval_names: list[str] = dataclasses.field(default_factory=list)
|
|
67
|
+
additional_requirements: list[str] = dataclasses.field(default_factory=list)
|
|
68
|
+
base_image: str | None = dataclasses.field(default=None)
|
|
69
|
+
func_code_hash: str | None = dataclasses.field(default=None)
|
|
70
|
+
|
|
71
|
+
def __post_init__(self):
|
|
72
|
+
super().__post_init__()
|
|
73
|
+
for arg in self.func_args + tuple(self.func_kwargs.values()):
|
|
74
|
+
if isinstance(arg, RetValWrapper):
|
|
75
|
+
if arg.sci.step_id not in [ps.step_id for ps in self._prev_steps]: # TODO: traverse
|
|
76
|
+
self._prev_steps.append(arg.sci)
|
|
77
|
+
arg.sci._next_steps.append(self)
|
|
78
|
+
self._first.extend(arg.sci._first)
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
def with_additional_requirements(self, additional_requirements: list[str]) -> Self:
|
|
82
|
+
self.additional_requirements = additional_requirements
|
|
83
|
+
return self
|
|
84
|
+
|
|
85
|
+
def with_base_image(self, base_image: str) -> Self:
|
|
86
|
+
self.base_image = base_image
|
|
87
|
+
return self
|
|
88
|
+
|
|
89
|
+
@property
|
|
90
|
+
def outputs(self):
|
|
91
|
+
ret_vars = inspect.signature(self.func).return_annotation
|
|
92
|
+
if ret_vars is tuple:
|
|
93
|
+
# TODO
|
|
94
|
+
pass
|
|
95
|
+
else:
|
|
96
|
+
parameters = {
|
|
97
|
+
self.retval_names[0]: RetValWrapper(value=ret_vars, sci=self, name=self.retval_names[0])}
|
|
98
|
+
name = self.func.__name__.capitalize()
|
|
99
|
+
cls = type(name, (), parameters)
|
|
100
|
+
return dataclasses.dataclass(cls)
|
|
101
|
+
|
|
102
|
+
def run(self):
|
|
103
|
+
processed_args = []
|
|
104
|
+
processed_kwargs = {}
|
|
105
|
+
func_signature = inspect.signature(self.func)
|
|
106
|
+
|
|
107
|
+
for i, arg in enumerate(self.func_args):
|
|
108
|
+
if isinstance(arg, RetValWrapper):
|
|
109
|
+
processed_args.append(os.environ[f"ORC_PARAM_{list(func_signature.parameters.keys())[i]}"])
|
|
110
|
+
else:
|
|
111
|
+
processed_args.append(arg)
|
|
112
|
+
|
|
113
|
+
for key, value in self.func_kwargs.items():
|
|
114
|
+
if isinstance(value, RetValWrapper):
|
|
115
|
+
processed_kwargs[key] = os.environ[f"ORC_PARAM_{key}"]
|
|
116
|
+
else:
|
|
117
|
+
processed_kwargs[key] = value
|
|
118
|
+
|
|
119
|
+
return self.func(*processed_args, **processed_kwargs)
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
FUNC_NAME_COUNTER = defaultdict(lambda: 0)
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
def task(retval_names: list[str] | None = None) -> Callable[[Callable[..., Any]], Callable[..., FuncStep]]:
|
|
126
|
+
def decorator(function: Callable[..., Any]) -> Callable[..., FuncStep]:
|
|
127
|
+
@wraps(function)
|
|
128
|
+
def wrapper(*args, **kwargs) -> FuncStep:
|
|
129
|
+
FUNC_NAME_COUNTER[function.__name__] += 1
|
|
130
|
+
sro_id = function.__name__ + "_" + str(FUNC_NAME_COUNTER[function.__name__])
|
|
131
|
+
func_code_hash = hashlib.md5(inspect.getsource(function).encode()).hexdigest()
|
|
132
|
+
|
|
133
|
+
nonlocal retval_names
|
|
134
|
+
if retval_names is None and inspect.signature(function).return_annotation is not inspect._empty:
|
|
135
|
+
num_outputs = 1 if not inspect.signature(function).return_annotation is tuple else len(inspect.signature(function).return_annotation)
|
|
136
|
+
retval_names = [f"output_{i}" for i in range(1, num_outputs + 1)]
|
|
137
|
+
|
|
138
|
+
sro = FuncStep(
|
|
139
|
+
step_id=sro_id, func=function, func_args=args, func_kwargs=kwargs,
|
|
140
|
+
retval_names=retval_names or [], func_code_hash=func_code_hash,
|
|
141
|
+
)
|
|
142
|
+
return sro
|
|
143
|
+
|
|
144
|
+
return wrapper
|
|
145
|
+
|
|
146
|
+
return decorator
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
import dataclasses
|
|
2
|
+
from typing import Self
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
@dataclasses.dataclass
|
|
6
|
+
class StepChainItem:
|
|
7
|
+
step_id: str
|
|
8
|
+
|
|
9
|
+
_first: list["StepChainItem"] = dataclasses.field(init=False)
|
|
10
|
+
_next_steps: list["StepChainItem"] = dataclasses.field(init=False)
|
|
11
|
+
_prev_steps: list["StepChainItem"] = dataclasses.field(init=False)
|
|
12
|
+
|
|
13
|
+
def __post_init__(self):
|
|
14
|
+
self._first = [self]
|
|
15
|
+
self._next_steps = []
|
|
16
|
+
self._prev_steps = []
|
|
17
|
+
|
|
18
|
+
def __rshift__(self, other: Self | list[Self]) -> Self:
|
|
19
|
+
self.set_downstream(other)
|
|
20
|
+
return other
|
|
21
|
+
|
|
22
|
+
def __lshift__(self, other: Self | list[Self]) -> Self | list[Self]:
|
|
23
|
+
self.set_upstream(other)
|
|
24
|
+
return other
|
|
25
|
+
|
|
26
|
+
def __rrshift__(self, other: Self | list[Self]) -> Self:
|
|
27
|
+
self.__lshift__(other)
|
|
28
|
+
return self
|
|
29
|
+
|
|
30
|
+
def __rlshift__(self, other: Self | list[Self]):
|
|
31
|
+
self.__rrshift__(other)
|
|
32
|
+
return self
|
|
33
|
+
|
|
34
|
+
def set_downstream(self, other: Self | list[Self]):
|
|
35
|
+
if isinstance(other, list):
|
|
36
|
+
for sro in other:
|
|
37
|
+
sro._first = self._first
|
|
38
|
+
self._next_steps.append(sro)
|
|
39
|
+
sro._prev_steps.append(self)
|
|
40
|
+
else:
|
|
41
|
+
other._first = self._first
|
|
42
|
+
self._next_steps.append(other)
|
|
43
|
+
other._prev_steps.append(self)
|
|
44
|
+
|
|
45
|
+
def set_upstream(self, other: Self | list[Self]):
|
|
46
|
+
if isinstance(other, list):
|
|
47
|
+
for sro in other:
|
|
48
|
+
sro._next_steps.append(self)
|
|
49
|
+
self._prev_steps.append(sro)
|
|
50
|
+
self._first = other[0]._first
|
|
51
|
+
else:
|
|
52
|
+
other._next_steps.append(self)
|
|
53
|
+
self._first = other._first
|
|
54
|
+
self._prev_steps.append(other)
|
|
55
|
+
|
|
56
|
+
def with_id(self, step_id: str) -> Self:
|
|
57
|
+
self.step_id = step_id
|
|
58
|
+
return self
|
|
@@ -0,0 +1,82 @@
|
|
|
1
|
+
import inspect
|
|
2
|
+
import dataclasses
|
|
3
|
+
from typing import Any
|
|
4
|
+
from functools import wraps
|
|
5
|
+
|
|
6
|
+
from orc_sdk.step_chain import StepChainItem
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
@dataclasses.dataclass
|
|
10
|
+
class SCIInfo:
|
|
11
|
+
sci: StepChainItem
|
|
12
|
+
depends_on: set[str]
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class WorkflowRuntimeObject:
|
|
16
|
+
def __init__(
|
|
17
|
+
self, workflow_path: str, triggers: list[Any],
|
|
18
|
+
additional_requirements: list[str]
|
|
19
|
+
):
|
|
20
|
+
self.workflow_path = workflow_path
|
|
21
|
+
self.triggers = triggers
|
|
22
|
+
|
|
23
|
+
self.additional_requirements = additional_requirements
|
|
24
|
+
|
|
25
|
+
self.first_steps = []
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def register_first_step(self, step: StepChainItem) -> StepChainItem:
|
|
29
|
+
self.first_steps.append(step)
|
|
30
|
+
return step
|
|
31
|
+
|
|
32
|
+
def get_steps(self):
|
|
33
|
+
steps: dict[str, SCIInfo] = {}
|
|
34
|
+
|
|
35
|
+
scis_to_inspect: list[tuple[StepChainItem, str | None]] = [(sci, None) for sci in self.first_steps]
|
|
36
|
+
|
|
37
|
+
while scis_to_inspect:
|
|
38
|
+
sci, parent = scis_to_inspect.pop(0)
|
|
39
|
+
|
|
40
|
+
if sci.step_id in steps:
|
|
41
|
+
assert parent is not None
|
|
42
|
+
steps[sci.step_id].depends_on.add(parent)
|
|
43
|
+
elif parent is not None:
|
|
44
|
+
steps[sci.step_id] = SCIInfo(sci, {parent})
|
|
45
|
+
else:
|
|
46
|
+
steps[sci.step_id] = SCIInfo(sci, set())
|
|
47
|
+
|
|
48
|
+
for next_sro in sci._next_steps:
|
|
49
|
+
scis_to_inspect.append((next_sro, sci.step_id))
|
|
50
|
+
|
|
51
|
+
return steps
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
@dataclasses.dataclass
|
|
55
|
+
class WfArgWrapper:
|
|
56
|
+
value: Any
|
|
57
|
+
name: str
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def workflow(workflow_path: str, triggers: list[Any] | None = None, additional_requirements: list[str] | None = None):
|
|
61
|
+
def decorator(function):
|
|
62
|
+
@wraps(function)
|
|
63
|
+
def wrapper(*args, **kwargs):
|
|
64
|
+
wfro = WorkflowRuntimeObject(
|
|
65
|
+
workflow_path,
|
|
66
|
+
triggers=triggers or [],
|
|
67
|
+
additional_requirements=additional_requirements or [],
|
|
68
|
+
)
|
|
69
|
+
|
|
70
|
+
wrapped_args = []
|
|
71
|
+
for arg in inspect.signature(function).parameters.values():
|
|
72
|
+
wrapped_args.append(WfArgWrapper(value=arg.default, name=arg.name))
|
|
73
|
+
|
|
74
|
+
wrapped_args[0] = wfro
|
|
75
|
+
step_chain_item = function(*wrapped_args)
|
|
76
|
+
wfro.step_chain_item = step_chain_item
|
|
77
|
+
return wfro
|
|
78
|
+
|
|
79
|
+
wrapper.is_workflow = True
|
|
80
|
+
return wrapper
|
|
81
|
+
|
|
82
|
+
return decorator
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
pyproject.toml
|
|
2
|
+
orc_sdk/__init__.py
|
|
3
|
+
orc_sdk/processor.py
|
|
4
|
+
orc_sdk/run_step.py
|
|
5
|
+
orc_sdk/step.py
|
|
6
|
+
orc_sdk/step_chain.py
|
|
7
|
+
orc_sdk/utils.py
|
|
8
|
+
orc_sdk/workflow.py
|
|
9
|
+
orchestracto_sdk.egg-info/PKG-INFO
|
|
10
|
+
orchestracto_sdk.egg-info/SOURCES.txt
|
|
11
|
+
orchestracto_sdk.egg-info/dependency_links.txt
|
|
12
|
+
orchestracto_sdk.egg-info/entry_points.txt
|
|
13
|
+
orchestracto_sdk.egg-info/requires.txt
|
|
14
|
+
orchestracto_sdk.egg-info/top_level.txt
|
|
15
|
+
tests/test_main.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
orc_sdk
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "orchestracto-sdk"
|
|
3
|
+
authors = [{name = "TractoAI team"}]
|
|
4
|
+
requires-python = ">=3.11"
|
|
5
|
+
dependencies = [
|
|
6
|
+
"orchestracto-client>=0.5.0",
|
|
7
|
+
"PyYAML~=6.0.2",
|
|
8
|
+
"requests~=2.32.3",
|
|
9
|
+
]
|
|
10
|
+
dynamic = ["version"]
|
|
11
|
+
|
|
12
|
+
[tool.setuptools]
|
|
13
|
+
packages = ["orc_sdk"]
|
|
14
|
+
|
|
15
|
+
[tool.setuptools.dynamic]
|
|
16
|
+
version = {attr = "orc_sdk.__version__"}
|
|
17
|
+
|
|
18
|
+
[project.scripts]
|
|
19
|
+
orcsdk = "orc_sdk.processor:main_cli"
|
|
20
|
+
orc_run_step = "orc_sdk.run_step:main"
|
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
from orc_sdk.main import workflow, task, SROBlock, get_steps
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
def test_main():
|
|
5
|
+
|
|
6
|
+
@task()
|
|
7
|
+
def foo():
|
|
8
|
+
print("foo")
|
|
9
|
+
|
|
10
|
+
@task()
|
|
11
|
+
def bar(key: str):
|
|
12
|
+
print(f"bar {key}")
|
|
13
|
+
|
|
14
|
+
@task()
|
|
15
|
+
def xyz():
|
|
16
|
+
print("xyz")
|
|
17
|
+
|
|
18
|
+
@task()
|
|
19
|
+
def abc():
|
|
20
|
+
print("abc")
|
|
21
|
+
|
|
22
|
+
#
|
|
23
|
+
# /foo \
|
|
24
|
+
# abc/ xyz\
|
|
25
|
+
# / \ bar/ \
|
|
26
|
+
# foo -> bar / \
|
|
27
|
+
# \ foo-------------bar
|
|
28
|
+
# xyz/
|
|
29
|
+
# \bar
|
|
30
|
+
#
|
|
31
|
+
|
|
32
|
+
foo_step = foo()
|
|
33
|
+
|
|
34
|
+
bar_step = foo_step >> bar(key="some_key")
|
|
35
|
+
|
|
36
|
+
abc_step = bar_step >> abc()
|
|
37
|
+
xyz_step = bar_step >> xyz()
|
|
38
|
+
|
|
39
|
+
last_xyz_step = abc_step >> SROBlock(foo(), bar()) >> xyz()
|
|
40
|
+
|
|
41
|
+
xyz_step >> bar()
|
|
42
|
+
last_foo_step = xyz_step >> foo()
|
|
43
|
+
|
|
44
|
+
last_bar_step = SROBlock(last_xyz_step, last_foo_step) >> bar()
|
|
45
|
+
|
|
46
|
+
stlist = get_steps(foo_step)
|
|
47
|
+
pass
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
def test_simple():
|
|
51
|
+
@task()
|
|
52
|
+
def foo():
|
|
53
|
+
print("foo")
|
|
54
|
+
|
|
55
|
+
@task()
|
|
56
|
+
def bar():
|
|
57
|
+
print("bar")
|
|
58
|
+
|
|
59
|
+
@task()
|
|
60
|
+
def fin():
|
|
61
|
+
print("fin")
|
|
62
|
+
|
|
63
|
+
foobar = SROBlock(foo().with_id("the_foo"), bar().with_id("the_bar"))
|
|
64
|
+
foobar >> fin().with_id("the_fin")
|
|
65
|
+
|
|
66
|
+
pass
|