mlrun 1.7.2rc3__py3-none-any.whl → 1.8.0rc1__py3-none-any.whl

mlrun/frameworks/xgboost/__init__.py

@@ -11,9 +11,8 @@
 # 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
-from typing import Union
+
+from typing import Optional, Union
 
 import xgboost as xgb
 
@@ -34,26 +33,28 @@ def apply_mlrun(
     model: xgb.XGBModel = None,
     model_name: str = "model",
     tag: str = "",
-    model_path: str = None,
-    modules_map: Union[dict[str, Union[None, str, list[str]]], str] = None,
-    custom_objects_map: Union[dict[str, Union[str, list[str]]], str] = None,
-    custom_objects_directory: str = None,
+    model_path: Optional[str] = None,
+    modules_map: Optional[Union[dict[str, Union[None, str, list[str]]], str]] = None,
+    custom_objects_map: Optional[Union[dict[str, Union[str, list[str]]], str]] = None,
+    custom_objects_directory: Optional[str] = None,
     context: mlrun.MLClientCtx = None,
-    artifacts: Union[list[MLPlan], list[str], dict[str, dict]] = None,
-    metrics: Union[
-        list[Metric],
-        list[XGBoostTypes.MetricEntryType],
-        dict[str, XGBoostTypes.MetricEntryType],
+    artifacts: Optional[Union[list[MLPlan], list[str], dict[str, dict]]] = None,
+    metrics: Optional[
+        Union[
+            list[Metric],
+            list[XGBoostTypes.MetricEntryType],
+            dict[str, XGBoostTypes.MetricEntryType],
+        ]
     ] = None,
     x_test: XGBoostTypes.DatasetType = None,
     y_test: XGBoostTypes.DatasetType = None,
     sample_set: Union[XGBoostTypes.DatasetType, mlrun.DataItem, str] = None,
-    y_columns: Union[list[str], list[int]] = None,
-    feature_vector: str = None,
-    feature_weights: list[float] = None,
-    labels: dict[str, Union[str, int, float]] = None,
-    parameters: dict[str, Union[str, int, float]] = None,
-    extra_data: dict[str, XGBoostTypes.ExtraDataType] = None,
+    y_columns: Optional[Union[list[str], list[int]]] = None,
+    feature_vector: Optional[str] = None,
+    feature_weights: Optional[list[float]] = None,
+    labels: Optional[dict[str, Union[str, int, float]]] = None,
+    parameters: Optional[dict[str, Union[str, int, float]]] = None,
+    extra_data: Optional[dict[str, XGBoostTypes.ExtraDataType]] = None,
     auto_log: bool = True,
     **kwargs,
 ) -> XGBoostModelHandler:

mlrun/frameworks/xgboost/model_handler.py

@@ -14,7 +14,7 @@
 #
 import os
 import pickle
-from typing import Union
+from typing import Optional, Union
 
 import cloudpickle
 
@@ -45,11 +45,15 @@ class XGBoostModelHandler(MLModelHandler):
     def __init__(
         self,
         model: XGBoostTypes.ModelType = None,
-        model_path: str = None,
-        model_name: str = None,
-        modules_map: Union[dict[str, Union[None, str, list[str]]], str] = None,
-        custom_objects_map: Union[dict[str, Union[str, list[str]]], str] = None,
-        custom_objects_directory: str = None,
+        model_path: Optional[str] = None,
+        model_name: Optional[str] = None,
+        modules_map: Optional[
+            Union[dict[str, Union[None, str, list[str]]], str]
+        ] = None,
+        custom_objects_map: Optional[
+            Union[dict[str, Union[str, list[str]]], str]
+        ] = None,
+        custom_objects_directory: Optional[str] = None,
         context: mlrun.MLClientCtx = None,
         model_format: str = ModelFormats.PKL,
         **kwargs,
@@ -152,7 +156,7 @@ class XGBoostModelHandler(MLModelHandler):
                 )
 
     @without_mlrun_interface(interface=XGBModelMLRunInterface)
-    def save(self, output_path: str = None, **kwargs):
+    def save(self, output_path: Optional[str] = None, **kwargs):
         """
         Save the handled model at the given output path. If a MLRun context is available, the saved model files will be
         logged and returned as artifacts.
@@ -186,10 +190,10 @@ class XGBoostModelHandler(MLModelHandler):
 
     def to_onnx(
         self,
-        model_name: str = None,
+        model_name: Optional[str] = None,
         optimize: bool = True,
         input_sample: XGBoostTypes = None,
-        log: bool = None,
+        log: Optional[bool] = None,
     ):
         """
         Convert the model in this handler to an ONNX model. The inputs names are optional, they do not change the

mlrun/launcher/base.py

@@ -18,8 +18,6 @@ import os
 import uuid
 from typing import Any, Callable, Optional, Union
 
-import mlrun_pipelines.common.ops
-
 import mlrun.common.schemas
 import mlrun.config
 import mlrun.errors
@@ -27,6 +25,7 @@ import mlrun.lists
 import mlrun.model
 import mlrun.runtimes
 import mlrun.utils.regex
+import mlrun_pipelines.common.ops
 from mlrun.utils import logger
 
 run_modes = ["pass"]
@@ -62,7 +61,7 @@ class BaseLauncher(abc.ABC):
         schedule: Optional[
             Union[str, mlrun.common.schemas.schedule.ScheduleCronTrigger]
         ] = None,
-        hyperparams: dict[str, list] = None,
+        hyperparams: Optional[dict[str, list]] = None,
         hyper_param_options: Optional[mlrun.model.HyperParamOptions] = None,
         verbose: Optional[bool] = None,
         scrape_metrics: Optional[bool] = None,
@@ -238,7 +237,7 @@ class BaseLauncher(abc.ABC):
         out_path=None,
         artifact_path=None,
         workdir=None,
-        notifications: list[mlrun.model.Notification] = None,
+        notifications: Optional[list[mlrun.model.Notification]] = None,
         state_thresholds: Optional[dict[str, int]] = None,
     ):
         run.spec.handler = (

mlrun/launcher/local.py

@@ -59,7 +59,7 @@ class ClientLocalLauncher(launcher.ClientBaseLauncher):
         schedule: Optional[
             Union[str, mlrun.common.schemas.schedule.ScheduleCronTrigger]
         ] = None,
-        hyperparams: dict[str, list] = None,
+        hyperparams: Optional[dict[str, list]] = None,
         hyper_param_options: Optional[mlrun.model.HyperParamOptions] = None,
         verbose: Optional[bool] = None,
         scrape_metrics: Optional[bool] = None,

mlrun/launcher/remote.py

@@ -49,7 +49,7 @@ class ClientRemoteLauncher(launcher.ClientBaseLauncher):
         schedule: Optional[
             Union[str, mlrun.common.schemas.schedule.ScheduleCronTrigger]
         ] = None,
-        hyperparams: dict[str, list] = None,
+        hyperparams: Optional[dict[str, list]] = None,
         hyper_param_options: Optional[mlrun.model.HyperParamOptions] = None,
         verbose: Optional[bool] = None,
         scrape_metrics: Optional[bool] = None,

mlrun/lists.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from copy import copy
+from typing import Optional
 
 import pandas as pd
 
@@ -129,11 +130,11 @@ class RunList(list):
     def compare(
         self,
         hide_identical: bool = True,
-        exclude: list = None,
-        show: bool = None,
+        exclude: Optional[list] = None,
+        show: Optional[bool] = None,
         extend_iterations=True,
         filename=None,
-        colorscale: str = None,
+        colorscale: Optional[str] = None,
     ):
         """return/show parallel coordinates plot + table to compare between the list of runs
 

mlrun/model.py

@@ -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
@@ -141,7 +144,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 +187,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 +199,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 +218,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 +435,7 @@ class Credentials(ModelObj):
 
     def __init__(
         self,
-        access_key: str = None,
+        access_key: Optional[str] = None,
     ):
         self.access_key = access_key
 
@@ -500,7 +505,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 +554,7 @@ class ImageBuilder(ModelObj):
         self,
         image="",
         base_image=None,
-        commands: list = None,
+        commands: Optional[list] = None,
         secret=None,
         source=None,
         extra=None,
@@ -750,14 +755,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 +978,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":
@@ -1275,9 +1280,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 +1290,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 +1298,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 +1658,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 +1701,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 +1712,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 +1932,7 @@ class EntrypointParam(ModelObj):
         default=None,
         doc="",
         required=None,
-        choices: list = None,
+        choices: Optional[list] = None,
     ):
         self.name = name
         self.type = type
@@ -2065,12 +2127,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 +2154,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 +2186,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 +2204,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 +2215,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 +2272,7 @@ class DataTarget(DataTargetBase):
 
     def __init__(
         self,
-        kind: str = None,
+        kind: Optional[str] = None,
         name: str = "",
         path=None,
         online=None,
@@ -2237,12 +2301,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,8 +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.
-#
-# 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

mlrun/model_monitoring/api.py

@@ -45,7 +45,7 @@ def get_or_create_model_endpoint(
     endpoint_id: str = "",
     function_name: str = "",
     context: mlrun.MLClientCtx = None,
-    sample_set_statistics: dict[str, typing.Any] = 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,
@@ -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`
@@ -387,7 +387,7 @@ def _generate_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 +445,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

mlrun/model_monitoring/applications/_application_steps.py

@@ -26,7 +26,11 @@ from mlrun.serving.utils import StepToDict
 from mlrun.utils import logger
 
 from .context import MonitoringApplicationContext
-from .results import ModelMonitoringApplicationMetric, ModelMonitoringApplicationResult
+from .results import (
+    ModelMonitoringApplicationMetric,
+    ModelMonitoringApplicationResult,
+    _ModelMonitoringApplicationStats,
+)
 
 
 class _PushToMonitoringWriter(StepToDict):
@@ -61,7 +65,9 @@ class _PushToMonitoringWriter(StepToDict):
         event: tuple[
             list[
                 Union[
-                    ModelMonitoringApplicationResult, ModelMonitoringApplicationMetric
+                    ModelMonitoringApplicationResult,
+                    ModelMonitoringApplicationMetric,
+                    _ModelMonitoringApplicationStats,
                 ]
             ],
             MonitoringApplicationContext,
@@ -90,21 +96,15 @@ class _PushToMonitoringWriter(StepToDict):
                 writer_event[mm_constant.WriterEvent.EVENT_KIND] = (
                     mm_constant.WriterEventKind.RESULT
                 )
-                data[mm_constant.ResultData.CURRENT_STATS] = json.dumps(
-                    application_context.sample_df_stats
+            elif isinstance(result, _ModelMonitoringApplicationStats):
+                writer_event[mm_constant.WriterEvent.EVENT_KIND] = (
+                    mm_constant.WriterEventKind.STATS
                 )
-                writer_event[mm_constant.WriterEvent.DATA] = json.dumps(data)
             else:
                 writer_event[mm_constant.WriterEvent.EVENT_KIND] = (
                     mm_constant.WriterEventKind.METRIC
                 )
-                writer_event[mm_constant.WriterEvent.DATA] = json.dumps(data)
-
-            writer_event[mm_constant.WriterEvent.EVENT_KIND] = (
-                mm_constant.WriterEventKind.RESULT
-                if isinstance(result, ModelMonitoringApplicationResult)
-                else mm_constant.WriterEventKind.METRIC
-            )
+            writer_event[mm_constant.WriterEvent.DATA] = json.dumps(data)
             logger.info(
                 f"Pushing data = {writer_event} \n to stream = {self.stream_uri}"
             )
@@ -113,9 +113,7 @@ class _PushToMonitoringWriter(StepToDict):
 
     def _lazy_init(self):
         if self.output_stream is None:
-            self.output_stream = mlrun.datastore.get_stream_pusher(
-                self.stream_uri,
-            )
+            self.output_stream = mlrun.datastore.get_stream_pusher(self.stream_uri)
 
 
 class _PrepareMonitoringEvent(StepToDict):