Bind thread with model

app.py

@@ -1,11 +1,13 @@
-import asyncio
+import atexit
 import functools
-import uuid
+from queue import Queue
+from threading import Thread
 
 from paddleocr import PaddleOCR, draw_ocr
 from PIL import Image
 import gradio as gr
 
+
 LANG_CONFIG = {
     "ch": {"num_workers": 4},
     "en": {"num_workers": 4},
@@ -17,95 +19,53 @@ LANG_CONFIG = {
 CONCURRENCY_LIMIT = 8
 
 
-class PaddleOCRModelWrapper(object):
-    def __init__(self, model, name=None):
-        super().__init__()
-        self._model = model
-        self._name = name or self._get_random_name()
-        self._state = "IDLE"
-
-    @property
-    def name(self):
-        return self._name
-
-    @property
-    def state(self):
-        return self._state
-
-    @state.setter
-    def state(self, state):
-        self._state = state
-
-    def infer(self, **kwargs):
-        img_path = kwargs["img"]
-        result = self._model.ocr(**kwargs)[0]
-        image = Image.open(img_path).convert("RGB")
-        boxes = [line[0] for line in result]
-        txts = [line[1][0] for line in result]
-        scores = [line[1][1] for line in result]
-        im_show = draw_ocr(image, boxes, txts, scores,
-                        font_path="./simfang.ttf")
-        return im_show
-
-    def _get_random_name(self):
-        return str(uuid.uuid4())
-
-
 class PaddleOCRModelManager(object):
     def __init__(self,
-                 num_models,
-                 model_factory,
-                 *,
-                 polling_interval=0.1):
+                 num_workers,
+                 model_factory):
         super().__init__()
-        self._num_models = num_models
         self._model_factory = model_factory
-        self._polling_interval = polling_interval
-        self._models = {}
-        self.new_models()
-
-    def new_models(self):
-        self._models.clear()
-        for _ in range(self._num_models):
-            model = self._new_model()
-            self._models[model.name] = model
-
-    async def infer(self, **kwargs):
+        self._queue = Queue()
+        self._workers = []
+        for _ in range(num_workers):
+            worker = Thread(target=self._worker, daemon=False)
+            worker.start()
+            self._workers.append(worker)
+
+    def infer(self, *args, **kwargs):
+        # XXX: Should I use a more lightweight data structure, say, a future?
+        result_queue = Queue(maxsize=1)
+        self._queue.put((args, kwargs, result_queue))
+        success, payload = result_queue.get()
+        if success:
+            return payload
+        else:
+            raise payload
+
+    def close(self):
+        for _ in self._workers:
+            self._queue.put(None)
+        for worker in self._workers:
+            worker.join()
+
+    def _worker(self):
+        model = self._model_factory()
         while True:
-            model = self._get_available_model()
-            if not model:
-                await asyncio.sleep(self._polling_interval)
-                continue
-            model.state = "RUNNING"
-            # NOTE: I take an optimistic approach here, assuming that the model
-            # is not broken even if inference fails.
+            item = self._queue.get()
+            if item is None:
+                break
+            args, kwargs, result_queue = item
             try:
-                result = await self._new_inference_task(model, **kwargs)
+                result = model.ocr(*args, **kwargs)
+                result_queue.put((True, result))
+            except Exception as e:
+                result_queue.put((False, e))
             finally:
-                model.state = "IDLE"
-            return result
-
-    def _new_model(self):
-        real_model = self._model_factory()
-        model = PaddleOCRModelWrapper(real_model)
-        return model
-
-    def _get_available_model(self):
-        if not self._models:
-            raise RuntimeError("No living models")
-        for model in self._models.values():
-            if model.state == "IDLE":
-                return model
-        return None
-
-    def _new_inference_task(self, model,
-                            **kwargs):
-        return asyncio.get_running_loop().run_in_executor(
-            None, functools.partial(model.infer, **kwargs))
+                self._queue.task_done()
 
 
 def create_model(lang):
-    return PaddleOCR(lang=lang, use_angle_cls=True, use_gpu=False) 
+    return PaddleOCR(lang=lang, use_angle_cls=True, use_gpu=False)
 
 
 model_managers = {}
@@ -114,10 +74,26 @@ for lang, config in LANG_CONFIG.items():
     model_managers[lang] = model_manager
 
 
-async def inference(img, lang):
+def close_model_managers():
+    for manager in model_managers.values():
+        manager.close()
+
+
+# XXX: Not sure if gradio allows adding custom teardown logic
+atexit.register(close_model_managers)
+
+
+def inference(img, lang):
     ocr = model_managers[lang]
-    result = await ocr.infer(img=img, cls=True)
-    return result
+    result = ocr.infer(img, cls=True)[0]
+    img_path = img
+    image = Image.open(img_path).convert("RGB")
+    boxes = [line[0] for line in result]
+    txts = [line[1][0] for line in result]
+    scores = [line[1][1] for line in result]
+    im_show = draw_ocr(image, boxes, txts, scores,
+                    font_path="./simfang.ttf")
+    return im_show
 
 
 title = 'PaddleOCR'