hojichar.core.async_composition

  1from __future__ import annotations
  2
  3import asyncio
  4import logging
  5from concurrent.futures import Executor, ThreadPoolExecutor
  6from typing import Any, AsyncGenerator, AsyncIterable, Iterable, Sequence
  7
  8import numpy as np
  9
 10from hojichar.core import inspection
 11from hojichar.core.async_filter_interface import AsyncFilter
 12from hojichar.core.composition import Compose
 13from hojichar.core.filter_interface import Filter
 14from hojichar.core.models import Document, Statistics, get_doc_info
 15from hojichar.utils.async_handlers import handle_stream_as_async
 16
 17
 18class AsyncFilterAdapter(AsyncFilter):
 19    """
 20    Adapter class for executing hojichar.Filter asynchronously.
 21    """
 22
 23    def __init__(
 24        self,
 25        sync_filter: Filter,
 26        *args: Any,
 27        executor: Executor | None = None,
 28        use_batch: bool = True,
 29        **kwargs: Any,
 30    ):
 31        """
 32        Adapter class for executing hojichar.Filter asynchronously.
 33        Used to incorporate Filter into AsyncCompose.
 34
 35        To reduce the overhead of asynchronous context switching,
 36        use_batch is set to True by default to process in batches
 37        in apply_stream, regardless of the sync_filter's use_batch setting.
 38
 39        If performing CPU-bound and heavy processing, you can specify an executor
 40        to offload the processing to the executor. However, due to Python's GIL
 41        constraints, using ThreadPoolExecutor will not parallelize CPU-bound
 42        processing, and the entire process will be locked.
 43
 44        By using ProcessPoolExecutor as the executor, it may be possible to
 45        parallelize CPU-bound processing. However, for parallelizing CPU-bound
 46        processing, it is recommended to use the hojichar.Parallel class to
 47        parallelize synchronous Compose pipeline.
 48        """
 49        super().__init__(*args, use_batch=use_batch, **kwargs)
 50        self.sync_filter = sync_filter
 51        self._has_external_executor = executor is not None
 52        self._executor = executor or ThreadPoolExecutor()
 53        self.batch_size = sync_filter.batch_size
 54
 55    async def apply(self, document: Document) -> Document:
 56        loop = asyncio.get_running_loop()
 57        return await loop.run_in_executor(self._executor, self.sync_filter.apply, document)
 58
 59    async def apply_batch(self, batch: Sequence[Document]) -> list[Document]:
 60        loop = asyncio.get_running_loop()
 61        return await loop.run_in_executor(
 62            self._executor,
 63            lambda: self.sync_filter.apply_batch(batch),
 64        )
 65
 66    async def shutdown(self) -> None:
 67        self.sync_filter.shutdown()
 68        if not self._has_external_executor:
 69            self._executor.shutdown()
 70
 71
 72class AsyncCompose(AsyncFilter):
 73    def __init__(
 74        self,
 75        filters: list[AsyncFilter | Filter],
 76        random_state: int | np.random.Generator | None = None,
 77        executor: ThreadPoolExecutor | None = None,
 78        *args: Any,
 79        **kwargs: Any,
 80    ):
 81        super().__init__(random_state=random_state, *args, **kwargs)
 82        self.logger = logging.getLogger(f"{self.__module__}.{self.__class__.__name__}")
 83        self._statistics.name = "Total"
 84        self._has_external_executor = executor is not None
 85        self._executor = executor or ThreadPoolExecutor()
 86        self.set_filters(filters)
 87
 88    def set_filters(self, filters: list[AsyncFilter | Filter]) -> None:
 89        self.filters: list[AsyncFilter] = []
 90        filter_idx = 0
 91        for f in filters:
 92            if isinstance(f, (AsyncCompose, Compose)):
 93                for sub in f.filters:
 94                    name = f"{filter_idx}-{sub.__class__.__name__}"
 95                    if isinstance(sub, Filter):
 96                        name = f"{filter_idx}-{sub.__class__.__name__}"
 97                        sub = AsyncFilterAdapter(sub, executor=self._executor)
 98
 99                    sub._set_rng_if_not_initialized(self._rng)
100                    sub.name = name
101                    sub._statistics.name = name
102                    self.filters.append(sub)
103                    filter_idx += 1
104            else:
105                name = f"{filter_idx}-{f.__class__.__name__}"
106                if isinstance(f, Filter):
107                    name = f"{filter_idx}-{f.__class__.__name__}"
108                    f = AsyncFilterAdapter(f, executor=self._executor)
109                f._set_rng_if_not_initialized(self._rng)
110                f.name = name
111                f._statistics.name = name
112                self.filters.append(f)
113                filter_idx += 1
114
115    async def apply(self, document: Document) -> Document:
116        stat = get_doc_info(document)
117        for filter_idx, filt in enumerate(self.filters):
118            document = await filt._apply(document)
119        new_stat = get_doc_info(document)
120        async with self._stats_lock:
121            self._statistics.update_by_diff(stat, new_stat)
122        return document
123
124    async def apply_batch(self, batch: Sequence[Document]) -> list[Document]:
125        stats = [get_doc_info(doc) for doc in batch]
126        for i, filt in enumerate(self.filters):
127            batch = await filt._apply_batch(batch)
128        batch = await self._finalize_batch(batch, stats)
129        return list(batch)
130
131    async def apply_stream(
132        self,
133        stream: AsyncIterable[Document] | Iterable[Document],
134    ) -> AsyncGenerator[Document, None]:
135        async_stream = handle_stream_as_async(stream, chunk_size=1000, executor=self._executor)
136        async_stream = self._count_input_stats(async_stream)
137
138        for i, filt in enumerate(self.filters):
139            async_stream = filt.apply_stream(async_stream)
140
141        async for doc in async_stream:
142            in_stat = doc._get_initial_stats()
143            if in_stat is None:
144                in_stat = get_doc_info(doc)
145                self.logger.debug(
146                    "Initial stats missing for document during async stream aggregation; "
147                    "using current stats as fallback"
148                )
149            out_stat = get_doc_info(doc)
150            async with self._stats_lock:
151                self._statistics.update_by_diff(in_stat, out_stat)
152            doc._clear_initial_stats()
153            yield doc
154
155    async def _count_input_stats(
156        self, async_stream: AsyncIterable[Document]
157    ) -> AsyncGenerator[Document, None]:
158        async for doc in async_stream:
159            doc._set_initial_stats(get_doc_info(doc))
160            yield doc
161
162    def get_total_statistics(self) -> list[Statistics]:
163        """
164        Get the statistics of the Compose object and sub filters.
165
166        The statistics of the Compose class are stored in an object with the name "Total",
167        and sub-filters's are stored with names in the format {filter_index}-{filter class name}.
168        """
169        stats = []
170        stats.append(self.get_statistics())
171        for i, filt in enumerate(self.filters):
172            stats.append(filt.get_statistics())
173        return stats
174
175    def get_total_statistics_map(self) -> list[dict[str, Any]]:
176        """
177        Get the statistics of the Compose object and sub filters as a list of dictionaries.
178        """
179        stats = self.get_total_statistics()
180        return [stat.to_dict() for stat in stats]
181
182    @property
183    def statistics(self) -> dict:
184        """
185        Deprecated
186
187        Get the statistics of the Compose object and sub filters.
188
189        This property is retained for compatibility with previous versions.
190        Please use `get_total_statistics` or `get_total_statistics_map` instead.
191        """
192        return inspection.statistics_obj_adapter(  # type: ignore
193            self.get_total_statistics()
194        ).get_human_readable_values()
195
196    @property
197    def statistics_obj(self) -> inspection.StatsContainer:
198        """
199        Deprecated
200
201        Get the statistics of the AsyncCompose object and sub filters.
202        This method returns a StatsContainer object which contains the statistics
203        of the AsyncCompose object and sub filters.
204
205        This property is retained for compatibility with previous versions.
206        Please use `get_total_statistics` or `get_total_statistics_map` instead.
207        """
208        return inspection.statistics_obj_adapter(self.get_total_statistics())  # type: ignore
209
210    async def shutdown(self) -> None:
211        for filt in self.filters:
212            await filt.shutdown()
213        if not self._has_external_executor:
214            self._executor.shutdown()
215
216    async def __aenter__(self) -> "AsyncCompose":
217        return self
218
219    async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
220        await self.shutdown()
221        if exc_type is not None:
222            raise exc_value