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

langfun/env/__init__.py

@@ -13,16 +13,12 @@
 # limitations under the License.
 """Environment for LLM agents."""
 
-# pylint: disable=g-importing-member, g-bad-import-order
-from langfun.env.interface import EnvironmentId
-from langfun.env.interface import SandboxId
-
+# pylint: disable=g-importing-member, g-bad-import-order, g-import-not-at-top
 from langfun.env.interface import EnvironmentError   # pylint: disable=redefined-builtin
 from langfun.env.interface import EnvironmentOutageError
 from langfun.env.interface import EnvironmentOverloadError
 from langfun.env.interface import SandboxError
 from langfun.env.interface import SandboxStateError
-from langfun.env.interface import EnvironmentEventHandler
 
 from langfun.env.interface import Environment
 from langfun.env.interface import Sandbox
@@ -35,4 +31,7 @@ from langfun.env.base_feature import BaseFeature
 from langfun.env import load_balancers
 from langfun.env.load_balancers import LoadBalancer
 
+from langfun.env import event_handlers
+EventHandler = event_handlers.EventHandler
+
 # Google-internal imports.

langfun/env/base_environment.py

@@ -34,6 +34,7 @@ import langfun.core as lf
 from langfun.env import base_sandbox
 from langfun.env import interface
 from langfun.env import load_balancers
+from langfun.env.event_handlers import base as event_handler_base
 import pyglove as pg
 
 
@@ -65,10 +66,20 @@ class BaseEnvironment(interface.Environment):
   load_balancer: Annotated[
       load_balancers.LoadBalancer,
       (
-          'The load balancer for the environment.'
+          'The load balancer for the environment to acquire sandboxes.'
       )
   ] = load_balancers.RoundRobin()
 
+  sandbox_keepalive_interval: Annotated[
+      float | None,
+      (
+          'The interval in seconds to send keepalive pings to sandboxes. '
+          'If None, sandbox keepalive is disabled. Please note that sandbox '
+          'keepalive is different from feature housekeeping. Usually sandbox '
+          'keepalive and feature housekeeping are different operations.'
+      )
+  ] = None
+
   proactive_session_setup: Annotated[
       bool,
       (
@@ -79,6 +90,13 @@ class BaseEnvironment(interface.Environment):
       )
   ] = True
 
+  event_handlers: Annotated[
+      list[event_handler_base.EventHandler],
+      (
+          'User handler for the environment events.'
+      )
+  ] = []
+
   outage_grace_period: Annotated[
       float,
       (
@@ -97,13 +115,15 @@ class BaseEnvironment(interface.Environment):
       )
   ] = 10.0
 
