mlrun 1.10.0rc13__py3-none-any.whl → 1.10.0rc42__py3-none-any.whl

mlrun/runtimes/pod.py

@@ -17,14 +17,17 @@ import os
 import re
 import time
 import typing
+import warnings
 from collections.abc import Iterable
 from enum import Enum
+from typing import Optional
 
 import dotenv
 import kubernetes.client as k8s_client
 from kubernetes.client import V1Volume, V1VolumeMount
 
 import mlrun.common.constants
+import mlrun.common.secrets
 import mlrun.errors
 import mlrun.runtimes.mounts
 import mlrun.utils.regex
@@ -35,6 +38,7 @@ from mlrun.common.schemas import (
 
 from ..config import config as mlconf
 from ..k8s_utils import (
+    generate_preemptible_nodes_affinity_terms,
     validate_node_selectors,
 )
 from ..utils import logger, update_in
@@ -104,6 +108,10 @@ class KubeResourceSpec(FunctionSpec):
         "security_context",
         "state_thresholds",
         "serving_spec",
+        "track_models",
+        "parameters",
+        "graph",
+        "filename",
     ]
     _default_fields_to_strip = FunctionSpec._default_fields_to_strip + [
         "volumes",
@@ -180,6 +188,9 @@ class KubeResourceSpec(FunctionSpec):
         security_context=None,
         state_thresholds=None,
         serving_spec=None,
+        track_models=None,
+        parameters=None,
+        graph=None,
     ):
         super().__init__(
             command=command,
@@ -226,6 +237,10 @@ class KubeResourceSpec(FunctionSpec):
             or mlrun.mlconf.function.spec.state_thresholds.default.to_dict()
         )
         self.serving_spec = serving_spec
+        self.track_models = track_models
+        self.parameters = parameters
+        self._graph = None
+        self.graph = graph
         # Termination grace period is internal for runtimes that have a pod termination hook hence it is not in the
         # _dict_fields and doesn't have a setter.
         self._termination_grace_period_seconds = None
@@ -303,6 +318,17 @@ class KubeResourceSpec(FunctionSpec):
     def termination_grace_period_seconds(self) -> typing.Optional[int]:
         return self._termination_grace_period_seconds
 
+    @property
+    def graph(self):
+        """states graph, holding the serving workflow/DAG topology"""
+        return self._graph
+
+    @graph.setter
+    def graph(self, graph):
+        from ..serving.states import graph_root_setter
+
+        graph_root_setter(self, graph)
+
     def _serialize_field(
         self, struct: dict, field_name: typing.Optional[str] = None, strip: bool = False
     ) -> typing.Any:
@@ -684,19 +710,45 @@ class KubeResource(BaseRuntime):
     def spec(self, spec):
         self._spec = self._verify_dict(spec, "spec", KubeResourceSpec)
 
-    def set_env_from_secret(self, name, secret=None, secret_key=None):
-        """set pod environment var from secret"""
-        secret_key = secret_key or name
+    def set_env_from_secret(
+        self,
+        name: str,
+        secret: Optional[str] = None,
+        secret_key: Optional[str] = None,
+    ):
+        """
+        Set an environment variable from a Kubernetes Secret.
+        Client-side guard forbids MLRun internal auth/project secrets; no-op on API.
+        """
+        mlrun.common.secrets.validate_not_forbidden_secret(secret)
+        key = secret_key or name
         value_from = k8s_client.V1EnvVarSource(
-            secret_key_ref=k8s_client.V1SecretKeySelector(name=secret, key=secret_key)
+            secret_key_ref=k8s_client.V1SecretKeySelector(name=secret, key=key)
         )
-        return self._set_env(name, value_from=value_from)
+        return self._set_env(name=name, value_from=value_from)
 
-    def set_env(self, name, value=None, value_from=None):
-        """set pod environment var from value"""
-        if value is not None:
-            return self._set_env(name, value=str(value))
-        return self._set_env(name, value_from=value_from)
+    def set_env(
+        self,
+        name: str,
+        value: Optional[str] = None,
+        value_from: Optional[typing.Any] = None,
+    ):
+        """
+        Set an environment variable.
+        If value comes from a Secret, validate on client-side only.
+        """
+        if value_from is not None:
+            secret_name = self._extract_secret_name_from_value_from(
+                value_from=value_from
+            )
+            if secret_name:
+                mlrun.common.secrets.validate_not_forbidden_secret(secret_name)
+            return self._set_env(name=name, value_from=value_from)
+
+        # Plain literal value path
+        return self._set_env(
+            name=name, value=(str(value) if value is not None else None)
+        )
 
     def with_annotations(self, annotations: dict):
         """set a key/value annotations in the metadata of the pod"""
@@ -853,6 +905,133 @@ class KubeResource(BaseRuntime):
         """
         self.spec.with_requests(mem, cpu, patch=patch)
 
+    @staticmethod
+    def detect_preemptible_node_selector(node_selector: dict[str, str]) -> list[str]:
+        """
+        Check whether any provided node selector matches preemptible selectors.
+
+        :param node_selector: User-provided node selector mapping.
+        :return: List of `"key='value'"` strings that match a preemptible selector.
+        """
+        preemptible_node_selector = mlconf.get_preemptible_node_selector()
+
+        return [
+            f"'{key}': '{val}'"
+            for key, val in node_selector.items()
+            if preemptible_node_selector.get(key) == val
+        ]
+
+    def detect_preemptible_tolerations(
+        self, tolerations: list[k8s_client.V1Toleration]
+    ) -> list[str]:
+        """
+        Check whether any provided toleration matches preemptible tolerations.
+
+        :param tolerations: User-provided tolerations.
+        :return: List of formatted toleration strings that are considered preemptible.
+        """
+        preemptible_tolerations = [
+            k8s_client.V1Toleration(
+                key=toleration.get("key"),
+                value=toleration.get("value"),
+                effect=toleration.get("effect"),
+            )
+            for toleration in mlconf.get_preemptible_tolerations()
+        ]
+
+        def _format_toleration(toleration):
+            return f"'{toleration.key}'='{toleration.value}' (effect: '{toleration.effect}')"
+
+        return [
+            _format_toleration(toleration)
+            for toleration in tolerations
+            if toleration in preemptible_tolerations
+        ]
+
+    def detect_preemptible_affinity(self, affinity: k8s_client.V1Affinity) -> list[str]:
+        """
+        Check whether any provided affinity rules match preemptible affinity configs.
+
+        :param affinity: User-provided affinity object.
+        :return: List of formatted expressions that overlap with preemptible terms.
+        """
+        preemptible_affinity_terms = generate_preemptible_nodes_affinity_terms()
+        conflicting_affinities = []
+
+        if (
+            affinity
+            and affinity.node_affinity
+            and affinity.node_affinity.required_during_scheduling_ignored_during_execution
+        ):
+            user_terms = affinity.node_affinity.required_during_scheduling_ignored_during_execution.node_selector_terms
+            for user_term in user_terms:
+                user_expressions = {
+                    (expr.key, expr.operator, tuple(expr.values or []))
+                    for expr in user_term.match_expressions or []
+                }
+
+                for preemptible_term in preemptible_affinity_terms:
+                    preemptible_expressions = {
+                        (expr.key, expr.operator, tuple(expr.values or []))
+                        for expr in preemptible_term.match_expressions or []
+                    }
+
+                    # Ensure operators match and preemptible expressions are present
+                    common_exprs = user_expressions & preemptible_expressions
+                    if common_exprs:
+                        formatted = ", ".join(
+                            f"'{key}  {operator}  {list(values)}'"
+                            for key, operator, values in common_exprs
+                        )
+                        conflicting_affinities.append(formatted)
+        return conflicting_affinities
+
+    def raise_preemptible_warning(
+        self,
+        node_selector: typing.Optional[dict[str, str]],
+        tolerations: typing.Optional[list[k8s_client.V1Toleration]],
+        affinity: typing.Optional[k8s_client.V1Affinity],
+    ) -> None:
+        """
+        Detect conflicts and emit a single consolidated warning if needed.
+
+        :param node_selector: User-provided node selector.
+        :param tolerations: User-provided tolerations.
+        :param affinity: User-provided affinity.
+        :warns: PreemptionWarning - Emitted when any of the provided selectors,
+                tolerations, or affinity terms match the configured preemptible
+                settings. The message lists the conflicting items.
+        """
+        conflict_messages = []
+
+        if node_selector:
+            ns_conflicts = ", ".join(
+                self.detect_preemptible_node_selector(node_selector)
+            )
+            if ns_conflicts:
+                conflict_messages.append(f"Node selectors: {ns_conflicts}")
+
+        if tolerations:
+            tol_conflicts = ", ".join(self.detect_preemptible_tolerations(tolerations))
+            if tol_conflicts:
+                conflict_messages.append(f"Tolerations: {tol_conflicts}")
+
+        if affinity:
+            affinity_conflicts = ", ".join(self.detect_preemptible_affinity(affinity))
+            if affinity_conflicts:
+                conflict_messages.append(f"Affinity: {affinity_conflicts}")
+
+        if conflict_messages:
+            warning_componentes = "; \n".join(conflict_messages)
+            warnings.warn(
+                f"Warning: based on MLRun's preemptible node configuration, the following components \n"
+                f"may be removed or adjusted at runtime:\n"
+                f"{warning_componentes}.\n"
+                "This adjustment depends on the function's preemption mode. \n"
+                "The list of potential adjusted preemptible selectors can be viewed here: "
+                "mlrun.mlconf.get_preemptible_node_selector() and mlrun.mlconf.get_preemptible_tolerations()."
+            )
+
     def with_node_selection(
         self,
         node_name: typing.Optional[str] = None,
@@ -861,18 +1040,26 @@ class KubeResource(BaseRuntime):
         tolerations: typing.Optional[list[k8s_client.V1Toleration]] = None,
     ):
         """
-        Enables to control on which k8s node the job will run
-
-        :param node_name:       The name of the k8s node
-        :param node_selector:   Label selector, only nodes with matching labels will be eligible to be picked
-        :param affinity:        Expands the types of constraints you can express - see
-                                https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/#affinity-and-anti-affinity
-                                for details
-        :param tolerations:     Tolerations are applied to pods, and allow (but do not require) the pods to schedule
-                                onto nodes with matching taints - see
-                                https://kubernetes.io/docs/concepts/scheduling-eviction/taint-and-toleration
-                                for details
+        Configure Kubernetes node scheduling for this function.
+
+        Updates one or more scheduling hints: exact node pinning, label-based selection,
+        affinity/anti-affinity rules, and taint tolerations. Passing ``None`` leaves the
+        current value unchanged; pass an empty dict/list (e.g., ``{}``, ``[]``) to clear.
+
+        :param node_name: Exact Kubernetes node name to pin the pod to.
+        :param node_selector: Mapping of label selectors. Use ``{}`` to clear.
+        :param affinity: :class:`kubernetes.client.V1Affinity` constraints.
+        :param tolerations: List of :class:`kubernetes.client.V1Toleration`. Use ``[]`` to clear.
+        :warns: PreemptionWarning - Emitted if provided selectors/tolerations/affinity
+                conflict with the function's preemption mode.
 
+        Example usage:
+            Prefer a GPU pool and allow scheduling on spot nodes::
+
+                job.with_node_selection(
+                    node_selector={"nodepool": "gpu"},
+                    tolerations=[k8s_client.V1Toleration(key="spot", operator="Exists")],
+                )
         """
         if node_name:
             self.spec.node_name = node_name
@@ -883,6 +1070,11 @@ class KubeResource(BaseRuntime):
             self.spec.affinity = affinity
         if tolerations is not None:
             self.spec.tolerations = tolerations
+        self.raise_preemptible_warning(
+            node_selector=self.spec.node_selector,
+            tolerations=self.spec.tolerations,
+            affinity=self.spec.affinity,
+        )
 
     def with_priority_class(self, name: typing.Optional[str] = None):
         """
@@ -1202,6 +1394,27 @@ class KubeResource(BaseRuntime):
 
         return self.status.state
 
+    @staticmethod
+    def _extract_secret_name_from_value_from(
+        value_from: typing.Any,
+    ) -> Optional[str]:
+        """Extract secret name from a V1EnvVarSource or dict representation."""
+        if isinstance(value_from, k8s_client.V1EnvVarSource):
+            if value_from.secret_key_ref:
+                return value_from.secret_key_ref.name
+        elif isinstance(value_from, dict):
+            value_from = (
+                value_from.get("valueFrom")
+                or value_from.get("value_from")
+                or value_from
+            )
+            secret_key_ref = (value_from or {}).get("secretKeyRef") or (
+                value_from or {}
+            ).get("secret_key_ref")
+            if isinstance(secret_key_ref, dict):
+                return secret_key_ref.get("name")
+        return None
+
 
 def _resolve_if_type_sanitized(attribute_name, attribute):
     attribute_config = sanitized_attributes[attribute_name]

mlrun/runtimes/remotesparkjob.py

@@ -59,6 +59,9 @@ class RemoteSparkSpec(KubeResourceSpec):
         security_context=None,
         state_thresholds=None,
         serving_spec=None,
+        graph=None,
+        parameters=None,
+        track_models=None,
     ):
         super().__init__(
             command=command,
@@ -89,6 +92,9 @@ class RemoteSparkSpec(KubeResourceSpec):
             security_context=security_context,
             state_thresholds=state_thresholds,
             serving_spec=serving_spec,
+            graph=graph,
+            parameters=parameters,
+            track_models=track_models,
         )
         self.provider = provider
 

mlrun/runtimes/sparkjob/spark3job.py

@@ -169,6 +169,9 @@ class Spark3JobSpec(KubeResourceSpec):
         security_context=None,
         state_thresholds=None,
         serving_spec=None,
+        graph=None,
+        parameters=None,
+        track_models=None,
     ):
         super().__init__(
             command=command,
@@ -199,6 +202,9 @@ class Spark3JobSpec(KubeResourceSpec):
             security_context=security_context,
             state_thresholds=state_thresholds,
             serving_spec=serving_spec,
+            graph=graph,
+            parameters=parameters,
+            track_models=track_models,
         )
 
         self.driver_resources = driver_resources or {}

mlrun/runtimes/utils.py

@@ -26,6 +26,7 @@ import pandas as pd
 import mlrun
 import mlrun.common.constants
 import mlrun.common.constants as mlrun_constants
+import mlrun.common.runtimes.constants
 import mlrun.common.schemas
 import mlrun.utils.regex
 from mlrun.artifacts import TableArtifact
@@ -153,6 +154,7 @@ def results_to_iter(results, runspec, execution):
 
     iter = []
     failed = 0
+    pending_retry = 0
     running = 0
     for task in results:
         if task:
@@ -164,17 +166,26 @@ def results_to_iter(results, runspec, execution):
                 "state": state,
                 "iter": id,
             }
-            if state == "error":
+            if state == mlrun.common.runtimes.constants.RunStates.error:
                 failed += 1
                 err = get_in(task, ["status", "error"], "")
-                logger.error(f"error in task  {execution.uid}:{id} - {err_to_str(err)}")
-            elif state != "completed":
+                logger.error(f"error in task {execution.uid}:{id} - {err_to_str(err)}")
+            elif state == mlrun.common.runtimes.constants.RunStates.pending_retry:
+                pending_retry += 1
+                err = get_in(task, ["status", "error"], "")
+                retry_count = get_in(task, ["status", "retry_count"], 0)
+                logger.warning(
+                    f"pending retry in task {execution.uid}:{id} - {err_to_str(err)}. Retry count: {retry_count}"
+                )
+            elif state != mlrun.common.runtimes.constants.RunStates.completed:
                 running += 1
 
             iter.append(struct)
 
     if not iter:
-        execution.set_state("completed", commit=True)
+        execution.set_state(
+            mlrun.common.runtimes.constants.RunStates.completed, commit=True
+        )
         logger.warning("Warning!, zero iteration results")
         return
     if hasattr(pd, "json_normalize"):
@@ -204,8 +215,14 @@ def results_to_iter(results, runspec, execution):
             error=f"{failed} of {len(results)} tasks failed, check logs in db for details",
             commit=False,
         )
+    elif pending_retry:
+        execution.set_state(
+            mlrun.common.runtimes.constants.RunStates.pending_retry, commit=False
+        )
     elif running == 0:
-        execution.set_state("completed", commit=False)
+        execution.set_state(
+            mlrun.common.runtimes.constants.RunStates.completed, commit=False
+        )
     execution.commit()
 
 
@@ -431,24 +448,45 @@ def enrich_function_from_dict(function, function_dict):
     return function
 
 
+def resolve_owner(
+    labels: dict,
+    owner_to_enrich: Optional[str] = None,
+):
+    """
+    Resolve the owner label value
+    :param labels: The run labels dict
+    :param auth_username: The authenticated username
+    :return: The resolved owner label value
+    """
+
+    if owner_to_enrich and (
+        labels.get("job-type") == mlrun.common.constants.JOB_TYPE_WORKFLOW_RUNNER
+        or labels.get("job-type")
+        == mlrun.common.constants.JOB_TYPE_RERUN_WORKFLOW_RUNNER
+    ):
+        return owner_to_enrich
+    else:
+        return os.environ.get("V3IO_USERNAME") or getpass.getuser()
+
+
 def enrich_run_labels(
     labels: dict,
     labels_to_enrich: Optional[list[mlrun_constants.MLRunInternalLabels]] = None,
+    owner_to_enrich: Optional[str] = None,
 ):
     """
-    Enrich the run labels with the internal labels and the labels enrichment extension
+    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
+    :param owner_to_enrich: Optional owner to enrich the labels with, if not provided will try to resolve it.
     :return: The enriched labels dict
     """
     # Merge the labels with the labels enrichment extension
     labels_enrichment = {
-        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"),
+        mlrun_constants.MLRunInternalLabels.owner: resolve_owner(
+            labels, owner_to_enrich
+        ),
     }
-
     # Resolve which label keys to enrich
     if labels_to_enrich is None:
         labels_to_enrich = (

mlrun/secrets.py

@@ -11,9 +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 json
 from ast import literal_eval
-from os import environ, getenv
+from os import environ
 from typing import Callable, Optional, Union
 
 from .utils import AzureVaultStore, list2dict
@@ -161,6 +161,9 @@ def get_secret_or_env(
     4. An MLRun-generated env. variable, mounted from a project secret (to be used in MLRun runtimes)
     5. The default value
 
+    Also supports discovering the value inside any environment variable that contains a JSON-encoded list
+    of dicts with fields: {'name': 'KEY', 'value': 'VAL', 'value_from': ...}. This fallback is applied
+    after checking normal environment variables and before returning the default.
     Example::
 
         secrets = {"KEY1": "VALUE1"}
@@ -187,18 +190,56 @@ def get_secret_or_env(
     if prefix:
         key = f"{prefix}_{key}"
 
-    value = None
     if secret_provider:
         if isinstance(secret_provider, (dict, SecretsStore)):
-            value = secret_provider.get(key)
+            secret_value = secret_provider.get(key)
         else:
-            value = secret_provider(key)
-        if value:
-            return value
+            secret_value = secret_provider(key)
+        if secret_value:
+            return secret_value
+
+    direct_environment_value = environ.get(key)
+    if direct_environment_value:
+        return direct_environment_value
+
+    json_list_value = _find_value_in_json_env_lists(key)
+    if json_list_value is not None:
+        return json_list_value
+
+    mlrun_env_key = SecretsStore.k8s_env_variable_name_for_secret(key)
+    mlrun_env_value = environ.get(mlrun_env_key)
+    if mlrun_env_value:
+        return mlrun_env_value
 
-    return (
-        value
-        or getenv(key)
-        or getenv(SecretsStore.k8s_env_variable_name_for_secret(key))
-        or default
-    )
+    return default
+
+
+def _find_value_in_json_env_lists(
+    secret_name: str,
+) -> Optional[str]:
+    """
+    Scan all environment variables. If any env var contains a JSON-encoded list
+    of dicts shaped like {'name': str, 'value': str|None, 'value_from': ...},
+    return the 'value' for the entry whose 'name' matches secret_name.
+    """
+    for environment_variable_value in environ.values():
+        if not environment_variable_value or not isinstance(
+            environment_variable_value, str
+        ):
+            continue
+        # Fast precheck to skip obvious non-JSON strings
+        first_char = environment_variable_value.lstrip()[:1]
+        if first_char not in ("[", "{"):
+            continue
+        try:
+            parsed_value = json.loads(environment_variable_value)
+        except ValueError:
+            continue
+        if isinstance(parsed_value, list):
+            for entry in parsed_value:
+                if isinstance(entry, dict) and entry.get("name") == secret_name:
+                    value_in_entry = entry.get("value")
+                    # Match original semantics: empty string is treated as "not found"
+                    if value_in_entry:
+                        return value_in_entry
+    return None

mlrun/serving/__init__.py

@@ -28,6 +28,7 @@ __all__ = [
     "Model",
     "ModelSelector",
     "MonitoredStep",
+    "LLModel",
 ]
 
 from .routers import ModelRouter, VotingEnsemble  # noqa
@@ -47,6 +48,7 @@ from .states import (
     Model,
     ModelSelector,
     MonitoredStep,
+    LLModel,
 )  # noqa
 from .v1_serving import MLModelServer, new_v1_model_server  # noqa
 from .v2_serving import V2ModelServer  # noqa