langfun 0.1.2.dev202510230805__py3-none-any.whl → 0.1.2.dev202511270805__py3-none-any.whl

langfun/env/interface_test.py

@@ -11,6 +11,8 @@
 # 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 contextlib
+from typing import Iterator
 import unittest
 from langfun.env import interface
 
@@ -29,15 +31,123 @@ class IdTest(unittest.TestCase):
   def test_sandbox_id(self):
     sandbox_id = interface.Sandbox.Id(
         environment_id=interface.Environment.Id('env'),
+        image_id='image:2025_01_01_00_00_00',
         sandbox_id='sandbox'
     )
-    self.assertEqual(str(sandbox_id), 'env/sandbox')
+    self.assertEqual(str(sandbox_id), 'env/image:2025_01_01_00_00_00:sandbox')
     self.assertEqual(
         sandbox_id.working_dir(root_dir='/tmp'),
-        '/tmp/env/sandbox'
+        '/tmp/env/image_2025_01_01_00_00_00/sandbox'
     )
     self.assertIsNone(sandbox_id.working_dir(root_dir=None))
 
+  def test_feature_id(self):
+    # For non-sandboxed feature.
+    feature_id = interface.Feature.Id(
+        container_id=interface.Environment.Id('env'),
+        feature_name='feature'
+    )
+    self.assertEqual(str(feature_id), 'env/feature')
+    self.assertEqual(
+        feature_id.working_dir(root_dir='/tmp'),
+        '/tmp/env/feature'
+    )
+    self.assertIsNone(feature_id.working_dir(root_dir=None))
+
+    # For sandboxed feature.
+    feature_id = interface.Feature.Id(
+        container_id=interface.Sandbox.Id(
+            environment_id=interface.Environment.Id('env'),
+            image_id='image1',
+            sandbox_id='0'
+        ),
+        feature_name='feature'
+    )
+    self.assertEqual(str(feature_id), 'env/image1:0/feature')
+    self.assertEqual(
+        feature_id.working_dir(root_dir='/tmp'),
+        '/tmp/env/image1/0/feature'
+    )
+    self.assertIsNone(feature_id.working_dir(root_dir=None))
+
+
+class TestingSandbox(interface.Sandbox):
+
+  id: interface.Sandbox.Id = interface.Sandbox.Id(
+      environment_id=interface.Environment.Id('env'),
+      image_id='test_image',
+      sandbox_id='0:0'
+  )
+  image_id: str = 'test_image'
+  features: dict[str, interface.Feature] = {}
+  status: interface.Sandbox.Status = interface.Sandbox.Status.READY
+  session_id: str | None = None
+
+  __test__ = False
+
+  def environment(self) -> interface.Environment:
+    pass
+
+  def _on_bound(self) -> None:
+    self.activities = []
+
+  def report_state_error(self, error: interface.SandboxStateError) -> None:
+    pass
+
+  def start(self) -> None:
+    pass
+
+  def shutdown(self) -> None:
+    pass
+
+  def start_session(self, session_id: str) -> None:
+    pass
+
+  def end_session(self, shutdown_sandbox: bool = False) -> None:
+    pass
+
+  @contextlib.contextmanager
+  def track_activity(
+      self,
+      name: str,
+      feature: interface.Feature | None = None,
+      **kwargs
+  ) -> Iterator[None]:
+    error = None
+    try:
+      yield
+    except BaseException as e:
+      error = e
+      raise
+    finally:
+      self.activities.append((name, error, kwargs))
+
+
+class DecoratorTest(unittest.TestCase):
+
+  def test_treat_as_sandbox_state_error(self):
+
+    class SandboxA(TestingSandbox):
+
+      @interface.treat_as_sandbox_state_error(errors=(ValueError,))
+      def foo(self, bar: str) -> None:
+        raise ValueError(bar)
+
+    with self.assertRaises(interface.SandboxStateError):
+      SandboxA().foo('foo')
+
+  def test_log_sandbox_activity(self):
+
+    class SandboxB(TestingSandbox):
+
+      @interface.log_activity()
+      def bar(self, x: str) -> None:
+        pass
+
+    sb = SandboxB()
+    sb.bar('foo')
+    self.assertEqual(sb.activities, [('bar', None, {'x': 'foo'})])
+
 
 if __name__ == '__main__':
   unittest.main()

langfun/env/load_balancers_test.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import concurrent.futures
+import contextlib
 import time
+from typing import Any, Iterator
 import unittest
 
 from langfun.env import interface
@@ -22,6 +24,9 @@ from langfun.env import load_balancers
 class TestingSandbox(interface.Sandbox):
   sandbox_id: str
   status: interface.Sandbox.Status = interface.Sandbox.Status.READY
