mlrun 1.3.3__py3-none-any.whl → 1.4.0__py3-none-any.whl

mlrun/utils/notifications/notification_pusher.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,123 +12,402 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import ast
+import asyncio
+import datetime
 import os
 import typing
+from concurrent.futures import ThreadPoolExecutor
 
+from fastapi.concurrency import run_in_threadpool
+
+import mlrun.api.db.base
+import mlrun.api.db.session
+import mlrun.common.schemas
+import mlrun.config
 import mlrun.lists
 import mlrun.model
 import mlrun.utils.helpers
+from mlrun.utils import logger
+from mlrun.utils.condition_evaluator import evaluate_condition_in_separate_process
 
-from .notification import NotificationBase, NotificationSeverity, NotificationTypes
+from .notification import NotificationBase, NotificationTypes
 
 
 class NotificationPusher(object):
 
     messages = {
-        "success": "Run completed",
+        "completed": "Run completed",
         "error": "Run failed",
+        "aborted": "Run aborted",
     }
 
     def __init__(self, runs: typing.Union[mlrun.lists.RunList, list]):
         self._runs = runs
-        self._notification_data = []
-        self._notifications = {}
+        self._sync_notifications = []
+        self._async_notifications = []
 
         for run in self._runs:
-            for notification_config in run.get("spec", {}).get("notifications", []):
-                if self._should_notify(run, notification_config):
-                    self._notification_data.append((run, notification_config))
-
-    def push(self):
-        for notification_data in self._notification_data:
-            self._send_notification(
-                self._load_notification(notification_data[1]), *notification_data
-            )
+            if isinstance(run, dict):
+                run = mlrun.model.RunObject.from_dict(run)
+
+            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
+                    )
+
+                if self._should_notify(run, notification):
+                    self._load_notification(run, notification)
+
+    def push(
+        self,
+        db: mlrun.api.db.base.DBInterface = 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
+        wait for all notifications to be pushed before returning.
+        """
+
+        if not len(self._sync_notifications) and not len(self._async_notifications):
+            return
+
+        def _sync_push():
+            for notification_data in self._sync_notifications:
+                self._push_notification_sync(
+                    notification_data[0],
+                    notification_data[1],
+                    notification_data[2],
+                    db,
+                )
+
+        async def _async_push():
+            tasks = []
+            for notification_data in self._async_notifications:
+                tasks.append(
+                    self._push_notification_async(
+                        notification_data[0],
+                        notification_data[1],
+                        notification_data[2],
+                        db,
+                    )
+                )
+
+            # return exceptions to "best-effort" fire all notifications
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+        logger.debug(
+            "Pushing notifications",
+            notifications_amount=len(self._sync_notifications)
+            + len(self._async_notifications),
+        )
+
+        # first push async notifications
+        main_event_loop = asyncio.get_event_loop()
+        if not main_event_loop.is_running():
+            # If running mlrun SDK locally (not from jupyter), there isn't necessarily an event loop.
+            # We create a new event loop and run the async push function in it.
+            main_event_loop.run_until_complete(_async_push())
+        elif mlrun.utils.helpers.is_running_in_jupyter_notebook():
+            # Running in Jupyter notebook.
+            # In this case, we need to create a new thread, run a separate event loop in
+            # that thread, and use it instead of the main_event_loop.
+            # This is necessary because Jupyter Notebook has its own event loop,
+            # but it runs in the main thread. As long as a cell is running,
+            # the event loop will not execute properly
+            _run_coroutine_in_jupyter_notebook(coroutine_method=_async_push)
+        else:
+            # Running in mlrun api, we are in a separate thread from the one in which
+            # the main event loop, so we can just send the notifications to that loop
+            asyncio.run_coroutine_threadsafe(_async_push(), main_event_loop)
+
+        # then push sync notifications
+        if not mlrun.config.is_running_as_api():
+            _sync_push()
 
     @staticmethod
     def _should_notify(
-        run: typing.Union[mlrun.model.RunObject, dict], notification_config: dict
+        run: mlrun.model.RunObject,
+        notification: mlrun.model.Notification,
     ) -> bool:
-        when_states = notification_config.get("when", [])
-        condition = notification_config.get("condition", "")
-        run_state = run.get("status", {}).get("state", "")
+        when_states = notification.when
+        run_state = run.state()
+
+        # if the notification isn't pending, don't push it
+        if (
+            notification.status
+            and notification.status != mlrun.common.schemas.NotificationStatus.PENDING
+        ):
+            return False
 
         # if at least one condition is met, notify
         for when_state in when_states:
-            if (
-                when_state == "success"
-                and run_state == "success"
-                and (not condition or ast.literal_eval(condition))
-            ) or (when_state == "failure" and run_state == "error"):
-                return True
+            if when_state == run_state:
+                if (
+                    run_state == "completed"
+                    and evaluate_condition_in_separate_process(
+                        notification.condition,
+                        context={
+                            "run": run.to_dict(),
+                            "notification": notification.to_dict(),
+                        },
+                    )
+                ) or run_state in ["error", "aborted"]:
+                    return True
 
         return False
 
-    def _load_notification(self, notification_config: dict) -> NotificationBase:
-        params = notification_config.get("params", {})
+    def _load_notification(
+        self, run: mlrun.model.RunObject, notification_object: mlrun.model.Notification
+    ) -> NotificationBase:
+        name = notification_object.name
         notification_type = NotificationTypes(
-            notification_config.get("type", "console")
+            notification_object.kind or NotificationTypes.console
         )
-        if notification_type not in self._notifications:
-            self._notifications[
-                notification_type
-            ] = notification_type.get_notification()(params)
+        notification = notification_type.get_notification()(
+            name, notification_object.params
+        )
+        if notification.is_async:
+            self._async_notifications.append((notification, run, notification_object))
         else:
-            self._notifications[notification_type].load_notification(params)
+            self._sync_notifications.append((notification, run, notification_object))
 
-        return self._notifications[notification_type]
+        logger.debug(
+            "Loaded notification", notification=name, type=notification_type.value
+        )
+        return notification
 
-    def _send_notification(
-        self, notification: NotificationBase, run: dict, notification_config: dict
+    def _prepare_notification_args(
+        self, run: mlrun.model.RunObject, notification_object: mlrun.model.Notification
     ):
-        message = self.messages.get(run.get("status", {}).get("state", ""), "")
-        severity = notification_config.get("severity", NotificationSeverity.INFO)
-        notification.send(message, severity, [run])
+        custom_message = (
+            f": {notification_object.message}" if notification_object.message else ""
+        )
+        message = self.messages.get(run.state(), "") + custom_message
+
+        severity = (
+            notification_object.severity
+            or mlrun.common.schemas.NotificationSeverity.INFO
+        )
+        return message, severity, [run.to_dict()]
+
+    def _push_notification_sync(
+        self,
+        notification: NotificationBase,
+        run: mlrun.model.RunObject,
+        notification_object: mlrun.model.Notification,
+        db: mlrun.api.db.base.DBInterface,
+    ):
+        message, severity, runs = self._prepare_notification_args(
+            run, notification_object
+        )
+        logger.debug(
+            "Pushing sync notification",
+            notification=_sanitize_notification(notification_object),
+            run_uid=run.metadata.uid,
+        )
+        try:
+            notification.push(message, severity, runs)
+            self._update_notification_status(
+                db,
+                run.metadata.uid,
+                run.metadata.project,
+                notification_object,
+                status=mlrun.common.schemas.NotificationStatus.SENT,
+                sent_time=datetime.datetime.now(tz=datetime.timezone.utc),
+            )
+        except Exception as exc:
+            self._update_notification_status(
+                db,
+                run.metadata.uid,
+                run.metadata.project,
+                notification_object,
+                status=mlrun.common.schemas.NotificationStatus.ERROR,
+            )
+            raise exc
+
+    async def _push_notification_async(
+        self,
+        notification: NotificationBase,
+        run: mlrun.model.RunObject,
+        notification_object: mlrun.model.Notification,
+        db: mlrun.api.db.base.DBInterface,
+    ):
+        message, severity, runs = self._prepare_notification_args(
+            run, notification_object
+        )
+        logger.debug(
+            "Pushing async notification",
+            notification=_sanitize_notification(notification_object),
+            run_uid=run.metadata.uid,
+        )
+        try:
+            await notification.push(message, severity, runs)
+
+            await run_in_threadpool(
+                self._update_notification_status,
+                db,
+                run.metadata.uid,
+                run.metadata.project,
+                notification_object,
+                status=mlrun.common.schemas.NotificationStatus.SENT,
+                sent_time=datetime.datetime.now(tz=datetime.timezone.utc),
+            )
+        except Exception as exc:
+            await run_in_threadpool(
+                self._update_notification_status,
+                db,
+                run.metadata.uid,
+                run.metadata.project,
+                notification_object,
+                status=mlrun.common.schemas.NotificationStatus.ERROR,
+            )
+            raise exc
+
+    @staticmethod
+    def _update_notification_status(
+        db: mlrun.api.db.base.DBInterface,
+        run_uid: str,
+        project: str,
+        notification: mlrun.model.Notification,
+        status: str = None,
+        sent_time: typing.Optional[datetime.datetime] = None,
+    ):
+
+        # nothing to update if not running as api
+        # note, the notification mechanism may run "locally" for certain runtimes
+        if not mlrun.config.is_running_as_api():
+            return
+
+        # TODO: move to api side
+        db_session = mlrun.api.db.session.create_session()
+        notification.status = status or notification.status
+        notification.sent_time = sent_time or notification.sent_time
+
+        # store directly in db, no need to use crud as the secrets are already loaded
+        db.store_run_notifications(
+            db_session,
+            [notification],
+            run_uid,
+            project,
+        )
 
 
 class CustomNotificationPusher(object):
     def __init__(self, notification_types: typing.List[str] = None):
-        self._notifications = {
+        notifications = {
             notification_type: NotificationTypes(notification_type).get_notification()()
             for notification_type in notification_types
         }
+        self._sync_notifications = {
+            notification_type: notification
+            for notification_type, notification in notifications.items()
+            if not notification.is_async
+        }
+        self._async_notifications = {
+            notification_type: notification
+            for notification_type, notification in notifications.items()
+            if notification.is_async
+        }
 
     def push(
         self,
         message: str,
-        severity: typing.Union[NotificationSeverity, str] = NotificationSeverity.INFO,
+        severity: typing.Union[
+            mlrun.common.schemas.NotificationSeverity, str
+        ] = mlrun.common.schemas.NotificationSeverity.INFO,
         runs: typing.Union[mlrun.lists.RunList, list] = None,
         custom_html: str = None,
     ):
-        for notification_type, notification in self._notifications.items():
-            if self.should_send_notification(notification_type):
-                notification.send(message, severity, runs, custom_html)
+        def _sync_push():
+            for notification_type, notification in self._sync_notifications.items():
+                if self.should_push_notification(notification_type):
+                    notification.push(message, severity, runs, custom_html)
+
+        async def _async_push():
+            tasks = []
+            for notification_type, notification in self._async_notifications.items():
+                if self.should_push_notification(notification_type):
+                    tasks.append(
+                        notification.push(message, severity, runs, custom_html)
+                    )
+            # return exceptions to "best-effort" fire all notifications
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+        # first push async notifications
+        main_event_loop = asyncio.get_event_loop()
+        if not main_event_loop.is_running():
+            # If running mlrun SDK locally (not from jupyter), there isn't necessarily an event loop.
+            # We create a new event loop and run the async push function in it.
+            main_event_loop.run_until_complete(_async_push())
+        elif mlrun.utils.helpers.is_running_in_jupyter_notebook():
+            # Running in Jupyter notebook.
+            # In this case, we need to create a new thread, run a separate event loop in
+            # that thread, and use it instead of the main_event_loop.
+            # This is necessary because Jupyter Notebook has its own event loop,
+            # but it runs in the main thread. As long as a cell is running,
+            # the event loop will not execute properly
+            _run_coroutine_in_jupyter_notebook(coroutine_method=_async_push)
+        else:
+            # Running in mlrun api, we are in a separate thread from the one in which
+            # the main event loop, so we can just send the notifications to that loop
+            asyncio.run_coroutine_threadsafe(_async_push(), main_event_loop)
+
+        # then push sync notifications
+        if not mlrun.config.is_running_as_api():
+            _sync_push()
 
     def add_notification(
         self, notification_type: str, params: typing.Dict[str, str] = None
     ):
-        if notification_type in self._notifications:
-            self._notifications[notification_type].load_notification(params)
+        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:
-            self._notifications[notification_type] = NotificationTypes(
-                notification_type
-            ).get_notification()(params)
+            notification = NotificationTypes(notification_type).get_notification()(
+                params=params,
+            )
+            if notification.is_async:
+                self._async_notifications[notification_type] = notification
+            else:
+                self._sync_notifications[notification_type] = notification
 
-    def should_send_notification(self, notification_type):
-        notification = self._notifications.get(notification_type)
+    def remove_notification(self, notification_type: str):
+        if notification_type in self._async_notifications:
+            del self._async_notifications[notification_type]
+
+        elif notification_type in self._sync_notifications:
+            del self._sync_notifications[notification_type]
+
+        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
+    ):
+        self.remove_notification(notification_type)
+        self.add_notification(notification_type, params)
+
+    def should_push_notification(self, notification_type):
+        notifications = {}
+        notifications.update(self._sync_notifications)
+        notifications.update(self._async_notifications)
+        notification = notifications.get(notification_type)
         if not notification or not notification.active:
             return False
 
-        # get notification's inverse dependencies, and only send the notification if
+        # get notification's inverse dependencies, and only push the notification if
         # none of its inverse dependencies are being sent
         inverse_dependencies = NotificationTypes(
             notification_type
         ).inverse_dependencies()
         for inverse_dependency in inverse_dependencies:
-            inverse_dependency_notification = self._notifications.get(
-                inverse_dependency
-            )
+            inverse_dependency_notification = notifications.get(inverse_dependency)
             if (
                 inverse_dependency_notification
                 and inverse_dependency_notification.active
@@ -144,7 +423,7 @@ class CustomNotificationPusher(object):
         pipeline_id: str = None,
         has_workflow_url: bool = False,
     ):
-        message = f"Pipeline started in project {project}"
+        message = f"Workflow started in project {project}"
         if pipeline_id:
             message += f" id={pipeline_id}"
         commit_id = (
@@ -194,3 +473,45 @@ class CustomNotificationPusher(object):
         if state:
             text += f", state={state}"
         self.push(text, "info", runs=runs_list)
+
+
+def _sanitize_notification(notification: mlrun.model.Notification):
+    notification_dict = notification.to_dict()
+    notification_dict.pop("params", None)
+    return notification_dict
+
+
+def _separate_sync_notifications(
+    notifications: typing.List[NotificationBase],
+) -> typing.Tuple[typing.List[NotificationBase], typing.List[NotificationBase]]:
+    sync_notifications = []
+    async_notifications = []
+    for notification in notifications:
+        if notification.is_async:
+            async_notifications.append(notification)
+        else:
+            sync_notifications.append(notification)
+    return sync_notifications, async_notifications
+
+
+def _run_coroutine_in_jupyter_notebook(coroutine_method):
+    """
+    Execute a coroutine in a Jupyter Notebook environment.
+
+    This function creates a new thread pool executor with a single thread and a new event loop.
+    It sets the created event loop as the current event loop.
+    Then, it submits the coroutine to the event loop and waits for its completion.
+
+    This approach is used in Jupyter Notebook to ensure the proper execution of the event loop in a separate thread,
+    allowing for the asynchronous push operation to be executed while the notebook is running.
+
+    :param coroutine_method: The coroutine method to be executed.
+    :return: The result of the executed coroutine.
+    """
+    thread_pool_executer = ThreadPoolExecutor(1)
+    async_event_loop = asyncio.new_event_loop()
+    thread_pool_executer.submit(asyncio.set_event_loop, async_event_loop).result()
+    result = thread_pool_executer.submit(
+        async_event_loop.run_until_complete, coroutine_method()
+    ).result()
+    return result

mlrun/utils/regex.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,6 +13,11 @@
 # limitations under the License.
 #
 
+# pipeline param format which is passed when running a pipeline (e.g. {{pipelineparam:op=;name=mem}})
+# https://github.com/kubeflow/pipelines/blob/16edebf4eaf84cd7478e2601ef4878ab339a7854/sdk/python/kfp/dsl/_pipeline_param.py#L213
+# this is expected to be resolved at runtime
+pipeline_param = [r"{{pipelineparam:op=([\w\s_-]*);name=([\w\s_-]+)}}"]
+
 # k8s character limit is for 63 characters
 k8s_character_limit = [r"^.{0,63}$"]
 

mlrun/utils/singleton.py

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

mlrun/utils/v3io_clients.py

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