0.6.0: shared_context(threads=)

CHANGELOG.md

@@ -1,5 +1,22 @@
 # Changelog
 
+## 0.6.0 (2026-10-23)
+
+One new argument, at chunkmirage's request. No placement change and no unit table change;
+nothing changes for a caller that does not pass it.
+
+- `wa.shared_context(cache_bytes, threads=)`: `None` (the default) limits the decoding and
+  file I/O pools to `thread_budget()` as in 0.5, a positive number limits them to that many,
+  and `0` writes no limit, so tensorstore sizes them itself (every CPU on the node). Each
+  `(cache_bytes, threads)` is one watched context for the process, so a server can share
+  one cache pool across every store it opens at the thread count it chose, without building
+  and watching its own context. A pickled handle on such a context comes back on the same
+  one; `threads=None` still lets the other process size by its own job.
+- Documented what a pool of one thread costs (spec Behaviour 7, `capped_context_spec`):
+  nothing on one-chunk blocks; on 8-chunk blocks, half the speed of 4 in a two-slot job
+  but only 7 to 19% slower in a one-slot one (c-elegans on /nrs), so a worker per slot
+  loses little at the default.
+
 ## 0.5.0 (2026-10-23)
 
 Two bugs that hung or crashed cellmap-analyze 0.6 on /nrs (salamander's and c-elegans'

README.md

@@ -149,8 +149,9 @@ wa.open_pyramid(path, ...) -> Pyramid          # .names .levels .by_name() .sele
 wa.metadata(path, level=None, units=None, open=True) -> ArrayMeta   # open=False: from documents alone
 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.shared_context(cache_bytes=0, threads=None) ; wa.capped_context(cache_bytes=0, concurrency=1, network=True) ; wa.Retry(...)
+wa.thread_budget()                              # the shared context's default thread limit: OMP_NUM_THREADS, the job's slots, or CPUs;
+                                                #   threads=N limits it to N, threads=0 leaves tensorstore's own pool sizes
 wa.list_chunk_keys(path, separator=".")        # chunk indices from keys alone, no metadata needed
 wa.Box ; wa.units.TABLE
 

docs/spec.md

@@ -185,7 +185,8 @@ wa.open_tensorstore(path, level=None, context=None) -> tensorstore.TensorStore
 wa.corner_to_centre(corner, voxel_size, axes=None) ; wa.default_axes(ndim) ; wa.index_range(begin, end, voxel_size, translation)
 wa.units.TABLE      # the unit tables, public: {"length": {...metres per unit}, "time": {...seconds per unit}}
 wa.select_level(levels, voxel_size, mode) ; wa.Box ; wa.Retry(max_attempts=10, timeout=None, base_delay=1.0)
-wa.shared_context(cache_bytes=0) ; wa.capped_context(cache_bytes=0, concurrency=1, network=True) ; wa.capped_context_spec(...)
+wa.shared_context(cache_bytes=0, threads=None) ; wa.capped_context(cache_bytes=0, concurrency=1, network=True) ; wa.capped_context_spec(...)
+                    # threads: None = thread_budget(), N = that many, 0 = tensorstore's own pool sizes (0.6)
                     # network=False caps decoding and local file I/O only, leaving tensorstore's request pools alone
 wa.list_chunk_keys(path, separator=".", ndim=None)   # chunk indices present under a path, from its keys alone: no
                     #   metadata document needed, so a directory of chunk files with no .zarray lists too
@@ -229,9 +230,11 @@ returned one would raise on inputs cellmap-analyze reads every day. A consumer t
 in integer Rois builds them from `corner`, `voxel_size` and `shape` in its own subclass,
 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. Its
-decoding and file I/O pools are sized by `thread_budget()` (0.5, Behaviour 7).
+`shared_context(cache_bytes, threads=None)` is one `tensorstore.Context` per budget and
+thread limit for the process, so 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), or by `threads` when given, `0` leaving tensorstore's own sizes
+(0.6, 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
@@ -684,13 +687,27 @@ this document and that one differ, fix one of them.
    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).
