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

flyte/_internal/controllers/remote/_controller.py

@@ -1,46 +1,62 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent.futures
+import os
+import threading
 from collections import defaultdict
+from collections.abc import Callable
 from pathlib import Path
 from typing import Any, Awaitable, DefaultDict, Tuple, TypeVar
 
+from flyteidl2.common import identifier_pb2, phase_pb2
+
 import flyte
 import flyte.errors
 import flyte.storage as storage
-import flyte.types as types
 from flyte._code_bundle import build_pkl_bundle
 from flyte._context import internal_ctx
-from flyte._datastructures import ActionID, NativeInterface, SerializationContext
 from flyte._internal.controllers import TraceInfo
 from flyte._internal.controllers.remote._action import Action
 from flyte._internal.controllers.remote._core import Controller
 from flyte._internal.controllers.remote._service_protocol import ClientSet
 from flyte._internal.runtime import convert, io
 from flyte._internal.runtime.task_serde import translate_task_to_wire
+from flyte._internal.runtime.types_serde import transform_native_to_typed_interface
 from flyte._logging import logger
-from flyte._protos.workflow import run_definition_pb2, task_definition_pb2
+from flyte._metrics import Stopwatch
 from flyte._task import TaskTemplate
+from flyte._utils.helpers import _selector_policy
+from flyte.models import MAX_INLINE_IO_BYTES, ActionID, NativeInterface, SerializationContext
+from flyte.remote._task import TaskDetails
 
 R = TypeVar("R")
 
+MAX_TRACE_BYTES = MAX_INLINE_IO_BYTES
+
 
-async def upload_inputs_with_retry(inputs: convert.Inputs, inputs_uri: str) -> None:
+async def upload_inputs_with_retry(serialized_inputs: bytes, inputs_uri: str, max_bytes: int) -> None:
     """
     Upload inputs to the specified URI with error handling.
 
     Args:
-        inputs: The inputs to upload
+        serialized_inputs: The serialized inputs to upload
         inputs_uri: The destination URI
+        max_bytes: Maximum number of bytes to read from the input stream
 
     Raises:
         RuntimeSystemError: If the upload fails
     """
+    if len(serialized_inputs) > max_bytes:
+        raise flyte.errors.InlineIOMaxBytesBreached(
+            f"Inputs exceed max_bytes limit of {max_bytes / 1024 / 1024} MB,"
+            f" actual size: {len(serialized_inputs) / 1024 / 1024} MB"
+        )
     try:
         # TODO Add retry decorator to this
-        await io.upload_inputs(inputs, inputs_uri)
+        await storage.put_stream(serialized_inputs, to_path=inputs_uri)
     except Exception as e:
-        logger.exception("Failed to upload inputs", e)
+        logger.exception("Failed to upload inputs")
         raise flyte.errors.RuntimeSystemError(type(e).__name__, str(e)) from e
 
 
@@ -56,7 +72,7 @@ async def handle_action_failure(action: Action, task_name: str) -> Exception:
         Exception: The converted native exception or RuntimeSystemError
     """
     err = action.err or action.client_err
-    if not err and action.phase == run_definition_pb2.PHASE_FAILED:
+    if not err and action.phase == phase_pb2.ACTION_PHASE_FAILED:
         logger.error(f"Server reported failure for action {action.name}, checking error file.")
         try:
             error_path = io.error_path(f"{action.run_output_base}/{action.action_id.name}/1")
@@ -73,22 +89,27 @@ async def handle_action_failure(action: Action, task_name: str) -> Exception:
     return exc
 
 
-async def load_and_convert_outputs(iface: NativeInterface, realized_outputs_uri: str) -> Any:
+async def load_and_convert_outputs(iface: NativeInterface, realized_outputs_uri: str, max_bytes: int) -> Any:
     """
     Load outputs from the given URI and convert them to native format.
 
     Args:
         iface: The Native interface
         realized_outputs_uri: The URI where outputs are stored
+        max_bytes: Maximum number of bytes to read from the output file
 
     Returns:
         The converted native outputs
     """
     outputs_file_path = io.outputs_path(realized_outputs_uri)
-    outputs = await io.load_outputs(outputs_file_path)
+    outputs = await io.load_outputs(outputs_file_path, max_bytes=max_bytes)
     return await convert.convert_outputs_to_native(iface, outputs)
 
 
+def unique_action_name(action_id: ActionID) -> str:
+    return f"{action_id.name}_{action_id.run_name}"
+
+
 class RemoteController(Controller):
     """
     This a specialized controller that wraps the core controller and performs IO, serialization and deserialization
@@ -97,9 +118,8 @@ class RemoteController(Controller):
     def __init__(
         self,
         client_coro: Awaitable[ClientSet],
-        workers: int,
-        max_system_retries: int,
-        default_parent_concurrency: int = 100,
+        workers: int = 20,
+        max_system_retries: int = 10,
     ):
         """ """
         super().__init__(
@@ -107,35 +127,56 @@ class RemoteController(Controller):
             workers=workers,
             max_system_retries=max_system_retries,
         )
+        default_parent_concurrency = int(os.getenv("_F_P_CNC", "1000"))
         self._default_parent_concurrency = default_parent_concurrency
         self._parent_action_semaphore: DefaultDict[str, asyncio.Semaphore] = defaultdict(
             lambda: asyncio.Semaphore(default_parent_concurrency)
         )
+        self._parent_action_task_call_sequence: DefaultDict[str, DefaultDict[int, int]] = defaultdict(
+            lambda: defaultdict(int)
+        )
+        self._submit_loop: asyncio.AbstractEventLoop | None = None
+        self._submit_thread: threading.Thread | None = None
 
-    async def _submit(self, _task: TaskTemplate, *args, **kwargs) -> Any:
+    def generate_task_call_sequence(self, task_obj: object, action_id: ActionID) -> int:
+        """
+        Generate a task call sequence for the given task object and action ID.
+        This is used to track the number of times a task is called within an action.
+        """
+        uniq = unique_action_name(action_id)
+        current_action_sequencer = self._parent_action_task_call_sequence[uniq]
+        current_task_id = id(task_obj)
+        v = current_action_sequencer[current_task_id]
+        new_seq = v + 1
+        current_action_sequencer[current_task_id] = new_seq
+        name = ""
+        if hasattr(task_obj, "__name__"):
+            name = task_obj.__name__
+        elif hasattr(task_obj, "name"):
+            name = task_obj.name
+        logger.info(f"For action {uniq}, task {name} call sequence is {new_seq}")
+        return new_seq
+
+    async def _submit(self, _task_call_seq: int, _task: TaskTemplate, *args, **kwargs) -> Any:
         ctx = internal_ctx()
         tctx = ctx.data.task_context
         if tctx is None:
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
         current_action_id = tctx.action
 
-        inputs = await convert.convert_from_native_to_inputs(_task.native_interface, *args, **kwargs)
-        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(tctx, _task.name, inputs)
-
         # In the case of a regular code bundle, we will just pass it down as it is to the downstream tasks
         # It is not allowed to change the code bundle (for regular code bundles) in the middle of a run.
         code_bundle = tctx.code_bundle
 
-        if code_bundle and code_bundle.pkl:
-            logger.debug(f"Building new pkl bundle for task {sub_action_id.name}")
+        if tctx.interactive_mode or (code_bundle and code_bundle.pkl):
+            logger.debug(f"Building new pkl bundle for task {_task.name}")
             code_bundle = await build_pkl_bundle(
                 _task,
                 upload_to_controlplane=False,
-                upload_from_dataplane_path=io.pkl_path(sub_action_output_path),
+                upload_from_dataplane_base_path=tctx.run_base_dir,
             )
 
-        inputs_uri = io.inputs_path(sub_action_output_path)
-        await upload_inputs_with_retry(inputs, inputs_uri)
+        inputs = await convert.convert_from_native_to_inputs(_task.native_interface, *args, **kwargs)
 
         root_dir = Path(code_bundle.destination).absolute() if code_bundle else Path.cwd()
         # Don't set output path in sec context because node executor will set it
@@ -146,17 +187,46 @@ class RemoteController(Controller):
             code_bundle=code_bundle,
             version=tctx.version,
             # supplied version.
-            input_path=inputs_uri,
+            # input_path=inputs_uri,
             image_cache=tctx.compiled_image_cache,
             root_dir=root_dir,
         )
 
         task_spec = translate_task_to_wire(_task, new_serialization_context)
+        inputs_hash = convert.generate_inputs_hash_from_proto(inputs.proto_inputs)
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, task_spec, inputs_hash, _task_call_seq
+        )
+        logger.info(f"Sub action {sub_action_id} output path {sub_action_output_path}")
+
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+        inputs_uri = io.inputs_path(sub_action_output_path)
+        await upload_inputs_with_retry(serialized_inputs, inputs_uri, max_bytes=_task.max_inline_io_bytes)
+
+        md = task_spec.task_template.metadata
+        ignored_input_vars = []
+        if len(md.cache_ignore_input_vars) > 0:
+            ignored_input_vars = list(md.cache_ignore_input_vars)
+        cache_key = None
+        if task_spec.task_template.metadata and task_spec.task_template.metadata.discoverable:
+            discovery_version = task_spec.task_template.metadata.discovery_version
+            cache_key = convert.generate_cache_key_hash(
+                _task.name,
+                inputs_hash,
+                task_spec.task_template.interface,
+                discovery_version,
+                ignored_input_vars,
+                inputs.proto_inputs,
+            )
+
+        # Clear to free memory
+        serialized_inputs = None  # type: ignore
+        inputs_hash = None  # type: ignore
 
         action = Action.from_task(
-            sub_action_id=run_definition_pb2.ActionIdentifier(
+            sub_action_id=identifier_pb2.ActionIdentifier(
                 name=sub_action_id.name,
-                run=run_definition_pb2.RunIdentifier(
+                run=identifier_pb2.RunIdentifier(
                     name=current_action_id.run_name,
                     project=current_action_id.project,
                     domain=current_action_id.domain,
@@ -168,6 +238,8 @@ class RemoteController(Controller):
             task_spec=task_spec,
             inputs_uri=inputs_uri,
             run_output_base=tctx.run_base_dir,
+            cache_key=cache_key,
+            queue=_task.queue,
         )
 
         try:
@@ -183,7 +255,22 @@ class RemoteController(Controller):
             await self.cancel_action(action)
             raise
 
-        if n.has_error() or n.phase == run_definition_pb2.PHASE_FAILED:
+        # If the action is aborted, we should abort the controller as well
+        if n.phase == phase_pb2.ACTION_PHASE_ABORTED:
+            logger.warning(f"Action {n.action_id.name} was aborted, aborting current Action{current_action_id.name}")
+            raise flyte.errors.RunAbortedError(
+                f"Action {n.action_id.name} was aborted, aborting current Action {current_action_id.name}"
+            )
+
+        if n.phase == phase_pb2.ACTION_PHASE_TIMED_OUT:
+            logger.warning(
+                f"Action {n.action_id.name} timed out, raising timeout exception Action {current_action_id.name}"
+            )
+            raise flyte.errors.TaskTimeoutError(
+                f"Action {n.action_id.name} timed out, raising exception in current Action {current_action_id.name}"
+            )
+
+        if n.has_error() or n.phase == phase_pb2.ACTION_PHASE_FAILED:
             exc = await handle_action_failure(action, _task.name)
             raise exc
 
@@ -193,7 +280,9 @@ class RemoteController(Controller):
                     "RuntimeError",
                     f"Task {n.action_id.name} did not return an output path, but the task has outputs defined.",
                 )
-            return await load_and_convert_outputs(_task.native_interface, n.realized_outputs_uri)
+            return await load_and_convert_outputs(
+                _task.native_interface, n.realized_outputs_uri, max_bytes=_task.max_inline_io_bytes
+            )
         return None
 
     async def submit(self, _task: TaskTemplate, *args, **kwargs) -> Any:
@@ -205,31 +294,81 @@ class RemoteController(Controller):
         if tctx is None:
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
         current_action_id = tctx.action
-        async with self._parent_action_semaphore[current_action_id.name]:
-            return await self._submit(_task, *args, **kwargs)
+        task_call_seq = self.generate_task_call_sequence(_task, current_action_id)
+        async with self._parent_action_semaphore[unique_action_name(current_action_id)]:
+            sw = Stopwatch(f"controller-submit-{unique_action_name(current_action_id)}")
+            sw.start()
+            result = await self._submit(task_call_seq, _task, *args, **kwargs)
+            sw.stop()
+            return result
+
+    def _sync_thread_loop_runner(self) -> None:
+        """This method runs the event loop and should be invoked in a separate thread."""
+
+        loop = self._submit_loop
+        assert loop is not None
+        try:
+            loop.run_forever()
+        finally:
+            loop.close()
+
+    def submit_sync(self, _task: TaskTemplate, *args, **kwargs) -> concurrent.futures.Future:
+        """
+        This function creates a cached thread and loop for the purpose of calling the submit method synchronously,
+        returning a concurrent Future that can be awaited. There's no need for a lock because this function itself is
+        single threaded and non-async. This pattern here is basically the trivial/degenerate case of the thread pool
+        in the LocalController.
+        Please see additional comments in protocol.
+
+        :param _task:
+        :param args:
+        :param kwargs:
+        :return:
+        """
+        if self._submit_thread is None:
+            # Please see LocalController for the general implementation of this pattern.
+            def exc_handler(loop, context):
+                logger.error(f"Remote controller submit sync loop caught exception in {loop}: {context}")
+
+            with _selector_policy():
+                self._submit_loop = asyncio.new_event_loop()
+                self._submit_loop.set_exception_handler(exc_handler)
+
+            self._submit_thread = threading.Thread(
+                name=f"remote-controller-{os.getpid()}-submitter",
+                daemon=True,
+                target=self._sync_thread_loop_runner,
+            )
+            self._submit_thread.start()
+
+        coro = self.submit(_task, *args, **kwargs)
+        assert self._submit_loop is not None, "Submit loop should always have been initialized by now"
+        fut = asyncio.run_coroutine_threadsafe(coro, self._submit_loop)
+        return fut
 
     async def finalize_parent_action(self, action_id: ActionID):
         """
         This method is invoked when the parent action is finished. It will finalize the run and upload the outputs
         to the control plane.
         """
-        run_id = run_definition_pb2.RunIdentifier(
+        run_id = identifier_pb2.RunIdentifier(
             name=action_id.run_name,
             project=action_id.project,
             domain=action_id.domain,
             org=action_id.org,
         )
         await super()._finalize_parent_action(run_id=run_id, parent_action_name=action_id.name)
-        self._parent_action_semaphore.pop(action_id.name, None)
+        self._parent_action_semaphore.pop(unique_action_name(action_id), None)
+        self._parent_action_task_call_sequence.pop(unique_action_name(action_id), None)
 
     async def get_action_outputs(
-        self, _interface: NativeInterface, _func_name: str, *args, **kwargs
+        self, _interface: NativeInterface, _func: Callable, *args, **kwargs
     ) -> Tuple[TraceInfo, bool]:
         """
         This method returns the outputs of the action, if it is available.
         If not available it raises a NotFoundError.
         :param _interface: NativeInterface
-        :param _func_name: Function name
+        :param _func: Function name
         :param args: Arguments
         :param kwargs: Keyword arguments
         :return:
@@ -240,16 +379,26 @@ class RemoteController(Controller):
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
         current_action_id = tctx.action
 
+        func_name = _func.__name__
+        invoke_seq_num = self.generate_task_call_sequence(_func, current_action_id)
+
         inputs = await convert.convert_from_native_to_inputs(_interface, *args, **kwargs)
-        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(tctx, _func_name, inputs)
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+        inputs_hash = convert.generate_inputs_hash_from_proto(inputs.proto_inputs)
+
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, func_name, inputs_hash, invoke_seq_num
+        )
 
         inputs_uri = io.inputs_path(sub_action_output_path)
-        await upload_inputs_with_retry(inputs, inputs_uri)
+        await upload_inputs_with_retry(serialized_inputs, inputs_uri, max_bytes=MAX_TRACE_BYTES)
+        # Clear to free memory
+        serialized_inputs = None  # type: ignore
 
         prev_action = await self.get_action(
-            run_definition_pb2.ActionIdentifier(
+            identifier_pb2.ActionIdentifier(
                 name=sub_action_id.name,
-                run=run_definition_pb2.RunIdentifier(
+                run=identifier_pb2.RunIdentifier(
                     name=current_action_id.run_name,
                     project=current_action_id.project,
                     domain=current_action_id.domain,
@@ -260,21 +409,26 @@ class RemoteController(Controller):
         )
 
         if prev_action is None:
-            return TraceInfo(sub_action_id, _interface, inputs_uri), False
+            return TraceInfo(func_name, sub_action_id, _interface, inputs_uri), False
 
-        if prev_action.phase == run_definition_pb2.PHASE_FAILED:
+        if prev_action.phase == phase_pb2.ACTION_PHASE_FAILED:
             if prev_action.has_error():
                 exc = convert.convert_error_to_native(prev_action.err)
-                return TraceInfo(sub_action_id, _interface, inputs_uri, error=exc), True
+                return (
+                    TraceInfo(func_name, sub_action_id, _interface, inputs_uri, error=exc),
+                    True,
+                )
             else:
                 logger.warning(f"Action {prev_action.action_id.name} failed, but no error was found, re-running trace!")
         elif prev_action.realized_outputs_uri is not None:
-            outputs_file_path = io.outputs_path(prev_action.realized_outputs_uri)
-            o = await io.load_outputs(outputs_file_path)
+            o = await io.load_outputs(prev_action.realized_outputs_uri, max_bytes=MAX_TRACE_BYTES)
             outputs = await convert.convert_outputs_to_native(_interface, o)
-            return TraceInfo(sub_action_id, _interface, inputs_uri, output=outputs), True
+            return (
+                TraceInfo(func_name, sub_action_id, _interface, inputs_uri, output=outputs),
+                True,
+            )
 
-        return TraceInfo(sub_action_id, _interface, inputs_uri), False
+        return TraceInfo(func_name, sub_action_id, _interface, inputs_uri), False
 
     async def record_trace(self, info: TraceInfo):
         """
@@ -287,40 +441,100 @@ class RemoteController(Controller):
         if tctx is None:
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
 
-        current_output_path = tctx.output_path
-        sub_run_output_path = storage.join(current_output_path, info.action.name)
+        current_action_id = tctx.action
+        sub_run_output_path = storage.join(tctx.run_base_dir, info.action.name)
+        outputs_file_path: str = ""
 
         if info.interface.has_outputs():
-            if info.output:
-                outputs = await convert.convert_from_native_to_outputs(info.output, info.interface)
-                outputs_file_path = io.outputs_path(sub_run_output_path)
-                await io.upload_outputs(outputs, outputs_file_path)
-            elif info.error:
+            if info.error:
                 err = convert.convert_from_native_to_error(info.error)
-                error_path = io.error_path(sub_run_output_path)
-                await io.upload_error(err.err, error_path)
+                await io.upload_error(err.err, sub_run_output_path)
             else:
-                raise flyte.errors.RuntimeSystemError("BadTraceInfo", "Trace info does not have output or error")
+                outputs = await convert.convert_from_native_to_outputs(info.output, info.interface)
+                outputs_file_path = io.outputs_path(sub_run_output_path)
+                await io.upload_outputs(outputs, sub_run_output_path, max_bytes=MAX_TRACE_BYTES)
+
+        typed_interface = transform_native_to_typed_interface(info.interface)
+
+        trace_action = Action.from_trace(
+            parent_action_name=current_action_id.name,
+            action_id=identifier_pb2.ActionIdentifier(
+                name=info.action.name,
+                run=identifier_pb2.RunIdentifier(
+                    name=current_action_id.run_name,
+                    project=current_action_id.project,
+                    domain=current_action_id.domain,
+                    org=current_action_id.org,
+                ),
+            ),
+            inputs_uri=info.inputs_path,
+            outputs_uri=outputs_file_path,
+            friendly_name=info.name,
+            group_data=tctx.group_data,
+            run_output_base=tctx.run_base_dir,
+            start_time=info.start_time,
+            end_time=info.end_time,
+            typed_interface=typed_interface if typed_interface else None,
+        )
 
-    async def submit_task_ref(self, _task: task_definition_pb2.TaskDetails, *args, **kwargs) -> Any:
+        async with self._parent_action_semaphore[unique_action_name(current_action_id)]:
+            try:
+                logger.info(
+                    f"Submitting Trace action Run:[{trace_action.run_name},"
+                    f" Parent:[{trace_action.parent_action_name}],"
+                    f" Trace fn:[{info.name}], action:[{info.action.name}]"
+                )
+                await self.submit_action(trace_action)
+                logger.info(f"Trace Action for [{info.name}] action id: {info.action.name}, completed!")
+            except asyncio.CancelledError:
+                # If the action is cancelled, we need to cancel the action on the server as well
+                raise
+
+    async def _submit_task_ref(self, invoke_seq_num: int, _task: TaskDetails, *args, **kwargs) -> Any:
         ctx = internal_ctx()
         tctx = ctx.data.task_context
         if tctx is None:
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
         current_action_id = tctx.action
-        task_name = _task.spec.task_template.id.name
+        task_name = _task.name
+
+        native_interface = _task.interface
+        pb_interface = _task.pb2.spec.task_template.interface
 
-        native_interface = types.guess_interface(_task.spec.task_template.interface)
         inputs = await convert.convert_from_native_to_inputs(native_interface, *args, **kwargs)
-        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(tctx, task_name, inputs)
+        inputs_hash = convert.generate_inputs_hash_from_proto(inputs.proto_inputs)
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, task_name, inputs_hash, invoke_seq_num
+        )
 
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
         inputs_uri = io.inputs_path(sub_action_output_path)
-        await upload_inputs_with_retry(inputs, inputs_uri)
+        await upload_inputs_with_retry(serialized_inputs, inputs_uri, _task.max_inline_io_bytes)
+        # cache key - task name, task signature, inputs, cache version
+        cache_key = None
+        md = _task.pb2.spec.task_template.metadata
+        ignored_input_vars = []
+        if len(md.cache_ignore_input_vars) > 0:
+            ignored_input_vars = list(md.cache_ignore_input_vars)
+        if md and md.discoverable:
+            discovery_version = md.discovery_version
+            cache_key = convert.generate_cache_key_hash(
+                task_name,
+                inputs_hash,
+                pb_interface,
+                discovery_version,
+                ignored_input_vars,
+                inputs.proto_inputs,
+            )
+
+        # Clear to free memory
+        serialized_inputs = None  # type: ignore
+        inputs_hash = None  # type: ignore
 
         action = Action.from_task(
-            sub_action_id=run_definition_pb2.ActionIdentifier(
+            sub_action_id=identifier_pb2.ActionIdentifier(
                 name=sub_action_id.name,
-                run=run_definition_pb2.RunIdentifier(
+                run=identifier_pb2.RunIdentifier(
                     name=current_action_id.run_name,
                     project=current_action_id.project,
                     domain=current_action_id.domain,
@@ -329,9 +543,11 @@ class RemoteController(Controller):
             ),
             parent_action_name=current_action_id.name,
             group_data=tctx.group_data,
-            task_spec=_task.spec,
+            task_spec=_task.pb2.spec,
             inputs_uri=inputs_uri,
             run_output_base=tctx.run_base_dir,
+            cache_key=cache_key,
+            queue=None,
         )
 
         try:
@@ -347,7 +563,7 @@ class RemoteController(Controller):
             await self.cancel_action(action)
             raise
 
-        if n.has_error() or n.phase == run_definition_pb2.PHASE_FAILED:
+        if n.has_error() or n.phase == phase_pb2.ACTION_PHASE_FAILED:
             exc = await handle_action_failure(action, task_name)
             raise exc
 
@@ -357,5 +573,15 @@ class RemoteController(Controller):
                     "RuntimeError",
                     f"Task {n.action_id.name} did not return an output path, but the task has outputs defined.",
                 )
-            return await load_and_convert_outputs(native_interface, n.realized_outputs_uri)
+            return await load_and_convert_outputs(native_interface, n.realized_outputs_uri, _task.max_inline_io_bytes)
         return None
+
+    async def submit_task_ref(self, _task: TaskDetails, *args, **kwargs) -> Any:
+        ctx = internal_ctx()
+        tctx = ctx.data.task_context
+        if tctx is None:
+            raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
+        current_action_id = tctx.action
+        task_call_seq = self.generate_task_call_sequence(_task, current_action_id)
+        async with self._parent_action_semaphore[unique_action_name(current_action_id)]:
+            return await self._submit_task_ref(task_call_seq, _task, *args, **kwargs)