Overview

The Flyte SDK is an asynchronous, type-safe Python SDK for building, orchestrating, and serving scalable data pipelines, machine learning workloads, and compound AI systems. It turns standard Python functions into containerized, distributed execution units with built-in support for resource management, container reuse, dynamic fanout, and interactive serving.

Why Flyte SDK?

Modern data and AI workloads suffer from an impedance mismatch between local development and distributed cloud execution:

• Infrastructure overhead: Developers often rewrite local Python code into bespoke orchestrator DSLs or complex container manifests to get GPU acceleration, memory limits, and distributed scaling.
• Slow iteration cycles: Traditional orchestrators incur cold-start image pulls and container initialization overhead for every sub-step.
• Complex dependency boundaries: Mixing heterogeneous workloads (e.g., lightweight CPU data extraction and heavy GPU model inference) in a single workflow usually requires separate codebases or fragile multi-stage builds.
• Disconnected serving and batching: Running offline batch workflows and real-time interactive microservices typically requires maintaining two disparate infrastructure stacks.

The Flyte SDK solves these challenges by treating Python functions and environments as first-class primitives. You write standard asynchronous or synchronous Python, annotate environment requirements (images, hardware accelerators, secrets, caching), and execute locally or seamlessly trigger remote execution on Flyte clusters.

Core Concepts & Mental Model

Think of Flyte as an Environment-First Execution Graph: instead of defining static DAG configuration files, you define target environments and compose Python callables inside them.

                  ┌─────────────────────────────────────┐
                  │           TaskEnvironment           │
                  │   Image, Caching, Secrets, Reuse    │
                  └──────────────────┬──────────────────┘
                                     │ binds to
                                     ▼
                  ┌─────────────────────────────────────┐
                  │              @env.task              │
                  │     Standard Python / Async Coro    │
                  └──────────────────┬──────────────────┘
                                     │
           ┌─────────────────────────┴─────────────────────────┐
           ▼                                                   ▼
┌──────────────────────┐                           ┌──────────────────────┐
│  Distributed Fanout  │                           │  Interactive Serving │
│ asyncio / flyte.map  │                           │   flyte.serve(...)   │
└──────────────────────┘                           └──────────────────────┘

• TaskEnvironment: Defines the runtime container boundaries, base Docker images, caching strategies, and secrets for a group of tasks.
• @env.task: Decorates Python functions or coroutines to register them as managed units of execution. Tasks can call other tasks dynamically across different container environments.
• flyte.io (Managed I/O Types): Handles data flow across container boundaries without manual upload/download steps using native types like File, Dir, and DataFrame.
• flyte.run & flyte.deploy: Enables launching runs directly from Python scripts, CLI tools, or deploying versioned environments and tasks to a Flyte cluster.
• flyte.serve: Extends execution to long-running interactive applications and microservices.

Architecture & How It Works

1. Environment Binding: You declare one or more TaskEnvironment objects specifying runtime images, secrets, and environment settings.
2. Task Composition: Tasks are written as standard Python functions (def or async def). Sub-tasks can be awaited directly, mapped in parallel using flyte.map(), or fanned out using standard asyncio.gather().
3. Type Serialization & Offloading: Flyte's type engine inspects type annotations (Pydantic models, dataframes, primitives, files). Large objects and files are transparently serialized, hashed, and offloaded to object storage.
4. Execution & Container Reuse: During execution (local or remote), Flyte spins up containers matching the environment definition. With ReusePolicy, containers stay warm between task executions to minimize startup latency.
5. Observability & Tracing: Every execution step captures action IDs, logs, metrics, custom reports, and visual artifacts that stream live to the terminal or web console.

Code Examples

1. Basic Task Definition and Remote Execution

Tasks are regular functions attached to a TaskEnvironment.

import flyte

# Define execution environment
env = flyte.TaskEnvironment(name="data-prep")


@env.task
def compute_metrics(values: list[float]) -> dict[str, float]:
    mean_val = sum(values) / len(values)
    return {"mean": mean_val, "count": float(len(values))}


if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(compute_metrics, values=[10.5, 20.0, 30.5])
    print(f"Run URL: {run.url}")
    run.wait()

2. Multi-Environment Workflows and Async Fanout

Tasks in different environments can call each other seamlessly using standard asyncio.

import asyncio
import flyte

worker_env = flyte.TaskEnvironment(name="worker-env")
driver_env = flyte.TaskEnvironment(name="driver-env")


