mlrun 1.7.2rc3__py3-none-any.whl → 1.8.0__py3-none-any.whl

mlrun/serving/utils.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 #
 import inspect
+from typing import Optional
 
 from mlrun.utils import get_in, update_in
 
@@ -55,7 +56,12 @@ class StepToDict:
         "kwargs",
     ]
 
-    def to_dict(self, fields: list = None, exclude: list = None, strip: bool = False):
+    def to_dict(
+        self,
+        fields: Optional[list] = None,
+        exclude: Optional[list] = None,
+        strip: bool = False,
+    ):
         """convert the step object to a python dictionary"""
         fields = fields or getattr(self, "_dict_fields", None)
         if not fields:
@@ -105,5 +111,10 @@ class MonitoringApplicationToDict(StepToDict):
 class RouterToDict(StepToDict):
     _STEP_KIND = "router"
 
-    def to_dict(self, fields: list = None, exclude: list = None, strip: bool = False):
+    def to_dict(
+        self,
+        fields: Optional[list] = None,
+        exclude: Optional[list] = None,
+        strip: bool = False,
+    ):
         return super().to_dict(exclude=["routes"], strip=strip)

mlrun/serving/v1_serving.py

@@ -18,6 +18,7 @@ import socket
 from copy import deepcopy
 from datetime import datetime
 from io import BytesIO
+from typing import Optional
 from urllib.request import urlopen
 
 import nuclio
