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

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/__init__.py

@@ -63,7 +63,7 @@ class ObjectStoreFactory(enum.Enum):
         :param value: Provided enum (invalid) value.
         """
         valid_values = list(cls.__members__.keys())
-        raise mlrun.errors.MLRunInvalidMMStoreType(
+        raise mlrun.errors.MLRunInvalidMMStoreTypeError(
             f"{value} is not a valid endpoint store, please choose a valid value: %{valid_values}."
         )
 
@@ -101,7 +101,7 @@ def get_store_object(
 
     :return: `StoreBase` object. Using this object, the user can apply different operations such as write, update, get
              and delete a model endpoint record.
-    :raise: `MLRunInvalidMMStoreType` if the user didn't provide store connection
+    :raise: `MLRunInvalidMMStoreTypeError` if the user didn't provide store connection
              or the provided store connection is invalid.
     """
 
@@ -123,7 +123,7 @@ def get_store_object(
             mlrun.common.schemas.model_monitoring.ModelEndpointTarget.V3IO_NOSQL
         )
     else:
-        raise mlrun.errors.MLRunInvalidMMStoreType(
+        raise mlrun.errors.MLRunInvalidMMStoreTypeError(
             "You must provide a valid store connection by using "
             "set_model_monitoring_credentials API."
         )

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/__init__.py

@@ -57,7 +57,7 @@ class ObjectTSDBFactory(enum.Enum):
         :param value: Provided enum (invalid) value.
         """
         valid_values = list(cls.__members__.keys())
-        raise mlrun.errors.MLRunInvalidMMStoreType(
+        raise mlrun.errors.MLRunInvalidMMStoreTypeError(
             f"{value} is not a valid tsdb, please choose a valid value: %{valid_values}."
         )
 
@@ -76,7 +76,7 @@ def get_tsdb_connector(
 
     :return: `TSDBConnector` object. The main goal of this object is to handle different operations on the
              TSDB connector such as updating drift metrics or write application record result.
-    :raise: `MLRunInvalidMMStoreType` if the user didn't provide TSDB connection
+    :raise: `MLRunInvalidMMStoreTypeError` if the user didn't provide TSDB connection
             or the provided TSDB connection is invalid.
     """
 
@@ -93,7 +93,7 @@ def get_tsdb_connector(
     elif tsdb_connection_string and tsdb_connection_string == "v3io":
         tsdb_connector_type = mlrun.common.schemas.model_monitoring.TSDBTarget.V3IO_TSDB
     else:
-        raise mlrun.errors.MLRunInvalidMMStoreType(
+        raise mlrun.errors.MLRunInvalidMMStoreTypeError(
             "You must provide a valid tsdb store connection by using "
             "set_model_monitoring_credentials API."
         )

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:
@@ -61,7 +68,7 @@ class TDEngineConnector(TSDBConnector):
         try:
             conn.execute(f"USE {self.database}")
         except taosws.QueryError as e:
-            raise mlrun.errors.MLRunTSDBConnectionFailure(
+            raise mlrun.errors.MLRunTSDBConnectionFailureError(
                 f"Failed to use TDEngine database {self.database}, {mlrun.errors.err_to_str(e)}"
             )
         return conn
@@ -84,7 +91,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.execute(create_table_query)
 
     def write_application_event(
         self,
@@ -128,10 +135,10 @@ class TDEngineConnector(TSDBConnector):
         create_table_query = table._create_subtable_query(
             subtable=table_name, values=event
         )
-        self._connection.execute(create_table_query)
+        self.connection.execute(create_table_query)
 
         insert_statement = table._insert_subtable_query(
-            self._connection,
+            self.connection,
             subtable=table_name,
             values=event,
         )
@@ -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()
@@ -196,12 +204,12 @@ class TDEngineConnector(TSDBConnector):
             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.query(get_subtable_names_query)
             for subtable in subtables:
                 drop_query = self.tables[table]._drop_subtable_query(
                     subtable=subtable[0]
                 )
-                self._connection.execute(drop_query)
+                self.connection.execute(drop_query)
         logger.info(
             f"Deleted all project resources in the TSDB connector for project {self.project}"
         )
@@ -273,7 +281,7 @@ class TDEngineConnector(TSDBConnector):
             database=self.database,
         )
         try:
-            query_result = self._connection.query(full_query)
+            query_result = self.connection.query(full_query)
         except taosws.QueryError as e:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 f"Failed to query table {table} in database {self.database}, {str(e)}"