From caac53f135a58f852ef9e079c4ee7432d31b7ea5 Mon Sep 17 00:00:00 2001
From: Jakob Langdal <jakob.langdal@alexandra.dk>
Date: Tue, 23 Apr 2024 08:41:32 +0200
Subject: [PATCH] Control worker timeout

---
 README.md                 | 20 +++++++++++++-------
 optimizerapi/optimizer.py |  7 ++++++-
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 682e535..c2dec66 100644
--- a/README.md
+++ b/README.md
@@ -62,6 +62,13 @@ worker threads using the following commands:
 
 The Redis server can be controlled through the environment variable `REDIS_URL` which defaults to `redis://localhost:6379`
 
+Time to live and timeout of the workers can be controlled with the following environment variables
+
+| Name           | Description                                 |
+| -------------- | ------------------------------------------- |
+| REDIS_TTL      | Time to keep results in redis (default=500) |
+| WORKER_TIMEOUT | Timeout in seconds (default=180)            |
+
 # Use [CORS](https://flask-cors.readthedocs.io/en/latest/index.html)
 
 The API server supports exposing its functionality to other origins than its own.
@@ -92,13 +99,12 @@ The static API key is configured by the environment variable `AUTH_API_KEY`
 
 Keycloak is configured using the following environement variables
 
-
-|Name               |Description                        |
-|-------------------|-----------------------------------|
-|AUTH_SERVER        |Base url of your Keycloak server   |
-|AUTH_REALM_NAME    |OAuth realm name                   |
-|AUTH_CLIENT_ID     |Client ID                          |
-|AUTH_CLIENT_SECRET |Client secret                      |
+| Name               | Description                      |
+| ------------------ | -------------------------------- |
+| AUTH_SERVER        | Base url of your Keycloak server |
+| AUTH_REALM_NAME    | OAuth realm name                 |
+| AUTH_CLIENT_ID     | Client ID                        |
+| AUTH_CLIENT_SECRET | Client secret                    |
 
 # Adding or updating dependencies
 
diff --git a/optimizerapi/optimizer.py b/optimizerapi/optimizer.py
index 633778e..ef4090e 100644
--- a/optimizerapi/optimizer.py
+++ b/optimizerapi/optimizer.py
@@ -43,6 +43,10 @@
     TTL = int(os.environ["REDIS_TTL"])
 else:
     TTL = 500
+if "WORKER_TIMEOUT" in os.environ:
+    WORKER_TIMEOUT = int(os.environ["WORKER_TIMEOUT"])
+else:
+    WORKER_TIMEOUT = 180
 
 plt.switch_backend("Agg")
 
@@ -77,7 +81,7 @@ def disconnect_check():
             print("Found existing job")
         except NoSuchJobError:
             print("Creating new job")
-            job = queue.enqueue(do_run_work, body, job_id=job_id, result_ttl=TTL)
+            job = queue.enqueue(do_run_work, body, job_id=job_id, result_ttl=TTL, timeout=WORKER_TIMEOUT)
         while job.return_value is None:
             if disconnect_check():
                 try:
@@ -102,6 +106,7 @@ def do_run_work(body) -> dict:
     except TypeError as err:
         return ({"message": "Type error", "error": str(err)}, 400)
     except ValueError as err:
+        traceback.print_exc()
         return ({"message": "Validation error", "error": str(err)}, 400)
     except Exception as err:
         # Log unknown exceptions to support debugging