diff --git a/CHANGELOG.md b/CHANGELOG.md index de0202d4..988cad59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## [Unreleased] +- `steps`: `Fit` primitive — fit one artifact over a cohort of items and transform the items accordingly. + ## 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..ee83b918 100644 --- a/docs/steps/items.md +++ b/docs/steps/items.md @@ -133,6 +133,37 @@ 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. `fit_many` fits; `run_many` 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.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 +``` + +`_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 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 +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..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: @@ -90,3 +90,12 @@ Helpers .. autoclass:: exca.steps.helpers.Func :members: + + +Fitting over items +------------------ + +.. autoclass:: exca.steps.Fit + :show-inheritance: + :members: fitted, ARTIFACT_CACHE_TYPE + :private-members: _fit, _cohort_uids 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..ede0aed0 100644 --- a/exca/steps/__init__.py +++ b/exca/steps/__init__.py @@ -29,3 +29,4 @@ from . import helpers as helpers from .base import Chain as Chain from .base import Step as Step +from .fit import Fit as Fit diff --git a/exca/steps/backends.py b/exca/steps/backends.py index ea15f65a..9a539ec3 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 @@ -199,6 +200,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 @@ -306,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) @@ -323,9 +340,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, ) @@ -335,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) @@ -419,6 +439,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.""" @@ -445,6 +470,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, *, @@ -467,7 +495,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 @@ -578,9 +606,35 @@ 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. + """ + group = self._active_group() + if group is not None: # nested sweep: the outer group submits it + yield + return + group = [] + groups = _ACTIVE_GROUPS.get() + token = _ACTIVE_GROUPS.set({**groups, id(self): group}) + try: + yield + finally: + _ACTIVE_GROUPS.reset(token) + with self._claim(group) as claimed: + if claimed.ready: + self._execute(claimed.ready) + def _run(self, step: Step, batch: items.StepItems) -> items.StepItems: """Execute *step* for uncached items, caching per uid.""" cbatch = self._prepare(step, batch) + 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: self._execute(claimed.ready) @@ -628,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 = [ @@ -648,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: @@ -667,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. @@ -677,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.""" @@ -709,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) @@ -749,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( @@ -769,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 @@ -901,6 +967,10 @@ def _run(self, step: Step, batch: items.StepItems) -> items.StepItems: lazy carrier instead of blocking. """ cbatch = self._prepare(step, batch) + 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 try: @@ -912,13 +982,13 @@ 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 - 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, ) @@ -931,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): @@ -949,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") @@ -958,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/base.py b/exca/steps/base.py index 6c3f7cd5..124e2b5d 100644 --- a/exca/steps/base.py +++ b/exca/steps/base.py @@ -245,7 +245,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.""" @@ -255,7 +257,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 @@ -400,12 +402,34 @@ def run_many(self, values: tp.Iterable[tp.Any]) -> items.StepItems: 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 = items.StepItems(source=dict(zip(uids, values)), uids=uids) 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 new file mode 100644 index 00000000..7c462252 --- /dev/null +++ b/exca/steps/fit.py @@ -0,0 +1,218 @@ +# 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 + +from exca import utils as xkutils + +from . import backends, identity, items, utils +from .base import Step + + +def _fingerprint(ordered: tp.Sequence[str]) -> str: + """Identity of a cohort that was not named, as ``,``.""" + digest = hashlib.sha256() + for uid in ordered: + digest.update(uid.encode("utf8")) + digest.update(b"\0") + return f"{digest.hexdigest()[:8]},{len(ordered)}" + + +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, uids: tp.Sequence[str], cohort: str | None = None +) -> None: + """Name the cohort in every ``Fit`` of *step*, privately, before anything runs.""" + 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: + 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): + """Fit one artifact over a cohort of items, then transform each item (N->1->N). + + .. warning:: Experimental -- API may change. + + Example:: + + class Normalize(Fit): + def _fit(self, values): # the cohort, streamed + return np.stack(list(values)).mean(0) + + def _run(self, value): # one item + return value - self.fitted + + norm = Normalize(infra={"backend": "Cached", "folder": cache}) + norm.fit_many(train) # fits on these, then transforms them + norm.run_many(test) # transforms with the same artifact + + 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 + fit, once per item), so give an expensive upstream its own ``infra``. + + 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 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. + """ + + ARTIFACT_CACHE_TYPE: tp.ClassVar[str | None] = "Auto" + + cohort: str | None = None + + _fitted: tp.Any = pydantic.PrivateAttr(None) + _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: + 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* re-iterates the cohort (one upstream read per pass). + """ + 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: call fit_many() first" + ) + 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 + 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 (no-op if held).""" + 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: 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 + 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) + fitted_for = (uid, identity.step_uid([*upstream, artifact])) + 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 + 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: 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)) + if handed < batch._total_size: + raise RuntimeError( + 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 + ) + # dispatch, not lookup: goes through infra's mode and caching + self._fitted = next(iter(artifact._dispatch(carrier))) + self._fitted_for = fitted_for + + +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(self.owner._cohort_uids(batch.uids))) diff --git a/exca/steps/helpers.py b/exca/steps/helpers.py index 3d68939c..65ad8ab5 100644 --- a/exca/steps/helpers.py +++ b/exca/steps/helpers.py @@ -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}" ) + # 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]: return [] # no identity of its own @@ -198,27 +196,36 @@ 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: + # 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 " + "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_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: 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] + ) -> 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/items.py b/exca/steps/items.py index 33fd8eef..e2b10a8d 100644 --- a/exca/steps/items.py +++ b/exca/steps/items.py @@ -140,6 +140,8 @@ def __init__( self._upstream = tuple(upstream) self._pending = tuple(pending) self._mode = mode + self._total_size = len(set(self.uids)) + self._cohort = False def __len__(self) -> int: return len(self.uids) @@ -148,14 +150,25 @@ 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 *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, + } + 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.""" - 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 +182,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 529c32c8..15f25c07 100644 --- a/exca/steps/patterns.py +++ b/exca/steps/patterns.py @@ -229,17 +229,18 @@ 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, 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 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_backends.py b/exca/steps/test_backends.py index 72b3296f..5ed2e157 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 @@ -286,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}") @@ -346,6 +372,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 new file mode 100644 index 00000000..1215a9cb --- /dev/null +++ b/exca/steps/test_fit.py @@ -0,0 +1,332 @@ +# 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 typing as tp +from pathlib import Path + +import pydantic +import pytest + +import exca + +from .. import steps +from . import conftest +from .patterns import Scatter + + +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) + _runs: 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 self.scale * sum(vals) + + def _run(self, value: float) -> float: + self._runs.append(value) + return value + self.fitted + + +class _SumOffsetSequence(_SumOffset): + 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: + 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.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.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" + + with pytest.raises(RuntimeError, match="no cohort"): + _SumOffset(infra=infra).run(10.0) + + infra["mode"] = "force" + forced = _SumOffset(infra=infra) + 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) + 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.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) + 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 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).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).fit_many(cohort="train") + scaled = step.clone(scale=10) + 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).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.fit_many([1.0, 3.0])) + assert output == [5.0, 7.0], "retry must allow the repaired fit" + + +@pytest.mark.parametrize( + "Variant,fit_values,reorder_fits", + [ + (_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[_SumOffset], + fit_values: list[float], + reorder_fits: list[list[float]], +) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + step = Variant(infra=infra) + 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.fit_many([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 = _SumOffset(infra=local if on_upstream else infra) + chain = steps.Chain( + steps=[conftest.Mult(coeff=2, infra=infra if on_upstream else local), fit] + ) + 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" + + +@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 = [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.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.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" + 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.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.fit_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)) + 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.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.fit_many(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"} + step: steps.Step = _SumOffset(infra=infra) + if nested: # the chain keys its cache before the fit resolves + step = steps.Chain(steps=[step], infra=infra) + + 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.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" + + +@pytest.mark.parametrize("cached", [False, True]) +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.fit_many(cohort) + with pytest.raises(RuntimeError, match="already holds artifact"): + steps.Chain(steps=[conftest.Mult(coeff=10), fit]).fit_many(cohort) + cloned = fit.clone() + 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" + assert cloned._fits == [[10.0, 20.0]], "the clone must fit scaled values" + + +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), + _SumOffset(infra=infra), + _SumOffsetSequence(infra=infra), + conftest.Add(value=1, infra=infra), + ] + ) + 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" + 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=_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" + + +def test_fit_variants_in_parallel(tmp_path: Path) -> None: + 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=[ + _SumOffset(infra=infra), + _SumOffsetSequence(infra=infra), + _Keyed(infra=infra), + ] + ) + output = sweep.fit_many([1.0, 1.0, 4.0]) + assert output == [None] * 3, "each Fit variant must consume the cohort" + 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 sweep.steps] + assert output == [9.0, 10.0, 9.0], "variants cache their own results" + + +def test_cohort_names_every_fit(tmp_path: Path) -> None: + infra: tp.Any = {"folder": tmp_path, "backend": "Cached"} + + 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, wrapped=_ResolvedOffset() + ) + ) + output = list(chain.fit_many([1.0, 3.0])) + assert output == [34.0, 38.0], "every discovered Fit contributes its artifact" + assert offset.cohort is None, "the config handed over is left as it was" + + with pytest.raises(TypeError, match="requires at least one Fit"): + conftest.Mult(coeff=2, infra=infra).fit_many([1.0]) diff --git a/exca/steps/test_helpers.py b/exca/steps/test_helpers.py index 489f6634..bdb10ab0 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 ResolvedAdd(Step): + value: float + + def _resolve_step(self) -> Step: + return conftest.Add(value=self.value, infra=infra) + + 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 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/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]