mlrun 1.10.0rc18__py3-none-any.whl → 1.11.0rc16__py3-none-any.whl

mlrun/serving/steps.py

@@ -0,0 +1,62 @@
+# 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.
+
+from typing import Union
+
+import storey
+
+import mlrun.errors
+
+
+class ChoiceByField(storey.Choice):
+    """
+    Selects downstream outlets to route each event based on a predetermined field.
+    :param field_name: event field name that contains the step name or names of the desired outlet or outlets
+    """
+
+    def __init__(self, field_name: Union[str, list[str]], **kwargs):
+        self.field_name = field_name
+        super().__init__(**kwargs)
+
+    def select_outlets(self, event):
+        # Case 1: Missing field
+        if self.field_name not in event:
+            raise mlrun.errors.MLRunRuntimeError(
+                f"Field '{self.field_name}' is not contained in the event keys {list(event.keys())}."
+            )
+
+        outlet = event[self.field_name]
+
+        # Case 2: Field exists but is None
+        if outlet is None:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"Field '{self.field_name}' exists but its value is None."
+            )
+
+        # Case 3: Invalid type
+        if not isinstance(outlet, str | list | tuple):
+            raise mlrun.errors.MLRunInvalidArgumentTypeError(
+                f"Field '{self.field_name}' must be a string or list of strings "
+                f"but is instead of type '{type(outlet).__name__}'."
+            )
+
+        outlets = [outlet] if isinstance(outlet, str) else outlet
+
+        # Case 4: Empty list or tuple
+        if not outlets:
+            raise mlrun.errors.MLRunRuntimeError(
+                f"The value of the key '{self.field_name}' cannot be an empty {type(outlets).__name__}."
+            )
+
+        return outlets

mlrun/serving/system_steps.py

@@ -11,8 +11,9 @@
 # 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 random
+import typing
+from copy import copy
 from datetime import timedelta
 from typing import Any, Optional, Union
 
@@ -22,11 +23,29 @@ import storey
 import mlrun
 import mlrun.artifacts
 import mlrun.common.schemas.model_monitoring as mm_schemas
+import mlrun.feature_store
 import mlrun.serving
+from mlrun.common.model_monitoring.helpers import (
+    get_model_endpoints_creation_task_status,
+)
 from mlrun.common.schemas import MonitoringData
 from mlrun.utils import get_data_from_path, logger
 
 
+class MatchingEndpointsState(mlrun.common.types.StrEnum):
+    all_matched = "all_matched"
+    not_all_matched = "not_all_matched"
+    no_check_needed = "no_check_needed"
+    not_yet_checked = "not_yet_matched"
+
+    @staticmethod
+    def success_states() -> list[str]:
+        return [
+            MatchingEndpointsState.all_matched,
+            MatchingEndpointsState.no_check_needed,
+        ]
+
+
 class MonitoringPreProcessor(storey.MapClass):
     """preprocess step, reconstructs the serving output event body to StreamProcessingEvent schema"""
 
@@ -45,33 +64,24 @@ class MonitoringPreProcessor(storey.MapClass):
         result_path = model_monitoring_data.get(MonitoringData.RESULT_PATH)
         input_path = model_monitoring_data.get(MonitoringData.INPUT_PATH)
 
-        result = get_data_from_path(result_path, event.body.get(model, event.body))
         output_schema = model_monitoring_data.get(MonitoringData.OUTPUTS)
         input_schema = model_monitoring_data.get(MonitoringData.INPUTS)
-        logger.debug("output schema retrieved", output_schema=output_schema)
-        if isinstance(result, dict):
-            # transpose by key the outputs:
-            outputs = self.transpose_by_key(result, output_schema)
-            if not output_schema:
-                logger.warn(
-                    "Output schema was not provided using Project:log_model or by ModelRunnerStep:add_model order "
-                    "may not preserved"
-                )
-        else:
-            outputs = result
-
-        event_inputs = event._metadata.get("inputs", {})
-        event_inputs = get_data_from_path(input_path, event_inputs)
-        if isinstance(event_inputs, dict):
-            # transpose by key the inputs:
-            inputs = self.transpose_by_key(event_inputs, input_schema)
-            if not input_schema:
-                logger.warn(
-                    "Input schema was not provided using by ModelRunnerStep:add_model, order "
-                    "may not preserved"
-                )
+        logger.debug(
+            "output and input schema retrieved",
+            output_schema=output_schema,
+            input_schema=input_schema,
+        )
+        if event.body and isinstance(event.body, list):
+            outputs, new_output_schema = self.get_listed_data(
+                event.body, result_path, output_schema
+            )
         else:
-            inputs = event_inputs
+            outputs, new_output_schema = self.get_listed_data(
+                event.body.get(model, event.body), result_path, output_schema
+            )
+        inputs, new_input_schema = self.get_listed_data(
+            event._metadata.get("inputs", {}), input_path, input_schema
+        )
 
         if outputs and isinstance(outputs[0], list):
             if output_schema and len(output_schema) != len(outputs[0]):
@@ -96,15 +106,43 @@ class MonitoringPreProcessor(storey.MapClass):
                     "outputs and inputs are not in the same length check 'input_path' and "
                     "'output_path' was specified if needed"
                 )
-        request = {"inputs": inputs, "id": getattr(event, "id", None)}
-        resp = {"outputs": outputs}
+        request = {
+            "inputs": inputs,
+            "id": getattr(event, "id", None),
+            "input_schema": new_input_schema,
+        }
+        resp = {"outputs": outputs, "output_schema": new_output_schema}
 
         return request, resp
 
+    def get_listed_data(
+        self,
+        raw_data: typing.Union[dict, list],
+        data_path: Optional[Union[list[str], str]] = None,
+        schema: Optional[list[str]] = None,
+    ):
+        """Get data from a path and transpose it by keys if dict is provided."""
+        new_schema = None
+        data_from_path = get_data_from_path(data_path, raw_data)
+        if isinstance(data_from_path, dict):
+            # transpose by key the inputs:
+            listed_data, new_schema = self.transpose_by_key(data_from_path, schema)
+            new_schema = new_schema or schema
+            if not schema:
+                logger.warn(
+                    f"No schema provided through add_model(); the order of {data_from_path} "
+                    "may not be preserved."
+                )
+        elif not isinstance(data_from_path, list):
+            listed_data = [data_from_path]
+        else:
+            listed_data = data_from_path
+        return listed_data, new_schema
+
     @staticmethod
     def transpose_by_key(
         data: dict, schema: Optional[Union[str, list[str]]] = None
-    ) -> Union[list[Any], list[list[Any]]]:
+    ) -> tuple[Union[list[Any], list[list[Any]]], list[str]]:
         """
         Transpose values from a dictionary by keys.
 
@@ -136,25 +174,32 @@ class MonitoringPreProcessor(storey.MapClass):
                          * If result is a matrix, returns a list of lists.
 
         :raises ValueError: If the values include a mix of scalars and lists, or if the list lengths do not match.
+                mlrun.MLRunInvalidArgumentError if the schema keys are not contained in the data keys.
         """
-
+        new_schema = None
+        # Normalize keys in data:
+        normalize_data = {
+            mlrun.feature_store.api.norm_column_name(k): copy(v)
+            for k, v in data.items()
+        }
         # Normalize schema to list
         if not schema:
-            keys = list(data.keys())
+            keys = list(normalize_data.keys())
+            new_schema = keys
         elif isinstance(schema, str):
-            keys = [schema]
+            keys = [mlrun.feature_store.api.norm_column_name(schema)]
         else:
-            keys = schema
+            keys = [mlrun.feature_store.api.norm_column_name(key) for key in schema]
 
-        values = [data[key] for key in keys if key in data]
+        values = [normalize_data[key] for key in keys if key in normalize_data]
         if len(values) != len(keys):
             raise mlrun.MLRunInvalidArgumentError(
-                f"Schema keys {keys} do not match the data keys {list(data.keys())}."
+                f"Schema keys {keys} are not contained in the data keys {list(data.keys())}."
             )
 
         # Detect if all are scalars ie: int,float,str
