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

mlrun/runtimes/sparkjob/spark3job.py

@@ -11,10 +11,10 @@
 # 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
+from mlrun_pipelines.mounts import mount_v3io, mount_v3iod
 
 import mlrun.common.schemas.function
 import mlrun.errors
@@ -23,7 +23,6 @@ from mlrun.config import config
 
 from ...execution import MLClientCtx
 from ...model import RunObject
-from ...platforms.iguazio import mount_v3io, mount_v3iod
 from ...utils import update_in, verify_field_regex
 from ..kubejob import KubejobRuntime
 from ..pod import KubeResourceSpec
@@ -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,22 +16,21 @@ import hashlib
 import json
 import os
 import re
-import typing
 from io import StringIO
 from sys import stderr
 
 import pandas as pd
-from kubernetes import client
 
 import mlrun
 import mlrun.common.constants
+import mlrun.common.constants as mlrun_constants
 import mlrun.common.schemas
 import mlrun.utils.regex
 from mlrun.artifacts import TableArtifact
+from mlrun.common.runtimes.constants import RunLabels
 from mlrun.config import config
 from mlrun.errors import err_to_str
 from mlrun.frameworks.parallel_coordinates import gen_pcp_plot
-from mlrun.runtimes.constants import RunLabels
 from mlrun.runtimes.generators import selector
 from mlrun.utils import get_in, helpers, logger, verify_field_regex
 
@@ -40,9 +39,6 @@ class RunError(Exception):
     pass
 
 
-mlrun_key = "mlrun/"
-
-
 class _ContextStore:
     def __init__(self):
         self._context = None
@@ -281,43 +277,6 @@ def get_item_name(item, attr="name"):
         return getattr(item, attr, None)
 
 
-def apply_kfp(modify, cop, runtime):
-    modify(cop)
-
-    # Have to do it here to avoid circular dependencies
-    from .pod import AutoMountType
-
-    if AutoMountType.is_auto_modifier(modify):
-        runtime.spec.disable_auto_mount = True
-
-    api = client.ApiClient()
-    for k, v in cop.pod_labels.items():
-        runtime.metadata.labels[k] = v
-    for k, v in cop.pod_annotations.items():
-        runtime.metadata.annotations[k] = v
-    if cop.container.env:
-        env_names = [
-            e.name if hasattr(e, "name") else e["name"] for e in runtime.spec.env
-        ]
-        for e in api.sanitize_for_serialization(cop.container.env):
-            name = e["name"]
-            if name in env_names:
-                runtime.spec.env[env_names.index(name)] = e
-            else:
-                runtime.spec.env.append(e)
-                env_names.append(name)
-        cop.container.env.clear()
-
-    if cop.volumes and cop.container.volume_mounts:
-        vols = api.sanitize_for_serialization(cop.volumes)
-        mounts = api.sanitize_for_serialization(cop.container.volume_mounts)
-        runtime.spec.update_vols_and_mounts(vols, mounts)
-        cop.volumes.clear()
-        cop.container.volume_mounts.clear()
-
-    return runtime
-
-
 def verify_limits(
     resources_field_name,
     mem=None,
@@ -411,41 +370,13 @@ def generate_resources(mem=None, cpu=None, gpus=None, gpu_type="nvidia.com/gpu")
 
 
 def get_func_selector(project, name=None, tag=None):
-    s = [f"{mlrun_key}project={project}"]
+    s = [f"{mlrun_constants.MLRunInternalLabels.project}={project}"]
     if name:
-        s.append(f"{mlrun_key}function={name}")
-        s.append(f"{mlrun_key}tag={tag or 'latest'}")
+        s.append(f"{mlrun_constants.MLRunInternalLabels.function}={name}")
+        s.append(f"{mlrun_constants.MLRunInternalLabels.tag}={tag or 'latest'}")
     return s
 
 
-class k8s_resource:
-    kind = ""
-    per_run = False
-    per_function = False
-    k8client = None
-
-    def deploy_function(self, function):
-        pass
-
-    def release_function(self, function):
-        pass
-
-    def submit_run(self, function, runobj):
-        pass
-
-    def get_object(self, name, namespace=None):
-        return None
-
-    def get_status(self, name, namespace=None):
-        return None
-
-    def del_object(self, name, namespace=None):
-        pass
-
-    def get_pods(self, name, namespace=None, master=False):
-        return {}
-
-
 def enrich_function_from_dict(function, function_dict):
     override_function = mlrun.new_function(runtime=function_dict, kind=function.kind)
     for attribute in [
@@ -501,10 +432,11 @@ 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(),
+        # TODO: remove this in 1.9.0
         RunLabels.v3io_user: os.environ.get("V3IO_USERNAME"),
     }
     labels_to_enrich = labels_to_enrich or RunLabels.all()

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:
@@ -163,15 +163,19 @@ def get_secret_or_env(
 
     Example::
 
-        secrets = { "KEY1": "VALUE1" }
+        secrets = {"KEY1": "VALUE1"}
         secret = get_secret_or_env("KEY1", secret_provider=secrets)
 
+
         # Using a function to retrieve a secret
         def my_secret_provider(key):
             # some internal logic to retrieve secret
             return value
 
-        secret = get_secret_or_env("KEY1", secret_provider=my_secret_provider, default="TOO-MANY-SECRETS")
+
+        secret = get_secret_or_env(
+            "KEY1", secret_provider=my_secret_provider, default="TOO-MANY-SECRETS"
+        )
 
     :param key: Secret key to look for
     :param secret_provider: Dictionary, callable or `SecretsStore` to extract the secret value from. If using a
@@ -185,7 +189,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

@@ -37,8 +37,6 @@ default_backoff_factor = 1
 
 
 class RemoteStep(storey.SendToHttp):
-    """class for calling remote endpoints"""
-
     def __init__(
         self,
         url: str,
@@ -174,8 +172,7 @@ class RemoteStep(storey.SendToHttp):
         if not self._session:
             self._session = mlrun.utils.HTTPSessionWithRetry(
                 self.retries,
-                self.backoff_factor
-                or mlrun.config.config.http_retry_defaults.backoff_factor,
+                self.backoff_factor or mlrun.mlconf.http_retry_defaults.backoff_factor,
                 retry_on_exception=False,
                 retry_on_status=self.retries > 0,
                 retry_on_post=True,
@@ -187,7 +184,7 @@ class RemoteStep(storey.SendToHttp):
             resp = self._session.request(
                 method,
                 url,
-                verify=mlrun.config.config.httpdb.http.verify,
+                verify=mlrun.mlconf.httpdb.http.verify,
                 headers=headers,
                 data=body,
                 timeout=self.timeout,
@@ -242,8 +239,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
@@ -28,6 +28,7 @@ import numpy as np
 import mlrun
 import mlrun.common.model_monitoring
 import mlrun.common.schemas.model_monitoring
+from mlrun.errors import err_to_str
 from mlrun.utils import logger, now_date
 
 from ..common.helpers import parse_versioned_object_uri
@@ -271,7 +272,9 @@ class ParallelRun(BaseModelRouter):
             fn = mlrun.new_function("parallel", kind="serving")
             graph = fn.set_topology(
                 "router",
-                mlrun.serving.routers.ParallelRun(extend_event=True, executor_type=executor),
+                mlrun.serving.routers.ParallelRun(
+                    extend_event=True, executor_type=executor
+                ),
             )
             graph.add_route("child1", class_name="Cls1")
             graph.add_route("child2", class_name="Cls2", my_arg={"c": 7})
@@ -485,7 +488,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 +706,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 +730,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 +744,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 +960,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.
@@ -1013,7 +1016,7 @@ 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
@@ -1089,12 +1092,12 @@ def _init_endpoint_record(
         except Exception as exc:
             logger.warning(
                 "Failed creating model endpoint record",
-                exc=exc,
+                exc=err_to_str(exc),
                 traceback=traceback.format_exc(),
             )
 
     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 endpoint_uid
 

mlrun/serving/server.py

@@ -23,6 +23,7 @@ import uuid
 from typing import Optional, Union
 
 import mlrun
+import mlrun.common.constants
 import mlrun.common.helpers
 import mlrun.model_monitoring
 from mlrun.config import config
@@ -52,7 +53,7 @@ class _StreamContext:
         Initialize _StreamContext object.
         :param enabled:      A boolean indication for applying the stream context
         :param parameters:   Dictionary of optional parameters, such as `log_stream` and `stream_args`. Note that these
-                             parameters might be relevant to the output source such as `kafka_bootstrap_servers` if
+                             parameters might be relevant to the output source such as `kafka_brokers` if
                              the output source is from type Kafka.
         :param function_uri: Full value of the function uri, usually it's <project-name>/<function-name>
         """
@@ -188,11 +189,6 @@ class GraphServer(ModelObj):
 
     def init_object(self, namespace):
         self.graph.init_object(self.context, namespace, self.load_mode, reset=True)
-        return (
-            v2_serving_async_handler
-            if config.datastore.async_source_mode == "enabled"
-            else v2_serving_handler
-        )
 
     def test(
         self,
@@ -310,17 +306,14 @@ class GraphServer(ModelObj):
 
     def wait_for_completion(self):
         """wait for async operation to complete"""
-        self.graph.wait_for_completion()
+        return self.graph.wait_for_completion()
 
 
 def v2_serving_init(context, namespace=None):
     """hook for nuclio init_context()"""
 
-    data = os.environ.get("SERVING_SPEC_ENV", "")
-    if not data:
-        raise MLRunInvalidArgumentError("failed to find spec env var")
-    spec = json.loads(data)
     context.logger.info("Initializing server from spec")
+    spec = mlrun.utils.get_serving_spec()
     server = GraphServer.from_dict(spec)
     if config.log_level.lower() == "debug":
         server.verbose = True
@@ -334,11 +327,18 @@ def v2_serving_init(context, namespace=None):
     context.logger.info_with(
         "Initializing states", namespace=namespace or get_caller_globals()
     )
-    server.init_states(context, namespace or get_caller_globals())
+    kwargs = {}
+    if hasattr(context, "is_mock"):
+        kwargs["is_mock"] = context.is_mock
+    server.init_states(
+        context,
+        namespace or get_caller_globals(),
+        **kwargs,
+    )
     context.logger.info("Initializing graph steps")
-    serving_handler = server.init_object(namespace or get_caller_globals())
+    server.init_object(namespace or get_caller_globals())
     # set the handler hook to point to our handler
-    setattr(context, "mlrun_handler", serving_handler)
+    setattr(context, "mlrun_handler", v2_serving_handler)
     setattr(context, "_server", server)
     context.logger.info_with("Serving was initialized", verbose=server.verbose)
     if server.verbose:
@@ -351,7 +351,7 @@ def v2_serving_init(context, namespace=None):
             "Setting termination callback to terminate graph on worker shutdown"
         )
 
-        def termination_callback():
+        async def termination_callback():
             context.logger.info("Termination callback called")
             server.wait_for_completion()
             context.logger.info("Termination of async flow is completed")
@@ -363,7 +363,7 @@ def v2_serving_init(context, namespace=None):
             "Setting drain callback to terminate and restart the graph on a drain event (such as rebalancing)"
         )
 
-        def drain_callback():
+        async def drain_callback():
             context.logger.info("Drain callback called")
             server.wait_for_completion()
             context.logger.info(
@@ -386,13 +386,8 @@ def v2_serving_handler(context, event, get_body=False):
     return context._server.run(event, context, get_body)
 
 
-async def v2_serving_async_handler(context, event, get_body=False):
-    """hook for nuclio handler()"""
-    return await context._server.run(event, context, get_body)
-
-
 def create_graph_server(
-    parameters={},
+    parameters=None,
     load_mode=None,
     graph=None,
     verbose=False,
@@ -408,6 +403,7 @@ def create_graph_server(
         server.graph.add_route("my", class_name=MyModelClass, model_path="{path}", z=100)
         print(server.test("/v2/models/my/infer", testdata))
     """
+    parameters = parameters or {}
     server = GraphServer(graph, parameters, load_mode, verbose=verbose, **kwargs)
     server.set_current_function(
         current_function or os.environ.get("SERVING_CURRENT_FUNCTION", "")
@@ -415,7 +411,7 @@ def create_graph_server(
     return server
 
 
-class MockTrigger(object):
+class MockTrigger:
     """mock nuclio event trigger"""
 
     def __init__(self, kind="", name=""):
@@ -423,7 +419,7 @@ class MockTrigger(object):
         self.name = name
 
 
-class MockEvent(object):
+class MockEvent:
     """mock basic nuclio event object"""
 
     def __init__(
@@ -456,7 +452,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 +559,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(