langfun 0.1.2.dev202510310805__tar.gz → 0.1.2.dev202511020804__tar.gz

{langfun-0.1.2.dev202510310805 → langfun-0.1.2.dev202511020804}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langfun
-Version: 0.1.2.dev202510310805
+Version: 0.1.2.dev202511020804
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.2.dev202510310805 → langfun-0.1.2.dev202511020804}/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-0.1.2.dev202510310805 → langfun-0.1.2.dev202511020804}/langfun/env/__init__.py

@@ -24,9 +24,11 @@ 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
+from langfun.env.interface import log_activity
 
 from langfun.env.base_environment import BaseEnvironment
 from langfun.env.base_sandbox import BaseSandbox
@@ -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-0.1.2.dev202510310805 → langfun-0.1.2.dev202511020804}/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)
     )
@@ -178,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()
@@ -191,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: '
@@ -266,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):
@@ -311,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():
@@ -322,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,
@@ -555,8 +575,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()
@@ -636,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():
@@ -659,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,
@@ -667,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(
@@ -675,6 +730,7 @@ class BaseEnvironment(interface.Environment):
             ),
             **kwargs
         )
+        self.shutdown()
 
   def _replace_dead_sandboxes(
       self,
@@ -738,16 +794,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 +810,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 +824,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)