Manage Environment Dependencies and Reuse

Environment composition starts with TaskEnvironment

A TaskEnvironment is the user-facing execution boundary for a group of tasks. It extends Environment, whose shared fields include image, resources, env_vars, secrets, pod_template, description, interruptible, include, and deployment dependencies. TaskEnvironment adds task-oriented defaults such as cache, reusable, plugin_config, and queue.

Environment settings are propagated when a function is decorated. The resulting task receives the environment image, resources, environment variables, secrets, reuse policy, queue, interruptibility, and pod template (unless a task-level value replaces a setting). Its fully qualified name is formed from the environment name and the Python function name, while short_name is only a friendly display name. The task is also registered in env.tasks.

import flyte

env = flyte.TaskEnvironment(
    name="my_env",
    image=flyte.Image.from_debian_base(python="3.12").with_pip_packages("pandas"),
)

@env.task
async def my_task():
    pass

For this definition, the task template name is my_env.my_task. The TaskEnvironment.task() implementation stores that template under the same key in the environment's task mapping. Configuration therefore has a useful hierarchy: the environment supplies defaults, @env.task(...) supplies task-definition settings, and task.override(...) can supply invocation-time settings where the task contract permits them.

Deploy dependent environments together

Use depends_on when one environment should be part of another environment's deployment graph. In the Wikipedia embedding example, the driver has different resources from the GPU worker but declares the worker as a dependency:

image = flyte.Image.from_uv_script(__file__, name="embed_wikipedia_image").with_pip_packages(
    "unionai-reuse>=0.1.9",
)

N_GPUS = 1
worker = flyte.TaskEnvironment(
    name="embed_wikipedia_worker",
    image=image,
    resources=flyte.Resources(cpu=2, memory="8Gi", gpu=1),
    env_vars={"HF_HUB_ENABLE_HF_TRANSFER": "1"},
    reusable=flyte.ReusePolicy(replicas=16, concurrency=1, idle_ttl=120, scaledown_ttl=120),
    secrets="HF_HUB_TOKEN",
)

driver = flyte.TaskEnvironment(
    name="embed_wikipedia_driver",
    image=image,
    resources=flyte.Resources(cpu=1, memory="4Gi", disk="16Gi"),
    secrets="HF_HUB_TOKEN",
    depends_on=[worker],
)

Deployment planning recursively visits depends_on: _recursive_discover() adds the current environment and then visits each dependent environment. Consequently, deploying driver discovers worker as part of the same plan. plan_deploy() also tracks already visited names, so passing a dependency both through depends_on and as a top-level deployment argument does not create a second plan when it is the same environment object.

Environment names are the keys used by this traversal. Two different environment objects with the same name are checked as duplicates and raise an error; each TaskEnvironment therefore needs a unique valid name. Environment._validate_name() accepts lowercase snake_case or kebab-case names with numbers, not uppercase names, spaces, or leading separators.

Add dependencies after construction

add_dependency(*env) is the late-bound equivalent of constructor-time depends_on. It is useful when independent modules export environments and a deployment entrypoint assembles them after importing them:

from root_env import env
from spark_env import env as spark_env
from torch_env import env as torch_env

import flyte

env.add_dependency(torch_env, spark_env)

if __name__ == "__main__":
    flyte.init_from_config(root_dir=Path(__file__).parent)
    v = flyte.deploy(env)
    print(v[0].summary_repr())

The actual example also imports Path before this code and uses it to initialize the deployment root. add_dependency() requires every argument to be an Environment and rejects an environment whose name matches the receiver's name. The method's duplicate check does not remove duplicates from the original argument tuple before extending depends_on, despite its docstring saying duplicates are ignored; avoid adding the same dependency repeatedly.

Clone shared configuration without cloning tasks

TaskEnvironment.clone_with(name=...) uses dataclasses.replace() to construct a new environment. The base _get_kwargs() supplies the inherited infrastructure configuration: dependencies, image, resources, secrets, environment variables, pod template, description, and included files. Passing an argument replaces that field rather than merging it. The clone receives a new name and does not copy the original _tasks mapping, so tasks must be decorated in the clone if it is to own tasks.

The reuse example uses this pattern to make a non-reusable orchestration environment from a reusable worker. It explicitly makes the worker a dependency:

# Clone the environment with a different name and no reuse policy.
cloned_env = env.clone_with(name="nonreusable", reusable=None, depends_on=[env])

# This task will run in the cloned environment without reuse.
@cloned_env.task
async def main(n: int) -> List[int]:
    """
    Run square and cube tasks in parallel for the range of x_list.
    """
    square_coros = []
    cube_coros = []
    for x in range(n):
        square_coros.append(square(x))
        cube_coros.append(cube(x))
    return await asyncio.gather(*square_coros) + await asyncio.gather(*cube_coros)

Task-specific values deserve particular attention. clone_with() accepts cache and reusable as keyword overrides, but _get_kwargs() does not add either field to the inherited arguments. Thus a clone created without those overrides uses TaskEnvironment's defaults (cache="disable" and reusable=None). This is why the example explicitly passes reusable=None; a clone can also explicitly receive a new ReusePolicy. The tests verify both inheritance of image/resources and that clone.tasks == {}.

Reuse warm containers with ReusePolicy

Set reusable on a TaskEnvironment when task initialization is expensive relative to task execution. ReusePolicy describes a pool of warm replicas, and the runtime may reuse the Python process and container for later invocations. The available concurrent capacity is:

max_replicas * concurrency

For example, (1, 3) replicas with concurrency 2 gives a maximum capacity of six concurrent tasks. Reuse is shared state at the replica level: mutable process state can persist across invocations on one replica, but it is not shared across replicas or restarts. It should not be treated as a database or distributed global state.

Replica and timeout settings

Setting	Accepted values	Runtime meaning
replicas	An int, or (min, max)	An integer is normalized to (n, n); a tuple supplies autoscaling bounds.
concurrency	An integer, default 1	Maximum concurrent tasks per replica. Values above 1 require an async task function.
idle_ttl	Seconds as int, or timedelta	The environment-wide idle timeout. After global inactivity, all replicas are shut down. Minimum: 30 seconds.
scaledown_ttl	Seconds as int, or timedelta	The delay before an individual idle replica is removed during autoscaling. Minimum: 30 seconds.

min_replicas and max_replicas expose the normalized bounds. get_scaledown_ttl() returns the scale-down timeout as a timedelta (or None if the value is unset). The source validates the timeout minimums and the replica tuple length, but does not validate tuple ordering, tuple element types, positivity, or concurrency bounds; provide sensible values.

A reusable worker commonly looks like this:

actor_image = flyte.Image.from_debian_base().with_pip_packages("unionai-reuse")

env = flyte.TaskEnvironment(
    name="reusable",
    resources=flyte.Resources(memory="500Mi", cpu=1),
    reusable=flyte.ReusePolicy(
        replicas=(1, 4),
        idle_ttl=300,
    ),
    image=actor_image,
)

@env.task
async def square(x: int) -> int:
    return x**2

@env.task
async def cube(x: int) -> int:
    return x**3

The reusable examples add unionai-reuse to the image; that package setup is part of those examples' image configuration. A GPU worker can instead use fixed replicas and explicit environment settings, as in embed_wikipedia_worker, where 16 replicas each have concurrency 1, 120-second environment idle and scale-down TTLs, GPU resources, an environment variable, and a secret.

The SDK logs a warning for replicas=(1, 1) with concurrency=1. A parent task occupying the only replica can leave no replica available for child work. The ReusePolicy docstring recommends at least two replicas for this situation, or increasing concurrency, increasing replicas, or disabling reuse for the parent task that launches children.

Constraints and escape hatches

Reuse is deliberately narrower than ordinary task execution:

• TaskEnvironment rejects plugin_config together with reusable.
• A reusable environment rejects a task-level pod_template passed to @env.task(...).
• If concurrency > 1, decorating a synchronous function raises ValueError; use an async function or set concurrency to 1.
• Reusable task resources, environment variables, and secrets cannot be changed through task.override() while reuse remains enabled. To make an exceptional invocation with different resources, turn reuse off in the same override call:

async def high_mem_examples():
    files = ["20231101.sv/train-00000-of-00005.parquet", "20231101.war/train-00000-of-00001.parquet"]
    large_task = embed_shard_to_file.override(
        resources=flyte.Resources(cpu=2, memory="8Gi", gpu=1),
        reusable="off",
    )

    coros = []
    for file in files:
        coros.append(flyte.run.aio(large_task, "wikimedia/wikipedia", file, "all-MiniLM-L6-v2", batch_size=256))

This creates an override without changing the original reusable task template. The same reusable="off" requirement applies when overriding env_vars or secrets.

At serialization time, flyte._internal.runtime.reuse.add_reusable() supports container tasks and pods. An existing custom plugin configuration raises RuntimeUserError. For a supported reusable task, the policy is emitted as actor configuration: parallelism is concurrency, min_replica_count and replica_count are the normalized replica bounds, ttl_seconds is idle_ttl.total_seconds(), and scaledown_ttl_seconds is the normalized scale-down timeout. The serialized task type is set to "actor".

Troubleshooting deployment and reuse configuration

• Invalid environment name: rename it to lowercase snake_case or kebab-case, using only lowercase letters, digits, _, and - in the accepted pattern.
• Duplicate environment name: ensure distinct environment objects do not share a name. A dependency discovered recursively is already included; do not create a second object with that name.
• Dependency missing from a deployment: put it in depends_on=[...], or call add_dependency(...) in the deployment entrypoint before flyte.deploy(...).
• from_task() image error: all supplied task templates must have the same image by Python object identity, not merely equivalent image values; tasks already assigned to environments are rejected.
• Reusable image setup: the reusable examples install unionai-reuse in the image used by the worker.
• Unexpected clone behavior: clone_with() copies the base environment configuration but not its tasks. Cache and reuse are task-environment-specific defaults unless explicitly passed to clone_with(); redecorate tasks in the clone.
• Override rejected: when a task is reusable, include reusable="off" in the same override() call as a resources, environment-variable, or secret override.
• Starvation warning: a single replica with concurrency one is allowed but logs a warning; increase the replica minimum/capacity or disable reuse for the parent task that launches children.