mlrun 1.10.0rc40__py3-none-any.whl → 1.11.0rc16__py3-none-any.whl

mlrun/runtimes/nuclio/function.py

@@ -17,8 +17,10 @@ import copy
 import json
 import typing
 import warnings
+from dataclasses import dataclass
 from datetime import datetime
 from time import sleep
+from urllib.parse import urlparse, urlunparse
 
 import inflection
 import nuclio
@@ -30,13 +32,14 @@ from kubernetes import client
 from nuclio.deploy import find_dashboard_url, get_deploy_status
 from nuclio.triggers import V3IOStreamTrigger
 
+import mlrun.auth.nuclio
 import mlrun.common.constants
 import mlrun.db
 import mlrun.errors
 import mlrun.k8s_utils
 import mlrun.utils
 import mlrun.utils.helpers
-from mlrun.common.schemas import AuthInfo
+from mlrun.common.schemas import AuthInfo, BatchingSpec
 from mlrun.config import config as mlconf
 from mlrun.errors import err_to_str
 from mlrun.lists import RunList
@@ -96,6 +99,13 @@ def min_nuclio_versions(*versions):
     return decorator
 
 
+@dataclass
+class AsyncSpec:
+    enabled: bool = True
+    max_connections: typing.Optional[int] = None
+    connection_availability_timeout: typing.Optional[int] = None
+
+
 class NuclioSpec(KubeResourceSpec):
     _dict_fields = KubeResourceSpec._dict_fields + [
         "min_replicas",
@@ -113,6 +123,7 @@ class NuclioSpec(KubeResourceSpec):
         "service_type",
         "add_templated_ingress_host_mode",
         "disable_default_http_trigger",
+        "auth",
     ]
 
     def __init__(
@@ -160,6 +171,7 @@ class NuclioSpec(KubeResourceSpec):
         graph=None,
         parameters=None,
         track_models=None,
+        auth=None,
     ):
         super().__init__(
             command=command,
@@ -216,6 +228,7 @@ class NuclioSpec(KubeResourceSpec):
         # When True it will set Nuclio spec.noBaseImagesPull to False (negative logic)
         # indicate that the base image should be pulled from the container registry (not cached)
         self.base_image_pull = False
+        self.auth = auth or {}
 
     def generate_nuclio_volumes(self):
         nuclio_volumes = []
@@ -300,29 +313,16 @@ class RemoteRuntime(KubeResource):
             return {}
 
         raw_config = copy.deepcopy(self.spec.config)
-
         for key, value in self.spec.config.items():
             if key.startswith("spec.triggers"):
-                trigger_name = key.split(".")[-1]
-
-                for path in SENSITIVE_PATHS_IN_TRIGGER_CONFIG:
-                    # Handle nested keys
-                    nested_keys = path.split("/")
-                    target = value
-                    for sub_key in nested_keys[:-1]:
-                        target = target.get(sub_key, {})
-
-                    last_key = nested_keys[-1]
-                    if last_key in target:
-                        sensitive_field = target[last_key]
-                        if sensitive_field.startswith(
-                            mlrun.model.Credentials.secret_reference_prefix
-                        ):
-                            # already masked
-                            continue
-                        target[last_key] = (
-                            f"{mlrun.model.Credentials.secret_reference_prefix}/spec/triggers/{trigger_name}/{path}"
-                        )
+                # support both types depending on the way how it was set
+                # sometimes trigger name is in the same key, sometimes it's nested in the value dict
+                if key == "spec.triggers":
+                    for trigger_name, trigger_config in value.items():
+                        self._mask_trigger_config(trigger_name, trigger_config)
+                else:
+                    trigger_name = key.split(".")[-1]
+                    self._mask_trigger_config(trigger_name, value)
 
         return raw_config
 
@@ -464,7 +464,7 @@ class RemoteRuntime(KubeResource):
 
     def with_http(
         self,
-        workers: typing.Optional[int] = 8,
+        workers: typing.Optional[int] = None,
         port: typing.Optional[int] = None,
         host: typing.Optional[str] = None,
         paths: typing.Optional[list[str]] = None,
@@ -475,6 +475,8 @@ class RemoteRuntime(KubeResource):
         trigger_name: typing.Optional[str] = None,
         annotations: typing.Optional[typing.Mapping[str, str]] = None,
         extra_attributes: typing.Optional[typing.Mapping[str, str]] = None,
+        batching_spec: typing.Optional[BatchingSpec] = None,
+        async_spec: typing.Optional[AsyncSpec] = None,
     ):
         """update/add nuclio HTTP trigger settings
 
@@ -482,7 +484,8 @@ class RemoteRuntime(KubeResource):
         if the max time a request will wait for until it will start processing, gateway_timeout must be greater than
         the worker_timeout.
 
-        :param workers:    number of worker processes (default=8). set 0 to use Nuclio's default workers count
+        :param workers: Number of worker processes. Defaults to 8 in synchronous mode and
+                        1 in asynchronous mode. Set to 0 to use Nuclio’s default worker count.
         :param port:       TCP port to listen on. by default, nuclio will choose a random port as long as
                            the function service is NodePort. if the function service is ClusterIP, the port
                            is ignored.
@@ -496,6 +499,12 @@ class RemoteRuntime(KubeResource):
         :param trigger_name:    alternative nuclio trigger name
         :param annotations:     key/value dict of ingress annotations
         :param extra_attributes: key/value dict of extra nuclio trigger attributes
+        :param batching_spec: BatchingSpec object that defines batching configuration.
+            By default, batching is disabled.
+
+        :param async_spec: AsyncSpec object defines async configuration. If number of max connections
+            won't be set, the default value will be set to 1000 according to nuclio default.
+
         :return: function object (self)
         """
         if self.disable_default_http_trigger:
@@ -503,11 +512,15 @@ class RemoteRuntime(KubeResource):
                 "Adding HTTP trigger despite the default HTTP trigger creation being disabled"
             )
 
