Skip to main content

Task authoring and execution

Flyte tasks are the fundamental building blocks of flytekit. They represent a single unit of work, characterized by a strong interface (typed inputs and outputs) and independent executability. In flytekit, tasks are primarily authored using the @task decorator, which transforms a standard Python function into a PythonFunctionTask.

Declaring Tasks

To define a task, you decorate a Python function with @task. flytekit uses the function's type hints to automatically derive the task's interface.

from flytekit import task

@task
def add_one(x: int) -> int:
return x + 1

When you apply the @task decorator, flytekit creates an instance of PythonFunctionTask (found in flytekit/core/python_function_task.py). This class wraps your function and manages its metadata, interface transformation, and execution logic.

Constraints on Task Declaration

Tasks must be accessible at the module level so that the Flyte container can re-import and execute them. The PythonFunctionTask constructor explicitly checks for nested or local functions:

# From flytekit/core/python_function_task.py
if (
not istestfunction(func=task_function)
and isnested(func=task_function)
and not is_functools_wrapped_module_level(task_function)
):
raise ValueError(
"TaskFunction cannot be a nested/inner or local function. "
"It should be accessible at a module level for Flyte to execute it."
)

Configuring Task Behavior

The @task decorator accepts several parameters that control how the task is executed on the Flyte platform. These settings are encapsulated in the TaskMetadata class (defined in flytekit/core/base_task.py).

Retries and Timeouts

You can specify how many times a task should be retried on failure and the maximum duration it is allowed to run.

import datetime
from flytekit import task

@task(retries=3, timeout=datetime.timedelta(minutes=5))
def flaky_task(x: int) -> int:
...

Internally, TaskMetadata handles the conversion of these values into the Flyte IDL models. For example, timeout can be passed as an int (seconds) or a timedelta, and it is normalized in __post_init__.

Caching

Caching allows you to skip task execution if the inputs haven't changed. To enable it, you must provide a cache_version.

@task(cache=True, cache_version="1.0")
def expensive_computation(x: int) -> int:
...

If you change the logic inside the function but keep the signature the same, you should manually bump the cache_version to invalidate existing cached results.

Resource Requirements

You can request specific compute resources like CPU, memory, or GPUs using the requests and limits parameters.

from flytekit import task, Resources

@task(requests=Resources(cpu="2", mem="500Mi"), limits=Resources(cpu="4", mem="1Gi"))
def memory_intensive_task(data: list) -> int:
...

Task Execution and Promises

When you call a task inside a @workflow, it does not return the actual result of the function. Instead, it returns a Promise object (defined in flytekit/core/promise.py).

The Role of Promises

A Promise represents a future value that will be produced by a Node in the Flyte graph. This allows flytekit to build the execution DAG without actually running the code.

from flytekit import task, workflow

@task
def get_val() -> int:
return 10

@workflow
def my_wf() -> int:
# result is a Promise, not an int
result = get_val()
return result

Because result is a Promise, you cannot use it for standard Python control flow (like if result > 0:) inside a workflow. The Promise class overrides comparison operators (like __gt__, __eq__) to return a ComparisonExpression instead of a boolean, which is used for Flyte's internal conditional logic.

Overriding Task Behavior at Call Time

You can override task configurations for a specific invocation within a workflow using the .with_overrides() method on the returned Promise.

@workflow
def my_wf(x: int):
# Override the name and resource limits for this specific call
t1(x=x).with_overrides(node_name="custom-node-name", limits=Resources(cpu="1"))

This method interacts with the underlying Node (defined in flytekit/core/node.py) to update the metadata that will be sent to the Flyte backend.

Local Execution vs. Remote Execution

flytekit is designed to run the same code locally and on a remote Flyte cluster.

  1. Local Execution: When you call a task directly (e.g., in a unit test), Task.__call__ triggers local_execute. This method translates Python inputs into Flyte literals, checks the LocalTaskCache, and then calls execute (your decorated function).
  2. Remote Execution: On the cluster, the pyflyte-execute entrypoint uses a TaskResolver (like default_task_resolver in flytekit/core/python_auto_container.py) to load the task object and call dispatch_execute. This method handles the translation of literals from the Flyte engine back into Python native types before running your code.

Ignoring Outputs

In some scenarios, such as distributed training where only the rank-0 process should return results, you can raise IgnoreOutputs.

from flytekit.core.base_task import IgnoreOutputs

@task
def distributed_task():
if not_rank_zero:
raise IgnoreOutputs()
return "result"

This exception is caught by the flytekit execution engine to signal that the task completed successfully but produced no outputs to be recorded.