mlrun 1.7.0rc17__py3-none-any.whl → 1.7.0rc19__py3-none-any.whl

mlrun/model_monitoring/db/stores/v3io_kv/kv_store.py

@@ -15,23 +15,79 @@
 import json
 import os
 import typing
+from dataclasses import dataclass
 from http import HTTPStatus
 
 import v3io.dataplane
+import v3io.dataplane.output
 import v3io.dataplane.response
 
 import mlrun.common.model_monitoring.helpers
-import mlrun.common.schemas.model_monitoring as mm_constants
+import mlrun.common.schemas.model_monitoring as mm_schemas
 import mlrun.model_monitoring.db
 import mlrun.utils.v3io_clients
 from mlrun.utils import logger
 
 # Fields to encode before storing in the KV table or to decode after retrieving
 fields_to_encode_decode = [
-    mm_constants.EventFieldType.FEATURE_STATS,
-    mm_constants.EventFieldType.CURRENT_STATS,
+    mm_schemas.EventFieldType.FEATURE_STATS,
+    mm_schemas.EventFieldType.CURRENT_STATS,
 ]
 
+_METRIC_FIELDS: list[str] = [
+    mm_schemas.WriterEvent.APPLICATION_NAME,
+    mm_schemas.MetricData.METRIC_NAME,
+    mm_schemas.MetricData.METRIC_VALUE,
+    mm_schemas.WriterEvent.START_INFER_TIME,
+    mm_schemas.WriterEvent.END_INFER_TIME,
+]
+
+
+class SchemaField(typing.TypedDict):
+    name: str
+    type: str
+    nullable: bool
+
+
+@dataclass
+class SchemaParams:
+    key: str
+    fields: list[SchemaField]
+
+
+_RESULT_SCHEMA: list[SchemaField] = [
+    SchemaField(
+        name=mm_schemas.ResultData.RESULT_NAME,
+        type=mm_schemas.GrafanaColumnType.STRING,
+        nullable=False,
+    )
+]
+
+_METRIC_SCHEMA: list[SchemaField] = [
+    SchemaField(
+        name=mm_schemas.WriterEvent.APPLICATION_NAME,
+        type=mm_schemas.GrafanaColumnType.STRING,
+        nullable=False,
+    ),
+    SchemaField(
+        name=mm_schemas.MetricData.METRIC_NAME,
+        type=mm_schemas.GrafanaColumnType.STRING,
+        nullable=False,
+    ),
+]
+
+
+_KIND_TO_SCHEMA_PARAMS: dict[mm_schemas.WriterEventKind, SchemaParams] = {
+    mm_schemas.WriterEventKind.RESULT: SchemaParams(
+        key=mm_schemas.WriterEvent.APPLICATION_NAME, fields=_RESULT_SCHEMA
+    ),
+    mm_schemas.WriterEventKind.METRIC: SchemaParams(
+        key="metric_id", fields=_METRIC_SCHEMA
+    ),
+}
+
+_EXCLUDE_SCHEMA_FILTER_EXPRESSION = '__name!=".#schema"'
+
 
 class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
     """
@@ -64,7 +120,7 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
         self.client.kv.put(
             container=self.container,
             table_path=self.path,
-            key=endpoint[mm_constants.EventFieldType.UID],
+            key=endpoint[mm_schemas.EventFieldType.UID],
             attributes=endpoint,
         )
 
@@ -151,7 +207,7 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
         """Getting path and container based on the model monitoring configurations"""
         path = mlrun.mlconf.model_endpoint_monitoring.store_prefixes.default.format(
             project=self.project,
-            kind=mlrun.common.schemas.ModelMonitoringStoreKinds.ENDPOINTS,
+            kind=mm_schemas.ModelMonitoringStoreKinds.ENDPOINTS,
         )
         (
             _,
@@ -217,11 +273,11 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
         if uids is None:
             uids = []
             for item in items:
-                if mm_constants.EventFieldType.UID not in item:
+                if mm_schemas.EventFieldType.UID not in item:
                     # This is kept for backwards compatibility - in old versions the key column named endpoint_id
-                    uids.append(item[mm_constants.EventFieldType.ENDPOINT_ID])
+                    uids.append(item[mm_schemas.EventFieldType.ENDPOINT_ID])
                 else:
-                    uids.append(item[mm_constants.EventFieldType.UID])
+                    uids.append(item[mm_schemas.EventFieldType.UID])
 
         # Add each relevant model endpoint to the model endpoints list
         for endpoint_id in uids:
@@ -241,11 +297,11 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
 
         # Delete model endpoint record from KV table
         for endpoint_dict in endpoints:
-            if mm_constants.EventFieldType.UID not in endpoint_dict:
+            if mm_schemas.EventFieldType.UID not in endpoint_dict:
                 # This is kept for backwards compatibility - in old versions the key column named endpoint_id
-                endpoint_id = endpoint_dict[mm_constants.EventFieldType.ENDPOINT_ID]
+                endpoint_id = endpoint_dict[mm_schemas.EventFieldType.ENDPOINT_ID]
             else:
-                endpoint_id = endpoint_dict[mm_constants.EventFieldType.UID]
+                endpoint_id = endpoint_dict[mm_schemas.EventFieldType.UID]
             self.delete_model_endpoint(
                 endpoint_id,
             )
@@ -268,11 +324,19 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
                 raise_for_status=v3io.dataplane.RaiseForStatus.never,
             )
 
+    @staticmethod
+    def _get_results_table_path(endpoint_id: str) -> str:
+        return endpoint_id
+
+    @staticmethod
+    def _get_metrics_table_path(endpoint_id: str) -> str:
+        return f"{endpoint_id}_metrics"
+
     def write_application_event(
         self,
         event: dict[str, typing.Any],
-        kind: mm_constants.WriterEventKind = mm_constants.WriterEventKind.RESULT,
-    ):
+        kind: mm_schemas.WriterEventKind = mm_schemas.WriterEventKind.RESULT,
+    ) -> None:
         """
         Write a new application event in the target table.
 
