Overrides

Most task configuration is set when you define a task: on the TaskEnvironment or in the @env.task decorator. But you often need to change some of that configuration for a single invocation of a task: give one call more memory, point it at a different secret, or bump its retries.

The task.override() method does exactly this. It returns a new task with the specified parameters changed, leaving the original task definition untouched, so you can invoke the overridden task in place of the original:

import flyte

env = flyte.TaskEnvironment(
    name="training",
    resources=flyte.Resources(cpu=1, memory="512Mi"),
)

@env.task
async def train(data: str) -> str:
    return f"trained on {data}"

@env.task
async def main() -> str:
    # Invoke train with its environment-level resources.
    baseline = await train("small.csv")

    # Invoke train with overridden resources for this call only.
    heavy = await train.override(
        resources=flyte.Resources(cpu="4", memory="24Gi"),
    )("large.csv")

    return heavy

The key idiom is task.override(...)(args): override() returns a callable task, which you then invoke with the task’s arguments. Note the two sets of parentheses.

What you can override

override() accepts the parameters that are settable at the task-invocation level:

For the full parameter interaction matrix showing which parameters can be set at which level, see Task configuration levels.

What you cannot override

name, image, docs, and the task’s interface (its input and output types) cannot be overridden. Attempting to override them raises an error. To run a task with a different image, define it in a separate TaskEnvironment (see Container images and Multiple environments).

Overriding when a task is reusable

When a task uses a reusable container (reusable is set), its resources, env_vars, and secrets come from the parent environment and cannot be overridden while reuse is active. The container is already running.

To override any of these on a reusable task, turn reuse off in the same override() call by passing reusable="off":

result = await my_task.override(
    reusable="off",
    resources=flyte.Resources(cpu="4", memory="8Gi"),
)(data)

Overriding secrets

Just as you can override resources per invocation, you can override the secrets injected into a task for a single call. This is useful when the same task needs different credentials depending on how it’s invoked: for example, calling an external API with a different key per tenant, or supplying a secret that the task’s environment doesn’t declare.

Pass secrets to override() exactly as you would to the TaskEnvironment: a secret key, a Secret object, or a list of either.

import flyte
from flyte import Secret

env = flyte.TaskEnvironment(
    name="model_calls",
    secrets=Secret("openai-key", as_env_var="LLM_API_KEY"),
)

@env.task
async def call_model(prompt: str) -> str:
    import os
    api_key = os.environ["LLM_API_KEY"]
    ...  # call the model using api_key
    return "response"

@env.task
async def main() -> str:
    # Use the environment's default secret.
    default = await call_model("hello")

    # Override the secret for this invocation, mounting a different
    # store key into the same LLM_API_KEY environment variable.
    alternate = await call_model.override(
        secrets=Secret("anthropic-key", as_env_var="LLM_API_KEY"),
    )("hello")

    return alternate

As with the environment-level secrets parameter, the secret is injected at runtime and accessed inside the task: typically as an environment variable via os.environ (or mounted as a file). See Secrets for how to create secrets and how they are injected, and remember that the overriding secrets value replaces the environment’s secrets for that invocation rather than adding to them.