How to Stage ML Model Workflows with Python Classes

Defines a Stage class to model ML pipeline stages with variants and mocks, printing grammar for Model, Staging, and Production stages.

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

ml-pipelines stages model-deployment mock workflow class-design

Python code

38 lines

Python 3.9+

class Stage:
    def __init__(self, name):
        self.name = name
        self.mocks = []
        self.variants = []

    def add_mock(self, mock_name):
        self.mocks.append(mock_name)

    def add_variant(self, variant_name, productions=()):
        self.variants.append((variant_name, list(productions)))

    def get_grammar(self):
        result = "grammar:\n"
        for variant_name, productions in self.variants:
            result += f"  {self.name} ::= {variant_name} -> {', '.join(productions)}\n"
        return result

def stages_workflow():
    model_stage = Stage("Model")
    staging = Stage("Staging")
    production = Stage("Production")

    model_stage.add_variant("DSL", ["parse", "evaluate", "infer"])
    staging.add_variant("Preflight", ["validate", "sanitize", "compile"])
    production.add_variant("Release", ["package", "deploy", "monitor"])

    model_stage.add_mock("mock_schema_transformer")
    staging.add_mock("mock_schema_transformer")
    production.add_mock("mock_schema_transformer")

    print(model_stage.get_grammar())
    print(staging.get_grammar())
    print(production.get_grammar())
    print("Mocks enabled:", all(len(s.mocks) > 0 for s in [model_stage, staging, production]))

if __name__ == "__main__":
    stages_workflow()

Output

stdout

grammar:
  Model ::= DSL -> parse, evaluate, infer
grammar:
  Staging ::= Preflight -> validate, sanitize, compile
grammar:
  Production ::= Release -> package, deploy, monitor
Mocks enabled: True

How it works

The Stage class encapsulates the name, mock list, and variant productions, keeping stage-specific state in one place. add_variant stores production names as a list per variant, while add_mock collects mock identifiers. get_grammar formats the variants into a simple grammar-like string output for each stage. The stages_workflow function instantiates three stages, populates them with realistic ML pipeline actions, and verifies that all stages have mocks enabled. This pattern makes it easy to extend stages or add new pipeline components without coupling them together.

Common mistakes

Forgetting to convert tuple productions to a list, causing mutation issues later
Not initializing `mocks` and `variants` in `__init__`, leading to AttributeError
Overwriting the `name` property by naming a variable the same as the attribute

Variations

Use dataclasses to reduce boilerplate for the Stage class
Add a `run_pipeline` method to execute all stage variants in sequence

Real-world use cases

Model registry pipelines that promote candidate models through dev, staging, and production gates.
Feature engineering frameworks where each stage defines transformation variants and test mocks.
CI/CD for ML services that stage validation, sanitization, and deployment as separate executable steps.

How to Stage ML Model Workflows with Python Classes

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