Port a model in 15 minutes¶

You have a model that works in a notebook. This guide changes it into a plugin. The platform can fit the plugin, score it on new data, forecast with it, keep its versions and plot it. The contract kit proves that the plugin works. You write the statistics. The platform writes all the other code.

flowchart LR
    A["0–3 min<br/><b>Pick a base class</b>"] --> B["3–6 min<br/><b>Declare the data contract</b>"]
    B --> C["6–11 min<br/><b>Write build / estimator / fit+score</b>"]
    C --> D["11–13 min<br/><b>Write a tiny spec</b>"]
    D --> E["13–15 min<br/><b>Run the contract kit</b>"]

What you write, what you get¶

	You write	The platform does
`PyMCModel`	`build(ds, settings, ctx) -> pm.Model`	Engines (nutpie, PyMC with segments that can continue, MAP, ADVI), convergence diagnostics, scores that build the model again and predict, new levels from the hyperprior, counterfactuals with common random numbers, fast scores
`SklearnModel`	`estimator(settings)`	It stacks labelled dims into rows and features. Also coordinate features, safe persistence (skops, xgboost JSON, LightGBM text, CatBoost JSON), quantile / bootstrap / conformal uncertainty, pooled new levels, contributions
`CustomModel`	`fit(...)` and `score(...)`	Binding, preprocessing, splits, context tails, the canonical form of coordinates, inverse transforms, summaries, bundles

The three base classes get these items at no cost:

Validated settings with JSON Schemas
Static validation of the spec
An evaluation, then a new fit
CoordPlan checks
Summaries that obey the representation
Plots and bundles

1. Pick a base class (0–3 min)¶

If your model is a PyMC model, use mapi_pymc.PyMCModel.
If your model is an sklearn, xgboost, LightGBM or CatBoost estimator on a design matrix, use mapi_sklearn.SklearnModel (or the quantile, bootstrap, conformal or classifier variants). LightGBM and CatBoost need the extra mapi-sklearn[gbm].
For all other models (statsmodels, JAX, a Kalman filter, closed-form mathematics), use mapi_sdk.CustomModel. A neural network that trains on many images also uses CustomModel, with the hooks fit_batches and score_batches (see Train on batches and Deep learning).

2. Declare the data contract (3–6 min)¶

Write the dims and the variables that your model needs, in canonical names. For each dim, write what can occur to it at score time. This is the schema_template of the model. In the spec, users bind the names of their own dataset to it.

Three questions for each dim

Can a score request bring a new level (a new county, store or market)? If yes, the dim is new_levels, with a strategy. Does the dim grow at forecast time (time, observations)? If yes, the dim is extend. If not, the dim is fixed.

3. Write the model (6–11 min)¶

PyMCscikit-learn / xgboostCustom

from typing import Any, ClassVar

import pymc as pm
import xarray as xr
from mapi_pymc import BuildContext, PyMCModel
from mapi_pymc.examples import RADON_SCHEMA  # dims obs (extend) + county (new_levels)
from mapi_sdk import ModelSettings
from mapi_spec import Capabilities, Representation
from pydantic import Field


class RadonSettings(ModelSettings):
    beta_sigma: float = Field(1.0, gt=0)


class Radon(PyMCModel[RadonSettings]):
    """log radon ~ Normal(a[county] + b * floor, sigma), with partially pooled a."""

    id = "acme.radon"
    version = "1.0.0"
    Settings = RadonSettings
    capabilities = Capabilities(
        inference="mcmc",
        representation=Representation.DRAWS,
        score_modes=frozenset({"predict", "counterfactual", "evaluate"}),
    )
    schema_template = RADON_SCHEMA
    outputs: ClassVar[dict[str, Any]] = {"y": {"dims": ["obs"], "units_of": "y"}}
    expected: ClassVar[dict[str, str]] = {"y": "mu"}  # output -> Deterministic for E[y]
    pad_extend = False  # long layout: no positional time axis

    def build(self, ds: xr.Dataset, s: RadonSettings, ctx: BuildContext) -> pm.Model:
        with pm.Model(coords=ctx.coords()) as model:
            floor = ctx.data("floor")
            county = ctx.data("county_idx")
            a = ctx.hier_effect("a", "county")  # new counties come from the hyperprior
            b = pm.Normal("b", 0.0, s.beta_sigma)
            mu = pm.Deterministic("mu", a[county] + b * floor, dims="obs")
            sigma = pm.HalfNormal("sigma", 1.0)
            pm.Normal("y", mu, sigma, observed=ctx.observed("y"), dims="obs")
        return model

