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

flyte/_run.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import pathlib
+import sys
 import uuid
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
@@ -12,13 +13,13 @@ from flyte._environment import Environment
 from flyte._initialize import (
     _get_init_config,
     get_client,
-    get_common_config,
+    get_init_config,
     get_storage,
     requires_initialization,
     requires_storage,
 )
-from flyte._logging import logger
-from flyte._task import P, R, TaskTemplate
+from flyte._logging import LogFormat, logger
+from flyte._task import F, P, R, TaskTemplate
 from flyte.models import (
     ActionID,
     Checkpoints,
@@ -29,6 +30,8 @@ from flyte.models import (
 )
 from flyte.syncify import syncify
 
+from ._constants import FLYTE_SYS_PATH
+
 if TYPE_CHECKING:
     from flyte.remote import Run
     from flyte.remote._task import LazyEntity
@@ -90,9 +93,12 @@ class _Runner:
         env_vars: Dict[str, str] | None = None,
         labels: Dict[str, str] | None = None,
         annotations: Dict[str, str] | None = None,
-        interruptible: bool = False,
+        interruptible: bool | None = None,
         log_level: int | None = None,
+        log_format: LogFormat = "console",
         disable_run_cache: bool = False,
+        queue: Optional[str] = None,
+        custom_context: Dict[str, str] | None = None,
     ):
         from flyte._tools import ipython_check
 
@@ -111,8 +117,8 @@ class _Runner:
         self._copy_bundle_to = copy_bundle_to
         self._interactive_mode = interactive_mode if interactive_mode else ipython_check()
         self._raw_data_path = raw_data_path
-        self._metadata_path = metadata_path or "/tmp"
-        self._run_base_dir = run_base_dir or "/tmp/base"
+        self._metadata_path = metadata_path
+        self._run_base_dir = run_base_dir
         self._overwrite_cache = overwrite_cache
         self._project = project
         self._domain = domain
@@ -121,12 +127,18 @@ class _Runner:
         self._annotations = annotations
         self._interruptible = interruptible
         self._log_level = log_level
+        self._log_format = log_format
         self._disable_run_cache = disable_run_cache
+        self._queue = queue
+        self._custom_context = custom_context or {}
 
     @requires_initialization
-    async def _run_remote(self, obj: TaskTemplate[P, R] | LazyEntity, *args: P.args, **kwargs: P.kwargs) -> Run:
+    async def _run_remote(self, obj: TaskTemplate[P, R, F] | LazyEntity, *args: P.args, **kwargs: P.kwargs) -> Run:
         import grpc
-        from flyteidl.core import literals_pb2
+        from flyteidl2.common import identifier_pb2
+        from flyteidl2.core import literals_pb2, security_pb2
+        from flyteidl2.task import run_pb2
+        from flyteidl2.workflow import run_definition_pb2, run_service_pb2
         from google.protobuf import wrappers_pb2
 
         from flyte.remote import Run
@@ -136,20 +148,21 @@ class _Runner:
         from ._deploy import build_images
         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()
+        cfg = get_init_config()
         project = self._project or cfg.project
         domain = self._domain or cfg.domain
 
         if isinstance(obj, LazyEntity):
             task = await obj.fetch.aio()
             task_spec = task.pb2.spec
-            inputs = await convert_from_native_to_inputs(task.interface, *args, **kwargs)
+            inputs = await convert_from_native_to_inputs(
+                task.interface, *args, custom_context=self._custom_context, **kwargs
+            )
             version = task.pb2.task_id.version
             code_bundle = None
         else:
+            task = cast(TaskTemplate[P, R, F], obj)
             if obj.parent_env is None:
                 raise ValueError("Task is not attached to an environment. Please attach the task to an environment")
 
@@ -161,7 +174,10 @@ class _Runner:
                 code_bundle = cached_value.code_bundle
                 image_cache = cached_value.image_cache
             else:
-                image_cache = await build_images.aio(cast(Environment, obj.parent_env()))
+                if not self._dry_run:
+                    image_cache = await build_images.aio(cast(Environment, obj.parent_env()))
+                else:
+                    image_cache = None
 
                 if self._interactive_mode:
                     code_bundle = await build_pkl_bundle(
@@ -196,13 +212,23 @@ class _Runner:
                 root_dir=cfg.root_dir,
             )
             task_spec = translate_task_to_wire(obj, s_ctx)
