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.

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 log_metrics
from chalkcompute import Image, training

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

@training(
    data="sms_spam_v3",
    cpu="8",
    memory="32Gi",
    gpu="nvidia-a10g",
    image=Image.debian_slim("3.12").pip_install([
        "chalkcompute==2.11.10",
        "chalkpy[runtime,training]==2.157.22",
        "torch",
        "torchvision",
        "safetensors",
        "polars",
    ]),
    max_retries=2,
)
def train_spam(df, config: SpamConfig):
    if hasattr(df, "collect"):
        df = df.collect()
    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)})
    return model

sms_spam_v3 is a placeholder. Replace it with the name of a Dataset in your environment that has an integer label column with values 0 and 1, and numeric values in every other column. to_torch() reads every column other than label as a model input. 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.

​

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 training function’s name, 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.

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.

---

​

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:

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
from chalkcompute import Image, training

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

@training(
    data="user_churn_v2",
    cpu="16",
    memory="64Gi",
    image=Image.debian_slim("3.12").pip_install([
        "chalkcompute==2.11.10",
        "chalkpy[runtime,training]==2.157.22",
        "xgboost",
        "pandas",
        "scikit-learn",
    ]),
)
def train_churn(df, config: ChurnConfig):
    if hasattr(df, "collect"):
        df = df.collect()
    pdf = df.to_pandas()
    X = pdf.drop(columns=["churned"])
    y = pdf["churned"]
    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

As in the PyTorch example, user_churn_v2 is a placeholder. Replace it with a Dataset that has a churned column with values 0 and 1, and numeric values in every other column.

---

​

Environment variables and secrets

from chalkcompute import Secret, training

@training(
    env={"NCCL_DEBUG": "INFO", "TOKENIZERS_PARALLELISM": "false"},
    secrets=[
        Secret.from_chalk_env("hf_token"),
        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.

---

​

See also

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