langfun 0.1.2.dev202509240805__py3-none-any.whl → 0.1.2.dev202509260805__py3-none-any.whl

langfun/env/__init__.py

@@ -13,16 +13,12 @@
 # limitations under the License.
 """Environment for LLM agents."""
 
-# pylint: disable=g-importing-member, g-bad-import-order
-from langfun.env.interface import EnvironmentId
-from langfun.env.interface import SandboxId
-
+# pylint: disable=g-importing-member, g-bad-import-order, g-import-not-at-top
 from langfun.env.interface import EnvironmentError   # pylint: disable=redefined-builtin
 from langfun.env.interface import EnvironmentOutageError
 from langfun.env.interface import EnvironmentOverloadError
 from langfun.env.interface import SandboxError
 from langfun.env.interface import SandboxStateError
-# from langfun.env.base import SandboxOverloadError
 
 from langfun.env.interface import Environment
 from langfun.env.interface import Sandbox
@@ -35,4 +31,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
+
 # Google-internal imports.

langfun/env/base_environment.py

@@ -34,6 +34,7 @@ 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
 
 
@@ -60,15 +61,25 @@ class BaseEnvironment(interface.Environment):
           'will be used as both min and max size. If 0, sandboxes will be '
           'created on demand and shutdown when user session ends.'
       )
-  ] = 1
+  ] = (0, 256)
 
   load_balancer: Annotated[
       load_balancers.LoadBalancer,
       (
-          'The load balancer for the environment.'
+          'The load balancer for the environment to acquire sandboxes.'
       )
   ] = load_balancers.RoundRobin()
 
+  sandbox_keepalive_interval: Annotated[
+      float | None,
+      (
+          'The interval in seconds to send keepalive pings to sandboxes. '
+          'If None, sandbox keepalive is disabled. Please note that sandbox '
+          'keepalive is different from feature housekeeping. Usually sandbox '
+          'keepalive and feature housekeeping are different operations.'
+      )
+  ] = None
+
   proactive_session_setup: Annotated[
       bool,
       (
@@ -79,6 +90,13 @@ class BaseEnvironment(interface.Environment):
       )
   ] = True
 
+  event_handlers: Annotated[
+      list[event_handler_base.EventHandler],
+      (
+          'User handler for the environment events.'
+      )
+  ] = []
+
   outage_grace_period: Annotated[
       float,
       (
@@ -97,13 +115,15 @@ class BaseEnvironment(interface.Environment):
       )
   ] = 10.0
 
-  stats_report_interval: Annotated[
+  housekeep_interval: Annotated[
       float,
       (
-          'The interval in seconds for reporting the environment stats. '
-          'If 0, stats will not be reported.'
+          'The interval in seconds for environment housekeeping. It recycles '
+          'the dead sandboxes in the pool. This interval is the minimal time '
+          'to detect outage while there is no request to obtain new sandboxes.'
+          'This is applicable only when the environment enables pooling.'
       )
-  ] = 60.0
+  ] = 10.0
 
   pool_operation_max_parallelism: Annotated[
       int,
@@ -132,7 +152,7 @@ class BaseEnvironment(interface.Environment):
         random if self.random_seed is None else random.Random(self.random_seed)
     )
 
-    self._maintenance_thread = None
+    self._housekeep_thread = None
     self._offline_start_time = None
 
   #
@@ -145,6 +165,7 @@ class BaseEnvironment(interface.Environment):
       sandbox_id: str,
       reusable: bool,
       proactive_session_setup: bool,
