cherry-pick: PENG-2678 Fixed issue with handle_errors_async

jobbergate-agent/CHANGELOG.md

@@ -4,12 +4,16 @@ This file keeps track of all notable changes to jobbergate-agent
 
 ## Unreleased
 
+- Upgraded py-buzz dependency to ^5.0 to use builtin `handle_errors_async` [[PENG-2678](https://sharing.clickup.com/t/h/c/18022949/PENG-2678/M7PTQRUNIW5T2NK)]
 
 ## 5.4.3 -- 2025-01-16
+
 - Added custom settings for configuring Sentry's sample rates [[PENG-2592](https://sharing.clickup.com/t/h/c/18022949/PENG-2592/QQUQ1ABLAP6QSYX)]
 
 ## 5.4.2 -- 2024-12-16
+
 ## 5.4.1 -- 2024-12-13
+
 ## 5.4.0 -- 2024-11-18
 
 - Changed auto-update task to reuse current scheduler instead of creating a new one
@@ -22,6 +26,7 @@ This file keeps track of all notable changes to jobbergate-agent
 - Cache slurm submissions to avoid the resubmission of the same job if the job status update fails [PENG-2342]
 
 ## 5.2.0 -- 2024-07-01
+
 - Change pydantic.BaseSettings config to use `extra=ignore`
 - Migrated to Pydantic version 2 [PENG-2278]
   - Upgraded pydantic to 2.7

jobbergate-agent/jobbergate_agent/jobbergate/submit.py

@@ -7,7 +7,7 @@
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
-from buzz import DoExceptParams
+from buzz import DoExceptParams, handle_errors_async
 from jobbergate_core.tools.sbatch import InfoHandler, SubmissionHandler, SubprocessHandler, inject_sbatch_params
 from loguru import logger
 
@@ -16,7 +16,7 @@
 from jobbergate_agent.jobbergate.schemas import JobScriptFile, PendingJobSubmission, SlurmJobData
 from jobbergate_agent.jobbergate.update import fetch_job_data
 from jobbergate_agent.settings import SETTINGS
-from jobbergate_agent.utils.exception import JobbergateApiError, JobSubmissionError, handle_errors_async
+from jobbergate_agent.utils.exception import JobbergateApiError, JobSubmissionError
 from jobbergate_agent.utils.logging import log_error
 from jobbergate_agent.utils.user_mapper import SlurmUserMapper, manufacture
 

jobbergate-agent/jobbergate_agent/utils/exception.py

@@ -1,11 +1,6 @@
 """Core module for exception related operations"""
 
-import contextlib
-from asyncio import iscoroutinefunction
-from typing import Any, AsyncIterator, Callable, Coroutine, Iterable, Mapping, Optional, Tuple, Type, Union
-
-from buzz import Buzz, get_traceback, reformat_exception
-from buzz.tools import DoExceptParams, noop
+from buzz import Buzz
 
 
 class ClusterAgentError(Buzz):
@@ -34,79 +29,3 @@ class JobSubmissionError(ClusterAgentError):
 
 class SlurmParameterParserError(ClusterAgentError):
     """Raise exception when Slurm mapper or SBATCH parser face any error"""
-
-
-@contextlib.asynccontextmanager
-async def handle_errors_async(
-    message: str,
-    raise_exc_class: Union[Type[Exception], None] = Exception,
-    raise_args: Optional[Iterable[Any]] = None,
-    raise_kwargs: Optional[Mapping[str, Any]] = None,
-    handle_exc_class: Union[Type[Exception], Tuple[Type[Exception], ...]] = Exception,
-    do_finally: Callable[[], None] | Callable[[], Coroutine[Any, Any, None]] = noop,
-    do_except: Callable[[DoExceptParams], None] | Callable[[DoExceptParams], Coroutine[Any, Any, None]] = noop,
-    do_else: Callable[[], None] | Callable[[], Coroutine[Any, Any, None]] = noop,
-) -> AsyncIterator[None]:
-    """
-    Async context manager that will intercept exceptions and repackage them with a message attached.
-
-    Example:
-        ``` python
-        with handle_errors("It didn't work"):
-            some_code_that_might_raise_an_exception()
-        ```
-
-    Args:
-        message: The message to attach to the raised exception.
-        raise_exc_class: The exception type to raise with the constructed message
-            if an exception is caught in the managed context.
-
-            Defaults to Exception.
-
-            If ``None`` is passed, no new exception will be raised and only the
-            ``do_except``, ``do_else``, and ``do_finally``
-            functions will be called.
-        raise_args: Additional positional args (after the constructed message) that will
-            be passed when raising an instance of the ``raise_exc_class``.
-        raise_kwargs: Keyword args that will be passed when raising an instance of the ``raise_exc_class``.
-        handle_exc_class: Limits the class of exceptions that will be intercepted.
-            Any other exception types will not be caught and re-packaged.
-            Defaults to Exception (will handle all exceptions). May also be
-            provided as a tuple of multiple exception types to handle.
-        do_finally: A function that should always be called at the end of the block. Should take no parameters.
-        do_except: A function that should be called only if there was an exception.
-            Must accept one parameter that is an instance of the
-            ``DoExceptParams`` dataclass. Note that the ``do_except``
-            method is passed the *original exception*.
-        do_else: A function that should be called only if there were no
-            exceptions encountered.
-    """
-    try:
-        yield
-    except handle_exc_class as err:
-        try:
-            final_message = reformat_exception(message, err)
-        except Exception as msg_err:
-            raise RuntimeError(f"Failed while formatting message: {repr(msg_err)}")
-
-        trace = get_traceback()
-
-        if iscoroutinefunction(do_except):
-            await do_except(DoExceptParams(err, final_message, trace))
-        else:
-            do_except(DoExceptParams(err, final_message, trace))
-
-        if raise_exc_class is not None:
-            args = raise_args or []
-            kwargs = raise_kwargs or {}
-            raise raise_exc_class(final_message, *args, **kwargs).with_traceback(trace)
-    else:
-        if iscoroutinefunction(do_else):
-            await do_else()
-        else:
-            do_else()
-    finally:
-        if iscoroutinefunction(do_finally):
-            await do_finally()
-        else:
-            do_finally()