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

mlrun/projects/pipelines.py

@@ -20,16 +20,15 @@ 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,
@@ -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
@@ -453,7 +452,7 @@ 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,
     ) -> _PipelineRunStatus:
         pass
 
@@ -463,7 +462,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
 
@@ -577,7 +576,7 @@ class _KFPRunner(_PipelineRunner):
         artifact_path=None,
         namespace=None,
         source=None,
-        notifications: list[mlrun.model.Notification] = None,
+        notifications: typing.Optional[list[mlrun.model.Notification]] = None,
     ) -> _PipelineRunStatus:
         pipeline_context.set(project, workflow_spec)
         workflow_handler = _PipelineRunner._get_handler(
@@ -641,7 +640,7 @@ class _KFPRunner(_PipelineRunner):
         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 +685,7 @@ class _LocalRunner(_PipelineRunner):
         artifact_path=None,
         namespace=None,
         source=None,
-        notifications: list[mlrun.model.Notification] = None,
+        notifications: typing.Optional[list[mlrun.model.Notification]] = None,
     ) -> _PipelineRunStatus:
         pipeline_context.set(project, workflow_spec)
         workflow_handler = _PipelineRunner._get_handler(
@@ -763,13 +762,13 @@ 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,
     ) -> typing.Optional[_PipelineRunStatus]:
         workflow_name = normalize_workflow_name(name=name, project_name=project.name)
         workflow_id = None
@@ -890,7 +889,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:
@@ -984,43 +983,34 @@ def github_webhook(request):
     return {"msg": "pushed"}
 
 
