flyte 0.2.0b1__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/_internal/resolvers/default.py

@@ -1,6 +1,6 @@
 import importlib
 from pathlib import Path
-from typing import List, Optional
+from typing import List
 
 from flyte._internal.resolvers._task_module import extract_task_module
 from flyte._internal.resolvers.common import Resolver
@@ -23,6 +23,6 @@ class DefaultTaskResolver(Resolver):
         task_def = getattr(task_module, task_name)
         return task_def
 
-    def loader_args(self, task: TaskTemplate, root_dir: Optional[Path] = None) -> List[str]:  # type:ignore
+    def loader_args(self, task: TaskTemplate, root_dir: Path) -> List[str]:  # type:ignore
         t, m = extract_task_module(task, root_dir)
         return ["mod", m, "instance", t]

flyte/_internal/runtime/convert.py

@@ -1,30 +1,40 @@
 from __future__ import annotations
 
+import asyncio
+import base64
+import hashlib
+import inspect
 from dataclasses import dataclass
-from typing import Any, Dict, Tuple, Union
+from types import NoneType
+from typing import Any, Dict, List, Tuple, Union, get_args
 
-from flyteidl.core import execution_pb2, literals_pb2
+from flyteidl2.core import execution_pb2, interface_pb2, literals_pb2
+from flyteidl2.task import common_pb2, task_definition_pb2
 
 import flyte.errors
 import flyte.storage as storage
-from flyte._datastructures import ActionID, NativeInterface, TaskContext
-from flyte._internal.controllers import pbhash
-from flyte._protos.workflow import run_definition_pb2
-from flyte.types import TypeEngine
+from flyte._context import ctx
+from flyte.models import ActionID, NativeInterface, TaskContext
+from flyte.types import TypeEngine, TypeTransformerFailedError
 
 
 @dataclass(frozen=True)
 class Inputs:
-    proto_inputs: run_definition_pb2.Inputs
+    proto_inputs: common_pb2.Inputs
 
     @classmethod
     def empty(cls) -> "Inputs":
-        return cls(proto_inputs=run_definition_pb2.Inputs())
+        return cls(proto_inputs=common_pb2.Inputs())
+
+    @property
+    def context(self) -> Dict[str, str]:
+        """Get the context as a dictionary."""
+        return {kv.key: kv.value for kv in self.proto_inputs.context}
 
 
 @dataclass(frozen=True)
 class Outputs:
-    proto_outputs: run_definition_pb2.Outputs
+    proto_outputs: common_pb2.Outputs
 
 
 @dataclass
@@ -56,43 +66,160 @@ async def convert_inputs_to_native(inputs: Inputs, python_interface: NativeInter
     return native_vals
 
 
-async def convert_from_native_to_inputs(interface: NativeInterface, *args, **kwargs) -> Inputs:
+async def convert_upload_default_inputs(interface: NativeInterface) -> List[common_pb2.NamedParameter]:
+    """
+    Converts the default inputs of a NativeInterface to a list of NamedParameters for upload.
+    This is used to upload default inputs to the Flyte backend.
+    """
+    if not interface.inputs:
+        return []
+
+    vars = []
+    literal_coros = []
+    for input_name, (input_type, default_value) in interface.inputs.items():
+        if default_value and default_value is not inspect.Parameter.empty:
+            lt = TypeEngine.to_literal_type(input_type)
+            literal_coros.append(TypeEngine.to_literal(default_value, input_type, lt))
+            vars.append((input_name, lt))
+
+    literals: List[literals_pb2.Literal] = await asyncio.gather(*literal_coros, return_exceptions=True)
+    named_params = []
+    for (name, lt), literal in zip(vars, literals):
+        if isinstance(literal, Exception):
+            raise RuntimeError(f"Failed to convert default value for parameter '{name}'") from literal
+        param = interface_pb2.Parameter(
+            var=interface_pb2.Variable(
+                type=lt,
+            ),
+            default=literal,
+        )
+        named_params.append(
+            common_pb2.NamedParameter(
+                name=name,
+                parameter=param,
+            ),
+        )
+    return named_params
+
+
+def is_optional_type(tp) -> bool:
+    """
+    True if the *annotation* `tp` is equivalent to Optional[…].
+    Works for Optional[T], Union[T, None], and T | None.
+    """
+    return NoneType in get_args(tp)  # fastest check
+
+
+async def convert_from_native_to_inputs(
+    interface: NativeInterface, *args, custom_context: Dict[str, str] | None = None, **kwargs
+) -> Inputs:
     kwargs = interface.convert_to_kwargs(*args, **kwargs)
-    if len(kwargs) == 0:
-        return Inputs.empty()
+
+    missing = [key for key in interface.required_inputs() if key not in kwargs]
+    if missing:
+        raise ValueError(f"Missing required inputs: {', '.join(missing)}")
+
+    # Read custom_context from TaskContext if available (inside task execution)
+    # Otherwise use the passed parameter (for remote run initiation)
+    context_kvs = None
+    tctx = ctx()
+    if tctx and tctx.custom_context:
+        # Inside a task - read from TaskContext
+        context_to_use = tctx.custom_context
+        context_kvs = [literals_pb2.KeyValuePair(key=k, value=v) for k, v in context_to_use.items()]
+    elif custom_context:
+        # Remote run initiation
+        context_kvs = [literals_pb2.KeyValuePair(key=k, value=v) for k, v in custom_context.items()]
+
+    if len(interface.inputs) == 0:
+        # Handle context even for empty inputs
+        return Inputs(proto_inputs=common_pb2.Inputs(context=context_kvs))
+
     # fill in defaults if missing
+    type_hints: Dict[str, type] = {}
+    already_converted_kwargs: Dict[str, literals_pb2.Literal] = {}
     for input_name, (input_type, default_value) in interface.inputs.items():
-        if input_name not in kwargs:
-            if default_value is not None:
+        if input_name in kwargs:
+            type_hints[input_name] = input_type
+        elif (
+            (default_value is not None and default_value is not inspect.Signature.empty)
+            or (default_value is None and is_optional_type(input_type))
+            or input_type is None
+            or input_type is type(None)
+        ):
+            if default_value == NativeInterface.has_default:
+                if interface._remote_defaults is None or input_name not in interface._remote_defaults:
+                    raise ValueError(f"Input '{input_name}' has a default value but it is not set in the interface.")
+                already_converted_kwargs[input_name] = interface._remote_defaults[input_name]
+            elif input_type is None or input_type is type(None):
+                # If the type is 'None' or 'class<None>', we assume it's a placeholder for no type
+                kwargs[input_name] = None
+                type_hints[input_name] = NoneType
+            else:
                 kwargs[input_name] = default_value
-            # todo: fill in Nones for optional inputs
-    if len(kwargs) < len(interface.inputs):
-        raise ValueError(
-            f"Received {len(kwargs)} inputs but interface has {len(interface.inputs)}. "
-            f"Please provide all required inputs."
-        )
-    literal_map = await TypeEngine.dict_to_literal_map(kwargs, interface.get_input_types())
+                type_hints[input_name] = input_type
+
+    literal_map = await TypeEngine.dict_to_literal_map(kwargs, type_hints)
+    if len(already_converted_kwargs) > 0:
+        copied_literals: Dict[str, literals_pb2.Literal] = {}
+        for k, v in literal_map.literals.items():
+            copied_literals[k] = v
+        # Add the already converted kwargs to the literal map
+        for k, v in already_converted_kwargs.items():
+            copied_literals[k] = v
+        literal_map = literals_pb2.LiteralMap(literals=copied_literals)
+
+    # Make sure we the interface, not literal_map or kwargs, because those may have a different order
     return Inputs(
-        proto_inputs=run_definition_pb2.Inputs(
-            literals=[run_definition_pb2.NamedLiteral(name=k, value=v) for k, v in literal_map.literals.items()]
+        proto_inputs=common_pb2.Inputs(
+            literals=[common_pb2.NamedLiteral(name=k, value=literal_map.literals[k]) for k in interface.inputs.keys()],
+            context=context_kvs,
         )
     )
 
 
-async def convert_from_native_to_outputs(o: Any, interface: NativeInterface) -> Outputs:
+async def convert_from_inputs_to_native(native_interface: NativeInterface, inputs: Inputs) -> Dict[str, Any]:
+    """
+    Converts the inputs from a run definition proto to a native Python dictionary.
+    :param native_interface: The native interface of the task.
+    :param inputs: The run definition inputs proto.
+    :return: A dictionary of input names to their native Python values.
+    """
+    if not inputs or not inputs.proto_inputs or not inputs.proto_inputs.literals:
+        return {}
+
+    literals = {named_literal.name: named_literal.value for named_literal in inputs.proto_inputs.literals}
+    return await TypeEngine.literal_map_to_kwargs(
+        literals_pb2.LiteralMap(literals=literals), native_interface.get_input_types()
+    )
+
+
+async def convert_from_native_to_outputs(o: Any, interface: NativeInterface, task_name: str = "") -> Outputs:
     # Always make it a tuple even if it's just one item to simplify logic below
     if not isinstance(o, tuple):
         o = (o,)
 
-    assert len(interface.outputs) == len(interface.outputs), (
-        f"Received {len(o)} outputs but interface has {len(interface.outputs)}"
-    )
+    if len(interface.outputs) == 0:
+        if len(o) != 0:
+            if len(o) == 1 and o[0] is not None:
+                raise flyte.errors.RuntimeDataValidationError(
+                    "o0",
+                    f"Expected no outputs but got {o},did you miss a return type annotation?",
+                    task_name,
+                )
+    else:
+        assert len(o) == len(interface.outputs), (
+            f"Received {len(o)} outputs but return annotation has {len(interface.outputs)} outputs specified. "
+        )
     named = []
     for (output_name, python_type), v in zip(interface.outputs.items(), o):
-        lit = await TypeEngine.to_literal(v, python_type, TypeEngine.to_literal_type(python_type))
-        named.append(run_definition_pb2.NamedLiteral(name=output_name, value=lit))
+        try:
+            lit = await TypeEngine.to_literal(v, python_type, TypeEngine.to_literal_type(python_type))
+            named.append(common_pb2.NamedLiteral(name=output_name, value=lit))
+        except TypeTransformerFailedError as e:
+            raise flyte.errors.RuntimeDataValidationError(output_name, e, task_name)
 
-    return Outputs(proto_outputs=run_definition_pb2.Outputs(literals=named))
+    return Outputs(proto_outputs=common_pb2.Outputs(literals=named))
 
 
 async def convert_outputs_to_native(interface: NativeInterface, outputs: Outputs) -> Union[Any, Tuple[Any, ...]]:
@@ -119,7 +246,7 @@ def convert_error_to_native(err: execution_pb2.ExecutionError | Exception | Erro
     if isinstance(err, Error):
         err = err.err
 
-    user_code, server_code = _clean_error_code(err.code)
+    user_code, _server_code = _clean_error_code(err.code)
     match err.kind:
         case execution_pb2.ExecutionError.UNKNOWN:
             return flyte.errors.RuntimeUnknownError(code=user_code, message=err.message, worker=err.worker)
@@ -185,21 +312,177 @@ def convert_from_native_to_error(err: BaseException) -> Error:
         )
 
 
-def generate_sub_action_id_and_output_path(tctx: TaskContext, task_name: str, inputs: Inputs) -> Tuple[ActionID, str]:
+def hash_data(data: Union[str, bytes]) -> str:
+    """
+    Generate a hash for the given data. If the data is a string, it will be encoded to bytes before hashing.
+    :param data: The data to hash, can be a string or bytes.
+    :return: A hexadecimal string representation of the hash.
+    """
+    if isinstance(data, str):
+        data = data.encode("utf-8")
+    digest = hashlib.sha256(data).digest()
+    return base64.b64encode(digest).decode("utf-8")
+
+
+def generate_inputs_hash(serialized_inputs: str | bytes) -> str:
+    """
+    Generate a hash for the inputs. This is used to uniquely identify the inputs for a task.
+    :return: A hexadecimal string representation of the hash.
+    """
+    return hash_data(serialized_inputs)
+
+
+def generate_inputs_repr_for_literal(literal: literals_pb2.Literal) -> bytes:
+    """
+    Generate a byte representation for a single literal that is meant to be hashed as part of the cache key
+    computation for an Action. This function should just serialize the literal deterministically, but will
+    use an existing hash value if present in the Literal.  This is trivial, except we need to handle nested literals
+    (inside collections and maps), that may have the hash property set.
+
+    :param literal: The literal to get a hashable representation for.
+    :return: byte representation of the literal that can be fed into a hash function.
+    """
+    # If the literal has a hash value, use that instead of serializing the full literal
+    if literal.hash:
+        return literal.hash.encode("utf-8")
+
+    if literal.HasField("collection"):
+        buf = bytearray()
+        for nested_literal in literal.collection.literals:
+            if nested_literal.hash:
+                buf += nested_literal.hash.encode("utf-8")
+            else:
+                buf += generate_inputs_repr_for_literal(nested_literal)
+
+        b = bytes(buf)
+        return b
+
+    elif literal.HasField("map"):
+        buf = bytearray()
+        # Sort keys to ensure deterministic ordering
+        for key in sorted(literal.map.literals.keys()):
+            nested_literal = literal.map.literals[key]
+            buf += key.encode("utf-8")
+            if nested_literal.hash:
+                buf += nested_literal.hash.encode("utf-8")
+            else:
+                buf += generate_inputs_repr_for_literal(nested_literal)
+
+        b = bytes(buf)
+        return b
+
+    # For all other cases (scalars, etc.), just serialize the literal normally
+    return literal.SerializeToString(deterministic=True)
+
+
+def generate_inputs_hash_for_named_literals(inputs: list[common_pb2.NamedLiteral]) -> str:
+    """
+    Generate a hash for the inputs using the new literal representation approach that respects
+    hash values already present in literals. This is used to uniquely identify the inputs for a task
+    when some literals may have precomputed hash values.
+
+    :param inputs: List of NamedLiteral inputs to hash.
+    :return: A base64-encoded string representation of the hash.
+    """
+    if not inputs:
+        return ""
+
+    # Build the byte representation by concatenating each literal's representation
+    combined_bytes = b""
+    for named_literal in inputs:
+        # Add the name to ensure order matters
+        name_bytes = named_literal.name.encode("utf-8")
+        literal_bytes = generate_inputs_repr_for_literal(named_literal.value)
+        # Combine name and literal bytes with a separator to avoid collisions
+        combined_bytes += name_bytes + b":" + literal_bytes + b";"
+
+    return hash_data(combined_bytes)
+
+
+def generate_inputs_hash_from_proto(inputs: common_pb2.Inputs) -> str:
+    """
+    Generate a hash for the inputs. This is used to uniquely identify the inputs for a task.
+    :param inputs: The inputs to hash.
+    :return: A hexadecimal string representation of the hash.
+    """
+    if not inputs or not inputs.literals:
+        return ""
+    return generate_inputs_hash_for_named_literals(list(inputs.literals))
+
+
+def generate_interface_hash(task_interface: interface_pb2.TypedInterface) -> str:
+    """
+    Generate a hash for the task interface. This is used to uniquely identify the task interface.
+    :param task_interface: The interface of the task.
+    :return: A hexadecimal string representation of the hash.
+    """
+    if not task_interface:
+        return ""
+    serialized_interface = task_interface.SerializeToString(deterministic=True)
+    return hash_data(serialized_interface)
+
+
+def generate_cache_key_hash(
+    task_name: str,
+    inputs_hash: str,
+    task_interface: interface_pb2.TypedInterface,
+    cache_version: str,
+    ignored_input_vars: List[str],
+    proto_inputs: common_pb2.Inputs,
+) -> str:
+    """
+    Generate a cache key hash based on the inputs hash, task name, task interface, and cache version.
+    This is used to uniquely identify the cache key for a task.
+
+    :param task_name: The name of the task.
+    :param inputs_hash: The hash of the inputs.
+    :param task_interface: The interface of the task.
+    :param cache_version: The version of the cache.
+    :param ignored_input_vars: A list of input variable names to ignore when generating the cache key.
+    :param proto_inputs: The proto inputs for the task, only used if there are ignored inputs.
+    :return: A hexadecimal string representation of the cache key hash.
+    """
+    if ignored_input_vars:
+        filtered = [named_lit for named_lit in proto_inputs.literals if named_lit.name not in ignored_input_vars]
+        final = common_pb2.Inputs(literals=filtered)
+        final_inputs = generate_inputs_hash_from_proto(final)
+    else:
+        final_inputs = inputs_hash
+
+    interface_hash = generate_interface_hash(task_interface)
+
+    data = f"{final_inputs}{task_name}{interface_hash}{cache_version}"
+    return hash_data(data)
+
+
+def generate_sub_action_id_and_output_path(
+    tctx: TaskContext,
+    task_spec_or_name: task_definition_pb2.TaskSpec | str,
+    inputs_hash: str,
+    invoke_seq: int,
+) -> Tuple[ActionID, str]:
     """
     Generate a sub-action ID and output path based on the current task context, task name, and inputs.
+
+    action name = current action name + task name + input hash + group name (if available)
     :param tctx:
-    :param task_name:
-    :param inputs:
+    :param task_spec_or_name: task specification or task name. Task name is only used in case of trace actions.
+    :param inputs_hash: Consistent hash string of the inputs
+    :param invoke_seq: The sequence number of the invocation, used to differentiate between multiple invocations.
     :return:
     """
     current_action_id = tctx.action
     current_output_path = tctx.run_base_dir
-    inputs_hash = pbhash.compute_hash_string(inputs.proto_inputs)
+    if isinstance(task_spec_or_name, task_definition_pb2.TaskSpec):
+        task_spec_or_name.task_template.interface
+        task_hash = hash_data(task_spec_or_name.SerializeToString(deterministic=True))
+    else:
+        task_hash = task_spec_or_name
     sub_action_id = current_action_id.new_sub_action_from(
-        task_name=task_name,
+        task_hash=task_hash,
         input_hash=inputs_hash,
         group=tctx.group_data.name if tctx.group_data else None,
+        task_call_seq=invoke_seq,
     )
     sub_run_output_path = storage.join(current_output_path, sub_action_id.name)
     return sub_action_id, sub_run_output_path

flyte/_internal/runtime/entrypoints.py

@@ -1,16 +1,20 @@
-from typing import List, Optional, Tuple
+import importlib
+import os
+import traceback
+from typing import List, Optional, Tuple, Type
 
 import flyte.errors
 from flyte._code_bundle import download_bundle
 from flyte._context import contextual_run
-from flyte._datastructures import ActionID, Checkpoints, CodeBundle, RawDataPath
 from flyte._internal import Controller
 from flyte._internal.imagebuild.image_builder import ImageCache
 from flyte._logging import log, logger
+from flyte._metrics import Stopwatch
 from flyte._task import TaskTemplate
+from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
 
+from ..._utils import adjust_sys_path
 from .convert import Error, Inputs, Outputs
-from .task_serde import load_task
 from .taskrunner import (
     convert_and_run,
     extract_download_run_upload,
@@ -51,25 +55,95 @@ async def direct_dispatch(
     )
 
 
+def load_class(qualified_name) -> Type:
+    """
+    Load a class from a qualified name. The qualified name should be in the format 'module.ClassName'.
+    :param qualified_name: The qualified name of the class to load.
+    :return: The class object.
+    """
+    module_name, class_name = qualified_name.rsplit(".", 1)  # Split module and class
+    module = importlib.import_module(module_name)  # Import the module
+    return getattr(module, class_name)  # Retrieve the class
+
+
+def load_task(resolver: str, *resolver_args: str) -> TaskTemplate:
+    """
+    Load a task from a resolver. This is a placeholder function.
+
+    :param resolver: The resolver to use to load the task.
+    :param resolver_args: Arguments to pass to the resolver.
+    :return: The loaded task.
+    """
+    resolver_class = load_class(resolver)
+    resolver_instance = resolver_class()
+    try:
+        return resolver_instance.load_task(resolver_args)
+    except ModuleNotFoundError as e:
+        cwd = os.getcwd()
+        files = []
+        try:
+            for root, dirs, filenames in os.walk(cwd):
+                for name in dirs + filenames:
+                    rel_path = os.path.relpath(os.path.join(root, name), cwd)
+                    files.append(rel_path)
+        except Exception as list_err:
+            files = [f"(Failed to list directory: {list_err})"]
+
+        msg = (
+            "\n\nFull traceback:\n" + "".join(traceback.format_exc()) + f"\n[ImportError Diagnostics]\n"
+            f"Module '{e.name}' not found in either the Python virtual environment or the current working directory.\n"
+            f"Current working directory: {cwd}\n"
+            f"Files found under current directory:\n" + "\n".join(f"  - {f}" for f in files)
+        )
+        raise ModuleNotFoundError(msg) from e
+
+
+def load_pkl_task(code_bundle: CodeBundle) -> TaskTemplate:
+    """
+    Loads a task from a pickled code bundle.
+    :param code_bundle: The code bundle to load the task from.
+    :return: The loaded task template.
+    """
+    logger.debug(f"Loading task from pkl: {code_bundle.downloaded_path}")
+    try:
+        import gzip
+
+        import cloudpickle
+
+        with gzip.open(str(code_bundle.downloaded_path), "rb") as f:
+            return cloudpickle.load(f)
+    except Exception as e:
+        logger.exception(f"Failed to load pickled task from {code_bundle.downloaded_path}. Reason: {e!s}")
+        raise
+
+
+async def download_code_bundle(code_bundle: CodeBundle) -> CodeBundle:
+    """
+    Downloads the code bundle if it is not already downloaded.
+    :param code_bundle: The code bundle to download.
+    :return: The code bundle with the downloaded path.
+    """
+    adjust_sys_path([str(code_bundle.destination)])
+    logger.debug(f"Downloading {code_bundle}")
+    sw = Stopwatch("download_code_bundle")
+    sw.start()
+    downloaded_path = await download_bundle(code_bundle)
+    sw.stop()
+    return code_bundle.with_downloaded_path(downloaded_path)
+
+
 async def _download_and_load_task(
     code_bundle: CodeBundle | None, resolver: str | None = None, resolver_args: List[str] | None = None
 ) -> TaskTemplate:
     if code_bundle and (code_bundle.tgz or code_bundle.pkl):
         logger.debug(f"Downloading {code_bundle}")
-        downloaded_path = await download_bundle(code_bundle)
-        code_bundle = code_bundle.with_downloaded_path(downloaded_path)
+        code_bundle = await download_code_bundle(code_bundle)
         if code_bundle.pkl:
-            try:
-                logger.debug(f"Loading task from pkl: {code_bundle.downloaded_path}")
-                import gzip
-
-                import cloudpickle
-
-                with gzip.open(str(code_bundle.downloaded_path), "rb") as f:
-                    return cloudpickle.load(f)
-            except Exception as e:
-                logger.exception(f"Failed to load pickled task from {code_bundle.downloaded_path}. Reason: {e!s}")
-                raise
+            sw = Stopwatch("load_pkl_task")
+            sw.start()
+            result = load_pkl_task(code_bundle)
+            sw.stop()
+            return result
 
         if not resolver or not resolver_args:
             raise flyte.errors.RuntimeSystemError(
@@ -78,11 +152,19 @@ async def _download_and_load_task(
         logger.debug(
             f"Loading task from tgz: {code_bundle.downloaded_path}, resolver: {resolver}, args: {resolver_args}"
         )
-        return load_task(resolver, *resolver_args)
+        sw = Stopwatch("load_task_from_tgz")
+        sw.start()
+        result = load_task(resolver, *resolver_args)
+        sw.stop()
+        return result
     if not resolver or not resolver_args:
         raise flyte.errors.RuntimeSystemError("MalformedCommand", "Resolver and resolver args are required. for task")
     logger.debug(f"No code bundle provided, loading task from resolver: {resolver}, args: {resolver_args}")
-    return load_task(resolver, *resolver_args)
+    sw = Stopwatch("load_task_from_resolver")
+    sw.start()
+    result = load_task(resolver, *resolver_args)
+    sw.stop()
+    return result
 
 
 @log
@@ -99,6 +181,7 @@ async def load_and_run_task(
     code_bundle: CodeBundle | None = None,
     input_path: str | None = None,
     image_cache: ImageCache | None = None,
+    interactive_mode: bool = False,
 ):
     """
     This method is invoked from the runtime/CLI and is used to run a task. This creates the context tree,
@@ -116,7 +199,10 @@ async def load_and_run_task(
     :param code_bundle: The code bundle to use for the task.
     :param input_path: The input path to use for the task.
     :param image_cache: Mappings of Image identifiers to image URIs.
+    :param interactive_mode: Whether to run the task in interactive mode.
     """
+    sw = Stopwatch("load_and_run_task_total")
+    sw.start()
     task = await _download_and_load_task(code_bundle, resolver, resolver_args)
 
     await contextual_run(
@@ -132,4 +218,6 @@ async def load_and_run_task(
         code_bundle=code_bundle,
         input_path=input_path,
         image_cache=image_cache,
+        interactive_mode=interactive_mode,
     )
+    sw.stop()