Add async ubatch

examples/fastapi.py

@@ -0,0 +1,44 @@
+import random
+import numpy as np
+from typing import Dict, List
+
+from fastapi import FastAPI  # type: ignore
+from joblib import load
+
+from sklearn.datasets import fetch_20newsgroups
+from sklearn.model_selection import train_test_split
+
+from microbatch import AsyncUBatch
+
+
+ngd = fetch_20newsgroups(subset="all")
+
+X = ngd.data
+y = ngd.target
+_, X_test, _, _ = train_test_split(X, y, test_size=0.33)
+
+
+model = load("xgbregressor.joblib")
+
+
+def predict(data: List[np.array]) -> List[np.float32]:
+    return model.predict(np.array(data))  # type: ignore
+
+
+preidct_ubatch = AsyncUBatch[List[np.array], np.float32](max_size=100, timeout=0.01)
+preidct_ubatch.set_handler(predict)
+preidct_ubatch.start()
+
+app = FastAPI()
+
+
+@app.post("/predict_ubatch")  # type: ignore
+async def predict_ubatch_post() -> Dict[str, float]:
+    output = await preidct_ubatch.ubatch(random.choice(X_test))
+    return {"prediction": float(output)}
+
+
+@app.post("/predict")  # type: ignore
+def predict_post() -> Dict[str, float]:
+    output = predict([random.choice(X_test)])[0]
+    return {"prediction": float(output)}

examples/flask_app.py

@@ -6,11 +6,8 @@
 from flask import Flask
 from flask_restx import Resource, Api
 
-# from numpy import genfromtxt
-
 from ubatch import ubatch_decorator
 
-# from keras.models import load_model
 from sklearn.datasets import fetch_20newsgroups
 from sklearn.model_selection import train_test_split
 
@@ -25,7 +22,6 @@
 
 
 model = load("xgbregressor.joblib")
-# X_test = genfromtxt("xgbregressor_inputs.csv", delimiter=",")
 
 app = Flask(__name__)
 api = Api(app)
@@ -37,14 +33,14 @@ def predict(data: List[np.array]) -> List[np.float32]:
 
 
 @api.route("/predict_ubatch")
-class BatchPredict(Resource):
+class BatchPredict(Resource):  # type: ignore
     def post(self) -> Dict[str, float]:
-        output = predict.ubatch(random.choice(X_test))
+        output: np.array = predict.ubatch(random.choice(X_test))
         return {"prediction": float(output)}
 
 
 @api.route("/predict")
-class Predict(Resource):
+class Predict(Resource):  # type: ignore
     def post(self) -> Dict[str, float]:
-        output = predict([random.choice(X_test)])[0]
+        output: np.array = predict([random.choice(X_test)])[0]
         return {"prediction": float(output)}

pyproject.toml

@@ -30,6 +30,8 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = ">= 3.6, < 3.9"
 flask-restx = { version = "^0.2", optional = true }
+fastapi = { version = "^0.61", optional = true }
+uvicorn = { version = "^0.12", optional = true }
 scikit-learn = { version = "^0.23", optional = true }
 gunicorn = { version = "^20.0", optional = true }
 xgboost = { version = "^1.2", optional = true }
@@ -50,7 +52,7 @@ pytest-timeout = "^1.4"
 pytest-reraise = "^1.0"
 
 [tool.poetry.extras]
-benchmark = ["flask-restx", "scikit-learn", "gunicorn", "xgboost", "keras", "tensorflow"]
+benchmark = ["flask-restx", "scikit-learn", "gunicorn", "xgboost", "keras", "tensorflow", "fastapi", "uvicorn"]
 
 [tool.black]
 line-length = 88

ubatch/__init__.py

@@ -1,7 +1,10 @@
+from ubatch.async_ubatch import AsyncUBatch
 from ubatch.decorators import ubatch_decorator
-from ubatch.ubatch import BadBatchOutputSize, HandlerAlreadySet, HandlerNotSet, UBatch
+from ubatch.exceptions import BadBatchOutputSize, HandlerAlreadySet, HandlerNotSet
+from ubatch.ubatch import UBatch
 
 __all__ = [
+    "AsyncUBatch",
     "UBatch",
     "HandlerNotSet",
     "HandlerAlreadySet",