-  stats_report_interval: Annotated[
+  housekeep_interval: Annotated[
       float,
       (
-          'The interval in seconds for reporting the environment stats. '
-          'If 0, stats will not be reported.'
+          'The interval in seconds for environment housekeeping. It recycles '
+          'the dead sandboxes in the pool. This interval is the minimal time '
+          'to detect outage while there is no request to obtain new sandboxes.'
+          'This is applicable only when the environment enables pooling.'
       )
-  ] = 60.0
+  ] = 10.0
 
   pool_operation_max_parallelism: Annotated[
       int,
@@ -145,6 +165,7 @@ class BaseEnvironment(interface.Environment):
       sandbox_id: str,
       reusable: bool,
       proactive_session_setup: bool,
+      keepalive_interval: float | None,
   ) -> base_sandbox.BaseSandbox:
     """Creates a sandbox with the given identifier.
 
@@ -153,6 +174,8 @@ class BaseEnvironment(interface.Environment):
       reusable: Whether the sandbox is reusable across user sessions.
       proactive_session_setup: Whether the sandbox performs session setup work
         before a user session is started.
+      keepalive_interval: Interval to ping the sandbox for keeping it alive.
+        If None, the sandbox will not be pinged.
 
     Returns:
       The created sandbox.
@@ -194,6 +217,7 @@ class BaseEnvironment(interface.Environment):
   def _start(self) -> None:
     """Implementation of starting the environment."""
     if self.min_pool_size > 0:
+      # Pre-allocate the sandbox pool before usage.
       self._sandbox_pool = [None] * self.min_pool_size
       for i, sandbox, _ in lf.concurrent_map(
           lambda i: self._bring_up_sandbox_with_retry(
@@ -207,12 +231,15 @@ class BaseEnvironment(interface.Environment):
           ),
       ):
         self._sandbox_pool[i] = sandbox
+
     self._next_sandbox_id = len(self._sandbox_pool)
-    self._housekeep_thread = threading.Thread(
-        target=self._housekeep_loop, daemon=True
-    )
-    self._housekeep_counter = 0
-    self._housekeep_thread.start()
+
+    if self.enable_pooling:
+      self._housekeep_thread = threading.Thread(
+          target=self._housekeep_loop, daemon=True
+      )
+      self._housekeep_counter = 0
+      self._housekeep_thread.start()
 
   def _shutdown(self) -> None:
     """Implementation of shutting down the environment."""
@@ -256,7 +283,7 @@ class BaseEnvironment(interface.Environment):
   @property
   def enable_pooling(self) -> bool:
     """Returns whether the environment enables pooling."""
-    return self.min_pool_size > 0
+    return self.max_pool_size > 0
 
   @property
   def status(self) -> interface.Environment.Status:
@@ -314,11 +341,6 @@ class BaseEnvironment(interface.Environment):
       self._start_time = time.time()
       self._set_status(self.Status.ONLINE)
       self.on_start(duration=time.time() - starting_time)
-
-      pg.logging.info(
-          '[%s]: %s started in %.2f seconds.',
-          self.id, self.__class__.__name__, time.time() - starting_time
-      )
     except BaseException as e:
       self.on_start(duration=time.time() - starting_time, error=e)
       self.shutdown()
@@ -338,13 +360,8 @@ class BaseEnvironment(interface.Environment):
     self._set_status(self.Status.SHUTTING_DOWN)
 
     try:
-      with pg.timeit('env.shutdown') as t:
-        self._shutdown()
+      self._shutdown()
       self.on_shutdown()
-      pg.logging.info(
-          '[%s]: %s shutdown in %.2f seconds.',
-          self.id, self.__class__.__name__, t.elapse
-      )
     except BaseException as e:  # pylint: disable=broad-except
       self.on_shutdown(error=e)
       raise e
@@ -416,6 +433,7 @@ class BaseEnvironment(interface.Environment):
         sandbox_id=sandbox_id,
         reusable=self.enable_pooling,
         proactive_session_setup=self.proactive_session_setup,
+        keepalive_interval=self.sandbox_keepalive_interval,
     )
     for handler in self.event_handlers:
       sandbox.add_event_handler(handler)
@@ -482,19 +500,8 @@ class BaseEnvironment(interface.Environment):
 
   def _housekeep_loop(self) -> None:
     """Housekeeping loop for the environment."""
-    pg.logging.info(
-        '[%s]: %s maintenance thread started.', self.id, self.__class__.__name__
-    )
-    stats_report_time = time.time()
     while self._status not in (self.Status.SHUTTING_DOWN, self.Status.OFFLINE):
       housekeep_start_time = time.time()
-      if time.time() - stats_report_time > self.stats_report_interval:
-        pg.logging.info(
-            '[%s] %s stats: %s.',
-            self.id, self.__class__.__name__, self.stats()
-        )
-        stats_report_time = time.time()
-
       dead_pool_indices = [
           i for i, s in enumerate(self._sandbox_pool)
           if s.status == interface.Sandbox.Status.OFFLINE
@@ -514,7 +521,7 @@ class BaseEnvironment(interface.Environment):
         break
       self._housekeep_counter += 1
       self.on_housekeep(time.time() - housekeep_start_time)
-      time.sleep(1)
+      time.sleep(self.housekeep_interval)
 
   def _replace_dead_sandboxes(self, dead_pool_indices: list[int]) -> bool:
     """Replaces a dead sandbox with a new one.

langfun/env/base_feature.py

