Model Training

---

​

Overview

Chalk Model Training is a declarative Python API for training models on your Chalk data. You define a training function once, then run it in-process while developing, or on a managed container with the CPU, memory, and GPU you ask for.

The API covers:

• Declaring data, resources, and container configuration with a decorator
• Running remotely on managed containers, or locally in-process
• Checkpointing to managed storage that survives container restarts
• Logging metrics that feed the training dashboard
• Registering trained models in the Model Registry

The API spans two packages. The @training decorator and the run handle come from chalkcompute, and the helpers you call from inside a run (checkpoint, last_checkpoint_path, log_metrics, log_rollout) come from chalkpy. Install both where you write and launch your training code:

pip install --upgrade chalkcompute 'chalkpy[runtime,training]'

Your training container needs the same two packages, and there it is worth pinning exact versions. A container image is cached and reused, so an unpinned image can go on running an SDK version you have long since upgraded past locally.

---

​

Defining a training function

Use the @training decorator to declare what the run needs. The function body is ordinary PyTorch.

The example trains on a Dataset that must already exist in your environment and have at least one revision. run_training() deploys the training function to a container and reads the Dataset’s data, but it does not create the Dataset. To create one, run an offline query with a dataset_name, as described in Datasets from offline query, and pass that name as data. The example uses the placeholder name sms_spam_v3.

import os
from dataclasses import dataclass

import polars as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

from chalk.ml.chalk_train import checkpoint, log_metrics
import chalkcompute

FEATURES = ["sms.length", "sms.num_links", "sms.num_digits"]
LABEL = "sms.is_spam"

@dataclass
class SpamConfig:
    lr: float = 1e-3
    epochs: int = 10
    batch_size: int = 64

class SpamNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(nn.LazyLinear(64), nn.ReLU(), nn.Linear(64, 2))

    def forward(self, x):
        return self.layers(x)

@chalkcompute.training(
    data="sms_spam_v3",
    cpu="8",
    memory="32Gi",
    gpu="nvidia-a10g",
    image=chalkcompute.Image.debian_slim("3.12").pip_install([
        "chalkcompute==2.13.8",
        "chalkpy[runtime,training]==2.165.0",
        "torch",
        "torchvision",
        "safetensors",
        "polars",
    ]),
    max_retries=2,
)
def train_spam(df, config: SpamConfig):
    if hasattr(df, "collect"):
        df = df.collect()
    df = df.select(FEATURES + [LABEL])
    loader = DataLoader(
        df.to_torch("dataset", label=LABEL, dtype=pl.Float32),
        batch_size=config.batch_size,
        shuffle=True,
    )
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = SpamNet().to(device)
    opt = torch.optim.Adam(model.parameters(), lr=config.lr)
    for epoch in range(config.epochs):
        for x, y in loader:
            x, y = x.to(device), y.to(device).long()
            loss = F.cross_entropy(model(x), y)
            opt.zero_grad()
            loss.backward()
            opt.step()
        if os.environ.get("CHALK_TRAINING_RUN_ID"):
            log_metrics({"loss": loss.item()}, tags={"epoch": str(epoch)})
    if os.environ.get("CHALK_TRAINING_RUN_ID"):
        checkpoint(model)
    return model

sms_spam_v3 is a placeholder. Replace it with the name of a Dataset in your environment, and set FEATURES and LABEL to its column names. A Dataset built from an offline query names each column after its feature, such as sms.length. The label column holds integers 0 and 1, and the feature columns hold numbers.

The select line keeps only the columns the model trains on. In a training run, the data also includes the primary key, a timestamp, and Chalk metadata columns, which local() does not return. dtype=pl.Float32 matches PyTorch’s default weight type, and it converts the label to a float as well, which is why the loop converts y back to an integer with .long() for cross_entropy. nn.LazyLinear sets its input width from the first batch, so SpamNet does not need to know the number of feature columns in advance. The final checkpoint() call saves the trained model to the run, so that run.latest_checkpoint() in Remotely returns it.

​

Decorator parameters

data is required, as is at least one of cpu, memory, or gpu. A run declaring no resources at all is rejected when it is created.

Parameter	Description
data	Name of an existing Dataset to train on. The run pins an immutable revision of it, so a rerun reads the same rows.
cpu, memory, gpu	Resource requirements, each written as a string: cpu="8", memory="32Gi", gpu="nvidia-a10g". See Chalk Compute for the values these accept.
image	Container image for the run. Use Chalk’s Image builder, or a Docker image with your dependencies already installed.
env	Environment variables injected into the container at runtime.
secrets	A list of Secret references, built with Secret.from_chalk_env() or Secret.from_chalk_integration().
name	Run name. Generated if omitted.
metadata	Key-value metadata stored on the run, for tracking and filtering.
max_retries	How many times to retry on failure or container crash.

