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

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:
@@ -509,6 +556,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 +617,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

@@ -148,6 +148,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
@@ -63,26 +62,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/projects/pipelines.py

@@ -39,7 +39,12 @@ from mlrun.utils import (
 
 from ..common.helpers import parse_versioned_object_uri
 from ..config import config
-from ..run import _run_pipeline, retry_pipeline, wait_for_pipeline_completion
+from ..run import (
+    _run_pipeline,
+    retry_pipeline,
+    terminate_pipeline,
+    wait_for_pipeline_completion,
+)
 from ..runtimes.pod import AutoMountType
 
 
@@ -696,6 +701,24 @@ class _KFPRunner(_PipelineRunner):
         )
         return run_id
 
+    @classmethod
+    def terminate(
+        cls,
+        run: "_PipelineRunStatus",
+        project: typing.Optional["mlrun.projects.MlrunProject"] = None,
+    ) -> str:
+        project_name = project.metadata.name if project else ""
+        logger.info(
+            "Terminating pipeline",
+            run_id=run.run_id,
+            project=project_name,
+        )
+        run_id = terminate_pipeline(
+            run.run_id,
+            project=project_name,
+        )
+        return run_id
+
     @staticmethod
     def wait_for_completion(
         run: "_PipelineRunStatus",
@@ -1145,7 +1168,9 @@ def load_and_run_workflow(
         notification.when = ["running"]
 
     workflow_log_message = workflow_name or workflow_path
-    context.logger.info(f"Running workflow {workflow_log_message} from remote")
+    context.logger.info(
+        "Running workflow from remote", workflow_log_message=workflow_log_message
+    )
     run = project.run(
         name=workflow_name,
         workflow_path=workflow_path,
@@ -1162,6 +1187,11 @@ def load_and_run_workflow(
         notifications=start_notifications,
         context=context,
     )
+    # Patch the current run object (the workflow-runner) with the workflow-id label
+    context.logger.info(
+        "Associating workflow-runner with workflow ID", run_id=run.run_id
+    )
+    context.set_label("workflow-id", run.run_id)
     context.log_result(key="workflow_id", value=run.run_id)
     context.log_result(key="engine", value=run._engine.engine, commit=True)
 
@@ -1321,4 +1351,4 @@ def import_remote_project(
         sync_functions=True,
     )
 
-    context.logger.info(f"Loaded project {project.name} successfully")
+    context.logger.info("Loaded project successfully", project_name=project.name)

mlrun/projects/project.py

@@ -2518,7 +2518,6 @@ class MlrunProject(ModelObj):
 
     def enable_model_monitoring(
         self,
-        default_controller_image: str = "mlrun/mlrun",
         base_period: int = 10,
         image: str = "mlrun/mlrun",
         *,
@@ -2534,7 +2533,6 @@ class MlrunProject(ModelObj):
         The stream function goal is to monitor the log of the data stream. It is triggered when a new log entry
         is detected. It processes the new events into statistics that are then written to statistics databases.
 
-        :param default_controller_image:          Deprecated.
         :param base_period:                       The time period in minutes in which the model monitoring controller
                                                   function is triggered. By default, the base period is 10 minutes
                                                   (which is also the minimum value for production environments).
@@ -2562,14 +2560,6 @@ class MlrunProject(ModelObj):
                                                   background, including the histogram data drift app if selected.
         :param fetch_credentials_from_sys_config: If true, fetch the credentials from the system configuration.
         """
-        if default_controller_image != "mlrun/mlrun":
-            # TODO: Remove this in 1.10.0
-            warnings.warn(
-                "'default_controller_image' is deprecated in 1.7.0 and will be removed in 1.10.0, "
-                "use 'image' instead",
-                FutureWarning,
-            )
-            image = default_controller_image
         if base_period < 10:
             logger.warn(
                 "enable_model_monitoring: 'base_period' < 10 minutes is not supported in production environments",
@@ -3848,7 +3838,8 @@ class MlrunProject(ModelObj):
             )
 
         The replication factor and timeout configuration might need to be adjusted according to your Confluent cluster
-        type and settings.
+        type and settings. Nuclio annotations for the model monitoring infrastructure and application functions are
+        supported through ``kwargs_public={"nuclio_annotations": {...}, ...}``.
 
         :param tsdb_profile_name:         The datastore profile name of the time-series database to be used in model
                                           monitoring. The supported profiles are:
@@ -4278,11 +4269,17 @@ class MlrunProject(ModelObj):
         function = mlrun.new_function("mlrun--project--image--builder", kind="job")
 
         if self.spec.source and not self.spec.load_source_on_run:
-            function.with_source_archive(
-                source=self.spec.source,
-                target_dir=target_dir,
-                pull_at_runtime=False,
-            )
+            if self.spec.source.startswith("db://"):
+                logger.debug(
+                    "Project source is 'db://', which refers to metadata stored in the MLRun DB."
+                    " Skipping source archive setup for image build"
+                )
+            else:
+                function.with_source_archive(
+                    source=self.spec.source,
+                    target_dir=target_dir,
+                    pull_at_runtime=False,
+                )
 
         build = self.spec.build
         result = self.build_function(

mlrun/run.py

@@ -894,7 +894,6 @@ def _run_pipeline(
 def retry_pipeline(
     run_id: str,
     project: str,
-    namespace: Optional[str] = None,
 ) -> str:
     """Retry a pipeline run.
 
@@ -903,7 +902,6 @@ def retry_pipeline(
 
     :param run_id: ID of the pipeline run to retry.
     :param project: name of the project associated with the pipeline run.
-    :param namespace: Optional; Kubernetes namespace to use if not the default.
 
     :returns: ID of the retried pipeline run or the ID of a cloned run if the original run is not retryable.
     :raises ValueError: If access to the remote API service is not available.
@@ -918,7 +916,6 @@ def retry_pipeline(
     pipeline_run_id = mldb.retry_pipeline(
         run_id=run_id,
         project=project,
-        namespace=namespace,
     )
     if pipeline_run_id == run_id:
         logger.info(
@@ -931,6 +928,35 @@ def retry_pipeline(
     return pipeline_run_id
 
 
+def terminate_pipeline(
+    run_id: str,
+    project: str,
+) -> str:
+    """Terminate a pipeline run.
+
+    This function terminates a running pipeline with the specified run ID. If the run is not in a
+    terminable state, an error is raised.
+
+    :param run_id: ID of the pipeline run to terminate.
+    :param project: name of the project associated with the pipeline run.
+
+    :returns: ID of the terminate pipeline run background task.
+    :raises ValueError: If access to the remote API service is not available.
+    """
+    mldb = mlrun.db.get_run_db()
+    if mldb.kind != "http":
+        raise ValueError(
+            "Terminating a pipeline requires access to remote API service. "
+            "Please set the dbpath URL."
+        )
+
+    pipeline_run_task = mldb.terminate_pipeline(
+        run_id=run_id,
+        project=project,
+    )
+    return pipeline_run_task["metadata"]["id"]
+
+
 def wait_for_pipeline_completion(
     run_id,
     timeout=60 * 60,
@@ -997,7 +1023,10 @@ def wait_for_pipeline_completion(
             _wait_for_pipeline_completion,
         )
     else:
-        client = mlrun_pipelines.utils.get_client(namespace=namespace)
+        client = mlrun_pipelines.utils.get_client(
+            logger=logger,
+            namespace=namespace,
+        )
         resp = client.wait_for_run_completion(run_id, timeout)
         if resp:
             resp = resp.to_dict()
@@ -1058,7 +1087,10 @@ def get_pipeline(
         )
 
     else:
-        client = mlrun_pipelines.utils.get_client(namespace=namespace)
+        client = mlrun_pipelines.utils.get_client(
+            logger=logger,
+            namespace=namespace,
+        )
         resp = client.get_run(run_id)
         if resp:
             resp = resp.to_dict()

mlrun/runtimes/nuclio/serving.py

@@ -11,7 +11,6 @@
 # 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 copy
 import json
 import os
 import warnings
@@ -482,6 +481,7 @@ class ServingRuntime(RemoteRuntime):
                 state = TaskStep(
                     class_name,
                     class_args,
+                    name=key,
                     handler=handler,
                     function=child_function,
                     model_endpoint_creation_strategy=creation_strategy,
@@ -751,13 +751,10 @@ class ServingRuntime(RemoteRuntime):
             set_paths(workdir)
             os.chdir(workdir)
 
-        system_graph = None
-        if isinstance(self.spec.graph, RootFlowStep):
-            system_graph = add_system_steps_to_graph(copy.deepcopy(self.spec.graph))
         server = create_graph_server(
             parameters=self.spec.parameters,
             load_mode=self.spec.load_mode,
-            graph=system_graph or self.spec.graph,
+            graph=self.spec.graph,
             verbose=self.verbose,
             current_function=current_function,
             graph_initializer=self.spec.graph_initializer,
@@ -778,6 +775,18 @@ class ServingRuntime(RemoteRuntime):
             monitoring_mock=self.spec.track_models,
         )
 
+        if (
+            isinstance(self.spec.graph, RootFlowStep)
+            and self.spec.graph.include_monitored_step()
+        ):
+            server.graph = add_system_steps_to_graph(
+                server.project,
+                server.graph,
+                self.spec.track_models,
+                server.context,
+                self.spec,
+            )
+
         if workdir:
             os.chdir(old_workdir)
 

mlrun/serving/__init__.py

@@ -27,6 +27,7 @@ __all__ = [
     "ModelRunner",
     "Model",
     "ModelSelector",
+    "MonitoredStep",
 ]
 
 from .routers import ModelRouter, VotingEnsemble  # noqa
@@ -45,6 +46,7 @@ from .states import (
     ModelRunner,
     Model,
     ModelSelector,
+    MonitoredStep,
 )  # noqa
 from .v1_serving import MLModelServer, new_v1_model_server  # noqa
 from .v2_serving import V2ModelServer  # noqa