Skip to content

Commit d1cfcd0

Browse files
michaltakacDimensionLab OSS Sync
authored andcommitted
perf(arrayroute): update arrayroute
GitOrigin-RevId: 51b82cef8c9aafa59e84e8a6ddf3e6710126e030
1 parent 5d5962e commit d1cfcd0

9 files changed

Lines changed: 517 additions & 16 deletions

File tree

‎packages/arrayroute/README.md‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1113,3 +1113,53 @@ buffer authority and reports native cost; pass `buffer_bytes=` to select a dense
11131113
plan. `plan_concurrency` reports each candidate's strategy (or `per-tile` when
11141114
stream subdivision is needed). Snapshot and remote session constructors do not
11151115
currently expose this opt-in.
1116+
1117+
## Parallel gathering with bounded intermediates
1118+
1119+
For independently decodable Zarr shards, a local Zarrista session can opt into
1120+
parallel region reads without allocating a dense bounding array:
1121+
1122+
```python
1123+
session = LocalReadSession(
1124+
"zarrista", read_strategy="parallel", gather_workers=8, buffer_bytes=8 << 20,
1125+
)
1126+
reader = session.open("velocity.zarr", "")
1127+
request = reader.plan(selection)
1128+
print(request.strategy, request.gather_workers, request.intermediate_bytes_upper_bound)
1129+
preview = reader.read(selection)
1130+
```
1131+
1132+
The planner considers up to `gather_workers` workers (1–256), bounded further by the
1133+
number of selected shards and the available buffer. It charges the final output,
1134+
worker-owned selected regions, and one decoded chunk per active worker. It lowers
1135+
the worker count before choosing ordinary native selection when fewer than two
1136+
workers fit. Contiguous requests and requests within one region use native selection.
1137+
Streaming, materialization and prefetch re-plan with their actual remaining buffers.
1138+
This is an explicit policy; small requests can get slower from scheduling overhead.
1139+
1140+
Selected shards are generated lazily. Each worker gathers only touched subchunks
1141+
inside its shard; results are copied into disjoint locations in the final compact
1142+
array. The bounded executor retains at most one region per scheduled slot, including
1143+
the region being consumed. Every call joins its pool, including on failure, so no
1144+
background worker outlives its read reservation. It cannot interrupt an already
1145+
running native call. Native codec threads, shard indexes, caches and scratch remain
1146+
outside the explicit array-buffer allowance. Outer read workers multiply inner
1147+
worker pools; benchmark the whole concurrency configuration rather than assuming
1148+
more threads improve throughput.
1149+
1150+
`ParallelGatherBackend(base, factory, GatherWindow(region_shape, workers))` is the
1151+
generic extension point. The factory must supply independent mutable reader state
1152+
with the same metadata as `base`. A read-only native array handle may be shared.
1153+
Worker results must not retain arrays larger than their selected region or one
1154+
decoded chunk; workers must not reserve from the parent reader's shared `ReadBudget`.
1155+
The base must guarantee a one-chunk decode window for strided reads; regions must
1156+
align to whole chunks on every axis. This capability is not inferred for arbitrary
1157+
engines. `LocalReadSession` currently exposes it only for Zarrista partial-shard
1158+
reads; other engines and outer-codec shards are rejected for this strategy.
1159+
1160+
`ReadStats.read_strategies` counts completed tile executions by strategy, and
1161+
`gather_workers_peak` reports the largest admitted inner worker limit. This is
1162+
not a measurement of active threads; a cache hit may launch none.
1163+
`probe_reads` includes these in every trial, so the requested maximum can be
1164+
compared with actual budget-reduced counts. Native reads report one gather worker;
1165+
that does not count the engine's own internal threads.

‎packages/arrayroute/pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "arrayroute"
3-
version = "0.1.11"
3+
version = "0.1.12"
44
description = "Measured, bounded selections over scientific chunked arrays"
55
requires-python = ">=3.11"
66
license = "Apache-2.0"
Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
# SPDX-License-Identifier: Apache-2.0
22
"""Selection planning independent of a storage format, scheduler or cloud provider."""
33

