feat: implement dynamic string metrics

Author: d-perl

bec_lib/bec_lib/messages.py

@@ -9,7 +9,7 @@
 from enum import Enum, auto
 from importlib.metadata import PackageNotFoundError
 from importlib.metadata import version as importlib_version
-from typing import Annotated, Any, ClassVar, Literal, Self, Union
+from typing import Annotated, Any, ClassVar, Literal, Mapping, Self, TypedDict, Union
 from uuid import uuid4
 
 import numpy as np
@@ -1192,8 +1192,19 @@ class ServiceMetricMessage(BECMessage):
 
 class _StrDynamicMetricValue(BaseModel):
     value: str
+    possible_values: list[str]
     type_name: Literal["str"] = "str"
 
+    @model_validator(mode="after")
+    def _validate(self) -> Self:
+        if self.value not in self.possible_values:
+            raise ValueError(
+                f"Invalid string metric: value '{self.value}' not in possible values: '{self.possible_values}'. Provide the possible values for string metrics."
+            )
+        if len(set(self.possible_values)) != len(self.possible_values):
+            raise ValueError(f"Duplicates in possible values: {self.possible_values}")
+        return self
+
 
 class _IntDynamicMetricValue(BaseModel):
     value: int
@@ -1219,6 +1230,14 @@ class _BoolDynamicMetricValue(BaseModel):
 ]
 
 
