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

mlrun/projects/project.py

@@ -1873,6 +1873,34 @@ class MlrunProject(ModelObj):
         vector_store: "VectorStore",  # noqa: F821
         collection_name: Optional[str] = None,
     ) -> VectorStoreCollection:
+        """
+        Create a VectorStoreCollection wrapper for a given vector store instance.
+
+        This method wraps a vector store implementation (like Milvus, Chroma) with MLRun
+        integration capabilities. The wrapper provides access to the underlying vector
+        store's functionality while adding MLRun-specific features like document and
+        artifact management.
+
+        Args:
+            vector_store: The vector store instance to wrap (e.g., Milvus, Chroma).
+                        This is the underlying implementation that will handle
+                        vector storage and retrieval.
+            collection_name: Optional name for the collection. If not provided,
+                            will attempt to extract it from the vector_store object
+                            by looking for 'collection_name', '_collection_name',
+                            'index_name', or '_index_name' attributes.
+
+        Returns:
+            VectorStoreCollection: A wrapped vector store instance with MLRun integration.
+                                This wrapper provides both access to the original vector
+                                store's capabilities and additional MLRun functionality.
+
+        Example:
+            >>> vector_store = Chroma(embedding_function=embeddings)
+            >>> collection = project.get_vector_store_collection(
+            ...     vector_store, collection_name="my_collection"
+            ... )
+        """
         return VectorStoreCollection(
             self,
             vector_store,
@@ -1899,12 +1927,39 @@ class MlrunProject(ModelObj):
         :param local_path:    path to the local file we upload, will also be use
                               as the destination subpath (under "artifact_path")
         :param artifact_path: Target path for artifact storage
-        :param document_loader_spec: Spec to use to load the artifact as langchain document
+        :param document_loader_spec: Spec to use to load the artifact as langchain document.
+
+            By default, uses DocumentLoaderSpec() which initializes with:
+
+            * loader_class_name="langchain_community.document_loaders.TextLoader"
+            * src_name="file_path"
+            * kwargs=None
+
+            Can be customized for different document types, e.g.::
+
+                DocumentLoaderSpec(
+                    loader_class_name="langchain_community.document_loaders.PDFLoader",
+                    src_name="file_path",
+                    kwargs={"extract_images": True}
+                )
         :param upload: Whether to upload the artifact
         :param labels: Key-value labels
         :param target_path: Target file path
         :param kwargs: Additional keyword arguments
         :return: DocumentArtifact object
+
+        Example:
+            >>> # Log a PDF document with custom loader
+            >>> project.log_document(
+            ...     key="my_doc",
+            ...     local_path="path/to/doc.pdf",
+            ...     document_loader=DocumentLoaderSpec(
+            ...         loader_class_name="langchain_community.document_loaders.PDFLoader",
+            ...         src_name="file_path",
+            ...         kwargs={"extract_images": True},
+            ...     ),
+            ... )
+
         """
         doc_artifact = DocumentArtifact(
             key=key,
@@ -2117,13 +2172,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 +2190,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 +2199,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 +2237,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 +2255,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)
@@ -2587,6 +2641,24 @@ class MlrunProject(ModelObj):
         self._set_function(resolved_function_name, tag, function_object, func)
         return function_object
 
+    def push_run_notifications(
+        self,
+        uid,
+        timeout=45,
+    ):
+        """
+        Push notifications for a run.
+
+        :param uid: Unique ID of the run.
+        :returns: :py:class:`~mlrun.common.schemas.BackgroundTask`.
+        """
+        db = mlrun.db.get_run_db(secrets=self._secrets)
+        return db.push_run_notifications(
+            project=self.name,
+            uid=uid,
+            timeout=timeout,
+        )
+
     def _instantiate_function(
         self,
         func: typing.Union[str, mlrun.runtimes.BaseRuntime] = None,
@@ -3240,6 +3312,7 @@ class MlrunProject(ModelObj):
         cleanup_ttl: Optional[int] = None,
         notifications: Optional[list[mlrun.model.Notification]] = None,
         workflow_runner_node_selector: typing.Optional[dict[str, str]] = None,
+        context: typing.Optional[mlrun.execution.MLClientCtx] = None,
     ) -> _PipelineRunStatus:
         """Run a workflow using kubeflow pipelines
 
@@ -3282,6 +3355,7 @@ class MlrunProject(ModelObj):
                           This allows you to control and specify where the workflow runner pod will be scheduled.
                           This setting is only relevant when the engine is set to 'remote' or for scheduled workflows,
                           and it will be ignored if the workflow is not run on a remote engine.
+        :param context:             mlrun context.
         :returns: ~py:class:`~mlrun.projects.pipelines._PipelineRunStatus` instance
         """
 
@@ -3368,6 +3442,7 @@ class MlrunProject(ModelObj):
             namespace=namespace,
             source=source,
             notifications=notifications,
+            context=context,
         )
         # run is None when scheduling
         if run and run.state == mlrun_pipelines.common.models.RunStatuses.failed:
@@ -3551,6 +3626,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 +3641,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 +3676,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/function.py

@@ -1036,9 +1036,10 @@ class RemoteRuntime(KubeResource):
         if args and sidecar.get("command"):
             sidecar["args"] = mlrun.utils.helpers.as_list(args)
 
-        # populate the sidecar resources from the function spec
+        # put the configured resources on the sidecar container instead of the reverse proxy container
         if self.spec.resources:
             sidecar["resources"] = self.spec.resources
+            self.spec.resources = None
 
     def _set_sidecar(self, name: str) -> dict:
         self.spec.config.setdefault("spec.sidecars", [])

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,16 @@ 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: 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 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 +432,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 +450,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 +609,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,12 @@ 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"),
+                endpoint_type=kwargs.get("endpoint_type"),
+            )
 
         self._update_weights(self.weights)
 
@@ -1001,7 +1005,10 @@ class VotingEnsemble(ParallelRun):
 
 
 def _init_endpoint_record(
-    graph_server: GraphServer, voting_ensemble: VotingEnsemble
+    graph_server: GraphServer,
+    voting_ensemble: VotingEnsemble,
+    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
@@ -1011,61 +1018,50 @@ 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: 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.
     """
 
     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 +1070,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 +1141,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 +1291,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."
             )