mlrun 1.7.0rc48__py3-none-any.whl → 1.7.0rc52__py3-none-any.whl

mlrun/model_monitoring/controller.py

@@ -15,28 +15,22 @@
 import concurrent.futures
 import datetime
 import json
-import multiprocessing
 import os
 import re
 from collections.abc import Iterator
-from typing import Any, NamedTuple, Optional, Union, cast
+from typing import NamedTuple, Optional, Union, cast
 
 import nuclio
 
 import mlrun
 import mlrun.common.schemas.model_monitoring.constants as mm_constants
 import mlrun.data_types.infer
-import mlrun.feature_store as fstore
 import mlrun.model_monitoring.db.stores
-from mlrun.common.model_monitoring.helpers import FeatureStats, pad_features_hist
 from mlrun.datastore import get_stream_pusher
-from mlrun.datastore.targets import ParquetTarget
 from mlrun.errors import err_to_str
 from mlrun.model_monitoring.helpers import (
     _BatchDict,
     batch_dict2timedelta,
-    calculate_inputs_statistics,
-    get_monitoring_parquet_path,
     get_stream_path,
 )
 from mlrun.utils import datetime_now, logger
@@ -292,15 +286,9 @@ class MonitoringApplicationController:
         )
 
         self.model_monitoring_access_key = self._get_model_monitoring_access_key()
