mlrun 1.9.0rc4__py3-none-any.whl → 1.10.0rc2__py3-none-any.whl

mlrun/runtimes/pod.py

@@ -29,17 +29,12 @@ import mlrun.errors
 import mlrun.runtimes.mounts
 import mlrun.utils.regex
 from mlrun.common.schemas import (
-    NodeSelectorOperator,
     PreemptionModes,
     SecurityContextEnrichmentModes,
 )
 
 from ..config import config as mlconf
 from ..k8s_utils import (
-    generate_preemptible_node_selector_requirements,
-    generate_preemptible_nodes_affinity_terms,
-    generate_preemptible_nodes_anti_affinity_terms,
-    generate_preemptible_tolerations,
     validate_node_selectors,
 )
 from ..utils import logger, update_in
@@ -292,7 +287,6 @@ class KubeResourceSpec(FunctionSpec):
     @preemption_mode.setter
     def preemption_mode(self, mode):
         self._preemption_mode = mode or mlconf.function_defaults.preemption_mode
-        self.enrich_function_preemption_spec()
 
     @property
     def security_context(self) -> k8s_client.V1SecurityContext:
@@ -560,300 +554,6 @@ class KubeResourceSpec(FunctionSpec):
             return {}
         return resources
 
-    def _merge_node_selector(self, node_selector: dict[str, str]):
-        if not node_selector:
-            return
-
-        # merge node selectors - precedence to existing node selector
-        self.node_selector = mlrun.utils.helpers.merge_dicts_with_precedence(
-            node_selector, self.node_selector
-        )
-
-    def _merge_tolerations(
-        self,
-        tolerations: list[k8s_client.V1Toleration],
-        tolerations_field_name: str,
-    ):
-        if not tolerations:
-            return
-        # In case function has no toleration, take all from input
-        self_tolerations = getattr(self, tolerations_field_name)
-        if not self_tolerations:
-            setattr(self, tolerations_field_name, tolerations)
-            return
-        tolerations_to_add = []
-
-        # Only add non-matching tolerations to avoid duplications
-        for toleration in tolerations:
-            to_add = True
-            for function_toleration in self_tolerations:
-                if function_toleration == toleration:
-                    to_add = False
-                    break
-            if to_add:
-                tolerations_to_add.append(toleration)
-
-        if len(tolerations_to_add) > 0:
-            self_tolerations.extend(tolerations_to_add)
-
-    def _override_required_during_scheduling_ignored_during_execution(
-        self,
-        node_selector: k8s_client.V1NodeSelector,
-        affinity_field_name: str,
-    ):
-        self._initialize_affinity(affinity_field_name)
-        self._initialize_node_affinity(affinity_field_name)
-
-        self_affinity = getattr(self, affinity_field_name)
-        self_affinity.node_affinity.required_during_scheduling_ignored_during_execution = node_selector
-
-    def enrich_function_preemption_spec(
-        self,
-        preemption_mode_field_name: str = "preemption_mode",
-        tolerations_field_name: str = "tolerations",
-        affinity_field_name: str = "affinity",
-        node_selector_field_name: str = "node_selector",
-    ):
-        """
-        Enriches function pod with the below described spec.
-        If no preemptible node configuration is provided, do nothing.
-            `allow` 	- Adds Tolerations if configured.
-                          otherwise, assume pods can be scheduled on preemptible nodes.
-                        > Purges any `affinity` / `anti-affinity` preemption related configuration
-                        > Purges preemptible node selector
-            `constrain` - Uses node-affinity to make sure pods are assigned using OR on the configured
-                          node label selectors.
-                        > Merges tolerations with preemptible tolerations.
-                        > Purges any `anti-affinity` preemption related configuration
-            `prevent`	- Prevention is done either using taints (if Tolerations were configured) or anti-affinity.
-                        > Purges any `tolerations` preemption related configuration
-                        > Purges any `affinity` preemption related configuration
-                        > Purges preemptible node selector
-                        > Sets anti-affinity and overrides any affinity if no tolerations were configured
-            `none`      - Doesn't apply any preemptible node selection configuration.
-        """
-        # nothing to do here, configuration is not populated
-        if not mlconf.is_preemption_nodes_configured():
-            return
-
-        if not getattr(self, preemption_mode_field_name):
-            # We're not supposed to get here, but if we do, we'll set the private attribute to
-            # avoid triggering circular enrichment.
-            setattr(
-                self,
-                f"_{preemption_mode_field_name}",
-                mlconf.function_defaults.preemption_mode,
-            )
-            logger.debug(
-                "No preemption mode was given, using the default preemption mode",
-                default_preemption_mode=getattr(self, preemption_mode_field_name),
-            )
-        self_preemption_mode = getattr(self, preemption_mode_field_name)
-        # don't enrich with preemption configuration.
-        if self_preemption_mode == PreemptionModes.none.value:
-            return
-        # remove preemptible tolerations and remove preemption related configuration
-        # and enrich with anti-affinity if preemptible tolerations configuration haven't been provided
-        if self_preemption_mode == PreemptionModes.prevent.value:
-            # ensure no preemptible node tolerations
-            self._prune_tolerations(
-                generate_preemptible_tolerations(),
-                tolerations_field_name=tolerations_field_name,
-            )
-
-            # purge affinity preemption related configuration
-            self._prune_affinity_node_selector_requirement(
-                generate_preemptible_node_selector_requirements(
-                    NodeSelectorOperator.node_selector_op_in.value
-                ),
-                affinity_field_name=affinity_field_name,
-            )
-            # remove preemptible nodes constrain
-            self._prune_node_selector(
-                mlconf.get_preemptible_node_selector(),
-                node_selector_field_name=node_selector_field_name,
-            )
-
-            # if tolerations are configured, simply pruning tolerations is sufficient because functions
-            # cannot be scheduled without tolerations on tainted nodes.
-            # however, if preemptible tolerations are not configured, we must use anti-affinity on preemptible nodes
-            # to ensure that the function is not scheduled on the nodes.
-            if not generate_preemptible_tolerations():
-                # using a single term with potentially multiple expressions to ensure anti-affinity
-                self._override_required_during_scheduling_ignored_during_execution(
-                    k8s_client.V1NodeSelector(
-                        node_selector_terms=generate_preemptible_nodes_anti_affinity_terms()
-                    ),
-                    affinity_field_name=affinity_field_name,
-                )
-        # enrich tolerations and override all node selector terms with preemptible node selector terms
-        elif self_preemption_mode == PreemptionModes.constrain.value:
-            # enrich with tolerations
-            self._merge_tolerations(
-                generate_preemptible_tolerations(),
-                tolerations_field_name=tolerations_field_name,
-            )
-
-            # setting required_during_scheduling_ignored_during_execution
-            # overriding other terms that have been set, and only setting terms for preemptible nodes
-            # when having multiple terms, pod scheduling is succeeded if at least one term is satisfied
-            self._override_required_during_scheduling_ignored_during_execution(
-                k8s_client.V1NodeSelector(
-                    node_selector_terms=generate_preemptible_nodes_affinity_terms()
-                ),
-                affinity_field_name=affinity_field_name,
-            )
-        # purge any affinity / anti-affinity preemption related configuration and enrich with preemptible tolerations
-        elif self_preemption_mode == PreemptionModes.allow.value:
-            # remove preemptible anti-affinity
-            self._prune_affinity_node_selector_requirement(
-                generate_preemptible_node_selector_requirements(
-                    NodeSelectorOperator.node_selector_op_not_in.value
-                ),
-                affinity_field_name=affinity_field_name,
-            )
-            # remove preemptible affinity
-            self._prune_affinity_node_selector_requirement(
-                generate_preemptible_node_selector_requirements(
-                    NodeSelectorOperator.node_selector_op_in.value
-                ),
-                affinity_field_name=affinity_field_name,
-            )
-
-            # remove preemptible nodes constrain
-            self._prune_node_selector(
-                mlconf.get_preemptible_node_selector(),
-                node_selector_field_name=node_selector_field_name,
-            )
-
-            # enrich with tolerations
-            self._merge_tolerations(
-                generate_preemptible_tolerations(),
-                tolerations_field_name=tolerations_field_name,
-            )
-
-        self._clear_affinity_if_initialized_but_empty(
-            affinity_field_name=affinity_field_name
-        )
-        self._clear_tolerations_if_initialized_but_empty(
-            tolerations_field_name=tolerations_field_name
-        )
-
-    def _clear_affinity_if_initialized_but_empty(self, affinity_field_name: str):
-        self_affinity = getattr(self, affinity_field_name)
-        if not getattr(self, affinity_field_name):
-            setattr(self, affinity_field_name, None)
-        elif (
-            not self_affinity.node_affinity
-            and not self_affinity.pod_affinity
-            and not self_affinity.pod_anti_affinity
-        ):
-            setattr(self, affinity_field_name, None)
-
-    def _clear_tolerations_if_initialized_but_empty(self, tolerations_field_name: str):
-        if not getattr(self, tolerations_field_name):
-            setattr(self, tolerations_field_name, None)
-
-    def _merge_node_selector_term_to_node_affinity(
-        self,
-        node_selector_terms: list[k8s_client.V1NodeSelectorTerm],
-        affinity_field_name: str,
-    ):
-        if not node_selector_terms:
-            return
-
-        self._initialize_affinity(affinity_field_name)
-        self._initialize_node_affinity(affinity_field_name)
-
-        self_affinity = getattr(self, affinity_field_name)
-        if not self_affinity.node_affinity.required_during_scheduling_ignored_during_execution:
-            self_affinity.node_affinity.required_during_scheduling_ignored_during_execution = k8s_client.V1NodeSelector(
-                node_selector_terms=node_selector_terms
-            )
-            return
-
-        node_selector = self_affinity.node_affinity.required_during_scheduling_ignored_during_execution
-        new_node_selector_terms = []
-
-        for node_selector_term_to_add in node_selector_terms:
-            to_add = True
-            for node_selector_term in node_selector.node_selector_terms:
-                if node_selector_term == node_selector_term_to_add:
-                    to_add = False
-                    break
-            if to_add:
-                new_node_selector_terms.append(node_selector_term_to_add)
-
-        if new_node_selector_terms:
-            node_selector.node_selector_terms += new_node_selector_terms
-
-    def _initialize_affinity(self, affinity_field_name: str):
-        if not getattr(self, affinity_field_name):
-            setattr(self, affinity_field_name, k8s_client.V1Affinity())
-
-    def _initialize_node_affinity(self, affinity_field_name: str):
-        if not getattr(getattr(self, affinity_field_name), "node_affinity"):
-            # self.affinity.node_affinity:
-            getattr(
-                self, affinity_field_name
-            ).node_affinity = k8s_client.V1NodeAffinity()
-            # self.affinity.node_affinity = k8s_client.V1NodeAffinity()
-
-    def _prune_affinity_node_selector_requirement(
-        self,
-        node_selector_requirements: list[k8s_client.V1NodeSelectorRequirement],
-        affinity_field_name: str = "affinity",
-    ):
-        """
-        Prunes given node selector requirements from affinity.
-        We are only editing required_during_scheduling_ignored_during_execution because the scheduler can't schedule
-        the pod unless the rule is met.
-        :param node_selector_requirements:
-        :return:
-        """
-        # both needs to exist to prune required affinity from spec affinity
-        self_affinity = getattr(self, affinity_field_name)
-        if not self_affinity or not node_selector_requirements:
-            return
-        if self_affinity.node_affinity:
-            node_affinity: k8s_client.V1NodeAffinity = self_affinity.node_affinity
-
-            new_required_during_scheduling_ignored_during_execution = None
-            if node_affinity.required_during_scheduling_ignored_during_execution:
-                node_selector: k8s_client.V1NodeSelector = (
-                    node_affinity.required_during_scheduling_ignored_during_execution
-                )
-                new_node_selector_terms = (
-                    self._prune_node_selector_requirements_from_node_selector_terms(
-                        node_selector_terms=node_selector.node_selector_terms,
-                        node_selector_requirements_to_prune=node_selector_requirements,
-                    )
-                )
-                # check whether there are node selector terms to add to the new list of required terms
-                if len(new_node_selector_terms) > 0:
-                    new_required_during_scheduling_ignored_during_execution = (
-                        k8s_client.V1NodeSelector(
-                            node_selector_terms=new_node_selector_terms
-                        )
-                    )
-            # if both preferred and new required are empty, clean node_affinity
-            if (
-                not node_affinity.preferred_during_scheduling_ignored_during_execution
-                and not new_required_during_scheduling_ignored_during_execution
-            ):
-                setattr(self_affinity, "node_affinity", None)
-                # self.affinity.node_affinity = None
-                return
-
-            self._initialize_affinity(affinity_field_name)
-            self._initialize_node_affinity(affinity_field_name)
-
-            # fmt: off
-            self_affinity.node_affinity.required_during_scheduling_ignored_during_execution = (
-                new_required_during_scheduling_ignored_during_execution
-            )
-            # fmt: on
-
     @staticmethod
     def _prune_node_selector_requirements_from_node_selector_terms(
         node_selector_terms: list[k8s_client.V1NodeSelectorTerm],
@@ -894,55 +594,6 @@ class KubeResourceSpec(FunctionSpec):
                 )
         return new_node_selector_terms
 
-    def _prune_tolerations(
-        self,
-        tolerations: list[k8s_client.V1Toleration],
-        tolerations_field_name: str = "tolerations",
-    ):
-        """
-        Prunes given tolerations from function spec
-        :param tolerations: tolerations to prune
-        """
-        self_tolerations = getattr(self, tolerations_field_name)
-        # both needs to exist to prune required tolerations from spec tolerations
-        if not tolerations or not self_tolerations:
-            return
-
-        # generate a list of tolerations without tolerations to prune
-        new_tolerations = []
-        for toleration in self_tolerations:
-            to_prune = False
-            for toleration_to_delete in tolerations:
-                if toleration == toleration_to_delete:
-                    to_prune = True
-                    # no need to keep going over the list provided for the current toleration
-                    break
-            if not to_prune:
-                new_tolerations.append(toleration)
-
-        # Set tolerations without tolerations to prune
-        setattr(self, tolerations_field_name, new_tolerations)
-
-    def _prune_node_selector(
-        self,
-        node_selector: dict[str, str],
-        node_selector_field_name: str,
-    ):
-        """
-        Prunes given node_selector key from function spec if their key and value are matching
-        :param node_selector: node selectors to prune
-        """
-        self_node_selector = getattr(self, node_selector_field_name)
-        # both needs to exists to prune required node_selector from the spec node selector
-        if not node_selector or not self_node_selector:
-            return
-
-        for key, value in node_selector.items():
-            if value:
-                spec_value = self_node_selector.get(key)
-                if spec_value and spec_value == value:
-                    self_node_selector.pop(key)
-
 
 class AutoMountType(str, Enum):
     none = "none"
@@ -1646,6 +1297,10 @@ def get_sanitized_attribute(spec, attribute_name: str):
         if _resolve_if_type_sanitized(attribute_name, attribute[0]):
             return attribute
 
+    return sanitize_attribute(attribute)
+
+
+def sanitize_attribute(attribute):
     api = k8s_client.ApiClient()
     return api.sanitize_for_serialization(attribute)
 

mlrun/runtimes/sparkjob/spark3job.py

@@ -288,12 +288,6 @@ class Spark3JobSpec(KubeResourceSpec):
         self._driver_preemption_mode = (
             mode or mlrun.mlconf.function_defaults.preemption_mode
         )
-        self.enrich_function_preemption_spec(
-            preemption_mode_field_name="driver_preemption_mode",
-            tolerations_field_name="driver_tolerations",
-            affinity_field_name="driver_affinity",
-            node_selector_field_name="driver_node_selector",
-        )
 
     @property
     def executor_preemption_mode(self) -> str:
@@ -304,12 +298,6 @@ class Spark3JobSpec(KubeResourceSpec):
         self._executor_preemption_mode = (
             mode or mlrun.mlconf.function_defaults.preemption_mode
         )
-        self.enrich_function_preemption_spec(
-            preemption_mode_field_name="executor_preemption_mode",
-            tolerations_field_name="executor_tolerations",
-            affinity_field_name="executor_affinity",
-            node_selector_field_name="executor_node_selector",
-        )
 
     @property
     def driver_volume_mounts(self) -> list:

mlrun/runtimes/utils.py

@@ -11,6 +11,7 @@
 # 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 enum
 import getpass
 import hashlib
 import json
@@ -28,7 +29,6 @@ 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
@@ -433,18 +433,35 @@ def enrich_function_from_dict(function, function_dict):
 
 def enrich_run_labels(
     labels: dict,
-    labels_to_enrich: Optional[list[RunLabels]] = None,
+    labels_to_enrich: Optional[list[mlrun_constants.MLRunInternalLabels]] = None,
 ):
+    """
+    Enrich the run labels with the internal labels and the labels enrichment extension
+    :param labels: The run labels dict
+    :param labels_to_enrich: The label keys to enrich from MLRunInternalLabels.default_run_labels_to_enrich
+    :return: The enriched labels dict
+    """
+    # Merge the labels with the labels enrichment extension
     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"),
+        mlrun_constants.MLRunInternalLabels.owner: os.environ.get("V3IO_USERNAME")
+        or getpass.getuser(),
+        # TODO: remove this in 1.10.0
+        mlrun_constants.MLRunInternalLabels.v3io_user: os.environ.get("V3IO_USERNAME"),
     }
-    labels_to_enrich = labels_to_enrich or RunLabels.all()
+
+    # Resolve which label keys to enrich
+    if labels_to_enrich is None:
+        labels_to_enrich = (
+            mlrun_constants.MLRunInternalLabels.default_run_labels_to_enrich()
+        )
+
+    # Enrich labels
     for label in labels_to_enrich:
+        if isinstance(label, enum.Enum):
+            label = label.value
         enrichment = labels_enrichment.get(label)
-        if label.value not in labels and enrichment:
-            labels[label.value] = enrichment
+        if label not in labels and enrichment:
+            labels[label] = enrichment
     return labels
 
 

mlrun/serving/merger.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.
-#
 
 from typing import Optional
 

mlrun/serving/remote.py

@@ -11,7 +11,7 @@
 # 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 asyncio
 import json
 from copy import copy

mlrun/serving/serving_wrapper.py

@@ -11,7 +11,7 @@
 # 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.
-#
+
 # serving runtime hooks, used in empty serving functions
 from mlrun.runtimes import nuclio_init_hook