mlrun 1.7.2rc3__py3-none-any.whl → 1.8.0__py3-none-any.whl

mlrun/utils/notifications/notification/webhook.py

@@ -16,6 +16,7 @@ import re
 import typing
 
 import aiohttp
+import orjson
 
 import mlrun.common.schemas
 import mlrun.lists
@@ -38,13 +39,13 @@ class WebhookNotification(NotificationBase):
     async def push(
         self,
         message: str,
-        severity: typing.Union[
-            mlrun.common.schemas.NotificationSeverity, str
+        severity: typing.Optional[
+            typing.Union[mlrun.common.schemas.NotificationSeverity, str]
         ] = 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,
+        runs: typing.Optional[typing.Union[mlrun.lists.RunList, list]] = None,
+        custom_html: typing.Optional[typing.Optional[str]] = None,
+        alert: typing.Optional[mlrun.common.schemas.AlertConfig] = None,
+        event_data: typing.Optional[mlrun.common.schemas.Event] = None,
     ):
         url = self.params.get("url", None)
         method = self.params.get("method", "post").lower()
@@ -86,9 +87,14 @@ class WebhookNotification(NotificationBase):
         # we automatically handle it as `ssl=None` for their convenience.
         verify_ssl = verify_ssl and None if url.startswith("https") else None
 
-        async with aiohttp.ClientSession() as session:
+        async with aiohttp.ClientSession(
+            json_serialize=self._encoder,
+        ) as session:
             response = await getattr(session, method)(
-                url, headers=headers, json=request_body, ssl=verify_ssl
+                url,
+                headers=headers,
+                json=request_body,
+                ssl=verify_ssl,
             )
             response.raise_for_status()
 
@@ -128,3 +134,13 @@ class WebhookNotification(NotificationBase):
                     )
 
         return override_body
+
+    @property
+    def _encoder(self):
+        return lambda body: orjson.dumps(
+            body,
+            option=orjson.OPT_NAIVE_UTC
+            | orjson.OPT_SERIALIZE_NUMPY
+            | orjson.OPT_NON_STR_KEYS
+            | orjson.OPT_SORT_KEYS,
+        ).decode()

mlrun/utils/notifications/notification_pusher.py

@@ -15,17 +15,12 @@
 import asyncio
 import datetime
 import os
-import re
 import traceback
 import typing
 from concurrent.futures import ThreadPoolExecutor
 
-import mlrun_pipelines.common.ops
-import mlrun_pipelines.models
-import mlrun_pipelines.utils
-
 import mlrun.common.constants as mlrun_constants
-import mlrun.common.runtimes.constants
+import mlrun.common.runtimes.constants as runtimes_constants
 import mlrun.common.schemas
 import mlrun.config
 import mlrun.db.base
@@ -33,11 +28,11 @@ import mlrun.errors
 import mlrun.lists
 import mlrun.model
 import mlrun.utils.helpers
-from mlrun.utils import logger
+import mlrun.utils.notifications.notification as notification_module
+import mlrun.utils.notifications.notification.base as base
+from mlrun.utils import Workflow, logger
 from mlrun.utils.condition_evaluator import evaluate_condition_in_separate_process
 
-from .notification import NotificationBase, NotificationTypes
-
 
 class _NotificationPusherBase:
     def _push(
@@ -60,6 +55,7 @@ class _NotificationPusherBase:
                 event_loop = asyncio.get_event_loop()
             except RuntimeError:
                 event_loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(event_loop)
 
             if not event_loop.is_running():
                 event_loop.run_until_complete(async_push_callback())
@@ -98,35 +94,71 @@ class NotificationPusher(_NotificationPusherBase):
         "completed": "{resource} completed",
         "error": "{resource} failed",
         "aborted": "{resource} aborted",
+        "running": "{resource} started",
     }
 
-    def __init__(self, runs: typing.Union[mlrun.lists.RunList, list]):
+    def __init__(
+        self,
+        runs: typing.Union[mlrun.lists.RunList, list],
+        default_params: typing.Optional[dict] = None,
+    ):
         self._runs = runs
+        self._default_params = default_params or {}
         self._sync_notifications: list[
-            tuple[NotificationBase, mlrun.model.RunObject, mlrun.model.Notification]
+            tuple[
+                base.NotificationBase, mlrun.model.RunObject, mlrun.model.Notification
+            ]
         ] = []
         self._async_notifications: list[
-            tuple[NotificationBase, mlrun.model.RunObject, mlrun.model.Notification]
+            tuple[
+                base.NotificationBase, mlrun.model.RunObject, mlrun.model.Notification
+            ]
         ] = []
 
         for run in self._runs:
-            if isinstance(run, dict):
-                run = mlrun.model.RunObject.from_dict(run)
+            try:
+                self._process_run(run)
+            except Exception as exc:
+                logger.warning(
+                    "Failed to process run",
+                    run_uid=run.metadata.uid,
+                    error=mlrun.errors.err_to_str(exc),
+                )
 
-            for notification in run.spec.notifications:
-                try:
-                    notification.status = run.status.notifications.get(
-                        notification.name
-                    ).get("status", mlrun.common.schemas.NotificationStatus.PENDING)
-                except (AttributeError, KeyError):
-                    notification.status = (
-                        mlrun.common.schemas.NotificationStatus.PENDING
-                    )
+    def _process_run(self, run):
+        if isinstance(run, dict):
+            run = mlrun.model.RunObject.from_dict(run)
 
-                if self._should_notify(run, notification):
-                    self._load_notification(run, notification)
+        for notification in run.spec.notifications:
+            try:
+                self._process_notification(notification, run)
+            except Exception as exc:
+                logger.warning(
+                    "Failed to process notification",
+                    run_uid=run.metadata.uid,
+                    notification=notification,
+                    error=mlrun.errors.err_to_str(exc),
+                )
 
-    def push(self):
+    def _process_notification(self, notification_object, run):
+        notification_object.status = run.status.notifications.get(
+            notification_object.name, {}
+        ).get(
+            "status",
+            mlrun.common.schemas.NotificationStatus.PENDING,
+        )
+        if self._should_notify(run, notification_object):
+            notification = self._load_notification(notification_object)
+            if notification.is_async:
+                self._async_notifications.append(
+                    (notification, run, notification_object)
+                )
+            else:
+                self._sync_notifications.append(
+                    (notification, run, notification_object)
+                )
+
+    def push(self, sync_push_callback=None, async_push_callback=None):
         """
         Asynchronously push notifications for all runs in the initialized runs list (if they should be pushed).
         When running from a sync environment, the notifications will be pushed asynchronously however the function will
@@ -169,7 +201,7 @@ class NotificationPusher(_NotificationPusherBase):
                         "Failed to push notification async",
                         error=mlrun.errors.err_to_str(result),
                         traceback=traceback.format_exception(
-                            etype=type(result),
+                            result,
                             value=result,
                             tb=result.__traceback__,
                         ),
@@ -180,8 +212,9 @@ class NotificationPusher(_NotificationPusherBase):
             notifications_amount=len(self._sync_notifications)
             + len(self._async_notifications),
         )
-
-        self._push(sync_push, async_push)
+        sync_push_callback = sync_push_callback or sync_push
+        async_push_callback = async_push_callback or async_push
+        self._push(sync_push_callback, async_push_callback)
 
     @staticmethod
     def _should_notify(
@@ -202,7 +235,7 @@ class NotificationPusher(_NotificationPusherBase):
         for when_state in when_states:
             if when_state == run_state:
                 if (
-                    run_state == "completed"
+                    run_state == runtimes_constants.RunStates.completed
                     and evaluate_condition_in_separate_process(
                         notification.condition,
                         context={
@@ -210,27 +243,29 @@ class NotificationPusher(_NotificationPusherBase):
                             "notification": notification.to_dict(),
                         },
                     )
-                ) or run_state in ["error", "aborted"]:
+                ) or run_state in [
+                    runtimes_constants.RunStates.error,
+                    runtimes_constants.RunStates.aborted,
+                    runtimes_constants.RunStates.running,
+                ]:
                     return True
 
         return False
 
     def _load_notification(
-        self, run: mlrun.model.RunObject, notification_object: mlrun.model.Notification
-    ) -> NotificationBase:
+        self, notification_object: mlrun.model.Notification
+    ) -> base.NotificationBase:
         name = notification_object.name
-        notification_type = NotificationTypes(
-            notification_object.kind or NotificationTypes.console
+        notification_type = notification_module.NotificationTypes(
+            notification_object.kind or notification_module.NotificationTypes.console
         )
         params = {}
-        params.update(notification_object.secret_params)
-        params.update(notification_object.params)
-        notification = notification_type.get_notification()(name, params)
-        if notification.is_async:
-            self._async_notifications.append((notification, run, notification_object))
-        else:
-            self._sync_notifications.append((notification, run, notification_object))
-
+        params.update(notification_object.secret_params or {})
+        params.update(notification_object.params or {})
+        default_params = self._default_params.get(notification_type.value, {})
+        notification = notification_type.get_notification()(
+            name, params, default_params
+        )
         logger.debug(
             "Loaded notification", notification=name, type=notification_type.value
         )
@@ -250,10 +285,14 @@ class NotificationPusher(_NotificationPusherBase):
             custom_message = (
                 f" (workflow: {run.metadata.labels['workflow']}){custom_message}"
             )
-            runs.extend(self.get_workflow_steps(run))
+            project = run.metadata.project
+            workflow_id = run.status.results.get("workflow_id", None)
+            db = mlrun.get_run_db()
+            runs.extend(Workflow.get_workflow_steps(db, workflow_id, project))
 
         message = (
             self.messages.get(run.state(), "").format(resource=resource)
+            + f" in project {run.metadata.project}"
             + custom_message
         )
 
@@ -265,13 +304,14 @@ class NotificationPusher(_NotificationPusherBase):
 
     def _push_notification_sync(
         self,
-        notification: NotificationBase,
+        notification: base.NotificationBase,
         run: mlrun.model.RunObject,
         notification_object: mlrun.model.Notification,
     ):
         message, severity, runs = self._prepare_notification_args(
             run, notification_object
         )
+
         logger.debug(
             "Pushing sync notification",
             notification=sanitize_notification(notification_object.to_dict()),
@@ -282,6 +322,7 @@ class NotificationPusher(_NotificationPusherBase):
             "project": run.metadata.project,
             "notification": notification_object,
             "status": mlrun.common.schemas.NotificationStatus.SENT,
+            "run_state": run.state(),
         }
         try:
             notification.push(message, severity, runs)
@@ -313,13 +354,14 @@ class NotificationPusher(_NotificationPusherBase):
 
     async def _push_notification_async(
         self,
-        notification: NotificationBase,
+        notification: base.NotificationBase,
         run: mlrun.model.RunObject,
         notification_object: mlrun.model.Notification,
     ):
         message, severity, runs = self._prepare_notification_args(
             run, notification_object
         )
+
         logger.debug(
             "Pushing async notification",
             notification=sanitize_notification(notification_object.to_dict()),
@@ -329,6 +371,7 @@ class NotificationPusher(_NotificationPusherBase):
             "run_uid": run.metadata.uid,
             "project": run.metadata.project,
             "notification": notification_object,
+            "run_state": run.state(),
             "status": mlrun.common.schemas.NotificationStatus.SENT,
         }
         try:
@@ -366,10 +409,29 @@ class NotificationPusher(_NotificationPusherBase):
         run_uid: str,
         project: str,
         notification: mlrun.model.Notification,
-        status: str = None,
+        run_state: runtimes_constants.RunStates,
+        status: typing.Optional[str] = None,
         sent_time: typing.Optional[datetime.datetime] = None,
         reason: typing.Optional[str] = None,
     ):
+        # Skip update the notification state if the following conditions are met:
+        # 1. the run is not in a terminal state
+        # 2. the when contains only one state (which is the current state)
+        # Skip updating because currently each notification has only one row in the db, even if it has multiple when.
+        # This means that if the notification is updated to sent for running state for example, it will not send for
+        # The terminal state
+        # TODO: Change this behavior after implementing ML-8723
+        if (
+            run_state not in runtimes_constants.RunStates.terminal_states()
+            and len(notification.when) > 1
+        ):
+            logger.debug(
+                "Skip updating notification status - run not in terminal state",
+                run_uid=run_uid,
+                state=run_state,
+            )
+            return
+
         db = mlrun.get_run_db()
         notification.status = status or notification.status
         notification.sent_time = sent_time or notification.sent_time
@@ -394,136 +456,13 @@ class NotificationPusher(_NotificationPusherBase):
             mask_params=False,
         )
 
-    def get_workflow_steps(self, run: mlrun.model.RunObject) -> list:
-        steps = []
-        db = mlrun.get_run_db()
-
-        def _add_run_step(_step: mlrun_pipelines.models.PipelineStep):
-            try:
-                _run = db.list_runs(
-                    project=run.metadata.project,
-                    labels=f"{mlrun_constants.MLRunInternalLabels.runner_pod}={_step.node_name}",
-                )[0]
-            except IndexError:
-                _run = {
-                    "metadata": {
-                        "name": _step.display_name,
-                        "project": run.metadata.project,
-                    },
-                }
-            _run["step_kind"] = _step.step_type
-            if _step.skipped:
-                _run.setdefault("status", {})["state"] = (
-                    mlrun.common.runtimes.constants.RunStates.skipped
-                )
-            steps.append(_run)
-
-        def _add_deploy_function_step(_step: mlrun_pipelines.models.PipelineStep):
-            project, name, hash_key = self._extract_function_uri(
-                _step.get_annotation("mlrun/function-uri")
-            )
-            if name:
-                try:
-                    function = db.get_function(
-                        project=project, name=name, hash_key=hash_key
-                    )
-                except mlrun.errors.MLRunNotFoundError:
-                    # If the function is not found (if build failed for example), we will create a dummy
-                    # function object for the notification to display the function name
-                    function = {
-                        "metadata": {
-                            "name": name,
-                            "project": project,
-                            "hash_key": hash_key,
-                        },
-                    }
-                pod_phase = _step.phase
-                if _step.skipped:
-                    state = mlrun.common.schemas.FunctionState.skipped
-                else:
-                    state = mlrun.common.runtimes.constants.PodPhases.pod_phase_to_run_state(
-                        pod_phase
-                    )
-                function["status"] = {"state": state}
-                if isinstance(function["metadata"].get("updated"), datetime.datetime):
-                    function["metadata"]["updated"] = function["metadata"][
-                        "updated"
-                    ].isoformat()
-                function["step_kind"] = _step.step_type
-                steps.append(function)
-
-        step_methods = {
-            mlrun_pipelines.common.ops.PipelineRunType.run: _add_run_step,
-            mlrun_pipelines.common.ops.PipelineRunType.build: _add_deploy_function_step,
-            mlrun_pipelines.common.ops.PipelineRunType.deploy: _add_deploy_function_step,
-        }
-
-        workflow_id = run.status.results.get("workflow_id", None)
-        if not workflow_id:
-            return steps
-
-        workflow_manifest = self._get_workflow_manifest(workflow_id)
-        if not workflow_manifest:
-            return steps
-
-        try:
-            for step in workflow_manifest.get_steps():
-                step_method = step_methods.get(step.step_type)
-                if step_method:
-                    step_method(step)
-            return steps
-        except Exception:
-            # If we fail to read the pipeline steps, we will return the list of runs that have the same workflow id
-            logger.warning(
-                "Failed to extract workflow steps from workflow manifest, "
-                "returning all runs with the workflow id label",
-                workflow_id=workflow_id,
-                traceback=traceback.format_exc(),
-            )
-            return db.list_runs(
-                project=run.metadata.project,
-                labels=f"workflow={workflow_id}",
-            )
-
-    @staticmethod
-    def _get_workflow_manifest(
-        workflow_id: str,
-    ) -> typing.Optional[mlrun_pipelines.models.PipelineManifest]:
-        kfp_client = mlrun_pipelines.utils.get_client(mlrun.mlconf.kfp_url)
-
-        # arbitrary timeout of 5 seconds, the workflow should be done by now
-        kfp_run = kfp_client.wait_for_run_completion(workflow_id, 5)
-        if not kfp_run:
-            return None
-
-        kfp_run = mlrun_pipelines.models.PipelineRun(kfp_run)
-        return kfp_run.workflow_manifest()
-
-    def _extract_function_uri(self, function_uri: str) -> tuple[str, str, str]:
-        """
-        Extract the project, name, and hash key from a function uri.
-        Examples:
-            - "project/name@hash_key" returns project, name, hash_key
-            - "project/name returns" project, name, ""
-        """
-        project, name, hash_key = None, None, None
-        hashed_pattern = r"^(.+)/(.+)@(.+)$"
-        pattern = r"^(.+)/(.+)$"
-        match = re.match(hashed_pattern, function_uri)
-        if match:
-            project, name, hash_key = match.groups()
-        else:
-            match = re.match(pattern, function_uri)
-            if match:
-                project, name = match.groups()
-                hash_key = ""
-        return project, name, hash_key
-
 
 class CustomNotificationPusher(_NotificationPusherBase):
-    def __init__(self, notification_types: list[str] = None):
+    def __init__(self, notification_types: typing.Optional[list[str]] = None):
         notifications = {
-            notification_type: NotificationTypes(notification_type).get_notification()()
+            notification_type: notification_module.NotificationTypes(
+                notification_type
+            ).get_notification()()
             for notification_type in notification_types
         }
         self._sync_notifications = {
@@ -536,6 +475,7 @@ class CustomNotificationPusher(_NotificationPusherBase):
             for notification_type, notification in notifications.items()
             if notification.is_async
         }
+        self._server_notifications = []
 
     @property
     def notifications(self):
@@ -543,6 +483,10 @@ class CustomNotificationPusher(_NotificationPusherBase):
         notifications.update(self._async_notifications)
         return notifications
 
+    @property
+    def server_notifications(self):
+        return self._server_notifications
+
     def push(
         self,
         message: str,
@@ -550,7 +494,7 @@ class CustomNotificationPusher(_NotificationPusherBase):
             mlrun.common.schemas.NotificationSeverity, str
         ] = mlrun.common.schemas.NotificationSeverity.INFO,
         runs: typing.Union[mlrun.lists.RunList, list] = None,
-        custom_html: str = None,
+        custom_html: typing.Optional[str] = None,
     ):
         def sync_push():
             for notification_type, notification in self._sync_notifications.items():
@@ -572,14 +516,43 @@ class CustomNotificationPusher(_NotificationPusherBase):
     def add_notification(
         self,
         notification_type: str,
-        params: dict[str, str] = None,
+        params: typing.Optional[dict[str, str]] = None,
+        name: typing.Optional[str] = None,
+        message: typing.Optional[str] = None,
+        severity: mlrun.common.schemas.notification.NotificationSeverity = (
+            mlrun.common.schemas.notification.NotificationSeverity.INFO
+        ),
+        when: typing.Optional[list[str]] = None,
+        condition: typing.Optional[str] = None,
+        secret_params: typing.Optional[dict[str, str]] = None,
     ):
+        if notification_type not in [
+            notification_module.NotificationTypes.console,
+            notification_module.NotificationTypes.ipython,
+        ]:
+            # We want that only the console and ipython notifications will be notified by the client.
+            # The rest of the notifications will be notified by the BE.
+            self._server_notifications.append(
+                mlrun.model.Notification(
+                    kind=notification_type,
+                    name=name,
+                    message=message,
+                    severity=severity,
+                    when=when or runtimes_constants.RunStates.notification_states(),
+                    params=params,
+                    secret_params=secret_params,
+                )
+            )
+            return
+
         if notification_type in self._async_notifications:
             self._async_notifications[notification_type].load_notification(params)
         elif notification_type in self._sync_notifications:
             self._sync_notifications[notification_type].load_notification(params)
         else:
-            notification = NotificationTypes(notification_type).get_notification()(
+            notification = notification_module.NotificationTypes(
+                notification_type
+            ).get_notification()(
                 params=params,
             )
             if notification.is_async:
@@ -597,7 +570,9 @@ class CustomNotificationPusher(_NotificationPusherBase):
         else:
             logger.warning(f"No notification of type {notification_type} in project")
 
-    def edit_notification(self, notification_type: str, params: dict[str, str] = None):
+    def edit_notification(
+        self, notification_type: str, params: typing.Optional[dict[str, str]] = None
+    ):
         self.remove_notification(notification_type)
         self.add_notification(notification_type, params)
 
@@ -611,7 +586,7 @@ class CustomNotificationPusher(_NotificationPusherBase):
 
         # get notification's inverse dependencies, and only push the notification if
         # none of its inverse dependencies are being sent
-        inverse_dependencies = NotificationTypes(
+        inverse_dependencies = notification_module.NotificationTypes(
             notification_type
         ).inverse_dependencies()
         for inverse_dependency in inverse_dependencies:
@@ -627,36 +602,28 @@ class CustomNotificationPusher(_NotificationPusherBase):
     def push_pipeline_start_message(
         self,
         project: str,
-        commit_id: str = None,
-        pipeline_id: str = None,
+        pipeline_id: typing.Optional[str] = None,
+    ):
+        db = mlrun.get_run_db()
+        db.push_run_notifications(pipeline_id, project)
+
+    def push_pipeline_start_message_from_client(
+        self,
+        project: str,
+        commit_id: typing.Optional[str] = None,
+        pipeline_id: typing.Optional[str] = None,
         has_workflow_url: bool = False,
     ):
-        message = f"Workflow started in project {project}"
-        if pipeline_id:
-            message += f" id={pipeline_id}"
-        commit_id = (
-            commit_id or os.environ.get("GITHUB_SHA") or os.environ.get("CI_COMMIT_SHA")
+        html, message = self.generate_start_message(
+            commit_id, has_workflow_url, pipeline_id, project
         )
-        if commit_id:
-            message += f", commit={commit_id}"
-        if has_workflow_url:
-            url = mlrun.utils.helpers.get_workflow_url(project, pipeline_id)
-        else:
-            url = mlrun.utils.helpers.get_ui_url(project)
-        html = ""
-        if url:
-            html = (
-                message
-                + f'<div><a href="{url}" target="_blank">click here to view progress</a></div>'
-            )
-            message = message + f", check progress in {url}"
         self.push(message, "info", custom_html=html)
 
     def push_pipeline_run_results(
         self,
         runs: typing.Union[mlrun.lists.RunList, list],
         push_all: bool = False,
-        state: str = None,
+        state: typing.Optional[str] = None,
     ):
         """
         push a structured table with run results to notification targets
@@ -682,6 +649,30 @@ class CustomNotificationPusher(_NotificationPusherBase):
             text += f", state={state}"
         self.push(text, "info", runs=runs_list)
 
+    def generate_start_message(
+        self, commit_id=None, has_workflow_url=None, pipeline_id=None, project=None
+    ):
+        message = f"Workflow started in project {project}"
+        if pipeline_id:
+            message += f" id={pipeline_id}"
+        commit_id = (
+            commit_id or os.environ.get("GITHUB_SHA") or os.environ.get("CI_COMMIT_SHA")
+        )
+        if commit_id:
+            message += f", commit={commit_id}"
+        if has_workflow_url:
+            url = mlrun.utils.helpers.get_workflow_url(project, pipeline_id)
+        else:
+            url = mlrun.utils.helpers.get_runs_url(project)
+        html = ""
+        if url:
+            html = (
+                message
+                + f'<div><a href="{url}" target="_blank">click here to view progress</a></div>'
+            )
+            message = message + f", check progress in {url}"
+        return html, message
+
 
 def sanitize_notification(notification_dict: dict):
     notification_dict.pop("secret_params", None)