Skip to content

flyteorg/flyte-sdk

Repository files navigation

Important

Flyte 2 Devbox is now available!

Check out the guide here to get started.


Flyte 2 SDK

Reliably orchestrate ML pipelines, models, and agents at scale — in pure Python.

Version Python License Try in Browser Docs SDK Reference CLI Reference

Install

pip install flyte

Example

Create a file called flyte_intro.py with the following content:

import asyncio
import flyte

env = flyte.TaskEnvironment(
    name="hello_world",
    image=flyte.Image.from_debian_base(python_version=(3, 12)),
)

@env.task(retries=3, cache="auto")  # 👈 add retries and caching
async def predict(x: int) -> int:
    return 2 * x + 5

@env.task
async def main(data: list[int]) -> float:
    xs = await asyncio.gather(*(predict(x) for x in data))
    return sum(xs) / len(xs)
PythonFlyte CLI
python flyte_intro.py
flyte run flyte_intro.py main --data '[1,2,3]'
ℹ️ Synchronous Python

Flyte 2 also supports synchronous python, although async is recommended for more control over concurrency and parallelism.

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

@env.task
def main(data: list[int]) -> float:
    xs = list(flyte.map(predict, data))
    return sum(xs) / len(xs)

Serve a Model

# serving.py
from fastapi import FastAPI
import flyte
from flyte.app.extras import FastAPIAppEnvironment

app = FastAPI()
env = FastAPIAppEnvironment(
    name="my-model",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi", "uvicorn"
    ),
)

@app.get("/predict")
async def predict(x: float) -> dict:
    return {"result": x * 2 + 5}

if __name__ == "__main__":
    flyte.init_from_config()
    flyte.serve(env)
PythonFlyte CLI
python serving.py
flyte serve serving.py env

Local Development Experience

Install the TUI for a rich local development experience:

pip install flyte[tui]

Run flyte_intro.py on the TUI:

flyte run --tui --local flyte_intro.py main --data '[1,2,3]'

Flyte TUI

Flyte Devbox

Flyte Devbox is a docker- and k3s-based local development environment for Flyte. It allows you to run Flyte workflows and services locally.

flyte start devbox

Create the configuration file for the devbox:

flyte create config \
    --endpoint localhost:30080 \
    --project flytesnacks \
    --domain development \
    --builder local \
    --insecure

Run on the devbox:

flyte run flyte_intro.py main --data '[1,2,3]'

Flyte Start Devbox

Flyte Hello World

Rust Controller (experimental)

Developer Guide

The Rust controller is an alternative implementation of the remote controller written in Rust and exposed to Python via maturin / pyo3. Distributed as a separate flyte_controller_base wheel so the main SDK does not need to switch its build toolchain to rust/maturin. Keep important dependencies (notably flyteidl2) in lockstep between pyproject.toml, rs_controller/pyproject.toml, and rs_controller/Cargo.toml.

Installing the Rust controller

flyte_controller_base is an optional dependency — pip install flyte does not pull it in automatically. The default Flyte task image (as of #1083) bundles the Rust controller, so tasks running on the default image work out of the box. You only need the extra when:

  • running locally (e.g. flyte run --local examples/basics/hello_v2.py), or
  • bringing your own task image.
pip install flyte[rust-controller]

Running with the Rust controller

The Rust controller is gated behind an env var. Set it to 1 (also accepts true / yes):

_F_USE_RUST_CONTROLLER=1 python examples/basics/hello_v2.py

The driver propagates this env var to all sub-task pods, so both the driver and child actions use the Rust controller for that run.

The Rust controller is currently under rapid development and contains gaps: abort RPC on cancel, Code.ABORTED fast-fail, tunable retries / QPS, graceful stop(). See PR #675.

Dev iteration requires the local image builder. The flyte_controller_base wheel is not on PyPI until release, and the remote image builder installs all wheels in a layer at once, so it cannot resolve flyte_controller_base from a sibling layer. Use the local image builder while developing the Rust controller:

# .flyte/config.yaml
image:
  builder: local

Developing the Rust controller

One-time setup

Build the manylinux builder images. They are cached, so you only need to rebuild them when the build tooling itself changes:

cd rs_controller
make build-builders
cd ..

Iteration loop

After every Rust change, run the all-in-one dev target from the repo root:

REGISTRY=<your-registry> make dev-rs-dist

dev-rs-dist does four things:

  1. cd rs_controller && make build-wheels — build manylinux x86_64 + aarch64 wheels (use make build-wheel-local if you only need a macOS wheel for the driver).
  2. make dist — build the main flyte SDK wheel.
  3. uv run python maint_tools/build_default_image.py --registry $(REGISTRY) — build the default image with both wheels baked in and push it to your registry.
  4. uv pip install --find-links ./rs_controller/dist --no-index --force-reinstall --no-deps flyte_controller_base — refresh the wheel in your local venv so the driver picks up the new build.

After this, any flyte.TaskEnvironment that does not pass an explicit image= will resolve to the default debian image and automatically have the Rust wheel layered in. If you do pass an explicit image=, the auto-bake is skipped; in that case, chain .with_local_rs_controller() onto the image to bake the Rust wheel manually.

If you only changed Python (not Rust), you can skip the wheel rebuild and just run make dist plus the rebuild image step. The Rust wheel is reused.

Build configuration summary

The Rust crate ships with two cargo features so the same project can produce a Rust rlib and a Python extension wheel:

[features]
default = ["pyo3/auto-initialize"]            # Rust crate users; links libpython
extension-module = ["pyo3/extension-module"]  # Python wheels; no libpython linking

[lib]
crate-type = ["rlib", "cdylib"]               # Both Rust and Python usage
  • pyo3/auto-initialize embeds Python into Rust (works locally on macOS, fails inside the manylinux builder because libpython is unavailable there).
  • pyo3/extension-module extends Python with Rust (must not link libpython for portable wheels).

So local cargo run --bin <name> uses default features, and the manylinux builder explicitly disables defaults and turns on extension-module:

# rs_controller/pyproject.toml
[tool.maturin]
no-default-features = true
features = ["extension-module"]

Learn More

License

Flyte 2 is licensed under the Apache 2.0 License.

About

Type-safe, distributed orchestration of agents, ML pipelines, and real-time inference — in pure Python with async/await.

Topics

Resources

License

Contributing

Stars

Watchers

Forks

Packages

 
 
 

Contributors

Languages