diff --git a/CHANGELOG.md b/CHANGELOG.md index dc321bf..a3e45a3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Add ewoksjob options under `EWOKSJOB_OPTIONS`: ``log_memory_usage`` and ``detect_memory_leaks``. - Added ability to handle remote ewoksutils exception types on the client side. +- Add `client.task_utils.TaskSubmitter` helper to execute a single task. ### Fixed diff --git a/src/ewoksjob/client/task_utils.py b/src/ewoksjob/client/task_utils.py index 17c6d79..9d9cb40 100644 --- a/src/ewoksjob/client/task_utils.py +++ b/src/ewoksjob/client/task_utils.py @@ -1,17 +1,32 @@ import re +from concurrent.futures import Future from copy import deepcopy from typing import Literal from typing import Optional +from uuid import uuid4 from ewoksutils.task_utils import task_inputs -from .celery import submit as _submit_celery -from .celery.futures import CeleryFuture +from . import celery from .futures import FutureInterface -from .local import pool as _local_pool -from .local import submit as _submit_local -from .local.futures import LocalFuture -from .local.pool import _LocalPool + +# ewoks is optional for celery execution; it is required for: synchronous, thread, process and slurm +try: + import ewoks +except ImportError: + ewoks = None +else: + from . import local + from .local.pool import _LocalPool + from .local.pool import active_pool_context + + class _MethodLocalFuture(local.Future): + """Local future that unwraps method task result""" + + def result(self, timeout: Optional[float] = None): + task_result = super().result(timeout) + return task_result["return_value"] + _QUALNAME_RE = re.compile(r"^[A-Za-z_]\w*(\.[A-Za-z_]\w*)*$") @@ -36,7 +51,7 @@ def _guess_task_type(task_identifier: str): return "script" -class _MethodCeleryFuture(CeleryFuture): +class _MethodCeleryFuture(celery.Future): """Celery future that unwraps method task result""" def result(self, timeout: Optional[float] = None, interval: Optional[float] = None): @@ -44,20 +59,14 @@ def result(self, timeout: Optional[float] = None, interval: Optional[float] = No return task_result["return_value"] -class _MethodLocalFuture(LocalFuture): - """Local future that unwraps method task result""" - - def result(self, timeout: Optional[float] = None): - task_result = super().result(timeout) - return task_result["return_value"] - - class TaskSubmitter: def __init__( self, task_identifier: str, task_type: Optional[str] = None, - execution_mode: Literal["celery", "process", "thread", "slurm"] = "celery", + execution_mode: Literal[ + "celery", "process", "thread", "slurm", "synchronous" + ] = "celery", **submit_options, ): """Wrapper function that execute a task asynchronously when called. @@ -76,6 +85,7 @@ def __init__( - "thread": Execute the task in a thread - "process": Execute the task in a different process - "slurm": Execute the task in a session on a SLURM cluster + - "synchronous": Execute the task synchronously when called :param submit_options: Extra arguments are passed to the choosen task queue/executor. Arguments depend on the `execution_mode`: @@ -83,6 +93,11 @@ def __init__( - For "celery", see [celery documentation](https://docs.celeryq.dev/en/stable/reference/celery.app.task.html#celery.app.task.Task.apply_async) - For "slurm", see [pyslumrutils documentation(https://pyslurmutils.readthedocs.io/en/stable/reference/_generated/pyslurmutils.concurrent.rest.SlurmRestExecutor.html#pyslurmutils.concurrent.rest.SlurmRestExecutor) """ + if ewoks is None and execution_mode != "celery": + raise RuntimeError( + f"Please install the 'ewoks' package to use the '{execution_mode}' execution mode." + ) + self._task_identifier = task_identifier if task_type is None: task_type = _guess_task_type(task_identifier) @@ -90,7 +105,7 @@ def __init__( self._execution_mode = execution_mode self._submit_options = deepcopy(submit_options) - if self._execution_mode == "celery": + if self._execution_mode in ("celery", "synchronous"): self._local_executor = None else: self._local_executor = _LocalPool( @@ -106,7 +121,7 @@ def task_type(self) -> str: return self._task_type def shutdown(self, **kwargs): - """Shutdown local executor, do nothing for remote execution""" + """Shutdown local executor, do nothing for remote and synchronous execution""" if self._local_executor is not None: self._local_executor.shutdown(**kwargs) @@ -139,16 +154,26 @@ def __call__(self, *args, **kwargs) -> FutureInterface: } if self._execution_mode == "celery": - future = _submit_celery( - args=(graph,), kwargs=kwargs, **self._submit_options - ) + future = celery.submit(args=(graph,), kwargs=kwargs, **self._submit_options) if self._task_type == "method": return _MethodCeleryFuture(future.uuid) else: return future + elif self._execution_mode == "synchronous": + future = Future() + try: + result = ewoks.execute_graph(graph, **kwargs) + except Exception as e: + future.set_exception(e) + else: + future.set_result(result) + if self._task_type == "method": + return _MethodLocalFuture(str(uuid4()), future) + else: + return local.Future(str(uuid4()), future) else: - with _local_pool.active_pool_context(self._local_executor): - future = _submit_local(args=(graph,), kwargs=kwargs) + with active_pool_context(self._local_executor): + future = local.submit(args=(graph,), kwargs=kwargs) if self._task_type == "method": return _MethodLocalFuture(future.uuid) else: diff --git a/src/ewoksjob/tests/test_task_utils.py b/src/ewoksjob/tests/test_task_utils.py index 1084172..7e93b74 100644 --- a/src/ewoksjob/tests/test_task_utils.py +++ b/src/ewoksjob/tests/test_task_utils.py @@ -32,6 +32,14 @@ def test_task_submitter_process(skip_if_gevent, task, args, kwargs, result): assert submitter(*args, **kwargs).result(timeout=10) == result +@pytest.mark.parametrize("task,args,kwargs,result", TASK_TESTS) +def test_task_submitter_synchronous(task, args, kwargs, result): + submitter = TaskSubmitter(task, execution_mode="synchronous") + future = submitter(*args, **kwargs) + assert future.done() + assert future.result() == result + + def test_function_from_script_submitter_worker(ewoks_worker, tmp_path): script_path = tmp_path / "test.py" script_path.write_text("def function(a, b): return a + b\n") @@ -62,6 +70,18 @@ def test_function_from_script_submitter_process(skip_if_gevent, tmp_path): assert submit_script_function(a=1, b=2).result(timeout=10) == 3 +def test_function_from_script_submitter_synchronous(tmp_path): + script_path = tmp_path / "test.py" + script_path.write_text("def function(a, b): return a + b\n") + + submit_script_function = TaskSubmitter( + f"{script_path}::function", execution_mode="synchronous" + ) + future = submit_script_function(a=1, b=2) + assert future.done() + assert future.result() == 3 + + def test_script_submitter_worker(ewoks_worker, tmp_path): script_path = tmp_path / "test.py" script_path.write_text("import sys; sys.exit(2)") @@ -96,3 +116,13 @@ def test_script_submitter_process(skip_if_gevent, tmp_path): "out": None, "return_code": 2, } + + +def test_script_submitter_synchronous(tmp_path): + script_path = tmp_path / "test.py" + script_path.write_text("import sys; sys.exit(2)") + + submit_script = TaskSubmitter(str(script_path), execution_mode="synchronous") + future = submit_script() + assert future.done() + assert future.result() == {"err": None, "out": None, "return_code": 2}