async-batcher 0.1.0__py3-none-any.whl

async_batcher/__init__.py

@@ -0,0 +1 @@
+

async_batcher/aws/dynamodb/get.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import aioboto3
+
+from async_batcher.batcher import AsyncBatcher
+
+if TYPE_CHECKING:
+    from aiobotocore.config import AioConfig
+    from types_aiobotocore_dynamodb import DynamoDBServiceResource
+    from types_aiobotocore_dynamodb.type_defs import TableAttributeValueTypeDef
+
+
+@dataclass(kw_only=True)
+class GetItem:
+    table_name: str
+    key: dict[str, TableAttributeValueTypeDef]
+
+
+class AsyncDynamoDbGetBatcher(AsyncBatcher[GetItem, dict[str, Any]]):
+    """Batcher for DynamoDB GetItem operation. It uses aioboto3 to interact with DynamoDB.
+
+    Args:
+        region_name: The region to use.
+        use_ssl: Whether to use SSL/TLS.
+        verify: Whether to verify SSL certificates.
+        endpoint_url: The complete URL to use for the constructed client. This is useful for local testing.
+        config: The configuration for the session.
+        aioboto3_session: The aioboto3 session to use. If not provided, a new session is created.
+        batch_size: The maximum number of items to process in a single batch. The default is 100 items,
+            which is the maximum number of items that can be processed in a single batch.
+        sleep_time (float, optional): The time to sleep between checking if the result is ready in seconds.
+            Defaults to 0.01. Set it to a value close to the expected time to process a batch
+        buffering_time (float, optional): The time to sleep after processing a batch or checking the buffer
+            in seconds. Defaults to 0.001.
+            You can increase this value if you don't need a low latency, but want to reduce the number of
+            processed batches.
+    """
+
+    def __init__(
+        self,
+        *,
+        region_name: str | None = None,
+        use_ssl: bool | None = None,
+        verify: bool | None = None,
+        endpoint_url: str | None = None,
+        config: AioConfig | None = None,
+        aioboto3_session: aioboto3.Session | None = None,
+        batch_size: int = 100,
+        sleep_time: float = 0.01,
+        buffering_time: float = 0.001,
+    ):
+        super().__init__(batch_size=batch_size, sleep_time=sleep_time, buffering_time=buffering_time)
+        self.region_name = region_name
+        self.use_ssl = use_ssl
+        self.verify = verify
+        self.endpoint_url = endpoint_url
+        self.config = config
+        self.aioboto3_session = aioboto3_session or aioboto3.Session()
+
+    async def process_batch(self, batch: list[GetItem]) -> list[dict[str, TableAttributeValueTypeDef]]:
+        indexed_items: dict[tuple, int] = {}
+        request_items = {}
+        tables_keys = {}
+        for ind, item in enumerate(batch):
+            if item.table_name not in request_items:
+                request_items[item.table_name] = {"Keys": []}
+                tables_keys[item.table_name] = sorted(item.key.keys())
+            # TODO: support ProjectionExpression and ConsistentRead
+            request_items[item.table_name]["Keys"].append(item.key)
+            indexed_items[(item.table_name, *[item.key[key] for key in tables_keys[item.table_name]])] = ind
+
+        dynamodb: DynamoDBServiceResource
+        async with self.aioboto3_session.resource(
+            "dynamodb",
+            region_name=self.region_name,
+            use_ssl=self.use_ssl,
+            verify=self.verify,
+            endpoint_url=self.endpoint_url,
+            config=self.config,
+        ) as dynamodb:
+            response = await dynamodb.batch_get_item(
+                RequestItems=request_items,
+                ReturnConsumedCapacity="NONE",
+            )
+            result: list[None | dict[str, TableAttributeValueTypeDef]] = [None] * len(batch)
+            # TODO: handle UnprocessedKeys
+            for table in response["Responses"]:
+                for item in response["Responses"][table]:
+                    index = indexed_items[(table, *[item[key] for key in tables_keys[table]])]
+                    result[index] = item
+            return result

