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

langfun/env/base_sandbox.py

@@ -29,6 +29,7 @@ import time
 from typing import Annotated, Any, Callable, Iterator, Sequence, Type
 
 from langfun.env import interface
+from langfun.env.event_handlers import base as event_handler_base
 import pyglove as pg
 
 
@@ -36,7 +37,7 @@ class BaseSandbox(interface.Sandbox):
   """Base class for a sandbox."""
 
   id: Annotated[
-      interface.SandboxId,
+      interface.Sandbox.Id,
       'The identifier for the sandbox.'
   ]
 
@@ -92,10 +93,12 @@ class BaseSandbox(interface.Sandbox):
     assert self._status != status, (self._status, status)
     self.on_status_change(self._status, status)
     self._status = status
+    self._status_start_time = time.time()
 
-  def _maybe_report_state_error(self, e: BaseException | None) -> None:
+  def report_maybe_state_error(self, e: BaseException | None) -> None:
     """Reports sandbox state errors."""
-    if isinstance(e, interface.SandboxStateError):
+    if (isinstance(e, interface.SandboxStateError)
+        and e not in self._state_errors):
       self._state_errors.append(e)
 
   def _setup_features(self) -> None:
@@ -133,7 +136,7 @@ class BaseSandbox(interface.Sandbox):
         try:
           feature.teardown()
         except BaseException as e:  # pylint: disable=broad-except
-          self._maybe_report_state_error(e)
+          self.report_maybe_state_error(e)
           errors[feature.name] = e
     if errors:
       return interface.FeatureTeardownError(sandbox=self, errors=errors)
@@ -167,7 +170,7 @@ class BaseSandbox(interface.Sandbox):
         try:
           feature.teardown_session()
         except BaseException as e:  # pylint: disable=broad-except
-          self._maybe_report_state_error(e)
+          self.report_maybe_state_error(e)
           feature_teardown_errors[name] = e
 
     return interface.SessionTeardownError(
@@ -190,7 +193,6 @@ class BaseSandbox(interface.Sandbox):
         for name, feature in self.environment.features.items()
     })
     self._event_handlers = []
-
     self._enable_pre_session_setup = (
         self.reusable and self.proactive_session_setup
     )
@@ -202,18 +204,24 @@ class BaseSandbox(interface.Sandbox):
         )
     )
     self._housekeep_thread = None
-    self._housekeep_count = 0
+    self._housekeep_counter = 0
 
     # Runtime state.
     self._status = self.Status.CREATED
+    self._status_start_time = time.time()
+
     self._start_time = None
     self._state_errors = []
+
     self._features_with_setup_called = set()
     self._features_with_setup_session_called = set()
 
     self._session_id = None
     self._session_start_time = None
 
+    # Thread local state for this sandbox.
+    self._tls_state = threading.local()
+
   @functools.cached_property
   def working_dir(self) -> str | None:
     """Returns the working directory for the sandbox."""
@@ -228,16 +236,21 @@ class BaseSandbox(interface.Sandbox):
     """Marks the sandbox as acquired."""
     self._set_status(self.Status.ACQUIRED)
 
+  @property
+  def housekeep_counter(self) -> int:
+    """Returns the housekeeping counter."""
+    return self._housekeep_counter
+
   def add_event_handler(
       self,
-      event_handler: interface.EnvironmentEventHandler | None
+      event_handler: event_handler_base.EventHandler | None
   ) -> None:
     """Sets the event handler for the sandbox."""
     self._event_handlers.append(event_handler)
 
   def remove_event_handler(
       self,
-      event_handler: interface.EnvironmentEventHandler | None
+      event_handler: event_handler_base.EventHandler | None
   ) -> None:
     """Removes the event handler for the sandbox."""
     self._event_handlers.remove(event_handler)
@@ -247,11 +260,40 @@ class BaseSandbox(interface.Sandbox):
     """Returns all errors encountered during sandbox lifecycle."""
     return self._state_errors
 
+  @property
+  def is_shutting_down(self) -> bool:
+    """Returns True if the sandbox is shutting down."""
+    return self._status == self.Status.SHUTTING_DOWN or (
+        self._state_errors and self._status == self.Status.EXITING_SESSION
+    )
+
   @property
   def features(self) -> dict[str, interface.Feature]:
     """Returns the features in the sandbox."""
     return self._features
 
