fix(api): prevent workers from blocking on their progress queues

ssube · Mar 6, 2023 · 3a4928e · 3a4928e
1 parent edc55ae
commit 3a4928e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/api/onnx_web/worker/context.py b/api/onnx_web/worker/context.py
@@ -15,7 +15,7 @@ class WorkerContext:
     cancel: "Value[bool]"
     job: str
     pending: "Queue[Tuple[str, Callable[..., None], Any, Any]]"
-    progress: "Value[int]"
+    progress: "Queue[Tuple[str, str, int]]"
 
     def __init__(
         self,

diff --git a/api/onnx_web/worker/worker.py b/api/onnx_web/worker/worker.py
@@ -18,6 +18,10 @@ def worker_main(context: WorkerContext, server: ServerContext):
 
     logger.info("checking in from worker, %s", get_available_providers())
 
+    # make leaking workers easier to recycle
+    context.progress.cancel_join_thread()
+    context.finished.cancel_join_thread()
+
     while True:
         try:
             name, fn, args, kwargs = context.pending.get(timeout=1.0)