mlrun 1.7.2rc4__py3-none-any.whl → 1.8.0__py3-none-any.whl

mlrun/runtimes/pod.py

@@ -17,14 +17,16 @@ import os
 import re
 import time
 import typing
+from collections.abc import Iterable
 from enum import Enum
 
 import dotenv
 import kubernetes.client as k8s_client
-import mlrun_pipelines.mounts
-from mlrun_pipelines.mixins import KfpAdapterMixin
+from kubernetes.client import V1Volume, V1VolumeMount
 
+import mlrun.common.constants
 import mlrun.errors
+import mlrun.runtimes.mounts
 import mlrun.utils.regex
 from mlrun.common.schemas import (
     NodeSelectorOperator,
@@ -212,9 +214,7 @@ class KubeResourceSpec(FunctionSpec):
         # default service account is set in mlrun.utils.process_function_service_account
         # due to project specific defaults
         self.service_account = service_account
-        self.image_pull_secret = (
-            image_pull_secret or mlrun.mlconf.function.spec.image_pull_secret.default
-        )
+        self.image_pull_secret = image_pull_secret
         self.node_name = node_name
         self.node_selector = node_selector or {}
         self._affinity = affinity
@@ -309,7 +309,7 @@ class KubeResourceSpec(FunctionSpec):
         return self._termination_grace_period_seconds
 
     def _serialize_field(
-        self, struct: dict, field_name: str = None, strip: bool = False
+        self, struct: dict, field_name: typing.Optional[str] = None, strip: bool = False
     ) -> typing.Any:
         """
         Serialize a field to a dict, list, or primitive type.
@@ -321,7 +321,7 @@ class KubeResourceSpec(FunctionSpec):
         return super()._serialize_field(struct, field_name, strip)
 
     def _enrich_field(
-        self, struct: dict, field_name: str = None, strip: bool = False
+        self, struct: dict, field_name: typing.Optional[str] = None, strip: bool = False
     ) -> typing.Any:
         k8s_api = k8s_client.ApiClient()
         if strip:
@@ -366,6 +366,35 @@ class KubeResourceSpec(FunctionSpec):
                 + f"service accounts {allowed_service_accounts}"
             )
 
+    def with_volumes(
+        self,
+        volumes: typing.Union[list[dict], dict, V1Volume],
+    ) -> "KubeResourceSpec":
+        """Add volumes to the volumes dictionary, only used as part of the mlrun_pipelines mount functions."""
+        if isinstance(volumes, dict):
+            set_named_item(self._volumes, volumes)
+        elif isinstance(volumes, Iterable):
+            for volume in volumes:
+                set_named_item(self._volumes, volume)
+        else:
+            set_named_item(self._volumes, volumes)
+        return self
+
+    def with_volume_mounts(
+        self,
+        volume_mounts: typing.Union[list[dict], dict, V1VolumeMount],
+    ) -> "KubeResourceSpec":
+        """Add volume mounts to the volume mounts dictionary,
+        only used as part of the mlrun_pipelines mount functions."""
+        if isinstance(volume_mounts, dict):
+            self._set_volume_mount(volume_mounts)
+        elif isinstance(volume_mounts, Iterable):
+            for volume_mount in volume_mounts:
+                self._set_volume_mount(volume_mount)
+        else:
+            self._set_volume_mount(volume_mounts)
+        return self
+
     def _set_volume_mount(
         self, volume_mount, volume_mounts_field_name="_volume_mounts"
     ):
@@ -380,9 +409,9 @@ class KubeResourceSpec(FunctionSpec):
     def _verify_and_set_limits(
         self,
         resources_field_name,
-        mem: str = None,
-        cpu: str = None,
-        gpus: int = None,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
+        gpus: typing.Optional[int] = None,
         gpu_type: str = "nvidia.com/gpu",
         patch: bool = False,
     ):
@@ -430,8 +459,8 @@ class KubeResourceSpec(FunctionSpec):
     def _verify_and_set_requests(
         self,
         resources_field_name,
-        mem: str = None,
-        cpu: str = None,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
         patch: bool = False,
     ):
         resources = verify_requests(resources_field_name, mem=mem, cpu=cpu)
@@ -456,9 +485,9 @@ class KubeResourceSpec(FunctionSpec):
 
     def with_limits(
         self,
-        mem: str = None,
-        cpu: str = None,
-        gpus: int = None,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
+        gpus: typing.Optional[int] = None,
         gpu_type: str = "nvidia.com/gpu",
         patch: bool = False,
     ):
@@ -474,7 +503,12 @@ class KubeResourceSpec(FunctionSpec):
         """
         self._verify_and_set_limits("resources", mem, cpu, gpus, gpu_type, patch=patch)
 
-    def with_requests(self, mem: str = None, cpu: str = None, patch: bool = False):
+    def with_requests(
+        self,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
+        patch: bool = False,
+    ):
         """
         Set requested (desired) pod cpu/memory resources
 
@@ -936,12 +970,12 @@ class AutoMountType(str, Enum):
     @classmethod
     def all_mount_modifiers(cls):
         return [
-            mlrun_pipelines.mounts.v3io_cred.__name__,
-            mlrun_pipelines.mounts.mount_v3io.__name__,
-            mlrun_pipelines.mounts.mount_pvc.__name__,
-            mlrun_pipelines.mounts.auto_mount.__name__,
-            mlrun_pipelines.mounts.mount_s3.__name__,
-            mlrun_pipelines.mounts.set_env_variables.__name__,
+            mlrun.runtimes.mounts.v3io_cred.__name__,
+            mlrun.runtimes.mounts.mount_v3io.__name__,
+            mlrun.runtimes.mounts.mount_pvc.__name__,
+            mlrun.runtimes.mounts.auto_mount.__name__,
+            mlrun.runtimes.mounts.mount_s3.__name__,
+            mlrun.runtimes.mounts.set_env_variables.__name__,
         ]
 
     @classmethod
@@ -958,27 +992,27 @@ class AutoMountType(str, Enum):
     def _get_auto_modifier():
         # If we're running on Iguazio - use v3io_cred
         if mlconf.igz_version != "":
-            return mlrun_pipelines.mounts.v3io_cred
+            return mlrun.runtimes.mounts.v3io_cred
         # Else, either pvc mount if it's configured or do nothing otherwise
         pvc_configured = (
             "MLRUN_PVC_MOUNT" in os.environ
             or "pvc_name" in mlconf.get_storage_auto_mount_params()
         )
-        return mlrun_pipelines.mounts.mount_pvc if pvc_configured else None
+        return mlrun.runtimes.mounts.mount_pvc if pvc_configured else None
 
     def get_modifier(self):
         return {
             AutoMountType.none: None,
-            AutoMountType.v3io_credentials: mlrun_pipelines.mounts.v3io_cred,
-            AutoMountType.v3io_fuse: mlrun_pipelines.mounts.mount_v3io,
-            AutoMountType.pvc: mlrun_pipelines.mounts.mount_pvc,
+            AutoMountType.v3io_credentials: mlrun.runtimes.mounts.v3io_cred,
+            AutoMountType.v3io_fuse: mlrun.runtimes.mounts.mount_v3io,
+            AutoMountType.pvc: mlrun.runtimes.mounts.mount_pvc,
             AutoMountType.auto: self._get_auto_modifier(),
-            AutoMountType.s3: mlrun_pipelines.mounts.mount_s3,
-            AutoMountType.env: mlrun_pipelines.mounts.set_env_variables,
+            AutoMountType.s3: mlrun.runtimes.mounts.mount_s3,
+            AutoMountType.env: mlrun.runtimes.mounts.set_env_variables,
         }[self]
 
 
-class KubeResource(BaseRuntime, KfpAdapterMixin):
+class KubeResource(BaseRuntime):
     """
     A parent class for runtimes that generate k8s resources when executing.
     """
@@ -1020,7 +1054,8 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
 
     def get_env(self, name, default=None):
         """Get the pod environment variable for the given name, if not found return the default
-        If it's a scalar value, will return it, if the value is from source, return the k8s struct (V1EnvVarSource)"""
+        If it's a scalar value, will return it, if the value is from source, return the k8s struct (V1EnvVarSource)
+        """
         for env_var in self.spec.env:
             if get_item_name(env_var) == name:
                 # valueFrom is a workaround for now, for some reason the envs aren't getting sanitized
@@ -1050,7 +1085,11 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
         self.spec.env.append(new_var)
         return self
 
-    def set_envs(self, env_vars: dict = None, file_path: str = None):
+    def set_envs(
+        self,
+        env_vars: typing.Optional[dict] = None,
+        file_path: typing.Optional[str] = None,
+    ):
         """set pod environment var from key/value dict or .env file
 
         :param env_vars:  dict with env key/values
@@ -1070,11 +1109,16 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
             else:
                 raise mlrun.errors.MLRunNotFoundError(f"{file_path} does not exist")
         for name, value in env_vars.items():
-            self.set_env(name, value)
+            if isinstance(value, dict) and "valueFrom" in value:
+                self.set_env(name, value_from=value["valueFrom"])
+            else:
+                self.set_env(name, value)
         return self
 
     def set_image_pull_configuration(
-        self, image_pull_policy: str = None, image_pull_secret_name: str = None
+        self,
+        image_pull_policy: typing.Optional[str] = None,
+        image_pull_secret_name: typing.Optional[str] = None,
     ):
         """
         Configure the image pull parameters for the runtime.
@@ -1123,9 +1167,9 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
 
     def with_limits(
         self,
-        mem: str = None,
-        cpu: str = None,
-        gpus: int = None,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
+        gpus: typing.Optional[int] = None,
         gpu_type: str = "nvidia.com/gpu",
         patch: bool = False,
     ):
@@ -1141,7 +1185,12 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
         """
         self.spec.with_limits(mem, cpu, gpus, gpu_type, patch=patch)
 
-    def with_requests(self, mem: str = None, cpu: str = None, patch: bool = False):
+    def with_requests(
+        self,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
+        patch: bool = False,
+    ):
         """
         Set requested (desired) pod cpu/memory resources
 
@@ -1253,6 +1302,36 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
             )
         self.spec.security_context = security_context
 
+    def apply(
+        self,
+        modifier: typing.Callable[["KubeResource"], "KubeResource"],
+    ) -> "KubeResource":
+        """
+        Apply a modifier to the runtime which is used to change the runtimes k8s object's spec.
+        All modifiers accept Kube, apply some changes on its spec and return it so modifiers can be chained
+        one after the other.
+
+        :param modifier: a modifier callable object
+        :return: the runtime (self) after the modifications
+        """
+        modifier(self)
+        if AutoMountType.is_auto_modifier(modifier):
+            self.spec.disable_auto_mount = True
+
+        api_client = k8s_client.ApiClient()
+        if self.spec.env:
+            for index, env in enumerate(
+                api_client.sanitize_for_serialization(self.spec.env)
+            ):
+                self.spec.env[index] = env
+
+        if self.spec.volumes and self.spec.volume_mounts:
+            vols = api_client.sanitize_for_serialization(self.spec.volumes)
+            mounts = api_client.sanitize_for_serialization(self.spec.volume_mounts)
+            self.spec.update_vols_and_mounts(vols, mounts)
+
+        return self
+
     def list_valid_priority_class_names(self):
         return mlconf.get_valid_function_priority_class_names()
 
@@ -1393,15 +1472,13 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
                 f"Started building image: {data.get('data', {}).get('spec', {}).get('build', {}).get('image')}"
             )
         if watch and not ready:
-            state = self._build_watch(
+            self.status.state = self._build_watch(
                 watch=watch,
                 show_on_failure=show_on_failure,
             )
-            ready = state == "ready"
-            self.status.state = state
-
-        if watch and not ready:
-            raise mlrun.errors.MLRunRuntimeError("Deploy failed")
+            ready = self.status.state == "ready"
+            if not ready:
+                raise mlrun.errors.MLRunRuntimeError("Deploy failed")
         return ready
 
     def _build_watch(
@@ -1412,20 +1489,32 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
     ):
         db = self._get_db()
         offset = 0
+        events_offset = 0
         try:
-            text, _ = db.get_builder_status(self, 0, logs=logs)
+            text, _, deploy_status_text_kind = db.get_builder_status(
+                self,
+                offset=0,
+                logs=logs,
+                events_offset=0,
+            )
         except mlrun.db.RunDBError:
-            raise ValueError("function or build process not found")
+            raise ValueError("Function or build process not found")
 
-        def print_log(text):
-            if text and (
+        def print_log(_text):
+            if _text and (
                 not show_on_failure
                 or self.status.state == mlrun.common.schemas.FunctionState.error
             ):
-                print(text, end="")
+                print(_text, end="")
 
         print_log(text)
-        offset += len(text)
+        if (
+            deploy_status_text_kind
+            == mlrun.common.constants.DeployStatusTextKind.events
+        ):
+            events_offset += len(text)
+        else:
+            offset += len(text)
         if watch:
             while self.status.state in [
                 mlrun.common.schemas.FunctionState.pending,
@@ -1434,14 +1523,30 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
                 time.sleep(2)
                 if show_on_failure:
                     text = ""
-                    db.get_builder_status(self, 0, logs=False)
+                    db.get_builder_status(self, offset=0, logs=False, events_offset=0)
                     if self.status.state == mlrun.common.schemas.FunctionState.error:
                         # re-read the full log on failure
-                        text, _ = db.get_builder_status(self, offset, logs=logs)
+                        text, _, deploy_status_text_kind = db.get_builder_status(
+                            self,
+                            offset=offset,
+                            logs=logs,
+                            events_offset=events_offset,
+                        )
                 else:
-                    text, _ = db.get_builder_status(self, offset, logs=logs)
+                    text, _, deploy_status_text_kind = db.get_builder_status(
+                        self,
+                        offset=offset,
+                        logs=logs,
+                        events_offset=events_offset,
+                    )
                 print_log(text)
-                offset += len(text)
+                if (
+                    deploy_status_text_kind
+                    == mlrun.common.constants.DeployStatusTextKind.events
+                ):
+                    events_offset += len(text)
+                else:
+                    offset += len(text)
 
         return self.status.state
 

mlrun/runtimes/remotesparkjob.py

@@ -13,12 +13,13 @@
 # limitations under the License.
 import re
 from subprocess import run
+from typing import Optional
 
 import kubernetes.client
-from mlrun_pipelines.mounts import mount_v3io, mount_v3iod
 
 import mlrun.errors
 from mlrun.config import config
+from mlrun.runtimes.mounts import mount_v3io, mount_v3iod
 
 from .kubejob import KubejobRuntime
 from .pod import KubeResourceSpec
@@ -179,7 +180,7 @@ class RemoteSparkRuntime(KubejobRuntime):
         skip_deployed=False,
         is_kfp=False,
         mlrun_version_specifier=None,
-        builder_env: dict = None,
+        builder_env: Optional[dict] = None,
         show_on_failure: bool = False,
         force_build: bool = False,
     ):

mlrun/runtimes/sparkjob/__init__.py

@@ -12,6 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
-
 from .spark3job import Spark3Runtime

mlrun/runtimes/sparkjob/spark3job.py

@@ -14,13 +14,13 @@
 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 mlrun.runtimes.mounts import mount_v3io, mount_v3iod
 
 from ...execution import MLClientCtx
 from ...model import RunObject
@@ -405,8 +405,8 @@ class Spark3JobSpec(KubeResourceSpec):
     def _verify_and_set_requests(
         self,
         resources_field_name,
-        mem: str = None,
-        cpu: str = None,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
         patch: bool = False,
     ):
         # Spark operator uses JVM notation for memory, so we must verify it separately
@@ -774,7 +774,11 @@ class Spark3Runtime(KubejobRuntime):
                 exporter_jar="/spark/jars/jmx_prometheus_javaagent-0.16.1.jar",
             )
 
-    def with_cores(self, executor_cores: int = None, driver_cores: int = None):
+    def with_cores(
+        self,
+        executor_cores: typing.Optional[int] = None,
+        driver_cores: typing.Optional[int] = None,
+    ):
         """
         Allows to configure spark.executor.cores and spark.driver.cores parameters. The values must be integers
         greater than or equal to 1. If a parameter is not specified, it defaults to 1.
@@ -785,7 +789,7 @@ class Spark3Runtime(KubejobRuntime):
         Spark itself uses the spec.[executor|driver].cores parameter to set the parallelism of tasks and cores
         assigned to each task within the pod. This function sets the .cores parameters for the job executed.
 
-        See https://github.com/GoogleCloudPlatform/spark-on-k8s-operator/issues/581 for a discussion about those
+        See https://github.com/kubeflow/spark-operator/issues/581 for a discussion about those
         parameters and their meaning in Spark operator.
 
         :param executor_cores: Number of cores to use for executor (spark.executor.cores)
@@ -853,7 +857,7 @@ class Spark3Runtime(KubejobRuntime):
         skip_deployed=False,
         is_kfp=False,
         mlrun_version_specifier=None,
-        builder_env: dict = None,
+        builder_env: typing.Optional[dict] = None,
         show_on_failure: bool = False,
         force_build: bool = False,
     ):
