Model Registry

The Model Registry integration provides version control, staging management, and lifecycle tracking for machine learning models. It acts as the central catalog for all trained models within a tenant, supporting model versioning, stage transitions, metadata tagging, and artifact storage backed by MLflow.

Registry Architecture

The Model Registry is backed by MLflow and accessed through the ML Service:

AI Service (Registry API) --> ML Service --> MLflow Tracking Server --> Artifact Store (MinIO/S3)

Model Lifecycle Stages

Stage	Description	Transition
`none`	Model artifact exists but is not staged	Automatic on training completion
`staging`	Model promoted for validation and testing	Manual or automated promotion
`production`	Model approved for live serving	Approval workflow
`archived`	Model retired from active use	Manual archival

Register Model

Registers a trained model artifact in the registry:

POST /api/v1/ml/registry/models

{
  "name": "churn-predictor",
  "artifact_id": "model-xyz789",
  "description": "XGBoost churn prediction model",
  "tags": {
    "algorithm": "xgboost",
    "dataset": "customer_churn_v2",
    "owner": "ml-team"
  },
  "metrics": {
    "f1_score": 0.895,
    "auc_roc": 0.96
  }
}

List Models

GET /api/v1/ml/registry/models?tenant_id=acme-corp

Response

{
  "models": [
    {
      "name": "churn-predictor",
      "latest_version": 3,
      "production_version": 2,
      "staging_version": 3,
      "created_at": "2025-02-01T08:00:00Z",
      "last_updated": "2025-03-15T10:00:00Z"
    }
  ]
}

Get Model Version

GET /api/v1/ml/registry/models/:name/versions/:version

Response

{
  "name": "churn-predictor",
  "version": 3,
  "stage": "staging",
  "artifact_uri": "s3://ml-artifacts/churn-predictor/v3/model.pkl",
  "metrics": {
    "f1_score": 0.912,
    "auc_roc": 0.97
  },
  "parameters": {
    "n_estimators": 200,
    "max_depth": 8,
    "learning_rate": 0.05
  },
  "tags": {
    "algorithm": "xgboost",
    "training_job": "train-abc123"
  },
  "created_at": "2025-03-15T10:00:00Z"
}

Transition Stage

Promotes or demotes a model version to a different lifecycle stage:

POST /api/v1/ml/registry/models/:name/versions/:version/stage

{
  "stage": "production",
  "comment": "Validated on holdout set with 0.91 F1 score"
}

Compare Versions

Compares metrics across model versions for selection:

GET /api/v1/ml/registry/models/:name/compare?versions=2,3

Response

{
  "comparison": [
    {
      "version": 2,
      "metrics": {"f1_score": 0.895, "auc_roc": 0.96},
      "stage": "production"
    },
    {
      "version": 3,
      "metrics": {"f1_score": 0.912, "auc_roc": 0.97},
      "stage": "staging"
    }
  ],
  "recommendation": "Version 3 shows improvement across all metrics"
}

Artifact Storage

Model artifacts are stored in the object store (MinIO in development, S3/Azure Blob in production):

Artifact Type	File	Purpose
Model binary	`model.pkl` or `model.pt`	Serialized model weights
Preprocessor	`preprocessor.pkl`	Feature preprocessing pipeline
Config	`config.json`	Model hyperparameters and metadata
Metrics	`metrics.json`	Training and validation metrics
Requirements	`requirements.txt`	Python dependency snapshot

Configuration

Environment Variable	Default	Description
`MLFLOW_TRACKING_URI`	`http://mlflow:5000`	MLflow tracking server URL
`MLFLOW_ARTIFACT_ROOT`	`s3://ml-artifacts`	Artifact storage root
`ML_REGISTRY_CACHE_TTL`	`300`	Model metadata cache TTL in seconds

Model Serving Experiment Tracking