Skip to content

Custom trainers

Write a custom trainer when the place that trains your model isn't already covered by trl() or vertex(): an internal GPU platform, another provider's fine-tuning API, or a bespoke remote runner. Don't write one to customize datasets or evals; those stay the same whichever trainer runs, in Datasets and Evals.

A trainer is a small adapter. It's a plain object that needs nothing but the types from june/train, and the whole contract is:

  • a stable id, which namespaces its options, metadata and job refs;
  • declared capabilities: the methods it runs and the artifact types it returns;
  • start(request), which starts the job and returns a serializable job ref;
  • status(job) and cancel(job), which work from that ref alone, in any process;
  • on success, a typed artifact: an adapter, checkpoint or provider-model.

Here it is for a GPU cluster with its own API:

import type { Trainer } from "june/train";

declare const cluster: {
  submit(spec: object): Promise<string>;
  get(id: string): Promise<{
    phase: "pending" | "running" | "done" | "error" | "stopped";
    step?: number;
    steps?: number;
    loss?: number;
    weights?: string;
    message?: string;
  }>;
  stop(id: string): Promise<void>;
};

export const clusterTrainer: Trainer = {
  id: "cluster",
  capabilities: { methods: ["lora"], artifacts: ["adapter"] },

  async start(request) {
    const id = await cluster.submit({
      base: request.model,
      examples: await request.train.examples(),
      rank: request.method.rank ?? 16,
      epochs: request.config.epochs ?? 1,
      ...request.providerOptions.cluster,
    });
    return {
      id,
      trainer: "cluster",
      metadata: { cluster: { base: request.model } },
    };
  },

  async status(job) {
    const remote = await cluster.get(job.id);
    const progress = {
      step: remote.step,
      totalSteps: remote.steps,
      ...(remote.loss !== undefined && { metrics: { loss: remote.loss } }),
    };

    switch (remote.phase) {
      case "pending":
        return { state: "queued" };
      case "running":
        return { state: "running", progress };
      case "stopped":
        return { state: "canceled" };
      case "error":
        return {
          state: "failed",
          error: { name: "ClusterError", message: remote.message ?? "failed" },
        };
      case "done":
        return {
          state: "succeeded",
          result: {
            artifact: {
              type: "adapter",
              format: "peft",
              uri: remote.weights!,
              baseModel: String(job.metadata!.cluster.base),
            },
            ...(remote.loss !== undefined && {
              metrics: { loss: remote.loss },
            }),
          },
        };
    }
  },

  async cancel(job) {
    await cluster.stop(job.id);
  },
};

Capabilities

capabilities.methods lists the training methods the trainer runs, and capabilities.artifacts the artifact types it can return. fineTune() rejects a method that isn't listed before calling start, and fails a job whose artifact type isn't listed, so callers can rely on them. Methods are sft, lora, qlora, dpo, preference and reinforcement; a method config's rank, alpha, dropout, targetModules and beta are hints a trainer uses or ignores.

Check anything else a trainer can't do (an example type, a message with images) in start, and throw before doing any work. Errors should say what to change.

Requests

start(request) gets:

Field
modelthe base model, as the caller wrote it
traina Dataset; call examples() and convert them to your format
validationa Dataset, if the caller split one off
methodalways a config object: { type: "lora", rank: 32 }
configepochs, learningRate, batchSize, maxSequenceLength, seed; all optional
providerOptionseverything callers passed; read only providerOptions[yourId]
namethe run's name, if given

New optional fields may be added to the request, so ignore fields you don't know rather than rejecting them. The Trainer contract is one of June's stable contracts.

Jobs

start returns a TrainingJob: { id, trainer, metadata? }. It plays the part a job ref plays for any job: startFineTune() returns a Job whose backend is your trainer. It's the only thing June keeps, and it's saved as JSON, so put whatever you need to find the job again from another process in metadata, under your trainer's id: a region, a directory, a resource name. Never put credentials in it; it ends up in manifests and logs.

status(job) must work from that ref alone, in a process that didn't call start, so look the job up in your backend rather than in memory. It returns a TrainingStatus:

  • state: queued, running, succeeded, failed or canceled. Map your backend's states onto these.
  • progress: fraction, step, totalSteps, message and numeric metrics, whatever you know. It shows up on the trace.
  • result: once succeeded, the artifact and final metrics.
  • error: once failed, { name, message, details }. Put useful context, such as a log tail, in details.

cancel(job) asks the backend to stop, and can return before the job has stopped. When fineTune() or <FineTune> cancels because its signal aborted, June sends the cancel and rejects with the abort reason without waiting. A caller that runs job.cancel() and then job.result() keeps polling until status reports canceled, so it's fine for the job to report running for a while.

June polls status, so it should be cheap. A trainer doesn't need its own polling, retries or event emitters; add retries only for errors your backend documents as transient.

Artifacts

Normalize what your backend produces into one of June's artifact types:

  • adapter with a format (e.g. "peft") and a uri, for weights applied to the base model;
  • checkpoint with a format and a uri, for complete weights;
  • provider-model with your trainer's id as provider and the id to call it by as modelId.

Every artifact carries baseModel. Anything else (resource names, endpoints, build ids) goes in metadata[yourId], which June stores and never reads.

Testing a trainer

Run fineTune() against it, the way users will. For an HTTP backend, point it at a local fixture server (Bun.serve) and exercise real fetch, as June's Vertex tests do. For a subprocess, run a real fixture script, as the TRL tests do. Rebuild jobs from JSON.parse(JSON.stringify(ref)) with fineTuneJob() to check they survive a restart.