Skip to content

Commit

Permalink
cherry-pick: PENG-2678 Fixed issue with handle_errors_async
Browse files Browse the repository at this point in the history
The implementation of `handle_errors_async()` that was embedded in the
`jobbergate-agent` was not compatible with a recent update to `py-buzz`
version `4.1`.

This change removes the embedded implementation in favor of using the
upstream implementation of the same that is now a part of `py-buzz`
version `5.0`.
  • Loading branch information
dusktreader authored and fschuch committed Jan 30, 2025
1 parent d794e87 commit f94adda
Show file tree
Hide file tree
Showing 7 changed files with 498 additions and 424 deletions.
5 changes: 5 additions & 0 deletions jobbergate-agent/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
4 changes: 2 additions & 2 deletions jobbergate-agent/jobbergate_agent/jobbergate/submit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

Expand Down
83 changes: 1 addition & 82 deletions jobbergate-agent/jobbergate_agent/utils/exception.py
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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()
Loading

0 comments on commit f94adda

Please sign in to comment.