Flyte 2 is the durable runtime built for open source

@env.task(retries=3)
async def validate_and_process(x: int) -> str:
    if x < 0:
        raise flyte.errors.NonRecoverableError(
            f"Input x={x} is negative, retrying will not help."
        )
    return f"processed({x})"

@env.task(cache="auto")
async def expensive_computation(data: str) -> str:
    # This result will be cached and reused for identical inputs
    ...

import flyte

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

@env.task(retries=3)
async def plan_step(query: str) -> Plan:
    ...

@env.task(retries=3)
async def execute_step(action: Action) -> StepResult:
    ...

@env.task
async def agent_loop(query: str) -> str:
    plan = await plan_step(query)
    while not plan.is_done:
        result = await execute_step(plan.next_action)
        plan = await plan_step(query, prior_result=result)
    return plan.final_answer

@env.task
async def main(queries: list[str]) -> list[str]:
    # each query recovers independently
    return list(await flyte.map(agent_loop, queries))

from flyte.extras import DynamicBatcher

async def process(batch: list[dict]) -> list[str]:
    """Your batch processing function. Must return results in the same order as the input."""
    return [heavy_computation(item) for item in batch]

async with DynamicBatcher(
    process_fn=process,
    target_batch_cost=1000,   # cost budget per batch
    max_batch_size=64,        # hard cap on records per batch
    batch_timeout_s=0.05,     # max wait time before dispatching a partial batch
    max_queue_size=5_000,     # queue size for backpressure
) as batcher:
    futures = []
    for record in my_records:
        future = await batcher.submit(record, estimated_cost=10)
        futures.append(future)
    results = await asyncio.gather(*futures)

app = FastAPI(title="ML Model API")

# Define request/response models
class PredictionRequest(BaseModel):
    feature1: float
    feature2: float
    feature3: float

class PredictionResponse(BaseModel):
    prediction: float
    probability: float

# Load model (you would typically load this from storage)
model = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global model
    model_path = os.getenv("MODEL_PATH", "/app/models/model.joblib")
    # In production, load from your storage
    if os.path.exists(model_path):
        with open(model_path, "rb") as f:
            model = joblib.load(f)
    yield

@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    # Make prediction
    # prediction = model.predict([[request.feature1, request.feature2, request.feature3]])

    # Dummy prediction for demo
    prediction = 0.85
    probability = 0.92

    return PredictionResponse(
        prediction=prediction,
        probability=probability,
    )

env = FastAPIAppEnvironment(
    name="ml-model-api",
    app=app,
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "fastapi",
        "uvicorn",
        "scikit-learn",
        "pydantic",
        "joblib",
    ),
    parameters=[
        flyte.app.Parameter(
            name="model_file",
            value=flyte.io.File.from_existing_remote("s3://bucket/models/model.joblib"),
            mount="/app/models",
            env_var="MODEL_PATH",
        ),
    ],
    resources=flyte.Resources(cpu=2, memory="2Gi"),
    requires_auth=False,
)

scaling=flyte.app.Scaling(
    replicas=(min_replicas, max_replicas),
    scaledown_after=idle_ttl_seconds,
)

@env.task(interruptible=True)
async def data_processing(batch: Batch) -> Result:
    ...