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 four 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 and encoding. Schema and dependency inference are partial: input_schema and output_schema are inferred for scikit-learn, PyTorch, XGBoost, and CatBoost, and dependencies for scikit-learn, PyTorch, and XGBoost. Pass these explicitly when you use another framework, or whenever you want named output columns or pinned versions.

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],
)

For the frameworks listed above, input_schema and output_schema are inferred from the model. Inferred schemas are positional (no column names), so pass them explicitly when you want named columns or when you use a framework that isn’t inferred.

​

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 the @model_handler Decorator

For the fastest path to a deployable model, decorate a class with @model_handler, write a predict method, and pass a trained model object — Chalk builds the serving image, ships your class, and mounts the serialized model for you. See The @model_handler decorator in the Model Deployments guide.

​

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",
)

​

Registering With a Chalk Image

Instead of building and pushing a Docker image yourself, you can submit a chalkcompute Image object and Chalk will build the compute image for you when you register your model. You must provide input_schema and output_schema to define the model’s data contract. See the Model Deployments guide for details on deploying.

client.register_model_version(
    name="ner-model",
    input_schema={"text": str},
    output_schema={"entities": str},
    model_image=chalkcompute.Image().debian_slim("3.11")
        .pip_install(["chalk-remote-call-python"])
        .add_local_file(HANDLER_FILE_PATH, "/app/handler.py", strategy="copy")
        .env({"PYTHONPATH": "/app"})
        .workdir("/app")
        .entrypoint(
            [
                "chalk-remote-call",
                "--handler",
                "handler.handler",
                "--port",
                "8080",
            ]
        )
)

​

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],
)

input_schema and output_schema are always required, but for scikit-learn, PyTorch, XGBoost, and CatBoost objects Chalk infers them so you can omit them; input_features / output_features are optional. When registering from file paths, Docker images, Chalk images, or any other framework, provide input_schema and output_schema explicitly.

​

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]
)