Merge pull request #1792 from roboflow/fix/webrtc-worker-do-not-return-in-finally

When encapsulating webrtc worker within exception block, do not return in finally

Author: grzegorz-roboflow

inference/core/interfaces/webrtc_worker/modal.py

@@ -204,16 +204,12 @@ async def run_rtc_peer_connection_with_watchdog(
         webrtc_request: WebRTCWorkerRequest,
         send_answer: Callable[[WebRTCWorkerResult], None],
         model_manager: ModelManager,
-    ) -> bool:
+        watchdog: Watchdog,
+    ):
         from inference.core.interfaces.webrtc_worker.webrtc import (
             init_rtc_peer_connection_with_loop,
         )
 
-        watchdog = Watchdog(
-            api_key=webrtc_request.api_key,
-            timeout_seconds=WEBRTC_MODAL_WATCHDOG_TIMEMOUT,
-        )
-
         rtc_peer_connection_task = asyncio.create_task(
             init_rtc_peer_connection_with_loop(
                 webrtc_request=webrtc_request,
@@ -238,12 +234,11 @@ def on_timeout(message: Optional[str] = ""):
         except modal.exception.InputCancellation:
             logger.warning("Modal function was cancelled")
         except asyncio.CancelledError as exc:
-            logger.info("WebRTC connection task was cancelled (%s)", exc)
+            logger.warning("WebRTC connection task was cancelled (%s)", exc)
         except Exception as exc:
             logger.error(exc)
         finally:
             watchdog.stop()
-            return watchdog.heartbeat_occurred
 
     class RTCPeerConnectionModal:
         _model_manager: Optional[ModelManager] = modal.parameter(
@@ -338,6 +333,10 @@ def rtc_peer_connection_modal(
 
             def send_answer(obj: WebRTCWorkerResult):
                 logger.info("Sending webrtc answer")
+                if obj.error_message:
+                    logger.error(
+                        "Error: %s (%s)", obj.error_message, obj.exception_type
+                    )
                 # Queue with no limit, below will never block
                 q.put(obj)
 
@@ -356,23 +355,35 @@ def send_answer(obj: WebRTCWorkerResult):
                 send_answer(WebRTCWorkerResult(error_message=error_msg))
                 return
 
+            watchdog = Watchdog(
+                api_key=webrtc_request.api_key,
+                timeout_seconds=WEBRTC_MODAL_WATCHDOG_TIMEMOUT,
+            )
+
             try:
-                heartbeat_occurred = asyncio.run(
+                asyncio.run(
                     run_rtc_peer_connection_with_watchdog(
                         webrtc_request=webrtc_request,
                         send_answer=send_answer,
                         model_manager=self._model_manager,
+                        watchdog=watchdog,
                     )
                 )
-            except (asyncio.CancelledError, modal.exception.InputCancellation):
+            except modal.exception.InputCancellation:
                 logger.warning("Modal function was cancelled")
+            except asyncio.CancelledError as exc:
+                logger.warning("WebRTC connection task was cancelled (%s)", exc)
+            except Exception as exc:
+                logger.warning("Unhandled exception: %s", exc)
+            finally:
+                watchdog.stop()
 
             _exec_session_stopped = datetime.datetime.now()
             logger.info(
                 "WebRTC session stopped at %s",
                 _exec_session_stopped.isoformat(),
             )
-            if not heartbeat_occurred:
+            if watchdog.total_heartbeats == 0:
                 raise Exception(
                     "WebRTC worker was terminated before processing a single frame"
                 )

inference/core/interfaces/webrtc_worker/watchdog.py

@@ -24,7 +24,11 @@ def __init__(
         self._last_log_ts = datetime.datetime.now()
         self._log_interval_seconds = 10
         self._heartbeats = 0
-        self.heartbeat_occurred = False
+        self._total_heartbeats = 0
+
+    @property
+    def total_heartbeats(self) -> int:
+        return self._total_heartbeats
 
     def start(self):
         logger.info("Starting watchdog with timeout %s", self.timeout_seconds)
@@ -44,25 +48,27 @@ def _watchdog_thread(self):
         while not self._stopping:
             if not self.is_alive():
                 logger.error(
-                    "Watchdog timeout reached, heartbeats: %s", self._heartbeats
+                    "Watchdog timeout reached, heartbeats: %s", self._total_heartbeats
                 )
                 self.on_timeout(
-                    message=f"Timeout reached, heartbeats: {self._heartbeats}"
+                    message=f"Timeout reached, heartbeats: {self._total_heartbeats}"
                 )
                 break
             if WEBRTC_MODAL_USAGE_QUOTA_ENABLED and is_over_quota(self._api_key):
-                logger.error("API key over quota, heartbeats: %s", self._heartbeats)
+                logger.error(
+                    "API key over quota, heartbeats: %s", self._total_heartbeats
+                )
                 self.on_timeout(
-                    message=f"API key over quota, heartbeats: {self._heartbeats}"
+                    message=f"API key over quota, heartbeats: {self._total_heartbeats}"
                 )
                 break
             time.sleep(1)
-        logger.info("Watchdog thread stopped, heartbeats: %s", self._heartbeats)
+        logger.info("Watchdog thread stopped, heartbeats: %s", self._total_heartbeats)
 
     def heartbeat(self):
         self.last_heartbeat = datetime.datetime.now()
         self._heartbeats += 1
-        self.heartbeat_occurred = True
+        self._total_heartbeats += 1
         if (
             datetime.datetime.now() - self._last_log_ts
         ).total_seconds() > self._log_interval_seconds:

inference/core/interfaces/webrtc_worker/webrtc.py

@@ -63,7 +63,7 @@
 from inference.core.workflows.core_steps.common.serializers import (
     serialize_wildcard_kind,
 )
-from inference.core.workflows.errors import WorkflowSyntaxError
+from inference.core.workflows.errors import WorkflowError, WorkflowSyntaxError
 from inference.core.workflows.execution_engine.entities.base import WorkflowImageData
 from inference.usage_tracking.collector import usage_collector
 
@@ -927,6 +927,8 @@ async def init_rtc_peer_connection_with_loop(
         KeyError,
         NotImplementedError,
     ) as error:
+        # heartbeat to indicate caller error
+        heartbeat_callback()
         send_answer(
             WebRTCWorkerResult(
                 exception_type=error.__class__.__name__,
@@ -935,6 +937,8 @@ async def init_rtc_peer_connection_with_loop(
         )
         return
     except WebRTCConfigurationError as error:
+        # heartbeat to indicate caller error
+        heartbeat_callback()
         send_answer(
             WebRTCWorkerResult(
                 exception_type=error.__class__.__name__,
@@ -943,6 +947,8 @@ async def init_rtc_peer_connection_with_loop(
         )
         return
     except RoboflowAPINotAuthorizedError:
+        # heartbeat to indicate caller error
+        heartbeat_callback()
         send_answer(
             WebRTCWorkerResult(
                 exception_type=RoboflowAPINotAuthorizedError.__name__,
@@ -951,6 +957,8 @@ async def init_rtc_peer_connection_with_loop(
         )
         return
     except RoboflowAPINotNotFoundError:
+        # heartbeat to indicate caller error
+        heartbeat_callback()
         send_answer(
             WebRTCWorkerResult(
                 exception_type=RoboflowAPINotNotFoundError.__name__,
@@ -959,6 +967,8 @@ async def init_rtc_peer_connection_with_loop(
         )
         return
     except WorkflowSyntaxError as error:
+        # heartbeat to indicate caller error
+        heartbeat_callback()
         send_answer(
             WebRTCWorkerResult(
                 exception_type=WorkflowSyntaxError.__name__,
@@ -968,6 +978,16 @@ async def init_rtc_peer_connection_with_loop(
             )
         )
         return
+    except WorkflowError as error:
+        # heartbeat to indicate caller error
+        heartbeat_callback()
+        send_answer(
+            WebRTCWorkerResult(
+                exception_type=WorkflowError.__name__,
+                error_message=str(error),
+            )
+        )
+        return
     except Exception as error:
         send_answer(
             WebRTCWorkerResult(