Skip to main content

Task Base Classes and Execution Model

Task metadata is the execution policy​

A task can have the right inputs and still be unsuitable for execution if its retry, timeout, cache, or scheduling policy is incomplete. In flytekit, TaskMetadata is the dataclass that carries these per-task directives. You pass it as metadata= when constructing a task; the @task decorator also creates one internally from its task options.

sql = SQLTask(
    "my-query",
    query_template="SELECT * FROM hive.city.fact_airport_sessions 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"),
)

TaskMetadata defines these fields:

  • cache enables caching; cache_version identifies the cached value, and cache_ignore_input_vars names inputs omitted from the cache-key calculation.
  • cache_serialize requests serialized execution of identical cached task instances.
  • interruptible marks whether the task may run on interruptible or pre-emptible capacity.
  • deprecated supplies the warning message for a deprecated task.
  • retries specifies the retry count. The retry_strategy property converts it to flytekit's RetryStrategy model.
  • timeout is the maximum duration for one execution. An integer is interpreted as seconds and converted to datetime.timedelta during construction.
  • pod_template_name selects an existing cluster PodTemplate resource.
  • generates_deck records whether the task generates a Deck URI.
  • is_eager marks an eager task.

The constructor validates related settings immediately. cache=True without a non-empty cache_version raises ValueError; cache_serialize=True and cache_ignore_input_vars are each rejected unless caching is enabled; and a timeout that is neither an integer nor a datetime.timedelta raises ValueError.

For example, this is valid because the cache version accompanies cache=True, and timeout=300 becomes a five-minute duration:

metadata = TaskMetadata(
    cache=True,
    cache_version="0.1",
    retries=2,
    timeout=300,
    interruptible=True,
    pod_template_name="gpu-pod",
)

At serialization time, to_taskmetadata_model() maps these Python fields into FlyteIDL's TaskMetadata. It sets discoverable from cache, creates runtime metadata from flytekit's version and the Python SDK, maps retries through retry_strategy, and carries through timeout, interruptibility, cache settings, deprecation text, Deck generation, pod-template name, and eagerness.

The older cache_serialize, cache_version, and cache_ignore_input_vars arguments on @task are deprecated. The decorator accepts a Cache object instead; when a Cache object is supplied, combining it with any of those three arguments raises ValueError. The decorator resolves the cache version and ignored inputs from the Cache object before constructing TaskMetadata.

The task class hierarchy​

flytekit separates the FlyteIDL representation of a task from Python-native typing and from container execution:

Task
└── PythonTask[T]
    └── PythonAutoContainerTask[T]
        └── PythonFunctionTask[T] / other task extensions

Task is the lowest-level abstraction and is closest to FlyteIDL TaskTemplate. Its constructor receives a task type, name, IDL-level TypedInterface, optional TaskMetadata, security context, and documentation. It does not assume Python-native input or output types. At this level, python_interface and get_input_types() return None by default.

PythonTask[T] adds a Python-native Interface and a plugin-specific task_config. It converts that interface into the IDL typed interface, exposes Python input and output types, and provides the lifecycle used by Python-backed tasks. It also owns Deck configuration: enable_deck, the deprecated disable_deck, and deck_fields. Supplying both enable_deck and disable_deck raises ValueError; use enable_deck for new code.

PythonAutoContainerTask[T] adds the hosted-container contract. It captures the image, resources, environment, resolver, pod template, accelerator, and shared-memory settings, and produces the pyflyte-execute command used in the task container. PythonFunctionTask is the usual function-backed extension; plugin tasks such as SQL, Spark, Ray, and sensor tasks also build on these bases.

Every Task instance appends itself to FlyteEntities.entities in Task.__init__. Consequently, creating a task makes it available to flytekit's serialization discovery; there is no separate registration call in the base class.

The Python execution lifecycle​

At runtime, the container entrypoint reads an IDL LiteralMap and invokes the task's dispatch_execute() method. The central flow is:

