Let a job script size each worker's threads (0.9.0)

CHANGELOG.md

@@ -1,5 +1,21 @@
 # Changelog
 
+## 0.9.0 (2026-10-23)
+
+A job script can size each worker's threads. No placement change, no unit table change, and
+no default changes: with `WORLDARRAY_THREADS` unset and no `file_io_floor`, every context
+is what 0.8 made.
+
+- `thread_budget()` reads `WORLDARRAY_THREADS` ahead of `OMP_NUM_THREADS` and the job's
+  slots, for a job script to give each dask worker its share of the job's cores where
+  `OMP_NUM_THREADS` stays 1 for numpy and `LSB_DJOB_NUMPROC` counts the whole job (for
+  example `max(1, ncpus // processes)`). A handle opened on the driver on the shared
+  context is sized in each worker by that worker's value. Anything but a positive integer
+  raises `ValueError`; blank is unset.
+- `shared_context(file_io_floor=)` sets the least the file I/O pool takes under the thread
+  budget (default `FILE_IO_FLOOR`, 8); `1` keeps file I/O at the budget. A pickle carries
+  it, and with `threads=N` or `threads=0` it is ignored.
+
 ## 0.8.0 (2026-10-23)
 
 Faster writes from the contexts the package makes, and the create open timed out. No

README.md

@@ -87,11 +87,12 @@ s1 = wa.open_pyramid("out.zarr", mode="r+").add_level("s1", factors=2)   # place
   interpolation anchored where you say. Adjoining block reads agree with one read across
   both.
 - **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.
+  from `WORLDARRAY_THREADS` (a job script's share for each worker), `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
@@ -150,10 +151,11 @@ wa.metadata(path, level=None, units=None, open=True, retries=None) -> ArrayMeta
 wa.detect(path, missing_ok=False, retries=None) -> Node      # format, array or group, raw attrs
 wa.open_tensorstore(path, level=None, context=None, retries=None)
                     # retries= times out and retries document reads and the open, as it does data reads
-wa.shared_context(cache_bytes=0, threads=None, fsync=False, locking="none") ; wa.Retry(...)
+wa.shared_context(cache_bytes=0, threads=None, fsync=False, locking="none", file_io_floor=8) ; wa.Retry(...)
 wa.capped_context(cache_bytes=0, concurrency=1, network=True, file_io=None, fsync=False, locking="none")
-wa.thread_budget()                              # the shared context's decoding threads: OMP_NUM_THREADS, the job's slots, or CPUs;
-                                                #   file I/O gets at least 8; threads=N limits both to N, threads=0 leaves tensorstore's sizes
+wa.thread_budget()                              # the shared context's decoding threads: WORLDARRAY_THREADS, OMP_NUM_THREADS,
+                                                #   the job's slots, or CPUs; file I/O gets at least file_io_floor (8);
+                                                #   threads=N limits both to N, threads=0 leaves tensorstore's sizes
                                                 # contexts made here write without fsync or OS file locks; fsync=True, locking="os" restore them
 wa.list_chunk_keys(path, separator=".")        # chunk indices from keys alone, no metadata needed
 wa.Box ; wa.units.TABLE

docs/spec.md

@@ -185,11 +185,12 @@ wa.open_tensorstore(path, level=None, context=None, retries=None) -> 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, threads=None, fsync=False, locking="none")
+wa.shared_context(cache_bytes=0, threads=None, fsync=False, locking="none", file_io_floor=8)
 wa.capped_context(cache_bytes=0, concurrency=1, network=True, file_io=None, fsync=False, locking="none") ; wa.capped_context_spec(...)
-                    # threads: None = thread_budget() for decoding and at least 8 for file I/O (0.8), N = that many,
-                    #   0 = tensorstore's own pool sizes (0.6); fsync and locking: tensorstore's file_io_sync and
-                    #   file_io_locking mode, off and "none" in every context the package makes (0.8, Behaviour 7)
+                    # threads: None = thread_budget() for decoding and at least file_io_floor for file I/O (0.8; the
+                    #   argument 0.9), N = that many, 0 = tensorstore's own pool sizes (0.6); fsync and locking:
+                    #   tensorstore's file_io_sync and file_io_locking mode, off and "none" in every context the
+                    #   package makes (0.8, Behaviour 7)
                     # 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
@@ -236,8 +237,8 @@ where it also chooses the integer scale factor that makes the corner whole.
 `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 pool is sized by `thread_budget()` (0.5) and
-its file I/O pool by that or 8, whichever is larger (0.8), or both by `threads` when given,
-`0` leaving tensorstore's own sizes (0.6, Behaviour 7).
+its file I/O pool by that or `file_io_floor` (8), whichever is larger (0.8; the argument
+0.9), or both 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
@@ -366,8 +367,8 @@ wa.create(..., attrs={...}, exists="adopt")   # an adopt also compares the keys
 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)
+wa.thread_budget() -> int # threads for tensorstore's decoding and file I/O in this process: WORLDARRAY_THREADS (0.9),
+                          #   else 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)
 ```
