hojichar.core.parallel

  1from __future__ import annotations
  2
  3import functools
  4import logging
  5import os
  6import signal
  7import threading
  8from copy import copy
  9from multiprocessing.pool import Pool
 10from typing import Iterator, List
 11
 12import hojichar
 13from hojichar.core import inspection
 14from hojichar.core.models import Statistics
 15
 16logger = logging.getLogger(__name__)
 17
 18
 19PARALLEL_BASE_FILTER: hojichar.Compose
 20WORKER_PARAM_IGNORE_ERRORS: bool
 21
 22
 23def _init_worker(filter: hojichar.Compose, ignore_errors: bool) -> None:
 24    signal.signal(signal.SIGINT, signal.SIG_IGN)
 25    global PARALLEL_BASE_FILTER, WORKER_PARAM_IGNORE_ERRORS
 26    PARALLEL_BASE_FILTER = hojichar.Compose(copy(filter.filters))  # TODO random state treatment
 27    WORKER_PARAM_IGNORE_ERRORS = ignore_errors
 28
 29
 30def _worker(
 31    doc: hojichar.Document,
 32) -> tuple[hojichar.Document, int, List[Statistics], str | None]:
 33    global PARALLEL_BASE_FILTER, WORKER_PARAM_IGNORE_ERRORS
 34    ignore_errors = WORKER_PARAM_IGNORE_ERRORS
 35    error_message = None
 36    try:
 37        result = PARALLEL_BASE_FILTER.apply(doc)
 38    except Exception as e:
 39        if ignore_errors:
 40            logger.error(e)
 41            error_message = str(e)
 42            result = hojichar.Document("", is_rejected=True)
 43        else:
 44            raise e  # If we're not ignoring errors, let this one propagate
 45    return result, os.getpid(), PARALLEL_BASE_FILTER.get_total_statistics(), error_message
 46
 47
 48class _InFlightGate:
 49    """Bounds documents drawn from the input but not yet returned to the caller.
 50
 51    ``feed`` wraps the input iterator and acquires one permit *before* each
 52    document is drawn (acquiring afterwards would hold one extra pre-fetched
 53    document while waiting). ``imap_apply`` releases the permit only after the
 54    corresponding result has been handed back to the caller, so at most
 55    ``max_in_flight`` documents exist anywhere between the input iterator and
 56    the caller at any moment.
 57
 58    ``feed`` runs inside Pool's task-handler thread. The polling acquire
 59    observes :meth:`stop` before and after each successful acquire (handing
 60    the permit back if stopped), and abnormal exits stop the gate before
 61    returning their permit, so the feeder never draws from the source after
 62    consumption ended and pool shutdown never blocks on the gate.
 63    """
 64
 65    def __init__(self, max_in_flight: int) -> None:
 66        self._semaphore = threading.Semaphore(max_in_flight)
 67        self._stop_feeding = threading.Event()
 68
 69    def feed(self, docs: Iterator[hojichar.Document]) -> Iterator[hojichar.Document]:
 70        iterator = iter(docs)
 71        while self._acquire():
 72            try:
 73                doc = next(iterator)
 74            except StopIteration:
 75                self.release()
 76                return
 77            yield doc
 78
 79    def release(self) -> None:
 80        self._semaphore.release()
 81
 82    def stop(self) -> None:
 83        """Unblock the feeder; called when consumption ends and at shutdown."""
 84        self._stop_feeding.set()
 85
 86    def _acquire(self) -> bool:
 87        while not self._stop_feeding.is_set():
 88            if self._semaphore.acquire(timeout=0.1):
 89                if self._stop_feeding.is_set():
 90                    # Stopped while blocked on (or right after) the acquire:
 91                    # hand the permit back instead of drawing one more
 92                    # document from a possibly blocking source.
 93                    self._semaphore.release()
 94                    return False
 95                return True
 96        return False
 97
 98
 99class Parallel:
100    """
101    The Parallel class provides a way to apply a hojichar.Compose filter
102    to an iterator of documents in a parallel manner using a specified
103    number of worker processes. This class should be used as a context
104    manager with a 'with' statement.
105
106    Example:
107
108    doc_iter = (hojichar.Document(d) for d in open("my_text.txt"))
109    with Parallel(my_filter, num_jobs=8) as pfilter:
110        for doc in pfilter.imap_apply(doc_iter):
111            pass  # Process the filtered document as needed.
112    """
113
114    def __init__(
115        self,
116        filter: hojichar.Compose,
117        num_jobs: int | None = None,
118        ignore_errors: bool = False,
119        ordered: bool = False,
120        max_in_flight: int | None = None,
121    ):
122        """
123        Initializes a new instance of the Parallel class.
124
125        Args:
126            filter (hojichar.Compose): A composed filter object that specifies the
127                processing operations to apply to each document in parallel.
128                A copy of the filter is made within a 'with' statement. When the 'with'
129                block terminates,the statistical information obtained through `filter.statistics`
130                or`filter.statistics_obj` is replaced with the total value of the statistical
131                information processed within the 'with' block.
132
133            num_jobs (int | None, optional): The number of worker processes to use.
134                If None, then the number returned by os.cpu_count() is used. Defaults to None.
135            ignore_errors (bool, optional): If set to True, any exceptions thrown during
136                the processing of a document will be caught and logged, but will not
137                stop the processing of further documents. If set to False, the first
138                exception thrown will terminate the entire parallel processing operation.
139                Defaults to False.
140            ordered (bool, optional): If set to True, processed documents are yielded in
141                the same order as the input documents. If set to False, documents are
142                yielded as soon as their processing completes. Defaults to False.
143            max_in_flight (int | None, optional): Strict upper bound on the number of
144                documents drawn from the input iterator whose results have not yet been
145                handed back to the caller. A yielded document keeps its permit until the
146                caller requests the next one, so with `max_in_flight=1` the pool holds a
147                single document end to end (no pipelining). Without a bound, Pool's
148                task-handler thread drains the input as fast as the worker pipe accepts,
149                so a producer that outruns the filters can buffer a large number of
150                documents. If None, no explicit bound is applied. Defaults to None.
151        """
152        if max_in_flight is not None and max_in_flight < 1:
153            raise ValueError("max_in_flight must be at least 1")
154        self.filter = filter
155        self.num_jobs = num_jobs
156        self.ignore_errors = ignore_errors
157        self.ordered = ordered
158        self.max_in_flight = max_in_flight
159
160        self._pool: Pool | None = None
161        self._pid_stats: dict[int, List[Statistics]] | None = None
162        self._gates: list[_InFlightGate] = []
163
164    def __enter__(self) -> Parallel:
165        self._pool = Pool(
166            processes=self.num_jobs,
167            initializer=_init_worker,
168            initargs=(self.filter, self.ignore_errors),
169        )
170        self._pid_stats = dict()
171        self._gates = []
172        return self
173
174    def imap_apply(self, docs: Iterator[hojichar.Document]) -> Iterator[hojichar.Document]:
175        """
176        Takes an iterator of Documents and applies the Compose filter to
177        each Document in a parallel manner. This is a generator method
178        that yields processed Documents.
179
180        Args:
181            docs (Iterator[hojichar.Document]): An iterator of Documents to be processed.
182
183        Raises:
184            RuntimeError: If the Parallel instance is not properly initialized. This
185                generally happens when the method is called outside of a 'with' statement.
186            Exception: If any exceptions are raised within the worker processes.
187
188        Yields:
189            Iterator[hojichar.Document]: An iterator that yields processed Documents.
190        """
191        if self._pool is None or self._pid_stats is None:
192            raise RuntimeError(
193                "Parallel instance not properly initialized. Use within a 'with' statement."
194            )
195        gate: _InFlightGate | None = None
196        if self.max_in_flight is not None:
197            gate = _InFlightGate(self.max_in_flight)
198            self._gates.append(gate)
199            docs = gate.feed(docs)
200        try:
201            results = (
202                self._pool.imap(_worker, docs)
203                if self.ordered
204                else self._pool.imap_unordered(_worker, docs)
205            )
206            for doc, pid, stat, err_msg in results:
207                self._pid_stats[pid] = stat
208                if err_msg is not None:
209                    logger.error(f"Error in worker {pid}: {err_msg}")
210                # The permit is returned only once the caller has taken the
211                # document: releasing before the yield would let the feeder
212                # momentarily draw max_in_flight + 1 documents. On an abnormal
213                # exit (close/throw at the yield point) the gate is stopped
214                # *before* the permit is returned, so a feeder woken by the
215                # release always observes the stop and cannot draw again.
216                try:
217                    yield doc
218                except BaseException:
219                    if gate is not None:
220                        gate.stop()
221                        gate.release()
222                    raise
223                else:
224                    if gate is not None:
225                        gate.release()
226        except Exception:
227            self.__exit__(None, None, None)
228            raise
229        finally:
230            if gate is not None:
231                gate.stop()
232
233    def __exit__(self, exc_type, exc_value, traceback) -> None:  # type: ignore
234        # Feeders blocked on a max_in_flight gate must exit before the pool is
235        # joined, or shutdown would wait on them forever.
236        for gate in self._gates:
237            gate.stop()
238        if self._pool:
239            self._pool.terminate()
240            self._pool.join()
241        if self._pid_stats:
242            total_stats = functools.reduce(
243                lambda x, y: Statistics.add_list_of_stats(x, y), self._pid_stats.values()
244            )
245            self.filter._statistics.update(Statistics.get_filter("Total", total_stats))
246            for stat in total_stats:
247                for filt in self.filter.filters:
248                    if stat.name == filt.name:
249                        filt._statistics.update(stat)
250                        break
251
252    def get_total_statistics(self) -> List[Statistics]:
253        """
254        Returns a statistics object of the total statistical
255        values processed within the Parallel block.
256
257        Returns:
258            StatsContainer: Statistics object
259        """
260        if self._pid_stats:
261            total_stats = functools.reduce(
262                lambda x, y: Statistics.add_list_of_stats(x, y), self._pid_stats.values()
263            )
264            return total_stats
265        else:
266            return []
267
268    def get_total_statistics_map(self) -> List[dict]:
269        return [stat.to_dict() for stat in self.get_total_statistics()]
270
271    @property
272    def statistics_obj(self) -> inspection.StatsContainer:
273        """
274        Returns the statistics object of the Parallel instance.
275        This is a StatsContainer object which contains the statistics
276        of the Parallel instance and sub filters.
277
278        Returns:
279            StatsContainer: Statistics object
280        """
281        return inspection.statistics_obj_adapter(self.get_total_statistics())  # type: ignore