​

Function signature

def train_fn(df, config: YourConfigClass):
    ...
    return model

• df is the resolved training data, as a Polars frame. run_training() passes a LazyFrame and local() passes an eager DataFrame, and the conversion methods differ between the two: to_torch(), to_pandas(), and to_arrow() exist only on the eager frame, while collect() exists only on the lazy one. Open with the line below so that the rest of the function receives an eager DataFrame either way:

  if hasattr(df, "collect"):
      df = df.collect()
  

• config is a @dataclass you define, or a plain dict. It is serialized to a protobuf Struct and passed to the function. In a remote run, the function receives an instance of your dataclass only when its config parameter is annotated with it, as in config: SpamConfig. Without the annotation, it receives a dict.

• The return value is the trained model.

---

​

Running training

​

Remotely

run = train_spam.run_training(config=SpamConfig(lr=1e-3, epochs=10))
run.wait()

artifact = run.latest_checkpoint()

run_training() builds a deploy image containing your training function’s source, provisions a managed container with the resources, environment variables, and secrets you declared, and resolves the data source inside it. It returns a TrainingRunHandle immediately, without waiting for the run to finish.

​

Locally

model = train_spam.local(config=SpamConfig(lr=1e-3, epochs=2))

local() runs the function in-process, with no container and no remote infrastructure, and returns the model directly. Use it while developing, then switch to run_training() once the code works.

checkpoint() and log_metrics() raise an error when called outside a training run, and local() does not start one. Chalk sets the CHALK_TRAINING_RUN_ID environment variable in every training container, so the example above calls log_metrics() only when that variable is set. Guard checkpoint() the same way in a function you run with local().

It reads the latest revision of the named Dataset, where run_training() pins an immutable revision for the life of the run.

---

​

Run lifecycle

run_training() returns a TrainingRunHandle for observing and controlling the run.

Member	Description
run_id	Property holding the run’s id.
status()	Returns the current status.
wait(poll_interval=5.0, timeout=None)	Blocks until the run reaches a terminal state, polling every poll_interval seconds, and returns the handle. Raises TrainingError if the run failed or was canceled, or TimeoutError if timeout seconds elapse first.
latest_checkpoint()	Returns the most recent checkpoint artifact for the run, or None if none has been uploaded.
cancel()	Cancels the run. The platform terminates the container.

---

​

Checkpointing and resume

Checkpoints let a run pick up where it left off after a container restart.

from safetensors.torch import load_file

from chalk.ml.chalk_train import checkpoint, last_checkpoint_path

ckpt = last_checkpoint_path()
if ckpt is not None:
    model.load_state_dict(load_file(ckpt))

for epoch in range(config.epochs):
    # ... training loop ...
    if epoch % 5 == 0:
        checkpoint(model, metadata={"epoch": epoch})

checkpoint() serializes the model, uploads it, and records it as a model artifact on the run. It takes the model plus four optional arguments: additional_files to store alongside the model, metadata recorded on the artifact, run_name, and path.

last_checkpoint_path() returns the path to the most recent checkpoint file, or None if there is not one yet. checkpoint() saves a PyTorch model in the safetensors format, so load it with safetensors.torch.load_file(), and include safetensors in the training image. Checkpoints outlive the container that wrote them, so a run that restarts picks up from the last one rather than starting over.

​

Naming checkpoints for experiment tracking

Checkpoint storage is scoped to the name set in @training, not to an individual run. Every run sharing a name writes into the same place, which is what lets you run several experiments against one model at the same time and compare them afterwards. When name is omitted, each run gets a generated name, and so its own checkpoint storage.

Each checkpoint lands under its own name inside that storage. By default it is the run’s id, which is accurate but tells you nothing later. Pass path to name it yourself:

checkpoint(model, path="experiment-one")

Run a sweep with path="experiment-one", path="experiment-two", and so on, and the checkpoints are labelled by what you were testing rather than by an opaque id. A latest entry alongside them records the most recent checkpoint, and that is what last_checkpoint_path() reads.

​

Continuing from an earlier run

To continue training from an earlier run’s model, set name in @training to a fixed value:

@chalkcompute.training(name="spam-classifier", data="sms_spam_v3", cpu="8", memory="32Gi")

Each run of the function then uses the same checkpoint storage, and last_checkpoint_path() returns the most recent checkpoint, whichever run wrote it. The loading code above then starts each run from the previous run’s model. To train on newer data, change data and run the function again. To start from a model in the Model Registry instead, see Fine-tuning an existing model.

---

​

Metrics