+        if async_spec and async_spec.enabled:
+            workers = 1 if workers is None else workers
+        else:
+            workers = 8 if workers is None else workers
+
         annotations = annotations or {}
         if worker_timeout:
             gateway_timeout = gateway_timeout or (worker_timeout + 60)
-        if workers is None:
-            workers = 0
+
         if gateway_timeout:
             if worker_timeout and worker_timeout >= gateway_timeout:
                 raise ValueError(
@@ -531,6 +544,28 @@ class RemoteRuntime(KubeResource):
             trigger._struct["workerAvailabilityTimeoutMilliseconds"] = (
                 worker_timeout
             ) * 1000
+
+        if batching_spec and (
+            batching_config := batching_spec.get_nuclio_batch_config()
+        ):
+            if not validate_nuclio_version_compatibility("1.14.0"):
+                raise mlrun.errors.MLRunValueError(
+                    "Batching is only supported on Nuclio 1.14.0 and higher"
+                )
+            trigger._struct["batch"] = batching_config
+
+        if async_spec:
+            if not validate_nuclio_version_compatibility("1.15.3"):
+                raise mlrun.errors.MLRunValueError(
+                    "Async spec is only supported on Nuclio 1.15.3 and higher"
+                )
+            if async_spec.enabled:
+                trigger._struct["mode"] = "async"
+                trigger._struct["async"] = {
+                    "maxConnectionsNumber": async_spec.max_connections,
+                    "connectionAvailabilityTimeout": async_spec.connection_availability_timeout,
+                }
+
         self.add_trigger(trigger_name or "http", trigger)
         return self
 
@@ -655,8 +690,6 @@ class RemoteRuntime(KubeResource):
         if tag:
             self.metadata.tag = tag
 
-        mlrun.utils.helpers.validate_function_name(self.metadata.name)
-
         # Attempt auto-mounting, before sending to remote build
         self.try_auto_mount_based_on_config()
         self._fill_credentials()
@@ -843,22 +876,6 @@ class RemoteRuntime(KubeResource):
             raise ValueError("function or deploy process not found")
         return self.status.state, text, last_log_timestamp
 
-    def _get_runtime_env(self):
-        # for runtime specific env var enrichment (before deploy)
-        runtime_env = {
-            mlrun.common.constants.MLRUN_ACTIVE_PROJECT: self.metadata.project
-            or mlconf.active_project,
-        }
-        if mlconf.httpdb.api_url:
-            runtime_env["MLRUN_DBPATH"] = mlconf.httpdb.api_url
-        if mlconf.namespace:
-            runtime_env["MLRUN_NAMESPACE"] = mlconf.namespace
-        if self.metadata.credentials.access_key:
-            runtime_env[
-                mlrun.common.runtimes.constants.FunctionEnvironmentVariables.auth_session
-            ] = self.metadata.credentials.access_key
-        return runtime_env
-
     def _get_serving_spec(self):
         return None
 
@@ -883,8 +900,9 @@ class RemoteRuntime(KubeResource):
             if value_from is not None:
                 external_source_env_dict[sanitized_env_var.get("name")] = value_from
 
-        for key, value in self._get_runtime_env().items():
-            env_dict[key] = value
+        envs, external_source_envs = self._generate_runtime_env()
+        env_dict.update(envs)
+        external_source_env_dict.update(external_source_envs)
 
         return env_dict, external_source_env_dict
 
@@ -941,7 +959,7 @@ class RemoteRuntime(KubeResource):
     def invoke(
         self,
         path: str,
-        body: typing.Optional[typing.Union[str, bytes, dict]] = None,
+        body: typing.Optional[typing.Union[str, bytes, dict, list]] = None,
         method: typing.Optional[str] = None,
         headers: typing.Optional[dict] = None,
         force_external_address: bool = False,
@@ -996,7 +1014,7 @@ class RemoteRuntime(KubeResource):
         if not http_client_kwargs:
             http_client_kwargs = {}
         if body:
-            if isinstance(body, (str, bytes)):
+            if isinstance(body, str | bytes):
                 http_client_kwargs["data"] = body
             else:
                 http_client_kwargs["json"] = body
@@ -1070,6 +1088,20 @@ class RemoteRuntime(KubeResource):
             sidecar["resources"] = self.spec.resources
             self.spec.resources = None
 
+    def set_probe(self, *args, **kwargs):
+        """Set a Kubernetes probe configuration for the sidecar container
+
+        This method is only available for ApplicationRuntime.
+        """
+        raise ValueError("set_probe() is only supported for ApplicationRuntime. ")
+
+    def delete_probe(self, *args, **kwargs):
+        """Delete a Kubernetes probe configuration from the sidecar container
+
+        This method is only available for ApplicationRuntime.
+        """
+        raise ValueError("delete_probe() is only supported for ApplicationRuntime.")
+
     def _set_sidecar(self, name: str) -> dict:
         self.spec.config.setdefault("spec.sidecars", [])
         sidecars = self.spec.config["spec.sidecars"]
@@ -1080,6 +1112,79 @@ class RemoteRuntime(KubeResource):
         sidecars.append({"name": name})
         return sidecars[-1]
 
+    def _mask_trigger_config(self, trigger_name, trigger_config):
+        self._mask_rabbitmq_url(trigger=trigger_config)
+        for path in SENSITIVE_PATHS_IN_TRIGGER_CONFIG:
+            # Handle nested keys
+            nested_keys = path.split("/")
+            target = trigger_config
+            for sub_key in nested_keys[:-1]:
+                target = target.get(sub_key, {})
+
+            last_key = nested_keys[-1]
+            if last_key in target:
+                sensitive_field = target[last_key]
+                if sensitive_field.startswith(
+                    mlrun.model.Credentials.secret_reference_prefix
+                ):
+                    # already masked
+                    continue
+                target[last_key] = (
+                    f"{mlrun.model.Credentials.secret_reference_prefix}/spec/triggers/{trigger_name}/{path}"
+                )
+
+    @staticmethod
+    def _mask_rabbitmq_url(trigger):
+        """
+        Extract credentials from RabbitMQ URL and move them to attributes dict.
+        This ensures credentials are not exposed in the URL.
+        """
+
+        # supported only for nuclio higher than 1.14.15
+        if not validate_nuclio_version_compatibility("1.14.15"):
+            return
+        if not isinstance(trigger, dict):
+            return
+
+        if trigger.get("kind") != "rabbit-mq":
+            return
+
+        url = trigger.get("url")
+        if not url or not isinstance(url, str):
+            return
+
+        try:
+            parsed = urlparse(url)
+        except Exception:
+            raise mlrun.errors.MLRunValueError("invalid URL format")
+
+        # Only process if credentials are present in the URL
+        if not (parsed.username or parsed.password):
+            return
+
+        # Extract credentials
+        username = parsed.username or ""
+        password = parsed.password or ""
+
+        # Reconstruct clean URL
+        hostname = parsed.hostname or ""
+        netloc = f"{hostname}:{parsed.port}" if parsed.port else hostname
+
+        clean_url = urlunparse(
+            (
+                parsed.scheme,
+                netloc,
+                parsed.path,
+                parsed.params,
+                parsed.query,
+                parsed.fragment,
+            )
+        )
+
+        # Update trigger safely
+        trigger["url"] = clean_url
+        trigger.update({"username": username, "password": password})
+
     def _trigger_of_kind_exists(self, kind: str) -> bool:
         if not self.spec.config:
             return False
@@ -1456,7 +1561,7 @@ def get_nuclio_deploy_status(
             verbose,
             resolve_address,
             return_function_status=True,
-            auth_info=auth_info.to_nuclio_auth_info() if auth_info else None,
+            auth_info=mlrun.auth.nuclio.NuclioAuthInfo.from_auth_info(auth_info),
         )
     except requests.exceptions.ConnectionError as exc:
         mlrun.errors.raise_for_status(

mlrun/runtimes/nuclio/serving.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 import json
 import os
-import warnings
 from base64 import b64decode
 from copy import deepcopy
 from typing import Optional, Union
@@ -23,7 +22,11 @@ from nuclio import KafkaTrigger
 
 import mlrun
 import mlrun.common.schemas as schemas
+import mlrun.common.secrets
 import mlrun.datastore.datastore_profile as ds_profile
+import mlrun.runtimes.kubejob as kubejob_runtime
+import mlrun.runtimes.nuclio.function as nuclio_function
+import mlrun.runtimes.pod as pod_runtime
 from mlrun.datastore import get_kafka_brokers_from_dict, parse_kafka_url
 from mlrun.model import ObjectList
 from mlrun.runtimes.function_reference import FunctionReference
@@ -44,10 +47,6 @@ from mlrun.serving.states import (
 )
 from mlrun.utils import get_caller_globals, logger, set_paths
 
-from .. import KubejobRuntime
-from ..pod import KubeResourceSpec
-from .function import NuclioSpec, RemoteRuntime, min_nuclio_versions
-
 serving_subkind = "serving_v2"
 
 
@@ -85,8 +84,8 @@ def new_v2_model_server(
     return f
 
 
-class ServingSpec(NuclioSpec):
-    _dict_fields = NuclioSpec._dict_fields + [
+class ServingSpec(nuclio_function.NuclioSpec):
+    _dict_fields = nuclio_function.NuclioSpec._dict_fields + [
         "graph",
         "load_mode",
         "graph_initializer",
@@ -154,6 +153,7 @@ class ServingSpec(NuclioSpec):
         disable_default_http_trigger=None,
         model_endpoint_creation_task_name=None,
         serving_spec=None,
+        auth=None,
     ):
         super().__init__(
             command=command,
@@ -195,6 +195,7 @@ class ServingSpec(NuclioSpec):
             add_templated_ingress_host_mode=add_templated_ingress_host_mode,
             disable_default_http_trigger=disable_default_http_trigger,
             serving_spec=serving_spec,
+            auth=auth,
         )
 
         self.models = models or {}
@@ -231,7 +232,7 @@ class ServingSpec(NuclioSpec):
         self._function_refs = ObjectList.from_list(FunctionReference, function_refs)
 
 
-class ServingRuntime(RemoteRuntime):
+class ServingRuntime(nuclio_function.RemoteRuntime):
     """MLRun Serving Runtime"""
 
     kind = "serving"
@@ -250,6 +251,8 @@ class ServingRuntime(RemoteRuntime):
         class_name=None,
         engine=None,
         exist_ok=False,
+        allow_cyclic: bool = False,
+        max_iterations: Optional[int] = None,
         **class_args,
     ) -> Union[RootFlowStep, RouterStep]:
         """set the serving graph topology (router/flow) and root class or params
@@ -280,14 +283,23 @@ class ServingRuntime(RemoteRuntime):
         :param class_name:   - optional for router, router class name/path or router object
         :param engine:       - optional for flow, sync or async engine
         :param exist_ok:     - allow overriding existing topology
+        :param allow_cyclic: - allow cyclic graphs (only for async flow)
+        :param max_iterations: - optional, max iterations for cyclic graphs (only for async flow)
         :param class_args:   - optional, router/flow class init args
 
-        :return graph object (fn.spec.graph)
+        :return: graph object (fn.spec.graph)
         """
         topology = topology or StepKinds.router
         if self.spec.graph and not exist_ok:
             raise mlrun.errors.MLRunInvalidArgumentError(
-                "graph topology is already set, cannot be overwritten"
+                "graph topology is already set, graph was initialized, use exist_ok=True to override"
+            )
+        if allow_cyclic and (
+            topology == StepKinds.router
+            or (topology == StepKinds.flow and engine == "sync")
+        ):
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "cyclic graphs are only supported in flow topology with async engine"
             )
 
         if topology == StepKinds.router:
@@ -301,7 +313,11 @@ class ServingRuntime(RemoteRuntime):
                 step = RouterStep(class_name=class_name, class_args=class_args)
             self.spec.graph = step
         elif topology == StepKinds.flow:
-            self.spec.graph = RootFlowStep(engine=engine or "async")
+            self.spec.graph = RootFlowStep(
+                engine=engine or "async",
+                allow_cyclic=allow_cyclic,
+                max_iterations=max_iterations,
+            )
             self.spec.graph.track_models = self.spec.track_models
         else:
             raise mlrun.errors.MLRunInvalidArgumentError(
@@ -312,7 +328,6 @@ class ServingRuntime(RemoteRuntime):
     def set_tracking(
         self,
         stream_path: Optional[str] = None,
-        batch: Optional[int] = None,
         sampling_percentage: float = 100,
         stream_args: Optional[dict] = None,
         enable_tracking: bool = True,
@@ -322,7 +337,6 @@ class ServingRuntime(RemoteRuntime):
 
         :param stream_path:                Path/url of the tracking stream e.g. v3io:///users/mike/mystream
                                            you can use the "dummy://" path for test/simulation.
-        :param batch:                      Deprecated. Micro batch size (send micro batches of N records at a time).
         :param sampling_percentage:        Down sampling events that will be pushed to the monitoring stream based on
                                            a specified percentage. e.g. 50 for 50%. By default, all events are pushed.
         :param stream_args:                Stream initialization parameters, e.g. shards, retention_in_hours, ..
@@ -370,13 +384,6 @@ class ServingRuntime(RemoteRuntime):
 
         if stream_path:
             self.spec.parameters["log_stream"] = stream_path
-        if batch:
-            warnings.warn(
-                "The `batch` size parameter was deprecated in version 1.8.0 and is no longer used. "
-                "It will be removed in 1.11.",
-                # TODO: Remove this in 1.11
-                FutureWarning,
-            )
         if stream_args:
             self.spec.parameters["stream_args"] = stream_args
 
@@ -395,7 +402,7 @@ class ServingRuntime(RemoteRuntime):
         outputs: Optional[list[str]] = None,
         **class_args,
     ):
-        """add ml model and/or route to the function.
+        """Add ml model and/or route to the function.
 
         Example, create a function (from the notebook), add a model class, and deploy::
 
@@ -403,7 +410,7 @@ class ServingRuntime(RemoteRuntime):
             fn.add_model("boost", model_path, model_class="MyClass", my_arg=5)
             fn.deploy()
 
-        only works with router topology, for nested topologies (model under router under flow)
+        Only works with router topology. For nested topologies (model under router under flow)
         need to add router to flow and use router.add_route()
 
         :param key:         model api key (or name:version), will determine the relative url/path
@@ -416,18 +423,19 @@ class ServingRuntime(RemoteRuntime):
                             with multiple router steps)
         :param child_function: child function name, when the model runs in a child function
         :param creation_strategy: Strategy for creating or updating the model endpoint:
-            * **overwrite**:
-            1. If model endpoints with the same name exist, delete the `latest` one.
-            2. Create a new model endpoint entry and set it as `latest`.
-            * **inplace** (default):
-            1. If model endpoints with the same name exist, update the `latest` entry.
-            2. Otherwise, create a new entry.
-            * **archive**:
-            1. If model endpoints with the same name exist, preserve them.
-            2. Create a new model endpoint with the same name and set it to `latest`.
-        :param outputs: list of the model outputs (e.g. labels) ,if provided will override the outputs that been
-                        configured in the model artifact, please note that those outputs need to be equal to the
-                        model serving function outputs (length, and order)
+
+                          * **overwrite**: If model endpoints with the same name exist, delete the `latest`
+                            one. Create a new model endpoint entry and set it as `latest`.
+
+                          * **inplace** (default): If model endpoints with the same name exist, update the
+                            `latest` entry. Otherwise, create a new entry.
+
+                          * **archive**: If model endpoints with the same name exist, preserve them.
+                            Create a new model endpoint with the same name and set it to `latest`.
+
+        :param outputs: list of the model outputs (e.g. labels), if provided will override the outputs that were
+                        configured in the model artifact. Note that those outputs need to be equal to the
+                        model serving function outputs (length, and order).
         :param class_args:  extra kwargs to pass to the model serving class __init__
                             (can be read in the model using .get_param(key) method)
         """
@@ -520,7 +528,7 @@ class ServingRuntime(RemoteRuntime):
         :param requirements: py package requirements file path OR list of packages
         :param kind:   mlrun function/runtime kind
 
-        :return function object
+        :return: function object
         """
         function_reference = FunctionReference(
             url,
@@ -635,14 +643,19 @@ class ServingRuntime(RemoteRuntime):
 
         :returns: The Runtime (function) object
         """
-
+        if kind == "azure_vault" and isinstance(source, dict):
+            candidate_secret_name = (source.get("k8s_secret") or "").strip()
+            if candidate_secret_name:
+                mlrun.common.secrets.validate_not_forbidden_secret(
+                    candidate_secret_name
+                )
         if kind == "vault" and isinstance(source, list):
             source = {"project": self.metadata.project, "secrets": source}
 
         self.spec.secret_sources.append({"kind": kind, "source": source})
         return self
 
-    @min_nuclio_versions("1.12.10")
+    @nuclio_function.min_nuclio_versions("1.12.10")
     def deploy(
         self,
         project="",
@@ -659,8 +672,6 @@ class ServingRuntime(RemoteRuntime):
         :param builder_env: env vars dict for source archive config/credentials e.g. builder_env={"GIT_TOKEN": token}
         :param force_build: set True for force building the image
         """
-        # Validate function name before deploying to k8s
-        mlrun.utils.helpers.validate_function_name(self.metadata.name)
 
         load_mode = self.spec.load_mode
         if load_mode and load_mode not in ["sync", "async"]:
@@ -858,7 +869,9 @@ class ServingRuntime(RemoteRuntime):
         )
         self._mock_server = self.to_mock_server()
 
-    def to_job(self, func_name: Optional[str] = None) -> KubejobRuntime:
+    def to_job(
+        self, func_name: Optional[str] = None
+    ) -> "kubejob_runtime.KubejobRuntime":
         """Convert this ServingRuntime to a KubejobRuntime, so that the graph can be run as a standalone job.
 
         Args:
@@ -877,7 +890,7 @@ class ServingRuntime(RemoteRuntime):
                 f"Cannot convert function '{self.metadata.name}' to a job because it has child functions"
             )
 
-        spec = KubeResourceSpec(
+        spec = pod_runtime.KubeResourceSpec(
             image=self.spec.image,
             mode=self.spec.mode,
             volumes=self.spec.volumes,
@@ -947,7 +960,7 @@ class ServingRuntime(RemoteRuntime):
                     suffix=suffix,
                 )
 
-        job = KubejobRuntime(
+        job = kubejob_runtime.KubejobRuntime(
             spec=spec,
             metadata=job_metadata,
         )