Ray Server & Deployment

Ray Server & Deployment

Overview

The Arctic RL platform provides ArcticRLRayServer to manage distributed training and inference workloads using Ray. The server orchestrates:

• Training workers using DeepSpeed
• Sampling inference via ArcticInference ReplicaPools (vLLM-based)
• Log probability computation via DeepSpeed or vLLM-based engines

The server runs as a Ray actor and coordinates GPU resource allocation, weight synchronization, and job lifecycle management across multiple GPUs and nodes.

Quick Start

Start the Ray server with specified GPU allocations:

python -m arctic_platform.rl.server \
    --training-gpus 4 --sampling-gpus 2 --log-prob-gpus 2

Architecture

Server State

The ArcticRLRayServerState class manages all server operations. It is instantiated as a Ray remote actor via create_arctic_rl_ray_server_state():

from arctic_platform.rl.ray_server import create_arctic_rl_ray_server_state

server_state = create_arctic_rl_ray_server_state(
    training_gpus=4,
    sampling_gpus=2,
    log_prob_gpus=2,
    log_prob_engine="vllm",  # or "deepspeed"
    colocate=True
)

Key Parameters:

At least one GPU allocation must be greater than zero.

Job Types

The server manages three job types via the initialize() method:

Training Jobs

Initializes DeepSpeed workers across the specified GPU count:

await server_state.initialize.remote({
    "job_type": "training",
    "model_name": "Snowflake/arctic-embed-large",
    "checkpoint_path": "/path/to/checkpoints",
    "ds_config": {...},  # DeepSpeed config
})

Behavior:

• Creates one DeepSpeed worker per GPU
• Establishes NCCL communication with a rendezvous master port (default 29500)
• Stores checkpoint path and weight sync path
• Raises ValueError if training is already running or no training GPUs configured

Rendezvous Port: Set via MASTER_PORT environment variable (default: 29500). For concurrent jobs on the same host, override this to avoid port conflicts.

Sampling Jobs

Initializes ArcticInference ReplicaPool for inference via vLLM:

await server_state.initialize.remote({
    "job_type": "sampling",
    "model_name": "Snowflake/arctic",
    "vllm_config": {"gpu_memory_utilization": 0.9},
    "arctic_inference_config": {...},
})

Behavior:

• Creates vLLM replicas with tensor parallelism (TP) support
• In colocate mode, enables vLLM's sleep mode for memory management
• Configures distributed Ray execution backend when TP > 1
• Sets environment variables for Arctic Inference integration

GPU Fractions in Colocate Mode:

When colocate=True, sampling uses a fractional GPU claim (0.33) within its bundle, allowing time-sharing with training and log_prob tasks on the same physical GPU.

Log Probability Jobs

Initializes either DeepSpeed or vLLM engine for log probability computation:

# DeepSpeed engine
await server_state.initialize.remote({
    "job_type": "log_prob",
    "model_name": "Snowflake/arctic",
    "ds_config": {...},
})

# vLLM engine
await server_state.initialize.remote({
    "job_type": "log_prob",
    "model_name": "Snowflake/arctic",
    "vllm_config": {"gpu_memory_utilization": 0.9},
})

DeepSpeed Engine:

• Creates DeepSpeed workers on bundle offset 0 (overlaps training bundles)
• Reference engine starts resident on GPU
• Offloading controlled by client via fsdp_config.param_offload
• Rendezvous port default: 29501

vLLM Engine:

• Creates ReplicaPool with tensor parallelism support
• Enables sleep mode in colocate configuration
• Shares bundles with training via offset 0

Colocated GPU Resource Sharing

When colocate=True, the server uses placement groups to pack training, sampling, and log_prob tasks on the same physical GPUs:

_COLOCATE_GPU_FRACTIONS = {
    "sampling": 0.33,
    "log_prob": 0.33,
    "training": 0.34
}

Layout (deterministic):

• Training rank r → bundle r
• Sampling replica r (with tensor parallel factor tp) → bundles [r*tp .. r*tp+tp-1]
• Log prob rank r → bundle r

Each physical GPU (bundle) hosts one training rank, one sampling replica, and one log_prob rank. The fractional GPU allocations are Ray scheduling accounting only—real VRAM is time-shared via sleep() and wake() calls.

Memory Management

In colocate mode:

• Sleep Mode: Inference engines offload weights to CPU, freeing GPU memory
• Wake Mode: Restores weights and GPU state for inference
• Cache Emptying: Training workers empty CUDA cache before wake to allow memory remapping

# Sleep inference engines
await server_state.sleep_inference.remote(job_id=1, level=1)

# Wake inference engines
await server_state.wake_inference.remote(
    tags=["weights"],
    restore_weights=True
)

Weight Synchronization

The server manages weight synchronization for training jobs:

Bucket Size: 256 MB (configurable via weight_sync_bucket_size)

Sync Port: Controlled via ARL_WEIGHT_SYNC_PORT environment variable (default: 29600). Override for concurrent jobs on the same host to avoid port conflicts.

Sync Path: Training jobs write synchronized weights to:

{checkpoint_path}/arctic_rl_job_{job_id}/weight_sync.pt

Inference Pool Operations

Reset Prefix Cache

Clears KV cache for sampling and log_prob engines:

await server_state.reset_prefix_cache.remote(job_id=1)

Sleep Inference

Reduces GPU memory footprint by offloading weights:

results = await server_state.sleep_inference.remote(job_id=1, level=1)
# Returns: {"job_id": 1, "sampling": {...}, "log_prob": {...}}

Wake Inference

Restores GPU memory and re-initializes engines:

results = await server_state.wake_inference.remote(
    tags=None,  # or ["weights"], ["config"], etc.
    restore_weights=True
)

The restore_weights parameter allows callers about to overwrite all weights (e.g., CUDA IPC sync) to skip the redundant CPU→GPU copy.

Job Lifecycle

Initialize

Creates workers/replicas and registers the job:

result = await server_state.initialize.remote(job_config_dict)
# Returns: {"job_id": 1, "job_type": "training", "running": true}

Destroy

Cleans up workers/replicas and removes the job:

await server_state.destroy.remote(job_id=1, job_type="training")

Destroys:

• Training workers and clears worker list
• Sampling ReplicaPool via shutdown()
• Log prob workers (DeepSpeed) or ReplicaPool (vLLM)

Environment Variables

Querying Server State

The server exposes async getter methods for monitoring:

# Get all jobs
jobs = await server_state.get_jobs.remote()

# Get training workers
workers = await server_state.get_training_workers.remote()

# Get sampling/log_prob pools
pool = await server_state.get_sampling_pool.remote()
lp_pool = await server_state.get_log_prob_pool.remote()

# Get configuration
colocate = await server_state.get_colocate.remote()
bucket_size = await server_state.get_weight_sync_bucket_size.remote()

Logging

The server initializes logging at INFO level with timestamp, logger name, level, and message:

2025-01-15 10:30:45,123 arctic_platform.rl.ray_server INFO [ArcticRLRayServer] initializing ray cluster

All output is prefixed with pr0() (rank-0 only) to avoid duplicate logs across distributed workers.