diff --git a/CHANGELOG.md b/CHANGELOG.md index 51095e13..eefdd229 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,9 @@ ## [Unreleased] -- `DiscriminatedModel`: optimized look-up. +- `DiscriminatedModel`: optimized look-up. [#313] +- `steps`: fixed nested infra claim deadlock. [#323] + ## 0.5.29 - 26-07-28 diff --git a/exca/steps/backends.py b/exca/steps/backends.py index ea15f65a..62a92fa5 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 @@ -282,15 +283,18 @@ def result(self) -> tp.Any: raise RuntimeError(f"No cached entry for {self._uid}") +# cache entries claimed by the running work or its ancestors +_HELD_ENTRIES: contextvars.ContextVar[frozenset[tuple[str, str]]] = ( + contextvars.ContextVar("exca_held_entries", default=frozenset()) +) + + @dataclasses.dataclass class CoordinationInfo: - """Driver-only per-run state for a ``ComputeBatch``; stripped from the - worker pickle (``ComputeBatch.__getstate__``). - """ - mode: identity.ModeType = "cached" upstream: tuple[Step, ...] = () # this step + everything before it claim: inflight.InflightClaim | None = None + held_entries: frozenset[tuple[str, str]] = frozenset() # {(folder, uid),...} @dataclasses.dataclass @@ -303,8 +307,20 @@ class ComputeBatch: items: items.StepItems info: CoordinationInfo = dataclasses.field(default_factory=CoordinationInfo) + def claimed_uids(self) -> list[str]: + """This batch's uids claimed by its own session (entries held by an ancestor + are excluded, so their rows keep pointing at the ancestor's job).""" + claim = self.info.claim + if claim is None: + raise RuntimeError(f"batch was never claimed: {self.paths.step_uid}") + owned = set(claim.uids) + return [uid for uid in self.items.uids if uid in owned] + def __getstate__(self) -> dict[str, tp.Any]: - return {**self.__dict__, "info": CoordinationInfo()} + return { + **self.__dict__, + "info": CoordinationInfo(held_entries=self.info.held_entries), + } def select(self, uids: tp.Sequence[str]) -> ComputeBatch: """Sub-batch over *uids*, sharing step/paths/cache; copies ``info`` @@ -335,9 +351,10 @@ def run_and_cache(self) -> None: folder = self.cache_dict.folder if folder is not None: folder.mkdir(parents=True, exist_ok=True) - result_items = self.step._run_items(self.items) written_uids: list[str] = [] + token = _HELD_ENTRIES.set(self.info.held_entries) try: + result_items = self.step._run_items(self.items) with self.cache_dict.write(): for i, result in enumerate(result_items): uid = self.items.uids[i] @@ -366,6 +383,8 @@ def run_and_cache(self) -> None: for uid in inflight: reg.record(uid, e, tb) raise + finally: + _HELD_ENTRIES.reset(token) def _multi_run_and_cache(batches: list[ComputeBatch]) -> None: @@ -554,12 +573,16 @@ def _clear_caches( # Other backends may have left inflight rows for this step folder. if paths.step_folder.exists(): try: + held = _HELD_ENTRIES.get() + folder_key = str(paths.step_folder) with inflight.InflightRegistry(paths.step_folder) as reg: info = reg.get(uids) jobs: dict[str, str] = {} for uid, worker in info.items(): if worker.job_id is None or worker.job_folder is None: continue # not submitit + if (folder_key, uid) in held: + continue # an ancestor's job — cancelling it kills us # Slurm array tasks share a scheduler job; avoid per-task cancels. job_id = worker.job_id.split("_", 1)[0] jobs[job_id] = worker.job_folder @@ -617,6 +640,7 @@ def _claim(self, cbatches: list[ComputeBatch]) -> _Claimed: if len(set(step_uids)) != len(step_uids): raise ValueError(f"one batch per step_uid required, got {step_uids}") claimed = _Claimed() + held = _HELD_ENTRIES.get() try: # sort by step_uid: concurrent dispatches claim in the same order for cb in sorted(cbatches, key=lambda cb: cb.paths.step_uid): @@ -628,10 +652,16 @@ def _claim(self, cbatches: list[ComputeBatch]) -> _Claimed: reg: inflight.InflightRegistry | None = None if self._concurrent: reg = inflight.InflightRegistry(cb.paths.step_folder) + # ancestors already hold their entries: claiming them self-deadlocks + folder_key = str(cb.paths.step_folder) + request = {u for u in pending if (folder_key, u) not in held} cb = cb.select(list(pending)) cb.info.claim = claimed.stack.enter_context( - inflight.inflight_session(reg, set(pending)) + inflight.inflight_session(reg, request) ) + cb.info.held_entries = held + if reg is not None: # registry-less claims hold nothing to inherit + cb.info.held_entries |= {(folder_key, u) for u in cb.info.claim.uids} claimed.batches.append(cb) claimed.ready = [ n @@ -769,7 +799,7 @@ 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) + batch.info.claim.record_worker_info(job, uids=batch.claimed_uids()) folder = batch.paths.step_folder by_folder.setdefault(folder, {})[job.job_id] = batch.items.uids for folder, records in by_folder.items(): @@ -971,7 +1001,7 @@ def _submit_pool( 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 72b3296f..29019929 100644 --- a/exca/steps/test_backends.py +++ b/exca/steps/test_backends.py @@ -365,3 +365,11 @@ def run(value: float, fail: bool, mode: str) -> float: # retry both; a uid-only _recomputed would make the 2nd re-raise the 1st's error assert run(1.0, fail=False, mode="retry") == 3.0 # 2 + 1 assert run(5.0, fail=False, mode="retry") == 7.0 # 2 + 5 + + +def test_nested_dispatch_on_shared_cell(tmp_path: Path) -> None: + infra: tp.Any = {"backend": "LocalProcess", "folder": tmp_path} + inner = Chain(steps=[conftest.Mult(coeff=3.0, infra=infra)], infra=infra) + chain = Chain(steps=[conftest.Add(value=1.0), inner], infra=infra) + # identity flattens recursively: the 3 infras share one cache entry + assert chain.run(1.0) == 6.0