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

langfun/env/base_test.py

@@ -280,7 +280,11 @@ class EnvironmentTests(unittest.TestCase):
         with self.assertRaises(interface.SandboxStateError):
           sb.shell('bad command', raise_error=interface.SandboxStateError)
       self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
-      env.wait_for_housekeeping()
+      sb_offline_time = time.time()
+      while time.time() - sb_offline_time < 10:
+        if not env.is_online:
+          break
+        time.sleep(0.5)
       self.assertFalse(env.is_online)
 
 
@@ -381,34 +385,34 @@ class SandboxStatusTests(unittest.TestCase):
       self.assertEqual(
           self.event_handler.logs,
           [
-              '[testing-env/0] shell: "feature1" setup',
-              '[testing-env/0/feature1] feature setup',
-              '[testing-env/0] shell: "feature2" setup',
-              '[testing-env/0/feature2] feature setup',
-              '[testing-env/0] shell: "feature1" setup session',
-              '[testing-env/0/feature1] feature setup session',
-              '[testing-env/0] shell: "feature2" setup session',
-              '[testing-env/0/feature2] feature setup session',
-              '[testing-env/0] created -> ready',
-              '[testing-env/0] sandbox started',
+              '[testing-env/0:0] shell: "feature1" setup',
+              '[testing-env/0:0/feature1] feature setup',
+              '[testing-env/0:0] shell: "feature2" setup',
+              '[testing-env/0:0/feature2] feature setup',
+              '[testing-env/0:0] shell: "feature1" setup session',
+              '[testing-env/0:0/feature1] feature setup session',
+              '[testing-env/0:0] shell: "feature2" setup session',
+              '[testing-env/0:0/feature2] feature setup session',
+              '[testing-env/0:0] created -> ready',
+              '[testing-env/0:0] sandbox started',
               '[testing-env] environment started',
-              '[testing-env/0] ready -> acquired',
-              '[testing-env/0] acquired -> setting_up',
-              '[testing-env/0] setting_up -> in_session',
-              "[testing-env/0] session 'session1' started",
-              '[testing-env/0/session1] shell: echo "hello"',
-              '[testing-env/0] in_session -> exiting_session',
-              '[testing-env/0/session1] shell: "feature1" teardown session',
-              '[testing-env/0/feature1] feature teardown session',
-              '[testing-env/0/session1] shell: "feature2" teardown session',
-              '[testing-env/0/feature2] feature teardown session',
-              "[testing-env/0] session 'session1' ended",
-              '[testing-env/0] exiting_session -> setting_up',
-              '[testing-env/0] shell: "feature1" setup session',
-              '[testing-env/0/feature1] feature setup session',
-              '[testing-env/0] shell: "feature2" setup session',
-              '[testing-env/0/feature2] feature setup session',
-              '[testing-env/0] setting_up -> ready'
+              '[testing-env/0:0] ready -> acquired',
+              '[testing-env/0:0] acquired -> setting_up',
+              '[testing-env/0:0] setting_up -> in_session',
+              "[testing-env/0:0] session 'session1' started",
+              '[testing-env/0:0/session1] shell: echo "hello"',
+              '[testing-env/0:0] in_session -> exiting_session',
+              '[testing-env/0:0/session1] shell: "feature1" teardown session',
+              '[testing-env/0:0/feature1] feature teardown session',
+              '[testing-env/0:0/session1] shell: "feature2" teardown session',
+              '[testing-env/0:0/feature2] feature teardown session',
+              "[testing-env/0:0] session 'session1' ended",
+              '[testing-env/0:0] exiting_session -> setting_up',
+              '[testing-env/0:0] shell: "feature1" setup session',
+              '[testing-env/0:0/feature1] feature setup session',
+              '[testing-env/0:0] shell: "feature2" setup session',
+              '[testing-env/0:0/feature2] feature setup session',
+              '[testing-env/0:0] setting_up -> ready'
           ]
       )
 
@@ -440,17 +444,19 @@ class SandboxStatusTests(unittest.TestCase):
       self.assertEqual(
           event_handler.logs,
           [
-              '[testing-env/0] in_session -> exiting_session',
-              '[testing-env/0/session1] shell: "test_feature" teardown session',
-              '[testing-env/0/test_feature] feature teardown session',
-              "[testing-env/0] session 'session1' ended",
-              '[testing-env/0] exiting_session -> setting_up',
-              '[testing-env/0/test_feature] feature setup session with SandboxStateError',  # pylint: disable=line-too-long
-              '[testing-env/0] setting_up -> shutting_down',
-              '[testing-env/0] shell: "test_feature" teardown',
-              '[testing-env/0/test_feature] feature teardown',
-              '[testing-env/0] shutting_down -> offline',
-              '[testing-env/0] sandbox shutdown'
+              # pylint: disable=line-too-long
+              '[testing-env/0:0] in_session -> exiting_session',
+              '[testing-env/0:0/session1] shell: "test_feature" teardown session',
+              '[testing-env/0:0/test_feature] feature teardown session',
+              "[testing-env/0:0] session 'session1' ended",
+              '[testing-env/0:0] exiting_session -> setting_up',
+              '[testing-env/0:0/test_feature] feature setup session with SandboxStateError',  # pylint: disable=line-too-long
+              '[testing-env/0:0] setting_up -> shutting_down',
+              '[testing-env/0:0] shell: "test_feature" teardown',
+              '[testing-env/0:0/test_feature] feature teardown',
+              '[testing-env/0:0] shutting_down -> offline',
+              '[testing-env/0:0] sandbox shutdown'
+              # pylint: enable=line-too-long
           ]
       )
 
