Skip to main content

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)
  • Task is the root class closest to the FlyteIDL spec. It stores task_type, name, interface (as a TypedInterface), metadata (TaskMetadata), a security context, and docs. It auto-registers each instance via FlyteEntities.entities.append(self).
  • PythonTask adds a Python-native Interface (inputs and output types as Python classes), a generic task_config for plugin-specific configuration, environment variables, deck controls (enable_deck, deck_fields), and automatic Documentation generation from the function's docstring.
  • PythonAutoContainerTask adds container-image management, resource specs, secrets, pod templates, accelerators, and the task-resolver plumbing that produces the pyflyte-execute command line at serialization time.
  • PythonFunctionTask wraps a decorated Python function with auto-detected input/output types and supports three execution modes (standard, dynamic, eager).
  • PythonInstanceTask is 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:

ModeEnum valueBehavior
DefaultDEFAULTCalls self._task_function(**kwargs) directly
DynamicDYNAMICCompiles the function body into a workflow at execution time, returns a DynamicJobSpec
EagerEAGERRuns 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:

  1. Creates a PythonFunctionWorkflow from the task function
  2. Serializes it via get_serializable
  3. 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 @eager entities. Conditionals are not supported — use plain Python if statements. For remote execution, client_secret_group and client_secret_key are required for authentication (unless using a local sandbox started with flytectl 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__:

AttributeTypeDefaultDescription
cacheboolFalseEnable output caching. When True, cache_version must be set.
cache_serializeboolFalseSerialize execution of identical-input instances when caching is enabled. Requires cache=True.
cache_versionstr""Version string for cached outputs.
cache_ignore_input_varsTuple[str, ...]()Input variables to exclude from cache hash calculation. Requires cache=True.
interruptibleOptional[bool]NoneAllow scheduling on lower-QoS/preemptible nodes.
deprecatedstr""Warning message for deprecated tasks. Empty string means active.
retriesint0Number of retries on failure.
timeoutOptional[Union[timedelta, int]]NoneMax duration for one execution. int is treated as seconds.
pod_template_nameOptional[str]NoneName of an existing PodTemplate resource.
generates_deckboolFalseWhether the task generates a Deck URI.
is_eagerboolFalseTreat 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:

  1. pre_execute — Called before input conversion. Returns modified ExecutionParameters. The base implementation is a no-op.
  2. Input conversion_literal_map_to_python_input converts the Flyte LiteralMap to Python-native kwargs using TypeEngine.literal_map_to_kwargs. Exceptions during conversion are wrapped as FlyteUserRuntimeException in remote mode.
  3. 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 as FlyteUserRuntimeException so the failure is recorded without crashing the Flyte engine.
  4. post_execute — Receives the return value and can transform it. The base implementation returns the value unchanged. The IgnoreOutputs exception can be raised here to signal that outputs should be discarded.
  5. Output conversion_output_to_literal_map converts native outputs back to a LiteralMap. It uses TypeEngine.async_to_literal with asyncio.gather. If a single output is a NamedTuple (detected via output_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.
  6. Short-circuit — If native_outputs is already a LiteralMap or DynamicJobSpec (as returned by dynamic tasks), steps 4–5 are skipped.
  7. Deck writing_write_decks generates HTML decks for inputs, outputs, source code, and dependencies when enable_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:

  1. Translates incoming values (Promises or native constants) to a LiteralMap via translate_inputs_to_literals
  2. Checks self.metadata.cache and LocalConfig.cache_enabled: if caching is enabled, it looks up the LocalTaskCache using the task name, cache version, and inputs (minus ignored inputs). On a cache hit, the saved LiteralMap is returned without executing. On a cache miss (or if cache_overwrite is set), execution proceeds and the result is stored back to the cache.
  3. Calls sandbox_execute which wraps the execution in a sandbox context.
  4. Wraps output literals back into Promise objects 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 class
  • loader_args(settings, t) — arguments that identify the specific task (e.g., module name and task name)
  • load_task(loader_args) — rehydrates the Task from those arguments
  • get_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_resolver is provided via the compilation context, it overrides any resolver passed directly to @task. The PythonAutoContainerTask.__init__ checks FlyteContextManager.current_context().compilation_state.task_resolver first (lines 134–144 of python_auto_container.py).

Gotchas and Pitfalls

  • Nested functions are rejected: PythonFunctionTask.__init__ raises ValueError if the decorated function is nested (not accessible at module level) and using default_task_resolver. Test functions (names starting with test_) and functions wrapped with functools.wraps at module level are exempt.
  • node_dependency_hints is only for dynamic tasks: Setting this on a DEFAULT-mode task raises ValueError. 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_deck is deprecated: Setting disable_deck triggers a FutureWarning. Use enable_deck=True instead. Setting both raises ValueError.
  • ReferenceTask unsupported in dynamic: compile_into_workflow raises ValueError("Reference tasks are currently unsupported within dynamic tasks").
  • Async + DYNAMIC not implemented: AsyncPythonFunctionTask.async_execute raises NotImplementedError for ExecutionBehavior.DYNAMIC.
  • PythonAutoContainerTask initialization order: The _container_image attribute must be set before super().__init__() calls Task.__init__(), which appends the task to FlyteEntities. The translator iterates over FlyteEntities calling container_image(), so the attribute must exist at that point.
  • Output tuple handling: Single-element NamedTuple outputs require special handling in both dynamic_execute and _output_to_literal_map. The code checks self.python_interface.output_tuple_name as a proxy and unpacks the tuple accordingly.
  • Eager remote authentication: Running eager workflows against a remote Flyte installation requires client_secret_group and client_secret_key in the FlyteRemote config, unless using a local sandbox (demostack).