Skip to main content

Tasks: Base Classes and Execution Model

The Task Class Hierarchy​

Every entity you write with @task in flytekit sits somewhere in a five-level class hierarchy. Understanding where a class lives tells you what capabilities it has:

Task                          # IDL-level abstraction, no Python types
└── PythonTask[T]             # Adds Python Interface, dispatch_execute pipeline
    └── PythonAutoContainerTask[T]   # Adds container_image, resources, task_resolver
        ├── PythonFunctionTask[T]    # Wraps a Python callable (the @task decorator target)
        │   └── AsyncPythonFunctionTask[T]    # Async function support
        │       └── EagerAsyncPythonFunctionTask[T]  # @eager workflows
        └── PythonInstanceTask[T]   # Tasks without a user function body

The separation between levels is deliberate: Task knows nothing about Python types; PythonTask adds the Python type system but nothing about deployment; PythonAutoContainerTask adds the runtime container machinery. Plugin authors typically extend one of the middle layers — PythonFunctionTask when they need a user function, PythonInstanceTask when they supply their own execute().

Task and TaskMetadata​

Task (in flytekit/core/base_task.py) is the lowest-level class in flytekit, closest to the FlyteIDL TaskTemplate. It holds:

  • task_type — a string identifier that backend plugins use to route execution (e.g., "python-task", "spark", "ray")
  • name — globally unique across the Flyte installation
  • interface — a TypedInterface (IDL-level types, not Python types)
  • metadata — a TaskMetadata instance
  • security_ctx — a SecurityContext for secrets

On construction, every Task appends itself to FlyteEntities.entities, which the serialization tooling later iterates to discover all tasks in a module.

Configuring tasks with TaskMetadata​

TaskMetadata is a dataclass that carries all task-level behavioral configuration:

from flytekit import task

@task(
    cache=True,
    cache_version="1.0",
    cache_serialize=True,
    retries=3,
    timeout=300,          # seconds, converted to timedelta internally
    interruptible=True,
    deprecated="Use new_task instead",
)
def my_task(x: int) -> int:
    return x * 2

TaskMetadata.__post_init__ enforces several constraints that will raise ValueError immediately at decoration time rather than at runtime:

  • cache=True requires cache_version to be non-empty
  • cache_serialize=True requires cache=True
  • cache_ignore_input_vars requires cache=True
  • An integer timeout is converted to datetime.timedelta(seconds=timeout)
from flytekit.core.base_task import TaskMetadata

# This raises at construction - cache_version is missing
TaskMetadata(cache=True)  # ValueError: Caching is enabled but cache_version is not set.

# This also raises - cache_serialize requires cache
TaskMetadata(cache_serialize=True)  # ValueError

TaskMetadata.to_taskmetadata_model() converts the dataclass to the _task_model.TaskMetadata protobuf for serialization. It also reads the flytekit version string to populate the RuntimeMetadata field.

PythonTask: The Execution Pipeline​

PythonTask (still in base_task.py) adds everything needed to work with Python types: a python_interface (an Interface instance holding Python-native input/output types), a generic task_config: T, and an environment dict for container environment variables.

The dispatch_execute pipeline​

When a task runs — locally or remotely — the call chain is:

local_execute()       ← called locally (wraps with sandbox context)
    └── sandbox_execute()
            └── dispatch_execute()    ← the real workhorse
                    ├── pre_execute()         # plugin setup hook
                    ├── _literal_map_to_python_input()  # IDL → Python
                    ├── execute()             # user code
                    ├── post_execute()        # plugin teardown hook
                    └── _output_to_literal_map()  # Python → IDL

dispatch_execute in PythonTask does the full translation. It first calls the overridable pre_execute() hook (which plugins like Spark use to set up a SparkSession), then converts the incoming LiteralMap to Python kwargs via TypeEngine.literal_map_to_kwargs, calls execute(**native_inputs), calls post_execute(), and finally converts Python outputs back to a LiteralMap via TypeEngine.async_to_literal.

Error handling in dispatch_execute differs between local and remote execution. Locally, the original exception propagates directly (with an added context message). Remotely, exceptions are wrapped in FlyteUserRuntimeException or FlyteNonRecoverableSystemException so the Flyte control plane can classify them correctly.

Local caching​

When running locally, local_execute checks LocalConfig.auto() for cache_enabled. If the task has metadata.cache=True and local caching is enabled, it queries LocalTaskCache before calling sandbox_execute. Cache hits skip execution entirely; cache misses store the result after execution:

# from base_task.py — the local caching path inside local_execute
if self.metadata.cache and local_config.cache_enabled:
    outputs_literal_map = LocalTaskCache.get(
        self.name, self.metadata.cache_version, input_literal_map,
        self.metadata.cache_ignore_input_vars
    )
    if outputs_literal_map is None:
        outputs_literal_map = self.sandbox_execute(ctx, input_literal_map)
        LocalTaskCache.set(
            self.name, self.metadata.cache_version, input_literal_map,
            self.metadata.cache_ignore_input_vars, outputs_literal_map
        )

The pre_execute / post_execute hooks​

PythonTask.pre_execute() is a no-op by default — it just returns the user_params it receives. Plugins override it to perform setup before type conversion happens. The Spark plugin, for example, creates and caches a SparkSession and injects it as SPARK_SESSION in the returned ExecutionParameters:

# from plugins/flytekit-spark/flytekitplugins/spark/task.py
def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters:
    import pyspark as _pyspark

    sess_builder = _pyspark.sql.SparkSession.builder.appName(
        f"FlyteSpark: {user_params.execution_id}"
    )
    # ... configure spark conf ...
    self.sess = sess_builder.getOrCreate()
    return user_params.builder().add_attr("SPARK_SESSION", self.sess).build()

post_execute() is similarly a no-op unless overridden. It receives the return value from execute() and can transform it. This is the correct place for cleanup logic that runs even when execution succeeds.

TaskResolverMixin: The Two-Phase Lifecycle​

Every PythonAutoContainerTask participates in a two-phase lifecycle: serialization, which happens locally when you register and push workflows, and execution, which happens inside the container that Flyte spins up in the cloud.

TaskResolverMixin (in base_task.py) defines the contract:

class TaskResolverMixin:
    @property
    @abstractmethod
    def location(self) -> str: ...      # importable dotted path to this resolver instance

    @abstractmethod
    def loader_args(self, settings: SerializationSettings, t: Task) -> List[str]: ...
    # Returns opaque list of strings that identify the task

    @abstractmethod
    def load_task(self, loader_args: List[str]) -> Task: ...
    # At execution time: reconstructs the task object from loader_args

At serialization time, PythonAutoContainerTask.get_default_command() builds the pyflyte-execute command that the container will run:

# from python_auto_container.py
container_args = [
    "pyflyte-execute",
    "--inputs", "{{.input}}",
    "--output-prefix", "{{.outputPrefix}}",
    "--raw-output-data-prefix", "{{.rawOutputDataPrefix}}",
    "--checkpoint-path", "{{.checkpointOutputPrefix}}",
    "--prev-checkpoint", "{{.prevCheckpointPrefix}}",
    "--resolver", self.task_resolver.location,
    "--",
    *self.task_resolver.loader_args(settings, self),
]

The {{.input}} and similar tokens are template variables that Flyte propeller fills with actual storage paths at runtime.

DefaultTaskResolver​

The default resolver (default_task_resolver, an instance of DefaultTaskResolver in python_auto_container.py) handles the common case of module-level functions:

# loader_args at serialization time returns:
["task-module", "mypackage.workflows.example", "task-name", "my_task"]

# load_task at execution time does:
task_module = importlib.import_module("mypackage.workflows.example")
task_def = getattr(task_module, "my_task")
return task_def

This is why PythonFunctionTask raises ValueError if the decorated function is nested or local — importlib.import_module can reach module-level names but not closures:

def outer():
    @task  # ValueError: TaskFunction cannot be a nested/inner function
    def inner(x: int) -> int:
        return x

The exception is test modules: functions in modules whose names start with test_ are permitted to be nested, and using functools.wraps on a wrapper function also satisfies the check because it copies the __module__ and __qualname__ attributes.

DefaultNotebookTaskResolver​

Tasks defined in Jupyter notebooks can't be imported by module path because notebooks aren't Python modules. DefaultNotebookTaskResolver handles this case by reading from a gzip-compressed cloudpickle file (pkl.gz):

# from python_auto_container.py
def load_task(self, loader_args: List[str]) -> PythonAutoContainerTask:
    _, entity_name, *_ = loader_args
    with gzip.open(PICKLE_FILE_PATH, "r") as f:
        loaded_data = cloudpickle.load(f)
    # validates Python major.minor version match
    pickled_object: PickledEntity = loaded_data
    return pickled_object.entities[entity_name]

PickledEntity is a dataclass holding a dict of entity names to PythonAutoContainerTask instances, plus PickledEntityMetadata with the Python version string. The resolver validates that the major and minor Python version used to create the pickle matches the current runtime, raising RuntimeError otherwise.

PythonAutoContainerTask: The Container Layer​

PythonAutoContainerTask (in python_auto_container.py) extends PythonTask with everything needed to run inside a container on Flyte: the image, compute resources, task resolver, and pod configuration.

Container image​

The container_image parameter accepts either a plain string or an ImageSpec object. A plain string can use template variables like {{.images.default.fqn}}:{{.images.default.tag}}. When serializing, get_image() resolves these via get_registerable_container_image():

@task(container_image="myrepo/my-image:latest")
def specialized_task(x: int) -> int:
    return x

If container_image=None, flytekit uses the default image from SerializationSettings.image_config, which itself reads from the FLYTE_INTERNAL_IMAGE environment variable.

Note that _container_image is set before super().__init__() is called. This ordering matters: the parent constructor registers the task in FlyteEntities.entities, and translation code may immediately call container_image() on it while iterating that list.

Resources​

Compute resources are managed through ResourceSpec, constructed from either the resources shorthand or separate requests/limits parameters:

from flytekit import Resources

# Using the combined shorthand - cpu request=1, cpu limit=2, mem both=1Gi
@task(resources=Resources(cpu=("1", "2"), mem="1Gi"))
def compute_heavy(x: int) -> int: ...

# Using separate requests and limits
@task(
    requests=Resources(cpu="1", mem="1Gi"),
    limits=Resources(cpu="2", mem="2Gi"),
)
def also_valid(x: int) -> int: ...

# Mixing is not allowed - raises ValueError
@task(resources=Resources(cpu="1"), limits=Resources(mem="1Gi"))  # ValueError
def invalid(x: int) -> int: ...

PodTemplate vs plain container​

When pod_template is not set, get_container() assembles a Container definition with the image, resource spec, environment variables (merged from settings.env and self.environment), and the pyflyte-execute args. When pod_template is set, get_container() returns None and get_k8s_pod() returns a full K8sPod with the pod's spec — including injecting the pyflyte-execute args into the primary container:

@task(
    container_image="repo/image:0.0.0",
    pod_template=PodTemplate(
        primary_container_name="primary",
        labels={"team": "ml"},
        pod_spec=V1PodSpec(containers=[V1Container(name="primary")]),
    ),
)
def custom_pod_task(i: str) -> str:
    return i

When serialized, custom_pod_task.get_container(settings) returns None, but custom_pod_task.get_k8s_pod(settings) returns the pod spec with pyflyte-execute args in the primary container.

Overriding the resolver and command​

set_resolver() replaces the task resolver after construction — useful when you need a task to be resolved differently depending on context (e.g., notebook environments call set_resolver(default_notebook_task_resolver)).

set_command_fn() replaces the function that generates the container args. Map tasks use this to switch to pyflyte-map-execute. reset_command_fn() restores the default.

PythonFunctionTask and the @task Decorator​

PythonFunctionTask (in python_function_task.py) is what the @task decorator creates in the common case. It wraps a Python callable and auto-derives the Interface from the function's type annotations via transform_function_to_interface().

The @task decorator​

from flytekit import task, Resources
from flytekit.models.security import Secret

@task(
    cache=True,
    cache_version="2.1",
    retries=2,
    timeout=600,
    container_image="myorg/myimage:latest",
    environment={"PYTHONPATH": "/app"},
    requests=Resources(cpu="2", mem="4Gi"),
    limits=Resources(cpu="4", mem="8Gi"),
    secret_requests=[Secret(group="my-secrets", key="api-key")],
    interruptible=True,
    enable_deck=True,
)
def train_model(data_path: str, epochs: int) -> float:
    ...

The decorator creates a TaskMetadata from cache/retry/timeout/interruptible arguments, calls TaskPlugins.find_pythontask_plugin(type(task_config)) to select the right task class, and instantiates it. For task_config=None (the default), this returns plain PythonFunctionTask. For async functions, the decorator switches to AsyncPythonFunctionTask automatically.

TaskPlugins: the plugin registry​

When you write @task(task_config=Spark(...)), the task_config type (Spark) maps to PysparkFunctionTask through TaskPlugins. Plugins register themselves at module import time:

# from plugins/flytekit-spark/flytekitplugins/spark/task.py
TaskPlugins.register_pythontask_plugin(Spark, PysparkFunctionTask)
TaskPlugins.register_pythontask_plugin(Databricks, PysparkFunctionTask)

When you then write:

from flytekitplugins.spark import Spark

@task(task_config=Spark(spark_conf={"spark.executor.cores": "4"}))
def spark_task(data: str) -> int:
    ...

TaskPlugins.find_pythontask_plugin(Spark) returns PysparkFunctionTask, and the decorator instantiates that class instead of the default PythonFunctionTask. If a config type has no registered plugin, find_pythontask_plugin returns PythonFunctionTask as the default.

