mlrun 1.7.0rc5__py3-none-any.whl → 1.7.2__py3-none-any.whl

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

@@ -0,0 +1,522 @@
+# Copyright 2024 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import typing
+from datetime import datetime
+from typing import Union
+
+import pandas as pd
+import taosws
+from taoswswrap.tdengine_connection import (
+    Statement,
+    TDEngineConnection,
+)
+
+import mlrun.common.schemas.model_monitoring as mm_schemas
+import mlrun.model_monitoring.db.tsdb.tdengine.schemas as tdengine_schemas
+import mlrun.model_monitoring.db.tsdb.tdengine.stream_graph_steps
+from mlrun.model_monitoring.db import TSDBConnector
+from mlrun.model_monitoring.db.tsdb.tdengine.schemas import TDEngineSchema
+from mlrun.model_monitoring.helpers import get_invocations_fqn
+from mlrun.utils import logger
+
+
+class TDEngineConnector(TSDBConnector):
+    """
+    Handles the TSDB operations when the TSDB connector is of type TDEngine.
+    """
+
+    type: str = mm_schemas.TSDBTarget.TDEngine
+
+    def __init__(
+        self,
+        project: str,
+        database: str = tdengine_schemas._MODEL_MONITORING_DATABASE,
+        **kwargs,
+    ):
+        super().__init__(project=project)
+        if "connection_string" not in kwargs:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "connection_string is a required parameter for TDEngineConnector."
+            )
+        self._tdengine_connection_string = kwargs.get("connection_string")
+        self.database = database
+
+        self._connection = None
+        self._init_super_tables()
+
+        self._timeout = mlrun.mlconf.model_endpoint_monitoring.tdengine.timeout
+        self._retries = mlrun.mlconf.model_endpoint_monitoring.tdengine.retries
+
+    @property
+    def connection(self) -> TDEngineConnection:
+        if not self._connection:
+            self._connection = self._create_connection()
+        return self._connection
+
+    def _create_connection(self) -> TDEngineConnection:
+        """Establish a connection to the TSDB server."""
+        logger.debug("Creating a new connection to TDEngine", project=self.project)
+        conn = TDEngineConnection(self._tdengine_connection_string)
+        conn.run(
+            statements=f"CREATE DATABASE IF NOT EXISTS {self.database}",
+            timeout=self._timeout,
+            retries=self._retries,
+        )
+        conn.prefix_statements = [f"USE {self.database}"]
+        logger.debug("Connected to TDEngine", project=self.project)
+        return conn
+
+    def _init_super_tables(self):
+        """Initialize the super tables for the TSDB."""
+        self.tables = {
+            mm_schemas.TDEngineSuperTables.APP_RESULTS: tdengine_schemas.AppResultTable(
+                project=self.project, database=self.database
+            ),
+            mm_schemas.TDEngineSuperTables.METRICS: tdengine_schemas.Metrics(
+                project=self.project, database=self.database
+            ),
+            mm_schemas.TDEngineSuperTables.PREDICTIONS: tdengine_schemas.Predictions(
+                project=self.project, database=self.database
+            ),
+        }
+
+    def create_tables(self):
+        """Create TDEngine supertables."""
+        for table in self.tables:
+            create_table_query = self.tables[table]._create_super_table_query()
+            self.connection.run(
+                statements=create_table_query,
+                timeout=self._timeout,
+                retries=self._retries,
+            )
+
+    def write_application_event(
+        self,
+        event: dict,
+        kind: mm_schemas.WriterEventKind = mm_schemas.WriterEventKind.RESULT,
+    ) -> None:
+        """
+        Write a single result or metric to TSDB.
+        """
+
+        table_name = (
+            f"{event[mm_schemas.WriterEvent.ENDPOINT_ID]}_"
+            f"{event[mm_schemas.WriterEvent.APPLICATION_NAME]}"
+        )
+
+        if kind == mm_schemas.WriterEventKind.RESULT:
+            # Write a new result
+            table = self.tables[mm_schemas.TDEngineSuperTables.APP_RESULTS]
+            table_name = (
+                f"{table_name}_{event[mm_schemas.ResultData.RESULT_NAME]}"
+            ).replace("-", "_")
+            event.pop(mm_schemas.ResultData.CURRENT_STATS, None)
+
+        else:
+            # Write a new metric
+            table = self.tables[mm_schemas.TDEngineSuperTables.METRICS]
+            table_name = (
+                f"{table_name}_{event[mm_schemas.MetricData.METRIC_NAME]}"
+            ).replace("-", "_")
+
+        # Escape the table name for case-sensitivity (ML-7908)
+        # https://github.com/taosdata/taos-connector-python/issues/260
+        table_name = f"`{table_name}`"
+
+        # Convert the datetime strings to datetime objects
+        event[mm_schemas.WriterEvent.END_INFER_TIME] = self._convert_to_datetime(
+            val=event[mm_schemas.WriterEvent.END_INFER_TIME]
+        )
+        event[mm_schemas.WriterEvent.START_INFER_TIME] = self._convert_to_datetime(
+            val=event[mm_schemas.WriterEvent.START_INFER_TIME]
+        )
+
+        create_table_sql = table._create_subtable_sql(subtable=table_name, values=event)
+
+        insert_statement = Statement(
+            TDEngineSchema._insert_subtable_stmt,
+            dict(columns=table.columns, subtable=table_name, values=event),
+        )
+
+        self.connection.run(
+            statements=[
+                create_table_sql,
+                insert_statement,
+            ],
+            timeout=self._timeout,
+            retries=self._retries,
+        )
+
+    @staticmethod
+    def _convert_to_datetime(val: typing.Union[str, datetime]) -> datetime:
+        return datetime.fromisoformat(val) if isinstance(val, str) else val
+
+    def apply_monitoring_stream_steps(self, graph):
+        """
+        Apply TSDB steps on the provided monitoring graph. Throughout these steps, the graph stores live data of
+        different key metric dictionaries. This data is being used by the monitoring dashboards in
+        grafana. At the moment, we store two types of data:
+        - prediction latency.
+        - custom metrics.
+        """
+
+        def apply_process_before_tsdb():
+            graph.add_step(
+                "mlrun.model_monitoring.db.tsdb.tdengine.stream_graph_steps.ProcessBeforeTDEngine",
+                name="ProcessBeforeTDEngine",
+                after="MapFeatureNames",
+            )
+
+        def apply_tdengine_target(name, after):
+            graph.add_step(
+                "storey.TDEngineTarget",
+                name=name,
+                after=after,
+                url=self._tdengine_connection_string,
+                supertable=self.tables[
+                    mm_schemas.TDEngineSuperTables.PREDICTIONS
+                ].super_table,
+                table_col=mm_schemas.EventFieldType.TABLE_COLUMN,
+                time_col=mm_schemas.EventFieldType.TIME,
+                database=self.database,
+                columns=[
+                    mm_schemas.EventFieldType.LATENCY,
+                    mm_schemas.EventKeyMetrics.CUSTOM_METRICS,
+                ],
+                tag_cols=[
+                    mm_schemas.EventFieldType.PROJECT,
+                    mm_schemas.EventFieldType.ENDPOINT_ID,
+                ],
+                max_events=1000,
+                flush_after_seconds=30,
+            )
+
+        apply_process_before_tsdb()
+        apply_tdengine_target(
+            name="TDEngineTarget",
+            after="ProcessBeforeTDEngine",
+        )
+
+    def handle_model_error(self, graph, **kwargs) -> None:
+        pass
+
+    def delete_tsdb_resources(self):
+        """
+        Delete all project resources in the TSDB connector, such as model endpoints data and drift results.
+        """
+        logger.debug(
+            "Deleting all project resources using the TDEngine connector",
+            project=self.project,
+        )
+        drop_statements = []
+        for table in self.tables:
+            drop_statements.append(self.tables[table].drop_supertable_query())
+
+        try:
+            self.connection.run(
+                statements=drop_statements,
+                timeout=self._timeout,
+                retries=self._retries,
+            )
+        except Exception as e:
+            logger.warning(
+                "Failed to drop TDEngine tables. You may need to drop them manually. "
+                "These can be found under the following supertables: app_results, "
+                "metrics, and predictions.",
+                project=self.project,
+                error=mlrun.errors.err_to_str(e),
+            )
+        logger.debug(
+            "Deleted all project resources using the TDEngine connector",
+            project=self.project,
+        )
+
+    def get_model_endpoint_real_time_metrics(
+        self,
+        endpoint_id: str,
+        metrics: list[str],
+        start: str,
+        end: str,
+    ) -> dict[str, list[tuple[str, float]]]:
+        # Not implemented, use get_records() instead
+        pass
+
+    def _get_records(
+        self,
+        table: str,
+        start: datetime,
+        end: datetime,
+        columns: typing.Optional[list[str]] = None,
+        filter_query: typing.Optional[str] = None,
+        interval: typing.Optional[str] = None,
+        agg_funcs: typing.Optional[list] = None,
+        limit: typing.Optional[int] = None,
+        sliding_window_step: typing.Optional[str] = None,
+        timestamp_column: str = mm_schemas.EventFieldType.TIME,
+    ) -> pd.DataFrame:
+        """
+        Getting records from TSDB data collection.
+        :param table:                 Either a supertable or a subtable name.
+        :param start:                 The start time of the metrics.
+        :param end:                   The end time of the metrics.
+        :param columns:               Columns to include in the result.
+        :param filter_query:          Optional filter expression as a string. TDengine supports SQL-like syntax.
+        :param interval:              The interval to aggregate the data by. Note that if interval is provided,
+                                      `agg_funcs` must bg provided as well. Provided as a string in the format of '1m',
+                                      '1h', etc.
+        :param agg_funcs:             The aggregation functions to apply on the columns. Note that if `agg_funcs` is
+                                      provided, `interval` must bg provided as well. Provided as a list of strings in
+                                      the format of ['sum', 'avg', 'count', ...].
+        :param limit:                 The maximum number of records to return.
+        :param sliding_window_step:   The time step for which the time window moves forward. Note that if
+                                      `sliding_window_step` is provided, interval must be provided as well. Provided
+                                      as a string in the format of '1m', '1h', etc.
+        :param timestamp_column:      The column name that holds the timestamp index.
+
+        :return: DataFrame with the provided attributes from the data collection.
+        :raise:  MLRunInvalidArgumentError if query the provided table failed.
+        """
+
+        full_query = tdengine_schemas.TDEngineSchema._get_records_query(
+            table=table,
+            start=start,
+            end=end,
+            columns_to_filter=columns,
+            filter_query=filter_query,
+            interval=interval,
+            limit=limit,
+            agg_funcs=agg_funcs,
+            sliding_window_step=sliding_window_step,
+            timestamp_column=timestamp_column,
+            database=self.database,
+        )
+        logger.debug("Querying TDEngine", query=full_query)
+        try:
+            query_result = self.connection.run(
+                query=full_query, timeout=self._timeout, retries=self._retries
+            )
+        except taosws.QueryError as e:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"Failed to query table {table} in database {self.database}, {str(e)}"
+            )
+
+        df_columns = [field.name for field in query_result.fields]
+        return pd.DataFrame(query_result.data, columns=df_columns)
+
+    def read_metrics_data(
+        self,
+        *,
+        endpoint_id: str,
+        start: datetime,
+        end: datetime,
+        metrics: list[mm_schemas.ModelEndpointMonitoringMetric],
+        type: typing.Literal["metrics", "results"],
+    ) -> typing.Union[
+        list[
+            typing.Union[
+                mm_schemas.ModelEndpointMonitoringResultValues,
+                mm_schemas.ModelEndpointMonitoringMetricNoData,
+            ],
+        ],
+        list[
+            typing.Union[
+                mm_schemas.ModelEndpointMonitoringMetricValues,
+                mm_schemas.ModelEndpointMonitoringMetricNoData,
+            ],
+        ],
+    ]:
+        timestamp_column = mm_schemas.WriterEvent.END_INFER_TIME
+        columns = [timestamp_column, mm_schemas.WriterEvent.APPLICATION_NAME]
+        if type == "metrics":
+            table = self.tables[mm_schemas.TDEngineSuperTables.METRICS].super_table
+            name = mm_schemas.MetricData.METRIC_NAME
+            columns += [name, mm_schemas.MetricData.METRIC_VALUE]
+            df_handler = self.df_to_metrics_values
+        elif type == "results":
+            table = self.tables[mm_schemas.TDEngineSuperTables.APP_RESULTS].super_table
+            name = mm_schemas.ResultData.RESULT_NAME
+            columns += [
+                name,
+                mm_schemas.ResultData.RESULT_VALUE,
+                mm_schemas.ResultData.RESULT_STATUS,
+                mm_schemas.ResultData.RESULT_KIND,
+            ]
+            df_handler = self.df_to_results_values
+        else:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"Invalid type {type}, must be either 'metrics' or 'results'."
+            )
+
+        metrics_condition = " OR ".join(
+            [
+                f"({mm_schemas.WriterEvent.APPLICATION_NAME}='{metric.app}' AND {name}='{metric.name}')"
+                for metric in metrics
+            ]
+        )
+        filter_query = f"(endpoint_id='{endpoint_id}') AND ({metrics_condition})"
+
+        df = self._get_records(
+            table=table,
+            start=start,
+            end=end,
+            filter_query=filter_query,
+            timestamp_column=timestamp_column,
+            columns=columns,
+        )
+
+        df[mm_schemas.WriterEvent.END_INFER_TIME] = pd.to_datetime(
+            df[mm_schemas.WriterEvent.END_INFER_TIME]
+        )
+        df.set_index(mm_schemas.WriterEvent.END_INFER_TIME, inplace=True)
+
+        logger.debug(
+            "Converting a DataFrame to a list of metrics or results values",
+            table=table,
+            project=self.project,
+            endpoint_id=endpoint_id,
+            is_empty=df.empty,
+        )
+
+        return df_handler(df=df, metrics=metrics, project=self.project)
+
+    def read_predictions(
+        self,
+        *,
+        endpoint_id: str,
+        start: datetime,
+        end: datetime,
+        aggregation_window: typing.Optional[str] = None,
+        agg_funcs: typing.Optional[list] = None,
+        limit: typing.Optional[int] = None,
+    ) -> typing.Union[
+        mm_schemas.ModelEndpointMonitoringMetricValues,
+        mm_schemas.ModelEndpointMonitoringMetricNoData,
+    ]:
+        if (agg_funcs and not aggregation_window) or (
+            aggregation_window and not agg_funcs
+        ):
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "both or neither of `aggregation_window` and `agg_funcs` must be provided"
+            )
+        df = self._get_records(
+            table=self.tables[mm_schemas.TDEngineSuperTables.PREDICTIONS].super_table,
+            start=start,
+            end=end,
+            columns=[mm_schemas.EventFieldType.LATENCY],
+            filter_query=f"endpoint_id='{endpoint_id}'",
+            agg_funcs=agg_funcs,
+            interval=aggregation_window,
+            limit=limit,
+        )
+
+        full_name = get_invocations_fqn(self.project)
+
+        if df.empty:
+            return mm_schemas.ModelEndpointMonitoringMetricNoData(
+                full_name=full_name,
+                type=mm_schemas.ModelEndpointMonitoringMetricType.METRIC,
+            )
+
+        if aggregation_window:
+            # _wend column, which represents the end time of each window, will be used as the time index
+            df["_wend"] = pd.to_datetime(df["_wend"])
+            df.set_index("_wend", inplace=True)
+
+        latency_column = (
+            f"{agg_funcs[0]}({mm_schemas.EventFieldType.LATENCY})"
+            if agg_funcs
+            else mm_schemas.EventFieldType.LATENCY
+        )
+
+        return mm_schemas.ModelEndpointMonitoringMetricValues(
+            full_name=full_name,
+            values=list(
+                zip(
+                    df.index,
+                    df[latency_column],
+                )
+            ),  # pyright: ignore[reportArgumentType]
+        )
+
+    def get_last_request(
+        self,
+        endpoint_ids: Union[str, list[str]],
+        start: Union[datetime, str] = "0",
+        end: Union[datetime, str] = "now",
+    ) -> pd.DataFrame:
+        pass
+
+    def get_drift_status(
+        self,
+        endpoint_ids: Union[str, list[str]],
+        start: Union[datetime, str] = "now-24h",
+        end: Union[datetime, str] = "now",
+    ) -> pd.DataFrame:
+        pass
+
+    def get_metrics_metadata(
+        self,
+        endpoint_id: str,
+        start: Union[datetime, str] = "0",
+        end: Union[datetime, str] = "now",
+    ) -> pd.DataFrame:
+        pass
+
+    def get_results_metadata(
+        self,
+        endpoint_id: str,
+        start: Union[datetime, str] = "0",
+        end: Union[datetime, str] = "now",
+    ) -> pd.DataFrame:
+        pass
+
+    def get_error_count(
+        self,
+        endpoint_ids: Union[str, list[str]],
+        start: Union[datetime, str] = "0",
+        end: Union[datetime, str] = "now",
+    ) -> pd.DataFrame:
+        pass
+
+    def get_avg_latency(
+        self,
+        endpoint_ids: Union[str, list[str]],
+        start: Union[datetime, str] = "0",
+        end: Union[datetime, str] = "now",
+    ) -> pd.DataFrame:
+        pass
+
+    # Note: this function serves as a reference for checking the TSDB for the existence of a metric.
+    #
+    # def read_prediction_metric_for_endpoint_if_exists(
+    #     self, endpoint_id: str
+    # ) -> typing.Optional[mm_schemas.ModelEndpointMonitoringMetric]:
+    #     """
+    #     Read the "invocations" metric for the provided model endpoint, and return the metric object
+    #     if it exists.
+    #
+    #     :param endpoint_id: The model endpoint identifier.
+    #     :return:            `None` if the invocations metric does not exist, otherwise return the
+    #                         corresponding metric object.
+    #     """
+    #     # Read just one record, because we just want to check if there is any data for this endpoint_id
+    #     predictions = self.read_predictions(
+    #         endpoint_id=endpoint_id,
+    #         start=datetime.min,
+    #         end=mlrun.utils.now_date(),
+    #         limit=1,
+    #     )
+    #     if predictions:
+    #         return get_invocations_metric(self.project)

mlrun/model_monitoring/db/tsdb/v3io/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2024 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .v3io_connector import V3IOTSDBConnector

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

@@ -0,0 +1,158 @@
+# Copyright 2024 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from datetime import datetime
+from typing import Any
+
+import mlrun.feature_store.steps
+from mlrun.common.schemas.model_monitoring import (
+    EventFieldType,
+    EventKeyMetrics,
+    EventLiveStats,
+)
+from mlrun.utils import logger
+
+
+def _normalize_dict_for_v3io_frames(event: dict[str, Any]) -> dict[str, Any]:
+    """
+    Normalize user defined keys - input data to a model and its predictions,
+    to a form V3IO frames tolerates.
+
+    The dictionary keys should conform to '^[a-zA-Z_:]([a-zA-Z0-9_:])*$'.
+    """
+    prefix = "_"
+
+    def norm_key(key: str) -> str:
+        key = key.replace("-", "_")  # hyphens `-` are not allowed
+        if key and key[0].isdigit():  # starting with a digit is not allowed
+            return prefix + key
+        return key
+
+    return {norm_key(k): v for k, v in event.items()}
+
+
+class ProcessBeforeTSDB(mlrun.feature_store.steps.MapClass):
+    def __init__(self, **kwargs):
+        """
+        Process the data before writing to TSDB. This step creates a dictionary that includes 3 different dictionaries
+        that each one of them contains important details and stats about the events:
+        1. base_metrics: stats about the average latency and the amount of predictions over time. It is based on
+           storey.AggregateByKey which was executed in step 5.
+        2. endpoint_features: feature names and values along with the prediction names and value.
+        3. custom_metric (opt): optional metrics provided by the user.
+        :returns: Dictionary of 2-3 dictionaries that contains stats and details about the events.
+        """
+        super().__init__(**kwargs)
+
+    def do(self, event):
+        # Compute prediction per second
+        event[EventLiveStats.PREDICTIONS_PER_SECOND] = (
+            float(event[EventLiveStats.PREDICTIONS_COUNT_5M]) / 300
+        )
+        base_fields = [
+            EventFieldType.TIMESTAMP,
+            EventFieldType.ENDPOINT_ID,
+            EventFieldType.ENDPOINT_TYPE,
+        ]
+
+        # Getting event timestamp and endpoint_id
+        base_event = {k: event[k] for k in base_fields}
+
+        # base_metrics includes the stats about the average latency and the amount of predictions over time
+        base_metrics = {
+            EventFieldType.RECORD_TYPE: EventKeyMetrics.BASE_METRICS,
+            EventLiveStats.PREDICTIONS_PER_SECOND: event[
+                EventLiveStats.PREDICTIONS_PER_SECOND
+            ],
+            EventLiveStats.PREDICTIONS_COUNT_5M: event[
+                EventLiveStats.PREDICTIONS_COUNT_5M
+            ],
+            EventLiveStats.PREDICTIONS_COUNT_1H: event[
+                EventLiveStats.PREDICTIONS_COUNT_1H
+            ],
+            EventLiveStats.LATENCY_AVG_5M: event[EventLiveStats.LATENCY_AVG_5M],
+            EventLiveStats.LATENCY_AVG_1H: event[EventLiveStats.LATENCY_AVG_1H],
+            **base_event,
+        }
+
+        # endpoint_features includes the event values of each feature and prediction
+        endpoint_features = {
+            EventFieldType.RECORD_TYPE: EventKeyMetrics.ENDPOINT_FEATURES,
+            **_normalize_dict_for_v3io_frames(event[EventFieldType.NAMED_PREDICTIONS]),
+            **_normalize_dict_for_v3io_frames(event[EventFieldType.NAMED_FEATURES]),
+            **base_event,
+        }
+        # Create a dictionary that includes both base_metrics and endpoint_features
+        processed = {
+            EventKeyMetrics.BASE_METRICS: base_metrics,
+            EventKeyMetrics.ENDPOINT_FEATURES: endpoint_features,
+        }
+
+        # If metrics provided, add another dictionary if custom_metrics values
+        if event[EventFieldType.METRICS]:
+            processed[EventKeyMetrics.CUSTOM_METRICS] = {
+                EventFieldType.RECORD_TYPE: EventKeyMetrics.CUSTOM_METRICS,
+                **event[EventFieldType.METRICS],
+                **base_event,
+            }
+
+        return processed
+
+
+class FilterAndUnpackKeys(mlrun.feature_store.steps.MapClass):
+    def __init__(self, keys, **kwargs):
+        """
+        Create unpacked event dictionary based on provided key metrics (base_metrics, endpoint_features,
+        or custom_metric). Please note that the next step of the TSDB target requires an unpacked dictionary.
+        :param keys: list of key metrics.
+        :returns: An unpacked dictionary of event filtered by the provided key metrics.
+        """
+        super().__init__(**kwargs)
+        self.keys = keys
+
+    def do(self, event):
+        # Keep only the relevant dictionary based on the provided keys
+        new_event = {}
+        for key in self.keys:
+            if key in event:
+                new_event[key] = event[key]
+
+        # Create unpacked dictionary
+        unpacked = {}
+        for key in new_event.keys():
+            if key in self.keys:
+                unpacked = {**unpacked, **new_event[key]}
+            else:
+                unpacked[key] = new_event[key]
+        return unpacked if unpacked else None
+
+
+class ErrorExtractor(mlrun.feature_store.steps.MapClass):
+    def __init__(self, **kwargs):
+        """
+        Prepare the event for insertion into the errors TSDB table.
+        """
+        super().__init__(**kwargs)
+
+    def do(self, event):
+        error = event.get("error")
+        timestamp = datetime.fromisoformat(event.get("when"))
+        endpoint_id = event[EventFieldType.ENDPOINT_ID]
+        event = {
+            EventFieldType.MODEL_ERROR: str(error),
+            EventFieldType.ENDPOINT_ID: endpoint_id,
+            EventFieldType.TIMESTAMP: timestamp,
+            EventFieldType.ERROR_COUNT: 1.0,
+        }
+        logger.info("Write error to errors TSDB table", event=event)
+        return event