mlrun 1.7.2rc4__py3-none-any.whl → 1.8.0__py3-none-any.whl

mlrun/model_monitoring/db/__init__.py

@@ -12,7 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .stores import ObjectStoreFactory, get_store_object
-from .stores.base import StoreBase
 from .tsdb import get_tsdb_connector
 from .tsdb.base import TSDBConnector

mlrun/model_monitoring/db/_schedules.py

@@ -0,0 +1,242 @@
+# Copyright 2024 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 json
+from abc import ABC, abstractmethod
+from contextlib import AbstractContextManager
+from types import TracebackType
+from typing import Final, Optional
+
+import botocore.exceptions
+
+import mlrun.common.schemas as schemas
+import mlrun.errors
+import mlrun.model_monitoring.helpers
+from mlrun.utils import logger
+
+
+class ModelMonitoringSchedulesFileBase(AbstractContextManager, ABC):
+    DEFAULT_SCHEDULES: Final = {}
+    INITIAL_CONTENT = json.dumps(DEFAULT_SCHEDULES)
+    ENCODING = "utf-8"
+
+    def __init__(self):
+        self._item = self.get_data_item_object()
+        if self._item:
+            self._path = self._item.url
+            self._fs = self._item.store.filesystem
+            # `self._schedules` is an in-memory copy of the DB for all the applications for
+            # the same model endpoint.
+            self._schedules = self.DEFAULT_SCHEDULES.copy()
+            # Does `self._schedules` hold the content of `self._item`?
+            self._open_schedules = False
+
+    @abstractmethod
+    def get_data_item_object(self) -> mlrun.DataItem:
+        pass
+
+    def create(self) -> None:
+        """Create a schedules file with initial content - an empty dictionary"""
+        logger.debug("Creating model monitoring schedules file", path=self._item.url)
+        self._item.put(self.INITIAL_CONTENT)
+
+    def delete(self) -> None:
+        """Delete schedules file if it exists"""
+        if (
+            self._fs is None  # In-memory store
+            or self._fs.exists(self._path)
+        ):
+            logger.debug(
+                "Deleting model monitoring schedules file", path=self._item.url
+            )
+            self._item.delete()
+        else:
+            logger.debug(
+                "Model monitoring schedules file does not exist, nothing to delete",
+                path=self._item.url,
+            )
+
+    def _open(self) -> None:
+        try:
+            content = self._item.get()
+        except (
+            mlrun.errors.MLRunNotFoundError,
+            # Different errors are raised for S3 or local storage, see ML-8042
+            botocore.exceptions.ClientError,
+            FileNotFoundError,
+        ) as err:
+            if (
+                isinstance(err, botocore.exceptions.ClientError)
+                # Add a log only to "NoSuchKey" errors codes - equivalent to `FileNotFoundError`
+                and err.response["Error"]["Code"] != "NoSuchKey"
+            ):
+                raise
+
+            logger.exception(
+                "The schedules file was not found. It should have been created "
+                "as a part of the model endpoint's creation",
+                path=self._path,
+            )
+            raise
+
+        if isinstance(content, bytes):
+            content = content.decode(encoding=self.ENCODING)
+        self._schedules = json.loads(content)
+        self._open_schedules = True
+
+    def _close(self) -> None:
+        self._item.put(json.dumps(self._schedules))
+        self._schedules = self.DEFAULT_SCHEDULES
+        self._open_schedules = False
+
+    def __enter__(self) -> "ModelMonitoringSchedulesFileBase":
+        self._open()
+        return super().__enter__()
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        self._close()
+
+    def _check_open_schedules(self) -> None:
+        if not self._open_schedules:
+            raise mlrun.errors.MLRunValueError(
+                "Open the schedules file as a context manager first"
+            )
+
+
+class ModelMonitoringSchedulesFileEndpoint(ModelMonitoringSchedulesFileBase):
+    def __init__(self, project: str, endpoint_id: str) -> None:
+        """
+        Initialize applications monitoring schedules file object.
+        The JSON file stores a dictionary of registered application name as key and Unix timestamp as value.
+        When working with the schedules data, use this class as a context manager to read and write the data.
+
+        :param project:     The project name.
+        :param endpoint_id: The endpoint ID.
+        """
+        # `self._item` is the persistent version of the monitoring schedules.
+        self._project = project
+        self._endpoint_id = endpoint_id
+        super().__init__()
+
+    def get_data_item_object(self) -> mlrun.DataItem:
+        return mlrun.model_monitoring.helpers.get_monitoring_schedules_endpoint_data(
+            project=self._project, endpoint_id=self._endpoint_id
+        )
+
+    @classmethod
+    def from_model_endpoint(
+        cls, model_endpoint: schemas.ModelEndpoint
+    ) -> "ModelMonitoringSchedulesFileEndpoint":
+        return cls(
+            project=model_endpoint.metadata.project,
+            endpoint_id=model_endpoint.metadata.uid,
+        )
+
+    def get_application_time(self, application: str) -> Optional[int]:
+        self._check_open_schedules()
+        return self._schedules.get(application)
+
+    def update_application_time(self, application: str, timestamp: int) -> None:
+        self._check_open_schedules()
+        self._schedules[application] = timestamp
+
+    def get_application_list(self) -> set[str]:
+        self._check_open_schedules()
+        return set(self._schedules.keys())
+
+    def get_min_timestamp(self) -> Optional[int]:
+        self._check_open_schedules()
+        return min(self._schedules.values(), default=None)
+
+
+class ModelMonitoringSchedulesFileChief(ModelMonitoringSchedulesFileBase):
+    def __init__(self, project: str) -> None:
+        """
+        Initialize applications monitoring schedules chief file object.
+        The JSON file stores a dictionary of registered model endpoints uid as key and point to a dictionary of
+        "last_request" and "last_analyzed" mapped to two Unix timestamps as values.
+        When working with the schedules data, use this class as a context manager to read and write the data.
+
+        :param project:     The project name.
+        """
+        # `self._item` is the persistent version of the monitoring schedules.
+        self._project = project
+        super().__init__()
+
+    def get_data_item_object(self) -> mlrun.DataItem:
+        return mlrun.model_monitoring.helpers.get_monitoring_schedules_chief_data(
+            project=self._project
+        )
+
+    def get_endpoint_last_request(self, endpoint_uid: str) -> Optional[int]:
+        self._check_open_schedules()
+        if endpoint_uid in self._schedules:
+            return self._schedules[endpoint_uid].get(
+                schemas.model_monitoring.constants.ScheduleChiefFields.LAST_REQUEST
+            )
+        else:
+            return None
+
+    def update_endpoint_timestamps(
+        self, endpoint_uid: str, last_request: int, last_analyzed: int
+    ) -> None:
+        self._check_open_schedules()
+        self._schedules[endpoint_uid] = {
+            schemas.model_monitoring.constants.ScheduleChiefFields.LAST_REQUEST: last_request,
+            schemas.model_monitoring.constants.ScheduleChiefFields.LAST_ANALYZED: last_analyzed,
+        }
+
+    def get_endpoint_last_analyzed(self, endpoint_uid: str) -> Optional[int]:
+        self._check_open_schedules()
+        if endpoint_uid in self._schedules:
+            return self._schedules[endpoint_uid].get(
+                schemas.model_monitoring.constants.ScheduleChiefFields.LAST_ANALYZED
+            )
+        else:
+            return None
+
+    def get_endpoint_list(self) -> set[str]:
+        self._check_open_schedules()
+        return set(self._schedules.keys())
+
+    def get_or_create(self) -> None:
+        try:
+            self._open()
+        except (
+            mlrun.errors.MLRunNotFoundError,
+            # Different errors are raised for S3 or local storage, see ML-8042
+            botocore.exceptions.ClientError,
+            FileNotFoundError,
+        ):
+            self.create()
+
+
+def delete_model_monitoring_schedules_folder(project: str) -> None:
+    """Delete the model monitoring schedules folder of the project"""
+    folder = mlrun.model_monitoring.helpers._get_monitoring_schedules_folder_path(
+        project
+    )
+    fs = mlrun.datastore.store_manager.object(folder).store.filesystem
+    if fs and fs.exists(folder):
+        logger.debug("Deleting model monitoring schedules folder", folder=folder)
+        fs.rm(folder, recursive=True)
+    elif fs is None:  # In-memory store
+        raise mlrun.errors.MLRunValueError(
+            "Cannot delete a folder without a file-system"
+        )

mlrun/model_monitoring/db/_stats.py

@@ -0,0 +1,189 @@
+# Copyright 2024 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 abc
+import json
+from abc import abstractmethod
+from datetime import datetime, timezone
+from typing import cast
+
+import botocore.exceptions
+import fsspec
+
+import mlrun.datastore.base
+from mlrun.common.schemas.model_monitoring.constants import StatsKind
+from mlrun.model_monitoring.helpers import (
+    get_monitoring_current_stats_data,
+    get_monitoring_drift_measures_data,
+    get_monitoring_stats_directory_path,
+)
+from mlrun.utils import logger
+
+
+class ModelMonitoringStatsFile(abc.ABC):
+    """
+    Abstract class
+    Initialize applications monitoring stats file object.
+    The JSON file stores a dictionary of registered application name as key and Unix timestamp as value.
+    When working with the schedules data, use this class as a context manager to read and write the data.
+    """
+
+    def __init__(self, item: mlrun.datastore.base.DataItem, file_type: str):
+        self._path = item.url
+        self._item = item
+        self._file_type = file_type
+        self._fs = cast(fsspec.AbstractFileSystem, self._item.store.filesystem)
+
+    def create(self) -> None:
+        """Create a json file with initial content - an empty dictionary"""
+        logger.debug(
+            f"Creating model monitoring {self._file_type} file", path=self._item.url
+        )
+        self._item.put(
+            json.dumps(
+                {
+                    "data": dict(),
+                    "timestamp": mlrun.utils.datetime_now().isoformat(
+                        sep=" ", timespec="microseconds"
+                    ),
+                }
+            )
+        )
+
+    def delete(self) -> None:
+        """Delete json file if it exists"""
+        if self._fs.exists(self._path):
+            logger.debug(
+                f"Deleting model monitoring {self._file_type} file", path=self._item.url
+            )
+            self._item.delete()
+        else:
+            logger.debug(
+                f"Model monitoring {self._file_type} file does not exist, nothing to delete",
+                path=self._item.url,
+            )
+
+    def read(self) -> tuple[dict, datetime]:
+        """
+        Read the stats data and timestamp saved in file
+        :return: tuple[dict, str] dictionary with stats data and timestamp saved in file
+        """
+        try:
+            content = json.loads(self._item.get().decode())
+            timestamp = content.get("timestamp")
+            if timestamp is not None:
+                timestamp = datetime.fromisoformat(timestamp).astimezone(
+                    tz=timezone.utc
+                )
+            return content.get("data"), timestamp
+        except (
+            mlrun.errors.MLRunNotFoundError,
+            # Different errors are raised for S3 or local storage, see ML-8042
+            botocore.exceptions.ClientError,
+            FileNotFoundError,
+        ) as err:
+            if (
+                isinstance(err, botocore.exceptions.ClientError)
+                # Add a log only to "NoSuchKey" errors codes - equivalent to `FileNotFoundError`
+                and err.response["Error"]["Code"] != "NoSuchKey"
+            ):
+                raise
+
+            logger.exception(
+                "The Stats file was not found. It should have been created "
+                "as a part of the model endpoint's creation",
+                path=self._path,
+                error=err,
+            )
+            raise
+
+    def write(self, stats: dict, timestamp: datetime) -> None:
+        """
+        Write stats data to file overwrite the existing file
+        :param stats: dictionary with the stats data
+        :param timestamp: datetime object with the timestamp of last entry point for the stats calculation
+        """
+        content = {
+            "data": stats,
+            "timestamp": timestamp.isoformat(sep=" ", timespec="microseconds"),
+        }
+        self._item.put(json.dumps(content))
+
+    @classmethod
+    @abstractmethod
+    def from_model_endpoint(
+        cls, model_endpoint: mlrun.common.schemas.ModelEndpoint
+    ) -> "ModelMonitoringStatsFile":
+        """
+        Return ModelMonitoringStatsFile child object using ModelEndpoint metadata
+        :param model_endpoint: The current model endpoint to get a stats object for
+        :return: ModelMonitoringStatsFile child object instance
+        """
+        pass
+
+
+class ModelMonitoringCurrentStatsFile(ModelMonitoringStatsFile):
+    def __init__(self, project: str, endpoint_id: str) -> None:
+        """
+        Initialize File object specific for current stats.
+        :param project:         (str) Project name
+        :param endpoint_id:     (str) Endpoint name
+        """
+        super().__init__(
+            get_monitoring_current_stats_data(project, endpoint_id),
+            StatsKind.CURRENT_STATS.value,
+        )
+
+    @classmethod
+    def from_model_endpoint(
+        cls, model_endpoint: mlrun.common.schemas.ModelEndpoint
+    ) -> "ModelMonitoringCurrentStatsFile":
+        return cls(
+            project=model_endpoint.metadata.project,
+            endpoint_id=model_endpoint.metadata.uid,
+        )
+
+
+class ModelMonitoringDriftMeasuresFile(ModelMonitoringStatsFile):
+    def __init__(self, project: str, endpoint_id: str) -> None:
+        """
+        Initialize File object specific for drift measures.
+        :param project:         (str) Project name
+        :param endpoint_id:     (str) Endpoint name
+        """
+        super().__init__(
+            get_monitoring_drift_measures_data(project, endpoint_id),
+            StatsKind.DRIFT_MEASURES.value,
+        )
+
+    @classmethod
+    def from_model_endpoint(
+        cls, model_endpoint: mlrun.common.schemas.ModelEndpoint
+    ) -> "ModelMonitoringDriftMeasuresFile":
+        return cls(
+            project=model_endpoint.metadata.project,
+            endpoint_id=model_endpoint.metadata.uid,
+        )
+
+
+def delete_model_monitoring_stats_folder(project: str) -> None:
+    """Delete the model monitoring schedules folder of the project"""
+    folder = get_monitoring_stats_directory_path(project)
+    fs = mlrun.datastore.store_manager.object(folder).store.filesystem
+    if fs and fs.exists(folder):
+        logger.debug("Deleting model monitoring stats folder", folder=folder)
+        fs.rm(folder, recursive=True)
+    elif fs is None:  # In-memory store
+        raise mlrun.errors.MLRunValueError(
+            "Cannot delete a folder without a file-system"
+        )

mlrun/model_monitoring/db/tsdb/__init__.py

@@ -16,7 +16,10 @@ import enum
 import typing
 
 import mlrun.common.schemas.secret
+import mlrun.datastore.datastore_profile
 import mlrun.errors
+import mlrun.model_monitoring.helpers
+from mlrun.datastore.datastore_profile import DatastoreProfile
 
 from .base import TSDBConnector
 
@@ -27,10 +30,13 @@ class ObjectTSDBFactory(enum.Enum):
     v3io_tsdb = "v3io-tsdb"
     tdengine = "tdengine"
 
-    def to_tsdb_connector(self, project: str, **kwargs) -> TSDBConnector:
+    def to_tsdb_connector(
+        self, project: str, profile: DatastoreProfile, **kwargs
+    ) -> TSDBConnector:
         """
         Return a TSDBConnector object based on the provided enum value.
         :param project: The name of the project.
+        :param profile: Datastore profile containing DSN and credentials for TSDB connection
         :return: `TSDBConnector` object.
         """
 
@@ -49,7 +55,7 @@ class ObjectTSDBFactory(enum.Enum):
 
         from .tdengine.tdengine_connector import TDEngineConnector
 
-        return TDEngineConnector(project=project, **kwargs)
+        return TDEngineConnector(project=project, profile=profile, **kwargs)
 
     @classmethod
     def _missing_(cls, value: typing.Any):
@@ -65,41 +71,46 @@ class ObjectTSDBFactory(enum.Enum):
 def get_tsdb_connector(
     project: str,
     secret_provider: typing.Optional[typing.Callable[[str], str]] = None,
-    tsdb_connection_string: typing.Optional[str] = None,
-    **kwargs,
+    profile: typing.Optional[mlrun.datastore.datastore_profile.DatastoreProfile] = None,
 ) -> TSDBConnector:
     """
     Get TSDB connector object.
     :param project:                 The name of the project.
     :param secret_provider:         An optional secret provider to get the connection string secret.
-    :param tsdb_connection_string:  An optional explicit connection string to the TSDB.
+    :param profile:                 An optional profile to initialize the TSDB connector from.
 
-    :return: `TSDBConnector` object. The main goal of this object is to handle different operations on the
+    :return: ``TSDBConnector`` object. The main goal of this object is to handle different operations on the
              TSDB connector such as updating drift metrics or write application record result.
-    :raise: `MLRunInvalidMMStoreTypeError` if the user didn't provide TSDB connection
-            or the provided TSDB connection is invalid.
+    :raise: ``MLRunNotFoundError`` if the user didn't set the TSDB datastore profile and didn't provide it through
+            the ``profile`` parameter.
+    :raise: ``MLRunInvalidMMStoreTypeError`` if the TSDB datastore profile is of an invalid type.
     """
-
-    tsdb_connection_string = (
-        tsdb_connection_string
-        or mlrun.model_monitoring.helpers.get_tsdb_connection_string(
-            secret_provider=secret_provider
-        )
+    profile = profile or mlrun.model_monitoring.helpers._get_tsdb_profile(
+        project=project, secret_provider=secret_provider
     )
-
-    if tsdb_connection_string and tsdb_connection_string.startswith("taosws"):
-        tsdb_connector_type = mlrun.common.schemas.model_monitoring.TSDBTarget.TDEngine
-        kwargs["connection_string"] = tsdb_connection_string
-    elif tsdb_connection_string and tsdb_connection_string == "v3io":
+    kwargs = {}
+    if isinstance(profile, mlrun.datastore.datastore_profile.DatastoreProfileV3io):
         tsdb_connector_type = mlrun.common.schemas.model_monitoring.TSDBTarget.V3IO_TSDB
+    elif isinstance(
+        profile, mlrun.datastore.datastore_profile.DatastoreProfileTDEngine
+    ):
+        tsdb_connector_type = mlrun.common.schemas.model_monitoring.TSDBTarget.TDEngine
     else:
+        extra_message = (
+            ""
+            if profile
+            else " by using `project.set_model_monitoring_credentials` API"
+        )
         raise mlrun.errors.MLRunInvalidMMStoreTypeError(
-            "You must provide a valid tsdb store connection by using "
-            "set_model_monitoring_credentials API."
+            "You must provide a valid TSDB datastore profile"
+            f"{extra_message}. "
+            f"Found an unexpected profile of class: {type(profile)}"
         )
 
     # Get connector type value from ObjectTSDBFactory enum class
     tsdb_connector_factory = ObjectTSDBFactory(tsdb_connector_type)
 
     # Convert into TSDB connector object
-    return tsdb_connector_factory.to_tsdb_connector(project=project, **kwargs)
+    return tsdb_connector_factory.to_tsdb_connector(
+        project=project, profile=profile, **kwargs
+    )