inference-server

Inference Server

Starting the server

Always use the inference entry point — never vllm serve or python -m vllm.entrypoints.openai.api_server directly. The entry point runs setup_vllm_env() which configures environment variables (LoRA, multiprocessing) before vLLM is imported.

# With a TOML config
uv run inference @ path/to/config.toml
 
# With CLI overrides
uv run inference --model.name Qwen/Qwen3-0.6B --model.max_model_len 2048 --model.enforce_eager
 
# Combined
uv run inference @ path/to/config.toml

SLURM scheduling

The inference entrypoint supports optional SLURM scheduling, following the same patterns as SFT and RL.

Single-node SLURM

Multi-node SLURM (independent vLLM replicas)

Each node runs an independent vLLM instance. No cross-node parallelism — TP and DP must fit within a single node's GPUs.

Dry run

Add dry_run = true to generate the sbatch script without submitting:

Custom endpoints

The server extends vLLM with:

• /v1/chat/completions/tokens — accepts token IDs as prompt input (used by multi-turn RL rollouts)
• /update_weights — hot-reload model weights from the trainer
• /load_lora_adapter — load LoRA adapters at runtime
• /init_broadcaster — initialize weight broadcast for distributed training

Testing the server

Key files

• src/prime_rl/entrypoints/inference.py — entrypoint with local/SLURM routing
• src/prime_rl/inference/server.py — vLLM env setup
• src/prime_rl/configs/inference.py — InferenceConfig and all sub-configs
• src/prime_rl/inference/vllm/server.py — FastAPI routes and vLLM monkey-patches
• src/prime_rl/templates/inference.sbatch.j2 — SLURM template (handles both single and multi-node)
• configs/debug/infer.toml — minimal debug config