Skip to main content

Introduction to Tasks

From a Python function to a task​

A function decorated with @env.task is not merely called by flyte.run as a generic Python function. TaskEnvironment.task creates an AsyncFunctionTaskTemplate, keeps the original callable in its func field, and stores the resulting template in the environment's _tasks mapping. That template is the entity Flyte can invoke, submit, serialize, and execute.

Use a TaskEnvironment as the authoring surface:

import flyte

env = flyte.TaskEnvironment("example")

@env.task
async def example_task(x: int, y: str) -> str:
    return f"{x} {y}"

if __name__ == "__main__":
    flyte.run(example_task, 1, y="hello")

This is the example in _run.py. flyte.run accepts a TaskTemplate (including the object produced by @env.task), rather than an arbitrary function. The decorated object supplies the typed task interface and metadata that the runner needs.

What decoration populates​

For a non-plugin environment, TaskEnvironment.task selects AsyncFunctionTaskTemplate and constructs it with the following values from _task_environment.py:

  • name is <environment name>.<function name>. For the example above, it is example.example_task.
  • short_name is the function name unless you pass short_name=....
  • interface comes from NativeInterface.from_callable(func), so the callable's annotated inputs and output are used to build the native task interface.
  • image, resources, env_vars, secrets, reusable, pod_template, and plugin_config come from the environment, while task-level options can supply cache, retries, timeout, docs, report, and max_inline_io_bytes.
  • parent_env is a weak reference to the TaskEnvironment.

For example, environment-level image and resources are attached to every task created by that environment:

import flyte
from flyte import Resources

env = flyte.TaskEnvironment(
    name="my_env",
    image="my_image",
    resources=Resources(cpu="1", memory="1Gi"),
)

@env.task
async def my_task():
    pass

The environment's declared cache value defaults to "disable" in TaskEnvironment, while a task decorator uses cache when supplied and otherwise falls back to self.cache. retries defaults to 0, timeout to 0, report to False, and max_inline_io_bytes to MAX_INLINE_IO_BYTES in the decorator signature. TaskTemplate.__post_init__, which runs before the subclass initialization finishes, normalizes metadata such as images, cache requests, retry values, and the default short name.

You can set task-specific execution metadata without changing the function body:

from datetime import timedelta

from flyte import Resources, TaskEnvironment, Timeout

env = TaskEnvironment("example")

timeout = Timeout(max_runtime=timedelta(minutes=5), max_queued_time=timedelta(minutes=10))

@env.task(
    resources=Resources(cpu=1, memory="1GiB", gpu="T4:1"),
    timeout=timeout,
    retries=2,
    short_name="compute",
)
async def example_task() -> int:
    return 42

The TaskEnvironment.task signature accepts short_name, cache, retries, timeout, docs, pod_template, report, and max_inline_io_bytes. Resources in the example are an inherited TaskTemplate field; they are configured on the environment in the SDK's environment model rather than accepted by this decorator signature.

One template, two invocation paths​

AsyncFunctionTaskTemplate has “Async” in its name because it is the standard function-task template, not because every wrapped function must be asynchronous. Its func type is a union of an async callable and a synchronous callable. In __post_init__, it sets _call_as_synchronous = True when inspect.iscoroutinefunction(self.func) is false.

The inherited TaskTemplate.__call__ chooses the path based on the Flyte context:

no task context  ──> TaskTemplate.__call__ ──> AsyncFunctionTaskTemplate.forward ──> func
                 
inside task context ──> TaskTemplate.__call__ ──> controller.submit / submit_sync

Outside a task context, __call__ delegates to forward. forward directly evaluates self.func(*args, **kwargs) and does not await it. Consequently, an async task call returns the coroutine for its caller to await; a synchronous task call returns its ordinary value. This distinction is important: forward is the local path, not a coroutine wrapper that always awaits.

Inside a task context, __call__ obtains the controller. Async tasks use controller.submit, while templates marked _call_as_synchronous use controller.submit_sync and wait for its result. If the context indicates a task execution but no controller is initialized, the call raises RuntimeSystemError("BadContext", "Controller is not initialized.").

For an async parent that needs to use a synchronous task, use the inherited .aio bridge:

import asyncio
from typing import List

@env.task
def my_legacy_task(x: int) -> int:
    return x

@env.task
async def my_new_parent_task(n: int) -> List[int]:
    collect = []
    for x in range(n):
        collect.append(my_legacy_task.aio(x))
    return asyncio.gather(*collect)

The .aio method submits synchronous tasks with submit_sync, wraps the returned future with asyncio.wrap_future, and awaits it when already in a task context. Outside a task context, it still calls forward; using .aio makes the call awaitable-compatible for migration from synchronous tasks.

Runtime execution: pre, function, post​

The controller/runtime path does not use forward to run the function body. _internal/runtime/taskrunner.py calls await task.execute(**inputs). AsyncFunctionTaskTemplate.execute requires an existing Flyte task context and asserts that ctx.data.task_context is not None. It then performs this sequence:

execute
  ├─ await pre(*args, **kwargs)
  ├─ replace the task context with the returned data
  ├─ await func(...)       # async callable
  │  or func(...)          # synchronous callable
  └─ await post(value)

The implementation invokes pre, replaces the context for the duration of the function call, invokes the original func, calls post(v), and returns v. Calling execute directly as a normal local-call API is therefore not appropriate: without a task context it fails its assertion. Use task(...) for normal invocation, allowing TaskTemplate.__call__ to choose local execution or controller submission.

source_file exposes the wrapped callable's __code__.co_filename, or None when the callable has no usable code object. Deployment error handling uses this property when associating failures with the task's source file.

Serialization into a runnable container command​

When Flyte serializes the template, _internal/runtime/task_serde.py calls its container_args method while constructing the task specification. AsyncFunctionTaskTemplate.container_args starts the standard Python runner command with a0 and supplies the serialized input/output locations and runtime identifiers:

[a0,
 --inputs, <input_path>,
 --outputs-path, <output_path>,
 --version, <version>,
 --raw-data-path, {{.rawOutputDataPrefix}},
 --checkpoint-path, {{.checkpointOutputPrefix}},
 --prev-checkpoint, {{.prevCheckpointPrefix}},
 --run-name, {{.runName}},
 --name, {{.actionName}}]

It appends image-cache data when present. A code bundle contributes either --tgz <bundle> or --pkl <bundle>, followed by --dest <destination>. If there is no code bundle, or the bundle is not a pickle bundle, the method also appends --resolver with DefaultTaskResolver.import_path and that resolver's loader arguments for the task and serialization root directory. The method finishes by asserting that every command argument is a string.

Cache serialization also uses the preserved function. When caching is enabled, task_serde uses the pickle bundle's computed version when a pickle bundle exists. Otherwise, for an AsyncFunctionTaskTemplate, it creates VersionParameters(func=task.func, image=task.image) and asks the cache for a version. Thus the original callable remains relevant after decoration; it participates in both reconstruction and cache-version calculation.

_internal/resolvers/_task_module.py resolves that callable by inspecting task.func. The function must have an importable module with a file located relative to the serialization source_dir; otherwise resolution raises a ValueError for a missing module or the relative-path calculation can fail. A __main__ module is handled specially using the main script's file. Dynamically defined or out-of-tree functions therefore need attention before remote serialization.

Extensions through flyte.extend​

Application code normally uses @env.task; it does not manually instantiate AsyncFunctionTaskTemplate. Extension authors can import the template from flyte.extend, which re-exports it from the private task module along with the process-local TaskPluginRegistry:

from flyte.extend import AsyncFunctionTaskTemplate, TaskPluginRegistry

An environment with plugin_config follows a different construction branch. TaskEnvironment.task looks up an exact configuration type with TaskPluginRegistry.find(config_type=type(self.plugin_config)). A plugin is registered by mapping a configuration type to an AsyncFunctionTaskTemplate subclass through TaskPluginRegistry.register(config_type, plugin). If no exact registration exists, decoration raises ValueError and instructs you to register a plugin with flyte.extend.TaskPluginRegistry.register().

Without plugin_config, the factory uses AsyncFunctionTaskTemplate directly. With it, the registered subclass receives the same function and task metadata construction inputs, including the plugin configuration. A reusable environment cannot also specify plugin_config; TaskEnvironment.__post_init__ rejects that combination. Reusable environments also reject a task-level pod_template, and synchronous tasks in a reusable environment with concurrency greater than one are rejected because that configuration requires an async function.

Configuration and pitfalls​

  • Keep local calls and nested task calls distinct. Outside a task context, a call reaches the Python function through forward. Inside one, the same syntax submits to the controller rather than directly running the function.
  • Do not assume the template is always awaitable. Async functions return a coroutine through forward; synchronous functions return a value and set _call_as_synchronous. Use .aio(...) when a synchronous task is consumed from an async parent.
  • Do not call execute to simulate a local invocation. It asserts that a task context already exists and is intended for runtime execution through run_task.
  • Treat names as task identity inputs. The full name is formed from the environment name and function name, while the short name defaults to the function name. Changing either changes the constructed metadata.
  • Check reusable-environment constraints. TaskEnvironment rejects plugin_config with reusable; it rejects a pod template on a reusable task; and it rejects synchronous functions when reusable concurrency is greater than one.
  • Make the function resolvable. Serialization inspects the original function's module and requires its source file to be relative to the serialization root. A function that cannot satisfy this requirement cannot be reconstructed by the default resolver.
  • Know which bundle mode you are using. Pickle bundles add --pkl and use the bundle's computed cache version; non-pickle or absent bundles add the default resolver and its loader arguments.

flyte.extend is therefore the extension-facing namespace, while TaskEnvironment.task is the normal boundary where a Python callable becomes a tracked AsyncFunctionTaskTemplate with an interface, execution metadata, runtime lifecycle, and serialization behavior.