mlrun 1.7.0rc21__py3-none-any.whl → 1.7.0rc23__py3-none-any.whl

mlrun/alerts/alert.py

@@ -30,6 +30,11 @@ class AlertConfig(ModelObj):
         "reset_policy",
         "state",
     ]
+    _fields_to_serialize = ModelObj._fields_to_serialize + [
+        "entities",
+        "notifications",
+        "trigger",
+    ]
 
     def __init__(
         self,
@@ -71,24 +76,44 @@ class AlertConfig(ModelObj):
         if not self.project or not self.name:
             raise mlrun.errors.MLRunBadRequestError("Project and name must be provided")
 
+    def _serialize_field(
+        self, struct: dict, field_name: str = None, strip: bool = False
+    ):
+        if field_name == "entities":
+            if self.entities:
+                return (
+                    self.entities.dict()
+                    if not isinstance(self.entities, dict)
+                    else self.entities
+                )
+            return None
+        if field_name == "notifications":
+            if self.notifications:
+                return [
+                    notification_data.dict()
+                    if not isinstance(notification_data, dict)
+                    else notification_data
+                    for notification_data in self.notifications
+                ]
+            return None
+        if field_name == "trigger":
+            if self.trigger:
+                return (
+                    self.trigger.dict()
+                    if not isinstance(self.trigger, dict)
+                    else self.trigger
+                )
+            return None
+        return super()._serialize_field(struct, field_name, strip)
+
     def to_dict(self, fields: list = None, exclude: list = None, strip: bool = False):
-        data = super().to_dict(self._dict_fields)
-
-        data["entities"] = (
-            self.entities.dict()
-            if not isinstance(self.entities, dict)
-            else self.entities
-        )
-        data["notifications"] = [
-            notification_data.dict()
-            if not isinstance(notification_data, dict)
-            else notification_data
-            for notification_data in self.notifications
-        ]
-        data["trigger"] = (
-            self.trigger.dict() if not isinstance(self.trigger, dict) else self.trigger
-        )
-        return data
+        if self.entities is None:
+            raise mlrun.errors.MLRunBadRequestError("Alert entity field is missing")
+        if not self.notifications:
+            raise mlrun.errors.MLRunBadRequestError(
+                "Alert must have at least one notification"
+            )
+        return super().to_dict(self._dict_fields)
 
     @classmethod
     def from_dict(cls, struct=None, fields=None, deprecated_fields: dict = None):

mlrun/common/schemas/__init__.py

@@ -82,6 +82,7 @@ from .events import (
 )
 from .feature_store import (
     EntitiesOutput,
+    EntitiesOutputV2,
     Entity,
     EntityListOutput,
     EntityRecord,
@@ -98,6 +99,7 @@ from .feature_store import (
     FeatureSetSpec,
     FeatureSetsTagsOutput,
     FeaturesOutput,
+    FeaturesOutputV2,
     FeatureVector,
     FeatureVectorRecord,
     FeatureVectorsOutput,

mlrun/common/schemas/feature_store.py

@@ -14,7 +14,7 @@
 #
 from typing import Optional
 
-from pydantic import BaseModel, Extra, Field
+import pydantic
 
 from .auth import AuthorizationResourceTypes, Credentials
 from .object import (
@@ -27,32 +27,62 @@ from .object import (
 )
 
 
-class Feature(BaseModel):
+class FeatureStoreBaseModel(pydantic.BaseModel):
+    """
+    Intermediate base class, in order to override pydantic's configuration, as per
+    https://docs.pydantic.dev/1.10/usage/model_config/#change-behaviour-globally
+    """
+
+    class Config:
+        copy_on_model_validation = "none"
+
+
+class Feature(FeatureStoreBaseModel):
+    name: str
+    value_type: str
+    labels: Optional[dict] = {}
+
+    class Config:
+        extra = pydantic.Extra.allow
+
+
+class QualifiedFeature(FeatureStoreBaseModel):
+    name: str
+    value_type: str
+    feature_set_index: int
+    labels: Optional[dict] = {}
+
+    class Config:
+        extra = pydantic.Extra.allow
+
+
+class Entity(FeatureStoreBaseModel):
     name: str
     value_type: str
     labels: Optional[dict] = {}
 
     class Config:
-        extra = Extra.allow
+        extra = pydantic.Extra.allow
 
 
-class Entity(BaseModel):
+class QualifiedEntity(FeatureStoreBaseModel):
     name: str
     value_type: str
+    feature_set_index: int
     labels: Optional[dict] = {}
 
     class Config:
-        extra = Extra.allow
+        extra = pydantic.Extra.allow
 
 
 class FeatureSetSpec(ObjectSpec):
     entities: list[Entity] = []
     features: list[Feature] = []
-    engine: Optional[str] = Field(default="storey")
+    engine: Optional[str] = pydantic.Field(default="storey")
 
 
-class FeatureSet(BaseModel):
-    kind: ObjectKind = Field(ObjectKind.feature_set, const=True)
+class FeatureSet(FeatureStoreBaseModel):
+    kind: ObjectKind = pydantic.Field(ObjectKind.feature_set, const=True)
     metadata: ObjectMetadata
     spec: FeatureSetSpec
     status: ObjectStatus
@@ -62,7 +92,7 @@ class FeatureSet(BaseModel):
         return AuthorizationResourceTypes.feature_set
 
 
-class EntityRecord(BaseModel):
+class EntityRecord(FeatureStoreBaseModel):
     name: str
     value_type: str
     labels: list[LabelRecord]
@@ -71,7 +101,7 @@ class EntityRecord(BaseModel):
         orm_mode = True
 
 
-class FeatureRecord(BaseModel):
+class FeatureRecord(FeatureStoreBaseModel):
     name: str
     value_type: str
     labels: list[LabelRecord]
@@ -88,44 +118,64 @@ class FeatureSetRecord(ObjectRecord):
         orm_mode = True
 
 
-class FeatureSetsOutput(BaseModel):
+class FeatureSetsOutput(FeatureStoreBaseModel):
     feature_sets: list[FeatureSet]
 
 
-class FeatureSetsTagsOutput(BaseModel):
+class FeatureSetsTagsOutput(FeatureStoreBaseModel):
     tags: list[str] = []
 
 
-class FeatureSetDigestSpec(BaseModel):
+class FeatureSetDigestSpec(FeatureStoreBaseModel):
     entities: list[Entity]
     features: list[Feature]
 
 
-class FeatureSetDigestOutput(BaseModel):
+class FeatureSetDigestOutput(FeatureStoreBaseModel):
     metadata: ObjectMetadata
     spec: FeatureSetDigestSpec
 
 
-class FeatureListOutput(BaseModel):
+class FeatureSetDigestSpecV2(FeatureStoreBaseModel):
+    entities: list[Entity]
+
+
+class FeatureSetDigestOutputV2(FeatureStoreBaseModel):
+    feature_set_index: int
+    metadata: ObjectMetadata
+    spec: FeatureSetDigestSpecV2
+
+
+class FeatureListOutput(FeatureStoreBaseModel):
     feature: Feature
     feature_set_digest: FeatureSetDigestOutput
 
 
-class FeaturesOutput(BaseModel):
+class FeaturesOutput(FeatureStoreBaseModel):
     features: list[FeatureListOutput]
 
 
-class EntityListOutput(BaseModel):
+class FeaturesOutputV2(FeatureStoreBaseModel):
+    features: list[QualifiedFeature]
+    feature_set_digests: list[FeatureSetDigestOutputV2]
+
+
+class EntityListOutput(FeatureStoreBaseModel):
     entity: Entity
     feature_set_digest: FeatureSetDigestOutput
 
 
-class EntitiesOutput(BaseModel):
+class EntitiesOutputV2(FeatureStoreBaseModel):
+    entities: list[QualifiedEntity]
+    feature_set_digests: list[FeatureSetDigestOutputV2]
+
+
+class EntitiesOutput(FeatureStoreBaseModel):
     entities: list[EntityListOutput]
 
 
-class FeatureVector(BaseModel):
-    kind: ObjectKind = Field(ObjectKind.feature_vector, const=True)
+class FeatureVector(FeatureStoreBaseModel):
+    kind: ObjectKind = pydantic.Field(ObjectKind.feature_vector, const=True)
     metadata: ObjectMetadata
     spec: ObjectSpec
     status: ObjectStatus
@@ -139,39 +189,39 @@ class FeatureVectorRecord(ObjectRecord):
     pass
 
 
-class FeatureVectorsOutput(BaseModel):
+class FeatureVectorsOutput(FeatureStoreBaseModel):
     feature_vectors: list[FeatureVector]
 
 
-class FeatureVectorsTagsOutput(BaseModel):
+class FeatureVectorsTagsOutput(FeatureStoreBaseModel):
     tags: list[str] = []
 
 
-class DataSource(BaseModel):
+class DataSource(FeatureStoreBaseModel):
     kind: str
     name: str
     path: str
 
     class Config:
-        extra = Extra.allow
+        extra = pydantic.Extra.allow
 
 
-class DataTarget(BaseModel):
+class DataTarget(FeatureStoreBaseModel):
     kind: str
     name: str
     path: Optional[str]
 
     class Config:
-        extra = Extra.allow
+        extra = pydantic.Extra.allow
 
 
-class FeatureSetIngestInput(BaseModel):
+class FeatureSetIngestInput(FeatureStoreBaseModel):
     source: Optional[DataSource]
     targets: Optional[list[DataTarget]]
     infer_options: Optional[int]
     credentials: Credentials = Credentials()
 
 
-class FeatureSetIngestOutput(BaseModel):
+class FeatureSetIngestOutput(FeatureStoreBaseModel):
     feature_set: FeatureSet
     run_object: dict

mlrun/config.py

@@ -229,6 +229,9 @@ default_config = {
                     "executing": "24h",
                 }
             },
+            # When the module is reloaded, the maximum depth recursion configuration for the recursive reload
+            # function is used to prevent infinite loop
+            "reload_max_recursion_depth": 100,
         },
         "databricks": {
             "artifact_directory_path": "/mlrun_databricks_runtime/artifacts_dictionaries"

mlrun/db/base.py

@@ -838,6 +838,7 @@ class RunDBInterface(ABC):
         base_period: int = 10,
         image: str = "mlrun/mlrun",
         deploy_histogram_data_drift_app: bool = True,
+        rebuild_images: bool = False,
     ) -> None:
         pass
 

mlrun/db/httpdb.py

@@ -3320,6 +3320,7 @@ class HTTPRunDB(RunDBInterface):
         base_period: int = 10,
         image: str = "mlrun/mlrun",
         deploy_histogram_data_drift_app: bool = True,
+        rebuild_images: bool = False,
     ) -> None:
         """
         Deploy model monitoring application controller, writer and stream functions.
@@ -3329,13 +3330,14 @@ class HTTPRunDB(RunDBInterface):
         The stream function goal is to monitor the log of the data stream. It is triggered when a new log entry
         is detected. It processes the new events into statistics that are then written to statistics databases.
 
-        :param project:     Project name.
-        :param base_period: The time period in minutes in which the model monitoring controller function
-                            triggers. By default, the base period is 10 minutes.
-        :param image:       The image of the model monitoring controller, writer & monitoring
-                            stream functions, which are real time nuclio functions.
-                            By default, the image is mlrun/mlrun.
+        :param project:                         Project name.
+        :param base_period:                     The time period in minutes in which the model monitoring controller
+                                                function triggers. By default, the base period is 10 minutes.
+        :param image:                           The image of the model monitoring controller, writer & monitoring
+                                                stream functions, which are real time nuclio functions.
+                                                By default, the image is mlrun/mlrun.
         :param deploy_histogram_data_drift_app: If true, deploy the default histogram-based data drift application.
+        :param rebuild_images:                  If true, force rebuild of model monitoring infrastructure images.
         """
         self.api_call(
             method=mlrun.common.types.HTTPMethod.POST,
@@ -3344,6 +3346,7 @@ class HTTPRunDB(RunDBInterface):
                 "base_period": base_period,
                 "image": image,
                 "deploy_histogram_data_drift_app": deploy_histogram_data_drift_app,
+                "rebuild_images": rebuild_images,
             },
         )
 

mlrun/db/nopdb.py

@@ -675,6 +675,7 @@ class NopDB(RunDBInterface):
         base_period: int = 10,
         image: str = "mlrun/mlrun",
         deploy_histogram_data_drift_app: bool = True,
+        rebuild_images: bool = False,
     ) -> None:
         pass
 

mlrun/errors.py

@@ -92,9 +92,7 @@ def raise_for_status(
     try:
         response.raise_for_status()
     except (requests.HTTPError, aiohttp.ClientResponseError) as exc:
-        error_message = err_to_str(exc)
-        if message:
-            error_message = f"{error_message}: {message}"
+        error_message = err_to_str(exc) if not message else message
         status_code = (
             response.status_code
             if hasattr(response, "status_code")

mlrun/execution.py

@@ -111,6 +111,7 @@ class MLClientCtx:
 
         self._project_object = None
         self._allow_empty_resources = None
+        self._reset_on_run = None
 
     def __enter__(self):
         return self
@@ -389,6 +390,7 @@ class MLClientCtx:
             self._state_thresholds = spec.get(
                 "state_thresholds", self._state_thresholds
             )
+            self._reset_on_run = spec.get("reset_on_run", self._reset_on_run)
 
         self._init_dbs(rundb)
 

mlrun/launcher/local.py

@@ -69,6 +69,7 @@ class ClientLocalLauncher(launcher.ClientBaseLauncher):
         notifications: Optional[list[mlrun.model.Notification]] = None,
         returns: Optional[list[Union[str, dict[str, str]]]] = None,
         state_thresholds: Optional[dict[str, int]] = None,
+        reset_on_run: Optional[bool] = None,
     ) -> "mlrun.run.RunObject":
         # do not allow local function to be scheduled
         if self._is_run_local and schedule is not None:
@@ -88,6 +89,7 @@ class ClientLocalLauncher(launcher.ClientBaseLauncher):
                 name=name,
                 workdir=workdir,
                 handler=handler,
+                reset_on_run=reset_on_run,
             )
 
         # sanity check
@@ -212,6 +214,7 @@ class ClientLocalLauncher(launcher.ClientBaseLauncher):
         name: Optional[str] = "",
         workdir: Optional[str] = "",
         handler: Optional[str] = None,
+        reset_on_run: Optional[bool] = None,
     ):
         project = project or runtime.metadata.project
         function_name = name or runtime.metadata.name
@@ -250,6 +253,7 @@ class ClientLocalLauncher(launcher.ClientBaseLauncher):
             fn.spec.build = runtime.spec.build
 
         run.spec.handler = handler
+        run.spec.reset_on_run = reset_on_run
         return fn
 
     @staticmethod

mlrun/launcher/remote.py

@@ -59,6 +59,7 @@ class ClientRemoteLauncher(launcher.ClientBaseLauncher):
         notifications: Optional[list[mlrun.model.Notification]] = None,
         returns: Optional[list[Union[str, dict[str, str]]]] = None,
         state_thresholds: Optional[dict[str, int]] = None,
+        reset_on_run: Optional[bool] = None,
     ) -> "mlrun.run.RunObject":
         self.enrich_runtime(runtime, project)
         run = self._create_run_object(task)

mlrun/model.py

@@ -872,6 +872,7 @@ class RunSpec(ModelObj):
         returns=None,
         notifications=None,
         state_thresholds=None,
+        reset_on_run=None,
     ):
         # A dictionary of parsing configurations that will be read from the inputs the user set. The keys are the inputs
         # keys (parameter names) and the values are the type hint given in the input keys after the colon.
@@ -908,6 +909,7 @@ class RunSpec(ModelObj):
         self.allow_empty_resources = allow_empty_resources
         self._notifications = notifications or []
         self.state_thresholds = state_thresholds or {}
+        self.reset_on_run = reset_on_run
 
     def _serialize_field(
         self, struct: dict, field_name: str = None, strip: bool = False

mlrun/model_monitoring/api.py

@@ -645,6 +645,7 @@ def _create_model_monitoring_function_base(
         app_step = prepare_step.to(class_name=application_class, **application_kwargs)
     else:
         app_step = prepare_step.to(class_name=application_class)
+    app_step.__class__ = mlrun.serving.MonitoringApplicationStep
     app_step.to(
         class_name="mlrun.model_monitoring.applications._application_steps._PushToMonitoringWriter",
         name="PushToMonitoringWriter",

mlrun/model_monitoring/applications/base.py

@@ -21,10 +21,10 @@ import pandas as pd
 import mlrun
 import mlrun.model_monitoring.applications.context as mm_context
 import mlrun.model_monitoring.applications.results as mm_results
-from mlrun.serving.utils import StepToDict
+from mlrun.serving.utils import MonitoringApplicationToDict
 
 
-class ModelMonitoringApplicationBaseV2(StepToDict, ABC):
+class ModelMonitoringApplicationBaseV2(MonitoringApplicationToDict, ABC):
     """
     A base class for a model monitoring application.
     Inherit from this class to create a custom model monitoring application.
@@ -112,7 +112,7 @@ class ModelMonitoringApplicationBaseV2(StepToDict, ABC):
         raise NotImplementedError
 
 
-class ModelMonitoringApplicationBase(StepToDict, ABC):
+class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
     """
     A base class for a model monitoring application.
     Inherit from this class to create a custom model monitoring application.

mlrun/model_monitoring/db/stores/__init__.py

@@ -31,17 +31,12 @@ class ObjectStoreFactory(enum.Enum):
     def to_object_store(
         self,
         project: str,
-        access_key: str = None,
-        secret_provider: typing.Callable = None,
+        **kwargs,
     ) -> StoreBase:
         """
         Return a StoreBase object based on the provided enum value.
 
         :param project:                   The name of the project.
-        :param access_key:                Access key with permission to the DB table. Note that if access key is None
-                                          and the endpoint target is from type KV then the access key will be
-                                          retrieved from the environment variable.
-        :param secret_provider:           An optional secret provider to get the connection string secret.
 
         :return: `StoreBase` object.
 
@@ -50,10 +45,7 @@ class ObjectStoreFactory(enum.Enum):
         if self == self.v3io_nosql:
             from mlrun.model_monitoring.db.stores.v3io_kv.kv_store import KVStoreBase
 
-            # Get V3IO access key from env
-            access_key = access_key or mlrun.mlconf.get_v3io_access_key()
-
-            return KVStoreBase(project=project, access_key=access_key)
+            return KVStoreBase(project=project)
 
         # Assuming SQL store target if store type is not KV.
         # Update these lines once there are more than two store target types.
@@ -62,7 +54,7 @@ class ObjectStoreFactory(enum.Enum):
 
         return SQLStoreBase(
             project=project,
-            secret_provider=secret_provider,
+            **kwargs,
         )
 
     @classmethod
@@ -79,7 +71,7 @@ class ObjectStoreFactory(enum.Enum):
 def get_model_endpoint_store(
     project: str,
     access_key: str = None,
-    secret_provider: typing.Callable = None,
+    secret_provider: typing.Optional[typing.Callable[[str], str]] = None,
 ) -> StoreBase:
     # Leaving here for backwards compatibility
     warnings.warn(
@@ -95,24 +87,38 @@ def get_model_endpoint_store(
 
 def get_store_object(
     project: str,
-    access_key: str = None,
-    secret_provider: typing.Callable = None,
+    secret_provider: typing.Optional[typing.Callable[[str], str]] = None,
+    **kwargs,
 ) -> StoreBase:
     """
-    Getting the DB target type based on mlrun.config.model_endpoint_monitoring.store_type.
+    Generate a store object. If a connection string is provided, the store type will be updated according to the
+    connection string. Currently, the supported store types are SQL and v3io-nosql.
 
     :param project:         The name of the project.
-    :param access_key:      Access key with permission to the DB table.
     :param secret_provider: An optional secret provider to get the connection string secret.
 
-    :return: `StoreBase` object. Using this object, the user can apply different operations on the
-             model monitoring record such as write, update, get and delete a model endpoint.
+    :return: `StoreBase` object. Using this object, the user can apply different operations such as write, update, get
+    and delete a model endpoint record.
     """
 
+    store_connection_string = mlrun.model_monitoring.helpers.get_connection_string(
+        secret_provider=secret_provider
+    )
+
+    if store_connection_string and (
+        store_connection_string.startswith("mysql")
+        or store_connection_string.startswith("sqlite")
+    ):
+        store_type = mlrun.common.schemas.model_monitoring.ModelEndpointTarget.SQL
+        kwargs["store_connection_string"] = store_connection_string
+    else:
+        # Set the default store type if no connection has been set
+        store_type = mlrun.mlconf.model_endpoint_monitoring.store_type
+
     # Get store type value from ObjectStoreFactory enum class
-    store_type = ObjectStoreFactory(mlrun.mlconf.model_endpoint_monitoring.store_type)
+    store_type_fact = ObjectStoreFactory(store_type)
 
     # Convert into store target object
-    return store_type.to_object_store(
-        project=project, access_key=access_key, secret_provider=secret_provider
+    return store_type_fact.to_object_store(
+        project=project, secret_provider=secret_provider, **kwargs
     )

mlrun/model_monitoring/db/stores/base/store.py

@@ -19,6 +19,7 @@ import mlrun.common.schemas.model_monitoring as mm_schemas
 
 
 class StoreBase(ABC):
+    type: typing.ClassVar[str]
     """
     An abstract class to handle the store object in the DB target.
     """

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

@@ -25,14 +25,15 @@ from sqlalchemy.sql.elements import BinaryExpression
 
 import mlrun.common.model_monitoring.helpers
 import mlrun.common.schemas.model_monitoring as mm_schemas
-import mlrun.model_monitoring.db
 import mlrun.model_monitoring.db.stores.sqldb.models
 import mlrun.model_monitoring.helpers
 from mlrun.common.db.sql_session import create_session, get_engine
+from mlrun.model_monitoring.db import StoreBase
 from mlrun.utils import datetime_now, logger
 
 
-class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
+class SQLStoreBase(StoreBase):
+    type: typing.ClassVar[str] = mm_schemas.ModelEndpointTarget.SQL
     """
     Handles the DB operations when the DB target is from type SQL. For the SQL operations, we use SQLAlchemy, a Python
     SQL toolkit that handles the communication with the database.  When using SQL for storing the model monitoring
@@ -44,23 +45,22 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
     def __init__(
         self,
         project: str,
-        secret_provider: typing.Callable = None,
+        **kwargs,
     ):
         """
         Initialize SQL store target object.
 
         :param project:               The name of the project.
-        :param secret_provider:       An optional secret provider to get the connection string secret.
         """
 
         super().__init__(project=project)
 
-        self._sql_connection_string = (
-            mlrun.model_monitoring.helpers.get_connection_string(
-                secret_provider=secret_provider
+        if "store_connection_string" not in kwargs:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "connection_string is a required parameter for SQLStoreBase."
             )
-        )
 
+        self._sql_connection_string = kwargs.get("store_connection_string")
         self._engine = get_engine(dsn=self._sql_connection_string)
 
     def _init_tables(self):