The ctx helpers make the scores work. You do not have to think about them:

Helper	Why you use it
`ctx.coords()`	The coords of the model. At score time, it includes new levels.
`ctx.data(name)`	A `pm.Data` over the training values at fit time and over the score values at score time (the platform builds the model again)
`ctx.observed(name)`	The target at fit time. A placeholder at score time (never imputed).
`ctx.hier_effect(name, dim)`	A partially pooled effect. Its new levels come from the hyperprior, for each posterior draw.
`ctx.position(dim)`	0, 1, 2, … along an `extend` dim. The count continues into the forecast horizon.

expected names the Deterministic that is the expected value of each output. Thus output.kind: expected returns it without observation noise.

from typing import ClassVar

from mapi_sklearn import PointSettings, SklearnModel, capabilities_for
from mapi_spec import Capabilities
from sklearn.linear_model import Ridge


class LinearDemand(SklearnModel[PointSettings]):
    id = "t.acme.linear_demand"
    version = "1.0.0"
    capabilities: ClassVar[Capabilities] = capabilities_for("none").model_copy(
        update={"contribution_scale": "additive"}  # a linear model decomposes exactly
    )

    def estimator(self, s: PointSettings) -> Ridge:
        return Ridge(alpha=0.5)  # must be exactly an allow-listed class

estimator() returns an estimator instance from the allow-list (or a {kind, params} mapping). The platform rejects subclasses and parameters that are not in the params model of the estimator. Thus the platform can always load the bundle safely. Rows, columns and coordinate features are settings, not code. See Tabular models.

from mapi_sdk import PREDICTIONS, CustomModel, FitOutput, ModelSettings, new_level_mask
from mapi_spec import Capabilities, DataSchema, Representation


class BaselineSettings(ModelSettings):
    ridge: float = 1e-6


class GeoBaseline(CustomModel[BaselineSettings]):
    """y = intercept[geo] + beta * spend, by least squares (point predictions)."""

    id = "acme.geo_baseline"
    version = "1.0.0"
    Settings = BaselineSettings
    capabilities = Capabilities(
        inference="point",
        representation=Representation.POINT,
        score_modes=frozenset({"predict", "forecast", "evaluate"}),
    )
    # the data contract the model needs, in canonical names (a DataSchema, or a
    # DataSchemaTemplate whose names come from settings)
    schema_template = DataSchema.model_validate(
        {
            "dims": {
                "geo": {"kind": "categorical", "policy": "new_levels", "new_level_strategy": "pooled"},
                "time": {"kind": "time", "policy": "extend"},
            },
            "vars": {
                "y": {"dims": ["geo", "time"], "dtype": "float64", "role": "target"},
                "spend": {"dims": ["geo", "time"], "dtype": "float64", "role": "covariate"},
            },
        }
    )
    outputs: ClassVar[dict] = {"y_hat": {"dims": ["geo", "time"], "units_of": "y"}}
    pad_extend = False  # time is not indexed by position

    def fit(self, ds, s, cfg, ctx):
        g, t = ds.sizes["geo"], ds.sizes["time"]
        x = np.concatenate(
            [np.repeat(np.eye(g), t, axis=0), ds["spend"].values.reshape(-1, 1)], axis=1
        )
        coef = np.linalg.solve(x.T @ x + s.ridge * np.eye(g + 1), x.T @ ds["y"].values.ravel())
        params = xr.Dataset(
            {"intercept": ("geo", coef[:g]), "beta": coef[g]}, coords={"geo": ds["geo"].values}
        )
        return FitOutput(groups={"params": params})

    def score(self, fitted, ds, req, ctx):
        # `ds` arrives canonicalised: fit order, unseen geos appended and flagged
        p = fitted.group("params")
        new = new_level_mask(ds, "geo")
        known = p["intercept"].reindex(geo=ds["geo"].values)
        pooled = 0.0 if ctx.strategy("geo") == "zero" else float(p["intercept"].mean())
        intercept = xr.where(new, pooled, known.fillna(pooled))
        y_hat = (intercept + p["beta"] * ds["spend"]).transpose("geo", "time")
        return xr.DataTree.from_dict({PREDICTIONS: xr.Dataset({"y_hat": y_hat})})

