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

mlrun/runtimes/sparkjob/spark3job.py

@@ -11,19 +11,19 @@
 # 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
+import mlrun.k8s_utils
 import mlrun.runtimes.pod
 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 +69,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 +231,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 +244,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
@@ -428,7 +452,7 @@ class Spark3JobSpec(KubeResourceSpec):
 class Spark3Runtime(KubejobRuntime):
     group = "sparkoperator.k8s.io"
     version = "v1beta2"
-    apiVersion = group + "/" + version
+    apiVersion = group + "/" + version  # noqa: N815
     kind = "spark"
     plural = "sparkapplications"
 
@@ -461,11 +485,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(
@@ -484,22 +506,18 @@ class Spark3Runtime(KubejobRuntime):
             raise NotImplementedError(
                 "Setting node name is not supported for spark runtime"
             )
-        # TODO add affinity support
-        # https://github.com/GoogleCloudPlatform/spark-on-k8s-operator/blob/master/pkg/apis/sparkoperator.k8s.io/v1beta2/types.go#L491
-        if affinity:
-            raise NotImplementedError(
-                "Setting affinity is not supported for spark runtime"
-            )
-        super().with_node_selection(node_name, node_selector, affinity, tolerations)
+        mlrun.k8s_utils.validate_node_selectors(node_selector, raise_on_error=False)
+        self.with_driver_node_selection(node_name, node_selector, affinity, tolerations)
+        self.with_executor_node_selection(
+            node_name, node_selector, affinity, tolerations
+        )
 
     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.
@@ -518,21 +536,20 @@ class Spark3Runtime(KubejobRuntime):
             raise NotImplementedError(
                 "Setting node name is not supported for spark runtime"
             )
-        if affinity:
+        if affinity is not None:
             self.spec.driver_affinity = affinity
-        if node_selector:
+        if node_selector is not None:
+            mlrun.k8s_utils.validate_node_selectors(node_selector, raise_on_error=False)
             self.spec.driver_node_selector = node_selector
-        if tolerations:
+        if tolerations is not None:
             self.spec.driver_tolerations = tolerations
 
     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.
@@ -551,11 +568,12 @@ class Spark3Runtime(KubejobRuntime):
             raise NotImplementedError(
                 "Setting node name is not supported for spark runtime"
             )
-        if affinity:
+        if affinity is not None:
             self.spec.executor_affinity = affinity
-        if node_selector:
+        if node_selector is not None:
+            mlrun.k8s_utils.validate_node_selectors(node_selector, raise_on_error=False)
             self.spec.executor_node_selector = node_selector
-        if tolerations:
+        if tolerations is not None:
             self.spec.executor_tolerations = tolerations
 
     def with_preemption_mode(
@@ -794,9 +812,7 @@ class Spark3Runtime(KubejobRuntime):
 
     @classmethod
     def deploy_default_image(cls, with_gpu=False):
-        from mlrun.run import new_function
-
-        sj = new_function(kind=cls.kind, name="spark-default-image-deploy-temp")
+        sj = mlrun.new_function(kind=cls.kind, name="spark-default-image-deploy-temp")
         sj.spec.build.image = cls._get_default_deployed_mlrun_image_name(with_gpu)
 
         # setting required resources

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()
@@ -513,3 +445,37 @@ def enrich_run_labels(
         if label.value not in labels and enrichment:
             labels[label.value] = enrichment
     return labels
+
+
+def resolve_node_selectors(
+    project_node_selector: dict, instance_node_selector: dict
+) -> dict:
+    config_node_selector = mlrun.mlconf.get_default_function_node_selector()
+    if project_node_selector or config_node_selector:
+        mlrun.utils.logger.debug(
+            "Enriching node selector from project and mlrun config",
+            project_node_selector=project_node_selector,
+            config_node_selector=config_node_selector,
+        )
+        return mlrun.utils.helpers.merge_dicts_with_precedence(
+            config_node_selector,
+            project_node_selector,
+            instance_node_selector,
+        )
+    return instance_node_selector
+
+
+def enrich_gateway_timeout_annotations(annotations: dict, gateway_timeout: int):
+    """
+    Set gateway proxy connect/read/send timeout annotations
+    :param annotations:     The annotations to enrich
+    :param gateway_timeout: The timeout to set
+    """
+    if not gateway_timeout:
+        return
+    gateway_timeout_str = str(gateway_timeout)
+    annotations["nginx.ingress.kubernetes.io/proxy-connect-timeout"] = (
+        gateway_timeout_str
+    )
+    annotations["nginx.ingress.kubernetes.io/proxy-read-timeout"] = gateway_timeout_str
+    annotations["nginx.ingress.kubernetes.io/proxy-send-timeout"] = gateway_timeout_str

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/__init__.py

@@ -22,10 +22,17 @@ __all__ = [
     "RouterStep",
     "QueueStep",
     "ErrorStep",
+    "MonitoringApplicationStep",
 ]
 
 from .routers import ModelRouter, VotingEnsemble  # noqa
 from .server import GraphContext, GraphServer, create_graph_server  # noqa
-from .states import ErrorStep, QueueStep, RouterStep, TaskStep  # noqa
+from .states import (
+    ErrorStep,
+    QueueStep,
+    RouterStep,
+    TaskStep,
+    MonitoringApplicationStep,
+)  # noqa
 from .v1_serving import MLModelServer, new_v1_model_server  # noqa
 from .v2_serving import V2ModelServer  # noqa

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,