diff --git a/pyroengine/engine.py b/pyroengine/engine.py index a0412a9..ff5fe16 100644 --- a/pyroengine/engine.py +++ b/pyroengine/engine.py @@ -9,6 +9,7 @@ import logging import os import shutil +import signal import time from collections import deque from datetime import datetime, timedelta, timezone @@ -30,6 +31,23 @@ logging.basicConfig(format="%(asctime)s | %(levelname)s: %(message)s", level=logging.INFO, force=True) +def handler(signum, frame): + raise TimeoutError("Heartbeat check timed out") + + +def heartbeat_with_timeout(api_instance, cam_id, timeout=1): + signal.signal(signal.SIGALRM, handler) + signal.alarm(timeout) + try: + api_instance.heartbeat(cam_id) + except TimeoutError: + logging.warning(f"Heartbeat check timed out for {cam_id}") + except ConnectionError: + logging.warning(f"Unable to reach the pyro-api with {cam_id}") + finally: + signal.alarm(0) + + class Engine: """This implements an object to manage predictions and API interactions for wildfire alerts. @@ -253,10 +271,7 @@ def predict(self, frame: Image.Image, cam_id: Optional[str] = None) -> float: # Heartbeat if len(self.api_client) > 0 and isinstance(cam_id, str): - try: - self.heartbeat(cam_id) - except ConnectionError: - logging.warning(f"Unable to reach the pyro-api with {cam_id}") + heartbeat_with_timeout(self, cam_id, timeout=1) cam_key = cam_id or "-1" # Reduce image size to save bandwidth