mlrun 1.7.0rc6__py3-none-any.whl → 1.7.0rc7__py3-none-any.whl

mlrun/runtimes/nuclio/api_gateway.py

@@ -12,15 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import base64
+import typing
 from typing import Optional, Union
 from urllib.parse import urljoin
 
 import requests
+from requests.auth import HTTPBasicAuth
 
 import mlrun
 import mlrun.common.schemas
 
-from .function import RemoteRuntime
+from .function import RemoteRuntime, get_fullname
 from .serving import ServingRuntime
 
 NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_BASIC_AUTH = "basicAuth"
@@ -28,6 +30,67 @@ NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_NONE = "none"
 PROJECT_NAME_LABEL = "nuclio.io/project-name"
 
 
+class APIGatewayAuthenticator(typing.Protocol):
+    @property
+    def authentication_mode(self) -> str:
+        return NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_NONE
+
+    @classmethod
+    def from_scheme(cls, api_gateway_spec: mlrun.common.schemas.APIGatewaySpec):
+        if (
+            api_gateway_spec.authenticationMode
+            == NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_BASIC_AUTH
+        ):
+            if api_gateway_spec.authentication:
+                return BasicAuth(
+                    username=api_gateway_spec.authentication.get("username", ""),
+                    password=api_gateway_spec.authentication.get("password", ""),
+                )
+            else:
+                return BasicAuth()
+        else:
+            return NoneAuth()
+
+    def to_scheme(
+        self,
+    ) -> Optional[dict[str, Optional[mlrun.common.schemas.APIGatewayBasicAuth]]]:
+        return None
+
+
+class NoneAuth(APIGatewayAuthenticator):
+    """
+    An API gateway authenticator with no authentication.
+    """
+
+    pass
+
+
+class BasicAuth(APIGatewayAuthenticator):
+    """
+    An API gateway authenticator with basic authentication.
+
+    :param username: (str) The username for basic authentication.
+    :param password: (str) The password for basic authentication.
+    """
+
+    def __init__(self, username=None, password=None):
+        self._username = username
+        self._password = password
+
+    @property
+    def authentication_mode(self) -> str:
+        return NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_BASIC_AUTH
+
+    def to_scheme(
+        self,
+    ) -> Optional[dict[str, Optional[mlrun.common.schemas.APIGatewayBasicAuth]]]:
+        return {
+            "authentication": mlrun.common.schemas.APIGatewayBasicAuth(
+                username=self._username, password=self._password
+            )
+        }
+
+
 class APIGateway:
     def __init__(
         self,
@@ -47,22 +110,34 @@ class APIGateway:
         ],
         description: str = "",
         path: str = "/",
-        authentication_mode: Optional[
-            str
-        ] = NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_NONE,
+        authentication: Optional[APIGatewayAuthenticator] = NoneAuth(),
         host: Optional[str] = None,
         canary: Optional[list[int]] = None,
-        username: Optional[str] = None,
-        password: Optional[str] = None,
     ):
+        """
+        Initialize the APIGateway instance.
+
+        :param project: The project name
+        :param name: The name of the API gateway
+        :param functions: The list of functions associated with the API gateway
+            Can be a list of function names (["my-func1", "my-func2"])
+            or a list or a single entity of
+            :py:class:`~mlrun.runtimes.nuclio.function.RemoteRuntime` OR
+            :py:class:`~mlrun.runtimes.nuclio.serving.ServingRuntime`
+
+        :param description: Optional description of the API gateway
+        :param path: Optional path of the API gateway, default value is "/"
+        :param authentication: The authentication for the API gateway of type
+                :py:class:`~mlrun.runtimes.nuclio.api_gateway.BasicAuth`
+        :param host:  The host of the API gateway (optional). If not set, it will be automatically generated
+        :param canary: The canary percents for the API gateway of type list[int]; for instance: [20,80]
+        """
         self.functions = None
         self._validate(
             project=project,
             functions=functions,
             name=name,
             canary=canary,
-            username=username,
-            password=password,
         )
         self.project = project
         self.name = name
@@ -70,14 +145,8 @@ class APIGateway:
 
         self.path = path
         self.description = description
