mlrun 1.8.0rc9__py3-none-any.whl → 1.8.0rc12__py3-none-any.whl

mlrun/projects/project.py

@@ -2117,13 +2117,13 @@ class MlrunProject(ModelObj):
 
     def set_model_monitoring_function(
         self,
+        name: str,
         func: typing.Union[str, mlrun.runtimes.RemoteRuntime, None] = None,
         application_class: typing.Union[
             str, mm_app.ModelMonitoringApplicationBase, None
         ] = None,
-        name: Optional[str] = None,
         image: Optional[str] = None,
-        handler=None,
+        handler: Optional[str] = None,
         with_repo: Optional[bool] = None,
         tag: Optional[str] = None,
         requirements: Optional[typing.Union[str, list[str]]] = None,
@@ -2135,7 +2135,7 @@ class MlrunProject(ModelObj):
         Note: to deploy the function after linking it to the project,
         call `fn.deploy()` where `fn` is the object returned by this method.
 
-        examples::
+        Example::
 
             project.set_model_monitoring_function(
                 name="myApp", application_class="MyApp", image="mlrun/mlrun"
@@ -2144,8 +2144,7 @@ class MlrunProject(ModelObj):
         :param func:                    Remote function object or spec/code URL. :code:`None` refers to the current
                                         notebook.
         :param name:                    Name of the function (under the project), can be specified with a tag to support
-                                        versions (e.g. myfunc:v1)
-                                        Default: job
+                                        versions (e.g. myfunc:v1).
         :param image:                   Docker image to be used, can also be specified in
                                         the function object/yaml
         :param handler:                 Default function handler to invoke (can only be set with .py/.ipynb files)
@@ -2183,12 +2182,13 @@ class MlrunProject(ModelObj):
 
     def create_model_monitoring_function(
         self,
+        name: str,
         func: Optional[str] = None,
         application_class: typing.Union[
             str,
             mm_app.ModelMonitoringApplicationBase,
+            None,
         ] = None,
-        name: Optional[str] = None,
         image: Optional[str] = None,
         handler: Optional[str] = None,
         with_repo: Optional[bool] = None,
@@ -2200,16 +2200,15 @@ class MlrunProject(ModelObj):
         """
         Create a monitoring function object without setting it to the project
 
-        examples::
+        Example::
 
             project.create_model_monitoring_function(
-                application_class_name="MyApp", image="mlrun/mlrun", name="myApp"
+                name="myApp", application_class="MyApp", image="mlrun/mlrun"
             )
 
         :param func:                    The function's code URL. :code:`None` refers to the current notebook.
         :param name:                    Name of the function, can be specified with a tag to support
-                                        versions (e.g. myfunc:v1)
-                                        Default: job
+                                        versions (e.g. myfunc:v1).
         :param image:                   Docker image to be used, can also be specified in
                                         the function object/yaml
         :param handler:                 Default function handler to invoke (can only be set with .py/.ipynb files)
@@ -3551,6 +3550,7 @@ class MlrunProject(ModelObj):
         self,
         name: Optional[str] = None,
         model_name: Optional[str] = None,
+        model_tag: Optional[str] = None,
         function_name: Optional[str] = None,
         function_tag: Optional[str] = None,
         labels: Optional[list[str]] = None,
@@ -3565,12 +3565,13 @@ class MlrunProject(ModelObj):
         model endpoint. This functions supports filtering by the following parameters:
         1) name
         2) model_name
-        3) function_name
-        4) function_tag
-        5) labels
-        6) top level
-        7) uids
-        8) start and end time, corresponding to the `created` field.
+        3) model_tag
+        4) function_name
+        5) function_tag
+        6) labels
+        7) top level
+        8) uids
+        9) start and end time, corresponding to the `created` field.
         By default, when no filters are applied, all available endpoints for the given project will be listed.
 
         In addition, this functions provides a facade for listing endpoint related metrics. This facade is time-based
@@ -3599,6 +3600,7 @@ class MlrunProject(ModelObj):
             project=self.name,
             name=name,
             model_name=model_name,
+            model_tag=model_tag,
             function_name=function_name,
             function_tag=function_tag,
             labels=labels,

mlrun/runtimes/nuclio/serving.py

@@ -22,7 +22,7 @@ import nuclio
 from nuclio import KafkaTrigger
 
 import mlrun
-import mlrun.common.schemas
+import mlrun.common.schemas as schemas
 from mlrun.datastore import get_kafka_brokers_from_dict, parse_kafka_url
 from mlrun.model import ObjectList
 from mlrun.runtimes.function_reference import FunctionReference
@@ -362,6 +362,9 @@ class ServingRuntime(RemoteRuntime):
         handler: Optional[str] = None,
         router_step: Optional[str] = None,
         child_function: Optional[str] = None,
+        creation_strategy: Optional[
+            schemas.ModelEndpointCreationStrategy
+        ] = schemas.ModelEndpointCreationStrategy.INPLACE,
         **class_args,
     ):
         """add ml model and/or route to the function.
@@ -384,6 +387,11 @@ class ServingRuntime(RemoteRuntime):
         :param router_step: router step name (to determine which router we add the model to in graphs
                             with multiple router steps)
         :param child_function: child function name, when the model runs in a child function
+        :param creation_strategy: model endpoint creation strategy :
+                                    * overwrite - Create a new model endpoint and delete the last old one if it exists.
+                                    * inplace - Use the existing model endpoint if it already exists (default).
+                                    * archive - Preserve the old model endpoint and create a new one,
+                                    tagging it as the latest.
         :param class_args:  extra kwargs to pass to the model serving class __init__
                             (can be read in the model using .get_param(key) method)
         """
@@ -419,7 +427,12 @@ class ServingRuntime(RemoteRuntime):
         if class_name and hasattr(class_name, "to_dict"):
             if model_path:
                 class_name.model_path = model_path
-            key, state = params_to_step(class_name, key)
+            key, state = params_to_step(
+                class_name,
+                key,
+                model_endpoint_creation_strategy=creation_strategy,
+                endpoint_type=schemas.EndpointType.LEAF_EP,
+            )
         else:
             class_name = class_name or self.spec.default_class
             if class_name and not isinstance(class_name, str):
@@ -432,12 +445,22 @@ class ServingRuntime(RemoteRuntime):
                 model_path = str(model_path)
 
             if model_url:
-                state = new_remote_endpoint(model_url, **class_args)
+                state = new_remote_endpoint(
+                    model_url,
+                    creation_strategy=creation_strategy,
+                    endpoint_type=schemas.EndpointType.LEAF_EP,
+                    **class_args,
+                )
             else:
                 class_args = deepcopy(class_args)
                 class_args["model_path"] = model_path
                 state = TaskStep(
-                    class_name, class_args, handler=handler, function=child_function
+                    class_name,
+                    class_args,
+                    handler=handler,
+                    function=child_function,
+                    model_endpoint_creation_strategy=creation_strategy,
+                    endpoint_type=schemas.EndpointType.LEAF_EP,
                 )
 
         return graph.add_route(key, state)
@@ -581,7 +604,7 @@ class ServingRuntime(RemoteRuntime):
         project="",
         tag="",
         verbose=False,
-        auth_info: mlrun.common.schemas.AuthInfo = None,
+        auth_info: schemas.AuthInfo = None,
         builder_env: Optional[dict] = None,
         force_build: bool = False,
     ):

mlrun/serving/__init__.py

@@ -23,6 +23,10 @@ __all__ = [
     "QueueStep",
     "ErrorStep",
     "MonitoringApplicationStep",
+    "ModelRunnerStep",
+    "ModelRunner",
+    "Model",
+    "ModelSelector",
 ]
 
 from .routers import ModelRouter, VotingEnsemble  # noqa
@@ -33,6 +37,10 @@ from .states import (
     RouterStep,
     TaskStep,
     MonitoringApplicationStep,
+    ModelRunnerStep,
+    ModelRunner,
+    Model,
+    ModelSelector,
 )  # noqa
 from .v1_serving import MLModelServer, new_v1_model_server  # noqa
 from .v2_serving import V2ModelServer  # noqa

mlrun/serving/merger.py

@@ -74,7 +74,7 @@ class Merge(storey.Flow):
         self._queue_len = max_behind or 64  # default queue is 64 entries
         self._keys_queue = []
 
-    def post_init(self, mode="sync"):
+    def post_init(self, mode="sync", **kwargs):
         # auto detect number of uplinks or use user specified value
         self._uplinks = self.expected_num_events or (
             len(self._graph_step.after) if self._graph_step else 0

mlrun/serving/remote.py

@@ -14,6 +14,7 @@
 #
 import asyncio
 import json
+from copy import copy
 from typing import Optional
 
 import aiohttp
@@ -53,6 +54,7 @@ class RemoteStep(storey.SendToHttp):
         retries=None,
         backoff_factor=None,
         timeout=None,
+        headers_expression: Optional[str] = None,
         **kwargs,
     ):
         """class for calling remote endpoints
@@ -86,6 +88,7 @@ class RemoteStep(storey.SendToHttp):
         :param retries:     number of retries (in exponential backoff)
         :param backoff_factor: A backoff factor in seconds to apply between attempts after the second try
         :param timeout:     How long to wait for the server to send data before giving up, float in seconds
+        :param headers_expression: an expression for getting the request headers from the event, e.g. "event['headers']"
         """
         # init retry args for storey
         retries = default_retries if retries is None else retries
@@ -102,6 +105,7 @@ class RemoteStep(storey.SendToHttp):
         self.url = url
         self.url_expression = url_expression
         self.body_expression = body_expression
+        self.headers_expression = headers_expression
         self.headers = headers
         self.method = method
         self.return_json = return_json
@@ -114,8 +118,9 @@ class RemoteStep(storey.SendToHttp):
         self._session = None
         self._url_function_handler = None
         self._body_function_handler = None
+        self._headers_function_handler = None
 
-    def post_init(self, mode="sync"):
+    def post_init(self, mode="sync", **kwargs):
         self._endpoint = self.url
         if self.url and self.context:
             self._endpoint = self.context.get_remote_endpoint(self.url).strip("/")
@@ -131,6 +136,12 @@ class RemoteStep(storey.SendToHttp):
                 {"endpoint": self._endpoint, "context": self.context},
                 {},
             )
+        if self.headers_expression:
+            self._headers_function_handler = eval(
+                "lambda event: " + self.headers_expression,
+                {"context": self.context},
+                {},
+            )
         elif self.subpath:
             self._append_event_path = self.subpath == "$path"
             if not self._append_event_path:
@@ -205,7 +216,10 @@ class RemoteStep(storey.SendToHttp):
 
     def _generate_request(self, event, body):
         method = self.method or event.method or "POST"
-        headers = self.headers or {}
+        if self._headers_function_handler:
+            headers = self._headers_function_handler(body)
+        else:
+            headers = copy(self.headers) or {}
 
         if self._url_function_handler:
             url = self._url_function_handler(body)
@@ -216,10 +230,8 @@ class RemoteStep(storey.SendToHttp):
                 url = url + "/" + striped_path
             if striped_path:
                 headers[event_path_key] = event.path
-
         if event.id:
             headers[event_id_key] = event.id
-
         if method == "GET":
             body = None
         elif body is not None and not isinstance(body, (str, bytes)):
@@ -334,7 +346,7 @@ class BatchHttpRequests(_ConcurrentJobExecution):
     async def _cleanup(self):
         await self._client_session.close()
 
-    def post_init(self, mode="sync"):
+    def post_init(self, mode="sync", **kwargs):
         self._endpoint = self.url
         if self.url and self.context:
             self._endpoint = self.context.get_remote_endpoint(self.url).strip("/")

mlrun/serving/routers.py

@@ -30,7 +30,6 @@ import mlrun.common.model_monitoring
 import mlrun.common.schemas.model_monitoring
 from mlrun.utils import logger, now_date
 
-from ..common.schemas.model_monitoring import ModelEndpointSchema
 from .server import GraphServer
 from .utils import RouterToDict, _extract_input_data, _update_result_body
 from .v2_serving import _ModelLogPusher
@@ -110,7 +109,7 @@ class BaseModelRouter(RouterToDict):
 
         return parsed_event
 
-    def post_init(self, mode="sync"):
+    def post_init(self, mode="sync", **kwargs):
         self.context.logger.info(f"Loaded {list(self.routes.keys())}")
 
     def get_metadata(self):
@@ -610,7 +609,7 @@ class VotingEnsemble(ParallelRun):
         self.model_endpoint_uid = None
         self.shard_by_endpoint = shard_by_endpoint
 
-    def post_init(self, mode="sync"):
+    def post_init(self, mode="sync", **kwargs):
         server = getattr(self.context, "_server", None) or getattr(
             self.context, "server", None
         )
@@ -619,7 +618,9 @@ class VotingEnsemble(ParallelRun):
             return
 
         if not self.context.is_mock or self.context.monitoring_mock:
-            self.model_endpoint_uid = _init_endpoint_record(server, self)
+            self.model_endpoint_uid = _init_endpoint_record(
+                server, self, creation_strategy=kwargs.get("creation_strategy")
+            )
 
         self._update_weights(self.weights)
 
@@ -1001,7 +1002,10 @@ class VotingEnsemble(ParallelRun):
 
 
 def _init_endpoint_record(
-    graph_server: GraphServer, voting_ensemble: VotingEnsemble
+    graph_server: GraphServer,
+    voting_ensemble: VotingEnsemble,
+    creation_strategy: str,
+    endpoint_type: mlrun.common.schemas.EndpointType,
 ) -> Union[str, None]:
     """
     Initialize model endpoint record and write it into the DB. In general, this method retrieve the unique model
@@ -1011,61 +1015,44 @@ def _init_endpoint_record(
     :param graph_server:    A GraphServer object which will be used for getting the function uri.
     :param voting_ensemble: Voting ensemble serving class. It contains important details for the model endpoint record
                             such as model name, model path, model version, and the ids of the children model endpoints.
-
+    :param creation_strategy: model endpoint creation strategy :
+                                * overwrite - Create a new model endpoint and delete the last old one if it exists.
+                                * inplace - Use the existing model endpoint if it already exists (default).
+                                * archive - Preserve the old model endpoint and create a new one,
+                                    tagging it as the latest.
+    :param endpoint_type:    model endpoint type
     :return: Model endpoint unique ID.
     """
 
     logger.info("Initializing endpoint records")
-    try:
-        model_endpoint = mlrun.get_run_db().get_model_endpoint(
-            project=graph_server.project,
-            name=voting_ensemble.name,
-            function_name=graph_server.function_name,
-            function_tag=graph_server.function_tag or "latest",
-        )
-    except mlrun.errors.MLRunNotFoundError:
-        model_endpoint = None
-    except mlrun.errors.MLRunBadRequestError as err:
-        logger.info(
-            "Cannot get the model endpoints store", err=mlrun.errors.err_to_str(err)
-        )
-        return
-
-    function = mlrun.get_run_db().get_function(
-        name=graph_server.function_name,
-        project=graph_server.project,
-        tag=graph_server.function_tag or "latest",
-    )
-    function_uid = function.get("metadata", {}).get("uid")
-    # Get the children model endpoints ids
     children_uids = []
     children_names = []
     for _, c in voting_ensemble.routes.items():
         if hasattr(c, "endpoint_uid"):
             children_uids.append(c.endpoint_uid)
             children_names.append(c.name)
-    if not model_endpoint and voting_ensemble.context.server.track_models:
+    try:
         logger.info(
-            "Creating a new model endpoint record",
+            "Creating Or Updating a new model endpoint record",
             name=voting_ensemble.name,
             project=graph_server.project,
             function_name=graph_server.function_name,
             function_tag=graph_server.function_tag or "latest",
-            function_uid=function_uid,
             model_class=voting_ensemble.__class__.__name__,
+            creation_strategy=creation_strategy,
         )
         model_endpoint = mlrun.common.schemas.ModelEndpoint(
             metadata=mlrun.common.schemas.ModelEndpointMetadata(
                 project=graph_server.project,
                 name=voting_ensemble.name,
-                endpoint_type=mlrun.common.schemas.model_monitoring.EndpointType.ROUTER,
+                endpoint_type=endpoint_type,
             ),
             spec=mlrun.common.schemas.ModelEndpointSpec(
                 function_name=graph_server.function_name,
-                function_uid=function_uid,
                 function_tag=graph_server.function_tag or "latest",
                 model_class=voting_ensemble.__class__.__name__,
-                children_uids=list(voting_ensemble.routes.keys()),
+                children_uids=children_uids,
+                children=children_names,
             ),
             status=mlrun.common.schemas.ModelEndpointStatus(
                 monitoring_mode=mlrun.common.schemas.model_monitoring.ModelMonitoringMode.enabled
@@ -1074,59 +1061,12 @@ def _init_endpoint_record(
             ),
         )
         db = mlrun.get_run_db()
-        db.create_model_endpoint(model_endpoint=model_endpoint)
-
-    elif model_endpoint:
-        attributes = {}
-        if function_uid != model_endpoint.spec.function_uid:
-            attributes[ModelEndpointSchema.FUNCTION_UID] = function_uid
-        if children_uids != model_endpoint.spec.children_uids:
-            attributes[ModelEndpointSchema.CHILDREN_UIDS] = children_uids
-        if (
-            model_endpoint.status.monitoring_mode
-            == mlrun.common.schemas.model_monitoring.ModelMonitoringMode.enabled
-        ) != voting_ensemble.context.server.track_models:
-            attributes[ModelEndpointSchema.MONITORING_MODE] = (
-                mlrun.common.schemas.model_monitoring.ModelMonitoringMode.enabled
-                if voting_ensemble.context.server.track_models
-                else mlrun.common.schemas.model_monitoring.ModelMonitoringMode.disabled
-            )
-        if attributes:
-            db = mlrun.get_run_db()
-            logger.info(
-                "Updating model endpoint attributes",
-                attributes=attributes,
-                project=model_endpoint.metadata.project,
-                name=model_endpoint.metadata.name,
-                function_name=model_endpoint.spec.function_name,
-            )
-            model_endpoint = db.patch_model_endpoint(
-                project=model_endpoint.metadata.project,
-                name=model_endpoint.metadata.name,
-                endpoint_id=model_endpoint.metadata.uid,
-                attributes=attributes,
-            )
-    else:
-        logger.info(
-            "Did not create a new model endpoint record, monitoring is disabled"
+        db.create_model_endpoint(
+            model_endpoint=model_endpoint, creation_strategy=creation_strategy
         )
+    except mlrun.errors.MLRunInvalidArgumentError as e:
+        logger.info("Failed to create model endpoint record", error=e)
         return None
-
-    # Update model endpoint children type
-    logger.info(
-        "Updating children model endpoint type",
-        children_uids=children_uids,
-        children_names=children_names,
-    )
-    for uid, name in zip(children_uids, children_names):
-        mlrun.get_run_db().patch_model_endpoint(
-            name=name,
-            project=graph_server.project,
-            endpoint_id=uid,
-            attributes={
-                ModelEndpointSchema.ENDPOINT_TYPE: mlrun.common.schemas.model_monitoring.EndpointType.LEAF_EP
-            },
-        )
     return model_endpoint.metadata.uid
 
 
@@ -1192,7 +1132,7 @@ class EnrichmentModelRouter(ModelRouter):
 
         self._feature_service = None
 
-    def post_init(self, mode="sync"):
+    def post_init(self, mode="sync", **kwargs):
         from ..feature_store import get_feature_vector
 
         super().post_init(mode)
@@ -1342,7 +1282,7 @@ class EnrichmentVotingEnsemble(VotingEnsemble):
 
         self._feature_service = None
 
-    def post_init(self, mode="sync"):
+    def post_init(self, mode="sync", **kwargs):
         from ..feature_store import get_feature_vector
 
         super().post_init(mode)

mlrun/serving/server.py

@@ -367,7 +367,9 @@ def _set_callbacks(server, context):
 
         async def termination_callback():
             context.logger.info("Termination callback called")
-            server.wait_for_completion()
+            maybe_coroutine = server.wait_for_completion()
+            if asyncio.iscoroutine(maybe_coroutine):
+                await maybe_coroutine
             context.logger.info("Termination of async flow is completed")
 
         context.platform.set_termination_callback(termination_callback)
@@ -379,7 +381,9 @@ def _set_callbacks(server, context):
 
         async def drain_callback():
             context.logger.info("Drain callback called")
-            server.wait_for_completion()
+            maybe_coroutine = server.wait_for_completion()
+            if asyncio.iscoroutine(maybe_coroutine):
+                await maybe_coroutine
             context.logger.info(
                 "Termination of async flow is completed. Rerunning async flow."
             )