Local vs. Remote Execution

When developing workflows in flyte-sdk, waiting for cluster orchestration, container builds, and remote pod scheduling creates unnecessary iteration latency. You need your decorated functions to behave like plain Python functions during local testing while still serializing into containerized workloads during remote cluster execution.

flyte-sdk achieves this dual behavior through AsyncFunctionTaskTemplate and dynamic execution context resolution via internal_ctx().

Defining Tasks for Local and Remote Use

When you decorate a function with @task or @env.task, flyte-sdk inspects the function and encapsulates it in an AsyncFunctionTaskTemplate instance (or a registered plugin template).

import asyncio
from flyte import task

@task
def add(a: int, b: int) -> int:
    return a + b

@task
async def fetch_data(url: str) -> str:
    # Simulating asynchronous I/O
    await asyncio.sleep(0.01)
    return f"payload from {url}"

# Local execution: tasks execute directly without remote setup
result_sync = add(2, 3)
assert result_sync == 5

async def main():
    result_async = await fetch_data("https://example.com")
    assert result_async == "payload from https://example.com"

asyncio.run(main())

Both synchronous and asynchronous task signatures are supported natively.

How Local Execution Works

When you invoke a decorated task (task_instance(*args, **kwargs)), Python calls TaskTemplate.__call__ in _task.py. The framework determines execution mode by checking the execution context:

ctx = internal_ctx()
if ctx.is_task_context():
    # Submit to controller for orchestrating runs
    ...
else:
    # If not in task context, purely function run, stay out of the way
    return self.forward(*args, **kwargs)

In local development scripts or unit tests, ctx.is_task_context() returns False. The call routes straight to AsyncFunctionTaskTemplate.forward():

def forward(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
    # In local execution, we want to just call the function. Note we're not awaiting anything here.
    # If the function was a coroutine function, the coroutine is returned and the await that the caller has
    # in front of the task invocation will handle the awaiting.
    return self.func(*args, **kwargs)

By delegating immediately to self.func(*args, **kwargs):

• Synchronous tasks return their return value directly with zero engine or overhead latency.
• Asynchronous tasks return a coroutine that your calling code awaits directly.
• Exceptions raised in func are caught by __call__ and wrapped into a RuntimeUserError(type(e).__name__, str(e)) from e unless already a RuntimeSystemError or RuntimeUserError.

Bridging Synchronous Tasks in Async Workflows with .aio()

If you have synchronous legacy tasks that need to run concurrently inside an asynchronous parent task, call task.aio(*args, **kwargs):

import asyncio
from flyte import task

@task
def compute_square(x: int) -> int:
    return x * x

@task
async def process_batch(numbers: list[int]) -> list[int]:
    coros = [compute_square.aio(n) for n in numbers]
    return await asyncio.gather(*coros)

In local execution outside a task context, aio() invokes self.forward(*args, **kwargs) to stay out of the way. When executed inside an active controller context, it submits the task synchronously via controller.submit_sync() and wraps the future with asyncio.wrap_future(fut) so it can be gathered or awaited in the event loop.

How Remote and In-Cluster Execution Works

Remote execution involves two distinct phases: submitting tasks to a controller during workflow execution, and executing the task function inside the worker container.

Local Invocation (No task context):
  TaskTemplate.__call__() ──> forward() ──> func(*args, **kwargs)

Parent Task Context (Active controller):
  TaskTemplate.__call__() ──> controller.submit() / submit_sync()

Worker Container Runtime:
  Worker Entrypoint ──> AsyncFunctionTaskTemplate.execute()
                            ├── await self.pre(*args, **kwargs)
                            ├── ctx.replace_task_context(...)
                            ├── await/call self.func(*args, **kwargs)
                            └── await self.post(v)

1. Task Invocation Under a Controller

When running as part of an orchestrated workflow or inside another task, ctx.is_task_context() returns True. TaskTemplate.__call__ retrieves the active controller via get_controller():

• For synchronous tasks (self._call_as_synchronous is True), it calls controller.submit_sync(self, *args, **kwargs) and waits for the result (fut.result(None)).
• For asynchronous tasks, it returns controller.submit(self, *args, **kwargs).

If no controller is initialized when is_task_context() is True, flyte-sdk raises a RuntimeSystemError("BadContext", "Controller is not initialized.").

2. Container Execution via execute()

When a worker container starts on a remote cluster node to run the task, the runtime loader invokes AsyncFunctionTaskTemplate.execute():

async def execute(self, *args: P.args, **kwargs: P.kwargs) -> R:
    ctx = internal_ctx()
    assert ctx.data.task_context is not None, "Function should have already returned if not in a task context"
    ctx_data = await self.pre(*args, **kwargs)
    tctx = ctx.data.task_context.replace(data=ctx_data)
    with ctx.replace_task_context(tctx):
        if iscoroutinefunction(self.func):
            v = await self.func(*args, **kwargs)
        else:
            v = self.func(*args, **kwargs)
        await self.post(v)
    return v

The execution flow enforces task lifecycle management:

1. Calls await self.pre(*args, **kwargs) to prepare input data and context hooks.
2. Replaces the task context using ctx.replace_task_context(...).
3. Calls the underlying function self.func (awaiting it if iscoroutinefunction(self.func) is True).
4. Invokes await self.post(v) to process outputs.
5. Returns the final value v.

3. Container Serialization and Resolvers

When serializing tasks for cluster deployment, AsyncFunctionTaskTemplate.container_args(serialize_context) constructs the command-line arguments needed by the worker entrypoint container. It configures paths for --inputs, --outputs-path, --version, --raw-data-path, --checkpoint-path, and --run-name.

If no pickled code bundle is specified, it appends the DefaultTaskResolver to dynamically resolve and import the task function in the remote container:

from flyte._internal.resolvers.default import DefaultTaskResolver

_task_resolver = DefaultTaskResolver()
args = [
    *args,
    *[
        "--resolver",
        _task_resolver.import_path,
        *_task_resolver.loader_args(task=self, root_dir=serialize_context.root_dir),
    ],
]

Extending Task Templates

To customize lifecycle behavior (such as adding custom telemetry, modifying pre/post execution hooks, or handling custom hardware plugins), create a subclass of AsyncFunctionTaskTemplate in flyte.extend and register it with TaskPluginRegistry:

from dataclasses import dataclass
from typing import Any
from flyte.extend import AsyncFunctionTaskTemplate, TaskPluginRegistry

@dataclass
class CustomPluginConfig:
    timeout_seconds: int = 60

@dataclass(kw_only=True)
class CustomTaskTemplate(AsyncFunctionTaskTemplate):
    async def pre(self, *args: Any, **kwargs: Any) -> Any:
        # Custom logic before task execution
        return await super().pre(*args, **kwargs)

    async def post(self, result: Any) -> Any:
        # Custom logic after task execution
        return await super().post(result)

# Register the plugin configuration to use this template
TaskPluginRegistry.register(CustomPluginConfig, CustomTaskTemplate)

When defining a task environment with TaskEnvironment(plugin_config=CustomPluginConfig()), TaskEnvironment.task looks up CustomTaskTemplate in TaskPluginRegistry and instantiates it, ensuring that both local calls (forward) and remote runs (execute) execute according to your custom template specifications.