-        self.parquet_directory = get_monitoring_parquet_path(
-            self.project_obj,
-            kind=mm_constants.FileTargetKind.APPS_PARQUET,
+        self.tsdb_connector = mlrun.model_monitoring.get_tsdb_connector(
+            project=self.project
         )
-        self.storage_options = None
-        if not mlrun.mlconf.is_ce_mode():
-            self._initialize_v3io_configurations()
-        elif self.parquet_directory.startswith("s3://"):
-            self.storage_options = mlrun.mlconf.get_s3_storage_options()
 
     @staticmethod
     def _get_model_monitoring_access_key() -> Optional[str]:
@@ -310,12 +298,6 @@ class MonitoringApplicationController:
             access_key = mlrun.mlconf.get_v3io_access_key()
         return access_key
 
-    def _initialize_v3io_configurations(self) -> None:
-        self.storage_options = dict(
-            v3io_access_key=self.model_monitoring_access_key,
-            v3io_api=mlrun.mlconf.v3io_api,
-        )
-
     def run(self) -> None:
         """
         Main method for run all the relevant monitoring applications on each endpoint.
@@ -367,11 +349,8 @@ class MonitoringApplicationController:
             )
             return
         # Initialize a process pool that will be used to run each endpoint applications on a dedicated process
-        with concurrent.futures.ProcessPoolExecutor(
+        with concurrent.futures.ThreadPoolExecutor(
             max_workers=min(len(endpoints), 10),
-            # On Linux, the default is "fork" (this is set to change in Python 3.14), which inherits the current heap
-            # and resources (such as sockets), which is not what we want (ML-7160)
-            mp_context=multiprocessing.get_context("spawn"),
         ) as pool:
             for endpoint in endpoints:
                 if (
@@ -395,13 +374,10 @@ class MonitoringApplicationController:
                         applications_names=applications_names,
                         batch_window_generator=self._batch_window_generator,
                         project=self.project,
-                        parquet_directory=self.parquet_directory,
-                        storage_options=self.storage_options,
                         model_monitoring_access_key=self.model_monitoring_access_key,
+                        tsdb_connector=self.tsdb_connector,
                     )
 
-        self._delete_old_parquet(endpoints=endpoints)
-
     @classmethod
     def model_endpoint_process(
         cls,
@@ -409,9 +385,8 @@ class MonitoringApplicationController:
         applications_names: list[str],
         batch_window_generator: _BatchWindowGenerator,
         project: str,
-        parquet_directory: str,
-        storage_options: dict,
         model_monitoring_access_key: str,
+        tsdb_connector: mlrun.model_monitoring.db.tsdb.TSDBConnector,
     ) -> None:
         """
         Process a model endpoint and trigger the monitoring applications. This function running on different process
@@ -422,16 +397,13 @@ class MonitoringApplicationController:
         :param applications_names:          (list[str]) List of application names to push results to.
         :param batch_window_generator:      (_BatchWindowGenerator) An object that generates _BatchWindow objects.
         :param project:                     (str) Project name.
-        :param parquet_directory:           (str) Directory to store application parquet files
-        :param storage_options:             (dict) Storage options for writing ParquetTarget.
         :param model_monitoring_access_key: (str) Access key to apply the model monitoring process.
+        :param tsdb_connector:              (mlrun.model_monitoring.db.tsdb.TSDBConnector) TSDB connector
         """
         endpoint_id = endpoint[mm_constants.EventFieldType.UID]
+        # if false the endpoint represent batch infer step.
+        has_stream = endpoint[mm_constants.EventFieldType.STREAM_PATH] != ""
         try:
-            m_fs = fstore.get_feature_set(
-                endpoint[mm_constants.EventFieldType.FEATURE_SET_URI]
-            )
-
             for application in applications_names:
                 batch_window = batch_window_generator.get_batch_window(
                     project=project,
@@ -439,158 +411,70 @@ class MonitoringApplicationController:
                     application=application,
                     first_request=endpoint[mm_constants.EventFieldType.FIRST_REQUEST],
                     last_request=endpoint[mm_constants.EventFieldType.LAST_REQUEST],
-                    has_stream=endpoint[mm_constants.EventFieldType.STREAM_PATH] != "",
+                    has_stream=has_stream,
                 )
 
                 for start_infer_time, end_infer_time in batch_window.get_intervals():
-                    # start - TODO : delete in 1.9.0 (V1 app deprecation)
-                    try:
-                        # Get application sample data
-                        offline_response = cls._get_sample_df(
-                            feature_set=m_fs,
+                    prediction_metric = tsdb_connector.read_predictions(
+                        endpoint_id=endpoint_id,
+                        start=start_infer_time,
+                        end=end_infer_time,
+                    )
+                    if not prediction_metric.data and has_stream:
+                        logger.info(
+                            "No data found for the given interval",
+                            start=start_infer_time,
+                            end=end_infer_time,
+                            endpoint_id=endpoint_id,
+                        )
+                    else:
+                        logger.info(
+                            "Data found for the given interval",
+                            start=start_infer_time,
+                            end=end_infer_time,
                             endpoint_id=endpoint_id,
+                        )
+                        cls._push_to_applications(
                             start_infer_time=start_infer_time,
                             end_infer_time=end_infer_time,
-                            parquet_directory=parquet_directory,
-                            storage_options=storage_options,
-                            application_name=application,
-                        )
-
-                        df = offline_response.to_dataframe()
-                        parquet_target_path = offline_response.vector.get_target_path()
-
-                        if len(df) == 0:
-                            logger.info(
-                                "During this time window, the endpoint has not received any data",
-                                endpoint=endpoint[mm_constants.EventFieldType.UID],
-                                start_time=start_infer_time,
-                                end_time=end_infer_time,
-                            )
-                            continue
-
-                    except FileNotFoundError:
-                        logger.warn(
-                            "No parquets were written yet",
-                            endpoint=endpoint[mm_constants.EventFieldType.UID],
+                            endpoint_id=endpoint_id,
+                            project=project,
+                            applications_names=[application],
+                            model_monitoring_access_key=model_monitoring_access_key,
                         )
-                        continue
-
-                    # Get the timestamp of the latest request:
-                    latest_request = df[mm_constants.EventFieldType.TIMESTAMP].iloc[-1]
-
-                    # Get the feature stats from the model endpoint for reference data
-                    feature_stats = json.loads(
-                        endpoint[mm_constants.EventFieldType.FEATURE_STATS]
-                    )
-
-                    # Pad the original feature stats to accommodate current
-                    # data out of the original range (unless already padded)
-                    pad_features_hist(FeatureStats(feature_stats))
-
-                    # Get the current stats:
-                    current_stats = calculate_inputs_statistics(
-                        sample_set_statistics=feature_stats, inputs=df
-                    )
-                    # end - TODO : delete in 1.9.0 (V1 app deprecation)
-                    cls._push_to_applications(
-                        current_stats=current_stats,
-                        feature_stats=feature_stats,
-                        start_infer_time=start_infer_time,
-                        end_infer_time=end_infer_time,
-                        endpoint_id=endpoint_id,
-                        latest_request=latest_request,
-                        project=project,
-                        applications_names=[application],
-                        model_monitoring_access_key=model_monitoring_access_key,
-                        parquet_target_path=parquet_target_path,
-                    )
         except Exception:
             logger.exception(
                 "Encountered an exception",
                 endpoint_id=endpoint[mm_constants.EventFieldType.UID],
             )
 
-    def _delete_old_parquet(self, endpoints: list[dict[str, Any]], days: int = 1):
-        """
-        Delete application parquets older than the argument days.
-
-        :param endpoints: A list of dictionaries of model endpoints records.
-        """
-        if self.parquet_directory.startswith("v3io:///"):
-            # create fs with access to the user side (under projects)
-            store, _, _ = mlrun.store_manager.get_or_create_store(
-                self.parquet_directory,
-                {"V3IO_ACCESS_KEY": self.model_monitoring_access_key},
-            )
-            fs = store.filesystem
-
-            # calculate time threshold (keep only files from the last 24 hours)
-            time_to_keep = (
-                datetime.datetime.now(tz=datetime.timezone.utc)
-                - datetime.timedelta(days=days)
-            ).timestamp()
-
-            for endpoint in endpoints:
-                try:
-                    apps_parquet_directories = fs.listdir(
-                        path=f"{self.parquet_directory}"
-                        f"/key={endpoint[mm_constants.EventFieldType.UID]}"
-                    )
-                    for directory in apps_parquet_directories:
-                        if directory["mtime"] < time_to_keep:
-                            # Delete files
-                            fs.rm(path=directory["name"], recursive=True)
-                            # Delete directory
-                            fs.rmdir(path=directory["name"])
-                except FileNotFoundError:
-                    logger.info(
-                        "Application parquet directory is empty, "
-                        "probably parquets have not yet been created for this app",
-                        endpoint=endpoint[mm_constants.EventFieldType.UID],
-                        path=f"{self.parquet_directory}"
-                        f"/key={endpoint[mm_constants.EventFieldType.UID]}",
-                    )
-
     @staticmethod
     def _push_to_applications(
-        current_stats,
-        feature_stats,
-        start_infer_time,
-        end_infer_time,
-        endpoint_id,
-        latest_request,
-        project,
-        applications_names,
-        model_monitoring_access_key,
-        parquet_target_path,
+        start_infer_time: datetime.datetime,
+        end_infer_time: datetime.datetime,
+        endpoint_id: str,
+        project: str,
+        applications_names: list[str],
+        model_monitoring_access_key: str,
     ):
         """
         Pushes data to multiple stream applications.
 
-        :param current_stats:       Current statistics of input data.
-        :param feature_stats:       Statistics of train features.
-        :param start_infer_time:    The beginning of the infer interval window.
-        :param end_infer_time:      The end of the infer interval window.
-        :param endpoint_id:         Identifier for the model endpoint.
-        :param latest_request:      Timestamp of the latest model request.
-        :param project: mlrun       Project name.
-        :param applications_names:  List of application names to which data will be pushed.
+        :param start_infer_time:            The beginning of the infer interval window.
+        :param end_infer_time:              The end of the infer interval window.
+        :param endpoint_id:                 Identifier for the model endpoint.
+        :param project: mlrun               Project name.
+        :param applications_names:          List of application names to which data will be pushed.
+        :param model_monitoring_access_key: Access key to apply the model monitoring process.
 
         """
-
         data = {
-            mm_constants.ApplicationEvent.CURRENT_STATS: json.dumps(current_stats),
-            mm_constants.ApplicationEvent.FEATURE_STATS: json.dumps(feature_stats),
-            mm_constants.ApplicationEvent.SAMPLE_PARQUET_PATH: parquet_target_path,
             mm_constants.ApplicationEvent.START_INFER_TIME: start_infer_time.isoformat(
                 sep=" ", timespec="microseconds"
             ),
             mm_constants.ApplicationEvent.END_INFER_TIME: end_infer_time.isoformat(
                 sep=" ", timespec="microseconds"
             ),
-            mm_constants.ApplicationEvent.LAST_REQUEST: latest_request.isoformat(
-                sep=" ", timespec="microseconds"
-            ),
             mm_constants.ApplicationEvent.ENDPOINT_ID: endpoint_id,
             mm_constants.ApplicationEvent.OUTPUT_STREAM_URI: get_stream_path(
                 project=project,
@@ -608,53 +492,6 @@ class MonitoringApplicationController:
                 [data]
             )
 
-    @staticmethod
-    def _get_sample_df(
-        feature_set: mlrun.common.schemas.FeatureSet,
-        endpoint_id: str,
-        start_infer_time: datetime.datetime,
-        end_infer_time: datetime.datetime,
-        parquet_directory: str,
-        storage_options: dict,
-        application_name: str,
-    ) -> mlrun.feature_store.OfflineVectorResponse:
-        """
-        Retrieves a sample DataFrame of the current input according to the provided infer interval window.
-
-        :param feature_set:         The main feature set.
-        :param endpoint_id:         Identifier for the model endpoint.
-        :param start_infer_time:    The beginning of the infer interval window.
-        :param end_infer_time:      The end of the infer interval window.
-        :param parquet_directory:   Directory where Parquet files are stored.
-        :param storage_options:     Storage options for accessing the data.
-        :param application_name:    Current application name.
-
-        :return: OfflineVectorResponse that can be used for generating a sample DataFrame for the specified endpoint.
-
-        """
-        features = [f"{feature_set.metadata.name}.*"]
-        vector = fstore.FeatureVector(
-            name=f"{endpoint_id}_vector",
-            features=features,
-            with_indexes=True,
-        )
-        vector.metadata.tag = application_name
-        vector.feature_set_objects = {feature_set.metadata.name: feature_set}
-
-        # get offline features based on application start and end time.
-        # store the result parquet by partitioning by controller end processing time
-        offline_response = vector.get_offline_features(
-            start_time=start_infer_time,
-            end_time=end_infer_time,
-            timestamp_for_filtering=mm_constants.EventFieldType.TIMESTAMP,
-            target=ParquetTarget(
-                path=parquet_directory
-                + f"/key={endpoint_id}/{int(start_infer_time.timestamp())}/{application_name}.parquet",
-                storage_options=storage_options,
-            ),
-        )
-        return offline_response
-
 
 def handler(context: nuclio.Context, event: nuclio.Event) -> None:
     """

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

@@ -107,25 +107,26 @@ class TDEngineSchema:
             )
         return f"CREATE TABLE if NOT EXISTS {self.database}.{subtable} USING {self.super_table} TAGS ({tags});"
 
