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

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

@@ -0,0 +1,207 @@
+# 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 datetime
+from dataclasses import dataclass
+from io import StringIO
+from typing import Union
+
+import mlrun.common.schemas.model_monitoring as mm_schemas
+import mlrun.common.types
+
+_MODEL_MONITORING_DATABASE = "mlrun_model_monitoring"
+
+
+class _TDEngineColumnType:
+    def __init__(self, data_type: str, length: int = None):
+        self.data_type = data_type
+        self.length = length
+
+    def __str__(self):
+        if self.length is not None:
+            return f"{self.data_type}({self.length})"
+        else:
+            return self.data_type
+
+
+class _TDEngineColumn(mlrun.common.types.StrEnum):
+    TIMESTAMP = _TDEngineColumnType("TIMESTAMP")
+    FLOAT = _TDEngineColumnType("FLOAT")
+    INT = _TDEngineColumnType("INT")
+    BINARY_40 = _TDEngineColumnType("BINARY", 40)
+    BINARY_64 = _TDEngineColumnType("BINARY", 64)
+    BINARY_10000 = _TDEngineColumnType("BINARY", 10000)
+
+
+@dataclass
+class TDEngineSchema:
+    """
+    A class to represent a supertable schema in TDengine. Using this schema, you can generate the relevant queries to
+    create, insert, delete and query data from TDengine. At the moment, there are 3 schemas: AppResultTable,
+    Metrics, and Predictions.
+    """
+
+    def __init__(
+        self,
+        super_table: str,
+        columns: dict[str, str],
+        tags: dict[str, str],
+    ):
+        self.super_table = super_table
+        self.columns = columns
+        self.tags = tags
+        self.database = _MODEL_MONITORING_DATABASE
+
+    def _create_super_table_query(self) -> str:
+        columns = ", ".join(f"{col} {val}" for col, val in self.columns.items())
+        tags = ", ".join(f"{col} {val}" for col, val in self.tags.items())
+        return f"CREATE STABLE if NOT EXISTS {self.database}.{self.super_table} ({columns}) TAGS ({tags});"
+
+    def _create_subtable_query(
+        self,
+        subtable: str,
+        values: dict[str, Union[str, int, float, datetime.datetime]],
+    ) -> str:
+        try:
+            values = ", ".join(f"'{values[val]}'" for val in self.tags)
+        except KeyError:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"values must contain all tags: {self.tags.keys()}"
+            )
+        return f"CREATE TABLE if NOT EXISTS {self.database}.{subtable} USING {self.super_table} TAGS ({values});"
+
+    def _insert_subtable_query(
+        self,
+        subtable: str,
+        values: dict[str, Union[str, int, float, datetime.datetime]],
+    ) -> str:
+        values = ", ".join(f"'{values[val]}'" for val in self.columns)
+        return f"INSERT INTO {self.database}.{subtable} VALUES ({values});"
+
+    def _delete_subtable_query(
+        self,
+        subtable: str,
+        values: dict[str, Union[str, int, float, datetime.datetime]],
+    ) -> str:
+        values = " AND ".join(
+            f"{val} like '{values[val]}'" for val in self.tags if val in values
+        )
+        if not values:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"values must contain at least one tag: {self.tags.keys()}"
+            )
+        return f"DELETE FROM {self.database}.{subtable} WHERE {values};"
+
+    def _drop_subtable_query(
+        self,
+        subtable: str,
+    ) -> str:
+        return f"DROP TABLE if EXISTS {self.database}.{subtable};"
+
+    def _get_subtables_query(
+        self,
+        values: dict[str, Union[str, int, float, datetime.datetime]],
+    ) -> str:
+        values = " AND ".join(
+            f"{val} like '{values[val]}'" for val in self.tags if val in values
+        )
+        if not values:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"values must contain at least one tag: {self.tags.keys()}"
+            )
+        return f"SELECT tbname FROM {self.database}.{self.super_table} WHERE {values};"
+
+    @staticmethod
+    def _get_records_query(
+        table: str,
+        start: str,
+        end: str,
+        columns_to_filter: list[str] = None,
+        filter_query: str = "",
+        timestamp_column: str = "time",
+        database: str = _MODEL_MONITORING_DATABASE,
+    ) -> str:
+        with StringIO() as query:
+            query.write("SELECT ")
+            if columns_to_filter:
+                query.write(", ".join(columns_to_filter))
+            else:
+                query.write("*")
+            query.write(f" from {database}.{table}")
+
+            if any([filter_query, start, end]):
+                query.write(" where ")
+                if filter_query:
+                    query.write(f"{filter_query} and ")
+                if start:
+                    query.write(f"{timestamp_column} >= '{start}'" + " and ")
+                if end:
+                    query.write(f"{timestamp_column} <= '{end}'")
+                full_query = query.getvalue()
+                if full_query.endswith(" and "):
+                    full_query = full_query[:-5]
+            return full_query + ";"
+
+
+@dataclass
+class AppResultTable(TDEngineSchema):
+    super_table = mm_schemas.TDEngineSuperTables.APP_RESULTS
+    columns = {
+        mm_schemas.WriterEvent.END_INFER_TIME: _TDEngineColumn.TIMESTAMP,
+        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.CURRENT_STATS: _TDEngineColumn.BINARY_10000,
+    }
+
+    tags = {
+        mm_schemas.EventFieldType.PROJECT: _TDEngineColumn.BINARY_64,
+        mm_schemas.WriterEvent.ENDPOINT_ID: _TDEngineColumn.BINARY_64,
+        mm_schemas.WriterEvent.APPLICATION_NAME: _TDEngineColumn.BINARY_64,
+        mm_schemas.ResultData.RESULT_NAME: _TDEngineColumn.BINARY_64,
+    }
+    database = _MODEL_MONITORING_DATABASE
+
+
+@dataclass
+class Metrics(TDEngineSchema):
+    super_table = mm_schemas.TDEngineSuperTables.METRICS
+    columns = {
+        mm_schemas.WriterEvent.END_INFER_TIME: _TDEngineColumn.TIMESTAMP,
+        mm_schemas.WriterEvent.START_INFER_TIME: _TDEngineColumn.TIMESTAMP,
+        mm_schemas.MetricData.METRIC_VALUE: _TDEngineColumn.FLOAT,
+    }
+
+    tags = {
+        mm_schemas.EventFieldType.PROJECT: _TDEngineColumn.BINARY_64,
+        mm_schemas.WriterEvent.ENDPOINT_ID: _TDEngineColumn.BINARY_64,
+        mm_schemas.WriterEvent.APPLICATION_NAME: _TDEngineColumn.BINARY_64,
+        mm_schemas.MetricData.METRIC_NAME: _TDEngineColumn.BINARY_64,
+    }
+    database = _MODEL_MONITORING_DATABASE
+
+
+@dataclass
+class Predictions(TDEngineSchema):
+    super_table = mm_schemas.TDEngineSuperTables.PREDICTIONS
+    columns = {
+        mm_schemas.EventFieldType.TIME: _TDEngineColumn.TIMESTAMP,
+        mm_schemas.EventFieldType.LATENCY: _TDEngineColumn.FLOAT,
+        mm_schemas.EventKeyMetrics.CUSTOM_METRICS: _TDEngineColumn.BINARY_10000,
+    }
+    tags = {
+        mm_schemas.EventFieldType.PROJECT: _TDEngineColumn.BINARY_64,
+        mm_schemas.WriterEvent.ENDPOINT_ID: _TDEngineColumn.BINARY_64,
+    }
+    database = _MODEL_MONITORING_DATABASE

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