+   `bsub -n 1` say, asks for more slots or says so: `shared_context(cache_bytes,
+   threads=N)` limits both pools to `N`, and `threads=0` writes no limit, leaving
+   tensorstore's own sizes (decided 2026-10-23, over chunkmirage's objection, once the
+   watchdog had removed the hang risk; the argument is 0.6, at chunkmirage's request, so
+   its server keeps one watched context per cache budget for every store it opens
+   instead of building its own). Each `(cache_bytes, threads)` is one context for the
+   process, made on first use and watched; `threads=None` and the number
+   `thread_budget()` returns are two contexts, as a pickle carries `None` for the other
+   process to size by its own job and a number as that number.
    `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 watchdog. 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, but a read or write spanning
+   several chunks then handles them one at a time. On /nrs, from one process alternating
+   limits op by op (c-elegans, 2026-10-23): a one-chunk block (128^3 uint64) read in 18 to
+   19 ms and wrote in 51 ms at both 1 and 4. An 8-chunk block (256^3, what a downsampling
+   step reads) read in 157 ms against 73 and wrote in 312 ms against 158 in a two-slot
+   job, but in a one-slot job read in 147 ms against 119 and wrote in 325 ms against 302
+   (and that job still used 1.3 cores, LSF not pinning it). Most of the gain is a core the
+   worker was not given, so a worker per slot, as a dask job runs, loses little at 1. 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).
 
@@ -769,8 +786,8 @@ Behaviours 10 to 16 are 0.2; 17 to 19 are 0.3.
     unchanged, and a handle pickles by its arguments and its context's spec, so each worker
     process opens its own cheaply (a local path goes resolved against the working directory
     at open, so a worker that changes directory before unpickling, as cellmap-analyze's dask
-    workers do, opens the same store; `reopen` uses it too), on the shared context when the handle had one, else on
-    a context rebuilt once per process from the same spec, so a capped context stays
+    workers do, opens the same store; `reopen` uses it too), on the shared context when the handle had one (the same
+    `threads`, and the same cache budget), else on a context rebuilt once per process from the same spec, so a capped context stays
     capped (mesh-n-bone's per-process cache and thread caps).
 
 12. **Group attributes belong to everyone.** The package rewrites only the `multiscales`
@@ -1000,7 +1017,9 @@ 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. Opt-in (`WORLDARRAY_STRESS=1`, about a minute,
+shared context's limits. 0.6: `shared_context(threads=)` at `None`, `N` and `0`, one
+context per pair, a negative count refused, and a pickled handle coming back on the same
+shared context. 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)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "worldarray"
-version = "0.5.0"
+version = "0.6.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

@@ -28,7 +28,7 @@
     parse_chunk_key,
 )
 
-__version__ = "0.5.0"
+__version__ = "0.6.0"
 
 __all__ = [
     "ArrayMeta",

src/worldarray/context.py

@@ -1,12 +1,14 @@
 """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``).
+``shared_context(cache_bytes, threads=)`` is one ``tensorstore.Context`` per budget and
+thread limit 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. By default 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; ``threads=0`` leaves them at tensorstore's sizes (chunkmirage's server). 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
@@ -26,7 +28,7 @@
 
 import tensorstore as ts
 
-_shared: dict[int, ts.Context] = {}
+_shared: dict[tuple[int, int | None], ts.Context] = {}
 _rebuilt: dict[str, ts.Context] = {}
 _lock = threading.Lock()
 
@@ -136,28 +138,33 @@ def thread_budget() -> int:
         return os.cpu_count() or 1
 
 
-def shared_context(cache_bytes: int = 0) -> ts.Context:
+def shared_context(cache_bytes: int = 0, *, threads: int | None = None) -> ts.Context:
+    """The process's one context for ``cache_bytes`` and ``threads``, made on first use and
+    watched. ``threads`` limits its decoding and file I/O pools: ``None`` to
+    ``thread_budget()``, a positive number to that many, ``0`` not at all (tensorstore's own
+    sizes, every CPU on the node)."""
+    key = (int(cache_bytes), None if threads is None else int(threads))
+    if key[1] is not None and key[1] < 0:
+        raise ValueError(f"threads must be None, 0 or positive, not {threads!r}")
     with _lock:
-        ctx = _shared.get(int(cache_bytes))
+        ctx = _shared.get(key)
         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
+            spec: dict = {"cache_pool": {"total_bytes_limit": key[0]}}
+            limit = thread_budget() if threads is None else key[1]
+            if limit:
+                spec["data_copy_concurrency"] = spec["file_io_concurrency"] = {"limit": limit}
+            ctx = watch(ts.Context(spec))
+            _shared[key] = 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``)."""
+    the context is watched (see ``KICK_EVERY``). It is not free: a read or write spanning
+    several chunks then decodes and stores them one at a time, half the speed of 4 on
+    8-chunk blocks in a two-slot job, 7 to 19% slower in a one-slot one (spec, Behaviour 7)."""
     limit = {"limit": int(concurrency)}
     spec = {
         "data_copy_concurrency": limit,
@@ -179,12 +186,14 @@ def capped_context(cache_bytes: int = 0, concurrency: int = 1, network: bool = T
 
 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)."""
+    context at the default limit (the other process takes its own ``thread_budget()``),
+    ``("shared", (cache_bytes, threads))`` for one with ``threads`` given, 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():
+        for (cache_bytes, threads), ctx in _shared.items():
             if ctx is context:
-                return ("shared", cache_bytes)
+                return ("shared", cache_bytes if threads is None else (cache_bytes, threads))
     spec = context.spec
     return ("spec", json.dumps(spec.to_json() if spec is not None else {}, sort_keys=True))
 
@@ -196,7 +205,8 @@ def from_recipe(recipe: tuple[str, Any] | None) -> ts.Context | None:
         return None
     kind, value = recipe
     if kind == "shared":
-        return shared_context(int(value))
+        cache_bytes, threads = value if isinstance(value, (tuple, list)) else (value, None)
+        return shared_context(int(cache_bytes), threads=threads)
     with _lock:
         ctx = _rebuilt.get(value)
         if ctx is None:

tests/test_access.py

@@ -301,3 +301,23 @@ def test_gce_guard_is_left_alone_on_a_google_machine(monkeypatch):
     monkeypatch.setattr(store, "_on_google_cloud", lambda: False)
     store.kvstore_specs("gs://bucket/vol")
     assert os.environ["GCE_METADATA_ROOT"] == "metadata.google.internal.invalid"
+
+
+def test_a_shared_context_takes_the_callers_thread_count(monkeypatch):
+    monkeypatch.setenv("OMP_NUM_THREADS", "1")
+    monkeypatch.setattr(wa.context, "_shared", {})
+    default = wa.shared_context(7)
+    assert default.spec.to_json()["data_copy_concurrency"] == {"limit": 1}
+    four = wa.shared_context(7, threads=4).spec.to_json()
+    assert four["data_copy_concurrency"] == four["file_io_concurrency"] == {"limit": 4}
+    machine = wa.shared_context(7, threads=0)
+    assert machine is wa.shared_context(7, threads=0) and machine is not default
+    spec = machine.spec.to_json()
+    assert "data_copy_concurrency" not in spec and "file_io_concurrency" not in spec
+    assert spec["cache_pool"] == {"total_bytes_limit": 7}
+    assert machine in wa.context._watched
+    assert wa.context.recipe_of(machine) == ("shared", (7, 0))
+    assert wa.context.from_recipe(("shared", (7, 0))) is machine
+    assert wa.context.from_recipe(("shared", 7)) is default
+    with pytest.raises(ValueError):
+        wa.shared_context(threads=-1)

tests/test_write.py

@@ -826,6 +826,8 @@ def test_a_pickled_handle_keeps_its_context(tmp_path):
     )  # one rebuilt context per spec
     shared = pickle.loads(pickle.dumps(wa.open(path)))
     assert shared.context is wa.shared_context(0)
+    machine = wa.shared_context(0, threads=0)
+    assert pickle.loads(pickle.dumps(wa.open(path, context=machine))).context is machine
 
 
 def test_codecs_of_a_path_and_the_fill_write_that_removes_a_key(tmp_path):