Client API

Client API

The ArcticRL client offers a unified frontend for RL training workflows. It works identically against remote deployments and local server instances, with the same API regardless of whether you use HTTP or Ray communication. All training, sampling, and log-probability jobs are initialized automatically at construction time.

Creating a Client

Use the factory function to create a client based on your configuration:

from arctic_platform.rl.client import create_arctic_rl_client
from arctic_platform.rl.config import ArcticRLClientConfig

config = ArcticRLClientConfig(
    backend="local",
    comm_protocol="ray",  # or "http"
    model_name="meta-llama/Llama-3-8B",
    training_gpus=4,
    sampling_gpus=2,
    log_prob_gpus=2,
)

client = create_arctic_rl_client(config)

The factory automatically selects ArcticRLHTTPClient for HTTP protocol or ArcticRLRayClient for Ray protocol.

Job ID Properties

Access the initialized job IDs for training, sampling, and log-probability computation:

training_id = client.training_job_id    # int
sampling_id = client.sampling_job_id    # int
log_prob_id = client.log_prob_job_id    # int or None

All three raise ValueError if accessed before their respective jobs are initialized (e.g., if the corresponding GPU allocation was zero).

Reconnection (Ray Clients Only)

For Ray clients, you can serialize the client's configuration to reconnect from another process without re-initializing jobs:

# driver process
client = create_arctic_rl_client(config)
reconnect_cfg = client.reconnect_config()  # returns ArcticRLClientConfig (Ray-serializable)

# Ray actor or different process
actor_client = create_arctic_rl_client(reconnect_cfg)  # connects to same jobs, no /initialize

The returned config has backend="local" and comm_protocol="ray", with job IDs populated.

Training Operations

Forward-Backward Pass

Perform a forward-backward pass on the training engine:

response = await client.fwd_bwd(
    batch={
        "args": [input_ids],
        "kwargs": {"attention_mask": attention_mask},
        "context": {
            "old_log_probs_shifted": old_lps,
            "prox_logp_shifted": prox_lps,
            "ref_log_probs_shifted": ref_lps,
        }
    },
    processing={
        "loss_fn": "arctic_platform.rl.processors.grpo_loss",
        "config": {"eps_clip": 0.2},
        "post": [],
    }
)

Parameters:

• batch: Dict containing args, kwargs (model inputs), and optional context (RL tensors).
• processing (optional): Processing descriptor with loss_fn, config, and post fields. Pass None to use the server's legacy loss configuration.

Log-Prob Convention: All log-probability tensors in context must follow the torch.roll(input_ids, shifts=-1) convention (the _shifted suffix encodes this):

• old_log_probs_shifted — behavioral policy log-probs after roll
• prox_logp_shifted — proximal log-probs after roll
• ref_log_probs_shifted — reference policy log-probs after roll

Forward (No Gradient)

Forward-only pass without gradient computation, used for log-probability or reference model inference:

response = await client.fwd_no_grad(
    batch={"args": [input_ids], "kwargs": {"attention_mask": mask}},
    reference_model=True  # True → log_prob_job_id, False → training_job_id
)

Optimizer Step

Execute one optimizer step on the training engine:

response = await client.step()

Save Checkpoint

Save the current training checkpoint:

response = await client.save_checkpoint()

Sampling Operations

Generate Text

Generate text completions using the sampling engine:

results = await client.generate(
    prompts=["Hello, world!", "The future of AI is"],
    sampling_params={"max_tokens": 100, "temperature": 0.8}
)

Returns: List of result dicts from the sampling engine.

Log Probability Computation

Compute Log Probabilities

Compute per-token log probabilities for prompts and completions:

response = await client.log_probs(
    prompts=["The capital of France is"],
    completions=["Paris"],
    top_k=1
)

Parameters:

• prompts: List of prompt strings.
• completions (optional): List of completion strings. If None, computes log-probs for prompts only.
• top_k: Number of top tokens to return log-probs for (default: 1).

Weight Synchronization

Sync Training Weights to Sampling

Copy training model weights to the sampling engine:

response = await client.sync_weights(
    cuda_ipc=False,      # True: zero-copy CUDA IPC, False: CPU file path
    low_memory=False     # Only applies to cuda_ipc=True path
)

In colocated mode:

• cuda_ipc=True: Zero-copy CUDA IPC (training weights must be on GPU)
• cuda_ipc=False: CPU file path (works when training is offloaded)
• low_memory=True: Stream one parameter at a time to reduce peak GPU memory usage

Weight Norm

Compute and compare L2 weight norms across training and sampling engines:

response = await client.weight_norm()
# Returns: {"training_norm": float, "sampling_norm": float, ...}

After a weight sync, both norms must match. Intended for testing sync correctness.

Lifecycle Management (Colocated Mode)

All methods below are for managing GPU memory in colocated training/sampling setups.

Inference Engine

# Sleep the sampling inference engine (free GPU memory)
response = await client.sleep_inference(level=1)

# Wake the sampling inference engine
response = await client.wake_inference(tags=None)

# Reset the prefix cache on the sampling engine
response = await client.reset_prefix_cache()

Training Engine

# Offload training state to CPU
response = await client.sleep_training(mode="all")
# mode="all":       Everything (training → inference transition)
# mode="non_lp":    Keep bf16 params, offload rest (before CUDA IPC)
# mode="lp_params": Offload bf16 params only (after CUDA IPC)

# Reload training state to GPU
response = await client.wake_training()

# Release PyTorch cached GPU memory on all training workers
response = await client.empty_training_cache()

Log-Probability Engine

# Offload the reference (log-prob) DeepSpeed engine to CPU
response = await client.sleep_log_prob()

# Reload the reference (log-prob) DeepSpeed engine to GPU
response = await client.wake_log_prob()

Both methods return {"status": "no_log_prob_job"} if no log-probability job is initialized.

Shutdown

Clean up all jobs and stop the local server if applicable:

await client.shutdown()

This destroys all training, sampling, and log-probability jobs in order and resets their IDs to None.