@worker_env.task
async def score_sample(sample_id: int) -> float:
    return sample_id * 0.42


@driver_env.task
async def driver_pipeline(sample_ids: list[int]) -> list[float]:
    with flyte.group("scoring-group"):
        scores = await asyncio.gather(*(score_sample(sid) for sid in sample_ids))
    return list(scores)

3. File and Artifact I/O

Use flyte.io.File and flyte.io.Dir to stream artifacts between remote tasks.

import flyte
from flyte.io import File

env = flyte.TaskEnvironment(name="io-pipeline")


@env.task
async def generate_report(output_file: File) -> File:
    async with output_file.open("w") as f:
        await f.write("Execution Summary: Success")
    return output_file


@env.task
async def read_report(report_file: File) -> str:
    async with report_file.open("r") as f:
        return await f.read()

4. Parallel Task Mapping

Run batch operations efficiently with flyte.map.

import flyte

env = flyte.TaskEnvironment(name="batch-pipeline")


@env.task
def transform(x: int) -> int:
    return x * 2 + 1


@env.task
def run_batch(items: list[int]) -> list[int]:
    return list(map(transform, items))

When to Use vs. When Not to Use

When to Use

• Heterogeneous Workflows: When distinct steps in your pipeline require drastically different environments or dependencies.
• Compound AI Systems: Building AI pipelines that combine batch data processing, LLM generation, sandboxed code execution, and model serving.
• Containerized Reproducibility: You need strict image and dependency isolation across tasks without maintaining raw Kubernetes pod manifests.
• Performance & Warm Containers: Pipelines where container initialization latency must be mitigated via warm container reuse.

When Not to Use

• Simple Single-Node Scripts: Small one-off scripts where local execution in a single virtual environment is sufficient.
• Sub-Millisecond Stream Processing: Ultra-low-latency event streaming (e.g. sub-millisecond trading or Kafka stream joins) is better served by streaming engines like Flink.
• In-Process Multiprocessing Only: When all compute fits on one machine and doesn't require distributed orchestration or external artifact tracking.

Integration & Stack Compatibility

• Python Ecosystem: Fully compatible with Python 3.10+ and integrates natively with asyncio, Pydantic, pandas, pyarrow, polars, and fsspec.
• Infrastructure & Compute: Runs on Kubernetes, with specialized hardware accelerators supported out-of-the-box (flyte.GPU, flyte.TPU, flyte.AMD_GPU, flyte.Neuron, and flyte.HABANA_GAUDI).
• Data & Object Storage: Integrated with obstore and fsspec for S3, GCS, Azure Blob, and local file storage.
• AI & ML Plugins: Extensible plugin architecture with first-party support for Ray, Spark, Dask, PyTorch, MLflow, Weights & Biases (W&B), Pandera, Papermill, Anthropic, and OpenAI.
• Packaging: Native support for uv, pip, and Docker/OCI container builds.

Getting Started Pointers

• Declare dependencies and build container environments using Define Task Environments.
• Structure compute pipelines using Tasks and Environments and execute them with Remote Execution and Inspection.
• Offload and stream large datasets across tasks with Data I/O.
• Set up warm execution pods and reduce startup overhead using ReusePolicy.
• Deploy interactive webhooks and inference endpoints via Serving Applications.

Limitations & Assumptions

• Python Version: Requires Python >= 3.10.
• Container Runtime: Distributed execution requires container-based execution backends (e.g., Docker daemon locally, or a Flyte/Kubernetes cluster remotely).
• Type Annotations: Flyte relies heavily on Python type hints for interface validation and automatic serialization/deserialization. Untyped arguments fallback to pickle where possible but limit engine optimizations.

Frequently Asked Questions

How does Flyte SDK differ from legacy Flytekit?

Flyte SDK (Flyte 2) introduces a streamlined TaskEnvironment model, native asyncio support, unified task and microservice authoring, warm container reuse, and faster import times with modular plugins.

Can I test my tasks locally without a Flyte cluster?

Yes. Tasks decorated with @env.task remain valid Python callables. You can execute them directly in unit tests or run them locally without initializing a remote session.

How does container reuse work?

By configuring ReusePolicy on a TaskEnvironment, Flyte keeps container instances alive after a task completes, allowing subsequent tasks with matching environments to execute immediately without re-pulling images or re-initializing dependencies.

Do I need to build Docker images manually?

No. Flyte provides programmatic image building utilities (flyte.Image.from_debian_base(...), flyte.build_images()) that automatically assemble your dependencies and bundle local code.