+  def _enter_service_call(self) -> bool:
+    """Enters a service call.
+
+    Returns:
+      True if the service call is at the top of the call stack.
+    """
+    v = getattr(self._tls_state, 'service_call_depth', None)
+    if v is None:
+      v = 0
+    setattr(self._tls_state, 'service_call_depth', v + 1)
+    return v == 0
+
+  def _exit_service_call(self) -> bool:
+    """Exits a service call.
+
+    Returns:
+      True if the service call is at the top of the call stack.
+    """
+    v = getattr(self._tls_state, 'service_call_depth')
+    setattr(self._tls_state, 'service_call_depth', v - 1)
+    return v == 1
+
   #
   # Sandbox start/shutdown.
   #
@@ -289,7 +331,7 @@ class BaseSandbox(interface.Sandbox):
         f'it is in {self._status} status.'
     )
 
-    t = time.time()
+    starting_time = time.time()
     self._state = self.Status.SETTING_UP
 
     try:
@@ -314,18 +356,16 @@ class BaseSandbox(interface.Sandbox):
       # Mark the sandbox as ready when all setup succeeds.
       self._set_status(self.Status.READY)
 
-      self.on_start()
-      pg.logging.info(
-          '[%s]: Sandbox started in %.2f seconds.',
-          self.id, time.time() - t
-      )
+      duration = time.time() - starting_time
+      self.on_start(duration)
     except BaseException as e:  # pylint: disable=broad-except
+      duration = time.time() - starting_time
       pg.logging.error(
-          '[%s]: Sandbox failed to start: %s',
-          self.id, e
+          '[%s]: Sandbox failed to start in %.2f seconds: %s',
+          self.id, duration, e
       )
-      self._maybe_report_state_error(e)
-      self.on_start(e)
+      self.report_maybe_state_error(e)
+      self.on_start(duration, e)
       self.shutdown()
       raise e
 
@@ -379,7 +419,6 @@ class BaseSandbox(interface.Sandbox):
       return
 
     self._set_status(interface.Sandbox.Status.SHUTTING_DOWN)
-    shutdown_start_time = time.time()
 
     if (self._housekeep_thread is not None
         and threading.current_thread() is not self._housekeep_thread):
@@ -390,20 +429,11 @@ class BaseSandbox(interface.Sandbox):
     try:
       self._shutdown()
       self._set_status(interface.Sandbox.Status.OFFLINE)
-
-      pg.logging.info(
-          '[%s]: Sandbox shutdown in %.2f seconds. '
-          '(lifetime: %.2f seconds, teardown errors: %s)',
-          self.id,
-          time.time() - shutdown_start_time,
-          time.time() - self._start_time if self._start_time else 0,
-          teardown_error
-      )
       self.on_shutdown(teardown_error)
       shutdown_error = None
     except BaseException as e:  # pylint: disable=broad-except
       shutdown_error = e
-      self._maybe_report_state_error(e)
+      self.report_maybe_state_error(e)
       self._set_status(interface.Sandbox.Status.OFFLINE)
       pg.logging.error(
           '[%s]: Sandbox shutdown with error: %s',
@@ -496,10 +526,12 @@ class BaseSandbox(interface.Sandbox):
     try:
       self._start_session()
       self._set_status(self.Status.IN_SESSION)
-      self.on_session_start(session_id)
+      self.on_session_start(session_id, time.time() - self._session_start_time)
     except BaseException as e:  # pylint: disable=broad-except
-      self._maybe_report_state_error(e)
-      self.on_session_start(session_id, e)
+      self.report_maybe_state_error(e)
+      self.on_session_start(
+          session_id, time.time() - self._session_start_time, e
+      )
       self.shutdown()
       raise e
 
@@ -507,15 +539,19 @@ class BaseSandbox(interface.Sandbox):
     """Ends the user session with the sandbox.
 
     State transitions:
-      IN_SESSION -> READY: When user session exits normally, and sandbox is set
-        to reuse.
-      IN_SESSION -> SHUTTING_DOWN -> OFFLINE: When user session exits while
+      IN_SESSION -> EXITING_SESSION -> READY: When user session exits normally,
+        and sandbox is set to reuse.
+      IN_SESSION -> EXITING_SESSION -> SHUTTING_DOWN -> OFFLINE: When user
+        session exits while
         sandbox is set not to reuse, or session teardown fails.
-      IN_SESSION -> SETTING_UP -> READY: When user session exits normally, and
-        sandbox is set to reuse, and proactive session setup is enabled.
-      IN_SESSION -> SETTING_UP -> SHUTTING_DOWN -> OFFLINE: When user session
-        exits normally, and proactive session setup is enabled but fails.
-      not (IN_SESSION) -> same state: No operation
+      IN_SESSION -> EXITING_SESSION -> SETTING_UP -> READY: When user session
+        exits normally, and sandbox is set to reuse, and proactive session setup
+        is enabled.
+      IN_SESSION -> EXITING_SESSION -> SETTING_UP -> SHUTTING_DOWN -> OFFLINE:
+        When user session exits normally, and proactive session setup is enabled
+        but fails.
+      EXITING_SESSION -> EXITING_SESSION: No operation.
+      not IN_SESSION -> same state: No operation
 
     `end_session` should always be called for each `start_session` call, even
     when the session fails to start, to ensure proper cleanup.
@@ -541,6 +577,9 @@ class BaseSandbox(interface.Sandbox):
     Raises:
       BaseException: If session teardown failed with user-defined errors.
     """
+    if self._status == self.Status.EXITING_SESSION:
+      return
+
     if self._status not in (
         self.Status.IN_SESSION,
     ):
@@ -549,6 +588,8 @@ class BaseSandbox(interface.Sandbox):
     assert self._session_id is not None, (
         'No user session is active for this sandbox'
     )
+    # Set sandbox status to EXITING_SESSION to avoid re-entry.
+    self._set_status(self.Status.EXITING_SESSION)
     shutdown_sandbox = shutdown_sandbox or not self.reusable
 
     # Teardown features for the current session.
@@ -572,7 +613,7 @@ class BaseSandbox(interface.Sandbox):
                 self.id,
                 e
             )
