Skip to main content

Composable Pipelines & Callbacks

How to build custom execution graphs and log metrics to external platforms.

node
advanced
Learn how to snap custom analytical solvers into the MLSys·im Pipeline and use middleware to log data.

The Question

The Engine.evaluate() function serves standard roofline analysis well. Some workflows, however, need to insert a proprietary power model into the middle of the evaluation, or to log every analytical result to a database or an MLOps platform like Weights & Biases.

How do we extend the core framework without forking it?

NoteLearning Objectives
  • Compose custom execution graphs using the Pipeline class.
  • Inject custom analytical solvers into the evaluation flow.
  • Register callback middleware to log or intercept data.

1. Setup

First, import the necessary modules. This tutorial uses the Pipeline orchestrator directly instead of relying on the high-level Engine.

import mlsysim
from mlsysim.engine.pipeline import Pipeline
from mlsysim.solvers import DistributedModel, EconomicsModel

2. The Composable Pipeline

The Pipeline object accepts a list of resolvers. Calling .run() executes them in order, and the output of an early solver is automatically made available as input to later solvers.

A standard evaluation pipeline looks like this:

# 1. Define the execution graph
pipe = Pipeline([
    DistributedModel(),
    EconomicsModel()
])

# 2. Provide the initial inputs (Demand & Supply)
results = pipe.run(
    model=mlsysim.Models.Language.Llama3_8B,
    fleet=mlsysim.Systems.Clusters.Frontier_8K,
    batch_size=32,
    precision="fp16",
    efficiency=0.45,
    duration_days=30
)

# 3. Access results by stage name
tco = results["EconomicsModel"].tco_usd
print(f"Total Cost of Ownership: ${tco:,.2f}")
Total Cost of Ownership: $7,007,509.16

For Llama-3.1-8B on the 8,192-GPU Frontier Cluster (8192 GPUs) cluster over 30 days, DistributedModel computes a scaling efficiency of 64.4% and feeds into EconomicsModel, which reports a 30-day TCO of $7,007,509.16.

3. Injecting Custom Solvers

Because the pipeline is dynamic, it accepts custom analytical solvers. Suppose a custom plugin called PowerEfficiencyModel has been written (as described in the Extending the Engine guide).

It slots directly into the DAG:

# Assuming a custom ForwardModel subclass is defined:
# class PowerEfficiencyModel(ForwardModel): ...

custom_pipe = Pipeline([
    DistributedModel(),
    PowerEfficiencyModel(), # Injected right into the middle
    EconomicsModel()
])

# Use it just like before
# custom_pipe.run(...)

4. Callbacks (Middleware)

MLSys·im supports middleware hooks. To log metrics to an external platform (like a database, MLflow, or Weights & Biases), write a Callback class with an on_stage_end method.

The pipeline will trigger this method every time a solver finishes.

class WandbLogger:
    def on_stage_end(self, stage_name: str, result):
        print(f"[W&B Callback] Logging data for {stage_name}...")

        # If the result is a Pydantic model, dump it to a dict
        if hasattr(result, "model_dump"):
            metrics = result.model_dump(exclude_none=True)

            # Example: wandb.log({f"{stage_name}/{k}": v for k, v in metrics.items()})
            print(f"  -> Logged {len(metrics)} keys")

# Attach the logger
pipe.register_callback(WandbLogger())

# Run the pipeline again to see the callback fire
print("Running Pipeline with Callback:")
_ = pipe.run(
    model=mlsysim.Models.Language.Llama3_8B,
    fleet=mlsysim.Systems.Clusters.Frontier_8K,
    batch_size=32,
    precision="fp16",
    efficiency=0.45,
    duration_days=30
)
Running Pipeline with Callback:
[W&B Callback] Logging data for DistributedModel...
  -> Logged 12 keys
[W&B Callback] Logging data for EconomicsModel...
  -> Logged 11 keys

Next Steps

Back to top