-def load_and_run(context, *args, **kwargs):
-    """
-    This function serves as an alias to `load_and_run_workflow`,
-    allowing to continue using `load_and_run` without modifying existing workflows or exported runs.
-    This approach ensures backward compatibility,
-    while directing all new calls to the updated `load_and_run_workflow` function.
-    """
-    kwargs.pop("load_only", None)
-    kwargs.pop("save", None)
-    load_and_run_workflow(context, *args, **kwargs)
-
-
-def load_and_run_workflow(
+def load_and_run(
     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,
+    save: bool = True,
+    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,
+    load_only: bool = False,
     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.
     This function loads a project from a given remote source and then runs the workflow.
+
     :param context:             mlrun context.
     :param url:                 remote url that represents the project's source.
                                 See 'mlrun.load_project()' for details
@@ -1028,6 +1018,7 @@ def load_and_run_workflow(
     :param init_git:            if True, will git init the context dir
     :param subpath:             project subpath (within the archive)
     :param clone:               if True, always clone (delete any existing content)
+    :param save:                whether to save the created project and artifact in the DB
     :param workflow_name:       name of the workflow
     :param workflow_path:       url to a workflow file, if not a project workflow
     :param workflow_arguments:  kubeflow pipelines arguments (parameters)
@@ -1043,31 +1034,48 @@ def load_and_run_workflow(
     :param schedule:            ScheduleCronTrigger class instance or a standard crontab expression string
     :param cleanup_ttl:         pipeline cleanup ttl in secs (time to wait after workflow completion, at which point the
                                 workflow and all its resources are deleted)
+    :param load_only:           for just loading the project, inner use.
     :param wait_for_completion: wait for workflow completion before returning
     :param project_context:     project context path (used for loading the project)
     """
-    project_context = project_context or f"./{project_name}"
-
-    # Load the project to fetch files which the runner needs, such as remote source files
-    pull_remote_project_files(
-        context=context,
-        project_context=project_context,
-        url=url,
-        project_name=project_name,
-        init_git=init_git,
-        subpath=subpath,
-        clone=clone,
-        schedule=schedule,
-        workflow_name=workflow_name,
-    )
+    try:
+        project = mlrun.load_project(
+            context=project_context or f"./{project_name}",
+            url=url,
+            name=project_name,
+            init_git=init_git,
+            subpath=subpath,
+            clone=clone,
+            save=save,
+            sync_functions=True,
+        )
+    except Exception as error:
+        if schedule:
+            notification_pusher = mlrun.utils.notifications.CustomNotificationPusher(
+                ["slack"]
+            )
+            url = get_ui_url(project_name, context.uid)
+            link = f"<{url}|*view workflow job details*>"
+            message = (
+                f":x: Failed to run scheduled workflow {workflow_name} in Project {project_name} !\n"
+                f"error: ```{error}```\n{link}"
+            )
+            # Sending Slack Notification without losing the original error:
+            try:
+                notification_pusher.push(
+                    message=message,
+                    severity=mlrun.common.schemas.NotificationSeverity.ERROR,
+                )
 
-    # Retrieve the project object:
-    # - If the project exists in the MLRun database, it will be loaded from there.
-    # - If it doesn't exist in the database, it will be created from the previously loaded local directory.
-    project = mlrun.get_or_create_project(
-        context=project_context or f"./{project_name}",
-        name=project_name,
-    )
+            except Exception as exc:
+                logger.error("Failed to send slack notification", exc=err_to_str(exc))
+
+        raise error
+
+    context.logger.info(f"Loaded project {project.name} successfully")
+
+    if load_only:
+        return
 
     # extract "start" notification if exists
     start_notifications = [
@@ -1100,156 +1108,18 @@ def load_and_run_workflow(
         raise RuntimeError(f"Workflow {workflow_log_message} failed") from run.exc
 
     if wait_for_completion:
-        handle_workflow_completion(
-            run=run,
-            project=project,
-            context=context,
-            workflow_log_message=workflow_log_message,
-        )
-
-
-def pull_remote_project_files(
-    context: mlrun.execution.MLClientCtx,
-    project_context: str,
-    url: str,
-    project_name: str,
-    init_git: typing.Optional[bool],
-    subpath: typing.Optional[str],
-    clone: bool,
-    schedule: typing.Optional[
-        typing.Union[str, mlrun.common.schemas.ScheduleCronTrigger]
-    ],
-    workflow_name: typing.Optional[str],
-) -> None:
-    """
-    Load the project to clone remote files if they exist.
-    If an exception occurs during project loading, send a notification if the workflow is scheduled.
-
-    :param context:        MLRun execution context.
-    :param project_context: Path to the project context.
-    :param url:            URL of the project repository.
-    :param project_name:   Name of the project.
-    :param init_git:       Initialize a git repository.
-    :param subpath:        Project subpath within the repository.
-    :param clone:          Whether to clone the repository.
-    :param schedule:       Schedule for running the workflow.
-    :param workflow_name:  Name of the workflow to run.
-    """
-    try:
-        # Load the project to clone remote files if they exist.
-        # Using save=False to avoid overriding changes from the database if it already exists.
-        mlrun.load_project(
-            context=project_context,
-            url=url,
-            name=project_name,
-            init_git=init_git,
-            subpath=subpath,
-            clone=clone,
-            save=False,
-        )
-    except Exception as error:
-        notify_scheduled_workflow_failure(
-            schedule=schedule,
-            project_name=project_name,
-            workflow_name=workflow_name,
-            error=error,
-            context_uid=context.uid,
-        )
-        raise error
-
-
-def notify_scheduled_workflow_failure(
-    schedule,
-    project_name: str,
-    workflow_name: str,
-    error: Exception,
-    context_uid: str,
-) -> None:
-    if schedule:
-        notification_pusher = mlrun.utils.notifications.CustomNotificationPusher(
-            ["slack"]
-        )
-        url = get_ui_url(project_name, context_uid)
-        link = f"<{url}|*view workflow job details*>"
-        message = (
-            f":x: Failed to run scheduled workflow {workflow_name} "
-            f"in Project {project_name}!\n"
-            f"Error: ```{err_to_str(error)}```\n{link}"
-        )
-        # Sending Slack Notification without losing the original error:
         try:
-            notification_pusher.push(
-                message=message,
-                severity=mlrun.common.schemas.NotificationSeverity.ERROR,
-            )
-
+            run.wait_for_completion()
         except Exception as exc:
-            logger.error("Failed to send slack notification", exc=err_to_str(exc))
-
-
-def handle_workflow_completion(
-    run: _PipelineRunStatus,
-    project,
-    context: mlrun.execution.MLClientCtx,
-    workflow_log_message: str,
-) -> None:
-    """
-    Handle workflow completion by waiting for it to finish and logging the final state.
-
-    :param run:                 Run object containing workflow execution details.
-    :param project:             MLRun project object.
-    :param context:             MLRun execution context.
-    :param workflow_log_message: Message used for logging.
-    """
-    try:
-        run.wait_for_completion()
-    except Exception as exc:
-        mlrun.utils.logger.error(
-            "Failed waiting for workflow completion",
-            workflow=workflow_log_message,
-            exc=err_to_str(exc),
-        )
-
-    pipeline_state, _, _ = project.get_run_status(run)
-    context.log_result(key="workflow_state", value=pipeline_state, commit=True)
-    if pipeline_state != mlrun_pipelines.common.models.RunStatuses.succeeded:
-        raise RuntimeError(
-            f"Workflow {workflow_log_message} failed, state={pipeline_state}"
-        )
-
-
-def import_remote_project(
-    context: mlrun.execution.MLClientCtx,
-    url: str = None,
-    project_name: str = "",
-    init_git: bool = None,
-    subpath: str = None,
-    clone: bool = False,
-    save: bool = True,
-    project_context: str = None,
-):
-    """
-    This function loads a project from a given remote source.
-
-    :param context:             mlrun context.
-    :param url:                 remote url that represents the project's source.
-                                See 'mlrun.load_project()' for details
-    :param project_name:        project name
-    :param init_git:            if True, will git init the context dir
-    :param subpath:             project subpath (within the archive)
-    :param clone:               if True, always clone (delete any existing content)
-    :param save:                whether to save the created project and artifact in the DB
-    :param project_context:     project context path (used for loading the project)
-    """
-    project = mlrun.load_project(
-        context=project_context or f"./{project_name}",
-        url=url,
-        name=project_name,
-        init_git=init_git,
-        subpath=subpath,
-        clone=clone,
-        save=save,
-        sync_functions=True,
-    )
+            logger.error(
+                "Failed waiting for workflow completion",
+                workflow=workflow_log_message,
+                exc=err_to_str(exc),
+            )
 
-    context.logger.info(f"Loaded project {project.name} successfully")
+        pipeline_state, _, _ = project.get_run_status(run)
+        context.log_result(key="workflow_state", value=pipeline_state, commit=True)
+        if pipeline_state != mlrun_pipelines.common.models.RunStatuses.succeeded:
+            raise RuntimeError(
+                f"Workflow {workflow_log_message} failed, state={pipeline_state}"
+            )