Skip to main content

Model Registry Guide

Save, version, and manage trained models using Michelangelo AI's model packaging system.

Overview​

Michelangelo AI's model packager turns your trained model into self-contained, versioned artifacts ready for serving or sharing. The packager handles dependency bundling, schema validation, and Triton configuration generation automatically.

What the model packager provides:

  • Dual format packaging -- deployable (Triton-ready) and raw (developer-facing) artifacts
  • Schema validation -- input/output contracts enforced at packaging time
  • Dependency bundling -- auto-packages Python modules your model needs at inference time
  • Built-in testing -- raw packages are validated with sample data before they are written
  • Triton compatibility -- deployable packages work directly with NVIDIA Triton Inference Server

Core Concepts​

Before diving in, here are the key components you will work with:

ConceptWhat It IsModule
ModelAbstract base class your model must implement (save, load, predict)michelangelo.lib.model_manager.interface.custom_model
ModelSchemaDefines input/output feature names, types, and shapesmichelangelo.lib.model_manager.schema
CustomTritonPackagerCreates deployable and raw model packagesmichelangelo.lib.model_manager.packager.custom_triton
load_raw_modelLoads a raw model package for testing or fine-tuningmichelangelo.lib.model_manager.serde.model

The packager produces two complementary artifacts:

ArtifactPurposeCreated By
Deployable model packageTriton Inference Server deploymentcreate_model_package()
Raw model packageTesting, fine-tuning, reproducibilitycreate_raw_model_package()

Quick Start: Package Your First Model​

This end-to-end walkthrough takes you from a trained model to a verified package in four steps.

Step 1: Implement the Model interface​

Create a class that extends Model with three methods: save, load, and predict.

import os
import numpy as np
from michelangelo.lib.model_manager.interface.custom_model import Model