pyflyte-execute
  → resolver.load_task()
  → Task.dispatch_execute()
  → pre_execute()
  → LiteralMap → Python values
  → execute(**kwargs)
  → post_execute()
  → Python values → LiteralMap
  → outputs.pb (or a dynamic-job result)

flytekit.bin.entrypoint._dispatch_execute() first calls its load_task callback, downloads inputs.pb, converts it to a model LiteralMap, and calls task_def.dispatch_execute(ctx, idl_input_literals). It writes a VoidPromise as an empty output map, writes a returned LiteralMap, or handles a DynamicJobSpec.

PythonTask.dispatch_execute() implements the Python-specific part of the contract:

  1. Call pre_execute(ctx.user_space_params) and use the returned ExecutionParameters.
  2. Convert the input literal map to native Python arguments with TypeEngine.
  3. Call the subclass's execute(**native_inputs).
  4. Call post_execute(new_user_params, native_outputs).
  5. Convert native outputs back to a literal map with _output_to_literal_map().

pre_execute() is a no-op in PythonTask unless a plugin overrides it. The Spark task demonstrates the intended extension point: it creates or obtains a SparkSession, stores it in ExecutionParameters under SPARK_SESSION, and returns the updated parameters.

def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters:
    import pyspark as _pyspark

    ctx = FlyteContextManager.current_context()
    sess_builder = _pyspark.sql.SparkSession.builder.appName(f"FlyteSpark: {user_params.execution_id}")
    self.sess = sess_builder.getOrCreate()
    return user_params.builder().add_attr("SPARK_SESSION", self.sess).build()

Your execute() implementation returns Python-native values. post_execute() then receives both the updated user parameters and that return value; by default it returns the value unchanged, but task extensions can use it for cleanup or output adjustment. During output conversion, PythonTask maps declared output names to the returned value or tuple and asks TypeEngine to create each output literal.

Local calls use a related path. Task.local_execute() translates native values and Promise inputs into a literal map, optionally checks LocalTaskCache, and calls sandbox_execute(). sandbox_execute() creates a sandbox execution context and delegates to dispatch_execute(). Afterward, local_execute() wraps returned literals in Promise objects, or returns a VoidPromise when the task declares no outputs.

A distributed task can signal that its worker's outputs should not be uploaded by raising IgnoreOutputs from inside execute(). The PyTorch plugin uses this pattern: rank zero returns the result, while other ranks raise IgnoreOutputs(). The entrypoint recognizes FlyteUserRuntimeException containing IgnoreOutputs and returns without uploading outputs.pb.

Resolving a task inside its container​

Serialization and runtime loading are separate phases. TaskResolverMixin defines the contract:

  • location identifies the resolver in the serialized command.
  • loader_args(settings, task) converts a task into strings that identify it.
  • load_task(loader_args) reconstructs the task at runtime.
  • get_all_tasks() exposes resolver-owned tasks where supported.
  • task_name(task) optionally supplies a custom serialized task name.

PythonAutoContainerTask.get_default_command() combines these pieces into a pyflyte-execute command:

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

With the standard resolver, the serialized tail identifies a module and an attribute:

task-module tests.flytekit.unit.core.test_python_auto_container
 task-name task

DefaultTaskResolver.loader_args() obtains those values with extract_task_module(). At runtime, load_task() unpacks the arguments, calls importlib.import_module(task_module), and returns getattr(task_module, task_name). This resolver therefore expects a task available as a module-level attribute; PythonFunctionTask rejects nested or inner functions for this default loading strategy.

You can replace the resolver with set_resolver(), or pass task_resolver= when constructing a PythonAutoContainerTask. If a compilation context already has a resolver, that context resolver takes precedence and flytekit logs that it is ignoring the passed resolver.

Resolvers can encode identity differently. ClassStorageTaskResolver stores tasks in mapping, serializes the task's list index, and loads it by converting the single loader argument to an integer:

class ClassStorageTaskResolver(TrackedInstance, TaskResolverMixin):
    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]

