mlrun 1.7.1rc4__py3-none-any.whl → 1.8.0rc8__py3-none-any.whl

mlrun/model.py

@@ -24,7 +24,7 @@ from datetime import datetime
 from os import environ
 from typing import Any, Optional, Union
 
-import pydantic.error_wrappers
+import pydantic.v1.error_wrappers
 
 import mlrun
 import mlrun.common.constants as mlrun_constants
@@ -74,7 +74,10 @@ class ModelObj:
 
     @mlrun.utils.filter_warnings("ignore", FutureWarning)
     def to_dict(
-        self, fields: list = None, exclude: list = None, strip: bool = False
+        self,
+        fields: Optional[list] = None,
+        exclude: Optional[list] = None,
+        strip: bool = False,
     ) -> dict:
         """
         Convert the object to a dict
@@ -114,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
@@ -141,7 +146,7 @@ class ModelObj:
         self._apply_enrichment_before_to_dict_completion(struct, strip=strip)
         return struct
 
-    def _resolve_initial_to_dict_fields(self, fields: list = None) -> list:
+    def _resolve_initial_to_dict_fields(self, fields: Optional[list] = None) -> list:
         """
         Resolve fields to be used in to_dict method.
         If fields is None, use `_dict_fields` attribute of the object.
@@ -184,7 +189,7 @@ class ModelObj:
         self,
         struct: dict,
         method: typing.Callable,
-        fields: typing.Union[list, set] = None,
+        fields: Optional[typing.Union[list, set]] = None,
         strip: bool = False,
     ) -> dict:
         for field_name in fields:
@@ -196,14 +201,14 @@ class ModelObj:
         return struct
 
     def _serialize_field(
-        self, struct: dict, field_name: str = None, strip: bool = False
+        self, struct: dict, field_name: Optional[str] = None, strip: bool = False
     ) -> typing.Any:
         # We pull the field from self and not from struct because it was excluded from the struct when looping over
         # the fields to save.
         return getattr(self, field_name, None)
 
     def _enrich_field(
-        self, struct: dict, field_name: str = None, strip: bool = False
+        self, struct: dict, field_name: Optional[str] = None, strip: bool = False
     ) -> typing.Any:
         # We first try to pull from struct because the field might have been already serialized and if not,
         # we pull from self
@@ -215,7 +220,9 @@ class ModelObj:
         return struct
 
     @classmethod
-    def from_dict(cls, struct=None, fields=None, deprecated_fields: dict = None):
+    def from_dict(
+        cls, struct=None, fields=None, deprecated_fields: Optional[dict] = None
+    ):
         """create an object from a python dictionary"""
         struct = {} if struct is None else struct
         deprecated_fields = deprecated_fields or {}
@@ -430,7 +437,7 @@ class Credentials(ModelObj):
 
     def __init__(
         self,
-        access_key: str = None,
+        access_key: Optional[str] = None,
     ):
         self.access_key = access_key
 
@@ -438,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",
@@ -460,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 {}
@@ -500,7 +510,7 @@ class ImageBuilder(ModelObj):
         origin_filename=None,
         with_mlrun=None,
         auto_build=None,
-        requirements: list = None,
+        requirements: Optional[list] = None,
         extra_args=None,
         builder_env=None,
         source_code_target_dir=None,
@@ -549,7 +559,7 @@ class ImageBuilder(ModelObj):
         self,
         image="",
         base_image=None,
-        commands: list = None,
+        commands: Optional[list] = None,
         secret=None,
         source=None,
         extra=None,
@@ -734,7 +744,7 @@ class Notification(ModelObj):
     def validate_notification(self):
         try:
             mlrun.common.schemas.notification.Notification(**self.to_dict())
-        except pydantic.error_wrappers.ValidationError as exc:
+        except pydantic.v1.error_wrappers.ValidationError as exc:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "Invalid notification object"
             ) from exc
@@ -750,14 +760,14 @@ class Notification(ModelObj):
                 "Notification params size exceeds max size of 1 MB"
             )
 
-    def validate_notification_params(self):
-        notification_class = mlrun.utils.notifications.NotificationTypes(
-            self.kind
-        ).get_notification()
-
+    def validate_notification_params(self, default_notification_params=None):
+        default_notification_params = default_notification_params or {}
+        notification_type = mlrun.utils.notifications.NotificationTypes(self.kind)
+        notification_class = notification_type.get_notification()
         secret_params = self.secret_params or {}
         params = self.params or {}
-
+        default_params = default_notification_params.get(notification_type, {})
+        params = notification_class.enrich_default_params(params, default_params)
         # if the secret_params are already masked - no need to validate
         params_secret = secret_params.get("secret", "")
         if params_secret:
@@ -973,7 +983,7 @@ class RunSpec(ModelObj):
         self.node_selector = node_selector or {}
 
     def _serialize_field(
-        self, struct: dict, field_name: str = None, strip: bool = False
+        self, struct: dict, field_name: Optional[str] = None, strip: bool = False
     ) -> Optional[str]:
         # We pull the field from self and not from struct because it was excluded from the struct
         if field_name == "handler":
@@ -1262,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,
@@ -1275,9 +1287,9 @@ class RunStatus(ModelObj):
         last_update=None,
         iterations=None,
         ui_url=None,
-        reason: str = None,
-        notifications: dict[str, Notification] = None,
-        artifact_uris: dict[str, str] = None,
+        reason: Optional[str] = None,
+        notifications: Optional[dict[str, Notification]] = None,
+        artifact_uris: Optional[dict[str, str]] = None,
     ):
         self.state = state or "created"
         self.status_text = status_text
@@ -1285,7 +1297,7 @@ class RunStatus(ModelObj):
         self.host = host
         self.commit = commit
         self.results = results
-        self.artifacts = artifacts
+        self._artifacts = artifacts
         self.start_time = start_time
         self.last_update = last_update
         self.iterations = iterations
@@ -1293,7 +1305,59 @@ class RunStatus(ModelObj):
         self.reason = reason
         self.notifications = notifications or {}
         # Artifact key -> URI mapping, since the full artifacts are not stored in the runs DB table
-        self.artifact_uris = artifact_uris or {}
+        self._artifact_uris = artifact_uris or {}
+
+    @classmethod
+    def from_dict(
+        cls, struct=None, fields=None, deprecated_fields: Optional[dict] = None
+    ):
+        deprecated_fields = {
+            # Set artifacts as deprecated for lazy loading
+            "artifacts": "artifact_uris"
+        }
+        return super().from_dict(
+            struct, fields=fields, deprecated_fields=deprecated_fields
+        )
+
+    @property
+    def artifacts(self):
+        """
+        Artifacts are lazy loaded to reduce memory consumption.
+        We keep artifact_uris (key -> store URI dictionary) to be able to get the run artifacts easily.
+        If the artifact is not already in the cache, we get it from the store (DB).
+        :return: List of artifact dictionaries
+        """
+        self._artifacts = self._artifacts or []
+        existing_artifact_keys = {
+            artifact["metadata"]["key"] for artifact in self._artifacts
+        }
+        for key, uri in self.artifact_uris.items():
+            if key not in existing_artifact_keys:
+                artifact = mlrun.datastore.get_store_resource(uri)
+                self._artifacts.append(artifact.to_dict())
+        return self._artifacts
+
+    @artifacts.setter
+    def artifacts(self, artifacts):
+        self._artifacts = artifacts
+
+    @property
+    def artifact_uris(self):
+        return self._artifact_uris
+
+    @artifact_uris.setter
+    def artifact_uris(self, artifact_uris):
+        resolved_artifact_uris = {}
+        if isinstance(artifact_uris, list):
+            # artifact_uris is the deprecated list of artifacts - convert to new form
+            for artifact in artifact_uris:
+                if isinstance(artifact, dict):
+                    artifact = mlrun.artifacts.dict_to_artifact(artifact)
+                resolved_artifact_uris[artifact.key] = artifact.uri
+        else:
+            resolved_artifact_uris = artifact_uris
+
+        self._artifact_uris = resolved_artifact_uris
 
     def is_failed(self) -> Optional[bool]:
         """
@@ -1601,7 +1665,7 @@ class RunObject(RunTemplate):
 
         return outputs
 
-    def artifact(self, key: str) -> "mlrun.DataItem":
+    def artifact(self, key: str) -> typing.Optional["mlrun.DataItem"]:
         """Return artifact DataItem by key.
 
         This method waits for the outputs to complete, searches for the artifact matching the given key,
@@ -1644,7 +1708,7 @@ class RunObject(RunTemplate):
         :param key: The key of the artifact to retrieve.
         :return: The last artifact DataItem with the given key, or None if no such artifact is found.
         """
-        if not self.status.artifacts:
+        if not self.status.artifacts and not self.status.artifact_uris:
             return None
 
         # Collect artifacts that match the key
@@ -1655,7 +1719,12 @@ class RunObject(RunTemplate):
         ]
 
         if not matching_artifacts:
-            return None
+            if key not in self.status.artifact_uris:
+                return None
+
+            # Get artifact by store URI sanity (should have been enriched by now in status.artifacts property)
+            artifact_uri = self.status.artifact_uris[key]
+            return mlrun.datastore.get_store_resource(artifact_uri)
 
         # Sort matching artifacts by creation date in ascending order.
         # The last element in the list will be the one created most recently.
@@ -1870,7 +1939,7 @@ class EntrypointParam(ModelObj):
         default=None,
         doc="",
         required=None,
-        choices: list = None,
+        choices: Optional[list] = None,
     ):
         self.name = name
         self.type = type
@@ -2065,12 +2134,12 @@ class DataSource(ModelObj):
 
     def __init__(
         self,
-        name: str = None,
-        path: str = None,
-        attributes: dict[str, object] = None,
-        key_field: str = None,
-        time_field: str = None,
-        schedule: str = None,
+        name: Optional[str] = None,
+        path: Optional[str] = None,
+        attributes: Optional[dict[str, object]] = None,
+        key_field: Optional[str] = None,
+        time_field: Optional[str] = None,
+        schedule: Optional[str] = None,
         start_time: Optional[Union[datetime, str]] = None,
         end_time: Optional[Union[datetime, str]] = None,
     ):
@@ -2092,7 +2161,7 @@ class DataSource(ModelObj):
         self._secrets = secrets
 
     def _serialize_field(
-        self, struct: dict, field_name: str = None, strip: bool = False
+        self, struct: dict, field_name: Optional[str] = None, strip: bool = False
     ) -> typing.Any:
         value = super()._serialize_field(struct, field_name, strip)
         # We pull the field from self and not from struct because it was excluded from the struct when looping over
@@ -2124,7 +2193,9 @@ class DataTargetBase(ModelObj):
     ]
 
     @classmethod
-    def from_dict(cls, struct=None, fields=None, deprecated_fields: dict = None):
+    def from_dict(
+        cls, struct=None, fields=None, deprecated_fields: Optional[dict] = None
+    ):
         return super().from_dict(struct, fields=fields)
 
     def get_path(self):
@@ -2140,10 +2211,10 @@ class DataTargetBase(ModelObj):
 
     def __init__(
         self,
-        kind: str = None,
+        kind: Optional[str] = None,
         name: str = "",
         path=None,
-        attributes: dict[str, str] = None,
+        attributes: Optional[dict[str, str]] = None,
         after_step=None,
         partitioned: bool = False,
         key_bucketing_number: Optional[int] = None,
@@ -2151,8 +2222,8 @@ class DataTargetBase(ModelObj):
         time_partitioning_granularity: Optional[str] = None,
         max_events: Optional[int] = None,
         flush_after_seconds: Optional[int] = None,
-        storage_options: dict[str, str] = None,
-        schema: dict[str, Any] = None,
+        storage_options: Optional[dict[str, str]] = None,
+        schema: Optional[dict[str, Any]] = None,
         credentials_prefix=None,
     ):
         self.name = name
@@ -2208,7 +2279,7 @@ class DataTarget(DataTargetBase):
 
     def __init__(
         self,
-        kind: str = None,
+        kind: Optional[str] = None,
         name: str = "",
         path=None,
         online=None,
@@ -2237,12 +2308,12 @@ class DataTarget(DataTargetBase):
 class VersionedObjMetadata(ModelObj):
     def __init__(
         self,
-        name: str = None,
-        tag: str = None,
-        uid: str = None,
-        project: str = None,
-        labels: dict[str, str] = None,
-        annotations: dict[str, str] = None,
+        name: Optional[str] = None,
+        tag: Optional[str] = None,
+        uid: Optional[str] = None,
+        project: Optional[str] = None,
+        labels: Optional[dict[str, str]] = None,
+        annotations: Optional[dict[str, str]] = None,
         updated=None,
     ):
         self.name = name

mlrun/model_monitoring/__init__.py

@@ -11,11 +11,11 @@
 # 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.
-#
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
+
 # 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,
-    sample_set_statistics: dict[str, typing.Any] = None,
-    drift_threshold: typing.Optional[float] = None,
-    possible_drift_threshold: typing.Optional[float] = None,
+    context: typing.Optional["mlrun.MLClientCtx"] = None,
+    sample_set_statistics: typing.Optional[dict[str, typing.Any]] = 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,
@@ -234,7 +234,7 @@ def record_results(
 def _model_endpoint_validations(
     model_endpoint: ModelEndpoint,
     model_path: str = "",
-    sample_set_statistics: dict[str, typing.Any] = None,
+    sample_set_statistics: typing.Optional[dict[str, typing.Any]] = None,
 ) -> None:
     """
     Validate that provided model endpoint configurations match the stored fields of the provided `ModelEndpoint`
@@ -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,39 +353,45 @@ 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(
     sample_set: DatasetType = None,
-    model_artifact_feature_stats: dict = None,
+    model_artifact_feature_stats: typing.Optional[dict] = None,
     sample_set_columns: typing.Optional[list] = None,
     sample_set_drop_columns: typing.Optional[list] = None,
     sample_set_label_columns: typing.Optional[list] = None,
@@ -445,9 +449,9 @@ def get_sample_set_statistics(
 
 def read_dataset_as_dataframe(
     dataset: DatasetType,
-    feature_columns: typing.Union[str, list[str]] = None,
-    label_columns: typing.Union[str, list[str]] = None,
-    drop_columns: typing.Union[str, list[str], int, list[int]] = None,
+    feature_columns: typing.Optional[typing.Union[str, list[str]]] = None,
+    label_columns: typing.Optional[typing.Union[str, list[str]]] = None,
+    drop_columns: typing.Optional[typing.Union[str, list[str], int, list[int]]] = None,
 ) -> tuple[pd.DataFrame, list[str]]:
     """
     Parse the given dataset into a DataFrame and drop the columns accordingly. In addition, the label columns will be
@@ -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,