Skip to main content

Task authoring and execution

Flyte tasks are the fundamental building blocks of a pipeline. In flytekit, tasks are typically defined by decorating a Python function with the @task decorator, which transforms the function into a PythonFunctionTask. This abstraction handles the translation between Python's native types and the Flyte IDL, manages execution metadata, and facilitates both local and remote execution.

Declaring Tasks​

The most common way to author a task is using the @task decorator. flytekit uses Python type hints to automatically infer the task's interface (inputs and outputs).

from flytekit import task

@task
def square(n: int) -> int:
    return n * n

@task
def greet(name: str = "Flyte") -> str:
    return f"Hello, {name}!"

When you call square(n=5), flytekit performs the following:

  1. Interface Inference: It inspects the function signature to create a TypedInterface.
  2. Type Conversion: It uses the TypeEngine to map Python types (like int or str) to Flyte literal types.
  3. Metadata Association: It attaches default metadata (like retries or timeouts) defined in the TaskMetadata class.

Task Configuration and Metadata​

The @task decorator accepts several parameters that control the task's behavior on the Flyte platform. These settings are encapsulated in the TaskMetadata class.

  • Caching: To enable caching, you must provide both cache=True and a cache_version. Changing the version string invalidates previous cache entries.
  • Retries: The retries parameter defines how many times Flyte should attempt to re-run the task upon failure.
  • Resources: You can specify resource requests and limits for CPU, memory, and GPU.
from flytekit import task, Resources

@task(
    cache=True, 
    cache_version="1.0", 
    retries=3,
    requests=Resources(cpu="1", mem="500Mi"),
    limits=Resources(cpu="2", mem="1Gi"),
    environment={"STAGE": "production"}
)
def heavy_computation(data: list[int]) -> int:
    return sum(data)

Internally, TaskMetadata.__post_init__ ensures that configuration is valid—for example, it will raise a ValueError if cache=True is set without a cache_version.

Task Execution Flow​

Flytekit tasks support two primary execution modes: local execution (for testing) and remote execution (on a Flyte cluster).

Local Execution​

When you run a task locally, flytekit bypasses the containerization layer. The Task.local_execute method handles the conversion of inputs into Flyte literals, checks the LocalTaskCache if caching is enabled, and then calls dispatch_execute.

# In base_task.py, Task.local_execute handles the local lifecycle
def local_execute(self, ctx: FlyteContext, **kwargs):
    # 1. Translate inputs to literals
    literals = translate_inputs_to_literals(...)
    
    # 2. Check local cache
    if self.metadata.cache and local_config.cache_enabled:
        outputs_literal_map = LocalTaskCache.get(...)
        
    # 3. Execute the task
    if outputs_literal_map is None:
        outputs_literal_map = self.sandbox_execute(ctx, input_literal_map)
    
    # 4. Wrap outputs back into Promises
    return create_task_output(vals, self.python_interface)

Remote Execution and Serialization​

For remote execution, flytekit serializes the task into a TaskTemplate. A key part of this is the TaskResolverMixin. By default, flytekit uses the default_task_resolver, which records the module and function name.

When the task runs on a cluster, the container starts with a command like pyflyte-execute. This command uses the resolver to rehydrate the task object:

pyflyte-execute --inputs s3://... --output-prefix s3://... \
  --resolver flytekit.core.python_auto_container.default_task_resolver \
  -- task-module my_project.tasks task-name square

The PythonFunctionTask.dispatch_execute method is the entry point during runtime. It converts the input LiteralMap back into Python native types, calls your function, and then converts the return values back into a LiteralMap to be sent back to the Flyte engine.

Specialized Task Behaviors​

The PythonFunctionTask class supports different execution behaviors via the ExecutionBehavior enum:

  1. DEFAULT: Standard task execution where the function body runs once.
  2. DYNAMIC: The task function returns a dynamic workflow. flytekit compiles the returned entities into a DynamicJobSpec at runtime.
  3. EAGER: Used for "Eager Workflows" where the Python code acts as a coordinator, potentially awaiting other Flyte tasks as if they were local async calls.

Dynamic Tasks​

If a task is marked as dynamic (usually via the @dynamic decorator, which sets execution_mode=ExecutionBehavior.DYNAMIC), dispatch_execute calls compile_into_workflow. This allows you to generate new tasks and workflows based on runtime data.

# python_function_task.py
def execute(self, **kwargs) -> Any:
    if self.execution_mode == self.ExecutionBehavior.DEFAULT:
        return self._task_function(**kwargs)
    elif self.execution_mode == self.ExecutionBehavior.DYNAMIC:
        return self.dynamic_execute(self._task_function, **kwargs)

Decks and Observability​

Tasks can generate "Flyte Decks"—HTML visualizations of task data. The PythonTask class manages this via enable_deck and deck_fields. By default, tasks can generate decks for SOURCE_CODE, DEPENDENCIES, INPUT, and OUTPUT. These are rendered during post_execute using various Renderer classes.