Skip to main content

Task authoring and execution

Flyte tasks are the fundamental execution units in Flyte workflows. This section explains how to declare, configure, and execute tasks using flytekit, covering the core task abstractions, Python function task behavior, and valid runnable examples.

Core task abstractions​

All Flyte tasks inherit from the Task base class, which defines the core interface and execution hooks. The Task class captures metadata, interface, task type, and provides core execution methods like local_execute, dispatch_execute, execute, pre_execute, and post_execute. It does not have Python-native interfaces by itself.

The PythonTask class extends Task to support Python-native interfaces. It adds _python_interface, task_config, environment, and deck configuration (enable_deck, deck_fields). It implements compile, construct_node_metadata, _literal_map_to_python_input, _output_to_literal_map, _write_decks, and dispatch_execute. It is the base for all Python-based tasks that don't necessarily wrap a user function.

The TaskResolverMixin class helps resolve a task implementation. At execution time, for most tasks that generate a container target, the container image containing the task needs to be spun up again. The container needs to know which task it's supposed to run and how to rehydrate the task object. The TaskResolverMixin provides methods like location, name, load_task, loader_args, and get_all_tasks to support this.

Task metadata and configuration​

The TaskMetadata class holds task-level configuration. It includes cache, cache_serialize, cache_version, cache_ignore_input_vars, interruptible, deprecated, retries, timeout, pod_template_name, generates_deck, and is_eager.

The __post_init__ method enforces validation rules:

  • If cache=True, then cache_version must be set; otherwise, a ValueError is raised.
  • If cache_serialize=True, then cache=True must be set; otherwise, a ValueError is raised.
  • If cache_ignore_input_vars is specified, then cache=True must be set; otherwise, a ValueError is raised.

The timeout parameter can be specified as an int (seconds) or datetime.timedelta; otherwise, a ValueError is raised.

Task authoring with the @task decorator​

The @task decorator is the primary way to declare tasks in flytekit. It constructs TaskMetadata and selects the appropriate task plugin (including AsyncPythonFunctionTask for coroutines).

For a simple Python task:

@task
def my_task(x: int, y: typing.Dict[str, str]) -> str:
    ...

For specific task types:

@task(task_config=Spark(), retries=3)
def my_task(x: int, y: typing.Dict[str, str]) -> str:
    ...

The @task decorator supports many configuration options:

  • cache: Boolean or Cache object indicating caching behavior
  • cache_version: Version string for cache key; required if cache=True
  • cache_serialize: Boolean indicating if identical cache instances should run serially
  • cache_ignore_input_vars: Tuple of input variable names to exclude from cache hash
  • retries: Number of times to retry the task on failure
  • interruptible: Boolean indicating if the task can be preempted
  • deprecated: Warning message for deprecated tasks
  • timeout: Maximum execution duration (int seconds or timedelta)
  • container_image: Optional custom image for this task
  • environment: Environment variables for task execution
  • requests: Compute resource requests (CPU, memory)
  • limits: Compute resource limits (CPU, memory)
  • secret_requests: List of secret keys to inject at runtime
  • execution_mode: Execution behavior (DEFAULT, DYNAMIC, EAGER)
  • node_dependency_hints: List of tasks/workflows/launch plans this task depends on (for dynamic tasks)
  • task_resolver: Custom task resolver for serialization/deserialization
  • enable_deck: Boolean to enable deck generation
  • deck_fields: Tuple of DeckField values to include in deck
  • pod_template: Custom PodTemplate for this task
  • pod_template_name: Name of existing PodTemplate resource
  • accelerator: Hardware accelerator to use (e.g., GPU)
  • pickle_untyped: Boolean to allow untyped outputs to be pickled

PythonFunctionTask: The core task type for Python functions​

The PythonFunctionTask class is the core task type for wrapping user-defined Python functions. It auto-detects the function interface, supports DEFAULT, DYNAMIC, and EAGER execution modes, and handles both synchronous and asynchronous functions.

The execute method dispatches based on execution mode:

def execute(self, **kwargs) -> Any:
    if self.execution_mode == self.ExecutionBehavior.DEFAULT:
        return self._task_function(**kwargs)
    elif self.execution_mode == self.ExecutionBehavior.DYNAMIC:
        return self.dynamic_execute(self._task_function, **kwargs)

The PythonFunctionTask enforces that the task function cannot be a nested/inner/local function unless it's a test function or wrapped with functools.wraps/update_wrapper. The default task resolver requires module-level accessibility.

Execution modes​

Flyte tasks support three execution modes: DEFAULT, DYNAMIC, and EAGER.

DEFAULT execution mode​

The DEFAULT execution mode is the standard task execution mode. The task function is executed directly when execute is called.

DYNAMIC execution mode​

The DYNAMIC execution mode is used for dynamic tasks. The task function constitutes a workflow, so it must be compiled at runtime and then executed.

The dynamic_execute method handles both local and remote execution:

def dynamic_execute(self, task_function: Callable, **kwargs) -> Any:
    ctx = FlyteContextManager.current_context()
    if ctx.execution_state and ctx.execution_state.is_local_execution():
        # The rest of this function mimics the local_execute of the workflow. We can't use the workflow
        # local_execute directly though since that converts inputs into Promises.
        logger.debug(f"Executing Dynamic workflow, using raw inputs {kwargs}")
        self._create_and_cache_dynamic_workflow()
        if self.execution_mode == self.ExecutionBehavior.DYNAMIC:
            es = ctx.new_execution_state().with_params(mode=ExecutionState.Mode.LOCAL_DYNAMIC_TASK_EXECUTION)
        else:
            es = cast(ExecutionState, ctx.execution_state)
        with FlyteContextManager.with_context(ctx.with_execution_state(es)):
            function_outputs = cast(PythonFunctionWorkflow, self._wf).execute(**kwargs)

        if isinstance(function_outputs, VoidPromise) or function_outputs is None:
            return VoidPromise(self.name)

        if len(cast(PythonFunctionWorkflow, self._wf).python_interface.outputs) == 0:
            raise FlyteValueException(function_outputs, "Interface output should've been VoidPromise or None.")

        # TODO: This will need to be cleaned up when we revisit top-level tuple support.
        expected_output_names = list(self.python_interface.outputs.keys())
        if len(expected_output_names) == 1:
            # Here we have to handle the fact that the wf could've been declared with a typing.NamedTuple of
            # length one. That convention is used for naming outputs - and single-length-NamedTuples are
            # particularly troublesome but elegant handling of them is not a high priority
            # Again, we're using the output_tuple_name as a proxy.
            if self.python_interface.output_tuple_name and isinstance(function_outputs, tuple):
                wf_outputs_as_map = {expected_output_names[0]: function_outputs[0]}
            else:
                wf_outputs_as_map = {expected_output_names[0]: function_outputs}
        else:
            wf_outputs_as_map = {
                expected_output_names[i]: function_outputs[i] for i, _ in enumerate(function_outputs)
            }

        # In a normal workflow, we'd repackage the promises coming from tasks into new Promises matching the
        # workflow's interface. For a dynamic workflow, just return the literal map.
        wf_outputs_as_literal_dict = translate_inputs_to_literals(
            ctx,
            wf_outputs_as_map,
            flyte_interface_types=self.interface.outputs,
            native_types=self.python_interface.outputs,
        )
        return _literal_models.LiteralMap(literals=wf_outputs_as_literal_dict)

    if ctx.execution_state and ctx.execution_state.mode == ExecutionState.Mode.TASK_EXECUTION:
        return self.compile_into_workflow(ctx, task_function, **kwargs)

    if ctx.execution_state and ctx.execution_state.mode == ExecutionState.Mode.LOCAL_TASK_EXECUTION:
        return task_function(**kwargs)

    raise ValueError(f"Invalid execution provided, execution state: {ctx.execution_state}")

Dynamic tasks must be registered separately if they reference launch plans. The node_dependency_hints parameter can be used to specify tasks, launch plans, or workflows that this task depends on. This is only for dynamic tasks/workflows, where flyte cannot automatically determine the dependencies prior to runtime.

@workflow
def workflow0():
    ...

launchplan0 = LaunchPlan.get_or_create(workflow0)

# Specify node_dependency_hints so that launchplan0 will be registered on flyteadmin, despite this being a
# dynamic task.
@dynamic(node_dependency_hints=[launchplan0])
def launch_dynamically():
    # To run a sub-launchplan it must have previously been registered on flyteadmin.
    return [launchplan0]*10

