mlrun 1.7.0rc14__py3-none-any.whl → 1.7.0rc21__py3-none-any.whl

mlrun/runtimes/nuclio/api_gateway.py

@@ -17,31 +17,30 @@ from typing import Optional, Union
 from urllib.parse import urljoin
 
 import requests
-from requests.auth import HTTPBasicAuth
+from nuclio.auth import AuthInfo as NuclioAuthInfo
+from nuclio.auth import AuthKinds as NuclioAuthKinds
 
 import mlrun
-import mlrun.common.schemas
+import mlrun.common.constants as mlrun_constants
+import mlrun.common.schemas as schemas
+import mlrun.common.types
+from mlrun.model import ModelObj
+from mlrun.platforms.iguazio import min_iguazio_versions
+from mlrun.utils import logger
 
-from ...model import ModelObj
-from ..utils import logger
-from .function import RemoteRuntime, get_fullname, min_nuclio_versions
-from .serving import ServingRuntime
+from .function import min_nuclio_versions
 
-NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_BASIC_AUTH = "basicAuth"
-NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_NONE = "none"
-PROJECT_NAME_LABEL = "nuclio.io/project-name"
 
-
-class APIGatewayAuthenticator(typing.Protocol):
+class Authenticator(typing.Protocol):
     @property
     def authentication_mode(self) -> str:
-        return NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_NONE
+        return schemas.APIGatewayAuthenticationMode.none.value
 
     @classmethod
-    def from_scheme(cls, api_gateway_spec: mlrun.common.schemas.APIGatewaySpec):
+    def from_scheme(cls, api_gateway_spec: schemas.APIGatewaySpec):
         if (
             api_gateway_spec.authenticationMode
-            == NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_BASIC_AUTH
+            == schemas.APIGatewayAuthenticationMode.basic.value
         ):
             if api_gateway_spec.authentication:
                 return BasicAuth(
@@ -50,15 +49,24 @@ class APIGatewayAuthenticator(typing.Protocol):
                 )
             else:
                 return BasicAuth()
+        elif (
+            api_gateway_spec.authenticationMode
+            == schemas.APIGatewayAuthenticationMode.access_key.value
+        ):
+            return AccessKeyAuth()
         else:
             return NoneAuth()
 
     def to_scheme(
         self,
-    ) -> Optional[dict[str, Optional[mlrun.common.schemas.APIGatewayBasicAuth]]]:
+    ) -> Optional[dict[str, Optional[schemas.APIGatewayBasicAuth]]]:
         return None
 
 
