[Step] Add Items distribution (array/pools)

Author: jrapin

exca/cachedict/inflight.py

@@ -356,28 +356,18 @@ def wait_for_inflight(
 def inflight_session(
     reg: InflightRegistry | None,
     item_uids: tp.Collection[str],
-    *,
-    local: bool = False,
 ) -> tp.Iterator[list[str]]:
     """Wait for in-flight items, claim available ones, release+close on exit.
 
     When *reg* is ``None`` (no cache folder), yields all *item_uids*
     unchanged so that callers never need a ``None`` guard.
 
-    Parameters
-    ----------
-    local:
-        Set to ``True`` when items will be processed locally (no Slurm
-        submission). Stamps claims with ``_LOCAL_JOB_ID`` so ``is_alive``
-        can distinguish "local work in progress" from "Slurm submission
-        that never completed ``update_worker_info``".
-
     Self-deadlock is prevented internally: ``wait_for_inflight`` skips items
     owned by the current PID, and ``claim`` treats same-PID rows as already
     ours.
 
-    The registry connection is closed on exit; callers must perform any
-    ``record_worker_info`` calls inside the ``with`` block.
+    Callers should call ``record_worker_info`` inside the ``with`` block
+    to stamp claimed items with an appropriate liveness signal.
     """
     if reg is None:
         yield list(item_uids)
@@ -406,8 +396,6 @@ def inflight_session(
         msg = "Claim race: got %d/%d items, re-waiting"
         logger.info(msg, len(claimed), len(item_uids))
         time.sleep(random.uniform(0.5, 2.0))
-    if local:
-        reg.update_worker_info(claimed, job_id=_LOCAL_JOB_ID)
     try:
         yield claimed
     finally:
@@ -417,10 +405,10 @@ def inflight_session(
 
 
 def record_worker_info(
-    reg: InflightRegistry, item_uids: list[str], job: submitit.Job[tp.Any]
+    reg: InflightRegistry, item_uids: list[str], job: tp.Any = None
 ) -> None:
-    """Stamp a submitit *job*'s worker info on *item_uids*. Slurm jobs get
-    job_id + folder; other backends get the local sentinel."""
+    """Stamp worker info on *item_uids*. Slurm jobs get job_id + folder;
+    everything else (including no job) gets the local PID sentinel."""
     if isinstance(job, submitit.SlurmJob):
         reg.update_worker_info(
             item_uids, job_id=str(job.job_id), job_folder=str(job.paths.folder)

exca/cachedict/test_inflight.py

@@ -79,9 +79,11 @@ def fresh() -> inflight.InflightRegistry:
             raise ValueError("boom")
     assert seen(["a"]) == {}
 
-    # Local: claims marked with job_id="local".
-    with inflight.inflight_session(fresh(), ["loc"], local=True) as claimed:
+    # Local: record_worker_info without job stamps _LOCAL_JOB_ID.
+    reg = fresh()
+    with inflight.inflight_session(reg, ["loc"]) as claimed:
         assert claimed == ["loc"]
+        inflight.record_worker_info(reg, claimed)
         assert seen(["loc"])["loc"].job_id == inflight._LOCAL_JOB_ID
 
     # Nested: inner session must NOT release outer's claim.

exca/map.py

@@ -18,9 +18,8 @@
 import numpy as np
 import pydantic
 
-from . import base, slurm
+from . import base, slurm, utils
 from .cachedict import CacheDict, inflight
-from .utils import ShortItemUid
 
 
 @dataclasses.dataclass
@@ -70,46 +69,7 @@ def __call__(self, items: tp.Sequence[tp.Any]) -> tp.Iterator[tp.Any]:
         return self.infra._method_override(items)
 
 
-def _make_pool_executor(pool: str, max_workers: int) -> futures.Executor:
-    """Falls back to ThreadPoolExecutor (with a warning) if ``pool="processpool"``
-    and ``ProcessPoolExecutor`` cannot be created (eg ``sem_open`` EPERM)."""
-    if pool == "processpool":
-        try:
-            return futures.ProcessPoolExecutor(max_workers=max_workers)
-        except PermissionError as e:
-            logger.warning(
-                "ProcessPoolExecutor unavailable (%s); falling back to "
-                "ThreadPoolExecutor. Set cluster='threadpool' or cluster=None "
-                "to silence this warning.",
-                e,
-            )
-    return futures.ThreadPoolExecutor(max_workers=max_workers)
-
-
-def to_chunks(
-    items: list[X], *, max_chunks: int | None, min_items_per_chunk: int = 1
-) -> tp.Iterator[list[X]]:
-    """Split a list of items into several smaller list of items
-
-    Parameters
-    ----------
-    max_chunks: optional int
-        maximum number of chunks to create
-    min_items_per_chunk: int
-        minimum number of items per chunk
-
-    Yields
-    ------
-    list of items
-    """
-    splits = min(
-        len(items) if max_chunks is None else max_chunks,
-        int(np.ceil(len(items) / min_items_per_chunk)),
-    )
-    items_per_chunk = int(np.ceil(len(items) / splits))
-    for k in range(splits):
-        # select a batch/chunk of samples_per_job items to send to a job
-        yield items[k * items_per_chunk : (k + 1) * items_per_chunk]
+_make_pool_executor = utils.make_pool_executor  # deprecated
 
 
 class MapInfra(base.BaseInfra, slurm.SubmititMixin):
@@ -281,7 +241,9 @@ def apply(
         params.pop("self")
         max_length = params.pop("item_uid_max_length")
         if max_length is not None:
-            params["item_uid"] = ShortItemUid(params["item_uid"], max_length=max_length)
+            params["item_uid"] = utils.ShortItemUid(
+                params["item_uid"], max_length=max_length
+            )
         if self._infra_method is not None:
             raise RuntimeError(f"Infra was already applied: {self._infra_method}")
 
@@ -392,7 +354,7 @@ def _method_override(self, *args: tp.Any, **kwargs: tp.Any) -> tp.Iterator[tp.An
                 if missing:
                     jobs: list[tp.Any] = []
                     uid_item_chunks = list(
-                        to_chunks(
+                        utils.to_chunks(
                             missing,
                             max_chunks=self.max_jobs,
                             min_items_per_chunk=self.min_samples_per_job,
@@ -455,9 +417,10 @@ def _method_override_futures(self, items: tp.Sequence[tp.Any]) -> tp.Iterator[tp
                 pool = None
             # avoid processing same files at same time if several jobs overlap
             np.random.shuffle(missing)
-            with inflight.inflight_session(
-                self._inflight_registry(), [k for k, _ in missing], local=True
-            ) as claimed_uids:
+            reg = self._inflight_registry()
+            with inflight.inflight_session(reg, [k for k, _ in missing]) as claimed_uids:
+                if reg is not None:
+                    inflight.record_worker_info(reg, claimed_uids)
                 claimed_set = set(claimed_uids)
                 # Re-check cache after wait: other workers may have completed
                 # items while we were blocked in inflight_session.
@@ -487,7 +450,7 @@ def _method_override_futures(self, items: tp.Sequence[tp.Any]) -> tp.Iterator[tp
                         max_workers = min(max_workers, self.max_jobs)
                     with _make_pool_executor(pool, max_workers) as ex:
                         mitems = [ki[1] for ki in missing]
-                        chunks = to_chunks(mitems, max_chunks=3 * max_workers)
+                        chunks = utils.to_chunks(mitems, max_chunks=3 * max_workers)
                         for chunk in chunks:
                             j = ex.submit(
                                 self._call_and_store,

exca/steps/backends.py

@@ -14,16 +14,20 @@
 
 import dataclasses
 import logging
+import os
+import random
 import shutil
 import traceback
 import typing as tp
 import warnings
+from concurrent import futures
 from pathlib import Path
 
 import pydantic
 import submitit
 
 import exca
+from exca import utils
 from exca.cachedict import inflight
 
 from . import errors, identity, items
@@ -280,7 +284,7 @@ def _exclude_from_cls_uid(cls) -> list[str]:
         return ["."]  # force ignored in uid
 
     # Used by Backend._run for the inflight registry (concurrent worker safety).
-    _is_off_process: tp.ClassVar[bool] = False
+    _concurrent: tp.ClassVar[bool] = False
 
     folder: Path | None = None
     # deprecated: declare `CACHE_TYPE` on the Step subclass.
@@ -364,7 +368,7 @@ def _clear_cache(
         uid: str,
     ) -> None:
         """Drop everything cached for this uid (cd row, error row, job folder)."""
-        if self._is_off_process and paths.cache_folder.exists():
+        if self._concurrent and paths.cache_folder.exists():
             try:
                 reg = inflight.InflightRegistry(paths.cache_folder)
                 info = reg.get([uid])
@@ -421,23 +425,19 @@ def _run(self, step: Step, batch: items.StepItems) -> items.StepItems:
             for uid in to_compute:
                 paths._ensure_folders(uid)
             reg: inflight.InflightRegistry | None = None
-            if self._is_off_process:
+            if self._concurrent:
                 reg = inflight.InflightRegistry(paths.cache_folder)
             with inflight.inflight_session(reg, to_compute):
-                # Re-check after lock (another worker may have finished).
                 recheck = _CachedEntry.lookup_statuses(cd, to_compute)
-                still_uids = [u for u in to_compute if recheck[u] != "success"]
-                if still_uids:
-                    filtered = batch.select(still_uids)
+                pending_uids = [u for u in to_compute if recheck[u] != "success"]
+                if pending_uids:
+                    filtered = batch.select(pending_uids)
                     wrapper = _CachingCall(step, cd, paths.step_uid)
                     try:
-                        job = self._submit(wrapper, filtered, paths=paths)
-                        if reg is not None:
-                            inflight.record_worker_info(reg, still_uids, job)
-                        job.result()
+                        self._execute(wrapper, filtered, paths=paths, reg=reg)
                     finally:
                         if mode in ("force", "retry"):
-                            self._recomputed.update(still_uids)
+                            self._recomputed.update(pending_uids)
                 verify = _CachedEntry.lookup_statuses(cd, to_compute)
                 for uid in to_compute:
                     if verify[uid] != "success":
@@ -452,17 +452,16 @@ def _run(self, step: Step, batch: items.StepItems) -> items.StepItems:
             mode=mode,
         )
 
-    def _submit(self, wrapper: _CachingCall, *args: tp.Any, paths: StepPaths) -> tp.Any:
-        """Submit wrapper for execution. Default: inline execution."""
-        wrapper(*args)
-        return _InlineJob()
-
-
-class _InlineJob:
-    """Completion signal for inline execution; the value lives in cache."""
-
-    def result(self) -> None:
-        return None
+    def _execute(
+        self,
+        wrapper: _CachingCall,
+        pending: items.StepItems,
+        *,
+        paths: StepPaths,
+        reg: inflight.InflightRegistry | None,
+    ) -> None:
+        """Run *wrapper* on *pending* items. Override for chunking/pools/arrays."""
+        wrapper(pending)
 
 
 class Cached(Backend):
@@ -473,7 +472,7 @@ class _SubmititBackend(Backend):
     """Base for submitit backends."""
 
     # Submitit cloud-pickles inputs to a worker (`SubmititDebug` overrides: inline).
-    _is_off_process: tp.ClassVar[bool] = True
+    _concurrent: tp.ClassVar[bool] = True
 
     job_name: str | None = None
     timeout_min: int | None = None
@@ -482,24 +481,51 @@ class _SubmititBackend(Backend):
     cpus_per_task: int | None = None
     gpus_per_node: int | None = None
     mem_gb: float | None = None
+    max_jobs: int = 128
+    min_items_per_job: int = 1
 
     # passed as `cluster=` to submitit.AutoExecutor; subclasses pin it.
     _CLUSTER: tp.ClassVar[str | None] = None
 
     def _submitit_params(self) -> dict[str, tp.Any]:
         """Build the kwargs dict forwarded to ``AutoExecutor.update_parameters``."""
         fields = set(type(self).model_fields) - set(Backend.model_fields)
-        params = {k: getattr(self, k) for k in fields if getattr(self, k) is not None}
+        skip = {"max_jobs", "min_items_per_job"}
+        params = {
+            k: getattr(self, k) for k in fields - skip if getattr(self, k) is not None
+        }
         if "job_name" in params:
             params["name"] = params.pop("job_name")
         return params
 
-    def _submit(self, wrapper: _CachingCall, *args: tp.Any, paths: StepPaths) -> tp.Any:
+    def _execute(
+        self,
+        wrapper: _CachingCall,
+        pending: items.StepItems,
+        *,
+        paths: StepPaths,
+        reg: inflight.InflightRegistry | None,
+    ) -> None:
+        uids = list(pending.uids)
+        random.shuffle(uids)
+        chunks = [
+            pending.select(c)
+            for c in utils.to_chunks(
+                uids, max_chunks=self.max_jobs, min_items_per_chunk=self.min_items_per_job
+            )
+        ]
         executor = submitit.AutoExecutor(folder=paths._logs_folder, cluster=self._CLUSTER)
-        executor.update_parameters(**self._submitit_params())
-        with submitit.helpers.clean_env():
-            job = executor.submit(wrapper, *args)
-        return job
+        params = self._submitit_params()
+        if self._CLUSTER in ("slurm", None):
+            params["slurm_array_parallelism"] = len(chunks)
+        executor.update_parameters(**params)
+        with submitit.helpers.clean_env(), executor.batch():
+            jobs = [executor.submit(wrapper, c) for c in chunks]
+        if reg is not None:
+            for c, j in zip(chunks, jobs):
+                inflight.record_worker_info(reg, c.uids, j)
+        for j in jobs:
+            j.result()
 
 
 class LocalProcess(_SubmititBackend):
@@ -512,7 +538,7 @@ class SubmititDebug(_SubmititBackend):
     """Debug executor (inline but simulates submitit)."""
 
     _CLUSTER: tp.ClassVar[str | None] = "debug"
-    _is_off_process: tp.ClassVar[bool] = False
+    _concurrent: tp.ClassVar[bool] = False
 
 
 class Slurm(_SubmititBackend):
@@ -542,3 +568,54 @@ class Auto(Slurm):
     """Auto-detect executor (local or Slurm). Slurm fields only apply on slurm."""
 
     _CLUSTER: tp.ClassVar[str | None] = None
+
+
+class _PoolBackend(Backend):
+    """Base for concurrent.futures pool backends."""
+
+    max_jobs: int | None = 128
+    _POOL_TYPE: tp.ClassVar[str]
+
+    def _execute(
+        self,
+        wrapper: _CachingCall,
+        pending: items.StepItems,
+        *,
+        paths: StepPaths,
+        reg: inflight.InflightRegistry | None,
+    ) -> None:
+        uids = list(pending.uids)
+        random.shuffle(uids)
+        cpus = max(1, (os.cpu_count() or 1) - 1)
+        max_workers = min(len(uids), cpus)
+        if self.max_jobs is not None:
+            max_workers = min(max_workers, self.max_jobs)
+        chunks = [
+            pending.select(c) for c in utils.to_chunks(uids, max_chunks=3 * max_workers)
+        ]
+        if reg is not None:
+            for c in chunks:
+                inflight.record_worker_info(reg, c.uids)
+        with utils.make_pool_executor(self._POOL_TYPE, max_workers) as pool:
+            futs = [pool.submit(wrapper, c) for c in chunks]
+            try:
+                for f in futures.as_completed(futs):
+                    f.result()
+            except BaseException:
+                for f in futs:
+                    f.cancel()
+                raise
+
+
+class ProcessPool(_PoolBackend):
+    """Process pool execution + caching."""
+
+    _POOL_TYPE: tp.ClassVar[str] = "processpool"
+    _concurrent: tp.ClassVar[bool] = True
+
+
+class ThreadPool(_PoolBackend):
+    """Thread pool execution + caching."""
+
+    _POOL_TYPE: tp.ClassVar[str] = "threadpool"
+    _concurrent: tp.ClassVar[bool] = True