mlrun 1.7.0rc43__py3-none-any.whl → 1.7.0rc56__py3-none-any.whl

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,13 +89,13 @@ 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,
         event: dict,
         kind: mm_schemas.WriterEventKind = mm_schemas.WriterEventKind.RESULT,
-    ):
+    ) -> None:
         """
         Write a single result or metric to TSDB.
         """
@@ -113,7 +111,7 @@ class TDEngineConnector(TSDBConnector):
             # Write a new result
             table = self.tables[mm_schemas.TDEngineSuperTables.APP_RESULTS]
             table_name = (
-                f"{table_name}_" f"{event[mm_schemas.ResultData.RESULT_NAME]}"
+                f"{table_name}_{event[mm_schemas.ResultData.RESULT_NAME]}"
             ).replace("-", "_")
             event.pop(mm_schemas.ResultData.CURRENT_STATS, None)
 
@@ -121,9 +119,13 @@ class TDEngineConnector(TSDBConnector):
             # Write a new metric
             table = self.tables[mm_schemas.TDEngineSuperTables.METRICS]
             table_name = (
-                f"{table_name}_" f"{event[mm_schemas.MetricData.METRIC_NAME]}"
+                f"{table_name}_{event[mm_schemas.MetricData.METRIC_NAME]}"
             ).replace("-", "_")
 
+        # Escape the table name for case-sensitivity (ML-7908)
+        # https://github.com/taosdata/taos-connector-python/issues/260
+        table_name = f"`{table_name}`"
+
         # 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]
@@ -132,18 +134,19 @@ class TDEngineConnector(TSDBConnector):
             val=event[mm_schemas.WriterEvent.START_INFER_TIME]
         )
 
-        create_table_query = table._create_subtable_query(
-            subtable=table_name, values=event
+        create_table_sql = table._create_subtable_sql(subtable=table_name, values=event)
+
+        insert_statement = Statement(
+            TDEngineSchema._insert_subtable_stmt,
+            dict(columns=table.columns, subtable=table_name, values=event),
         )
-        self.connection.execute(create_table_query)
 
-        insert_statement = table._insert_subtable_query(
-            self.connection,
-            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(
@@ -262,7 +271,7 @@ class TDEngineConnector(TSDBConnector):
 
         project_condition = f"project = '{self.project}'"
         filter_query = (
-            f"{filter_query} AND {project_condition}"
+            f"({filter_query}) AND ({project_condition})"
             if filter_query
             else project_condition
         )
@@ -280,15 +289,16 @@ class TDEngineConnector(TSDBConnector):
             timestamp_column=timestamp_column,
             database=self.database,
         )
+        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,
@@ -336,11 +346,11 @@ class TDEngineConnector(TSDBConnector):
 
         metrics_condition = " OR ".join(
             [
-                f"({mm_schemas.WriterEvent.APPLICATION_NAME} = '{metric.app}' AND {name} = '{metric.name}')"
+                f"({mm_schemas.WriterEvent.APPLICATION_NAME}='{metric.app}' AND {name}='{metric.name}')"
                 for metric in metrics
             ]
         )
-        filter_query = f"endpoint_id='{endpoint_id}' AND ({metrics_condition})"
+        filter_query = f"(endpoint_id='{endpoint_id}') AND ({metrics_condition})"
 
         df = self._get_records(
             table=table,

mlrun/model_monitoring/helpers.py

@@ -18,6 +18,10 @@ import typing
 import numpy as np
 import pandas as pd
 
+if typing.TYPE_CHECKING:
+    from mlrun.db.base import RunDBInterface
+    from mlrun.projects import MlrunProject
+
 import mlrun
 import mlrun.artifacts
 import mlrun.common.model_monitoring.helpers
@@ -26,16 +30,11 @@ import mlrun.data_types.infer
 import mlrun.model_monitoring
 from mlrun.common.schemas.model_monitoring.model_endpoints import (
     ModelEndpointMonitoringMetric,
-    ModelEndpointMonitoringMetricType,
     _compose_full_name,
 )
 from mlrun.model_monitoring.model_endpoint import ModelEndpoint
 from mlrun.utils import logger
 
-if typing.TYPE_CHECKING:
-    from mlrun.db.base import RunDBInterface
-    from mlrun.projects import MlrunProject
-
 
 class _BatchDict(typing.TypedDict):
     minutes: int
@@ -63,7 +62,6 @@ def get_stream_path(
     )
 
     if not stream_uri or stream_uri == "v3io":
-        # TODO : remove the first part of this condition in 1.9.0
         stream_uri = mlrun.mlconf.get_model_monitoring_file_target_path(
             project=project,
             kind=mm_constants.FileTargetKind.STREAM,
@@ -71,8 +69,6 @@ def get_stream_path(
             function_name=function_name,
         )
 
-    if isinstance(stream_uri, list):  # ML-6043 - user side gets only the new stream uri
-        stream_uri = stream_uri[1]  # get new stream path, under projects
     return mlrun.common.model_monitoring.helpers.parse_monitoring_stream_path(
         stream_uri=stream_uri, project=project, function_name=function_name
     )
@@ -179,7 +175,7 @@ def _get_monitoring_time_window_from_controller_run(
 def update_model_endpoint_last_request(
     project: str,
     model_endpoint: ModelEndpoint,
-    current_request: datetime,
+    current_request: datetime.datetime,
     db: "RunDBInterface",
 ) -> None:
     """
@@ -190,7 +186,8 @@ def update_model_endpoint_last_request(
     :param current_request: current request time
     :param db:              DB interface.
     """
-    if model_endpoint.spec.stream_path != "":
+    is_model_server_endpoint = model_endpoint.spec.stream_path != ""
+    if is_model_server_endpoint:
         current_request = current_request.isoformat()
         logger.info(
             "Update model endpoint last request time (EP with serving)",
@@ -204,12 +201,13 @@ def update_model_endpoint_last_request(
             endpoint_id=model_endpoint.metadata.uid,
             attributes={mm_constants.EventFieldType.LAST_REQUEST: current_request},
         )
-    else:
+    else:  # model endpoint without any serving function - close the window "manually"
         try:
             time_window = _get_monitoring_time_window_from_controller_run(project, db)
         except mlrun.errors.MLRunNotFoundError:
-            logger.debug(
-                "Not bumping model endpoint last request time - the monitoring controller isn't deployed yet"
+            logger.warn(
+                "Not bumping model endpoint last request time - the monitoring controller isn't deployed yet.\n"
+                "Call `project.enable_model_monitoring()` first."
             )
             return
 
@@ -302,7 +300,7 @@ def get_invocations_fqn(project: str) -> str:
         project=project,
         app=mm_constants.SpecialApps.MLRUN_INFRA,
         name=mm_constants.PredictionsQueryConstants.INVOCATIONS,
-        type=ModelEndpointMonitoringMetricType.METRIC,
+        type=mm_constants.ModelEndpointMonitoringMetricType.METRIC,
     )
 
 
@@ -316,7 +314,7 @@ def get_invocations_metric(project: str) -> ModelEndpointMonitoringMetric:
     return ModelEndpointMonitoringMetric(
         project=project,
         app=mm_constants.SpecialApps.MLRUN_INFRA,
-        type=ModelEndpointMonitoringMetricType.METRIC,
+        type=mm_constants.ModelEndpointMonitoringMetricType.METRIC,
         name=mm_constants.PredictionsQueryConstants.INVOCATIONS,
         full_name=get_invocations_fqn(project),
     )

mlrun/model_monitoring/writer.py

@@ -160,7 +160,9 @@ class ModelMonitoringWriter(StepToDict):
             event_kind = f"{event_kind}_detected"
         else:
             event_kind = f"{event_kind}_suspected"
-        return alert_objects.EventKind(value=event_kind)
+        return alert_objects.EventKind(
+            value=mlrun.utils.helpers.normalize_name(event_kind)
+        )
 
     @staticmethod
     def _reconstruct_event(event: _RawEvent) -> tuple[_AppResultEvent, WriterEventKind]:

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:

mlrun/projects/pipelines.py

@@ -27,6 +27,8 @@ import mlrun_pipelines.utils
 import mlrun
 import mlrun.common.runtimes.constants
 import mlrun.common.schemas
+import mlrun.common.schemas.function
+import mlrun.common.schemas.workflow
 import mlrun.utils.notifications
 from mlrun.errors import err_to_str
 from mlrun.utils import (
@@ -44,21 +46,21 @@ from ..runtimes.pod import AutoMountType
 
 def get_workflow_engine(engine_kind, local=False):
     if pipeline_context.is_run_local(local):
-        if engine_kind == "kfp":
+        if engine_kind == mlrun.common.schemas.workflow.EngineType.KFP:
             logger.warning(
                 "Running kubeflow pipeline locally, note some ops may not run locally!"
             )
-        elif engine_kind == "remote":
+        elif engine_kind == mlrun.common.schemas.workflow.EngineType.REMOTE:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "Cannot run a remote pipeline locally using `kind='remote'` and `local=True`. "
                 "in order to run a local pipeline remotely, please use `engine='remote:local'` instead"
             )
         return _LocalRunner
-    if not engine_kind or engine_kind == "kfp":
+    if not engine_kind or engine_kind == mlrun.common.schemas.workflow.EngineType.KFP:
         return _KFPRunner
-    if engine_kind == "local":
+    if engine_kind == mlrun.common.schemas.workflow.EngineType.LOCAL:
         return _LocalRunner
-    if engine_kind == "remote":
+    if engine_kind == mlrun.common.schemas.workflow.EngineType.REMOTE:
         return _RemoteRunner
     raise mlrun.errors.MLRunInvalidArgumentError(
         f"Provided workflow engine is not supported. engine_kind={engine_kind}"
@@ -80,6 +82,7 @@ class WorkflowSpec(mlrun.model.ModelObj):
         schedule: typing.Union[str, mlrun.common.schemas.ScheduleCronTrigger] = None,
         cleanup_ttl: typing.Optional[int] = None,
         image: typing.Optional[str] = None,
+        workflow_runner_node_selector: typing.Optional[dict[str, str]] = None,
     ):
         self.engine = engine
         self.code = code
@@ -93,6 +96,7 @@ class WorkflowSpec(mlrun.model.ModelObj):
         self._tmp_path = None
         self.schedule = schedule
         self.image = image
+        self.workflow_runner_node_selector = workflow_runner_node_selector
 
     def get_source_file(self, context=""):
         if not self.code and not self.path:
@@ -311,7 +315,11 @@ def get_db_function(project, key) -> mlrun.runtimes.BaseRuntime:
 
 
 def enrich_function_object(
-    project, function, decorator=None, copy_function=True, try_auto_mount=True
+    project: mlrun.common.schemas.Project,
+    function: mlrun.runtimes.BaseRuntime,
+    decorator: typing.Callable = None,
+    copy_function: bool = True,
+    try_auto_mount: bool = True,
 ) -> mlrun.runtimes.BaseRuntime:
     if hasattr(function, "_enriched"):
         return function
@@ -352,7 +360,6 @@ def enrich_function_object(
         f.enrich_runtime_spec(
             project.spec.default_function_node_selector,
         )
-
     if try_auto_mount:
         if (
             decorator and AutoMountType.is_auto_modifier(decorator)
@@ -452,7 +459,12 @@ class _PipelineRunner(abc.ABC):
 
     @staticmethod
     @abc.abstractmethod
-    def wait_for_completion(run_id, project=None, timeout=None, expected_statuses=None):
+    def wait_for_completion(
+        run: "_PipelineRunStatus",
+        project: typing.Optional["mlrun.projects.MlrunProject"] = None,
+        timeout: typing.Optional[int] = None,
+        expected_statuses: list[str] = None,
+    ):
         pass
 
     @staticmethod
@@ -581,13 +593,13 @@ class _KFPRunner(_PipelineRunner):
             logger.warning(
                 "Setting notifications on kfp pipeline runner uses old notification behavior. "
                 "Notifications will only be sent if you wait for pipeline completion. "
-                "To use the new notification behavior, use the remote pipeline runner."
+                "Some of the features (like setting message or severity level) are not supported."
             )
             # for start message, fallback to old notification behavior
             for notification in notifications or []:
-                project.notifiers.add_notification(
-                    notification.kind, notification.params
-                )
+                params = notification.params
+                params.update(notification.secret_params)
+                project.notifiers.add_notification(notification.kind, params)
 
         run_id = _run_pipeline(
             workflow_handler,
@@ -625,12 +637,19 @@ class _KFPRunner(_PipelineRunner):
         return _PipelineRunStatus(run_id, cls, project=project, workflow=workflow_spec)
 
     @staticmethod
-    def wait_for_completion(run, project=None, timeout=None, expected_statuses=None):
+    def wait_for_completion(
+        run: "_PipelineRunStatus",
+        project: typing.Optional["mlrun.projects.MlrunProject"] = None,
+        timeout: typing.Optional[int] = None,
+        expected_statuses: list[str] = None,
+    ):
+        project_name = project.metadata.name if project else ""
         logger.info(
-            "Waiting for pipeline run completion", run_id=run.run_id, project=project
+            "Waiting for pipeline run completion",
+            run_id=run.run_id,
+            project=project_name,
         )
         timeout = timeout or 60 * 60
-        project_name = project.metadata.name if project else ""
         run_info = wait_for_pipeline_completion(
             run.run_id,
             timeout=timeout,
@@ -1062,7 +1081,7 @@ def load_and_run(
     # extract "start" notification if exists
     start_notifications = [
         notification
-        for notification in context.get_notifications()
+        for notification in context.get_notifications(unmask_secret_params=True)
         if "running" in notification.when
     ]