Define Task Environments

What you will build

You will define a named flyte.TaskEnvironment, put Python tasks in it, inspect their fully qualified names, and verify shared defaults. You will also compose environments for deployment, derive a reusable/non-reusable pair with clone_with(), group a standalone ContainerTask with from_task(), and see how flyte.extend.TaskTemplate supports plugin tasks.

Prerequisites

You need Python code using flyte-sdk's public flyte package. The examples use flyte.TaskEnvironment, flyte.Resources, flyte.Image, flyte.ReusePolicy, and flyte.extras.ContainerTask. To run a remote deployment or run, configure Flyte first; the repository examples call flyte.init_from_config() before creating a remote run.

1. Create a named task environment

Start with an environment and two tasks:

import flyte

env = flyte.TaskEnvironment(
    name="hello_world",
    resources=flyte.Resources(),
)


@env.task
async def say_hello(data: str, lt: list[int]) -> str:
    return f"Hello {data} {lt}"


@env.task
async def square(i: int = 3) -> int:
    return i * i


assert env.tasks["hello_world.say_hello"] is say_hello
assert env.tasks["hello_world.square"] is square
assert say_hello.name == "hello_world.say_hello"
assert say_hello.resources == env.resources

TaskEnvironment.task() creates a task template, stores it in the environment's task mapping, and names it <environment name>.<function name>. Thus say_hello is registered as hello_world.say_hello. The environment name and function-name component are not changed by task overrides.

TaskEnvironment inherits common fields from flyte.Environment and adds task-oriented defaults. Its environment-level surface includes image, resources, env_vars, secrets, interruptible, queue, pod_template, description, depends_on, include, cache, reusable, and plugin_config. image defaults to "auto"; TaskTemplate.__post_init__() converts that request to an Image. resources supplies CPU, memory, GPU, or disk settings. env_vars is a dictionary copied into task templates, and secrets is an SDK secret request used at runtime. interruptible defaults to False, while cache defaults to "disable".

The environment also supplies a shared queue and interruptibility setting to its tasks:

import flyte

env = flyte.TaskEnvironment(
    name="batch_workers",
    image=flyte.Image.from_base("python:3.12"),
    resources=flyte.Resources(),
    env_vars={"LOG_LEVEL": "INFO"},
    queue="gpu-cluster",
    interruptible=True,
)


@env.task
async def process(value: int) -> int:
    return value * 2


assert process.image == env.image
assert process.env_vars == env.env_vars
assert process.queue == "gpu-cluster"
assert process.interruptible is True

Secrets use the SDK's SecretRequest forms and are copied to task templates. The connector example uses flyte.Secret with as_env_var to map a control-plane secret to a task/container environment variable.

2. Apply the three configuration levels

Task configuration proceeds from general to specific: TaskEnvironment, then @env.task(...), then task.override(...).

import flyte

env = flyte.TaskEnvironment(
    name="processing",
    resources=flyte.Resources(),
    cache="disable",
    queue="standard",
    interruptible=False,
)


@env.task(
    cache="auto",
    queue="priority",
    interruptible=True,
    short_name="visible_process",
)
async def process(value: int) -> int:
    return value + 1


per_call = process.override(
    resources=flyte.Resources(),
    queue="urgent",
    interruptible=False,
)

assert process.name == "processing.process"
assert process.short_name == "visible_process"
assert per_call.name == process.name
assert per_call.queue == "urgent"
assert per_call.interruptible is False

The decorator's cache, queue, and interruptible values are passed to the created TaskTemplate; omitted values use the environment settings. It also accepts retries, timeout, pod template, report, inline-I/O size, triggers, links, a task resolver, and entrypoint. task.override() returns a replaced template for invocation-specific resources, cache, retries, timeout, reusable policy, environment variables, secrets, inline-I/O size, pod template, queue, interruptibility, entrypoint, and links.

short_name is a friendly name only. The example's short name is visible_process, while its canonical name remains processing.process. TaskTemplate.override() rejects attempts to change name, image, docs, or interface.

