mlrun 1.8.0rc4__py3-none-any.whl → 1.8.0rc7__py3-none-any.whl

mlrun/execution.py

@@ -17,7 +17,7 @@ import os
 import uuid
 import warnings
 from copy import deepcopy
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 import numpy as np
 import yaml
@@ -42,6 +42,7 @@ from .features import Feature
 from .model import HyperParamOptions
 from .secrets import SecretsStore
 from .utils import (
+    Logger,
     RunKeys,
     dict_to_json,
     dict_to_yaml,
@@ -158,7 +159,7 @@ class MLClientCtx:
         return self._project
 
     @property
-    def logger(self):
+    def logger(self) -> Logger:
         """Built-in logger interface
 
         Example::
@@ -500,11 +501,11 @@ class MLClientCtx:
             return default
         return self._parameters[key]
 
-    def get_project_object(self):
+    def get_project_object(self) -> Optional["mlrun.MlrunProject"]:
         """
         Get the MLRun project object by the project name set in the context.
 
-        :return: The project object or None if it couldn't be retrieved.
+        :returns: The project object or None if it couldn't be retrieved.
         """
         return self._load_project_object()
 
@@ -628,7 +629,7 @@ class MLClientCtx:
         format=None,
         db_key=None,
         **kwargs,
-    ):
+    ) -> Artifact:
         """Log an output artifact and optionally upload it to datastore
 
         Example::
@@ -698,7 +699,7 @@ class MLClientCtx:
         extra_data=None,
         label_column: Optional[str] = None,
         **kwargs,
-    ):
+    ) -> DatasetArtifact:
         """Log a dataset artifact and optionally upload it to datastore
 
         If the dataset exists with the same key and tag, it will be overwritten.
@@ -736,7 +737,7 @@ class MLClientCtx:
         :param db_key:        The key to use in the artifact DB table, by default its run name + '_' + key
                               db_key=False will not register it in the artifacts table
 
-        :returns: Artifact object
+        :returns: Dataset artifact object
         """
         ds = DatasetArtifact(
             key,
@@ -749,16 +750,19 @@ class MLClientCtx:
             **kwargs,
         )
 
-        item = self._artifacts_manager.log_artifact(
-            self,
-            ds,
-            local_path=local_path,
-            artifact_path=extend_artifact_path(artifact_path, self.artifact_path),
-            target_path=target_path,
-            tag=tag,
-            upload=upload,
-            db_key=db_key,
-            labels=labels,
+        item = cast(
+            DatasetArtifact,
+            self._artifacts_manager.log_artifact(
+                self,
+                ds,
+                local_path=local_path,
+                artifact_path=extend_artifact_path(artifact_path, self.artifact_path),
+                target_path=target_path,
+                tag=tag,
+                upload=upload,
+                db_key=db_key,
+                labels=labels,
+            ),
         )
         self._update_run()
         return item
@@ -786,7 +790,7 @@ class MLClientCtx:
         extra_data=None,
         db_key=None,
         **kwargs,
-    ):
+    ) -> ModelArtifact:
         """Log a model artifact and optionally upload it to datastore
 
         Example::
@@ -828,7 +832,7 @@ class MLClientCtx:
         :param db_key:          The key to use in the artifact DB table, by default its run name + '_' + key
                                 db_key=False will not register it in the artifacts table
 
-        :returns: Artifact object
+        :returns: Model artifact object
         """
 
         if training_set is not None and inputs:
@@ -855,14 +859,17 @@ class MLClientCtx:
         if training_set is not None:
             model.infer_from_df(training_set, label_column)
 
-        item = self._artifacts_manager.log_artifact(
-            self,
-            model,
-            artifact_path=extend_artifact_path(artifact_path, self.artifact_path),
-            tag=tag,
-            upload=upload,
-            db_key=db_key,
-            labels=labels,
+        item = cast(
+            ModelArtifact,
+            self._artifacts_manager.log_artifact(
+                self,
+                model,
+                artifact_path=extend_artifact_path(artifact_path, self.artifact_path),
+                tag=tag,
+                upload=upload,
+                db_key=db_key,
+                labels=labels,
+            ),
         )
         self._update_run()
         return item
@@ -870,28 +877,35 @@ class MLClientCtx:
     def log_document(
         self,
         key: str,
+        tag: str = "",
+        local_path: str = "",
         artifact_path: Optional[str] = None,
         document_loader: DocumentLoaderSpec = DocumentLoaderSpec(),
-        tag: str = "",
         upload: Optional[bool] = False,
         labels: Optional[dict[str, str]] = None,
+        target_path: Optional[str] = None,
         **kwargs,
     ) -> DocumentArtifact:
         """
         Log a document as an artifact.
 
         :param key: Artifact key
-        :param target_path: Path to the local file
-        :param artifact_path: Target path for artifact storage
-        :param document_loader: Spec to use to load the artifact as langchain document
         :param tag: Version tag
+        :param local_path:    path to the local file we upload, will also be use
+                              as the destination subpath (under "artifact_path")
+        :param artifact_path:   Target artifact path (when not using the default)
+                                to define a subpath under the default location use:
+                                `artifact_path=context.artifact_subpath('data')`
+        :param document_loader: Spec to use to load the artifact as langchain document
         :param upload: Whether to upload the artifact
         :param labels: Key-value labels
+        :param target_path: Path to the local file
         :param kwargs: Additional keyword arguments
         :return: DocumentArtifact object
         """
         doc_artifact = DocumentArtifact(
             key=key,
+            original_source=local_path or target_path,
             document_loader=document_loader,
             **kwargs,
         )
@@ -1193,7 +1207,7 @@ class MLClientCtx:
         self._data_stores = store_manager.set(self._secrets_manager, db=self._rundb)
         self._artifacts_manager = ArtifactManager(db=self._rundb)
 
-    def _load_project_object(self):
+    def _load_project_object(self) -> Optional["mlrun.MlrunProject"]:
         if not self._project_object:
             if not self._project:
                 self.logger.warning(

mlrun/feature_store/api.py

@@ -11,6 +11,7 @@
 # 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 copy
 import importlib.util
 import pathlib

mlrun/model.py

@@ -117,6 +117,8 @@ class ModelObj:
                 # If one of the attributes is a third party object that has to_dict method (such as k8s objects), then
                 # add it to the object's _fields_to_serialize attribute and handle it in the _serialize_field method.
                 if hasattr(field_value, "to_dict"):
+                    # TODO: Allow passing fields to exclude from the parent object to the child object
+                    #  e.g.: run.to_dict(exclude=["status.artifacts"])
                     field_value = field_value.to_dict(strip=strip)
                     if self._is_valid_field_value_for_serialization(
                         field_name, field_value, strip
@@ -443,6 +445,7 @@ class Credentials(ModelObj):
 class BaseMetadata(ModelObj):
     _default_fields_to_strip = ModelObj._default_fields_to_strip + [
         "hash",
+        "uid",
         # Below are environment specific fields, no need to keep when stripping
         "namespace",
         "project",
@@ -465,10 +468,12 @@ class BaseMetadata(ModelObj):
         categories=None,
         updated=None,
         credentials=None,
+        uid=None,
     ):
         self.name = name
         self.tag = tag
         self.hash = hash
+        self.uid = uid
         self.namespace = namespace
         self.project = project or ""
         self.labels = labels or {}
@@ -1267,6 +1272,8 @@ class RunSpec(ModelObj):
 class RunStatus(ModelObj):
     """Run status"""
 
+    _default_fields_to_strip = ModelObj._default_fields_to_strip + ["artifacts"]
+
     def __init__(
         self,
         state=None,

mlrun/model_monitoring/__init__.py

@@ -14,7 +14,8 @@
 
 # for backwards compatibility
 
-from .db import get_store_object, get_tsdb_connector
+from mlrun.common.schemas import ModelEndpoint, ModelEndpointList
+
+from .db import get_tsdb_connector
 from .helpers import get_stream_path
-from .model_endpoint import ModelEndpoint
 from .tracking_policy import TrackingPolicy

mlrun/model_monitoring/api.py

@@ -23,18 +23,28 @@ import pandas as pd
 import mlrun.artifacts
 import mlrun.common.helpers
 import mlrun.common.schemas.model_monitoring.constants as mm_constants
+import mlrun.datastore.base
 import mlrun.feature_store
 import mlrun.model_monitoring.applications as mm_app
 import mlrun.serving
+from mlrun.common.schemas import ModelEndpoint
+from mlrun.common.schemas.model_monitoring import (
+    FunctionURI,
+)
 from mlrun.data_types.infer import InferOptions, get_df_stats
 from mlrun.utils import datetime_now, logger
 
 from .helpers import update_model_endpoint_last_request
-from .model_endpoint import ModelEndpoint
 
 # A union of all supported dataset types:
 DatasetType = typing.Union[
-    mlrun.DataItem, list, dict, pd.DataFrame, pd.Series, np.ndarray, typing.Any
+    mlrun.datastore.base.DataItem,
+    list,
+    dict,
+    pd.DataFrame,
+    pd.Series,
+    np.ndarray,
+    typing.Any,
 ]
 
 
@@ -44,10 +54,8 @@ def get_or_create_model_endpoint(
     model_endpoint_name: str = "",
     endpoint_id: str = "",
     function_name: str = "",
-    context: mlrun.MLClientCtx = None,
+    context: typing.Optional["mlrun.MLClientCtx"] = None,
     sample_set_statistics: typing.Optional[dict[str, typing.Any]] = None,
-    drift_threshold: typing.Optional[float] = None,
-    possible_drift_threshold: typing.Optional[float] = None,
     monitoring_mode: mm_constants.ModelMonitoringMode = mm_constants.ModelMonitoringMode.disabled,
     db_session=None,
 ) -> ModelEndpoint:
@@ -68,10 +76,6 @@ def get_or_create_model_endpoint(
                                      full function hash.
     :param sample_set_statistics:    Dictionary of sample set statistics that will be used as a reference data for
                                      the new model endpoint (applicable only to new endpoint_id).
-    :param drift_threshold:          (deprecated) The threshold of which to mark drifts (applicable only to new
-                                     endpoint_id).
-    :param possible_drift_threshold: (deprecated) The threshold of which to mark possible drifts (applicable only to new
-                                     endpoint_id).
     :param monitoring_mode:          If enabled, apply model monitoring features on the provided endpoint id
                                      (applicable only to new endpoint_id).
     :param db_session:               A runtime session that manages the current dialog with the database.
@@ -79,18 +83,15 @@ def get_or_create_model_endpoint(
     :return: A ModelEndpoint object
     """
 
-    if not endpoint_id:
-        # Generate a new model endpoint id based on the project name and model name
-        endpoint_id = hashlib.sha1(
-            f"{project}_{model_endpoint_name}".encode()
-        ).hexdigest()
-
     if not db_session:
         # Generate a runtime database
         db_session = mlrun.get_run_db()
     try:
         model_endpoint = db_session.get_model_endpoint(
-            project=project, endpoint_id=endpoint_id
+            project=project,
+            name=model_endpoint_name,
+            endpoint_id=endpoint_id,
+            function_name=function_name,
         )
         # If other fields provided, validate that they are correspond to the existing model endpoint data
         _model_endpoint_validations(
@@ -104,7 +105,6 @@ def get_or_create_model_endpoint(
         model_endpoint = _generate_model_endpoint(
             project=project,
             db_session=db_session,
-            endpoint_id=endpoint_id,
             model_path=model_path,
             model_endpoint_name=model_endpoint_name,
             function_name=function_name,
@@ -121,7 +121,7 @@ def record_results(
     model_endpoint_name: str,
     endpoint_id: str = "",
     function_name: str = "",
-    context: typing.Optional[mlrun.MLClientCtx] = None,
+    context: typing.Optional["mlrun.MLClientCtx"] = None,
     infer_results_df: typing.Optional[pd.DataFrame] = None,
     sample_set_statistics: typing.Optional[dict[str, typing.Any]] = None,
     monitoring_mode: mm_constants.ModelMonitoringMode = mm_constants.ModelMonitoringMode.enabled,
@@ -208,13 +208,13 @@ def record_results(
         monitoring_mode=monitoring_mode,
         db_session=db,
     )
-    logger.debug("Model endpoint", endpoint=model_endpoint.to_dict())
+    logger.debug("Model endpoint", endpoint=model_endpoint)
 
     timestamp = datetime_now()
     if infer_results_df is not None:
         # Write the monitoring parquet to the relevant model endpoint context
         write_monitoring_df(
-            feature_set_uri=model_endpoint.status.monitoring_feature_set_uri,
+            feature_set_uri=model_endpoint.spec.monitoring_feature_set_uri,
             infer_datetime=timestamp,
             endpoint_id=model_endpoint.metadata.uid,
             infer_results_df=infer_results_df,
@@ -278,7 +278,7 @@ def _model_endpoint_validations(
     # Feature stats
     if (
         sample_set_statistics
-        and sample_set_statistics != model_endpoint.status.feature_stats
+        and sample_set_statistics != model_endpoint.spec.feature_stats
     ):
         logger.warning(
             "Provided sample set statistics is different from the registered statistics. "
@@ -290,7 +290,7 @@ def write_monitoring_df(
     endpoint_id: str,
     infer_results_df: pd.DataFrame,
     infer_datetime: datetime,
-    monitoring_feature_set: typing.Optional[mlrun.feature_store.FeatureSet] = None,
+    monitoring_feature_set: typing.Optional["mlrun.feature_store.FeatureSet"] = None,
     feature_set_uri: str = "",
 ) -> None:
     """Write infer results dataframe to the monitoring parquet target of the current model endpoint. The dataframe will
@@ -330,11 +330,10 @@ def write_monitoring_df(
 def _generate_model_endpoint(
     project: str,
     db_session,
-    endpoint_id: str,
     model_path: str,
     model_endpoint_name: str,
     function_name: str,
-    context: mlrun.MLClientCtx,
+    context: "mlrun.MLClientCtx",
     sample_set_statistics: dict[str, typing.Any],
     monitoring_mode: mm_constants.ModelMonitoringMode = mm_constants.ModelMonitoringMode.disabled,
 ) -> ModelEndpoint:
@@ -344,7 +343,6 @@ def _generate_model_endpoint(
     :param project:                  Project name.
 
     :param db_session:               A session that manages the current dialog with the database.
-    :param endpoint_id:              Model endpoint unique ID.
     :param model_path:               The model Store path.
     :param model_endpoint_name:      Model endpoint name will be presented under the new model endpoint.
     :param function_name:            If a new model endpoint is created, use this function name for generating the
@@ -355,34 +353,40 @@ def _generate_model_endpoint(
                                      the current model endpoint. Will be stored under
                                      `model_endpoint.status.feature_stats`.
 
-    :return `mlrun.model_monitoring.model_endpoint.ModelEndpoint` object.
+    :return `mlrun.common.schemas.ModelEndpoint` object.
     """
-    model_endpoint = ModelEndpoint()
-    model_endpoint.metadata.project = project
-    model_endpoint.metadata.uid = endpoint_id
-    if function_name:
-        model_endpoint.spec.function_uri = project + "/" + function_name
-    elif not context:
-        raise mlrun.errors.MLRunInvalidArgumentError(
-            "Please provide either a function name or a valid MLRun context"
+    if not function_name and context:
+        function_name = FunctionURI.from_string(
+            context.to_dict()["spec"]["function"]
+        ).function
+    model_obj = None
+    if model_path:
+        model_obj: mlrun.artifacts.ModelArtifact = (
+            mlrun.datastore.store_resources.get_store_resource(
+                model_path, db=db_session
+            )
         )
-    else:
-        model_endpoint.spec.function_uri = context.to_dict()["spec"]["function"]
-    model_endpoint.spec.model_uri = model_path
-    model_endpoint.spec.model = model_endpoint_name
-    model_endpoint.spec.model_class = "drift-analysis"
-    model_endpoint.spec.monitoring_mode = monitoring_mode
-    model_endpoint.status.first_request = model_endpoint.status.last_request = (
-        datetime_now().isoformat()
-    )
-    if sample_set_statistics:
-        model_endpoint.status.feature_stats = sample_set_statistics
-
-    db_session.create_model_endpoint(
-        project=project, endpoint_id=endpoint_id, model_endpoint=model_endpoint
+    current_time = datetime_now()
+    model_endpoint = mlrun.common.schemas.ModelEndpoint(
+        metadata=mlrun.common.schemas.ModelEndpointMetadata(
+            project=project,
+            name=model_endpoint_name,
+            endpoint_type=mlrun.common.schemas.model_monitoring.EndpointType.BATCH_EP,
+        ),
+        spec=mlrun.common.schemas.ModelEndpointSpec(
+            function_name=function_name,
+            model_name=model_obj.metadata.key if model_path else None,
+            model_uid=model_obj.metadata.uid if model_path else None,
+            model_class="drift-analysis",
+        ),
+        status=mlrun.common.schemas.ModelEndpointStatus(
+            monitoring_mode=monitoring_mode,
+            first_request=current_time,
+            last_request=current_time,
+        ),
     )
 
-    return db_session.get_model_endpoint(project=project, endpoint_id=endpoint_id)
+    return db_session.create_model_endpoint(model_endpoint=model_endpoint)
 
 
 def get_sample_set_statistics(
@@ -531,7 +535,7 @@ def read_dataset_as_dataframe(
 
 
 def log_result(
-    context: mlrun.MLClientCtx,
+    context: "mlrun.MLClientCtx",
     result_set_name: str,
     result_set: pd.DataFrame,
     artifacts_tag: str,
@@ -559,9 +563,7 @@ def _create_model_monitoring_function_base(
     project: str,
     func: typing.Union[str, None] = None,
     application_class: typing.Union[
-        str,
-        mm_app.ModelMonitoringApplicationBase,
-        None,
+        str, "mm_app.ModelMonitoringApplicationBase", None
     ] = None,
     name: typing.Optional[str] = None,
     image: typing.Optional[str] = None,

mlrun/model_monitoring/applications/_application_steps.py

@@ -16,6 +16,7 @@ import json
 import traceback
 from typing import Any, Optional, Union
 
+import mlrun.common.schemas
 import mlrun.common.schemas.alert as alert_objects
 import mlrun.common.schemas.model_monitoring.constants as mm_constant
 import mlrun.datastore
@@ -81,6 +82,7 @@ class _PushToMonitoringWriter(StepToDict):
         self._lazy_init()
         application_results, application_context = event
         writer_event = {
+            mm_constant.WriterEvent.ENDPOINT_NAME: application_context.endpoint_name,
             mm_constant.WriterEvent.APPLICATION_NAME: application_context.application_name,
             mm_constant.WriterEvent.ENDPOINT_ID: application_context.endpoint_id,
             mm_constant.WriterEvent.START_INFER_TIME: application_context.start_infer_time.isoformat(
@@ -125,7 +127,7 @@ class _PrepareMonitoringEvent(StepToDict):
         """
         self.graph_context = context
         self.application_name = application_name
-        self.model_endpoints: dict[str, mlrun.model_monitoring.ModelEndpoint] = {}
+        self.model_endpoints: dict[str, mlrun.common.schemas.ModelEndpoint] = {}
 
     def do(self, event: dict[str, Any]) -> MonitoringApplicationContext:
         """
@@ -135,10 +137,10 @@ class _PrepareMonitoringEvent(StepToDict):
         :return: Application context.
         """
         application_context = MonitoringApplicationContext(
-            graph_context=self.graph_context,
             application_name=self.application_name,
             event=event,
             model_endpoint_dict=self.model_endpoints,
+            graph_context=self.graph_context,
         )
 
         self.model_endpoints.setdefault(