The rules for score:

The dataset comes in the canonical form. It has the fit order. New levels come at the end, with the flag <dim>__is_new (new_level_mask(ds, dim)). Thus parameters that use the position as an index stay aligned.
Return an xr.DataTree with a predictions group in the representation that you declared.
Never look at the targets in the payload. The evaluate mode sends them for the metrics, not for you.

FitOutput.groups uses the ArviZ group names (posterior, sample_stats, …). learned contains small JSON. native contains bytes in a safe format only (skops, xgboost JSON, safetensors). The platform rejects pickle.

What the notebook models of science needed¶

The 15 sci.* models came from notebooks in the same way. Some of them needed more than the three steps above. If your model is similar, see these sections:

Your model	What to use	See
Its rules join two or more settings	A pydantic `model_validator` on the settings. The SDK refuses bad settings before it uses the data.	Settings rules that need no data
An output is an image, a map or a vector field	`grid` in the output declaration	Outputs that are images or maps
An output is a direction or a time of day	`circular: true`, and a `period` attribute on the scored array	Outputs that are angles
It makes new data from the posterior	The `simulate` mode, and real predictive draws for `kind: predictive`	Score modes
You compare variants by `elpd_loo`	A `log_likelihood` group, or `loo_valid=False` with a reason	Models that give `elpd_loo`
NUTS gets stuck at the default start	`EnginePlan.initvals` and `jitter=False`	Start every chain at one point
It has its own sampler (for example Gibbs)	A `CustomModel` with `inference="mcmc"` and a `diagnostics` hook	A model with its own sampler
Some NaN of an input are part of the design	`structural_missing`	A model
Users must upload the data of a case study	An export case in `mapi_plugin_science.export.EXPORTS`	Data for a case study

What the platform can show from your model v0.10–v0.15¶

These items are optional. Each one makes the run page or the model version show more:

You want	What to use	See
Standard plots for the model	The class attribute `plots`, a list of recipe ids. The Insert panel of the spec editor shows them first. A fit draws only the plots in `outputs.plots`.	Default plots and metrics
Curves and values on the run page	`ctx.log_metric`, `ctx.log_params`, `ctx.set_tag` and `ctx.log_line` in `fit` and `score`	Log metrics, params and tags
The progress of each chain of your own sampler	`ctx.chains([...])`, and step metrics with the label `chain`	A model with its own sampler
Images, figures and small tables on the run page	`ctx.log_image`, `ctx.log_figure` and `ctx.log_table`	Log images, figures and tables
Training on many images without all of them in memory	`Capabilities(input="collection")`, `fit_batches` and `score_batches`	Train on batches
Checkpoints that a run can continue from	`ctx.fit_key` and `ctx.checkpoint_sink`	Train on batches
A side table that only the model reads, for example pretrained weights	`side_table_params` and `ctx.side_tables`	Read a side table in a model

4. Write a tiny spec (11–13 min)¶

The contract kit fits this spec. Keep it small and fast: use few draws and few trees.

PyMCCustom (with a step)

from mapi_pymc.examples import radon_dataset
from mapi_sdk import Registry
from mapi_sdk.testing import ModelFixture, check_model

SPEC = """
apiVersion: mapi/v1
kind: ModelSpec
metadata: {name: radon, project: demo}
data: {dataset: dataset://demo/radon@latest}
model:
  use: acme.radon@^1
  fit: {mode: train_only, engine: pymc, tune: 100, draws: 100, chains: 2, seed: 3}
"""
report = check_model(
    Radon, ModelFixture(spec=SPEC, data=radon_dataset(), raw_schema=RADON_SCHEMA, registry=Registry())
)
print(report.summary())

from mapi_sdk import default_registry
from mapi_sdk.testing import ModelFixture, check_model

registry = default_registry().copy()
registry.register(ScaleStep)  # the model is registered by check_model itself