+      keepalive_interval: float | None,
   ) -> base_sandbox.BaseSandbox:
     """Creates a sandbox with the given identifier.
 
@@ -153,6 +174,8 @@ class BaseEnvironment(interface.Environment):
       reusable: Whether the sandbox is reusable across user sessions.
       proactive_session_setup: Whether the sandbox performs session setup work
         before a user session is started.
+      keepalive_interval: Interval to ping the sandbox for keeping it alive.
+        If None, the sandbox will not be pinged.
 
     Returns:
       The created sandbox.
@@ -170,6 +193,11 @@ class BaseEnvironment(interface.Environment):
     ).hex[:7]
     return f'session-{suffix}'
 
+  @property
+  def housekeep_counter(self) -> int:
+    """Returns the housekeeping counter."""
+    return self._housekeep_counter
+
   #
   # Subclasses can override:
   #
@@ -189,36 +217,40 @@ class BaseEnvironment(interface.Environment):
   def _start(self) -> None:
     """Implementation of starting the environment."""
     if self.min_pool_size > 0:
-      self._sandbox_pool = [
-          sandbox
-          for _, sandbox, _ in lf.concurrent_map(
-              lambda i: self._bring_up_sandbox_with_retry(
-                  sandbox_id=str(i), shutdown_env_upon_outage=False
-              ),
-              range(self.min_pool_size),
-              silence_on_errors=None,
-              max_workers=min(
-                  self.pool_operation_max_parallelism,
-                  self.min_pool_size
-              ),
-          )
-      ]
+      # Pre-allocate the sandbox pool before usage.
+      self._sandbox_pool = [None] * self.min_pool_size
+      for i, sandbox, _ in lf.concurrent_map(
+          lambda i: self._bring_up_sandbox_with_retry(
+              sandbox_id=str(i), shutdown_env_upon_outage=False
+          ),
+          range(self.min_pool_size),
+          silence_on_errors=None,
+          max_workers=min(
+              self.pool_operation_max_parallelism,
+              self.min_pool_size
+          ),
+      ):
+        self._sandbox_pool[i] = sandbox
+
     self._next_sandbox_id = len(self._sandbox_pool)
-    self._maintenance_thread = threading.Thread(
-        target=self._maintenance_loop, daemon=True
-    )
-    self._maintenance_count = 0
-    self._maintenance_thread.start()
+
+    if self.enable_pooling:
+      self._housekeep_thread = threading.Thread(
+          target=self._housekeep_loop, daemon=True
+      )
+      self._housekeep_counter = 0
+      self._housekeep_thread.start()
 
   def _shutdown(self) -> None:
     """Implementation of shutting down the environment."""
-    if (self._maintenance_thread is not None
-        and threading.current_thread() is not self._maintenance_thread):
-      self._maintenance_thread.join()
-      self._maintenance_thread = None
+    if (self._housekeep_thread is not None
+        and threading.current_thread() is not self._housekeep_thread):
+      self._housekeep_thread.join()
+      self._housekeep_thread = None
 
     def _shutdown_sandbox(sandbox: base_sandbox.BaseSandbox) -> None:
-      sandbox.shutdown()
+      if sandbox is not None:
+        sandbox.shutdown()
 
     if self._sandbox_pool:
       _ = list(
@@ -251,7 +283,7 @@ class BaseEnvironment(interface.Environment):
   @property
   def enable_pooling(self) -> bool:
     """Returns whether the environment enables pooling."""
-    return self.min_pool_size > 0
+    return self.max_pool_size > 0
 
   @property
   def status(self) -> interface.Environment.Status:
@@ -303,18 +335,14 @@ class BaseEnvironment(interface.Environment):
         f'it is in {self._status.value!r} status.'
     )
 
+    starting_time = time.time()
     try:
-      with pg.timeit('env.start') as t:
-        self._start()
-        self._start_time = time.time()
-      pg.logging.info(
-          '[%s]: %s started in %.2f seconds.',
-          self.id, self.__class__.__name__, t.elapse
-      )
+      self._start()
+      self._start_time = time.time()
       self._set_status(self.Status.ONLINE)
-      self.on_start()
+      self.on_start(duration=time.time() - starting_time)
     except BaseException as e:
-      self.on_start(error=e)
+      self.on_start(duration=time.time() - starting_time, error=e)
       self.shutdown()
       raise e
 
@@ -332,13 +360,8 @@ class BaseEnvironment(interface.Environment):
     self._set_status(self.Status.SHUTTING_DOWN)
 
     try:
-      with pg.timeit('env.shutdown') as t:
-        self._shutdown()
+      self._shutdown()
       self.on_shutdown()
-      pg.logging.info(
-          '[%s]: %s shutdown in %.2f seconds.',
-          self.id, self.__class__.__name__, t.elapse
-      )
     except BaseException as e:  # pylint: disable=broad-except
       self.on_shutdown(error=e)
       raise e
@@ -410,6 +433,7 @@ class BaseEnvironment(interface.Environment):
         sandbox_id=sandbox_id,
         reusable=self.enable_pooling,
         proactive_session_setup=self.proactive_session_setup,
+        keepalive_interval=self.sandbox_keepalive_interval,
     )
     for handler in self.event_handlers:
       sandbox.add_event_handler(handler)
@@ -474,20 +498,10 @@ class BaseEnvironment(interface.Environment):
   # Environment maintenance loop.
   #
 
-  def _maintenance_loop(self) -> None:
-    """Maintains the server pool."""
-    pg.logging.info(
-        '[%s]: %s maintenance thread started.', self.id, self.__class__.__name__
-    )
-    stats_report_time = time.time()
+  def _housekeep_loop(self) -> None:
+    """Housekeeping loop for the environment."""
     while self._status not in (self.Status.SHUTTING_DOWN, self.Status.OFFLINE):
-      if time.time() - stats_report_time > self.stats_report_interval:
-        pg.logging.info(
-            '[%s] %s stats: %s.',
-            self.id, self.__class__.__name__, self.stats()
-        )
-        stats_report_time = time.time()
-
+      housekeep_start_time = time.time()
       dead_pool_indices = [
           i for i, s in enumerate(self._sandbox_pool)
           if s.status == interface.Sandbox.Status.OFFLINE
@@ -496,10 +510,18 @@ class BaseEnvironment(interface.Environment):
           dead_pool_indices
       ):
         self.shutdown()
-        self._maintenance_count += 1
+        self._housekeep_counter += 1
+        self.on_housekeep(
+            time.time() - housekeep_start_time,
+            interface.EnvironmentOutageError(
+                environment=self,
+                offline_duration=self.offline_duration
+            )
+        )
         break
-      self._maintenance_count += 1
-      time.sleep(1)
+      self._housekeep_counter += 1
+      self.on_housekeep(time.time() - housekeep_start_time)
+      time.sleep(self.housekeep_interval)
 
   def _replace_dead_sandboxes(self, dead_pool_indices: list[int]) -> bool:
     """Replaces a dead sandbox with a new one.
