Skip to main content

Tasks: Base Classes and Execution Model

Flytekit's task system is built on a small hierarchy of base classes that handle everything from IDL serialization to local execution and distributed training coordination. Most users interact only with the @task decorator, but anyone extending flytekit — writing a new plugin, a custom task type, or a non-standard resolver — works directly with these base abstractions.

The inheritance chain looks like this:

Task  (flytekit/core/base_task.py)
  └── PythonTask  [+ TrackedInstance, Generic[T]]
        └── PythonAutoContainerTask  (flytekit/core/python_auto_container.py)
              ├── PythonFunctionTask         ← what @task returns for sync functions
              └── AsyncPythonFunctionTask    ← what @task returns for async functions

The Task Base Class​

Task, defined in flytekit/core/base_task.py, is the lowest-level abstraction. Its docstring describes it as "closest to the FlyteIDL TaskTemplate" — it captures the structural information that Flyte's control plane needs without assuming the task is Python-native.

Every Task constructor accepts:

def __init__(
    self,
    task_type: str,         # e.g., "python-task", "spark", "raw-container"
    name: str,              # unique fully-qualified name, e.g., "mywf.module.my_task"
    interface: _interface_models.TypedInterface,  # IDL-level typed interface
    metadata: Optional[TaskMetadata] = None,
    task_type_version=0,
    security_ctx: Optional[SecurityContext] = None,
    docs: Optional[Documentation] = None,
    **kwargs,
)

One side effect of construction is automatic global registration:

FlyteEntities.entities.append(self)

FlyteEntities.entities is the global list that the serialization toolchain scans when generating workflow specs. Constructing a task — even outside a workflow body — adds it to this list. This is what makes the @task decorator able to work at module import time.

Abstract Methods​

Three methods are abstract and must be implemented by every concrete task:

  • execute(**kwargs) -> Any — runs the actual task logic, receiving Python-native inputs.
  • dispatch_execute(ctx, input_literal_map) -> LiteralMap — translates Flyte's type-system inputs into Python values, calls execute, and converts outputs back to a LiteralMap. This is called both locally and at runtime inside a container.
  • pre_execute(user_params) -> ExecutionParameters — runs before any inputs are converted, giving subclasses a chance to set up execution context (e.g. initializing a SparkSession before type transformers run).

Serialization Hooks​

Task also defines a family of optional serialization hooks that subclasses override to describe how they run on the Flyte platform:

MethodPurpose
get_container(settings)Returns image + command for container-based tasks
get_k8s_pod(settings)Returns a Kubernetes pod spec for pod-based tasks
get_sql(settings)Returns SQL definition for SQL tasks
get_custom(settings)Plugin-specific data serialized into the TaskTemplate
get_config(settings)Additional string key-value config
get_extended_resources(settings)GPU accelerators, shared memory, etc.

All of these return None by default. A task that doesn't override get_container produces a TaskTemplate with no container target — appropriate for tasks like Athena SQL or Databricks, whose execution target is specified differently.

Local vs. Remote Execution​

When a task is called inside a local Python process (either in a test or via pyflyte run), the flow goes through local_execute:

task(x=1)
  └── __call__ → flyte_entity_call_handler
        └── local_execute(ctx, x=1)
              ├── translate_inputs_to_literals(...)          # Python → LiteralMap
              ├── [optional] LocalTaskCache.get(...)         # check local cache
              ├── sandbox_execute(ctx, input_literal_map)    # sets up sandbox context
              │     └── dispatch_execute(ctx, input_literal_map)
              └── [output_names] → [Promise, ...]           # LiteralMap → Promises

sandbox_execute creates a temporary working directory for the task before handing off to dispatch_execute. It is not called during remote runtime execution — only the dispatch_execute call happens there.

TaskMetadata​

TaskMetadata is a @dataclass that groups the behavioral knobs for a task: caching, retries, timeouts, interruptibility, and more.

@dataclass
class TaskMetadata(object):
    cache: bool = False
    cache_serialize: bool = False
    cache_version: str = ""
    cache_ignore_input_vars: Tuple[str, ...] = ()
    interruptible: Optional[bool] = None
    deprecated: str = ""
    retries: int = 0
    timeout: Optional[Union[datetime.timedelta, int]] = None
    pod_template_name: Optional[str] = None
    generates_deck: bool = False
    is_eager: bool = False

Validation​

__post_init__ enforces several rules that surface as ValueError at decoration time, not at runtime:

# cache=True requires cache_version to be set
@task(cache=True)                        # ← ValueError: cache_version not set
def missing_version(i: str): ...

# cache_serialize requires cache=True
@task(cache_serialize=True)             # ← ValueError: cache is not enabled
def serialize_without_cache(i: str): ...

# cache_ignore_input_vars requires cache=True
@task(cache_ignore_input_vars=("ts",))  # ← ValueError: cache is not enabled
def ignore_without_cache(ts: str): ...

# This is valid:
@task(cache=True, cache_serialize=True, cache_version="1.0")
def well_configured(i: str): ...

Integer timeouts are automatically converted to datetime.timedelta in __post_init__:

metadata = TaskMetadata(timeout=300)
# metadata.timeout is now datetime.timedelta(seconds=300)

Caching Fields​

cache_ignore_input_vars is a tuple of input variable names that should be excluded when computing the cache key. This is useful for timestamp or run-ID parameters that vary but don't affect the output:

@task(cache=True, cache_version="v1", cache_ignore_input_vars=("run_timestamp",))
def featurize(dataset: str, run_timestamp: str) -> pd.DataFrame:
    ...

IDL Serialization​

to_taskmetadata_model() converts the dataclass to the IDL protobuf-backed _task_model.TaskMetadata. It attaches the current flytekit version into the RuntimeMetadata:

def to_taskmetadata_model(self) -> _task_model.TaskMetadata:
    from flytekit import __version__
    return _task_model.TaskMetadata(
        discoverable=self.cache,
        runtime=_task_model.RuntimeMetadata(
            _task_model.RuntimeMetadata.RuntimeType.FLYTE_SDK, __version__, "python"
        ),
        timeout=self.timeout,
        retries=self.retry_strategy,
        ...
    )

Interaction with PythonAutoContainerTask​

When you pass pod_template_name directly to a task constructor (or via @task(pod_template_name=...)), it overwrites any pod_template_name set inside a metadata=TaskMetadata(...) argument. This is consistent behavior across ContainerTask and PythonAutoContainerTask:

task_with_pod_template = DummyAutoContainerTask(
    name="x",
    metadata=TaskMetadata(
        pod_template_name="podTemplateB",  # will be overwritten
        retries=3,
    ),
    task_config=None,
    task_type="t",
    pod_template_name="podTemplateA",  # this wins
)

PythonTask​

PythonTask extends both TrackedInstance and Task and is generic over T (the task config type). It's the foundation for all tasks that have a Python-native typed interface — meaning inputs and outputs described as Python type annotations rather than raw IDL types.

Two Interface Objects​

Every PythonTask carries two representations of its interface:

  • self._python_interface (Interface) — stores Python-native types, used by the type engine for serialization and deserialization.
  • Inherited self._interface (TypedInterface) — the IDL-level representation, produced by transform_interface_to_typed_interface(interface) at construction.

The python_interface property exposes the former; the interface property (inherited from Task) exposes the latter. Type dispatch in dispatch_execute uses python_interface.inputs to know how to convert each literal.

Task Config​

The generic T parameter represents a plugin-specific configuration object surfaced through task_config:

class SparkTask(PythonTask[Spark]):
    def __init__(self, ..., task_config: Spark, ...):
        super().__init__(task_type="spark", ..., task_config=task_config)

    def execute(self, **kwargs) -> Any:
        spark = flytekit.current_context().spark_session
        ...

The task config is passed through to get_custom() / get_config() by plugin authors to describe backend execution requirements.

The dispatch_execute Flow​

When a task runs — locally or in a container — PythonTask.dispatch_execute orchestrates the entire lifecycle:

dispatch_execute(ctx, input_literal_map)
  │
  ├── 1. pre_execute(ctx.user_space_params)
  │       → returns (possibly modified) ExecutionParameters
  │
  ├── 2. _literal_map_to_python_input(input_literal_map, exec_ctx)
  │       → TypeEngine.literal_map_to_kwargs(...)
  │       → produces Dict[str, Python_native_value]
  │
  ├── 3. execute(**native_inputs)
  │       → user code runs here
  │       → returns Python-native outputs
  │
  ├── 4. post_execute(new_user_params, native_outputs)
  │       → no-op by default; may raise IgnoreOutputs
  │
  └── 5. _output_to_literal_map(native_outputs, exec_ctx)  [async]
          → TypeEngine.async_to_literal(ctx, v, py_type, literal_type) per output
          → returns LiteralMap

