From 943281feb53f716fa4c805b8cc7f580ba1969fdb Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sat, 25 Feb 2023 23:55:30 -0600
Subject: [PATCH] wire up worker jobs

---
 api/onnx_web/worker/context.py | 3 ++-
 api/onnx_web/worker/worker.py  | 8 ++++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/api/onnx_web/worker/context.py b/api/onnx_web/worker/context.py
index bf927f03..8dfb7715 100644
--- a/api/onnx_web/worker/context.py
+++ b/api/onnx_web/worker/context.py
@@ -1,6 +1,6 @@
 from logging import getLogger
 from torch.multiprocessing import Queue, Value
-from typing import Any, Callable
+from typing import Any, Callable, Tuple
 
 from ..params import DeviceParams
 
@@ -12,6 +12,7 @@ ProgressCallback = Callable[[int, int, Any], None]
 class WorkerContext:
     cancel: "Value[bool]" = None
     key: str = None
+    pending: "Queue[Tuple[Callable, Any, Any]]" = None
     progress: "Value[int]" = None
 
     def __init__(
diff --git a/api/onnx_web/worker/worker.py b/api/onnx_web/worker/worker.py
index 8f6adccb..cf47f85e 100644
--- a/api/onnx_web/worker/worker.py
+++ b/api/onnx_web/worker/worker.py
@@ -1,6 +1,7 @@
 from logging import getLogger
 from torch.multiprocessing import Lock
 from time import sleep
+from traceback import print_exception
 
 from .context import WorkerContext
 
@@ -30,3 +31,10 @@ def worker_init(lock: Lock, context: WorkerContext):
         else:
             job = context.pending.get()
             logger.info("got job: %s", job)
+            try:
+                fn, args, kwargs = job
+                fn(context, *args, **kwargs)
+                logger.info("finished job")
+            except Exception as e:
+                print_exception(type(e), e, e.__traceback__)
+