Skip to main content
Shared data types passed to and returned from the ServiceClient, TrainingClient, and SamplingClient.

Training inputs

A single training example: a ModelInput paired with a dict of TensorData loss function inputs. with_provenance(model_input, loss_fn_inputs) returns a copy that records where each input token and each target_tokens entry came from, as lists of provenance spans, for router replay. Datum rejects fields it doesn’t define.
A run of tokens from a recorded sample’s prompt (PromptProvenanceSpan) or its completion (SampledProvenanceSpan). sequence_id names the sample, offset is the position of the first token in that part and defaults to 0, and length is the number of tokens, at least 1. Pass lists of spans to Datum.with_provenance().
A tokenized prompt, represented as a list of ModelInputChunk objects. Construct with ModelInput.from_ints(token_ids) for the common case.
A discriminated union of EncodedTextChunk (a list of token IDs), ImageChunk (a base64-encoded image with an expected token count), and ImageAssetPointerChunk.
A pointer to an out-of-band image asset, present for Tinker import parity only. Baseten samplers don’t resolve out-of-band assets, so constructing one raises ValueError. Pass images inline as ImageChunk instead. See Image inputs must be inline.
A serializable tensor with a flat data list, a dtype string, and a shape. Convert to and from torch.Tensor with TensorData.to_torch() and TensorData.from_torch(tensor).

Configuration

Controls for text generation: temperature, top_p, top_k, max_tokens, seed, stop, and structured_outputs.structured_outputs is a dict that constrains generation, for example to a JSON schema or a regular expression. The SDK forwards it unchanged as the vLLM structured_outputs request field, so use vLLM’s keys, such as {"json": schema} or {"regex": pattern}. It defaults to None, which leaves generation unconstrained.
Optimizer hyperparameters: learning_rate, beta1, beta2, eps, weight_decay, and grad_clip_norm.grad_clip_norm defaults to 1.0. Set it to 0.0 to disable gradient clipping.
A StrEnum that sets trainer and sampler capacity. AvailabilityModel.dedicated uses capacity that Baseten doesn’t preempt. AvailabilityModel.spot uses idle capacity that Baseten can reclaim. Pass it to ServiceClient(availability_model=...) or set LOOPS_AVAILABILITY_MODEL. See Spot capacity.
Optional Weights & Biases settings (project and an optional run name) passed to create_lora_training_client to stream training metrics.
String literal aliases. LossFnType names the built-in loss functions: "cross_entropy", "importance_sampling", "ppo", "cispo", "dro", "dpo", "dppo". See Loss functions for each loss’s inputs and loss_fn_config defaults. StopReason is "stop" or "length", the value of SampledSequence.stop_reason.

Results and handles

The full response from sample(): a list of SampledSequence objects in sequences, the policy_version the sampler replica was running, and prompt_logprobs / topk_prompt_logprobs populated when the matching sample() flags are set.
A single generated sequence: a list of output token IDs, optional per-token log-probabilities, a stop reason, and optional decoded text. sequence_id refers to the sample’s record when the sampling client was created with record_stability_info=True, and is None otherwise.
Metadata for a saved checkpoint, populated by list_checkpoints(). id addresses the checkpoint in the API; checkpoint_id is the name it was saved under. base_model identifies the model, and lora_adapter_config contains its saved LoRA configuration.
The LoRA configuration stored in a Checkpoint. rank is the adapter rank, read from the r field in the checkpoint payload.
Which side saved a named checkpoint: "sampler" or "trainer". Pass it to get_checkpoint_archive_url() or download_checkpoint() to pick one when a name was saved for both.
A paginated list of presigned file URLs for an S3-backed checkpoint, populated by get_checkpoint_archive_url().
One entry in a CheckpointFilesResponse.presigned_urls list: a presigned URL plus relative_file_name, node_rank, size_bytes, and last_modified metadata.
The result of download_checkpoint(): dest_dir is the pathlib.Path where the checkpoint landed, file_count is the number of files downloaded, and bytes_written is their total size.
Returned by ServiceClient.get_server_capabilities(); describe which base models the control plane can provision and on which GPU classes.
A handle to a long-running training operation. Call .result() or .result(timeout=seconds) to block until the operation completes and return the result. The forward and forward_backward methods return a ForwardBackwardFuture with the same .result() contract.
The result returned by forward() and forward_backward(). It contains the aggregate loss, a metrics mapping, and one loss_fn_outputs entry per input datum.Each entry in loss_fn_outputs contains a logprobs TensorData value for the target tokens. When the request sets top_k_logprobs=K, each entry also contains topk_token_ids and topk_logprobs values with shape [sequence_length, K].
The result returned by optim_step(). Its metrics field, a dict[str, float] or None, describes the update. With a WandbConfig, calling .result() on the future that optim_step() returns also logs metrics to Weights & Biases.
Response payloads returned by the matching TrainingClient and SamplingClient methods.