register_pythontask_plugin raises TypeError if you try to register a second different plugin for the same config type — duplicate registration of the same plugin class is a no-op.

ExecutionBehavior​

PythonFunctionTask.ExecutionBehavior is an enum with three values that control what happens in execute():

  • DEFAULT: calls self._task_function(**kwargs) directly
  • DYNAMIC: calls self.dynamic_execute(), which compiles the function body into a DynamicJobSpec at runtime
  • EAGER: used by EagerAsyncPythonFunctionTask

The execution_mode parameter to @task is mostly for internal use (the @dynamic decorator sets it to DYNAMIC). You will not normally set it yourself.

PythonInstanceTask: Tasks Without a User Function​

When the task logic is platform-defined rather than user-defined, PythonInstanceTask is the right base. It extends PythonAutoContainerTask and is instantiated as a module-level variable rather than used as a function decorator. The DBT plugin is a typical example:

# from plugins/flytekit-dbt/flytekitplugins/dbt/task.py
from flytekit import kwtypes
from flytekit.core.python_function_task import PythonInstanceTask

class DBTRun(PythonInstanceTask):
    def __init__(self, name: str, **kwargs):
        super().__init__(
            task_type="dbt-run",
            name=name,
            task_config=None,
            interface=Interface(
                inputs=kwtypes(input=DBTRunInput),
                outputs=kwtypes(output=DBTRunOutput),
            ),
            **kwargs,
        )

    def execute(self, **kwargs) -> DBTRunOutput:
        ...  # custom execution logic here

# Instantiated at module level, not used as decorator
dbt_run_task = DBTRun(name="my-dbt-run")

The workflow calls it like any other task:

@workflow
def my_workflow() -> DBTRunOutput:
    return dbt_run_task(input=DBTRunInput(project_dir="./jaffle_shop"))

PythonInstanceTask exists to ensure the module loader can find the task object by name (dbt_run_task) in the module, because the DefaultTaskResolver uses getattr(module, task_name) to rehydrate the task at runtime.

Async and Eager Tasks​

AsyncPythonFunctionTask​

When the decorated function is an async def, the @task decorator switches to AsyncPythonFunctionTask:

@task
async def fetch_data(url: str) -> bytes:
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            return await response.read()

AsyncPythonFunctionTask overrides __call__ to be async (routing through async_flyte_entity_call_handler) and provides async_execute() which awaits the task function. The synchronous execute is generated by loop_manager.synced(async_execute) — so the same task can be called from both sync and async contexts.

EagerAsyncPythonFunctionTask and @eager​

Eager workflows are a different execution model entirely. Rather than compiling a static workflow DAG, every @task or @workflow call inside an @eager function creates a Flyte execution on the backend and awaits the result. This means you can use arbitrary Python control flow:

from flytekit import task, eager
import asyncio

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

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

@eager
async def compute(x: int) -> int:
    out = add_one(x=x)
    if out > 10:              # regular Python branch — allowed in eager
        return out
    return double(x=out)

# Run locally:
# result = asyncio.run(compute(x=5))

Parallel execution is possible with asyncio.gather:

@eager
async def parallel_compute(a: int, b: int) -> tuple:
    t1 = asyncio.create_task(compute(x=a))
    t2 = asyncio.create_task(compute(x=b))
    r1, r2 = await asyncio.gather(t1, t2)
    return r1, r2

EagerAsyncPythonFunctionTask always sets metadata.is_eager=True and forces enable_deck=True. Passing execution_mode is ignored — the class deletes it from kwargs in __init__. The decorator requires an async def function; the eager() function in task.py simply instantiates EagerAsyncPythonFunctionTask directly.

At runtime, EagerAsyncPythonFunctionTask.execute() sets up a Controller (the worker queue), installs SIGINT/SIGTERM signal handlers (which must happen on the main thread), and runs async_execute via the event loop. If execution is local, async_execute runs the function directly with EAGER_LOCAL_EXECUTION mode set. If on the backend, it calls run_with_backend(), which sets EAGER_EXECUTION mode and delegates to the task function.

EagerFailureHandlerTask​

When an eager workflow fails, flytekit needs to terminate any sub-executions it launched. EagerAsyncPythonFunctionTask.get_as_workflow() wraps the eager function in an ImperativeWorkflow and registers an EagerFailureHandlerTask as the on_failure handler. At failure time, EagerFailureHandlerTask.dispatch_execute() queries Flyteadmin for all executions tagged with the parent execution's ID and terminates them.

