mlrun 1.10.0rc18__py3-none-any.whl → 1.11.0rc16__py3-none-any.whl

mlrun/serving/states.py

@@ -24,12 +24,15 @@ import inspect
 import os
 import pathlib
 import traceback
+import warnings
 from abc import ABC
+from collections.abc import Collection
 from copy import copy, deepcopy
 from inspect import getfullargspec, signature
 from typing import Any, Optional, Union, cast
 
 import storey.utils
+from deprecated import deprecated
 from storey import ParallelExecutionMechanisms
 
 import mlrun
@@ -38,17 +41,21 @@ import mlrun.common.schemas as schemas
 from mlrun.artifacts.llm_prompt import LLMPromptArtifact, PlaceholderDefaultDict
 from mlrun.artifacts.model import ModelArtifact
 from mlrun.datastore.datastore_profile import (
-    DatastoreProfileKafkaSource,
+    DatastoreProfileKafkaStream,
     DatastoreProfileKafkaTarget,
     DatastoreProfileV3io,
     datastore_profile_read,
 )
-from mlrun.datastore.model_provider.model_provider import ModelProvider
+from mlrun.datastore.model_provider.model_provider import (
+    InvokeResponseFormat,
+    ModelProvider,
+    UsageResponseKeys,
+)
 from mlrun.datastore.storeytargets import KafkaStoreyTarget, StreamStoreyTarget
-from mlrun.utils import get_data_from_path, logger, split_path
+from mlrun.utils import get_data_from_path, logger, set_data_by_path, split_path
 
 from ..config import config
-from ..datastore import get_stream_pusher
+from ..datastore import _DummyStream, get_stream_pusher
 from ..datastore.utils import (
     get_kafka_brokers_from_dict,
     parse_kafka_url,
@@ -85,25 +92,6 @@ class StepKinds:
     model_runner = "model_runner"
 
 
-_task_step_fields = [
-    "kind",
-    "class_name",
-    "class_args",
-    "handler",
-    "skip_context",
-    "after",
-    "function",
-    "comment",
-    "shape",
-    "full_event",
-    "on_error",
-    "responder",
-    "input_path",
-    "result_path",
-    "model_endpoint_creation_strategy",
-    "endpoint_type",
-]
-
 _default_fields_to_strip_from_step = [
     "model_endpoint_creation_strategy",
     "endpoint_type",
@@ -129,7 +117,14 @@ def new_remote_endpoint(
 class BaseStep(ModelObj):
     kind = "BaseStep"
     default_shape = "ellipse"
-    _dict_fields = ["kind", "comment", "after", "on_error"]
+    _dict_fields = [
+        "kind",
+        "comment",
+        "after",
+        "on_error",
+        "max_iterations",
+        "cycle_from",
+    ]
     _default_fields_to_strip = _default_fields_to_strip_from_step
 
     def __init__(
@@ -137,6 +132,7 @@ class BaseStep(ModelObj):
         name: Optional[str] = None,
         after: Optional[list] = None,
         shape: Optional[str] = None,
+        max_iterations: Optional[int] = None,
     ):
         self.name = name
         self._parent = None
@@ -150,6 +146,8 @@ class BaseStep(ModelObj):
         self.model_endpoint_creation_strategy = (
             schemas.ModelEndpointCreationStrategy.SKIP
         )
+        self._max_iterations = max_iterations
+        self.cycle_from = []
 
     def get_shape(self):
         """graphviz shape"""
@@ -343,6 +341,8 @@ class BaseStep(ModelObj):
         model_endpoint_creation_strategy: Optional[
             schemas.ModelEndpointCreationStrategy
         ] = None,
+        cycle_to: Optional[list[str]] = None,
+        max_iterations: Optional[int] = None,
         **class_args,
     ):
         """add a step right after this step and return the new step
@@ -372,21 +372,17 @@ class BaseStep(ModelObj):
                             to event["y"] resulting in {"x": 5, "y": <result>}
         :param model_endpoint_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`.
+                            * **overwrite**: If model endpoints with the same name exist, delete the `latest` one;
+                              create a new model endpoint entry and set it as `latest`.
 
-                            * **inplace** (default):
+                            * **inplace** (default): If model endpoints with the same name exist, update the `latest`
+                              entry; otherwise, create a new entry.
 
-                            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`.
+                            * **archive**: If model endpoints with the same name exist, preserve them;
+                              create a new model endpoint with the same name and set it to `latest`.
 
+        :param cycle_to:    list of step names to create a cycle to (for cyclic graphs)
+        :param max_iterations: maximum number of iterations for this step in case of a cycle graph
         :param class_args:  class init arguments
         """
         if hasattr(self, "steps"):
@@ -421,8 +417,39 @@ class BaseStep(ModelObj):
             # check that its not the root, todo: in future may gave nested flows
             step.after_step(self.name)
         parent._last_added = step
+        step.cycle_to(cycle_to or [])
+        step._max_iterations = max_iterations
         return step
 
+    def cycle_to(self, step_names: Union[str, list[str]]):
+        """create a cycle in the graph to the specified step names
+
+        example:
+            in the below example, a cycle is created from 'step3' to 'step1':
+            graph.to('step1')\
+                 .to('step2')\
+                 .to('step3')\
+                 .cycle_to(['step1'])  # creates a cycle from step3 to step1
+
+        :param step_names: list of step names to create a cycle to (for cyclic graphs)
+        """
+        root = self._extract_root_step()
+        if not isinstance(root, RootFlowStep):
+            raise GraphError("cycle_to() can only be called on a step within a graph")
+        if not root.allow_cyclic and step_names:
+            raise GraphError("cyclic graphs are not allowed, enable allow_cyclic")
+        step_names = [step_names] if isinstance(step_names, str) else step_names
+
+        for step_name in step_names:
+            if step_name not in root:
+                raise GraphError(
+                    f"step {step_name} doesnt exist in the graph under {self._parent.fullname}"
+                )
+            root[step_name].after_step(self.name, append=True)
+            root[step_name].cycle_from.append(self.name)
+
+        return self
+
     def set_flow(
         self,
         steps: list[Union[str, StepToDict, dict[str, Any]]],
@@ -517,7 +544,9 @@ class BaseStep(ModelObj):
 
         root = self._extract_root_step()
 
-        if not isinstance(root, RootFlowStep):
+        if not isinstance(root, RootFlowStep) or (
+            isinstance(root, RootFlowStep) and root.engine != "async"
+        ):
             raise GraphError(
                 "ModelRunnerStep can be added to 'Flow' topology graph only"
             )
@@ -541,8 +570,8 @@ class BaseStep(ModelObj):
         # Update model endpoints names in the root step
         root.update_model_endpoints_names(step_model_endpoints_names)
 
-    @staticmethod
     def _verify_shared_models(
+        self,
         root: "RootFlowStep",
         step: "ModelRunnerStep",
         step_model_endpoints_names: list[str],
@@ -571,35 +600,41 @@ class BaseStep(ModelObj):
             prefix, _ = mlrun.datastore.parse_store_uri(model_artifact_uri)
             # if the model artifact is a prompt, we need to get the model URI
             # to ensure that the shared runnable name is correct
+            llm_artifact_uri = None
             if prefix == mlrun.utils.StorePrefix.LLMPrompt:
                 llm_artifact, _ = mlrun.store_manager.get_store_artifact(
                     model_artifact_uri
                 )
+                llm_artifact_uri = llm_artifact.uri
                 model_artifact_uri = mlrun.utils.remove_tag_from_artifact_uri(
                     llm_artifact.spec.parent_uri
                 )
-            actual_shared_name = root.get_shared_model_name_by_artifact_uri(
-                model_artifact_uri
+            actual_shared_name, shared_model_class, shared_model_params = (
+                root.get_shared_model_by_artifact_uri(model_artifact_uri)
             )
 
-            if not shared_runnable_name:
-                if not actual_shared_name:
-                    raise GraphError(
-                        f"Can't find shared model for {name} model endpoint"
-                    )
-                else:
-                    step.class_args[schemas.ModelRunnerStepData.MODELS][name][
-                        schemas.ModelsData.MODEL_PARAMETERS.value
-                    ]["shared_runnable_name"] = actual_shared_name
-                    shared_models.append(actual_shared_name)
+            if not actual_shared_name:
+                raise GraphError(
+                    f"Can't find shared model named {shared_runnable_name}"
+                )
+            elif not shared_runnable_name:
+                step.class_args[schemas.ModelRunnerStepData.MODELS][name][
+                    schemas.ModelsData.MODEL_PARAMETERS.value
+                ]["shared_runnable_name"] = actual_shared_name
             elif actual_shared_name != shared_runnable_name:
                 raise GraphError(
                     f"Model endpoint {name} shared runnable name mismatch: "
                     f"expected {actual_shared_name}, got {shared_runnable_name}"
                 )
-            else:
-                shared_models.append(actual_shared_name)
-
+            shared_models.append(actual_shared_name)
+            self._edit_proxy_model_data(
+                step,
+                name,
+                actual_shared_name,
+                shared_model_params,
+                shared_model_class,
+                llm_artifact_uri or model_artifact_uri,
+            )
         undefined_shared_models = list(
             set(shared_models) - set(root.shared_models.keys())
         )
@@ -608,12 +643,71 @@ class BaseStep(ModelObj):
                 f"The following shared models are not defined in the graph: {undefined_shared_models}."
             )
 
+    @staticmethod
+    def _edit_proxy_model_data(
+        step: "ModelRunnerStep",
+        name: str,
+        actual_shared_name: str,
+        shared_model_params: dict,
+        shared_model_class: Any,
+        artifact: Union[ModelArtifact, LLMPromptArtifact, str],
+    ):
+        monitoring_data = step.class_args.setdefault(
+            schemas.ModelRunnerStepData.MONITORING_DATA, {}
+        )
+
+        # edit monitoring data according to the shared model parameters
+        monitoring_data[name][schemas.MonitoringData.INPUT_PATH] = shared_model_params[
+            "input_path"
+        ]
+        monitoring_data[name][schemas.MonitoringData.RESULT_PATH] = shared_model_params[
+            "result_path"
+        ]
+        monitoring_data[name][schemas.MonitoringData.INPUTS] = shared_model_params[
+            "inputs"
+        ]
+        monitoring_data[name][schemas.MonitoringData.OUTPUTS] = shared_model_params[
+            "outputs"
+        ]
+        monitoring_data[name][schemas.MonitoringData.MODEL_CLASS] = (
+            shared_model_class
+            if isinstance(shared_model_class, str)
+            else shared_model_class.__class__.__name__
+        )
+        if actual_shared_name and actual_shared_name not in step._shared_proxy_mapping:
+            step._shared_proxy_mapping[actual_shared_name] = {
+                name: artifact.uri
+                if isinstance(artifact, ModelArtifact | LLMPromptArtifact)
+                else artifact
+            }
+        elif actual_shared_name:
+            step._shared_proxy_mapping[actual_shared_name].update(
+                {
+                    name: artifact.uri
+                    if isinstance(artifact, ModelArtifact | LLMPromptArtifact)
+                    else artifact
+                }
+            )
+
 
 class TaskStep(BaseStep):
     """task execution step, runs a class or handler"""
 
     kind = "task"
-    _dict_fields = _task_step_fields
+    _dict_fields = BaseStep._dict_fields + [
+        "class_name",
+        "class_args",
+        "handler",
+        "skip_context",
+        "function",
+        "shape",
+        "full_event",
+        "responder",
+        "input_path",
+        "result_path",
+        "model_endpoint_creation_strategy",
+        "endpoint_type",
+    ]
     _default_class = ""
 
     def __init__(
@@ -639,6 +733,7 @@ class TaskStep(BaseStep):
         self.handler = handler
         self.function = function
         self._handler = None
+        self._outlets_selector = None
         self._object = None
         self._async_object = None
         self.skip_context = None
@@ -706,6 +801,8 @@ class TaskStep(BaseStep):
                     handler = "do"
             if handler:
                 self._handler = getattr(self._object, handler, None)
+            if hasattr(self._object, "select_outlets"):
+                self._outlets_selector = self._object.select_outlets
 
         self._set_error_handler()
         if mode != "skip":
@@ -879,7 +976,7 @@ class ErrorStep(TaskStep):
     """error execution step, runs a class or handler"""
 
     kind = "error_step"
-    _dict_fields = _task_step_fields + ["before", "base_step"]
+    _dict_fields = TaskStep._dict_fields + ["before", "base_step"]
     _default_class = ""
 
     def __init__(
@@ -916,7 +1013,7 @@ class RouterStep(TaskStep):
 
     kind = "router"
     default_shape = "doubleoctagon"
-    _dict_fields = _task_step_fields + ["routes", "name"]
+    _dict_fields = TaskStep._dict_fields + ["routes", "name"]
     _default_class = "mlrun.serving.ModelRouter"
 
     def __init__(
@@ -983,20 +1080,14 @@ class RouterStep(TaskStep):
         :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):
+                           * **overwrite**: If model endpoints with the same name exist, delete the `latest` one;
+                             create a new model endpoint entry and set it as `latest`.
 
-                           1. If model endpoints with the same name exist, update the `latest` entry.
-                           2. Otherwise, create a new entry.
+                           * **inplace** (default): If model endpoints with the same name exist, update the `latest`
+                             entry;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`.
+                           * **archive**: If model endpoints with the same name exist, preserve them;
+                             create a new model endpoint with the same name and set it to `latest`.
 
         """
         if len(self.routes.keys()) >= MAX_MODELS_PER_ROUTER and key not in self.routes:
@@ -1090,6 +1181,7 @@ class Model(storey.ParallelExecutionRunnable, ModelObj):
         "artifact_uri",
         "shared_runnable_name",
         "shared_proxy_mapping",
+        "execution_mechanism",
     ]
     kind = "model"
 
@@ -1111,6 +1203,8 @@ class Model(storey.ParallelExecutionRunnable, ModelObj):
         self.invocation_artifact: Optional[LLMPromptArtifact] = None
         self.model_artifact: Optional[ModelArtifact] = None
         self.model_provider: Optional[ModelProvider] = None
+        self._artifact_were_loaded = False
+        self._execution_mechanism = None
 
     def __init_subclass__(cls):
         super().__init_subclass__()
@@ -1130,13 +1224,33 @@ class Model(storey.ParallelExecutionRunnable, ModelObj):
                 raise_missing_schema_exception=False,
             )
 
-    def _load_artifacts(self) -> None:
-        artifact = self._get_artifact_object()
-        if isinstance(artifact, LLMPromptArtifact):
-            self.invocation_artifact = artifact
-            self.model_artifact = self.invocation_artifact.model_artifact
+        # Check if the relevant predict method is implemented when trying to initialize the model
+        if self._execution_mechanism == storey.ParallelExecutionMechanisms.asyncio:
+            if self.__class__.predict_async is Model.predict_async:
+                raise mlrun.errors.ModelRunnerError(
+                    {
+                        self.name: f"is running with {self._execution_mechanism} "
+                        f"execution_mechanism but predict_async() is not implemented"
+                    }
+                )
         else:
-            self.model_artifact = artifact
+            if self.__class__.predict is Model.predict:
+                raise mlrun.errors.ModelRunnerError(
+                    {
+                        self.name: f"is running with {self._execution_mechanism} execution_mechanism but predict() "
+                        f"is not implemented"
+                    }
+                )
+
+    def _load_artifacts(self) -> None:
+        if not self._artifact_were_loaded:
+            artifact = self._get_artifact_object()
+            if isinstance(artifact, LLMPromptArtifact):
+                self.invocation_artifact = artifact
+                self.model_artifact = self.invocation_artifact.model_artifact
+            else:
+                self.model_artifact = artifact
+            self._artifact_were_loaded = True
 
     def _get_artifact_object(
         self, proxy_uri: Optional[str] = None
@@ -1144,7 +1258,9 @@ class Model(storey.ParallelExecutionRunnable, ModelObj):
         uri = proxy_uri or self.artifact_uri
         if uri:
             if mlrun.datastore.is_store_uri(uri):
-                artifact, _ = mlrun.store_manager.get_store_artifact(uri)
+                artifact, _ = mlrun.store_manager.get_store_artifact(
+                    uri, allow_empty_resources=True
+                )
                 return artifact
             else:
                 raise ValueError(
@@ -1158,13 +1274,15 @@ class Model(storey.ParallelExecutionRunnable, ModelObj):
 
     def predict(self, body: Any, **kwargs) -> Any:
         """Override to implement prediction logic. If the logic requires asyncio, override predict_async() instead."""
-        return body
+        raise NotImplementedError("predict() method not implemented")
 
     async def predict_async(self, body: Any, **kwargs) -> Any:
         """Override to implement prediction logic if the logic requires asyncio."""
-        return body
+        raise NotImplementedError("predict_async() method not implemented")
 
     def run(self, body: Any, path: str, origin_name: Optional[str] = None) -> Any:
+        if isinstance(body, list):
+            body = self.format_batch(body)
         return self.predict(body)
 
     async def run_async(
@@ -1203,28 +1321,117 @@ class Model(storey.ParallelExecutionRunnable, ModelObj):
             return model_file, extra_dataitems
         return None, None
 
+    @staticmethod
+    def format_batch(body: Any):
+        return body
+
 
 class LLModel(Model):
+    """
+    A model wrapper for handling LLM (Large Language Model) prompt-based inference.
+
+    This class extends the base `Model` to provide specialized handling for
+    `LLMPromptArtifact` objects, enabling both synchronous and asynchronous
+    invocation of language models.
+
+    **Model Invocation**:
+
+    - The execution of enriched prompts is delegated to the `model_provider`
+      configured for the model (e.g., **Hugging Face** or **OpenAI**).
+    - The `model_provider` is responsible for sending the prompt to the correct
+      backend API and returning the generated output.
+    - Users can override the `predict` and `predict_async` methods to customize
+      the behavior of the model invocation.
+
+    **Prompt Enrichment Overview**:
+
+    - If an `LLMPromptArtifact` is found, load its prompt template and fill in
+      placeholders using values from the request body.
+    - If the artifact is not an `LLMPromptArtifact`, skip formatting and attempt
+      to retrieve `messages` directly from the request body using the input path.
+
+    **Simplified Example**:
+
+    Input body::
+
+        {"city": "Paris", "days": 3}
+
+    Prompt template in artifact::
+
+        [
+            {"role": "system", "content": "You are a travel planning assistant."},
+            {"role": "user", "content": "Create a {{days}}-day itinerary for {{city}}."},
+        ]
+
+    Result after enrichment::
+
+        [
+            {"role": "system", "content": "You are a travel planning assistant."},
+            {"role": "user", "content": "Create a 3-day itinerary for Paris."},
+        ]
+
+    :param name: Name of the model.
+    :param input_path: Path in the request body where input data is located.
+    :param result_path: Path in the response body where model outputs and the statistics
+                        will be stored.
+    """
+
+    _dict_fields = Model._dict_fields + ["result_path", "input_path"]
+
     def __init__(
-        self, name: str, input_path: Optional[Union[str, list[str]]] = None, **kwargs
+        self,
+        name: str,
+        input_path: Optional[Union[str, list[str]]] = None,
+        result_path: Optional[Union[str, list[str]]] = None,
+        **kwargs,
     ):
         super().__init__(name, **kwargs)
         self._input_path = split_path(input_path)
+        self._result_path = split_path(result_path)
+        logger.info(
+            "LLModel initialized",
+            model_name=name,
+            input_path=input_path,
+            result_path=result_path,
+        )
 
     def predict(
         self,
         body: Any,
         messages: Optional[list[dict]] = None,
-        model_configuration: Optional[dict] = None,
+        invocation_config: Optional[dict] = None,
         **kwargs,
     ) -> Any:
+        llm_prompt_artifact = kwargs.get("llm_prompt_artifact")
         if isinstance(
-            self.invocation_artifact, mlrun.artifacts.LLMPromptArtifact
+            llm_prompt_artifact, mlrun.artifacts.LLMPromptArtifact
         ) and isinstance(self.model_provider, ModelProvider):
-            body["result"] = self.model_provider.invoke(
+            logger.debug(
+                "Invoking model provider",
+                model_name=self.name,
+                messages=messages,
+                invocation_config=invocation_config,
+            )
+            response_with_stats = self.model_provider.invoke(
                 messages=messages,
-                as_str=True,
-                **(model_configuration or {}),
+                invoke_response_format=InvokeResponseFormat.USAGE,
+                **(invocation_config or {}),
+            )
+            set_data_by_path(
+                path=self._result_path, data=body, value=response_with_stats
+            )
+            logger.debug(
+                "LLModel prediction completed",
+                model_name=self.name,
+                answer=response_with_stats.get("answer"),
+                usage=response_with_stats.get("usage"),
+            )
+        else:
+            logger.warning(
+                "LLModel invocation artifact or model provider not set, skipping prediction",
+                model_name=self.name,
+                invocation_artifact_type=type(llm_prompt_artifact).__name__,
+                model_provider_type=type(self.model_provider).__name__,
             )
         return body
 
@@ -1232,61 +1439,130 @@ class LLModel(Model):
         self,
         body: Any,
         messages: Optional[list[dict]] = None,
-        model_configuration: Optional[dict] = None,
+        invocation_config: Optional[dict] = None,
         **kwargs,
     ) -> Any:
+        llm_prompt_artifact = kwargs.get("llm_prompt_artifact")
         if isinstance(
-            self.invocation_artifact, mlrun.artifacts.LLMPromptArtifact
+            llm_prompt_artifact, mlrun.artifacts.LLMPromptArtifact
         ) and isinstance(self.model_provider, ModelProvider):
-            body["result"] = await self.model_provider.async_invoke(
+            logger.debug(
+                "Async invoking model provider",
+                model_name=self.name,
                 messages=messages,
-                as_str=True,
-                **(model_configuration or {}),
+                invocation_config=invocation_config,
+            )
+            response_with_stats = await self.model_provider.async_invoke(
+                messages=messages,
+                invoke_response_format=InvokeResponseFormat.USAGE,
+                **(invocation_config or {}),
+            )
+            set_data_by_path(
+                path=self._result_path, data=body, value=response_with_stats
+            )
+            logger.debug(
+                "LLModel async prediction completed",
+                model_name=self.name,
+                answer=response_with_stats.get("answer"),
+                usage=response_with_stats.get("usage"),
+            )
+        else:
+            logger.warning(
+                "LLModel invocation artifact or model provider not set, skipping async prediction",
+                model_name=self.name,
+                invocation_artifact_type=type(llm_prompt_artifact).__name__,
+                model_provider_type=type(self.model_provider).__name__,
             )
         return body
 
+    def init(self):
+        super().init()
+
+        if not self.model_provider:
+            if self._execution_mechanism != storey.ParallelExecutionMechanisms.asyncio:
+                unchanged_predict = self.__class__.predict is LLModel.predict
+                predict_function_name = "predict"
+            else:
+                unchanged_predict = (
+                    self.__class__.predict_async is LLModel.predict_async
+                )
+                predict_function_name = "predict_async"
+            if unchanged_predict:
+                raise mlrun.errors.MLRunRuntimeError(
+                    f"Model provider could not be determined for model '{self.name}',"
+                    f" and the {predict_function_name} function was not overridden."
+                )
+
     def run(self, body: Any, path: str, origin_name: Optional[str] = None) -> Any:
-        messages, model_configuration = self.enrich_prompt(body, origin_name)
+        llm_prompt_artifact = self._get_invocation_artifact(origin_name)
+        messages, invocation_config = self.enrich_prompt(
+            body, origin_name, llm_prompt_artifact
+        )
+        logger.info(
+            "Calling LLModel predict",
+            model_name=self.name,
+            model_endpoint_name=origin_name,
+            messages_len=len(messages) if messages else 0,
+        )
         return self.predict(
-            body, messages=messages, model_configuration=model_configuration
+            body,
+            messages=messages,
+            invocation_config=invocation_config,
+            llm_prompt_artifact=llm_prompt_artifact,
         )
 
     async def run_async(
         self, body: Any, path: str, origin_name: Optional[str] = None
     ) -> Any:
-        messages, model_configuration = self.enrich_prompt(body, origin_name)
+        llm_prompt_artifact = self._get_invocation_artifact(origin_name)
+        messages, invocation_config = self.enrich_prompt(
+            body, origin_name, llm_prompt_artifact
+        )
+        logger.info(
+            "Calling LLModel async predict",
+            model_name=self.name,
+            model_endpoint_name=origin_name,
+            messages_len=len(messages) if messages else 0,
+        )
         return await self.predict_async(
-            body, messages=messages, model_configuration=model_configuration
+            body,
+            messages=messages,
+            invocation_config=invocation_config,
+            llm_prompt_artifact=llm_prompt_artifact,
         )
 
     def enrich_prompt(
-        self, body: dict, origin_name: str
+        self,
+        body: dict,
+        origin_name: str,
+        llm_prompt_artifact: Optional[LLMPromptArtifact] = None,
     ) -> Union[tuple[list[dict], dict], tuple[None, None]]:
-        if origin_name and self.shared_proxy_mapping:
-            llm_prompt_artifact = self.shared_proxy_mapping.get(origin_name)
-            if isinstance(llm_prompt_artifact, str):
-                llm_prompt_artifact = self._get_artifact_object(llm_prompt_artifact)
-                self.shared_proxy_mapping[origin_name] = llm_prompt_artifact
-        else:
-            llm_prompt_artifact = (
-                self.invocation_artifact or self._get_artifact_object()
-            )
-        if not (
+        logger.info(
+            "Enriching prompt",
+            model_name=self.name,
+            model_endpoint_name=origin_name,
+        )
+        if not llm_prompt_artifact or not (
             llm_prompt_artifact and isinstance(llm_prompt_artifact, LLMPromptArtifact)
         ):
             logger.warning(
-                "LLMModel must be provided with LLMPromptArtifact",
+                "LLModel must be provided with LLMPromptArtifact",
+                model_name=self.name,
+                artifact_type=type(llm_prompt_artifact).__name__,
                 llm_prompt_artifact=llm_prompt_artifact,
             )
-            return None, None
-        prompt_legend = llm_prompt_artifact.spec.prompt_legend
-        prompt_template = deepcopy(llm_prompt_artifact.read_prompt())
+            prompt_legend, prompt_template, invocation_config = {}, [], {}
+        else:
+            prompt_legend = llm_prompt_artifact.spec.prompt_legend
+            prompt_template = deepcopy(llm_prompt_artifact.read_prompt())
+            invocation_config = llm_prompt_artifact.spec.invocation_config
         input_data = copy(get_data_from_path(self._input_path, body))
-        if isinstance(input_data, dict):
+        if isinstance(input_data, dict) and prompt_template:
             kwargs = (
                 {
                     place_holder: input_data.get(body_map["field"])
                     for place_holder, body_map in prompt_legend.items()
+                    if input_data.get(body_map["field"])
                 }
                 if prompt_legend
                 else {}
@@ -1298,23 +1574,124 @@ class LLModel(Model):
                     message["content"] = message["content"].format(**input_data)
                 except KeyError as e:
                     logger.warning(
-                        "Input data was missing a placeholder, placeholder stay unformatted",
-                        key_error=e,
+                        "Input data missing placeholder, content stays unformatted",
+                        model_name=self.name,
+                        key_error=mlrun.errors.err_to_str(e),
                     )
                     message["content"] = message["content"].format_map(
                         default_place_holders
                     )
+        elif isinstance(input_data, dict) and not prompt_template:
+            # If there is no prompt template, we assume the input data is already in the correct format.
+            logger.debug("Attempting to retrieve messages from the request body.")
+            prompt_template = input_data.get("messages", [])
         else:
             logger.warning(
-                f"Expected input data to be a dict, but received input data from type {type(input_data)} prompt "
-                f"template stay unformatted",
+                "Expected input data to be a dict, prompt template stays unformatted",
+                model_name=self.name,
+                input_data_type=type(input_data).__name__,
             )
-        return prompt_template, llm_prompt_artifact.spec.model_configuration
+        return prompt_template, invocation_config
+
+    def _get_invocation_artifact(
+        self, origin_name: Optional[str] = None
+    ) -> Union[LLMPromptArtifact, None]:
+        """
+        Get the LLMPromptArtifact object for this model.
+
+        :param proxy_uri: Optional; URI to the proxy artifact.
+        :return: LLMPromptArtifact object or None if not found.
+        """
+        if origin_name and self.shared_proxy_mapping:
+            llm_prompt_artifact = self.shared_proxy_mapping.get(origin_name)
+            if isinstance(llm_prompt_artifact, str):
+                llm_prompt_artifact = self._get_artifact_object(llm_prompt_artifact)
+                self.shared_proxy_mapping[origin_name] = llm_prompt_artifact
+        elif self._artifact_were_loaded:
+            llm_prompt_artifact = self.invocation_artifact
+        else:
+            self._load_artifacts()
+            llm_prompt_artifact = self.invocation_artifact
+        return llm_prompt_artifact
+
+
+class ModelRunnerSelector(ModelObj):
+    """
+    Strategy for controlling model selection and output routing in ModelRunnerStep.
+
+    Subclass this to implement custom logic for agent workflows:
+    - `select_models()`: Called BEFORE execution to choose which models run
+    - `select_outlets()`: Called AFTER execution to route output to downstream steps
+
+    Return `None` from either method to use default behavior (all models / all outlets).
+
+    Example::
+
+        class ToolSelector(ModelRunnerSelector):
+            def select_outlets(self, event):
+                tool = event.get("tool_call")
+                return [tool] if tool else ["final"]
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__()
+
+    def __init_subclass__(cls):
+        super().__init_subclass__()
+        cls._dict_fields = list(
+            set(cls._dict_fields)
+            | set(inspect.signature(cls.__init__).parameters.keys())
+        )
+        cls._dict_fields.remove("self")
+
+    def select_models(
+        self,
+        event: Any,
+        available_models: list[Model],
+    ) -> Optional[Union[list[str], list[Model]]]:
+        """
+        Called before model execution.
+
+        :param event: The full event
+        :param available_models: List of available models
+
+        Returns the models to execute (by name or Model objects).
+        """
+        return None
+
+    def select_outlets(
+        self,
+        event: Any,
+    ) -> Optional[list[str]]:
+        """
+        Called after model execution.
+
+        :param event: The event body after model execution
+        :return: Returns the downstream outlets to route the event to.
+        """
+        return None
 
 
-class ModelSelector:
+# TODO: Remove in 1.13.0
+@deprecated(
+    version="1.11.0",
+    reason="ModelSelector will be removed in 1.13.0, use ModelRunnerSelector instead",
+    category=FutureWarning,
+)
+class ModelSelector(ModelObj):
     """Used to select which models to run on each event."""
 
+    def __init__(self, **kwargs):
+        super().__init__()
+
+    def __init_subclass__(cls):
+        super().__init_subclass__()
+        cls._dict_fields = list(
+            set(cls._dict_fields)
+            | set(inspect.signature(cls.__init__).parameters.keys())
+        )
+        cls._dict_fields.remove("self")
+
     def select(
         self, event, available_models: list[Model]
     ) -> Union[list[str], list[Model]]:
@@ -1332,16 +1709,22 @@ 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.
+    :param model_runner_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, context, model_selector: Optional[ModelSelector] = None, **kwargs
+        self,
+        *args,
+        context,
+        model_runner_selector: Optional[ModelRunnerSelector] = None,
+        raise_exception: bool = True,
+        **kwargs,
     ):
         super().__init__(*args, **kwargs)
-        self.model_selector = model_selector or ModelSelector()
+        self.model_runner_selector = model_runner_selector or ModelRunnerSelector()
         self.context = context
+        self._raise_exception = raise_exception
 
     def preprocess_event(self, event):
         if not hasattr(event, "_metadata"):
@@ -1354,7 +1737,31 @@ class ModelRunner(storey.ParallelExecution):
 
     def select_runnables(self, event):
         models = cast(list[Model], self.runnables)
-        return self.model_selector.select(event, models)
+        return self.model_runner_selector.select_models(event, models)
+
+    def select_outlets(self, event) -> Optional[Collection[str]]:
+        sys_outlets = [f"{self.name}_error_raise"]
+        if "background_task_status_step" in self._name_to_outlet:
+            sys_outlets.append("background_task_status_step")
+        if self._raise_exception and self._is_error(event):
+            return sys_outlets
+        user_outlets = self.model_runner_selector.select_outlets(event)
+        if user_outlets:
+            return (
+                user_outlets if isinstance(user_outlets, list) else [user_outlets]
+            ) + sys_outlets
+        return None
+
+    def _is_error(self, event: dict) -> bool:
+        if len(self.runnables) == 1:
+            if isinstance(event, dict):
+                return event.get("error") is not None
+        else:
+            for model in event:
+                body_by_model = event.get(model)
+                if isinstance(body_by_model, dict) and "error" in body_by_model:
+                    return True
+        return False
 
 
 class MonitoredStep(ABC, TaskStep, StepToDict):
@@ -1406,34 +1813,122 @@ class ModelRunnerStep(MonitoredStep):
         model_runner_step.add_model(..., model_class=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.
-    :param raise_exception:  If True, an error will be raised when model selection fails or if one of the models raised
-      an error. If False, the error will appear in the output event.
+    Note when ModelRunnerStep is used in a graph, MLRun automatically imports
+    the default language model class (LLModel) during function deployment.
+
+    Note ModelRunnerStep can only be added to a graph that has the flow topology and running with async engine.
 
-    :raise ModelRunnerError - when a model raise an error the ModelRunnerStep will handle it, collect errors and outputs
-                              from added models, If raise_exception is True will raise ModelRunnerError Else will add
-                              the error msg as part of the event body mapped by model name if more than one model was
-                              added to the ModelRunnerStep
+    Note see configure_pool_resource method documentation for default number of max threads and max processes.
+
+    :raise ModelRunnerError: when a model raises an error the ModelRunnerStep will handle it, collect errors and
+                              outputs from added models. If raise_exception is True will raise ModelRunnerError. Else
+                              will add the error msg as part of the event body mapped by model name if more than
+                              one model was added to the ModelRunnerStep
     """
 
     kind = "model_runner"
-    _dict_fields = MonitoredStep._dict_fields + ["_shared_proxy_mapping"]
+    _dict_fields = MonitoredStep._dict_fields + [
+        "_shared_proxy_mapping",
+        "max_processes",
+        "max_threads",
+        "pool_factor",
+    ]
 
     def __init__(
         self,
         *args,
         name: Optional[str] = None,
+        model_runner_selector: Optional[Union[str, ModelRunnerSelector]] = None,
+        model_runner_selector_parameters: Optional[dict] = None,
         model_selector: Optional[Union[str, ModelSelector]] = None,
+        model_selector_parameters: Optional[dict] = None,
         raise_exception: bool = True,
         **kwargs,
     ):
+        """
+
+        :param name:                                The name of the ModelRunnerStep.
+        :param model_runner_selector:               ModelRunnerSelector instance whose select_models()
+                                                    and select_outlets() methods will be used
+                                                    to select models to run on each event and outlets to
+                                                    route the event to.
+        :param model_runner_selector_parameters:    Parameters for the model_runner_selector, if model_runner_selector
+                                                    is the class name we will use this param when
+                                                    initializing the selector.
+        :param model_selector:                      (Deprecated)
+        :param model_selector_parameters:           (Deprecated)
+        :param raise_exception:                     Determines whether to raise ModelRunnerError when one or more models
+                                                    raise an error during execution.
+                                                    If False, errors will be added to the event body.
+        """
+        self.max_processes = None
+        self.max_threads = None
+        self.pool_factor = None
+
+        if (model_selector or model_selector_parameters) and (
+            model_runner_selector or model_runner_selector_parameters
+        ):
+            raise GraphError(
+                "Cannot provide both `model_selector`/`model_selector_parameters` "
+                "and `model_runner_selector`/`model_runner_selector_parameters`. "
+                "Please use only the latter pair."
+            )
+        if model_selector or model_selector_parameters:
+            warnings.warn(
+                "`model_selector` and `model_selector_parameters` are deprecated, "
+                "please use `model_runner_selector` and `model_runner_selector_parameters` instead.",
+                # TODO: Remove this in 1.13.0
+                FutureWarning,
+            )
+            if isinstance(model_selector, ModelSelector) and model_selector_parameters:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "Cannot provide a model_selector object as argument to `model_selector` and also provide "
+                    "`model_selector_parameters`."
+                )
+            if model_selector:
+                model_selector_parameters = model_selector_parameters or (
+                    model_selector.to_dict()
+                    if isinstance(model_selector, ModelSelector)
+                    else {}
+                )
+                model_selector = (
+                    model_selector
+                    if isinstance(model_selector, str)
+                    else model_selector.__class__.__name__
+                )
+        else:
+            if (
+                isinstance(model_runner_selector, ModelRunnerSelector)
+                and model_runner_selector_parameters
+            ):
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "Cannot provide a model_runner_selector object as argument to `model_runner_selector` "
+                    "and also provide `model_runner_selector_parameters`."
+                )
+            if model_runner_selector:
+                model_runner_selector_parameters = model_runner_selector_parameters or (
+                    model_runner_selector.to_dict()
+                    if isinstance(model_runner_selector, ModelRunnerSelector)
+                    else {}
+                )
+                model_runner_selector = (
+                    model_runner_selector
+                    if isinstance(model_runner_selector, str)
+                    else model_runner_selector.__class__.__name__
+                )
+
         super().__init__(
             *args,
             name=name,
             raise_exception=raise_exception,
             class_name="mlrun.serving.ModelRunner",
-            class_args=dict(model_selector=model_selector),
+            class_args=dict(
+                model_selector=(model_selector, model_selector_parameters),
+                model_runner_selector=(
+                    model_runner_selector,
+                    model_runner_selector_parameters,
+                ),
+            ),
             **kwargs,
         )
         self.raise_exception = raise_exception
@@ -1449,10 +1944,6 @@ class ModelRunnerStep(MonitoredStep):
         model_endpoint_creation_strategy: Optional[
             schemas.ModelEndpointCreationStrategy
         ] = schemas.ModelEndpointCreationStrategy.INPLACE,
-        inputs: Optional[list[str]] = None,
-        outputs: Optional[list[str]] = None,
-        input_path: Optional[str] = None,
-        result_path: Optional[str] = None,
         override: bool = False,
     ) -> None:
         """
@@ -1465,28 +1956,18 @@ class ModelRunnerStep(MonitoredStep):
         :param shared_model_name:   str, the name of the shared model that is already defined within the graph
         :param labels:              model endpoint labels, should be list of str or mapping of str:str
         :param model_endpoint_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 inputs:              list of the model inputs (e.g. features) ,if provided will override the inputs
-                                    that been configured in the model artifact, please note that those inputs need to
-                                    be equal in length and order to the inputs that model_class predict method expects
-        :param outputs:             list of the model outputs (e.g. labels) ,if provided will override the outputs
-                                    that been configured in the model artifact, please note that those outputs need to
-                                    be equal to the model_class predict method outputs (length, and order)
-        :param input_path:          input path inside the user event, expect scopes to be defined by dot notation
-                                    (e.g "inputs.my_model_inputs"). expects list or dictionary type object in path.
-        :param result_path:         result path inside the user output event, expect scopes to be defined by dot
-                                    notation (e.g "outputs.my_model_outputs") expects list or dictionary type object
-                                    in path.
+                 * **overwrite**: If model endpoints with the same name exist, delete the `latest` one;
+                   create a new model endpoint entry and set it as `latest`.
+
+                 * **inplace** (default): If model endpoints with the same name exist, update the `latest` entry;
+                   otherwise, create a new entry.
+
+                 * **archive**: If model endpoints with the same name exist, preserve them;
+                   create a new model endpoint with the same name and set it to `latest`.
+
         :param override:            bool allow override existing model on the current ModelRunnerStep.
+        :raise GraphError:  when the shared model is not found in the root flow step shared models.
         """
         model_class, model_params = (
             "mlrun.serving.Model",
@@ -1503,11 +1984,21 @@ class ModelRunnerStep(MonitoredStep):
                 "model_artifact must be a string, ModelArtifact or LLMPromptArtifact"
             )
         root = self._extract_root_step()
+        shared_model_params = {}
         if isinstance(root, RootFlowStep):
-            shared_model_name = (
-                shared_model_name
-                or root.get_shared_model_name_by_artifact_uri(model_artifact_uri)
+            actual_shared_model_name, shared_model_class, shared_model_params = (
+                root.get_shared_model_by_artifact_uri(model_artifact_uri)
             )
+            if not actual_shared_model_name or (
+                shared_model_name and actual_shared_model_name != shared_model_name
+            ):
+                raise GraphError(
+                    f"ModelRunnerStep can only add proxy models that were added to the root flow step, "
+                    f"model {shared_model_name} is not in the shared models."
+                )
+            elif not shared_model_name:
+                shared_model_name = actual_shared_model_name
+                model_params["shared_runnable_name"] = shared_model_name
             if not root.shared_models or (
                 root.shared_models
                 and shared_model_name
@@ -1517,17 +2008,31 @@ class ModelRunnerStep(MonitoredStep):
                     f"ModelRunnerStep can only add proxy models that were added to the root flow step, "
                     f"model {shared_model_name} is not in the shared models."
                 )
-        if shared_model_name not in self._shared_proxy_mapping:
+            monitoring_data = self.class_args.get(
+                schemas.ModelRunnerStepData.MONITORING_DATA, {}
+            )
+            monitoring_data.setdefault(endpoint_name, {})[
+                schemas.MonitoringData.MODEL_CLASS
+            ] = (
+                shared_model_class
+                if isinstance(shared_model_class, str)
+                else shared_model_class.__class__.__name__
+            )
+            self.class_args[schemas.ModelRunnerStepData.MONITORING_DATA] = (
+                monitoring_data
+            )
+
+        if shared_model_name and shared_model_name not in self._shared_proxy_mapping:
             self._shared_proxy_mapping[shared_model_name] = {
                 endpoint_name: model_artifact.uri
-                if isinstance(model_artifact, (ModelArtifact, LLMPromptArtifact))
+                if isinstance(model_artifact, ModelArtifact | LLMPromptArtifact)
                 else model_artifact
             }
-        else:
+        elif override and shared_model_name:
             self._shared_proxy_mapping[shared_model_name].update(
                 {
                     endpoint_name: model_artifact.uri
-                    if isinstance(model_artifact, (ModelArtifact, LLMPromptArtifact))
+                    if isinstance(model_artifact, ModelArtifact | LLMPromptArtifact)
                     else model_artifact
                 }
             )
@@ -1538,11 +2043,11 @@ class ModelRunnerStep(MonitoredStep):
             model_artifact=model_artifact,
             labels=labels,
             model_endpoint_creation_strategy=model_endpoint_creation_strategy,
+            inputs=shared_model_params.get("inputs"),
+            outputs=shared_model_params.get("outputs"),
+            input_path=shared_model_params.get("input_path"),
+            result_path=shared_model_params.get("result_path"),
             override=override,
-            inputs=inputs,
-            outputs=outputs,
-            input_path=input_path,
-            result_path=result_path,
             **model_params,
         )
 
@@ -1567,48 +2072,52 @@ class ModelRunnerStep(MonitoredStep):
         Add a Model to this ModelRunner.
 
         :param endpoint_name:       str, will identify the model in the ModelRunnerStep, and assign model endpoint name
-        :param model_class:         Model class name
+        :param model_class:         Model class name. If LLModel is chosen
+                                    (either by name `LLModel` or by its full path, e.g. mlrun.serving.states.LLModel),
+                                    outputs will be overridden with UsageResponseKeys fields.
         :param execution_mechanism: Parallel execution mechanism to be used to execute this model. Must be one of:
-            * "process_pool" – To run in a separate process from a process pool. This is appropriate for CPU or GPU
-                intensive tasks as they would otherwise block the main process by holding Python's Global Interpreter
-                Lock (GIL).
-            * "dedicated_process" – To run in a separate dedicated process. This is appropriate for CPU or GPU intensive
-                tasks that also require significant Runnable-specific initialization (e.g. a large model).
-            * "thread_pool" – To run in a separate thread. This is appropriate for blocking I/O tasks, as they would
-                otherwise block the main event loop thread.
-            * "asyncio" – To run in an asyncio task. This is appropriate for I/O tasks that use asyncio, allowing the
-                event loop to continue running while waiting for a response.
-            * "shared_executor" – Reuses an external executor (typically managed by the flow or context) to execute the
-                runnable. Should be used only if you have multiply `ParallelExecution` in the same flow and especially
-                useful when:
-                - You want to share a heavy resource like a large model loaded onto a GPU.
-                - You want to centralize task scheduling or coordination for multiple lightweight tasks.
-                - You aim to minimize overhead from creating new executors or processes/threads per runnable.
-                The runnable is expected to be pre-initialized and reused across events, enabling efficient use of
-                memory and hardware accelerators.
-            * "naive" – To run in the main event loop. This is appropriate only for trivial computation and/or file I/O.
-                It means that the runnable will not actually be run in parallel to anything else.
-
-            :param model_artifact:      model artifact or mlrun model artifact uri
-            :param labels:              model endpoint labels, should be list of str or mapping of str:str
-            :param model_endpoint_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 inputs:              list of the model inputs (e.g. features) ,if provided will override the inputs
+
+                            * **process_pool**: To run in a separate process from a process pool. This is appropriate
+                              for CPU or GPU intensive tasks as they would otherwise block the main process by holding
+                              Python's Global Interpreter Lock (GIL).
+
+                            * **dedicated_process**: To run in a separate dedicated process. This is appropriate for CPU
+                              or GPU intensive tasks that also require significant Runnable-specific initialization
+                              (e.g. a large model).
+
+                            * **thread_pool**: To run in a separate thread. This is appropriate for blocking I/O tasks,
+                              as they would otherwise block the main event loop thread.
+
+                            * **asyncio**: To run in an asyncio task. This is appropriate for I/O tasks that use
+                              asyncio, allowing the event loop to continue running while waiting for a response.
+
+                            * **naive**: To run in the main event loop. This is appropriate only for trivial computation
+                              and/or file I/O. It means that the runnable will not actually be run in parallel to
+                              anything else.
+
+        :param model_artifact:      model artifact or mlrun model artifact uri
+        :param labels:              model endpoint labels, should be list of str or mapping of str:str
+        :param model_endpoint_creation_strategy:   Strategy for creating or updating the model endpoint:
+
+                            * **overwrite**: If model endpoints with the same name exist, delete the `latest` one;
+                              create a new model endpoint entry and set it as `latest`.
+
+                            * **inplace** (default): If model endpoints with the same name exist, update the `latest`
+                              entry; otherwise, create a new entry.
+
+                            * **archive**: If model endpoints with the same name exist, preserve them;
+                              create a new model endpoint with the same name and set it to `latest`.
+
+        :param inputs:              list of the model inputs (e.g. features) ,if provided will override the inputs
                                       that been configured in the model artifact, please note that those inputs need to
                                       be equal in length and order to the inputs that model_class predict method expects
-          :param outputs:             list of the model outputs (e.g. labels) ,if provided will override the outputs
+        :param outputs:             list of the model outputs (e.g. labels) ,if provided will override the outputs
                                       that been configured in the model artifact, please note that those outputs need to
                                       be equal to the model_class predict method outputs (length, and order)
-          :param input_path:          when specified selects the key/path in the event to use as model monitoring inputs
+
+                                      When using LLModel, the output will be overridden with UsageResponseKeys.fields().
+
+        :param input_path:          when specified selects the key/path in the event to use as model monitoring inputs
                                       this require that the event body will behave like a dict, expects scopes to be
                                       defined by dot notation (e.g "data.d").
                                       examples: input_path="data.b"
@@ -1618,7 +2127,7 @@ class ModelRunnerStep(MonitoredStep):
                                       be {"f0": [1, 2]}.
                                       if a ``list`` or ``list of lists`` is provided, it must follow the order and
                                       size defined by the input schema.
-          :param result_path:         when specified selects the key/path in the output event to use as model monitoring
+        :param result_path:         when specified selects the key/path in the output event to use as model monitoring
                                       outputs this require that the output event body will behave like a dict,
                                       expects scopes to be defined by dot notation (e.g "data.d").
                                       examples: result_path="out.b"
@@ -1629,14 +2138,22 @@ class ModelRunnerStep(MonitoredStep):
                                       if a ``list`` or ``list of lists`` is provided, it must follow the order and
                                       size defined by the output schema.
 
-          :param override:            bool allow override existing model on the current ModelRunnerStep.
-          :param model_parameters:    Parameters for model instantiation
+        :param override:            bool allow override existing model on the current ModelRunnerStep.
+        :param model_parameters:    Parameters for model instantiation
         """
         if isinstance(model_class, Model) and model_parameters:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "Cannot provide a model object as argument to `model_class` and also provide `model_parameters`."
             )
-
+        if type(model_class) is LLModel or (
+            isinstance(model_class, str)
+            and model_class.split(".")[-1] == LLModel.__name__
+        ):
+            if outputs:
+                warnings.warn(
+                    "LLModel with existing outputs detected, overriding to default"
+                )
+            outputs = UsageResponseKeys.fields()
         model_parameters = model_parameters or (
             model_class.to_dict() if isinstance(model_class, Model) else {}
         )
@@ -1652,8 +2169,6 @@ class ModelRunnerStep(MonitoredStep):
             except mlrun.errors.MLRunNotFoundError:
                 raise mlrun.errors.MLRunInvalidArgumentError("Artifact not found.")
 
-        outputs = outputs or self._get_model_output_schema(model_artifact)
-
         model_artifact = (
             model_artifact.uri
             if isinstance(model_artifact, mlrun.artifacts.Artifact)
@@ -1719,28 +2234,13 @@ class ModelRunnerStep(MonitoredStep):
         self.class_args[schemas.ModelRunnerStepData.MONITORING_DATA] = monitoring_data
 
     @staticmethod
-    def _get_model_output_schema(
-        model_artifact: Union[ModelArtifact, LLMPromptArtifact],
-    ) -> Optional[list[str]]:
-        if isinstance(
-            model_artifact,
-            ModelArtifact,
-        ):
-            return [feature.name for feature in model_artifact.spec.outputs]
-        elif isinstance(
-            model_artifact,
-            LLMPromptArtifact,
-        ):
-            _model_artifact = model_artifact.model_artifact
-            return [feature.name for feature in _model_artifact.spec.outputs]
-
-    @staticmethod
-    def _get_model_endpoint_output_schema(
+    def _get_model_endpoint_schema(
         name: str,
         project: str,
         uid: str,
-    ) -> list[str]:
+    ) -> tuple[list[str], list[str]]:
         output_schema = None
+        input_schema = None
         try:
             model_endpoint: mlrun.common.schemas.model_monitoring.ModelEndpoint = (
                 mlrun.db.get_run_db().get_model_endpoint(
@@ -1751,6 +2251,7 @@ class ModelRunnerStep(MonitoredStep):
                 )
             )
             output_schema = model_endpoint.spec.label_names
+            input_schema = model_endpoint.spec.feature_names
         except (
             mlrun.errors.MLRunNotFoundError,
             mlrun.errors.MLRunInvalidArgumentError,
@@ -1759,7 +2260,7 @@ class ModelRunnerStep(MonitoredStep):
                 f"Model endpoint not found, using default output schema for model {name}",
                 error=f"{type(ex).__name__}: {ex}",
             )
-        return output_schema
+        return input_schema, output_schema
 
     def _calculate_monitoring_data(self) -> dict[str, dict[str, str]]:
         monitoring_data = deepcopy(
@@ -1775,47 +2276,154 @@ class ModelRunnerStep(MonitoredStep):
                 monitoring_data[model][schemas.MonitoringData.RESULT_PATH] = split_path(
                     monitoring_data[model][schemas.MonitoringData.RESULT_PATH]
                 )
+
+                mep_output_schema, mep_input_schema = None, None
+
+                output_schema = self.class_args[
+                    mlrun.common.schemas.ModelRunnerStepData.MONITORING_DATA
+                ][model][schemas.MonitoringData.OUTPUTS]
+                input_schema = self.class_args[
+                    mlrun.common.schemas.ModelRunnerStepData.MONITORING_DATA
+                ][model][schemas.MonitoringData.INPUTS]
+                if not output_schema or not input_schema:
+                    # if output or input schema is not provided, try to get it from the model endpoint
+                    mep_input_schema, mep_output_schema = (
+                        self._get_model_endpoint_schema(
+                            model,
+                            self.context.project,
+                            monitoring_data[model].get(
+                                schemas.MonitoringData.MODEL_ENDPOINT_UID, ""
+                            ),
+                        )
+                    )
+                self.class_args[
+                    mlrun.common.schemas.ModelRunnerStepData.MONITORING_DATA
+                ][model][schemas.MonitoringData.OUTPUTS] = (
+                    output_schema or mep_output_schema
+                )
+                self.class_args[
+                    mlrun.common.schemas.ModelRunnerStepData.MONITORING_DATA
+                ][model][schemas.MonitoringData.INPUTS] = (
+                    input_schema or mep_input_schema
+                )
             return monitoring_data
         else:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "Monitoring data must be a dictionary."
             )
 
+    def configure_pool_resource(
+        self,
+        max_processes: Optional[int] = None,
+        max_threads: Optional[int] = None,
+        pool_factor: Optional[int] = None,
+    ) -> None:
+        """
+        Configure the resource limits for the shared models in the graph.
+
+        :param max_processes: Maximum number of processes to spawn (excluding dedicated processes).
+            Defaults to the number of CPUs or 16 if undetectable.
+        :param max_threads: Maximum number of threads to spawn. Defaults to 32.
+        :param pool_factor: Multiplier to scale the number of process/thread workers per runnable. Defaults to 1.
+        """
+        self.max_processes = max_processes
+        self.max_threads = max_threads
+        self.pool_factor = pool_factor
+
     def init_object(self, context, namespace, mode="sync", reset=False, **extra_kwargs):
         self.context = context
         if not self._is_local_function(context):
             # skip init of non local functions
             return
-        model_selector = self.class_args.get("model_selector")
+        model_selector, model_selector_params = self.class_args.get(
+            "model_selector", (None, None)
+        )
+        model_runner_selector, model_runner_selector_params = self.class_args.get(
+            "model_runner_selector", (None, None)
+        )
         execution_mechanism_by_model_name = self.class_args.get(
             schemas.ModelRunnerStepData.MODEL_TO_EXECUTION_MECHANISM
         )
         models = self.class_args.get(schemas.ModelRunnerStepData.MODELS, {})
-        if isinstance(model_selector, str):
-            model_selector = get_class(model_selector, namespace)()
+        if model_selector:
+            model_selector = get_class(model_selector, namespace).from_dict(
+                model_selector_params, init_with_params=True
+            )
+            model_runner_selector = (
+                self._convert_model_selector_to_model_runner_selector(
+                    model_selector=model_selector
+                )
+            )
+        elif model_runner_selector:
+            model_runner_selector = get_class(
+                model_runner_selector, namespace
+            ).from_dict(model_runner_selector_params, init_with_params=True)
         model_objects = []
         for model, model_params in models.values():
+            model_name = model_params.get("name")
             model_params[schemas.MonitoringData.INPUT_PATH] = (
                 self.class_args.get(
                     mlrun.common.schemas.ModelRunnerStepData.MONITORING_DATA, {}
                 )
-                .get(model_params.get("name"), {})
+                .get(model_name, {})
                 .get(schemas.MonitoringData.INPUT_PATH)
             )
+            model_params[schemas.MonitoringData.RESULT_PATH] = (
+                self.class_args.get(
+                    mlrun.common.schemas.ModelRunnerStepData.MONITORING_DATA, {}
+                )
+                .get(model_name, {})
+                .get(schemas.MonitoringData.RESULT_PATH)
+            )
             model = get_class(model, namespace).from_dict(
                 model_params, init_with_params=True
             )
             model._raise_exception = False
+            model._execution_mechanism = execution_mechanism_by_model_name.get(
+                model_name
+            )
             model_objects.append(model)
         self._async_object = ModelRunner(
-            model_selector=model_selector,
+            model_runner_selector=model_runner_selector,
             runnables=model_objects,
             execution_mechanism_by_runnable_name=execution_mechanism_by_model_name,
             shared_proxy_mapping=self._shared_proxy_mapping or None,
             name=self.name,
             context=context,
+            max_processes=self.max_processes,
+            max_threads=self.max_threads,
+            pool_factor=self.pool_factor,
+            raise_exception=self.raise_exception,
+            **extra_kwargs,
         )
 
+    def _convert_model_selector_to_model_runner_selector(
+        self,
+        model_selector,
+    ) -> "ModelRunnerSelector":
+        """
+        Wrap a ModelSelector into a ModelRunnerSelector for backward compatibility.
+        """
+
+        class Adapter(ModelRunnerSelector):
+            def __init__(self):
+                self.selector = model_selector
+
+            def select_models(
+                self, event, available_models
+            ) -> Union[list[str], list[Model]]:
+                # Call old ModelSelector logic
+                return self.selector.select(event, available_models)
+
+            def select_outlets(
+                self,
+                event,
+            ) -> Optional[list[str]]:
+                # By default, return all outlets (old ModelSelector didn't control routing)
+                return None
+
+        return Adapter()
+
 
 class ModelRunnerErrorRaiser(storey.MapClass):
     def __init__(self, raise_exception: bool, models_names: list[str], **kwargs):
@@ -1828,11 +2436,15 @@ class ModelRunnerErrorRaiser(storey.MapClass):
             errors = {}
             should_raise = False
             if len(self._models_names) == 1:
-                should_raise = event.body.get("error") is not None
-                errors[self._models_names[0]] = event.body.get("error")
+                if isinstance(event.body, dict):
+                    should_raise = event.body.get("error") is not None
+                    errors[self._models_names[0]] = event.body.get("error")
             else:
                 for model in event.body:
-                    errors[model] = event.body.get(model).get("error")
+                    body_by_model = event.body.get(model)
+                    errors[model] = None
+                    if isinstance(body_by_model, dict):
+                        errors[model] = body_by_model.get("error")
                     if errors[model] is not None:
                         should_raise = True
             if should_raise:
@@ -1902,6 +2514,8 @@ class QueueStep(BaseStep, StepToDict):
         model_endpoint_creation_strategy: Optional[
             schemas.ModelEndpointCreationStrategy
         ] = None,
+        cycle_to: Optional[list[str]] = None,
+        max_iterations: Optional[int] = None,
         **class_args,
     ):
         if not function:
@@ -1919,6 +2533,8 @@ class QueueStep(BaseStep, StepToDict):
             input_path,
             result_path,
             model_endpoint_creation_strategy,
+            cycle_to,
+            max_iterations,
             **class_args,
         )
 
@@ -1954,8 +2570,10 @@ class FlowStep(BaseStep):
         after: Optional[list] = None,
         engine=None,
         final_step=None,
+        allow_cyclic: bool = False,
+        max_iterations: Optional[int] = None,
     ):
-        super().__init__(name, after)
+        super().__init__(name, after, max_iterations=max_iterations)
         self._steps = None
         self.steps = steps
         self.engine = engine
@@ -1967,6 +2585,7 @@ class FlowStep(BaseStep):
         self._wait_for_result = False
         self._source = None
         self._start_steps = []
+        self._allow_cyclic = allow_cyclic
 
     def get_children(self):
         return self._steps.values()
@@ -2000,6 +2619,8 @@ class FlowStep(BaseStep):
         model_endpoint_creation_strategy: Optional[
             schemas.ModelEndpointCreationStrategy
         ] = None,
+        cycle_to: Optional[list[str]] = None,
+        max_iterations: Optional[int] = None,
         **class_args,
     ):
         """add task, queue or router step/class to the flow
@@ -2033,21 +2654,17 @@ class FlowStep(BaseStep):
                             to event["y"] resulting in {"x": 5, "y": <result>}
         :param model_endpoint_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`.
+                             * **overwrite**: If model endpoints with the same name exist, delete the `latest` one;
+                                create a new model endpoint entry and set it as `latest`.
 
-                            * **inplace** (default):
+                            * **inplace** (default): If model endpoints with the same name exist, update the `latest`
+                                entry; otherwise, create a new entry.
 
-                            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`.
+                            * **archive**: If model endpoints with the same name exist, preserve them;
+                              create a new model endpoint with the same name and set it to `latest`.
 
+        :param cycle_to:    list of step names to create a cycle to (for cyclic graphs)
+        :param max_iterations: maximum number of iterations for this step in case of a cycle graph
         :param class_args:  class init arguments
         """
 
@@ -2073,6 +2690,8 @@ class FlowStep(BaseStep):
         after_list = after if isinstance(after, list) else [after]
         for after in after_list:
             self.insert_step(name, step, after, before)
+        step.cycle_to(cycle_to or [])
+        step._max_iterations = max_iterations
         return step
 
     def insert_step(self, key, step, after, before=None):
@@ -2165,13 +2784,24 @@ class FlowStep(BaseStep):
         for step in self._steps.values():
             step._next = None
             step._visited = False
-            if step.after:
+            if step.after and not step.cycle_from:
+                has_illegal_branches = len(step.after) > 1 and self.engine == "sync"
+                if has_illegal_branches:
+                    raise GraphError(
+                        f"synchronous flow engine doesnt support branches use async for step {step.name}"
+                    )
                 loop_step = has_loop(step, [])
-                if loop_step:
+                if loop_step and not self.allow_cyclic:
                     raise GraphError(
                         f"Error, loop detected in step {loop_step}, graph must be acyclic (DAG)"
                     )
-            else:
+            elif (
+                step.after
+                and step.cycle_from
+                and set(step.after) == set(step.cycle_from)
+            ):
+                start_steps.append(step.name)
+            elif not step.cycle_from:
                 start_steps.append(step.name)
 
         responders = []
@@ -2268,6 +2898,9 @@ class FlowStep(BaseStep):
         def process_step(state, step, root):
             if not state._is_local_function(self.context) or state._visited:
                 return
+            state._visited = (
+                True  # mark visited to avoid re-visit in case of multiple uplinks
+            )
             for item in state.next or []:
                 next_state = root[item]
                 if next_state.async_object:
@@ -2278,7 +2911,7 @@ class FlowStep(BaseStep):
             )
 
         default_source, self._wait_for_result = _init_async_objects(
-            self.context, self._steps.values()
+            self.context, self._steps.values(), self
         )
 
         source = self._source or default_source
@@ -2509,6 +3142,8 @@ class RootFlowStep(FlowStep):
         "shared_models",
         "shared_models_mechanism",
         "pool_factor",
+        "allow_cyclic",
+        "max_iterations",
     ]
 
     def __init__(
@@ -2518,13 +3153,11 @@ class RootFlowStep(FlowStep):
         after: Optional[list] = None,
         engine=None,
         final_step=None,
+        allow_cyclic: bool = False,
+        max_iterations: Optional[int] = 10_000,
     ):
         super().__init__(
-            name,
-            steps,
-            after,
-            engine,
-            final_step,
+            name, steps, after, engine, final_step, allow_cyclic, max_iterations
         )
         self._models = set()
         self._route_models = set()
@@ -2535,48 +3168,102 @@ class RootFlowStep(FlowStep):
         self._shared_max_threads = None
         self._pool_factor = None
 
+    @property
+    def max_iterations(self) -> int:
+        return self._max_iterations
+
+    @max_iterations.setter
+    def max_iterations(self, max_iterations: int):
+        self._max_iterations = max_iterations
+
+    @property
+    def allow_cyclic(self) -> bool:
+        return self._allow_cyclic
+
+    @allow_cyclic.setter
+    def allow_cyclic(self, allow_cyclic: bool):
+        self._allow_cyclic = allow_cyclic
+
     def add_shared_model(
         self,
         name: str,
         model_class: Union[str, Model],
         execution_mechanism: Union[str, ParallelExecutionMechanisms],
         model_artifact: Union[str, ModelArtifact],
+        inputs: Optional[list[str]] = None,
+        outputs: Optional[list[str]] = None,
+        input_path: Optional[str] = None,
+        result_path: Optional[str] = None,
         override: bool = False,
         **model_parameters,
     ) -> None:
         """
         Add a shared model to the graph, this model will be available to all the ModelRunners in the graph
         :param name:                Name of the shared model (should be unique in the graph)
-        :param model_class:         Model class name
+        :param model_class:         Model class name. If LLModel is chosen
+        (either by name `LLModel` or by its full path, e.g. mlrun.serving.states.LLModel),
+        outputs will be overridden with UsageResponseKeys fields.
         :param execution_mechanism: Parallel execution mechanism to be used to execute this model. Must be one of:
-            * "process_pool" – To run in a separate process from a process pool. This is appropriate for CPU or GPU
+
+            * **process_pool**: To run in a separate process from a process pool. This is appropriate for CPU or GPU
                 intensive tasks as they would otherwise block the main process by holding Python's Global Interpreter
                 Lock (GIL).
-            * "dedicated_process" – To run in a separate dedicated process. This is appropriate for CPU or GPU intensive
-                tasks that also require significant Runnable-specific initialization (e.g. a large model).
-            * "thread_pool" – To run in a separate thread. This is appropriate for blocking I/O tasks, as they would
+
+            * **dedicated_process**: To run in a separate dedicated process. This is appropriate for CPU or GPU
+                intensive tasks that also require significant Runnable-specific initialization (e.g. a large model).
+
+            * **thread_pool**: To run in a separate thread. This is appropriate for blocking I/O tasks, as they would
                 otherwise block the main event loop thread.
-            * "asyncio" – To run in an asyncio task. This is appropriate for I/O tasks that use asyncio, allowing the
+
+            * **asyncio**: To run in an asyncio task. This is appropriate for I/O tasks that use asyncio, allowing the
                 event loop to continue running while waiting for a response.
-            * "shared_executor" – Reuses an external executor (typically managed by the flow or context) to execute the
-                runnable. Should be used only if you have multiply `ParallelExecution` in the same flow and especially
-                useful when:
+
+            * **shared_executor**:  Reuses an external executor (typically managed by the flow or context) to execute
+                the runnable. Should be used only if you have multiple `ParallelExecution` in the same flow and
+                especially useful when:
+
                 - You want to share a heavy resource like a large model loaded onto a GPU.
+
                 - You want to centralize task scheduling or coordination for multiple lightweight tasks.
+
                 - You aim to minimize overhead from creating new executors or processes/threads per runnable.
+
                 The runnable is expected to be pre-initialized and reused across events, enabling efficient use of
                 memory and hardware accelerators.
-            * "naive" – To run in the main event loop. This is appropriate only for trivial computation and/or file I/O.
-                It means that the runnable will not actually be run in parallel to anything else.
 
-            :param model_artifact:      model artifact or mlrun model artifact uri
-            :param override:            bool allow override existing model on the current ModelRunnerStep.
-            :param model_parameters:    Parameters for model instantiation
+            * **naive**: To run in the main event loop. This is appropriate only for trivial computation and/or file
+                I/O. It means that the runnable will not actually be run in parallel to anything else.
+
+        :param model_artifact:      model artifact or mlrun model artifact uri
+        :param inputs:              list of the model inputs (e.g. features) ,if provided will override the inputs
+                                    that been configured in the model artifact, please note that those inputs need
+                                    to be equal in length and order to the inputs that model_class
+                                    predict method expects
+        :param outputs:             list of the model outputs (e.g. labels) ,if provided will override the outputs
+                                    that been configured in the model artifact, please note that those outputs need
+                                    to be equal to the model_class
+                                    predict method outputs (length, and order)
+        :param input_path:          input path inside the user event, expect scopes to be defined by dot notation
+                                    (e.g "inputs.my_model_inputs"). expects list or dictionary type object in path.
+        :param result_path:         result path inside the user output event, expect scopes to be defined by dot
+                                    notation (e.g "outputs.my_model_outputs") expects list or dictionary type object
+                                    in path.
+        :param override:            bool allow override existing model on the current ModelRunnerStep.
+        :param model_parameters:    Parameters for model instantiation
         """
         if isinstance(model_class, Model) and model_parameters:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "Cannot provide a model object as argument to `model_class` and also provide `model_parameters`."
             )
+        if type(model_class) is LLModel or (
+            isinstance(model_class, str)
+            and model_class.split(".")[-1] == LLModel.__name__
+        ):
+            if outputs:
+                warnings.warn(
+                    "LLModel with existing outputs detected, overriding to default"
+                )
+            outputs = UsageResponseKeys.fields()
 
         if execution_mechanism == ParallelExecutionMechanisms.shared_executor:
             raise mlrun.errors.MLRunInvalidArgumentError(
@@ -2604,6 +3291,14 @@ class RootFlowStep(FlowStep):
                 "Inconsistent name for the added model."
             )
         model_parameters["name"] = name
+        model_parameters["inputs"] = inputs or model_parameters.get("inputs", [])
+        model_parameters["outputs"] = outputs or model_parameters.get("outputs", [])
+        model_parameters["input_path"] = input_path or model_parameters.get(
+            "input_path"
+        )
+        model_parameters["result_path"] = result_path or model_parameters.get(
+            "result_path"
+        )
 
         if name in self.shared_models and not override:
             raise mlrun.errors.MLRunInvalidArgumentError(
@@ -2618,7 +3313,9 @@ class RootFlowStep(FlowStep):
         self.shared_models[name] = (model_class, model_parameters)
         self.shared_models_mechanism[name] = execution_mechanism
 
-    def get_shared_model_name_by_artifact_uri(self, artifact_uri: str) -> Optional[str]:
+    def get_shared_model_by_artifact_uri(
+        self, artifact_uri: str
+    ) -> Union[tuple[str, str, dict], tuple[None, None, None]]:
         """
         Get a shared model by its artifact URI.
         :param artifact_uri: The artifact URI of the model.
@@ -2626,10 +3323,10 @@ class RootFlowStep(FlowStep):
         """
         for model_name, (model_class, model_params) in self.shared_models.items():
             if model_params.get("artifact_uri") == artifact_uri:
-                return model_name
-        return None
+                return model_name, model_class, model_params
+        return None, None, None
 
-    def config_pool_resource(
+    def configure_shared_pool_resource(
         self,
         max_processes: Optional[int] = None,
         max_threads: Optional[int] = None,
@@ -2637,8 +3334,9 @@ class RootFlowStep(FlowStep):
     ) -> None:
         """
         Configure the resource limits for the shared models in the graph.
+
         :param max_processes: Maximum number of processes to spawn (excluding dedicated processes).
-                             Defaults to the number of CPUs or 16 if undetectable.
+            Defaults to the number of CPUs or 16 if undetectable.
         :param max_threads: Maximum number of threads to spawn. Defaults to 32.
         :param pool_factor: Multiplier to scale the number of process/thread workers per runnable. Defaults to 1.
         """
@@ -2677,6 +3375,7 @@ class RootFlowStep(FlowStep):
                     model_params, init_with_params=True
                 )
                 model._raise_exception = False
+                model._execution_mechanism = self._shared_models_mechanism[model.name]
                 self.context.executor.add_runnable(
                     model, self._shared_models_mechanism[model.name]
                 )
@@ -2796,12 +3495,10 @@ def _add_graphviz_router(graph, step, source=None, **kwargs):
         graph.edge(step.fullname, route.fullname)
 
 
-def _add_graphviz_model_runner(graph, step, source=None):
+def _add_graphviz_model_runner(graph, step, source=None, is_monitored=False):
     if source:
         graph.node("_start", source.name, shape=source.shape, style="filled")
         graph.edge("_start", step.fullname)
-
-    is_monitored = step._extract_root_step().track_models
     m_cell = '<FONT POINT-SIZE="9">🄼</FONT>' if is_monitored else ""
 
     number_of_models = len(
@@ -2840,6 +3537,7 @@ def _add_graphviz_flow(
         allow_empty=True
     )
     graph.node("_start", source.name, shape=source.shape, style="filled")
+    is_monitored = step.track_models if isinstance(step, RootFlowStep) else False
     for start_step in start_steps:
         graph.edge("_start", start_step.fullname)
     for child in step.get_children():
@@ -2848,7 +3546,7 @@ def _add_graphviz_flow(
             with graph.subgraph(name="cluster_" + child.fullname) as sg:
                 _add_graphviz_router(sg, child)
         elif kind == StepKinds.model_runner:
-            _add_graphviz_model_runner(graph, child)
+            _add_graphviz_model_runner(graph, child, is_monitored=is_monitored)
         else:
             graph.node(child.fullname, label=child.name, shape=child.get_shape())
         _add_edges(child.after or [], step, graph, child)
@@ -3034,7 +3732,7 @@ def params_to_step(
     return name, step
 
 
-def _init_async_objects(context, steps):
+def _init_async_objects(context, steps, root):
     try:
         import storey
     except ImportError:
@@ -3049,6 +3747,7 @@ def _init_async_objects(context, steps):
 
     for step in steps:
         if hasattr(step, "async_object") and step._is_local_function(context):
+            max_iterations = step._max_iterations or root.max_iterations
             if step.kind == StepKinds.queue:
                 skip_stream = context.is_mock and step.next
                 if step.path and not skip_stream:
@@ -3067,23 +3766,25 @@ def _init_async_objects(context, steps):
                         datastore_profile = datastore_profile_read(stream_path)
                         if isinstance(
                             datastore_profile,
-                            (DatastoreProfileKafkaTarget, DatastoreProfileKafkaSource),
+                            DatastoreProfileKafkaTarget | DatastoreProfileKafkaStream,
                         ):
                             step._async_object = KafkaStoreyTarget(
                                 path=stream_path,
                                 context=context,
+                                max_iterations=max_iterations,
                                 **options,
                             )
                         elif isinstance(datastore_profile, DatastoreProfileV3io):
                             step._async_object = StreamStoreyTarget(
                                 stream_path=stream_path,
                                 context=context,
+                                max_iterations=max_iterations,
                                 **options,
                             )
                         else:
                             raise mlrun.errors.MLRunValueError(
                                 f"Received an unexpected stream profile type: {type(datastore_profile)}\n"
-                                "Expects `DatastoreProfileV3io` or `DatastoreProfileKafkaSource`."
+                                "Expects `DatastoreProfileV3io` or `DatastoreProfileKafkaStream`."
                             )
                     elif stream_path.startswith("kafka://") or kafka_brokers:
                         topic, brokers = parse_kafka_url(stream_path, kafka_brokers)
@@ -3097,6 +3798,13 @@ def _init_async_objects(context, steps):
                             brokers=brokers,
                             producer_options=kafka_producer_options,
                             context=context,
+                            max_iterations=max_iterations,
+                            **options,
+                        )
+                    elif stream_path.startswith("dummy://"):
+                        step._async_object = _DummyStream(
+                            context=context,
+                            max_iterations=max_iterations,
                             **options,
                         )
                     else:
@@ -3107,10 +3815,14 @@ def _init_async_objects(context, steps):
                             storey.V3ioDriver(endpoint or config.v3io_api),
                             stream_path,
                             context=context,
+                            max_iterations=max_iterations,
                             **options,
                         )
                 else:
-                    step._async_object = storey.Map(lambda x: x)
+                    step._async_object = storey.Map(
+                        lambda x: x,
+                        max_iterations=max_iterations,
+                    )
 
             elif not step.async_object or not hasattr(step.async_object, "_outlets"):
                 # if regular class, wrap with storey Map
@@ -3122,6 +3834,8 @@ def _init_async_objects(context, steps):
                     name=step.name,
                     context=context,
                     pass_context=step._inject_context,
+                    fn_select_outlets=step._outlets_selector,
+                    max_iterations=max_iterations,
                 )
             if (
                 respond_supported