Skip to main content

Core Task Abstractions

Every flytekit task lives two separate lives: once at serialization time (when your code is packaged and registered with Flyte Admin), and again at execution time (when a container is spun up inside a Kubernetes pod to actually run the work). Understanding how flytekit bridges these two moments — and which class is responsible for what — makes the whole system much easier to reason about.

The Two-Phase Task Lifecycle

When you register a workflow, flytekit serializes each task into a TaskTemplate protobuf. Part of that serialization is baking a CLI command into the container spec. When Flyte later executes the task, it launches a container that runs this command:

pyflyte-execute \
--inputs s3://path/inputs.pb \
--output-prefix s3://outputs/location \
--raw-output-data-prefix /tmp/data \
--resolver flytekit.core.python_auto_container.default_task_resolver \
-- \
task-module myapp.workflows.pipeline task-name my_task

The --resolver argument is a dotted import path to a resolver singleton. The -- separator is followed by arguments that the resolver interprets to reconstruct the task object. At runtime, _execute_task in flytekit/bin/entrypoint.py loads the resolver via importlib and calls its load_task() method:

# flytekit/bin/entrypoint.py
resolver_obj = load_object_from_module(resolver)

def load_task():
return resolver_obj.load_task(loader_args=resolver_args)

_dispatch_execute(ctx, load_task, inputs, output_prefix)

Every class and abstraction described in this section exists to make this serialization + rehydration cycle reliable.


TaskMetadata: Controlling Task Execution Semantics

TaskMetadata (in flytekit/core/base_task.py) is a Python dataclass that captures the execution-level configuration for a task. It maps directly to the FlyteIDL TaskMetadata protobuf via to_taskmetadata_model().

from flytekit import task, TaskMetadata

@task(
cache=True,
cache_version="v2",
retries=3,
timeout=600, # int seconds, auto-converted to timedelta
interruptible=True,
)
def my_task(n: int) -> int:
return n * 2

The @task decorator constructs a TaskMetadata from these arguments and passes it to the underlying task plugin:

# flytekit/core/task.py
_metadata = TaskMetadata(
cache=cache,
cache_serialize=cache_serialize,
cache_version=cache_version,
cache_ignore_input_vars=cache_ignore_input_vars,
retries=retries,
interruptible=interruptible,
deprecated=deprecated,
timeout=timeout,
)

Caching Fields

All four caching-related fields are validated together in __post_init__:

FieldTypeDefaultNotes
cacheboolFalseEnables result caching
cache_versionstr""Required when cache=True
cache_serializeboolFalseSerializes same-input executions; requires cache=True
cache_ignore_input_varsTuple[str, ...]()Excludes named inputs from cache key; requires cache=True

Violating any of these rules raises a ValueError at construction time — not at runtime:

# Raises: "Caching is enabled cache=True but cache_version is not set."
TaskMetadata(cache=True)

# Raises: "Cache serialize is enabled cache_serialize=True but cache is not enabled."
TaskMetadata(cache_serialize=True)

# Raises: "Cache ignore input vars are specified ... but cache is not enabled."
TaskMetadata(cache_ignore_input_vars=("timestamp",))

Retries and Timeouts

retries sets the number of times Flyte will retry the task on failure. timeout accepts either an int (seconds, silently converted to datetime.timedelta) or a timedelta directly. Any other type raises a ValueError.

import datetime
from flytekit import TaskMetadata

# These are equivalent:
TaskMetadata(retries=2, timeout=300)
TaskMetadata(retries=2, timeout=datetime.timedelta(seconds=300))

Other Metadata Fields

  • interruptible: When True, Flyte may schedule this task on preemptible nodes. None defers to the platform default.
  • deprecated: Non-empty string attaches a deprecation warning to the task in the Admin UI.
  • pod_template_name: Names an existing PodTemplate resource in the Kubernetes cluster. In PythonAutoContainerTask, the constructor argument pod_template_name always overwrites whatever pod_template_name was set on the passed TaskMetadata object — even if TaskMetadata.pod_template_name was already set.
  • generates_deck: Set automatically when deck output is enabled.
  • is_eager: Marks the task for eager execution semantics.

