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

mlrun/common/schemas/alert.py

@@ -49,7 +49,7 @@ class Event(pydantic.BaseModel):
     kind: EventKind
     timestamp: Union[str, datetime] = None  # occurrence time
     entity: EventEntity
-    value: Optional[Union[float, str]] = None
+    value_dict: Optional[dict] = pydantic.Field(default_factory=dict)
 
     def is_valid(self):
         return self.entity.kind in _event_kind_entity_map[self.kind]

mlrun/common/schemas/project.py

@@ -110,6 +110,7 @@ class ProjectSummary(pydantic.BaseModel):
     files_count: int
     feature_sets_count: int
     models_count: int
+    runs_completed_recent_count: int
     runs_failed_recent_count: int
     runs_running_count: int
     schedules_count: int

mlrun/config.py

@@ -188,6 +188,7 @@ default_config = {
     "background_tasks": {
         # enabled / disabled
         "timeout_mode": "enabled",
+        "function_deletion_batch_size": 10,
         # timeout in seconds to wait for background task to be updated / finished by the worker responsible for the task
         "default_timeouts": {
             "operations": {
@@ -196,6 +197,7 @@ default_config = {
                 "run_abortion": "600",
                 "abort_grace_period": "10",
                 "delete_project": "900",
+                "delete_function": "900",
             },
             "runtimes": {"dask": "600"},
         },
@@ -552,6 +554,7 @@ default_config = {
             "nosql": "v3io:///projects/{project}/FeatureStore/{name}/nosql",
             # "authority" is optional and generalizes [userinfo "@"] host [":" port]
             "redisnosql": "redis://{authority}/projects/{project}/FeatureStore/{name}/nosql",
+            "dsnosql": "ds://{ds_profile_name}/projects/{project}/FeatureStore/{name}/nosql",
         },
         "default_targets": "parquet,nosql",
         "default_job_image": "mlrun/mlrun",
@@ -691,6 +694,10 @@ default_config = {
         # supported modes: "enabled", "disabled".
         "mode": "disabled"
     },
+    "auth_with_client_id": {
+        "enabled": False,
+        "request_timeout": 5,
+    },
 }
 
 _is_running_as_api = None
@@ -1395,7 +1402,11 @@ def read_env(env=None, prefix=env_prefix):
         log_formatter = mlrun.utils.create_formatter_instance(
             mlrun.utils.FormatterKinds(log_formatter_name)
         )
-        mlrun.utils.logger.get_handler("default").setFormatter(log_formatter)
+        current_handler = mlrun.utils.logger.get_handler("default")
+        current_formatter_name = current_handler.formatter.__class__.__name__
+        desired_formatter_name = log_formatter.__class__.__name__
+        if current_formatter_name != desired_formatter_name:
+            current_handler.setFormatter(log_formatter)
 
     # The default function pod resource values are of type str; however, when reading from environment variable numbers,
     # it converts them to type int if contains only number, so we want to convert them to str.

mlrun/datastore/datastore_profile.py

@@ -83,14 +83,28 @@ class DatastoreProfileBasic(DatastoreProfile):
 class DatastoreProfileKafkaTarget(DatastoreProfile):
     type: str = pydantic.Field("kafka_target")
     _private_attributes = "kwargs_private"
-    bootstrap_servers: str
-    brokers: str
+    bootstrap_servers: typing.Optional[str] = None
+    brokers: typing.Optional[str] = None
     topic: str
     kwargs_public: typing.Optional[dict]
     kwargs_private: typing.Optional[dict]
 
-    def __pydantic_post_init__(self):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        if not self.brokers and not self.bootstrap_servers:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "DatastoreProfileKafkaTarget requires the 'brokers' field to be set"
+            )
+
         if self.bootstrap_servers:
+            if self.brokers:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "DatastoreProfileKafkaTarget cannot be created with both 'brokers' and 'bootstrap_servers'"
+                )
+            else:
+                self.brokers = self.bootstrap_servers
+                self.bootstrap_servers = None
             warnings.warn(
                 "'bootstrap_servers' parameter is deprecated in 1.7.0 and will be removed in 1.9.0, "
                 "use 'brokers' instead.",

mlrun/datastore/hdfs.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+from urllib.parse import urlparse
 
 import fsspec
 
@@ -49,3 +50,7 @@ class HdfsStore(DataStore):
     @property
     def spark_url(self):
         return f"hdfs://{self.host}:{self.port}"
+
+    def rm(self, url, recursive=False, maxdepth=None):
+        path = urlparse(url).path
+        self.filesystem.rm(path=path, recursive=recursive, maxdepth=maxdepth)

mlrun/datastore/targets.py

@@ -1390,6 +1390,39 @@ class RedisNoSqlTarget(NoSqlBaseTarget):
     support_spark = True
     writer_step_name = "RedisNoSqlTarget"
 
+    @property
+    def _target_path_object(self):
+        url = self.path or mlrun.mlconf.redis.url
+        if self._resource and url:
+            parsed_url = urlparse(url)
+            if not parsed_url.path or parsed_url.path == "/":
+                kind_prefix = (
+                    "sets"
+                    if self._resource.kind
+                    == mlrun.common.schemas.ObjectKind.feature_set
+                    else "vectors"
+                )
+                kind = self.kind
+                name = self._resource.metadata.name
+                project = (
+                    self._resource.metadata.project or mlrun.mlconf.default_project
+                )
+                data_prefix = get_default_prefix_for_target(kind).format(
+                    ds_profile_name=parsed_url.netloc,
+                    authority=parsed_url.netloc,
+                    project=project,
+                    kind=kind,
+                    name=name,
+                )
+                if url.startswith("rediss://"):
+                    data_prefix = data_prefix.replace("redis://", "rediss://", 1)
+                if not self.run_id:
+                    version = self._resource.metadata.tag or "latest"
+                    name = f"{name}-{version}"
+                url = f"{data_prefix}/{kind_prefix}/{name}"
+                return TargetPathObject(url, self.run_id, False)
+        return super()._target_path_object
+
     # Fetch server url from the RedisNoSqlTarget::__init__() 'path' parameter.
     # If not set fetch it from 'mlrun.mlconf.redis.url' (MLRUN_REDIS__URL environment variable).
     # Then look for username and password at REDIS_xxx secrets
@@ -1516,6 +1549,25 @@ class StreamTarget(BaseStoreTarget):
 
 
 class KafkaTarget(BaseStoreTarget):
+    """
+    Kafka target storage driver, used to write data into kafka topics.
+    example::
+        # define target
+        kafka_target = KafkaTarget(
+            name="kafka", path="my_topic", brokers="localhost:9092"
+        )
+        # ingest
+        stocks_set.ingest(stocks, [kafka_target])
+    :param name:                target name
+    :param path:                topic name e.g. "my_topic"
+    :param after_step:          optional, after what step in the graph to add the target
+    :param columns:             optional, which columns from data to write
+    :param bootstrap_servers:   Deprecated. Use the brokers parameter instead
+    :param producer_options:    additional configurations for kafka producer
+    :param brokers:             kafka broker as represented by a host:port pair, or a list of kafka brokers, e.g.
+        "localhost:9092", or ["kafka-broker-1:9092", "kafka-broker-2:9092"]
+    """
+
     kind = TargetTypes.kafka
     is_table = False
     is_online = False

mlrun/datastore/v3io.py

@@ -29,7 +29,7 @@ from .base import (
 )
 
 V3IO_LOCAL_ROOT = "v3io"
-V3IO_DEFAULT_UPLOAD_CHUNK_SIZE = 1024 * 1024 * 100
+V3IO_DEFAULT_UPLOAD_CHUNK_SIZE = 1024 * 1024 * 10
 
 
 class V3ioStore(DataStore):

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

@@ -802,7 +802,7 @@ class RunDBInterface(ABC):
         project: str,
         base_period: int = 10,
         image: str = "mlrun/mlrun",
-    ):
+    ) -> None:
         pass
 
     @abstractmethod

mlrun/db/httpdb.py

@@ -38,6 +38,7 @@ import mlrun.platforms
 import mlrun.projects
 import mlrun.runtimes.nuclio.api_gateway
 import mlrun.utils
+from mlrun.db.auth_utils import OAuthClientIDTokenProvider, StaticTokenProvider
 from mlrun.errors import MLRunInvalidArgumentError, err_to_str
 
 from ..artifacts import Artifact
@@ -138,17 +139,28 @@ class HTTPRunDB(RunDBInterface):
             endpoint += f":{parsed_url.port}"
         base_url = f"{parsed_url.scheme}://{endpoint}{parsed_url.path}"
 
+        self.base_url = base_url
         username = parsed_url.username or config.httpdb.user
         password = parsed_url.password or config.httpdb.password
+        self.token_provider = None
 
-        username, password, token = mlrun.platforms.add_or_refresh_credentials(
-            parsed_url.hostname, username, password, config.httpdb.token
-        )
+        if config.auth_with_client_id.enabled:
+            self.token_provider = OAuthClientIDTokenProvider(
+                token_endpoint=mlrun.get_secret_or_env("MLRUN_AUTH_TOKEN_ENDPOINT"),
+                client_id=mlrun.get_secret_or_env("MLRUN_AUTH_CLIENT_ID"),
+                client_secret=mlrun.get_secret_or_env("MLRUN_AUTH_CLIENT_SECRET"),
+                timeout=config.auth_with_client_id.request_timeout,
+            )
+        else:
+            username, password, token = mlrun.platforms.add_or_refresh_credentials(
+                parsed_url.hostname, username, password, config.httpdb.token
+            )
+
+            if token:
+                self.token_provider = StaticTokenProvider(token)
 
-        self.base_url = base_url
         self.user = username
         self.password = password
-        self.token = token
 
     def __repr__(self):
         cls = self.__class__.__name__
@@ -218,17 +230,19 @@ class HTTPRunDB(RunDBInterface):
 
         if self.user:
             kw["auth"] = (self.user, self.password)
-        elif self.token:
-            # Iguazio auth doesn't support passing token through bearer, so use cookie instead
-            if mlrun.platforms.iguazio.is_iguazio_session(self.token):
-                session_cookie = f'j:{{"sid": "{self.token}"}}'
-                cookies = {
-                    "session": session_cookie,
-                }
-                kw["cookies"] = cookies
-            else:
-                if "Authorization" not in kw.setdefault("headers", {}):
-                    kw["headers"].update({"Authorization": "Bearer " + self.token})
+        elif self.token_provider:
+            token = self.token_provider.get_token()
+            if token:
+                # Iguazio auth doesn't support passing token through bearer, so use cookie instead
+                if self.token_provider.is_iguazio_session():
+                    session_cookie = f'j:{{"sid": "{token}"}}'
+                    cookies = {
+                        "session": session_cookie,
+                    }
+                    kw["cookies"] = cookies
+                else:
+                    if "Authorization" not in kw.setdefault("headers", {}):
+                        kw["headers"].update({"Authorization": "Bearer " + token})
 
         if mlrun.common.schemas.HeaderNames.client_version not in kw.setdefault(
             "headers", {}
@@ -1142,7 +1156,29 @@ class HTTPRunDB(RunDBInterface):
         project = project or config.default_project
         path = f"projects/{project}/functions/{name}"
         error_message = f"Failed deleting function {project}/{name}"
-        self.api_call("DELETE", path, error_message)
+        response = self.api_call("DELETE", path, error_message, version="v2")
+        if response.status_code == http.HTTPStatus.ACCEPTED:
+            logger.info(
+                "Function is being deleted", project_name=project, function_name=name
+            )
+            background_task = mlrun.common.schemas.BackgroundTask(**response.json())
+            background_task = self._wait_for_background_task_to_reach_terminal_state(
+                background_task.metadata.name, project=project
+            )
+            if (
+                background_task.status.state
+                == mlrun.common.schemas.BackgroundTaskState.succeeded
+            ):
+                logger.info(
+                    "Function deleted", project_name=project, function_name=name
+                )
+            elif (
+                background_task.status.state
+                == mlrun.common.schemas.BackgroundTaskState.failed
+            ):
+                logger.info(
+                    "Function deletion failed", project_name=project, function_name=name
+                )
 
     def list_functions(self, name=None, project=None, tag=None, labels=None):
         """Retrieve a list of functions, filtered by specific criteria.
@@ -1488,16 +1524,15 @@ class HTTPRunDB(RunDBInterface):
         """
 
         try:
+            normalized_name = normalize_name(func.metadata.name)
             params = {
-                "name": normalize_name(func.metadata.name),
+                "name": normalized_name,
                 "project": func.metadata.project,
                 "tag": func.metadata.tag,
                 "last_log_timestamp": str(last_log_timestamp),
                 "verbose": bool2str(verbose),
             }
-            _path = (
-                f"projects/{func.metadata.project}/nuclio/{func.metadata.name}/deploy"
-            )
+            _path = f"projects/{func.metadata.project}/nuclio/{normalized_name}/deploy"
             resp = self.api_call("GET", _path, params=params)
         except OSError as err:
             logger.error(f"error getting deploy status: {err_to_str(err)}")
@@ -3214,7 +3249,7 @@ class HTTPRunDB(RunDBInterface):
         project: str,
         base_period: int = 10,
         image: str = "mlrun/mlrun",
-    ):
+    ) -> None:
         """
         Redeploy model monitoring application controller function.
 
@@ -3224,13 +3259,14 @@ class HTTPRunDB(RunDBInterface):
         :param image: The image of the model monitoring controller function.
                                          By default, the image is mlrun/mlrun.
         """
-
-        params = {
-            "image": image,
-            "base_period": base_period,
-        }
-        path = f"projects/{project}/model-monitoring/model-monitoring-controller"
-        self.api_call(method="POST", path=path, params=params)
+        self.api_call(
+            method=mlrun.common.types.HTTPMethod.POST,
+            path=f"projects/{project}/model-monitoring/model-monitoring-controller",
+            params={
+                "base_period": base_period,
+                "image": image,
+            },
+        )
 
     def enable_model_monitoring(
         self,

mlrun/model.py

@@ -766,6 +766,11 @@ class RunMetadata(ModelObj):
     def iteration(self, iteration):
         self._iteration = iteration
 
+    def is_workflow_runner(self):
+        if not self.labels:
+            return False
+        return self.labels.get("job-type", "") == "workflow-runner"
+
 
 class HyperParamStrategies:
     grid = "grid"
@@ -1218,6 +1223,19 @@ class RunStatus(ModelObj):
         self.reason = reason
         self.notifications = notifications or {}
 
+    def is_failed(self) -> Optional[bool]:
+        """
+        This method returns whether a run has failed.
+        Returns none if state has yet to be defined. callee is responsible for handling None.
+        (e.g wait for state to be defined)
+        """
+        if not self.state:
+            return None
+        return self.state.casefold() in [
+            mlrun.run.RunStatuses.failed.casefold(),
+            mlrun.run.RunStatuses.error.casefold(),
+        ]
+
 
 class RunTemplate(ModelObj):
     """Run template"""

mlrun/model_monitoring/helpers.py

@@ -253,3 +253,10 @@ def calculate_inputs_statistics(
             )
 
     return inputs_statistics
+
+
+def get_endpoint_record(project: str, endpoint_id: str):
+    model_endpoint_store = mlrun.model_monitoring.get_store_object(
+        project=project,
+    )
+    return model_endpoint_store.get_model_endpoint(endpoint_id=endpoint_id)

mlrun/model_monitoring/stream_processing.py

@@ -40,6 +40,7 @@ from mlrun.common.schemas.model_monitoring.constants import (
     ProjectSecretKeys,
     PrometheusEndpoints,
 )
+from mlrun.model_monitoring.helpers import get_endpoint_record
 from mlrun.utils import logger
 
 
@@ -1233,13 +1234,6 @@ def update_endpoint_record(
     )
 
 
-def get_endpoint_record(project: str, endpoint_id: str):
-    model_endpoint_store = mlrun.model_monitoring.get_store_object(
-        project=project,
-    )
-    return model_endpoint_store.get_model_endpoint(endpoint_id=endpoint_id)
-
-
 def update_monitoring_feature_set(
     endpoint_record: dict[str, typing.Any],
     feature_names: list[str],

mlrun/model_monitoring/writer.py

@@ -27,8 +27,13 @@ import mlrun.common.schemas.alert as alert_constants
 import mlrun.model_monitoring
 import mlrun.model_monitoring.db.stores
 import mlrun.utils.v3io_clients
-from mlrun.common.schemas.model_monitoring.constants import ResultStatusApp, WriterEvent
+from mlrun.common.schemas.model_monitoring.constants import (
+    EventFieldType,
+    ResultStatusApp,
+    WriterEvent,
+)
 from mlrun.common.schemas.notification import NotificationKind, NotificationSeverity
+from mlrun.model_monitoring.helpers import get_endpoint_record
 from mlrun.serving.utils import StepToDict
 from mlrun.utils import logger
 from mlrun.utils.notifications.notification_pusher import CustomNotificationPusher
@@ -112,6 +117,7 @@ class ModelMonitoringWriter(StepToDict):
             notification_types=[NotificationKind.slack]
         )
         self._create_tsdb_table()
+        self._endpoints_records = {}
 
     @staticmethod
     def get_v3io_container(project_name: str) -> str:
@@ -174,7 +180,7 @@ class ModelMonitoringWriter(StepToDict):
 
     @staticmethod
     def _generate_event_on_drift(
-        uid: str, drift_status: str, drift_value: float, project_name: str
+        uid: str, drift_status: str, event_value: dict, project_name: str
     ):
         if (
             drift_status == ResultStatusApp.detected
@@ -191,7 +197,7 @@ class ModelMonitoringWriter(StepToDict):
                 else alert_constants.EventKind.DRIFT_SUSPECTED
             )
             event_data = mlrun.common.schemas.Event(
-                kind=event_kind, entity=entity, value=drift_value
+                kind=event_kind, entity=entity, value_dict=event_value
             )
             mlrun.get_run_db().generate_event(event_kind, event_data)
 
@@ -227,10 +233,22 @@ class ModelMonitoringWriter(StepToDict):
         _Notifier(event=event, notification_pusher=self._custom_notifier).notify()
 
         if mlrun.mlconf.alerts.mode == mlrun.common.schemas.alert.AlertsModes.enabled:
+            endpoint_id = event[WriterEvent.ENDPOINT_ID]
+            endpoint_record = self._endpoints_records.setdefault(
+                endpoint_id,
+                get_endpoint_record(project=self.project, endpoint_id=endpoint_id),
+            )
+            event_value = {
+                "app_name": event[WriterEvent.APPLICATION_NAME],
+                "model": endpoint_record.get(EventFieldType.MODEL),
+                "model_endpoint_id": event[WriterEvent.ENDPOINT_ID],
+                "result_name": event[WriterEvent.RESULT_NAME],
+                "result_value": event[WriterEvent.RESULT_VALUE],
+            }
             self._generate_event_on_drift(
                 event[WriterEvent.ENDPOINT_ID],
                 event[WriterEvent.RESULT_STATUS],
-                event[WriterEvent.RESULT_VALUE],
+                event_value,
                 self.project,
             )
         logger.info("Completed event DB writes")