diff --git a/pylabrobot/error_handling/__init__.py b/pylabrobot/error_handling/__init__.py new file mode 100644 index 0000000000..cbe02dbf65 --- /dev/null +++ b/pylabrobot/error_handling/__init__.py @@ -0,0 +1,2 @@ +from .handlers import choose_handler, serial_error_handler, until_success, basic_retry_handler +from .with_error_handler import with_error_handler diff --git a/pylabrobot/error_handling/handlers/__init__.py b/pylabrobot/error_handling/handlers/__init__.py new file mode 100644 index 0000000000..317d0ab021 --- /dev/null +++ b/pylabrobot/error_handling/handlers/__init__.py @@ -0,0 +1,4 @@ +from .choose_handler import choose_handler +from .serial_handler import serial_error_handler +from .until_success import until_success +from .retry import basic_retry_handler diff --git a/pylabrobot/error_handling/handlers/choose_handler.py b/pylabrobot/error_handling/handlers/choose_handler.py new file mode 100644 index 0000000000..d84afc623e --- /dev/null +++ b/pylabrobot/error_handling/handlers/choose_handler.py @@ -0,0 +1,12 @@ +from typing import Callable, Dict + + +def choose_handler(handlers: Dict[Exception, Callable]) -> Callable: + """Choose the appropriate error handler based on the type of error.""" + + async def handler(func, exception, **kwargs): + for exc_type, handler in handlers.items(): + if isinstance(exception, exc_type): + return await handler(func, exception, **kwargs) + + return handler diff --git a/pylabrobot/error_handling/handlers/retry.py b/pylabrobot/error_handling/handlers/retry.py new file mode 100644 index 0000000000..69adfb8525 --- /dev/null +++ b/pylabrobot/error_handling/handlers/retry.py @@ -0,0 +1,3 @@ +async def basic_retry_handler(func, error, **kwargs): + """Will simply retry the function call with the same arguments.""" + return await func(**kwargs) diff --git a/pylabrobot/error_handling/handlers/serial_handler.py b/pylabrobot/error_handling/handlers/serial_handler.py new file mode 100644 index 0000000000..4927e2d1cc --- /dev/null +++ b/pylabrobot/error_handling/handlers/serial_handler.py @@ -0,0 +1,11 @@ +class serial_error_handler: + def __init__(self, child_handlers: list): + self.child_handlers = child_handlers + self.index = 0 + + async def __call__(self, func, exception, **kwargs): + if self.index >= len(self.child_handlers): + raise RuntimeError("No more child handlers to call") + handler = self.child_handlers[self.index] + self.index += 1 + return await handler(func, exception, **kwargs) diff --git a/pylabrobot/error_handling/handlers/until_success.py b/pylabrobot/error_handling/handlers/until_success.py new file mode 100644 index 0000000000..3cca0eb192 --- /dev/null +++ b/pylabrobot/error_handling/handlers/until_success.py @@ -0,0 +1,24 @@ +from typing import Callable +from typing import Callable, Optional + + +class until_success: + """ + Error handler that retries the given handler until the main function does not raise + an exception, or until the maximum number of tries is reached. + + Args: + handler: The async function to be executed. + max_tries: Maximum number of retries. Default is None, which means infinite retries. + """ + + def __init__(self, handler: Callable, max_tries: Optional[int] = None): + self.handler = handler + self.max_tries = max_tries + self.attempts = 0 + + async def __call__(self, *args, **kwargs): + if self.max_tries is not None and self.attempts >= self.max_tries: + raise RuntimeError("Maximum number of retries reached") + self.attempts += 1 + return await self.handler(*args, **kwargs) diff --git a/pylabrobot/error_handling/with_error_handler.py b/pylabrobot/error_handling/with_error_handler.py new file mode 100644 index 0000000000..f78b7249ed --- /dev/null +++ b/pylabrobot/error_handling/with_error_handler.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +import functools +import inspect +from typing import Any, Awaitable, Callable, Optional, ParamSpec, TypeVar + +_P = ParamSpec("_P") +_R = TypeVar("_R", bound=Awaitable[Any]) + +Handler: Callable[[Callable[_P, _R], Exception, dict[str, Any]], Awaitable[Any]] + + +def with_error_handler(func: Callable[_P, _R]) -> Callable[_P, _R]: + @functools.wraps(func) + async def wrapper(self, *args, error_handler: Optional[Handler] = None, **kwargs): + try: + return await func(self, *args, **kwargs) + except Exception as error: + print("caught error", error) + if error_handler is not None: + bound = wrapper.__get__(self, type(self)) + + # convert all args to kwargs, remove self + sig = inspect.signature(func) + bound_args = sig.bind(self, *args, **kwargs) + bound_args = {k: v for k, v in bound_args.arguments.items() if k != "self"} + bound_args["error_handler"] = error_handler + + return await error_handler(bound, error, **bound_args) + raise + + return wrapper diff --git a/pylabrobot/liquid_handling/error_handlers.py b/pylabrobot/liquid_handling/error_handlers.py new file mode 100644 index 0000000000..92fecb42d6 --- /dev/null +++ b/pylabrobot/liquid_handling/error_handlers.py @@ -0,0 +1,24 @@ +from pylabrobot.liquid_handling.errors import ChannelizedError + + +def try_next_tip_spot(try_tip_spots): + async def handler(func, error: Exception, **kwargs): + assert isinstance(error, ChannelizedError) + + new_tip_spots, new_use_channels = [], [] + + tip_spots = kwargs.pop("tip_spots") + if "use_channels" not in kwargs: + use_channels = list(range(len(tip_spots))) + else: + use_channels = kwargs.pop("use_channels") + + for idx, channel_idx in zip(tip_spots, use_channels): + if channel_idx in error.errors.keys(): + new_tip_spots.append(next(try_tip_spots)) + new_use_channels.append(channel_idx) + + print(f"Retrying with tip spots: {new_tip_spots} and use channels: {new_use_channels}") + return await func(tip_spots=new_tip_spots, use_channels=new_use_channels, **kwargs) + + return handler diff --git a/pylabrobot/liquid_handling/liquid_handler.py b/pylabrobot/liquid_handling/liquid_handler.py index 0224803212..561ff2f6c9 100644 --- a/pylabrobot/liquid_handling/liquid_handler.py +++ b/pylabrobot/liquid_handling/liquid_handler.py @@ -23,6 +23,7 @@ cast, ) +from pylabrobot.error_handling import with_error_handler from pylabrobot.liquid_handling.errors import ChannelizedError from pylabrobot.liquid_handling.strictness import ( Strictness, @@ -333,7 +334,7 @@ def _make_sure_channels_exist(self, channels: List[int]): if not len(invalid_channels) == 0: raise ValueError(f"Invalid channels: {invalid_channels}") - @need_setup_finished + @with_error_handler async def pick_up_tips( self, tip_spots: List[TipSpot],