langfun 0.1.2.dev202510200805__py3-none-any.whl → 0.1.2.dev202511160804__py3-none-any.whl

langfun/env/event_handlers/metric_writer.py

@@ -14,14 +14,14 @@
 """Environment event logger."""
 
 from typing import Annotated
-from langfun.env.event_handlers import base
+from langfun.env import interface
 import pyglove as pg
 
 
 _METRIC_NAMESPACE = '/langfun/env'
 
 
-class MetricWriter(pg.Object, base.EventHandler):
+class MetricWriter(pg.Object, interface.EventHandler):
   """Event handler for streamz metrics."""
 
   app: Annotated[
@@ -100,6 +100,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'error': str,
         }
     )
@@ -109,6 +110,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'error': str
         }
     )
@@ -118,6 +120,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'status': str,
         },
     )
@@ -127,6 +130,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'error': str,
         }
     )
@@ -136,6 +140,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'activity': str,
             'error': str,
         }
@@ -148,6 +153,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'error': str,
         }
     )
@@ -157,6 +163,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'error': str,
         }
     )
@@ -166,6 +173,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'error': str,
         }
     )
@@ -175,6 +183,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'error': str,
         }
     )
@@ -184,6 +193,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'status': str,
         }
     )
@@ -193,11 +203,47 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'activity': str,
             'error': str,
         }
     )
 
+    #
+    # Sandbox session metrics.
+    #
+
+    self._sandbox_session_start_duration_ms = self._get_distribution(
+        'sandbox_session_start_duration_ms',
+        description='Sandbox session start duration in milliseconds',
+        parameters={
+            'app': str,
+            'environment_id': str,
+            'image_id': str,
+            'error': str,
+        }
+    )
+    self._sandbox_session_end_duration_ms = self._get_distribution(
+        'sandbox_session_end_duration_ms',
+        description='Sandbox session end duration in milliseconds',
+        parameters={
+            'app': str,
+            'environment_id': str,
+            'image_id': str,
+            'error': str,
+        }
+    )
+    self._sandbox_session_lifetime_ms = self._get_distribution(
+        'sandbox_session_lifetime_ms',
+        description='Sandbox session lifetime in milliseconds',
+        parameters={
+            'app': str,
+            'environment_id': str,
+            'image_id': str,
+            'error': str,
+        }
+    )
+
     #
     # Feature metrics.
     #
@@ -209,6 +255,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'feature_name': str,
             'error': str,
         }
@@ -219,6 +266,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'feature_name': str,
             'error': str,
         }
@@ -229,6 +277,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'feature_name': str,
             'error': str,
         }
@@ -239,16 +288,29 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'feature_name': str,
             'error': str,
         }
     )
+    self._feature_activity = self._get_counter(
+        'feature_activity',
+        description='Feature activity counter',
+        parameters={
+            'app': str,
+            'environment_id': str,
+            'image_id': str,
+            'activity': str,
+            'error': str,
+        }
+    )
     self._feature_housekeep = self._get_counter(
         'feature_housekeep',
         description='Feature housekeeping counter',
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'feature_name': str,
             'error': str,
         }
@@ -261,6 +323,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'feature_name': str,
             'error': str,
         }
@@ -271,6 +334,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'feature_name': str,
             'error': str,
         }
@@ -281,6 +345,7 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'feature_name': str,
             'error': str,
         }
@@ -291,63 +356,44 @@ class MetricWriter(pg.Object, base.EventHandler):
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
             'feature_name': str,
             'error': str,
         }
     )
-    self._feature_housekeep_duration_ms = self._get_distribution(
-        'feature_housekeep_duration_ms',
-        description='Feature housekeeping duration in milliseconds',
-        parameters={
-            'app': str,
-            'environment_id': str,
-            'feature_name': str,
-            'error': str,
-        }
-    )
-
-    #
-    # Session metrics.
-    #
-
-    self._session_start_duration_ms = self._get_distribution(
-        'session_start_duration_ms',
-        description='Session start duration in milliseconds',
-        parameters={
-            'app': str,
-            'environment_id': str,
-            'error': str,
-        }
-    )
-    self._session_end_duration_ms = self._get_distribution(
-        'session_end_duration_ms',
-        description='Session end duration in milliseconds',
+    self._feature_activity_duration_ms = self._get_distribution(
+        'feature_activity_duration_ms',
+        description='Feature activity duration in milliseconds',
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
+            'activity': str,
             'error': str,
         }
     )