4-
from .plan import ArraySpec, DecodeWindow, ReadPlan, Tile, plan, tiles
4+
from .plan import ArraySpec, DecodeWindow, GatherWindow, ReadPlan, Tile, plan, tiles
55
from .reader import ArrayReader, ReadStats
66

77
__all__ = [
8-
"ArrayReader", "ArraySpec", "DecodeWindow", "ReadPlan", "ReadStats", "Tile", "plan", "tiles",
8+
"ArrayReader", "ArraySpec", "DecodeWindow", "GatherWindow",
9+
"ReadPlan", "ReadStats", "Tile", "plan", "tiles",
910
]
Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
"""Bounded region gathering over independent, single-chunk decoding readers."""
3+
from __future__ import annotations
4+
5+
from collections.abc import Callable
6+
from dataclasses import replace
7+
8+
import numpy as np
9+
10+
from .backends import _grid_regions
11+
from .execution import ordered_map
12+
from .plan import GatherWindow, ReadPlan, Selection
13+
from .reader import Backend
14+
15+
16+
class ParallelGatherBackend:
17+
"""Read disjoint chunk-aligned regions with bounded, worker-owned intermediates.
18+
19+
Each factory result must have the same metadata and single-chunk decode guarantee
20+
as base. It must own its mutable state; read-only native array handles may be shared.
21+
Results must not retain arrays larger than their selected region or one decoded
22+
chunk. Workers must not reserve from the parent reader's shared ReadBudget.
23+
Only planned strided calls parallelize. No thread pool survives a call, including
24+
failures. Native reads are joined before the output reservation is released.
25+
"""
26+
27+
def __init__(self, base: Backend, factory: Callable[[], Backend], window: GatherWindow):
28+
if base.spec.gather_window is not None:
29+
raise ValueError("Nested gather backends are not supported")
30+
self.base, self.factory, self.window = base, factory, window
31+
self.dtype = base.dtype
32+
self.spec = replace(base.spec, gather_window=window)
33+
34+
def read(self, selection: Selection) -> np.ndarray:
35+
return self.base.read(selection)
36+
37+
def read_plan(self, request: ReadPlan) -> np.ndarray:
38+
if request.strategy == 'native':
39+
return self.base.read(request.selection)
40+
if (request.strategy != 'parallel-gather'
41+
or not 2 <= request.gather_workers <= self.window.workers):
42+
raise ValueError('Parallel gather requires its planned worker count')
43+
output = np.empty(request.output_shape, dtype=self.dtype)
44+
45+
def make_worker():
46+
backend = self.factory()
47+
if backend.spec != self.base.spec or backend.dtype != self.dtype:
48+
raise ValueError('Gather factory changed array metadata')
49+
50+
def read(region):
51+
return region, backend.read(region)
52+
53+
return read
54+
55+
regions = _grid_regions(request.selection, self.window.region_shape)
56+
with ordered_map(make_worker, regions, workers=request.gather_workers,
57+
max_pending=request.gather_workers) as rows:
58+
for region, value in rows:
59+
targets = []
60+
shape = []
61+
for part, original in zip(region, request.selection, strict=True):
62+
if isinstance(original, slice):
63+
assert isinstance(part, slice)
64+
count = len(range(part.start, part.stop, part.step))
65+
begin = (part.start - original.start) // original.step
66+
targets.append(slice(begin, begin + count))
67+
shape.append(count)
68+
if value.shape != tuple(shape) or value.dtype != self.dtype:
69+
raise ValueError('Gather worker returned an unexpected shape or dtype')
70+
output[tuple(targets)] = value
71+
del value
72+
return output

‎packages/arrayroute/src/arrayroute/plan.py‎

Lines changed: 52 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# SPDX-License-Identifier: Apache-2.0
2-
"""O(rank) selection costs; lazy subdivision without a global chunk inventory."""
2+
"""Bounded selection costs; lazy subdivision without a global chunk inventory."""
33