-        self.authentication_mode = (
-            authentication_mode
-            if authentication_mode
-            else self._enrich_authentication_mode(username=username, password=password)
-        )
         self.canary = canary
-        self._username = username
-        self._password = password
+        self.authentication = authentication
 
     def invoke(
         self,
@@ -86,24 +155,94 @@ class APIGateway:
         auth: Optional[tuple[str, str]] = None,
         **kwargs,
     ):
+        """
+        Invoke the API gateway.
+
+        :param method: (str, optional) The HTTP method for the invocation.
+        :param headers: (dict, optional) The HTTP headers for the invocation.
+        :param auth: (Optional[tuple[str, str]], optional) The authentication creds for the invocation if required.
+        :param kwargs: (dict) Additional keyword arguments.
+
+        :return: The response from the API gateway invocation.
+        """
         if not self.invoke_url:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                "Invocation url is not set. Set up gateway's `invoke_url` attribute."
-            )
+            # try to resolve invoke_url before fail
+            self.sync()
+            if not self.invoke_url:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "Invocation url is not set. Set up gateway's `invoke_url` attribute."
+                )
         if (
-            self.authentication_mode
+            self.authentication.authentication_mode
             == NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_BASIC_AUTH
             and not auth
         ):
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "API Gateway invocation requires authentication. Please pass credentials"
             )
-        if auth:
-            headers["Authorization"] = self._generate_basic_auth(*auth)
         return requests.request(
-            method=method, url=self.invoke_url, headers=headers, **kwargs
+            method=method,
+            url=self.invoke_url,
+            headers=headers,
+            **kwargs,
+            auth=HTTPBasicAuth(*auth) if auth else None,
         )
 
+    def sync(self):
+        """
+        Synchronize the API gateway from the server.
+        """
+        synced_gateway = mlrun.get_run_db().get_api_gateway(self.name, self.project)
+        synced_gateway = self.from_scheme(synced_gateway)
+
+        self.host = synced_gateway.host
+        self.path = synced_gateway.path
+        self.authentication = synced_gateway.authentication
+        self.functions = synced_gateway.functions
+        self.canary = synced_gateway.canary
+        self.description = synced_gateway.description
+
+    def with_basic_auth(self, username: str, password: str):
+        """
+        Set basic authentication for the API gateway.
+
+        :param username: (str) The username for basic authentication.
+        :param password: (str) The password for basic authentication.
+        """
+        self.authentication = BasicAuth(username=username, password=password)
+
+    def with_canary(
+        self,
+        functions: Union[
+            list[str],
+            list[
+                Union[
+                    RemoteRuntime,
+                    ServingRuntime,
+                ]
+            ],
+        ],
+        canary: list[int],
+    ):
+        """
+        Set canary function for the API gateway
+
+        :param functions: The list of functions associated with the API gateway
+            Can be a list of function names (["my-func1", "my-func2"])
+            or a list of nuclio functions of types
+            :py:class:`~mlrun.runtimes.nuclio.function.RemoteRuntime` OR
+            :py:class:`~mlrun.runtimes.nuclio.serving.ServingRuntime`
+        :param canary: The canary percents for the API gateway of type list[int]; for instance: [20,80]
+
+        """
+        if len(functions) != 2:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"Gateway with canary can be created only with two functions, "
+                f"the number of functions passed is {len(functions)}"
+            )
+        self.functions = self._validate_functions(self.project, functions)
+        self.canary = self._validate_canary(canary)
+
     @classmethod
     def from_scheme(cls, api_gateway: mlrun.common.schemas.APIGateway):
         project = api_gateway.metadata.labels.get(PROJECT_NAME_LABEL)
@@ -114,7 +253,7 @@ class APIGateway:
             name=api_gateway.spec.name,
             host=api_gateway.spec.host,
             path=api_gateway.spec.path,
-            authentication_mode=str(api_gateway.spec.authenticationMode),
+            authentication=APIGatewayAuthenticator.from_scheme(api_gateway.spec),
             functions=functions,
             canary=canary,
         )
