langfun 0.1.2.dev202509290805__py3-none-any.whl → 0.1.2.dev202510010805__py3-none-any.whl

langfun/env/event_handlers/event_logger.py

@@ -23,6 +23,13 @@ import pyglove as pg
 class EventLogger(pg.Object, base.EventHandler):
   """Event handler for logging debugger."""
 
+  colored: Annotated[
+      bool,
+      (
+          'If True, log events with colors.'
+      )
+  ] = False
+
   regex: Annotated[
       str | list[str] | None,
       (
@@ -85,13 +92,22 @@ class EventLogger(pg.Object, base.EventHandler):
     self._regex_exps = [re.compile(x) for x in regex_exps]
     self._last_stats_report_time = None
 
+  def _maybe_colored(
+      self, message: str, color: str, styles: list[str] | None = None
+  ) -> str:
+    if self.colored:
+      return pg.colored(message, color, styles=styles)
+    return message
+
   def _format_message(
       self,
       message: str,
       error: BaseException | None,
   ) -> str:
     if error is not None:
-      message = f'{message} with error: {error}'
+      message = (
+          f'{message} with error: {pg.utils.ErrorInfo.from_exception(error)}'
+      )
     return message
 
   def _keep(
@@ -107,6 +123,32 @@ class EventLogger(pg.Object, base.EventHandler):
       return False
     return True
 
+  def on_environment_starting(
+      self,
+      environment: base.Environment,
+  ) -> None:
+    """Called when the environment is starting."""
+    self._print(
+        f'[{environment.id}] environment starting',
+        error=None,
+        color='green',
+        styles=['bold'],
+    )
+
+  def on_environment_shutting_down(
+      self,
+      environment: base.Environment,
+      offline_duration: float,
+  ) -> None:
+    """Called when the environment is shutting down."""
+    self._print(
+        f'[{environment.id}] environment shutting down '
+        f'(offline_duration={offline_duration:.2f} seconds)',
+        error=None,
+        color='green',
+        styles=['bold'],
+    )
+
   def on_environment_start(
       self,
       environment: base.Environment,
@@ -127,13 +169,15 @@ class EventLogger(pg.Object, base.EventHandler):
       environment: base.Environment,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs
   ) -> None:
     """Called when the environment is housekeeping."""
     if self.housekeep_status:
       self._print(
-          f'[{environment.id}] environment housekeeping complete'
-          f'(counter={counter}, duration={duration:.2f} seconds)',
+          f'[{environment.id}] environment housekeeping complete '
+          f'(counter={counter}, duration={duration:.2f} seconds, '
+          f'housekeep_info={kwargs})',
           error=error,
           color='green',
       )
@@ -151,13 +195,14 @@ class EventLogger(pg.Object, base.EventHandler):
   def on_environment_shutdown(
       self,
       environment: base.Environment,
+      duration: float,
       lifetime: float,
       error: BaseException | None
   ) -> None:
     """Called when the environment is shutdown."""
     self._print(
         f'[{environment.id}] environment shutdown '
-        f'(lifetime={lifetime:.2f} seconds)',
+        f'(duration={duration:.2f} seconds), lifetime={lifetime:.2f} seconds)',
         error=error,
         color='green',
         styles=['bold'],
@@ -199,13 +244,15 @@ class EventLogger(pg.Object, base.EventHandler):
       self,
       environment: base.Environment,
       sandbox: base.Sandbox,
+      duration: float,
       lifetime: float,
       error: BaseException | None
   ) -> None:
     if self.sandbox_status:
       self._print(
           f'[{sandbox.id}] sandbox shutdown '
-          f'(lifetime={lifetime:.2f} seconds)',
+          f'(duration={duration:.2f} seconds), '
+          f'lifetime={lifetime:.2f} seconds)',
           error=error,
           color='white',
           styles=['bold'],
@@ -217,13 +264,15 @@ class EventLogger(pg.Object, base.EventHandler):
       sandbox: base.Sandbox,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs
   ) -> None:
     """Called when a sandbox feature is housekeeping."""
     if self.sandbox_status and self.housekeep_status:
       self._print(
           f'[{sandbox.id}] sandbox housekeeping complete '
-          f'(counter={counter}, duration={duration:.2f} seconds)',
+          f'(counter={counter}, duration={duration:.2f} seconds, '
+          f'housekeep_info={kwargs})',
           error=error,
           color='white',
       )
@@ -305,13 +354,15 @@ class EventLogger(pg.Object, base.EventHandler):
       feature: base.Feature,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs
   ) -> None:
     """Called when a sandbox feature is housekeeping."""
     if self.feature_status and self.housekeep_status:
       self._print(
           f'[{sandbox.id}/<idle>/{feature.name}] feature housekeeping complete '
-          f'(counter={counter}, (duration={duration:.2f} seconds)',
+          f'(counter={counter}, (duration={duration:.2f} seconds, '
+          f'housekeep_info={kwargs})',
           error=error,
           color='white',
       )
@@ -338,6 +389,7 @@ class EventLogger(pg.Object, base.EventHandler):
       environment: base.Environment,
       sandbox: base.Sandbox,
       session_id: str,
+      duration: float,
       lifetime: float,
       error: BaseException | None
   ) -> None:
@@ -345,7 +397,8 @@ class EventLogger(pg.Object, base.EventHandler):
     if self.session_status:
       self._print(
           f'[{sandbox.id}/{session_id}] session ended '
-          f'(lifetime={lifetime:.2f} seconds)',
+          f'(duration={duration:.2f} seconds), '
+          f'lifetime={lifetime:.2f} seconds)',
           error=error,
           color='blue',
       )
@@ -394,22 +447,29 @@ class EventLogger(pg.Object, base.EventHandler):
       color: str | None = None,
       styles: list[str] | None = None,
   ):
+    message = self._maybe_colored(
+        message, color if error is None else 'red', styles=styles
+    )
     if error is not None:
-      pg.logging.error(pg.colored(message, 'red', styles=styles))
+      pg.logging.error(message)
     else:
-      pg.logging.info(pg.colored(message, color, styles=styles))
+      pg.logging.info(message)
 
 
 class ConsoleEventLogger(EventLogger):
   """Event handler for console debugger."""
 
+  colored = True
+
   def _write_log(
       self,
       message: str,
       error: BaseException | None,
       color: str | None = None,
-      styles: list[str] | None = None,
+      styles: list[str] | None = None
   ):
     print(
-        pg.colored(message, color if error is None else 'red', styles=styles)
+        self._maybe_colored(
+            message, color if error is None else 'red', styles=styles
+        )
     )

langfun/env/event_handlers/event_logger_test.py

@@ -92,6 +92,7 @@ class EventLoggerTest(unittest.TestCase):
     return self._test_logger(
         test_name='test_all_flags_on',
         expected_substrings=[
+            'environment starting',
             'environment started',
             'environment shutdown',
             'environment housekeeping',
@@ -122,7 +123,9 @@ class EventLoggerTest(unittest.TestCase):
             'RuntimeError',
         ],
         unexpected_substrings=[
+            'environment starting',
             'environment started',
+            'environment shutting down',
             'environment shutdown',
             'environment housekeeping',
             'sandbox started',
@@ -139,7 +142,9 @@ class EventLoggerTest(unittest.TestCase):
     return self._test_logger(
         test_name='test_regex',
         expected_substrings=[
+            'environment starting',
             'environment started',
+            'environment shutting down',
             'environment shutdown',
             'environment housekeeping',
         ],
@@ -161,7 +166,9 @@ class EventLoggerTest(unittest.TestCase):
     return self._test_logger(
         test_name='test_sandbox_status_off',
         expected_substrings=[
+            'environment starting',
             'environment started',
+            'environment shutting down',
             'environment shutdown',
             'environment housekeeping',
             'feature setup complete',
@@ -185,7 +192,9 @@ class EventLoggerTest(unittest.TestCase):
     return self._test_logger(
         test_name='test_feature_status_off',
         expected_substrings=[
+            'environment starting',
             'environment started',
+            'environment shutting down',
             'environment shutdown',
             'environment housekeeping',
             'environment stats',
@@ -210,7 +219,9 @@ class EventLoggerTest(unittest.TestCase):
     return self._test_logger(
         test_name='test_session_status_off',
         expected_substrings=[
+            'environment starting',
             'environment started',
+            'environment shutting down',
             'environment shutdown',
             'environment stats',
             'environment housekeeping',
@@ -235,7 +246,9 @@ class EventLoggerTest(unittest.TestCase):
     return self._test_logger(
         test_name='test_housekeep_status_off',
         expected_substrings=[
+            'environment starting',
             'environment started',
+            'environment shutting down',
             'environment shutdown',
             'environment stats',
             'sandbox started',
@@ -260,7 +273,9 @@ class EventLoggerTest(unittest.TestCase):
     return self._test_logger(
         test_name='test_housekeep_status_off',
         expected_substrings=[
+            'environment starting',
             'environment started',
+            'environment shutting down',
             'environment shutdown',
             'environment housekeeping',
         ],