diff --git a/CHANGELOG.md b/CHANGELOG.md index 46380e1..e894870 100644 --- a/CHANGELOG.md +++ b/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' diff --git a/README.md b/README.md index bae857c..a012444 100644 --- a/README.md +++ b/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 diff --git a/docs/spec.md b/docs/spec.md index 927eb34..072ecc8 100644 --- a/docs/spec.md +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index 1c5be8d..f624a9c 100644 --- a/pyproject.toml +++ b/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" } diff --git a/src/worldarray/__init__.py b/src/worldarray/__init__.py index ea54e79..2fc7822 100644 --- a/src/worldarray/__init__.py +++ b/src/worldarray/__init__.py @@ -28,7 +28,7 @@ parse_chunk_key, ) -__version__ = "0.5.0" +__version__ = "0.6.0" __all__ = [ "ArrayMeta", diff --git a/src/worldarray/context.py b/src/worldarray/context.py index 0508557..1226375 100644 --- a/src/worldarray/context.py +++ b/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,20 +138,23 @@ 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 @@ -157,7 +162,9 @@ def capped_context_spec(cache_bytes: int = 0, concurrency: int = 1, network: boo """The spec dict behind ``capped_context``, for callers building their own opens. ``network=False`` caps only decoding and local file I/O, leaving tensorstore's request pools alone (cellmap-analyze's choice for remote reads). A ``concurrency`` of 1 is safe: - 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: diff --git a/tests/test_access.py b/tests/test_access.py index 2eea1dd..b9f2102 100644 --- a/tests/test_access.py +++ b/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) diff --git a/tests/test_write.py b/tests/test_write.py index a4290f4..339ced6 100644 --- a/tests/test_write.py +++ b/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): diff --git a/uv.lock b/uv.lock index 8c27fbe..2c5e0df 100644 --- a/uv.lock +++ b/uv.lock @@ -1247,7 +1247,7 @@ wheels = [ [[package]] name = "worldarray" -version = "0.5.0" +version = "0.6.0" source = { editable = "." } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },