Skip to main content

Defining tasks with the core task model

What you will build​

You will trace one annotated @task from definition to execution and serialization, then inspect the same task at Flyte's dispatch boundary. The path is:

Task
└── PythonTask
    └── PythonAutoContainerTask
        ├── PythonFunctionTask   (@task and callable-based extensions)
        └── PythonInstanceTask   (class-based extensions)

By the end, you will be able to identify which layer owns the Python interface, which layer invokes user code, how local calls become Promise values, and how hosted execution becomes a pyflyte-execute container command.

Prerequisites​

Use an environment in which flytekit imports successfully. The examples use these public APIs and core execution/serialization APIs:

  • task, workflow, dynamic, and PythonFunctionTask from flytekit.
  • FlyteContextManager, ExecutionState, and TypeEngine from flytekit.core.
  • SerializationSettings, Image, and ImageConfig from flytekit.configuration.
  • get_serializable from flytekit.tools.translator.

Annotated function inputs and outputs are important: PythonFunctionTask derives its native Interface from the callable's annotations.

1. Start with the task hierarchy​

Task is the IDL-oriented root. Its constructor stores the task type, name, Flyte TypedInterface, TaskMetadata, task type version, security context, and documentation, and appends the instance to FlyteEntities.entities. It does not provide a Python-native interface: its python_interface property returns None and its get_input_types() returns None.

Task also defines the call and runtime boundaries:

  • __call__ delegates to flyte_entity_call_handler.
  • local_execute() translates native values and Promise values into a LiteralMap, calls sandbox_execute(), and wraps the output literals as Promise values or a VoidPromise.
  • sandbox_execute() changes the user-space parameters to task-sandbox parameters and calls dispatch_execute().
  • dispatch_execute(), pre_execute(), and execute() are abstract at this level.
  • get_container(), get_k8s_pod(), get_sql(), get_custom(), get_config(), and get_extended_resources() are serialization hooks that return None by default.

PythonTask adds the Python-native layer. It accepts a Python Interface, transforms it to the Flyte typed interface used by Task, and retains the original interface as python_interface. It also stores task_config and environment, provides Python input/output type lookups, and implements compile() by calling create_and_link_node().

PythonTask is intended for task types that have a native Python interface but do not themselves represent a user function. Subclasses must implement execute(**kwargs); the inherited dispatch implementation supplies the conversion and lifecycle around it.

PythonAutoContainerTask adds hosted-container behavior. It accepts container_image or ImageSpec, resource settings, environment variables, a task resolver, secrets, pod-template settings, an accelerator, and shared memory. It chooses a resolver from the compilation context when one is present, otherwise from the explicit task_resolver argument or default_task_resolver.

The two concrete extension directions are:

  • PythonFunctionTask wraps a user Callable. The @task decorator produces this function-oriented kind of entity, and the class derives the name, interface, and docstring from the function.
  • PythonInstanceTask is for a task object with no user-defined function body. A subclass supplies the platform-defined execute() implementation; the base class supplies automatic container serialization, resolver handling, native interface handling, and the inherited dispatch lifecycle.

PythonFunctionTask and PythonInstanceTask are exported from flytekit. Task, PythonTask, and PythonAutoContainerTask are extension bases available from their core modules.

2. Define an annotated function task​

Use the same annotated task shape exercised by tests/flytekit/unit/core/test_array_node_map_task.py:

from typing import List

from flytekit import map_task, task, workflow


@task
def say_hello(name: str) -> str:
    return f"hello {name}!"


@workflow
def wf() -> List[str]:
    return map_task(say_hello)(name=["abc", "def"])

Calling wf() verifies that the task has a usable interface and can be placed in a workflow. In the task construction path, PythonFunctionTask calls transform_function_to_interface() with the function and its Docstring, removes any names listed in ignore_input_vars, and derives the task name with extract_task_module().

The callable is stored as task_function. In ExecutionBehavior.DEFAULT, PythonFunctionTask.execute() invokes that callable with the converted keyword arguments. The task's python_interface therefore describes name: str as an input and str as its output, while Task.interface contains the transformed Flyte typed representation.

With the default resolver, the callable must be accessible at module level. PythonFunctionTask rejects nested or local functions unless they are test functions or module-level functions wrapped in a way recognized by flytekit. A custom TaskResolverMixin can be supplied when a different loading strategy is required.

3. Follow a local call through conversion and dispatch​

At the task boundary, Task.__call__() routes the call through flyte_entity_call_handler. During local execution, Task.local_execute() performs these operations:

  1. translate_inputs_to_literals() converts native constants, Promise values, and collections containing them into the task's input LiteralMap.
  2. If task caching is configured and enabled in LocalConfig, LocalTaskCache is checked using the task name, cache version, inputs, and ignored input variables.
  3. A cache miss calls sandbox_execute(), which calls dispatch_execute().
  4. The output LiteralMap is converted into Promise values. A task with no declared outputs returns VoidPromise.

For a regular PythonFunctionTask, PythonTask.dispatch_execute() then runs the lifecycle below:

pre_execute(user_space_params)
    │
    ├─ convert input LiteralMap → native Python keyword arguments
    │
    ├─ execute(**native_inputs)
    │
    ├─ post_execute(user_params, native_outputs)
    │
    ├─ native outputs → output LiteralMap through TypeEngine
    │
    └─ write configured decks

pre_execute() runs before input conversion. Its documented use is to modify user-space parameters before type transformers run, for example when execution setup must exist before conversion. The default PythonTask.pre_execute() returns the same parameters.

execute() runs with native Python values. PythonFunctionTask.execute() calls the stored function in DEFAULT mode. PythonInstanceTask leaves the implementation to its subclass. post_execute() receives the possibly modified user parameters and the return value; the default implementation returns the value unchanged.

If the execution state is local, conversion and user exceptions are re-raised with the task name added to the error context. In non-local execution, input conversion errors are wrapped as FlyteNonRecoverableSystemException and user-code errors as FlyteUserRuntimeException. If execution returns a LiteralMap or DynamicJobSpec, dispatch returns it directly instead of attempting ordinary Python-output conversion.

The following test verifies a local dynamic call all the way to native results. It uses the public decorators and the same task behavior described above:

import typing

from flytekit import dynamic, task


@task
def t1(a: int) -> str:
    a = a + 2
    return "fast-" + str(a)


@dynamic
def ranged_int_to_str(a: int) -> typing.List[str]:
    s = []
    for i in range(a):
        s.append(t1(a=i))
    return s


res = ranged_int_to_str(a=5)
assert res == ["fast-2", "fast-3", "fast-4", "fast-5", "fast-6"]

The assertion is the visible result: local dynamic execution evaluates the generated workflow and returns native list values. At the lower task level, local_execute() still uses literal conversion and Promise wrapping; the workflow/dynamic execution layer unwraps the results for the local call.

Output shapes are part of the interface​

PythonTask._output_to_literal_map() uses the declared output names to interpret the return value:

  • Zero declared outputs produce an empty literal map.
  • One output uses the return value directly, with a special case for a one-element NamedTuple.
  • Multiple outputs are selected by position and assigned to the declared output names.
  • A tuple returned for an individual declared output raises TypeError rather than treating that tuple as another output shape.

Each value is converted asynchronously with TypeEngine.async_to_literal(). If output metadata was recorded in the execution context, dynamic partition and time-partition metadata is encoded onto the corresponding literal.

4. Exercise the runtime dispatch boundary directly​

Hosted execution enters the task through dispatch_execute() with a Flyte LiteralMap, rather than with the native dictionary used by the Python function. The test suite constructs that map using TypeEngine.dict_to_literal_map() and changes the context to TASK_EXECUTION:

import typing

from flytekit import map_task, task
from flytekit.core import context_manager
from flytekit.core.type_engine import TypeEngine


@task
def say_hello(name: str) -> str:
    return f"hello {name}!"


ctx = context_manager.FlyteContextManager.current_context()
with context_manager.FlyteContextManager.with_context(
    ctx.with_execution_state(
        ctx.execution_state.with_params(
            mode=context_manager.ExecutionState.Mode.TASK_EXECUTION
        )
    )
) as ctx:
    t = map_task(say_hello)
    lm = TypeEngine.dict_to_literal_map(
        ctx,
        {"name": ["earth", "mars"]},
        type_hints={"name": typing.List[str]},
    )
    res = t.dispatch_execute(ctx, lm)
    assert len(res.literals) == 1
    assert res.literals["o0"].scalar.primitive.string_value == "hello earth!"

This verifies the serialized side of the boundary: the task receives literals, PythonTask translates them using its Python interface, invokes the function, and returns serialized literals. Task.local_execute() is the convenience path for native values; dispatch_execute() is the common boundary used by local sandbox execution and runtime execution.

5. Understand dynamic function behavior​

PythonFunctionTask.ExecutionBehavior defines DEFAULT, DYNAMIC, and EAGER. The ordinary function path is DEFAULT; a DYNAMIC task uses dynamic_execute() instead of directly returning the callable's result. node_dependency_hints may be supplied only for DYNAMIC tasks. Supplying them for a static task raises ValueError because static dependencies can be found during compilation.

dynamic_execute() changes behavior according to the execution state:

  • During local execution, it creates or reuses a PythonFunctionWorkflow, executes that generated workflow locally, and translates its results to a literal map.
  • During TASK_EXECUTION, it calls compile_into_workflow() and returns a DynamicJobSpec containing generated nodes, task templates, outputs, and subworkflows.
  • During LOCAL_TASK_EXECUTION, it directly calls the function with the native keyword arguments.
  • Any other execution state raises ValueError.

The runtime compilation path serializes the generated workflow with get_serializable(). It rejects ReferenceTask dependencies inside a dynamic task and gathers the generated task templates into the DynamicJobSpec.

The test suite verifies the TASK_EXECUTION result with the same dynamic shape:

import typing

from flytekit import dynamic, task
from flytekit.configuration import (
    FastSerializationSettings,
    Image,
    ImageConfig,
    SerializationSettings,
)
from flytekit.core import context_manager
from flytekit.core.context_manager import ExecutionState
from flytekit.core.type_engine import TypeEngine


@task
def t1(a: int) -> str:
    a = a + 2
    return "fast-" + str(a)


@dynamic
def my_subwf(a: int) -> typing.List[str]:
    s = []
    for i in range(a):
        s.append(t1(a=i))
    return s


settings = SerializationSettings(
    project="project",
    domain="domain",
    version="version",
    image_config=ImageConfig(
        default_image=Image(name="default", fqn="test", tag="tag"),
        images=[Image(name="default", fqn="test", tag="tag")],
    ),
    fast_serialization_settings=FastSerializationSettings(
        enabled=True,
        destination_dir="/User/flyte/workflows",
        distribution_location="s3://my-s3-bucket/fast/123",
    ),
)

ctx = context_manager.FlyteContextManager.current_context()
with context_manager.FlyteContextManager.with_context(
    ctx.with_serialization_settings(settings)
) as ctx:
    with context_manager.FlyteContextManager.with_context(
        ctx.with_execution_state(
            ctx.execution_state.with_params(
                mode=ExecutionState.Mode.TASK_EXECUTION,
            )
        )
    ) as ctx:
        input_literal_map = TypeEngine.dict_to_literal_map(ctx, {"a": 5})
        dynamic_job_spec = my_subwf.dispatch_execute(ctx, input_literal_map)
        assert len(dynamic_job_spec._nodes) == 6
        assert len(dynamic_job_spec.tasks) == 1
        args = " ".join(dynamic_job_spec.tasks[0].container.args)
        assert args.startswith(
            "pyflyte-fast-execute --additional-distribution "
            "s3://my-s3-bucket/fast/123 --dest-dir /User/flyte/workflows"
        )

The dynamic function is therefore not just a special return type: the same dispatch boundary either evaluates its generated workflow locally or materializes a DynamicJobSpec when the task is executing in Flyte's task runtime.

6. Inspect hosted serialization​

PythonAutoContainerTask supplies the hosted representation inherited by both function and instance tasks. Its default command is constructed by get_default_command():

pyflyte-execute
--inputs {{.input}}
--output-prefix {{.outputPrefix}}
--raw-output-data-prefix {{.rawOutputDataPrefix}}
--checkpoint-path {{.checkpointOutputPrefix}}
--prev-checkpoint {{.prevCheckpointPrefix}}
--resolver <resolver location>
-- <resolver loader args>

get_command() returns the configured command function, which is the default unless set_command_fn() has replaced it. reset_command_fn() restores the default. The resolver's location and loader_args() are included so the hosted process can resolve the task entity and invoke it.

get_container() calls _get_container() unless a pod_template is configured. _get_container() obtains the image through get_image(), uses the ResourceSpec, adds the command arguments, and merges SerializationSettings.env with the task's environment. If an ImageSpec has runtime_packages, those packages are joined and placed in RUNTIME_PACKAGES_ENV_NAME.

A pod template changes the serialization shape: get_container() returns None, while get_k8s_pod() returns a pod containing the serialized task container and the pod labels and annotations. get_config() records the pod template's primary container name. Accelerators and shared memory are emitted through get_extended_resources().

The following setup is the exact image configuration pattern used by the serialization tests:

from flytekit.configuration import Image, ImageConfig, SerializationSettings


def serialization_settings():
    default_img = Image(name="default", fqn="test", tag="tag")
    return SerializationSettings(
        project="project",
        domain="domain",
        version="version",
        env=None,
        image_config=ImageConfig(
            default_image=default_img,
            images=[default_img],
        ),
    )

You can serialize the annotated task with the translator and inspect the resulting task template:

from collections import OrderedDict

from flytekit import task
from flytekit.tools.translator import get_serializable


@task
def t1(a: int) -> int:
    return a + 1


settings = serialization_settings()
task_spec = get_serializable(OrderedDict(), settings, t1)
assert task_spec.template.type == "python-task"
assert task_spec.template.container.args[0] == "pyflyte-execute"

The PythonAutoContainerTask layer is what turns the function task into this container-backed template; the PythonFunctionTask layer supplies the callable and resolver metadata that the container command must load.

Map tasks customize this command. In tests/flytekit/unit/core/test_array_node_map_task.py, serialization produces pyflyte-map-execute and uses ArrayNodeMapTaskResolver as the resolver. map_task() accepts PythonFunctionTask instances only in DEFAULT execution mode, or PythonInstanceTask instances; dynamic and eager function tasks are rejected.

7. Add local caching without changing the lifecycle​

Caching is task metadata, not a separate execution implementation. Task.local_execute() checks the metadata and LocalConfig before calling sandbox_execute(). A cache hit returns the stored output literal map; a miss dispatches normally and stores the resulting map.

The cache configuration and its serialized metadata are exercised like this:

from collections import OrderedDict

from flytekit import Cache, task
from flytekit.core.cache import SaltCachePolicy
from flytekit.tools.translator import get_serializable_task


@task(cache=Cache(policies=SaltCachePolicy()))
def t1(a: int) -> int:
    return a


serialized_t1 = get_serializable_task(
    OrderedDict(),
    default_serialization_settings,
    t1,
)
assert serialized_t1.template.metadata.discoverable is True
assert serialized_t1.template.metadata.discovery_version == (
    "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
)

Here default_serialization_settings is the serialization fixture from tests/flytekit/unit/core/test_cache.py. The same metadata path supports an explicit cache version and ignored input variables; the test verifies that Cache(version="a-version", ignored_inputs=("a",)) serializes those values into the task template. At local execution time, LocalConfig.cache_enabled and LocalConfig.cache_overwrite control whether the lookup is used or bypassed.

8. Choose the correct extension base​

Use PythonFunctionTask when the task's behavior is a Python callable. The @task decorator is the normal authoring path, and the class automatically derives the interface, module/name information, and docstring. Override execute() only when you are deliberately extending this behavior; the class docstring warns that an override must also handle dynamic tasks if the task is to remain a dynamic task generator.

Use PythonInstanceTask when the task is represented by an instantiated object and the execution method belongs to the task class rather than to a separate user function. Its class documentation describes the intended call shape as an instantiated task object followed by a call with arguments. A concrete subclass must implement execute() because PythonInstanceTask itself is abstract.

Use PythonTask directly for a Python-native interface whose execution is not a user function, or when building another task family. Implement execute() and override pre_execute() or post_execute() only for behavior that belongs at those lifecycle points. Use PythonAutoContainerTask as the intermediate base when the extension needs the user's code in an automatically configured container but is neither a normal function task nor an instance-task specialization.

9. Check these constraints before serialization​

  • Set _container_image before PythonAutoContainerTask calls super().__init__(). Task registration occurs during initialization, and flytekit's translator can inspect registered entities while initialization and serialization overlap.
  • Pass either resources or requests/limits, not both. PythonAutoContainerTask raises ValueError when the combined resources argument is used with either separate setting.
  • Every entry in secret_requests must be a Secret; otherwise construction raises AssertionError. Valid secrets are placed in a SecurityContext.
  • Supply exactly one of enable_deck and the deprecated disable_deck if either is specified. Decks are disabled by default; enable_deck=True enables the selected deck_fields.
  • Keep ordinary function tasks module-level when using the default resolver. Nested/local functions need the permitted test/wrapper cases or a custom resolver.
  • Treat dynamic execution state as significant. A dynamic task does not have one universal return path: local execution evaluates a generated workflow, TASK_EXECUTION returns a DynamicJobSpec, and LOCAL_TASK_EXECUTION calls the function directly.
  • Match the declared output interface. local_execute() asserts that the number of output literals matches the declared output names, and output conversion applies the single-output and multiple-output rules described above.

Complete result and next steps​

You now have one complete mental and executable path: an annotated function becomes a PythonFunctionTask; PythonAutoContainerTask supplies its resolver, image, resources, and hosted command; PythonTask performs native/literal conversion and lifecycle dispatch; and Task supplies the call boundary, local sandbox path, output promises, and generic serialization hooks.

To continue, inspect flytekit/core/array_node_map_task.py to see how a task resolver and command are replaced for mapped execution, flytekit/tools/translator.py to follow task-template serialization, and flytekit/core/dynamic_workflow_task.py for the @dynamic authoring wrapper. For a class-based extension, start from PythonInstanceTask, implement the required execute() method, and retain the inherited interface, dispatch, resolver, and container behavior.