[Notifications] Display Step Kind In Pipeline Notification (#5538)

* Add step kind to pipeline notification * better system test * redundant comment * fix ut --------- Co-authored-by: quaark <a.melnick@icloud.com>
mlrun · May 9, 2024 · e3eafff · e3eafff
1 parent fb3637c
commit e3eafff
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 21 deletions.
diff --git a/mlrun/utils/notifications/notification/slack.py b/mlrun/utils/notifications/notification/slack.py
@@ -154,18 +154,21 @@ def _get_run_line(self, run: dict) -> dict:
         url = mlrun.utils.helpers.get_ui_url(meta.get("project"), meta.get("uid"))
 
         # Only show the URL if the run is not a function (serving or mlrun function)
-        if run.get("kind") not in ["serving", None] and url:
+        kind = run.get("step_kind")
+        if 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(

diff --git a/mlrun/utils/notifications/notification_pusher.py b/mlrun/utils/notifications/notification_pusher.py
@@ -392,15 +392,15 @@ def get_workflow_steps(self, run: mlrun.model.RunObject) -> list:
         steps = []
         db = mlrun.get_run_db()
 
-        def _add_run_step(_node_name, _):
-            steps.append(
-                db.list_runs(
-                    project=run.metadata.project,
-                    labels=f"mlrun/runner-pod={_node_name}",
-                )[0]
-            )
+        def _add_run_step(_node_name, _node_template, _step_kind):
+            _run = db.list_runs(
+                project=run.metadata.project,
+                labels=f"mlrun/runner-pod={_node_name}",
+            )[0]
+            _run["step_kind"] = _step_kind
+            steps.append(_run)
 
-        def _add_deploy_function_step(_, _node_template):
+        def _add_deploy_function_step(_, _node_template, _step_kind):
             project, name, hash_key = self._extract_function_uri(
                 _node_template["metadata"]["annotations"]["mlrun/function-uri"]
             )
@@ -428,6 +428,7 @@ def _add_deploy_function_step(_, _node_template):
                     function["metadata"]["updated"] = function["metadata"][
                         "updated"
                     ].isoformat()
+                function["step_kind"] = _step_kind
                 steps.append(function)
 
         step_methods = {
@@ -464,7 +465,7 @@ def _add_deploy_function_step(_, _node_template):
                 )
                 step_method = step_methods.get(step_type)
                 if step_method:
-                    step_method(node_name, node_template)
+                    step_method(node_name, node_template, step_type)
             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

diff --git a/tests/system/projects/test_project.py b/tests/system/projects/test_project.py
@@ -650,14 +650,12 @@ def test_remote_pipeline_with_kfp_engine_from_github(self):
             engine="remote",
             watch=True,
             notification_steps={
-                # workflow runner, summary, train, test and model testing steps
-                "run": 5,
-                # gen data step
-                "job": 1,
+                # gen data function build step
+                "build": 1,
+                # workflow runner, gen data, summary, train, test and model testing steps
+                "run": 6,
                 # serving step
-                "serving": 1,
-                # gen data function build step (doesn't have a `kind` field)
-                None: 1,
+                "deploy": 1,
             },
         )
         self._test_remote_pipeline_from_github(
@@ -1585,7 +1583,7 @@ def _assert_pipeline_notification_steps(
         )[0]
         notification_data_steps = {}
         for step in notification_data:
-            notification_data_steps.setdefault(step.get("kind"), 0)
-            notification_data_steps[step.get("kind")] += 1
+            notification_data_steps.setdefault(step.get("step_kind"), 0)
+            notification_data_steps[step.get("step_kind")] += 1
 
         assert notification_data_steps == notification_steps
diff --git a/tests/utils/test_notifications.py b/tests/utils/test_notifications.py
@@ -297,7 +297,7 @@ def set_result(result):
                             {"text": "*Runs*", "type": "mrkdwn"},
                             {"text": "*Results*", "type": "mrkdwn"},
                             {"text": ":x:  test-run", "type": "mrkdwn"},
-                            {"text": "**", "type": "mrkdwn"},
+                            {"text": "*error*", "type": "mrkdwn"},
                         ],
                         "type": "section",
                     },