cluster_tools 4.2.4__tar.gz → 4.3.0__tar.gz

{cluster_tools-4.2.4 → cluster_tools-4.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cluster_tools
-Version: 4.2.4
+Version: 4.3.0
 Summary: Utility library for easily distributing code execution on clusters
 Author-email: scalable minds <hello@scalableminds.com>
 License-Expression: MIT

{cluster_tools-4.2.4 → cluster_tools-4.3.0}/cluster_tools/__init__.py

@@ -7,6 +7,10 @@ from cluster_tools.executors.multiprocessing_ import MultiprocessingExecutor
 from cluster_tools.executors.multiprocessing_pickle import MultiprocessingPickleExecutor
 from cluster_tools.executors.sequential import SequentialExecutor
 from cluster_tools.executors.sequential_pickle import SequentialPickleExecutor
+from cluster_tools.output_store import (
+    FileOutputStore,  # noqa: F401 unused import
+    OutputStore,  # noqa: F401 unused import
+)
 from cluster_tools.schedulers.cluster_executor import (
     ClusterExecutor,  # noqa:  F401 `cluster_tools.schedulers.cluster_executor.ClusterExecutor` imported but unused;
     RemoteOutOfMemoryException,  # noqa:  F401 `cluster_tools.schedulers.cluster_executor.ClusterExecutor` imported but unused;
@@ -108,7 +112,9 @@ def get_executor(environment: str, **kwargs: Any) -> "Executor":
         return KubernetesExecutor(**kwargs)
     elif environment == "dask":
         if "client" in kwargs:
-            return DaskExecutor(kwargs["client"])
+            return DaskExecutor(
+                kwargs["client"], output_store=kwargs.get("output_store")
+            )
         else:
             return DaskExecutor.from_config(**kwargs)
     elif environment == "multiprocessing":

cluster_tools-4.3.0/cluster_tools/_utils/cfut_options.py

@@ -0,0 +1,108 @@
+"""Handling of the `__cfut_options` kwarg and the `map_to_futures` getters, which
+select where a job's output is persisted. Shared by all executors."""
+
+import logging
+import os
+import warnings
+from collections.abc import Callable
+from typing import Any, TypedDict, TypeVar, cast
+
+from typing_extensions import NotRequired, ParamSpec
+
+from cluster_tools._utils import pickling
+from cluster_tools.output_store import OutputStore
+
+_T = TypeVar("_T")
+_P = ParamSpec("_P")
+_S = TypeVar("_S")
+
+
+# A job's output is persisted as the pickled tuple `(True, result)` in the executor's
+# `OutputStore` (only in the success case, so that it can serve as a checkpoint).
+# `output_key` selects where; for the default `FileOutputStore` it is a file path.
+# `output_pickle_path` is a deprecated alias.
+class CFutDict(TypedDict):
+    output_key: NotRequired[str]
+    output_pickle_path: NotRequired[str | os.PathLike]
+
+
+def parse_cfut_options(kwargs: dict[str, Any]) -> str | None:
+    """Removes `__cfut_options` from kwargs and returns the output key, if any."""
+    if "__cfut_options" not in kwargs:
+        return None
+    options = cast(CFutDict, kwargs["__cfut_options"])
+    del kwargs["__cfut_options"]
+    output_key = options.get("output_key")
+    output_pickle_path = options.get("output_pickle_path")
+    if output_key is None and output_pickle_path is None:
+        raise ValueError("__cfut_options must contain output_key.")
+    if output_key is not None and output_pickle_path is not None:
+        raise ValueError(
+            "__cfut_options must not contain both output_key and output_pickle_path."
+        )
+    if output_key is not None:
+        return output_key
+    warnings.warn(
+        "The output_pickle_path option of __cfut_options is deprecated, "
+        "use output_key instead.",
+        DeprecationWarning,
+        stacklevel=3,
+    )
+    return str(output_pickle_path)
+
+
+def resolve_output_key_getter(
+    output_pickle_path_getter: Callable[[_S], os.PathLike] | None,
+    output_key_getter: Callable[[_S], str] | None,
+) -> Callable[[_S], str] | None:
+    """Merges the `map_to_futures` getters into a single key getter."""
+    if output_key_getter is not None and output_pickle_path_getter is not None:
+        raise ValueError(
+            "Specify either output_key_getter or output_pickle_path_getter, not both."
+        )
+    if output_key_getter is not None:
+        return output_key_getter
+    if output_pickle_path_getter is not None:
+        warnings.warn(
+            "The output_pickle_path_getter argument of map_to_futures is deprecated, "
+            "use output_key_getter instead.",
+            DeprecationWarning,
+            stacklevel=3,
+        )
+        getter = output_pickle_path_getter
+        return lambda arg: str(getter(arg))
+    return None
+
+
+def cfut_options_kwargs(
+    arg: _S,
+    output_pickle_path_getter: Callable[[_S], os.PathLike] | None,
+    output_key_getter: Callable[[_S], str] | None,
+) -> dict[str, CFutDict]:
+    """Builds the `__cfut_options` kwarg for `submit` from the `map_to_futures` getters."""
+    key_getter = resolve_output_key_getter(output_pickle_path_getter, output_key_getter)
+    if key_getter is None:
+        return {}
+    return {"__cfut_options": {"output_key": key_getter(arg)}}
+
+
+def execute_and_persist(
+    output_store: OutputStore,
+    output_key: str,
+    fn: Callable[_P, _T],
+    *args: _P.args,
+    **kwargs: _P.kwargs,
+) -> _T:
+    """Runs `fn` in the job and stores its result. Used by the in-process executors;
+    the cluster executors do the same in `cluster_tools.remote`."""
+    try:
+        result = fn(*args, **kwargs)
+    except Exception as exc:
+        logging.warning(f"Job computation failed with:\n{exc.__repr__()}")
+        raise exc
+    else:
+        # Only store the result in the success case, since the output
+        # is used as a checkpoint. The cluster executor also stores failures,
+        # since it needs them to transport the exception back.
+        output_store.write(output_key, pickling.dumps((True, result)), success=True)
+        return result

cluster_tools-4.3.0/cluster_tools/_utils/output_wait_thread.py

@@ -0,0 +1,127 @@
+import logging
+import threading
+import time
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from cluster_tools.output_store import OutputStore
+    from cluster_tools.schedulers.cluster_executor import ClusterExecutor
+
+
+class OutputWaitThread(threading.Thread):
+    """A thread that polls the executor's output store waiting for a set of
+    output keys to be written. When a key is written, it invokes a callback.
+    """
+
+    MAX_RETRY = 30
+
+    def __init__(
+        self,
+        callback: Callable[[str, bool], None],
+        executor: "ClusterExecutor",
+        interval: int = 2,
+    ):
+        """The callable ``callback`` will be invoked with the value
+        associated with each key that is written.
+        ``interval`` specifies the polling rate.
+        """
+        threading.Thread.__init__(self)
+        self.callback = callback
+        self.interval = interval
+        self.waiting: dict[str, str] = {}
+        self.retryMap: dict[str, int] = {}
+        self.lock = threading.Lock()
+        self.shutdown = False
+        self.executor = executor
+
+    def stop(self) -> None:
+        """Stop the thread soon."""
+        with self.lock:
+            self.shutdown = True
+
+    def waitFor(self, key: str, value: str) -> None:
+        """Adds a new output key (and its associated callback value) to
+        the set of keys being waited upon.
+        """
+        with self.lock:
+            self.waiting[key] = value
+
+    @staticmethod
+    def _poll_single(store: "OutputStore", key: str) -> set[str]:
+        try:
+            return store.poll([key])
+        except Exception:
+            logging.warning(
+                f"Polling the output store for {key} failed.", exc_info=True
+            )
+            return set()
+
+    def run(self) -> None:
+        def handle_completed_job(job_id: str, key: str, failed_early: bool) -> None:
+            # The callback must not raise, since that would stop this thread and
+            # leave all other jobs pending.
+            try:
+                self.callback(job_id, failed_early)
+            except Exception:
+                logging.exception(f"Handling the completion of job {job_id} failed.")
+            del self.waiting[key]
+
+        while True:
+            with self.lock:
+                if self.shutdown:
+                    return
+
+                pending_tasks = self.executor.get_pending_tasks()
+                store = self.executor.output_store
+
+                # Don't check status of pending tasks, since this
+                # can vastly slow down the polling.
+                keys_to_check = [
+                    key
+                    for key, job_id in self.waiting.items()
+                    if job_id not in pending_tasks
+                ]
+                # Check for outputs in one batch as a fast indicator for job completion.
+                # A store can be remote, so a failing poll must not kill this thread;
+                # it is retried on the next iteration.
+                try:
+                    written_keys = store.poll(keys_to_check)
+                except Exception:
+                    logging.warning(
+                        "Polling the output store failed, retrying.", exc_info=True
+                    )
+                    written_keys = set()
+
+                for key in keys_to_check:
+                    job_id = self.waiting[key]
+
+                    if key in written_keys:
+                        handle_completed_job(job_id, key, False)
+                    elif self.executor is not None:
+                        status = self.executor.check_job_state(job_id)
+
+                        # We have to re-check for the output since this could be written in the mean time
+                        if key in self._poll_single(store, key):
+                            handle_completed_job(job_id, key, False)
+                        else:
+                            if status == "completed":
+                                self.retryMap[key] = self.retryMap.get(key, 0)
+                                self.retryMap[key] += 1
+
+                                if self.retryMap[key] <= OutputWaitThread.MAX_RETRY:
+                                    # Retry by looping again
+                                    logging.warning(
+                                        f"Job state is completed, but output {key} couldn't be found. Retrying {self.retryMap[key]}/{OutputWaitThread.MAX_RETRY}"
+                                    )
+                                else:
+                                    logging.error(
+                                        f"Job state is completed, but output {key} couldn't be found."
+                                    )
+                                    handle_completed_job(job_id, key, True)
+
+                            elif status == "failed":
+                                handle_completed_job(job_id, key, True)
+                            elif status == "ignore":
+                                pass
+            time.sleep(self.interval)

{cluster_tools-4.2.4 → cluster_tools-4.3.0}/cluster_tools/_utils/string_.py

@@ -18,7 +18,3 @@ def random_string(
     length: int = 32, chars: str = (string.ascii_letters + string.digits)
 ) -> str:
     return "".join(random_generator.choice(chars) for i in range(length))
-
-
-def with_preliminary_postfix(name: str) -> str:
-    return f"{name}.preliminary"

{cluster_tools-4.2.4 → cluster_tools-4.3.0}/cluster_tools/executor_protocol.py

@@ -31,6 +31,7 @@ class Executor(Protocol, AbstractContextManager["Executor"]):
         fn: Callable[[_S], _T],
         args: Iterable[_S],
         output_pickle_path_getter: Callable[[_S], PathLike] | None = None,
+        output_key_getter: Callable[[_S], str] | None = None,
     ) -> list[Future[_T]]: ...
 
     def map(

{cluster_tools-4.2.4 → cluster_tools-4.3.0}/cluster_tools/executors/batching.py

@@ -111,10 +111,11 @@ class BatchingExecutor:
         fn: Callable[[_S], _T],
         args: Iterable[_S],
         output_pickle_path_getter: Callable[[_S], PathLike] | None = None,
+        output_key_getter: Callable[[_S], str] | None = None,
     ) -> list[Future[_T]]:
-        if output_pickle_path_getter is not None:
+        if output_key_getter is not None or output_pickle_path_getter is not None:
             raise NotImplementedError(
-                "BatchingExecutor does not support output_pickle_path_getter"
+                "BatchingExecutor does not support output_key_getter or output_pickle_path_getter"
             )
 
         items = list(args)

{cluster_tools-4.2.4 → cluster_tools-4.3.0}/cluster_tools/executors/dask.py

@@ -7,7 +7,6 @@ from concurrent import futures
 from concurrent.futures import Future
 from functools import partial
 from multiprocessing import Queue, get_context
-from pathlib import Path
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -19,8 +18,13 @@ from weakref import ReferenceType, ref
 
 from typing_extensions import ParamSpec
 
+from cluster_tools._utils.cfut_options import (
+    cfut_options_kwargs,
+    execute_and_persist,
+    parse_cfut_options,
+)
 from cluster_tools._utils.warning import enrich_future_with_uncaught_warning
-from cluster_tools.executors.multiprocessing_ import CFutDict, MultiprocessingExecutor
+from cluster_tools.output_store import FileOutputStore, OutputStore
 
 if TYPE_CHECKING:
     from distributed import Client
@@ -108,7 +112,10 @@ class DaskExecutor(futures.Executor):
     is_shutting_down = False
 
     def __init__(
-        self, client: "Client", job_resources: dict[str, Any] | None = None
+        self,
+        client: "Client",
+        job_resources: dict[str, Any] | None = None,
+        output_store: OutputStore | None = None,
     ) -> None:
         try:
             import distributed  # noqa: F401 unused import
@@ -127,6 +134,7 @@ class DaskExecutor(futures.Executor):
         self.client = client
         self.pending_futures = set()
         self.job_resources = job_resources
+        self.output_store = FileOutputStore() if output_store is None else output_store
 
         if self.job_resources is not None:
             # `mem` needs to be a number for dask, so we need to parse it
@@ -151,6 +159,7 @@ class DaskExecutor(futures.Executor):
     def from_config(
         cls,
         job_resources: dict[str, str],
+        output_store: OutputStore | None = None,
         **_kwargs: Any,
     ) -> "DaskExecutor":
         from distributed import Client
@@ -161,7 +170,7 @@ class DaskExecutor(futures.Executor):
             address = os.environ.get("DASK_ADDRESS", None)
 
         client = Client(address=address)
-        return cls(client, job_resources=job_resources)
+        return cls(client, job_resources=job_resources, output_store=output_store)
 
     @classmethod
     def as_completed(cls, futures: list[Future[_T]]) -> Iterator[Future[_T]]:
@@ -175,17 +184,14 @@ class DaskExecutor(futures.Executor):
         *args: _P.args,
         **kwargs: _P.kwargs,
     ) -> Future[_T]:
-        if "__cfut_options" in kwargs:
-            output_pickle_path = cast(CFutDict, kwargs["__cfut_options"])[
-                "output_pickle_path"
-            ]
-            del kwargs["__cfut_options"]
-
+        output_key = parse_cfut_options(kwargs)
+        if output_key is not None:
             __fn = cast(
                 Callable[_P, _T],
                 partial(
-                    MultiprocessingExecutor._execute_and_persist_function,
-                    Path(output_pickle_path),
+                    execute_and_persist,
+                    self.output_store,
+                    output_key,
                     __fn,
                 ),
             )
@@ -245,20 +251,18 @@ class DaskExecutor(futures.Executor):
             _S
         ],  # TODO change: allow more than one arg per call # noqa FIX002 Line contains TODO
         output_pickle_path_getter: Callable[[_S], os.PathLike] | None = None,
+        output_key_getter: Callable[[_S], str] | None = None,
     ) -> list[Future[_T]]:
-        if output_pickle_path_getter is not None:
-            futs = [
-                self.submit(  # type: ignore[call-arg]
-                    fn,
-                    arg,
-                    __cfut_options={
-                        "output_pickle_path": output_pickle_path_getter(arg)
-                    },
-                )
-                for arg in args
-            ]
-        else:
-            futs = [self.submit(fn, arg) for arg in args]
+        futs = [
+            self.submit(  # type: ignore[call-arg]
+                fn,
+                arg,
+                **cfut_options_kwargs(
+                    arg, output_pickle_path_getter, output_key_getter
+                ),
+            )
+            for arg in args
+        ]
 
         return futs
 

