mlrun 1.10.0rc13__py3-none-any.whl → 1.10.0rc42__py3-none-any.whl

mlrun/serving/server.py

@@ -15,22 +15,28 @@
 __all__ = ["GraphServer", "create_graph_server", "GraphContext", "MockEvent"]
 
 import asyncio
+import base64
 import copy
+import importlib
 import json
 import os
 import socket
 import traceback
 import uuid
+from collections import defaultdict
+from datetime import datetime, timezone
 from typing import Any, Optional, Union
 
+import pandas as pd
 import storey
 from nuclio import Context as NuclioContext
 from nuclio.request import Logger as NuclioLogger
 
 import mlrun
-import mlrun.common.constants
 import mlrun.common.helpers
 import mlrun.common.schemas
+import mlrun.common.schemas.model_monitoring.constants as mm_constants
+import mlrun.datastore.datastore_profile as ds_profile
 import mlrun.model_monitoring
 import mlrun.utils
 from mlrun.config import config
@@ -39,12 +45,13 @@ from mlrun.secrets import SecretsStore
 
 from ..common.helpers import parse_versioned_object_uri
 from ..common.schemas.model_monitoring.constants import FileTargetKind
+from ..common.schemas.serving import MAX_BATCH_JOB_DURATION
 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 ..utils import get_caller_globals, get_relative_module_name_from_path
 from .states import (
     FlowStep,
     MonitoredStep,
@@ -76,7 +83,6 @@ class _StreamContext:
         self.hostname = socket.gethostname()
         self.function_uri = function_uri
         self.output_stream = None
-        stream_uri = None
         log_stream = parameters.get(FileTargetKind.LOG_STREAM, "")
 
         if (enabled or log_stream) and function_uri:
@@ -87,20 +93,16 @@ class _StreamContext:
 
             stream_args = parameters.get("stream_args", {})
 
-            if log_stream == DUMMY_STREAM:
-                # Dummy stream used for testing, see tests/serving/test_serving.py
-                stream_uri = DUMMY_STREAM
-            elif not stream_args.get("mock"):  # if not a mock: `context.is_mock = True`
-                stream_uri = mlrun.model_monitoring.get_stream_path(project=project)
-
             if log_stream:
-                # Update the stream path to the log stream value
-                stream_uri = log_stream.format(project=project)
-                self.output_stream = get_stream_pusher(stream_uri, **stream_args)
+                # Get the output stream from the log stream path
+                stream_path = log_stream.format(project=project)
+                self.output_stream = get_stream_pusher(stream_path, **stream_args)
             else:
                 # Get the output stream from the profile
                 self.output_stream = mlrun.model_monitoring.helpers.get_output_stream(
-                    project=project, mock=stream_args.get("mock", False)
+                    project=project,
+                    profile=parameters.get("stream_profile"),
+                    mock=stream_args.get("mock", False),
                 )
 
 
@@ -178,11 +180,12 @@ class GraphServer(ModelObj):
         self,
         context,
         namespace,
-        resource_cache: ResourceCache = None,
+        resource_cache: Optional[ResourceCache] = None,
         logger=None,
         is_mock=False,
         monitoring_mock=False,
-    ):
+        stream_profile: Optional[ds_profile.DatastoreProfile] = None,
+    ) -> None:
         """for internal use, initialize all steps (recursively)"""
 
         if self.secret_sources:
@@ -197,6 +200,20 @@ class GraphServer(ModelObj):
         context.monitoring_mock = monitoring_mock
         context.root = self.graph
 
+        if is_mock and monitoring_mock:
+            if stream_profile:
+                # Add the user-defined stream profile to the parameters
+                self.parameters["stream_profile"] = stream_profile
+            elif not (
+                self.parameters.get(FileTargetKind.LOG_STREAM)
+                or mlrun.get_secret_or_env(
+                    mm_constants.ProjectSecretKeys.STREAM_PROFILE_NAME
+                )
+            ):
+                # Set a dummy log stream for mocking purposes if there is no direct
+                # user-defined stream profile and no information in the environment
+                self.parameters[FileTargetKind.LOG_STREAM] = DUMMY_STREAM
+
         context.stream = _StreamContext(
             self.track_models, self.parameters, self.function_uri
         )
@@ -349,33 +366,34 @@ def add_error_raiser_step(
     monitored_steps_raisers = {}
     user_steps = list(graph.steps.values())
     for monitored_step in monitored_steps.values():
-        if monitored_step.raise_exception:
-            error_step = graph.add_step(
-                class_name="mlrun.serving.states.ModelRunnerErrorRaiser",
-                name=f"{monitored_step.name}_error_raise",
-                after=monitored_step.name,
-                full_event=True,
-                raise_exception=monitored_step.raise_exception,
-                models_names=list(monitored_step.class_args["models"].keys()),
-                model_endpoint_creation_strategy=mlrun.common.schemas.ModelEndpointCreationStrategy.SKIP,
-            )
-            if monitored_step.responder:
-                monitored_step.responder = False
-                error_step.respond()
-            monitored_steps_raisers[monitored_step.name] = error_step.name
-            error_step.on_error = monitored_step.on_error
-    for step in user_steps:
-        if step.after:
-            if isinstance(step.after, list):
-                for i in range(len(step.after)):
-                    if step.after[i] in monitored_steps_raisers:
-                        step.after[i] = monitored_steps_raisers[step.after[i]]
-            else:
-                if (
-                    isinstance(step.after, str)
-                    and step.after in monitored_steps_raisers
-                ):
-                    step.after = monitored_steps_raisers[step.after]
+        error_step = graph.add_step(
+            class_name="mlrun.serving.states.ModelRunnerErrorRaiser",
+            name=f"{monitored_step.name}_error_raise",
+            after=monitored_step.name,
+            full_event=True,
+            raise_exception=monitored_step.raise_exception,
+            models_names=list(monitored_step.class_args["models"].keys()),
+            model_endpoint_creation_strategy=mlrun.common.schemas.ModelEndpointCreationStrategy.SKIP,
+            function=monitored_step.function,
+        )
+        if monitored_step.responder:
+            monitored_step.responder = False
+            error_step.respond()
+        monitored_steps_raisers[monitored_step.name] = error_step.name
+        error_step.on_error = monitored_step.on_error
+    if monitored_steps_raisers:
+        for step in user_steps:
+            if step.after:
+                if isinstance(step.after, list):
+                    for i in range(len(step.after)):
+                        if step.after[i] in monitored_steps_raisers:
+                            step.after[i] = monitored_steps_raisers[step.after[i]]
+                else:
+                    if (
+                        isinstance(step.after, str)
+                        and step.after in monitored_steps_raisers
+                    ):
+                        step.after = monitored_steps_raisers[step.after]
     return graph
 
 
@@ -384,6 +402,7 @@ def add_monitoring_general_steps(
     graph: RootFlowStep,
     context,
     serving_spec,
+    pause_until_background_task_completion: bool,
 ) -> tuple[RootFlowStep, FlowStep]:
     """
     Adding the monitoring flow connection steps, this steps allow the graph to reconstruct the serving event enrich it
@@ -392,18 +411,23 @@ def add_monitoring_general_steps(
         "background_task_status_step" --> "filter_none" --> "monitoring_pre_processor_step" --> "flatten_events"
         --> "sampling_step" --> "filter_none_sampling" --> "model_monitoring_stream"
     """
+    background_task_status_step = None
+    if pause_until_background_task_completion:
+        background_task_status_step = graph.add_step(
+            "mlrun.serving.system_steps.BackgroundTaskStatus",
+            "background_task_status_step",
+            model_endpoint_creation_strategy=mlrun.common.schemas.ModelEndpointCreationStrategy.SKIP,
+            full_event=True,
+        )
     monitor_flow_step = graph.add_step(
-        "mlrun.serving.system_steps.BackgroundTaskStatus",
-        "background_task_status_step",
-        model_endpoint_creation_strategy=mlrun.common.schemas.ModelEndpointCreationStrategy.SKIP,
-    )
-    graph.add_step(
         "storey.Filter",
         "filter_none",
         _fn="(event is not None)",
-        after="background_task_status_step",
+        after="background_task_status_step" if background_task_status_step else None,
         model_endpoint_creation_strategy=mlrun.common.schemas.ModelEndpointCreationStrategy.SKIP,
     )
+    if background_task_status_step:
+        monitor_flow_step = background_task_status_step
     graph.add_step(
         "mlrun.serving.system_steps.MonitoringPreProcessor",
         "monitoring_pre_processor_step",
@@ -466,14 +490,28 @@ def add_monitoring_general_steps(
 
 
 def add_system_steps_to_graph(
-    project: str, graph: RootFlowStep, track_models: bool, context, serving_spec
+    project: str,
+    graph: RootFlowStep,
+    track_models: bool,
+    context,
+    serving_spec,
+    pause_until_background_task_completion: bool = True,
 ) -> RootFlowStep:
+    if not (isinstance(graph, RootFlowStep) and graph.include_monitored_step()):
+        return graph
     monitored_steps = graph.get_monitored_steps()
     graph = add_error_raiser_step(graph, monitored_steps)
     if track_models:
+        background_task_status_step = None
         graph, monitor_flow_step = add_monitoring_general_steps(
-            project, graph, context, serving_spec
+            project,
+            graph,
+            context,
+            serving_spec,
+            pause_until_background_task_completion,
         )
+        if background_task_status_step:
+            monitor_flow_step = background_task_status_step
         # Connect each model runner to the monitoring step:
         for step_name, step in monitored_steps.items():
             if monitor_flow_step.after:
@@ -494,18 +532,13 @@ def v2_serving_init(context, namespace=None):
     context.logger.info("Initializing server from spec")
     spec = mlrun.utils.get_serving_spec()
     server = GraphServer.from_dict(spec)
-    if isinstance(server.graph, RootFlowStep) and server.graph.include_monitored_step():
-        server.graph = add_system_steps_to_graph(
-            server.project,
-            copy.deepcopy(server.graph),
-            spec.get("track_models"),
-            context,
-            spec,
-        )
-        context.logger.info_with(
-            "Server graph after adding system steps",
-            graph=str(server.graph.steps),
-        )
+    server.graph = add_system_steps_to_graph(
+        server.project,
+        copy.deepcopy(server.graph),
+        spec.get("track_models"),
+        context,
+        spec,
+    )
 
     if config.log_level.lower() == "debug":
         server.verbose = True
@@ -542,22 +575,120 @@ def v2_serving_init(context, namespace=None):
 async def async_execute_graph(
     context: MLClientCtx,
     data: DataItem,
+    timestamp_column: Optional[str],
     batching: bool,
     batch_size: Optional[int],
-) -> list[Any]:
+    read_as_lists: bool,
+    nest_under_inputs: bool,
+) -> None:
+    # Validate that data parameter is a DataItem and not passed via params
+    if not isinstance(data, DataItem):
+        raise MLRunInvalidArgumentError(
+            f"Parameter 'data' has type hint 'DataItem' but got {type(data).__name__} instead. "
+            f"Data files and artifacts must be passed via the 'inputs' parameter, not 'params'. "
+            f"The 'params' parameter is for simple configuration values (strings, numbers, booleans), "
+            f"while 'inputs' is for data files that need to be loaded. "
+            f"Example: run_function(..., inputs={{'data': 'path/to/data.csv'}}, params={{other_config: value}})"
+        )
+    run_call_count = 0
     spec = mlrun.utils.get_serving_spec()
+    modname = None
+    code = os.getenv("MLRUN_EXEC_CODE")
+    if code:
+        code = base64.b64decode(code).decode("utf-8")
+        with open("user_code.py", "w") as fp:
+            fp.write(code)
+        modname = "user_code"
+    else:
+        # TODO: find another way to get the local file path, or ensure that MLRUN_EXEC_CODE
+        #  gets set in local flow and not just in the remote pod
+        source_file_path = spec.get("filename", None)
+        if source_file_path:
+            source_file_path_object, working_dir_path_object = (
+                mlrun.utils.helpers.get_source_and_working_dir_paths(source_file_path)
+            )
+            if not source_file_path_object.is_relative_to(working_dir_path_object):
+                raise mlrun.errors.MLRunRuntimeError(
+                    f"Source file path '{source_file_path}' is not under the current working directory "
+                    f"(which is required when running with local=True)"
+                )
+            modname = get_relative_module_name_from_path(
+                source_file_path_object, working_dir_path_object
+            )
 
-    source_filename = spec.get("filename", None)
     namespace = {}
-    if source_filename:
-        with open(source_filename) as f:
-            exec(f.read(), namespace)
+    if modname:
+        mod = importlib.import_module(modname)
+        namespace = mod.__dict__
 
     server = GraphServer.from_dict(spec)
 
+    if server.model_endpoint_creation_task_name:
+        context.logger.info(
+            f"Waiting for model endpoint creation task '{server.model_endpoint_creation_task_name}'..."
+        )
+        background_task = (
+            mlrun.get_run_db().wait_for_background_task_to_reach_terminal_state(
+                project=server.project,
+                name=server.model_endpoint_creation_task_name,
+            )
+        )
+        task_state = background_task.status.state
+        if task_state == mlrun.common.schemas.BackgroundTaskState.failed:
+            raise mlrun.errors.MLRunRuntimeError(
+                "Aborting job due to model endpoint creation background task failure"
+            )
+        elif task_state != mlrun.common.schemas.BackgroundTaskState.succeeded:
+            # this shouldn't happen, but we need to know if it does
+            raise mlrun.errors.MLRunRuntimeError(
+                "Aborting job because the model endpoint creation background task did not succeed "
+                f"(status='{task_state}')"
+            )
+
+    df = data.as_df()
+
+    if df.empty:
+        context.logger.warn("Job terminated due to empty inputs (0 rows)")
+        return []
+
+    track_models = spec.get("track_models")
+
+    if track_models and timestamp_column:
+        context.logger.info(f"Sorting dataframe by {timestamp_column}")
+        df[timestamp_column] = pd.to_datetime(  # in case it's a string
+            df[timestamp_column]
+        )
+        df.sort_values(by=timestamp_column, inplace=True)
+        if len(df) > 1:
+            start_time = df[timestamp_column].iloc[0]
+            end_time = df[timestamp_column].iloc[-1]
+            time_range = end_time - start_time
+            start_time = start_time.isoformat()
+            end_time = end_time.isoformat()
+            # TODO: tie this to the controller's base period
+            if time_range > pd.Timedelta(MAX_BATCH_JOB_DURATION):
+                raise mlrun.errors.MLRunRuntimeError(
+                    f"Dataframe time range is too long: {time_range}. "
+                    "Please disable tracking or reduce the input dataset's time range below the defined limit "
+                    f"of {MAX_BATCH_JOB_DURATION}."
+                )
+        else:
+            start_time = end_time = df["timestamp"].iloc[0].isoformat()
+    else:
+        # end time will be set from clock time when the batch completes
+        start_time = datetime.now(tz=timezone.utc).isoformat()
+
+    server.graph = add_system_steps_to_graph(
+        server.project,
+        copy.deepcopy(server.graph),
+        track_models,
+        context,
+        spec,
+        pause_until_background_task_completion=False,  # we've already awaited it
+    )
+
     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
@@ -574,57 +705,166 @@ async def async_execute_graph(
     if server.verbose:
         context.logger.info(server.to_yaml())
 
-    df = data.as_df()
-
-    responses = []
-
     async def run(body):
+        nonlocal run_call_count
         event = storey.Event(id=index, body=body)
-        response = await server.run(event, context)
-        responses.append(response)
+        if timestamp_column:
+            if batching:
+                # we use the first row in the batch to determine the timestamp for the whole batch
+                body = body[0]
+            if not isinstance(body, dict):
+                raise mlrun.errors.MLRunRuntimeError(
+                    f"When timestamp_column=True, event body must be a dict – got {type(body).__name__} instead"
+                )
+            if timestamp_column not in body:
+                raise mlrun.errors.MLRunRuntimeError(
+                    f"Event body '{body}' did not contain timestamp column '{timestamp_column}'"
+                )
+            event._original_timestamp = body[timestamp_column]
+        run_call_count += 1
+        return await server.run(event, context)
 
     if batching and not batch_size:
         batch_size = len(df)
 
     batch = []
+    tasks = []
     for index, row in df.iterrows():
-        data = row.to_dict()
+        data = row.to_list() if read_as_lists else row.to_dict()
+        if nest_under_inputs:
+            data = {"inputs": data}
         if batching:
             batch.append(data)
             if len(batch) == batch_size:
-                await run(batch)
+                tasks.append(asyncio.create_task(run(batch)))
                 batch = []
         else:
-            await run(data)
+            tasks.append(asyncio.create_task(run(data)))
 
     if batch:
-        await run(batch)
+        tasks.append(asyncio.create_task(run(batch)))
+
+    responses = await asyncio.gather(*tasks)
 
     termination_result = server.wait_for_completion()
     if asyncio.iscoroutine(termination_result):
         await termination_result
 
-    return responses
+    model_endpoint_uids = spec.get("model_endpoint_uids", [])
+
+    # needed for output_stream to be created
+    server = GraphServer.from_dict(spec)
+    server.init_states(None, namespace)
+
+    batch_completion_time = datetime.now(tz=timezone.utc).isoformat()
+
+    if not timestamp_column:
+        end_time = batch_completion_time
+
+    mm_stream_record = dict(
+        kind="batch_complete",
+        project=context.project,
+        first_timestamp=start_time,
+        last_timestamp=end_time,
+        batch_completion_time=batch_completion_time,
+    )
+    output_stream = server.context.stream.output_stream
+    for mep_uid in spec.get("model_endpoint_uids", []):
+        mm_stream_record["endpoint_id"] = mep_uid
+        output_stream.push(mm_stream_record, partition_key=mep_uid)
+
+    context.logger.info(
+        f"Job completed processing {len(df)} rows",
+        timestamp_column=timestamp_column,
+        model_endpoint_uids=model_endpoint_uids,
+    )
+
+    # log the results as artifacts
+    num_of_meps_in_the_graph = len(server.graph.model_endpoints_names)
+    artifact_path = None
+    if (
+        "{{run.uid}}" not in context.artifact_path
+    ):  # TODO: delete when IG-22841 is resolved
+        artifact_path = "+/{{run.uid}}"  # will be concatenated to the context's path in extend_artifact_path
+    if num_of_meps_in_the_graph <= 1:
+        context.log_dataset(
+            "prediction", df=pd.DataFrame(responses), artifact_path=artifact_path
+        )
+    else:
+        # turn this list of samples into a dict of lists, one per model endpoint
+        grouped = defaultdict(list)
+        for sample in responses:
+            for model_name, features in sample.items():
+                grouped[model_name].append(features)
+        # create a dataframe per model endpoint and log it
+        for model_name, features in grouped.items():
+            context.log_dataset(
+                f"prediction_{model_name}",
+                df=pd.DataFrame(features),
+                artifact_path=artifact_path,
+            )
+    context.log_result("num_rows", run_call_count)
+
+
+def _is_inside_asyncio_loop():
+    try:
+        asyncio.get_running_loop()
+        return True
+    except RuntimeError:
+        return False
+
+
+# Workaround for running with local=True in Jupyter (ML-10620)
+def _workaround_asyncio_nesting():
+    try:
+        import nest_asyncio
+    except ImportError:
+        raise mlrun.errors.MLRunRuntimeError(
+            "Cannot execute graph from within an already running asyncio loop. "
+            "Attempt to import nest_asyncio as a workaround failed as well."
+        )
+    nest_asyncio.apply()
 
 
 def execute_graph(
     context: MLClientCtx,
     data: DataItem,
+    timestamp_column: Optional[str] = None,
     batching: bool = False,
     batch_size: Optional[int] = None,
+    read_as_lists: bool = False,
+    nest_under_inputs: bool = False,
 ) -> (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 timestamp_column: The name of the column that will be used as the timestamp for model monitoring purposes.
+        when timestamp_column is used in conjunction with batching, the first timestamp will be used for the entire
+        batch.
     :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.
+    :param read_as_lists: Whether to read each row as a list instead of a dictionary.
+    :param nest_under_inputs: Whether to wrap each row with {"inputs": ...}.
 
     :return: A list of responses.
     """
-    return asyncio.run(async_execute_graph(context, data, batching, batch_size))
+    if _is_inside_asyncio_loop():
+        _workaround_asyncio_nesting()
+
+    return asyncio.run(
+        async_execute_graph(
+            context,
+            data,
+            timestamp_column,
+            batching,
+            batch_size,
+            read_as_lists,
+            nest_under_inputs,
+        )
+    )
 
 
 def _set_callbacks(server, context):