Sample DAG orechestrator using Prefect

prefect_dag.py

from prefect import flow, task
from sklearn.datasets import load_diabetes
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
import pandas as pd


@task
def load_data():
    data = load_diabetes(as_frame=True)
    df = data.frame
    print(f"✅ Dataset loaded with shape: {df.shape}")
    return df


@task
def train_model(df: pd.DataFrame):
    X = df.drop(columns=["target"])
    y = df["target"]
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    model = LinearRegression()
    model.fit(X_train, y_train)
    preds = model.predict(X_test)
    mse = mean_squared_error(y_test, preds)
    print(f"📉 Model trained. MSE = {mse:.3f}")
    return mse


@task
def evaluate_performance(mse: float):
    if mse < 3000:
        print(f"🎯 Great performance! MSE = {mse:.2f}")
    else:
        print(f"⚠️ Needs improvement. MSE = {mse:.2f}")


@flow(name="Intro ML Orchestration Flow")
def diabetes_flow():
    df = load_data()
    mse = train_model(df)
    evaluate_performance(mse)


if __name__ == "__main__":
    diabetes_flow()