mlrun 1.8.0rc10__py3-none-any.whl → 1.8.0rc13__py3-none-any.whl

mlrun/serving/states.py

@@ -25,11 +25,12 @@ import pathlib
 import traceback
 from copy import copy, deepcopy
 from inspect import getfullargspec, signature
-from typing import Any, Optional, Union
+from typing import Any, Optional, Union, cast
 
 import storey.utils
 
 import mlrun
+import mlrun.common.schemas as schemas
 
 from ..config import config
 from ..datastore import get_stream_pusher
@@ -81,22 +82,28 @@ _task_step_fields = [
     "responder",
     "input_path",
     "result_path",
+    "model_endpoint_creation_strategy",
+    "endpoint_type",
 ]
 
 
 MAX_ALLOWED_STEPS = 4500
 
 
-def new_model_endpoint(class_name, model_path, handler=None, **class_args):
-    class_args = deepcopy(class_args)
-    class_args["model_path"] = model_path
-    return TaskStep(class_name, class_args, handler=handler)
-
-
-def new_remote_endpoint(url, **class_args):
+def new_remote_endpoint(
+    url: str,
+    creation_strategy: schemas.ModelEndpointCreationStrategy,
+    endpoint_type: schemas.EndpointType,
+    **class_args,
+):
     class_args = deepcopy(class_args)
     class_args["url"] = url
-    return TaskStep("$remote", class_args)
+    return TaskStep(
+        "$remote",
+        class_args=class_args,
+        model_endpoint_creation_strategy=creation_strategy,
+        endpoint_type=endpoint_type,
+    )
 
 
 class BaseStep(ModelObj):
@@ -419,6 +426,10 @@ class TaskStep(BaseStep):
         responder: Optional[bool] = None,
         input_path: Optional[str] = None,
         result_path: Optional[str] = None,
+        model_endpoint_creation_strategy: Optional[
+            schemas.ModelEndpointCreationStrategy
+        ] = schemas.ModelEndpointCreationStrategy.INPLACE,
+        endpoint_type: Optional[schemas.EndpointType] = schemas.EndpointType.NODE_EP,
     ):
         super().__init__(name, after)
         self.class_name = class_name
@@ -438,6 +449,8 @@ class TaskStep(BaseStep):
         self.on_error = None
         self._inject_context = False
         self._call_with_event = False
+        self.model_endpoint_creation_strategy = model_endpoint_creation_strategy
+        self.endpoint_type = endpoint_type
 
     def init_object(self, context, namespace, mode="sync", reset=False, **extra_kwargs):
         self.context = context
@@ -554,7 +567,11 @@ class TaskStep(BaseStep):
 
     def _post_init(self, mode="sync"):
         if self._object and hasattr(self._object, "post_init"):
-            self._object.post_init(mode)
+            self._object.post_init(
+                mode,
+                creation_strategy=self.model_endpoint_creation_strategy,
+                endpoint_type=self.endpoint_type,
+            )
             if hasattr(self._object, "model_endpoint_uid"):
                 self.endpoint_uid = self._object.model_endpoint_uid
             if hasattr(self._object, "name"):
@@ -705,6 +722,7 @@ class RouterStep(TaskStep):
         )
         self._routes: ObjectDict = None
         self.routes = routes
+        self.endpoint_type = schemas.EndpointType.ROUTER
 
     def get_children(self):
         """get child steps (routes)"""
@@ -726,6 +744,7 @@ class RouterStep(TaskStep):
         class_name=None,
         handler=None,
         function=None,
+        creation_strategy: schemas.ModelEndpointCreationStrategy = schemas.ModelEndpointCreationStrategy.INPLACE,
         **class_args,
     ):
         """add child route step or class to the router
@@ -736,12 +755,31 @@ class RouterStep(TaskStep):
         :param class_args: class init arguments
         :param handler:    class handler to invoke on run/event
         :param function:   function this step should run in
+        :param creation_strategy: Strategy for creating or updating the model endpoint:
+            * **overwrite**:
+            1. If model endpoints with the same name exist, delete the `latest` one.
+            2. Create a new model endpoint entry and set it as `latest`.
+            * **inplace** (default):
+            1. If model endpoints with the same name exist, update the `latest` entry.
+            2. Otherwise, create a new entry.
+            * **archive**:
+            1. If model endpoints with the same name exist, preserve them.
+            2. Create a new model endpoint with the same name and set it to `latest`.
+
         """
 
         if not route and not class_name and not handler:
             raise MLRunInvalidArgumentError("route or class_name must be specified")
         if not route:
-            route = TaskStep(class_name, class_args, handler=handler)
+            route = TaskStep(
+                class_name,
+                class_args,
+                handler=handler,
+                model_endpoint_creation_strategy=creation_strategy,
+                endpoint_type=schemas.EndpointType.LEAF_EP
+                if self.class_name and "serving.VotingEnsemble" in self.class_name
+                else schemas.EndpointType.NODE_EP,
+            )
         route.function = function or route.function
 
         if len(self._routes) >= MAX_ALLOWED_STEPS:
@@ -805,6 +843,106 @@ class RouterStep(TaskStep):
         )
 
 
+class Model(storey.ParallelExecutionRunnable):
+    def load(self) -> None:
+        """Override to load model if needed."""
+        pass
+
+    def init(self):
+        self.load()
+
+    def predict(self, body: Any) -> Any:
+        """Override to implement prediction logic. If the logic requires asyncio, override predict_async() instead."""
+        return body
+
+    async def predict_async(self, body: Any) -> Any:
+        """Override to implement prediction logic if the logic requires asyncio."""
+        return body
+
+    def run(self, body: Any, path: str) -> Any:
+        return self.predict(body)
+
+    async def run_async(self, body: Any, path: str) -> Any:
+        return self.predict(body)
+
+
+class ModelSelector:
+    """Used to select which models to run on each event."""
+
+    def select(
+        self, event, available_models: list[Model]
+    ) -> Union[list[str], list[Model]]:
+        """
+        Given an event, returns a list of model names or a list of model objects to run on the event.
+        If None is returned, all models will be run.
+
+        :param event: The full event
+        :param available_models: List of available models
+        """
+        pass
+
+
+class ModelRunner(storey.ParallelExecution):
+    """
+    Runs multiple Models on each event. See ModelRunnerStep.
+
+    :param model_selector: ModelSelector instance whose select() method will be used to select models to run on each
+      event. Optional. If not passed, all models will be run.
+    """
+
+    def __init__(self, *args, model_selector: Optional[ModelSelector] = None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.model_selector = model_selector or ModelSelector()
+
+    def select_runnables(self, event):
+        models = cast(list[Model], self.runnables)
+        return self.model_selector.select(event, models)
+
+
+class ModelRunnerStep(TaskStep):
+    """
+    Runs multiple Models on each event.
+
+    example::
+
+        model_runner_step = ModelRunnerStep(name="my_model_runner")
+        model_runner_step.add_model(MyModel(name="my_model"))
+        graph.to(model_runner_step)
+
+    :param model_selector: ModelSelector instance whose select() method will be used to select models to run on each
+      event. Optional. If not passed, all models will be run.
+    """
+
+    kind = "model_runner"
+
+    def __init__(
+        self,
+        *args,
+        model_selector: Optional[Union[str, ModelSelector]] = None,
+        **kwargs,
+    ):
+        self._models = []
+        super().__init__(
+            *args,
+            class_name="mlrun.serving.ModelRunner",
+            class_args=dict(runnables=self._models, model_selector=model_selector),
+            **kwargs,
+        )
+
+    def add_model(self, model: Model) -> None:
+        """Add a Model to this ModelRunner."""
+        self._models.append(model)
+
+    def init_object(self, context, namespace, mode="sync", reset=False, **extra_kwargs):
+        model_selector = self.class_args.get("model_selector")
+        if isinstance(model_selector, str):
+            model_selector = get_class(model_selector, namespace)()
+        self._async_object = ModelRunner(
+            self.class_args.get("runnables"),
+            model_selector=model_selector,
+        )
+
+
 class QueueStep(BaseStep):
     """queue step, implement an async queue or represent a stream"""
 
@@ -1344,8 +1482,9 @@ class FlowStep(BaseStep):
 
         if self._controller:
             if hasattr(self._controller, "terminate"):
-                self._controller.terminate()
-            return self._controller.await_termination()
+                return self._controller.terminate(wait=True)
+            else:
+                return self._controller.await_termination()
 
     def plot(self, filename=None, format=None, source=None, targets=None, **kw):
         """plot/save graph using graphviz
@@ -1433,6 +1572,7 @@ classes_map = {
     "queue": QueueStep,
     "error_step": ErrorStep,
     "monitoring_application": MonitoringApplicationStep,
+    "model_runner": ModelRunnerStep,
 }
 
 
@@ -1572,6 +1712,10 @@ def params_to_step(
     input_path: Optional[str] = None,
     result_path: Optional[str] = None,
     class_args=None,
+    model_endpoint_creation_strategy: Optional[
+        schemas.ModelEndpointCreationStrategy
+    ] = None,
+    endpoint_type: Optional[schemas.EndpointType] = None,
 ):
     """return step object from provided params or classes/objects"""
 
@@ -1587,6 +1731,9 @@ def params_to_step(
         step.full_event = full_event or step.full_event
         step.input_path = input_path or step.input_path
         step.result_path = result_path or step.result_path
+        if kind == StepKinds.task:
+            step.model_endpoint_creation_strategy = model_endpoint_creation_strategy
+            step.endpoint_type = endpoint_type
 
     elif class_name and class_name in queue_class_names:
         if "path" not in class_args:
@@ -1627,6 +1774,8 @@ def params_to_step(
             full_event=full_event,
             input_path=input_path,
             result_path=result_path,
+            model_endpoint_creation_strategy=model_endpoint_creation_strategy,
+            endpoint_type=endpoint_type,
         )
     else:
         raise MLRunInvalidArgumentError("class_name or handler must be provided")

mlrun/serving/v2_serving.py

@@ -23,7 +23,6 @@ import mlrun.common.schemas.model_monitoring
 import mlrun.model_monitoring
 from mlrun.utils import logger, now_date
 
-from ..common.schemas.model_monitoring import ModelEndpointSchema
 from .server import GraphServer
 from .utils import StepToDict, _extract_input_data, _update_result_body
 
@@ -130,7 +129,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":
@@ -149,7 +148,10 @@ class V2ModelServer(StepToDict):
 
         if not self.context.is_mock or self.context.monitoring_mock:
             self.model_endpoint_uid = _init_endpoint_record(
-                graph_server=server, model=self
+                graph_server=server,
+                model=self,
+                creation_strategy=kwargs.get("creation_strategy"),
+                endpoint_type=kwargs.get("endpoint_type"),
             )
         self._model_logger = (
             _ModelLogPusher(self, self.context)
@@ -554,7 +556,10 @@ class _ModelLogPusher:
 
 
 def _init_endpoint_record(
-    graph_server: GraphServer, model: V2ModelServer
+    graph_server: GraphServer,
+    model: V2ModelServer,
+    creation_strategy: mlrun.common.schemas.ModelEndpointCreationStrategy,
+    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
@@ -564,6 +569,17 @@ def _init_endpoint_record(
     :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.
+    :param creation_strategy: Strategy for creating or updating the model endpoint:
+        * **overwrite**:
+        1. If model endpoints with the same name exist, delete the `latest` one.
+        2. Create a new model endpoint entry and set it as `latest`.
+        * **inplace** (default):
+        1. If model endpoints with the same name exist, update the `latest` entry.
+        2. Otherwise, create a new entry.
+        * **archive**:
+        1. If model endpoints with the same name exist, preserve them.
+        2. Create a new model endpoint with the same name and set it to `latest`.
+    :param endpoint_type    model endpoint type
 
     :return: Model endpoint unique ID.
     """
@@ -583,51 +599,30 @@ def _init_endpoint_record(
         model_uid = None
         model_tag = None
         model_labels = {}
-    try:
-        model_ep = mlrun.get_run_db().get_model_endpoint(
-            project=graph_server.project,
-            name=model.name,
-            function_name=graph_server.function_name,
-            function_tag=graph_server.function_tag or "latest",
-        )
-    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
-
-    function = mlrun.get_run_db().get_function(
-        name=graph_server.function_name,
+    logger.info(
+        "Creating Or Updating a new model endpoint record",
+        name=model.name,
         project=graph_server.project,
-        tag=graph_server.function_tag or "latest",
+        function_name=graph_server.function_name,
+        function_tag=graph_server.function_tag or "latest",
+        model_name=model_name,
+        model_tag=model_tag,
+        model_db_key=model_db_key,
+        model_uid=model_uid,
+        model_class=model.__class__.__name__,
+        creation_strategy=creation_strategy,
+        endpoint_type=endpoint_type,
     )
-    function_uid = function.get("metadata", {}).get("uid")
-    if not model_ep and model.context.server.track_models:
-        logger.info(
-            "Creating a new model endpoint record",
-            name=model.name,
-            project=graph_server.project,
-            function_name=graph_server.function_name,
-            function_tag=graph_server.function_tag or "latest",
-            function_uid=function_uid,
-            model_name=model_name,
-            model_tag=model_tag,
-            model_db_key=model_db_key,
-            model_uid=model_uid,
-            model_class=model.__class__.__name__,
-        )
+    try:
         model_ep = mlrun.common.schemas.ModelEndpoint(
             metadata=mlrun.common.schemas.ModelEndpointMetadata(
                 project=graph_server.project,
                 labels=model_labels,
                 name=model.name,
-                endpoint_type=mlrun.common.schemas.model_monitoring.EndpointType.NODE_EP,
+                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_name=model_name,
                 model_db_key=model_db_key,
@@ -642,49 +637,11 @@ def _init_endpoint_record(
             ),
         )
         db = mlrun.get_run_db()
-        model_ep = db.create_model_endpoint(model_endpoint=model_ep)
-
-    elif model_ep:
-        attributes = {}
-        if function_uid != model_ep.spec.function_uid:
-            attributes[ModelEndpointSchema.FUNCTION_UID] = function_uid
-        if model_name != model_ep.spec.model_name:
-            attributes[ModelEndpointSchema.MODEL_NAME] = model_name
-        if model_uid != model_ep.spec.model_uid:
-            attributes[ModelEndpointSchema.MODEL_UID] = model_uid
-        if model_tag != model_ep.spec.model_tag:
-            attributes[ModelEndpointSchema.MODEL_TAG] = model_tag
-        if model_db_key != model_ep.spec.model_db_key:
-            attributes[ModelEndpointSchema.MODEL_DB_KEY] = model_db_key
-        if model_labels != model_ep.metadata.labels:
-            attributes[ModelEndpointSchema.LABELS] = model_labels
-        if model.__class__.__name__ != model_ep.spec.model_class:
-            attributes[ModelEndpointSchema.MODEL_CLASS] = model.__class__.__name__
-        if (
-            model_ep.status.monitoring_mode
-            == mlrun.common.schemas.model_monitoring.ModelMonitoringMode.enabled
-        ) != model.context.server.track_models:
-            attributes[ModelEndpointSchema.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:
-            logger.info(
-                "Updating model endpoint attributes",
-                attributes=attributes,
-                project=model_ep.metadata.project,
-                name=model_ep.metadata.name,
-                function_name=model_ep.spec.function_name,
-            )
-            db = mlrun.get_run_db()
-            model_ep = db.patch_model_endpoint(
-                project=model_ep.metadata.project,
-                name=model_ep.metadata.name,
-                endpoint_id=model_ep.metadata.uid,
-                attributes=attributes,
-            )
-    else:
+        model_ep = db.create_model_endpoint(
+            model_endpoint=model_ep, creation_strategy=creation_strategy
+        )
+    except mlrun.errors.MLRunBadRequestError as e:
+        logger.info("Failed to create model endpoint record", error=e)
         return None
 
     return model_ep.metadata.uid

mlrun/utils/helpers.py

@@ -111,15 +111,21 @@ def get_artifact_target(item: dict, project=None):
     project_str = project or item["metadata"].get("project")
     tree = item["metadata"].get("tree")
     tag = item["metadata"].get("tag")
+    iter = item["metadata"].get("iter")
     kind = item.get("kind")
+    uid = item["metadata"].get("uid")
 
     if kind in {"dataset", "model", "artifact"} and db_key:
         target = (
             f"{DB_SCHEMA}://{StorePrefix.kind_to_prefix(kind)}/{project_str}/{db_key}"
         )
+        if iter:
+            target = f"{target}#{iter}"
         target += f":{tag}" if tag else ":latest"
         if tree:
             target += f"@{tree}"
+        if uid:
+            target += f"^{uid}"
         return target
 
     return item["spec"].get("target_path")

mlrun/utils/notifications/notification/base.py

@@ -57,7 +57,7 @@ class NotificationBase:
             typing.Union[mlrun.common.schemas.NotificationSeverity, str]
         ] = mlrun.common.schemas.NotificationSeverity.INFO,
         runs: typing.Optional[typing.Union[mlrun.lists.RunList, list]] = None,
-        custom_html: typing.Optional[typing.Optional[str]] = None,
+        custom_html: typing.Optional[str] = None,
         alert: typing.Optional[mlrun.common.schemas.AlertConfig] = None,
         event_data: typing.Optional[mlrun.common.schemas.Event] = None,
     ):

mlrun/utils/notifications/notification/webhook.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import re
 import typing
 
 import aiohttp
@@ -93,7 +94,6 @@ class WebhookNotification(NotificationBase):
 
     @staticmethod
     def _serialize_runs_in_request_body(override_body, runs):
-        str_parsed_runs = ""
         runs = runs or []
 
         def parse_runs():
@@ -105,22 +105,23 @@ class WebhookNotification(NotificationBase):
                     parsed_run = {
                         "project": run["metadata"]["project"],
                         "name": run["metadata"]["name"],
-                        "host": run["metadata"]["labels"]["host"],
                         "status": {"state": run["status"]["state"]},
                     }
-                    if run["status"].get("error", None):
-                        parsed_run["status"]["error"] = run["status"]["error"]
-                    elif run["status"].get("results", None):
-                        parsed_run["status"]["results"] = run["status"]["results"]
+                    if host := run["metadata"].get("labels", {}).get("host", ""):
+                        parsed_run["host"] = host
+                    if error := run["status"].get("error"):
+                        parsed_run["status"]["error"] = error
+                    elif results := run["status"].get("results"):
+                        parsed_run["status"]["results"] = results
                     parsed_runs.append(parsed_run)
             return str(parsed_runs)
 
         if isinstance(override_body, dict):
             for key, value in override_body.items():
-                if "{{ runs }}" or "{{runs}}" in value:
-                    if not str_parsed_runs:
-                        str_parsed_runs = parse_runs()
-                    override_body[key] = value.replace(
-                        "{{ runs }}", str_parsed_runs
-                    ).replace("{{runs}}", str_parsed_runs)
+                if re.search(r"{{\s*runs\s*}}", value):
+                    str_parsed_runs = parse_runs()
+                    override_body[key] = re.sub(
+                        r"{{\s*runs\s*}}", str_parsed_runs, value
+                    )
+
         return override_body

mlrun/utils/notifications/notification_pusher.py

@@ -14,7 +14,6 @@
 
 import asyncio
 import datetime
-import os
 import re
 import traceback
 import typing
@@ -97,6 +96,7 @@ class NotificationPusher(_NotificationPusherBase):
         "completed": "{resource} completed",
         "error": "{resource} failed",
         "aborted": "{resource} aborted",
+        "running": "{resource} started",
     }
 
     def __init__(
@@ -285,6 +285,7 @@ class NotificationPusher(_NotificationPusherBase):
 
         message = (
             self.messages.get(run.state(), "").format(resource=resource)
+            + f" in project {run.metadata.project}"
             + custom_message
         )
 
@@ -303,6 +304,7 @@ class NotificationPusher(_NotificationPusherBase):
         message, severity, runs = self._prepare_notification_args(
             run, notification_object
         )
+
         logger.debug(
             "Pushing sync notification",
             notification=sanitize_notification(notification_object.to_dict()),
@@ -313,6 +315,7 @@ class NotificationPusher(_NotificationPusherBase):
             "project": run.metadata.project,
             "notification": notification_object,
             "status": mlrun.common.schemas.NotificationStatus.SENT,
+            "run_state": run.state(),
         }
         try:
             notification.push(message, severity, runs)
@@ -351,6 +354,7 @@ class NotificationPusher(_NotificationPusherBase):
         message, severity, runs = self._prepare_notification_args(
             run, notification_object
         )
+
         logger.debug(
             "Pushing async notification",
             notification=sanitize_notification(notification_object.to_dict()),
@@ -360,6 +364,7 @@ class NotificationPusher(_NotificationPusherBase):
             "run_uid": run.metadata.uid,
             "project": run.metadata.project,
             "notification": notification_object,
+            "run_state": run.state(),
             "status": mlrun.common.schemas.NotificationStatus.SENT,
         }
         try:
@@ -397,10 +402,20 @@ class NotificationPusher(_NotificationPusherBase):
         run_uid: str,
         project: str,
         notification: mlrun.model.Notification,
+        run_state: runtimes_constants.RunStates,
         status: typing.Optional[str] = None,
         sent_time: typing.Optional[datetime.datetime] = None,
         reason: typing.Optional[str] = None,
     ):
+        if run_state not in runtimes_constants.RunStates.terminal_states():
+            # we want to update the notification status only if the run is in a terminal state for BC
+            logger.debug(
+                "Skip updating notification status - run not in terminal state",
+                run_uid=run_uid,
+                state=run_state,
+            )
+            return
+
         db = mlrun.get_run_db()
         notification.status = status or notification.status
         notification.sent_time = sent_time or notification.sent_time
@@ -664,30 +679,10 @@ class CustomNotificationPusher(_NotificationPusherBase):
     def push_pipeline_start_message(
         self,
         project: str,
-        commit_id: typing.Optional[str] = None,
         pipeline_id: typing.Optional[str] = None,
-        has_workflow_url: bool = False,
     ):
-        message = f"Workflow started in project {project}"
-        if pipeline_id:
-            message += f" id={pipeline_id}"
-        commit_id = (
-            commit_id or os.environ.get("GITHUB_SHA") or os.environ.get("CI_COMMIT_SHA")
-        )
-        if commit_id:
-            message += f", commit={commit_id}"
-        if has_workflow_url:
-            url = mlrun.utils.helpers.get_workflow_url(project, pipeline_id)
-        else:
-            url = mlrun.utils.helpers.get_ui_url(project)
-        html = ""
-        if url:
-            html = (
-                message
-                + f'<div><a href="{url}" target="_blank">click here to view progress</a></div>'
-            )
-            message = message + f", check progress in {url}"
-        self.push(message, "info", custom_html=html)
+        db = mlrun.get_run_db()
+        db.push_run_notifications(pipeline_id, project)
 
     def push_pipeline_run_results(
         self,

mlrun/utils/version/version.json

@@ -1,4 +1,4 @@
 {
-  "git_commit": "9e2d1e195daed90072d1cd33a5eee339577dc35a",
-  "version": "1.8.0-rc10"
+  "git_commit": "6b23c9dc1a77d33bf56dd6e91623278630bb46c4",
+  "version": "1.8.0-rc13"
 }