@@ -493,10 +499,10 @@ class SandboxStatusTests(unittest.TestCase):
     self.assertEqual(
         self.event_handler.logs,
         [
-            '[testing-env/0] sandbox started with SandboxStateError',
-            '[testing-env/0] created -> shutting_down',
-            '[testing-env/0] shutting_down -> offline',
-            '[testing-env/0] sandbox shutdown',
+            '[testing-env/0:0] sandbox started with SandboxStateError',
+            '[testing-env/0:0] created -> shutting_down',
+            '[testing-env/0:0] shutting_down -> offline',
+            '[testing-env/0:0] sandbox shutdown',
             '[testing-env] environment started with EnvironmentOutageError',
             '[testing-env] environment shutdown'
         ]
@@ -569,24 +575,24 @@ class SandboxStatusTests(unittest.TestCase):
     self.assertEqual(
         self.event_handler.logs,
         [
-            '[testing-env/0] shell: "feature1" setup',
-            '[testing-env/0/feature1] feature setup',
-            '[testing-env/0] shell: "feature2" setup',
-            '[testing-env/0/feature2] feature setup',
-            '[testing-env/0] shell: "feature1" setup session',
-            '[testing-env/0/feature1] feature setup session',
-            '[testing-env/0] shell: "feature2" setup session',
-            '[testing-env/0/feature2] feature setup session',
-            '[testing-env/0] created -> ready',
-            '[testing-env/0] sandbox started',
+            '[testing-env/0:0] shell: "feature1" setup',
+            '[testing-env/0:0/feature1] feature setup',
+            '[testing-env/0:0] shell: "feature2" setup',
+            '[testing-env/0:0/feature2] feature setup',
+            '[testing-env/0:0] shell: "feature1" setup session',
+            '[testing-env/0:0/feature1] feature setup session',
+            '[testing-env/0:0] shell: "feature2" setup session',
+            '[testing-env/0:0/feature2] feature setup session',
+            '[testing-env/0:0] created -> ready',
+            '[testing-env/0:0] sandbox started',
             '[testing-env] environment started',
-            '[testing-env/0] ready -> shutting_down',
-            '[testing-env/0] shell: "feature1" teardown',
-            '[testing-env/0/feature1] feature teardown',
-            '[testing-env/0] shell: "feature2" teardown',
-            '[testing-env/0/feature2] feature teardown',
-            '[testing-env/0] shutting_down -> offline',
-            '[testing-env/0] sandbox shutdown with ValueError',
+            '[testing-env/0:0] ready -> shutting_down',
+            '[testing-env/0:0] shell: "feature1" teardown',
+            '[testing-env/0:0/feature1] feature teardown',
+            '[testing-env/0:0] shell: "feature2" teardown',
+            '[testing-env/0:0/feature2] feature teardown',
+            '[testing-env/0:0] shutting_down -> offline',
+            '[testing-env/0:0] sandbox shutdown with ValueError',
             '[testing-env] environment shutdown with ValueError'
         ]
     )
@@ -998,27 +1004,27 @@ class SandboxStatusTests(unittest.TestCase):
       self.assertEqual(
           self.event_handler.logs,
           [
-              '[testing-env/0] shell: "feature1" setup',
-              '[testing-env/0/feature1] feature setup',
-              '[testing-env/0] shell: "feature1" setup session',
-              '[testing-env/0/feature1] feature setup session',
-              '[testing-env/0] created -> ready',
-              '[testing-env/0] sandbox started',
+              '[testing-env/0:0] shell: "feature1" setup',
+              '[testing-env/0:0/feature1] feature setup',
+              '[testing-env/0:0] shell: "feature1" setup session',
+              '[testing-env/0:0/feature1] feature setup session',
+              '[testing-env/0:0] created -> ready',
+              '[testing-env/0:0] sandbox started',
               '[testing-env] environment started',
-              '[testing-env/0] ready -> acquired',
-              '[testing-env/0] acquired -> setting_up',
-              '[testing-env/0] setting_up -> in_session',
-              "[testing-env/0] session 'session1' started",
-              '[testing-env/0/session1] shell: echo foo with ValueError',
-              '[testing-env/0/session1] shell: echo bar',
-              '[testing-env/0] in_session -> exiting_session',
-              '[testing-env/0/session1] shell: "feature1" teardown session',
-              '[testing-env/0/feature1] feature teardown session',
-              "[testing-env/0] session 'session1' ended",
-              '[testing-env/0] exiting_session -> setting_up',
-              '[testing-env/0] shell: "feature1" setup session',
-              '[testing-env/0/feature1] feature setup session',
-              '[testing-env/0] setting_up -> ready',
+              '[testing-env/0:0] ready -> acquired',
+              '[testing-env/0:0] acquired -> setting_up',
+              '[testing-env/0:0] setting_up -> in_session',
+              "[testing-env/0:0] session 'session1' started",
+              '[testing-env/0:0/session1] shell: echo foo with ValueError',
+              '[testing-env/0:0/session1] shell: echo bar',
+              '[testing-env/0:0] in_session -> exiting_session',
+              '[testing-env/0:0/session1] shell: "feature1" teardown session',
+              '[testing-env/0:0/feature1] feature teardown session',
+              "[testing-env/0:0] session 'session1' ended",
+              '[testing-env/0:0] exiting_session -> setting_up',
+              '[testing-env/0:0] shell: "feature1" setup session',
+              '[testing-env/0:0/feature1] feature setup session',
+              '[testing-env/0:0] setting_up -> ready',
           ]
       )
 
@@ -1039,29 +1045,31 @@ class SandboxStatusTests(unittest.TestCase):
       self.assertEqual(
           self.event_handler.logs,
           [
-              '[testing-env/0] shell: "feature1" setup',
-              '[testing-env/0/feature1] feature setup',
-              '[testing-env/0] shell: "feature1" setup session',
-              '[testing-env/0/feature1] feature setup session',
-              '[testing-env/0] created -> ready',
-              '[testing-env/0] sandbox started',
+              # pylint: disable=line-too-long
+              '[testing-env/0:0] shell: "feature1" setup',
+              '[testing-env/0:0/feature1] feature setup',
+              '[testing-env/0:0] shell: "feature1" setup session',
+              '[testing-env/0:0/feature1] feature setup session',
+              '[testing-env/0:0] created -> ready',
+              '[testing-env/0:0] sandbox started',
               '[testing-env] environment started',
-              '[testing-env/0] ready -> acquired',
-              '[testing-env/0] acquired -> setting_up',
-              '[testing-env/0] setting_up -> in_session',
-              "[testing-env/0] session 'session1' started",
-              '[testing-env/0/session1] shell: echo foo with RuntimeError',
-              '[testing-env/0] in_session -> exiting_session',
-              '[testing-env/0/session1] shell: "feature1" teardown session',
-              '[testing-env/0/feature1] feature teardown session',
-              "[testing-env/0] session 'session1' ended with SandboxStateError",
-              '[testing-env/0] exiting_session -> acquired',
-              '[testing-env/0] acquired -> shutting_down',
-              '[testing-env/0] shell: "feature1" teardown',
-              '[testing-env/0/feature1] feature teardown',
-              '[testing-env/0] shutting_down -> offline',
-              '[testing-env/0] sandbox shutdown',
-              '[testing-env/0] shell: echo bar',
+              '[testing-env/0:0] ready -> acquired',
+              '[testing-env/0:0] acquired -> setting_up',
+              '[testing-env/0:0] setting_up -> in_session',
+              "[testing-env/0:0] session 'session1' started",
+              '[testing-env/0:0/session1] shell: echo foo with RuntimeError',
+              '[testing-env/0:0] in_session -> exiting_session',
+              '[testing-env/0:0/session1] shell: "feature1" teardown session',
+              '[testing-env/0:0/feature1] feature teardown session',
+              "[testing-env/0:0] session 'session1' ended with SandboxStateError",
+              '[testing-env/0:0] exiting_session -> acquired',
+              '[testing-env/0:0] acquired -> shutting_down',
+              '[testing-env/0:0] shell: "feature1" teardown',
+              '[testing-env/0:0/feature1] feature teardown',
+              '[testing-env/0:0] shutting_down -> offline',
+              '[testing-env/0:0] sandbox shutdown',
+              '[testing-env/0:0] shell: echo bar',
+              # pylint: enable=line-too-long
           ]
       )
 
@@ -1107,12 +1115,15 @@ class SandboxActivityTests(unittest.TestCase):
         self.assertEqual(sb.status, sb.Status.OFFLINE)
 
   def test_housekeep_error(self):
+    event_handler = TestingEventHandler(log_housekeep=False)
     env = TestingEnvironment(
         features={'test_feature': TestingFeature(housekeep_interval=0)},
         pool_size=1,
+        housekeep_interval=1.0,
         outage_grace_period=0,
-        outage_retry_interval=0,
+        outage_retry_interval=0.1,
         sandbox_keepalive_interval=0,
+        event_handlers=[event_handler],
     )
     with env:
       with env.sandbox('session1') as sb:
@@ -1129,6 +1140,31 @@ class SandboxActivityTests(unittest.TestCase):
         ):
           time.sleep(0.01)
         self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
