flyte 0.2.0b1__py3-none-any.whl → 0.2.0b3__py3-none-any.whl

flyte/_internal/controllers/remote/_controller.py

@@ -2,8 +2,9 @@ from __future__ import annotations
 
 import asyncio
 from collections import defaultdict
+from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Awaitable, DefaultDict, Tuple, TypeVar
+from typing import Any, AsyncIterable, Awaitable, DefaultDict, Tuple, TypeVar
 
 import flyte
 import flyte.errors
@@ -11,7 +12,6 @@ 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
@@ -21,16 +21,17 @@ from flyte._internal.runtime.task_serde import translate_task_to_wire
 from flyte._logging import logger
 from flyte._protos.workflow import run_definition_pb2, task_definition_pb2
 from flyte._task import TaskTemplate
+from flyte.models import ActionID, NativeInterface, SerializationContext
 
 R = TypeVar("R")
 
 
-async def upload_inputs_with_retry(inputs: convert.Inputs, inputs_uri: str) -> None:
+async def upload_inputs_with_retry(serialized_inputs: AsyncIterable[bytes] | bytes, inputs_uri: str) -> 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
 
     Raises:
@@ -38,9 +39,9 @@ async def upload_inputs_with_retry(inputs: convert.Inputs, inputs_uri: str) -> N
     """
     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
 
 
@@ -89,6 +90,10 @@ async def load_and_convert_outputs(iface: NativeInterface, realized_outputs_uri:
     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
@@ -111,31 +116,42 @@ class RemoteController(Controller):
         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)
+        )
 
-    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.
+        """
+        current_action_sequencer = self._parent_action_task_call_sequence[unique_action_name(action_id)]
+        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
+        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}")
+            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,13 +162,23 @@ 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)
 
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, task_spec, serialized_inputs, _task_call_seq
+        )
+
+        inputs_uri = io.inputs_path(sub_action_output_path)
+        await upload_inputs_with_retry(serialized_inputs, inputs_uri)
+        # Clear to free memory
+        serialized_inputs = None  # type: ignore
+
         action = Action.from_task(
             sub_action_id=run_definition_pb2.ActionIdentifier(
                 name=sub_action_id.name,
@@ -205,8 +231,9 @@ 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)]:
+            return await self._submit(task_call_seq, _task, *args, **kwargs)
 
     async def finalize_parent_action(self, action_id: ActionID):
         """
@@ -220,16 +247,17 @@ class RemoteController(Controller):
             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,11 +268,19 @@ 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)
+
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, func_name, serialized_inputs, 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)
+        # Clear to free memory
+        serialized_inputs = None  # type: ignore
 
         prev_action = await self.get_action(
             run_definition_pb2.ActionIdentifier(
@@ -310,12 +346,20 @@ class RemoteController(Controller):
         current_action_id = tctx.action
         task_name = _task.spec.task_template.id.name
 
+        invoke_seq_num = self.generate_task_call_sequence(_task, current_action_id)
+
         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)
+        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+
+        sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
+            tctx, task_name, serialized_inputs, 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)
+        # Clear to free memory
+        serialized_inputs = None  # type: ignore
 
         action = Action.from_task(
             sub_action_id=run_definition_pb2.ActionIdentifier(

flyte/_internal/controllers/remote/_core.py

@@ -32,7 +32,7 @@ class Controller:
         max_system_retries: int = 5,
         resource_log_interval_sec: float = 10.0,
         min_backoff_on_err_sec: float = 0.1,
-        thread_wait_timeout_sec: float = 10.0,
+        thread_wait_timeout_sec: float = 0.5,
         enqueue_timeout_sec: float = 5.0,
     ):
         """

flyte/_internal/runtime/convert.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import base64
+import hashlib
 from dataclasses import dataclass
 from typing import Any, Dict, Tuple, Union
 
@@ -7,9 +9,8 @@ from flyteidl.core import execution_pb2, literals_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._protos.workflow import run_definition_pb2, task_definition_pb2
+from flyte.models import ActionID, NativeInterface, TaskContext
 from flyte.types import TypeEngine
 
 
@@ -185,21 +186,46 @@ 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_sub_action_id_and_output_path(
+    tctx: TaskContext,
+    task_spec_or_name: task_definition_pb2.TaskSpec | str,
+    serialized_inputs: str | bytes,
+    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 serialized_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)
+    inputs_hash = hash_data(serialized_inputs)
+    if isinstance(task_spec_or_name, task_definition_pb2.TaskSpec):
+        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

@@ -3,11 +3,11 @@ from typing import List, Optional, Tuple
 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._task import TaskTemplate
+from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
 
 from .convert import Error, Inputs, Outputs
 from .task_serde import load_task

flyte/_internal/runtime/io.py

@@ -21,11 +21,11 @@ _OUTPUTS_FILE_NAME = "outputs.pb"
 _CHECKPOINT_FILE_NAME = "_flytecheckpoints"
 _ERROR_FILE_NAME = "error.pb"
 _REPORT_FILE_NAME = "report.html"
-_PKL_FILE_NAME = "code_bundle.pkl.gz"
+_PKL_EXT = ".pkl.gz"
 
 
-def pkl_path(base_path: str) -> str:
-    return storage.join(base_path, _PKL_FILE_NAME)
+def pkl_path(base_path: str, pkl_name: str) -> str:
+    return storage.join(base_path, f"{pkl_name}{_PKL_EXT}")
 
 
 def inputs_path(base_path: str) -> str:

flyte/_internal/runtime/task_serde.py

@@ -12,11 +12,11 @@ from google.protobuf import duration_pb2, wrappers_pb2
 
 import flyte.errors
 from flyte._cache.cache import VersionParameters, cache_from_request
-from flyte._datastructures import SerializationContext
 from flyte._logging import logger
 from flyte._protos.workflow import task_definition_pb2
 from flyte._secret import SecretRequest, secrets_from_request
 from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
+from flyte.models import CodeBundle, SerializationContext
 
 from ..._retry import RetryStrategy
 from ..._timeout import TimeoutType, timeout_from_request
@@ -208,3 +208,33 @@ def _get_urun_container(
         data_config=task_template.data_loading_config(serialize_context),
         config=task_template.config(serialize_context),
     )
+
+
+def extract_code_bundle(task_spec: task_definition_pb2.TaskSpec) -> Optional[CodeBundle]:
+    """
+    Extract the code bundle from the task spec.
+    :param task_spec: The task spec to extract the code bundle from.
+    :return: The extracted code bundle or None if not present.
+    """
+    container = task_spec.task_template.container
+    if container and container.args:
+        pkl_path = None
+        tgz_path = None
+        dest_path: str = "."
+        version = ""
+        for i, v in enumerate(container.args):
+            if v == "--pkl":
+                # Extract the code bundle path from the argument
+                pkl_path = container.args[i + 1] if i + 1 < len(container.args) else None
+            elif v == "--tgz":
+                # Extract the code bundle path from the argument
+                tgz_path = container.args[i + 1] if i + 1 < len(container.args) else None
+            elif v == "--dest":
+                # Extract the destination path from the argument
+                dest_path = container.args[i + 1] if i + 1 < len(container.args) else "."
+            elif v == "--version":
+                # Extract the version from the argument
+                version = container.args[i + 1] if i + 1 < len(container.args) else ""
+        if pkl_path or tgz_path:
+            return CodeBundle(destination=dest_path, tgz=tgz_path, pkl=pkl_path, computed_version=version)
+    return None

flyte/_internal/runtime/taskrunner.py

@@ -8,11 +8,11 @@ from typing import Any, Dict, List, Optional, Tuple
 
 import flyte.report
 from flyte._context import internal_ctx
-from flyte._datastructures import ActionID, Checkpoints, CodeBundle, RawDataPath, TaskContext
 from flyte._internal.imagebuild.image_builder import ImageCache
 from flyte._logging import log, logger
 from flyte._task import TaskTemplate
 from flyte.errors import CustomError, RuntimeSystemError, RuntimeUnknownError, RuntimeUserError
+from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath, TaskContext
 
 from .. import Controller
 from .convert import (

flyte/_internal/runtime/types_serde.py

@@ -2,7 +2,7 @@ from typing import Dict, Optional, TypeVar
 
 from flyteidl.core import interface_pb2
 
-from flyte._datastructures import NativeInterface
+from flyte.models import NativeInterface
 from flyte.types._type_engine import TypeEngine
 
 T = TypeVar("T")

flyte/_run.py

@@ -4,17 +4,12 @@ import pathlib
 import uuid
 from typing import TYPE_CHECKING, Any, Literal, Optional, Tuple, Union, cast
 
-import flyte
-import flyte.report
-from flyte import S3
+from flyte.errors import InitializationError
 
 from ._api_commons import syncer
 from ._context import contextual_run, internal_ctx
-from ._datastructures import ActionID, Checkpoints, RawDataPath, SerializationContext, TaskContext
 from ._environment import Environment
 from ._initialize import (
-    ABFS,
-    GCS,
     _get_init_config,
     get_client,
     get_common_config,
@@ -22,14 +17,10 @@ from ._initialize import (
     requires_initialization,
     requires_storage,
 )
-from ._internal import create_controller
-from ._internal.runtime.io import _CHECKPOINT_FILE_NAME
-from ._internal.runtime.taskrunner import run_task
 from ._logging import logger
-from ._protos.common import identifier_pb2
 from ._task import P, R, TaskTemplate
 from ._tools import ipython_check
-from .errors import InitializationError
+from .models import ActionID, Checkpoints, CodeBundle, RawDataPath, SerializationContext, TaskContext
 
 if TYPE_CHECKING:
     from flyte.remote import Run
@@ -39,6 +30,22 @@ if TYPE_CHECKING:
 Mode = Literal["local", "remote", "hybrid"]
 
 
+async def _get_code_bundle_for_run(name: str) -> CodeBundle | None:
+    """
+    Get the code bundle for the run with the given name.
+    This is used to get the code bundle for the run when running in hybrid mode.
+    """
+    from flyte._internal.runtime.task_serde import extract_code_bundle
+    from flyte.remote import Run
+
+    run = await Run.get.aio(Run, name=name)
+    if run:
+        run_details = await run.details()
+        spec = run_details.action_details.pb2.resolved_task_spec
+        return extract_code_bundle(spec)
+    return None
+
+
 @syncer.wrap
 class _Runner:
     def __init__(
@@ -81,6 +88,7 @@ class _Runner:
         from ._deploy import build_images, plan_deploy
         from ._internal.runtime.convert import convert_from_native_to_inputs
         from ._internal.runtime.task_serde import translate_task_to_wire
+        from ._protos.common import identifier_pb2
         from ._protos.workflow import run_definition_pb2, run_service_pb2
 
         cfg = get_common_config()
@@ -187,9 +195,14 @@ class _Runner:
         run in the cluster remotely. This is currently only used for testing,
         over the longer term we will productize this.
         """
+        import flyte.report
         from flyte._code_bundle import build_code_bundle, build_pkl_bundle
-        from flyte._datastructures import RawDataPath
         from flyte._deploy import build_images, plan_deploy
+        from flyte.models import RawDataPath
+        from flyte.storage import ABFS, GCS, S3
+
+        from ._internal import create_controller
+        from ._internal.runtime.taskrunner import run_task
 
         cfg = get_common_config()
 
@@ -199,17 +212,23 @@ class _Runner:
         deploy_plan = plan_deploy(cast(Environment, obj.parent_env()))
         image_cache = await build_images(deploy_plan)
 
-        if self._interactive_mode:
-            code_bundle = await build_pkl_bundle(
-                obj, upload_to_controlplane=not self._dry_run, copy_bundle_to=self._copy_bundle_to
-            )
-        else:
-            if self._copy_files != "none":
-                code_bundle = await build_code_bundle(
-                    from_dir=cfg.root_dir, dryrun=self._dry_run, copy_bundle_to=self._copy_bundle_to
+        code_bundle = None
+        if self._name is not None:
+            # Check if remote run service has this run name already and if exists, then extract the code bundle from it.
+            code_bundle = await _get_code_bundle_for_run(name=self._name)
+
+        if not code_bundle:
+            if self._interactive_mode:
+                code_bundle = await build_pkl_bundle(
+                    obj, upload_to_controlplane=not self._dry_run, copy_bundle_to=self._copy_bundle_to
                 )
             else:
-                code_bundle = None
+                if self._copy_files != "none":
+                    code_bundle = await build_code_bundle(
+                        from_dir=cfg.root_dir, dryrun=self._dry_run, copy_bundle_to=self._copy_bundle_to
+                    )
+                else:
+                    code_bundle = None
 
         version = self._version or (
             code_bundle.computed_version if code_bundle and code_bundle.computed_version else None
@@ -217,14 +236,14 @@ class _Runner:
         if not version:
             raise ValueError("Version is required when running a task")
 
-        project = cfg.project or "testproject"
-        domain = cfg.domain or "development"
-        org = cfg.org or "testorg"
+        project = cfg.project
+        domain = cfg.domain
+        org = cfg.org
         action_name = "a0"
         run_name = self._name
         random_id = str(uuid.uuid4())[:6]
 
-        controller = create_controller(ct="remote", endpoint="localhost:8090", insecure=True)
+        controller = create_controller("remote", endpoint="localhost:8090", insecure=True)
         action = ActionID(name=action_name, run_name=run_name, project=project, domain=domain, org=org)
 
         inputs = obj.native_interface.convert_to_kwargs(*args, **kwargs)
@@ -242,7 +261,7 @@ class _Runner:
         output_path = self._run_base_dir
         raw_data_path = f"{output_path}/rd/{random_id}"
         raw_data_path_obj = RawDataPath(path=raw_data_path)
-        checkpoint_path = f"{raw_data_path}/{_CHECKPOINT_FILE_NAME}"
+        checkpoint_path = f"{raw_data_path}/checkpoint"
         prev_checkpoint = f"{raw_data_path}/prev_checkpoint"
         checkpoints = Checkpoints(checkpoint_path, prev_checkpoint)
 
@@ -253,7 +272,7 @@ class _Runner:
                 checkpoints=checkpoints,
                 code_bundle=code_bundle,
                 output_path=output_path,
-                version=version,
+                version=version if version else "na",
                 raw_data_path=raw_data_path_obj,
                 compiled_image_cache=image_cache,
                 run_base_dir=self._run_base_dir,
@@ -276,7 +295,7 @@ class _Runner:
         )
         from flyte._internal.runtime.entrypoints import direct_dispatch
 
-        controller = create_controller(ct="local")
+        controller = create_controller("local")
         inputs = await convert_from_native_to_inputs(obj.native_interface, *args, **kwargs)
         if self._name is None:
             action = ActionID.create_random()

flyte/_task.py

@@ -25,7 +25,6 @@ from flyte.errors import RuntimeSystemError, RuntimeUserError
 
 from ._cache import Cache, CacheRequest
 from ._context import internal_ctx
-from ._datastructures import NativeInterface, SerializationContext
 from ._doc import Documentation
 from ._image import Image
 from ._resources import Resources
@@ -33,6 +32,7 @@ from ._retry import RetryStrategy
 from ._reusable_environment import ReusePolicy
 from ._secret import SecretRequest
 from ._timeout import TimeoutType
+from .models import NativeInterface, SerializationContext
 
 if TYPE_CHECKING:
     from kubernetes.client import V1PodTemplate
@@ -224,7 +224,7 @@ class TaskTemplate(Generic[P, R]):
                 # We will also check if we are not initialized, It is not expected to be not initialized
                 from ._internal.controllers import get_controller
 
-                controller = await get_controller()
+                controller = get_controller()
                 if controller:
                     return await controller.submit(self, *args, **kwargs)
             return await self.execute(*args, **kwargs)

flyte/_task_environment.py

@@ -10,7 +10,6 @@ from typing import TYPE_CHECKING, Awaitable, Callable, Dict, List, Literal, Opti
 import rich.repr
 
 from ._cache import CacheRequest
-from ._datastructures import NativeInterface
 from ._doc import Documentation
 from ._environment import Environment
 from ._image import Image
@@ -19,6 +18,7 @@ from ._retry import RetryStrategy
 from ._reusable_environment import ReusePolicy
 from ._secret import SecretRequest
 from ._task import AsyncFunctionTaskTemplate, TaskTemplate
+from .models import NativeInterface
 
 if TYPE_CHECKING:
     from kubernetes.client import V1PodTemplate

flyte/_trace.py

@@ -4,7 +4,7 @@ import time
 from datetime import timedelta
 from typing import Any, AsyncGenerator, AsyncIterator, Awaitable, Callable, TypeGuard, TypeVar, Union, cast
 
-from flyte._datastructures import NativeInterface
+from flyte.models import NativeInterface
 
 T = TypeVar("T")
 
@@ -14,7 +14,6 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
     A decorator that traces function execution with timing information.
     Works with regular functions, async functions, and async generators/iterators.
     """
-    func_name = func.__name__
 
     @functools.wraps(func)
     def wrapper_sync(*args: Any, **kwargs: Any) -> Any:
@@ -31,9 +30,9 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
             # We will also check if we are not initialized, It is not expected to be not initialized
             from ._internal.controllers import get_controller
 
-            controller = await get_controller()
+            controller = get_controller()
             iface = NativeInterface.from_callable(func)
-            info, ok = await controller.get_action_outputs(iface, func_name, *args, **kwargs)
+            info, ok = await controller.get_action_outputs(iface, func, *args, **kwargs)
             if ok:
                 if info.output:
                     return info.output
@@ -74,9 +73,9 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
             # We will also check if we are not initialized, It is not expected to be not initialized
             from ._internal.controllers import get_controller
 
-            controller = await get_controller()
+            controller = get_controller()
             iface = NativeInterface.from_callable(func)
-            info, ok = await controller.get_action_outputs(iface, func_name, *args, **kwargs)
+            info, ok = await controller.get_action_outputs(iface, func, *args, **kwargs)
             if ok:
                 if info.output:
                     for item in info.output:

flyte/_utils/__init__.py

@@ -4,6 +4,7 @@ Internal utility functions.
 Except for logging, modules in this package should not depend on any other part of the repo.
 """
 
+from .async_cache import AsyncLRUCache
 from .coro_management import run_coros
 from .file_handling import filehash_update, update_hasher_for_source
 from .helpers import get_cwd_editable_install
@@ -11,6 +12,7 @@ from .lazy_module import lazy_module
 from .uv_script_parser import parse_uv_script_file
 
 __all__ = [
+    "AsyncLRUCache",
     "filehash_update",
     "get_cwd_editable_install",
     "lazy_module",