@@ -281,66 +345,63 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
                       object.
         :param kind: The type of the event, can be either "result" or "metric".
         """
-        if kind == mm_constants.WriterEventKind.METRIC:
-            # TODO : Implement the logic for writing metrics to KV
-            return
-
-        endpoint_id = event.pop(mm_constants.WriterEvent.ENDPOINT_ID)
-        app_name = event.pop(mm_constants.WriterEvent.APPLICATION_NAME)
-        metric_name = event.pop(mm_constants.ResultData.RESULT_NAME)
-        attributes = {metric_name: json.dumps(event)}
 
-        v3io_monitoring_apps_container = self.get_v3io_monitoring_apps_container(
-            project_name=self.project
-        )
+        container = self.get_v3io_monitoring_apps_container(project_name=self.project)
+        endpoint_id = event.pop(mm_schemas.WriterEvent.ENDPOINT_ID)
+
+        if kind == mm_schemas.WriterEventKind.METRIC:
+            table_path = self._get_metrics_table_path(endpoint_id)
+            key = f"{event[mm_schemas.WriterEvent.APPLICATION_NAME]}.{event[mm_schemas.MetricData.METRIC_NAME]}"
+            attributes = {event_key: event[event_key] for event_key in _METRIC_FIELDS}
+        elif kind == mm_schemas.WriterEventKind.RESULT:
+            table_path = self._get_results_table_path(endpoint_id)
+            key = event.pop(mm_schemas.WriterEvent.APPLICATION_NAME)
+            metric_name = event.pop(mm_schemas.ResultData.RESULT_NAME)
+            attributes = {metric_name: json.dumps(event)}
+        else:
+            raise ValueError(f"Invalid {kind = }")
 
         self.client.kv.update(
-            container=v3io_monitoring_apps_container,
-            table_path=endpoint_id,
-            key=app_name,
+            container=container,
+            table_path=table_path,
+            key=key,
             attributes=attributes,
         )
 
         schema_file = self.client.kv.new_cursor(
-            container=v3io_monitoring_apps_container,
-            table_path=endpoint_id,
+            container=container,
+            table_path=table_path,
             filter_expression='__name==".#schema"',
         )
 
         if not schema_file.all():
             logger.info(
-                "Generate a new V3IO KV schema file",
-                container=v3io_monitoring_apps_container,
-                endpoint_id=endpoint_id,
+                "Generating a new V3IO KV schema file",
+                container=container,
+                table_path=table_path,
+            )
+            self._generate_kv_schema(
+                container=container, table_path=table_path, kind=kind
             )
-            self._generate_kv_schema(endpoint_id, v3io_monitoring_apps_container)
-        logger.info("Updated V3IO KV successfully", key=app_name)
+        logger.info("Updated V3IO KV successfully", key=key)
 
     def _generate_kv_schema(
-        self, endpoint_id: str, v3io_monitoring_apps_container: str
-    ):
+        self, *, container: str, table_path: str, kind: mm_schemas.WriterEventKind
+    ) -> None:
         """Generate V3IO KV schema file which will be used by the model monitoring applications dashboard in Grafana."""
-        fields = [
-            {
-                "name": mm_constants.ResultData.RESULT_NAME,
-                "type": "string",
-                "nullable": False,
-            }
-        ]
+        schema_params = _KIND_TO_SCHEMA_PARAMS[kind]
         res = self.client.kv.create_schema(
-            container=v3io_monitoring_apps_container,
-            table_path=endpoint_id,
-            key=mm_constants.WriterEvent.APPLICATION_NAME,
-            fields=fields,
+            container=container,
+            table_path=table_path,
+            key=schema_params.key,
+            fields=schema_params.fields,
         )
         if res.status_code != HTTPStatus.OK:
             raise mlrun.errors.MLRunBadRequestError(
-                f"Couldn't infer schema for endpoint {endpoint_id} which is required for Grafana dashboards"
+                f"Couldn't infer schema for endpoint {table_path} which is required for Grafana dashboards"
             )
         else:
-            logger.info(
-                "Generated V3IO KV schema successfully", endpoint_id=endpoint_id
-            )
+            logger.info("Generated V3IO KV schema successfully", table_path=table_path)
 
     def get_last_analyzed(self, endpoint_id: str, application_name: str) -> int:
         """
@@ -361,7 +422,7 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
                 table_path=endpoint_id,
                 key=application_name,
             )
-            return data.output.item[mm_constants.SchedulingKeys.LAST_ANALYZED]
+            return data.output.item[mm_schemas.SchedulingKeys.LAST_ANALYZED]
         except v3io.dataplane.response.HttpResponseError as err:
             logger.debug("Error while getting last analyzed time", err=err)
             raise mlrun.errors.MLRunNotFoundError(
@@ -386,7 +447,7 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
             ),
             table_path=endpoint_id,
             key=application_name,
-            attributes={mm_constants.SchedulingKeys.LAST_ANALYZED: last_analyzed},
+            attributes={mm_schemas.SchedulingKeys.LAST_ANALYZED: last_analyzed},
         )
 
     def _generate_tsdb_paths(self) -> tuple[str, str]:
@@ -399,7 +460,7 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
         full_path = (
             mlrun.mlconf.model_endpoint_monitoring.store_prefixes.default.format(
                 project=self.project,
-                kind=mlrun.common.schemas.ModelMonitoringStoreKinds.EVENTS,
+                kind=mm_schemas.ModelMonitoringStoreKinds.EVENTS,
             )
         )
 
@@ -495,8 +556,8 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
         # Apply top_level filter (remove endpoints that considered a child of a router)
         if top_level:
             filter_expression.append(
-                f"(endpoint_type=='{str(mm_constants.EndpointType.NODE_EP.value)}' "
-                f"OR  endpoint_type=='{str(mm_constants.EndpointType.ROUTER.value)}')"
+                f"(endpoint_type=='{str(mm_schemas.EndpointType.NODE_EP.value)}' "
+                f"OR  endpoint_type=='{str(mm_schemas.EndpointType.ROUTER.value)}')"
             )
 
         return " AND ".join(filter_expression)
@@ -516,30 +577,30 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
         # Validate default value for `error_count`
         # For backwards compatibility reasons, we validate that the model endpoint includes the `error_count` key
         if (
-            mm_constants.EventFieldType.ERROR_COUNT in endpoint
-            and endpoint[mm_constants.EventFieldType.ERROR_COUNT] == "null"
+            mm_schemas.EventFieldType.ERROR_COUNT in endpoint
+            and endpoint[mm_schemas.EventFieldType.ERROR_COUNT] == "null"
         ):
-            endpoint[mm_constants.EventFieldType.ERROR_COUNT] = "0"
+            endpoint[mm_schemas.EventFieldType.ERROR_COUNT] = "0"
 
         # Validate default value for `metrics`
         # For backwards compatibility reasons, we validate that the model endpoint includes the `metrics` key
         if (
-            mm_constants.EventFieldType.METRICS in endpoint
-            and endpoint[mm_constants.EventFieldType.METRICS] == "null"
+            mm_schemas.EventFieldType.METRICS in endpoint
+            and endpoint[mm_schemas.EventFieldType.METRICS] == "null"
         ):
-            endpoint[mm_constants.EventFieldType.METRICS] = json.dumps(
+            endpoint[mm_schemas.EventFieldType.METRICS] = json.dumps(
                 {
-                    mm_constants.EventKeyMetrics.GENERIC: {
-                        mm_constants.EventLiveStats.LATENCY_AVG_1H: 0,
-                        mm_constants.EventLiveStats.PREDICTIONS_PER_SECOND: 0,
+                    mm_schemas.EventKeyMetrics.GENERIC: {
+                        mm_schemas.EventLiveStats.LATENCY_AVG_1H: 0,
+                        mm_schemas.EventLiveStats.PREDICTIONS_PER_SECOND: 0,
                     }
                 }
             )
         # Validate key `uid` instead of `endpoint_id`
         # For backwards compatibility reasons, we replace the `endpoint_id` with `uid` which is the updated key name
-        if mm_constants.EventFieldType.ENDPOINT_ID in endpoint:
-            endpoint[mm_constants.EventFieldType.UID] = endpoint[
-                mm_constants.EventFieldType.ENDPOINT_ID
+        if mm_schemas.EventFieldType.ENDPOINT_ID in endpoint:
+            endpoint[mm_schemas.EventFieldType.UID] = endpoint[
+                mm_schemas.EventFieldType.ENDPOINT_ID
             ]
 
     @staticmethod
@@ -566,57 +627,95 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
     def _get_monitoring_schedules_container(project_name: str) -> str:
         return f"users/pipelines/{project_name}/monitoring-schedules/functions"
 
-    def _extract_metrics_from_items(
+    def _extract_results_from_items(
         self, app_items: list[dict[str, str]]
-    ) -> list[mm_constants.ModelEndpointMonitoringMetric]:
-        metrics: list[mm_constants.ModelEndpointMonitoringMetric] = []
+    ) -> list[mm_schemas.ModelEndpointMonitoringMetric]:
+        """Assuming .#schema items are filtered out"""
+        metrics: list[mm_schemas.ModelEndpointMonitoringMetric] = []
         for app_item in app_items:
-            # See https://www.iguazio.com/docs/latest-release/services/data-layer/reference/system-attributes/#sys-attr-__name
             app_name = app_item.pop("__name")
-            if app_name == ".#schema":
-                continue
             for result_name in app_item:
                 metrics.append(
-                    mm_constants.ModelEndpointMonitoringMetric(
+                    mm_schemas.ModelEndpointMonitoringMetric(
                         project=self.project,
                         app=app_name,
-                        type=mm_constants.ModelEndpointMonitoringMetricType.RESULT,
+                        type=mm_schemas.ModelEndpointMonitoringMetricType.RESULT,
                         name=result_name,
-                        full_name=mlrun.common.schemas.model_monitoring.model_endpoints._compose_full_name(
+                        full_name=mm_schemas.model_endpoints._compose_full_name(
                             project=self.project, app=app_name, name=result_name
                         ),
                     )
                 )
         return metrics
 
+    def _extract_metrics_from_items(
+        self, result_items: list[dict[str, str]]
+    ) -> list[mm_schemas.ModelEndpointMonitoringMetric]:
+        metrics: list[mm_schemas.ModelEndpointMonitoringMetric] = []
+        logger.debug("Result items", result_items=result_items)
+        for result_item in result_items:
+            app = result_item[mm_schemas.WriterEvent.APPLICATION_NAME]
+            name = result_item[mm_schemas.MetricData.METRIC_NAME]
+            metrics.append(
+                mm_schemas.ModelEndpointMonitoringMetric(
+                    project=self.project,
+                    app=app,
+                    type=mm_schemas.ModelEndpointMonitoringMetricType.METRIC,
+                    name=name,
+                    full_name=mm_schemas.model_endpoints._compose_full_name(
+                        project=self.project,
+                        app=app,
+                        name=name,
+                        type=mm_schemas.ModelEndpointMonitoringMetricType.METRIC,
+                    ),
+                )
+            )
+        return metrics
+
     def get_model_endpoint_metrics(
-        self, endpoint_id: str
-    ) -> list[mm_constants.ModelEndpointMonitoringMetric]:
+        self, endpoint_id: str, type: mm_schemas.ModelEndpointMonitoringMetricType
+    ) -> list[mm_schemas.ModelEndpointMonitoringMetric]:
         """Get model monitoring results and metrics on the endpoint"""
-        metrics: list[mm_constants.ModelEndpointMonitoringMetric] = []
+        metrics: list[mm_schemas.ModelEndpointMonitoringMetric] = []
         container = self.get_v3io_monitoring_apps_container(self.project)
+        if type == mm_schemas.ModelEndpointMonitoringMetricType.METRIC:
+            table_path = self._get_metrics_table_path(endpoint_id)
+            items_extractor = self._extract_metrics_from_items
+        elif type == mm_schemas.ModelEndpointMonitoringMetricType.RESULT:
+            table_path = self._get_results_table_path(endpoint_id)
+            items_extractor = self._extract_results_from_items
+        else:
+            raise ValueError(f"Invalid metric {type = }")
+
+        def scan(
+            marker: typing.Optional[str] = None,
+        ) -> v3io.dataplane.response.Response:
+            # TODO: Use AIO client: `v3io.aio.dataplane.client.Client`
+            return self.client.kv.scan(
+                container=container,
+                table_path=table_path,
+                marker=marker,
+                filter_expression=_EXCLUDE_SCHEMA_FILTER_EXPRESSION,
+            )
+
         try:
-            response = self.client.kv.scan(container=container, table_path=endpoint_id)
+            response = scan()
         except v3io.dataplane.response.HttpResponseError as err:
             if err.status_code == HTTPStatus.NOT_FOUND:
                 logger.warning(
-                    "Attempt getting metrics and results - no data. Check the "
+                    f"Attempt getting {type}s - no data. Check the "
                     "project name, endpoint, or wait for the applications to start.",
                     container=container,
-                    table_path=endpoint_id,
+                    table_path=table_path,
                 )
                 return []
             raise
 
         while True:
-            metrics.extend(self._extract_metrics_from_items(response.output.items))
-            if response.output.last:
+            output = typing.cast(v3io.dataplane.output.GetItemsOutput, response.output)
+            metrics.extend(items_extractor(output.items))
+            if output.last:
                 break
-            # TODO: Use AIO client: `v3io.aio.dataplane.client.Client`
-            response = self.client.kv.scan(
-                container=container,
-                table_path=endpoint_id,
-                marker=response.output.next_marker,
-            )
+            response = scan(marker=output.next_marker)
 
         return metrics

mlrun/model_monitoring/db/tsdb/__init__.py

@@ -25,6 +25,7 @@ class ObjectTSDBFactory(enum.Enum):
     """Enum class to handle the different TSDB connector type values for storing real time metrics"""
 
     v3io_tsdb = "v3io-tsdb"
+    tdengine = "tdengine"
 
     def to_tsdb_connector(self, project: str, **kwargs) -> TSDBConnector:
         """
@@ -43,6 +44,13 @@ class ObjectTSDBFactory(enum.Enum):
 
             return V3IOTSDBConnector(project=project, **kwargs)
 
+        # Assuming TDEngine connector if connector type is not V3IO TSDB.
+        # Update these lines once there are more than two connector types.
+
+        from .tdengine.tdengine_connector import TDEngineConnector
+
+        return TDEngineConnector(project=project, **kwargs)
+
     @classmethod
     def _missing_(cls, value: typing.Any):
         """A lookup function to handle an invalid value.
@@ -54,18 +62,39 @@ class ObjectTSDBFactory(enum.Enum):
         )
 
 
-def get_tsdb_connector(project: str, **kwargs) -> TSDBConnector:
+def get_tsdb_connector(
+    project: str,
+    tsdb_connector_type: str = "",
+    secret_provider: typing.Optional[typing.Callable] = None,
+    **kwargs,
+) -> TSDBConnector:
     """
-    Get the TSDB connector type based on mlrun.config.model_endpoint_monitoring.tsdb_connector_type.
+    Get TSDB connector object.
     :param project: The name of the project.
+    :param tsdb_connector_type: The type of the TSDB connector. See mlrun.model_monitoring.db.tsdb.ObjectTSDBFactory
+                                for available options.
+    :param secret_provider: An optional secret provider to get the connection string secret.
+
     :return: `TSDBConnector` object. The main goal of this object is to handle different operations on the
              TSDB connector such as updating drift metrics or write application record result.
     """
 
-    # Get store type value from ObjectTSDBFactory enum class
-    tsdb_connector_type = ObjectTSDBFactory(
-        mlrun.mlconf.model_endpoint_monitoring.tsdb_connector_type
+    tsdb_connection_string = mlrun.model_monitoring.helpers.get_tsdb_connection_string(
+        secret_provider=secret_provider
     )
 
+    if tsdb_connection_string and tsdb_connection_string.startswith("taosws"):
+        tsdb_connector_type = mlrun.common.schemas.model_monitoring.TSDBTarget.TDEngine
+        kwargs["connection_string"] = tsdb_connection_string
+
+    # Set the default TSDB connector type if no connection has been set
+    tsdb_connector_type = (
+        tsdb_connector_type
+        or mlrun.mlconf.model_endpoint_monitoring.tsdb_connector_type
+    )
+
+    # Get connector type value from ObjectTSDBFactory enum class
+    tsdb_connector_factory = ObjectTSDBFactory(tsdb_connector_type)
+
     # Convert into TSDB connector object
-    return tsdb_connector_type.to_tsdb_connector(project=project, **kwargs)
+    return tsdb_connector_factory.to_tsdb_connector(project=project, **kwargs)