The serialized form sets is_eager=True in the task template metadata:

# from tests/flytekit/unit/core/test_async.py
def test_serialization():
    se_spec = get_serializable(OrderedDict(), serialization_settings, simple_eager_workflow)
    assert se_spec.template.metadata.is_eager
    assert se_spec.template.metadata.generates_deck

IgnoreOutputs: Distributed Training Pattern​

IgnoreOutputs is a bare Exception subclass in base_task.py. When a task's execute() raises it, _dispatch_execute in flytekit/bin/entrypoint.py catches it as a FlyteUserRuntimeException and skips writing outputs.pb:

# from bin/entrypoint.py
except FlyteUserRuntimeException as e:
    if isinstance(e.value, IgnoreOutputs):
        logger.warning("User-scoped IgnoreOutputs received! Outputs.pb will not be uploaded.")
        return

The KFPyTorch plugin uses this for distributed training: only the rank-0 worker returns a real value; all other workers raise IgnoreOutputs after their local computation completes:

# from plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py
if 0 in out:
    return out[0].return_value   # rank 0 returns the result
else:
    raise IgnoreOutputs()        # other ranks: no outputs uploaded

Runtime Execution Flow​

When Flyte propeller launches a container for a task, the entrypoint is pyflyte-execute. The command line was baked in at serialization time by get_default_command() and looks like:

pyflyte-execute \
  --inputs s3://bucket/inputs.pb \
  --output-prefix s3://bucket/outputs/ \
  --raw-output-data-prefix /tmp/data \
  --resolver flytekit.core.python_auto_container.default_task_resolver \
  -- \
  task-module mypackage.workflows.example task-name my_task

The _dispatch_execute function in bin/entrypoint.py orchestrates the three-step runtime:

  1. Load: calls load_task() — for DefaultTaskResolver, this is importlib.import_module("mypackage.workflows.example") then getattr(module, "my_task"). If this fails, it raises FlyteUserRuntimeException.
  2. Execute: downloads inputs.pb, deserializes it into a LiteralMap, and calls task_def.dispatch_execute(ctx, idl_input_literals).
  3. Store: checks the return type — LiteralMap writes outputs.pb, DynamicJobSpec writes futures.pb, VoidPromise writes an empty outputs.pb. Exceptions write error.pb.

Propeller injects several environment variables into the container at launch: FLYTE_INTERNAL_EXECUTION_PROJECT, FLYTE_INTERNAL_EXECUTION_DOMAIN, FLYTE_INTERNAL_EXECUTION_ID, FLYTE_INTERNAL_TASK_PROJECT, FLYTE_INTERNAL_TASK_NAME, and FLYTE_INTERNAL_TASK_VERSION. These are read back during execution to populate ExecutionParameters.

Writing Custom Task Plugins​

The standard extension pattern for a custom task that wraps a third-party framework:

  1. Define a config type to carry plugin-specific settings
  2. Subclass PythonFunctionTask and override pre_execute, execute, and optionally get_custom
  3. Register the config type with TaskPlugins
from dataclasses import dataclass
from typing import Optional
from flytekit import PythonFunctionTask, TaskMetadata
from flytekit.core.task import TaskPlugins
from flytekit.core.context_manager import ExecutionParameters

@dataclass
class MyFrameworkConfig:
    workers: int = 4
    timeout_seconds: int = 3600

class MyFrameworkTask(PythonFunctionTask[MyFrameworkConfig]):
    def __init__(self, task_config: MyFrameworkConfig, task_function, **kwargs):
        super().__init__(
            task_config=task_config,
            task_function=task_function,
            task_type="my-framework",
            **kwargs,
        )

    def pre_execute(self, user_params: Optional[ExecutionParameters]) -> Optional[ExecutionParameters]:
        # Set up framework-specific context
        framework_client = initialize_my_framework(self.task_config.workers)
        return user_params.builder().add_attr("MY_FRAMEWORK", framework_client).build()

    def get_custom(self, settings) -> dict:
        # Return plugin-specific data included in the TaskTemplate
        return {"workers": self.task_config.workers}

# Register at module import time
TaskPlugins.register_pythontask_plugin(MyFrameworkConfig, MyFrameworkTask)

After registration, users write:

@task(task_config=MyFrameworkConfig(workers=8))
def run_training(data: str) -> float:
    ctx = flytekit.current_context()
    client = ctx.MY_FRAMEWORK  # injected by pre_execute
    return client.train(data)

For tasks that don't wrap a user function at all — where the execution logic is entirely in the plugin — subclass PythonInstanceTask instead and instantiate it as a module-level variable (not as a decorator).