mlrun 1.7.2rc3__py3-none-any.whl → 1.8.0rc2__py3-none-any.whl

mlrun/model_monitoring/db/tsdb/base.py

@@ -15,10 +15,9 @@
 import typing
 from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import Union
 
 import pandas as pd
-import pydantic
+import pydantic.v1
 
 import mlrun.common.schemas.model_monitoring as mm_schemas
 import mlrun.model_monitoring.db.tsdb.helpers
@@ -132,6 +131,7 @@ class TSDBConnector(ABC):
         end: datetime,
         metrics: list[mm_schemas.ModelEndpointMonitoringMetric],
         type: typing.Literal["metrics", "results"],
+        with_result_extra_data: bool,
     ) -> typing.Union[
         list[
             typing.Union[
@@ -150,11 +150,13 @@ class TSDBConnector(ABC):
         Read metrics OR results from the TSDB and return as a list.
 
         :param endpoint_id: The model endpoint identifier.
-        :param start:       The start time of the query.
-        :param end:         The end time of the query.
-        :param metrics:     The list of metrics to get the values for.
-        :param type:        "metrics" or "results" - the type of each item in metrics.
-        :return:            A list of result values or a list of metric values.
+        :param start:                  The start time of the query.
+        :param end:                    The end time of the query.
+        :param metrics:                The list of metrics to get the values for.
+        :param type:                   "metrics" or "results" - the type of each item in metrics.
+        :param with_result_extra_data: Whether to include the extra data in the results, relevant only when
+                                       `type="results"`.
+        :return:                        A list of result values or a list of metric values.
         """
 
     @abstractmethod
@@ -193,9 +195,9 @@ class TSDBConnector(ABC):
     @abstractmethod
     def get_last_request(
         self,
-        endpoint_ids: Union[str, list[str]],
-        start: Union[datetime, str] = "0",
-        end: Union[datetime, str] = "now",
+        endpoint_ids: typing.Union[str, list[str]],
+        start: typing.Optional[datetime] = None,
+        end: typing.Optional[datetime] = None,
     ) -> pd.DataFrame:
         """
         Fetches data from the predictions TSDB table and returns the most recent request
@@ -212,9 +214,9 @@ class TSDBConnector(ABC):
     @abstractmethod
     def get_drift_status(
         self,
-        endpoint_ids: Union[str, list[str]],
-        start: Union[datetime, str] = "now-24h",
-        end: Union[datetime, str] = "now",
+        endpoint_ids: typing.Union[str, list[str]],
+        start: typing.Optional[datetime] = None,
+        end: typing.Optional[datetime] = None,
     ) -> pd.DataFrame:
         """
         Fetches data from the app-results TSDB table and returns the highest status among all
@@ -233,8 +235,8 @@ class TSDBConnector(ABC):
     def get_metrics_metadata(
         self,
         endpoint_id: str,
-        start: Union[datetime, str] = "0",
-        end: Union[datetime, str] = "now",
+        start: typing.Optional[datetime] = None,
+        end: typing.Optional[datetime] = None,
     ) -> pd.DataFrame:
         """
         Fetches distinct metrics metadata from the metrics TSDB table for a specified model endpoint.
@@ -251,8 +253,8 @@ class TSDBConnector(ABC):
     def get_results_metadata(
         self,
         endpoint_id: str,
-        start: Union[datetime, str] = "0",
-        end: Union[datetime, str] = "now",
+        start: typing.Optional[datetime] = None,
+        end: typing.Optional[datetime] = None,
     ) -> pd.DataFrame:
         """
         Fetches distinct results metadata from the app-results TSDB table for a specified model endpoint.
@@ -268,9 +270,9 @@ class TSDBConnector(ABC):
     @abstractmethod
     def get_error_count(
         self,
-        endpoint_ids: Union[str, list[str]],
-        start: Union[datetime, str] = "0",
-        end: Union[datetime, str] = "now",
+        endpoint_ids: typing.Union[str, list[str]],
+        start: typing.Optional[datetime] = None,
+        end: typing.Optional[datetime] = None,
     ) -> pd.DataFrame:
         """
         Fetches data from the error TSDB table and returns the error count for each specified endpoint.
@@ -286,9 +288,9 @@ class TSDBConnector(ABC):
     @abstractmethod
     def get_avg_latency(
         self,
-        endpoint_ids: Union[str, list[str]],
-        start: Union[datetime, str] = "0",
-        end: Union[datetime, str] = "now",
+        endpoint_ids: typing.Union[str, list[str]],
+        start: typing.Optional[datetime] = None,
+        end: typing.Optional[datetime] = None,
     ) -> pd.DataFrame:
         """
         Fetches data from the predictions TSDB table and returns the average latency for each specified endpoint
@@ -420,11 +422,12 @@ class TSDBConnector(ABC):
                                 sub_df.index,
                                 sub_df[mm_schemas.ResultData.RESULT_VALUE],
                                 sub_df[mm_schemas.ResultData.RESULT_STATUS],
+                                sub_df[mm_schemas.ResultData.RESULT_EXTRA_DATA],
                             )
                         ),  # pyright: ignore[reportArgumentType]
                     )
                 )
-            except pydantic.ValidationError:
+            except pydantic.v1.ValidationError:
                 logger.exception(
                     "Failed to convert data-frame into `ModelEndpointMonitoringResultValues`",
                     full_name=full_name,
@@ -446,3 +449,52 @@ class TSDBConnector(ABC):
             )
 
         return metrics_values
+
+    @staticmethod
+    def df_to_metrics_list(
+        *,
+        df: pd.DataFrame,
+        project: str,
+        type: str,
+    ) -> list[mm_schemas.ModelEndpointMonitoringMetric]:
+        """
+        Parse a DataFrame of metrics from the TSDB into a list of mm metrics objects.
+
+        :param df:      The DataFrame to parse.
+        :param project: The project name.
+        :param type:    The type of the metrics (either "result" or "metric").
+
+        :return:        A list of mm metrics objects.
+        """
+        return list(
+            map(
+                lambda record: mm_schemas.ModelEndpointMonitoringMetric(
+                    project=project,
+                    type=type,
+                    app=record.get(mm_schemas.WriterEvent.APPLICATION_NAME),
+                    name=record.get(mm_schemas.ResultData.RESULT_NAME)
+                    or record.get(mm_schemas.MetricData.METRIC_NAME),
+                    kind=record.get(mm_schemas.ResultData.RESULT_KIND),
+                ),
+                df.to_dict("records"),
+            )
+        )
+
+    @staticmethod
+    def _get_start_end(
+        start: typing.Union[datetime, None],
+        end: typing.Union[datetime, None],
+    ) -> tuple[datetime, datetime]:
+        """
+        static utils function for tsdb start end format
+        :param start:       Either None or datetime, None is handled as datetime.min(tz=timezone.utc)
+        :param end:         Either None or datetime, None is handled as datetime.now(tz=timezone.utc)
+        :return:            start datetime, end datetime
+        """
+        start = start or mlrun.utils.datetime_min()
+        end = end or mlrun.utils.datetime_now()
+        if not (isinstance(start, datetime) and isinstance(end, datetime)):
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "Both start and end must be datetime objects"
+            )
+        return start, end

mlrun/model_monitoring/db/tsdb/tdengine/schemas.py

@@ -26,7 +26,7 @@ _MODEL_MONITORING_DATABASE = "mlrun_model_monitoring"
 
 
 class _TDEngineColumnType:
-    def __init__(self, data_type: str, length: int = None):
+    def __init__(self, data_type: str, length: Optional[int] = None):
         self.data_type = data_type
         self.length = length
 
@@ -46,7 +46,7 @@ class _TDEngineColumn(mlrun.common.types.StrEnum):
     INT = _TDEngineColumnType("INT")
     BINARY_40 = _TDEngineColumnType("BINARY", 40)
     BINARY_64 = _TDEngineColumnType("BINARY", 64)
-    BINARY_10000 = _TDEngineColumnType("BINARY", 10000)
+    BINARY_1000 = _TDEngineColumnType("BINARY", 1000)
 
 
 def values_to_column(values, column_type):
@@ -61,7 +61,7 @@ def values_to_column(values, column_type):
         return taosws.binary_to_column(values)
     if column_type == _TDEngineColumn.BINARY_64:
         return taosws.binary_to_column(values)
-    if column_type == _TDEngineColumn.BINARY_10000:
+    if column_type == _TDEngineColumn.BINARY_1000:
         return taosws.binary_to_column(values)
 
     raise mlrun.errors.MLRunInvalidArgumentError(
@@ -170,7 +170,7 @@ class TDEngineSchema:
         table: str,
         start: datetime.datetime,
         end: datetime.datetime,
-        columns_to_filter: list[str] = None,
+        columns_to_filter: Optional[list[str]] = None,
         filter_query: Optional[str] = None,
         interval: Optional[str] = None,
         limit: int = 0,
@@ -178,6 +178,10 @@ class TDEngineSchema:
         sliding_window_step: Optional[str] = None,
         timestamp_column: str = "time",
         database: str = _MODEL_MONITORING_DATABASE,
+        group_by: Optional[Union[list[str], str]] = None,
+        preform_agg_funcs_columns: Optional[list[str]] = None,
+        order_by: Optional[str] = None,
+        desc: Optional[bool] = None,
     ) -> str:
         if agg_funcs and not columns_to_filter:
             raise mlrun.errors.MLRunInvalidArgumentError(
@@ -194,15 +198,37 @@ class TDEngineSchema:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "`interval` must be provided when using sliding window"
             )
+        if group_by and not agg_funcs:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "aggregate functions must be provided when using group by"
+            )
+        if desc and not order_by:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "`order_by` must be provided when using descending"
+            )
 
         with StringIO() as query:
             query.write("SELECT ")
             if interval:
                 query.write("_wstart, _wend, ")
             if agg_funcs:
+                preform_agg_funcs_columns = (
+                    columns_to_filter
+                    if preform_agg_funcs_columns is None
+                    else preform_agg_funcs_columns
+                )
                 query.write(
                     ", ".join(
-                        [f"{a}({col})" for a in agg_funcs for col in columns_to_filter]
+                        [
+                            f"{a}({col})"
+                            if col.upper()
+                            in map(
+                                str.upper, preform_agg_funcs_columns
+                            )  # Case-insensitive check
+                            else f"{col}"
+                            for a in agg_funcs
+                            for col in columns_to_filter
+                        ]
                     )
                 )
             elif columns_to_filter:
@@ -219,6 +245,13 @@ class TDEngineSchema:
                     query.write(f"{timestamp_column} >= '{start}' AND ")
                 if end:
                     query.write(f"{timestamp_column} <= '{end}'")
+            if group_by:
+                if isinstance(group_by, list):
+                    group_by = ", ".join(group_by)
+                query.write(f" GROUP BY {group_by}")
+            if order_by:
+                desc = " DESC" if desc else ""
+                query.write(f" ORDER BY {order_by}{desc}")
             if interval:
                 query.write(f" INTERVAL({interval})")
             if sliding_window_step:
@@ -238,6 +271,7 @@ class AppResultTable(TDEngineSchema):
             mm_schemas.WriterEvent.START_INFER_TIME: _TDEngineColumn.TIMESTAMP,
             mm_schemas.ResultData.RESULT_VALUE: _TDEngineColumn.FLOAT,
             mm_schemas.ResultData.RESULT_STATUS: _TDEngineColumn.INT,
+            mm_schemas.ResultData.RESULT_EXTRA_DATA: _TDEngineColumn.BINARY_1000,
         }
         tags = {
             mm_schemas.WriterEvent.ENDPOINT_ID: _TDEngineColumn.BINARY_64,
@@ -284,10 +318,31 @@ class Predictions(TDEngineSchema):
         columns = {
             mm_schemas.EventFieldType.TIME: _TDEngineColumn.TIMESTAMP,
             mm_schemas.EventFieldType.LATENCY: _TDEngineColumn.FLOAT,
-            mm_schemas.EventKeyMetrics.CUSTOM_METRICS: _TDEngineColumn.BINARY_10000,
+            mm_schemas.EventKeyMetrics.CUSTOM_METRICS: _TDEngineColumn.BINARY_1000,
+        }
+        tags = {
+            mm_schemas.WriterEvent.ENDPOINT_ID: _TDEngineColumn.BINARY_64,
+        }
+        super().__init__(
+            super_table=super_table,
+            columns=columns,
+            tags=tags,
+            database=database,
+            project=project,
+        )
+
+
+@dataclass
+class Errors(TDEngineSchema):
+    def __init__(self, project: str, database: Optional[str] = None):
+        super_table = mm_schemas.TDEngineSuperTables.ERRORS
+        columns = {
+            mm_schemas.EventFieldType.TIME: _TDEngineColumn.TIMESTAMP,
+            mm_schemas.EventFieldType.MODEL_ERROR: _TDEngineColumn.BINARY_1000,
         }
         tags = {
             mm_schemas.WriterEvent.ENDPOINT_ID: _TDEngineColumn.BINARY_64,
+            mm_schemas.EventFieldType.ERROR_TYPE: _TDEngineColumn.BINARY_64,
         }
         super().__init__(
             super_table=super_table,

mlrun/model_monitoring/db/tsdb/tdengine/stream_graph_steps.py

@@ -13,12 +13,14 @@
 # limitations under the License.
 
 import json
+from datetime import datetime
 
 import mlrun.feature_store.steps
 from mlrun.common.schemas.model_monitoring import (
     EventFieldType,
     EventKeyMetrics,
 )
+from mlrun.utils import logger
 
 
 class ProcessBeforeTDEngine(mlrun.feature_store.steps.MapClass):
@@ -40,3 +42,34 @@ class ProcessBeforeTDEngine(mlrun.feature_store.steps.MapClass):
         event[EventFieldType.TABLE_COLUMN] = "_" + event.get(EventFieldType.ENDPOINT_ID)
 
         return event
+
+
+class ErrorExtractor(mlrun.feature_store.steps.MapClass):
+    def __init__(self, **kwargs):
+        """
+        Prepare the event for insertion into the TDEngine error table
+        """
+        super().__init__(**kwargs)
+
+    def do(self, event):
+        error = str(event.get("error"))
+        if len(error) > 1000:
+            error = error[-1000:]
+            logger.warning(
+                f"Error message exceeds 1000 chars: The error message writen to TSDB will be it last "
+                f"1000 chars, Error:  {error}",
+                event=event,
+            )
+        timestamp = datetime.fromisoformat(event.get("when"))
+        endpoint_id = event[EventFieldType.ENDPOINT_ID]
+        event = {
+            EventFieldType.MODEL_ERROR: error,
+            EventFieldType.ERROR_TYPE: EventFieldType.INFER_ERROR,
+            EventFieldType.ENDPOINT_ID: endpoint_id,
+            EventFieldType.TIME: timestamp,
+            EventFieldType.PROJECT: event[EventFieldType.FUNCTION_URI].split("/")[0],
+            EventFieldType.TABLE_COLUMN: "_err_"
+            + event.get(EventFieldType.ENDPOINT_ID),
+        }
+        logger.info("Write error to errors TSDB table", event=event)
+        return event