0.5.0: retry writes, watchdog for tensorstore's lost wakeup, size threads to the job

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

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

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
 

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" }

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",