@@ -944,7 +948,10 @@ class Spark3Runtime(KubejobRuntime):
         )
 
     def with_executor_requests(
-        self, mem: str = None, cpu: str = None, patch: bool = False
+        self,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
+        patch: bool = False,
     ):
         """
         set executor pod required cpu/memory/gpu resources
@@ -954,8 +961,8 @@ class Spark3Runtime(KubejobRuntime):
 
     def with_executor_limits(
         self,
-        cpu: str = None,
-        gpus: int = None,
+        cpu: typing.Optional[str] = None,
+        gpus: typing.Optional[int] = None,
         gpu_type: str = "nvidia.com/gpu",
         patch: bool = False,
     ):
@@ -970,7 +977,10 @@ class Spark3Runtime(KubejobRuntime):
         )
 
     def with_driver_requests(
-        self, mem: str = None, cpu: str = None, patch: bool = False
+        self,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
+        patch: bool = False,
     ):
         """
         set driver pod required cpu/memory/gpu resources
@@ -980,8 +990,8 @@ class Spark3Runtime(KubejobRuntime):
 
     def with_driver_limits(
         self,
-        cpu: str = None,
-        gpus: int = None,
+        cpu: typing.Optional[str] = None,
+        gpus: typing.Optional[int] = None,
         gpu_type: str = "nvidia.com/gpu",
         patch: bool = False,
     ):

mlrun/runtimes/utils.py

@@ -18,6 +18,7 @@ import os
 import re
 from io import StringIO
 from sys import stderr
+from typing import Optional
 
 import pandas as pd
 
@@ -147,7 +148,7 @@ def add_code_metadata(path=""):
 
 def results_to_iter(results, runspec, execution):
     if not results:
-        logger.error("got an empty results list in to_iter")
+        logger.error("Got an empty results list in to_iter")
         return
 
     iter = []
@@ -174,7 +175,7 @@ def results_to_iter(results, runspec, execution):
 
     if not iter:
         execution.set_state("completed", commit=True)
-        logger.warning("warning!, zero iteration results")
+        logger.warning("Warning!, zero iteration results")
         return
     if hasattr(pd, "json_normalize"):
         df = pd.json_normalize(iter).sort_values("iter")
@@ -189,10 +190,10 @@ def results_to_iter(results, runspec, execution):
     item, id = selector(results, criteria)
     if runspec.spec.selector and not id:
         logger.warning(
-            f"no best result selected, check selector ({criteria}) or results"
+            f"No best result selected, check selector ({criteria}) or results"
         )
     if id:
-        logger.info(f"best iteration={id}, used criteria {criteria}")
+        logger.info(f"Best iteration={id}, used criteria {criteria}")
     task = results[item] if id and results else None
     execution.log_iteration_results(id, summary, task)
 
@@ -432,11 +433,11 @@ def enrich_function_from_dict(function, function_dict):
 
 def enrich_run_labels(
     labels: dict,
-    labels_to_enrich: list[RunLabels] = None,
+    labels_to_enrich: Optional[list[RunLabels]] = None,
 ):
     labels_enrichment = {
         RunLabels.owner: os.environ.get("V3IO_USERNAME") or getpass.getuser(),
-        # TODO: remove this in 1.9.0
+        # TODO: remove this in 1.10.0
         RunLabels.v3io_user: os.environ.get("V3IO_USERNAME"),
     }
     labels_to_enrich = labels_to_enrich or RunLabels.all()

mlrun/secrets.py

@@ -134,7 +134,7 @@ class SecretsStore:
     def k8s_env_variable_name_for_secret(secret_name):
         from mlrun.config import config
 
-        return config.secret_stores.kubernetes.env_variable_prefix + secret_name.upper()
+        return config.secret_stores.kubernetes.env_variable_prefix + secret_name
 
     def get_k8s_secrets(self):
         for source in self._hidden_sources:
@@ -151,7 +151,7 @@ def get_secret_or_env(
     secret_provider: Union[dict, SecretsStore, Callable, None] = None,
     default: Optional[str] = None,
     prefix: Optional[str] = None,
-) -> str:
+) -> Optional[str]:
     """Retrieve value of a secret, either from a user-provided secret store, or from environment variables.
     The function will retrieve a secret value, attempting to find it according to the following order:
 

mlrun/serving/__init__.py

@@ -23,6 +23,10 @@ __all__ = [
     "QueueStep",
     "ErrorStep",
     "MonitoringApplicationStep",
+    "ModelRunnerStep",
+    "ModelRunner",
+    "Model",
+    "ModelSelector",
 ]
 
 from .routers import ModelRouter, VotingEnsemble  # noqa
@@ -33,6 +37,10 @@ from .states import (
     RouterStep,
     TaskStep,
     MonitoringApplicationStep,
+    ModelRunnerStep,
+    ModelRunner,
+    Model,
+    ModelSelector,
 )  # noqa
 from .v1_serving import MLModelServer, new_v1_model_server  # noqa
 from .v2_serving import V2ModelServer  # noqa

mlrun/serving/merger.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 #
 
+from typing import Optional
+
 import storey
 
 
@@ -37,10 +39,10 @@ class CacheEntry:
 class Merge(storey.Flow):
     def __init__(
         self,
-        full_event: bool = None,
-        key_path: str = None,
-        max_behind: int = None,
-        expected_num_events: int = None,
+        full_event: Optional[bool] = None,
+        key_path: Optional[str] = None,
+        max_behind: Optional[int] = None,
+        expected_num_events: Optional[int] = None,
         **kwargs,
     ):
         """Merge multiple events based on event id or provided key path
@@ -72,7 +74,7 @@ class Merge(storey.Flow):
         self._queue_len = max_behind or 64  # default queue is 64 entries
         self._keys_queue = []
 
-    def post_init(self, mode="sync"):
+    def post_init(self, mode="sync", **kwargs):
         # auto detect number of uplinks or use user specified value
         self._uplinks = self.expected_num_events or (
             len(self._graph_step.after) if self._graph_step else 0