@@ -703,8 +704,20 @@ 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. Its `file_io_concurrency` limit is that or
-   `context.FILE_IO_FLOOR` (8), whichever is larger (0.8): file I/O threads wait on storage
+   `LSB_DJOB_NUMPROC` elsewhere. Ahead of them all comes `WORLDARRAY_THREADS` (0.9), the
+   package's own variable, for a job script to give each worker its share of the job's
+   cores where `OMP_NUM_THREADS` stays 1 for numpy and `LSB_DJOB_NUMPROC` counts the whole
+   job: cellmap-analyze's prologue exports its CPUs over its worker processes, at least 1,
+   so a worker with spare cores uses them (in a two-slot job, pools of four against one
+   halved an 8-chunk block's time, below) and a worker per slot stays at one. A handle opened on
+   the driver before the cluster starts is then sized in each worker by that worker's
+   share, as a pickle carries `threads=None`, not a number (Behaviour 11). A value that is
+   not a positive integer raises `ValueError` rather than being skipped: nothing else sets
+   this variable, so a value worldarray cannot use is a mistake to report. Its
+   `file_io_concurrency` limit is that or `file_io_floor`, by default
+   `context.FILE_IO_FLOOR` (8), whichever is larger (0.8; the argument 0.9, `1` keeping
+   file I/O at the budget, cellmap-analyze's choice until eight is measured on /nrs at
+   scale): file I/O threads wait on storage
    rather than compute, so they take no core from the job, and one of them serialises every
    chunk file of a block on NFS's round trips. From one decoding thread, eight 64^3 label
    chunks written to NFS took 166 to 176 ms with one file I/O thread and 83 to 89 ms with
@@ -720,10 +733,11 @@ this document and that one differ, fix one of them.
    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.
+   instead of building its own). Each set of arguments 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. `file_io_floor` sizes only the budget's pool, so with
+   a thread count of its own it is ignored and no second context is made.
 
    Writes (0.8). Every context the package makes (`shared_context`, `capped_context`)
    sets tensorstore's `file_io_sync` to false and its `file_io_locking` mode to `"none"`,
@@ -1072,7 +1086,11 @@ shared context. 0.7: `metadata`, `open` and `detect` against a server
 that accepts and never answers raising `TimeoutError` after the attempts. 0.8: the default
 and capped contexts' file I/O limit, sync and locking, `fsync` and `locking` carried
 through a pickle, and a create whose first attempt lands after its timeout adopting its own