+  image_id: str = 'test_image'
+
+  __test__ = False
 
   def _on_bound(self) -> None:
     super()._on_bound()
@@ -31,6 +36,7 @@ class TestingSandbox(interface.Sandbox):
   def id(self) -> interface.Sandbox.Id:
     return interface.Sandbox.Id(
         environment_id=interface.Environment.Id('testing-env'),
+        image_id=self.image_id,
         sandbox_id=self.sandbox_id
     )
 
@@ -46,6 +52,9 @@ class TestingSandbox(interface.Sandbox):
   def state_errors(self) -> list[interface.SandboxStateError]:
     return []
 
+  def report_state_error(self, error: interface.SandboxStateError) -> None:
+    pass
+
   def set_status(self, status: interface.Sandbox.Status) -> None:
     self.rebind(status=status, skip_notification=True)
 
@@ -53,10 +62,10 @@ class TestingSandbox(interface.Sandbox):
     self.set_status(self.Status.ACQUIRED)
 
   def start(self) -> None:
-    self.set_alive()
+    pass
 
   def shutdown(self) -> None:
-    self.set_alive(False)
+    pass
 
   def start_session(self, session_id: str) -> None:
     self._session_id = session_id
@@ -68,6 +77,18 @@ class TestingSandbox(interface.Sandbox):
   def session_id(self) -> str | None:
     return self._session_id
 
+  @contextlib.contextmanager
+  def track_activity(
+      self,
+      name: str,
+      feature: interface.Feature | None = None,
+      **kwargs: Any
+  ) -> Iterator[None]:
+    try:
+      yield
+    finally:
+      pass
+
 
 class RoundRobinTest(unittest.TestCase):
 

langfun/env/test_utils.py

@@ -21,19 +21,20 @@ from langfun.env import base_environment
 from langfun.env import base_feature
 from langfun.env import base_sandbox
 from langfun.env import interface
-from langfun.env.event_handlers import base as event_handler_base
 import pyglove as pg
 
 
 class TestingEnvironment(base_environment.BaseEnvironment):
   """Testing environment for unit tests."""
-
+  image_ids: list[str] = ['test_image']
   housekeep_interval: float = 0.0
   simulate_start_error: Type[BaseException] | None = None
   simulate_shutdown_error: Type[BaseException] | None = None
   simulate_ping_error: Type[BaseException] | None = None
   offline: bool = False
 
+  __test__ = False
+
   @property
   def id(self) -> interface.Environment.Id:
     return interface.Environment.Id('testing-env')