class EchoModel(Model):
    """Minimal model that returns inputs unchanged."""

    def save(self, path: str) -> None:
        os.makedirs(path, exist_ok=True)
        with open(os.path.join(path, "config.txt"), "w") as f:
            f.write("echo-model-v1")

    @classmethod
    def load(cls, path: str) -> "EchoModel":
        with open(os.path.join(path, "config.txt")) as f:
            f.read()
        return cls()

    def predict(self, inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
        a = inputs["a"].astype(np.int32)
        return {"response": a, "doubled": (2 * a).astype(np.int32)}

Step 2: Define the schema​

Declare what your model expects as input and what it produces as output.

from michelangelo.lib.model_manager.schema import DataType, ModelSchema, ModelSchemaItem

schema = ModelSchema(
    input_schema=[
        ModelSchemaItem(name="a", data_type=DataType.INT, shape=[1]),
    ],
    output_schema=[
        ModelSchemaItem(name="response", data_type=DataType.INT, shape=[1]),
        ModelSchemaItem(name="doubled", data_type=DataType.INT, shape=[1]),
    ],
)

Step 3: Package the model​

Save your model artifacts, then create both package types.

from michelangelo.lib.model_manager.packager.custom_triton import CustomTritonPackager

# Save model artifacts first
model = EchoModel()
model.save("/tmp/echo-artifacts")

# Create the packager
packager = CustomTritonPackager()

# Deployable package (for Triton serving)
deployable_path = packager.create_model_package(
    model_path="/tmp/echo-artifacts",
    model_class="myproject.models.EchoModel",
    model_schema=schema,
    model_name="echo-model",
    dest_model_path="/tmp/echo-deployable",
)

# Raw package (for testing and fine-tuning)
sample_data = [
    {"a": np.array([1], dtype=np.int32)},
    {"a": np.array([5], dtype=np.int32)},
]

raw_path = packager.create_raw_model_package(
    model_path="/tmp/echo-artifacts",
    model_class="myproject.models.EchoModel",
    model_schema=schema,
    sample_data=sample_data,
    dest_model_path="/tmp/echo-raw",
    requirements=["numpy"],
)

Step 4: Verify the package​

Load the raw package and run a prediction to confirm everything works.

from michelangelo.lib.model_manager.serde.model import load_raw_model

loaded = load_raw_model("/tmp/echo-raw")
result = loaded.predict({"a": np.array([42], dtype=np.int32)})
print(result)
# {'response': array([42], dtype=int32), 'doubled': array([84], dtype=int32)}

You now have a deployable Triton package at /tmp/echo-deployable and a verified raw package at /tmp/echo-raw.

API Reference​

Model Interface​

All custom models must extend the Model abstract base class:

from michelangelo.lib.model_manager.interface.custom_model import Model
MethodSignatureDescription
savesave(self, path: str)Serialize model artifacts to a directory
loadload(cls, path: str) -> ModelClass method that loads and returns a ready-to-use model instance
predictpredict(self, inputs: dict[str, ndarray]) -> dict[str, ndarray]Run inference; keys must match the model schema
tip

Avoid using pickle or torch.save directly for persistence. Prefer format-specific serialization methods (e.g., state_dict for PyTorch, SavedModel for TensorFlow) for better compatibility and security.

ModelSchema​

Defines the contract between your model and the serving infrastructure.

from michelangelo.lib.model_manager.schema import DataType, ModelSchema, ModelSchemaItem

ModelSchemaItem fields:

FieldTypeDefaultDescription
namestr(required)Feature name, used as the key in input/output dictionaries
data_typeDataTypeDataType.UNKNOWNThe data type of the feature
shapelist[int]NoneShape following NumPy conventions; use -1 for variable-length dimensions
optionalboolNoneIf True, the feature may be omitted from input data

Supported data types:

DataTypePython/NumPy TypeDescription
BOOLEANboolBoolean values
STRINGstr / bytesText data. NumPy string, bytes, or object arrays are accepted; maps to Triton BYTES
BYTEint88-bit signed integer
CHARuint88-bit unsigned integer
SHORTint1616-bit signed integer
INTint3232-bit signed integer
LONGint6464-bit signed integer
FLOATfloat3232-bit floating point
DOUBLEfloat6464-bit floating point

Shape examples:

ShapeMeaning
[1]Scalar value
[10]1D array of length 10
[10, 5]2D array (10 rows, 5 columns)
[-1]Variable-length 1D array

CustomTritonPackager​

from michelangelo.lib.model_manager.packager.custom_triton import CustomTritonPackager

Constructor:

ParameterDefaultDescription
custom_batch_processingFalseIf True, your model handles batching internally

create_model_package() parameters:

ParameterRequiredDefaultDescription
model_pathYes--Path to saved model artifacts
model_classYes--Fully qualified Python class name (e.g., "mypackage.models.MyModel")
model_schemaYes--ModelSchema instance defining inputs and outputs
model_nameNoDerived from classDisplay name in Michelangelo AI Studio
dest_model_pathNoAuto temp dirOutput directory for the package
model_revisionNoNoneRevision number for versioning
model_path_source_typeNoStorageType.LOCALStorage backend type
include_import_prefixesNoNone (all imports)List of module prefixes to bundle
additional_import_prefixesNoNoneExtra module prefixes whose source files are bundled, for dynamic imports that static analysis misses
triton_parametersNoNoneKey-value strings written as parameters blocks in config.pbtxt
sample_dataNoNoneSample inputs written to metadata/sample_data.json for reference

create_raw_model_package() parameters:

ParameterRequiredDefaultDescription
model_pathYes--Path to saved model artifacts
model_classYes--Fully qualified Python class name
model_schemaYes--ModelSchema instance defining inputs and outputs
sample_dataYes--List of sample inputs for validation
dest_model_pathNoAuto temp dirOutput directory for the package
model_path_source_typeNoStorageType.LOCALStorage backend type
requirementsNoNoneDependencies as a list or path to requirements.txt
include_import_prefixesNoNone (all imports)List of module prefixes to bundle
additional_import_prefixesNoNoneExtra module prefixes whose source files are bundled, for dynamic imports that static analysis misses

load_raw_model​

from michelangelo.lib.model_manager.serde.model import load_raw_model

model = load_raw_model("/path/to/raw/package")

Returns an instance of your Model subclass, fully loaded and ready for inference.

note

load_raw_model supports Custom Python models (RawModelType.CUSTOM_PYTHON) and PyTorch/Lightning raw packages (RawModelType.TORCH, RawModelType.LIGHTNING). Other types, such as HuggingFace, raise NotImplementedError.

Advanced Topics​

PyTorch Model Example​

The packager works with any framework. Here is an example using PyTorch internally while conforming to the numpy-based Model interface:

import numpy as np
import torch
from michelangelo.lib.model_manager.interface.custom_model import Model


class TorchClassifier(Model):
    """PyTorch model with numpy I/O for Model Manager."""

    def __init__(self):
        self.net = torch.nn.Linear(4, 2)

    def save(self, path: str) -> None:
        import os
        os.makedirs(path, exist_ok=True)
        torch.save(self.net.state_dict(), os.path.join(path, "model.pt"))

    @classmethod
    def load(cls, path: str) -> "TorchClassifier":
        import os
        obj = cls()
        state = torch.load(os.path.join(path, "model.pt"), weights_only=True)
        obj.net.load_state_dict(state)
        obj.net.eval()
        return obj

    def predict(self, inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
        x = torch.from_numpy(inputs["x"].astype(np.float32))
        with torch.no_grad():
            out = self.net(x)
        return {"prediction": out.numpy()}

Package it with a matching schema:

schema = ModelSchema(
    input_schema=[
        ModelSchemaItem(name="x", data_type=DataType.FLOAT, shape=[1, 4]),
    ],
    output_schema=[
        ModelSchemaItem(name="prediction", data_type=DataType.FLOAT, shape=[1, 2]),
    ],
)

sample_data = [
    {"x": np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)},
]

packager = CustomTritonPackager()

raw_path = packager.create_raw_model_package(
    model_path="/tmp/torch-artifacts",
    model_class="myproject.models.TorchClassifier",
    model_schema=schema,
    sample_data=sample_data,
    requirements=["numpy", "torch"],
)

Custom Batch Processing​

By default, Triton handles batching automatically and your predict method receives individual samples. If your model handles batching internally, enable custom batch processing:

packager = CustomTritonPackager(custom_batch_processing=True)

When enabled, inputs include an additional leading batch dimension. For example, if the schema specifies shape [n, m], the actual input shape will be [batch_size, n, m].

Model Package Formats​

Deployable Format (Triton-Compatible)​

model_name/
├── 0/
│   ├── model.py                    # Triton Python backend entry point
│   ├── user_model.py               # Your model implementation
│   ├── model_class.txt             # Fully qualified Python class path
│   ├── model/                      # Model binaries (your saved artifacts)
│   └── myproject/models/...        # Auto-packaged runtime dependencies
├── metadata/
│   ├── schema.yaml                 # Input/output schema
│   └── sample_data.json            # Only when sample_data is passed
└── config.pbtxt                    # Triton configuration (I/O schema, batching)
FilePurpose
model.pyTriton Python backend entry point
user_model.pyYour model's forward pass and inference logic
model_class.txtFully qualified Python class path
model/Your saved model artifacts
metadata/schema.yamlThe model schema
config.pbtxtTriton configuration (I/O schema, batching, instances)

Raw Format (Developer-Facing)​

raw_package/
├── metadata/
│   ├── type.yaml                   # Model type (custom-python)
│   ├── schema.yaml                 # Input/output schema
│   └── sample_data.json            # Sample data for testing
├── model/                          # Model binaries (your saved artifacts)
├── defs/
│   ├── model_class.txt             # Fully qualified class path
│   └── myproject/models/...        # Runtime code dependencies
└── dependencies/
    └── requirements.txt            # Only when requirements are passed

Integration with Uniflow Workflows​

Package model registration as a task in your ML pipeline:

import michelangelo.uniflow.core as uniflow
from michelangelo.lib.model_manager.packager.custom_triton import CustomTritonPackager
from michelangelo.lib.model_manager.schema import DataType, ModelSchema, ModelSchemaItem
from michelangelo.uniflow.plugins.ray import RayTask


@uniflow.task(config=RayTask(head_cpu=2, head_memory="4Gi"))
def package_model(model_path: str, model_class: str):
    """Package a trained model for deployment."""
    packager = CustomTritonPackager()

    schema = ModelSchema(
        input_schema=[
            ModelSchemaItem(name="feature", data_type=DataType.STRING, shape=[1]),
        ],
        output_schema=[
            ModelSchemaItem(name="response", data_type=DataType.STRING, shape=[1]),
        ],
    )

    deployable_path = packager.create_model_package(
        model_path=model_path,
        model_class=model_class,
        model_schema=schema,
    )

    return deployable_path

This task can be chained after a training task in a workflow:

@uniflow.workflow()
def train_and_package(dataset_id: str):
    model_path = train_model(dataset_id)
    package_path = package_model(model_path, "myproject.models.MyModel")
    return package_path

Troubleshooting​

ValueError: model_class is required​

The model_class parameter must be a non-empty string containing the fully qualified Python class path (e.g., "mypackage.models.MyModel").

ValueError: model_schema is required​

A non-empty model_schema argument must be provided (passing None raises this error).

Schema validation errors​

Ensure your sample data matches the schema exactly:

  • Each sample must contain exactly the input features named in input_schema -- missing or extra keys are rejected
  • Array shapes must match the schema's shape field
  • Array dtypes must be compatible with the schema's data_type
note

For DataType.STRING, sample data can use NumPy string (np.array(["hello"])), bytes (np.array([b"hello"])), or object arrays. Numeric or boolean arrays are rejected.

NotImplementedError: The loader for ... model is not supported yet​

load_raw_model does not yet have a loader for the model type recorded in the package's metadata/type.yaml (for example, HuggingFace). Custom Python, PyTorch, and Lightning raw packages are supported.

Model class validation fails​

Your model class must:

  • Be importable from the current Python environment
  • Extend michelangelo.lib.model_manager.interface.custom_model.Model
  • Implement all three abstract methods: save, load, and predict

Register a Revision​

A Revision is a versioned snapshot of a Model resource. One Model can have many Revisions — each one represents a distinct training run or artifact version. When you deploy a model, you target a specific Revision, not the Model directly. This lets you roll out new versions and roll back independently of the model definition itself.

Create a revision.yaml​

apiVersion: michelangelo.api/v2
kind: Revision
metadata:
  name: my-model-v1
  namespace: my-project
spec:
  baseType:
    kind: Model
    apiVersion: michelangelo.api/v2
  baseResource:
    name: my-model
    namespace: my-project
  owner:
    name: <your-username>

The baseResource.name and baseResource.namespace must match an existing Model resource in your project. owner.name is the username configured by your platform operator.

Apply and verify​

Apply the revision to the control plane:

ma revision apply -f revision.yaml

List revisions in your namespace to confirm it was registered:

ma revision get -n my-project

When to create a new Revision vs. update​

Create a new Revision whenever you have a new training run or a new artifact — each Revision should correspond to a distinct, reproducible artifact. Update an existing Revision only to fix metadata (such as the owner field); changing the underlying artifact warrants a new Revision so your deployment history stays traceable.

Next Steps​