The Task Class Hierarchy

Flytekit organizes tasks into three layers, each adding progressively more Python-specific functionality:

Task                         (flytekit/core/base_task.py)
└── PythonTask (flytekit/core/base_task.py)
└── PythonAutoContainerTask (flytekit/core/python_auto_container.py)
├── PythonFunctionTask
├── PythonInstanceTask
├── ContainerTask
└── (all user-visible task types)

Task works directly with FlyteIDL types (TypedInterface, LiteralMap) and knows nothing about Python's type system. It provides the local_execute → sandbox_execute → dispatch_execute execution chain and handles local caching via LocalTaskCache. Every Task instance appends itself to FlyteEntities.entities at init time, which is how the serialization tools discover all tasks in a module.

PythonTask adds a Python-native Interface (inputs and outputs typed with actual Python types), the task_config generic parameter, environment variable support, and the full deck reporting system. It also mixes in TrackedInstance for automatic name resolution.

PythonAutoContainerTask adds everything needed to deploy a task as a container: image selection, resource requests/limits, pod templates, secret injection, accelerator configuration, and — crucially — the task resolver that makes runtime rehydration possible.

Most application developers only interact with the @task decorator, but plugin authors target this hierarchy directly.


PythonTask: The Execution Lifecycle

When a task runs, PythonTask.dispatch_execute orchestrates the full execution:

  1. pre_execute(user_params) — called before inputs are deserialized. Plugins override this to set up external resources (e.g., a Spark session) before type transformers run. The default implementation is a no-op that returns user_params unchanged.

  2. Input conversion_literal_map_to_python_input calls TypeEngine.literal_map_to_kwargs to convert the LiteralMap from Flyte's type system into Python native values.

  3. execute(**kwargs) — the abstract method that subclasses implement. This is where user code runs.

  4. post_execute(user_params, rval) — called with the return value from execute. Plugins can use this for cleanup or to transform outputs. The default is a no-op.

  5. Output conversion_output_to_literal_map converts Python native return values back into a LiteralMap.

  6. Deck writing_write_decks outputs HTML for the configured deck fields.

The IgnoreOutputs Pattern

In distributed training scenarios, only one worker (typically rank 0) produces outputs. All other workers should execute but not write results. The IgnoreOutputs exception (also in flytekit/core/base_task.py) handles this:

from flytekit import task
from flytekit.core.base_task import IgnoreOutputs

@task
def distributed_worker(rank: int, data: str) -> str:
result = do_work(data)
if rank != 0:
raise IgnoreOutputs() # non-primary workers skip output writing
return result

When IgnoreOutputs propagates through dispatch_execute to the entrypoint's _dispatch_execute, the entrypoint catches it and skips writing outputs.pb. The PyTorch plugin (flytekitplugins/kfpytorch/task.py) uses this pattern for all non-rank-0 workers.

Deck Configuration

Decks are controlled via enable_deck and deck_fields:

from flytekit import task
from flytekit.deck import DeckField

@task(
enable_deck=True,
deck_fields=[DeckField.INPUT, DeckField.OUTPUT, DeckField.TIMELINE],
)
def documented_task(x: int) -> int:
return x + 1

disable_deck was deprecated in flytekit 1.10.0. Setting both disable_deck and enable_deck raises a ValueError. When enable_deck=True is not set, _disable_deck defaults to True and _deck_fields is an empty list, so no deck HTML is produced.


PythonAutoContainerTask: Container Configuration

PythonAutoContainerTask (in flytekit/core/python_auto_container.py) is an abstract class that adds container-deployment concerns on top of PythonTask. You typically reach it through PythonFunctionTask (via @task) or PythonInstanceTask, but plugin authors subclass it directly.

Resource Configuration

Two approaches exist for specifying resources, and they cannot be mixed:

from flytekit import task, Resources

# Approach 1: separate requests and limits
@task(requests=Resources(cpu="1", mem="2Gi"), limits=Resources(cpu="2", mem="4Gi"))
def task_a(): ...

# Approach 2: combined resources (request and limit together)
@task(resources=Resources(cpu=("1", "2"), mem="2Gi"))
def task_b(): ...

# This raises ValueError:
# `resource can not be used together with the limits or requests`
@task(resources=Resources(cpu="1"), limits=Resources(cpu="2"))
def task_c(): ...

Resources instances flow into a ResourceSpec (holding both requests and limits) stored on _resources.

Container Image Selection

The container_image parameter accepts a string, a template, or an ImageSpec:

from flytekit import task
from flytekit.image_spec import ImageSpec

# Explicit string — used as-is
@task(container_image="ghcr.io/myorg/myimage:v1.2.3")
def task_a(): ...

# Template — interpolated from ImageConfig at serialization time
@task(container_image="{{.image.default.fqn}}:{{.image.default.version}}")
def task_b(): ...

# ImageSpec — built or looked up from registry
my_image = ImageSpec(packages=["pandas", "scikit-learn"])

@task(container_image=my_image)
def task_c(): ...

get_image() resolves the final image string via get_registerable_container_image(), which handles all three cases against the ImageConfig from SerializationSettings. If no image is configured and no default is set, it raises a ValueError.

Pod Templates

When a pod_template is passed, the container spec is embedded inside a full Kubernetes pod spec:

from flytekit import task
from flytekit.core.pod_template import PodTemplate

@task(pod_template=PodTemplate(primary_container_name="main"))
def pod_task(): ...

With pod_template set, get_container() returns None and get_k8s_pod() returns the merged K8sPod. Serialization tools must check get_k8s_pod() in this case. get_config() returns {"primary_container_name": ...} to tell the Flyte agent which container carries the task.

The pod_template_name constructor argument (referencing an existing cluster-side PodTemplate resource) always overwrites metadata.pod_template_name, even if TaskMetadata was constructed with a different value:

# metadata.pod_template_name ends up as "podTemplateA", not "podTemplateB"
task_with_pod = DummyAutoContainerTask(
name="x",
metadata=TaskMetadata(pod_template_name="podTemplateB"),
pod_template_name="podTemplateA", # wins
...
)
assert task_with_pod.metadata.pod_template_name == "podTemplateA"

Secrets, Accelerators, and Shared Memory

from flytekit import task, Secret
from flytekit.extras.accelerators import T4

@task(
secret_requests=[Secret(group="my-secrets", key="api-key")],
accelerator=T4,
shared_memory="2Gi",
)
def gpu_task(): ...

Secrets are validated to be flytekit.Secret instances and collected into a SecurityContext.


TaskResolverMixin and Task Rehydration

TaskResolverMixin (in flytekit/core/base_task.py) is the interface that connects serialization time to execution time. It has four abstract methods:

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

@abstractmethod
def loader_args(self, settings, t) -> List[str]: ... # called at serialization time

@abstractmethod
def load_task(self, loader_args) -> Task: ... # called at execution time

@abstractmethod
def get_all_tasks(self) -> List[Task]: ...

loader_args is called once when the task's container spec is built. The returned list is embedded after the -- separator in the container command. load_task receives that same list when the container starts and must reconstruct the exact task object.

DefaultTaskResolver

DefaultTaskResolver (singleton: default_task_resolver) handles the standard case of a module-level task function:

# loader_args output for a task defined in myapp.workflows.pipeline named my_task:
# ["task-module", "myapp.workflows.pipeline", "task-name", "my_task"]

# load_task implementation:
def load_task(self, loader_args):
_, task_module, _, task_name, *_ = loader_args
task_module = importlib.import_module(name=task_module)
return getattr(task_module, task_name)

This is why nested functions cannot be used as tasks with the default resolverimportlib.import_module can import the module, but getattr cannot retrieve a function defined inside another function. Flytekit enforces this at construction time via the isnested() check in PythonFunctionTask. Two exceptions exist: functions defined in test modules (prefixed with test_ or suffixed with _test) and functions decorated with functools.wraps at module level (detected via is_functools_wrapped_module_level()).

The command embedded in the container spec looks like this (from get_default_command):

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),
]

ClassStorageTaskResolver

For tasks defined inline inside a dynamic workflow, flytekit uses ClassStorageTaskResolver (in flytekit/core/class_based_resolver.py). It stores tasks in a list and identifies them by index:

class ClassStorageTaskResolver(TrackedInstance, TaskResolverMixin):
def __init__(self):
self.mapping = []

def add(self, t: PythonAutoContainerTask):
self.mapping.append(t)

def load_task(self, loader_args: List[str]) -> PythonAutoContainerTask:
idx = int(loader_args[0])
return self.mapping[idx]

def loader_args(self, settings, t) -> List[str]:
return [f"{self.mapping.index(t)}"]

When a task is defined inside a @dynamic workflow, the compilation context sets compilation_state.task_resolver to the workflow's ClassStorageTaskResolver. Any task_resolver passed explicitly to PythonAutoContainerTask.__init__ is silently ignored in this case — the context wins.

DefaultNotebookTaskResolver

In Jupyter notebook and interactive environments where module import is unavailable, DefaultNotebookTaskResolver reconstructs tasks from a compressed cloudpickle file at PICKLE_FILE_PATH. It wraps tasks in a PickledEntity dataclass:

from flytekit.core.python_auto_container import (
PickledEntity,
PickledEntityMetadata,
default_notebook_task_resolver,
)

# loader_args identifies tasks by their fully-qualified name:
# ["entity-name", "mymodule.my_task"]

At load_task time, DefaultNotebookTaskResolver reads the gzip-compressed cloudpickle file, validates that the Python major and minor version match exactly, and returns the named entity. A mismatch raises:

RuntimeError: The Python version used to create the pickle file is different from the current Python version.
Current Python version: 3.11. Python version used to create the pickle file: 3.12.0.

PythonAutoContainerTask.set_resolver() is how you attach this resolver to a task after construction:

task.set_resolver(default_notebook_task_resolver)

Map tasks use the complementary set_command_fn() to switch from pyflyte-execute to pyflyte-map-execute. reset_command_fn() restores the default.


TrackedInstance and Automatic Naming

TrackedInstance (in flytekit/core/tracker.py) uses the InstanceTrackingMeta metaclass to capture where each instance is created. The metaclass walks the call stack during __call__ to find the innermost module-level frame (skipping __main__) and records both the module name and the file path:

# InstanceTrackingMeta.__call__ sets these:
o._instantiated_in = mod_name # e.g., "myapp.workflows.pipeline"
o._module_file = Path(mod_file).resolve()

TrackedInstance.find_lhs() then searches the module's namespace — iterating dir(m) and comparing by identity — to find which variable name the instance was assigned to:

# If you write:
my_resolver = MyCustomResolver()
# then:
assert my_resolver.lhs == "my_resolver"
assert my_resolver.instantiated_in == "myapp.resolvers"

The location property calls extract_task_module(self) to produce the fully-qualified dotted path (e.g., "myapp.resolvers.my_resolver"), which is what gets embedded in the container's --resolver argument.

TrackedInstance serves two distinct roles:

  1. Task resolver singletons — because resolvers are instances (not classes), Flyte must be able to locate them by their dotted path at runtime. location provides this.
  2. Non-function PythonAutoContainerTasks — tasks like SQLite3Task or SQLTask don't get their names from a function's __name__, so TrackedInstance provides the naming mechanism instead.

extract_task_module() handles several edge cases: __main__ module (scripts run directly), Jupyter notebooks (where module resolution returns an empty string), and functools.wraps-decorated functions.

FLYTE_PYTHON_PACKAGE_ROOT

By default (auto), extract_task_module walks up the directory tree looking for the root Python package (the directory lacking an __init__.py). You can override this via the FLYTE_PYTHON_PACKAGE_ROOT configuration:

  • auto — detect package root by scanning for __init__.py (default)
  • . — treat the current working directory as the root
  • /path/to/root — use the given path as the absolute package root

kwtypes: Typed Interfaces Without Python Functions