@@ -33,7 +34,7 @@ serving_handler = "handler"
 def new_v1_model_server(
     name,
     model_class: str,
-    models: dict = None,
+    models: Optional[dict] = None,
     filename="",
     protocol="",
     image="",
@@ -68,7 +69,7 @@ def new_v1_model_server(
 
 
 class MLModelServer:
-    def __init__(self, name: str, model_dir: str = None, model=None):
+    def __init__(self, name: str, model_dir: Optional[str] = None, model=None):
         self.name = name
         self.ready = False
         self.model_dir = model_dir

mlrun/serving/v2_serving.py

@@ -12,20 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import random
 import threading
 import time
 import traceback
-from typing import Optional, Union
+from typing import Optional
 
 import mlrun.artifacts
 import mlrun.common.model_monitoring.helpers
 import mlrun.common.schemas.model_monitoring
 import mlrun.model_monitoring
-from mlrun.errors import err_to_str
 from mlrun.utils import logger, now_date
 
-from ..common.helpers import parse_versioned_object_uri
-from .server import GraphServer
 from .utils import StepToDict, _extract_input_data, _update_result_body
 
 
@@ -33,12 +31,12 @@ class V2ModelServer(StepToDict):
     def __init__(
         self,
         context=None,
-        name: str = None,
-        model_path: str = None,
+        name: Optional[str] = None,
+        model_path: Optional[str] = None,
         model=None,
         protocol=None,
-        input_path: str = None,
-        result_path: str = None,
+        input_path: Optional[str] = None,
+        result_path: Optional[str] = None,
         shard_by_endpoint: Optional[bool] = None,
         **kwargs,
     ):
@@ -97,9 +95,6 @@ class V2ModelServer(StepToDict):
         :param kwargs:     extra arguments (can be accessed using self.get_param(key))
         """
         self.name = name
-        self.version = ""
-        if name and ":" in name:
-            self.name, self.version = name.split(":", 1)
         self.context = context
         self.ready = False
         self.error = ""
@@ -110,21 +105,17 @@ class V2ModelServer(StepToDict):
         self._result_path = result_path
         self._kwargs = kwargs  # for to_dict()
         self._params = kwargs
-        self._model_logger = (
-            _ModelLogPusher(self, context)
-            if context and context.stream.enabled
-            else None
-        )
-
         self.metrics = {}
         self.labels = {}
         self.model = None
         if model:
             self.model = model
             self.ready = True
-        self._versioned_model_name = None
-        self.model_endpoint_uid = None
+        self.model_endpoint_uid = kwargs.get("model_endpoint_uid", None)
         self.shard_by_endpoint = shard_by_endpoint
+        self._model_logger = None
+        self.initialized = False
+        self.output_schema = kwargs.get("outputs", [])
 
     def _load_and_update_state(self):
         try:
@@ -136,7 +127,7 @@ class V2ModelServer(StepToDict):
         self.ready = True
         self.context.logger.info(f"model {self.name} was loaded")
 
-    def post_init(self, mode="sync"):
+    def post_init(self, mode="sync", **kwargs):
         """sync/async model loading, for internal use"""
         if not self.ready:
             if mode == "async":
@@ -146,17 +137,35 @@ class V2ModelServer(StepToDict):
             else:
                 self._load_and_update_state()
 
-        server = getattr(self.context, "_server", None) or getattr(
-            self.context, "server", None
-        )
-        if not server:
-            logger.warn("GraphServer not initialized for VotingEnsemble instance")
-            return
+        if self.ready and not self.context.is_mock and not self.model_spec:
+            self.get_model()
 
-        if not self.context.is_mock or self.context.monitoring_mock:
-            self.model_endpoint_uid = _init_endpoint_record(
-                graph_server=server, model=self
-            )
+        if self.model_spec:
+            self.output_schema = self.output_schema or [
+                feature.name for feature in self.model_spec.outputs
+            ]
+
+        if (
+            kwargs.get("endpoint_type", mlrun.common.schemas.EndpointType.LEAF_EP)
+            == mlrun.common.schemas.EndpointType.NODE_EP
+        ):
+            self._initialize_model_logger()
+
+    def _lazy_init(self, event):
+        if event and isinstance(event, dict) and not self.initialized:
+            background_task_state = event.get("background_task_state", None)
+            if (
+                background_task_state
+                == mlrun.common.schemas.BackgroundTaskState.succeeded
+            ):
+                self._model_logger = (
+                    _ModelLogPusher(self, self.context)
+                    if self.context
+                    and self.context.stream.enabled
+                    and self.model_endpoint_uid
+                    else None
+                )
+                self.initialized = True
 
     def get_param(self, key: str, default=None):
         """get param by key (specified in the model or the function)"""
@@ -198,13 +207,15 @@ class V2ModelServer(StepToDict):
             extra dataitems dictionary
 
         """
-        model_file, self.model_spec, extra_dataitems = mlrun.artifacts.get_model(
-            self.model_path, suffix
-        )
-        if self.model_spec and self.model_spec.parameters:
-            for key, value in self.model_spec.parameters.items():
-                self._params[key] = value
-        return model_file, extra_dataitems
+        if self.model_path:
+            model_file, self.model_spec, extra_dataitems = mlrun.artifacts.get_model(
+                self.model_path, suffix
+            )
+            if self.model_spec and self.model_spec.parameters:
+                for key, value in self.model_spec.parameters.items():
+                    self._params[key] = value
+            return model_file, extra_dataitems
+        return None, None
 
     def load(self):
         """model loading function, see also .get_model() method"""
@@ -230,25 +241,10 @@ class V2ModelServer(StepToDict):
         request = self.preprocess(event_body, op)
         return self.validate(request, op)
 
-    @property
-    def versioned_model_name(self):
-        if self._versioned_model_name:
-            return self._versioned_model_name
-
-        # Generating version model value based on the model name and model version
-        if self.model_path and self.model_path.startswith("store://"):
-            # Enrich the model server with the model artifact metadata
-            self.get_model()
-            if not self.version:
-                # Enrich the model version with the model artifact tag
-                self.version = self.model_spec.tag
-            self.labels = self.model_spec.labels
-        version = self.version or "latest"
-        self._versioned_model_name = f"{self.name}:{version}"
-        return self._versioned_model_name
-
     def do_event(self, event, *args, **kwargs):
         """main model event handler method"""
+        if not self.initialized:
+            self._lazy_init(event.body)
         start = now_date()
         original_body = event.body
         event_body = _extract_input_data(self._input_path, event.body)
@@ -289,13 +285,12 @@ class V2ModelServer(StepToDict):
 
             response = {
                 "id": event_id,
-                "model_name": self.name,
+                "model_name": self.name.split(":")[0],
                 "outputs": outputs,
                 "timestamp": start.isoformat(sep=" ", timespec="microseconds"),
             }
-            if self.version:
-                response["model_version"] = self.version
-
+            if self.model_endpoint_uid:
+                response["model_endpoint_uid"] = self.model_endpoint_uid
         elif op == "ready" and event.method == "GET":
             # get model health operation
             setattr(event, "terminated", True)
@@ -320,8 +315,8 @@ class V2ModelServer(StepToDict):
             # get model metadata operation
             setattr(event, "terminated", True)
             event_body = {
-                "name": self.name,
-                "version": self.version or "",
+                "name": self.name.split(":")[0],
+                "model_endpoint_uid": self.model_endpoint_uid or "",
                 "inputs": [],
                 "outputs": [],
             }
@@ -355,8 +350,8 @@ class V2ModelServer(StepToDict):
                 "model_name": self.name,
                 "outputs": outputs,
             }
-            if self.version:
-                response["model_version"] = self.version
+            if self.model_endpoint_uid:
+                response["model_endpoint_uid"] = self.model_endpoint_uid
 
         elif hasattr(self, "op_" + op):
             # custom operation (child methods starting with "op_")
@@ -472,30 +467,70 @@ class V2ModelServer(StepToDict):
         request["inputs"] = new_inputs
         return request
 
+    def _initialize_model_logger(self):
+        server: mlrun.serving.GraphServer = getattr(
+            self.context, "_server", None
+        ) 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,
+                )
+                if (
+                    background_task.status.state
+                    in mlrun.common.schemas.BackgroundTaskState.terminal_states()
+                ):
+                    logger.debug(
+                        f"Model endpoint creation task completed with state {background_task.status.state}"
+                    )
+                    if (
+                        background_task.status.state
+                        == mlrun.common.schemas.BackgroundTaskState.succeeded
+                    ):
+                        self._model_logger = (
+                            _ModelLogPusher(self, self.context)
+                            if self.context
+                            and self.context.stream.enabled
+                            and self.model_endpoint_uid
+                            else None
+                        )
+                        self.initialized = True
+
+                else:  # in progress
+                    logger.debug(
+                        f"Model endpoint creation task is still in progress with the current state: "
+                        f"{background_task.status.state}.",
+                        name=self.name,
+                    )
+            else:
+                logger.debug(
+                    "Model endpoint creation task name not provided",
+                )
+
 
 class _ModelLogPusher:
-    def __init__(self, model, context, output_stream=None):
+    def __init__(self, model: V2ModelServer, context, output_stream=None):
         self.model = model
         self.verbose = context.verbose
         self.hostname = context.stream.hostname
         self.function_uri = context.stream.function_uri
-        self.stream_path = context.stream.stream_uri
-        self.stream_batch = int(context.get_param("log_stream_batch", 1))
-        self.stream_sample = int(context.get_param("log_stream_sample", 1))
+        self.sampling_percentage = float(context.get_param("sampling_percentage", 100))
         self.output_stream = output_stream or context.stream.output_stream
         self._worker = context.worker_id
-        self._sample_iter = 0
-        self._batch_iter = 0
-        self._batch = []
 
     def base_data(self):
         base_data = {
             "class": self.model.__class__.__name__,
             "worker": self._worker,
             "model": self.model.name,
-            "version": self.model.version,
             "host": self.hostname,
             "function_uri": self.function_uri,
+            "endpoint_id": self.model.model_endpoint_uid,
+            "sampling_percentage": self.sampling_percentage,
         }
         if getattr(self.model, "labels", None):
             base_data["labels"] = self.model.labels
@@ -515,143 +550,66 @@ class _ModelLogPusher:
             self.output_stream.push([data], partition_key=partition_key)
             return
 
-        self._sample_iter = (self._sample_iter + 1) % self.stream_sample
-        if self.output_stream and self._sample_iter == 0:
+        if self.output_stream:
+            # Ensure that the inputs are a list of lists
+            request["inputs"] = (
+                request["inputs"]
+                if not any(not isinstance(req, list) for req in request["inputs"])
+                else [request["inputs"]]
+            )
             microsec = (now_date() - start).microseconds
 