-            inputs = await convert_from_native_to_inputs(obj.native_interface, *args, **kwargs)
+            inputs = await convert_from_native_to_inputs(
+                obj.native_interface, *args, custom_context=self._custom_context, **kwargs
+            )
 
         env = self._env_vars or {}
-        if self._log_level:
-            env["LOG_LEVEL"] = str(self._log_level)
-        else:
-            env["LOG_LEVEL"] = str(logger.getEffectiveLevel())
+        if env.get("LOG_LEVEL") is None:
+            if self._log_level:
+                env["LOG_LEVEL"] = str(self._log_level)
+            else:
+                env["LOG_LEVEL"] = str(logger.getEffectiveLevel())
+        env["LOG_FORMAT"] = self._log_format
+
+        # These paths will be appended to sys.path at runtime.
+        if cfg.sync_local_sys_paths:
+            env[FLYTE_SYS_PATH] = ":".join(
+                f"./{pathlib.Path(p).relative_to(cfg.root_dir)}" for p in sys.path if p.startswith(str(cfg.root_dir))
+            )
 
         if not self._dry_run:
             if get_client() is None:
@@ -245,9 +271,17 @@ class _Runner:
                     raise ValueError(f"Environment variable {k} must be a string, got {type(v)}")
                 kv_pairs.append(literals_pb2.KeyValuePair(key=k, value=v))
 
-            env_kv = run_definition_pb2.Envs(values=kv_pairs)
-            annotations = run_definition_pb2.Annotations(values=self._annotations)
-            labels = run_definition_pb2.Labels(values=self._labels)
+            env_kv = run_pb2.Envs(values=kv_pairs)
+            annotations = run_pb2.Annotations(values=self._annotations)
+            labels = run_pb2.Labels(values=self._labels)
+            raw_data_storage = (
+                run_pb2.RawDataStorage(raw_data_prefix=self._raw_data_path) if self._raw_data_path else None
+            )
+            security_context = (
+                security_pb2.SecurityContext(run_as=security_pb2.Identity(k8s_service_account=self._service_account))
+                if self._service_account
+                else None
+            )
 
             try:
                 resp = await get_client().run_service.CreateRun(
@@ -256,12 +290,17 @@ class _Runner:
                         project_id=project_id,
                         task_spec=task_spec,
                         inputs=inputs.proto_inputs,
-                        run_spec=run_definition_pb2.RunSpec(
+                        run_spec=run_pb2.RunSpec(
                             overwrite_cache=self._overwrite_cache,
-                            interruptible=wrappers_pb2.BoolValue(value=self._interruptible),
+                            interruptible=wrappers_pb2.BoolValue(value=self._interruptible)
+                            if self._interruptible is not None
+                            else None,
                             annotations=annotations,
                             labels=labels,
                             envs=env_kv,
+                            cluster=self._queue or task.queue,
+                            raw_data_storage=raw_data_storage,
+                            security_context=security_context,
                         ),
                     ),
                 )
@@ -306,7 +345,7 @@ class _Runner:
 
     @requires_storage
     @requires_initialization
