prime rl

The Problem

Large‑scale reinforcement‑learning experiments need a way to run many asynchronous agents across dozens or thousands of GPUs while keeping the training loop simple enough for research iteration. Existing RL libraries are either single‑process or require heavyweight orchestration code that is hard to extend.

What This Does

prime‑rl provides an end‑to‑end framework for asynchronous RL at scale.

• Core training logic lives under src/prime_rl/ (e.g., trainer/models/, trainer/ckpt.py, orchestrator/).
• Configuration is declarative: a family of TOML files in configs/ (e.g., configs/rl.toml, configs/orchestrator.toml) drive the trainer, the inference server, and the multi‑node deployment.
• Two server entry points exist: src/prime_rl/inference/server.py (standard FastAPI‑style server) and src/prime_rl/inference/vllm/server.py (vLLM‑backed inference). Both expose a main function that starts the HTTP API used by downstream agents.

The framework also ships Docker support (Dockerfile.cuda) and Helm‑style Kubernetes manifests (k8s/prime-rl/) for production clusters.

How It Is Wired

Execution begins at src/prime_rl/inference/server.py:main (line 17). main builds a logger (get_logger – called from 46 places) and registers the FastAPI app. The request handler eventually calls interleave_rollout, which is the highest‑frequency call path (20 callers).

Key shared utilities:

The import graph is shallow (187 modules, 37 edges, no cycles), but a few “hub” modules (tests/conftest.py, tests/utils.py) are referenced by many tests, indicating test‑side churn rather than production risk.

External effects are limited to:

• Filesystem reads/writes (e.g., load_json_dir in utils/pathing.py).
• Network calls from src/prime_rl/utils/elastic.py (discovering server IPs).
• One outbound call to a model inference backend in the vLLM server.

No persistent database or message‑queue code is present in the repository itself.

How To Use It

# Clone (use the exact URL the client requested)
git clone https://github.com/moses-y/prime-rl
cd prime-rl

# Install the UV toolchain (recommended by the README)
curl -LsSf https://astral.sh/uv/install.sh | sh
source $HOME/.local/bin/env

# Install dependencies (the repo ships a pyproject.toml but no lockfile)
uv sync --all-extras          # pulls all optional extras

# Optional: build the Docker image for GPU workers
docker build -f Dockerfile.cuda -t prime-rl:latest .

# Run the inference server (example)
uv run src/prime_rl/inference/server.py --config configs/inference.toml

Configuration files such as configs/debug/rl/train.toml define the model, environment, and scaling parameters. The orchestrator can be launched with uv run src/prime_rl/orchestrator/orchestrator.py -c configs/ci/integration/rl/start.toml.

Real‑World Use

A research team wanting to train a policy on 128 GPUs could:

1. Create a custom my_experiment.toml under configs/ that points to their environment and model checkpoint.
2. Build the Docker image and push it to the cluster registry.
3. Deploy the Helm chart in k8s/prime-rl/, mounting the config as a ConfigMap.
4. Start the orchestrator (uv run src/prime_rl/orchestrator/orchestrator.py -c my_experiment.toml).
5. Clients query the inference server (src/prime_rl/inference/vllm/server.py) for action predictions during rollouts.

All components communicate via the filesystem (shared checkpoint directory) and lightweight HTTP calls; no external message broker is required.

Code Health & Issues

Static analysis identified the following concrete problems (no speculation added):

• HIGH – GitHub Actions use mutable tags (.github/workflows/*). Pin each action to a 40‑character commit SHA.
• MEDIUM – No least‑privilege GITHUB_TOKEN permissions declared (cpu_tests.yaml). Add permissions: contents: read.
• MEDIUM – Dependabot / Renovate not configured. Add .github/dependabot.yml.
• MEDIUM – Docker base images are unpinned (Dockerfile.cuda). Use digest tags (...@sha256:…).
• MEDIUM – No dependency‑vulnerability scan in CI. Add dependency-review-action or osv-scanner.
• LOW – Workflow jobs lack timeout-minutes. Set reasonable limits.

Additional hygiene notes: the repository has a license, tests, and CI, but lacks a lockfile for reproducible Python builds and contains several oversized or deeply nested files (e.g., src/prime_rl/orchestrator/orchestrator.py – 1853 lines, nesting depth 8). Duplicate configuration blocks appear across many configs/*.py files.

The Bottom Line

prime‑rl delivers a functional, scalable async RL stack with clear entry points and a modular config system, making it suitable for research teams operating multi‑GPU clusters. However, the codebase suffers from maintainability problems (large monolithic files, deep nesting, duplicated config logic) and a few production‑readiness gaps (unpinned dependencies, missing lockfile, limited CI safety checks). Addressing these hygiene issues will improve reliability and ease future extensions.