Configure Task-Level Settings

Flyte tasks offer a flexible configuration hierarchy, allowing you to define settings at different scopes. This enables you to establish global defaults for an environment, override them for specific tasks, or even dynamically adjust them at invocation time. The flyte-sdk provides three primary levels for task configuration:

1. TaskEnvironment: Sets default configurations for all tasks within that environment.
2. @env.task decorator: Overrides settings for individual tasks defined within a TaskEnvironment.
3. task.override(): Provides dynamic, invocation-specific overrides for a task.

More specific levels always take precedence over more general ones.

TaskEnvironment: Establishing Global Defaults

The flyte.TaskEnvironment class serves as the foundational layer for defining shared settings that apply to a collection of tasks. When you initialize a TaskEnvironment, you can specify common parameters such as the Docker image, compute resources, container reuse policy, caching behavior, and more. Any task subsequently defined using the @env.task decorator will inherit these settings by default.

For example, to define an environment with a specific image and resource allocation:

Unverified example

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

import flyte
from flyte import Resources

my_image = flyte.Image.from_debian_base(python="3.12").with_pip_packages("pandas")

env = flyte.TaskEnvironment(
    name="my_env",
    image=my_image,
    resources=Resources(1, "1Gi"),
    cache="auto",
    interruptible=True,
)

@env.task
def my_task_one(a: int, b: int) -> int:
    # This task inherits image, resources, cache, and interruptible from 'env'
    return a + b

@env.task
def my_task_two(x: str) -> str:
    # This task also inherits from 'env'
    return x.upper()

The TaskEnvironment can configure: name, image, depends_on, description, plugin_config, resources, env_vars, secrets, cache, reusable, queue, pod_template, and interruptible.

@env.task: Overriding Per-Task Settings

While TaskEnvironment provides sensible defaults, you often need to customize settings for individual tasks. The @env.task decorator allows you to specify parameters that override the TaskEnvironment's defaults for that particular task.

You can override settings like cache, retries, timeout, pod_template, queue, and interruptible directly within the decorator.

Unverified example

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

import flyte
from flyte import Resources, Timeout
from datetime import timedelta

my_image = flyte.Image.from_debian_base(python="3.12")

env = flyte.TaskEnvironment(
    name="my_env",
    image=my_image,
    resources=Resources(1, "1Gi"),
    cache="disable", # Environment default is disabled
    interruptible=False, # Environment default is not interruptible
)

@env.task(
    cache="auto", # Override: enable caching for this task
    retries=3, # Set retries for this task
    timeout=Timeout(300), # Set a 5-minute runtime timeout
    interruptible=True, # Override: make this task interruptible
    queue="high-priority", # Assign to a specific queue
)
def my_special_task(data: list) -> int:
    # This task uses its own cache, retries, timeout, interruptible, and queue settings
    return len(data)

task.override(): Dynamic Invocation-Time Overrides

For scenarios where you need to adjust task settings dynamically at the point of invocation, flyte-sdk provides the task.override() method. This is particularly useful for ad-hoc adjustments or when you want to run the same task with different configurations within a workflow.

task.override() returns a new TaskTemplate instance with the specified parameters changed, leaving the original task definition untouched. This allows you to create variations of a task without modifying its base definition.

Parameters that can be overridden at invocation time include short_name, resources, cache, retries, timeout, reusable, env_vars, secrets, max_inline_io_bytes, pod_template, queue, and interruptible.

Unverified example

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

import flyte
from flyte import Resources, Timeout
from datetime import timedelta

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

@env.task(retries=1)
def process_data(value: int) -> int:
    return value * 2

# Run with default retries (1)
result_default = process_data(10)

# Override retries to 5 for a specific invocation
result_with_more_retries = process_data.override(retries=5)(20)

# Override timeout for a specific invocation
result_with_timeout = process_data.override(timeout=Timeout(60))(30)

# Override resources for a specific invocation
result_with_resources = process_data.override(resources=Resources(2, "4Gi"))(40)

Interaction with reusable

When a TaskEnvironment is configured with a flyte.ReusePolicy, overriding resources, env_vars, or secrets via task.override() requires explicitly setting reusable="off" in the same call. This ensures that resource or environment variable changes do not conflict with the reusable container's fixed configuration.

Unverified example

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

import flyte
from flyte import Resources, ReusePolicy

reusable_env = flyte.TaskEnvironment(
    name="reusable_env",
    reusable=ReusePolicy(replicas=(1, 3)),
    resources=Resources(1, "1Gi"),
)

@reusable_env.task
def my_reusable_task(x: int) -> int:
    return x + 1

# This will raise a ValueError because resources cannot be overridden when reusable is set
# my_reusable_task.override(resources=Resources(cpu="2"))(5)

# To override resources, you must disable reusability for that invocation
result = my_reusable_task.override(resources=Resources(2, None), reusable="off")(5)

Specific Configuration Aspects

Timeouts with flyte.Timeout

The flyte.Timeout class allows you to define precise time limits for task execution. This includes max_runtime (the maximum time a task can actively run) and max_queued_time (the maximum time a task can wait in a queue before starting).

Unverified example

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

from flyte import Timeout
from datetime import timedelta

# Define a timeout with a 5-minute max runtime and 10-minute max queued time
timeout_config = Timeout(300, 600)

# You can also use integers for seconds
short_timeout = Timeout(300) # 300 seconds = 5 minutes

You can apply Timeout objects at the @env.task decorator level or dynamically with task.override():

Unverified example

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

import flyte
from flyte import Timeout
from datetime import timedelta

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

@env.task(timeout=Timeout(3600))
def long_running_task() -> str:
    return "Completed within an hour"

@env.task
def flexible_task(duration_minutes: int) -> str:
    return f"Processed in {duration_minutes} minutes"

# Override the timeout for a specific call
flexible_task.override(timeout=Timeout(duration_minutes * 60))(15)

Container Reuse with flyte.ReusePolicy

For tasks with expensive initialization, flyte.ReusePolicy enables the reuse of containers across multiple task invocations, significantly reducing cold-start overhead. You configure ReusePolicy with the number of replicas to maintain and idle_ttl for how long idle containers are kept.

replicas can be a fixed integer or a (min, max) tuple for auto-scaling. idle_ttl and scaledown_ttl can be int (seconds) or timedelta.

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: 'flyte.ReusePolicy' does not accept keyword(s): concurrency, idle_ttl, replicas

import flyte
from flyte import ReusePolicy
from datetime import timedelta

# Configure a reusable environment with auto-scaling replicas and a 30-minute idle TTL
reusable_env = flyte.TaskEnvironment(
    name="reusable_processing_env",
    reusable=ReusePolicy(
        replicas=(1, 3), # Auto-scale between 1 and 3 replicas
        idle_ttl=timedelta(minutes=30),
        concurrency=2, # Allow 2 concurrent tasks per replica
    ),
)

@reusable_env.task
def process_batch(data: list) -> int:
    # This task will run in a potentially reused container
    return sum(data)

Custom Kubernetes Pods with flyte.PodTemplate

When you need fine-grained control over the underlying Kubernetes Pod that executes your task, flyte.PodTemplate allows you to specify a custom V1PodSpec. This is useful for advanced configurations like adding sidecar containers, mounting volumes, setting specific node selectors, or injecting environment variables directly into the pod.

import flyte
from flyte import PodTemplate
from kubernetes.client import V1PodSpec, V1Container, V1EnvVar

# Define a PodTemplate to inject an environment variable into the primary container
custom_pod_template = PodTemplate(
    pod_spec=V1PodSpec(
        containers=[
            V1Container(
                name="primary", # The name of your task's main container
                env=[V1EnvVar(name="CUSTOM_SETTING", value="my_value")]
            )
        ],
    ),
    primary_container_name="primary", # Ensure this matches the container name in pod_spec
)

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

@env.task(pod_template=custom_pod_template)
def task_with_custom_env() -> str:
    import os
    return os.getenv("CUSTOM_SETTING", "Not Set")

# You can also override the pod template at invocation time
# This example is from tests/user_api/test_override.py
# oomer.override(pod_template=pod, short_name="first_override")

Caching

Caching can significantly speed up workflows by reusing results from previous task executions. The cache parameter can be set at the TaskEnvironment, @env.task, or task.override() levels. It accepts string values like "auto", "override", "disable", or a flyte.Cache object.

import flyte

env = flyte.TaskEnvironment(
    name="my_env",
    cache="disable", # Environment-level caching is disabled by default
)

@env.task(cache="auto") # Enable caching for this specific task
def cached_computation(x: int) -> int:
    return x * x

@env.task
def non_cached_task(y: int) -> int:
    return y + 1

# Temporarily enable caching for a specific invocation
result = non_cached_task.override(cache="auto")(10)

Retries

Tasks can be configured to automatically retry upon failure. The retries parameter accepts an integer for the number of retries or a flyte.RetryStrategy object for more advanced retry logic. This can be set at the @env.task or task.override() levels.

import flyte

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

@env.task(retries=3) # Retry this task up to 3 times on failure
def flaky_operation(attempt: int) -> str:
    if attempt < 2:
        raise ValueError("Simulated failure")
    return "Success"

@env.task
def critical_task() -> str:
    return "Always succeeds"

# Override retries for a specific invocation
result = critical_task.override(retries=5)()

Queues

Tasks can be assigned to specific queues, which often correspond to different compute resources or priorities in your Flyte deployment. The queue parameter accepts a string and can be set at the TaskEnvironment, @env.task, or task.override() levels.

import flyte

env_high_priority = flyte.TaskEnvironment(
    name="high_priority_env",
    queue="high-priority-queue", # All tasks in this env go to high-priority
)

env_default = flyte.TaskEnvironment(name="default_env")

@env_high_priority.task
def urgent_task() -> str:
    return "Runs on high-priority queue"

@env_default.task(queue="low-priority-queue") # Override env default to low-priority
def background_task() -> str:
    return "Runs on low-priority queue"

@env_default.task
def adhoc_task() -> str:
    return "Runs on default queue"

# Override queue for a specific invocation
result = adhoc_task.override(queue="urgent-adhoc-queue")()

Interruptibility

Tasks can be marked as interruptible, allowing them to run on spot or preemptible instances, which can be more cost-effective but may be interrupted. The interruptible parameter accepts a boolean and can be set at the TaskEnvironment, @env.task, or task.override() levels.

import flyte

env = flyte.TaskEnvironment(
    name="my_env",
    interruptible=True, # Environment tasks are interruptible by default
)

@env.task
def cost_optimized_task() -> str:
    return "Runs on a spot instance"

@env.task(interruptible=False) # Override: this task is not interruptible
def critical_non_interruptible_task() -> str:
    return "Runs on a dedicated instance"

@env.task
def flexible_interruptible_task() -> str:
    return "Can be interruptible"

# Override interruptibility for a specific invocation
result = flexible_interruptible_task.override(interruptible=False)()

Underlying Mechanism: TaskTemplate

Internally, all these configuration settings ultimately modify attributes of a flyte.extend.TaskTemplate object. When you define a task using @env.task, the decorator constructs a TaskTemplate instance, populating its fields with the inherited TaskEnvironment settings and any explicit overrides from the decorator. Similarly, task.override() creates a new TaskTemplate instance based on the original, applying the specified changes. While you typically interact with tasks through decorators and override() calls, understanding that TaskTemplate is the underlying data structure helps clarify how these configurations are managed within flyte-sdk.