[feat] Add exception-free mode to Aim (#2561)

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Changelog
 
 ## 3.16.2
+- Add exception-free mode to Aim (alberttorosyan)
 - Expose `capture_terminal_logs` argument for `aim.sdk.adapters` classes (mihran113)
 - Handle inconsistency between Sequence data and metadata (alberttorosyan) 
 

aim/ext/exception_resistant.py

@@ -1,6 +1,11 @@
+import logging
+
 from functools import wraps
 
 
+logger = logging.getLogger(__name__)
+
+
 def exception_resistant(silent: bool):
     def inner(func):
         if not silent:
@@ -30,3 +35,38 @@ def wrapper(*args, **kwargs):
                     pass
         return wrapper
     return inner
+
+
+class _SafeModeConfig:
+    @staticmethod
+    def log_exception(e: Exception, func: callable):
+        logger.warning(f'Exception "{str(e)}" raised in function "{func.__name__}"')
+
+    @staticmethod
+    def reraise_exception(e: Exception, func: callable):
+        raise e
+
+    exception_callback = reraise_exception
+
+
+def enable_safe_mode():
+    _SafeModeConfig.exception_callback = _SafeModeConfig.log_exception
+
+
+def disable_safe_mode():
+    _SafeModeConfig.exception_callback = _SafeModeConfig.reraise_exception
+
+
+def set_exception_callback(callback: callable):
+    _SafeModeConfig.exception_callback = callback
+
+
+def noexcept(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            _SafeModeConfig.exception_callback(e, func)
+
+    return wrapper

aim/sdk/run.py

@@ -35,6 +35,7 @@
     get_installed_packages,
     get_git_info,
 )
+from aim.ext.exception_resistant import noexcept
 
 from typing import Any, Dict, Iterator, Optional, Tuple, Union
 from typing import TYPE_CHECKING
@@ -334,6 +335,7 @@ def __hash__(self) -> int:
     def idx_to_ctx(self, idx: int) -> Context:
         return self._tracker.idx_to_ctx(idx)
 
+    @noexcept
     def __setitem__(self, key: str, val: Any):
         """Set Run top-level meta-parameter.
 
@@ -364,6 +366,7 @@ def __getitem__(self, key):
         """
         return self._collect(key)
 
+    @noexcept
     def set(self, key, val: Any, strict: bool = True):
         self.meta_run_attrs_tree.set(key, val, strict)
         self.meta_attrs_tree.set(key, val, strict)
@@ -385,6 +388,7 @@ def __delitem__(self, key: str):
         del self.meta_attrs_tree[key]
         del self.meta_run_attrs_tree[key]
 
+    @noexcept
     def track(
         self,
         value,
@@ -812,6 +816,7 @@ class Run(BasicRun):
             git info, environment variables, etc.
     """
 
+    @noexcept
     def __init__(self, run_hash: Optional[str] = None, *,
                  repo: Optional[Union[str, 'Repo']] = None,
                  read_only: bool = False,

aim/sdk/training_flow.py

@@ -60,6 +60,21 @@ def training_successfully_finished(
     ):
         """Is called after the training phase is successfully finished."""
 
+    @event
+    def exception_raised(
+        self, *,
+        exception: Exception,
+        **kwargs
+    ):
+        """Is called when exception is raised from Aim codebase. """
+
+    def handle_exceptions(self):
+        from aim.ext.exception_resistant import set_exception_callback
+
+        def callback(e: Exception, func: callable):
+            self.exception_raised(exception=e, function=func)
+        set_exception_callback(callback)
+
     # @event
     # def run_available_experimental(
     #     self, *,

aim/utils/__init__.py

@@ -0,0 +1 @@
+from aim.ext.exception_resistant import enable_safe_mode, disable_safe_mode  # noqa

pyproject.toml

@@ -1,2 +1,2 @@
 [build-system]
-requires = ["setuptools", "cython == 3.0.0a11", "aimrocks == 0.3.1"]
+requires = ["setuptools", "cython == 3.0.0a11", "aimrocks == 0.4.0"]

setup.py

@@ -52,7 +52,7 @@ def package_files(directory):
 REQUIRED = [
     f'aim-ui=={__version__}',
     'aimrecords==0.0.7',
-    'aimrocks==0.3.1',
+    'aimrocks==0.4.0',
     'cachetools>=4.0.0',
     'click>=7.0',
     'cryptography>=3.0',