+    @staticmethod
     def _insert_subtable_stmt(
-        self,
-        connection: taosws.Connection,
+        statement: taosws.TaosStmt,
+        columns: dict[str, _TDEngineColumn],
         subtable: str,
         values: dict[str, Union[str, int, float, datetime.datetime]],
     ) -> taosws.TaosStmt:
-        stmt = connection.statement()
-        question_marks = ", ".join("?" * len(self.columns))
-        stmt.prepare(f"INSERT INTO ? VALUES ({question_marks});")
-        stmt.set_tbname(subtable)
+        question_marks = ", ".join("?" * len(columns))
+        statement.prepare(f"INSERT INTO ? VALUES ({question_marks});")
+        statement.set_tbname(subtable)
 
         bind_params = []
 
-        for col_name, col_type in self.columns.items():
+        for col_name, col_type in columns.items():
             val = values[col_name]
             bind_params.append(values_to_column([val], col_type))
 
-        stmt.bind_param(bind_params)
-        return stmt
+        statement.bind_param(bind_params)
+        statement.add_batch()
+        return statement
 
     def _delete_subtable_query(
         self,

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

@@ -18,11 +18,16 @@ 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
 
@@ -52,25 +57,18 @@ class TDEngineConnector(TSDBConnector):
         self._init_super_tables()
 
     @property
-    def connection(self) -> taosws.Connection:
+    def connection(self) -> TDEngineConnection:
         if not self._connection:
             self._connection = self._create_connection()
         return self._connection
 
-    def _create_connection(self) -> taosws.Connection:
+    def _create_connection(self) -> TDEngineConnection:
         """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
-        try:
-            conn.execute(f"USE {self.database}")
-        except taosws.QueryError as e:
-            raise mlrun.errors.MLRunTSDBConnectionFailureError(
-                f"Failed to use TDEngine database {self.database}, {mlrun.errors.err_to_str(e)}"
-            )
+        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}")
+        conn.prefix_statements = [f"USE {self.database}"]
+        logger.debug("Connected to TDEngine", project=self.project)
         return conn
 
     def _init_super_tables(self):
@@ -91,7 +89,7 @@ class TDEngineConnector(TSDBConnector):
         """Create TDEngine supertables."""
         for table in self.tables:
             create_table_query = self.tables[table]._create_super_table_query()
-            self.connection.execute(create_table_query)
+            self.connection.run(statements=create_table_query)
 
     def write_application_event(
         self,
@@ -137,13 +135,18 @@ class TDEngineConnector(TSDBConnector):
         )
 
         create_table_sql = table._create_subtable_sql(subtable=table_name, values=event)
-        self.connection.execute(create_table_sql)
 
-        insert_statement = table._insert_subtable_stmt(
-            self.connection, 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,
+            ]
         )
-        insert_statement.add_batch()
-        insert_statement.execute()
 
     @staticmethod
     def _convert_to_datetime(val: typing.Union[str, datetime]) -> datetime:
@@ -200,18 +203,24 @@ class TDEngineConnector(TSDBConnector):
         """
         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,
+        )
         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)
