Skip to main content

Infrastructure & Security

Infrastructure and security configuration in this SDK is managed primarily through the TaskEnvironment, which serves as the central hub for defining container images, compute resources, and sensitive credentials.

Task Environments​

The TaskEnvironment class (found in src/flyte/_task_environment.py) is the primary entry point for configuring infrastructure. It sets defaults for all tasks defined within it.

import flyte

env = flyte.TaskEnvironment(
    name="my-environment",
    image=flyte.Image.from_debian_base(python="3.12"),
    resources=flyte.Resources(cpu="1", memory="1Gi"),
)

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

Configuration Hierarchy​

Infrastructure settings can be defined at three levels, with more specific levels overriding more general ones:

  1. TaskEnvironment: Sets defaults for all tasks in the environment.
  2. @env.task decorator: Overrides settings for a specific task (e.g., cache, interruptible).
  3. task.override(): Overrides settings at invocation time (e.g., resources, env_vars).

Note: When reusable is enabled on the environment, resources, env_vars, and secrets can only be overridden via task.override() if reusable="off" is passed in the same call.

Compute Resources​

The Resources class (in src/flyte/_resources.py) allows you to specify CPU, memory, disk, and accelerator requirements.

from flyte import Resources, TaskEnvironment

env = TaskEnvironment(
    name="compute-env",
    resources=Resources(
        cpu="2",             # 2 cores
        memory="4Gi",        # 4 GiB RAM
        disk="10Gi",         # 10 GiB ephemeral storage
        shm="auto",          # Automatically set max shared memory
        gpu="A100 80G:1"     # 1 NVIDIA A100 80GB GPU
    )
)

Resource Formats​

  • CPU: Supports integers (1), floats (0.5), Kubernetes-style strings ("500m"), or tuples for request/limit ranges ((1, 2)).
  • Memory: Supports strings with units ("1Gi", "512Mi") or tuples (("1Gi", "2Gi")).
  • Shared Memory (SHM): Can be a specific size ("2Gi") or "auto" to use the maximum available on the node.

Accelerators (GPU, TPU, Neuron)​

For advanced accelerator configurations, use the helper functions provided in flyte.resources:

from flyte.resources import GPU, TPU, Neuron

# Specific GPU with MIG partitioning
res_gpu = Resources(gpu=GPU(device="A100", quantity=1, partition="1g.5gb"))

# Google Cloud TPU
res_tpu = Resources(gpu=TPU(device="V5P", partition="2x2x1"))

# AWS Inferentia/Trainium
res_neuron = Resources(gpu=Neuron(device="Inf2"))

Container Images​

The Image class (in src/flyte/_image.py) provides a fluent API for building container images. You start with a base constructor and chain with_* methods to add layers.

Building from a Base Image​

from flyte import Image

image = (
    Image.from_debian_base(python="3.12")
    .with_apt_packages("vim", "git")
    .with_pip_packages("pandas", "numpy")
    .with_env_vars({"MY_VAR": "value"})
)

Building from a UV Script​

If you are using a script with inline uv metadata, you can build an image directly from it:

# From examples/image/uv_image.py
image = Image.from_uv_script(__file__, name="my-app", registry="ghcr.io/my-org")

Local Development​

When developing locally, use with_local_v2() to include your local source code in the image build context:

image = Image.from_debian_base().with_local_v2()

Important Constraints​

  • Hashability: Image layers must be hashable for caching. Avoid passing lists to with_* methods; use tuples or multiple arguments instead.
  • Registry: The default registry is ghcr.io/flyteorg. If your endpoint contains localhost, it defaults to localhost:30000.

Security and Secrets​

The Secret class (in src/flyte/_secret.py) manages sensitive data injection. Secrets can be injected as environment variables or mounted as files.

Environment Variables​

By default, secrets are injected as environment variables. The SDK automatically formats the name to be uppercase with underscores.

from flyte import Secret, TaskEnvironment

# Injected as AWS_ACCESS_KEY_ID
aws_secret = Secret(group="aws", key="id", as_env_var="AWS_ACCESS_KEY_ID")

env = TaskEnvironment(name="secure-env", secrets=[aws_secret])

@env.task
async def secure_task():
    import os
    key = os.environ["AWS_ACCESS_KEY_ID"]

File Mounts​

Secrets can also be mounted as files. Currently, the SDK only supports mounting to the /etc/flyte/secrets directory.

import pathlib
from flyte import Secret, TaskEnvironment

SECRET_PATH = "/etc/flyte/secrets"

env = TaskEnvironment(
    name="file-secret-env",
    secrets=Secret(
        group="my-group", 
        key="my-key", 
        mount=pathlib.Path(SECRET_PATH)
    )
)

@env.task
def read_secret() -> str:
    # Path format: /etc/flyte/secrets/<group>/<key_lowercase>
    path = pathlib.Path(f"{SECRET_PATH}/my-group/my-key")
    return path.read_text()

Build-time Secrets​

Secrets can also be used during the image build process (e.g., for private pip indices):

image = Image.from_debian_base().with_pip_packages(
    "private-pkg",
    secret_mounts=[Secret(key="pip-conf", as_env_var="PIP_CONFIG_FILE")]
)