-        all_scalars = all(not isinstance(v, (list, tuple, np.ndarray)) for v in values)
-        all_lists = all(isinstance(v, (list, tuple, np.ndarray)) for v in values)
+        all_scalars = all(not isinstance(v, list | tuple | np.ndarray) for v in values)
+        all_lists = all(isinstance(v, list | tuple | np.ndarray) for v in values)
 
         if not (all_scalars or all_lists):
             raise ValueError(
@@ -168,12 +213,12 @@ class MonitoringPreProcessor(storey.MapClass):
             mat = np.stack(arrays, axis=0)
             transposed = mat.T
         else:
-            return values[0]
+            return values[0], new_schema
 
         if transposed.shape[1] == 1 and transposed.shape[0] == 1:
             # Transform [[0]] -> [0]:
-            return transposed[:, 0].tolist()
-        return transposed.tolist()
+            return transposed[:, 0].tolist(), new_schema
+        return transposed.tolist(), new_schema
 
     def do(self, event):
         monitoring_event_list = []
@@ -202,6 +247,21 @@ class MonitoringPreProcessor(storey.MapClass):
                         when = event._metadata.get(model, {}).get(
                             mm_schemas.StreamProcessingEvent.WHEN
                         )
+                    #  if the body is not a dict, use empty labels, error and metrics
+                    if isinstance(event.body[model], dict):
+                        body_by_model = event.body[model]
+                        labels = body_by_model.get("labels") or {}
+                        error = body_by_model.get(
+                            mm_schemas.StreamProcessingEvent.ERROR
+                        )
+                        metrics = body_by_model.get(
+                            mm_schemas.StreamProcessingEvent.METRICS
+                        )
+                    else:
+                        labels = {}
+                        error = None
+                        metrics = None
+
                     monitoring_event_list.append(
                         {
                             mm_schemas.StreamProcessingEvent.MODEL: model,
@@ -217,26 +277,14 @@ class MonitoringPreProcessor(storey.MapClass):
                             ].get(
                                 mlrun.common.schemas.MonitoringData.MODEL_ENDPOINT_UID
                             ),
-                            mm_schemas.StreamProcessingEvent.LABELS: monitoring_data[
-                                model
-                            ].get(mlrun.common.schemas.MonitoringData.OUTPUTS),
+                            mm_schemas.StreamProcessingEvent.LABELS: labels,
                             mm_schemas.StreamProcessingEvent.FUNCTION_URI: self.server.function_uri
                             if self.server
                             else None,
                             mm_schemas.StreamProcessingEvent.REQUEST: request,
                             mm_schemas.StreamProcessingEvent.RESPONSE: resp,
-                            mm_schemas.StreamProcessingEvent.ERROR: event.body[model][
-                                mm_schemas.StreamProcessingEvent.ERROR
-                            ]
-                            if mm_schemas.StreamProcessingEvent.ERROR
-                            in event.body[model]
-                            else None,
-                            mm_schemas.StreamProcessingEvent.METRICS: event.body[model][
-                                mm_schemas.StreamProcessingEvent.METRICS
-                            ]
-                            if mm_schemas.StreamProcessingEvent.METRICS
-                            in event.body[model]
-                            else None,
+                            mm_schemas.StreamProcessingEvent.ERROR: error,
+                            mm_schemas.StreamProcessingEvent.METRICS: metrics,
                         }
                     )
         elif monitoring_data:
@@ -248,6 +296,15 @@ class MonitoringPreProcessor(storey.MapClass):
                 when = event._original_timestamp
             else:
                 when = event._metadata.get(mm_schemas.StreamProcessingEvent.WHEN)
+            #  if the body is not a dict, use empty labels, error and metrics
+            if isinstance(event.body, dict):
+                labels = event.body.get("labels") or {}
+                error = event.body.get(mm_schemas.StreamProcessingEvent.ERROR)
+                metrics = event.body.get(mm_schemas.StreamProcessingEvent.METRICS)
+            else:
+                labels = {}
+                error = None
+                metrics = None
             monitoring_event_list.append(
                 {
                     mm_schemas.StreamProcessingEvent.MODEL: model,
@@ -261,24 +318,14 @@ class MonitoringPreProcessor(storey.MapClass):
                     mm_schemas.StreamProcessingEvent.ENDPOINT_ID: monitoring_data[
                         model
                     ].get(mlrun.common.schemas.MonitoringData.MODEL_ENDPOINT_UID),
-                    mm_schemas.StreamProcessingEvent.LABELS: monitoring_data[model].get(
-                        mlrun.common.schemas.MonitoringData.OUTPUTS
-                    ),
+                    mm_schemas.StreamProcessingEvent.LABELS: labels,
                     mm_schemas.StreamProcessingEvent.FUNCTION_URI: self.server.function_uri
                     if self.server
                     else None,
                     mm_schemas.StreamProcessingEvent.REQUEST: request,
                     mm_schemas.StreamProcessingEvent.RESPONSE: resp,
-                    mm_schemas.StreamProcessingEvent.ERROR: event.body[
-                        mm_schemas.StreamProcessingEvent.ERROR
-                    ]
-                    if mm_schemas.StreamProcessingEvent.ERROR in event.body
-                    else None,
-                    mm_schemas.StreamProcessingEvent.METRICS: event.body[
-                        mm_schemas.StreamProcessingEvent.METRICS
-                    ]
-                    if mm_schemas.StreamProcessingEvent.METRICS in event.body
-                    else None,
+                    mm_schemas.StreamProcessingEvent.ERROR: error,
+                    mm_schemas.StreamProcessingEvent.METRICS: metrics,
                 }
             )
         event.body = monitoring_event_list
@@ -293,6 +340,9 @@ class BackgroundTaskStatus(storey.MapClass):
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
+        self.matching_endpoints = MatchingEndpointsState.not_yet_checked
+        self.graph_model_endpoint_uids: set = set()
+        self.listed_model_endpoint_uids: set = set()
         self.server: mlrun.serving.GraphServer = (
             getattr(self.context, "server", None) if self.context else None
         )
@@ -313,43 +363,47 @@ class BackgroundTaskStatus(storey.MapClass):
                 )
             )
         ):
-            background_task = mlrun.get_run_db().get_project_background_task(
-                self.server.project, self.server.model_endpoint_creation_task_name
-            )
-            self._background_task_check_timestamp = mlrun.utils.now_date()
-            self._log_background_task_state(background_task.status.state)
-            self._background_task_state = background_task.status.state
+            (
+                self._background_task_state,
+                self._background_task_check_timestamp,
+                self.listed_model_endpoint_uids,
+            ) = get_model_endpoints_creation_task_status(self.server)
+        if (
+            self.listed_model_endpoint_uids
+            and self.matching_endpoints == MatchingEndpointsState.not_yet_checked
+        ):
+            if not self.graph_model_endpoint_uids:
+                self.graph_model_endpoint_uids = collect_model_endpoint_uids(
+                    self.server
+                )
+
+            if self.graph_model_endpoint_uids.issubset(self.listed_model_endpoint_uids):
+                self.matching_endpoints = MatchingEndpointsState.all_matched
+        elif self.listed_model_endpoint_uids is None:
+            self.matching_endpoints = MatchingEndpointsState.no_check_needed
 
         if (
             self._background_task_state
             == mlrun.common.schemas.BackgroundTaskState.succeeded
+            and self.matching_endpoints in MatchingEndpointsState.success_states()
         ):
             return event
         else:
             return None
 
-    def _log_background_task_state(
-        self, background_task_state: mlrun.common.schemas.BackgroundTaskState
-    ):
-        logger.info(
-            "Checking model endpoint creation task status",
-            task_name=self.server.model_endpoint_creation_task_name,
-        )
-        if (
-            background_task_state
-            in mlrun.common.schemas.BackgroundTaskState.terminal_states()
-        ):
-            logger.info(
-                f"Model endpoint creation task completed with state {background_task_state}"
-            )
-        else:  # in progress
-            logger.info(
-                f"Model endpoint creation task is still in progress with the current state: "
-                f"{background_task_state}. Events will not be monitored for the next "
-                f"{mlrun.mlconf.model_endpoint_monitoring.model_endpoint_creation_check_period} seconds",
-                name=self.name,
-                background_task_check_timestamp=self._background_task_check_timestamp.isoformat(),
-            )
+
+def collect_model_endpoint_uids(server: mlrun.serving.GraphServer) -> set[str]:
+    """Collects all model endpoint UIDs from the server's graph steps."""
+    model_endpoint_uids = set()
+    for step in server.graph.steps.values():
+        if hasattr(step, "monitoring_data"):
+            for model in step.monitoring_data.keys():
+                uid = step.monitoring_data[model].get(
+                    mlrun.common.schemas.MonitoringData.MODEL_ENDPOINT_UID
+                )
+                if uid:
+                    model_endpoint_uids.add(uid)
+    return model_endpoint_uids
 
 
 class SamplingStep(storey.MapClass):

mlrun/serving/v2_serving.py

@@ -24,6 +24,9 @@ import mlrun.common.schemas.model_monitoring
 import mlrun.model_monitoring
 from mlrun.utils import logger, now_date
 
+from ..common.model_monitoring.helpers import (
+    get_model_endpoints_creation_task_status,
+)
 from .utils import StepToDict, _extract_input_data, _update_result_body
 
 
@@ -474,22 +477,18 @@ class V2ModelServer(StepToDict):
         ) or getattr(self.context, "server", None)
         if not self.context.is_mock or self.context.monitoring_mock:
             if server.model_endpoint_creation_task_name:
-                background_task = mlrun.get_run_db().get_project_background_task(
-                    server.project, server.model_endpoint_creation_task_name
-                )
-                logger.debug(
-                    "Checking model endpoint creation task status",
-                    task_name=server.model_endpoint_creation_task_name,
+                background_task_state, _, _ = get_model_endpoints_creation_task_status(
+                    server
                 )
                 if (
-                    background_task.status.state
+                    background_task_state
                     in mlrun.common.schemas.BackgroundTaskState.terminal_states()
                 ):
                     logger.debug(
-                        f"Model endpoint creation task completed with state {background_task.status.state}"
+                        f"Model endpoint creation task completed with state {background_task_state}"
                     )
                     if (
-                        background_task.status.state
+                        background_task_state
                         == mlrun.common.schemas.BackgroundTaskState.succeeded
                     ):
                         self._model_logger = (
@@ -504,7 +503,7 @@ class V2ModelServer(StepToDict):
                 else:  # in progress
                     logger.debug(
                         f"Model endpoint creation task is still in progress with the current state: "
-                        f"{background_task.status.state}.",
+                        f"{background_task_state}.",
                         name=self.name,
                     )
             else:

mlrun/track/trackers/mlflow_tracker.py

@@ -217,7 +217,7 @@ class MLFlowTracker(Tracker):
             handler=handler,
             run_name=run.info.run_name,
             project_name=project.name,
-            uid=run.info.run_uuid,
+            uid=run.info.run_id,
         )
 
         # Create a context from the run object:
@@ -373,7 +373,7 @@ class MLFlowTracker(Tracker):
             # Import the MLFlow run's artifacts to MLRun (model are logged after the rest of artifacts
             # so the artifacts can be registered as extra data in the models):
             artifacts = {}
-            model_paths = []
+            model_uris = []
             for artifact in client.list_artifacts(run_id=run.info.run_id):
                 # Get the artifact's local path (MLFlow suggests that if the artifact is already in the local filesystem
                 # its local path will be returned:
@@ -381,29 +381,29 @@ class MLFlowTracker(Tracker):
                     run_id=run.info.run_id,
                     artifact_path=artifact.path,
                 )
-                # Check if the artifact is a model (will be logged after the artifacts):
-                if artifact.is_dir and os.path.exists(
-                    os.path.join(
-                        artifact_local_path, "MLmodel"
-                    )  # Add tag to show model dir
-                ):
-                    model_paths.append(artifact_local_path)
-                else:
-                    # Log the artifact:
-                    artifact = MLFlowTracker._log_artifact(
-                        context=context,
-                        key=pathlib.Path(artifact.path).name.replace(".", "_"),
-                        # Mlflow has the same name for files but with different extensions, so we add extension to name
-                        local_path=artifact_local_path,
-                        tmp_path=tmp_dir,
-                    )
-                    artifacts[artifact.key] = artifact
+                # Log the artifact:
+                artifact = MLFlowTracker._log_artifact(
+                    context=context,
+                    key=pathlib.Path(artifact.path).name.replace(".", "_"),
+                    # Mlflow has the same name for files but with different extensions, so we add extension to name
+                    local_path=artifact_local_path,
+                    tmp_path=tmp_dir,
+                )
+                artifacts[artifact.key] = artifact
+
+            # get all run model's uri's (artifact_location in mlflow 3.0.0).
+            logged_models = mlflow.search_logged_models(
+                filter_string=f"source_run_id = '{run.info.run_id}'",
+                output_format="list",
+            )
+            for logged_model in logged_models:
+                model_uris.append(logged_model.artifact_location)
 
-            for model_path in model_paths:
+            for model_uri in model_uris:
                 MLFlowTracker._log_model(
                     context=context,
-                    model_uri=model_path,
-                    key=pathlib.Path(model_path).stem,
+                    model_uri=model_uri,
+                    key=pathlib.Path(model_uri).stem,
                     metrics=results,
                     extra_data=artifacts,
                     tmp_path=tmp_dir,
@@ -439,20 +439,18 @@ class MLFlowTracker(Tracker):
 
         # Get the model info from MLFlow:
         model_info = mlflow.models.get_model_info(model_uri=model_uri)
+        # Download the model and set the path to local path:
+        local_model_path = mlflow.artifacts.download_artifacts(
+            artifact_uri=str(model_uri)
+        )
+        model_path = pathlib.Path(local_model_path)
 
         # Prepare the archive path:
-        model_uri = pathlib.Path(model_uri)
-        archive_path = pathlib.Path(tmp_path) / f"{model_uri.stem}.zip"
-        if not os.path.exists(model_uri):
-            local_path = mlflow.artifacts.download_artifacts(
-                artifact_uri=str(model_uri)
-            )
-            model_uri = pathlib.Path(local_path)
-
+        archive_path = pathlib.Path(tmp_path) / f"{model_path.name}.zip"
         # TODO add progress bar for the case of large files
         # Zip the artifact:
         with zipfile.ZipFile(archive_path, "w") as zip_file:
-            for path in model_uri.rglob("*"):
+            for path in model_path.rglob("*"):
                 zip_file.write(filename=path, arcname=path.relative_to(model_uri))
 
         # Get inputs and outputs info: