langfun 0.1.2.dev202509150805__py3-none-any.whl → 0.1.2.dev202509170804__py3-none-any.whl

langfun/env/interface_test.py

@@ -0,0 +1,43 @@
+# 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 unittest
+from langfun.env import interface
+
+
+class IdTest(unittest.TestCase):
+
+  def test_environment_id(self):
+    env_id = interface.EnvironmentId('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.SandboxId(
+        environment_id=interface.EnvironmentId('env'),
+        sandbox_id='sandbox'
+    )
+    self.assertEqual(str(sandbox_id), 'env/sandbox')
+    self.assertEqual(
+        sandbox_id.working_dir(root_dir='/tmp'),
+        '/tmp/env/sandbox'
+    )
+    self.assertIsNone(sandbox_id.working_dir(root_dir=None))
+
+
+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.is_alive and not sandbox.is_busy and not sandbox.is_pending:
+          # Mark the sandbox as pending so that it will not be acquired by other
+          # threads.
+          sandbox.set_pending()
+          return sandbox
+    raise IndexError('No free sandbox in the pool.')

langfun/env/load_balancers_test.py

@@ -0,0 +1,157 @@
+# 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 time
+import unittest
+
+from langfun.env import interface
+from langfun.env import load_balancers
+
+
+class TestingSandbox(interface.Sandbox):
+  sandbox_id: str
+  is_alive: bool = True
+  is_pending: bool = False
+  is_busy: bool = False
+
+  def _on_bound(self) -> None:
+    super()._on_bound()
+    self._session_id = None
+
+  @property
+  def id(self) -> interface.SandboxId:
+    return interface.SandboxId(
+        environment_id=interface.EnvironmentId('testing-env'),
+        sandbox_id=self.sandbox_id
+    )
+
+  @property
+  def environment(self) -> interface.Environment:
+    raise NotImplementedError()
+
+  @property
+  def features(self) -> dict[str, interface.Feature]:
+    raise NotImplementedError()
+
+  def set_pending(self, pending: bool = True) -> None:
+    self.rebind(
+        is_pending=pending, skip_notification=True, raise_on_no_change=False
+    )
+
+  def set_busy(self, busy: bool = True) -> None:
+    self.rebind(
+        is_busy=busy, skip_notification=True, raise_on_no_change=False
+    )
+
+  def set_alive(self, alive: bool = True) -> None:
+    self.rebind(
+        is_alive=alive, skip_notification=True, raise_on_no_change=False
+    )
+
+  def start(self) -> None:
+    self.set_alive()
+
+  def shutdown(self) -> None:
+    self.set_alive(False)
+
+  def ping(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
+
+
+class RoundRobinTest(unittest.TestCase):
+
+  def test_basic(self):
+    sandbox_pool = [
+        TestingSandbox(
+            '0',
+            is_alive=False,
+            is_pending=False,
+            is_busy=False,
+        ),
+        TestingSandbox(
+            '1',
+            is_alive=True,
+            is_pending=True,
+            is_busy=False,
+        ),
+        TestingSandbox(
+            '2',
+            is_alive=True,
+            is_pending=False,
+            is_busy=True,
+        ),
+        TestingSandbox(
+            '3',
+            is_alive=True,
+            is_pending=False,
+            is_busy=False,
+        ),
+        TestingSandbox(
+            '4',
+            is_alive=True,
+            is_pending=False,
+            is_busy=False,
+        ),
+    ]
+    lb = load_balancers.RoundRobin()
+    sandbox = lb.acquire(sandbox_pool)
+    self.assertIs(sandbox, sandbox_pool[3])
+    self.assertTrue(sandbox.is_pending)
+
+    sandbox = lb.acquire(sandbox_pool)
+    self.assertIs(sandbox, sandbox_pool[4])
+    self.assertTrue(sandbox.is_pending)
+
+    sandbox_pool[0].set_alive()
+    sandbox = lb.acquire(sandbox_pool)
+    self.assertIs(sandbox, sandbox_pool[0])
+    self.assertTrue(sandbox.is_pending)
+
+    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_busy()
+      sandbox.set_pending(False)
+      time.sleep(0.1)
+      sandbox.set_busy(False)
+      sandbox.set_alive(False)
+      time.sleep(0.1)
+      sandbox.set_alive()
+      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()

{langfun-0.1.2.dev202509150805.dist-info → langfun-0.1.2.dev202509170804.dist-info}/METADATA

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

{langfun-0.1.2.dev202509150805.dist-info → langfun-0.1.2.dev202509170804.dist-info}/RECORD

@@ -165,8 +165,17 @@ langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fik
 langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3ABKzup8jReCljHeK4,2162
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=Ot__1P1M8oJfoTp-M9-PQ6HUXqZKyMwvZ5f7yQ3yfyM,2326
-langfun-0.1.2.dev202509150805.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202509150805.dist-info/METADATA,sha256=WGTqnnkjfg9kGoytP4lQ57CpSlQAHLULz2jkcppH3rk,7380
-langfun-0.1.2.dev202509150805.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-langfun-0.1.2.dev202509150805.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202509150805.dist-info/RECORD,,
+langfun/env/__init__.py,sha256=zgiONH9pIpNTEL623jRmW2mMI-GhThzEegZoys_znrQ,1554
+langfun/env/base_environment.py,sha256=BMdSpsk0uMCPklrZ7XQPpJ03IL3ZEzCIs3RlUfG8yNo,14215
+langfun/env/base_feature.py,sha256=SVvgoA68vafZZct5ZuHjPHblrCVv2Es4bLg5wgVTMww,4653
+langfun/env/base_sandbox.py,sha256=oSuUSd3zVgBxH14r9uX9SelyoXHzQQWzGHEiain8eKo,13690
+langfun/env/base_test.py,sha256=FuytEQJeKnxX_VEW0axdEoCg1A1d2rIVQss0qxoPPAQ,26739
+langfun/env/interface.py,sha256=4Q6ewOtj8XJyy8OEQqSYwogjzfvivFOljGS5txhlx_U,21534
+langfun/env/interface_test.py,sha256=nrqFLtpHd0vnOjoIX9cjV5WD1c-7J3WW29aogNu_94A,1395
+langfun/env/load_balancers.py,sha256=jbT3pB7QpPVovbXeJK1PNohUcrYrJNtWU8KaWzxO9nI,1963
+langfun/env/load_balancers_test.py,sha256=bAjC89zxmROpkl56BqvMU_yKOJ03Gtu8x8j56HO5HNc,4124
+langfun-0.1.2.dev202509170804.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202509170804.dist-info/METADATA,sha256=Ps_Bme20ro1JnbT4nQj2Cs-euv4K-DpwcTg4BQr37u4,7380
+langfun-0.1.2.dev202509170804.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+langfun-0.1.2.dev202509170804.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202509170804.dist-info/RECORD,,