mlrun 1.6.0rc35__py3-none-any.whl → 1.7.0rc2__py3-none-any.whl

mlrun/runtimes/sparkjob/spark3job.py

@@ -11,7 +11,6 @@
 # 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 typing
 
 import kubernetes.client
@@ -69,6 +68,48 @@ class Spark3JobSpec(KubeResourceSpec):
         "driver_cores",
         "executor_cores",
     ]
+    _default_fields_to_strip = KubeResourceSpec._default_fields_to_strip + [
+        "driver_node_selector",
+        "executor_node_selector",
+        "driver_tolerations",
+        "executor_tolerations",
+        "driver_affinity",
+        "executor_affinity",
+        "driver_volume_mounts",
+        "executor_volume_mounts",
+        "driver_cores",
+        "executor_cores",
+    ]
+
+    __k8s_fields_to_serialize = [
+        "driver_volume_mounts",
+        "executor_volume_mounts",
+        "driver_node_selector",
+        "executor_node_selector",
+        "executor_affinity",
+        "executor_tolerations",
+        "driver_affinity",
+        "driver_tolerations",
+    ]
+    _k8s_fields_to_serialize = (
+        KubeResourceSpec._k8s_fields_to_serialize + __k8s_fields_to_serialize
+    )
+    _fields_to_serialize = (
+        KubeResourceSpec._fields_to_serialize + __k8s_fields_to_serialize
+    )
+    _fields_to_skip_validation = KubeResourceSpec._fields_to_skip_validation + [
+        # TODO: affinity, tolerations and node_selector are skipped due to preemption mode transitions.
+        #  Preemption mode 'none' depends on the previous mode while the default mode may enrich these values.
+        #  When we allow 'None' values for these attributes we get their true values and they will undo the default
+        #  enrichment when creating the runtime from dict.
+        #  The enrichment should move to the server side and then this can be removed.
+        "driver_node_selector",
+        "executor_node_selector",
+        "executor_affinity",
+        "executor_tolerations",
+        "driver_affinity",
+        "driver_tolerations",
+    ]
 
     def __init__(
         self,
@@ -189,26 +230,8 @@ class Spark3JobSpec(KubeResourceSpec):
         self.driver_cores = driver_cores
         self.executor_cores = executor_cores
 
-    def to_dict(self, fields=None, exclude=None):
-        exclude = exclude or []
-        _exclude = [
-            "affinity",
-            "tolerations",
-            "security_context",
-            "executor_affinity",
-            "executor_tolerations",
-            "driver_affinity",
-            "driver_tolerations",
-        ]
-        struct = super().to_dict(fields, exclude=list(set(exclude + _exclude)))
-        api = kubernetes.client.ApiClient()
-        for field in _exclude:
-            if field not in exclude:
-                struct[field] = api.sanitize_for_serialization(getattr(self, field))
-        return struct
-
     @property
-    def executor_tolerations(self) -> typing.List[kubernetes.client.V1Toleration]:
+    def executor_tolerations(self) -> list[kubernetes.client.V1Toleration]:
         return self._executor_tolerations
 
     @executor_tolerations.setter
@@ -220,7 +243,7 @@ class Spark3JobSpec(KubeResourceSpec):
         )
 
     @property
-    def driver_tolerations(self) -> typing.List[kubernetes.client.V1Toleration]:
+    def driver_tolerations(self) -> list[kubernetes.client.V1Toleration]:
         return self._driver_tolerations
 
     @driver_tolerations.setter
