mlrun 1.6.0rc20__py3-none-any.whl → 1.6.0rc22__py3-none-any.whl

mlrun/model_monitoring/controller.py

@@ -17,7 +17,7 @@ import datetime
 import json
 import os
 import re
-from typing import Any, Iterator, Optional, Tuple, Union, cast
+from typing import Any, Iterator, NamedTuple, Optional, Union, cast
 
 from v3io.dataplane.response import HttpResponseError
 
@@ -29,11 +29,21 @@ from mlrun.common.model_monitoring.helpers import FeatureStats, pad_features_his
 from mlrun.datastore import get_stream_pusher
 from mlrun.datastore.targets import ParquetTarget
 from mlrun.model_monitoring.batch import calculate_inputs_statistics
-from mlrun.model_monitoring.helpers import get_monitoring_parquet_path, get_stream_path
-from mlrun.utils import logger
+from mlrun.model_monitoring.helpers import (
+    _BatchDict,
+    batch_dict2timedelta,
+    get_monitoring_parquet_path,
+    get_stream_path,
+)
+from mlrun.utils import create_logger, datetime_now, logger
 from mlrun.utils.v3io_clients import get_v3io_client
 
 
+class _Interval(NamedTuple):
+    start: datetime.datetime
+    end: datetime.datetime
+
+
 class _BatchWindow:
     V3IO_CONTAINER_FORMAT = "users/pipelines/{project}/monitoring-schedules/functions"
 
@@ -55,11 +65,15 @@ class _BatchWindow:
         self._endpoint = endpoint
         self._application = application
         self._first_request = first_request
-        self._kv_storage = get_v3io_client(endpoint=mlrun.mlconf.v3io_api).kv
+        self._kv_storage = get_v3io_client(
+            endpoint=mlrun.mlconf.v3io_api,
+            # Avoid noisy warning logs before the KV table is created
+            logger=create_logger(name="v3io_client", level="error"),
+        ).kv
         self._v3io_container = self.V3IO_CONTAINER_FORMAT.format(project=project)
-        self._start = self._get_last_analyzed()
         self._stop = last_updated
         self._step = timedelta_seconds
+        self._start = self._get_last_analyzed()
 
     def _get_last_analyzed(self) -> Optional[int]:
         try:
@@ -69,15 +83,26 @@ class _BatchWindow:
                 key=self._application,
             )
         except HttpResponseError as err:
-            logger.warn(
-                "Failed to get the last analyzed time for this endpoint and application, "
-                "as this is probably the first time this application is running. "
-                "Using the first request time instead.",
+            logger.info(
+                "No last analyzed time was found for this endpoint and "
+                "application, as this is probably the first time this "
+                "application is running. Using the latest between first "
+                "request time or last update time minus one day instead",
                 endpoint=self._endpoint,
                 application=self._application,
                 first_request=self._first_request,
-                error=err,
+                last_updated=self._stop,
             )
+            logger.debug("Error while getting last analyzed time", err=err)
+            if self._first_request and self._stop:
+                # TODO : Change the timedelta according to the policy.
+                first_period_in_seconds = max(
+                    int(datetime.timedelta(days=1).total_seconds()), self._step
+                )  # max between one day and the base period
+                return max(
+                    self._first_request,
+                    self._stop - first_period_in_seconds,
+                )
             return self._first_request
 
         last_analyzed = data.output.item[mm_constants.SchedulingKeys.LAST_ANALYZED]
@@ -105,20 +130,29 @@ class _BatchWindow:
 
     def get_intervals(
         self,
-    ) -> Iterator[Tuple[datetime.datetime, datetime.datetime]]:
+    ) -> Iterator[_Interval]:
         """Generate the batch interval time ranges."""
         if self._start is not None and self._stop is not None:
             entered = False
-            for timestamp in range(self._start, self._stop, self._step):
+            # Iterate timestamp from start until timestamp <= stop - step
+            # so that the last interval will end at (timestamp + step) <= stop.
+            # Add 1 to stop - step to get <= and not <.
+            for timestamp in range(
+                self._start, self._stop - self._step + 1, self._step
+            ):
                 entered = True
