whileai.simulations.training

Training runs, trainer callbacks, TRL export. 14 public names. import whileai.simulations as wai, then wai.name.

Name	What it does
`TrainerCallback`	One line on a Transformers or TRL trainer: `trainer.add_callback(wai.TrainerCallback(run))`.
`TrainingRun`	One fine-tune, as the platform sees it.
`attach_delta`	Compute `delta_report` for a finished run and put it on the run page: the summary is re-sent with `delta` added, status unchanged.
`attach_holdout`	Did it work? Put the held-out pass rate before and after on a run that has already finished — the two numbers its page opens with::
`delete_model`	Stop hosting `name`: removes the model row from the account, so `models()` no longer lists it and its endpoint stops answering for that name.
`delete_run`
`get_run`	The run plus `series`: its points, oldest first.
`list_runs`
`models`	The account’s hosted models: `name`, `baseModel`, `adapter`, `adapterRunId`, `version`, `endpoint` (an OpenAI-compatible base URL; send the account key as the bearer and `name` as the model).
`reward_model`	A judge backed by a finished reward-model run.
`serve`	Host a finished run’s adapter under `name`.
`train`	Start a hosted fine-tune on a pushed dataset and return the run.
`training_run`	Create a run on the platform and return the handle to log into.
`unserve`	Stop hosting `name`: removes the model row from the account, so `models()` no longer lists it and its endpoint stops answering for that name.

TrainerCallback

TrainerCallback(run: TrainingRun, finish: bool = True)

One line on a Transformers or TRL trainer: trainer.add_callback(wai.TrainerCallback(run)). Logs every on_log to the run (loss, lr, eval loss, epoch, grad norm, token accuracy), takes the step count from the trainer at on_train_begin, and finishes the run at on_train_end. If the trainer raises, finish the run yourself with run.fail(...) or use the run as a context manager around trainer.train().

TrainingRun

TrainingRun(
    run_id: str,
    name: str,
    api_key: str | None = None,
    total_steps: int | None = None,
    flush_every: int = 25,
    flush_seconds: float = 15.0,
    max_batch: int = 500,
    transport: Callable[..., Any] | None = None,
)

One fine-tune, as the platform sees it. Create with training_run. log buffers; flush sends. A send that fails is retried on the next flush and counted in errors; the training loop is never interrupted by the dashboard. finish flushes first.

TrainingRun.delta

delta(
    self,
    before: Sequence[dict],
    after: Sequence[dict],
    target: str | None = 'pass_at_1',
    must_not_regress: Sequence[str] = (),
    by: str | Callable[[dict], Any] | None = None,
    proxy: str | None = None,
) -> dict[str, Any]

Did the training move the behavior? delta_report over the rollouts before and after, kept on the run and sent with finish under summary["delta"] (sent right away when the run is already finished). The run page draws it, including the per-group table when by names a row key or marker. proxy names the training reward’s marker so the report can call the run over-optimized when the proxy moved and the target did not.

TrainingRun.finish

finish(
    self,
    status: str = 'done',
    summary: Mapping[str, Any] | None = None,
    adapter: str | None = None,
    error: str | None = None,
) -> dict[str, Any]

Flush, then mark the run done, failed, or stopped.

TrainingRun.flush

flush(self) -> bool

Send buffered points. Returns True when nothing is left unsent.

TrainingRun.holdout

holdout(
    self,
    before: float,
    after: float,
    metric: str = 'pass',
) -> dict[str, float]

Did it work? The held-out pass rate before and after, which is what the run’s page opens with. metric="loss" for held-out loss (SFT), where lower is better. Pass rates are 0 to 1.

TrainingRun.log

log(self, step: int, **metrics: float) -> None

Record one point. Any finite numeric keyword is a metric (loss, eval_loss, lr, epoch, grad_norm, …).

TrainingRun.note

note(self, **fields: Any) -> None

Put fields on the run’s summary ahead of finish: whichever callback finishes the run, the summary carries them. Sent right away when the run is already finished.

TrainingRun.progress

progress(self, step: int, total_steps: int | None = None) -> None

Advance the bar without a metric. total_steps (re)sets the denominator; a trainer that learns its length late can call this.

TrainingRun.refresh

refresh(self) -> str

Read a hosted run’s state from the platform: running, done or failed. Fills adapter, training (before, after, seconds, rows) and error once it has ended.

TrainingRun.wait

wait(self, timeout: float | None = None, poll: float = 15.0) -> str

Block until a hosted run ends, reading its state every poll seconds (never under TRAINING_POLL_MIN_S). Returns the final status; raises TimeoutError when timeout seconds pass first.

attach_delta

attach_delta(
    run_id: str,
    before: Sequence[dict],
    after: Sequence[dict],
    target: str | None = 'pass_at_1',
    must_not_regress: Sequence[str] = (),
    by: str | Callable[[dict], Any] | None = None,
    proxy: str | None = None,
    api_key: str | None = None,
) -> dict[str, Any]

Compute delta_report for a finished run and put it on the run page: the summary is re-sent with delta added, status unchanged.

attach_holdout

attach_holdout(
    run_id: str,
    before: float,
    after: float,
    metric: str = 'pass',
    api_key: str | None = None,
) -> dict[str, float]

Did it work? Put the held-out pass rate before and after on a run that has already finished — the two numbers its page opens with:: wai.attach_holdout(“run_…”, before=0.42, after=0.58) Pass rates are 0 to 1. metric="loss" sends held-out loss instead (SFT), where lower is better. The summary is re-sent with the two keys added and the status unchanged; sending again is a correction.

delete_model

delete_model(
    name: str,
    api_key: str | None = None,
    transport: Callable[..., Any] | None = None,
) -> dict[str, Any]

Stop hosting name: removes the model row from the account, so models() no longer lists it and its endpoint stops answering for that name. The inverse of serve, the way delete_dataset is the inverse of push. The adapter weights and the training run stay; serve the run again to bring it back (at version 1). Returns {"name": ..., "deleted": True}.

delete_run

delete_run(run_id: str, api_key: str | None = None) -> dict[str, Any]

get_run

get_run(run_id: str, api_key: str | None = None) -> dict[str, Any]

The run plus series: its points, oldest first.

list_runs

list_runs(api_key: str | None = None) -> list[dict[str, Any]]

models

models(api_key: str | None = None) -> list[dict[str, Any]]

The account’s hosted models: name, baseModel, adapter, adapterRunId, version, endpoint (an OpenAI-compatible base URL; send the account key as the bearer and name as the model). A row here is a registry entry, not a running GPU: the endpoint behind it idles to zero on its own and an unused model costs nothing. The row stays until unserve(name) removes it; serving the same name again bumps its version rather than adding a row.

reward_model

reward_model(
    run: TrainingRun | str,
    threshold: float | None = None,
    api_key: str | None = None,
    transport: Callable[..., Any] | None = None,
    batch: int = 256,
) -> RewardModel

A judge backed by a finished reward-model run. run = wai.train("ds_...", method="rm", wait=True) trains a sequence-classification head on the set’s pass-vs-fail pairs and picks the score threshold that best separates the held-out pairs. judge = wai.reward_model(run) then scores any rollout row: data.grade(judge=judge), wai.evaluate(rollouts, judge), wai.judge_trust(scored.rows, judge=judge). Pass threshold= to override the run’s own cut. The scores are the model’s; a reward model trained on one agent’s pairs says nothing about another agent.

serve

serve(
    name: str,
    run: TrainingRun | str | None = None,
    base_model: str | None = None,
    api_key: str | None = None,
    transport: Callable[..., Any] | None = None,
) -> dict[str, Any]

Host a finished run’s adapter under name. Returns the model row; endpoint is the OpenAI-compatible base URL and name the model id to send. Posting an existing name bumps version. run is a TrainingRun, the record get_run returns, or the run id; the adapter and base model come from the run record unless base_model is given. No run serves the bare base (base_model required). unserve is the inverse.

train

train(
    dataset: str,
    method: str = 'sft',
    steps: int | None = None,
    epochs: float | None = None,
    holdout: str | None = None,
    base_model: str | None = None,
    generations: int | None = None,
    learning_rate: float | None = None,
    beta: float | None = None,
    seed: int | None = None,
    max_completion_length: int | None = None,
    loss_type: str | None = None,
    temperature: float | None = None,
    truncated: str | None = None,
    config: Mapping[str, Any] | None = None,
    wait: bool = False,
    timeout: float | None = None,
    poll: float = 15.0,
    api_key: str | None = None,
    transport: Callable[..., Any] | None = None,
) -> TrainingRun

Start a hosted fine-tune on a pushed dataset and return the run. method is "sft" (LoRA on the passing rows), "grpo" (the reference-first-action reward over the graded rows), "dpo" (a pass against a fail per prompt, length matched) or "rm" (a reward model on those same pairs; reward_model(run) is then a judge). steps sets the optimizer steps for GRPO, DPO and RM, epochs the SFT epochs; each method has a default. holdout names the eval set; it defaults to the train set’s split sibling from datasets.cut. base_model overrides the trainer’s base; only SERVED_BASES can be served afterwards, and train warns when the run will not be. The run is the same record training_run makes, so run.url is the loss curve, run.delta and get_run work unchanged, and the trainer finishes it. run.refresh() reads where it is; run.wait() (or wait=True) blocks until done or failed, after which run.adapter names the weights and run.training carries before, after, rows and seconds. serve puts the adapter on an endpoint. The knobs a run is reproduced and compared by (rlhf-book ch. 6, 7): generations is the group size per prompt for GRPO (the k the advantage is taken over; a pushed set’s repeats is the natural value), beta the KL coefficient for GRPO and DPO, learning_rate the optimizer step for every method, seed the sampling and data order seed, max_completion_length the token cap on a sampled reply (GRPO, DPO), loss_type the objective variant (GRPO: bnpo, grpo, dr_grpo; DPO: any TRL loss). temperature is the sampling temperature the trainer rolls out at (GRPO); the dataset’s rows say what they were measured at under sampling.temperature, and train says so when the two differ, since a before/after comparison across temperatures is not like for like. truncated says what GRPO does with a sampled reply the token cap cut: "mask" (the default) gives it no gradient, "zero" scores it 0 the old way. A cut reply scored 0 teaches shorter thinking before it teaches the task, so "zero" is the knob to reach for only when the cap itself is the behavior under training (#253). Each has a trainer default when left None; the range each is accepted in and the value the cited paper used are in TRAINING_KNOBS (defaults.py: DAPO, Dr. GRPO, ProRL, DPO and the rlhf-book chapters), and a rejected value is told the reference. config passes further host keys as given (epsilonHigh, scaleRewards, balance). Every knob lands on the run’s config so the run page shows it. A dataset already training answers with that run instead of a second.

training_run

training_run(
    name: str,
    dataset: str | None = None,
    after_dataset: str | None = None,
    base_model: str | None = None,
    trainer: str | None = None,
    total_steps: int | None = None,
    config: Mapping[str, Any] | None = None,
    api_key: str | None = None,
    flush_every: int = 25,
    flush_seconds: float = 15.0,
    max_batch: int = 500,
    transport: Callable[..., Any] | None = None,
) -> TrainingRun

Create a run on the platform and return the handle to log into. dataset is the ds_... id trained on; after_dataset the set of post-training rollouts, when you have one, so the run page can show the before/after. config is anything JSON-shaped you want on the run page (hyperparameters, the command). api_key defaults to the usual credential chain.

unserve

unserve(
    name: str,
    api_key: str | None = None,
    transport: Callable[..., Any] | None = None,
) -> dict[str, Any]

Get started

Guides

Concepts

Reference

API

whileai.simulations.training

TrainerCallback

TrainingRun

TrainingRun.delta

TrainingRun.finish

TrainingRun.flush

TrainingRun.holdout

TrainingRun.log

TrainingRun.note

TrainingRun.progress

TrainingRun.refresh

TrainingRun.wait

attach_delta

attach_holdout

delete_model

delete_run

get_run

list_runs

models

reward_model

serve

train

training_run

unserve

​TrainerCallback

​TrainingRun

​TrainingRun.delta

​TrainingRun.finish

​TrainingRun.flush

​TrainingRun.holdout

​TrainingRun.log

​TrainingRun.note

​TrainingRun.progress

​TrainingRun.refresh

​TrainingRun.wait

​attach_delta

​attach_holdout

​delete_model

​delete_run

​get_run

​list_runs

​models

​reward_model

​serve

​train

​training_run

​unserve

TrainerCallback

TrainingRun

TrainingRun.delta

TrainingRun.finish

TrainingRun.flush

TrainingRun.holdout

TrainingRun.log

TrainingRun.note

TrainingRun.progress

TrainingRun.refresh

TrainingRun.wait

attach_delta

attach_holdout

delete_model

delete_run

get_run

list_runs

models

reward_model

serve

train

training_run

unserve