mlrun 1.6.2rc6__py3-none-any.whl → 1.6.3__py3-none-any.whl

mlrun/datastore/v3io.py

@@ -12,32 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import mmap
-import os
 import time
-from copy import deepcopy
 from datetime import datetime
 
 import fsspec
-import requests
-import v3io.dataplane
+import v3io
+from v3io.dataplane.response import HttpResponseError
 
 import mlrun
-from mlrun.datastore.helpers import ONE_GB, ONE_MB
 
 from ..platforms.iguazio import parse_path, split_path
 from .base import (
     DataStore,
     FileStats,
     basic_auth_header,
-    get_range,
-    http_get,
-    http_head,
-    http_put,
-    http_upload,
 )
 
 V3IO_LOCAL_ROOT = "v3io"
+V3IO_DEFAULT_UPLOAD_CHUNK_SIZE = 1024 * 1024 * 100
 
 
 class V3ioStore(DataStore):
@@ -47,17 +39,18 @@ class V3ioStore(DataStore):
 
         self.headers = None
         self.secure = self.kind == "v3ios"
+
+        token = self._get_secret_or_env("V3IO_ACCESS_KEY")
+        username = self._get_secret_or_env("V3IO_USERNAME")
+        password = self._get_secret_or_env("V3IO_PASSWORD")
         if self.endpoint.startswith("https://"):
             self.endpoint = self.endpoint[len("https://") :]
             self.secure = True
         elif self.endpoint.startswith("http://"):
             self.endpoint = self.endpoint[len("http://") :]
             self.secure = False
-
-        token = self._get_secret_or_env("V3IO_ACCESS_KEY")
-        username = self._get_secret_or_env("V3IO_USERNAME")
-        password = self._get_secret_or_env("V3IO_PASSWORD")
-
+        self.client = v3io.dataplane.Client(access_key=token, endpoint=self.url)
+        self.object = self.client.object
         self.auth = None
         self.token = token
         if token:
@@ -65,6 +58,16 @@ class V3ioStore(DataStore):
         elif username and password:
             self.headers = basic_auth_header(username, password)
 
+    @staticmethod
+    def _do_object_request(function: callable, *args, **kwargs):
+        try:
+            return function(*args, **kwargs)
+        except HttpResponseError as http_response_error:
+            raise mlrun.errors.err_for_status_code(
+                status_code=http_response_error.status_code,
+                message=mlrun.errors.err_to_str(http_response_error),
+            )
+
     @staticmethod
     def uri_to_ipython(endpoint, subpath):
         return V3IO_LOCAL_ROOT + subpath
@@ -89,80 +92,79 @@ class V3ioStore(DataStore):
         )
         return self._sanitize_storage_options(res)
 
-    def _upload(self, key: str, src_path: str, max_chunk_size: int = ONE_GB):
+    def _upload(
+        self,
+        key: str,
+        src_path: str,
+        max_chunk_size: int = V3IO_DEFAULT_UPLOAD_CHUNK_SIZE,
+    ):
         """helper function for upload method, allows for controlling max_chunk_size in testing"""
-        file_size = os.path.getsize(src_path)  # in bytes
-        if file_size <= ONE_MB:
-            http_upload(self.url + self._join(key), src_path, self.headers, None)
-            return
-        append_header = deepcopy(self.headers)
-        append_header["Range"] = "-1"
-
-        # chunk must be a multiple of the ALLOCATIONGRANULARITY
-        # https://docs.python.org/3/library/mmap.html
-        if residue := max_chunk_size % mmap.ALLOCATIONGRANULARITY:
-            # round down to the nearest multiple of ALLOCATIONGRANULARITY
-            max_chunk_size -= residue
-
+        container, path = split_path(self._join(key))
         with open(src_path, "rb") as file_obj:
-            file_offset = 0
-            while file_offset < file_size:
-                chunk_size = min(file_size - file_offset, max_chunk_size)
-                with mmap.mmap(
-                    file_obj.fileno(),
-                    length=chunk_size,
-                    access=mmap.ACCESS_READ,
-                    offset=file_offset,
-                ) as mmap_obj:
-                    http_put(
-                        self.url + self._join(key),
-                        mmap_obj,
-                        append_header if file_offset else self.headers,
-                        None,
-                    )
-                    file_offset += chunk_size
+            append = False
+            while True:
+                data = memoryview(file_obj.read(max_chunk_size))
+                if not data:
+                    break
+                self._do_object_request(
+                    self.object.put,
+                    container=container,
+                    path=path,
+                    body=data,
+                    append=append,
+                )
+                append = True
 
     def upload(self, key, src_path):
         return self._upload(key, src_path)
 
     def get(self, key, size=None, offset=0):
-        headers = self.headers
-        if size or offset:
-            headers = deepcopy(headers)
-            headers["Range"] = get_range(size, offset)
-        return http_get(self.url + self._join(key), headers)
-
-    def _put(self, key, data, max_chunk_size: int = ONE_GB):
+        container, path = split_path(self._join(key))
+        return self._do_object_request(
+            function=self.object.get,
+            container=container,
+            path=path,
+            offset=offset,
+            num_bytes=size,
+        ).body
+
+    def _put(
+        self,
+        key,
+        data,
+        append=False,
+        max_chunk_size: int = V3IO_DEFAULT_UPLOAD_CHUNK_SIZE,
+    ):
         """helper function for put method, allows for controlling max_chunk_size in testing"""
+        container, path = split_path(self._join(key))
         buffer_size = len(data)  # in bytes
-        if buffer_size <= ONE_MB:
-            http_put(self.url + self._join(key), data, self.headers, None)
-            return
-        append_header = deepcopy(self.headers)
-        append_header["Range"] = "-1"
         buffer_offset = 0
         try:
             data = memoryview(data)
         except TypeError:
             pass
 
-        with requests.Session() as requests_session:
-            while buffer_offset < buffer_size:
-                chunk_size = min(buffer_size - buffer_offset, max_chunk_size)
-                http_put(
-                    self.url + self._join(key),
-                    data[buffer_offset : buffer_offset + chunk_size],
-                    append_header if buffer_offset else self.headers,
-                    None,
-                    requests_session,
-                )
-                buffer_offset += chunk_size
+        while buffer_offset < buffer_size:
+            chunk_size = min(buffer_size - buffer_offset, max_chunk_size)
+            append = True if buffer_offset or append else False
+            self._do_object_request(
+                self.object.put,
+                container=container,
+                path=path,
+                body=data[buffer_offset : buffer_offset + chunk_size],
+                append=append,
+            )
+            buffer_offset += chunk_size
 
     def put(self, key, data, append=False):
-        return self._put(key, data)
+        return self._put(key, data, append)
 
     def stat(self, key):