-                start_time = datetime.datetime.utcfromtimestamp(timestamp)
-                end_time = datetime.datetime.utcfromtimestamp(timestamp + self._step)
-                yield start_time, end_time
+                start_time = datetime.datetime.fromtimestamp(
+                    timestamp, tz=datetime.timezone.utc
+                )
+                end_time = datetime.datetime.fromtimestamp(
+                    timestamp + self._step, tz=datetime.timezone.utc
+                )
+                yield _Interval(start_time, end_time)
                 self._update_last_analyzed(timestamp + self._step)
             if not entered:
                 logger.info(
                     "All the data is set, but no complete intervals were found. "
-                    "Wait for last_updated to be updated.",
+                    "Wait for last_updated to be updated",
                     endpoint=self._endpoint,
                     application=self._application,
                     start=self._start,
@@ -127,8 +161,8 @@ class _BatchWindow:
                 )
         else:
             logger.warn(
-                "The first request time is not not found for this endpoint. "
-                "No intervals will be generated.",
+                "The first request time is not found for this endpoint. "
+                "No intervals will be generated",
                 endpoint=self._endpoint,
                 application=self._application,
                 start=self._start,
@@ -165,38 +199,44 @@ class _BatchWindowGenerator:
             self._batch_dict[pair_list[0]] = float(pair_list[1])
 
     def _get_timedelta(self) -> int:
-        """Get the timedelta from a batch dictionary"""
-        self._batch_dict = cast(dict[str, int], self._batch_dict)
-        minutes, hours, days = (
-            self._batch_dict[mm_constants.EventFieldType.MINUTES],
-            self._batch_dict[mm_constants.EventFieldType.HOURS],
-            self._batch_dict[mm_constants.EventFieldType.DAYS],
-        )
+        """Get the timedelta in seconds from the batch dictionary"""
         return int(
-            datetime.timedelta(minutes=minutes, hours=hours, days=days).total_seconds()
+            batch_dict2timedelta(cast(_BatchDict, self._batch_dict)).total_seconds()
         )
 
     @classmethod
-    def _get_last_updated_time(cls, last_request: Optional[str]) -> Optional[int]:
+    def _get_last_updated_time(
+        cls, last_request: Optional[str], has_stream: bool
+    ) -> Optional[int]:
         """
         Get the last updated time of a model endpoint.
         """
         if not last_request:
             return None
-        return int(
+        last_updated = int(
             cls._date_string2timestamp(last_request)
             - cast(
                 float,
                 mlrun.mlconf.model_endpoint_monitoring.parquet_batching_timeout_secs,
             )
         )
+        if not has_stream:
+            # If the endpoint does not have a stream, `last_updated` should be
+            # the minimum between the current time and the last updated time.
+            # This compensates for the bumping mechanism - see
+            # `bump_model_endpoint_last_request`.
+            last_updated = min(int(datetime_now().timestamp()), last_updated)
+            logger.debug(
+                "The endpoint does not have a stream", last_updated=last_updated
+            )
+        return last_updated
 
     @classmethod
     def _normalize_first_request(
         cls, first_request: Optional[str], endpoint: str
     ) -> Optional[int]:
         if not first_request:
-            logger.warn(
+            logger.debug(
                 "There is no first request time for this endpoint.",
                 endpoint=endpoint,
                 first_request=first_request,
@@ -215,6 +255,7 @@ class _BatchWindowGenerator:
         application: str,
         first_request: Optional[str],
         last_request: Optional[str],
+        has_stream: bool,
     ) -> _BatchWindow:
         """
         Get the batch window for a specific endpoint and application.
@@ -226,7 +267,7 @@ class _BatchWindowGenerator:
             endpoint=endpoint,
             application=application,
             timedelta_seconds=self._timedelta,
-            last_updated=self._get_last_updated_time(last_request),
+            last_updated=self._get_last_updated_time(last_request, has_stream),
             first_request=self._normalize_first_request(first_request, endpoint),
         )
 
@@ -251,20 +292,12 @@ class MonitoringApplicationController:
         """
         self.context = context
         self.project = project
+        self.project_obj = mlrun.get_or_create_project(project)
 
-        logger.info(
-            "Initializing MonitoringApplicationController",
-            project=project,
-        )
-
-        # Get a runtime database
+        context.logger.debug(f"Initializing {self.__class__.__name__}", project=project)
 
         self.db = mlrun.model_monitoring.get_model_endpoint_store(project=project)
 
-        # If an error occurs, it will be raised using the following argument
-        self.endpoints_exceptions = {}
-
-        # The batch window
         self._batch_window_generator = _BatchWindowGenerator(
             batch_dict=context.parameters[
                 mm_constants.EventFieldType.BATCH_INTERVALS_DICT
@@ -277,7 +310,7 @@ class MonitoringApplicationController:
         )
         self.model_monitoring_access_key = self._get_model_monitoring_access_key()
         self.parquet_directory = get_monitoring_parquet_path(
-            project=project,
+            self.project_obj,
             kind=mm_constants.FileTargetKind.APPS_PARQUET,
         )
         self.storage_options = None
@@ -303,21 +336,23 @@ class MonitoringApplicationController:
 
     def run(self):
         """
-        Main method for run all the relevant monitoring application on each endpoint
+        Main method for run all the relevant monitoring applications on each endpoint
         """
         try:
             endpoints = self.db.list_model_endpoints(uids=self.model_endpoints)
-            application = mlrun.get_or_create_project(
-                self.project
-            ).list_model_monitoring_functions()
-            if application:
-                applications_names = list({app.metadata.name for app in application})
+            monitoring_functions = self.project_obj.list_model_monitoring_functions()
+            if monitoring_functions:
+                applications_names = list(
+                    {app.metadata.name for app in monitoring_functions}
+                )
             else:
-                logger.info("There are no monitoring application found in this project")
+                self.context.logger.info(
+                    "No monitoring functions found", project=self.project
+                )
                 applications_names = []
 
         except Exception as e:
-            logger.error("Failed to list endpoints", exc=e)
+            self.context.logger.error("Failed to list endpoints", exc=e)
             return
         if endpoints and applications_names:
             # Initialize a process pool that will be used to run each endpoint applications on a dedicated process
@@ -354,9 +389,7 @@ class MonitoringApplicationController:
                     futures.append(future)
 
             for future in concurrent.futures.as_completed(futures):
-                res = future.result()
-                if res:
-                    self.endpoints_exceptions[res[0]] = res[1]
+                future.result()
 
             self._delete_old_parquet(endpoints=endpoints)
 
@@ -370,7 +403,7 @@ class MonitoringApplicationController:
         parquet_directory: str,
         storage_options: dict,
         model_monitoring_access_key: str,
-    ) -> Optional[Tuple[str, Exception]]:
+    ) -> None:
         """
         Process a model endpoint and trigger the monitoring applications. This function running on different process
         for each endpoint. In addition, this function will generate a parquet file that includes the relevant data
@@ -405,6 +438,7 @@ class MonitoringApplicationController:
                     application=application,
                     first_request=endpoint[mm_constants.EventFieldType.FIRST_REQUEST],
                     last_request=endpoint[mm_constants.EventFieldType.LAST_REQUEST],
+                    has_stream=endpoint[mm_constants.EventFieldType.STREAM_PATH] != "",
                 )
 
                 for start_infer_time, end_infer_time in batch_window.get_intervals():
@@ -424,22 +458,18 @@ class MonitoringApplicationController:
                         parquet_target_path = offline_response.vector.get_target_path()
 
                         if len(df) == 0:
-                            logger.warn(
-                                "Not enough model events since the beginning of the batch interval",
-                                featureset_name=m_fs.metadata.name,
+                            logger.info(
+                                "During this time window, the endpoint has not received any data",
                                 endpoint=endpoint[mm_constants.EventFieldType.UID],
-                                min_required_events=mlrun.mlconf.model_endpoint_monitoring.parquet_batching_max_events,
                                 start_time=start_infer_time,
                                 end_time=end_infer_time,
                             )
                             continue
 
-                    # Continue if not enough events provided since the deployment of the model endpoint
                     except FileNotFoundError:
                         logger.warn(
-                            "Parquet not found, probably due to not enough model events",
+                            "No parquets were written yet",
                             endpoint=endpoint[mm_constants.EventFieldType.UID],
-                            min_required_events=mlrun.mlconf.model_endpoint_monitoring.parquet_batching_max_events,
                         )
                         continue
 
@@ -473,12 +503,11 @@ class MonitoringApplicationController:
                         model_monitoring_access_key=model_monitoring_access_key,
                         parquet_target_path=parquet_target_path,
                     )
-        except Exception as e:
-            logger.error(
+        except Exception:
+            logger.exception(
                 "Encountered an exception",
                 endpoint_id=endpoint[mm_constants.EventFieldType.UID],
             )
-            return endpoint_id, e
 
     def _delete_old_parquet(self, endpoints: list[dict[str, Any]], days: int = 1):
         """
@@ -492,12 +521,14 @@ class MonitoringApplicationController:
                 self.parquet_directory,
                 {"V3IO_ACCESS_KEY": self.model_monitoring_access_key},
             )
-            fs = store.get_filesystem()
+            fs = store.filesystem
 
             # calculate time threshold (keep only files from the last 24 hours)
-            time_to_keep = float(
-                (datetime.datetime.now() - datetime.timedelta(days=days)).strftime("%s")
-            )
+            time_to_keep = (
+                datetime.datetime.now(tz=datetime.timezone.utc)
+                - datetime.timedelta(days=days)
+            ).timestamp()
+
             for endpoint in endpoints:
                 try:
                     apps_parquet_directories = fs.listdir(
@@ -611,14 +642,13 @@ class MonitoringApplicationController:
 
         # get offline features based on application start and end time.
         # store the result parquet by partitioning by controller end processing time
-        offline_response = fstore.get_offline_features(
-            feature_vector=vector,
+        offline_response = vector.get_offline_features(
             start_time=start_infer_time,
             end_time=end_infer_time,
             timestamp_for_filtering=mm_constants.EventFieldType.TIMESTAMP,
             target=ParquetTarget(
                 path=parquet_directory
-                + f"/key={endpoint_id}/{start_infer_time.strftime('%s')}/{application_name}.parquet",
+                + f"/key={endpoint_id}/{int(start_infer_time.timestamp())}/{application_name}.parquet",
                 storage_options=storage_options,
             ),
         )

mlrun/model_monitoring/controller_handler.py

@@ -16,7 +16,7 @@ import mlrun
 from mlrun.model_monitoring.controller import MonitoringApplicationController
 
 
-def handler(context: mlrun.run.MLClientCtx):
+def handler(context: mlrun.run.MLClientCtx) -> None:
     """
     Run model monitoring application processor
 
@@ -27,5 +27,3 @@ def handler(context: mlrun.run.MLClientCtx):
         project=context.project,
     )
     monitor_app_controller.run()
-    if monitor_app_controller.endpoints_exceptions:
-        context.logger.error(monitor_app_controller.endpoints_exceptions)

mlrun/model_monitoring/helpers.py

@@ -12,20 +12,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import datetime
 import typing
 
 import mlrun
 import mlrun.common.model_monitoring.helpers
 import mlrun.common.schemas
-from mlrun.common.schemas.model_monitoring import EventFieldType
-from mlrun.errors import MLRunInvalidArgumentError
+from mlrun.common.schemas.model_monitoring import (
+    EventFieldType,
+    MonitoringFunctionNames,
+)
+from mlrun.errors import MLRunValueError
 from mlrun.model_monitoring.model_endpoint import ModelEndpoint
 from mlrun.utils import logger
 
 if typing.TYPE_CHECKING:
     from mlrun.db.base import RunDBInterface
+    from mlrun.projects import MlrunProject
+
+
+class _BatchDict(typing.TypedDict):
+    minutes: int
+    hours: int
+    days: int
+
+
+class _MLRunNoRunsFoundError(Exception):
+    pass
 
 
 def get_stream_path(project: str = None, application_name: str = None):
@@ -55,24 +68,22 @@ def get_stream_path(project: str = None, application_name: str = None):
 
 
 def get_monitoring_parquet_path(
-    project: str,
+    project: "MlrunProject",
     kind: str = mlrun.common.schemas.model_monitoring.FileTargetKind.PARQUET,
 ) -> str:
     """Get model monitoring parquet target for the current project and kind. The parquet target path is based on the
     project artifact path. If project artifact path is not defined, the parquet target path will be based on MLRun
     artifact path.
 
-    :param project:     Project name.
+    :param project:     Project object.
     :param kind:        indicate the kind of the parquet path, can be either stream_parquet or stream_controller_parquet
 
     :return:           Monitoring parquet target path.
     """
-
-    project_obj = mlrun.get_or_create_project(name=project)
-    artifact_path = project_obj.spec.artifact_path
+    artifact_path = project.spec.artifact_path
     # Generate monitoring parquet path value
     parquet_path = mlrun.mlconf.get_model_monitoring_file_target_path(
-        project=project,
+        project=project.name,
         kind=kind,
         target="offline",
         artifact_path=artifact_path,
@@ -99,12 +110,46 @@ def get_connection_string(secret_provider: typing.Callable = None) -> str:
     )
 
 
+def batch_dict2timedelta(batch_dict: _BatchDict) -> datetime.timedelta:
+    """
+    Convert a batch dictionary to timedelta.
+
+    :param batch_dict:  Batch dict.
+
+    :return:            Timedelta.
+    """
+    return datetime.timedelta(**batch_dict)
+
+
+def _get_monitoring_time_window_from_controller_run(
+    project: str, db: "RunDBInterface"
+) -> datetime.timedelta:
+    """
+    Get timedelta for the controller to run.
+
+    :param project: Project name.
+    :param db:      DB interface.
+
+    :return:    Timedelta for the controller to run.
+    """
+    run_name = MonitoringFunctionNames.APPLICATION_CONTROLLER
+    runs = db.list_runs(project=project, name=run_name, sort=True)
+    if not runs:
+        raise _MLRunNoRunsFoundError(f"No {run_name} runs were found")
+    last_run = runs[0]
+    try:
+        batch_dict = last_run["spec"]["parameters"]["batch_intervals_dict"]
+    except KeyError:
+        raise MLRunValueError(
+            f"Could not find `batch_intervals_dict` in {run_name} run"
+        )
+    return batch_dict2timedelta(batch_dict)
+
+
 def bump_model_endpoint_last_request(
     project: str,
     model_endpoint: ModelEndpoint,
     db: "RunDBInterface",
-    minutes_delta: int = 10,  # TODO: move to config - should be the same as `batch_interval`
-    seconds_delta: int = 1,
 ) -> None:
     """
     Update the last request field of the model endpoint to be after the current last request time.
@@ -112,10 +157,6 @@ def bump_model_endpoint_last_request(
     :param project:         Project name.
     :param model_endpoint:  Model endpoint object.
     :param db:              DB interface.
-    :param minutes_delta:   Minutes delta to add to the last request time.
-    :param seconds_delta:   Seconds delta to add to the last request time. This is mainly to ensure that the last
-                            request time is strongly greater than the previous one (with respect to the window time)
-                            after adding the minutes delta.
     """
     if not model_endpoint.status.last_request:
         logger.error(
@@ -123,14 +164,18 @@ def bump_model_endpoint_last_request(
             project=project,
             endpoint_id=model_endpoint.metadata.uid,
         )
-        raise MLRunInvalidArgumentError("Model endpoint last request time is empty")
+        raise MLRunValueError("Model endpoint last request time is empty")
+    try:
+        time_window = _get_monitoring_time_window_from_controller_run(project, db)
+    except _MLRunNoRunsFoundError:
+        logger.debug(
+            "Not bumping model endpoint last request time - no controller runs were found"
+        )
+        return
 
     bumped_last_request = (
         datetime.datetime.fromisoformat(model_endpoint.status.last_request)
-        + datetime.timedelta(
-            minutes=minutes_delta,
-            seconds=seconds_delta,
-        )
+        + time_window
         + datetime.timedelta(
             seconds=mlrun.mlconf.model_endpoint_monitoring.parquet_batching_timeout_secs
         )

mlrun/model_monitoring/stream_processing.py

@@ -528,9 +528,6 @@ class ProcessBeforeTSDB(mlrun.feature_store.steps.MapClass):
 
         # Getting event timestamp and endpoint_id
         base_event = {k: event[k] for k in base_fields}
-        base_event[EventFieldType.TIMESTAMP] = datetime.datetime.fromisoformat(
-            base_event[EventFieldType.TIMESTAMP]
-        )
 
         # base_metrics includes the stats about the average latency and the amount of predictions over time
         base_metrics = {

mlrun/projects/operations.py

@@ -274,7 +274,7 @@ def build_function(
     if not overwrite_build_params:
         # TODO: change overwrite_build_params default to True in 1.8.0
         warnings.warn(
-            "The `overwrite_build_params` parameter default will change from 'False' to 'True in 1.8.0.",
+            "The `overwrite_build_params` parameter default will change from 'False' to 'True' in 1.8.0.",
             mlrun.utils.OverwriteBuildParamsWarning,
         )
 

mlrun/projects/project.py

@@ -2167,7 +2167,7 @@ class MlrunProject(ModelObj):
         self.spec.remove_function(name)
 
     def remove_model_monitoring_function(self, name):
-        """remove the specified model-monitoring-app function from the project
+        """remove the specified model-monitoring-app function from the project and from the db
 
         :param name: name of the model-monitoring-app function (under the project)
         """
@@ -2177,6 +2177,7 @@ class MlrunProject(ModelObj):
             == mm_constants.ModelMonitoringAppLabel.VAL
         ):
             self.remove_function(name=name)
+            mlrun.db.get_run_db().delete_function(name=name.lower())
             logger.info(f"{name} function has been removed from {self.name} project")
         else:
             raise logger.error(
@@ -2753,6 +2754,11 @@ class MlrunProject(ModelObj):
             project_file_path = path.join(
                 self.spec.context, self.spec.subpath or "", "project.yaml"
             )
+        if filepath and "://" in str(filepath) and not archive_code:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "URLs are only applicable to archives"
+            )
+
         project_dir = pathlib.Path(project_file_path).parent
         project_dir.mkdir(parents=True, exist_ok=True)
         with open(project_file_path, "w") as fp:
@@ -3011,7 +3017,7 @@ class MlrunProject(ModelObj):
         if not overwrite_build_params:
             # TODO: change overwrite_build_params default to True in 1.8.0
             warnings.warn(
-                "The `overwrite_build_params` parameter default will change from 'False' to 'True in 1.8.0.",
+                "The `overwrite_build_params` parameter default will change from 'False' to 'True' in 1.8.0.",
                 mlrun.utils.OverwriteBuildParamsWarning,
             )
         default_image_name = mlrun.mlconf.default_project_image_name.format(
@@ -3097,7 +3103,7 @@ class MlrunProject(ModelObj):
         if not overwrite_build_params:
             # TODO: change overwrite_build_params default to True in 1.8.0
             warnings.warn(
-                "The `overwrite_build_params` parameter default will change from 'False' to 'True in 1.8.0.",
+                "The `overwrite_build_params` parameter default will change from 'False' to 'True' in 1.8.0.",
                 mlrun.utils.OverwriteBuildParamsWarning,
             )
 
@@ -3402,7 +3408,7 @@ class MlrunProject(ModelObj):
         :param state: List only runs whose state is specified.
         :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.
+        :param last: Deprecated - currently not used (will be removed in 1.8.0).
         :param iter: If ``True`` return runs from all iterations. Otherwise, return only runs whose ``iter`` is 0.
         :param start_time_from: Filter by run start time in ``[start_time_from, start_time_to]``.
         :param start_time_to: Filter by run start time in ``[start_time_from, start_time_to]``.

mlrun/runtimes/base.py

@@ -550,7 +550,12 @@ class BaseRuntime(ModelObj):
             if err:
                 updates["status.error"] = err_to_str(err)
 
-        elif not was_none and last_state != "completed":
+        elif (
+            not was_none
+            and last_state != mlrun.runtimes.constants.RunStates.completed
+            and last_state
+            not in mlrun.runtimes.constants.RunStates.error_and_abortion_states()
+        ):
             try:
                 runtime_cls = mlrun.runtimes.get_runtime_class(kind)
                 updates = runtime_cls._get_run_completion_updates(resp)

mlrun/runtimes/constants.py

@@ -165,6 +165,17 @@ class RunStates(object):
             RunStates.aborted,
         ]
 
+    @staticmethod
+    def abortion_states():
+        return [
+            RunStates.aborted,
+            RunStates.aborting,
+        ]
+
+    @staticmethod
+    def error_and_abortion_states():
+        return list(set(RunStates.error_states()) | set(RunStates.abortion_states()))
+
     @staticmethod
     def non_terminal_states():
         return list(set(RunStates.all()) - set(RunStates.terminal_states()))

mlrun/runtimes/databricks_job/databricks_runtime.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 from ast import FunctionDef, parse, unparse
 from base64 import b64decode, b64encode
 from typing import Callable, Dict, List, Optional, Union
@@ -197,14 +196,13 @@ class DatabricksRuntime(kubejob.KubejobRuntime):
             if value:
                 task_parameters[key] = value  # in order to handle reruns.
         runspec.spec.parameters["task_parameters"] = task_parameters
-        current_file = os.path.abspath(__file__)
-        current_dir = os.path.dirname(current_file)
-        databricks_runtime_wrap_path = os.path.join(
-            current_dir, "databricks_wrapper.py"
-        )
-        with open(databricks_runtime_wrap_path, "r") as databricks_runtime_wrap_file:
-            wrap_code = databricks_runtime_wrap_file.read()
-            wrap_code = b64encode(wrap_code.encode("utf-8")).decode("utf-8")
+        wrap_code = b"""
+from mlrun.runtimes.databricks_job import databricks_wrapper
+
+def run_mlrun_databricks_job(context,task_parameters: dict, **kwargs):
+        databricks_wrapper.run_mlrun_databricks_job(context, task_parameters, **kwargs)
+"""
+        wrap_code = b64encode(wrap_code).decode("utf-8")
         self.spec.build.functionSourceCode = wrap_code
         runspec.spec.handler = "run_mlrun_databricks_job"