mlrun 1.6.0rc6__py3-none-any.whl → 1.6.0rc8__py3-none-any.whl

mlrun/package/packagers/python_standard_library_packagers.py

@@ -45,8 +45,7 @@ class NonePackager(DefaultPackager):
     DEFAULT_PACKING_ARTIFACT_TYPE = ArtifactType.RESULT
 
     # TODO: `None` as pickle will be available from Python 3.10, so this method can be removed once we move to 3.10.
-    @classmethod
-    def get_supported_artifact_types(cls) -> List[str]:
+    def get_supported_artifact_types(self) -> List[str]:
         """
         Get all the supported artifact types on this packager. It will be the same as `DefaultPackager` but without the
         'object' artifact type support (None cannot be pickled, only from Python 3.10, and it should not be pickled
@@ -95,9 +94,8 @@ class StrPackager(DefaultPackager):
     DEFAULT_PACKING_ARTIFACT_TYPE = ArtifactType.RESULT
     DEFAULT_UNPACKING_ARTIFACT_TYPE = ArtifactType.PATH
 
-    @classmethod
     def pack_path(
-        cls, obj: str, key: str, archive_format: str = DEFAULT_ARCHIVE_FORMAT
+        self, obj: str, key: str, archive_format: str = DEFAULT_ARCHIVE_FORMAT
     ) -> Tuple[Artifact, dict]:
         """
         Pack a path string value content (pack the file or directory in that path).
@@ -138,9 +136,11 @@ class StrPackager(DefaultPackager):
 
         return artifact, instructions
 
-    @classmethod
     def unpack_path(
-        cls, data_item: DataItem, is_directory: bool = False, archive_format: str = None
+        self,
+        data_item: DataItem,
+        is_directory: bool = False,
+        archive_format: str = None,
     ) -> str:
         """
         Unpack a data item representing a path string. If the path is of a file, the file is downloaded to a local
@@ -155,11 +155,8 @@ class StrPackager(DefaultPackager):
 
         :return: The unpacked string.
         """
-        # Get the file to a local temporary directory:
-        path = data_item.local()
-
-        # Mark the downloaded file for future clear:
-        cls.add_future_clearing_path(path=path)
+        # Get the file:
+        path = self.get_data_item_local_path(data_item=data_item)
 
         # If it's not a directory, return the file path. Otherwise, it should be extracted according to the archive
         # format:
@@ -182,7 +179,7 @@ class StrPackager(DefaultPackager):
         )
 
         # Mark the extracted content for future clear:
-        cls.add_future_clearing_path(path=directory_path)
+        self.add_future_clearing_path(path=directory_path)
 
         # Return the extracted directory path:
         return directory_path
@@ -196,9 +193,8 @@ class _BuiltinCollectionPackager(DefaultPackager):
     DEFAULT_PACKING_ARTIFACT_TYPE = ArtifactType.RESULT
     DEFAULT_UNPACKING_ARTIFACT_TYPE = ArtifactType.FILE
 