-    async def _run_hybrid(self, obj: TaskTemplate[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
+    async def _run_hybrid(self, obj: TaskTemplate[P, R, F], *args: P.args, **kwargs: P.kwargs) -> R:
         """
         Run a task in hybrid mode. This means that the parent action will be run locally, but the child actions will be
         run in the cluster remotely. This is currently only used for testing,
@@ -321,7 +360,7 @@ class _Runner:
         from ._internal import create_controller
         from ._internal.runtime.taskrunner import run_task
 
-        cfg = get_common_config()
+        cfg = get_init_config()
 
         if obj.parent_env is None:
             raise ValueError("Task is not attached to an environment. Please attach the task to an environment.")
@@ -380,6 +419,7 @@ class _Runner:
                 " flyte.with_runcontext(run_base_dir='s3://bucket/metadata/outputs')",
             )
         output_path = self._run_base_dir
+        run_base_dir = 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"
@@ -396,8 +436,9 @@ class _Runner:
                 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,
+                run_base_dir=run_base_dir,
                 report=flyte.report.Report(name=action.name),
+                custom_context=self._custom_context,
             )
             async with ctx.replace_task_context(tctx):
                 return await run_task(tctx=tctx, controller=controller, task=obj, inputs=inputs)
@@ -407,10 +448,11 @@ class _Runner:
             raise err
         return outputs
 
-    async def _run_local(self, obj: TaskTemplate[P, R], *args: P.args, **kwargs: P.kwargs) -> Run:
+    async def _run_local(self, obj: TaskTemplate[P, R, F], *args: P.args, **kwargs: P.kwargs) -> Run:
+        from flyteidl2.common import identifier_pb2
+
         from flyte._internal.controllers import create_controller
         from flyte._internal.controllers._local_controller import LocalController
-        from flyte._protos.common import identifier_pb2
         from flyte.remote import Run
         from flyte.report import Report
 
@@ -421,6 +463,18 @@ class _Runner:
         else:
             action = ActionID(name=self._name)
 
+        metadata_path = self._metadata_path
+        if metadata_path is None:
+            metadata_path = pathlib.Path("/") / "tmp" / "flyte" / "metadata" / action.name
+        else:
+            metadata_path = pathlib.Path(metadata_path) / action.name
+        output_path = metadata_path / "a0"
+        if self._raw_data_path is None:
+            path = pathlib.Path("/") / "tmp" / "flyte" / "raw_data" / action.name
+            raw_data_path = RawDataPath(path=str(path))
+        else:
+            raw_data_path = RawDataPath(path=self._raw_data_path)
+
         ctx = internal_ctx()
         tctx = TaskContext(
             action=action,
@@ -429,14 +483,16 @@ class _Runner:
                 checkpoint_path=internal_ctx().raw_data.path,
             ),
             code_bundle=None,
-            output_path=self._metadata_path,
-            run_base_dir=self._metadata_path,
+            output_path=str(output_path),
+            run_base_dir=str(metadata_path),
             version="na",
-            raw_data_path=internal_ctx().raw_data,
+            raw_data_path=raw_data_path,
             compiled_image_cache=None,
             report=Report(name=action.name),
             mode="local",
+            custom_context=self._custom_context,
         )
+
         with ctx.replace_task_context(tctx):
             # make the local version always runs on a different thread, returns a wrapped future.
             if obj._call_as_synchronous:
@@ -448,7 +504,7 @@ class _Runner:
 
         class _LocalRun(Run):
             def __init__(self, outputs: Tuple[Any, ...] | Any):
-                from flyte._protos.workflow import run_definition_pb2
+                from flyteidl2.workflow import run_definition_pb2
 
                 self._outputs = outputs
                 super().__init__(
@@ -464,7 +520,7 @@ class _Runner:
 
             @property
             def url(self) -> str:
-                return "local-run"
+                return str(metadata_path)
 
             def wait(
                 self,
@@ -481,7 +537,7 @@ class _Runner:
     @syncify
     async def run(
         self,
-        task: TaskTemplate[P, Union[R, Run]] | LazyEntity,
+        task: TaskTemplate[P, Union[R, Run], F] | LazyEntity,
         *args: P.args,
         **kwargs: P.kwargs,
     ) -> Union[R, Run]:
@@ -545,9 +601,12 @@ def with_runcontext(
     env_vars: Dict[str, str] | None = None,
     labels: Dict[str, str] | None = None,
     annotations: Dict[str, str] | None = None,
-    interruptible: bool = False,
+    interruptible: bool | None = None,
     log_level: int | None = None,
+    log_format: LogFormat = "console",
     disable_run_cache: bool = False,
+    queue: Optional[str] = None,
+    custom_context: Dict[str, str] | None = None,
 ) -> _Runner:
     """
     Launch a new run with the given parameters as the context.
@@ -575,8 +634,8 @@ def with_runcontext(
     :param interactive_mode: Optional, can be forced to True or False.
          If not provided, it will be set based on the current environment. For example Jupyter notebooks are considered
          interactive mode, while scripts are not. This is used to determine how the code bundle is created.
-    :param raw_data_path: Use this path to store the raw data for the run. Currently only supported for local runs,
-      and can be used to store raw data in specific locations. TODO coming soon for remote runs as well.
+    :param raw_data_path: Use this path to store the raw data for the run for local and remote, and can be used to
+         store raw data in specific locations.
     :param run_base_dir: Optional The base directory to use for the run. This is used to store the metadata for the run,
      that is passed between tasks.
     :param overwrite_cache: Optional If true, the cache will be overwritten for the run
@@ -585,15 +644,24 @@ def with_runcontext(
     :param env_vars: Optional Environment variables to set for the run
     :param labels: Optional Labels to set for the run
     :param annotations: Optional Annotations to set for the run
-    :param interruptible: Optional If true, the run can be interrupted by the user.
+    :param interruptible: Optional If true, the run can be scheduled on interruptible instances and false implies
+        that all tasks in the run should only be scheduled on non-interruptible instances. If not specified the
+        original setting on all tasks is retained.
     :param log_level: Optional Log level to set for the run. If not provided, it will be set to the default log level
         set using `flyte.init()`
+    :param log_format: Optional Log format to set for the run. If not provided, it will be set to the default log format
     :param disable_run_cache: Optional If true, the run cache will be disabled. This is useful for testing purposes.
+    :param queue: Optional The queue to use for the run. This is used to specify the cluster to use for the run.
+    :param custom_context: Optional global input context to pass to the task. This will be available via
+        get_custom_context() within the task and will automatically propagate to sub-tasks.
+        Acts as base/default values that can be overridden by context managers in the code.
 
     :return: runner
     """
     if mode == "hybrid" and not name and not run_base_dir:
         raise ValueError("Run name and run base dir are required for hybrid mode")
+    if copy_style == "none" and not version:
+        raise ValueError("Version is required when copy_style is 'none'")
     return _Runner(
         force_mode=mode,
         name=name,
@@ -613,12 +681,15 @@ def with_runcontext(
         project=project,
         domain=domain,
         log_level=log_level,
+        log_format=log_format,
         disable_run_cache=disable_run_cache,
+        queue=queue,
+        custom_context=custom_context,
     )
 
 
 @syncify
-async def run(task: TaskTemplate[P, R], *args: P.args, **kwargs: P.kwargs) -> Union[R, Run]:
+async def run(task: TaskTemplate[P, R, F], *args: P.args, **kwargs: P.kwargs) -> Union[R, Run]:
     """
     Run a task with the given parameters
     :param task: task to run

flyte/_task.py

@@ -15,10 +15,12 @@ from typing import (
     Literal,
     Optional,
     ParamSpec,
+    Tuple,
     TypeAlias,
     TypeVar,
     Union,
     cast,
+    overload,
 )
 
 from flyte._pod import PodTemplate
@@ -33,10 +35,11 @@ from ._retry import RetryStrategy
 from ._reusable_environment import ReusePolicy
 from ._secret import SecretRequest
 from ._timeout import TimeoutType
+from ._trigger import Trigger
 from .models import MAX_INLINE_IO_BYTES, NativeInterface, SerializationContext
 
 if TYPE_CHECKING:
-    from flyteidl.core.tasks_pb2 import DataLoadingConfig
+    from flyteidl2.core.tasks_pb2 import DataLoadingConfig
 
     from ._task_environment import TaskEnvironment
 
@@ -45,11 +48,12 @@ R = TypeVar("R")  # return type
 
 AsyncFunctionType: TypeAlias = Callable[P, Coroutine[Any, Any, R]]
 SyncFunctionType: TypeAlias = Callable[P, R]
-FunctionTypes: TypeAlias = Union[AsyncFunctionType, SyncFunctionType]
+FunctionTypes: TypeAlias = AsyncFunctionType | SyncFunctionType
+F = TypeVar("F", bound=FunctionTypes)
 
 
 @dataclass(kw_only=True)
-class TaskTemplate(Generic[P, R]):
+class TaskTemplate(Generic[P, R, F]):
     """
     Task template is a template for a task that can be executed. It defines various parameters for the task, which
     can be defined statically at the time of task definition or dynamically at the time of task invocation using
@@ -69,8 +73,8 @@ class TaskTemplate(Generic[P, R]):
     version with flyte installed
     :param resources: Optional The resources to use for the task
     :param cache: Optional The cache policy for the task, defaults to auto, which will cache the results of the task.
-    :param interruptable: Optional The interruptable policy for the task, defaults to False, which means the task
-     will not be scheduled on interruptable nodes. If set to True, the task will be scheduled on interruptable nodes,
+    :param interruptible: Optional The interruptible policy for the task, defaults to False, which means the task
+     will not be scheduled on interruptible nodes. If set to True, the task will be scheduled on interruptible nodes,
      and the code should handle interruptions and resumptions.
     :param retries: Optional The number of retries for the task, defaults to 0, which means no retries.
     :param reusable: Optional The reusability policy for the task, defaults to None, which means the task environment
@@ -81,17 +85,21 @@ class TaskTemplate(Generic[P, R]):
     :param timeout: Optional The timeout for the task.
     :param max_inline_io_bytes: Maximum allowed size (in bytes) for all inputs and outputs passed directly to the task
         (e.g., primitives, strings, dicts). Does not apply to files, directories, or dataframes.
+    :param pod_template: Optional The pod template to use for the task.
+    :param report: Optional Whether to report the task execution to the Flyte console, defaults to False.
+    :param queue: Optional The queue to use for the task. If not provided, the default queue will be used.
+    :param debuggable: Optional Whether the task supports debugging capabilities, defaults to False.
     """
 
     name: str
     interface: NativeInterface
-    friendly_name: str = ""
+    short_name: str = ""
     task_type: str = "python"
     task_type_version: int = 0
     image: Union[str, Image, Literal["auto"]] = "auto"
     resources: Optional[Resources] = None
-    cache: CacheRequest = "auto"
-    interruptable: bool = False
+    cache: CacheRequest = "disable"
+    interruptible: bool = False
     retries: Union[int, RetryStrategy] = 0
     reusable: Union[ReusePolicy, None] = None
     docs: Optional[Documentation] = None
@@ -100,10 +108,14 @@ class TaskTemplate(Generic[P, R]):
     timeout: Optional[TimeoutType] = None
     pod_template: Optional[Union[str, PodTemplate]] = None
     report: bool = False
+    queue: Optional[str] = None
+    debuggable: bool = False
 
     parent_env: Optional[weakref.ReferenceType[TaskEnvironment]] = None
+    parent_env_name: Optional[str] = None
     ref: bool = field(default=False, init=False, repr=False, compare=False)
     max_inline_io_bytes: int = MAX_INLINE_IO_BYTES
+    triggers: Tuple[Trigger, ...] = field(default_factory=tuple)
 
     # Only used in python 3.10 and 3.11, where we cannot use markcoroutinefunction
     _call_as_synchronous: bool = False
@@ -129,9 +141,9 @@ class TaskTemplate(Generic[P, R]):
         if isinstance(self.retries, int):
             self.retries = RetryStrategy(count=self.retries)
 
-        if self.friendly_name == "":
-            # If friendly_name is not set, use the name of the task
-            self.friendly_name = self.name
+        if self.short_name == "":
+            # If short_name is not set, use the name of the task
+            self.short_name = self.name
 
     def __getstate__(self):
         """
@@ -217,6 +229,14 @@ class TaskTemplate(Generic[P, R]):
     def native_interface(self) -> NativeInterface:
         return self.interface
 
+    @overload
+    async def aio(self: TaskTemplate[P, R, SyncFunctionType], *args: P.args, **kwargs: P.kwargs) -> R: ...
+
+    @overload
+    async def aio(
+        self: TaskTemplate[P, R, AsyncFunctionType], *args: P.args, **kwargs: P.kwargs
+    ) -> Coroutine[Any, Any, R]: ...
+
     async def aio(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
         """
         The aio function allows executing "sync" tasks, in an async context. This helps with migrating v1 defined sync
@@ -240,7 +260,6 @@ class TaskTemplate(Generic[P, R]):
         :param kwargs:
         :return:
         """
-
         ctx = internal_ctx()
         if ctx.is_task_context():
             from ._internal.controllers import get_controller
@@ -258,10 +277,21 @@ class TaskTemplate(Generic[P, R]):
             else:
                 raise RuntimeSystemError("BadContext", "Controller is not initialized.")
         else:
+            from flyte._logging import logger
+
+            logger.warning(f"Task {self.name} running aio outside of a task context.")
             # Local execute, just stay out of the way, but because .aio is used, we want to return an awaitable,
             # even for synchronous tasks. This is to support migration.
             return self.forward(*args, **kwargs)
 
+    @overload
+    def __call__(self: TaskTemplate[P, R, SyncFunctionType], *args: P.args, **kwargs: P.kwargs) -> R: ...
+
+    @overload
+    def __call__(
+        self: TaskTemplate[P, R, AsyncFunctionType], *args: P.args, **kwargs: P.kwargs
+    ) -> Coroutine[Any, Any, R]: ...
+
     def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
         """
         This is the entrypoint for an async function task at runtime. It will be called during an execution.
@@ -314,7 +344,7 @@ class TaskTemplate(Generic[P, R]):
     def override(
         self,
         *,
-        friendly_name: Optional[str] = None,
+        short_name: Optional[str] = None,
         resources: Optional[Resources] = None,
         cache: Optional[CacheRequest] = None,
         retries: Union[int, RetryStrategy] = 0,
@@ -324,11 +354,30 @@ class TaskTemplate(Generic[P, R]):
         secrets: Optional[SecretRequest] = None,
         max_inline_io_bytes: int | None = None,
         pod_template: Optional[Union[str, PodTemplate]] = None,
+        queue: Optional[str] = None,
+        interruptible: Optional[bool] = None,
         **kwargs: Any,
     ) -> TaskTemplate:
         """
         Override various parameters of the task template. This allows for dynamic configuration of the task
         when it is called, such as changing the image, resources, cache policy, etc.
+
+        :param short_name: Optional override for the short name of the task.
+        :param resources: Optional override for the resources to use for the task.
+        :param cache: Optional override for the cache policy for the task.
+        :param retries: Optional override for the number of retries for the task.
+        :param timeout: Optional override for the timeout for the task.
+        :param reusable: Optional override for the reusability policy for the task.
+        :param env_vars: Optional override for the environment variables to set for the task.
+        :param secrets: Optional override for the secrets that will be injected into the task at runtime.
+        :param max_inline_io_bytes: Optional override for the maximum allowed size (in bytes) for all inputs and outputs
+         passed directly to the task.
+        :param pod_template: Optional override for the pod template to use for the task.
+        :param queue: Optional override for the queue to use for the task.
+        :param kwargs: Additional keyword arguments for further overrides. Some fields like name, image, docs,
+         and interface cannot be overridden.
+
+        :return: A new TaskTemplate instance with the overridden parameters.
         """
         cache = cache or self.cache
         retries = retries or self.retries
@@ -363,6 +412,8 @@ class TaskTemplate(Generic[P, R]):
         env_vars = env_vars or self.env_vars
         secrets = secrets or self.secrets
 
+        interruptible = interruptible if interruptible is not None else self.interruptible
+
         for k, v in kwargs.items():
             if k == "name":
                 raise ValueError("Name cannot be overridden")
@@ -375,7 +426,7 @@ class TaskTemplate(Generic[P, R]):
 
         return replace(
             self,
-            friendly_name=friendly_name or self.friendly_name,
+            short_name=short_name or self.short_name,
             resources=resources,
             cache=cache,
             retries=retries,
@@ -385,19 +436,22 @@ class TaskTemplate(Generic[P, R]):
             secrets=secrets,
             max_inline_io_bytes=max_inline_io_bytes,
             pod_template=pod_template,
+            interruptible=interruptible,
+            queue=queue or self.queue,
             **kwargs,
         )
 
 
 @dataclass(kw_only=True)
-class AsyncFunctionTaskTemplate(TaskTemplate[P, R]):
+class AsyncFunctionTaskTemplate(TaskTemplate[P, R, F]):
     """
     A task template that wraps an asynchronous functions. This is automatically created when an asynchronous function
     is decorated with the task decorator.
     """
 
-    func: FunctionTypes
+    func: F
     plugin_config: Optional[Any] = None  # This is used to pass plugin specific configuration
+    debuggable: bool = True
 
     def __post_init__(self):
         super().__post_init__()
@@ -476,6 +530,11 @@ class AsyncFunctionTaskTemplate(TaskTemplate[P, R]):
 
             from flyte._internal.resolvers.default import DefaultTaskResolver
 
+            if not serialize_context.root_dir:
+                raise RuntimeSystemError(
+                    "SerializationError",
+                    "Root dir is required for default task resolver when no code bundle is provided.",
+                )
             _task_resolver = DefaultTaskResolver()
             args = [
                 *args,