cluster_tools 3.2.4__tar.gz → 3.4.0__tar.gz

{cluster_tools-3.2.4 → cluster_tools-3.4.0}/PKG-INFO

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

{cluster_tools-3.2.4 → cluster_tools-3.4.0}/cluster_tools/__init__.py

@@ -1,6 +1,7 @@
 from typing import Any, Literal, overload
 
 from cluster_tools.executor_protocol import Executor
+from cluster_tools.executors.batching import BatchingExecutor
 from cluster_tools.executors.dask import DaskExecutor
 from cluster_tools.executors.multiprocessing_ import MultiprocessingExecutor
 from cluster_tools.executors.multiprocessing_pickle import MultiprocessingPickleExecutor
@@ -86,6 +87,12 @@ def get_executor(
 ) -> SequentialExecutor: ...
 
 
+@overload
+def get_executor(
+    environment: Literal["batching"], **kwargs: Any
+) -> BatchingExecutor: ...
+
+
 @overload
 def get_executor(
     environment: Literal["sequential_with_pickling"], **kwargs: Any
@@ -113,6 +120,30 @@ def get_executor(environment: str, **kwargs: Any) -> "Executor":
         return MultiprocessingExecutor(**kwargs)
     elif environment == "sequential":
         return SequentialExecutor(**kwargs)
+    elif environment == "batching":
+        executor_config = kwargs.get("executor")
+        if not isinstance(executor_config, dict):
+            raise ValueError(
+                "The 'batching' executor requires a nested 'executor' config."
+            )
+
+        inner_executor_config = executor_config.copy()
+        name = inner_executor_config.pop("name", None)
+        if name is None:
+            raise ValueError(
+                "The 'batching' executor's nested 'executor' config requires a 'name' key."
+            )
+
+        inner_executor = get_executor(name, **inner_executor_config)
+        batch_size = kwargs.get("batch_size")
+        target_job_count = kwargs.get("target_job_count")
+        if batch_size is None and target_job_count is None:
+            raise ValueError(
+                "The 'batching' executor requires either a 'batch_size' or 'target_job_count' key."
+            )
+        return BatchingExecutor(
+            inner_executor, batch_size=batch_size, target_job_count=target_job_count
+        )
     elif environment == "sequential_with_pickling":
         return SequentialPickleExecutor(**kwargs)
     elif environment == "multiprocessing_with_pickling":

cluster_tools-3.4.0/cluster_tools/executors/batching.py

@@ -0,0 +1,177 @@
+import math
+import warnings
+from collections.abc import Callable, Iterable, Iterator
+from concurrent.futures import Future, as_completed
+from functools import partial
+from itertools import islice
+from os import PathLike
+from types import TracebackType
+from typing import TypeVar
+
+from typing_extensions import ParamSpec
+
+from cluster_tools.executor_protocol import Executor
+
+_T = TypeVar("_T")
+_S = TypeVar("_S")
+_P = ParamSpec("_P")
+
+
+def _apply_fn_to_batch(fn: Callable, batch: list) -> list:
+    return [fn(item) for item in batch]
+
+
+def _iter_batches(iterable: Iterable[_S], batch_size: int) -> Iterator[list[_S]]:
+    it = iter(iterable)
+    while batch := list(islice(it, batch_size)):
+        yield batch
+
+
+class BatchingExecutor:
+    """
+    Wraps another executor and groups items into batches before submission.
+
+    Specify exactly one of:
+    - batch_size: fixed number of items per job (must be > 0)
+    - target_job_count: desired number of jobs; batch size is computed as
+      ceil(n_items / target_job_count) for each map call (must be > 0)
+    """
+
+    _executor: Executor
+    _batch_size: int | None
+    _target_job_count: int | None
+
+    def __init__(
+        self,
+        executor: Executor,
+        *,
+        batch_size: int | None = None,
+        target_job_count: int | None = None,
+    ) -> None:
+        if batch_size is not None and target_job_count is not None:
+            raise ValueError("Specify either batch_size or target_job_count, not both")
+        if batch_size is None and target_job_count is None:
+            raise ValueError("Either batch_size or target_job_count must be specified")
+        if batch_size is not None and batch_size <= 0:
+            raise ValueError("batch_size must be greater than 0")
+        if target_job_count is not None and target_job_count <= 0:
+            raise ValueError("target_job_count must be greater than 0")
+        self._executor = executor
+        self._batch_size = batch_size
+        self._target_job_count = target_job_count
+
+    @property
+    def batch_size(self) -> int | None:
+        return self._batch_size
+
+    @property
+    def target_job_count(self) -> int | None:
+        return self._target_job_count
+
+    def _resolve_batch_size(self, n_items: int) -> int:
+        if self._batch_size is not None:
+            return self._batch_size
+        assert self._target_job_count is not None
+        if n_items == 0:
+            return 1
+        return math.ceil(n_items / self._target_job_count)
+
+    def __enter__(self) -> "BatchingExecutor":
+        self._executor.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        self._executor.__exit__(exc_type, exc_val, exc_tb)
+
+    @classmethod
+    def as_completed(cls, futures: list[Future[_T]]) -> Iterator[Future[_T]]:
+        return as_completed(futures)
+
+    def submit(
+        self,
+        __fn: Callable[_P, _T],
+        /,
+        *args: _P.args,
+        **kwargs: _P.kwargs,
+    ) -> Future[_T]:
+        warnings.warn(
+            "BatchingExecutor.submit() bypasses batching and submits a single job directly.",
+            stacklevel=2,
+        )
+        return self._executor.submit(__fn, *args, **kwargs)
+
+    def map_to_futures(
+        self,
+        fn: Callable[[_S], _T],
+        args: Iterable[_S],
+        output_pickle_path_getter: Callable[[_S], PathLike] | None = None,
+    ) -> list[Future[_T]]:
+        if output_pickle_path_getter is not None:
+            raise NotImplementedError(
+                "BatchingExecutor does not support output_pickle_path_getter"
+            )
+
+        items = list(args)
+        batch_size = self._resolve_batch_size(len(items))
+        all_item_futures: list[Future[_T]] = []
+
+        for batch in _iter_batches(items, batch_size):
+            batch_fn = partial(_apply_fn_to_batch, fn)
+            batch_fn.__name__ = f"batch_{fn.__name__}"  # type: ignore[attr-defined]
+            (batch_future,) = self._executor.map_to_futures(batch_fn, [batch])
+            item_futures: list[Future[_T]] = [Future() for _ in batch]
+            all_item_futures.extend(item_futures)
+
+            def on_batch_done(
+                batch_future: Future[list[_T]],
+                _item_futures: list[Future[_T]] = item_futures,
+            ) -> None:
+                try:
+                    results = batch_future.result()
+                    for future, result in zip(_item_futures, results):
+                        future.set_result(result)
+                except Exception as e:
+                    for future in _item_futures:
+                        future.set_exception(e)
+
+            batch_future.add_done_callback(on_batch_done)
+
+        return all_item_futures
+
+    def map(
+        self,
+        fn: Callable[[_S], _T],
+        iterables: Iterable[_S],
+        timeout: float | None = None,
+        chunksize: int | None = None,
+    ) -> Iterator[_T]:
+        if chunksize is not None:
+            raise ValueError(
+                "BatchingExecutor does not support chunksize. Use batch_size instead."
+            )
+
+        items = list(iterables)
+        batch_size = self._resolve_batch_size(len(items))
+
+        def result_generator() -> Iterator[_T]:
+            batch_fn = partial(_apply_fn_to_batch, fn)
+            batch_fn.__name__ = f"batch_{fn.__name__}"  # type: ignore[attr-defined]
+            for batch_results in self._executor.map(
+                batch_fn,
+                _iter_batches(items, batch_size),
+                timeout=timeout,
+            ):
+                yield from batch_results
+
+        return result_generator()
+
+    def forward_log(self, fut: Future[_T]) -> _T:
+        return self._executor.forward_log(fut)
+
+    def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None:
+        self._executor.shutdown(wait=wait, cancel_futures=cancel_futures)

cluster_tools-3.4.0/cluster_tools/version.py

@@ -0,0 +1 @@
+__version__ = '3.4.0'

{cluster_tools-3.2.4 → cluster_tools-3.4.0}/cluster_tools.egg-info/PKG-INFO

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

{cluster_tools-3.2.4 → cluster_tools-3.4.0}/cluster_tools.egg-info/SOURCES.txt

@@ -18,6 +18,7 @@ cluster_tools/_utils/reflection.py
 cluster_tools/_utils/string_.py
 cluster_tools/_utils/tailf.py
 cluster_tools/_utils/warning.py
+cluster_tools/executors/batching.py
 cluster_tools/executors/dask.py
 cluster_tools/executors/multiprocessing_.py
 cluster_tools/executors/multiprocessing_pickle.py
@@ -29,6 +30,7 @@ cluster_tools/schedulers/kube.py
 cluster_tools/schedulers/pbs.py
 cluster_tools/schedulers/slurm.py
 tests/test_all.py
+tests/test_batching.py
 tests/test_dask.py
 tests/test_deref_main.py
 tests/test_kubernetes.py

{cluster_tools-3.2.4 → cluster_tools-3.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cluster_tools"
-version = "3.2.4"  # filled by dunamai
+version = "3.4.0"  # filled by dunamai
 description = "Utility library for easily distributing code execution on clusters"
 authors = [{name= "scalable minds", email="hello@scalableminds.com"}]
 readme = "README.md"

cluster_tools-3.4.0/tests/test_batching.py

@@ -0,0 +1,200 @@
+from concurrent.futures import CancelledError, Future
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+import cluster_tools
+from cluster_tools.executors.batching import BatchingExecutor
+from cluster_tools.executors.multiprocessing_ import MultiprocessingExecutor
+from cluster_tools.executors.sequential import SequentialExecutor
+
+
+def double(x: int) -> int:
+    return x * 2
+
+
+def raise_runtime_error(_: int) -> int:
+    raise RuntimeError("failure")
+
+
+@pytest.fixture(params=["sequential", "multiprocessing"])
+def inner_executor(request: pytest.FixtureRequest) -> cluster_tools.Executor:
+    if request.param == "sequential":
+        return SequentialExecutor()
+    return MultiprocessingExecutor(max_workers=2)
+
+
+def test_map_basic(inner_executor: cluster_tools.Executor) -> None:
+    with BatchingExecutor(inner_executor, batch_size=3) as executor:
+        results = list(executor.map(double, [1, 2, 3, 4, 5]))
+    assert results == [2, 4, 6, 8, 10]
+
+
+def test_map_exact_batch_boundary() -> None:
+    with BatchingExecutor(SequentialExecutor(), batch_size=3) as executor:
+        results = list(executor.map(double, [1, 2, 3]))
+    assert results == [2, 4, 6]
+
+
+def test_map_single_item() -> None:
+    with BatchingExecutor(SequentialExecutor(), batch_size=3) as executor:
+        results = list(executor.map(double, [5]))
+    assert results == [10]
+
+
+def test_map_empty() -> None:
+    with BatchingExecutor(SequentialExecutor(), batch_size=3) as executor:
+        results = list(executor.map(double, []))
+    assert results == []
+
+
+def test_map_calls_inner_with_batches() -> None:
+    """Verify the inner executor receives batches, not individual items."""
+    inner = SequentialExecutor()
+    received_batches: list[list[int]] = []
+    original_map = inner.map
+
+    def spy(fn, iterables, **kwargs):  # type: ignore[no-untyped-def]
+        batches = list(iterables)
+        received_batches.extend(batches)
+        return original_map(fn, iter(batches), **kwargs)
+
+    with patch.object(inner, "map", side_effect=spy):
+        with BatchingExecutor(inner, batch_size=3) as executor:
+            results = list(executor.map(double, range(7)))
+
+    assert results == [0, 2, 4, 6, 8, 10, 12]
+    assert received_batches == [[0, 1, 2], [3, 4, 5], [6]]
+
+
+def test_map_to_futures(inner_executor: cluster_tools.Executor) -> None:
+    with BatchingExecutor(inner_executor, batch_size=3) as executor:
+        futures = executor.map_to_futures(double, [1, 2, 3, 4, 5])
+    assert len(futures) == 5
+    assert [f.result() for f in futures] == [2, 4, 6, 8, 10]
+
+
+def test_map_to_futures_empty() -> None:
+    with BatchingExecutor(SequentialExecutor(), batch_size=3) as executor:
+        futures = executor.map_to_futures(double, [])
+    assert futures == []
+
+
+def test_map_chunksize_raises() -> None:
+    with BatchingExecutor(SequentialExecutor(), batch_size=3) as executor:
+        with pytest.raises(ValueError):
+            list(executor.map(double, [1, 2, 3], chunksize=2))
+
+
+def test_map_to_futures_output_pickle_path_getter_raises() -> None:
+    with BatchingExecutor(SequentialExecutor(), batch_size=3) as executor:
+        with pytest.raises(NotImplementedError):
+            executor.map_to_futures(
+                double, [1], output_pickle_path_getter=lambda _: Path("/tmp/x")
+            )
+
+
+def test_map_to_futures_exception_propagates_to_all_item_futures() -> None:
+    # SequentialExecutor raises synchronously, so use MultiprocessingExecutor
+    # to exercise the on_batch_done exception path (failure arrives via future).
+    with BatchingExecutor(
+        MultiprocessingExecutor(max_workers=1), batch_size=3
+    ) as executor:
+        futures = executor.map_to_futures(raise_runtime_error, [1, 2, 3])
+
+    assert len(futures) == 3
+    for fut in futures:
+        with pytest.raises(RuntimeError, match="failure"):
+            fut.result()
+
+
+def test_submit_warns() -> None:
+    with BatchingExecutor(SequentialExecutor(), batch_size=3) as executor:
+        with pytest.warns(UserWarning, match="bypasses batching"):
+            fut = executor.submit(double, 5)
+    assert fut.result() == 10
+
+
+def test_init_requires_one_of_batch_size_or_target_job_count() -> None:
+    with pytest.raises(ValueError, match="Either batch_size or target_job_count"):
+        BatchingExecutor(SequentialExecutor())
+
+
+def test_init_rejects_both() -> None:
+    with pytest.raises(ValueError, match="not both"):
+        BatchingExecutor(SequentialExecutor(), batch_size=3, target_job_count=2)
+
+
+def test_init_rejects_zero_batch_size() -> None:
+    with pytest.raises(ValueError, match="batch_size must be greater than 0"):
+        BatchingExecutor(SequentialExecutor(), batch_size=0)
+
+
+def test_init_rejects_zero_target_job_count() -> None:
+    with pytest.raises(ValueError, match="target_job_count must be greater than 0"):
+        BatchingExecutor(SequentialExecutor(), target_job_count=0)
+
+
+def test_target_job_count_map() -> None:
+    """With target_job_count=2 and 7 items, batch_size=ceil(7/2)=4 → 2 jobs."""
+    inner = SequentialExecutor()
+    received_batches: list[list[int]] = []
+    original_map = inner.map
+
+    def spy(fn, iterables, **kwargs):  # type: ignore[no-untyped-def]
+        batches = list(iterables)
+        received_batches.extend(batches)
+        return original_map(fn, iter(batches), **kwargs)
+
+    with patch.object(inner, "map", side_effect=spy):
+        with BatchingExecutor(inner, target_job_count=2) as executor:
+            results = list(executor.map(double, range(7)))
+
+    assert results == [0, 2, 4, 6, 8, 10, 12]
+    assert received_batches == [[0, 1, 2, 3], [4, 5, 6]]
+
+
+def test_target_job_count_map_to_futures() -> None:
+    with BatchingExecutor(SequentialExecutor(), target_job_count=2) as executor:
+        futures = executor.map_to_futures(double, [1, 2, 3, 4, 5])
+    assert [f.result() for f in futures] == [2, 4, 6, 8, 10]
+
+
+def test_target_job_count_empty() -> None:
+    with BatchingExecutor(SequentialExecutor(), target_job_count=3) as executor:
+        results = list(executor.map(double, []))
+    assert results == []
+
+
+def test_get_executor() -> None:
+    executor = cluster_tools.get_executor(
+        "batching",
+        executor={
+            "name": "multiprocessing",
+            "max_workers": 3,
+        },
+        batch_size=10,
+    )
+    assert isinstance(executor, BatchingExecutor)
+    assert isinstance(executor._executor, MultiprocessingExecutor)
+    assert executor._executor._max_workers == 3  # type: ignore[attr-defined]
+    assert executor.batch_size == 10
+
+
+def test_get_executor_target_job_count() -> None:
+    executor = cluster_tools.get_executor(
+        "batching",
+        executor={"name": "sequential"},
+        target_job_count=4,
+    )
+    assert isinstance(executor, BatchingExecutor)
+    assert executor.target_job_count == 4
+
+
+def test_get_executor_missing_sizing_raises() -> None:
+    with pytest.raises(ValueError, match="batch_size.*target_job_count"):
+        cluster_tools.get_executor(
+            "batching",
+            executor={"name": "sequential"},
+        )

cluster_tools-3.2.4/cluster_tools/version.py

@@ -1 +0,0 @@
-__version__ = '3.2.4'