# Copyright 2023 Iguazio
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
__all__ = ["GraphServer", "create_graph_server", "GraphContext", "MockEvent"]
import asyncio
import base64
import copy
import importlib
import inspect
import json
import os
import socket
import traceback
import uuid
from collections import defaultdict
from datetime import UTC, datetime
from http import HTTPMethod
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.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.errors
import mlrun.model_monitoring
import mlrun.runtimes.nuclio.serving
import mlrun.utils
from mlrun.config import config
from mlrun.errors import err_to_str
from mlrun.secrets import SecretsStore
from mlrun.serving.endpoint_mapping import APIHandlerConfig
from mlrun.serving.result_handler import ResultHandler
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, get_relative_module_name_from_path
from .states import (
FlowStep,
MonitoredStep,
RootFlowStep,
RouterStep,
get_function,
graph_root_setter,
)
from .utils import event_id_key, event_path_key
DUMMY_STREAM = "dummy://"
class _StreamContext:
"""Handles the stream context for the events stream process. Includes the configuration for the output stream
that will be used for pushing the events from the nuclio model serving function"""
def __init__(self, enabled: bool, parameters: dict, function_uri: str):
"""
Initialize _StreamContext object.
:param enabled: A boolean indication for applying the stream context
:param parameters: Dictionary of optional parameters, such as `log_stream` and `stream_args`. Note that these
parameters might be relevant to the output source such as `kafka_brokers` if
the output source is from type Kafka.
:param function_uri: Full value of the function uri, usually it's <project-name>/<function-name>
"""
self.enabled = False
self.hostname = socket.gethostname()
self.function_uri = function_uri
self.output_stream = None
log_stream = parameters.get(FileTargetKind.LOG_STREAM, "")
if (enabled or log_stream) and function_uri:
self.enabled = True
project, _, _, _ = parse_versioned_object_uri(
function_uri, config.active_project
)
stream_args = parameters.get("stream_args", {})
if log_stream:
# 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,
profile=parameters.get("stream_profile"),
mock=stream_args.get("mock", False),
)
[docs]
class GraphServer(ModelObj):
kind = "server"
def __init__(
self,
graph=None,
parameters=None,
load_mode=None,
function_uri=None,
verbose=False,
version=None,
functions=None,
graph_initializer=None,
error_stream=None,
track_models=None,
secret_sources=None,
default_content_type=None,
function_name=None,
function_tag=None,
project=None,
model_endpoint_creation_task_name=None,
api_handler_config: "APIHandlerConfig | None" = None,
):
self._graph = None
self.graph: Union[RouterStep, RootFlowStep] = graph
self.function_uri = function_uri
self.parameters = parameters or {}
self.verbose = verbose
self.load_mode = load_mode or "sync"
self.version = version or "v2"
self.context = None
self._current_function = None
self.functions = functions or {}
self.graph_initializer = graph_initializer
self.error_stream = error_stream
self.track_models = track_models
self._error_stream_object = None
self.secret_sources = secret_sources
self._secrets = SecretsStore.from_list(secret_sources)
self._db_conn = None
self.resource_cache = None
self.default_content_type = default_content_type
self.http_trigger = True
self.function_name = function_name
self.function_tag = function_tag
self.project = project
self.model_endpoint_creation_task_name = model_endpoint_creation_task_name
self.streaming = False
self.api_handler_config = api_handler_config
@property
def api_handler_config(self) -> "APIHandlerConfig | None":
return self._api_handler_config
@api_handler_config.setter
def api_handler_config(self, value: "APIHandlerConfig | dict | None") -> None:
if isinstance(value, dict):
value = APIHandlerConfig.from_dict(value)
self._api_handler_config = value
self.result_handler = ResultHandler(value) if value else None
[docs]
def set_current_function(self, function):
"""set which child function this server is currently running on"""
self._current_function = function
@property
def graph(self) -> Union[RootFlowStep, RouterStep]:
return self._graph
@graph.setter
def graph(self, graph):
graph_root_setter(self, graph)
[docs]
def set_error_stream(self, error_stream):
"""set/initialize the error notification stream"""
self.error_stream = error_stream
if error_stream:
self._error_stream_object = get_stream_pusher(error_stream)
else:
self._error_stream_object = None
def _get_db(self):
return mlrun.get_run_db(secrets=self._secrets)
[docs]
def init_states(
self,
context,
namespace,
resource_cache: ResourceCache | None = None,
logger=None,
is_mock=False,
monitoring_mock=False,
stream_profile: ds_profile.DatastoreProfile | None = None,
) -> None:
"""for internal use, initialize all steps (recursively)"""
if self.secret_sources:
self._secrets = SecretsStore.from_list(self.secret_sources)
if self.error_stream:
self._error_stream_object = get_stream_pusher(self.error_stream)
self.resource_cache = resource_cache or ResourceCache()
context = GraphContext(server=self, nuclio_context=context, logger=logger)
context.is_mock = is_mock
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
)
context.current_function = self._current_function
context.get_store_resource = self.resource_cache.resource_getter(
self._get_db(), self._secrets
)
context.get_table = self.resource_cache.get_table
context.verbose = self.verbose
self.context = context
if self.graph_initializer:
if callable(self.graph_initializer):
handler = self.graph_initializer
else:
handler = get_function(self.graph_initializer, namespace or [])
handler(self)
context.root = self.graph
[docs]
def init_object(self, namespace):
self.graph.init_object(self.context, namespace, self.load_mode, reset=True)
[docs]
def test(
self,
path: str = "/",
body: Union[str, bytes, dict] | None = None,
method: str = "",
headers: str | None = None,
content_type: str | None = None,
silent: bool = False,
get_body: bool = True,
event_id: str | None = None,
trigger: "MockTrigger" = None,
offset=None,
time=None,
):
"""invoke a test event into the server to simulate/test server behavior
example::
server = create_graph_server()
server.add_model("my", class_name=MyModelClass, model_path="{path}", z=100)
print(server.test("my/infer", testdata))
:param path: api path, e.g. (/{router.url_prefix}/{model-name}/..) path
:param body: message body (dict or json str/bytes)
:param method: optional, GET, POST, ..
:param headers: optional, request headers, ..
:param content_type: optional, http mime type
:param silent: don't raise on error responses (when not 20X)
:param get_body: return the body as py object (vs serialize response into json)
:param event_id: specify the unique event ID (by default a random value will be generated)
:param trigger: nuclio trigger info or mlrun.serving.server.MockTrigger class (holds kind and name)
:param offset: trigger offset (for streams)
:param time: event time Datetime or str, default to now()
"""
if not self.graph:
raise MLRunInvalidArgumentError(
"no models or steps were set, use function.set_topology() and add steps"
)
if not method:
method = HTTPMethod.POST if body else HTTPMethod.GET
event = MockEvent(
body=body,
path=path,
method=method,
headers=headers,
content_type=content_type,
event_id=event_id,
trigger=trigger,
offset=offset,
time=time,
)
resp = self.run(event, get_body=get_body)
if hasattr(resp, "status_code") and resp.status_code >= 300 and not silent:
raise RuntimeError(f"failed ({resp.status_code}): {resp.body}")
return resp
[docs]
def run(self, event, context=None, get_body: bool = False, extra_args=None):
server_context = self.context
context = context or server_context
event.content_type = event.content_type or self.default_content_type or ""
if event.headers:
if event_id_key in event.headers:
event.id = event.headers.get(event_id_key)
if event_path_key in event.headers:
event.path = event.headers.get(event_path_key)
if isinstance(event.body, str | bytes) and (
not event.content_type or event.content_type in ["json", "application/json"]
):
# assume it is json and try to load
try:
body = json.loads(event.body)
event.body = body
except (json.decoder.JSONDecodeError, UnicodeDecodeError) as exc:
if event.content_type in ["json", "application/json"]:
# if its json type and didnt load, raise exception
message = f"failed to json decode event, {err_to_str(exc)}"
context.logger.error(message)
server_context.push_error(event, message, source="_handler")
return context.Response(
body=message, content_type="text/plain", status_code=400
)
try:
response = self.graph.run(event, **(extra_args or {}))
# TODO: this is only relevant in certain flows (MockServer, sync...)
if hasattr(response, "body"):
response = response.body
if self.http_trigger and self.result_handler:
method = getattr(event, "method", None)
path = getattr(event, "path", None)
if method and path:
response = self.result_handler.apply(method, path, response)
except Exception as exc:
# Extract appropriate status code from MLRunHTTPStatusError exceptions
# For backwards compatibility, default to 400 for other exceptions
if isinstance(exc, mlrun.errors.MLRunHTTPStatusError):
status_code = exc.error_status_code
else:
status_code = 400
message = f"{exc.__class__.__name__}: {err_to_str(exc)}"
if server_context.verbose:
message += "\n" + str(traceback.format_exc())
context.logger.error(f"run error, {traceback.format_exc()}")
server_context.push_error(event, message, source="_handler")
return context.Response(
body=message, content_type="text/plain", status_code=status_code
)
return response
[docs]
def wait_for_completion(self):
"""wait for async operation to complete"""
return self.graph.wait_for_completion()
def add_error_raiser_step(
graph: RootFlowStep, monitored_steps: dict[str, MonitoredStep]
) -> RootFlowStep:
for monitored_step in monitored_steps.values():
unpack_step = f"{monitored_step.name}_unpacker"
graph.add_step(
class_name="storey.FlatMap",
name=unpack_step,
_fn="(event.body)",
after=monitored_step.name,
full_event=True,
model_endpoint_creation_strategy=mlrun.common.schemas.ModelEndpointCreationStrategy.SKIP,
function=monitored_step.function,
)
# Add error raiser step after the unpacker
error_step = graph.add_step(
class_name="mlrun.serving.states.ModelRunnerErrorRaiser",
name=f"{monitored_step.name}_error_raise",
after=[monitored_step.name, unpack_step],
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()
error_step.on_error = monitored_step.on_error
return graph
def add_monitoring_general_steps(
project: str,
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
and push it to the model monitoring stream
system_steps structure -
"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(
"storey.Filter",
"filter_none",
_fn="(event is not None)",
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",
after="filter_none",
full_event=True,
model_endpoint_creation_strategy=mlrun.common.schemas.ModelEndpointCreationStrategy.SKIP,
)
# flatten the events
graph.add_step(
"storey.FlatMap",
"flatten_events",
_fn="(event)",
after="monitoring_pre_processor_step",
model_endpoint_creation_strategy=mlrun.common.schemas.ModelEndpointCreationStrategy.SKIP,
)
graph.add_step(
"mlrun.serving.system_steps.SamplingStep",
"sampling_step",
after="flatten_events",
sampling_percentage=float(
serving_spec.get("parameters", {}).get("sampling_percentage", 100.0)
if isinstance(serving_spec, dict)
else getattr(serving_spec, "parameters", {}).get(
"sampling_percentage", 100.0
),
),
model_endpoint_creation_strategy=mlrun.common.schemas.ModelEndpointCreationStrategy.SKIP,
)
graph.add_step(
"storey.Filter",
"filter_none_sampling",
_fn="(event is not None)",
after="sampling_step",
model_endpoint_creation_strategy=mlrun.common.schemas.ModelEndpointCreationStrategy.SKIP,
)
if getattr(context, "is_mock", False):
graph.add_step(
"mlrun.serving.system_steps.MockStreamPusher",
"model_monitoring_stream",
after="filter_none_sampling",
model_endpoint_creation_strategy=mlrun.common.schemas.ModelEndpointCreationStrategy.SKIP,
)
else:
stream_uri = mlrun.model_monitoring.get_stream_path(
project=project,
function_name=mlrun.common.schemas.MonitoringFunctionNames.STREAM,
)
context.logger.info_with(
"Creating Model Monitoring stream target using uri:", uri=stream_uri
)
graph.add_step(
">>",
"model_monitoring_stream",
path=stream_uri,
sharding_func=mlrun.common.schemas.model_monitoring.constants.StreamProcessingEvent.ENDPOINT_ID,
after="filter_none_sampling",
)
return graph, monitor_flow_step
def _add_api_handler_step_to_graph(
graph: RootFlowStep,
serving_spec: Optional["mlrun.runtimes.nuclio.serving.ServingSpec"],
context: "GraphContext",
) -> RootFlowStep:
"""Add API handler step to graph if api_handler_config is present"""
if isinstance(serving_spec, dict):
# Nuclio runtime
api_handler_config = serving_spec.get("api_handler_config")
elif isinstance(serving_spec, mlrun.runtimes.nuclio.serving.ServingSpec):
# Mock server
api_handler_config = getattr(serving_spec, "api_handler_config", None)
else:
raise mlrun.errors.MLRunValueError(
f"serving_spec must be dict or ServingSpec, got {type(serving_spec)}"
)
if api_handler_config:
context.logger.info(
"Adding API handler step to graph based on serving spec config"
)
# Check if _APIHandlerStep already exists to avoid duplicates
existing_api_handler = None
for step_name, step in graph.steps.items():
if (
hasattr(step, "class_name")
and step.class_name == "mlrun.serving.api_handler._APIHandlerStep"
):
existing_api_handler = step
break
if not existing_api_handler:
# Find current starting steps (using same logic as check_and_process_graph)
current_start_steps = []
for step_name, step in graph.steps.items():
# A step is a starting step if:
# 1. It has no 'after' and no 'cycle_from' (simple starting step)
# 2. It has both 'after' and 'cycle_from', and they match (cyclic starting step)
if not step.after and not getattr(step, "cycle_from", None):
current_start_steps.append(step_name)
elif (
step.after
and getattr(step, "cycle_from", None)
and set(step.after) == set(step.cycle_from)
):
current_start_steps.append(step_name)
# Add _APIHandlerStep as the first step
graph.add_step(
class_name="mlrun.serving.api_handler._APIHandlerStep",
name="api-handler",
graph_shape="diamond",
config=api_handler_config,
context=context,
after=None, # First step
full_event=True,
)
# Chain all existing starting steps to come after the API handler step
for step_name in current_start_steps:
step = graph[step_name]
step.after = step.after or []
if isinstance(step.after, str):
step.after = [step.after]
if "api-handler" not in step.after:
step.after.append("api-handler")
return graph
def add_system_steps_to_graph(
project: str,
graph: RootFlowStep,
track_models: bool,
context,
serving_spec: Optional["mlrun.runtimes.nuclio.serving.ServingSpec"],
pause_until_background_task_completion: bool = True,
) -> RootFlowStep:
# Always add API handler step if configured
graph = _add_api_handler_step_to_graph(graph, serving_spec, context)
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,
pause_until_background_task_completion,
)
if background_task_status_step:
monitor_flow_step = background_task_status_step
# Check if streaming is enabled for this function
streaming_enabled = (
serving_spec.get("streaming", False)
if isinstance(serving_spec, dict)
else getattr(serving_spec, "streaming", False)
)
# Connect each model runner to the monitoring step.
# For streaming functions, add a Collector step to aggregate streaming
# chunks into a single event for MM. For non-streaming, connect directly.
for step_name, step in monitored_steps.items():
if streaming_enabled:
# Add a Collector step after each monitored step
collector_name = f"{step_name}_collector"
graph.add_step(
"storey.Collector",
collector_name,
after=step_name,
model_endpoint_creation_strategy=mlrun.common.schemas.ModelEndpointCreationStrategy.SKIP,
)
source_step = collector_name
else:
source_step = step_name
# Connect monitor_flow_step to receive from source
if monitor_flow_step.after:
if isinstance(monitor_flow_step.after, list):
monitor_flow_step.after.append(source_step)
elif isinstance(monitor_flow_step.after, str):
monitor_flow_step.after = [monitor_flow_step.after, source_step]
else:
monitor_flow_step.after = [source_step]
return graph
def v2_serving_init(context, namespace=None):
"""hook for nuclio init_context()"""
context.logger.info("Initializing server from spec")
spec = mlrun.utils.get_serving_spec()
server = GraphServer.from_dict(spec)
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
if hasattr(context, "trigger"):
server.http_trigger = getattr(context.trigger, "kind", "http") == "http"
context.logger.info_with(
"Setting current function",
current_function=os.getenv("SERVING_CURRENT_FUNCTION", ""),
)
server.set_current_function(os.getenv("SERVING_CURRENT_FUNCTION", ""))
# Set streaming mode before init_states so it's available during graph initialization
server.streaming = spec.get("streaming", False)
context.logger.info_with(
"Initializing states", namespace=namespace or get_caller_globals()
)
kwargs = {}
if hasattr(context, "is_mock"):
kwargs["is_mock"] = context.is_mock
server.init_states(
context,
namespace or get_caller_globals(),
**kwargs,
)
context.logger.info("Initializing graph steps")
server.init_object(namespace or get_caller_globals())
# Select the appropriate handler based on streaming mode
if server.streaming:
# Validate that trigger is HTTP when streaming is enabled
if (
hasattr(context, "trigger")
and getattr(context.trigger, "kind", "http") != "http"
):
raise ValueError(
f"Streaming is only supported with HTTP triggers, but trigger kind is "
f"'{context.trigger.kind}'. Disable streaming or use an HTTP trigger."
)
context.logger.info("Streaming mode enabled, using streaming handler")
setattr(context, "mlrun_handler", v2_serving_streaming_handler)
else:
setattr(context, "mlrun_handler", v2_serving_handler)
setattr(context, "_server", server)
context.logger.info_with("Serving was initialized", verbose=server.verbose)
if server.verbose:
context.logger.info(server.to_yaml())
_set_callbacks(server, context)
async def async_execute_graph(
context: MLClientCtx,
data: DataItem | None = None,
data_object: dict | None = None,
timestamp_column: str | None = None,
batching: bool = False,
batch_size: int | None = None,
read_as_lists: bool = False,
nest_under_inputs: bool = False,
) -> Any:
"""See :func:`execute_graph` for parameter documentation."""
# Fail-fast argument validation
if data is None and data_object is None:
raise MLRunInvalidArgumentError(
"exactly one of 'data' or 'data_object' must be provided"
)
if data is not None and data_object is not None:
raise MLRunInvalidArgumentError(
"'data' and 'data_object' are mutually exclusive — provide exactly one"
)
if data_object is not None and not isinstance(data_object, dict):
raise MLRunInvalidArgumentError(
f"data_object must be a dict, got {type(data_object).__name__}. "
f"If you have a Pydantic model, call .model_dump(); "
f"if you have a JSON string, json.loads() it before passing."
)
# Validate that data parameter is a DataItem and not passed via params
# (only enforced when data_object is None — the data_object path bypasses the DataItem contract).
if data_object is None and 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}})"
)
is_dict_path = data_object is not None
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
)
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}')"
)
track_models = spec.get("track_models")
# end_time may be overwritten further down (in the dict branch or from batch_completion_time)
end_time = None
if is_dict_path:
# Single-event path: no DataFrame, no sort, no MAX_BATCH_JOB_DURATION guard.
if track_models and timestamp_column:
if timestamp_column not in data_object:
raise mlrun.errors.MLRunRuntimeError(
f"Event body '{data_object}' did not contain timestamp column '{timestamp_column}'"
)
# Single event => first_timestamp == last_timestamp.
start_time = end_time = str(data_object[timestamp_column])
else:
# end_time will be set from clock time when the batch completes (same as the batch path).
start_time = datetime.now(tz=UTC).isoformat()
else:
df = data.as_df()
if df.empty:
context.logger.warn("Job terminated due to empty inputs (0 rows)")
return
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=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
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())
async def run(body, idx):
nonlocal run_call_count
event = storey.Event(id=idx, body=body)
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)
tasks = []
if is_dict_path:
# `batching` / `batch_size` are ignored on the data_object path.
if batching or batch_size:
context.logger.debug(
"ignoring batch params on data_object path",
batching=batching,
batch_size=batch_size,
)
body = list(data_object.values()) if read_as_lists else data_object
if nest_under_inputs:
body = {"inputs": body}
tasks.append(asyncio.create_task(run(body, 0)))
else:
if batching and not batch_size:
batch_size = len(df)
batch = []
for index, row in df.iterrows():
body = row.to_list() if read_as_lists else row.to_dict()
if nest_under_inputs:
body = {"inputs": body}
if batching:
batch.append(body)
if len(batch) == batch_size:
tasks.append(asyncio.create_task(run(batch, index)))
batch = []
else:
tasks.append(asyncio.create_task(run(body, index)))
if batch:
tasks.append(asyncio.create_task(run(batch, index)))
responses = await asyncio.gather(*tasks)
termination_result = server.wait_for_completion()
if asyncio.iscoroutine(termination_result):
await termination_result
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=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(
"Job completed",
rows=run_call_count,
timestamp_column=timestamp_column,
model_endpoint_uids=model_endpoint_uids,
)
has_responder = False
for step in server.graph.steps.values():
if getattr(step, "responder", False):
has_responder = True
break
if has_responder:
# log the results as a dataset artifact
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
context.log_dataset(
"prediction", df=pd.DataFrame(responses), artifact_path=artifact_path
)
# if we got responses that appear to be in the right format, try to log per-model datasets too
if (
responses
and responses[0]
and isinstance(responses[0], dict)
and isinstance(next(iter(responses[0].values())), dict | list)
):
try:
# 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,
)
except Exception as e:
context.logger.warning(
"Failed to log per-model prediction datasets",
error=err_to_str(e),
)
context.log_result("num_rows", run_call_count)
if is_dict_path:
return responses[0]
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 | None = None,
data_object: dict | None = None,
timestamp_column: str | None = None,
batching: bool = False,
batch_size: int | None = None,
read_as_lists: bool = False,
nest_under_inputs: bool = False,
) -> Any:
"""
Execute graph as a job, from start to finish.
Exactly one of ``data`` or ``data_object`` must be provided. The two modes
are mutually exclusive:
- **Batch mode** (``data``): the existing ``DataItem`` path — loads a
DataFrame from the input artifact and feeds it into the graph row by row
(or batched). Returns ``None``; the responses are surfaced as the
``prediction`` (and per-model ``prediction_<model>``) dataset artifacts.
- **Single-instance mode** (``data_object``): runs the graph **exactly
once** with the provided dict as the event body. Returns the graph
response as-is so the caller can read it.
``batching`` and ``batch_size`` are ignored (a DEBUG log entry is emitted
when they are set). If ``timestamp_column`` is set, the value is read
from the root of ``data_object`` (e.g. ``timestamp_column="ts"`` with
``data_object={"value": 8, "ts": "2020-01-01T00:00:00"}``); a missing key
raises :class:`mlrun.errors.MLRunRuntimeError`. If ``read_as_lists=True``
the body is ``list(data_object.values())`` (insertion order). If
``nest_under_inputs=True`` the body is wrapped as ``{"inputs": body}``.
:param context: The job's execution client context.
:param data: The input data (``DataItem``) to be pushed into the graph row
by row, or in batches. Mutually exclusive with ``data_object``.
:param data_object: A single dict instance to run the graph with exactly
once. Must be a ``dict`` (subclasses accepted). Mutually exclusive with
``data``. If you have a Pydantic model, call ``.model_dump()`` first;
if you have a JSON string, ``json.loads()`` it before passing.
:param timestamp_column: The name of the column that will be used as the
timestamp for model monitoring purposes. On the ``data`` path, when
used in conjunction with ``batching``, the first timestamp will be
used for the entire batch. On the ``data_object`` path, the value is
read from the root of the dict - in this case it should be a string
that is a serialized ``datetime`` (for example using ``isoformat()``.
:param batching: Whether to push one or more batches into the graph rather
than row by row. Ignored on the ``data_object`` path.
: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. Ignored on the ``data_object`` path.
:param read_as_lists: Whether to read each row (or, on the ``data_object``
path, the dict's first-level values) as a list instead of a dictionary.
:param nest_under_inputs: Whether to wrap each event body with
``{"inputs": ...}``.
:return: On the ``data_object`` path, the graph's response. On the ``data`` path, ``None``.
"""
if _is_inside_asyncio_loop():
_workaround_asyncio_nesting()
return asyncio.run(
async_execute_graph(
context,
data,
data_object,
timestamp_column,
batching,
batch_size,
read_as_lists,
nest_under_inputs,
)
)
def _set_callbacks(server, context):
if not server.graph.supports_termination() or not hasattr(context, "platform"):
return
if hasattr(context.platform, "set_termination_callback"):
context.logger.info(
"Setting termination callback to terminate graph on worker shutdown"
)
async def termination_callback():
context.logger.info("Termination callback called")
maybe_coroutine = server.wait_for_completion()
if asyncio.iscoroutine(maybe_coroutine):
await maybe_coroutine
context.logger.info("Termination of async flow is completed")
context.platform.set_termination_callback(termination_callback)
if hasattr(context.platform, "set_drain_callback"):
context.logger.info(
"Setting drain callback to terminate and restart the graph on a drain event (such as rebalancing)"
)
async def drain_callback():
context.logger.info("Drain callback called")
maybe_coroutine = server.wait_for_completion()
if asyncio.iscoroutine(maybe_coroutine):
await maybe_coroutine
context.logger.info(
"Termination of async flow is completed. Rerunning async flow."
)
# Rerun the flow without reconstructing it
server.graph._run_async_flow()
context.logger.info("Async flow restarted")
context.platform.set_drain_callback(drain_callback)
def _preprocess_event(context, event):
"""Preprocess event before running through the graph.
Handles Nuclio workarounds for empty body and stream path setup.
"""
if context._server.http_trigger:
# Workaround for a Nuclio bug where it sometimes passes b'' instead of None due to dirty memory
if event.body == b"":
event.body = None
# original path is saved in stream_path so it can be used by explicit ack, but path is reset to / as a
# workaround for NUC-178
# nuclio 1.12.12 added the topic attribute, and we must use it as part of the fix for NUC-233
# TODO: Remove fallback on event.path once support for nuclio<1.12.12 is dropped
event.stream_path = getattr(event, "topic", event.path)
if hasattr(event, "trigger") and event.trigger.kind in (
"kafka",
"kafka-cluster",
"v3ioStream",
"v3io-stream",
"rabbit-mq",
"rabbitMq",
):
event.path = "/"
def _process_single_response(context, response, get_body):
if (
isinstance(context, MLClientCtx)
or isinstance(response, context.Response)
or get_body
):
return response
if response and not isinstance(response, str | bytes):
body = json.dumps(response)
return context.Response(
body=body, content_type="application/json", status_code=200
)
return response
async def _process_single_async_response(context, response, get_body):
return _process_single_response(context, await response, get_body)
def v2_serving_handler(context, event, get_body=False):
"""Standard handler for non-streaming serving functions."""
_preprocess_event(context, event)
response = context._server.run(event, context, get_body)
if asyncio.iscoroutine(response):
return _process_single_async_response(context, response, get_body)
return _process_single_response(context, response, get_body)
async def v2_serving_streaming_handler(context, event, get_body=False):
"""Async streaming handler for nuclio that yields results as they arrive.
This handler is used when streaming mode is enabled on the serving function.
It yields results from streaming steps in the graph as they are produced,
allowing for real-time streaming responses (e.g., for LLM token streaming).
The handler is an async generator function that nuclio recognizes and handles
appropriately, streaming responses back to the HTTP client.
"""
_preprocess_event(context, event)
response = context._server.run(event, context, get_body)
# Unwrap coroutines to get the actual result
if asyncio.iscoroutine(response):
response = await response
# Yield chunks from the response
if inspect.isasyncgen(response):
async for chunk in response:
yield chunk
elif inspect.isgenerator(response):
for chunk in response:
yield chunk
else:
yield response
[docs]
def create_graph_server(
parameters=None,
load_mode=None,
graph=None,
verbose=False,
current_function=None,
**kwargs,
) -> GraphServer:
"""create graph server host/emulator for local or test runs
Usage example::
server = create_graph_server(graph=RouterStep(), parameters={})
server.init(None, globals())
server.graph.add_route(
"my", class_name=MyModelClass, model_path="{path}", z=100
)
print(server.test("/v2/models/my/infer", testdata))
"""
parameters = parameters or {}
server = GraphServer(graph, parameters, load_mode, verbose=verbose, **kwargs)
server.set_current_function(
current_function or os.getenv("SERVING_CURRENT_FUNCTION", "")
)
return server
class MockTrigger:
"""mock nuclio event trigger"""
def __init__(self, kind="", name=""):
self.kind = kind
self.name = name
class MockEvent:
"""mock basic nuclio event object"""
def __init__(
self,
body=None,
content_type=None,
headers=None,
method=None,
path=None,
event_id=None,
trigger: MockTrigger = None,
offset=None,
time=None,
):
self.id = event_id or uuid.uuid4().hex
self.key = ""
self.body = body
# optional
self.headers = headers or {}
self.method = method
self.path = path or "/"
self.content_type = content_type
self.error = None
self.trigger = trigger or MockTrigger()
self.offset = offset or 0
def __str__(self):
error = f", error={self.error}" if self.error else ""
return f"Event(id={self.id}, body={self.body}, method={self.method}, path={self.path}{error})"
class Response:
def __init__(self, headers=None, body=None, content_type=None, status_code=200):
self.headers = headers or {}
self.body = body
self.status_code = status_code
self.content_type = content_type or "text/plain"
def __repr__(self):
cls = self.__class__.__name__
items = self.__dict__.items()
args = (f"{key}={repr(value)}" for key, value in items)
args_str = ", ".join(args)
return f"{cls}({args_str})"
[docs]
class GraphContext:
"""Graph context object"""
def __init__(
self,
level="info", # Unused argument
logger=None,
server=None,
nuclio_context: NuclioContext | None = None,
) -> None:
self.state = None
self.logger = logger
self.worker_id = 0
self.Response = Response
self.verbose = False
self.stream = None
self.root = None
self.executor: storey.flow.RunnableExecutor | None = None
if nuclio_context:
self.logger: NuclioLogger = nuclio_context.logger
self.Response = nuclio_context.Response
if hasattr(nuclio_context, "trigger") and hasattr(
nuclio_context.trigger, "kind"
):
self.trigger = nuclio_context.trigger.kind
self.worker_id = nuclio_context.worker_id
if hasattr(nuclio_context, "platform"):
self.platform = nuclio_context.platform
elif not logger:
self.logger: mlrun.utils.Logger = mlrun.utils.logger
self._server = server
self.current_function = None
self.get_store_resource = None
self.get_table = None
self.is_mock = False
self.monitoring_mock = False
self._project_obj = None
@property
def server(self):
return self._server
@property
def project_obj(self):
if not self._project_obj:
self._project_obj = mlrun.get_run_db().get_project(name=self.project)
return self._project_obj
@property
def project(self) -> str:
"""current project name (for the current function)"""
project, _, _, _ = mlrun.common.helpers.parse_versioned_object_uri(
self._server.function_uri
)
return project
[docs]
def push_error(self, event, message, source=None, **kwargs):
if self.verbose:
self.logger.error(
f"got error from {source} state:\n{event.body}\n{message}"
)
if self._server and self._server._error_stream_object:
try:
message = format_error(
self._server, self, source, event, message, kwargs
)
self._server._error_stream_object.push(message)
except Exception as ex:
message = traceback.format_exc()
self.logger.error(f"failed to write to error stream: {ex}\n{message}")
[docs]
def get_param(self, key: str, default=None):
if self._server and self._server.parameters:
return self._server.parameters.get(key, default)
return default
[docs]
def get_secret(self, key: str):
if self._server and self._server._secrets:
return self._server._secrets.get(key)
return None
[docs]
def get_remote_endpoint(self, name, external=True):
"""return the remote nuclio/serving function http(s) endpoint given its name
:param name: the function name/uri in the form [project/]function-name[:tag]
:param external: return the external url (returns the external url by default)
"""
if "://" in name:
return name
project, uri, tag, _ = mlrun.common.helpers.parse_versioned_object_uri(
self._server.function_uri
)
if name.startswith("."):
name = f"{uri}-{name[1:]}"
else:
project, name, tag, _ = mlrun.common.helpers.parse_versioned_object_uri(
name, project
)
(
state,
fullname,
_,
_,
_,
function_status,
) = mlrun.runtimes.nuclio.function.get_nuclio_deploy_status(name, project, tag)
if state in ["error", "unhealthy"]:
raise ValueError(
f"Nuclio function {fullname} is in error state, cannot be accessed"
)
key = "externalInvocationUrls" if external else "internalInvocationUrls"
urls = function_status.get(key)
if not urls:
raise ValueError(f"cannot read {key} for nuclio function {fullname}")
return f"http://{urls[0]}"
def format_error(server, context, source, event, message, args):
return {
"function_uri": server.function_uri,
"worker": context.worker_id,
"host": socket.gethostname(),
"source": source,
"event": {"id": event.id, "body": event.body},
"message": message,
"args": args,
}
