mlrun 1.6.4rc2__py3-none-any.whl → 1.7.0rc20__py3-none-any.whl

mlrun/serving/states.py

@@ -19,13 +19,16 @@ import pathlib
 import traceback
 from copy import copy, deepcopy
 from inspect import getfullargspec, signature
-from typing import Union
+from typing import Any, Union
 
 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
@@ -324,7 +327,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 +349,36 @@ 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")
+
 
 class TaskStep(BaseStep):
     """task execution step, runs a class or handler"""
@@ -556,6 +589,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"""
@@ -1132,19 +1193,11 @@ 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
@@ -1182,6 +1235,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 +1288,27 @@ 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
+
 
 class RootFlowStep(FlowStep):
     """root flow step"""
@@ -1473,13 +1548,11 @@ def _init_async_objects(context, steps):
                     endpoint = None
                     options = {}
                     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 +1560,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 +1587,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 +1600,9 @@ 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()
 
+    # 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,7 @@ 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):
+    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:
@@ -97,5 +97,5 @@ class StepToDict:
 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

mlrun/serving/v2_serving.py

@@ -15,12 +15,13 @@
 import threading
 import time
 import traceback
-from typing import Dict, Union
+from typing import Union
 
 import mlrun.common.model_monitoring
 import mlrun.common.schemas.model_monitoring
 from mlrun.artifacts import ModelArtifact  # noqa: F401
 from mlrun.config import config
+from mlrun.errors import err_to_str
 from mlrun.utils import logger, now_date
 
 from ..common.helpers import parse_versioned_object_uri
@@ -29,8 +30,6 @@ from .utils import StepToDict, _extract_input_data, _update_result_body
 
 
 class V2ModelServer(StepToDict):
-    """base model serving class (v2), using similar API to KFServing v2 and Triton"""
-
     def __init__(
         self,
         context=None,
@@ -64,11 +63,11 @@ class V2ModelServer(StepToDict):
             class MyClass(V2ModelServer):
                 def load(self):
                     # load and initialize the model and/or other elements
-                    model_file, extra_data = self.get_model(suffix='.pkl')
+                    model_file, extra_data = self.get_model(suffix=".pkl")
                     self.model = load(open(model_file, "rb"))
 
                 def predict(self, request):
-                    events = np.array(request['inputs'])
+                    events = np.array(request["inputs"])
                     dmatrix = xgb.DMatrix(events)
                     result: xgb.DMatrix = self.model.predict(dmatrix)
                     return {"outputs": result.tolist()}
@@ -177,9 +176,9 @@ class V2ModelServer(StepToDict):
         ::
 
             def load(self):
-                model_file, extra_data = self.get_model(suffix='.pkl')
+                model_file, extra_data = self.get_model(suffix=".pkl")
                 self.model = load(open(model_file, "rb"))
-                categories = extra_data['categories'].as_df()
+                categories = extra_data["categories"].as_df()
 
         Parameters
         ----------
@@ -221,6 +220,8 @@ class V2ModelServer(StepToDict):
 
     def _pre_event_processing_actions(self, event, event_body, op):
         self._check_readiness(event)
+        if "_dict" in op:
+            event_body = self._inputs_to_list(event_body)
         request = self.preprocess(event_body, op)
         return self.validate(request, op)
 
@@ -237,7 +238,12 @@ class V2ModelServer(StepToDict):
         if not op and event.method != "GET":
             op = "infer"
 
-        if op == "predict" or op == "infer":
+        if (
+            op == "predict"
+            or op == "infer"
+            or op == "infer_dict"
+            or op == "predict_dict"
+        ):
             # predict operation
             request = self._pre_event_processing_actions(event, event_body, op)
             try:
@@ -362,22 +368,59 @@ class V2ModelServer(StepToDict):
 
         return request
 
-    def preprocess(self, request: Dict, operation) -> Dict:
+    def preprocess(self, request: dict, operation) -> dict:
         """preprocess the event body before validate and action"""
         return request
 
-    def postprocess(self, request: Dict) -> Dict:
+    def postprocess(self, request: dict) -> dict:
         """postprocess, before returning response"""
         return request
 
-    def predict(self, request: Dict) -> Dict:
+    def predict(self, request: dict) -> dict:
         """model prediction operation"""
         raise NotImplementedError()
 
-    def explain(self, request: Dict) -> Dict:
+    def explain(self, request: dict) -> dict:
         """model explain operation"""
         raise NotImplementedError()
 
+    def _inputs_to_list(self, request: dict) -> dict:
+        """
+        Convert the inputs from list of dictionary / dictionary to list of lists / list
+        where the internal list order is according to the ArtifactModel inputs.
+
+        :param request: event
+        :return: evnet body converting the inputs to be list of lists
+        """
+        if self.model_spec and self.model_spec.inputs:
+            input_order = [feature.name for feature in self.model_spec.inputs]
+        else:
+            raise mlrun.MLRunInvalidArgumentError(
+                "In order to use predict_dict or infer_dict operation you have to provide `model_path` "
+                "to the model server and to load it by `load()` function"
+            )
+        inputs = request.get("inputs")
+        try:
+            if isinstance(inputs, list) and all(
+                isinstance(item, dict) for item in inputs
+            ):
+                new_inputs = [
+                    [input_dict[key] for key in input_order] for input_dict in inputs
+                ]
+            elif isinstance(inputs, dict):
+                new_inputs = [inputs[key] for key in input_order]
+            else:
+                raise mlrun.MLRunInvalidArgumentError(
+                    "When using predict_dict or infer_dict operation the inputs must be "
+                    "of type `list[dict]` or `dict`"
+                )
+        except KeyError:
+            raise mlrun.MLRunInvalidArgumentError(
+                f"Input dictionary don't contain all the necessary input keys : {input_order}"
+            )
+        request["inputs"] = new_inputs
+        return request
+
 
 class _ModelLogPusher:
     def __init__(self, model, context, output_stream=None):
@@ -481,15 +524,11 @@ def _init_endpoint_record(
             graph_server.function_uri
         )
     except Exception as e:
-        logger.error("Failed to parse function URI", exc=e)
+        logger.error("Failed to parse function URI", exc=err_to_str(e))
         return None
 
     # Generating version model value based on the model name and model version
-    if model.model_path and model.model_path.startswith("store://"):
-        # Enrich the model server with the model artifact metadata
-        model.get_model()
-        model.version = model.model_spec.tag
-        model.labels = model.model_spec.labels
+    if model.version:
         versioned_model_name = f"{model.name}:{model.version}"
     else:
         versioned_model_name = f"{model.name}:latest"
@@ -538,9 +577,9 @@ def _init_endpoint_record(
             )
 
         except Exception as e:
-            logger.error("Failed to create endpoint record", exc=e)
+            logger.error("Failed to create endpoint record", exc=err_to_str(e))
 
     except Exception as e:
-        logger.error("Failed to retrieve model endpoint object", exc=e)
+        logger.error("Failed to retrieve model endpoint object", exc=err_to_str(e))
 
     return uid

mlrun/track/tracker.py

@@ -31,8 +31,9 @@ class Tracker(ABC):
     * Offline: Manually importing models and artifacts into an MLRun project using the `import_x` methods.
     """
 
+    @staticmethod
     @abstractmethod
-    def is_enabled(self) -> bool:
+    def is_enabled() -> bool:
         """
         Checks if tracker is enabled.
 

mlrun/track/tracker_manager.py

@@ -14,7 +14,7 @@
 
 import importlib
 import inspect
-from typing import List, Union
+from typing import Union
 
 import mlrun.errors
 from mlrun.config import config as mlconf
@@ -28,7 +28,7 @@ _TRACKERS = ["mlflow"]
 
 # A list for the available trackers during runtime. It will be setup at the beginning of the run by the function
 # `_collect_available_trackers`:
-_AVAILABLE_TRACKERS: List[Tracker] = None
+_AVAILABLE_TRACKERS: list[Tracker] = None
 
 
 class TrackerManager(metaclass=Singleton):
@@ -41,7 +41,7 @@ class TrackerManager(metaclass=Singleton):
         """
         Initialize a new empty tracker manager.
         """
-        self._trackers: List[Tracker] = []
+        self._trackers: list[Tracker] = []
 
         # Check general config for tracking usage, if false we return an empty manager
         if mlconf.external_platform_tracking.enabled:

mlrun/track/trackers/mlflow_tracker.py

@@ -15,7 +15,6 @@ import os
 import pathlib
 import tempfile
 import zipfile
-from typing import List
 
 import mlflow
 import mlflow.entities
@@ -526,7 +525,7 @@ class MLFlowTracker(Tracker):
         )
 
     @staticmethod
-    def _schema_to_feature(schema: mlflow.types.Schema) -> List[Feature]:
+    def _schema_to_feature(schema: mlflow.types.Schema) -> list[Feature]:
         """
         Cast MLFlow schema to MLRun features.
 

mlrun/utils/async_http.py

@@ -16,7 +16,7 @@
 import asyncio
 import logging
 import typing
-from typing import List, Optional
+from typing import Optional
 
 import aiohttp
 import aiohttp.http_exceptions
@@ -38,12 +38,10 @@ class AsyncClientWithRetry(RetryClient):
         self,
         max_retries: int = config.http_retry_defaults.max_retries,
         retry_backoff_factor: float = config.http_retry_defaults.backoff_factor,
-        retry_on_status_codes: typing.List[
-            int
-        ] = config.http_retry_defaults.status_codes,
+        retry_on_status_codes: list[int] = config.http_retry_defaults.status_codes,
         retry_on_exception: bool = True,
         raise_for_status: bool = True,
-        blacklisted_methods: typing.Optional[typing.List[str]] = None,
+        blacklisted_methods: typing.Optional[list[str]] = None,
         logger: logging.Logger = None,
         *args,
         **kwargs,
@@ -82,7 +80,7 @@ class AsyncClientWithRetry(RetryClient):
 
     def _make_requests(
         self,
-        params_list: List[RequestParams],
+        params_list: list[RequestParams],
         retry_options: Optional[RetryOptionsBase] = None,
         raise_for_status: Optional[bool] = None,
     ) -> "_CustomRequestContext":
@@ -117,7 +115,7 @@ class ExponentialRetryOverride(ExponentialRetry):
     def __init__(
         self,
         retry_on_exception: bool,
-        blacklisted_methods: typing.List[str],
+        blacklisted_methods: list[str],
         *args,
         **kwargs,
     ):

mlrun/utils/azure_vault.py

@@ -63,7 +63,7 @@ class AzureVaultStore:
             mlconf.secret_stores.azure_vault.secret_path + "/" + file_name
         )
         if os.path.isfile(full_path):
-            with open(full_path, "r") as secret_file:
+            with open(full_path) as secret_file:
                 contents = secret_file.read()
             return contents
         return None

mlrun/utils/clones.py

@@ -18,7 +18,6 @@ import tarfile
 import tempfile
 import zipfile
 from os import path, remove
-from typing import Tuple
 from urllib.parse import urlparse
 
 from git import Repo
@@ -91,7 +90,7 @@ def get_repo_url(repo):
     return url
 
 
-def add_credentials_git_remote_url(url: str, secrets=None) -> Tuple[str, bool]:
+def add_credentials_git_remote_url(url: str, secrets=None) -> tuple[str, bool]:
     """Enrich a Git remote URL with credential related secrets, if any are available
     If no secrets are supplied, or if the secrets are insufficient, the original URL is returned
     Besides the URL, this function also returns a bool indicating if any enrichment was done

mlrun/utils/condition_evaluator.py

@@ -19,7 +19,7 @@ from mlrun.utils import logger
 
 
 def evaluate_condition_in_separate_process(
-    condition: str, context: typing.Dict[str, typing.Any], timeout: int = 5
+    condition: str, context: dict[str, typing.Any], timeout: int = 5
 ):
     if not condition:
         return True
@@ -44,13 +44,13 @@ def evaluate_condition_in_separate_process(
 
 
 def _evaluate_condition_wrapper(
-    connection, condition: str, context: typing.Dict[str, typing.Any]
+    connection, condition: str, context: dict[str, typing.Any]
 ):
     connection.send(_evaluate_condition(condition, context))
     return connection.close()
 
 
-def _evaluate_condition(condition: str, context: typing.Dict[str, typing.Any]):
+def _evaluate_condition(condition: str, context: dict[str, typing.Any]):
     import jinja2.sandbox
 
     jinja_env = jinja2.sandbox.SandboxedEnvironment()

mlrun/utils/db.py

@@ -19,7 +19,7 @@ from sqlalchemy.orm import class_mapper
 
 
 class BaseModel:
-    def to_dict(self, exclude=None):
+    def to_dict(self, exclude=None, strip: bool = False):
         """
         NOTE - this function (currently) does not handle serializing relationships
         """
@@ -44,10 +44,10 @@ class HasStruct(BaseModel):
     def struct(self, value):
         self.body = pickle.dumps(value)
 
-    def to_dict(self, exclude=None):
+    def to_dict(self, exclude=None, strip: bool = False):
         """
         NOTE - this function (currently) does not handle serializing relationships
         """
         exclude = exclude or []
         exclude.append("body")
-        return super().to_dict(exclude)
+        return super().to_dict(exclude, strip=strip)