@@ -526,6 +548,8 @@ class BaseEnvironment(interface.Environment):
           str(i), shutdown_env_upon_outage=False
       )
 
+    # TODO(daiyip): Consider to loose the condition to allow some dead
+    # sandboxes to be replaced successfully.
     return not any([
         error for _, _, error in lf.concurrent_map(
             _replace, dead_pool_indices,
@@ -540,12 +564,26 @@ class BaseEnvironment(interface.Environment):
   # Event handlers subclasses can override.
   #
 
-  def on_start(self, error: BaseException | None = None) -> None:
+  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, error)
+      handler.on_environment_start(self, duration, error)
+
+  def on_housekeep(
+      self,
+      duration: float,
+      error: BaseException | None = None
+  ) -> 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)
 
   def on_shutdown(self, 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, error)
+      handler.on_environment_shutdown(self, lifetime, error)

langfun/env/base_feature.py

@@ -23,6 +23,8 @@ the `Environment` and `Sandbox` interfaces directly.
 """
 
 import functools
+import os
+import time
 from typing import Annotated, Callable
 
 from langfun.env import interface
@@ -82,6 +84,7 @@ class BaseFeature(interface.Feature):
     """Called when the feature is bound."""
     super()._on_bound()
     self._sandbox = None
+    self._housekeep_counter = 0
 
   @functools.cached_property
   def name(self) -> str:
@@ -104,6 +107,14 @@ class BaseFeature(interface.Feature):
     assert self._sandbox is not None, 'Feature has not been set up yet.'
     return self._sandbox
 
+  @property
+  def working_dir(self) -> str | None:
+    """Returns the working directory of the feature."""
+    sandbox_workdir = self.sandbox.working_dir
+    if sandbox_workdir is None:
+      return None
+    return os.path.join(sandbox_workdir, self.name)
+
   #
   # Setup and teardown of the feature.
   #
@@ -115,13 +126,14 @@ class BaseFeature(interface.Feature):
   ) -> None:
     """Triggers an event handler."""
     error = None
+    start_time = time.time()
     try:
       action()
     except BaseException as e:  # pylint: disable=broad-except
       error = e
       raise
     finally:
-      event_handler(error=error)
+      event_handler(duration=time.time() - start_time, error=error)
 
   def setup(self, sandbox: interface.Sandbox) -> None:
     """Sets up the feature."""
@@ -146,7 +158,10 @@ class BaseFeature(interface.Feature):
 
   def housekeep(self) -> None:
     """Performs housekeeping for the feature."""
-    self._do(self._housekeep, self.on_housekeep)
+    try:
+      self._do(self._housekeep, self.on_housekeep)
+    finally:
+      self._housekeep_counter += 1
 
   #
   # Event handlers subclasses can override.
@@ -154,51 +169,58 @@ class BaseFeature(interface.Feature):
 
   def on_setup(
       self,
+      duration: float,
       error: BaseException | None = None
   ) -> None:
     """Called when the feature is setup."""
-    self.sandbox.on_feature_setup(self, error)
+    self.sandbox.on_feature_setup(self, duration, error)
 
   def on_teardown(
       self,
+      duration: float,
       error: BaseException | None = None
   ) -> None:
     """Called when the feature is teardown."""
-    self.sandbox.on_feature_teardown(self, error)
+    self.sandbox.on_feature_teardown(self, duration, error)
 
   def on_housekeep(
       self,
+      duration: float,
       error: BaseException | None = None
   ) -> None:
     """Called when the feature has done housekeeping."""
-    self.sandbox.on_feature_housekeep(self, error)
+    self.sandbox.on_feature_housekeep(
+        self, self._housekeep_counter, duration, error
+    )
 
   def on_setup_session(
       self,
+      duration: float,
       error: BaseException | None = None,
   ) -> None:
     """Called when the feature is setup for a user session."""
-    self.sandbox.on_feature_setup_session(self, error)
+    self.sandbox.on_feature_setup_session(self, duration, error)
 
   def on_teardown_session(
       self,
+      duration: float,
       error: BaseException | None = None,
   ) -> None:
     """Called when the feature is teardown for a user session."""
-    self.sandbox.on_feature_teardown_session(self, error)
+    self.sandbox.on_feature_teardown_session(self, duration, error)
 
-  def on_session_activity(
+  def on_activity(
       self,
-      session_id: str,
       name: str,
-      error: BaseException | None,
+      duration: float,
+      error: BaseException | None = None,
       **kwargs
   ) -> None:
     """Called when a sandbox activity is performed."""
-    self.sandbox.on_session_activity(
-        session_id=session_id,
-        name=name,
+    self.sandbox.on_activity(
+        name=f'{self.name}.{name}',
         feature=self,
         error=error,
+        duration=duration,
         **kwargs
     )