-            self._maybe_report_state_error(e)
+            self.report_maybe_state_error(e)
             self.shutdown()
 
         # End session before setting up the next session.
@@ -602,14 +643,6 @@ class BaseSandbox(interface.Sandbox):
       self._set_status(interface.Sandbox.Status.ACQUIRED)
       shutdown_sandbox = True
 
-    pg.logging.info(
-        '[%s]: User session %s ended. '
-        '(lifetime: %.2f seconds, teardown errors: %s).',
-        self.id,
-        self._session_id,
-        time.time() - self._session_start_time,
-        end_session_error
-    )
     self._session_start_time = None
     self._session_event_handler = None
 
@@ -633,8 +666,31 @@ class BaseSandbox(interface.Sandbox):
     last_ping = now
     last_housekeep_time = {name: now for name in self._features.keys()}
 
+    def _next_housekeep_wait_time() -> float:
+      # Decide how long to sleep for the next housekeeping.
+      next_housekeep_time = None
+      if self.keepalive_interval is not None:
+        next_housekeep_time = last_ping + self.keepalive_interval
+
+      for name, feature in self._features.items():
+        if feature.housekeep_interval is None:
+          continue
+        next_feature_housekeep_time = (
+            last_housekeep_time[name] + feature.housekeep_interval
+        )
+        if (next_housekeep_time is None
+            or next_housekeep_time > next_feature_housekeep_time):
+          next_housekeep_time = next_feature_housekeep_time
+
+      # Housekeep loop is installed when at least one feature requires
+      # housekeeping or the sandbox has a keepalive interval.
+      assert next_housekeep_time is not None
+      return max(0, next_housekeep_time - time.time())
+
     while self._status not in (self.Status.SHUTTING_DOWN, self.Status.OFFLINE):
+      housekeep_start = time.time()
       if self.keepalive_interval is not None:
+
         if time.time() - last_ping > self.keepalive_interval:
           try:
             self.ping()
@@ -645,8 +701,9 @@ class BaseSandbox(interface.Sandbox):
                 self.id,
                 str(e)
             )
-            self._housekeep_count += 1
-            self._maybe_report_state_error(e)
+            self._housekeep_counter += 1
+            self.report_maybe_state_error(e)
+            self.on_housekeep(time.time() - housekeep_start, e)
             self.shutdown()
             break
           last_ping = time.time()
@@ -667,20 +724,28 @@ class BaseSandbox(interface.Sandbox):
                 feature.name,
                 e,
             )
-            self._maybe_report_state_error(e)
+            self.report_maybe_state_error(e)
+            self._housekeep_counter += 1
+            self.on_housekeep(time.time() - housekeep_start, e)
             self.shutdown()
             break
-      self._housekeep_count += 1
-      time.sleep(1)
+
+      self._housekeep_counter += 1
+      self.on_housekeep(time.time() - housekeep_start)
+      time.sleep(_next_housekeep_wait_time())
 
   #
   # 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 sandbox is started."""
     for handler in self._event_handlers:
-      handler.on_sandbox_start(self.environment, self, error)
+      handler.on_sandbox_start(self.environment, self, duration, error)
 
   def on_status_change(
       self,
@@ -690,96 +755,124 @@ class BaseSandbox(interface.Sandbox):
     """Called when the sandbox status changes."""
     for handler in self._event_handlers:
       handler.on_sandbox_status_change(
-          self.environment, self, old_status, new_status
+          self.environment,
+          self,
+          old_status,
+          new_status,
+          time.time() - self._status_start_time
       )
 
   def on_shutdown(self, error: BaseException | None = None) -> None:
     """Called when the sandbox is shutdown."""
+    if self._start_time is None:
+      lifetime = 0.0
+    else:
+      lifetime = time.time() - self._start_time
     for handler in self._event_handlers:
-      handler.on_sandbox_shutdown(self.environment, self, error)
+      handler.on_sandbox_shutdown(self.environment, self, lifetime, error)
+
+  def on_housekeep(
+      self,
+      duration: float,
+      error: BaseException | None = None
+  ) -> None:
+    """Called when the sandbox finishes a round of housekeeping."""
+    counter = self._housekeep_counter
+    for handler in self._event_handlers:
+      handler.on_sandbox_housekeep(
+          self.environment, self, counter, duration, error
+      )
 
   def on_feature_setup(
       self,
       feature: interface.Feature,
+      duration: float,
       error: BaseException | None = None
   ) -> None:
     """Called when a feature is setup."""
     for handler in self._event_handlers:
       handler.on_feature_setup(
-          self.environment, self, feature, error
+          self.environment, self, feature, duration, error
       )
 
   def on_feature_teardown(
       self,
       feature: interface.Feature,
+      duration: float,
       error: BaseException | None = None
   ) -> None:
     """Called when a feature is teardown."""
     for handler in self._event_handlers:
       handler.on_feature_teardown(
-          self.environment, self, feature, error
+          self.environment, self, feature, duration, error
       )
 
   def on_feature_setup_session(
       self,
       feature: interface.Feature,
+      duration: float,
       error: BaseException | None = None
   ) -> None:
     """Called when a feature is setup for a user session."""
     for handler in self._event_handlers:
       handler.on_feature_setup_session(
-          self.environment, self, feature, self.session_id, error
+          self.environment, self, feature, self.session_id, duration, error
       )
 
   def on_feature_teardown_session(
       self,
       feature: interface.Feature,
+      duration: float,
       error: BaseException | None = None
   ) -> None:
     """Called when a feature is teardown for a user session."""
     for handler in self._event_handlers:
       handler.on_feature_teardown_session(
-          self.environment, self, feature, self.session_id, error
+          self.environment, self, feature, self.session_id, duration, error
       )
 
   def on_feature_housekeep(
       self,
       feature: interface.Feature,
+      counter: int,
+      duration: float,
       error: BaseException | None = None
   ) -> None:
     """Called when a feature is housekeeping."""
     for handler in self._event_handlers:
       handler.on_feature_housekeep(
-          self.environment, self, feature, error
+          self.environment, self, feature, counter, duration, error
       )
 
   def on_session_start(
       self,
       session_id: str,
+      duration: float,
       error: BaseException | None = None
   ) -> None:
     """Called when the user session starts."""
     for handler in self._event_handlers:
       handler.on_session_start(
-          self.environment, self, session_id, error
+          self.environment, self, session_id, duration, error
       )
 
-  def on_session_activity(
+  def on_activity(
       self,
-      session_id: str,
       name: str,
+      duration: float,
       feature: interface.Feature | None = None,
       error: BaseException | None = None,
       **kwargs
   ) -> None:
     """Called when a sandbox activity is performed."""
     for handler in self._event_handlers:
-      handler.on_session_activity(
-          session_id=session_id,
+      handler.on_sandbox_activity(
           name=name,
           environment=self.environment,
           sandbox=self,
           feature=feature,
+          session_id=self.session_id,
+          duration=duration,
           error=error,
           **kwargs
       )
@@ -790,9 +883,10 @@ class BaseSandbox(interface.Sandbox):
       error: BaseException | None = None
   ) -> None:
     """Called when the user session ends."""
+    lifetime = time.time() - self._session_start_time
     for handler in self._event_handlers:
       handler.on_session_end(
-          self.environment, self, session_id, error
+          self.environment, self, session_id, lifetime, error
       )
 
 
@@ -876,9 +970,14 @@ def sandbox_service(
     @functools.wraps(func)
     def method_wrapper(self, *args, **kwargs) -> Any:
       """Helper function to safely execute logics in the sandbox."""
+
       assert isinstance(self, (BaseSandbox, interface.Feature)), self
       sandbox = self.sandbox if isinstance(self, interface.Feature) else self
 
+      # We count the service call stack depth so we could shutdown the sandbox
+      # at the top upon sandbox state error.
+      sandbox._enter_service_call()  # pylint: disable=protected-access
+
       # When a capability is directly accessed from the environment,
       # we create a new session for the capability call. This
       # prevents the sandbox from being reused for other feature calls.
@@ -895,70 +994,102 @@ def sandbox_service(
         new_session = False
 
       kwargs.pop('session_id', None)
-      session_id = sandbox.session_id
       result = None
-      state_error = None
       error = None
+      start_time = time.time()
 
       try:
         # Execute the service function.
         result = func(self, *args, **kwargs)
 
-        # If the result is a context manager, use it and end the session
-        # afterwards.
-        if new_session and isinstance(
-            result, contextlib.AbstractContextManager
-        ):
-          return _end_session_when_exit(result, sandbox)
+        # If the result is a context manager, wrap it with a context manager
+        # to end the session when exiting.
+        if isinstance(result, contextlib.AbstractContextManager):
+          return _service_context_manager_wrapper(
+              service=result,
+              sandbox_or_feature=self,
+              sandbox=sandbox,
+              name=func.__name__,
+              kwargs=to_kwargs(*args, **kwargs),
+              start_time=start_time,
+              new_session=new_session
+          )
 
         # Otherwise, return the result and end the session in the finally block.
         return result
-      except interface.SandboxStateError as e:
-        sandbox._maybe_report_state_error(e)  # pylint: disable=protected-access
-        state_error = e
-        error = e
-        raise
       except BaseException as e:
         error = e
+        sandbox.report_maybe_state_error(e)
         if pg.match_error(e, critical_errors):
           state_error = interface.SandboxStateError(
               'Sandbox encountered an unexpected error executing '
               f'`{func.__name__}` (args={args!r}, kwargs={kwargs!r}): {e}',
               sandbox=self
           )
-          sandbox._maybe_report_state_error(state_error)  # pylint: disable=protected-access
+          sandbox.report_maybe_state_error(state_error)
           raise state_error from e
         raise
       finally:
-        if session_id is not None:
-          self.on_session_activity(
+        is_topmost_call = sandbox._exit_service_call()  # pylint: disable=protected-access
+        if not isinstance(result, contextlib.AbstractContextManager):
+          self.on_activity(
               name=func.__name__,
-              session_id=session_id,
+              duration=time.time() - start_time,
               error=error,
               **to_kwargs(*args, **kwargs),
           )
-
-        if state_error is not None:
+          if new_session:
+            assert is_topmost_call
+
+            # End the session if it's from a feature method and the result
+            # is not a context manager.
+            sandbox.end_session()
+
+        # Shutdown the sandbox if it is at the top of the service call stack and
+        # has state errors.
+        if (is_topmost_call
+            and sandbox.state_errors
+            # Sandbox service method might be called during shutting down, in
+            # that case we don't want to shutdown the sandbox again.
+            and not sandbox.is_shutting_down):
           sandbox.shutdown()
-        elif (new_session
-              and not isinstance(result, contextlib.AbstractContextManager)):
-          # End the session if it's from a feature method and the result is not
-          # a context manager.
-          sandbox.end_session(
-              shutdown_sandbox=isinstance(error, interface.SandboxStateError)
-          )
+
     return method_wrapper
   return decorator
 
 
 @contextlib.contextmanager
-def _end_session_when_exit(
+def _service_context_manager_wrapper(
     service: contextlib.AbstractContextManager[Any],
-    sandbox: interface.Sandbox
+    sandbox_or_feature: BaseSandbox | interface.Feature,
+    sandbox: interface.Sandbox,
+    name: str,
+    kwargs: dict[str, Any],
+    new_session: bool,
+    start_time: float,
 ) -> Iterator[Any]:
   """Context manager wrapper for ending a sandbox session when exiting."""
+  error = None
+  sandbox._enter_service_call()  # pylint: disable=protected-access
+
   try:
     with service as result:
       yield result
+  except BaseException as e:
+    error = e
+    sandbox.report_maybe_state_error(error)
+    raise
   finally:
-    sandbox.end_session()
+    sandbox_or_feature.on_activity(
+        name=name,
+        error=error,
+        duration=time.time() - start_time,
+        **kwargs,
+    )
+    is_topmost_call = sandbox._exit_service_call()  # pylint: disable=protected-access
+
+    if new_session:
+      assert is_topmost_call
+      sandbox.end_session()
+    elif isinstance(error, interface.SandboxStateError):
+      sandbox.shutdown()