From 351c13212ca699537faaeaac85bbfb4bd251428c Mon Sep 17 00:00:00 2001 From: David Ackerman Date: Fri, 23 Oct 2026 00:17:49 -0400 Subject: [PATCH 1/7] Retry writes, kick stranded tensorstore pools, size threads to the job (0.5.0) cellmap-analyze 0.6 hung and crashed on /nrs (salamander's and c-elegans' runs): - A read whose first attempt timed out raised IndexError: the retry warning took the first line of tensorstore's TimeoutError(), which has no message. - Dozens of workers per run waited forever in write_index with no tensorstore thread alive. tensorstore up to 0.1.85 strands a task queued while a pool worker exits idle (fixed upstream in ed51718085, unreleased); pools capped at 1 meet it most. One more task on the same pool wakes it, which both sessions proved on live stuck workers. Every wait on a data read, write or open now kicks its context's decoding and file I/O pools with a no-op task after each second pending; writes are timed out and retried as reads are; and the shared context's pools take thread_budget() (OMP_NUM_THREADS, else LSB_DJOB_NUMPROC, else the CPUs allowed) instead of every core on the node. Spec open question 14 says when the kick goes. Claude-Session: https://claude.ai/code/session_0141wgF8mkGbjGfBFKfKhp64 --- CHANGELOG.md | 33 ++++++++++++++ README.md | 12 +++-- docs/spec.md | 76 ++++++++++++++++++++++++++----- pyproject.toml | 2 +- src/worldarray/__init__.py | 5 ++- src/worldarray/array.py | 37 +++++++++++---- src/worldarray/context.py | 87 +++++++++++++++++++++++++++++++++-- src/worldarray/retry.py | 79 +++++++++++++++++++++++++------- src/worldarray/store.py | 4 +- src/worldarray/write.py | 3 +- tests/test_access.py | 92 +++++++++++++++++++++++++++++++++++++- tests/test_write.py | 43 ++++++++++++++++++ uv.lock | 2 +- 13 files changed, 426 insertions(+), 49 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4205cb3..4d2e9e1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,38 @@ # Changelog +## 0.5.0 (unreleased) + +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 wait on a data read, a data write or a store open nudges its context's thread + pools after each second the operation is still pending: one no-op task on the decoding + pool and one on the file I/O pool (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 operation now finishes + within about a second at any concurrency, a limit of 1 included. 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 + `LSB_DJOB_NUMPROC` (an LSF job's slots), else the CPUs the process may run on. + 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. Writes stay 2 to + 8 times faster than cellmap-analyze's zarr-python writer on NFS, at one scale, measured + on 128^3 uint64 label chunks. + ## 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..3481287 100644 --- a/README.md +++ b/README.md @@ -86,9 +86,11 @@ 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 a waiting operation's thread pools nudged + every second, 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 +149,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, LSB_DJOB_NUMPROC or CPUs wa.list_chunk_keys(path, separator=".") # chunk indices from keys alone, no metadata needed wa.Box ; wa.units.TABLE @@ -174,7 +177,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..1ee2a93 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 wait on a data read, write or open kicks its context's pools after each second + # it is still pending (Behaviour 6) +``` + ## 3. Conventions These are the rules the package exists to hold in one place. They match chunkmirage's @@ -620,25 +633,53 @@ 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 kick (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). So every wait on a data read, a data write or a store open + gives the context's decoding and file I/O pools one no-op task each (an in-memory array + copy and a local read of a key that does not exist; no store is touched) after every + second the future is still pending. A stranded operation finishes within about a + second, at any concurrency, and the retry stays the outer net. Open question 14 says + when the kick 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 kick + makes safe. `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 kick. 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 +796,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 +984,11 @@ 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 pending future kicked until it resolves, a wait's timeout, 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. 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 +1172,14 @@ 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 kick goes.** Behaviour 6's kick 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"), require that release, remove `context.kick` and `KICK_AFTER` + (waits go back to `future.result(timeout)`), 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..fe3320b 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,14 @@ 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, + context=self.context, + ) 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 +608,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 +621,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 +638,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 +651,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, @@ -719,7 +736,9 @@ def _read_block( transform: Transform | None = None, ) -> np.ndarray: view = self._store[tuple(slice(int(a), int(c)) for a, c in zip(lo, hi))] - data = read_with_retries(view, retries or self.retries, remote=self.remote, what=self.path) + data = read_with_retries( + view, retries or self.retries, remote=self.remote, what=self.path, context=self.context + ) data = np.asarray(data) if not data.dtype.isnative: data = data.astype(data.dtype.newbyteorder("=")) diff --git a/src/worldarray/context.py b/src/worldarray/context.py index 8d283f2..6ca90bb 100644 --- a/src/worldarray/context.py +++ b/src/worldarray/context.py @@ -3,7 +3,9 @@ ``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. ``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 +16,10 @@ from __future__ import annotations import json +import os +import tempfile import threading +import time from typing import Any import tensorstore as ts @@ -24,11 +29,86 @@ _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). Any later task on the same pool wakes it, so ``wait`` gives a context's pools +# a no-op task (``kick``) every KICK_AFTER seconds a future is pending. 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 ``kick`` +# (spec, open question 14): `gh api repos/google/tensorstore/compare/ed51718085...vX.Y.Z` +# says "ahead" or "identical" for a release that has it. +KICK_AFTER = 1.0 +_KICK_DIR = tempfile.gettempdir() +_KICK_KEY = "worldarray-kick-no-such-key" + + +def kick(context: ts.Context) -> list: + """One no-op task on each of ``context``'s decoding and file I/O pools: an in-memory + array copy and a local read of a key that does not exist. No store is touched. Returns + the futures, which the caller holds until done (a dropped future may be cancelled).""" + futures = [] + try: + array = ts.open( + {"driver": "array", "array": [0], "dtype": "uint8"}, context=context + ).result(timeout=KICK_AFTER) + futures.append(array.read()) + kv = ts.KvStore.open({"driver": "file", "path": _KICK_DIR + "/"}, context=context).result( + timeout=KICK_AFTER + ) + futures.append(kv.read(_KICK_KEY)) + except Exception: # noqa: BLE001 a kick that cannot start changes nothing + pass + return futures + + +def wait(future, context: ts.Context | None, timeout: float | None = None): + """``future.result()``, kicking ``context``'s pools after every ``KICK_AFTER`` seconds it + is still pending; ``TimeoutError`` once ``timeout`` seconds have passed (never when + None). A future that is not a tensorstore one, or no context, is simply waited on.""" + deadline = None if timeout is None else time.monotonic() + timeout + kicks: list = [] + while True: + step = KICK_AFTER if deadline is None else min(KICK_AFTER, deadline - time.monotonic()) + try: + return future.result(timeout=max(step, 0.0)) + except TimeoutError: + if deadline is not None and time.monotonic() >= deadline: + raise + if context is not None: + kicks = [f for f in kicks if not f.done()] + kick(context) + + +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 + ``LSB_DJOB_NUMPROC`` (an LSF job's slots), else the CPUs this process may run on.""" + for name in ("OMP_NUM_THREADS", "LSB_DJOB_NUMPROC"): + 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, + } + ) _shared[int(cache_bytes)] = ctx return ctx @@ -36,7 +116,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: + every wait in the package kicks a pool that stranded a task (see ``KICK_AFTER``).""" limit = {"limit": int(concurrency)} spec = { "data_copy_concurrency": limit, diff --git a/src/worldarray/retry.py b/src/worldarray/retry.py index 3eeb1c2..254b138 100644 --- a/src/worldarray/retry.py +++ b/src/worldarray/retry.py @@ -1,10 +1,11 @@ -"""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). +jitter. A timed-out future is cancelled before the next attempt (spec, Behaviour 6). While +an attempt waits, its context's pools are kicked every second (``context.wait``). """ from __future__ import annotations @@ -17,6 +18,8 @@ import numpy as np +from .context import wait + logger = logging.getLogger(__name__) _TRANSIENT_PREFIXES = ("UNAVAILABLE", "DEADLINE_EXCEEDED", "RESOURCE_EXHAUSTED") @@ -43,17 +46,27 @@ 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, context=None +): + """``start().result()``, retried per ``retry``. ``start`` begins the operation and + returns a future with ``result(timeout)`` and ``cancel()``; ``nbytes`` sizes the timeout, + ``verb`` ("read", "write") names the operation in the warning and the error, and + ``context`` (the operation's) is kicked while it waits.""" 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) + return wait(future, context, timeout * attempt) except TimeoutError as e: future.cancel() last = e @@ -61,16 +74,52 @@ 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 = "", context=None +) -> 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", context=context + ) + + +def write_with_retries( + view, data: np.ndarray, retry: Retry, *, remote: bool, what: str = "", context=None +) -> 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", + context=context, + ) diff --git a/src/worldarray/store.py b/src/worldarray/store.py index ea0acf9..3daa5c2 100644 --- a/src/worldarray/store.py +++ b/src/worldarray/store.py @@ -20,6 +20,8 @@ import tensorstore as ts +from .context import wait + 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 wait(ts.open(spec, read=True, write=write, context=context), context) 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..43d8d9c 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 wait 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 wait(ts.open(spec, create=True, open=False, context=context), context) try: with_access(kvstore_specs(path), attempt) diff --git a/tests/test_access.py b/tests/test_access.py index 7dfd5d8..dcd5365 100644 --- a/tests/test_access.py +++ b/tests/test_access.py @@ -2,14 +2,16 @@ from __future__ import annotations +import logging import os 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 +114,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) + assert wa.thread_budget() == len(os.sched_getaffinity(0)) + 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.setattr(wa.context, "_shared", {}) + spec = wa.shared_context(5).spec.to_json() + assert spec["data_copy_concurrency"] == spec["file_io_concurrency"] == {"limit": 12} + 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_a_pending_future_is_kicked_until_it_resolves(monkeypatch): + # tensorstore's lost wakeup: the stranded task runs once anything else uses its pool + monkeypatch.setattr(wa.context, "KICK_AFTER", 0.01) + promise, future = ts.Promise.new() + kicked = [] + + def kick(context): + kicked.append(context) + if len(kicked) == 3: + promise.set_result(7) + return [] + + monkeypatch.setattr(wa.context, "kick", kick) + assert wa.context.wait(future, "ctx") == 7 + assert kicked == ["ctx"] * 3 + stuck = ts.Promise.new() + with pytest.raises(TimeoutError): + wa.context.wait(stuck[1], "ctx", timeout=0.05) + assert wa.context.wait(_ready(), None) == 1 + + +def _ready(): + promise, future = ts.Promise.new() + promise.set_result(1) + return future + + +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_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'" }, From 5ade167c039ee0b600d89739b1142efcccfafed9 Mon Sep 17 00:00:00 2001 From: David Ackerman Date: Fri, 23 Oct 2026 00:20:15 -0400 Subject: [PATCH 2/7] Changelog: write speed at one thread per pool Claude-Session: https://claude.ai/code/session_0141wgF8mkGbjGfBFKfKhp64 --- CHANGELOG.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4d2e9e1..bdf3c24 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,9 +29,10 @@ unit table change. new `wa.thread_budget()`: `OMP_NUM_THREADS` when set (dask sets 1 in each worker), else `LSB_DJOB_NUMPROC` (an LSF job's slots), else the CPUs the process may run on. 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. Writes stay 2 to - 8 times faster than cellmap-analyze's zarr-python writer on NFS, at one scale, measured - on 128^3 uint64 label chunks. + a dask worker's tensorstore now takes the same one thread as its numpy. Writes stay faster + than cellmap-analyze's zarr-python writer on NFS, at one scale on 128^3 uint64 label + chunks: 2.3 to 5 times at one thread per pool (7 threads per process against + zarr-python's 24 to 31), 2 to 7 times with the pools at the machine's CPUs. ## 0.4.0 (2026-10-22) From c3a2ca4b9f86b7c9ccc36173fe858d800c32e597 Mon Sep 17 00:00:00 2001 From: David Ackerman Date: Fri, 23 Oct 2026 00:22:19 -0400 Subject: [PATCH 3/7] Test thread_budget where there is no sched_getaffinity (macOS) Claude-Session: https://claude.ai/code/session_0141wgF8mkGbjGfBFKfKhp64 --- tests/test_access.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_access.py b/tests/test_access.py index dcd5365..f8c58ba 100644 --- a/tests/test_access.py +++ b/tests/test_access.py @@ -117,7 +117,8 @@ def test_contexts_and_retry_timeouts(): def test_the_thread_budget_follows_the_job(monkeypatch): monkeypatch.delenv("OMP_NUM_THREADS", raising=False) monkeypatch.delenv("LSB_DJOB_NUMPROC", raising=False) - assert wa.thread_budget() == len(os.sched_getaffinity(0)) + 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 From d09ad4bc38804ceb36fd615cb48277ebd56b23f6 Mon Sep 17 00:00:00 2001 From: David Ackerman Date: Fri, 23 Oct 2026 00:30:35 -0400 Subject: [PATCH 4/7] Kick every watched context from one watchdog thread, not inside waits chunkmirage and cellmap-flow read through vol.ts directly and with await, so a kick inside worldarray's own waits missed them. Every context the package makes or opens a store on is now held weakly by one daemon thread that gives its decoding and file I/O pools a no-op task once a second; a stranded task runs within about a second whoever waits on it. Waits go back to future.result(timeout). thread_budget also reads SLURM_CPUS_PER_TASK and NSLOTS. cellmap-analyze's read retries of June 2025 were working around the same tensorstore bug. Claude-Session: https://claude.ai/code/session_0141wgF8mkGbjGfBFKfKhp64 --- CHANGELOG.md | 26 ++++---- README.md | 7 ++- docs/spec.md | 45 +++++++------- src/worldarray/array.py | 5 +- src/worldarray/context.py | 123 +++++++++++++++++++++++++------------- src/worldarray/retry.py | 36 +++-------- src/worldarray/store.py | 4 +- src/worldarray/write.py | 4 +- tests/test_access.py | 52 ++++++++-------- 9 files changed, 164 insertions(+), 138 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bdf3c24..d4edb52 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,24 +10,24 @@ unit table change. 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 wait on a data read, a data write or a store open nudges its context's thread - pools after each second the operation is still pending: one no-op task on the decoding - pool and one on the file I/O pool (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 operation now finishes - within about a second at any concurrency, a limit of 1 included. Spec open question 14 - says when this goes. +- 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`. 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 - `LSB_DJOB_NUMPROC` (an LSF job's slots), else the CPUs the process may run on. + the job's slots (`LSB_DJOB_NUMPROC`, `SLURM_CPUS_PER_TASK`, `NSLOTS`), else the CPUs the + process may run on. 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. Writes stay faster than cellmap-analyze's zarr-python writer on NFS, at one scale on 128^3 uint64 label diff --git a/README.md b/README.md index 3481287..bae857c 100644 --- a/README.md +++ b/README.md @@ -89,8 +89,9 @@ s1 = wa.open_pyramid("out.zarr", mode="r+").add_level("s1", factors=2) # place - **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 a waiting operation's thread pools nudged - every second, so a stuck operation recovers or fails loudly instead of hanging the job. + 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 @@ -149,7 +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, LSB_DJOB_NUMPROC or CPUs +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 diff --git a/docs/spec.md b/docs/spec.md index 1ee2a93..f9e0adc 100644 --- a/docs/spec.md +++ b/docs/spec.md @@ -362,8 +362,8 @@ vol.write(box, data, partial=False, retries=None) ; vol.write_chunk(index, data, # 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 wait on a data read, write or open kicks its context's pools after each second - # it is still pending (Behaviour 6) + # every context the package makes or opens a store on is kicked once a second, whoever + # waits on it (Behaviour 6) ``` ## 3. Conventions @@ -650,19 +650,22 @@ this document and that one differ, fix one of them. 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 kick (0.5): tensorstore up to at least 0.1.85 can strand a queued task. A pool + 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). So every wait on a data read, a data write or a store open - gives the context's decoding and file I/O pools one no-op task each (an in-memory array - copy and a local read of a key that does not exist; no store is touched) after every - second the future is still pending. A stranded operation finishes within about a - second, at any concurrency, and the retry stays the outer net. Open question 14 says - when the kick goes. + 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. Its @@ -673,13 +676,14 @@ this document and that one differ, fix one of them. 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 kick - makes safe. + 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. `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; a - `concurrency` of 1 is safe under Behaviour 6's kick. Both are plain tensorstore contexts; a caller may pass its own, at construction or through + `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). @@ -986,9 +990,10 @@ path reopening after a `chdir`; `codecs_of` by path; an all-fill write removing a scalar filling a box through `write_index`, `write` and `write_chunk`, refused when it 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 pending future kicked until it resolves, a wait's timeout, 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. +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. 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`; @@ -1172,12 +1177,12 @@ 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 kick goes.** Behaviour 6's kick 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 +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"), require that release, remove `context.kick` and `KICK_AFTER` - (waits go back to `future.result(timeout)`), and keep the write retries. That release + "ahead" or "identical"), 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. diff --git a/src/worldarray/array.py b/src/worldarray/array.py index fe3320b..35c71bc 100644 --- a/src/worldarray/array.py +++ b/src/worldarray/array.py @@ -589,7 +589,6 @@ def write_index( retries or self.retries, remote=self.remote, what=self.path, - context=self.context, ) def _scalar(self, value: np.ndarray) -> np.ndarray: @@ -736,9 +735,7 @@ def _read_block( transform: Transform | None = None, ) -> np.ndarray: view = self._store[tuple(slice(int(a), int(c)) for a, c in zip(lo, hi))] - data = read_with_retries( - view, retries or self.retries, remote=self.remote, what=self.path, context=self.context - ) + data = read_with_retries(view, retries or self.retries, remote=self.remote, what=self.path) data = np.asarray(data) if not data.dtype.isnative: data = data.astype(data.dtype.newbyteorder("=")) diff --git a/src/worldarray/context.py b/src/worldarray/context.py index 6ca90bb..23a50b3 100644 --- a/src/worldarray/context.py +++ b/src/worldarray/context.py @@ -5,7 +5,8 @@ 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. 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. +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 @@ -20,6 +21,7 @@ import tempfile import threading import time +import weakref from typing import Any import tensorstore as ts @@ -33,58 +35,94 @@ # 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). Any later task on the same pool wakes it, so ``wait`` gives a context's pools -# a no-op task (``kick``) every KICK_AFTER seconds a future is pending. 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 ``kick`` -# (spec, open question 14): `gh api repos/google/tensorstore/compare/ed51718085...vX.Y.Z` -# says "ahead" or "identical" for a release that has it. -KICK_AFTER = 1.0 +# 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. +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: an in-memory - array copy and a local read of a key that does not exist. No store is touched. Returns - the futures, which the caller holds until done (a dropped future may be cancelled).""" - futures = [] + """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: - array = ts.open( - {"driver": "array", "array": [0], "dtype": "uint8"}, context=context - ).result(timeout=KICK_AFTER) - futures.append(array.read()) - kv = ts.KvStore.open({"driver": "file", "path": _KICK_DIR + "/"}, context=context).result( - timeout=KICK_AFTER - ) - futures.append(kv.read(_KICK_KEY)) + 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 - pass - return futures + return [] -def wait(future, context: ts.Context | None, timeout: float | None = None): - """``future.result()``, kicking ``context``'s pools after every ``KICK_AFTER`` seconds it - is still pending; ``TimeoutError`` once ``timeout`` seconds have passed (never when - None). A future that is not a tensorstore one, or no context, is simply waited on.""" - deadline = None if timeout is None else time.monotonic() + timeout - kicks: list = [] +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: - step = KICK_AFTER if deadline is None else min(KICK_AFTER, deadline - time.monotonic()) - try: - return future.result(timeout=max(step, 0.0)) - except TimeoutError: - if deadline is not None and time.monotonic() >= deadline: - raise - if context is not None: - kicks = [f for f in kicks if not f.done()] + kick(context) + 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 - ``LSB_DJOB_NUMPROC`` (an LSF job's slots), else the CPUs this process may run on.""" - for name in ("OMP_NUM_THREADS", "LSB_DJOB_NUMPROC"): + 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: @@ -109,6 +147,7 @@ def shared_context(cache_bytes: int = 0) -> ts.Context: "file_io_concurrency": limit, } ) + watch(ctx) _shared[int(cache_bytes)] = ctx return ctx @@ -117,7 +156,7 @@ def capped_context_spec(cache_bytes: int = 0, concurrency: int = 1, network: boo """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). A ``concurrency`` of 1 is safe: - every wait in the package kicks a pool that stranded a task (see ``KICK_AFTER``).""" + the context is watched (see ``KICK_EVERY``).""" limit = {"limit": int(concurrency)} spec = { "data_copy_concurrency": limit, @@ -134,7 +173,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]: @@ -160,6 +199,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 254b138..352c05f 100644 --- a/src/worldarray/retry.py +++ b/src/worldarray/retry.py @@ -4,8 +4,7 @@ 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). While -an attempt waits, its context's pools are kicked every second (``context.wait``). +jitter. A timed-out future is cancelled before the next attempt (spec, Behaviour 6). """ from __future__ import annotations @@ -18,8 +17,6 @@ import numpy as np -from .context import wait - logger = logging.getLogger(__name__) _TRANSIENT_PREFIXES = ("UNAVAILABLE", "DEADLINE_EXCEEDED", "RESOURCE_EXHAUSTED") @@ -53,20 +50,17 @@ def describe(error: BaseException) -> str: return lines[0][:120] if lines else type(error).__name__ -def with_retries( - start, nbytes: int, retry: Retry, *, remote: bool, what: str, verb: str, context=None -): +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, - ``verb`` ("read", "write") names the operation in the warning and the error, and - ``context`` (the operation's) is kicked while it waits.""" + 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, attempts + 1): future = start() try: - return wait(future, context, timeout * attempt) + return future.result(timeout=timeout * attempt) except TimeoutError as e: future.cancel() last = e @@ -95,31 +89,19 @@ def with_retries( raise last -def read_with_retries( - view, retry: Retry, *, remote: bool, what: str = "", context=None -) -> np.ndarray: +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", context=context - ) + 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 = "", context=None + 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", - context=context, - ) + 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 3daa5c2..b8aee74 100644 --- a/src/worldarray/store.py +++ b/src/worldarray/store.py @@ -20,7 +20,7 @@ import tensorstore as ts -from .context import wait +from .context import watch logger = logging.getLogger(__name__) T = TypeVar("T") @@ -440,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 wait(ts.open(spec, read=True, write=write, context=context), context) + 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 43d8d9c..01cc960 100644 --- a/src/worldarray/write.py +++ b/src/worldarray/write.py @@ -27,7 +27,7 @@ import numpy as np import tensorstore as ts -from .context import wait +from .context import watch from .geometry import snap_integral from .placement import spatial_axes from .store import ( @@ -768,7 +768,7 @@ def attempt(kvstore: dict) -> ts.TensorStore: separator, store_fill, ) - return wait(ts.open(spec, create=True, open=False, context=context), context) + 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 f8c58ba..2eea1dd 100644 --- a/tests/test_access.py +++ b/tests/test_access.py @@ -2,8 +2,11 @@ from __future__ import annotations +import gc import logging import os +import time +import weakref import numpy as np import pytest @@ -127,9 +130,12 @@ def test_the_thread_budget_follows_the_job(monkeypatch): 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": 12} + assert spec["data_copy_concurrency"] == spec["file_io_concurrency"] == {"limit": 6} assert spec["cache_pool"] == {"total_bytes_limit": 5} @@ -140,31 +146,27 @@ def test_a_kick_runs_on_the_contexts_own_pools(): assert [type(f.result(timeout=10)).__name__ for f in futures] == ["ndarray", "ReadResult"] -def test_a_pending_future_is_kicked_until_it_resolves(monkeypatch): - # tensorstore's lost wakeup: the stranded task runs once anything else uses its pool - monkeypatch.setattr(wa.context, "KICK_AFTER", 0.01) - promise, future = ts.Promise.new() +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 = [] - - def kick(context): - kicked.append(context) - if len(kicked) == 3: - promise.set_result(7) - return [] - - monkeypatch.setattr(wa.context, "kick", kick) - assert wa.context.wait(future, "ctx") == 7 - assert kicked == ["ctx"] * 3 - stuck = ts.Promise.new() - with pytest.raises(TimeoutError): - wa.context.wait(stuck[1], "ctx", timeout=0.05) - assert wa.context.wait(_ready(), None) == 1 - - -def _ready(): - promise, future = ts.Promise.new() - promise.set_result(1) - return future + 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: From 04bc94a70468ff801ed43efa45a067a2f4c252ca Mon Sep 17 00:00:00 2001 From: David Ackerman Date: Fri, 23 Oct 2026 00:37:05 -0400 Subject: [PATCH 5/7] Spec: why the shared context caps threads by default Claude-Session: https://claude.ai/code/session_0141wgF8mkGbjGfBFKfKhp64 --- CHANGELOG.md | 3 ++- docs/spec.md | 9 ++++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d4edb52..a2b5577 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,7 +27,8 @@ unit table change. - 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. + 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. Writes stay faster than cellmap-analyze's zarr-python writer on NFS, at one scale on 128^3 uint64 label diff --git a/docs/spec.md b/docs/spec.md index f9e0adc..337210d 100644 --- a/docs/spec.md +++ b/docs/spec.md @@ -678,7 +678,14 @@ this document and that one differ, fix one of them. 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. + `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 From 079ac1eed90fc85fef9e4741d2c0577d546a93a6 Mon Sep 17 00:00:00 2001 From: David Ackerman Date: Fri, 23 Oct 2026 00:56:13 -0400 Subject: [PATCH 6/7] Reproduce tensorstore's lost wakeup on purpose; opt-in stress test Tasks aimed 19.5-20.5 ms after a one-thread pool's last one strand about 0.3% of the time on a plain context (tensorstore 0.1.75, 0.1.78 and 0.1.85) and never on a watched one, slowest under a second. cellmap-analyze's block loop on its 0.6.1 context gave 29 IndexErrors and 55 writes stuck past 20 s on worldarray 0.4.0, none on this branch. WORLDARRAY_STRESS=1 runs it; it is also the check for removing the watchdog (open question 14). Claude-Session: https://claude.ai/code/session_0141wgF8mkGbjGfBFKfKhp64 --- CHANGELOG.md | 5 ++- docs/spec.md | 14 ++++++-- src/worldarray/context.py | 3 +- tests/test_lost_wakeup.py | 76 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 93 insertions(+), 5 deletions(-) create mode 100644 tests/test_lost_wakeup.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a2b5577..b21fa90 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,7 +20,10 @@ unit table change. 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`. Spec open question 14 says when this goes. + 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. diff --git a/docs/spec.md b/docs/spec.md index 337210d..927eb34 100644 --- a/docs/spec.md +++ b/docs/spec.md @@ -1000,7 +1000,13 @@ nothing resolves) cancelled and retried, the warning naming `TimeoutError`, and 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. +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`; @@ -1188,8 +1194,10 @@ over funlib geometry, by this spec's own assignment. 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"), require that release, remove `watch`, `kick` and the watchdog - thread from `context.py`, and keep the write retries. That release + "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. diff --git a/src/worldarray/context.py b/src/worldarray/context.py index 23a50b3..0508557 100644 --- a/src/worldarray/context.py +++ b/src/worldarray/context.py @@ -43,7 +43,8 @@ # "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. +# 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" 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)" + ) From 95ab2f62dcdbc7d7df6aa0a7691bb553753c5127 Mon Sep 17 00:00:00 2001 From: David Ackerman Date: Fri, 23 Oct 2026 10:13:47 -0400 Subject: [PATCH 7/7] Date 0.5.0 and correct its speed claim The "2.3 to 5 times" figure counted `import zarr` (2.7 s) as write time. Measured against cellmap-analyze 0.5.0's own block I/O instead. Claude-Session: https://claude.ai/code/session_0141wgF8mkGbjGfBFKfKhp64 --- CHANGELOG.md | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b21fa90..46380e1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Changelog -## 0.5.0 (unreleased) +## 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 @@ -33,10 +33,12 @@ unit table change. 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. Writes stay faster - than cellmap-analyze's zarr-python writer on NFS, at one scale on 128^3 uint64 label - chunks: 2.3 to 5 times at one thread per pool (7 threads per process against - zarr-python's 24 to 31), 2 to 7 times with the pools at the machine's CPUs. + 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)