mlrun 1.7.1rc10__py3-none-any.whl → 1.8.0rc11__py3-none-any.whl

mlrun/model_monitoring/writer.py

@@ -13,23 +13,29 @@
 # limitations under the License.
 
 import json
-from typing import Any, Callable, NewType
+from datetime import datetime, timezone
+from typing import Any, Callable, NewType, Optional
 
 import mlrun.common.model_monitoring
 import mlrun.common.schemas
 import mlrun.common.schemas.alert as alert_objects
 import mlrun.model_monitoring
 from mlrun.common.schemas.model_monitoring.constants import (
-    EventFieldType,
     HistogramDataDriftApplicationConstants,
     MetricData,
     ResultData,
     ResultKindApp,
     ResultStatusApp,
+    StatsData,
+    StatsKind,
     WriterEvent,
     WriterEventKind,
 )
 from mlrun.common.schemas.notification import NotificationKind, NotificationSeverity
+from mlrun.model_monitoring.db._stats import (
+    ModelMonitoringCurrentStatsFile,
+    ModelMonitoringDriftMeasuresFile,
+)
 from mlrun.model_monitoring.helpers import get_result_instance_fqn
 from mlrun.serving.utils import StepToDict
 from mlrun.utils import logger
@@ -105,7 +111,7 @@ class ModelMonitoringWriter(StepToDict):
     def __init__(
         self,
         project: str,
-        secret_provider: Callable = None,
+        secret_provider: Optional[Callable] = None,
     ) -> None:
         self.project = project
         self.name = project  # required for the deployment process
@@ -114,9 +120,6 @@ class ModelMonitoringWriter(StepToDict):
             notification_types=[NotificationKind.slack]
         )
 
-        self._app_result_store = mlrun.model_monitoring.get_store_object(
-            project=self.project, secret_provider=secret_provider
-        )
         self._tsdb_connector = mlrun.model_monitoring.get_tsdb_connector(
             project=self.project, secret_provider=secret_provider
         )
@@ -190,6 +193,8 @@ class ModelMonitoringWriter(StepToDict):
             expected_keys.extend(MetricData.list())
         elif kind == WriterEventKind.RESULT:
             expected_keys.extend(ResultData.list())
+        elif kind == WriterEventKind.STATS:
+            expected_keys.extend(StatsData.list())
         else:
             raise _WriterEventValueError(
                 f"Unknown event kind: {kind}, expected one of: {WriterEventKind.list()}"
@@ -198,16 +203,50 @@ class ModelMonitoringWriter(StepToDict):
         if missing_keys:
             raise _WriterEventValueError(
                 f"The received event misses some keys compared to the expected "
-                f"monitoring application event schema: {missing_keys}"
+                f"monitoring application event schema: {missing_keys} for event kind {kind}"
             )
 
         return result_event, kind
 
+    def write_stats(self, event: _AppResultEvent) -> None:
+        """
+        Write to file the application stats event
+        :param event: application stats event
+        """
+        endpoint_id = event[WriterEvent.ENDPOINT_ID]
+        logger.debug(
+            "Updating the model endpoint with stats",
+            endpoint_id=endpoint_id,
+        )
+        stat_kind = event.get(StatsData.STATS_NAME)
+        data, timestamp_str = event.get(StatsData.STATS), event.get(StatsData.TIMESTAMP)
+        timestamp = datetime.fromisoformat(timestamp_str).astimezone(tz=timezone.utc)
+        if stat_kind == StatsKind.CURRENT_STATS.value:
+            ModelMonitoringCurrentStatsFile(self.project, endpoint_id).write(
+                data, timestamp
+            )
+        elif stat_kind == StatsKind.DRIFT_MEASURES.value:
+            ModelMonitoringDriftMeasuresFile(self.project, endpoint_id).write(
+                data, timestamp
+            )
+        logger.info(
+            "Updating the model endpoint statistics",
+            endpoint_id=endpoint_id,
+            stats_kind=stat_kind,
+        )
+
     def do(self, event: _RawEvent) -> None:
         event, kind = self._reconstruct_event(event)
         logger.info("Starting to write event", event=event)
+        if (
+            kind == WriterEventKind.STATS
+            and event[WriterEvent.APPLICATION_NAME]
+            == HistogramDataDriftApplicationConstants.NAME
+        ):
+            self.write_stats(event)
+            logger.info("Model monitoring writer finished handling event")
+            return
         self._tsdb_connector.write_application_event(event=event.copy(), kind=kind)
-        self._app_result_store.write_application_event(event=event.copy(), kind=kind)
 
         logger.info("Completed event DB writes")
 
@@ -223,14 +262,9 @@ class ModelMonitoringWriter(StepToDict):
                 == ResultStatusApp.potential_detection.value
             )
         ):
-            endpoint_id = event[WriterEvent.ENDPOINT_ID]
-            endpoint_record = self._endpoints_records.setdefault(
-                endpoint_id,
-                self._app_result_store.get_model_endpoint(endpoint_id=endpoint_id),
-            )
             event_value = {
                 "app_name": event[WriterEvent.APPLICATION_NAME],
-                "model": endpoint_record.get(EventFieldType.MODEL),
+                "model": event[WriterEvent.ENDPOINT_NAME],
                 "model_endpoint_id": event[WriterEvent.ENDPOINT_ID],
                 "result_name": event[ResultData.RESULT_NAME],
                 "result_value": event[ResultData.RESULT_VALUE],
@@ -247,26 +281,4 @@ class ModelMonitoringWriter(StepToDict):
                 result_kind=event[ResultData.RESULT_KIND],
             )
 
-        if (
-            kind == WriterEventKind.RESULT
-            and event[WriterEvent.APPLICATION_NAME]
-            == HistogramDataDriftApplicationConstants.NAME
-            and event[ResultData.RESULT_NAME]
-            == HistogramDataDriftApplicationConstants.GENERAL_RESULT_NAME
-        ):
-            endpoint_id = event[WriterEvent.ENDPOINT_ID]
-            logger.info(
-                "Updating the model endpoint with metadata specific to the histogram "
-                "data drift app",
-                endpoint_id=endpoint_id,
-            )
-            attributes = json.loads(event[ResultData.RESULT_EXTRA_DATA])
-            attributes[EventFieldType.DRIFT_STATUS] = str(
-                attributes[EventFieldType.DRIFT_STATUS]
-            )
-            self._app_result_store.update_model_endpoint(
-                endpoint_id=endpoint_id,
-                attributes=attributes,
-            )
-
         logger.info("Model monitoring writer finished handling event")

mlrun/package/__init__.py

@@ -11,14 +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
 
 import functools
 import inspect
 from collections import OrderedDict
-from typing import Callable, Union
+from typing import Callable, Optional, Union
 
 from ..config import config
 from .context_handler import ContextHandler
@@ -40,8 +37,8 @@ from .utils import (
 
 
 def handler(
-    labels: dict[str, str] = None,
-    outputs: list[Union[str, dict[str, str]]] = None,
+    labels: Optional[dict[str, str]] = None,
+    outputs: Optional[list[Union[str, dict[str, str]]]] = None,
     inputs: Union[bool, dict[str, Union[str, type]]] = True,
 ):
     """

mlrun/package/context_handler.py

@@ -216,7 +216,7 @@ class ContextHandler:
                     )
             # Link packages:
             self._packagers_manager.link_packages(
-                additional_artifacts=self._context.artifacts,
+                additional_artifact_uris=self._context.artifact_uris,
                 additional_results=self._context.results,
             )
             # Log the packed results and artifacts:

mlrun/package/packager.py

@@ -14,7 +14,7 @@
 #
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 from mlrun.artifacts import Artifact
 from mlrun.datastore import DataItem
@@ -144,9 +144,9 @@ class Packager(ABC):
     def pack(
         self,
         obj: Any,
-        key: str = None,
-        artifact_type: str = None,
-        configurations: dict = None,
+        key: Optional[str] = None,
+        artifact_type: Optional[str] = None,
+        configurations: Optional[dict] = None,
     ) -> Union[tuple[Artifact, dict], dict]:
         """
         Pack an object as the given artifact type using the provided configurations.
@@ -165,8 +165,8 @@ class Packager(ABC):
     def unpack(
         self,
         data_item: DataItem,
-        artifact_type: str = None,
-        instructions: dict = None,
+        artifact_type: Optional[str] = None,
+        instructions: Optional[dict] = None,
     ) -> Any:
         """
         Unpack the data item's artifact by the provided type using the given instructions.
@@ -180,7 +180,10 @@ class Packager(ABC):
         pass
 
     def is_packable(
-        self, obj: Any, artifact_type: str = None, configurations: dict = None
+        self,
+        obj: Any,
+        artifact_type: Optional[str] = None,
+        configurations: Optional[dict] = None,
     ) -> bool:
         """
         Check if this packager can pack an object of the provided type as the provided artifact type.
@@ -212,7 +215,7 @@ class Packager(ABC):
         return True
 
     def is_unpackable(
-        self, data_item: DataItem, type_hint: type, artifact_type: str = None
+        self, data_item: DataItem, type_hint: type, artifact_type: Optional[str] = None
     ) -> bool:
         """
         Check if this packager can unpack an input according to the user-given type hint and the provided artifact type.
@@ -315,7 +318,7 @@ class Packager(ABC):
         )
 
     def get_data_item_local_path(
-        self, data_item: DataItem, add_to_future_clearing_path: bool = None
+        self, data_item: DataItem, add_to_future_clearing_path: Optional[bool] = None
     ) -> str:
         """
         Get the local path to the item handled by the data item provided. The local path can be the same as the data

mlrun/package/packagers/__init__.py

@@ -11,8 +11,6 @@
 # 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 .default_packager import DefaultPackager
 from .numpy_packagers import NumPySupportedFormat

mlrun/package/packagers/default_packager.py

@@ -15,7 +15,7 @@
 import inspect
 from abc import ABCMeta
 from types import MethodType
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 import docstring_parser
 
@@ -323,9 +323,9 @@ class DefaultPackager(Packager, metaclass=_DefaultPackagerMeta):
     def pack(
         self,
         obj: Any,
-        key: str = None,
-        artifact_type: str = None,
-        configurations: dict = None,
+        key: Optional[str] = None,
+        artifact_type: Optional[str] = None,
+        configurations: Optional[dict] = None,
     ) -> Union[tuple[Artifact, dict], dict]:
         """
         Pack an object as the given artifact type using the provided configurations.
@@ -361,8 +361,8 @@ class DefaultPackager(Packager, metaclass=_DefaultPackagerMeta):
     def unpack(
         self,
         data_item: DataItem,
-        artifact_type: str = None,
-        instructions: dict = None,
+        artifact_type: Optional[str] = None,
+        instructions: Optional[dict] = None,
     ) -> Any:
         """
         Unpack the data item's artifact by the provided type using the given instructions.
@@ -399,7 +399,10 @@ class DefaultPackager(Packager, metaclass=_DefaultPackagerMeta):
         return unpack_method(data_item, **instructions)
 
     def is_packable(
-        self, obj: Any, artifact_type: str = None, configurations: dict = None
+        self,
+        obj: Any,
+        artifact_type: Optional[str] = None,
+        configurations: Optional[dict] = None,
     ) -> bool:
         """
         Check if this packager can pack an object of the provided type as the provided artifact type.
@@ -480,10 +483,10 @@ class DefaultPackager(Packager, metaclass=_DefaultPackagerMeta):
         self,
         data_item: DataItem,
         pickle_module_name: str = DEFAULT_PICKLE_MODULE,
-        object_module_name: str = None,
-        python_version: str = None,
-        pickle_module_version: str = None,
-        object_module_version: str = None,
+        object_module_name: Optional[str] = None,
+        python_version: Optional[str] = None,
+        pickle_module_version: Optional[str] = None,
+        object_module_version: Optional[str] = None,
     ) -> Any:
         """
         Unpack the data item's object, unpickle it using the instructions, and return.

mlrun/package/packagers/numpy_packagers.py

@@ -16,7 +16,7 @@ import os
 import pathlib
 import tempfile
 from abc import ABC, abstractmethod
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -371,7 +371,10 @@ class NumPyNDArrayPackager(DefaultPackager):
         return artifact, {}
 
     def unpack_file(
-        self, data_item: DataItem, file_format: str = None, allow_pickle: bool = False
+        self,
+        data_item: DataItem,
+        file_format: Optional[str] = None,
+        allow_pickle: bool = False,
     ) -> np.ndarray:
         """
         Unpack a numppy array from file.
@@ -474,7 +477,7 @@ class _NumPyNDArrayCollectionPackager(DefaultPackager):
     def unpack_file(
         self,
         data_item: DataItem,
-        file_format: str = None,
+        file_format: Optional[str] = None,
         allow_pickle: bool = False,
     ) -> dict[str, np.ndarray]:
         """
@@ -548,7 +551,10 @@ class NumPyNDArrayDictPackager(_NumPyNDArrayCollectionPackager):
     PACKABLE_OBJECT_TYPE = dict[str, np.ndarray]
 
     def is_packable(
-        self, obj: Any, artifact_type: str = None, configurations: dict = None
+        self,
+        obj: Any,
+        artifact_type: Optional[str] = None,
+        configurations: Optional[dict] = None,
     ) -> bool:
         """
         Check if the object provided is a dictionary of numpy arrays.
@@ -602,7 +608,7 @@ class NumPyNDArrayDictPackager(_NumPyNDArrayCollectionPackager):
     def unpack_file(
         self,
         data_item: DataItem,
-        file_format: str = None,
+        file_format: Optional[str] = None,
         allow_pickle: bool = False,
     ) -> dict[str, np.ndarray]:
         """
@@ -633,7 +639,10 @@ class NumPyNDArrayListPackager(_NumPyNDArrayCollectionPackager):
     PACKABLE_OBJECT_TYPE = list[np.ndarray]
 
     def is_packable(
-        self, obj: Any, artifact_type: str = None, configurations: dict = None
+        self,
+        obj: Any,
+        artifact_type: Optional[str] = None,
+        configurations: Optional[dict] = None,
     ) -> bool:
         """
         Check if the object provided is a list of numpy arrays.
@@ -679,7 +688,7 @@ class NumPyNDArrayListPackager(_NumPyNDArrayCollectionPackager):
     def unpack_file(
         self,
         data_item: DataItem,
-        file_format: str = None,
+        file_format: Optional[str] = None,
         allow_pickle: bool = False,
     ) -> list[np.ndarray]:
         """

mlrun/package/packagers/pandas_packagers.py

@@ -17,7 +17,7 @@ import os
 import pathlib
 import tempfile
 from abc import ABC, abstractmethod
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 import pandas as pd
 
@@ -56,7 +56,7 @@ class _Formatter(ABC):
     @classmethod
     @abstractmethod
     def read(
-        cls, file_path: str, unflatten_kwargs: dict = None, **read_kwargs
+        cls, file_path: str, unflatten_kwargs: Optional[dict] = None, **read_kwargs
     ) -> pd.DataFrame:
         """
         Read the dataframe from the given file path.
@@ -173,7 +173,7 @@ class _ParquetFormatter(_Formatter):
 
     @classmethod
     def read(
-        cls, file_path: str, unflatten_kwargs: dict = None, **read_kwargs
+        cls, file_path: str, unflatten_kwargs: Optional[dict] = None, **read_kwargs
     ) -> pd.DataFrame:
         """
         Read the dataframe from the given parquet file path.
@@ -221,7 +221,7 @@ class _CSVFormatter(_Formatter):
 
     @classmethod
     def read(
-        cls, file_path: str, unflatten_kwargs: dict = None, **read_kwargs
+        cls, file_path: str, unflatten_kwargs: Optional[dict] = None, **read_kwargs
     ) -> pd.DataFrame:
         """
         Read the dataframe from the given csv file path.
@@ -275,7 +275,7 @@ class _H5Formatter(_Formatter):
 
     @classmethod
     def read(
-        cls, file_path: str, unflatten_kwargs: dict = None, **read_kwargs
+        cls, file_path: str, unflatten_kwargs: Optional[dict] = None, **read_kwargs
     ) -> pd.DataFrame:
         """
         Read the dataframe from the given h5 file path.
@@ -332,7 +332,7 @@ class _XMLFormatter(_Formatter):
 
     @classmethod
     def read(
-        cls, file_path: str, unflatten_kwargs: dict = None, **read_kwargs
+        cls, file_path: str, unflatten_kwargs: Optional[dict] = None, **read_kwargs
     ) -> pd.DataFrame:
         """
         Read the dataframe from the given xml file path.
@@ -391,7 +391,7 @@ class _XLSXFormatter(_Formatter):
 
     @classmethod
     def read(
-        cls, file_path: str, unflatten_kwargs: dict = None, **read_kwargs
+        cls, file_path: str, unflatten_kwargs: Optional[dict] = None, **read_kwargs
     ) -> pd.DataFrame:
         """
         Read the dataframe from the given xlsx file path.
@@ -449,7 +449,7 @@ class _HTMLFormatter(_Formatter):
 
     @classmethod
     def read(
-        cls, file_path: str, unflatten_kwargs: dict = None, **read_kwargs
+        cls, file_path: str, unflatten_kwargs: Optional[dict] = None, **read_kwargs
     ) -> pd.DataFrame:
         """
         Read dataframes from the given html file path.
@@ -510,7 +510,7 @@ class _JSONFormatter(_Formatter):
 
     @classmethod
     def read(
-        cls, file_path: str, unflatten_kwargs: dict = None, **read_kwargs
+        cls, file_path: str, unflatten_kwargs: Optional[dict] = None, **read_kwargs
     ) -> pd.DataFrame:
         """
         Read dataframes from the given json file path.
@@ -565,7 +565,7 @@ class _FeatherFormatter(_Formatter):
 
     @classmethod
     def read(
-        cls, file_path: str, unflatten_kwargs: dict = None, **read_kwargs
+        cls, file_path: str, unflatten_kwargs: Optional[dict] = None, **read_kwargs
     ) -> pd.DataFrame:
         """
         Read dataframes from the given feather file path.
@@ -620,7 +620,7 @@ class _ORCFormatter(_Formatter):
 
     @classmethod
     def read(
-        cls, file_path: str, unflatten_kwargs: dict = None, **read_kwargs
+        cls, file_path: str, unflatten_kwargs: Optional[dict] = None, **read_kwargs
     ) -> pd.DataFrame:
         """
         Read dataframes from the given orc file path.
@@ -730,7 +730,7 @@ class PandasDataFramePackager(DefaultPackager):
         self,
         obj: pd.DataFrame,
         key: str,
-        file_format: str = None,
+        file_format: Optional[str] = None,
         flatten: bool = True,
         **to_kwargs,
     ) -> tuple[Artifact, dict]:
@@ -785,8 +785,8 @@ class PandasDataFramePackager(DefaultPackager):
     def unpack_file(
         self,
         data_item: DataItem,
-        file_format: str = None,
-        read_kwargs: dict = None,
+        file_format: Optional[str] = None,
+        read_kwargs: Optional[dict] = None,
     ) -> pd.DataFrame:
         """
         Unpack a pandas dataframe from file.
@@ -883,7 +883,7 @@ class PandasSeriesPackager(PandasDataFramePackager):
         self,
         obj: pd.Series,
         key: str,
-        file_format: str = None,
+        file_format: Optional[str] = None,
         flatten: bool = True,
         **to_kwargs,
     ) -> tuple[Artifact, dict]:
@@ -919,9 +919,9 @@ class PandasSeriesPackager(PandasDataFramePackager):
     def unpack_file(
         self,
         data_item: DataItem,
-        file_format: str = None,
-        read_kwargs: dict = None,
-        column_name: Union[str, int] = None,
+        file_format: Optional[str] = None,
+        read_kwargs: Optional[dict] = None,
+        column_name: Optional[Union[str, int]] = None,
     ) -> pd.Series:
         """
         Unpack a pandas series from file.

mlrun/package/packagers/python_standard_library_packagers.py

@@ -15,7 +15,7 @@
 import os
 import pathlib
 import tempfile
-from typing import Union
+from typing import Optional, Union
 
 from mlrun.artifacts import Artifact
 from mlrun.datastore import DataItem
@@ -140,7 +140,7 @@ class StrPackager(DefaultPackager):
         self,
         data_item: DataItem,
         is_directory: bool = False,
-        archive_format: str = None,
+        archive_format: Optional[str] = None,
     ) -> str:
         """
         Unpack a data item representing a path string. If the path is of a file, the file is downloaded to a local
@@ -222,7 +222,7 @@ class _BuiltinCollectionPackager(DefaultPackager):
         return artifact, instructions
 
     def unpack_file(
-        self, data_item: DataItem, file_format: str = None
+        self, data_item: DataItem, file_format: Optional[str] = None
     ) -> Union[dict, list]:
         """
         Unpack a builtin collection from file.
@@ -259,7 +259,9 @@ class DictPackager(_BuiltinCollectionPackager):
 
     PACKABLE_OBJECT_TYPE = dict
 
-    def unpack_file(self, data_item: DataItem, file_format: str = None) -> dict:
+    def unpack_file(
+        self, data_item: DataItem, file_format: Optional[str] = None
+    ) -> dict:
         """
         Unpack a dictionary from file.
 
@@ -285,7 +287,9 @@ class ListPackager(_BuiltinCollectionPackager):
 
     PACKABLE_OBJECT_TYPE = list
 
-    def unpack_file(self, data_item: DataItem, file_format: str = None) -> list:
+    def unpack_file(
+        self, data_item: DataItem, file_format: Optional[str] = None
+    ) -> list:
         """
         Unpack a list from file.
 
@@ -355,7 +359,9 @@ class TuplePackager(ListPackager):
         """
         return super().pack_file(obj=list(obj), key=key, file_format=file_format)
 
-    def unpack_file(self, data_item: DataItem, file_format: str = None) -> tuple:
+    def unpack_file(
+        self, data_item: DataItem, file_format: Optional[str] = None
+    ) -> tuple:
         """
         Unpack a tuple from file.
 
@@ -400,7 +406,9 @@ class SetPackager(ListPackager):
         """
         return super().pack_file(obj=list(obj), key=key, file_format=file_format)
 
-    def unpack_file(self, data_item: DataItem, file_format: str = None) -> set:
+    def unpack_file(
+        self, data_item: DataItem, file_format: Optional[str] = None
+    ) -> set:
         """
         Unpack a set from file.
 
@@ -434,7 +442,9 @@ class FrozensetPackager(SetPackager):
         """
         return super().pack_file(obj=set(obj), key=key, file_format=file_format)
 
-    def unpack_file(self, data_item: DataItem, file_format: str = None) -> frozenset:
+    def unpack_file(
+        self, data_item: DataItem, file_format: Optional[str] = None
+    ) -> frozenset:
         """
         Unpack a frozenset from file.
 
@@ -481,7 +491,9 @@ class BytesPackager(ListPackager):
         """
         return super().pack_file(obj=list(obj), key=key, file_format=file_format)
 
-    def unpack_file(self, data_item: DataItem, file_format: str = None) -> bytes:
+    def unpack_file(
+        self, data_item: DataItem, file_format: Optional[str] = None
+    ) -> bytes:
         """
         Unpack a bytes from file.
 
@@ -526,7 +538,9 @@ class BytearrayPackager(BytesPackager):
         """
         return super().pack_file(obj=bytes(obj), key=key, file_format=file_format)
 
-    def unpack_file(self, data_item: DataItem, file_format: str = None) -> bytearray:
+    def unpack_file(
+        self, data_item: DataItem, file_format: Optional[str] = None
+    ) -> bytearray:
         """
         Unpack a bytearray from file.
 
@@ -585,7 +599,7 @@ class PathPackager(StrPackager):
         self,
         data_item: DataItem,
         is_directory: bool = False,
-        archive_format: str = None,
+        archive_format: Optional[str] = None,
     ) -> pathlib.Path:
         """
         Unpack a data item representing a `Path`. If the path is of a file, the file is downloaded to a local