worldarray 0.4.0__tar.gz → 0.5.0__tar.gz

{worldarray-0.4.0 → worldarray-0.5.0}/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

{worldarray-0.4.0 → worldarray-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: worldarray
-Version: 0.4.0
+Version: 0.5.0
 Summary: Read zarr, N5 and Neuroglancer precomputed volumes in world coordinates through tensorstore: correct voxel placement, pyramid levels, padded region reads, resampling. numpy + tensorstore only.
 Project-URL: Homepage, https://janeliascicomp.github.io/worldarray/
 Project-URL: Documentation, https://janeliascicomp.github.io/worldarray/
@@ -145,9 +145,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
@@ -206,6 +209,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
 
@@ -233,7 +237,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

{worldarray-0.4.0 → worldarray-0.5.0}/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

{worldarray-0.4.0 → worldarray-0.5.0}/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
 

{worldarray-0.4.0 → worldarray-0.5.0}/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" }

{worldarray-0.4.0 → worldarray-0.5.0}/src/worldarray/__init__.py

@@ -12,7 +12,7 @@ from __future__ import annotations
 
 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 @@ from .write import (
     parse_chunk_key,
 )
 
-__version__ = "0.4.0"
+__version__ = "0.5.0"
 
 __all__ = [
     "ArrayMeta",
@@ -62,6 +62,7 @@ __all__ = [
     "open_tensorstore",
     "select_level",
     "shared_context",
+    "thread_budget",
     "store",
     "units",
     "write",

{worldarray-0.4.0 → worldarray-0.5.0}/src/worldarray/array.py

@@ -35,7 +35,7 @@ from .placement import (
     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 @@ class WorldArray:
     # --- 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 @@ class WorldArray:
         read-modify-write of the touched units, unsafe while anything else writes the same
         unit. ``data`` may be a scalar, which fills the box: it is converted to the array's
         dtype only when the value survives (``0`` into ``uint8``, not ``300`` or ``0.5``).
-        No other cast, no padding, no resampling; the write commits before returning."""
+        No other cast, no padding, no resampling; the write commits before returning, timed
+        out and retried as reads are (``retries``, else the handle's)."""
         if not self.writable:
             raise ValueError(
                 f"{self.path} was opened read-only; open it with mode='r+' or create it"
@@ -577,7 +583,13 @@ class WorldArray:
                         f"{self.path}: write [{b}, {e}) on axis {axis} ({self.axes[axis]}) is not whole write "
                         f"units of {u}; align the box to write_shape {self.write_shape} or pass partial=True"
                     )
-        self._store[tuple(slice(b, e) for b, e in zip(begin, end))].write(data).result()
+        write_with_retries(
+            self._store[tuple(slice(b, e) for b, e in zip(begin, end))],
+            data,
+            retries or self.retries,
+            remote=self.remote,
+            what=self.path,
+        )
 
     def _scalar(self, value: np.ndarray) -> np.ndarray:
         """A 0-d ``value`` in the array's dtype, or ``ValueError`` when converting changes it:
@@ -595,7 +607,9 @@ class WorldArray:
             raise ValueError(f"{self.path}: {value.item()!r} does not fit {dtype}")
         return out
 
-    def write(self, box, data: np.ndarray, partial: bool = False) -> None:
+    def write(
+        self, box, data: np.ndarray, partial: bool = False, retries: Retry | None = None
+    ) -> None:
         """Write ``data`` over a world box: the voxels Behaviour 2's centre rule selects, then
         ``write_index``. A box reaching outside the array raises; nothing is padded or
         resampled on write."""
@@ -606,7 +620,7 @@ class WorldArray:
                     f"{self.path}: the box reaches outside the array on axis {axis} ({self.axes[axis]}): "
                     f"voxels [{b}, {e}) of {n}; nothing is padded on write"
                 )
-        self.write_index(begin, end, data, partial=partial)
+        self.write_index(begin, end, data, partial=partial, retries=retries)
 
     def blocks(self, block_shape=None, context=0, box=None, clip: bool = True):
         """The array cut into blocks (Behaviour 17): write boxes tiling the stored grid (or a
@@ -623,7 +637,9 @@ class WorldArray:
             clip=clip,
         )
 
-    def write_chunk(self, index: Sequence[int], data: np.ndarray) -> None:
+    def write_chunk(
+        self, index: Sequence[int], data: np.ndarray, retries: Retry | None = None
+    ) -> None:
         """Write the whole write unit at ``index`` on the storage grid (the shard of a sharded
         array); ``data`` has that unit's shape, clipped at the array's end."""
         unit, shape = self.write_shape, self._grid.shape
@@ -634,7 +650,7 @@ class WorldArray:
         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,

worldarray-0.5.0/src/worldarray/context.py

@@ -0,0 +1,205 @@
+"""Tensorstore contexts: one shared per cache budget, or a capped one per call.
+
+``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. Its decoding and file I/O
+pools are sized by ``thread_budget()``, from the variables a cluster job sets, where
+tensorstore alone would size them to every CPU on the node. Every context here, and every
+one a store is opened on, is watched for tensorstore's lost wakeup (``watch``).
+
+``capped_context`` is mesh-n-bone's worker-process context: every tensorstore concurrency
+resource limited, so dozens of worker processes on one host do not multiply into thousands
+of threads. It returns a **new** context each call, as those workers inline the spec into
+every open today; pass the same object twice to share pools (spec, Behaviour 7).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+import threading
+import time
+import weakref
+from typing import Any
+
+import tensorstore as ts
+
+_shared: dict[int, ts.Context] = {}
+_rebuilt: dict[str, ts.Context] = {}
+_lock = threading.Lock()
+
+
+# tensorstore's lost wakeup. In every release up to at least 0.1.85, a pool worker idle for
+# 20 ms exits, and a task queued as it exits gets no thread while the pool counts itself
+# full: the task never runs and its future waits forever, with no tensorstore thread alive.
+# A pool limited to one thread meets it most (cellmap-analyze's workers on /nrs,
+# 2026-10-22; its read retries of June 2025 were working around the same thing). Any later
+# task on the same pool wakes it. So every context the package makes or opens a store on is
+# watched: a daemon thread gives each one's decoding and file I/O pools a no-op task every
+# KICK_EVERY seconds, and a stranded task, whoever waits on it (worldarray, ``vol.ts``,
+# ``await``), runs within that. Fixed upstream in google/tensorstore ed51718085 (2026-09-30,
+# "Fix lost wakeups in TaskGroup when workers block or exit"). Once a release contains it,
+# require that release and remove the watchdog (spec, open question 14):
+# `gh api repos/google/tensorstore/compare/ed51718085...vX.Y.Z` says "ahead" or "identical"
+# for a release that has it, and `WORLDARRAY_STRESS=1 pytest tests/test_lost_wakeup.py -s`
+# then strands nothing on a plain context.
+KICK_EVERY = 1.0
+_KICK_DIR = tempfile.gettempdir()
+_KICK_KEY = "worldarray-kick-no-such-key"
+_watched: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
+_watch_lock = threading.Lock()
+_watchdog: threading.Thread | None = None
+
+
+def _kickers(context: ts.Context) -> tuple:
+    """An in-memory array and a local kvstore on ``context``: reading the one copies on its
+    decoding pool, reading a missing key from the other runs on its file I/O pool."""
+    array = ts.open({"driver": "array", "array": [0], "dtype": "uint8"}, context=context)
+    kv = ts.KvStore.open({"driver": "file", "path": _KICK_DIR + "/"}, context=context)
+    return array.result(timeout=10), kv.result(timeout=10)
+
+
+def kick(context: ts.Context) -> list:
+    """One no-op task on each of ``context``'s decoding and file I/O pools; no store is
+    touched. Returns the futures, which the caller holds until done (a dropped future may
+    be cancelled before it is queued)."""
+    try:
+        with _watch_lock:
+            kickers = _watched.get(context)
+        if kickers is None:
+            kickers = _kickers(context)
+            with _watch_lock:
+                if context in _watched:
+                    _watched[context] = kickers
+        array, kv = kickers
+        return [array.read(), kv.read(_KICK_KEY)]
+    except Exception:  # noqa: BLE001  a kick that cannot start changes nothing
+        return []
+
+
+def watch(context: ts.Context | None) -> ts.Context | None:
+    """Kick ``context``'s pools every ``KICK_EVERY`` seconds for as long as it lives, starting
+    the watchdog thread on first use. Returns ``context``."""
+    global _watchdog
+    if context is None:
+        return None
+    with _watch_lock:
+        if context not in _watched:
+            _watched[context] = None
+        if _watchdog is None or not _watchdog.is_alive():
+            _watchdog = threading.Thread(target=_watch, name="worldarray-kick", daemon=True)
+            _watchdog.start()
+    return context
+
+
+def _watch() -> None:
+    held: list = []
+    while True:
+        time.sleep(KICK_EVERY)
+        with _watch_lock:
+            contexts = list(_watched.keys())
+        held = [f for f in held if not f.done()]
+        for context in contexts:
+            held += kick(context)
+        del contexts
+
+
+def _forget_watchdog() -> None:
+    global _watchdog
+    _watchdog = None  # a forked child has no such thread; the next watch() starts one
+
+
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(after_in_child=_forget_watchdog)
+
+
+_SLOTS = ("OMP_NUM_THREADS", "LSB_DJOB_NUMPROC", "SLURM_CPUS_PER_TASK", "NSLOTS")
+
+
+def thread_budget() -> int:
+    """Threads for this process's tensorstore decoding and file I/O: ``OMP_NUM_THREADS``
+    when set (dask sets 1 in every worker it spawns; a job script sets its slots), else the
+    job's slots (LSF's ``LSB_DJOB_NUMPROC``, SLURM's ``SLURM_CPUS_PER_TASK``, SGE's
+    ``NSLOTS``), else the CPUs this process may run on."""
+    for name in _SLOTS:
+        try:
+            n = int(os.environ.get(name, ""))
+        except ValueError:
+            continue
+        if n > 0:
+            return n
+    try:
+        return len(os.sched_getaffinity(0))
+    except (AttributeError, OSError):
+        return os.cpu_count() or 1
+
+
+def shared_context(cache_bytes: int = 0) -> ts.Context:
+    with _lock:
+        ctx = _shared.get(int(cache_bytes))
+        if ctx is None:
+            limit = {"limit": thread_budget()}
+            ctx = ts.Context(
+                {
+                    "cache_pool": {"total_bytes_limit": int(cache_bytes)},
+                    "data_copy_concurrency": limit,
+                    "file_io_concurrency": limit,
+                }
+            )
+            watch(ctx)
+            _shared[int(cache_bytes)] = ctx
+        return ctx
+
+
+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). A ``concurrency`` of 1 is safe:
+    the context is watched (see ``KICK_EVERY``)."""
+    limit = {"limit": int(concurrency)}
+    spec = {
+        "data_copy_concurrency": limit,
+        "file_io_concurrency": limit,
+        "cache_pool": {"total_bytes_limit": int(cache_bytes)},
+    }
+    if network:
+        spec.update(
+            gcs_request_concurrency=limit,
+            s3_request_concurrency=limit,
+            http_request_concurrency=limit,
+        )
+    return spec
+
+
+def capped_context(cache_bytes: int = 0, concurrency: int = 1, network: bool = True) -> ts.Context:
+    return watch(ts.Context(capped_context_spec(cache_bytes, concurrency, network)))
+
+
+def recipe_of(context: ts.Context) -> tuple[str, Any]:
+    """How another process rebuilds ``context``: ``("shared", cache_bytes)`` for a shared
+    context, else ``("spec", json)`` from the context's own spec, so a capped context stays
+    capped across a pickle (Behaviour 11, mesh-n-bone's workers)."""
+    with _lock:
+        for cache_bytes, ctx in _shared.items():
+            if ctx is context:
+                return ("shared", cache_bytes)
+    spec = context.spec
+    return ("spec", json.dumps(spec.to_json() if spec is not None else {}, sort_keys=True))
+
+
+def from_recipe(recipe: tuple[str, Any] | None) -> ts.Context | None:
+    """``recipe_of`` inverted; one rebuilt context per spec per process, so handles that
+    shared a context keep sharing one."""
+    if recipe is None:
+        return None
+    kind, value = recipe
+    if kind == "shared":
+        return shared_context(int(value))
+    with _lock:
+        ctx = _rebuilt.get(value)
+        if ctx is None:
+            ctx = watch(ts.Context(ts.Context.Spec(json.loads(value))))
+            _rebuilt[value] = ctx
+        return ctx

{worldarray-0.4.0 → worldarray-0.5.0}/src/worldarray/retry.py

@@ -1,7 +1,7 @@
-"""Retrying a data read on a transient failure.
+"""Retrying a data read or write on a transient failure.
 
 mesh-n-bone's policy, which its tests pin: attempt ``k`` (from 1) waits ``timeout * k``
-seconds; ``timeout`` is given, or sized from the read and where it comes from:
+seconds; ``timeout`` is given, or sized from the bytes moved and where they go:
 ``clamp(MiB / 16, 5, 120)`` for a remote store and ``clamp(MiB / 256, 5, 30)`` for a local
 one. Between attempts the sleep is ``base_delay * 1.3 ** (k - 1)`` plus up to one second of
 jitter. A timed-out future is cancelled before the next attempt (spec, Behaviour 6).
@@ -43,15 +43,22 @@ def is_transient(error: BaseException) -> bool:
     return message.startswith(_TRANSIENT_PREFIXES) or bool(_HTTP_TRANSIENT.search(message))
 
 
-def read_with_retries(view, retry: Retry, *, remote: bool, what: str = "") -> np.ndarray:
-    """``view.read()`` as a numpy array, retried per ``retry``. ``view`` is a tensorstore
-    (or any object whose ``read()`` returns a future with ``result(timeout)`` and
-    ``cancel()``)."""
-    nbytes = int(np.prod(view.shape, dtype=np.int64)) * np.dtype(view.dtype.numpy_dtype).itemsize
+def describe(error: BaseException) -> str:
+    """The first line of ``error``'s message, or its type's name when the message is empty
+    (a timed-out tensorstore future raises ``TimeoutError()``)."""
+    lines = str(error).splitlines()
+    return lines[0][:120] if lines else type(error).__name__
+
+
+def with_retries(start, nbytes: int, retry: Retry, *, remote: bool, what: str, verb: str):
+    """``start().result()``, retried per ``retry``. ``start`` begins the operation and
+    returns a future with ``result(timeout)`` and ``cancel()``; ``nbytes`` sizes the timeout
+    and ``verb`` ("read", "write") names the operation in the warning and the error."""
     timeout = retry.timeout_for(nbytes, remote)
+    attempts = max(1, retry.max_attempts)
     last: BaseException | None = None
-    for attempt in range(1, max(1, retry.max_attempts) + 1):
-        future = view.read()
+    for attempt in range(1, attempts + 1):
+        future = start()
         try:
             return future.result(timeout=timeout * attempt)
         except TimeoutError as e:
@@ -61,16 +68,40 @@ def read_with_retries(view, retry: Retry, *, remote: bool, what: str = "") -> np
             if not is_transient(e):
                 raise
             last = e
-        if attempt < retry.max_attempts:
+        if attempt < attempts:
             delay = retry.base_delay * 1.3 ** (attempt - 1) + random.random()
             logger.warning(
-                "read %s failed on attempt %d/%d (%s); retrying in %.1fs",
+                "%s %s failed on attempt %d/%d (%s); retrying in %.1fs",
+                verb,
                 what,
                 attempt,
-                retry.max_attempts,
-                str(last).splitlines()[0][:120],
+                attempts,
+                describe(last),
                 delay,
             )
             time.sleep(delay)
     assert last is not None
+    if isinstance(last, TimeoutError):
+        raise TimeoutError(
+            f"{verb} {what} timed out on all {attempts} attempts, the last after "
+            f"{timeout * attempts:g}s"
+        ) from last
     raise last
+
+
+def read_with_retries(view, retry: Retry, *, remote: bool, what: str = "") -> np.ndarray:
+    """``view.read()`` as a numpy array, retried per ``retry``. ``view`` is a tensorstore
+    (or any object whose ``read()`` returns a future with ``result(timeout)`` and
+    ``cancel()``)."""
+    nbytes = int(np.prod(view.shape, dtype=np.int64)) * np.dtype(view.dtype.numpy_dtype).itemsize
+    return with_retries(view.read, nbytes, retry, remote=remote, what=what, verb="read")
+
+
+def write_with_retries(
+    view, data: np.ndarray, retry: Retry, *, remote: bool, what: str = ""
+) -> None:
+    """``view.write(data)`` committed, retried per ``retry`` like a read. A timed-out write
+    is cancelled, but may still land; it writes the same bytes, so the next attempt is
+    safe."""
+    nbytes = int(np.prod(view.shape, dtype=np.int64)) * np.dtype(view.dtype.numpy_dtype).itemsize
+    with_retries(lambda: view.write(data), nbytes, retry, remote=remote, what=what, verb="write")

{worldarray-0.4.0 → worldarray-0.5.0}/src/worldarray/store.py

@@ -20,6 +20,8 @@ from typing import Any, TypeVar
 
 import tensorstore as ts
 
+from .context import watch
+
 logger = logging.getLogger(__name__)
 T = TypeVar("T")
 
@@ -438,7 +440,7 @@ def open_array(
             spec["assume_metadata"] = True
         if node.format == "precomputed" and scale_index is not None:
             spec["scale_index"] = scale_index
-        return ts.open(spec, read=True, write=write, context=context).result()
+        return ts.open(spec, read=True, write=write, context=watch(context)).result()
 
     store = with_access(kvstore_specs(node.path), attempt)
     if node.format in ("n5", "precomputed"):

{worldarray-0.4.0 → worldarray-0.5.0}/src/worldarray/write.py

@@ -27,6 +27,7 @@ from typing import Any
 import numpy as np
 import tensorstore as ts
 
+from .context import watch
 from .geometry import snap_integral
 from .placement import spatial_axes
 from .store import (
@@ -767,7 +768,7 @@ def create_array(
             separator,
             store_fill,
         )
-        return ts.open(spec, create=True, open=False, context=context).result()
+        return ts.open(spec, create=True, open=False, context=watch(context)).result()
 
     try:
         with_access(kvstore_specs(path), attempt)

{worldarray-0.4.0 → worldarray-0.5.0}/tests/test_access.py

@@ -2,14 +2,19 @@
 
 from __future__ import annotations
 
+import gc
+import logging
 import os
+import time
+import weakref
 
 import numpy as np
 import pytest
+import tensorstore as ts
 
 import worldarray as wa
 from conftest import pattern, write_zarr2, write_zarr2_group
-from worldarray import store
+from worldarray import retry, store
 
 
 @pytest.fixture(autouse=True)
@@ -112,6 +117,94 @@ def test_contexts_and_retry_timeouts():
     assert wa.Retry(timeout=7).timeout_for(10**12, remote=True) == 7.0
 
 
+def test_the_thread_budget_follows_the_job(monkeypatch):
+    monkeypatch.delenv("OMP_NUM_THREADS", raising=False)
+    monkeypatch.delenv("LSB_DJOB_NUMPROC", raising=False)
+    allowed = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count()
+    assert wa.thread_budget() == allowed  # macOS has no affinity call
+    monkeypatch.setenv("LSB_DJOB_NUMPROC", "12")
+    assert wa.thread_budget() == 12
+    monkeypatch.setenv("OMP_NUM_THREADS", "24")  # a job script's 2 * slots
+    assert wa.thread_budget() == 24
+    monkeypatch.setenv("OMP_NUM_THREADS", "1")  # what dask sets in each worker
+    assert wa.thread_budget() == 1
+    monkeypatch.setenv("OMP_NUM_THREADS", "x")
+    assert wa.thread_budget() == 12
+    monkeypatch.delenv("LSB_DJOB_NUMPROC")
+    monkeypatch.setenv("SLURM_CPUS_PER_TASK", "6")
+    assert wa.thread_budget() == 6
+    monkeypatch.setattr(wa.context, "_shared", {})
+    spec = wa.shared_context(5).spec.to_json()
+    assert spec["data_copy_concurrency"] == spec["file_io_concurrency"] == {"limit": 6}
+    assert spec["cache_pool"] == {"total_bytes_limit": 5}
+
+
+def test_a_kick_runs_on_the_contexts_own_pools():
+    ctx = wa.capped_context(concurrency=1)
+    futures = wa.context.kick(ctx)
+    assert len(futures) == 2
+    assert [type(f.result(timeout=10)).__name__ for f in futures] == ["ndarray", "ReadResult"]
+
+
+def test_every_context_the_package_touches_is_kicked(monkeypatch, tmp_path):
+    # tensorstore's lost wakeup: a stranded task runs once anything else uses its pool, so
+    # a watched context's pools get a no-op task every KICK_EVERY seconds, whoever waits
+    kicked = []
+    monkeypatch.setattr(wa.context, "_watched", weakref.WeakKeyDictionary())
+    monkeypatch.setattr(wa.context, "kick", lambda c: kicked.append(id(c)) or [])
+    monkeypatch.setattr(wa.context, "KICK_EVERY", 0.01)
+    monkeypatch.setattr(wa.context, "_watchdog", None)
+    capped = wa.capped_context(concurrency=1)
+    own = ts.Context({"data_copy_concurrency": {"limit": 1}})
+    write_zarr2(tmp_path / "a.zarr", pattern((4, 4)), chunks=(2, 2))
+    wa.open_tensorstore(str(tmp_path / "a.zarr"), context=own)
+    deadline = time.monotonic() + 5
+    while not {id(capped), id(own)} <= set(kicked) and time.monotonic() < deadline:
+        time.sleep(0.01)
+    assert {id(capped), id(own)} <= set(kicked)
+    assert wa.context._watchdog.daemon and wa.context._watchdog.name == "worldarray-kick"
+    n = len(wa.context._watched)
+    del capped
+    gc.collect()
+    assert len(wa.context._watched) == n - 1  # a context that died is no longer watched
+
+
+class _StalledView:
+    """A view whose first ``stalls`` reads never finish (a future nothing resolves, as a
+    read stuck on NFS looks from Python); later reads return ``0, 1, 2, 3``."""
+
+    shape, dtype = (4,), ts.uint8
+
+    def __init__(self, stalls):
+        self.stalls, self.held = stalls, []
+
+    def read(self):
+        promise, future = ts.Promise.new()
+        if len(self.held) < self.stalls:
+            self.held.append((promise, future))
+        else:
+            promise.set_result(np.arange(4, dtype=np.uint8))
+        return future
+
+
+def test_a_timed_out_read_is_cancelled_and_retried(monkeypatch, caplog):
+    # tensorstore's TimeoutError has an empty message; the warning names its type
+    # (salamander's IndexError on /nrs under 0.4.0)
+    monkeypatch.setattr(retry.time, "sleep", lambda s: None)
+    view = _StalledView(stalls=2)
+    with caplog.at_level(logging.WARNING, logger="worldarray.retry"):
+        got = retry.read_with_retries(
+            view, wa.Retry(timeout=0.01, max_attempts=3), remote=False, what="v"
+        )
+    assert got.tolist() == [0, 1, 2, 3]
+    assert [f.cancelled() for _, f in view.held] == [True, True]
+    assert "read v failed on attempt 1/3 (TimeoutError)" in caplog.text
+    with pytest.raises(TimeoutError, match="read v timed out on all 2 attempts"):
+        retry.read_with_retries(
+            _StalledView(stalls=9), wa.Retry(timeout=0.01, max_attempts=2), remote=False, what="v"
+        )
+
+
 def test_metadata_reads_carry_capped_retries():
     spec = store.metadata_spec({"driver": "s3", "bucket": "b", "path": "p/"})
     assert spec["context"]["s3_request_retries"] == store.METADATA_RETRIES

worldarray-0.5.0/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)"
+        )

{worldarray-0.4.0 → worldarray-0.5.0}/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)

{worldarray-0.4.0 → worldarray-0.5.0}/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'" },

worldarray-0.4.0/src/worldarray/context.py

@@ -1,84 +0,0 @@
-"""Tensorstore contexts: one shared per cache budget, or a capped one per call.
-
-``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.
-
-``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
-of threads. It returns a **new** context each call, as those workers inline the spec into
-every open today; pass the same object twice to share pools (spec, Behaviour 7).
-"""
-
-from __future__ import annotations
-
-import json
-import threading
-from typing import Any
-
-import tensorstore as ts
-
-_shared: dict[int, ts.Context] = {}
-_rebuilt: dict[str, ts.Context] = {}
-_lock = threading.Lock()
-
-
-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)}})
-            _shared[int(cache_bytes)] = ctx
-        return ctx
-
-
-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)."""
-    limit = {"limit": int(concurrency)}
-    spec = {
-        "data_copy_concurrency": limit,
-        "file_io_concurrency": limit,
-        "cache_pool": {"total_bytes_limit": int(cache_bytes)},
-    }
-    if network:
-        spec.update(
-            gcs_request_concurrency=limit,
-            s3_request_concurrency=limit,
-            http_request_concurrency=limit,
-        )
-    return spec
-
-
-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))
-
-
-def recipe_of(context: ts.Context) -> tuple[str, Any]:
-    """How another process rebuilds ``context``: ``("shared", cache_bytes)`` for a shared
-    context, else ``("spec", json)`` from the context's own spec, so a capped context stays
-    capped across a pickle (Behaviour 11, mesh-n-bone's workers)."""
-    with _lock:
-        for cache_bytes, ctx in _shared.items():
-            if ctx is context:
-                return ("shared", cache_bytes)
-    spec = context.spec
-    return ("spec", json.dumps(spec.to_json() if spec is not None else {}, sort_keys=True))
-
-
-def from_recipe(recipe: tuple[str, Any] | None) -> ts.Context | None:
-    """``recipe_of`` inverted; one rebuilt context per spec per process, so handles that
-    shared a context keep sharing one."""
-    if recipe is None:
-        return None
-    kind, value = recipe
-    if kind == "shared":
-        return shared_context(int(value))
-    with _lock:
-        ctx = _rebuilt.get(value)
-        if ctx is None:
-            ctx = ts.Context(ts.Context.Spec(json.loads(value)))
-            _rebuilt[value] = ctx
-        return ctx