async_batcher/aws/dynamodb/write.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Literal
+
+import aioboto3
+
+from async_batcher.batcher import AsyncBatcher
+
+if TYPE_CHECKING:
+    from aiobotocore.config import AioConfig
+    from types_aiobotocore_dynamodb import DynamoDBServiceResource
+    from types_aiobotocore_dynamodb.type_defs import TableAttributeValueTypeDef
+
+
+@dataclass(kw_only=True)
+class WriteOperation:
+    operation: Literal["PUT", "DELETE"]
+    table_name: str
+    data: dict[str, TableAttributeValueTypeDef]
+
+
+class AsyncDynamoDbWriteBatcher(AsyncBatcher[WriteOperation, None]):
+    """Batcher for DynamoDB WriteOperation. It uses aioboto3 to interact with DynamoDB.
+
+    Args:
+        region_name: The region to use.
+        use_ssl: Whether to use SSL/TLS.
+        verify: Whether to verify SSL certificates.
+        endpoint_url: The complete URL to use for the constructed client. This is useful for local testing.
+        config: The configuration for the session.
+        aioboto3_session: The aioboto3 session to use. If not provided, a new session is created.
+        batch_size: The maximum number of items to process in a single batch. The default is 25 items,
+            which is the maximum number of items that can be processed in a single batch.
+        sleep_time (float, optional): The time to sleep between checking if the result is ready in seconds.
+            Defaults to 0.01. Set it to a value close to the expected time to process a batch
+        buffering_time (float, optional): The time to sleep after processing a batch or checking the buffer
+            in seconds. Defaults to 0.001.
+            You can increase this value if you don't need a low latency, but want to reduce the number of
+            processed batches.
+    """
+
+    def __init__(
+        self,
+        *,
+        region_name: str | None = None,
+        use_ssl: bool | None = None,
+        verify: bool | None = None,
+        endpoint_url: str | None = None,
+        config: AioConfig | None = None,
+        aioboto3_session: aioboto3.Session | None = None,
+        batch_size: int = 25,
+        sleep_time: float = 0.01,
+        buffering_time: float = 0.001,
+    ):
+        super().__init__(batch_size=batch_size, sleep_time=sleep_time, buffering_time=buffering_time)
+        self.region_name = region_name
+        self.use_ssl = use_ssl
+        self.verify = verify
+        self.endpoint_url = endpoint_url
+        self.config = config
+        self.aioboto3_session = aioboto3_session or aioboto3.Session()
+
+    async def process_batch(self, batch: list[WriteOperation]) -> None:
+        request_items = {}
+        for operation in batch:
+            if operation.table_name not in request_items:
+                request_items[operation.table_name] = []
+            request = {}
+            if operation.operation == "PUT":
+                request["PutRequest"] = {"Item": operation.data}
+            elif operation.operation == "DELETE":
+                request["DeleteRequest"] = {"Key": operation.data}
+            request_items[operation.table_name].append(request)
+
+        dynamodb: DynamoDBServiceResource
+        async with self.aioboto3_session.resource(
+            "dynamodb",
+            region_name=self.region_name,
+            use_ssl=self.use_ssl,
+            verify=self.verify,
+            endpoint_url=self.endpoint_url,
+            config=self.config,
+        ) as dynamodb:
+            # TODO: return something useful
+            await dynamodb.batch_write_item(
+                RequestItems=request_items,
+                ReturnConsumedCapacity="NONE",
+                ReturnItemCollectionMetrics="NONE",
+            )

