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

mlrun/projects/pipelines.py

@@ -20,19 +20,18 @@ import tempfile
 import typing
 import uuid
 
-import mlrun_pipelines.common.models
-import mlrun_pipelines.patcher
-import mlrun_pipelines.utils
-
 import mlrun
 import mlrun.common.runtimes.constants
 import mlrun.common.schemas
 import mlrun.common.schemas.function
 import mlrun.common.schemas.workflow
 import mlrun.utils.notifications
+import mlrun_pipelines.common.models
+import mlrun_pipelines.patcher
+import mlrun_pipelines.utils
 from mlrun.errors import err_to_str
 from mlrun.utils import (
-    get_ui_url,
+    get_workflow_url,
     logger,
     normalize_workflow_name,
     retry_until_successful,
@@ -40,7 +39,7 @@ from mlrun.utils import (
 
 from ..common.helpers import parse_versioned_object_uri
 from ..config import config
-from ..run import _run_pipeline, wait_for_pipeline_completion
+from ..run import _run_pipeline, retry_pipeline, wait_for_pipeline_completion
 from ..runtimes.pod import AutoMountType
 
 
@@ -317,7 +316,7 @@ def get_db_function(project, key) -> mlrun.runtimes.BaseRuntime:
 def enrich_function_object(
     project: mlrun.common.schemas.Project,
     function: mlrun.runtimes.BaseRuntime,
-    decorator: typing.Callable = None,
+    decorator: typing.Optional[typing.Callable] = None,
     copy_function: bool = True,
     try_auto_mount: bool = True,
 ) -> mlrun.runtimes.BaseRuntime:
@@ -380,7 +379,7 @@ class _PipelineRunStatus:
         project: "mlrun.projects.MlrunProject",
         workflow: WorkflowSpec = None,
         state: mlrun_pipelines.common.models.RunStatuses = "",
-        exc: Exception = None,
+        exc: typing.Optional[Exception] = None,
     ):
         """
         :param run_id:      unique id of the pipeline run
@@ -422,6 +421,13 @@ class _PipelineRunStatus:
             self._state = returned_state
         return self._state
 
+    def retry(self) -> str:
+        run_id = self._engine.retry(
+            self,
+            project=self.project,
+        )
+        return run_id
+
     def __str__(self):
         return str(self.run_id)
 
@@ -441,6 +447,17 @@ class _PipelineRunner(abc.ABC):
             f"Save operation not supported in {cls.engine} pipeline engine"
         )
 
+    @classmethod
+    @abc.abstractmethod
+    def retry(
+        cls,
+        run: "_PipelineRunStatus",
+        project: typing.Optional["mlrun.projects.MlrunProject"] = None,
+    ) -> str:
+        raise NotImplementedError(
+            f"Retry operation not supported in {cls.engine} pipeline engine"
+        )
+
     @classmethod
     @abc.abstractmethod
     def run(
@@ -453,7 +470,8 @@ class _PipelineRunner(abc.ABC):
         artifact_path=None,
         namespace=None,
         source=None,
-        notifications: list[mlrun.model.Notification] = None,
+        notifications: typing.Optional[list[mlrun.model.Notification]] = None,
+        context: typing.Optional[mlrun.execution.MLClientCtx] = None,
     ) -> _PipelineRunStatus:
         pass
 
@@ -463,7 +481,7 @@ class _PipelineRunner(abc.ABC):
         run: "_PipelineRunStatus",
         project: typing.Optional["mlrun.projects.MlrunProject"] = None,
         timeout: typing.Optional[int] = None,
-        expected_statuses: list[str] = None,
+        expected_statuses: typing.Optional[list[str]] = None,
     ):
         pass
 
@@ -505,11 +523,12 @@ class _PipelineRunner(abc.ABC):
         text = _PipelineRunner._generate_workflow_finished_message(
             run.run_id, errors_counter, run._state
         )
-
         notifiers = notifiers or project.notifiers
         if notifiers:
             notifiers.push(text, "info", runs)
 
+        project.push_pipeline_notification_kfp_runner(run.run_id, run._state, text)
+
         if raise_error:
             raise raise_error
         return state or run._state, errors_counter, text
@@ -577,7 +596,8 @@ class _KFPRunner(_PipelineRunner):
         artifact_path=None,
         namespace=None,
         source=None,
-        notifications: list[mlrun.model.Notification] = None,
+        notifications: typing.Optional[list[mlrun.model.Notification]] = None,
+        context: typing.Optional[mlrun.execution.MLClientCtx] = None,
     ) -> _PipelineRunStatus:
         pipeline_context.set(project, workflow_spec)
         workflow_handler = _PipelineRunner._get_handler(
@@ -595,11 +615,21 @@ class _KFPRunner(_PipelineRunner):
                 "Notifications will only be sent if you wait for pipeline completion. "
                 "Some of the features (like setting message or severity level) are not supported."
             )
-            # for start message, fallback to old notification behavior
             for notification in notifications or []:
                 params = notification.params
                 params.update(notification.secret_params)
-                project.notifiers.add_notification(notification.kind, params)
+                project.notifiers.add_notification(
+                    notification_type=notification.kind,
+                    params=params,
+                    name=notification.name,
+                    message=notification.message,
+                    severity=notification.severity,
+                    when=notification.when,
+                    condition=notification.condition,
+                    secret_params=notification.secret_params,
+                )
+
+        project.spec.notifications = project.notifiers.server_notifications
 
         run_id = _run_pipeline(
             workflow_handler,
@@ -627,21 +657,51 @@ class _KFPRunner(_PipelineRunner):
                     func_name=func.metadata.name,
                     exc_info=err_to_str(exc),
                 )
-        project.notifiers.push_pipeline_start_message(
-            project.metadata.name,
-            project.get_param("commit_id", None),
-            run_id,
-            True,
+
+        # Pushing only relevant notification for the client (ipython and console)
+        project.notifiers.push_pipeline_start_message_from_client(
+            project.metadata.name, pipeline_id=run_id
         )
+
+        if context:
+            project.notifiers.push_pipeline_start_message(
+                project.metadata.name,
+                context.uid,
+            )
+        else:
+            project.push_pipeline_notification_kfp_runner(
+                run_id,
+                mlrun_pipelines.common.models.RunStatuses.running,
+                f"Workflow {run_id} started in project {project.metadata.name}",
+                notifications,
+            )
         pipeline_context.clear()
         return _PipelineRunStatus(run_id, cls, project=project, workflow=workflow_spec)
 
+    @classmethod
+    def retry(
+        cls,
+        run: "_PipelineRunStatus",
+        project: typing.Optional["mlrun.projects.MlrunProject"] = None,
+    ) -> str:
+        project_name = project.metadata.name if project else ""
+        logger.info(
+            "Retrying pipeline",
+            run_id=run.run_id,
+            project=project_name,
+        )
+        run_id = retry_pipeline(
+            run.run_id,
+            project=project_name,
+        )
+        return run_id
+
     @staticmethod
     def wait_for_completion(
         run: "_PipelineRunStatus",
         project: typing.Optional["mlrun.projects.MlrunProject"] = None,
         timeout: typing.Optional[int] = None,
-        expected_statuses: list[str] = None,
+        expected_statuses: typing.Optional[list[str]] = None,
     ):
         project_name = project.metadata.name if project else ""
         logger.info(
@@ -686,7 +746,8 @@ class _LocalRunner(_PipelineRunner):
         artifact_path=None,
         namespace=None,
         source=None,
-        notifications: list[mlrun.model.Notification] = None,
+        notifications: typing.Optional[list[mlrun.model.Notification]] = None,
+        context: typing.Optional[mlrun.execution.MLClientCtx] = None,
     ) -> _PipelineRunStatus:
         pipeline_context.set(project, workflow_spec)
         workflow_handler = _PipelineRunner._get_handler(
@@ -708,7 +769,8 @@ class _LocalRunner(_PipelineRunner):
             project.set_source(source=source)
         pipeline_context.workflow_artifact_path = artifact_path
 
-        project.notifiers.push_pipeline_start_message(
+        # TODO: we should create endpoint for sending custom notification from BE
+        project.notifiers.push_pipeline_start_message_from_client(
             project.metadata.name, pipeline_id=workflow_id
         )
         err = None
@@ -763,13 +825,14 @@ class _RemoteRunner(_PipelineRunner):
         cls,
         project: "mlrun.projects.MlrunProject",
         workflow_spec: WorkflowSpec,
-        name: str = None,
-        workflow_handler: typing.Union[str, typing.Callable] = None,
+        name: typing.Optional[str] = None,
+        workflow_handler: typing.Optional[typing.Union[str, typing.Callable]] = None,
         secrets: mlrun.secrets.SecretsStore = None,
-        artifact_path: str = None,
-        namespace: str = None,
-        source: str = None,
-        notifications: list[mlrun.model.Notification] = None,
+        artifact_path: typing.Optional[str] = None,
+        namespace: typing.Optional[str] = None,
+        source: typing.Optional[str] = None,
+        notifications: typing.Optional[list[mlrun.model.Notification]] = None,
+        context: typing.Optional[mlrun.execution.MLClientCtx] = None,
     ) -> typing.Optional[_PipelineRunStatus]:
         workflow_name = normalize_workflow_name(name=name, project_name=project.name)
         workflow_id = None
@@ -890,7 +953,7 @@ class _RemoteRunner(_PipelineRunner):
         timeout=None,
         expected_statuses=None,
         notifiers: mlrun.utils.notifications.CustomNotificationPusher = None,
-        inner_engine: type[_PipelineRunner] = None,
+        inner_engine: typing.Optional[type[_PipelineRunner]] = None,
     ):
         inner_engine = inner_engine or _KFPRunner
         if inner_engine.engine == _KFPRunner.engine:
@@ -998,25 +1061,25 @@ def load_and_run(context, *args, **kwargs):
 
 def load_and_run_workflow(
     context: mlrun.execution.MLClientCtx,
-    url: str = None,
+    url: typing.Optional[str] = None,
     project_name: str = "",
-    init_git: bool = None,
-    subpath: str = None,
+    init_git: typing.Optional[bool] = None,
+    subpath: typing.Optional[str] = None,
     clone: bool = False,
-    workflow_name: str = None,
-    workflow_path: str = None,
-    workflow_arguments: dict[str, typing.Any] = None,
-    artifact_path: str = None,
-    workflow_handler: typing.Union[str, typing.Callable] = None,
-    namespace: str = None,
+    workflow_name: typing.Optional[str] = None,
+    workflow_path: typing.Optional[str] = None,
+    workflow_arguments: typing.Optional[dict[str, typing.Any]] = None,
+    artifact_path: typing.Optional[str] = None,
+    workflow_handler: typing.Optional[typing.Union[str, typing.Callable]] = None,
+    namespace: typing.Optional[str] = None,
     sync: bool = False,
     dirty: bool = False,
-    engine: str = None,
-    local: bool = None,
+    engine: typing.Optional[str] = None,
+    local: typing.Optional[bool] = None,
     schedule: typing.Union[str, mlrun.common.schemas.ScheduleCronTrigger] = None,
-    cleanup_ttl: int = None,
+    cleanup_ttl: typing.Optional[int] = None,
     wait_for_completion: bool = False,
-    project_context: str = None,
+    project_context: typing.Optional[str] = None,
 ):
     """
     Auxiliary function that the RemoteRunner run once or run every schedule.
@@ -1076,6 +1139,11 @@ def load_and_run_workflow(
         if "running" in notification.when
     ]
 
+    # Prevent redundant notifications for run completion by ensuring that notifications are only triggered when the run
+    # reaches the "running" state, as the server already handles the completion notifications.
+    for notification in start_notifications:
+        notification.when = ["running"]
+
     workflow_log_message = workflow_name or workflow_path
     context.logger.info(f"Running workflow {workflow_log_message} from remote")
     run = project.run(
@@ -1092,6 +1160,7 @@ def load_and_run_workflow(
         engine=engine,
         local=local,
         notifications=start_notifications,
+        context=context,
     )
     context.log_result(key="workflow_id", value=run.run_id)
     context.log_result(key="engine", value=run._engine.engine, commit=True)
@@ -1169,7 +1238,7 @@ def notify_scheduled_workflow_failure(
         notification_pusher = mlrun.utils.notifications.CustomNotificationPusher(
             ["slack"]
         )
-        url = get_ui_url(project_name, context_uid)
+        url = get_workflow_url(project_name, context_uid)
         link = f"<{url}|*view workflow job details*>"
         message = (
             f":x: Failed to run scheduled workflow {workflow_name} "
@@ -1220,13 +1289,13 @@ def handle_workflow_completion(
 
 def import_remote_project(
     context: mlrun.execution.MLClientCtx,
-    url: str = None,
+    url: typing.Optional[str] = None,
     project_name: str = "",
-    init_git: bool = None,
-    subpath: str = None,
+    init_git: typing.Optional[bool] = None,
+    subpath: typing.Optional[str] = None,
     clone: bool = False,
     save: bool = True,
-    project_context: str = None,
+    project_context: typing.Optional[str] = None,
 ):
     """
     This function loads a project from a given remote source.