Task authoring and execution

Flyte tasks are the fundamental building blocks of a workflow. In flytekit, tasks are versioned, strongly typed, and independently executable. They represent a discrete unit of work that can be executed locally or on a Flyte cluster.

Declaring Tasks with the @task Decorator

The most common way to define a task in flytekit is by using the @task decorator on a Python function. This wraps the function in a PythonFunctionTask object, which automatically extracts the task's interface (inputs and outputs) from Python type hints.

from flytekit import task

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

When you call this function, flytekit handles the translation between Python native types and the Flyte IDL (Interface Definition Language).

Task Configuration and Metadata

The @task decorator accepts several parameters to control execution behavior, resource allocation, and caching. These are stored internally in the TaskMetadata class (found in flytekit/core/base_task.py).

• Caching: Enable caching to avoid re-running tasks with the same inputs.

  from flytekit import task, Cache
  
  @task(cache=True, cache_version="1.0")
  def expensive_computation(x: int) -> int:
      ...
• Retries: Specify how many times Flyte should retry the task on failure.

  @task(retries=3)
  def flaky_task(x: int) -> int:
      ...
• Resources: Request specific CPU, memory, or GPU resources using the Resources class.

  from flytekit import task, Resources
  
  @task(requests=Resources(cpu="2", mem="500Mi"), limits=Resources(cpu="4", mem="1Gi"))
  def resource_intensive_task(x: int) -> int:
      ...

Task Execution Modes

Flytekit supports different execution behaviors via the PythonFunctionTask.ExecutionBehavior enum.

Default Execution

In the DEFAULT mode, the task function is executed as a standard containerized process. The inputs are downloaded, the function is run, and the outputs are uploaded.

Dynamic Workflows

Dynamic workflows (declared with the @dynamic decorator) allow you to generate a workflow structure at runtime based on task inputs. Internally, a @dynamic task is a PythonFunctionTask with execution_mode=ExecutionBehavior.DYNAMIC.

from flytekit import task, dynamic

@task
def process_item(item: int) -> int:
    return item * 2

@dynamic
def dynamic_wf(items: list[int]) -> list[int]:
    return [process_item(item=i) for i in items]

When a dynamic task runs, it executes the function body to produce a DynamicJobSpec (defined in flytekit/core/python_function_task.py). This spec contains a list of new nodes and tasks that Flyte Propeller then schedules as a subworkflow.

Eager Tasks

Eager tasks (using EagerAsyncPythonFunctionTask) allow for a more imperative, async-style execution where Python code acts as the orchestrator. Every task call within an eager task creates a new execution on the Flyte cluster rather than running in the local memory stack. This is useful for complex logic that requires immediate feedback from task executions.

Internal Task Abstractions

Flytekit uses a hierarchy of classes to manage task definitions and execution:

1. Task (flytekit/core/base_task.py): The base class that captures the Flyte IDL TaskTemplate. It defines the core interface for dispatch_execute, pre_execute, and execute.
2. PythonTask (flytekit/core/base_task.py): A specialization for tasks with Python-native interfaces. It handles the conversion between Flyte literals and Python objects using the TypeEngine.
3. PythonFunctionTask (flytekit/core/python_function_task.py): The implementation for tasks defined by a Python function. It manages the task_function and handles different ExecutionBehavior modes.

The Execution Flow

When a task is executed (either locally or on the cluster), the following sequence occurs within dispatch_execute:

1. pre_execute: Sets up the execution environment (e.g., initializing a Spark session or configuring decks).
2. Input Translation: The input_literal_map (Flyte IDL) is converted to Python native kwargs using _literal_map_to_python_input.
3. execute: The actual user-defined function is called with the translated inputs.
4. post_execute: Performs any necessary cleanup or output modification.
5. Output Translation: Python return values are converted back into a LiteralMap via _output_to_literal_map.

Task Resolvers

When a task runs on a hosted Flyte platform, the container needs to know how to find and load the specific Python task object. This is handled by TaskResolverMixin.

The default_task_resolver (in flytekit/core/python_auto_container.py) serializes the task's module path and function name into the container's command-line arguments. At runtime, pyflyte-execute uses these arguments to import the module and retrieve the task instance.

# Example of how a task is invoked in a container
# pyflyte-execute --resolver flytekit.core.python_auto_container.default_task_resolver \
# -- task-module my_module task-name my_task

Parallel Execution with map_task

To run a task over a collection of inputs in parallel, use map_task. This creates a specialized task that Flyte can optimize for large-scale parallel processing.

from flytekit import task, workflow, map_task

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

@workflow
def my_workflow(inputs: list[int]) -> list[int]:
    return map_task(square)(x=inputs)

The map_task function supports configuration for concurrency (limiting the number of parallel instances) and min_success_ratio (allowing the task to succeed even if some sub-tasks fail). Internally, this uses either the legacy MapTask or the newer ArrayNode depending on the platform configuration.