EAGER execution mode​

The EAGER execution mode is used for eager tasks. It sets is_eager=True in metadata, uses ExecutionBehavior.EAGER, and implements run_with_backend for remote eager execution. It supports calling other Flyte entities (tasks, workflows) during local execution via worker_queue.

The async_execute method handles both local and remote execution:

async def async_execute(self, *args, **kwargs) -> Any:
    """
    Overrides the base execute function. This function does not handle dynamic at all. Eager and dynamic don't mix.

    Some notes on the different call scenarios since it's a little different than other tasks.
    a) starting local execution - eager_task()
        -> last condition of call handler,
        -> set execution mode and self.local_execute()
        -> self.execute(native_vals)
          -> 1) -> task function() or 2) -> self.run_with_backend()  # fn name will be changed.
    b) inside an eager task local execution - calling normal_task()
        -> call handler detects in eager local execution (middle part of call handler)
        -> call normal_task's local_execute()
    c) inside an eager task local execution - calling async_normal_task()
        -> produces a coro, which when awaited/run
            -> call handler detects in eager local execution (middle part of call handler)
            -> call async_normal_task's local_execute()
            -> call AsyncPythonFunctionTask's async_execute(), which awaits the task function
    d) inside an eager task local execution - calling another_eager_task()
        -> produces a coro, which when awaited/run
            -> call handler detects in eager local execution (middle part of call handler)
            -> call another_eager_task's local_execute()
            -> results are returned instead of being passed to create_native_named_tuple
    d) eager_task, starting backend execution from entrypoint.py
        -> eager_task.dispatch_execute(literals)
        -> eager_task.execute(native values)
        -> awaits eager_task.run_with_backend()  # fn name will be changed
    e) in an eager task during backend execution, calling any flyte_entity()
        -> add the entity to the worker queue and await the result.
    """
    # Args is present because the asyn helper function passes it, but everything should be in kwargs by this point
    assert len(args) == 1
    ctx = FlyteContextManager.current_context()
    is_local_execution = cast(ExecutionState, ctx.execution_state).is_local_execution()
    if not is_local_execution:
        # a real execution
        return await self.run_with_backend(**kwargs)
    else:
        # set local mode and proceed with running the function.  This makes the
        mode = self.local_execution_mode()
        with FlyteContextManager.with_context(
            ctx.with_execution_state(cast(ExecutionState, ctx.execution_state).with_params(mode=mode))
        ):
            return await self._task_function(**kwargs)

Eager tasks require worker_queue to be set during remote execution; otherwise, a default remote is constructed.

AsyncPythonFunctionTask: For async tasks​

The AsyncPythonFunctionTask class is the base task for async tasks (coroutines). It overrides __call__ and async_execute to support async task functions. execute is bound to loop_manager.synced(async_execute). It does not support dynamic mode.

The @task decorator automatically selects AsyncPythonFunctionTask for coroutine functions unless a non-default plugin is used:

if inspect.iscoroutinefunction(fn):
    if task_plugin is PythonFunctionTask:
        task_plugin = AsyncPythonFunctionTask
    else:
        if not issubclass(task_plugin, AsyncPythonFunctionTask):
            raise AssertionError(f"Task plugin {task_plugin} is not compatible with async functions")

ArrayNodeMapTask: For parallel execution​

The ArrayNodeMapTask class is a PythonTask that wraps a PythonFunctionTask or PythonInstanceTask to enable parallel execution over list inputs. It transforms the interface to List[T] and executes the underlying task multiple times. It supports concurrency, min_successes, min_success_ratio, and bound_inputs. It only supports single-output tasks and tasks with DEFAULT execution mode (not dynamic/eager).

The input validation enforces these constraints:

# TODO: add support for other Flyte entities
if not (
    (
        isinstance(actual_task, PythonFunctionTask)
        and actual_task.execution_mode == PythonFunctionTask.ExecutionBehavior.DEFAULT
    )
    or isinstance(actual_task, PythonInstanceTask)
):
    raise ValueError(
        "Only PythonFunctionTask with default execution mode (not @dynamic or @eager) and PythonInstanceTask are supported in map tasks."
    )

