mlrun 1.6.4rc8__py3-none-any.whl → 1.7.0__py3-none-any.whl

mlrun/serving/states.py

@@ -12,20 +12,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__all__ = ["TaskStep", "RouterStep", "RootFlowStep", "ErrorStep"]
+__all__ = [
+    "TaskStep",
+    "RouterStep",
+    "RootFlowStep",
+    "ErrorStep",
+    "MonitoringApplicationStep",
+]
 
 import os
 import pathlib
 import traceback
 from copy import copy, deepcopy
 from inspect import getfullargspec, signature
-from typing import Union
+from typing import Any, Union
+
+import storey.utils
 
 import mlrun
 
 from ..config import config
 from ..datastore import get_stream_pusher
-from ..datastore.utils import parse_kafka_url
+from ..datastore.utils import (
+    get_kafka_brokers_from_dict,
+    parse_kafka_url,
+)
 from ..errors import MLRunInvalidArgumentError, err_to_str
 from ..model import ModelObj, ObjectDict
 from ..platforms.iguazio import parse_path
@@ -52,6 +63,7 @@ class StepKinds:
     choice = "choice"
     root = "root"
     error_step = "error_step"
+    monitoring_application = "monitoring_application"
 
 
 _task_step_fields = [
@@ -72,6 +84,9 @@ _task_step_fields = [
 ]
 
 
+MAX_ALLOWED_STEPS = 4500
+
+
 def new_model_endpoint(class_name, model_path, handler=None, **class_args):
     class_args = deepcopy(class_args)
     class_args["model_path"] = model_path
@@ -324,7 +339,7 @@ class BaseStep(ModelObj):
             parent = self._parent
         else:
             raise GraphError(
-                f"step {self.name} parent is not set or its not part of a graph"
+                f"step {self.name} parent is not set or it's not part of a graph"
             )
 
         name, step = params_to_step(
@@ -346,6 +361,39 @@ class BaseStep(ModelObj):
         parent._last_added = step
         return step
 
+    def set_flow(
+        self,
+        steps: list[Union[str, StepToDict, dict[str, Any]]],
+        force: bool = False,
+    ):
+        """set list of steps as downstream from this step, in the order specified. This will overwrite any existing
+        downstream steps.
+
+        :param steps: list of steps to follow this one
+        :param force: whether to overwrite existing downstream steps. If False, this method will fail if any downstream
+        steps have already been defined. Defaults to False.
+        :return: the last step added to the flow
+
+        example:
+            The below code sets the downstream nodes of step1 by using a list of steps (provided to `set_flow()`) and a
+            single step (provided to `to()`), resulting in the graph (step1 -> step2 -> step3 -> step4).
+            Notice that using `force=True` is required in case step1 already had downstream nodes (e.g. if the existing
+            graph is step1 -> step2_old) and that following the execution of this code the existing downstream steps
+            are removed. If the intention is to split the graph (and not to overwrite), please use `to()`.
+
+            step1.set_flow(
+                [
+                    dict(name="step2", handler="step2_handler"),
+                    dict(name="step3", class_name="Step3Class"),
+                ],
+                force=True,
+            ).to(dict(name="step4", class_name="Step4Class"))
+        """
+        raise NotImplementedError("set_flow() can only be called on a FlowStep")
+
+    def supports_termination(self):
+        return False
+
 
 class TaskStep(BaseStep):
     """task execution step, runs a class or handler"""
@@ -452,13 +500,15 @@ class TaskStep(BaseStep):
                 class_args[key] = arg
         class_args.update(extra_kwargs)
 
-        # add common args (name, context, ..) only if target class can accept them
-        argspec = getfullargspec(class_object)
-        for key in ["name", "context", "input_path", "result_path", "full_event"]:
-            if argspec.varkw or key in argspec.args:
-                class_args[key] = getattr(self, key)
-        if argspec.varkw or "graph_step" in argspec.args:
-            class_args["graph_step"] = self
+        if not isinstance(self, MonitoringApplicationStep):
+            # add common args (name, context, ..) only if target class can accept them
+            argspec = getfullargspec(class_object)
+
+            for key in ["name", "context", "input_path", "result_path", "full_event"]:
+                if argspec.varkw or key in argspec.args:
+                    class_args[key] = getattr(self, key)
+            if argspec.varkw or "graph_step" in argspec.args:
+                class_args["graph_step"] = self
         return class_args
 
     def get_step_class_object(self, namespace):
@@ -549,6 +599,39 @@ class TaskStep(BaseStep):
         return event
 
 
+class MonitoringApplicationStep(TaskStep):
+    """monitoring application execution step, runs users class code"""
+
+    kind = "monitoring_application"
+    _default_class = ""
+
+    def __init__(
+        self,
+        class_name: Union[str, type] = None,
+        class_args: dict = None,
+        handler: str = None,
+        name: str = None,
+        after: list = None,
+        full_event: bool = None,
+        function: str = None,
+        responder: bool = None,
+        input_path: str = None,
+        result_path: str = None,
+    ):
+        super().__init__(
+            class_name=class_name,
+            class_args=class_args,
+            handler=handler,
+            name=name,
+            after=after,
+            full_event=full_event,
+            function=function,
+            responder=responder,
+            input_path=input_path,
+            result_path=result_path,
+        )
+
+
 class ErrorStep(TaskStep):
     """error execution step, runs a class or handler"""
 
@@ -556,6 +639,34 @@ class ErrorStep(TaskStep):
     _dict_fields = _task_step_fields + ["before", "base_step"]
     _default_class = ""
 
+    def __init__(
+        self,
+        class_name: Union[str, type] = None,
+        class_args: dict = None,
+        handler: str = None,
+        name: str = None,
+        after: list = None,
+        full_event: bool = None,
+        function: str = None,
+        responder: bool = None,
+        input_path: str = None,
+        result_path: str = None,
+    ):
+        super().__init__(
+            class_name=class_name,
+            class_args=class_args,
+            handler=handler,
+            name=name,
+            after=after,
+            full_event=full_event,
+            function=function,
+            responder=responder,
+            input_path=input_path,
+            result_path=result_path,
+        )
+        self.before = None
+        self.base_step = None
+
 
 class RouterStep(TaskStep):
     """router step, implement routing logic for running child routes"""
@@ -625,6 +736,11 @@ class RouterStep(TaskStep):
         if not route:
             route = TaskStep(class_name, class_args, handler=handler)
         route.function = function or route.function
+
+        if len(self._routes) >= MAX_ALLOWED_STEPS:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"Cannot create the serving graph: the maximum number of steps is {MAX_ALLOWED_STEPS}"
+            )
         route = self._routes.update(key, route)
         route.set_parent(self)
         return route
@@ -729,13 +845,45 @@ class QueueStep(BaseStep):
     def async_object(self):
         return self._async_object
 
+    def to(
+        self,
+        class_name: Union[str, StepToDict] = None,
+        name: str = None,
+        handler: str = None,
+        graph_shape: str = None,
+        function: str = None,
+        full_event: bool = None,
+        input_path: str = None,
+        result_path: str = None,
+        **class_args,
+    ):
+        if not function:
+            name = get_name(name, class_name)
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"step '{name}' must specify a function, because it follows a queue step"
+            )
+        return super().to(
+            class_name,
+            name,
+            handler,
+            graph_shape,
+            function,
+            full_event,
+            input_path,
+            result_path,
+            **class_args,
+        )
+
     def run(self, event, *args, **kwargs):
         data = event.body
         if not data:
             return event
 
         if self._stream:
-            self._stream.push({"id": event.id, "body": data, "path": event.path})
+            full_event = self.options.get("full_event")
+            if full_event or full_event is None and self.next:
+                data = storey.utils.wrap_event_for_serialization(event, data)
+            self._stream.push(data)
             event.terminated = True
             event.body = None
         return event
@@ -1132,23 +1280,17 @@ class FlowStep(BaseStep):
         if self._controller:
             # async flow (using storey)
             event._awaitable_result = None
-            if config.datastore.async_source_mode == "enabled":
-                resp_awaitable = self._controller.emit(
-                    event, await_result=self._wait_for_result
-                )
-                if self._wait_for_result:
-                    return resp_awaitable
-                return self._await_and_return_id(resp_awaitable, event)
-            else:
-                resp = self._controller.emit(
-                    event, return_awaitable_result=self._wait_for_result
-                )
-                if self._wait_for_result and resp:
-                    return resp.await_result()
+            resp = self._controller.emit(
+                event, return_awaitable_result=self._wait_for_result
+            )
+            if self._wait_for_result and resp:
+                return resp.await_result()
             event = copy(event)
             event.body = {"id": event.id}
             return event
 
+        event = storey.utils.unpack_event_if_wrapped(event)
+
         if len(self._start_steps) == 0:
             return event
         next_obj = self._start_steps[0]
@@ -1182,6 +1324,7 @@ class FlowStep(BaseStep):
 
     def wait_for_completion(self):
         """wait for completion of run in async flows"""
+
         if self._controller:
             if hasattr(self._controller, "terminate"):
                 self._controller.terminate()
@@ -1234,6 +1377,30 @@ class FlowStep(BaseStep):
                 )
             self[step_name].after_step(name)
 
+    def set_flow(
+        self,
+        steps: list[Union[str, StepToDict, dict[str, Any]]],
+        force: bool = False,
+    ):
+        if not force and self.steps:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "set_flow() called on a step that already has downstream steps. "
+                "If you want to overwrite existing steps, set force=True."
+            )
+
+        self.steps = None
+        step = self
+        for next_step in steps:
+            if isinstance(next_step, dict):
+                step = step.to(**next_step)
+            else:
+                step = step.to(next_step)
+
+        return step
+
+    def supports_termination(self):
+        return self.engine != "sync"
+
 
 class RootFlowStep(FlowStep):
     """root flow step"""
@@ -1248,6 +1415,7 @@ classes_map = {
     "flow": FlowStep,
     "queue": QueueStep,
     "error_step": ErrorStep,
+    "monitoring_application": MonitoringApplicationStep,
 }
 
 
@@ -1471,15 +1639,17 @@ def _init_async_objects(context, steps):
                 if step.path and not skip_stream:
                     stream_path = step.path
                     endpoint = None
-                    options = {}
+                    # in case of a queue, we default to a full_event=True
+                    full_event = step.options.get("full_event")
+                    options = {
+                        "full_event": full_event or full_event is None and step.next
+                    }
                     options.update(step.options)
-                    kafka_bootstrap_servers = options.pop(
-                        "kafka_bootstrap_servers", None
-                    )
-                    if stream_path.startswith("kafka://") or kafka_bootstrap_servers:
-                        topic, bootstrap_servers = parse_kafka_url(
-                            stream_path, kafka_bootstrap_servers
-                        )
+
+                    kafka_brokers = get_kafka_brokers_from_dict(options, pop=True)
+
+                    if stream_path.startswith("kafka://") or kafka_brokers:
+                        topic, brokers = parse_kafka_url(stream_path, kafka_brokers)
 
                         kafka_producer_options = options.pop(
                             "kafka_producer_options", None
@@ -1487,7 +1657,7 @@ def _init_async_objects(context, steps):
 
                         step._async_object = storey.KafkaTarget(
                             topic=topic,
-                            bootstrap_servers=bootstrap_servers,
+                            brokers=brokers,
                             producer_options=kafka_producer_options,
                             context=context,
                             **options,
@@ -1514,6 +1684,7 @@ def _init_async_objects(context, steps):
                     result_path=step.result_path,
                     name=step.name,
                     context=context,
+                    pass_context=step._inject_context,
                 )
             if (
                 respond_supported
@@ -1526,9 +1697,11 @@ def _init_async_objects(context, steps):
                 wait_for_result = True
 
     source_args = context.get_param("source_args", {})
+    explicit_ack = (
+        is_explicit_ack_supported(context) and mlrun.mlconf.is_explicit_ack_enabled()
+    )
 
-    explicit_ack = is_explicit_ack_supported(context) and mlrun.mlconf.is_explicit_ack()
-
+    # TODO: Change to AsyncEmitSource once we can drop support for nuclio<1.12.10
     default_source = storey.SyncEmitSource(
         context=context,
         explicit_ack=explicit_ack,

mlrun/serving/utils.py

@@ -46,7 +46,16 @@ def _update_result_body(result_path, event_body, result):
 class StepToDict:
     """auto serialization of graph steps to a python dictionary"""
 
-    def to_dict(self, fields=None, exclude=None):
+    meta_keys = [
+        "context",
+        "name",
+        "input_path",
+        "result_path",
+        "full_event",
+        "kwargs",
+    ]
+
+    def to_dict(self, fields: list = None, exclude: list = None, strip: bool = False):
         """convert the step object to a python dictionary"""
         fields = fields or getattr(self, "_dict_fields", None)
         if not fields:
@@ -54,24 +63,16 @@ class StepToDict:
         if exclude:
             fields = [field for field in fields if field not in exclude]
 
-        meta_keys = [
-            "context",
-            "name",
-            "input_path",
-            "result_path",
-            "full_event",
-            "kwargs",
-        ]
         args = {
             key: getattr(self, key)
             for key in fields
-            if getattr(self, key, None) is not None and key not in meta_keys
+            if getattr(self, key, None) is not None and key not in self.meta_keys
         }
         # add storey kwargs or extra kwargs
         if "kwargs" in fields and (hasattr(self, "kwargs") or hasattr(self, "_kwargs")):
             kwargs = getattr(self, "kwargs", {}) or getattr(self, "_kwargs", {})
             for key, value in kwargs.items():
-                if key not in meta_keys:
+                if key not in self.meta_keys:
                     args[key] = value
 
         mod_name = self.__class__.__module__
@@ -80,7 +81,9 @@ class StepToDict:
             class_path = f"{mod_name}.{class_path}"
         struct = {
             "class_name": class_path,
-            "name": self.name or self.__class__.__name__,
+            "name": self.name
+            if hasattr(self, "name") and self.name
+            else self.__class__.__name__,
             "class_args": args,
         }
         if hasattr(self, "_STEP_KIND"):
@@ -94,8 +97,13 @@ class StepToDict:
         return struct
 
 
+class MonitoringApplicationToDict(StepToDict):
+    _STEP_KIND = "monitoring_application"
+    meta_keys = []
+
+
 class RouterToDict(StepToDict):
     _STEP_KIND = "router"
 
-    def to_dict(self, fields=None, exclude=None):
-        return super().to_dict(exclude=["routes"])
+    def to_dict(self, fields: list = None, exclude: list = None, strip: bool = False):
+        return super().to_dict(exclude=["routes"], strip=strip)

mlrun/serving/v1_serving.py

@@ -18,7 +18,6 @@ import socket
 from copy import deepcopy
 from datetime import datetime
 from io import BytesIO
-from typing import Dict
 from urllib.request import urlopen
 
 import nuclio
@@ -26,7 +25,7 @@ import nuclio
 import mlrun
 from mlrun.errors import err_to_str
 from mlrun.platforms.iguazio import OutputStream
-from mlrun.runtimes import RemoteRuntime
+from mlrun.runtimes.nuclio.function import RemoteRuntime
 
 serving_handler = "handler"
 
@@ -97,16 +96,16 @@ class MLModelServer:
         if not self.ready and not self.model:
             raise ValueError("please specify a load method or a model object")
 
-    def preprocess(self, request: Dict) -> Dict:
+    def preprocess(self, request: dict) -> dict:
         return request
 
-    def postprocess(self, request: Dict) -> Dict:
+    def postprocess(self, request: dict) -> dict:
         return request
 
-    def predict(self, request: Dict) -> Dict:
+    def predict(self, request: dict) -> dict:
         raise NotImplementedError()
 
-    def explain(self, request: Dict) -> Dict:
+    def explain(self, request: dict) -> dict:
         raise NotImplementedError()
 
 
@@ -200,7 +199,7 @@ class _ServerInfo:
 class HTTPHandler:
     kind = ""
 
-    def __init__(self, models: Dict, server: _ServerInfo = None):
+    def __init__(self, models: dict, server: _ServerInfo = None):
         self.models = models
         self.srvinfo = server
         self.context = None