{cluster_tools-4.2.4 → cluster_tools-4.3.0}/cluster_tools/executors/multiprocessing_.py

@@ -1,4 +1,3 @@
-import logging
 import multiprocessing
 import os
 from collections.abc import Callable, Iterable, Iterator
@@ -6,39 +5,38 @@ from concurrent import futures
 from concurrent.futures import Future, ProcessPoolExecutor
 from functools import partial
 from multiprocessing.context import BaseContext
-from pathlib import Path
 from typing import (
     Any,
-    TypedDict,
     TypeVar,
     cast,
 )
 
 from typing_extensions import ParamSpec
 
-from cluster_tools._utils import pickling
+from cluster_tools._utils.cfut_options import (
+    cfut_options_kwargs,
+    execute_and_persist,
+    parse_cfut_options,
+)
 from cluster_tools._utils.multiprocessing_logging_handler import (
     _MultiprocessingLoggingHandlerPool,
 )
 from cluster_tools._utils.warning import enrich_future_with_uncaught_warning
-
-# The module name includes a _-suffix to avoid name clashes with the standard library multiprocessing module.
-
-
-class CFutDict(TypedDict):
-    output_pickle_path: str
-
+from cluster_tools.output_store import FileOutputStore, OutputStore
 
 _T = TypeVar("_T")
 _P = ParamSpec("_P")
 _S = TypeVar("_S")
 
 
+# The module name includes a _-suffix to avoid name clashes with the standard library multiprocessing module.
+
+
 class MultiprocessingExecutor(ProcessPoolExecutor):
     """
     Wraps the ProcessPoolExecutor to add various features:
     - map_to_futures method
-    - pickling of job's output (see output_pickle_path_getter and output_pickle_path)
+    - pickling of job's output (see output_key_getter, output_store and CFutDict)
     """
 
     _mp_context: BaseContext
@@ -51,8 +49,10 @@ class MultiprocessingExecutor(ProcessPoolExecutor):
         mp_context: BaseContext | None = None,
         initializer: Callable | None = None,
         initargs: tuple = (),
+        output_store: OutputStore | None = None,
         **__kwargs: Any,
     ) -> None:
+        self.output_store = FileOutputStore() if output_store is None else output_store
         if mp_context is None:
             if start_method is not None:
                 mp_context = multiprocessing.get_context(start_method)
@@ -88,26 +88,21 @@ class MultiprocessingExecutor(ProcessPoolExecutor):
         *args: _P.args,
         **kwargs: _P.kwargs,
     ) -> Future[_T]:
-        if "__cfut_options" in kwargs:
-            output_pickle_path = cast(CFutDict, kwargs["__cfut_options"])[
-                "output_pickle_path"
-            ]
-            del kwargs["__cfut_options"]
-        else:
-            output_pickle_path = None
+        output_key = parse_cfut_options(kwargs)
 