@@ -23,6 +23,7 @@ the `Environment` and `Sandbox` interfaces directly.
 """
 
 import functools
+import os
 import time
 from typing import Annotated, Callable
 
@@ -106,6 +107,14 @@ class BaseFeature(interface.Feature):
     assert self._sandbox is not None, 'Feature has not been set up yet.'
     return self._sandbox
 
+  @property
+  def working_dir(self) -> str | None:
+    """Returns the working directory of the feature."""
+    sandbox_workdir = self.sandbox.working_dir
+    if sandbox_workdir is None:
+      return None
+    return os.path.join(sandbox_workdir, self.name)
+
   #
   # Setup and teardown of the feature.
   #
@@ -204,7 +213,7 @@ class BaseFeature(interface.Feature):
       self,
       name: str,
       duration: float,
-      error: BaseException | None,
+      error: BaseException | None = None,
       **kwargs
   ) -> None:
     """Called when a sandbox activity is performed."""

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.'
   ]
 
@@ -242,14 +243,14 @@ class BaseSandbox(interface.Sandbox):
 
   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)
@@ -357,10 +358,6 @@ class BaseSandbox(interface.Sandbox):
 
       duration = time.time() - starting_time
       self.on_start(duration)
-      pg.logging.info(
-          '[%s]: Sandbox started in %.2f seconds.',
-          self.id, duration
-      )
     except BaseException as e:  # pylint: disable=broad-except
       duration = time.time() - starting_time
       pg.logging.error(
@@ -422,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):
@@ -433,15 +429,6 @@ 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
@@ -656,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
 
@@ -687,9 +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()
@@ -731,7 +732,7 @@ class BaseSandbox(interface.Sandbox):
 
       self._housekeep_counter += 1
       self.on_housekeep(time.time() - housekeep_start)
-      time.sleep(1)
+      time.sleep(_next_housekeep_wait_time())
 
   #
   # Event handlers subclasses can override.

langfun/env/base_test.py

@@ -18,12 +18,13 @@ import unittest
 from langfun.env import base_sandbox
 from langfun.env import interface
 from langfun.env import test_utils
+from langfun.env.event_handlers import base as event_handler_base
 
 
 TestingEnvironment = test_utils.TestingEnvironment
 TestingSandbox = test_utils.TestingSandbox
 TestingFeature = test_utils.TestingFeature
-TestingEnvironmentEventHandler = test_utils.TestingEnvironmentEventHandler
+TestingEventHandler = test_utils.TestingEventHandler
 
 
 class EnvironmentTests(unittest.TestCase):
@@ -42,9 +43,8 @@ class EnvironmentTests(unittest.TestCase):
     self.assertEqual(env.min_pool_size, 0)
     self.assertEqual(env.max_pool_size, 0)
     self.assertEqual(env.sandbox_pool, [])
-    self.assertEqual(env.id, interface.EnvironmentId('testing-env'))
+    self.assertEqual(env.id, interface.Environment.Id('testing-env'))
     self.assertEqual(env.outage_grace_period, 1)
-    self.assertEqual(env.stats_report_interval, 60)
     self.assertEqual(env.features['test_feature'].name, 'test_feature')
 
     self.assertIsNone(env.start_time)
@@ -60,12 +60,16 @@ class EnvironmentTests(unittest.TestCase):
 
       with env.sandbox('session1') as sb:
         self.assertEqual(
-            sb.id, interface.SandboxId(environment_id=env.id, sandbox_id='0')
+            sb.id, interface.Sandbox.Id(environment_id=env.id, sandbox_id='0')
         )
         self.assertEqual(sb.session_id, 'session1')
         self.assertEqual(sb.working_dir, '/tmp/testing-env/0')
         self.assertTrue(sb.is_online)
         self.assertIs(sb.test_feature, sb.features['test_feature'])
+        self.assertEqual(
+            sb.test_feature.working_dir,
+            '/tmp/testing-env/0/test_feature'
+        )
         with self.assertRaises(AttributeError):
           _ = sb.test_feature2
       self.assertFalse(sb.is_online)
@@ -74,14 +78,26 @@ class EnvironmentTests(unittest.TestCase):
       with self.assertRaises(AttributeError):
         _ = env.test_feature2
 
+  def test_del(self):
+    env = TestingEnvironment(
+        features={'test_feature': TestingFeature()},
+        pool_size=0,
+        outage_grace_period=1,
+        outage_retry_interval=0,
+        sandbox_keepalive_interval=0,
+    )
+    env.start()
+    sb = env.acquire()
+    del sb
+    del env
+
   def test_acquire_env_offline(self):
     env = TestingEnvironment(
         features={'test_feature': TestingFeature()},
         pool_size=0,
         outage_grace_period=1,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with self.assertRaises(interface.EnvironmentOutageError):
       env.acquire()
@@ -92,12 +108,14 @@ class EnvironmentTests(unittest.TestCase):
         pool_size=0,
         outage_grace_period=1,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with env:
       sb = env.acquire()
       self.assertEqual(sb.status, interface.Sandbox.Status.ACQUIRED)
+      self.assertIsNone(env.working_dir)
+      self.assertIsNone(sb.working_dir)
+      self.assertIsNone(sb.test_feature.working_dir)
 
   def test_acquire_no_pooling_with_error(self):
     env = TestingEnvironment(
@@ -109,8 +127,7 @@ class EnvironmentTests(unittest.TestCase):
         pool_size=0,
         outage_grace_period=1,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with env:
       with self.assertRaises(interface.EnvironmentOutageError):
@@ -122,8 +139,7 @@ class EnvironmentTests(unittest.TestCase):
         pool_size=1,
         outage_grace_period=1,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with env:
       sb = env.acquire()
@@ -135,8 +151,7 @@ class EnvironmentTests(unittest.TestCase):
         pool_size=1,
         outage_grace_period=1,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with env:
       sb = env.acquire()
@@ -150,17 +165,76 @@ class EnvironmentTests(unittest.TestCase):
         pool_size=(1, 3),
         outage_grace_period=1,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with env:
       self.assertEqual(len(env.sandbox_pool), 1)
+      self.assertEqual(
+          env.stats(),
+          {
+              'sandbox': {
+                  'created': 0,
+                  'setting_up': 0,
+                  'ready': 1,
+                  'acquired': 0,
+                  'in_session': 0,
+                  'exiting_session': 0,
+                  'shutting_down': 0,
+                  'offline': 0,
+              }
+          }
+      )
       sb = env.acquire()
       self.assertEqual(sb.status, interface.Sandbox.Status.ACQUIRED)
+      self.assertEqual(
+          env.stats(),
+          {
+              'sandbox': {
+                  'created': 0,
+                  'setting_up': 0,
+                  'ready': 0,
+                  'acquired': 1,
+                  'in_session': 0,
+                  'exiting_session': 0,
+                  'shutting_down': 0,
+                  'offline': 0,
+              }
+          }
+      )
       self.assertEqual(len(env.sandbox_pool), 1)
       sb2 = env.acquire()
       self.assertEqual(sb2.status, interface.Sandbox.Status.ACQUIRED)
       self.assertEqual(len(env.sandbox_pool), 2)
+      self.assertEqual(
+          env.stats(),
+          {
+              'sandbox': {
+                  'created': 0,
+                  'setting_up': 0,
+                  'ready': 0,
+                  'acquired': 2,
+                  'in_session': 0,
+                  'exiting_session': 0,
+                  'shutting_down': 0,
+                  'offline': 0,
+              }
+          }
+      )
+    self.assertEqual(
+        env.stats(),
+        {
+            'sandbox': {
+                'created': 0,
+                'setting_up': 0,
+                'ready': 0,
+                'acquired': 0,
+                'in_session': 0,
+                'exiting_session': 0,
+                'shutting_down': 0,
+                'offline': 0,
+            }
+        }
+    )
 
   def test_acquire_with_growing_pool_failure(self):
     env = TestingEnvironment(
@@ -168,8 +242,7 @@ class EnvironmentTests(unittest.TestCase):
         pool_size=(1, 3),
         outage_grace_period=1,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with env:
       self.assertEqual(len(env.sandbox_pool), 1)
@@ -184,15 +257,14 @@ class EnvironmentTests(unittest.TestCase):
       with self.assertRaises(interface.EnvironmentOutageError):
         env.acquire()
 
-  def test_maintenance_error(self):
+  def test_housekeep_error(self):
     env = TestingEnvironment(
         features={'test_feature': TestingFeature()},
         pool_size=1,
         proactive_session_setup=True,
         outage_grace_period=1,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with env:
       self.assertEqual(len(env.sandbox_pool), 1)
@@ -216,7 +288,7 @@ class SandboxStatusTests(unittest.TestCase):
 
   def setUp(self):
     super().setUp()
-    self.event_handler = TestingEnvironmentEventHandler(
+    self.event_handler = TestingEventHandler(
         log_sandbox_status=True,
         log_feature_setup=True,
         log_session_setup=True,
@@ -246,6 +318,7 @@ class SandboxStatusTests(unittest.TestCase):
             'feature2': TestingFeature(),
         },
     )
+    self.assertFalse(env.enable_pooling)
     with env:
       with env.sandbox('session1') as sb:
         sb.shell('echo "hello"')
@@ -294,6 +367,7 @@ class SandboxStatusTests(unittest.TestCase):
         pool_size=1,
         proactive_session_setup=True,
     )
+    self.assertTrue(env.enable_pooling)
     with env:
       with env.sandbox('session1') as sb:
         sb.shell('echo "hello"')
@@ -343,7 +417,7 @@ class SandboxStatusTests(unittest.TestCase):
         features={'test_feature': TestingFeature(setup_session_delay=0.5)},
         pool_size=1,
     )
-    event_handler = TestingEnvironmentEventHandler(
+    event_handler = TestingEventHandler(
         log_sandbox_status=True,
         log_feature_setup=True,
         log_session_setup=True,
@@ -919,6 +993,7 @@ class SandboxStatusTests(unittest.TestCase):
         self.assertEqual(len(sb.state_errors), 0)
         sb.shell('echo bar')
         self.assertEqual(sb.status, interface.Sandbox.Status.IN_SESSION)
+      sb.wait_until_not(interface.Sandbox.Status.SETTING_UP)
       self.assertEqual(sb.status, interface.Sandbox.Status.READY)
       self.assertEqual(
           self.event_handler.logs,
@@ -1020,7 +1095,7 @@ class SandboxActivityTests(unittest.TestCase):
     env = TestingEnvironment(
         features={'test_feature': TestingFeature(housekeep_interval=0)},
         pool_size=1,
-        keepalive_interval=0,
+        sandbox_keepalive_interval=0,
     )
     with env:
       with env.sandbox('session1') as sb:
@@ -1037,7 +1112,7 @@ class SandboxActivityTests(unittest.TestCase):
         pool_size=1,
         outage_grace_period=0,
         outage_retry_interval=0,
-        keepalive_interval=0,
+        sandbox_keepalive_interval=0,
     )
     with env:
       with env.sandbox('session1') as sb:
@@ -1052,7 +1127,7 @@ class SandboxActivityTests(unittest.TestCase):
         while sb.housekeep_counter == housekeep_count or (
             sb.status == interface.Sandbox.Status.IN_SESSION
         ):
-          time.sleep(0.1)
+          time.sleep(0.01)
         self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
 
   def test_remove_event_handler(self):
@@ -1061,10 +1136,9 @@ class SandboxActivityTests(unittest.TestCase):
         pool_size=1,
         outage_grace_period=0,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
-    event_handler = TestingEnvironmentEventHandler()
+    event_handler = TestingEventHandler()
     with env:
       with env.sandbox('session1') as sb:
         sb.add_event_handler(event_handler)
@@ -1082,21 +1156,20 @@ class SandboxServiceTests(unittest.TestCase):
   def setUp(self):
     super().setUp()
     self.maxDiff = None
-    self.event_handler = TestingEnvironmentEventHandler()
+    self.event_handler = TestingEventHandler()
     self.env = TestingEnvironment(
         features={'test_feature': TestingFeature()},
         pool_size=0,
         outage_grace_period=0,
         outage_retry_interval=0,
-        keepalive_interval=0,
+        sandbox_keepalive_interval=0,
         event_handlers=[self.event_handler],
-        stats_report_interval=1,
         random_seed=1,
     )
 
   def test_service_call_activity_log(self):
 
-    class CustomEventHandler(interface.EnvironmentEventHandler):
+    class CustomEventHandler(event_handler_base.EventHandler):
 
       def __init__(self):
         self.calls = []

langfun/env/event_handlers/__init__.py

@@ -0,0 +1,12 @@
+"""Environment event handlers."""
+
+# pylint: disable=g-importing-member
+# pylint: disable=g-bad-import-order
+
+from langfun.env.event_handlers.base import EventHandler
+
+from langfun.env.event_handlers.event_logger import EventLogger
+from langfun.env.event_handlers.event_logger import ConsoleEventLogger
+
+# pylint: enable=g-importing-member
+# pylint: enable=g-bad-import-order