mlrun 1.7.0rc22__py3-none-any.whl → 1.7.0rc28__py3-none-any.whl

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

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import datetime
-import json
 import typing
 import uuid
 
@@ -21,18 +20,20 @@ import pandas as pd
 import sqlalchemy
 import sqlalchemy.exc
 import sqlalchemy.orm
+from sqlalchemy.engine import make_url
 from sqlalchemy.sql.elements import BinaryExpression
 
 import mlrun.common.model_monitoring.helpers
 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
 from mlrun.common.db.sql_session import create_session, get_engine
+from mlrun.model_monitoring.db import StoreBase
 from mlrun.utils import datetime_now, logger
 
 
-class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
+class SQLStoreBase(StoreBase):
+    type: typing.ClassVar[str] = mm_schemas.ModelEndpointTarget.SQL
     """
     Handles the DB operations when the DB target is from type SQL. For the SQL operations, we use SQLAlchemy, a Python
     SQL toolkit that handles the communication with the database.  When using SQL for storing the model monitoring
@@ -44,24 +45,27 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
     def __init__(
         self,
         project: str,
-        secret_provider: typing.Callable = None,
+        **kwargs,
     ):
         """
         Initialize SQL store target object.
 
         :param project:               The name of the project.
-        :param secret_provider:       An optional secret provider to get the connection string secret.
         """
 
         super().__init__(project=project)
 
-        self._sql_connection_string = (
-            mlrun.model_monitoring.helpers.get_connection_string(
-                secret_provider=secret_provider
+        if "store_connection_string" not in kwargs:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "connection_string is a required parameter for SQLStoreBase."
             )
-        )
 
+        self._sql_connection_string = kwargs.get("store_connection_string")
         self._engine = get_engine(dsn=self._sql_connection_string)
+        self._init_tables()
+
+    def create_tables(self):
+        self._create_tables_if_not_exist()
 
     def _init_tables(self):
         self._init_model_endpoints_table()
@@ -70,13 +74,13 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         self._init_monitoring_schedules_table()
 
     def _init_model_endpoints_table(self):
-        self.ModelEndpointsTable = (
+        self.model_endpoints_table = (
             mlrun.model_monitoring.db.stores.sqldb.models._get_model_endpoints_table(
                 connection_string=self._sql_connection_string
             )
         )
         self._tables[mm_schemas.EventFieldType.MODEL_ENDPOINTS] = (
-            self.ModelEndpointsTable
+            self.model_endpoints_table
         )
 
     def _init_application_results_table(self):
@@ -150,22 +154,17 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         :param criteria: A list of binary expressions that filter the query.
         """
         with create_session(dsn=self._sql_connection_string) as session:
-            try:
-                logger.debug(
-                    "Querying the DB",
-                    table=table.__name__,
-                    criteria=[str(criterion) for criterion in criteria],
-                )
-                # Generate the get query
-                return (
-                    session.query(table)  # pyright: ignore[reportOptionalCall]
-                    .filter(*criteria)
-                    .one_or_none()
-                )
-            except sqlalchemy.exc.ProgrammingError:
-                # Probably table doesn't exist, try to create tables
-                self._create_tables_if_not_exist()
-                return
+            logger.debug(
+                "Querying the DB",
+                table=table.__name__,
+                criteria=[str(criterion) for criterion in criteria],
+            )
+            # Generate the get query
+            return (
+                session.query(table)  # pyright: ignore[reportOptionalCall]
+                .filter(*criteria)
+                .one_or_none()
+            )
 
     def _delete(
         self,
@@ -213,14 +212,13 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
                            of the attributes dictionary should exist in the SQL table.
 
         """
-        self._init_model_endpoints_table()
 
         attributes.pop(mm_schemas.EventFieldType.ENDPOINT_ID, None)
 
         self._update(
             attributes=attributes,
-            table=self.ModelEndpointsTable,
-            criteria=[self.ModelEndpointsTable.uid == endpoint_id],
+            table=self.model_endpoints_table,
+            criteria=[self.model_endpoints_table.uid == endpoint_id],
         )
 
     def delete_model_endpoint(self, endpoint_id: str) -> None:
@@ -229,11 +227,10 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
 
         :param endpoint_id: The unique id of the model endpoint.
         """
-        self._init_model_endpoints_table()
         # Delete the model endpoint record using sqlalchemy ORM
         self._delete(
-            table=self.ModelEndpointsTable,
-            criteria=[self.ModelEndpointsTable.uid == endpoint_id],
+            table=self.model_endpoints_table,
+            criteria=[self.model_endpoints_table.uid == endpoint_id],
         )
 
     def get_model_endpoint(
@@ -249,12 +246,11 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
 
         :raise MLRunNotFoundError: If the model endpoints table was not found or the model endpoint id was not found.
         """
-        self._init_model_endpoints_table()
 
         # Get the model endpoint record
         endpoint_record = self._get(
-            table=self.ModelEndpointsTable,
-            criteria=[self.ModelEndpointsTable.uid == endpoint_id],
+            table=self.model_endpoints_table,
+            criteria=[self.model_endpoints_table.uid == endpoint_id],
         )
 
         if not endpoint_record:
@@ -286,23 +282,22 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
 
         :return: A list of model endpoint dictionaries.
         """
-        self._init_model_endpoints_table()
         # Generate an empty model endpoints that will be filled afterwards with model endpoint dictionaries
         endpoint_list = []
 
         model_endpoints_table = (
-            self.ModelEndpointsTable.__table__  # pyright: ignore[reportAttributeAccessIssue]
+            self.model_endpoints_table.__table__  # pyright: ignore[reportAttributeAccessIssue]
         )
-
         # Get the model endpoints records using sqlalchemy ORM
         with create_session(dsn=self._sql_connection_string) as session:
             # Generate the list query
-            query = session.query(self.ModelEndpointsTable).filter_by(
+            query = session.query(self.model_endpoints_table).filter_by(
                 project=self.project
             )
 
             # Apply filters
             if model:
+                model = model if ":" in model else f"{model}:latest"
                 query = self._filter_values(
                     query=query,
                     model_endpoints_table=model_endpoints_table,
@@ -310,11 +305,12 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
                     filtered_values=[model],
                 )
             if function:
+                function_uri = f"{self.project}/{function}"
                 query = self._filter_values(
                     query=query,
                     model_endpoints_table=model_endpoints_table,
-                    key_filter=mm_schemas.EventFieldType.FUNCTION,
-                    filtered_values=[function],
+                    key_filter=mm_schemas.EventFieldType.FUNCTION_URI,
+                    filtered_values=[function_uri],
                 )
             if uids:
                 query = self._filter_values(
@@ -364,11 +360,9 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         """
 
         if kind == mm_schemas.WriterEventKind.METRIC:
-            self._init_application_metrics_table()
             table = self.application_metrics_table
             table_name = mm_schemas.FileTargetKind.APP_METRICS
         elif kind == mm_schemas.WriterEventKind.RESULT:
-            self._init_application_results_table()
             table = self.application_results_table
             table_name = mm_schemas.FileTargetKind.APP_RESULTS
         else:
@@ -442,7 +436,6 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         :return: Timestamp as a Unix time.
         :raise:  MLRunNotFoundError if last analyzed value is not found.
         """
-        self._init_monitoring_schedules_table()
         monitoring_schedule_record = self._get(
             table=self.MonitoringSchedulesTable,
             criteria=self._get_filter_criteria(
@@ -469,8 +462,6 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         :param last_analyzed:    Timestamp as a Unix time that represents the last analyzed time of a certain
                                  application and model endpoint.
         """
-        self._init_monitoring_schedules_table()
-
         criteria = self._get_filter_criteria(
             table=self.MonitoringSchedulesTable,
             endpoint_id=endpoint_id,
@@ -500,7 +491,6 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
     def _delete_last_analyzed(
         self, endpoint_id: str, application_name: typing.Optional[str] = None
     ) -> None:
-        self._init_monitoring_schedules_table()
         criteria = self._get_filter_criteria(
             table=self.MonitoringSchedulesTable,
             endpoint_id=endpoint_id,
@@ -512,7 +502,6 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
     def _delete_application_result(
         self, endpoint_id: str, application_name: typing.Optional[str] = None
     ) -> None:
-        self._init_application_results_table()
         criteria = self._get_filter_criteria(
             table=self.application_results_table,
             endpoint_id=endpoint_id,
@@ -524,7 +513,6 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
     def _delete_application_metrics(
         self, endpoint_id: str, application_name: typing.Optional[str] = None
     ) -> None:
-        self._init_application_metrics_table()
         criteria = self._get_filter_criteria(
             table=self.application_metrics_table,
             endpoint_id=endpoint_id,
@@ -538,8 +526,12 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
 
         for table in self._tables:
             # Create table if not exist. The `metadata` contains the `ModelEndpointsTable`
+            db_name = make_url(self._sql_connection_string).database
             if not self._engine.has_table(table):
+                logger.info(f"Creating table {table} on {db_name} db.")
                 self._tables[table].metadata.create_all(bind=self._engine)
+            else:
+                logger.info(f"Table {table} already exists on {db_name} db.")
 
     @staticmethod
     def _filter_values(
@@ -581,41 +573,6 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         # Apply AND operator on the SQL query object with the filters tuple
         return query.filter(sqlalchemy.and_(*filter_query))
 
-    @staticmethod
-    def _validate_labels(
-        endpoint_dict: dict,
-        labels: list,
-    ) -> bool:
-        """Validate that the model endpoint dictionary has the provided labels. There are 2 possible cases:
-        1 - Labels were provided as a list of key-values pairs (e.g. ['label_1=value_1', 'label_2=value_2']): Validate
-            that each pair exist in the endpoint dictionary.
-        2 - Labels were provided as a list of key labels (e.g. ['label_1', 'label_2']): Validate that each key exist in
-            the endpoint labels dictionary.
-
-        :param endpoint_dict: Dictionary of the model endpoint records.
-        :param labels:        List of dictionary of required labels.
-
-        :return: True if the labels exist in the endpoint labels dictionary, otherwise False.
-        """
-
-        # Convert endpoint labels into dictionary
-        endpoint_labels = json.loads(
-            endpoint_dict.get(mm_schemas.EventFieldType.LABELS)
-        )
-
-        for label in labels:
-            # Case 1 - label is a key=value pair
-            if "=" in label:
-                lbl, value = list(map(lambda x: x.strip(), label.split("=")))
-                if lbl not in endpoint_labels or str(endpoint_labels[lbl]) != value:
-                    return False
-            # Case 2 - label is just a key
-            else:
-                if label not in endpoint_labels:
-                    return False
-
-        return True
-
     def delete_model_endpoints_resources(self) -> None:
         """
         Delete all the model monitoring resources of the project in the SQL tables.
@@ -650,11 +607,9 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
             type=type,
         )
         if type == mm_schemas.ModelEndpointMonitoringMetricType.METRIC:
-            self._init_application_metrics_table()
             table = self.application_metrics_table
             name_col = mm_schemas.MetricData.METRIC_NAME
         else:
-            self._init_application_results_table()
             table = self.application_results_table
             name_col = mm_schemas.ResultData.RESULT_NAME
 

mlrun/model_monitoring/db/stores/v3io_kv/kv_store.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import json
-import os
 import typing
 from dataclasses import dataclass
 from http import HTTPStatus
@@ -24,8 +23,8 @@ import v3io.dataplane.response
 
 import mlrun.common.model_monitoring.helpers
 import mlrun.common.schemas.model_monitoring as mm_schemas
-import mlrun.model_monitoring.db
 import mlrun.utils.v3io_clients
+from mlrun.model_monitoring.db import StoreBase
 from mlrun.utils import logger
 
 # Fields to encode before storing in the KV table or to decode after retrieving
@@ -89,18 +88,21 @@ _KIND_TO_SCHEMA_PARAMS: dict[mm_schemas.WriterEventKind, SchemaParams] = {
 _EXCLUDE_SCHEMA_FILTER_EXPRESSION = '__name!=".#schema"'
 
 
-class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
+class KVStoreBase(StoreBase):
+    type: typing.ClassVar[str] = "v3io-nosql"
     """
     Handles the DB operations when the DB target is from type KV. For the KV operations, we use an instance of V3IO
     client and usually the KV table can be found under v3io:///users/pipelines/project-name/model-endpoints/endpoints/.
     """
 
-    def __init__(self, project: str, access_key: typing.Optional[str] = None) -> None:
+    def __init__(
+        self,
+        project: str,
+    ) -> None:
         super().__init__(project=project)
         # Initialize a V3IO client instance
-        self.access_key = access_key or os.environ.get("V3IO_ACCESS_KEY")
         self.client = mlrun.utils.v3io_clients.get_v3io_client(
-            endpoint=mlrun.mlconf.v3io_api, access_key=self.access_key
+            endpoint=mlrun.mlconf.v3io_api,
         )
         # Get the KV table path and container
         self.path, self.container = self._get_path_and_container()
@@ -186,7 +188,6 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
             table_path=self.path,
             key=endpoint_id,
             raise_for_status=v3io.dataplane.RaiseForStatus.never,
-            access_key=self.access_key,
         )
         endpoint = endpoint.output.item
 
@@ -255,7 +256,6 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
                     self.project,
                     function,
                     model,
-                    labels,
                     top_level,
                 ),
                 raise_for_status=v3io.dataplane.RaiseForStatus.never,
@@ -268,7 +268,6 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
                 exc=mlrun.errors.err_to_str(exc),
             )
             return endpoint_list
-
         # Create a list of model endpoints unique ids
         if uids is None:
             uids = []
@@ -281,10 +280,16 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
 
         # Add each relevant model endpoint to the model endpoints list
         for endpoint_id in uids:
-            endpoint = self.get_model_endpoint(
+            endpoint_dict = self.get_model_endpoint(
                 endpoint_id=endpoint_id,
             )
-            endpoint_list.append(endpoint)
+
+            if labels and not self._validate_labels(
+                endpoint_dict=endpoint_dict, labels=labels
+            ):
+                continue
+
+            endpoint_list.append(endpoint_dict)
 
         return endpoint_list
 
@@ -499,7 +504,6 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
 
     def _get_frames_client(self):
         return mlrun.utils.v3io_clients.get_frames_client(
-            token=self.access_key,
             address=mlrun.mlconf.v3io_framesd,
             container=self.container,
         )
@@ -509,20 +513,16 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
         project: str,
         function: str = None,
         model: str = None,
-        labels: list[str] = None,
         top_level: bool = False,
     ) -> str:
         """
         Convert the provided filters into a valid filter expression. The expected filter expression includes different
         conditions, divided by ' AND '.
 
-        :param project:    The name of the project.
-        :param model:      The name of the model to filter by.
-        :param function:   The name of the function to filter by.
-        :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 top_level:  If True will return only routers and endpoint that are NOT children of any router.
+        :param project:         The name of the project.
+        :param model:           The name of the model to filter by.
+        :param function:        The name of the function to filter by.
+        :param top_level:       If True will return only routers and endpoint that are NOT children of any router.
 
         :return: A valid filter expression as a string.
 
@@ -533,25 +533,17 @@ class KVStoreBase(mlrun.model_monitoring.db.StoreBase):
             raise mlrun.errors.MLRunInvalidArgumentError("project can't be empty")
 
         # Add project filter
-        filter_expression = [f"project=='{project}'"]
+        filter_expression = [f"{mm_schemas.EventFieldType.PROJECT}=='{project}'"]
 
         # Add function and model filters
         if function:
-            filter_expression.append(f"function=='{function}'")
+            function_uri = f"{project}/{function}" if function else None
+            filter_expression.append(
+                f"{mm_schemas.EventFieldType.FUNCTION_URI}=='{function_uri}'"
+            )
         if model:
-            filter_expression.append(f"model=='{model}'")
-
-        # Add labels filters
-        if labels:
-            for label in labels:
-                if not label.startswith("_"):
-                    label = f"_{label}"
-
-                if "=" in label:
-                    lbl, value = list(map(lambda x: x.strip(), label.split("=")))
-                    filter_expression.append(f"{lbl}=='{value}'")
-                else:
-                    filter_expression.append(f"exists({label})")
+            model = model if ":" in model else f"{model}:latest"
+            filter_expression.append(f"{mm_schemas.EventFieldType.MODEL}=='{model}'")
 
         # Apply top_level filter (remove endpoints that considered a child of a router)
         if top_level:

mlrun/model_monitoring/db/tsdb/__init__.py

@@ -57,41 +57,41 @@ class ObjectTSDBFactory(enum.Enum):
         :param value: Provided enum (invalid) value.
         """
         valid_values = list(cls.__members__.keys())
-        raise mlrun.errors.MLRunInvalidArgumentError(
+        raise mlrun.errors.MLRunInvalidMMStoreType(
             f"{value} is not a valid tsdb, please choose a valid value: %{valid_values}."
         )
 
 
 def get_tsdb_connector(
     project: str,
-    tsdb_connector_type: str = "",
-    secret_provider: typing.Optional[typing.Callable] = None,
+    secret_provider: typing.Optional[typing.Callable[[str], str]] = None,
+    tsdb_connection_string: typing.Optional[str] = None,
     **kwargs,
 ) -> TSDBConnector:
     """
     Get TSDB connector object.
-    :param project: The name of the project.
-    :param tsdb_connector_type: The type of the TSDB connector. See mlrun.model_monitoring.db.tsdb.ObjectTSDBFactory
-                                for available options.
-    :param secret_provider: An optional secret provider to get the connection string secret.
+    :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.
 
     :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.
     """
 
-    tsdb_connection_string = mlrun.model_monitoring.helpers.get_tsdb_connection_string(
-        secret_provider=secret_provider
+    tsdb_connection_string = (
+        tsdb_connection_string
+        or mlrun.model_monitoring.helpers.get_tsdb_connection_string(
+            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
-
-    # Set the default TSDB connector type if no connection has been set
-    tsdb_connector_type = (
-        tsdb_connector_type
-        or mlrun.mlconf.model_endpoint_monitoring.tsdb_connector_type
-    )
+    elif tsdb_connection_string and tsdb_connection_string == "v3io":
+        tsdb_connector_type = mlrun.common.schemas.model_monitoring.TSDBTarget.V3IO_TSDB
+    else:
+        tsdb_connector_type = None
 
     # Get connector type value from ObjectTSDBFactory enum class
     tsdb_connector_factory = ObjectTSDBFactory(tsdb_connector_type)

mlrun/model_monitoring/db/tsdb/base.py

@@ -25,7 +25,7 @@ from mlrun.utils import logger
 
 
 class TSDBConnector(ABC):
-    type: str = ""
+    type: typing.ClassVar[str]
 
     def __init__(self, project: str):
         """

mlrun/model_monitoring/db/tsdb/v3io/v3io_connector.py

@@ -38,7 +38,10 @@ def _is_no_schema_error(exc: v3io_frames.ReadError) -> bool:
     In case of a nonexistent TSDB table - a `v3io_frames.ReadError` error is raised.
     Check if the error message contains the relevant string to verify the cause.
     """
-    return "No TSDB schema file found" in str(exc)
+    msg = str(exc)
+    # https://github.com/v3io/v3io-tsdb/blob/v0.14.1/pkg/tsdb/v3iotsdb.go#L205
+    # https://github.com/v3io/v3io-tsdb/blob/v0.14.1/pkg/partmgr/partmgr.go#L238
+    return "No TSDB schema file found" in msg or "Failed to read schema at path" in msg
 
 
 class V3IOTSDBConnector(TSDBConnector):
@@ -418,9 +421,8 @@ class V3IOTSDBConnector(TSDBConnector):
                 f"Available tables: {list(self.tables.keys())}"
             )
 
-        if agg_funcs:
-            # Frames client expects the aggregators to be a comma-separated string
-            aggregators = ",".join(agg_funcs)
+        # Frames client expects the aggregators to be a comma-separated string
+        aggregators = ",".join(agg_funcs) if agg_funcs else None
         table_path = self.tables[table]
         try:
             df = self._frames_client.read(

mlrun/model_monitoring/helpers.py

@@ -59,13 +59,17 @@ def get_stream_path(
 
     stream_uri = mlrun.get_secret_or_env(
         mlrun.common.schemas.model_monitoring.ProjectSecretKeys.STREAM_PATH
-    ) or mlrun.mlconf.get_model_monitoring_file_target_path(
-        project=project,
-        kind=mlrun.common.schemas.model_monitoring.FileTargetKind.STREAM,
-        target="online",
-        function_name=function_name,
     )
 
+    if not stream_uri or stream_uri == "v3io":
+        # TODO : remove the first part of this condition in 1.9.0
+        stream_uri = mlrun.mlconf.get_model_monitoring_file_target_path(
+            project=project,
+            kind=mlrun.common.schemas.model_monitoring.FileTargetKind.STREAM,
+            target="online",
+            function_name=function_name,
+        )
+
     if isinstance(stream_uri, list):  # ML-6043 - user side gets only the new stream uri
         stream_uri = stream_uri[1]  # get new stream path, under projects
     return mlrun.common.model_monitoring.helpers.parse_monitoring_stream_path(
@@ -97,7 +101,7 @@ def get_monitoring_parquet_path(
     return parquet_path
 
 
-def get_connection_string(secret_provider: typing.Callable = None) -> str:
+def get_connection_string(secret_provider: typing.Callable[[str], str] = None) -> str:
     """Get endpoint store connection string from the project secret. If wasn't set, take it from the system
     configurations.
 
@@ -117,7 +121,7 @@ def get_connection_string(secret_provider: typing.Callable = None) -> str:
 
 
 def get_tsdb_connection_string(
-    secret_provider: typing.Optional[typing.Callable] = None,
+    secret_provider: typing.Optional[typing.Callable[[str], str]] = None,
 ) -> str:
     """Get TSDB connection string from the project secret. If wasn't set, take it from the system
     configurations.
@@ -278,9 +282,13 @@ def calculate_inputs_statistics(
     return inputs_statistics
 
 
-def get_endpoint_record(project: str, endpoint_id: str):
+def get_endpoint_record(
+    project: str,
+    endpoint_id: str,
+    secret_provider: typing.Optional[typing.Callable[[str], str]] = None,
+) -> dict[str, typing.Any]:
     model_endpoint_store = mlrun.model_monitoring.get_store_object(
-        project=project,
+        project=project, secret_provider=secret_provider
     )
     return model_endpoint_store.get_model_endpoint(endpoint_id=endpoint_id)