async_batcher/batcher.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+import abc
+import asyncio
+import logging
+import uuid
+from collections import deque
+from threading import Thread
+from typing import Generic, TypeVar
+
+T = TypeVar("T")
+S = TypeVar("S")
+
+
+class AsyncBatcher(Generic[T, S], abc.ABC, Thread):
+    """A generic class for batching and processing items asynchronously.
+
+    Args:
+        batch_size (int, optional): The max number of items to process in a batch. Defaults to -1 (no limit).
+        sleep_time (float, optional): The time to sleep between checking if the result is ready in seconds.
+            Defaults to 0.01. Set it to a value close to the expected time to process a batch
+        buffering_time (float, optional): The time to sleep after processing a batch or checking the buffer
+            in seconds. Defaults to 0.001.
+            You can increase this value if you don't need a low latency, but want to reduce the number of
+            processed batches.
+    """
+
+    _buffer = deque[tuple[uuid.UUID, T]]
+    _results = dict[uuid.UUID, S]
+    logger = logging.getLogger(__name__)
+
+    def __init__(
+        self,
+        *,
+        batch_size: int = -1,
+        sleep_time: float = 0.01,
+        buffering_time: float = 0.001,
+    ):
+        super().__init__()
+        self.batch_size = batch_size
+        self.sleep_time = sleep_time
+        self.buffering_time = buffering_time
+        self._buffer = deque()
+        self._results = {}
+        self._current_batch = 0
+        self._should_stop = False
+        self._force_stop = False
+        self._is_running = False
+
+    @abc.abstractmethod
+    async def process_batch(self, batch: list[T]) -> list[S] | None:
+        """Process a batch of items.
+
+        This method should be overridden by the user to define how to process a batch of items.
+        """
+
+    async def _process_single(self, *, item: T) -> S:
+        """Process a single item.
+
+        This method adds the item to the buffer and waits for the result to be ready.
+        It's used by the `process` method to add an item to the buffer and get the result when it's ready.
+
+        Args:
+            item (T): The item to process.
+
+        Returns:
+            S: The result of processing the item.
+        """
+        started_at = asyncio.get_event_loop().time()
+        query_id = uuid.uuid4()
+        last_checked_batch = self._current_batch
+        self.logger.debug(f"Adding item {query_id} to the buffer.")
+        self._buffer.append((query_id, item))
+        while True:
+            if self._current_batch > last_checked_batch:
+                last_checked_batch = self._current_batch
+                if query_id in self._results:
+                    elapsed_time = asyncio.get_event_loop().time() - started_at
+                    self.logger.debug(f"Item {query_id} is ready after {elapsed_time} seconds.")
+                    return self._results.pop(query_id)
+            await asyncio.sleep(self.sleep_time)
+
+    async def process(self, item: T) -> S:
+        """Add an item to the buffer and get the result when it's ready.
+
+        Args:
+            item (T): The item to process.
+
+        Returns:
+            S: The result of processing the item.
+        """
+        if not self._is_running:
+            raise RuntimeError("The batcher is not running.")
+        result = await asyncio.get_event_loop().create_task(self._process_single(item=item))
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    async def _arun(self):
+        """Run the batcher asynchronously."""
+        self._is_running = True
+        while not self._should_stop or (not self._force_stop and len(self._buffer) > 0):
+            ids = []
+            batch = []
+            while self._buffer and (len(batch) < self.batch_size or self.batch_size == -1):
+                query_id, item = self._buffer.popleft()
+                batch.append(item)
+                ids.append(query_id)
+            if batch:
+                started_at = asyncio.get_event_loop().time()
+                try:
+                    results = await asyncio.get_event_loop().create_task(self.process_batch(batch=batch))
+                    if results is None:
+                        results = [None] * len(batch)
+                    if len(results) != len(batch):
+                        raise ValueError(f"Expected to get {len(batch)} results, but got {len(results)}.")
+                except Exception as e:
+                    self.logger.error("Error processing batch", exc_info=True)
+                    results = [e] * len(batch)
+                for query_id, result in zip(ids, results):
+                    self._results[query_id] = result
+                elapsed_time = asyncio.get_event_loop().time() - started_at
+                self.logger.debug(
+                    f"Processed batch {self._current_batch} of {len(batch)} elements"
+                    f" in {elapsed_time} seconds."
+                )
+                self._current_batch += 1
+                await asyncio.sleep(max(self.buffering_time - elapsed_time, 0))
+            else:
+                self.logger.debug("No items to process. Sleeping.")
+                await asyncio.sleep(self.buffering_time)
+        self._is_running = False
+
+    def run(self):
+        """Run the batcher thread."""
+        asyncio.run(self._arun())
+
+    def stop(self, force: bool = False):
+        """Stop the batcher thread.
+
+        Args:
+            force (bool, optional): Whether to force stop the batcher without waiting for processing
+                the remaining buffer items. Defaults to False.
+        """
+        if force:
+            self._force_stop = True
+        self._should_stop = True

async_batcher/ml/keras/model.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING
+
+from async_batcher.batcher import AsyncBatcher
+
+if TYPE_CHECKING:
+    from concurrent.futures import Executor
+
+    from keras import Model
+
+
+class KerasAsyncBatcher(AsyncBatcher):
+    """Batcher for Keras models.
+
+    Args:
+        model: The Keras model to use for prediction.
+        executor: The executor to use for running the prediction.
+        batch_size (int, optional): The max number of items to process in a batch. Defaults to -1 (no limit).
+        sleep_time (float, optional): The time to sleep between checking if the result is ready in seconds.
+            Defaults to 0.01. Set it to a value close to the expected time to process a batch
+        buffering_time (float, optional): The time to sleep after processing a batch or checking the buffer
+            in seconds. Defaults to 0.001.
+            You can increase this value if you don't need a low latency, but want to reduce the number of
+            processed batches.
+    """
+
+    def __init__(
+        self,
+        *,
+        model: Model,
+        executor: Executor | None = None,
+        batch_size: int = -1,
+        sleep_time: float = 0.01,
+        buffering_time: float = 0.001,
+    ):
+        super().__init__(batch_size=batch_size, sleep_time=sleep_time, buffering_time=buffering_time)
+        self.model = model
+        self.executor = executor
+
+    async def process_batch(self, batch):
+        result = await asyncio.get_event_loop().run_in_executor(
+            self.executor, lambda _batch: self.model.predict(_batch, batch_size=len(_batch)), batch
+        )
+        return result