DefaultNotebookTaskResolver instead serializes an entity-name, reads a gzip-compressed cloudpickle file, verifies that the pickled and current Python major and minor versions match, and returns the named entity. A mismatch raises RuntimeError; the notebook task must be created and run with matching Python major and minor versions.

Container configuration with PythonAutoContainerTask​

Use PythonAutoContainerTask when the user's code belongs in the task container and flytekit should generate the container definition. A minimal extension only needs to implement execute():

class DummyAutoContainerTask(PythonAutoContainerTask):
    def execute(self, **kwargs) -> Any:
        pass


task = DummyAutoContainerTask(name="x", task_config=None, task_type="t")

The constructor accepts a string or ImageSpec as container_image, requests and limits as separate resource settings, or a unified resources value. Do not combine resources with either requests or limits; flytekit raises ValueError. It also accepts environment, pod_template, accelerator, shared_memory, secret_requests, and pod_template_name.

For a normal container, get_container() resolves the image, merges serialization-settings environment variables with the task environment, adds runtime-package information for an ImageSpec, and uses get_command() for the container arguments. get_image() handles ImageSpec updates and obtains the registerable image string. When a pod_template is supplied, get_container() returns None and get_k8s_pod() returns the pod definition with flytekit's generated container merged into the pod template. This makes the pod-template and direct-container representations mutually exclusive.

The pod_template_name constructor argument has a separate precedence rule: PythonAutoContainerTask creates or uses the metadata object and assigns metadata.pod_template_name = pod_template_name. Thus, when both are supplied, the constructor argument overwrites the value in TaskMetadata while preserving other metadata such as retries.

Typed interfaces without a Python function​

Function-backed tasks normally obtain their Python-native interface from annotations. For non-function-based tasks such as ContainerTask and SQLTask, use kwtypes() to create the named type mapping that becomes the task interface:

square = ContainerTask(
    name="square",
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs=kwtypes(val=int),
    outputs=kwtypes(out=int),
    image="alpine",
    environment={"a": "b"},
    command=["sh", "-c", "echo $(( {{.Inputs.val}} * {{.Inputs.val}} )) | tee /var/outputs/out"],
)

sum = ContainerTask(
    name="sum",
    input_data_dir="/var/flyte/inputs",
    output_data_dir="/var/flyte/outputs",
    inputs=kwtypes(x=int, y=int),
    outputs=kwtypes(out=int),
    image="alpine",
    command=["sh", "-c", "echo $(( {{.Inputs.x}} + {{.Inputs.y}} )) | tee /var/flyte/outputs/out"],
)

Here kwtypes(val=int) and kwtypes(out=int) provide the Python types needed to transform the task's IDL literals, while the container command performs the user operation. This contrasts with @task functions, where flytekit's task decorator builds a PythonFunctionTask (or a registered plugin task) from the function signature.

Implementation constraints to keep in mind​

  • Construct TaskMetadata(cache=True, ...) with a non-empty cache_version; validation happens immediately.
  • Use the Cache object rather than the deprecated cache-related decorator arguments. Do not mix the two forms.
  • Set PythonAutoContainerTask._container_image before calling super().__init__(). The superclass registers the task in FlyteEntities.entities, and serialization may inspect the entity's image while construction is in progress.
  • Raise IgnoreOutputs from execute(), not from outside the Python task lifecycle, so dispatch_execute() and the entrypoint can handle it as the documented user-runtime signal.
  • Keep default-resolver tasks at module scope. Nested functions cannot be reconstructed by DefaultTaskResolver.
  • Do not manually rely on TaskMetadata.generates_deck at definition time. During serialization, flytekit's translator sets it when the entity has Decks enabled.
  • For notebook resolution, match Python major and minor versions between the environment that produced pkl.gz and the task container.
  • At the base Task level, get_container(), get_k8s_pod(), get_sql(), get_custom(), get_config(), and get_extended_resources() return None; task extensions override only the representation they support.