mlrun 1.10.0rc7__py3-none-any.whl → 1.10.0rc9__py3-none-any.whl

mlrun/db/sql_types.py

@@ -0,0 +1,160 @@
+# Copyright 2025 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This module provides SQLAlchemy TypeDecorator subclasses that are aware of
+database dialects (MySQL, PostgreSQL, SQLite) and automatically select
+appropriate native types (e.g., UUID, BLOB, TIMESTAMP with precision) or
+fallbacks (e.g., hex-string storage) to ensure consistent behavior across
+different database backends.
+"""
+
+import uuid
+from typing import Any, Optional, Union
+
+import sqlalchemy.types
+from sqlalchemy import CHAR, Text
+from sqlalchemy.dialects.mysql import DATETIME as MYSQL_DATETIME
+from sqlalchemy.dialects.mysql import MEDIUMBLOB
+from sqlalchemy.dialects.postgresql import BYTEA
+from sqlalchemy.dialects.postgresql import TIMESTAMP as PG_TIMESTAMP
+from sqlalchemy.dialects.postgresql import UUID as PG_UUID
+from sqlalchemy.engine.interfaces import Dialect
+from sqlalchemy.types import TypeDecorator
+
+import mlrun.common.db.dialects
+
+
+class DateTime(TypeDecorator):
+    impl = sqlalchemy.types.DateTime
+    cache_ok = True
+    precision: int = 3
+
+    def load_dialect_impl(
+        self,
+        dialect: Dialect,
+    ) -> sqlalchemy.types.TypeEngine:
+        if dialect.name == mlrun.common.db.dialects.Dialects.MYSQL:
+            return dialect.type_descriptor(
+                MYSQL_DATETIME(
+                    fsp=self.precision,
+                    timezone=True,
+                )
+            )
+        if dialect.name == mlrun.common.db.dialects.Dialects.POSTGRESQL:
+            return dialect.type_descriptor(
+                PG_TIMESTAMP(
+                    precision=self.precision,
+                    timezone=True,
+                )
+            )
+        return dialect.type_descriptor(sqlalchemy.types.DateTime)
+
+
+class MicroSecondDateTime(DateTime):
+    cache_ok = True
+    precision: int = 6
+
+
+class Blob(TypeDecorator):
+    impl = sqlalchemy.types.LargeBinary
+    cache_ok = True
+
+    def load_dialect_impl(
+        self,
+        dialect: Dialect,
+    ) -> sqlalchemy.types.TypeEngine:
+        if dialect.name == mlrun.common.db.dialects.Dialects.MYSQL:
+            return dialect.type_descriptor(MEDIUMBLOB)
+        if dialect.name == mlrun.common.db.dialects.Dialects.POSTGRESQL:
+            return dialect.type_descriptor(BYTEA)
+        return dialect.type_descriptor(self.impl)
+
+
+class Utf8BinText(TypeDecorator):
+    impl = Text
+    cache_ok = True
+
+    def load_dialect_impl(
+        self,
+        dialect: Dialect,
+    ) -> sqlalchemy.types.TypeEngine:
+        if dialect.name == mlrun.common.db.dialects.Dialects.MYSQL:
+            return dialect.type_descriptor(
+                sqlalchemy.dialects.mysql.VARCHAR(
+                    collation="utf8_bin",
+                    length=255,
+                )
+            )
+        if dialect.name == mlrun.common.db.dialects.Dialects.POSTGRESQL:
+            # This collation is created as part of the database creation
+            return dialect.type_descriptor(
+                Text(
+                    collation="utf8_bin",
+                )
+            )
+        if dialect.name == mlrun.common.db.dialects.Dialects.SQLITE:
+            return dialect.type_descriptor(
+                Text(
+                    collation="BINARY",
+                )
+            )
+        return dialect.type_descriptor(self.impl)
+
+
+class UuidType(TypeDecorator):
+    """
+    A UUID type which stores as native UUID on Postgres (as_uuid=True)
+    and as 32-char hex strings on other dialects.
+    """
+
+    impl = CHAR(32)
+    cache_ok = True
+
+    def load_dialect_impl(self, dialect: Dialect) -> sqlalchemy.types.TypeEngine:
+        if dialect.name == mlrun.common.db.dialects.Dialects.POSTGRESQL:
+            return dialect.type_descriptor(PG_UUID(as_uuid=True))
+        return dialect.type_descriptor(CHAR(32))
+
+    def process_bind_param(
+        self,
+        value: Optional[Union[uuid.UUID, str]],
+        dialect: Dialect,
+    ) -> Optional[Union[uuid.UUID, str]]:
+        if value is None:
+            return None
+        if isinstance(value, uuid.UUID):
+            return (
+                value
+                if dialect.name == mlrun.common.db.dialects.Dialects.POSTGRESQL
+                else value.hex
+            )
+        if isinstance(value, str):
+            u = uuid.UUID(value)
+            return (
+                u
+                if dialect.name == mlrun.common.db.dialects.Dialects.POSTGRESQL
+                else u.hex
+            )
+        raise ValueError(f"Cannot bind UUID value {value!r}")
+
+    def process_result_value(
+        self, value: Optional[Union[uuid.UUID, bytes, str]], dialect: Dialect
+    ) -> Optional[uuid.UUID]:
+        if value is None:
+            return None
+        return value if isinstance(value, uuid.UUID) else uuid.UUID(value)
+
+    def coerce_compared_value(self, op: Any, value: Any) -> TypeDecorator:
+        # ensure STR comparisons are coerced through this type
+        return self

mlrun/frameworks/tf_keras/__init__.py

@@ -14,7 +14,7 @@
 
 from typing import Any, Optional, Union
 
-from tensorflow import keras
+import tensorflow as tf
 
 import mlrun
 import mlrun.common.constants as mlrun_constants
@@ -27,11 +27,11 @@ from .utils import TFKerasTypes, TFKerasUtils
 
 
 def apply_mlrun(
-    model: keras.Model = None,
+    model: tf.keras.Model = None,
     model_name: Optional[str] = None,
     tag: str = "",
     model_path: Optional[str] = None,
-    model_format: str = TFKerasModelHandler.ModelFormats.SAVED_MODEL,
+    model_format: Optional[str] = None,
     save_traces: bool = False,
     modules_map: Optional[Union[dict[str, Union[None, str, list[str]]], str]] = None,
     custom_objects_map: Optional[Union[dict[str, Union[str, list[str]]], str]] = None,
@@ -54,7 +54,7 @@ def apply_mlrun(
     :param model_path:                  The model's store object path. Mandatory for evaluation (to know which model to
                                         update). If model is not provided, it will be loaded from this path.
     :param model_format:                The format to use for saving and loading the model. Should be passed as a
-                                        member of the class 'ModelFormats'. Default: 'ModelFormats.SAVED_MODEL'.
+                                        member of the class 'ModelFormats'.
     :param save_traces:                 Whether or not to use functions saving (only available for the 'SavedModel'
                                         format) for loading the model later without the custom objects dictionary. Only
                                         from tensorflow version >= 2.4.0. Using this setting will increase the model

mlrun/frameworks/tf_keras/callbacks/logging_callback.py

@@ -16,14 +16,14 @@ from typing import Callable, Optional, Union
 
 import numpy as np
 import tensorflow as tf
-from tensorflow import Tensor, Variable
+from tensorflow import keras
 from tensorflow.python.keras.callbacks import Callback
 
 import mlrun
 
 from ..._common import LoggingMode
 from ..._dl_common.loggers import Logger
-from ..utils import TFKerasTypes
+from ..utils import TFKerasTypes, is_keras_3
 
 
 class LoggingCallback(Callback):
@@ -70,7 +70,7 @@ class LoggingCallback(Callback):
                                         {
                                             "epochs": 7
                                         }
-        :param auto_log:                Whether or not to enable auto logging, trying to track common static and dynamic
+        :param auto_log:                Whether to enable auto logging, trying to track common static and dynamic
                                         hyperparameters.
         """
         super().__init__()
