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.

Publish the model as an artifact

Return the weights from the training task and wrap them with flyte.artifacts.new(). Every run registers a new version of the artifact:

train_artifact.py
                
                    
                
            
                
            

                
                    
                
            
import flyte
from flyte import artifacts
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(produces_artifacts=True)
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")

    weights = await Dir.from_local("model")
    return artifacts.new(
        weights,
        artifacts.Metadata.create_model_metadata(
            name="support-llm",
            framework="PyTorch",
            task="text-generation",
            serial_format="safetensors",
        ),
    )

The task must set produces_artifacts=True. Without it, no artifact is registered and the app has nothing to resolve. See Task outputs as artifacts.

Serve an LLM with vLLM

Pass the artifact to VLLMAppEnvironment as model_path. SGLangAppEnvironment takes the same argument:

serve_artifact_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.ArtifactValue(name="support-llm", type="directory"),
    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()

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, publish it the same way:

train_artifact.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(produces_artifacts=True)
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")

    weights = await Dir.from_local("classifier")
    return artifacts.new(
        weights,
        artifacts.Metadata.create_model_metadata(
            name="churn-classifier",
            framework="scikit-learn",
            task="classification",
            serial_format="joblib",
        ),
    )

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_artifact.py train_classifier

Then mount the artifact into a FastAPIAppEnvironment with a parameter:

serve_artifact_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.ArtifactValue(name="churn-classifier", type="directory"),
            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 artifact is downloaded into it, and a flyte.io.Dir artifact appears at it. An artifact used as an app parameter must store one of those two types. See Use artifacts in apps.

Pin and roll back a version

An app resolves the latest artifact version when you deploy it and keeps that version until you deploy again. A version published later doesn’t change a running app.

To roll back, deploy with an explicit version:

serve_artifact_vllm.py
                
def roll_back(version: str) -> None:
    pinned_env = app_env.clone_with(
        name=app_env.name,
        model_path=flyte.app.ArtifactValue(name="support-llm", type="directory", version=version),
    )
    flyte.deploy(pinned_env)
    App.get(name=pinned_env.name).activate()

List the versions with flyte get artifact support-llm. The artifact’s lineage view shows which app deployment serves each version and which run produced it.

Automate training through deployment

To retrain, check, and deploy in one step, with an accuracy gate or a person’s approval before release, use a factory. See End in an app.

Choose an app type

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