n_outputs = len(actual_task.python_interface.outputs)
if n_outputs > 1:
    raise ValueError("Only tasks with a single output are supported in map tasks.")

The _raw_execute method handles local execution:

def _raw_execute(self, **kwargs) -> Any:
    """
    This is called during locally run executions. Unlike array task execution on the Flyte platform, _raw_execute
    produces the full output collection.
    """
    outputs_expected = True
    if not self.interface.outputs:
        outputs_expected = False
    outputs = []

    mapped_tasks_count = 0
    if self.python_function_task.interface.inputs.items():
        for k in self.python_function_task.interface.inputs.keys():
            v = kwargs[k]
            if isinstance(v, list) and k not in self.bound_inputs:
                mapped_tasks_count = len(v)
                break

    failed_count = 0
    min_successes = mapped_tasks_count
    if self._min_successes:
        min_successes = self._min_successes
    elif self._min_success_ratio:
        min_successes = math.ceil(min_successes * self._min_success_ratio)

    for i in range(mapped_tasks_count):
        single_instance_inputs = {}
        for k in self.interface.inputs.keys():
            v = kwargs[k]
            if isinstance(v, list) and k not in self._bound_inputs:
                single_instance_inputs[k] = kwargs[k][i]
            else:
                single_instance_inputs[k] = kwargs[k]
        try:
            o = self._run_task.execute(**single_instance_inputs)
            if outputs_expected:
                outputs.append(o)
        except Exception as exc:
            outputs.append(None)
            failed_count += 1
            if mapped_tasks_count - failed_count < min_successes:
                logger.error("The number of successful tasks is lower than the minimum ratio")
                raise exc

    return outputs

Task execution flow​

The local_execute method in the Task class handles local execution with caching support:

def local_execute(
    self, ctx: FlyteContext, **kwargs
) -> Union[Tuple[Promise], Promise, VoidPromise, Coroutine, None]:
    """
    This function is used only in the local execution path and is responsible for calling dispatch execute.
    Use this function when calling a task with native values (or Promises containing Flyte literals derived from
    Python native values).
    """
    # Unwrap the kwargs values. After this, we essentially have a LiteralMap
    # The reason why we need to do this is because the inputs during local execute can be of 2 types
    #  - Promises or native constants
    #  Promises as essentially inputs from previous task executions
    #  native constants are just bound to this specific task (default values for a task input)
    #  Also along with promises and constants, there could be dictionary or list of promises or constants
    try:
        literals = translate_inputs_to_literals(
            ctx,
            incoming_values=kwargs,
            flyte_interface_types=self.interface.inputs,
            native_types=self.get_input_types(),  # type: ignore
        )
    except TypeTransformerFailedError as exc:
        exc.args = (f"Failed to convert inputs of task '{self.name}':\n  {exc.args[0]}",)
        raise
    input_literal_map = _literal_models.LiteralMap(literals=literals)

    # if metadata.cache is set, check memoized version
    local_config = LocalConfig.auto()
    if self.metadata.cache and local_config.cache_enabled:
        if local_config.cache_overwrite:
            outputs_literal_map = None
            logger.info("Cache overwrite, task will be executed now")
        else:
            logger.info(
                f"Checking cache for task named {self.name}, cache version {self.metadata.cache_version} "
                f", inputs: {kwargs}, and ignore input vars: {self.metadata.cache_ignore_input_vars}"
            )
            outputs_literal_map = LocalTaskCache.get(
                self.name, self.metadata.cache_version, input_literal_map, self.metadata.cache_ignore_input_vars
            )
            # The cache returns None iff the key does not exist in the cache
            if outputs_literal_map is None:
                logger.info("Cache miss, task will be executed now")
            else:
                logger.info("Cache hit")
        if outputs_literal_map is None:
            outputs_literal_map = self.sandbox_execute(ctx, input_literal_map)
            # TODO: need `native_inputs`
            LocalTaskCache.set(
                self.name,
                self.metadata.cache_version,
                input_literal_map,
                self.metadata.cache_ignore_input_vars,
                outputs_literal_map,
            )
            logger.info(
                f"Cache set for task named {self.name}, cache version {self.metadata.cache_version} "
                f", inputs: {kwargs}, and ignore input vars: {self.metadata.cache_ignore_input_vars}"
            )
    else:
        # This code should mirror the call to `sandbox_execute` in the above cache case.
        # Code is simpler with duplication and less metaprogramming, but introduces regressions
        # if one is changed and not the other.
        outputs_literal_map = self.sandbox_execute(ctx, input_literal_map)

    outputs_literals = outputs_literal_map.literals

    # TODO maybe this is the part that should be done for local execution, we pass the outputs to some special
    #    location, otherwise we dont really need to right? The higher level execute could just handle literalMap
    # After running, we again have to wrap the outputs, if any, back into Promise objects
    output_names = list(self.interface.outputs.keys())  # type: ignore
    if len(output_names) != len(outputs_literals):
        # Length check, clean up exception
        raise AssertionError(f"Length difference {len(output_names)} {len(outputs_literals)}")

    # Tasks that don't return anything still return a VoidPromise
    if len(output_names) == 0:
        return VoidPromise(self.name)

    vals = [Promise(var, outputs_literals[var]) for var in output_names]
    return create_task_output(vals, self.python_interface)

The dispatch_execute method translates Flyte's Type system based input values and invokes the actual call to the executor. It is invoked during runtime.

The pre_execute method is invoked directly before executing the task method and before all the inputs are converted. One particular case where this is useful is if the context is to be modified for the user process to get some user space parameters. This also ensures that things like SparkSession are already correctly setup before the type transformers are called.

The post_execute method is called after the execution has completed, with the user_params and can be used to clean-up, or alter the outputs to match the intended tasks outputs. If not overridden, then this function is a No-op.

The IgnoreOutputs exception can be raised to indicate that the outputs generated by the task can be safely ignored. This is useful in case of distributed training or peer-to-peer parallel algorithms.

Deck generation​

Tasks can generate HTML decks to provide additional context about the task execution. The enable_deck parameter controls deck generation. The deck_fields parameter specifies which decks to generate.

The DeckField enum includes:

  • SOURCE_CODE: The source code of the task function
  • DEPENDENCIES: Python dependencies
  • TIMELINE: Timeline of the task execution
  • INPUT: Inputs to the task
  • OUTPUT: Outputs from the task

The disable_deck parameter is deprecated; use enable_deck instead. Setting both disable_deck and enable_deck raises a ValueError.

Gotchas and warnings​

  • TaskFunction cannot be a nested/inner/local function unless it's a test function or wrapped with functools.wraps/update_wrapper. The default task resolver requires module-level accessibility.
  • Caching requires cache_version to be set when cache=True. Otherwise, ValueError is raised in TaskMetadata.__post_init__.
  • cache_serialize and cache_ignore_input_vars require cache=True; otherwise, ValueError is raised.
  • ArrayNodeMapTask only supports tasks with a single output and only DEFAULT execution mode (not @dynamic or @eager).
  • Dynamic tasks cannot be used inside eager tasks (and vice versa). EagerAsyncPythonFunctionTask raises NotImplementedError for dynamic mode.
  • Async tasks cannot be used as dynamic tasks. AsyncPythonFunctionTask raises NotImplementedError for dynamic mode.
  • When using functools.partial with map tasks, bound inputs must be handled carefully to avoid interface mismatch.
  • Tasks that don't return anything still return a VoidPromise.
  • During local execution, inputs can be Promises or native constants; local_execute handles both.
  • The task name is auto-generated from the module and function name unless explicitly provided.
  • Eager tasks require worker_queue to be set during remote execution; otherwise, a default remote is constructed.
  • Dynamic tasks compile a workflow at runtime and must be registered separately if they reference launch plans.
  • Reference tasks must have an interface that exactly matches the remote task; otherwise, compilation fails.
  • Tasks with multiple outputs must be declared with a typing.NamedTuple or tuple return type.
  • Timeout can be specified as int (seconds) or datetime.timedelta; otherwise, ValueError is raised.
  • Tasks decorated with @task that are coroutines are automatically wrapped in AsyncPythonFunctionTask unless a non-default plugin is used.
  • Tasks with deck generation enabled will write decks only during local execution if ctx.user_space_params is available.
  • Map tasks in local execution produce full output collections; remote execution produces per-index outputs.