Update abstra-lib

abstra_internals/controllers/execution_consumer.py

@@ -17,6 +17,10 @@ class NonCleanExit(Exception):
     pass
 
 
+# should be smaller than the RabbitMQ consumer timeout (24h now)
+PROCESS_TIMEOUT = 6 * 60 * 60
+
+
 def PreExecController(
     *,
     controller: MainController,
@@ -49,16 +53,19 @@ def PreExecController(
         )
 
         p.start()
-        p.join()
+        p.join(timeout=PROCESS_TIMEOUT)
+        consumer.done_callback(msg)
+
+        if p.is_alive():
+            p.terminate()
+            raise NonCleanExit(f"Worker took too long to complete ({PROCESS_TIMEOUT})")
 
         if p.exitcode != 0:
             err_msg = f"Worker exited with status ({p.exitcode})"
             if p.exitcode == -9:
                 err_msg += ": Server reached its memory limit"
 
             raise NonCleanExit(err_msg)
-
-        consumer.done_callback(msg)
     except Exception as e:
         AbstraLogger.error(f"[{head_id}] PreExecController ERROR: aborting consumer")
         AbstraLogger.capture_exception(e)
@@ -89,11 +96,14 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 def ExecutionConsumer(consumer: Consumer, controller: MainController):
     with Arbiter(controller) as arbiter:
         with ThreadPoolExecutor(max_workers=QUEUE_CONCURRENCY) as executor:
-            for msg in consumer.iter():
-                executor.submit(
-                    PreExecController,
-                    arbiter_uuid=arbiter.uuid,
-                    controller=controller,
-                    consumer=consumer,
-                    msg=msg,
-                )
+            try:
+                for msg in consumer.iter():
+                    executor.submit(
+                        PreExecController,
+                        arbiter_uuid=arbiter.uuid,
+                        controller=controller,
+                        consumer=consumer,
+                        msg=msg,
+                    )
+            except Exception as e:
+                AbstraLogger.capture_exception(e)