mlrun 1.10.0rc13__py3-none-any.whl → 1.10.0rc42__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,20 +11,38 @@
 # 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
-from copy import deepcopy
+from copy import copy
 from datetime import timedelta
 from typing import Any, Optional, Union
 
+import numpy as np
 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 logger
+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):
@@ -45,46 +63,20 @@ class MonitoringPreProcessor(storey.MapClass):
         result_path = model_monitoring_data.get(MonitoringData.RESULT_PATH)
         input_path = model_monitoring_data.get(MonitoringData.INPUT_PATH)
 
-        result = self._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):
-            if len(result) > 1:
-                # transpose by key the outputs:
-                outputs = self.transpose_by_key(result, output_schema)
-            elif len(result) == 1:
-                outputs = (
-                    result[output_schema[0]]
-                    if output_schema
-                    else list(result.values())[0]
-                )
-            else:
-                outputs = []
-            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 = self._get_data_from_path(input_path, event_inputs)
-        if isinstance(event_inputs, dict):
-            if len(event_inputs) > 1:
-                # transpose by key the inputs:
-                inputs = self.transpose_by_key(event_inputs, input_schema)
-            else:
-                inputs = (
-                    event_inputs[input_schema[0]]
-                    if input_schema
-                    else list(result.values())[0]
-                )
-        else:
-            inputs = event_inputs
+        logger.debug(
+            "output and input schema retrieved",
+            output_schema=output_schema,
+            input_schema=input_schema,
+        )
+
+        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]):
@@ -104,48 +96,124 @@ class MonitoringPreProcessor(storey.MapClass):
                     output_len=len(outputs),
                     schema_len=len(output_schema),
                 )
-        request = {"inputs": inputs, "id": getattr(event, "id", None)}
-        resp = {"outputs": outputs}
+            if len(inputs) != len(outputs):
+                logger.warn(
+                    "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),
+            "input_schema": new_input_schema,
+        }
+        resp = {"outputs": outputs, "output_schema": new_output_schema}
 
         return request, resp
 
