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

langfun/env/interface.py

@@ -22,60 +22,6 @@ from typing import Annotated, Any, ContextManager, ClassVar, Iterator, Optional
 
 import pyglove as pg
 
-#
-# Environemnt identifiers.
-#
-
-
-@dataclasses.dataclass(frozen=True)
-class EnvironmentId:
-  """Identifier for an environment."""
-  environment_id: str
-
-  def __str__(self) -> str:
-    return self.environment_id
-
-  def working_dir(self, root_dir: str | None) -> str | None:
-    """Returns the download directory for the service."""
-    if root_dir is None:
-      return None
-    return os.path.join(root_dir, _make_path_compatible(self.environment_id))
-
-# Enable automatic conversion from str to EnvironmentId.
-pg.typing.register_converter(str, EnvironmentId, EnvironmentId)
-
-
-@dataclasses.dataclass(frozen=True)
-class SandboxId:
-  """Identifier for a sandbox."""
-  environment_id: EnvironmentId
-  sandbox_id: str
-
-  def __str__(self) -> str:
-    return f'{self.environment_id}/{self.sandbox_id}'
-
-  def working_dir(self, root_dir: str | None) -> str | None:
-    """Returns the download directory for the sandbox."""
-    if root_dir is None:
-      return None
-    return os.path.join(
-        self.environment_id.working_dir(root_dir),
-        _make_path_compatible(self.sandbox_id)
-    )
-
-
-def _make_path_compatible(id_str: str) -> str:
-  """Makes a path compatible with CNS."""
-  return id_str.translate(
-      str.maketrans({
-          '@': '_',
-          ':': '_',
-          '#': '_',
-          ' ': '',
-      })
-  )
-
-
 #
 # Environment errors.
 #
@@ -216,144 +162,6 @@ class SessionTeardownError(SandboxError):
         not isinstance(e, SandboxStateError) for e in self.errors.values()
     )
 
-
-#
-# Event handler.
-#
-
-
-class SessionEventHandler:
-  """Base class for session event handlers."""
-
-  def on_session_start(
-      self,
-      environment: 'Environment',
-      sandbox: 'Sandbox',
-      session_id: str,
-      error: BaseException | None
-  ) -> None:
-    """Called when a sandbox session starts."""
-
-  def on_session_activity(
-      self,
-      session_id: str,
-      name: str,
-      environment: 'Environment',
-      sandbox: 'Sandbox',
-      feature: Optional['Feature'],
-      error: BaseException | None,
-      **kwargs
-  ) -> None:
-    """Called when a sandbox activity is performed."""
-
-  def on_session_end(
-      self,
-      environment: 'Environment',
-      sandbox: 'Sandbox',
-      session_id: str,
-      error: BaseException | None
-  ) -> None:
-    """Called when a sandbox session ends."""
-
-
-class FeatureEventHandler:
-  """Base class for feature event handlers."""
-
-  def on_feature_setup(
-      self,
-      environment: 'Environment',
-      sandbox: 'Sandbox',
-      feature: 'Feature',
-      error: BaseException | None
-  ) -> None:
-    """Called when a sandbox feature is setup."""
-
-  def on_feature_teardown(
-      self,
-      environment: 'Environment',
-      sandbox: 'Sandbox',
-      feature: 'Feature',
-      error: BaseException | None
-  ) -> None:
-    """Called when a sandbox feature is teardown."""
-
-  def on_feature_teardown_session(
-      self,
-      environment: 'Environment',
-      sandbox: 'Sandbox',
-      feature: 'Feature',
-      session_id: str,
-      error: BaseException | None
-  ) -> None:
-    """Called when a feature is teardown with a session."""
-
-  def on_feature_setup_session(
-      self,
-      environment: 'Environment',
-      sandbox: 'Sandbox',
-      feature: 'Feature',
-      session_id: str | None,
-      error: BaseException | None
-  ) -> None:
-    """Called when a feature is setup with a session."""
-
-  def on_feature_housekeep(
-      self,
-      environment: 'Environment',
-      sandbox: 'Sandbox',
-      feature: 'Feature',
-      error: BaseException | None
-  ) -> None:
-    """Called when a sandbox feature is housekeeping."""
-
-
-class SandboxEventHandler(FeatureEventHandler, SessionEventHandler):
-  """Base class for sandbox event handlers."""
-
-  def on_sandbox_start(
-      self,
-      environment: 'Environment',
-      sandbox: 'Sandbox',
-      error: BaseException | None
-  ) -> None:
-    """Called when a sandbox is started."""
-
-  def on_sandbox_status_change(
-      self,
-      environment: 'Environment',
-      sandbox: 'Sandbox',
-      old_status: 'Sandbox.Status',
-      new_status: 'Sandbox.Status',
-  ) -> None:
-    """Called when a sandbox status changes."""
-
-  def on_sandbox_shutdown(
-      self,
-      environment: 'Environment',
-      sandbox: 'Sandbox',
-      error: BaseException | None
-  ) -> None:
-    """Called when a sandbox is shutdown."""
-
-
-class EnvironmentEventHandler(SandboxEventHandler):
-  """Base class for environment event handlers."""
-
-  def on_environment_start(
-      self,
-      environment: 'Environment',
-      error: BaseException | None
-  ) -> None:
-    """Called when the environment is started."""
-
-  def on_environment_shutdown(
-      self,
-      environment: 'Environment',
-      error: BaseException | None
-  ) -> None:
-    """Called when the environment is shutdown."""
-
-
 #
 # Interface for sandbox-based environment.
 #