@@ -385,18 +385,24 @@ class LoggingCallback(Callback):
             self._logger.log_context_parameters()
 
         # Add learning rate:
-        learning_rate_key = "lr"
-        learning_rate_key_chain = ["optimizer", "lr"]
-        if learning_rate_key not in self._dynamic_hyperparameters_keys and hasattr(
-            self.model, "optimizer"
-        ):
-            try:
-                self._get_hyperparameter(key_chain=learning_rate_key_chain)
+        learning_rate_keys = [
+            "learning_rate",
+            "lr",
+        ]  # "lr" is for backward compatibility in older keras versions.
+        if all(
+            learning_rate_key not in self._dynamic_hyperparameters_keys
+            for learning_rate_key in learning_rate_keys
+        ) and hasattr(self.model, "optimizer"):
+            for learning_rate_key in learning_rate_keys:
+                learning_rate_key_chain = ["optimizer", learning_rate_key]
+                try:
+                    self._get_hyperparameter(key_chain=learning_rate_key_chain)
+                except (KeyError, IndexError, AttributeError, ValueError):
+                    continue
                 self._dynamic_hyperparameters_keys[learning_rate_key] = (
                     learning_rate_key_chain
                 )
-            except (KeyError, IndexError, ValueError):
-                pass
+                break
 
     def _get_hyperparameter(
         self,
@@ -427,7 +433,7 @@ class LoggingCallback(Callback):
                         value = value[key]
                     else:
                         value = getattr(value, key)
-                except KeyError or IndexError as KeyChainError:
+                except KeyError or IndexError or AttributeError as KeyChainError:
                     raise KeyChainError(
                         f"Error during getting a hyperparameter value with the key chain {key_chain}. "
                         f"The {value.__class__} in it does not have the following key/index from the key provided: "
@@ -435,7 +441,9 @@ class LoggingCallback(Callback):
                     )
 
         # Parse the value:
-        if isinstance(value, Tensor) or isinstance(value, Variable):
+        if isinstance(value, (tf.Tensor, tf.Variable)) or (
+            is_keras_3() and isinstance(value, (keras.KerasTensor, keras.Variable))
+        ):
             if int(tf.size(value)) == 1:
                 value = float(value)
             else:
@@ -451,12 +459,7 @@ class LoggingCallback(Callback):
                     f"The parameter with the following key chain: {key_chain} is a numpy.ndarray with {value.size} "
                     f"elements. numpy arrays are trackable only if they have 1 element."
                 )
-        elif not (
-            isinstance(value, float)
-            or isinstance(value, int)
-            or isinstance(value, str)
-            or isinstance(value, bool)
-        ):
+        elif not (isinstance(value, (float, int, str, bool))):
             raise mlrun.errors.MLRunInvalidArgumentError(
                 f"The parameter with the following key chain: {key_chain} is of type '{type(value)}'. The only "
                 f"trackable types are: float, int, str and bool."

mlrun/frameworks/tf_keras/mlrun_interface.py

@@ -280,7 +280,10 @@ class TFKerasMLRunInterface(MLRunInterface, ABC):
             print(f"Horovod worker #{self._hvd.rank()} is using CPU")
 
         # Adjust learning rate based on the number of GPUs:
-        optimizer.lr = optimizer.lr * self._hvd.size()
+        if hasattr(self.optimizer, "lr"):
+            optimizer.lr *= self._hvd.size()
+        else:
+            optimizer.learning_rate *= self._hvd.size()
 
         # Wrap the optimizer in horovod's distributed optimizer: 'hvd.DistributedOptimizer'.
         optimizer = self._hvd.DistributedOptimizer(optimizer)

mlrun/frameworks/tf_keras/model_handler.py

@@ -29,7 +29,7 @@ from mlrun.features import Feature
 from .._common import without_mlrun_interface
 from .._dl_common import DLModelHandler
 from .mlrun_interface import TFKerasMLRunInterface
-from .utils import TFKerasUtils
+from .utils import TFKerasUtils, is_keras_3
 
 
 class TFKerasModelHandler(DLModelHandler):
@@ -40,8 +40,8 @@ class TFKerasModelHandler(DLModelHandler):
     # Framework name:
     FRAMEWORK_NAME = "tensorflow.keras"
 
-    # Declare a type of an input sample:
-    IOSample = Union[tf.Tensor, tf.TensorSpec, np.ndarray]
+    # Declare a type of input sample (only from keras v3 there is a KerasTensor type):
+    IOSample = Union[tf.Tensor, tf.TensorSpec, "keras.KerasTensor", np.ndarray]
 
     class ModelFormats:
         """
@@ -49,9 +49,19 @@ class TFKerasModelHandler(DLModelHandler):
         """
 
         SAVED_MODEL = "SavedModel"
+        KERAS = "keras"
         H5 = "h5"
         JSON_ARCHITECTURE_H5_WEIGHTS = "json_h5"
 
+        @classmethod
+        def default(cls) -> str:
+            """
+            Get the default model format to use for saving and loading the model based on the keras version.
+
+            :return: The default model format to use.
+            """
+            return cls.KERAS if is_keras_3() else cls.SAVED_MODEL
+
     class _LabelKeys:
         """
         Required labels keys to log with the model.
@@ -65,7 +75,7 @@ class TFKerasModelHandler(DLModelHandler):
         model: keras.Model = None,
         model_path: Optional[str] = None,
         model_name: Optional[str] = None,
-        model_format: str = ModelFormats.SAVED_MODEL,
+        model_format: Optional[str] = None,
         context: mlrun.MLClientCtx = None,
         modules_map: Optional[
             Union[dict[str, Union[None, str, list[str]]], str]
@@ -98,7 +108,7 @@ class TFKerasModelHandler(DLModelHandler):
                                          * If given a loaded model object and the model name is None, the name will be
                                            set to the model's object name / class.
         :param model_format:             The format to use for saving and loading the model. Should be passed as a
-                                         member of the class 'ModelFormats'. Default: 'ModelFormats.SAVED_MODEL'.
+                                         member of the class 'ModelFormats'.
         :param context:                  MLRun context to work with for logging the model.
         :param modules_map:              A dictionary of all the modules required for loading the model. Each key
                                          is a path to a module and its value is the object name to import from it. All
@@ -144,8 +154,11 @@ class TFKerasModelHandler(DLModelHandler):
                                           * 'save_traces' parameter was miss-used.
         """
         # Validate given format:
+        if not model_format:
+            model_format = TFKerasModelHandler.ModelFormats.default()
         if model_format not in [
             TFKerasModelHandler.ModelFormats.SAVED_MODEL,
+            TFKerasModelHandler.ModelFormats.KERAS,
             TFKerasModelHandler.ModelFormats.H5,
             TFKerasModelHandler.ModelFormats.JSON_ARCHITECTURE_H5_WEIGHTS,
         ]:
@@ -153,6 +166,22 @@ class TFKerasModelHandler(DLModelHandler):
                 f"Unrecognized model format: '{model_format}'. Please use one of the class members of "
                 "'TFKerasModelHandler.ModelFormats'"
             )
+        if not is_keras_3():
+            if model_format == TFKerasModelHandler.ModelFormats.KERAS:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "The 'keras' model format is only supported in Keras 3.0.0 and above. "
+                    f"Current version is {keras.__version__}."
+                )
+        else:
+            if (
+                model_format == TFKerasModelHandler.ModelFormats.SAVED_MODEL
+                or model_format
+                == TFKerasModelHandler.ModelFormats.JSON_ARCHITECTURE_H5_WEIGHTS
+            ):
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    f"The '{model_format}' model format is not supported in Keras 3.0.0 and above. "
+                    f"Current version is {keras.__version__}."
+                )
 
         # Validate 'save_traces':
         if save_traces:
@@ -239,11 +268,19 @@ class TFKerasModelHandler(DLModelHandler):
             self._model_file = f"{self._model_name}.h5"
             self._model.save(self._model_file)
 
+        # ModelFormats.keras - Save as a keras file:
+        elif self._model_format == self.ModelFormats.KERAS:
+            self._model_file = f"{self._model_name}.keras"
+            self._model.save(self._model_file)
+
         # ModelFormats.SAVED_MODEL - Save as a SavedModel directory and zip its file:
         elif self._model_format == TFKerasModelHandler.ModelFormats.SAVED_MODEL:
             # Save it in a SavedModel format directory:
+            # Note: Using keras>=3.0.0 can save in this format via `model.export` but then it won't be able to load it
+            # back, only for inference. So, we use the `save` method instead for keras 2 and validate the user won't use
+            # keras 3 and this model format.
             if self._save_traces is True:
-                # Save traces can only be used in versions >= 2.4, so only if its true we use it in the call:
+                # Save traces can only be used in versions >= 2.4, so only if it's true, we use it in the call:
                 self._model.save(self._model_name, save_traces=self._save_traces)
             else:
                 self._model.save(self._model_name)
@@ -303,6 +340,12 @@ class TFKerasModelHandler(DLModelHandler):
                 self._model_file, custom_objects=self._custom_objects
             )
 
+        # ModelFormats.KERAS - Load from a keras file:
+        elif self._model_format == TFKerasModelHandler.ModelFormats.KERAS:
+            self._model = keras.models.load_model(
+                self._model_file, custom_objects=self._custom_objects
+            )
+
         # ModelFormats.SAVED_MODEL - Load from a SavedModel directory:
         elif self._model_format == TFKerasModelHandler.ModelFormats.SAVED_MODEL:
             self._model = keras.models.load_model(
@@ -434,7 +477,10 @@ class TFKerasModelHandler(DLModelHandler):
             )
 
         # Read the inputs:
-        input_signature = [input_layer.type_spec for input_layer in self._model.inputs]
+        input_signature = [
+            getattr(input_layer, "type_spec", input_layer)
+            for input_layer in self._model.inputs
+        ]
 
         # Set the inputs:
         self.set_inputs(from_sample=input_signature)
@@ -453,7 +499,8 @@ class TFKerasModelHandler(DLModelHandler):
 
         # Read the outputs:
         output_signature = [
-            output_layer.type_spec for output_layer in self._model.outputs
+            getattr(output_layer, "type_spec", output_layer)
+            for output_layer in self._model.outputs
         ]
 
         # Set the outputs:
@@ -480,11 +527,22 @@ class TFKerasModelHandler(DLModelHandler):
             self._model_file = os.path.join(
                 os.path.dirname(self._model_file), self._model_name
             )
+        elif self._model_format == TFKerasModelHandler.ModelFormats.KERAS:
+            # When keras tried to load it, it validates the suffix. The `artifacts.model.get_model` function is
+            # downloading the keras file to a temp file with a `pkl` suffix, so it needs to be replaced:
+            self._model_file = self._model_file.rsplit(".pkl", 1)[0] + ".keras"
+        elif self._model_format == TFKerasModelHandler.ModelFormats.H5:
+            # When keras tried to load it, it validates the suffix. The `artifacts.model.get_model` function is
+            # downloading the keras file to a temp file with a `pkl` suffix, so it needs to be replaced:
+            self._model_file = self._model_file.rsplit(".pkl", 1)[0] + ".h5"
         # # ModelFormats.JSON_ARCHITECTURE_H5_WEIGHTS - Get the weights file:
         elif (
             self._model_format
             == TFKerasModelHandler.ModelFormats.JSON_ARCHITECTURE_H5_WEIGHTS
         ):
+            # When keras tried to load it, it validates the suffix. The `artifacts.model.get_model` function is
+            # downloading the keras file to a temp file with a `pkl` suffix, so it needs to be replaced:
+            self._model_file = self._model_file.rsplit(".pkl", 1)[0] + ".json"
             # Get the weights file:
             self._weights_file = self._extra_data[
                 self._get_weights_file_artifact_name()
@@ -509,6 +567,17 @@ class TFKerasModelHandler(DLModelHandler):
                     f"'{self._model_path}'"
                 )
 
+        # ModelFormats.KERAS - Get the keras model file:
+        elif self._model_format == TFKerasModelHandler.ModelFormats.KERAS:
+            self._model_file = os.path.join(
+                self._model_path, f"{self._model_name}.keras"
+            )
+            if not os.path.exists(self._model_file):
+                raise mlrun.errors.MLRunNotFoundError(
+                    f"The model file '{self._model_name}.keras' was not found within the given 'model_path': "
+                    f"'{self._model_path}'"
+                )
+
         # ModelFormats.SAVED_MODEL - Get the zip file and extract it, or simply locate the directory:
         elif self._model_format == TFKerasModelHandler.ModelFormats.SAVED_MODEL:
             self._model_file = os.path.join(self._model_path, f"{self._model_name}.zip")
@@ -559,7 +628,9 @@ class TFKerasModelHandler(DLModelHandler):
         # Supported types:
         if isinstance(sample, np.ndarray):
             return super()._read_sample(sample=sample)
-        elif isinstance(sample, tf.TensorSpec):
+        elif isinstance(sample, tf.TensorSpec) or (
+            is_keras_3() and isinstance(sample, keras.KerasTensor)
+        ):
             return Feature(
                 name=sample.name,
                 value_type=TFKerasUtils.convert_tf_dtype_to_value_type(

mlrun/frameworks/tf_keras/utils.py

@@ -11,8 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 import tensorflow as tf
+from packaging import version
 from tensorflow import keras
 
 import mlrun
@@ -117,3 +117,14 @@ class TFKerasUtils(DLUtils):
         raise mlrun.errors.MLRunInvalidArgumentError(
             f"MLRun value type is not supporting the given tensorflow data type: '{tf_dtype}'."
         )
+
+
+def is_keras_3() -> bool:
+    """
+    Check if the current Keras version is 3.x.
+
+    :return: True if Keras version is 3.x, False otherwise.
+    """
+    return hasattr(keras, "__version__") and version.parse(
+        keras.__version__
+    ) >= version.parse("3.0.0")

mlrun/launcher/base.py

@@ -82,7 +82,6 @@ class BaseLauncher(abc.ABC):
         runtime: "mlrun.runtimes.base.BaseRuntime",
         project_name: Optional[str] = "",
         full: bool = True,
-        client_version: str = "",
     ):
         pass
 
@@ -148,6 +147,12 @@ class BaseLauncher(abc.ABC):
         self._validate_run_params(run.spec.parameters)
         self._validate_output_path(runtime, run)
 
+        for image in [
+            runtime.spec.image,
+            getattr(runtime.spec.build, "base_image", None),
+        ]:
+            mlrun.utils.helpers.warn_on_deprecated_image(image)
+
     @staticmethod
     def _validate_output_path(
         runtime: "mlrun.runtimes.BaseRuntime",

mlrun/launcher/client.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import abc
-import warnings
 from typing import Optional
 
 import IPython.display
@@ -37,7 +36,6 @@ class ClientBaseLauncher(launcher.BaseLauncher, abc.ABC):
         runtime: "mlrun.runtimes.base.BaseRuntime",
         project_name: Optional[str] = "",
         full: bool = True,
-        client_version: str = "",
     ):
         runtime.try_auto_mount_based_on_config()
         runtime._fill_credentials()
@@ -63,26 +61,7 @@ class ClientBaseLauncher(launcher.BaseLauncher, abc.ABC):
         ):
             image = mlrun.mlconf.function_defaults.image_by_kind.to_dict()[runtime.kind]
 
-        # Warn if user explicitly set the deprecated mlrun/ml-base image
-        if image and "mlrun/ml-base" in image:
-            client_version = mlrun.utils.version.Version().get()["version"]
-            auto_replaced = mlrun.utils.validate_component_version_compatibility(
-                "mlrun-client", "1.10.0", mlrun_client_version=client_version
-            )
-            message = (
-                "'mlrun/ml-base' image is deprecated in 1.10.0 and will be removed in 1.12.0, "
-                "use 'mlrun/mlrun' instead."
-            )
-            if auto_replaced:
-                message += (
-                    " Since your client version is >= 1.10.0, the image will be automatically "
-                    "replaced with mlrun/mlrun."
-                )
-            warnings.warn(
-                message,
-                # TODO: Remove this in 1.12.0
-                FutureWarning,
-            )
+        mlrun.utils.helpers.warn_on_deprecated_image(image)
 
         # TODO: need a better way to decide whether a function requires a build
         if require_build and image and not runtime.spec.build.base_image:

mlrun/launcher/local.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 import os
 import pathlib
-from os import environ
 from typing import Callable, Optional, Union
 
 import mlrun.common.constants as mlrun_constants
@@ -252,9 +251,6 @@ class ClientLocalLauncher(launcher.ClientBaseLauncher):
             # copy the code/base-spec to the local function (for the UI and code logging)
             fn.spec.description = runtime.spec.description
             fn.spec.build = runtime.spec.build
-            serving_spec = getattr(runtime.spec, "serving_spec", None)
-            if serving_spec:
-                environ["SERVING_SPEC_ENV"] = serving_spec
 
         run.spec.handler = handler
         run.spec.reset_on_run = reset_on_run

mlrun/model_monitoring/applications/base.py

@@ -166,13 +166,29 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         return result
 
     @staticmethod
+    def _check_writer_is_up(project: "mlrun.MlrunProject") -> None:
+        try:
+            project.get_function(
+                mm_constants.MonitoringFunctionNames.WRITER, ignore_cache=True
+            )
+        except mlrun.errors.MLRunNotFoundError:
+            raise mlrun.errors.MLRunValueError(
+                "Writing outputs to the databases is blocked as the model monitoring infrastructure is disabled.\n"
+                "To unblock, enable model monitoring with `project.enable_model_monitoring()`."
+            )
+
+    @classmethod
     @contextmanager
     def _push_to_writer(
+        cls,
         *,
         write_output: bool,
         stream_profile: Optional[ds_profile.DatastoreProfile],
+        project: "mlrun.MlrunProject",
     ) -> Iterator[dict[str, list[tuple]]]:
         endpoints_output: dict[str, list[tuple]] = defaultdict(list)
+        if write_output:
+            cls._check_writer_is_up(project)
         try:
             yield endpoints_output
         finally:
@@ -220,6 +236,9 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         for an MLRun job.
         This method should not be called directly.
         """
+        project = context.get_project_object()
+        if not project:
+            raise mlrun.errors.MLRunValueError("Could not load project from context")
 
         if write_output and (
             not endpoints or sample_data is not None or reference_data is not None
@@ -236,7 +255,7 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         )
 
         with self._push_to_writer(
-            write_output=write_output, stream_profile=stream_profile
+            write_output=write_output, stream_profile=stream_profile, project=project
         ) as endpoints_output:
 
             def call_do_tracking(event: Optional[dict] = None):
@@ -249,6 +268,7 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
                         event=event,
                         application_name=self.__class__.__name__,
                         context=context,
+                        project=project,
                         sample_df=sample_data,
                         feature_stats=feature_stats,
                     )

mlrun/model_monitoring/applications/context.py

@@ -137,13 +137,14 @@ class MonitoringApplicationContext:
         cls,
         context: "mlrun.MLClientCtx",
         *,
+        project: Optional["mlrun.MlrunProject"] = None,
         application_name: str,
         event: dict[str, Any],
         model_endpoint_dict: Optional[dict[str, ModelEndpoint]] = None,
         sample_df: Optional[pd.DataFrame] = None,
         feature_stats: Optional[FeatureStats] = None,
     ) -> "MonitoringApplicationContext":
-        project = context.get_project_object()
+        project = project or context.get_project_object()
         if not project:
             raise mlrun.errors.MLRunValueError("Could not load project from context")
         logger = context.logger