async_batcher/ml/sklearn/model.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING
+
+from async_batcher.batcher import AsyncBatcher
+
+if TYPE_CHECKING:
+    from concurrent.futures import Executor
+
+    from sklearn.base import BaseEstimator
+
+
+class SklearnAsyncBatcher(AsyncBatcher):
+    """Batcher for scikit-learn models.
+
+    Args:
+        model: The scikit-learn model to use for prediction.
+        executor: The executor to use for running the prediction.
+        batch_size (int, optional): The max number of items to process in a batch. Defaults to -1 (no limit).
+        sleep_time (float, optional): The time to sleep between checking if the result is ready in seconds.
+            Defaults to 0.01. Set it to a value close to the expected time to process a batch
+        buffering_time (float, optional): The time to sleep after processing a batch or checking the buffer
+            in seconds. Defaults to 0.001.
+            You can increase this value if you don't need a low latency, but want to reduce the number of
+            processed batches.
+    """
+
+    def __init__(
+        self,
+        *,
+        model: BaseEstimator,
+        executor: Executor | None = None,
+        batch_size: int = -1,
+        sleep_time: float = 0.01,
+        buffering_time: float = 0.001,
+    ):
+        super().__init__(batch_size=batch_size, sleep_time=sleep_time, buffering_time=buffering_time)
+        self.model = model
+        self.executor = executor
+
+    async def process_batch(self, batch):
+        if hasattr(self.model, "predict"):
+            result = await asyncio.get_event_loop().run_in_executor(self.executor, self.model.predict, batch)
+            return result
+        else:
+            raise AttributeError("Model does not have a predict method")

async_batcher/scylladb/update.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal
+
+from cassandra.cqlengine.query import BatchQuery
+
+from async_batcher.batcher import AsyncBatcher
+
+if TYPE_CHECKING:
+    from cassandra.cqlengine.models import Model
+
+
+@dataclass(kw_only=True)
+class WriteOperation:
+    operation: Literal["INSERT", "UPDATE", "DELETE"]
+    model: Model
+    key: dict[str, Any] | None = None
+    data: dict[str, Any] | None = None
+
+
+class AsyncScyllaDbWriteBatcher(AsyncBatcher[WriteOperation, None]):
+    async def process_batch(self, *, batch: list[WriteOperation]) -> list[None | Exception]:
+        results = []
+        with BatchQuery() as b:
+            for op in batch:
+                if op.operation == "INSERT":
+                    if op.data is None:
+                        results.append(ValueError("data must be provided for INSERT operations"))
+                    else:
+                        op.model.batch(b).create(**op.data)
+                        results.append(None)
+                elif op.operation == "UPDATE":
+                    if op.key is None or op.data is None:
+                        results.append(ValueError("key and data must be provided for UPDATE operations"))
+                    else:
+                        op.model.objects(**op.key).batch(b).update(**op.data)
+                        results.append(None)
+                elif op.operation == "DELETE":
+                    if op.key is None:
+                        results.append(ValueError("key must be provided for DELETE operations"))
+                    else:
+                        op.model.objects(**op.key).batch(b).delete()
+                        results.append(None)
+        return results

async_batcher-0.1.0.dist-info/METADATA

