Configure tasks

As we saw in Quickstart, you can run any Python function as a task in Flyte just by decorating it with @env.task.

This allows you to run your Python code in a distributed manner, with each function running in its own container. Flyte manages the spinning up of the containers, the execution of the code, and the passing of data between the tasks.

The simplest possible case is a TaskEnvironment with only a name parameter, and an env.task decorator, with no parameters:

env = flyte.TaskEnvironment(name="my_env")

@env.task
async def my_task(name:str) -> str:
    return f"Hello {name}!"

This will run your task in the default container environment with default settings.

But, of course, one of the key advantages of Flyte is the ability to control the software environment, hardware environment, and other execution parameters for each task, right in your Python code.

Task configuration levels

Task configuration is done at three levels. From most general to most specific, they are:

• The TaskEnvironment level: setting parameters when defining the TaskEnvironment object.
• The @env.task decorator level: Setting parameters in the @env.task decorator when defining a task function.
• The task invocation level: Using the task.override() method when invoking task execution.

Each level has its own set of parameters, and some parameters are shared across levels. For shared parameters, the more specific level will override the more general one.

Example

Here is an example of how these levels work together, showing each level with all available parameters:

# Level 1: TaskEnvironment - Base configuration
env_2 = flyte.TaskEnvironment(
    name="data_processing_env",
    image=flyte.Image.from_debian_base(),
    resources=flyte.Resources(cpu=1, memory="512Mi"),
    env_vars={"MY_VAR": "value"},
    # secrets=flyte.Secret(key="openapi_key", as_env_var="MY_API_KEY"),
    cache="disable",
    # pod_template=my_pod_template,
    # reusable=flyte.ReusePolicy(replicas=2, idle_ttl=300),
    depends_on=[another_env],
    description="Data processing task environment",
    # plugin_config=my_plugin_config
)

# Level 2: Decorator - Override some environment settings
@env_2.task(
    short_name="process",
    # secrets=flyte.Secret(key="openapi_key", as_env_var="MY_API_KEY_2"),
    cache="auto",
    # pod_template=my_pod_template,
    report=True,
    max_inline_io_bytes=100 * 1024,
    retries=3,
    timeout=60,
    docs="This task processes data and generates a report."
)
async def process_data(data_path: str) -> str:
    return f"Processed {data_path}"

@env_2.task
async def invoke_process_data() -> str:
    result = await process_data.override(
        resources=flyte.Resources(cpu=4, memory="2Gi"),
        env_vars={"MY_VAR": "new_value"},
        # secrets=flyte.Secret(key="openapi_key", as_env_var="MY_API_KEY_3"),
        cache="auto",
        max_inline_io_bytes=100 * 1024,
        retries=3,
        timeout=60
    )("input.csv")
    return result

Task configuration parameters

Each parameter is documented in detail on its dedicated page in this section. For the complete parameter interaction matrix showing which parameters can be set at which level, and for full type signatures and constraints, see the TaskEnvironment API reference.

*When reusable is set, resources, env_vars, and secrets can only be overridden via task.override() with reusable="off" in the same call.

Distributed tasks with a clustered environment

Everything above configures a task that runs in a single container. Some tasks instead need a gang of pods working together as one distributed job — the classic case being multi-node, multi-GPU model training. For those, use a specialized flyte.clustered.ClusteredTaskEnvironment instead of a plain TaskEnvironment. It inherits all the configuration above and adds a few fields that describe the cluster shape: replicas (pods) and nproc_per_node (worker processes per pod). When you run the task, the backend launches a Kubernetes JobSet of identical pods, sets up a torchrun rendezvous across them, and runs your task body once per worker.

import flyte
from flyte.clustered import ClusteredTaskEnvironment, TorchRun

env = ClusteredTaskEnvironment(
    name="ddp_env",
    image=flyte.Image.from_debian_base().with_pip_packages("torch"),
    resources=flyte.Resources(cpu=(2, 4), memory=("4Gi", "8Gi"), gpu="L4:1"),
    replicas=2,            # number of pods (nodes)
    nproc_per_node=1,      # worker processes per pod => world_size = 2 x 1
    runtime=TorchRun(),
)

See Clustered task environment for the full guide, including DDP, FSDP, and framework-integration examples.