-    self._session_lifetime_ms = self._get_distribution(
-        'session_lifetime_ms',
-        description='Session lifetime in milliseconds',
+    self._feature_housekeep_duration_ms = self._get_distribution(
+        'feature_housekeep_duration_ms',
+        description='Feature housekeeping duration in milliseconds',
         parameters={
             'app': str,
             'environment_id': str,
+            'image_id': str,
+            'feature_name': str,
             'error': str,
         }
     )
 
   def on_environment_starting(
       self,
-      environment: base.Environment
+      environment: interface.Environment
   ) -> None:
     """Called when the environment is starting."""
     self._initialize_metrics()
 
   def on_environment_housekeep(
       self,
-      environment: base.Environment,
+      environment: interface.Environment,
       counter: int,
       duration: float,
       error: BaseException | None,
@@ -363,81 +409,151 @@ class MetricWriter(pg.Object, base.EventHandler):
 
   def on_sandbox_start(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
+      sandbox: interface.Sandbox,
       duration: float,
       error: BaseException | None
   ) -> None:
     self._sandbox_start.increment(
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(sandbox.environment.id),
+        image_id=sandbox.image_id,
         error=self._error_tag(error)
     )
     self._sandbox_start_duration_ms.record(
         int(duration * 1000),
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(sandbox.environment.id),
+        image_id=sandbox.image_id,
         error=self._error_tag(error)
     )
 
   def on_sandbox_status_change(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      old_status: base.Sandbox.Status,
-      new_status: base.Sandbox.Status,
+      sandbox: interface.Sandbox,
+      old_status: interface.Sandbox.Status,
+      new_status: interface.Sandbox.Status,
       span: float
   ) -> None:
     self._sandbox_status_duration_ms.record(
         int(span * 1000),
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(sandbox.environment.id),
+        image_id=sandbox.image_id,
         status=old_status.value
     )
-    if old_status != base.Sandbox.Status.CREATED:
+    if old_status != interface.Sandbox.Status.CREATED:
       self._sandbox_count.increment(
           delta=-1,
           app=self.app,
-          environment_id=str(environment.id),
+          environment_id=str(sandbox.environment.id),
+          image_id=sandbox.image_id,
           status=old_status.value
       )
-    if new_status != base.Sandbox.Status.OFFLINE:
+    if new_status != interface.Sandbox.Status.OFFLINE:
       self._sandbox_count.increment(
           app=self.app,
-          environment_id=str(environment.id),
+          environment_id=str(sandbox.environment.id),
+          image_id=sandbox.image_id,
           status=new_status.value
       )
 
   def on_sandbox_shutdown(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
+      sandbox: interface.Sandbox,
       duration: float,
       lifetime: float,
       error: BaseException | None
   ) -> None:
     self._sandbox_shutdown.increment(
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(sandbox.environment.id),
+        image_id=sandbox.image_id,
         error=self._error_tag(error)
     )
     self._sandbox_shutdown_duration_ms.record(
         int(duration * 1000),
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(sandbox.environment.id),
+        image_id=sandbox.image_id,
         error=self._error_tag(error)
     )
     self._sandbox_lifetime_ms.record(
         int(lifetime * 1000),
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(sandbox.environment.id),
+        image_id=sandbox.image_id,
+        error=self._error_tag(error)
+    )
+
+  def on_sandbox_session_start(
+      self,
+      sandbox: interface.Sandbox,
+      session_id: str,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox session starts."""
+    self._sandbox_session_start_duration_ms.record(
+        int(duration * 1000),
+        app=self.app,
+        environment_id=str(sandbox.environment.id),
+        image_id=sandbox.image_id,
+        error=self._error_tag(error)
+    )
+
+  def on_sandbox_session_end(
+      self,
+      sandbox: interface.Sandbox,
+      session_id: str,
+      duration: float,
+      lifetime: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox session ends."""
+    self._sandbox_session_end_duration_ms.record(
+        int(duration * 1000),
+        app=self.app,
+        environment_id=str(sandbox.environment.id),
+        image_id=sandbox.image_id,
+        error=self._error_tag(error)
+    )
+    self._sandbox_session_lifetime_ms.record(
+        int(lifetime * 1000),
+        app=self.app,
+        environment_id=str(sandbox.environment.id),
+        image_id=sandbox.image_id,
+        error=self._error_tag(error)
+    )
+
+  def on_sandbox_activity(
+      self,
+      name: str,
+      sandbox: interface.Sandbox,
+      session_id: str | None,
+      duration: float,
+      error: BaseException | None,
+      **kwargs
+  ) -> None:
+    """Called when a sandbox activity is performed."""
+    self._sandbox_activity.increment(
+        app=self.app,
+        environment_id=str(sandbox.environment.id),
+        image_id=sandbox.image_id,
+        activity=name,
+        error=self._error_tag(error)
+    )
+    self._sandbox_activity_duration_ms.record(
+        int(duration * 1000),
+        app=self.app,
+        environment_id=str(sandbox.environment.id),
+        image_id=sandbox.image_id,
+        activity=name,
         error=self._error_tag(error)
     )
 
   def on_sandbox_housekeep(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
+      sandbox: interface.Sandbox,
       counter: int,
       duration: float,
       error: BaseException | None,
@@ -446,196 +562,165 @@ class MetricWriter(pg.Object, base.EventHandler):
     """Called when a sandbox feature is housekeeping."""
     self._sandbox_housekeep.increment(
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(sandbox.environment.id),
+        image_id=sandbox.image_id,
         error=self._error_tag(error)
     )
     self._sandbox_housekeep_duration_ms.record(
         int(duration * 1000),
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(sandbox.environment.id),
+        image_id=sandbox.image_id,
         error=self._error_tag(error)
     )
 
   def on_feature_setup(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      feature: base.Feature,
+      feature: interface.Feature,
       duration: float,
       error: BaseException | None
   ) -> None:
     """Called when a sandbox feature is setup."""
+    image_id = feature.sandbox.image_id if feature.sandbox else ''
     self._feature_setup.increment(
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(feature.environment.id),
+        image_id=image_id,
         feature_name=feature.name,
         error=self._error_tag(error)
     )
     self._feature_setup_duration_ms.record(
         int(duration * 1000),
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(feature.environment.id),
+        image_id=image_id,
         feature_name=feature.name,
         error=self._error_tag(error)
     )
 
   def on_feature_teardown(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      feature: base.Feature,
+      feature: interface.Feature,
       duration: float,
       error: BaseException | None
   ) -> None:
     """Called when a sandbox feature is teardown."""
+    image_id = feature.sandbox.image_id if feature.sandbox else ''
     self._feature_teardown.increment(
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(feature.environment.id),
+        image_id=image_id,
         feature_name=feature.name,
         error=self._error_tag(error)
     )
     self._feature_teardown_duration_ms.record(
         int(duration * 1000),
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(feature.environment.id),
+        image_id=image_id,
         feature_name=feature.name,
         error=self._error_tag(error)
     )
 
   def on_feature_setup_session(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      feature: base.Feature,
+      feature: interface.Feature,
       session_id: str | None,
       duration: float,
       error: BaseException | None
   ) -> None:
     """Called when a sandbox feature is setup."""
+    image_id = feature.sandbox.image_id if feature.sandbox else ''
     self._feature_setup_session.increment(
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(feature.environment.id),
+        image_id=image_id,
         feature_name=feature.name,
         error=self._error_tag(error)
     )
     self._feature_setup_session_duration_ms.record(
         int(duration * 1000),
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(feature.environment.id),
+        image_id=image_id,
         feature_name=feature.name,
         error=self._error_tag(error)
     )
 
   def on_feature_teardown_session(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      feature: base.Feature,
+      feature: interface.Feature,
       session_id: str,
       duration: float,
       error: BaseException | None
   ) -> None:
     """Called when a sandbox feature is teardown."""
+    image_id = feature.sandbox.image_id if feature.sandbox else ''
     self._feature_teardown_session.increment(
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(feature.environment.id),
+        image_id=image_id,
         feature_name=feature.name,
         error=self._error_tag(error)
     )
     self._feature_teardown_session_duration_ms.record(
         int(duration * 1000),
         app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(feature.environment.id),
+        image_id=image_id,
         feature_name=feature.name,
         error=self._error_tag(error)
     )
 
-  def on_feature_housekeep(
+  def on_feature_activity(
       self,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      feature: base.Feature,
-      counter: int,
+      name: str,
+      feature: interface.Feature,
+      session_id: str | None,
       duration: float,
       error: BaseException | None,
       **kwargs
   ) -> None:
-    """Called when a sandbox feature is housekeeping."""
-    self._feature_housekeep.increment(
+    """Called when a feature activity is performed."""
+    image_id = feature.sandbox.image_id if feature.sandbox else ''
+    self._feature_activity.increment(
         app=self.app,
-        environment_id=str(environment.id),
-        feature_name=feature.name,
-        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,
-        error=self._error_tag(error)
-    )
-
-  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),
-        app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(feature.environment.id),
+        image_id=image_id,
+        activity=name,
         error=self._error_tag(error)
     )
-
-  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(
+    self._feature_activity_duration_ms.record(
         int(duration * 1000),
         app=self.app,
-        environment_id=str(environment.id),
-        error=self._error_tag(error)
-    )
-    self._session_lifetime_ms.record(
-        int(lifetime * 1000),
-        app=self.app,
-        environment_id=str(environment.id),
+        environment_id=str(feature.environment.id),
+        image_id=image_id,
+        activity=name,
         error=self._error_tag(error)
     )
 
-  def on_sandbox_activity(
+  def on_feature_housekeep(
       self,
-      name: str,
-      environment: base.Environment,
-      sandbox: base.Sandbox,
-      feature: base.Feature | None,
-      session_id: str | None,
+      feature: interface.Feature,
+      counter: int,
       duration: float,
       error: BaseException | None,
       **kwargs
   ) -> None:
-    """Called when a sandbox activity is performed."""
-    self._sandbox_activity.increment(
+    """Called when a sandbox feature is housekeeping."""
+    image_id = feature.sandbox.image_id if feature.sandbox else ''
+    self._feature_housekeep.increment(
         app=self.app,
-        environment_id=str(environment.id),
-        activity=name,
+        environment_id=str(feature.environment.id),
+        image_id=image_id,
+        feature_name=feature.name,
         error=self._error_tag(error)
     )
-    self._sandbox_activity_duration_ms.record(
+    self._feature_housekeep_duration_ms.record(
         int(duration * 1000),
         app=self.app,
-        environment_id=str(environment.id),
-        activity=name,
+        environment_id=str(feature.environment.id),
+        image_id=image_id,
+        feature_name=feature.name,
         error=self._error_tag(error)
     )