mlrun 1.4.0rc25__py3-none-any.whl → 1.5.0rc2__py3-none-any.whl

mlrun/model_monitoring/stores/models/sqlite.py

@@ -16,83 +16,101 @@
 
 from sqlalchemy import TIMESTAMP, Boolean, Column, Integer, String, Text
 
-import mlrun.common.model_monitoring as model_monitoring_constants
+import mlrun.common.schemas.model_monitoring
 from mlrun.utils.db import BaseModel
 
 from .base import Base
 
 
 class ModelEndpointsTable(Base, BaseModel):
-    __tablename__ = model_monitoring_constants.EventFieldType.MODEL_ENDPOINTS
+    __tablename__ = mlrun.common.schemas.model_monitoring.EventFieldType.MODEL_ENDPOINTS
 
     uid = Column(
-        model_monitoring_constants.EventFieldType.UID,
+        mlrun.common.schemas.model_monitoring.EventFieldType.UID,
         String(40),
         primary_key=True,
     )
-    state = Column(model_monitoring_constants.EventFieldType.STATE, String(10))
-    project = Column(model_monitoring_constants.EventFieldType.PROJECT, String(40))
+    state = Column(
+        mlrun.common.schemas.model_monitoring.EventFieldType.STATE, String(10)
+    )
+    project = Column(
+        mlrun.common.schemas.model_monitoring.EventFieldType.PROJECT, String(40)
+    )
     function_uri = Column(
-        model_monitoring_constants.EventFieldType.FUNCTION_URI,
+        mlrun.common.schemas.model_monitoring.EventFieldType.FUNCTION_URI,
         String(255),
     )
-    model = Column(model_monitoring_constants.EventFieldType.MODEL, String(255))
+    model = Column(
+        mlrun.common.schemas.model_monitoring.EventFieldType.MODEL, String(255)
+    )
     model_class = Column(
-        model_monitoring_constants.EventFieldType.MODEL_CLASS,
+        mlrun.common.schemas.model_monitoring.EventFieldType.MODEL_CLASS,
         String(255),
     )
-    labels = Column(model_monitoring_constants.EventFieldType.LABELS, Text)
-    model_uri = Column(model_monitoring_constants.EventFieldType.MODEL_URI, String(255))
-    stream_path = Column(model_monitoring_constants.EventFieldType.STREAM_PATH, Text)
+    labels = Column(mlrun.common.schemas.model_monitoring.EventFieldType.LABELS, Text)
+    model_uri = Column(
+        mlrun.common.schemas.model_monitoring.EventFieldType.MODEL_URI, String(255)
+    )
+    stream_path = Column(
+        mlrun.common.schemas.model_monitoring.EventFieldType.STREAM_PATH, Text
+    )
     algorithm = Column(
-        model_monitoring_constants.EventFieldType.ALGORITHM,
+        mlrun.common.schemas.model_monitoring.EventFieldType.ALGORITHM,
         String(255),
     )
-    active = Column(model_monitoring_constants.EventFieldType.ACTIVE, Boolean)
+    active = Column(
+        mlrun.common.schemas.model_monitoring.EventFieldType.ACTIVE, Boolean
+    )
     monitoring_mode = Column(
-        model_monitoring_constants.EventFieldType.MONITORING_MODE,
+        mlrun.common.schemas.model_monitoring.EventFieldType.MONITORING_MODE,
         String(10),
     )
     feature_stats = Column(
-        model_monitoring_constants.EventFieldType.FEATURE_STATS, Text
+        mlrun.common.schemas.model_monitoring.EventFieldType.FEATURE_STATS, Text
     )
     current_stats = Column(
-        model_monitoring_constants.EventFieldType.CURRENT_STATS, Text
+        mlrun.common.schemas.model_monitoring.EventFieldType.CURRENT_STATS, Text
     )
     feature_names = Column(
-        model_monitoring_constants.EventFieldType.FEATURE_NAMES, Text
+        mlrun.common.schemas.model_monitoring.EventFieldType.FEATURE_NAMES, Text
+    )
+    children = Column(
+        mlrun.common.schemas.model_monitoring.EventFieldType.CHILDREN, Text
+    )
+    label_names = Column(
+        mlrun.common.schemas.model_monitoring.EventFieldType.LABEL_NAMES, Text
     )
-    children = Column(model_monitoring_constants.EventFieldType.CHILDREN, Text)
-    label_names = Column(model_monitoring_constants.EventFieldType.LABEL_NAMES, Text)
     endpoint_type = Column(
-        model_monitoring_constants.EventFieldType.ENDPOINT_TYPE,
+        mlrun.common.schemas.model_monitoring.EventFieldType.ENDPOINT_TYPE,
         String(10),
     )
     children_uids = Column(
-        model_monitoring_constants.EventFieldType.CHILDREN_UIDS, Text
+        mlrun.common.schemas.model_monitoring.EventFieldType.CHILDREN_UIDS, Text
     )
     drift_measures = Column(
-        model_monitoring_constants.EventFieldType.DRIFT_MEASURES, Text
+        mlrun.common.schemas.model_monitoring.EventFieldType.DRIFT_MEASURES, Text
     )
     drift_status = Column(
-        model_monitoring_constants.EventFieldType.DRIFT_STATUS,
+        mlrun.common.schemas.model_monitoring.EventFieldType.DRIFT_STATUS,
         String(40),
     )
     monitor_configuration = Column(
-        model_monitoring_constants.EventFieldType.MONITOR_CONFIGURATION,
+        mlrun.common.schemas.model_monitoring.EventFieldType.MONITOR_CONFIGURATION,
         Text,
     )
     monitoring_feature_set_uri = Column(
-        model_monitoring_constants.EventFieldType.FEATURE_SET_URI,
+        mlrun.common.schemas.model_monitoring.EventFieldType.FEATURE_SET_URI,
         String(255),
     )
     first_request = Column(
-        model_monitoring_constants.EventFieldType.FIRST_REQUEST,
+        mlrun.common.schemas.model_monitoring.EventFieldType.FIRST_REQUEST,
         TIMESTAMP,
     )
     last_request = Column(
-        model_monitoring_constants.EventFieldType.LAST_REQUEST,
+        mlrun.common.schemas.model_monitoring.EventFieldType.LAST_REQUEST,
         TIMESTAMP,
     )
-    error_count = Column(model_monitoring_constants.EventFieldType.ERROR_COUNT, Integer)
-    metrics = Column(model_monitoring_constants.EventFieldType.METRICS, Text)
+    error_count = Column(
+        mlrun.common.schemas.model_monitoring.EventFieldType.ERROR_COUNT, Integer
+    )
+    metrics = Column(mlrun.common.schemas.model_monitoring.EventFieldType.METRICS, Text)

mlrun/model_monitoring/stores/sql_model_endpoint_store.py

@@ -20,12 +20,10 @@ from datetime import datetime, timezone
 import pandas as pd
 import sqlalchemy as db
 
-import mlrun
-import mlrun.common.model_monitoring as model_monitoring_constants
-import mlrun.model_monitoring.model_endpoint
-import mlrun.utils.model_monitoring
-import mlrun.utils.v3io_clients
-from mlrun.api.db.sqldb.session import create_session, get_engine
+import mlrun.common.model_monitoring.helpers
+import mlrun.common.schemas.model_monitoring
+import mlrun.model_monitoring.helpers
+from mlrun.common.db.sql_session import create_session, get_engine
 from mlrun.utils import logger
 
 from .model_endpoint_store import ModelEndpointStore
@@ -47,22 +45,28 @@ class SQLModelEndpointStore(ModelEndpointStore):
         self,
         project: str,
         sql_connection_string: str = None,
+        secret_provider: typing.Callable = None,
     ):
         """
         Initialize SQL store target object.
 
         :param project:               The name of the project.
         :param sql_connection_string: Valid connection string or a path to SQL database with model endpoints table.
+        :param secret_provider:       An optional secret provider to get the connection string secret.
         """
 
         super().__init__(project=project)
 
         self.sql_connection_string = (
             sql_connection_string
-            or mlrun.utils.model_monitoring.get_connection_string(project=self.project)
+            or mlrun.model_monitoring.helpers.get_connection_string(
+                secret_provider=secret_provider
+            )
         )
 
-        self.table_name = model_monitoring_constants.EventFieldType.MODEL_ENDPOINTS
+        self.table_name = (
+            mlrun.common.schemas.model_monitoring.EventFieldType.MODEL_ENDPOINTS
+        )
 
         self._engine = get_engine(dsn=self.sql_connection_string)
         self.ModelEndpointsTable = get_ModelEndpointsTable(
@@ -84,10 +88,10 @@ class SQLModelEndpointStore(ModelEndpointStore):
         with self._engine.connect() as connection:
             # Adjust timestamps fields
             endpoint[
-                model_monitoring_constants.EventFieldType.FIRST_REQUEST
+                mlrun.common.schemas.model_monitoring.EventFieldType.FIRST_REQUEST
             ] = datetime.now(timezone.utc)
             endpoint[
-                model_monitoring_constants.EventFieldType.LAST_REQUEST
+                mlrun.common.schemas.model_monitoring.EventFieldType.LAST_REQUEST
             ] = datetime.now(timezone.utc)
 
             # Convert the result into a pandas Dataframe and write it into the database
@@ -112,7 +116,9 @@ class SQLModelEndpointStore(ModelEndpointStore):
         # Update the model endpoint record using sqlalchemy ORM
         with create_session(dsn=self.sql_connection_string) as session:
             # Remove endpoint id (foreign key) from the update query
-            attributes.pop(model_monitoring_constants.EventFieldType.ENDPOINT_ID, None)
+            attributes.pop(
+                mlrun.common.schemas.model_monitoring.EventFieldType.ENDPOINT_ID, None
+            )
 
             # Generate and commit the update session query
             session.query(self.ModelEndpointsTable).filter(
@@ -201,32 +207,36 @@ class SQLModelEndpointStore(ModelEndpointStore):
                 query = self._filter_values(
                     query=query,
                     model_endpoints_table=self.model_endpoints_table,
-                    key_filter=model_monitoring_constants.EventFieldType.MODEL,
+                    key_filter=mlrun.common.schemas.model_monitoring.EventFieldType.MODEL,
                     filtered_values=[model],
                 )
             if function:
                 query = self._filter_values(
                     query=query,
                     model_endpoints_table=self.model_endpoints_table,
-                    key_filter=model_monitoring_constants.EventFieldType.FUNCTION,
+                    key_filter=mlrun.common.schemas.model_monitoring.EventFieldType.FUNCTION,
                     filtered_values=[function],
                 )
             if uids:
                 query = self._filter_values(
                     query=query,
                     model_endpoints_table=self.model_endpoints_table,
-                    key_filter=model_monitoring_constants.EventFieldType.UID,
+                    key_filter=mlrun.common.schemas.model_monitoring.EventFieldType.UID,
                     filtered_values=uids,
                     combined=False,
                 )
             if top_level:
-                node_ep = str(mlrun.common.model_monitoring.EndpointType.NODE_EP.value)
-                router_ep = str(mlrun.common.model_monitoring.EndpointType.ROUTER.value)
+                node_ep = str(
+                    mlrun.common.schemas.model_monitoring.EndpointType.NODE_EP.value
+                )
+                router_ep = str(
+                    mlrun.common.schemas.model_monitoring.EndpointType.ROUTER.value
+                )
                 endpoint_types = [node_ep, router_ep]
                 query = self._filter_values(
                     query=query,
                     model_endpoints_table=self.model_endpoints_table,
-                    key_filter=model_monitoring_constants.EventFieldType.ENDPOINT_TYPE,
+                    key_filter=mlrun.common.schemas.model_monitoring.EventFieldType.ENDPOINT_TYPE,
                     filtered_values=endpoint_types,
                     combined=False,
                 )
@@ -303,7 +313,9 @@ class SQLModelEndpointStore(ModelEndpointStore):
 
         # Convert endpoint labels into dictionary
         endpoint_labels = json.loads(
-            endpoint_dict.get(model_monitoring_constants.EventFieldType.LABELS)
+            endpoint_dict.get(
+                mlrun.common.schemas.model_monitoring.EventFieldType.LABELS
+            )
         )
 
         for label in labels:
@@ -331,7 +343,7 @@ class SQLModelEndpointStore(ModelEndpointStore):
         for endpoint_dict in endpoints:
             # Delete model endpoint record from SQL table
             self.delete_model_endpoint(
-                endpoint_dict[model_monitoring_constants.EventFieldType.UID],
+                endpoint_dict[mlrun.common.schemas.model_monitoring.EventFieldType.UID],
             )
 
     def get_endpoint_real_time_metrics(