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

mlrun/datastore/utils.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import math
 import tarfile
 import tempfile
 import typing
+import warnings
 from urllib.parse import parse_qs, urlparse
 
 import pandas as pd
@@ -23,24 +25,29 @@ import semver
 import mlrun.datastore
 
 
-def parse_kafka_url(url: str, bootstrap_servers: list = None) -> tuple[str, list]:
+def parse_kafka_url(
+    url: str, brokers: typing.Union[list, str] = None
+) -> tuple[str, list]:
     """Generating Kafka topic and adjusting a list of bootstrap servers.
 
     :param url:               URL path to parse using urllib.parse.urlparse.
-    :param bootstrap_servers: List of bootstrap servers for the kafka brokers.
+    :param brokers: List of kafka brokers.
 
     :return: A tuple of:
          [0] = Kafka topic value
          [1] = List of bootstrap servers
     """
-    bootstrap_servers = bootstrap_servers or []
+    brokers = brokers or []
+
+    if isinstance(brokers, str):
+        brokers = brokers.split(",")
 
     # Parse the provided URL into six components according to the general structure of a URL
     url = urlparse(url)
 
     # Add the network location to the bootstrap servers list
     if url.netloc:
-        bootstrap_servers = [url.netloc] + bootstrap_servers
+        brokers = [url.netloc] + brokers
 
     # Get the topic value from the parsed url
     query_dict = parse_qs(url.query)
@@ -49,7 +56,7 @@ def parse_kafka_url(url: str, bootstrap_servers: list = None) -> tuple[str, list
     else:
         topic = url.path
         topic = topic.lstrip("/")
-    return topic, bootstrap_servers
+    return topic, brokers
 
 
 def upload_tarball(source_dir, target, secrets=None):
@@ -159,3 +166,59 @@ def _generate_sql_query_with_time_filter(
             query = query.filter(getattr(table.c, time_column) <= end_time)
 
     return query, parse_dates
+
+
+def get_kafka_brokers_from_dict(options: dict, pop=False) -> typing.Optional[str]:
+    get_or_pop = options.pop if pop else options.get
+    kafka_brokers = get_or_pop("kafka_brokers", None)
+    if kafka_brokers:
+        return kafka_brokers
+    kafka_bootstrap_servers = get_or_pop("kafka_bootstrap_servers", None)
+    if kafka_bootstrap_servers:
+        warnings.warn(
+            "The 'kafka_bootstrap_servers' parameter is deprecated and will be removed in "
+            "1.9.0. Please pass the 'kafka_brokers' parameter instead.",
+            FutureWarning,
+        )
+    return kafka_bootstrap_servers
+
+
+def transform_list_filters_to_tuple(additional_filters):
+    tuple_filters = []
+    if not additional_filters:
+        return tuple_filters
+    validate_additional_filters(additional_filters)
+    for additional_filter in additional_filters:
+        tuple_filters.append(tuple(additional_filter))
+    return tuple_filters
+
+
+def validate_additional_filters(additional_filters):
+    nan_error_message = "using NaN in additional_filters is not supported"
+    if additional_filters in [None, [], ()]:
+        return
+    for filter_tuple in additional_filters:
+        if filter_tuple == () or filter_tuple == []:
+            continue
+        if not isinstance(filter_tuple, (list, tuple)):
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"mlrun supports additional_filters only as a list of tuples."
+                f" Current additional_filters: {additional_filters}"
+            )
+        if isinstance(filter_tuple[0], (list, tuple)):
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"additional_filters does not support nested list inside filter tuples except in -in- logic."
+                f" Current filter_tuple: {filter_tuple}."
+            )
+        if len(filter_tuple) != 3:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"illegal filter tuple length, {filter_tuple} in additional filters:"
+                f" {additional_filters}"
+            )
+        col_name, op, value = filter_tuple
+        if isinstance(value, float) and math.isnan(value):
+            raise mlrun.errors.MLRunInvalidArgumentError(nan_error_message)
+        elif isinstance(value, (list, tuple)):
+            for sub_value in value:
+                if isinstance(sub_value, float) and math.isnan(sub_value):
+                    raise mlrun.errors.MLRunInvalidArgumentError(nan_error_message)

mlrun/datastore/v3io.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import mmap
-import os
 import time
 from datetime import datetime
 
