API reference

Reference for river_client and its rl library, generated from river-client 0.11.0. The guide covers the API as a walkthrough; this page lists public classes, methods, and types.

pip install river-client

import river_client as river

Core client

Use these primitives for sampling, sessions, models, and training operations.

Client

River API client.

Connects to the River API server over gRPC with automatic retry on transient failures and connection keepalive.

Client(
    api_key: str,
    endpoint: str = 'api.river.ai',
    port: int = 443,
    timeout: float = 86400.0,
    use_ssl: bool = True,
    enable_retries: bool = True,
    image_upload_concurrency: int = 4,
)

Methods: aclose, session, sample, health_check, get_capabilities, get_server_capabilities, create_deployment, get_deployment, list_deployments, scale_on_target, delete_deployment, get_deployment_usage, wait_for_deployment, chat_complete, chat_complete_from_checkpoint, chat_complete_from_training, close

Client.aclose

async Client.aclose() -> None

Drain started image uploads and close connections off the event loop.

Client.session

Client.session(
    *,
    timeout: float = 86400.0,
    before_poll: Callable[[], None] | None = None,
    on_session_creation_attempted: Callable[[], None] | None = None,
    on_session_created: Callable[[], None] | None = None,
    on_session_closed: Callable[[], None] | None = None,
    **tags: str,
) -> SessionContext

Create a session context manager.

Returns: SessionContext — Context manager that yields a Session

Client.sample

Client.sample(
    prompts: str | list[str] | None = None,
    *,
    base_model: str,
    num_samples: int = 1,
    max_tokens: int = 256,
    temperature: float = 1.0,
    top_p: float = 1.0,
    top_k: int = -1,
    stop: list[str] | None = None,
    seed: int | None = None,
    return_prompt_logprobs: bool = False,
    logprobs: int | None = None,
    images: list[Image] | list[list[Image]] | None = None,
    prompt_token_ids: list[int] | list[list[int]] | None = None,
    model_input: list[dict] | list[list[dict]] | None = None,
    tokenizer: Any | None = None,
    metrics_type: str = '',
    timeout: float | None = None,
) -> list[Sample]

Sample from a base model (no session required).

Returns: list[Sample] — Flat list[Sample] — all samples across all prompts. For a single prompt with num_samples=1 (the default), this is a list with one element: result[0].text.

Client.health_check

Client.health_check() -> bool

Check API health.

Returns: bool — True if healthy

Client.get_capabilities

Client.get_capabilities() -> list[str]

Get supported model names; preserves the original list-returning API.

Client.get_server_capabilities

Client.get_server_capabilities(*, model_name: str | None = None) -> ServerCapabilities

Read public feature contracts, optionally for one authorized model.

This checks protocol support, not current model capacity or queue state. Results are fetched fresh so a new run detects an endpoint change.

Client.create_deployment

Client.create_deployment(
    checkpoint: str | Checkpoint | None = None,
    *,
    base_model: str | None = None,
    unified_replicas: int | None = None,
    prefill_replicas: int | None = None,
    decode_replicas: int | None = None,
    idempotency_key: str | None = None,
    wait: bool = False,
    wait_timeout: float = 1800.0,
    timeout: float | None = None,
) -> Deployment

Create a dedicated deployment of a checkpoint or base model.

Supply exactly one of checkpoint or base_model. Base-model creation requires a server advertising base_model_creation.

This is a gated feature, disabled by default. Contact River to enable dedicated deployments for your team and the checkpoint's base model before calling this method. Use a team API key with that access; personal API keys cannot create deployments. Installing or upgrading the client does not enable access. The server rejects creation when the feature is disabled or the required team/model access is absent.

The checkpoint path may be republished with new weights. New or restarted workers load the available weights; serving workers keep their loaded adapter. Republishing does not automatically reload the fleet or provide an atomic update, including across prefill/decode workers. For a consistent update at the same URL, scale all roles to zero, wait for workers to drain and stop, finish saving, then restore the targets. Use a fresh checkpoint name and a new deployment when the old deployment must keep serving during the update.

Supply exactly one capacity group: unified_replicas for unified serving, or both prefill_replicas and decode_replicas for prefill/decode disaggregated serving. The topology is inferred from the group and is immutable afterwards. The backend selects the hardware specification, placement and ports from the checkpoint's base model.

Provisioning is asynchronous: the returned deployment is usually phase="accepted". Use the standard OpenAI client with base_url=deployment.base_url once it is ready; the URL binds the checkpoint, so the existing model argument can stay unchanged. Pass wait=True to block until it serves (or fails). An idempotency_key makes a retried call return the same accepted deployment; each explicit call without one creates a new deployment.

Client.get_deployment

Client.get_deployment(deployment_id: str, *, timeout: float | None = None) -> Deployment

Fetch a deployment by id, including its tombstone after deletion.

Dedicated deployments are gated and disabled by default. See create_deployment for access requirements.

Client.list_deployments

Client.list_deployments(
    *,
    include_deleted: bool = False,
    timeout: float | None = None,
) -> list[Deployment]

List the caller's deployments with their per-role replica counts.

Dedicated deployments are gated and disabled by default. See create_deployment for access requirements.

Client.scale_on_target

Client.scale_on_target(
    deployment_id: str,
    *,
    unified_replicas: int | None = None,
    prefill_replicas: int | None = None,
    decode_replicas: int | None = None,
    idempotency_key: str | None = None,
    timeout: float | None = None,
) -> Deployment

Set the absolute desired count of each supplied role.

Dedicated deployments are gated and disabled by default. Increasing capacity, including resuming from zero, requires active access for your team and this deployment's base model; contact River to request it. Reducing capacity remains available after that access is revoked.

Omitted roles keep their current target. A unified deployment accepts only unified_replicas; a prefill/decode deployment accepts prefill_replicas and/or decode_replicas, and a target with one role at zero and the other positive is rejected — scale both to zero together to stop, and both up to resume.

Client.delete_deployment

Client.delete_deployment(
    deployment_id: str,
    *,
    idempotency_key: str | None = None,
    wait: bool = False,
    wait_timeout: float = 600.0,
    timeout: float | None = None,
) -> Deployment

Delete a deployment in any live state; repeating is a no-op success.

Dedicated deployments are gated and disabled by default. Deleting an existing deployment remains available after your team's access to create or increase capacity for its model is revoked.

Client.get_deployment_usage

Client.get_deployment_usage(
    deployment_id: str,
    *,
    timeout: float | None = None,
) -> dict[str, Any]

Requested replica/GPU-hours per role and accepted-operation history.

Dedicated deployments are gated and disabled by default. See create_deployment for access requirements.

Client.wait_for_deployment

Client.wait_for_deployment(
    deployment_id: str,
    *,
    timeout: float = 1800.0,
    poll_interval: float = 5.0,
    until: tuple[str, ...] = ('ready', 'degraded'),
) -> Deployment

Wait for a requested phase and its serving capacity (or failure).

Dedicated deployments are gated and disabled by default. See create_deployment for access requirements.

The default accepts partially available capacity. Pass until=("ready",) to wait for every requested replica to be ready.

Client.chat_complete

Client.chat_complete(
    messages: list[dict],
    *,
    base_model: str,
    timeout: float | None = None,
    **kwargs,
) -> ChatCompleteResult

Chat completion from a base model (no LoRA).

Builds an OpenAI-format request body and sends it through the gRPC ChatCompleteFromBase RPC.

Returns: ChatCompleteResult — ChatCompleteResult with response_json and status_code.

Client.chat_complete_from_checkpoint

Client.chat_complete_from_checkpoint(
    messages: list[dict],
    *,
    checkpoint_path: str,
    base_model: str = '',
    timeout: float | None = None,
    **kwargs,
) -> ChatCompleteResult

Chat completion from a saved checkpoint.

Returns: ChatCompleteResult — ChatCompleteResult with response_json and status_code.

Client.chat_complete_from_training

Client.chat_complete_from_training(
    messages: list[dict],
    *,
    model_id: str,
    timeout: float | None = None,
    **kwargs,
) -> ChatCompleteResult

Chat completion from in-memory training weights.

Returns: ChatCompleteResult — ChatCompleteResult with response_json and status_code.

Client.close

Client.close() -> None

Close connections after any started image uploads finish.

Session

A training session with GPU allocation.

Entered through Client.session; it owns the models you train.

Methods: upload_image, upload_image_async, release_image, release_image_async, restore_images, get_server_capabilities, attest_training_data, create_model, sample, submit_sampling_batch, submit_sample

Session.upload_image

Session.upload_image(
    data: bytes,
    *,
    idempotency_key: str | None = None,
    timeout: float | None = None,
) -> ImageHandle

Upload into this session; bytes remain recoverable in the local image cache.

Session.upload_image_async

async Session.upload_image_async(
    data: bytes,
    *,
    idempotency_key: str | None = None,
    timeout: float | None = None,
) -> ImageHandle

Bounded non-blocking upload, automatically released when this session ends.

Session.release_image

Session.release_image(
    image: ImageHandle | None = None,
    *,
    idempotency_key: str | None = None,
) -> None

Session.release_image_async

async Session.release_image_async(
    image: ImageHandle | None = None,
    *,
    idempotency_key: str | None = None,
) -> None

Session.restore_images

async Session.restore_images(value: Any, *, image_store: ImageStore) -> Any

Re-upload checkpoint references once per content hash and remap every occurrence.

Session.get_server_capabilities

Session.get_server_capabilities(*, model_name: str | None = None) -> ServerCapabilities

Read protocol support, optionally for one authorized model.

Session.attest_training_data

Session.attest_training_data(
    artifacts: list[TrainingDataArtifact],
    timeout: float = 86400.0,
) -> TrainingDataAttestation

Ask the API to hash source artifacts and retain their manifest.

The source bytes are discarded by the API after hashing. Pass the result to create_model to make the server fail closed before every forward/backward request if that manifest disappears or no longer belongs to the model's session.

Session.create_model

Session.create_model(
    base_model: str,
    lora: LoraConfig,
    tokenizer: str | Any | None = None,
    checkpoint: str | Checkpoint | None = None,
    timeout: float = 86400.0,
    training_data_attestation: TrainingDataAttestation | str | None = None,
) -> Model

Create a new model for training.

Returns: Model — Model object for training

Session.sample

Session.sample(
    prompts: str | list[str] | None = None,
    *,
    base_model: str,
    checkpoint: str | Checkpoint | None = None,
    num_samples: int = 1,
    max_tokens: int = 256,
    temperature: float = 1.0,
    top_p: float = 1.0,
    top_k: int = -1,
    stop: list[str] | None = None,
    seed: int | None = None,
    seeds: list[int] | None = None,
    return_prompt_logprobs: bool = False,
    logprobs: int | None = None,
    return_expert_routing: bool = False,
    images: list[Image] | list[list[Image]] | None = None,
    prompt_token_ids: list[int] | list[list[int]] | None = None,
    model_input: list[dict] | list[list[dict]] | None = None,
    tokenizer: Any | None = None,
    metrics_type: str = '',
    timeout: float = 86400.0,
) -> list[list[Sample]]

Sample from a base model or checkpoint.

When checkpoint is provided, the server loads the LoRA from the saved checkpoint, generates text, then unloads.

Returns: list[list[Sample]] — list[list[Sample]] — outer list is per-prompt, inner list is per-sample.

Session.submit_sampling_batch

Session.submit_sampling_batch(
    prompts: str | list[str] | None = None,
    *,
    base_model: str,
    checkpoint: str | Checkpoint | None = None,
    num_samples: int = 1,
    max_tokens: int = 256,
    temperature: float = 1.0,
    top_p: float = 1.0,
    top_k: int = -1,
    stop: list[str] | None = None,
    seed: int | None = None,
    seeds: list[int] | None = None,
    return_prompt_logprobs: bool = False,
    return_prompt_token_ids: bool = False,
    logprobs: int | None = None,
    return_expert_routing: bool = False,
    images: list[Image] | list[list[Image]] | None = None,
    prompt_token_ids: list[int] | list[list[int]] | None = None,
    model_input: list[dict] | list[list[dict]] | None = None,
    tokenizer: Any | None = None,
    idempotency_key: str | None = None,
    metrics_type: str = '',
    timeout: float = 86400.0,
) -> PendingSamplingBatch

Submit up to 128 independent samples in one RPC.

Uses the new sampling transport; requires server support. Consume pending.as_completed() to receive individual successes/failures. Other arguments match sample. Existing submit_sample retains its whole-batch future behavior.

Session.submit_sample

Session.submit_sample(
    prompts: str | list[str] | None = None,
    *,
    base_model: str,
    checkpoint: str | Checkpoint | None = None,
    num_samples: int = 1,
    max_tokens: int = 256,
    temperature: float = 1.0,
    top_p: float = 1.0,
    top_k: int = -1,
    stop: list[str] | None = None,
    seed: int | None = None,
    seeds: list[int] | None = None,
    return_prompt_logprobs: bool = False,
    logprobs: int | None = None,
    return_expert_routing: bool = False,
    images: list[Image] | list[list[Image]] | None = None,
    prompt_token_ids: list[int] | list[list[int]] | None = None,
    model_input: list[dict] | list[list[dict]] | None = None,
    tokenizer: Any | None = None,
    metrics_type: str = '',
    timeout: float = 86400.0,
) -> PendingSample

Submit sampling from a base model or checkpoint without waiting.

See sample for the full kwarg reference.

Session.session_id

Session.session_id: str

SessionContext

Context manager for Session with auto-heartbeat.

The context manager returned by Client.session.

Model

A training model with mutable in-memory weights.

Created by Session.create_model.

Methods: get_server_capabilities, get_policy_version, forward, forward_backward, optim_step, train_step, submit_forward_backward, submit_optim_step, submit_train_step, sample, submit_sampling_batch, submit_sample, chat_complete, save_weights, load_weights

Model.get_server_capabilities

Model.get_server_capabilities() -> ServerCapabilities

Read the endpoint's public protocol support (used by RL preflight).

Model.get_policy_version

Model.get_policy_version() -> PolicyVersion | None

Read the last server-committed policy; legacy servers return None.

Model.forward

Model.forward(
    data: list[dict],
    loss_fn: str = 'cross_entropy',
    timeout: float = 86400.0,
    expected_policy_id: str | None = None,
    **loss_config: float,
) -> ForwardResult

Forward pass only (compute loss, no gradients).

Returns: ForwardResult — ForwardResult with metrics

Model.forward_backward

Model.forward_backward(
    data: list[dict],
    loss_fn: str = 'cross_entropy',
    timeout: float = 86400.0,
    return_logprobs: bool = False,
    zero_out: bool = True,
    compute_expert_flip_metric: bool = False,
    force_routing_replay: bool = False,
    expected_policy_id: str | None = None,
    **loss_config: float,
) -> ForwardResult

Forward + backward pass (compute gradients).

Returns: ForwardResult — ForwardResult with metrics and logprobs when returned by the worker.

Requests larger than the 1 GiB upload limit are split at datum boundaries into sequential gradient-accumulation sub-batches. The first sub-batch honors zero_out; later sub-batches retain its gradients. A failed sub-batch stops the sequence and this method does not submit an optimizer step. timeout applies to each submitted sub-batch, so the total operation can take up to the number of chunks times timeout. Split requests return only metrics the client can combine exactly across sub-batches; per-sub-batch-only metrics are omitted.

Model.optim_step

Model.optim_step(
    lr: float,
    beta1: float = 0.9,
    beta2: float = 0.999,
    eps: float = 1e-08,
    weight_decay: float = 0.0,
    grad_clip_norm: float | None = None,
    timeout: float = 86400.0,
    expected_policy_id: str | None = None,
    idempotency_key: str | None = None,
    gradient_scale: float | None = None,
) -> OptimStepResult

Apply gradients with Adam optimizer.

Returns: OptimStepResult — OptimStepResult with metrics (step, lr, grad_norm, grad_norm_finite)

Model.train_step

Model.train_step(
    data: list[dict],
    lr: float,
    *,
    loss_fn: str = 'cross_entropy',
    beta1: float = 0.9,
    beta2: float = 0.999,
    eps: float = 1e-08,
    weight_decay: float = 0.0,
    grad_clip_norm: float | None = None,
    compute_expert_flip_metric: bool = False,
    force_routing_replay: bool = False,
    timeout: float = 86400.0,
    **loss_config: float,
) -> tuple[ForwardResult, OptimStepResult]

Complete training step: forward+backward plus optimizer update.

Submits forward+backward and the optimizer step back-to-back — the server runs them in order as one pipelined unit, without a client round trip in between — then waits for both results. See forward_backward and optim_step for the full parameter reference.

The error path differs from calling forward_backward() then optim_step(): the optimizer step is already submitted when forward-backward resolves, so if forward-backward fails, this call raises its error while the optimizer step still runs server-side and Model.step has already advanced. Callers that need to inspect both outcomes should use submit_train_step.

A train step is a complete step: gradients are always cleared first. For micro-batch gradient accumulation, use submit_forward_backward(zero_out=...) and submit_optim_step directly.

For a request larger than the 1 GiB upload limit, the client waits for each accumulated forward/backward sub-batch before submitting the optimizer step. Smaller requests retain the normal pipelined path. The forward/backward result from a split request includes only metrics the client can combine exactly across its sub-batches.

Returns: tuple[ForwardResult, OptimStepResult] — Tuple of (ForwardResult, OptimStepResult).

Model.submit_forward_backward

Model.submit_forward_backward(
    data: list[dict],
    loss_fn: str = 'cross_entropy',
    timeout: float = 86400.0,
    return_logprobs: bool = False,
    zero_out: bool = True,
    compute_expert_flip_metric: bool = False,
    force_routing_replay: bool = False,
    expected_policy_id: str | None = None,
    **loss_config: float,
) -> PendingOp

Submit forward+backward without blocking. Returns a PendingOp.

Call pending.result() later to get the ForwardResult. This enables pipelining: submit step N+1 while step N is still running. See forward_backward for the full kwarg reference.

When submitting multiple micro-batches before submit_optim_step, use zero_out=True for the first one and False for later submissions.

timeout bounds the submit RPC and, separately, the wait inside pending.result().

Requests larger than the 1 GiB upload limit must use the synchronous forward_backward method so the client can keep later sub-batches and any optimizer step behind successful earlier ones.

Model.submit_optim_step

Model.submit_optim_step(
    lr: float,
    beta1: float = 0.9,
    beta2: float = 0.999,
    eps: float = 1e-08,
    weight_decay: float = 0.0,
    grad_clip_norm: float | None = None,
    timeout: float = 86400.0,
    expected_policy_id: str | None = None,
    idempotency_key: str | None = None,
    gradient_scale: float | None = None,
) -> PendingOp

Submit optimizer step without blocking. Returns a PendingOp.

Call pending.result() later to get the OptimStepResult. Model.step advances at submit time, even if the operation later fails.

Model.submit_train_step

Model.submit_train_step(
    data: list[dict],
    lr: float,
    *,
    loss_fn: str = 'cross_entropy',
    beta1: float = 0.9,
    beta2: float = 0.999,
    eps: float = 1e-08,
    weight_decay: float = 0.0,
    grad_clip_norm: float | None = None,
    compute_expert_flip_metric: bool = False,
    force_routing_replay: bool = False,
    timeout: float = 86400.0,
    **loss_config: float,
) -> tuple[PendingOp, PendingOp]

Submit a complete training step without blocking.

Fires forward+backward and the optimizer step back-to-back; the server runs them in submission order per model. Returns the two PendingOps as (forward_backward, optim_step).

Because both are submitted up front, a failed forward-backward does not cancel the already-submitted optimizer step. Model.step advances at submit time. See train_step for the full kwarg reference.

Requests larger than the 1 GiB upload limit must use train_step, which waits for all accumulated sub-batches before submitting the optimizer step.

Model.sample

Model.sample(
    prompts: str | list[str] | None = None,
    *,
    num_samples: int = 1,
    max_tokens: int = 256,
    temperature: float = 1.0,
    top_p: float = 1.0,
    top_k: int = -1,
    stop: list[str] | None = None,
    seed: int | None = None,
    seeds: list[int] | None = None,
    return_prompt_logprobs: bool = False,
    logprobs: int | None = None,
    images: list[Image] | list[list[Image]] | None = None,
    return_expert_routing: bool = False,
    prompt_token_ids: list[int] | list[list[int]] | None = None,
    model_input: list[dict] | list[list[dict]] | None = None,
    metrics_type: str = '',
    timeout: float = 86400.0,
    poll_interval: float = 1.0,
) -> list[list[Sample]]

Sample from the current in-memory training weights.

Generates text using the model's current weights, without needing to save a checkpoint first. Per-token logprobs are always returned.

Returns: list[list[Sample]] — list[list[Sample]] — outer list is per-prompt, inner list is per-sample. Each Sample has .tokens, .text, .logprobs, and .stop_reason.

Model.submit_sampling_batch

Model.submit_sampling_batch(
    prompts: str | list[str] | None = None,
    *,
    num_samples: int = 1,
    max_tokens: int = 256,
    temperature: float = 1.0,
    top_p: float = 1.0,
    top_k: int = -1,
    stop: list[str] | None = None,
    seed: int | None = None,
    seeds: list[int] | None = None,
    return_prompt_logprobs: bool = False,
    return_prompt_token_ids: bool = False,
    logprobs: int | None = None,
    images: list[Image] | list[list[Image]] | None = None,
    return_expert_routing: bool = False,
    prompt_token_ids: list[int] | list[list[int]] | None = None,
    model_input: list[dict] | list[list[dict]] | None = None,
    idempotency_key: str | None = None,
    pinned_policy_id: str | None = None,
    policy_selection: str = 'ordered',
    retained_kv_groups: list[str] | None = None,
    kv_cache_policy: dict | None = None,
    metrics_type: str = '',
    timeout: float = 86400.0,
    poll_interval: float = 0.1,
) -> PendingSamplingBatch

Submit up to 128 independent samples in one RPC.

Uses the new sampling transport; requires server support. Consume pending.as_completed() to receive individual successes/failures. policy_selection="ordered" preserves training-queue ordering. policy_selection="latest_snapshot" samples the newest available immutable committed snapshot without waiting for queued training work. Publish one with save_weights(mode="inference", immutable=True) first; missing snapshots fail explicitly. This mode may lag the committed head. pinned_policy_id continues a previously returned policy snapshot and cannot be combined with latest_snapshot. Alternatively, retained_kv_groups supplies one trajectory UUID per independent sample (including the num_samples expansion) and permits earlier-policy KV within kv_cache_policy: max_staleness (updates), anchor_policy_id, previous_policy_id, and refill_on_limit (bool). Continuations echo the prior result's cache origin and sampling policy. Check Sample.retained_kv for worker acknowledgement. return_prompt_token_ids echoes canonical IDs without requesting prompt logprobs, preserving prefix-cache eligibility. Other arguments match sample. Existing submit_sample retains its whole-batch future behavior.

Model.submit_sample