When a task doesn't wrap a Python function, there's no function signature to inspect for type information. kwtypes (in flytekit/core/base_task.py) builds an OrderedDict[str, Type] from keyword arguments to fill this gap:

from flytekit import kwtypes
from flytekit.types.schema import FlyteSchema
import datetime

sql = SQLTask(
"my-query",
query_template="SELECT * FROM events WHERE ds = '{{ .Inputs.ds }}' LIMIT 10",
inputs=kwtypes(ds=datetime.datetime),
outputs=kwtypes(results=FlyteSchema),
metadata=TaskMetadata(retries=2, cache=True, cache_version="0.1"),
)

kwtypes is simply:

def kwtypes(**kwargs) -> OrderedDict[str, Type]:
d = collections.OrderedDict()
for k, v in kwargs.items():
d[k] = v
return d

ContainerTask, SQLTask, and any other task that defines its interface declaratively rather than by function introspection uses this pattern.


Extension Patterns

Extending PythonAutoContainerTask

The minimal pattern for a custom container task: subclass PythonAutoContainerTask, implement execute, and pass the name and type:

from typing import Any
from flytekit.core.python_auto_container import PythonAutoContainerTask
from flytekit.core.interface import Interface

class MyPluginTask(PythonAutoContainerTask):
def __init__(self, name: str, config: MyConfig, **kwargs):
super().__init__(
name=name,
task_config=config,
task_type="my-plugin",
interface=Interface(inputs={"x": int}, outputs={"result": str}),
**kwargs,
)

def execute(self, **kwargs) -> Any:
x = kwargs["x"]
return {"result": str(self.task_config.do_work(x))}

For tasks that wrap a Python function (where the function itself defines the interface), use PythonFunctionTask as the base or, more commonly, register a plugin with the @task decorator via TaskPlugins.register_pythontask_plugin(MyConfig, MyPluginTask).

Implementing a Custom TaskResolverMixin

A custom resolver needs loader_args + load_task to be inverses of each other, and a location that can be imported at runtime. The TrackedInstance base provides location automatically:

from typing import List
from flytekit.core.base_task import TaskResolverMixin, Task
from flytekit.core.tracker import TrackedInstance
from flytekit.configuration import SerializationSettings


class MyResolver(TrackedInstance, TaskResolverMixin):
def name(self) -> str:
return "MyResolver"

def loader_args(self, settings: SerializationSettings, t: Task) -> List[str]:
# embed whatever is needed to find this task later
return ["task-key", t.name]

def load_task(self, loader_args: List[str]) -> Task:
_, task_key = loader_args
return MY_TASK_REGISTRY[task_key]

def get_all_tasks(self) -> List[Task]:
return list(MY_TASK_REGISTRY.values())


# The singleton instance — its variable name becomes part of `location`
my_resolver = MyResolver()

Because my_resolver is assigned at module level, TrackedInstance will find lhs == "my_resolver" and location == "mypackage.resolvers.my_resolver". That string is what Flyte embeds in --resolver and what load_object_from_module loads at container startup.


Configuration Reference

Environment VariableDefaultEffect
FLYTE_PYTHON_PACKAGE_ROOTautoControls module path resolution. auto = scan for __init__.py. . = current directory is root. Path = explicit root.
FLYTE_LOCAL_CACHE_ENABLEDtrueSet to false to disable local task caching during local_execute. Tasks with cache=True run on every invocation.
FLYTE_LOCAL_CACHE_OVERWRITEfalseSet to true to force overwrite of local cache entries, re-executing even on cache hits.
FLYTE_INTERNAL_IMAGEOverrides the default container image for tasks that don't specify container_image.
_F_RUNTIME_PACKAGESSpace-separated packages to pip-install at container startup. Set automatically when container_image is an ImageSpec with runtime_packages.
FLYTE_FAIL_ON_ERRORfalseWhen true, the container exits with a non-zero status code if an error file is written. Useful for backends like Databricks or AWS Batch that don't read error.pb.
SERIALIZED_CONTEXT_ENV_VARCarries compressed SerializationSettings from the workflow to the container at execution time. Consumed by setup_execution() in entrypoint.py.