Mock a Flyte ML workflow in Python

Build a lightweight mock of a Flyte ML pipeline with dataclasses and a simple execution loop that passes outputs between tasks.

Medium Python 3.9+ Aug 9, 2026 ML engineering pipelines 16 views 0 copies

flyte ml-pipeline dataclass workflow mock

Python code

69 lines

Python 3.9+

from dataclasses import dataclass, field
from typing import List, Dict, Optional
import time


@dataclass
class FlyteTask:
    name: str
    inputs: Dict = field(default_factory=dict)
    outputs: Dict = field(default_factory=dict)

    def run(self) -> Dict:
        time.sleep(0.1)  # simulate work
        return self.outputs


class FlyteWorkflow:
    """Simple mock of a Flyte ML pipeline."""

    def __init__(self, name: str):
        self.name = name
        self.tasks: List[FlyteTask] = []

    def add_task(self, task: FlyteTask) -> None:
        self.tasks.append(task)

    def execute(self) -> Dict:
        results = {}
        for task in self.tasks:
            results[task.name] = task.run()
            # Mock passing outputs to next task that references them
            for other_task in self.tasks:
                for key, value in task.outputs.items():
                    if key in other_task.inputs:
                        other_task.inputs[key] = value
        return results


def create_ml_workflow() -> FlyteWorkflow:
    workflow = FlyteWorkflow("iris-classifier")

    preprocess = FlyteTask(
        name="preprocess",
        outputs={"X_train": [1.0, 2.0, 3.0], "y_train": [0, 1, 0]}
    )

    train = FlyteTask(
        name="train",
        inputs={"X_train": None, "y_train": None},
        outputs={"model": "random_forest.pkl"}
    )

    evaluate = FlyteTask(
        name="evaluate",
        inputs={"model": None},
        outputs={"accuracy": 0.94}
    )

    workflow.add_task(preprocess)
    workflow.add_task(train)
    workflow.add_task(evaluate)
    return workflow


if __name__ == "__main__":
    pipeline = create_ml_workflow()
    results = pipeline.execute()
    for task_name, outputs in results.items():
        print(f"{task_name}: {outputs}")

Output

stdout

preprocess: {'X_train': [1.0, 2.0, 3.0], 'y_train': [0, 1, 0]}
train: {'model': 'random_forest.pkl'}
evaluate: {'accuracy': 0.94}

How it works

This mock models a Flyte workflow as a list of task dataclasses. Each FlyteTask holds its inputs and outputs, and run() simulates work with a short sleep. The execute() method iterates through tasks, stores each result, and manually propagates outputs into downstream task inputs — mimicking how Flyte resolves dependencies between nodes. Keeping the mock isolated means you can test pipeline structure and data flow logic without standing up a real Flyte cluster.

Common mistakes

Forgetting to propagate outputs to downstream tasks before executing them
Using mutable default arguments like `inputs={}` instead of `field(default_factory=dict)`
Assuming execution order always matches task order without explicit dependencies
Not handling cases where a task input is missing a corresponding output

Variations

Use a topological sort to execute tasks in dependency order
Replace `time.sleep` with actual ML operations like `sklearn` fitting or inference

Real-world use cases

Prototyping a pipeline design before implementing the full Flyte workflow with real tasks.
Unit-testing data flow or business logic without spinning up Flyte's backend or containers.
Creating a fast local dev loop for ML engineers to verify orchestration logic before deployment.

Mock a Flyte ML workflow in Python

Python code

Output

How it works

Common mistakes

Variations

Real-world use cases

More from ML engineering pipelines

Tutorials

Quizzes

Python code

Output

How it works

Common mistakes

Variations

Real-world use cases

More from ML engineering pipelines

Keep learning

Tutorials

Quizzes