Tasks and Environments

Define an execution environment

When several tasks share an image and compute defaults, put those settings on one flyte.TaskEnvironment and decorate functions with @env.task:

import flyte

resources = flyte.Resources()
resources.cpu = "1"
resources.memory = "1Gi"
resources.shm = "auto"
resources.disk = "1Gi"
resources.gpu = "A100 80G:8"

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


@env.task
async def my_task(x: int) -> int:
    return x + 1


@env.task
async def my_task2(x: int) -> int:
    return await my_task(x) + 1

TaskEnvironment defines three configuration levels: environment defaults, @env.task(...) settings, and invocation-time task.override(...) settings. The decorator creates the fixed fully qualified name <environment>.<function_name>—resources.my_task above. short_name changes only the friendly/UI name. override() cannot replace a task's name, image, docs, or interface.

Internally, TaskEnvironment.task() captures the function signature with NativeInterface.from_callable(), constructs an AsyncFunctionTaskTemplate (or a plugin-selected subclass), propagates environment values, and records the template in env.tasks. Environment also holds depends_on, pod_template, env_vars, secrets, interruptible, description, and include settings. Names are validated as lower-case snake_case or kebab-case names with numbers; descriptions longer than 255 characters are parsed/truncated. Relative include paths are anchored to the file that created the environment.

Use clone_with() to preserve common settings while changing the environment name or selected values:

gpu_env = flyte.TaskEnvironment(name="gpu_env", image=my_image)
gpu_env_cached = gpu_env.clone_with("gpu_env_cached", cache="auto")

The method uses dataclasses.replace, inherits unspecified settings, and creates a fresh task registry.

Configure compute resources

flyte.Resources models CPU, memory, accelerators, ephemeral disk, and shared memory. CPU and memory can be scalar values or two-element request/limit tuples. The resource fields are mutable dataclass fields, so the following also shows the values represented by the public object without relying on constructor keyword parsing:

resources = flyte.Resources()
resources.cpu = (1, 2)
resources.memory = ("1Gi", "4Gi")
resources.disk = "10Gi"

CPU tuples and memory tuples must contain exactly two elements, and numeric CPU values cannot be negative. GPU values may be an integer count, a recognized "type:quantity" string, or a Device made by flyte.GPU(...) or flyte.TPU(...):

resources = flyte.Resources()
resources.gpu = "A100:2"
assert resources.get_device().device == "A100"
assert resources.get_device().quantity == 2

resources.gpu = flyte.GPU(device="A100", quantity=1, partition="1g.5gb")
tpu_resources = flyte.Resources()
tpu_resources.gpu = flyte.TPU(device="V5P", partition="2x2x1")

Resources.get_device() converts the typed string to a Device and returns accelerator objects unchanged. Unknown accelerator names, missing colons, non-integer quantities, and negative integer counts raise ValueError. Device requires a quantity of at least one. shm="auto" remains "auto" on the public object but get_shared_memory() returns "", which is the representation used to request maximum available shared memory:

resources = flyte.Resources()
resources.shm = "auto"
assert resources.get_shared_memory() == ""

flyte.extend.get_proto_resources() translates resources into FlyteIDL requests and limits, separating tuple request/limit values and omitting empty resources. flyte.extend.pod_spec_from_resources() translates singular values into Kubernetes pod resource requests and limits; it rejects tuple/list values. The Ray plugin uses this conversion for its head and worker pod templates.

Retries and caching

Pass an integer retry count or a RetryStrategy to @env.task. TaskTemplate.__post_init__ converts an integer to RetryStrategy(count=...); the current RetryStrategy contains only count, not backoff settings.

import random

import flyte

env = flyte.TaskEnvironment(name="caching-retries-example")


@env.task(retries=10)
def flaky_lookup(user_id: int) -> str:
    if random.random() < 0.5:
        raise RuntimeError(f"Transient upstream error for user {user_id}")
    return f"user-{user_id}"

flyte.errors.NonRecoverableError bypasses configured retries. In the repository example, this task has retries=3 but negative input fails on the first attempt:

import flyte
import flyte.errors

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


@env.task(retries=3)
async def validate_and_process(x: int) -> str:
    if x < 0:
        raise flyte.errors.NonRecoverableError(
            f"Input x={x} is negative. Negative inputs are always invalid — retrying would not help."
        )
    return f"processed({x})"

Cache requests are "auto", "override", or "disable", or a Cache object. They are accepted at environment, decorator, and override levels and are converted to Cache instances by TaskTemplate.__post_init__. The caching example keeps a transient lookup retryable and caches the downstream computation:

@env.task(cache="auto")
def cached_compute(user: str) -> dict[str, str]:
    return {
        "user": user,
        "email": f"{user}@example.com",
        "phone": "+1234567890",
    }


