Skip to content

Commit 04c1be4

Browse files
michaltakacDimensionLab OSS Sync
authored andcommitted
feat(arrayroute): update arrayroute
GitOrigin-RevId: c2c5827b69413012e1d210cd144d4dbff3eb3e0b
1 parent d8f834b commit 04c1be4

5 files changed

Lines changed: 283 additions & 12 deletions

File tree

‎packages/arrayroute/README.md‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1038,6 +1038,37 @@ cannot be forcibly cancelled by this API; hard deadlines require process isolati
10381038
The existing `ordered_map` convenience function uses a fresh pool for one pass and
10391039
releases it on exit. Neither API depends on a machine-learning framework or cloud.
10401040

1041+
For uneven work sizes, also constrain the planned bytes held by submitted tasks,
1042+
completed results and the currently yielded result:
1043+
1044+
```python
1045+
with OrderedPool(make_processor, workers=4, max_pending=8,
1046+
max_pending_bytes=128 << 20,
1047+
estimate_bytes=lambda request: request.peak_buffer_bytes) as pool:
1048+
with pool.map(requests) as results:
1049+
for result in results:
1050+
consume(result)
1051+
del result
1052+
```
1053+
1054+
The estimator runs once per input before submission. It must return a nonnegative
1055+
integer and account for every buffer covered by the claimed allowance. A task larger
1056+
than the allowance fails before execution; reduce its selection or batch size, or
1057+
increase the allowance. Tasks wait in input order instead of bypassing a large task.
1058+
Admission remains charged while a result is yielded and releases when iteration resumes.
1059+
Estimation and input-iteration errors abort the pass immediately; worker failures still
1060+
surface at their input position.
1061+
1062+
One input awaiting admission may be retained as lookahead. The input iterable,
1063+
worker caches, unaccounted scratch, allocator overhead and consumer-retained values
1064+
are outside this planned-byte allowance. Use lightweight selection descriptors as
1065+
inputs when bounding actual data-buffer lifetime. This is not a process RSS limit.
1066+
1067+
`python -m arrayroute.benchmark_admission --output admission.json` compares count-only
1068+
and byte admission in fresh processes with uneven synthetic buffers, verified output
1069+
hashes, traced allocations and peak RSS. It measures queue memory, not scientific
1070+
I/O or training throughput.
1071+
10411072
## Measure reader concurrency on your own data
10421073

10431074
`arrayroute.probe.probe_reads` runs bounded pilots through any `ArrayReader` factory:

‎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.21"
3+
version = "0.1.22"
44
description = "Measured, bounded selections over scientific chunked arrays"
55
requires-python = ">=3.11"
66
license = "Apache-2.0"
Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
"""Measure ordered queue payloads with uneven buffers and a paced consumer."""
3+
from __future__ import annotations
4+
5+
import argparse
6+
import hashlib
7+
import json
8+
import platform
9+
import random
10+
import resource
11+
import subprocess
12+
import sys
13+
import time
14+
import tracemalloc
15+
from pathlib import Path
16+
17+
import numpy as np
18+
19+
from .execution import ordered_map
20+
21+
22+
def run_trial(policy: str, unit_bytes: int) -> dict:
23+
if policy not in {'count', 'bytes'} or unit_bytes < 1:
24+
raise ValueError('Choose count/bytes admission and a positive buffer unit')
25+
sizes = [unit_bytes * n for n in [1, 4, 2, 6] * 4]
26+
expected = []
27+
for i, size in enumerate(sizes):
28+
digest = hashlib.sha256()
29+
block = bytes([i]) * min(size, 65536)
30+
whole, remainder = divmod(size, len(block))
31+
for _ in range(whole):
32+
digest.update(block)
33+
digest.update(block[:remainder])
34+
expected.append(digest.hexdigest())
35+
36+
def factory():
37+
return lambda i: np.full(sizes[i], i, dtype=np.uint8)
38+
39+
digests = []
40+
tracemalloc.start()
41+
start = time.perf_counter()
42+
with ordered_map(factory, range(len(sizes)), workers=4, max_pending=8,
43+
max_pending_bytes=unit_bytes * 8 if policy == 'bytes' else None,
44+
estimate_bytes=(lambda i: sizes[i]) if policy == 'bytes' else None) as outputs:
45+
for value in outputs:
46+
digests.append(hashlib.sha256(memoryview(value)).hexdigest())
47+
time.sleep(.01)
48+
del value
49+
seconds = time.perf_counter() - start
50+
_, peak = tracemalloc.get_traced_memory()
51+
tracemalloc.stop()
52+
if digests != expected:
53+
raise ValueError('Ordered payload bytes differ from the independent reference')
54+
rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
55+
return {
56+
'policy': policy, 'unit_bytes': unit_bytes, 'payload_bytes': sum(sizes),
57+
'planned_queue_bytes': unit_bytes * 8 if policy == 'bytes' else None,
58+
'seconds': seconds, 'peak_traced_bytes': peak,
59+
'peak_rss_bytes': rss if sys.platform == 'darwin' else rss * 1024,
60+
'digest_sha256': hashlib.sha256(json.dumps(digests).encode()).hexdigest(),
61+
'verified': True, 'workers': 4, 'max_pending': 8, 'consumer_delay_seconds': .01,
62+
}
63+
64+
65+
def main() -> None:
66+
parser = argparse.ArgumentParser(description=__doc__)
67+
parser.add_argument('--output', type=Path)
68+
parser.add_argument('--unit-mib', type=int, default=16)
69+
parser.add_argument('--repeats', type=int, default=3)
70+
parser.add_argument('--worker', choices=['count', 'bytes'])
71+
args = parser.parse_args()
72+
if args.unit_mib < 1 or args.repeats < 1:
73+
parser.error('unit-mib and repeats must be positive')
74+
if args.worker:
75+
print(json.dumps(run_trial(args.worker, args.unit_mib * (1 << 20))))
76+
return
77+
if args.output is None or args.output.exists():
78+
parser.error('Choose a new --output path')
79+
rows = []
80+
for repeat in range(args.repeats):
81+
policies = ['count', 'bytes']
82+
random.Random(20261023 + repeat).shuffle(policies)
83+
for policy in policies:
84+
completed = subprocess.run(
85+
[sys.executable, '-m', 'arrayroute.benchmark_admission',
86+
'--worker', policy, '--unit-mib', str(args.unit_mib)],
87+
capture_output=True, text=True, check=True, timeout=300,
88+
)
89+
rows.append({'repeat': repeat, **json.loads(completed.stdout)})
90+
if len({row['digest_sha256'] for row in rows}) != 1:
91+
raise ValueError('Policies or repeats changed payload bytes/order')
92+
report = {'scope': 'Synthetic variable-buffer queue; not scientific I/O or training speed',
93+
'platform': platform.platform(), 'python': platform.python_version(),
94+
'numpy': np.__version__, 'trials': rows}
95+
args.output.parent.mkdir(parents=True, exist_ok=True)
96+
with args.output.open('x') as output:
97+
json.dump(report, output, indent=2)
98+
output.write('\n')
99+
print(json.dumps(report, indent=2))
100+
101+
102+
if __name__ == '__main__':
103+
main()

‎packages/arrayroute/src/arrayroute/execution.py‎

