mlrun 1.9.0rc3__py3-none-any.whl → 1.10.0rc1__py3-none-any.whl

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

@@ -0,0 +1,213 @@
+# Copyright 2025 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 traceback
+from collections.abc import Callable
+from enum import Enum
+from typing import Any, Final, Optional, Union
+
+import taosws
+from taosws import TaosStmt
+
+
+class _StrEnum(str, Enum):
+    pass
+
+
+class TimestampPrecision(_StrEnum):
+    ms = "ms"  # milliseconds
+    us = "us"  # microseconds
+    ns = "ns"  # nanoseconds
+
+
+_TS_PRECISION_TO_FACTOR_AND_FUNC: Final[
+    dict[TimestampPrecision, tuple[int, Callable[[list[int]], taosws.PyColumnView]]]
+] = {
+    TimestampPrecision.ms: (10**3, taosws.millis_timestamps_to_column),
+    TimestampPrecision.us: (10**6, taosws.micros_timestamps_to_column),
+    TimestampPrecision.ns: (10**9, taosws.nanos_timestamps_to_column),
+}
+
+
+class QueryResult:
+    def __init__(self, data, fields):
+        self.data = data
+        self.fields = fields
+
+    def __eq__(self, other):
+        return self.data == other.data and self.fields == other.fields
+
+    def __repr__(self):
+        return f"QueryResult({self.data}, {self.fields})"
+
+
+class Field:
+    def __init__(self, name, type, bytes):
+        self.name = name
+        self.type = type
+        self.bytes = bytes
+
+    def __eq__(self, other):
+        return (
+            self.name == other.name
+            and self.type == other.type
+            and self.bytes == other.bytes
+        )
+
+    def __repr__(self):
+        return f"Field({self.name}, {self.type}, {self.bytes})"
+
+
+class TDEngineError(Exception):
+    pass
+
+
+class ErrorResult:
+    def __init__(self, tb, err):
+        self.tb = tb
+        self.err = err
+
+
+def _get_timestamp_column(
+    values: list, timestamp_precision: TimestampPrecision
+) -> taosws.PyColumnView:
+    factor, to_col_func = _TS_PRECISION_TO_FACTOR_AND_FUNC[timestamp_precision]
+    timestamps = [round(timestamp.timestamp() * factor) for timestamp in values]
+    return to_col_func(timestamps)
+
+
+def values_to_column(
+    values: list,
+    column_type: str,
+    timestamp_precision: TimestampPrecision = TimestampPrecision.ms,
+) -> taosws.PyColumnView:
+    if column_type == "TIMESTAMP":
+        return _get_timestamp_column(values, timestamp_precision)
+    if column_type == "FLOAT":
+        return taosws.floats_to_column(values)
+    if column_type == "INT":
+        return taosws.ints_to_column(values)
+    if column_type.startswith("BINARY"):
+        return taosws.binary_to_column(values)
+
+    raise NotImplementedError(f"Unsupported column type '{column_type}'")
+
+
+class Statement:
+    def __init__(
+        self,
+        columns: dict[str, str],
+        subtable: str,
+        values: dict[str, Any],
+        timestamp_precision: str = TimestampPrecision.ms,
+    ) -> None:
+        self.columns = columns
+        self.subtable = subtable
+        self.values = values
+        self.timestamp_precision = TimestampPrecision[timestamp_precision]
+
+    def prepare(self, statement: TaosStmt) -> TaosStmt:
+        question_marks = ", ".join("?" * len(self.columns))
+        statement.prepare(f"INSERT INTO ? VALUES ({question_marks});")
+        statement.set_tbname(self.subtable)
+
+        bind_params = []
+
+        for col_name, col_type in self.columns.items():
+            val = self.values[col_name]
+            bind_params.append(
+                values_to_column(
+                    [val], col_type, timestamp_precision=self.timestamp_precision
+                )
+            )
+
+        statement.bind_param(bind_params)
+        statement.add_batch()
+        return statement
+
+
+def _run(connection_string, prefix_statements, q, statements, query):
+    try:
+        conn = taosws.connect(connection_string)
+
+        for statement in prefix_statements + statements:
+            if isinstance(statement, Statement):
+                prepared_statement = statement.prepare(conn.statement())
+                prepared_statement.execute()
+            else:
+                conn.execute(statement)
+
+        if not query:
+            q.put(None)
+            return
+
+        res = conn.query(query)
+
+        # taosws.TaosField is not serializable
+        fields = [
+            Field(field.name(), field.type(), field.bytes()) for field in res.fields
+        ]
+
+        q.put(QueryResult(list(res), fields))
+    except Exception as e:
+        tb = traceback.format_exc()
+        q.put(ErrorResult(tb, e))
+
+
+class TDEngineConnection:
+    def __init__(self, connection_string):
+        self._connection_string = connection_string
+        self.prefix_statements = []
+
+        self._conn = taosws.connect(self._connection_string)
+
+    def run(
+        self,
+        statements: Optional[Union[str, Statement, list[Union[str, Statement]]]] = None,
+        query: Optional[str] = None,
+    ) -> Optional[QueryResult]:
+        statements = statements or []
+        if not isinstance(statements, list):
+            statements = [statements]
+
+        for statement in self.prefix_statements + statements:
+            if isinstance(statement, Statement):
+                try:
+                    prepared_statement = statement.prepare(self._conn.statement())
+                    prepared_statement.execute()
+                except taosws.Error as e:
+                    raise TDEngineError(
+                        f"Failed to run prepared statement `{self._conn.statement()}`: {e}"
+                    ) from e
+            else:
+                try:
+                    self._conn.execute(statement)
+                except taosws.Error as e:
+                    raise TDEngineError(
+                        f"Failed to run statement `{statement}`: {e}"
+                    ) from e
+
+        if not query:
+            return None
+
+        try:
+            res = self._conn.query(query)
+        except taosws.Error as e:
+            raise TDEngineError(f"Failed to run query `{query}`: {e}") from e
+
+        fields = [
+            Field(field.name(), field.type(), field.bytes()) for field in res.fields
+        ]
+
+        return QueryResult(list(res), fields)

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

@@ -12,17 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import asyncio
 from datetime import datetime, timedelta
 from threading import Lock
 from typing import Callable, Final, Literal, Optional, 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.common.types
@@ -30,6 +25,10 @@ import mlrun.model_monitoring.db.tsdb.tdengine.schemas as tdengine_schemas
 import mlrun.model_monitoring.db.tsdb.tdengine.stream_graph_steps
 from mlrun.datastore.datastore_profile import DatastoreProfile
 from mlrun.model_monitoring.db import TSDBConnector
+from mlrun.model_monitoring.db.tsdb.tdengine.tdengine_connection import (
+    Statement,
+    TDEngineConnection,
+)
 from mlrun.model_monitoring.helpers import get_invocations_fqn
 from mlrun.utils import logger
 
@@ -75,12 +74,6 @@ class TDEngineConnector(TSDBConnector):
 
         self._init_super_tables()
 
-        self._run_directly = (
-            mlrun.mlconf.model_endpoint_monitoring.tdengine.run_directly
-        )
-        self._timeout = mlrun.mlconf.model_endpoint_monitoring.tdengine.timeout
-        self._retries = mlrun.mlconf.model_endpoint_monitoring.tdengine.retries
-
     @property
     def connection(self) -> TDEngineConnection:
         global _connection
@@ -98,7 +91,7 @@ class TDEngineConnector(TSDBConnector):
         """Establish a connection to the TSDB server."""
         logger.debug("Creating a new connection to TDEngine", project=self.project)
         conn = TDEngineConnection(
-            self._tdengine_connection_profile.dsn(), run_directly=self._run_directly
+            self._tdengine_connection_profile.dsn(),
         )
         conn.prefix_statements = [f"USE {self.database}"]
 
@@ -126,8 +119,6 @@ class TDEngineConnector(TSDBConnector):
         self.connection.prefix_statements = []
         self.connection.run(
             statements=f"CREATE DATABASE IF NOT EXISTS {self.database} PRECISION '{self._timestamp_precision}'",
-            timeout=self._timeout,
-            retries=self._retries,
         )
         self.connection.prefix_statements = [f"USE {self.database}"]
         logger.debug(
@@ -147,8 +138,6 @@ class TDEngineConnector(TSDBConnector):
             conn = self.connection
             conn.run(
                 statements=create_table_query,
-                timeout=self._timeout,
-                retries=self._retries,
             )
 
     def write_application_event(
@@ -208,8 +197,6 @@ class TDEngineConnector(TSDBConnector):
                 create_table_sql,
                 insert_statement,
             ],
-            timeout=self._timeout,
-            retries=self._retries,
         )
 
     @staticmethod