-    @classmethod
     def pack_file(
-        cls,
+        self,
         obj: Union[dict, list],
         key: str,
         file_format: str = DEFAULT_STRUCT_FILE_FORMAT,
@@ -215,7 +211,7 @@ class _BuiltinCollectionPackager(DefaultPackager):
         # Write to file:
         formatter = StructFileSupportedFormat.get_format_handler(fmt=file_format)
         temp_directory = pathlib.Path(tempfile.mkdtemp())
-        cls.add_future_clearing_path(path=temp_directory)
+        self.add_future_clearing_path(path=temp_directory)
         file_path = temp_directory / f"{key}.{file_format}"
         formatter.write(obj=obj, file_path=str(file_path))
 
@@ -225,9 +221,8 @@ class _BuiltinCollectionPackager(DefaultPackager):
 
         return artifact, instructions
 
-    @classmethod
     def unpack_file(
-        cls, data_item: DataItem, file_format: str = None
+        self, data_item: DataItem, file_format: str = None
     ) -> Union[dict, list]:
         """
         Unpack a builtin collection from file.
@@ -239,8 +234,7 @@ class _BuiltinCollectionPackager(DefaultPackager):
         :return: The unpacked builtin collection.
         """
         # Get the file:
-        file_path = data_item.local()
-        cls.add_future_clearing_path(path=file_path)
+        file_path = self.get_data_item_local_path(data_item=data_item)
 
         # Get the archive format by the file extension if needed:
         if file_format is None:
@@ -265,8 +259,7 @@ class DictPackager(_BuiltinCollectionPackager):
 
     PACKABLE_OBJECT_TYPE = dict
 
-    @classmethod
-    def unpack_file(cls, data_item: DataItem, file_format: str = None) -> dict:
+    def unpack_file(self, data_item: DataItem, file_format: str = None) -> dict:
         """
         Unpack a dictionary from file.
 
@@ -292,8 +285,7 @@ class ListPackager(_BuiltinCollectionPackager):
 
     PACKABLE_OBJECT_TYPE = list
 
-    @classmethod
-    def unpack_file(cls, data_item: DataItem, file_format: str = None) -> list:
+    def unpack_file(self, data_item: DataItem, file_format: str = None) -> list:
         """
         Unpack a list from file.
 
@@ -338,8 +330,7 @@ class TuplePackager(ListPackager):
 
     PACKABLE_OBJECT_TYPE = tuple
 
-    @classmethod
-    def pack_result(cls, obj: tuple, key: str) -> dict:
+    def pack_result(self, obj: tuple, key: str) -> dict:
         """
         Pack a tuple as a result.
 
@@ -350,9 +341,8 @@ class TuplePackager(ListPackager):
         """
         return super().pack_result(obj=list(obj), key=key)
 
-    @classmethod
     def pack_file(
-        cls, obj: tuple, key: str, file_format: str = DEFAULT_STRUCT_FILE_FORMAT
+        self, obj: tuple, key: str, file_format: str = DEFAULT_STRUCT_FILE_FORMAT
     ) -> Tuple[Artifact, dict]:
         """
         Pack a tuple as a file by the given format.
@@ -365,8 +355,7 @@ class TuplePackager(ListPackager):
         """
         return super().pack_file(obj=list(obj), key=key, file_format=file_format)
 
-    @classmethod
-    def unpack_file(cls, data_item: DataItem, file_format: str = None) -> tuple:
+    def unpack_file(self, data_item: DataItem, file_format: str = None) -> tuple:
         """
         Unpack a tuple from file.
 
@@ -386,8 +375,7 @@ class SetPackager(ListPackager):
 
     PACKABLE_OBJECT_TYPE = set
 
-    @classmethod
-    def pack_result(cls, obj: set, key: str) -> dict:
+    def pack_result(self, obj: set, key: str) -> dict:
         """
         Pack a set as a result.
 
@@ -398,9 +386,8 @@ class SetPackager(ListPackager):
         """
         return super().pack_result(obj=list(obj), key=key)
 
-    @classmethod
     def pack_file(
-        cls, obj: set, key: str, file_format: str = DEFAULT_STRUCT_FILE_FORMAT
+        self, obj: set, key: str, file_format: str = DEFAULT_STRUCT_FILE_FORMAT
     ) -> Tuple[Artifact, dict]:
         """
         Pack a set as a file by the given format.
@@ -413,8 +400,7 @@ class SetPackager(ListPackager):
         """
         return super().pack_file(obj=list(obj), key=key, file_format=file_format)
 
-    @classmethod
-    def unpack_file(cls, data_item: DataItem, file_format: str = None) -> set:
+    def unpack_file(self, data_item: DataItem, file_format: str = None) -> set:
         """
         Unpack a set from file.
 
@@ -434,9 +420,8 @@ class FrozensetPackager(SetPackager):
 
     PACKABLE_OBJECT_TYPE = frozenset
 
-    @classmethod
     def pack_file(
-        cls, obj: frozenset, key: str, file_format: str = DEFAULT_STRUCT_FILE_FORMAT
+        self, obj: frozenset, key: str, file_format: str = DEFAULT_STRUCT_FILE_FORMAT
     ) -> Tuple[Artifact, dict]:
         """
         Pack a frozenset as a file by the given format.
@@ -449,8 +434,7 @@ class FrozensetPackager(SetPackager):
         """
         return super().pack_file(obj=set(obj), key=key, file_format=file_format)
 
-    @classmethod
-    def unpack_file(cls, data_item: DataItem, file_format: str = None) -> frozenset:
+    def unpack_file(self, data_item: DataItem, file_format: str = None) -> frozenset:
         """
         Unpack a frozenset from file.
 
@@ -472,8 +456,7 @@ class BytesPackager(ListPackager):
 
     PACKABLE_OBJECT_TYPE = bytes
 
-    @classmethod
-    def pack_result(cls, obj: bytes, key: str) -> dict:
+    def pack_result(self, obj: bytes, key: str) -> dict:
         """
         Pack bytes as a result.
 
@@ -484,9 +467,8 @@ class BytesPackager(ListPackager):
         """
         return {key: obj}
 
-    @classmethod
     def pack_file(
-        cls, obj: bytes, key: str, file_format: str = DEFAULT_STRUCT_FILE_FORMAT
+        self, obj: bytes, key: str, file_format: str = DEFAULT_STRUCT_FILE_FORMAT
     ) -> Tuple[Artifact, dict]:
         """
         Pack a bytes as a file by the given format.
@@ -499,8 +481,7 @@ class BytesPackager(ListPackager):
         """
         return super().pack_file(obj=list(obj), key=key, file_format=file_format)
 
-    @classmethod
-    def unpack_file(cls, data_item: DataItem, file_format: str = None) -> bytes:
+    def unpack_file(self, data_item: DataItem, file_format: str = None) -> bytes:
         """
         Unpack a bytes from file.
 
@@ -520,8 +501,7 @@ class BytearrayPackager(BytesPackager):
 
     PACKABLE_OBJECT_TYPE = bytearray
 
-    @classmethod
-    def pack_result(cls, obj: bytearray, key: str) -> dict:
+    def pack_result(self, obj: bytearray, key: str) -> dict:
         """
         Pack a bytearray as a result.
 
@@ -532,9 +512,8 @@ class BytearrayPackager(BytesPackager):
         """
         return {key: bytes(obj)}
 
-    @classmethod
     def pack_file(
-        cls, obj: bytearray, key: str, file_format: str = DEFAULT_STRUCT_FILE_FORMAT
+        self, obj: bytearray, key: str, file_format: str = DEFAULT_STRUCT_FILE_FORMAT
     ) -> Tuple[Artifact, dict]:
         """
         Pack a bytearray as a file by the given format.
@@ -547,8 +526,7 @@ class BytearrayPackager(BytesPackager):
         """
         return super().pack_file(obj=bytes(obj), key=key, file_format=file_format)
 
-    @classmethod
-    def unpack_file(cls, data_item: DataItem, file_format: str = None) -> bytearray:
+    def unpack_file(self, data_item: DataItem, file_format: str = None) -> bytearray:
         """
         Unpack a bytearray from file.
 
@@ -578,8 +556,7 @@ class PathPackager(StrPackager):
     PACK_SUBCLASSES = True
     DEFAULT_PACKING_ARTIFACT_TYPE = "path"
 
-    @classmethod
-    def pack_result(cls, obj: pathlib.Path, key: str) -> dict:
+    def pack_result(self, obj: pathlib.Path, key: str) -> dict:
         """
         Pack the `Path` as a string result.
 
@@ -590,9 +567,8 @@ class PathPackager(StrPackager):
         """
         return super().pack_result(obj=str(obj), key=key)
 
-    @classmethod
     def pack_path(
-        cls, obj: pathlib.Path, key: str, archive_format: str = DEFAULT_ARCHIVE_FORMAT
+        self, obj: pathlib.Path, key: str, archive_format: str = DEFAULT_ARCHIVE_FORMAT
     ) -> Tuple[Artifact, dict]:
         """
         Pack a `Path` value (pack the file or directory in that path).
@@ -605,9 +581,11 @@ class PathPackager(StrPackager):
         """
         return super().pack_path(obj=str(obj), key=key, archive_format=archive_format)
 
-    @classmethod
     def unpack_path(
-        cls, data_item: DataItem, is_directory: bool = False, archive_format: str = None
+        self,
+        data_item: DataItem,
+        is_directory: bool = False,
+        archive_format: str = None,
     ) -> pathlib.Path:
         """
         Unpack a data item representing a `Path`. If the path is of a file, the file is downloaded to a local

mlrun/package/packagers_manager.py

@@ -51,10 +51,10 @@ class PackagersManager:
                                  object or data item. Default to ``mlrun.DefaultPackager``.
         """
         # Set the default packager:
-        self._default_packager = default_packager or DefaultPackager
+        self._default_packager = (default_packager or DefaultPackager)()
 
         # Initialize the packagers list (with the default packager in it):
-        self._packagers: List[Type[Packager]] = []
+        self._packagers: List[Packager] = []
 
         # Set an artifacts list and results dictionary to collect all packed objects (will be used later to write extra
         # data if noted by the user using the log hint key "extra_data")
@@ -80,7 +80,7 @@ class PackagersManager:
         return self._results
 
     def collect_packagers(
-        self, packagers: List[Union[Type, str]], default_priority: int = 5
+        self, packagers: List[Union[Type[Packager], str]], default_priority: int = 5
     ):
         """
         Collect the provided packagers. Packagers passed as module paths are imported and validated to be of type
@@ -155,9 +155,11 @@ class PackagersManager:
                 raise MLRunPackageCollectionError(
                     f"The packager '{packager.__name__}' could not be collected as it is not a `mlrun.Packager`."
                 )
+            # Initialize the packager class:
+            packager = packager()
             # Set default priority in case it is not set in the packager's class:
-            if packager.PRIORITY is ...:
-                packager.PRIORITY = default_priority
+            if packager.priority is ...:
+                packager.priority = default_priority
             # Collect the packager (putting him first in the list for highest priority:
             self._packagers.insert(0, packager)
             # For debugging, we'll print the collected packager:
@@ -350,13 +352,14 @@ class PackagersManager:
         artifacts, to ensure that files that require uploading have already been uploaded.
         """
         for packager in self._get_packagers_with_default_packager():
-            for path in packager.get_future_clearing_path_list():
+            for path in packager.future_clearing_path_list:
                 if not os.path.exists(path):
                     continue
                 if os.path.isdir(path):
                     shutil.rmtree(path)
                 else:
                     os.remove(path)
+            packager.future_clearing_path_list.clear()
 
     class _InstructionsNotesKey:
         """
@@ -368,7 +371,7 @@ class PackagersManager:
         ARTIFACT_TYPE = "artifact_type"
         INSTRUCTIONS = "instructions"
 
-    def _get_packagers_with_default_packager(self) -> List[Type[Packager]]:
+    def _get_packagers_with_default_packager(self) -> List[Packager]:
         """
         Get the full list of packagers - the collected packagers and the default packager (located at last place in the
         list - the lowest priority).
@@ -377,7 +380,7 @@ class PackagersManager:
         """
         return [*self._packagers, self._default_packager]
 
-    def _get_packager_by_name(self, name: str) -> Union[Type[Packager], None]:
+    def _get_packager_by_name(self, name: str) -> Union[Packager, None]:
         """
         Look for a packager with the given name and return it.
 
@@ -389,7 +392,7 @@ class PackagersManager:
         """
         # Look for a packager by exact name:
         for packager in self._get_packagers_with_default_packager():
-            if packager.__name__ == name:
+            if packager.__class__.__name__ == name:
                 return packager
 
         # No packager was found:
@@ -401,7 +404,7 @@ class PackagersManager:
         obj: Any,
         artifact_type: str = None,
         configurations: dict = None,
-    ) -> Union[Type[Packager], None]:
+    ) -> Union[Packager, None]:
         """
         Look for a packager that can pack the provided object as the provided artifact type.
 
@@ -428,7 +431,7 @@ class PackagersManager:
         data_item: Any,
         type_hint: type,
         artifact_type: str = None,
-    ) -> Union[Type[Packager], None]:
+    ) -> Union[Packager, None]:
         """
         Look for a packager that can unpack the data item of the given type hint as the provided artifact type.
 
@@ -495,7 +498,7 @@ class PackagersManager:
 
         # Prepare the manager's unpackaging instructions:
         unpackaging_instructions = {
-            self._InstructionsNotesKey.PACKAGER_NAME: packager.__name__,
+            self._InstructionsNotesKey.PACKAGER_NAME: packager.__class__.__name__,
             self._InstructionsNotesKey.OBJECT_TYPE: self._get_type_name(typ=type(obj)),
             self._InstructionsNotesKey.ARTIFACT_TYPE: (
                 artifact_type
@@ -646,7 +649,7 @@ class PackagersManager:
         :raise MLRunPackageUnpackingError: If there is no packager that supports the provided type hint.
         """
         # Prepare a list of a packager and exception string for all the failures in case there was no fitting packager:
-        found_packagers: List[Tuple[Type[Packager], str]] = []
+        found_packagers: List[Tuple[Packager, str]] = []
 
         # Try to unpack as one of the possible types in the type hint:
         possible_type_hints = {type_hint}

mlrun/package/utils/_pickler.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 #
 import importlib
+import importlib.metadata as importlib_metadata
 import os
 import sys
 import tempfile
@@ -188,26 +189,15 @@ class Pickler:
         """
         # First we'll try to get the module version from `importlib`:
         try:
-            # Since Python 3.8, `version` is part of `importlib.metadata`. Before 3.8, we'll use the module
-            # `importlib_metadata` to get `version`.
-            if (
-                sys.version_info[1] > 7
-            ):  # TODO: Remove once Python 3.7 is not supported.
-                from importlib.metadata import version
-            else:
-                from importlib_metadata import version
-
-            return version(module_name)
-        except (ModuleNotFoundError, importlib.metadata.PackageNotFoundError):
-            # User won't necessarily have the `importlib_metadata` module, so we will ignore it by catching
-            # `ModuleNotFoundError`. `PackageNotFoundError` is ignored as well as this is raised when `version` could
-            # not find the package related to the module.
+            return importlib_metadata.version(module_name)
+        except importlib.metadata.PackageNotFoundError:
+            # `PackageNotFoundError` is ignored this is raised when `version` could not find the package related to the
+            # module.
             pass
 
-        # Secondly, if importlib could not get the version (most likely 'importlib_metadata' is not installed), we'll
-        # try to use `pkg_resources` to get the version (the version will be found only if the package name is equal to
-        # the module name. For example, if the module name is 'x' then the way we installed the package must be
-        # 'pip install x'):
+        # Secondly, if importlib could not get the version, we'll try to use `pkg_resources` to get the version (the
+        # version will be found only if the package name is equal to the module name. For example, if the module name is
+        # 'x' then the way we installed the package must be 'pip install x'):
         import pkg_resources
 
         with warnings.catch_warnings():

mlrun/package/utils/_supported_format.py

@@ -24,7 +24,7 @@ class SupportedFormat(ABC, Generic[FileHandlerType]):
     Library of supported formats by some builtin MLRun packagers.
     """
 
-    # Add here the all the supported formats in ALL CAPS and their value as a string:
+    # Add here all the supported formats in ALL CAPS and their value as a string:
     ...
 
     # The map to use in the method `get_format_handler`. A dictionary of string key to a class type to handle that

mlrun/projects/pipelines.py

@@ -133,7 +133,7 @@ class WorkflowSpec(mlrun.model.ModelObj):
                     required.remove(k)
         if required:
             raise mlrun.errors.MLRunInvalidArgumentError(
-                f"workflow argument(s) {','.join(required)} are required and were not specified"
+                f"Workflow argument(s) {','.join(required)} are required and were not specified"
             )
 
     def clear_tmp(self):
@@ -246,7 +246,7 @@ class _PipelineContext:
             return True
         if raise_exception:
             raise ValueError(
-                "pipeline context is not initialized, must be used inside a pipeline"
+                "Pipeline context is not initialized, must be used inside a pipeline"
             )
         return False
 
@@ -496,6 +496,7 @@ class _PipelineRunner(abc.ABC):
         artifact_path=None,
         namespace=None,
         source=None,
+        notifications: typing.List[mlrun.model.Notification] = None,
     ) -> _PipelineRunStatus:
         pass
 
@@ -573,6 +574,7 @@ class _KFPRunner(_PipelineRunner):
         artifact_path=None,
         namespace=None,
         source=None,
+        notifications: typing.List[mlrun.model.Notification] = None,
     ) -> _PipelineRunStatus:
         pipeline_context.set(project, workflow_spec)
         workflow_handler = _PipelineRunner._get_handler(
@@ -582,6 +584,19 @@ class _KFPRunner(_PipelineRunner):
             project.set_source(source=source)
 
         namespace = namespace or config.namespace
+
+        # fallback to old notification behavior
+        if notifications:
+            logger.warning(
+                "Setting notifications on kfp pipeline runner uses old notification behavior. "
+                "Notifications will only be sent if you wait for pipeline completion. "
+                "To use the new notification behavior, use the remote pipeline runner."
+            )
+            for notification in notifications:
+                project.notifiers.add_notification(
+                    notification.kind, notification.params
+                )
+
         run_id = _run_pipeline(
             workflow_handler,
             project=project.metadata.name,
@@ -655,7 +670,7 @@ class _KFPRunner(_PipelineRunner):
         raise_error = None
         try:
             if timeout:
-                logger.info("waiting for pipeline run completion")
+                logger.info("Waiting for pipeline run completion")
                 state = run.wait_for_completion(
                     timeout=timeout, expected_statuses=expected_statuses
                 )
@@ -666,6 +681,7 @@ class _KFPRunner(_PipelineRunner):
         mldb = mlrun.db.get_run_db(secrets=project._secrets)
         runs = mldb.list_runs(project=project.name, labels=f"workflow={run.run_id}")
 
+        # TODO: The below section duplicates notifiers.push_pipeline_run_results() logic. We should use it instead.
         had_errors = 0
         for r in runs:
             if r["status"].get("state", "") == "error":
@@ -701,12 +717,17 @@ class _LocalRunner(_PipelineRunner):
         artifact_path=None,
         namespace=None,
         source=None,
+        notifications: typing.List[mlrun.model.Notification] = None,
     ) -> _PipelineRunStatus:
         pipeline_context.set(project, workflow_spec)
         workflow_handler = _PipelineRunner._get_handler(
             workflow_handler, workflow_spec, project, secrets
         )
 
+        # fallback to old notification behavior
+        for notification in notifications or []:
+            project.notifiers.add_notification(notification.kind, notification.params)
+
         workflow_id = uuid.uuid4().hex
         pipeline_context.workflow_id = workflow_id
         # When using KFP, it would do this replacement. When running locally, we need to take care of it.
@@ -785,10 +806,15 @@ class _RemoteRunner(_PipelineRunner):
         artifact_path: str = None,
         namespace: str = None,
         source: str = None,
+        notifications: typing.List[mlrun.model.Notification] = None,
     ) -> typing.Optional[_PipelineRunStatus]:
         workflow_name = normalize_workflow_name(name=name, project_name=project.name)
         workflow_id = None
 
+        # for start message, fallback to old notification behavior
+        for notification in notifications or []:
+            project.notifiers.add_notification(notification.kind, notification.params)
+
         # The returned engine for this runner is the engine of the workflow.
         # In this way wait_for_completion/get_run_status would be executed by the correct pipeline runner.
         inner_engine = get_workflow_engine(workflow_spec.engine)
@@ -827,6 +853,7 @@ class _RemoteRunner(_PipelineRunner):
                     workflow_name
                 ),
                 namespace=namespace,
+                notifications=notifications,
             )
             if workflow_spec.schedule:
                 logger.info(
@@ -873,11 +900,33 @@ class _RemoteRunner(_PipelineRunner):
             exc=err,
         )
 
+    @staticmethod
+    def get_run_status(
+        project,
+        run,
+        timeout=None,
+        expected_statuses=None,
+        notifiers: mlrun.utils.notifications.CustomNotificationPusher = None,
+    ):
+        # ignore notifiers, as they are handled by the remote pipeline notifications,
+        # so overriding with CustomNotificationPusher with empty list of notifiers
+        state, had_errors, text = _KFPRunner.get_run_status(
+            project,
+            run,
+            timeout,
+            expected_statuses,
+            notifiers=mlrun.utils.notifications.CustomNotificationPusher([]),
+        )
+
+        # indicate the pipeline status since we don't push the notifications in the remote runner
+        logger.info(text)
+        return state, had_errors, text
+
 
 def create_pipeline(project, pipeline, functions, secrets=None, handler=None):
     spec = imputil.spec_from_file_location("workflow", pipeline)
     if spec is None:
-        raise ImportError(f"cannot import workflow {pipeline}")
+        raise ImportError(f"Cannot import workflow {pipeline}")
     mod = imputil.module_from_spec(spec)
     spec.loader.exec_module(mod)
 
@@ -938,6 +987,7 @@ def load_and_run(
     schedule: typing.Union[str, mlrun.common.schemas.ScheduleCronTrigger] = None,
     cleanup_ttl: int = None,
     load_only: bool = False,
+    wait_for_completion: bool = False,
 ):
     """
     Auxiliary function that the RemoteRunner run once or run every schedule.
@@ -967,6 +1017,7 @@ def load_and_run(
     :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
     """
     try:
         project = mlrun.load_project(
@@ -1028,3 +1079,11 @@ def load_and_run(
 
     if run.state == mlrun.run.RunStatuses.failed:
         raise RuntimeError(f"Workflow {workflow_log_message} failed") from run.exc
+
+    if wait_for_completion:
+        pipeline_state, _, _ = project.get_run_status(run)
+        context.log_result(key="workflow_state", value=pipeline_state, commit=True)
+        if pipeline_state != mlrun.run.RunStatuses.succeeded:
+            raise RuntimeError(
+                f"Workflow {workflow_log_message} failed, state={pipeline_state}"
+            )