mlrun 1.7.0rc37__py3-none-any.whl → 1.7.0rc39__py3-none-any.whl

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

@@ -17,6 +17,8 @@ from dataclasses import dataclass
 from io import StringIO
 from typing import Optional, Union
 
+import taosws
+
 import mlrun.common.schemas.model_monitoring as mm_schemas
 import mlrun.common.types
 
@@ -28,6 +30,9 @@ class _TDEngineColumnType:
         self.data_type = data_type
         self.length = length
 
+    def values_to_column(self, values):
+        raise NotImplementedError()
+
     def __str__(self):
         if self.length is not None:
             return f"{self.data_type}({self.length})"
@@ -44,6 +49,26 @@ class _TDEngineColumn(mlrun.common.types.StrEnum):
     BINARY_10000 = _TDEngineColumnType("BINARY", 10000)
 
 
+def values_to_column(values, column_type):
+    if column_type == _TDEngineColumn.TIMESTAMP:
+        timestamps = [round(timestamp.timestamp() * 1000) for timestamp in values]
+        return taosws.millis_timestamps_to_column(timestamps)
+    if column_type == _TDEngineColumn.FLOAT:
+        return taosws.floats_to_column(values)
+    if column_type == _TDEngineColumn.INT:
+        return taosws.ints_to_column(values)
+    if column_type == _TDEngineColumn.BINARY_40:
+        return taosws.binary_to_column(values)
+    if column_type == _TDEngineColumn.BINARY_64:
+        return taosws.binary_to_column(values)
+    if column_type == _TDEngineColumn.BINARY_10000:
+        return taosws.binary_to_column(values)
+
+    raise mlrun.errors.MLRunInvalidArgumentError(
+        f"unsupported column type '{column_type}'"
+    )
+
+
 @dataclass
 class TDEngineSchema:
     """
@@ -55,13 +80,14 @@ class TDEngineSchema:
     def __init__(
         self,
         super_table: str,
-        columns: dict[str, str],
+        columns: dict[str, _TDEngineColumn],
         tags: dict[str, str],
+        database: Optional[str] = None,
     ):
         self.super_table = super_table
         self.columns = columns
         self.tags = tags
-        self.database = _MODEL_MONITORING_DATABASE
+        self.database = database or _MODEL_MONITORING_DATABASE
 
     def _create_super_table_query(self) -> str:
         columns = ", ".join(f"{col} {val}" for col, val in self.columns.items())
@@ -83,11 +109,23 @@ class TDEngineSchema:
 
     def _insert_subtable_query(
         self,
+        connection: taosws.Connection,
         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});"
+    ) -> taosws.TaosStmt:
+        stmt = connection.statement()
+        question_marks = ", ".join("?" * len(self.columns))
+        stmt.prepare(f"INSERT INTO ? VALUES ({question_marks});")
+        stmt.set_tbname_tags(subtable, [])
+
+        bind_params = []
+
+        for col_name, col_type in self.columns.items():
+            val = values[col_name]
+            bind_params.append(values_to_column([val], col_type))
+
+        stmt.bind_param(bind_params)
+        return stmt
 
     def _delete_subtable_query(
         self,
@@ -188,53 +226,53 @@ class TDEngineSchema:
 
 @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,
-        mm_schemas.ResultData.RESULT_KIND: _TDEngineColumn.INT,
-    }
-    database = _MODEL_MONITORING_DATABASE
+    def __init__(self, database: Optional[str] = None):
+        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,
+        }
+        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,
+            mm_schemas.ResultData.RESULT_KIND: _TDEngineColumn.INT,
+        }
+        super().__init__(super_table, columns, tags, 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
+    def __init__(self, database: Optional[str] = None):
+        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,
+        }
+        super().__init__(super_table, columns, tags, 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
+    def __init__(self, database: Optional[str] = None):
+        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,
+        }
+        super().__init__(super_table, columns, tags, database)

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

@@ -14,6 +14,7 @@
 
 import typing
 from datetime import datetime
+from typing import Union
 
 import pandas as pd
 import taosws
@@ -57,15 +58,26 @@ class TDEngineConnector(TSDBConnector):
         except taosws.QueryError:
             # Database already exists
             pass
-        conn.execute(f"USE {self.database}")
+        try:
+            conn.execute(f"USE {self.database}")
+        except taosws.QueryError as e:
+            raise mlrun.errors.MLRunTSDBConnectionFailure(
+                f"Failed to use TDEngine database {self.database}, {mlrun.errors.err_to_str(e)}"
+            )
         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(),
+            mm_schemas.TDEngineSuperTables.APP_RESULTS: tdengine_schemas.AppResultTable(
+                self.database
+            ),
+            mm_schemas.TDEngineSuperTables.METRICS: tdengine_schemas.Metrics(
+                self.database
+            ),
+            mm_schemas.TDEngineSuperTables.PREDICTIONS: tdengine_schemas.Predictions(
+                self.database
+            ),
         }
 
     def create_tables(self):
@@ -96,6 +108,7 @@ class TDEngineConnector(TSDBConnector):
             table_name = (
                 f"{table_name}_" f"{event[mm_schemas.ResultData.RESULT_NAME]}"
             ).replace("-", "_")
+            event.pop(mm_schemas.ResultData.CURRENT_STATS, None)
 
         else:
             # Write a new metric
@@ -104,14 +117,30 @@ class TDEngineConnector(TSDBConnector):
                 f"{table_name}_" f"{event[mm_schemas.MetricData.METRIC_NAME]}"
             ).replace("-", "_")
 
+        # 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_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
+
+        insert_statement = table._insert_subtable_query(
+            self._connection,
+            subtable=table_name,
+            values=event,
         )
-        self._connection.execute(insert_table_query)
+        insert_statement.add_batch()
+        insert_statement.execute()
+
+    @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):
         """
@@ -156,6 +185,9 @@ class TDEngineConnector(TSDBConnector):
             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.
@@ -246,11 +278,9 @@ class TDEngineConnector(TSDBConnector):
             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)
+        df_columns = [field.name() for field in query_result.fields]
+        return pd.DataFrame(query_result, columns=df_columns)
 
     def read_metrics_data(
         self,
@@ -274,13 +304,22 @@ class TDEngineConnector(TSDBConnector):
             ],
         ],
     ]:
+        timestamp_column = mm_schemas.WriterEvent.END_INFER_TIME
+        columns = [timestamp_column, mm_schemas.WriterEvent.APPLICATION_NAME]
         if type == "metrics":
             table = mm_schemas.TDEngineSuperTables.METRICS
             name = mm_schemas.MetricData.METRIC_NAME
+            columns += [name, mm_schemas.MetricData.METRIC_VALUE]
             df_handler = self.df_to_metrics_values
         elif type == "results":
             table = mm_schemas.TDEngineSuperTables.APP_RESULTS
             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(
@@ -300,7 +339,8 @@ class TDEngineConnector(TSDBConnector):
             start=start,
             end=end,
             filter_query=filter_query,
-            timestamp_column=mm_schemas.WriterEvent.END_INFER_TIME,
+            timestamp_column=timestamp_column,
+            columns=columns,
         )
 
         df[mm_schemas.WriterEvent.END_INFER_TIME] = pd.to_datetime(
@@ -377,6 +417,54 @@ class TDEngineConnector(TSDBConnector):
             ),  # 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(

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

@@ -11,7 +11,7 @@
 # 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
@@ -20,6 +20,7 @@ from mlrun.common.schemas.model_monitoring import (
     EventKeyMetrics,
     EventLiveStats,
 )
+from mlrun.utils import logger
 
 
 def _normalize_dict_for_v3io_frames(event: dict[str, Any]) -> dict[str, Any]:
@@ -134,3 +135,24 @@ class FilterAndUnpackKeys(mlrun.feature_store.steps.MapClass):
             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