+            subtables = self.connection.run(query=get_subtable_names_query).data
+            drop_statements = []
             for subtable in subtables:
-                drop_query = self.tables[table]._drop_subtable_query(
-                    subtable=subtable[0]
+                drop_statements.append(
+                    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}"
+            self.connection.run(statements=[drop_statements])
+        logger.debug(
+            "Deleted all project resources using the TDEngine connector",
+            project=self.project,
         )
 
     def get_model_endpoint_real_time_metrics(
@@ -282,14 +291,14 @@ class TDEngineConnector(TSDBConnector):
         )
         logger.debug("Querying TDEngine", query=full_query)
         try:
-            query_result = self.connection.query(full_query)
+            query_result = self.connection.run(query=full_query)
         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, columns=df_columns)
+        df_columns = [field.name for field in query_result.fields]
+        return pd.DataFrame(query_result.data, columns=df_columns)
 
     def read_metrics_data(
         self,

mlrun/projects/operations.py

@@ -15,10 +15,13 @@
 import warnings
 from typing import Optional, Union
 
-from mlrun_pipelines.models import PipelineNodeWrapper
+import mlrun_pipelines.common.models
+import mlrun_pipelines.models
 
 import mlrun
 import mlrun.common.constants as mlrun_constants
+import mlrun.common.schemas.function
+import mlrun.common.schemas.workflow
 from mlrun.utils import hub_prefix
 
 from .pipelines import enrich_function_object, pipeline_context
@@ -49,7 +52,7 @@ def _get_engine_and_function(function, project=None):
         function = enrich_function_object(project, function, copy_function=False)
 
     if not pipeline_context.workflow:
-        return "local", function
+        return mlrun.common.schemas.workflow.EngineType.LOCAL, function
 
     return pipeline_context.workflow.engine, function
 
@@ -78,7 +81,7 @@ def run_function(
     returns: Optional[list[Union[str, dict[str, str]]]] = None,
     builder_env: Optional[list] = None,
     reset_on_run: Optional[bool] = None,
-) -> Union[mlrun.model.RunObject, PipelineNodeWrapper]:
+) -> Union[mlrun.model.RunObject, mlrun_pipelines.models.PipelineNodeWrapper]:
     """Run a local or remote task as part of a local/kubeflow pipeline
 
     run_function() allow you to execute a function locally, on a remote cluster, or as part of an automated workflow
@@ -186,7 +189,7 @@ def run_function(
     )
     task.spec.verbose = task.spec.verbose or verbose
 
-    if engine == "kfp":
+    if engine == mlrun.common.schemas.workflow.EngineType.KFP:
         if schedule:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "Scheduling jobs is not supported when running a workflow with the kfp engine."
@@ -266,7 +269,7 @@ def build_function(
     overwrite_build_params: bool = False,
     extra_args: str = None,
     force_build: bool = False,
-) -> Union[BuildStatus, PipelineNodeWrapper]:
+) -> Union[BuildStatus, mlrun_pipelines.models.PipelineNodeWrapper]:
     """deploy ML function, build container with its dependencies
 
     :param function:        Name of the function (in the project) or function object
@@ -302,7 +305,7 @@ def build_function(
         raise mlrun.errors.MLRunInvalidArgumentError(
             "Cannot build use deploy_function()"
         )
-    if engine == "kfp":
+    if engine == mlrun.common.schemas.workflow.EngineType.KFP:
         if overwrite_build_params:
             function.spec.build.commands = None
         if requirements or requirements_file:
@@ -375,7 +378,7 @@ def deploy_function(
     builder_env: dict = None,
     project_object=None,
     mock: bool = None,
-) -> Union[DeployStatus, PipelineNodeWrapper]:
+) -> Union[DeployStatus, mlrun_pipelines.models.PipelineNodeWrapper]:
     """deploy real-time (nuclio based) functions
 
     :param function:   name of the function (in the project) or function object
@@ -392,7 +395,7 @@ def deploy_function(
         raise mlrun.errors.MLRunInvalidArgumentError(
             "deploy is used with real-time functions, for other kinds use build_function()"
         )
-    if engine == "kfp":
+    if engine == mlrun.common.schemas.workflow.EngineType.KFP:
         return function.deploy_step(models=models, env=env, tag=tag, verbose=verbose)
     else:
         if env: