diff --git a/CHANGELOG.md b/CHANGELOG.md index 4205cb3..46380e1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,45 @@ # Changelog +## 0.5.0 (2026-10-23) + +Two bugs that hung or crashed cellmap-analyze 0.6 on /nrs (salamander's and c-elegans' +runs, 2026-10-22), and tensorstore's threads sized to the job. No placement change and no +unit table change. + +- A read whose first attempt timed out raised `IndexError` instead of retrying: the + warning took the first line of the error's message, and tensorstore's `TimeoutError()` + has none. The warning now names the error's type when the message is empty, and the + error after the last attempt names the operation, the path and the attempts. +- Every context the package makes or opens a store on is watched: once a second a daemon + thread gives its decoding and file I/O pools one no-op task each (an in-memory copy and + a local read of a missing key; no store is touched). tensorstore up to at least 0.1.85 + can strand a queued task: a pool worker idle for 20 ms exits as the task arrives, and + the task never runs (fixed upstream in google/tensorstore ed51718085, in no release + yet). With `data_copy` and `file_io` capped at 1, as cellmap-analyze's workers were, + dozens of workers per run waited forever in `write_index` with no tensorstore thread + alive, and reads timed out into the `IndexError` above. One more task on the same pool + wakes it, which salamander's and c-elegans' sessions proved on live stuck workers, so a + stranded task now runs within about a second, at any concurrency, whoever waits on it: + worldarray, `vol.ts` directly or `await`. Reproduced on purpose on tensorstore 0.1.75, + 0.1.78 and 0.1.85 (opt-in `tests/test_lost_wakeup.py`): cellmap-analyze's block loop + gave 29 `IndexError`s and 55 writes stuck past 20 s on 0.4.0, none on 0.5.0, slowest + 0.98 s. Spec open question 14 says when this goes. +- `write_index`, `write` and `write_chunk` are timed out and retried as reads are + (Behaviour 6), by a new `retries=` argument or the handle's `retries`. A write used to + wait on its future with no timeout at all. +- Without `context`, the shared context's decoding and file I/O pools are limited to the + new `wa.thread_budget()`: `OMP_NUM_THREADS` when set (dask sets 1 in each worker), else + the job's slots (`LSB_DJOB_NUMPROC`, `SLURM_CPUS_PER_TASK`, `NSLOTS`), else the CPUs the + process may run on. A server that wants more threads than its job's slots asks for more + slots or passes its own context. + tensorstore alone sizes both to every CPU on the node, whatever `bsub -n` asked for, so + a dask worker's tensorstore now takes the same one thread as its numpy (pinned to one + CPU, tensorstore's default context started 43 to 45 pool threads, the capped one 4 to 5). + Against cellmap-analyze 0.5.0's own block I/O on NFS (128^3 uint64 label blocks, each + one chunk), its 0.6.1 on this release reads a block in 26 ms instead of 40 and writes it + in 45 ms instead of 62 from one process at one thread per pool, and about 1.15 times + faster from 8. Speed gains grow with blocks of several chunks and more threads. + ## 0.4.0 (2026-10-22) cellmap-analyze's three asks, so it can call worldarray directly with no wrapper, and two diff --git a/README.md b/README.md index 556360c..bae857c 100644 --- a/README.md +++ b/README.md @@ -86,9 +86,12 @@ s1 = wa.open_pyramid("out.zarr", mode="r+").add_level("s1", factors=2) # place `transform` applied before padding (or `post_transform` after resampling), and resampling by nearest neighbour or linear interpolation anchored where you say. Adjoining block reads agree with one read across both. -- **Behaves on a cluster**: one shared tensorstore context per process or a capped one per - worker, retries with backoff and size-aware timeouts on transient store errors, and a - cancelled future on timeout. +- **Behaves on a cluster**: one shared tensorstore context per process, its threads sized + from `OMP_NUM_THREADS` or the LSF job's slots instead of the node's cores, or a capped one + per worker; reads and writes retried with backoff and size-aware timeouts on transient + store errors, a cancelled future on timeout, and every context's thread pools nudged once + a second against a tensorstore bug that can leave a queued task with no thread, so a + stuck operation recovers or fails loudly instead of hanging the job. - **Carries the geometry**: `Box` with `intersect`, `union`, `contains`, `grow`, `shift`, `snap` to a grid, and the voxel indices a world box selects. - **Writes placed arrays** (0.2): the OME-NGFF 0.4 or 0.5 document from voxel size, centre @@ -147,6 +150,7 @@ wa.metadata(path, level=None, units=None, open=True) -> ArrayMeta # open=False wa.detect(path, missing_ok=False) -> Node # format, array or group, raw attrs wa.open_tensorstore(path, level=None, context=None) wa.shared_context(cache_bytes=0) ; wa.capped_context(cache_bytes=0, concurrency=1, network=True) ; wa.Retry(...) +wa.thread_budget() # the shared context's thread limit: OMP_NUM_THREADS, the job's slots, or CPUs wa.list_chunk_keys(path, separator=".") # chunk indices from keys alone, no metadata needed wa.Box ; wa.units.TABLE @@ -174,7 +178,8 @@ wa.units.abbreviation("nanometer") ; wa.store.join(path, "s0") ; wa.store.parent `read(box, voxel_size=, padding=, method=, mask=, anchor=, transform=, post_transform=, select=, strict=)`, `read_index(begin, end, padding=, mask=, select=, transform=)`, `read_onto(grid, begin=, end=, ...)`, `index_box(box)`, `chunks_covering(begin, end, write=)`, `list_chunks()`, `chunk_key(index)`, `parse_chunk_key(key)`, -`write_index(begin, end, data, partial=)` and `write(box, data, partial=)` (`data` an array or a scalar), `reopen(context=)`. +`write_index(begin, end, data, partial=, retries=)` and `write(box, data, partial=, retries=)` (`data` an array or a +scalar), `reopen(context=)`. A box is a `(begin, end)` pair of floats in your units, half-open per axis, over the spatial axes or every axis; any object with `.begin`/`.end` or `.offset`/`.shape` also diff --git a/docs/spec.md b/docs/spec.md index c656faa..927eb34 100644 --- a/docs/spec.md +++ b/docs/spec.md @@ -230,7 +230,8 @@ in integer Rois builds them from `corner`, `voxel_size` and `shape` in its own s where it also chooses the integer scale factor that makes the corner whole. `shared_context(cache_bytes)` is one `tensorstore.Context` per budget for the process, so -the levels of a pyramid and repeat opens share one cache pool; the default budget is 0. +the levels of a pyramid and repeat opens share one cache pool; the default budget is 0. Its +decoding and file I/O pools are sized by `thread_budget()` (0.5, Behaviour 7). `capped_context` is mesh-n-bone's worker-process context (every tensorstore concurrency resource limited, cache off); it returns a **new** context per call, as those workers inline the spec into every open today, and `capped_context_spec` exports the dict for @@ -353,6 +354,18 @@ vol.source_translation # the stored level's translation in `units`, beside so wa.create(..., attrs={...}, exists="adopt") # an adopt also compares the keys of `attrs` (Behaviour 11) ``` +### 0.5: writes that cannot hang + +```python +vol.write_index(begin, end, data, partial=False, retries=None) +vol.write(box, data, partial=False, retries=None) ; vol.write_chunk(index, data, retries=None) + # timed out and retried as reads are: `retries`, else the handle's (Behaviour 6) +wa.thread_budget() -> int # threads for tensorstore's decoding and file I/O in this process: OMP_NUM_THREADS, else + # LSB_DJOB_NUMPROC, else the CPUs it may run on (Behaviour 7) + # every context the package makes or opens a store on is kicked once a second, whoever + # waits on it (Behaviour 6) +``` + ## 3. Conventions These are the rules the package exists to hold in one place. They match chunkmirage's @@ -620,25 +633,64 @@ this document and that one differ, fix one of them. `on_mismatch="relabel"` and picks `method` with its own `is_label_dtype` rule; cellmap-analyze passes `method` from its `interpolation_order`. -6. **Retries** apply to data reads only, with mesh-n-bone's formula, which its tests - already pin: attempt `k` (from 1) waits `timeout * k` seconds on `.result()`, where +6. **Retries** apply to data reads and data writes, with mesh-n-bone's formula, which its + tests already pin: attempt `k` (from 1) waits `timeout * k` seconds on `.result()`, where `timeout` is the given value or, when `None`, `clamp(MiB / 16, 5, 120)` for a remote store (`gs://`, `s3://`, `http(s)://`) and `clamp(MiB / 256, 5, 30)` for a local one (`file`, `memory`), `MiB` being the size of the in-bounds part of the read after - clipping; between attempts the sleep is `base_delay * 1.3 ** (k - 1)` plus `U(0, 1)` + clipping, or of the box written; between attempts the sleep is `base_delay * 1.3 ** (k - 1)` plus `U(0, 1)` seconds; at most `max_attempts` (default 10). Retried: `TimeoutError`, and a tensorstore error whose message starts with `UNAVAILABLE`, `DEADLINE_EXCEEDED` or `RESOURCE_EXHAUSTED` or carries an HTTP 5xx or 429. Any other error is raised at once. A - timed-out tensorstore future is cancelled before the next attempt starts. The policy is a - per-handle attribute and a per-read argument, never only constructor state. + timed-out tensorstore future is cancelled before the next attempt starts; a write may + still land after its cancel, with the same bytes the next attempt writes. Each retry is + logged as a warning naming the error's first line, or its type when the message is + empty (tensorstore's timeout is `TimeoutError()`); when every attempt times out the + `TimeoutError` names the operation, the path and the attempts. Metadata documents are not + retried this way (Convention 5). The policy is a per-handle attribute and a per-read and + per-write argument, never only constructor state. + + The watchdog (0.5): tensorstore up to at least 0.1.85 can strand a queued task. A pool + worker idle for 20 ms exits, and a task queued as it exits gets no thread while the pool + counts itself full (fixed upstream in google/tensorstore ed51718085, in no release yet). + The future then waits forever with no tensorstore thread alive, on a read, a write or an + open. A pool limited to one thread meets it most; cellmap-analyze 0.6 hung on /nrs with + `data_copy` and `file_io` at 1, and one more task on the same context woke the pool + within seconds (salamander's and c-elegans' live tests, 2026-10-22, with un-kicked + workers staying stuck). cellmap-analyze's read retries of June 2025 ("it sometimes + hangs") worked for the same reason: each attempt is a new task. So every context the + package makes (`shared_context`, `capped_context`, a rebuilt one) or opens a store on is + watched, weakly: one daemon thread, started on first use, gives each live watched + context's decoding and file I/O pools one no-op task every second (an in-memory array + copy and a local read of a key that does not exist; no store is touched, about 0.3 ms + of CPU). A stranded task then runs within about a second whoever waits on it: the + package, `vol.ts` directly, or `await`, at any concurrency. The retry stays the outer + net. Open question 14 says when the watchdog goes. 7. **Contexts.** Without `context`, every open in the process shares `shared_context(0)`: - tensorstore's cache is off and the levels of one pyramid share one context. + tensorstore's cache is off and the levels of one pyramid share one context. Its + `data_copy_concurrency` and `file_io_concurrency` limits are `thread_budget()`: the + first of `OMP_NUM_THREADS` (dask sets 1 in every worker it spawns; a cluster job script + sets its slots) and `LSB_DJOB_NUMPROC` (an LSF job's slots) that holds a positive + integer, else the CPUs `sched_getaffinity` allows: the cap every other threaded library + in the job already takes, as Janelia's cluster guidance asks. tensorstore alone sizes + both pools to every CPU of the node, whatever the job's slots, so each of twelve dask + workers in a twelve-slot job on a 96-core node could start 96 threads (0.5). A dask + worker therefore runs tensorstore on one thread of each pool, which Behaviour 6's + watchdog makes safe. SLURM's `SLURM_CPUS_PER_TASK` and SGE's `NSLOTS` stand in for + `LSB_DJOB_NUMPROC` elsewhere. The cap is the default because the process that passes + no context is the one nobody tuned: a quick dask script over `wa.open(path)` is how a + job came to start a thousand threads on twelve slots. One thread per pool slows one + task (a 256^3 block write took 2.1 s against 1.3 s uncapped) but not a job of many + workers, whose throughput its cores bound either way (four writers: 1.8 s against + 1.9 s). A process that wants more threads than its job has slots, a server under + `bsub -n 1` say, asks for more slots or passes its own context (decided 2026-10-23, + over chunkmirage's objection, once the watchdog had removed the hang risk). `capped_context` sets `data_copy_concurrency`, `file_io_concurrency`, `gcs_request_concurrency`, `s3_request_concurrency` and `http_request_concurrency` to `concurrency` and the cache to `cache_bytes`, and returns a new context each call, so two - handles never share one capped pool unless the caller passes the same object twice. Both - are plain tensorstore contexts; a caller may pass its own, at construction or through + handles never share one capped pool unless the caller passes the same object twice; a + `concurrency` of 1 is safe under Behaviour 6's watchdog. Both are plain tensorstore contexts; a caller may pass its own, at construction or through `reopen(context=)` when a long-lived worker changes its concurrency (cellmap-analyze reopens tensorstore on every read today for that reason). @@ -755,7 +807,8 @@ Behaviours 10 to 16 are 0.2; 17 to 19 are 0.3. construction. `partial=True` lifts the check and does tensorstore's read-modify-write of the touched units, documented as unsafe while anything else writes the same unit (cellmap-flow's paint slabs, which it serialises itself). A write commits before the - call returns and is not retried by the package; `retries` applies to reads. `write(box, + call returns, timed out and retried by Behaviour 6 (0.5; before, it was not, and a + stranded write waited forever). `write(box, data)` maps a world box by Behaviour 2 and then calls `write_index`; a box reaching outside the array raises (nothing is padded on write) and `voxel_size` is not an argument (nothing is resampled on write). @@ -942,7 +995,18 @@ and the error after the attempts; `ensure_group_entry` on a foreign array, idemp pickled handle keeping a capped context's spec and sharing the shared one, and a relative path reopening after a `chdir`; `codecs_of` by path; an all-fill write removing its own key; a scalar filling a box through `write_index`, `write` and `write_chunk`, refused when it -does not fit the dtype. +does not fit the dtype. 0.5: a read and a write whose first attempts never finish (a future +nothing resolves) cancelled and retried, the warning naming `TimeoutError`, and the error +after the last attempt; a capped context and a caller's context passed to an open both +kicked by the watchdog, a dead context no longer watched, and a kick landing on a capped +context's own pools; `thread_budget` from each variable, a bad value skipped, and the +shared context's limits. Opt-in (`WORLDARRAY_STRESS=1`, about a minute, +`tests/test_lost_wakeup.py`): worker processes aim tasks at the moment a one-thread pool's +worker exits idle, so a plain context strands some for good (13 of about 7,000 on +tensorstore 0.1.78; 40 to 44 of 12,000 to 15,000 with 16 workers on 0.1.75 and 0.1.85) +and a watched one strands none, its slowest task under a second. Run against cellmap-analyze's +block loop (a chunk read, then a whole-chunk `write_index`, on its 0.6.1 context), 0.4.0 +gave 29 `IndexError`s and 55 writes stuck past 20 s, and 0.5.0 none, slowest 0.98 s. Grids (0.3): `blocks` partitioning with and without a box, `context` as an int, per axis and as before/after pairs, clipping on and off, `len`, `blocks[id]` and `index_of`; @@ -1126,6 +1190,16 @@ over funlib geometry, by this spec's own assignment. boxes on one grid. `Grid.blocks` could carry world boxes too (`b.box`, `b.read_box` as `Box`, a context in world units). Two lines on cellmap-flow's side today. Not planned (2026-10-22): daisy does cellmap-flow's tiling and stays; reopen if cellmap-flow drops it. +14. **When the watchdog goes.** Behaviour 6's watchdog works around tensorstore's lost + wakeup, fixed upstream in google/tensorstore ed51718085 (2026-09-30, "Fix lost wakeups + in TaskGroup when workers block or exit") and in no release up to 0.1.85. When a release + contains it (`gh api repos/google/tensorstore/compare/ed51718085...vX.Y.Z` reports + "ahead" or "identical"), confirm it with `WORLDARRAY_STRESS=1 pytest + tests/test_lost_wakeup.py -s` (a plain context should strand nothing), require that + release, remove `watch`, `kick` and the watchdog thread from `context.py`, and keep the + write retries. That release + needs Python 3.11 or later, as tensorstore has since 0.1.81, so it also ends support for + 3.10. ## 8. Constraints for whoever builds on this diff --git a/pyproject.toml b/pyproject.toml index dbf4ea0..1c5be8d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "worldarray" -version = "0.4.0" +version = "0.5.0" description = "Read zarr, N5 and Neuroglancer precomputed volumes in world coordinates through tensorstore: correct voxel placement, pyramid levels, padded region reads, resampling. numpy + tensorstore only." readme = "README.md" license = { file = "LICENSE" } diff --git a/src/worldarray/__init__.py b/src/worldarray/__init__.py index 37e8d67..ea54e79 100644 --- a/src/worldarray/__init__.py +++ b/src/worldarray/__init__.py @@ -12,7 +12,7 @@ from . import grid, store, units, write from .array import Pyramid, WorldArray, create, metadata, open, open_pyramid, open_tensorstore -from .context import capped_context, capped_context_spec, shared_context +from .context import capped_context, capped_context_spec, shared_context, thread_budget from .geometry import Box, index_range from .grid import Block, Grid, blocks, level_factors from .levels import select_level @@ -28,7 +28,7 @@ parse_chunk_key, ) -__version__ = "0.4.0" +__version__ = "0.5.0" __all__ = [ "ArrayMeta", @@ -62,6 +62,7 @@ "open_tensorstore", "select_level", "shared_context", + "thread_budget", "store", "units", "write", diff --git a/src/worldarray/array.py b/src/worldarray/array.py index cff63c7..35c71bc 100644 --- a/src/worldarray/array.py +++ b/src/worldarray/array.py @@ -35,7 +35,7 @@ read_meta, read_meta_light, ) -from .retry import Retry, read_with_retries +from .retry import Retry, read_with_retries, write_with_retries from .store import ( Node, detect, @@ -533,7 +533,12 @@ def list_chunks(self) -> list[tuple[int, ...]]: # --- writing (Behaviour 13) ---------------------------------------------------- def write_index( - self, begin: Sequence[int], end: Sequence[int], data: np.ndarray, partial: bool = False + self, + begin: Sequence[int], + end: Sequence[int], + data: np.ndarray, + partial: bool = False, + retries: Retry | None = None, ) -> None: """Write ``data`` over the voxel box ``[begin, end)`` of the stored grid. The box must be a union of whole write units (every bound a multiple of ``write_shape`` on its axis, @@ -542,7 +547,8 @@ def write_index( read-modify-write of the touched units, unsafe while anything else writes the same unit. ``data`` may be a scalar, which fills the box: it is converted to the array's dtype only when the value survives (``0`` into ``uint8``, not ``300`` or ``0.5``). - No other cast, no padding, no resampling; the write commits before returning.""" + No other cast, no padding, no resampling; the write commits before returning, timed + out and retried as reads are (``retries``, else the handle's).""" if not self.writable: raise ValueError( f"{self.path} was opened read-only; open it with mode='r+' or create it" @@ -577,7 +583,13 @@ def write_index( f"{self.path}: write [{b}, {e}) on axis {axis} ({self.axes[axis]}) is not whole write " f"units of {u}; align the box to write_shape {self.write_shape} or pass partial=True" ) - self._store[tuple(slice(b, e) for b, e in zip(begin, end))].write(data).result() + write_with_retries( + self._store[tuple(slice(b, e) for b, e in zip(begin, end))], + data, + retries or self.retries, + remote=self.remote, + what=self.path, + ) def _scalar(self, value: np.ndarray) -> np.ndarray: """A 0-d ``value`` in the array's dtype, or ``ValueError`` when converting changes it: @@ -595,7 +607,9 @@ def _scalar(self, value: np.ndarray) -> np.ndarray: raise ValueError(f"{self.path}: {value.item()!r} does not fit {dtype}") return out - def write(self, box, data: np.ndarray, partial: bool = False) -> None: + def write( + self, box, data: np.ndarray, partial: bool = False, retries: Retry | None = None + ) -> None: """Write ``data`` over a world box: the voxels Behaviour 2's centre rule selects, then ``write_index``. A box reaching outside the array raises; nothing is padded or resampled on write.""" @@ -606,7 +620,7 @@ def write(self, box, data: np.ndarray, partial: bool = False) -> None: f"{self.path}: the box reaches outside the array on axis {axis} ({self.axes[axis]}): " f"voxels [{b}, {e}) of {n}; nothing is padded on write" ) - self.write_index(begin, end, data, partial=partial) + self.write_index(begin, end, data, partial=partial, retries=retries) def blocks(self, block_shape=None, context=0, box=None, clip: bool = True): """The array cut into blocks (Behaviour 17): write boxes tiling the stored grid (or a @@ -623,7 +637,9 @@ def blocks(self, block_shape=None, context=0, box=None, clip: bool = True): clip=clip, ) - def write_chunk(self, index: Sequence[int], data: np.ndarray) -> None: + def write_chunk( + self, index: Sequence[int], data: np.ndarray, retries: Retry | None = None + ) -> None: """Write the whole write unit at ``index`` on the storage grid (the shard of a sharded array); ``data`` has that unit's shape, clipped at the array's end.""" unit, shape = self.write_shape, self._grid.shape @@ -634,7 +650,7 @@ def write_chunk(self, index: Sequence[int], data: np.ndarray) -> None: if any(i < 0 or b >= n for i, b, n in zip(index, begin, shape)): raise ValueError(f"{self.path}: no write unit at {index}") end = tuple(min(b + u, n) for b, u, n in zip(begin, unit, shape)) - self.write_index(begin, end, data) + self.write_index(begin, end, data, retries=retries) def read_index( self, diff --git a/src/worldarray/context.py b/src/worldarray/context.py index 8d283f2..0508557 100644 --- a/src/worldarray/context.py +++ b/src/worldarray/context.py @@ -3,7 +3,10 @@ ``shared_context(cache_bytes)`` is one ``tensorstore.Context`` per budget for the whole process, so the levels of a pyramid and repeat opens of an array draw on one cache pool. The default budget is 0: tensorstore's decoded-chunk cache is off, because every consumer -either reads each region once or caches decoded chunks itself. +either reads each region once or caches decoded chunks itself. Its decoding and file I/O +pools are sized by ``thread_budget()``, from the variables a cluster job sets, where +tensorstore alone would size them to every CPU on the node. Every context here, and every +one a store is opened on, is watched for tensorstore's lost wakeup (``watch``). ``capped_context`` is mesh-n-bone's worker-process context: every tensorstore concurrency resource limited, so dozens of worker processes on one host do not multiply into thousands @@ -14,7 +17,11 @@ from __future__ import annotations import json +import os +import tempfile import threading +import time +import weakref from typing import Any import tensorstore as ts @@ -24,11 +31,124 @@ _lock = threading.Lock() +# tensorstore's lost wakeup. In every release up to at least 0.1.85, a pool worker idle for +# 20 ms exits, and a task queued as it exits gets no thread while the pool counts itself +# full: the task never runs and its future waits forever, with no tensorstore thread alive. +# A pool limited to one thread meets it most (cellmap-analyze's workers on /nrs, +# 2026-10-22; its read retries of June 2025 were working around the same thing). Any later +# task on the same pool wakes it. So every context the package makes or opens a store on is +# watched: a daemon thread gives each one's decoding and file I/O pools a no-op task every +# KICK_EVERY seconds, and a stranded task, whoever waits on it (worldarray, ``vol.ts``, +# ``await``), runs within that. Fixed upstream in google/tensorstore ed51718085 (2026-09-30, +# "Fix lost wakeups in TaskGroup when workers block or exit"). Once a release contains it, +# require that release and remove the watchdog (spec, open question 14): +# `gh api repos/google/tensorstore/compare/ed51718085...vX.Y.Z` says "ahead" or "identical" +# for a release that has it, and `WORLDARRAY_STRESS=1 pytest tests/test_lost_wakeup.py -s` +# then strands nothing on a plain context. +KICK_EVERY = 1.0 +_KICK_DIR = tempfile.gettempdir() +_KICK_KEY = "worldarray-kick-no-such-key" +_watched: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() +_watch_lock = threading.Lock() +_watchdog: threading.Thread | None = None + + +def _kickers(context: ts.Context) -> tuple: + """An in-memory array and a local kvstore on ``context``: reading the one copies on its + decoding pool, reading a missing key from the other runs on its file I/O pool.""" + array = ts.open({"driver": "array", "array": [0], "dtype": "uint8"}, context=context) + kv = ts.KvStore.open({"driver": "file", "path": _KICK_DIR + "/"}, context=context) + return array.result(timeout=10), kv.result(timeout=10) + + +def kick(context: ts.Context) -> list: + """One no-op task on each of ``context``'s decoding and file I/O pools; no store is + touched. Returns the futures, which the caller holds until done (a dropped future may + be cancelled before it is queued).""" + try: + with _watch_lock: + kickers = _watched.get(context) + if kickers is None: + kickers = _kickers(context) + with _watch_lock: + if context in _watched: + _watched[context] = kickers + array, kv = kickers + return [array.read(), kv.read(_KICK_KEY)] + except Exception: # noqa: BLE001 a kick that cannot start changes nothing + return [] + + +def watch(context: ts.Context | None) -> ts.Context | None: + """Kick ``context``'s pools every ``KICK_EVERY`` seconds for as long as it lives, starting + the watchdog thread on first use. Returns ``context``.""" + global _watchdog + if context is None: + return None + with _watch_lock: + if context not in _watched: + _watched[context] = None + if _watchdog is None or not _watchdog.is_alive(): + _watchdog = threading.Thread(target=_watch, name="worldarray-kick", daemon=True) + _watchdog.start() + return context + + +def _watch() -> None: + held: list = [] + while True: + time.sleep(KICK_EVERY) + with _watch_lock: + contexts = list(_watched.keys()) + held = [f for f in held if not f.done()] + for context in contexts: + held += kick(context) + del contexts + + +def _forget_watchdog() -> None: + global _watchdog + _watchdog = None # a forked child has no such thread; the next watch() starts one + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_forget_watchdog) + + +_SLOTS = ("OMP_NUM_THREADS", "LSB_DJOB_NUMPROC", "SLURM_CPUS_PER_TASK", "NSLOTS") + + +def thread_budget() -> int: + """Threads for this process's tensorstore decoding and file I/O: ``OMP_NUM_THREADS`` + when set (dask sets 1 in every worker it spawns; a job script sets its slots), else the + job's slots (LSF's ``LSB_DJOB_NUMPROC``, SLURM's ``SLURM_CPUS_PER_TASK``, SGE's + ``NSLOTS``), else the CPUs this process may run on.""" + for name in _SLOTS: + try: + n = int(os.environ.get(name, "")) + except ValueError: + continue + if n > 0: + return n + try: + return len(os.sched_getaffinity(0)) + except (AttributeError, OSError): + return os.cpu_count() or 1 + + def shared_context(cache_bytes: int = 0) -> ts.Context: with _lock: ctx = _shared.get(int(cache_bytes)) if ctx is None: - ctx = ts.Context({"cache_pool": {"total_bytes_limit": int(cache_bytes)}}) + limit = {"limit": thread_budget()} + ctx = ts.Context( + { + "cache_pool": {"total_bytes_limit": int(cache_bytes)}, + "data_copy_concurrency": limit, + "file_io_concurrency": limit, + } + ) + watch(ctx) _shared[int(cache_bytes)] = ctx return ctx @@ -36,7 +156,8 @@ def shared_context(cache_bytes: int = 0) -> ts.Context: def capped_context_spec(cache_bytes: int = 0, concurrency: int = 1, network: bool = True) -> dict: """The spec dict behind ``capped_context``, for callers building their own opens. ``network=False`` caps only decoding and local file I/O, leaving tensorstore's request - pools alone (cellmap-analyze's choice for remote reads).""" + pools alone (cellmap-analyze's choice for remote reads). A ``concurrency`` of 1 is safe: + the context is watched (see ``KICK_EVERY``).""" limit = {"limit": int(concurrency)} spec = { "data_copy_concurrency": limit, @@ -53,7 +174,7 @@ def capped_context_spec(cache_bytes: int = 0, concurrency: int = 1, network: boo def capped_context(cache_bytes: int = 0, concurrency: int = 1, network: bool = True) -> ts.Context: - return ts.Context(capped_context_spec(cache_bytes, concurrency, network)) + return watch(ts.Context(capped_context_spec(cache_bytes, concurrency, network))) def recipe_of(context: ts.Context) -> tuple[str, Any]: @@ -79,6 +200,6 @@ def from_recipe(recipe: tuple[str, Any] | None) -> ts.Context | None: with _lock: ctx = _rebuilt.get(value) if ctx is None: - ctx = ts.Context(ts.Context.Spec(json.loads(value))) + ctx = watch(ts.Context(ts.Context.Spec(json.loads(value)))) _rebuilt[value] = ctx return ctx diff --git a/src/worldarray/retry.py b/src/worldarray/retry.py index 3eeb1c2..352c05f 100644 --- a/src/worldarray/retry.py +++ b/src/worldarray/retry.py @@ -1,7 +1,7 @@ -"""Retrying a data read on a transient failure. +"""Retrying a data read or write on a transient failure. mesh-n-bone's policy, which its tests pin: attempt ``k`` (from 1) waits ``timeout * k`` -seconds; ``timeout`` is given, or sized from the read and where it comes from: +seconds; ``timeout`` is given, or sized from the bytes moved and where they go: ``clamp(MiB / 16, 5, 120)`` for a remote store and ``clamp(MiB / 256, 5, 30)`` for a local one. Between attempts the sleep is ``base_delay * 1.3 ** (k - 1)`` plus up to one second of jitter. A timed-out future is cancelled before the next attempt (spec, Behaviour 6). @@ -43,15 +43,22 @@ def is_transient(error: BaseException) -> bool: return message.startswith(_TRANSIENT_PREFIXES) or bool(_HTTP_TRANSIENT.search(message)) -def read_with_retries(view, retry: Retry, *, remote: bool, what: str = "") -> np.ndarray: - """``view.read()`` as a numpy array, retried per ``retry``. ``view`` is a tensorstore - (or any object whose ``read()`` returns a future with ``result(timeout)`` and - ``cancel()``).""" - nbytes = int(np.prod(view.shape, dtype=np.int64)) * np.dtype(view.dtype.numpy_dtype).itemsize +def describe(error: BaseException) -> str: + """The first line of ``error``'s message, or its type's name when the message is empty + (a timed-out tensorstore future raises ``TimeoutError()``).""" + lines = str(error).splitlines() + return lines[0][:120] if lines else type(error).__name__ + + +def with_retries(start, nbytes: int, retry: Retry, *, remote: bool, what: str, verb: str): + """``start().result()``, retried per ``retry``. ``start`` begins the operation and + returns a future with ``result(timeout)`` and ``cancel()``; ``nbytes`` sizes the timeout + and ``verb`` ("read", "write") names the operation in the warning and the error.""" timeout = retry.timeout_for(nbytes, remote) + attempts = max(1, retry.max_attempts) last: BaseException | None = None - for attempt in range(1, max(1, retry.max_attempts) + 1): - future = view.read() + for attempt in range(1, attempts + 1): + future = start() try: return future.result(timeout=timeout * attempt) except TimeoutError as e: @@ -61,16 +68,40 @@ def read_with_retries(view, retry: Retry, *, remote: bool, what: str = "") -> np if not is_transient(e): raise last = e - if attempt < retry.max_attempts: + if attempt < attempts: delay = retry.base_delay * 1.3 ** (attempt - 1) + random.random() logger.warning( - "read %s failed on attempt %d/%d (%s); retrying in %.1fs", + "%s %s failed on attempt %d/%d (%s); retrying in %.1fs", + verb, what, attempt, - retry.max_attempts, - str(last).splitlines()[0][:120], + attempts, + describe(last), delay, ) time.sleep(delay) assert last is not None + if isinstance(last, TimeoutError): + raise TimeoutError( + f"{verb} {what} timed out on all {attempts} attempts, the last after " + f"{timeout * attempts:g}s" + ) from last raise last + + +def read_with_retries(view, retry: Retry, *, remote: bool, what: str = "") -> np.ndarray: + """``view.read()`` as a numpy array, retried per ``retry``. ``view`` is a tensorstore + (or any object whose ``read()`` returns a future with ``result(timeout)`` and + ``cancel()``).""" + nbytes = int(np.prod(view.shape, dtype=np.int64)) * np.dtype(view.dtype.numpy_dtype).itemsize + return with_retries(view.read, nbytes, retry, remote=remote, what=what, verb="read") + + +def write_with_retries( + view, data: np.ndarray, retry: Retry, *, remote: bool, what: str = "" +) -> None: + """``view.write(data)`` committed, retried per ``retry`` like a read. A timed-out write + is cancelled, but may still land; it writes the same bytes, so the next attempt is + safe.""" + nbytes = int(np.prod(view.shape, dtype=np.int64)) * np.dtype(view.dtype.numpy_dtype).itemsize + with_retries(lambda: view.write(data), nbytes, retry, remote=remote, what=what, verb="write") diff --git a/src/worldarray/store.py b/src/worldarray/store.py index ea0acf9..b8aee74 100644 --- a/src/worldarray/store.py +++ b/src/worldarray/store.py @@ -20,6 +20,8 @@ import tensorstore as ts +from .context import watch + logger = logging.getLogger(__name__) T = TypeVar("T") @@ -438,7 +440,7 @@ def attempt(kvstore: dict) -> ts.TensorStore: spec["assume_metadata"] = True if node.format == "precomputed" and scale_index is not None: spec["scale_index"] = scale_index - return ts.open(spec, read=True, write=write, context=context).result() + return ts.open(spec, read=True, write=write, context=watch(context)).result() store = with_access(kvstore_specs(node.path), attempt) if node.format in ("n5", "precomputed"): diff --git a/src/worldarray/write.py b/src/worldarray/write.py index f7fa190..01cc960 100644 --- a/src/worldarray/write.py +++ b/src/worldarray/write.py @@ -27,6 +27,7 @@ import numpy as np import tensorstore as ts +from .context import watch from .geometry import snap_integral from .placement import spatial_axes from .store import ( @@ -767,7 +768,7 @@ def attempt(kvstore: dict) -> ts.TensorStore: separator, store_fill, ) - return ts.open(spec, create=True, open=False, context=context).result() + return ts.open(spec, create=True, open=False, context=watch(context)).result() try: with_access(kvstore_specs(path), attempt) diff --git a/tests/test_access.py b/tests/test_access.py index 7dfd5d8..2eea1dd 100644 --- a/tests/test_access.py +++ b/tests/test_access.py @@ -2,14 +2,19 @@ from __future__ import annotations +import gc +import logging import os +import time +import weakref import numpy as np import pytest +import tensorstore as ts import worldarray as wa from conftest import pattern, write_zarr2, write_zarr2_group -from worldarray import store +from worldarray import retry, store @pytest.fixture(autouse=True) @@ -112,6 +117,94 @@ def test_contexts_and_retry_timeouts(): assert wa.Retry(timeout=7).timeout_for(10**12, remote=True) == 7.0 +def test_the_thread_budget_follows_the_job(monkeypatch): + monkeypatch.delenv("OMP_NUM_THREADS", raising=False) + monkeypatch.delenv("LSB_DJOB_NUMPROC", raising=False) + allowed = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count() + assert wa.thread_budget() == allowed # macOS has no affinity call + monkeypatch.setenv("LSB_DJOB_NUMPROC", "12") + assert wa.thread_budget() == 12 + monkeypatch.setenv("OMP_NUM_THREADS", "24") # a job script's 2 * slots + assert wa.thread_budget() == 24 + monkeypatch.setenv("OMP_NUM_THREADS", "1") # what dask sets in each worker + assert wa.thread_budget() == 1 + monkeypatch.setenv("OMP_NUM_THREADS", "x") + assert wa.thread_budget() == 12 + monkeypatch.delenv("LSB_DJOB_NUMPROC") + monkeypatch.setenv("SLURM_CPUS_PER_TASK", "6") + assert wa.thread_budget() == 6 + monkeypatch.setattr(wa.context, "_shared", {}) + spec = wa.shared_context(5).spec.to_json() + assert spec["data_copy_concurrency"] == spec["file_io_concurrency"] == {"limit": 6} + assert spec["cache_pool"] == {"total_bytes_limit": 5} + + +def test_a_kick_runs_on_the_contexts_own_pools(): + ctx = wa.capped_context(concurrency=1) + futures = wa.context.kick(ctx) + assert len(futures) == 2 + assert [type(f.result(timeout=10)).__name__ for f in futures] == ["ndarray", "ReadResult"] + + +def test_every_context_the_package_touches_is_kicked(monkeypatch, tmp_path): + # tensorstore's lost wakeup: a stranded task runs once anything else uses its pool, so + # a watched context's pools get a no-op task every KICK_EVERY seconds, whoever waits + kicked = [] + monkeypatch.setattr(wa.context, "_watched", weakref.WeakKeyDictionary()) + monkeypatch.setattr(wa.context, "kick", lambda c: kicked.append(id(c)) or []) + monkeypatch.setattr(wa.context, "KICK_EVERY", 0.01) + monkeypatch.setattr(wa.context, "_watchdog", None) + capped = wa.capped_context(concurrency=1) + own = ts.Context({"data_copy_concurrency": {"limit": 1}}) + write_zarr2(tmp_path / "a.zarr", pattern((4, 4)), chunks=(2, 2)) + wa.open_tensorstore(str(tmp_path / "a.zarr"), context=own) + deadline = time.monotonic() + 5 + while not {id(capped), id(own)} <= set(kicked) and time.monotonic() < deadline: + time.sleep(0.01) + assert {id(capped), id(own)} <= set(kicked) + assert wa.context._watchdog.daemon and wa.context._watchdog.name == "worldarray-kick" + n = len(wa.context._watched) + del capped + gc.collect() + assert len(wa.context._watched) == n - 1 # a context that died is no longer watched + + +class _StalledView: + """A view whose first ``stalls`` reads never finish (a future nothing resolves, as a + read stuck on NFS looks from Python); later reads return ``0, 1, 2, 3``.""" + + shape, dtype = (4,), ts.uint8 + + def __init__(self, stalls): + self.stalls, self.held = stalls, [] + + def read(self): + promise, future = ts.Promise.new() + if len(self.held) < self.stalls: + self.held.append((promise, future)) + else: + promise.set_result(np.arange(4, dtype=np.uint8)) + return future + + +def test_a_timed_out_read_is_cancelled_and_retried(monkeypatch, caplog): + # tensorstore's TimeoutError has an empty message; the warning names its type + # (salamander's IndexError on /nrs under 0.4.0) + monkeypatch.setattr(retry.time, "sleep", lambda s: None) + view = _StalledView(stalls=2) + with caplog.at_level(logging.WARNING, logger="worldarray.retry"): + got = retry.read_with_retries( + view, wa.Retry(timeout=0.01, max_attempts=3), remote=False, what="v" + ) + assert got.tolist() == [0, 1, 2, 3] + assert [f.cancelled() for _, f in view.held] == [True, True] + assert "read v failed on attempt 1/3 (TimeoutError)" in caplog.text + with pytest.raises(TimeoutError, match="read v timed out on all 2 attempts"): + retry.read_with_retries( + _StalledView(stalls=9), wa.Retry(timeout=0.01, max_attempts=2), remote=False, what="v" + ) + + def test_metadata_reads_carry_capped_retries(): spec = store.metadata_spec({"driver": "s3", "bucket": "b", "path": "p/"}) assert spec["context"]["s3_request_retries"] == store.METADATA_RETRIES diff --git a/tests/test_lost_wakeup.py b/tests/test_lost_wakeup.py new file mode 100644 index 0000000..3b6cb5a --- /dev/null +++ b/tests/test_lost_wakeup.py @@ -0,0 +1,76 @@ +"""tensorstore's lost wakeup, reproduced on purpose (spec, Behaviour 6 and open question 14). + +Opt-in, about a minute: ``WORLDARRAY_STRESS=1 pytest tests/test_lost_wakeup.py -s``. Worker +processes give a pool limited to one thread its next task 19.5-20.5 ms after the last one +finished, when the pool's only worker exits idle. On an affected tensorstore (every release +up to at least 0.1.85) about 0.3% of those tasks are stranded on a plain context and never +run; on a context the package watches, every one finishes within about a second. When a +release says it carries google/tensorstore ed51718085, run this: no strands on the plain +context means the watchdog can go. +""" + +from __future__ import annotations + +import multiprocessing as mp +import os +import random +import tempfile +import time + +import pytest + +pytestmark = pytest.mark.skipif( + not os.environ.get("WORLDARRAY_STRESS"), reason="set WORLDARRAY_STRESS=1 to run" +) + +PROCS, SECONDS, GIVE_UP = 8, 20.0, 10.0 + + +def _aim(args): + worker, watched = args + import tensorstore as ts + + import worldarray as wa + + limits = {"data_copy_concurrency": {"limit": 1}, "file_io_concurrency": {"limit": 1}} + ctx = wa.capped_context(concurrency=1, network=False) if watched else ts.Context(limits) + kv = ts.KvStore.open({"driver": "file", "path": tempfile.mkdtemp() + "/"}, context=ctx) + kv = kv.result() + rng = random.Random(worker) + slowest, stranded = 0.0, 0 + end = time.monotonic() + SECONDS + while time.monotonic() < end: + time.sleep(rng.uniform(0.0195, 0.0205)) + future, start = kv.read("missing"), time.monotonic() + try: + future.result(timeout=GIVE_UP) + except TimeoutError: + stranded += 1 + kv.read("other").result(timeout=GIVE_UP) # any task on the pool frees it + future.result(timeout=GIVE_UP) + slowest = max(slowest, time.monotonic() - start) + return stranded, slowest + + +def _run(watched: bool) -> tuple[int, float]: + with mp.get_context("spawn").Pool(PROCS) as pool: + results = pool.map(_aim, [(w, watched) for w in range(PROCS)]) + return sum(r[0] for r in results), max(r[1] for r in results) + + +def test_a_watched_context_never_strands_a_task(): + import importlib.metadata + + version = importlib.metadata.version("tensorstore") + stranded, _ = _run(watched=False) + watched_stranded, slowest = _run(watched=True) + print( + f"\ntensorstore {version}: plain context stranded {stranded} tasks for {GIVE_UP:g}s+; " + f"watched context stranded {watched_stranded}, slowest {slowest:.2f}s" + ) + assert watched_stranded == 0 and slowest < 2.0 + if stranded == 0: + pytest.skip( + f"tensorstore {version} stranded nothing: if it carries ed51718085, " + "the watchdog can go (open question 14)" + ) diff --git a/tests/test_write.py b/tests/test_write.py index daff2cb..a4290f4 100644 --- a/tests/test_write.py +++ b/tests/test_write.py @@ -3,8 +3,10 @@ from __future__ import annotations import json +import logging import multiprocessing as mp import pickle +from types import SimpleNamespace import numpy as np import pytest @@ -748,6 +750,47 @@ def test_a_pickled_relative_path_reopens_after_a_chdir(tmp_path, monkeypatch): assert vol.reopen().level_path == again.level_path +class _StallingStore: + """A tensorstore whose first ``stalls`` writes never finish, as a write stuck on NFS + looks from Python; later writes go through.""" + + def __init__(self, store, stalls): + self.store, self.stalls, self.held = store, stalls, [] + + def __getitem__(self, index): + view = self.store[index] + + def write(data): + if len(self.held) >= self.stalls: + return view.write(data) + promise, future = ts.Promise.new() + self.held.append((promise, future)) + return future + + return SimpleNamespace(shape=view.shape, dtype=view.dtype, write=write) + + +def test_a_stalled_write_is_cancelled_and_retried(tmp_path, monkeypatch, caplog): + # salamander's 0.4.0 runs on /nrs: workers parked forever in write().result() + monkeypatch.setattr(wa.retry.time, "sleep", lambda s: None) + vol = wa.create( + str(tmp_path / "vol.zarr/s0"), (8, 8, 8), "uint8", (1, 1, 1), (0.5,) * 3, chunks=(4, 4, 4) + ) + store = vol._store + vol._store = _StallingStore(store, stalls=2) + with caplog.at_level(logging.WARNING, logger="worldarray.retry"): + vol.write_index((0, 0, 0), (4, 4, 4), 9, retries=wa.Retry(timeout=0.01, max_attempts=3)) + assert [f.cancelled() for _, f in vol._store.held] == [True, True] + assert "write " + vol.path + " failed on attempt 1/3 (TimeoutError)" in caplog.text + vol.retries = wa.Retry(timeout=0.01, max_attempts=2) + vol._store = _StallingStore(store, stalls=9) + with pytest.raises(TimeoutError, match="timed out on all 2 attempts"): + vol.write_chunk((1, 1, 1), 5) + vol._store = store + got = vol.read_index((0, 0, 0), (8, 8, 8)) + assert np.all(got[:4, :4, :4] == 9) and np.all(got[4:, 4:, 4:] == 0) + + def test_a_scalar_fills_the_box(tmp_path): vol = wa.create( str(tmp_path / "vol.zarr/s0"), (8, 8, 8), "uint8", (1, 1, 1), (0.5,) * 3, chunks=(4, 4, 4) diff --git a/uv.lock b/uv.lock index 1780bd0..8c27fbe 100644 --- a/uv.lock +++ b/uv.lock @@ -1247,7 +1247,7 @@ wheels = [ [[package]] name = "worldarray" -version = "0.4.0" +version = "0.5.0" source = { editable = "." } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },