Task authoring and execution
Task Hierarchy Overview
Flytekit's task system is built on a layered class hierarchy in flytekit/core/base_task.py and flytekit/core/python_function_task.py. Each layer adds specific capabilities:
Task (base_task.py, line 200)
└── PythonTask (base_task.py, line 459)
└── PythonAutoContainerTask (python_auto_container.py, line 33)
├── PythonFunctionTask (python_function_task.py, line 90)
└── PythonInstanceTask (python_function_task.py, line 61)
Taskis the root class closest to the FlyteIDL spec. It storestask_type,name,interface(as aTypedInterface),metadata(TaskMetadata), a security context, and docs. It auto-registers each instance viaFlyteEntities.entities.append(self).PythonTaskadds a Python-nativeInterface(inputs and output types as Python classes), a generictask_configfor plugin-specific configuration, environment variables, deck controls (enable_deck,deck_fields), and automaticDocumentationgeneration from the function's docstring.PythonAutoContainerTaskadds container-image management, resource specs, secrets, pod templates, accelerators, and the task-resolver plumbing that produces thepyflyte-executecommand line at serialization time.PythonFunctionTaskwraps a decorated Python function with auto-detected input/output types and supports three execution modes (standard, dynamic, eager).PythonInstanceTaskis the abstract base for tasks without a user-defined function body — it captures the module and variable name so the loader can rehydrate it automatically at runtime.
Declaring Tasks with @task
The primary user-facing API is the @task decorator defined in flytekit/core/task.py (line 174). It inspects the decorated function's type annotations and docstring to build the task's interface automatically.
from flytekit import task
@task
def greet(name: str) -> str:
return f"Hello, {name}!"
This produces a PythonFunctionTask instance whose inputs are {"name": str} and outputs are str. The decorator supports extensive configuration:
from flytekit import task
from flytekit import Resources
@task(
retries=3,
timeout=600, # seconds, or a datetime.timedelta
interruptible=True,
cache=True,
cache_version="1.0",
container_image="my-registry/my-image:latest",
environment={"MY_VAR": "value"},
requests=Resources(cpu="2", mem="4Gi"),
limits=Resources(cpu="4", mem="8Gi"),
)
def compute(data: list[int]) -> float:
...
Cache configuration
The cache parameter accepts either a bool or a Cache object (from flytekit.core.cache). The Cache object is the modern approach and replaces the deprecated cache_serialize, cache_version, and cache_ignore_input_vars parameters (line 354–398 of task.py):
from flytekit import task, Cache
@task(cache=Cache(serialize=True, ignored_inputs=["context_var"]))
def expensive_query(date: str, context_var: str) -> pd.DataFrame:
...
When cache=True is passed without cache_version, the decorator creates a Cache with default settings and derives a version from the function signature, container image, pod template, and pod template name via Cache.get_version(VersionParameters(...)).
If you pass cache_serialize, cache_version, or cache_ignore_input_vars alongside a Cache object, the decorator raises ValueError — those three are deprecated (line 375–378).
Auto-detection of async functions
When the decorated function is a coroutine function (inspect.iscoroutinefunction(fn) is True), the decorator instantiates AsyncPythonFunctionTask instead of PythonFunctionTask (lines 411–424 of task.py):
@task
async def fetch_data(url: str) -> bytes:
...
Custom task plugins that register for async functions must subclass AsyncPythonFunctionTask; otherwise the decorator raises AssertionError.
VSCode debugging
If the environment variable FLYTE_ENABLE_VSCODE is set to "True", the decorate_function helper (line 520 of task.py) wraps the task function with vscode debugging support from flytekit.interactive.vscode.
Execution Modes
PythonFunctionTask.ExecutionBehavior (line 104 of python_function_task.py) defines three modes:
| Mode | Enum value | Behavior |
|---|---|---|
| Default | DEFAULT | Calls self._task_function(**kwargs) directly |
| Dynamic | DYNAMIC | Compiles the function body into a workflow at execution time, returns a DynamicJobSpec |
| Eager | EAGER | Runs eagerly via async; each Flyte entity call becomes a remote execution |
DEFAULT mode
This is the standard mode used by @task. The execute method (line 252 of python_function_task.py) simply calls the wrapped function with the resolved keyword arguments:
def execute(self, **kwargs) -> Any:
if self.execution_mode == self.ExecutionBehavior.DEFAULT:
return self._task_function(**kwargs)
DYNAMIC mode
The @dynamic decorator, defined in flytekit/core/dynamic_workflow_task.py, is a functools.partial of @task with execution_mode=DYNAMIC:
dynamic = functools.partial(task.task, execution_mode=PythonFunctionTask.ExecutionBehavior.DYNAMIC)
A dynamic task's function runs at execution time (not compilation time). Inside it, you can use Python control flow — loops, conditionals — to call other tasks. The result is compiled into a DynamicJobSpec that the Flyte engine executes as a subworkflow.
from flytekit import dynamic, task
@task
def t1(a: int) -> str:
return f"item-{a}"
@dynamic
def my_dynamic_subwf(a: int) -> tuple[list[str], int]:
s = []
for i in range(a):
s.append(t1(a=i))
return s, 5
The execution path goes through dynamic_execute (line 277 of python_function_task.py). When running locally, it compiles the function into a PythonFunctionWorkflow and executes it. When running in production (TASK_EXECUTION mode), it calls compile_into_workflow which:
- Creates a
PythonFunctionWorkflowfrom the task function - Serializes it via
get_serializable - Collects all referenced task templates into a
DynamicJobSpec
Dynamic tasks have a recommended size limit of under 50 tasks — large loops produce unwieldy workflow specs. Reference tasks (ReferenceTask) are not supported inside dynamic tasks and raise ValueError (line 152 of python_function_task.py).
EAGER mode
The @eager decorator (line 578 of task.py) creates an EagerAsyncPythonFunctionTask. Eager workflows treat every Flyte entity call as a remote execution — Python becomes propeller. Each task(), workflow(), or nested eager() call creates a real Flyte execution on the backend.
from flytekit import task, eager
@task
def add_one(x: int) -> int:
return x + 1
@task
def double(x: int) -> int:
return x * 2
@eager
async def eager_workflow(x: int) -> int:
out = add_one(x=x)
return double(x=out)
# run locally with asyncio
if __name__ == "__main__":
import asyncio
result = asyncio.run(eager_workflow(x=1))
print(f"Result: {result}") # "Result: 4"
Under the hood, EagerAsyncPythonFunctionTask.execute (line 539 of python_function_task.py) constructs a Controller from flytekit.core.worker_queue, which manages a queue of sub-executions against a FlyteRemote backend. The controller installs signal handlers for SIGINT and SIGTERM so that in-flight executions can be cleaned up on interruption.
In remote execution, run_with_backend (line 599) runs the task function with the worker queue context, and on failure it renders an "Eager Executions" deck and raises FlyteNonRecoverableSystemException. A companion EagerFailureHandlerTask (line 695) terminates any still-running sub-executions if the parent eager workflow fails.
Note: Eager workflows only support
@task,@workflow, and@eagerentities. Conditionals are not supported — use plain Pythonifstatements. For remote execution,client_secret_groupandclient_secret_keyare required for authentication (unless using a local sandbox started withflytectl demo start).
Task Configuration and Metadata
TaskMetadata (base_task.py, line 114) is a dataclass that holds the operational metadata for a task. All fields have defaults and are validated in __post_init__:
| Attribute | Type | Default | Description |
|---|---|---|---|
cache | bool | False | Enable output caching. When True, cache_version must be set. |
cache_serialize | bool | False | Serialize execution of identical-input instances when caching is enabled. Requires cache=True. |
cache_version | str | "" | Version string for cached outputs. |
cache_ignore_input_vars | Tuple[str, ...] | () | Input variables to exclude from cache hash calculation. Requires cache=True. |
interruptible | Optional[bool] | None | Allow scheduling on lower-QoS/preemptible nodes. |
deprecated | str | "" | Warning message for deprecated tasks. Empty string means active. |
retries | int | 0 | Number of retries on failure. |
timeout | Optional[Union[timedelta, int]] | None | Max duration for one execution. int is treated as seconds. |
pod_template_name | Optional[str] | None | Name of an existing PodTemplate resource. |
generates_deck | bool | False | Whether the task generates a Deck URI. |
is_eager | bool | False | Treat the task as eager. |
Validation rules enforced in __post_init__:
def __post_init__(self):
if self.timeout:
if isinstance(self.timeout, int):
self.timeout = datetime.timedelta(seconds=self.timeout)
elif not isinstance(self.timeout, datetime.timedelta):
raise ValueError(
"timeout should be duration represented as either a datetime.timedelta or int seconds"
)
if self.cache and not self.cache_version:
raise ValueError("Caching is enabled ``cache=True`` but ``cache_version`` is not set.")
if self.cache_serialize and not self.cache:
raise ValueError("Cache serialize is enabled but ``cache`` is not enabled.")
if self.cache_ignore_input_vars and not self.cache:
raise ValueError(
f"Cache ignore input vars are specified but ``cache`` is not enabled."
)
The retry_strategy property converts self.retries to a _literal_models.RetryStrategy object. The to_taskmetadata_model() method (line 172) produces the protobuf _task_model.TaskMetadata, embedding the Flyte SDK version and runtime metadata.
Task Plugins System
Flytekit provides a plugin registry in flytekit/core/task.py (line 31). The TaskPlugins class maps config types to task implementation classes:
from flytekit import TaskPlugins
class MyTaskConfig:
database: str = "default"
class MyCustomTask(PythonFunctionTask[MyTaskConfig]):
def execute(self, **kwargs) -> Any:
... # custom execution logic
TaskPlugins.register_pythontask_plugin(MyTaskConfig, MyCustomTask)
Once registered, you can use the config with @task:
@task(task_config=MyTaskConfig(database="analytics"))
def query_data(sql: str) -> pd.DataFrame:
...
The @task decorator calls TaskPlugins.find_pythontask_plugin(type(task_config)) (line 418 of task.py). If no plugin is registered for the config type, it defaults to PythonFunctionTask. If the function is async and the found plugin does not subclass AsyncPythonFunctionTask, the decorator raises an error.
Existing plugins include Athena, Hive, PyTorch, TensorFlow, and Pod task plugins. Each plugin defines its own task_config class and a PythonFunctionTask subclass that implements execute and serialization logic.
Execution Lifecycle (dispatch_execute)
The core execution pipeline lives in PythonTask.dispatch_execute (base_task.py, line 714). It is invoked both locally and on the Flyte platform:
pre_execute— Called before input conversion. Returns modifiedExecutionParameters. The base implementation is a no-op.- Input conversion —
_literal_map_to_python_inputconverts the FlyteLiteralMapto Python-native kwargs usingTypeEngine.literal_map_to_kwargs. Exceptions during conversion are wrapped asFlyteUserRuntimeExceptionin remote mode. execute(**native_inputs)— Runs the user's task function. Exceptions are caught: in local execution the original traceback is preserved; in remote execution they are wrapped asFlyteUserRuntimeExceptionso the failure is recorded without crashing the Flyte engine.post_execute— Receives the return value and can transform it. The base implementation returns the value unchanged. TheIgnoreOutputsexception can be raised here to signal that outputs should be discarded.- Output conversion —
_output_to_literal_mapconverts native outputs back to aLiteralMap. It usesTypeEngine.async_to_literalwithasyncio.gather. If a single output is aNamedTuple(detected viaoutput_tuple_name), the tuple is unpacked to match the declared output names. Output metadata (dynamic partitions, time partitions) is attached to the literal during this step. - Short-circuit — If
native_outputsis already aLiteralMaporDynamicJobSpec(as returned by dynamic tasks), steps 4–5 are skipped. - Deck writing —
_write_decksgenerates HTML decks for inputs, outputs, source code, and dependencies whenenable_deck=True.
# Simplified flow from dispatch_execute (base_task.py, line 714)
new_user_params = self.pre_execute(ctx.user_space_params)
native_inputs = self._literal_map_to_python_input(input_literal_map, exec_ctx)
native_outputs = self.execute(**native_inputs)
native_outputs = self.post_execute(new_user_params, native_outputs)
if isinstance(native_outputs, (_literal_models.LiteralMap, _dynamic_job.DynamicJobSpec)):
return native_outputs
literals_map, native_outputs_as_map = self._output_to_literal_map(native_outputs, exec_ctx)
self._write_decks(native_inputs, native_outputs_as_map, ctx, new_user_params)
return literals_map
Local Execution and Caching
The Task.local_execute method (base_task.py, line 282) handles the local execution path. It:
- Translates incoming values (Promises or native constants) to a
LiteralMapviatranslate_inputs_to_literals - Checks
self.metadata.cacheandLocalConfig.cache_enabled: if caching is enabled, it looks up theLocalTaskCacheusing the task name, cache version, and inputs (minus ignored inputs). On a cache hit, the savedLiteralMapis returned without executing. On a cache miss (or ifcache_overwriteis set), execution proceeds and the result is stored back to the cache. - Calls
sandbox_executewhich wraps the execution in a sandbox context. - Wraps output literals back into
Promiseobjects for the caller.
# From Task.local_execute (base_task.py, line 282)
if self.metadata.cache and local_config.cache_enabled:
if local_config.cache_overwrite:
outputs_literal_map = None # force execution
else:
outputs_literal_map = LocalTaskCache.get(
self.name, self.metadata.cache_version,
input_literal_map, self.metadata.cache_ignore_input_vars
)
if outputs_literal_map is None:
outputs_literal_map = self.sandbox_execute(ctx, input_literal_map)
LocalTaskCache.set(...)
else:
outputs_literal_map = self.sandbox_execute(ctx, input_literal_map)
Testing Tasks
Flytekit provides task_mock in flytekit/core/testing.py (line 13) for unit-testing tasks. It is a context manager that patches task.execute with a MagicMock:
from flytekit import task
from flytekit.testing import task_mock
@task
def t1(i: int) -> int:
... # real implementation, not used in test
with task_mock(t1) as m:
m.side_effect = lambda x: x # replace implementation
result = t1(10)
# result == 10
# The mock is automatically removed when the context exits
The patch decorator (line 50) provides the same functionality as a test decorator, injecting the MagicMock as an argument to the test function. Both tools work on any PythonTask, WorkflowBase, or ReferenceEntity instance.
Task Resolution
At serialization time, PythonAutoContainerTask.get_default_command (python_auto_container.py, line 164) builds the pyflyte-execute command line:
pyflyte-execute --inputs {{.input}} --output-prefix {{.outputPrefix}} \
--raw-output-data-prefix {{.rawOutputDataPrefix}} \
--checkpoint-path {{.checkpointOutputPrefix}} \
--prev-checkpoint {{.prevCheckpointPrefix}} \
--resolver <resolver_location> -- <loader_args>
The resolver's location and loader_args identify which task to rehydrate at runtime. The TaskResolverMixin abstract class (base_task.py, line 858) defines the contract:
location— a fully-qualified Python path to the resolver classloader_args(settings, t)— arguments that identify the specific task (e.g., module name and task name)load_task(loader_args)— rehydrates theTaskfrom those argumentsget_all_tasks()— future-proof method for enumerating tasks
The default resolver (default_task_resolver in python_auto_container.py) stores the module and the task's variable name. When load_task is called, it does importlib.import_module(module) and looks up the task by name. Custom resolvers can be passed via task_resolver= in @task or set globally through the compilation context.
Gotcha: When a
task_resolveris provided via the compilation context, it overrides any resolver passed directly to@task. ThePythonAutoContainerTask.__init__checksFlyteContextManager.current_context().compilation_state.task_resolverfirst (lines 134–144 ofpython_auto_container.py).
Gotchas and Pitfalls
- Nested functions are rejected:
PythonFunctionTask.__init__raisesValueErrorif the decorated function is nested (not accessible at module level) and usingdefault_task_resolver. Test functions (names starting withtest_) and functions wrapped withfunctools.wrapsat module level are exempt. node_dependency_hintsis only for dynamic tasks: Setting this on aDEFAULT-mode task raisesValueError. It is an optional hint for dynamic tasks when Flyte cannot statically determine the dependencies (e.g., when calling launch plans inside a dynamic workflow).disable_deckis deprecated: Settingdisable_decktriggers aFutureWarning. Useenable_deck=Trueinstead. Setting both raisesValueError.- ReferenceTask unsupported in dynamic:
compile_into_workflowraisesValueError("Reference tasks are currently unsupported within dynamic tasks"). - Async + DYNAMIC not implemented:
AsyncPythonFunctionTask.async_executeraisesNotImplementedErrorforExecutionBehavior.DYNAMIC. PythonAutoContainerTaskinitialization order: The_container_imageattribute must be set beforesuper().__init__()callsTask.__init__(), which appends the task toFlyteEntities. The translator iterates overFlyteEntitiescallingcontainer_image(), so the attribute must exist at that point.- Output tuple handling: Single-element
NamedTupleoutputs require special handling in bothdynamic_executeand_output_to_literal_map. The code checksself.python_interface.output_tuple_nameas a proxy and unpacks the tuple accordingly. - Eager remote authentication: Running eager workflows against a remote Flyte installation requires
client_secret_groupandclient_secret_keyin theFlyteRemoteconfig, unless using a local sandbox (demostack).