Model Registry

With Chalk you can easily register and load machine learning models into your deployments. You can then run inference on these models. This guide covers how to integrate models into your Chalk applications, including loading existing models and running inference.

Chalk has a few key model concepts:

• Model Name: A namespace for a model (which might have multiple implementations).
• Model Version: A model that has been registered with a specific version number—these are models that can be loaded into Chalk deployments.
• Model Artifact: The actual files that make up a model, such as weights, configuration files, and tokenizers.
• Model Reference: A reference to a specific model version. Used to load models into deployments and track model performance and feature distributions for a deployed model.

Generally running inference on models in Chalk involves three steps:

1. Registering models in Chalk
2. Including models in Chalk deployments
3. Connecting your models to features for inference

​

Registering Models in Chalk

To add a model to the Chalk model registry, you can either register an existing model or run training code to create a new model. Once registered, models can be versioned and tracked over time. Registering models through training jobs is covered in the Chalk model training docs.

To add a new model namespace to the registry, you can use the client.register_model_namespace method. This namespace can then be used when registering a model version to assign it to this namespace.

from chalk.client import ChalkClient

client = ChalkClient()

client.register_model_namespace(
    name="RiskScoreModel",
    description="Risk score model developed in pytorch"
)

​

Registration Modes

There are three ways to register a model version, depending on how your model is stored and where you want to run it:

​

Registering From Python Model Objects

Models can be registered directly from Python objects, such as scikit-learn or PyTorch models. When you pass a Python model object, Chalk automatically infers the model type, encoding, input/output schema, and package dependencies.

from chalk.client import ChalkClient
from sklearn.ensemble import RandomForestClassifier

client = ChalkClient()

rfc = RandomForestClassifier()
rfc.fit(X_train, y_train)

client.register_model_version(
    name="RiskScoreModel",
    aliases=["v1.0.0"],
    model=rfc,
    metadata={
        "framework": "sklearn",
        "training_date": "2025-07-29",
        "accuracy": 0.89,
        "precision": 0.91
    },
    input_features=[User.age, User.income],
    output_features=[User.churn_risk],
)

When registering from a Python object, input_schema and output_schema are automatically inferred from the model. You can still provide them explicitly if you want to override the inferred schemas.

​

Registering From Local Files

To register a model from local files, provide the path to the model files using model_paths. When registering from files, you must explicitly specify the model_encoding and provide input_schema and output_schema since these cannot be inferred.

from chalk.client import ChalkClient
from chalk.ml import ModelType, ModelEncoding

client = ChalkClient()

client.register_model_version(
    name="FraudDetector",
    model_paths=["./fraud_model.pkl"],
    model_type=ModelType.SKLEARN,
    model_encoding=ModelEncoding.PICKLE,
    input_schema={
        "transaction_amount": float,
        "merchant_category": str,
    },
    output_schema={
        "fraud_probability": float,
    },
    aliases=["v1.0", "production"],
    metadata={"accuracy": 0.96, "training_date": "2025-03-15"},
)

If your model requires additional files for inference, such as tokenizers or configuration files, include them with additional_files. You can also pin the dependencies needed to load and run the model with the dependencies parameter:

client.register_model_version(
    name="SentimentModel",
    model_paths=["./sentiment_model.pt"],
    model_type=ModelType.PYTORCH,
    model_encoding=ModelEncoding.PICKLE,
    additional_files=["./tokenizer.json", "./config.json"],
    input_schema={"text": str},
    output_schema={"sentiment_score": float},
    dependencies=["torch==2.7.1", "transformers==4.40.0"],
)

​

Registering From Cloud Storage

You can register models stored in cloud storage by providing cloud URIs in model_paths along with a source_config for authentication.

From S3:

from chalk.ml import S3SourceConfig, ModelType, ModelEncoding

client.register_model_version(
    name="FraudDetector",
    model_paths=["s3://my-bucket/models/fraud_model.pkl"],
    model_type=ModelType.SKLEARN,
    model_encoding=ModelEncoding.PICKLE,
    input_schema={"transaction_amount": float},
    output_schema={"fraud_probability": float},
    source_config=S3SourceConfig(
        aws_profile="my-profile",
        aws_region="us-east-1",
    ),
)

From GCS:

from chalk.ml import GCSSourceConfig