+    def get_listed_data(
+        self,
+        raw_data: dict,
+        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_to_transpose, schema: Optional[list[str]] = None
-    ) -> list[list[float]]:
-        values = (
-            list(data_to_transpose.values())
-            if not schema
-            else [data_to_transpose[key] for key in schema]
-        )
-        if values and not isinstance(values[0], list):
-            values = [values]
-        transposed = (
-            list(map(list, zip(*values)))
-            if all(isinstance(v, list) for v in values) and len(values) > 1
-            else values
-        )
-        return transposed
+        data: dict, schema: Optional[Union[str, list[str]]] = None
+    ) -> tuple[Union[list[Any], list[list[Any]]], list[str]]:
+        """
+        Transpose values from a dictionary by keys.
 
-    @staticmethod
-    def _get_data_from_path(
-        path: Union[str, list[str], None], data: dict
-    ) -> dict[str, Any]:
-        if isinstance(path, str):
-            output_data = data.get(path)
-        elif isinstance(path, list):
-            output_data = deepcopy(data)
-            for key in path:
-                output_data = output_data.get(key, {})
-        elif path is None:
-            output_data = data
+        Given a dictionary and an optional schema (a key or list of keys), this function:
+        - Extracts the values for the specified keys (or all keys if no schema is provided).
+        - Ensures the data is represented as a list of rows, then transposes it (i.e., switches rows to columns).
+        - Handles edge cases:
+            * If a single scalar or single-element list is provided, returns a flat list.
+            * If a single key is provided (as a string or a list with one element), handles it properly.
+            * If only one row with len of one remains after transposition, unwraps it to avoid nested list-of-one.
+
+        Example::
+
+            transpose_by_key({"a": 1})
+            # returns: [1]
+
+            transpose_by_key({"a": [1, 2]})
+            # returns: [1 ,2]
+
+            transpose_by_key({"a": [1, 2], "b": [3, 4]})
+            # returns: [[1, 3], [2, 4]]
+
+        :param data:     Dictionary with values that are either scalars or lists.
+        :param schema:   Optional key or list of keys to extract. If not provided, all keys are used.
+                         Can be a string (single key) or a list of strings.
+
+        :return:         Transposed values:
+                         * If result is a single column or row, returns a flat list.
+                         * 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(normalize_data.keys())
+            new_schema = keys
+        elif isinstance(schema, str):
+            keys = [mlrun.feature_store.api.norm_column_name(schema)]
         else:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                "Expected path be of type str or list of str or None"
+            keys = [mlrun.feature_store.api.norm_column_name(key) for key in schema]
+
+        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} are not contained in the data keys {list(data.keys())}."
             )
-        if isinstance(output_data, (int, float)):
-            output_data = [output_data]
-        return output_data
+
+        # 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)
+
+        if not (all_scalars or all_lists):
+            raise ValueError(
+                "All values must be either scalars or lists of equal length."
+            )
+
+        if all_scalars:
+            transposed = np.array([values], dtype=object)
+        elif all_lists and len(keys) > 1:
+            arrays = [np.array(v, dtype=object) for v in values]
+            mat = np.stack(arrays, axis=0)
+            transposed = mat.T
+        else:
+            return values[0], new_schema
+
+        if transposed.shape[1] == 1 and transposed.shape[0] == 1:
+            # Transform [[0]] -> [0]:
+            return transposed[:, 0].tolist(), new_schema
+        return transposed.tolist(), new_schema
 
     def do(self, event):
         monitoring_event_list = []
@@ -168,6 +236,12 @@ class MonitoringPreProcessor(storey.MapClass):
                     request, resp = self.reconstruct_request_resp_fields(
                         event, model, monitoring_data[model]
                     )
+                    if hasattr(event, "_original_timestamp"):
+                        when = event._original_timestamp
+                    else:
+                        when = event._metadata.get(model, {}).get(
+                            mm_schemas.StreamProcessingEvent.WHEN
+                        )
                     monitoring_event_list.append(
                         {
                             mm_schemas.StreamProcessingEvent.MODEL: model,
@@ -177,17 +251,16 @@ class MonitoringPreProcessor(storey.MapClass):
                             mm_schemas.StreamProcessingEvent.MICROSEC: event._metadata.get(
                                 model, {}
                             ).get(mm_schemas.StreamProcessingEvent.MICROSEC),
-                            mm_schemas.StreamProcessingEvent.WHEN: event._metadata.get(
-                                model, {}
-                            ).get(mm_schemas.StreamProcessingEvent.WHEN),
+                            mm_schemas.StreamProcessingEvent.WHEN: when,
                             mm_schemas.StreamProcessingEvent.ENDPOINT_ID: monitoring_data[
                                 model
                             ].get(
                                 mlrun.common.schemas.MonitoringData.MODEL_ENDPOINT_UID
                             ),
-                            mm_schemas.StreamProcessingEvent.LABELS: monitoring_data[
+                            mm_schemas.StreamProcessingEvent.LABELS: event.body[
                                 model
-                            ].get(mlrun.common.schemas.MonitoringData.OUTPUTS),
+                            ].get("labels")
+                            or {},
                             mm_schemas.StreamProcessingEvent.FUNCTION_URI: self.server.function_uri
                             if self.server
                             else None,
@@ -212,6 +285,10 @@ class MonitoringPreProcessor(storey.MapClass):
             request, resp = self.reconstruct_request_resp_fields(
                 event, model, monitoring_data[model]
             )
+            if hasattr(event, "_original_timestamp"):
+                when = event._original_timestamp
+            else:
+                when = event._metadata.get(mm_schemas.StreamProcessingEvent.WHEN)
             monitoring_event_list.append(
                 {
                     mm_schemas.StreamProcessingEvent.MODEL: model,
@@ -221,25 +298,20 @@ class MonitoringPreProcessor(storey.MapClass):
                     mm_schemas.StreamProcessingEvent.MICROSEC: event._metadata.get(
                         mm_schemas.StreamProcessingEvent.MICROSEC
                     ),
-                    mm_schemas.StreamProcessingEvent.WHEN: event._metadata.get(
-                        mm_schemas.StreamProcessingEvent.WHEN
-                    ),
+                    mm_schemas.StreamProcessingEvent.WHEN: when,
                     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: event.body.get("labels")
+                    or {},
                     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: event.body.get(
                         mm_schemas.StreamProcessingEvent.ERROR
-                    ]
-                    if mm_schemas.StreamProcessingEvent.ERROR in event.body
-                    else None,
+                    ),
                     mm_schemas.StreamProcessingEvent.METRICS: event.body[
                         mm_schemas.StreamProcessingEvent.METRICS
                     ]
@@ -259,6 +331,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
         )
@@ -279,43 +354,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):
@@ -337,7 +416,9 @@ class SamplingStep(storey.MapClass):
             event=event,
             sampling_percentage=self.sampling_percentage,
         )
-        if self.sampling_percentage != 100:
+        if self.sampling_percentage != 100 and not event.get(
+            mm_schemas.StreamProcessingEvent.ERROR
+        ):
             request = event[mm_schemas.StreamProcessingEvent.REQUEST]
             num_of_inputs = len(request["inputs"])
             sampled_requests_indices = self._pick_random_requests(

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: