mlrun 1.8.0rc30__py3-none-any.whl → 1.8.0rc32__py3-none-any.whl

mlrun/model_monitoring/db/tsdb/base.py

@@ -12,12 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import typing
 from abc import ABC, abstractmethod
 from datetime import datetime
+from typing import Callable, ClassVar, Literal, Optional, Union
 
 import pandas as pd
 import pydantic.v1
+import v3io_frames.client
 
 import mlrun.common.schemas.model_monitoring as mm_schemas
 import mlrun.model_monitoring.db.tsdb.helpers
@@ -26,7 +27,7 @@ from mlrun.utils import logger
 
 
 class TSDBConnector(ABC):
-    type: typing.ClassVar[str]
+    type: ClassVar[str]
 
     def __init__(self, project: str) -> None:
         """
@@ -130,17 +131,17 @@ class TSDBConnector(ABC):
         start: datetime,
         end: datetime,
         metrics: list[mm_schemas.ModelEndpointMonitoringMetric],
-        type: typing.Literal["metrics", "results"],
+        type: Literal["metrics", "results"],
         with_result_extra_data: bool,
-    ) -> typing.Union[
+    ) -> Union[
         list[
-            typing.Union[
+            Union[
                 mm_schemas.ModelEndpointMonitoringResultValues,
                 mm_schemas.ModelEndpointMonitoringMetricNoData,
             ],
         ],
         list[
-            typing.Union[
+            Union[
                 mm_schemas.ModelEndpointMonitoringMetricValues,
                 mm_schemas.ModelEndpointMonitoringMetricNoData,
             ],
@@ -166,10 +167,10 @@ class TSDBConnector(ABC):
         endpoint_id: str,
         start: datetime,
         end: datetime,
-        aggregation_window: typing.Optional[str] = None,
-        agg_funcs: typing.Optional[list[str]] = None,
-        limit: typing.Optional[int] = None,
-    ) -> typing.Union[
+        aggregation_window: Optional[str] = None,
+        agg_funcs: Optional[list[str]] = None,
+        limit: Optional[int] = None,
+    ) -> Union[
         mm_schemas.ModelEndpointMonitoringMetricValues,
         mm_schemas.ModelEndpointMonitoringMetricNoData,
     ]:
@@ -195,10 +196,11 @@ class TSDBConnector(ABC):
     @abstractmethod
     def get_last_request(
         self,
-        endpoint_ids: typing.Union[str, list[str]],
-        start: typing.Optional[datetime] = None,
-        end: typing.Optional[datetime] = None,
-    ) -> pd.DataFrame:
+        endpoint_ids: Union[str, list[str]],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
+        get_raw: bool = False,
+    ) -> Union[pd.DataFrame, list[v3io_frames.client.RawFrame]]:
         """
         Fetches data from the predictions TSDB table and returns the most recent request
         timestamp for each specified endpoint.
@@ -206,6 +208,8 @@ class TSDBConnector(ABC):
         :param endpoint_ids:    A list of model endpoint identifiers.
         :param start:           The start time for the query.
         :param end:             The end time for the query.
+        :param get_raw:         Whether to return the request as raw frames rather than a pandas dataframe. Defaults
+          to False. This can greatly improve performance when a dataframe isn't needed.
 
         :return: A pd.DataFrame containing the columns [endpoint_id, last_request, last_latency].
         If an endpoint has not been invoked within the specified time range, it will not appear in the result.
@@ -214,10 +218,11 @@ class TSDBConnector(ABC):
     @abstractmethod
     def get_drift_status(
         self,
-        endpoint_ids: typing.Union[str, list[str]],
-        start: typing.Optional[datetime] = None,
-        end: typing.Optional[datetime] = None,
-    ) -> pd.DataFrame:
+        endpoint_ids: Union[str, list[str]],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
+        get_raw: bool = False,
+    ) -> Union[pd.DataFrame, list[v3io_frames.client.RawFrame]]:
         """
         Fetches data from the app-results TSDB table and returns the highest status among all
         the result in the provided time range, which by default is the last 24 hours, for each specified endpoint.
@@ -225,6 +230,8 @@ class TSDBConnector(ABC):
         :param endpoint_ids:    A list of model endpoint identifiers.
         :param start:           The start time for the query.
         :param end:             The end time for the query.
+        :param get_raw:         Whether to return the request as raw frames rather than a pandas dataframe. Defaults
+          to False. This can greatly improve performance when a dataframe isn't needed.
 
         :return: A pd.DataFrame containing the columns [result_status, endpoint_id].
         If an endpoint has not been monitored within the specified time range (last 24 hours),
@@ -234,9 +241,9 @@ class TSDBConnector(ABC):
     @abstractmethod
     def get_metrics_metadata(
         self,
-        endpoint_id: typing.Union[str, list[str]],
-        start: typing.Optional[datetime] = None,
-        end: typing.Optional[datetime] = None,
+        endpoint_id: Union[str, list[str]],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
     ) -> pd.DataFrame:
         """
         Fetches distinct metrics metadata from the metrics TSDB table for a specified model endpoints.
@@ -252,9 +259,9 @@ class TSDBConnector(ABC):
     @abstractmethod
     def get_results_metadata(
         self,
-        endpoint_id: typing.Union[str, list[str]],
-        start: typing.Optional[datetime] = None,
-        end: typing.Optional[datetime] = None,
+        endpoint_id: Union[str, list[str]],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
     ) -> pd.DataFrame:
         """
         Fetches distinct results metadata from the app-results TSDB table for a specified model endpoints.
@@ -270,16 +277,19 @@ class TSDBConnector(ABC):
     @abstractmethod
     def get_error_count(
         self,
-        endpoint_ids: typing.Union[str, list[str]],
-        start: typing.Optional[datetime] = None,
-        end: typing.Optional[datetime] = None,
-    ) -> pd.DataFrame:
+        endpoint_ids: Union[str, list[str]],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
+        get_raw: bool = False,
+    ) -> Union[pd.DataFrame, list[v3io_frames.client.RawFrame]]:
         """
         Fetches data from the error TSDB table and returns the error count for each specified endpoint.
 
         :param endpoint_ids:    A list of model endpoint identifiers.
         :param start:           The start time for the query.
         :param end:             The end time for the query.
+        :param get_raw:         Whether to return the request as raw frames rather than a pandas dataframe. Defaults
+          to False. This can greatly improve performance when a dataframe isn't needed.
 
         :return: A pd.DataFrame containing the columns [error_count, endpoint_id].
         If an endpoint have not raised error within the specified time range, it will not appear in the result.
@@ -288,10 +298,11 @@ class TSDBConnector(ABC):
     @abstractmethod
     def get_avg_latency(
         self,
-        endpoint_ids: typing.Union[str, list[str]],
-        start: typing.Optional[datetime] = None,
-        end: typing.Optional[datetime] = None,
-    ) -> pd.DataFrame:
+        endpoint_ids: Union[str, list[str]],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
+        get_raw: bool = False,
+    ) -> Union[pd.DataFrame, list[v3io_frames.client.RawFrame]]:
         """
         Fetches data from the predictions TSDB table and returns the average latency for each specified endpoint
         in the provided time range, which by default is the last 24 hours.
@@ -299,11 +310,21 @@ class TSDBConnector(ABC):
         :param endpoint_ids:    A list of model endpoint identifiers.
         :param start:           The start time for the query.
         :param end:             The end time for the query.
+        :param get_raw:         Whether to return the request as raw frames rather than a pandas dataframe. Defaults
+          to False. This can greatly improve performance when a dataframe isn't needed.
 
         :return: A pd.DataFrame containing the columns [avg_latency, endpoint_id].
         If an endpoint has not been invoked within the specified time range, it will not appear in the result.
         """
 
+    async def add_basic_metrics(
+        self,
+        model_endpoint_objects: list[mlrun.common.schemas.ModelEndpoint],
+        project: str,
+        run_in_threadpool: Callable,
+    ) -> list[mlrun.common.schemas.ModelEndpoint]:
+        raise NotImplementedError()
+
     @staticmethod
     def df_to_metrics_values(
         *,
@@ -311,7 +332,7 @@ class TSDBConnector(ABC):
         metrics: list[mm_schemas.ModelEndpointMonitoringMetric],
         project: str,
     ) -> list[
-        typing.Union[
+        Union[
             mm_schemas.ModelEndpointMonitoringMetricValues,
             mm_schemas.ModelEndpointMonitoringMetricNoData,
         ]
@@ -324,7 +345,7 @@ class TSDBConnector(ABC):
         metrics_without_data = {metric.full_name: metric for metric in metrics}
 
         metrics_values: list[
-            typing.Union[
+            Union[
                 mm_schemas.ModelEndpointMonitoringMetricValues,
                 mm_schemas.ModelEndpointMonitoringMetricNoData,
             ]
@@ -377,7 +398,7 @@ class TSDBConnector(ABC):
         metrics: list[mm_schemas.ModelEndpointMonitoringMetric],
         project: str,
     ) -> list[
-        typing.Union[
+        Union[
             mm_schemas.ModelEndpointMonitoringResultValues,
             mm_schemas.ModelEndpointMonitoringMetricNoData,
         ]
@@ -390,7 +411,7 @@ class TSDBConnector(ABC):
         metrics_without_data = {metric.full_name: metric for metric in metrics}
 
         metrics_values: list[
-            typing.Union[
+            Union[
                 mm_schemas.ModelEndpointMonitoringResultValues,
                 mm_schemas.ModelEndpointMonitoringMetricNoData,
             ]
@@ -536,7 +557,7 @@ class TSDBConnector(ABC):
         *,
         df: pd.DataFrame,
         project: str,
-        type: typing.Union[str, mm_schemas.ModelEndpointMonitoringMetricType],
+        type: Union[str, mm_schemas.ModelEndpointMonitoringMetricType],
     ) -> dict[str, list[mm_schemas.ModelEndpointMonitoringMetric]]:
         """
         Parse a DataFrame of metrics from the TSDB into a dict of intersection metrics/results by name and application
@@ -591,8 +612,8 @@ class TSDBConnector(ABC):
 
     @staticmethod
     def _get_start_end(
-        start: typing.Union[datetime, None],
-        end: typing.Union[datetime, None],
+        start: Union[datetime, None],
+        end: Union[datetime, None],
     ) -> tuple[datetime, datetime]:
         """
         static utils function for tsdb start end format

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

@@ -11,9 +11,9 @@
 # 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 asyncio
 from datetime import datetime, timedelta
+from typing import Callable, Literal, Optional, Union
 
 import pandas as pd
 import taosws
@@ -25,6 +25,7 @@ from taoswswrap.tdengine_connection import (
 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.datastore.datastore_profile import DatastoreProfile
 from mlrun.model_monitoring.db import TSDBConnector
 from mlrun.model_monitoring.helpers import get_invocations_fqn
 from mlrun.utils import logger
@@ -40,16 +41,13 @@ class TDEngineConnector(TSDBConnector):
     def __init__(
         self,
         project: str,
-        database: typing.Optional[str] = None,
+        profile: DatastoreProfile,
+        database: Optional[str] = None,
         **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._tdengine_connection_profile = profile
         self.database = (
             database
             or f"{tdengine_schemas._MODEL_MONITORING_DATABASE}_{mlrun.mlconf.system_id}"
@@ -70,7 +68,7 @@ class TDEngineConnector(TSDBConnector):
     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 = TDEngineConnection(self._tdengine_connection_profile.dsn())
         conn.run(
             statements=f"CREATE DATABASE IF NOT EXISTS {self.database}",
             timeout=self._timeout,
@@ -168,11 +166,11 @@ class TDEngineConnector(TSDBConnector):
         )
 
     @staticmethod
-    def _convert_to_datetime(val: typing.Union[str, datetime]) -> datetime:
+    def _convert_to_datetime(val: Union[str, datetime]) -> datetime:
         return datetime.fromisoformat(val) if isinstance(val, str) else val
 
     @staticmethod
-    def _get_endpoint_filter(endpoint_id: typing.Union[str, list[str]]) -> str:
+    def _get_endpoint_filter(endpoint_id: Union[str, list[str]]) -> str:
         if isinstance(endpoint_id, str):
             return f"endpoint_id='{endpoint_id}'"
         elif isinstance(endpoint_id, list):
@@ -200,10 +198,10 @@ class TDEngineConnector(TSDBConnector):
 
         def apply_tdengine_target(name, after):
             graph.add_step(
-                "storey.TDEngineTarget",
+                "mlrun.datastore.storeytargets.TDEngineStoreyTarget",
                 name=name,
                 after=after,
-                url=self._tdengine_connection_string,
+                url=f"ds://{self._tdengine_connection_profile.name}",
                 supertable=self.tables[
                     mm_schemas.TDEngineSuperTables.PREDICTIONS
                 ].super_table,
@@ -242,10 +240,10 @@ class TDEngineConnector(TSDBConnector):
             after="ForwardError",
         )
         graph.add_step(
-            "storey.TDEngineTarget",
+            "mlrun.datastore.storeytargets.TDEngineStoreyTarget",
             name="tsdb_error",
             after="error_extractor",
-            url=self._tdengine_connection_string,
+            url=f"ds://{self._tdengine_connection_profile.name}",
             supertable=self.tables[mm_schemas.TDEngineSuperTables.ERRORS].super_table,
             table_col=mm_schemas.EventFieldType.TABLE_COLUMN,
             time_col=mm_schemas.EventFieldType.TIME,
@@ -307,17 +305,17 @@ class TDEngineConnector(TSDBConnector):
         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,
+        columns: Optional[list[str]] = None,
+        filter_query: Optional[str] = None,
+        interval: Optional[str] = None,
+        agg_funcs: Optional[list] = None,
+        limit: Optional[int] = None,
+        sliding_window_step: Optional[str] = None,
         timestamp_column: str = mm_schemas.EventFieldType.TIME,
-        group_by: typing.Optional[typing.Union[list[str], str]] = None,
-        preform_agg_columns: typing.Optional[list] = None,
-        order_by: typing.Optional[str] = None,
-        desc: typing.Optional[bool] = None,
+        group_by: Optional[Union[list[str], str]] = None,
+        preform_agg_columns: Optional[list] = None,
+        order_by: Optional[str] = None,
+        desc: Optional[bool] = None,
     ) -> pd.DataFrame:
         """
         Getting records from TSDB data collection.
@@ -387,17 +385,17 @@ class TDEngineConnector(TSDBConnector):
         start: datetime,
         end: datetime,
         metrics: list[mm_schemas.ModelEndpointMonitoringMetric],
-        type: typing.Literal["metrics", "results"],
+        type: Literal["metrics", "results"],
         with_result_extra_data: bool = False,
-    ) -> typing.Union[
+    ) -> Union[
         list[
-            typing.Union[
+            Union[
                 mm_schemas.ModelEndpointMonitoringResultValues,
                 mm_schemas.ModelEndpointMonitoringMetricNoData,
             ],
         ],
         list[
-            typing.Union[
+            Union[
                 mm_schemas.ModelEndpointMonitoringMetricValues,
                 mm_schemas.ModelEndpointMonitoringMetricNoData,
             ],
@@ -475,10 +473,10 @@ class TDEngineConnector(TSDBConnector):
         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[
+        aggregation_window: Optional[str] = None,
+        agg_funcs: Optional[list] = None,
+        limit: Optional[int] = None,
+    ) -> Union[
         mm_schemas.ModelEndpointMonitoringMetricValues,
         mm_schemas.ModelEndpointMonitoringMetricNoData,
     ]:
@@ -530,9 +528,10 @@ class TDEngineConnector(TSDBConnector):
 
     def get_last_request(
         self,
-        endpoint_ids: typing.Union[str, list[str]],
-        start: typing.Optional[datetime] = None,
-        end: typing.Optional[datetime] = None,
+        endpoint_ids: Union[str, list[str]],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
+        get_raw: bool = False,
     ) -> pd.DataFrame:
         filter_query = self._get_endpoint_filter(endpoint_id=endpoint_ids)
         start, end = self._get_start_end(start, end)
@@ -570,9 +569,10 @@ class TDEngineConnector(TSDBConnector):
 
     def get_drift_status(
         self,
-        endpoint_ids: typing.Union[str, list[str]],
-        start: typing.Optional[datetime] = None,
-        end: typing.Optional[datetime] = None,
+        endpoint_ids: Union[str, list[str]],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
+        get_raw: bool = False,
     ) -> pd.DataFrame:
         filter_query = self._get_endpoint_filter(endpoint_id=endpoint_ids)
         start = start or (mlrun.utils.datetime_now() - timedelta(hours=24))
@@ -603,9 +603,9 @@ class TDEngineConnector(TSDBConnector):
 
     def get_metrics_metadata(
         self,
-        endpoint_id: typing.Union[str, list[str]],
-        start: typing.Optional[datetime] = None,
-        end: typing.Optional[datetime] = None,
+        endpoint_id: Union[str, list[str]],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
     ) -> pd.DataFrame:
         start, end = self._get_start_end(start, end)
         df = self._get_records(
@@ -640,9 +640,9 @@ class TDEngineConnector(TSDBConnector):
 
     def get_results_metadata(
         self,
-        endpoint_id: typing.Union[str, list[str]],
-        start: typing.Optional[datetime] = None,
-        end: typing.Optional[datetime] = None,
+        endpoint_id: Union[str, list[str]],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
     ) -> pd.DataFrame:
         start, end = self._get_start_end(start, end)
         df = self._get_records(
@@ -679,9 +679,10 @@ class TDEngineConnector(TSDBConnector):
 
     def get_error_count(
         self,
-        endpoint_ids: typing.Union[str, list[str]],
-        start: typing.Optional[datetime] = None,
-        end: typing.Optional[datetime] = None,
+        endpoint_ids: Union[str, list[str]],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
+        get_raw: bool = False,
     ) -> pd.DataFrame:
         filter_query = self._get_endpoint_filter(endpoint_id=endpoint_ids)
         filter_query += f"AND {mm_schemas.EventFieldType.ERROR_TYPE} = '{mm_schemas.EventFieldType.INFER_ERROR}'"
@@ -709,9 +710,10 @@ class TDEngineConnector(TSDBConnector):
 
     def get_avg_latency(
         self,
-        endpoint_ids: typing.Union[str, list[str]],
-        start: typing.Optional[datetime] = None,
-        end: typing.Optional[datetime] = None,
+        endpoint_ids: Union[str, list[str]],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
+        get_raw: bool = False,
     ) -> pd.DataFrame:
         endpoint_ids = (
             endpoint_ids if isinstance(endpoint_ids, list) else [endpoint_ids]
@@ -739,11 +741,74 @@ class TDEngineConnector(TSDBConnector):
             df.dropna(inplace=True)
         return df
 
+    async def add_basic_metrics(
+        self,
+        model_endpoint_objects: list[mlrun.common.schemas.ModelEndpoint],
+        project: str,
+        run_in_threadpool: Callable,
+    ) -> list[mlrun.common.schemas.ModelEndpoint]:
+        """
+        Add basic metrics to the model endpoint object.
+
+        :param model_endpoint_objects: A list of `ModelEndpoint` objects that will
+                                        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.
+
+        :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)
+
+        def add_metrics(
+            mep: mlrun.common.schemas.ModelEndpoint,
+            df_dictionary: dict[str, pd.DataFrame],
+        ):
+            for metric in df_dictionary.keys():
+                df = df_dictionary.get(metric, pd.DataFrame())
+                if not df.empty:
+                    line = df[df["endpoint_id"] == mep.metadata.uid]
+                    if not line.empty and metric in line:
+                        value = line[metric].item()
+                        if isinstance(value, pd.Timestamp):
+                            value = value.to_pydatetime()
+                        setattr(mep.status, metric, value)
+
+            return mep
+
+        return list(
+            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,
+                    },
+                ),
+                model_endpoint_objects,
+            )
+        )
+
     # 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]:
+    # ) -> Optional[mm_schemas.ModelEndpointMonitoringMetric]:
     #     """
     #     Read the "invocations" metric for the provided model endpoint, and return the metric object
     #     if it exists.