@@ -362,6 +170,20 @@ class EnvironmentEventHandler(SandboxEventHandler):
 class Environment(pg.Object):
   """Base class for an environment."""
 
+  @dataclasses.dataclass(frozen=True)
+  class Id:
+    """Identifier for an environment."""
+    environment_id: str
+
+    def __str__(self) -> str:
+      return self.environment_id
+
+    def working_dir(self, root_dir: str | None) -> str | None:
+      """Returns the download directory for the service."""
+      if root_dir is None:
+        return None
+      return os.path.join(root_dir, _make_path_compatible(self.environment_id))
+
   class Status(enum.Enum):
     """Environment state.
 
@@ -384,13 +206,6 @@ class Environment(pg.Object):
       'Features to be exposed by the environment.'
   ] = {}
 
-  event_handlers: Annotated[
-      list[EnvironmentEventHandler],
-      (
-          'User handler for the environment events.'
-      )
-  ] = []
-
   _ENV_STACK: Annotated[
       ClassVar[list['Environment']],
       'Recording the environments stacked through context managers.'
@@ -402,7 +217,7 @@ class Environment(pg.Object):
 
   @property
   @abc.abstractmethod
-  def id(self) -> EnvironmentId:
+  def id(self) -> Id:
     """Returns the identifier for the environment."""
 
   @property
@@ -466,6 +281,10 @@ class Environment(pg.Object):
     Environment._ENV_STACK.pop()
     self.shutdown()
 
+  def __del__(self):
+    """Deletes the environment."""
+    self.shutdown()
+
   @classmethod
   def current(cls) -> Optional['Environment']:
     """Returns the current environment."""
@@ -506,49 +325,78 @@ class Environment(pg.Object):
     raise AttributeError(name)
 
 
+# Enable automatic conversion from str to Environment.Id.
+pg.typing.register_converter(str, Environment.Id, Environment.Id)
+
+
 class Sandbox(pg.Object):
   """Interface for sandboxes."""
 
+  @dataclasses.dataclass(frozen=True, slots=True)
+  class Id:
+    """Identifier for a sandbox."""
+    environment_id: Environment.Id
+    sandbox_id: str
+
+    def __str__(self) -> str:
+      return f'{self.environment_id}/{self.sandbox_id}'
+
+    def working_dir(self, root_dir: str | None) -> str | None:
+      """Returns the download directory for the sandbox."""
+      if root_dir is None:
+        return None
+      return os.path.join(
+          self.environment_id.working_dir(root_dir),
+          _make_path_compatible(self.sandbox_id)
+      )
+
   class Status(enum.Enum):
-    """Sandbox state.
+    r"""Sandbox state.
 
     State transitions:
 
-            +---------------+          +---------------+
-            |  <OFFLINE>    | <------  | SHUTTING_DOWN |
-            +---------------+          +---------------+
-                                           ^    ^
-                                           |    |
-                                 (shutdown)|    +------------------------+
-                                           |                             |
-    +-----------+    (call start) +------------+                         |
-    | <CREATED> | ------------->  | SETTING_UP | <----------------+      |
-    +-----------+                 +------------+                  |      |
-                                       |                          |      |
-                                       | (start succeeded)        |      |
-                                       | OR (_setup_session)      |      |
-                                       v                          |      |
-                                   +---------+                    |      |
-                                   |  READY  |                    |      |
-                                   +---------+                    |      |
-                                       |                          |      |
-                                       | (set_acquired)           |      |
-                                       v                          |      |
-                                    +----------+                  |      |
-                                    | ACQUIRED |                  |      |
-                                    +----------+                  |      |
-                                       |                          |      |
-                                       | (call start_session)     |      |
-                                    +------------+                |      |
-                                    | SETTING_UP | --(failed) -----------+
-                                    +------------+                |      |
-                                       |                          |      |
-                                       v (start_session succeeded)|      |
-                                 +--------------+                 |      |
-                                 |  IN_SESSION  |--(end_session)--+      |
-                                 +--------------+                        |
-                                         |                               |
-                                         +------(shutdown)---------------+
+                    (sandbox / feature
+    +------------+   teardown)            +---------------+
+    | <OFFLINE>  | <--------------------- | SHUTTING_DOWN |
+    +------------+                        +---------------+
+                                            ^     ^
+                                           /       \
+                          (setup failed)  /         \
+                                         /           \
+    +-----------+   (start)  +------------+           \
+    | <CREATED> | -------->  | SETTING_UP |            \
+    +-----------+         ^  +------------+             \
+                         /        |                      \
+                        /         | (sandbox /            \
+                       /          |  feature /session      \
+                      /           v  setup succeeded)       \
+                     /        +---------+                    \
+                    /         |  READY  |                     \
+                   /          +---------+                      \
+                  /                |                            \
+                 /                 |  (acquire)                  \
+                /                  v                              \
+               /              +----------+                         \
+              |               | ACQUIRED |                          \
+              |               +----------+                           |
+              |                    |                                 |
+              |                    |  (start_session)                |
+              |               +------------+                         |
+              |               | SETTING_UP |-- (setup failed) ------>+
+              |               +------------+                         |
+              |                    |                                 |
+              |                    v  (succeeded)                    |
+              |               +--------------+                       |
+              |               |  IN_SESSION  |- (op failed) -------->+
+              |               +--------------+                       |
+              |                    |                                 |
+              |                    |  (end_session)                  |
+              |                    |                                 |
+              |                    v             (session teardown   |
+      (setup next            +-----------------+  failed OR          |
+      session for  <---------| EXITING_SESSION |- non-reusable  -----+
+      reusable sandbox)      +-----------------+  sandbox)
+
     """
 
     # The sandbox is created, but not yet started.