Model.submit_sample(
    prompts: str | list[str] | None = None,
    *,
    num_samples: int = 1,
    max_tokens: int = 256,
    temperature: float = 1.0,
    top_p: float = 1.0,
    top_k: int = -1,
    stop: list[str] | None = None,
    seed: int | None = None,
    seeds: list[int] | None = None,
    return_prompt_logprobs: bool = False,
    logprobs: int | None = None,
    images: list[Image] | list[list[Image]] | None = None,
    return_expert_routing: bool = False,
    prompt_token_ids: list[int] | list[list[int]] | None = None,
    model_input: list[dict] | list[list[dict]] | None = None,
    metrics_type: str = '',
    timeout: float = 86400.0,
    poll_interval: float = 1.0,
) -> PendingSample

Submit sampling from current training weights without waiting.

See sample for the full kwarg reference.

Model.chat_complete

Model.chat_complete(
    messages: list[dict],
    *,
    timeout: float | None = None,
    **kwargs,
) -> ChatCompleteResult

OpenAI-compatible chat completion from current training weights.

Like model.sample() but using the OpenAI chat-completions format instead of raw prompts.

Returns: ChatCompleteResult — ChatCompleteResult with response_json (full OpenAI-format JSON string) and status_code.

Model.save_weights

Model.save_weights(
    name: str,
    mode: str = 'training',
    timeout: float = 86400.0,
    ttl: datetime.timedelta | None = None,
    immutable: bool = False,
    expected_policy_id: str | None = None,
) -> Checkpoint

Save a checkpoint of the current model weights.

Returns: Checkpoint — Checkpoint object with river:// path.

Model.load_weights

Model.load_weights(
    checkpoint: str | Checkpoint,
    load_optimizer: bool = True,
    timeout: float = 86400.0,
) -> None

Load weights from a checkpoint.

Model.model_id

Model.model_id: str

Model.base_model

Model.base_model: str

Model.training_run_id

Model.training_run_id: str

Model.step

Model.step: int

Current training step.

Advances when an optimizer step is submitted (so pipelined submissions observe a consistent value), not when it resolves — a failed or timed-out optimizer step leaves it advanced.

Data types

Values returned by the client and the configuration objects you pass to it.

LoraConfig

LoRA adapter configuration.

Sample

A generated sample.

tokens contains generated token IDs when the response includes them; older responses may fall back to retokenizing text client-side.

prompt_token_ids / top_logprobs / prompt_top_logprobs are None when the server did not provide them or the feature was not requested.

expert_routing is populated when return_expert_routing=True was passed on the request.

metrics carries per-result scalar metrics; empty unless the originating request passed a recognized metrics_type token.

Sample.routing_datum_keys

Sample.routing_datum_keys(*, required: bool = False) -> dict[str, bytes | str]

Return the per-datum keys to splat into a forward_backward datum when enabling force_routing_replay or compute_expert_flip_metric.

Returns an empty dict when this sample carries no routing capture and required is false. Raises ValueError when required is true and no routing handle is available.

Example:

data = []
for sample in samples:
    datum = {
        "input_ids": ...,
        "advantages": ...,
        **sample.routing_datum_keys(required=True),
    }
    data.append(datum)

ChatCompleteResult

Result of a chat completion request.

ForwardResult

Result of forward or forward_backward pass.

OptimStepResult

Result of an optimizer step.

Checkpoint

A saved model checkpoint.

TopLogprob

One top-K candidate token at a single position.

PendingOp

A submitted but not-yet-resolved async operation.

Returned by Model.submit_forward_backward() and Model.submit_optim_step(). Call .result() to block until complete.

PendingOp.result

PendingOp.result(
    *,
    before_poll: Callable[[], None] | None = None,
) -> ForwardResult | OptimStepResult

Block until the operation completes and return the result.

PendingSample

A submitted but not-yet-resolved sampling operation.

PendingSample.result

PendingSample.result(
    *,
    retry_connection_errors: bool = True,
    before_poll: Callable[[], None] | None = None,
) -> list[list[Sample]]

Block until sampling completes and return grouped samples.

retry_connection_errors=False exposes the first RetrieveFuture transport failure to callers that own a sealed retry policy.

PendingSample.result_async

async PendingSample.result_async(
    *,
    executor: Executor | None = None,
    retry_connection_errors: bool = True,
    before_poll: Callable[[], None] | None = None,
) -> list[list[Sample]]

Wait with bounded RPC workers, releasing them between polls.

Pass a shared executor to bound simultaneous RPCs and result decoding independently from the number of outstanding sampling operations.

ExpertRouting

Captured MoE expert routing for one sample.

Populated when the caller passed return_expert_routing=True on the sample request and routing capture was available.

To enable router replay / flip-rate metrics, splat the canonical per-datum keys via Sample.routing_datum_keys() rather than building them by hand:

datum = {..., **sample.routing_datum_keys(required=True)}

Current servers return routing captures as an opaque handle; the splat yields expert_routing_handle. topk_ids is retained for inspection and proto round-trip compatibility only; replay always recomputes the routing weights in the trainer.

The shape header fields (num_decoder_layers / top_k / layer_indices) are informational only and let the user inspect capture metadata.

Byte layout (when ids are present): [num_tokens, num_decoder_layers, top_k] row-major int16 little-endian. num_tokens = seqlen - 1 because routing capture omits the trailing position.

Deployment

A dedicated streaming deployment: one checkpoint on its own engines.

Dedicated deployments are gated and disabled by default. Contact River to enable access for your team and base model before creating capacity.

base_url binds the checkpoint when used with the standard OpenAI client, so its existing model argument can stay unchanged. model is the deployment identity (equal to id), also accepted by global /v1. replicas is keyed by role: unified for unified serving, prefill and decode for disaggregated serving. phase is one of accepted, provisioning, ready, degraded, unavailable, scaled_to_zero, failed, deleting, deleted.

DeploymentReplicas

Desired, allocated, ready and pending counts for one role of a deployment.

pending are allocated replicas the cell cannot place yet (no capacity).

ServerCapabilities

Supported public protocols and model names, without server internals.

ServerCapabilities.require_model

ServerCapabilities.require_model(model: str, *features: str) -> None

ServerCapabilities.require

ServerCapabilities.require(*features: str) -> None

PolicyVersion

Server identity of a committed logical weight state (not a byte hash).

PolicyVersion.from_proto

PolicyVersion.from_proto(value)

ImageHandle

Immutable metadata for an image owned by one live API session.

Checkpoints retain bytes in a client-side ImageStore and re-upload on resume.

ImageStore

Content-addressed image bytes on a local/shared durable filesystem.

ImageStore.put

ImageStore.put(data)

ImageStore.read

ImageStore.read(image)

ImageStore.copy_images

ImageStore.copy_images(value, source)

ImageStore.prune

ImageStore.prune(retained)

PendingSamplingBatch

Independent operations accepted together; results arrive in completion order.

as_completed reports failures per sample. Cancelling local waiting does not cancel accepted server operations. Retain request_ids to retrieve them.

PendingSamplingBatch.as_completed

async PendingSamplingBatch.as_completed(
    *,
    executor: Executor | None = None,
    collector: SamplingResultCollector | None = None,
) -> AsyncIterator[SamplingCompletion]

Consume independently; share a collector across concurrent batches.

SamplingCompletion

One sample or failure, identified in the original prompt/sample grid.

SamplingResultCollector

Coalesce result snapshots across batches on one event loop.

At most max_poll_rpcs short streams run concurrently, each querying up to 128 handles. Backpressure permits one decoded result buffered per stream. Executors are caller-owned; no thread is occupied between snapshots.

SamplingResultCollector.register

SamplingResultCollector.register(batch: PendingSamplingBatch) -> list[asyncio.Future]

SamplingResultCollector.wake

SamplingResultCollector.wake()

SamplingResultCollector.aclose

async SamplingResultCollector.aclose()

RunMetricsLogger

Buffers and sends run metrics without interrupting training.

RunMetricsLogger.log

RunMetricsLogger.log(metric: str, step: int, value: float) -> None

Queue one metric point and return without waiting for the network.

RunMetricsLogger.flush

RunMetricsLogger.flush() -> None

Request prompt delivery of the points currently in the buffer.

RunMetricsLogger.close

RunMetricsLogger.close() -> None

Drain buffered points when possible, then stop the worker.

TrainingDataArtifact

Source bytes and expected digest for server-side integrity verification.

The API hashes content itself and retains only the resulting manifest. Callers should verify their source files locally before constructing these objects, then bind the returned TrainingDataAttestation to the model they create.

AttestedTrainingDataArtifact

Server-observed source-artifact digest and size.

TrainingDataAttestation

A server-verified source-artifact manifest bound to one session.

Errors

Every server and transport failure derives from RiverError, so a single except river.RiverError catches all of them. Invalid arguments still raise the standard ValueError and TypeError.

RiverError

Inherits Exception.

Base exception for River client errors.

AuthenticationError

Inherits RiverError.

Authentication failed.

CapacityError

Inherits RiverError.

No capacity available for the operation.

ModelNotFoundError

Inherits RiverError.

Model not found.

RiverConnectionError

Inherits RiverError.

Connection or communication error with the River API server.

This wraps gRPC errors with a more user-friendly message while preserving the original error details for debugging.

RiverConnectionError.from_grpc_error

RiverConnectionError.from_grpc_error(
    error: Exception,
    context: str = 'API call',
) -> RiverConnectionError

Create a RiverConnectionError from a gRPC RpcError.

RiverTimeoutError

Inherits RiverError.

Operation timed out while retaining its recoverable future ID.

SessionHeartbeatError

Inherits RiverConnectionError.

Session heartbeat was lost or rejected while a session was active.

Tokenizers

Helpers for tokenizing prompts yourself, for example when you build prompt_token_ids for training data.

load_tokenizer

load_tokenizer(
    tokenizer: str | Any | None = None,
    *,
    base_model: str | None = None,
    revision: str | None = None,
    local_files_only: bool = False,
    resolve_aliases: bool = True,
)

Load or return a tokenizer for River client result parsing.

base_model remains the public River model name used for routing. When it is a hyphen-suffixed deployment alias of a known canonical model, this helper resolves it to the underlying Hugging Face tokenizer id before loading. local_files_only keeps sealed jobs from reaching Hugging Face after their tokenizer revision has been frozen. resolve_aliases=False retains a deployment's own tokenizer source for unpinned compatibility paths.

resolve_tokenizer_name

resolve_tokenizer_name(model_name: str) -> str

Return the tokenizer for a canonical River model or one of its aliases.

Deployment aliases must append a hyphen-delimited suffix to a canonical model name. Longest-root matching keeps the result deterministic if a future canonical model name extends another root.

MODEL_TOKENIZER_ALIASES

MODEL_TOKENIZER_ALIASES: dict[str, str]

Reinforcement learning

The river_client.rl library manages multi-turn rollouts, rewards, policy staleness, and optimizer updates on top of Client, Session, and Model.

from river_client import rl

rl.AsyncTrainer

Train a model from rollout groups while sampling continues.

Computes advantages from completed groups, submits forward/backward work, and applies optimizer updates. Set max_staleness=0 to train synchronously or allow bounded policy lag for asynchronous training. Use one trainer to update a given model.

rl.AsyncTrainer(
    *,
    engine,
    optimizer: Adam,
    completion: GroupCompletion,
    normalize: str,
    groups_per_step: int,
    group_size: int = 16,
    advantage=None,
    loss: str = 'cispo',
    max_staleness: int = 2,
    stale_policy: str = 'wait',
    allow_trajectory_loss: bool = False,
    allow_unbounded_staleness: bool = False,
    truncation: Truncation = Truncation(),
    checkpoint: Checkpointing | None = None,
    init_checkpoint=None,
    evaluator=None,
    run_config=None,
    max_importance_weight: float = 5.0,
    routing_replay: bool = False,
    repeat_dataset: bool = True,
    batch_order: str = 'completion',
    validate_batch=None,
    forward_backward_groups: int | None = None,
    forward_backward_batch: ForwardBackwardBatch | Literal['auto'] | None = 'auto',
    max_pending_forward_backward: int = 16,
    **loss_config,
)

Methods: progress, run

rl.AsyncTrainer.progress

rl.AsyncTrainer.progress()

Return current batch, model-update, and activity counters.

model_step reports the last completed optimizer update, including while another update is in flight.

rl.AsyncTrainer.run

async rl.AsyncTrainer.run(dataset, *, steps: int, after_recovery=None)

Yield a Step after each completed training batch.

rl.RolloutEngine

Generate multi-turn trajectories from an environment and sampling model.

Coordinates rollout groups, sampling concurrency, and policy/KV-cache continuity. Use it directly for custom algorithms or pass it to AsyncTrainer for managed training.

rl.RolloutEngine(
    model,
    *,
    env,
    renderer,
    budget: Budget = Budget(),
    schedule: Schedule = Schedule(),
    sampling_policy: str = 'trajectory',
    kv_cache: KVCache = KVCache(),
    temperature: float = 1.0,
    top_p: float = 1.0,
    top_k: int = -1,
    seed: int = 0,
    sample_timeout: float = 1800.0,
    environment_timeout: float | None = None,
    return_expert_routing: bool = False,
)

Methods: submit, next_group, accept, acknowledge, rollout, progress, sampling_metrics, state_dict, load_state_dict

rl.RolloutEngine.submit

rl.RolloutEngine.submit(
    row,
    *,
    group_size: int,
    completion: GroupCompletion,
    id: str | None = None,
    priority: int = 0,
)

Start a rollout group for one dataset row and return its group id.

Call inside async with engine. Each member gets an independent sampling seed.

rl.RolloutEngine.next_group

async rl.RolloutEngine.next_group() -> CompletedGroup

Wait for the next completed group or deadline partial result.

rl.RolloutEngine.accept

rl.RolloutEngine.accept(result)

Mark a delivered result as consumed by the caller.

rl.RolloutEngine.acknowledge

rl.RolloutEngine.acknowledge(group_id)

Release a group only after all of its emitted portions are consumed.

rl.RolloutEngine.rollout

async rl.RolloutEngine.rollout(rows, *, group_size: int, completion: GroupCompletion)

Yield completed trajectory groups; custom algorithms can own updates.

rl.RolloutEngine.progress