Spaces:
Running
Running
| """Bounded process pool for the nightly PDF parse stages (certificates and audit reports). | |
| Workers only fetch and parse; the parent records every outcome, so the SQLite work DB keeps a | |
| single writer. Design: docs/superpowers/specs/2026-10-05-nightly-parallel-pdf-parsing-design.md. | |
| """ | |
| from __future__ import annotations | |
| import time | |
| from collections import deque | |
| from concurrent.futures import FIRST_COMPLETED, Future, ProcessPoolExecutor, wait | |
| from concurrent.futures.process import BrokenProcessPool | |
| from dataclasses import asdict, dataclass, field | |
| from typing import Callable, Generic, Mapping, Sequence, TypeVar | |
| T = TypeVar("T") | |
| R = TypeVar("R") | |
| # A pool broken by a dying worker fails every in-flight future, so the culprit is unknown. | |
| # Each affected item gets this many fresh attempts before it is recorded as a crash. | |
| BROKEN_POOL_RETRIES = 1 | |
| class PoolStats: | |
| workers: int = 0 | |
| submitted: int = 0 | |
| completed: int = 0 | |
| timeouts: int = 0 | |
| crashes: int = 0 | |
| recycles: int = 0 | |
| stopped_no_time: int = 0 | |
| wall_s: float = 0.0 | |
| peak_child_rss_mb: float | None = None | |
| peak_in_flight: int = 0 | |
| peak_by_capacity: dict[str, int] = field(default_factory=dict) | |
| submitted_by_capacity: dict[str, int] = field(default_factory=dict) | |
| completed_by_capacity: dict[str, int] = field(default_factory=dict) | |
| stopped_by_capacity: dict[str, int] = field(default_factory=dict) | |
| first_admission_s: dict[str, float] = field(default_factory=dict) | |
| last_completion_s: dict[str, float] = field(default_factory=dict) | |
| def as_dict(self) -> dict: | |
| return asdict(self) | |
| class _Pending(Generic[T]): | |
| item: T | |
| attempts: int = 0 | |
| deadline: float | None = None | |
| def peak_child_rss_mb() -> float | None: | |
| """Largest resident set of any reaped child process (Linux/macOS); None on Windows.""" | |
| try: | |
| import resource | |
| except ImportError: | |
| return None | |
| peak = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss | |
| # Linux reports KiB, macOS bytes. | |
| import sys | |
| return round(peak / (1024 * 1024 if sys.platform == "darwin" else 1024), 1) | |
| def _terminate(pool: ProcessPoolExecutor) -> None: | |
| """Kill the pool's worker processes; shutdown(wait=False) alone leaves a wedged one running.""" | |
| # _processes is private but has been the executor's pid -> Process map since Python 3.2. | |
| for process in list((getattr(pool, "_processes", None) or {}).values()): | |
| try: | |
| process.terminate() | |
| except Exception: # noqa: BLE001 -- already gone | |
| pass | |
| pool.shutdown(wait=False, cancel_futures=True) | |
| def run_bounded_pool( | |
| items: Sequence[T], | |
| work: Callable[[T], R], | |
| record: Callable[[T, R], None], | |
| *, | |
| workers: int, | |
| on_timeout: Callable[[T], R], | |
| on_crash: Callable[[T, BaseException], R], | |
| should_continue: Callable[[], bool] | None = None, | |
| lane: Callable[[T], str] | None = None, | |
| lane_limits: Mapping[str, int] | None = None, | |
| capacity_keys: Callable[[T], Sequence[str]] | None = None, | |
| capacity_limits: Mapping[str, int] | None = None, | |
| item_timeout_s: float = 180.0, | |
| task_timeout_s: Callable[[T], float] | None = None, | |
| max_tasks_per_child: int = 25, | |
| initializer: Callable[..., None] | None = None, | |
| initargs: tuple = (), | |
| ) -> PoolStats: | |
| """Run ``work`` over ``items`` in worker processes, recording each outcome in the parent. | |
| ``work`` must be a picklable top-level function. Submission follows ``items`` order, except | |
| that an item whose lane is at its limit waits while later items from other lanes start. | |
| Alternatively ``capacity_keys`` supplies multiple independent capacity names per item; | |
| every matching capacity must have space. The legacy lane arguments remain unchanged. | |
| ``should_continue`` is asked before every submission; once it says no, nothing new starts, | |
| in-flight items finish, and the never-started count lands in ``stopped_no_time``. An item | |
| running longer than ``item_timeout_s`` is recorded through ``on_timeout`` and the worker | |
| processes are killed; innocent in-flight items are re-queued at the front. ``initializer`` | |
| runs once in every worker process (also after a recycle) to set up per-process state. | |
| ``task_timeout_s``, when supplied, returns the fixed timeout for each item in the parent; | |
| omission retains the legacy shared ``item_timeout_s`` deadline. | |
| """ | |
| workers = max(1, int(workers)) | |
| if capacity_keys is not None and (lane is not None or lane_limits is not None): | |
| raise ValueError("use capacity_keys/capacity_limits or lane/lane_limits, not both") | |
| limits = {name: max(1, int(limit)) for name, limit in (lane_limits or {}).items()} | |
| lane_of = lane or (lambda _item: "") | |
| if capacity_keys is not None: | |
| limits = dict(capacity_limits or {}) | |
| if any(type(limit) is not int or limit < 1 for limit in limits.values()): | |
| raise ValueError("capacity limits must be positive integers") | |
| def keys_of(item: T) -> tuple[str, ...]: | |
| return tuple(dict.fromkeys(capacity_keys(item))) if capacity_keys else (lane_of(item),) | |
| def expires_at(entry: _Pending[T], submitted_at: float) -> float: | |
| # Mixed deadlines belong to the logical document, including innocent preemptions. | |
| # Legacy shared-timeout mode deliberately retains its original per-submission clock. | |
| return entry.deadline if task_timeout_s is not None else submitted_at + item_timeout_s | |
| stats = PoolStats(workers=workers) | |
| started = time.monotonic() | |
| pending: deque[_Pending[T]] = deque(_Pending(item) for item in items) | |
| stopping = False | |
| def _next_submittable(running_lanes: dict[str, int]) -> _Pending[T] | None: | |
| for index, entry in enumerate(pending): | |
| if all(running_lanes.get(name, 0) < limits.get(name, workers) | |
| for name in keys_of(entry.item)): | |
| del pending[index] | |
| return entry | |
| return None | |
| def _finish(item: T, result: R) -> None: | |
| record(item, result) | |
| stats.completed += 1 | |
| for name in keys_of(item): | |
| stats.completed_by_capacity[name] = stats.completed_by_capacity.get(name, 0) + 1 | |
| stats.last_completion_s[name] = time.monotonic() - started | |
| while pending and not stopping: | |
| pool = ProcessPoolExecutor( | |
| max_workers=workers, | |
| max_tasks_per_child=max_tasks_per_child, | |
| initializer=initializer, | |
| initargs=initargs, | |
| ) | |
| in_flight: dict[Future, tuple[_Pending[T], float]] = {} | |
| running_lanes: dict[str, int] = {} | |
| recycle = False | |
| try: | |
| while (pending and not stopping) or in_flight: | |
| while not stopping and pending and len(in_flight) < workers: | |
| if should_continue is not None and not should_continue(): | |
| stopping = True | |
| break | |
| entry = _next_submittable(running_lanes) | |
| if entry is None: | |
| break | |
| if task_timeout_s is not None: | |
| now = time.monotonic() | |
| if entry.deadline is None: | |
| entry.deadline = now + task_timeout_s(entry.item) | |
| elif entry.deadline <= now: | |
| stats.timeouts += 1 | |
| _finish(entry.item, on_timeout(entry.item)) | |
| continue | |
| in_flight[pool.submit(work, entry.item)] = (entry, time.monotonic()) | |
| for name in keys_of(entry.item): | |
| running_lanes[name] = running_lanes.get(name, 0) + 1 | |
| stats.peak_by_capacity[name] = max(stats.peak_by_capacity.get(name, 0), running_lanes[name]) | |
| stats.submitted_by_capacity[name] = stats.submitted_by_capacity.get(name, 0) + 1 | |
| stats.first_admission_s.setdefault(name, time.monotonic() - started) | |
| stats.submitted += 1 | |
| stats.peak_in_flight = max(stats.peak_in_flight, len(in_flight)) | |
| if not in_flight: | |
| break | |
| next_deadline = min(expires_at(entry, submitted_at) | |
| for entry, submitted_at in in_flight.values()) | |
| wait_s = max(0.0, next_deadline - time.monotonic()) | |
| finished, _ = wait(set(in_flight), timeout=wait_s, return_when=FIRST_COMPLETED) | |
| broken: list[_Pending[T]] = [] | |
| for future in finished: | |
| entry, _submitted_at = in_flight.pop(future) | |
| for name in keys_of(entry.item): | |
| running_lanes[name] -= 1 | |
| try: | |
| result = future.result() | |
| except BrokenProcessPool: | |
| broken.append(entry) | |
| continue | |
| except Exception as exc: # noqa: BLE001 -- work() raised inside the worker | |
| stats.crashes += 1 | |
| _finish(entry.item, on_crash(entry.item, exc)) | |
| continue | |
| _finish(entry.item, result) | |
| if broken: | |
| # Everything still in flight is lost with the pool too. | |
| broken.extend(entry for entry, _submitted_at in in_flight.values()) | |
| in_flight.clear() | |
| for entry in reversed(broken): | |
| if entry.attempts < BROKEN_POOL_RETRIES: | |
| entry.attempts += 1 | |
| pending.appendleft(entry) | |
| else: | |
| stats.crashes += 1 | |
| _finish(entry.item, on_crash(entry.item, BrokenProcessPool("worker process died"))) | |
| recycle = True | |
| break | |
| now = time.monotonic() | |
| overdue = [ | |
| future for future, (entry, submitted_at) in in_flight.items() | |
| if now >= expires_at(entry, submitted_at) | |
| ] | |
| if overdue: | |
| for future in overdue: | |
| entry, _submitted_at = in_flight.pop(future) | |
| stats.timeouts += 1 | |
| _finish(entry.item, on_timeout(entry.item)) | |
| for entry, _submitted_at in reversed(list(in_flight.values())): | |
| pending.appendleft(entry) | |
| in_flight.clear() | |
| recycle = True | |
| break | |
| except BaseException: | |
| _terminate(pool) | |
| raise | |
| if recycle: | |
| stats.recycles += 1 | |
| _terminate(pool) | |
| else: | |
| pool.shutdown(wait=True) | |
| stats.stopped_no_time = len(pending) | |
| for entry in pending: | |
| for name in keys_of(entry.item): | |
| stats.stopped_by_capacity[name] = stats.stopped_by_capacity.get(name, 0) + 1 | |
| stats.wall_s = round(time.monotonic() - started, 1) | |
| stats.peak_child_rss_mb = peak_child_rss_mb() | |
| return stats | |