-        # Depending on the start_method and output_pickle_path, setup functions may need to be
+        # Depending on the start_method and output_key, setup functions may need to be
         # executed in the new process context, before the actual code is ran.
         # These wrapper functions consume their arguments from *args, **kwargs and assume
         # that the next and last argument will be another function that is then called.
         # Eventually, the actually submitted function will be called.
 
-        if output_pickle_path is not None:
+        if output_key is not None:
             __fn = cast(
                 Callable[_P, _T],
                 partial(
-                    MultiprocessingExecutor._execute_and_persist_function,
-                    Path(output_pickle_path),
+                    execute_and_persist,
+                    self.output_store,
+                    output_key,
                     __fn,
                 ),
             )
@@ -153,51 +148,23 @@ class MultiprocessingExecutor(ProcessPoolExecutor):
         multiprocessing_logging_setup_fn()
         return fn(*args, **kwargs)
 
-    @staticmethod
-    def _execute_and_persist_function(
-        output_pickle_path: Path,
-        fn: Callable[_P, _T],
-        *args: _P.args,
-        **kwargs: _P.kwargs,
-    ) -> _T:
-        try:
-            result = fn(*args, **kwargs)
-        except Exception as exc:
-            logging.warning(f"Job computation failed with:\n{exc.__repr__()}")
-            raise exc
-        else:
-            # Only pickle the result in the success case, since the output
-            # is used as a checkpoint.
-            # Note that this behavior differs a bit from the cluster executor
-            # which will always serialize the output (even exceptions) to
-            # disk. However, the output will have a .preliminary prefix at first
-            # which is only removed in the success case so that a checkpoint at
-            # the desired target only exists if the job was successful.
-            with output_pickle_path.open("wb") as file:
-                pickling.dump((True, result), file)
-            return result
-
     def map_to_futures(
         self,
         fn: Callable[[_S], _T],
         args: Iterable[_S],
         output_pickle_path_getter: Callable[[_S], os.PathLike] | None = None,
+        output_key_getter: Callable[[_S], str] | None = None,
     ) -> list[Future[_T]]:
-        if output_pickle_path_getter is not None:
-            futs = [
-                self.submit(  # type: ignore[call-arg]
-                    fn,
-                    arg,
-                    __cfut_options={
-                        "output_pickle_path": output_pickle_path_getter(arg)
-                    },
-                )
-                for arg in args
-            ]
-        else:
-            futs = [self.submit(fn, arg) for arg in args]
-
-        return futs
+        return [
+            self.submit(  # type: ignore[call-arg]
+                fn,
+                arg,
+                **cfut_options_kwargs(
+                    arg, output_pickle_path_getter, output_key_getter
+                ),
+            )
+            for arg in args
+        ]
 
     def forward_log(self, fut: Future[_T]) -> _T:
         """

{cluster_tools-4.2.4 → cluster_tools-4.3.0}/cluster_tools/executors/sequential.py

@@ -2,13 +2,17 @@ import warnings
 from collections.abc import Callable, Iterable, Iterator
 from concurrent.futures import Executor, Future, as_completed
 from os import PathLike
-from pathlib import Path
-from typing import Any, TypeVar, cast
+from typing import Any, TypeVar
 
 from typing_extensions import ParamSpec
 
+from cluster_tools._utils.cfut_options import (
+    cfut_options_kwargs,
+    execute_and_persist,
+    parse_cfut_options,
+)
 from cluster_tools._utils.warning import enrich_future_with_uncaught_warning
-from cluster_tools.executors.multiprocessing_ import CFutDict, MultiprocessingExecutor
+from cluster_tools.output_store import FileOutputStore, OutputStore
 
 _T = TypeVar("_T")
 _S = TypeVar("_S")
@@ -22,9 +26,10 @@ class SequentialExecutor(Executor):
 
     def __init__(
         self,
+        output_store: OutputStore | None = None,
         **__kwargs: Any,
     ) -> None:
-        pass
+        self.output_store = FileOutputStore() if output_store is None else output_store
 
     def submit(  # type: ignore[override]
         self,
@@ -33,13 +38,11 @@ class SequentialExecutor(Executor):
         **kwargs: _P.kwargs,
     ) -> Future[_T]:
         fut: Future[_T] = Future()
-        if "__cfut_options" in kwargs:
-            output_pickle_path = cast(CFutDict, kwargs["__cfut_options"])[
-                "output_pickle_path"
-            ]
-            del kwargs["__cfut_options"]
-            result = MultiprocessingExecutor._execute_and_persist_function(
-                Path(output_pickle_path),
+        output_key = parse_cfut_options(kwargs)
+        if output_key is not None:
+            result = execute_and_persist(
+                self.output_store,
+                output_key,
                 __fn,
                 *args,
                 **kwargs,
@@ -60,22 +63,18 @@ class SequentialExecutor(Executor):
         fn: Callable[[_S], _T],
         args: Iterable[_S],
         output_pickle_path_getter: Callable[[_S], PathLike] | None = None,
+        output_key_getter: Callable[[_S], str] | None = None,
     ) -> list[Future[_T]]:
-        if output_pickle_path_getter is not None:
-            futs = [
-                self.submit(  # type: ignore[call-arg]
-                    fn,
-                    arg,
-                    __cfut_options={
-                        "output_pickle_path": output_pickle_path_getter(arg)
-                    },
-                )
-                for arg in args
-            ]
-        else:
-            futs = [self.submit(fn, arg) for arg in args]
-
-        return futs
+        return [
+            self.submit(  # type: ignore[call-arg]
+                fn,
+                arg,
+                **cfut_options_kwargs(
+                    arg, output_pickle_path_getter, output_key_getter
+                ),
+            )
+            for arg in args
+        ]
 
     def map(  # type: ignore[override]
         self,