From b5d41b4306ef7c90e5ee72b595c95f6ae542a9fe Mon Sep 17 00:00:00 2001 From: Ivan Krakhmaliuk Date: Thu, 26 Jun 2025 15:28:37 +0300 Subject: [PATCH 1/4] Minimal arbiter that has working unit tests --- .gitignore | 41 +- README.md | 78 +- arbiter/__init__.py | 27 + arbiter/arbiter.py | 376 ++++++ .../shared => arbiter/eventnode}/__init__.py | 6 +- arbiter/eventnode/base.py | 241 ++++ arbiter/eventnode/hooks.py | 23 + arbiter/eventnode/mock.py | 50 + arbiter/eventnode/tools.py | 63 + arbiter/log.py | 98 ++ arbiter/minion.py | 91 ++ arbiter/task.py | 55 + .../test_api => arbiter/tasknode}/__init__.py | 4 +- arbiter/tasknode/housekeeper.py | 61 + arbiter/tasknode/tasknode.py | 1168 +++++++++++++++++ arbiter/tasknode/tools.py | 48 + arbiter/tasknode/watcher.py | 320 +++++ docker-compose.yml | 59 - pylon/configs/.gitkeep | 0 pylon/plugins/shared/metadata.json | 6 - pylon/plugins/shared/module.py | 59 - pylon/plugins/shared/requirements.txt | 3 - pylon/plugins/shared/tools/config.py | 109 -- pylon/plugins/shared/tools/db.py | 77 -- pylon/plugins/shared/tools/db_tools.py | 88 -- pylon/plugins/shared/tools/patterns.py | 58 - pylon/plugins/test_api/api/__init__.py | 0 pylon/plugins/test_api/api/v1/__init__.py | 0 pylon/plugins/test_api/api/v1/metadata.py | 24 - pylon/plugins/test_api/init_db.py | 6 - pylon/plugins/test_api/metadata.json | 6 - pylon/plugins/test_api/models/__init__.py | 0 pylon/plugins/test_api/models/metadata.py | 31 - pylon/plugins/test_api/models/pd/__init__.py | 0 pylon/plugins/test_api/models/pd/metadata.py | 5 - pylon/plugins/test_api/module.py | 44 - pylon/pylon.yml | 64 - pylon/requirements/.gitkeep | 0 requirements.txt | 18 + setup.py | 49 + tests/README.md | 22 + .../shared/tools => tests}/__init__.py | 0 tests/minion.py | 45 + tests/test_arbiter.py | 178 +++ version.txt | 1 + 45 files changed, 3019 insertions(+), 683 deletions(-) create mode 100644 arbiter/__init__.py create mode 100644 arbiter/arbiter.py rename {pylon/plugins/shared => arbiter/eventnode}/__init__.py (87%) create mode 100644 arbiter/eventnode/base.py create mode 100644 arbiter/eventnode/hooks.py create mode 100644 arbiter/eventnode/mock.py create mode 100644 arbiter/eventnode/tools.py create mode 100644 arbiter/log.py create mode 100644 arbiter/minion.py create mode 100644 arbiter/task.py rename {pylon/plugins/test_api => arbiter/tasknode}/__init__.py (92%) create mode 100644 arbiter/tasknode/housekeeper.py create mode 100644 arbiter/tasknode/tasknode.py create mode 100644 arbiter/tasknode/tools.py create mode 100644 arbiter/tasknode/watcher.py delete mode 100644 docker-compose.yml delete mode 100644 pylon/configs/.gitkeep delete mode 100644 pylon/plugins/shared/metadata.json delete mode 100644 pylon/plugins/shared/module.py delete mode 100644 pylon/plugins/shared/requirements.txt delete mode 100644 pylon/plugins/shared/tools/config.py delete mode 100644 pylon/plugins/shared/tools/db.py delete mode 100644 pylon/plugins/shared/tools/db_tools.py delete mode 100644 pylon/plugins/shared/tools/patterns.py delete mode 100644 pylon/plugins/test_api/api/__init__.py delete mode 100644 pylon/plugins/test_api/api/v1/__init__.py delete mode 100644 pylon/plugins/test_api/api/v1/metadata.py delete mode 100644 pylon/plugins/test_api/init_db.py delete mode 100644 pylon/plugins/test_api/metadata.json delete mode 100644 pylon/plugins/test_api/models/__init__.py delete mode 100644 pylon/plugins/test_api/models/metadata.py delete mode 100644 pylon/plugins/test_api/models/pd/__init__.py delete mode 100644 pylon/plugins/test_api/models/pd/metadata.py delete mode 100644 pylon/plugins/test_api/module.py delete mode 100644 pylon/pylon.yml delete mode 100644 pylon/requirements/.gitkeep create mode 100644 requirements.txt create mode 100644 setup.py create mode 100644 tests/README.md rename {pylon/plugins/shared/tools => tests}/__init__.py (100%) create mode 100644 tests/minion.py create mode 100644 tests/test_arbiter.py create mode 100644 version.txt diff --git a/.gitignore b/.gitignore index 3d984c8..b6e4761 100644 --- a/.gitignore +++ b/.gitignore @@ -20,6 +20,7 @@ parts/ sdist/ var/ wheels/ +pip-wheel-metadata/ share/python-wheels/ *.egg-info/ .installed.cfg @@ -49,7 +50,6 @@ coverage.xml *.py,cover .hypothesis/ .pytest_cache/ -cover/ # Translations *.mo @@ -72,7 +72,6 @@ instance/ docs/_build/ # PyBuilder -.pybuilder/ target/ # Jupyter Notebook @@ -83,9 +82,7 @@ profile_default/ ipython_config.py # pyenv -# For a library or package, you might want to ignore these files since the code is -# intended to run in multiple environments; otherwise, check them in: -# .python-version +.python-version # pipenv # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. @@ -94,22 +91,7 @@ ipython_config.py # install all needed dependencies. #Pipfile.lock -# poetry -# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. -# This is especially recommended for binary packages to ensure reproducibility, and is more -# commonly ignored for libraries. -# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control -#poetry.lock - -# pdm -# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. -#pdm.lock -# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it -# in version control. -# https://pdm.fming.dev/#use-with-ide -.pdm.toml - -# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +# PEP 582; used by e.g. github.com/David-OConnor/pyflow __pypackages__/ # Celery stuff @@ -145,20 +127,3 @@ dmypy.json # Pyre type checker .pyre/ - -# pytype static type analyzer -.pytype/ - -# Cython debug symbols -cython_debug/ - -# PyCharm -# JetBrains specific template is maintained in a separate JetBrains.gitignore that can -# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore -# and can be added to the global gitignore or merged into this file. For a more nuclear -# option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ - -# project-specific -/pylon/requirements/* -!/pylon/requirements/.gitkeep diff --git a/README.md b/README.md index 6d24850..81ef2e2 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,76 @@ -# test_project -Test task repository (temporary) +# Arbiter +Distributed tasks queue use Redis as broker. Consists of arbiter and minion. + +## Installation + +Clone git repo +```bash +git clone https://github.com/carrier-io/arbiter.git +``` +and install by running +```bash +cd arbiter +python setup.py install +``` + +## Basic scenario + +Launch Redis, as it required for everything to work +```bash +docker run -d --rm --hostname arbiter-redis --name arbiter-redis \ + -p 6379:6379 redis:alpine redis-server +``` + +### Create simple task +You need to initiate minion and provide connection details: +```python +from arbiter import RedisEventNode +from arbiter import Minion + +event_node = RedisEventNode(host="localhost", port=6379, password="", event_queue="tasks") +app = Minion(event_node, queue="default") +``` +then you can declare tasks by decorating callable with `@app.task` +```python +@app.task(name="simple_add") +def adds(x, y): + return x + y +``` +Every task need to have a name, which it will be referred by when initiated from arbiter. +this is pretty much it to create first task + +Now we need to create execution point +```python +if __name__ == "__main__": + app.run(workers=3) +``` +where `workers` is a quantity of worker slots to do the job(s) + +Run created script. Minion is ready to accept work orders. + +### Call created task from arbiter +Arbiter is job initiator, it maintain the state of all jobs it created and can retrieve results. + +Each arbiter have it's own communication channel, so job results won't mess between two different arbiters + +Declaring the arbiter +```python +from arbiter import RedisEventNode +from arbiter import Arbiter + +event_node = RedisEventNode(host="localhost", port=6379, password="", event_queue="tasks") +arbiter = Arbiter(event_node) +``` +to call the task and track it till it done (tasks are obviously async) +```python +task_keys = arbiter.apply("simple_add", tasks_count=1, task_args=[1, 2]) # will return array of task ids + +# while loop with returns results of each task once it done +for message in arbiter.wait_for_tasks(task_keys): + print(message) +``` +Alternatively you can get task result by calling +```python +arbiter.status(task_keys[0]) +``` +it will return `json` where `result` will be one of the keys diff --git a/arbiter/__init__.py b/arbiter/__init__.py new file mode 100644 index 0000000..b5009fb --- /dev/null +++ b/arbiter/__init__.py @@ -0,0 +1,27 @@ +#!/usr/bin/python3 +# coding=utf-8 +# pylint: disable=C0114 + +# Copyright 2023 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from arbiter.arbiter import Arbiter +from arbiter.minion import Minion +from arbiter.task import Task + +from arbiter.eventnode import make_event_node +from arbiter.eventnode import MockEventNode + +from arbiter.tasknode import TaskNode diff --git a/arbiter/arbiter.py b/arbiter/arbiter.py new file mode 100644 index 0000000..b5dadc6 --- /dev/null +++ b/arbiter/arbiter.py @@ -0,0 +1,376 @@ +#!/usr/bin/python3 +# coding=utf-8 +# pylint: disable=C0114,C0115,C0116 + +# Copyright 2023 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import time +import threading + +from uuid import uuid4 + +from arbiter import log + +from .task import Task +from .tasknode import TaskNode + + +class Arbiter: # pylint: disable=R0902 + def __init__(self, event_node, finalizer_check_interval=10): + self.raw_task_node = TaskNode(event_node, task_limit=0) + self.finalizer_check_interval = finalizer_check_interval + # + self.task_state = {} + self.group_state = {} + self.callbacks = {} + self.finalizers = {} + # + self.lock = threading.Lock() + self.stop_event = threading.Event() + self.no_waiting_tasks = threading.Event() + + @property + def task_node(self): + if not self.raw_task_node.started: + self.stop_event.clear() + self.no_waiting_tasks.clear() + self.raw_task_node.start() + self.raw_task_node.subscribe_to_task_statuses(self.on_task_change) + ArbiterFinalizer(self).start() + # + return self.raw_task_node + + def on_task_change(self, event, data): + _ = event + # + task_id = data.get("task_id", None) + status = data.get("status", "unknown") + # + if task_id is None or status == "unknown": + return + # + status_map = { + "pending": "initiated", + "running": "running", + "stopped": "done", + } + # + with self.lock: + if task_id not in self.task_state: + self.task_state[task_id] = { + "task_type": "task", + } + # + self.task_state[task_id]["state"] = status_map[status] + # + if status == "stopped": + result = self.task_node.get_task_result(task_id) + self.task_state[task_id]["result"] = result + + def wait_for_tasks(self, tasks): + for task in tasks: + result = self.task_node.join_task(task) + # + with self.lock: + if task not in self.task_state: + self.task_state[task] = { + "task_type": "task", + } + # + self.task_state[task]["state"] = "done" + self.task_state[task]["result"] = result + # + yield self.task_state[task].copy() + + def add_task(self, task, sync=False): + tasks = [] + # + for _ in range(task.tasks_count): + task_key = self.task_node.start_task( + name=task.name, + args=task.task_args, + kwargs=task.task_kwargs, + pool=task.queue + ) + # + with self.lock: + if task_key not in self.task_state: + self.task_state[task_key] = { + "task_type": task.task_type, + "state": "initiated", + } + else: + self.task_state[task_key]["task_type"] = task.task_type + # + tasks.append(task_key) + yield task_key + # + if sync: + for message in self.wait_for_tasks(tasks): + yield message + + def apply(self, task_name, queue="default", tasks_count=1, task_args=None, task_kwargs=None, sync=False): # pylint: disable=C0301,R0913 + task = Task(name=task_name, queue=queue, tasks_count=tasks_count, + task_args=task_args, task_kwargs=task_kwargs) + return list(self.add_task(task, sync=sync)) + + def kill(self, task_key, sync=True): + self.task_node.stop_task(task_key) + if sync: + self.task_node.wait_for_task(task_key) + + def kill_group(self, group_id): + tasks = [] + # + for task_id in self.group_state[group_id]: + if task_id in self.task_state: + tasks.append(task_id) + self.kill(task_id, sync=False) + # + log.info("Terminating ...") + for task in tasks: + self.task_node.wait_for_task(task) + + def status(self, task_key): + if task_key in self.task_state: + return self.task_state[task_key] + # + if task_key in self.group_state: + group_results = { + "state": "done", + "initiated": 0, + "running": 0, + "done": 0, + "tasks": [] + } + # + for task_id in self.group_state[task_key]: + if task_id in self.task_state: + if self.task_state[task_id]["state"] in ["running", "initiated"]: + group_results["state"] = self.task_state[task_id]["state"] + group_results[self.task_state[task_id]["state"]] += 1 + group_results["tasks"].append(self.task_state[task_id]) + else: + log.info(f"[Group status] {task_id} is missing") + group_results["state"] = "running" + # + with self.lock: + for callback in self.callbacks.values(): + if callback["group_id"] == task_key: + group_results["state"] = "running" + group_results["initiated"] += 1 + # + for finalizer in self.finalizers.values(): + if finalizer["group_id"] == task_key: + group_results["state"] = "running" + group_results["initiated"] += 1 + # + return group_results + # + raise NameError("Task or Group not found") + + def close(self, waiting_tasks_timeout=None): + self.no_waiting_tasks.wait(waiting_tasks_timeout) + self.stop_event.set() + if self.raw_task_node.started: + self.raw_task_node.stop() + + def workers(self): + self.task_node.query_pool_state() + result = {} + # + for pool, nodes in self.task_node.global_pool_state.items(): + if pool not in result: + result[pool] = { + "total": 0, + "active": 0, + "available": 0, + } + # + for state in nodes.values(): + total = state.get("task_limit", None) + if total is None: + total = 1000 # FIXME: unlimited by task_node, using some constant here + # + active = state.get("running_tasks", 0) + available = total - active + # + result[pool]["total"] += total + result[pool]["active"] += active + result[pool]["available"] += available + # + return result + + def squad(self, tasks, callback=None): + """ + Set of tasks that need to be executed together + """ + workers_count = {} + for each in tasks: + if each.task_type != "finalize": + if each.queue not in list(workers_count.keys()): # pylint: disable=C0201 + workers_count[each.queue] = 0 + workers_count[each.queue] += each.tasks_count + # + stats = self.workers() + log.info(f"Workers: {stats}") + log.info(f"Tests to run {workers_count}") + # + for key in workers_count.keys(): # pylint: disable=C0201,C0206 + if not stats.get(key) or stats[key]["available"] < workers_count[key]: + raise NameError(f"Not enough of {key} workers") + # + return self.group(tasks, callback) + + def group(self, tasks, callback=None): + """ + Set of tasks that need to be executed regardless of order + """ + group_id = str(uuid4()) + self.group_state[group_id] = [] + # + finalizers = [] + # + for each in tasks: + if each.task_type == "finalize": + finalizers.append(each) + continue + # + for task in self.add_task(each): + self.group_state[group_id].append(task) + # + if callback: + callback_id = f'callback-{str(uuid4())}' + callback.task_type = "callback" + # + with self.lock: + self.callbacks[callback_id] = { + "callback": callback, + "group_id": group_id, + } + # + with self.lock: + for finalizer in finalizers: + finalizer_id = f'finalizer-{str(uuid4())}' + # + self.finalizers[finalizer_id] = { + "finalizer": finalizer, + "group_id": group_id, + } + # + return group_id + + def pipe(self, tasks, persistent_args=None, persistent_kwargs=None): + """ + Set of tasks that need to be executed sequentially + NOTE: Persistent args always before the task args + Task itself need to have **kwargs if you want to ignore upstream results + """ + pipe_id = str(uuid4()) + self.group_state[pipe_id] = [] + # + if not persistent_args: + persistent_args = [] + if not persistent_kwargs: + persistent_kwargs = {} + # + res = {} + yield {"pipe_id": pipe_id} + # + for task in tasks: + task.task_args = persistent_args + task.task_args + # + for key, value in persistent_kwargs: + if key not in task.task_kwargs: + task.task_kwargs[key] = value + # + if res: + task.task_kwargs['upstream'] = res.get("result") + # + res = list(self.add_task(task, sync=True)) + self.group_state[pipe_id].append(res[0]) + # + res = res[1] + yield res + + +class ArbiterFinalizer(threading.Thread): # pylint: disable=R0903 + """ Perform finalizer timeout checks """ + + def __init__(self, arbiter): + super().__init__(daemon=True) + self.arbiter = arbiter + + def run(self): + """ Run checker thread """ + while not self.arbiter.stop_event.is_set(): + time.sleep(self.arbiter.finalizer_check_interval) + # + # Callbacks + # + with self.arbiter.lock: + active_callbacks = list(self.arbiter.callbacks) + # + for callback_id in active_callbacks: + callback = self.arbiter.callbacks[callback_id] + # + if not all( + self.arbiter.status(task_id)["state"] == "done" + for task_id in self.arbiter.group_state[callback["group_id"]] + ): + continue + # + for task_id in self.arbiter.add_task(callback["callback"]): + with self.arbiter.lock: + self.arbiter.group_state[callback["group_id"]].append(task_id) + # + with self.arbiter.lock: + self.arbiter.callbacks.pop(callback_id, None) + # + # Finalizers + # + with self.arbiter.lock: + active_finalizers = list(self.arbiter.finalizers) + # + for finalizer_id in active_finalizers: + finalizer = self.arbiter.finalizers[finalizer_id] + # + if not all( + self.arbiter.status(task_id)["state"] == "done" + for task_id in self.arbiter.group_state[finalizer["group_id"]] + ): + continue + # + if any( + callback["group_id"] == finalizer["group_id"] + for callback in self.arbiter.callbacks.values() + ): + continue + # + for task_id in self.arbiter.add_task(finalizer["finalizer"]): + with self.arbiter.lock: + self.arbiter.group_state[finalizer["group_id"]].append(task_id) + # + with self.arbiter.lock: + self.arbiter.finalizers.pop(finalizer_id, None) + # + # Check event + # + with self.arbiter.lock: + if not self.arbiter.callbacks and not self.arbiter.finalizers: + self.arbiter.no_waiting_tasks.set() + else: + self.arbiter.no_waiting_tasks.clear() diff --git a/pylon/plugins/shared/__init__.py b/arbiter/eventnode/__init__.py similarity index 87% rename from pylon/plugins/shared/__init__.py rename to arbiter/eventnode/__init__.py index ad8ac8a..225be5c 100644 --- a/pylon/plugins/shared/__init__.py +++ b/arbiter/eventnode/__init__.py @@ -1,5 +1,6 @@ #!/usr/bin/python3 # coding=utf-8 +# pylint: disable=C0114 # Copyright 2024 getcarrier.io # @@ -15,6 +16,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" Module init """ -from .module import Module +from .tools import make_event_node + +from .mock import MockEventNode diff --git a/arbiter/eventnode/base.py b/arbiter/eventnode/base.py new file mode 100644 index 0000000..b16eafa --- /dev/null +++ b/arbiter/eventnode/base.py @@ -0,0 +1,241 @@ +#!/usr/bin/python3 +# coding=utf-8 + +# Copyright 2024 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" + Event node +""" + +import hmac +import gzip +import queue +import pickle +import threading + +from arbiter import log + +from .tools import make_event_node +from . import hooks + + +class EventNodeBase: # pylint: disable=R0902 + """ Event node (base) - allows to subscribe to events and to emit new events """ + + def __init__( + self, + hmac_key=None, hmac_digest="sha512", + callback_workers=1, + log_errors=True, + ): # pylint: disable=R0913 + self.clone_config = None + # + self.log_errors = log_errors + self.event_callbacks = {} # event_name -> [callbacks] + self.catch_all_callbacks = [] + # + self.before_callback_hooks = [] + self.after_callback_hooks = [] + # + self.hmac_key = hmac_key + self.hmac_digest = hmac_digest + if self.hmac_key is not None and isinstance(self.hmac_key, str): + self.hmac_key = self.hmac_key.encode("utf-8") + # + self.stop_event = threading.Event() + self.event_lock = threading.Lock() + self.sync_queue = queue.SimpleQueue() + # + self.listening_thread = threading.Thread(target=self.listening_worker, daemon=True) + self.callback_threads = [] + for _ in range(callback_workers): + self.callback_threads.append( + threading.Thread(target=self.callback_worker, daemon=True) + ) + # + self.ready_event = threading.Event() + self.can_emit = True + self.started = False + + def clone(self): + """ Make new event node with same config """ + if self.clone_config is None: + raise NotImplementedError + # + return make_event_node(config=self.clone_config) + + def start(self, emit_only=False): + """ Start event node """ + if self.started: + return + # + if emit_only: + self.ready_event.set() + else: + self.listening_thread.start() + for callback_thread in self.callback_threads: + callback_thread.start() + # + self.ready_event.wait() + self.started = True + + def stop(self): + """ Stop event node """ + self.stop_event.set() + + @property + def running(self): + """ Check if it is time to stop """ + return not self.stop_event.is_set() + + def subscribe(self, event_name, callback): + """ Subscribe to event """ + with self.event_lock: + if event_name is ...: + if callback not in self.catch_all_callbacks: + self.catch_all_callbacks.append(callback) + return + # + if event_name not in self.event_callbacks: + self.event_callbacks[event_name] = [] + if callback not in self.event_callbacks[event_name]: + self.event_callbacks[event_name].append(callback) + + def unsubscribe(self, event_name, callback): + """ Unsubscribe from event """ + with self.event_lock: + if event_name is ...: + if callback in self.catch_all_callbacks: + self.catch_all_callbacks.remove(callback) + return + # + if event_name not in self.event_callbacks: + return + if callback not in self.event_callbacks[event_name]: + return + self.event_callbacks[event_name].remove(callback) + + def emit(self, event_name, payload=None): + """ Emit event with payload data """ + if not self.can_emit: + return + # + data = self.make_event_data(event_name, payload) + self.emit_data(data) + + def make_event_data(self, event_name, payload=None): + """ Make event data """ + event = { + "name": event_name, + "payload": payload, + } + # + data = gzip.compress(pickle.dumps( + event, protocol=pickle.HIGHEST_PROTOCOL + )) + # + if self.hmac_key is not None: + digest = hmac.digest(self.hmac_key, data, self.hmac_digest) + data = data + digest + # + return data + + def emit_data(self, data): + """ Emit event data """ + raise NotImplementedError + + def listening_worker(self): + """ Listening thread: push event data to sync_queue """ + raise NotImplementedError + + def add_before_callback_hook(self, hook): + """ Register pre-callback hook """ + with self.event_lock: + if hook not in self.before_callback_hooks: + self.before_callback_hooks.append(hook) + + def remove_before_callback_hook(self, hook): + """ De-register pre-callback hook """ + with self.event_lock: + while hook in self.before_callback_hooks: + self.before_callback_hooks.remove(hook) + + def add_after_callback_hook(self, hook): + """ Register post-callback hook """ + with self.event_lock: + if hook not in self.after_callback_hooks: + self.after_callback_hooks.append(hook) + + def remove_after_callback_hook(self, hook): + """ De-register post-callback hook """ + with self.event_lock: + while hook in self.after_callback_hooks: + self.after_callback_hooks.remove(hook) + + def callback_worker(self): # pylint: disable=R0912 + """ Callback thread: call subscribers """ + while self.running: + try: + body = self.sync_queue.get() + # + if self.hmac_key is not None: + hmac_obj = hmac.new(self.hmac_key, digestmod=self.hmac_digest) + hmac_size = hmac_obj.digest_size + # + body_digest = body[-hmac_size:] + body = body[:-hmac_size] + # + digest = hmac.digest(self.hmac_key, body, self.hmac_digest) + # + if not hmac.compare_digest(body_digest, digest): + if self.log_errors: + log.error("Invalid event digest, skipping") + continue + # + event = pickle.loads(gzip.decompress(body)) + # + event_name = event.get("name") + event_payload = event.get("payload") + # + with self.event_lock: + callbacks = self.catch_all_callbacks.copy() + if event_name in self.event_callbacks: + callbacks.extend(self.event_callbacks[event_name]) + # + for callback in callbacks: + for hook in hooks.before_callback_hooks + self.before_callback_hooks: + try: + hook(callback, event_name, event_payload) + except: # pylint: disable=W0702 + if self.log_errors: + log.exception("Before callback hook failed, skipping") + # + try: + callback_result = callback(event_name, event_payload) + except: # pylint: disable=W0702 + if self.log_errors: + log.exception("Event callback failed, skipping") + # + callback_result = None # FIXME: pass exceptions to after_callback_hooks? + # + for hook in hooks.after_callback_hooks + self.after_callback_hooks: + try: + hook(callback, callback_result, event_name, event_payload) + except: # pylint: disable=W0702 + if self.log_errors: + log.exception("After callback hook failed, skipping") + except: # pylint: disable=W0702 + if self.log_errors: + log.exception("Error during event processing, skipping") diff --git a/arbiter/eventnode/hooks.py b/arbiter/eventnode/hooks.py new file mode 100644 index 0000000..626748b --- /dev/null +++ b/arbiter/eventnode/hooks.py @@ -0,0 +1,23 @@ +#!/usr/bin/python3 +# coding=utf-8 + +# Copyright 2025 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" + Event node +""" + +before_callback_hooks = [] +after_callback_hooks = [] diff --git a/arbiter/eventnode/mock.py b/arbiter/eventnode/mock.py new file mode 100644 index 0000000..e59da0a --- /dev/null +++ b/arbiter/eventnode/mock.py @@ -0,0 +1,50 @@ +#!/usr/bin/python3 +# coding=utf-8 + +# Copyright 2024 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" + Event node +""" + +from .base import EventNodeBase + + +class MockEventNode(EventNodeBase): # pylint: disable=R0902 + """ Event node (local-only mock) - allows to subscribe to events and to emit new events """ + + def __init__( + self, + hmac_key=None, hmac_digest="sha512", + callback_workers=1, + log_errors=True, + ): # pylint: disable=R0913 + super().__init__(hmac_key, hmac_digest, callback_workers, log_errors) + # + self.clone_config = { + "type": "MockEventNode", + "hmac_key": hmac_key, + "hmac_digest": hmac_digest, + "callback_workers": callback_workers, + "log_errors": log_errors, + } + + def emit_data(self, data): + """ Emit event data """ + self.sync_queue.put(data) + + def listening_worker(self): + """ Listening thread: push event data to sync_queue """ + self.ready_event.set() diff --git a/arbiter/eventnode/tools.py b/arbiter/eventnode/tools.py new file mode 100644 index 0000000..0e505ef --- /dev/null +++ b/arbiter/eventnode/tools.py @@ -0,0 +1,63 @@ +#!/usr/bin/python3 +# coding=utf-8 + +# Copyright 2024 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" + Event node +""" + +import os +import importlib + + +def make_event_node(config=None, env_prefix="EVENTNODE_"): + """ Make *EventNode instance """ + if config is None: + config = {} + # + int_vars = [ + "port", + "callback_workers", + "mute_first_failed_connections", + ] + # + bool_vars = [ + "use_ssl", + "ssl_verify", + ] + # + for key, value in os.environ.items(): + if key.startswith(env_prefix): + config_key = key[len(env_prefix):].lower() + config_value = value + # + if config_key in int_vars: + config_value = int(value) + # + if config_key in bool_vars: + config_value = value.lower() in ["true", "yes"] + # + config[config_key] = config_value + # + eventnode_cfg = config.copy() + eventnode_type = eventnode_cfg.pop("type") + # + eventnode_cls = getattr( + importlib.import_module("arbiter.eventnode"), + eventnode_type + ) + # + return eventnode_cls(**eventnode_cfg) diff --git a/arbiter/log.py b/arbiter/log.py new file mode 100644 index 0000000..b2bb89a --- /dev/null +++ b/arbiter/log.py @@ -0,0 +1,98 @@ +#!/usr/bin/python3 +# coding=utf-8 + +# Copyright 2020 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" + Logging tool +""" + +import logging +import inspect + +LOG_FORMAT = "%(asctime)s - %(levelname)8s - %(name)s - %(message)s" +LOG_DATE_FORMAT = "%Y.%m.%d %H:%M:%S %Z" + +initialized = False # pylint: disable=C0103 + + +def init(level=logging.INFO): + """ Initialize logging """ + global initialized # pylint: disable=W0603 + if initialized: + return + # + logging.basicConfig( + level=level, + datefmt=LOG_DATE_FORMAT, + format=LOG_FORMAT, + ) + logging.raiseExceptions = False + logging.getLogger("pika").setLevel(logging.WARNING) + # + initialized = True + + +def get_logger(): + """ Get logger for caller context """ + global initialized # pylint: disable=W0602,W0603 + if not initialized: + init() + # + return logging.getLogger( + inspect.currentframe().f_back.f_globals["__name__"] + ) + + +def get_outer_logger(): + """ Get logger for callers context (for use in this module) """ + return logging.getLogger( + inspect.currentframe().f_back.f_back.f_globals["__name__"] + ) + + +def debug(msg, *args, **kwargs): + """ Logs a message with level DEBUG """ + return get_outer_logger().debug(msg, *args, **kwargs) + + +def info(msg, *args, **kwargs): + """ Logs a message with level INFO """ + return get_outer_logger().info(msg, *args, **kwargs) + + +def warning(msg, *args, **kwargs): + """ Logs a message with level WARNING """ + return get_outer_logger().warning(msg, *args, **kwargs) + + +def error(msg, *args, **kwargs): + """ Logs a message with level ERROR """ + return get_outer_logger().error(msg, *args, **kwargs) + + +def critical(msg, *args, **kwargs): + """ Logs a message with level CRITICAL """ + return get_outer_logger().critical(msg, *args, **kwargs) + + +def log(lvl, msg, *args, **kwargs): + """ Logs a message with integer level lvl """ + return get_outer_logger().log(lvl, msg, *args, **kwargs) + + +def exception(msg, *args, **kwargs): + """ Logs a message with level ERROR inside exception handler """ + return get_outer_logger().exception(msg, *args, **kwargs) diff --git a/arbiter/minion.py b/arbiter/minion.py new file mode 100644 index 0000000..3adb963 --- /dev/null +++ b/arbiter/minion.py @@ -0,0 +1,91 @@ +#!/usr/bin/python3 +# coding=utf-8 +# pylint: disable=C0114,C0115,C0116 + +# Copyright 2023 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from arbiter import log + +from .task import Task +from .tasknode import TaskNode + + +class Minion: + def __init__(self, event_node, queue="default"): + self.queue = queue + self.raw_task_node = TaskNode(event_node, pool=self.queue, task_limit=0) + + @property + def task_node(self): + if not self.raw_task_node.started: + self.raw_task_node.start() + return self.raw_task_node + + def wait_for_tasks(self, tasks): + for task in tasks: + result = self.task_node.join_task(task) + yield { + "task_type": "task", + "state": "done", + "result": result, + } + + def add_task(self, task, sync=False): + tasks = [] + for _ in range(task.tasks_count): + task_key = self.task_node.start_task( + name=task.name, + args=task.task_args, + kwargs=task.task_kwargs, + pool=task.queue + ) + tasks.append(task_key) + yield task_key + if sync: + for message in self.wait_for_tasks(tasks): + yield message + + def apply(self, task_name, queue=None, tasks_count=1, task_args=None, task_kwargs=None, sync=True): # pylint: disable=C0301,R0913 + task = Task(task_name, queue=queue if queue else self.queue, + tasks_count=tasks_count, task_args=task_args, task_kwargs=task_kwargs) + for message in self.add_task(task, sync=sync): + yield message + + def task(self, *args, **kwargs): # pylint: disable=W0613 + """ Task decorator """ + def inner_task(func): + def create_task(**kwargs): + def _create_task(func): + return self._create_task_from_callable(func, **kwargs) + return _create_task + if callable(func): + return create_task(**kwargs)(func) + raise TypeError('@task decorated function must be callable') + return inner_task + + def _create_task_from_callable(self, func, name=None, **kwargs): # pylint: disable=W0613 + name = name if name else f"{func.__name__}.{func.__module__}" + self.raw_task_node.register_task(func, name) + return func + + def run(self, workers, block=True): + log.info("Starting '%s' worker", self.queue) + # + self.raw_task_node.task_limit = workers + self.raw_task_node.start() + # + if block: + self.raw_task_node.stop_event.wait() diff --git a/arbiter/task.py b/arbiter/task.py new file mode 100644 index 0000000..03c6228 --- /dev/null +++ b/arbiter/task.py @@ -0,0 +1,55 @@ +#!/usr/bin/python3 +# coding=utf-8 +# pylint: disable=C0114,C0115,C0116 + +# Copyright 2020 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from time import time + + +class Task: # pylint: disable=R0902,R0903 + def __init__(self, name, queue='default', tasks_count=1, task_key="", task_type="task", # pylint: disable=R0913 + task_args=None, task_kwargs=None, callback=False, callback_queue=None, timeout=-1): + if not task_args: + task_args = [] + if not task_kwargs: + task_kwargs = {} + self.task_type = task_type + self.task_key = task_key + self.name = name + self.queue = queue + self.tasks_count = tasks_count + self.task_args = task_args + self.task_kwargs = task_kwargs + self.callback = callback + self.callback_queue = callback_queue + self.tasks_array = [] # this is for a task ids that need to be verified to be done before callback # pylint: disable=C0301 + self.timeout = timeout # timeout in seconds. Works only with task_type=finalize + self.start_time = int(time()) + + def to_json(self): + return { + "type": self.task_type, + "queue": self.queue, + "task_name": self.name, + "task_key": self.task_key, + "args": self.task_args, + "kwargs": self.task_kwargs, + "arbiter": self.callback_queue, + "callback": self.callback, + "tasks_array": self.tasks_array, + "timeout": self.timeout, + "start_time": self.start_time + } diff --git a/pylon/plugins/test_api/__init__.py b/arbiter/tasknode/__init__.py similarity index 92% rename from pylon/plugins/test_api/__init__.py rename to arbiter/tasknode/__init__.py index ad8ac8a..d60452a 100644 --- a/pylon/plugins/test_api/__init__.py +++ b/arbiter/tasknode/__init__.py @@ -1,5 +1,6 @@ #!/usr/bin/python3 # coding=utf-8 +# pylint: disable=C0114 # Copyright 2024 getcarrier.io # @@ -15,6 +16,5 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" Module init """ -from .module import Module +from .tasknode import TaskNode diff --git a/arbiter/tasknode/housekeeper.py b/arbiter/tasknode/housekeeper.py new file mode 100644 index 0000000..73e1050 --- /dev/null +++ b/arbiter/tasknode/housekeeper.py @@ -0,0 +1,61 @@ +#!/usr/bin/python3 +# coding=utf-8 + +# Copyright 2024 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" + Task node +""" + +import time +import datetime +import threading + + +class TaskNodeHousekeeper(threading.Thread): # pylint: disable=R0903 + """ Perform cleanup """ + + def __init__(self, node): + super().__init__(daemon=True) + self.node = node + + def run(self): + """ Run housekeeper thread """ + while not self.node.stop_event.is_set(): + time.sleep(self.node.housekeeping_interval) + # + with self.node.lock: + for task_id in list(self.node.state_events): + data = self.node.state_events[task_id] + # + if not data["event"].is_set(): + continue + # + age = (datetime.datetime.now() - data["timestamp"]).total_seconds() + # + if age < self.node.task_retention_period: + continue + # + self.node.state_events.pop(task_id, None) + self.node.global_task_state.pop(task_id, None) + self.node.known_task_ids.discard(task_id) + # + self.node.event_node.emit( + "task_status_change", + { + "task_id": task_id, + "status": "pruned", + } + ) diff --git a/arbiter/tasknode/tasknode.py b/arbiter/tasknode/tasknode.py new file mode 100644 index 0000000..c759ddd --- /dev/null +++ b/arbiter/tasknode/tasknode.py @@ -0,0 +1,1168 @@ +#!/usr/bin/python3 +# coding=utf-8 +# pylint: disable=C0116,C0302 + +# Copyright 2023-2025 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" + Task node + + Allows to start, register, query tasks and workers + + Uses existing EventNode as a transport +""" + +import os +import gzip +import time +import uuid +import queue +import ctypes +import pickle +import datetime +import threading +import functools +import traceback +import multiprocessing +import multiprocessing.connection + +from arbiter import log + +from ..eventnode.tools import make_event_node +from .housekeeper import TaskNodeHousekeeper +from .watcher import TaskNodeWatcher +from .tools import InterruptTaskThread +from .tools import reap_zombies + + +class TaskNode: # pylint: disable=R0902,R0904 + """ Task node - start, register, query tasks and workers """ + + def __init__( # pylint: disable=R0913,R0914 + self, event_node, + pool=None, task_limit=None, ident_prefix="", + multiprocessing_context="fork", kill_on_stop=False, + task_retention_period=3600, housekeeping_interval=60, + start_max_wait=3, query_wait=3, + watcher_max_wait=3, stop_node_task_wait=3, result_max_wait=3, + tmp_path="/tmp/tasknode", result_transport="memory", + start_attempts=3, thread_scan_interval=1, + task_approver=None, + ): + self.event_node = event_node + self.event_node_was_started = False + # + self.ident_prefix = ident_prefix + self.ident = None + self.pool = pool + # + self.sync_queues = {} + self.task_registry = {} + self.running_tasks = {} + self.local_tasks = {} + # + self.global_pool_state = {} + self.global_task_state = {} + # + self.have_running_tasks = threading.Event() + self.known_task_ids = set() + self.state_events = {} + self.task_status_subscribers = [] + # + self.multiprocessing_context = multiprocessing_context + self.kill_on_stop = kill_on_stop + self.task_limit = task_limit + self.task_retention_period = task_retention_period + # + self.tmp_path = tmp_path + self.result_transport = result_transport + # + self.housekeeping_interval = housekeeping_interval + self.start_max_wait = start_max_wait + self.query_wait = query_wait + self.watcher_max_wait = watcher_max_wait + self.stop_node_task_wait = stop_node_task_wait + self.result_max_wait = result_max_wait + # + self.start_attempts = start_attempts + self.thread_scan_interval = thread_scan_interval + # + self.lock = threading.Lock() + self.stop_event = threading.Event() + self.started = False + # + self.task_approver = task_approver + + # + # Node start and stop + # + + def start(self, block=False): + """ Start task node """ + if self.started: + return + # + self.stop_event.clear() + # + if not self.event_node.started: + self.event_node.start() + self.event_node_was_started = True + # + self.ident = f'{self.ident_prefix}{str(uuid.uuid4())}' + # + if self.result_transport == "files": + os.makedirs(self.tmp_path, exist_ok=True) + elif self.result_transport == "events": + self.event_node.subscribe("task_result_payload", self.on_result_payload) + # + self.event_node.subscribe("task_node_announce", self.on_node_announce) + self.event_node.subscribe("task_node_withhold", self.on_node_withhold) + # + self.event_node.subscribe("task_start_query", self.on_start_query) + self.event_node.subscribe("task_start_candidate", self.on_sync_reply) + self.event_node.subscribe("task_start_request", self.on_start_request) + self.event_node.subscribe("task_start_ack", self.on_sync_reply) + # + self.event_node.subscribe("task_stop_request", self.on_stop_request) + self.event_node.subscribe("task_state_announce", self.on_state_announce) + # + self.event_node.subscribe("task_state_query", self.on_state_query) + self.event_node.subscribe("task_state_reply", self.on_state_reply) + self.event_node.subscribe("task_pool_query", self.on_pool_query) + self.event_node.subscribe("task_pool_reply", self.on_pool_reply) + # + TaskNodeWatcher(self).start() + TaskNodeHousekeeper(self).start() + # + self.event_node.emit( + "task_node_announce", + { + "ident": self.ident, + "pool": self.pool, + "task_limit": self.task_limit, + "running_tasks": 0, + } + ) + # + self.started = True + # + if block: + self.stop_event.wait() + + def stop(self, block=True): + """ Stop task node """ + self.event_node.unsubscribe("task_node_announce", self.on_node_announce) + self.event_node.unsubscribe("task_node_withhold", self.on_node_withhold) + # + self.event_node.emit( + "task_node_withhold", + { + "ident": self.ident, + } + ) + # + self.event_node.unsubscribe("task_start_query", self.on_start_query) + self.event_node.unsubscribe("task_start_candidate", self.on_sync_reply) + self.event_node.unsubscribe("task_start_request", self.on_start_request) + self.event_node.unsubscribe("task_start_ack", self.on_sync_reply) + # + for task_id in list(self.running_tasks): + self.stop_task(task_id) + if block: + self.wait_for_task(task_id, self.stop_node_task_wait) + # + self.event_node.unsubscribe("task_stop_request", self.on_stop_request) + self.event_node.unsubscribe("task_state_announce", self.on_state_announce) + # + if self.result_transport == "events": + self.event_node.unsubscribe("task_result_payload", self.on_result_payload) + # + self.event_node.unsubscribe("task_state_query", self.on_state_query) + self.event_node.unsubscribe("task_state_reply", self.on_state_reply) + self.event_node.unsubscribe("task_pool_query", self.on_pool_query) + self.event_node.unsubscribe("task_pool_reply", self.on_pool_reply) + # + while self.task_status_subscribers: + subscriber = self.task_status_subscribers.pop() + self.event_node.unsubscribe("task_status_change", subscriber) + # + if self.event_node_was_started: + self.event_node.stop() + # + self.started = False + self.stop_event.set() + + # + # Task registration + # + + def register_task(self, func, name=None, approver=None): + """ Register task function """ + if name is None: + name = self.get_callable_name(func) + # + with self.lock: + self.task_registry[name] = [func, approver] + + def unregister_task(self, func=None, name=None): + """ Unregister task function """ + if name is None and func is None: + raise ValueError("Missing name or func") + # + if name is None: + name = self.get_callable_name(func) + # + with self.lock: + if name in self.task_registry: + self.task_registry.pop(name) + + # + # Task start and stop + # + + def start_task(self, name, args=None, kwargs=None, pool=None, meta=None, durable=False): # pylint: disable=R0913 + """ Start task execution """ + for _ in range(self.start_attempts): + task_id = self.start_task_attempt(name, args, kwargs, pool, meta, durable) + if task_id is not None: + return task_id + # + return None + + def start_task_attempt(self, name, args=None, kwargs=None, pool=None, meta=None, durable=False): # pylint: disable=R0913 + """ Try to start task execution """ + if meta is not None and not isinstance(meta, dict): + raise ValueError("Meta must be None or dict") + # + task_id = self.generate_task_id() + # + self.event_node.emit( + "task_state_announce", + { + "task_id": task_id, + "requestor": self.ident, + "runner": None, + "status": "pending", + "result": None, + "meta": meta, + } + ) + # + self.event_node.emit( + "task_status_change", + { + "task_id": task_id, + "status": "pending", + } + ) + # + query_queue = f'task_start_query_{task_id}' + ack_queue = f'task_start_ack_{task_id}' + # + with self.lock: + self.sync_queues[query_queue] = queue.SimpleQueue() + self.sync_queues[ack_queue] = queue.SimpleQueue() + # + self.event_node.emit( + "task_start_query", + { + "name": name, + "pool": pool, + "task_id": task_id, + "requestor": self.ident, + "sync_queue": query_queue, + } + ) + # + try: + while True: + try: + candidate = self.sync_queues[query_queue].get(timeout=self.start_max_wait) + # + self.event_node.emit( + "task_start_request", + { + "name": name, + "meta": meta, + "args": args, + "kwargs": kwargs, + "durable": durable, + "pool": pool, + "task_id": task_id, + "runner": candidate.get("ident"), + "requestor": self.ident, + "sync_queue": ack_queue, + } + ) + # + try: + self.sync_queues[ack_queue].get(timeout=self.start_max_wait) + except: # pylint: disable=W0702 + continue # try next candidate if present + # + return task_id + except: # pylint: disable=W0702 + self.event_node.emit( + "task_state_announce", + { + "task_id": task_id, + "requestor": self.ident, + "runner": None, + "status": "stopped", + "result": None, + "meta": meta, + } + ) + # + self.event_node.emit( + "task_status_change", + { + "task_id": task_id, + "status": "stopped", + } + ) + # + return None + finally: + with self.lock: + self.sync_queues.pop(query_queue) + self.sync_queues.pop(ack_queue) + + def stop_task(self, task_id): + """ Stop running task """ + self.event_node.emit( + "task_stop_request", + { + "task_id": task_id, + "requestor": self.ident, + } + ) + + # + # Wait for task / join on task + # + + def wait_for_task(self, task_id, timeout=None): + """ Wait for task to stop """ + if task_id not in self.state_events: + self.query_task_state(task_id) + # + if task_id not in self.state_events: + raise RuntimeError("Unknown task") + # + self.state_events[task_id]["event"].wait(timeout) + + def join_task(self, task_id, timeout=None): + """ Wait for task to stop and get task result """ + self.wait_for_task(task_id, timeout) + return self.get_task_result(task_id) + + # + # Task status, meta and result + # + + def get_task_status(self, task_id): + """ Get task status """ + if task_id not in self.global_task_state: + self.query_task_state(task_id) + # + if task_id not in self.global_task_state: + raise RuntimeError("Unknown task") + # + return self.global_task_state[task_id].get("status", "unknown") + + def get_task_meta(self, task_id): + """ Get task meta """ + if task_id not in self.global_task_state: + self.query_task_state(task_id) + # + if task_id not in self.global_task_state: + raise RuntimeError("Unknown task") + # + meta = self.global_task_state[task_id].get("meta", None) + if meta is None: + meta = {} + # + return meta.copy() + + def get_task_result(self, task_id): + """ Get task result """ + if task_id not in self.global_task_state: + self.query_task_state(task_id) + # + if task_id not in self.global_task_state: + raise RuntimeError("Unknown task") + # + result = self.global_task_state[task_id].get("result", None) + # + if result is None: + return ... # invalid result or task is still running + # + result = pickle.loads(gzip.decompress(result)) + # + if "return" in result: + return result["return"] + # + if "raise" in result: + exception_data = "\n".join(["", result["raise"]]) + # + if exception_data.rstrip().endswith("arbiter.tasknode.tools.InterruptTaskThread"): + return ... # task was stopped by stop_task + # + raise Exception(exception_data) # pylint: disable=E0012,W0719 + # + return ... # invalid result + + def subscribe_to_task_statuses(self, func): + """ Subscribe to task status changes """ + self.event_node.subscribe("task_status_change", func) + with self.lock: + self.task_status_subscribers.append(func) + + # + # Node network queries + # + + def query_task_state(self, task_id=None): + """ Sync info from other nodes """ + self.event_node.emit( + "task_state_query", + { + "task_id": task_id, + "requestor": self.ident, + } + ) + # + time.sleep(self.query_wait) + + def query_pool_state(self, pool=None): + """ Sync info from other nodes """ + self.event_node.emit( + "task_pool_query", + { + "pool": pool, + "requestor": self.ident, + } + ) + # + time.sleep(self.query_wait) + + def count_free_workers(self, pool=None): + """ Get task limit (how many can we start now) """ + self.query_pool_state(pool) + # + if pool not in self.global_pool_state: + return 0 + # + free = 0 + # + for data in self.global_pool_state[pool].values(): + if data["task_limit"] is None: + return ... # unlimited + # + free += (data["task_limit"] - data["running_tasks"]) + # + return free + + # + # Event handlers + # + + def on_node_announce(self, event_name, event_payload): + _ = event_name + # + if "for_requestor" in event_payload and event_payload.get("for_requestor") != self.ident: + return + # + if "ident" not in event_payload: + return + # + ident = event_payload.get("ident") + pool = event_payload.get("pool", None) + # + with self.lock: + if pool not in self.global_pool_state: + self.global_pool_state[pool] = {} + # + self.global_pool_state[pool][ident] = event_payload.copy() + + def on_node_withhold(self, event_name, event_payload): + _ = event_name + # + if "ident" not in event_payload: + return + # + ident = event_payload.get("ident") + # + with self.lock: + for _, nodes in self.global_pool_state.items(): + nodes.pop(ident, None) + + def on_stop_request(self, event_name, event_payload): + _ = event_name + # + if "task_id" not in event_payload: + return + # + task_id = event_payload.get("task_id") + # + with self.lock: + if task_id in self.local_tasks: + self.local_tasks[task_id]["durable"] = False + # + if self.multiprocessing_context in ["threading"]: + self._stop_task__threading(task_id) + else: + self._stop_task__multiprocessing(task_id) + + def _stop_task__threading(self, task_id): + if task_id not in self.running_tasks: + return + # + with self.lock: + data = self.running_tasks.get(task_id, {}) + thread = data.get("thread", None) + # + if thread is not None: + # Note: this way will not stop running blocking system calls (e.g. sleep()) + # May try to use pthread_kill to interrupt if possible in the future + # Also can do some error checks (e.g. if exception was set to multiple threads) + ctypes.pythonapi.PyThreadState_SetAsyncExc( + ctypes.c_ulong(thread.ident), + ctypes.py_object(InterruptTaskThread), + ) + + def _stop_task__multiprocessing(self, task_id): + if task_id not in self.running_tasks: + return + # + with self.lock: + data = self.running_tasks.get(task_id, {}) + process = data.get("process", None) + # + if process is not None: + if self.kill_on_stop: + process.kill() + else: + process.terminate() + + def on_state_announce(self, event_name, event_payload): + _ = event_name + # + if "for_requestor" in event_payload and event_payload.get("for_requestor") != self.ident: + return + # + if "task_id" not in event_payload: + return + # + task_id = event_payload.get("task_id") + task_status = event_payload.get("status", "unknown") + # + with self.lock: + self.global_task_state[task_id] = event_payload.copy() + self.known_task_ids.add(task_id) + # + if task_id not in self.state_events: + self.state_events[task_id] = { + "event": threading.Event(), + } + # + self.state_events[task_id]["timestamp"] = datetime.datetime.now() + # + if task_status == "stopped": + self.state_events[task_id]["event"].set() + + def on_result_payload(self, event_name, event_payload): + _ = event_name + # + task_id = event_payload.get("task_id") + payload = event_payload.get("payload") + # + with self.lock: + if task_id not in self.running_tasks: + return + # + self.running_tasks[task_id]["result"] = payload + + def on_state_query(self, event_name, event_payload): + _ = event_name + # + if event_payload.get("requestor", None) == self.ident: + return + # + if event_payload.get("task_id", None) is not None: + task_id = event_payload.get("task_id") + # + if task_id not in self.global_task_state: + return + # + task_state = self.global_task_state[task_id].copy() + task_state["for_requestor"] = event_payload.get("requestor", None) + # + self.event_node.emit( + "task_state_announce", + task_state + ) + else: + self.event_node.emit( + "task_state_reply", + { + "for_requestor": event_payload.get("requestor", None), + "global_task_state": self.global_task_state, + } + ) + + def on_state_reply(self, event_name, event_payload): + _ = event_name + # + if event_payload.get("for_requestor", None) != self.ident: + return + # + if "global_task_state" not in event_payload: + return + # + global_task_state = event_payload.get("global_task_state") + # + with self.lock: + for task_id in list(self.global_task_state): + if task_id in self.running_tasks: + global_task_state.pop(task_id, None) + else: + self.global_task_state.pop(task_id, None) + # + self.global_task_state.update(global_task_state) + + def on_pool_query(self, event_name, event_payload): + _ = event_name + # + if event_payload.get("requestor", None) == self.ident: + return + # + if event_payload.get("pool", None) is not None: + pool = event_payload.get("pool") + # + if pool not in self.global_pool_state: + return + # + global_pool_state = { + pool: self.global_pool_state[pool] + } + else: + global_pool_state = self.global_pool_state + # + self.event_node.emit( + "task_pool_reply", + { + "for_requestor": event_payload.get("requestor", None), + "global_pool_state": global_pool_state, + } + ) + + def on_pool_reply(self, event_name, event_payload): + _ = event_name + # + if event_payload.get("for_requestor", None) != self.ident: + return + # + if "global_pool_state" not in event_payload: + return + # + global_pool_state = event_payload.get("global_pool_state") + # + with self.lock: + for pool, state in global_pool_state.items(): + if pool not in self.global_pool_state: + self.global_pool_state[pool] = state + continue + # + for ident in list(self.global_pool_state[pool]): + if ident == self.ident: + continue + # + self.global_pool_state[pool].pop(ident, None) + # + state.pop(self.ident, None) + self.global_pool_state[pool].update(state) + + def on_sync_reply(self, event_name, event_payload): + _ = event_name + # + if event_payload.get("for_requestor", None) != self.ident: + return + # + if event_payload.get("sync_queue", None) not in self.sync_queues: + return + # + self.sync_queues[event_payload.get("sync_queue")].put(event_payload.copy()) + + def on_start_query(self, event_name, event_payload): # pylint: disable=R0911 + _ = event_name + # + task_name = event_payload.get("name", None) + # + if task_name not in self.task_registry: + return + # + if event_payload.get("pool", None) != self.pool: + return + # + if self.task_limit is not None and len(self.running_tasks) >= self.task_limit: + return + # + if self.task_approver is not None: + try: + if not self.task_approver(event_name, event_payload): + return + except: # pylint: disable=W0702 + return + # + approver = self.task_registry[task_name][1] + # + if approver is not None: + try: + if not approver(event_name, event_payload): + return + except: # pylint: disable=W0702 + return + # + self.event_node.emit( + "task_start_candidate", + { + "ident": self.ident, + "for_requestor": event_payload.get("requestor", None), + "sync_queue": event_payload.get("sync_queue", None), + } + ) + + def on_start_request(self, event_name, event_payload): # pylint: disable=R0911 + _ = event_name + # + if event_payload.get("runner", None) != self.ident: + return + # + task_name = event_payload.get("name", None) + # + if task_name not in self.task_registry: + return + # + if event_payload.get("pool", None) != self.pool: + return + # + if self.task_limit is not None and len(self.running_tasks) >= self.task_limit: + return + # + if self.task_approver is not None: + try: + if not self.task_approver(event_name, event_payload): + return + except: # pylint: disable=W0702 + return + # + approver = self.task_registry[task_name][1] + # + if approver is not None: + try: + if not approver(event_name, event_payload): + return + except: # pylint: disable=W0702 + return + # + self.event_node.emit( + "task_start_ack", + { + "for_requestor": event_payload.get("requestor", None), + "sync_queue": event_payload.get("sync_queue", None), + } + ) + # + self.event_node.emit( + "task_state_announce", + { + "task_id": event_payload.get("task_id", None), + "requestor": event_payload.get("requestor", None), + "runner": self.ident, + "status": "running", + "result": None, + "meta": event_payload.get("meta", None), + } + ) + # + self.event_node.emit( + "task_status_change", + { + "task_id": event_payload.get("task_id", None), + "status": "running", + } + ) + # + self.execute_local_task( + event_payload.get("task_id", None), + event_payload.get("name", None), + event_payload.get("meta", None), + event_payload.get("args", None), + event_payload.get("kwargs", None), + event_payload.get("durable", False), + event_payload.get("pool", None), + ) + + # + # Tools + # + + def generate_task_id(self): + """ Get 'mostly' safe new task_id """ + with self.lock: + while True: + task_id = str(uuid.uuid4()) + # + if task_id in self.known_task_ids: + continue + # + self.known_task_ids.add(task_id) + break + # + return task_id + + def execute_local_task( # pylint: disable=R0913 + self, task_id, name, meta, args=None, kwargs=None, durable=False, pool=None, + ): + """ Start task from task registry """ + with self.lock: + self.local_tasks[task_id] = { + "name": name, + "meta": meta, + "args": args, + "kwargs": kwargs, + "durable": durable, + "pool": pool, + } + # + if self.multiprocessing_context in ["threading"]: + self._execute_local_task__threading(task_id, name, meta, args, kwargs, pool) + else: + self._execute_local_task__multiprocessing(task_id, name, meta, args, kwargs, pool) + + def _execute_local_task__threading( # pylint: disable=R0913 + self, task_id, name, meta, args=None, kwargs=None, pool=None, + ): + if name not in self.task_registry: + raise RuntimeError("Task not found") + # + if meta is None: + meta = {} + if args is None: + args = () + if kwargs is None: + kwargs = {} + # + result = None + if self.result_transport == "files": + result_config = self.tmp_path + elif self.result_transport == "events": + result_config = self.event_node.clone_config.copy() + elif self.result_transport == "memory": + result_config = queue.SimpleQueue() + result = result_config + else: + raise RuntimeError(f"Invalid result transport: {self.result_transport}") + # + with self.lock: + import sys # pylint: disable=C0415 + if "tasknode_task" not in sys.modules: + sys.modules["tasknode_task"] = threading.local() + # + thread = threading.Thread( + target=self.executor, + name=f'tasknode_task {task_id}', + args=(), + kwargs={ + "name": name, + "target": self.task_registry[name][0], + "task_id": task_id, + "meta": meta, + "args": args, + "kwargs": kwargs, + "result_transport": self.result_transport, + "result_config": result_config, + "multiprocessing_context": self.multiprocessing_context, + "pool": pool, + }, + daemon=True, + ) + thread.start() + # + with self.lock: + self.running_tasks[task_id] = { + "thread": thread, + "result": result, + } + self.have_running_tasks.set() + # + self.event_node.emit( + "task_node_announce", + { + "ident": self.ident, + "pool": self.pool, + "task_limit": self.task_limit, + "running_tasks": len(self.running_tasks), + } + ) + + def _execute_local_task__multiprocessing( # pylint: disable=R0913 + self, task_id, name, meta, args=None, kwargs=None, pool=None, + ): + if name not in self.task_registry: + raise RuntimeError("Task not found") + # + if meta is None: + meta = {} + if args is None: + args = () + if kwargs is None: + kwargs = {} + # + multiprocessing_ctx = multiprocessing.get_context(self.multiprocessing_context) + # + result = None + if self.result_transport == "files": + result_config = self.tmp_path + elif self.result_transport == "events": + result_config = self.event_node.clone_config.copy() + elif self.result_transport == "memory": + result_config = multiprocessing_ctx.Queue() + result = result_config + else: + raise RuntimeError(f"Invalid result transport: {self.result_transport}") + # + process = multiprocessing_ctx.Process( + target=self.executor, + args=(), + kwargs={ + "name": name, + "target": self.task_registry[name][0], + "task_id": task_id, + "meta": meta, + "args": args, + "kwargs": kwargs, + "result_transport": self.result_transport, + "result_config": result_config, + "multiprocessing_context": self.multiprocessing_context, + "pool": pool, + }, + daemon=False, + ) + process.start() + # + process_pid = process.pid + if process_pid is not None: + try: + import pylon # pylint: disable=C0415,E0401,W0611 + from tools import context # pylint: disable=C0415,E0401 + # + context.zombie_reaper.external_pids.add(process_pid) + except: # pylint: disable=W0702 + pass + # + with self.lock: + self.running_tasks[task_id] = { + "process": process, + "result": result, + } + self.have_running_tasks.set() + # + self.event_node.emit( + "task_node_announce", + { + "ident": self.ident, + "pool": self.pool, + "task_limit": self.task_limit, + "running_tasks": len(self.running_tasks), + } + ) + + def executor( + self, + name, target, task_id, meta, args, kwargs, + result_transport, result_config, multiprocessing_context, + pool, + ): # pylint: disable=R0913,R0914 + """ Task executor """ + if multiprocessing_context in ["threading"]: + self._executor__threading( + name, target, task_id, meta, args, kwargs, + result_transport, result_config, multiprocessing_context, + pool, + ) + else: + self._executor__multiprocessing( + name, target, task_id, meta, args, kwargs, + result_transport, result_config, multiprocessing_context, + pool, + ) + + def _executor__threading( + self, + name, target, task_id, meta, args, kwargs, + result_transport, result_config, multiprocessing_context, + pool, + ): # pylint: disable=R0913,R0914 + try: + import setproctitle # pylint: disable=C0415,E0401 + setproctitle.setthreadtitle(f'tasknode_task {task_id}') + # + import sys # pylint: disable=C0415 + sys.modules["tasknode_task"].id = task_id + sys.modules["tasknode_task"].meta = meta.copy() + sys.modules["tasknode_task"].name = name + sys.modules["tasknode_task"].pool = pool + sys.modules["tasknode_task"].multiprocessing_context = multiprocessing_context + # + try: + output = target(*args, **kwargs) + data = {"return": output} + except: # pylint: disable=W0702 + error = traceback.format_exc() + data = {"raise": error} + # + result = gzip.compress(pickle.dumps( + data, protocol=pickle.HIGHEST_PROTOCOL + )) + # + if result_transport == "files": + with open(os.path.join(result_config, f'{task_id}.bin'), "wb") as file: + file.write(result) + # + elif result_transport == "events": + result_event_node = make_event_node(config=result_config) + result_event_node.start(emit_only=True) + result_event_node.emit("task_result_payload", { + "task_id": task_id, + "payload": result, + }) + result_event_node.stop() + # + elif result_transport == "memory": + result_config.put(result) + # + else: + raise RuntimeError(f"Invalid result transport: {result_transport}") + except: # pylint: disable=W0702 + log.exception("Task execution failed") + # + raise + + def _executor__multiprocessing( + self, + name, target, task_id, meta, args, kwargs, + result_transport, result_config, multiprocessing_context, + pool, + ): # pylint: disable=R0912,R0913,R0914,R0915 + try: + if multiprocessing_context == "fork": + # Clear TaskNode->EventNode. Do not attempt to close connections + self.event_node.can_emit = False + self.event_node.event_callbacks = {} + self.event_node.catch_all_callbacks = [] + self.running_tasks = {} + # Drop importlib locks + import importlib._bootstrap # pylint: disable=C0415 + importlib._bootstrap._module_locks = {} # pylint: disable=W0212 + # Re-init for SSL + import ssl # pylint: disable=C0415 + ssl.RAND_bytes(1) + # Signals + import signal # pylint: disable=C0415 + for sig in [signal.SIGTERM, signal.SIGINT]: + signal.signal(sig, lambda *x, **y: os._exit(0)) # pylint: disable=W0212 + # Re-open stderr to try to mitigate logging locks + import sys # pylint: disable=C0415 + prev_stderr = sys.stderr + new_stderr = open(os.dup(sys.stderr.fileno()), sys.stderr.mode) # pylint: disable=W1514,R1732 + sys.stderr = new_stderr + # Change logging stderr streams to new one + import logging # pylint: disable=C0415 + for handler in list(logging.root.handlers): + if not isinstance(handler, logging.StreamHandler): + continue + if handler.stream == prev_stderr: + handler.stream = new_stderr + # Think more about gevent? Logging re-init? Base pylon re-init here? + # + import setproctitle # pylint: disable=C0415,E0401 + setproctitle.setproctitle(f'tasknode_task {task_id}') + # + import sys # pylint: disable=C0415 + import types # pylint: disable=C0415 + sys.modules["tasknode_task"] = types.ModuleType("tasknode_task") + sys.modules["tasknode_task"].__path__ = [] + setattr(sys.modules["tasknode_task"], "id", task_id) + setattr(sys.modules["tasknode_task"], "meta", meta.copy()) + setattr(sys.modules["tasknode_task"], "name", name) + setattr(sys.modules["tasknode_task"], "pool", pool) + setattr( + sys.modules["tasknode_task"], "multiprocessing_context", multiprocessing_context + ) + # + try: + output = target(*args, **kwargs) + data = {"return": output} + except: # pylint: disable=W0702 + error = traceback.format_exc() + data = {"raise": error} + # + result = gzip.compress(pickle.dumps( + data, protocol=pickle.HIGHEST_PROTOCOL + )) + # + if result_transport == "files": + with open(os.path.join(result_config, f'{task_id}.bin'), "wb") as file: + file.write(result) + # + elif result_transport == "events": + result_event_node = make_event_node(config=result_config) + result_event_node.start(emit_only=True) + result_event_node.emit("task_result_payload", { + "task_id": task_id, + "payload": result, + }) + result_event_node.stop() + # + elif result_transport == "memory": + result_config.put(result) + result_config.close() + result_config.join_thread() + # + else: + raise RuntimeError(f"Invalid result transport: {result_transport}") + except: # pylint: disable=W0702 + log.exception("Task execution failed") + # + if multiprocessing_context == "fork": + reap_zombies() + os._exit(1) # pylint: disable=W0212 + # + raise + # + if multiprocessing_context == "fork": + reap_zombies() + os._exit(0) # pylint: disable=W0212 + + def get_callable_name(self, func): + """ Get callable name """ + if hasattr(func, "__name__"): + return func.__name__ + if isinstance(func, functools.partial): + return self.get_callable_name(func.func) + raise ValueError("Cannot guess callable name") diff --git a/arbiter/tasknode/tools.py b/arbiter/tasknode/tools.py new file mode 100644 index 0000000..61c535f --- /dev/null +++ b/arbiter/tasknode/tools.py @@ -0,0 +1,48 @@ +#!/usr/bin/python3 +# coding=utf-8 + +# Copyright 2024 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" + Task node +""" + +import os + +from arbiter import log + + +def reap_zombies(): + """ Reap zombie processes """ + while True: + try: + child_siginfo = os.waitid(os.P_ALL, os.getpid(), os.WEXITED | os.WNOHANG) # pylint: disable=E1101 + # + if child_siginfo is None: + break + # + log.info( + "Reaped child process: %s -> %s -> %s", + child_siginfo.si_pid, + child_siginfo.si_code, + child_siginfo.si_status, + ) + except: # pylint: disable=W0702 + break + + +class InterruptTaskThread(Exception): + """ Special exception sent to thread in stop_task """ + pass # pylint: disable=W0107 diff --git a/arbiter/tasknode/watcher.py b/arbiter/tasknode/watcher.py new file mode 100644 index 0000000..bae6bd2 --- /dev/null +++ b/arbiter/tasknode/watcher.py @@ -0,0 +1,320 @@ +#!/usr/bin/python3 +# coding=utf-8 + +# Copyright 2024 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" + Task node +""" + +import os +import time +import datetime +import threading +import multiprocessing + +from arbiter import log + + +class TaskNodeWatcher(threading.Thread): # pylint: disable=R0903 + """ Watch running tasks """ + + def __init__(self, node): + super().__init__(daemon=True) + self.node = node + + def run(self): + """ Run watcher thread """ + # + while not self.node.stop_event.is_set(): + try: + if self.node.multiprocessing_context in ["threading"]: + self._watch_stopped_tasks__threading() + else: + self._watch_stopped_tasks__multiprocessing() + except: # pylint: disable=W0702 + log.exception("Exception in watcher thread, continuing") + + def _watch_stopped_tasks__threading(self): # pylint: disable=R0912,R0914,R0915 + # + # Wait until we have tasks to watch + # + watcher_max_wait = self.node.watcher_max_wait + self.node.have_running_tasks.wait(watcher_max_wait) + # + time.sleep(self.node.thread_scan_interval) # Additional wait interval to avoid busy looping + # + if self.node.result_transport == "events": + # + # Check stopped tasks with missing results + # + with self.node.lock: + done_tasks = [] + # + for task_id, data in self.node.running_tasks.items(): + if data["thread"] is not None: + continue + # + if data["result"] is not None: + done_task = (task_id, data["result"]) + done_tasks.append(done_task) + continue + # + age = (datetime.datetime.now() - data["timestamp"]).total_seconds() + if age < self.node.result_max_wait: + continue + # + done_task = (task_id, data["result"]) + done_tasks.append(done_task) + # + # Announce late and expired + # + for task_id, result in done_tasks: + self._announce_task_stopped(task_id, result) + # + # Collect task IDs of stopped tasks + # + with self.node.lock: + stopped_tasks = [] + # + for task_id, data in self.node.running_tasks.items(): + if data["thread"] is None: + continue # task is stopped, awaiting result + # + if data["thread"].is_alive(): + continue + # + stopped_tasks.append(task_id) + # + # Process newly stopped tasks + # + for task_id in stopped_tasks: + task_data = self.node.running_tasks.get(task_id, None) + # + if task_data is None: + continue + # + try: + task_data["thread"].join(1) + except: # pylint: disable=W0702 + log.exception("Failed to join thread, continuing") + finally: + task_data["thread"] = None + # + task_payload = self.node.local_tasks.get(task_id, None) + # + if task_payload is None: + task_payload = { + "durable": False, + } + # + if task_payload.get("durable", False): + # Task is durable, restart it + self.node.execute_local_task( + task_id, + name=task_payload.get("name"), + meta=task_payload.get("meta"), + args=task_payload.get("args"), + kwargs=task_payload.get("kwargs"), + durable=task_payload.get("durable"), + pool=task_payload.get("pool"), + ) + # + continue + # + if self.node.result_transport == "files": + try: + result_path = os.path.join(self.node.tmp_path, f'{task_id}.bin') + with open(result_path, "rb") as file: + task_data["result"] = file.read() + os.remove(result_path) + except: # pylint: disable=W0702 + log.exception("Failed to load/remove result, continuing") + # + elif self.node.result_transport == "events" and task_data["result"] is None: + # Result event is not processed (or process crashed badly) + task_data["timestamp"] = datetime.datetime.now() + continue + elif self.node.result_transport == "memory": + try: + result = task_data["result"].get(timeout=self.node.result_max_wait) + except: # pylint: disable=W0702 + result = None + # + task_data["result"] = result + # + self._announce_task_stopped(task_id, task_data["result"]) + + def _watch_stopped_tasks__multiprocessing(self): # pylint: disable=R0912,R0914,R0915 + # + # Wait until we have tasks to watch + # + watcher_max_wait = self.node.watcher_max_wait + self.node.have_running_tasks.wait(watcher_max_wait) + # + if self.node.result_transport == "events": + # + # Check stopped tasks with missing results + # + with self.node.lock: + done_tasks = [] + # + for task_id, data in self.node.running_tasks.items(): + if data["process"] is not None: + continue + # + if data["result"] is not None: + done_task = (task_id, data["result"]) + done_tasks.append(done_task) + continue + # + age = (datetime.datetime.now() - data["timestamp"]).total_seconds() + if age < self.node.result_max_wait: + continue + # + done_task = (task_id, data["result"]) + done_tasks.append(done_task) + # + # Announce late and expired + # + for task_id, result in done_tasks: + self._announce_task_stopped(task_id, result) + # + # Collect sentinels of running tasks + # + with self.node.lock: + sentinel_map = {} + # + for task_id, data in self.node.running_tasks.items(): + if data["process"] is None: + continue # task is stopped, awaiting result + sentinel_map[data["process"].sentinel] = task_id + # + # Wait for tasks to stop + # + ready_sentinels = multiprocessing.connection.wait( + list(sentinel_map), timeout=watcher_max_wait + ) + # + # Process newly stopped tasks + # + for sentinel in ready_sentinels: + task_id = sentinel_map[sentinel] + task_data = self.node.running_tasks.get(task_id, None) + # + if task_data is None: + continue + # + process_pid = task_data["process"].pid + if process_pid is not None: + try: + import pylon # pylint: disable=C0415,E0401,W0611 + from tools import context # pylint: disable=C0415,E0401 + # + context.zombie_reaper.external_pids.discard(process_pid) + except: # pylint: disable=W0702 + pass + # + try: + task_data["process"].join(1) + task_data["process"].close() + except: # pylint: disable=W0702 + log.exception("Failed to close process, continuing") + finally: + task_data["process"] = None + # + task_payload = self.node.local_tasks.get(task_id, None) + # + if task_payload is None: + task_payload = { + "durable": False, + } + # + if task_payload.get("durable", False): + # Task is durable, restart it + self.node.execute_local_task( + task_id, + name=task_payload.get("name"), + meta=task_payload.get("meta"), + args=task_payload.get("args"), + kwargs=task_payload.get("kwargs"), + durable=task_payload.get("durable"), + pool=task_payload.get("pool"), + ) + # + continue + # + if self.node.result_transport == "files": + try: + result_path = os.path.join(self.node.tmp_path, f'{task_id}.bin') + with open(result_path, "rb") as file: + task_data["result"] = file.read() + os.remove(result_path) + except: # pylint: disable=W0702 + log.exception("Failed to load/remove result, continuing") + # + elif self.node.result_transport == "events" and task_data["result"] is None: + # Result event is not processed (or process crashed badly) + task_data["timestamp"] = datetime.datetime.now() + continue + elif self.node.result_transport == "memory": + try: + result = task_data["result"].get(timeout=self.node.result_max_wait) + except: # pylint: disable=W0702 + result = None + # + try: + task_data["result"].close() + except: # pylint: disable=W0702 + log.exception("Failed to close result, continuing") + # + task_data["result"] = result + # + self._announce_task_stopped(task_id, task_data["result"]) + + def _announce_task_stopped(self, task_id, result): + task_state = self.node.global_task_state[task_id].copy() + # + task_state["status"] = "stopped" + task_state["result"] = result + # + with self.node.lock: + self.node.local_tasks.pop(task_id, None) + self.node.running_tasks.pop(task_id, None) + if not self.node.running_tasks: + self.node.have_running_tasks.clear() + # + self.node.event_node.emit( + "task_node_announce", + { + "ident": self.node.ident, + "pool": self.node.pool, + "task_limit": self.node.task_limit, + "running_tasks": len(self.node.running_tasks), + } + ) + # + self.node.event_node.emit( + "task_state_announce", + task_state + ) + # + self.node.event_node.emit( + "task_status_change", + { + "task_id": task_id, + "status": "stopped", + } + ) diff --git a/docker-compose.yml b/docker-compose.yml deleted file mode 100644 index 7cc9a16..0000000 --- a/docker-compose.yml +++ /dev/null @@ -1,59 +0,0 @@ -version: "3" - - -services: - redis: - image: redis:alpine - command: redis-server --appendonly yes --requirepass redis - networks: - - pylon - - postgres: - image: postgres:15.1 - logging: - driver: "json-file" - options: - max-file: "5" - max-size: "10m" - restart: unless-stopped - volumes: - - postgres-data:/var/lib/postgresql/data - environment: - - POSTGRES_DB=pylon - - POSTGRES_USER=pylon - - POSTGRES_PASSWORD=pylon - - POSTGRES_INITDB_ARGS=--data-checksums - networks: - - pylon - - pylon: - image: getcarrier/pylon:tasknode - logging: - driver: "json-file" - options: - max-file: "5" - max-size: "10m" - restart: unless-stopped - environment: - - PYLON_CONFIG_SEED=file:/data/pylon.yml - - REDIS_PASSWORD=redis - - POSTGRES_DB=pylon - - POSTGRES_USER=pylon - - POSTGRES_PASSWORD=pylon - volumes: - - ./pylon:/data - depends_on: - - redis - - postgres - networks: - - pylon - ports: - - 8080:8080 - - -volumes: - postgres-data: - - -networks: - pylon: diff --git a/pylon/configs/.gitkeep b/pylon/configs/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/pylon/plugins/shared/metadata.json b/pylon/plugins/shared/metadata.json deleted file mode 100644 index 7c8a97d..0000000 --- a/pylon/plugins/shared/metadata.json +++ /dev/null @@ -1,6 +0,0 @@ -{ - "name": "Shared plugin", - "version": "0.1", - "depends_on": [], - "init_after": [] -} diff --git a/pylon/plugins/shared/module.py b/pylon/plugins/shared/module.py deleted file mode 100644 index 732fbc4..0000000 --- a/pylon/plugins/shared/module.py +++ /dev/null @@ -1,59 +0,0 @@ -#!/usr/bin/python3 -# coding=utf-8 - -# Copyright 2024 getcarrier.io -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" Module """ - -from pylon.core.tools import log # pylint: disable=E0611,E0401 -from pylon.core.tools import module # pylint: disable=E0611,E0401 - - -class Module(module.ModuleModel): - """ Pylon module """ - - def __init__(self, context, descriptor): - self.context = context - self.descriptor = descriptor - - def init(self): - """ Init module """ - log.info("Initializing module") - # Init - self.descriptor.init_all( - url_prefix="/", - static_url_prefix="/", - ) - - from .tools.config import Config - _config = Config(self) - self.descriptor.register_tool('constants', _config) - self.descriptor.register_tool('config', _config) - - from .tools import db - self.descriptor.register_tool('db', db) - - from .tools import db_tools - self.descriptor.register_tool('db_tools', db_tools) - - @self.context.app.teardown_appcontext - def shutdown_session(exception=None): - db.session.remove() - - def deinit(self): - """ De-init module """ - log.info("De-initializing module") - # De-init - self.descriptor.deinit_all() diff --git a/pylon/plugins/shared/requirements.txt b/pylon/plugins/shared/requirements.txt deleted file mode 100644 index de56622..0000000 --- a/pylon/plugins/shared/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -Flask-SQLAlchemy==3.0.2 -pydantic==1.10.11 -alembic==1.11.1 diff --git a/pylon/plugins/shared/tools/config.py b/pylon/plugins/shared/tools/config.py deleted file mode 100644 index 407b725..0000000 --- a/pylon/plugins/shared/tools/config.py +++ /dev/null @@ -1,109 +0,0 @@ -# pylint: disable=E1101,E0203,C0103 -# -# Copyright 2023 getcarrier.io -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" Config """ - -import json -from pylon.core.tools import log # pylint: disable=E0401 -from .patterns import SingletonABC - - -class Config(metaclass=SingletonABC): # pylint: disable=R0903 - """ Config singleton """ - - def __init__(self, module): - module_cfg = module.descriptor.config - self.load_settings( - module_cfg.get("settings", {}), - ( - ("DATABASE_VENDOR", "str", "postgres"), - ("POSTGRES_HOST", "str", "postgres"), - ("POSTGRES_PORT", "int", 5432), - ("POSTGRES_USER", "str", ""), - ("POSTGRES_PASSWORD", "str", ""), - ("POSTGRES_DB", "str", ""), - ("POSTGRES_SCHEMA", "str", "pylon"), - ("POSTGRES_TENANT_SCHEMA", "str", "tenant"), - ("DATABASE_URI", "str", None), - ("DATABASE_ENGINE_OPTIONS", "dict", None), - ) - ) - # - # Make DB URI if not set - # - if self.DATABASE_ENGINE_OPTIONS is None: - self.DATABASE_ENGINE_OPTIONS = {} - # - if self.DATABASE_URI is None: - if self.DATABASE_VENDOR == "sqlite": # Probably is not supported with tenant schemas now # pylint: disable=C0301 - self.DATABASE_URI = f"sqlite:///{self.SQLITE_DB}" - self.DATABASE_ENGINE_OPTIONS["isolation_level"] = "SERIALIZABLE" - elif self.DATABASE_VENDOR == "postgres": - self.DATABASE_URI = 'postgresql://{username}:{password}@{host}:{port}/{database}'.format( # pylint: disable=C0301 - host=self.POSTGRES_HOST, - port=self.POSTGRES_PORT, - username=self.POSTGRES_USER, - password=self.POSTGRES_PASSWORD, - database=self.POSTGRES_DB - ) - if not self.DATABASE_ENGINE_OPTIONS: - self.DATABASE_ENGINE_OPTIONS = { - "isolation_level": "READ COMMITTED", - "echo": False, - "pool_size": 50, - "max_overflow": 100, - "pool_pre_ping": True - } - else: - raise RuntimeError(f"Unsupported DB vendor: {self.DATABASE_VENDOR}") - # - log.info('Initialized config %s', self) - - def load_settings(self, settings, schema): - """ Load and set config vars """ - processors = { - "str": lambda item: item if isinstance(item, str) else str(item), - "int": lambda item: item if isinstance(item, int) else int(item), - "bool": lambda item: item if isinstance(item, bool) else item.lower() in ["true", "yes"], # pylint: disable=C0301 - "dict": lambda item: item if isinstance(item, dict) else json.loads(item), - } - # - for item in schema: - if len(item) == 3: - key, kind, default = item - elif len(item) == 2: - key, kind = item - default = ... - else: - raise RuntimeError(f"Invalid config schema: {item}") - # - if isinstance(default, set): - default = getattr(self, list(default)[0]) - # - data = ... - for variant in [key, key.lower(), key.upper()]: - if variant in settings: - data = settings[variant] - # - if data is ... and default is ...: - raise RuntimeError(f"Required config value is not set: {key}") - # - if data is ...: - data = default - elif kind in processors: - data = processors[kind](data) - # - setattr(self, key, data) diff --git a/pylon/plugins/shared/tools/db.py b/pylon/plugins/shared/tools/db.py deleted file mode 100644 index a3f9107..0000000 --- a/pylon/plugins/shared/tools/db.py +++ /dev/null @@ -1,77 +0,0 @@ -import time - -from contextlib import contextmanager -from flask_sqlalchemy import BaseQuery -from sqlalchemy import create_engine, MetaData -from sqlalchemy.schema import CreateSchema -from sqlalchemy.orm import sessionmaker, scoped_session, declarative_base - -from tools import config as c - -engine = create_engine(c.DATABASE_URI, **c.DATABASE_ENGINE_OPTIONS) -while True: - try: - with engine.connect() as conn: - conn.execute(CreateSchema(c.POSTGRES_SCHEMA, if_not_exists=True)) - conn.commit() - break - except: - time.sleep(1) -# -session = scoped_session(sessionmaker(bind=engine)) -Base = declarative_base() -Base.query = session.query_property(query_cls=BaseQuery) - - -def get_all_metadata(): - meta = MetaData() - for table in Base.metadata.tables.values(): - table.tometadata(meta) - return meta - - -def get_shared_metadata(): - meta = MetaData() - for table in Base.metadata.tables.values(): - if table.schema != c.POSTGRES_TENANT_SCHEMA: - table.tometadata(meta) - return meta - - -def get_tenant_specific_metadata(): - meta = MetaData(schema=c.POSTGRES_TENANT_SCHEMA) - for table in Base.metadata.tables.values(): - if table.schema == c.POSTGRES_TENANT_SCHEMA: - table.tometadata(meta) - return meta - - -def get_schema_translate_map(project_id: int | None) -> dict | None: - if project_id: - from tools import project_constants as pc - template = pc['PROJECT_SCHEMA_TEMPLATE'] - return { - c.POSTGRES_TENANT_SCHEMA: template.format(project_id), - # c.POSTGRES_SCHEMA: c.POSTGRES_SCHEMA - } - return None - - -def get_project_schema_session(project_id: int | None): - schema_translate_map = get_schema_translate_map(project_id) - connectable = engine.execution_options(schema_translate_map=schema_translate_map) - return scoped_session(sessionmaker(bind=connectable)) - - -@contextmanager -def with_project_schema_session(project_id: int | None): - # schema_translate_map = get_schema_translate_map(project_id) - # connectable = engine.execution_options(schema_translate_map=schema_translate_map) - db = None - try: - db = get_project_schema_session(project_id) - # db = scoped_session(sessionmaker(bind=connectable)) - yield db - finally: - if db: - db.close() diff --git a/pylon/plugins/shared/tools/db_tools.py b/pylon/plugins/shared/tools/db_tools.py deleted file mode 100644 index 9033598..0000000 --- a/pylon/plugins/shared/tools/db_tools.py +++ /dev/null @@ -1,88 +0,0 @@ -#!/usr/bin/python3 -# coding=utf-8 - -# Copyright 2022 getcarrier.io -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" DB tools """ -import json - -from datetime import datetime -from typing import Optional - -from pylon.core.tools import log - -from .db import session -from tools import config as c - - -def sqlalchemy_mapping_to_dict(obj): - """ Make dict from sqlalchemy mappings().one() object """ - return {str(key): value for key, value in dict(obj).items()} - - -class AbstractBaseMixin: - _session = session - __table__ = None - __table_args__ = {"schema": c.POSTGRES_SCHEMA} - - def __repr__(self) -> str: - return json.dumps(self.to_json(), indent=2) - - def to_json(self, exclude_fields: tuple = ()) -> dict: - log.debug('Be cautious "to_json()". Better write your own serialization for %s', getattr(self, '__tablename__')) - result = dict() - for column in self.__table__.columns: - if column.name not in set(exclude_fields): - value = getattr(self, column.name) - if isinstance(value, datetime): - value = value.isoformat() - result[column.name] = value - return result - - @staticmethod - def commit() -> None: - try: - session.commit() - except: # pylint: disable=W0702 - self.rollback() - raise - - def add(self, with_session: Optional = None) -> None: - session.add(self) - - def insert(self, with_session: Optional = None) -> None: - self.add() - self.commit() - - def delete(self, commit: bool = True, with_session: Optional = None) -> None: - session.delete(self) - if commit: - self.commit() - - def rollback(self, with_session: Optional = None): - session.rollback() - - @property - def serialized(self): - raise NotImplementedError - - -def bulk_save(objects): - session.bulk_save_objects(objects) - try: - session.commit() - except: # pylint: disable=W0702 - session.rollback() - raise diff --git a/pylon/plugins/shared/tools/patterns.py b/pylon/plugins/shared/tools/patterns.py deleted file mode 100644 index 6b209db..0000000 --- a/pylon/plugins/shared/tools/patterns.py +++ /dev/null @@ -1,58 +0,0 @@ -# Copyright 2020 getcarrier.io -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import inspect -from abc import ABCMeta -from threading import Lock -from typing import Optional, Dict, Tuple, Callable - - -class SingletonMeta(type): - _instance: Optional["SingletonABC"] = None - _lock: Lock = Lock() - - def __call__(cls, *args, **kwargs): - with cls._lock: - if not cls._instance: - cls._instance = super().__call__(*args, **kwargs) - return cls._instance - - -class SingletonABC(SingletonMeta, ABCMeta): - ... - - -class SingletonParametrizedMeta(type): - _instances: Dict[Tuple[str, frozenset], "SingletonParametrizedABC"] = {} - _init: Dict[str, Callable] = {} - _lock: Lock = Lock() - - def __init__(cls, name, bases, attrs): - cls._init[cls.__name__] = attrs.get('__init__', None) - super().__init__(name, bases, attrs) - - def __call__(cls, *args, **kwargs): - init = cls._init[cls.__name__] - if init is not None: - key = (cls.__name__, - frozenset(inspect.getcallargs(init, None, *args, **kwargs).items())) - else: - key = cls.__name__ - - if key not in cls._instances: - cls._instances[key] = super().__call__(*args, **kwargs) - return cls._instances[key] - - -class SingletonParametrizedABC(SingletonParametrizedMeta, ABCMeta): - ... diff --git a/pylon/plugins/test_api/api/__init__.py b/pylon/plugins/test_api/api/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/pylon/plugins/test_api/api/v1/__init__.py b/pylon/plugins/test_api/api/v1/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/pylon/plugins/test_api/api/v1/metadata.py b/pylon/plugins/test_api/api/v1/metadata.py deleted file mode 100644 index 1d92478..0000000 --- a/pylon/plugins/test_api/api/v1/metadata.py +++ /dev/null @@ -1,24 +0,0 @@ -from flask_restful import Resource - -from pylon.core.tools import log - -# from ...models.metadata import MetadataEntry - - -class API(Resource): - """ API implementation """ - - url_params = [ - '', - '', - ] - - def __init__(self, module): - self.module = module - - def get(self, key=None): - """ List all metadata keys or get metadata vales for specific key """ - if key is None: - return [] # TODO: list present keys - # TODO: get data for key - return {"error": "not implemented yet"}, 418 diff --git a/pylon/plugins/test_api/init_db.py b/pylon/plugins/test_api/init_db.py deleted file mode 100644 index 57f8259..0000000 --- a/pylon/plugins/test_api/init_db.py +++ /dev/null @@ -1,6 +0,0 @@ -from tools import db - - -def init_db(): - from .models import metadata - db.get_shared_metadata().create_all(bind=db.engine) diff --git a/pylon/plugins/test_api/metadata.json b/pylon/plugins/test_api/metadata.json deleted file mode 100644 index 4f90110..0000000 --- a/pylon/plugins/test_api/metadata.json +++ /dev/null @@ -1,6 +0,0 @@ -{ - "name": "Testing plugin: API", - "version": "0.1", - "depends_on": ["shared"], - "init_after": [] -} diff --git a/pylon/plugins/test_api/models/__init__.py b/pylon/plugins/test_api/models/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/pylon/plugins/test_api/models/metadata.py b/pylon/plugins/test_api/models/metadata.py deleted file mode 100644 index cbfeeb1..0000000 --- a/pylon/plugins/test_api/models/metadata.py +++ /dev/null @@ -1,31 +0,0 @@ -#!/usr/bin/python -# coding=utf-8 - -# Copyright 2024 getcarrier.io -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" Secrets DB model """ - -# from sqlalchemy import Something # pylint: disable=E0401 - -from tools import db, db_tools - - -# class MetadataEntry(db_tools.AbstractBaseMixin, db.Base): # pylint: disable=C0111 -# -# # TODO: add something here -# -# @property -# def serialized(self): -# raise RuntimeError("Not supported") diff --git a/pylon/plugins/test_api/models/pd/__init__.py b/pylon/plugins/test_api/models/pd/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/pylon/plugins/test_api/models/pd/metadata.py b/pylon/plugins/test_api/models/pd/metadata.py deleted file mode 100644 index 6cdee01..0000000 --- a/pylon/plugins/test_api/models/pd/metadata.py +++ /dev/null @@ -1,5 +0,0 @@ -from pydantic import BaseModel - - -# class MetadataValidatorModel(BaseModel): -# # FIXME: can add something here for input schema validation diff --git a/pylon/plugins/test_api/module.py b/pylon/plugins/test_api/module.py deleted file mode 100644 index 2e2c09c..0000000 --- a/pylon/plugins/test_api/module.py +++ /dev/null @@ -1,44 +0,0 @@ -#!/usr/bin/python3 -# coding=utf-8 - -# Copyright 2024 getcarrier.io -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" Module """ - -from pylon.core.tools import log # pylint: disable=E0611,E0401 -from pylon.core.tools import module # pylint: disable=E0611,E0401 - - -class Module(module.ModuleModel): - """ Pylon module """ - - def __init__(self, context, descriptor): - self.context = context - self.descriptor = descriptor - - def init(self): - """ Init module """ - log.info("Initializing module") - # Init - self.descriptor.init_all() - # DB - from .init_db import init_db - init_db() - - def deinit(self): - """ De-init module """ - log.info("De-initializing module") - # De-init - self.descriptor.deinit_all() diff --git a/pylon/pylon.yml b/pylon/pylon.yml deleted file mode 100644 index 88f9685..0000000 --- a/pylon/pylon.yml +++ /dev/null @@ -1,64 +0,0 @@ -server: - path: / - host: "0.0.0.0" - port: 8080 - -modules: - plugins: - provider: - type: folder - path: /data/plugins - # - requirements: - mode: relaxed - activation: bulk - provider: - type: folder - path: /data/requirements - # - config: - provider: - type: folder - path: /data/configs - -configs: - shared: - settings: - postgres_db: pylon - postgres_user: pylon - postgres_password: pylon - -sessions: - redis: - host: redis - password: $REDIS_PASSWORD - prefix: pylon_session_ - -events: - redis: - host: redis - password: $REDIS_PASSWORD - queue: events - hmac_key: events_hmac_key - hmac_digest: sha512 - callback_workers: 16 - -rpc: - redis: - host: redis - password: $REDIS_PASSWORD - queue: rpc - hmac_key: rpc_hmac_key - hmac_digest: sha512 - callback_workers: 16 - id_prefix: "pylon_" - -socketio: - redis: - host: redis - password: $REDIS_PASSWORD - queue: socketio - -application: - SECRET_KEY: pylon_key - SESSION_COOKIE_NAME: pylon_session_id diff --git a/pylon/requirements/.gitkeep b/pylon/requirements/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..b46a3d4 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,18 @@ +requests>=2.31.0 +PyYAML>=6.0 +pika>=1.1.0 +pytest>=6.2.4 +coverage>=5.5 +redis>=4.5.5 + +# SocketIO + deps +python-engineio[client]>=4.11.2 +python-socketio[client]>=5.12.1 +simple-websocket>=1.1.0 +# +bidict>=0.23.1 +wsproto>=1.2.0 +h11>=0.14.0 + +# TaskNode +setproctitle>=1.3.3 diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..23c3103 --- /dev/null +++ b/setup.py @@ -0,0 +1,49 @@ +#!/usr/bin/python3 +# coding=utf-8 +# pylint: disable=C0103,C0413 + +# Copyright (c) 2020 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" + Setup script +""" + +with open("requirements.txt", "r", encoding="utf-8") as f: + required = f.read().splitlines() + +with open("version.txt", "r", encoding="utf-8") as f: + version = f.read().splitlines()[0].strip() + +try: + import subprocess + tag = subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]) + version = f"{version}+git.{tag.decode('utf-8').strip()}" +except: # pylint: disable=W0702 + pass + +from setuptools import setup, find_packages + +setup( + name="arbiter", + version=version, + description="Distributed queues, RPCs, events, tasks", + long_description="Lightweight distributed task management framework", + url="https://getcarrier.io", + license="Apache License 2.0", + author="arozumenko, LifeDJIK", + author_email="artem_rozumenko@epam.com, ivan_krakhmaliuk@epam.com", + packages=find_packages(), + install_requires=required +) diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..6130f85 --- /dev/null +++ b/tests/README.md @@ -0,0 +1,22 @@ +## Testing app for arbiter + +Launch redis container +``` +docker run -d --rm --hostname arbiter-redis --name arbiter-redis \ + -p 6379:6379 redis:alpine redis-server +``` + +Launch minion app with `python minion.py` + +## Running tests + +to run tests you need to execution `pytest -q tests/` + +## Test coverage + +to check test coverage you need to run +``` +coverage run --source=arbiter -m pytest -q tests/ +coverage report -m +coverage html +``` diff --git a/pylon/plugins/shared/tools/__init__.py b/tests/__init__.py similarity index 100% rename from pylon/plugins/shared/tools/__init__.py rename to tests/__init__.py diff --git a/tests/minion.py b/tests/minion.py new file mode 100644 index 0000000..c71aee1 --- /dev/null +++ b/tests/minion.py @@ -0,0 +1,45 @@ +from arbiter import Minion +from time import sleep +import logging + + +def start_minion(event_node): + app = Minion(event_node, queue="default") + # + app.raw_task_node.multiprocessing_context = "threading" + app.raw_task_node.result_transport = "memory" + # + @app.task(name="add") + def add(x, y): + logging.info("Running task 'add'") + # task that initiate new task within same app + increment = 0 + for message in app.apply('simple_add', task_args=[3, 4]): + if isinstance(message, dict): + increment = message["result"] + logging.info("sleep done") + return x + y + increment + # + @app.task(name="simple_add") + def adds(x, y): + logging.info(f"Running task 'add_small' with params {x}, {y}") + return x + y + # + @app.task(name="add_in_pipe") + def addp(x, y, upstream=0): + logging.info("Running task 'add_in_pipe'") + return x + y + upstream + # + @app.task(name="long_running") + def long_task(): + for _ in range(180): + sleep(1) + return "Long Task" + # + app.run(workers=10, block=False) + # + return app + + +def stop_minion(app): + app.raw_task_node.stop() diff --git a/tests/test_arbiter.py b/tests/test_arbiter.py new file mode 100644 index 0000000..8d6d505 --- /dev/null +++ b/tests/test_arbiter.py @@ -0,0 +1,178 @@ +#!/usr/bin/python3 +# coding=utf-8 +# pylint: disable=C0114,C0115,C0116,C0411,C0103 + +# Copyright 2023 getcarrier.io +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import pytest # pylint: disable=E0401,W0611 + +from time import sleep, time +from arbiter import Arbiter, Task, MockEventNode +from tests.minion import stop_minion, start_minion + +arbiter_queue = "default" + + +class TestArbiter: + event_node = None + minion = None + + @classmethod + def setup_class(cls): + cls.event_node = MockEventNode() + cls.event_node.start() + cls.minion = start_minion(cls.event_node) + + @classmethod + def teardown_class(cls): + stop_minion(cls.minion) + cls.event_node.stop() + + @staticmethod + def test_task_in_task(): + tasks_in_task_in_task = 5 + arbiter = Arbiter(event_node=TestArbiter.event_node) + assert arbiter.workers()[arbiter_queue]['total'] == 10 + task_keys = [] + for _ in range(tasks_in_task_in_task): + task_keys.append(arbiter.apply("simple_add", task_args=[1, 2])[0]) + for task_key in task_keys: + assert arbiter.status(task_key)['state'] in ('initiated', 'running', 'done') + for message in arbiter.wait_for_tasks(task_keys): + assert message['state'] == 'done' + assert message['result'] == 3 + for task_key in task_keys: + assert arbiter.status(task_key)['state'] == 'done' + assert arbiter.status(task_key)['result'] == 3 + assert arbiter.workers()[arbiter_queue]['available'] == 10 + arbiter.close() + + @staticmethod + def test_squad(): + tasks_in_squad = 3 + arbiter = Arbiter(event_node=TestArbiter.event_node) + tasks = [] + for _ in range(tasks_in_squad): + tasks.append(Task("simple_add", task_args=[1, 2])) + squad_id = arbiter.squad(tasks) + while arbiter.status(squad_id).get("state") != "done": + sleep(1) + status = arbiter.status(squad_id) + assert status["done"] == tasks_in_squad + assert len(status["tasks"]) == tasks_in_squad + assert arbiter.workers()[arbiter_queue]['available'] == 10 + arbiter.close() + + @staticmethod + def test_pipe(): + tasks_in_pipe = 20 + arbiter = Arbiter(event_node=TestArbiter.event_node) + tasks = [] + for _ in range(tasks_in_pipe): + tasks.append(Task("add_in_pipe", task_args=[2])) + pipe_id = None + _loop_result = 0 + _loop_id = 1 + for message in arbiter.pipe(tasks, persistent_args=[2]): + if "pipe_id" in message: + pipe_id = message["pipe_id"] + else: + _loop_result = message['result'] + assert _loop_result == 4 * _loop_id + _loop_id += 1 + status = arbiter.status(pipe_id) + assert status["done"] == tasks_in_pipe + assert len(status["tasks"]) == tasks_in_pipe + assert arbiter.workers()[arbiter_queue]['available'] == 10 + arbiter.close() + + @staticmethod + def test_kill_task(): + arbiter = Arbiter(event_node=TestArbiter.event_node) + start = time() + tasks = arbiter.apply("long_running") + for task_key in tasks: + assert arbiter.status(task_key)['state'] in ['initiated', 'running'] + sleep(2) # time for task to settle + arbiter.kill(tasks[0], sync=True) + for message in arbiter.wait_for_tasks(tasks): + assert message['state'] == 'done' + assert time()-start < 180 # 180 sec is a length of task + assert arbiter.workers()[arbiter_queue]['available'] == 10 + arbiter.close() + + @staticmethod + def test_kill_group(): + tasks_in_squad = 3 + start = time() + arbiter = Arbiter(event_node=TestArbiter.event_node) + tasks = [] + for _ in range(tasks_in_squad): + tasks.append(Task("long_running")) + squad_id = arbiter.squad(tasks) + sleep(5) # time for squad to settle + arbiter.kill_group(squad_id) + while arbiter.status(squad_id).get("state") != "done": + sleep(1) + assert time() - start < 180 + assert arbiter.workers()[arbiter_queue]['available'] == 10 + arbiter.close() + + @staticmethod + def test_squad_callback(): + tasks_in_squad = 3 + arbiter = Arbiter(event_node=TestArbiter.event_node) + tasks = [] + for _ in range(tasks_in_squad): + tasks.append(Task("simple_add", task_args=[1, 2])) + squad_id = arbiter.squad(tasks, callback=Task("simple_add", task_args=[5, 4])) + while arbiter.status(squad_id).get("state") != "done": + sleep(1) + status = arbiter.status(squad_id) + assert status["done"] == tasks_in_squad + 1 + assert len(status["tasks"]) == tasks_in_squad + 1 + assert status["tasks"][-1]['task_type'] == "callback" + assert status["tasks"][-1]['result'] == 9 + assert arbiter.workers()[arbiter_queue]['available'] == 10 + arbiter.close() + + @staticmethod + def test_squad_finalyzer(): + tasks_in_squad = 3 + arbiter = Arbiter(event_node=TestArbiter.event_node) + tasks = [] + for _ in range(tasks_in_squad): + tasks.append(Task("simple_add", task_args=[1, 2])) + tasks.append(Task("simple_add", task_args=[5, 5], task_type='finalize')) + squad_id = arbiter.squad(tasks, callback=Task("simple_add", task_args=[5, 4])) + while arbiter.status(squad_id).get("state") != "done": + sleep(1) + status = arbiter.status(squad_id) + assert status["done"] == tasks_in_squad + 2 # callback + finalizer + assert len(status["tasks"]) == tasks_in_squad + 2 # callback + finalizer + assert status["tasks"][-1]['task_type'] == "finalize" + assert status["tasks"][-1]['result'] == 10 + assert arbiter.workers()[arbiter_queue]['available'] == 10 + arbiter.close() + + @staticmethod + def test_sync_task(): + arbiter = Arbiter(event_node=TestArbiter.event_node) + for result in arbiter.add_task(Task("simple_add", task_args=[1, 2]), sync=True): + if isinstance(result, dict): + assert result['state'] == 'done' + assert result['result'] == 3 + assert arbiter.workers()[arbiter_queue]['available'] == 10 diff --git a/version.txt b/version.txt new file mode 100644 index 0000000..23aa839 --- /dev/null +++ b/version.txt @@ -0,0 +1 @@ +1.2.2 From b6d0678b5210bcd4dfe861228b811f1228993647 Mon Sep 17 00:00:00 2001 From: Ivan Krakhmaliuk Date: Thu, 26 Jun 2025 16:36:20 +0300 Subject: [PATCH 2/4] add simple agents.md --- AGENTS.md | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 AGENTS.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..300775d --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,2 @@ +# VirtualEnv +Use venv from .venv From 3f2c7949ba268f3b9680d84073bd547fe8df6323 Mon Sep 17 00:00:00 2001 From: Ivan Krakhmaliuk Date: Fri, 27 Jun 2025 19:53:47 +0300 Subject: [PATCH 3/4] Implement TaskNode unit tests - Create pytest.ini --- pytest.ini | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 pytest.ini diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..9038ff2 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +# Only run TaskNode unit tests to avoid slow/minion tests +python_files = test_tasknode.py \ No newline at end of file From ace814262932e8b0add85fad0edf175b036584bf Mon Sep 17 00:00:00 2001 From: Ivan Krakhmaliuk Date: Fri, 27 Jun 2025 19:53:49 +0300 Subject: [PATCH 4/4] Implement TaskNode unit tests - Create tests/test_tasknode.py --- tests/test_tasknode.py | 145 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 145 insertions(+) create mode 100644 tests/test_tasknode.py diff --git a/tests/test_tasknode.py b/tests/test_tasknode.py new file mode 100644 index 0000000..93e127e --- /dev/null +++ b/tests/test_tasknode.py @@ -0,0 +1,145 @@ +#!/usr/bin/env pytest +""" +Unit tests for TaskNode: initialization, execution, error handling, +dependency management, cancellation, and state transitions. +""" +import time +import threading +import pytest + +from arbiter.tasknode.tasknode import TaskNode +from arbiter.eventnode.mock import MockEventNode + + +@pytest.fixture +def tasknode(): + # Create a TaskNode with fast intervals for testing + event_node = MockEventNode() + # configure short waits for watcher and housekeeping + node = TaskNode( + event_node=event_node, + multiprocessing_context="threading", + watcher_max_wait=0.1, + thread_scan_interval=0.05, + result_max_wait=0.1, + housekeeping_interval=0.1, + ) + # start node + node.start() + yield node + # stop node and event_node + node.stop() + if event_node.started: + event_node.stop() + + +def test_initialization(tasknode): + # After start, node should be marked started and have an ident + assert tasknode.started is True + assert isinstance(tasknode.ident, str) and tasknode.ident + # stop_event should not be set + assert not tasknode.stop_event.is_set() + + +def test_register_and_unregister(tasknode): + # define dummy task + def dummy(): + return 'ok' + + # register + tasknode.register_task(dummy, name='dummy') + assert 'dummy' in tasknode.task_registry + # unregister + tasknode.unregister_task(name='dummy') + assert 'dummy' not in tasknode.task_registry + # unregister with func + tasknode.register_task(dummy) + name = tasknode.get_callable_name(dummy) + tasknode.unregister_task(func=dummy) + assert name not in tasknode.task_registry + + +def test_task_success_and_result(tasknode): + # register add + def add(x, y): + return x + y + tasknode.register_task(add, name='add') + # start task + task_id = tasknode.start_task('add', args=[2, 3]) + assert task_id is not None + # wait for completion + tasknode.wait_for_task(task_id, timeout=1) + # get status and result + status = tasknode.get_task_status(task_id) + # Completed tasks report status 'stopped' + assert status == 'stopped' + result = tasknode.get_task_result(task_id) + assert result == 5 + + +def test_task_exception(tasknode): + # register failing task + def fail(): + raise ValueError('oops') + tasknode.register_task(fail, name='fail') + task_id = tasknode.start_task('fail') + assert task_id + tasknode.wait_for_task(task_id, timeout=1) + # status should be stopped (task completed with error) + assert tasknode.get_task_status(task_id) == 'stopped' + # retrieving result should raise Exception wrapping the original + with pytest.raises(Exception) as ei: + _ = tasknode.get_task_result(task_id) + assert 'oops' in str(ei.value) + + +def test_task_chaining(tasknode): + # chain tasks: produce a value then use it + def produce(): + return 7 + def consume(v): + return v * 2 + tasknode.register_task(produce, name='produce') + tasknode.register_task(consume, name='consume') + # first task + t1 = tasknode.start_task('produce') + tasknode.wait_for_task(t1, timeout=1) + v = tasknode.get_task_result(t1) + # second task uses result + t2 = tasknode.start_task('consume', args=[v]) + tasknode.wait_for_task(t2, timeout=1) + assert tasknode.get_task_result(t2) == 14 + + +def test_task_cancellation(tasknode): + # long-running task that sleeps + def long_task(): + time.sleep(2) + return 'done' + tasknode.register_task(long_task, name='long') + t = tasknode.start_task('long') + # give it a moment to start + time.sleep(0.1) + tasknode.stop_task(t) + tasknode.wait_for_task(t, timeout=1) + # cancelled tasks return Ellipsis + res = tasknode.get_task_result(t) + assert res is ... + + +def test_status_subscriptions(tasknode): + events = [] + # subscriber collects status changes (event_name, payload) + def on_status(event_name, data): + events.append((data.get('task_id'), data.get('status'))) + tasknode.subscribe_to_task_statuses(on_status) + # register quick task + def quick(): + return 1 + tasknode.register_task(quick, name='quick') + tid = tasknode.start_task('quick') + tasknode.wait_for_task(tid, timeout=1) + # expect at least pending and stopped statuses + statuses = [s for (_id, s) in events if _id == tid] + assert 'pending' in statuses + assert 'stopped' in statuses \ No newline at end of file