From 920569700d2b9e2f0a221f7606568d5b8df56d04 Mon Sep 17 00:00:00 2001 From: Teo Date: Thu, 31 Oct 2024 21:29:01 -0600 Subject: [PATCH 1/2] feat(session.py): implement graceful shutdown handling for sessions using atexit and signal refactor(session.py): streamline end_session logic to improve readability and maintainability --- agentops/session.py | 45 ++++++++++++++++++++++++++++++++------------- 1 file changed, 32 insertions(+), 13 deletions(-) diff --git a/agentops/session.py b/agentops/session.py index 58225b6cc..c1064a7f1 100644 --- a/agentops/session.py +++ b/agentops/session.py @@ -8,6 +8,8 @@ from typing import Optional, List, Union from uuid import UUID, uuid4 from datetime import datetime +import atexit +import signal from .exceptions import ApiServerException from .enums import EndState @@ -71,6 +73,10 @@ def __init__( self.stop_flag.set() self.thread.join(timeout=1) + self._cleanup_done = False + atexit.register(self._cleanup) + signal.signal(signal.SIGINT, self._signal_handler) + def set_video(self, video: str) -> None: """ Sets a url to the video recording of the session. @@ -86,20 +92,15 @@ def end_session( end_state_reason: Optional[str] = None, video: Optional[str] = None, ) -> Union[Decimal, None]: + with self.lock: + if not self.is_running or self.end_timestamp: + return - if not self.is_running: - return - - if not any(end_state == state.value for state in EndState): - return logger.warning( - "Invalid end_state. Please use one of the EndState enums" - ) - - self.end_timestamp = get_ISO_time() - self.end_state = end_state - self.end_state_reason = end_state_reason - if video is not None: - self.video = video + self.end_timestamp = get_ISO_time() + self.end_state = end_state + self.end_state_reason = end_state_reason + if video is not None: + self.video = video self.stop_flag.set() self.thread.join(timeout=1) @@ -396,5 +397,23 @@ def wrapper(*args, **kwargs): return wrapper + def _signal_handler(self, signum, frame): + """Handle interrupt signals gracefully""" + self._cleanup() + # Re-raise the signal after cleanup + signal.signal(signum, signal.default_int_handler) + signal.raise_signal(signum) + + def _cleanup(self): + """Ensure clean shutdown of session""" + if not self._cleanup_done: + with self.lock: + if self.is_running and not self.end_timestamp: + self.end_session( + end_state="Indeterminate", + end_state_reason="Session interrupted" + ) + self._cleanup_done = True + active_sessions: List[Session] = [] From 8abfda619242495a80d8870ee096e1f115e171f1 Mon Sep 17 00:00:00 2001 From: Teo Date: Fri, 15 Nov 2024 13:34:31 -0600 Subject: [PATCH 2/2] black Signed-off-by: Teo --- agentops/session.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agentops/session.py b/agentops/session.py index 7ea3f9930..4948d55a3 100644 --- a/agentops/session.py +++ b/agentops/session.py @@ -406,7 +406,7 @@ def _cleanup(self): if self.is_running and not self.end_timestamp: self.end_session( end_state="Indeterminate", - end_state_reason="Session interrupted" + end_state_reason="Session interrupted", ) self._cleanup_done = True