@@ -22,7 +20,6 @@ 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 (
@@ -32,6 +29,7 @@ from .base import (
 )
 
 V3IO_LOCAL_ROOT = "v3io"
+V3IO_DEFAULT_UPLOAD_CHUNK_SIZE = 1024 * 1024 * 10
 
 
 class V3ioStore(DataStore):
@@ -98,46 +96,28 @@ 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"""
         container, path = split_path(self._join(key))
-        file_size = os.path.getsize(src_path)  # in bytes
-        if file_size <= ONE_MB:
-            with open(src_path, "rb") as source_file:
-                data = source_file.read()
-            self._do_object_request(
-                self.object.put,
-                container=container,
-                path=path,
-                body=data,
-                append=False,
-            )
-            return
-        # 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
-
         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:
-                    append = file_offset != 0
-                    self._do_object_request(
-                        self.object.put,
-                        container=container,
-                        path=path,
-                        body=mmap_obj,
-                        append=append,
-                    )
-                    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)
@@ -152,19 +132,16 @@ class V3ioStore(DataStore):
             num_bytes=size,
         ).body
 
-    def _put(self, key, data, append=False, max_chunk_size: int = ONE_GB):
+    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:
-            self._do_object_request(
-                self.object.put,
-                container=container,
-                path=path,
-                body=data,
-                append=append,
-            )
-            return
         buffer_offset = 0
         try:
             data = memoryview(data)

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

@@ -16,8 +16,11 @@ import datetime
 from abc import ABC, abstractmethod
 from typing import Optional, Union
 
+import mlrun.alerts
+import mlrun.common.formatters
+import mlrun.common.runtimes.constants
 import mlrun.common.schemas
-import mlrun.model_monitoring.model_endpoint
+import mlrun.model_monitoring
 
 
 class RunDBError(Exception):
@@ -62,7 +65,10 @@ class RunDBInterface(ABC):
         uid: Optional[Union[str, list[str]]] = None,
         project: Optional[str] = None,
         labels: Optional[Union[str, list[str]]] = None,
-        state: Optional[str] = None,
+        state: Optional[
+            mlrun.common.runtimes.constants.RunStates
+        ] = None,  # Backward compatibility
+        states: Optional[list[mlrun.common.runtimes.constants.RunStates]] = None,
         sort: bool = True,
         last: int = 0,
         iter: bool = False,
@@ -117,7 +123,18 @@ class RunDBInterface(ABC):
         pass
 
     @abstractmethod
-    def del_artifact(self, key, tag="", project="", tree=None, uid=None):
+    def del_artifact(
+        self,
+        key,
+        tag="",
+        project="",
+        tree=None,
+        uid=None,
+        deletion_strategy: mlrun.common.schemas.artifact.ArtifactsDeletionStrategies = (
+            mlrun.common.schemas.artifact.ArtifactsDeletionStrategies.metadata_only
+        ),
+        secrets: dict = None,
+    ):
         pass
 
     @abstractmethod
@@ -251,7 +268,7 @@ class RunDBInterface(ABC):
     def list_projects(
         self,
         owner: str = None,
-        format_: mlrun.common.schemas.ProjectsFormat = mlrun.common.schemas.ProjectsFormat.name_only,
+        format_: mlrun.common.formatters.ProjectFormat = mlrun.common.formatters.ProjectFormat.name_only,
         labels: list[str] = None,
         state: mlrun.common.schemas.ProjectState = None,
     ) -> mlrun.common.schemas.ProjectsOutput:
@@ -427,8 +444,8 @@ class RunDBInterface(ABC):
         namespace: str = None,
         timeout: int = 30,
         format_: Union[
-            str, mlrun.common.schemas.PipelinesFormat
-        ] = mlrun.common.schemas.PipelinesFormat.summary,
+            str, mlrun.common.formatters.PipelineFormat
+        ] = mlrun.common.formatters.PipelineFormat.summary,
         project: str = None,
     ):
         pass
@@ -442,8 +459,8 @@ class RunDBInterface(ABC):
         page_token: str = "",
         filter_: str = "",
         format_: Union[
-            str, mlrun.common.schemas.PipelinesFormat
-        ] = mlrun.common.schemas.PipelinesFormat.metadata_only,
+            str, mlrun.common.formatters.PipelineFormat
+        ] = mlrun.common.formatters.PipelineFormat.metadata_only,
         page_size: int = None,
     ) -> mlrun.common.schemas.PipelinesOutput:
         pass
@@ -509,9 +526,7 @@ class RunDBInterface(ABC):
         self,
         project: str,
         endpoint_id: str,
-        model_endpoint: Union[
-            mlrun.model_monitoring.model_endpoint.ModelEndpoint, dict
-        ],
+        model_endpoint: Union[mlrun.model_monitoring.ModelEndpoint, dict],
     ):
         pass
 
@@ -545,7 +560,7 @@ class RunDBInterface(ABC):
         end: Optional[str] = None,
         metrics: Optional[list[str]] = None,
         features: bool = False,
-    ):
+    ) -> mlrun.model_monitoring.ModelEndpoint:
         pass
 
     @abstractmethod
@@ -616,6 +631,86 @@ class RunDBInterface(ABC):
     ):
         pass
 
+    @abstractmethod
+    def store_api_gateway(
+        self,
+        api_gateway: mlrun.common.schemas.APIGateway,
+        project: str = None,
+    ):
+        pass
+
+    @abstractmethod
+    def list_api_gateways(self, project=None) -> mlrun.common.schemas.APIGatewaysOutput:
+        pass
+
+    @abstractmethod
+    def get_api_gateway(self, name, project=None) -> mlrun.common.schemas.APIGateway:
+        pass
+
+    @abstractmethod
+    def delete_api_gateway(self, name, project=None):
+        pass
+
+    @abstractmethod
+    def remote_builder(
+        self,
+        func: "mlrun.runtimes.BaseRuntime",
+        with_mlrun: bool,
+        mlrun_version_specifier: Optional[str] = None,
+        skip_deployed: bool = False,
+        builder_env: Optional[dict] = None,
+        force_build: bool = False,
+    ):
+        pass
+
+    @abstractmethod
+    def deploy_nuclio_function(
+        self,
+        func: "mlrun.runtimes.RemoteRuntime",
+        builder_env: Optional[dict] = None,
+    ):
+        pass
+
+    @abstractmethod
+    def generate_event(
+        self, name: str, event_data: Union[dict, mlrun.common.schemas.Event], project=""
+    ):
+        pass
+
+    @abstractmethod
+    def store_alert_config(
+        self,
+        alert_name: str,
+        alert_data: Union[dict, mlrun.alerts.alert.AlertConfig],
+        project="",
+    ):
+        pass
+
+    @abstractmethod
+    def get_alert_config(self, alert_name: str, project=""):
+        pass
+
+    @abstractmethod
+    def list_alerts_configs(self, project=""):
+        pass
+
+    @abstractmethod
+    def delete_alert_config(self, alert_name: str, project=""):
+        pass
+
+    @abstractmethod
+    def reset_alert_config(self, alert_name: str, project=""):
+        pass
+
+    @abstractmethod
+    def get_alert_template(self, template_name: str):
+        pass
+
+    @abstractmethod
+    def list_alert_templates(self):
+        pass
+
+    @abstractmethod
     def get_builder_status(
         self,
         func: "mlrun.runtimes.BaseRuntime",
@@ -626,6 +721,16 @@ class RunDBInterface(ABC):
     ):
         pass
 
+    @abstractmethod
+    def get_nuclio_deploy_status(
+        self,
+        func: "mlrun.runtimes.RemoteRuntime",
+        last_log_timestamp: float = 0.0,
+        verbose: bool = False,
+    ):
+        pass
+
+    @abstractmethod
     def set_run_notifications(
         self,
         project: str,
@@ -634,6 +739,7 @@ class RunDBInterface(ABC):
     ):
         pass
 
+    @abstractmethod
     def store_run_notifications(
         self,
         notification_objects: list[mlrun.model.Notification],
@@ -643,40 +749,60 @@ class RunDBInterface(ABC):
     ):
         pass
 
+    @abstractmethod
     def get_log_size(self, uid, project=""):
         pass
 
+    @abstractmethod
+    def store_alert_notifications(
+        self,
+        session,
+        notification_objects: list[mlrun.model.Notification],
+        alert_id: str,
+        project: str,
+        mask_params: bool = True,
+    ):
+        pass
+
+    @abstractmethod
     def watch_log(self, uid, project="", watch=True, offset=0):
         pass
 
+    @abstractmethod
     def get_datastore_profile(
         self, name: str, project: str
     ) -> Optional[mlrun.common.schemas.DatastoreProfile]:
         pass
 
+    @abstractmethod
     def delete_datastore_profile(
         self, name: str, project: str
     ) -> mlrun.common.schemas.DatastoreProfile:
         pass
 
+    @abstractmethod
     def list_datastore_profiles(
         self, project: str
     ) -> list[mlrun.common.schemas.DatastoreProfile]:
         pass
 
+    @abstractmethod
     def store_datastore_profile(
         self, profile: mlrun.common.schemas.DatastoreProfile, project: str
     ):
         pass
 
+    @abstractmethod
     def function_status(self, project, name, kind, selector):
         pass
 
+    @abstractmethod
     def start_function(
         self, func_url: str = None, function: "mlrun.runtimes.BaseRuntime" = None
     ):
         pass
 
+    @abstractmethod
     def submit_workflow(
         self,
         project: str,
@@ -695,18 +821,27 @@ class RunDBInterface(ABC):
     ) -> "mlrun.common.schemas.WorkflowResponse":
         pass
 
+    @abstractmethod
     def update_model_monitoring_controller(
         self,
         project: str,
         base_period: int = 10,
         image: str = "mlrun/mlrun",
-    ):
+    ) -> None:
         pass
 
+    @abstractmethod
     def enable_model_monitoring(
         self,
         project: str,
         base_period: int = 10,
         image: str = "mlrun/mlrun",
-    ):
+        deploy_histogram_data_drift_app: bool = True,
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def deploy_histogram_data_drift_app(
+        self, project: str, image: str = "mlrun/mlrun"
+    ) -> None:
         pass

mlrun/db/factory.py

@@ -1,4 +1,4 @@
-# Copyright 2023 MLRun Authors
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.