castform
start training book demo
pricing docs learn blog careers
start training book demo
  1. what benchmax offers
  2. step-by-step overview
  3. prerequisites
  4. 1. prepare your dataset
  5. 2. define your environment
  6. 3. launch
  7. what happens next
  8. learn more

sdk overview

start Apr 8, 2026 5 min read

benchmax is castform’s open-source python sdk and the companion to the training platform. it’s the library the platform runs under the hood, and the one you reach for when you want full control: you define your dataset, environment, and rewards in code, then launch a run from your own machine. it ships as two packages: benchmax (the environment runtime: envs, rewards, bundling) and castform (the platform sdk and cli: validation, uploads, launches, corpus and trace workflows).

prefer not to write it by hand? see see our quickstart guide here to scaffold a project with the castform cli (castform setup) and let your coding agent build it for you. the scaffolded main.py owns the whole workflow: python main.py prepares data and validates, python main.py launch starts a run.

what benchmax offers

  • environments as code: extend BaseEnv to define an environment with rewards and tools for your task, or run existing agent harnesses through harbor. see environments.
  • flexible rewards: our rewards and rubrics api allows you to easily define the success criteria, optimized for training. see rewards.
  • automated dataset generation: generate synthetic qa pairs from a corpus.
  • managed training + eval: the castform package provides the interface to the platform to validate, upload and launch runs. you can launch runs on gpus castform provisions, then evaluate against baselines.

step-by-step overview

prerequisites

  • python 3.12 (3.13 is not supported)
  • a castform account and an active castform login session
pip install castform   # pulls in benchmax

1. prepare your dataset

training data is JSONL rows of dicts. the simplest format uses prompt and ground_truth columns (for multi-turn conversations, use messages instead of prompt; see dataset):

dataset = generate_my_examples(n=400)  # your data generation logic
# each example looks like: {"prompt": "...", "ground_truth": "..."}

split = int(len(dataset) * 0.8)
train_data, eval_data = dataset[:split], dataset[split:]

see dataset for format details and data sources.

2. define your environment

an environment tells the platform how to score the model’s output. at minimum you map your dataset rows to prompts and score completed rollouts:

from pathlib import Path

from benchmax.envs import (
    BaseEnv, BaseRollout, DatasetSplit, Example, JsonlDataset,
    canonical_example_id,
)

class MyEnv(BaseEnv):
    max_turns = 1

    async def create_dataset(
        self, split: DatasetSplit, base_dir: Path, *, max_examples: int | None = None
    ):
        return JsonlDataset(
            base_dir / f"{split}.jsonl",
            row_to_example=self._example,
            max_examples=max_examples,
        )

    def _example(self, row: dict) -> Example:
        payload = {
            "prompt_messages": [
                {"role": "system", "content": "answer concisely."},
                {"role": "user", "content": row["prompt"]},
            ],
            "ground_truth": row["ground_truth"],
        }
        return Example(id=canonical_example_id(payload), payload=payload)

    async def compute_reward(self, rollout: BaseRollout) -> dict[str, float]:
        ...  # score rollout.messages against rollout.example_args, return {"correct": ...}

see base env for the full contract, rewards for reward patterns, and tools for adding tool use.

3. launch

import dataclasses

from benchmax.bundle import dump_bundle
from castform.platform import ensure_session
from castform.platform.environment_assets import upload_assets
from castform.platform.client import TrainerClient

ensure_session()
bundle = dump_bundle(MyEnv)
uploaded = upload_assets(
    bundle=bundle, train_dataset=train_data,
    eval_dataset=eval_data, run_name="my-first-run",
)
trainer = TrainerClient()
run_id = trainer.launch_training_run(**dataclasses.asdict(uploaded))

view your run at https://app.castform.com/experiments/{run_id}. see launching for full api reference.

complete runnable script
import dataclasses
import re
from pathlib import Path

from benchmax.bundle import dump_bundle
from benchmax.envs import (
    BaseEnv, BaseRollout, DatasetSplit, Example, JsonlDataset,
    canonical_example_id,
)
from benchmax.rewards import extract_completion_text
from castform.platform import ensure_session
from castform.platform.environment_assets import upload_assets
from castform.platform.client import TrainerClient

ensure_session()

# --- dataset ---
dataset = [
    {"prompt": "what is 2+2?", "ground_truth": "4"},
    {"prompt": "capital of France?", "ground_truth": "Paris"},
    # ... add 200+ examples for real training
]
split = int(len(dataset) * 0.8)
train_data, eval_data = dataset[:split], dataset[split:]

# --- environment ---
class MyEnv(BaseEnv):
    max_turns = 1

    async def create_dataset(
        self, split: DatasetSplit, base_dir: Path, *, max_examples: int | None = None
    ):
        return JsonlDataset(
            base_dir / f"{split}.jsonl",
            row_to_example=self._example,
            max_examples=max_examples,
        )

    def _example(self, row: dict) -> Example:
        payload = {
            "prompt_messages": [
                {
                    "role": "system",
                    "content": "answer the question concisely. put your answer in <answer> tags.",
                },
                {"role": "user", "content": row["prompt"]},
            ],
            "ground_truth": row["ground_truth"],
        }
        return Example(id=canonical_example_id(payload), payload=payload)

    async def compute_reward(self, rollout: BaseRollout) -> dict[str, float]:
        text = extract_completion_text(rollout.messages)
        match = re.search(r"<answer>(.*?)</answer>", text, re.DOTALL)
        answer = match.group(1).strip().lower() if match else ""
        expected = str(rollout.example_args.get("ground_truth", "")).strip().lower()
        return {"correct": 1.0 if answer == expected else 0.0}

# --- bundle, upload and launch ---
bundle = dump_bundle(MyEnv)
uploaded = upload_assets(
    bundle=bundle, train_dataset=train_data,
    eval_dataset=eval_data, run_name="my-first-run",
)
trainer = TrainerClient()
run_id = trainer.launch_training_run(**dataclasses.asdict(uploaded))
print(f"https://app.castform.com/experiments/{run_id}")

what happens next

once your run launches:

  1. gpus warm up (a few minutes). status shows “pending”.
  2. metrics start flowing. reward curves & model responses appear on the train tab.
  3. inspect completions. expand rollouts to see what the model is generating at each step.

don’t draw conclusions from the first dozen steps. rewards will fluctuate early as the model explores. see monitoring a run for how to interpret metrics and spot problems.

once training converges, evaluate your model against baselines and test it in the playground.

learn more