@@ -141,26 +280,26 @@ class APIGateway:
             spec=mlrun.common.schemas.APIGatewaySpec(
                 name=self.name,
                 description=self.description,
+                host=self.host,
                 path=self.path,
-                authentication_mode=mlrun.common.schemas.APIGatewayAuthenticationMode.from_str(
-                    self.authentication_mode
+                authenticationMode=mlrun.common.schemas.APIGatewayAuthenticationMode.from_str(
+                    self.authentication.authentication_mode
                 ),
                 upstreams=upstreams,
             ),
         )
-        if (
-            self.authentication_mode
-            is NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_BASIC_AUTH
-        ):
-            api_gateway.spec.authentication = mlrun.common.schemas.APIGatewayBasicAuth(
-                username=self._username, password=self._password
-            )
+        api_gateway.spec.authentication = self.authentication.to_scheme()
         return api_gateway
 
     @property
     def invoke_url(
         self,
     ):
+        """
+        Get the invoke URL.
+
+        :return: (str) The invoke URL.
+        """
         return urljoin(self.host, self.path)
 
     def _validate(
@@ -180,8 +319,6 @@ class APIGateway:
             ],
         ],
         canary: Optional[list[int]] = None,
-        username: Optional[str] = None,
-        password: Optional[str] = None,
     ):
         if not name:
             raise mlrun.errors.MLRunInvalidArgumentError(
@@ -192,26 +329,23 @@ class APIGateway:
 
         # validating canary
         if canary:
-            if len(self.functions) != len(canary):
-                raise mlrun.errors.MLRunInvalidArgumentError(
-                    "Function and canary lists lengths do not match"
-                )
-            for canary_percent in canary:
-                if canary_percent < 0 or canary_percent > 100:
-                    raise mlrun.errors.MLRunInvalidArgumentError(
-                        "The percentage value must be in the range from 0 to 100"
-                    )
-            if sum(canary) != 100:
+            self._validate_canary(canary)
+
+    def _validate_canary(self, canary: list[int]):
+        if len(self.functions) != len(canary):
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "Function and canary lists lengths do not match"
+            )
+        for canary_percent in canary:
+            if canary_percent < 0 or canary_percent > 100:
                 raise mlrun.errors.MLRunInvalidArgumentError(
-                    "The sum of canary function percents should be equal to 100"
+                    "The percentage value must be in the range from 0 to 100"
                 )
-
-        # validating auth
-        if username and not password:
-            raise mlrun.errors.MLRunInvalidArgumentError("Password is not specified")
-
-        if password and not username:
-            raise mlrun.errors.MLRunInvalidArgumentError("Username is not specified")
+        if sum(canary) != 100:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "The sum of canary function percents should be equal to 100"
+            )
+        return canary
 
     @staticmethod
     def _validate_functions(
@@ -257,17 +391,10 @@ class APIGateway:
                     f"input function {function_name} "
                     f"does not belong to this project"
                 )
-            function_names.append(func.uri)
+            nuclio_name = get_fullname(function_name, project, func.metadata.tag)
+            function_names.append(nuclio_name)
         return function_names
 
-    @staticmethod
-    def _enrich_authentication_mode(username, password):
-        return (
-            NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_NONE
-            if username is not None and password is not None
-            else NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_BASIC_AUTH
-        )
-
     @staticmethod
     def _generate_basic_auth(username: str, password: str):
         token = base64.b64encode(f"{username}:{password}".encode()).decode()

mlrun/runtimes/nuclio/function.py

@@ -775,6 +775,9 @@ class RemoteRuntime(KubeResource):
             ] = self.metadata.credentials.access_key
         return runtime_env
 
+    def _get_serving_spec(self):
+        return None
+
     def _get_nuclio_config_spec_env(self):
         env_dict = {}
         external_source_env_dict = {}

mlrun/runtimes/nuclio/serving.py

@@ -14,8 +14,9 @@
 
 import json
 import os
+import warnings
 from copy import deepcopy
-from typing import Union
+from typing import TYPE_CHECKING, Optional, Union
 
 import nuclio
 from nuclio import KafkaTrigger
@@ -24,7 +25,6 @@ import mlrun
 import mlrun.common.schemas
 from mlrun.datastore import parse_kafka_url
 from mlrun.model import ObjectList
-from mlrun.model_monitoring.tracking_policy import TrackingPolicy
 from mlrun.runtimes.function_reference import FunctionReference
 from mlrun.secrets import SecretsStore
 from mlrun.serving.server import GraphServer, create_graph_server
@@ -43,6 +43,10 @@ from .function import NuclioSpec, RemoteRuntime
 
 serving_subkind = "serving_v2"
 
+if TYPE_CHECKING:
+    # remove this block in 1.9.0
+    from mlrun.model_monitoring import TrackingPolicy
+
 
 def new_v2_model_server(
     name,
@@ -291,7 +295,9 @@ class ServingRuntime(RemoteRuntime):
                         "provided class is not a router step, must provide a router class in router topology"
                     )
             else:
-                step = RouterStep(class_name=class_name, class_args=class_args)
+                step = RouterStep(
+                    class_name=class_name, class_args=class_args, engine=engine
+                )
             self.spec.graph = step
         elif topology == StepKinds.flow:
             self.spec.graph = RootFlowStep(engine=engine)
@@ -303,12 +309,12 @@ class ServingRuntime(RemoteRuntime):
 
     def set_tracking(
         self,
-        stream_path: str = None,
-        batch: int = None,
-        sample: int = None,
-        stream_args: dict = None,
-        tracking_policy: Union[TrackingPolicy, dict] = None,
-    ):
+        stream_path: Optional[str] = None,
+        batch: Optional[int] = None,
+        sample: Optional[int] = None,
+        stream_args: Optional[dict] = None,
+        tracking_policy: Optional[Union["TrackingPolicy", dict]] = None,
+    ) -> None:
         """apply on your serving function to monitor a deployed model, including real-time dashboards to detect drift
            and analyze performance.
 
@@ -317,31 +323,17 @@ class ServingRuntime(RemoteRuntime):
         :param batch:           Micro batch size (send micro batches of N records at a time).
         :param sample:          Sample size (send only one of N records).
         :param stream_args:     Stream initialization parameters, e.g. shards, retention_in_hours, ..
-        :param tracking_policy: Tracking policy object or a dictionary that will be converted into a tracking policy
-                                object. By using TrackingPolicy, the user can apply his model monitoring requirements,
-                                such as setting the scheduling policy of the model monitoring batch job or changing
-                                the image of the model monitoring stream.
 
                                 example::
 
                                     # initialize a new serving function
                                     serving_fn = mlrun.import_function("hub://v2-model-server", new_name="serving")
-                                    # apply model monitoring and set monitoring batch job to run every 3 hours
-                                    tracking_policy = {'default_batch_intervals':"0 */3 * * *"}
-                                    serving_fn.set_tracking(tracking_policy=tracking_policy)
+                                    # apply model monitoring
+                                    serving_fn.set_tracking()
 
         """
-
         # Applying model monitoring configurations
         self.spec.track_models = True
-        self.spec.tracking_policy = None
-        if tracking_policy:
-            if isinstance(tracking_policy, dict):
-                # Convert tracking policy dictionary into `model_monitoring.TrackingPolicy` object
-                self.spec.tracking_policy = TrackingPolicy.from_dict(tracking_policy)
-            else:
-                # Tracking_policy is already a `model_monitoring.TrackingPolicy` object
-                self.spec.tracking_policy = tracking_policy
 
         if stream_path:
             self.spec.parameters["log_stream"] = stream_path
@@ -351,6 +343,14 @@ class ServingRuntime(RemoteRuntime):
             self.spec.parameters["log_stream_sample"] = sample
         if stream_args:
             self.spec.parameters["stream_args"] = stream_args
+        if tracking_policy is not None:
+            warnings.warn(
+                "The `tracking_policy` argument is deprecated from version 1.7.0 "
+                "and has no effect. It will be removed in 1.9.0.\n"
+                "To set the desired model monitoring time window and schedule, use "
+                "the `base_period` argument in `project.enable_model_monitoring()`.",
+                FutureWarning,
+            )
 
     def add_model(
         self,
@@ -644,8 +644,7 @@ class ServingRuntime(RemoteRuntime):
             force_build=force_build,
         )
 
-    def _get_runtime_env(self):
-        env = super()._get_runtime_env()
+    def _get_serving_spec(self):
         function_name_uri_map = {f.name: f.uri(self) for f in self.spec.function_refs}
 
         serving_spec = {
@@ -658,9 +657,7 @@ class ServingRuntime(RemoteRuntime):
             "graph_initializer": self.spec.graph_initializer,
             "error_stream": self.spec.error_stream,
             "track_models": self.spec.track_models,
-            "tracking_policy": self.spec.tracking_policy.to_dict()
-            if self.spec.tracking_policy
-            else None,
+            "tracking_policy": None,
             "default_content_type": self.spec.default_content_type,
         }
 
@@ -668,8 +665,7 @@ class ServingRuntime(RemoteRuntime):
             self._secrets = SecretsStore.from_list(self.spec.secret_sources)
             serving_spec["secret_sources"] = self._secrets.to_serial()
 
-        env["SERVING_SPEC_ENV"] = json.dumps(serving_spec)
-        return env
+        return json.dumps(serving_spec)
 
     def to_mock_server(
         self,

mlrun/runtimes/pod.py

@@ -1057,6 +1057,32 @@ 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

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
@@ -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."
             )

mlrun/serving/states.py

@@ -14,7 +14,6 @@
 
 __all__ = ["TaskStep", "RouterStep", "RootFlowStep", "ErrorStep"]
 
-import asyncio
 import os
 import pathlib
 import traceback
@@ -591,7 +590,7 @@ class RouterStep(TaskStep):
 
     kind = "router"
     default_shape = "doubleoctagon"
-    _dict_fields = _task_step_fields + ["routes"]
+    _dict_fields = _task_step_fields + ["routes", "engine"]
     _default_class = "mlrun.serving.ModelRouter"
 
     def __init__(
@@ -604,6 +603,7 @@ class RouterStep(TaskStep):
         function: str = None,
         input_path: str = None,
         result_path: str = None,
+        engine: str = None,
     ):
         super().__init__(
             class_name,
@@ -616,6 +616,8 @@ class RouterStep(TaskStep):
         )
         self._routes: ObjectDict = None
         self.routes = routes
+        self.engine = engine
+        self._controller = None
 
     def get_children(self):
         """get child steps (routes)"""
@@ -685,6 +687,33 @@ class RouterStep(TaskStep):
         self._set_error_handler()
         self._post_init(mode)
 
+        if self.engine == "async":
+            self._build_async_flow()
+            self._run_async_flow()
+
+    def _build_async_flow(self):
+        """initialize and build the async/storey DAG"""
+
+        self.respond()
+        source, self._wait_for_result = _init_async_objects(self.context, [self])
+        source.to(self.async_object)
+
+        self._async_flow = source
+
+    def _run_async_flow(self):
+        self._controller = self._async_flow.run()
+
+    def run(self, event, *args, **kwargs):
+        if self._controller:
+            # async flow (using storey)
+            event._awaitable_result = None
+            resp = self._controller.emit(
+                event, return_awaitable_result=self._wait_for_result
+            )
+            return resp.await_result()
+
+        return super().run(event, *args, **kwargs)
+
     def __getitem__(self, name):
         return self._routes[name]
 
@@ -1205,18 +1234,9 @@ class FlowStep(BaseStep):
         """wait for completion of run in async flows"""
 
         if self._controller:
-            if asyncio.iscoroutinefunction(self._controller.await_termination):
-
-                async def terminate_and_await_termination():
-                    if hasattr(self._controller, "terminate"):
-                        await self._controller.terminate()
-                    return await self._controller.await_termination()
-
-                return terminate_and_await_termination()
-            else:
-                if hasattr(self._controller, "terminate"):
-                    self._controller.terminate()
-                return self._controller.await_termination()
+            if hasattr(self._controller, "terminate"):
+                self._controller.terminate()
+            return self._controller.await_termination()
 
     def plot(self, filename=None, format=None, source=None, targets=None, **kw):
         """plot/save graph using graphviz