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

langfun/env/event_handlers/metric_writer.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Environment event logger."""
 
+from typing import Annotated
 from langfun.env.event_handlers import base
 import pyglove as pg
 
@@ -23,6 +24,11 @@ _METRIC_NAMESPACE = '/langfun/env'
 class MetricWriter(pg.Object, base.EventHandler):
   """Event handler for streamz metrics."""
 
+  app: Annotated[
+      str,
+      'Application name that will be used as a metric parameter.'
+  ] = ''
+
   def _get_counter(
       self,
       name: str,
@@ -47,6 +53,11 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters=parameters
     )
 
+  def _error_tag(self, error: BaseException | None) -> str:
+    if error is None:
+      return ''
+    return pg.utils.ErrorInfo.from_exception(error).tag
+
   def _initialize_metrics(self) -> None:
     """Initializes metrics."""
 
@@ -60,8 +71,9 @@ class MetricWriter(pg.Object, base.EventHandler):
         'environment_housekeep_duration_ms',
         description='Environment housekeeping duration in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
-            'has_error': bool,
+            'error': str,
         }
     )
 
@@ -74,33 +86,37 @@ class MetricWriter(pg.Object, base.EventHandler):
         'sandbox_start',
         description='Sandbox start counter',
         parameters={
+            'app': str,
             'environment_id': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._sandbox_shutdown = self._get_counter(
         'sandbox_shutdown',
         description='Sandbox shutdown counter',
         parameters={
+            'app': str,
             'environment_id': str,
-            'has_error': bool
+            'error': str
         }
     )
     self._sandbox_housekeep = self._get_counter(
         'sandbox_housekeep',
         description='Sandbox housekeeping counter',
         parameters={
+            'app': str,
             'environment_id': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._sandbox_activity = self._get_counter(
         'sandbox_activity',
         description='Sandbox activity counter',
         parameters={
+            'app': str,
             'environment_id': str,
             'activity': str,
-            'has_error': bool,
+            'error': str,
         }
     )
 
@@ -109,38 +125,43 @@ class MetricWriter(pg.Object, base.EventHandler):
         'sandbox_lifetime_ms',
         description='Sandbox life time in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._sandbox_start_duration_ms = self._get_scalar(
         'sandbox_start_duration_ms',
         description='Sandbox start duration in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._sandbox_shutdown_duration_ms = self._get_scalar(
         'sandbox_shutdown_duration_ms',
         description='Sandbox shutdown duration in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._sandbox_housekeep_duration_ms = self._get_scalar(
         'sandbox_housekeep_duration_ms',
         description='Sandbox housekeeping duration in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._sandbox_status_duration_ms = self._get_scalar(
         'sandbox_status_duration_ms',
         description='Sandbox duration of specific status in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
             'status': str,
         }
@@ -149,9 +170,10 @@ class MetricWriter(pg.Object, base.EventHandler):
         'sandbox_activity_duration_ms',
         description='Sandbox activity duration in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
             'activity': str,
-            'has_error': bool,
+            'error': str,
         }
     )
 
@@ -164,45 +186,50 @@ class MetricWriter(pg.Object, base.EventHandler):
         'feature_setup',
         description='Feature setup counter',
         parameters={
+            'app': str,
             'environment_id': str,
             'feature_name': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._feature_teardown = self._get_counter(
         'feature_teardown',
         description='Feature teardown counter',
         parameters={
+            'app': str,
             'environment_id': str,
             'feature_name': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._feature_setup_session = self._get_counter(
         'feature_setup_session',
         description='Feature setup session counter',
         parameters={
+            'app': str,
             'environment_id': str,
             'feature_name': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._feature_teardown_session = self._get_counter(
         'feature_teardown_session',
         description='Feature teardown session counter',
         parameters={
+            'app': str,
             'environment_id': str,
             'feature_name': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._feature_housekeep = self._get_counter(
         'feature_housekeep',
         description='Feature housekeeping counter',
         parameters={
+            'app': str,
             'environment_id': str,
             'feature_name': str,
-            'has_error': bool,
+            'error': str,
         }
     )
 
@@ -211,45 +238,50 @@ class MetricWriter(pg.Object, base.EventHandler):
         'feature_setup_duration_ms',
         description='Feature setup duration in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
             'feature_name': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._feature_teardown_duration_ms = self._get_scalar(
         'feature_teardown_duration_ms',
         description='Feature teardown duration in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
             'feature_name': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._feature_setup_session_duration_ms = self._get_scalar(
         'feature_setup_session_duration_ms',
         description='Feature setup session duration in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
             'feature_name': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._feature_teardown_session_duration_ms = self._get_scalar(
         'feature_teardown_session_duration_ms',
         description='Feature teardown session duration in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
             'feature_name': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._feature_housekeep_duration_ms = self._get_scalar(
         'feature_housekeep_duration_ms',
         description='Feature housekeeping duration in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
             'feature_name': str,
-            'has_error': bool,
+            'error': str,
         }
     )
 
@@ -261,24 +293,27 @@ class MetricWriter(pg.Object, base.EventHandler):
         'session_start_duration_ms',
         description='Session start duration in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._session_end_duration_ms = self._get_scalar(
         'session_end_duration_ms',
         description='Session end duration in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
-            'has_error': bool,
+            'error': str,
         }
     )
     self._session_lifetime_ms = self._get_scalar(
         'session_lifetime_ms',
         description='Session lifetime in milliseconds',
         parameters={
+            'app': str,
             'environment_id': str,
-            'has_error': bool,
+            'error': str,
         }
     )
 
@@ -294,13 +329,15 @@ class MetricWriter(pg.Object, base.EventHandler):
       environment: base.Environment,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs
   ) -> None:
     """Called when the environment is housekeeping."""
     self._environment_housekeep_duration_ms.record(
         int(duration * 1000),
+        app=self.app,
         environment_id=str(environment.id),
-        has_error=error is not None
+        error=self._error_tag(error)
     )
 
   def on_sandbox_start(
@@ -311,13 +348,15 @@ class MetricWriter(pg.Object, base.EventHandler):
       error: BaseException | None
   ) -> None:
     self._sandbox_start.increment(
+        app=self.app,
         environment_id=str(environment.id),
-        has_error=error is not None
+        error=self._error_tag(error)
     )
     self._sandbox_start_duration_ms.record(
         int(duration * 1000),
+        app=self.app,
         environment_id=str(environment.id),
-        has_error=error is not None
+        error=self._error_tag(error)
     )
 
   def on_sandbox_status_change(
@@ -330,6 +369,7 @@ class MetricWriter(pg.Object, base.EventHandler):
   ) -> None:
     self._sandbox_status_duration_ms.record(
         int(span * 1000),
+        app=self.app,
         environment_id=str(environment.id),
         status=old_status.value
     )
@@ -343,18 +383,21 @@ class MetricWriter(pg.Object, base.EventHandler):
       error: BaseException | None
   ) -> None:
     self._sandbox_shutdown.increment(
+        app=self.app,
         environment_id=str(environment.id),
-        has_error=error is not None
+        error=self._error_tag(error)
     )
     self._sandbox_shutdown_duration_ms.record(
         int(duration * 1000),
+        app=self.app,
         environment_id=str(environment.id),
-        has_error=error is not None
+        error=self._error_tag(error)
     )
     self._sandbox_lifetime_ms.record(
         int(lifetime * 1000),
+        app=self.app,
         environment_id=str(environment.id),
-        has_error=error is not None
+        error=self._error_tag(error)
     )
 
   def on_sandbox_housekeep(
@@ -363,17 +406,20 @@ class MetricWriter(pg.Object, base.EventHandler):
       sandbox: base.Sandbox,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs
   ) -> None:
     """Called when a sandbox feature is housekeeping."""
     self._sandbox_housekeep.increment(
+        app=self.app,
         environment_id=str(environment.id),
-        has_error=error is not None
+        error=self._error_tag(error)
     )
     self._sandbox_housekeep_duration_ms.record(
         int(duration * 1000),
+        app=self.app,
         environment_id=str(environment.id),
-        has_error=error is not None
+        error=self._error_tag(error)
     )
 
   def on_feature_setup(
@@ -386,15 +432,17 @@ class MetricWriter(pg.Object, base.EventHandler):
   ) -> None:
     """Called when a sandbox feature is setup."""
     self._feature_setup.increment(
+        app=self.app,
         environment_id=str(environment.id),
         feature_name=feature.name,
-        has_error=error is not None
+        error=self._error_tag(error)
     )
     self._feature_setup_duration_ms.record(
         int(duration * 1000),
+        app=self.app,
         environment_id=str(environment.id),
         feature_name=feature.name,
-        has_error=error is not None
+        error=self._error_tag(error)
     )
 
   def on_feature_teardown(
@@ -407,15 +455,17 @@ class MetricWriter(pg.Object, base.EventHandler):
   ) -> None:
     """Called when a sandbox feature is teardown."""
     self._feature_teardown.increment(
+        app=self.app,
         environment_id=str(environment.id),
         feature_name=feature.name,
-        has_error=error is not None
+        error=self._error_tag(error)
     )
     self._feature_teardown_duration_ms.record(
         int(duration * 1000),
+        app=self.app,
         environment_id=str(environment.id),
         feature_name=feature.name,
-        has_error=error is not None
+        error=self._error_tag(error)
     )
 
   def on_feature_setup_session(
@@ -429,15 +479,17 @@ class MetricWriter(pg.Object, base.EventHandler):
   ) -> None:
     """Called when a sandbox feature is setup."""
     self._feature_setup_session.increment(
+        app=self.app,
         environment_id=str(environment.id),
         feature_name=feature.name,
-        has_error=error is not None
+        error=self._error_tag(error)
     )
     self._feature_setup_session_duration_ms.record(
         int(duration * 1000),
+        app=self.app,
         environment_id=str(environment.id),
         feature_name=feature.name,
-        has_error=error is not None
+        error=self._error_tag(error)
     )
 
   def on_feature_teardown_session(
@@ -451,15 +503,17 @@ class MetricWriter(pg.Object, base.EventHandler):
   ) -> None:
     """Called when a sandbox feature is teardown."""
     self._feature_teardown_session.increment(
+        app=self.app,
         environment_id=str(environment.id),
         feature_name=feature.name,
-        has_error=error is not None
+        error=self._error_tag(error)
     )
     self._feature_teardown_session_duration_ms.record(
         int(duration * 1000),
+        app=self.app,
         environment_id=str(environment.id),
         feature_name=feature.name,
-        has_error=error is not None
+        error=self._error_tag(error)
     )
 
   def on_feature_housekeep(
@@ -469,19 +523,22 @@ class MetricWriter(pg.Object, base.EventHandler):
       feature: base.Feature,
       counter: int,
       duration: float,
-      error: BaseException | None
+      error: BaseException | None,
+      **kwargs
   ) -> None:
     """Called when a sandbox feature is housekeeping."""
     self._feature_housekeep.increment(
+        app=self.app,
         environment_id=str(environment.id),
         feature_name=feature.name,
-        has_error=error is not None
+        error=self._error_tag(error)
     )
     self._feature_housekeep_duration_ms.record(
         int(duration * 1000),
+        app=self.app,
         environment_id=str(environment.id),
         feature_name=feature.name,
-        has_error=error is not None
+        error=self._error_tag(error)
     )
 
   def on_session_start(
@@ -495,8 +552,9 @@ class MetricWriter(pg.Object, base.EventHandler):
     """Called when a sandbox session starts."""
     self._session_start_duration_ms.record(
         int(duration * 1000),
+        app=self.app,
         environment_id=str(environment.id),
-        has_error=error is not None
+        error=self._error_tag(error)
     )
 
   def on_session_end(
@@ -511,13 +569,15 @@ class MetricWriter(pg.Object, base.EventHandler):
     """Called when a sandbox session ends."""
     self._session_end_duration_ms.record(
         int(duration * 1000),
+        app=self.app,
         environment_id=str(environment.id),
-        has_error=error is not None
+        error=self._error_tag(error)
     )
     self._session_lifetime_ms.record(
         int(lifetime * 1000),
+        app=self.app,
         environment_id=str(environment.id),
-        has_error=error is not None
+        error=self._error_tag(error)
     )
 
   def on_sandbox_activity(
@@ -533,13 +593,15 @@ class MetricWriter(pg.Object, base.EventHandler):
   ) -> None:
     """Called when a sandbox activity is performed."""
     self._sandbox_activity.increment(
+        app=self.app,
         environment_id=str(environment.id),
         activity=name,
-        has_error=error is not None
+        error=self._error_tag(error)
     )
     self._sandbox_activity_duration_ms.record(
         int(duration * 1000),
+        app=self.app,
         environment_id=str(environment.id),
         activity=name,
-        has_error=error is not None
+        error=self._error_tag(error)
     )

langfun/env/event_handlers/metric_writer_test.py

@@ -22,7 +22,7 @@ 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()
+    writer = metric_writer_lib.MetricWriter(app='test_app')
     env = test_utils.TestingEnvironment(
         features={
             'test_feature1': test_utils.TestingFeature(housekeep_interval=0),
@@ -45,112 +45,122 @@ class MetricWriterTest(unittest.TestCase):
 
     self.assertEqual(
         writer._sandbox_start.value(
-            environment_id='testing-env', has_error=False
+            app='test_app',
+            environment_id='testing-env',
+            error=''
         ),
         2
     )
     self.assertGreater(
         writer._sandbox_housekeep.value(
+            app='test_app',
             environment_id='testing-env',
-            has_error=False
+            error=''
         ),
         0,
     )
     self.assertEqual(
         writer._sandbox_shutdown.value(
-            environment_id='testing-env', has_error=False
+            app='test_app',
+            environment_id='testing-env',
+            error=''
         ),
         2
     )
     self.assertEqual(
         writer._feature_setup.value(
+            app='test_app',
             environment_id='testing-env',
-            feature_name='test_feature1',
-            has_error=False
+            feature_name='test_feature1', error=''
         ),
         2,
     )
     self.assertEqual(
         writer._feature_setup.value(
+            app='test_app',
             environment_id='testing-env',
-            feature_name='test_feature2',
-            has_error=False
+            feature_name='test_feature2', error=''
         ),
         2,
     )
     self.assertEqual(
         writer._feature_setup_session.value(
+            app='test_app',
             environment_id='testing-env',
-            feature_name='test_feature1',
-            has_error=False
+            feature_name='test_feature1', error=''
         ),
         3,
     )
     self.assertEqual(
         writer._feature_setup_session.value(
+            app='test_app',
             environment_id='testing-env',
-            feature_name='test_feature2',
-            has_error=False
+            feature_name='test_feature2', error=''
         ),
         3,
     )
     self.assertEqual(
         writer._feature_teardown_session.value(
+            app='test_app',
             environment_id='testing-env',
-            feature_name='test_feature1',
-            has_error=False
+            feature_name='test_feature1', error=''
         ),
         2,
     )
     self.assertEqual(
         writer._feature_teardown_session.value(
+            app='test_app',
             environment_id='testing-env',
-            feature_name='test_feature2',
-            has_error=False
+            feature_name='test_feature2', error=''
         ),
         2,
     )
     self.assertEqual(
         writer._feature_teardown.value(
+            app='test_app',
             environment_id='testing-env',
-            feature_name='test_feature1',
-            has_error=False
+            feature_name='test_feature1', error=''
         ),
         2,
     )
     self.assertEqual(
         writer._feature_teardown.value(
+            app='test_app',
             environment_id='testing-env',
-            feature_name='test_feature2',
-            has_error=False
+            feature_name='test_feature2', error=''
         ),
         2,
     )
     self.assertGreater(
         writer._feature_housekeep.value(
+            app='test_app',
             environment_id='testing-env',
-            feature_name='test_feature1',
-            has_error=False
+            feature_name='test_feature1', error=''
         ),
         0,
     )
     self.assertEqual(
         writer._feature_housekeep.value(
+            app='test_app',
             environment_id='testing-env',
-            feature_name='test_feature2',
-            has_error=False
+            feature_name='test_feature2', error=''
         ),
         0,
     )
     self.assertEqual(
         writer._sandbox_activity.value(
-            environment_id='testing-env', activity='shell', has_error=False
+            app='test_app',
+            environment_id='testing-env',
+            activity='shell', error=''
         ),
         18
     )
     self.assertEqual(
         writer._sandbox_activity.value(
-            environment_id='testing-env', activity='shell', has_error=True
+            app='test_app',
+            environment_id='testing-env',
+            activity='shell',
+            error='RuntimeError'
         ),
         1
     )