[Notifications] Fix missed notifications in notifications pusher (#3607)

mlrun · May 24, 2023 · 9771d06 · 9771d06
1 parent 2efa698
commit 9771d06
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/mlrun/api/main.py b/mlrun/api/main.py
@@ -566,7 +566,6 @@ def _push_terminal_run_notifications(db: mlrun.api.db.base.DBInterface, db_sessi
     Get all runs with notification configs which became terminal since the last call to the function
     and push their notifications if they haven't been pushed yet.
     """
-
     # Import here to avoid circular import
     import mlrun.api.api.utils
 
@@ -576,6 +575,8 @@ def _push_terminal_run_notifications(db: mlrun.api.db.base.DBInterface, db_sessi
     # and their notifications haven't been sent yet.
     global _last_notification_push_time
 
+    now = datetime.datetime.now(datetime.timezone.utc)
+
     runs = db.list_runs(
         db_session,
         project="*",
@@ -599,7 +600,7 @@ def _push_terminal_run_notifications(db: mlrun.api.db.base.DBInterface, db_sessi
     )
     mlrun.utils.notifications.NotificationPusher(unmasked_runs).push(db)
 
-    _last_notification_push_time = datetime.datetime.now(datetime.timezone.utc)
+    _last_notification_push_time = now
 
 
 async def _stop_logs():

diff --git a/tests/api/utils/test_scheduler.py b/tests/api/utils/test_scheduler.py
@@ -56,7 +56,10 @@ async def scheduler(db: Session) -> typing.Generator:
 
 
 call_counter: int = 0
-schedule_end_time_margin = 0.5
+
+# TODO: The margin will need to rise for each additional CPU-consuming operation added along the flow,
+#  we need to consider how to decouple in the future
+schedule_end_time_margin = 0.7
 
 
 async def bump_counter():