mlrun 1.7.0rc21__py3-none-any.whl → 1.7.0rc23__py3-none-any.whl

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
 
@@ -499,7 +500,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,
         )

mlrun/model_monitoring/db/tsdb/__init__.py

@@ -65,7 +65,7 @@ class ObjectTSDBFactory(enum.Enum):
 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,
     **kwargs,
 ) -> TSDBConnector:
     """

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):
         """
@@ -177,19 +177,6 @@ class TSDBConnector(ABC):
         :return:                   Metric values object or no data object.
         """
 
-    @abstractmethod
-    def read_prediction_metric_for_endpoint_if_exists(
-        self, endpoint_id: str
-    ) -> typing.Optional[mm_schemas.ModelEndpointMonitoringMetric]:
-        """
-        Read the "invocations" metric for the provided model endpoint, and return the metric object
-        if it exists.
-
-        :param endpoint_id: The model endpoint identifier.
-        :return:            `None` if the invocations metric does not exist, otherwise return the
-                            corresponding metric object.
-        """
-
     @staticmethod
     def df_to_metrics_values(
         *,

mlrun/model_monitoring/db/tsdb/tdengine/tdengine_connector.py

@@ -377,21 +377,25 @@ class TDEngineConnector(TSDBConnector):
             ),  # pyright: ignore[reportArgumentType]
         )
 
-    def read_prediction_metric_for_endpoint_if_exists(
-        self, endpoint_id: str
-    ) -> typing.Optional[mm_schemas.ModelEndpointMonitoringMetric]:
-        # Read just one record, because we just want to check if there is any data for this endpoint_id
-        predictions = self.read_predictions(
-            endpoint_id=endpoint_id,
-            start=datetime.min,
-            end=mlrun.utils.now_date(),
-            limit=1,
-        )
-        if predictions:
-            return mm_schemas.ModelEndpointMonitoringMetric(
-                project=self.project,
-                app=mm_schemas.SpecialApps.MLRUN_INFRA,
-                type=mm_schemas.ModelEndpointMonitoringMetricType.METRIC,
-                name=mm_schemas.PredictionsQueryConstants.INVOCATIONS,
-                full_name=get_invocations_fqn(self.project),
-            )
+    # Note: this function serves as a reference for checking the TSDB for the existence of a metric.
+    #
+    # def read_prediction_metric_for_endpoint_if_exists(
+    #     self, endpoint_id: str
+    # ) -> typing.Optional[mm_schemas.ModelEndpointMonitoringMetric]:
+    #     """
+    #     Read the "invocations" metric for the provided model endpoint, and return the metric object
+    #     if it exists.
+    #
+    #     :param endpoint_id: The model endpoint identifier.
+    #     :return:            `None` if the invocations metric does not exist, otherwise return the
+    #                         corresponding metric object.
+    #     """
+    #     # Read just one record, because we just want to check if there is any data for this endpoint_id
+    #     predictions = self.read_predictions(
+    #         endpoint_id=endpoint_id,
+    #         start=datetime.min,
+    #         end=mlrun.utils.now_date(),
+    #         limit=1,
+    #     )
+    #     if predictions:
+    #         return get_invocations_metric(self.project)

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

@@ -12,15 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import typing
 from datetime import datetime
 from io import StringIO
 from typing import Literal, Optional, Union
 
 import pandas as pd
+import v3io_frames
 import v3io_frames.client
-import v3io_frames.errors
-from v3io_frames.frames_pb2 import IGNORE
 
 import mlrun.common.model_monitoring
 import mlrun.common.schemas.model_monitoring as mm_schemas
@@ -35,6 +33,14 @@ _TSDB_RATE = "1/s"
 _CONTAINER = "users"
 
 
