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

mlrun/model_monitoring/applications/_application_steps.py

@@ -11,19 +11,16 @@
 # 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 json
-import typing
-from typing import Optional
+from typing import Any, Optional, Union
 
-import mlrun.common.helpers
-import mlrun.common.model_monitoring.helpers
 import mlrun.common.schemas.alert as alert_objects
 import mlrun.common.schemas.model_monitoring.constants as mm_constant
 import mlrun.datastore
-import mlrun.serving
-import mlrun.utils.helpers
-import mlrun.utils.v3io_clients
+import mlrun.model_monitoring
 from mlrun.model_monitoring.helpers import get_stream_path
+from mlrun.serving import GraphContext
 from mlrun.serving.utils import StepToDict
 from mlrun.utils import logger
 
@@ -62,7 +59,7 @@ class _PushToMonitoringWriter(StepToDict):
         self,
         event: tuple[
             list[
-                typing.Union[
+                Union[
                     ModelMonitoringApplicationResult, ModelMonitoringApplicationMetric
                 ]
             ],
@@ -121,50 +118,35 @@ class _PushToMonitoringWriter(StepToDict):
 
 
 class _PrepareMonitoringEvent(StepToDict):
-    def __init__(self, application_name: str):
+    def __init__(self, context: GraphContext, application_name: str) -> None:
         """
         Class for preparing the application event for the application step.
 
         :param application_name: Application name.
         """
+        self.graph_context = context
+        self.application_name = application_name
+        self.model_endpoints: dict[str, mlrun.model_monitoring.ModelEndpoint] = {}
 
-        self.context = self._create_mlrun_context(application_name)
-        self.model_endpoints = {}
-
-    def do(self, event: dict[str, dict]) -> MonitoringApplicationContext:
+    def do(self, event: dict[str, Any]) -> MonitoringApplicationContext:
         """
         Prepare the application event for the application step.
 
         :param event: Application event.
-        :return: Application event.
+        :return: Application context.
         """
-        if not event.get("mlrun_context"):
-            application_context = MonitoringApplicationContext().from_dict(
-                event,
-                context=self.context,
-                model_endpoint_dict=self.model_endpoints,
-            )
-        else:
-            application_context = MonitoringApplicationContext().from_dict(event)
+        application_context = MonitoringApplicationContext(
+            graph_context=self.graph_context,
+            application_name=self.application_name,
+            event=event,
+            model_endpoint_dict=self.model_endpoints,
+        )
+
         self.model_endpoints.setdefault(
             application_context.endpoint_id, application_context.model_endpoint
         )
-        return application_context
 
-    @staticmethod
-    def _create_mlrun_context(app_name: str):
-        artifact_path = mlrun.utils.helpers.template_artifact_path(
-            mlrun.mlconf.artifact_path, mlrun.mlconf.default_project
-        )
-        context = mlrun.get_or_create_ctx(
-            f"{app_name}-logger",
-            spec={
-                "metadata": {"labels": {"kind": mlrun.runtimes.RuntimeKinds.serving}},
-                "spec": {mlrun.utils.helpers.RunKeys.output_path: artifact_path},
-            },
-        )
-        context.__class__ = MonitoringApplicationContext
-        return context
+        return application_context
 
 
 class _ApplicationErrorHandler(StepToDict):
@@ -181,13 +163,13 @@ class _ApplicationErrorHandler(StepToDict):
 
         logger.error(f"Error in application step: {event}")
 
-        event_data = mlrun.common.schemas.Event(
+        event_data = alert_objects.Event(
             kind=alert_objects.EventKind.MM_APP_FAILED,
-            entity={
-                "kind": alert_objects.EventEntityKind.MODEL_MONITORING_APPLICATION,
-                "project": self.project,
-                "ids": [f"{self.project}_{event.body.application_name}"],
-            },
+            entity=alert_objects.EventEntities(
+                kind=alert_objects.EventEntityKind.MODEL_MONITORING_APPLICATION,
+                project=self.project,
+                ids=[f"{self.project}_{event.body.application_name}"],
+            ),
             value_dict={
                 "Error": event.error,
                 "Timestamp": event.timestamp,

mlrun/model_monitoring/applications/context.py

@@ -11,19 +11,22 @@
 # 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 json
-import typing
+import socket
+from typing import Any, Optional, cast
 
 import numpy as np
 import pandas as pd
 
-import mlrun.common.helpers
-import mlrun.common.model_monitoring.helpers
+import mlrun.common.constants as mlrun_constants
 import mlrun.common.schemas.model_monitoring.constants as mm_constants
 import mlrun.feature_store as fstore
-from mlrun.artifacts.model import ModelArtifact, get_model
+import mlrun.features
+import mlrun.serving
+import mlrun.utils
+from mlrun.artifacts import Artifact, DatasetArtifact, ModelArtifact, get_model
 from mlrun.common.model_monitoring.helpers import FeatureStats, pad_features_hist
-from mlrun.execution import MLClientCtx
 from mlrun.model_monitoring.helpers import (
     calculate_inputs_statistics,
     get_endpoint_record,
@@ -31,13 +34,17 @@ from mlrun.model_monitoring.helpers import (
 from mlrun.model_monitoring.model_endpoint import ModelEndpoint
 
 
-class MonitoringApplicationContext(MLClientCtx):
+class MonitoringApplicationContext:
     """
     The monitoring context holds all the relevant information for the monitoring application,
     and also it can be used for logging artifacts and results.
     The monitoring context has the following attributes:
 
-    :param application_name:        (str) the app name
+    :param application_name:        (str) The model monitoring application name.
+    :param project_name:            (str) The project name.
+    :param project:                 (MlrunProject) The project object.
+    :param logger:                  (mlrun.utils.Logger) MLRun logger.
+    :param nuclio_logger:           (nuclio.request.Logger) Nuclio logger.
     :param sample_df_stats:         (FeatureStats) The new sample distribution dictionary.
     :param feature_stats:           (FeatureStats) The train sample distribution dictionary.
     :param sample_df:               (pd.DataFrame) The new sample DataFrame.
@@ -49,79 +56,89 @@ class MonitoringApplicationContext(MLClientCtx):
     :param model_endpoint:          (ModelEndpoint) The model endpoint object.
     :param feature_names:           (list[str]) List of models feature names.
     :param label_names:             (list[str]) List of models label names.
-    :param model:                   (tuple[str, ModelArtifact, dict]) The model file, model spec object, and list of
-
+    :param model:                   (tuple[str, ModelArtifact, dict]) The model file, model spec object,
+                                    and a list of extra data items.
     """
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def _enrich_data(self):
-        self.application_name: typing.Optional[str] = None
-        self.start_infer_time: typing.Optional[pd.Timestamp] = None
-        self.end_infer_time: typing.Optional[pd.Timestamp] = None
-        self.latest_request: typing.Optional[pd.Timestamp] = None
-        self.endpoint_id: typing.Optional[str] = None
-        self.output_stream_uri: typing.Optional[str] = None
-
-        self._sample_df: typing.Optional[pd.DataFrame] = None
-        self._model_endpoint: typing.Optional[ModelEndpoint] = None
-        self._feature_stats: typing.Optional[FeatureStats] = None
-        self._sample_df_stats: typing.Optional[FeatureStats] = None
-
-    @classmethod
-    def from_dict(
-        cls,
-        attrs: dict,
-        context=None,
-        model_endpoint_dict=None,
-        **kwargs,
-    ) -> "MonitoringApplicationContext":
+    def __init__(
+        self,
+        *,
+        graph_context: mlrun.serving.GraphContext,
+        application_name: str,
+        event: dict[str, Any],
+        model_endpoint_dict: dict[str, ModelEndpoint],
+    ) -> None:
         """
-        Create an instance of the MonitoringApplicationContext from a dictionary.
+        Initialize a `MonitoringApplicationContext` object.
+        Note: this object should not be instantiated manually.
 
-        :param attrs:               The instance data dictionary.
-        :param context:             The current application context.
+        :param application_name:    The application name.
+        :param event:               The instance data dictionary.
         :param model_endpoint_dict: Dictionary of model endpoints.
-
         """
+        self.application_name = application_name
 
-        if not context:
-            ctx = (
-                super().from_dict(
-                    attrs=attrs.get(mm_constants.ApplicationEvent.MLRUN_CONTEXT, {}),
-                    **kwargs,
-                ),
-            )
-        else:
-            ctx = context
-            cls._enrich_data(ctx)
+        self.project_name = graph_context.project
+        self.project = mlrun.load_project(url=self.project_name)
+
+        # MLRun Logger
+        self.logger = mlrun.utils.create_logger(
+            level=mlrun.mlconf.log_level,
+            formatter_kind=mlrun.mlconf.log_formatter,
+            name="monitoring-application",
+        )
+        # Nuclio logger - `nuclio.request.Logger`.
+        # Note: this logger does not accept keyword arguments.
+        self.nuclio_logger = graph_context.logger
 
-        ctx.start_infer_time = pd.Timestamp(
-            attrs.get(mm_constants.ApplicationEvent.START_INFER_TIME)
+        # event data
+        self.start_infer_time = pd.Timestamp(
+            cast(str, event.get(mm_constants.ApplicationEvent.START_INFER_TIME))
         )
-        ctx.end_infer_time = pd.Timestamp(
-            attrs.get(mm_constants.ApplicationEvent.END_INFER_TIME)
+        self.end_infer_time = pd.Timestamp(
+            cast(str, event.get(mm_constants.ApplicationEvent.END_INFER_TIME))
         )
-        ctx.latest_request = pd.Timestamp(
-            attrs.get(mm_constants.ApplicationEvent.LAST_REQUEST)
+        self.latest_request = pd.Timestamp(
+            cast(str, event.get(mm_constants.ApplicationEvent.LAST_REQUEST))
         )
-        ctx.application_name = attrs.get(mm_constants.ApplicationEvent.APPLICATION_NAME)
-        ctx._feature_stats = json.loads(
-            attrs.get(mm_constants.ApplicationEvent.FEATURE_STATS, "{}")
+        self.endpoint_id = cast(
+            str, event.get(mm_constants.ApplicationEvent.ENDPOINT_ID)
         )
-        ctx._sample_df_stats = json.loads(
-            attrs.get(mm_constants.ApplicationEvent.CURRENT_STATS, "{}")
+        self.output_stream_uri = cast(
+            str, event.get(mm_constants.ApplicationEvent.OUTPUT_STREAM_URI)
         )
 
-        ctx.endpoint_id = attrs.get(mm_constants.ApplicationEvent.ENDPOINT_ID)
-        ctx._model_endpoint = model_endpoint_dict.get(ctx.endpoint_id)
+        self._feature_stats: Optional[FeatureStats] = json.loads(
+            event.get(mm_constants.ApplicationEvent.FEATURE_STATS, "{}")
+        )
+        self._sample_df_stats: Optional[FeatureStats] = json.loads(
+            event.get(mm_constants.ApplicationEvent.CURRENT_STATS, "{}")
+        )
 
-        return ctx
+        # Default labels for the artifacts
+        self._default_labels = self._get_default_labels()
+
+        # Persistent data - fetched when needed
+        self._sample_df: Optional[pd.DataFrame] = None
+        self._model_endpoint: Optional[ModelEndpoint] = model_endpoint_dict.get(
+            self.endpoint_id
+        )
+
+    def _get_default_labels(self) -> dict[str, str]:
+        return {
+            mlrun_constants.MLRunInternalLabels.runner_pod: socket.gethostname(),
+            mlrun_constants.MLRunInternalLabels.producer_type: "model-monitoring-app",
+            mlrun_constants.MLRunInternalLabels.app_name: self.application_name,
+            mlrun_constants.MLRunInternalLabels.endpoint_id: self.endpoint_id,
+        }
+
+    def _add_default_labels(self, labels: Optional[dict[str, str]]) -> dict[str, str]:
+        """Add the default labels to logged artifacts labels"""
+        return (labels or {}) | self._default_labels
 
     @property
     def sample_df(self) -> pd.DataFrame:
-        if not hasattr(self, "_sample_df") or self._sample_df is None:
+        if self._sample_df is None:
             feature_set = fstore.get_feature_set(
                 self.model_endpoint.status.monitoring_feature_set_uri
             )
@@ -144,15 +161,15 @@ class MonitoringApplicationContext(MLClientCtx):
 
     @property
     def model_endpoint(self) -> ModelEndpoint:
-        if not hasattr(self, "_model_endpoint") or not self._model_endpoint:
+        if not self._model_endpoint:
             self._model_endpoint = ModelEndpoint.from_flat_dict(
-                get_endpoint_record(self.project, self.endpoint_id)
+                get_endpoint_record(self.project_name, self.endpoint_id)
             )
         return self._model_endpoint
 
     @property
     def feature_stats(self) -> FeatureStats:
-        if not hasattr(self, "_feature_stats") or not self._feature_stats:
+        if not self._feature_stats:
             self._feature_stats = json.loads(self.model_endpoint.status.feature_stats)
             pad_features_hist(self._feature_stats)
         return self._feature_stats
@@ -160,7 +177,7 @@ class MonitoringApplicationContext(MLClientCtx):
     @property
     def sample_df_stats(self) -> FeatureStats:
         """statistics of the sample dataframe"""
-        if not hasattr(self, "_sample_df_stats") or not self._sample_df_stats:
+        if not self._sample_df_stats:
             self._sample_df_stats = calculate_inputs_statistics(
                 self.feature_stats, self.sample_df
             )
@@ -184,13 +201,11 @@ class MonitoringApplicationContext(MLClientCtx):
 
     @property
     def model(self) -> tuple[str, ModelArtifact, dict]:
-        """return model file, model spec object, and list of extra data items"""
+        """The model file, model spec object, and a list of extra data items"""
         return get_model(self.model_endpoint.spec.model_uri)
 
     @staticmethod
-    def dict_to_histogram(
-        histogram_dict: mlrun.common.model_monitoring.helpers.FeatureStats,
-    ) -> pd.DataFrame:
+    def dict_to_histogram(histogram_dict: FeatureStats) -> pd.DataFrame:
         """
         Convert histogram dictionary to pandas DataFrame with feature histograms as columns
 
@@ -210,3 +225,124 @@ class MonitoringApplicationContext(MLClientCtx):
         histograms = pd.DataFrame(histograms)
 
         return histograms
+
+    def log_artifact(
+        self,
+        item,
+        body=None,
+        tag: str = "",
+        local_path: str = "",
+        artifact_path: Optional[str] = None,
+        format: Optional[str] = None,
+        upload: Optional[bool] = None,
+        labels: Optional[dict[str, str]] = None,
+        target_path: Optional[str] = None,
+        **kwargs,
+    ) -> Artifact:
+        """
+        Log an artifact.
+        See :func:`~mlrun.projects.MlrunProject.log_artifact` for the documentation.
+        """
+        labels = self._add_default_labels(labels)
+        return self.project.log_artifact(
+            item,
+            body=body,
+            tag=tag,
+            local_path=local_path,
+            artifact_path=artifact_path,
+            format=format,
+            upload=upload,
+            labels=labels,
+            target_path=target_path,
+            **kwargs,
+        )
+
+    def log_dataset(
+        self,
+        key,
+        df,
+        tag="",
+        local_path=None,
+        artifact_path=None,
+        upload=None,
+        labels=None,
+        format="",
+        preview=None,
+        stats=None,
+        target_path="",
+        extra_data=None,
+        label_column: Optional[str] = None,
+        **kwargs,
+    ) -> DatasetArtifact:
+        """
+        Log a dataset artifact.
+        See :func:`~mlrun.projects.MlrunProject.log_dataset` for the documentation.
+        """
+        labels = self._add_default_labels(labels)
+        return self.project.log_dataset(
+            key,
+            df,
+            tag=tag,
+            local_path=local_path,
+            artifact_path=artifact_path,
+            upload=upload,
+            labels=labels,
+            format=format,
+            preview=preview,
+            stats=stats,
+            target_path=target_path,
+            extra_data=extra_data,
+            label_column=label_column,
+            **kwargs,
+        )
+
+    def log_model(
+        self,
+        key,
+        body=None,
+        framework="",
+        tag="",
+        model_dir=None,
+        model_file=None,
+        algorithm=None,
+        metrics=None,
+        parameters=None,
+        artifact_path=None,
+        upload=None,
+        labels=None,
+        inputs: Optional[list[mlrun.features.Feature]] = None,
+        outputs: Optional[list[mlrun.features.Feature]] = None,
+        feature_vector: Optional[str] = None,
+        feature_weights: Optional[list] = None,
+        training_set=None,
+        label_column=None,
+        extra_data=None,
+        **kwargs,
+    ) -> ModelArtifact:
+        """
+        Log a model artifact.
+        See :func:`~mlrun.projects.MlrunProject.log_model` for the documentation.
+        """
+        labels = self._add_default_labels(labels)
+        return self.project.log_model(
+            key,
+            body=body,
+            framework=framework,
+            tag=tag,
+            model_dir=model_dir,
+            model_file=model_file,
+            algorithm=algorithm,
+            metrics=metrics,
+            parameters=parameters,
+            artifact_path=artifact_path,
+            upload=upload,
+            labels=labels,
+            inputs=inputs,
+            outputs=outputs,
+            feature_vector=feature_vector,
+            feature_weights=feature_weights,
+            training_set=training_set,
+            label_column=label_column,
+            extra_data=extra_data,
+            **kwargs,
+        )

mlrun/model_monitoring/controller.py

@@ -596,7 +596,6 @@ class MonitoringApplicationController:
                 project=project,
                 function_name=mm_constants.MonitoringFunctionNames.WRITER,
             ),
-            mm_constants.ApplicationEvent.MLRUN_CONTEXT: {},  # TODO : for future use by ad-hoc batch infer
         }
         for app_name in applications_names:
             data.update({mm_constants.ApplicationEvent.APPLICATION_NAME: app_name})

mlrun/model_monitoring/db/stores/sqldb/sql_store.py

@@ -20,7 +20,7 @@ import pandas as pd
 import sqlalchemy
 import sqlalchemy.exc
 import sqlalchemy.orm
-from sqlalchemy.engine import make_url
+from sqlalchemy.engine import Engine, make_url
 from sqlalchemy.sql.elements import BinaryExpression
 
 import mlrun.common.model_monitoring.helpers
@@ -61,9 +61,15 @@ class SQLStoreBase(StoreBase):
             )
 
         self._sql_connection_string = kwargs.get("store_connection_string")
-        self._engine = get_engine(dsn=self._sql_connection_string)
+        self._engine = None
         self._init_tables()
 
+    @property
+    def engine(self) -> Engine:
+        if not self._engine:
+            self._engine = get_engine(dsn=self._sql_connection_string)
+        return self._engine
+
     def create_tables(self):
         self._create_tables_if_not_exist()
 
@@ -116,7 +122,7 @@ class SQLStoreBase(StoreBase):
         :param table_name: Target table name.
         :param event:      Event dictionary that will be written into the DB.
         """
-        with self._engine.connect() as connection:
+        with self.engine.connect() as connection:
             # Convert the result into a pandas Dataframe and write it into the database
             event_df = pd.DataFrame([event])
             event_df.to_sql(table_name, con=connection, index=False, if_exists="append")
@@ -177,7 +183,7 @@ class SQLStoreBase(StoreBase):
         param table:     SQLAlchemy declarative table.
         :param criteria: A list of binary expressions that filter the query.
         """
-        if not self._engine.has_table(table.__tablename__):
+        if not self.engine.has_table(table.__tablename__):
             logger.debug(
                 f"Table {table.__tablename__} does not exist in the database. Skipping deletion."
             )
@@ -524,9 +530,9 @@ class SQLStoreBase(StoreBase):
         for table in self._tables:
             # Create table if not exist. The `metadata` contains the `ModelEndpointsTable`
             db_name = make_url(self._sql_connection_string).database
-            if not self._engine.has_table(table):
+            if not self.engine.has_table(table):
                 logger.info(f"Creating table {table} on {db_name} db.")
-                self._tables[table].metadata.create_all(bind=self._engine)
+                self._tables[table].metadata.create_all(bind=self.engine)
             else:
                 logger.info(f"Table {table} already exists on {db_name} db.")
 
@@ -574,8 +580,11 @@ class SQLStoreBase(StoreBase):
         """
         Delete all the model monitoring resources of the project in the SQL tables.
         """
+        logger.debug(
+            "Deleting model monitoring endpoints resources from the SQL tables",
+            project=self.project,
+        )
         endpoints = self.list_model_endpoints()
-        logger.debug("Deleting model monitoring resources", project=self.project)
 
         for endpoint_dict in endpoints:
             endpoint_id = endpoint_dict[mm_schemas.EventFieldType.UID]
@@ -612,7 +621,7 @@ class SQLStoreBase(StoreBase):
 
         # Note: the block below does not use self._get, as we need here all the
         # results, not only `one_or_none`.
-        with sqlalchemy.orm.Session(self._engine) as session:
+        with sqlalchemy.orm.Session(self.engine) as session:
             metric_rows = (
                 session.query(table)  # pyright: ignore[reportOptionalCall]
                 .filter(table.endpoint_id == endpoint_id)

mlrun/model_monitoring/db/stores/v3io_kv/kv_store.py

@@ -20,6 +20,7 @@ from http import HTTPStatus
 import v3io.dataplane
 import v3io.dataplane.output
 import v3io.dataplane.response
+from v3io.dataplane import Client as V3IOClient
 
 import mlrun.common.model_monitoring.helpers
 import mlrun.common.schemas.model_monitoring as mm_schemas
@@ -100,13 +101,18 @@ class KVStoreBase(StoreBase):
         project: str,
     ) -> None:
         super().__init__(project=project)
-        # Initialize a V3IO client instance
-        self.client = mlrun.utils.v3io_clients.get_v3io_client(
-            endpoint=mlrun.mlconf.v3io_api,
-        )
+        self._client = None
         # Get the KV table path and container
         self.path, self.container = self._get_path_and_container()
 
+    @property
+    def client(self) -> V3IOClient:
+        if not self._client:
+            self._client = mlrun.utils.v3io_clients.get_v3io_client(
+                endpoint=mlrun.mlconf.v3io_api,
+            )
+        return self._client
+
     def write_model_endpoint(self, endpoint: dict[str, typing.Any]):
         """
         Create a new endpoint record in the KV table.
@@ -285,6 +291,10 @@ class KVStoreBase(StoreBase):
         """
         Delete all model endpoints resources in V3IO KV.
         """
+        logger.debug(
+            "Deleting model monitoring endpoints resources in V3IO KV",
+            project=self.project,
+        )
 
         endpoints = self.list_model_endpoints()
 

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

@@ -47,10 +47,17 @@ class TDEngineConnector(TSDBConnector):
             )
         self._tdengine_connection_string = kwargs.get("connection_string")
         self.database = database
-        self._connection = self._create_connection()
+
+        self._connection = None
         self._init_super_tables()
 
-    def _create_connection(self):
+    @property
+    def connection(self) -> taosws.Connection:
+        if not self._connection:
+            self._connection = self._create_connection()
+        return self._connection
+
+    def _create_connection(self) -> taosws.Connection:
         """Establish a connection to the TSDB server."""
         conn = taosws.connect(self._tdengine_connection_string)
         try:
@@ -176,7 +183,8 @@ class TDEngineConnector(TSDBConnector):
                     mm_schemas.EventFieldType.PROJECT,
                     mm_schemas.EventFieldType.ENDPOINT_ID,
                 ],
-                max_events=10,
+                max_events=1000,
+                flush_after_seconds=30,
             )
 
         apply_process_before_tsdb()