@@ -0,0 +1,45 @@
+# 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 json
+
+import mlrun.feature_store.steps
+from mlrun.common.schemas.model_monitoring import (
+    EventFieldType,
+    EventKeyMetrics,
+)
+
+_TABLE_COLUMN = "table_column"
+
+
+class ProcessBeforeTDEngine(mlrun.feature_store.steps.MapClass):
+    def __init__(self, **kwargs):
+        """
+        Process the data before writing to TDEngine. This step create the relevant keys for the TDEngine table,
+        including project name, custom metrics, time column, and table name column.
+
+        :returns: Event as a dictionary which will be written into the TDEngine Predictions table.
+        """
+        super().__init__(**kwargs)
+
+    def do(self, event):
+        event[EventFieldType.PROJECT] = event[EventFieldType.FUNCTION_URI].split("/")[0]
+        event[EventKeyMetrics.CUSTOM_METRICS] = json.dumps(
+            event.get(EventFieldType.METRICS, {})
+        )
+        event[EventFieldType.TIME] = event.get(EventFieldType.TIMESTAMP)
+        event[EventFieldType.TABLE_COLUMN] = "_" + event.get(EventFieldType.ENDPOINT_ID)
+
+        return event

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

@@ -0,0 +1,231 @@
+# 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
+
+import pandas as pd
+import taosws
+
+import mlrun.common.schemas.model_monitoring as mm_schemas
+import mlrun.model_monitoring.db
+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.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 = self._create_connection()
+        self._init_super_tables()
+
+    def _create_connection(self):
+        """Establish a connection to the TSDB server."""
+        conn = taosws.connect(self._tdengine_connection_string)
+        try:
+            conn.execute(f"CREATE DATABASE {self.database}")
+        except taosws.QueryError:
+            # Database already exists
+            pass
+        conn.execute(f"USE {self.database}")
+        return conn
+
+    def _init_super_tables(self):
+        """Initialize the super tables for the TSDB."""
+        self.tables = {
+            mm_schemas.TDEngineSuperTables.APP_RESULTS: tdengine_schemas.AppResultTable(),
+            mm_schemas.TDEngineSuperTables.METRICS: tdengine_schemas.Metrics(),
+            mm_schemas.TDEngineSuperTables.PREDICTIONS: tdengine_schemas.Predictions(),
+        }
+
+    def create_tables(self):
+        """Create TDEngine supertables."""
+        for table in self.tables:
+            create_table_query = self.tables[table]._create_super_table_query()
+            self._connection.execute(create_table_query)
+
+    def write_application_event(
+        self,
+        event: dict,
+        kind: mm_schemas.WriterEventKind = mm_schemas.WriterEventKind.RESULT,
+    ):
+        """
+        Write a single result or metric to TSDB.
+        """
+
+        table_name = (
+            f"{self.project}_"
+            f"{event[mm_schemas.WriterEvent.ENDPOINT_ID]}_"
+            f"{event[mm_schemas.WriterEvent.APPLICATION_NAME]}_"
+        )
+        event[mm_schemas.EventFieldType.PROJECT] = self.project
+
+        if kind == mm_schemas.WriterEventKind.RESULT:
+            # Write a new result
+            table = self.tables[mm_schemas.TDEngineSuperTables.APP_RESULTS]
+            table_name = (
+                f"{table_name}_" f"{event[mm_schemas.ResultData.RESULT_NAME]}"
+            ).replace("-", "_")
+
+        else:
+            # Write a new metric
+            table = self.tables[mm_schemas.TDEngineSuperTables.METRICS]
+            table_name = (
+                f"{table_name}_" f"{event[mm_schemas.MetricData.METRIC_NAME]}"
+            ).replace("-", "_")
+
+        create_table_query = table._create_subtable_query(
+            subtable=table_name, values=event
+        )
+        self._connection.execute(create_table_query)
+        insert_table_query = table._insert_subtable_query(
+            subtable=table_name, values=event
+        )
+        self._connection.execute(insert_table_query)
+
+    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=mm_schemas.TDEngineSuperTables.PREDICTIONS,
+                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=10,
+            )
+
+        apply_process_before_tsdb()
+        apply_tdengine_target(
+            name="TDEngineTarget",
+            after="ProcessBeforeTDEngine",
+        )
+
+    def delete_tsdb_resources(self):
+        """
+        Delete all project resources in the TSDB connector, such as model endpoints data and drift results.
+        """
+        for table in self.tables:
+            get_subtable_names_query = self.tables[table]._get_subtables_query(
+                values={mm_schemas.EventFieldType.PROJECT: self.project}
+            )
+            subtables = self._connection.query(get_subtable_names_query)
+            for subtable in subtables:
+                drop_query = self.tables[table]._drop_subtable_query(
+                    subtable=subtable[0]
+                )
+                self._connection.execute(drop_query)
+        logger.info(
+            f"Deleted all project resources in the TSDB connector for 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: str,
+        end: str,
+        columns: typing.Optional[list[str]] = None,
+        filter_query: str = "",
+        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 columns:          Columns to include in the result.
+        :param filter_query:     Optional filter expression as a string. The filter structure depends on the TSDB
+                                 connector type.
+        :param start:            The start time of the metrics.
+        :param end:              The end time of the metrics.
+        :param timestamp_column: The column name that holds the timestamp.
+
+        :return: DataFrame with the provided attributes from the data collection.
+        :raise:  MLRunInvalidArgumentError if query the provided table failed.
+        """
+
+        filter_query += f" project = '{self.project}'"
+
+        full_query = tdengine_schemas.TDEngineSchema._get_records_query(
+            table=table,
+            columns_to_filter=columns,
+            filter_query=filter_query,
+            start=start,
+            end=end,
+            timestamp_column=timestamp_column,
+            database=self.database,
+        )
+        try:
+            query_result = self._connection.query(full_query)
+        except taosws.QueryError as e:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"Failed to query table {table} in database {self.database}, {str(e)}"
+            )
+        columns = []
+        for column in query_result.fields:
+            columns.append(column.name())
+
+        return pd.DataFrame(query_result, columns=columns)