flyte 2.0.0b13__py3-none-any.whl → 2.0.0b30__py3-none-any.whl

flyte/_initialize.py

@@ -3,14 +3,14 @@ from __future__ import annotations
 import functools
 import threading
 import typing
-from dataclasses import dataclass, replace
+from dataclasses import dataclass, field, replace
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable, List, Literal, Optional, TypeVar
 
 from flyte.errors import InitializationError
 from flyte.syncify import syncify
 
-from ._logging import initialize_logger, logger
+from ._logging import LogFormat, initialize_logger, logger
 
 if TYPE_CHECKING:
     from flyte._internal.imagebuild import ImageBuildEngine
@@ -33,6 +33,8 @@ class CommonInit:
     project: str | None = None
     domain: str | None = None
     batch_size: int = 1000
+    source_config_path: Optional[Path] = None  # Only used for documentation
+    sync_local_sys_paths: bool = True
 
 
 @dataclass(init=True, kw_only=True, repr=True, eq=True, frozen=True)
@@ -40,6 +42,7 @@ class _InitConfig(CommonInit):
     client: Optional[ClientSet] = None
     storage: Optional[Storage] = None
     image_builder: "ImageBuildEngine.ImageBuilderType" = "local"
+    images: typing.Dict[str, str] = field(default_factory=dict)
 
     def replace(self, **kwargs) -> _InitConfig:
         return replace(self, **kwargs)
@@ -110,6 +113,10 @@ async def _initialize_client(
     )
 
 
+def _initialize_logger(log_level: int | None = None, log_format: LogFormat | None = None) -> None:
+    initialize_logger(log_level=log_level, log_format=log_format, enable_rich=True)
+
+
 @syncify
 async def init(
     org: str | None = None,
@@ -117,6 +124,7 @@ async def init(
     domain: str | None = None,
     root_dir: Path | None = None,
     log_level: int | None = None,
+    log_format: LogFormat | None = None,
     endpoint: str | None = None,
     headless: bool = False,
     insecure: bool = False,
@@ -134,6 +142,10 @@ async def init(
     storage: Storage | None = None,
     batch_size: int = 1000,
     image_builder: ImageBuildEngine.ImageBuilderType = "local",
+    images: typing.Dict[str, str] | None = None,
+    source_config_path: Optional[Path] = None,
+    sync_local_sys_paths: bool = True,
+    load_plugin_type_transformers: bool = True,
 ) -> None:
     """
     Initialize the Flyte system with the given configuration. This method should be called before any other Flyte
@@ -146,6 +158,7 @@ async def init(
       also use to determine all the code that needs to be copied to the remote location.
       defaults to the editable install directory if the cwd is in a Python editable install, else just the cwd.
     :param log_level: Optional logging level for the logger, default is set using the default initialization policies
+    :param log_format: Optional logging format for the logger, default is "console"
     :param api_key: Optional API key for authentication
     :param endpoint: Optional API endpoint URL
     :param headless: Optional Whether to run in headless mode
@@ -162,24 +175,26 @@ async def init(
     :param ca_cert_file_path: [optional] str Root Cert to be loaded and used to verify admin
     :param http_proxy_url: [optional] HTTP Proxy to be used for OAuth requests
     :param rpc_retries: [optional] int Number of times to retry the platform calls
-    :param audience: oauth2 audience for the token request. This is used to validate the token
     :param insecure: insecure flag for the client
     :param storage: Optional blob store (S3, GCS, Azure) configuration if needed to access (i.e. using Minio)
     :param org: Optional organization override for the client. Should be set by auth instead.
     :param batch_size: Optional batch size for operations that use listings, defaults to 1000, so limit larger than
       batch_size will be split into multiple requests.
     :param image_builder: Optional image builder configuration, if not provided, the default image builder will be used.
-
+    :param images: Optional dict of images that can be used by referencing the image name.
+    :param source_config_path: Optional path to the source configuration file (This is only used for documentation)
+    :param sync_local_sys_paths: Whether to include and synchronize local sys.path entries under the root directory
+      into the remote container (default: True).
+    :param load_plugin_type_transformers: If enabled (default True), load the type transformer plugins registered under
+      the "flyte.plugins.types" entry point group.
     :return: None
     """
-    from flyte._tools import ipython_check
     from flyte._utils import get_cwd_editable_install, org_from_endpoint, sanitize_endpoint
+    from flyte.types import _load_custom_type_transformers
 
-    interactive_mode = ipython_check()
-
-    initialize_logger(enable_rich=interactive_mode)
-    if log_level:
-        initialize_logger(log_level=log_level, enable_rich=interactive_mode)
+    _initialize_logger(log_level=log_level, log_format=log_format)
+    if load_plugin_type_transformers:
+        _load_custom_type_transformers()
 
     global _init_config  # noqa: PLW0603
 
@@ -205,7 +220,15 @@ async def init(
                 http_proxy_url=http_proxy_url,
             )
 
-        root_dir = root_dir or get_cwd_editable_install() or Path.cwd()
+        if not root_dir:
+            editable_root = get_cwd_editable_install()
+            if editable_root:
+                logger.info(f"Using editable install as root directory: {editable_root}")
+                root_dir = editable_root
+            else:
+                logger.info("No editable install found, using current working directory as root directory.")
+                root_dir = Path.cwd()
+
         _init_config = _InitConfig(
             root_dir=root_dir,
             project=project,
@@ -215,14 +238,21 @@ async def init(
             org=org or org_from_endpoint(endpoint),
             batch_size=batch_size,
             image_builder=image_builder,
+            images=images or {},
+            source_config_path=source_config_path,
+            sync_local_sys_paths=sync_local_sys_paths,
         )
 
 
 @syncify
 async def init_from_config(
-    path_or_config: str | Config | None = None,
+    path_or_config: str | Path | Config | None = None,
     root_dir: Path | None = None,
     log_level: int | None = None,
+    log_format: LogFormat = "console",
+    storage: Storage | None = None,
+    images: tuple[str, ...] | None = None,
+    sync_local_sys_paths: bool = True,
 ) -> None:
     """
     Initialize the Flyte system using a configuration file or Config object. This method should be called before any
@@ -235,29 +265,43 @@ async def init_from_config(
         if not available, the current working directory.
     :param log_level: Optional logging level for the framework logger,
         default is set using the default initialization policies
+    :param log_format: Optional logging format for the logger, default is "console"
+    :param storage: Optional blob store (S3, GCS, Azure) configuration if needed to access (i.e. using Minio)
+    :param images: List of image strings in format "imagename=imageuri" or just "imageuri".
+    :param sync_local_sys_paths: Whether to include and synchronize local sys.path entries under the root directory
+     into the remote container (default: True).
     :return: None
     """
+    from rich.highlighter import ReprHighlighter
+
     import flyte.config as config
+    from flyte.cli._common import parse_images
 
     cfg: config.Config
-    if path_or_config is None or isinstance(path_or_config, str):
-        # If a string is passed, treat it as a path to the config file
-        if path_or_config:
-            if not Path(path_or_config).exists():
-                raise InitializationError(
-                    "ConfigFileNotFoundError",
-                    "user",
-                    f"Configuration file '{path_or_config}' does not exist., current working directory is {Path.cwd()}",
-                )
-        if root_dir and path_or_config:
-            cfg = config.auto(str(root_dir / path_or_config))
+    cfg_path: Optional[Path] = None
+    if path_or_config is None:
+        # If no path is provided, use the default config file
+        cfg = config.auto()
+    elif isinstance(path_or_config, (str, Path)):
+        if root_dir:
+            cfg_path = root_dir.expanduser() / path_or_config
         else:
-            cfg = config.auto(path_or_config)
+            cfg_path = Path(path_or_config).expanduser()
+        if not Path(cfg_path).exists():
+            raise InitializationError(
+                "ConfigFileNotFoundError",
+                "user",
+                f"Configuration file '{cfg_path}' does not exist., current working directory is {Path.cwd()}",
+            )
+        cfg = config.auto(cfg_path)
     else:
-        # If a Config object is passed, use it directly
         cfg = path_or_config
 
-    logger.debug(f"Flyte config initialized as {cfg}")
+    logger.info(f"Flyte config initialized as {cfg}", extra={"highlighter": ReprHighlighter()})
+
+    # parse image, this will overwrite the image_refs set in the config file
+    parse_images(cfg, images)
+
     await init.aio(
         org=cfg.task.org,
         project=cfg.task.project,
@@ -273,7 +317,12 @@ async def init_from_config(
         client_credentials_secret=cfg.platform.client_credentials_secret,
         root_dir=root_dir,
         log_level=log_level,
+        log_format=log_format,
         image_builder=cfg.image.builder,
+        images=cfg.image.image_refs,
+        storage=storage,
+        source_config_path=cfg_path,
+        sync_local_sys_paths=sync_local_sys_paths,
     )
 
 
@@ -287,7 +336,7 @@ def _get_init_config() -> Optional[_InitConfig]:
         return _init_config
 
 
-def get_common_config() -> CommonInit:
+def get_init_config() -> _InitConfig:
     """
     Get the current initialization configuration. Thread-safe implementation.
 
@@ -374,30 +423,6 @@ def ensure_client():
         )
 
 
-def requires_client(func: T) -> T:
-    """
-    Decorator that checks if the client has been initialized before executing the function.
-    Raises InitializationError if the client is not initialized.
-
-    :param func: Function to decorate
-    :return: Decorated function that checks for initialization
-    """
-
-    @functools.wraps(func)
-    async def wrapper(*args, **kwargs) -> T:
-        init_config = _get_init_config()
-        if init_config is None or init_config.client is None:
-            raise InitializationError(
-                "ClientNotInitializedError",
-                "user",
-                f"Function '{func.__name__}' requires client to be initialized. "
-                f"Call flyte.init() with a valid endpoint or api-key before using this function.",
-            )
-        return func(*args, **kwargs)
-
-    return typing.cast(T, wrapper)
-
-
 def requires_storage(func: T) -> T:
     """
     Decorator that checks if the storage has been initialized before executing the function.
@@ -469,6 +494,34 @@ def requires_initialization(func: T) -> T:
     return typing.cast(T, wrapper)
 
 
+def require_project_and_domain(func):
+    """
+    Decorator that ensures the current Flyte configuration defines
+    both 'project' and 'domain'. Raises a clear error if not found.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        cfg = get_init_config()
+        if cfg.project is None:
+            raise ValueError(
+                "Project must be provided to initialize the client. "
+                "Please set 'project' in the 'task' section of your config file, "
+                "or pass it directly to flyte.init(project='your-project-name')."
+            )
+
+        if cfg.domain is None:
+            raise ValueError(
+                "Domain must be provided to initialize the client. "
+                "Please set 'domain' in the 'task' section of your config file, "
+                "or pass it directly to flyte.init(domain='your-domain-name')."
+            )
+
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
 async def _init_for_testing(
     project: str | None = None,
     domain: str | None = None,
@@ -498,3 +551,31 @@ def replace_client(client):
 
     with _init_lock:
         _init_config = _init_config.replace(client=client)
+
+
+def current_domain() -> str:
+    """
+    Returns the current domain from Runtime environment (on the cluster) or from the initialized configuration.
+    This is safe to be used during `deploy`, `run` and within `task` code.
+
+    NOTE: This will not work if you deploy a task to a domain and then run it in another domain.
+
+    Raises InitializationError if the configuration is not initialized or domain is not set.
+    :return: The current domain
+    """
+    from ._context import ctx
+
+    tctx = ctx()
+    if tctx is not None:
+        domain = tctx.action.domain
+        if domain is not None:
+            return domain
+
+    cfg = _get_init_config()
+    if cfg is None or cfg.domain is None:
+        raise InitializationError(
+            "DomainNotInitializedError",
+            "user",
+            "Domain has not been initialized. Call flyte.init() with a valid domain before using this function.",
+        )
+    return cfg.domain

flyte/_interface.py

@@ -1,7 +1,13 @@
 from __future__ import annotations
 
 import inspect
-from typing import Dict, Generator, Tuple, Type, TypeVar, Union, cast, get_args, get_type_hints
+import typing
+from enum import Enum
+from typing import Dict, Generator, Literal, Tuple, Type, TypeVar, Union, cast, get_args, get_origin, get_type_hints
+
+from flyte._logging import logger
+
+LITERAL_ENUM = "LiteralEnum"
 
 
 def default_output_name(index: int = 0) -> str:
@@ -45,6 +51,8 @@ def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Di
     Note that Options 1 and 2 are identical, just syntactic sugar. In the NamedTuple case, we'll use the names in the
     definition. In all other cases, we'll automatically generate output names, indexed starting at 0.
     """
+    if isinstance(return_annotation, str):
+        raise TypeError("String return annotations are not supported.")
 
     # Handle Option 6
     # We can think about whether we should add a default output name with type None in the future.
@@ -69,7 +77,15 @@ def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Di
         if len(return_annotation.__args__) == 1:  # type: ignore
             raise TypeError("Tuples should be used to indicate multiple return values, found only one return variable.")
         ra = get_args(return_annotation)
-        return dict(zip(list(output_name_generator(len(ra))), ra))
+        annotations = {}
+        for i, r in enumerate(ra):
+            if r is Ellipsis:
+                raise TypeError("Variable length tuples are not supported as return types.")
+            if get_origin(r) is Literal:
+                annotations[default_output_name(i)] = literal_to_enum(cast(Type, r))
+            else:
+                annotations[default_output_name(i)] = r
+        return annotations
 
     elif isinstance(return_annotation, tuple):
         if len(return_annotation) == 1:
@@ -79,4 +95,25 @@ def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Di
     else:
         # Handle all other single return types
         # Task returns unnamed native tuple
+        if get_origin(return_annotation) is Literal:
+            return {default_output_name(): literal_to_enum(cast(Type, return_annotation))}
         return {default_output_name(): cast(Type, return_annotation)}
+
+
+def literal_to_enum(literal_type: Type) -> Type[Enum | typing.Any]:
+    """Convert a Literal[...] into Union[str, Enum]."""
+
+    if get_origin(literal_type) is not Literal:
+        raise TypeError(f"{literal_type} is not a Literal")
+
+    values = get_args(literal_type)
+    if not all(isinstance(v, str) for v in values):
+        logger.warning(f"Literal type {literal_type} contains non-string values, using Any instead of Enum")
+        return typing.Any
+    # Deduplicate & keep order
+    enum_dict = {str(v).upper(): v for v in values}
+
+    # Dynamically create an Enum
+    literal_enum = Enum(LITERAL_ENUM, enum_dict)  # type: ignore
+
+    return literal_enum  # type: ignore

flyte/_internal/controllers/__init__.py

@@ -5,12 +5,13 @@ from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Protocol, Tu
 from flyte._task import TaskTemplate
 from flyte.models import ActionID, NativeInterface
 
+if TYPE_CHECKING:
+    from flyte.remote._task import TaskDetails
+
 from ._trace import TraceInfo
 
 __all__ = ["Controller", "ControllerType", "TraceInfo", "create_controller", "get_controller"]
 
-from ..._protos.workflow import task_definition_pb2
-
 if TYPE_CHECKING:
     import concurrent.futures
 
@@ -41,9 +42,7 @@ class Controller(Protocol):
         """
         ...
 
-    async def submit_task_ref(
-        self, _task: task_definition_pb2.TaskDetails, max_inline_io_bytes: int, *args, **kwargs
-    ) -> Any:
+    async def submit_task_ref(self, _task: "TaskDetails", *args, **kwargs) -> Any:
         """
         Submit a task reference to the controller asynchronously and wait for the result. This is async and will block
         the current coroutine until the result is available.

flyte/_internal/controllers/_local_controller.py

@@ -2,19 +2,23 @@ import asyncio
 import atexit
 import concurrent.futures
 import os
+import pathlib
 import threading
 from typing import Any, Callable, Tuple, TypeVar
 
 import flyte.errors
+from flyte._cache.cache import VersionParameters, cache_from_request
+from flyte._cache.local_cache import LocalTaskCache
 from flyte._context import internal_ctx
 from flyte._internal.controllers import TraceInfo
 from flyte._internal.runtime import convert
 from flyte._internal.runtime.entrypoints import direct_dispatch
+from flyte._internal.runtime.types_serde import transform_native_to_typed_interface
 from flyte._logging import log, logger
-from flyte._protos.workflow import task_definition_pb2
-from flyte._task import TaskTemplate
+from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
 from flyte._utils.helpers import _selector_policy
 from flyte.models import ActionID, NativeInterface
+from flyte.remote._task import TaskDetails
 
 R = TypeVar("R")
 
@@ -81,31 +85,67 @@ class LocalController:
             raise flyte.errors.RuntimeSystemError("BadContext", "Task context not initialized")
 
         inputs = await convert.convert_from_native_to_inputs(_task.native_interface, *args, **kwargs)
-        serialized_inputs = inputs.proto_inputs.SerializeToString(deterministic=True)
+        inputs_hash = convert.generate_inputs_hash_from_proto(inputs.proto_inputs)
+        task_interface = transform_native_to_typed_interface(_task.interface)
 
         sub_action_id, sub_action_output_path = convert.generate_sub_action_id_and_output_path(
-            tctx, _task.name, serialized_inputs, 0
+            tctx, _task.name, inputs_hash, 0
         )
         sub_action_raw_data_path = tctx.raw_data_path
-
-        out, err = await direct_dispatch(
-            _task,
-            controller=self,
-            action=sub_action_id,
-            raw_data_path=sub_action_raw_data_path,
-            inputs=inputs,
-            version=tctx.version,
-            checkpoints=tctx.checkpoints,
-            code_bundle=tctx.code_bundle,
-            output_path=sub_action_output_path,
-            run_base_dir=tctx.run_base_dir,
+        # Make sure the output path exists
+        pathlib.Path(sub_action_output_path).mkdir(parents=True, exist_ok=True)
+        pathlib.Path(sub_action_raw_data_path.path).mkdir(parents=True, exist_ok=True)
+
+        task_cache = cache_from_request(_task.cache)
+        cache_enabled = task_cache.is_enabled()
+        if isinstance(_task, AsyncFunctionTaskTemplate):
+            version_parameters = VersionParameters(func=_task.func, image=_task.image)
+        else:
+            version_parameters = VersionParameters(func=None, image=_task.image)
+        cache_version = task_cache.get_version(version_parameters)
+        cache_key = convert.generate_cache_key_hash(
+            _task.name,
+            inputs_hash,
+            task_interface,
+            cache_version,
+            list(task_cache.get_ignored_inputs()),
+            inputs.proto_inputs,
         )
-        if err:
-            exc = convert.convert_error_to_native(err)
-            if exc:
-                raise exc
-            else:
-                raise flyte.errors.RuntimeSystemError("BadError", "Unknown error")
+
+        out = None
+        # We only get output from cache if the cache behavior is set to auto
+        if task_cache.behavior == "auto":
+            out = await LocalTaskCache.get(cache_key)
+            if out is not None:
+                logger.info(
+                    f"Cache hit for task '{_task.name}' (version: {cache_version}), getting result from cache..."
+                )
+
+        if out is None:
+            out, err = await direct_dispatch(
+                _task,
+                controller=self,
+                action=sub_action_id,
+                raw_data_path=sub_action_raw_data_path,
+                inputs=inputs,
+                version=cache_version,
+                checkpoints=tctx.checkpoints,
+                code_bundle=tctx.code_bundle,
+                output_path=sub_action_output_path,
+                run_base_dir=tctx.run_base_dir,
+            )
+
+            if err:
+                exc = convert.convert_error_to_native(err)
+                if exc:
+                    raise exc
+                else:
+                    raise flyte.errors.RuntimeSystemError("BadError", "Unknown error")
+
+            # store into cache
+            if cache_enabled and out is not None:
+                await LocalTaskCache.set(cache_key, out)
+
         if _task.native_interface.outputs:
             if out is None:
                 raise flyte.errors.RuntimeSystemError("BadOutput", "Task output not captured.")
@@ -129,7 +169,7 @@ class LocalController:
         pass
 
     async def stop(self):
-        pass
+        await LocalTaskCache.close()
 
     async def watch_for_errors(self):
         pass
@@ -146,16 +186,17 @@ class LocalController:
         tctx = ctx.data.task_context
         if not tctx:
             raise flyte.errors.NotInTaskContextError("BadContext", "Task context not initialized")
+
         converted_inputs = convert.Inputs.empty()
         if _interface.inputs:
             converted_inputs = await convert.convert_from_native_to_inputs(_interface, *args, **kwargs)
             assert converted_inputs
 
-        serialized_inputs = converted_inputs.proto_inputs.SerializeToString(deterministic=True)
+        inputs_hash = convert.generate_inputs_hash_from_proto(converted_inputs.proto_inputs)
         action_id, action_output_path = convert.generate_sub_action_id_and_output_path(
             tctx,
             _func.__name__,
-            serialized_inputs,
+            inputs_hash,
             0,
         )
         assert action_output_path
@@ -192,7 +233,7 @@ class LocalController:
         assert info.start_time
         assert info.end_time
 
-    async def submit_task_ref(
-        self, _task: task_definition_pb2.TaskDetails, max_inline_io_bytes: int, *args, **kwargs
-    ) -> Any:
-        raise flyte.errors.ReferenceTaskError("Reference tasks cannot be executed locally, only remotely.")
+    async def submit_task_ref(self, _task: TaskDetails, max_inline_io_bytes: int, *args, **kwargs) -> Any:
+        raise flyte.errors.ReferenceTaskError(
+            f"Reference tasks cannot be executed locally, only remotely. Found remote task {_task.name}"
+        )

flyte/_internal/controllers/_trace.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
 from typing import Any, Optional
 
-from flyteidl.core import interface_pb2
+from flyteidl2.core import interface_pb2
 
 from flyte.models import ActionID, NativeInterface
 

flyte/_internal/controllers/remote/__init__.py

@@ -54,7 +54,5 @@ def create_remote_controller(
 
     controller = RemoteController(
         client_coro=client_coro,
-        workers=10,
-        max_system_retries=5,
     )
     return controller

flyte/_internal/controllers/remote/_action.py

@@ -1,18 +1,17 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Literal
+from typing import Literal, Optional
 
-from flyteidl.core import execution_pb2, interface_pb2, tasks_pb2
-from google.protobuf import timestamp_pb2
-
-from flyte._protos.common import identifier_pb2
-from flyte._protos.workflow import (
-    queue_service_pb2,
+from flyteidl2.common import identifier_pb2
+from flyteidl2.core import execution_pb2, interface_pb2
+from flyteidl2.task import common_pb2, task_definition_pb2
+from flyteidl2.workflow import (
     run_definition_pb2,
     state_service_pb2,
-    task_definition_pb2,
 )
+from google.protobuf import timestamp_pb2
+
 from flyte.models import GroupData
 
 ActionType = Literal["task", "trace"]
@@ -31,7 +30,7 @@ class Action:
     friendly_name: str | None = None
     group: GroupData | None = None
     task: task_definition_pb2.TaskSpec | None = None
-    trace: queue_service_pb2.TraceAction | None = None
+    trace: run_definition_pb2.TraceAction | None = None
     inputs_uri: str | None = None
     run_output_base: str | None = None
     realized_outputs_uri: str | None = None
@@ -39,6 +38,7 @@ class Action:
     phase: run_definition_pb2.Phase | None = None
     started: bool = False
     retries: int = 0
+    queue: Optional[str] = None  # The queue to which this action was submitted.
     client_err: Exception | None = None  # This error is set when something goes wrong in the controller.
     cache_key: str | None = None  # None means no caching, otherwise it is the version of the cache.
 
@@ -122,6 +122,7 @@ class Action:
         inputs_uri: str,
         run_output_base: str,
         cache_key: str | None = None,
+        queue: Optional[str] = None,
     ) -> Action:
         return cls(
             action_id=sub_action_id,
@@ -132,6 +133,7 @@ class Action:
             inputs_uri=inputs_uri,
             run_output_base=run_output_base,
             cache_key=cache_key,
+            queue=queue,
         )
 
     @classmethod
@@ -183,11 +185,7 @@ class Action:
         et.FromSeconds(int(end_time))
         et.nanos = int((end_time % 1) * 1e9)
 
-        spec = (
-            task_definition_pb2.TaskSpec(task_template=tasks_pb2.TaskTemplate(interface=typed_interface))
-            if typed_interface
-            else None
-        )
+        spec = task_definition_pb2.TraceSpec(interface=typed_interface) if typed_interface else None
 
         return cls(
             action_id=action_id,
@@ -199,12 +197,12 @@ class Action:
             realized_outputs_uri=outputs_uri,
             phase=run_definition_pb2.Phase.PHASE_SUCCEEDED,
             run_output_base=run_output_base,
-            trace=queue_service_pb2.TraceAction(
+            trace=run_definition_pb2.TraceAction(
                 name=friendly_name,
                 phase=run_definition_pb2.Phase.PHASE_SUCCEEDED,
                 start_time=st,
                 end_time=et,
-                outputs=run_definition_pb2.OutputReferences(
+                outputs=common_pb2.OutputReferences(
                     output_uri=outputs_uri,
                     report_uri=report_uri,
                 ),

flyte/_internal/controllers/remote/_client.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 import grpc.aio
+from flyteidl2.workflow import queue_service_pb2_grpc, state_service_pb2_grpc
 
-from flyte._protos.workflow import queue_service_pb2_grpc, state_service_pb2_grpc
 from flyte.remote import create_channel
 
 from ._service_protocol import QueueService, StateService