flyte 2.0.0b3__py3-none-any.whl → 2.0.0b5__py3-none-any.whl

flyte/_code_bundle/bundle.py

@@ -7,6 +7,7 @@ import tempfile
 from pathlib import Path
 from typing import ClassVar, Type
 
+from async_lru import alru_cache
 from flyteidl.core.tasks_pb2 import TaskTemplate
 
 from flyte._logging import log, logger
@@ -109,6 +110,7 @@ async def build_pkl_bundle(
             return CodeBundle(pkl=str(dest), computed_version=str_digest)
 
 
+@alru_cache
 async def build_code_bundle(
     from_dir: Path,
     *ignore: Type[Ignore],

flyte/_deploy.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import asyncio
+import typing
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 
+import grpc.aio
 import rich.repr
 
 import flyte.errors
@@ -90,27 +92,39 @@ async def _deploy_task(
 
     image_uri = task.image.uri if isinstance(task.image, Image) else task.image
 
-    if dryrun:
-        return translate_task_to_wire(task, serialization_context)
-
-    default_inputs = await convert_upload_default_inputs(task.interface)
-    spec = translate_task_to_wire(task, serialization_context, default_inputs=default_inputs)
-
-    msg = f"Deploying task {task.name}, with image {image_uri} version {serialization_context.version}"
-    if spec.task_template.HasField("container") and spec.task_template.container.args:
-        msg += f" from {spec.task_template.container.args[-3]}.{spec.task_template.container.args[-1]}"
-    logger.info(msg)
-    task_id = task_definition_pb2.TaskIdentifier(
-        org=spec.task_template.id.org,
-        project=spec.task_template.id.project,
-        domain=spec.task_template.id.domain,
-        version=spec.task_template.id.version,
-        name=spec.task_template.id.name,
-    )
+    try:
+        if dryrun:
+            return translate_task_to_wire(task, serialization_context)
+
+        default_inputs = await convert_upload_default_inputs(task.interface)
+        spec = translate_task_to_wire(task, serialization_context, default_inputs=default_inputs)
+
+        msg = f"Deploying task {task.name}, with image {image_uri} version {serialization_context.version}"
+        if spec.task_template.HasField("container") and spec.task_template.container.args:
+            msg += f" from {spec.task_template.container.args[-3]}.{spec.task_template.container.args[-1]}"
+        logger.info(msg)
+        task_id = task_definition_pb2.TaskIdentifier(
+            org=spec.task_template.id.org,
+            project=spec.task_template.id.project,
+            domain=spec.task_template.id.domain,
+            version=spec.task_template.id.version,
+            name=spec.task_template.id.name,
+        )
 
-    await get_client().task_service.DeployTask(task_service_pb2.DeployTaskRequest(task_id=task_id, spec=spec))
-    logger.info(f"Deployed task {task.name} with version {task_id.version}")
-    return spec
+        try:
+            await get_client().task_service.DeployTask(task_service_pb2.DeployTaskRequest(task_id=task_id, spec=spec))
+            logger.info(f"Deployed task {task.name} with version {task_id.version}")
+        except grpc.aio.AioRpcError as e:
+            if e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                logger.info(f"Task {task.name} with image {image_uri} already exists, skipping deployment.")
+                return spec
+            raise
+        return spec
+    except Exception as e:
+        logger.error(f"Failed to deploy task {task.name} with image {image_uri}: {e}")
+        raise flyte.errors.DeploymentError(
+            f"Failed to deploy task {task.name} file{task.source_file} with image {image_uri}, Error: {e!s}"
+        ) from e
 
 
 async def _build_image_bg(env_name: str, image: Image) -> Tuple[str, str]:
@@ -151,13 +165,14 @@ async def _build_images(deployment: DeploymentPlan) -> ImageCache:
 
 
 @requires_initialization
-async def apply(deployment: DeploymentPlan, copy_style: CopyFiles, dryrun: bool = False) -> Deployment:
+async def apply(deployment_plan: DeploymentPlan, copy_style: CopyFiles, dryrun: bool = False) -> Deployment:
     from ._code_bundle import build_code_bundle
 
     cfg = get_common_config()
-    image_cache = await _build_images(deployment)
 
-    version = deployment.version
+    image_cache = await _build_images(deployment_plan)
+
+    version = deployment_plan.version
     if copy_style == "none" and not version:
         raise flyte.errors.DeploymentError("Version must be set when copy_style is none")
     else:
@@ -178,41 +193,44 @@ async def apply(deployment: DeploymentPlan, copy_style: CopyFiles, dryrun: bool
     )
 
     tasks = []
-    for env_name, env in deployment.envs.items():
+
+    for env_name, env in deployment_plan.envs.items():
         logger.info(f"Deploying environment {env_name}")
         # TODO Make this pluggable based on the environment type
         if isinstance(env, TaskEnvironment):
             for task in env.tasks.values():
                 tasks.append(_deploy_task(task, dryrun=dryrun, serialization_context=sc))
-    return Deployment(envs=deployment.envs, deployed_tasks=await asyncio.gather(*tasks))
+    return Deployment(envs=deployment_plan.envs, deployed_tasks=await asyncio.gather(*tasks))
 
 
-def _recursive_discover(
-    planned_envs: Dict[str, Environment], envs: Environment | List[Environment]
-) -> Dict[str, Environment]:
+def _recursive_discover(planned_envs: Dict[str, Environment], env: Environment) -> Dict[str, Environment]:
     """
     Recursively deploy the environment and its dependencies, if not already deployed (present in env_tasks) and
     return the updated env_tasks.
     """
-    if isinstance(envs, Environment):
-        envs = [envs]
-    for env in envs:
-        # Skip if the environment is already planned
-        if env.name in planned_envs:
-            continue
-        # Recursively discover dependent environments
-        for dependent_env in env.depends_on:
-            _recursive_discover(planned_envs, dependent_env)
-        # Add the environment to the existing envs
-        planned_envs[env.name] = env
+    # Skip if the environment is already planned
+    if env.name in planned_envs:
+        return planned_envs
+    # Recursively discover dependent environments
+    for dependent_env in env.depends_on:
+        _recursive_discover(planned_envs, dependent_env)
+    # Add the environment to the existing envs
+    planned_envs[env.name] = env
     return planned_envs
 
 
-def plan_deploy(*envs: Environment, version: Optional[str] = None) -> DeploymentPlan:
+def plan_deploy(*envs: Environment, version: Optional[str] = None) -> List[DeploymentPlan]:
     if envs is None:
-        return DeploymentPlan({})
-    planned_envs = _recursive_discover({}, *envs)
-    return DeploymentPlan(planned_envs, version=version)
+        return [DeploymentPlan({})]
+    deployment_plans = []
+    visited_envs: typing.Set[str] = set()
+    for env in envs:
+        if env.name in visited_envs:
+            continue
+        planned_envs = _recursive_discover({}, env)
+        deployment_plans.append(DeploymentPlan(planned_envs, version=version))
+        visited_envs.update(planned_envs.keys())
+    return deployment_plans
 
 
 @syncify
@@ -222,7 +240,7 @@ async def deploy(
     version: str | None = None,
     interactive_mode: bool | None = None,
     copy_style: CopyFiles = "loaded_modules",
-) -> Deployment:
+) -> List[Deployment]:
     """
     Deploy the given environment or list of environments.
     :param envs: Environment or list of environments to deploy.
@@ -238,16 +256,19 @@ async def deploy(
     """
     if interactive_mode:
         raise NotImplementedError("Interactive mode not yet implemented for deployment")
-    deployment = plan_deploy(*envs, version=version)
-    return await apply(deployment, copy_style=copy_style, dryrun=dryrun)
+    deployment_plans = plan_deploy(*envs, version=version)
+    deployments = []
+    for deployment_plan in deployment_plans:
+        deployments.append(apply(deployment_plan, copy_style=copy_style, dryrun=dryrun))
+    return await asyncio.gather(*deployments)
 
 
 @syncify
-async def build_images(*envs: Environment) -> ImageCache:
+async def build_images(envs: Environment) -> ImageCache:
     """
     Build the images for the given environments.
-    :param envs: Environment or list of environments to build images for.
+    :param envs: Environment to build images for.
     :return: ImageCache containing the built images.
     """
-    deployment = plan_deploy(*envs)
-    return await _build_images(deployment)
+    deployment = plan_deploy(envs)
+    return await _build_images(deployment[0])

flyte/_environment.py

@@ -6,14 +6,24 @@ from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
 
 import rich.repr
 
-from flyte._secret import SecretRequest
-
 from ._image import Image
 from ._resources import Resources
+from ._secret import SecretRequest
 
 if TYPE_CHECKING:
     from kubernetes.client import V1PodTemplate
 
+# Global registry to track all Environment instances in load order
+_ENVIRONMENT_REGISTRY: List[Environment] = []
+
+
+def list_loaded_environments() -> List[Environment]:
+    """
+    Return a list of all Environment objects in the order they were loaded.
+    This is useful for deploying environments in the order they were defined.
+    """
+    return _ENVIRONMENT_REGISTRY
+
 
 def is_snake_or_kebab_with_numbers(s: str) -> bool:
     return re.fullmatch(r"^[a-z0-9]+([_-][a-z0-9]+)*$", s) is not None
@@ -44,6 +54,8 @@ class Environment:
     def __post_init__(self):
         if not is_snake_or_kebab_with_numbers(self.name):
             raise ValueError(f"Environment name '{self.name}' must be in snake_case or kebab-case format.")
+        # Automatically register this environment instance in load order
+        _ENVIRONMENT_REGISTRY.append(self)
 
     def add_dependency(self, *env: Environment):
         """

flyte/_image.py

@@ -8,11 +8,14 @@ from abc import abstractmethod
 from dataclasses import asdict, dataclass, field
 from functools import cached_property
 from pathlib import Path
-from typing import Callable, ClassVar, Dict, List, Literal, Optional, Tuple, TypeVar, Union
+from typing import TYPE_CHECKING, Callable, ClassVar, Dict, List, Literal, Optional, Tuple, TypeVar, Union
 
 import rich.repr
 from packaging.version import Version
 
+if TYPE_CHECKING:
+    from flyte import Secret, SecretRequest
+
 # Supported Python versions
 PYTHON_3_10 = (3, 10)
 PYTHON_3_11 = (3, 11)
@@ -71,6 +74,7 @@ class PipOption:
     extra_index_urls: Optional[Tuple[str] | Tuple[str, ...] | List[str]] = None
     pre: bool = False
     extra_args: Optional[str] = None
+    secret_mounts: Optional[Tuple[str | Secret, ...]] = None
 
     def get_pip_install_args(self) -> List[str]:
         pip_install_args = []
@@ -101,6 +105,9 @@ class PipOption:
             hash_input += str(self.pre)
         if self.extra_args:
             hash_input += self.extra_args
+        if self.secret_mounts:
+            for secret_mount in self.secret_mounts:
+                hash_input += str(secret_mount)
 
         hasher.update(hash_input.encode("utf-8"))
 
@@ -110,9 +117,6 @@ class PipOption:
 class PipPackages(PipOption, Layer):
     packages: Optional[Tuple[str, ...]] = None
 
-    # todo: to be implemented
-    # secret_mounts: Optional[List[Tuple[str, str]]] = None
-
     def update_hash(self, hasher: hashlib._Hash):
         """
         Update the hash with the pip packages
@@ -161,6 +165,15 @@ class UVProject(PipOption, Layer):
     pyproject: Path
     uvlock: Path
 
+    def validate(self):
+        if not self.pyproject.exists():
+            raise FileNotFoundError(f"pyproject.toml file {self.pyproject} does not exist")
+        if not self.pyproject.is_file():
+            raise ValueError(f"Pyproject file {self.pyproject} is not a file")
+        if not self.uvlock.exists():
+            raise ValueError(f"UVLock file {self.uvlock} does not exist")
+        super().validate()
+
     def update_hash(self, hasher: hashlib._Hash):
         from ._utils import filehash_update
 
@@ -168,13 +181,40 @@ class UVProject(PipOption, Layer):
         filehash_update(self.uvlock, hasher)
 
 
+@rich.repr.auto
+@dataclass(frozen=True, repr=True)
+class UVScript(PipOption, Layer):
+    script: Path
+
+    def validate(self):
+        if not self.script.exists():
+            raise FileNotFoundError(f"UV script {self.script} does not exist")
+        if not self.script.is_file():
+            raise ValueError(f"UV script {self.script} is not a file")
+        if not self.script.suffix == ".py":
+            raise ValueError(f"UV script {self.script} must have a .py extension")
+        super().validate()
+
+    def update_hash(self, hasher: hashlib._Hash):
+        from ._utils import filehash_update
+
+        super().update_hash(hasher)
+        filehash_update(self.script, hasher)
+
+
 @rich.repr.auto
 @dataclass(frozen=True, repr=True)
 class AptPackages(Layer):
     packages: Tuple[str, ...]
+    secret_mounts: Optional[Tuple[str | Secret, ...]] = None
 
     def update_hash(self, hasher: hashlib._Hash):
-        hasher.update("".join(self.packages).encode("utf-8"))
+        hash_input = "".join(self.packages)
+
+        if self.secret_mounts:
+            for secret_mount in self.secret_mounts:
+                hash_input += str(secret_mount)
+        hasher.update(hash_input.encode("utf-8"))
 
 
 @rich.repr.auto
@@ -394,8 +434,8 @@ class Image:
         )
         labels_and_user = _DockerLines(
             (
-                "LABEL org.opencontainers.image.authors='Union.AI <sales@union.ai>'",
-                "LABEL org.opencontainers.image.source=https://github.com/unionai/unionv2",
+                "LABEL org.opencontainers.image.authors='Union.AI <info@union.ai>'",
+                "LABEL org.opencontainers.image.source=https://github.com/flyteorg/flyte",
                 "RUN useradd --create-home --shell /bin/bash flytekit &&"
                 " chown -R flytekit /root && chown -R flytekit /home",
                 "WORKDIR /root",
@@ -501,6 +541,7 @@ class Image:
         pre: bool = False,
         extra_args: Optional[str] = None,
         platform: Optional[Tuple[Architecture, ...]] = None,
+        secret_mounts: Optional[SecretRequest] = None,
     ) -> Image:
         """
         Use this method to create a new image with the specified uv script.
@@ -534,36 +575,18 @@ class Image:
 
         :return: Image
         """
-        from ._utils import parse_uv_script_file
-
-        if isinstance(script, str):
-            script = Path(script)
-        if not script.exists():
-            raise FileNotFoundError(f"UV script {script} does not exist")
-        if not script.is_file():
-            raise ValueError(f"UV script {script} is not a file")
-        if not script.suffix == ".py":
-            raise ValueError(f"UV script {script} must have a .py extension")
-        header = parse_uv_script_file(script)
+        ll = UVScript(
+            script=Path(script),
+            index_url=index_url,
+            extra_index_urls=_ensure_tuple(extra_index_urls) if extra_index_urls else None,
+            pre=pre,
+            extra_args=extra_args,
+            secret_mounts=_ensure_tuple(secret_mounts) if secret_mounts else None,
+        )
 
-        # todo: arch
         img = cls.from_debian_base(registry=registry, name=name, python_version=python_version, platform=platform)
 
-        # add ca-certificates to the image by default
-        img = img.with_apt_packages("ca-certificates")
-
-        if header.dependencies:
-            return img.with_pip_packages(
-                *header.dependencies,
-                index_url=index_url,
-                extra_index_urls=extra_index_urls,
-                pre=pre,
-                extra_args=extra_args,
-            )
-
-        # todo: override the _identifier_override to be the script name or a hash of the script contents
-        # This is needed because inside the image, the identifier will be computed to be something different.
-        return img
+        return img.clone(addl_layer=ll)
 
     def clone(
         self,
@@ -690,19 +713,26 @@ class Image:
         new_image = self.clone(addl_layer=WorkDir(workdir=workdir))
         return new_image
 
-    def with_requirements(self, file: str | Path) -> Image:
+    def with_requirements(
+        self,
+        file: str | Path,
+        secret_mounts: Optional[SecretRequest] = None,
+    ) -> Image:
         """
         Use this method to create a new image with the specified requirements file layered on top of the current image
         Cannot be used in conjunction with conda
 
         :param file: path to the requirements file, must be a .txt file
+        :param secret_mounts: list of secret to mount for the build process.
         :return:
         """
         if isinstance(file, str):
             file = Path(file)
         if file.suffix != ".txt":
             raise ValueError(f"Requirements file {file} must have a .txt extension")
-        new_image = self.clone(addl_layer=Requirements(file=file))
+        new_image = self.clone(
+            addl_layer=Requirements(file=file, secret_mounts=_ensure_tuple(secret_mounts) if secret_mounts else None)
+        )
         return new_image
 
     def with_pip_packages(
@@ -712,6 +742,7 @@ class Image:
         extra_index_urls: Union[str, List[str], Tuple[str, ...], None] = None,
         pre: bool = False,
         extra_args: Optional[str] = None,
+        secret_mounts: Optional[SecretRequest] = None,
     ) -> Image:
         """
         Use this method to create a new image with the specified pip packages layered on top of the current image
@@ -732,8 +763,8 @@ class Image:
         :param extra_index_urls: extra index urls to use for pip install, default is None
         :param pre: whether to allow pre-release versions, default is False
         :param extra_args: extra arguments to pass to pip install, default is None
-        # :param secret_mounts: todo
         :param extra_args: extra arguments to pass to pip install, default is None
+        :param secret_mounts: list of secret to mount for the build process.
         :return: Image
         """
         new_packages: Optional[Tuple] = packages or None
@@ -745,6 +776,7 @@ class Image:
             extra_index_urls=new_extra_index_urls,
             pre=pre,
             extra_args=extra_args,
+            secret_mounts=_ensure_tuple(secret_mounts) if secret_mounts else None,
         )
         new_image = self.clone(addl_layer=ll)
         return new_image
@@ -790,11 +822,13 @@ class Image:
 
     def with_uv_project(
         self,
-        pyproject_file: Path,
+        pyproject_file: str | Path,
+        uvlock: Path | None = None,
         index_url: Optional[str] = None,
-        extra_index_urls: Union[str, List[str], Tuple[str, ...], None] = None,
+        extra_index_urls: Union[List[str], Tuple[str, ...], None] = None,
         pre: bool = False,
         extra_args: Optional[str] = None,
+        secret_mounts: Optional[SecretRequest] = None,
     ) -> Image:
         """
         Use this method to create a new image with the specified uv.lock file layered on top of the current image
@@ -803,26 +837,44 @@ class Image:
         In the Union builders, using this will change the virtual env to /root/.venv
 
         :param pyproject_file: path to the pyproject.toml file, needs to have a corresponding uv.lock file
-        :return:
+        :param uvlock: path to the uv.lock file, if not specified, will use the default uv.lock file in the same
+        directory as the pyproject.toml file. (pyproject.parent / uv.lock)
+        :param index_url: index url to use for pip install, default is None
+        :param extra_index_urls: extra index urls to use for pip install, default is None
+        :param pre: whether to allow pre-release versions, default is False
+        :param extra_args: extra arguments to pass to pip install, default is None
+        :param secret_mounts: list of secret mounts to use for the build process.
+        :return: Image
         """
-        if not pyproject_file.exists():
-            raise FileNotFoundError(f"UVLock file {pyproject_file} does not exist")
-        if not pyproject_file.is_file():
-            raise ValueError(f"UVLock file {pyproject_file} is not a file")
-        lock = pyproject_file.parent / "uv.lock"
-        if not lock.exists():
-            raise ValueError(f"UVLock file {lock} does not exist")
-        new_image = self.clone(addl_layer=UVProject(pyproject=pyproject_file, uvlock=lock))
+        if isinstance(pyproject_file, str):
+            pyproject_file = Path(pyproject_file)
+        new_image = self.clone(
+            addl_layer=UVProject(
+                pyproject=pyproject_file,
+                uvlock=uvlock or (pyproject_file.parent / "uv.lock"),
+                index_url=index_url,
+                extra_index_urls=extra_index_urls,
+                pre=pre,
+                extra_args=extra_args,
+                secret_mounts=_ensure_tuple(secret_mounts) if secret_mounts else None,
+            )
+        )
         return new_image
 
-    def with_apt_packages(self, *packages: str) -> Image:
+    def with_apt_packages(self, *packages: str, secret_mounts: Optional[SecretRequest] = None) -> Image:
         """
         Use this method to create a new image with the specified apt packages layered on top of the current image
 
         :param packages: list of apt packages to install
+        :param secret_mounts: list of secret mounts to use for the build process.
         :return: Image
         """
-        new_image = self.clone(addl_layer=AptPackages(packages=packages))
+        new_image = self.clone(
+            addl_layer=AptPackages(
+                packages=packages,
+                secret_mounts=_ensure_tuple(secret_mounts) if secret_mounts else None,
+            )
+        )
         return new_image
 
     def with_commands(self, commands: List[str]) -> Image:

flyte/_internal/controllers/remote/_controller.py

@@ -153,7 +153,7 @@ class RemoteController(Controller):
             name = task_obj.__name__
         elif hasattr(task_obj, "name"):
             name = task_obj.name
-        logger.warning(f"For action {uniq}, task {name} call sequence is {new_seq}")
+        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:
@@ -196,7 +196,7 @@ class RemoteController(Controller):
         sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
             tctx, task_spec, inputs_hash, _task_call_seq
         )
-        logger.warning(f"Sub action {sub_action_id} output path {sub_action_output_path}")
+        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)
@@ -421,17 +421,12 @@ class RemoteController(Controller):
 
         current_action_id = tctx.action
         sub_run_output_path = storage.join(tctx.run_base_dir, info.action.name)
-        print(f"Sub run output path for {info.name} is {sub_run_output_path}", flush=True)
 
         if info.interface.has_outputs():
             outputs_file_path: str = ""
             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)
-                print(
-                    f"Uploading outputs for {info.name} Outputs file path: {outputs_file_path}",
-                    flush=True,
-                )
                 await io.upload_outputs(outputs, sub_run_output_path, max_bytes=MAX_TRACE_BYTES)
             elif info.error:
                 err = convert.convert_from_native_to_error(info.error)
@@ -461,19 +456,22 @@ class RemoteController(Controller):
                 end_time=info.end_time,
                 typed_interface=typed_interface if typed_interface else None,
             )
-            try:
-                logger.info(
-                    f"Submitting Trace action Run:[{trace_action.run_name}, 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, _task: task_definition_pb2.TaskDetails, max_inline_io_bytes: int, *args, **kwargs
+            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: task_definition_pb2.TaskDetails, max_inline_io_bytes: int, *args, **kwargs
     ) -> Any:
         ctx = internal_ctx()
         tctx = ctx.data.task_context
@@ -482,8 +480,6 @@ 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, default_inputs=_task.spec.default_inputs
         )
@@ -560,3 +556,15 @@ class RemoteController(Controller):
                 )
             return await load_and_convert_outputs(native_interface, n.realized_outputs_uri, max_inline_io_bytes)
         return None
+
+    async def submit_task_ref(
+        self, _task: task_definition_pb2.TaskDetails, max_inline_io_bytes: int, *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, max_inline_io_bytes, *args, **kwargs)

flyte/_internal/controllers/remote/_core.py

@@ -86,10 +86,7 @@ class Controller:
 
     async def get_action(self, action_id: identifier_pb2.ActionIdentifier, parent_action_name: str) -> Optional[Action]:
         """Get the action from the informer"""
-        informer = await self._informers.get(run_name=action_id.run.name, parent_action_name=parent_action_name)
-        if informer:
-            return await informer.get(action_id.name)
-        return None
+        return await self._run_coroutine_in_controller_thread(self._bg_get_action(action_id, parent_action_name))
 
     @log
     async def cancel_action(self, action: Action):