mlrun 1.8.0rc9__py3-none-any.whl → 1.8.0rc12__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,23 @@ 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: 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.
         """
 
         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.NODE_EP,
+            )
         route.function = function or route.function
 
         if len(self._routes) >= MAX_ALLOWED_STEPS:
@@ -805,6 +835,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 +1474,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 +1564,7 @@ classes_map = {
     "queue": QueueStep,
     "error_step": ErrorStep,
     "monitoring_application": MonitoringApplicationStep,
+    "model_runner": ModelRunnerStep,
 }
 
 
@@ -1572,6 +1704,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 +1723,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 +1766,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: 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
@@ -564,6 +569,12 @@ 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: 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.
     """
@@ -573,62 +584,46 @@ def _init_endpoint_record(
         model.get_model()
     if model.model_spec:
         model_name = model.model_spec.metadata.key
+        model_db_key = model.model_spec.spec.db_key
         model_uid = model.model_spec.metadata.uid
         model_tag = model.model_spec.tag
         model_labels = model.model_spec.labels  # todo : check if we still need this
     else:
         model_name = None
+        model_db_key = None
         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_uid=model_uid,
-            model_class=model.__class__.__name__,
-            model_tag=model_tag,
-        )
+    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,
                 model_uid=model_uid,
                 model_class=model.__class__.__name__,
+                model_tag=model_tag,
             ),
             status=mlrun.common.schemas.ModelEndpointStatus(
                 monitoring_mode=mlrun.common.schemas.model_monitoring.ModelMonitoringMode.enabled
@@ -637,47 +632,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_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/version/version.json

@@ -1,4 +1,4 @@
 {
-  "git_commit": "5be219a3a066786e96ee697c208cd0f59ef63362",
-  "version": "1.8.0-rc9"
+  "git_commit": "e190e91298c56dce1ac5f5cab61ad39495936fe3",
+  "version": "1.8.0-rc12"
 }

{mlrun-1.8.0rc9.dist-info → mlrun-1.8.0rc12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mlrun
-Version: 1.8.0rc9
+Version: 1.8.0rc12
 Summary: Tracking and config of machine learning runs
 Home-page: https://github.com/mlrun/mlrun
 Author: Yaron Haviv
@@ -44,7 +44,7 @@ Requires-Dist: semver~=3.0
 Requires-Dist: dependency-injector~=4.41
 Requires-Dist: fsspec<2024.7,>=2023.9.2
 Requires-Dist: v3iofs~=0.1.17
-Requires-Dist: storey~=1.8.0
+Requires-Dist: storey~=1.8.6
 Requires-Dist: inflection~=0.5.0
 Requires-Dist: python-dotenv~=1.0
 Requires-Dist: setuptools>=75.2
@@ -103,16 +103,16 @@ Requires-Dist: taoswswrap~=0.2.0; extra == "tdengine"
 Provides-Extra: snowflake
 Requires-Dist: snowflake-connector-python~=3.7; extra == "snowflake"
 Provides-Extra: api
-Requires-Dist: uvicorn~=0.27.1; extra == "api"
+Requires-Dist: uvicorn~=0.32.1; extra == "api"
 Requires-Dist: dask-kubernetes~=0.11.0; extra == "api"
 Requires-Dist: apscheduler<4,>=3.11; extra == "api"
 Requires-Dist: objgraph~=3.6; extra == "api"
 Requires-Dist: igz-mgmt~=0.4.1; extra == "api"
 Requires-Dist: humanfriendly~=10.0; extra == "api"
-Requires-Dist: fastapi~=0.110.0; extra == "api"
+Requires-Dist: fastapi~=0.115.6; extra == "api"
 Requires-Dist: sqlalchemy~=1.4; extra == "api"
-Requires-Dist: pymysql~=1.0; extra == "api"
-Requires-Dist: alembic~=1.9; extra == "api"
+Requires-Dist: pymysql~=1.1; extra == "api"
+Requires-Dist: alembic~=1.14; extra == "api"
 Requires-Dist: timelength~=1.1; extra == "api"
 Requires-Dist: memray~=1.12; sys_platform != "win32" and extra == "api"
 Requires-Dist: aiosmtplib~=3.0; extra == "api"
@@ -181,7 +181,7 @@ Provides-Extra: complete-api
 Requires-Dist: adlfs==2023.9.0; extra == "complete-api"
 Requires-Dist: aiobotocore<2.16,>=2.5.0; extra == "complete-api"
 Requires-Dist: aiosmtplib~=3.0; extra == "complete-api"
-Requires-Dist: alembic~=1.9; extra == "complete-api"
+Requires-Dist: alembic~=1.14; extra == "complete-api"
 Requires-Dist: apscheduler<4,>=3.11; extra == "complete-api"
 Requires-Dist: avro~=1.11; extra == "complete-api"
 Requires-Dist: azure-core~=1.24; extra == "complete-api"
@@ -192,7 +192,7 @@ Requires-Dist: dask-kubernetes~=0.11.0; extra == "complete-api"
 Requires-Dist: dask~=2023.12.1; extra == "complete-api"
 Requires-Dist: databricks-sdk~=0.13.0; extra == "complete-api"
 Requires-Dist: distributed~=2023.12.1; extra == "complete-api"
-Requires-Dist: fastapi~=0.110.0; extra == "complete-api"
+Requires-Dist: fastapi~=0.115.6; extra == "complete-api"
 Requires-Dist: gcsfs<2024.7,>=2023.9.2; extra == "complete-api"
 Requires-Dist: google-cloud-bigquery-storage~=2.17; extra == "complete-api"
 Requires-Dist: google-cloud-bigquery[bqstorage,pandas]==3.14.1; extra == "complete-api"
@@ -209,7 +209,7 @@ Requires-Dist: objgraph~=3.6; extra == "complete-api"
 Requires-Dist: oss2==2.18.1; extra == "complete-api"
 Requires-Dist: ossfs==2023.12.0; extra == "complete-api"
 Requires-Dist: plotly~=5.23; extra == "complete-api"
-Requires-Dist: pymysql~=1.0; extra == "complete-api"
+Requires-Dist: pymysql~=1.1; extra == "complete-api"
 Requires-Dist: pyopenssl>=23; extra == "complete-api"
 Requires-Dist: redis~=4.3; extra == "complete-api"
 Requires-Dist: s3fs<2024.7,>=2023.9.2; extra == "complete-api"
@@ -218,7 +218,7 @@ Requires-Dist: sqlalchemy~=1.4; extra == "complete-api"
 Requires-Dist: taos-ws-py==0.3.2; extra == "complete-api"
 Requires-Dist: taoswswrap~=0.2.0; extra == "complete-api"
 Requires-Dist: timelength~=1.1; extra == "complete-api"
-Requires-Dist: uvicorn~=0.27.1; extra == "complete-api"
+Requires-Dist: uvicorn~=0.32.1; extra == "complete-api"
 
 <a id="top"></a>
 [![Build Status](https://github.com/mlrun/mlrun/actions/workflows/build.yaml/badge.svg?branch=development)](https://github.com/mlrun/mlrun/actions/workflows/build.yaml?query=branch%3Adevelopment)