mlrun 1.6.4rc8__py3-none-any.whl → 1.7.0__py3-none-any.whl

mlrun/utils/http.py

@@ -14,7 +14,6 @@
 #
 
 import time
-import typing
 
 import requests
 import requests.adapters
@@ -96,7 +95,7 @@ class HTTPSessionWithRetry(requests.Session):
                     total=self.max_retries,
                     backoff_factor=self.retry_backoff_factor,
                     status_forcelist=config.http_retry_defaults.status_codes,
-                    method_whitelist=self._retry_methods,
+                    allowed_methods=self._retry_methods,
                     # we want to retry but not to raise since we do want that last response (to parse details on the
                     # error from response body) we'll handle raising ourselves
                     raise_on_status=False,
@@ -123,7 +122,7 @@ class HTTPSessionWithRetry(requests.Session):
 
                 self._logger.warning(
                     "Error during request handling, retrying",
-                    exc=str(exc),
+                    exc=err_to_str(exc),
                     retry_count=retry_count,
                     url=url,
                     method=method,
@@ -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,23 +12,51 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
 import logging
+import os
 import typing
 from enum import Enum
+from functools import cached_property
 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))
@@ -39,14 +67,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,18 +94,67 @@ class HumanReadableFormatter(logging.Formatter):
 
 
 class HumanReadableExtendedFormatter(HumanReadableFormatter):
+    _colors = {
+        logging.NOTSET: "",
+        logging.DEBUG: "\x1b[34m",
+        logging.INFO: "\x1b[36m",
+        logging.WARNING: "\x1b[33m",
+        logging.ERROR: "\x1b[0;31m",
+        logging.CRITICAL: "\x1b[1;31m",
+    }
+    _color_reset = "\x1b[0m"
+
     def format(self, record) -> str:
+        more = ""
         record_with = self._record_with(record)
-        more = f": {record_with}" if record_with else ""
+        if record_with:
+
+            def _format_value(val):
+                formatted_val = (
+                    val
+                    if isinstance(val, str)
+                    else str(orjson.loads(self._json_dump(val)))
+                )
+                return (
+                    formatted_val.replace("\n", "\n\t\t")
+                    if len(formatted_val) < 4096
+                    else repr(formatted_val)
+                )
+
+            more = "\n\t" + "\n\t".join(
+                [f"{key}: {_format_value(val)}" for key, val in record_with.items()]
+            )
         return (
-            "> "
+            f"{self._get_message_color(record.levelno)}> "
             f"{self.formatTime(record, self.datefmt)} "
             f"[{record.name}:{record.levelname.lower()}] "
-            f"{record.getMessage()}{more}"
+            f"{record.getMessage()}{more}{self._get_color_reset()}"
         )
 
+    def _get_color_reset(self):
+        if not self._have_color_support:
+            return ""
 
-class Logger(object):
+        return self._color_reset
+
+    def _get_message_color(self, levelno):
+        if not self._have_color_support:
+            return ""
+
+        return self._colors[levelno]
+
+    @cached_property
+    def _have_color_support(self):
+        if os.environ.get("PYCHARM_HOSTED"):
+            return True
+        if os.environ.get("NO_COLOR"):
+            return False
+        if os.environ.get("CLICOLOR_FORCE"):
+            return True
+        return stdout.isatty()
+
+
+class Logger:
     def __init__(
         self,
         level,
@@ -189,7 +276,7 @@ class FormatterKinds(Enum):
 
 def resolve_formatter_by_kind(
     formatter_kind: FormatterKinds,
-) -> typing.Type[
+) -> type[
     typing.Union[HumanReadableFormatter, HumanReadableExtendedFormatter, JSONFormatter]
 ]:
     return {
@@ -199,6 +286,15 @@ def resolve_formatter_by_kind(
     }[formatter_kind]
 
 
+def create_test_logger(name: str = "mlrun", stream: IO[str] = stdout) -> Logger:
+    return create_logger(
+        level="debug",
+        formatter_kind=FormatterKinds.HUMAN_EXTENDED.name,
+        name=name,
+        stream=stream,
+    )
+
+
 def create_logger(
     level: Optional[str] = None,
     formatter_kind: str = FormatterKinds.HUMAN.name,

mlrun/utils/notifications/notification/__init__.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import enum
-import typing
 
 from mlrun.common.schemas.notification import NotificationKind
 
@@ -32,7 +31,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 +40,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,13 +51,18 @@ class NotificationTypes(str, enum.Enum):
         }.get(self, [])
 
     @classmethod
-    def all(cls) -> typing.List[str]:
-        return list(
-            [
-                cls.console,
-                cls.git,
-                cls.ipython,
-                cls.slack,
-                cls.webhook,
-            ]
-        )
+    def local(cls) -> list[str]:
+        return [
+            cls.console,
+            cls.ipython,
+        ]
+
+    @classmethod
+    def all(cls) -> list[str]:
+        return [
+            cls.console,
+            cls.git,
+            cls.ipython,
+            cls.slack,
+            cls.webhook,
+        ]

mlrun/utils/notifications/notification/base.py

@@ -23,11 +23,15 @@ 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 {}
 
+    @classmethod
+    def validate_params(cls, params):
+        pass
+
     @property
     def active(self) -> bool:
         return True
@@ -44,12 +48,14 @@ class NotificationBase:
         ] = mlrun.common.schemas.NotificationSeverity.INFO,
         runs: typing.Union[mlrun.lists.RunList, list] = None,
         custom_html: str = None,
+        alert: mlrun.common.schemas.AlertConfig = None,
+        event_data: mlrun.common.schemas.Event = None,
     ):
         raise NotImplementedError()
 
     def load_notification(
         self,
-        params: typing.Dict[str, str],
+        params: dict[str, str],
     ) -> None:
         self.params = params or {}
 
@@ -61,10 +67,31 @@ class NotificationBase:
         ] = mlrun.common.schemas.NotificationSeverity.INFO,
         runs: typing.Union[mlrun.lists.RunList, list] = None,
         custom_html: str = None,
+        alert: mlrun.common.schemas.AlertConfig = None,
+        event_data: mlrun.common.schemas.Event = None,
     ) -> str:
         if custom_html:
             return custom_html
 
+        if alert:
+            if not event_data:
+                return f"[{severity}] {message}"
+
+            html = f"<h3>[{severity}] {message}</h3>"
+            html += f"<br>{alert.name} alert has occurred<br>"
+            html += f"<br><h4>Project:</h4>{alert.project}<br>"
+            html += f"<br><h4>ID:</h4>{event_data.entity.ids[0]}<br>"
+            html += f"<br><h4>Summary:</h4>{mlrun.utils.helpers.format_alert_summary(alert, event_data)}<br>"
+
+            if event_data.value_dict:
+                html += "<br><h4>Event data:</h4>"
+                for key, value in event_data.value_dict.items():
+                    html += f"{key}: {value}<br>"
+
+            overview_type, url = self._get_overview_type_and_url(alert, event_data)
+            html += f"<br><h4>Overview:</h4><a href={url}>{overview_type}</a>"
+            return html
+
         if self.name:
             message = f"{self.name}: {message}"
 
@@ -78,3 +105,24 @@ class NotificationBase:
         html += "<br>click the hyper links below to see detailed results<br>"
         html += runs.show(display=False, short=True)
         return html
+
+    def _get_overview_type_and_url(
+        self,
+        alert: mlrun.common.schemas.AlertConfig,
+        event_data: mlrun.common.schemas.Event,
+    ) -> (str, str):
+        if (
+            event_data.entity.kind == mlrun.common.schemas.alert.EventEntityKind.JOB
+        ):  # JOB entity
+            uid = event_data.value_dict.get("uid")
+            url = mlrun.utils.helpers.get_ui_url(alert.project, uid)
+            overview_type = "Job overview"
+        else:  # MODEL entity
+            model_name = event_data.value_dict.get("model")
+            model_endpoint_id = event_data.value_dict.get("model_endpoint_id")
+            url = mlrun.utils.helpers.get_model_endpoint_url(
+                alert.project, model_name, model_endpoint_id
+            )
+            overview_type = "Model endpoint"
+
+        return overview_type, url

mlrun/utils/notifications/notification/console.py

@@ -36,6 +36,8 @@ class ConsoleNotification(NotificationBase):
         ] = mlrun.common.schemas.NotificationSeverity.INFO,
         runs: typing.Union[mlrun.lists.RunList, list] = None,
         custom_html: str = None,
+        alert: mlrun.common.schemas.AlertConfig = None,
+        event_data: mlrun.common.schemas.Event = None,
     ):
         severity = self._resolve_severity(severity)
         print(f"[{severity}] {message}")

mlrun/utils/notifications/notification/git.py

@@ -30,6 +30,27 @@ class GitNotification(NotificationBase):
     API/Client notification for setting a rich run statuses git issue comment (github/gitlab)
     """
 
+    @classmethod
+    def validate_params(cls, params):
+        git_repo = params.get("repo", None)
+        git_issue = params.get("issue", None)
+        git_merge_request = params.get("merge_request", None)
+        token = (
+            params.get("token", None)
+            or params.get("GIT_TOKEN", None)
+            or params.get("GITHUB_TOKEN", None)
+        )
+        if not git_repo:
+            raise ValueError("Parameter 'repo' is required for GitNotification")
+
+        if not token:
+            raise ValueError("Parameter 'token' is required for GitNotification")
+
+        if not git_issue and not git_merge_request:
+            raise ValueError(
+                "At least one of 'issue' or 'merge_request' is required for GitNotification"
+            )
+
     async def push(
         self,
         message: str,
@@ -38,6 +59,8 @@ class GitNotification(NotificationBase):
         ] = mlrun.common.schemas.NotificationSeverity.INFO,
         runs: typing.Union[mlrun.lists.RunList, list] = None,
         custom_html: str = None,
+        alert: mlrun.common.schemas.AlertConfig = None,
+        event_data: mlrun.common.schemas.Event = None,
     ):
         git_repo = self.params.get("repo", None)
         git_issue = self.params.get("issue", None)
@@ -50,7 +73,7 @@ class GitNotification(NotificationBase):
         server = self.params.get("server", None)
         gitlab = self.params.get("gitlab", False)
         await self._pr_comment(
-            self._get_html(message, severity, runs, custom_html),
+            self._get_html(message, severity, runs, custom_html, alert, event_data),
             git_repo,
             git_issue,
             merge_request=git_merge_request,

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
@@ -53,6 +53,8 @@ class IPythonNotification(NotificationBase):
         ] = mlrun.common.schemas.NotificationSeverity.INFO,
         runs: typing.Union[mlrun.lists.RunList, list] = None,
         custom_html: str = None,
+        alert: mlrun.common.schemas.AlertConfig = None,
+        event_data: mlrun.common.schemas.Event = None,
     ):
         if not self._ipython:
             mlrun.utils.helpers.logger.debug(

mlrun/utils/notifications/notification/slack.py

@@ -32,8 +32,17 @@ class SlackNotification(NotificationBase):
         "completed": ":smiley:",
         "running": ":man-running:",
         "error": ":x:",
+        "skipped": ":zzz:",
     }
 
+    @classmethod
+    def validate_params(cls, params):
+        webhook = params.get("webhook", None) or mlrun.get_secret_or_env(
+            "SLACK_WEBHOOK"
+        )
+        if not webhook:
+            raise ValueError("Parameter 'webhook' is required for SlackNotification")
+
     async def push(
         self,
         message: str,
@@ -42,6 +51,8 @@ class SlackNotification(NotificationBase):
         ] = mlrun.common.schemas.NotificationSeverity.INFO,
         runs: typing.Union[mlrun.lists.RunList, list] = None,
         custom_html: str = None,
+        alert: mlrun.common.schemas.AlertConfig = None,
+        event_data: mlrun.common.schemas.Event = None,
     ):
         webhook = self.params.get("webhook", None) or mlrun.get_secret_or_env(
             "SLACK_WEBHOOK"
@@ -53,7 +64,7 @@ class SlackNotification(NotificationBase):
             )
             return
 
-        data = self._generate_slack_data(message, severity, runs)
+        data = self._generate_slack_data(message, severity, runs, alert, event_data)
 
         async with aiohttp.ClientSession() as session:
             async with session.post(webhook, json=data) as response:
@@ -66,57 +77,121 @@ class SlackNotification(NotificationBase):
             mlrun.common.schemas.NotificationSeverity, str
         ] = mlrun.common.schemas.NotificationSeverity.INFO,
         runs: typing.Union[mlrun.lists.RunList, list] = None,
+        alert: mlrun.common.schemas.AlertConfig = None,
+        event_data: mlrun.common.schemas.Event = None,
     ) -> dict:
         data = {
-            "blocks": [
-                {
-                    "type": "section",
-                    "text": self._get_slack_row(f"[{severity}] {message}"),
-                },
-            ]
+            "blocks": self._generate_slack_header_blocks(severity, message),
         }
         if self.name:
             data["blocks"].append(
                 {"type": "section", "text": self._get_slack_row(self.name)}
             )
 
-        if not runs:
-            return data
+        if alert:
+            fields = self._get_alert_fields(alert, event_data)
 
-        if isinstance(runs, list):
-            runs = mlrun.lists.RunList(runs)
+            for i in range(len(fields)):
+                data["blocks"].append({"type": "section", "text": fields[i]})
+        else:
+            if not runs:
+                return data
+
+            if isinstance(runs, list):
+                runs = mlrun.lists.RunList(runs)
 
-        fields = [self._get_slack_row("*Runs*"), self._get_slack_row("*Results*")]
-        for run in runs:
-            fields.append(self._get_run_line(run))
-            fields.append(self._get_run_result(run))
+            fields = [self._get_slack_row("*Runs*"), self._get_slack_row("*Results*")]
+            for run in runs:
+                fields.append(self._get_run_line(run))
+                fields.append(self._get_run_result(run))
 
-        for i in range(0, len(fields), 8):
-            data["blocks"].append({"type": "section", "fields": fields[i : i + 8]})
+            for i in range(0, len(fields), 8):
+                data["blocks"].append({"type": "section", "fields": fields[i : i + 8]})
 
         return data
 
+    def _generate_slack_header_blocks(self, severity: str, message: str):
+        header_text = block_text = f"[{severity}] {message}"
+        section_text = None
+
+        # Slack doesn't allow headers to be longer than 150 characters
+        # If there's a comma in the message, split the message at the comma
+        # Otherwise, split the message at 150 characters
+        if len(block_text) > 150:
+            if ", " in block_text and block_text.index(", ") < 149:
+                header_text = block_text.split(",")[0]
+                section_text = block_text[len(header_text) + 2 :]
+            else:
+                header_text = block_text[:150]
+                section_text = block_text[150:]
+        blocks = [
+            {"type": "header", "text": {"type": "plain_text", "text": header_text}}
+        ]
+        if section_text:
+            blocks.append(
+                {
+                    "type": "section",
+                    "text": self._get_slack_row(section_text),
+                }
+            )
+        return blocks
+
+    def _get_alert_fields(
+        self,
+        alert: mlrun.common.schemas.AlertConfig,
+        event_data: mlrun.common.schemas.Event,
+    ) -> list:
+        line = [
+            self._get_slack_row(f":bell: {alert.name} alert has occurred"),
+            self._get_slack_row(f"*Project:*\n{alert.project}"),
+            self._get_slack_row(f"*ID:*\n{event_data.entity.ids[0]}"),
+        ]
+
+        if alert.summary:
+            line.append(
+                self._get_slack_row(
+                    f"*Summary:*\n{mlrun.utils.helpers.format_alert_summary(alert, event_data)}"
+                )
+            )
+
+        if event_data.value_dict:
+            data_lines = []
+            for key, value in event_data.value_dict.items():
+                data_lines.append(f"{key}: {value}")
+            data_text = "\n".join(data_lines)
+            line.append(self._get_slack_row(f"*Event data:*\n{data_text}"))
+
+        overview_type, url = self._get_overview_type_and_url(alert, event_data)
+        line.append(self._get_slack_row(f"*Overview:*\n<{url}|*{overview_type}*>"))
+
+        return line
+
     def _get_run_line(self, run: dict) -> dict:
         meta = run["metadata"]
         url = mlrun.utils.helpers.get_ui_url(meta.get("project"), meta.get("uid"))
-        if url:
+
+        # Only show the URL if the run is not a function (serving or mlrun function)
+        kind = run.get("step_kind")
+        state = run["status"].get("state", "")
+        if state != "skipped" and (url and not kind or kind == "run"):
             line = f'<{url}|*{meta.get("name")}*>'
         else:
             line = meta.get("name")
-        state = run["status"].get("state", "")
+        if kind:
+            line = f'{line} *({run.get("step_kind", run.get("kind", ""))})*'
         line = f'{self.emojis.get(state, ":question:")}  {line}'
         return self._get_slack_row(line)
 
     def _get_run_result(self, run: dict) -> dict:
         state = run["status"].get("state", "")
         if state == "error":
-            error_status = run["status"].get("error", "")
+            error_status = run["status"].get("error", "") or state
             result = f"*{error_status}*"
         else:
             result = mlrun.utils.helpers.dict_to_str(
                 run["status"].get("results", {}), ", "
             )
-        return self._get_slack_row(result or "None")
+        return self._get_slack_row(result or state)
 
     @staticmethod
     def _get_slack_row(text: str) -> dict: