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

mlrun/model.py

@@ -935,6 +935,41 @@ class HyperParamOptions(ModelObj):
             )
 
 
+class RetryBackoff(ModelObj):
+    """Backoff strategy for retries."""
+
+    def __init__(self, base_delay: Optional[str] = None):
+        # The base_delay time string must conform to timelength python package standards and be at least
+        # mlrun.mlconf.function.spec.retry.backoff.min_base_delay (e.g. 1000s, 1 hour 30m, 1h etc.).
+        self.base_delay = (
+            base_delay or mlrun.mlconf.function.spec.retry.backoff.default_base_delay
+        )
+
+
+class Retry(ModelObj):
+    """Retry configuration"""
+
+    def __init__(
+        self,
+        count: int = 0,
+        backoff: typing.Union[RetryBackoff, dict] = None,
+    ):
+        # Set to None if count is 0 to eliminate the retry configuration from the dictionary representation.
+        self.count = count or None
+        self.backoff = backoff
+
+    @property
+    def backoff(self) -> Optional[RetryBackoff]:
+        if not self.count:
+            # Retry is not configured, return None
+            return None
+        return self._backoff
+
+    @backoff.setter
+    def backoff(self, backoff):
+        self._backoff = self._verify_dict(backoff, "backoff", RetryBackoff)
+
+
 class RunSpec(ModelObj):
     """Run specification"""
 
@@ -971,6 +1006,7 @@ class RunSpec(ModelObj):
         node_selector=None,
         tolerations=None,
         affinity=None,
+        retry=None,
     ):
         # A dictionary of parsing configurations that will be read from the inputs the user set. The keys are the inputs
         # keys (parameter names) and the values are the type hint given in the input keys after the colon.
@@ -1011,6 +1047,7 @@ class RunSpec(ModelObj):
         self.node_selector = node_selector or {}
         self.tolerations = tolerations or {}
         self.affinity = affinity or {}
+        self.retry = retry or {}
 
     def _serialize_field(
         self, struct: dict, field_name: Optional[str] = None, strip: bool = False
@@ -1212,6 +1249,14 @@ class RunSpec(ModelObj):
         self._verify_dict(state_thresholds, "state_thresholds")
         self._state_thresholds = state_thresholds
 
+    @property
+    def retry(self) -> Retry:
+        return self._retry
+
+    @retry.setter
+    def retry(self, retry: typing.Union[Retry, dict]):
+        self._retry = self._verify_dict(retry, "retry", Retry)
+
     def extract_type_hints_from_inputs(self):
         """
         This method extracts the type hints from the input keys in the input dictionary.
@@ -1329,6 +1374,7 @@ class RunStatus(ModelObj):
         reason: Optional[str] = None,
         notifications: Optional[dict[str, Notification]] = None,
         artifact_uris: Optional[dict[str, str]] = None,
+        retry_count: Optional[int] = None,
     ):
         self.state = state or "created"
         self.status_text = status_text
@@ -1346,6 +1392,7 @@ class RunStatus(ModelObj):
         self.notifications = notifications or {}
         # Artifact key -> URI mapping, since the full artifacts are not stored in the runs DB table
         self._artifact_uris = artifact_uris or {}
+        self._retry_count = retry_count or None
 
     @classmethod
     def from_dict(
@@ -1399,6 +1446,21 @@ class RunStatus(ModelObj):
 
         self._artifact_uris = resolved_artifact_uris
 
+    @property
+    def retry_count(self) -> Optional[int]:
+        """
+        The number of retries that were made for this run.
+        """
+        return self._retry_count
+
+    @retry_count.setter
+    def retry_count(self, retry_count: int):
+        """
+        Set the number of retries that were made for this run.
+        :param retry_count: The number of retries.
+        """
+        self._retry_count = retry_count
+
     def is_failed(self) -> Optional[bool]:
         """
         This method returns whether a run has failed.
@@ -2026,6 +2088,7 @@ def new_task(
     secrets=None,
     base=None,
     returns=None,
+    retry=None,
 ) -> RunTemplate:
     """Creates a new task
 
@@ -2061,6 +2124,7 @@ def new_task(
                             * A dictionary of configurations to use when logging. Further info per object type and
                               artifact type can be given there. The artifact key must appear in the dictionary as
                               "key": "the_key".
+    :param retry:           Retry configuration for the run, can be a dict or an instance of mlrun.model.Retry.
     """
 
     if base:
@@ -2086,6 +2150,7 @@ def new_task(
     run.spec.hyper_param_options.selector = (
         selector or run.spec.hyper_param_options.selector
     )
+    run.spec.retry = retry or run.spec.retry
     return run
 
 

mlrun/package/packagers_manager.py

@@ -21,6 +21,7 @@ from typing import Any, Optional, Union
 
 import mlrun.errors
 from mlrun.artifacts import Artifact
+from mlrun.artifacts.base import verify_target_path
 from mlrun.datastore import DataItem, get_store_resource, store_manager
 from mlrun.errors import MLRunInvalidArgumentError
 from mlrun.utils import logger
@@ -276,6 +277,7 @@ class PackagersManager:
         if data_item.get_artifact_type():
             # Get the artifact object in the data item:
             artifact, _ = store_manager.get_store_artifact(url=data_item.artifact_url)
+            verify_target_path(artifact)
             # Get the key from the artifact's metadata and instructions from the artifact's spec:
             artifact_key = artifact.metadata.key
             packaging_instructions = artifact.spec.unpackaging_instructions

mlrun/projects/operations.py

@@ -20,7 +20,6 @@ import mlrun
 import mlrun.common.constants as mlrun_constants
 import mlrun.common.schemas.function
 import mlrun.common.schemas.workflow
-import mlrun_pipelines.common.models
 import mlrun_pipelines.models
 from mlrun.utils import hub_prefix
 
@@ -82,6 +81,7 @@ def run_function(
     builder_env: Optional[list] = None,
     reset_on_run: Optional[bool] = None,
     output_path: Optional[str] = None,
+    retry: Optional[Union[mlrun.model.Retry, dict]] = None,
 ) -> Union[mlrun.model.RunObject, mlrun_pipelines.models.PipelineNodeWrapper]:
     """Run a local or remote task as part of a local/kubeflow pipeline
 
@@ -177,6 +177,7 @@ def run_function(
                             This ensures latest code changes are executed. This argument must be used in
                             conjunction with the local=True argument.
     :param output_path:     path to store artifacts, when running in a workflow this will be set automatically
+    :param retry:           Retry configuration for the run, can be a dict or an instance of mlrun.model.Retry.
     :return: MLRun RunObject or PipelineNodeWrapper
     """
     if artifact_path:
@@ -197,6 +198,7 @@ def run_function(
         returns=returns,
         base=base_task,
         selector=selector,
+        retry=retry,
     )
     task.spec.verbose = task.spec.verbose or verbose
 
@@ -205,6 +207,11 @@ def run_function(
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "Scheduling jobs is not supported when running a workflow with the kfp engine."
             )
+        if retry:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "Retrying jobs is not supported when running a workflow with the kfp engine. "
+                "Use KFP set_retry instead."
+            )
         return function.as_step(
             name=name, runspec=task, workdir=workdir, outputs=outputs, labels=labels
         )

mlrun/projects/project.py

@@ -159,7 +159,8 @@ def new_project(
     parameters: Optional[dict] = None,
     default_function_node_selector: Optional[dict] = None,
 ) -> "MlrunProject":
-    """Create a new MLRun project, optionally load it from a yaml/zip/git template
+    """Create a new MLRun project, optionally load it from a yaml/zip/git template.
+    The project will become the active project for the current session.
 
     A new project is created and returned, you can customize the project by placing a project_setup.py file
     in the project root dir, it will be executed upon project creation or loading.
@@ -326,7 +327,8 @@ def load_project(
     parameters: Optional[dict] = None,
     allow_cross_project: Optional[bool] = None,
 ) -> "MlrunProject":
-    """Load an MLRun project from git or tar or dir
+    """Load an MLRun project from git or tar or dir. The project will become the active project for
+    the current session.
 
     MLRun looks for a project.yaml file with project definition and objects in the project root path
     and use it to initialize the project, in addition it runs the project_setup.py file (if it exists)
@@ -2688,8 +2690,8 @@ class MlrunProject(ModelObj):
         requirements_file: str = "",
     ) -> mlrun.runtimes.BaseRuntime:
         """
-        | Update or add a function object to the project.
-        | Function can be provided as an object (func) or a .py/.ipynb/.yaml URL.
+        Update or add a function object to the project.
+        Function can be provided as an object (func) or a .py/.ipynb/.yaml URL.
 
         | Creating a function from a single file is done by specifying ``func`` and disabling ``with_repo``.
         | Creating a function with project source (specify ``with_repo=True``):
@@ -2734,6 +2736,20 @@ class MlrunProject(ModelObj):
             # By providing a path to a pip requirements file
             proj.set_function("my.py", requirements="requirements.txt")
 
+        One of the most important parameters is 'kind', used to specify the chosen runtime. The options are:
+           - local: execute a local python or shell script
+           - job: insert the code into a Kubernetes pod and execute it
+           - nuclio: insert the code into a real-time serverless nuclio function
+           - serving: insert code into orchestrated nuclio function(s) forming a DAG
+           - dask: run the specified python code / script as Dask Distributed job
+           - mpijob: run distributed Horovod jobs over the MPI job operator
+           - spark: run distributed Spark job using Spark Kubernetes Operator
+           - remote-spark: run distributed Spark job on remote Spark service
+           - databricks: run code on Databricks cluster (python scripts, Spark etc.)
+           - application: run a long living application (e.g. a web server, UI, etc.)
+
+        Learn more about :doc:`../../concepts/functions-overview`.
+
         :param func:                Function object or spec/code url, None refers to current Notebook
         :param name:                Name of the function (under the project), can be specified with a tag to support
                                     Versions (e.g. myfunc:v1). If the `tag` parameter is provided, the tag in the name
@@ -3967,6 +3983,7 @@ class MlrunProject(ModelObj):
         builder_env: Optional[dict] = None,
         reset_on_run: Optional[bool] = None,
         output_path: Optional[str] = None,
+        retry: Optional[Union[mlrun.model.Retry, dict]] = None,
     ) -> typing.Union[mlrun.model.RunObject, PipelineNodeWrapper]:
         """Run a local or remote task as part of a local/kubeflow pipeline
 
@@ -4029,7 +4046,7 @@ class MlrunProject(ModelObj):
                                 This ensures latest code changes are executed. This argument must be used in
                                 conjunction with the local=True argument.
         :param output_path:     path to store artifacts, when running in a workflow this will be set automatically
-
+        :param retry:           Retry configuration for the run, can be a dict or an instance of mlrun.model.Retry.
         :return: MLRun RunObject or PipelineNodeWrapper
         """
         if artifact_path:
@@ -4068,6 +4085,7 @@ class MlrunProject(ModelObj):
                 returns=returns,
                 builder_env=builder_env,
                 reset_on_run=reset_on_run,
+                retry=retry,
             )
 
     def build_function(

mlrun/run.py

@@ -36,6 +36,7 @@ import mlrun.common.schemas
 import mlrun.errors
 import mlrun.utils.helpers
 import mlrun_pipelines.utils
+from mlrun.datastore.model_provider.model_provider import ModelProvider
 from mlrun_pipelines.common.models import RunStatuses
 from mlrun_pipelines.common.ops import format_summary_from_kfp_run, show_kfp_run
 
@@ -1152,6 +1153,22 @@ def get_dataitem(url, secrets=None, db=None) -> "DataItem":
     return stores.object(url=url)
 
 
+def get_model_provider(
+    url,
+    secrets=None,
+    db=None,
+    default_invoke_kwargs: Optional[dict] = None,
+    raise_missing_schema_exception=True,
+) -> ModelProvider:
+    """get mlrun dataitem object (from path/url)"""
+    store_manager.set(secrets, db=db)
+    return store_manager.model_provider_object(
+        url=url,
+        default_invoke_kwargs=default_invoke_kwargs,
+        raise_missing_schema_exception=raise_missing_schema_exception,
+    )
+
+
 def download_object(url, target, secrets=None):
     """download mlrun dataitem (from path/url to target path)"""
     stores = store_manager.set(secrets)

mlrun/runtimes/__init__.py

@@ -148,6 +148,12 @@ class RuntimeKinds:
             "",
         ]
 
+    @staticmethod
+    def retriable_runtimes():
+        return [
+            RuntimeKinds.job,
+        ]
+
     @staticmethod
     def nuclio_runtimes():
         return [

mlrun/runtimes/base.py

@@ -33,6 +33,7 @@ import mlrun.launcher.factory
 import mlrun.utils.helpers
 import mlrun.utils.notifications
 import mlrun.utils.regex
+from mlrun.common.runtimes.constants import RunStates
 from mlrun.model import (
     BaseMetadata,
     HyperParamOptions,
@@ -319,6 +320,7 @@ class BaseRuntime(ModelObj):
         state_thresholds: Optional[dict[str, int]] = None,
         reset_on_run: Optional[bool] = None,
         output_path: Optional[str] = "",
+        retry: Optional[Union[mlrun.model.Retry, dict]] = None,
         **launcher_kwargs,
     ) -> RunObject:
         """
@@ -377,6 +379,7 @@ class BaseRuntime(ModelObj):
                              This ensures latest code changes are executed. This argument must be used in
                              conjunction with the local=True argument.
         :param output_path:    Default artifact output path.
+        :param retry:          Retry configuration for the run, can be a dict or an instance of mlrun.model.Retry.
         :return: Run context object (RunObject) with run metadata, results and status
         """
         if artifact_path or out_path:
@@ -414,6 +417,7 @@ class BaseRuntime(ModelObj):
             returns=returns,
             state_thresholds=state_thresholds,
             reset_on_run=reset_on_run,
+            retry=retry,
         )
 
     def _get_db_run(
@@ -570,12 +574,27 @@ class BaseRuntime(ModelObj):
         updates = None
         last_state = get_in(resp, "status.state", "")
         kind = get_in(resp, "metadata.labels.kind", "")
-        if last_state == "error" or err:
+        if last_state in RunStates.error_states() or err:
+            new_state = RunStates.error
+            status_text = None
+            max_retries = get_in(resp, "spec.retry.count", 0)
+            retry_count = get_in(resp, "status.retry_count", 0) or 0
+            attempts = retry_count + 1
+            if max_retries:
+                if retry_count < max_retries:
+                    new_state = RunStates.pending_retry
+                    status_text = f"Run failed attempt {attempts} of {max_retries + 1}"
+                elif retry_count >= max_retries:
+                    status_text = f"Run failed after {attempts} attempts"
+
             updates = {
                 "status.last_update": now_date().isoformat(),
-                "status.state": "error",
+                "status.state": new_state,
             }
-            update_in(resp, "status.state", "error")
+            update_in(resp, "status.state", new_state)
+            if status_text:
+                updates["status.status_text"] = status_text
+                update_in(resp, "status.status_text", status_text)
             if err:
                 update_in(resp, "status.error", err_to_str(err))
             err = get_in(resp, "status.error")
@@ -584,9 +603,8 @@ class BaseRuntime(ModelObj):
 
         elif (
             not was_none
-            and last_state != mlrun.common.runtimes.constants.RunStates.completed
-            and last_state
-            not in mlrun.common.runtimes.constants.RunStates.error_and_abortion_states()
+            and last_state != RunStates.completed
+            and last_state not in RunStates.error_and_abortion_states()
         ):
             try:
                 runtime_cls = mlrun.runtimes.get_runtime_class(kind)

mlrun/runtimes/daskjob.py

@@ -505,6 +505,7 @@ class DaskCluster(KubejobRuntime):
         state_thresholds: Optional[dict[str, int]] = None,
         reset_on_run: Optional[bool] = None,
         output_path: Optional[str] = "",
+        retry: Optional[Union[mlrun.model.Retry, dict]] = None,
         **launcher_kwargs,
     ) -> RunObject:
         if state_thresholds:

mlrun/runtimes/databricks_job/databricks_runtime.py

@@ -233,6 +233,7 @@ def run_mlrun_databricks_job(context,task_parameters: dict, **kwargs):
         state_thresholds: Optional[dict[str, int]] = None,
         reset_on_run: Optional[bool] = None,
         output_path: Optional[str] = "",
+        retry: Optional[Union[mlrun.model.Retry, dict]] = None,
         **launcher_kwargs,
     ) -> RunObject:
         if local:

mlrun/runtimes/local.py

@@ -34,6 +34,7 @@ from nuclio import Event
 
 import mlrun
 import mlrun.common.constants as mlrun_constants
+import mlrun.common.runtimes.constants
 from mlrun.lists import RunList
 
 from ..errors import err_to_str
@@ -315,15 +316,9 @@ class LocalRuntime(BaseRuntime, ParallelRunner):
                 return context.to_dict()
 
             # if RunError was raised it means that the error was raised as part of running the function
-            # ( meaning the state was already updated to error ) therefore we just re-raise the error
             except RunError as err:
                 raise err
-            # this exception handling is for the case where we fail on pre-loading or post-running the function
-            # and the state was not updated to error yet, therefore we update the state to error and raise as RunError
             except Exception as exc:
-                # set_state here is mainly for sanity, as we will raise RunError which is expected to be handled
-                # by the caller and will set the state to error ( in `update_run_state` )
-                context.set_state(error=err_to_str(exc), commit=True)
                 logger.error(f"Run error, {traceback.format_exc()}")
                 raise RunError(
                     "Failed on pre-loading / post-running of the function"

mlrun/serving/server.py

@@ -395,7 +395,6 @@ def add_monitoring_general_steps(
     monitor_flow_step = graph.add_step(
         "mlrun.serving.system_steps.BackgroundTaskStatus",
         "background_task_status_step",
-        context=context,
         model_endpoint_creation_strategy=mlrun.common.schemas.ModelEndpointCreationStrategy.SKIP,
     )
     graph.add_step(
@@ -410,7 +409,6 @@ def add_monitoring_general_steps(
         "monitoring_pre_processor_step",
         after="filter_none",
         full_event=True,
-        context=context,
         model_endpoint_creation_strategy=mlrun.common.schemas.ModelEndpointCreationStrategy.SKIP,
     )
     # flatten the events

mlrun/serving/states.py

@@ -35,6 +35,7 @@ from storey import ParallelExecutionMechanisms
 import mlrun
 import mlrun.artifacts
 import mlrun.common.schemas as schemas
+from mlrun.artifacts.llm_prompt import LLMPromptArtifact
 from mlrun.artifacts.model import ModelArtifact
 from mlrun.datastore.datastore_profile import (
     DatastoreProfileKafkaSource,
@@ -42,6 +43,7 @@ from mlrun.datastore.datastore_profile import (
     DatastoreProfileV3io,
     datastore_profile_read,
 )
+from mlrun.datastore.model_provider.model_provider import ModelProvider
 from mlrun.datastore.store_resources import get_store_resource
 from mlrun.datastore.storeytargets import KafkaStoreyTarget, StreamStoreyTarget
 from mlrun.utils import logger
@@ -1019,6 +1021,9 @@ class Model(storey.ParallelExecutionRunnable, ModelObj):
         if artifact_uri is not None and not isinstance(artifact_uri, str):
             raise MLRunInvalidArgumentError("'artifact_uri' argument must be a string")
         self.artifact_uri = artifact_uri
+        self.invocation_artifact: Optional[LLMPromptArtifact] = None
+        self.model_artifact: Optional[ModelArtifact] = None
+        self.model_provider: Optional[ModelProvider] = None
 
     def __init_subclass__(cls):
         super().__init_subclass__()
@@ -1030,12 +1035,27 @@ class Model(storey.ParallelExecutionRunnable, ModelObj):
 
     def load(self) -> None:
         """Override to load model if needed."""
-        pass
+        self._load_artifacts()
+        if self.model_artifact:
+            self.model_provider = mlrun.get_model_provider(
+                url=self.model_artifact.model_url,
+                default_invoke_kwargs=self.model_artifact.default_config,
+                raise_missing_schema_exception=False,
+            )
+
+    def _load_artifacts(self) -> None:
+        artifact = self._get_artifact_object()
+        if isinstance(artifact, LLMPromptArtifact):
+            self.invocation_artifact = artifact
+            self.model_artifact = self.invocation_artifact.model_artifact
+        else:
+            self.model_artifact = artifact
 
-    def _get_artifact_object(self) -> Union[ModelArtifact, None]:
+    def _get_artifact_object(self) -> Union[ModelArtifact, LLMPromptArtifact, None]:
         if self.artifact_uri:
             if mlrun.datastore.is_store_uri(self.artifact_uri):
-                return get_store_resource(self.artifact_uri)
+                artifact, _ = mlrun.store_manager.get_store_artifact(self.artifact_uri)
+                return artifact
             else:
                 raise ValueError(
                     "Could not get artifact, 'artifact_uri' must be a valid artifact store URI"
@@ -1058,7 +1078,7 @@ class Model(storey.ParallelExecutionRunnable, ModelObj):
         return self.predict(body)
 
     async def run_async(self, body: Any, path: str) -> Any:
-        return self.predict(body)
+        return await self.predict_async(body)
 
     def get_local_model_path(self, suffix="") -> (str, dict):
         """
@@ -1223,7 +1243,9 @@ class ModelRunnerStep(MonitoredStep):
         endpoint_name: str,
         model_class: Union[str, Model],
         execution_mechanism: Union[str, ParallelExecutionMechanisms],
-        model_artifact: Optional[Union[str, mlrun.artifacts.ModelArtifact]] = None,
+        model_artifact: Optional[
+            Union[str, mlrun.artifacts.ModelArtifact, mlrun.artifacts.LLMPromptArtifact]
+        ] = None,
         labels: Optional[Union[list[str], dict[str, str]]] = None,
         creation_strategy: Optional[
             schemas.ModelEndpointCreationStrategy
@@ -1407,6 +1429,9 @@ class ModelRunnerStep(MonitoredStep):
             return monitoring_data
 
     def init_object(self, context, namespace, mode="sync", reset=False, **extra_kwargs):
+        if not self._is_local_function(context):
+            # skip init of non local functions
+            return
         model_selector = self.class_args.get("model_selector")
         execution_mechanism_by_model_name = self.class_args.get(
             schemas.ModelRunnerStepData.MODEL_TO_EXECUTION_MECHANISM

mlrun/serving/system_steps.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import random
-from copy import copy, deepcopy
+from copy import deepcopy
 from datetime import timedelta
 from typing import Any, Optional, Union
 
@@ -32,11 +32,12 @@ class MonitoringPreProcessor(storey.MapClass):
 
     def __init__(
         self,
-        context,
         **kwargs,
     ):
         super().__init__(**kwargs)
-        self.context = copy(context)
+        self.server: mlrun.serving.GraphServer = (
+            getattr(self.context, "server", None) if self.context else None
+        )
 
     def reconstruct_request_resp_fields(
         self, event, model: str, model_monitoring_data: dict
@@ -148,9 +149,8 @@ class MonitoringPreProcessor(storey.MapClass):
 
     def do(self, event):
         monitoring_event_list = []
-        server: mlrun.serving.GraphServer = getattr(self.context, "server", None)
         model_runner_name = event._metadata.get("model_runner_name", "")
-        step = server.graph.steps[model_runner_name] if server else {}
+        step = self.server.graph.steps[model_runner_name] if self.server else {}
         monitoring_data = step.monitoring_data
         logger.debug(
             "monitoring preprocessor started",
@@ -184,8 +184,8 @@ class MonitoringPreProcessor(storey.MapClass):
                             mm_schemas.StreamProcessingEvent.LABELS: monitoring_data[
                                 model
                             ].get(mlrun.common.schemas.MonitoringData.OUTPUTS),
-                            mm_schemas.StreamProcessingEvent.FUNCTION_URI: server.function_uri
-                            if server
+                            mm_schemas.StreamProcessingEvent.FUNCTION_URI: self.server.function_uri
+                            if self.server
                             else None,
                             mm_schemas.StreamProcessingEvent.REQUEST: request,
                             mm_schemas.StreamProcessingEvent.RESPONSE: resp,
@@ -226,8 +226,8 @@ class MonitoringPreProcessor(storey.MapClass):
                     mm_schemas.StreamProcessingEvent.LABELS: monitoring_data[model].get(
                         mlrun.common.schemas.MonitoringData.OUTPUTS
                     ),
-                    mm_schemas.StreamProcessingEvent.FUNCTION_URI: server.function_uri
-                    if server
+                    mm_schemas.StreamProcessingEvent.FUNCTION_URI: self.server.function_uri
+                    if self.server
                     else None,
                     mm_schemas.StreamProcessingEvent.REQUEST: request,
                     mm_schemas.StreamProcessingEvent.RESPONSE: resp,
@@ -253,19 +253,17 @@ class BackgroundTaskStatus(storey.MapClass):
     creation failed or in progress
     """
 
-    def __init__(self, context, **kwargs):
-        self.context = copy(context)
-        self.server: mlrun.serving.GraphServer = getattr(self.context, "server", None)
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.server: mlrun.serving.GraphServer = (
+            getattr(self.context, "server", None) if self.context else None
+        )
         self._background_task_check_timestamp = None
         self._background_task_state = mlrun.common.schemas.BackgroundTaskState.running
-        super().__init__(**kwargs)
 
     def do(self, event):
-        if (self.context and self.context.is_mock) or self.context is None:
-            return event
         if self.server is None:
             return None
-
         if (
             self._background_task_state
             == mlrun.common.schemas.BackgroundTaskState.running
@@ -283,19 +281,14 @@ class BackgroundTaskStatus(storey.MapClass):
             self._background_task_check_timestamp = mlrun.utils.now_date()
             self._log_background_task_state(background_task.status.state)
             self._background_task_state = background_task.status.state
-            if (
-                background_task.status.state
-                == mlrun.common.schemas.BackgroundTaskState.succeeded
-            ):
-                return event
-            else:
-                return None
-        elif (
+
+        if (
             self._background_task_state
-            == mlrun.common.schemas.BackgroundTaskState.failed
+            == mlrun.common.schemas.BackgroundTaskState.succeeded
         ):
+            return event
+        else:
             return None
-        return event
 
     def _log_background_task_state(
         self, background_task_state: mlrun.common.schemas.BackgroundTaskState
@@ -382,9 +375,10 @@ class SamplingStep(storey.MapClass):
 
 
 class MockStreamPusher(storey.MapClass):
-    def __init__(self, context, output_stream=None, **kwargs):
+    def __init__(self, output_stream=None, **kwargs):
         super().__init__(**kwargs)
-        self.output_stream = output_stream or context.stream.output_stream
+        stream = self.context.stream if self.context else None
+        self.output_stream = output_stream or stream.output_stream
 
     def do(self, event):
         self.output_stream.push(