+class APIGatewayAuthenticator(Authenticator, ModelObj):
+    _dict_fields = ["authentication_mode"]
+
+
 class NoneAuth(APIGatewayAuthenticator):
     """
     An API gateway authenticator with no authentication.
@@ -81,18 +89,28 @@ class BasicAuth(APIGatewayAuthenticator):
 
     @property
     def authentication_mode(self) -> str:
-        return NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_BASIC_AUTH
+        return schemas.APIGatewayAuthenticationMode.basic.value
 
     def to_scheme(
         self,
-    ) -> Optional[dict[str, Optional[mlrun.common.schemas.APIGatewayBasicAuth]]]:
+    ) -> Optional[dict[str, Optional[schemas.APIGatewayBasicAuth]]]:
         return {
-            "basicAuth": mlrun.common.schemas.APIGatewayBasicAuth(
+            "basicAuth": schemas.APIGatewayBasicAuth(
                 username=self._username, password=self._password
             )
         }
 
 
+class AccessKeyAuth(APIGatewayAuthenticator):
+    """
+    An API gateway authenticator with access key authentication.
+    """
+
+    @property
+    def authentication_mode(self) -> str:
+        return schemas.APIGatewayAuthenticationMode.access_key.value
+
+
 class APIGatewayMetadata(ModelObj):
     _dict_fields = ["name", "namespace", "labels", "annotations", "creation_timestamp"]
 
@@ -138,17 +156,17 @@ class APIGatewaySpec(ModelObj):
     def __init__(
         self,
         functions: Union[
-            list[str],
-            Union[
-                list[
-                    Union[
-                        RemoteRuntime,
-                        ServingRuntime,
-                    ]
-                ],
-                RemoteRuntime,
-                ServingRuntime,
+            list[
+                Union[
+                    str,
+                    "mlrun.runtimes.nuclio.function.RemoteRuntime",
+                    "mlrun.runtimes.nuclio.serving.ServingRuntime",
+                    "mlrun.runtimes.nuclio.application.ApplicationRuntime",
+                ]
             ],
+            "mlrun.runtimes.nuclio.function.RemoteRuntime",
+            "mlrun.runtimes.nuclio.serving.ServingRuntime",
+            "mlrun.runtimes.nuclio.application.ApplicationRuntime",
         ],
         project: str = None,
         description: str = "",
@@ -156,20 +174,24 @@ class APIGatewaySpec(ModelObj):
         path: str = "/",
         authentication: Optional[APIGatewayAuthenticator] = NoneAuth(),
         canary: Optional[list[int]] = None,
+        ports: Optional[list[int]] = None,
     ):
         """
         :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`
+            :py:class:`~mlrun.runtimes.nuclio.serving.ServingRuntime` OR
+            :py:class:`~mlrun.runtimes.nuclio.application.ApplicationRuntime`
         :param project: The project name
         :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]
+        :param canary: The canary percents for the API gateway of type list[int]; for instance: [20,80] (optional)
+        :param ports: The ports of the API gateway, as a list of integers that correspond to the functions in the
+            functions list. for instance: [8050] or [8050, 8081] (optional)
         """
         self.description = description
         self.host = host
@@ -178,26 +200,28 @@ class APIGatewaySpec(ModelObj):
         self.functions = functions
         self.canary = canary
         self.project = project
+        self.ports = ports
 
-        self.validate(project=project, functions=functions, canary=canary)
+        self.validate(project=project, functions=functions, canary=canary, ports=ports)
 
     def validate(
         self,
         project: str,
         functions: Union[
-            list[str],
-            Union[
-                list[
-                    Union[
-                        RemoteRuntime,
-                        ServingRuntime,
-                    ]
-                ],
-                RemoteRuntime,
-                ServingRuntime,
+            list[
+                Union[
+                    str,
+                    "mlrun.runtimes.nuclio.function.RemoteRuntime",
+                    "mlrun.runtimes.nuclio.serving.ServingRuntime",
+                    "mlrun.runtimes.nuclio.application.ApplicationRuntime",
+                ]
             ],
+            "mlrun.runtimes.nuclio.function.RemoteRuntime",
+            "mlrun.runtimes.nuclio.serving.ServingRuntime",
+            "mlrun.runtimes.nuclio.application.ApplicationRuntime",
         ],
         canary: Optional[list[int]] = None,
+        ports: Optional[list[int]] = None,
     ):
         self.functions = self._validate_functions(project=project, functions=functions)
 
@@ -205,6 +229,10 @@ class APIGatewaySpec(ModelObj):
         if canary:
             self.canary = self._validate_canary(canary)
 
+        # validating ports
+        if ports:
+            self.ports = self._validate_ports(ports)
+
     def _validate_canary(self, canary: list[int]):
         if len(self.functions) != len(canary):
             raise mlrun.errors.MLRunInvalidArgumentError(
@@ -221,20 +249,29 @@ class APIGatewaySpec(ModelObj):
             )
         return canary
 
+    def _validate_ports(self, ports):
+        if len(self.functions) != len(ports):
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "Function and port lists lengths do not match"
+            )
+
+        return ports
+
     @staticmethod
     def _validate_functions(
         project: str,
         functions: Union[
-            list[str],
-            Union[
-                list[
-                    Union[
-                        RemoteRuntime,
-                        ServingRuntime,
-                    ]
-                ],
-                Union[RemoteRuntime, ServingRuntime],
+            list[
+                Union[
+                    str,
+                    "mlrun.runtimes.nuclio.function.RemoteRuntime",
+                    "mlrun.runtimes.nuclio.serving.ServingRuntime",
+                    "mlrun.runtimes.nuclio.application.ApplicationRuntime",
+                ]
             ],
+            "mlrun.runtimes.nuclio.function.RemoteRuntime",
+            "mlrun.runtimes.nuclio.serving.ServingRuntime",
+            "mlrun.runtimes.nuclio.application.ApplicationRuntime",
         ],
     ):
         if not isinstance(functions, list):
@@ -250,7 +287,21 @@ class APIGatewaySpec(ModelObj):
         function_names = []
         for func in functions:
             if isinstance(func, str):
-                function_names.append(func)
+                # check whether the function was passed as a URI or just a name
+                parsed_project, function_name, _, _ = (
+                    mlrun.common.helpers.parse_versioned_object_uri(func)
+                )
+
+                if parsed_project and function_name:
+                    # check that parsed project and passed project are the same
+                    if parsed_project != project:
+                        raise mlrun.errors.MLRunInvalidArgumentError(
+                            "Function doesn't belong to passed project"
+                        )
+                    function_uri = func
+                else:
+                    function_uri = mlrun.utils.generate_object_uri(project, func)
+                function_names.append(function_uri)
                 continue
 
             function_name = (
@@ -265,8 +316,13 @@ class APIGatewaySpec(ModelObj):
                     f"input function {function_name} "
                     f"does not belong to this project"
                 )
-            nuclio_name = get_fullname(function_name, project, func.metadata.tag)
-            function_names.append(nuclio_name)
+            function_uri = mlrun.utils.generate_object_uri(
+                project,
+                function_name,
+                func.metadata.tag,
+                func.metadata.hash,
+            )
+            function_names.append(function_uri)
         return function_names
 
 
@@ -312,8 +368,9 @@ class APIGateway(ModelObj):
     def invoke(
         self,
         method="POST",
-        headers: dict = {},
-        auth: Optional[tuple[str, str]] = None,
+        headers: dict = None,
+        credentials: Optional[tuple[str, str]] = None,
+        path: Optional[str] = None,
         **kwargs,
     ):
         """
@@ -321,7 +378,9 @@ class APIGateway(ModelObj):
 
         :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 credentials: (Optional[tuple[str, str]], optional) The (username,password) for the invocation if required
+            can also be set by the environment variable (_, V3IO_ACCESS_KEY) for access key authentication.
+        :param path: (str, optional) The sub-path for the invocation.
         :param kwargs: (dict) Additional keyword arguments.
 
         :return: The response from the API gateway invocation.
@@ -338,20 +397,44 @@ class APIGateway(ModelObj):
                 f"API gateway is not ready. " f"Current state: {self.state}"
             )
 
+        auth = None
+
         if (
             self.spec.authentication.authentication_mode
-            == NUCLIO_API_GATEWAY_AUTHENTICATION_MODE_BASIC_AUTH
-            and not auth
+            == schemas.APIGatewayAuthenticationMode.basic.value
         ):
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                "API Gateway invocation requires authentication. Please pass credentials"
-            )
+            if not credentials:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "API Gateway invocation requires authentication. Please pass credentials"
+                )
+            auth = NuclioAuthInfo(
+                username=credentials[0], password=credentials[1]
+            ).to_requests_auth()
+
+        if (
+            self.spec.authentication.authentication_mode
+            == schemas.APIGatewayAuthenticationMode.access_key.value
+        ):
+            # inject access key from env
+            if credentials:
+                auth = NuclioAuthInfo(
+                    username=credentials[0],
+                    password=credentials[1],
+                    mode=NuclioAuthKinds.iguazio,
+                ).to_requests_auth()
+            else:
+                auth = NuclioAuthInfo().from_envvar().to_requests_auth()
+            if not auth:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "API Gateway invocation requires authentication. Please set V3IO_ACCESS_KEY env var"
+                )
+        url = urljoin(self.invoke_url, path or "")
         return requests.request(
             method=method,
-            url=self.invoke_url,
-            headers=headers,
+            url=url,
+            headers=headers or {},
+            auth=auth,
             **kwargs,
-            auth=HTTPBasicAuth(*auth) if auth else None,
         )
 
     def wait_for_readiness(self, max_wait_time=90):
@@ -376,10 +459,10 @@ class APIGateway(ModelObj):
         )
 
     def is_ready(self):
-        if self.state is not mlrun.common.schemas.api_gateway.APIGatewayState.ready:
+        if self.state is not schemas.api_gateway.APIGatewayState.ready:
             # try to sync the state
             self.sync()
-        return self.state == mlrun.common.schemas.api_gateway.APIGatewayState.ready
+        return self.state == schemas.api_gateway.APIGatewayState.ready
 
     def sync(self):
         """
@@ -407,16 +490,27 @@ class APIGateway(ModelObj):
         """
         self.spec.authentication = BasicAuth(username=username, password=password)
 
+    @min_iguazio_versions("3.5.5")
+    def with_access_key_auth(self):
+        """
+        Set access key authentication for the API gateway.
+        """
+        self.spec.authentication = AccessKeyAuth()
+
     def with_canary(
         self,
         functions: Union[
-            list[str],
             list[
                 Union[
-                    RemoteRuntime,
-                    ServingRuntime,
+                    str,
+                    "mlrun.runtimes.nuclio.function.RemoteRuntime",
+                    "mlrun.runtimes.nuclio.serving.ServingRuntime",
+                    "mlrun.runtimes.nuclio.application.ApplicationRuntime",
                 ]
             ],
+            "mlrun.runtimes.nuclio.function.RemoteRuntime",
+            "mlrun.runtimes.nuclio.serving.ServingRuntime",
+            "mlrun.runtimes.nuclio.application.ApplicationRuntime",
         ],
         canary: list[int],
     ):
@@ -427,7 +521,8 @@ class APIGateway(ModelObj):
             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`
+            :py:class:`~mlrun.runtimes.nuclio.serving.ServingRuntime` OR
+            :py:class:`~mlrun.runtimes.nuclio.application.ApplicationRuntime`
         :param canary: The canary percents for the API gateway of type list[int]; for instance: [20,80]
 
         """
@@ -440,14 +535,27 @@ class APIGateway(ModelObj):
             project=self.spec.project, functions=functions, canary=canary
         )
 
+    def with_ports(self, ports: list[int]):
+        """
+        Set ports for the API gateway
+
+        :param ports: The ports of the API gateway, as a list of integers that correspond to the functions in the
+            functions list. for instance: [8050] or [8050, 8081]
+        """
+        self.spec.validate(
+            project=self.spec.project, functions=self.spec.functions, ports=ports
+        )
+
     @classmethod
-    def from_scheme(cls, api_gateway: mlrun.common.schemas.APIGateway):
-        project = api_gateway.metadata.labels.get(PROJECT_NAME_LABEL)
+    def from_scheme(cls, api_gateway: schemas.APIGateway):
+        project = api_gateway.metadata.labels.get(
+            mlrun_constants.MLRunInternalLabels.nuclio_project_name
+        )
         functions, canary = cls._resolve_canary(api_gateway.spec.upstreams)
         state = (
             api_gateway.status.state
             if api_gateway.status
-            else mlrun.common.schemas.APIGatewayState.none
+            else schemas.APIGatewayState.none
         )
         new_api_gateway = cls(
             metadata=APIGatewayMetadata(
@@ -466,14 +574,14 @@ class APIGateway(ModelObj):
         new_api_gateway.state = state
         return new_api_gateway
 
-    def to_scheme(self) -> mlrun.common.schemas.APIGateway:
+    def to_scheme(self) -> schemas.APIGateway:
         upstreams = (
             [
-                mlrun.common.schemas.APIGatewayUpstream(
+                schemas.APIGatewayUpstream(
                     nucliofunction={"name": self.spec.functions[0]},
                     percentage=self.spec.canary[0],
                 ),
-                mlrun.common.schemas.APIGatewayUpstream(
+                schemas.APIGatewayUpstream(
                     # do not set percent for the second function,
                     # so we can define which function to display as a primary one in UI
                     nucliofunction={"name": self.spec.functions[1]},
@@ -481,22 +589,24 @@ class APIGateway(ModelObj):
             ]
             if self.spec.canary
             else [
-                mlrun.common.schemas.APIGatewayUpstream(
+                schemas.APIGatewayUpstream(
                     nucliofunction={"name": function_name},
                 )
                 for function_name in self.spec.functions
             ]
         )
-        api_gateway = mlrun.common.schemas.APIGateway(
-            metadata=mlrun.common.schemas.APIGatewayMetadata(
-                name=self.metadata.name, labels={}
-            ),
-            spec=mlrun.common.schemas.APIGatewaySpec(
+        if self.spec.ports:
+            for i, port in enumerate(self.spec.ports):
+                upstreams[i].port = port
+
+        api_gateway = schemas.APIGateway(
+            metadata=schemas.APIGatewayMetadata(name=self.metadata.name, labels={}),
+            spec=schemas.APIGatewaySpec(
                 name=self.metadata.name,
                 description=self.spec.description,
                 host=self.spec.host,
                 path=self.spec.path,
-                authenticationMode=mlrun.common.schemas.APIGatewayAuthenticationMode.from_str(
+                authenticationMode=schemas.APIGatewayAuthenticationMode.from_str(
                     self.spec.authentication.authentication_mode
                 ),
                 upstreams=upstreams,
@@ -526,7 +636,7 @@ class APIGateway(ModelObj):
 
     @staticmethod
     def _resolve_canary(
-        upstreams: list[mlrun.common.schemas.APIGatewayUpstream],
+        upstreams: list[schemas.APIGatewayUpstream],
     ) -> tuple[Union[list[str], None], Union[list[int], None]]:
         if len(upstreams) == 1:
             return [upstreams[0].nucliofunction.get("name")], None