From 0f2423ce0307735a3f23e61bdd201ffbf304ce05 Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Wed, 26 Aug 2026 11:56:01 +0200 Subject: [PATCH 01/11] Add a Fit API --- CHANGELOG.md | 4 + docs/internal/steps/caching.md | 6 + docs/steps/items.md | 40 ++++ docs/steps/reference.rst | 11 + exca/cachedict/handlers.py | 9 +- exca/cachedict/test_dumpcontext.py | 13 ++ exca/steps/__init__.py | 2 + exca/steps/backends.py | 42 +++- exca/steps/base.py | 24 +- exca/steps/fit.py | 211 +++++++++++++++++ exca/steps/helpers.py | 60 ++--- exca/steps/items.py | 28 ++- exca/steps/patterns.py | 7 +- exca/steps/test_fit.py | 351 +++++++++++++++++++++++++++++ exca/steps/test_helpers.py | 39 +++- pyproject.toml | 3 +- 16 files changed, 793 insertions(+), 57 deletions(-) create mode 100644 exca/steps/fit.py create mode 100644 exca/steps/test_fit.py diff --git a/CHANGELOG.md b/CHANGELOG.md index de0202d4..d33e0d5f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +- `steps`: `Fit` primitive — fit one artifact over a cohort of items, then transform each item with it (N->1->N). The cohort is declared by wrapping the items in `FitCohort`; its identity (the items fingerprint, or the name the config already carries) is written to the `cohort` field before the run, scoping the artifact and the outputs. A step that ran is frozen, so fitting another cohort takes a fresh config. +- `steps`: `Parallel` now dispatches its variants through `Step._dispatch` (grouping their submission through the new `Backend._grouped`), so a variant can be a `Fit` or resolve through the `build()` convention — the latter used to run and cache as its unresolved wrapper. +- `cachedict`: `Auto` sequences all load back as lists, and the contents of a tuple now resolve — a nested tuple used to load back as raw dump info, permanently so on the pickle path. + ## 0.5.29 - 26-07-28 - `ConfDict`: added `ConfDict` operations `ConfDict.ops.DELETE`, `ConfDict.ops.BEFORE`, and `ConfDict.ops.AFTER`; `ConfDict.ops.REPLACE` replaces `ConfDict.OVERRIDE`. [#310] diff --git a/docs/internal/steps/caching.md b/docs/internal/steps/caching.md index bc1d79b7..195170c7 100644 --- a/docs/internal/steps/caching.md +++ b/docs/internal/steps/caching.md @@ -94,6 +94,12 @@ whole batch set, so a sweep (many step variants dispatched together) packs into one submitit array (one pool for pool backends). `_mark_recomputed` records force/retry batches per attempt. +`Backend._grouped()` is how a sweep gets there: inside it, `_run` only +prepares and queues its batch, and the group is claimed and executed once +on exit. `Parallel` uses it so its variants still go through the regular +`Step._dispatch` (step resolution, `Fit` cohort resolution) instead of +driving the backend primitives itself. + The session locks `inflight.db` only — direct user calls to `LookupHandle.clear_cache()` race against in-flight workers. Results are not round-tripped through the job pickle (would be wasteful under diff --git a/docs/steps/items.md b/docs/steps/items.md index f6629af7..d503e49c 100644 --- a/docs/steps/items.md +++ b/docs/steps/items.md @@ -133,6 +133,46 @@ at a time (never the full set in memory, never round-tripped through the job pickle). Execution order within a batch is non-deterministic; output order matches input order. +## Fitting on the items + +A `Fit` step derives one artifact from many items, then transforms +each item with it — normalization statistics, a PCA basis, a trained +model. Only a batch wrapped in a `FitCohort` is fitted on; any other +call transforms with what is already fitted: + +```python +class Normalize(steps.Fit): + def _fit(self, values): # the cohort + stacked = np.concatenate(list(values)) + return stacked.mean(0), stacked.std(0) + + def _run(self, value): # one item + mean, std = self.fitted + return (value - mean) / std + + +norm = Normalize(infra={"backend": "Cached", "folder": cache}) +for value in norm.run_many(steps.FitCohort(train_paths)): + train(value) # fitted on this cohort, then transformed +for value in norm.run_many(test_paths): + evaluate(value) # same artifact, novel items +``` + +`_fit` receives the cohort as an iterable it can stream, and iterate +again (one upstream read per pass). + +The cohort's identity — the fingerprint of its items, or the name the +config already carries (`Normalize(cohort="train")`) — is written to +the `cohort` field before anything runs, so the artifact and every +downstream cache are scoped to it. A named cohort is recoverable from +the config alone, for a pipeline that never presents the items it was +fitted on. A step that ran is frozen, so fitting another cohort takes +a fresh config (`clone({"cohort": None})`). + +The fit runs where the step is dispatched from, ahead of any split, +and is cached under `infra`. A `Fit` under a backend that shards the +cohort raises rather than fitting on a shard. + ## What's stable Pinned by tests — safe to rely on: diff --git a/docs/steps/reference.rst b/docs/steps/reference.rst index 57d471fe..c19ab46d 100644 --- a/docs/steps/reference.rst +++ b/docs/steps/reference.rst @@ -90,3 +90,14 @@ Helpers .. autoclass:: exca.steps.helpers.Func :members: + + +Fitting over items +------------------ + +.. autoclass:: exca.steps.Fit + :show-inheritance: + :members: fitted, COHORT_KEY, ARTIFACT_CACHE_TYPE + :private-members: _fit + +.. autoclass:: exca.steps.FitCohort diff --git a/exca/cachedict/handlers.py b/exca/cachedict/handlers.py index 4020972b..a3542bb2 100644 --- a/exca/cachedict/handlers.py +++ b/exca/cachedict/handlers.py @@ -383,9 +383,8 @@ def _dump_value(cls, ctx: DumpContext, val: tp.Any, key: str) -> tp.Any: f"Found #type={val['#type']!r} which is not a registered handler." ) return {k: cls._dump_value(ctx, v, f"{key}[{k}]") for k, v in val.items()} - if isinstance(val, (list, tuple)): - items = [cls._dump_value(ctx, v, f"{key}[{i}]") for i, v in enumerate(val)] - return tuple(items) if isinstance(val, tuple) else items + if isinstance(val, (list, tuple)): # as a list: `_load_value` walks lists only + return [cls._dump_value(ctx, v, f"{key}[{i}]") for i, v in enumerate(val)] handler = DumpContext._find_handler(type(val)) if handler is not None or hasattr(val, "__dump_info__"): ctx.key = key @@ -416,7 +415,7 @@ def _delete_value(cls, ctx: DumpContext, val: tp.Any) -> None: else: for v in val.values(): cls._delete_value(ctx, v) - elif isinstance(val, list): + elif isinstance(val, (list, tuple)): # tuple: only in legacy entries for item in val: cls._delete_value(ctx, item) @@ -426,7 +425,7 @@ def _load_value(cls, ctx: DumpContext, val: tp.Any) -> tp.Any: if "#type" in val: return ctx.load(val) return {k: cls._load_value(ctx, v) for k, v in val.items()} - if isinstance(val, list): + if isinstance(val, (list, tuple)): # tuple: only in legacy entries return [cls._load_value(ctx, item) for item in val] return val diff --git a/exca/cachedict/test_dumpcontext.py b/exca/cachedict/test_dumpcontext.py index ad9642a9..ee1d64e7 100644 --- a/exca/cachedict/test_dumpcontext.py +++ b/exca/cachedict/test_dumpcontext.py @@ -300,6 +300,19 @@ def tracking_dump( np.testing.assert_array_almost_equal(loaded["weights"], [[1.0, 2.0], [3.0, 4.0]]) +def test_auto_tuple_loads_as_list(tmp_path: Path) -> None: + ctx = DumpContext(tmp_path, key="test") + with ctx: + info = ctx.dump((np.array([1.0, 2.0]), 3), cache_type="Auto") + loaded = ctx.load(info) + assert isinstance(loaded, list), "json storage cannot hold a tuple either" + np.testing.assert_array_almost_equal(loaded[0], [1.0, 2.0]) + assert loaded[1] == 3 + + legacy = ctx.load({**info, "content": tuple(info["content"])}) # tuple: old entry + np.testing.assert_array_almost_equal(legacy[0], [1.0, 2.0]) + + class _Opaque: """Module-level class so it's picklable (local classes are not).""" diff --git a/exca/steps/__init__.py b/exca/steps/__init__.py index 958aa85e..91bdd2a4 100644 --- a/exca/steps/__init__.py +++ b/exca/steps/__init__.py @@ -29,3 +29,5 @@ from . import helpers as helpers from .base import Chain as Chain from .base import Step as Step +from .fit import Fit as Fit +from .fit import FitCohort as FitCohort diff --git a/exca/steps/backends.py b/exca/steps/backends.py index a308be6c..ba16a9e0 100644 --- a/exca/steps/backends.py +++ b/exca/steps/backends.py @@ -323,9 +323,9 @@ def shuffled(self) -> ComputeBatch: def cached_items(self) -> StepItems: """Lazy cache-backed carrier; use on a top-level batch, not a chunk.""" - return items.StepItems( + return self.items._replace( source=self.cache_dict, - uids=self.items.uids, + pending=(), upstream=self.info.upstream, mode=self.info.mode, ) @@ -436,6 +436,7 @@ def _exclude_from_cls_uid(cls) -> list[str]: # Force/retry: recompute each (step_folder, uid) at most once per lifetime _recomputed: set[tuple[Path, str]] = pydantic.PrivateAttr(default_factory=set) _checked_configs: set[Path] = pydantic.PrivateAttr(default_factory=set) + _group: list[ComputeBatch] | None = pydantic.PrivateAttr(default=None) def __getstate__(self) -> dict[str, tp.Any]: recomputed = self._recomputed @@ -578,9 +579,39 @@ def _clear_caches( ereg.clear(uids) self._checked_configs.discard(paths.step_folder) + @contextlib.contextmanager + def _grouped(self) -> tp.Iterator[None]: + """Hold back ``_run`` calls, then claim and execute them as one submission. + + Carriers returned while grouping only hold results once the group ran. + Nesting joins the enclosing group. + """ + if self._group is not None: # nested sweep: the outer group submits it + yield + return + group: list[ComputeBatch] = [] + self._group = group + try: + yield + finally: + self._group = None + with self._claim(group) as claimed: + if claimed.ready: + self._execute(claimed.ready) + + def _defer(self, cbatch: ComputeBatch) -> items.StepItems | None: + """Inside ``_grouped``, queue *cbatch* and return its pending carrier.""" + if self._group is None: + return None + self._group.append(cbatch) + return cbatch.cached_items() + def _run(self, step: Step, batch: items.StepItems) -> items.StepItems: """Execute *step* for uncached items, caching per uid.""" cbatch = self._prepare(step, batch) + deferred = self._defer(cbatch) + if deferred is not None: + return deferred with self._claim([cbatch]) as claimed: if claimed.ready: self._execute(claimed.ready) @@ -901,6 +932,9 @@ def _run(self, step: Step, batch: items.StepItems) -> items.StepItems: lazy carrier instead of blocking. """ cbatch = self._prepare(step, batch) + deferred = self._defer(cbatch) + if deferred is not None: + return deferred claimed = self._claim([cbatch]) transferred = False try: @@ -916,9 +950,9 @@ def _run(self, step: Step, batch: items.StepItems) -> items.StepItems: } ctx = _PoolContext(cbatch.cache_dict, cbatch.paths.step_uid, pool, claimed) transferred = True # _PoolContext closes `claimed`, not the finally - return items.StepItems( + return cbatch.items._replace( source=_PoolSource(uid_to_future, ctx), - uids=cbatch.items.uids, + pending=(), upstream=cbatch.info.upstream, mode=cbatch.info.mode, ) diff --git a/exca/steps/base.py b/exca/steps/base.py index 6c3f7cd5..6e519589 100644 --- a/exca/steps/base.py +++ b/exca/steps/base.py @@ -22,6 +22,9 @@ from . import backends, identity, items, utils +if tp.TYPE_CHECKING: + from .fit import FitCohort + logger = logging.getLogger(__name__) @@ -376,33 +379,44 @@ def run(self, value: tp.Any = identity.NoValue()) -> tp.Any: """ return next(iter(self.run_many([value]))) - def run_many(self, values: tp.Iterable[tp.Any]) -> items.StepItems: + def run_many(self, values: tp.Iterable[tp.Any] | FitCohort) -> items.StepItems: """Execute the step over many inputs, one cache entry per input. Parameters ---------- values: - Inputs to run; one result is produced per input, in order. + Inputs to run; one result is produced per input, in order. Wrap them in a + :class:`~exca.steps.FitCohort` to let a :class:`~exca.steps.Fit` fit on them. Returns ------- StepItems Iterator yielding one result per input, in input order. """ + from .fit import FitCohort, _stamp_cohorts # local import: fit builds on Step + built = utils.resolved_step(self) if built is not self: return built.run_many(values) + cohort: FitCohort | None = None + if isinstance(values, FitCohort): + cohort, values = values, values.items values = list(values) # eager: uid computation needs all values upfront uids = [identity.materialize_uid(self, v) for v in values] - warm = self._warm_items(uids) - if warm is not None: - return warm # extra-fast path -> avoid StepItems + _dispatch overhead + if cohort is None: # a declared cohort must reach the steps + warm = self._warm_items(uids) + if warm is not None: + return warm # extra-fast path -> avoid StepItems + _dispatch overhead + else: + cohort.uids = uids + _stamp_cohorts(self, cohort) boundary = items.StepItems( source=dict(zip(uids, values)), uids=uids, + cohort=cohort, ) return boundary.apply_step(self) diff --git a/exca/steps/fit.py b/exca/steps/fit.py new file mode 100644 index 00000000..262c1727 --- /dev/null +++ b/exca/steps/fit.py @@ -0,0 +1,211 @@ +# 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. + +from __future__ import annotations + +import hashlib +import typing as tp + +import pydantic + +import exca + +from . import backends, items +from .base import Step + +CohortKey = tp.Literal["set", "multiset", "sequence"] + + +class FitCohort: + """The items a :class:`Fit` fits on, passed to ``run_many`` in their place. + + .. warning:: Experimental -- API may change. + + Parameters + ---------- + items + Items to fit on, then transform. + + Note + ---- + ``fitted_by`` lists the steps the cohort identified, to check it was not + passed in vain. + """ + + def __init__(self, items: tp.Iterable[tp.Any]) -> None: + self.items = list(items) + self.uids: list[str] = [] # the declared cohort, set by run_many + self.fitted_by: list[str] = [] + + def __repr__(self) -> str: + return f"{type(self).__name__}({len(self.items)} items)" + + def __getstate__(self) -> dict[str, tp.Any]: + return {**self.__dict__, "items": []} # values travel with the carrier + + +def _cohort_uids(uids: tp.Sequence[str], key: CohortKey) -> list[str]: + """The cohort's uids as its fit sees them, ordered so the fit is reproducible.""" + if key == "sequence": + return list(uids) + return sorted(set(uids) if key == "set" else uids) + + +def _fingerprint(uids: tp.Sequence[str], key: CohortKey) -> str: + """Identity of a cohort that was not named, as ``,``.""" + ordered = _cohort_uids(uids, key) + digest = hashlib.sha256() + for uid in ordered: + digest.update(uid.encode("utf8")) + digest.update(b"\0") + return f"{digest.hexdigest()[:8]},{len(ordered)}" + + +def _stamp_cohorts(step: Step, cohort: FitCohort) -> None: + """Name the cohort in every ``Fit`` of *step*, before anything runs. + + The name is part of the configuration, hence of the cache identity, so it must + be settled before a composite step keys its own cache. + """ + fits = exca.utils.find_models(step, Fit, include_private=False) + for path, fit in fits.items(): + uid = fit.cohort + if uid is None or fit._stamped: # a name given in the config is kept + uid = _fingerprint(cohort.uids, fit.COHORT_KEY) + if fit.cohort != uid: + try: + fit.cohort = uid + except Exception as e: + raise RuntimeError( + f"{type(fit).__name__} at {path or '.'} already ran on cohort " + f"{fit.cohort!r}: refitting in place is not supported, run " + "clone({'cohort': None}) on the new cohort" + ) from e + fit._stamped = True + cohort.fitted_by.append(f"{path or '.'}({uid})") + + +class Fit(Step): + """Fit one artifact over a cohort of items, then transform each item (N->1->N). + + .. warning:: Experimental -- API may change. + + Override :meth:`_fit` for the artifact, and :meth:`_run` for one item's output + (reading :attr:`fitted`). + + Only a batch passed as a :class:`FitCohort` is fitted on; any other batch + transforms with what is already fitted, and raises if that is nothing. The + cohort's identity -- its name, or the fingerprint of its items -- is written to + :attr:`cohort` before the run, so the artifact and every downstream cache are + scoped to it. A step that ran is frozen, so its cohort cannot be replaced. + + The fit runs where the step is dispatched from, ahead of a backend splitting the + batch, and is cached under ``infra``. The upstream is read twice (once for the + fit, once per item), so give an expensive upstream its own ``infra``. + + ``COHORT_KEY`` states whether order and repetitions define the cohort, and + ``ARTIFACT_CACHE_TYPE`` the artifact's cache format (``CACHE_TYPE`` stays the + per-item outputs'). Prefer fitting arrays or tensors -- e.g. a state dict over a + model -- as they cache natively instead of pickling. + + Parameters + ---------- + cohort + Name of the artifact, to fit it under a name or to use it in a run that + never presents the cohort (a config-only pipeline). Left unset, the + fingerprint of the cohort's items is written here instead. + """ + + COHORT_KEY: tp.ClassVar[CohortKey] = "set" + ARTIFACT_CACHE_TYPE: tp.ClassVar[str | None] = "Auto" + + cohort: str | None = None + + _fitted: tp.Any = pydantic.PrivateAttr(None) + _fitted_for: str | None = pydantic.PrivateAttr(None) # cohort of `_fitted` + _stamped: bool = pydantic.PrivateAttr(False) # `cohort` written by a declaration + + def _fit(self, values: tp.Iterable[tp.Any]) -> tp.Any: + """The artifact for the cohort, from the values this step receives. + + *values* streams the cohort, and can be iterated more than once (e.g. one + pass per epoch) -- at the cost of re-reading the upstream each time. + """ + raise NotImplementedError + + @property + def fitted(self) -> tp.Any: + """The artifact :meth:`_fit` produced, for :meth:`_run` to transform with.""" + if self._fitted_for is None: + raise RuntimeError( + f"{type(self).__name__} is not fitted: run it on a FitCohort first" + ) + return self._fitted + + def _dispatch(self, batch: items.StepItems) -> items.StepItems: + # before super(): a split ships the artifact along with the step + if self.cohort is None or self._fitted_for != self.cohort: + self._resolve(batch) + return super()._dispatch(batch) + + def _resolve(self, batch: items.StepItems) -> None: + """Read the artifact for this step's cohort back, or fit it.""" + kind = type(self).__name__ + uid = self.cohort + if uid is None: + raise RuntimeError( + f"{kind} has no cohort to fit on or to read back: run it on a " + "FitCohort, or set its 'cohort' name" + ) + cohort = batch._cohort + mode = backends._fold_modes(batch._mode, backends._effective_mode(self)) + # cohort cleared: all cohorts of this Fit share one folder, one entry each + owner = self.model_copy(update={"infra": None, "cohort": None}) + infra = None if self.infra is None else self.infra.derive(mode=mode) + artifact = _Artifact(owner=owner, infra=infra) + upstream = tuple(batch._upstream) + handle = artifact.lookup(_upstream=upstream, _uid=uid) + status = handle.status + # same rule as `_pending_statuses`: a cached error still raises in "cached" mode + if status is None or mode == "force" or (mode == "retry" and status == "error"): + if cohort is None: + hint = "drop the force mode" if mode == "force" else "check its name" + raise RuntimeError( + f"{kind} must fit cohort {uid!r} but was handed no items to fit " + f"on: run it on a FitCohort, or {hint}" + ) + # counts, not uids: a step re-keying items has its own uid space + handed, declared_n = len(set(batch.uids)), len(set(cohort.uids)) + if handed < declared_n: + raise RuntimeError( + f"{kind} was handed {handed} of the {declared_n} items of cohort " + f"{uid!r}, too few to fit it -- an enclosing backend sharded them; " + "fit it before distributing, or move that backend onto this step" + ) + carrier = items.StepItems( + source={uid: (batch,)}, uids=[uid], upstream=upstream, mode=mode + ) + # dispatch, not lookup: goes through infra's mode and caching + self._fitted = next(iter(artifact._dispatch(carrier))) + self._fitted_for = uid + + +class _Artifact(Step): + """One cohort's artifact for a :class:`Fit`, cached as a single entry. + + Its input is the cohort's carrier, wrapped in a tuple to make it one item. + """ + + owner: Fit + + def _infer_cache_type(self) -> str | None: + return self.owner.ARTIFACT_CACHE_TYPE + + def _run(self, value: tuple[items.StepItems]) -> tp.Any: + batch = value[0] # the carrier itself: re-iterable, applies its pending steps + return self.owner._fit( + batch.select(_cohort_uids(batch.uids, self.owner.COHORT_KEY)) + ) diff --git a/exca/steps/helpers.py b/exca/steps/helpers.py index 3d68939c..ca48f656 100644 --- a/exca/steps/helpers.py +++ b/exca/steps/helpers.py @@ -11,7 +11,7 @@ import pydantic -from . import identity, items, utils +from . import fit, identity, items, utils from .base import Step @@ -127,8 +127,9 @@ class Parallel(Step): The variants run together under one shared backend, each caching under its own identity. ``run`` is for effect — read results back per variant via - ``parallel.steps[k].lookup(value)``. It has no composable output (yields - ``None`` per input), so use it standalone, not as a non-terminal chain step. + ``parallel.steps[k].lookup(value)``, which are re-parented copies of the + steps handed over. It has no composable output (yields ``None`` per input), + so use it standalone, not as a non-terminal chain step. Example:: @@ -165,19 +166,16 @@ def _unify_infra(self) -> None: base = next((i.folder for i in infras if i.folder is not None), None) if self.infra.folder is None: self.infra.folder = base - infra_dict = self.infra.model_dump() - self.steps = [ - s if s.infra is not None else s.clone({"infra": infra_dict}) - for s in self.steps - ] if base is not None: utils.propagate_folder(self, base) for step in self.steps: - if step.infra != self.infra: + if step.infra is not None and step.infra != self.infra: raise ValueError( "Parallel requires one shared backend across itself and its " f"steps; {self.infra!r} differs from {step.infra!r}" ) + # the very object: one grouped dispatch then covers every variant + self.steps = [s.model_copy(update={"infra": self.infra}) for s in self.steps] def _uid_steps(self) -> list[Step]: return [] # no identity of its own @@ -198,27 +196,31 @@ def _run_items(self, batch: items.StepItems) -> items.StepItems: f"Parallel needs a cache folder; set infra.folder (on Parallel or " f"a step), got {self.infra!r}" ) - cbatches = [] - for child in self.steps: - uids = [identity.materialize_uid(child, v) for v in batch] - child_batch = items.StepItems(source=dict(zip(uids, batch)), uids=uids) - cbatches.append(self.infra._prepare(child, child_batch)) - with self.infra._claim(cbatches) as claimed: - if claimed.ready: - self.infra._execute(claimed.ready) - return items.StepItems( - source={uid: None for uid in batch.uids}, - uids=batch.uids, - upstream=batch._upstream, - mode=batch._mode, - ) + values = list(batch) # one read, shared by every variant + with self.infra._grouped(): + for child in self.steps: + child = utils.resolved_step( + child + ) # item uids key on it, as does its cache + if child.infra is None: + raise ValueError( + f"Parallel variant {type(child).__name__} resolves to a step " + "with no infra: it would run uncached, leaving nothing to look up" + ) + if child.infra is not self.infra and child.infra == self.infra: + child = child.model_copy(update={"infra": self.infra}) # one group + uids = [identity.materialize_uid(child, v) for v in values] + child._dispatch( + items.StepItems( + source=dict(zip(uids, values)), uids=uids, cohort=batch._cohort + ) + ) + return batch._replace(source={uid: None for uid in batch.uids}, pending=()) def run(self, value: tp.Any = identity.NoValue()) -> None: self.run_many([value]) - def run_many(self, values: tp.Iterable[tp.Any]) -> list[None]: # type: ignore[override] - values = list(values) - uids = [identity.materialize_uid(self, v) for v in values] - batch = items.StepItems(source=dict(zip(uids, values)), uids=uids) - self._dispatch(batch) - return [None] * len(values) + def run_many( # type: ignore[override] + self, values: tp.Iterable[tp.Any] | fit.FitCohort + ) -> list[None]: + return list(super().run_many(values)) diff --git a/exca/steps/items.py b/exca/steps/items.py index 33fd8eef..e8681467 100644 --- a/exca/steps/items.py +++ b/exca/steps/items.py @@ -21,6 +21,7 @@ if tp.TYPE_CHECKING: from .base import Step + from .fit import FitCohort class _Source(tp.Protocol): @@ -128,6 +129,7 @@ def __init__( upstream: tp.Sequence[Step] = (), pending: tp.Sequence[Step] = (), mode: identity.ModeType = "cached", + cohort: FitCohort | None = None, ) -> None: if uids is None: if not isinstance(source, dict): @@ -140,6 +142,7 @@ def __init__( self._upstream = tuple(upstream) self._pending = tuple(pending) self._mode = mode + self._cohort = cohort def __len__(self) -> int: return len(self.uids) @@ -148,14 +151,27 @@ def apply_step(self, step: Step) -> StepItems: """Run *step* over the carrier, honoring its infra/caching (leaf or ``Chain``).""" return step._dispatch(self) + def _replace(self, **changes: tp.Any) -> StepItems: + """Copy with some parts changed; everything else (mode, cohort, ...) carries over. + + Pass ``pending=()`` when *source* becomes computed results, or its steps + would run a second time. + """ + params: dict[str, tp.Any] = { + "source": self._source, + "uids": self.uids, + "upstream": self._upstream, + "pending": self._pending, + "mode": self._mode, + "cohort": self._cohort, + } + return StepItems(**{**params, **changes}) + def _append(self, step: Step) -> StepItems: """Append a single leaf step's computation and identity.""" - return StepItems( - source=self._source, - uids=self.uids, + return self._replace( upstream=self._upstream + tuple(step._uid_steps()), pending=self._pending + (step,), - mode=self._mode, ) def select( @@ -169,11 +185,9 @@ def select( source = {uid: source[uid] for uid in dict.fromkeys(uids)} elif hasattr(source, "select"): # subset lazy sources before pickle source = source.select(uids) - return StepItems( + return self._replace( source=source, uids=uids, - upstream=self._upstream, - pending=self._pending, mode=mode if mode is not None else self._mode, ) diff --git a/exca/steps/patterns.py b/exca/steps/patterns.py index 800ee539..5aa65e03 100644 --- a/exca/steps/patterns.py +++ b/exca/steps/patterns.py @@ -228,7 +228,7 @@ def _run_items(self, batch: items.StepItems) -> items.StepItems: for uid, m in plan.items() for branch_uid, branch in m.items() } - carrier = items.StepItems( + carrier = items.StepItems( # branch uids: not the cohort's identity space source=_Parts(batch, self.take, origin), uids=uids, upstream=branch_upstream, @@ -236,9 +236,8 @@ def _run_items(self, batch: items.StepItems) -> items.StepItems: ) # one dispatch over all branches lets a backend submit them together dispatched = self._body()._dispatch(carrier) - return items.StepItems( + return batch._replace( source=_Gather(dispatched, plan, self.gather), - uids=batch.uids, + pending=(), upstream=output_upstream, - mode=batch._mode, ) diff --git a/exca/steps/test_fit.py b/exca/steps/test_fit.py new file mode 100644 index 00000000..8bf96eea --- /dev/null +++ b/exca/steps/test_fit.py @@ -0,0 +1,351 @@ +# 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. + +import collections +import pickle +import typing as tp +from pathlib import Path + +import numpy as np +import pydantic +import pytest +import sklearn.decomposition +import torch + +from .. import steps +from . import conftest + + +class Offset(steps.Fit): + """Subtracts the cohort's mean from each item, recording the cohorts it fit on.""" + + _fits: list[list[float]] = pydantic.PrivateAttr(default_factory=list) + + def _fit(self, values: tp.Iterable[float]) -> float: + vals = list(values) + self._fits.append(vals) + return sum(vals) / len(vals) + + def _run(self, value: float) -> float: + return value - self.fitted + + +class OffsetMultiset(Offset): + COHORT_KEY = "multiset" + + +class OffsetSequence(Offset): + COHORT_KEY = "sequence" + + +def test_fit_cohort_then_novel_items(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + step = Offset(infra=infra) + assert list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) == [-1.0, 0.0, 1.0] + assert list(step.run_many([10.0])) == [8.0], "novel item must use the fitted mean" + assert len(step._fits) == 1, "a second call must not refit" + + read_back = Offset(infra=infra) + assert list(read_back.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) == [-1.0, 0.0, 1.0] + assert not read_back._fits, "the artifact must be read back from the cache" + + unfitted = Offset(infra=infra) + with pytest.raises(RuntimeError, match="no cohort"): + unfitted.run_many([10.0]) + + forced_infra: tp.Any = {**infra, "mode": "force"} + forced = Offset(infra=forced_infra) + forced.run_many(steps.FitCohort([1.0, 2.0, 3.0])) + assert forced._fits == [[1.0, 2.0, 3.0]], "force must refit instead of reading back" + + +def test_cohort_pickles_without_its_values() -> None: + cohort = steps.FitCohort([1.0, 2.0]) + cohort.uids = ["a", "b"] + loaded = pickle.loads(pickle.dumps(cohort)) + assert not loaded.items, "workers read the values from the carrier, not the cohort" + assert loaded.uids == ["a", "b"] + + +def test_named_cohort(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + step = Offset(infra=infra, cohort="train") + assert list(step.run_many(steps.FitCohort([1.0, 3.0]))) == [-1.0, 1.0] + assert step.cohort == "train", "a declared cohort must not rename it" + + configured = Offset(infra=infra, cohort="train") + assert list(configured.run_many([10.0])) == [8.0], "the config name must recover it" + assert not configured._fits + + forced_infra: tp.Any = {**infra, "mode": "force"} + forced = Offset(infra=forced_infra, cohort="train") + with pytest.raises(RuntimeError, match="drop the force mode"): + forced.run_many([10.0]) + + with pytest.raises(RuntimeError, match="must fit cohort 'test'"): + Offset(infra=infra, cohort="test").run_many([10.0]) + + +class Flaky(Offset): + """Fails its first fit, to leave an error cached for the artifact.""" + + broken: bool = True + + @classmethod + def _exclude_from_cls_uid(cls) -> list[str]: + return super()._exclude_from_cls_uid() + ["broken"] + + def _fit(self, values: tp.Iterable[float]) -> float: + if self.broken: + raise ValueError("Triggered an error") + return super()._fit(values) + + +def test_retry_of_a_failed_fit(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + with pytest.raises(ValueError, match="Triggered an error"): + Flaky(infra=infra, cohort="train").run_many(steps.FitCohort([1.0, 3.0])) + + retry: tp.Any = {**infra, "mode": "retry"} + with pytest.raises(RuntimeError, match="was handed no items"): + Flaky(infra=retry, cohort="train", broken=False).run_many([10.0]) + + fixed = Flaky(infra=retry, cohort="train", broken=False) + assert list(fixed.run_many(steps.FitCohort([1.0, 3.0]))) == [-1.0, 1.0] + + +@pytest.mark.parametrize( + "Variant,fit_values,reorder_fits", + [ + (Offset, [1.0, 4.0], []), + (OffsetMultiset, [1.0, 1.0, 4.0], []), + (OffsetSequence, [1.0, 1.0, 4.0], [[4.0, 1.0, 1.0]]), + ], +) +def test_cohort_key( + tmp_path: Path, + Variant: type[Offset], + fit_values: list[float], + reorder_fits: list[list[float]], +) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + step = Variant(infra=infra) + step.run_many(steps.FitCohort([1.0, 1.0, 4.0])) + assert step._fits == [fit_values], "the fit sees the cohort as its key defines it" + + reordered = Variant(infra=infra) + reordered.run_many(steps.FitCohort([4.0, 1.0, 1.0])) + assert reordered._fits == reorder_fits, "only a sequence cohort is order-sensitive" + + +@pytest.mark.parametrize("backend", ["ThreadPool", "ProcessPool"]) +@pytest.mark.parametrize("on_upstream", [False, True]) +def test_fit_with_distributed_items( + tmp_path: Path, on_upstream: bool, backend: str +) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": backend, "max_jobs": 2} + local: tp.Any = {"folder": tmp_path, "backend": "Cached"} + fit = Offset(infra=local if on_upstream else infra) + chain = steps.Chain( + steps=[conftest.Mult(coeff=2, infra=infra if on_upstream else local), fit] + ) + out = list(chain.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) + assert out == [-2.0, 0.0, 2.0] + assert fit._fits == [[2.0, 4.0, 6.0]], "the fit must see the whole cohort, once" + + +@pytest.mark.parametrize("max_jobs", [1, 2]) +def test_fit_under_distributed_chain(tmp_path: Path, max_jobs: int) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "ThreadPool", "max_jobs": max_jobs} + chain = steps.Chain(steps=[Offset()], infra=infra) + cohort = steps.FitCohort([1.0, 2.0, 3.0, 4.0]) + if max_jobs == 1: # one worker holds the whole cohort, so it can fit + list(chain.run_many(cohort)) # read: the pool is awaited lazily + else: + with pytest.raises(Exception, match="too few to fit"): + list(chain.run_many(cohort)) + + +@pytest.mark.parametrize("nested", [False, True]) +def test_fit_scopes_outputs_per_cohort(tmp_path: Path, nested: bool) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + + def pipeline() -> steps.Step: + if nested: # the chain keys its cache before the fit resolves + return steps.Chain(steps=[Offset()], infra=infra) + return Offset(infra=infra) + + assert list(pipeline().run_many(steps.FitCohort([1.0, 2.0, 3.0, 100.0])))[0] == -25.5 + assert list(pipeline().run_many(steps.FitCohort([1.0, 2.0, 3.0])))[0] == -1.0 + assert list(pipeline().run_many(steps.FitCohort([1.0, 5.0]))) == [-2.0, 2.0] + + +def test_fit_folder_structure(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + chain = steps.Chain( + steps=[ + conftest.Mult(coeff=2, infra=infra), + Offset(infra=infra), + OffsetMultiset(infra=infra), + conftest.Add(value=1, infra=infra), + ] + ) + assert list(chain.run_many(steps.FitCohort([1.0, 1.0, 4.0]))) == [-1.0, -1.0, 5.0] + + mult = "type=Mult-b9b7a7a5" + offset = f"{mult}/type=Offset,cohort=ddcbb484,2-2e44a0e3" + multiset = f"{offset}/cohort=230495c7,3,type=OffsetMultiset-d8dbd9b1" + assert conftest.extract_cache_folders(tmp_path) == ( + mult, # upstream of every fit: one folder for all cohorts + offset, # "set" key: the 2 distinct items + multiset, # "multiset" key: the 3 items, and the cohort of the fit above + f"{multiset}/value=1,type=Add-c1a6f4c8", # downstream: scoped by both fits + f"{offset}/type=_Artifact,owner.type=OffsetMultiset-74b84d3f", + f"{mult}/type=_Artifact,owner.type=Offset-21454ccf", # one entry per cohort + ) + + +def test_refit_needs_a_fresh_config(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + step = Offset(infra=infra) + assert list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0, 100.0])))[0] == -25.5 + with pytest.raises(RuntimeError, match="refitting in place"): + step.run_many(steps.FitCohort([1.0, 2.0, 3.0])) + fresh = step.clone({"cohort": None}) + assert list(fresh.run_many(steps.FitCohort([1.0, 2.0, 3.0])))[0] == -1.0 + + +def test_fit_variants_in_parallel(tmp_path: Path) -> None: + class Keyed(Offset): # re-keys items, so its uids are not the cohort's + def item_uid(self, value: tp.Any) -> str | None: + return f"v{value}" if isinstance(value, float) else None + + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + sweep = steps.helpers.Parallel( + steps=[Offset(infra=infra), OffsetMultiset(infra=infra), Keyed(infra=infra)] + ) + assert sweep.run_many(steps.FitCohort([1.0, 1.0, 4.0])) == [None] * 3 + variants = tp.cast(list[Offset], list(sweep.steps)) + fits = [v._fits for v in variants] + assert fits == [[[1.0, 4.0]], [[1.0, 1.0, 4.0]], [[1.0, 4.0]]], "one key each" + assert [v.lookup(4.0).result() for v in variants] == [1.5, 2.0, 1.5] + + +def test_cohort_names_every_fit(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + offset = Offset(infra=infra) + chain = steps.Chain( + steps=collections.OrderedDict(scale=conftest.Mult(coeff=2), offset=offset) + ) + assert offset.cohort is None, "unfitted, so nothing keys an artifact yet" + cohort = steps.FitCohort([1.0, 3.0]) + chain.run_many(cohort) + assert offset.cohort is not None, "the run names the cohort in the config" + assert cohort.fitted_by == [f"steps.offset({offset.cohort})"] + + vain = steps.FitCohort([1.0]) + conftest.Mult(coeff=2, infra=infra).run_many(vain) + assert not vain.fitted_by, "no Fit to name it, and nothing to fit" + + +# ============================================================================= +# Use cases +# ============================================================================= + + +class Normalize(steps.Fit): + """Standardizes each array with the cohort's mean and std.""" + + def _fit(self, values: tp.Iterable[np.ndarray]) -> tuple[np.ndarray, np.ndarray]: + total = np.zeros(()) + squares = np.zeros(()) + count = 0 + for value in values: # streamed: the cohort need not fit in memory + total = total + value.sum(axis=0) + squares = squares + (value**2).sum(axis=0) + count += value.shape[0] + mean = total / count + return mean, np.sqrt(squares / count - mean**2) + + def _run(self, value: np.ndarray) -> np.ndarray: + mean, std = self.fitted + return (value - mean) / std + + +class PCA(steps.Fit): + """Projects each array on the cohort's principal components.""" + + n_components: int = 2 + + def _fit(self, values: tp.Iterable[np.ndarray]) -> tuple[np.ndarray, np.ndarray]: + model = sklearn.decomposition.PCA(n_components=self.n_components) + model.fit(np.concatenate(list(values), axis=0)) + return model.mean_, model.components_ + + def _run(self, value: np.ndarray) -> np.ndarray: + mean, components = self.fitted + return (value - mean) @ components.T + + +class TrainLinear(steps.Fit): + """Trains a linear model to predict 2x+1, then predicts for each item.""" + + epochs: int = 300 + + def _fit(self, values: tp.Iterable[float]) -> dict[str, torch.Tensor]: + torch.manual_seed(12) + model = torch.nn.Linear(1, 1) + optimizer = torch.optim.SGD(model.parameters(), lr=0.05) + for _ in range(self.epochs): + x = torch.tensor([[value] for value in values]) # re-read, one pass/epoch + loss = torch.nn.functional.mse_loss(model(x), 2 * x + 1) + optimizer.zero_grad() + loss.backward() + optimizer.step() + return model.state_dict() # a Module would pickle, tensors cache natively + + def _run(self, value: float) -> float: + model = torch.nn.Linear(1, 1) + model.load_state_dict(self.fitted) + with torch.no_grad(): + return float(model(torch.tensor([value]))) + + +def test_normalize(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + rng = np.random.default_rng(12) + cohort = [rng.normal(3.0, 2.0, size=(8, 3)) for _ in range(4)] + step = Normalize(infra=infra) + out = np.concatenate(list(step.run_many(steps.FitCohort(cohort))), axis=0) + np.testing.assert_allclose(out.mean(axis=0), np.zeros(3), atol=1e-10) + np.testing.assert_allclose(out.std(axis=0), np.ones(3), atol=1e-10) + + +def test_pca(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + rng = np.random.default_rng(12) + base = rng.normal(size=(24, 2)) @ rng.normal(size=(2, 5)) + cohort = [base[i : i + 8] for i in range(0, 24, 8)] + step = PCA(n_components=2, infra=infra) + out = list(step.run_many(steps.FitCohort(cohort))) + assert [x.shape for x in out] == [(8, 2)] * 3 + + read_back = PCA(n_components=2, infra=infra) + again = list(read_back.run_many(steps.FitCohort(cohort))) + np.testing.assert_allclose(again[0], out[0], atol=1e-10) + novel = rng.normal(size=(8, 5)) + assert list(read_back.run_many([novel]))[0].shape == (8, 2) + + +def test_torch_train(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + step = TrainLinear(infra=infra) + out = list(step.run_many(steps.FitCohort([0.0, 1.0, 2.0, 3.0]))) + np.testing.assert_allclose(out, [1.0, 3.0, 5.0, 7.0], atol=0.2) + + read_back = TrainLinear(infra=infra) + read_back.run_many(steps.FitCohort([0.0, 1.0, 2.0, 3.0])) + assert read_back.run(10.0) == pytest.approx(21.0, abs=0.5) diff --git a/exca/steps/test_helpers.py b/exca/steps/test_helpers.py index 489f6634..7c3c8e74 100644 --- a/exca/steps/test_helpers.py +++ b/exca/steps/test_helpers.py @@ -179,6 +179,25 @@ def test_backend_on_steps_is_adopted_with_folder( assert [s.lookup(5.0).result() for s in sweep.steps] == [10.0, 15.0] +def test_resolved_variants(tmp_path: Path) -> None: + infra: tp.Any = {"backend": "Cached", "folder": tmp_path} + + class Configured(Step): + def _resolve_step(self) -> Step: + return conftest.Mult(coeff=3.0, infra=infra) + + sweep = Parallel(steps=[Configured()], infra=infra) + sweep.run(5.0) + assert sweep.steps[0].lookup(5.0).result() == 15.0, "cached as the resolved step" + + class Bare(Step): + def _resolve_step(self) -> Step: + return conftest.Mult(coeff=3.0) + + with pytest.raises(ValueError, match="no infra"): + Parallel(steps=[Bare()], infra=infra).run(5.0) + + def test_one_variant_errors_others_still_cache(tmp_path: Path) -> None: infra: tp.Any = {"backend": "LocalProcess", "folder": tmp_path} ok, bad = conftest.Add(value=2.0), conftest.Add(value=5.0, fail_on="all") @@ -190,13 +209,29 @@ def test_one_variant_errors_others_still_cache(tmp_path: Path) -> None: sweep.steps[1].lookup().result() # the error itself is cached +@pytest.mark.parametrize("resolved", [False, True]) def test_single_array_across_variants( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, resolved: bool ) -> None: _CapturingAutoExecutor.captured = [] monkeypatch.setattr(submitit, "AutoExecutor", _CapturingAutoExecutor) infra: tp.Any = {"backend": "Slurm", "folder": tmp_path} - sweep = Parallel(steps=[conftest.Add(value=v) for v in (1.0, 2.0, 3.0)], infra=infra) + + class Wrapper(Step): + value: float + + def _resolve_step(self) -> Step: + return conftest.Add(value=self.value, infra=infra) + + variant: tp.Any = Wrapper if resolved else conftest.Add + sweep = Parallel(steps=[variant(value=v) for v in (1.0, 2.0, 3.0)], infra=infra) sweep.run() [(_, params)] = _CapturingAutoExecutor.captured assert params["slurm_array_parallelism"] == 3, "one array spans all variants" + + +def test_nested_sweep(tmp_path: Path) -> None: + infra: tp.Any = {"backend": "Cached", "folder": tmp_path} + inner = Parallel(steps=[conftest.Add(value=1.0)], infra=infra) + Parallel(steps=[inner, conftest.Add(value=2.0)], infra=infra).run() + assert inner.steps[0].lookup().result() == 1.0 diff --git a/pyproject.toml b/pyproject.toml index e3186df1..765acac8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ dependencies = [ "mne>=1.4.0", "pybv>=0.7.6", "nibabel>=5.1.0", + "scikit-learn>=1.3", "pyarrow>=17.0.0", # Test "pytest>=7.4.0", @@ -115,7 +116,7 @@ addopts = ["--ignore=build", "--ignore=dist"] show_error_codes = true [[tool.mypy.overrides]] - module = ['pytest', 'setuptools', 'cloudpickle', 'mne', 'mne.*', 'nibabel', 'neuralset', 'pyarrow', 'pybv', 'excatest'] + module = ['pytest', 'setuptools', 'cloudpickle', 'mne', 'mne.*', 'nibabel', 'neuralset', 'pyarrow', 'pybv', 'excatest', 'sklearn', 'sklearn.*'] ignore_missing_imports = true [[tool.mypy.overrides]] # some packages we do not install From 42cb8a910a811f5d4ce9b90329ca89e20d4a21ca Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Tue, 1 Sep 2026 09:25:30 +0200 Subject: [PATCH 02/11] simplify --- CHANGELOG.md | 4 +- docs/steps/items.md | 12 ++-- docs/steps/reference.rst | 4 +- exca/steps/backends.py | 26 ++++--- exca/steps/base.py | 13 ++-- exca/steps/fit.py | 145 +++++++++++++++++++++------------------ exca/steps/helpers.py | 13 ++-- exca/steps/items.py | 17 ++--- exca/steps/test_fit.py | 70 +++++++++++-------- 9 files changed, 157 insertions(+), 147 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d33e0d5f..988cad59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,9 +2,7 @@ ## [Unreleased] -- `steps`: `Fit` primitive — fit one artifact over a cohort of items, then transform each item with it (N->1->N). The cohort is declared by wrapping the items in `FitCohort`; its identity (the items fingerprint, or the name the config already carries) is written to the `cohort` field before the run, scoping the artifact and the outputs. A step that ran is frozen, so fitting another cohort takes a fresh config. -- `steps`: `Parallel` now dispatches its variants through `Step._dispatch` (grouping their submission through the new `Backend._grouped`), so a variant can be a `Fit` or resolve through the `build()` convention — the latter used to run and cache as its unresolved wrapper. -- `cachedict`: `Auto` sequences all load back as lists, and the contents of a tuple now resolve — a nested tuple used to load back as raw dump info, permanently so on the pickle path. +- `steps`: `Fit` primitive — fit one artifact over a cohort of items and transform the items accordingly. ## 0.5.29 - 26-07-28 diff --git a/docs/steps/items.md b/docs/steps/items.md index d503e49c..e6b62dfb 100644 --- a/docs/steps/items.md +++ b/docs/steps/items.md @@ -162,12 +162,12 @@ for value in norm.run_many(test_paths): again (one upstream read per pass). The cohort's identity — the fingerprint of its items, or the name the -config already carries (`Normalize(cohort="train")`) — is written to -the `cohort` field before anything runs, so the artifact and every -downstream cache are scoped to it. A named cohort is recoverable from -the config alone, for a pipeline that never presents the items it was -fitted on. A step that ran is frozen, so fitting another cohort takes -a fresh config (`clone({"cohort": None})`). +config already carries (`Normalize(cohort="train")`) — settles before +anything runs, and the step that runs is a copy carrying it, so the +artifact and every downstream cache are scoped to it. A named cohort is +recoverable from the config alone, for a pipeline that never presents +the items it was fitted on. Fitting another cohort takes another +config, so `clone()` it. The fit runs where the step is dispatched from, ahead of any split, and is cached under `infra`. A `Fit` under a backend that shards the diff --git a/docs/steps/reference.rst b/docs/steps/reference.rst index c19ab46d..8ea61ceb 100644 --- a/docs/steps/reference.rst +++ b/docs/steps/reference.rst @@ -97,7 +97,7 @@ Fitting over items .. autoclass:: exca.steps.Fit :show-inheritance: - :members: fitted, COHORT_KEY, ARTIFACT_CACHE_TYPE - :private-members: _fit + :members: fitted, ARTIFACT_CACHE_TYPE + :private-members: _fit, _cohort_uids .. autoclass:: exca.steps.FitCohort diff --git a/exca/steps/backends.py b/exca/steps/backends.py index ba16a9e0..7f490e50 100644 --- a/exca/steps/backends.py +++ b/exca/steps/backends.py @@ -199,6 +199,11 @@ def _fold_modes(*modes: identity.ModeType) -> identity.ModeType: return acc +def _must_recompute(status: LookupStatus, mode: identity.ModeType) -> bool: + """Whether a cached entry runs again under *mode* (a cached error still raises).""" + return mode == "force" or (mode == "retry" and status == "error") + + def _effective_mode(step: Step) -> identity.ModeType: """The mode in effect for ``step`` once its sub-steps are folded in.""" from . import utils # lazy — backends is imported by utils at module level @@ -468,7 +473,7 @@ def _pending_statuses( if status == "error": _CachedEntry.lookup(cd, uid).result() # loads + re-raises continue - elif mode == "force" or (mode == "retry" and status == "error"): + elif _must_recompute(status, mode): pending[uid] = status elif status == "error": _CachedEntry.lookup(cd, uid).result() # loads + re-raises @@ -599,19 +604,12 @@ def _grouped(self) -> tp.Iterator[None]: if claimed.ready: self._execute(claimed.ready) - def _defer(self, cbatch: ComputeBatch) -> items.StepItems | None: - """Inside ``_grouped``, queue *cbatch* and return its pending carrier.""" - if self._group is None: - return None - self._group.append(cbatch) - return cbatch.cached_items() - def _run(self, step: Step, batch: items.StepItems) -> items.StepItems: """Execute *step* for uncached items, caching per uid.""" cbatch = self._prepare(step, batch) - deferred = self._defer(cbatch) - if deferred is not None: - return deferred + if self._group is not None: # submitted when the group exits + self._group.append(cbatch) + return cbatch.cached_items() with self._claim([cbatch]) as claimed: if claimed.ready: self._execute(claimed.ready) @@ -932,9 +930,9 @@ def _run(self, step: Step, batch: items.StepItems) -> items.StepItems: lazy carrier instead of blocking. """ cbatch = self._prepare(step, batch) - deferred = self._defer(cbatch) - if deferred is not None: - return deferred + if self._group is not None: # submitted when the group exits + self._group.append(cbatch) + return cbatch.cached_items() claimed = self._claim([cbatch]) transferred = False try: diff --git a/exca/steps/base.py b/exca/steps/base.py index 6e519589..e750c49a 100644 --- a/exca/steps/base.py +++ b/exca/steps/base.py @@ -393,14 +393,14 @@ def run_many(self, values: tp.Iterable[tp.Any] | FitCohort) -> items.StepItems: StepItems Iterator yielding one result per input, in input order. """ - from .fit import FitCohort, _stamp_cohorts # local import: fit builds on Step + from . import fit # circular built = utils.resolved_step(self) if built is not self: return built.run_many(values) cohort: FitCohort | None = None - if isinstance(values, FitCohort): + if isinstance(values, fit.FitCohort): cohort, values = values, values.items values = list(values) # eager: uid computation needs all values upfront uids = [identity.materialize_uid(self, v) for v in values] @@ -411,13 +411,10 @@ def run_many(self, values: tp.Iterable[tp.Any] | FitCohort) -> items.StepItems: return warm # extra-fast path -> avoid StepItems + _dispatch overhead else: cohort.uids = uids - _stamp_cohorts(self, cohort) + fit.declare_cohorts(self, cohort) - boundary = items.StepItems( - source=dict(zip(uids, values)), - uids=uids, - cohort=cohort, - ) + boundary = items.StepItems(source=dict(zip(uids, values)), uids=uids) + boundary._cohort = cohort is not None return boundary.apply_step(self) def forward(self, *args: tp.Any, **kwargs: tp.Any) -> tp.NoReturn: # removed diff --git a/exca/steps/fit.py b/exca/steps/fit.py index 262c1727..05478baa 100644 --- a/exca/steps/fit.py +++ b/exca/steps/fit.py @@ -11,13 +11,11 @@ import pydantic -import exca +from exca import utils as xkutils -from . import backends, items +from . import backends, items, utils from .base import Step -CohortKey = tp.Literal["set", "multiset", "sequence"] - class FitCohort: """The items a :class:`Fit` fits on, passed to ``run_many`` in their place. @@ -43,20 +41,9 @@ def __init__(self, items: tp.Iterable[tp.Any]) -> None: def __repr__(self) -> str: return f"{type(self).__name__}({len(self.items)} items)" - def __getstate__(self) -> dict[str, tp.Any]: - return {**self.__dict__, "items": []} # values travel with the carrier - - -def _cohort_uids(uids: tp.Sequence[str], key: CohortKey) -> list[str]: - """The cohort's uids as its fit sees them, ordered so the fit is reproducible.""" - if key == "sequence": - return list(uids) - return sorted(set(uids) if key == "set" else uids) - -def _fingerprint(uids: tp.Sequence[str], key: CohortKey) -> str: +def _fingerprint(ordered: tp.Sequence[str]) -> str: """Identity of a cohort that was not named, as ``,``.""" - ordered = _cohort_uids(uids, key) digest = hashlib.sha256() for uid in ordered: digest.update(uid.encode("utf8")) @@ -64,27 +51,35 @@ def _fingerprint(uids: tp.Sequence[str], key: CohortKey) -> str: return f"{digest.hexdigest()[:8]},{len(ordered)}" -def _stamp_cohorts(step: Step, cohort: FitCohort) -> None: - """Name the cohort in every ``Fit`` of *step*, before anything runs. - - The name is part of the configuration, hence of the cache identity, so it must - be settled before a composite step keys its own cache. - """ - fits = exca.utils.find_models(step, Fit, include_private=False) - for path, fit in fits.items(): - uid = fit.cohort - if uid is None or fit._stamped: # a name given in the config is kept - uid = _fingerprint(cohort.uids, fit.COHORT_KEY) - if fit.cohort != uid: - try: - fit.cohort = uid - except Exception as e: +def _find_fits(step: Step, root: str = "") -> dict[str, Fit]: + """Every ``Fit`` of *step*, including those a ``_resolve_step`` builds.""" + found: dict[str, Fit] = {} + for path, sub in xkutils.find_models(step, Step, include_private=False).items(): + key = f"{root}{path}" if path else root + if isinstance(sub, Fit): # not its resolution: that one carries a cohort + found[key] = sub + elif (built := utils.resolved_step(sub)) is not sub: + found.update(_find_fits(built, key)) + return found + + +def declare_cohorts(step: Step, cohort: FitCohort) -> None: + """Name the cohort in every ``Fit`` of *step*, privately, before anything runs.""" + for path, fit in _find_fits(step).items(): + if fit.cohort is not None and fit._declared is None: + uid = fit.cohort # a name given in the config is kept + else: + uid = _fingerprint(fit._cohort_uids(cohort.uids)) + # the config that ran: this copy, or the one it resolved to + memo = tp.cast("Fit | None", fit._resolution_cache) + ran = fit if fit.cohort is not None else memo + if ran is not None and ran.cohort != uid: raise RuntimeError( f"{type(fit).__name__} at {path or '.'} already ran on cohort " - f"{fit.cohort!r}: refitting in place is not supported, run " - "clone({'cohort': None}) on the new cohort" - ) from e - fit._stamped = True + f"{ran.cohort!r}: refitting in place is not supported, run " + "clone() to fit another cohort" + ) + fit._declared = uid cohort.fitted_by.append(f"{path or '.'}({uid})") @@ -93,46 +88,62 @@ class Fit(Step): .. warning:: Experimental -- API may change. - Override :meth:`_fit` for the artifact, and :meth:`_run` for one item's output - (reading :attr:`fitted`). + Example:: + + class Normalize(Fit): + def _fit(self, values): # the cohort, streamed + return np.stack(list(values)).mean(0) - Only a batch passed as a :class:`FitCohort` is fitted on; any other batch - transforms with what is already fitted, and raises if that is nothing. The - cohort's identity -- its name, or the fingerprint of its items -- is written to - :attr:`cohort` before the run, so the artifact and every downstream cache are - scoped to it. A step that ran is frozen, so its cohort cannot be replaced. + def _run(self, value): # one item + return value - self.fitted + + norm = Normalize(infra={"backend": "Cached", "folder": cache}) + norm.run_many(FitCohort(train)) # fits on these, then transforms them + norm.run_many(test) # transforms with the same artifact + + Only a :class:`FitCohort` is fitted on; its name -- :attr:`cohort`, else the + fingerprint of its items -- scopes the artifact and every downstream cache. + Fitting another cohort takes another config (:meth:`clone`). The fit runs where the step is dispatched from, ahead of a backend splitting the batch, and is cached under ``infra``. The upstream is read twice (once for the fit, once per item), so give an expensive upstream its own ``infra``. - ``COHORT_KEY`` states whether order and repetitions define the cohort, and - ``ARTIFACT_CACHE_TYPE`` the artifact's cache format (``CACHE_TYPE`` stays the - per-item outputs'). Prefer fitting arrays or tensors -- e.g. a state dict over a - model -- as they cache natively instead of pickling. + Override :meth:`_cohort_uids` if order or repetitions define the cohort. + ``ARTIFACT_CACHE_TYPE`` is the artifact's cache format -- prefer fitting arrays + or tensors (e.g. a state dict over a model), which cache natively. Parameters ---------- cohort - Name of the artifact, to fit it under a name or to use it in a run that - never presents the cohort (a config-only pipeline). Left unset, the - fingerprint of the cohort's items is written here instead. + Name for the artifact, to fit it under a name or to read it back in a run + that never presents the cohort. Unset, the items' fingerprint names it. """ - COHORT_KEY: tp.ClassVar[CohortKey] = "set" ARTIFACT_CACHE_TYPE: tp.ClassVar[str | None] = "Auto" cohort: str | None = None _fitted: tp.Any = pydantic.PrivateAttr(None) _fitted_for: str | None = pydantic.PrivateAttr(None) # cohort of `_fitted` - _stamped: bool = pydantic.PrivateAttr(False) # `cohort` written by a declaration + _declared: str | None = pydantic.PrivateAttr(None) # cohort handed by `run_many` + + def _resolve_step(self) -> Step: + if self.cohort is not None or self._declared is None: + return self + return self.model_copy(update={"cohort": self._declared}) + + def _cohort_uids(self, uids: tp.Sequence[str]) -> list[str]: + """The cohort's uids as :meth:`_fit` reads them, and as they identify it. + + Deduplicated and sorted; override with ``list(uids)`` for a sequence fit. + """ + return sorted(set(uids)) def _fit(self, values: tp.Iterable[tp.Any]) -> tp.Any: """The artifact for the cohort, from the values this step receives. - *values* streams the cohort, and can be iterated more than once (e.g. one - pass per epoch) -- at the cost of re-reading the upstream each time. + *values* re-iterates the cohort (one upstream read per pass). """ raise NotImplementedError @@ -146,12 +157,15 @@ def fitted(self) -> tp.Any: return self._fitted def _dispatch(self, batch: items.StepItems) -> items.StepItems: + built = utils.resolved_step(self) + if built is not self: + return built._dispatch(batch) # before super(): a split ships the artifact along with the step if self.cohort is None or self._fitted_for != self.cohort: - self._resolve(batch) + self._resolve_artifact(batch) return super()._dispatch(batch) - def _resolve(self, batch: items.StepItems) -> None: + def _resolve_artifact(self, batch: items.StepItems) -> None: """Read the artifact for this step's cohort back, or fit it.""" kind = type(self).__name__ uid = self.cohort @@ -160,30 +174,29 @@ def _resolve(self, batch: items.StepItems) -> None: f"{kind} has no cohort to fit on or to read back: run it on a " "FitCohort, or set its 'cohort' name" ) - cohort = batch._cohort mode = backends._fold_modes(batch._mode, backends._effective_mode(self)) # cohort cleared: all cohorts of this Fit share one folder, one entry each owner = self.model_copy(update={"infra": None, "cohort": None}) + owner._declared = None # or it would resolve the cohort back in infra = None if self.infra is None else self.infra.derive(mode=mode) artifact = _Artifact(owner=owner, infra=infra) upstream = tuple(batch._upstream) handle = artifact.lookup(_upstream=upstream, _uid=uid) status = handle.status - # same rule as `_pending_statuses`: a cached error still raises in "cached" mode - if status is None or mode == "force" or (mode == "retry" and status == "error"): - if cohort is None: + if status is None or backends._must_recompute(status, mode): + if not batch._cohort: hint = "drop the force mode" if mode == "force" else "check its name" raise RuntimeError( f"{kind} must fit cohort {uid!r} but was handed no items to fit " f"on: run it on a FitCohort, or {hint}" ) # counts, not uids: a step re-keying items has its own uid space - handed, declared_n = len(set(batch.uids)), len(set(cohort.uids)) - if handed < declared_n: + handed = len(set(batch.uids)) + if handed < batch._total_size: raise RuntimeError( - f"{kind} was handed {handed} of the {declared_n} items of cohort " - f"{uid!r}, too few to fit it -- an enclosing backend sharded them; " - "fit it before distributing, or move that backend onto this step" + f"{kind} was handed {handed} of the {batch._total_size} items of " + f"cohort {uid!r}, too few to fit it -- an enclosing backend sharded " + "them; fit it before distributing, or move that backend onto this step" ) carrier = items.StepItems( source={uid: (batch,)}, uids=[uid], upstream=upstream, mode=mode @@ -206,6 +219,4 @@ def _infer_cache_type(self) -> str | None: def _run(self, value: tuple[items.StepItems]) -> tp.Any: batch = value[0] # the carrier itself: re-iterable, applies its pending steps - return self.owner._fit( - batch.select(_cohort_uids(batch.uids, self.owner.COHORT_KEY)) - ) + return self.owner._fit(batch.select(self.owner._cohort_uids(batch.uids))) diff --git a/exca/steps/helpers.py b/exca/steps/helpers.py index ca48f656..fc4fdcc1 100644 --- a/exca/steps/helpers.py +++ b/exca/steps/helpers.py @@ -199,9 +199,8 @@ def _run_items(self, batch: items.StepItems) -> items.StepItems: values = list(batch) # one read, shared by every variant with self.infra._grouped(): for child in self.steps: - child = utils.resolved_step( - child - ) # item uids key on it, as does its cache + # item uids key on the resolved step, as does its cache + child = utils.resolved_step(child) if child.infra is None: raise ValueError( f"Parallel variant {type(child).__name__} resolves to a step " @@ -210,11 +209,9 @@ def _run_items(self, batch: items.StepItems) -> items.StepItems: if child.infra is not self.infra and child.infra == self.infra: child = child.model_copy(update={"infra": self.infra}) # one group uids = [identity.materialize_uid(child, v) for v in values] - child._dispatch( - items.StepItems( - source=dict(zip(uids, values)), uids=uids, cohort=batch._cohort - ) - ) + child_batch = items.StepItems(source=dict(zip(uids, values)), uids=uids) + child_batch._cohort = batch._cohort # same items, own uid space + child._dispatch(child_batch) return batch._replace(source={uid: None for uid in batch.uids}, pending=()) def run(self, value: tp.Any = identity.NoValue()) -> None: diff --git a/exca/steps/items.py b/exca/steps/items.py index e8681467..e2b10a8d 100644 --- a/exca/steps/items.py +++ b/exca/steps/items.py @@ -21,7 +21,6 @@ if tp.TYPE_CHECKING: from .base import Step - from .fit import FitCohort class _Source(tp.Protocol): @@ -129,7 +128,6 @@ def __init__( upstream: tp.Sequence[Step] = (), pending: tp.Sequence[Step] = (), mode: identity.ModeType = "cached", - cohort: FitCohort | None = None, ) -> None: if uids is None: if not isinstance(source, dict): @@ -142,7 +140,8 @@ def __init__( self._upstream = tuple(upstream) self._pending = tuple(pending) self._mode = mode - self._cohort = cohort + self._total_size = len(set(self.uids)) + self._cohort = False def __len__(self) -> int: return len(self.uids) @@ -152,20 +151,18 @@ def apply_step(self, step: Step) -> StepItems: return step._dispatch(self) def _replace(self, **changes: tp.Any) -> StepItems: - """Copy with some parts changed; everything else (mode, cohort, ...) carries over. - - Pass ``pending=()`` when *source* becomes computed results, or its steps - would run a second time. - """ + """Copy with *changes* applied; pass ``pending=()`` if *source* holds results.""" params: dict[str, tp.Any] = { "source": self._source, "uids": self.uids, "upstream": self._upstream, "pending": self._pending, "mode": self._mode, - "cohort": self._cohort, } - return StepItems(**{**params, **changes}) + new = StepItems(**{**params, **changes}) + new._total_size = self._total_size + new._cohort = self._cohort + return new def _append(self, step: Step) -> StepItems: """Append a single leaf step's computation and identity.""" diff --git a/exca/steps/test_fit.py b/exca/steps/test_fit.py index 8bf96eea..3cfacb09 100644 --- a/exca/steps/test_fit.py +++ b/exca/steps/test_fit.py @@ -5,7 +5,6 @@ # LICENSE file in the root directory of this source tree. import collections -import pickle import typing as tp from pathlib import Path @@ -15,6 +14,8 @@ import sklearn.decomposition import torch +import exca + from .. import steps from . import conftest @@ -33,12 +34,11 @@ def _run(self, value: float) -> float: return value - self.fitted -class OffsetMultiset(Offset): - COHORT_KEY = "multiset" - - class OffsetSequence(Offset): - COHORT_KEY = "sequence" + ARTIFACT_CACHE_TYPE = "Pickle" + + def _cohort_uids(self, uids: tp.Sequence[str]) -> list[str]: + return list(uids) def test_fit_cohort_then_novel_items(tmp_path: Path) -> None: @@ -62,14 +62,6 @@ def test_fit_cohort_then_novel_items(tmp_path: Path) -> None: assert forced._fits == [[1.0, 2.0, 3.0]], "force must refit instead of reading back" -def test_cohort_pickles_without_its_values() -> None: - cohort = steps.FitCohort([1.0, 2.0]) - cohort.uids = ["a", "b"] - loaded = pickle.loads(pickle.dumps(cohort)) - assert not loaded.items, "workers read the values from the carrier, not the cohort" - assert loaded.uids == ["a", "b"] - - def test_named_cohort(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} step = Offset(infra=infra, cohort="train") @@ -121,11 +113,10 @@ def test_retry_of_a_failed_fit(tmp_path: Path) -> None: "Variant,fit_values,reorder_fits", [ (Offset, [1.0, 4.0], []), - (OffsetMultiset, [1.0, 1.0, 4.0], []), (OffsetSequence, [1.0, 1.0, 4.0], [[4.0, 1.0, 1.0]]), ], ) -def test_cohort_key( +def test_cohort_uids( tmp_path: Path, Variant: type[Offset], fit_values: list[float], @@ -134,7 +125,7 @@ def test_cohort_key( infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} step = Variant(infra=infra) step.run_many(steps.FitCohort([1.0, 1.0, 4.0])) - assert step._fits == [fit_values], "the fit sees the cohort as its key defines it" + assert step._fits == [fit_values], "the fit sees the uids it asked for" reordered = Variant(infra=infra) reordered.run_many(steps.FitCohort([4.0, 1.0, 1.0])) @@ -189,7 +180,7 @@ def test_fit_folder_structure(tmp_path: Path) -> None: steps=[ conftest.Mult(coeff=2, infra=infra), Offset(infra=infra), - OffsetMultiset(infra=infra), + OffsetSequence(infra=infra), conftest.Add(value=1, infra=infra), ] ) @@ -197,15 +188,17 @@ def test_fit_folder_structure(tmp_path: Path) -> None: mult = "type=Mult-b9b7a7a5" offset = f"{mult}/type=Offset,cohort=ddcbb484,2-2e44a0e3" - multiset = f"{offset}/cohort=230495c7,3,type=OffsetMultiset-d8dbd9b1" + sequence = f"{offset}/cohort=230495c7,3,type=OffsetSequence-f1b2246a" assert conftest.extract_cache_folders(tmp_path) == ( mult, # upstream of every fit: one folder for all cohorts - offset, # "set" key: the 2 distinct items - multiset, # "multiset" key: the 3 items, and the cohort of the fit above - f"{multiset}/value=1,type=Add-c1a6f4c8", # downstream: scoped by both fits - f"{offset}/type=_Artifact,owner.type=OffsetMultiset-74b84d3f", + offset, # deduplicated: the 2 distinct items + sequence, # as a sequence: the 3 items, and the cohort of the fit above + f"{sequence}/value=1,type=Add-c1a6f4c8", # downstream: scoped by both fits + f"{offset}/type=_Artifact,owner.type=OffsetSequence-1cddb5b8", f"{mult}/type=_Artifact,owner.type=Offset-21454ccf", # one entry per cohort ) + [dumped] = tmp_path.glob("**/type=_Artifact*/cache/data/*") + assert dumped.suffix == ".pkl", "only ARTIFACT_CACHE_TYPE dumps beside the entry" def test_refit_needs_a_fresh_config(tmp_path: Path) -> None: @@ -214,8 +207,27 @@ def test_refit_needs_a_fresh_config(tmp_path: Path) -> None: assert list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0, 100.0])))[0] == -25.5 with pytest.raises(RuntimeError, match="refitting in place"): step.run_many(steps.FitCohort([1.0, 2.0, 3.0])) - fresh = step.clone({"cohort": None}) - assert list(fresh.run_many(steps.FitCohort([1.0, 2.0, 3.0])))[0] == -1.0 + assert list(step.clone().run_many(steps.FitCohort([1.0, 2.0, 3.0])))[0] == -1.0 + + +def test_fit_in_a_frozen_config(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + step = Offset(infra=infra) + exca.utils.recursive_freeze(step) + out = list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) + assert out == [-1.0, 0.0, 1.0], "a config frozen by an enclosing one still fits" + + +def test_fit_from_a_nested_resolve_step(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + + class Wrapper(steps.Step): + def _resolve_step(self) -> steps.Step: + return Offset(infra=infra) + + chain = steps.Chain(steps=[Wrapper()]) + out = list(chain.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) + assert out == [-1.0, 0.0, 1.0], "a Fit only a resolution builds must be named too" def test_fit_variants_in_parallel(tmp_path: Path) -> None: @@ -225,7 +237,7 @@ def item_uid(self, value: tp.Any) -> str | None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} sweep = steps.helpers.Parallel( - steps=[Offset(infra=infra), OffsetMultiset(infra=infra), Keyed(infra=infra)] + steps=[Offset(infra=infra), OffsetSequence(infra=infra), Keyed(infra=infra)] ) assert sweep.run_many(steps.FitCohort([1.0, 1.0, 4.0])) == [None] * 3 variants = tp.cast(list[Offset], list(sweep.steps)) @@ -240,11 +252,11 @@ def test_cohort_names_every_fit(tmp_path: Path) -> None: chain = steps.Chain( steps=collections.OrderedDict(scale=conftest.Mult(coeff=2), offset=offset) ) - assert offset.cohort is None, "unfitted, so nothing keys an artifact yet" cohort = steps.FitCohort([1.0, 3.0]) chain.run_many(cohort) - assert offset.cohort is not None, "the run names the cohort in the config" - assert cohort.fitted_by == [f"steps.offset({offset.cohort})"] + [named] = cohort.fitted_by + assert named.startswith("steps.offset("), "the run names the cohort in every Fit" + assert offset.cohort is None, "the config handed over is left as it was" vain = steps.FitCohort([1.0]) conftest.Mult(coeff=2, infra=infra).run_many(vain) From f760a94fe855451d9073e9da5797468fd504f74d Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Tue, 1 Sep 2026 11:54:41 +0200 Subject: [PATCH 03/11] fix --- exca/steps/base.py | 2 +- exca/steps/fit.py | 6 +++--- exca/steps/helpers.py | 2 +- exca/steps/test_fit.py | 9 +-------- 4 files changed, 6 insertions(+), 13 deletions(-) diff --git a/exca/steps/base.py b/exca/steps/base.py index e750c49a..d4104144 100644 --- a/exca/steps/base.py +++ b/exca/steps/base.py @@ -410,7 +410,7 @@ def run_many(self, values: tp.Iterable[tp.Any] | FitCohort) -> items.StepItems: if warm is not None: return warm # extra-fast path -> avoid StepItems + _dispatch overhead else: - cohort.uids = uids + cohort._uids = uids fit.declare_cohorts(self, cohort) boundary = items.StepItems(source=dict(zip(uids, values)), uids=uids) diff --git a/exca/steps/fit.py b/exca/steps/fit.py index 05478baa..ef634587 100644 --- a/exca/steps/fit.py +++ b/exca/steps/fit.py @@ -35,7 +35,7 @@ class FitCohort: def __init__(self, items: tp.Iterable[tp.Any]) -> None: self.items = list(items) - self.uids: list[str] = [] # the declared cohort, set by run_many + self._uids: list[str] = [] # the declared cohort, set by run_many self.fitted_by: list[str] = [] def __repr__(self) -> str: @@ -69,9 +69,9 @@ def declare_cohorts(step: Step, cohort: FitCohort) -> None: if fit.cohort is not None and fit._declared is None: uid = fit.cohort # a name given in the config is kept else: - uid = _fingerprint(fit._cohort_uids(cohort.uids)) + uid = _fingerprint(fit._cohort_uids(cohort._uids)) # the config that ran: this copy, or the one it resolved to - memo = tp.cast("Fit | None", fit._resolution_cache) + memo = tp.cast(Fit | None, fit._resolution_cache) ran = fit if fit.cohort is not None else memo if ran is not None and ran.cohort != uid: raise RuntimeError( diff --git a/exca/steps/helpers.py b/exca/steps/helpers.py index fc4fdcc1..e9179284 100644 --- a/exca/steps/helpers.py +++ b/exca/steps/helpers.py @@ -174,7 +174,7 @@ def _unify_infra(self) -> None: "Parallel requires one shared backend across itself and its " f"steps; {self.infra!r} differs from {step.infra!r}" ) - # the very object: one grouped dispatch then covers every variant + # one infra instance: its _grouped() then batches every variant together self.steps = [s.model_copy(update={"infra": self.infra}) for s in self.steps] def _uid_steps(self) -> list[Step]: diff --git a/exca/steps/test_fit.py b/exca/steps/test_fit.py index 3cfacb09..9566e162 100644 --- a/exca/steps/test_fit.py +++ b/exca/steps/test_fit.py @@ -341,16 +341,9 @@ def test_pca(tmp_path: Path) -> None: rng = np.random.default_rng(12) base = rng.normal(size=(24, 2)) @ rng.normal(size=(2, 5)) cohort = [base[i : i + 8] for i in range(0, 24, 8)] - step = PCA(n_components=2, infra=infra) - out = list(step.run_many(steps.FitCohort(cohort))) + out = list(PCA(n_components=2, infra=infra).run_many(steps.FitCohort(cohort))) assert [x.shape for x in out] == [(8, 2)] * 3 - read_back = PCA(n_components=2, infra=infra) - again = list(read_back.run_many(steps.FitCohort(cohort))) - np.testing.assert_allclose(again[0], out[0], atol=1e-10) - novel = rng.normal(size=(8, 5)) - assert list(read_back.run_many([novel]))[0].shape == (8, 2) - def test_torch_train(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} From 9734da0e8ac0776b081a6d455f7f94d41503dbde Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Tue, 1 Sep 2026 13:05:29 +0200 Subject: [PATCH 04/11] rmpca --- exca/steps/test_fit.py | 25 ------------------------- pyproject.toml | 3 +-- 2 files changed, 1 insertion(+), 27 deletions(-) diff --git a/exca/steps/test_fit.py b/exca/steps/test_fit.py index 9566e162..a49f2cef 100644 --- a/exca/steps/test_fit.py +++ b/exca/steps/test_fit.py @@ -11,7 +11,6 @@ import numpy as np import pydantic import pytest -import sklearn.decomposition import torch import exca @@ -287,21 +286,6 @@ def _run(self, value: np.ndarray) -> np.ndarray: return (value - mean) / std -class PCA(steps.Fit): - """Projects each array on the cohort's principal components.""" - - n_components: int = 2 - - def _fit(self, values: tp.Iterable[np.ndarray]) -> tuple[np.ndarray, np.ndarray]: - model = sklearn.decomposition.PCA(n_components=self.n_components) - model.fit(np.concatenate(list(values), axis=0)) - return model.mean_, model.components_ - - def _run(self, value: np.ndarray) -> np.ndarray: - mean, components = self.fitted - return (value - mean) @ components.T - - class TrainLinear(steps.Fit): """Trains a linear model to predict 2x+1, then predicts for each item.""" @@ -336,15 +320,6 @@ def test_normalize(tmp_path: Path) -> None: np.testing.assert_allclose(out.std(axis=0), np.ones(3), atol=1e-10) -def test_pca(tmp_path: Path) -> None: - infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} - rng = np.random.default_rng(12) - base = rng.normal(size=(24, 2)) @ rng.normal(size=(2, 5)) - cohort = [base[i : i + 8] for i in range(0, 24, 8)] - out = list(PCA(n_components=2, infra=infra).run_many(steps.FitCohort(cohort))) - assert [x.shape for x in out] == [(8, 2)] * 3 - - def test_torch_train(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} step = TrainLinear(infra=infra) diff --git a/pyproject.toml b/pyproject.toml index 765acac8..e3186df1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,6 @@ dependencies = [ "mne>=1.4.0", "pybv>=0.7.6", "nibabel>=5.1.0", - "scikit-learn>=1.3", "pyarrow>=17.0.0", # Test "pytest>=7.4.0", @@ -116,7 +115,7 @@ addopts = ["--ignore=build", "--ignore=dist"] show_error_codes = true [[tool.mypy.overrides]] - module = ['pytest', 'setuptools', 'cloudpickle', 'mne', 'mne.*', 'nibabel', 'neuralset', 'pyarrow', 'pybv', 'excatest', 'sklearn', 'sklearn.*'] + module = ['pytest', 'setuptools', 'cloudpickle', 'mne', 'mne.*', 'nibabel', 'neuralset', 'pyarrow', 'pybv', 'excatest'] ignore_missing_imports = true [[tool.mypy.overrides]] # some packages we do not install From d2d2c0e3501aa537c0f977a80552bfd2e49e8f2e Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Tue, 1 Sep 2026 13:46:47 +0200 Subject: [PATCH 05/11] fix --- exca/steps/test_fit.py | 211 ++++++++++++++++--------------------- exca/steps/test_helpers.py | 4 +- 2 files changed, 94 insertions(+), 121 deletions(-) diff --git a/exca/steps/test_fit.py b/exca/steps/test_fit.py index a49f2cef..d42553d8 100644 --- a/exca/steps/test_fit.py +++ b/exca/steps/test_fit.py @@ -19,21 +19,29 @@ from . import conftest -class Offset(steps.Fit): - """Subtracts the cohort's mean from each item, recording the cohorts it fit on.""" +class _SumOffset(steps.Fit): + """Adds the cohort's sum to each item, recording the cohorts it fit on.""" + + broken: bool = False _fits: list[list[float]] = pydantic.PrivateAttr(default_factory=list) + @classmethod + def _exclude_from_cls_uid(cls) -> list[str]: + return super()._exclude_from_cls_uid() + ["broken"] # a fix reuses the entry + def _fit(self, values: tp.Iterable[float]) -> float: + if self.broken: + raise ValueError("Triggered an error") vals = list(values) self._fits.append(vals) - return sum(vals) / len(vals) + return sum(vals) def _run(self, value: float) -> float: - return value - self.fitted + return value + self.fitted -class OffsetSequence(Offset): +class _SumOffsetSequence(_SumOffset): ARTIFACT_CACHE_TYPE = "Pickle" def _cohort_uids(self, uids: tp.Sequence[str]) -> list[str]: @@ -42,82 +50,68 @@ def _cohort_uids(self, uids: tp.Sequence[str]) -> list[str]: def test_fit_cohort_then_novel_items(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} - step = Offset(infra=infra) - assert list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) == [-1.0, 0.0, 1.0] - assert list(step.run_many([10.0])) == [8.0], "novel item must use the fitted mean" + step = _SumOffset(infra=infra) + exca.utils.recursive_freeze(step) # as an enclosing config would + assert list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) == [7.0, 8.0, 9.0] + assert step.run(10.0) == 16.0, "novel item must use the fitted sum" assert len(step._fits) == 1, "a second call must not refit" - read_back = Offset(infra=infra) - assert list(read_back.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) == [-1.0, 0.0, 1.0] + read_back = _SumOffset(infra=infra) + assert list(read_back.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) == [7.0, 8.0, 9.0] assert not read_back._fits, "the artifact must be read back from the cache" - unfitted = Offset(infra=infra) with pytest.raises(RuntimeError, match="no cohort"): - unfitted.run_many([10.0]) + _SumOffset(infra=infra).run(10.0) - forced_infra: tp.Any = {**infra, "mode": "force"} - forced = Offset(infra=forced_infra) + infra["mode"] = "force" + forced = _SumOffset(infra=infra) forced.run_many(steps.FitCohort([1.0, 2.0, 3.0])) assert forced._fits == [[1.0, 2.0, 3.0]], "force must refit instead of reading back" def test_named_cohort(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} - step = Offset(infra=infra, cohort="train") - assert list(step.run_many(steps.FitCohort([1.0, 3.0]))) == [-1.0, 1.0] + step = _SumOffset(infra=infra, cohort="train") + assert list(step.run_many(steps.FitCohort([1.0, 3.0]))) == [5.0, 7.0] assert step.cohort == "train", "a declared cohort must not rename it" - configured = Offset(infra=infra, cohort="train") - assert list(configured.run_many([10.0])) == [8.0], "the config name must recover it" + configured = _SumOffset(infra=infra, cohort="train") + assert configured.run(10.0) == 14.0, "the config name must recover it" assert not configured._fits - forced_infra: tp.Any = {**infra, "mode": "force"} - forced = Offset(infra=forced_infra, cohort="train") - with pytest.raises(RuntimeError, match="drop the force mode"): - forced.run_many([10.0]) - with pytest.raises(RuntimeError, match="must fit cohort 'test'"): - Offset(infra=infra, cohort="test").run_many([10.0]) - - -class Flaky(Offset): - """Fails its first fit, to leave an error cached for the artifact.""" - - broken: bool = True - - @classmethod - def _exclude_from_cls_uid(cls) -> list[str]: - return super()._exclude_from_cls_uid() + ["broken"] + _SumOffset(infra=infra, cohort="test").run(10.0) - def _fit(self, values: tp.Iterable[float]) -> float: - if self.broken: - raise ValueError("Triggered an error") - return super()._fit(values) + infra["mode"] = "force" + with pytest.raises(RuntimeError, match="drop the force mode"): + _SumOffset(infra=infra, cohort="train").run(10.0) def test_retry_of_a_failed_fit(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} with pytest.raises(ValueError, match="Triggered an error"): - Flaky(infra=infra, cohort="train").run_many(steps.FitCohort([1.0, 3.0])) + _SumOffset(infra=infra, cohort="train", broken=True).run_many( + steps.FitCohort([1.0, 3.0]) + ) - retry: tp.Any = {**infra, "mode": "retry"} + infra["mode"] = "retry" with pytest.raises(RuntimeError, match="was handed no items"): - Flaky(infra=retry, cohort="train", broken=False).run_many([10.0]) + _SumOffset(infra=infra, cohort="train").run(10.0) - fixed = Flaky(infra=retry, cohort="train", broken=False) - assert list(fixed.run_many(steps.FitCohort([1.0, 3.0]))) == [-1.0, 1.0] + fixed = _SumOffset(infra=infra, cohort="train") + assert list(fixed.run_many(steps.FitCohort([1.0, 3.0]))) == [5.0, 7.0] @pytest.mark.parametrize( "Variant,fit_values,reorder_fits", [ - (Offset, [1.0, 4.0], []), - (OffsetSequence, [1.0, 1.0, 4.0], [[4.0, 1.0, 1.0]]), + (_SumOffset, [1.0, 4.0], []), + (_SumOffsetSequence, [1.0, 1.0, 4.0], [[4.0, 1.0, 1.0]]), ], ) def test_cohort_uids( tmp_path: Path, - Variant: type[Offset], + Variant: type[_SumOffset], fit_values: list[float], reorder_fits: list[list[float]], ) -> None: @@ -138,19 +132,18 @@ def test_fit_with_distributed_items( ) -> None: infra: tp.Any = {"folder": tmp_path, "backend": backend, "max_jobs": 2} local: tp.Any = {"folder": tmp_path, "backend": "Cached"} - fit = Offset(infra=local if on_upstream else infra) + fit = _SumOffset(infra=local if on_upstream else infra) chain = steps.Chain( steps=[conftest.Mult(coeff=2, infra=infra if on_upstream else local), fit] ) - out = list(chain.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) - assert out == [-2.0, 0.0, 2.0] + assert list(chain.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) == [14.0, 16.0, 18.0] assert fit._fits == [[2.0, 4.0, 6.0]], "the fit must see the whole cohort, once" @pytest.mark.parametrize("max_jobs", [1, 2]) def test_fit_under_distributed_chain(tmp_path: Path, max_jobs: int) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "ThreadPool", "max_jobs": max_jobs} - chain = steps.Chain(steps=[Offset()], infra=infra) + chain = steps.Chain(steps=[_SumOffset()], infra=infra) cohort = steps.FitCohort([1.0, 2.0, 3.0, 4.0]) if max_jobs == 1: # one worker holds the whole cohort, so it can fit list(chain.run_many(cohort)) # read: the pool is awaited lazily @@ -160,17 +153,16 @@ def test_fit_under_distributed_chain(tmp_path: Path, max_jobs: int) -> None: @pytest.mark.parametrize("nested", [False, True]) -def test_fit_scopes_outputs_per_cohort(tmp_path: Path, nested: bool) -> None: +def test_a_config_fits_one_cohort(tmp_path: Path, nested: bool) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + step: steps.Step = _SumOffset(infra=infra) + if nested: # the chain keys its cache before the fit resolves + step = steps.Chain(steps=[step], infra=infra) - def pipeline() -> steps.Step: - if nested: # the chain keys its cache before the fit resolves - return steps.Chain(steps=[Offset()], infra=infra) - return Offset(infra=infra) - - assert list(pipeline().run_many(steps.FitCohort([1.0, 2.0, 3.0, 100.0])))[0] == -25.5 - assert list(pipeline().run_many(steps.FitCohort([1.0, 2.0, 3.0])))[0] == -1.0 - assert list(pipeline().run_many(steps.FitCohort([1.0, 5.0]))) == [-2.0, 2.0] + assert list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) == [7.0, 8.0, 9.0] + with pytest.raises(RuntimeError, match="refitting in place"): + step.run_many(steps.FitCohort([1.0, 4.0])) + assert list(step.clone().run_many(steps.FitCohort([1.0, 4.0]))) == [6.0, 9.0] def test_fit_folder_structure(tmp_path: Path) -> None: @@ -178,83 +170,65 @@ def test_fit_folder_structure(tmp_path: Path) -> None: chain = steps.Chain( steps=[ conftest.Mult(coeff=2, infra=infra), - Offset(infra=infra), - OffsetSequence(infra=infra), + _SumOffset(infra=infra), + _SumOffsetSequence(infra=infra), conftest.Add(value=1, infra=infra), ] ) - assert list(chain.run_many(steps.FitCohort([1.0, 1.0, 4.0]))) == [-1.0, -1.0, 5.0] + assert list(chain.run_many(steps.FitCohort([1.0, 1.0, 4.0]))) == [55.0, 55.0, 61.0] mult = "type=Mult-b9b7a7a5" - offset = f"{mult}/type=Offset,cohort=ddcbb484,2-2e44a0e3" - sequence = f"{offset}/cohort=230495c7,3,type=OffsetSequence-f1b2246a" + offset = f"{mult}/type=_SumOffset,cohort=ddcbb484,2-814895f3" + sequence = f"{offset}/cohort=230495c7,3,type=_SumOffsetSequence-38eb5fab" assert conftest.extract_cache_folders(tmp_path) == ( mult, # upstream of every fit: one folder for all cohorts + f"{mult}/type=_Artifact,owner.type=_SumOffset-47fc70c6", # one per cohort offset, # deduplicated: the 2 distinct items sequence, # as a sequence: the 3 items, and the cohort of the fit above f"{sequence}/value=1,type=Add-c1a6f4c8", # downstream: scoped by both fits - f"{offset}/type=_Artifact,owner.type=OffsetSequence-1cddb5b8", - f"{mult}/type=_Artifact,owner.type=Offset-21454ccf", # one entry per cohort + f"{offset}/type=_Artifact,owner.type=_SumOffsetSequence-8f5c57ef", ) [dumped] = tmp_path.glob("**/type=_Artifact*/cache/data/*") assert dumped.suffix == ".pkl", "only ARTIFACT_CACHE_TYPE dumps beside the entry" -def test_refit_needs_a_fresh_config(tmp_path: Path) -> None: - infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} - step = Offset(infra=infra) - assert list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0, 100.0])))[0] == -25.5 - with pytest.raises(RuntimeError, match="refitting in place"): - step.run_many(steps.FitCohort([1.0, 2.0, 3.0])) - assert list(step.clone().run_many(steps.FitCohort([1.0, 2.0, 3.0])))[0] == -1.0 - - -def test_fit_in_a_frozen_config(tmp_path: Path) -> None: - infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} - step = Offset(infra=infra) - exca.utils.recursive_freeze(step) - out = list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) - assert out == [-1.0, 0.0, 1.0], "a config frozen by an enclosing one still fits" - - -def test_fit_from_a_nested_resolve_step(tmp_path: Path) -> None: - infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} - - class Wrapper(steps.Step): - def _resolve_step(self) -> steps.Step: - return Offset(infra=infra) - - chain = steps.Chain(steps=[Wrapper()]) - out = list(chain.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) - assert out == [-1.0, 0.0, 1.0], "a Fit only a resolution builds must be named too" - - def test_fit_variants_in_parallel(tmp_path: Path) -> None: - class Keyed(Offset): # re-keys items, so its uids are not the cohort's + class _Keyed(_SumOffset): # re-keys items, so its uids are not the cohort's def item_uid(self, value: tp.Any) -> str | None: return f"v{value}" if isinstance(value, float) else None infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} sweep = steps.helpers.Parallel( - steps=[Offset(infra=infra), OffsetSequence(infra=infra), Keyed(infra=infra)] + steps=[ + _SumOffset(infra=infra), + _SumOffsetSequence(infra=infra), + _Keyed(infra=infra), + ] ) assert sweep.run_many(steps.FitCohort([1.0, 1.0, 4.0])) == [None] * 3 - variants = tp.cast(list[Offset], list(sweep.steps)) + variants = tp.cast(list[_SumOffset], list(sweep.steps)) fits = [v._fits for v in variants] assert fits == [[[1.0, 4.0]], [[1.0, 1.0, 4.0]], [[1.0, 4.0]]], "one key each" - assert [v.lookup(4.0).result() for v in variants] == [1.5, 2.0, 1.5] + assert [v.lookup(4.0).result() for v in variants] == [9.0, 10.0, 9.0] def test_cohort_names_every_fit(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} - offset = Offset(infra=infra) + + class _ResolvedOffset(steps.Step): + def _resolve_step(self) -> steps.Step: + return _SumOffset(infra=infra) + + offset = _SumOffset(infra=infra) chain = steps.Chain( - steps=collections.OrderedDict(scale=conftest.Mult(coeff=2), offset=offset) + steps=collections.OrderedDict( + scale=conftest.Mult(coeff=2), offset=offset, wrapped=_ResolvedOffset() + ) ) cohort = steps.FitCohort([1.0, 3.0]) - chain.run_many(cohort) - [named] = cohort.fitted_by - assert named.startswith("steps.offset("), "the run names the cohort in every Fit" + assert list(chain.run_many(cohort)) == [34.0, 38.0] + named = [n.split("(")[0] for n in cohort.fitted_by] + assert named == ["steps.offset", "steps.wrapped"], "even one a resolution builds" assert offset.cohort is None, "the config handed over is left as it was" vain = steps.FitCohort([1.0]) @@ -267,7 +241,7 @@ def test_cohort_names_every_fit(tmp_path: Path) -> None: # ============================================================================= -class Normalize(steps.Fit): +class _Normalize(steps.Fit): """Standardizes each array with the cohort's mean and std.""" def _fit(self, values: tp.Iterable[np.ndarray]) -> tuple[np.ndarray, np.ndarray]: @@ -286,7 +260,17 @@ def _run(self, value: np.ndarray) -> np.ndarray: return (value - mean) / std -class TrainLinear(steps.Fit): +def test_normalize(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + rng = np.random.default_rng(12) + cohort = [rng.normal(3.0, 2.0, size=(8, 3)) for _ in range(4)] + step = _Normalize(infra=infra) + out = np.concatenate(list(step.run_many(steps.FitCohort(cohort))), axis=0) + np.testing.assert_allclose(out.mean(axis=0), np.zeros(3), atol=1e-10) + np.testing.assert_allclose(out.std(axis=0), np.ones(3), atol=1e-10) + + +class _TrainLinear(steps.Fit): """Trains a linear model to predict 2x+1, then predicts for each item.""" epochs: int = 300 @@ -310,22 +294,11 @@ def _run(self, value: float) -> float: return float(model(torch.tensor([value]))) -def test_normalize(tmp_path: Path) -> None: - infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} - rng = np.random.default_rng(12) - cohort = [rng.normal(3.0, 2.0, size=(8, 3)) for _ in range(4)] - step = Normalize(infra=infra) - out = np.concatenate(list(step.run_many(steps.FitCohort(cohort))), axis=0) - np.testing.assert_allclose(out.mean(axis=0), np.zeros(3), atol=1e-10) - np.testing.assert_allclose(out.std(axis=0), np.ones(3), atol=1e-10) - - def test_torch_train(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} - step = TrainLinear(infra=infra) - out = list(step.run_many(steps.FitCohort([0.0, 1.0, 2.0, 3.0]))) + out = list(_TrainLinear(infra=infra).run_many(steps.FitCohort([0.0, 1.0, 2.0, 3.0]))) np.testing.assert_allclose(out, [1.0, 3.0, 5.0, 7.0], atol=0.2) - read_back = TrainLinear(infra=infra) + read_back = _TrainLinear(infra=infra) read_back.run_many(steps.FitCohort([0.0, 1.0, 2.0, 3.0])) assert read_back.run(10.0) == pytest.approx(21.0, abs=0.5) diff --git a/exca/steps/test_helpers.py b/exca/steps/test_helpers.py index 7c3c8e74..bdb10ab0 100644 --- a/exca/steps/test_helpers.py +++ b/exca/steps/test_helpers.py @@ -217,13 +217,13 @@ def test_single_array_across_variants( monkeypatch.setattr(submitit, "AutoExecutor", _CapturingAutoExecutor) infra: tp.Any = {"backend": "Slurm", "folder": tmp_path} - class Wrapper(Step): + class ResolvedAdd(Step): value: float def _resolve_step(self) -> Step: return conftest.Add(value=self.value, infra=infra) - variant: tp.Any = Wrapper if resolved else conftest.Add + variant: tp.Any = ResolvedAdd if resolved else conftest.Add sweep = Parallel(steps=[variant(value=v) for v in (1.0, 2.0, 3.0)], infra=infra) sweep.run() [(_, params)] = _CapturingAutoExecutor.captured From 30e1e6220e69d5309f7678755094df93a9c8338a Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Tue, 1 Sep 2026 17:14:08 +0200 Subject: [PATCH 06/11] fix --- exca/steps/fit.py | 15 +++++--- exca/steps/test_fit.py | 80 +++++----------------------------------- exca/steps/test_items.py | 1 + 3 files changed, 20 insertions(+), 76 deletions(-) diff --git a/exca/steps/fit.py b/exca/steps/fit.py index ef634587..217bfd96 100644 --- a/exca/steps/fit.py +++ b/exca/steps/fit.py @@ -13,7 +13,7 @@ from exca import utils as xkutils -from . import backends, items, utils +from . import backends, identity, items, utils from .base import Step @@ -125,7 +125,7 @@ def _run(self, value): # one item cohort: str | None = None _fitted: tp.Any = pydantic.PrivateAttr(None) - _fitted_for: str | None = pydantic.PrivateAttr(None) # cohort of `_fitted` + _fitted_for: tuple[str, str] | None = pydantic.PrivateAttr(None) # cohort, upstream _declared: str | None = pydantic.PrivateAttr(None) # cohort handed by `run_many` def _resolve_step(self) -> Step: @@ -161,12 +161,11 @@ def _dispatch(self, batch: items.StepItems) -> items.StepItems: if built is not self: return built._dispatch(batch) # before super(): a split ships the artifact along with the step - if self.cohort is None or self._fitted_for != self.cohort: - self._resolve_artifact(batch) + self._resolve_artifact(batch) return super()._dispatch(batch) def _resolve_artifact(self, batch: items.StepItems) -> None: - """Read the artifact for this step's cohort back, or fit it.""" + """Read the artifact for this step's cohort back, or fit it (no-op if held).""" kind = type(self).__name__ uid = self.cohort if uid is None: @@ -174,10 +173,14 @@ def _resolve_artifact(self, batch: items.StepItems) -> None: f"{kind} has no cohort to fit on or to read back: run it on a " "FitCohort, or set its 'cohort' name" ) + fitted_for = (uid, identity.step_uid(list(batch._upstream))) + if self._fitted_for == fitted_for: + return mode = backends._fold_modes(batch._mode, backends._effective_mode(self)) # cohort cleared: all cohorts of this Fit share one folder, one entry each owner = self.model_copy(update={"infra": None, "cohort": None}) owner._declared = None # or it would resolve the cohort back in + owner._fitted = owner._fitted_for = None # never read by _fit, and heavy infra = None if self.infra is None else self.infra.derive(mode=mode) artifact = _Artifact(owner=owner, infra=infra) upstream = tuple(batch._upstream) @@ -203,7 +206,7 @@ def _resolve_artifact(self, batch: items.StepItems) -> None: ) # dispatch, not lookup: goes through infra's mode and caching self._fitted = next(iter(artifact._dispatch(carrier))) - self._fitted_for = uid + self._fitted_for = fitted_for class _Artifact(Step): diff --git a/exca/steps/test_fit.py b/exca/steps/test_fit.py index d42553d8..5d2dfa65 100644 --- a/exca/steps/test_fit.py +++ b/exca/steps/test_fit.py @@ -8,10 +8,8 @@ import typing as tp from pathlib import Path -import numpy as np import pydantic import pytest -import torch import exca @@ -165,6 +163,16 @@ def test_a_config_fits_one_cohort(tmp_path: Path, nested: bool) -> None: assert list(step.clone().run_many(steps.FitCohort([1.0, 4.0]))) == [6.0, 9.0] +def test_one_config_under_two_upstreams(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + fit = _SumOffset(infra=infra) + cohort = [1.0, 2.0] # same items, so the same cohort uid names both fits + assert list(fit.run_many(steps.FitCohort(cohort))) == [4.0, 5.0] + chain = steps.Chain(steps=[conftest.Mult(coeff=10), fit]) + assert list(chain.run_many(steps.FitCohort(cohort))) == [40.0, 50.0] + assert fit._fits == [[1.0, 2.0], [10.0, 20.0]], "the upstream re-keys the artifact" + + def test_fit_folder_structure(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} chain = steps.Chain( @@ -234,71 +242,3 @@ def _resolve_step(self) -> steps.Step: vain = steps.FitCohort([1.0]) conftest.Mult(coeff=2, infra=infra).run_many(vain) assert not vain.fitted_by, "no Fit to name it, and nothing to fit" - - -# ============================================================================= -# Use cases -# ============================================================================= - - -class _Normalize(steps.Fit): - """Standardizes each array with the cohort's mean and std.""" - - def _fit(self, values: tp.Iterable[np.ndarray]) -> tuple[np.ndarray, np.ndarray]: - total = np.zeros(()) - squares = np.zeros(()) - count = 0 - for value in values: # streamed: the cohort need not fit in memory - total = total + value.sum(axis=0) - squares = squares + (value**2).sum(axis=0) - count += value.shape[0] - mean = total / count - return mean, np.sqrt(squares / count - mean**2) - - def _run(self, value: np.ndarray) -> np.ndarray: - mean, std = self.fitted - return (value - mean) / std - - -def test_normalize(tmp_path: Path) -> None: - infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} - rng = np.random.default_rng(12) - cohort = [rng.normal(3.0, 2.0, size=(8, 3)) for _ in range(4)] - step = _Normalize(infra=infra) - out = np.concatenate(list(step.run_many(steps.FitCohort(cohort))), axis=0) - np.testing.assert_allclose(out.mean(axis=0), np.zeros(3), atol=1e-10) - np.testing.assert_allclose(out.std(axis=0), np.ones(3), atol=1e-10) - - -class _TrainLinear(steps.Fit): - """Trains a linear model to predict 2x+1, then predicts for each item.""" - - epochs: int = 300 - - def _fit(self, values: tp.Iterable[float]) -> dict[str, torch.Tensor]: - torch.manual_seed(12) - model = torch.nn.Linear(1, 1) - optimizer = torch.optim.SGD(model.parameters(), lr=0.05) - for _ in range(self.epochs): - x = torch.tensor([[value] for value in values]) # re-read, one pass/epoch - loss = torch.nn.functional.mse_loss(model(x), 2 * x + 1) - optimizer.zero_grad() - loss.backward() - optimizer.step() - return model.state_dict() # a Module would pickle, tensors cache natively - - def _run(self, value: float) -> float: - model = torch.nn.Linear(1, 1) - model.load_state_dict(self.fitted) - with torch.no_grad(): - return float(model(torch.tensor([value]))) - - -def test_torch_train(tmp_path: Path) -> None: - infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} - out = list(_TrainLinear(infra=infra).run_many(steps.FitCohort([0.0, 1.0, 2.0, 3.0]))) - np.testing.assert_allclose(out, [1.0, 3.0, 5.0, 7.0], atol=0.2) - - read_back = _TrainLinear(infra=infra) - read_back.run_many(steps.FitCohort([0.0, 1.0, 2.0, 3.0])) - assert read_back.run(10.0) == pytest.approx(21.0, abs=0.5) diff --git a/exca/steps/test_items.py b/exca/steps/test_items.py index 95a1f32f..ceede4db 100644 --- a/exca/steps/test_items.py +++ b/exca/steps/test_items.py @@ -32,6 +32,7 @@ def source_abc(request: pytest.FixtureRequest, tmp_path: Path) -> items.StepItem def test_step_items_iteration_and_select(source_abc: items.StepItems) -> None: assert list(source_abc) == [1, 2, 3] + assert list(source_abc) == [1, 2, 3], "re-iterable: one source read per pass" assert list(source_abc.uids) == ["a", "b", "c"] sub = source_abc.select(["c", "a"]) assert list(sub) == [3, 1] From 915f5ad2830a3563251dbe14e6a8684a08c14982 Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Tue, 1 Sep 2026 19:23:34 +0200 Subject: [PATCH 07/11] wip --- exca/steps/backends.py | 30 +++++++++++++++------- exca/steps/base.py | 6 +++-- exca/steps/fit.py | 8 +++--- exca/steps/test_backends.py | 21 ++++++++++++++++ exca/steps/test_fit.py | 50 +++++++++++++++++++++++++++++-------- 5 files changed, 90 insertions(+), 25 deletions(-) diff --git a/exca/steps/backends.py b/exca/steps/backends.py index f5a518a9..84ff1b74 100644 --- a/exca/steps/backends.py +++ b/exca/steps/backends.py @@ -14,6 +14,7 @@ import collections import contextlib +import contextvars import dataclasses import datetime import logging @@ -424,6 +425,11 @@ def __exit__(self, *exc: tp.Any) -> None: self.close() +_ACTIVE_GROUPS: contextvars.ContextVar[dict[int, list[ComputeBatch]]] = ( + contextvars.ContextVar("exca_step_backend_groups", default={}) +) + + class Backend(exca.helpers.DiscriminatedModel, discriminator_key="backend"): """Base class for execution backends with integrated caching.""" @@ -441,7 +447,6 @@ def _exclude_from_cls_uid(cls) -> list[str]: # Force/retry: recompute each (step_folder, uid) at most once per lifetime _recomputed: set[tuple[Path, str]] = pydantic.PrivateAttr(default_factory=set) _checked_configs: set[Path] = pydantic.PrivateAttr(default_factory=set) - _group: list[ComputeBatch] | None = pydantic.PrivateAttr(default=None) def __getstate__(self) -> dict[str, tp.Any]: recomputed = self._recomputed @@ -451,6 +456,9 @@ def __getstate__(self) -> dict[str, tp.Any]: finally: self._recomputed = recomputed + def _active_group(self) -> list[ComputeBatch] | None: + return _ACTIVE_GROUPS.get().get(id(self)) + def _pending_statuses( self, *, @@ -591,15 +599,17 @@ def _grouped(self) -> tp.Iterator[None]: Carriers returned while grouping only hold results once the group ran. Nesting joins the enclosing group. """ - if self._group is not None: # nested sweep: the outer group submits it + group = self._active_group() + if group is not None: # nested sweep: the outer group submits it yield return - group: list[ComputeBatch] = [] - self._group = group + group = [] + groups = _ACTIVE_GROUPS.get() + token = _ACTIVE_GROUPS.set({**groups, id(self): group}) try: yield finally: - self._group = None + _ACTIVE_GROUPS.reset(token) with self._claim(group) as claimed: if claimed.ready: self._execute(claimed.ready) @@ -607,8 +617,9 @@ def _grouped(self) -> tp.Iterator[None]: def _run(self, step: Step, batch: items.StepItems) -> items.StepItems: """Execute *step* for uncached items, caching per uid.""" cbatch = self._prepare(step, batch) - if self._group is not None: # submitted when the group exits - self._group.append(cbatch) + group = self._active_group() + if group is not None: # submitted when the group exits + group.append(cbatch) return cbatch.cached_items() with self._claim([cbatch]) as claimed: if claimed.ready: @@ -930,8 +941,9 @@ def _run(self, step: Step, batch: items.StepItems) -> items.StepItems: lazy carrier instead of blocking. """ cbatch = self._prepare(step, batch) - if self._group is not None: # submitted when the group exits - self._group.append(cbatch) + group = self._active_group() + if group is not None: # submitted when the group exits + group.append(cbatch) return cbatch.cached_items() claimed = self._claim([cbatch]) transferred = False diff --git a/exca/steps/base.py b/exca/steps/base.py index d4104144..5c525b8d 100644 --- a/exca/steps/base.py +++ b/exca/steps/base.py @@ -248,7 +248,9 @@ def _warm_items(self, uids: tp.Sequence[str]) -> items.StepItems | None: with cached._source.frozen_cache_folder(): if not all(uid in cached._source for uid in uids): return None - return cached.select(uids) + warm = cached.select(uids) + warm._cohort = False + return warm def _dispatch(self, batch: items.StepItems) -> items.StepItems: """Resolve, then route *batch*: reuse/remember warm carrier, run inline or via backend.""" @@ -258,7 +260,7 @@ def _dispatch(self, batch: items.StepItems) -> items.StepItems: standalone = ( not batch._upstream and not batch._pending and batch._mode == "cached" ) - if standalone: # _output_items only valid with no upstream + if standalone and not batch._cohort: # _output_items only valid with no upstream warm = self._warm_items(batch.uids) if warm is not None: return warm diff --git a/exca/steps/fit.py b/exca/steps/fit.py index 217bfd96..90ea94d1 100644 --- a/exca/steps/fit.py +++ b/exca/steps/fit.py @@ -125,7 +125,7 @@ def _run(self, value): # one item cohort: str | None = None _fitted: tp.Any = pydantic.PrivateAttr(None) - _fitted_for: tuple[str, str] | None = pydantic.PrivateAttr(None) # cohort, upstream + _fitted_for: tuple[str, str] | None = pydantic.PrivateAttr(None) # cohort, artifact _declared: str | None = pydantic.PrivateAttr(None) # cohort handed by `run_many` def _resolve_step(self) -> Step: @@ -173,9 +173,6 @@ def _resolve_artifact(self, batch: items.StepItems) -> None: f"{kind} has no cohort to fit on or to read back: run it on a " "FitCohort, or set its 'cohort' name" ) - fitted_for = (uid, identity.step_uid(list(batch._upstream))) - if self._fitted_for == fitted_for: - return mode = backends._fold_modes(batch._mode, backends._effective_mode(self)) # cohort cleared: all cohorts of this Fit share one folder, one entry each owner = self.model_copy(update={"infra": None, "cohort": None}) @@ -184,6 +181,9 @@ def _resolve_artifact(self, batch: items.StepItems) -> None: infra = None if self.infra is None else self.infra.derive(mode=mode) artifact = _Artifact(owner=owner, infra=infra) upstream = tuple(batch._upstream) + fitted_for = (uid, identity.step_uid([*upstream, artifact])) + if self._fitted_for == fitted_for: + return handle = artifact.lookup(_upstream=upstream, _uid=uid) status = handle.status if status is None or backends._must_recompute(status, mode): diff --git a/exca/steps/test_backends.py b/exca/steps/test_backends.py index 72b3296f..bfb29dea 100644 --- a/exca/steps/test_backends.py +++ b/exca/steps/test_backends.py @@ -10,8 +10,10 @@ import logging import stat import sys +import threading import time import typing as tp +from concurrent import futures from pathlib import Path import pydantic @@ -346,6 +348,25 @@ def prepare(step: Step, value: float) -> backends.ComputeBatch: assert key not in backend._recomputed, "cb_ok never ran, so it must be unmarked" +def test_grouped_is_context_local(tmp_path: Path) -> None: + backend = backends.Cached(folder=tmp_path) + barrier = threading.Barrier(2) + groups: list[list[backends.ComputeBatch]] = [] + + def capture_group() -> None: + with backend._grouped(): + group = backend._active_group() + assert group is not None + groups.append(group) + barrier.wait(timeout=5) + + with futures.ThreadPoolExecutor(max_workers=2) as pool: + jobs = [pool.submit(capture_group) for _ in range(2)] + for job in jobs: + job.result(timeout=5) + assert groups[0] is not groups[1] + + def test_recomputed_keyed_by_step(tmp_path: Path) -> None: backend = backends.Cached(folder=tmp_path) # two distinct steps (distinct value → distinct folder), same input → same uid diff --git a/exca/steps/test_fit.py b/exca/steps/test_fit.py index 5d2dfa65..c0611133 100644 --- a/exca/steps/test_fit.py +++ b/exca/steps/test_fit.py @@ -21,6 +21,7 @@ class _SumOffset(steps.Fit): """Adds the cohort's sum to each item, recording the cohorts it fit on.""" broken: bool = False + scale: float = 1 _fits: list[list[float]] = pydantic.PrivateAttr(default_factory=list) @@ -33,7 +34,7 @@ def _fit(self, values: tp.Iterable[float]) -> float: raise ValueError("Triggered an error") vals = list(values) self._fits.append(vals) - return sum(vals) + return self.scale * sum(vals) def _run(self, value: float) -> float: return value + self.fitted @@ -67,6 +68,17 @@ def test_fit_cohort_then_novel_items(tmp_path: Path) -> None: assert forced._fits == [[1.0, 2.0, 3.0]], "force must refit instead of reading back" +def test_fit_cohort_marker_is_request_scoped(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + upstream = conftest.Mult(coeff=2, infra=infra) + upstream.run_many(steps.FitCohort([1.0, 2.0])) + assert upstream._output_items is not None, "Fitting should warm cache as well" + chain = steps.Chain(steps=[upstream, _SumOffset(cohort="train", infra=infra)]) + with pytest.raises(RuntimeError, match="handed no items"): + chain.run_many([1.0, 2.0]) + assert list(chain.run_many(steps.FitCohort([1.0, 2.0]))) == [8.0, 10.0] + + def test_named_cohort(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} step = _SumOffset(infra=infra, cohort="train") @@ -83,6 +95,8 @@ def test_named_cohort(tmp_path: Path) -> None: infra["mode"] = "force" with pytest.raises(RuntimeError, match="drop the force mode"): _SumOffset(infra=infra, cohort="train").run(10.0) + copied = step.model_copy(update={"scale": 10}) + assert list(copied.run_many(steps.FitCohort([1.0, 3.0]))) == [41.0, 43.0] def test_retry_of_a_failed_fit(tmp_path: Path) -> None: @@ -138,13 +152,26 @@ def test_fit_with_distributed_items( assert fit._fits == [[2.0, 4.0, 6.0]], "the fit must see the whole cohort, once" +@pytest.mark.parametrize( + "Variant,values,expected", + [ + (_SumOffset, [1.0, 2.0, 3.0, 4.0], [11.0, 12.0, 13.0, 14.0]), + (_SumOffsetSequence, [1.0, 1.0, 4.0], [7.0, 7.0, 10.0]), + ], +) @pytest.mark.parametrize("max_jobs", [1, 2]) -def test_fit_under_distributed_chain(tmp_path: Path, max_jobs: int) -> None: +def test_fit_under_distributed_chain( + tmp_path: Path, + max_jobs: int, + Variant: type[_SumOffset], + values: list[float], + expected: list[float], +) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "ThreadPool", "max_jobs": max_jobs} - chain = steps.Chain(steps=[_SumOffset()], infra=infra) - cohort = steps.FitCohort([1.0, 2.0, 3.0, 4.0]) + chain = steps.Chain(steps=[Variant()], infra=infra) + cohort = steps.FitCohort(values) if max_jobs == 1: # one worker holds the whole cohort, so it can fit - list(chain.run_many(cohort)) # read: the pool is awaited lazily + assert list(chain.run_many(cohort)) == expected else: with pytest.raises(Exception, match="too few to fit"): list(chain.run_many(cohort)) @@ -163,13 +190,16 @@ def test_a_config_fits_one_cohort(tmp_path: Path, nested: bool) -> None: assert list(step.clone().run_many(steps.FitCohort([1.0, 4.0]))) == [6.0, 9.0] -def test_one_config_under_two_upstreams(tmp_path: Path) -> None: - infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} - fit = _SumOffset(infra=infra) +@pytest.mark.parametrize("cached", [False, True]) +def test_one_config_under_two_upstreams(tmp_path: Path, cached: bool) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} if cached else None + fit = _SumOffset(infra=infra, cohort=None if cached else "train") cohort = [1.0, 2.0] # same items, so the same cohort uid names both fits - assert list(fit.run_many(steps.FitCohort(cohort))) == [4.0, 5.0] + first = fit.run_many(steps.FitCohort(cohort)) chain = steps.Chain(steps=[conftest.Mult(coeff=10), fit]) - assert list(chain.run_many(steps.FitCohort(cohort))) == [40.0, 50.0] + second = chain.run_many(steps.FitCohort(cohort)) + assert list(first) == [4.0, 5.0] + assert list(second) == [40.0, 50.0] assert fit._fits == [[1.0, 2.0], [10.0, 20.0]], "the upstream re-keys the artifact" From 0cf6a4d35466c54d39962642dbd25fbb2afcf31e Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Tue, 1 Sep 2026 19:30:08 +0200 Subject: [PATCH 08/11] more --- docs/steps/items.md | 2 +- exca/steps/fit.py | 9 +++++++-- exca/steps/test_fit.py | 40 +++++++++++++++++----------------------- 3 files changed, 25 insertions(+), 26 deletions(-) diff --git a/docs/steps/items.md b/docs/steps/items.md index e6b62dfb..7d025172 100644 --- a/docs/steps/items.md +++ b/docs/steps/items.md @@ -166,7 +166,7 @@ config already carries (`Normalize(cohort="train")`) — settles before anything runs, and the step that runs is a copy carrying it, so the artifact and every downstream cache are scoped to it. A named cohort is recoverable from the config alone, for a pipeline that never presents -the items it was fitted on. Fitting another cohort takes another +the items it was fitted on. Another cohort or upstream takes another config, so `clone()` it. The fit runs where the step is dispatched from, ahead of any split, diff --git a/exca/steps/fit.py b/exca/steps/fit.py index 90ea94d1..d1514101 100644 --- a/exca/steps/fit.py +++ b/exca/steps/fit.py @@ -103,7 +103,7 @@ def _run(self, value): # one item Only a :class:`FitCohort` is fitted on; its name -- :attr:`cohort`, else the fingerprint of its items -- scopes the artifact and every downstream cache. - Fitting another cohort takes another config (:meth:`clone`). + Another cohort or upstream takes another config (:meth:`clone`). The fit runs where the step is dispatched from, ahead of a backend splitting the batch, and is cached under ``infra``. The upstream is read twice (once for the @@ -182,7 +182,12 @@ def _resolve_artifact(self, batch: items.StepItems) -> None: artifact = _Artifact(owner=owner, infra=infra) upstream = tuple(batch._upstream) fitted_for = (uid, identity.step_uid([*upstream, artifact])) - if self._fitted_for == fitted_for: + if self._fitted_for is not None: + if self._fitted_for != fitted_for: + raise RuntimeError( + f"{kind} already holds artifact {self._fitted_for!r}, cannot replace " + f"it with {fitted_for!r}; use clone() for another upstream or config" + ) return handle = artifact.lookup(_upstream=upstream, _uid=uid) status = handle.status diff --git a/exca/steps/test_fit.py b/exca/steps/test_fit.py index c0611133..9eb3d01d 100644 --- a/exca/steps/test_fit.py +++ b/exca/steps/test_fit.py @@ -95,8 +95,8 @@ def test_named_cohort(tmp_path: Path) -> None: infra["mode"] = "force" with pytest.raises(RuntimeError, match="drop the force mode"): _SumOffset(infra=infra, cohort="train").run(10.0) - copied = step.model_copy(update={"scale": 10}) - assert list(copied.run_many(steps.FitCohort([1.0, 3.0]))) == [41.0, 43.0] + scaled = step.clone(scale=10) + assert list(scaled.run_many(steps.FitCohort([1.0, 3.0]))) == [41.0, 43.0] def test_retry_of_a_failed_fit(tmp_path: Path) -> None: @@ -152,26 +152,13 @@ def test_fit_with_distributed_items( assert fit._fits == [[2.0, 4.0, 6.0]], "the fit must see the whole cohort, once" -@pytest.mark.parametrize( - "Variant,values,expected", - [ - (_SumOffset, [1.0, 2.0, 3.0, 4.0], [11.0, 12.0, 13.0, 14.0]), - (_SumOffsetSequence, [1.0, 1.0, 4.0], [7.0, 7.0, 10.0]), - ], -) @pytest.mark.parametrize("max_jobs", [1, 2]) -def test_fit_under_distributed_chain( - tmp_path: Path, - max_jobs: int, - Variant: type[_SumOffset], - values: list[float], - expected: list[float], -) -> None: +def test_fit_under_distributed_chain(tmp_path: Path, max_jobs: int) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "ThreadPool", "max_jobs": max_jobs} - chain = steps.Chain(steps=[Variant()], infra=infra) - cohort = steps.FitCohort(values) + chain = steps.Chain(steps=[_SumOffset()], infra=infra) + cohort = steps.FitCohort([1.0, 2.0, 3.0, 4.0]) if max_jobs == 1: # one worker holds the whole cohort, so it can fit - assert list(chain.run_many(cohort)) == expected + list(chain.run_many(cohort)) # read: the pool is awaited lazily else: with pytest.raises(Exception, match="too few to fit"): list(chain.run_many(cohort)) @@ -191,16 +178,23 @@ def test_a_config_fits_one_cohort(tmp_path: Path, nested: bool) -> None: @pytest.mark.parametrize("cached", [False, True]) -def test_one_config_under_two_upstreams(tmp_path: Path, cached: bool) -> None: +def test_fit_clone_for_another_upstream(tmp_path: Path, cached: bool) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} if cached else None fit = _SumOffset(infra=infra, cohort=None if cached else "train") cohort = [1.0, 2.0] # same items, so the same cohort uid names both fits first = fit.run_many(steps.FitCohort(cohort)) - chain = steps.Chain(steps=[conftest.Mult(coeff=10), fit]) - second = chain.run_many(steps.FitCohort(cohort)) + with pytest.raises(RuntimeError, match="already holds artifact"): + steps.Chain(steps=[conftest.Mult(coeff=10), fit]).run_many( + steps.FitCohort(cohort) + ) + cloned = fit.clone() + second = steps.Chain(steps=[conftest.Mult(coeff=10), cloned]).run_many( + steps.FitCohort(cohort) + ) assert list(first) == [4.0, 5.0] assert list(second) == [40.0, 50.0] - assert fit._fits == [[1.0, 2.0], [10.0, 20.0]], "the upstream re-keys the artifact" + assert fit._fits == [[1.0, 2.0]] + assert cloned._fits == [[10.0, 20.0]] def test_fit_folder_structure(tmp_path: Path) -> None: From 677d7390e6fbf6fc1dcd3becf6041bc6b0d3ed52 Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Wed, 2 Sep 2026 11:33:07 +0200 Subject: [PATCH 09/11] wip --- exca/steps/patterns.py | 2 ++ exca/steps/test_fit.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/exca/steps/patterns.py b/exca/steps/patterns.py index 9e05fccd..15f25c07 100644 --- a/exca/steps/patterns.py +++ b/exca/steps/patterns.py @@ -235,6 +235,8 @@ def _run_items(self, batch: items.StepItems) -> items.StepItems: upstream=branch_upstream, mode=batch._mode, ) + if batch._cohort and len(set(batch.uids)) == batch._total_size: + carrier._cohort = True # one dispatch over all branches lets a backend submit them together dispatched = self._body()._dispatch(carrier) return batch._replace( diff --git a/exca/steps/test_fit.py b/exca/steps/test_fit.py index 9eb3d01d..f18b3c91 100644 --- a/exca/steps/test_fit.py +++ b/exca/steps/test_fit.py @@ -15,6 +15,7 @@ from .. import steps from . import conftest +from .patterns import Scatter class _SumOffset(steps.Fit): @@ -164,6 +165,37 @@ def test_fit_under_distributed_chain(tmp_path: Path, max_jobs: int) -> None: list(chain.run_many(cohort)) +class _ScatterValues(Scatter): + body: steps.Step + + def branches(self, item: dict[str, float]) -> list[str]: + return list(item) + + def take(self, item: dict[str, float], branch: str) -> float: + return item[branch] + + +def test_fit_inside_scatter(tmp_path: Path) -> None: + cohort = [{"a": 1.0, "b": 2.0}, {"c": 3.0}] + expected = [{"a": 7.0, "b": 8.0}, {"c": 9.0}] + cache: tp.Any = {"folder": tmp_path / "cached", "backend": "Cached"} + scatter = _ScatterValues(body=_SumOffset(infra=cache)) + assert list(scatter.run_many(steps.FitCohort(cohort))) == expected, ( + "a complete Scatter must fit over all branches" + ) + + outer: tp.Any = {**cache, "backend": "ThreadPool", "max_jobs": 2} + chain = steps.Chain(steps=[scatter.clone()], infra=outer) + assert list(chain.run_many(steps.FitCohort(cohort))) == expected, ( + "an incomplete Scatter may reuse an existing artifact" + ) + + cold: tp.Any = {**outer, "folder": tmp_path / "cold"} + chain = steps.Chain(steps=[_ScatterValues(body=_SumOffset())], infra=cold) + with pytest.raises(RuntimeError, match="handed no items"): + list(chain.run_many(steps.FitCohort(cohort))) + + @pytest.mark.parametrize("nested", [False, True]) def test_a_config_fits_one_cohort(tmp_path: Path, nested: bool) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} From bf5264624f05e2b3b37a637a8a118d3805ba12b7 Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Wed, 2 Sep 2026 13:51:24 +0200 Subject: [PATCH 10/11] fix --- exca/steps/backends.py | 80 +++++++++++++++++++++---------- exca/steps/test_backends.py | 24 ++++++++++ exca/steps/test_fit.py | 96 +++++++++++++++++++++++++------------ 3 files changed, 146 insertions(+), 54 deletions(-) diff --git a/exca/steps/backends.py b/exca/steps/backends.py index 84ff1b74..9a539ec3 100644 --- a/exca/steps/backends.py +++ b/exca/steps/backends.py @@ -312,11 +312,22 @@ class ComputeBatch: def __getstate__(self) -> dict[str, tp.Any]: return {**self.__dict__, "info": CoordinationInfo()} + def _claimed_uids(self) -> list[str]: + claim = self.info.claim + if claim is None: + raise RuntimeError("ComputeBatch has no claim") + return list(claim.uids) + def select(self, uids: tp.Sequence[str]) -> ComputeBatch: """Sub-batch over *uids*, sharing step/paths/cache; copies ``info`` (avoid aliasing the parent's claim). """ info = dataclasses.replace(self.info) + if info.claim is not None: + selected = set(uids) + info.claim = dataclasses.replace( + info.claim, uids=tuple(uid for uid in info.claim.uids if uid in selected) + ) items_ = self.items.select(uids, mode=self.info.mode) return dataclasses.replace(self, items=items_, info=info) @@ -341,12 +352,15 @@ def run_and_cache(self) -> None: folder = self.cache_dict.folder if folder is not None: folder.mkdir(parents=True, exist_ok=True) + unique_uids = list(dict.fromkeys(self.items.uids)) + with self.cache_dict.frozen_cache_folder(): + run_uids = [uid for uid in unique_uids if uid not in self.cache_dict] result_items = self.step._run_items(self.items) written_uids: list[str] = [] try: with self.cache_dict.write(): - for i, result in enumerate(result_items): - uid = self.items.uids[i] + for i, result in enumerate(result_items.read(run_uids)): + uid = run_uids[i] if uid not in self.cache_dict: self.cache_dict[uid] = result written_uids.append(uid) @@ -668,9 +682,10 @@ def _claim(self, cbatches: list[ComputeBatch]) -> _Claimed: reg: inflight.InflightRegistry | None = None if self._concurrent: reg = inflight.InflightRegistry(cb.paths.step_folder) - cb = cb.select(list(pending)) + uids = cb.items.uids if cb.items._cohort else pending + cb = cb.select(list(uids)) cb.info.claim = claimed.stack.enter_context( - inflight.inflight_session(reg, set(pending)) + inflight.inflight_session(reg, pending) ) claimed.batches.append(cb) claimed.ready = [ @@ -688,11 +703,12 @@ def _recheck_and_clear(self, cbatch: ComputeBatch) -> ComputeBatch | None: or ``None`` if fully populated by a competitor. """ mode = cbatch.info.mode + claimed_uids = cbatch._claimed_uids() pending_statuses = self._pending_statuses( - paths=cbatch.paths, uids=cbatch.items.uids, mode=mode + paths=cbatch.paths, uids=claimed_uids, mode=mode ) inflight.after_wait_log( - cbatch.paths.step_uid, len(cbatch.items.uids), len(pending_statuses) + cbatch.paths.step_uid, len(claimed_uids), len(pending_statuses) ) retry_count = sum(status == "error" for status in pending_statuses.values()) if retry_count: @@ -707,7 +723,12 @@ def _recheck_and_clear(self, cbatch: ComputeBatch) -> ComputeBatch | None: self._clear_caches(paths=cbatch.paths, cd=cbatch.cache_dict, uids=clear_uids) if not pending_statuses: return None - return cbatch.select(list(pending_statuses)) + if cbatch.info.claim is not None: + cbatch.info.claim = dataclasses.replace( + cbatch.info.claim, uids=tuple(pending_statuses) + ) + uids = cbatch.items.uids if cbatch.items._cohort else pending_statuses + return cbatch.select(list(uids)) def _mark_recomputed(self, cbatch: ComputeBatch) -> None: """Record *cbatch*'s uids as recomputed-this-lifetime. @@ -717,7 +738,7 @@ def _mark_recomputed(self, cbatch: ComputeBatch) -> None: """ if cbatch.info.mode in ("force", "retry"): folder = cbatch.paths.step_folder - self._recomputed.update((folder, uid) for uid in cbatch.items.uids) + self._recomputed.update((folder, uid) for uid in cbatch._claimed_uids()) def _execute(self, cbatches: list[ComputeBatch]) -> None: """Run *cbatches* (filtered+claimed) blocking; override for pools/arrays.""" @@ -749,7 +770,7 @@ def _execute(self, cbatches: list[ComputeBatch]) -> None: if log_folder is not None: time = datetime.datetime.now(datetime.UTC).isoformat(timespec="seconds") step_uids = ", ".join(cbatch.paths.step_uid for cbatch in cbatches) - n_items = sum(len(cbatch.items.uids) for cbatch in cbatches) + n_items = sum(len(cbatch._claimed_uids()) for cbatch in cbatches) header = f"{time} - Running {n_items} items for steps: {step_uids}" print(header) print(header, file=sys.stderr) @@ -789,10 +810,14 @@ def _execute(self, cbatches: list[ComputeBatch]) -> None: if cbatch.info.claim is None: raise RuntimeError("_execute runs only on claimed batches") self._mark_recomputed(cbatch) # all tasks submitted together below - tasks = _tasks_from_batches( - [cb.shuffled() for cb in cbatches], - max_chunks=self.max_jobs, - min_items_per_chunk=self.min_items_per_job, + tasks = ( + [cbatches] + if self.max_jobs == 1 + else _tasks_from_batches( + [cb.select(cb._claimed_uids()).shuffled() for cb in cbatches], + max_chunks=self.max_jobs, + min_items_per_chunk=self.min_items_per_job, + ) ) # one array → one logs folder; jobs.db still records per step_folder executor = submitit.AutoExecutor( @@ -809,13 +834,14 @@ def _execute(self, cbatches: list[ComputeBatch]) -> None: for task, job in zip(tasks, jobs): for batch in task: assert batch.info.claim is not None # inherited from its variant - batch.info.claim.record_worker_info(job, uids=batch.items.uids) + uids = batch._claimed_uids() + batch.info.claim.record_worker_info(job, uids=uids) folder = batch.paths.step_folder - by_folder.setdefault(folder, {})[job.job_id] = batch.items.uids + by_folder.setdefault(folder, {})[job.job_id] = uids for folder, records in by_folder.items(): with jobregistry.JobRegistry(folder) as reg: reg.record(records, cluster=executor.cluster) - n_items = sum(len(cb.items.uids) for cb in cbatches) + n_items = sum(len(cb._claimed_uids()) for cb in cbatches) msg = "Sent %s items for %s steps into %s jobs on cluster '%s' (eg: %s)" logger.info( msg, n_items, len(cbatches), len(tasks), self._CLUSTER, jobs[0].job_id @@ -956,7 +982,7 @@ def _run(self, step: Step, batch: items.StepItems) -> items.StepItems: uid: fut for fut, task in task_futs.items() for b in task - for uid in b.items.uids + for uid in b._claimed_uids() } ctx = _PoolContext(cbatch.cache_dict, cbatch.paths.step_uid, pool, claimed) transferred = True # _PoolContext closes `claimed`, not the finally @@ -975,7 +1001,7 @@ def _execute(self, cbatches: list[ComputeBatch]) -> None: if submission is None: # ran inline (single worker) return pool, task_futs = submission - n_items = sum(len(cb.items.uids) for cb in cbatches) + n_items = sum(len(cb._claimed_uids()) for cb in cbatches) with pool: try: for f in futures.as_completed(task_futs): @@ -993,7 +1019,7 @@ def _submit_pool( its task->future map, or ``None`` if the work ran inline (single worker). """ # one pool across variants: heterogeneous variants overlap (load balance) - n_items = sum(len(cb.items.uids) for cb in cbatches) + n_items = sum(len(cb._claimed_uids()) for cb in cbatches) for cbatch in cbatches: if cbatch.info.claim is None: raise RuntimeError("_submit_pool runs only on claimed batches") @@ -1002,20 +1028,26 @@ def _submit_pool( max_workers = min(n_items, cpus) if self.max_jobs is not None: max_workers = min(max_workers, self.max_jobs) - if max_workers <= 1: - for cbatch in cbatches: + one_task = self.max_jobs == 1 + batches = ( + cbatches if one_task else [cb.select(cb._claimed_uids()) for cb in cbatches] + ) + if max_workers <= 1 and not ( + one_task and any(cb.items._cohort for cb in cbatches) + ): + for cbatch in batches: cbatch.run_and_cache() return None # ~3x as many tasks as workers, run in one pool tasks = _tasks_from_batches( - [cb.shuffled() for cb in cbatches], - max_chunks=3 * max_workers, + batches if one_task else [cb.shuffled() for cb in batches], + max_chunks=1 if one_task else 3 * max_workers, min_items_per_chunk=1, ) for task in tasks: for batch in task: assert batch.info.claim is not None # inherited from its variant - batch.info.claim.record_worker_info(uids=batch.items.uids) + batch.info.claim.record_worker_info(uids=batch._claimed_uids()) pool = utils.make_pool_executor(self._POOL_TYPE, max_workers) logger.info("Sent %s items for %s steps into a %s", n_items, len(cbatches), pool) task_futs = {pool.submit(_multi_run_and_cache, task): task for task in tasks} diff --git a/exca/steps/test_backends.py b/exca/steps/test_backends.py index bfb29dea..5ed2e157 100644 --- a/exca/steps/test_backends.py +++ b/exca/steps/test_backends.py @@ -288,6 +288,30 @@ def test_pool_backend(tmp_path: Path, backend: str) -> None: assert step.lookup(1.0).paths.cache_folder.exists() +def test_multiworker_cohort_chunks_do_not_duplicate_payload( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + source = {str(i): i for i in range(10)} + batch = items.StepItems(source=source, uids=list(source)) + batch._cohort = True + backend = backends.ThreadPool(folder=tmp_path, max_jobs=2) + cbatch = backend._prepare(conftest.Add(value=1, infra=backend), batch) + tasks: list[list[backends.ComputeBatch]] = [] + monkeypatch.setattr(backends.os, "cpu_count", lambda: 4) + monkeypatch.setattr(backends, "_multi_run_and_cache", tasks.append) + + with backend._claim([cbatch]) as claimed: + backend._execute(claimed.ready) + + assert len(tasks) > 1, "multiworker cohort must be split" + task_size = 0 + for task in tasks: + for cb in task: + assert isinstance(cb.items._source, dict) + task_size += len(cb.items._source) + assert task_size == len(source), "task sources must partition payload" + + class _PrintStep(Step): def _run(self, value: str) -> str: print(f"stdout:{value}") diff --git a/exca/steps/test_fit.py b/exca/steps/test_fit.py index f18b3c91..b5aeeeaf 100644 --- a/exca/steps/test_fit.py +++ b/exca/steps/test_fit.py @@ -25,6 +25,7 @@ class _SumOffset(steps.Fit): scale: float = 1 _fits: list[list[float]] = pydantic.PrivateAttr(default_factory=list) + _runs: list[float] = pydantic.PrivateAttr(default_factory=list) @classmethod def _exclude_from_cls_uid(cls) -> list[str]: @@ -38,6 +39,7 @@ def _fit(self, values: tp.Iterable[float]) -> float: return self.scale * sum(vals) def _run(self, value: float) -> float: + self._runs.append(value) return value + self.fitted @@ -52,12 +54,14 @@ def test_fit_cohort_then_novel_items(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} step = _SumOffset(infra=infra) exca.utils.recursive_freeze(step) # as an enclosing config would - assert list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) == [7.0, 8.0, 9.0] + output = list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) + assert output == [7.0, 8.0, 9.0], "cohort items use their fitted sum" assert step.run(10.0) == 16.0, "novel item must use the fitted sum" assert len(step._fits) == 1, "a second call must not refit" read_back = _SumOffset(infra=infra) - assert list(read_back.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) == [7.0, 8.0, 9.0] + output = list(read_back.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) + assert output == [7.0, 8.0, 9.0], "cached artifact must reproduce outputs" assert not read_back._fits, "the artifact must be read back from the cache" with pytest.raises(RuntimeError, match="no cohort"): @@ -77,18 +81,20 @@ def test_fit_cohort_marker_is_request_scoped(tmp_path: Path) -> None: chain = steps.Chain(steps=[upstream, _SumOffset(cohort="train", infra=infra)]) with pytest.raises(RuntimeError, match="handed no items"): chain.run_many([1.0, 2.0]) - assert list(chain.run_many(steps.FitCohort([1.0, 2.0]))) == [8.0, 10.0] + output = list(chain.run_many(steps.FitCohort([1.0, 2.0]))) + assert output == [8.0, 10.0], "current request's cohort marker must reach Fit" def test_named_cohort(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} step = _SumOffset(infra=infra, cohort="train") - assert list(step.run_many(steps.FitCohort([1.0, 3.0]))) == [5.0, 7.0] + output = list(step.run_many(steps.FitCohort([1.0, 3.0]))) + assert output == [5.0, 7.0], "named cohort must fit and transform" assert step.cohort == "train", "a declared cohort must not rename it" configured = _SumOffset(infra=infra, cohort="train") assert configured.run(10.0) == 14.0, "the config name must recover it" - assert not configured._fits + assert not configured._fits, "recovering a named artifact must not refit" with pytest.raises(RuntimeError, match="must fit cohort 'test'"): _SumOffset(infra=infra, cohort="test").run(10.0) @@ -97,7 +103,8 @@ def test_named_cohort(tmp_path: Path) -> None: with pytest.raises(RuntimeError, match="drop the force mode"): _SumOffset(infra=infra, cohort="train").run(10.0) scaled = step.clone(scale=10) - assert list(scaled.run_many(steps.FitCohort([1.0, 3.0]))) == [41.0, 43.0] + output = list(scaled.run_many(steps.FitCohort([1.0, 3.0]))) + assert output == [41.0, 43.0], "cloned config must fit its own artifact" def test_retry_of_a_failed_fit(tmp_path: Path) -> None: @@ -112,7 +119,8 @@ def test_retry_of_a_failed_fit(tmp_path: Path) -> None: _SumOffset(infra=infra, cohort="train").run(10.0) fixed = _SumOffset(infra=infra, cohort="train") - assert list(fixed.run_many(steps.FitCohort([1.0, 3.0]))) == [5.0, 7.0] + output = list(fixed.run_many(steps.FitCohort([1.0, 3.0]))) + assert output == [5.0, 7.0], "retry must allow the repaired fit" @pytest.mark.parametrize( @@ -149,17 +157,40 @@ def test_fit_with_distributed_items( chain = steps.Chain( steps=[conftest.Mult(coeff=2, infra=infra if on_upstream else local), fit] ) - assert list(chain.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) == [14.0, 16.0, 18.0] + output = list(chain.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) + assert output == [14.0, 16.0, 18.0], "distribution must preserve cohort results" assert fit._fits == [[2.0, 4.0, 6.0]], "the fit must see the whole cohort, once" -@pytest.mark.parametrize("max_jobs", [1, 2]) -def test_fit_under_distributed_chain(tmp_path: Path, max_jobs: int) -> None: - infra: tp.Any = {"folder": tmp_path, "backend": "ThreadPool", "max_jobs": max_jobs} +@pytest.mark.parametrize( + ("backend", "max_jobs"), + [("ThreadPool", 1), ("ProcessPool", 1), ("ThreadPool", 2)], +) +def test_fit_under_distributed_chain(tmp_path: Path, backend: str, max_jobs: int) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": backend, "max_jobs": max_jobs} chain = steps.Chain(steps=[_SumOffset()], infra=infra) cohort = steps.FitCohort([1.0, 2.0, 3.0, 4.0]) if max_jobs == 1: # one worker holds the whole cohort, so it can fit - list(chain.run_many(cohort)) # read: the pool is awaited lazily + output = list(chain.run_many(cohort)) + assert output == [11.0, 12.0, 13.0, 14.0], "one task keeps the cohort" + fit = _SumOffsetSequence() + sequence = steps.Chain( + steps=[conftest.Mult(coeff=2), fit, conftest.Add(value=1)], + infra=infra, + ) + output = list(sequence.run_many(steps.FitCohort([1.0, 1.0, 4.0]))) + assert output == [15.0, 15.0, 21.0], "enclosing backend preserves the sequence" + if backend == "ThreadPool": # worker shares the observable Fit instance + assert fit._fits == [[2.0, 2.0, 8.0]], "fitting keeps repeated values" + runs = sorted(fit._runs) + assert runs == [2.0, 8.0], "transforming deduplicates cache keys" + fit._runs.clear() + + sequence.lookup(4.0).clear_cache(recursive=False) + output = list(sequence.run_many(steps.FitCohort([1.0, 1.0, 4.0]))) + assert output == [15.0, 15.0, 21.0], "partial cache keeps cohort semantics" + if backend == "ThreadPool": + assert fit._runs == [8.0], "only the missing cache key is transformed" else: with pytest.raises(Exception, match="too few to fit"): list(chain.run_many(cohort)) @@ -180,15 +211,13 @@ def test_fit_inside_scatter(tmp_path: Path) -> None: expected = [{"a": 7.0, "b": 8.0}, {"c": 9.0}] cache: tp.Any = {"folder": tmp_path / "cached", "backend": "Cached"} scatter = _ScatterValues(body=_SumOffset(infra=cache)) - assert list(scatter.run_many(steps.FitCohort(cohort))) == expected, ( - "a complete Scatter must fit over all branches" - ) + output = list(scatter.run_many(steps.FitCohort(cohort))) + assert output == expected, "a complete Scatter must fit over all branches" outer: tp.Any = {**cache, "backend": "ThreadPool", "max_jobs": 2} chain = steps.Chain(steps=[scatter.clone()], infra=outer) - assert list(chain.run_many(steps.FitCohort(cohort))) == expected, ( - "an incomplete Scatter may reuse an existing artifact" - ) + output = list(chain.run_many(steps.FitCohort(cohort))) + assert output == expected, "an incomplete Scatter may reuse an existing artifact" cold: tp.Any = {**outer, "folder": tmp_path / "cold"} chain = steps.Chain(steps=[_ScatterValues(body=_SumOffset())], infra=cold) @@ -203,10 +232,12 @@ def test_a_config_fits_one_cohort(tmp_path: Path, nested: bool) -> None: if nested: # the chain keys its cache before the fit resolves step = steps.Chain(steps=[step], infra=infra) - assert list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) == [7.0, 8.0, 9.0] + output = list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) + assert output == [7.0, 8.0, 9.0], "one config must fit its first cohort" with pytest.raises(RuntimeError, match="refitting in place"): step.run_many(steps.FitCohort([1.0, 4.0])) - assert list(step.clone().run_many(steps.FitCohort([1.0, 4.0]))) == [6.0, 9.0] + output = list(step.clone().run_many(steps.FitCohort([1.0, 4.0]))) + assert output == [6.0, 9.0], "clone must fit another cohort independently" @pytest.mark.parametrize("cached", [False, True]) @@ -223,10 +254,10 @@ def test_fit_clone_for_another_upstream(tmp_path: Path, cached: bool) -> None: second = steps.Chain(steps=[conftest.Mult(coeff=10), cloned]).run_many( steps.FitCohort(cohort) ) - assert list(first) == [4.0, 5.0] - assert list(second) == [40.0, 50.0] - assert fit._fits == [[1.0, 2.0]] - assert cloned._fits == [[10.0, 20.0]] + assert list(first) == [4.0, 5.0], "the original must keep its first upstream" + assert list(second) == [40.0, 50.0], "the clone must use its new upstream" + assert fit._fits == [[1.0, 2.0]], "the original must fit unscaled values" + assert cloned._fits == [[10.0, 20.0]], "the clone must fit scaled values" def test_fit_folder_structure(tmp_path: Path) -> None: @@ -239,7 +270,8 @@ def test_fit_folder_structure(tmp_path: Path) -> None: conftest.Add(value=1, infra=infra), ] ) - assert list(chain.run_many(steps.FitCohort([1.0, 1.0, 4.0]))) == [55.0, 55.0, 61.0] + output = list(chain.run_many(steps.FitCohort([1.0, 1.0, 4.0]))) + assert output == [55.0, 55.0, 61.0], "stacked Fits must use both artifacts" mult = "type=Mult-b9b7a7a5" offset = f"{mult}/type=_SumOffset,cohort=ddcbb484,2-814895f3" @@ -251,7 +283,7 @@ def test_fit_folder_structure(tmp_path: Path) -> None: sequence, # as a sequence: the 3 items, and the cohort of the fit above f"{sequence}/value=1,type=Add-c1a6f4c8", # downstream: scoped by both fits f"{offset}/type=_Artifact,owner.type=_SumOffsetSequence-8f5c57ef", - ) + ), "artifact and downstream caches must use their respective cohort identities" [dumped] = tmp_path.glob("**/type=_Artifact*/cache/data/*") assert dumped.suffix == ".pkl", "only ARTIFACT_CACHE_TYPE dumps beside the entry" @@ -269,11 +301,14 @@ def item_uid(self, value: tp.Any) -> str | None: _Keyed(infra=infra), ] ) - assert sweep.run_many(steps.FitCohort([1.0, 1.0, 4.0])) == [None] * 3 + output = sweep.run_many(steps.FitCohort([1.0, 1.0, 4.0])) + assert output == [None] * 3, "each Fit variant must consume the cohort" variants = tp.cast(list[_SumOffset], list(sweep.steps)) fits = [v._fits for v in variants] - assert fits == [[[1.0, 4.0]], [[1.0, 1.0, 4.0]], [[1.0, 4.0]]], "one key each" - assert [v.lookup(4.0).result() for v in variants] == [9.0, 10.0, 9.0] + expected = [[[1.0, 4.0]], [[1.0, 1.0, 4.0]], [[1.0, 4.0]]] + assert fits == expected, "each variant must fit on its selected uid sequence" + output = [v.lookup(4.0).result() for v in variants] + assert output == [9.0, 10.0, 9.0], "variants cache their own results" def test_cohort_names_every_fit(tmp_path: Path) -> None: @@ -290,7 +325,8 @@ def _resolve_step(self) -> steps.Step: ) ) cohort = steps.FitCohort([1.0, 3.0]) - assert list(chain.run_many(cohort)) == [34.0, 38.0] + output = list(chain.run_many(cohort)) + assert output == [34.0, 38.0], "every discovered Fit contributes its artifact" named = [n.split("(")[0] for n in cohort.fitted_by] assert named == ["steps.offset", "steps.wrapped"], "even one a resolution builds" assert offset.cohort is None, "the config handed over is left as it was" From 1d2446bd5cae8aaa1fef8763a5fac629c587cb6f Mon Sep 17 00:00:00 2001 From: Jeremy Rapin Date: Wed, 2 Sep 2026 14:26:33 +0200 Subject: [PATCH 11/11] simplify --- docs/steps/items.md | 15 ++----- docs/steps/reference.rst | 4 +- exca/steps/__init__.py | 1 - exca/steps/base.py | 49 ++++++++++++--------- exca/steps/fit.py | 88 +++++++++++++++++--------------------- exca/steps/helpers.py | 12 +++++- exca/steps/test_fit.py | 92 +++++++++++++++++++--------------------- 7 files changed, 126 insertions(+), 135 deletions(-) diff --git a/docs/steps/items.md b/docs/steps/items.md index 7d025172..ee83b918 100644 --- a/docs/steps/items.md +++ b/docs/steps/items.md @@ -135,10 +135,7 @@ output order matches input order. ## Fitting on the items -A `Fit` step derives one artifact from many items, then transforms -each item with it — normalization statistics, a PCA basis, a trained -model. Only a batch wrapped in a `FitCohort` is fitted on; any other -call transforms with what is already fitted: +A `Fit` step derives one artifact from many items, then transforms each item with it — normalization statistics, a PCA basis, a trained model. `fit_many` fits; `run_many` transforms with what is already fitted: ```python class Normalize(steps.Fit): @@ -152,7 +149,7 @@ class Normalize(steps.Fit): norm = Normalize(infra={"backend": "Cached", "folder": cache}) -for value in norm.run_many(steps.FitCohort(train_paths)): +for value in norm.fit_many(train_paths): train(value) # fitted on this cohort, then transformed for value in norm.run_many(test_paths): evaluate(value) # same artifact, novel items @@ -161,13 +158,7 @@ for value in norm.run_many(test_paths): `_fit` receives the cohort as an iterable it can stream, and iterate again (one upstream read per pass). -The cohort's identity — the fingerprint of its items, or the name the -config already carries (`Normalize(cohort="train")`) — settles before -anything runs, and the step that runs is a copy carrying it, so the -artifact and every downstream cache are scoped to it. A named cohort is -recoverable from the config alone, for a pipeline that never presents -the items it was fitted on. Another cohort or upstream takes another -config, so `clone()` it. +The cohort's identity — the fingerprint of its items, or the name passed to `fit_many(..., cohort="train")` — settles before anything runs, and the step that runs is a copy carrying it, so the artifact and every downstream cache are scoped to it. Bind an existing named cohort with `fit_many(cohort="train")`. Another cohort or upstream takes another config, so `clone()` it. The fit runs where the step is dispatched from, ahead of any split, and is cached under `infra`. A `Fit` under a backend that shards the diff --git a/docs/steps/reference.rst b/docs/steps/reference.rst index 8ea61ceb..302a2e2c 100644 --- a/docs/steps/reference.rst +++ b/docs/steps/reference.rst @@ -11,7 +11,7 @@ Core classes ------------ .. autoclass:: exca.steps.Step - :members: run, run_many, lookup, clone, item_uid, CACHE_TYPE + :members: run, run_many, fit_many, lookup, clone, item_uid, CACHE_TYPE .. autoclass:: exca.steps.Chain :show-inheritance: @@ -99,5 +99,3 @@ Fitting over items :show-inheritance: :members: fitted, ARTIFACT_CACHE_TYPE :private-members: _fit, _cohort_uids - -.. autoclass:: exca.steps.FitCohort diff --git a/exca/steps/__init__.py b/exca/steps/__init__.py index 91bdd2a4..ede0aed0 100644 --- a/exca/steps/__init__.py +++ b/exca/steps/__init__.py @@ -30,4 +30,3 @@ from .base import Chain as Chain from .base import Step as Step from .fit import Fit as Fit -from .fit import FitCohort as FitCohort diff --git a/exca/steps/base.py b/exca/steps/base.py index 5c525b8d..124e2b5d 100644 --- a/exca/steps/base.py +++ b/exca/steps/base.py @@ -22,9 +22,6 @@ from . import backends, identity, items, utils -if tp.TYPE_CHECKING: - from .fit import FitCohort - logger = logging.getLogger(__name__) @@ -381,44 +378,58 @@ def run(self, value: tp.Any = identity.NoValue()) -> tp.Any: """ return next(iter(self.run_many([value]))) - def run_many(self, values: tp.Iterable[tp.Any] | FitCohort) -> items.StepItems: + def run_many(self, values: tp.Iterable[tp.Any]) -> items.StepItems: """Execute the step over many inputs, one cache entry per input. Parameters ---------- values: - Inputs to run; one result is produced per input, in order. Wrap them in a - :class:`~exca.steps.FitCohort` to let a :class:`~exca.steps.Fit` fit on them. + Inputs to run; one result is produced per input, in order. Returns ------- StepItems Iterator yielding one result per input, in input order. """ - from . import fit # circular - built = utils.resolved_step(self) if built is not self: return built.run_many(values) - cohort: FitCohort | None = None - if isinstance(values, fit.FitCohort): - cohort, values = values, values.items values = list(values) # eager: uid computation needs all values upfront uids = [identity.materialize_uid(self, v) for v in values] - if cohort is None: # a declared cohort must reach the steps - warm = self._warm_items(uids) - if warm is not None: - return warm # extra-fast path -> avoid StepItems + _dispatch overhead - else: - cohort._uids = uids - fit.declare_cohorts(self, cohort) + warm = self._warm_items(uids) + if warm is not None: + return warm # extra-fast path -> avoid StepItems + _dispatch overhead boundary = items.StepItems(source=dict(zip(uids, values)), uids=uids) - boundary._cohort = cohort is not None return boundary.apply_step(self) + def fit_many( + self, + values: tp.Iterable[tp.Any] = (), + *, + cohort: str | None = None, + ) -> items.StepItems: + """Fit contained :class:`Fit` steps, then return their transformed inputs. + + Parameters + ---------- + values: + Cohort items. Omit to bind an existing named cohort. + cohort: + Cohort name. Unset, the item fingerprint identifies it. + """ + from . import fit # circular + + built = utils.resolved_step(self) + values = list(values) + uids = [identity.materialize_uid(built, value) for value in values] + fit._declare_cohorts(self, uids, cohort) + boundary = items.StepItems(source=dict(zip(uids, values)), uids=uids) + boundary._cohort = bool(values) + return boundary.apply_step(built) + def forward(self, *args: tp.Any, **kwargs: tp.Any) -> tp.NoReturn: # removed raise AttributeError("Step.forward() was removed; use run() instead") diff --git a/exca/steps/fit.py b/exca/steps/fit.py index d1514101..7c462252 100644 --- a/exca/steps/fit.py +++ b/exca/steps/fit.py @@ -17,31 +17,6 @@ from .base import Step -class FitCohort: - """The items a :class:`Fit` fits on, passed to ``run_many`` in their place. - - .. warning:: Experimental -- API may change. - - Parameters - ---------- - items - Items to fit on, then transform. - - Note - ---- - ``fitted_by`` lists the steps the cohort identified, to check it was not - passed in vain. - """ - - def __init__(self, items: tp.Iterable[tp.Any]) -> None: - self.items = list(items) - self._uids: list[str] = [] # the declared cohort, set by run_many - self.fitted_by: list[str] = [] - - def __repr__(self) -> str: - return f"{type(self).__name__}({len(self.items)} items)" - - def _fingerprint(ordered: tp.Sequence[str]) -> str: """Identity of a cohort that was not named, as ``,``.""" digest = hashlib.sha256() @@ -63,24 +38,37 @@ def _find_fits(step: Step, root: str = "") -> dict[str, Fit]: return found -def declare_cohorts(step: Step, cohort: FitCohort) -> None: +def _declare_cohorts( + step: Step, uids: tp.Sequence[str], cohort: str | None = None +) -> None: """Name the cohort in every ``Fit`` of *step*, privately, before anything runs.""" - for path, fit in _find_fits(step).items(): - if fit.cohort is not None and fit._declared is None: - uid = fit.cohort # a name given in the config is kept + found = _find_fits(step) + if not found: + raise TypeError(f"{type(step).__name__}.fit_many() requires at least one Fit") + for path, fit in found.items(): + if cohort is not None and fit.cohort not in (None, cohort): + raise ValueError( + f"{type(fit).__name__} at {path or '.'} has cohort " + f"{fit.cohort!r}, incompatible with {cohort!r}" + ) + if cohort is not None: + uid = cohort + elif fit.cohort is not None: + uid = fit.cohort + elif uids: + uid = _fingerprint(fit._cohort_uids(uids)) + elif fit._declared is not None: + uid = fit._declared else: - uid = _fingerprint(fit._cohort_uids(cohort._uids)) - # the config that ran: this copy, or the one it resolved to - memo = tp.cast(Fit | None, fit._resolution_cache) - ran = fit if fit.cohort is not None else memo - if ran is not None and ran.cohort != uid: - raise RuntimeError( - f"{type(fit).__name__} at {path or '.'} already ran on cohort " - f"{ran.cohort!r}: refitting in place is not supported, run " - "clone() to fit another cohort" - ) - fit._declared = uid - cohort.fitted_by.append(f"{path or '.'}({uid})") + raise ValueError("fit_many() without values requires a named cohort") + bound = fit.cohort if fit.cohort is not None else fit._declared + if bound is not None and bound != uid: + raise RuntimeError( + f"{type(fit).__name__} at {path or '.'} already uses cohort " + f"{bound!r}: refitting in place is not supported, use clone() " + f"for {uid!r}" + ) + fit._declared = uid class Fit(Step): @@ -98,12 +86,12 @@ def _run(self, value): # one item return value - self.fitted norm = Normalize(infra={"backend": "Cached", "folder": cache}) - norm.run_many(FitCohort(train)) # fits on these, then transforms them - norm.run_many(test) # transforms with the same artifact + norm.fit_many(train) # fits on these, then transforms them + norm.run_many(test) # transforms with the same artifact - Only a :class:`FitCohort` is fitted on; its name -- :attr:`cohort`, else the - fingerprint of its items -- scopes the artifact and every downstream cache. - Another cohort or upstream takes another config (:meth:`clone`). + Only :meth:`~exca.steps.Step.fit_many` fits; its cohort name, else the fingerprint + of its items, scopes the artifact and every downstream cache. Another cohort or + upstream takes another config (:meth:`clone`). The fit runs where the step is dispatched from, ahead of a backend splitting the batch, and is cached under ``infra``. The upstream is read twice (once for the @@ -152,7 +140,7 @@ def fitted(self) -> tp.Any: """The artifact :meth:`_fit` produced, for :meth:`_run` to transform with.""" if self._fitted_for is None: raise RuntimeError( - f"{type(self).__name__} is not fitted: run it on a FitCohort first" + f"{type(self).__name__} is not fitted: call fit_many() first" ) return self._fitted @@ -170,8 +158,8 @@ def _resolve_artifact(self, batch: items.StepItems) -> None: uid = self.cohort if uid is None: raise RuntimeError( - f"{kind} has no cohort to fit on or to read back: run it on a " - "FitCohort, or set its 'cohort' name" + f"{kind} has no cohort to fit on or to read back: call fit_many(), " + "or set its 'cohort' name" ) mode = backends._fold_modes(batch._mode, backends._effective_mode(self)) # cohort cleared: all cohorts of this Fit share one folder, one entry each @@ -196,7 +184,7 @@ def _resolve_artifact(self, batch: items.StepItems) -> None: hint = "drop the force mode" if mode == "force" else "check its name" raise RuntimeError( f"{kind} must fit cohort {uid!r} but was handed no items to fit " - f"on: run it on a FitCohort, or {hint}" + f"on: pass values to fit_many(), or {hint}" ) # counts, not uids: a step re-keying items has its own uid space handed = len(set(batch.uids)) diff --git a/exca/steps/helpers.py b/exca/steps/helpers.py index e9179284..65ad8ab5 100644 --- a/exca/steps/helpers.py +++ b/exca/steps/helpers.py @@ -11,7 +11,7 @@ import pydantic -from . import fit, identity, items, utils +from . import identity, items, utils from .base import Step @@ -218,6 +218,14 @@ def run(self, value: tp.Any = identity.NoValue()) -> None: self.run_many([value]) def run_many( # type: ignore[override] - self, values: tp.Iterable[tp.Any] | fit.FitCohort + self, values: tp.Iterable[tp.Any] ) -> list[None]: return list(super().run_many(values)) + + def fit_many( # type: ignore[override] + self, + values: tp.Iterable[tp.Any] = (), + *, + cohort: str | None = None, + ) -> list[None]: + return list(super().fit_many(values, cohort=cohort)) diff --git a/exca/steps/test_fit.py b/exca/steps/test_fit.py index b5aeeeaf..1215a9cb 100644 --- a/exca/steps/test_fit.py +++ b/exca/steps/test_fit.py @@ -54,13 +54,13 @@ def test_fit_cohort_then_novel_items(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} step = _SumOffset(infra=infra) exca.utils.recursive_freeze(step) # as an enclosing config would - output = list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) + output = list(step.fit_many([1.0, 2.0, 3.0])) assert output == [7.0, 8.0, 9.0], "cohort items use their fitted sum" assert step.run(10.0) == 16.0, "novel item must use the fitted sum" assert len(step._fits) == 1, "a second call must not refit" read_back = _SumOffset(infra=infra) - output = list(read_back.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) + output = list(read_back.fit_many([1.0, 2.0, 3.0])) assert output == [7.0, 8.0, 9.0], "cached artifact must reproduce outputs" assert not read_back._fits, "the artifact must be read back from the cache" @@ -69,57 +69,62 @@ def test_fit_cohort_then_novel_items(tmp_path: Path) -> None: infra["mode"] = "force" forced = _SumOffset(infra=infra) - forced.run_many(steps.FitCohort([1.0, 2.0, 3.0])) + forced.fit_many([1.0, 2.0, 3.0]) assert forced._fits == [[1.0, 2.0, 3.0]], "force must refit instead of reading back" def test_fit_cohort_marker_is_request_scoped(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} upstream = conftest.Mult(coeff=2, infra=infra) - upstream.run_many(steps.FitCohort([1.0, 2.0])) + steps.Chain(steps=[upstream, _SumOffset(infra=infra)]).fit_many([1.0, 2.0]) assert upstream._output_items is not None, "Fitting should warm cache as well" chain = steps.Chain(steps=[upstream, _SumOffset(cohort="train", infra=infra)]) with pytest.raises(RuntimeError, match="handed no items"): chain.run_many([1.0, 2.0]) - output = list(chain.run_many(steps.FitCohort([1.0, 2.0]))) + output = list(chain.fit_many([1.0, 2.0])) assert output == [8.0, 10.0], "current request's cohort marker must reach Fit" def test_named_cohort(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} - step = _SumOffset(infra=infra, cohort="train") - output = list(step.run_many(steps.FitCohort([1.0, 3.0]))) + step = _SumOffset(infra=infra) + output = list(step.fit_many([1.0, 3.0], cohort="train")) assert output == [5.0, 7.0], "named cohort must fit and transform" - assert step.cohort == "train", "a declared cohort must not rename it" + assert step.cohort is None, "a declared cohort must not alter the config" configured = _SumOffset(infra=infra, cohort="train") assert configured.run(10.0) == 14.0, "the config name must recover it" assert not configured._fits, "recovering a named artifact must not refit" + loaded = _SumOffset(infra=infra) + assert list(loaded.fit_many(cohort="train")) == [], "loading has no item outputs" + assert loaded.run(10.0) == 14.0, "fit_many must bind a prefitted cohort" + assert not loaded._fits, "loading a named artifact must not refit" + with pytest.raises(RuntimeError, match="must fit cohort 'test'"): - _SumOffset(infra=infra, cohort="test").run(10.0) + _SumOffset(infra=infra).fit_many(cohort="test") + with pytest.raises(ValueError, match="requires a named cohort"): + _SumOffset(infra=infra).fit_many() infra["mode"] = "force" with pytest.raises(RuntimeError, match="drop the force mode"): - _SumOffset(infra=infra, cohort="train").run(10.0) + _SumOffset(infra=infra).fit_many(cohort="train") scaled = step.clone(scale=10) - output = list(scaled.run_many(steps.FitCohort([1.0, 3.0]))) + output = list(scaled.fit_many([1.0, 3.0], cohort="train")) assert output == [41.0, 43.0], "cloned config must fit its own artifact" def test_retry_of_a_failed_fit(tmp_path: Path) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} with pytest.raises(ValueError, match="Triggered an error"): - _SumOffset(infra=infra, cohort="train", broken=True).run_many( - steps.FitCohort([1.0, 3.0]) - ) + _SumOffset(infra=infra, cohort="train", broken=True).fit_many([1.0, 3.0]) infra["mode"] = "retry" with pytest.raises(RuntimeError, match="was handed no items"): _SumOffset(infra=infra, cohort="train").run(10.0) fixed = _SumOffset(infra=infra, cohort="train") - output = list(fixed.run_many(steps.FitCohort([1.0, 3.0]))) + output = list(fixed.fit_many([1.0, 3.0])) assert output == [5.0, 7.0], "retry must allow the repaired fit" @@ -138,11 +143,11 @@ def test_cohort_uids( ) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} step = Variant(infra=infra) - step.run_many(steps.FitCohort([1.0, 1.0, 4.0])) + step.fit_many([1.0, 1.0, 4.0]) assert step._fits == [fit_values], "the fit sees the uids it asked for" reordered = Variant(infra=infra) - reordered.run_many(steps.FitCohort([4.0, 1.0, 1.0])) + reordered.fit_many([4.0, 1.0, 1.0]) assert reordered._fits == reorder_fits, "only a sequence cohort is order-sensitive" @@ -157,7 +162,7 @@ def test_fit_with_distributed_items( chain = steps.Chain( steps=[conftest.Mult(coeff=2, infra=infra if on_upstream else local), fit] ) - output = list(chain.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) + output = list(chain.fit_many([1.0, 2.0, 3.0])) assert output == [14.0, 16.0, 18.0], "distribution must preserve cohort results" assert fit._fits == [[2.0, 4.0, 6.0]], "the fit must see the whole cohort, once" @@ -169,16 +174,16 @@ def test_fit_with_distributed_items( def test_fit_under_distributed_chain(tmp_path: Path, backend: str, max_jobs: int) -> None: infra: tp.Any = {"folder": tmp_path, "backend": backend, "max_jobs": max_jobs} chain = steps.Chain(steps=[_SumOffset()], infra=infra) - cohort = steps.FitCohort([1.0, 2.0, 3.0, 4.0]) + cohort = [1.0, 2.0, 3.0, 4.0] if max_jobs == 1: # one worker holds the whole cohort, so it can fit - output = list(chain.run_many(cohort)) + output = list(chain.fit_many(cohort)) assert output == [11.0, 12.0, 13.0, 14.0], "one task keeps the cohort" fit = _SumOffsetSequence() sequence = steps.Chain( steps=[conftest.Mult(coeff=2), fit, conftest.Add(value=1)], infra=infra, ) - output = list(sequence.run_many(steps.FitCohort([1.0, 1.0, 4.0]))) + output = list(sequence.fit_many([1.0, 1.0, 4.0])) assert output == [15.0, 15.0, 21.0], "enclosing backend preserves the sequence" if backend == "ThreadPool": # worker shares the observable Fit instance assert fit._fits == [[2.0, 2.0, 8.0]], "fitting keeps repeated values" @@ -187,13 +192,13 @@ def test_fit_under_distributed_chain(tmp_path: Path, backend: str, max_jobs: int fit._runs.clear() sequence.lookup(4.0).clear_cache(recursive=False) - output = list(sequence.run_many(steps.FitCohort([1.0, 1.0, 4.0]))) + output = list(sequence.fit_many([1.0, 1.0, 4.0])) assert output == [15.0, 15.0, 21.0], "partial cache keeps cohort semantics" if backend == "ThreadPool": assert fit._runs == [8.0], "only the missing cache key is transformed" else: with pytest.raises(Exception, match="too few to fit"): - list(chain.run_many(cohort)) + list(chain.fit_many(cohort)) class _ScatterValues(Scatter): @@ -211,18 +216,18 @@ def test_fit_inside_scatter(tmp_path: Path) -> None: expected = [{"a": 7.0, "b": 8.0}, {"c": 9.0}] cache: tp.Any = {"folder": tmp_path / "cached", "backend": "Cached"} scatter = _ScatterValues(body=_SumOffset(infra=cache)) - output = list(scatter.run_many(steps.FitCohort(cohort))) + output = list(scatter.fit_many(cohort)) assert output == expected, "a complete Scatter must fit over all branches" outer: tp.Any = {**cache, "backend": "ThreadPool", "max_jobs": 2} chain = steps.Chain(steps=[scatter.clone()], infra=outer) - output = list(chain.run_many(steps.FitCohort(cohort))) + output = list(chain.fit_many(cohort)) assert output == expected, "an incomplete Scatter may reuse an existing artifact" cold: tp.Any = {**outer, "folder": tmp_path / "cold"} chain = steps.Chain(steps=[_ScatterValues(body=_SumOffset())], infra=cold) with pytest.raises(RuntimeError, match="handed no items"): - list(chain.run_many(steps.FitCohort(cohort))) + list(chain.fit_many(cohort)) @pytest.mark.parametrize("nested", [False, True]) @@ -232,11 +237,11 @@ def test_a_config_fits_one_cohort(tmp_path: Path, nested: bool) -> None: if nested: # the chain keys its cache before the fit resolves step = steps.Chain(steps=[step], infra=infra) - output = list(step.run_many(steps.FitCohort([1.0, 2.0, 3.0]))) + output = list(step.fit_many([1.0, 2.0, 3.0])) assert output == [7.0, 8.0, 9.0], "one config must fit its first cohort" with pytest.raises(RuntimeError, match="refitting in place"): - step.run_many(steps.FitCohort([1.0, 4.0])) - output = list(step.clone().run_many(steps.FitCohort([1.0, 4.0]))) + step.fit_many([1.0, 4.0]) + output = list(step.clone().fit_many([1.0, 4.0])) assert output == [6.0, 9.0], "clone must fit another cohort independently" @@ -245,15 +250,11 @@ def test_fit_clone_for_another_upstream(tmp_path: Path, cached: bool) -> None: infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} if cached else None fit = _SumOffset(infra=infra, cohort=None if cached else "train") cohort = [1.0, 2.0] # same items, so the same cohort uid names both fits - first = fit.run_many(steps.FitCohort(cohort)) + first = fit.fit_many(cohort) with pytest.raises(RuntimeError, match="already holds artifact"): - steps.Chain(steps=[conftest.Mult(coeff=10), fit]).run_many( - steps.FitCohort(cohort) - ) + steps.Chain(steps=[conftest.Mult(coeff=10), fit]).fit_many(cohort) cloned = fit.clone() - second = steps.Chain(steps=[conftest.Mult(coeff=10), cloned]).run_many( - steps.FitCohort(cohort) - ) + second = steps.Chain(steps=[conftest.Mult(coeff=10), cloned]).fit_many(cohort) assert list(first) == [4.0, 5.0], "the original must keep its first upstream" assert list(second) == [40.0, 50.0], "the clone must use its new upstream" assert fit._fits == [[1.0, 2.0]], "the original must fit unscaled values" @@ -270,7 +271,7 @@ def test_fit_folder_structure(tmp_path: Path) -> None: conftest.Add(value=1, infra=infra), ] ) - output = list(chain.run_many(steps.FitCohort([1.0, 1.0, 4.0]))) + output = list(chain.fit_many([1.0, 1.0, 4.0])) assert output == [55.0, 55.0, 61.0], "stacked Fits must use both artifacts" mult = "type=Mult-b9b7a7a5" @@ -301,13 +302,12 @@ def item_uid(self, value: tp.Any) -> str | None: _Keyed(infra=infra), ] ) - output = sweep.run_many(steps.FitCohort([1.0, 1.0, 4.0])) + output = sweep.fit_many([1.0, 1.0, 4.0]) assert output == [None] * 3, "each Fit variant must consume the cohort" - variants = tp.cast(list[_SumOffset], list(sweep.steps)) - fits = [v._fits for v in variants] + fits = [v._fits for v in sweep.steps if isinstance(v, _SumOffset)] expected = [[[1.0, 4.0]], [[1.0, 1.0, 4.0]], [[1.0, 4.0]]] assert fits == expected, "each variant must fit on its selected uid sequence" - output = [v.lookup(4.0).result() for v in variants] + output = [v.lookup(4.0).result() for v in sweep.steps] assert output == [9.0, 10.0, 9.0], "variants cache their own results" @@ -324,13 +324,9 @@ def _resolve_step(self) -> steps.Step: scale=conftest.Mult(coeff=2), offset=offset, wrapped=_ResolvedOffset() ) ) - cohort = steps.FitCohort([1.0, 3.0]) - output = list(chain.run_many(cohort)) + output = list(chain.fit_many([1.0, 3.0])) assert output == [34.0, 38.0], "every discovered Fit contributes its artifact" - named = [n.split("(")[0] for n in cohort.fitted_by] - assert named == ["steps.offset", "steps.wrapped"], "even one a resolution builds" assert offset.cohort is None, "the config handed over is left as it was" - vain = steps.FitCohort([1.0]) - conftest.Mult(coeff=2, infra=infra).run_many(vain) - assert not vain.fitted_by, "no Fit to name it, and nothing to fit" + with pytest.raises(TypeError, match="requires at least one Fit"): + conftest.Mult(coeff=2, infra=infra).fit_many([1.0])