There are two truthiness details to account for. TaskEnvironment.task() uses cache or self.cache and queue or self.queue; a falsey decorator value falls back to the environment. TaskTemplate.override() uses similar fallback behavior for cache, retries, timeout, resources, environment variables, secrets, queue, and links, so a falsey value may not clear an existing default.

Reusable environments

Configure worker reuse with flyte.ReusePolicy:

import flyte

processing_env = flyte.TaskEnvironment(
    name="deltalake_processing",
    resources=flyte.Resources(),
    reusable=flyte.ReusePolicy(),
    image=flyte.Image.from_base("python:3.12"),
    cache="auto",
)

When reusable is set, TaskEnvironment.task() rejects a task-level pod_template. TaskTemplate.override() rejects resource, environment-variable, and secret overrides while reusability remains enabled. To make one invocation non-reusable and change those values, pass reusable="off" in the same call:

non_reusable_call = process.override(
    reusable="off",
    resources=flyte.Resources(),
)

A reusable environment with concurrency greater than one requires an async task function. Environment-level plugin_config also cannot be combined with reusable.

3. Compose environments for deployment

Use depends_on when one environment contains tasks that another environment needs deployed alongside it:

import flyte
from flyte.extras import ContainerTask

greeting_task = ContainerTask(
    name="echo_and_return_greeting",
    image=flyte.Image.from_base("alpine:3.18"),
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"name": str},
    outputs={"greeting": str},
    command=[
        "/bin/sh",
        "-c",
        "echo 'Hello, my name is {{.inputs.name}}.' | tee -a /var/outputs/greeting",
    ],
)

container_env = flyte.TaskEnvironment.from_task("container_env", greeting_task)
env = flyte.TaskEnvironment(name="hello_world", depends_on=[container_env])


@env.task
async def say_hello(name: str = "flyte") -> str:
    return await greeting_task(name=name)

plan_deploy() recursively follows Environment.depends_on, placing the root and its dependencies in a deployment plan. An environment already included through an earlier root is not planned a second time. Deployment also requires unique environment names: two distinct objects with the same name produce a duplicate-name error. For a src/ layout that imported the same file under two module names, the diagnostic recommends the appropriate --root-dir, such as --root-dir src.

You can add a dependency after construction:

import flyte

base_env = flyte.TaskEnvironment(name="base")
worker_env = flyte.TaskEnvironment(name="worker")
worker_env.add_dependency(base_env)
assert worker_env.depends_on == [base_env]

add_dependency() accepts Environment instances, rejects other types, and rejects a self-dependency. Although its docstring says duplicate dependencies are silently ignored, the current loop checks membership without skipping before extend, so repeated calls can append duplicates.

4. Derive environments without copying tasks

clone_with() keeps shared settings while changing selected fields:

import flyte

processing_env = flyte.TaskEnvironment(
    name="deltalake_processing",
    resources=flyte.Resources(),
    reusable=flyte.ReusePolicy(),
    image=flyte.Image.from_base("python:3.12"),
    cache="auto",
)

orchestrator_env = processing_env.clone_with(
    name="deltalake_orchestrator",
    reusable=None,
    depends_on=[processing_env],
)

assert orchestrator_env.image == processing_env.image
assert orchestrator_env.resources == processing_env.resources
assert orchestrator_env.reusable is None
assert orchestrator_env.depends_on == [processing_env]
assert orchestrator_env.tasks == {}

The new environment inherits settings not explicitly overridden, but its task registry is empty. Decorate functions on the returned environment if they belong there.

5. Group standalone templates with from_task()

TaskEnvironment.from_task() accepts standalone TaskTemplate objects, including ContainerTask and connector tasks:

import flyte
from flyte.extras import ContainerTask

greeting_task = ContainerTask(
    name="echo_and_return_greeting",
    image=flyte.Image.from_base("alpine:3.18"),
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs={"name": str},
    outputs={"greeting": str},
    command=[
        "/bin/sh",
        "-c",
        "echo 'Hello, my name is {{.inputs.name}}.' | tee -a /var/outputs/greeting",
    ],
)

container_env = flyte.TaskEnvironment.from_task("container_env", greeting_task)
assert container_env.tasks[greeting_task.name] is greeting_task
assert greeting_task.parent_env_name == "container_env"

The supplied templates must have compatible images and must not already belong to an environment. If no image is defined, the new environment uses "auto". Each template receives a weak reference to the new environment and its parent_env_name; deployment and serialization consume this metadata.

6. Understand the shared Environment base

flyte.Environment is the dataclass base shared by TaskEnvironment and AppEnvironment. It owns common image, resource, secret, environment-variable, pod-template, description, interruptibility, dependency, and include-file fields. Normally construct flyte.TaskEnvironment, not flyte.Environment: the base clone_with() raises NotImplementedError, while TaskEnvironment supplies task registration and cloning.

Names must match the SDK's lowercase snake-case or kebab-case grammar, including digits: valid_name, valid-name, and valid123_name are valid; invalid-name! raises ValueError. Descriptions longer than 255 characters are shortened with the SDK description parser. include must be a sequence of strings, not a bare string, and is normalized to a tuple. Relative include paths are anchored to the user file where the environment was instantiated and are added to normal code-bundle discovery.

Every environment is appended to the process-global load-order registry at construction. flyte._environment.list_loaded_environments() returns that registry directly, so temporary environments created during tests or imports appear in it too.

7. Extend TaskTemplate for plugin tasks

flyte.extend.TaskTemplate is the extensible dataclass for executable task templates. A normal TaskEnvironment.task() creates an AsyncFunctionTaskTemplate; when plugin_config is present, it looks up a matching class in flyte.extend.TaskPluginRegistry. Plugin authors can subclass TaskTemplate and provide a NativeInterface, task type, and serialization hooks.

The repository's connector example defines BatchJobTask like this:

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: 'flyte.models.NativeInterface' accepts at most 0 positional argument(s)

from dataclasses import dataclass
from typing import Any, Dict, Optional, Type

from flyte.connectors import AsyncConnectorExecutorMixin
from flyte.extend import TaskTemplate
from flyte.models import NativeInterface, SerializationContext


@dataclass
class BatchJobConfig:
    timeout_seconds: int = 300


class BatchJobTask(AsyncConnectorExecutorMixin, TaskTemplate):
    _TASK_TYPE = "batch_job"

    def __init__(
        self,
        name: str,
        plugin_config: BatchJobConfig,
        inputs: Optional[Dict[str, Type]] = None,
        outputs: Optional[Dict[str, Type]] = None,
        **kwargs,
    ):
        super().__init__(
            name=name,
            interface=NativeInterface(
                {k: (v, None) for k, v in inputs.items()} if inputs else {},
                outputs or {},
            ),
            task_type=self._TASK_TYPE,
            image=None,
            **kwargs,
        )
        self.plugin_config = plugin_config

    def custom_config(self, sctx: SerializationContext) -> Optional[Dict[str, Any]]:
        return {"timeout_seconds": self.plugin_config.timeout_seconds}

TaskTemplate stores identity and runtime defaults, including image, resources, cache, interruptibility, retries, reusable policy, environment variables, secrets, timeout, pod template, queue, and interface metadata. Initialization converts string image requests to Image, string cache requests to Cache, integer retries to RetryStrategy, and an empty short name to the task name. custom_config() is one of its plugin serialization hooks.

Common pitfalls

• Environment names must be lowercase snake-case or kebab-case and use digits only where the name grammar permits them.
• Environment names must be unique at deployment; dual imports in a src/ layout may require --root-dir.
• TaskEnvironment caching defaults to "disable"; select cache="auto", "override", or a flyte.Cache explicitly.
• short_name changes the friendly name, not the fully qualified task name.
• Reusable environments restrict pod-template and resource, environment-variable, and secret overrides; use reusable="off" with the invocation override when needed.
• clone_with() does not copy tasks, and from_task() cannot combine already-associated templates or incompatible images.
• add_dependency() rejects invalid and self dependencies, but its current duplicate check does not prevent duplicate entries despite the docstring.