langfun 0.1.2.dev202509120804__py3-none-any.whl → 0.1.2.dev202512150805__py3-none-any.whl

langfun/env/interface_test.py

@@ -0,0 +1,153 @@
+# Copyright 2025 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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
+
+
+class IdTest(unittest.TestCase):
+
+  def test_environment_id(self):
+    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'),
+        '/tmp/env_1/ab_c_def'
+    )
+    self.assertIsNone(env_id.working_dir(root_dir=None))
+
+  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/image:2025_01_01_00_00_00:sandbox')
+    self.assertEqual(
+        sandbox_id.working_dir(root_dir='/tmp'),
+        '/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.py

@@ -0,0 +1,59 @@
+# Copyright 2025 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""Load balancers for environments."""
+
+import abc
+import threading
+
+from langfun.env import interface
+import pyglove as pg
+
+
+class LoadBalancer(pg.Object):
+  """Base class for load balancers."""
+
+  @abc.abstractmethod
+  def acquire(self, sandbox_pool: list[interface.Sandbox]) -> interface.Sandbox:
+    """Acquires a free sandbox from a pool of sandboxes.
+
+    The load balancer will pick a sandbox from the pool and mark it as pending.
+
+    Args:
+      sandbox_pool: The pool of sandboxes to pick from.
+
+    Raises:
+      IndexError: If all sandboxes in the pool are either busy or dead.
+    """
+
+
+class RoundRobin(LoadBalancer):
+  """Round robin load balancer."""
+
+  def _on_bound(self):
+    super()._on_bound()
+    self._counter = 0
+    self._acquire_lock = threading.Lock()
+
+  def acquire(self, sandbox_pool: list[interface.Sandbox]) -> interface.Sandbox:
+    """Returns a free sandbox from the pool."""
+    with self._acquire_lock:
+      for _ in range(len(sandbox_pool)):
+        sandbox = sandbox_pool[self._counter % len(sandbox_pool)]
+        self._counter = self._counter + 1
+        if sandbox.status == interface.Sandbox.Status.READY:
+          # Mark the sandbox as acquired to prevent it from being acquired by
+          # other threads.
+          sandbox.set_acquired()
+          return sandbox
+    raise IndexError('No free sandbox in the pool.')

langfun/env/load_balancers_test.py

@@ -0,0 +1,141 @@
+# Copyright 2025 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 concurrent.futures
+import contextlib
+import time
+from typing import Any, Iterator
+import unittest
+
+from langfun.env import interface
+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()
+    self._session_id = None
+
+  @property
+  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
+    )
+
+  @property
+  def environment(self) -> interface.Environment:
+    raise NotImplementedError()
+
+  @property
+  def features(self) -> dict[str, interface.Feature]:
+    raise NotImplementedError()
+
+  @property
+  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)
+
+  def set_acquired(self) -> None:
+    self.set_status(self.Status.ACQUIRED)
+
+  def start(self) -> None:
+    pass
+
+  def shutdown(self) -> None:
+    pass
+
+  def start_session(self, session_id: str) -> None:
+    self._session_id = session_id
+
+  def end_session(self, session_id: str) -> None:
+    self._session_id = None
+
+  @property
+  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):
+
+  def test_basic(self):
+    sandbox_pool = [
+        TestingSandbox('0', interface.Sandbox.Status.OFFLINE),
+        TestingSandbox('1', interface.Sandbox.Status.SETTING_UP),
+        TestingSandbox('2', interface.Sandbox.Status.IN_SESSION),
+        TestingSandbox('3', status=interface.Sandbox.Status.READY),
+        TestingSandbox('4', status=interface.Sandbox.Status.READY),
+    ]
+    lb = load_balancers.RoundRobin()
+    sandbox = lb.acquire(sandbox_pool)
+    self.assertIs(sandbox, sandbox_pool[3])
+    self.assertEqual(sandbox.status, interface.Sandbox.Status.ACQUIRED)
+
+    sandbox = lb.acquire(sandbox_pool)
+    self.assertIs(sandbox, sandbox_pool[4])
+    self.assertEqual(sandbox.status, interface.Sandbox.Status.ACQUIRED)
+
+    sandbox_pool[0].set_status(interface.Sandbox.Status.READY)
+    sandbox = lb.acquire(sandbox_pool)
+    self.assertIs(sandbox, sandbox_pool[0])
+    self.assertEqual(sandbox.status, interface.Sandbox.Status.ACQUIRED)
+
+    with self.assertRaisesRegex(IndexError, 'No free sandbox in the pool.'):
+      lb.acquire(sandbox_pool)
+
+  def test_thread_safety(self):
+    sandbox_pool = [TestingSandbox(str(i)) for i in range(64)]
+
+    lb = load_balancers.RoundRobin()
+
+    def _thread_func(i):
+      sandbox = lb.acquire(sandbox_pool)
+      time.sleep(0.1)
+      sandbox.set_status(interface.Sandbox.Status.IN_SESSION)
+      time.sleep(0.1)
+      sandbox.set_status(interface.Sandbox.Status.OFFLINE)
+      time.sleep(0.1)
+      sandbox.set_status(interface.Sandbox.Status.READY)
+      return i
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=64) as executor:
+      for i, o in enumerate(executor.map(_thread_func, range(1024))):
+        self.assertEqual(o, i)
+
+
+if __name__ == '__main__':
+  unittest.main()