+class StringMetricDict(TypedDict):
+    value: str
+    possible_values: list[str]
+
+
+DynamicMetricDict = Mapping[str, StringMetricDict | int | float | bool]
+
+
 class DynamicMetricMessage(BECMessage):
     """Message for propagating metrics to the log ingestor.
 
@@ -1230,14 +1249,16 @@ class DynamicMetricMessage(BECMessage):
     timestamp: float = Field(default_factory=time.time)
 
     @classmethod
-    def from_dict(cls, metrics: dict[str, str | int | float | bool], separator: str = "-"):
+    def from_dict(cls, metrics: DynamicMetricDict, separator: str = "-"):
+        def _fmt(value):
+            if isinstance(value, dict):
+                return {"type_name": "str", **value}
+            return {"value": value, "type_name": type(value).__name__}
+
         return cls.model_validate(
             {
-                "metrics": {
-                    # Metric names are concatenated with the group in the ingestor
-                    separator + k: {"value": v, "type_name": type(v).__name__}
-                    for k, v in metrics.items()
-                }
+                # Metric names are concatenated with the group in the ingestor
+                "metrics": {separator + k: _fmt(v) for k, v in metrics.items()}
             }
         )
 

bec_lib/bec_lib/redis_connector.py

@@ -54,6 +54,7 @@
     BECMessage,
     BundleMessage,
     ClientInfoMessage,
+    DynamicMetricDict,
     DynamicMetricMessage,
     ErrorInfo,
 )
@@ -1562,13 +1563,7 @@ def blocking_list_pop(
             return None
         return MsgpackSerialization.loads(raw_msg[1])  # type: ignore # list pop returns one item
 
-    def publish_metrics(
-        self, group_name: str, metrics: dict[str, str | int | float | bool], separator="_"
-    ):
-        if str in set(map(type, metrics.values())):
-            bec_logger.logger.warning(
-                f"String type found in dynamic metric: {group_name}. String metrics are currently not supported and will be ignored in the ingestor."
-            )
+    def publish_metrics(self, group_name: str, metrics: DynamicMetricDict, separator="_"):
         msg = DynamicMetricMessage.from_dict(metrics, separator=separator)
         ep = MessageEndpoints.dynamic_metric(group_name)
         self._redis_conn.publish(ep.endpoint, MsgpackSerialization.dumps(msg))

bec_lib/bec_lib/tests/utils.py

@@ -643,6 +643,8 @@ def redis_server_is_running(self):
     def get_last(self, topic, key):
         return None
 
-    def publish_metrics(self, group_name: str, metrics: dict[str, str | int | float | bool]):
+    def publish_metrics(
+        self, group_name: str, metrics: messages.DynamicMetricDict, separator: str = "_"
+    ):
         if not isinstance(group_name, str) or not isinstance(metrics, dict):
             raise TypeError("invalid arguments to publish_metrics")

bec_lib/tests/test_bec_messages.py

@@ -684,7 +684,12 @@ def test_valid_add_slice_various_indices(self, index):
 
 def test_dynamic_metric_message():
     message = messages.DynamicMetricMessage.from_dict(
-        {"m1": 5, "m2": 5.5, "m3": "test", "m4": True}
+        {
+            "m1": 5,
+            "m2": 5.5,
+            "m3": {"value": "test", "possible_values": ["prod", "test"]},
+            "m4": True,
+        }
     )
     assert isinstance(message.metrics["-m1"], messages._IntDynamicMetricValue)
     assert isinstance(message.metrics["-m2"], messages._FloatDynamicMetricValue)

bec_lib/tests/test_redis_connector_fakeredis.py

@@ -600,7 +600,15 @@ def cb(msg):
         data.append(msg.value)
 
     connected_connector.register(ep, cb=cb, start_thread=False)
-    connected_connector.publish_metrics("test", {"m1": 5, "m2": 5.5, "m3": "test", "m4": True})
+    connected_connector.publish_metrics(
+        "test",
+        {
+            "m1": 5,
+            "m2": 5.5,
+            "m3": {"value": "test", "possible_values": ["test", "prod"]},
+            "m4": True,
+        },
+    )
     connected_connector.poll_messages(timeout=1)
 
     stop = time.time()
@@ -610,6 +618,7 @@ def cb(msg):
     assert res.metrics["_m1"].value == 5
     assert res.metrics["_m2"].value == 5.5
     assert res.metrics["_m3"].value == "test"
+    assert set(res.metrics["_m3"].possible_values) == set(["prod", "test"])
     assert res.metrics["_m4"].value is True
 
 

bec_server/bec_server/procedures/manager.py

@@ -303,10 +303,6 @@ def _publish_metrics(self, msg: ProcedureQNotifMessage):
                 for name in self._helper.get.active_and_pending_queue_names()
             },
         )
-        self._conn.publish_metrics(
-            "active_procedures",
-            {msg.queue: msg.identifier for msg in self._helper.get.running_procedures()},
-        )
 
     def _startup(self):
         # If the server is restarted, clear any pending requests, they'll have to be resubmitted

bec_server/bec_server/scihub/atlas/atlas_active_client_emitter.py

@@ -19,10 +19,9 @@
 
 
 class AtlasActiveClientEmitter:
-
     def __init__(self, atlas_connector: AtlasConnector) -> None:
         self.atlas_connector = atlas_connector
-        self._last_active_client_metrics: dict[str, messages.DynamicMetricValue] | None = None
+        self._last_active_client_metrics: messages.DynamicMetricDict = {}
         self._start_service_status_subscription()
         self._emit_active_client_metrics()
 
@@ -41,20 +40,29 @@ def _emit_active_client_metrics(self) -> None:
         self._last_active_client_metrics = metrics
         self.atlas_connector.connector.publish_metrics("active_clients", metrics)
 
-    def _get_active_client_metrics(self) -> dict[str, str]:
-        active_clients = {}
-        for service_name, status_msg in self.atlas_connector.scihub.service_status.items():
-            if self._is_internal_service(service_name):
-                continue
-            if status_msg.status != messages.BECStatus.RUNNING:
-                continue
-            info = status_msg.info
-            hostname = (
-                info.hostname if isinstance(info, messages.ServiceInfo) else info.get("hostname")
-            )
-            if hostname:
-                active_clients[service_name] = hostname
-        return active_clients
+    def _get_active_client_metrics(self) -> messages.DynamicMetricDict:
+        clients = [
+            (name, status)
+            for (name, status) in self.atlas_connector.scihub.service_status.items()
+            if not (self._is_internal_service(name) or status.status != messages.BECStatus.RUNNING)
+        ]
+        ipython_clients = len(list(filter(lambda ns: "BECIPythonClient" in ns[0], clients)))
+        bec_clients = len(list(filter(lambda ns: "BECClient" in ns[0], clients)))
+
+        def _from_console(name_msg: tuple[str, messages.StatusMessage]):
+            msg = name_msg[1]
+            if not isinstance(msg.info, messages.ServiceInfo):
+                return False
+            return "cons" in msg.info.hostname
+
+        from_consoles = len(list(filter(_from_console, clients)))
+
+        return {
+            "total_connected": len(clients),
+            "bec_clients": bec_clients,
+            "ipython_clients": ipython_clients,
+            "from_consoles": from_consoles,
+        }
 
     @classmethod
     def _is_internal_service(cls, service_name: str) -> bool:

bec_server/bec_server/scihub/atlas/atlas_connector.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+from enum import StrEnum
 from typing import TYPE_CHECKING
 
 from dotenv import dotenv_values
@@ -22,8 +23,29 @@
 logger = bec_logger.logger
 
 
-class AtlasConnector:
+class AtlasHost(StrEnum):
+    NotSet = "none"
+    Local = "localhost"
+    Dev = "bec-atlas-dev"
+    Prod = "bec-atlas-prod"
+    Qa = "bec-atlas-qa"
+    Unknown = "unknown"
+
+    @classmethod
+    def translate_host(cls, hostname: str | None):
+        if not hostname:
+            return cls.NotSet
+        for opt in [cls.Local, cls.Dev, cls.Prod, cls.Qa]:
+            if opt in hostname:
+                return opt
+        return cls.Unknown
+
+    @classmethod
+    def fmt_for_metric(cls, hostname: str | None) -> messages.StringMetricDict:
+        return {"value": cls.translate_host(hostname), "possible_values": list(cls)}
 
+
+class AtlasConnector:
     def __init__(
         self, scihub: SciHub, connector: RedisConnector, redis_atlas: RedisConnector | None = None
     ) -> None:
@@ -54,7 +76,8 @@ def start(self):
         self.update_acls()
         self.update_available_endpoints()
         self.connector.publish_metrics(
-            "atlas_connection", {"connected": self.connected_to_atlas, "host": self.host or ""}
+            "atlas_connection",
+            {"connected": self.connected_to_atlas, "host": AtlasHost.fmt_for_metric(self.host)},
         )
 
     @property

bec_server/tests/tests_scihub/test_atlas_active_client_emitter.py

@@ -36,16 +36,18 @@ def test_get_active_client_metrics_filters_internal_services(atlas_connector):
     ) as mock_service_status:
         mock_service_status.return_value = service_status
         assert emitter._get_active_client_metrics() == {
-            "BECIPythonClient/client-1": "client-host-1",
-            "BECClient/client-2": "client-host-2",
+            "bec_clients": 1,
+            "from_consoles": 0,
+            "ipython_clients": 1,
+            "total_connected": 2,
         }
 
 
 def test_handle_service_status_emits_active_client_metric(atlas_connector):
     emitter = atlas_connector.active_client_emitter
     service_status = {
         "BECIPythonClient/client-1": create_status_message(
-            "BECIPythonClient/client-1", hostname="client-host-1"
+            "BECIPythonClient/client-1", hostname="client-cons-01"
         ),
         "ScanServer": create_status_message("ScanServer", hostname="server-host"),
     }
@@ -59,7 +61,8 @@ def test_handle_service_status_emits_active_client_metric(atlas_connector):
         mock_service_status.return_value = service_status
         emitter._handle_service_status(None)
     mock_publish_metrics.assert_called_once_with(
-        "active_clients", {"BECIPythonClient/client-1": "client-host-1"}
+        "active_clients",
+        {"total_connected": 1, "bec_clients": 0, "ipython_clients": 1, "from_consoles": 1},
     )
 
 

bec_server/tests/tests_scihub/test_atlas_connector.py

@@ -4,7 +4,7 @@
 
 from bec_lib import messages
 from bec_lib.endpoints import MessageEndpoints
-from bec_server.scihub.atlas.atlas_connector import AtlasConnector
+from bec_server.scihub.atlas.atlas_connector import AtlasConnector, AtlasHost
 
 
 def test_atlas_connector_load_env(SciHubMock, connected_atlas_connector):
@@ -103,3 +103,20 @@ def test_atlas_connector_retries_without_ssl_if_tls_fails(SciHubMock, connected_
                 mock_auth.side_effect = [Exception("SSL connection failed"), mock.DEFAULT]
                 atlas_connector.connect_to_atlas()
                 assert atlas_connector.use_tls is False
+
+
+def test_host_enum_can_be_used_for_string_metric():
+    poss_strs = ["none", "localhost", "bec-atlas-dev", "bec-atlas-prod", "bec-atlas-qa", "unknown"]
+    formatted = AtlasHost.fmt_for_metric("bec-atlas-dev.psi.ch")
+    assert formatted == {"value": AtlasHost.Dev, "possible_values": poss_strs}
+    msg_value = messages._StrDynamicMetricValue.model_validate(formatted)
+    assert msg_value.possible_values == poss_strs
+
+
+def test_translate_atlas_host():
+    assert AtlasHost.translate_host(None) == "none"
+    assert AtlasHost.translate_host("some_other_host") == "unknown"
+    assert AtlasHost.translate_host("bec-atlas-dev.psi.ch") == "bec-atlas-dev"
+    assert AtlasHost.translate_host("bec-atlas-prod.psi.ch") == "bec-atlas-prod"
+    assert AtlasHost.translate_host("bec-atlas-qa.psi.ch") == "bec-atlas-qa"
+    assert AtlasHost.translate_host("localhost") == "localhost"