44
from __future__ import annotations
55

@@ -34,13 +34,30 @@ def __post_init__(self) -> None:
3434
raise TypeError("strided_only must be a boolean")
3535

3636

37+
@dataclass(frozen=True)
38+
class GatherWindow:
39+
"""Independent grid regions with one live decoded chunk per worker."""
40+
41+
region_shape: tuple[int, ...]
42+
workers: int
43+
44+
def __post_init__(self) -> None:
45+
if (isinstance(self.workers, bool) or not isinstance(self.workers, int)
46+
or not 1 <= self.workers <= 256):
47+
raise ValueError("Gather workers must be an integer from 1 to 256")
48+
if (not self.region_shape or any(isinstance(v, bool) or not isinstance(v, int)
49+
or v < 1 for v in self.region_shape)):
50+
raise ValueError("Gather regions require positive integer dimensions")
51+
52+
3753
@dataclass(frozen=True)
3854
class ArraySpec:
3955
shape: tuple[int, ...]
4056
chunks: tuple[int, ...]
4157
itemsize: int
4258
decode_window: DecodeWindow | None = None
4359
dense_gather: bool = False
60+
gather_window: GatherWindow | None = None
4461

4562
def __post_init__(self) -> None:
4663
if not isinstance(self.dense_gather, bool):
@@ -53,6 +70,17 @@ def __post_init__(self) -> None:
5370
if any(isinstance(v, bool) or not isinstance(v, int) or v < minimum for v in values):
5471
raise ValueError("shape, chunks and itemsize require valid integer dimensions")
5572

73+
gather = self.gather_window
74+
if gather is not None:
75+
if not isinstance(gather, GatherWindow):
76+
raise TypeError("gather_window must be a GatherWindow or None")
77+
if (self.dense_gather or self.decode_window is None
78+
or self.decode_window.max_chunks != 1):
79+
raise ValueError("Parallel gathering requires a single-chunk native decode window")
80+
if (len(gather.region_shape) != len(self.chunks)
81+
or any(r % c for r, c in zip(gather.region_shape, self.chunks, strict=True))):
82+
raise ValueError("Gather regions must be a whole number of chunks on each axis")
83+
5684
@property
5785
def logical_bytes(self) -> int:
5886
return math.prod(self.shape) * self.itemsize
@@ -67,8 +95,9 @@ class ReadPlan:
6795
decoded_bytes_upper_bound: int
6896
peak_decoded_bytes_upper_bound: int
6997
intermediate_bytes_upper_bound: int = 0
70-
strategy: Literal["native", "dense-gather"] = "native"
98+
strategy: Literal["native", "dense-gather", "parallel-gather"] = "native"
7199
read_selection: Selection | None = None
100+
gather_workers: int = 1
72101