-            if self.stream_batch > 1:
-                if self._batch_iter == 0:
-                    self._batch = []
-                self._batch.append(
-                    [request, op, resp, str(start), microsec, self.model.metrics]
+            if self.sampling_percentage != 100:
+                # Randomly select a subset of the requests based on the percentage
+                num_of_inputs = len(request["inputs"])
+                sampled_requests_indices = self._pick_random_requests(
+                    num_of_inputs, self.sampling_percentage
                 )
-                self._batch_iter = (self._batch_iter + 1) % self.stream_batch
-
-                if self._batch_iter == 0:
-                    data = self.base_data()
-                    data["headers"] = [
-                        "request",
-                        "op",
-                        "resp",
-                        "when",
-                        "microsec",
-                        "metrics",
+                if not sampled_requests_indices:
+                    # No events were selected for sampling
+                    return
+
+                request["inputs"] = [
+                    request["inputs"][i] for i in sampled_requests_indices
+                ]
+
+                if resp and "outputs" in resp and isinstance(resp["outputs"], list):
+                    resp["outputs"] = [
+                        resp["outputs"][i] for i in sampled_requests_indices
                     ]
-                    data["values"] = self._batch
-                    self.output_stream.push([data], partition_key=partition_key)
-            else:
-                data = self.base_data()
-                data["request"] = request
-                data["op"] = op
-                data["resp"] = resp
-                data["when"] = start_str
-                data["microsec"] = microsec
-                if getattr(self.model, "metrics", None):
-                    data["metrics"] = self.model.metrics
-                self.output_stream.push([data], partition_key=partition_key)
-
-
-def _init_endpoint_record(
-    graph_server: GraphServer, model: V2ModelServer
-) -> Union[str, None]:
-    """
-    Initialize model endpoint record and write it into the DB. In general, this method retrieve the unique model
-    endpoint ID which is generated according to the function uri and the model version. If the model endpoint is
-    already exist in the DB, we skip the creation process. Otherwise, it writes the new model endpoint record to the DB.
-
-    :param graph_server: A GraphServer object which will be used for getting the function uri.
-    :param model:        Base model serving class (v2). It contains important details for the model endpoint record
-                         such as model name, model path, and model version.
-
-    :return: Model endpoint unique ID.
-    """
-
-    logger.info("Initializing endpoint records")
-
-    # Generate required values for the model endpoint record
-    try:
-        # Getting project name from the function uri
-        project, uri, tag, hash_key = parse_versioned_object_uri(
-            graph_server.function_uri
-        )
-    except Exception as e:
-        logger.error("Failed to parse function URI", exc=err_to_str(e))
-        return None
-
-    # Generating model endpoint ID based on function uri and model version
-    uid = mlrun.common.model_monitoring.create_model_endpoint_uid(
-        function_uri=graph_server.function_uri,
-        versioned_model=model.versioned_model_name,
-    ).uid
-
-    try:
-        model_ep = mlrun.get_run_db().get_model_endpoint(
-            project=project, endpoint_id=uid
-        )
-    except mlrun.errors.MLRunNotFoundError:
-        model_ep = None
-    except mlrun.errors.MLRunBadRequestError as err:
-        logger.info(
-            "Cannot get the model endpoints store", err=mlrun.errors.err_to_str(err)
-        )
-        return
-
-    if model.context.server.track_models and not model_ep:
-        logger.info("Creating a new model endpoint record", endpoint_id=uid)
-        model_endpoint = mlrun.common.schemas.ModelEndpoint(
-            metadata=mlrun.common.schemas.ModelEndpointMetadata(
-                project=project, labels=model.labels, uid=uid
-            ),
-            spec=mlrun.common.schemas.ModelEndpointSpec(
-                function_uri=graph_server.function_uri,
-                model=model.versioned_model_name,
-                model_class=model.__class__.__name__,
-                model_uri=model.model_path,
-                stream_path=model.context.stream.stream_uri,
-                active=True,
-                monitoring_mode=mlrun.common.schemas.model_monitoring.ModelMonitoringMode.enabled,
-            ),
-            status=mlrun.common.schemas.ModelEndpointStatus(
-                endpoint_type=mlrun.common.schemas.model_monitoring.EndpointType.NODE_EP
-            ),
-        )
+                if self.model.output_schema and len(self.model.output_schema) != len(
+                    resp["outputs"][0]
+                ):
+                    logger.info(
+                        "The number of outputs returned by the model does not match the number of outputs "
+                        "specified in the model endpoint.",
+                        model_endpoint=self.model.name,
+                        model_endpoint_id=self.model.model_endpoint_uid,
+                        output_len=len(resp["outputs"][0]),
+                        schema_len=len(self.model.output_schema),
+                    )
 
-        db = mlrun.get_run_db()
-        db.create_model_endpoint(
-            project=project,
-            endpoint_id=uid,
-            model_endpoint=model_endpoint.dict(),
-        )
+            data = self.base_data()
+            data["request"] = request
+            data["op"] = op
+            data["resp"] = resp
+            data["when"] = start_str
+            data["microsec"] = microsec
+            if getattr(self.model, "metrics", None):
+                data["metrics"] = self.model.metrics
+            data["effective_sample_count"] = len(request["inputs"])
+            self.output_stream.push([data], partition_key=partition_key)
 
-    elif model_ep:
-        attributes = {}
-        old_model_uri = model_ep.spec.model_uri
-        mlrun.model_monitoring.helpers.enrich_model_endpoint_with_model_uri(
-            model_endpoint=model_ep,
-            model_obj=model.model_spec,
-        )
-        if model_ep.spec.model_uri != old_model_uri:
-            attributes["model_uri"] = model_ep.spec.model_uri
-        if (
-            model_ep.spec.monitoring_mode
-            == mlrun.common.schemas.model_monitoring.ModelMonitoringMode.enabled
-        ) != model.context.server.track_models:
-            attributes["monitoring_mode"] = (
-                mlrun.common.schemas.model_monitoring.ModelMonitoringMode.enabled
-                if model.context.server.track_models
-                else mlrun.common.schemas.model_monitoring.ModelMonitoringMode.disabled
-            )
-        if attributes:
-            db = mlrun.get_run_db()
-            db.patch_model_endpoint(
-                project=project,
-                endpoint_id=uid,
-                attributes=attributes,
-            )
-            logger.info(
-                "Updating model endpoint attributes",
-                attributes=attributes,
-                endpoint_id=uid,
-            )
+    @staticmethod
+    def _pick_random_requests(num_of_reqs: int, percentage: float) -> list[int]:
+        """
+        Randomly selects indices of requests to sample based on the given percentage
+
+        :param num_of_reqs: Number of requests to select from
+        :param percentage: Sample percentage for each request
+        :return: A list containing the indices of the selected requests
+        """
 
-    return uid
+        return [
+            req for req in range(num_of_reqs) if random.random() < (percentage / 100)
+        ]

mlrun/track/__init__.py

@@ -11,6 +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.
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
+
 from mlrun.track.tracker import Tracker
 from mlrun.track.tracker_manager import TrackerManager

mlrun/track/tracker.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 from mlrun.artifacts import Artifact, ModelArtifact
 from mlrun.execution import MLClientCtx
@@ -63,7 +63,7 @@ class Tracker(ABC):
         project: MlrunProject,
         reference_id: Any,
         function_name: str,
-        handler: str = None,
+        handler: Optional[str] = None,
         **kwargs,
     ) -> RunObject:
         """

mlrun/track/trackers/mlflow_tracker.py

@@ -15,6 +15,7 @@ import os
 import pathlib
 import tempfile
 import zipfile
+from typing import Optional
 
 import mlflow
 import mlflow.entities
@@ -190,7 +191,7 @@ class MLFlowTracker(Tracker):
         project: MlrunProject,
         reference_id: str,
         function_name: str,
-        handler: str = None,
+        handler: Optional[str] = None,
         **kwargs,
     ) -> RunObject:
         """
@@ -251,9 +252,9 @@ class MLFlowTracker(Tracker):
         self,
         project: MlrunProject,
         reference_id: str,
-        key: str = None,
-        metrics: dict = None,
-        extra_data: dict = None,
+        key: Optional[str] = None,
+        metrics: Optional[dict] = None,
+        extra_data: Optional[dict] = None,
     ) -> ModelArtifact:
         """
         Import a model from MLFlow to MLRun.
@@ -290,7 +291,7 @@ class MLFlowTracker(Tracker):
             return model
 
     def import_artifact(
-        self, project: MlrunProject, reference_id: str, key: str = None
+        self, project: MlrunProject, reference_id: str, key: Optional[str] = None
     ) -> Artifact:
         """
         Import an artifact from MLFlow to MLRun.