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

mlrun/k8s_utils.py

@@ -11,7 +11,9 @@
 # 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 copy
 import re
+import typing
 import warnings
 
 import kubernetes.client
@@ -228,3 +230,341 @@ def validate_node_selectors(
             handle_invalid(str(err))
             return False
     return True
+
+
+def enrich_preemption_mode(
+    preemption_mode: typing.Optional[str],
+    node_selector: dict[str, str],
+    tolerations: list[kubernetes.client.V1Toleration],
+    affinity: typing.Optional[kubernetes.client.V1Affinity],
+) -> tuple[
+    dict[str, str],
+    list[kubernetes.client.V1Toleration],
+    typing.Optional[kubernetes.client.V1Affinity],
+]:
+    """
+    Enriches a pod spec's scheduling configuration (node selector, tolerations, affinity)
+    based on the provided preemption mode.
+
+    If no preemptible node configuration is defined in the system, or the mode is `none`,
+    the original values are returned unchanged.
+
+    Modes:
+        - allow: Adds tolerations, removes preemption constraints.
+        - constrain: Requires preemptible node affinity and adds tolerations.
+        - prevent: Enforces scheduling on non-preemptible nodes using taints or anti-affinity.
+        - none: No enrichment is applied.
+    """
+    if (
+        not mlconfig.is_preemption_nodes_configured()
+        or preemption_mode == mlrun.common.schemas.PreemptionModes.none.value
+    ):
+        return node_selector, tolerations, affinity
+
+    if not preemption_mode:
+        preemption_mode = mlconfig.function_defaults.preemption_mode
+        mlrun.utils.logger.debug(
+            "No preemption mode provided, using default",
+            default_preemption_mode=preemption_mode,
+        )
+
+    enriched_node_selector = copy.deepcopy(node_selector or {})
+    enriched_tolerations = copy.deepcopy(tolerations or [])
+    enriched_affinity = copy.deepcopy(affinity)
+    preemptible_tolerations = generate_preemptible_tolerations()
+
+    if handler := _get_mode_handler(preemption_mode):
+        enriched_node_selector, enriched_tolerations, enriched_affinity = handler(
+            enriched_node_selector,
+            enriched_tolerations,
+            enriched_affinity,
+            preemptible_tolerations,
+        )
+
+    return (
+        enriched_node_selector,
+        enriched_tolerations,
+        _prune_empty_affinity(enriched_affinity),
+    )
+
+
+def _get_mode_handler(mode: str):
+    return {
+        mlrun.common.schemas.PreemptionModes.prevent: _handle_prevent_mode,
+        mlrun.common.schemas.PreemptionModes.constrain: _handle_constrain_mode,
+        mlrun.common.schemas.PreemptionModes.allow: _handle_allow_mode,
+    }.get(mode)
+
+
+def _handle_prevent_mode(
+    node_selector: dict[str, str],
+    tolerations: list[kubernetes.client.V1Toleration],
+    affinity: typing.Optional[kubernetes.client.V1Affinity],
+    preemptible_tolerations: list[kubernetes.client.V1Toleration],
+) -> tuple[
+    dict[str, str],
+    list[kubernetes.client.V1Toleration],
+    typing.Optional[kubernetes.client.V1Affinity],
+]:
+    # Ensure no preemptible node tolerations
+    tolerations = [t for t in tolerations if t not in preemptible_tolerations]
+
+    # Purge affinity preemption-related configuration
+    affinity = _prune_affinity_node_selector_requirement(
+        generate_preemptible_node_selector_requirements(
+            mlrun.common.schemas.NodeSelectorOperator.node_selector_op_in.value
+        ),
+        affinity=affinity,
+    )
+
+    # Remove preemptible nodes constraint
+    node_selector = _prune_node_selector(
+        mlconfig.get_preemptible_node_selector(),
+        enriched_node_selector=node_selector,
+    )
+
+    # Use anti-affinity only if no tolerations configured
+    if not preemptible_tolerations:
+        affinity = _override_required_during_scheduling_ignored_during_execution(
+            kubernetes.client.V1NodeSelector(
+                node_selector_terms=generate_preemptible_nodes_anti_affinity_terms()
+            ),
+            affinity,
+        )
+
+    return node_selector, tolerations, affinity
+
+
+def _handle_constrain_mode(
+    node_selector: dict[str, str],
+    tolerations: list[kubernetes.client.V1Toleration],
+    affinity: typing.Optional[kubernetes.client.V1Affinity],
+    preemptible_tolerations: list[kubernetes.client.V1Toleration],
+) -> tuple[
+    dict[str, str],
+    list[kubernetes.client.V1Toleration],
+    typing.Optional[kubernetes.client.V1Affinity],
+]:
+    tolerations = _merge_tolerations(tolerations, preemptible_tolerations)
+
+    affinity = _override_required_during_scheduling_ignored_during_execution(
+        kubernetes.client.V1NodeSelector(
+            node_selector_terms=generate_preemptible_nodes_affinity_terms()
+        ),
+        affinity=affinity,
+    )
+
+    return node_selector, tolerations, affinity
+
+
+def _handle_allow_mode(
+    node_selector: dict[str, str],
+    tolerations: list[kubernetes.client.V1Toleration],
+    affinity: typing.Optional[kubernetes.client.V1Affinity],
+    preemptible_tolerations: list[kubernetes.client.V1Toleration],
+) -> tuple[
+    dict[str, str],
+    list[kubernetes.client.V1Toleration],
+    typing.Optional[kubernetes.client.V1Affinity],
+]:
+    for op in [
+        mlrun.common.schemas.NodeSelectorOperator.node_selector_op_not_in.value,
+        mlrun.common.schemas.NodeSelectorOperator.node_selector_op_in.value,
+    ]:
+        affinity = _prune_affinity_node_selector_requirement(
+            generate_preemptible_node_selector_requirements(op),
+            affinity=affinity,
+        )
+
+    node_selector = _prune_node_selector(
+        mlconfig.get_preemptible_node_selector(),
+        enriched_node_selector=node_selector,
+    )
+
+    tolerations = _merge_tolerations(tolerations, preemptible_tolerations)
+    return node_selector, tolerations, affinity
+
+
+def _merge_tolerations(
+    existing: list[kubernetes.client.V1Toleration],
+    to_add: list[kubernetes.client.V1Toleration],
+) -> list[kubernetes.client.V1Toleration]:
+    for toleration in to_add:
+        if toleration not in existing:
+            existing.append(toleration)
+    return existing
+
+
+def _prune_node_selector(
+    node_selector: dict[str, str],
+    enriched_node_selector: dict[str, str],
+):
+    """
+    Prunes given node_selector key from function spec if their key and value are matching
+    :param node_selector: node selectors to prune
+    """
+    # both needs to exists to prune required node_selector from the spec node selector
+    if not node_selector or not enriched_node_selector:
+        return
+
+    mlrun.utils.logger.debug("Pruning node selectors", node_selector=node_selector)
+    return {
+        key: value
+        for key, value in enriched_node_selector.items()
+        if node_selector.get(key) != value
+    }
+
+
+def _prune_affinity_node_selector_requirement(
+    node_selector_requirements: list[kubernetes.client.V1NodeSelectorRequirement],
+    affinity: typing.Optional[kubernetes.client.V1Affinity],
+):
+    """
+    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
+    if not affinity or not node_selector_requirements:
+        return
+    if affinity.node_affinity:
+        node_affinity: kubernetes.client.V1NodeAffinity = affinity.node_affinity
+
+        new_required_during_scheduling_ignored_during_execution = None
+        if node_affinity.required_during_scheduling_ignored_during_execution:
+            node_selector: kubernetes.client.V1NodeSelector = (
+                node_affinity.required_during_scheduling_ignored_during_execution
+            )
+            new_node_selector_terms = (
+                _prune_node_selector_requirements_from_node_selector_terms(
+                    node_selector_terms=node_selector.node_selector_terms,
+                    requirements_to_prune=node_selector_requirements,
+                )
+            )
+            # check whether there are node selector terms to add to the new list of required terms
+            if new_node_selector_terms:
+                new_required_during_scheduling_ignored_during_execution = (
+                    kubernetes.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
+        ):
+            affinity.node_affinity = None
+            return
+
+        _initialize_affinity(affinity=affinity)
+        _initialize_node_affinity(affinity=affinity)
+
+        affinity.node_affinity.required_during_scheduling_ignored_during_execution = (
+            new_required_during_scheduling_ignored_during_execution
+        )
+        return affinity
+
+
+def _prune_node_selector_requirements_from_node_selector_terms(
+    node_selector_terms: list[kubernetes.client.V1NodeSelectorTerm],
+    requirements_to_prune: list[kubernetes.client.V1NodeSelectorRequirement],
+) -> list[kubernetes.client.V1NodeSelectorTerm]:
+    """
+    Removes matching node selector requirements from the given list of node selector terms.
+
+    Each term may contain multiple match expressions. This function iterates over each expression,
+    and removes any that exactly match one of the requirements provided.
+
+    :param node_selector_terms: List of V1NodeSelectorTerm objects to be processed.
+    :param requirements_to_prune: List of V1NodeSelectorRequirement objects to remove.
+    :return: A new list of V1NodeSelectorTerm objects with the specified requirements pruned.
+    """
+    pruned_terms = []
+
+    for term in node_selector_terms:
+        remaining_requirements = [
+            expr
+            for expr in term.match_expressions or []
+            if expr not in requirements_to_prune
+        ]
+
+        # Only add term if there are remaining match expressions or match fields
+        if remaining_requirements or term.match_fields:
+            pruned_terms.append(
+                kubernetes.client.V1NodeSelectorTerm(
+                    match_expressions=remaining_requirements,
+                    match_fields=term.match_fields,
+                )
+            )
+
+    return pruned_terms
+
+
+def _override_required_during_scheduling_ignored_during_execution(
+    node_selector: kubernetes.client.V1NodeSelector,
+    affinity: typing.Optional[kubernetes.client.V1Affinity],
+):
+    affinity = _initialize_affinity(affinity)
+    affinity = _initialize_node_affinity(affinity)
+    affinity.node_affinity.required_during_scheduling_ignored_during_execution = (
+        node_selector
+    )
+    return affinity
+
+
+def _initialize_affinity(
+    affinity: typing.Optional[kubernetes.client.V1Affinity],
+) -> kubernetes.client.V1Affinity:
+    return affinity or kubernetes.client.V1Affinity()
+
+
+def _initialize_node_affinity(
+    affinity: typing.Optional[kubernetes.client.V1Affinity],
+) -> kubernetes.client.V1Affinity:
+    affinity = affinity or kubernetes.client.V1Affinity()
+    affinity.node_affinity = (
+        affinity.node_affinity or kubernetes.client.V1NodeAffinity()
+    )
+    return affinity
+
+
+def _prune_empty_affinity(
+    affinity: typing.Optional[kubernetes.client.V1Affinity],
+) -> typing.Optional[kubernetes.client.V1Affinity]:
+    """
+    Return None if the given affinity object has no meaningful constraints.
+
+    Keeps the affinity object only if it contains:
+    - Any pod affinity or pod anti-affinity
+    - Preferred node affinity
+    - Required node affinity with at least one match expression or match field
+    """
+    if not affinity:
+        return None
+
+    node_affinity = affinity.node_affinity
+    pod_affinity = affinity.pod_affinity
+    pod_anti_affinity = affinity.pod_anti_affinity
+
+    # If any pod affinity exists, keep the object
+    if pod_affinity or pod_anti_affinity:
+        return affinity
+
+    # If node affinity exists, check if it has any meaningful content
+    if node_affinity:
+        required = node_affinity.required_during_scheduling_ignored_during_execution
+        preferred = node_affinity.preferred_during_scheduling_ignored_during_execution
+
+        if preferred:
+            return affinity
+
+        if required and required.node_selector_terms:
+            for term in required.node_selector_terms:
+                if term.match_expressions or term.match_fields:
+                    return affinity  # at least one term has meaningful constraints
+
+    # At this point, none of the affinity sections contain meaningful constraints,
+    # so the affinity object is effectively empty and can be safely discarded.
+    return None

mlrun/launcher/base.py

@@ -57,6 +57,7 @@ class BaseLauncher(abc.ABC):
         out_path: Optional[str] = "",
         workdir: Optional[str] = "",
         artifact_path: Optional[str] = "",
+        output_path: Optional[str] = "",
         watch: Optional[bool] = True,
         schedule: Optional[
             Union[str, mlrun.common.schemas.schedule.ScheduleCronTrigger]
@@ -234,8 +235,7 @@ class BaseLauncher(abc.ABC):
         hyper_param_options=None,
         verbose=None,
         scrape_metrics=None,
-        out_path=None,
-        artifact_path=None,
+        output_path=None,
         workdir=None,
         notifications: Optional[list[mlrun.model.Notification]] = None,
         state_thresholds: Optional[dict[str, int]] = None,
@@ -301,7 +301,7 @@ class BaseLauncher(abc.ABC):
         meta = run.metadata
         meta.uid = meta.uid or uuid.uuid4().hex
 
-        run.spec.output_path = out_path or artifact_path or run.spec.output_path
+        run.spec.output_path = output_path or run.spec.output_path
 
         if not run.spec.output_path:
             if run.metadata.project:

mlrun/launcher/local.py

@@ -55,6 +55,7 @@ class ClientLocalLauncher(launcher.ClientBaseLauncher):
         out_path: Optional[str] = "",
         workdir: Optional[str] = "",
         artifact_path: Optional[str] = "",
+        output_path: Optional[str] = "",
         watch: Optional[bool] = True,
         schedule: Optional[
             Union[str, mlrun.common.schemas.schedule.ScheduleCronTrigger]
@@ -116,8 +117,7 @@ class ClientLocalLauncher(launcher.ClientBaseLauncher):
             hyper_param_options=hyper_param_options,
             verbose=verbose,
             scrape_metrics=scrape_metrics,
-            out_path=out_path,
-            artifact_path=artifact_path,
+            output_path=output_path,
             workdir=workdir,
             notifications=notifications,
             state_thresholds=state_thresholds,

mlrun/launcher/remote.py

@@ -45,6 +45,7 @@ class ClientRemoteLauncher(launcher.ClientBaseLauncher):
         out_path: Optional[str] = "",
         workdir: Optional[str] = "",
         artifact_path: Optional[str] = "",
+        output_path: Optional[str] = "",
         watch: Optional[bool] = True,
         schedule: Optional[
             Union[str, mlrun.common.schemas.schedule.ScheduleCronTrigger]
@@ -77,8 +78,7 @@ class ClientRemoteLauncher(launcher.ClientBaseLauncher):
             hyper_param_options=hyper_param_options,
             verbose=verbose,
             scrape_metrics=scrape_metrics,
-            out_path=out_path,
-            artifact_path=artifact_path,
+            output_path=output_path,
             workdir=workdir,
             notifications=notifications,
             state_thresholds=state_thresholds,

mlrun/model.py

@@ -929,6 +929,8 @@ class RunSpec(ModelObj):
 
     _fields_to_serialize = ModelObj._fields_to_serialize + [
         "handler",
+        "affinity",
+        "tolerations",
     ]
 
     def __init__(
@@ -956,6 +958,8 @@ class RunSpec(ModelObj):
         state_thresholds=None,
         reset_on_run=None,
         node_selector=None,
+        tolerations=None,
+        affinity=None,
     ):
         # A dictionary of parsing configurations that will be read from the inputs the user set. The keys are the inputs
         # keys (parameter names) and the values are the type hint given in the input keys after the colon.
@@ -994,6 +998,8 @@ class RunSpec(ModelObj):
         self.state_thresholds = state_thresholds or {}
         self.reset_on_run = reset_on_run
         self.node_selector = node_selector or {}
+        self.tolerations = tolerations or {}
+        self.affinity = affinity or {}
 
     def _serialize_field(
         self, struct: dict, field_name: Optional[str] = None, strip: bool = False
@@ -1003,6 +1009,14 @@ class RunSpec(ModelObj):
             if self.handler and isinstance(self.handler, str):
                 return self.handler
             return None
+
+        # Properly serialize known K8s objects
+        if field_name in {"affinity", "tolerations"}:
+            value = getattr(self, field_name, None)
+            if hasattr(value, "to_dict"):
+                return value.to_dict()
+            return value
+
         return super()._serialize_field(struct, field_name, strip)
 
     def is_hyper_job(self):

mlrun/model_monitoring/applications/__init__.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 .base import ModelMonitoringApplicationBase
 from .context import MonitoringApplicationContext

mlrun/model_monitoring/applications/_application_steps.py

@@ -96,7 +96,9 @@ class _PushToMonitoringWriter(StepToDict):
             logger.debug(
                 "Pushing data to output stream", writer_event=str(writer_event)
             )
-            self.output_stream.push([writer_event])
+            self.output_stream.push(
+                [writer_event], partition_key=application_context.endpoint_id
+            )
             logger.debug("Pushed data to output stream successfully")
 
     def _lazy_init(self):

mlrun/model_monitoring/controller.py

@@ -673,7 +673,9 @@ class MonitoringApplicationController:
         """
         logger.info("Starting monitoring controller chief")
         applications_names = []
-        endpoints = self.project_obj.list_model_endpoints(tsdb_metrics=True).endpoints
+        endpoints = self.project_obj.list_model_endpoints(
+            metric_list=["last_request"]
+        ).endpoints
         if not endpoints:
             logger.info("No model endpoints found", project=self.project)
             return

mlrun/model_monitoring/db/tsdb/base.py

@@ -82,7 +82,8 @@ class TSDBConnector(ABC):
 
     @abstractmethod
     def delete_tsdb_records(
-        self, endpoint_ids: list[str], delete_timeout: Optional[int] = None
+        self,
+        endpoint_ids: list[str],
     ) -> None:
         """
         Delete model endpoint records from the TSDB connector.
@@ -332,6 +333,7 @@ class TSDBConnector(ABC):
         model_endpoint_objects: list[mlrun.common.schemas.ModelEndpoint],
         project: str,
         run_in_threadpool: Callable,
+        metric_list: Optional[list[str]] = None,
     ) -> list[mlrun.common.schemas.ModelEndpoint]:
         raise NotImplementedError()