Model Contract: `UbunyeModel`¶

For anyone writing a model class that Ubunye can train, save, register and load.

UbunyeModel is the base class every model you write must fill in. It is the only thing the engine ever calls on your model. The engine itself imports no ML library.

The contract¶

from abc import ABC, abstractmethod
from typing import Any, Dict

class UbunyeModel(ABC):

    @abstractmethod
    def train(self, df: Any) -> Dict[str, Any]:
        """Train on df. Returns a metrics dict e.g. {"auc": 0.94, "f1": 0.87}."""

    @abstractmethod
    def predict(self, df: Any) -> Any:
        """Generate predictions. Returns df-like with prediction columns added."""

    @abstractmethod
    def save(self, path: str) -> None:
        """Serialize model to path. Pick any format: pickle, joblib, ONNX, …"""

    @classmethod
    @abstractmethod
    def load(cls, path: str) -> "UbunyeModel":
        """Deserialize from path. Returns a ready-to-predict instance."""

    @abstractmethod
    def metadata(self) -> Dict[str, Any]:
        """Return at minimum: library, library_version, features, params."""

    def validate(self, df: Any) -> Dict[str, Any]:
        """Optional holdout validation. Default raises NotImplementedError."""
        raise NotImplementedError("validate() not implemented")

Why this contract?¶

The engine never imports sklearn, torch, xgboost, or any ML library. It only calls train(), predict(), save(), and load(). The implementation is entirely yours.

This gives you:

Library freedom: wrap any ML framework without engine changes.
Testability: test your model class with plain unit tests and a small fake data frame (a MockDF, see below); no Spark needed.
Registry integration: ModelRegistry.register() calls save() and stores metadata() and the metrics train() returned.

Minimal implementation (sklearn)¶

# model.py  (in your task folder)
import joblib
from pathlib import Path
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import roc_auc_score, f1_score
from ubunye.models.base import UbunyeModel

class FraudRiskModel(UbunyeModel):

    def __init__(self):
        self._clf = GradientBoostingClassifier(n_estimators=100, max_depth=4)
        self._features = ["amount_sum_30d", "txn_count_30d", "avg_amount_30d", "risk_encoded"]

    def train(self, df) -> dict:
        pdf = df.toPandas()
        X, y = pdf[self._features], pdf["is_fraud"]
        self._clf.fit(X, y)
        proba = self._clf.predict_proba(X)[:, 1]
        preds = self._clf.predict(X)
        return {
            "auc": float(roc_auc_score(y, proba)),
            "f1":  float(f1_score(y, preds)),
        }

    def predict(self, df):
        pdf = df.toPandas()
        pdf["fraud_score"]     = self._clf.predict_proba(pdf[self._features])[:, 1]
        pdf["fraud_predicted"] = self._clf.predict(pdf[self._features])
        return df.sparkSession.createDataFrame(pdf)

    def save(self, path: str) -> None:
        Path(path).mkdir(parents=True, exist_ok=True)
        joblib.dump(self._clf, f"{path}/model.pkl")

    @classmethod
    def load(cls, path: str) -> "FraudRiskModel":
        m = cls()
        m._clf = joblib.load(f"{path}/model.pkl")
        return m

    def metadata(self) -> dict:
        return {
            "library":         "scikit-learn",
            "library_version": "1.4.0",
            "features":        self._features,
            "params":          self._clf.get_params(),
        }

`train()`: return value¶

train() must return a dict of metric name → numeric value. These metrics are stored in the registry and used by promotion gates.

return {
    "auc":      0.92,
    "f1":       0.87,
    "accuracy": 0.95,
    "loss":     0.08,
}

A promotion gate can only check a metric that train() returns. The sklearn model above returns auc and f1, so min_auc works and max_loss would block every promotion. See Promotion gates.

`predict()`: return value¶

predict() receives the same type of object train() received (a Spark DataFrame, a pandas DataFrame, or anything with the methods your model calls) and should return a DataFrame-like with prediction columns added. Return it from your task's transform() under an output name and the engine writes it like any other output.

`metadata()`: required keys¶

The metadata() dict must contain at minimum:

Key	Type	Description
`library`	string	ML library used (e.g. `"scikit-learn"`, `"xgboost"`, `"pytorch"`)
`library_version`	string	Version string
`features`	list of strings	Input feature column names
`params`	dict	Hyperparameters

Additional keys are stored but not validated.

features is also the model's input contract at predict. A ModelTransform predict compares the frame's columns with it before calling the model. Extra columns are fine. By default a missing column is a warning on the log and the run record; with registry.check_features: true the frame is refused with FeatureContractError, which names the missing columns, the model and the version. If your predict() builds some declared features itself, leave the check at its default or set check_features: false. A model that declares no features is not checked. Only names are checked: dtypes are not stored.

Why it works this way: ADR 012.

Testing your model (Spark-free)¶

A MockDF is a tiny stand-in for a Spark DataFrame that you write in your test. It has only the methods your model calls (here toPandas() and count()), so the model can't tell the difference and you need no Spark to test it:

import pytest
from model import FraudRiskModel

class MockDF:
    def __init__(self, rows):
        self._rows = rows
    def toPandas(self):
        import pandas as pd
        return pd.DataFrame(self._rows)
    def count(self): return len(self._rows)

def test_train_returns_metrics(tmp_path):
    model = FraudRiskModel()
    rows = [
        {"amount_sum_30d": 100, "txn_count_30d": 5, "avg_amount_30d": 20, "risk_encoded": 0, "is_fraud": 0},
        {"amount_sum_30d": 900, "txn_count_30d": 2, "avg_amount_30d": 450, "risk_encoded": 2, "is_fraud": 1},
    ] * 50
    metrics = model.train(MockDF(rows))
    assert "auc" in metrics
    assert "f1" in metrics
    model.save(str(tmp_path / "model"))
    loaded = FraudRiskModel.load(str(tmp_path / "model"))
    assert loaded is not None

Training in a task¶

Train and register the model from the task's transformations.py; model.py sits next to it and imports as from model import FraudRiskModel. A model registered inside a run records that run's id. The full example is in Transform: Training a model.

A task's config.yaml cannot run the model by naming type: model: a task always runs its transformations.py. The type is not run: ubunye run warns (FutureWarning), and ubunye plan and ubunye validate report it. To use the registry: settings from config.yaml, call ModelTransform from transformations.py; see Registry settings for a model transform.

See Model Registry for the full registry reference.

Abstract enforcement¶

UbunyeModel is a strict ABC. Any class that does not implement all five abstract methods (train, predict, save, load, metadata) will raise TypeError on instantiation:

class IncompleteModel(UbunyeModel):
    def train(self, df): return {}
    # missing predict, save, load, metadata

IncompleteModel()   # TypeError: Can't instantiate abstract class IncompleteModel

This is checked by unit tests and when the model class is instantiated.

Model Contract: UbunyeModel¶