@@ -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
0 commit comments