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

mlrun/runtimes/daskjob.py

@@ -16,7 +16,7 @@ import inspect
 import socket
 import time
 from os import environ
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable, Optional, Union
 
 import mlrun.common.schemas
 import mlrun.errors
@@ -461,7 +461,7 @@ class DaskCluster(KubejobRuntime):
 
     def set_state_thresholds(
         self,
-        state_thresholds: Dict[str, str],
+        state_thresholds: dict[str, str],
         patch: bool = True,
     ):
         raise NotImplementedError(
@@ -477,23 +477,23 @@ class DaskCluster(KubejobRuntime):
         name: Optional[str] = "",
         project: Optional[str] = "",
         params: Optional[dict] = None,
-        inputs: Optional[Dict[str, str]] = None,
+        inputs: Optional[dict[str, str]] = None,
         out_path: Optional[str] = "",
         workdir: Optional[str] = "",
         artifact_path: Optional[str] = "",
         watch: Optional[bool] = True,
         schedule: Optional[Union[str, mlrun.common.schemas.ScheduleCronTrigger]] = None,
-        hyperparams: Optional[Dict[str, list]] = None,
+        hyperparams: Optional[dict[str, list]] = None,
         hyper_param_options: Optional[mlrun.model.HyperParamOptions] = None,
         verbose: Optional[bool] = None,
         scrape_metrics: Optional[bool] = None,
         local: Optional[bool] = False,
         local_code_path: Optional[str] = None,
         auto_build: Optional[bool] = None,
-        param_file_secrets: Optional[Dict[str, str]] = None,
-        notifications: Optional[List[mlrun.model.Notification]] = None,
-        returns: Optional[List[Union[str, Dict[str, str]]]] = None,
-        state_thresholds: Optional[Dict[str, int]] = None,
+        param_file_secrets: Optional[dict[str, str]] = None,
+        notifications: Optional[list[mlrun.model.Notification]] = None,
+        returns: Optional[list[Union[str, dict[str, str]]]] = None,
+        state_thresholds: Optional[dict[str, int]] = None,
         **launcher_kwargs,
     ) -> RunObject:
         if state_thresholds:

mlrun/runtimes/databricks_job/databricks_cancel_task.py

@@ -33,7 +33,7 @@ def main():
                 "The Databricks credentials path does not exist."
                 " Please manually cancel the job from the Databricks environment."
             )
-    with open(credentials_path, "r") as yaml_file:
+    with open(credentials_path) as yaml_file:
         loaded_data = yaml.safe_load(yaml_file)
     # use for flat yaml only
     for key, value in loaded_data.items():

mlrun/runtimes/databricks_job/databricks_runtime.py

@@ -14,7 +14,7 @@
 
 from ast import FunctionDef, parse, unparse
 from base64 import b64decode, b64encode
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable, Optional, Union
 
 import mlrun
 import mlrun.runtimes.kubejob as kubejob
@@ -215,23 +215,23 @@ def run_mlrun_databricks_job(context,task_parameters: dict, **kwargs):
         name: Optional[str] = "",
         project: Optional[str] = "",
         params: Optional[dict] = None,
-        inputs: Optional[Dict[str, str]] = None,
+        inputs: Optional[dict[str, str]] = None,
         out_path: Optional[str] = "",
         workdir: Optional[str] = "",
         artifact_path: Optional[str] = "",
         watch: Optional[bool] = True,
         schedule: Optional[Union[str, mlrun.common.schemas.ScheduleCronTrigger]] = None,
-        hyperparams: Optional[Dict[str, list]] = None,
+        hyperparams: Optional[dict[str, list]] = None,
         hyper_param_options: Optional[HyperParamOptions] = None,
         verbose: Optional[bool] = None,
         scrape_metrics: Optional[bool] = None,
         local: Optional[bool] = False,
         local_code_path: Optional[str] = None,
         auto_build: Optional[bool] = None,
-        param_file_secrets: Optional[Dict[str, str]] = None,
-        notifications: Optional[List[mlrun.model.Notification]] = None,
-        returns: Optional[List[Union[str, Dict[str, str]]]] = None,
-        state_thresholds: Optional[Dict[str, int]] = None,
+        param_file_secrets: Optional[dict[str, str]] = None,
+        notifications: Optional[list[mlrun.model.Notification]] = None,
+        returns: Optional[list[Union[str, dict[str, str]]]] = None,
+        state_thresholds: Optional[dict[str, int]] = None,
         **launcher_kwargs,
     ) -> RunObject:
         if local:

mlrun/runtimes/function_reference.py

@@ -111,7 +111,7 @@ class FunctionReference(ModelObj):
             code = self.code
             if kind == mlrun.runtimes.RuntimeKinds.serving:
                 code = code + mlrun_footer.format(
-                    mlrun.runtimes.serving.serving_subkind
+                    mlrun.runtimes.nuclio.serving.serving_subkind
                 )
             func = mlrun.new_function(
                 self.name, kind=kind, image=self.image or default_image

mlrun/runtimes/local.py

@@ -436,7 +436,7 @@ def run_exec(cmd, args, env=None, cwd=None):
     return out, err
 
 
-class _DupStdout(object):
+class _DupStdout:
     def __init__(self):
         self.terminal = sys.stdout
         self.buf = StringIO()

mlrun/runtimes/mpijob/abstract.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 import abc
 import os
-import typing
 
 from mlrun.config import config
 from mlrun.runtimes.kubejob import KubejobRuntime
@@ -206,7 +205,7 @@ class AbstractMPIJobRuntime(KubejobRuntime, abc.ABC):
 
         self.set_envs(horovod_autotune_settings)
 
-    def set_mpi_args(self, args: typing.List[str]) -> None:
+    def set_mpi_args(self, args: list[str]) -> None:
         """Sets the runtime's mpi arguments to args.
 
         Parameters

mlrun/runtimes/nuclio/__init__.py

@@ -0,0 +1,20 @@
+# 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.
+
+from .serving import ServingRuntime, new_v2_model_server  # noqa
+from .nuclio import nuclio_init_hook  # noqa
+from .function import (
+    min_nuclio_versions,
+    RemoteRuntime,
+)  # noqa

mlrun/runtimes/{function.py → nuclio/function.py}

@@ -34,23 +34,22 @@ import mlrun.k8s_utils
 import mlrun.utils
 import mlrun.utils.helpers
 from mlrun.common.schemas import AuthInfo
-
-from ..config import config as mlconf
-from ..errors import err_to_str
-from ..kfpops import deploy_op
-from ..lists import RunList
-from ..model import RunObject
-from ..platforms.iguazio import (
+from mlrun.config import config as mlconf
+from mlrun.errors import err_to_str
+from mlrun.kfpops import deploy_op
+from mlrun.lists import RunList
+from mlrun.model import RunObject
+from mlrun.platforms.iguazio import (
     VolumeMount,
     mount_v3io,
     parse_path,
     split_path,
     v3io_cred,
 )
-from ..utils import get_in, logger, update_in
-from .base import FunctionStatus, RunError
-from .pod import KubeResource, KubeResourceSpec
-from .utils import get_item_name, log_std
+from mlrun.runtimes.base import FunctionStatus, RunError
+from mlrun.runtimes.pod import KubeResource, KubeResourceSpec
+from mlrun.runtimes.utils import get_item_name, log_std
+from mlrun.utils import get_in, logger, update_in
 
 
 def validate_nuclio_version_compatibility(*min_versions):
@@ -386,7 +385,7 @@ class RemoteRuntime(KubeResource):
         workers: typing.Optional[int] = 8,
         port: typing.Optional[int] = None,
         host: typing.Optional[str] = None,
-        paths: typing.Optional[typing.List[str]] = None,
+        paths: typing.Optional[list[str]] = None,
         canary: typing.Optional[float] = None,
         secret: typing.Optional[str] = None,
         worker_timeout: typing.Optional[int] = None,
@@ -629,9 +628,9 @@ class RemoteRuntime(KubeResource):
     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[client.V1Affinity] = None,
-        tolerations: typing.Optional[typing.List[client.V1Toleration]] = None,
+        tolerations: typing.Optional[list[client.V1Toleration]] = None,
     ):
         """k8s node selection attributes"""
         if tolerations and not validate_nuclio_version_compatibility("1.7.5"):
@@ -683,7 +682,7 @@ class RemoteRuntime(KubeResource):
 
     def set_state_thresholds(
         self,
-        state_thresholds: typing.Dict[str, int],
+        state_thresholds: dict[str, int],
         patch: bool = True,
     ):
         raise NotImplementedError(
@@ -716,7 +715,7 @@ class RemoteRuntime(KubeResource):
         raise_on_exception=True,
         resolve_address=True,
         auth_info: AuthInfo = None,
-    ) -> typing.Tuple[str, str, typing.Optional[float]]:
+    ) -> tuple[str, str, typing.Optional[float]]:
         if dashboard:
             (
                 state,

mlrun/runtimes/{nuclio.py → nuclio/nuclio.py}

@@ -17,13 +17,13 @@ import os
 import socket
 
 import mlrun.db
+from mlrun.errors import err_to_str
+from mlrun.execution import MLClientCtx
+from mlrun.model import RunTemplate
+from mlrun.runtimes.local import get_func_arg
+from mlrun.serving.server import v2_serving_init
+from mlrun.serving.v1_serving import nuclio_serving_init
 
-from ..errors import err_to_str
-from ..execution import MLClientCtx
-from ..model import RunTemplate
-from ..serving.server import v2_serving_init
-from ..serving.v1_serving import nuclio_serving_init
-from .local import get_func_arg
 from .serving import serving_subkind
 
 

mlrun/runtimes/{serving.py → nuclio/serving.py}

@@ -15,20 +15,20 @@
 import json
 import os
 from copy import deepcopy
-from typing import List, Union
+from typing import Union
 
 import nuclio
 from nuclio import KafkaTrigger
 
 import mlrun
 import mlrun.common.schemas
+from mlrun.datastore import parse_kafka_url
+from mlrun.model import ObjectList
 from mlrun.model_monitoring.tracking_policy import TrackingPolicy
-
-from ..datastore import parse_kafka_url
-from ..model import ObjectList
-from ..secrets import SecretsStore
-from ..serving.server import GraphServer, create_graph_server
-from ..serving.states import (
+from mlrun.runtimes.function_reference import FunctionReference
+from mlrun.secrets import SecretsStore
+from mlrun.serving.server import GraphServer, create_graph_server
+from mlrun.serving.states import (
     RootFlowStep,
     RouterStep,
     StepKinds,
@@ -37,9 +37,9 @@ from ..serving.states import (
     new_remote_endpoint,
     params_to_step,
 )
-from ..utils import get_caller_globals, logger, set_paths
+from mlrun.utils import get_caller_globals, logger, set_paths
+
 from .function import NuclioSpec, RemoteRuntime
-from .function_reference import FunctionReference
 
 serving_subkind = "serving_v2"
 
@@ -216,12 +216,12 @@ class ServingSpec(NuclioSpec):
         graph_root_setter(self, graph)
 
     @property
-    def function_refs(self) -> List[FunctionReference]:
+    def function_refs(self) -> list[FunctionReference]:
         """function references, list of optional child function refs"""
         return self._function_refs
 
     @function_refs.setter
-    def function_refs(self, function_refs: List[FunctionReference]):
+    def function_refs(self, function_refs: list[FunctionReference]):
         self._function_refs = ObjectList.from_list(FunctionReference, function_refs)
 
 
@@ -309,7 +309,8 @@ class ServingRuntime(RemoteRuntime):
         stream_args: dict = None,
         tracking_policy: Union[TrackingPolicy, dict] = None,
     ):
-        """set tracking parameters:
+        """apply on your serving function to monitor a deployed model, including real-time dashboards to detect drift
+           and analyze performance.
 
         :param stream_path:     Path/url of the tracking stream e.g. v3io:///users/mike/mystream
                                 you can use the "dummy://" path for test/simulation.

mlrun/runtimes/pod.py

@@ -105,6 +105,50 @@ class KubeResourceSpec(FunctionSpec):
         "security_context",
         "state_thresholds",
     ]
+    _default_fields_to_strip = FunctionSpec._default_fields_to_strip + [
+        "volumes",
+        "volume_mounts",
+        "resources",
+        "replicas",
+        "image_pull_policy",
+        "service_account",
+        "image_pull_secret",
+        "node_name",
+        "node_selector",
+        "affinity",
+        "priority_class_name",
+        "tolerations",
+        "preemption_mode",
+        "security_context",
+    ]
+    _k8s_fields_to_serialize = [
+        "volumes",
+        "volume_mounts",
+        "resources",
+        "env",
+        "image_pull_policy",
+        "service_account",
+        "image_pull_secret",
+        "node_name",
+        "node_selector",
+        "affinity",
+        "tolerations",
+        "security_context",
+    ]
+    _fields_to_serialize = FunctionSpec._fields_to_serialize + _k8s_fields_to_serialize
+    _fields_to_enrich = FunctionSpec._fields_to_enrich + [
+        "env",  # Removing sensitive data from env
+    ]
+    _fields_to_skip_validation = FunctionSpec._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.
+        "affinity",
+        "tolerations",
+        "node_selector",
+    ]
 
     def __init__(
         self,
@@ -222,7 +266,7 @@ class KubeResourceSpec(FunctionSpec):
         self._affinity = transform_attribute_to_k8s_class_instance("affinity", affinity)
 
     @property
-    def tolerations(self) -> typing.List[k8s_client.V1Toleration]:
+    def tolerations(self) -> list[k8s_client.V1Toleration]:
         return self._tolerations
 
     @tolerations.setter
@@ -264,15 +308,42 @@ class KubeResourceSpec(FunctionSpec):
     def termination_grace_period_seconds(self) -> typing.Optional[int]:
         return self._termination_grace_period_seconds
 
-    def to_dict(self, fields=None, exclude=None):
-        exclude = exclude or []
-        _exclude = ["affinity", "tolerations", "security_context"]
-        struct = super().to_dict(fields, exclude=list(set(exclude + _exclude)))
-        api = k8s_client.ApiClient()
-        for field in _exclude:
-            if field not in exclude:
-                struct[field] = api.sanitize_for_serialization(getattr(self, field))
-        return struct
+    def _serialize_field(
+        self, struct: dict, field_name: str = None, strip: bool = False
+    ) -> typing.Any:
+        """
+        Serialize a field to a dict, list, or primitive type.
+        If field_name is in _k8s_fields_to_serialize, we will apply k8s serialization
+        """
+        k8s_api = k8s_client.ApiClient()
+        if field_name in self._k8s_fields_to_serialize:
+            return k8s_api.sanitize_for_serialization(getattr(self, field_name))
+        return super()._serialize_field(struct, field_name, strip)
+
+    def _enrich_field(
+        self, struct: dict, field_name: str = None, strip: bool = False
+    ) -> typing.Any:
+        k8s_api = k8s_client.ApiClient()
+        if strip:
+            if field_name == "env":
+                # We first try to pull from struct because the field might have been already serialized and if not,
+                # we pull from self
+                envs = struct.get(field_name, None) or getattr(self, field_name, None)
+                if envs:
+                    serialized_envs = k8s_api.sanitize_for_serialization(envs)
+                    for env in serialized_envs:
+                        if env["name"].startswith("V3IO_"):
+                            env["value"] = ""
+                    return serialized_envs
+        return super()._enrich_field(struct=struct, field_name=field_name, strip=strip)
+
+    def _apply_enrichment_before_to_dict_completion(
+        self, struct: dict, strip: bool = False
+    ):
+        if strip:
+            # Reset this, since mounts and env variables were cleared.
+            struct["disable_auto_mount"] = False
+        return super()._apply_enrichment_before_to_dict_completion(struct, strip)
 
     def update_vols_and_mounts(
         self, volumes, volume_mounts, volume_mounts_field_name="_volume_mounts"
@@ -455,7 +526,7 @@ class KubeResourceSpec(FunctionSpec):
             return {}
         return resources
 
-    def _merge_node_selector(self, node_selector: typing.Dict[str, str]):
+    def _merge_node_selector(self, node_selector: dict[str, str]):
         if not node_selector:
             return
 
@@ -464,7 +535,7 @@ class KubeResourceSpec(FunctionSpec):
 
     def _merge_tolerations(
         self,
-        tolerations: typing.List[k8s_client.V1Toleration],
+        tolerations: list[k8s_client.V1Toleration],
         tolerations_field_name: str,
     ):
         if not tolerations:
@@ -649,7 +720,7 @@ class KubeResourceSpec(FunctionSpec):
 
     def _merge_node_selector_term_to_node_affinity(
         self,
-        node_selector_terms: typing.List[k8s_client.V1NodeSelectorTerm],
+        node_selector_terms: list[k8s_client.V1NodeSelectorTerm],
         affinity_field_name: str,
     ):
         if not node_selector_terms:
@@ -694,7 +765,7 @@ class KubeResourceSpec(FunctionSpec):
 
     def _prune_affinity_node_selector_requirement(
         self,
-        node_selector_requirements: typing.List[k8s_client.V1NodeSelectorRequirement],
+        node_selector_requirements: list[k8s_client.V1NodeSelectorRequirement],
         affinity_field_name: str = "affinity",
     ):
         """
@@ -749,20 +820,18 @@ class KubeResourceSpec(FunctionSpec):
 
     @staticmethod
     def _prune_node_selector_requirements_from_node_selector_terms(
-        node_selector_terms: typing.List[k8s_client.V1NodeSelectorTerm],
-        node_selector_requirements_to_prune: typing.List[
-            k8s_client.V1NodeSelectorRequirement
-        ],
-    ) -> typing.List[k8s_client.V1NodeSelectorTerm]:
+        node_selector_terms: list[k8s_client.V1NodeSelectorTerm],
+        node_selector_requirements_to_prune: list[k8s_client.V1NodeSelectorRequirement],
+    ) -> list[k8s_client.V1NodeSelectorTerm]:
         """
         Goes over each expression in all the terms provided and removes the expressions if it matches
         one of the requirements provided to remove
 
         :return: New list of terms without the provided node selector requirements
         """
-        new_node_selector_terms: typing.List[k8s_client.V1NodeSelectorTerm] = []
+        new_node_selector_terms: list[k8s_client.V1NodeSelectorTerm] = []
         for term in node_selector_terms:
-            new_node_selector_requirements: typing.List[
+            new_node_selector_requirements: list[
                 k8s_client.V1NodeSelectorRequirement
             ] = []
             for node_selector_requirement in term.match_expressions:
@@ -791,7 +860,7 @@ class KubeResourceSpec(FunctionSpec):
 
     def _prune_tolerations(
         self,
-        tolerations: typing.List[k8s_client.V1Toleration],
+        tolerations: list[k8s_client.V1Toleration],
         tolerations_field_name: str = "tolerations",
     ):
         """
@@ -820,7 +889,7 @@ class KubeResourceSpec(FunctionSpec):
 
     def _prune_node_selector(
         self,
-        node_selector: typing.Dict[str, str],
+        node_selector: dict[str, str],
         node_selector_field_name: str,
     ):
         """
@@ -927,28 +996,6 @@ class KubeResource(BaseRuntime):
     def spec(self, spec):
         self._spec = self._verify_dict(spec, "spec", KubeResourceSpec)
 
-    def to_dict(self, fields=None, exclude=None, strip=False):
-        struct = super().to_dict(fields, exclude, strip=strip)
-        api = k8s_client.ApiClient()
-        struct = api.sanitize_for_serialization(struct)
-        if strip:
-            spec = struct["spec"]
-            for attr in [
-                "volumes",
-                "volume_mounts",
-                "driver_volume_mounts",
-                "executor_volume_mounts",
-            ]:
-                if attr in spec:
-                    del spec[attr]
-            if "env" in spec and spec["env"]:
-                for ev in spec["env"]:
-                    if ev["name"].startswith("V3IO_"):
-                        ev["value"] = ""
-            # Reset this, since mounts and env variables were cleared.
-            spec["disable_auto_mount"] = False
-        return struct
-
     def apply(self, modify):
         """
         Apply a modifier to the runtime which is used to change the runtimes k8s object's spec.
@@ -1065,7 +1112,7 @@ class KubeResource(BaseRuntime):
 
     def set_state_thresholds(
         self,
-        state_thresholds: typing.Dict[str, str],
+        state_thresholds: dict[str, str],
         patch: bool = True,
     ):
         """
@@ -1126,9 +1173,9 @@ class KubeResource(BaseRuntime):
     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[k8s_client.V1Affinity] = None,
-        tolerations: typing.Optional[typing.List[k8s_client.V1Toleration]] = None,
+        tolerations: typing.Optional[list[k8s_client.V1Toleration]] = None,
     ):
         """
         Enables to control on which k8s node the job will run

mlrun/runtimes/remotesparkjob.py

@@ -92,7 +92,7 @@ class RemoteSparkSpec(KubeResourceSpec):
         self.provider = provider
 
 
-class RemoteSparkProviders(object):
+class RemoteSparkProviders:
     iguazio = "iguazio"