-        head = http_head(self.url + self._join(key), self.headers)
+        container, path = split_path(self._join(key))
+        response = self._do_object_request(
+            function=self.object.head, container=container, path=path
+        )
+        head = dict(response.headers)
         size = int(head.get("Content-Length", "0"))
         datestr = head.get("Last-Modified", "0")
         modified = time.mktime(
@@ -171,7 +173,6 @@ class V3ioStore(DataStore):
         return FileStats(size, modified)
 
     def listdir(self, key):
-        v3io_client = v3io.dataplane.Client(endpoint=self.url, access_key=self.token)
         container, subpath = split_path(self._join(key))
         if not subpath.endswith("/"):
             subpath += "/"
@@ -180,7 +181,7 @@ class V3ioStore(DataStore):
         subpath_length = len(subpath) - 1
 
         try:
-            response = v3io_client.container.list(
+            response = self.client.container.list(
                 container=container,
                 path=subpath,
                 get_all_attributes=False,

mlrun/db/auth_utils.py

@@ -0,0 +1,152 @@
+# Copyright 2024 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABC, abstractmethod
+from datetime import datetime, timedelta
+
+import requests
+
+import mlrun.errors
+from mlrun.utils import logger
+
+
+class TokenProvider(ABC):
+    @abstractmethod
+    def get_token(self):
+        pass
+
+    @abstractmethod
+    def is_iguazio_session(self):
+        pass
+
+
+class StaticTokenProvider(TokenProvider):
+    def __init__(self, token: str):
+        self.token = token
+
+    def get_token(self):
+        return self.token
+
+    def is_iguazio_session(self):
+        return mlrun.platforms.iguazio.is_iguazio_session(self.token)
+
+
+class OAuthClientIDTokenProvider(TokenProvider):
+    def __init__(
+        self, token_endpoint: str, client_id: str, client_secret: str, timeout=5
+    ):
+        if not token_endpoint or not client_id or not client_secret:
+            raise mlrun.errors.MLRunValueError(
+                "Invalid client_id configuration for authentication. Must provide token endpoint, client-id and secret"
+            )
+        self.token_endpoint = token_endpoint
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.timeout = timeout
+
+        # Since we're only issuing POST requests, which are actually a disguised GET, then it's ok to allow retries
+        # on them.
+        self._session = mlrun.utils.HTTPSessionWithRetry(
+            retry_on_post=True,
+            verbose=True,
+        )
+
+        self._cleanup()
+        self._refresh_token_if_needed()
+
+    def get_token(self):
+        self._refresh_token_if_needed()
+        return self.token
+
+    def is_iguazio_session(self):
+        return False
+
+    def _cleanup(self):
+        self.token = self.token_expiry_time = self.token_refresh_time = None
+
+    def _refresh_token_if_needed(self):
+        now = datetime.now()
+        if self.token:
+            if self.token_refresh_time and now <= self.token_refresh_time:
+                return self.token
+
+            # We only cleanup if token was really expired - even if we fail in refreshing the token, we can still
+            # use the existing one given that it's not expired.
+            if now >= self.token_expiry_time:
+                self._cleanup()
+
+        self._issue_token_request()
+        return self.token
+
+    def _issue_token_request(self, raise_on_error=False):
+        try:
+            headers = {"Content-Type": "application/x-www-form-urlencoded"}
+            request_body = {
+                "grant_type": "client_credentials",
+                "client_id": self.client_id,
+                "client_secret": self.client_secret,
+            }
+            response = self._session.request(
+                "POST",
+                self.token_endpoint,
+                timeout=self.timeout,
+                headers=headers,
+                data=request_body,
+            )
+        except requests.RequestException as exc:
+            error = f"Retrieving token failed: {mlrun.errors.err_to_str(exc)}"
+            if raise_on_error:
+                raise mlrun.errors.MLRunRuntimeError(error) from exc
+            else:
+                logger.warning(error)
+                return
+
+        if not response.ok:
+            error = "No error available"
+            if response.content:
+                try:
+                    data = response.json()
+                    error = data.get("error")
+                except Exception:
+                    pass
+            logger.warning(
+                "Retrieving token failed", status=response.status_code, error=error
+            )
+            if raise_on_error:
+                mlrun.errors.raise_for_status(response)
+            return
+
+        self._parse_response(response.json())
+
+    def _parse_response(self, data: dict):
+        # Response is described in https://datatracker.ietf.org/doc/html/rfc6749#section-4.4.3
+        # According to spec, there isn't a refresh token - just the access token and its expiry time (in seconds).
+        self.token = data.get("access_token")
+        expires_in = data.get("expires_in")
+        if not self.token or not expires_in:
+            token_str = "****" if self.token else "missing"
+            logger.warning(
+                "Failed to parse token response", token=token_str, expires_in=expires_in
+            )
+            return
+
+        now = datetime.now()
+        self.token_expiry_time = now + timedelta(seconds=expires_in)
+        self.token_refresh_time = now + timedelta(seconds=expires_in / 2)
+        logger.info(
+            "Successfully retrieved client-id token",
+            expires_in=expires_in,
+            expiry=str(self.token_expiry_time),
+            refresh=str(self.token_refresh_time),
+        )

mlrun/db/base.py

@@ -677,3 +677,21 @@ class RunDBInterface(ABC):
         self, func_url: str = None, function: "mlrun.runtimes.BaseRuntime" = None
     ):
         pass
+
+    def submit_workflow(
+        self,
+        project: str,
+        name: str,
+        workflow_spec: Union[
+            "mlrun.projects.pipelines.WorkflowSpec",
+            "mlrun.common.schemas.WorkflowSpec",
+            dict,
+        ],
+        arguments: Optional[dict] = None,
+        artifact_path: Optional[str] = None,
+        source: Optional[str] = None,
+        run_name: Optional[str] = None,
+        namespace: Optional[str] = None,
+        notifications: list["mlrun.model.Notification"] = None,
+    ) -> "mlrun.common.schemas.WorkflowResponse":
+        pass