Errors in steps 2 and 5 are categorized differently based on execution context. During local execution, the original exception is re-raised with a descriptive message. During remote (container) execution, input conversion failures become FlyteNonRecoverableSystemException and execute() failures become FlyteUserRuntimeException.

Pre/Post Execute Hooks​

Plugin authors override pre_execute to set up resources before inputs are converted — for instance, SparkTask initializes SparkSession in pre_execute so it's available when execute runs. The method receives ExecutionParameters and must return ExecutionParameters (either the same object or a modified copy).

post_execute receives the raw return value of execute and can transform or suppress it. It's a no-op in the base class:

def post_execute(self, user_params: Optional[ExecutionParameters], rval: Any) -> Any:
    return rval

Raising IgnoreOutputs from post_execute suppresses all output uploads — see the section below.

Deck Configuration​

By default, decks are disabled (_disable_deck = True). Pass enable_deck=True to enable the full set of automatically generated views:

@task(enable_deck=True)
def my_task(df: pd.DataFrame) -> pd.DataFrame:
    ...

The deck_fields parameter controls which panels are generated, using values from the DeckField enum (SOURCE_CODE, DEPENDENCIES, TIMELINE, INPUT, OUTPUT). Setting enable_deck=True without customizing deck_fields uses the full default set. The old disable_deck parameter is deprecated since flytekit 1.10.0; mixing both raises ValueError.

TaskResolverMixin and Task Rehydration​

When Flyte's control plane schedules a task for execution, it starts your container with specific arguments. The container doesn't know in advance which task it's supposed to run. TaskResolverMixin, defined in flytekit/core/base_task.py, is the protocol that bridges serialization time and execution time.

The contract has two sides:

At serialization time (loader_args): given a task object and serialization settings, produce a list of strings that identify the task.

At execution time (load_task): given those same strings, reconstruct the task object.

The Default Resolver​

DefaultTaskResolver (in flytekit/core/python_auto_container.py) handles the standard case: it identifies a task by its module path and variable name.

class DefaultTaskResolver(TrackedInstance, TaskResolverMixin):
    def load_task(self, loader_args: List[str]) -> PythonAutoContainerTask:
        _, task_module, _, task_name, *_ = loader_args
        task_module = importlib.import_module(name=task_module)
        task_def = getattr(task_module, task_name)
        return task_def

    def loader_args(self, settings: SerializationSettings, task: PythonAutoContainerTask) -> List[str]:
        _, m, t, _ = extract_task_module(task)
        return ["task-module", m, "task-name", t]

default_task_resolver = DefaultTaskResolver()

A simple task defined at module level serializes to a container command like:

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 repo_root.workflows.example task-name t1

The --resolver argument is the importable dotted path to the resolver instance (not the class). The -- separator marks the start of the loader_args that get passed to load_task.

Nested Tasks and ClassStorageTaskResolver​

Tasks defined inside a workflow function body can't be found by module import. ClassStorageTaskResolver (in flytekit/core/class_based_resolver.py) handles this by maintaining an ordered list and using the integer index as the loader argument:

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

    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: SerializationSettings, t: PythonAutoContainerTask) -> List[str]:
        if t not in self.mapping:
            raise ValueError("no such task")
        return [f"{self.mapping.index(t)}"]

Nested tasks serialized this way use the workflow object itself as the resolver. The container args end with --resolver tests.my_module.my_wf -- 0 (for the first nested task) and --resolver tests.my_module.my_wf -- 1 (for the second), as verified by the resolver tests:

srz_t0_spec = get_serializable(OrderedDict(), serialization_settings, workflows_tasks[0])
assert srz_t0_spec.template.container.args[-4:] == [
    "--resolver",
    "tests.flytekit.unit.core.test_resolver.my_wf",
    "--",
    "0",
]

The location Property​

The location property of a TaskResolverMixin must return an importable dotted path to the resolver instance — not the class. At runtime, _execute_task in flytekit/bin/entrypoint.py loads it like this:

resolver_obj = load_object_from_module(resolver)

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

load_object_from_module splits the dotted path into a module path and an attribute name. This means the resolver instance must be a module-level attribute or reachable as one. Tasks that are inner functions in test files are exempt from this constraint because test modules are not serialized and deployed.

Implementing a Custom Resolver​

Override all five abstract methods of TaskResolverMixin (location, name, load_task, loader_args, get_all_tasks) and pass your resolver instance via the task_resolver= argument of PythonAutoContainerTask (or via @task(task_resolver=...)).