73102
@property
74103
def decode_amplification(self) -> float:
@@ -129,9 +158,29 @@ def plan(spec: ArraySpec, selection: Selection, *, buffer_bytes: int | None = No
129158
touched * math.prod(spec.chunks) * spec.itemsize,
130159
live_chunks * math.prod(spec.chunks) * spec.itemsize,
131160
)
132-
if (not spec.dense_gather or buffer_bytes is None or not request.output_bytes
161+
if (buffer_bytes is None or not request.output_bytes
133162
or not any(isinstance(s, slice) and s.step > 1 for s in normalized)):
134163
return request
164+
gather = spec.gather_window
165+
if gather is not None:
166+
regions = plan(ArraySpec(spec.shape, gather.region_shape, spec.itemsize),
167+
request.selection).chunks_touched
168+
region_output = spec.itemsize * math.prod(
169+
min(_length(s), (r + s.step - 1) // s.step)
170+
for s, r in zip(request.selection, gather.region_shape, strict=True)
171+
if isinstance(s, slice)
172+
)
173+
for workers in range(min(gather.workers, regions), 1, -1):
174+
candidate = replace(
175+
request, strategy="parallel-gather", gather_workers=workers,
176+
peak_decoded_bytes_upper_bound=workers * request.peak_decoded_bytes_upper_bound,
177+
intermediate_bytes_upper_bound=min(request.output_bytes, workers * region_output),
178+
)
179+
if candidate.buffer_bytes_upper_bound <= buffer_bytes:
180+
return candidate
181+
return request
182+
if not spec.dense_gather:
183+
return request
135184
bounds = tuple(
136185
s if isinstance(s, int) else slice(s.start, s.start + (_length(s) - 1) * s.step + 1, 1)
137186
for s in normalized

‎packages/arrayroute/src/arrayroute/probe.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import random
99
import statistics
1010
import time
11+
from collections import Counter
1112
from collections.abc import Callable, Sequence
1213
from typing import cast
1314

@@ -131,7 +132,9 @@ def read(index: int) -> dict:
131132
buffersize=max(1, 32768 // dtype.itemsize)):
132133
digest.update(cast(np.ndarray, block).tobytes())
133134
return {'sha256': digest.hexdigest(), 'bytes': value.nbytes,
134-
'requests': stats.requests, 'read_seconds': stats.seconds}
135+
'requests': stats.requests, 'read_seconds': stats.seconds,
136+
'read_strategies': stats.read_strategies,
137+
'gather_workers_peak': stats.gather_workers_peak}
135138

136139
return read
137140

@@ -152,11 +155,16 @@ def read(index: int) -> dict:
152155
if any((r['sha256'], r['bytes']) != (e['sha256'], e['bytes'])
153156
for r, e in zip(rows, reference, strict=True)):
154157
raise ValueError("Pilot values changed from the sequential reference")
158+
strategies = Counter()
159+
for row in rows:
160+
strategies.update(row['read_strategies'])
155161
report['trials'].append({
156162
'repetition': repetition, 'workers': count, 'elapsed_seconds': elapsed,
157163
'output_bytes': sum(r['bytes'] for r in rows),
158164
'read_seconds_sum': sum(r['read_seconds'] for r in rows),
159165
'requests': sum(r['requests'] for r in rows), 'verified': True,
166+
'read_strategies': dict(strategies),
167+
'gather_workers_peak': max(r['gather_workers_peak'] for r in rows),
160168
})
161169
medians = {count: statistics.median(t['elapsed_seconds'] for t in report['trials']
162170
if t['workers'] == count) for count in allowed}

‎packages/arrayroute/src/arrayroute/reader.py‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from collections.abc import Callable, Generator, Iterator
88
from concurrent.futures import ThreadPoolExecutor
99
from contextlib import contextmanager, nullcontext
10-
from dataclasses import dataclass
10+
from dataclasses import dataclass, field
1111
from typing import Protocol
1212

1313
import numpy as np
@@ -25,10 +25,17 @@ def read(self, selection: Selection) -> np.ndarray: ...
2525

2626
@dataclass
2727
class ReadStats:
28+
"""Completed tile plans; worker limits are admitted maxima, not observed active threads.
29+
30+
A decoded cache hit can satisfy a parallel plan without starting any workers.
31+
"""
32+
2833
requests: int = 0
2934
output_bytes: int = 0
3035
decoded_bytes_upper_bound: int = 0
3136
seconds: float = 0.0
37+
read_strategies: dict[str, int] = field(default_factory=dict)
38+
gather_workers_peak: int = 0
3239

3340

3441
class ArrayReader:
@@ -42,8 +49,9 @@ def __init__(
4249
self.buffer_bytes = (min(buffer_bytes, shared_budget.max_bytes)
4350
if shared_budget is not None else buffer_bytes)
4451
self.shared_budget = shared_budget
45-
if backend.spec.dense_gather and not callable(getattr(backend, "read_plan", None)):
46-
raise TypeError("Dense gather metadata requires a backend with read_plan")
52+
if ((backend.spec.dense_gather or backend.spec.gather_window is not None)
53+
and not callable(getattr(backend, "read_plan", None))):
54+
raise TypeError("Gather metadata requires a backend with read_plan")
4755
self._execute: Callable[[ReadPlan], np.ndarray] = getattr(
4856
backend, "read_plan", lambda request: backend.read(request.selection))
4957
if not callable(self._execute):
@@ -73,6 +81,10 @@ def stream(
7381
raise ValueError("Backend returned an unexpected shape or dtype")
7482
if stats is not None:
7583
stats.requests += 1
84+
strategy = tile.plan.strategy
85+
stats.read_strategies[strategy] = stats.read_strategies.get(strategy, 0) + 1
86+
stats.gather_workers_peak = max(
87+
stats.gather_workers_peak, tile.plan.gather_workers)
7688
stats.seconds += time.perf_counter() - start
7789
stats.output_bytes += value.nbytes
7890
stats.decoded_bytes_upper_bound += tile.plan.decoded_bytes_upper_bound

‎packages/arrayroute/src/arrayroute/session.py‎

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,8 @@
1111

1212
from .backends import DenseGatherBackend, TensorStoreBackend, ZarrBackend, ZarristaBackend
1313
from .budget import ReadBudget
14-
from .plan import Selection
14+
from .gather import ParallelGatherBackend
15+
from .plan import GatherWindow, Selection
1516
from .ranges import RangePolicy
1617
from .reader import ArrayReader, Backend, ReadStats
1718

@@ -49,6 +50,7 @@ def __init__(
4950
range_policy: RangePolicy | None = None,
5051
read_budget: ReadBudget | None = None,
5152
read_strategy: str = "native",
53+
gather_workers: int = 4,
5254
):
5355
if not isinstance(immutable, bool):
5456
raise TypeError("immutable must be a boolean")
@@ -58,8 +60,16 @@ def __init__(
5860
)
5961
if engine not in {"zarr-python", "tensorstore", "zarrista"}:
6062
raise ValueError(f"Unknown array engine {engine!r}")
61-
if read_strategy not in {"native", "dense"}:
62-
raise ValueError("read_strategy must be native or dense")
63+
if read_strategy not in {"native", "dense", "parallel"}:
64+
raise ValueError("read_strategy must be native, dense or parallel")
65+
if (isinstance(gather_workers, bool) or not isinstance(gather_workers, int)
66+
or not 1 <= gather_workers <= 256):
67+
raise ValueError("Gather workers must be an integer from 1 to 256")
68+
if read_strategy == "parallel" and engine != "zarrista":
69+
raise ValueError("Parallel gathering currently requires the Zarrista engine")
70+
if read_strategy != "parallel" and gather_workers != 4:
71+
raise ValueError("gather_workers requires the parallel read strategy")
72+
self.gather_workers = gather_workers
6373
self.read_strategy = read_strategy
6474
if (
6575
min(buffer_bytes, max_open_arrays, tensorstore_threads) < 1
@@ -141,9 +151,16 @@ def open(self, store: str | Path, array_path: str) -> ArrayReader:
141151
else:
142152
import zarrista
143153

144-
backend = ZarristaBackend(
145-
zarrista.Array.open(zarrista.store.FilesystemStore(key[0]), path="/" + key[1])
146-
)
154+
array = zarrista.Array.open(zarrista.store.FilesystemStore(key[0]), path="/" + key[1])
155+
native = ZarristaBackend(array)
156+
backend = native
157+
if self.read_strategy == "parallel":
158+
if not native._partial_shards:
159+
raise ValueError("Parallel gathering requires independently decodable shards")
160+
backend = ParallelGatherBackend(
161+
native, lambda: ZarristaBackend(array),
162+
GatherWindow(tuple(array.chunk_shape([0] * array.ndim)), self.gather_workers),
163+
)
147164
if self.read_strategy == "dense":
148165
backend = DenseGatherBackend(backend)
149166
reader = ArrayReader(

0 commit comments

Comments
 (0)