mlrun 1.7.0rc58__py3-none-any.whl → 1.7.1__py3-none-any.whl

mlrun/projects/pipelines.py

@@ -984,14 +984,25 @@ def github_webhook(request):
     return {"msg": "pushed"}
 
 
-def load_and_run(
+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(
     context: mlrun.execution.MLClientCtx,
     url: str = None,
     project_name: str = "",
     init_git: bool = None,
     subpath: str = None,
     clone: bool = False,
-    save: bool = True,
     workflow_name: str = None,
     workflow_path: str = None,
     workflow_arguments: dict[str, typing.Any] = None,
@@ -1004,14 +1015,12 @@ def load_and_run(
     local: bool = None,
     schedule: typing.Union[str, mlrun.common.schemas.ScheduleCronTrigger] = None,
     cleanup_ttl: int = None,
-    load_only: bool = False,
     wait_for_completion: bool = False,
     project_context: 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
@@ -1019,7 +1028,6 @@ def load_and_run(
     :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)
@@ -1035,48 +1043,31 @@ def load_and_run(
     :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)
     """
-    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,
-                )
-
-            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")
+    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,
+    )
 
-    if load_only:
-        return
+    # 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,
+    )
 
     # extract "start" notification if exists
     start_notifications = [
@@ -1109,18 +1100,156 @@ def load_and_run(
         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:
-            run.wait_for_completion()
-        except Exception as exc:
-            logger.error(
-                "Failed waiting for workflow completion",
-                workflow=workflow_log_message,
-                exc=err_to_str(exc),
+            notification_pusher.push(
+                message=message,
+                severity=mlrun.common.schemas.NotificationSeverity.ERROR,
             )
 
-        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}"
-            )
+        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,
+    )
+
+    context.logger.info(f"Loaded project {project.name} successfully")

mlrun/projects/project.py

@@ -40,6 +40,7 @@ import requests
 import yaml
 from mlrun_pipelines.models import PipelineNodeWrapper
 
+import mlrun.common.formatters
 import mlrun.common.helpers
 import mlrun.common.runtimes.constants
 import mlrun.common.schemas.artifact
@@ -47,6 +48,7 @@ import mlrun.common.schemas.model_monitoring.constants as mm_constants
 import mlrun.db
 import mlrun.errors
 import mlrun.k8s_utils
+import mlrun.lists
 import mlrun.model_monitoring.applications as mm_app
 import mlrun.runtimes
 import mlrun.runtimes.nuclio.api_gateway
@@ -3799,6 +3801,9 @@ class MlrunProject(ModelObj):
         category: typing.Union[str, mlrun.common.schemas.ArtifactCategories] = None,
         tree: str = None,
         limit: int = None,
+        format_: Optional[
+            mlrun.common.formatters.ArtifactFormat
+        ] = mlrun.common.formatters.ArtifactFormat.full,
     ) -> mlrun.lists.ArtifactList:
         """List artifacts filtered by various parameters.
 
@@ -3829,6 +3834,7 @@ class MlrunProject(ModelObj):
         :param category: Return artifacts of the requested category.
         :param tree: Return artifacts of the requested tree.
         :param limit: Maximum number of artifacts to return.
+        :param format_: The format in which to return the artifacts. Default is 'full'.
         """
         db = mlrun.db.get_run_db(secrets=self._secrets)
         return db.list_artifacts(
@@ -3843,6 +3849,7 @@ class MlrunProject(ModelObj):
             kind=kind,
             category=category,
             tree=tree,
+            format_=format_,
             limit=limit,
         )
 
@@ -3856,6 +3863,10 @@ class MlrunProject(ModelObj):
         iter: int = None,
         best_iteration: bool = False,
         tree: str = None,
+        limit: int = None,
+        format_: Optional[
+            mlrun.common.formatters.ArtifactFormat
+        ] = mlrun.common.formatters.ArtifactFormat.full,
     ):
         """List models in project, filtered by various parameters.
 
@@ -3879,6 +3890,8 @@ class MlrunProject(ModelObj):
             artifacts generated from a hyper-param run. If only a single iteration exists, will return the artifact
             from that iteration. If using ``best_iter``, the ``iter`` parameter must not be used.
         :param tree: Return artifacts of the requested tree.
+        :param limit: Maximum number of artifacts to return.
+        :param format_: The format in which to return the artifacts. Default is 'full'.
         """
         db = mlrun.db.get_run_db(secrets=self._secrets)
         return db.list_artifacts(
@@ -3892,6 +3905,8 @@ class MlrunProject(ModelObj):
             best_iteration=best_iteration,
             kind="model",
             tree=tree,
+            limit=limit,
+            format_=format_,
         ).to_objects()
 
     def list_functions(self, name=None, tag=None, labels=None):

mlrun/runtimes/nuclio/serving.py

@@ -607,7 +607,7 @@ class ServingRuntime(RemoteRuntime):
         ):
             # initialize or create required streams/queues
             self.spec.graph.check_and_process_graph()
-            self.spec.graph.init_queues()
+            self.spec.graph.create_queue_streams()
             functions_in_steps = self.spec.graph.list_child_functions()
             child_functions = list(self._spec.function_refs.keys())
             for function in functions_in_steps:

mlrun/serving/routers.py

@@ -491,6 +491,7 @@ class VotingEnsemble(ParallelRun):
         executor_type: Union[ParallelRunnerModes, str] = ParallelRunnerModes.thread,
         format_response_with_col_name_flag: bool = False,
         prediction_col_name: str = "prediction",
+        shard_by_endpoint: typing.Optional[bool] = None,
         **kwargs,
     ):
         """Voting Ensemble
@@ -580,6 +581,8 @@ class VotingEnsemble(ParallelRun):
                               `{id: <id>, model_name: <name>, outputs: {..., prediction: [<predictions>], ...}}`
                               the prediction_col_name should be `prediction`.
                               by default, `prediction`
+        :param shard_by_endpoint: whether to use the endpoint as the partition/sharding key when writing to model
+                                  monitoring stream. Defaults to True.
         :param kwargs:        extra arguments
         """
         super().__init__(
@@ -606,6 +609,7 @@ class VotingEnsemble(ParallelRun):
         self.prediction_col_name = prediction_col_name or "prediction"
         self.format_response_with_col_name_flag = format_response_with_col_name_flag
         self.model_endpoint_uid = None
+        self.shard_by_endpoint = shard_by_endpoint
 
     def post_init(self, mode="sync"):
         server = getattr(self.context, "_server", None) or getattr(
@@ -907,7 +911,12 @@ class VotingEnsemble(ParallelRun):
         if self._model_logger and self.log_router:
             if "id" not in request:
                 request["id"] = response.body["id"]
-            self._model_logger.push(start, request, response.body)
+            partition_key = (
+                self.model_endpoint_uid if self.shard_by_endpoint is not False else None
+            )
+            self._model_logger.push(
+                start, request, response.body, partition_key=partition_key
+            )
         event.body = _update_result_body(
             self._result_path, original_body, response.body if response else None
         )

mlrun/serving/states.py

@@ -839,6 +839,8 @@ class QueueStep(BaseStep):
                 retention_in_hours=self.retention_in_hours,
                 **self.options,
             )
+            if hasattr(self._stream, "create_stream"):
+                self._stream.create_stream()
         self._set_error_handler()
 
     @property
@@ -1247,8 +1249,8 @@ class FlowStep(BaseStep):
                     links[next_step.function] = step
         return links
 
-    def init_queues(self):
-        """init/create the streams used in this flow"""
+    def create_queue_streams(self):
+        """create the streams used in this flow"""
         for step in self.get_children():
             if step.kind == StepKinds.queue:
                 step.init_object(self.context, None)

mlrun/serving/v2_serving.py

@@ -39,6 +39,7 @@ class V2ModelServer(StepToDict):
         protocol=None,
         input_path: str = None,
         result_path: str = None,
+        shard_by_endpoint: Optional[bool] = None,
         **kwargs,
     ):
         """base model serving class (v2), using similar API to KFServing v2 and Triton
@@ -91,6 +92,8 @@ class V2ModelServer(StepToDict):
                               this require that the event body will behave like a dict, example:
                               event: {"x": 5} , result_path="resp" means the returned response will be written
                               to event["y"] resulting in {"x": 5, "resp": <result>}
+        :param shard_by_endpoint: whether to use the endpoint as the partition/sharding key when writing to model
+                                  monitoring stream. Defaults to True.
         :param kwargs:     extra arguments (can be accessed using self.get_param(key))
         """
         self.name = name
@@ -119,7 +122,9 @@ class V2ModelServer(StepToDict):
         if model:
             self.model = model
             self.ready = True
+        self._versioned_model_name = None
         self.model_endpoint_uid = None
+        self.shard_by_endpoint = shard_by_endpoint
 
     def _load_and_update_state(self):
         try:
@@ -225,6 +230,23 @@ class V2ModelServer(StepToDict):
         request = self.preprocess(event_body, op)
         return self.validate(request, op)
 
+    @property
+    def versioned_model_name(self):
+        if self._versioned_model_name:
+            return self._versioned_model_name
+
+        # Generating version model value based on the model name and model version
+        if self.model_path and self.model_path.startswith("store://"):
+            # Enrich the model server with the model artifact metadata
+            self.get_model()
+            if not self.version:
+                # Enrich the model version with the model artifact tag
+                self.version = self.model_spec.tag
+            self.labels = self.model_spec.labels
+        version = self.version or "latest"
+        self._versioned_model_name = f"{self.name}:{version}"
+        return self._versioned_model_name
+
     def do_event(self, event, *args, **kwargs):
         """main model event handler method"""
         start = now_date()
@@ -232,6 +254,11 @@ class V2ModelServer(StepToDict):
         event_body = _extract_input_data(self._input_path, event.body)
         event_id = event.id
         op = event.path.strip("/")
+
+        partition_key = (
+            self.model_endpoint_uid if self.shard_by_endpoint is not False else None
+        )
+
         if event_body and isinstance(event_body, dict):
             op = op or event_body.get("operation")
             event_id = event_body.get("id", event_id)
@@ -251,7 +278,13 @@ class V2ModelServer(StepToDict):
             except Exception as exc:
                 request["id"] = event_id
                 if self._model_logger:
-                    self._model_logger.push(start, request, op=op, error=exc)
+                    self._model_logger.push(
+                        start,
+                        request,
+                        op=op,
+                        error=exc,
+                        partition_key=partition_key,
+                    )
                 raise exc
 
             response = {
@@ -288,7 +321,7 @@ class V2ModelServer(StepToDict):
             setattr(event, "terminated", True)
             event_body = {
                 "name": self.name,
-                "version": self.version,
+                "version": self.version or "",
                 "inputs": [],
                 "outputs": [],
             }
@@ -308,7 +341,13 @@ class V2ModelServer(StepToDict):
             except Exception as exc:
                 request["id"] = event_id
                 if self._model_logger:
-                    self._model_logger.push(start, request, op=op, error=exc)
+                    self._model_logger.push(
+                        start,
+                        request,
+                        op=op,
+                        error=exc,
+                        partition_key=partition_key,
+                    )
                 raise exc
 
             response = {
@@ -332,12 +371,20 @@ class V2ModelServer(StepToDict):
         if self._model_logger:
             inputs, outputs = self.logged_results(request, response, op)
             if inputs is None and outputs is None:
-                self._model_logger.push(start, request, response, op)
+                self._model_logger.push(
+                    start, request, response, op, partition_key=partition_key
+                )
             else:
                 track_request = {"id": event_id, "inputs": inputs or []}
                 track_response = {"outputs": outputs or []}
                 # TODO : check dict/list
-                self._model_logger.push(start, track_request, track_response, op)
+                self._model_logger.push(
+                    start,
+                    track_request,
+                    track_response,
+                    op,
+                    partition_key=partition_key,
+                )
         event.body = _update_result_body(self._result_path, original_body, response)
         return event
 
@@ -454,7 +501,7 @@ class _ModelLogPusher:
             base_data["labels"] = self.model.labels
         return base_data
 
-    def push(self, start, request, resp=None, op=None, error=None):
+    def push(self, start, request, resp=None, op=None, error=None, partition_key=None):
         start_str = start.isoformat(sep=" ", timespec="microseconds")
         if error:
             data = self.base_data()
@@ -465,7 +512,7 @@ class _ModelLogPusher:
             if self.verbose:
                 message = f"{message}\n{traceback.format_exc()}"
             data["error"] = message
-            self.output_stream.push([data])
+            self.output_stream.push([data], partition_key=partition_key)
             return
 
         self._sample_iter = (self._sample_iter + 1) % self.stream_sample
@@ -491,7 +538,7 @@ class _ModelLogPusher:
                         "metrics",
                     ]
                     data["values"] = self._batch
-                    self.output_stream.push([data])
+                    self.output_stream.push([data], partition_key=partition_key)
             else:
                 data = self.base_data()
                 data["request"] = request
@@ -501,7 +548,7 @@ class _ModelLogPusher:
                 data["microsec"] = microsec
                 if getattr(self.model, "metrics", None):
                     data["metrics"] = self.model.metrics
-                self.output_stream.push([data])
+                self.output_stream.push([data], partition_key=partition_key)
 
 
 def _init_endpoint_record(
@@ -531,21 +578,10 @@ def _init_endpoint_record(
         logger.error("Failed to parse function URI", exc=err_to_str(e))
         return None
 
-    # Generating version model value based on the model name and model version
-    if model.model_path and model.model_path.startswith("store://"):
-        # Enrich the model server with the model artifact metadata
-        model.get_model()
-        if not model.version:
-            # Enrich the model version with the model artifact tag
-            model.version = model.model_spec.tag
-        model.labels = model.model_spec.labels
-        versioned_model_name = f"{model.name}:{model.version}"
-    else:
-        versioned_model_name = f"{model.name}:latest"
-
     # Generating model endpoint ID based on function uri and model version
     uid = mlrun.common.model_monitoring.create_model_endpoint_uid(
-        function_uri=graph_server.function_uri, versioned_model=versioned_model_name
+        function_uri=graph_server.function_uri,
+        versioned_model=model.versioned_model_name,
     ).uid
 
     try:
@@ -568,7 +604,7 @@ def _init_endpoint_record(
             ),
             spec=mlrun.common.schemas.ModelEndpointSpec(
                 function_uri=graph_server.function_uri,
-                model=versioned_model_name,
+                model=model.versioned_model_name,
                 model_class=model.__class__.__name__,
                 model_uri=model.model_path,
                 stream_path=model.context.stream.stream_uri,

mlrun/utils/helpers.py

@@ -1226,14 +1226,24 @@ def datetime_to_iso(time_obj: Optional[datetime]) -> Optional[str]:
     return time_obj.isoformat()
 
 
-def enrich_datetime_with_tz_info(timestamp_string):
+def enrich_datetime_with_tz_info(timestamp_string) -> Optional[datetime]:
     if not timestamp_string:
         return timestamp_string
 
     if timestamp_string and not mlrun.utils.helpers.has_timezone(timestamp_string):
         timestamp_string += datetime.now(timezone.utc).astimezone().strftime("%z")
 
-    return datetime.strptime(timestamp_string, "%Y-%m-%d %H:%M:%S.%f%z")
+    for _format in [
+        # e.g: 2021-08-25 12:00:00.000Z
+        "%Y-%m-%d %H:%M:%S.%f%z",
+        # e.g: 2024-11-11 07:44:56+0000
+        "%Y-%m-%d %H:%M:%S%z",
+    ]:
+        try:
+            return datetime.strptime(timestamp_string, _format)
+        except ValueError as exc:
+            last_exc = exc
+    raise last_exc
 
 
 def has_timezone(timestamp):

mlrun/utils/version/version.json

@@ -1,4 +1,4 @@
 {
-  "git_commit": "b07f7ed6fa6f05047bece3472b5edc29733efd94",
-  "version": "1.7.0-rc58"
+  "git_commit": "582dbd4c8eef782ec66ab455530c1e3595f95b58",
+  "version": "1.7.1"
 }