DEV Community

Abdelrahman Adnan
Abdelrahman Adnan

Posted on

🚀 MLOps Zoomcamp 2025: Week 3 - Workflow Orchestration with Prefect 📊

#programming #python #learning #mlops

🔄 What is Workflow Orchestration?

Workflow orchestration is the process of coordinating, scheduling, and integrating various tasks in a data pipeline. For ML pipelines, this includes:

  • Data ingestion and preprocessing
  • Feature engineering
  • Model training and evaluation
  • Model deployment
  • Retraining schedules
  • Monitoring

Orchestration tools help manage complex dependencies, handle failures gracefully, and provide visibility into pipeline execution.

🧩 Prefect: A Modern Orchestration Tool for ML

Prefect is designed specifically for data-intensive applications with a Python-first approach. Unlike older orchestration tools, Prefect allows you to convert your existing Python code into production-ready data pipelines with minimal changes.

Why Prefect for MLOps?

  1. Pythonic: Works with your existing Python code
  2. Dynamic: DAGs can be created at runtime
  3. Resilient: Built-in error handling and recovery
  4. Observable: Comprehensive UI and monitoring
  5. Distributed: Can scale across multiple machines
  6. Modern: Actively maintained and feature-rich

🏗️ Prefect Architecture Explained

Prefect has a distributed architecture with distinct components:

┌─────────────────┐     ┌───────────────┐     ┌─────────────────┐
│                 │     │               │     │                 │
│  Prefect API    │◄────┤  Prefect UI   │     │  Storage        │
│  (Orchestrator) │     │  (Dashboard)  │     │  (S3, GCS, etc) │
│                 │     │               │     │                 │
└───────┬─────────┘     └───────────────┘     └────────┬────────┘
        │                                              │
        │                                              │
┌───────▼──────────────────────────────────────┐       │
│                                              │       │
│  Prefect Agents                              │◄──────┘
│  (Workers that execute flows)                │
│                                              │
└──────────────────────────────────────────────┘
Enter fullscreen mode Exit fullscreen mode

Core Components:

  1. Prefect API (Orion): The orchestration engine that:

    • Schedules and triggers flows
    • Tracks flow and task states
    • Stores execution history
    • Manages concurrency and dependencies

  2. Prefect UI: Web-based dashboard that provides:

    • Visual representation of flows and tasks
    • Execution history and logs
    • Performance metrics
    • Administrative controls

  3. Storage: Where flow code and artifacts are stored:

    • Local filesystem
    • Cloud object storage (S3, GCS, etc.)
    • Git repositories
    • Docker images

  4. Agents: Workers that:

    • Poll for scheduled flows
    • Pull flow code from storage
    • Execute flows in the specified environment
    • Report results back to the API

🧠 Prefect Core Concepts in Detail

1. Tasks 📋

Tasks are the atomic units of work in Prefect. They encapsulate individual operations and can be composed into larger workflows.

Task Definition: 

from prefect import task

@task(
    name="extract_data",  # Custom name for the task
    retries=3,            # Retry automatically on failure
    retry_delay_seconds=30,  # Wait between retries
    cache_key_fn=lambda context, **params: params["url"],  # Cache by URL parameter
    cache_expiration=timedelta(hours=12)  # Cache for 12 hours
)
def extract_data(url: str) -> pd.DataFrame:
    """Extract data from a given URL"""
    return pd.read_csv(url)
Enter fullscreen mode Exit fullscreen mode

Task Properties:

  • Name: Human-readable identifier for the task
  • Retries: Number of times to retry on failure
  • Timeout: Maximum execution time
  • Tags: Metadata for filtering and organization
  • Cache: Store and reuse task results
  • Result Storage: Where to store task outputs
  • Result Handlers: How to serialize/deserialize outputs

Task States:

Tasks can be in various states during execution:

  • Pending: Ready to run
  • Running: Currently executing
  • Completed: Successfully finished
  • Failed: Encountered an error
  • Retrying: Failed but will retry
  • Cancelled: Manually stopped

2. Flows 🌊

Flows are the main unit of work in Prefect. They coordinate the execution of tasks and manage dependencies.

Flow Definition: 

from prefect import flow
from prefect.task_runners import ConcurrentTaskRunner

