mlrun 1.7.0rc4__py3-none-any.whl → 1.7.0rc20__py3-none-any.whl

mlrun/runtimes/pod.py

@@ -15,12 +15,14 @@ import copy
 import inspect
 import os
 import re
+import time
 import typing
 from enum import Enum
 
 import dotenv
-import kfp.dsl
 import kubernetes.client as k8s_client
+import mlrun_pipelines.mounts
+from mlrun_pipelines.mixins import KfpAdapterMixin
 
 import mlrun.errors
 import mlrun.utils.regex
@@ -40,7 +42,6 @@ from ..k8s_utils import (
 from ..utils import logger, update_in
 from .base import BaseRuntime, FunctionSpec, spec_fields
 from .utils import (
-    apply_kfp,
     get_gpu_from_resource_requirement,
     get_item_name,
     set_named_item,
@@ -934,12 +935,12 @@ class AutoMountType(str, Enum):
     @classmethod
     def all_mount_modifiers(cls):
         return [
-            mlrun.v3io_cred.__name__,
-            mlrun.mount_v3io.__name__,
-            mlrun.platforms.other.mount_pvc.__name__,
-            mlrun.auto_mount.__name__,
-            mlrun.platforms.mount_s3.__name__,
-            mlrun.platforms.set_env_variables.__name__,
+            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__,
         ]
 
     @classmethod
@@ -956,27 +957,27 @@ class AutoMountType(str, Enum):
     def _get_auto_modifier():
         # If we're running on Iguazio - use v3io_cred
         if mlconf.igz_version != "":
-            return mlrun.v3io_cred
+            return mlrun_pipelines.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.platforms.other.mount_pvc if pvc_configured else None
+        return mlrun_pipelines.mounts.mount_pvc if pvc_configured else None
 
     def get_modifier(self):
         return {
             AutoMountType.none: None,
-            AutoMountType.v3io_credentials: mlrun.v3io_cred,
-            AutoMountType.v3io_fuse: mlrun.mount_v3io,
-            AutoMountType.pvc: mlrun.platforms.other.mount_pvc,
+            AutoMountType.v3io_credentials: mlrun_pipelines.mounts.v3io_cred,
+            AutoMountType.v3io_fuse: mlrun_pipelines.mounts.mount_v3io,
+            AutoMountType.pvc: mlrun_pipelines.mounts.mount_pvc,
             AutoMountType.auto: self._get_auto_modifier(),
-            AutoMountType.s3: mlrun.platforms.mount_s3,
-            AutoMountType.env: mlrun.platforms.set_env_variables,
+            AutoMountType.s3: mlrun_pipelines.mounts.mount_s3,
+            AutoMountType.env: mlrun_pipelines.mounts.set_env_variables,
         }[self]
 
 
-class KubeResource(BaseRuntime):
+class KubeResource(BaseRuntime, KfpAdapterMixin):
     """
     A parent class for runtimes that generate k8s resources when executing.
     """
@@ -985,7 +986,7 @@ class KubeResource(BaseRuntime):
     _is_nested = True
 
     def __init__(self, spec=None, metadata=None):
-        super().__init__(metadata, spec)
+        super().__init__(metadata=metadata, spec=spec)
         self.verbose = False
 
     @property
@@ -996,26 +997,6 @@ class KubeResource(BaseRuntime):
     def spec(self, spec):
         self._spec = self._verify_dict(spec, "spec", KubeResourceSpec)
 
-    def apply(self, modify):
-        """
-        Apply a modifier to the runtime which is used to change the runtimes k8s object's spec.
-        Modifiers can be either KFP modifiers or MLRun modifiers (which are compatible with KFP). All modifiers accept
-        a `kfp.dsl.ContainerOp` object, apply some changes on its spec and return it so modifiers can be chained
-        one after the other.
-
-        :param modify: a modifier runnable object
-        :return: the runtime (self) after the modifications
-        """
-
-        # Kubeflow pipeline have a hook to add the component to the DAG on ContainerOp init
-        # we remove the hook to suppress kubeflow op registration and return it after the apply()
-        old_op_handler = kfp.dsl._container_op._register_op_handler
-        kfp.dsl._container_op._register_op_handler = lambda x: self.metadata.name
-        cop = kfp.dsl.ContainerOp("name", "image")
-        kfp.dsl._container_op._register_op_handler = old_op_handler
-
-        return apply_kfp(modify, cop, self)
-
     def set_env_from_secret(self, name, secret=None, secret_key=None):
         """set pod environment var from secret"""
         secret_key = secret_key or name
@@ -1057,14 +1038,40 @@ class KubeResource(BaseRuntime):
                 return True
         return False
 
+    def enrich_runtime_spec(
+        self,
+        project_node_selector: dict[str, str],
+    ):
+        """
+        Enriches the runtime spec with the project-level node selector.
+
+        This method merges the project-level node selector with the existing function node_selector.
+        The merge logic used here combines the two dictionaries, giving precedence to
+        the keys in the runtime node_selector. If there are conflicting keys between the
+        two dictionaries, the values from self.spec.node_selector will overwrite the
+        values from project_node_selector.
+
+        Example:
+        Suppose self.spec.node_selector = {"type": "gpu", "zone": "us-east-1"}
+        and project_node_selector = {"type": "cpu", "environment": "production"}.
+        After the merge, the resulting node_selector will be:
+        {"type": "gpu", "zone": "us-east-1", "environment": "production"}
+
+        Note:
+        - The merge uses the ** operator, also known as the "unpacking" operator in Python,
+          combining key-value pairs from each dictionary. Later dictionaries take precedence
+          when there are conflicting keys.
+        """
+        self.spec.node_selector = {**project_node_selector, **self.spec.node_selector}
+
     def _set_env(self, name, value=None, value_from=None):
         new_var = k8s_client.V1EnvVar(name=name, value=value, value_from=value_from)
-        i = 0
-        for v in self.spec.env:
-            if get_item_name(v) == name:
-                self.spec.env[i] = new_var
+
+        # ensure we don't have duplicate env vars with the same name
+        for env_index, value_item in enumerate(self.spec.env):
+            if get_item_name(value_item) == name:
+                self.spec.env[env_index] = new_var
                 return self
-            i += 1
         self.spec.env.append(new_var)
         return self
 
@@ -1251,9 +1258,9 @@ class KubeResource(BaseRuntime):
             from kubernetes import client as k8s_client
 
             security_context = k8s_client.V1SecurityContext(
-                        run_as_user=1000,
-                        run_as_group=3000,
-                    )
+                run_as_user=1000,
+                run_as_group=3000,
+            )
             function.with_security_context(security_context)
 
         More info:
@@ -1312,6 +1319,150 @@ class KubeResource(BaseRuntime):
 
         self.spec.validate_service_account(allowed_service_accounts)
 
+    def _configure_mlrun_build_with_source(
+        self, source, workdir=None, handler=None, pull_at_runtime=True, target_dir=None
+    ):
+        mlrun.utils.helpers.validate_builder_source(source, pull_at_runtime, workdir)
+
+        self.spec.build.source = source
+        if handler:
+            self.spec.default_handler = handler
+        if workdir:
+            self.spec.workdir = workdir
+        if target_dir:
+            self.spec.build.source_code_target_dir = target_dir
+
+        self.spec.build.load_source_on_run = pull_at_runtime
+        if (
+            self.spec.build.base_image
+            and not self.spec.build.commands
+            and pull_at_runtime
+            and not self.spec.image
+        ):
+            # if we load source from repo and don't need a full build use the base_image as the image
+            self.spec.image = self.spec.build.base_image
+        elif not pull_at_runtime:
+            # clear the image so build will not be skipped
+            self.spec.build.base_image = self.spec.build.base_image or self.spec.image
+            self.spec.image = ""
+
+    def _resolve_build_with_mlrun(self, with_mlrun: typing.Optional[bool] = None):
+        build = self.spec.build
+        if with_mlrun is None:
+            if build.with_mlrun is not None:
+                with_mlrun = build.with_mlrun
+            else:
+                with_mlrun = build.base_image and not (
+                    build.base_image.startswith("mlrun/")
+                    or "/mlrun/" in build.base_image
+                )
+        if (
+            not build.source
+            and not build.commands
+            and not build.requirements
+            and not build.extra
+            and with_mlrun
+        ):
+            logger.info(
+                "Running build to add mlrun package, set "
+                "with_mlrun=False to skip if its already in the image"
+            )
+        return with_mlrun
+
+    def _build_image(
+        self,
+        builder_env,
+        force_build,
+        mlrun_version_specifier,
+        show_on_failure,
+        skip_deployed,
+        watch,
+        is_kfp,
+        with_mlrun,
+    ):
+        # When we're in pipelines context we must watch otherwise the pipelines pod will exit before the operation
+        # is actually done. (when a pipelines pod exits, the pipeline step marked as done)
+        if is_kfp:
+            watch = True
+
+        db = self._get_db()
+        data = db.remote_builder(
+            self,
+            with_mlrun,
+            mlrun_version_specifier,
+            skip_deployed,
+            builder_env=builder_env,
+            force_build=force_build,
+        )
+        self.status = data["data"].get("status", None)
+        self.spec.image = mlrun.utils.get_in(
+            data, "data.spec.image"
+        ) or mlrun.utils.get_in(data, "data.spec.build.image")
+        self.spec.build.base_image = self.spec.build.base_image or mlrun.utils.get_in(
+            data, "data.spec.build.base_image"
+        )
+        # Get the source target dir in case it was enriched due to loading source
+        self.spec.build.source_code_target_dir = mlrun.utils.get_in(
+            data, "data.spec.build.source_code_target_dir"
+        ) or mlrun.utils.get_in(data, "data.spec.clone_target_dir")
+        ready = data.get("ready", False)
+        if not ready:
+            logger.info(
+                f"Started building image: {data.get('data', {}).get('spec', {}).get('build', {}).get('image')}"
+            )
+        if watch and not ready:
+            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")
+        return ready
+
+    def _build_watch(
+        self,
+        watch: bool = True,
+        logs: bool = True,
+        show_on_failure: bool = False,
+    ):
+        db = self._get_db()
+        offset = 0
+        try:
+            text, _ = db.get_builder_status(self, 0, logs=logs)
+        except mlrun.db.RunDBError:
+            raise ValueError("function or build process not found")
+
+        def print_log(text):
+            if text and (
+                not show_on_failure
+                or self.status.state == mlrun.common.schemas.FunctionState.error
+            ):
+                print(text, end="")
+
+        print_log(text)
+        offset += len(text)
+        if watch:
+            while self.status.state in [
+                mlrun.common.schemas.FunctionState.pending,
+                mlrun.common.schemas.FunctionState.running,
+            ]:
+                time.sleep(2)
+                if show_on_failure:
+                    text = ""
+                    db.get_builder_status(self, 0, logs=False)
+                    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)
+                else:
+                    text, _ = db.get_builder_status(self, offset, logs=logs)
+                print_log(text)
+                offset += len(text)
+
+        return self.status.state
+
 
 def _resolve_if_type_sanitized(attribute_name, attribute):
     attribute_config = sanitized_attributes[attribute_name]

mlrun/runtimes/remotesparkjob.py

@@ -15,11 +15,11 @@ import re
 from subprocess import run
 
 import kubernetes.client
+from mlrun_pipelines.mounts import mount_v3io, mount_v3iod
 
 import mlrun.errors
 from mlrun.config import config
 
-from ..platforms.iguazio import mount_v3io, mount_v3iod
 from .kubejob import KubejobRuntime
 from .pod import KubeResourceSpec
 

mlrun/runtimes/sparkjob/spark3job.py

@@ -14,6 +14,7 @@
 import typing
 
 import kubernetes.client
+from mlrun_pipelines.mounts import mount_v3io, mount_v3iod
 
 import mlrun.common.schemas.function
 import mlrun.errors
@@ -22,7 +23,6 @@ from mlrun.config import config
 
 from ...execution import MLClientCtx
 from ...model import RunObject
-from ...platforms.iguazio import mount_v3io, mount_v3iod
 from ...utils import update_in, verify_field_regex
 from ..kubejob import KubejobRuntime
 from ..pod import KubeResourceSpec

mlrun/runtimes/utils.py

@@ -20,17 +20,17 @@ from io import StringIO
 from sys import stderr
 
 import pandas as pd
-from kubernetes import client
 
 import mlrun
 import mlrun.common.constants
