mlrun 1.7.0rc17__py3-none-any.whl → 1.7.0rc18__py3-none-any.whl

mlrun/db/httpdb.py

@@ -30,6 +30,7 @@ import semver
 from mlrun_pipelines.utils import compile_pipeline
 
 import mlrun
+import mlrun.common.runtimes
 import mlrun.common.schemas
 import mlrun.common.types
 import mlrun.model_monitoring.model_endpoint
@@ -530,6 +531,10 @@ class HTTPRunDB(RunDBInterface):
                 server_cfg.get("model_endpoint_monitoring_endpoint_store_connection")
                 or config.model_endpoint_monitoring.endpoint_store_connection
             )
+            config.model_endpoint_monitoring.tsdb_connection = (
+                server_cfg.get("model_monitoring_tsdb_connection")
+                or config.model_endpoint_monitoring.tsdb_connection
+            )
             config.packagers = server_cfg.get("packagers") or config.packagers
             server_data_prefixes = server_cfg.get("feature_store_data_prefixes") or {}
             for prefix in ["default", "nosql", "redisnosql"]:
@@ -752,7 +757,10 @@ class HTTPRunDB(RunDBInterface):
         uid: Optional[Union[str, list[str]]] = None,
         project: Optional[str] = None,
         labels: Optional[Union[str, list[str]]] = None,
-        state: Optional[str] = None,
+        state: Optional[
+            mlrun.common.runtimes.constants.RunStates
+        ] = None,  # Backward compatibility
+        states: typing.Optional[list[mlrun.common.runtimes.constants.RunStates]] = None,
         sort: bool = True,
         last: int = 0,
         iter: bool = False,
@@ -790,7 +798,8 @@ class HTTPRunDB(RunDBInterface):
         :param labels: A list of labels to filter by. Label filters work by either filtering a specific value
             of a label (i.e. list("key=value")) or by looking for the existence of a given
             key (i.e. "key").
-        :param state: List only runs whose state is specified.
+        :param state: Deprecated - List only runs whose state is specified (will be removed in 1.9.0)
+        :param states: List only runs whose state is one of the provided states.
         :param sort: Whether to sort the result according to their start time. Otherwise, results will be
             returned by their internal order in the DB (order will not be guaranteed).
         :param last: Deprecated - currently not used (will be removed in 1.8.0).
@@ -826,11 +835,19 @@ class HTTPRunDB(RunDBInterface):
                 FutureWarning,
             )
 
+        if state:
+            # TODO: Remove this in 1.9.0
+            warnings.warn(
+                "'state' is deprecated and will be removed in 1.9.0. Use 'states' instead.",
+                FutureWarning,
+            )
+
         if (
             not name
             and not uid
             and not labels
             and not state
+            and not states
             and not last
             and not start_time_from
             and not start_time_to
@@ -849,7 +866,9 @@ class HTTPRunDB(RunDBInterface):
             "name": name,
             "uid": uid,
             "label": labels or [],
-            "state": state,
+            "state": mlrun.utils.helpers.as_list(state)
+            if state is not None
+            else states or None,
             "sort": bool2str(sort),
             "iter": bool2str(iter),
             "start_time_from": datetime_to_iso(start_time_from),

mlrun/db/nopdb.py

@@ -17,6 +17,7 @@ import datetime
 from typing import Optional, Union
 
 import mlrun.alerts
+import mlrun.common.runtimes.constants
 import mlrun.common.schemas
 import mlrun.errors
 
@@ -80,7 +81,10 @@ class NopDB(RunDBInterface):
         uid: Optional[Union[str, list[str]]] = None,
         project: Optional[str] = None,
         labels: Optional[Union[str, list[str]]] = None,
-        state: Optional[str] = None,
+        state: Optional[
+            mlrun.common.runtimes.constants.RunStates
+        ] = None,  # Backward compatibility
+        states: Optional[list[mlrun.common.runtimes.constants.RunStates]] = None,
         sort: bool = True,
         last: int = 0,
         iter: bool = False,

mlrun/errors.py

@@ -155,6 +155,10 @@ class MLRunNotFoundError(MLRunHTTPStatusError):
     error_status_code = HTTPStatus.NOT_FOUND.value
 
 
+class MLRunPaginationEndOfResultsError(MLRunNotFoundError):
+    pass
+
+
 class MLRunBadRequestError(MLRunHTTPStatusError):
     error_status_code = HTTPStatus.BAD_REQUEST.value
 
@@ -240,3 +244,5 @@ STATUS_ERRORS = {
     HTTPStatus.SERVICE_UNAVAILABLE.value: MLRunServiceUnavailableError,
     HTTPStatus.NOT_IMPLEMENTED.value: MLRunNotImplementedServerError,
 }
+
+EXPECTED_ERRORS = (MLRunPaginationEndOfResultsError,)

mlrun/feature_store/retrieval/conversion.py

@@ -168,10 +168,10 @@ class PandasConversionMixin:
         column_counter = Counter(self.columns)
 
         dtype = [None] * len(self.schema)
-        for fieldIdx, field in enumerate(self.schema):
+        for field_idx, field in enumerate(self.schema):
             # For duplicate column name, we use `iloc` to access it.
             if column_counter[field.name] > 1:
-                pandas_col = pdf.iloc[:, fieldIdx]
+                pandas_col = pdf.iloc[:, field_idx]
             else:
                 pandas_col = pdf[field.name]
 
@@ -187,12 +187,12 @@ class PandasConversionMixin:
                 and field.nullable
                 and pandas_col.isnull().any()
             ):
-                dtype[fieldIdx] = pandas_type
+                dtype[field_idx] = pandas_type
             # Ensure we fall back to nullable numpy types, even when whole column is null:
             if isinstance(field.dataType, IntegralType) and pandas_col.isnull().any():
-                dtype[fieldIdx] = np.float64
+                dtype[field_idx] = np.float64
             if isinstance(field.dataType, BooleanType) and pandas_col.isnull().any():
-                dtype[fieldIdx] = object
+                dtype[field_idx] = object
 
         df = pd.DataFrame()
         for index, t in enumerate(dtype):

mlrun/feature_store/retrieval/job.py

@@ -198,7 +198,8 @@ import mlrun.feature_store.retrieval
 from mlrun.datastore.targets import get_target_driver
 def merge_handler(context, vector_uri, target, entity_rows=None, 
                   entity_timestamp_column=None, drop_columns=None, with_indexes=None, query=None,
-                  engine_args=None, order_by=None, start_time=None, end_time=None, timestamp_for_filtering=None):
+                  engine_args=None, order_by=None, start_time=None, end_time=None, timestamp_for_filtering=None,
+                  additional_filters=None):
     vector = context.get_store_resource(vector_uri)
     store_target = get_target_driver(target, vector)
     if entity_rows:
@@ -208,7 +209,7 @@ def merge_handler(context, vector_uri, target, entity_rows=None,
     merger = mlrun.feature_store.retrieval.{{{engine}}}(vector, **(engine_args or {}))
     merger.start(entity_rows, entity_timestamp_column, store_target, drop_columns, with_indexes=with_indexes, 
                  query=query, order_by=order_by, start_time=start_time, end_time=end_time,
-                 timestamp_for_filtering=timestamp_for_filtering)
+                 timestamp_for_filtering=timestamp_for_filtering, additional_filters=additional_filters)
 
     target = vector.status.targets[store_target.name].to_dict()
     context.log_result('feature_vector', vector.uri)

mlrun/feature_store/retrieval/spark_merger.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+
 import pandas as pd
 import semver
 
@@ -252,13 +253,13 @@ class SparkFeatureMerger(BaseMerger):
         # handling case where there are multiple feature sets and user creates vector where
         # entity_timestamp_column is from a specific feature set (can't be entity timestamp)
         source_driver = mlrun.datastore.sources.source_kind_to_driver[source_kind]
-
         source = source_driver(
             name=self.vector.metadata.name,
             path=source_path,
             time_field=time_column,
             start_time=start_time,
             end_time=end_time,
+            additional_filters=additional_filters,
             **source_kwargs,
         )
 

mlrun/frameworks/_dl_common/loggers/tensorboard_logger.py

@@ -648,13 +648,13 @@ class TensorboardLogger(Logger, Generic[DLTypes.WeightType]):
         if isinstance(value, list):
             if len(value) == 0:
                 return ""
-            text = "\n" + yaml.dump(value)
+            text = "\n" + yaml.safe_dump(value)
             text = "  \n".join(["  " * tabs + line for line in text.splitlines()])
             return text
         if isinstance(value, dict):
             if len(value) == 0:
                 return ""
-            text = yaml.dump(value)
+            text = yaml.safe_dump(value)
             text = "  \n".join(
                 ["  " * tabs + "- " + line for line in text.splitlines()]
             )

mlrun/model_monitoring/db/stores/base/store.py

@@ -15,7 +15,7 @@
 import typing
 from abc import ABC, abstractmethod
 
-import mlrun.common.schemas.model_monitoring.constants as mm_constants
+import mlrun.common.schemas.model_monitoring as mm_schemas
 
 
 class StoreBase(ABC):
@@ -115,8 +115,8 @@ class StoreBase(ABC):
     def write_application_event(
         self,
         event: dict[str, typing.Any],
-        kind: mm_constants.WriterEventKind = mm_constants.WriterEventKind.RESULT,
-    ):
+        kind: mm_schemas.WriterEventKind = mm_schemas.WriterEventKind.RESULT,
+    ) -> None:
         """
         Write a new event in the target table.
 
@@ -157,3 +157,16 @@ class StoreBase(ABC):
 
         """
         pass
+
+    @abstractmethod
+    def get_model_endpoint_metrics(
+        self, endpoint_id: str, type: mm_schemas.ModelEndpointMonitoringMetricType
+    ) -> list[mm_schemas.ModelEndpointMonitoringMetric]:
+        """
+        Get the model monitoring results and metrics of the requested model endpoint.
+
+        :param: endpoint_id: The model endpoint identifier.
+        :param: type:        The type of the requested metrics ("result" or "metric").
+
+        :return:            A list of the available metrics.
+        """

mlrun/model_monitoring/db/stores/sqldb/sql_store.py

@@ -21,7 +21,7 @@ import pandas as pd
 import sqlalchemy
 
 import mlrun.common.model_monitoring.helpers
-import mlrun.common.schemas.model_monitoring as mm_constants
+import mlrun.common.schemas.model_monitoring as mm_schemas
 import mlrun.model_monitoring.db
 import mlrun.model_monitoring.db.stores.sqldb.models
 import mlrun.model_monitoring.helpers
@@ -71,7 +71,7 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
                 connection_string=self._sql_connection_string
             )
         )
-        self._tables[mm_constants.EventFieldType.MODEL_ENDPOINTS] = (
+        self._tables[mm_schemas.EventFieldType.MODEL_ENDPOINTS] = (
             self.ModelEndpointsTable
         )
 
@@ -81,7 +81,7 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
                 connection_string=self._sql_connection_string
             )
         )
-        self._tables[mm_constants.FileTargetKind.APP_RESULTS] = (
+        self._tables[mm_schemas.FileTargetKind.APP_RESULTS] = (
             self.ApplicationResultsTable
         )
 
@@ -89,7 +89,7 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         self.MonitoringSchedulesTable = mlrun.model_monitoring.db.stores.sqldb.models._get_monitoring_schedules_table(
             connection_string=self._sql_connection_string
         )
-        self._tables[mm_constants.FileTargetKind.MONITORING_SCHEDULES] = (
+        self._tables[mm_schemas.FileTargetKind.MONITORING_SCHEDULES] = (
             self.MonitoringSchedulesTable
         )
 
@@ -182,12 +182,12 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         """
 
         # Adjust timestamps fields
-        endpoint[mm_constants.EventFieldType.FIRST_REQUEST] = (endpoint)[
-            mm_constants.EventFieldType.LAST_REQUEST
+        endpoint[mm_schemas.EventFieldType.FIRST_REQUEST] = (endpoint)[
+            mm_schemas.EventFieldType.LAST_REQUEST
         ] = mlrun.utils.datetime_now()
 
         self._write(
-            table=mm_constants.EventFieldType.MODEL_ENDPOINTS,
+            table=mm_schemas.EventFieldType.MODEL_ENDPOINTS,
             event=endpoint,
         )
 
@@ -204,9 +204,9 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         """
         self._init_model_endpoints_table()
 
-        attributes.pop(mm_constants.EventFieldType.ENDPOINT_ID, None)
+        attributes.pop(mm_schemas.EventFieldType.ENDPOINT_ID, None)
 
-        filter_endpoint = {mm_constants.EventFieldType.UID: endpoint_id}
+        filter_endpoint = {mm_schemas.EventFieldType.UID: endpoint_id}
 
         self._update(
             attributes=attributes, table=self.ModelEndpointsTable, **filter_endpoint
@@ -220,7 +220,7 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         """
         self._init_model_endpoints_table()
 
-        filter_endpoint = {mm_constants.EventFieldType.UID: endpoint_id}
+        filter_endpoint = {mm_schemas.EventFieldType.UID: endpoint_id}
         # Delete the model endpoint record using sqlalchemy ORM
         self._delete(table=self.ModelEndpointsTable, **filter_endpoint)
 
@@ -240,7 +240,7 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         self._init_model_endpoints_table()
 
         # Get the model endpoint record using sqlalchemy ORM
-        filter_endpoint = {mm_constants.EventFieldType.UID: endpoint_id}
+        filter_endpoint = {mm_schemas.EventFieldType.UID: endpoint_id}
         endpoint_record = self._get(table=self.ModelEndpointsTable, **filter_endpoint)
 
         if not endpoint_record:
@@ -292,32 +292,32 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
                 query = self._filter_values(
                     query=query,
                     model_endpoints_table=model_endpoints_table,
-                    key_filter=mm_constants.EventFieldType.MODEL,
+                    key_filter=mm_schemas.EventFieldType.MODEL,
                     filtered_values=[model],
                 )
             if function:
                 query = self._filter_values(
                     query=query,
                     model_endpoints_table=model_endpoints_table,
-                    key_filter=mm_constants.EventFieldType.FUNCTION,
+                    key_filter=mm_schemas.EventFieldType.FUNCTION,
                     filtered_values=[function],
                 )
             if uids:
                 query = self._filter_values(
                     query=query,
                     model_endpoints_table=model_endpoints_table,
-                    key_filter=mm_constants.EventFieldType.UID,
+                    key_filter=mm_schemas.EventFieldType.UID,
                     filtered_values=uids,
                     combined=False,
                 )
             if top_level:
-                node_ep = str(mm_constants.EndpointType.NODE_EP.value)
-                router_ep = str(mm_constants.EndpointType.ROUTER.value)
+                node_ep = str(mm_schemas.EndpointType.NODE_EP.value)
+                router_ep = str(mm_schemas.EndpointType.ROUTER.value)
                 endpoint_types = [node_ep, router_ep]
                 query = self._filter_values(
                     query=query,
                     model_endpoints_table=model_endpoints_table,
-                    key_filter=mm_constants.EventFieldType.ENDPOINT_TYPE,
+                    key_filter=mm_schemas.EventFieldType.ENDPOINT_TYPE,
                     filtered_values=endpoint_types,
                     combined=False,
                 )
@@ -338,8 +338,8 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
     def write_application_event(
         self,
         event: dict[str, typing.Any],
-        kind: mm_constants.WriterEventKind = mm_constants.WriterEventKind.RESULT,
-    ):
+        kind: mm_schemas.WriterEventKind = mm_schemas.WriterEventKind.RESULT,
+    ) -> None:
         """
         Write a new application event in the target table.
 
@@ -349,16 +349,14 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         :param kind: The type of the event, can be either "result" or "metric".
         """
 
-        if kind == mm_constants.WriterEventKind.METRIC:
+        if kind == mm_schemas.WriterEventKind.METRIC:
             # TODO : Implement the logic for writing metrics to MySQL
             return
 
         self._init_application_results_table()
 
         application_filter_dict = {
-            mm_constants.EventFieldType.UID: self._generate_application_result_uid(
-                event
-            )
+            mm_schemas.EventFieldType.UID: self._generate_application_result_uid(event)
         }
 
         application_record = self._get(
@@ -367,11 +365,11 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         if application_record:
             self._convert_to_datetime(
                 event=event,
-                key=mm_constants.WriterEvent.START_INFER_TIME,
+                key=mm_schemas.WriterEvent.START_INFER_TIME,
             )
             self._convert_to_datetime(
                 event=event,
-                key=mm_constants.WriterEvent.END_INFER_TIME,
+                key=mm_schemas.WriterEvent.END_INFER_TIME,
             )
             # Update an existing application result
             self._update(
@@ -381,12 +379,12 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
             )
         else:
             # Write a new application result
-            event[mm_constants.EventFieldType.UID] = application_filter_dict[
-                mm_constants.EventFieldType.UID
+            event[mm_schemas.EventFieldType.UID] = application_filter_dict[
+                mm_schemas.EventFieldType.UID
             ]
 
             self._write(
-                table=mm_constants.FileTargetKind.APP_RESULTS,
+                table=mm_schemas.FileTargetKind.APP_RESULTS,
                 event=event,
             )
 
@@ -398,11 +396,11 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
     @staticmethod
     def _generate_application_result_uid(event: dict[str, typing.Any]) -> str:
         return (
-            event[mm_constants.WriterEvent.ENDPOINT_ID]
+            event[mm_schemas.WriterEvent.ENDPOINT_ID]
             + "_"
-            + event[mm_constants.WriterEvent.APPLICATION_NAME]
+            + event[mm_schemas.WriterEvent.APPLICATION_NAME]
             + "_"
-            + event[mm_constants.ResultData.RESULT_NAME]
+            + event[mm_schemas.ResultData.RESULT_NAME]
         )
 
     def get_last_analyzed(self, endpoint_id: str, application_name: str) -> int:
@@ -452,17 +450,17 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         if not monitoring_schedule_record:
             # Add a new record with empty last analyzed value
             self._write(
-                table=mm_constants.FileTargetKind.MONITORING_SCHEDULES,
+                table=mm_schemas.FileTargetKind.MONITORING_SCHEDULES,
                 event={
-                    mm_constants.SchedulingKeys.UID: uuid.uuid4().hex,
-                    mm_constants.SchedulingKeys.APPLICATION_NAME: application_name,
-                    mm_constants.SchedulingKeys.ENDPOINT_ID: endpoint_id,
-                    mm_constants.SchedulingKeys.LAST_ANALYZED: last_analyzed,
+                    mm_schemas.SchedulingKeys.UID: uuid.uuid4().hex,
+                    mm_schemas.SchedulingKeys.APPLICATION_NAME: application_name,
+                    mm_schemas.SchedulingKeys.ENDPOINT_ID: endpoint_id,
+                    mm_schemas.SchedulingKeys.LAST_ANALYZED: last_analyzed,
                 },
             )
 
         self._update(
-            attributes={mm_constants.SchedulingKeys.LAST_ANALYZED: last_analyzed},
+            attributes={mm_schemas.SchedulingKeys.LAST_ANALYZED: last_analyzed},
             table=self.MonitoringSchedulesTable,
             **application_filter_dict,
         )
@@ -558,7 +556,7 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
 
         # Convert endpoint labels into dictionary
         endpoint_labels = json.loads(
-            endpoint_dict.get(mm_constants.EventFieldType.LABELS)
+            endpoint_dict.get(mm_schemas.EventFieldType.LABELS)
         )
 
         for label in labels:
@@ -585,11 +583,9 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
             )
         application_filter_dict = {}
         if endpoint_id:
-            application_filter_dict[mm_constants.SchedulingKeys.ENDPOINT_ID] = (
-                endpoint_id
-            )
+            application_filter_dict[mm_schemas.SchedulingKeys.ENDPOINT_ID] = endpoint_id
         if application_name:
-            application_filter_dict[mm_constants.SchedulingKeys.APPLICATION_NAME] = (
+            application_filter_dict[mm_schemas.SchedulingKeys.APPLICATION_NAME] = (
                 application_name
             )
         return application_filter_dict
@@ -603,7 +599,7 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         endpoints = self.list_model_endpoints()
 
         for endpoint_dict in endpoints:
-            endpoint_id = endpoint_dict[mm_constants.EventFieldType.UID]
+            endpoint_id = endpoint_dict[mm_schemas.EventFieldType.UID]
 
             # Delete last analyzed records
             self._delete_last_analyzed(endpoint_id=endpoint_id)
@@ -613,3 +609,8 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
 
             # Delete model endpoint record
             self.delete_model_endpoint(endpoint_id=endpoint_id)
+
+    def get_model_endpoint_metrics(
+        self, endpoint_id: str, type: mm_schemas.ModelEndpointMonitoringMetricType
+    ) -> list[mm_schemas.ModelEndpointMonitoringMetric]:
+        raise NotImplementedError