+      env.wait_for_housekeeping()
+    self.assertEqual(
+        event_handler.logs,
+        [
+            '[testing-env/0:0] shell: "test_feature" setup',
+            '[testing-env/0:0/test_feature] feature setup',
+            '[testing-env/0:0] shell: "test_feature" setup session',
+            '[testing-env/0:0] sandbox started',
+            '[testing-env] environment started',
+            "[testing-env/0:0] session 'session1' started",
+            '[testing-env/0:0/session1] shell: "test_feature" teardown session',
+            "[testing-env/0:0] session 'session1' ended with SandboxStateError",
+            '[testing-env/0:0] shell: "test_feature" teardown',
+            '[testing-env/0:0/test_feature] feature teardown',
+            '[testing-env/0:0] sandbox shutdown',
+            '[testing-env/0:1] shell: "test_feature" setup',
+            '[testing-env/0:1/test_feature] feature setup',
+            '[testing-env/0:1] shell: "test_feature" setup session',
+            '[testing-env/0:1] sandbox started',
+            '[testing-env/0:1] shell: "test_feature" teardown',
+            '[testing-env/0:1/test_feature] feature teardown',
+            '[testing-env/0:1] sandbox shutdown',
+            '[testing-env] environment shutdown'
+        ]
+    )
 
   def test_remove_event_handler(self):
     env = TestingEnvironment(
@@ -1145,6 +1181,7 @@ class SandboxActivityTests(unittest.TestCase):
         sb.shell('test_feature')
         sb.remove_event_handler(event_handler)
       events = list(event_handler.logs)
+      sb.wait_until_not(interface.Sandbox.Status.SETTING_UP)
       self.assertGreater(len(events), 0)
       with env.sandbox('session2') as sb:
         sb.shell('test_feature')

langfun/env/event_handlers/__init__.py

@@ -8,5 +8,7 @@ from langfun.env.event_handlers.base import EventHandler
 from langfun.env.event_handlers.event_logger import EventLogger
 from langfun.env.event_handlers.event_logger import ConsoleEventLogger
 
+from langfun.env.event_handlers.metric_writer import MetricWriter
+
 # pylint: enable=g-importing-member
 # pylint: enable=g-bad-import-order

langfun/env/event_handlers/base.py

@@ -48,6 +48,7 @@ class _SessionEventHandler:
       environment: Environment,
       sandbox: Sandbox,
       session_id: str,
+      duration: float,
       lifetime: float,
       error: BaseException | None
   ) -> None:
@@ -57,6 +58,7 @@ class _SessionEventHandler:
       environment: The environment.
       sandbox: The sandbox.
       session_id: The session ID.
+      duration: The time spent on ending the session.
       lifetime: The session lifetime in seconds.
       error: The error that caused the session to end. If None, the session
         ended normally.
@@ -74,7 +76,16 @@ class _FeatureEventHandler:
       duration: float,
       error: BaseException | None
   ) -> None:
-    """Called when a sandbox feature is setup."""
+    """Called when a sandbox feature is setup.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      feature: The feature.
+      duration: The feature setup duration in seconds.
+      error: The error happened during the feature setup. If None,
+        the feature setup performed normally.
+    """
 
   def on_feature_teardown(
       self,
@@ -84,7 +95,16 @@ class _FeatureEventHandler:
       duration: float,
       error: BaseException | None
   ) -> None:
-    """Called when a sandbox feature is teardown."""
+    """Called when a sandbox feature is teardown.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      feature: The feature.
+      duration: The feature teardown duration in seconds.
+      error: The error happened during the feature teardown. If None,
+        the feature teardown performed normally.
+    """
 
   def on_feature_teardown_session(
       self,
@@ -95,7 +115,17 @@ class _FeatureEventHandler:
       duration: float,
       error: BaseException | None
   ) -> None:
-    """Called when a feature is teardown with a session."""
+    """Called when a feature is teardown with a session.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      feature: The feature.
+      session_id: The session ID.
+      duration: The feature teardown session duration in seconds.
+      error: The error happened during the feature teardown session. If
+        None, the feature teardown session performed normally.
+    """
 
   def on_feature_setup_session(
       self,
@@ -104,9 +134,19 @@ class _FeatureEventHandler:
       feature: Feature,
       session_id: str | None,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
   ) -> None:
-    """Called when a feature is setup with a session."""
+    """Called when a feature is setup with a session.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      feature: The feature.
+      session_id: The session ID.
+      duration: The feature setup session duration in seconds.
+      error: The error happened during the feature setup session. If
+        None, the feature setup session performed normally.
+    """
 
   def on_feature_housekeep(
       self,
@@ -115,9 +155,21 @@ class _FeatureEventHandler:
       feature: Feature,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs,
   ) -> None:
-    """Called when a sandbox feature is housekeeping."""
+    """Called when a sandbox feature is housekeeping.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      feature: The feature.
+      counter: Zero-based counter of the housekeeping round.
+      duration: The feature housekeeping duration in seconds.
+      error: The error happened during the feature housekeeping. If None, the
+        feature housekeeping normally.
+      **kwargs: Feature-specific properties computed during housekeeping.
+    """
 
 
 class _SandboxEventHandler(_FeatureEventHandler, _SessionEventHandler):
@@ -162,6 +214,7 @@ class _SandboxEventHandler(_FeatureEventHandler, _SessionEventHandler):
       self,
       environment: Environment,
       sandbox: Sandbox,
+      duration: float,
       lifetime: float,
       error: BaseException | None
   ) -> None:
@@ -170,6 +223,7 @@ class _SandboxEventHandler(_FeatureEventHandler, _SessionEventHandler):
     Args:
       environment: The environment.
       sandbox: The sandbox.
+      duration: The time spent on shutting down the sandbox.
       lifetime: The sandbox lifetime in seconds.
       error: The error that caused the sandbox to shutdown. If None, the
         sandbox shutdown normally.
@@ -206,7 +260,8 @@ class _SandboxEventHandler(_FeatureEventHandler, _SessionEventHandler):
       sandbox: Sandbox,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs
   ) -> None:
     """Called when a sandbox finishes a round of housekeeping.
 
@@ -217,12 +272,20 @@ class _SandboxEventHandler(_FeatureEventHandler, _SessionEventHandler):
       duration: The sandbox housekeeping duration in seconds.
       error: The error that caused the sandbox to housekeeping. If None, the
         sandbox housekeeping normally.
+      **kwargs: Sandbox-specific properties computed during housekeeping.
     """
 
 
 class EventHandler(_SandboxEventHandler):
   """Base class for event handlers of an environment."""
 
+  def on_environment_starting(self, environment: Environment) -> None:
+    """Called when the environment is getting started.
+
+    Args:
+      environment: The environment.
+    """
+
   def on_environment_start(
       self,
       environment: Environment,
@@ -243,7 +306,8 @@ class EventHandler(_SandboxEventHandler):
       environment: Environment,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs
   ) -> None:
     """Called when the environment finishes a round of housekeeping.
 
@@ -253,11 +317,25 @@ class EventHandler(_SandboxEventHandler):
       duration: The environment start duration in seconds.
       error: The error that failed the housekeeping. If None, the
         housekeeping succeeded.
+      **kwargs: Environment-specific properties computed during housekeeping.
+    """
+
+  def on_environment_shutting_down(
+      self,
+      environment: Environment,
+      offline_duration: float,
+  ) -> None:
+    """Called when the environment is shutting down.
+
+    Args:
+      environment: The environment.
+      offline_duration: The environment offline duration in seconds.
     """
 
   def on_environment_shutdown(
       self,
       environment: Environment,
+      duration: float,
       lifetime: float,
       error: BaseException | None
   ) -> None:
@@ -265,6 +343,7 @@ class EventHandler(_SandboxEventHandler):
 
     Args:
       environment: The environment.
+      duration: The environment shutdown duration in seconds.
       lifetime: The environment lifetime in seconds.
       error: The error that caused the environment to shutdown. If None, the
         environment shutdown normally.