@env.task
def main(user_id: int) -> dict[str, str]:
    return cached_compute(user=flaky_lookup(user_id=user_id))

Use cache="disable" for work that must not reuse a result. timeout accepts an integer number of seconds or a Timeout; Timeout represents maximum runtime and optional maximum queued time.

Reuse warm containers

flyte.ReusePolicy keeps a pool of warm containers for environments where startup is expensive:

env = flyte.TaskEnvironment(
    name="fast_env",
    reusable=flyte.ReusePolicy(replicas=(1, 3), concurrency=2),
)

An integer replicas value normalizes to (n, n). A tuple supplies minimum and maximum replicas. idle_ttl shuts down the whole environment after it is idle; scaledown_ttl delays removal of an individual idle replica. Integer TTLs become timedelta values, and both TTLs require at least 30 seconds. Capacity is max_replicas * concurrency. Concurrency above one is rejected for synchronous task functions. A one-replica/one-concurrency policy logs a starvation warning.

Reuse shares the Python process and environment resources across invocations. Consequently, override() rejects resource, environment-variable, or secret changes while reuse remains active. Turn reuse off in the same call when changing resources:

new_task = oomer_with_reuse.override(
    reusable="off",
    resources=flyte.Resources(),
)
new_task.resources.cpu = 2
new_task.resources.memory = "500Mi"

Reusable environments cannot combine plugin_config with reuse, and TaskEnvironment.task() rejects a per-task pod template for a reusable environment.

Runtime behavior of a decorated task

A decorated function is an AsyncFunctionTaskTemplate, a concrete flyte.extend.TaskTemplate. TaskTemplate.__call__() calls forward() outside a Flyte task context; AsyncFunctionTaskTemplate.forward() calls the wrapped function directly. Inside a task context, __call__() submits through the controller. Nested calls inside flyte.trace raise TraceDoesNotAllowNestedTasksError; use .forward(...) for intentional ordinary-Python execution in a trace.

AsyncFunctionTaskTemplate.execute() runs the lifecycle in this order:

pre(...) -> wrapped async function (or run_sync_with_loop for def functions) -> post(result)

The template also exposes the function's source_file, derives json_schema from its NativeInterface, and builds container arguments containing input/output paths, version, raw/checkpoint paths, run/action names, image-cache data, code-bundle arguments, and a resolver when a pickle code bundle is unavailable.

Extend tasks with flyte.extend

Use flyte.extend to subclass AsyncFunctionTaskTemplate, register plugin task classes, serialize resources, or provide image-build integrations. The Ray plugin demonstrates the task-template extension point:

from dataclasses import dataclass
from typing import Any, Dict

import flyte
from flyte.extend import AsyncFunctionTaskTemplate


@dataclass(kw_only=True)
class RayFunctionTask(AsyncFunctionTaskTemplate):
    task_type: str = "ray"
    plugin_config: RayJobConfig
    debuggable: bool = True

    async def pre(self, *args, **kwargs) -> Dict[str, Any]:
        init_params = {"address": self.plugin_config.address}
        if flyte.ctx().is_in_cluster():
            init_params["runtime_env"] = {
                "working_dir": os.getcwd(),
                "excludes": ["script_mode.tar.gz", "fast*.tar.gz", ".python_history", ".code-server"],
            }
        if not ray.is_initialized():
            ray.init(**init_params)
        return {}

The actual Ray class implements custom_config() as well: it calls pod_spec_from_resources() for head/worker requests and limits, converts the result with PodTemplate.to_k8s_pod(), builds a FlyteIDL RayJob, and returns MessageToDict(ray_job). It registers with TaskPluginRegistry for the concrete RayJobConfig type. TaskEnvironment.task() performs the matching lookup with type(self.plugin_config), so registration must occur before decoration and is exact-type keyed.

Provide the task image

Image cannot be constructed directly. Start with a from_* constructor and chain immutable-style with_* methods:

image = (
    flyte.Image.from_debian_base(python_version=(3, 12))
    .with_pip_packages("pandas", "scikit-learn")
    .with_apt_packages("curl", "git")
)

env = flyte.TaskEnvironment(name="data_env", image=image)

Image.from_base() starts from an existing URI; from_dockerfile() starts from a Dockerfile; from_ref_name() uses a configured image reference. Dockerfile-based images cannot receive additional layers. Unmodified Flyte default images skip existence checks, while images tagged latest are always treated as needing a build.

Initialization supplies the process-global settings used by task execution and image building. Call flyte.init() or flyte.init_from_config() to configure root_dir, named images, and the default image_builder ("local"). flyte.extend.ImageBuildEngine.build() accepts builder="local" or builder="remote" for a per-build choice. The CI example clones a base image with registry/name metadata, sets its tag, and calls the engine; actual builds require registry access and the corresponding local or remote tooling.