@@ -0,0 +1,99 @@
+Metadata-Version: 2.1
+Name: async-batcher
+Version: 0.1.0
+Summary: A service to batch your http requests.
+Home-page: https://github.com/hussein-awala/async-batcher
+License: Apache-2.0
+Keywords: python,asyncio,rest,grpc,fastapi
+Author: Hussein Awala
+Author-email: hussein@awala.fr
+Requires-Python: >=3.8,<4.0
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Provides-Extra: aws
+Provides-Extra: keras
+Provides-Extra: scylla
+Provides-Extra: sklearn
+Requires-Dist: aioboto3 (>=12.3,<13.0) ; extra == "aws"
+Requires-Dist: keras (>=2.12,<3.0) ; extra == "keras"
+Requires-Dist: scikit-learn (>=1.2,<2.0) ; extra == "sklearn"
+Requires-Dist: scylla-driver (>=3,<4) ; extra == "scylla"
+Project-URL: Bug Tracker, https://github.com/hussein-awala/async-batcher/issues
+Project-URL: Documentation, https://github.com/hussein-awala/async-batcher/blob/main/README.md
+Project-URL: Repository, https://github.com/hussein-awala/async-batcher
+Description-Content-Type: text/markdown
+
+# Async Batcher
+
+This project provides a Python library to batch the asynchronous requests and handle them in batches.
+
+## How to use
+
+To use the library, you need to install the package in your environment. You can install the package using pip:
+
+```bash
+pip install async-batcher
+```
+
+Then, you can create a new `AsyncBatcher` class by implementing the `process_batch` method:
+
+```python
+from async_batcher.batcher import AsyncBatcher
+
+class MyAsyncBatcher(AsyncBatcher):
+    async def process_batch(self, batch):
+        # Process the batch
+        print(batch)
+
+# Create a new instance of the `MyAsyncBatcher` class
+async_batcher = MyAsyncBatcher(batch_size=20)
+async_batcher.start()
+```
+
+## Benchmark
+
+To evaluate the performance of the `AsyncBatcher` library, we used the [Keras example](examples/keras)
+with [locust](https://locust.io/) to simulate multiple users making requests to the server.
+
+In this example, we have a tensorflow model that takes ~11ms to make a single prediction and ~13ms to process
+a batch of 200 predictions.
+
+For this benchmark, we ran the FastAPI server with Uvicorn with 12 workers,
+on a MacBook Pro with Apple M2 Max chip (12 cores).
+
+Predict endpoint:
+
+![predict__total_requests_per_second.png](assets%2Fpredict__total_requests_per_second.png)
+
+Optimized Predict endpoint:
+
+![optimized_predict__total_requests_per_second.png](assets%2Foptimized_predict__total_requests_per_second.png)
+
+As we can see from the graphs, for the `/predict` endpoint, the failure rate started to increase linearly at the RPS
+of ~284 requests/second, and the 95th percentile response time was ~1100ms at this point. Then with the increase of
+the RPS, the number of successful requests was ~130 requests/second in average, with a high failure rate
+(~300 from ~400 requests/second).
+
+For the `/optimized_predict` endpoint, the failure was smaller than 3 requests/second (0.00625%) during the whole test,
+the average response time was increasing slightly with the increase of the RPS, and it reached ~120ms in average at
+the end of the test (> 480 requests/second), with a 95th percentile response time almost stable and smaller than 500ms.
+
+You can find check the reports for more details:
+- [predict__report.html](assets/predict__report.html)
+- [optimized_predict__report.html](assets/optimized_predict__report.html)
+## Use cases
+
+The `AsyncBatcher` library can be used in any application that needs to handle asynchronous requests in batches,
+such as:
+- Serving machine learning models that optimize the batch processing (e.g. TensorFlow, PyTorch, Scikit-learn, etc.)
+- Storing multiple records in a database in a single query to optimize the I/O operations (or to reduce the cost of the
+  database operations, e.g. AWS DynamoDB)
+- Sending multiple messages in a single request to optimize the network operations (or to reduce the cost of the network
+  operations, e.g. Kafka, RabbitMQ, AWS SQS, AWS SNS, etc.)
+
+

async_batcher-0.1.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+async_batcher/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+async_batcher/aws/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+async_batcher/aws/dynamodb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+async_batcher/aws/dynamodb/get.py,sha256=fkYjx4RhhiQ0uXF6Me5nTGZ9MSNYmaAG8rZ9_B681mI,4063
+async_batcher/aws/dynamodb/write.py,sha256=bWuSaF435nEhix9TeaQFNSBwqAhtY-8yE1lZYDtRXa0,3679
+async_batcher/batcher.py,sha256=M1CxKnoRW-0zadBFaVdclRTup9jD3Jw26O033cuUz_o,5683
+async_batcher/ml/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+async_batcher/ml/keras/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+async_batcher/ml/keras/model.py,sha256=x9t9YFcbR6BM7SjQyv1mh5Nikg9SigWF9wkdODLbNlU,1657
+async_batcher/ml/sklearn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+async_batcher/ml/sklearn/model.py,sha256=hh_xtR0o24V5Yk8OH7tQONR0C-FD4OAv8tPRKT6rctg,1765
+async_batcher/scylladb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+async_batcher/scylladb/update.py,sha256=a3X8Bea8n_bwJYw-SrZ-dMgSy2noQnBMNkXXM3AOs7Q,1745
+async_batcher-0.1.0.dist-info/METADATA,sha256=TA1GlST7JFBf_62ih3FURPXp7tJiEvPEzaRDYamFDMs,4205
+async_batcher-0.1.0.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+async_batcher-0.1.0.dist-info/RECORD,,

async_batcher-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 1.8.1
+Root-Is-Purelib: true
+Tag: py3-none-any