mlrun 1.10.0rc9__py3-none-any.whl → 1.10.0rc11__py3-none-any.whl

mlrun/model_monitoring/helpers.py

@@ -589,3 +589,51 @@ def _get_monitoring_schedules_file_chief_path(
     return os.path.join(
         _get_monitoring_schedules_folder_path(project), f"{project}.json"
     )
+
+
+def get_start_end(
+    start: Union[datetime.datetime, None],
+    end: Union[datetime.datetime, None],
+    delta: Optional[datetime.timedelta] = None,
+) -> tuple[datetime.datetime, datetime.datetime]:
+    """
+    static utils function for tsdb start end format
+    :param start:       Either None or datetime, None is handled as datetime.min(tz=timezone.utc) unless `delta`
+                        is provided.
+    :param end:         Either None or datetime, None is handled as datetime.now(tz=timezone.utc)
+    :param delta:       Optional timedelta to define a time span.
+                        - If both `start` and `end` are provided, `delta` is ignored.
+                        - If only one of `start` or `end` is provided, the other will be
+                          calculated using `delta`.
+                        - If neither `start` nor `end` is provided, `end` defaults to now,
+                          and `start` is calculated as `end - delta`.
+    :return:            start datetime, end datetime
+    """
+
+    if delta and start and end:
+        # If both start and end are provided, delta is ignored
+        pass
+    elif delta:
+        if start and not end:
+            end = start + delta
+        else:
+            end = end or mlrun.utils.datetime_now()
+            start = end - delta
+    else:
+        start = start or mlrun.utils.datetime_min()
+        end = end or mlrun.utils.datetime_now()
+
+    if not (
+        isinstance(start, datetime.datetime) and isinstance(end, datetime.datetime)
+    ):
+        raise mlrun.errors.MLRunInvalidArgumentError(
+            "Both start and end must be datetime objects"
+        )
+
+    if start > end:
+        raise mlrun.errors.MLRunInvalidArgumentError(
+            "The start time must be before the end time. Note that if end time is not provided, "
+            "the current time is used by default"
+        )
+
+    return start, end

mlrun/projects/__init__.py

@@ -32,6 +32,7 @@ from .pipelines import (
     load_and_run_workflow,
     load_and_run,
     pipeline_context,
+    rerun_workflow,
 )  # noqa
 from .project import (
     MlrunProject,

mlrun/projects/pipelines.py

@@ -21,6 +21,7 @@ import typing
 import uuid
 
 import mlrun
+import mlrun.common.constants as mlrun_constants
 import mlrun.common.runtimes.constants
 import mlrun.common.schemas
 import mlrun.common.schemas.function
@@ -1070,6 +1071,46 @@ def github_webhook(request):
     return {"msg": "pushed"}
 
 
+def rerun_workflow(
+    context: mlrun.execution.MLClientCtx, run_uid: str, project_name: str
+):
+    """
+    Re-run a workflow by retrying a previously failed KFP pipeline.
+
+    :param context:      MLRun context.
+    :param run_uid:      The run UID of the original workflow to retry.
+    :param project_name: The project name.
+    """
+
+    try:
+        # TODO in followups: handle start and running notifications
+
+        # Retry the pipeline  - TODO: add submit-direct flag when created
+        db = mlrun.get_run_db()
+        new_pipeline_id = db.retry_pipeline(
+            run_uid, project_name, submit_mode=mlrun_constants.WorkflowSubmitMode.direct
+        )
+
+        # Store result for observability
+        context.set_label(
+            mlrun_constants.MLRunInternalLabels.workflow_id, new_pipeline_id
+        )
+        context.update_run()
+
+        context.log_result("workflow_id", new_pipeline_id)
+
+        # wait for pipeline completion so monitor will push terminal notifications
+        wait_for_pipeline_completion(
+            new_pipeline_id,
+            project=project_name,
+        )
+
+    # Temporary exception
+    except Exception as exc:
+        context.logger.error("Failed to rerun workflow", exc=err_to_str(exc))
+        raise
+
+
 def load_and_run(context, *args, **kwargs):
     """
     This function serves as an alias to `load_and_run_workflow`,
@@ -1192,7 +1233,9 @@ def load_and_run_workflow(
     context.logger.info(
         "Associating workflow-runner with workflow ID", run_id=run.run_id
     )
-    context.set_label("workflow-id", run.run_id)
+    context.set_label(mlrun_constants.MLRunInternalLabels.workflow_id, run.run_id)
+    context.update_run()
+
     context.log_result(key="workflow_id", value=run.run_id)
     context.log_result(key="engine", value=run._engine.engine, commit=True)
 

mlrun/projects/project.py

@@ -4968,6 +4968,36 @@ class MlrunProject(ModelObj):
             include_infra=include_infra,
         )
 
+    def get_monitoring_function_summary(
+        self,
+        name: str,
+        start: Optional[datetime.datetime] = None,
+        end: Optional[datetime.datetime] = None,
+        include_latest_metrics: bool = False,
+    ) -> mlrun.common.schemas.model_monitoring.FunctionSummary:
+        """Get a monitoring function summary for the specified project and function name.
+        :param name:                   Name of the monitoring function to retrieve the summary for.
+        :param start:                  Start time for filtering the results (optional).
+        :param end:                    End time for filtering the results (optional).
+        :param include_latest_metrics: Whether to include the latest metrics in the response (default is False).
+
+        :return: A FunctionSummary object containing information about the monitoring function.
+        """
+        if start is not None and end is not None:
+            if start.tzinfo is None or end.tzinfo is None:
+                raise mlrun.errors.MLRunInvalidArgumentTypeError(
+                    "Custom start and end times must contain the timezone."
+                )
+
+        db = mlrun.db.get_run_db(secrets=self._secrets)
+        return db.get_monitoring_function_summary(
+            project=self.metadata.name,
+            function_name=name,
+            start=start,
+            end=end,
+            include_latest_metrics=include_latest_metrics,
+        )
+
     def list_runs(
         self,
         name: Optional[str] = None,

mlrun/runtimes/daskjob.py

@@ -92,6 +92,7 @@ class DaskSpec(KubeResourceSpec):
         preemption_mode=None,
         security_context=None,
         state_thresholds=None,
+        serving_spec=None,
     ):
         super().__init__(
             command=command,
@@ -121,6 +122,7 @@ class DaskSpec(KubeResourceSpec):
             preemption_mode=preemption_mode,
             security_context=security_context,
             state_thresholds=state_thresholds,
+            serving_spec=serving_spec,
         )
         self.args = args
 

mlrun/runtimes/kubejob.py

@@ -207,3 +207,7 @@ class KubejobRuntime(KubeResource):
         raise NotImplementedError(
             f"Running a {self.kind} function from the client is not supported. Use .run() to submit the job to the API."
         )
+
+    @property
+    def serving_spec(self):
+        return self.spec.serving_spec

mlrun/runtimes/mpijob/abstract.py

@@ -54,6 +54,7 @@ class MPIResourceSpec(KubeResourceSpec):
         preemption_mode=None,
         security_context=None,
         state_thresholds=None,
+        serving_spec=None,
     ):
         super().__init__(
             command=command,
@@ -83,6 +84,7 @@ class MPIResourceSpec(KubeResourceSpec):
             preemption_mode=preemption_mode,
             security_context=security_context,
             state_thresholds=state_thresholds,
+            serving_spec=serving_spec,
         )
         self.mpi_args = mpi_args or [
             "-x",

mlrun/runtimes/mpijob/v1.py

@@ -49,6 +49,7 @@ class MPIV1ResourceSpec(MPIResourceSpec):
         preemption_mode=None,
         security_context=None,
         state_thresholds=None,
+        serving_spec=None,
     ):
         super().__init__(
             command=command,
@@ -79,6 +80,7 @@ class MPIV1ResourceSpec(MPIResourceSpec):
             preemption_mode=preemption_mode,
             security_context=security_context,
             state_thresholds=state_thresholds,
+            serving_spec=serving_spec,
         )
         self.clean_pod_policy = clean_pod_policy or MPIJobV1CleanPodPolicies.default()
 

mlrun/runtimes/nuclio/function.py

@@ -154,6 +154,7 @@ class NuclioSpec(KubeResourceSpec):
         add_templated_ingress_host_mode=None,
         state_thresholds=None,
         disable_default_http_trigger=None,
+        serving_spec=None,
     ):
         super().__init__(
             command=command,
@@ -183,6 +184,7 @@ class NuclioSpec(KubeResourceSpec):
             preemption_mode=preemption_mode,
             security_context=security_context,
             state_thresholds=state_thresholds,
+            serving_spec=serving_spec,
         )
 
         self.base_spec = base_spec or {}

mlrun/runtimes/nuclio/serving.py

@@ -42,6 +42,8 @@ from mlrun.serving.states import (
 )
 from mlrun.utils import get_caller_globals, logger, set_paths
 
+from .. import KubejobRuntime
+from ..pod import KubeResourceSpec
 from .function import NuclioSpec, RemoteRuntime, min_nuclio_versions
 
 serving_subkind = "serving_v2"
@@ -149,6 +151,7 @@ class ServingSpec(NuclioSpec):
         state_thresholds=None,
         disable_default_http_trigger=None,
         model_endpoint_creation_task_name=None,
+        serving_spec=None,
     ):
         super().__init__(
             command=command,
@@ -189,6 +192,7 @@ class ServingSpec(NuclioSpec):
             service_type=service_type,
             add_templated_ingress_host_mode=add_templated_ingress_host_mode,
             disable_default_http_trigger=disable_default_http_trigger,
+            serving_spec=serving_spec,
         )
 
         self.models = models or {}
@@ -296,6 +300,7 @@ class ServingRuntime(RemoteRuntime):
             self.spec.graph = step
         elif topology == StepKinds.flow:
             self.spec.graph = RootFlowStep(engine=engine or "async")
+            self.spec.graph.track_models = self.spec.track_models
         else:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 f"unsupported topology {topology}, use 'router' or 'flow'"
@@ -331,6 +336,8 @@ class ServingRuntime(RemoteRuntime):
         """
         # Applying model monitoring configurations
         self.spec.track_models = enable_tracking
+        if self.spec.graph and isinstance(self.spec.graph, RootFlowStep):
+            self.spec.graph.track_models = enable_tracking
         if self._spec and self._spec.function_refs:
             logger.debug(
                 "Set tracking for children references", enable_tracking=enable_tracking
@@ -343,6 +350,16 @@ class ServingRuntime(RemoteRuntime):
                         name
                     ]._function.spec.track_models = enable_tracking
 
+                    if self._spec.function_refs[
+                        name
+                    ]._function.spec.graph and isinstance(
+                        self._spec.function_refs[name]._function.spec.graph,
+                        RootFlowStep,
+                    ):
+                        self._spec.function_refs[
+                            name
+                        ]._function.spec.graph.track_models = enable_tracking
+
         if not 0 < sampling_percentage <= 100:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "`sampling_percentage` must be greater than 0 and less or equal to 100."
@@ -703,6 +720,7 @@ class ServingRuntime(RemoteRuntime):
             "track_models": self.spec.track_models,
             "default_content_type": self.spec.default_content_type,
             "model_endpoint_creation_task_name": self.spec.model_endpoint_creation_task_name,
+            "filename": getattr(self.spec, "filename", None),
         }
 
         if self.spec.secret_sources:
@@ -711,6 +729,10 @@ class ServingRuntime(RemoteRuntime):
 
         return json.dumps(serving_spec)
 
+    @property
+    def serving_spec(self):
+        return self._get_serving_spec()
+
     def to_mock_server(
         self,
         namespace=None,
@@ -815,3 +837,40 @@ class ServingRuntime(RemoteRuntime):
             "Turn off the mock (mock=False) and make sure Nuclio is installed for real deployment to Nuclio"
         )
         self._mock_server = self.to_mock_server()
+
+    def to_job(self) -> KubejobRuntime:
+        """Convert this ServingRuntime to a KubejobRuntime, so that the graph can be run as a standalone job."""
+        if self.spec.function_refs:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"Cannot convert function '{self.metadata.name}' to a job because it has child functions"
+            )
+
+        spec = KubeResourceSpec(
+            image=self.spec.image,
+            mode=self.spec.mode,
+            volumes=self.spec.volumes,
+            volume_mounts=self.spec.volume_mounts,
+            env=self.spec.env,
+            resources=self.spec.resources,
+            default_handler="mlrun.serving.server.execute_graph",
+            pythonpath=self.spec.pythonpath,
+            entry_points=self.spec.entry_points,
+            description=self.spec.description,
+            workdir=self.spec.workdir,
+            image_pull_secret=self.spec.image_pull_secret,
+            node_name=self.spec.node_name,
+            node_selector=self.spec.node_selector,
+            affinity=self.spec.affinity,
+            disable_auto_mount=self.spec.disable_auto_mount,
+            priority_class_name=self.spec.priority_class_name,
+            tolerations=self.spec.tolerations,
+            preemption_mode=self.spec.preemption_mode,
+            security_context=self.spec.security_context,
+            state_thresholds=self.spec.state_thresholds,
+            serving_spec=self._get_serving_spec(),
+        )
+        job = KubejobRuntime(
+            spec=spec,
+            metadata=self.metadata,
+        )
+        return job

mlrun/runtimes/pod.py

@@ -103,6 +103,7 @@ class KubeResourceSpec(FunctionSpec):
         "preemption_mode",
         "security_context",
         "state_thresholds",
+        "serving_spec",
     ]
     _default_fields_to_strip = FunctionSpec._default_fields_to_strip + [
         "volumes",
@@ -178,6 +179,7 @@ class KubeResourceSpec(FunctionSpec):
         preemption_mode=None,
         security_context=None,
         state_thresholds=None,
+        serving_spec=None,
     ):
         super().__init__(
             command=command,
@@ -223,6 +225,7 @@ class KubeResourceSpec(FunctionSpec):
             state_thresholds
             or mlrun.mlconf.function.spec.state_thresholds.default.to_dict()
         )
+        self.serving_spec = serving_spec
         # Termination grace period is internal for runtimes that have a pod termination hook hence it is not in the
         # _dict_fields and doesn't have a setter.
         self._termination_grace_period_seconds = None

mlrun/runtimes/remotesparkjob.py

@@ -58,6 +58,7 @@ class RemoteSparkSpec(KubeResourceSpec):
         preemption_mode=None,
         security_context=None,
         state_thresholds=None,
+        serving_spec=None,
     ):
         super().__init__(
             command=command,
@@ -87,6 +88,7 @@ class RemoteSparkSpec(KubeResourceSpec):
             preemption_mode=preemption_mode,
             security_context=security_context,
             state_thresholds=state_thresholds,
+            serving_spec=serving_spec,
         )
         self.provider = provider
 

