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

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('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(image_id='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(image_id='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(image_id='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(image_id='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

@@ -11,9 +11,9 @@
 # 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.
-"""Common base class for sandbox-based features.
+"""Common base class for environment features.
 
-This module provides an base class `BaseFeature` for sandbox-based features,
+This module provides an base class `BaseFeature` for environment features,
 which provides event handlers for the feature lifecycle events, which can be
 overridden by subclasses to provide custom behaviors. Please note that this base
 class is intended to provide a convenient way to implement features, and not
@@ -22,17 +22,32 @@ coupled with `BaseEnvironment` and `BaseSandbox`, and is expected to work with
 the `Environment` and `Sandbox` interfaces directly.
 """
 
+import contextlib
 import functools
 import os
+import re
 import time
-from typing import Annotated, Callable
+from typing import Annotated, Any, Callable, Iterator
 
 from langfun.env import interface
 import pyglove as pg
 
 
 class BaseFeature(interface.Feature):
-  """Common base class for sandbox-based features."""
+  """Common base class for environment features."""
+
+  is_sandbox_based: Annotated[
+      bool,
+      'Whether the feature is sandbox-based.'
+  ] = True
+
+  applicable_images: Annotated[
+      list[str],
+      (
+          'A list of regular expressions for image IDs which enable '
+          'this feature. By default, all images are enabled.'
+      )
+  ] = ['.*']
 
   housekeep_interval: Annotated[
       float | None,
@@ -101,10 +116,21 @@ 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:
+  def sandbox(self) -> interface.Sandbox | None:
     """Returns the sandbox that the feature is running in."""
-    assert self._sandbox is not None, 'Feature has not been set up yet.'
+    assert self._sandbox is not None or not self.is_sandbox_based, (
+        'Feature has not been set up yet.'
+    )
     return self._sandbox
 
   @property
@@ -115,6 +141,12 @@ class BaseFeature(interface.Feature):
       return None
     return os.path.join(sandbox_workdir, self.name)
 
+  def is_applicable(self, image_id: str) -> bool:
+    """Returns True if the feature is applicable to the given image."""
+    return any(
+        re.fullmatch(regex, image_id) for regex in self.applicable_images
+    )
+
   #
   # Setup and teardown of the feature.
   #
@@ -135,7 +167,7 @@ class BaseFeature(interface.Feature):
     finally:
       event_handler(duration=time.time() - start_time, error=error)
 
-  def setup(self, sandbox: interface.Sandbox) -> None:
+  def setup(self, sandbox: interface.Sandbox | None = None) -> None:
     """Sets up the feature."""
     self._sandbox = sandbox
     self._do(self._setup, self.on_setup)
@@ -173,7 +205,11 @@ 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(
+        feature=self,
+        duration=duration,
+        error=error
+    )
 
   def on_teardown(
       self,
@@ -181,7 +217,11 @@ 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(
+        feature=self,
+        duration=duration,
+        error=error
+    )
 
   def on_housekeep(
       self,
@@ -190,8 +230,12 @@ 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(
+        feature=self,
+        counter=self._housekeep_counter,
+        duration=duration,
+        error=error,
+        **kwargs
     )
 
   def on_setup_session(
@@ -200,7 +244,12 @@ 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(
+        feature=self,
+        session_id=self.session_id,
+        duration=duration,
+        error=error
+    )
 
   def on_teardown_session(
       self,
@@ -208,7 +257,12 @@ 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(
+        feature=self,
+        session_id=self.session_id,
+        duration=duration,
+        error=error
+    )
 
   def on_activity(
       self,
@@ -218,10 +272,33 @@ class BaseFeature(interface.Feature):
       **kwargs
   ) -> None:
     """Called when a sandbox activity is performed."""
-    self.sandbox.on_activity(
+    self.environment.event_handler.on_feature_activity(
         name=f'{self.name}.{name}',
         feature=self,
-        error=error,
+        session_id=self.session_id,
         duration=duration,
+        error=error,
         **kwargs
     )
+
+  @contextlib.contextmanager
+  def track_activity(
+      self,
+      name: str,
+      **kwargs: Any
+  ) -> Iterator[None]:
+    """Context manager that tracks a feature activity."""
+    start_time = time.time()
+    error = None
+    try:
+      yield None
+    except BaseException as e:  # pylint: disable=broad-except
+      error = e
+      raise
+    finally:
+      self.on_activity(
+          name=name,
+          duration=time.time() - start_time,
+          error=error,
+          **kwargs
+      )