From 3a4928e59b72d8be126ed15ea765f31a180eef69 Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sun, 5 Mar 2023 20:07:06 -0600 Subject: [PATCH] fix(api): prevent workers from blocking on their progress queues --- api/onnx_web/worker/context.py | 2 +- api/onnx_web/worker/worker.py | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/api/onnx_web/worker/context.py b/api/onnx_web/worker/context.py index 83f2da1d..24d96c6a 100644 --- 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 index 94c57020..dc5d36c0 100644 --- 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)