@@ -313,7 +300,8 @@ class TDEngineConnector(TSDBConnector):
         )
 
     def delete_tsdb_records(
-        self, endpoint_ids: list[str], delete_timeout: Optional[int] = None
+        self,
+        endpoint_ids: list[str],
     ):
         """
         To delete subtables within TDEngine, we first query the subtables names with the provided endpoint_ids.
@@ -334,8 +322,6 @@ class TDEngineConnector(TSDBConnector):
                 )
                 subtables_result = self.connection.run(
                     query=get_subtable_query,
-                    timeout=self._timeout,
-                    retries=self._retries,
                 )
                 subtables.extend([subtable[0] for subtable in subtables_result.data])
         except Exception as e:
@@ -356,8 +342,6 @@ class TDEngineConnector(TSDBConnector):
         try:
             self.connection.run(
                 statements=drop_statements,
-                timeout=delete_timeout or self._timeout,
-                retries=self._retries,
             )
         except Exception as e:
             logger.warning(
@@ -388,8 +372,6 @@ class TDEngineConnector(TSDBConnector):
         try:
             self.connection.run(
                 statements=drop_statements,
-                timeout=self._timeout,
-                retries=self._retries,
             )
         except Exception as e:
             logger.warning(
@@ -413,8 +395,6 @@ class TDEngineConnector(TSDBConnector):
         try:
             table_name = self.connection.run(
                 query=query_random_table_name,
-                timeout=self._timeout,
-                retries=self._retries,
             )
             if len(table_name.data) == 0:
                 # no tables were found under the database
@@ -437,8 +417,6 @@ class TDEngineConnector(TSDBConnector):
             try:
                 self.connection.run(
                     statements=drop_database_query,
-                    timeout=self._timeout,
-                    retries=self._retries,
                 )
                 logger.debug(
                     "The TDEngine database has been successfully dropped",
@@ -531,7 +509,7 @@ class TDEngineConnector(TSDBConnector):
         logger.debug("Querying TDEngine", query=full_query)
         try:
             query_result = self.connection.run(
-                query=full_query, timeout=self._timeout, retries=self._retries
+                query=full_query,
             )
         except taosws.QueryError as e:
             raise mlrun.errors.MLRunInvalidArgumentError(
@@ -908,6 +886,7 @@ class TDEngineConnector(TSDBConnector):
         model_endpoint_objects: list[mlrun.common.schemas.ModelEndpoint],
         project: str,
         run_in_threadpool: Callable,
+        metric_list: Optional[list[str]] = None,
     ) -> list[mlrun.common.schemas.ModelEndpoint]:
         """
         Add basic metrics to the model endpoint object.
@@ -916,24 +895,28 @@ class TDEngineConnector(TSDBConnector):
                                         be filled with the relevant basic metrics.
         :param project:                The name of the project.
         :param run_in_threadpool:      A function that runs another function in a thread pool.
+        :param metric_list:            List of metrics to include from the time series DB. Defaults to all metrics.
 
         :return: A list of `ModelEndpointMonitoringMetric` objects.
         """
 
         uids = [mep.metadata.uid for mep in model_endpoint_objects]
-        coroutines = [
-            run_in_threadpool(self.get_error_count, endpoint_ids=uids),
-            run_in_threadpool(self.get_last_request, endpoint_ids=uids),
-            run_in_threadpool(self.get_avg_latency, endpoint_ids=uids),
-            run_in_threadpool(self.get_drift_status, endpoint_ids=uids),
-        ]
-
-        (
-            error_count_df,
-            last_request_df,
-            avg_latency_df,
-            drift_status_df,
-        ) = await asyncio.gather(*coroutines)
+
+        metric_name_to_function = {
+            "error_count": self.get_error_count,
+            "last_request": self.get_last_request,
+            "avg_latency": self.get_avg_latency,
+            "result_status": self.get_drift_status,
+        }
+        if metric_list is not None:
+            for metric_name in list(metric_name_to_function):
+                if metric_name not in metric_list:
+                    del metric_name_to_function[metric_name]
+
+        metric_name_to_df = {
+            metric_name: function(endpoint_ids=uids)
+            for metric_name, function in metric_name_to_function.items()
+        }
 
         def add_metrics(
             mep: mlrun.common.schemas.ModelEndpoint,
@@ -955,12 +938,7 @@ class TDEngineConnector(TSDBConnector):
             map(
                 lambda mep: add_metrics(
                     mep=mep,
-                    df_dictionary={
-                        "error_count": error_count_df,
-                        "last_request": last_request_df,
-                        "avg_latency": avg_latency_df,
-                        "result_status": drift_status_df,
-                    },
+                    df_dictionary=metric_name_to_df,
                 ),
                 model_endpoint_objects,
             )

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

@@ -470,7 +470,8 @@ class V3IOTSDBConnector(TSDBConnector):
         store.rm(tsdb_path, recursive=True)
 
     def delete_tsdb_records(
-        self, endpoint_ids: list[str], delete_timeout: Optional[int] = None
+        self,
+        endpoint_ids: list[str],
     ):
         logger.debug(
             "Deleting model endpoints resources using the V3IO TSDB connector",
@@ -1085,6 +1086,7 @@ class V3IOTSDBConnector(TSDBConnector):
         model_endpoint_objects: list[mlrun.common.schemas.ModelEndpoint],
         project: str,
         run_in_threadpool: Callable,
+        metric_list: Optional[list[str]] = None,
     ) -> list[mlrun.common.schemas.ModelEndpoint]:
         """
         Fetch basic metrics from V3IO TSDB and add them to MEP objects.
@@ -1093,6 +1095,7 @@ class V3IOTSDBConnector(TSDBConnector):
                                        be filled with the relevant basic metrics.
         :param project:                The name of the project.
         :param run_in_threadpool:      A function that runs another function in a thread pool.
+        :param metric_list:            List of metrics to include from the time series DB. Defaults to all metrics.
 
         :return: A list of `ModelEndpointMonitoringMetric` objects.
         """
@@ -1104,15 +1107,27 @@ class V3IOTSDBConnector(TSDBConnector):
             uids.append(uid)
             model_endpoint_objects_by_uid[uid] = model_endpoint_object
 
-        error_count_res = await run_in_threadpool(
-            self.get_error_count, endpoint_ids=uids, get_raw=True
-        )
-        avg_latency_res = await run_in_threadpool(
-            self.get_avg_latency, endpoint_ids=uids, get_raw=True
-        )
-        drift_status_res = await run_in_threadpool(
-            self.get_drift_status, endpoint_ids=uids, get_raw=True
-        )
+        metric_name_to_function_and_column_name = {
+            "error_count": (self.get_error_count, "count(error_count)"),
+            "avg_latency": (self.get_avg_latency, "avg(latency)"),
+            "result_status": (self.get_drift_status, "max(result_status)"),
+        }
+        if metric_list is not None:
+            for metric_name in list(metric_name_to_function_and_column_name):
+                if metric_name not in metric_list:
+                    del metric_name_to_function_and_column_name[metric_name]
+
+        metric_name_to_result = {}
+
+        for metric_name, (
+            function,
+            _,
+        ) in metric_name_to_function_and_column_name.items():
+            metric_name_to_result[metric_name] = await run_in_threadpool(
+                function,
+                endpoint_ids=uids,
+                get_raw=True,
+            )
 
         def add_metric(
             metric: str,
@@ -1128,26 +1143,16 @@ class V3IOTSDBConnector(TSDBConnector):
                     if mep and value is not None and not math.isnan(value):
                         setattr(mep.status, metric, value)
 
-        add_metric(
-            "error_count",
-            "count(error_count)",
-            error_count_res,
-        )
-
-        add_metric(
-            "avg_latency",
-            "avg(latency)",
-            avg_latency_res,
-        )
-        add_metric(
-            "result_status",
-            "max(result_status)",
-            drift_status_res,
-        )
-
-        self._enrich_mep_with_last_request(
-            model_endpoint_objects_by_uid=model_endpoint_objects_by_uid
-        )
+        for metric_name, result in metric_name_to_result.items():
+            add_metric(
+                metric_name,
+                metric_name_to_function_and_column_name[metric_name][1],
+                result,
+            )
+        if metric_list is None or "last_request" in metric_list:
+            self._enrich_mep_with_last_request(
+                model_endpoint_objects_by_uid=model_endpoint_objects_by_uid
+            )
 
         return list(model_endpoint_objects_by_uid.values())
 

mlrun/package/context_handler.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.
-#
+
 import inspect
 import os
 from collections import OrderedDict

mlrun/package/errors.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 mlrun.errors import MLRunBaseError
 
 

mlrun/package/packager.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 abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Any, Optional, Union

mlrun/package/packagers/default_packager.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.
-#
+
 import inspect
 from abc import ABCMeta
 from types import MethodType

mlrun/package/packagers/numpy_packagers.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.
-#
+
 import os
 import pathlib
 import tempfile

mlrun/package/packagers/pandas_packagers.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.
-#
+
 import importlib
 import os
 import pathlib

mlrun/package/packagers/python_standard_library_packagers.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.
-#
+
 import os
 import pathlib
 import tempfile

mlrun/package/packagers_manager.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.
-#
+
 import importlib
 import inspect
 import os

mlrun/package/utils/_archiver.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.
-#
+
 import os
 import tarfile
 import zipfile

mlrun/package/utils/_formatter.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.
-#
+
 import ast
 import json
 from abc import ABC, abstractmethod

mlrun/package/utils/_pickler.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.
-#
+
 import importlib
 import importlib.metadata as importlib_metadata
 import os

mlrun/package/utils/_supported_format.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 abc import ABC
 from typing import Generic, TypeVar, Union
 

mlrun/package/utils/log_hint_utils.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.
-#
+
 import typing
 
 from mlrun.errors import MLRunInvalidArgumentError

mlrun/package/utils/type_hint_utils.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.
-#
+
 import builtins
 import importlib
 import itertools