feat: timeout

Author: maciejmajek

src/rai_core/rai/__init__.py

@@ -20,6 +20,7 @@
     get_llm_model_direct,
     get_tracing_callbacks,
 )
+from .utils import timeout
 
 __all__ = [
     "AgentRunner",
@@ -29,4 +30,5 @@
     "get_llm_model_config_and_vendor",
     "get_llm_model_direct",
     "get_tracing_callbacks",
+    "timeout",
 ]

src/rai_core/rai/tools/__init__.py

@@ -11,3 +11,5 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from .timeout import timeout, timeout_method

src/rai_core/rai/tools/ros2/detection/tools.py

@@ -16,6 +16,7 @@
 
 from pydantic import BaseModel, Field
 
+from rai.tools import timeout
 from rai.tools.ros2.base import BaseROS2Tool
 from rai.tools.ros2.detection.pcl import (
     GrippingPointEstimator,
@@ -40,24 +41,44 @@ class GetGrippingPointTool(BaseROS2Tool):
     gripping_point_estimator: GrippingPointEstimator
     point_cloud_filter: PointCloudFilter
 
+    timeout_sec: float = Field(
+        default=10.0, description="Timeout in seconds to get the gripping point"
+    )
+
     args_schema: Type[GetGrippingPointToolInput] = GetGrippingPointToolInput
 
     def _run(self, object_name: str) -> str:
-        pcl = self.point_cloud_from_segmentation.run(object_name)
-        pcl = self.point_cloud_filter.run(pcl)
-        gps = self.gripping_point_estimator.run(pcl)
-
-        message = ""
-        if len(gps) == 0:
-            message += f"No gripping point found for the object {object_name}\n"
-        elif len(gps) == 1:
-            message += f"The gripping point of the object {object_name} is {gps[0]}\n"
-        else:
-            message += f"Multiple gripping points found for the object {object_name}\n"
-
-        for i, gp in enumerate(gps):
-            message += (
-                f"The gripping point of the object {i + 1} {object_name} is {gp}\n"
-            )
-
-        return message
+        @timeout(
+            self.timeout_sec,
+            f"Gripping point detection for object '{object_name}' exceeded {self.timeout_sec} seconds",
+        )
+        def _run_with_timeout():
+            pcl = self.point_cloud_from_segmentation.run(object_name)
+            pcl = self.point_cloud_filter.run(pcl)
+            gps = self.gripping_point_estimator.run(pcl)
+
+            message = ""
+            if len(gps) == 0:
+                message += f"No gripping point found for the object {object_name}\n"
+            elif len(gps) == 1:
+                message += (
+                    f"The gripping point of the object {object_name} is {gps[0]}\n"
+                )
+            else:
+                message += (
+                    f"Multiple gripping points found for the object {object_name}\n"
+                )
+
+            for i, gp in enumerate(gps):
+                message += (
+                    f"The gripping point of the object {i + 1} {object_name} is {gp}\n"
+                )
+
+            return message
+
+        try:
+            return _run_with_timeout()
+        except Exception as e:
+            if "timed out" in str(e).lower():
+                return f"Timeout: Gripping point detection for object '{object_name}' exceeded {self.timeout_sec} seconds"
+            raise

src/rai_core/rai/tools/timeout.py

@@ -0,0 +1,139 @@
+# Copyright (C) 2025 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import signal
+from functools import wraps
+from typing import Any, Callable, TypeVar
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+class TimeoutError(Exception):
+    """Raised when an operation times out."""
+
+    pass
+
+
+def timeout(seconds: float, timeout_message: str = None) -> Callable[[F], F]:
+    """
+    Decorator that adds timeout functionality to a function.
+
+    Parameters
+    ----------
+    seconds : float
+        Timeout duration in seconds
+    timeout_message : str, optional
+        Custom timeout message. If not provided, a default message will be used.
+
+    Returns
+    -------
+    Callable
+        Decorated function with timeout functionality
+
+    Raises
+    ------
+    TimeoutError
+        When the decorated function exceeds the specified timeout
+
+    Examples
+    --------
+    >>> @timeout(5.0, "Operation timed out")
+    ... def slow_operation():
+    ...     import time
+    ...     time.sleep(10)
+    ...     return "Done"
+    >>>
+    >>> try:
+    ...     result = slow_operation()
+    ... except TimeoutError as e:
+    ...     print(f"Timeout: {e}")
+    """
+
+    def decorator(func: F) -> F:
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            def timeout_handler(signum, frame):
+                message = (
+                    timeout_message
+                    or f"Function '{func.__name__}' timed out after {seconds} seconds"
+                )
+                raise TimeoutError(message)
+
+            # Set up timeout
+            old_handler = signal.signal(signal.SIGALRM, timeout_handler)
+            signal.alarm(int(seconds))
+
+            try:
+                return func(*args, **kwargs)
+            finally:
+                # Clean up timeout
+                signal.alarm(0)
+                signal.signal(signal.SIGALRM, old_handler)
+
+        return wrapper
+
+    return decorator
+
+
+def timeout_method(seconds: float, timeout_message: str = None) -> Callable[[F], F]:
+    """
+    Decorator that adds timeout functionality to a method.
+    Similar to timeout but designed for class methods.
+
+    Parameters
+    ----------
+    seconds : float
+        Timeout duration in seconds
+    timeout_message : str, optional
+        Custom timeout message. If not provided, a default message will be used.
+
+    Returns
+    -------
+    Callable
+        Decorated method with timeout functionality
+
+    Examples
+    --------
+    >>> class MyClass:
+    ...     @timeout_method(3.0, "Method timed out")
+    ...     def slow_method(self):
+    ...         import time
+    ...         time.sleep(5)
+    ...         return "Done"
+    """
+
+    def decorator(func: F) -> F:
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            def timeout_handler(signum, frame):
+                message = (
+                    timeout_message
+                    or f"Method '{func.__name__}' of {self.__class__.__name__} timed out after {seconds} seconds"
+                )
+                raise TimeoutError(message)
+
+            # Set up timeout
+            old_handler = signal.signal(signal.SIGALRM, timeout_handler)
+            signal.alarm(int(seconds))
+
+            try:
+                return func(self, *args, **kwargs)
+            finally:
+                # Clean up timeout
+                signal.alarm(0)
+                signal.signal(signal.SIGALRM, old_handler)
+
+        return wrapper
+
+    return decorator