From 35dc8a0bc4d64bf94e381d0a11be3ffdea3544a7 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sun, 5 Mar 2023 21:37:39 -0600
Subject: [PATCH] improve exit logging

---
 api/onnx_web/worker/context.py |  7 ++++---
 api/onnx_web/worker/worker.py  | 17 ++++++++++++-----
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/api/onnx_web/worker/context.py b/api/onnx_web/worker/context.py
index 411145a2..5d2013a2 100644
--- a/api/onnx_web/worker/context.py
+++ b/api/onnx_web/worker/context.py
@@ -43,10 +43,11 @@ class WorkerContext:
         return self.cancel.value
 
     def is_current(self) -> bool:
-        if self.current.value > 0:
-            return self.current.value == getpid()
+        return self.get_current() == getpid()
 
-        return True
+    def get_current(self) -> int:
+        with self.current.get_lock():
+            return self.current.value
 
     def get_device(self) -> DeviceParams:
         """
diff --git a/api/onnx_web/worker/worker.py b/api/onnx_web/worker/worker.py
index 640c6fe0..69d9a120 100644
--- a/api/onnx_web/worker/worker.py
+++ b/api/onnx_web/worker/worker.py
@@ -1,5 +1,6 @@
 from logging import getLogger
 from queue import Empty
+from os import getpid
 from sys import exit
 from traceback import format_exception
 
@@ -11,6 +12,12 @@ from .context import WorkerContext
 
 logger = getLogger(__name__)
 
+EXIT_ERROR = 1
+EXIT_INTERRUPT = 0
+EXIT_MEMORY = 2
+EXIT_REPLACED = 3
+EXIT_SUCCESS = 0
+
 
 def worker_main(context: WorkerContext, server: ServerContext):
     apply_patches(server)
@@ -25,8 +32,8 @@ def worker_main(context: WorkerContext, server: ServerContext):
     while True:
         try:
             if not context.is_current():
-                logger.warning("worker has been replaced, exiting")
-                exit(3)
+                logger.warning("worker %s has been replaced by %s, exiting", getpid(), context.get_current())
+                exit(EXIT_REPLACED)
 
             name, fn, args, kwargs = context.pending.get(timeout=1.0)
             logger.info("worker for %s got job: %s", context.device.device, name)
@@ -41,14 +48,14 @@ def worker_main(context: WorkerContext, server: ServerContext):
             pass
         except KeyboardInterrupt:
             logger.info("worker got keyboard interrupt")
-            exit(0)
+            exit(EXIT_INTERRUPT)
         except ValueError as e:
             logger.info("value error in worker, exiting: %s", e)
-            exit(1)
+            exit(EXIT_ERROR)
         except Exception as e:
             if "Failed to allocate memory" in str(e):
                 logger.error("detected out-of-memory error, exiting: %s", e)
-                exit(2)
+                exit(EXIT_MEMORY)
             else:
                 logger.error(
                     "error while running job: %s",