client.register_model_version(
    name="FraudDetector",
    model_paths=["gs://my-bucket/models/fraud_model.pkl"],
    model_type=ModelType.SKLEARN,
    model_encoding=ModelEncoding.PICKLE,
    input_schema={"transaction_amount": float},
    output_schema={"fraud_probability": float},
    source_config=GCSSourceConfig(gcp_project="my-project"),
)

​

Registering With a Docker Image

Models can also be registered as Docker images for deployment to scaling groups. This mode does not require model serialization—Chalk runs your container as-is. You must provide input_schema and output_schema to define the model’s data contract. See the Model Deployments guide for details on building images and deploying.

client.register_model_version(
    name="ner-model",
    input_schema={"text": str},
    output_schema={"entities": str},
    model_image="ghcr.io/my-org/ner-model:latest",
)

​

Connecting Models to Features

When registering a model version, there are two related but distinct ways to describe the model’s inputs and outputs:

• input_schema / output_schema: Define the model’s raw data contract — what columns or tensors the model expects and produces. Use a dict mapping column names to types for tabular data (float, int, str, bool, or PyArrow types), or a list of (shape, dtype) tuples for tensor data.
• input_features / output_features: Map the model’s inputs and outputs to Chalk features. This is used by make_model_resolver to wire features into the model. Pass feature references like [User.age, User.income] or strings like ["user.age", "user.income"].

client.register_model_version(
    name="ChurnPredictor",
    model=trained_model,
    input_schema={
        "age": int,
        "income": float,
    },
    output_schema={
        "churn_probability": float,
    },
    input_features=[User.age, User.income],
    output_features=[User.churn_risk],
)

When registering from a Python object, input_schema and output_schema are auto-inferred and input_features / output_features are optional. When registering from file paths or Docker images, input_schema and output_schema are required.

​

register_model_version Parameter Reference

​

Types and Enums

See the API reference for full details on ModelType, ModelEncoding, ModelClass, and SourceConfig.

​

Retrieving Model Information

Registered models are tracked in the Chalk Model Registry. They can be loaded through the client and queried for metadata such as versions and performance metrics. They can also be viewed in the Chalk Dashboard. If they’ve been included in a deployment, you will also see feature distributions and model performance over time.

# Get model metadata
model = client.get_model(name="RiskScoreModel")
print(f"Latest version: {model.latest_version}")
print(f"Available versions: {model.versions}")

# Get specific version details
model_v1 = client.get_model(name="RiskScoreModel", version=1)
print(f"Performance: {model_v1.metadata['training_metrics']}")

​

Including Models in Chalk Deployments

To include models in your Chalk deployment, use ModelReference objects. ModelReferences are code-defined objects which should be included in your Chalk code: they are used to connect model versions to model deployments.

from chalk.features import features, _
from chalk.ml import ModelReference
from chalk import functions as F
from chalk import make_model_resolver

# Load a model into the deployment
churn_risk = ModelReference.from_version(
    name="ChurnModel",
    version=2,
)

@features
class User:
    id: int
    age: int
    income: float

    churn_risk: float # define the prediction feature

resolver = make_model_resolver(
    name="churn_model_prediction",
    model=churn_risk,
    input=[User.age, User.income],
    output=[User.churn_risk],
)

Once the model is deployed, metrics for the model will be tracked in the Chalk dashboard, including feature distributions and performance over time.

​

Running Inference on Models

Once models are registered and loaded into your deployment, you can run inference on your models:

from chalk.features import feature, features, _
from chalk.ml import ModelReference
from chalk import functions as F
from chalk import make_model_resolver

# Load a model into the deployment using alias
risk_model_latest = ModelReference.from_alias(
    name="RiskScoreModel",
    alias="latest",
)

# Load a model into the deployment using version
risk_model_v1_0_0 = ModelReference.from_version(
    name="RiskScoreModel",
    version=1,
)

# Reference the loaded model(s) in inference
@features
class User:
    id: int
    age: int
    income: float
    risk_score: float = feature(versions=2) 

make_model_resolver(
    name="user_risk_score_v1_prediction",
    input=[User.age, User.income],
    output=[User.risk_score@1]
)

make_model_resolver(
    name="user_risk_score_v2_prediction",
    input=[User.age, User.income],
    output=[User.risk_score@2]
)