Deploy trained models

A model trained in a task can go straight to an endpoint. The task saves the weights to object storage, and an app mounts them when it starts, so you never copy weights by hand or bake them into an image.

Return the model from a task

Save the weights in the training task and return them as a flyte.io.Dir:

train_output.py
                
import flyte
from flyte.io import Dir

env = flyte.TaskEnvironment(
    name="training",
    image=flyte.Image.from_debian_base().with_pip_packages(
        "torch", "transformers", "accelerate"
    ),
    resources=flyte.Resources(cpu="4", memory="32Gi", gpu="L40s:1", disk="100Gi"),
)

@env.task
async def finetune(base_model: str) -> Dir:
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(base_model)
    model = AutoModelForCausalLM.from_pretrained(base_model)

    # Add your training loop here.

    model.save_pretrained("model")
    tokenizer.save_pretrained("model")
    return await Dir.from_local("model")

Serve an LLM with vLLM

Pass the task’s output to VLLMAppEnvironment as model_path with flyte.app.RunOutput. SGLangAppEnvironment takes the same argument:

serve_output_vllm.py
                
import flyte
import flyte.app
from flyte.remote import App
from flyteplugins.vllm import VLLMAppEnvironment

app_env = VLLMAppEnvironment(
    name="support-llm",
    model_path=flyte.app.RunOutput(type="directory", task_name="training.finetune"),
    model_id="support-llm",
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1", disk="50Gi"),
    scaling=flyte.app.Scaling(replicas=(0, 1), scaledown_after=600),
)

if __name__ == "__main__":
    flyte.init_from_config()
    flyte.deploy(app_env)
    App.get(name=app_env.name).activate()

task_name is <environment name>.<task name>, and the app uses the latest successful run of that task when you deploy it. To serve the output of one specific run, pass run_name instead.

The directory must be in a format the engine can load, such as the output of save_pretrained(). A PEFT adapter saved on its own is not enough, so merge it into the base model first.

Call the deployed app with any OpenAI client. See Using the OpenAI-compatible API.

Serve any other model with FastAPI

For a model vLLM and SGLang don’t serve, such as a classifier, return it the same way:

train_output.py
                
classifier_env = flyte.TaskEnvironment(
    name="classifier",
    image=flyte.Image.from_debian_base().with_pip_packages("joblib", "scikit-learn"),
    resources=flyte.Resources(cpu="2", memory="4Gi"),
)

@classifier_env.task
async def train_classifier() -> Dir:
    import os

    import joblib
    from sklearn.datasets import make_classification
    from sklearn.linear_model import LogisticRegression

    features, labels = make_classification(n_samples=500, n_features=4, random_state=0)
    model = LogisticRegression().fit(features, labels)

    os.makedirs("classifier", exist_ok=True)
    joblib.dump(model, "classifier/model.joblib")
    return await Dir.from_local("classifier")

Run the task once on your cluster before you deploy the app, because the app resolves its model at deploy time. Running finetune doesn’t train the classifier:

flyte run train_output.py train_classifier

Then mount the task’s output into a FastAPIAppEnvironment with a parameter:

serve_output_fastapi.py
                
                    
                
            
                
            

                
                    
                
            
from contextlib import asynccontextmanager

import flyte
import flyte.app
import joblib
from fastapi import FastAPI
from flyte.app import Parameter
from flyte.app.extras import FastAPIAppEnvironment
from flyte.remote import App

MODEL_DIR = "/tmp/model"
state = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    state["model"] = joblib.load(f"{MODEL_DIR}/model.joblib")
    yield

app = FastAPI(lifespan=lifespan)

@app.post("/predict")
async def predict(features: list[float]) -> dict:
    return {"prediction": state["model"].predict([features]).tolist()[0]}

app_env = FastAPIAppEnvironment(
    name="churn-classifier",
    app=app,
    image=flyte.Image.from_debian_base().with_pip_packages("fastapi", "uvicorn", "joblib", "scikit-learn"),
    resources=flyte.Resources(cpu="2", memory="4Gi"),
    parameters=[
        Parameter(
            name="model",
            value=flyte.app.RunOutput(type="directory", task_name="classifier.train_classifier"),
            mount=MODEL_DIR,
        ),
    ],
)

if __name__ == "__main__":
    flyte.init_from_config()
    flyte.deploy(app_env)
    App.get(name=app_env.name).activate()

mount is always a directory. A flyte.io.File output is downloaded into it, and a flyte.io.Dir output appears at it. See Passing parameters into app environments.

Choose an app type

Model	App	Details
LLM, high throughput	`VLLMAppEnvironment`	vLLM endpoint
LLM, structured generation	`SGLangAppEnvironment`	SGLang endpoint
Anything else	`FastAPIAppEnvironment`	FastAPI app