langfun 0.1.2.dev202510300805__py3-none-any.whl → 0.1.2.dev202511010804__py3-none-any.whl

langfun/env/__init__.py

@@ -24,6 +24,8 @@ from langfun.env.interface import Environment
 from langfun.env.interface import Sandbox
 from langfun.env.interface import Feature
 
+from langfun.env.interface import EventHandler
+
 # Decorators for sandbox/feature methods.
 from langfun.env.interface import treat_as_sandbox_state_error
 from langfun.env.interface import log_sandbox_activity
@@ -35,7 +37,7 @@ from langfun.env.base_feature import BaseFeature
 from langfun.env import load_balancers
 from langfun.env.load_balancers import LoadBalancer
 
-from langfun.env import event_handlers
-EventHandler = event_handlers.EventHandler
+# Import all event handlers.
+from langfun.env.event_handlers import *
 
 # Google-internal imports.

langfun/env/base_environment.py

@@ -36,7 +36,6 @@ import langfun.core as lf
 from langfun.env import base_sandbox
 from langfun.env import interface
 from langfun.env import load_balancers
-from langfun.env.event_handlers import base as event_handler_base
 import pyglove as pg
 
 
@@ -113,12 +112,13 @@ class BaseEnvironment(interface.Environment):
       )
   ] = True
 
-  event_handlers: Annotated[
-      list[event_handler_base.EventHandler],
+  event_handler: Annotated[
+      interface.EventHandler,
       (
           'User handler for the environment events.'
+          'By default, the no-op event handler is used.'
       )
-  ] = []
+  ] = interface.EventHandler()
 
   outage_grace_period: Annotated[
       float,
@@ -169,6 +169,7 @@ class BaseEnvironment(interface.Environment):
 
     self._status = self.Status.CREATED
     self._start_time = None
+
     self._sandbox_pool: dict[str, list[base_sandbox.BaseSandbox]] = (
         collections.defaultdict(list)
     )
@@ -555,8 +556,6 @@ class BaseEnvironment(interface.Environment):
           proactive_session_setup=self.proactive_session_setup,
           keepalive_interval=self.sandbox_keepalive_interval,
       )
-      for handler in self.event_handlers:
-        sandbox.add_event_handler(handler)
       sandbox.start()
       if set_acquired:
         sandbox.set_acquired()
@@ -738,16 +737,14 @@ class BaseEnvironment(interface.Environment):
 
   def on_starting(self) -> None:
     """Called when the environment is getting started."""
-    for handler in self.event_handlers:
-      handler.on_environment_starting(self)
+    self.event_handler.on_environment_starting(self)
 
   def on_start(
       self,
       duration: float, error: BaseException | None = None
   ) -> None:
     """Called when the environment is started."""
-    for handler in self.event_handlers:
-      handler.on_environment_start(self, duration, error)
+    self.event_handler.on_environment_start(self, duration, error)
 
   def on_housekeep(
       self,
@@ -756,16 +753,13 @@ class BaseEnvironment(interface.Environment):
       **kwargs
   ) -> None:
     """Called when the environment finishes a round of housekeeping."""
-    housekeep_counter = self.housekeep_counter
-    for handler in self.event_handlers:
-      handler.on_environment_housekeep(
-          self, housekeep_counter, duration, error, **kwargs
-      )
+    self.event_handler.on_environment_housekeep(
+        self, self.housekeep_counter, duration, error, **kwargs
+    )
 
   def on_shutting_down(self) -> None:
     """Called when the environment is shutting down."""
-    for handler in self.event_handlers:
-      handler.on_environment_shutting_down(self, self.offline_duration)
+    self.event_handler.on_environment_shutting_down(self, self.offline_duration)
 
   def on_shutdown(
       self,
@@ -773,5 +767,4 @@ class BaseEnvironment(interface.Environment):
       error: BaseException | None = None) -> None:
     """Called when the environment is shutdown."""
     lifetime = (time.time() - self.start_time) if self.start_time else 0.0
-    for handler in self.event_handlers:
-      handler.on_environment_shutdown(self, duration, lifetime, error)
+    self.event_handler.on_environment_shutdown(self, duration, lifetime, error)

langfun/env/base_environment_test.py

@@ -0,0 +1,473 @@
+# Copyright 2025 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import time
+import unittest
+
+from langfun.env import interface
+from langfun.env import test_utils
+
+TestingEnvironment = test_utils.TestingEnvironment
+TestingSandbox = test_utils.TestingSandbox
+TestingFeature = test_utils.TestingFeature
+TestingEventHandler = test_utils.TestingEventHandler
+
+
+class BaseEnvironmentTests(unittest.TestCase):
+
+  def test_basics(self):
+    env = TestingEnvironment(
+        image_ids=['test_image'],
+        root_dir='/tmp',
+        pool_size=0,
+        features={'test_feature': TestingFeature()},
+        outage_grace_period=1,
+        outage_retry_interval=0,
+    )
+    self.assertIsNone(interface.Environment.current())
+    self.assertEqual(env.image_ids, ['test_image'])
+    self.assertFalse(env.supports_dynamic_image_loading)
+    self.assertEqual(env.status, interface.Environment.Status.CREATED)
+    self.assertFalse(env.is_online)
+    self.assertEqual(env.min_pool_size('test_image'), 0)
+    self.assertEqual(env.max_pool_size('test_image'), 0)
+    self.assertEqual(env.sandbox_pool, {})
+    self.assertEqual(env.id, interface.Environment.Id('testing-env'))
+    self.assertEqual(env.outage_grace_period, 1)
+    self.assertEqual(env.features['test_feature'].name, 'test_feature')
+
+    self.assertIsNone(env.start_time)
+
+    with env:
+      self.assertEqual(env.status, interface.Environment.Status.ONLINE)
+      self.assertIs(interface.Environment.current(), env)
+      self.assertTrue(env.is_online)
+      self.assertIsNotNone(env.start_time)
+      self.assertEqual(env.offline_duration, 0.0)
+      self.assertEqual(env.sandbox_pool, {})
+      self.assertEqual(env.working_dir, '/tmp/testing-env')
+
+      with env.sandbox(session_id='session1') as sb:
+        self.assertEqual(
+            sb.id, interface.Sandbox.Id(
+                environment_id=env.id,
+                image_id=sb.image_id,
+                sandbox_id='0')
+        )
+        self.assertEqual(sb.session_id, 'session1')
+        self.assertEqual(sb.working_dir, '/tmp/testing-env/test_image/0')
+        self.assertTrue(sb.is_online)
+        self.assertIs(sb.test_feature, sb.features['test_feature'])
+        self.assertEqual(
+            sb.test_feature.working_dir,
+            '/tmp/testing-env/test_image/0/test_feature'
+        )
+        with self.assertRaises(AttributeError):
+          _ = sb.test_feature2
+      self.assertFalse(sb.is_online)
+
+      with self.assertRaisesRegex(
+          ValueError, 'Environment .* does not serve image ID .*'
+      ):
+        env.sandbox('test_image2')
+
+      with env.test_feature() as feature:
+        self.assertIsInstance(feature, TestingFeature)
+        self.assertEqual(
+            feature.sandbox.status, interface.Sandbox.Status.IN_SESSION
+        )
+        self.assertTrue(
+            feature.sandbox.session_id.startswith('test_feature-session')
+        )
+
+      with self.assertRaises(AttributeError):
+        _ = env.test_feature2
+
+  def test_dynamic_image_loading(self):
+    env = TestingEnvironment(
+        image_ids=[],
+        supports_dynamic_image_loading=True,
+        pool_size=0,
+        features={'test_feature': TestingFeature()},
+        outage_grace_period=1,
+        outage_retry_interval=0,
+    )
+    with env:
+      with env.sandbox(image_id='test_image2') as sb:
+        self.assertEqual(sb.image_id, 'test_image2')
+
+      with self.assertRaisesRegex(
+          ValueError, 'Environment .* does not have a default image ID.'
+      ):
+        env.sandbox()
+
+  def test_dynamic_image_loading_with_pooling(self):
+    env = TestingEnvironment(
+        image_ids=[],
+        supports_dynamic_image_loading=True,
+        pool_size=2,
+        features={'test_feature': TestingFeature()},
+        outage_grace_period=1,
+        outage_retry_interval=0,
+    )
+    with env:
+      with env.sandbox(image_id='test_image'):
+        self.assertEqual(len(env.sandbox_pool['test_image']), 1)
+
+        with env.sandbox(image_id='test_image'):
+          self.assertEqual(len(env.sandbox_pool['test_image']), 2)
+
+          with self.assertRaises(interface.EnvironmentOverloadError):
+            with env.sandbox(image_id='test_image'):
+              pass
+        self.assertEqual(len(env.sandbox_pool['test_image']), 2)
+
+        with env.sandbox(image_id='test_image'):
+          self.assertEqual(len(env.sandbox_pool['test_image']), 2)
+
+  def test_image_feature_mappings(self):
+    env = TestingEnvironment(
+        image_ids=[
+            'test_image1',
+            'test_image2',
+        ],
+        features={
+            'test_feature': TestingFeature(
+                applicable_images=['test_image1.*']
+            ),
+            'test_feature2': TestingFeature(
+                applicable_images=['test_image2.*']
+            ),
+            'test_feature3': TestingFeature(
+                applicable_images=['test_image.*']
+            ),
+        },
+        pool_size=0,
+        outage_grace_period=1,
+        outage_retry_interval=0,
+        sandbox_keepalive_interval=0,
+    )
+    with env:
+      with env.sandbox(image_id='test_image1') as sb:
+        self.assertIn('test_feature', sb.features)
+        self.assertNotIn('test_feature2', sb.features)
+        self.assertIn('test_feature3', sb.features)
+
+      with env.sandbox(image_id='test_image2') as sb:
+        self.assertNotIn('test_feature', sb.features)
+        self.assertIn('test_feature2', sb.features)
+        self.assertIn('test_feature3', sb.features)
+
+      with env.test_feature() as feature:
+        self.assertEqual(feature.sandbox.image_id, 'test_image1')
+
+      with self.assertRaisesRegex(
+          ValueError, 'Feature .* is not applicable to .*'
+      ):
+        with env.test_feature('test_image2'):
+          pass
+
+      with env.test_feature2() as feature:
+        self.assertEqual(feature.sandbox.image_id, 'test_image2')
+
+      with env.test_feature3() as feature:
+        self.assertEqual(feature.sandbox.image_id, 'test_image1')
+
+      with env.test_feature3('test_image2') as feature:
+        self.assertEqual(feature.sandbox.image_id, 'test_image2')
+
+  def test_feature_applicability_check(self):
+    with self.assertRaisesRegex(
+        ValueError, 'Feature .* is not applicable to .*'
+    ):
+      TestingEnvironment(
+          image_ids=[
+              'test_image1',
+          ],
+          features={
+              'test_feature2': TestingFeature(
+                  applicable_images=['test_image2.*']
+              ),
+          },
+      )
+    env = TestingEnvironment(
+        image_ids=[],
+        supports_dynamic_image_loading=True,
+        features={
+            'test_feature2': TestingFeature(
+                applicable_images=['test_image2.*']
+            ),
+        },
+        pool_size=0
+    )
+    with env:
+      with self.assertRaisesRegex(
+          ValueError, 'No image ID found for feature .*'
+      ):
+        with env.test_feature2():
+          pass
+
+      # Dynamically loaded IDs.
+      with env.test_feature2('test_image2') as feature:
+        self.assertEqual(feature.sandbox.image_id, 'test_image2')
+
+  def test_pool_size(self):
+    env = TestingEnvironment(
+        image_ids=['test_image'],
+        pool_size=1,
+        outage_grace_period=1,
+        outage_retry_interval=0,
+    )
+    self.assertEqual(env.min_pool_size('test_image'), 1)
+    self.assertEqual(env.max_pool_size('test_image'), 1)
+
+    env = TestingEnvironment(
+        image_ids=['test_image'],
+        pool_size=(0, 256),
+        outage_grace_period=1,
+        outage_retry_interval=0,
+    )
+    self.assertEqual(env.min_pool_size('test_image'), 0)
+    self.assertEqual(env.max_pool_size('test_image'), 256)
+
+    env = TestingEnvironment(
+        image_ids=['test_image'],
+        pool_size={
+            'test_.*': (0, 128),
+            'my.*': (5, 64),
+            'exact_image_name': 10,
+        },
+        outage_grace_period=1,
+        outage_retry_interval=0,
+    )
+    self.assertEqual(env.min_pool_size('test_image'), 0)
+    self.assertEqual(env.max_pool_size('test_image'), 128)
+    self.assertEqual(env.min_pool_size('my_image'), 5)
+    self.assertEqual(env.max_pool_size('my_image'), 64)
+    self.assertEqual(env.min_pool_size('exact_image_name'), 10)
+    self.assertEqual(env.max_pool_size('exact_image_name'), 10)
+    self.assertEqual(env.min_pool_size('some_image'), 0)  # default
+    self.assertEqual(env.max_pool_size('some_image'), 256)  # default
+
+  def test_del(self):
+    env = TestingEnvironment(
+        features={'test_feature': TestingFeature()},
+        pool_size=0,
+        outage_grace_period=1,
+        outage_retry_interval=0,
+        sandbox_keepalive_interval=0,
+    )
+    env.start()
+    sb = env.acquire()
+    del sb
+    del env
+
+  def test_acquire_env_offline(self):
+    env = TestingEnvironment(
+        features={'test_feature': TestingFeature()},
+        pool_size=0,
+        outage_grace_period=1,
+        outage_retry_interval=0,
+        sandbox_keepalive_interval=0,
+    )
+    with self.assertRaises(interface.EnvironmentOutageError):
+      env.acquire()
+
+  def test_acquire_no_pooling(self):
+    env = TestingEnvironment(
+        features={'test_feature': TestingFeature()},
+        pool_size=0,
+        outage_grace_period=1,
+        outage_retry_interval=0,
+        sandbox_keepalive_interval=0,
+    )
+    with env:
+      sb = env.acquire()
+      self.assertEqual(sb.status, interface.Sandbox.Status.ACQUIRED)
+      self.assertIsNone(env.working_dir)
+      self.assertIsNone(sb.working_dir)
+      self.assertIsNone(sb.test_feature.working_dir)
+
+  def test_acquire_no_pooling_with_error(self):
+    env = TestingEnvironment(
+        features={
+            'test_feature': TestingFeature(
+                simulate_setup_error=interface.SandboxStateError
+            )
+        },
+        pool_size=0,
+        outage_grace_period=1,
+        outage_retry_interval=0,
+        sandbox_keepalive_interval=0,
+    )
+    with env:
+      with self.assertRaises(interface.EnvironmentOutageError):
+        env.acquire()
+
+  def test_acquire_with_pooling(self):
+    env = TestingEnvironment(
+        features={'test_feature': TestingFeature()},
+        pool_size=1,
+        outage_grace_period=1,
+        outage_retry_interval=0,
+        sandbox_keepalive_interval=0,
+    )
+    with env:
+      sb = env.acquire()
+      self.assertEqual(sb.status, interface.Sandbox.Status.ACQUIRED)
+
+  def test_acquire_with_pooling_overload(self):
+    env = TestingEnvironment(
+        features={'test_feature': TestingFeature()},
+        pool_size=1,
+        outage_grace_period=1,
+        outage_retry_interval=0,
+        sandbox_keepalive_interval=0,
+    )
+    with env:
+      sb = env.acquire()
+      self.assertEqual(sb.status, interface.Sandbox.Status.ACQUIRED)
+      with self.assertRaises(interface.EnvironmentOverloadError):
+        env.acquire()
+
+  def test_acquire_with_growing_pool(self):
+    env = TestingEnvironment(
+        features={'test_feature': TestingFeature()},
+        pool_size=(1, 3),
+        outage_grace_period=1,
+        outage_retry_interval=0,
+        sandbox_keepalive_interval=0,
+    )
+    with env:
+      self.assertEqual(len(env.sandbox_pool['test_image']), 1)
+      self.assertEqual(
+          env.stats(),
+          {
+              'sandbox': {
+                  'test_image': {
+                      'created': 0,
+                      'setting_up': 0,
+                      'ready': 1,
+                      'acquired': 0,
+                      'in_session': 0,
+                      'exiting_session': 0,
+                      'shutting_down': 0,
+                      'offline': 0,
+                  }
+              }
+          }
+      )
+      sb = env.acquire()
+      self.assertEqual(sb.status, interface.Sandbox.Status.ACQUIRED)
+      self.assertEqual(
+          env.stats(),
+          {
+              'sandbox': {
+                  'test_image': {
+                      'created': 0,
+                      'setting_up': 0,
+                      'ready': 0,
+                      'acquired': 1,
+                      'in_session': 0,
+                      'exiting_session': 0,
+                      'shutting_down': 0,
+                      'offline': 0,
+                  }
+              }
+          }
+      )
+      self.assertEqual(len(env.sandbox_pool['test_image']), 1)
+      sb2 = env.acquire()
+      self.assertEqual(sb2.status, interface.Sandbox.Status.ACQUIRED)
+      self.assertEqual(len(env.sandbox_pool['test_image']), 2)
+      self.assertEqual(
+          env.stats(),
+          {
+              'sandbox': {
+                  'test_image': {
+                      'created': 0,
+                      'setting_up': 0,
+                      'ready': 0,
+                      'acquired': 2,
+                      'in_session': 0,
+                      'exiting_session': 0,
+                      'shutting_down': 0,
+                      'offline': 0,
+                  }
+              }
+          }
+      )
+    self.assertEqual(
+        env.stats(),
+        {
+            'sandbox': {}
+        }
+    )
+
+  def test_acquire_with_growing_pool_failure(self):
+    env = TestingEnvironment(
+        features={'test_feature': TestingFeature()},
+        pool_size=(1, 3),
+        outage_grace_period=1,
+        outage_retry_interval=0,
+        sandbox_keepalive_interval=0,
+    )
+    with env:
+      self.assertEqual(len(env.sandbox_pool), 1)
+      sb = env.acquire()
+      self.assertEqual(sb.status, interface.Sandbox.Status.ACQUIRED)
+
+      # Make future sandbox setup to fail.
+      env.features.test_feature.rebind(
+          simulate_setup_error=interface.SandboxStateError,
+          skip_notification=True
+      )
+      with self.assertRaises(interface.EnvironmentOutageError):
+        env.acquire()
+
+  def test_housekeep_error(self):
+    env = TestingEnvironment(
+        features={'test_feature': TestingFeature()},
+        pool_size=1,
+        proactive_session_setup=True,
+        outage_grace_period=1,
+        outage_retry_interval=0,
+        sandbox_keepalive_interval=0,
+    )
+    with env:
+      self.assertEqual(len(env.sandbox_pool), 1)
+      self.assertIn('test_image', env.sandbox_pool)
+      self.assertEqual(
+          env.sandbox_pool['test_image'][0].status,
+          interface.Sandbox.Status.READY
+      )
+      # Make future sandbox setup to fail.
+      env.features.test_feature.rebind(
+          simulate_setup_error=interface.SandboxStateError,
+          skip_notification=True
+      )
+      with self.assertRaises(interface.SandboxStateError):
+        with env.sandbox() as sb:
+          sb.shell('bad command', raise_error=interface.SandboxStateError)
+      self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
+      self.assertEqual(len(sb.state_errors), 1)
+      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)
+
+
+if __name__ == '__main__':
+  unittest.main()

langfun/env/base_feature.py

@@ -110,6 +110,15 @@ class BaseFeature(interface.Feature):
     super()._on_parent_change(old_parent, new_parent)
     self.__dict__.pop('name', None)
 
+  @functools.cached_property
+  def environment(self) -> interface.Environment:
+    """Returns the environment that the feature is running in."""
+    if self._sandbox is not None:
+      return self._sandbox.environment
+    env = self.sym_ancestor(lambda v: isinstance(v, interface.Environment))
+    assert env is not None, 'Feature is not put into an environment.'
+    return env
+
   @property
   def sandbox(self) -> interface.Sandbox:
     """Returns the sandbox that the feature is running in."""
@@ -188,7 +197,13 @@ class BaseFeature(interface.Feature):
       error: BaseException | None = None
   ) -> None:
     """Called when the feature is setup."""
-    self.sandbox.on_feature_setup(self, duration, error)
+    self.environment.event_handler.on_feature_setup(
+        environment=self.environment,
+        sandbox=self.sandbox,
+        feature=self,
+        duration=duration,
+        error=error
+    )
 
   def on_teardown(
       self,
@@ -196,7 +211,13 @@ class BaseFeature(interface.Feature):
       error: BaseException | None = None
   ) -> None:
     """Called when the feature is teardown."""
-    self.sandbox.on_feature_teardown(self, duration, error)
+    self.environment.event_handler.on_feature_teardown(
+        environment=self.environment,
+        sandbox=self.sandbox,
+        feature=self,
+        duration=duration,
+        error=error
+    )
 
   def on_housekeep(
       self,
@@ -205,8 +226,14 @@ class BaseFeature(interface.Feature):
       **kwargs
   ) -> None:
     """Called when the feature has done housekeeping."""
-    self.sandbox.on_feature_housekeep(
-        self, self._housekeep_counter, duration, error, **kwargs
+    self.environment.event_handler.on_feature_housekeep(
+        environment=self.environment,
+        sandbox=self.sandbox,
+        feature=self,
+        counter=self._housekeep_counter,
+        duration=duration,
+        error=error,
+        **kwargs
     )
 
   def on_setup_session(
@@ -215,7 +242,14 @@ class BaseFeature(interface.Feature):
       error: BaseException | None = None,
   ) -> None:
     """Called when the feature is setup for a user session."""
-    self.sandbox.on_feature_setup_session(self, duration, error)
+    self.environment.event_handler.on_feature_setup_session(
+        environment=self.environment,
+        sandbox=self.sandbox,
+        feature=self,
+        session_id=self.session_id,
+        duration=duration,
+        error=error
+    )
 
   def on_teardown_session(
       self,
@@ -223,7 +257,14 @@ class BaseFeature(interface.Feature):
       error: BaseException | None = None,
   ) -> None:
     """Called when the feature is teardown for a user session."""
-    self.sandbox.on_feature_teardown_session(self, duration, error)
+    self.environment.event_handler.on_feature_teardown_session(
+        environment=self.environment,
+        sandbox=self.sandbox,
+        feature=self,
+        session_id=self.session_id,
+        duration=duration,
+        error=error
+    )
 
   def on_activity(
       self,
@@ -233,10 +274,13 @@ class BaseFeature(interface.Feature):
       **kwargs
   ) -> None:
     """Called when a sandbox activity is performed."""
-    self.sandbox.on_activity(
+    self.environment.event_handler.on_sandbox_activity(
         name=f'{self.name}.{name}',
+        environment=self.environment,
+        sandbox=self.sandbox,
         feature=self,
-        error=error,
+        session_id=self.session_id,
         duration=duration,
+        error=error,
         **kwargs
     )