@flow(
    name="Data Processing Pipeline",
    description="Extracts, transforms, and loads data",
    version="1.0.0",
    task_runner=ConcurrentTaskRunner(),  # Run tasks concurrently
    retries=2,  # Retry the entire flow if it fails
)
def process_data(date: str = None):
    """Main flow to process data"""
    # If no date provided, use today's date
    if date is None:
        date = datetime.today().strftime("%Y-%m-%d")

    # Call tasks, which automatically creates dependencies
    data = extract_data(f"data/data-{date}.csv")
    processed = transform_data(data)
    load_data(processed)

    return processed
Enter fullscreen mode Exit fullscreen mode

Flow Features:

  • Task Dependencies: Automatically determined from function calls
  • Error Handling: Custom exception handling
  • Subflows: Nested flows for better organization
  • Parameters: Runtime configuration
  • State Handlers: Custom logic on state changes
  • Logging: Structured logging for monitoring
  • Concurrency: Parallel task execution

Flow Execution:

Flows can be executed in various ways:

  • Directly calling the function
  • Using flow.serve() for real-time API
  • Through deployments for scheduled runs
  • Via the Prefect UI or API

3. Task Runners 🏃‍♂️

Task runners determine how tasks within a flow are executed.

Types of Task Runners: 

from prefect.task_runners import SequentialTaskRunner  # Default, run tasks in sequence
from prefect.task_runners import ConcurrentTaskRunner  # Run tasks concurrently
from prefect.task_runners import DaskTaskRunner  # Distributed execution with Dask
Enter fullscreen mode Exit fullscreen mode

Use Cases:

  • SequentialTaskRunner: Simple workflows, easy debugging
  • ConcurrentTaskRunner: Independent tasks that can run in parallel
  • DaskTaskRunner: Large-scale data processing, distributed computing

4. Deployments 🚢

Deployments make flows executable outside your local Python environment.

Creating a Deployment: 

from prefect.deployments import Deployment
from prefect.orion.schemas.schedules import CronSchedule
from prefect.infrastructure.process import Process

# Create deployment from a flow
deployment = Deployment.build_from_flow(
    flow=process_data,
    name="daily-data-processing",
    version="1",
    schedule=CronSchedule("0 0 * * *"),  # Daily at midnight
    infrastructure=Process(),  # Run as a local process
    tags=["production", "data-pipeline"],
    parameters={"date": None},  # Default parameters
)

# Apply deployment to the Prefect API
deployment.apply()
Enter fullscreen mode Exit fullscreen mode

Deployment Properties:

  • Name: Identifier for the deployment
  • Schedule: When to execute the flow
  • Infrastructure: Where to run the flow
  • Storage: Where to store the flow code
  • Parameters: Default parameters for the flow
  • Tags: For organizing and filtering deployments

5. Schedules ⏰

Schedules determine when flows should be executed automatically.

Types of Schedules: 

from prefect.orion.schemas.schedules import (
    CronSchedule,          # Based on cron expressions
    IntervalSchedule,      # Run at regular intervals
    RRuleSchedule          # Based on iCalendar recurrence rules
)

# Examples
cron_schedule = CronSchedule(cron="0 9 * * 1-5")  # Weekdays at 9 AM
interval_schedule = IntervalSchedule(interval=timedelta(hours=4))  # Every 4 hours
Enter fullscreen mode Exit fullscreen mode

6. Results and Persistence 💾

Prefect can store and track results from task and flow runs.

Persisting Results: 

@task(
    persist_result=True,  # Store the result
    result_serializer=JSONSerializer(),  # How to serialize the result
    result_storage=S3Bucket.load("my-bucket"),  # Where to store the result
)
def calculate_metrics(data):
    # Process data and return metrics
    return {"accuracy": 0.95, "precision": 0.92}
Enter fullscreen mode Exit fullscreen mode

🛠️ Building an ML Pipeline with Prefect: Complete Example

The following example demonstrates a complete ML pipeline for NYC taxi trip duration prediction:

import pandas as pd
import pickle
from datetime import datetime, timedelta
from sklearn.feature_extraction import DictVectorizer
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error

from prefect import task, flow
from prefect.task_runners import SequentialTaskRunner
from prefect.deployments import Deployment
from prefect.orion.schemas.schedules import CronSchedule
from prefect.logging import get_run_logger

# Data ingestion task with retries
@task(retries=3, retry_delay_seconds=5, name="download_taxi_data")
def get_data(date):
    """Download taxi data for given date range"""
    logger = get_run_logger()

    train_date = datetime.strptime(date, '%Y-%m-%d')
    val_date = train_date + timedelta(days=28)

    train_month = train_date.month
    val_month = val_date.month
    train_year = train_date.year
    val_year = val_date.year

    train_file = f"yellow_tripdata_{train_year}-{train_month:02d}.parquet"
    val_file = f"yellow_tripdata_{val_year}-{val_month:02d}.parquet"

    logger.info(f"Downloading training data: {train_file}")
    train_url = f"https://d37ci6vzurychx.cloudfront.net/trip-data/{train_file}"
    train_data = pd.read_parquet(train_url)

    logger.info(f"Downloading validation data: {val_file}")
    val_url = f"https://d37ci6vzurychx.cloudfront.net/trip-data/{val_file}"
    val_data = pd.read_parquet(val_url)

    return train_data, val_data

# Data preprocessing task
@task(name="prepare_taxi_features")
def prepare_features(df, categorical_features):
    """Prepare the features for training"""
    logger = get_run_logger()

    # Calculate trip duration
    df['duration'] = (df.tpep_dropoff_datetime - df.tpep_pickup_datetime).dt.total_seconds() / 60

    # Filter outliers
    logger.info(f"Initial shape: {df.shape}")
    df = df[(df.duration >= 1) & (df.duration <= 60)]
    logger.info(f"Filtered shape: {df.shape}")

    # Create features
    df['PU_DO'] = df['PULocationID'].astype(str) + '_' + df['DOLocationID'].astype(str)

    # Process categoricals
    df[categorical_features] = df[categorical_features].fillna(-1).astype('str')

    return df

# Model training task
@task(name="train_regression_model")
def train_model(df, categorical_features):
    """Train the model with the given data"""
    logger = get_run_logger()

    dicts = df[categorical_features].to_dict(orient='records')
    dv = DictVectorizer()
    X_train = dv.fit_transform(dicts)
    y_train = df['duration'].values

    logger.info(f"Training on {X_train.shape[0]} examples")

    lr = LinearRegression()
    lr.fit(X_train, y_train)

    return lr, dv

# Model evaluation task
@task(name="evaluate_model_performance")
def evaluate_model(model, dv, df, categorical_features):
    """Evaluate the model performance"""
    logger = get_run_logger()

    dicts = df[categorical_features].to_dict(orient='records')
    X_val = dv.transform(dicts)
    y_val = df['duration'].values

    y_pred = model.predict(X_val)
    rmse = mean_squared_error(y_val, y_pred, squared=False)

    logger.info(f"RMSE: {rmse:.3f}")
    return rmse

# Model artifacts storage task
@task(name="save_model_artifacts")
def save_model(model, dv, date):
    """Save the model and DictVectorizer"""
    logger = get_run_logger()

    # Create directory if it doesn't exist
    os.makedirs("models", exist_ok=True)

    with open(f'models/model-{date}.pkl', 'wb') as f:
        pickle.dump(model, f)

    with open(f'models/dv-{date}.pkl', 'wb') as f:
        pickle.dump(dv, f)

    logger.info(f"Model saved: models/model-{date}.pkl")
    return True

# Main flow that orchestrates all tasks
@flow(name="taxi-duration-prediction", task_runner=SequentialTaskRunner())
def main(date=None):
    """Main flow for taxi duration prediction model training"""
    if date is None:
        date = datetime.today().strftime('%Y-%m-%d')

    logger = get_run_logger()
    logger.info(f"Starting training pipeline for date: {date}")

    categorical_features = ['PULocationID', 'DOLocationID', 'PU_DO']

    # 1. Get data
    train_data, val_data = get_data(date)

    # 2. Process training data
    df_train = prepare_features(df=train_data, categorical_features=categorical_features)

    # 3. Train model
    model, dv = train_model(df=df_train, categorical_features=categorical_features)

    # 4. Process validation data
    df_val = prepare_features(df=val_data, categorical_features=categorical_features)

    # 5. Evaluate model
    rmse = evaluate_model(
        model=model,
        dv=dv,
        df=df_val,
        categorical_features=categorical_features
    )

    # 6. Save model
    save_model(model=model, dv=dv, date=date)

    logger.info(f"Pipeline completed successfully!")
    return rmse

# Define deployment for production
def create_deployment():
    return Deployment.build_from_flow(
        flow=main,
        name="nyc-taxi-monthly-training",
        schedule=CronSchedule("0 0 1 * *"),  # First day of each month at midnight
        tags=["production", "mlops", "taxi-duration"]
    )

# For local development/testing
if __name__ == "__main__":
    main("2023-01-01")
Enter fullscreen mode Exit fullscreen mode

📊 MLOps Workflow Integration

In the MLOps lifecycle, Prefect serves as the workflow orchestration layer that connects:

  1. Data Engineering: Extracting and preparing data for ML
  2. Experimentation: Running and tracking model experiments
  3. Continuous Training: Automating regular model retraining
  4. Deployment: Pushing models to production
  5. Monitoring: Tracking model performance over time

Integration with MLflow

Prefect works well with MLflow for experiment tracking:

@task
def train_with_mlflow(X_train, y_train, params):
    with mlflow.start_run():
        # Log parameters
        mlflow.log_params(params)

        # Train model
        model = RandomForestRegressor(**params)
        model.fit(X_train, y_train)

        # Log metrics
        train_rmse = mean_squared_error(y_train, model.predict(X_train), squared=False)
        mlflow.log_metric("train_rmse", train_rmse)

        # Log model
        mlflow.sklearn.log_model(model, "model")

        return model
Enter fullscreen mode Exit fullscreen mode

🚦 Prefect vs. Other Orchestration Tools

Prefect vs. Airflow

Feature Prefect Airflow
Programming Model Python-native, tasks as functions DAGs as configuration
Dynamic Workflows Yes, DAGs can change at runtime Limited, mostly static DAGs
Error Handling Rich, automated retries and custom handlers Basic retry policies
UI/UX Modern, task-oriented Traditional, DAG-oriented
Parametrization First-class support More complex implementation
Learning Curve Lower for Python developers Steeper
Community Growing Large, established

Prefect vs. Kubeflow

Feature Prefect Kubeflow
Focus General workflow orchestration Kubernetes-native ML pipelines
Deployment Multiple options (local, cloud, k8s) Primarily Kubernetes
Integration Python ecosystem Container-based components
Complexity Lower barrier to entry Higher complexity
ML Specific General purpose, adaptable to ML Built specifically for ML

🛫 Getting Started with Prefect in 5 Minutes

Installation 

pip install prefect
Enter fullscreen mode Exit fullscreen mode

Start the Prefect Server 

prefect server start
Enter fullscreen mode Exit fullscreen mode

Create a Simple Flow 

from prefect import task, flow

@task
def say_hello(name):
    return f"Hello, {name}!"

@flow
def hello_flow(name="world"):
    result = say_hello(name)
    print(result)
    return result

if __name__ == "__main__":
    hello_flow("MLOps Zoomcamp")
Enter fullscreen mode Exit fullscreen mode

Create a Deployment 

prefect deployment build hello_flow.py:hello_flow -n my-first-deployment -q default
prefect deployment apply hello_flow-deployment.yaml
Enter fullscreen mode Exit fullscreen mode

Start an Agent 

prefect agent start -q default
Enter fullscreen mode Exit fullscreen mode

🧪 Best Practices for ML Workflows with Prefect

  1. Task Granularity: Create tasks at the right level - not too fine-grained, not too coarse
  2. Error Boundaries: Place tasks boundaries around operations that might fail
  3. Parameterize Flows: Make flows configurable with parameters
  4. Logging: Use the Prefect logger to capture important information
  5. Resource Management: Clean up resources in task teardown
  6. Caching Strategy: Cache expensive computations but be mindful of data changes
  7. Testing: Test flows and tasks independently
  8. Version Control: Track flow code in version control
  9. Documentation: Document flow purpose, inputs, and outputs
  10. Monitoring: Set up notifications for critical flow failures

🔮 Advanced Prefect Features for ML

Dask Integration for Distributed Training 

from prefect.task_runners import DaskTaskRunner

@flow(task_runner=DaskTaskRunner())
def distributed_training_flow():
    # Tasks will be executed in a distributed Dask cluster
    results = []
    for i in range(10):
        results.append(train_model_fold(fold_id=i))
    return results
Enter fullscreen mode Exit fullscreen mode

Storage Options for Large Models 

from prefect.filesystems import S3

# Register S3 block
s3_block = S3(bucket_path="my-model-registry")
s3_block.save("model-storage")

# Use in deployment
deployment = Deployment.build_from_flow(
    flow=train_flow,
    name="distributed-training",
    storage=S3.load("model-storage"),
)
Enter fullscreen mode Exit fullscreen mode

Notifications for Critical Failures 

from prefect.notifications import SlackWebhook

slack_webhook = SlackWebhook(url="https://hooks.slack.com/services/XXX/YYY/ZZZ")
slack_webhook.save("ml-alerts")

@flow(on_failure=[SlackWebhook.load("ml-alerts")])
def critical_training_flow():
    # This flow will send a Slack message if it fails
    ...
Enter fullscreen mode Exit fullscreen mode

Top comments (0)