+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
-from mlrun.runtimes.constants import RunLabels
 from mlrun.runtimes.generators import selector
 from mlrun.utils import get_in, helpers, logger, verify_field_regex
 
@@ -39,9 +39,6 @@ class RunError(Exception):
     pass
 
 
-mlrun_key = "mlrun/"
-
-
 class _ContextStore:
     def __init__(self):
         self._context = None
@@ -280,43 +277,6 @@ def get_item_name(item, attr="name"):
         return getattr(item, attr, None)
 
 
-def apply_kfp(modify, cop, runtime):
-    modify(cop)
-
-    # Have to do it here to avoid circular dependencies
-    from .pod import AutoMountType
-
-    if AutoMountType.is_auto_modifier(modify):
-        runtime.spec.disable_auto_mount = True
-
-    api = client.ApiClient()
-    for k, v in cop.pod_labels.items():
-        runtime.metadata.labels[k] = v
-    for k, v in cop.pod_annotations.items():
-        runtime.metadata.annotations[k] = v
-    if cop.container.env:
-        env_names = [
-            e.name if hasattr(e, "name") else e["name"] for e in runtime.spec.env
-        ]
-        for e in api.sanitize_for_serialization(cop.container.env):
-            name = e["name"]
-            if name in env_names:
-                runtime.spec.env[env_names.index(name)] = e
-            else:
-                runtime.spec.env.append(e)
-                env_names.append(name)
-        cop.container.env.clear()
-
-    if cop.volumes and cop.container.volume_mounts:
-        vols = api.sanitize_for_serialization(cop.volumes)
-        mounts = api.sanitize_for_serialization(cop.container.volume_mounts)
-        runtime.spec.update_vols_and_mounts(vols, mounts)
-        cop.volumes.clear()
-        cop.container.volume_mounts.clear()
-
-    return runtime
-
-
 def verify_limits(
     resources_field_name,
     mem=None,
@@ -410,41 +370,13 @@ def generate_resources(mem=None, cpu=None, gpus=None, gpu_type="nvidia.com/gpu")
 
 
 def get_func_selector(project, name=None, tag=None):
-    s = [f"{mlrun_key}project={project}"]
+    s = [f"{mlrun_constants.MLRunInternalLabels.project}={project}"]
     if name:
-        s.append(f"{mlrun_key}function={name}")
-        s.append(f"{mlrun_key}tag={tag or 'latest'}")
+        s.append(f"{mlrun_constants.MLRunInternalLabels.function}={name}")
+        s.append(f"{mlrun_constants.MLRunInternalLabels.tag}={tag or 'latest'}")
     return s
 
 
-class k8s_resource:
-    kind = ""
-    per_run = False
-    per_function = False
-    k8client = None
-
-    def deploy_function(self, function):
-        pass
-
-    def release_function(self, function):
-        pass
-
-    def submit_run(self, function, runobj):
-        pass
-
-    def get_object(self, name, namespace=None):
-        return None
-
-    def get_status(self, name, namespace=None):
-        return None
-
-    def del_object(self, name, namespace=None):
-        pass
-
-    def get_pods(self, name, namespace=None, master=False):
-        return {}
-
-
 def enrich_function_from_dict(function, function_dict):
     override_function = mlrun.new_function(runtime=function_dict, kind=function.kind)
     for attribute in [
@@ -504,6 +436,7 @@ def enrich_run_labels(
 ):
     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"),
     }
     labels_to_enrich = labels_to_enrich or RunLabels.all()

mlrun/secrets.py

@@ -163,15 +163,19 @@ def get_secret_or_env(
 
     Example::
 
-        secrets = { "KEY1": "VALUE1" }
+        secrets = {"KEY1": "VALUE1"}
         secret = get_secret_or_env("KEY1", secret_provider=secrets)
 
+
         # Using a function to retrieve a secret
         def my_secret_provider(key):
             # some internal logic to retrieve secret
             return value
 
-        secret = get_secret_or_env("KEY1", secret_provider=my_secret_provider, default="TOO-MANY-SECRETS")
+
+        secret = get_secret_or_env(
+            "KEY1", secret_provider=my_secret_provider, default="TOO-MANY-SECRETS"
+        )
 
     :param key: Secret key to look for
     :param secret_provider: Dictionary, callable or `SecretsStore` to extract the secret value from. If using a

mlrun/serving/remote.py

@@ -172,8 +172,7 @@ class RemoteStep(storey.SendToHttp):
         if not self._session:
             self._session = mlrun.utils.HTTPSessionWithRetry(
                 self.retries,
-                self.backoff_factor
-                or mlrun.config.config.http_retry_defaults.backoff_factor,
+                self.backoff_factor or mlrun.mlconf.http_retry_defaults.backoff_factor,
                 retry_on_exception=False,
                 retry_on_status=self.retries > 0,
                 retry_on_post=True,
@@ -185,7 +184,7 @@ class RemoteStep(storey.SendToHttp):
             resp = self._session.request(
                 method,
                 url,
-                verify=mlrun.config.config.httpdb.http.verify,
+                verify=mlrun.mlconf.httpdb.http.verify,
                 headers=headers,
                 data=body,
                 timeout=self.timeout,

mlrun/serving/routers.py

@@ -28,6 +28,7 @@ import numpy as np
 import mlrun
 import mlrun.common.model_monitoring
 import mlrun.common.schemas.model_monitoring
+from mlrun.errors import err_to_str
 from mlrun.utils import logger, now_date
 
 from ..common.helpers import parse_versioned_object_uri
@@ -271,7 +272,9 @@ class ParallelRun(BaseModelRouter):
             fn = mlrun.new_function("parallel", kind="serving")
             graph = fn.set_topology(
                 "router",
-                mlrun.serving.routers.ParallelRun(extend_event=True, executor_type=executor),
+                mlrun.serving.routers.ParallelRun(
+                    extend_event=True, executor_type=executor
+                ),
             )
             graph.add_route("child1", class_name="Cls1")
             graph.add_route("child2", class_name="Cls2", my_arg={"c": 7})
@@ -1013,7 +1016,7 @@ def _init_endpoint_record(
             graph_server.function_uri
         )
     except Exception as e:
-        logger.error("Failed to parse function URI", exc=e)
+        logger.error("Failed to parse function URI", exc=err_to_str(e))
         return None
 
     # Generating version model value based on the model name and model version
@@ -1089,12 +1092,12 @@ def _init_endpoint_record(
         except Exception as exc:
             logger.warning(
                 "Failed creating model endpoint record",
-                exc=exc,
+                exc=err_to_str(exc),
                 traceback=traceback.format_exc(),
             )
 
     except Exception as e:
-        logger.error("Failed to retrieve model endpoint object", exc=e)
+        logger.error("Failed to retrieve model endpoint object", exc=err_to_str(e))
 
     return endpoint_uid
 

mlrun/serving/server.py

@@ -23,6 +23,7 @@ import uuid
 from typing import Optional, Union
 
 import mlrun
+import mlrun.common.constants
 import mlrun.common.helpers
 import mlrun.model_monitoring
 from mlrun.config import config
@@ -52,7 +53,7 @@ class _StreamContext:
         Initialize _StreamContext object.
         :param enabled:      A boolean indication for applying the stream context
         :param parameters:   Dictionary of optional parameters, such as `log_stream` and `stream_args`. Note that these
-                             parameters might be relevant to the output source such as `kafka_bootstrap_servers` if
+                             parameters might be relevant to the output source such as `kafka_brokers` if
                              the output source is from type Kafka.
         :param function_uri: Full value of the function uri, usually it's <project-name>/<function-name>
         """
@@ -311,11 +312,8 @@ class GraphServer(ModelObj):
 def v2_serving_init(context, namespace=None):
     """hook for nuclio init_context()"""
 
-    data = os.environ.get("SERVING_SPEC_ENV", "")
-    if not data:
-        raise MLRunInvalidArgumentError("failed to find spec env var")
-    spec = json.loads(data)
     context.logger.info("Initializing server from spec")
+    spec = mlrun.utils.get_serving_spec()
     server = GraphServer.from_dict(spec)
     if config.log_level.lower() == "debug":
         server.verbose = True
@@ -355,7 +353,7 @@ def v2_serving_init(context, namespace=None):
 
         async def termination_callback():
             context.logger.info("Termination callback called")
-            await server.wait_for_completion()
+            server.wait_for_completion()
             context.logger.info("Termination of async flow is completed")
 
         context.platform.set_termination_callback(termination_callback)
@@ -367,7 +365,7 @@ def v2_serving_init(context, namespace=None):
 
         async def drain_callback():
             context.logger.info("Drain callback called")
-            await server.wait_for_completion()
+            server.wait_for_completion()
             context.logger.info(
                 "Termination of async flow is completed. Rerunning async flow."
             )
@@ -389,7 +387,7 @@ def v2_serving_handler(context, event, get_body=False):
 
 
 def create_graph_server(
-    parameters={},
+    parameters=None,
     load_mode=None,
     graph=None,
     verbose=False,
@@ -405,6 +403,7 @@ def create_graph_server(
         server.graph.add_route("my", class_name=MyModelClass, model_path="{path}", z=100)
         print(server.test("/v2/models/my/infer", testdata))
     """
+    parameters = parameters or {}
     server = GraphServer(graph, parameters, load_mode, verbose=verbose, **kwargs)
     server.set_current_function(
         current_function or os.environ.get("SERVING_CURRENT_FUNCTION", "")