IgnoreOutputs​

In distributed training scenarios, multiple workers execute the same task function simultaneously but only one (rank 0) should produce outputs. IgnoreOutputs is the signal for this:

class IgnoreOutputs(Exception):
    """
    This exception should be used to indicate that the outputs generated by this can be safely ignored.
    This is useful in case of distributed training or peer-to-peer parallel algorithms.
    """
    pass

Raise it from execute() or post_execute() when no output should be written. The PyTorch elastic task plugin demonstrates the pattern:

# In ElasticWorkerResult processing (kfpytorch/task.py)
except SignalException as e:
    logger.exception(f"Elastic launch agent process terminating: {e}")
    raise IgnoreOutputs()

if 0 in out:
    return out[0].return_value
else:
    raise IgnoreOutputs()  # non-rank-0 workers don't produce outputs

IgnoreOutputs raised from within execute() is caught by PythonTask.dispatch_execute as part of post_execute processing and propagates as a FlyteUserRuntimeException. The entrypoint (flytekit/bin/entrypoint.py) then catches it and returns early without writing any outputs.pb:

except FlyteUserRuntimeException as e:
    if isinstance(e.value, IgnoreOutputs):
        logger.warning(f"User-scoped IgnoreOutputs received! Outputs.pb will not be uploaded. reason {e}!!")
        return

IgnoreOutputs must be raised, not returned. It only suppresses output upload when it propagates through FlyteUserRuntimeException — which is how dispatch_execute wraps errors from execute() during remote execution.

kwtypes​

kwtypes is a small utility function in flytekit/core/base_task.py that builds an OrderedDict of types from keyword arguments. It's used to define typed interfaces for tasks that don't derive their interface from Python function annotations:

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

The most common use is with ContainerTask inputs and outputs:

from flytekit import kwtypes, ContainerTask

container_task = ContainerTask(
    name="raw.container.count",
    image="alpine",
    command=["/bin/sh", "-c", "echo 5 > /var/flyte/outputs/output_files/o0"],
    inputs=kwtypes(path=str),
    outputs=kwtypes(count=int),
)

Local Execution and Caching​

During local execution, Task.local_execute manages a lightweight caching layer backed by LocalTaskCache (stored at ~/.flyte/local-cache). When metadata.cache=True is set, the cache is keyed on the task name, cache version, and a hash of the input literals (minus any cache_ignore_input_vars):

@task(cache=True, cache_version="v1")
def is_even(n: int) -> bool:
    return n % 2 == 0

@workflow
def check_evenness(n: int) -> bool:
    return is_even(n=n)

check_evenness(n=1)  # executes is_even
check_evenness(n=1)  # cache hit — is_even is not called again
check_evenness(n=8)  # different input, cache miss — executes is_even

Two environment variables control local caching behavior:

  • FLYTE_LOCAL_CACHE_ENABLED (default: true) — set to false to disable local caching even for tasks with cache=True.
  • FLYTE_LOCAL_CACHE_OVERWRITE (default: false) — set to true to force re-execution and overwrite existing cache entries.
def test_cache_can_be_disabled(monkeypatch):
    monkeypatch.setenv("FLYTE_LOCAL_CACHE_ENABLED", "false")
    # tasks with cache=True will still execute every time

Cache misses execute via sandbox_execute, which wraps dispatch_execute in a temporary working directory context. Cache hits short-circuit directly to output Promise construction, bypassing execution entirely.

Runtime Entrypoint​

At execution time on the Flyte platform, the container runs pyflyte-execute, which calls _execute_task in flytekit/bin/entrypoint.py. This function parses --resolver and -- arguments from the container command and uses them to load and dispatch the task:

def _execute_task(inputs, output_prefix, test, raw_output_data_prefix,
                  resolver, resolver_args, ...):
    with setup_execution(raw_output_data_prefix, output_prefix, ...) as ctx:
        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)

_dispatch_execute calls load_task() to reconstruct the task object, then calls its dispatch_execute method with the deserialized input LiteralMap. Outputs are written to the output_prefix location. When IgnoreOutputs propagates up through FlyteUserRuntimeException, _dispatch_execute returns without writing any output file.

Flyte Propeller injects several environment variables into the container at execution time, including FLYTE_INTERNAL_EXECUTION_ID, FLYTE_INTERNAL_TASK_NAME, and _F_SS_C (compressed SerializationSettings). These are consumed by setup_execution to reconstruct the execution context before dispatch_execute is called.