From 652de88badcd123e10aff92d6158ba143c16cab9 Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Mon, 9 Feb 2026 15:23:55 +0100 Subject: [PATCH 1/2] update --- exca/steps/__init__.py | 1 + exca/steps/backends.py | 96 +++++++++++++++- exca/steps/base.py | 249 ++++++++++++++++++++++++++++++++++++++++- exca/steps/conftest.py | 11 ++ 4 files changed, 354 insertions(+), 3 deletions(-) diff --git a/exca/steps/__init__.py b/exca/steps/__init__.py index 80eb53c1..99a96dc3 100644 --- a/exca/steps/__init__.py +++ b/exca/steps/__init__.py @@ -27,4 +27,5 @@ from . import backends from .base import Chain as Chain +from .base import Items as Items from .base import Step as Step diff --git a/exca/steps/backends.py b/exca/steps/backends.py index 2860fc8b..70d7635d 100644 --- a/exca/steps/backends.py +++ b/exca/steps/backends.py @@ -18,6 +18,7 @@ import pickle import shutil import typing as tp +from concurrent import futures from pathlib import Path import pydantic @@ -76,13 +77,15 @@ def from_step(cls, folder: Path, step: "Step", value: tp.Any) -> "StepPaths": """Create StepPaths from a step and input value. step_uid is computed from _chain_hash() giving nested folder structure. - item_uid is computed from the input value (or sentinel for generators). + item_uid is computed via ``step.item_uid(value)`` — subclasses + can override ``item_uid`` for custom cache keys (affects both + ``forward()`` and ``map()``). Generators use a sentinel. """ step_uid = step._chain_hash() if isinstance(value, NoValue): item_uid = _NOINPUT_UID else: - item_uid = exca.ConfDict(value=value).to_uid() + item_uid = step.item_uid(value) return cls(base_folder=folder, step_uid=step_uid, item_uid=item_uid) @property @@ -421,6 +424,20 @@ def _submit(self, wrapper: _CachingCall, *args: tp.Any) -> tp.Any: wrapper(*args) return _InlineJob() + def _submit_map( + self, + fn: tp.Callable[..., None], + chunks: list[list[tp.Any]], + logs_folder: str, + ) -> None: + """Execute *fn(chunk)* for each chunk. + + Called by :meth:`Step.map` to process batches of items. + Override in subclasses for parallel / remote execution. + """ + for chunk in chunks: + fn(chunk) + class _InlineJob: """Dummy job for inline execution.""" @@ -468,6 +485,33 @@ def _submit(self, wrapper: _CachingCall, *args: tp.Any) -> tp.Any: return job + def _submit_map( + self, + fn: tp.Callable[..., None], + chunks: list[list[tp.Any]], + logs_folder: str, + ) -> None: + """Submit each chunk as a separate submitit job.""" + executor = self._EXECUTOR_CLS(folder=logs_folder) + + submitit_fields = set(type(self).model_fields) - set(Backend.model_fields) + params = { + k: getattr(self, k) for k in submitit_fields if getattr(self, k) is not None + } + if "job_name" in params: + params["name"] = params.pop("job_name") + executor.update_parameters(**params) + + jobs: list[tp.Any] = [] + with submitit.helpers.clean_env(): + with executor.batch(): + for chunk in chunks: + jobs.append(executor.submit(fn, chunk)) + + logger.info("Submitted %d map jobs (e.g. %s)", len(jobs), jobs[0].job_id) + for job in jobs: + job.result() + class LocalProcess(_SubmititBackend): """Subprocess execution + caching.""" @@ -498,3 +542,51 @@ class Auto(Slurm): """Auto-detect executor (local or Slurm).""" _EXECUTOR_CLS: tp.ClassVar[tp.Type[submitit.Executor]] = submitit.AutoExecutor + + +class ThreadPool(Backend): + """Thread-pool execution + caching. + + Uses ``concurrent.futures.ThreadPoolExecutor`` for :meth:`Step.map`. + For :meth:`Step.forward`, behaves like ``Cached`` (inline execution). + """ + + max_workers: int | None = None + + def _submit_map( + self, + fn: tp.Callable[..., None], + chunks: list[list[tp.Any]], + logs_folder: str, + ) -> None: + max_workers = self.max_workers + if max_workers is not None: + max_workers = min(len(chunks), max_workers) + with futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + futs = [executor.submit(fn, chunk) for chunk in chunks] + for f in futures.as_completed(futs): + f.result() + + +class ProcessPool(Backend): + """Process-pool execution + caching. + + Uses ``concurrent.futures.ProcessPoolExecutor`` for :meth:`Step.map`. + For :meth:`Step.forward`, behaves like ``Cached`` (inline execution). + """ + + max_workers: int | None = None + + def _submit_map( + self, + fn: tp.Callable[..., None], + chunks: list[list[tp.Any]], + logs_folder: str, + ) -> None: + max_workers = self.max_workers + if max_workers is not None: + max_workers = min(len(chunks), max_workers) + with futures.ProcessPoolExecutor(max_workers=max_workers) as executor: + futs = [executor.submit(fn, chunk) for chunk in chunks] + for f in futures.as_completed(futs): + f.result() diff --git a/exca/steps/base.py b/exca/steps/base.py index 56bc3fbd..fa4c2674 100644 --- a/exca/steps/base.py +++ b/exca/steps/base.py @@ -5,10 +5,13 @@ # LICENSE file in the root directory of this source tree. """ -Core step classes. +Core step classes and map/batch processing. Step handles computation logic, backends handles execution + caching. Backends holds a reference to its owning Step for cache key computation. + +Map support: ``step.map(Items([...]))`` processes multiple items with +per-item caching, delegated to the step's backend for parallelism. """ from __future__ import annotations @@ -16,12 +19,14 @@ import collections import inspect import logging +import math import typing as tp from pathlib import Path import pydantic import exca +from exca import cachedict as cachedict_mod from exca import utils from . import backends @@ -30,6 +35,104 @@ logger = logging.getLogger(__name__) +# ============================================================================= +# Items wrapper (public API for step.map) +# ============================================================================= + + +class Items: + """Batch of items for :meth:`Step.map`. + + Accepts any iterable, including generators. Generators are consumed + once during ``map()``; only uncached items are kept in memory. + + Parameters + ---------- + items: iterable + Items to process. + max_jobs: optional int + Maximum number of parallel jobs / chunks. ``None`` = no limit. + min_items_per_job: int + Minimum items per chunk. + """ + + def __init__( + self, + items: tp.Iterable[tp.Any], + *, + max_jobs: int | None = None, + min_items_per_job: int = 1, + ): + self._items = items + self.max_jobs = max_jobs + self.min_items_per_job = min_items_per_job + + def __iter__(self) -> tp.Iterator[tp.Any]: + return iter(self._items) + + def __repr__(self) -> str: + try: + n = len(self._items) # type: ignore[arg-type] + return f"Items({n} items, max_jobs={self.max_jobs})" + except TypeError: + return f"Items(max_jobs={self.max_jobs})" + + +# ============================================================================= +# Map internals +# ============================================================================= + + +def _to_chunks( + items: list[tp.Any], + *, + max_chunks: int | None = None, + min_items_per_chunk: int = 1, +) -> list[list[tp.Any]]: + """Split items into balanced chunks for parallel processing.""" + n = len(items) + if n == 0: + return [] + splits = min( + n if max_chunks is None else max_chunks, + math.ceil(n / min_items_per_chunk), + ) + splits = max(1, splits) + per_chunk = math.ceil(n / splits) + return [items[k * per_chunk : (k + 1) * per_chunk] for k in range(splits)] + + +class _ChunkProcessor: + """Picklable callable that processes a chunk of ``(uid, item)`` pairs. + + Creates a fresh ``CacheDict`` in the worker (necessary for remote + processes) and writes one result per item. Used by + ``Backend._submit_map`` and its overrides. + """ + + def __init__( + self, + step: "Step", + cache_folder: Path, + cache_type: str | None, + permissions: int | None, + ) -> None: + self.step = step.model_copy(deep=True) + self.cache_folder = cache_folder + self.cache_type = cache_type + self.permissions = permissions + + def __call__(self, chunk: list[tuple[str, tp.Any]]) -> None: + cd: cachedict_mod.CacheDict[tp.Any] = cachedict_mod.CacheDict( + folder=self.cache_folder, + cache_type=self.cache_type, + permissions=self.permissions, + ) + with cd.writer() as writer: + for uid, item in chunk: + writer[uid] = self.step._map_compute(item) + + def _set_mode_recursive(steps: tp.Iterable["Step"], mode: str) -> None: """Recursively set mode on steps and all nested chain steps.""" for step in steps: @@ -179,6 +282,126 @@ def forward(self, value: tp.Any = NoValue()) -> tp.Any: return result + def _map_compute(self, item: tp.Any) -> tp.Any: + """Compute one item for :meth:`map`. Override in Chain.""" + return self._forward(item) + + def map(self, items: tp.Any) -> tp.Iterator[tp.Any]: + """Process multiple items with per-item caching. + + Iterates through *items* **once** (generator-friendly): + + * Cached items: uid recorded, value discarded (memory-efficient). + * Uncached items: ``(uid, value)`` kept for processing. + + Processing of uncached items is delegated to the backend's + ``_submit_map`` method (sequential, thread-pool, Slurm …). + + Parameters + ---------- + items: Items + Batch of items wrapped in ``Items(iterable, ...)``. + Generators supported: only uncached items kept in memory. + + Returns + ------- + Iterator of results in the same order as input items. + """ + if not isinstance(items, Items): + raise TypeError( + f"map() requires an Items instance, got {type(items).__name__}. " + "Use step.map(Items([...]))" + ) + return self._map_iter(items) + + def _map_iter(self, items: Items) -> tp.Iterator[tp.Any]: + """Generator that implements :meth:`map` (separated for eager validation).""" + # --- No caching: pure streaming --- + if self.infra is None or self.infra.folder is None: + for item in items: + yield self._map_compute(item) + return + + # --- With caching --- + # Compute cache folder (with_input sets _previous = Input(NoValue) + # which is enough for step_uid / cache_folder — neither depends on input). + configured = self.with_input() + assert configured.infra is not None + cache_folder = configured.infra.paths.cache_folder + configured.infra._check_configs(write=True) + + cd: cachedict_mod.CacheDict[tp.Any] = cachedict_mod.CacheDict( + folder=cache_folder, + keep_in_ram=self.infra.keep_in_ram, + cache_type=self.infra.cache_type, + permissions=self.infra.permissions, + ) + + mode = self.infra.mode + + # Single pass through items (generator-safe). + uid_order: list[str] = [] + missing: list[tuple[str, tp.Any]] = [] + seen: set[str] = set() + + with cd.frozen_cache_folder(): + for item in items: + uid = self.item_uid(item) + uid_order.append(uid) + if uid in seen: + continue + seen.add(uid) + if mode in ("force", "force-forward"): + missing.append((uid, item)) + elif uid not in cd: + missing.append((uid, item)) + + # Mode: read-only + if mode == "read-only" and missing: + raise RuntimeError( + f"mode='read-only' but {len(missing)} items are not cached" + ) + + # Mode: force — clear existing entries before rewriting + if mode in ("force", "force-forward"): + for uid, _ in missing: + if uid in cd: + del cd[uid] + + # Process missing items via the backend + if missing: + logger.info( + "Processing %d/%d items for %s", + len(missing), + len(uid_order), + type(self).__name__, + ) + + chunks = _to_chunks( + missing, + max_chunks=items.max_jobs, + min_items_per_chunk=items.min_items_per_job, + ) + processor = _ChunkProcessor( + step=self, + cache_folder=cache_folder, + cache_type=self.infra.cache_type, + permissions=self.infra.permissions, + ) + self.infra._submit_map( + processor, chunks, logs_folder=configured.infra.paths.logs_folder + ) + + logger.info("Finished processing items for %s", type(self).__name__) + + # Reset force modes (consistent with forward()) + if mode in ("force", "force-forward"): + self.infra.mode = "cached" + + # Yield results in original order from CacheDict + for uid in uid_order: + yield cd[uid] + # ========================================================================= # Cache key computation # ========================================================================= @@ -196,6 +419,15 @@ def _chain_hash(self) -> str: opts = {"exclude_defaults": True, "uid": True} return "/".join(exca.ConfDict.from_model(s, **opts).to_uid() for s in steps) + def item_uid(self, value: tp.Any) -> str: + """Derive cache key from an input value. + + Override in subclass for custom cache keys. Used by both + ``forward()`` and ``map()`` — they share the same cache entries. + The default uses ``ConfDict(value=...).to_uid()``. + """ + return exca.ConfDict(value=value).to_uid() + # ========================================================================= # Cache operations (backend auto-configures generators, errors for transformers) # ========================================================================= @@ -365,6 +597,21 @@ def forward(self, value: tp.Any = NoValue()) -> tp.Any: return result + def _map_compute(self, item: tp.Any) -> tp.Any: + """Compute one item for :meth:`map`. + + Uses ``with_input(item)`` + ``_forward()`` to handle chain + initialisation and folder propagation while skipping chain-level + ``Backend.run()`` caching (whose key does not vary per item). + Internal steps still use their own backends for intermediate caching. + """ + configured = self.with_input(item) + return configured._forward() + + def item_uid(self, value: tp.Any) -> str: + """Delegate to first step's ``item_uid`` (the step receiving raw input).""" + return self._step_sequence()[0].item_uid(value) + def _aligned_step(self) -> list[Step]: # Flatten to contained steps - chain itself is not in the UID. # This means chain and its last step share the same step_uid (cache folder). diff --git a/exca/steps/conftest.py b/exca/steps/conftest.py index a20f481a..084310e8 100644 --- a/exca/steps/conftest.py +++ b/exca/steps/conftest.py @@ -88,6 +88,17 @@ def _forward(self, value: float = 0) -> float: # ============================================================================= +class CountMult(Step): + """Multiplies input; counts calls (for testing partial cache).""" + + coeff: float = 2.0 + _call_count: int = 0 # class-level, shared across instances + + def _forward(self, value: float) -> float: + type(self)._call_count += 1 + return value * self.coeff + + class RandomGenerator(Step): """Generates a random value - useful to verify caching. From 7c8b37ca24aabf88e3fdc995e952cb0158f5043b Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Mon, 9 Feb 2026 15:54:47 +0100 Subject: [PATCH 2/2] add --- exca/steps/test_map_step.py | 294 ++++++++++++++++++++++++++++++++++++ 1 file changed, 294 insertions(+) create mode 100644 exca/steps/test_map_step.py diff --git a/exca/steps/test_map_step.py b/exca/steps/test_map_step.py new file mode 100644 index 00000000..3a19988f --- /dev/null +++ b/exca/steps/test_map_step.py @@ -0,0 +1,294 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Tests for map/batch processing.""" + +import typing as tp +from pathlib import Path + +import pytest + +from . import conftest +from .base import Chain, Items, Step, _to_chunks + +# ============================================================================= +# Items and chunking utilities +# ============================================================================= + + +def test_items() -> None: + """Items wraps lists and generators with batch parameters.""" + items = Items([1, 2, 3]) + assert list(items) == [1, 2, 3] + assert items.max_jobs is None + assert "3 items" in repr(items) + + # Generator (exhausted after one iteration) + items = Items((x for x in [10, 20, 30]), max_jobs=2) + assert items.max_jobs == 2 + assert "max_jobs=2" in repr(items) + assert list(items) == [10, 20, 30] + assert list(items) == [] + + +@pytest.mark.parametrize( + "items,kwargs,expected", + [ + ([], {"max_chunks": 3}, []), + ([1, 2, 3], {"max_chunks": 2}, [[1, 2], [3]]), + ([1, 2, 3, 4], {"max_chunks": 2}, [[1, 2], [3, 4]]), + ([1, 2, 3], {"max_chunks": 10}, [[1], [2], [3]]), + ([1, 2, 3], {}, [[1], [2], [3]]), + ([1, 2, 3, 4, 5], {"min_items_per_chunk": 3}, [[1, 2, 3], [4, 5]]), + ], +) +def test_to_chunks(items: list, kwargs: dict, expected: list) -> None: # type: ignore + assert _to_chunks(items, **kwargs) == expected + + +# ============================================================================= +# item_uid +# ============================================================================= + + +def test_item_uid() -> None: + """Default item_uid is deterministic; subclass can override.""" + step = conftest.Mult() + assert step.item_uid(5) == step.item_uid(5) + assert step.item_uid(5) != step.item_uid(6) + + class ModStep(Step): + def item_uid(self, value: tp.Any) -> str: + return f"mod-{int(value) % 10}" + + def _forward(self, value: float) -> float: + return value * 2 + + mod = ModStep() + assert mod.item_uid(1) == mod.item_uid(11) == "mod-1" + + +# ============================================================================= +# step.map() — no infra (streaming) +# ============================================================================= + + +def test_map_no_infra() -> None: + """No infra: list, generator, empty all work; non-Items rejected.""" + step = conftest.Mult(coeff=3.0) + assert list(step.map(Items([1.0, 2.0, 3.0]))) == [3.0, 6.0, 9.0] + assert list(step.map(Items(x for x in [1.0, 2.0]))) == [3.0, 6.0] + assert list(step.map(Items([]))) == [] + with pytest.raises(TypeError, match="Items"): + step.map([1, 2, 3]) # type: ignore + + +# ============================================================================= +# step.map() — caching +# ============================================================================= + + +def test_map_caching(tmp_path: Path) -> None: + """Results are cached; generator and list inputs hit same cache.""" + infra: tp.Any = {"backend": "Cached", "folder": tmp_path} + step = conftest.Add(randomize=True, infra=infra) + + r1 = list(step.map(Items(x for x in [1.0, 2.0, 3.0]))) + r2 = list(step.map(Items([1.0, 2.0, 3.0]))) + assert r1 == r2 + + +def test_map_partial_cache_and_dedup(tmp_path: Path) -> None: + """Only missing items computed; duplicates deduplicated; generators work.""" + infra: tp.Any = {"backend": "Cached", "folder": tmp_path} + step = conftest.CountMult(coeff=2.0, infra=infra) + + # Initial batch: all 3 computed + conftest.CountMult._call_count = 0 + assert list(step.map(Items([1.0, 2.0, 3.0]))) == [2.0, 4.0, 6.0] + assert conftest.CountMult._call_count == 3 + + # Overlapping batch: only 4.0 and 5.0 need computing + conftest.CountMult._call_count = 0 + assert list(step.map(Items([2.0, 3.0, 4.0, 5.0]))) == [4.0, 6.0, 8.0, 10.0] + assert conftest.CountMult._call_count == 2 + + # Deduplication: duplicates resolved from cache + conftest.CountMult._call_count = 0 + assert list(step.map(Items([1.0, 2.0, 1.0, 2.0]))) == [2.0, 4.0, 2.0, 4.0] + assert conftest.CountMult._call_count == 0 + + # Generator partial cache: only uncached item computed + conftest.CountMult._call_count = 0 + assert list(step.map(Items(x for x in [1.0, 2.0, 6.0]))) == [2.0, 4.0, 12.0] + assert conftest.CountMult._call_count == 1 + + +def test_map_shares_cache_with_forward(tmp_path: Path) -> None: + """map() and forward() share cache via item_uid.""" + infra: tp.Any = {"backend": "Cached", "folder": tmp_path} + step = conftest.Add(randomize=True, infra=infra) + + result = step.forward(5.0) + assert list(step.map(Items([5.0]))) == [result] + + +# ============================================================================= +# Modes +# ============================================================================= + + +def test_map_force_mode(tmp_path: Path) -> None: + """Force mode recomputes; mode resets to cached after.""" + infra: tp.Any = {"backend": "Cached", "folder": tmp_path} + step = conftest.Add(randomize=True, infra=infra) + + r1 = list(step.map(Items([1.0, 2.0]))) + step.infra.mode = "force" # type: ignore + r2 = list(step.map(Items([1.0, 2.0]))) + assert r1 != r2 + assert step.infra.mode == "cached" # type: ignore + assert list(step.map(Items([1.0, 2.0]))) == r2 + + +def test_map_readonly_mode(tmp_path: Path) -> None: + """read-only fails without cache, succeeds with cache.""" + infra: tp.Any = {"backend": "Cached", "folder": tmp_path} + step = conftest.Mult(coeff=2.0, infra=infra) + + step.infra.mode = "read-only" # type: ignore + with pytest.raises(RuntimeError, match="read-only"): + list(step.map(Items([1.0]))) + + # Populate cache, then read-only succeeds + step.infra.mode = "cached" # type: ignore + list(step.map(Items([1.0, 2.0]))) + step.infra.mode = "read-only" # type: ignore + assert list(step.map(Items([1.0, 2.0]))) == [2.0, 4.0] + + +# ============================================================================= +# Custom item_uid with caching +# ============================================================================= + + +def test_map_custom_item_uid(tmp_path: Path) -> None: + """Custom item_uid controls cache keys — same uid shares result.""" + + class ModStep(Step): + def item_uid(self, value: tp.Any) -> str: + return str(int(value) % 10) + + def _forward(self, value: float) -> float: + return value * 2 + + infra: tp.Any = {"backend": "Cached", "folder": tmp_path} + step = ModStep(infra=infra) + + # 1, 11, 21 all map to uid "1" — only computed once + results = list(step.map(Items([1.0, 11.0, 21.0]))) + assert results == [2.0, 2.0, 2.0] + + +# ============================================================================= +# Backend-driven parallelism +# ============================================================================= + + +@pytest.mark.parametrize("backend", ("ThreadPool", "LocalProcess", "SubmititDebug")) +def test_map_backend(tmp_path: Path, backend: str) -> None: + """All backends compute correct results and cache them.""" + infra: tp.Any = {"backend": backend, "folder": tmp_path} + step = conftest.Mult(coeff=2.0, infra=infra) + + results = list(step.map(Items([1.0, 2.0, 3.0], max_jobs=2))) + assert results == [2.0, 4.0, 6.0] + assert list(step.map(Items([1.0, 2.0, 3.0]))) == results + + +def test_map_cross_backend_cache(tmp_path: Path) -> None: + """ThreadPool and Cached backends share the same cache.""" + infra: tp.Any = {"backend": "ThreadPool", "folder": tmp_path} + step = conftest.Add(randomize=True, infra=infra) + r1 = list(step.map(Items([1.0, 2.0, 3.0], max_jobs=2))) + + infra2: tp.Any = {"backend": "Cached", "folder": tmp_path} + step2 = conftest.Add(randomize=True, infra=infra2) + assert list(step2.map(Items([1.0, 2.0, 3.0]))) == r1 + + +# ============================================================================= +# Chain support +# ============================================================================= + + +@pytest.mark.parametrize("use_infra", (True, False), ids=("cached", "no-infra")) +def test_map_chain(tmp_path: Path, use_infra: bool) -> None: + """Chain.map() processes items through full chain; caches when infra set.""" + infra: tp.Any = {"backend": "Cached", "folder": tmp_path} if use_infra else None + chain = Chain( + steps=[conftest.Add(value=1), conftest.Mult(coeff=2)], + infra=infra, + ) + results = list(chain.map(Items([1.0, 2.0, 3.0]))) + assert results == [4.0, 6.0, 8.0] # (x + 1) * 2 + if use_infra: + assert list(chain.map(Items([1.0, 2.0, 3.0]))) == results + + +def test_map_chain_intermediate_cache(tmp_path: Path) -> None: + """Chain with intermediate caching: inner step results are cached.""" + infra: tp.Any = {"backend": "Cached", "folder": tmp_path} + chain = Chain( + steps=[conftest.Add(randomize=True, infra=infra), conftest.Mult(coeff=10)], + infra=infra, + ) + r1 = list(chain.map(Items([1.0, 2.0]))) + assert list(chain.map(Items([1.0, 2.0]))) == r1 + + +def test_map_chain_force_mode(tmp_path: Path) -> None: + """Chain with force mode recomputes; mode resets.""" + infra: tp.Any = {"backend": "Cached", "folder": tmp_path} + chain = Chain( + steps=[conftest.Add(randomize=True), conftest.Mult(coeff=1)], + infra=infra, + ) + r1 = list(chain.map(Items([1.0, 2.0]))) + assert chain.infra is not None + chain.infra.mode = "force" + r2 = list(chain.map(Items([1.0, 2.0]))) + assert r1 != r2 + assert chain.infra.mode == "cached" + + +def test_map_chain_custom_item_uid(tmp_path: Path) -> None: + """Chain uses first step's item_uid for cache keys.""" + + class ModAdd(Step): + value: float = 1.0 + + def item_uid(self, v: tp.Any) -> str: + return str(int(v) % 10) + + def _forward(self, v: float) -> float: + return v + self.value + + infra: tp.Any = {"backend": "Cached", "folder": tmp_path} + chain = Chain(steps=[ModAdd(), conftest.Mult(coeff=2)], infra=infra) + results = list(chain.map(Items([1.0, 11.0]))) + assert results == [4.0, 4.0] # same uid → same cached result + + +@pytest.mark.parametrize("backend", ("LocalProcess", "SubmititDebug")) +def test_map_chain_submitit(tmp_path: Path, backend: str) -> None: + """Chain.map() works with subprocess/submitit backends.""" + infra: tp.Any = {"backend": backend, "folder": tmp_path} + chain = Chain( + steps=[conftest.Add(value=1), conftest.Mult(coeff=2)], + infra=infra, + ) + assert list(chain.map(Items([1.0, 2.0, 3.0], max_jobs=2))) == [4.0, 6.0, 8.0]