mlrun/runtimes/sparkjob/spark3job.py

@@ -168,6 +168,7 @@ class Spark3JobSpec(KubeResourceSpec):
         executor_cores=None,
         security_context=None,
         state_thresholds=None,
+        serving_spec=None,
     ):
         super().__init__(
             command=command,
@@ -197,6 +198,7 @@ class Spark3JobSpec(KubeResourceSpec):
             preemption_mode=preemption_mode,
             security_context=security_context,
             state_thresholds=state_thresholds,
+            serving_spec=serving_spec,
         )
 
         self.driver_resources = driver_resources or {}

mlrun/serving/routers.py

@@ -80,10 +80,16 @@ class BaseModelRouter(RouterToDict):
         self._input_path = input_path
         self._result_path = result_path
         self._background_task_check_timestamp = None
-        self._background_task_terminate = False
         self._background_task_current_state = None
         self.kwargs = kwargs
 
+    @property
+    def background_task_reached_terminal_state(self):
+        return (
+            self._background_task_current_state
+            and self._background_task_current_state != "running"
+        )
+
     def parse_event(self, event):
         parsed_event = {}
         try:
@@ -185,35 +191,33 @@ class BaseModelRouter(RouterToDict):
                     background_task.status.state
                     in mlrun.common.schemas.BackgroundTaskState.terminal_states()
                 ):
-                    logger.debug(
+                    logger.info(
                         f"Model endpoint creation task completed with state {background_task.status.state}"
                     )
-                    self._background_task_terminate = True
                 else:  # in progress
-                    logger.debug(
+                    logger.info(
                         f"Model endpoint creation task is still in progress with the current state: "
-                        f"{background_task.status.state}. Events will not be monitored for the next 15 seconds",
+                        f"{background_task.status.state}. Events will not be monitored for the next "
+                        f"{mlrun.mlconf.model_endpoint_monitoring.model_endpoint_creation_check_period} seconds",
                         name=self.name,
                         background_task_check_timestamp=self._background_task_check_timestamp.isoformat(),
                     )
                 return background_task.status.state
             else:
-                logger.debug(
-                    "Model endpoint creation task name not provided",
+                logger.error(
+                    "Model endpoint creation task name not provided. This function is not being monitored.",
                 )
         elif self.context.monitoring_mock:
-            self._background_task_terminate = (
-                True  # If mock monitoring we return success and terminate task check.
-            )
             return mlrun.common.schemas.BackgroundTaskState.succeeded
-        self._background_task_terminate = True  # If mock without monitoring we return failed and terminate task check.
         return mlrun.common.schemas.BackgroundTaskState.failed
 
     def _update_background_task_state(self, event):
-        if not self._background_task_terminate and (
+        if not self.background_task_reached_terminal_state and (
             self._background_task_check_timestamp is None
             or now_date() - self._background_task_check_timestamp
-            >= timedelta(seconds=15)
+            >= timedelta(
+                seconds=mlrun.mlconf.model_endpoint_monitoring.model_endpoint_creation_check_period
+            )
         ):
             self._background_task_current_state = self._get_background_task_status()
         if event.body:

mlrun/serving/server.py

@@ -21,8 +21,9 @@ import os
 import socket
 import traceback
 import uuid
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
+import storey
 from nuclio import Context as NuclioContext
 from nuclio.request import Logger as NuclioLogger
 
@@ -38,9 +39,10 @@ from mlrun.secrets import SecretsStore
 
 from ..common.helpers import parse_versioned_object_uri
 from ..common.schemas.model_monitoring.constants import FileTargetKind
-from ..datastore import get_stream_pusher
+from ..datastore import DataItem, get_stream_pusher
 from ..datastore.store_resources import ResourceCache
 from ..errors import MLRunInvalidArgumentError
+from ..execution import MLClientCtx
 from ..model import ModelObj
 from ..utils import get_caller_globals
 from .states import (
@@ -322,7 +324,11 @@ class GraphServer(ModelObj):
 
     def _process_response(self, context, response, get_body):
         body = response.body
-        if isinstance(body, context.Response) or get_body:
+        if (
+            isinstance(context, MLClientCtx)
+            or isinstance(body, context.Response)
+            or get_body
+        ):
             return body
 
         if body and not isinstance(body, (str, bytes)):
@@ -535,6 +541,94 @@ def v2_serving_init(context, namespace=None):
     _set_callbacks(server, context)
 
 
+async def async_execute_graph(
+    context: MLClientCtx,
+    data: DataItem,
+    batching: bool,
+    batch_size: Optional[int],
+) -> list[Any]:
+    spec = mlrun.utils.get_serving_spec()
+
+    source_filename = spec.get("filename", None)
+    namespace = {}
+    if source_filename:
+        with open(source_filename) as f:
+            exec(f.read(), namespace)
+
+    server = GraphServer.from_dict(spec)
+
+    if config.log_level.lower() == "debug":
+        server.verbose = True
+    context.logger.info_with("Initializing states", namespace=namespace)
+    kwargs = {}
+    if hasattr(context, "is_mock"):
+        kwargs["is_mock"] = context.is_mock
+    server.init_states(
+        context=None,  # this context is expected to be a nuclio context, which we don't have in this flow
+        namespace=namespace,
+        **kwargs,
+    )
+    context.logger.info("Initializing graph steps")
+    server.init_object(namespace)
+
+    context.logger.info_with("Graph was initialized", verbose=server.verbose)
+
+    if server.verbose:
+        context.logger.info(server.to_yaml())
+
+    df = data.as_df()
+
+    responses = []
+
+    async def run(body):
+        event = storey.Event(id=index, body=body)
+        response = await server.run(event, context)
+        responses.append(response)
+
+    if batching and not batch_size:
+        batch_size = len(df)
+
+    batch = []
+    for index, row in df.iterrows():
+        data = row.to_dict()
+        if batching:
+            batch.append(data)
+            if len(batch) == batch_size:
+                await run(batch)
+                batch = []
+        else:
+            await run(data)
+
+    if batch:
+        await run(batch)
+
+    termination_result = server.wait_for_completion()
+    if asyncio.iscoroutine(termination_result):
+        await termination_result
+
+    return responses
+
+
+def execute_graph(
+    context: MLClientCtx,
+    data: DataItem,
+    batching: bool = False,
+    batch_size: Optional[int] = None,
+) -> (list[Any], Any):
+    """
+    Execute graph as a job, from start to finish.
+
+    :param context: The job's execution client context.
+    :param data: The input data to the job, to be pushed into the graph row by row, or in batches.
+    :param batching: Whether to push one or more batches into the graph rather than row by row.
+    :param batch_size: The number of rows to push per batch. If not set, and batching=True, the entire dataset will
+        be pushed into the graph in one batch.
+
+    :return: A list of responses.
+    """
+    return asyncio.run(async_execute_graph(context, data, batching, batch_size))
+
+
 def _set_callbacks(server, context):
     if not server.graph.supports_termination() or not hasattr(context, "platform"):
         return