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

langfun/core/eval/v2/runners_test.py

@@ -100,6 +100,9 @@ class RunnerTest(unittest.TestCase):
         print(i, pg.diff(x, y))
       self.assertIs(x, y)
 
+
+class SequentialRunnerTest(RunnerTest):
+
   def test_basic(self):
     plugin = TestPlugin()
     exp = eval_test_helper.test_experiment()

langfun/env/__init__.py

@@ -28,7 +28,7 @@ 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
+from langfun.env.interface import log_activity
 
 from langfun.env.base_environment import BaseEnvironment
 from langfun.env.base_sandbox import BaseSandbox

langfun/env/base_environment.py

@@ -179,6 +179,7 @@ class BaseEnvironment(interface.Environment):
     )
     self._housekeep_thread = None
     self._offline_start_time = None
+    self._non_sandbox_based_features_with_setup_called = set()
 
     # Check image IDs and feature requirements.
     self._check_image_ids()
@@ -192,7 +193,9 @@ class BaseEnvironment(interface.Environment):
     if self.supports_dynamic_image_loading:
       return
     for name, feature in self.features.items():
-      if any(feature.is_applicable(image_id) for image_id in self.image_ids):
+      if not feature.is_sandbox_based or any(
+          feature.is_applicable(image_id) for image_id in self.image_ids
+      ):
         continue
       raise ValueError(
           f'Feature {name!r} is not applicable to all available images: '
@@ -267,6 +270,13 @@ class BaseEnvironment(interface.Environment):
   def _start(self) -> None:
     """Implementation of starting the environment."""
     sandbox_startup_infos = []
+    self._non_sandbox_based_features_with_setup_called.clear()
+    # Setup all non-sandbox-based features.
+    for feature in self.non_sandbox_based_features():
+      self._non_sandbox_based_features_with_setup_called.add(feature.name)
+      feature.setup(sandbox=None)
+
+    # Setup sandbox pools.
     for image_id in self.image_ids:
       next_sandbox_id = 0
       if self.enable_pooling(image_id):
@@ -312,10 +322,15 @@ class BaseEnvironment(interface.Environment):
       self._housekeep_thread.join()
       self._housekeep_thread = None
 
-    def _shutdown_sandbox(sandbox: base_sandbox.BaseSandbox) -> None:
-      if sandbox is not None:
-        sandbox.shutdown()
+    # Teardown all non-sandbox-based features.
+    for feature in self.non_sandbox_based_features():
+      if feature.name in self._non_sandbox_based_features_with_setup_called:
+        try:
+          feature.teardown()
+        except BaseException:   # pylint: disable=broad-except
+          pass
 
+    # Shutdown sandbox pools.
     if self._sandbox_pool:
       sandboxes = []
       for sandbox in self._sandbox_pool.values():
@@ -323,6 +338,10 @@ class BaseEnvironment(interface.Environment):
       self._sandbox_pool = {}
 
       if sandboxes:
+        def _shutdown_sandbox(sandbox: base_sandbox.BaseSandbox) -> None:
+          if sandbox is not None:
+            sandbox.shutdown()
+
         _ = list(
             lf.concurrent_map(
                 _shutdown_sandbox,
@@ -635,9 +654,45 @@ class BaseEnvironment(interface.Environment):
         indices_by_image_id[image_id].append(i)
       return indices_by_image_id
 
+    last_housekeep_time = {
+        f.name: time.time() for f in self.non_sandbox_based_features()
+    }
+
     while self._status not in (self.Status.SHUTTING_DOWN, self.Status.OFFLINE):
       housekeep_start_time = time.time()
+      feature_housekeep_successes = []
+      feature_housekeep_failures = []
+
+      # Housekeeping non-sandbox-based features.
+      for feature in self.non_sandbox_based_features():
+        if feature.housekeep_interval is None:
+          continue
+        if (last_housekeep_time[feature.name]
+            + feature.housekeep_interval < time.time()):
+          try:
+            feature.housekeep()
+            last_housekeep_time[feature.name] = time.time()
+            feature_housekeep_successes.append(feature.name)
+          except BaseException as e:  # pylint: disable=broad-except
+            pg.logging.error(
+                '[%s/%s]: Feature housekeeping failed with error: %s.'
+                'Shutting down environment...',
+                self.id,
+                feature.name,
+                e,
+            )
+            feature_housekeep_failures.append(feature.name)
+            self._housekeep_counter += 1
+            self.on_housekeep(
+                duration=time.time() - housekeep_start_time,
+                error=e,
+                feature_housekeep_successes=feature_housekeep_successes,
+                feature_housekeep_failures=feature_housekeep_failures,
+            )
+            self.shutdown()
+            return
 
+      # Replace dead sandboxes.
       is_online = True
       dead_sandbox_entries = []
       for image_id, sandboxes in self._sandbox_pool.items():
@@ -658,6 +713,8 @@ class BaseEnvironment(interface.Environment):
       duration = time.time() - housekeep_start_time
 
       kwargs = dict(
+          feature_housekeep_successes=feature_housekeep_successes,
+          feature_housekeep_failures=feature_housekeep_failures,
           dead_sandboxes=_indices_by_image_id(dead_sandbox_entries),
           replaced_sandboxes=replaced_indices_by_image_id,
           offline_duration=self.offline_duration,
@@ -666,7 +723,6 @@ class BaseEnvironment(interface.Environment):
         self.on_housekeep(duration, **kwargs)
         time.sleep(self.housekeep_interval)
       else:
-        self.shutdown()
         self.on_housekeep(
             duration,
             interface.EnvironmentOutageError(
@@ -674,6 +730,7 @@ class BaseEnvironment(interface.Environment):
             ),
             **kwargs
         )
+        self.shutdown()
 
   def _replace_dead_sandboxes(
       self,

langfun/env/base_environment_test.py

@@ -57,7 +57,7 @@ class BaseEnvironmentTests(unittest.TestCase):
       self.assertEqual(env.sandbox_pool, {})
       self.assertEqual(env.working_dir, '/tmp/testing-env')
 
-      with env.sandbox(session_id='session1') as sb:
+      with env.sandbox('session1') as sb:
         self.assertEqual(
             sb.id, interface.Sandbox.Id(
                 environment_id=env.id,
@@ -79,7 +79,7 @@ class BaseEnvironmentTests(unittest.TestCase):
       with self.assertRaisesRegex(
           ValueError, 'Environment .* does not serve image ID .*'
       ):
-        env.sandbox('test_image2')
+        env.sandbox(image_id='test_image2')
 
       with env.test_feature() as feature:
         self.assertIsInstance(feature, TestingFeature)
@@ -174,7 +174,7 @@ class BaseEnvironmentTests(unittest.TestCase):
       with self.assertRaisesRegex(
           ValueError, 'Feature .* is not applicable to .*'
       ):
-        with env.test_feature('test_image2'):
+        with env.test_feature(image_id='test_image2'):
           pass
 
       with env.test_feature2() as feature:
@@ -183,7 +183,7 @@ class BaseEnvironmentTests(unittest.TestCase):
       with env.test_feature3() as feature:
         self.assertEqual(feature.sandbox.image_id, 'test_image1')
 
-      with env.test_feature3('test_image2') as feature:
+      with env.test_feature3(image_id='test_image2') as feature:
         self.assertEqual(feature.sandbox.image_id, 'test_image2')
 
   def test_feature_applicability_check(self):
@@ -218,7 +218,7 @@ class BaseEnvironmentTests(unittest.TestCase):
           pass
 
       # Dynamically loaded IDs.
-      with env.test_feature2('test_image2') as feature:
+      with env.test_feature2(image_id='test_image2') as feature:
         self.assertEqual(feature.sandbox.image_id, 'test_image2')
 
   def test_pool_size(self):

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,18 +22,24 @@ 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],
@@ -120,9 +126,11 @@ class BaseFeature(interface.Feature):
     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
@@ -159,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)
@@ -198,8 +206,6 @@ class BaseFeature(interface.Feature):
   ) -> None:
     """Called when the feature is setup."""
     self.environment.event_handler.on_feature_setup(
-        environment=self.environment,
-        sandbox=self.sandbox,
         feature=self,
         duration=duration,
         error=error
@@ -212,8 +218,6 @@ class BaseFeature(interface.Feature):
   ) -> None:
     """Called when the feature is teardown."""
     self.environment.event_handler.on_feature_teardown(
-        environment=self.environment,
-        sandbox=self.sandbox,
         feature=self,
         duration=duration,
         error=error
@@ -227,8 +231,6 @@ class BaseFeature(interface.Feature):
   ) -> None:
     """Called when the feature has done housekeeping."""
     self.environment.event_handler.on_feature_housekeep(
-        environment=self.environment,
-        sandbox=self.sandbox,
         feature=self,
         counter=self._housekeep_counter,
         duration=duration,
@@ -243,8 +245,6 @@ class BaseFeature(interface.Feature):
   ) -> None:
     """Called when the feature is setup for a user session."""
     self.environment.event_handler.on_feature_setup_session(
-        environment=self.environment,
-        sandbox=self.sandbox,
         feature=self,
         session_id=self.session_id,
         duration=duration,
@@ -258,8 +258,6 @@ class BaseFeature(interface.Feature):
   ) -> None:
     """Called when the feature is teardown for a user session."""
     self.environment.event_handler.on_feature_teardown_session(
-        environment=self.environment,
-        sandbox=self.sandbox,
         feature=self,
         session_id=self.session_id,
         duration=duration,
@@ -274,13 +272,33 @@ class BaseFeature(interface.Feature):
       **kwargs
   ) -> None:
     """Called when a sandbox activity is performed."""
-    self.environment.event_handler.on_sandbox_activity(
+    self.environment.event_handler.on_feature_activity(
         name=f'{self.name}.{name}',
-        environment=self.environment,
-        sandbox=self.sandbox,
         feature=self,
         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
+      )

langfun/env/base_feature_test.py

@@ -0,0 +1,228 @@
+# 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 unittest
+
+from langfun.env import test_utils
+
+TestingEnvironment = test_utils.TestingEnvironment
+TestingNonSandboxBasedFeature = test_utils.TestingNonSandboxBasedFeature
+TestingEventHandler = test_utils.TestingEventHandler
+
+
+class NonSandboxBasedFeatureTests(unittest.TestCase):
+
+  def test_basics(self):
+    feature = TestingNonSandboxBasedFeature()
+    event_handler = TestingEventHandler(
+        log_session_setup=True,
+        log_feature_setup=True,
+        log_sandbox_status=True
+    )
+    env = TestingEnvironment(
+        image_ids=[],
+        features={'test_feature': feature},
+        event_handler=event_handler,
+    )
+    self.assertFalse(env.is_online)
+    self.assertEqual(len(list(env.non_sandbox_based_features())), 1)
+    with env:
+      self.assertTrue(env.is_online)
+      with env.test_feature('session1') as feature:
+        self.assertIsNone(feature.sandbox)
+        self.assertEqual(feature.session_id, 'session1')
+
+    self.assertEqual(
+        event_handler.logs,
+        [
+            '[testing-env/test_feature] feature setup',
+            '[testing-env] environment started',
+            '[testing-env/test_feature@session1] feature setup session',
+            '[testing-env/test_feature@session1] feature teardown session',
+            '[testing-env/test_feature] feature teardown',
+            '[testing-env] environment shutdown'
+        ]
+    )
+
+  def test_feature_setup_error(self):
+    event_handler = TestingEventHandler(
+        log_session_setup=True,
+        log_feature_setup=True,
+        log_sandbox_status=True
+    )
+    env = TestingEnvironment(
+        image_ids=[],
+        features={
+            'test_feature': TestingNonSandboxBasedFeature(
+                simulate_setup_error=ValueError
+            )
+        },
+        event_handler=event_handler,
+    )
+    with self.assertRaises(ValueError):
+      with env:
+        pass
+    self.assertEqual(
+        event_handler.logs,
+        [
+            '[testing-env/test_feature] feature setup with ValueError',
+            '[testing-env] environment started with ValueError',
+            '[testing-env/test_feature] feature teardown',
+            '[testing-env] environment shutdown'
+        ]
+    )
+
+  def test_feature_teardown_error(self):
+    event_handler = TestingEventHandler(
+        log_session_setup=True,
+        log_feature_setup=True,
+        log_sandbox_status=True
+    )
+    env = TestingEnvironment(
+        image_ids=[],
+        features={
+            'test_feature': TestingNonSandboxBasedFeature(
+                simulate_teardown_error=ValueError
+            )
+        },
+        event_handler=event_handler,
+    )
+    with env:
+      pass
+    self.assertEqual(
+        event_handler.logs,
+        [
+            '[testing-env/test_feature] feature setup',
+            '[testing-env] environment started',
+            '[testing-env/test_feature] feature teardown with ValueError',
+            '[testing-env] environment shutdown'
+        ]
+    )
+
+  def test_feature_setup_session_error(self):
+    event_handler = TestingEventHandler(
+        log_session_setup=True,
+        log_feature_setup=True,
+        log_sandbox_status=True
+    )
+    env = TestingEnvironment(
+        image_ids=[],
+        features={
+            'test_feature': TestingNonSandboxBasedFeature(
+                simulate_setup_session_error=ValueError
+            )
+        },
+        event_handler=event_handler,
+    )
+    with env:
+      with self.assertRaises(ValueError):
+        with env.test_feature('session1'):
+          pass
+    self.assertEqual(
+        event_handler.logs,
+        [
+            # pylint: disable=line-too-long
+            '[testing-env/test_feature] feature setup',
+            '[testing-env] environment started',
+            '[testing-env/test_feature@session1] feature setup session with ValueError',
+            '[testing-env/test_feature@session1] feature teardown session',
+            '[testing-env/test_feature] feature teardown',
+            '[testing-env] environment shutdown',
+            # pylint: enable=line-too-long
+        ]
+    )
+
+  def test_feature_teardown_session_error(self):
+    event_handler = TestingEventHandler(
+        log_session_setup=True,
+        log_feature_setup=True,
+        log_sandbox_status=True
+    )
+    env = TestingEnvironment(
+        image_ids=[],
+        features={
+            'test_feature': TestingNonSandboxBasedFeature(
+                simulate_teardown_session_error=ValueError
+            )
+        },
+        event_handler=event_handler,
+    )
+    with env:
+      with env.test_feature('session1'):
+        pass
+    self.assertEqual(
+        event_handler.logs,
+        [
+            # pylint: disable=line-too-long
+            '[testing-env/test_feature] feature setup',
+            '[testing-env] environment started',
+            '[testing-env/test_feature@session1] feature setup session',
+            '[testing-env/test_feature@session1] feature teardown session with ValueError',
+            '[testing-env/test_feature] feature teardown',
+            '[testing-env] environment shutdown',
+            # pylint: enable=line-too-long
+        ]
+    )
+
+  def test_feature_housekeeping(self):
+    event_handler = TestingEventHandler(
+        log_sandbox_status=False,
+        log_feature_setup=False,
+        log_housekeep=True
+    )
+    env = TestingEnvironment(
+        image_ids=[],
+        features={
+            'test_feature': TestingNonSandboxBasedFeature(
+                housekeep_interval=0.1
+            )
+        },
+        event_handler=event_handler,
+        housekeep_interval=0.2
+    )
+    with env:
+      env.wait_for_housekeeping()
+    self.assertIn(
+        '[testing-env/test_feature] feature housekeeping 0',
+        event_handler.logs
+    )
+
+  def test_feature_housekeeping_error(self):
+    event_handler = TestingEventHandler(
+        log_sandbox_status=False,
+        log_feature_setup=False,
+        log_housekeep=True
+    )
+    env = TestingEnvironment(
+        image_ids=[],
+        features={
+            'test_feature': TestingNonSandboxBasedFeature(
+                simulate_housekeep_error=ValueError,
+                housekeep_interval=0.1
+            )
+        },
+        event_handler=event_handler,
+        housekeep_interval=0.2
+    )
+    with env:
+      env.wait_for_housekeeping()
+      self.assertFalse(env.is_online)
+    self.assertIn(
+        '[testing-env/test_feature] feature housekeeping 0 with ValueError',
+        event_handler.logs
+    )
+
+
+if __name__ == '__main__':
+  unittest.main()
+

langfun/env/base_sandbox.py

@@ -648,7 +648,6 @@ class BaseSandbox(interface.Sandbox):
   def track_activity(
       self,
       name: str,
-      feature: interface.Feature | None = None,
       **kwargs: Any
   ) -> Iterator[None]:
     """Tracks an activity for the sandbox."""
@@ -662,7 +661,6 @@ class BaseSandbox(interface.Sandbox):
     finally:
       self.on_activity(
           name=name,
-          feature=feature,
           duration=time.time() - start_time,
           error=error,
           **kwargs
@@ -755,9 +753,7 @@ class BaseSandbox(interface.Sandbox):
       error: BaseException | None = None
   ) -> None:
     """Called when the sandbox is started."""
-    self._event_handler.on_sandbox_start(
-        self.environment, self, duration, error
-    )
+    self._event_handler.on_sandbox_start(self, duration, error)
 
   def on_status_change(
       self,
@@ -767,7 +763,7 @@ class BaseSandbox(interface.Sandbox):
     """Called when the sandbox status changes."""
     status_duration = time.time() - self._status_start_time
     self._event_handler.on_sandbox_status_change(
-        self.environment, self, old_status, new_status, status_duration
+        self, old_status, new_status, status_duration
     )
 
   def on_shutdown(
@@ -777,11 +773,11 @@ class BaseSandbox(interface.Sandbox):
   ) -> None:
     """Called when the sandbox is shutdown."""
     self._event_handler.on_sandbox_shutdown(
-        self.environment,
-        self,
-        duration,
-        0.0 if self._start_time is None else (time.time() - self._start_time),
-        error
+        sandbox=self,
+        duration=duration,
+        lifetime=(0.0 if self._start_time is None
+                  else (time.time() - self._start_time)),
+        error=error
     )
 
   def on_housekeep(
@@ -792,11 +788,10 @@ class BaseSandbox(interface.Sandbox):
   ) -> None:
     """Called when the sandbox finishes a round of housekeeping."""
     self._event_handler.on_sandbox_housekeep(
-        self.environment,
-        self,
-        self._housekeep_counter,
-        duration,
-        error,
+        sandbox=self,
+        counter=self._housekeep_counter,
+        duration=duration,
+        error=error,
         **kwargs
     )
 
@@ -807,27 +802,27 @@ class BaseSandbox(interface.Sandbox):
       error: BaseException | None = None
   ) -> None:
     """Called when the user session starts."""
-    self._event_handler.on_session_start(
-        self.environment, self, session_id, duration, error
+    self._event_handler.on_sandbox_session_start(
+        sandbox=self,
+        session_id=session_id,
+        duration=duration,
+        error=error
     )
 
   def on_activity(
       self,
       name: str,
       duration: float,
-      feature: interface.Feature | None = None,
       error: BaseException | None = None,
       **kwargs
   ) -> None:
     """Called when a sandbox activity is performed."""
     self._event_handler.on_sandbox_activity(
-        name,
-        self.environment,
-        self,
-        feature,
-        self.session_id,
-        duration,
-        error,
+        name=name,
+        sandbox=self,
+        session_id=self.session_id,
+        duration=duration,
+        error=error,
         **kwargs
     )
 
@@ -838,11 +833,10 @@ class BaseSandbox(interface.Sandbox):
       error: BaseException | None = None
   ) -> None:
     """Called when the user session ends."""
-    self._event_handler.on_session_end(
-        self.environment,
-        self,
-        session_id,
-        duration,
-        time.time() - self._session_start_time,
-        error
+    self._event_handler.on_sandbox_session_end(
+        sandbox=self,
+        session_id=session_id,
+        duration=duration,
+        lifetime=time.time() - self._session_start_time,
+        error=error
     )