@@ -461,11 +484,9 @@ class Spark3Runtime(KubejobRuntime):
     def with_node_selection(
         self,
         node_name: typing.Optional[str] = None,
-        node_selector: typing.Optional[typing.Dict[str, str]] = None,
+        node_selector: typing.Optional[dict[str, str]] = None,
         affinity: typing.Optional[kubernetes.client.V1Affinity] = None,
-        tolerations: typing.Optional[
-            typing.List[kubernetes.client.V1Toleration]
-        ] = None,
+        tolerations: typing.Optional[list[kubernetes.client.V1Toleration]] = None,
     ):
         if node_name:
             raise NotImplementedError(
@@ -495,11 +516,9 @@ class Spark3Runtime(KubejobRuntime):
     def with_driver_node_selection(
         self,
         node_name: typing.Optional[str] = None,
-        node_selector: typing.Optional[typing.Dict[str, str]] = None,
+        node_selector: typing.Optional[dict[str, str]] = None,
         affinity: typing.Optional[kubernetes.client.V1Affinity] = None,
-        tolerations: typing.Optional[
-            typing.List[kubernetes.client.V1Toleration]
-        ] = None,
+        tolerations: typing.Optional[list[kubernetes.client.V1Toleration]] = None,
     ):
         """
         Enables control of which k8s node the spark executor will run on.
@@ -528,11 +547,9 @@ class Spark3Runtime(KubejobRuntime):
     def with_executor_node_selection(
         self,
         node_name: typing.Optional[str] = None,
-        node_selector: typing.Optional[typing.Dict[str, str]] = None,
+        node_selector: typing.Optional[dict[str, str]] = None,
         affinity: typing.Optional[kubernetes.client.V1Affinity] = None,
-        tolerations: typing.Optional[
-            typing.List[kubernetes.client.V1Toleration]
-        ] = None,
+        tolerations: typing.Optional[list[kubernetes.client.V1Toleration]] = None,
     ):
         """
         Enables control of which k8s node the spark executor will run on.

mlrun/runtimes/utils.py

@@ -16,7 +16,6 @@ import hashlib
 import json
 import os
 import re
-import typing
 from io import StringIO
 from sys import stderr
 
@@ -501,7 +500,7 @@ def enrich_function_from_dict(function, function_dict):
 
 def enrich_run_labels(
     labels: dict,
-    labels_to_enrich: typing.List[RunLabels] = None,
+    labels_to_enrich: list[RunLabels] = None,
 ):
     labels_enrichment = {
         RunLabels.owner: os.environ.get("V3IO_USERNAME") or getpass.getuser(),

mlrun/secrets.py

@@ -14,7 +14,7 @@
 
 from ast import literal_eval
 from os import environ, getenv
-from typing import Callable, Dict, Optional, Union
+from typing import Callable, Optional, Union
 
 from .utils import AzureVaultStore, list2dict
 
@@ -148,7 +148,7 @@ class SecretsStore:
 
 def get_secret_or_env(
     key: str,
-    secret_provider: Union[Dict, SecretsStore, Callable, None] = None,
+    secret_provider: Union[dict, SecretsStore, Callable, None] = None,
     default: Optional[str] = None,
     prefix: Optional[str] = None,
 ) -> str:
@@ -185,7 +185,7 @@ def get_secret_or_env(
 
     value = None
     if secret_provider:
-        if isinstance(secret_provider, (Dict, SecretsStore)):
+        if isinstance(secret_provider, (dict, SecretsStore)):
             value = secret_provider.get(key)
         else:
             value = secret_provider(key)

mlrun/serving/remote.py

@@ -36,8 +36,6 @@ default_backoff_factor = 1
 
 
 class RemoteStep(storey.SendToHttp):
-    """class for calling remote endpoints"""
-
     def __init__(
         self,
         url: str,
@@ -240,8 +238,6 @@ class RemoteStep(storey.SendToHttp):
 
 
 class BatchHttpRequests(_ConcurrentJobExecution):
-    """class for calling remote endpoints in parallel"""
-
     def __init__(
         self,
         url: str = None,

mlrun/serving/routers.py

@@ -20,7 +20,7 @@ import traceback
 import typing
 from enum import Enum
 from io import BytesIO
-from typing import Dict, List, Union
+from typing import Union
 
 import numpy
 import numpy as np
@@ -485,7 +485,7 @@ class VotingEnsemble(ParallelRun):
         url_prefix: str = None,
         health_prefix: str = None,
         vote_type: str = None,
-        weights: Dict[str, float] = None,
+        weights: dict[str, float] = None,
         executor_type: Union[ParallelRunnerModes, str] = ParallelRunnerModes.thread,
         format_response_with_col_name_flag: bool = False,
         prediction_col_name: str = "prediction",
@@ -703,7 +703,7 @@ class VotingEnsemble(ParallelRun):
             )
         return model, None, subpath
 
-    def _majority_vote(self, all_predictions: List[List[int]], weights: List[float]):
+    def _majority_vote(self, all_predictions: list[list[int]], weights: list[float]):
         """
         Returns most predicted class for each event
 
@@ -727,7 +727,7 @@ class VotingEnsemble(ParallelRun):
         weighted_res = one_hot_representation @ weights
         return np.argmax(weighted_res, axis=1).tolist()
 
-    def _mean_vote(self, all_predictions: List[List[float]], weights: List[float]):
+    def _mean_vote(self, all_predictions: list[list[float]], weights: list[float]):
         """
         Returns weighted mean of the predictions
 
@@ -741,7 +741,7 @@ class VotingEnsemble(ParallelRun):
     def _is_int(self, value):
         return float(value).is_integer()
 
-    def logic(self, predictions: List[List[Union[int, float]]], weights: List[float]):
+    def logic(self, predictions: list[list[Union[int, float]]], weights: list[float]):
         """
         Returns the final prediction of all the models after applying the desire logic
 
@@ -957,7 +957,7 @@ class VotingEnsemble(ParallelRun):
                 raise Exception('Expected "inputs" to be a list')
         return request
 
-    def _normalize_weights(self, weights_dict: Dict[str, float]):
+    def _normalize_weights(self, weights_dict: dict[str, float]):
         """
         Normalized all the weights such that abs(weights_sum - 1.0) <= 0.001
         and adding 0 weight to all the routes that doesn't appear in the dict.

mlrun/serving/server.py

@@ -415,7 +415,7 @@ def create_graph_server(
     return server
 
 
-class MockTrigger(object):
+class MockTrigger:
     """mock nuclio event trigger"""
 
     def __init__(self, kind="", name=""):
@@ -423,7 +423,7 @@ class MockTrigger(object):
         self.name = name
 
 
-class MockEvent(object):
+class MockEvent:
     """mock basic nuclio event object"""
 
     def __init__(
@@ -456,7 +456,7 @@ class MockEvent(object):
         return f"Event(id={self.id}, body={self.body}, method={self.method}, path={self.path}{error})"
 
 
-class Response(object):
+class Response:
     def __init__(self, headers=None, body=None, content_type=None, status_code=200):
         self.headers = headers or {}
         self.body = body
@@ -563,7 +563,7 @@ class GraphContext:
             _,
             _,
             function_status,
-        ) = mlrun.runtimes.function.get_nuclio_deploy_status(name, project, tag)
+        ) = mlrun.runtimes.nuclio.function.get_nuclio_deploy_status(name, project, tag)
 
         if state in ["error", "unhealthy"]:
             raise ValueError(

mlrun/serving/states.py

@@ -556,6 +556,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"""
@@ -1514,6 +1542,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

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,7 +15,7 @@
 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
@@ -29,8 +29,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,
@@ -221,6 +219,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 +237,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 +367,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):

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,
@@ -67,7 +65,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":
@@ -102,7 +100,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)