langfun 0.1.2.dev202510230805__py3-none-any.whl → 0.1.2.dev202511270805__py3-none-any.whl

langfun/env/event_handlers/event_logger.py

@@ -16,11 +16,11 @@
 import re
 import time
 from typing import Annotated
-from langfun.env.event_handlers import base
+from langfun.env import interface
 import pyglove as pg
 
 
-class EventLogger(pg.Object, base.EventHandler):
+class EventLogger(pg.Object, interface.EventHandler):
   """Event handler for logging debugger."""
 
   colored: Annotated[
@@ -125,7 +125,7 @@ class EventLogger(pg.Object, base.EventHandler):
 
   def on_environment_starting(
       self,
-      environment: base.Environment,
+      environment: interface.Environment,
   ) -> None:
     """Called when the environment is starting."""
     self._print(
@@ -137,7 +137,7 @@ class EventLogger(pg.Object, base.EventHandler):
 
   def on_environment_shutting_down(
       self,
-      environment: base.Environment,
+      environment: interface.Environment,
       offline_duration: float,
   ) -> None:
     """Called when the environment is shutting down."""
@@ -151,7 +151,7 @@ class EventLogger(pg.Object, base.EventHandler):
 
   def on_environment_start(
       self,
-      environment: base.Environment,
+      environment: interface.Environment,
       duration: float,
       error: BaseException | None
   ) -> None:
@@ -166,7 +166,7 @@ class EventLogger(pg.Object, base.EventHandler):
 
   def on_environment_housekeep(
       self,
-      environment: base.Environment,
+      environment: interface.Environment,
       counter: int,
       duration: float,
       error: BaseException | None,
@@ -194,7 +194,7 @@ class EventLogger(pg.Object, base.EventHandler):
 
   def on_environment_shutdown(
       self,
-      environment: base.Environment,
+      environment: interface.Environment,
       duration: float,
       lifetime: float,
       error: BaseException | None
@@ -210,8 +210,7 @@ class EventLogger(pg.Object, base.EventHandler):
 
   def on_sandbox_start(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
+      sandbox: interface.Sandbox,
       duration: float,
       error: BaseException | None
   ) -> None:
@@ -226,10 +225,9 @@ class EventLogger(pg.Object, base.EventHandler):
 
   def on_sandbox_status_change(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      old_status: base.Sandbox.Status,
-      new_status: base.Sandbox.Status,
+      sandbox: interface.Sandbox,
+      old_status: interface.Sandbox.Status,
+      new_status: interface.Sandbox.Status,
       span: float
   ) -> None:
     if self.sandbox_status:
@@ -242,8 +240,7 @@ class EventLogger(pg.Object, base.EventHandler):
 
   def on_sandbox_shutdown(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
+      sandbox: interface.Sandbox,
       duration: float,
       lifetime: float,
       error: BaseException | None
@@ -258,10 +255,62 @@ class EventLogger(pg.Object, base.EventHandler):
           styles=['bold'],
       )
 
+  def on_sandbox_session_start(
+      self,
+      sandbox: interface.Sandbox,
+      session_id: str,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox session starts."""
+    if self.session_status:
+      self._print(
+          f'[{sandbox.id}@{session_id}] sandbox session started '
+          f'(duration={duration:.2f} seconds)',
+          error=error,
+          color='blue',
+      )
+
+  def on_sandbox_session_end(
+      self,
+      sandbox: interface.Sandbox,
+      session_id: str,
+      duration: float,
+      lifetime: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox session ends."""
+    if self.session_status:
+      self._print(
+          f'[{sandbox.id}@{session_id}] sandbox session ended '
+          f'(duration={duration:.2f} seconds), '
+          f'lifetime={lifetime:.2f} seconds)',
+          error=error,
+          color='blue',
+      )
+
+  def on_sandbox_activity(
+      self,
+      name: str,
+      sandbox: interface.Sandbox,
+      session_id: str | None,
+      duration: float,
+      error: BaseException | None,
+      **kwargs
+  ) -> None:
+    """Called when a sandbox activity is performed."""
+    log_id = f'{sandbox.id}@{session_id or "<idle>"}'
+    color = 'yellow' if session_id is None else 'cyan'
+    self._print(
+        f'[{log_id}] sandbox call {name!r} '
+        f'(duration={duration:.2f} seconds, kwargs={kwargs}) ',
+        error,
+        color=color
+    )
+
   def on_sandbox_housekeep(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
+      sandbox: interface.Sandbox,
       counter: int,
       duration: float,
       error: BaseException | None,
@@ -279,16 +328,14 @@ class EventLogger(pg.Object, base.EventHandler):
 
   def on_feature_setup(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      feature: base.Feature,
+      feature: interface.Feature,
       duration: float,
       error: BaseException | None
   ) -> None:
     """Called when a sandbox feature is setup."""
     if self.feature_status:
       self._print(
-          f'[{sandbox.id}/<idle>/{feature.name}] feature setup complete '
+          f'[{feature.id}] feature setup complete '
           f'(duration={duration:.2f} seconds)',
           error=error,
           color='white',
@@ -296,16 +343,14 @@ class EventLogger(pg.Object, base.EventHandler):
 
   def on_feature_teardown(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      feature: base.Feature,
+      feature: interface.Feature,
       duration: float,
       error: BaseException | None
   ) -> None:
     """Called when a sandbox feature is teardown."""
     if self.feature_status:
       self._print(
-          f'[{sandbox.id}/<idle>/{feature.name}] feature teardown complete '
+          f'[{feature.id}] feature teardown complete '
           f'(duration={duration:.2f} seconds)',
           error=error,
           color='white',
@@ -313,9 +358,7 @@ class EventLogger(pg.Object, base.EventHandler):
 
   def on_feature_setup_session(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      feature: base.Feature,
+      feature: interface.Feature,
       session_id: str | None,
       duration: float,
       error: BaseException | None
@@ -323,7 +366,7 @@ class EventLogger(pg.Object, base.EventHandler):
     """Called when a sandbox feature is setup."""
     if self.feature_status:
       self._print(
-          f'[{sandbox.id}/{session_id or "<idle>"}/{feature.name}] '
+          f'[{feature.id}@{session_id or "<idle>"}] '
           f'feature setup complete (duration={duration:.2f} seconds)',
           error=error,
           color='yellow',
@@ -331,9 +374,7 @@ class EventLogger(pg.Object, base.EventHandler):
 
   def on_feature_teardown_session(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      feature: base.Feature,
+      feature: interface.Feature,
       session_id: str,
       duration: float,
       error: BaseException | None
@@ -341,17 +382,34 @@ class EventLogger(pg.Object, base.EventHandler):
     """Called when a sandbox feature is teardown."""
     if self.feature_status:
       self._print(
-          f'[{sandbox.id}/{session_id}>/{feature.name}] '
+          f'[{feature.id}@{session_id}] '
           f'feature teardown complete (duration={duration:.2f} seconds)',
           error=error,
           color='yellow',
       )
 
+  def on_feature_activity(
+      self,
+      name: str,
+      feature: interface.Feature,
+      session_id: str | None,
+      duration: float,
+      error: BaseException | None,
+      **kwargs
+  ) -> None:
+    """Called when a feature activity is performed."""
+    log_id = f'{feature.id}@{session_id or "<idle>"}'
+    color = 'yellow' if session_id is None else 'cyan'
+    self._print(
+        f'[{log_id}] feature call {name!r} '
+        f'(duration={duration:.2f} seconds, kwargs={kwargs}) ',
+        error,
+        color=color
+    )
+
   def on_feature_housekeep(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      feature: base.Feature,
+      feature: interface.Feature,
       counter: int,
       duration: float,
       error: BaseException | None,
@@ -360,74 +418,13 @@ class EventLogger(pg.Object, base.EventHandler):
     """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'[{feature.id}] feature housekeeping complete '
           f'(counter={counter}, (duration={duration:.2f} seconds, '
           f'housekeep_info={kwargs})',
           error=error,
           color='white',
       )
 
-  def on_session_start(
-      self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      session_id: str,
-      duration: float,
-      error: BaseException | None
-  ) -> None:
-    """Called when a sandbox session starts."""
-    if self.session_status:
-      self._print(
-          f'[{sandbox.id}/{session_id}] session started '
-          f'(duration={duration:.2f} seconds)',
-          error=error,
-          color='blue',
-      )
-
-  def on_session_end(
-      self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      session_id: str,
-      duration: float,
-      lifetime: float,
-      error: BaseException | None
-  ) -> None:
-    """Called when a sandbox session ends."""
-    if self.session_status:
-      self._print(
-          f'[{sandbox.id}/{session_id}] session ended '
-          f'(duration={duration:.2f} seconds), '
-          f'lifetime={lifetime:.2f} seconds)',
-          error=error,
-          color='blue',
-      )
-
-  def on_sandbox_activity(
-      self,
-      name: str,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      feature: base.Feature | None,
-      session_id: str | None,
-      duration: float,
-      error: BaseException | None,
-      **kwargs
-  ) -> None:
-    """Called when a sandbox activity is performed."""
-    del environment
-    log_id = f'{sandbox.id}/{session_id or "<idle>"}'
-    if feature is not None:
-      log_id = f'{log_id}/{feature.name}'
-
-    color = 'yellow' if session_id is None else 'cyan'
-    self._print(
-        f'[{log_id}] call {name!r} '
-        f'(duration={duration:.2f} seconds, kwargs={kwargs}) ',
-        error,
-        color=color
-    )
-
   def _print(
       self,
       message: str,

langfun/env/event_handlers/event_logger_test.py

@@ -71,15 +71,15 @@ class EventLoggerTest(unittest.TestCase):
         outage_retry_interval=0,
         housekeep_interval=1.0,
         sandbox_keepalive_interval=1.0,
-        event_handlers=[event_logger],
+        event_handler=event_logger,
     )
     with self._capture_logs(test_name) as stream:
       with env:
-        with env.sandbox('session1') as sb:
+        with env.sandbox(session_id='session1') as sb:
           self.assertEqual(sb.test_feature1.num_shell_calls(), 4)
 
         with self.assertRaises(interface.SandboxStateError):
-          with env.sandbox('session2') as sb:
+          with env.sandbox(session_id='session2') as sb:
             sb.shell('echo "bar"', raise_error=RuntimeError)
 
     stdout = stream.getvalue()
@@ -104,9 +104,9 @@ class EventLoggerTest(unittest.TestCase):
             'feature setup complete',
             'feature teardown complete',
             '/test_feature1] feature housekeeping',
-            'session started',
-            'session ended',
-            'call \'shell\'',
+            'sandbox session started',
+            'sandbox session ended',
+            'sandbox call \'shell\'',
             'RuntimeError',
         ],
         unexpected_substrings=[
@@ -154,9 +154,9 @@ class EventLoggerTest(unittest.TestCase):
             'sandbox shutdown',
             'feature setup complete',
             'feature teardown complete',
-            'session started',
-            'session ended',
-            'call \'shell\'',
+            'sandbox session started',
+            'sandbox session ended',
+            'sandbox call \'shell\'',
             'RuntimeError',
         ],
         regex='.*environment.*',
@@ -174,9 +174,9 @@ class EventLoggerTest(unittest.TestCase):
             'feature setup complete',
             'feature teardown complete',
             'feature housekeeping',
-            'session started',
-            'session ended',
-            'call \'shell\'',
+            'sandbox session started',
+            'sandbox session ended',
+            'sandbox call \'shell\'',
             'RuntimeError',
         ],
         unexpected_substrings=[
@@ -202,9 +202,9 @@ class EventLoggerTest(unittest.TestCase):
             '-> acquired',
             'sandbox shutdown',
             'sandbox housekeeping',
-            'session started',
-            'session ended',
-            'call \'shell\'',
+            'sandbox session started',
+            'sandbox session ended',
+            'sandbox call \'shell\'',
             'RuntimeError',
         ],
         unexpected_substrings=[
@@ -232,12 +232,12 @@ class EventLoggerTest(unittest.TestCase):
             'feature setup complete',
             'feature teardown complete',
             'feature housekeeping',
-            'call \'shell\'',
+            'sandbox call \'shell\'',
             'RuntimeError',
         ],
         unexpected_substrings=[
-            'session started',
-            'session ended',
+            'sandbox session started',
+            'sandbox session ended',
         ],
         session_status=False,
     )
@@ -256,9 +256,9 @@ class EventLoggerTest(unittest.TestCase):
             'sandbox shutdown',
             'feature setup complete',
             'feature teardown complete',
-            'session started',
-            'session ended',
-            'call \'shell\'',
+            'sandbox session started',
+            'sandbox session ended',
+            'sandbox call \'shell\'',
             'RuntimeError',
         ],
         unexpected_substrings=[