@@ -45,6 +46,7 @@ class TestingEnvironment(base_environment.BaseEnvironment):
 
   def _create_sandbox(
       self,
+      image_id: str,
       sandbox_id: str,
       reusable: bool,
       proactive_session_setup: bool,
@@ -54,8 +56,10 @@ class TestingEnvironment(base_environment.BaseEnvironment):
         environment=self,
         id=interface.Sandbox.Id(
             environment_id=self.id,
+            image_id=image_id,
             sandbox_id=sandbox_id
         ),
+        image_id=image_id,
         reusable=reusable,
         proactive_session_setup=proactive_session_setup,
         keepalive_interval=keepalive_interval,
@@ -72,6 +76,8 @@ class TestingSandbox(base_sandbox.BaseSandbox):
   simulate_shutdown_error: Type[BaseException] | None = None
   simulate_ping_error: Type[BaseException] | None = None
 
+  __test__ = False
+
   def _on_bound(self) -> None:
     super()._on_bound()
     self._shell_history = []
@@ -109,7 +115,8 @@ class TestingSandbox(base_sandbox.BaseSandbox):
       self._raise_error('Sandbox shutdown error', self.simulate_shutdown_error)
     super()._shutdown()
 
-  @base_sandbox.sandbox_service(critical_errors=(RuntimeError,))
+  @interface.treat_as_sandbox_state_error(errors=(RuntimeError,))
+  @interface.log_activity()
   def shell(
       self,
       code: str,
@@ -138,6 +145,8 @@ class TestingFeature(base_feature.BaseFeature):
   simulate_teardown_session_error: Type[BaseException] | None = None
   call_end_session_on_teardown_session: bool = False
 
+  __test__ = False
+
   class Service:
     """Sandbox."""
 
@@ -181,19 +190,19 @@ class TestingFeature(base_feature.BaseFeature):
     if self.call_end_session_on_teardown_session:
       self.sandbox.end_session()
 
-  @base_sandbox.sandbox_service()
+  @interface.log_activity()
   def num_shell_calls(self) -> int:
     return len(self.sandbox._shell_history)  # pylint: disable=protected-access
 
-  @base_sandbox.sandbox_service()
+  @interface.log_activity()
   def bad_shell_call(self) -> None:
     self.sandbox.shell('bad command', raise_error=RuntimeError)
 
-  @base_sandbox.sandbox_service()
+  @interface.log_activity()
   def show_session_id(self):
     return self.session_id
 
-  @base_sandbox.sandbox_service()
+  @interface.log_activity()
   def call_with_varargs(self, code: str, *args, **kwargs):
     del code, args, kwargs
     return 0
@@ -202,7 +211,6 @@ class TestingFeature(base_feature.BaseFeature):
     super()._on_bound()
     self._service = None
 
-  @base_sandbox.sandbox_service()
   @contextlib.contextmanager
   def my_service(self) -> Iterator[Service]:
     try:
@@ -218,9 +226,46 @@ class TestingFeature(base_feature.BaseFeature):
       )
 
 
-class TestingEventHandler(
-    pg.Object, event_handler_base.EventHandler
-):
+class TestingNonSandboxBasedFeature(base_feature.BaseFeature):
+  """Testing non-sandbox based feature for unit tests."""
+  is_sandbox_based: bool = False
+  simulate_setup_error: Type[BaseException] | None = None
+  simulate_teardown_error: Type[BaseException] | None = None
+  simulate_setup_session_error: Type[BaseException] | None = None
+  simulate_teardown_session_error: Type[BaseException] | None = None
+  simulate_housekeep_error: Type[BaseException] | None = None
+
+  __test__ = False
+
+  def _setup(self) -> None:
+    if self.simulate_setup_error:
+      raise self.simulate_setup_error('Feature setup error')
+
+  def _teardown(self) -> None:
+    if self.simulate_teardown_error:
+      raise self.simulate_teardown_error('Feature teardown error')
+
+  def _setup_session(self) -> None:
+    if self.simulate_setup_session_error:
+      raise self.simulate_setup_session_error('Feature session setup error')
+
+  def _teardown_session(self) -> None:
+    if self.simulate_teardown_session_error:
+      raise self.simulate_teardown_session_error(
+          'Feature session teardown error'
+      )
+
+  def _housekeep(self) -> None:
+    if self.simulate_housekeep_error:
+      raise self.simulate_housekeep_error('Feature housekeeping error')
+    _ = self.foo(1)
+
+  @interface.log_activity()
+  def foo(self, x: int) -> int:
+    return x + 1
+
+
+class TestingEventHandler(pg.Object, interface.EventHandler):
   """Testing environment event handler for unit tests."""
 
   log_sandbox_status: bool = False
@@ -228,6 +273,8 @@ class TestingEventHandler(
   log_session_setup: bool = False
   log_housekeep: bool = False
 
+  __test__ = False
+
   def _on_bound(self) -> None:
     super()._on_bound()
     self._logs = []
@@ -281,7 +328,6 @@ class TestingEventHandler(
 
   def on_sandbox_start(
       self,
-      environment: interface.Environment,
       sandbox: interface.Sandbox,
       duration: float,
       error: BaseException | None
@@ -291,7 +337,6 @@ class TestingEventHandler(
 
   def on_sandbox_status_change(
       self,
-      environment: interface.Environment,
       sandbox: interface.Sandbox,
       old_status: interface.Sandbox.Status,
       new_status: interface.Sandbox.Status,
@@ -306,7 +351,6 @@ class TestingEventHandler(
 
   def on_sandbox_shutdown(
       self,
-      environment: interface.Environment,
       sandbox: interface.Sandbox,
       duration: float,
       lifetime: float,
@@ -315,9 +359,53 @@ class TestingEventHandler(
     assert duration > 0 and lifetime is not None
     self._add_message(f'[{sandbox.id}] sandbox shutdown', error)
 
+  def on_sandbox_session_start(
+      self,
+      sandbox: interface.Sandbox,
+      session_id: str,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox session starts."""
+    assert duration > 0
+    self._add_message(
+        f'[{sandbox.id}] session {session_id!r} started', error
+    )
+
+  def on_sandbox_session_end(
+      self,
+      sandbox: interface.Sandbox,
+      session_id: str,
+      duration: float,
+      lifetime: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox session ends."""
+    assert duration > 0 and lifetime > 0
+    self._add_message(
+        f'[{sandbox.id}] session {session_id!r} ended', error
+    )
+
+  def on_sandbox_activity(
+      self,
+      name: str,
+      sandbox: interface.Sandbox,
+      session_id: str | None,
+      duration: float,
+      error: BaseException | None,
+      *,
+      code: str | None = None,
+      **kwargs
+  ) -> None:
+    """Called when a sandbox activity is performed."""
+    del kwargs
+    log_id = f'{sandbox.id}@{session_id or "<idle>"}'
+    self._add_message(
+        f'[{log_id}] {name}: {code}', error
+    )
+
   def on_sandbox_housekeep(
       self,
-      environment: interface.Environment,
       sandbox: interface.Sandbox,
       counter: int,
       duration: float,
@@ -332,8 +420,6 @@ class TestingEventHandler(
 
   def on_feature_setup(
       self,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
       feature: interface.Feature,
       duration: float,
       error: BaseException | None
@@ -342,13 +428,11 @@ class TestingEventHandler(
     assert duration > 0
     if self.log_feature_setup:
       self._add_message(
-          f'[{sandbox.id}/{feature.name}] feature setup', error
+          f'[{feature.id}] feature setup', error
       )
 
   def on_feature_teardown(
       self,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
       feature: interface.Feature,
       duration: float,
       error: BaseException | None
@@ -357,13 +441,11 @@ class TestingEventHandler(
     assert duration > 0
     if self.log_feature_setup:
       self._add_message(
-          f'[{sandbox.id}/{feature.name}] feature teardown', error
+          f'[{feature.id}] feature teardown', error
       )
 
   def on_feature_setup_session(
       self,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
       feature: interface.Feature,
       session_id: str | None,
       duration: float,
@@ -373,13 +455,12 @@ class TestingEventHandler(
     assert duration > 0
     if self.log_session_setup:
       self._add_message(
-          f'[{sandbox.id}/{feature.name}] feature setup session', error
+          f'[{feature.id}@{session_id or "<idle>"}] feature setup session',
+          error
       )
 
   def on_feature_teardown_session(
       self,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
       feature: interface.Feature,
       session_id: str,
       duration: float,
@@ -389,77 +470,38 @@ class TestingEventHandler(
     assert duration > 0
     if self.log_session_setup:
       self._add_message(
-          f'[{sandbox.id}/{feature.name}] feature teardown session', error
+          f'[{feature.id}@{session_id}] feature teardown session', error
       )
 
-  def on_feature_housekeep(
+  def on_feature_activity(
       self,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
+      name: str,
       feature: interface.Feature,
-      counter: int,
+      session_id: str | None,
       duration: float,
       error: BaseException | None,
+      *,
+      code: str | None = None,
       **kwargs
   ) -> None:
-    """Called when a sandbox feature is housekeeping."""
-    assert duration > 0
-    if self.log_housekeep:
-      self._add_message(
-          f'[{sandbox.id}/{feature.name}] feature housekeeping {counter}', error
-      )
-
-  def on_session_start(
-      self,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
-      session_id: str,
-      duration: float,
-      error: BaseException | None
-  ) -> None:
-    """Called when a sandbox session starts."""
-    assert duration > 0
-    self._add_message(
-        f'[{sandbox.id}] session {session_id!r} started', error
-    )
-
-  def on_session_end(
-      self,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
-      session_id: str,
-      duration: float,
-      lifetime: float,
-      error: BaseException | None
-  ) -> None:
-    """Called when a sandbox session ends."""
-    assert duration > 0 and lifetime > 0
+    """Called when a sandbox activity is performed."""
+    del kwargs
+    log_id = f'{feature.id}@{session_id or "<idle>"}'
     self._add_message(
-        f'[{sandbox.id}] session {session_id!r} ended', error
+        f'[{log_id}] {name}: {code}', error
     )
 
-  def on_sandbox_activity(
+  def on_feature_housekeep(
       self,
-      name: str,
-      environment: interface.Environment,
-      sandbox: interface.Sandbox,
-      feature: interface.Feature | None,
-      session_id: str | None,
+      feature: interface.Feature,
+      counter: int,
       duration: float,
       error: BaseException | None,
-      *,
-      code: str | None = None,
       **kwargs
   ) -> None:
-    """Called when a sandbox activity is performed."""
-    del environment, kwargs
-    if session_id is None:
-      log_id = sandbox.id
-    else:
-      log_id = f'{sandbox.id}/{session_id}'
-
-    if feature is not None:
-      log_id = f'{log_id}/{feature.name}'
-    self._add_message(
-        f'[{log_id}] {name}: {code}', error
-    )
+    """Called when a sandbox feature is housekeeping."""
+    assert duration > 0
+    if self.log_housekeep:
+      self._add_message(
+          f'[{feature.id}] feature housekeeping {counter}', error
+      )

{langfun-0.1.2.dev202510230805.dist-info → langfun-0.1.2.dev202511270805.dist-info}/METADATA

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