@@ -566,6 +414,9 @@ class Sandbox(pg.Object):
     # The sandbox is in a user session.
     IN_SESSION = 'in_session'
 
+    # The sandbox is exiting a user session.
+    EXITING_SESSION = 'exiting_session'
+
     # The sandbox is being shut down.
     SHUTTING_DOWN = 'shutting_down'
 
@@ -584,7 +435,7 @@ class Sandbox(pg.Object):
 
   @property
   @abc.abstractmethod
-  def id(self) -> SandboxId:
+  def id(self) -> Id:
     """Returns the identifier for the sandbox."""
 
   @property
@@ -607,24 +458,6 @@ class Sandbox(pg.Object):
     """Returns True if the sandbox is online."""
     return self.status.is_online
 
-  @abc.abstractmethod
-  def set_acquired(self) -> None:
-    """Marks the sandbox as acquired."""
-
-  @abc.abstractmethod
-  def add_event_handler(
-      self,
-      event_handler: EnvironmentEventHandler
-  ) -> None:
-    """Sets the status of the sandbox."""
-
-  @abc.abstractmethod
-  def remove_event_handler(
-      self,
-      event_handler: EnvironmentEventHandler
-  ) -> None:
-    """Removes the status of the sandbox."""
-
   @property
   @abc.abstractmethod
   def state_errors(self) -> list[SandboxStateError]:
@@ -742,14 +575,18 @@ class Sandbox(pg.Object):
     """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.
+      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
@@ -838,6 +675,10 @@ class Sandbox(pg.Object):
       return self.features[name]
     raise AttributeError(name)
 
+  def __del__(self):
+    """Deletes the sandbox."""
+    self.shutdown()
+
 
 class Feature(pg.Object):
   """Interface for sandbox features."""
@@ -969,3 +810,17 @@ class Feature(pg.Object):
     """Returns the current user session identifier."""
     assert self.sandbox is not None
     return self.sandbox.session_id
+
+
+def _make_path_compatible(id_str: str) -> str:
+  """Makes a path compatible with CNS."""
+  return id_str.translate(
+      str.maketrans({
+          '@': '_',
+          ':': '_',
+          '#': '_',
+          ' ': '',
+          '<': '',
+          '>': '',
+      })
+  )

langfun/env/interface_test.py

@@ -18,7 +18,7 @@ from langfun.env import interface
 class IdTest(unittest.TestCase):
 
   def test_environment_id(self):
-    env_id = interface.EnvironmentId('env@1/a b:c#def')
+    env_id = interface.Environment.Id('env@1/a b:c#def')
     self.assertEqual(str(env_id), 'env@1/a b:c#def')
     self.assertEqual(
         env_id.working_dir(root_dir='/tmp'),
@@ -27,8 +27,8 @@ class IdTest(unittest.TestCase):
     self.assertIsNone(env_id.working_dir(root_dir=None))
 
   def test_sandbox_id(self):
-    sandbox_id = interface.SandboxId(
-        environment_id=interface.EnvironmentId('env'),
+    sandbox_id = interface.Sandbox.Id(
+        environment_id=interface.Environment.Id('env'),
         sandbox_id='sandbox'
     )
     self.assertEqual(str(sandbox_id), 'env/sandbox')

langfun/env/load_balancers_test.py

@@ -28,9 +28,9 @@ class TestingSandbox(interface.Sandbox):
     self._session_id = None
 
   @property
-  def id(self) -> interface.SandboxId:
-    return interface.SandboxId(
-        environment_id=interface.EnvironmentId('testing-env'),
+  def id(self) -> interface.Sandbox.Id:
+    return interface.Sandbox.Id(
+        environment_id=interface.Environment.Id('testing-env'),
         sandbox_id=self.sandbox_id
     )
 
@@ -52,18 +52,6 @@ class TestingSandbox(interface.Sandbox):
   def set_acquired(self) -> None:
     self.set_status(self.Status.ACQUIRED)
 
-  def add_event_handler(
-      self,
-      event_handler: interface.EnvironmentEventHandler
-  ) -> None:
-    pass
-
-  def remove_event_handler(
-      self,
-      event_handler: interface.EnvironmentEventHandler
-  ) -> None:
-    pass
-
   def start(self) -> None:
     self.set_alive()