-array under `exists="error"`. Opt-in (`WORLDARRAY_STRESS=1`, about a minute,
+array under `exists="error"`. 0.9: `WORLDARRAY_THREADS` ahead of every other variable,
+blank as unset and a bad value refused; `file_io_floor` sizing file I/O under the budget,
+ignored with a thread count, and carried through a pickle; a handle pickled on the driver
+sized in a spawned worker by that worker's `WORLDARRAY_THREADS`.
+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.8.0"
+version = "0.9.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.8.0"
+__version__ = "0.9.0"
 
 __all__ = [
     "ArrayMeta",

src/worldarray/context.py

@@ -5,9 +5,10 @@
 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 pool is sized by ``thread_budget()``, from the variables a
-cluster job sets, where tensorstore alone would size it to every CPU on the node, and its
-file I/O pool to at least ``FILE_IO_FLOOR`` threads, which wait on storage rather than
-compute; ``threads=0`` leaves both at tensorstore's sizes (chunkmirage's server). Every
+cluster job sets (``WORLDARRAY_THREADS`` first, a job script's share for each worker), where
+tensorstore alone would size it to every CPU on the node, and its file I/O pool to at least
+``file_io_floor`` threads (``FILE_IO_FLOOR``), which wait on storage rather than compute;
+``threads=0`` leaves both at tensorstore's sizes (chunkmirage's server). Every
 context made here writes without fsync and without OS file locks (``write_settings``).
 Every context here, and every one a store is opened on, is watched for tensorstore's lost
 wakeup (``watch``).
@@ -30,7 +31,7 @@
 
 import tensorstore as ts
 
-_shared: dict[tuple[int, int | None, bool, str], ts.Context] = {}
+_shared: dict[tuple[int, int | None, bool, str, int], ts.Context] = {}
 _rebuilt: dict[str, ts.Context] = {}
 _lock = threading.Lock()
 
@@ -119,14 +120,26 @@ def _forget_watchdog() -> None:
     os.register_at_fork(after_in_child=_forget_watchdog)
 
 
+THREADS_VAR = "WORLDARRAY_THREADS"
 _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``
+    """Threads for this process's tensorstore decoding and file I/O: ``WORLDARRAY_THREADS``
+    when set, a job script's share for each worker where ``OMP_NUM_THREADS`` must stay 1 for
+    numpy (a positive integer, anything else a ``ValueError``), else ``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."""
+    own = os.environ.get(THREADS_VAR, "").strip()
+    if own:
+        try:
+            n = int(own)
+        except ValueError:
+            n = 0
+        if n < 1:
+            raise ValueError(f"{THREADS_VAR} must be a positive integer, not {own!r}")
+        return n
     for name in _SLOTS:
         try:
             n = int(os.environ.get(name, ""))
@@ -162,15 +175,27 @@ def shared_context(
     threads: int | None = None,
     fsync: bool = False,
     locking: str = "none",
+    file_io_floor: int = FILE_IO_FLOOR,
 ) -> ts.Context:
     """The process's one context for these arguments, made on first use and watched.
     ``threads`` limits its decoding and file I/O pools: ``None`` decoding to
-    ``thread_budget()`` and file I/O to that or ``FILE_IO_FLOOR``, whichever is larger, a
-    positive number both to that many, ``0`` neither (tensorstore's own sizes, every CPU on
-    the node). ``fsync`` and ``locking`` go to ``write_settings``."""
-    key = (int(cache_bytes), None if threads is None else int(threads), bool(fsync), str(locking))
-    if key[1] is not None and key[1] < 0:
+    ``thread_budget()``, read when the context is made, and file I/O to that or
+    ``file_io_floor``, whichever is larger (``1`` keeps file I/O at the budget), a positive
+    number both to that many, ``0`` neither (tensorstore's own sizes, every CPU on the
+    node); ``file_io_floor`` applies to ``threads=None`` alone. ``fsync`` and ``locking`` go
+    to ``write_settings``."""
+    if threads is not None and int(threads) < 0:
         raise ValueError(f"threads must be None, 0 or positive, not {threads!r}")
+    if int(file_io_floor) < 0:
+        raise ValueError(f"file_io_floor must be 0 or positive, not {file_io_floor!r}")
+    floor = max(int(file_io_floor), 1) if threads is None else FILE_IO_FLOOR
+    key = (
+        int(cache_bytes),
+        None if threads is None else int(threads),
+        bool(fsync),
+        str(locking),
+        floor,
+    )
     with _lock:
         ctx = _shared.get(key)
         if ctx is None:
@@ -181,7 +206,7 @@ def shared_context(
             if threads is None:
                 budget = thread_budget()
                 spec["data_copy_concurrency"] = {"limit": budget}
-                spec["file_io_concurrency"] = {"limit": max(budget, FILE_IO_FLOOR)}
+                spec["file_io_concurrency"] = {"limit": max(budget, floor)}
             elif key[1]:
                 spec["data_copy_concurrency"] = spec["file_io_concurrency"] = {"limit": key[1]}
             ctx = watch(ts.Context(spec))
@@ -238,16 +263,20 @@ def capped_context(
     return watch(ts.Context(spec))
 
 
+# shared_context's key with every argument at its default
+_SHARED_DEFAULTS = (0, None, False, "none", FILE_IO_FLOOR)
+
+
 def recipe_of(context: ts.Context) -> tuple[str, Any]:
     """How another process rebuilds ``context``: ``("shared", cache_bytes)`` for a shared
     context with every other argument at its default (the other process takes its own
-    ``thread_budget()``), ``("shared", (cache_bytes, threads, fsync, locking))`` otherwise,
-    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)."""
+    ``thread_budget()``), ``("shared", (cache_bytes, threads, fsync, locking, file_io_floor))``
+    otherwise, 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 key, ctx in _shared.items():
             if ctx is context:
-                return ("shared", key[0] if key[1:] == (None, False, "none") else key)
+                return ("shared", key[0] if key[1:] == _SHARED_DEFAULTS[1:] else key)
     spec = context.spec
     return ("spec", json.dumps(spec.to_json() if spec is not None else {}, sort_keys=True))
 
@@ -259,12 +288,14 @@ def from_recipe(recipe: tuple[str, Any] | None) -> ts.Context | None:
         return None
     kind, value = recipe
     if kind == "shared":
-        cache_bytes, threads, fsync, locking = (
-            (*value, False, "none")[:4]
-            if isinstance(value, (tuple, list))
-            else (value, None, False, "none")
+        given = tuple(value) if isinstance(value, (tuple, list)) else (value,)
+        cache_bytes, threads, fsync, locking, floor = (
+            *given,
+            *_SHARED_DEFAULTS[len(given) :],
+        )[:5]  # a 0.6 or 0.7 pickle carries two, a 0.8 one four
+        return shared_context(
+            int(cache_bytes), threads=threads, fsync=fsync, locking=locking, file_io_floor=floor
         )
-        return shared_context(int(cache_bytes), threads=threads, fsync=fsync, locking=locking)
     with _lock:
         ctx = _rebuilt.get(value)
         if ctx is None: