mlrun 1.3.3__py3-none-any.whl → 1.4.0__py3-none-any.whl

mlrun/serving/states.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__all__ = ["TaskStep", "RouterStep", "RootFlowStep"]
+__all__ = ["TaskStep", "RouterStep", "RootFlowStep", "ErrorStep"]
 
 import os
 import pathlib
@@ -28,7 +28,7 @@ from ..errors import MLRunInvalidArgumentError, err_to_str
 from ..model import ModelObj, ObjectDict
 from ..platforms.iguazio import parse_path
 from ..utils import get_class, get_function
-from .utils import _extract_input_data, _update_result_body
+from .utils import StepToDict, _extract_input_data, _update_result_body
 
 callable_prefix = "_"
 path_splitter = "/"
@@ -49,6 +49,7 @@ class StepKinds:
     queue = "queue"
     choice = "choice"
     root = "root"
+    error_step = "error_step"
 
 
 _task_step_fields = [
@@ -134,11 +135,82 @@ class BaseStep(ModelObj):
                 self.after.append(name)
         return self
 
-    def error_handler(self, step_name: str = None):
-        """set error handler step (on failure/raise of this step)"""
-        if not step_name:
-            raise MLRunInvalidArgumentError("Must specify step_name")
-        self.on_error = step_name
+    def error_handler(
+        self,
+        name: str = None,
+        class_name=None,
+        handler=None,
+        before=None,
+        function=None,
+        full_event: bool = None,
+        input_path: str = None,
+        result_path: str = None,
+        **class_args,
+    ):
+        """set error handler on a step or the entire graph (to be executed on failure/raise)
+
+        When setting the error_handler on the graph object, the graph completes after the error handler execution.
+
+        example:
+            in the below example, an 'error_catcher' step is set as the error_handler of the 'raise' step:
+            in case of error/raise in 'raise' step, the handle_error will be run. after that,
+            the 'echo' step will be run.
+            graph = function.set_topology('flow', engine='async')
+            graph.to(name='raise', handler='raising_step')\
+                .error_handler(name='error_catcher', handler='handle_error', full_event=True, before='echo')
+            graph.add_step(name="echo", handler='echo', after="raise").respond()
+
+        :param name:        unique name (and path) for the error handler step, default is class name
+        :param class_name:  class name or step object to build the step from
+                            the error handler step is derived from task step (ie no router/queue functionally)
+        :param handler:     class/function handler to invoke on run/event
+        :param before:      string or list of next step(s) names that will run after this step.
+                            the `before` param must not specify upstream steps as it will cause a loop.
+                            if `before` is not specified, the graph will complete after the error handler execution.
+        :param function:    function this step should run in
+        :param full_event:  this step accepts the full event (not just the body)
+        :param input_path:  selects the key/path in the event to use as input to the step
+                            this requires that the event body will behave like a dict, for example:
+                            event: {"data": {"a": 5, "b": 7}}, input_path="data.b" means the step will
+                            receive 7 as input
+        :param result_path: selects the key/path in the event to write the results to
+                            this requires that the event body will behave like a dict, for example:
+                            event: {"x": 5} , result_path="y" means the output of the step will be written
+                            to event["y"] resulting in {"x": 5, "y": <result>}
+        :param class_args:  class init arguments
+
+        """
+        if not (class_name or handler):
+            raise MLRunInvalidArgumentError("class_name or handler must be provided")
+        if isinstance(self, RootFlowStep) and before:
+            raise MLRunInvalidArgumentError(
+                "`before` arg can't be specified for graph error handler"
+            )
+
+        name = get_name(name, class_name)
+        step = ErrorStep(
+            class_name,
+            class_args,
+            handler,
+            name=name,
+            function=function,
+            full_event=full_event,
+            input_path=input_path,
+            result_path=result_path,
+        )
+        self.on_error = name
+        before = [before] if isinstance(before, str) else before
+        step.before = before or []
+        step.base_step = self.name
+        if hasattr(self, "_parent") and self._parent:
+            # when self is a step
+            step = self._parent._steps.update(name, step)
+            step.set_parent(self._parent)
+        else:
+            # when self is the graph
+            step = self._steps.update(name, step)
+            step.set_parent(self)
+
         return self
 
     def init_object(self, context, namespace, mode="sync", reset=False, **extra_kwargs):
@@ -186,10 +258,11 @@ class BaseStep(ModelObj):
 
     def _call_error_handler(self, event, err, **kwargs):
         """call the error handler if exist"""
-        if self._on_error_handler:
-            event.error = err_to_str(err)
-            event.origin_state = self.fullname
-            return self._on_error_handler(event)
+        if not event.error:
+            event.error = {}
+        event.error[self.name] = err_to_str(err)
+        event.origin_state = self.fullname
+        return self._on_error_handler(event)
 
     def path_to_step(self, path: str):
         """return step object from step relative/fullname"""
@@ -206,7 +279,7 @@ class BaseStep(ModelObj):
 
     def to(
         self,
-        class_name: Union[str, type] = None,
+        class_name: Union[str, StepToDict] = None,
         name: str = None,
         handler: str = None,
         graph_shape: str = None,
@@ -218,11 +291,12 @@ class BaseStep(ModelObj):
     ):
         """add a step right after this step and return the new step
 
-        example, a 4 step pipeline ending with a stream:
-        graph.to('URLDownloader')\
-             .to('ToParagraphs')\
-             .to(name='to_json', handler='json.dumps')\
-             .to('>>', 'to_v3io', path=stream_path)\
+        example:
+            a 4-step pipeline ending with a stream:
+            graph.to('URLDownloader')\
+                 .to('ToParagraphs')\
+                 .to(name='to_json', handler='json.dumps')\
+                 .to('>>', 'to_v3io', path=stream_path)\
 
         :param class_name:  class name or step object to build the step from
                             for router steps the class name should start with '*'
@@ -233,7 +307,7 @@ class BaseStep(ModelObj):
         :param function:    function this step should run in
         :param full_event:  this step accepts the full event (not just body)
         :param input_path:  selects the key/path in the event to use as input to the step
-                            this require that the event body will behave like a dict, example:
+                            this requires that the event body will behave like a dict, example:
                             event: {"data": {"a": 5, "b": 7}}, input_path="data.b" means the step will
                             receive 7 as input
         :param result_path: selects the key/path in the event to write the results to
@@ -327,6 +401,7 @@ class TaskStep(BaseStep):
             args = signature(self._handler).parameters
             if args and "context" in list(args.keys()):
                 self._inject_context = True
+            self._set_error_handler()
             return
 
         self._class_object, self.class_name = self.get_step_class_object(
@@ -464,14 +539,23 @@ class TaskStep(BaseStep):
             )
             event.body = _update_result_body(self.result_path, event.body, result)
         except Exception as exc:
-            self._log_error(event, exc)
-            handled = self._call_error_handler(event, exc)
-            if not handled:
+            if self._on_error_handler:
+                self._log_error(event, exc)
+                result = self._call_error_handler(event, exc)
+                event.body = _update_result_body(self.result_path, event.body, result)
+            else:
                 raise exc
-            event.terminated = True
         return event
 
 
+class ErrorStep(TaskStep):
+    """error execution step, runs a class or handler"""
+
+    kind = "error_step"
+    _dict_fields = _task_step_fields + ["before", "base_step"]
+    _default_class = ""
+
+
 class RouterStep(TaskStep):
     """router step, implement routing logic for running child routes"""
 
@@ -824,6 +908,7 @@ class FlowStep(BaseStep):
     def init_object(self, context, namespace, mode="sync", reset=False, **extra_kwargs):
         """initialize graph objects and classes"""
         self.context = context
+        self._insert_all_error_handlers()
         self.check_and_process_graph()
 
         for step in self._steps.values():
@@ -866,7 +951,11 @@ class FlowStep(BaseStep):
 
         responders = []
         for step in self._steps.values():
-            if hasattr(step, "responder") and step.responder:
+            if (
+                hasattr(step, "responder")
+                and step.responder
+                and step.kind != "error_step"
+            ):
                 responders.append(step.name)
             if step.on_error and step.on_error in start_steps:
                 start_steps.remove(step.on_error)
@@ -979,6 +1068,10 @@ class FlowStep(BaseStep):
                 # never set a step as its own error handler
                 if step != error_step:
                     step.async_object.set_recovery_step(error_step.async_object)
+                    for next_step in error_step.next or []:
+                        next_state = self[next_step]
+                        if next_state.async_object and error_step.async_object:
+                            error_step.async_object.to(next_state.async_object)
 
         self._controller = source.run()
 
@@ -1059,15 +1152,22 @@ class FlowStep(BaseStep):
             try:
                 event = next_obj.run(event, *args, **kwargs)
             except Exception as exc:
-                self._log_error(event, exc, failed_step=next_obj.name)
-                handled = self._call_error_handler(event, exc)
-                if not handled:
+                if self._on_error_handler:
+                    self._log_error(event, exc, failed_step=next_obj.name)
+                    event.body = self._call_error_handler(event, exc)
+                    event.terminated = True
+                    return event
+                else:
                     raise exc
-                event.terminated = True
-                return event
 
             if hasattr(event, "terminated") and event.terminated:
                 return event
+            if (
+                hasattr(event, "error")
+                and isinstance(event.error, dict)
+                and next_obj.name in event.error
+            ):
+                next_obj = self._steps[next_obj.on_error]
             next = next_obj.next
             if next and len(next) > 1:
                 raise GraphError(
@@ -1103,6 +1203,33 @@ class FlowStep(BaseStep):
             **kw,
         )
 
+    def _insert_all_error_handlers(self):
+        """
+        insert all error steps to the graph
+        run after deployment
+        """
+        for name, step in self._steps.items():
+            if step.kind == "error_step":
+                self._insert_error_step(name, step)
+
+    def _insert_error_step(self, name, step):
+        """
+        insert error step to the graph
+        run after deployment
+        """
+        if not step.before and not any(
+            [step.name in other_step.after for other_step in self._steps.values()]
+        ):
+            step.responder = True
+            return
+
+        for step_name in step.before:
+            if step_name not in self._steps.keys():
+                raise MLRunInvalidArgumentError(
+                    f"cant set before, there is no step named {step_name}"
+                )
+            self[step_name].after_step(name)
+
 
 class RootFlowStep(FlowStep):
     """root flow step"""
@@ -1116,6 +1243,7 @@ classes_map = {
     "router": RouterStep,
     "flow": FlowStep,
     "queue": QueueStep,
+    "error_step": ErrorStep,
 }
 
 
@@ -1155,15 +1283,8 @@ def _add_graphviz_flow(
                 _add_graphviz_router(sg, child)
         else:
             graph.node(child.fullname, label=child.name, shape=child.get_shape())
-        after = child.after or []
-        for item in after:
-            previous_object = step[item]
-            kw = (
-                {"ltail": "cluster_" + previous_object.fullname}
-                if previous_object.kind == StepKinds.router
-                else {}
-            )
-            graph.edge(previous_object.fullname, child.fullname, **kw)
+        _add_edges(child.after or [], step, graph, child)
+        _add_edges(getattr(child, "before", []), step, graph, child, after=False)
         if child.on_error:
             graph.edge(child.fullname, child.on_error, style="dashed")
 
@@ -1183,6 +1304,18 @@ def _add_graphviz_flow(
                 graph.edge(last_step, target.fullname)
 
 
+def _add_edges(items, step, graph, child, after=True):
+    for item in items:
+        next_or_prev_object = step[item]
+        kw = {}
+        if next_or_prev_object.kind == StepKinds.router:
+            kw["ltail"] = f"cluster_{next_or_prev_object.fullname}"
+        if after:
+            graph.edge(next_or_prev_object.fullname, child.fullname, **kw)
+        else:
+            graph.edge(child.fullname, next_or_prev_object.fullname, **kw)
+
+
 def _generate_graphviz(
     step,
     renderer,
@@ -1355,7 +1488,7 @@ def _init_async_objects(context, steps):
                             endpoint, stream_path = parse_path(step.path)
                             stream_path = stream_path.strip("/")
                         step._async_object = storey.StreamTarget(
-                            storey.V3ioDriver(endpoint),
+                            storey.V3ioDriver(endpoint or config.v3io_api),
                             stream_path,
                             context=context,
                             **options,

mlrun/serving/utils.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mlrun/serving/v1_serving.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mlrun/serving/v2_serving.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,23 +11,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 import threading
 import time
 import traceback
 from typing import Dict, Union
 
 import mlrun
-from mlrun.api.schemas import (
-    ModelEndpoint,
-    ModelEndpointMetadata,
-    ModelEndpointSpec,
-    ModelEndpointStatus,
-    ModelMonitoringMode,
-)
+import mlrun.common.model_monitoring
+import mlrun.common.schemas
 from mlrun.artifacts import ModelArtifact  # noqa: F401
 from mlrun.config import config
 from mlrun.utils import logger, now_date, parse_versioned_object_uri
-from mlrun.utils.model_monitoring import EndpointType
 
 from .server import GraphServer
 from .utils import StepToDict, _extract_input_data, _update_result_body
@@ -265,11 +260,20 @@ class V2ModelServer(StepToDict):
             # get model health operation
             setattr(event, "terminated", True)
             if self.ready:
-                event.body = self.context.Response()
+                # Generate a response, confirming that the model is ready
+                event.body = self.context.Response(
+                    status_code=200,
+                    body=bytes(
+                        f"Model {self.name} is ready (event_id = {event_id})",
+                        encoding="utf-8",
+                    ),
+                )
+
             else:
                 event.body = self.context.Response(
                     status_code=408, body=b"model not ready"
                 )
+
             return event
 
         elif op == "" and event.method == "GET":
@@ -404,11 +408,12 @@ class _ModelLogPusher:
         return base_data
 
     def push(self, start, request, resp=None, op=None, error=None):
+        start_str = start.isoformat(sep=" ", timespec="microseconds")
         if error:
             data = self.base_data()
             data["request"] = request
             data["op"] = op
-            data["when"] = str(start)
+            data["when"] = start_str
             message = str(error)
             if self.verbose:
                 message = f"{message}\n{traceback.format_exc()}"
@@ -445,7 +450,7 @@ class _ModelLogPusher:
                 data["request"] = request
                 data["op"] = op
                 data["resp"] = resp
-                data["when"] = str(start)
+                data["when"] = start_str
                 data["microsec"] = microsec
                 if getattr(self.model, "metrics", None):
                     data["metrics"] = self.model.metrics
@@ -486,7 +491,7 @@ def _init_endpoint_record(
         versioned_model_name = f"{model.name}:latest"
 
     # Generating model endpoint ID based on function uri and model version
-    uid = mlrun.utils.model_monitoring.create_model_endpoint_id(
+    uid = mlrun.common.model_monitoring.create_model_endpoint_uid(
         function_uri=graph_server.function_uri, versioned_model=versioned_model_name
     ).uid
 
@@ -498,11 +503,11 @@ def _init_endpoint_record(
         logger.info("Creating a new model endpoint record", endpoint_id=uid)
 
         try:
-            model_endpoint = ModelEndpoint(
-                metadata=ModelEndpointMetadata(
+            model_endpoint = mlrun.common.schemas.ModelEndpoint(
+                metadata=mlrun.common.schemas.ModelEndpointMetadata(
                     project=project, labels=model.labels, uid=uid
                 ),
-                spec=ModelEndpointSpec(
+                spec=mlrun.common.schemas.ModelEndpointSpec(
                     function_uri=graph_server.function_uri,
                     model=versioned_model_name,
                     model_class=model.__class__.__name__,
@@ -511,18 +516,21 @@ def _init_endpoint_record(
                         project=project, kind="stream"
                     ),
                     active=True,
-                    monitoring_mode=ModelMonitoringMode.enabled
+                    monitoring_mode=mlrun.common.model_monitoring.ModelMonitoringMode.enabled
                     if model.context.server.track_models
-                    else ModelMonitoringMode.disabled,
+                    else mlrun.common.model_monitoring.ModelMonitoringMode.disabled,
+                ),
+                status=mlrun.common.schemas.ModelEndpointStatus(
+                    endpoint_type=mlrun.common.model_monitoring.EndpointType.NODE_EP
                 ),
-                status=ModelEndpointStatus(endpoint_type=EndpointType.NODE_EP),
             )
 
             db = mlrun.get_run_db()
+
             db.create_model_endpoint(
                 project=project,
-                endpoint_id=model_endpoint.metadata.uid,
-                model_endpoint=model_endpoint,
+                endpoint_id=uid,
+                model_endpoint=model_endpoint.dict(),
             )
 
         except Exception as e:

mlrun/utils/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -18,4 +18,3 @@ from .clones import get_git_username_password_from_token  # noqa
 from .helpers import *  # noqa
 from .http import *  # noqa
 from .logger import *  # noqa
-from .vault import *  # noqa

mlrun/utils/async_http.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -134,6 +134,13 @@ class _CustomRequestContext(_RequestContext):
                     params = self._params_list[-1]
 
                 headers = {k: v for k, v in params.headers.items() if v is not None}
+
+                # enrich user agent
+                # will help traceability and debugging
+                headers[
+                    aiohttp.hdrs.USER_AGENT
+                ] = f"{aiohttp.http.SERVER_SOFTWARE} mlrun/{config.version}"
+
                 response: typing.Optional[
                     aiohttp.ClientResponse
                 ] = await self._request_func(

mlrun/utils/azure_vault.py

@@ -1,4 +1,4 @@
-# Copyright 2021 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mlrun/utils/clones.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -51,7 +51,7 @@ def get_git_username_password_from_token(token):
     # Github's access tokens have a known prefix according to their type. See
     # https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/about-authentication-to-github#githubs-token-formats
     # We distinguish new fine-grained access tokens (begin with "github_pat_" from classic tokens.
-    if token.startswith("github_pat_"):
+    if token.startswith("github_pat_") or token.startswith("glpat"):
         username = "oauth2"
         password = token
     else:

mlrun/utils/condition_evaluator.py

@@ -0,0 +1,65 @@
+# Copyright 2023 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import multiprocessing
+import typing
+
+from mlrun.utils import logger
+
+
+def evaluate_condition_in_separate_process(
+    condition: str, context: typing.Dict[str, typing.Any], timeout: int = 5
+):
+
+    if not condition:
+        return True
+
+    receiver, sender = multiprocessing.Pipe()
+    p = multiprocessing.Process(
+        target=_evaluate_condition_wrapper,
+        args=(sender, condition, context),
+    )
+    p.start()
+    if receiver.poll(timeout):
+        result = receiver.recv()
+        p.join()
+        return result
+    else:
+        p.kill()
+        logger.warning(
+            f"Condition evaluation timed out after {timeout} seconds. Ignoring condition",
+            condition=condition,
+        )
+        return True
+
+
+def _evaluate_condition_wrapper(
+    connection, condition: str, context: typing.Dict[str, typing.Any]
+):
+    connection.send(_evaluate_condition(condition, context))
+    return connection.close()
+
+
+def _evaluate_condition(condition: str, context: typing.Dict[str, typing.Any]):
+
+    import jinja2.sandbox
+
+    jinja_env = jinja2.sandbox.SandboxedEnvironment()
+    template = jinja_env.from_string(condition)
+    result = template.render(**context)
+    if result.lower() in ["0", "no", "n", "f", "false", "off"]:
+        return False
+
+    # if the condition is not a boolean, we ignore the condition
+    return True

mlrun/utils/db.py

@@ -0,0 +1,52 @@
+# Copyright 2023 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import pickle
+from datetime import datetime
+
+from sqlalchemy.orm import class_mapper
+
+
+class BaseModel:
+    def to_dict(self, exclude=None):
+        """
+        NOTE - this function (currently) does not handle serializing relationships
+        """
+        exclude = exclude or []
+        mapper = class_mapper(self.__class__)
+        columns = [column.key for column in mapper.columns if column.key not in exclude]
+        get_key_value = (
+            lambda c: (c, getattr(self, c).isoformat())
+            if isinstance(getattr(self, c), datetime)
+            else (c, getattr(self, c))
+        )
+        return dict(map(get_key_value, columns))
+
+
+class HasStruct(BaseModel):
+    @property
+    def struct(self):
+        return pickle.loads(self.body)
+
+    @struct.setter
+    def struct(self, value):
+        self.body = pickle.dumps(value)
+
+    def to_dict(self, exclude=None):
+        """
+        NOTE - this function (currently) does not handle serializing relationships
+        """
+        exclude = exclude or []
+        exclude.append("body")
+        return super().to_dict(exclude)