SPEC = """
apiVersion: mapi/v1
kind: ModelSpec
metadata: {name: acme-baseline, project: demo}
data:
  dataset: dataset://demo/sales@latest
  split: {kind: time_holdout, dim: time, last: 6}
pipeline:
  - {id: scale, use: acme.scale@^1, with: {vars: [spend], over: [time], unseen_levels: pooled}}
model:
  use: acme.geo_baseline@^1
  fit: {seed: 1}
evaluate: {metrics: [rmse, mape]}
"""
raw_schema = infer_schema(data, targets=["y"], policies={"geo": "new_levels"})
report = check_model(
    GeoBaseline,
    ModelFixture(spec=SPEC, data=data, raw_schema=raw_schema, registry=registry),
)
print(report.summary())
# contract kit: model acme.geo_baseline@1.0.0
#   [PASS] fit: ...
#   [PASS] bundle_roundtrip_no_pickle: ...
#   [PASS] shuffled_scoring: ...
#   ...

5. Run the contract kit (13–15 min)¶

check_model fits your fixture. Then it tries to break your plugin in the same ways that real users will:

contract kit: model acme.geo_baseline@1.0.0
  [PASS] manifest: acme.geo_baseline@1.0.0 (point)
  [PASS] fit: bundle ef4100375dbb, metrics ['mape', 'rmse']
  [PASS] bundle_roundtrip_no_pickle: scores match after write/read
  [PASS] shuffled_scoring: shuffled ['geo']
  [PASS] subset_scoring: subset of 'geo'
  [PASS] new_level_policies: geo: new level scored (pooled)
  [PASS] extend_typed_errors: raised NON_CONTIGUOUS_EXTENSION; raised COORD_NOT_ALLOWED
  [PASS] unsupported_mode: raised MODE_NOT_SUPPORTED
  [PASS] context_window_equivalence: holdout via stored context == in-sample full series
  [PASS] evaluate_path_equals_score_path: metrics ['rmse', 'mape'] reproduced through the score path
  [PASS] output_space_and_units: y_hat in raw (None)
  [PASS] representation_aggregation: point aggregation over 'geo' allowed
  [SKIP] additive_contributions: contribution_scale is not additive
  [PASS] seed_determinism: same seed, same fit and scores

Check	What it proves
`bundle_roundtrip_no_pickle`	The scores are the same after the kit writes and reads the bundle. The bundle contains no pickle.
`shuffled_scoring`, `subset_scoring`	Labels in a different sequence and subsets give the same answers for the same cells.
`new_level_policies`	Each strategy that your dims declare works (or raises `UNSEEN_LEVELS` when it must).
`extend_typed_errors`	Gaps and time labels out of range raise typed errors.
`unsupported_mode`	The model refuses modes that it does not declare.
`context_window_equivalence`	A forecast score of the last training periods is equal to the in-sample result.
`evaluate_path_equals_score_path`	The normal score path makes the same holdout metrics again.
`output_space_and_units`	The outputs are inverted. They have labels for the units and the output space.
`representation_aggregation`	The representation controls which aggregations the model allows or refuses.
`additive_contributions`	The sum of the contributions that are declared additive is equal to the prediction.
`seed_determinism`	If the plugin claims it, the same seed gives the same fit and scores.

If a check fails, the kit raises ContractViolation with the report. To get the report as a return value, pass raise_on_failure=False.

In your test suite¶

mapi-sdk installs a pytest plugin. Thus the kit is available as fixtures in all projects:

import pytest
from mapi_sdk.testing import ModelFixture, toy  # toy: small reference steps and models


@pytest.mark.mapi_contract
def test_step_follows_the_contract(mapi_check_step, mapi_panel):
    data, _fx = mapi_panel  # a daily market x channel panel and its FX side table
    report = mapi_check_step(
        toy.ScaleStep, data, {"vars": ["spend_local"], "over": ["date"], "unseen_levels": "pooled"}
    )
    assert report.ok, report.summary()


SPEC = """
apiVersion: mapi/v1
kind: ModelSpec
metadata: {name: toy, project: demo}
data:
  dataset: dataset://demo/panel@latest
  bind: {dims: {geo: market, time: date}, vars: {y: revenue_local, spend: spend_local}}
  split: {kind: time_holdout, dim: time, last: 14}
model: {use: toy.linear@^1, fit: {seed: 1}}
evaluate: {metrics: [rmse]}
"""


@pytest.mark.mapi_contract
def test_model_follows_the_contract(mapi_check_model, mapi_panel, mapi_toy_registry):
    data, _fx = mapi_panel
    fixture = ModelFixture(spec=SPEC, data=data, registry=mapi_toy_registry)
    report = mapi_check_model(toy.LinearModel, fixture)  # raises ContractViolation on failure
    assert report.ok

Bonus: a preprocessing step¶

Steps use the same pattern. A step has settings, an optional state, fit / transform / inverse, and a static effect_template. The effect_template tells the control plane how the schema changes.

from typing import ClassVar

from mapi_sdk import Step, StepContext, StepSettings, StepState
from mapi_spec import StepKind
from pydantic import Field, JsonValue


class ScaleSettings(StepSettings):
    vars: list[str]
    over: list[str] = Field(default=[])  # plain defaults: the exported JSON Schema keeps them
    unseen_levels: str = "error"  # error | pooled | fill: what a new level's scale becomes
    fill_value: float = 1.0  # the scale used by 'fill'


class ScaleStep(Step[ScaleSettings, StepState]):
    """Divide each variable by its mean absolute value over `over` (learned on train)."""

    id = "acme.scale"
    version = "1.0.0"
    kind = StepKind.VALUE
    Settings = ScaleSettings
    State = StepState  # stateful: fitted on the training data only
    invertible = True
    var_params = ("vars",)  # settings that name variables ...
    dim_params = ("over",)  # ... and dims, so the control plane can check them
    unseen_levels_param = "unseen_levels"
    effect_template: ClassVar[dict[str, JsonValue]] = {  # static: how the schema changes
        "template_params": ["vars"],
        "updates": {"{*vars as v}": {"units": "1"}},
    }
    summary = "Divide by the mean absolute value over the named dims"  # catalog and UI text
    summary_plain = "Puts each series on a comparable scale"

    def fit(self, ds: xr.Dataset, s: ScaleSettings, ctx: StepContext) -> StepState:
        return StepState(arrays=xr.Dataset({f"scale__{v}": abs(ds[v]).mean(s.over) for v in s.vars}))

    def _scale(self, v: str, ds: xr.Dataset, s: ScaleSettings, st: StepState, ctx: StepContext):
        state = st.array(f"scale__{v}")
        return ctx.align_state(state, ds, s.unseen_levels, fill_value=s.fill_value)

    def transform(self, ds, s, st, ctx):  # never learns; skips absent vars (the target at score)
        return ds.assign({v: ds[v] / self._scale(v, ds, s, st, ctx) for v in s.vars if v in ds})

    def inverse(self, ds, s, st, ctx):
        return ds.assign({v: ds[v] * self._scale(v, ds, s, st, ctx) for v in s.vars if v in ds})

from mapi_sdk import infer_schema
from mapi_sdk.testing import check_step

report = check_step(
    ScaleStep,
    data,
    {"vars": ["spend"], "over": ["time"], "unseen_levels": "pooled"},
    schema=infer_schema(data, targets=["y"]),
)
print(report.summary())
# contract kit: step acme.scale@1.0.0
#   [PASS] catalog_entry: exported acme.scale@1.0.0
#   [PASS] transform_never_learns: ...
#   ...

The step kit does a check of these items:

The exported catalog entry
A deterministic and pure fit
That transform never learns
The bundle round trip of the state
That the declared schema effect agrees with the real output
That the step accepts a missing target
The inverse round trip
The equivalence of the context window for steps that look back
Each policy for unseen levels

For more information, see Write a step or model plugin.

Ship it¶

Register your plugins through entry points in the pyproject.toml of your package. The value can be a class, a tuple of classes or a function that returns them:

[project.entry-points."mapi.models"]
acme = "acme_models.plugins:PLUGINS"      # PLUGINS = (Radon, GeoBaseline)

[project.entry-points."mapi.steps"]
acme = "acme_models.steps:STEPS"          # STEPS = (ScaleStep,)

mapi_sdk.default_registry() then finds them. The runner images of the platform also find them.

Checklist

For the full guides with examples, see PyMC, scikit-learn and xgboost and steps and custom models.