Skip to main content

Extending the Engine

The 3-Tier API Contract for building custom Models, Solvers, and Optimizers.

MLSys·im is designed to be fully extensible. Researchers and students can add custom analytical tools to resolve new constraints or search new design spaces.

All extensions must follow the 3-Tier API Contract, which requires each extension to declare what kind of math it does.

The 3-Tier API Contract

Before writing code, decide what kind of math the tool does.

1. ForwardModel: The Physics Engine

  • The Math: \(Y = f(X)\). Forward propagation.
  • When to use: To evaluate a physical or logical state. Given a fixed hardware config and a fixed workload, a Model predicts latency, cost, memory footprint, or energy.
  • Rule: A Model cannot make decisions or loop through options. It must be a deterministic calculation of a single state.

2. BaseSolver: The Math Engine

  • The Math: \(X = f^{-1}(Y)\) or \(\nabla f\). Algebraic inversion or calculus.
  • When to use: Given a specific target (like a latency SLA or a memory budget), solve algebraically for the exact hardware or model parameter required to hit it.
  • Rule: A Solver should yield a mathematically precise answer derived from inverting a Model’s equations.

3. BaseOptimizer: The Engineering Engine

  • The Math: \(\max_{x \in X} f(x) \text{ s.t. } g(x) \le c\). Constrained optimization.
  • When to use: To search a design space (discrete or continuous) for the “best” configuration, balancing competing trade-offs (e.g., maximizing throughput while minimizing carbon).
  • Rule: An Optimizer must internally call Models to evaluate candidates. It must return an OptimizerResult tracking the objective value and the size of the search space.

1. Building a Custom Model

Every resolver follows the same pattern: declare inputs (requires), declare outputs (produces), and implement solve().

The following example builds a custom PowerEfficiencyModel that calculates TFLOPs per Watt.

from mlsysim.solvers import ForwardModel
from mlsysim.engine.results import SolverResult
from mlsysim.core.units import Q_
from mlsysim.hardware.types import HardwareNode
from mlsysim.core.types import Quantity

# 1. Define the strictly typed Output
class PowerEfficiencyResult(SolverResult):
    flops_per_watt: Quantity
    is_efficient: bool

# 2. Implement the Model
class PowerEfficiencyModel(ForwardModel):
    """Evaluates the compute efficiency per watt of an accelerator."""
    requires = ("hardware",)
    produces = PowerEfficiencyResult

    def solve(self, hardware: HardwareNode) -> PowerEfficiencyResult:
        if hardware.tdp is None:
            raise ValueError(f"{hardware.name} has no TDP specified.")

        fpw = hardware.compute.peak_flops / hardware.tdp

        threshold = Q_("1 TFLOPs/s / W")
        is_eff = fpw > threshold

        return PowerEfficiencyResult(
            flops_per_watt=fpw.to("TFLOPs/s/W"),
            is_efficient=is_eff
        )

2. Building a Custom Solver

A solver algebraically inverts an equation. For example, if \(T = \frac{W}{BW}\), and we have a target \(T\), we solve for \(BW = \frac{W}{T}\).

from mlsysim.solvers import BaseSolver
from mlsysim.engine.results import SolverResult
from mlsysim.models.types import Workload
from mlsysim.core.types import Quantity

class RequiredBandwidthResult(SolverResult):
    required_bw: Quantity

class RequiredBandwidthSolver(BaseSolver):
    """Solves for the exact memory bandwidth needed to hit an SLA."""
    requires = ("workload", "target_latency")
    produces = RequiredBandwidthResult

    def solve(self, model: Workload, target_latency: Quantity) -> RequiredBandwidthResult:
        weight_bytes = model.size_in_bytes()
        t_target = target_latency.to("s")

        # Algebraic inversion
        required_bw = (weight_bytes / t_target).to("GB/s")

        return RequiredBandwidthResult(required_bw=required_bw)

3. Building a Custom Optimizer

An Optimizer explores a design space. It MUST inherit from BaseOptimizer and its result MUST inherit from OptimizerResult.

The following example builds a CheapestHardwareOptimizer that searches the HardwareZoo for the cheapest chip that satisfies a minimum TFLOP requirement.

from mlsysim.solvers import BaseOptimizer
from mlsysim.engine.results import OptimizerResult
from mlsysim.hardware.registry import Hardware
from typing import Dict, Any

# Inherit from OptimizerResult, which requires specific fields
class CheapestHardwareResult(OptimizerResult):
    cheapest_cost: float
    hardware_name: str

class CheapestHardwareOptimizer(BaseOptimizer):
    requires = ("min_tflops",)
    produces = CheapestHardwareResult

    def solve(self, min_tflops: float) -> CheapestHardwareResult:
        candidates = []
        target = Q_(min_tflops, "TFLOPs/s")

        # 1. Define Search Space
        for hw in Hardware.list():
            if hw.unit_cost is None:
                continue

            # 2. Evaluate Constraint
            if hw.compute.peak_flops >= target:
                candidates.append({
                    "name": hw.name,
                    "cost": hw.unit_cost.magnitude
                })

        if not candidates:
            raise ValueError("No hardware meets the requirement.")

        # 3. Optimize Objective (Minimize cost)
        best = min(candidates, key=lambda x: x["cost"])

        # 4. Return standard OptimizerResult structure
        return CheapestHardwareResult(
            objective_value=best["cost"], # Standard field
            best_config={"hardware": best["name"]}, # Standard field
            total_searched=len(Hardware.list()), # Standard field
            cheapest_cost=best["cost"],
            hardware_name=best["name"]
        )

4. Composable Pipelines & Callbacks

MLSys·im does not evaluate models only in isolation. It uses a Composable Pipeline architecture, in which custom solvers snap into a pipeline and execute sequentially, with outputs from earlier stages passing automatically to later stages.

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

# Build a pipeline with your custom solver in the middle
my_pipeline = Pipeline([
    DistributedModel(),
    PowerEfficiencyModel(), # Your custom model
    EconomicsModel()
])

# Run the pipeline
results = my_pipeline.run(
    model=mlsysim.Models.Language.Llama3_8B,
    fleet=mlsysim.Systems.Clusters.Frontier_8K,
    duration_days=30
)

# Access individual stage results
print(results["PowerEfficiencyModel"].flops_per_watt)

Callbacks (Middleware)

A callback intercepts pipeline execution to log metrics to MLOps platforms like Weights & Biases or Datadog.

class WandbLogger:
    def on_stage_end(self, stage_name: str, result: SolverResult):
        # Example pseudo-code
        wandb.log({f"{stage_name}/{k}": v for k, v in result.dict().items()})

my_pipeline.register_callback(WandbLogger())

5. Registering a Custom Resolver (Plugins)

Custom models do not require a fork of mlsysim. Package the code as a standard Python module and register it via pyproject.toml entry points. When a user installs the package, mlsysim automatically discovers and loads its custom solvers into the engine.

In the plugin’s pyproject.toml, add:

[project.entry-points."mlsysim.solvers"]
power_efficiency = "my_custom_package.solvers:PowerEfficiencyModel"
required_bandwidth = "my_custom_package.solvers:RequiredBandwidthSolver"
cheapest_hardware = "my_custom_package.solvers:CheapestHardwareOptimizer"

After installation via pip install ., check that mlsysim has discovered the models by running:

from mlsysim.engine.resolver_factory import ResolverFactory

# This will now include 'PowerEfficiencyModel', etc.
print(ResolverFactory.list_available().keys())

6. Testing an Extension

Every custom resolver should have at least one unit test that verifies dimensional correctness and expected output:

import pytest
from mlsysim import Hardware

def test_power_efficiency_h100():
    model = PowerEfficiencyModel()
    result = model.solve(hardware=Hardware.Cloud.H100)

    assert result.flops_per_watt.units == "TFLOPs/s/W"
    assert result.flops_per_watt.magnitude > 0
    assert isinstance(result.is_efficient, bool)

def test_power_efficiency_no_tdp():
    """Hardware without TDP should raise ValueError."""
    from mlsysim.hardware.types import HardwareNode, ComputeCore, MemoryHierarchy
    from mlsysim.core.units import Q_

    hw = HardwareNode(
        name="NoTDP",
        release_year=2024,
        compute=ComputeCore(peak_flops=Q_("100 TFLOPs/s")),
        memory=MemoryHierarchy(capacity=Q_("80 GB"), bandwidth=Q_("2 TB/s")),
        tdp=None,
    )
    with pytest.raises(ValueError, match="has no TDP"):
        PowerEfficiencyModel().solve(hardware=hw)

Run tests with:

pytest tests/test_my_extension.py -v

Why Strict Typing

By forcing inputs and outputs to use pint.Quantity, MLSys·im guarantees dimensional consistency. The Pipeline module uses these class signatures (requires and produces) to automatically stitch different Models, Solvers, and Optimizers together into a single execution DAG.

As a result, a custom extension automatically works with:

  • Engine.sweep(): sweep the model across hardware
  • SystemEvaluator.evaluate(): include the model in the full scorecard
  • The CLI: expose the extension via mlsysim eval with YAML input
  • Schema tooling: validate configuration files before they reach the solver
Back to top