Skip to main content

Task Abstractions

Tasks are the fundamental building blocks of flytekit workflows. They represent a single unit of work that is versioned, strongly typed, and independently executable. In flytekit, tasks are implemented through a hierarchy of classes that bridge Python-native code with the Flyte backend.

The Task Hierarchy​

Flytekit uses a layered abstraction to handle different types of tasks, ranging from raw IDL-based tasks to Python-native functions.

Base Task​

The Task class in flytekit.core.base_task is the root of all tasks. It is designed to be closest to the FlyteIDL TaskTemplate. It captures metadata like retries, timeouts, and caching but does not have a Python-native interface. You typically won't instantiate this class directly; instead, it serves as the foundation for specialized task types.

Python Task​

The PythonTask class (also in flytekit.core.base_task) extends the base Task by adding a python_interface. This interface allows flytekit to map Python type hints (like int, str, or pandas.DataFrame) to Flyte's internal type system.

Python Function Task​

The PythonFunctionTask in flytekit.core.python_function_task is the most common task type. It is the underlying class created when you use the @task decorator. It automatically detects the interface of a Python function and handles the serialization required to run that function on the Flyte platform.

Creating Tasks with the @task Decorator​

The primary way to define a task in flytekit is using the @task decorator. This decorator transforms a standard Python function into a PythonFunctionTask.

from flytekit import task

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

When you call greet("World") locally, flytekit executes the function directly. When used inside a @workflow, flytekit uses the PythonFunctionTask metadata to compile a execution graph.

Task Configuration and Metadata​

The @task decorator accepts several arguments to control execution behavior. These are stored internally in the TaskMetadata class.

from datetime import timedelta
from flytekit import task, Resources

@task(
    retries=3,
    timeout=timedelta(minutes=60),
    requests=Resources(cpu="2", mem="4Gi"),
    limits=Resources(cpu="4", mem="8Gi"),
    cache=True,
    cache_version="1.0"
)
def heavy_computation(x: int) -> int:
    return x ** 2

Common configuration options include:

  • retries: The number of times to retry the task on failure.
  • timeout: A datetime.timedelta or int (seconds) specifying the maximum runtime.
  • cache: Enables caching of results based on input values.
  • requests / limits: Defines the Resources (CPU, Memory, GPU) required for the task.

Advanced Task Types​

Dynamic Tasks​

A dynamic task is a task that can generate a new workflow at runtime based on its inputs. This is useful for scenarios like processing a variable number of files. You create a dynamic task by setting the execution_mode to DYNAMIC (or using the @dynamic decorator, which is a shortcut).

Internally, PythonFunctionTask.dispatch_execute detects the DYNAMIC execution mode and invokes compile_into_workflow to produce a DynamicJobSpec.

Python Instance Tasks​

For tasks that do not have a user-defined function body but instead have a platform-defined execute method, flytekit provides PythonInstanceTask. This is an abstract base class used by plugins to create tasks that behave like objects.

# Example of a conceptual Instance Task implementation
class MyCustomTask(PythonInstanceTask):
    def __init__(self, name: str, **kwargs):
        super().__init__(name=name, task_config=None, task_type="custom", **kwargs)

    def execute(self, **kwargs) -> Any:
        # Custom logic here
        return "result"

Eager Tasks​

Eager tasks (or eager workflows) allow you to run Flyte entities using standard Python control flow (like if statements and for loops) while still executing each step as a Flyte task. These are implemented via EagerAsyncPythonFunctionTask. When an eager task is executed, it uses a Controller to manage a worker queue that communicates with the Flyte backend.

Internal Execution Flow​

When a task is executed, flytekit follows a specific lifecycle managed by the Task and PythonTask classes:

  1. pre_execute: Invoked before the task method. It can modify the ExecutionParameters to set up the environment (e.g., initializing a Spark session).
  2. dispatch_execute: The main entry point for execution. It handles:
    • Translating input LiteralMap (Flyte types) to Python native values using the TypeEngine.
    • Calling the user's execute method (the decorated function).
    • Handling different execution modes (Default, Dynamic, or Eager).
  3. post_execute: Invoked after the task completes. It can be used for cleanup or to alter the return values.
  4. Output Translation: The _output_to_literal_map method converts Python return values back into a Flyte LiteralMap.

Task Resolvers​

When a task runs in a container on a Flyte cluster, the system needs to know how to find and load the Python code. This is handled by the TaskResolverMixin. The default resolver (default_task_resolver) captures the module name and the function name. At runtime, the pyflyte-execute command uses these arguments to rehydrate the task object and run the code.