Call log_metrics() from inside the training function with a dictionary of metric names and values, plus tags describing where in the run the measurement came from.

from chalk.ml.chalk_train import log_metrics

for epoch in range(config.epochs):
    # ... training loop ...
    log_metrics({"accuracy": accuracy, "loss": loss}, tags={"split": "test", "epoch": str(epoch)})

​

Supported names and tags

The dashboard recognizes a fixed vocabulary. These metric names are supported:

accuracy, loss, precision, recall, f1, mae, mse, rmse, r2, reward, reward_std, pass_rate, kl, entropy, grad_norm, learning_rate, response_length

And these tags:

Tag	Description
split	Which subset of the data the measurement covers, such as train or test.
class	The class the measurement covers, for per-class metrics.
threshold	The decision threshold the measurement was taken at.
optimizer	The optimizer in use, such as adam or sgd.
model_version	The model version the measurement belongs to.
epoch	The training epoch.
step	The training step.

Tagging epoch or step throughout a run lets the dashboard use either as the x-axis of its plots. Tagging split, class, threshold, optimizer, or model_version lets you compare across those values with Group by.

​

Custom metrics

chalk.log_custom_metric sends a metric to your own sink rather than the training dashboard. It emits through StatsD, so the environment needs StatsD export configured before any of these measurements leave the run. See Metrics Export.

import chalk

chalk.log_custom_metric(
    "grad_norm",
    3.14,
    tags={"split": "validation"},
)

That is emitted as chalk.custom_metrics.grad_norm(3.14), with its tags attached.

---

​

Rollouts

For reinforcement learning, log_rollout() records one episode as a trace: the conversation between the model and its environment, plus the reward it earned. Each rollout appears on the run’s Traces tab, where you can read what the model did on a task and why it scored what it did. This works for any model, including the one being trained.

from chalk.ml.chalk_train import log_rollout

log_rollout(
    messages,               # the episode, in order
    reward=1.0,
    step=7,
    task_id="qrdata/112",
    metrics={"correct": 1.0},
)

messages uses the OpenAI chat format: each message has a role and content, and may carry tool_calls or tool_call_id. Each assistant message becomes a model turn, and the messages that follow it are the environment’s responses. Pass rollout_id (32 hex characters) to reuse your framework’s own rollout id as the trace id.

---

​

Models other than PyTorch

@training works with any serializable model. This example trains an XGBoost classifier. As in the PyTorch example, it reads a Dataset that must already exist, here the placeholder user_churn_v2:

import os
from dataclasses import dataclass

import xgboost as xgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

from chalk.ml.chalk_train import log_metrics
import chalkcompute

FEATURES = ["user.sessions_30d", "user.tickets_90d", "user.tenure_days"]
LABEL = "user.churned"

@dataclass
class ChurnConfig:
    num_rounds: int = 200
    max_depth: int = 6
    learning_rate: float = 0.1

@chalkcompute.training(
    data="user_churn_v2",
    cpu="16",
    memory="64Gi",
    image=chalkcompute.Image.debian_slim("3.12").pip_install([
        "chalkcompute==2.13.8",
        "chalkpy[runtime,training]==2.165.0",
        "xgboost",
        "pandas",
        "scikit-learn",
    ]),
)
def train_churn(df, config: ChurnConfig):
    if hasattr(df, "collect"):
        df = df.collect()
    pdf = df.select(FEATURES + [LABEL]).to_pandas()
    X = pdf[FEATURES]
    y = pdf[LABEL]
    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2)

    dtrain = xgb.DMatrix(X_train, label=y_train)
    dval = xgb.DMatrix(X_val, label=y_val)

    model = xgb.train(
        {
            "max_depth": config.max_depth,
            "learning_rate": config.learning_rate,
            "objective": "binary:logistic",
            "eval_metric": "auc",
        },
        dtrain,
        num_boost_round=config.num_rounds,
        evals=[(dval, "val")],
    )

    if os.environ.get("CHALK_TRAINING_RUN_ID"):
        log_metrics(
            {"accuracy": accuracy_score(y_val, model.predict(dval) > 0.5)},
            tags={"split": "test", "step": str(config.num_rounds)},
        )

    return model

Replace user_churn_v2 with the name of a Dataset in your environment, and set FEATURES and LABEL to its column names, as in the PyTorch example. The label column holds 0 and 1, and the feature columns hold numbers.

---

​

Environment variables and secrets

import chalkcompute

@chalkcompute.training(
    env={"NCCL_DEBUG": "INFO", "TOKENIZERS_PARALLELISM": "false"},
    secrets=[
        chalkcompute.Secret.from_chalk_env("hf_token"),
        chalkcompute.Secret.from_chalk_env("wandb_api_key"),
    ],
    ...
)

Environment variables are injected directly. secrets takes Secret references rather than plain names: Secret.from_chalk_env() resolves a standalone secret from Chalk’s secret store, and Secret.from_chalk_integration() resolves the credentials of a named integration. Both are exposed as environment variables in the container, and their values are never written to the run record.

---

​

Registering a model

Promote an artifact from a completed run into the Model Registry with ChalkClient.promote_model_artifact:

from chalk.client import ChalkClient

client = ChalkClient()

client.promote_model_artifact(
    name="RiskModel",
    run_id=run.run_id,
    aliases=["latest"],
)

name is the model namespace to promote into, and is the only required argument. Identify the artifact with model_artifact_id, run_id, or run_name. When a run produced several checkpoints, criterion selects among them; the latest artifact in the run is used if you omit it. aliases assigns version labels such as v1.0 or latest.

---

​

Fine-tuning an existing model

To start training from an existing model’s weights, load the model’s files at the start of the training function, then train on new data as usual. This example fine-tunes a version of SpamModel that is already in the Model Registry, such as the artifact you promoted in Registering a model.

The example trains on a Dataset that must already exist, like the first example. Here the placeholder is sms_spam_v4. FEATURES must list the same columns, in the same order, as the data the base model was trained on, and the model class in the function must match the one that produced the registered weights.

import os
from dataclasses import dataclass

import polars as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file
from torch.utils.data import DataLoader

from chalk.client.client_grpc import ChalkGRPCClient
from chalk.ml.chalk_train import checkpoint, log_metrics
import chalkcompute

FEATURES = ["sms.length", "sms.num_links", "sms.num_digits"]
LABEL = "sms.is_spam"

@dataclass
class FineTuneConfig:
    base_model: str = "SpamModel"
    base_version: int = 1
    lr: float = 1e-4
    epochs: int = 3
    batch_size: int = 64

class SpamNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(nn.LazyLinear(64), nn.ReLU(), nn.Linear(64, 2))

    def forward(self, x):
        return self.layers(x)

@chalkcompute.training(
    data="sms_spam_v4",
    cpu="8",
    memory="32Gi",
    gpu="nvidia-a10g",
    image=chalkcompute.Image.debian_slim("3.12").pip_install([
        "chalkcompute==2.13.8",
        "chalkpy[runtime,training]==2.165.0",
        "torch",
        "torchvision",
        "safetensors",
        "polars",
    ]),
)
def fine_tune_spam(df, config: FineTuneConfig):
    if hasattr(df, "collect"):
        df = df.collect()
    df = df.select(FEATURES + [LABEL])
    loader = DataLoader(
        df.to_torch("dataset", label=LABEL, dtype=pl.Float32),
        batch_size=config.batch_size,
        shuffle=True,
    )

    base = ChalkGRPCClient().download_model_artifact(config.base_model, config.base_version)
    model = SpamNet()
    model.load_state_dict(load_file(base.downloaded_model_files[0]))

    opt = torch.optim.Adam(model.parameters(), lr=config.lr)
    for epoch in range(config.epochs):
        for x, y in loader:
            loss = F.cross_entropy(model(x), y.long())
            opt.zero_grad()
            loss.backward()
            opt.step()
        if os.environ.get("CHALK_TRAINING_RUN_ID"):
            log_metrics({"loss": loss.item()}, tags={"epoch": str(epoch)})

    if os.environ.get("CHALK_TRAINING_RUN_ID"):
        checkpoint(model, metadata={"base_model": config.base_model, "base_version": config.base_version})
    return model

Run it the same way as any other training function:

run = fine_tune_spam.run_training(config=FineTuneConfig(base_model="SpamModel", base_version=1))
run.wait()

How the example works:

• download_model_artifact(name, version) downloads the files of a registered model version. By default it writes them to a <name>-<version> folder in the current directory, and download_dir chooses a different location. downloaded_model_files lists the paths of the model files.
• checkpoint() saves a PyTorch model in the safetensors format, so a model promoted from a training run loads with safetensors.torch.load_file(). A model registered from other file formats needs the loader that matches its format.
• load_state_dict() copies the registered weights into a new SpamNet, and the training loop then updates them. nn.LazyLinear takes its input size from the loaded weights, so the model needs no forward pass first.
• The learning rate is 10 times lower than in the first example, so the new data adjusts the base weights gradually.
• checkpoint() runs only inside a training run, so the example guards it with CHALK_TRAINING_RUN_ID, as it does for log_metrics(). The metadata records which registered version the new checkpoint came from.

---

​

See also

• Model Registry: track and version trained models.
• Model Deployments: serve a registered model.
• Datasets: create the datasets a training run reads from.