Tasks: The Core Executable Unit

When you want to execute a specific piece of logic—like a data cleaning script or a machine learning training loop—within a Flyte workflow, you use a Task. In flytekit, a Task is the fundamental executable unit that encapsulates a specific operation, its inputs, its outputs, and its execution configuration.

Defining Tasks with @task

The most common way to create a task in flytekit is by using the @task decorator on a Python function. This automatically creates an instance of PythonFunctionTask (defined in flytekit.core.python_function_task).

from flytekit import task

@task(cache=True, cache_version="1.0", retries=3)
def square(i: int) -> int:
    return i * i

Internally, PythonFunctionTask uses transform_function_to_interface to inspect your function's type hints and docstrings, creating a TypedInterface that Flyte understands. When this task is serialized for the Flyte platform, it captures the module path and function name so the Flyte agent can rehydrate and execute the correct code in a container.

Configuring Task Behavior

You can control how a task behaves—such as how many times it retries on failure or whether it caches results—by passing arguments to the @task decorator. These arguments are stored in a TaskMetadata object (found in flytekit.core.base_task).

Caching and Retries

Flytekit enforces specific rules for metadata configuration. For example, if you enable caching, you must provide a version string. The TaskMetadata.__post_init__ method validates these requirements:

• Caching: Set cache=True and provide a cache_version. If cache_version is missing, flytekit raises a ValueError.
• Retries: Set retries=n. The Task.metadata.retry_strategy property converts this into a model the Flyte backend can process.
• Timeouts: You can pass an int (seconds) or a datetime.timedelta. TaskMetadata automatically converts integers to timedelta objects.

# Example of metadata validation in base_task.py
if self.cache and not self.cache_version:
    raise ValueError("Caching is enabled ``cache=True`` but ``cache_version`` is not set.")

Extending Tasks for Custom Logic

While @task covers most use cases, flytekit provides base classes for building specialized task types that don't map directly to a single Python function.

PythonTask

The PythonTask class (in base_task.py) serves as the base for tasks that have a Python-native interface but might not execute a standard Python function. A primary example is the SQLTask, which uses the Python interface to define inputs and outputs but executes a SQL query against a database instead of running Python code.

PythonInstanceTask

If you need to create a task where the logic is encapsulated within a class rather than a function, you inherit from PythonInstanceTask (in python_function_task.py). This is useful for tasks like ShellTask, where the task instance stores a script and overrides the execute method to run that script in a subprocess.

# Simplified example based on flytekit.extras.tasks.shell.ShellTask
class MyShellTask(PythonInstanceTask):
    def __init__(self, name, script, **kwargs):
        super().__init__(name=name, task_config=None, task_type="shell", **kwargs)
        self._script = script

    def execute(self, **kwargs) -> Any:
        # Custom execution logic here
        import subprocess
        return subprocess.run(self._script, shell=True)

The Task Execution Lifecycle

When a task is executed, flytekit follows a structured lifecycle managed by the Task.dispatch_execute method in base_task.py. This process ensures that data is correctly translated between Flyte's type system and Python native types.

1. Pre-execution: The pre_execute method is called. This is used to set up the environment, such as initializing a Spark session or modifying the ExecutionParameters.
2. Input Translation: The _literal_map_to_python_input method converts Flyte LiteralMap objects (the format Flyte uses to store data) into standard Python arguments (kwargs) using the TypeEngine.
3. Execution: The execute method is invoked with the translated inputs. For a PythonFunctionTask, this simply calls your decorated function.
4. Post-execution: The post_execute method allows for cleanup or final output modification.
5. Output Translation: The _output_to_literal_map method converts the Python return values back into Flyte LiteralMap objects to be passed to the next task in the workflow.

Local vs. Remote Execution

When you run a task locally (e.g., in a unit test), flytekit uses local_execute. This method checks the LocalTaskCache if caching is enabled and handles the wrapping of results into Promise objects, allowing you to call tasks directly in your Python environment:

# Local execution call
result = square(i=5) 

Internally, local_execute calls sandbox_execute, which prepares a local FlyteContext before triggering the standard dispatch_execute flow.