langfun 0.1.2.dev202509290805__py3-none-any.whl → 0.1.2.dev202509300805__py3-none-any.whl

langfun/env/event_handlers/metric_writer.py

@@ -0,0 +1,545 @@
+# 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.
+"""Environment event logger."""
+
+from langfun.env.event_handlers import base
+import pyglove as pg
+
+
+_METRIC_NAMESPACE = '/langfun/env'
+
+
+class MetricWriter(pg.Object, base.EventHandler):
+  """Event handler for streamz metrics."""
+
+  def _get_counter(
+      self,
+      name: str,
+      description: str,
+      parameters: dict[str, type[str]] | None = None
+  ) -> pg.monitoring.Counter:
+    return self._metric_collection.get_counter(
+        name=name,
+        description=description,
+        parameters=parameters
+    )
+
+  def _get_scalar(
+      self,
+      name: str,
+      description: str,
+      parameters: dict[str, type[str]] | None = None
+  ) -> pg.monitoring.Metric:
+    return self._metric_collection.get_scalar(
+        name=name,
+        description=description,
+        parameters=parameters
+    )
+
+  def _initialize_metrics(self) -> None:
+    """Initializes metrics."""
+
+    self._metric_collection = pg.monitoring.metric_collection(_METRIC_NAMESPACE)
+
+    #
+    # Environment metrics.
+    #
+
+    self._environment_housekeep_duration_ms = self._get_scalar(
+        'environment_housekeep_duration_ms',
+        description='Environment housekeeping duration in milliseconds',
+        parameters={
+            'environment_id': str,
+            'has_error': bool,
+        }
+    )
+
+    #
+    # Sandbox metrics.
+    #
+
+    # Sandbox counters.
+    self._sandbox_start = self._get_counter(
+        'sandbox_start',
+        description='Sandbox start counter',
+        parameters={
+            'environment_id': str,
+            'has_error': bool,
+        }
+    )
+    self._sandbox_shutdown = self._get_counter(
+        'sandbox_shutdown',
+        description='Sandbox shutdown counter',
+        parameters={
+            'environment_id': str,
+            'has_error': bool
+        }
+    )
+    self._sandbox_housekeep = self._get_counter(
+        'sandbox_housekeep',
+        description='Sandbox housekeeping counter',
+        parameters={
+            'environment_id': str,
+            'has_error': bool,
+        }
+    )
+    self._sandbox_activity = self._get_counter(
+        'sandbox_activity',
+        description='Sandbox activity counter',
+        parameters={
+            'environment_id': str,
+            'activity': str,
+            'has_error': bool,
+        }
+    )
+
+    # Sandbox scalars.
+    self._sandbox_lifetime_ms = self._get_scalar(
+        'sandbox_lifetime_ms',
+        description='Sandbox life time in milliseconds',
+        parameters={
+            'environment_id': str,
+            'has_error': bool,
+        }
+    )
+    self._sandbox_start_duration_ms = self._get_scalar(
+        'sandbox_start_duration_ms',
+        description='Sandbox start duration in milliseconds',
+        parameters={
+            'environment_id': str,
+            'has_error': bool,
+        }
+    )
+    self._sandbox_shutdown_duration_ms = self._get_scalar(
+        'sandbox_shutdown_duration_ms',
+        description='Sandbox shutdown duration in milliseconds',
+        parameters={
+            'environment_id': str,
+            'has_error': bool,
+        }
+    )
+    self._sandbox_housekeep_duration_ms = self._get_scalar(
+        'sandbox_housekeep_duration_ms',
+        description='Sandbox housekeeping duration in milliseconds',
+        parameters={
+            'environment_id': str,
+            'has_error': bool,
+        }
+    )
+    self._sandbox_status_duration_ms = self._get_scalar(
+        'sandbox_status_duration_ms',
+        description='Sandbox duration of specific status in milliseconds',
+        parameters={
+            'environment_id': str,
+            'status': str,
+        }
+    )
+    self._sandbox_activity_duration_ms = self._get_scalar(
+        'sandbox_activity_duration_ms',
+        description='Sandbox activity duration in milliseconds',
+        parameters={
+            'environment_id': str,
+            'activity': str,
+            'has_error': bool,
+        }
+    )
+
+    #
+    # Feature metrics.
+    #
+
+    # Feature counters.
+    self._feature_setup = self._get_counter(
+        'feature_setup',
+        description='Feature setup counter',
+        parameters={
+            'environment_id': str,
+            'feature_name': str,
+            'has_error': bool,
+        }
+    )
+    self._feature_teardown = self._get_counter(
+        'feature_teardown',
+        description='Feature teardown counter',
+        parameters={
+            'environment_id': str,
+            'feature_name': str,
+            'has_error': bool,
+        }
+    )
+    self._feature_setup_session = self._get_counter(
+        'feature_setup_session',
+        description='Feature setup session counter',
+        parameters={
+            'environment_id': str,
+            'feature_name': str,
+            'has_error': bool,
+        }
+    )
+    self._feature_teardown_session = self._get_counter(
+        'feature_teardown_session',
+        description='Feature teardown session counter',
+        parameters={
+            'environment_id': str,
+            'feature_name': str,
+            'has_error': bool,
+        }
+    )
+    self._feature_housekeep = self._get_counter(
+        'feature_housekeep',
+        description='Feature housekeeping counter',
+        parameters={
+            'environment_id': str,
+            'feature_name': str,
+            'has_error': bool,
+        }
+    )
+
+    # Feature scalars.
+    self._feature_setup_duration_ms = self._get_scalar(
+        'feature_setup_duration_ms',
+        description='Feature setup duration in milliseconds',
+        parameters={
+            'environment_id': str,
+            'feature_name': str,
+            'has_error': bool,
+        }
+    )
+    self._feature_teardown_duration_ms = self._get_scalar(
+        'feature_teardown_duration_ms',
+        description='Feature teardown duration in milliseconds',
+        parameters={
+            'environment_id': str,
+            'feature_name': str,
+            'has_error': bool,
+        }
+    )
+    self._feature_setup_session_duration_ms = self._get_scalar(
+        'feature_setup_session_duration_ms',
+        description='Feature setup session duration in milliseconds',
+        parameters={
+            'environment_id': str,
+            'feature_name': str,
+            'has_error': bool,
+        }
+    )
+    self._feature_teardown_session_duration_ms = self._get_scalar(
+        'feature_teardown_session_duration_ms',
+        description='Feature teardown session duration in milliseconds',
+        parameters={
+            'environment_id': str,
+            'feature_name': str,
+            'has_error': bool,
+        }
+    )
+    self._feature_housekeep_duration_ms = self._get_scalar(
+        'feature_housekeep_duration_ms',
+        description='Feature housekeeping duration in milliseconds',
+        parameters={
+            'environment_id': str,
+            'feature_name': str,
+            'has_error': bool,
+        }
+    )
+
+    #
+    # Session metrics.
+    #
+
+    self._session_start_duration_ms = self._get_scalar(
+        'session_start_duration_ms',
+        description='Session start duration in milliseconds',
+        parameters={
+            'environment_id': str,
+            'has_error': bool,
+        }
+    )
+    self._session_end_duration_ms = self._get_scalar(
+        'session_end_duration_ms',
+        description='Session end duration in milliseconds',
+        parameters={
+            'environment_id': str,
+            'has_error': bool,
+        }
+    )
+    self._session_lifetime_ms = self._get_scalar(
+        'session_lifetime_ms',
+        description='Session lifetime in milliseconds',
+        parameters={
+            'environment_id': str,
+            'has_error': bool,
+        }
+    )
+
+  def on_environment_starting(
+      self,
+      environment: base.Environment
+  ) -> None:
+    """Called when the environment is starting."""
+    self._initialize_metrics()
+
+  def on_environment_housekeep(
+      self,
+      environment: base.Environment,
+      counter: int,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when the environment is housekeeping."""
+    self._environment_housekeep_duration_ms.record(
+        int(duration * 1000),
+        environment_id=str(environment.id),
+        has_error=error is not None
+    )
+
+  def on_sandbox_start(
+      self,
+      environment: base.Environment,
+      sandbox: base.Sandbox,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    self._sandbox_start.increment(
+        environment_id=str(environment.id),
+        has_error=error is not None
+    )
+    self._sandbox_start_duration_ms.record(
+        int(duration * 1000),
+        environment_id=str(environment.id),
+        has_error=error is not None
+    )
+
+  def on_sandbox_status_change(
+      self,
+      environment: base.Environment,
+      sandbox: base.Sandbox,
+      old_status: base.Sandbox.Status,
+      new_status: base.Sandbox.Status,
+      span: float
+  ) -> None:
+    self._sandbox_status_duration_ms.record(
+        int(span * 1000),
+        environment_id=str(environment.id),
+        status=old_status.value
+    )
+
+  def on_sandbox_shutdown(
+      self,
+      environment: base.Environment,
+      sandbox: base.Sandbox,
+      duration: float,
+      lifetime: float,
+      error: BaseException | None
+  ) -> None:
+    self._sandbox_shutdown.increment(
+        environment_id=str(environment.id),
+        has_error=error is not None
+    )
+    self._sandbox_shutdown_duration_ms.record(
+        int(duration * 1000),
+        environment_id=str(environment.id),
+        has_error=error is not None
+    )
+    self._sandbox_lifetime_ms.record(
+        int(lifetime * 1000),
+        environment_id=str(environment.id),
+        has_error=error is not None
+    )
+
+  def on_sandbox_housekeep(
+      self,
+      environment: base.Environment,
+      sandbox: base.Sandbox,
+      counter: int,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox feature is housekeeping."""
+    self._sandbox_housekeep.increment(
+        environment_id=str(environment.id),
+        has_error=error is not None
+    )
+    self._sandbox_housekeep_duration_ms.record(
+        int(duration * 1000),
+        environment_id=str(environment.id),
+        has_error=error is not None
+    )
+
+  def on_feature_setup(
+      self,
+      environment: base.Environment,
+      sandbox: base.Sandbox,
+      feature: base.Feature,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox feature is setup."""
+    self._feature_setup.increment(
+        environment_id=str(environment.id),
+        feature_name=feature.name,
+        has_error=error is not None
+    )
+    self._feature_setup_duration_ms.record(
+        int(duration * 1000),
+        environment_id=str(environment.id),
+        feature_name=feature.name,
+        has_error=error is not None
+    )
+
+  def on_feature_teardown(
+      self,
+      environment: base.Environment,
+      sandbox: base.Sandbox,
+      feature: base.Feature,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox feature is teardown."""
+    self._feature_teardown.increment(
+        environment_id=str(environment.id),
+        feature_name=feature.name,
+        has_error=error is not None
+    )
+    self._feature_teardown_duration_ms.record(
+        int(duration * 1000),
+        environment_id=str(environment.id),
+        feature_name=feature.name,
+        has_error=error is not None
+    )
+
+  def on_feature_setup_session(
+      self,
+      environment: base.Environment,
+      sandbox: base.Sandbox,
+      feature: base.Feature,
+      session_id: str | None,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox feature is setup."""
+    self._feature_setup_session.increment(
+        environment_id=str(environment.id),
+        feature_name=feature.name,
+        has_error=error is not None
+    )
+    self._feature_setup_session_duration_ms.record(
+        int(duration * 1000),
+        environment_id=str(environment.id),
+        feature_name=feature.name,
+        has_error=error is not None
+    )
+
+  def on_feature_teardown_session(
+      self,
+      environment: base.Environment,
+      sandbox: base.Sandbox,
+      feature: base.Feature,
+      session_id: str,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox feature is teardown."""
+    self._feature_teardown_session.increment(
+        environment_id=str(environment.id),
+        feature_name=feature.name,
+        has_error=error is not None
+    )
+    self._feature_teardown_session_duration_ms.record(
+        int(duration * 1000),
+        environment_id=str(environment.id),
+        feature_name=feature.name,
+        has_error=error is not None
+    )
+
+  def on_feature_housekeep(
+      self,
+      environment: base.Environment,
+      sandbox: base.Sandbox,
+      feature: base.Feature,
+      counter: int,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox feature is housekeeping."""
+    self._feature_housekeep.increment(
+        environment_id=str(environment.id),
+        feature_name=feature.name,
+        has_error=error is not None
+    )
+    self._feature_housekeep_duration_ms.record(
+        int(duration * 1000),
+        environment_id=str(environment.id),
+        feature_name=feature.name,
+        has_error=error is not None
+    )
+
+  def on_session_start(
+      self,
+      environment: base.Environment,
+      sandbox: base.Sandbox,
+      session_id: str,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox session starts."""
+    self._session_start_duration_ms.record(
+        int(duration * 1000),
+        environment_id=str(environment.id),
+        has_error=error is not None
+    )
+
+  def on_session_end(
+      self,
+      environment: base.Environment,
+      sandbox: base.Sandbox,
+      session_id: str,
+      duration: float,
+      lifetime: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox session ends."""
+    self._session_end_duration_ms.record(
+        int(duration * 1000),
+        environment_id=str(environment.id),
+        has_error=error is not None
+    )
+    self._session_lifetime_ms.record(
+        int(lifetime * 1000),
+        environment_id=str(environment.id),
+        has_error=error is not None
+    )
+
+  def on_sandbox_activity(
+      self,
+      name: str,
+      environment: base.Environment,
+      sandbox: base.Sandbox,
+      feature: base.Feature | None,
+      session_id: str | None,
+      duration: float,
+      error: BaseException | None,
+      **kwargs
+  ) -> None:
+    """Called when a sandbox activity is performed."""
+    self._sandbox_activity.increment(
+        environment_id=str(environment.id),
+        activity=name,
+        has_error=error is not None
+    )
+    self._sandbox_activity_duration_ms.record(
+        int(duration * 1000),
+        environment_id=str(environment.id),
+        activity=name,
+        has_error=error is not None
+    )

langfun/env/event_handlers/metric_writer_test.py

@@ -0,0 +1,160 @@
+# 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
+from langfun.env import test_utils
+from langfun.env.event_handlers import metric_writer as metric_writer_lib
+
+
+class MetricWriterTest(unittest.TestCase):
+
+  def test_write_metric(self):
+    writer = metric_writer_lib.MetricWriter()
+    env = test_utils.TestingEnvironment(
+        features={
+            'test_feature1': test_utils.TestingFeature(housekeep_interval=0),
+            'test_feature2': test_utils.TestingFeature(housekeep_interval=None),
+        },
+        pool_size=2,
+        outage_grace_period=0,
+        outage_retry_interval=0,
+        housekeep_interval=0.0,
+        sandbox_keepalive_interval=1.0,
+        event_handlers=[writer],
+    )
+    with env:
+      with env.sandbox('session1') as sb:
+        self.assertEqual(sb.test_feature1.num_shell_calls(), 4)
+
+      with self.assertRaises(interface.SandboxStateError):
+        with env.sandbox('session2') as sb:
+          sb.shell('echo "bar"', raise_error=RuntimeError)
+
+    self.assertEqual(
+        writer._sandbox_start.value(
+            environment_id='testing-env', has_error=False
+        ),
+        2
+    )
+    self.assertGreater(
+        writer._sandbox_housekeep.value(
+            environment_id='testing-env',
+            has_error=False
+        ),
+        0,
+    )
+    self.assertEqual(
+        writer._sandbox_shutdown.value(
+            environment_id='testing-env', has_error=False
+        ),
+        2
+    )
+    self.assertEqual(
+        writer._feature_setup.value(
+            environment_id='testing-env',
+            feature_name='test_feature1',
+            has_error=False
+        ),
+        2,
+    )
+    self.assertEqual(
+        writer._feature_setup.value(
+            environment_id='testing-env',
+            feature_name='test_feature2',
+            has_error=False
+        ),
+        2,
+    )
+    self.assertEqual(
+        writer._feature_setup_session.value(
+            environment_id='testing-env',
+            feature_name='test_feature1',
+            has_error=False
+        ),
+        3,
+    )
+    self.assertEqual(
+        writer._feature_setup_session.value(
+            environment_id='testing-env',
+            feature_name='test_feature2',
+            has_error=False
+        ),
+        3,
+    )
+    self.assertEqual(
+        writer._feature_teardown_session.value(
+            environment_id='testing-env',
+            feature_name='test_feature1',
+            has_error=False
+        ),
+        2,
+    )
+    self.assertEqual(
+        writer._feature_teardown_session.value(
+            environment_id='testing-env',
+            feature_name='test_feature2',
+            has_error=False
+        ),
+        2,
+    )
+    self.assertEqual(
+        writer._feature_teardown.value(
+            environment_id='testing-env',
+            feature_name='test_feature1',
+            has_error=False
+        ),
+        2,
+    )
+    self.assertEqual(
+        writer._feature_teardown.value(
+            environment_id='testing-env',
+            feature_name='test_feature2',
+            has_error=False
+        ),
+        2,
+    )
+    self.assertGreater(
+        writer._feature_housekeep.value(
+            environment_id='testing-env',
+            feature_name='test_feature1',
+            has_error=False
+        ),
+        0,
+    )
+    self.assertEqual(
+        writer._feature_housekeep.value(
+            environment_id='testing-env',
+            feature_name='test_feature2',
+            has_error=False
+        ),
+        0,
+    )
+    self.assertEqual(
+        writer._sandbox_activity.value(
+            environment_id='testing-env', activity='shell', has_error=False
+        ),
+        18
+    )
+    self.assertEqual(
+        writer._sandbox_activity.value(
+            environment_id='testing-env', activity='shell', has_error=True
+        ),
+        1
+    )
+
+
+if __name__ == '__main__':
+  unittest.main()