+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)
+
+
 class V3IOTSDBConnector(TSDBConnector):
     """
     Handles the TSDB operations when the TSDB connector is of type V3IO. To manage these operations we use V3IO Frames
@@ -47,7 +53,7 @@ class V3IOTSDBConnector(TSDBConnector):
         self,
         project: str,
         container: str = _CONTAINER,
-        v3io_framesd: typing.Optional[str] = None,
+        v3io_framesd: Optional[str] = None,
         create_table: bool = False,
     ) -> None:
         super().__init__(project=project)
@@ -132,7 +138,7 @@ class V3IOTSDBConnector(TSDBConnector):
             self._frames_client.create(
                 backend=_TSDB_BE,
                 table=table,
-                if_exists=IGNORE,
+                if_exists=v3io_frames.IGNORE,
                 rate=_TSDB_RATE,
             )
 
@@ -162,7 +168,7 @@ class V3IOTSDBConnector(TSDBConnector):
             time_col=mm_schemas.EventFieldType.TIMESTAMP,
             container=self.container,
             v3io_frames=self.v3io_framesd,
-            columns=["latency"],
+            columns=[mm_schemas.EventFieldType.LATENCY],
             index_cols=[
                 mm_schemas.EventFieldType.ENDPOINT_ID,
             ],
@@ -280,7 +286,7 @@ class V3IOTSDBConnector(TSDBConnector):
                 index_cols=index_cols,
             )
             logger.info("Updated V3IO TSDB successfully", table=table)
-        except v3io_frames.errors.Error as err:
+        except v3io_frames.Error as err:
             logger.exception(
                 "Could not write drift measures to TSDB",
                 err=err,
@@ -291,7 +297,7 @@ class V3IOTSDBConnector(TSDBConnector):
                 f"Failed to write application result to TSDB: {err}"
             )
 
-    def delete_tsdb_resources(self, table: typing.Optional[str] = None):
+    def delete_tsdb_resources(self, table: Optional[str] = None):
         if table:
             # Delete a specific table
             tables = [table]
@@ -301,7 +307,7 @@ class V3IOTSDBConnector(TSDBConnector):
         for table_to_delete in tables:
             try:
                 self._frames_client.delete(backend=_TSDB_BE, table=table_to_delete)
-            except v3io_frames.errors.DeleteError as e:
+            except v3io_frames.DeleteError as e:
                 logger.warning(
                     f"Failed to delete TSDB table '{table}'",
                     err=mlrun.errors.err_to_str(e),
@@ -362,7 +368,7 @@ class V3IOTSDBConnector(TSDBConnector):
                 ]
                 metrics_mapping[metric] = values
 
-        except v3io_frames.errors.Error as err:
+        except v3io_frames.Error as err:
             logger.warn("Failed to read tsdb", err=err, endpoint=endpoint_id)
 
         return metrics_mapping
@@ -372,12 +378,11 @@ class V3IOTSDBConnector(TSDBConnector):
         table: str,
         start: Union[datetime, str],
         end: Union[datetime, str],
-        columns: typing.Optional[list[str]] = None,
+        columns: Optional[list[str]] = None,
         filter_query: str = "",
-        interval: typing.Optional[str] = None,
-        agg_funcs: typing.Optional[list] = None,
-        limit: typing.Optional[int] = None,
-        sliding_window_step: typing.Optional[str] = None,
+        interval: Optional[str] = None,
+        agg_funcs: Optional[list[str]] = None,
+        sliding_window_step: Optional[str] = None,
         **kwargs,
     ) -> pd.DataFrame:
         """
@@ -400,7 +405,6 @@ class V3IOTSDBConnector(TSDBConnector):
         :param agg_funcs:             The aggregation functions to apply on the columns. Note that if `agg_funcs` is
                                       provided, `interval` must bg provided as well. Provided as a list of strings in
                                       the format of ['sum', 'avg', 'count', ...].
-        :param limit:                 The maximum number of records to return.
         :param sliding_window_step:   The time step for which the time window moves forward. Note that if
                                       `sliding_window_step` is provided, interval must be provided as well. Provided
                                       as a string in the format of '1m', '1h', etc.
@@ -414,9 +418,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
-            agg_funcs = ",".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(
@@ -427,18 +430,16 @@ class V3IOTSDBConnector(TSDBConnector):
                 columns=columns,
                 filter=filter_query,
                 aggregation_window=interval,
-                aggregators=agg_funcs,
+                aggregators=aggregators,
                 step=sliding_window_step,
                 **kwargs,
             )
         except v3io_frames.ReadError as err:
-            if "No TSDB schema file found" in str(err):
+            if _is_no_schema_error(err):
                 return pd.DataFrame()
             else:
                 raise err
 
-        if limit:
-            df = df.head(limit)
         return df
 
     def _get_v3io_source_directory(self) -> str:
@@ -509,8 +510,8 @@ class V3IOTSDBConnector(TSDBConnector):
             raise ValueError(f"Invalid {type = }")
 
         query = self._get_sql_query(
-            endpoint_id,
-            [(metric.app, metric.name) for metric in metrics],
+            endpoint_id=endpoint_id,
+            metric_and_app_names=[(metric.app, metric.name) for metric in metrics],
             table_path=table_path,
             name=name,
         )
@@ -536,21 +537,28 @@ class V3IOTSDBConnector(TSDBConnector):
 
     @staticmethod
     def _get_sql_query(
+        *,
         endpoint_id: str,
-        names: list[tuple[str, str]],
         table_path: str,
         name: str = mm_schemas.ResultData.RESULT_NAME,
+        metric_and_app_names: Optional[list[tuple[str, str]]] = None,
+        columns: Optional[list[str]] = None,
     ) -> str:
         """Get the SQL query for the results/metrics table"""
+        if columns:
+            selection = ",".join(columns)
+        else:
+            selection = "*"
+
         with StringIO() as query:
             query.write(
-                f"SELECT * FROM '{table_path}' "
+                f"SELECT {selection} FROM '{table_path}' "
                 f"WHERE {mm_schemas.WriterEvent.ENDPOINT_ID}='{endpoint_id}'"
             )
-            if names:
+            if metric_and_app_names:
                 query.write(" AND (")
 
-                for i, (app_name, result_name) in enumerate(names):
+                for i, (app_name, result_name) in enumerate(metric_and_app_names):
                     sub_cond = (
                         f"({mm_schemas.WriterEvent.APPLICATION_NAME}='{app_name}' "
                         f"AND {name}='{result_name}')"
@@ -572,7 +580,6 @@ class V3IOTSDBConnector(TSDBConnector):
         end: Union[datetime, str],
         aggregation_window: Optional[str] = None,
         agg_funcs: Optional[list[str]] = None,
-        limit: Optional[int] = None,
     ) -> Union[
         mm_schemas.ModelEndpointMonitoringMetricNoData,
         mm_schemas.ModelEndpointMonitoringMetricValues,
@@ -591,7 +598,6 @@ class V3IOTSDBConnector(TSDBConnector):
             filter_query=f"endpoint_id=='{endpoint_id}'",
             interval=aggregation_window,
             agg_funcs=agg_funcs,
-            limit=limit,
             sliding_window_step=aggregation_window,
         )
 
@@ -619,18 +625,33 @@ class V3IOTSDBConnector(TSDBConnector):
             ),  # pyright: ignore[reportArgumentType]
         )
 
-    def read_prediction_metric_for_endpoint_if_exists(
-        self, endpoint_id: str
-    ) -> Optional[mm_schemas.ModelEndpointMonitoringMetric]:
-        # Read just one record, because we just want to check if there is any data for this endpoint_id
-        predictions = self.read_predictions(
-            endpoint_id=endpoint_id, start="0", end="now", limit=1
-        )
-        if predictions.data:
-            return mm_schemas.ModelEndpointMonitoringMetric(
-                project=self.project,
-                app=mm_schemas.SpecialApps.MLRUN_INFRA,
-                type=mm_schemas.ModelEndpointMonitoringMetricType.METRIC,
-                name=mm_schemas.PredictionsQueryConstants.INVOCATIONS,
-                full_name=get_invocations_fqn(self.project),
-            )
+    # Note: this function serves as a reference for checking the TSDB for the existence of a metric.
+    #
+    # def read_prediction_metric_for_endpoint_if_exists(
+    #     self, endpoint_id: str
+    # ) -> Optional[mm_schemas.ModelEndpointMonitoringMetric]:
+    #     """
+    #     Read the count of the latency column in the predictions table for the given endpoint_id.
+    #     We just want to check if there is any data for this endpoint_id.
+    #     """
+    #     query = self._get_sql_query(
+    #         endpoint_id=endpoint_id,
+    #         table_path=self.tables[mm_schemas.FileTargetKind.PREDICTIONS],
+    #         columns=[f"count({mm_schemas.EventFieldType.LATENCY})"],
+    #     )
+    #     try:
+    #         logger.debug("Checking TSDB", project=self.project, query=query)
+    #         df: pd.DataFrame = self._frames_client.read(
+    #             backend=_TSDB_BE, query=query, start="0", end="now"
+    #         )
+    #     except v3io_frames.ReadError as err:
+    #         if _is_no_schema_error(err):
+    #             logger.debug(
+    #                 "No predictions yet", project=self.project, endpoint_id=endpoint_id
+    #             )
+    #             return
+    #         else:
+    #             raise
+    #
+    #     if not df.empty:
+    #         return get_invocations_metric(self.project)

mlrun/model_monitoring/helpers.py

@@ -25,6 +25,7 @@ from mlrun.common.schemas.model_monitoring import (
     EventFieldType,
 )
 from mlrun.common.schemas.model_monitoring.model_endpoints import (
+    ModelEndpointMonitoringMetric,
     ModelEndpointMonitoringMetricType,
     _compose_full_name,
 )
@@ -96,7 +97,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.
 
@@ -116,7 +117,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.
@@ -277,9 +278,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)
 
@@ -305,3 +310,19 @@ def get_invocations_fqn(project: str) -> str:
         name=mm_constants.PredictionsQueryConstants.INVOCATIONS,
         type=ModelEndpointMonitoringMetricType.METRIC,
     )
+
+
+def get_invocations_metric(project: str) -> ModelEndpointMonitoringMetric:
+    """
+    Return the invocations metric of any model endpoint in the given project.
+
+    :param project: The project name.
+    :returns:       The model monitoring metric object.
+    """
+    return ModelEndpointMonitoringMetric(
+        project=project,
+        app=mm_constants.SpecialApps.MLRUN_INFRA,
+        type=ModelEndpointMonitoringMetricType.METRIC,
+        name=mm_constants.PredictionsQueryConstants.INVOCATIONS,
+        full_name=get_invocations_fqn(project),
+    )

mlrun/model_monitoring/stream_processing.py

@@ -66,10 +66,6 @@ class EventStreamProcessor:
         self.parquet_batching_max_events = parquet_batching_max_events
         self.parquet_batching_timeout_secs = parquet_batching_timeout_secs
 
-        self.model_endpoint_store_target = (
-            mlrun.mlconf.model_endpoint_monitoring.store_type
-        )
-
         logger.info(
             "Initializing model monitoring event stream processor",
             parquet_path=self.parquet_path,
@@ -139,7 +135,7 @@ class EventStreamProcessor:
     def apply_monitoring_serving_graph(
         self,
         fn: mlrun.runtimes.ServingRuntime,
-        tsdb_service_provider: typing.Optional[typing.Callable] = None,
+        secret_provider: typing.Optional[typing.Callable[[str], str]] = None,
     ) -> None:
         """
         Apply monitoring serving graph to a given serving function. The following serving graph includes about 4 main
@@ -167,7 +163,8 @@ class EventStreamProcessor:
            using CE, the parquet target path is based on the defined MLRun artifact path.
 
         :param fn: A serving function.
-        :param tsdb_service_provider: An optional callable function that provides the TSDB connection string.
+        :param secret_provider: An optional callable function that provides the connection string from the project
+                                secret.
         """
 
         graph = typing.cast(
@@ -293,7 +290,6 @@ class EventStreamProcessor:
                 name="UpdateEndpoint",
                 after="ProcessBeforeEndpointUpdate",
                 project=self.project,
-                model_endpoint_store_target=self.model_endpoint_store_target,
             )
 
         apply_update_endpoint()
@@ -310,7 +306,10 @@ class EventStreamProcessor:
                 table=self.kv_path,
             )
 
-        if self.model_endpoint_store_target == ModelEndpointTarget.V3IO_NOSQL:
+        store_object = mlrun.model_monitoring.get_store_object(
+            project=self.project, secret_provider=secret_provider
+        )
+        if store_object.type == ModelEndpointTarget.V3IO_NOSQL:
             apply_infer_schema()
 
         # Emits the event in window size of events based on sample_window size (10 by default)
@@ -328,7 +327,7 @@ class EventStreamProcessor:
         # TSDB branch (skip to Prometheus if in CE env)
         if not mlrun.mlconf.is_ce_mode():
             tsdb_connector = mlrun.model_monitoring.get_tsdb_connector(
-                project=self.project, secret_provider=tsdb_service_provider
+                project=self.project, secret_provider=secret_provider
             )
             tsdb_connector.apply_monitoring_stream_steps(graph=graph)
 
@@ -904,7 +903,7 @@ class MapFeatureNames(mlrun.feature_store.steps.MapClass):
 
 
 class UpdateEndpoint(mlrun.feature_store.steps.MapClass):
-    def __init__(self, project: str, model_endpoint_store_target: str, **kwargs):
+    def __init__(self, project: str, **kwargs):
         """
         Update the model endpoint record in the DB. Note that the event at this point includes metadata and stats about
         the average latency and the amount of predictions over time. This data will be used in the monitoring dashboards
@@ -914,7 +913,6 @@ class UpdateEndpoint(mlrun.feature_store.steps.MapClass):
         """
         super().__init__(**kwargs)
         self.project = project
-        self.model_endpoint_store_target = model_endpoint_store_target
 
     def do(self, event: dict):
         # Remove labels from the event

mlrun/model_monitoring/writer.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import json
-from typing import Any, NewType
+from typing import Any, Callable, NewType
 
 import mlrun.common.model_monitoring
 import mlrun.common.schemas
@@ -30,7 +30,7 @@ from mlrun.common.schemas.model_monitoring.constants import (
     WriterEventKind,
 )
 from mlrun.common.schemas.notification import NotificationKind, NotificationSeverity
-from mlrun.model_monitoring.helpers import get_endpoint_record, get_result_instance_fqn
+from mlrun.model_monitoring.helpers import get_result_instance_fqn
 from mlrun.serving.utils import StepToDict
 from mlrun.utils import logger
 from mlrun.utils.notifications.notification_pusher import CustomNotificationPusher
@@ -102,7 +102,11 @@ class ModelMonitoringWriter(StepToDict):
 
     kind = "monitoring_application_stream_pusher"
 
-    def __init__(self, project: str, tsdb_secret_provider=None) -> None:
+    def __init__(
+        self,
+        project: str,
+        secret_provider: Callable = None,
+    ) -> None:
         self.project = project
         self.name = project  # required for the deployment process
 
@@ -111,10 +115,10 @@ class ModelMonitoringWriter(StepToDict):
         )
 
         self._app_result_store = mlrun.model_monitoring.get_store_object(
-            project=self.project
+            project=self.project, secret_provider=secret_provider
         )
         self._tsdb_connector = mlrun.model_monitoring.get_tsdb_connector(
-            project=self.project, secret_provider=tsdb_secret_provider
+            project=self.project, secret_provider=secret_provider
         )
         self._endpoints_records = {}
 
@@ -223,7 +227,7 @@ class ModelMonitoringWriter(StepToDict):
             endpoint_id = event[WriterEvent.ENDPOINT_ID]
             endpoint_record = self._endpoints_records.setdefault(
                 endpoint_id,
-                get_endpoint_record(project=self.project, endpoint_id=endpoint_id),
+                self._app_result_store.get_model_endpoint(endpoint_id=endpoint_id),
             )
             event_value = {
                 "app_name": event[WriterEvent.APPLICATION_NAME],

mlrun/projects/operations.py

@@ -77,6 +77,7 @@ def run_function(
     notifications: list[mlrun.model.Notification] = None,
     returns: Optional[list[Union[str, dict[str, str]]]] = None,
     builder_env: Optional[list] = None,
+    reset_on_run: Optional[bool] = None,
 ) -> Union[mlrun.model.RunObject, PipelineNodeWrapper]:
     """Run a local or remote task as part of a local/kubeflow pipeline
 
@@ -167,6 +168,9 @@ def run_function(
                               artifact type can be given there. The artifact key must appear in the dictionary as
                               "key": "the_key".
     :param builder_env:     env vars dict for source archive config/credentials e.g. builder_env={"GIT_TOKEN": token}
+    :param reset_on_run:    When True, function python modules would reload prior to code execution.
+                            This ensures latest code changes are executed. This argument must be used in
+                            conjunction with the local=True argument.
     :return: MLRun RunObject or PipelineNodeWrapper
     """
     engine, function = _get_engine_and_function(function, project_object)
@@ -215,6 +219,7 @@ def run_function(
             schedule=schedule,
             notifications=notifications,
             builder_env=builder_env,
+            reset_on_run=reset_on_run,
         )
         if run_result:
             run_result._notified = False