mlrun 1.6.0rc35__py3-none-any.whl → 1.7.0rc2__py3-none-any.whl

mlrun/utils/helpers.py

@@ -22,14 +22,13 @@ import os
 import re
 import string
 import sys
-import time
 import typing
 import warnings
 from datetime import datetime, timezone
 from importlib import import_module
 from os import path
 from types import ModuleType
-from typing import Any, List, Optional, Tuple
+from typing import Any, Optional
 
 import anyio
 import git
@@ -51,9 +50,15 @@ import mlrun.errors
 import mlrun.utils.regex
 import mlrun.utils.version.version
 from mlrun.config import config
-from mlrun.errors import err_to_str
 
 from .logger import create_logger
+from .retryer import (  # noqa: F401
+    AsyncRetryer,
+    Retryer,
+    create_exponential_backoff,
+    create_linear_backoff,
+    create_step_backoff,
+)
 
 yaml.Dumper.ignore_aliases = lambda *args: True
 _missing = object()
@@ -276,12 +281,12 @@ def validate_v3io_stream_consumer_group(
     )
 
 
-def get_regex_list_as_string(regex_list: List) -> str:
+def get_regex_list_as_string(regex_list: list) -> str:
     """
     This function is used to combine a list of regex strings into a single regex,
     with and condition between them.
     """
-    return "".join(["(?={regex})".format(regex=regex) for regex in regex_list]) + ".*$"
+    return "".join([f"(?={regex})" for regex in regex_list]) + ".*$"
 
 
 def tag_name_regex_as_string() -> str:
@@ -698,7 +703,7 @@ def generate_artifact_uri(project, key, tag=None, iter=None, tree=None):
     return artifact_uri
 
 
-def extend_hub_uri_if_needed(uri) -> Tuple[str, bool]:
+def extend_hub_uri_if_needed(uri) -> tuple[str, bool]:
     """
     Retrieve the full uri of the item's yaml in the hub.
 
@@ -787,7 +792,7 @@ def gen_html_table(header, rows=None):
 def new_pipe_metadata(
     artifact_path: str = None,
     cleanup_ttl: int = None,
-    op_transformers: typing.List[typing.Callable] = None,
+    op_transformers: list[typing.Callable] = None,
 ):
     from kfp.dsl import PipelineConf
 
@@ -893,7 +898,7 @@ def get_docker_repository_or_default(repository: str) -> str:
     return repository
 
 
-def get_parsed_docker_registry() -> Tuple[Optional[str], Optional[str]]:
+def get_parsed_docker_registry() -> tuple[Optional[str], Optional[str]]:
     # according to https://stackoverflow.com/questions/37861791/how-are-docker-image-names-parsed
     docker_registry = config.httpdb.builder.docker_registry or ""
     first_slash_index = docker_registry.find("/")
@@ -947,65 +952,27 @@ def fill_function_hash(function_dict, tag=""):
     return fill_object_hash(function_dict, "hash", tag)
 
 
-def create_linear_backoff(base=2, coefficient=2, stop_value=120):
-    """
-    Create a generator of linear backoff. Check out usage example in test_helpers.py
-    """
-    x = 0
-    comparison = min if coefficient >= 0 else max
-
-    while True:
-        next_value = comparison(base + x * coefficient, stop_value)
-        yield next_value
-        x += 1
-
-
-def create_step_backoff(steps=None):
-    """
-    Create a generator of steps backoff.
-    Example: steps = [[2, 5], [20, 10], [120, None]] will produce a generator in which the first 5
-    values will be 2, the next 10 values will be 20 and the rest will be 120.
-    :param steps: a list of lists [step_value, number_of_iteration_in_this_step]
-    """
-    steps = steps if steps is not None else [[2, 10], [10, 10], [120, None]]
-    steps = iter(steps)
-
-    # Get first step
-    step = next(steps)
-    while True:
-        current_step_value, current_step_remain = step
-        if current_step_remain == 0:
-            # No more in this step, moving on
-            step = next(steps)
-        elif current_step_remain is None:
-            # We are in the last step, staying here forever
-            yield current_step_value
-        elif current_step_remain > 0:
-            # Still more remains in this step, just reduce the remaining number
-            step[1] -= 1
-            yield current_step_value
-
-
-def create_exponential_backoff(base=2, max_value=120, scale_factor=1):
+def retry_until_successful(
+    backoff: int, timeout: int, logger, verbose: bool, _function, *args, **kwargs
+):
     """
-    Create a generator of exponential backoff. Check out usage example in test_helpers.py
-    :param base: exponent base
-    :param max_value: max limit on the result
-    :param scale_factor: factor to be used as linear scaling coefficient
+    Runs function with given *args and **kwargs.
+    Tries to run it until success or timeout reached (timeout is optional)
+    :param backoff: can either be a:
+            - number (int / float) that will be used as interval.
+            - generator of waiting intervals. (support next())
+    :param timeout: pass None if timeout is not wanted, number of seconds if it is
+    :param logger: a logger so we can log the failures
+    :param verbose: whether to log the failure on each retry
+    :param _function: function to run
+    :param args: functions args
+    :param kwargs: functions kwargs
+    :return: function result
     """
-    exponent = 1
-    while True:
-        # This "complex" implementation (unlike the one in linear backoff) is to avoid exponent growing too fast and
-        # risking going behind max_int
-        next_value = scale_factor * (base**exponent)
-        if next_value < max_value:
-            exponent += 1
-            yield next_value
-        else:
-            yield max_value
+    return Retryer(backoff, timeout, logger, verbose, _function, *args, **kwargs).run()
 
 
-def retry_until_successful(
+async def retry_until_successful_async(
     backoff: int, timeout: int, logger, verbose: bool, _function, *args, **kwargs
 ):
     """
@@ -1022,64 +989,15 @@ def retry_until_successful(
     :param kwargs: functions kwargs
     :return: function result
     """
-    start_time = time.time()
-    last_exception = None
-
-    # Check if backoff is just a simple interval
-    if isinstance(backoff, int) or isinstance(backoff, float):
-        backoff = create_linear_backoff(base=backoff, coefficient=0)
-
-    first_interval = next(backoff)
-    if timeout and timeout <= first_interval:
-        logger.warning(
-            f"Timeout ({timeout}) must be higher than backoff ({first_interval})."
-            f" Set timeout to be higher than backoff."
-        )
-
-    # If deadline was not provided or deadline not reached
-    while timeout is None or time.time() < start_time + timeout:
-        next_interval = first_interval or next(backoff)
-        first_interval = None
-        try:
-            result = _function(*args, **kwargs)
-            return result
-
-        except mlrun.errors.MLRunFatalFailureError as exc:
-            raise exc.original_exception
-        except Exception as exc:
-            last_exception = exc
-
-            # If next interval is within allowed time period - wait on interval, abort otherwise
-            if timeout is None or time.time() + next_interval < start_time + timeout:
-                if logger is not None and verbose:
-                    logger.debug(
-                        f"Operation not yet successful, Retrying in {next_interval} seconds."
-                        f" exc: {err_to_str(exc)}"
-                    )
-
-                time.sleep(next_interval)
-            else:
-                break
-
-    if logger is not None:
-        logger.warning(
-            f"Operation did not complete on time. last exception: {last_exception}"
-        )
-
-    raise mlrun.errors.MLRunRetryExhaustedError(
-        f"Failed to execute command by the given deadline."
-        f" last_exception: {last_exception},"
-        f" function_name: {_function.__name__},"
-        f" timeout: {timeout}"
-    ) from last_exception
+    return await AsyncRetryer(
+        backoff, timeout, logger, verbose, _function, *args, **kwargs
+    ).run()
 
 
 def get_ui_url(project, uid=None):
     url = ""
     if mlrun.mlconf.resolve_ui_url():
-        url = "{}/{}/{}/jobs".format(
-            mlrun.mlconf.resolve_ui_url(), mlrun.mlconf.ui.projects_prefix, project
-        )
+        url = f"{mlrun.mlconf.resolve_ui_url()}/{mlrun.mlconf.ui.projects_prefix}/{project}/jobs"
         if uid:
             url += f"/monitor/{uid}/overview"
     return url
@@ -1095,7 +1013,7 @@ def get_workflow_url(project, id=None):
 
 
 def are_strings_in_exception_chain_messages(
-    exception: Exception, strings_list=typing.List[str]
+    exception: Exception, strings_list=list[str]
 ) -> bool:
     while exception is not None:
         if any([string in str(exception) for string in strings_list]):
@@ -1275,7 +1193,7 @@ def has_timezone(timestamp):
         return False
 
 
-def as_list(element: Any) -> List[Any]:
+def as_list(element: Any) -> list[Any]:
     return element if isinstance(element, list) else [element]
 
 
@@ -1618,7 +1536,7 @@ def is_ecr_url(registry: str) -> bool:
     return ".ecr." in registry and ".amazonaws.com" in registry
 
 
-def get_local_file_schema() -> List:
+def get_local_file_schema() -> list:
     # The expression `list(string.ascii_lowercase)` generates a list of lowercase alphabets,
     # which corresponds to drive letters in Windows file paths such as `C:/Windows/path`.
     return ["file"] + list(string.ascii_lowercase)

mlrun/utils/http.py

@@ -14,7 +14,6 @@
 #
 
 import time
-import typing
 
 import requests
 import requests.adapters
@@ -202,9 +201,7 @@ class HTTPSessionWithRetry(requests.Session):
     def _method_retryable(self, method: str):
         return method in self._retry_methods
 
-    def _resolve_retry_methods(
-        self, retry_on_post: bool = False
-    ) -> typing.FrozenSet[str]:
+    def _resolve_retry_methods(self, retry_on_post: bool = False) -> frozenset[str]:
         methods = urllib3.util.retry.Retry.DEFAULT_ALLOWED_METHODS
         methods = methods.union({"PATCH"})
         if retry_on_post:

mlrun/utils/logger.py

@@ -12,22 +12,48 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
 import logging
 from enum import Enum
 from sys import stdout
 from traceback import format_exception
 from typing import IO, Optional, Union
 
-from mlrun.config import config
+import orjson
+import pydantic
 
+from mlrun.config import config
 
-class JSONFormatter(logging.Formatter):
-    def __init__(self):
-        super(JSONFormatter, self).__init__()
-        self._json_encoder = json.JSONEncoder()
 
-    def format(self, record):
+class _BaseFormatter(logging.Formatter):
+    def _json_dump(self, json_object):
+        def default(obj):
+            if isinstance(obj, pydantic.BaseModel):
+                return obj.dict()
+
+            # EAFP all the way.
+            # Leave the unused "exc" in for debugging ease
+            try:
+                return obj.__log__()
+            except Exception as exc:  # noqa
+                try:
+                    return obj.__repr__()
+                except Exception as exc:  # noqa
+                    try:
+                        return str(obj)
+                    except Exception as exc:
+                        raise TypeError from exc
+
+        return orjson.dumps(
+            json_object,
+            option=orjson.OPT_NAIVE_UTC
+            | orjson.OPT_SERIALIZE_NUMPY
+            | orjson.OPT_SORT_KEYS,
+            default=default,
+        ).decode()
+
+
+class JSONFormatter(_BaseFormatter):
+    def format(self, record) -> str:
         record_with = getattr(record, "with", {})
         if record.exc_info:
             record_with.update(exc_info=format_exception(*record.exc_info))
@@ -38,14 +64,24 @@ class JSONFormatter(logging.Formatter):
             "with": record_with,
         }
 
-        return self._json_encoder.encode(record_fields)
+        return self._json_dump(record_fields)
 
 
-class HumanReadableFormatter(logging.Formatter):
+class HumanReadableFormatter(_BaseFormatter):
     def format(self, record) -> str:
+        more = self._resolve_more(record)
+        return (
+            f"> {self.formatTime(record, self.datefmt)} "
+            f"[{record.levelname.lower()}] "
+            f"{record.getMessage().rstrip()}"
+            f"{more}"
+        )
+
+    def _resolve_more(self, record):
         record_with = self._record_with(record)
-        more = f": {record_with}" if record_with else ""
-        return f"> {self.formatTime(record, self.datefmt)} [{record.levelname.lower()}] {record.getMessage()}{more}"
+        record_with_encoded = self._json_dump(record_with) if record_with else ""
+        more = f": {record_with_encoded}" if record_with_encoded else ""
+        return more
 
     def _record_with(self, record):
         record_with = getattr(record, "with", {})
@@ -56,8 +92,7 @@ class HumanReadableFormatter(logging.Formatter):
 
 class HumanReadableExtendedFormatter(HumanReadableFormatter):
     def format(self, record) -> str:
-        record_with = self._record_with(record)
-        more = f": {record_with}" if record_with else ""
+        more = self._resolve_more(record)
         return (
             "> "
             f"{self.formatTime(record, self.datefmt)} "
@@ -66,7 +101,7 @@ class HumanReadableExtendedFormatter(HumanReadableFormatter):
         )
 
 
-class Logger(object):
+class Logger:
     def __init__(
         self,
         level,

mlrun/utils/notifications/notification/__init__.py

@@ -32,7 +32,7 @@ class NotificationTypes(str, enum.Enum):
     slack = NotificationKind.slack.value
     webhook = NotificationKind.webhook.value
 
-    def get_notification(self) -> typing.Type[NotificationBase]:
+    def get_notification(self) -> type[NotificationBase]:
         return {
             self.console: ConsoleNotification,
             self.git: GitNotification,
@@ -41,7 +41,7 @@ class NotificationTypes(str, enum.Enum):
             self.webhook: WebhookNotification,
         }.get(self)
 
-    def inverse_dependencies(self) -> typing.List[str]:
+    def inverse_dependencies(self) -> list[str]:
         """
         Some notifications should only run if another notification type didn't run.
         Per given notification type, return a list of notification types that should not run in order for this
@@ -52,7 +52,7 @@ class NotificationTypes(str, enum.Enum):
         }.get(self, [])
 
     @classmethod
-    def all(cls) -> typing.List[str]:
+    def all(cls) -> list[str]:
         return list(
             [
                 cls.console,

mlrun/utils/notifications/notification/base.py

@@ -23,7 +23,7 @@ class NotificationBase:
     def __init__(
         self,
         name: str = None,
-        params: typing.Dict[str, str] = None,
+        params: dict[str, str] = None,
     ):
         self.name = name
         self.params = params or {}
@@ -49,7 +49,7 @@ class NotificationBase:
 
     def load_notification(
         self,
-        params: typing.Dict[str, str],
+        params: dict[str, str],
     ) -> None:
         self.params = params or {}
 

mlrun/utils/notifications/notification/ipython.py

@@ -29,7 +29,7 @@ class IPythonNotification(NotificationBase):
     def __init__(
         self,
         name: str = None,
-        params: typing.Dict[str, str] = None,
+        params: dict[str, str] = None,
     ):
         super().__init__(name, params)
         self._ipython = None

mlrun/utils/notifications/notification_pusher.py

@@ -32,7 +32,7 @@ from mlrun.utils.condition_evaluator import evaluate_condition_in_separate_proce
 from .notification import NotificationBase, NotificationTypes
 
 
-class _NotificationPusherBase(object):
+class _NotificationPusherBase:
     def _push(
         self, sync_push_callback: typing.Callable, async_push_callback: typing.Callable
     ):
@@ -95,15 +95,11 @@ class NotificationPusher(_NotificationPusherBase):
 
     def __init__(self, runs: typing.Union[mlrun.lists.RunList, list]):
         self._runs = runs
-        self._sync_notifications: typing.List[
-            typing.Tuple[
-                NotificationBase, mlrun.model.RunObject, mlrun.model.Notification
-            ]
+        self._sync_notifications: list[
+            tuple[NotificationBase, mlrun.model.RunObject, mlrun.model.Notification]
         ] = []
-        self._async_notifications: typing.List[
-            typing.Tuple[
-                NotificationBase, mlrun.model.RunObject, mlrun.model.Notification
-            ]
+        self._async_notifications: list[
+            tuple[NotificationBase, mlrun.model.RunObject, mlrun.model.Notification]
         ] = []
 
         for run in self._runs:
@@ -401,7 +397,7 @@ class NotificationPusher(_NotificationPusherBase):
 
 
 class CustomNotificationPusher(_NotificationPusherBase):
-    def __init__(self, notification_types: typing.List[str] = None):
+    def __init__(self, notification_types: list[str] = None):
         notifications = {
             notification_type: NotificationTypes(notification_type).get_notification()()
             for notification_type in notification_types
@@ -446,7 +442,7 @@ class CustomNotificationPusher(_NotificationPusherBase):
     def add_notification(
         self,
         notification_type: str,
-        params: typing.Dict[str, str] = None,
+        params: dict[str, str] = None,
     ):
         if notification_type in self._async_notifications:
             self._async_notifications[notification_type].load_notification(params)
@@ -471,9 +467,7 @@ class CustomNotificationPusher(_NotificationPusherBase):
         else:
             logger.warning(f"No notification of type {notification_type} in project")
 
-    def edit_notification(
-        self, notification_type: str, params: typing.Dict[str, str] = None
-    ):
+    def edit_notification(self, notification_type: str, params: dict[str, str] = None):
         self.remove_notification(notification_type)
         self.add_notification(notification_type, params)