Lines changed: 51 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,11 @@ class OrderedPool(Generic[T, R]):
2525
Each active thread constructs one callable and keeps it until this pool closes.
2626
A pass owns at most ``max_pending`` tasks/results, including its yielded result.
2727
Caller-retained values, native caches and scratch are additional memory. This is
28-
a task-count bound, not an RSS cap. Factories must not return shared mutable readers.
28+
a task-count bound, not an RSS cap. Optional byte admission also bounds the sum
29+
of caller-supplied task estimates until each result is consumed. Estimates must
30+
cover the memory they claim; caches, scratch and caller-retained values remain
31+
additional. One not-yet-admitted input may be retained for lookahead. Factories
32+
must not return shared mutable readers.
2933
3034
Map contexts are exclusive, including before consumption starts. Task failures
3135
surface in input order; no task is retried or skipped. Leaving a map context
@@ -40,15 +44,25 @@ def __init__(
4044
*,
4145
workers: int,
4246
max_pending: int | None = None,
47+
max_pending_bytes: int | None = None,
48+
estimate_bytes: Callable[[T], int] | None = None,
4349
) -> None:
4450
if isinstance(workers, bool) or isinstance(max_pending, bool):
4551
raise ValueError("workers and max_pending must be positive integers")
4652
workers = operator.index(workers)
4753
capacity = workers if max_pending is None else operator.index(max_pending)
4854
if workers < 1 or capacity < 1:
4955
raise ValueError("workers and max_pending must be positive integers")
56+
if (max_pending_bytes is None) != (estimate_bytes is None):
57+
raise ValueError("Byte admission requires both max_pending_bytes and estimate_bytes")
58+
if max_pending_bytes is not None:
59+
if isinstance(max_pending_bytes, bool) or operator.index(max_pending_bytes) < 1:
60+
raise ValueError("max_pending_bytes must be a positive integer")
61+
max_pending_bytes = operator.index(max_pending_bytes)
5062
self._make_worker = make_worker
5163
self._capacity = capacity
64+
self._byte_capacity = max_pending_bytes
65+
self._estimate_bytes = estimate_bytes
5266
self._local = threading.local()
5367
self._state_lock = threading.Lock()
5468
self._active = False
@@ -74,7 +88,7 @@ def map(
7488
if self._active:
7589
raise RuntimeError("Concurrent map contexts are not supported")
7690
self._active = True
77-
pending: deque[Future[R]] = deque()
91+
pending: deque[tuple[Future[R], int]] = deque()
7892

7993
def check() -> None:
8094
if self._closed:
@@ -85,38 +99,61 @@ def check() -> None:
8599
def iterate() -> Generator[R, None, None]:
86100
source = iter(items)
87101
exhausted = False
102+
deferred: tuple[T, int] | None = None
103+
admitted_bytes = 0
88104
while True:
89105
check()
90106
while not exhausted and len(pending) < self._capacity:
91107
check()
92-
try:
93-
item = next(source)
94-
except StopIteration:
95-
exhausted = True
108+
if deferred is None:
109+
try:
110+
item = next(source)
111+
except StopIteration:
112+
exhausted = True
113+
break
114+
amount = 0 if self._estimate_bytes is None else self._estimate_bytes(item)
115+
if isinstance(amount, bool) or operator.index(amount) < 0:
116+
raise ValueError("Task byte estimate must be a nonnegative integer")
117+
amount = operator.index(amount)
118+
if self._byte_capacity is not None and amount > self._byte_capacity:
119+
raise MemoryError(
120+
f"Task requires {amount} planned bytes; pending-byte allowance "
121+
f"is {self._byte_capacity}. "
122+
"Reduce the task size or raise the allowance."
123+
)
124+
else:
125+
item, amount = deferred
126+
deferred = None
127+
if (self._byte_capacity is not None
128+
and admitted_bytes + amount > self._byte_capacity):
129+
deferred = (item, amount)
130+
del item
96131
break
97-
pending.append(self._executor.submit(self._execute, item, check))
132+
pending.append((self._executor.submit(self._execute, item, check), amount))
133+
admitted_bytes += amount
98134
del item
99135
if not pending:
100136
return
101137
# Keep the awaited future in the deque so cleanup also waits for it
102138
# if the consumer is interrupted while future.result() is waiting.
103-
future = pending[0]
139+
future, amount = pending[0]
104140
result = future.result()
105141
pending.popleft()
106142
del future
107143
check()
108144
yield result
109145
del result
146+
admitted_bytes -= amount
110147

111148
output = iterate()
112149
try:
113150
yield output
114151
finally:
115152
try:
116-
for future in pending:
153+
for future, _ in pending:
117154
future.cancel()
118155
if pending:
119-
wait(pending)
156+
wait(future for future, _ in pending)
120157
output.close()
121158
pending.clear()
122159
finally:
@@ -147,13 +184,16 @@ def ordered_map(
147184
*,
148185
workers: int,
149186
max_pending: int | None = None,
187+
max_pending_bytes: int | None = None,
188+
estimate_bytes: Callable[[T], int] | None = None,
150189
cancelled: Callable[[], bool] | None = None,
151190
) -> Iterator[Iterator[R]]:
152191
"""Run one bounded pass and release its thread-owned workers on context exit.
153192
154193
See ``OrderedPool`` for ordering, memory and cancellation guarantees. Use a pool
155194
explicitly when repeated passes should retain the same native readers and caches.
156195
"""
157-
with OrderedPool(make_worker, workers=workers, max_pending=max_pending) as pool:
196+
with OrderedPool(make_worker, workers=workers, max_pending=max_pending,
197+
max_pending_bytes=max_pending_bytes, estimate_bytes=estimate_bytes) as pool:
158198
with pool.map(items, cancelled=cancelled) as output:
159199
yield output
Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
import threading
3+
4+
import pytest
5+
6+
from arrayroute.execution import OrderedPool, ProcessingCancelledError, ordered_map
7+
8+
9+
def test_bytes_limit_submission_without_reordering_or_reestimating():
10+
started, estimated = [], []
11+
second = threading.Event()
12+
weights = [3, 3, 5, 1, 0]
13+
14+
def estimate(i):
15+
estimated.append(i)
16+
return weights[i]
17+
18+
def factory():
19+
def process(i):
20+
started.append(i)
21+
if i == 0:
22+
assert second.wait(5)
23+
if i == 1:
24+
second.set()
25+
return i
26+
return process
27+
28+
with ordered_map(factory, range(5), workers=3, max_pending=8,
29+
max_pending_bytes=6, estimate_bytes=estimate) as values:
30+
assert next(values) == 0
31+
assert set(started) == {0, 1}
32+
assert estimated == [0, 1, 2]
33+
assert next(values) == 1
34+
assert set(started) == {0, 1}
35+
assert list(values) == [2, 3, 4]
36+
assert estimated == list(range(5))
37+
38+
39+
def test_one_full_budget_task_progresses_and_pool_can_be_reused():
40+
with OrderedPool(lambda: lambda x: x, workers=8, max_pending=16,
41+
max_pending_bytes=10, estimate_bytes=lambda x: 10) as pool:
42+
for _ in range(3):
43+
with pool.map(range(100)) as values:
44+
assert list(values) == list(range(100))
45+
46+
47+
def test_oversized_task_is_rejected_before_worker_creation():
48+
with ordered_map(lambda: pytest.fail('No worker should be created'), [11], workers=2,
49+
max_pending_bytes=10, estimate_bytes=lambda x: x) as values:
50+
with pytest.raises(MemoryError, match='11 planned bytes.*10'):
51+
next(values)
52+
53+
54+
@pytest.mark.parametrize('budget,estimate', [(None, lambda x: 1), (1, None),
55+
(0, lambda x: 1), (True, lambda x: 1)])
56+
def test_invalid_byte_admission(budget, estimate):
57+
with pytest.raises(ValueError):
58+
OrderedPool(lambda: lambda x: x, workers=1,
59+
max_pending_bytes=budget, estimate_bytes=estimate)
60+
61+
62+
@pytest.mark.parametrize('estimate', [-1, True, 1.5])
63+
def test_invalid_task_estimate_does_not_execute(estimate):
64+
with ordered_map(lambda: pytest.fail('Invalid estimate must not execute'), [0], workers=1,
65+
max_pending_bytes=10, estimate_bytes=lambda x: estimate) as values:
66+
with pytest.raises((ValueError, TypeError)):
67+
next(values)
68+
69+
70+
def test_cancelled_deferred_task_is_not_submitted_and_next_pass_works():
71+
cancelled = False
72+
started = []
73+
with OrderedPool(lambda: lambda x: started.append(x) or x, workers=2,
74+
max_pending_bytes=1, estimate_bytes=lambda x: 1) as pool:
75+
with pool.map(range(5), cancelled=lambda: cancelled) as values:
76+
assert next(values) == 0
77+
cancelled = True
78+
with pytest.raises(ProcessingCancelledError):
79+
next(values)
80+
assert started == [0]
81+
with pool.map([9]) as values:
82+
assert list(values) == [9]
83+
84+
85+
def test_failed_task_retains_order_and_releases_pass():
86+
def process(i):
87+
if i == 1:
88+
raise OSError('decoder failed')
89+
return i
90+
with OrderedPool(lambda: process, workers=3, max_pending_bytes=5,
91+
estimate_bytes=lambda x: 3) as pool:
92+
with pool.map(range(4)) as values:
93+
assert next(values) == 0
94+
with pytest.raises(OSError, match='decoder failed'):
95+
next(values)
96+
with pool.map([8]) as values:
97+
assert list(values) == [8]

0 commit comments

Comments
 (0)