diff --git a/.github/workflows/cpu_tests.yaml b/.github/workflows/cpu_tests.yaml index 4ed424bcb6c..12fdd881edc 100644 --- a/.github/workflows/cpu_tests.yaml +++ b/.github/workflows/cpu_tests.yaml @@ -77,8 +77,11 @@ jobs: from prime_rl.utils.config import BaseConfig, find_package_resource, rgetattr, rsetattr from prime_rl.utils.validation import validate_shared_ckpt_config + # `verifiers` (+ its `datasets` dep) is a declared slim dep: the v1 config types + # (EnvConfig, Task, ...) extend verifiers.v1, which is pure-pydantic and pulls no + # GPU/ML deps. We still forbid the actual heavy training deps below. forbidden = ["torch", "transformers", "vllm", "wandb", "ring_flash_attn", - "verifiers", "prime", "datasets", "liger_kernel", "loguru"] + "prime", "liger_kernel", "loguru"] leaked = [m for m in forbidden if m in sys.modules] if leaked: raise SystemExit(f"slim install leaked heavy deps into sys.modules: {leaked}") diff --git a/CHANGELOG.md b/CHANGELOG.md index 700bcdd925c..6981598e01e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,12 @@ Documenting **breaking** configuration changes — renamed, removed, or moved fields that require users to update existing configs. +- **verifiers v1 training integration** (collection of removals/renames). The orchestrator now trains against a v1 env server returning `vf.Trace` rollouts, renderer-only, with a unified token schema. + - **`[[orchestrator.train.env]]` is now a v1 `vf.EnvServerConfig`**: a v1 env is `taskset = { id = "" }` + `harness = { id = "default", runtime = { type = "subprocess" } }`. A bare `id = ""` (optionally with `args`) runs the v0 env through the verifiers legacy bridge — old flat `id`/`args`-only configs keep working as legacy envs; v1 envs must use the `taskset`/`harness` shape. + - **`orchestrator.train.env` no longer defaults to a `reverse-text` env**: the default is now empty (`[]`), so configs must declare `[[orchestrator.train.env]]` explicitly. + - **`...env[].num_workers` → `pool`**: per-env worker sizing moves to a `pool` config (`static`/`elastic`). Back-compat: `num_workers = N` maps to `pool = { type = "static", num_workers = N }` and `"auto"` to the default elastic pool; the group-level `orchestrator.train.num_workers` / `orchestrator.eval.num_workers` are removed. + - **`orchestrator.renderer` is now required (MITO removed)**: it was `RendererConfig | None`, where `renderer = None` selected the MITO path; training is now renderer-only, so `renderer` is non-optional. Drop `renderer = None` / `renderer = "None"`. + - **`...client.timeout` / `...client.connect_timeout` removed**: the v1 rollout client manages its own timeouts. Drop these from `[orchestrator.student.client]` and `[orchestrator.teacher.client]`. (2026-06-23) - **`orchestrator.advantage` length-penalty overhaul + length-weighted baseline**: `length_penalty` is now a single shape — a linear penalty that subtracts `coef * pass_rate * (model completion tokens / orchestrator.seq_len)` from each reward before the baseline subtraction (`pass_rate` = group mean reward, so solved-often problems get the strongest concision pressure and never-solved groups get none), with an optional `gate_by_correctness` flag (default `false`) that restricts the penalty to correct rollouts (`reward == 1`). The previous discriminated `type = "tokens" | "turns"` efficiency-shaping variants and the per-penalty `max_length` field have been removed — configure `[orchestrator.advantage.length_penalty]` with `coef` (and optionally `gate_by_correctness`) instead; the denominator is always `orchestrator.seq_len`. Also adds `orchestrator.advantage.length_weighted_baseline` (default `false`): when enabled, the GRPO baseline is the token-length-weighted mean reward (`sum(len_i * reward_i) / sum(len_i)`) instead of the plain mean. (2026-06-03) - **`trainer.experimental.token_export` removed → `trainer.enable_token_export`**: `TrainerExperimentalConfig` and the `[trainer.experimental] token_export = true` setting have been removed with no migration alias. Existing trainer configs must delete the `[trainer.experimental]` block and set the new top-level trainer flag instead: `enable_token_export = true` in a trainer config file, or `trainer.enable_token_export = true` when nested under `[trainer]` in an RL config. The trainer entrypoint CLI flag is `--enable-token-export`. (2026-06-10) - **`inference.kv_cache_offload.cpu_bytes` removed → discriminated `type` config**: The flat `[inference.kv_cache_offload]` block with a single `cpu_bytes` field is replaced by a backend-discriminated union with composable `cpu`/`disk` tiers. Migrate native CPU offload from `[inference.kv_cache_offload]\ncpu_bytes = N` to `[inference.kv_cache_offload]\ntype = "native"` plus `[inference.kv_cache_offload.cpu]\nnum_bytes = N`. A `type = "mooncake"` backend (per-node distributed store; multi-node/SLURM only) and an optional `[inference.kv_cache_offload.disk]\npath = "..."` tier (layered behind cpu) are also available. `extra="forbid"` rejects the old `cpu_bytes` key, so existing configs must migrate. (2026-06-02) diff --git a/configs/ci/integration/reverse_text_rl_sft/start.toml b/configs/ci/integration/reverse_text_rl_sft/start.toml index 6b26bb33351..e658fab5d4e 100644 --- a/configs/ci/integration/reverse_text_rl_sft/start.toml +++ b/configs/ci/integration/reverse_text_rl_sft/start.toml @@ -19,6 +19,10 @@ training_mode = "sft" batch_size = 128 group_size = 16 +# Teacher rolls out over plain chat-completions (no tokens); the renderer backfills them. +[orchestrator.renderer] +name = "qwen3" + [orchestrator.train.sampling] max_completion_tokens = 128 diff --git a/configs/debug/reverse_text_v1.toml b/configs/debug/reverse_text_v1.toml deleted file mode 100644 index 1d23b965ac0..00000000000 --- a/configs/debug/reverse_text_v1.toml +++ /dev/null @@ -1,31 +0,0 @@ -max_steps = 20 -seq_len = 2048 - -[model] -name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" - -[wandb] -project = "reverse-text" -name = "reverse-text-v1" - -[orchestrator] -batch_size = 128 -group_size = 16 - -[orchestrator.train.sampling] -max_completion_tokens = 128 - -[[orchestrator.train.env]] -id = "reverse-text" -args = { v1 = true } - -[trainer.optim] -lr = 3e-6 - -[ckpt] # Checkpoint at the end of training - -[inference] - -# Model not in MODEL_RENDERER_MAP — opt into DefaultRenderer (apply_chat_template). -[orchestrator.renderer] -name = "default" diff --git a/configs/debug/training_modes/README.md b/configs/debug/training_modes/README.md index 96ccebb0096..c789f43b187 100644 --- a/configs/debug/training_modes/README.md +++ b/configs/debug/training_modes/README.md @@ -9,9 +9,8 @@ Minimal end-to-end configs for the three training modes (`rl` / `opd` / `sft`) a | `opd_lora.toml` | `opd` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | trains a LoRA adapter (rank 8) | | `sft.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | | | `sft_lora.toml` | `sft` | local vLLM (`Qwen3-0.6B-Reverse-Text-RL`) | trains a LoRA adapter (rank 8) | -| `sft_external.toml` | `sft` | PI inference (`openai/gpt-5-mini`) | external OAI endpoint; no local teacher | -The student inference server is auto-launched on GPU 0 at `http://localhost:8000/v1` with `gpu_memory_utilization=0.5`. The local teacher (used by everything except `rl.toml` and `sft_external.toml`) is **not** auto-launched — start it manually on GPU 1. +The student inference server is auto-launched on GPU 0 at `http://localhost:8000/v1` with `gpu_memory_utilization=0.5`. The local teacher (used by everything except `rl.toml`) is **not** auto-launched — start it manually on GPU 1. ## Start the local teacher @@ -38,10 +37,6 @@ uv run rl @ configs/debug/training_modes/opd_lora.toml # SFT hard distill (needs teacher on port 8001) uv run rl @ configs/debug/training_modes/sft.toml uv run rl @ configs/debug/training_modes/sft_lora.toml - -# SFT hard distill from openai/gpt-5-mini via PI inference -# (requires PRIME_API_KEY + PRIME_TEAM_ID in env; no local teacher needed) -uv run rl @ configs/debug/training_modes/sft_external.toml ``` See [docs/training.md](../../docs/training.md#training-modes-rl--opd--sft-via-orchestrator) for what each mode does. diff --git a/configs/debug/training_modes/sft.toml b/configs/debug/training_modes/sft.toml index aed5b30cb32..46c2afd1e8a 100644 --- a/configs/debug/training_modes/sft.toml +++ b/configs/debug/training_modes/sft.toml @@ -20,6 +20,10 @@ training_mode = "sft" batch_size = 128 group_size = 4 +# Teacher rolls out over plain chat-completions (no tokens); the renderer backfills them. +[orchestrator.renderer] +name = "qwen3" + [orchestrator.train.sampling] max_completion_tokens = 128 diff --git a/configs/debug/training_modes/sft_external.toml b/configs/debug/training_modes/sft_external.toml deleted file mode 100644 index cb9ea8d09e1..00000000000 --- a/configs/debug/training_modes/sft_external.toml +++ /dev/null @@ -1,55 +0,0 @@ -# SFT from openai/gpt-5-mini via PI inference. -# Requires PRIME_API_KEY + PRIME_TEAM_ID in the environment. -# -# Run with: -# uv run rl @ configs/debug/training_modes/sft_external.toml - -max_steps = 20 -seq_len = 2048 - -[model] -name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" - -[wandb] -project = "reverse-text-debug" -name = "debug-sft-external" - -[orchestrator] -training_mode = "sft" -batch_size = 128 -group_size = 4 - -[orchestrator.train.sampling] -max_completion_tokens = 2048 -extra_body = { reasoning_effort = "minimal" } - -[[orchestrator.train.env]] -id = "reverse-text" - -[orchestrator.eval] -interval = 1 -num_examples = 128 - -[orchestrator.eval.sampling] -max_completion_tokens = 128 - -[[orchestrator.eval.env]] -id = "reverse-text" - -[orchestrator.teacher.model] -name = "openai/gpt-5-mini" - -[orchestrator.teacher.client] -base_url = ["https://api.pinference.ai/api/v1"] -api_key_var = "PRIME_API_KEY" - -[orchestrator.teacher.client.headers_from_env] -X-Prime-Team-ID = "PRIME_TEAM_ID" - -[trainer.optim] -lr = 3e-6 - -[ckpt] - -[inference] -gpu_memory_utilization = 0.5 diff --git a/configs/debug/training_modes/sft_lora.toml b/configs/debug/training_modes/sft_lora.toml index 687b45bbe36..b6ad2a8df60 100644 --- a/configs/debug/training_modes/sft_lora.toml +++ b/configs/debug/training_modes/sft_lora.toml @@ -20,6 +20,10 @@ training_mode = "sft" batch_size = 128 group_size = 4 +# Teacher rolls out over plain chat-completions (no tokens); the renderer backfills them. +[orchestrator.renderer] +name = "qwen3" + [orchestrator.train.sampling] max_completion_tokens = 128 diff --git a/configs/debug/v1/alphabet_sort.toml b/configs/debug/v1/alphabet_sort.toml new file mode 100644 index 00000000000..3429f57526f --- /dev/null +++ b/configs/debug/v1/alphabet_sort.toml @@ -0,0 +1,42 @@ +# v1 port of examples/alphabet_sort/rl.toml — identical except the env block, which loads +# the v1 `alphabet-sort-v1` taskset (multi-turn via a colocated vf.User) instead of the v0 +# `primeintellect/alphabet-sort` env. Harness runs on the subprocess runtime. + +max_steps = 200 +seq_len = 2048 + +[ckpt] # Checkpoint at the end of training + +[model] +name = "Qwen/Qwen3-4B-Instruct-2507" + +[wandb] +project = "alphabet-sort" +name = "alphabet-sort" + +[trainer.model] +impl = "auto" + +[trainer.model.ac] +freq = 1 + +[trainer.model.lora] +rank = 32 +alpha = 64 + +[trainer.optim] +lr = 1e-5 + +[orchestrator] +batch_size = 512 +group_size = 8 + +[orchestrator.train.sampling] +max_completion_tokens = 768 + +[[orchestrator.train.env]] +name = "alphabet-sort" +taskset = { id = "alphabet-sort-v1", min_turns = 3, max_turns = 5, power_per_turn = false } +harness = { id = "default", runtime = { type = "subprocess" } } + +[inference] # Default inference config diff --git a/configs/debug/v1/general_agent.toml b/configs/debug/v1/general_agent.toml new file mode 100644 index 00000000000..07779a40162 --- /dev/null +++ b/configs/debug/v1/general_agent.toml @@ -0,0 +1,55 @@ +# Debug v1 RL run: the v1 `general-agent-v1` taskset (GeneralAgentSolverTaskset) driven by the +# default harness, with the agent and its tool server colocated in a modal sandbox. Copied from +# configs/general_agent/rl_qwen3_4b.toml and rewired for v1 (taskset + harness + renderer; the +# v0 `id = "general-agent-solver-rlm"` env becomes the v1 taskset + a separate harness). + +max_steps = 200 +seq_len = 32768 + +[deployment] +num_train_gpus = 4 +num_infer_gpus = 4 + +[wandb] +project = "general-agent-debug" +name = "qwen3-4b-v1" + +[ckpt] +interval = 100 +keep_last = 1 + +[model] +name = "Qwen/Qwen3-4B-Instruct-2507" + +[trainer] + +[trainer.model] +cp = 2 + +[trainer.model.ac] +freq = 1 + +[trainer.model.compile] + +[orchestrator] +training_mode = "rl" +batch_size = 256 +rollouts_per_example = 8 +max_off_policy_steps = 32 + +[orchestrator.renderer] +name = "qwen3" + +[[orchestrator.train.env]] +# v1 general-agent solver taskset; the tool server is colocated in the harness's modal sandbox. +taskset = { id = "general-agent-v1", tools = { colocated = true } } +harness = { id = "default", runtime = { type = "modal" } } + +[inference] +gpu_memory_utilization = 0.85 + +[inference.model] +max_model_len = 32768 + +[inference.parallel] +dp = 4 diff --git a/configs/debug/v1/gsm8k.toml b/configs/debug/v1/gsm8k.toml new file mode 100644 index 00000000000..65dc9fe4deb --- /dev/null +++ b/configs/debug/v1/gsm8k.toml @@ -0,0 +1,42 @@ +# GSM8K (v1) — the v1 port of configs/gsm8k/rl.toml. Trains Qwen3-0.6B on GSM8K through the +# v1 env server: the orchestrator spawns a v1 EnvServer for the gsm8k-v1 taskset and trains on +# the returned Traces. Runs on 2 GPUs (1 inference, 1 training): +# uv run inference --model.name PrimeIntellect/Qwen3-0.6B # Inference pane (GPU 0) +# uv run rl @ configs/debug/v1/gsm8k.toml --ckpt # Trainer pane (GPU 1) + +max_steps = 100 +seq_len = 2048 + +[wandb] +project = "gsm8k-debug" +name = "gsm8k-v1" + +[model] +name = "Qwen/Qwen3-0.6B" + +[orchestrator] +training_mode = "rl" +batch_size = 512 +group_size = 16 + +[orchestrator.train.sampling] +max_completion_tokens = 2048 + +[[orchestrator.train.env]] +taskset = { id = "gsm8k-v1", split = "train" } +harness = { id = "default", runtime = { type = "subprocess" } } + +[orchestrator.eval] +interval = 20 +num_examples = 256 + +[orchestrator.eval.sampling] +max_completion_tokens = 2048 + +[[orchestrator.eval.env]] +taskset = { id = "gsm8k-v1", split = "test" } +harness = { id = "default", runtime = { type = "subprocess" } } + +[trainer] # Default trainer config + +[inference] # Default inference config \ No newline at end of file diff --git a/configs/debug/v1/hendrycks_sanity.toml b/configs/debug/v1/hendrycks_sanity.toml new file mode 100644 index 00000000000..9a3c97d452e --- /dev/null +++ b/configs/debug/v1/hendrycks_sanity.toml @@ -0,0 +1,55 @@ +# v1 analog of examples/hendrycks_sanity/rl.toml — identical config with only +# the env sections swapped to v1 taskset/harness syntax (math-env taskset, +# default harness, modal runtime). Submits to slurm by default; pass --no-slurm +# to run locally. + +output_dir = "/beegfs/mika/hendrycks-sanity-v1" +max_steps = 200 + +[wandb] +project = "hendrycks-sanity" +name = "v1" + +[deployment] +num_train_gpus = 4 +num_infer_gpus = 4 + +[slurm] +job_name = "hendrycks-modal" +partition = "all" + +[model] +name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" + +[orchestrator] +batch_size = 512 +group_size = 8 +seq_len = 8192 +max_inflight_rollouts = 256 + +[[orchestrator.train.env]] +name = "hendrycks-math" +taskset = { id = "math-env-v1", dataset_name = "mikasenghaas/Sanity-Test-R1D-1.5B", dataset_subset = "default" } +harness = { id = "default", runtime = { type = "modal" } } +timeout = { scoring = 300 } + +[orchestrator.eval] +interval = 50 + +[[orchestrator.eval.env]] +name = "aime2024" +taskset = { id = "aime24-v1" } +harness = { id = "default", runtime = { type = "modal" } } +timeout = { scoring = 300 } +group_size = 32 + +[trainer.model] +seq_len = 16384 + +[inference.model] +max_model_len = 8192 + +# Model not in MODEL_RENDERER_MAP — opt into DefaultRenderer (apply_chat_template). +[orchestrator.renderer] +name = "default" +reasoning_parser = "think" diff --git a/configs/debug/v1/r2e_gym.toml b/configs/debug/v1/r2e_gym.toml new file mode 100644 index 00000000000..2f45d8561fb --- /dev/null +++ b/configs/debug/v1/r2e_gym.toml @@ -0,0 +1,86 @@ +# v1 port of configs/rlm_swe/qwen35_4b.toml — identical except the env blocks, which load +# the v1 `r2e-gym-v1` taskset through the rlm harness on the prime runtime instead of the +# v0 `rlm_swe` composable env. Same model / trainer / inference / orchestrator knobs. + +output_dir = "/beegfs/mika/rlm-swe-qwen35-4b" +max_steps = 400 +seq_len = 65536 + +[slurm] +job_name = "rlm-swe-qwen35-4b" +project_dir = "." +pre_run_command = "prime sandbox delete --label rlm-swe-qwen35-4b -y --plain || true" + +[deployment] +type = "multi_node" +num_train_nodes = 1 +num_infer_nodes = 1 +num_infer_replicas = 2 + +[wandb] +project = "rlm-swe-debug" +name = "qwen35-4b" + +[weight_broadcast] +type = "nccl" + +[ckpt] +interval = 50 +keep_last = 1 +resume_step = -1 + +[model] +name = "Qwen/Qwen3.5-4B" + +# --- Trainer --- + +[trainer] + +[trainer.model] +cp = 4 +cp_style = "ulysses" + +[trainer.model.ac] +freq = 1 + +[trainer.model.compile] + +# --- Orchestrator --- + +[orchestrator] +batch_size = 256 +group_size = 8 +max_inflight_rollouts = 512 +max_off_policy_steps = 16 + +# Thinking enabled for the Qwen3.5 renderer. +[orchestrator.renderer] +name = "qwen3.5" +enable_thinking = true + +[orchestrator.train.sampling] +temperature = 1.0 + +[[orchestrator.train.env]] +name = "rlm-swe-r2e" +taskset = { id = "r2e-gym-v1" } +harness = { id = "rlm", runtime = { type = "prime", labels = ["rlm-swe-qwen35-4b"] } } + +[orchestrator.prime_monitor] + +# --- Inference --- + +[inference] +gpu_memory_utilization = 0.85 +enable_prefix_caching = true + +[inference.model] +max_model_len = 65536 + +[inference.parallel] +dp = 8 + +# Qwen3.5-4B is a VL model; skip the vision tower for text-only SWE. +# `language_model_only` is a vLLM MultiModalConfig arg (no prime-rl field) → pass via vllm_extra. +[inference.vllm_extra] +language_model_only = true diff --git a/configs/debug/v1/reverse_text.toml b/configs/debug/v1/reverse_text.toml new file mode 100644 index 00000000000..2bd761e5033 --- /dev/null +++ b/configs/debug/v1/reverse_text.toml @@ -0,0 +1,42 @@ +# Debug RL run on the v1 env server (reverse-text starter). +# The orchestrator spawns a v1 EnvServer per env (it never loads the env +# itself), dispatches rollouts by task index, and trains on the returned Traces +# (renderer-tokenized). Light settings for a quick end-to-end smoke. + +max_steps = 20 +seq_len = 2048 + +[wandb] +project = "reverse-text" +name = "v1" + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" + +[orchestrator] +training_mode = "rl" +batch_size = 128 +group_size = 16 + +[orchestrator.renderer] +name = "qwen3" + +[orchestrator.train.sampling] +max_completion_tokens = 128 + +[[orchestrator.train.env]] +taskset = { id = "reverse-text-v1" } +# reverse-text is a pure-text single-turn env: disable the bash tool (the model answers +# directly) and use the subprocess runtime (no docker). +harness = { id = "default", runtime = { type = "subprocess" } } + +# No eval block, mirroring examples/reverse_text/rl.toml — this is a train-only +# smoke. Add an [orchestrator.eval] block (with an interval) to exercise eval. + +[trainer.optim] +lr = 3e-6 + +[ckpt] + +[inference] +gpu_memory_utilization = 0.5 diff --git a/configs/debug/v1/scaleswe.toml b/configs/debug/v1/scaleswe.toml new file mode 100644 index 00000000000..df660d44ce6 --- /dev/null +++ b/configs/debug/v1/scaleswe.toml @@ -0,0 +1,118 @@ +# v1 RL run: Qwen3.6-35B-A3B (MoE) on the v1 `scaleswe-v1` SWE taskset, rlm harness on a prime +# sandbox runtime. Each scaleswe row ships its own Docker image (NEEDS_CONTAINER) which the +# Environment injects into the prime sandbox; the agent edits the repo + runs tests there (no MCP +# toolset). Router replay is enabled: trainer.enable_router_replay replays the routed-expert +# decisions returned by inference (enable_return_routed_experts). +# +# Built from configs/debug/v1/r2e_gym.toml (the v1 SWE structure) + the MoE parallelism from +# configs/general_agent/rl_qwen3_30b_a3b.toml. The deployment/parallelism is a starting point off +# the 30B-A3B reference — confirm/adjust for your cluster and the 35B / 65k-context footprint. + +max_steps = 400 +seq_len = 65536 + +[slurm] +job_name = "scaleswe-qwen36-35b-a3b-rlm" +project_dir = "." +partition = "all" +# Clean up this run's orphaned prime sandboxes before launch (label == SLURM job name). +pre_run_command = 'timeout 120 prime sandbox delete --label "$SLURM_JOB_NAME" -y || true' + +[deployment] +type = "multi_node" +num_train_nodes = 2 +num_infer_nodes = 2 + +[wandb] +project = "scaleswe-debug" +name = "qwen36-35b-a3b-rlm" + +[weight_broadcast] +type = "nccl" + +[ckpt] +interval = 50 +keep_last = 1 +resume_step = -1 + +[model] +name = "Qwen/Qwen3.6-35B-A3B" + +# --- Trainer --- + +[trainer] +# NCCL collective watchdog timeout — raised from default 600s: scaleswe's slow +# multi-node startup (step-0 swebench eval + slow first SWE batch) was tripping the +# 600s _ALLGATHER_BASE watchdog → SIGABRT trainer crash. +dist_timeout_seconds = 3600 +# Replay the routed-expert decisions returned by inference (enable_return_routed_experts) +# so the trainer's MoE routing matches the rollout's — cuts trainer/inference mismatch on A3B. +enable_router_replay = true + +[trainer.model] +cp = 2 +# Qwen3.6 is a hybrid/linear-attention model — ring CP is unsupported; ulysses +# all-to-all on Q/K/V works with the non-softmax kernels. +cp_style = "ulysses" +# Chunk the LM head over tokens to avoid materializing the full [seq, vocab] logits +# (65k seq × large vocab) — big activation-memory saver. 1024 matches the SWE configs. +fused_lm_head_token_chunk_size = 1024 + +[trainer.model.ac] +freq = 1 + +[trainer.model.compile] + +# --- Orchestrator --- + +[orchestrator] +batch_size = 256 +group_size = 16 +max_off_policy_steps = 16 + +[orchestrator.renderer] +name = "qwen3.6" +enable_thinking = true + +[orchestrator.train.sampling] +temperature = 1.0 + +[[orchestrator.train.env]] +name = "scaleswe" +# v1 scaleswe SWE taskset; rlm harness in a prime sandbox using the row's injected Docker image. +taskset = { id = "scaleswe-v1" } +harness = { id = "rlm", runtime = { type = "prime", labels = ["scaleswe-qwen36-35b-a3b-rlm"] } } + +[orchestrator.prime_monitor] + +# --- Eval: SWE-Bench Verified (first 100 tasks) every 20 steps, rlm harness on prime --- + +[orchestrator.eval] +interval = 20 + +[[orchestrator.eval.env]] +name = "swebench-verified" +taskset = { id = "swebench-verified-v1", use_prime_registry = true } +harness = { id = "rlm", runtime = { type = "prime", labels = ["scaleswe-qwen36-35b-a3b-rlm"] } } +num_examples = 100 +group_size = 1 +timeout = { rollout = 3600 } + +# --- Inference --- + +[inference] +gpu_memory_utilization = 0.85 +enable_prefix_caching = true +# Return per-token routed-expert decisions to the trainer for router replay +# (pairs with trainer.enable_router_replay). +enable_return_routed_experts = true + +[inference.model] +max_model_len = 65536 +# Qwen3.6 is not matched by the auto tool-parser resolver; set explicitly so vLLM +# enables auto tool choice (the scaleswe SWE agent sends tool_choice="auto"). +tool_call_parser = "qwen3_coder" + +[inference.parallel] +dp = 2 +tp = 4 diff --git a/configs/rlm_swe/qwen35_4b.toml b/configs/rlm_swe/qwen35_4b.toml index 626b62b8ef4..0ae908c7de4 100644 --- a/configs/rlm_swe/qwen35_4b.toml +++ b/configs/rlm_swe/qwen35_4b.toml @@ -60,7 +60,6 @@ temperature = 1.0 [[orchestrator.train.env]] id = "rlm_swe" name = "rlm-swe-r2e" -num_workers = 4 [orchestrator.train.env.args] labels = ["rlm-swe-qwen35-4b"] @@ -73,8 +72,7 @@ interval = 20 [[orchestrator.eval.env]] id = "rlm_swe" name = "rlm-swe-swebench-verified-quick" -num_workers = 4 -timeout = 3600 +timeout = { rollout = 3600 } [orchestrator.eval.env.args] task_type = "swebench" diff --git a/configs/wordle/rl.toml b/configs/wordle/rl.toml new file mode 100644 index 00000000000..826f3cd5198 --- /dev/null +++ b/configs/wordle/rl.toml @@ -0,0 +1,41 @@ +# Wordle (classic v0 TextArena env) on 2 GPUs (1 trainer + 1 inference), run via the +# legacy bridge. The Wordle-SFT model plays a full ~6-turn game per rollout; the bridge +# serves each multi-turn rollout to the orchestrator as a vf.Trace. + +max_steps = 200 +seq_len = 8192 + +[deployment] +num_train_gpus = 1 +num_infer_gpus = 1 + +[wandb] +project = "wordle" +name = "wordle" + +[ckpt] # Checkpoint at the end of training + +[model] +name = "PrimeIntellect/Qwen3-1.7B-Wordle-SFT" + +[orchestrator] +batch_size = 128 +group_size = 8 + +[[orchestrator.train.env]] +id = "primeintellect/wordle" +name = "wordle" + +[orchestrator.train.sampling] +max_completion_tokens = 1024 + +[trainer] # Default trainer config + +[inference] # single inference GPU (no data parallelism) + +# Qwen3 finetune with the standard PI template patch (byte-identical to +# PrimeIntellect/Qwen3-0.6B base); always re-emits prior blocks. +# Match that with the qwen3 renderer's preserve_all_thinking. +[orchestrator.renderer] +name = "qwen3" +preserve_all_thinking = true diff --git a/deps/pydantic-config b/deps/pydantic-config index 896ade4e69d..99f47c6cc04 160000 --- a/deps/pydantic-config +++ b/deps/pydantic-config @@ -1 +1 @@ -Subproject commit 896ade4e69d8d8dff2d4b0a431b7e1c7c12d638f +Subproject commit 99f47c6cc043e91c995fae225a495a92efe1126c diff --git a/deps/verifiers b/deps/verifiers index 78c1b3c545a..7c347f6110f 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 78c1b3c545a7df1f9cce99f72894883548c99212 +Subproject commit 7c347f6110f5759578acde7a7b6970678dec1d94 diff --git a/packages/prime-rl-configs/pyproject.toml b/packages/prime-rl-configs/pyproject.toml index 55ae0a5b47a..97110aa48eb 100644 --- a/packages/prime-rl-configs/pyproject.toml +++ b/packages/prime-rl-configs/pyproject.toml @@ -10,6 +10,7 @@ dependencies = [ "renderers>=0.1.8.dev28", "tomli>=2.2.1", "tomli-w>=1.2.0", + "verifiers>=0.1.15.dev357", ] [build-system] diff --git a/packages/prime-rl-configs/src/prime_rl/configs/env_server.py b/packages/prime-rl-configs/src/prime_rl/configs/env_server.py index 50c99adcff4..21f6264b1a6 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/env_server.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/env_server.py @@ -1,25 +1,14 @@ from pathlib import Path -from pydantic import model_validator - from prime_rl.configs.orchestrator import EnvConfig from prime_rl.configs.shared import LogConfig from prime_rl.utils.config import BaseConfig class EnvServerConfig(BaseConfig): - env: EnvConfig = EnvConfig() + env: EnvConfig log: LogConfig = LogConfig() - env_install_prerelease: bool = False - """Allow pre-release versions when installing environments. Passes ``--prerelease`` to ``prime env install``.""" - output_dir: Path = Path("outputs") """Directory to write outputs to — logs and any generated artifacts are written as subdirectories.""" - - @model_validator(mode="after") - def validate_num_workers(self): - if self.env.num_workers == "auto": - self.env.num_workers = 1 - return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index 65363862e71..1a960977e79 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -1,10 +1,9 @@ -import math import warnings from pathlib import Path from typing import Annotated, Any, Literal, TypeAlias -from pydantic import AliasChoices, Field, model_serializer, model_validator -from pydantic_core.core_schema import SerializerFunctionWrapHandler +import verifiers.v1 as vf +from pydantic import AliasChoices, Field, model_validator from renderers import AutoRendererConfig, RendererConfig from prime_rl.configs.shared import ( @@ -143,6 +142,69 @@ def _deprecate_max_tokens(cls, data: Any) -> Any: return data +class EnvConfig(vf.EnvServerConfig): + name: str | None = None + """Display name for this environment in logs, metrics, and buffer keys. Defaults to the taskset id. Must be unique across all envs in the same group.""" + + address: str | None = None + """ZMQ address of an external env server (e.g. ``tcp://host:5000``). When set, the orchestrator connects to this server instead of spawning one; when None, a subprocess env server is spawned automatically. The ``pool`` sizes the spawned server.""" + + ratio: float | None = Field(None, gt=0) + """Sampling weight for this environment in the buffer. When None for all envs, samples uniformly across all available problems. When set, must be set on all envs — values are relative weights normalized to probabilities (e.g. [1, 1] and [0.5, 0.5] are equivalent).""" + + max_retries: int = Field(3, ge=0) + """Times the env server retries a failed rollout before returning an error.""" + + @model_validator(mode="before") + @classmethod + def _migrate_num_workers(cls, data): + """Back-compat: the removed ``num_workers`` maps onto ``pool`` — an int becomes a + fixed ``static`` pool, ``"auto"`` falls through to the default ``elastic`` pool. An + explicit ``pool`` always wins.""" + if isinstance(data, dict) and "num_workers" in data: + num_workers = data.pop("num_workers") + if "pool" not in data and num_workers != "auto": + data["pool"] = {"type": "static", "num_workers": num_workers} + return data + + @property + def is_legacy(self) -> bool: + """A v0/legacy env (run via the bridge): an ``id`` is set and no v1 ``taskset`` is.""" + return not self.taskset.id + + @property + def env_id(self) -> str: + """The env identifier — the v1 taskset id (v1) or the legacy env id (v0).""" + return self.taskset.id or self.id or "" + + @property + def resolved_name(self) -> str: + return self.name or self.env_id + + @model_validator(mode="after") + def validate_env(self): + if not self.taskset.id and not self.id: + raise ValueError('no env configured — set taskset = { id = "" } (v1) or id = "" (v0/legacy)') + if self.resolved_name == "all": + raise ValueError( + 'Environment name "all" is reserved for global metric aggregation. Use a different name or id.' + ) + return self + + @model_validator(mode="after") + def resolve_legacy_env_kwargs(self): + """For a v0/legacy env, surface the v1 knobs the legacy bridge applies via + ``extra_env_kwargs`` (``env.set_kwargs(...)``): the per-rollout wall-clock timeout and + the multi-turn completion-token budget. (``max_seq_len`` is added per train run in + ``OrchestratorConfig.resolve_env_config``, which knows ``seq_len``.)""" + if self.is_legacy: + if self.timeout.rollout is not None: + self.extra_env_kwargs["timeout_seconds"] = self.timeout.rollout + if self.max_output_tokens is not None: + self.extra_env_kwargs["max_total_completion_tokens"] = self.max_output_tokens + return self + + class LinearLengthPenaltyConfig(BaseConfig): coef: float = Field(0.25, ge=0, allow_inf_nan=False) """Scale on the linear length penalty. Each reward is reduced by ``coef * pass_rate * (model completion tokens / orchestrator.seq_len)`` — where ``pass_rate`` is the group's mean reward — before the GRPO baseline subtraction. Finite and non-negative.""" @@ -177,69 +239,6 @@ class CustomAdvantageConfig(BaseConfig): ] -class EnvConfig(BaseConfig): - id: str = "reverse-text" - """Registered verifiers environment ID (e.g. ``math-env``, ``primeintellect/math-env``). May include an ``@version`` suffix for installation.""" - - name: str | None = None - """Display name for this environment in logs, metrics, and buffer keys. Defaults to the ``id`` without ``@version``. Must be unique across all envs in the same group.""" - - args: dict = {} - """Keyword arguments forwarded to ``vf.load_environment``. See the environment's docstring for accepted args.""" - - extra_env_kwargs: dict[str, Any] = {} - """Extra kwargs passed to the env (e.g. ``seq_len``, ``max_total_completion_tokens``). Auto-populated by the orchestrator; user overrides are generally discouraged. The main use case is matching ``extra_env_kwargs`` when running an env in an isolated environment server.""" - - address: str | None = None - """ZMQ address of an external env server (e.g. ``tcp://host:5000``). When set, the orchestrator connects to this server instead of spawning one; when None, a subprocess env server is spawned automatically.""" - - num_workers: int | Literal["auto"] = "auto" - """Worker processes for the spawned env server. ``auto`` scales to 1 worker per 256 concurrent rollouts. Ignored when ``address`` is set.""" - - ratio: float | None = Field(None, gt=0) - """Sampling weight for this environment in the buffer. When None for all envs, samples uniformly across all available problems. When set, must be set on all envs — values are relative weights normalized to probabilities (e.g. [1, 1] and [0.5, 0.5] are equivalent).""" - - max_retries: int = Field(3, ge=0) - """Times the env server retries a failed rollout before returning an error.""" - - max_total_completion_tokens: int = -1 - """Maximum total completion tokens across all turns in a multi-turn rollout. ``-1`` disables. Auto-populated into ``extra_env_kwargs``.""" - - timeout: float | None = Field(None, validation_alias=AliasChoices("timeout", "timeout_seconds")) - """Per-rollout wall-clock timeout in seconds. None disables.""" - - state_columns: list[str] = [] - """Extra ``State`` fields to persist into the saved rollout records (in addition to the always-saved ``trajectory`` and ``sampling_args``). Values must be JSON-serializable.""" - - @property - def stripped_id(self) -> str: - """Environment ID without the @version suffix.""" - return self.id.split("@")[0] - - @property - def resolved_name(self) -> str: - return self.name or self.stripped_id - - @model_validator(mode="after") - def validate_env_name(self): - if self.resolved_name == "all": - raise ValueError( - 'Environment name "all" is reserved for global metric aggregation. Use a different name or id.' - ) - return self - - @model_validator(mode="after") - def resolve_max_total_completion_tokens(self): - self.extra_env_kwargs["max_total_completion_tokens"] = self.max_total_completion_tokens - return self - - @model_validator(mode="after") - def resolve_timeout(self): - if self.timeout is not None: - self.extra_env_kwargs["timeout_seconds"] = self.timeout - return self - - class TrainEnvConfig(EnvConfig): sampling: TrainSamplingConfig = TrainSamplingConfig() """Per-env sampling overrides. Unset fields inherit from the group-level train sampling config.""" @@ -269,21 +268,19 @@ class EvalEnvConfig(EnvConfig): class TrainConfig(BaseConfig): - env: list[TrainEnvConfig] = [TrainEnvConfig()] + env: list[TrainEnvConfig] = Field(default_factory=list) """Training environments.""" sampling: TrainSamplingConfig = TrainSamplingConfig() """Shared training sampling configuration.""" - num_workers: int | Literal["auto"] = "auto" - """Default worker processes for env servers. Can be overridden per env.""" - max_retries: int = Field(3, ge=0) """Default retries for failed rollouts. Can be overridden per env.""" @model_validator(mode="after") def resolve_env_defaults(self): - """Resolve per-env overrides: inherit group-level sampling, num_workers, and max_retries.""" + """Resolve per-env overrides: inherit group-level sampling and max_retries (the + worker ``pool`` is configured per env, defaulting to elastic).""" group_sampling = self.sampling.model_dump() for env in self.env: if "sampling" not in env.model_fields_set: @@ -291,8 +288,6 @@ def resolve_env_defaults(self): else: merged = group_sampling | env.sampling.model_dump(exclude_unset=True) env.sampling = TrainSamplingConfig(**merged) - if "num_workers" not in env.model_fields_set: - env.num_workers = self.num_workers if "max_retries" not in env.model_fields_set: env.max_retries = self.max_retries return self @@ -318,7 +313,7 @@ def validate_env_ratios(self): class EvalConfig(BaseConfig): - env: list[EvalEnvConfig] = [EvalEnvConfig()] + env: list[EvalEnvConfig] = Field(default_factory=list) """Evaluation environments.""" sampling: EvalSamplingConfig = Field(default_factory=EvalSamplingConfig) @@ -330,9 +325,6 @@ class EvalConfig(BaseConfig): group_size: int = Field(1, ge=1, validation_alias=AliasChoices("group_size", "rollouts_per_example")) """Default rollouts per example. Can be overridden per env.""" - num_workers: int | Literal["auto"] = "auto" - """Default worker processes for env servers. Can be overridden per env.""" - max_retries: int = Field(3, ge=0) """Default retries for failed rollouts. Can be overridden per env.""" @@ -345,7 +337,8 @@ class EvalConfig(BaseConfig): @model_validator(mode="after") def resolve_env_defaults(self): - """Resolve per-env overrides: inherit group-level sampling, num_workers, max_retries, num_examples, group_size, and interval. Then resolve auto num_workers.""" + """Resolve per-env overrides: inherit group-level sampling, max_retries, num_examples, + group_size, and interval (the worker ``pool`` is configured per env, default elastic).""" group_sampling = self.sampling.model_dump() for env in self.env: if "sampling" not in env.model_fields_set: @@ -359,17 +352,8 @@ def resolve_env_defaults(self): env.group_size = self.group_size if "interval" not in env.model_fields_set: env.interval = self.interval - if "num_workers" not in env.model_fields_set: - env.num_workers = self.num_workers if "max_retries" not in env.model_fields_set: env.max_retries = self.max_retries - # Resolve auto num_workers now that num_examples and group_size are set - if env.num_workers == "auto": - if env.num_examples == -1: - env.num_workers = 4 - else: - max_concurrent = env.num_examples * env.group_size - env.num_workers = max(1, math.ceil(max_concurrent / 256)) return self @model_validator(mode="after") @@ -509,30 +493,16 @@ class OrchestratorConfig(BaseConfig): tokenizer: TokenizerConfig = TokenizerConfig() - renderer: RendererConfig | None = AutoRendererConfig() - """Typed renderer config (``renderers.RendererConfig`` discriminated - union). Defaults to ``"auto"``, which resolves from - ``tokenizer.name_or_path`` via ``MODEL_RENDERER_MAP``. ``None`` - opts into MITO (``openai_chat_completions``).""" + renderer: RendererConfig = AutoRendererConfig() + """Typed renderer config (``renderers.RendererConfig`` discriminated union), required — + training is renderer-only. Defaults to ``"auto"``, which resolves from + ``tokenizer.name_or_path`` via ``MODEL_RENDERER_MAP``. RL/OPD roll out through the renderer + client; SFT uses it to backfill tokens for its chat-completions teacher.""" pool_size: int | None = Field(None, ge=1) """Number of renderer slots shared across concurrent rollouts. Bump for long multi-turn prompts where client-side jinja tokenization - serializes. Only meaningful when ``renderer`` is not ``None``.""" - - @model_serializer(mode="wrap") - def _preserve_mito_renderer(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]: - """Emit ``renderer = "None"`` (string) when MITO so - ``model_dump(exclude_none=True)`` round-trips: dumped TOML has - ``renderer = "None"``, and on reload - ``BaseConfig._none_str_to_none`` coerces it back to ``None``. - Without this, a MITO orchestrator config saved to - ``control/orch.toml`` would lose the renderer key entirely and - reload as the default ``AutoRendererConfig()`` (TITO).""" - result = handler(self) - if self.renderer is None: - result["renderer"] = "None" - return result + serializes.""" optim: OptimizerConfig = OptimizerConfig() """Per-run optimizer configuration for multi-run training.""" @@ -621,9 +591,6 @@ def _preserve_mito_renderer(self, handler: SerializerFunctionWrapHandler) -> dic heartbeat: HeartbeatConfig | None = None """BetterStack heartbeat configuration for monitoring training progress.""" - env_install_prerelease: bool = False - """Allow pre-release versions when installing environments (e.g. ``verifiers>=0.1.12.dev5``). Passes ``--prerelease`` to ``prime env install``.""" - experimental: OrchestratorExperimentalConfig = OrchestratorExperimentalConfig() @model_validator(mode="before") @@ -748,51 +715,14 @@ def validate_unique_filter_types(self): ) return self - @model_validator(mode="after") - def _force_no_renderer_for_sft(self): - """Teacher-backed SFT rolls out via the teacher's plain chat-completions - endpoint; the renderer client doesn't apply. When no teacher is - configured, SFT uses the student rollout path and keeps the renderer.""" - if self.training_mode == "sft" and self.teacher is not None: - self.renderer = None - return self - @model_validator(mode="after") def validate_training_mode(self): """Enforce training mode invariants that involve only orchestrator fields.""" has_teacher = self.teacher is not None if self.training_mode == "rl" and has_teacher: raise ValueError("orchestrator.teacher must not be set when training_mode = 'rl'.") - if self.training_mode == "opd" and not has_teacher: - raise ValueError("orchestrator.teacher must be configured when training_mode = 'opd'.") - return self - - @model_validator(mode="after") - def validate_pool_size(self): - """``pool_size`` is only meaningful when the renderer is enabled - (``renderer is not None``). Reject otherwise so callers don't - silently pass it and wonder why it's ignored.""" - if self.renderer is None and self.pool_size is not None: - raise ValueError( - f"orchestrator.pool_size={self.pool_size!r} is set but " - "orchestrator.renderer is None (MITO mode). Either configure a renderer " - "or remove pool_size." - ) - return self - - @model_validator(mode="after") - def vlm_requires_renderer(self): - """VLMs (``[model.vlm]`` block set) must go through the renderer. - - The renderer owns the processor per-slot, produces byte-identical - tokens, and ships generic ``mm_kwargs`` keyed by whatever the - model's forward signature expects. - """ - if self.student.model.vlm is not None and self.renderer is None: - raise ValueError( - "orchestrator.renderer must be set when model.vlm is set. " - "VLMs must go through a renderer (e.g. Qwen3VLRenderer) that owns the processor." - ) + if self.training_mode in ("opd", "sft") and not has_teacher: + raise ValueError(f"orchestrator.teacher must be configured when training_mode = '{self.training_mode}'.") return self @model_validator(mode="after") @@ -808,7 +738,7 @@ def validate_renderer_auto_resolves(self): ``DefaultRendererConfig.tool_parser`` is configured. Surface at config time so ``--dry-run`` reports the error. """ - if self.renderer is None or self.renderer.name != "auto": + if self.renderer.name != "auto": return self from renderers.base import MODEL_RENDERER_MAP @@ -826,9 +756,7 @@ def validate_renderer_auto_resolves(self): f"(b) [orchestrator.renderer] name= — " f"if {model_id!r} is template-identical to a mapped family " f"(and ideally also add it upstream to " - f"renderers.base.MODEL_RENDERER_MAP). " - f"(c) orchestrator.renderer='none' — opt out of the renderer " - f"client entirely (MITO)." + f"renderers.base.MODEL_RENDERER_MAP)." ) @model_validator(mode="after") @@ -876,12 +804,6 @@ def resolve_batching(self): if "advantage" not in env_cfg.model_fields_set: env_cfg.advantage = self.advantage - # Resolve train env num_workers from max_inflight_rollouts - for env_cfg in self.train.env: - if env_cfg.num_workers == "auto": - assert self.max_inflight_rollouts is not None - env_cfg.num_workers = max(1, math.ceil(self.max_inflight_rollouts / 256)) - return self @model_validator(mode="after") @@ -900,12 +822,15 @@ def auto_setup_bench(self): @model_validator(mode="after") def resolve_env_config(self): - """Populate extra_env_kwargs and vLLM sampling defaults from top-level fields.""" - is_vllm = self.training_mode != "sft" + """Set vLLM sampling defaults on each train env from top-level fields.""" + if self.training_mode == "sft": + return self for env in self.train.env: - env.extra_env_kwargs.update(max_seq_len=self.seq_len) - if is_vllm: - env.sampling.extra_body.setdefault("top_k", -1) - env.sampling.extra_body.setdefault("min_p", 0.0) - env.sampling.extra_body.setdefault("return_token_ids", True) + env.sampling.extra_body.setdefault("top_k", -1) + env.sampling.extra_body.setdefault("min_p", 0.0) + env.sampling.extra_body.setdefault("return_token_ids", True) + if env.is_legacy: + # v0 env: cap per-turn response tokens to the training budget (the legacy + # bridge applies extra_env_kwargs via env.set_kwargs). + env.extra_env_kwargs["max_seq_len"] = self.seq_len return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index ff311f145df..5cbf598ec7a 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -100,12 +100,6 @@ class ElasticConfig(BaseConfig): class ClientConfig(BaseConfig): - timeout: int = 1200 - """Request timeout in seconds.""" - - connect_timeout: float = 30.0 - """TCP connect timeout in seconds for inference API requests.""" - wait_for_ready_timeout: int = 1800 """Seconds to wait at startup for the inference pool to become ready. Applies to both the static health check and elastic DNS-based discovery.""" diff --git a/pyproject.toml b/pyproject.toml index c6ed9f84e11..f6422ab5c54 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,13 +33,16 @@ dependencies = [ "setproctitle>=1.3.0", "uvloop>=0.21.0", "torchtitan", - "verifiers", + "verifiers>=0.1.15.dev357", "renderers", "dion", "tilelang>=0.1.8", "flash-linear-attention", "nvidia-ml-py>=12.575.51", "pybase64>=1.4.2", + "orjson>=3.10", + "pandas>=2.0", + "msgspec>=0.18", ] [project.scripts] @@ -94,6 +97,24 @@ envs = [ "tau2-bench", "wiki-search", "wordle", + "reverse-text-v1", + "gsm8k-v1", + "math-env-v1", + "aime24-v1", + "alphabet-sort-v1", + "scaleswe-v1", + "color-codeword-v1", + "r2e-gym-v1", + "terminal-bench-2-v1", + "code-golf-v1", + "deepwiki-v1", + "glossary-v1", + "swelego-v1", + "wiki-search-v1", + "wikispeedia-v1", + "wordle-v1", + "general-agent-v1", + "swebench-verified-v1", ] disagg = [ "deep-ep ; platform_machine == 'x86_64'", @@ -153,7 +174,7 @@ override-dependencies = [ "transformers==5.6.2", "torch>=2.9.0", "openenv-core", - "verifiers[packages]>=0.1.15.dev150", + "verifiers>=0.1.15.dev357", ] # ModelExpress 0.3.0 publishes protobuf<6 metadata, but its generated proto is @@ -188,8 +209,6 @@ prime-tunnel = false prime-evals = false verifiers = false renderers = false -tasksets = false -harnesses = false prime-pydantic-config = false # Trusted git/URL sources pinned by rev or wheel URL (defensive: belt-and-suspenders) vllm-router = false @@ -200,6 +219,10 @@ nixl-cu12 = false [tool.uv.sources] prime-rl-configs = { path = "packages/prime-rl-configs", editable = true } +# prime-rl depends on the one `verifiers` package (v0 + v1 + the legacy bridge + the built-in v1 +# harnesses/tasksets, now bundled in verifiers.v1). v0 envs resolve from +# deps/{verifiers,research-environments}/environments; v1 tasksets (`-v1`) and the `compact` +# example live under deps/verifiers/environments. verifiers = { path = "deps/verifiers", editable = true } renderers = { path = "deps/renderers", editable = true } prime-pydantic-config = { path = "deps/pydantic-config", editable = true } @@ -211,7 +234,6 @@ color-codeword = { path = "deps/research-environments/environments/color_codewor deepdive = { path = "deps/research-environments/environments/deepdive", editable = true } general-agent = { path = "deps/research-environments/environments/general_agent", editable = true } gpqa = { path = "deps/research-environments/environments/gpqa", editable = true } -harnesses = { path = "deps/verifiers/packages/harnesses", editable = true } hle = { path = "deps/research-environments/environments/hle", editable = true } ifeval = { path = "deps/research-environments/environments/ifeval", editable = true } livecodebench = { path = "deps/research-environments/environments/livecodebench", editable = true } @@ -233,9 +255,26 @@ rlm-uuid-ctf = { path = "deps/research-environments/environments/rlm_uuid_ctf", science-env = { path = "deps/research-environments/environments/science_env", editable = true } simpleqa-verified = { path = "deps/research-environments/environments/simpleqa_verified", editable = true } tau2-bench = { path = "deps/research-environments/environments/tau2_bench", editable = true } -tasksets = { path = "deps/verifiers/packages/tasksets", editable = true } wiki-search = { path = "deps/verifiers/environments/wiki_search", editable = true } wordle = { path = "deps/verifiers/environments/wordle", editable = true } +reverse-text-v1 = { path = "deps/verifiers/environments/reverse_text_v1", editable = true } +gsm8k-v1 = { path = "deps/verifiers/environments/gsm8k_v1", editable = true } +math-env-v1 = { path = "deps/verifiers/environments/math_env_v1", editable = true } +aime24-v1 = { path = "deps/verifiers/environments/aime24_v1", editable = true } +alphabet-sort-v1 = { path = "deps/verifiers/environments/alphabet_sort_v1", editable = true } +scaleswe-v1 = { path = "deps/verifiers/environments/scaleswe_v1", editable = true } +color-codeword-v1 = { path = "deps/verifiers/environments/color_codeword_v1", editable = true } +r2e-gym-v1 = { path = "deps/verifiers/environments/r2e_gym_v1", editable = true } +terminal-bench-2-v1 = { path = "deps/verifiers/environments/terminal_bench_2_v1", editable = true } +code-golf-v1 = { path = "deps/verifiers/environments/code_golf_v1", editable = true } +deepwiki-v1 = { path = "deps/verifiers/environments/deepwiki_v1", editable = true } +glossary-v1 = { path = "deps/verifiers/environments/glossary_v1", editable = true } +swelego-v1 = { path = "deps/verifiers/environments/swelego_v1", editable = true } +wiki-search-v1 = { path = "deps/verifiers/environments/wiki_search_v1", editable = true } +wikispeedia-v1 = { path = "deps/verifiers/environments/wikispeedia_v1", editable = true } +wordle-v1 = { path = "deps/verifiers/environments/wordle_v1", editable = true } +general-agent-v1 = { path = "deps/verifiers/environments/general_agent_v1", editable = true } +swebench-verified-v1 = { path = "deps/verifiers/environments/swebench_verified_v1", editable = true } torch = { index = "pytorch-cu128" } torchvision = { index = "pytorch-cu128" } torchaudio = { index = "pytorch-cu128" } diff --git a/scripts/tmux.sh b/scripts/tmux.sh index 3511d774584..12420f28eec 100755 --- a/scripts/tmux.sh +++ b/scripts/tmux.sh @@ -81,7 +81,7 @@ tmux send-keys -t "$SESSION_NAME:Logs.1" \ "tail -F ${LOG_DIR}/orchestrator.log 2>/dev/null" C-m tmux send-keys -t "$SESSION_NAME:Logs.2" \ - "tail -F ${LOG_DIR}/envs/*/*/*.log 2>/dev/null" C-m + "tail -F ${LOG_DIR}/envs/*/*.log 2>/dev/null" C-m tmux send-keys -t "$SESSION_NAME:Logs.3" \ "tail -F ${LOG_DIR}/inference.log 2>/dev/null" C-m @@ -95,8 +95,8 @@ tmux send-keys -t "$SESSION_NAME:Claude" \ All ranks: ${LOG_DIR}/trainer/torchrun/*/*/*/*.log Orchestrator: ${LOG_DIR}/orchestrator.log Inference: ${LOG_DIR}/inference.log - Envs: ${LOG_DIR}/envs/*/*/*.log - Train envs: ${LOG_DIR}/envs/train/*/*.log + Envs: ${LOG_DIR}/envs/*/*.log + Train envs: ${LOG_DIR}/envs/train/*.log You are running inside tmux session \"${SESSION_NAME}\". The Launcher window (window 0) is where the user runs launch commands. You can read its contents with: tmux capture-pane -t ${SESSION_NAME}:Launcher -p Help the user monitor and debug this run.'" C-m diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index 67e00b39d49..c0bcc346bc7 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -71,9 +71,7 @@ After a restart, verify all processes are back up and progress resumed before th ├── inference/ │ ├── node_*.log # per-node (multi-node only) │ └── router_0.log # vllm-router per replica (multi-node only) -└── envs/{train,eval}/{env_name}/ - ├── env_server.log - └── env_worker_*.log +└── envs/{train,eval}/{env_name}.log # one log file per env ``` Usually tailing `trainer.log`, `orchestrator.log`, and `inference.log` is enough. Drop into per-node or per-rank logs only when debugging. All logs are loguru with `HH:mm:ss LEVEL message`; levels: `DEBUG`, `INFO`, `SUCCESS`, `WARNING`, `ERROR`. @@ -82,7 +80,7 @@ Scan for problems: ```bash grep -E "WARNING|ERROR" {output_dir}/logs/{trainer,orchestrator,inference}.log -grep -E "WARNING|ERROR" {output_dir}/logs/envs/train/*/env_{server,worker_*}.log +grep -E "WARNING|ERROR" {output_dir}/logs/envs/{train,eval}/*.log ``` ### Metrics diff --git a/src/prime_rl/orchestrator/advantage.py b/src/prime_rl/orchestrator/advantage.py index f101d4037c3..29f7d1864d0 100644 --- a/src/prime_rl/orchestrator/advantage.py +++ b/src/prime_rl/orchestrator/advantage.py @@ -4,17 +4,17 @@ from typing import TYPE_CHECKING, Callable import torch -import verifiers as vf if TYPE_CHECKING: - from prime_rl.orchestrator.types import TrainRollout + import verifiers.v1 as vf + + from prime_rl.orchestrator.types import Rollout from prime_rl.configs.orchestrator import ( AdvantageConfig, CustomAdvantageConfig, LinearLengthPenaltyConfig, ) -from prime_rl.orchestrator.utils import get_model_completion_len from prime_rl.utils.utils import import_object @@ -22,7 +22,7 @@ class AdvantageInputs: """Inputs for advantage computation of a single group (one example × N rollouts).""" - rollouts: list[vf.RolloutOutput] + rollouts: list[vf.Trace] @dataclass @@ -57,8 +57,8 @@ def default_advantage_fn( (``reward == 1``) rollouts. ``length_weighted_baseline`` uses the token-length-weighted mean reward as the baseline instead of the plain mean. """ - rewards = torch.tensor([r["reward"] for r in inputs.rollouts], dtype=torch.float32) - lengths = torch.tensor([get_model_completion_len(r) for r in inputs.rollouts], dtype=rewards.dtype) + rewards = torch.tensor([r.reward for r in inputs.rollouts], dtype=torch.float32) + lengths = torch.tensor([r.completion_len for r in inputs.rollouts], dtype=rewards.dtype) if length_penalty is not None: if max_seq_len is None: @@ -98,19 +98,18 @@ def advantage_fn(inputs: AdvantageInputs) -> AdvantageOutputs: def assign_advantages( - rollouts: list["TrainRollout"], # noqa: F821 (forward ref) + rollouts: list[Rollout], advantage_fn: AdvantageFn | None, ) -> None: """Compute and assign advantages for one finished group of rollouts (``TrainSink.process_group`` hands in a single group's surviving rollouts). ``advantage_fn=None`` is the trivial case (advantage = reward); a custom - ``advantage_fn`` receives the raw ``vf.RolloutOutput``\\ s via - ``AdvantageInputs.rollouts``. + ``advantage_fn`` receives the ``vf.Trace``\\ s via ``AdvantageInputs.rollouts``. """ if advantage_fn is None: for rollout in rollouts: rollout.advantage = rollout.reward return - result = advantage_fn(AdvantageInputs(rollouts=[r.raw for r in rollouts])) + result = advantage_fn(AdvantageInputs(rollouts=[r for r in rollouts])) for rollout, advantage in zip(rollouts, result.advantages): rollout.advantage = advantage diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 8842807a356..eb08d35d5fb 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -6,7 +6,7 @@ - Emit-everything invariant: every dispatched rollout eventually reaches ``out_q`` exactly once as a ``TrainRollout`` / ``EvalRollout``. Failures (env error, empty trajectory, task exception, off-policy cancel) carry - ``raw["error"]`` set; sinks decide drop / partial-train policy. + ``trace.error`` set; sinks decide drop / partial-train policy. - ``DispatcherMode.PREFER_TRAIN`` / ``PREFER_EVAL`` controls which kind to schedule next. Transitions are level-triggered (driven by the eval source's emptiness), so in-flight rollouts of the opposite kind drain @@ -29,20 +29,18 @@ from enum import Enum, auto from typing import Literal -import verifiers as vf +import verifiers.v1 as vf from aiolimiter import AsyncLimiter from prime_rl.orchestrator.envs import EvalEnvs, TrainEnvs from prime_rl.orchestrator.eval_source import EvalSource from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( - EvalRollout, - FinishedRollout, GroupState, InflightRollout, Policy, + Rollout, RolloutKind, - TrainRollout, ) from prime_rl.utils.async_utils import safe_cancel, safe_cancel_all from prime_rl.utils.client import InferencePool, client_identity @@ -159,7 +157,7 @@ def __init__( self.groups: dict[uuid.UUID, GroupState] = {} # Bounded so the dispatcher backpressures on a slow sink - self.out_q: asyncio.Queue[FinishedRollout] = asyncio.Queue(maxsize=max(8, self.max_inflight)) + self.out_q: asyncio.Queue[Rollout] = asyncio.Queue(maxsize=max(8, self.max_inflight)) self.mode: DispatcherMode = DispatcherMode.PREFER_TRAIN # Set by the orchestrator after the final train step; pipeline then @@ -382,7 +380,7 @@ def next_fresh_group(self, kind: RolloutKind, envs) -> GroupState | None: return GroupState( kind=kind, env_name=env_name, - example=example, + task_idx=example["task_idx"], rollouts_to_schedule=group_size, target_rollouts=group_size, eval_step=eval_step, @@ -436,7 +434,7 @@ async def schedule_group_rollout(self, group_id: uuid.UUID, group: GroupState) - task: asyncio.Task = asyncio.create_task( env.run_group( client=client, - example=group.example, + task_idx=group.task_idx, model_name=model_name, group_size=permits, cache_salt=cache_salt, @@ -449,7 +447,7 @@ async def schedule_group_rollout(self, group_id: uuid.UUID, group: GroupState) - task = asyncio.create_task( env.run_rollout( client=client, - example=group.example, + task_idx=group.task_idx, model_name=model_name, cache_salt=cache_salt, ) @@ -493,97 +491,60 @@ async def handle_completed_rollout(self, task: asyncio.Task) -> None: is_synth_exception = False try: result = task.result() - rollouts: list[vf.RolloutOutput] = result if isinstance(result, list) else [result] + rollouts: list[Rollout] = result if isinstance(result, list) else [result] except asyncio.CancelledError: return except Exception as exc: get_logger().warning(f"Rollout task failed in group {meta.group_id} ({meta.env_name}): {exc!r}") - rollouts = [ - self.error_rollout_output(error_type=type(exc).__name__, error_repr=repr(exc)) - for _ in range(meta.rollout_count) - ] + task_idx = group.task_idx if group is not None else -1 + rollouts = [Rollout(task=vf.Task(idx=task_idx, prompt=None)) for _ in range(meta.rollout_count)] + for r in rollouts: + r.capture_error(exc) is_synth_exception = True for r in rollouts: - if r.get("error") is None and len(r.get("trajectory") or []) == 0: + if not r.has_error and r.num_turns == 0: # Empty trajectory: promote to an explicit error so the sink # treats it like any other failure - r["error"] = { - "error": "EmptyTrajectory", - "error_chain_repr": "Rollout returned with no trajectory steps", - "error_chain_str": "", - } + r.errors.append(vf.Error(type="EmptyTrajectory", message="Rollout returned with no trajectory steps")) get_logger().warning(f"Empty trajectory in group {meta.group_id} ({meta.env_name})") - if r.get("error") is not None: - err_type = r["error"].get("error", "Unknown") + if r.has_error: self.metrics.record_error(kind=meta.kind, env_name=meta.env_name) if not is_synth_exception: get_logger().warning( - f"Rollout failed in group {meta.group_id} ({meta.env_name}) — " - f"{r['error'].get('error_chain_repr', err_type)}" + f"Rollout failed in group {meta.group_id} ({meta.env_name}) — {r.error.type}: {r.error.message}" ) await self.emit_rollout(meta, group, r) - async def emit_rollout(self, meta: InflightRollout, group: GroupState | None, raw: vf.RolloutOutput) -> None: - """Build a ``TrainRollout`` / ``EvalRollout`` and put it on ``out_q``. + async def emit_rollout(self, meta: InflightRollout, group: GroupState | None, rollout: Rollout) -> None: + """Stamp prime-rl metadata onto the completed rollout and put it on ``out_q``. Pops the group from ``self.groups`` once every member has been emitted.""" eval_step = meta.eval_step policy_version = meta.policy_version - example_id = raw.get("example_id") if group is not None: eval_step = group.eval_step policy_version = group.policy_version_at_start - example_id = group.example["example_id"] group.emitted += 1 if group.emitted >= group.target_rollouts: self.groups.pop(meta.group_id, None) - common = dict( - raw=raw, - env_name=meta.env_name, - example_id=example_id if example_id is not None else -1, - group_id=meta.group_id, - policy_version=policy_version, - off_policy_steps=meta.off_policy_steps, - ) - rollout: FinishedRollout - if meta.kind == "train": - rollout = TrainRollout(**common) - else: + rollout.kind = meta.kind + rollout.env_name = meta.env_name + rollout.group_id = meta.group_id + rollout.policy_version = policy_version + rollout.off_policy_steps = meta.off_policy_steps + if meta.kind == "eval": assert eval_step is not None, "eval rollout missing eval_step" - rollout = EvalRollout(**common, eval_step=eval_step) + rollout.eval_step = eval_step await self.out_q.put(rollout) - @staticmethod - def error_rollout_output(*, error_type: str, error_repr: str) -> vf.RolloutOutput: - """Minimal ``vf.RolloutOutput`` for rollouts that never produced - real output (task exception, off-policy cancel).""" - out: vf.RolloutOutput = vf.RolloutOutput() - out["error"] = { - "error": error_type, - "error_chain_repr": error_repr, - "error_chain_str": error_repr, - } - out["trajectory"] = [] - out["completion"] = None - out["reward"] = 0.0 - out["is_truncated"] = False - out["metrics"] = {} - out["stop_condition"] = None - out["token_usage"] = { - "input_tokens": 0.0, - "output_tokens": 0.0, - "final_input_tokens": 0.0, - "final_output_tokens": 0.0, - } - return out - async def drop_group(self, group_id: uuid.UUID) -> int: """Cancel remaining in-flight tasks for this group and emit a ``Cancelled`` marker for every rollout it still owes the sink (both in-flight and not-yet-scheduled). Returns the count for off-policy metrics.""" group = self.groups.pop(group_id, None) + task_idx = group.task_idx if group is not None else -1 # Sync claim phase: pop matching tasks from ``self.inflight`` and # release their permits in one non-yielding sweep. After this loop @@ -603,8 +564,12 @@ async def drop_group(self, group_id: uuid.UUID) -> int: last_meta: InflightRollout | None = claimed[-1][1] if claimed else None for _, meta in claimed: for _ in range(meta.rollout_count): - raw = self.error_rollout_output(error_type="Cancelled", error_repr="Off-policy cancel") - await self.emit_rollout(meta, group, raw) + trace = Rollout( + task=vf.Task(idx=task_idx, prompt=None), + errors=[vf.Error(type="Cancelled", message="Off-policy cancel")], + stop_condition="error", + ) + await self.emit_rollout(meta, group, trace) # For non-group-scoring envs, the group may have rollouts that # were never dispatched (``rollouts_to_schedule > 0``). Emit @@ -625,8 +590,12 @@ async def drop_group(self, group_id: uuid.UUID) -> int: ) unscheduled_cancelled = group.rollouts_to_schedule for _ in range(unscheduled_cancelled): - raw = self.error_rollout_output(error_type="Cancelled", error_repr="Off-policy cancel") - await self.emit_rollout(fallback_meta, group, raw) + trace = Rollout( + task=vf.Task(idx=task_idx, prompt=None), + errors=[vf.Error(type="Cancelled", message="Off-policy cancel")], + stop_condition="error", + ) + await self.emit_rollout(fallback_meta, group, trace) cancelled = inflight_cancelled + unscheduled_cancelled if cancelled > 0: diff --git a/src/prime_rl/orchestrator/env_server/env_server.py b/src/prime_rl/orchestrator/env_server/env_server.py index 8ff057b9d97..18dbcb46cb6 100644 --- a/src/prime_rl/orchestrator/env_server/env_server.py +++ b/src/prime_rl/orchestrator/env_server/env_server.py @@ -1,38 +1,35 @@ -import asyncio +from functools import partial -from verifiers.serve import ZMQEnvServer +from verifiers.v1 import pool_serve_kwargs +from verifiers.v1.serve import serve_env from prime_rl.configs.env_server import EnvServerConfig +from prime_rl.orchestrator.utils import setup_env_server_logging from prime_rl.utils.config import cli -from prime_rl.utils.logger import setup_logger -from prime_rl.utils.pathing import get_log_dir from prime_rl.utils.process import set_proc_title -from prime_rl.utils.utils import clean_exit, get_env_ids_to_install, install_env +from prime_rl.utils.utils import clean_exit @clean_exit def run_server(config: EnvServerConfig): - setup_logger(config.log.level, json_logging=config.log.json_logging) - - # install environment if not already installed - env_ids_to_install = set() - env_ids_to_install.update(get_env_ids_to_install([config.env])) - for env_id in env_ids_to_install: - install_env(env_id, prerelease=config.env_install_prerelease) - - log_dir = (get_log_dir(config.output_dir) / config.env.resolved_name).as_posix() - - server = ZMQEnvServer( - env_id=config.env.stripped_id, - env_args=config.env.args, - extra_env_kwargs=config.env.extra_env_kwargs, - log_level=config.log.level, - log_dir=log_dir, - json_logging=config.log.json_logging, - num_workers=config.env.num_workers, - **{"address": config.env.address} if config.env.address is not None else {}, + env = config.env + address = env.address or "tcp://127.0.0.1:5000" + # The env's ``pool`` (static or elastic) sizes the server; a v0/legacy env runs through + # the bridge, a v1 env is a native taskset — both serve vf.Trace over the same protocol, + # so the orchestrator is agnostic. serve_env applies the logging setup in this process + # and in every spawned worker. + server_kwargs = ( + {"env_id": env.env_id, "env_args": env.args, "extra_env_kwargs": env.extra_env_kwargs} + if env.is_legacy + else {"config": env} + ) + serve_env( + **pool_serve_kwargs(env.pool), + legacy=env.is_legacy, + address=address, + log_setup=partial(setup_env_server_logging, config.log.level, config.log.json_logging), + **server_kwargs, ) - asyncio.run(server.run()) def main(): diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 0bf8dc91010..33fb3eca84d 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -1,33 +1,90 @@ +"""Env wrappers over a v1 env server. + +Each ``Env`` owns a v1 ``EnvServer`` (spawned as a child process, or an +external one given by ``config.address``) and an ``EnvClient`` to drive it. The +orchestrator never *runs* an environment: it asks the server for ``info`` +(``num_tasks`` + whether group scoring is needed), then runs rollouts purely by +**task index**. The server returns a ``Trace`` (a plain ``model_dump`` — derived values are +properties, not serialized) which we validate into a ``Trace[WireTask]`` — a real ``vf.Trace`` +(never a loose dict) whose task keeps the env's +task-specific fields as extras (``WireTask`` allows them). The orchestrator never imports the +env package: the env's *type* and *runtime* both live only in the server, and the orchestrator +drives it purely by task index. (Nothing here reads typed env task fields — only ``task.idx`` +and a full ``task.model_dump``, both of which ``WireTask`` preserves.) +""" + from __future__ import annotations +import asyncio import atexit import multiprocessing as mp +import os +import queue +import sys from collections.abc import Iterator, Sequence from multiprocessing.process import BaseProcess from pathlib import Path from typing import Generic, TypeVar -import verifiers as vf -from verifiers.serve import ZMQEnvClient, ZMQEnvServer -from verifiers.utils.serve_utils import get_free_port +import verifiers.v1 as vf +from verifiers.v1.serve import EnvClient from prime_rl.configs.orchestrator import EnvConfig, EvalEnvConfig, TrainEnvConfig from prime_rl.orchestrator.advantage import AdvantageFn, setup_advantage_fn +from prime_rl.orchestrator.types import Rollout from prime_rl.utils.logger import get_logger -REQUIRED_STATE_COLUMNS = ["trajectory"] +# Every wire trace validates into this type. WireTask (extra="allow") keeps the env's task +# fields without importing the env package — the orchestrator never reads them typed (only +# task.idx + task.model_dump). +ROLLOUT_TYPE = Rollout[vf.WireTask] + +# Max wait for a spawned env server to bind and report its address. The child +# loads the taskset (possibly downloading a dataset) before reporting, so this +# is generous. +ENV_SERVER_SPAWN_TIMEOUT = 600.0 + + +def _run_env_server( + *, + log_file: str, + log_level: str, + json_logging: bool, + legacy: bool = False, + **kwargs, +) -> None: + """Spawned-process entry point: redirect this process's output to ``log_file`` (the + server's logging + any subprocess-runtime output), then serve via ``serve_env``. The + worker-pool sizing arrives in ``kwargs`` (``max_workers`` / ``multiplex`` / ``elastic`` + from the env's ``pool``). ``serve_env`` applies ``log_setup`` here and in every spawned + worker; a worker inherits this process's redirected stdout/stderr, so its per-rollout + logs reach ``log_file`` too. Top-level so it stays picklable for the ``spawn`` start + method. ``legacy`` picks the v0 bridge.""" + from functools import partial + + from verifiers.v1.serve import serve_env + + from prime_rl.orchestrator.utils import setup_env_server_logging + + fh = open(log_file, "w", buffering=1) + os.dup2(fh.fileno(), sys.stdout.fileno()) + os.dup2(fh.fileno(), sys.stderr.fileno()) + serve_env( + legacy=legacy, + log_setup=partial(setup_env_server_logging, log_level, json_logging), + **kwargs, + ) class Env: - """Wraps a vf.Environment - only exposes features used in PRIME-RL.""" + """Wraps a v1 env server + client. The orchestrator never loads the env.""" def __init__(self, config: EnvConfig): self.config = config self.sampling_args: dict = {} - - get_logger().debug(f"Initializing {config.resolved_name} ({config})") - self._env: vf.Environment = vf.load_environment(config.stripped_id, **config.args) - self._env_client: ZMQEnvClient | None = None + self.num_tasks: int = 0 + self.requires_group_scoring: bool = False + self._env_client: EnvClient | None = None self._env_server_process: BaseProcess | None = None @property @@ -35,122 +92,105 @@ def name(self) -> str: return self.config.resolved_name @property - def env(self) -> vf.Environment: - return self._env - - @property - def env_client(self) -> ZMQEnvClient: - if not self._env_client: - raise RuntimeError( - f"Env {self.name} has no env client connected. Call connect() first to connect to an env server." - ) + def env_client(self) -> EnvClient: + if self._env_client is None: + raise RuntimeError(f"Env {self.name} not started — call start() first.") return self._env_client - @property - def requires_group_scoring(self) -> bool: - return any(self.env.rubric._is_group_func(func) for func in self.env.rubric._get_reward_funcs()) - - async def start( - self, - log_dir: Path, - log_level: str | None = None, - json_logging: bool = False, - ) -> None: - """Spawn an env server (if needed) and connect to it.""" - if self.config.address is None: - address = self._spawn(log_dir=log_dir, log_level=log_level, json_logging=json_logging) - else: - address = self.config.address + async def start(self, log_dir: Path, log_level: str | None = None, json_logging: bool = False) -> None: + """Spawn the env server (if needed), connect, and cache its ``info``.""" + external = self.config.address is not None + address = self.config.address or await self._spawn(log_dir, log_level or "INFO", json_logging) get_logger().debug(f"Connecting {self.name} to env server {address}") - self._env_client = ZMQEnvClient(address=address, name=self.name) - await self.env_client.wait_for_server_startup() - - def _spawn( - self, - log_dir: Path, - log_level: str | None = None, - json_logging: bool = False, - ) -> str: - assert isinstance(self.config.num_workers, int), ( - f"num_workers must be resolved before spawn, got {self.config.num_workers!r}" + self._env_client = EnvClient(address=address) + # A spawned server already reported its address *after* binding + loading, + # so it's up — the untimed ``info`` below is enough. An external server has + # no such handshake, so poll until it answers before we block on ``info``. + if external: + await self.env_client.wait_for_server_startup() + info = await self.env_client.info() + self.num_tasks = info.num_tasks + self.requires_group_scoring = info.requires_group_scoring + get_logger().info( + f"Env {self.name} ready: num_tasks={self.num_tasks} group_scoring={self.requires_group_scoring}" ) - num_workers = self.config.num_workers - address = f"tcp://127.0.0.1:{get_free_port()}" - get_logger().debug(f"Spawning env server {self.name} ({address=}, {num_workers=})") - process = mp.get_context("spawn").Process( - target=ZMQEnvServer.run_server, - args=( - self.config.stripped_id, - self.config.args, - self.config.extra_env_kwargs, - log_level, - (log_dir / self.name).as_posix(), - ), + + async def _spawn(self, log_dir: Path, log_level: str, json_logging: bool) -> str: + """Spawn a v1 EnvServer child process (it loads the env; we never do). + The server binds an OS-assigned port (``:0``) and reports the concrete + address back over a queue — no free-port guess, no TOCTOU race. Its output + goes to ``/.log`` (``log_dir`` is already the train/eval-split + ``.../logs/envs/{train,eval}`` the orchestrator passes in).""" + ctx = mp.get_context("spawn") + address_queue: mp.Queue = ctx.Queue() + log_file = log_dir / f"{self.name}.log" + log_file.parent.mkdir(parents=True, exist_ok=True) + get_logger().debug(f"Spawning env server {self.name} (id={self.config.env_id}, log={log_file})") + server_kwargs = ( + dict( + legacy=True, + env_id=self.config.env_id, + env_args=self.config.args, + extra_env_kwargs=self.config.extra_env_kwargs, + ) + if self.config.is_legacy + else dict(legacy=False, config=self.config) + ) + process = ctx.Process( + target=_run_env_server, kwargs=dict( - address=address, + log_file=str(log_file), + log_level=log_level, json_logging=json_logging, - console_logging=False, - num_workers=num_workers, + **vf.pool_serve_kwargs(self.config.pool), + address="tcp://127.0.0.1:0", + address_queue=address_queue, + **server_kwargs, ), daemon=False, ) process.start() self._env_server_process = process + try: + address = await asyncio.to_thread(address_queue.get, timeout=ENV_SERVER_SPAWN_TIMEOUT) + except queue.Empty: + raise RuntimeError(f"Env server {self.name} did not report its address within {ENV_SERVER_SPAWN_TIMEOUT}s") + finally: + address_queue.close() + address_queue.join_thread() + get_logger().debug(f"Env server {self.name} bound at {address}") return address - def _sampling_args_with_salt(self, cache_salt: str | None) -> dict: - sampling_args = {**self.sampling_args} - if cache_salt is None: - return sampling_args - extra_body = {**sampling_args.get("extra_body", {}), "cache_salt": cache_salt} - sampling_args["extra_body"] = extra_body - return sampling_args - - @property - def state_columns(self) -> list[str]: - """Required columns plus any extras configured on the env, deduped (required first).""" - merged: list[str] = [] - for col in (*REQUIRED_STATE_COLUMNS, *self.config.state_columns): - if col not in merged: - merged.append(col) - return merged + def _sampling(self, cache_salt: str | None) -> vf.SamplingConfig: + sampling = {**self.sampling_args} + if cache_salt is not None: + sampling["extra_body"] = {**sampling.get("extra_body", {}), "cache_salt": cache_salt} + return vf.SamplingConfig(**sampling) async def run_rollout( - self, - client: vf.ClientConfig, - example: dict, - model_name: str, - cache_salt: str | None, - ) -> vf.RolloutOutput: - """Run a single rollout for an example.""" - return await self.env.run_rollout( - vf.RolloutInput(**example), + self, client: vf.ClientConfig, task_idx: int, model_name: str, cache_salt: str | None + ) -> Rollout: + """Run a single rollout for ``task_idx``; return a typed Trace.""" + wire = await self.env_client.run_rollout( + task_idx=task_idx, client=client, model=model_name, - sampling_args=self._sampling_args_with_salt(cache_salt), - max_retries=self.config.max_retries, - state_columns=self.state_columns, - env_client=self.env_client, + sampling=self._sampling(cache_salt), ) + return ROLLOUT_TYPE.model_construct(**dict(wire)) async def run_group( - self, - client: vf.ClientConfig, - example: dict, - model_name: str, - group_size: int, - cache_salt: str | None, - ) -> list[vf.RolloutOutput]: - """Run a group of rollouts for an example. Required for group-scoring envs.""" - return await self.env.run_group( - [vf.RolloutInput(**example) for _ in range(group_size)], + self, client: vf.ClientConfig, task_idx: int, model_name: str, group_size: int, cache_salt: str | None + ) -> list[Rollout]: + """Run a group of rollouts for ``task_idx`` (group-scoring envs); return typed Traces.""" + wires = await self.env_client.run_group( + task_idx=task_idx, + n=group_size, client=client, model=model_name, - sampling_args=self._sampling_args_with_salt(cache_salt), - max_retries=self.config.max_retries, - state_columns=self.state_columns, - env_client=self.env_client, + sampling=self._sampling(cache_salt), ) + return [ROLLOUT_TYPE.model_construct(**dict(wire)) for wire in wires] def shutdown(self) -> None: if self._env_server_process is None: @@ -165,15 +205,11 @@ class TrainEnv(Env): def __init__(self, config: TrainEnvConfig, max_seq_len: int): super().__init__(config) self.sampling_args = config.sampling.to_sampling_args() - # Built once — custom advantage funcs do an ``import_object`` we don't - # want to pay per group. ``None`` = reward-only path. + # Built once — custom advantage funcs do an ``import_object`` we don't want to pay per group. self.advantage_fn: AdvantageFn | None = ( setup_advantage_fn(config.advantage, max_seq_len=max_seq_len) if config.advantage is not None else None ) - def get_dataset(self, seed: int | None = None): - return self.env.get_dataset(seed=seed) - class EvalEnv(Env): config: EvalEnvConfig @@ -181,7 +217,12 @@ class EvalEnv(Env): def __init__(self, config: EvalEnvConfig): super().__init__(config) self.sampling_args = config.sampling.to_sampling_args() - self.examples = self.env.get_eval_dataset(n=config.num_examples).to_list() + self.examples: list[dict] = [] + + async def start(self, log_dir: Path, log_level: str | None = None, json_logging: bool = False) -> None: + await super().start(log_dir=log_dir, log_level=log_level, json_logging=json_logging) + n = self.num_tasks if self.config.num_examples < 0 else min(self.config.num_examples, self.num_tasks) + self.examples = [{"task_idx": i} for i in range(n)] EnvT = TypeVar("EnvT", bound=Env) @@ -209,23 +250,15 @@ def __iter__(self) -> Iterator[EnvT]: def __len__(self) -> int: return len(self._envs) - async def start( - self, - log_dir: Path, - log_level: str | None = None, - json_logging: bool = False, - ) -> None: - """Spawn env servers (where needed) and connect env clients one at a time. - - Serialized to avoid a TOCTOU port race: get_free_port() only holds the port - until it returns, so parallel spawns can hand the same port to two children. - """ + async def start(self, log_dir: Path, log_level: str | None = None, json_logging: bool = False) -> None: + """Spawn env servers (where needed) and connect, one at a time. Each server + binds an OS-assigned port and reports it back, so there's no port race.""" for env in self: await env.start(log_dir=log_dir, log_level=log_level, json_logging=json_logging) atexit.register(self.shutdown) def shutdown(self) -> None: - """Terminate all spawned env server processes in parallel.""" + """Terminate all spawned env server processes.""" processes = [env._env_server_process for env in self if env._env_server_process is not None] if not processes: return diff --git a/src/prime_rl/orchestrator/eval_sink.py b/src/prime_rl/orchestrator/eval_sink.py index 2c218416789..f6ceda5f58c 100644 --- a/src/prime_rl/orchestrator/eval_sink.py +++ b/src/prime_rl/orchestrator/eval_sink.py @@ -18,7 +18,7 @@ from prime_rl.orchestrator.envs import EvalEnvs from prime_rl.orchestrator.eval_utils import compute_pass_at_k -from prime_rl.orchestrator.types import EvalBatch, EvalBatchMetrics, EvalRollout +from prime_rl.orchestrator.types import EvalBatch, EvalBatchMetrics, Rollout from prime_rl.utils.logger import get_logger @@ -27,12 +27,12 @@ class EvalSink: def __init__(self, *, eval_envs: EvalEnvs) -> None: self.eval_envs = eval_envs - self.pending_groups: dict[uuid.UUID, list[EvalRollout]] = defaultdict(list) + self.pending_groups: dict[uuid.UUID, list[Rollout]] = defaultdict(list) # Bucket size IS the arrival count — ``process_group`` flushes # everything in without filtering - self.pending_batches: dict[tuple[str, int], list[EvalRollout]] = defaultdict(list) + self.pending_batches: dict[tuple[str, int], list[Rollout]] = defaultdict(list) - def add(self, rollout: EvalRollout) -> EvalBatch | None: + def add(self, rollout: Rollout) -> EvalBatch | None: """Process one arrival; finalize the group on the ``group_size``-th arrival and the per-env epoch on the ``num_examples × group_size``-th.""" env_name = rollout.env_name @@ -82,7 +82,7 @@ def batch_progress(self) -> list[tuple[str, int, int, int, int]]: # ── level 1: per-rollout (no-op for eval) ───────────────────────────── - def process_rollout(self, rollout: EvalRollout) -> None: + def process_rollout(self, rollout: Rollout) -> None: """No-op. Eval rollouts don't need trainer-bound tokenization; the method exists to keep the three-level structure uniform with ``TrainSink``. @@ -96,17 +96,17 @@ def process_group(self, group_id: uuid.UUID) -> None: if not group: return env_name = group[0].env_name - example_id = group[0].example_id + task_idx = group[0].task.idx eval_step = group[0].eval_step bucket = self.pending_batches[(env_name, eval_step)] bucket.extend(group) - survivors = [r for r in group if r.error is None] + survivors = [r for r in group if not r.has_error] num_errored = len(group) - len(survivors) rewards = [r.reward for r in survivors] avg_reward = sum(rewards) / len(rewards) if rewards else 0.0 get_logger().debug( - f"Finished group | env={env_name} example_id={example_id} eval_step={eval_step} | " + f"Finished group | env={env_name} task_idx={task_idx} eval_step={eval_step} | " f"rollouts={len(group)} (errored={num_errored}) | reward={avg_reward:.4f}" ) @@ -120,9 +120,9 @@ def process_batch(self, key: tuple[str, int]) -> EvalBatch: rollouts = self.pending_batches.pop(key, []) n_total = len(rollouts) - n_cancelled = sum(1 for r in rollouts if (r.error or {}).get("error") == "Cancelled") - n_errored = sum(1 for r in rollouts if r.error is not None) - n_cancelled - valid = [r for r in rollouts if r.error is None] + n_cancelled = sum(1 for r in rollouts if r.has_error and r.error.type == "Cancelled") + n_errored = sum(1 for r in rollouts if r.has_error) - n_cancelled + valid = [r for r in rollouts if not r.has_error] metrics = EvalBatchMetrics( n_rollouts=n_total, n_cancelled=n_cancelled, @@ -131,23 +131,23 @@ def process_batch(self, key: tuple[str, int]) -> EvalBatch: if valid: rewards = [r.reward for r in valid] - lens = [r.raw["token_usage"]["final_output_tokens"] for r in valid] + lens = [r.completion_len for r in valid] metrics.group_size = self.group_size_for(env_name) metrics.reward_mean = float(sum(rewards) / len(rewards)) metrics.completion_len_mean = float(sum(lens) / len(lens)) metrics.completion_len_max = float(max(lens)) metrics.completion_len_min = float(min(lens)) metrics.truncation_rate = float(sum(1 for r in valid if r.is_truncated) / len(valid)) - metrics.no_response_rate = float(sum(1 for r in valid if not r.raw.get("completion")) / len(valid)) - num_turns = [len(r.raw.get("trajectory") or []) for r in valid] + metrics.no_response_rate = float(sum(1 for r in valid if not r.has_response) / len(valid)) + num_turns = [r.num_turns for r in valid] metrics.num_turns_mean = float(sum(num_turns) / len(num_turns)) metrics.num_turns_min = float(min(num_turns)) metrics.num_turns_max = float(max(num_turns)) # pass@k: errored attempts don't count toward k tries - by_example: dict[int | str, list[float]] = {} + by_example: dict[int, list[float]] = {} for r in valid: - by_example.setdefault(r.example_id, []).append(r.reward) + by_example.setdefault(r.task.idx, []).append(r.reward) metrics.n_examples = len(by_example) unique_rewards = {float(r) for r in rewards} if unique_rewards.issubset({0.0, 1.0}) and by_example: diff --git a/src/prime_rl/orchestrator/filters.py b/src/prime_rl/orchestrator/filters.py index f8deda1230b..b90e8b2537e 100644 --- a/src/prime_rl/orchestrator/filters.py +++ b/src/prime_rl/orchestrator/filters.py @@ -16,7 +16,7 @@ from prime_rl.utils.logger import get_logger if TYPE_CHECKING: - from prime_rl.orchestrator.types import TrainRollout + from prime_rl.orchestrator.types import Rollout @dataclass @@ -29,7 +29,7 @@ class RolloutFilter(Protocol): name: str enforce: bool - def check(self, rollout: "TrainRollout") -> FilterResult: ... + def check(self, rollout: Rollout) -> FilterResult: ... @dataclass @@ -49,13 +49,14 @@ class GibberishFilter: logprob_threshold: float enforce: bool = False - def check(self, rollout: "TrainRollout") -> FilterResult: + def check(self, rollout: Rollout) -> FilterResult: global_idx = 0 - for step in rollout.raw["trajectory"]: - tokens = step["tokens"] - if tokens is None: - continue - for token_id, logprob in zip(tokens["completion_ids"], tokens["completion_logprobs"]): + for node in rollout.nodes: + # token_ids / logprobs / mask are per-token aligned; check only the + # sampled completion tokens, pairing each with its own logprob. + for token_id, logprob, sampled in zip(node.token_ids, node.logprobs, node.mask): + if not sampled: + continue if token_id > self.token_id_threshold and logprob < self.logprob_threshold: return FilterResult(detected=True, detection_index=global_idx) global_idx += 1 @@ -79,14 +80,15 @@ class RepetitionFilter: logprob_threshold: float enforce: bool = False - def check(self, rollout: "TrainRollout") -> FilterResult: + def check(self, rollout: Rollout) -> FilterResult: consecutive = 0 global_idx = 0 - for step in rollout.raw["trajectory"]: - tokens = step["tokens"] - if tokens is None: - continue - for logprob in tokens["completion_logprobs"]: + for node in rollout.nodes: + # Streak only over sampled completion tokens — context/prompt tokens + # (mask=False) are not model-generated and must not feed the loop count. + for logprob, sampled in zip(node.logprobs, node.mask): + if not sampled: + continue if logprob > self.logprob_threshold: consecutive += 1 else: @@ -105,7 +107,7 @@ class ZeroAdvantageFilter: name: str enforce: bool = True - def check(self, rollout: "TrainRollout") -> FilterResult: + def check(self, rollout: Rollout) -> FilterResult: if rollout.advantage is not None and rollout.advantage == 0.0: return FilterResult(detected=True) return FilterResult(detected=False) @@ -147,8 +149,8 @@ def setup_filters(configs: list[FilterConfig], vocab_size: int, *, kind: str) -> return filters -def apply_filters(filters: list[RolloutFilter], rollouts: list["TrainRollout"]) -> None: # noqa: F821 (forward ref) - """Flag ``TrainRollout``\\ s in place with per-filter detection + drop decision. +def apply_filters(filters: list[RolloutFilter], rollouts: list[Rollout]) -> None: + """Flag ``Rollout``\\ s in place with per-filter detection + drop decision. Each rollout's ``filter_results`` dict records per-filter detection bools; ``is_filtered`` is True iff an enforcing filter detected it. First matching diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index d32dbc02def..459a674268d 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -1,9 +1,4 @@ -"""MetricsBuilder: assembles the per-step W&B dict. - -The only I/O / state is the trainer's token-export metrics, which lag the -orchestrator: ``build`` folds in the oldest unlogged stable step it finds on -disk and tracks the last step logged so it never re-logs one. -""" +"""MetricsBuilder: assembles the per-step W&B dict. No I/O, no side effects.""" from __future__ import annotations @@ -12,22 +7,18 @@ import pandas as pd from prime_rl.configs.orchestrator import OrchestratorConfig -from prime_rl.orchestrator.token_export_metrics import collect_next_token_export_metrics -from prime_rl.orchestrator.types import Progress, TrainBatchMetrics, TrainRollout +from prime_rl.orchestrator.types import Progress, Rollout, TrainBatchMetrics class MetricsBuilder: - def __init__(self, config: OrchestratorConfig, start_step: int = 0) -> None: + def __init__(self, config: OrchestratorConfig) -> None: self.config = config - # Token exports are read back one step behind the orchestrator; track the - # last run step whose metrics we've folded in so we never re-log a step. - self._last_token_export_step_logged = start_step - 1 def build( self, *, step: int, - rollouts: list[TrainRollout], + rollouts: list[Rollout], metrics: TrainBatchMetrics, progress: Progress, step_time: float, @@ -41,30 +32,25 @@ def build( existing dashboards / alerts keep working.""" num_rollouts = len(rollouts) num_unique_examples = len({r.group_id for r in rollouts}) - num_tokens = sum( - r.raw["token_usage"]["final_input_tokens"] + r.raw["token_usage"]["final_output_tokens"] for r in rollouts - ) + num_tokens = sum(r.total_tokens for r in rollouts) results_df = pd.DataFrame( { "group_id": [r.group_id for r in rollouts], - "example_id": [r.example_id for r in rollouts], + "task_idx": [r.task.idx for r in rollouts], "env_name": [r.env_name for r in rollouts], "reward": [r.reward for r in rollouts], "is_truncated": [r.is_truncated for r in rollouts], "is_filtered": [r.is_filtered for r in rollouts], - "stop_condition": [r.raw.get("stop_condition") for r in rollouts], - "seq_len": [ - r.raw["token_usage"]["final_input_tokens"] + r.raw["token_usage"]["final_output_tokens"] - for r in rollouts - ], + "stop_condition": [r.stop_condition for r in rollouts], + "seq_len": [r.total_tokens for r in rollouts], "prefill_len": metrics.rollout_prefill_lens, "decode_len": metrics.rollout_decode_lens, "samples_per_rollout": metrics.samples_per_rollout, - "num_turns": [len(r.raw["trajectory"]) for r in rollouts], + "num_turns": [r.num_turns for r in rollouts], } ) - metrics_df = pd.DataFrame([(r.raw.get("metrics") or {}) for r in rollouts]) + metrics_df = pd.DataFrame([r.metrics for r in rollouts]) filter_df = pd.DataFrame([r.filter_results for r in rollouts]) timing_df = self.timing_df(rollouts) @@ -188,42 +174,14 @@ def compute_solve_rates(df): for name, count in pre_filter_dropped_by_name.items(): to_log[f"pre_filters/all/{name}/rate"] = count / pre_filter_seen - # Fold in the trainer's token-export metrics for the oldest unlogged stable - # step (exports always lag the orchestrator, so this is a past step). - to_log.update(self.next_token_export_metrics()) - return to_log - @property - def last_token_export_step_logged(self) -> int: - return self._last_token_export_step_logged - - def next_token_export_metrics(self) -> dict[str, float | int]: - """Log-ready token-export metrics for the next unlogged stable step (empty if - none), advancing the cursor. They arrive under trainer/, including trainer/step - (the run step they belong to) for plotting against a lag-corrected axis. Shared - by build() and the orchestrator's end-of-run drain.""" - token_export = collect_next_token_export_metrics( - self.config.output_dir, - last_logged_step=self._last_token_export_step_logged, - ) - if token_export: - self._last_token_export_step_logged = token_export["trainer/step"] - return token_export - @staticmethod - def timing_df(rollouts: list[TrainRollout]) -> pd.DataFrame: - return pd.DataFrame( - [ - { - "total": r.raw["timing"]["total"], - "setup": r.raw["timing"]["setup"]["duration"], - "generation": r.raw["timing"]["generation"]["duration"], - "model": r.raw["timing"]["model"]["duration"], - "env": r.raw["timing"]["env"]["duration"], - "scoring": r.raw["timing"]["scoring"]["duration"], - "overhead": r.raw["timing"]["overhead"], - } - for r in rollouts - ] - ) + def timing_df(rollouts: list[Rollout]) -> pd.DataFrame: + """Per-rollout timing from the v1 Trace (`generation`/`scoring` spans).""" + rows = [] + for r in rollouts: + timing = r.timing + generation, scoring = timing.generation.duration, timing.scoring.duration + rows.append({"total": generation + scoring, "generation": generation, "scoring": scoring}) + return pd.DataFrame(rows) diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index b1c30902ccf..1344596cfd0 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -21,7 +21,6 @@ from __future__ import annotations import asyncio -import logging import os import time from typing import TYPE_CHECKING @@ -36,8 +35,6 @@ from prime_rl.transport.base import TrainingBatchSender from prime_rl.utils.client import InferencePool from prime_rl.utils.monitor.base import Monitor -from verifiers.utils.async_utils import EventLoopLagMonitor, EventLoopLagStats - import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before transitive imports from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.orchestrator.ckpt import setup_ckpt_manager @@ -57,12 +54,10 @@ from prime_rl.orchestrator.train_source import TrainSource from prime_rl.orchestrator.types import ( EvalBatch, - EvalRollout, - FinishedRollout, Policy, Progress, + Rollout, TrainBatch, - TrainRollout, ) from prime_rl.orchestrator.utils import ( compute_teacher_logprobs, @@ -76,7 +71,7 @@ from prime_rl.orchestrator.watcher import WeightWatcher from prime_rl.trainer.model import setup_tokenizer from prime_rl.transport import TrainingBatch, setup_training_batch_sender -from prime_rl.utils.async_utils import safe_cancel +from prime_rl.utils.async_utils import EventLoopLagMonitor, EventLoopLagStats, safe_cancel from prime_rl.utils.client import init_nccl_broadcast, setup_inference_pool from prime_rl.utils.heartbeat import Heartbeat from prime_rl.utils.logger import format_time, get_logger, setup_logger @@ -85,8 +80,6 @@ from prime_rl.utils.usage_reporter import UsageReporter from prime_rl.utils.utils import ( clean_exit, - get_env_ids_to_install, - install_env, resolve_latest_ckpt_step, ) @@ -108,10 +101,12 @@ # resumed when the watcher advances ``policy.version``. TARGET_LAG = 1 -# After max_steps, how long to wait for the trainer to finalize the trailing -# token-export steps it hadn't flushed yet, and how often to poll for them. -TOKEN_EXPORT_DRAIN_TIMEOUT_S = 300.0 -TOKEN_EXPORT_DRAIN_POLL_S = 2.0 +# Drop the per-node training tensors when dumping a Trace to disk (rollout jsonl / wandb +# tables): the multimodal `mm_kwargs` carrier and the router-replay `routed_experts` array are +# training inputs, not part of the rollout record, and would bloat every line. They also can't +# round-trip the json dump — their `__nd__` carriers hold raw bytes. `__all__` applies the +# exclude to every node in the list. +ROLLOUT_DUMP_EXCLUDE = {"nodes": {"__all__": {"multi_modal_data", "routed_experts"}}} class Orchestrator: @@ -158,10 +153,9 @@ class Orchestrator: def __init__(self, config: OrchestratorConfig) -> None: self.config = config setup_logger(config.log.level, json_logging=config.log.json_logging) - # Silence in-process ``verifiers.*`` library noise but keep - # ``verifiers.serve`` (env-server lifecycle) through our handler - logging.getLogger("verifiers").setLevel(logging.CRITICAL + 1) - intercept_vf_logging(logger="verifiers.serve", level="WARN") + # Route the in-process v1 library logging through our handler. The + # env server runs in a child process, so its logging is separate. + intercept_vf_logging(logger="verifiers.v1", level="WARN") get_logger().info(f"Starting orchestrator ({config.training_mode})") if config.bench: @@ -211,12 +205,6 @@ async def setup(self) -> None: with open(config_dir / "orch.toml", "wb") as f: tomli_w.dump(config.model_dump(exclude_none=True, mode="json"), f) - env_ids_to_install = set(get_env_ids_to_install(config.train.env)) - if config.eval is not None: - env_ids_to_install.update(get_env_ids_to_install(config.eval.env)) - for env_id in env_ids_to_install: - install_env(env_id, prerelease=config.env_install_prerelease) - get_logger().info(f"Initializing tokenizer ({config.tokenizer})") self.tokenizer = setup_tokenizer(config.tokenizer) @@ -242,7 +230,11 @@ async def setup(self) -> None: self.teacher_inference = await setup_inference_pool( config.teacher.client, model_name=config.teacher.model.name, - train_client_type="openai_chat_completions", + # SFT rolls the teacher out through the renderer client (token-in/out) so its + # rollouts carry tokens directly — training is renderer-only. (OPD reads teacher + # logprobs via prefill, so its pool client type is moot.) + train_client_type="renderer", + renderer_config=config.renderer, ) get_logger().info(f"Initializing monitor (wandb={config.wandb}, prime_monitor={config.prime_monitor})") @@ -387,11 +379,10 @@ async def setup(self) -> None: training_mode=config.training_mode, use_cache_salt=use_cache_salt, ) - self.metrics = MetricsBuilder(config, start_step=self.progress.step) + self.metrics = MetricsBuilder(config) self.train_sink = TrainSink( config, tokenizer=self.tokenizer, - renderer=self.renderer, train_envs=self.train_envs, mm_token_type_ids_mapping=self.mm_token_type_ids_mapping, batch_size=config.batch_size, @@ -465,7 +456,6 @@ async def start(self) -> None: clean_exit = False try: await self.main_loop() - await self._drain_token_export_metrics() clean_exit = True finally: elapsed = format_time(time.perf_counter() - start_time) @@ -495,53 +485,23 @@ async def main_loop(self) -> None: break try: - rollout: FinishedRollout = await asyncio.wait_for(self.dispatcher.out_q.get(), timeout=0.5) + rollout: Rollout = await asyncio.wait_for(self.dispatcher.out_q.get(), timeout=0.5) except asyncio.TimeoutError: continue - if isinstance(rollout, EvalRollout): + if rollout.kind == "eval": assert self.eval_sink is not None # eval rollouts only emitted when eval is configured eval_batch = self.eval_sink.add(rollout) if eval_batch is not None: await self.finalize_eval_batch(eval_batch) continue - assert isinstance(rollout, TrainRollout) train_batch = await self.train_sink.add(rollout) # In drain mode any late-arriving train batch is dropped — we # don't want to ship past ``max_steps`` if train_batch is not None and not self.draining and not self.stopped.is_set(): await self.finalize_train_batch(train_batch) - async def _drain_token_export_metrics(self) -> None: - """Token exports lag the orchestrator, so once the loop ends the trainer is - still finalizing the last shipped steps and ``build`` no longer runs to pick - them up. Poll for the remaining stable steps up to ``max_steps - 1`` and log - each under its own ``trainer/step`` (a local W&B step axis, so the final - checkpoint's ``progress.step`` is untouched).""" - if self.config.max_steps is None: - return - if not (self.config.output_dir / "token_exports").exists(): - return # token export not enabled for this run — don't block shutdown - target = self.config.max_steps - 1 - if self.metrics.last_token_export_step_logged >= target: - return # build() already kept up — nothing trailing to flush - wandb_step = self.progress.step - deadline = time.perf_counter() + TOKEN_EXPORT_DRAIN_TIMEOUT_S - while time.perf_counter() < deadline: - token_export = self.metrics.next_token_export_metrics() - if not token_export: - await asyncio.sleep(TOKEN_EXPORT_DRAIN_POLL_S) - continue - self.monitor.log(token_export, step=wandb_step) - wandb_step += 1 - if self.metrics.last_token_export_step_logged >= target: - return - get_logger().warning( - f"Token-export drain timed out before step {target} " - f"(last logged {self.metrics.last_token_export_step_logged})" - ) - async def finalize_train_batch(self, batch: TrainBatch) -> None: """Ship one ``TrainBatch`` out to the trainer and handle the I/O side-effects (ckpt, save_rollouts, teacher logprobs, sender.send, @@ -588,13 +548,11 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: f"({batch.metrics.n_trainable / len(batch.rollouts):.1%}) — consider reviewing task difficulty / filter config" ) - # Materialize at the I/O boundary so prime-rl metadata travels with - # the raw vf payload on disk + in wandb sample tables - rollout_dicts = [r.to_dict() for r in batch.rollouts] + # Serialize the typed Trace at the I/O boundary (disk + wandb sample tables); drop the + # per-node multimodal tensors — they're for training, not the rollout record, and bloat it. + rollout_dicts = [r.model_dump(mode="json", exclude=ROLLOUT_DUMP_EXCLUDE) for r in batch.rollouts] step_path = get_step_path(get_rollout_dir(config.output_dir), step) - await asyncio.to_thread( - save_rollouts, rollout_dicts, step_path / "train_rollouts.jsonl", exclude_keys={"trajectory"} - ) + await asyncio.to_thread(save_rollouts, rollout_dicts, step_path / "train_rollouts.jsonl") teacher_logprobs_time = 0.0 # opd only if config.training_mode == "opd" and self.teacher_inference is not None: @@ -626,7 +584,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: pre_filter_dropped_by_name=dict(self.train_sink.pre_filter_dropped_by_name), ) self.monitor.log(metrics, step=step) - self.monitor.log_samples(rollout_dicts, step=step) + self.monitor.log_samples(batch.rollouts, step=step) self.monitor.log_distributions( distributions={ "rewards": [r.reward for r in batch.rollouts], @@ -648,10 +606,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: num_rollouts = len(batch.rollouts) num_unique_examples = len({r.group_id for r in batch.rollouts}) - num_tokens = sum( - r.raw["token_usage"]["final_input_tokens"] + r.raw["token_usage"]["final_output_tokens"] - for r in batch.rollouts - ) + num_tokens = sum(r.total_tokens for r in batch.rollouts) self.progress.total_tokens += num_tokens self.progress.total_samples += num_rollouts self.progress.total_problems += num_unique_examples @@ -661,6 +616,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: self.train_sink.reset_pre_filter_stats() self.progress.step += 1 self.maybe_trigger_eval(self.progress.step) + trim_process_memory() def maybe_trigger_eval(self, step: int) -> None: """Fire eligible eval epochs and flip to ``PREFER_EVAL`` if anything @@ -758,13 +714,14 @@ def log_train_batch(self, batch: TrainBatch, *, step: int, step_time: float) -> trainable_rate = (n_trainable / n_survivors) if n_survivors else 0.0 reward_mean = sum(r.reward for r in batch.rollouts) / max(n_survivors, 1) max_off_policy = max((r.off_policy_steps for r in batch.rollouts), default=0) - turns_mean = sum(len(r.raw.get("trajectory") or []) for r in batch.rollouts) / max(n_survivors, 1) + turns_mean = sum(r.num_turns for r in batch.rollouts) / max(n_survivors, 1) + branches_mean = sum(r.num_branches for r in batch.rollouts) / max(n_survivors, 1) truncation_rate = sum(1 for r in batch.rollouts if r.is_truncated) / max(n_survivors, 1) head = ( f"Step {step} | {format_time(step_time):>7} | Reward {reward_mean:.4f} | " f"Trainable {n_trainable}/{n_survivors} ({trainable_rate:.1%}) | " - f"Turns {turns_mean:.1f} | Max Off-Policy {max_off_policy} | " + f"Turns {turns_mean:.1f} | Branches {branches_mean:.1f} | Max Off-Policy {max_off_policy} | " f"Error {error_rate:.1%} | Truncation {truncation_rate:.1%}" ) if len(self.train_envs) <= 1: @@ -782,15 +739,12 @@ def log_train_batch(self, batch: TrainBatch, *, step: int, step_time: float) -> env_error_rate = (n_env_errors / n_env_arrivals) if n_env_arrivals else 0.0 env_reward = (sum(r.reward for r in env_rollouts) / len(env_rollouts)) if env_rollouts else 0.0 env_max_off_policy = max((r.off_policy_steps for r in env_rollouts), default=0) - env_turns = ( - sum(len(r.raw.get("trajectory") or []) for r in env_rollouts) / len(env_rollouts) - if env_rollouts - else 0.0 - ) + env_turns = sum(r.num_turns for r in env_rollouts) / len(env_rollouts) if env_rollouts else 0.0 + env_branches = sum(r.num_branches for r in env_rollouts) / len(env_rollouts) if env_rollouts else 0.0 env_truncation = sum(1 for r in env_rollouts if r.is_truncated) / len(env_rollouts) if env_rollouts else 0.0 lines.append( f"╰─ {env_name:<{name_width}} | Ratio {ratio:.1%} | Reward {env_reward:.4f} | " - f"Turns {env_turns:.1f} | Max Off-Policy {env_max_off_policy} | " + f"Turns {env_turns:.1f} | Branches {env_branches:.1f} | Max Off-Policy {env_max_off_policy} | " f"Error {env_error_rate:.1%} | Truncation {env_truncation:.1%}" ) get_logger().success("\n\t\t ".join(lines)) @@ -802,15 +756,14 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: get_logger().warning(f"Eval @ step={batch.step} env={batch.env_name}: no surviving rollouts, skipping log") return - rollout_dicts = [r.to_dict() for r in batch.rollouts] + rollout_dicts = [r.model_dump(mode="json", exclude=ROLLOUT_DUMP_EXCLUDE) for r in batch.rollouts] step_path = get_step_path(get_rollout_dir(self.config.output_dir), batch.step) await asyncio.to_thread( save_rollouts, rollout_dicts, step_path / f"eval_rollouts_{batch.env_name}.jsonl", - exclude_keys={"trajectory"}, ) - self.monitor.log_eval_samples(rollout_dicts, env_name=batch.env_name, step=batch.step) + self.monitor.log_eval_samples(batch.rollouts, env_name=batch.env_name, step=batch.step) policy_versions = {r.policy_version for r in batch.rollouts} policy_version = min(policy_versions) if len(policy_versions) > 1: @@ -825,11 +778,12 @@ async def finalize_eval_batch(self, batch: EvalBatch) -> None: error_rate = ((batch.metrics.n_cancelled + batch.metrics.n_errored) / n_total) if n_total else 0.0 triggered_at = self.eval_triggered_at.pop((batch.env_name, batch.step), None) elapsed = (time.perf_counter() - triggered_at) if triggered_at is not None else 0.0 + branches_mean = sum(r.num_branches for r in batch.rollouts) / len(batch.rollouts) get_logger().success( f"Evaluated {batch.env_name} (Step {batch.step}) | " f"Policy v{policy_version} | {format_time(elapsed):>7} | Reward {batch.metrics.reward_mean:.4f} | " - f"Turns {batch.metrics.num_turns_mean:.1f} | " + f"Turns {batch.metrics.num_turns_mean:.1f} | Branches {branches_mean:.1f} | " f"Error {error_rate:.1%} | Truncation {batch.metrics.truncation_rate:.1%}" ) diff --git a/src/prime_rl/orchestrator/token_export_metrics.py b/src/prime_rl/orchestrator/token_export_metrics.py deleted file mode 100644 index c6bed05fc09..00000000000 --- a/src/prime_rl/orchestrator/token_export_metrics.py +++ /dev/null @@ -1,116 +0,0 @@ -from __future__ import annotations - -import json -import math -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -from prime_rl.utils.logger import get_logger - -TOKEN_EXPORT_METRIC_FIELDS = ("mismatch_kl", "entropy") - - -@dataclass -class _Stats: - total: float = 0.0 - count: int = 0 - maximum: float = -math.inf - - def add(self, value: Any) -> None: - if not isinstance(value, int | float): - return - value = float(value) - if not math.isfinite(value): - return - self.total += value - self.count += 1 - self.maximum = max(self.maximum, value) - - def as_metrics(self, prefix: str) -> dict[str, float]: - if self.count == 0: - return {} - return { - f"{prefix}/mean": self.total / self.count, - f"{prefix}/max": self.maximum, - } - - -def collect_next_token_export_metrics(output_dir: Path, *, last_logged_step: int) -> dict[str, float | int]: - """Log-ready metrics for the oldest unlogged stable step, namespaced under - ``trainer/`` and stamped with ``trainer/step`` (the run step they belong to). - Empty when there is no new stable step.""" - stable_steps = [step for step in _stable_token_export_steps(output_dir) if step > last_logged_step] - if not stable_steps: - return {} - - step = min(stable_steps) - metrics: dict[str, float | int] = { - f"trainer/{key}": value for key, value in collect_token_export_metrics(output_dir, step).items() - } - metrics["trainer/step"] = step - return metrics - - -def collect_token_export_metrics(output_dir: Path, step: int) -> dict[str, float | int]: - step_dir = output_dir / "token_exports" / f"step_{step}" - if not (step_dir / "STABLE").exists(): - return {} - - stats_by_key: dict[str, _Stats] = {} - sequence_count = 0 - skipped_count = 0 - - for export_file in sorted(step_dir.glob("rank_*.jsonl")): - with export_file.open(encoding="utf-8") as file: - for line in file: - if not line.strip(): - continue - # Token exports are best-effort monitoring; a malformed line or a - # length-mismatched record shouldn't abort the orchestrator step. - try: - record = json.loads(line) - values_by_field = {field: _loss_token_values(record, field) for field in TOKEN_EXPORT_METRIC_FIELDS} - except ValueError: - skipped_count += 1 - continue - sequence_count += 1 - for field, values in values_by_field.items(): - for value in values: - stats_by_key.setdefault(field, _Stats()).add(value) - - if skipped_count: - get_logger().warning(f"Skipped {skipped_count} malformed token-export record(s) in {step_dir}") - - metrics: dict[str, float | int] = {} - if sequence_count == 0: - return metrics - for key, stats in sorted(stats_by_key.items()): - metrics.update(stats.as_metrics(key)) - return metrics - - -def _stable_token_export_steps(output_dir: Path) -> list[int]: - token_exports_dir = output_dir / "token_exports" - if not token_exports_dir.exists(): - return [] - - steps = [] - for step_dir in token_exports_dir.glob("step_*"): - if not (step_dir / "STABLE").exists(): - continue - try: - steps.append(int(step_dir.name.removeprefix("step_"))) - except ValueError: - continue - return sorted(steps) - - -def _loss_token_values(record: dict[str, Any], field: str) -> list[Any]: - loss_mask = record.get("loss_mask", []) - values = record.get(field, []) - if len(loss_mask) != len(values): - raise ValueError( - f"Token export record has mismatched loss_mask and {field} lengths: {len(loss_mask)} != {len(values)}" - ) - return [value for value, keep in zip(values, loss_mask) if keep] diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index f79a0d5eff4..79fdc36a5d6 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -21,12 +21,8 @@ from prime_rl.orchestrator.advantage import assign_advantages from prime_rl.orchestrator.envs import TrainEnvs from prime_rl.orchestrator.filters import RolloutFilter, apply_filters -from prime_rl.orchestrator.trajectories import ( - backfill_rollout_tokens, - interleave_rollout, - offload_images_to_disk, -) -from prime_rl.orchestrator.types import TrainBatch, TrainBatchMetrics, TrainRollout +from prime_rl.orchestrator.trajectories import trace_to_samples +from prime_rl.orchestrator.types import Rollout, TrainBatch, TrainBatchMetrics from prime_rl.transport import TrainingSample from prime_rl.utils.logger import get_logger @@ -39,7 +35,6 @@ def __init__( config: OrchestratorConfig, *, tokenizer, - renderer, train_envs: TrainEnvs, mm_token_type_ids_mapping: dict[int, int] | None, batch_size: int | None, @@ -52,7 +47,6 @@ def __init__( ) self.config = config self.tokenizer = tokenizer - self.renderer = renderer self.train_envs = train_envs self.mm_token_type_ids_mapping = mm_token_type_ids_mapping self.batch_size = batch_size @@ -60,11 +54,15 @@ def __init__( self.pre_filters = pre_filters self.post_filters = post_filters - # Keyed by the dispatcher's group UUID. ``(env_name, example_id)`` - # isn't unique — the same example can be re-sampled while an + # Keyed by the dispatcher's group UUID. ``(env_name, task_idx)`` + # isn't unique — the same task can be re-sampled while an # earlier group is still in flight - self.pending_groups: dict[uuid.UUID, list[TrainRollout]] = defaultdict(list) - self.pending_batch: list[TrainRollout] = [] + self.pending_groups: dict[uuid.UUID, list[Rollout]] = defaultdict(list) + self.pending_batch: list[Rollout] = [] + # Running token total of ``pending_batch`` (token-batched runs), kept in + # sync on append/pop so the readiness check never re-walks the uncached + # ``Trace.total_tokens`` graph property per arrival. + self.pending_tokens: int = 0 # Reset by the orchestrator after each ship via ``reset_pre_filter_stats`` self.pre_filter_seen = 0 @@ -79,13 +77,13 @@ def __init__( def group_size_for(self, env_name: str) -> int: return self.train_envs.get(env_name).config.group_size - def in_progress_groups(self) -> list[list[TrainRollout]]: + def in_progress_groups(self) -> list[list[Rollout]]: """Per-rollout groups currently accumulating in ``pending_groups`` — i.e. groups that haven't hit ``group_size`` yet, so the pipeline log can reflect partial-group progress. Skips group-scoring envs (whose rollouts only make sense as a unit — the user expects per-group fill, not per-rollout, for those).""" - out: list[list[TrainRollout]] = [] + out: list[list[Rollout]] = [] for rollouts in self.pending_groups.values(): if not rollouts: continue @@ -103,11 +101,7 @@ def batch_progress(self) -> tuple[int, int, str]: if self.batch_size is not None: return len(self.pending_batch), self.batch_size, "rollouts" assert self.token_batch_size is not None - tokens = sum( - r.raw["token_usage"]["final_input_tokens"] + r.raw["token_usage"]["final_output_tokens"] - for r in self.pending_batch - ) - return tokens, self.token_batch_size, "tokens" + return self.pending_tokens, self.token_batch_size, "tokens" def buffered_count(self) -> int: """Rollouts that have arrived but sit in not-yet-complete groups @@ -122,13 +116,13 @@ def pending_batch_by_env(self) -> dict[str, int]: counts[r.env_name] += 1 return dict(counts) - async def add(self, rollout: TrainRollout) -> TrainBatch | None: + async def add(self, rollout: Rollout) -> TrainBatch | None: """Process one arrival; finalize the group on the ``group_size``-th arrival; return a ``TrainBatch`` if the batch threshold is met.""" await self.process_rollout(rollout) env_name = rollout.env_name self.arrivals_by_env[env_name] += 1 - if rollout.error is not None: + if rollout.has_error: self.errors_by_env[env_name] += 1 self.pending_groups[rollout.group_id].append(rollout) if len(self.pending_groups[rollout.group_id]) >= self.group_size_for(env_name): @@ -136,37 +130,26 @@ async def add(self, rollout: TrainRollout) -> TrainBatch | None: ready = ( len(self.pending_batch) >= self.batch_size if self.batch_size is not None - else sum( - r.raw["token_usage"]["final_input_tokens"] + r.raw["token_usage"]["final_output_tokens"] - for r in self.pending_batch - ) - >= (self.token_batch_size or 0) + else self.pending_tokens >= (self.token_batch_size or 0) ) if ready: return self.process_batch() return None - async def process_rollout(self, rollout: TrainRollout) -> None: - """Tokenize the rollout eagerly. Backfills tokens if the env didn't - return them (SFT against external teacher APIs); errored rollouts - skip tokenization and get dropped at the group level.""" - if rollout.error is not None: + async def process_rollout(self, rollout: Rollout) -> None: + """Build training samples from the rollout's Trace (one per branch), walking the + message graph. Training is renderer-only across all modes (RL/OPD student, SFT teacher), + so every node already carries its tokens. Errored rollouts are dropped at the group + level, so skip them here.""" + if rollout.has_error: return - raw = rollout.raw - needs_backfill = any(s["tokens"] is None for s in raw.get("trajectory") or []) - if needs_backfill: - await asyncio.to_thread(backfill_rollout_tokens, raw, self.tokenizer, renderer=self.renderer) samples = await asyncio.to_thread( - interleave_rollout, - raw, - mm_token_type_ids_mapping=self.mm_token_type_ids_mapping, + trace_to_samples, + rollout, env_name=rollout.env_name, + mm_token_type_ids_mapping=self.mm_token_type_ids_mapping, ) rollout.samples = samples or [] - # Offload base64 image bytes to disk as soon as the rollout is - # tokenized, so memory stays flat instead of holding every buffered - # rollout's images until the batch ships (no-op for text-only). - await asyncio.to_thread(offload_images_to_disk, [raw], self.config.output_dir) def process_group(self, group_id: uuid.UUID) -> None: """Finalize one GRPO group: drop errored rollouts (the whole group @@ -176,8 +159,8 @@ def process_group(self, group_id: uuid.UUID) -> None: if not group: return env_name = group[0].env_name - example_id = group[0].example_id - survivors = [r for r in group if r.error is None] + task_idx = group[0].task.idx + survivors = [r for r in group if not r.has_error] num_errored = len(group) - len(survivors) # Group-scoring envs: any failure makes survivors' rewards unsafe @@ -185,23 +168,23 @@ def process_group(self, group_id: uuid.UUID) -> None: env = self.train_envs.get(env_name) if num_errored > 0 and env.requires_group_scoring: get_logger().debug( - f"Finished group | env={env_name} example_id={example_id} | " + f"Finished group | env={env_name} task_idx={task_idx} | " f"rollouts={len(group)} (errored={num_errored}) | dropped: group-scored partial" ) return if not survivors: get_logger().debug( - f"Finished group | env={env_name} example_id={example_id} | " + f"Finished group | env={env_name} task_idx={task_idx} | " f"rollouts={len(group)} (errored={num_errored}) | dropped: all failed" ) return - assign_advantages(survivors, self.train_envs.get(env_name).advantage_fn) + assign_advantages(survivors, env.advantage_fn) # Propagate to the pre-tokenized samples so the orchestrator can # collect samples at ship time without re-walking rollouts. The env - # has a single sampling temperature; fan it out across each sample's - # completion tokens here (interleave leaves it empty). + # has a single sampling temperature; fan it out per token (context + # tokens are masked out, so their temperature is don't-care). temperature = env.sampling_args["temperature"] for r in survivors: for sample in r.samples: @@ -209,7 +192,7 @@ def process_group(self, group_id: uuid.UUID) -> None: sample.reward = r.reward sample.env_name = r.env_name sample.training_mode = self.config.training_mode - sample.completion_temperatures = [temperature] * len(sample.completion_ids) + sample.temperatures = [temperature] * len(sample.token_ids) if self.pre_filters: apply_filters(self.pre_filters, survivors) @@ -229,6 +212,8 @@ def process_group(self, group_id: uuid.UUID) -> None: r.filter_results = {} r.is_filtered = False self.pending_batch.append(r) + if self.token_batch_size is not None: + self.pending_tokens += r.total_tokens # Per-group summary. One line per finalized group; per-filter # detection breakdown lives at debug level in ``apply_filters`` @@ -236,7 +221,7 @@ def process_group(self, group_id: uuid.UUID) -> None: avg_reward = sum(rewards) / len(rewards) if rewards else 0.0 filter_str = ", ".join(f"{n}={c}" for n, c in filtered_by_name.items()) if filtered_by_name else "—" get_logger().debug( - f"Finished group | env={env_name} example_id={example_id} | " + f"Finished group | env={env_name} task_idx={task_idx} | " f"rollouts={len(group)} (errored={num_errored}, filtered={num_filtered}) | " f"reward={avg_reward:.4f} | filters: {filter_str}" ) @@ -255,12 +240,13 @@ def process_batch(self) -> TrainBatch: cut = 0 running = 0 for i, r in enumerate(self.pending_batch): - running += r.raw["token_usage"]["final_input_tokens"] + r.raw["token_usage"]["final_output_tokens"] + running += r.total_tokens cut = i + 1 if running >= self.token_batch_size: break cohort = self.pending_batch[:cut] self.pending_batch = self.pending_batch[cut:] + self.pending_tokens -= running if self.post_filters: apply_filters(self.post_filters, cohort) @@ -278,8 +264,8 @@ def process_batch(self) -> TrainBatch: prefill = 0 decode = 0 for sample in r.samples: - sample_decode = sum(sample.completion_mask) - sample_prefill = len(sample.prompt_ids) + len(sample.completion_mask) - sample_decode + sample_decode = sum(sample.mask) + sample_prefill = len(sample.token_ids) - sample_decode decode += sample_decode prefill += sample_prefill if not r.is_filtered: diff --git a/src/prime_rl/orchestrator/train_source.py b/src/prime_rl/orchestrator/train_source.py index db439f7539e..35c8165a89c 100644 --- a/src/prime_rl/orchestrator/train_source.py +++ b/src/prime_rl/orchestrator/train_source.py @@ -14,7 +14,7 @@ class TrainSource: """``next_example(available_permits)`` picks a weighted-RR env and returns its next example (or ``None`` when the env's per-call permit cost doesn't fit — the dispatch loop retries when permits free up). - Returned dicts carry ``env_name`` + ``example_id``.""" + Returned dicts carry ``env_name`` + ``task_idx``.""" def __init__(self, train_envs: TrainEnvs, *, seed: int | None) -> None: self.rng = random.Random(seed) @@ -28,11 +28,9 @@ def __init__(self, train_envs: TrainEnvs, *, seed: int | None) -> None: # per-rollout envs need 1 self.env_costs: dict[str, int] = {} for env in self.envs: - rows: list[dict] = [] - for row in env.get_dataset(seed=seed): - ex = dict(row) - ex["env_name"] = env.name - rows.append(ex) + # The orchestrator never loads the env: sample over the task-index + # range the server reported via info() (num_tasks). + rows: list[dict] = [{"task_idx": i, "env_name": env.name} for i in range(env.num_tasks)] self.rng.shuffle(rows) self.examples[env.name] = rows self.cursors[env.name] = 0 diff --git a/src/prime_rl/orchestrator/trajectories.py b/src/prime_rl/orchestrator/trajectories.py index 3e8431c12a8..ff9d3ef3e20 100644 --- a/src/prime_rl/orchestrator/trajectories.py +++ b/src/prime_rl/orchestrator/trajectories.py @@ -1,652 +1,115 @@ -import base64 -import hashlib -from pathlib import Path -from typing import Any - -import numpy as np -import pybase64 -import torch -import verifiers as vf -from transformers.tokenization_utils import PreTrainedTokenizer - -from prime_rl.transport import RoutedExperts, TrainingSample -from prime_rl.utils.chat_template import ( - common_prefix_len, - deserialize_tool_calls, - normalize_messages, - render_messages, - strip_message_content, -) -from prime_rl.utils.logger import get_logger - -# We use list() instead of deepcopy() for flat lists (token IDs, logprobs) - safe because -# primitives are immutable. mm_kwargs payloads are not mutated after creation. - - -def align_routed_experts( - routed_experts: np.ndarray | None, - expected_len: int, -) -> np.ndarray | None: - """Align routed_experts length with the expected token count. - - VLLM's capturer uses `num_tokens - 1` slot mappings because the final - generated token was never fed as input to a forward pass and has no - routing decision. Append zero-filled entries for the missing positions. - """ - if routed_experts is None: - return routed_experts - assert routed_experts.ndim == 3 - if routed_experts.shape[0] > expected_len: - return np.ascontiguousarray(routed_experts[:expected_len]) - deficit = expected_len - routed_experts.shape[0] - if deficit <= 0: - return routed_experts - padding = np.zeros((deficit, routed_experts.shape[1], routed_experts.shape[2]), dtype=routed_experts.dtype) - return np.concatenate((routed_experts, padding), axis=0) - - -def _common_prefix_len(a: list[int], b: list[int]) -> int: - return common_prefix_len(a, b) - - -def _normalize_messages(messages: Any, default_role: str) -> list[dict[str, Any]]: - return normalize_messages(messages, default_role) +"""Turn a v1 `Trace` (the env server's native, typed output) into training data. +The orchestrator holds a real `vf.Trace` (validated in `envs.py`), so everything here is +attribute access — no dicts. The trace is a message graph (`trace.nodes`); each `trace.branches` +entry (a root→leaf path) is first-class and carries its own flat token sequence +(`branch.token_ids` / `branch.sampled_mask` / `branch.logprobs`), so a branch yields one +training sample directly. Token-length readers (`completion_len`, `total_tokens`, `num_turns`) +live on `vf.Trace` itself. -def _deserialize_tool_calls(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: - return deserialize_tool_calls(messages) +Training is renderer-only across every mode (RL/OPD student, SFT teacher), so every node +always carries its tokens — no backfill needed. For multimodal rollouts the branch also carries +the images it introduced (`branch.multi_modal_data`), rebuilt here into the flat `mm_kwargs` / +`mm_token_type_ids` the trainer forwards. +""" +from __future__ import annotations -def _strip_message_content(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: - return strip_message_content(messages) - - -def _render_messages( - tokenizer: PreTrainedTokenizer, - messages: list[dict[str, Any]], - add_generation_prompt: bool = False, - tools: list[dict[str, Any]] | None = None, -) -> list[int]: - return render_messages( - tokenizer, - messages, - add_generation_prompt=add_generation_prompt, - tools=tools, - ) - - -def _tokenize_step_from_messages( - step: vf.TrajectoryStep, - tokenizer: PreTrainedTokenizer, - tools: list[dict[str, Any]] | None = None, -) -> dict[str, Any]: - prompt = _normalize_messages(step.get("prompt"), default_role="user") - completion = _normalize_messages(step.get("completion"), default_role="assistant") - - prompt = _strip_message_content(_deserialize_tool_calls(prompt)) - completion = _strip_message_content(_deserialize_tool_calls(completion)) - - assert all(m.get("role") == "assistant" for m in completion), ( - "Expected all completion messages to be assistant role for SFT distillation, " - f"got roles: {[m.get('role') for m in completion]}" - ) - - all_messages = prompt + completion - prompt_has_assistant_completion = len(completion) > 0 and completion[0].get("role") == "assistant" - prompt_ids = _render_messages( - tokenizer, - prompt, - add_generation_prompt=prompt_has_assistant_completion, - tools=tools, - ) - full_ids = _render_messages( - tokenizer, - all_messages, - tools=tools, - ) - - split_idx = _common_prefix_len(prompt_ids, full_ids) - original_prompt_len = len(prompt_ids) - - prompt_ids = full_ids[:split_idx] - completion_ids = full_ids[split_idx:] - completion_mask = [True] * len(completion_ids) - completion_logprobs = [0.0] * len(completion_ids) +import numpy as np +import verifiers.v1 as vf - return { - "prompt_ids": prompt_ids, - "prompt_mask": [False] * len(prompt_ids), - "completion_ids": completion_ids, - "completion_mask": completion_mask, - "completion_logprobs": completion_logprobs, - "routed_experts": None, - "prompt_prefix_len": split_idx, - "original_prompt_len": original_prompt_len, - } +from prime_rl.transport import TrainingSample +from prime_rl.transport.types import EncodedTensor, RoutedExperts +from prime_rl.utils.logger import get_logger -def _convert_tools_to_oai_format(tool_defs: list) -> list[dict[str, Any]] | None: - """Convert verifiers Tool objects or dicts to OAI function-calling format.""" - if not tool_defs: +def _to_numpy(val) -> np.ndarray: + """A renderer mm item value (torch tensor or numpy array) -> a contiguous numpy array.""" + if hasattr(val, "detach"): # torch tensor + val = val.detach().cpu().numpy() + return np.ascontiguousarray(val) + + +def _encode_mm_kwargs(mm_items: dict[str, list[dict]]) -> dict[str, EncodedTensor] | None: + """Concatenate the branch's per-image renderer items into the flat `mm_kwargs` the trainer + forwards — one `EncodedTensor` per kwarg key (e.g. `pixel_values`, `image_grid_thw`), images + cat'd along dim 0 in branch token order. Model-agnostic: the keys are whatever the processor + emits. Returns None when there are no items.""" + bins: dict[str, list[np.ndarray]] = {} + for items in mm_items.values(): # per modality + for item in items: # per image + for key, val in item.items(): + bins.setdefault(key, []).append(_to_numpy(val)) + encoded: dict[str, EncodedTensor] = {} + for key, arrs in bins.items(): + arr = np.concatenate(arrs, axis=0) + encoded[key] = EncodedTensor(dtype=str(arr.dtype), shape=list(arr.shape), data=arr.tobytes()) + return encoded or None + + +def _encode_routed_experts(arr: np.ndarray | None, num_tokens: int) -> RoutedExperts | None: + """The branch's router-replay array (`[tokens, layers, top_k]`) -> the transport + `RoutedExperts` the trainer replays. Defensively realigns the token axis to `num_tokens` + (the trainer asserts `routed_experts.shape[0] == len(token_ids)`): truncate if longer, + zero-pad the tail if shorter. `Branch.routed_experts` already guarantees alignment, so this + is a backstop.""" + if arr is None: return None - - def _get(tool: Any, key: str) -> Any: - if isinstance(tool, dict): - return tool.get(key) - return getattr(tool, key, None) - - return [ - { - "type": "function", - "function": { - "name": _get(tool, "name"), - "description": _get(tool, "description"), - "parameters": _get(tool, "parameters"), - **({} if _get(tool, "strict") is None else {"strict": _get(tool, "strict")}), - }, - } - for tool in tool_defs - ] + arr = np.ascontiguousarray(arr) + if arr.shape[0] > num_tokens: + arr = arr[:num_tokens] + elif arr.shape[0] < num_tokens: + pad = np.zeros((num_tokens - arr.shape[0], *arr.shape[1:]), dtype=arr.dtype) + arr = np.concatenate([arr, pad], axis=0) + return RoutedExperts(data=arr.tobytes(), shape=list(arr.shape), dtype=str(arr.dtype)) -def _tokenize_step_with_renderer( - step: vf.TrajectoryStep, - renderer, - tools: list[dict[str, Any]] | None = None, -) -> dict[str, Any]: - """Tokenize a trajectory step using a Renderer.""" - from renderers.base import build_trajectory_step - - prompt = _normalize_messages(step.get("prompt"), default_role="user") - completion = _normalize_messages(step.get("completion"), default_role="assistant") - prompt = _strip_message_content(_deserialize_tool_calls(prompt)) - completion = _strip_message_content(_deserialize_tool_calls(completion)) - return build_trajectory_step(renderer, prompt, completion, tools=tools) - - -def backfill_rollout_tokens( - output: vf.RolloutOutput, - tokenizer: PreTrainedTokenizer, - renderer=None, -) -> bool: - """Populate missing step tokens from prompt/completion messages. - - When a renderer is provided, uses it for tokenization (faster, deterministic). - Otherwise falls back to the tokenizer + apply_chat_template path. - """ - if all(step["tokens"] is not None for step in output["trajectory"]): - return True - - logger = get_logger() - tools = _convert_tools_to_oai_format(output.get("tool_defs", [])) - - for step_idx, step in enumerate(output["trajectory"]): - if step["tokens"] is not None: - continue - - if renderer is not None: - step["tokens"] = _tokenize_step_with_renderer(step, renderer, tools=tools) - else: - reconstructed = _tokenize_step_from_messages(step, tokenizer, tools=tools) - if reconstructed["prompt_prefix_len"] < reconstructed["original_prompt_len"]: - logger.debug( - f"Prompt tokenization was non-prefix for example {output['example_id']} step {step_idx}. " - f"Using longest common prefix length {reconstructed['prompt_prefix_len']} " - f"(original prompt had {reconstructed['original_prompt_len']} tokens)." - ) - reconstructed.pop("prompt_prefix_len") - reconstructed.pop("original_prompt_len") - step["tokens"] = reconstructed - - return True - - -def interleave_rollout( - output: vf.RolloutOutput, - mm_token_type_ids_mapping: dict[int, int] | None = None, +def trace_to_samples( + trace: vf.Trace, *, env_name: str = "", -) -> list[TrainingSample] | None: - """ - Convert vf.RolloutOutput to trainable rollouts by interleaving trajectory steps - where the extension property holds. - - When consecutive steps share token prefixes (extension property), they are - merged into a single sample. When extension breaks (e.g., due to context - compaction or a change in control-flow), a new sample is started. - - Supports multi-prefix matching to handle interleaved agents. For example, - [agent1-step1, agent1-step2, agent2-step1, agent1-step3] produces two samples: - agent1 steps merged together, agent2 step separate. - - Returns a list of samples - could be 1 (extension always held) or up to T - (extension never held). - - For VLM models, each renderer-produced trajectory step carries its - per-image processed tensors inline on ``multi_modal_data``; the last - merged step's sidecar covers every image in the sample. + mm_token_type_ids_mapping: dict[int, int] | None = None, +) -> list[TrainingSample]: + """Convert a v1 `Trace` into `TrainingSample`s — one per branch. + + Each `trace.branches` entry is already a flat token sequence (`branch.token_ids` / + `branch.sampled_mask` / `branch.logprobs`), so a sample carries it directly: `mask` marks + the trainable (model-sampled) tokens, the context tokens between completions stay masked + out. On a rollout error the whole completion is masked out. A branch carrying images also + gets `mm_kwargs` (the concatenated pixel tensors) and `mm_token_type_ids` (the renderer's + `mm_token_type_id_map` applied to the branch tokens). Branches with no sampled tokens + (e.g. an openai client carrying none) yield nothing. """ - logger = get_logger() - - trajectory = output["trajectory"] - if len(trajectory) == 0: - error = output.get("error") - stop = output.get("stop_condition") - logger.warning( - f"No trajectory steps for example {output['example_id']} (error={error}, stop={stop}). Skipping rollout." - ) - return None - - has_error = output["error"] is not None - - def prepare_step_tokens(step: vf.TrajectoryStep, step_idx: int) -> dict[str, Any] | None: - tokens = step["tokens"] - if tokens is not None: - routed_experts_payload = tokens.get("routed_experts") - routed_experts = None - routed_experts_start = None - if routed_experts_payload is not None: - decoded_routed_experts = pybase64.b64decode_as_bytearray(routed_experts_payload["data"]) - routed_experts = np.frombuffer(decoded_routed_experts, dtype=np.uint8).reshape( - routed_experts_payload["shape"] - ) - routed_experts_start = routed_experts_payload["start"] - - return { - "prompt_ids": list(tokens["prompt_ids"]), - "prompt_mask": list(map(bool, tokens["prompt_mask"])), - "completion_ids": list(tokens["completion_ids"]), - "completion_mask": list(map(bool, tokens["completion_mask"])), - "completion_logprobs": list(tokens["completion_logprobs"]), - "routed_experts": routed_experts, - "routed_experts_start": routed_experts_start, - # Renderer-emitted multimodal sidecar (placeholders + per-item - # processed tensors). Populated when the rollout went through - # a multimodal-aware renderer (e.g. Qwen3VLRenderer); absent - # for text-only rollouts. - "multi_modal_data": tokens.get("multi_modal_data"), - } - - logger.warning(f"Missing rollout tokens for example {output['example_id']} step {step_idx}.") - return None - - prepared_steps: list[dict[str, Any]] = [] - for step_idx, step in enumerate(trajectory): - prepared = prepare_step_tokens(step, step_idx) - if prepared is None: - return None - prepared_steps.append(prepared) - - # Deferred routed_experts state per sample: O(N) chunk list concatenated - # once at finalize, replacing the prior O(N²) per-extension unpack/repack. - sample_routed_state: dict[int, dict[str, Any]] = {} - routed_prefix_states: dict[int, list[tuple[list[int], list[int], dict[str, Any]]]] = {} - - # Track (prefix_tokens, sample, step_indices) per active sample. step_indices - # is the explicit list of prepared_steps positions merged into this sample — - # non-contiguous when other agents' steps interleave. - active_samples: list[tuple[list[int], TrainingSample, list[int]]] = [] - - def make_sample(tokens: dict[str, Any], step_idx: int) -> TrainingSample: - """Create a new TrainingSample from a trajectory step.""" - if has_error: - completion_mask = [False] * len(tokens["completion_mask"]) - else: - completion_mask = list(tokens["completion_mask"]) - completion_ids = list(tokens["completion_ids"]) - - prompt_ids = list(tokens["prompt_ids"]) - sample = TrainingSample( - prompt_ids=prompt_ids, - prompt_mask=list(tokens["prompt_mask"]), - completion_ids=completion_ids, - completion_mask=completion_mask, - completion_logprobs=list(tokens["completion_logprobs"]), - completion_temperatures=[], - teacher_logprobs=None, - advantage=None, - env_name=env_name, - mm_token_type_ids=None, - routed_experts=None, # deferred — finalized at end of interleave_rollout - ) - # Initialize routed-experts state for this sample. First chunk is the - # raw step routed_experts (no pad, no copy). running_len is the - # cumulative count across chunks; tracked so the boundary fix-up at - # each extension is a no-op append rather than a destructive write. - step_routed = tokens.get("routed_experts") - if step_routed is not None: - routed_start = tokens["routed_experts_start"] - assert routed_start is not None, f"Missing routed_experts_start for step {step_idx}" - chunks: list[np.ndarray] = [] - running_len = 0 - if routed_start > 0: - source_len = routed_start + 1 - assert source_len in routed_prefix_states, ( - f"Missing routed prefix state for step {step_idx}: " - f"routed_start={routed_start}, prompt_len={len(tokens['prompt_ids'])}" - ) - source_state = None - for prompt_ids, completion_ids, candidate_state in routed_prefix_states[source_len]: - prompt_len = len(prompt_ids) - if ( - tokens["prompt_ids"][:prompt_len] == prompt_ids - and tokens["prompt_ids"][prompt_len:source_len] == completion_ids - ): - source_state = candidate_state - break - assert source_state is not None, ( - f"No matching routed prefix for step {step_idx}: " - f"routed_start={routed_start}, prompt_len={len(tokens['prompt_ids'])}" - ) - assert source_state["running_len"] >= routed_start, ( - f"Routed prefix too short for step {step_idx}: " - f"running_len={source_state['running_len']}, routed_start={routed_start}" - ) - remaining = routed_start - for chunk in source_state["chunks"]: - if remaining == 0: - break - take = min(remaining, int(chunk.shape[0])) - chunks.append(chunk[:take]) - remaining -= take - assert remaining == 0, ( - f"Could not reconstruct routed prefix for step {step_idx}: " - f"remaining={remaining}, routed_start={routed_start}" - ) - running_len = routed_start - chunks.append(step_routed) - running_len += int(step_routed.shape[0]) - sample_routed_state[id(sample)] = { - "chunks": chunks, - "running_len": running_len, - } - return sample - - def extend_sample( - sample: TrainingSample, - prefix_len: int, - step_idx: int, - ) -> None: - """Extend an existing sample with a new trajectory step (extension property holds).""" - tokens = prepared_steps[step_idx] - - # Extend with new prompt tokens (mask=False, no gradient) - new_prompt_ids = tokens["prompt_ids"][prefix_len:] - sample.completion_ids.extend(new_prompt_ids) - sample.completion_mask.extend([False] * len(new_prompt_ids)) - sample.completion_logprobs.extend([0.0] * len(new_prompt_ids)) - - # Extend with new completion tokens - completion_ids = tokens["completion_ids"] - sample.completion_ids.extend(completion_ids) - if has_error: - sample.completion_mask.extend([False] * len(tokens["completion_mask"])) - else: - sample.completion_mask.extend(tokens["completion_mask"]) - sample.completion_logprobs.extend(tokens["completion_logprobs"]) - - step_routed = tokens.get("routed_experts") - state = sample_routed_state.get(id(sample)) - if state is not None: - assert step_routed is not None, f"Missing routed experts for routed sample extension at step {step_idx}" - if step_routed is not None: - assert state is not None, f"Unexpected routed experts for unrouted sample at step {step_idx}" - assert tokens["routed_experts_start"] == prefix_len - 1, ( - f"Routed experts delta start mismatch at step {step_idx}: " - f"start={tokens['routed_experts_start']}, expected={prefix_len - 1}, prefix_len={prefix_len}" + has_error = trace.has_error + samples: list[TrainingSample] = [] + for branch in trace.branches: + mask = branch.sampled_mask + if not any(mask): + continue + token_ids = branch.token_ids + mm_kwargs: dict[str, EncodedTensor] | None = None + mm_token_type_ids: list[int] | None = None + mmd = branch.multi_modal_data + if mmd is not None: + mm_kwargs = _encode_mm_kwargs(mmd.mm_items) + mapping = mm_token_type_ids_mapping or {} + mm_token_type_ids = [mapping.get(t, 0) for t in token_ids] + samples.append( + TrainingSample( + token_ids=token_ids, + mask=[m and not has_error for m in mask], + logprobs=branch.logprobs, + temperatures=[], # filled by TrainSink.process_group + teacher_logprobs=None, + advantage=None, + env_name=env_name, + mm_kwargs=mm_kwargs, + mm_token_type_ids=mm_token_type_ids, + routed_experts=_encode_routed_experts(branch.routed_experts, len(token_ids)), ) - # Delta payloads start at prefix_len - 1. Row 0 fills the boundary - # token missing from the previous request; the rest is the new suffix. - if prefix_len > 0: - boundary_chunk = step_routed[:1] - state["chunks"].append(boundary_chunk) - state["running_len"] += 1 - step_routed = step_routed[1:] - new_chunk = step_routed - state["chunks"].append(new_chunk) - state["running_len"] += int(new_chunk.shape[0]) - - first_tokens = prepared_steps[0] - first_prefix = first_tokens["prompt_ids"] + first_tokens["completion_ids"] - first_sample = make_sample(first_tokens, step_idx=0) - active_samples.append((first_prefix, first_sample, [0])) - first_routed_state = sample_routed_state.get(id(first_sample)) - if first_routed_state is not None: - routed_prefix_states.setdefault(len(first_prefix), []).append( - (first_tokens["prompt_ids"], first_tokens["completion_ids"], first_routed_state) ) - - for step_idx, _step in enumerate(trajectory[1:], start=1): - tokens = prepared_steps[step_idx] - step_prompt_ids = tokens["prompt_ids"] - - # Pick the *longest* matching active prefix. With compaction/rollback, - # one active sample's prefix can be a strict prefix of another (e.g. a - # later sample re-generated tokens that overlap an earlier sample's - # prefix). Both would satisfy the slice check; the shorter would - # silently absorb the longer sample's generated tokens as user input. - matched_idx = None - matched_len = -1 - matching_prefix_lens: list[int] = [] - for idx, (prefix_tokens, _, _) in enumerate(active_samples): - pl = len(prefix_tokens) - if step_prompt_ids[:pl] == prefix_tokens: - matching_prefix_lens.append(pl) - if pl > matched_len: - matched_idx = idx - matched_len = pl - - if len(matching_prefix_lens) > 1: - # Ambiguous extension: rare, but reachable via compaction/rollback - # where a new sample's prefix happens to start with an older - # sample's prefix. Longest-match is the correct choice; surface - # the ambiguity so we can audit if it shows up in real rollouts. - logger.warning( - f"Ambiguous prefix match at step {step_idx} for example {output['example_id']}: " - f"{len(matching_prefix_lens)} of {len(active_samples)} active prefixes match " - f"(lens={sorted(matching_prefix_lens)}, step_prompt_len={len(step_prompt_ids)}). " - f"Extending the longest (len={matched_len})." - ) - - if matched_idx is not None: - # Extension holds - merge into matched sample - prefix_tokens, sample, step_indices = active_samples[matched_idx] - extend_sample(sample, len(prefix_tokens), step_idx=step_idx) - new_prefix = tokens["prompt_ids"] + tokens["completion_ids"] - active_samples[matched_idx] = ( - new_prefix, - sample, - step_indices + [step_idx], - ) - routed_state = sample_routed_state.get(id(sample)) - if routed_state is not None: - routed_prefix_states.setdefault(len(new_prefix), []).append( - (tokens["prompt_ids"], tokens["completion_ids"], routed_state) - ) - else: - # No prefix matches - start a new sample - logger.debug( - f"Extension property broke at step {step_idx + 1} for example {output['example_id']}. " - f"Starting new sample (active_prefixes={len(active_samples)}, step_prompt_len={len(step_prompt_ids)})." - ) - new_prefix = tokens["prompt_ids"] + tokens["completion_ids"] - sample = make_sample(tokens, step_idx=step_idx) - active_samples.append((new_prefix, sample, [step_idx])) - routed_state = sample_routed_state.get(id(sample)) - if routed_state is not None: - routed_prefix_states.setdefault(len(new_prefix), []).append( - (tokens["prompt_ids"], tokens["completion_ids"], routed_state) - ) - - # Finalize routed_experts for each sample. One concat per sample (O(N) byte - # work) replaces the previous per-step unpack/concat/repack (O(N²)). The - # boundary entries between steps were already inserted as one-entry chunks - # during extend_sample, so a straight concat is correct. - for _, sample, _ in active_samples: - state = sample_routed_state.get(id(sample)) - if state is None: - continue - chunks = state["chunks"] - if not chunks: - continue - combined = np.concatenate(chunks, axis=0) if len(chunks) > 1 else np.ascontiguousarray(chunks[0]) - expected_len = len(sample.prompt_ids) + len(sample.completion_ids) - combined = align_routed_experts(combined, expected_len) - combined = np.ascontiguousarray(combined) - sample.routed_experts = RoutedExperts( - data=combined.tobytes(), - shape=list(combined.shape), - dtype=str(combined.dtype), + if not samples: + get_logger().warning( + f"No trainable samples (error={has_error}, stop={trace.stop_condition}, num_turns={trace.num_turns})." ) - - # Attach images by concatenating mm_items across every step the - # sample covers. verifiers' ``state_to_output`` ships per-step - # *delta* mm_data (each step contains only items not present in the - # prior step's cumulative set, with multiset-aware dedup), so - # reading the last step alone would miss every earlier-turn image. - # Concat in step order recovers the per-sample cumulative set; - # deduping again here would drop legitimate duplicate placeholders. - for _, sample, step_indices in active_samples: - renderer_mm = _union_step_mm_data(prepared_steps, step_indices) - if renderer_mm is not None: - mm_kwargs = _pack_mm_kwargs_from_renderer(renderer_mm) - if mm_kwargs is not None: - sample.mm_kwargs = mm_kwargs - # ``mm_token_type_ids``: 1 for image-placeholder tokens, 2 - # for video, 0 otherwise. Renderer-supplied via - # ``mm_token_type_id_map`` (single source of truth). - if mm_token_type_ids_mapping is not None: - sample.mm_token_type_ids = [ - mm_token_type_ids_mapping.get(token_id, 0) - for token_id in sample.prompt_ids + sample.completion_ids - ] - - return [sample for _, sample, _ in active_samples] - - -def _union_step_mm_data( - prepared_steps: list[dict[str, Any]], - step_indices: list[int], -) -> "dict[str, Any] | None": - """Concatenate renderer-emitted mm_items across this sample's owned steps. - - ``step_indices`` lists exactly the ``prepared_steps`` positions merged into - the sample — explicit, not a range, so interleaved-agent trajectories skip - steps owned by other agents. - - Verifiers ≥ c7731bbb ships per-step *delta* mm_data instead of - cumulative — see ``verifiers/utils/save_utils.py::_delta_intermediate_mm_data``. - The cross-step dedup is already done there with multiset semantics - (preserving multiplicity for an image that appears in multiple - placeholder runs in the token stream). We just concatenate in step - order to recover the per-sample cumulative; deduping again here - would drop legitimate duplicate placeholders. - """ - union_items: dict[str, list] = {} - union_hashes: dict[str, list] = {} - has_any = False - for i in step_indices: - mm = prepared_steps[i].get("multi_modal_data") - if mm is None: - continue - items = mm.mm_items if hasattr(mm, "mm_items") else (mm or {}).get("mm_items") or {} - hashes = mm.mm_hashes if hasattr(mm, "mm_hashes") else (mm or {}).get("mm_hashes") or {} - for modality, item_lst in items.items(): - hash_lst = hashes.get(modality, []) or [] - for j, item in enumerate(item_lst or []): - h = hash_lst[j] if j < len(hash_lst) else None - union_items.setdefault(modality, []).append(item) - union_hashes.setdefault(modality, []).append(h) - has_any = True - if not has_any: - return None - return {"mm_items": union_items, "mm_hashes": union_hashes} - - -def _pack_mm_kwargs_from_renderer(mm_data: Any) -> "dict[str, Any] | None": - """Batch the renderer's per-image ``mm_items`` into model-agnostic - forward kwargs. - - ``mm_data`` may arrive as a ``MultiModalData`` instance (in-process - for tests) or as a plain dict (after msgpack round-trip from the - env-worker). Each item is a dict keyed by the names the model's - ``forward`` expects (``pixel_values`` + ``image_grid_thw`` for - Qwen3-VL, just ``pixel_values`` for Gemma3-VL, etc.). We batch by - ``torch.cat(..., dim=0)`` per key — generic because every HF VLM - processor emits a leading batch/patch dimension, and the renderer - always processes one image per call. - - Returns a dict of ``EncodedTensor`` payloads keyed by kwarg name, - or ``None`` when no multimodal data is present. - """ - from verifiers.utils.serve_utils import decode_tensor_payload - - from prime_rl.transport.types import EncodedTensor - - mm_items = mm_data.mm_items if hasattr(mm_data, "mm_items") else (mm_data or {}).get("mm_items") or {} - # Flatten across modalities into one kwarg dict — the model's - # forward signature is the schema. ``mm_items`` is typically - # ``{"image": [...], "video": [...]}`` but each modality's keys - # don't collide for any HF VLM we ship today. - per_kwarg: dict[str, list] = {} - for _modality, items in mm_items.items(): - for item in items or []: - for key, payload in item.items(): - per_kwarg.setdefault(key, []).append(decode_tensor_payload(payload)) - if not per_kwarg: - return None - out: dict[str, EncodedTensor] = {} - for key, tensors in per_kwarg.items(): - cat = torch.cat(tensors, dim=0).contiguous() - arr = cat.detach().cpu().numpy() - out[key] = EncodedTensor( - dtype=str(arr.dtype), - shape=list(arr.shape), - data=arr.tobytes(), - ) - return out - - -_FILE_URL_PREFIX = "file://" - - -def offload_images_to_disk(rollouts: list[vf.RolloutOutput], output_dir: Path) -> int: - """Replace base64 image data in rollout trajectories with file paths on disk. - - Scans all trajectory step prompts for data:image URLs, writes the decoded - image bytes to ``{output_dir}/assets/images/{hash}.png``, and replaces the - URL in-place with ``file://{path}``. Deduplicates by content hash so each - unique image is written only once. - - Returns the number of unique images written to disk. - """ - images_dir = output_dir / "assets" / "images" - images_dir.mkdir(parents=True, exist_ok=True) - - written: set[str] = set() - - for output in rollouts: - for step in output.get("trajectory", []): - prompt = step.get("prompt") - if not prompt or not isinstance(prompt, list): - continue - for msg in prompt: - content = msg.get("content", []) - if not isinstance(content, list): - continue - for item in content: - if item.get("type") != "image_url": - continue - url = item.get("image_url", {}).get("url", "") - if not url.startswith("data:image"): - continue - b64_data = url.split(",", 1)[1] - content_hash = hashlib.sha256(b64_data.encode()).hexdigest()[:16] - path = images_dir / f"{content_hash}.png" - if content_hash not in written: - if not path.exists(): - path.write_bytes(base64.b64decode(b64_data)) - written.add(content_hash) - item["image_url"]["url"] = f"{_FILE_URL_PREFIX}{path}" - - return len(written) + return samples diff --git a/src/prime_rl/orchestrator/types.py b/src/prime_rl/orchestrator/types.py index 54c3f827b9a..7d785e32275 100644 --- a/src/prime_rl/orchestrator/types.py +++ b/src/prime_rl/orchestrator/types.py @@ -3,10 +3,12 @@ from __future__ import annotations import uuid -from dataclasses import dataclass, field, fields -from typing import Literal, Protocol +from dataclasses import dataclass, field +from typing import Generic, Literal, Protocol -import verifiers as vf +import verifiers.v1 as vf +from pydantic import ConfigDict, Field +from verifiers.v1.task import TaskT from prime_rl.transport import TrainingSample @@ -55,7 +57,7 @@ class GroupState: kind: RolloutKind env_name: str - example: dict + task_idx: int rollouts_to_schedule: int target_rollouts: int emitted: int = 0 @@ -64,60 +66,25 @@ class GroupState: policy_version_at_start: int = 0 -@dataclass -class FinishedRollout: - """A completed rollout the sink receives. ``raw`` is the env's untouched - ``vf.RolloutOutput``; prime-rl metadata lives on typed fields. Train vs - eval is discriminated via ``isinstance``. ``rollout_id`` is the only - safe key for tracing one rollout — ``(env_name, example_id)`` collides - on re-sampling and ``group_id`` covers a whole group.""" - - raw: vf.RolloutOutput - env_name: str - example_id: int | str - group_id: uuid.UUID - policy_version: int - off_policy_steps: int - rollout_id: uuid.UUID = field(default_factory=uuid.uuid4) - - @property - def error(self) -> dict | None: - return self.raw.get("error") - - @property - def reward(self) -> float: - return float(self.raw.get("reward", 0.0)) - - @property - def is_truncated(self) -> bool: - return bool(self.raw.get("is_truncated", False)) - - def to_dict(self) -> vf.RolloutOutput: - """``raw`` + metadata merged for I/O (``save_rollouts``, - ``monitor.log_samples``). Shallow copy; never mutates ``self.raw``.""" - out: vf.RolloutOutput = dict(self.raw) # type: ignore[assignment] - for f in fields(self): - if f.name in ("raw", "samples"): - continue - val = getattr(self, f.name) - if f.name == "filter_results": - out["filters"] = dict(val) - continue - out[f.name] = str(val) if isinstance(val, uuid.UUID) else val - return out +class Rollout(vf.Trace[TaskT], Generic[TaskT]): + """A completed rollout: the env's typed ``vf.Trace`` *is* the rollout — prime-rl's + orchestration metadata lives on it directly (set by the dispatcher once the rollout + returns), so there's no wrapper. Train vs eval is the ``kind`` discriminator. All metadata + fields are ``exclude=True``, so dumping a Rollout yields a plain trace — the on-disk + ``results.jsonl`` is unchanged.""" + model_config = ConfigDict(arbitrary_types_allowed=True) # ``samples`` holds msgspec structs -@dataclass -class TrainRollout(FinishedRollout): - samples: list[TrainingSample] = field(default_factory=list) - advantage: float | None = None - is_filtered: bool = False - filter_results: dict[str, bool] = field(default_factory=dict) - - -@dataclass -class EvalRollout(FinishedRollout): - eval_step: int = 0 + kind: RolloutKind = Field(default="train", exclude=True) + env_name: str = Field(default="", exclude=True) + group_id: uuid.UUID = Field(default_factory=uuid.uuid4, exclude=True) + policy_version: int = Field(default=0, exclude=True) + off_policy_steps: int = Field(default=0, exclude=True) + samples: list[TrainingSample] = Field(default_factory=list, exclude=True) + advantage: float | None = Field(default=None, exclude=True) + is_filtered: bool = Field(default=False, exclude=True) + filter_results: dict[str, bool] = Field(default_factory=dict, exclude=True) + eval_step: int | None = Field(default=None, exclude=True) @dataclass @@ -142,7 +109,7 @@ class TrainBatch: """``samples`` is the trainer-bound payload (post-filter survivors); ``rollouts`` is the full cohort kept for orchestrator-side I/O.""" - rollouts: list[TrainRollout] + rollouts: list[Rollout] samples: list[TrainingSample] metrics: TrainBatchMetrics @@ -197,7 +164,7 @@ class EvalBatch: env_name: str step: int - rollouts: list[EvalRollout] + rollouts: list[Rollout] metrics: EvalBatchMetrics diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 9c68f361dea..62671c1495a 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -2,20 +2,19 @@ import ctypes import gc import logging +import os import time from concurrent.futures import ThreadPoolExecutor from itertools import cycle from pathlib import Path import orjson -import verifiers as vf -from verifiers.utils.client_utils import setup_openai_client -from verifiers.utils.save_utils import make_serializable +import verifiers.v1 as vf from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.transport import TrainingSample from prime_rl.utils.client import setup_inference_pool -from prime_rl.utils.logger import InterceptHandler, get_logger +from prime_rl.utils.logger import InterceptHandler, get_logger, setup_logger from prime_rl.utils.utils import ( get_broadcast_dir, get_ckpt_dir, @@ -24,63 +23,38 @@ async def setup_student_inference_pool(*, config: OrchestratorConfig, tokenizer): - """Build the student inference pool + matching renderer. Returns - ``(renderer | None, inference_pool)``; ``renderer`` is ``None`` on the - MITO path (``config.renderer is None``).""" + """Build the student renderer and inference pool, returning ``(renderer, inference_pool)``. + + Training is renderer-only: RL/OPD roll out through the env server's renderer client + (token-in/out), and SFT — which rolls out against a chat-completions teacher that returns + no tokens — re-renders the conversation with this renderer to backfill them. The renderer + is built here from the (always-set) ``config.renderer`` and also supplies the multimodal + token-type-id map. The eval client is plain chat-completions (eval traces aren't trained).""" from renderers.base import create_renderer client_config = config.student.client model_name = config.student.model.name - - if config.renderer is not None: - renderer = create_renderer(tokenizer, config.renderer) - get_logger().info(f"Initialized {type(renderer).__name__} for {model_name}") - inference_pool = await setup_inference_pool( - client_config, - model_name=model_name, - train_client_type="renderer", - eval_client_type="openai_chat_completions", - renderer_config=config.renderer, - pool_size=config.pool_size, - ) - get_logger().info("Using direct renderer rollout client") - return renderer, inference_pool - - get_logger().info("Using MITO (openai_chat_completions) for rollouts") + renderer = create_renderer(tokenizer, config.renderer) + get_logger().info("Using renderer rollout client") inference_pool = await setup_inference_pool( client_config, model_name=model_name, - train_client_type="openai_chat_completions", + train_client_type="renderer", eval_client_type="openai_chat_completions", + renderer_config=config.renderer, + pool_size=config.pool_size, ) - return None, inference_pool - - -def get_model_completion_len(output: vf.RolloutOutput) -> int: - """Sum of model-generated completion tokens across all turns (excludes - environment-injected tokens between turns).""" - return sum(len(step["tokens"]["completion_ids"]) for step in output["trajectory"] if step.get("tokens")) + return renderer, inference_pool -def get_tool_response_len(output: vf.RolloutOutput) -> int: - """Total tool-response tokens consumed across the whole rollout, read from a - harness-emitted metric (e.g. RLM's `rlm_total_tool_response_tokens`, deduped - across turns/branches/sub-RLMs). Returns 0 when no such metric is present.""" - metrics = output.get("metrics") or {} - for key, value in metrics.items(): - if key.endswith("total_tool_response_tokens") and isinstance(value, (int, float)): - return int(value) - return 0 - - -def save_rollouts(rollouts: list[vf.RolloutOutput], path: Path, exclude_keys: set[str] | None = None) -> None: - """Save rollouts to a JSONL file using verifiers serialization.""" +def save_rollouts(rollouts: list[dict], path: Path, exclude_keys: set[str] | None = None) -> None: + """Save rollouts (Trace dicts, already JSON-serializable) to a JSONL file.""" path.parent.mkdir(parents=True, exist_ok=True) opts = orjson.OPT_APPEND_NEWLINE | orjson.OPT_SERIALIZE_NUMPY with open(path, "wb") as f: for rollout in rollouts: row = {k: v for k, v in rollout.items() if k not in exclude_keys} if exclude_keys else rollout - f.write(orjson.dumps(row, default=make_serializable, option=opts)) + f.write(orjson.dumps(row, default=str, option=opts)) def intercept_vf_logging(logger: str = "verifiers", level: str = "DEBUG", prefix: str | None = None): @@ -92,6 +66,15 @@ def intercept_vf_logging(logger: str = "verifiers", level: str = "DEBUG", prefix vf_logger.propagate = False +def setup_env_server_logging(log_level: str, json_logging: bool = False) -> None: + """Configure logging for an env-server process: prime-rl's logger + routing v1's stdlib + logs through it. Passed to verifiers' ``serve_env`` so it runs in the broker and in every + spawned worker — fresh ``spawn`` processes that otherwise have no handlers and would drop + their per-rollout logs.""" + setup_logger(log_level, json_logging=json_logging) + intercept_vf_logging(logger="verifiers.v1", level=log_level) + + def set_default_executor(max_workers: int = 64) -> None: """Scale the default asyncio thread pool so asyncio.to_thread has enough capacity.""" get_logger().info(f"Setting default executor to ThreadPoolExecutor(max_workers={max_workers})") @@ -114,10 +97,15 @@ async def compute_teacher_logprobs( ) -> list[list[float]]: """Compute teacher model logprobs for a batch of training samples via prefill.""" import httpx + from openai import AsyncOpenAI from vllm.entrypoints.serve.disagg.protocol import GenerateResponse async def _compute_single(client_config: vf.ClientConfig, sample: TrainingSample) -> list[float]: - client = setup_openai_client(client_config) + client = AsyncOpenAI( + base_url=client_config.base_url, + api_key=os.environ.get(client_config.api_key_var, "EMPTY"), + default_headers=client_config.headers or None, + ) # Two escape hatches from ``AsyncOpenAI.post``: # 1. URL — ``/inference/v1/generate`` is mounted at server root, not @@ -136,7 +124,7 @@ async def _compute_single(client_config: vf.ClientConfig, sample: TrainingSample cast_to=httpx.Response, body={ "model": model_name, - "token_ids": list(sample.prompt_ids) + list(sample.completion_ids), + "token_ids": list(sample.token_ids), "sampling_params": { "max_tokens": 1, "temperature": 1.0, diff --git a/src/prime_rl/templates/_launch_router.sh.j2 b/src/prime_rl/templates/_launch_router.sh.j2 index 03f02052aeb..62ae7da599c 100644 --- a/src/prime_rl/templates/_launch_router.sh.j2 +++ b/src/prime_rl/templates/_launch_router.sh.j2 @@ -57,6 +57,7 @@ LLMD_EOF local tag="vllm-router"; [ "$mode" = pd ] && tag="vllm-router (PD)" echo "Starting $tag on $LOCAL_IP:$port${replica:+ for replica $replica}" | tee "$log" local -a cmd=(vllm-router --policy "$policy" --host 0.0.0.0 --port "$port" + --request-id-headers x-session-id --worker-startup-timeout-secs 4200 --log-level debug) if [ "$mode" = pd ]; then cmd+=(--vllm-pd-disaggregation $worker_args) diff --git a/src/prime_rl/trainer/batch.py b/src/prime_rl/trainer/batch.py index 2c39a83e4cc..6da2ca3f8d3 100644 --- a/src/prime_rl/trainer/batch.py +++ b/src/prime_rl/trainer/batch.py @@ -2,8 +2,10 @@ from collections.abc import Callable, Sequence from dataclasses import dataclass +import numpy as np + from prime_rl.trainer.utils import balanced_partition -from prime_rl.transport.types import MicroBatch, RoutedExperts, TrainingSample +from prime_rl.transport.types import EncodedTensor, MicroBatch, RoutedExperts, TrainingSample ROUTED_EXPERTS_DTYPE_ITEMSIZE = { "uint8": 1, @@ -41,26 +43,77 @@ def _pad_routed_experts(micro_batch: MicroBatch, padding_size: int) -> None: routed_experts.shape[0] += padding_size +def _slice_encoded(tensor: EncodedTensor, n_rows: int) -> EncodedTensor: + """First `n_rows` rows of a dim-0-stacked encoded tensor (e.g. pixel_values, image_grid_thw).""" + row = int(np.prod(tensor.shape[1:])) if len(tensor.shape) > 1 else 1 + itemsize = np.dtype(tensor.dtype).itemsize + return EncodedTensor( + dtype=tensor.dtype, + shape=[n_rows, *tensor.shape[1:]], + data=tensor.data[: n_rows * row * itemsize], + ) + + +def _truncate_mm( + mm_token_type_ids: list[int], mm_kwargs: dict[str, EncodedTensor], seq_len: int +) -> tuple[int, dict[str, EncodedTensor] | None]: + """Truncating a sample must not split an image's placeholder block, else the surviving image + token count no longer matches the image embeddings in `mm_kwargs`. Returns the cut point + (<= seq_len, never inside an image block) and `mm_kwargs` sliced to the images whose + placeholders fully survive (None if no image survives).""" + grid = np.frombuffer(bytearray(mm_kwargs["image_grid_thw"].data), dtype=mm_kwargs["image_grid_thw"].dtype).reshape( + mm_kwargs["image_grid_thw"].shape + ) + patches_per_image = [int(g.prod()) for g in grid] + total_patches = mm_kwargs["pixel_values"].shape[0] + total_tokens = sum(1 for t in mm_token_type_ids if t) + ppt = total_patches // total_tokens if total_tokens else 1 # patches per token (merge^2) + tokens_per_image = [p // ppt for p in patches_per_image] + + surviving = sum(1 for t in mm_token_type_ids[:seq_len] if t) + kept = acc = 0 + for n in tokens_per_image: + if acc + n > surviving: + break + acc += n + kept += 1 + if acc == surviving: + cut = seq_len # surviving image tokens are exactly `kept` whole images + else: + # `surviving` lands inside image `kept`; cut to its first placeholder, dropping it. + seen, cut = 0, seq_len + for i, t in enumerate(mm_token_type_ids): + if t: + seen += 1 + if seen == acc + 1: + cut = i + break + if not kept: + return cut, None + kept_patches = sum(patches_per_image[:kept]) + sliced = {k: _slice_encoded(v, kept if k == "image_grid_thw" else kept_patches) for k, v in mm_kwargs.items()} + return cut, sliced + + def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch: """ Prepare a problem for sequence packing training. Tokenize and prepare tensors. """ - input_ids = training_example.prompt_ids + training_example.completion_ids - loss_mask = training_example.prompt_mask + training_example.completion_mask - inference_logprobs = [0.0] * len(training_example.prompt_ids) + training_example.completion_logprobs + input_ids = training_example.token_ids + loss_mask = training_example.mask + inference_logprobs = training_example.logprobs advantages = [training_example.advantage] * len(input_ids) reward = training_example.reward if training_example.reward is not None else float("nan") rewards = [reward] * len(input_ids) position_ids = list(range(len(input_ids))) mm_token_type_ids = training_example.mm_token_type_ids + mm_kwargs = training_example.mm_kwargs assert training_example.env_name != "all", "env_name='all' is reserved for aggregate metric keys" env_names = [training_example.env_name] * len(input_ids) - # Per-token temperatures: prompt tokens use first completion temp (masked out anyway) - # Default to 1.0 if completion is empty (e.g., model generated only tool calls with no text) - prompt_temp = training_example.completion_temperatures[0] if training_example.completion_temperatures else 1.0 - temperatures = [prompt_temp] * len(training_example.prompt_ids) + training_example.completion_temperatures + # Per-token sampling temperatures (context tokens are masked out, so theirs are don't-care). + temperatures = training_example.temperatures # Teacher logprobs already cover the full sequence (prompt + completion), # computed via prefill in the orchestrator when a teacher model is configured @@ -70,20 +123,25 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch ) if len(input_ids) > seq_len: - input_ids = input_ids[:seq_len] - loss_mask = loss_mask[:seq_len] - inference_logprobs = inference_logprobs[:seq_len] - position_ids = position_ids[:seq_len] - advantages = advantages[:seq_len] - rewards = rewards[:seq_len] - temperatures = temperatures[:seq_len] + # Multimodal: never split an image's placeholder block — cut to a whole-image boundary + # and slice mm_kwargs to match, so image-token count == image-embedding count. + cut = seq_len + if mm_token_type_ids is not None and mm_kwargs is not None: + cut, mm_kwargs = _truncate_mm(mm_token_type_ids, mm_kwargs, seq_len) + input_ids = input_ids[:cut] + loss_mask = loss_mask[:cut] + inference_logprobs = inference_logprobs[:cut] + position_ids = position_ids[:cut] + advantages = advantages[:cut] + rewards = rewards[:cut] + temperatures = temperatures[:cut] if teacher_logprobs is not None: - teacher_logprobs = teacher_logprobs[:seq_len] + teacher_logprobs = teacher_logprobs[:cut] if routed_experts is not None: - routed_experts = _slice_routed_experts(routed_experts, seq_len) + routed_experts = _slice_routed_experts(routed_experts, cut) if mm_token_type_ids is not None: - mm_token_type_ids = mm_token_type_ids[:seq_len] - env_names = env_names[:seq_len] + mm_token_type_ids = mm_token_type_ids[:cut] + env_names = env_names[:cut] assert ( len(input_ids) @@ -124,7 +182,7 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch routed_experts=routed_experts, mm_token_type_ids=mm_token_type_ids, env_names=env_names, - mm_kwargs=training_example.mm_kwargs, + mm_kwargs=mm_kwargs, training_mode=training_example.training_mode, ) diff --git a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py index 710659fd5db..d16102e9544 100644 --- a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py @@ -227,7 +227,11 @@ def forward( position_embeddings = self.rotary_emb(hidden_states, position_ids) for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): - routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None + # .contiguous(): the per-layer slice is a strided view of [tokens, layers, topk] + # (dim-1 stride = layers*topk); torch.compile's MoE kernel asserts a contiguous input. + routed_experts_layer = ( + routed_experts[:, :, layer_idx, :].contiguous() if routed_experts is not None else None + ) hidden_states = decoder_layer( hidden_states, position_embeddings=position_embeddings, diff --git a/src/prime_rl/trainer/rl/packer.py b/src/prime_rl/trainer/rl/packer.py index ae0d9863e14..7b52e6b06e6 100644 --- a/src/prime_rl/trainer/rl/packer.py +++ b/src/prime_rl/trainer/rl/packer.py @@ -160,27 +160,17 @@ def _on_run_data_deleted(self, idx: int, run_id: str) -> None: def _validate_sample(self, sample: TrainingSample) -> tuple[bool, str | None]: """Validate a sample to ensure it won't crash the trainer.""" - sample_length = len(sample.prompt_ids) + len(sample.completion_ids) - if len(sample.prompt_mask) != len(sample.prompt_ids): - return ( - False, - f"Run wrote a sample with prompt mask length != prompt ids length ({len(sample.prompt_mask)} != {len(sample.prompt_ids)})", - ) - if len(sample.completion_mask) != len(sample.completion_ids): - return ( - False, - f"Run wrote a sample with completion mask length != completion ids length ({len(sample.completion_mask)} != {len(sample.completion_ids)})", - ) - if len(sample.completion_logprobs) != len(sample.completion_ids): - return ( - False, - f"Run wrote a sample with completion logprobs length != completion ids length ({len(sample.completion_logprobs)} != {len(sample.completion_ids)})", - ) - if len(sample.completion_temperatures) != len(sample.completion_ids): - return ( - False, - f"Run wrote a sample with completion temperatures length != completion ids length ({len(sample.completion_temperatures)} != {len(sample.completion_ids)})", - ) + sample_length = len(sample.token_ids) + for name, arr in ( + ("mask", sample.mask), + ("logprobs", sample.logprobs), + ("temperatures", sample.temperatures), + ): + if len(arr) != sample_length: + return ( + False, + f"Run wrote a sample with {name} length != token_ids length ({len(arr)} != {sample_length})", + ) if sample_length == 0: return False, "Run wrote a sample with no tokens" if sample_length > self.seq_len: @@ -228,7 +218,7 @@ def _count_tokens(self, threshold: int | None = None) -> int: for sample, step in buffer: if step > current_step: break - tokens += len(sample.prompt_ids) + len(sample.completion_ids) + tokens += len(sample.token_ids) if threshold is not None and tokens >= threshold: return tokens return tokens @@ -265,10 +255,10 @@ def _select_samples_round_robin(self, token_budget: int) -> list[tuple[int, Trai if step > current_step: # Samples from different steps should be consumed later break - tokens_collected += len(sample.prompt_ids) + len(sample.completion_ids) + tokens_collected += len(sample.token_ids) if tokens_collected > token_budget: - if tokens_collected == (len(sample.prompt_ids) + len(sample.completion_ids)): - tokens_collected -= len(sample.prompt_ids) + len(sample.completion_ids) + if tokens_collected == (len(sample.token_ids)): + tokens_collected -= len(sample.token_ids) # This means we have a sample that has more tokens than max seqlen self.buffers[run_idx].popleft() continue @@ -322,7 +312,7 @@ def pack(self): assert steps_by_run[run_idx] == step, "Micro batches for a run must come from a single run step" samples_by_run[run_idx].append(sample) - num_tokens = len(sample.prompt_ids) + len(sample.completion_ids) + num_tokens = len(sample.token_ids) if run_idx in per_run_stats: cur_samples, cur_tokens = per_run_stats[run_idx] per_run_stats[run_idx] = (cur_samples + 1, cur_tokens + num_tokens) diff --git a/src/prime_rl/transport/types.py b/src/prime_rl/transport/types.py index 1666c510927..12aba3b101f 100644 --- a/src/prime_rl/transport/types.py +++ b/src/prime_rl/transport/types.py @@ -24,14 +24,17 @@ class RoutedExperts(msgspec.Struct, array_like=True, gc=False, omit_defaults=Tru # Orchestrator -> Packer class TrainingSample(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): - """A single training example.""" - - prompt_ids: list[int] - prompt_mask: list[bool] - completion_ids: list[int] - completion_mask: list[bool] - completion_logprobs: list[float] - completion_temperatures: list[float] # Per-token temperatures used during generation + """A single training example — one branch of a rollout as a flat token sequence. + + There is no prompt/completion split: an agentic, multi-turn branch interleaves context and + model-sampled spans, so ``mask`` marks which tokens are trainable (model-sampled) and + ``logprobs`` / ``temperatures`` are aligned per token. All four arrays share the length of + ``token_ids``.""" + + token_ids: list[int] + mask: list[bool] + logprobs: list[float] + temperatures: list[float] env_name: str teacher_logprobs: list[float] | None = None advantage: float | None = None @@ -74,9 +77,9 @@ class MicroBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): advantages: list[float] inference_logprobs: list[float] position_ids: list[int] + sequence_lengths: list[int] temperatures: list[float] # Per-token temperatures used during generation env_names: list[str] - sequence_lengths: list[int] teacher_logprobs: list[float] | None = None lora_num_tokens: list[int] | None = None routed_experts: RoutedExperts | None = None diff --git a/src/prime_rl/utils/async_utils.py b/src/prime_rl/utils/async_utils.py index 1de097592f4..49e03500e27 100644 --- a/src/prime_rl/utils/async_utils.py +++ b/src/prime_rl/utils/async_utils.py @@ -1,4 +1,11 @@ +from __future__ import annotations + import asyncio +from collections import deque +from time import perf_counter + +import numpy as np +from pydantic import BaseModel async def safe_cancel(task: asyncio.Task) -> None: @@ -13,3 +20,54 @@ async def safe_cancel(task: asyncio.Task) -> None: async def safe_cancel_all(tasks: list[asyncio.Task]) -> None: """Safely cancels and awaits all asyncio.Tasks.""" await asyncio.gather(*[safe_cancel(task) for task in tasks]) + + +class EventLoopLagMonitor: + """Monitors how busy the main event loop is by timing short sleeps. + + Vendored from verifiers.utils.async_utils (the orchestrator now runs on + v1 and no longer depends on v1 verifiers).""" + + def __init__(self, measure_interval: float = 0.1, max_measurements: int = 1000): + assert measure_interval > 0 and max_measurements > 0 + self.measure_interval = measure_interval + self.max_measurements = max_measurements + self.lags: deque[float] = deque(maxlen=max_measurements) + + async def measure_lag(self) -> float: + next_time = perf_counter() + self.measure_interval + await asyncio.sleep(self.measure_interval) + return perf_counter() - next_time + + async def run(self) -> None: + """Loop measuring event-loop lag; run as a background task.""" + while True: + self.lags.append(await self.measure_lag()) + + +class EventLoopLagStats(BaseModel): + """Snapshot of event-loop lag statistics.""" + + min: float = 0.0 + mean: float = 0.0 + median: float = 0.0 + p90: float = 0.0 + p99: float = 0.0 + max: float = 0.0 + n: int = 0 + + @classmethod + def from_monitor(cls, monitor: EventLoopLagMonitor) -> EventLoopLagStats: + n = len(monitor.lags) + if n == 0: + return cls(n=0) + arr = np.array(monitor.lags) + return cls( + min=float(arr.min()), + mean=float(arr.mean()), + median=float(np.median(arr)), + p90=float(np.percentile(arr, 90)), + p99=float(np.percentile(arr, 99)), + max=float(arr.max()), + n=n, + ) diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 533f6e27111..f2335aa3d6c 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -8,16 +8,17 @@ from typing import Protocol, runtime_checkable import httpx -import verifiers as vf +import verifiers.v1 as vf from httpx import AsyncClient from openai import NotFoundError from renderers import RendererConfig from tenacity import AsyncRetrying, retry, retry_if_exception, stop_after_attempt, stop_after_delay, wait_exponential +from verifiers.v1.clients.config import EvalClientConfig, TrainClientConfig from prime_rl.configs.shared import ClientConfig from prime_rl.utils.logger import get_logger -# Identity tuple used by ``select_train_client`` to key load counts. ``api_base_url`` +# Identity tuple used by ``select_train_client`` to key load counts. ``base_url`` # distinguishes servers; ``X-data-parallel-rank`` distinguishes DP shards within a # server, since the router uses that header to route to specific GPU ranks. ClientIdentity = tuple[str, str | None] @@ -25,7 +26,7 @@ def client_identity(client: vf.ClientConfig) -> ClientIdentity: """Stable identity for load balancing across inference clients.""" - return (client.api_base_url, client.extra_headers.get("X-data-parallel-rank")) + return (client.base_url, client.headers.get("X-data-parallel-rank")) @runtime_checkable @@ -185,42 +186,31 @@ def setup_clients( renderer_model_name: str | None = None, pool_size: int | None = None, ) -> list[vf.ClientConfig]: - clients = [] - client_idx = 0 - # Only forward the renderer config when the client actually uses a - # renderer — MITO/TITO clients ignore it. + """Build v1 client configs (one per base_url × DP rank). ``client_type`` + ``renderer`` → token-in/out (``TrainClientConfig``, with the renderer the env + server should use forwarded as a serialized config so it doesn't fall back to the + default renderer); otherwise plain chat-completions (``EvalClientConfig``).""" + is_renderer = client_type == "renderer" + config_cls = TrainClientConfig if is_renderer else EvalClientConfig renderer_extra: dict = {} - if client_type == "renderer": + if is_renderer: renderer_extra = { - "renderer_config": renderer_config, + "renderer": renderer_config, + "pool_size": pool_size or 1, "renderer_model_name": renderer_model_name, - "renderer_pool_size": pool_size, } env_headers = { k: v for k, v in ((k, os.getenv(v)) for k, v in client_config.headers_from_env.items()) if v is not None } + clients: list[vf.ClientConfig] = [] for base_url in client_config.base_url: for dp_rank in range(client_config.dp_rank_count): headers = {**client_config.headers, **env_headers} if client_config.dp_rank_count > 1: headers["X-data-parallel-rank"] = str(dp_rank) clients.append( - vf.ClientConfig( - client_idx=client_idx, - client_type=client_type, - api_base_url=base_url, - api_key_var=client_config.api_key_var, - timeout=client_config.timeout, - connect_timeout=client_config.connect_timeout, - max_connections=8192, - max_keepalive_connections=8192, - max_retries=10, - extra_headers=headers, - extra_headers_from_state=client_config.extra_headers_from_state, - **renderer_extra, - ) + config_cls(base_url=base_url, api_key_var=client_config.api_key_var, headers=headers, **renderer_extra) ) - client_idx += 1 return clients diff --git a/src/prime_rl/utils/elastic.py b/src/prime_rl/utils/elastic.py index 951b3673c14..2476560c1fc 100644 --- a/src/prime_rl/utils/elastic.py +++ b/src/prime_rl/utils/elastic.py @@ -17,7 +17,7 @@ from typing import Literal import httpx -import verifiers as vf +import verifiers.v1 as vf from httpx import AsyncClient from renderers import RendererConfig @@ -191,8 +191,6 @@ def _rebuild_clients(self) -> None: self._eval_index = 0 url_config = ClientConfig( - timeout=self.client_config.timeout, - connect_timeout=self.client_config.connect_timeout, base_url=urls, api_key_var=self.client_config.api_key_var, headers=self.client_config.headers, @@ -251,7 +249,6 @@ def num_ready_servers(self) -> int: async def _create_admin_client(self, ip: str) -> AsyncClient: url = self._build_url(ip) config = ClientConfig( - timeout=self.client_config.timeout, base_url=[f"{url}/v1"], api_key_var=self.client_config.api_key_var, headers=self.client_config.headers, diff --git a/src/prime_rl/utils/monitor/base.py b/src/prime_rl/utils/monitor/base.py index bc88080451d..f777461816c 100644 --- a/src/prime_rl/utils/monitor/base.py +++ b/src/prime_rl/utils/monitor/base.py @@ -1,8 +1,11 @@ +from __future__ import annotations + import random from abc import ABC, abstractmethod -from typing import Any +from typing import TYPE_CHECKING, Any -import verifiers as vf +if TYPE_CHECKING: + from prime_rl.orchestrator.types import Rollout def sample_items_for_logging(items: list[Any], sample_ratio: float | None) -> list[Any]: @@ -39,11 +42,11 @@ def log(self, metrics: dict[str, Any], step: int) -> None: pass @abstractmethod - def log_samples(self, rollouts: list[vf.RolloutOutput], step: int) -> None: + def log_samples(self, rollouts: list[Rollout], step: int) -> None: pass @abstractmethod - def log_eval_samples(self, rollouts: list[vf.RolloutOutput], env_name: str, step: int) -> None: + def log_eval_samples(self, rollouts: list[Rollout], env_name: str, step: int) -> None: pass @abstractmethod @@ -72,10 +75,10 @@ def log(self, metrics: dict[str, Any], step: int) -> None: else: self.history = [metrics] - def log_samples(self, rollouts: list[vf.RolloutOutput], step: int) -> None: + def log_samples(self, rollouts: list[Rollout], step: int) -> None: pass - def log_eval_samples(self, rollouts: list[vf.RolloutOutput], env_name: str, step: int) -> None: + def log_eval_samples(self, rollouts: list[Rollout], env_name: str, step: int) -> None: pass def save_final_summary(self, filename: str = "final_summary.json") -> None: diff --git a/src/prime_rl/utils/monitor/multi.py b/src/prime_rl/utils/monitor/multi.py index 0f32ea5906c..69f3eec965d 100644 --- a/src/prime_rl/utils/monitor/multi.py +++ b/src/prime_rl/utils/monitor/multi.py @@ -1,10 +1,13 @@ -from typing import Any +from __future__ import annotations -import verifiers as vf +from typing import TYPE_CHECKING, Any from prime_rl.utils.logger import get_logger from prime_rl.utils.monitor.base import Monitor +if TYPE_CHECKING: + from prime_rl.orchestrator.types import Rollout + class MultiMonitor(Monitor): """Monitor that wraps multiple monitors and delegates calls to all of them.""" @@ -26,14 +29,14 @@ def log(self, metrics: dict[str, Any], step: int) -> None: except Exception as e: self.logger.warning(f"Failed to log metrics to {monitor.__class__.__name__}: {e}") - def log_samples(self, rollouts: list[vf.RolloutOutput], step: int) -> None: + def log_samples(self, rollouts: list[Rollout], step: int) -> None: for monitor in self.monitors: try: monitor.log_samples(rollouts=rollouts, step=step) except Exception as e: self.logger.warning(f"Failed to log samples to {monitor.__class__.__name__}: {e}") - def log_eval_samples(self, rollouts: list[vf.RolloutOutput], env_name: str, step: int) -> None: + def log_eval_samples(self, rollouts: list[Rollout], env_name: str, step: int) -> None: for monitor in self.monitors: try: monitor.log_eval_samples(rollouts=rollouts, env_name=env_name, step=step) diff --git a/src/prime_rl/utils/monitor/prime.py b/src/prime_rl/utils/monitor/prime.py index cd1fa8375a9..a3139cddb82 100644 --- a/src/prime_rl/utils/monitor/prime.py +++ b/src/prime_rl/utils/monitor/prime.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import asyncio import io import json @@ -7,12 +9,11 @@ from datetime import datetime, timezone from pathlib import Path from threading import Thread -from typing import Any +from typing import TYPE_CHECKING, Any import httpx import pyarrow as pa import pyarrow.parquet as pq -import verifiers as vf from prime_cli.core.config import Config as PrimeConfig from transformers.tokenization_utils import PreTrainedTokenizer @@ -21,14 +22,8 @@ from prime_rl.utils.logger import get_logger from prime_rl.utils.monitor.base import Monitor, sample_items_for_logging - -def _json(val: Any) -> str: - """JSON-serialize dicts/lists, pass strings through, default to empty string for None.""" - if isinstance(val, str): - return val - if val is None: - return "" - return json.dumps(val) +if TYPE_CHECKING: + from prime_rl.orchestrator.types import Rollout _SAMPLE_SCHEMA = pa.schema( @@ -286,7 +281,7 @@ def log(self, metrics: dict[str, Any], step: int) -> None: }, ) - def log_samples(self, rollouts: list[vf.RolloutOutput], step: int) -> None: + def log_samples(self, rollouts: list[Rollout], step: int) -> None: """Logs rollouts to Prime Intellect API using presigned URLs for direct R2 upload.""" if not self.is_master: return @@ -329,35 +324,35 @@ def log_samples(self, rollouts: list[vf.RolloutOutput], step: int) -> None: f"Initiated samples upload at step {step} to Prime Intellect API in {time.perf_counter() - start_time:.2f}s" ) - def _rollouts_to_parquet_bytes(self, rollouts: list[vf.RolloutOutput], step: int) -> bytes | None: - """Convert rollouts directly to Parquet bytes for upload.""" + def _rollouts_to_parquet_bytes(self, rollouts: list[Rollout], step: int) -> bytes | None: + """Convert rollouts to Parquet bytes for upload. One row per rollout, built from the + message graph: the conversation is the unit (no prompt/completion split — meaningless in + a multi-turn branch), so `completion` carries the main (last) branch's full message list + and `trajectory` carries one message list per branch (`trace.branches`).""" now = datetime.now(timezone.utc) rows = [] for sample_id, rollout in enumerate(rollouts): - prompt = rollout.get("prompt") - completion = rollout.get("completion") - trajectory = rollout.get("trajectory") or [] - if prompt is None or completion is None or not trajectory: + branches = rollout.branches + if not branches: continue + main_messages = [m.model_dump(mode="json") for m in branches[-1].messages] - example_id = rollout.get("example_id") + task_idx = rollout.task.idx try: - problem_id = int(example_id) if example_id is not None else sample_id + problem_id = int(task_idx) if task_idx is not None else sample_id except (TypeError, ValueError): problem_id = sample_id trajectory_data = [ { - "prompt": ts["prompt"], - "completion": ts["completion"], - "reward": ts.get("reward"), - "advantage": ts.get("advantage"), - "extras": ts.get("extras", {}), - "num_input_tokens": len(ts["tokens"]["prompt_ids"]) if ts.get("tokens") else None, - "num_output_tokens": len(ts["tokens"]["completion_ids"]) if ts.get("tokens") else None, + "messages": [m.model_dump(mode="json") for m in branch.messages], + "reward": rollout.reward, + "advantage": rollout.advantage, + "num_input_tokens": branch.prompt_len, + "num_output_tokens": branch.completion_len, } - for ts in trajectory + for branch in branches ] rows.append( @@ -367,19 +362,19 @@ def _rollouts_to_parquet_bytes(self, rollouts: list[vf.RolloutOutput], step: int "tag": "", "problem_id": problem_id, "sample_id": sample_id, - "prompt": json.dumps(prompt), - "completion": json.dumps(completion), + "prompt": "", + "completion": json.dumps(main_messages), "trajectory": json.dumps(trajectory_data), - "answer": rollout.get("answer") or "", - "env_name": rollout.get("env_name") or "", - "task": rollout.get("task") or "", - "info": _json(rollout.get("info")), - "reward": rollout.get("reward"), - "advantage": rollout.get("advantage"), - "metrics": _json(rollout.get("metrics")), - "timing": _json(rollout.get("timing")), - "num_input_tokens": 0, - "num_output_tokens": 0, + "answer": "", + "env_name": rollout.env_name, + "task": rollout.task.model_dump_json(), + "info": json.dumps(rollout.info), + "reward": rollout.reward, + "advantage": rollout.advantage, + "metrics": json.dumps(rollout.metrics), + "timing": rollout.timing.model_dump_json(), + "num_input_tokens": branches[-1].prompt_len, + "num_output_tokens": branches[-1].completion_len, "created_at": now, } ) @@ -492,7 +487,7 @@ async def _confirm_samples_upload(self, step: int, s3_key: str, max_retries: int await asyncio.sleep(delay) return False - def log_eval_samples(self, rollouts: list[vf.RolloutOutput], env_name: str, step: int) -> None: + def log_eval_samples(self, rollouts: list[Rollout], env_name: str, step: int) -> None: pass def log_distributions(self, distributions: dict[str, list[float]], step: int) -> None: diff --git a/src/prime_rl/utils/monitor/wandb.py b/src/prime_rl/utils/monitor/wandb.py index 7935d13a6ae..8f0ac9c173f 100644 --- a/src/prime_rl/utils/monitor/wandb.py +++ b/src/prime_rl/utils/monitor/wandb.py @@ -1,22 +1,42 @@ +from __future__ import annotations + import json import os import sys import time from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any -import verifiers as vf import wandb from transformers.tokenization_utils import PreTrainedTokenizer from wandb.errors import CommError from wandb.sdk.mailbox.mailbox_handle import ServerResponseError from prime_rl.configs.shared import WandbConfig, WandbWithExtrasConfig -from prime_rl.utils.chat_template import deserialize_tool_calls from prime_rl.utils.config import BaseConfig from prime_rl.utils.logger import get_logger from prime_rl.utils.monitor.base import Monitor, sample_items_for_logging +if TYPE_CHECKING: + from prime_rl.orchestrator.types import Rollout + + +def _loggable_task(task) -> str: + """A Table-safe JSON string of the task for sample logging. Image content parts are elided to + a short placeholder — their base64 data bloats the table and breaks wandb Table's nested-type + inference on the variable-length content list (a plain dict would otherwise crash on it).""" + + def elide(obj): + if isinstance(obj, dict): + if obj.get("type") == "image_url": + return {"type": "image_url", "image_url": ""} + return {k: elide(v) for k, v in obj.items()} + if isinstance(obj, list): + return [elide(v) for v in obj] + return obj + + return json.dumps(elide(task.model_dump(mode="json"))) + class WandbMonitor(Monitor): """Logs to Weights and Biases.""" @@ -113,13 +133,13 @@ def init_wandb(max_retries: int): if config is not None and isinstance(config, WandbWithExtrasConfig) and config.log_extras: if config.log_extras.samples: self.last_log_samples_step = -1 - self.samples_cols = ["step", "env_name", "task", "example_id", "messages", "input_ids", "reward"] + self.samples_cols = ["step", "env_name", "task", "task_idx", "messages", "input_ids", "reward"] self.samples_table = wandb.Table( columns=self.samples_cols, log_mode="INCREMENTAL", ) self.tokenizer = tokenizer - self.eval_samples_cols = ["step", "env", "task", "example_id", "completion", "reward"] + self.eval_samples_cols = ["step", "env", "task", "task_idx", "completion", "reward"] self.eval_samples_table = wandb.Table( columns=self.eval_samples_cols, log_mode="INCREMENTAL", @@ -143,7 +163,7 @@ def log(self, metrics: dict[str, Any], step: int) -> None: return wandb.log({**metrics, "step": step}) - def log_samples(self, rollouts: list[vf.RolloutOutput], step: int) -> None: + def log_samples(self, rollouts: list[Rollout], step: int) -> None: """Logs rollouts to W&B table.""" if not self.is_master: return @@ -172,32 +192,30 @@ def log_samples(self, rollouts: list[vf.RolloutOutput], step: int) -> None: start_time = time.perf_counter() for rollout in rollouts: - trajectory = rollout["trajectory"] - if not trajectory: - continue - last_step = trajectory[-1] - tokens = last_step["tokens"] - full_ids = tokens["prompt_ids"] + tokens["completion_ids"] - messages_text = self.tokenizer.decode(full_ids) - sample = { - "step": step, - "env_name": rollout.get("env_name"), - "task": rollout.get("task"), - "example_id": rollout["example_id"], - "messages": messages_text, - "input_ids": str(full_ids), - "reward": rollout["reward"], - } - assert list(sample.keys()) == self.samples_cols, ( - "Order of columns in the table must be the same as order of the keys here" - ) - self.samples_table.add_data(*sample.values()) + trace = rollout + for branch in trace.branches: + token_ids = branch.token_ids + if not token_ids: + continue + sample = { + "step": step, + "env_name": rollout.env_name, + "task": _loggable_task(trace.task), + "task_idx": trace.task.idx, + "messages": self.tokenizer.decode(token_ids), + "input_ids": str(token_ids), + "reward": trace.reward, + } + assert list(sample.keys()) == self.samples_cols, ( + "Order of columns in the table must be the same as order of the keys here" + ) + self.samples_table.add_data(*sample.values()) wandb.log({"samples": self.samples_table, "step": step}) self.last_log_samples_step = step self.logger.debug(f"Logged samples at step {step} to W&B table in {time.perf_counter() - start_time:.2f}s") - def log_eval_samples(self, rollouts: list[vf.RolloutOutput], env_name: str, step: int) -> None: + def log_eval_samples(self, rollouts: list[Rollout], env_name: str, step: int) -> None: """Logs eval rollouts to a separate W&B table.""" if not self.is_master: return @@ -210,23 +228,22 @@ def log_eval_samples(self, rollouts: list[vf.RolloutOutput], env_name: str, step return for rollout in rollouts: - completion = rollout.get("completion") - if not completion: - continue - if isinstance(completion, list): - try: - completion = self.tokenizer.apply_chat_template(deserialize_tool_calls(completion), tokenize=False) - except Exception: - completion = str(completion) - sample = { - "step": step, - "env": env_name, - "task": rollout.get("task"), - "example_id": rollout["example_id"], - "completion": completion, - "reward": rollout["reward"], - } - self.eval_samples_table.add_data(*sample.values()) + trace = rollout + for branch in trace.branches: + # Eval runs the openai client (no token ids), so show the assistant message + # content rather than decoded tokens. + completion = "".join(m.content or "" for m in branch.messages if m.role == "assistant") + if not completion: + continue + sample = { + "step": step, + "env": env_name, + "task": _loggable_task(trace.task), + "task_idx": trace.task.idx, + "completion": completion, + "reward": trace.reward, + } + self.eval_samples_table.add_data(*sample.values()) wandb.log({"eval/samples": self.eval_samples_table, "step": step}) diff --git a/src/prime_rl/utils/utils.py b/src/prime_rl/utils/utils.py index dec33f625c6..64f430e97af 100644 --- a/src/prime_rl/utils/utils.py +++ b/src/prime_rl/utils/utils.py @@ -2,7 +2,6 @@ import functools import importlib import os -import subprocess import sys from collections import defaultdict from contextlib import contextmanager @@ -13,7 +12,6 @@ import torch.distributed as dist import wandb -from prime_rl.configs.orchestrator import EnvConfig, EvalEnvConfig from prime_rl.utils.logger import get_logger # TODO: Change all imports to use utils.pathing @@ -229,31 +227,3 @@ def default_dtype(dtype): yield finally: torch.set_default_dtype(prev) - - -def install_env(env_id: str, prerelease: bool = False) -> None: - """Install an environment in subprocess.""" - logger = get_logger() - logger.info(f"Installing environment {env_id}") - install_cmd = ["uv", "run", "--no-sync", "prime", "env", "install", env_id] - if prerelease: - install_cmd.insert(-1, "--prerelease") - result = subprocess.run(install_cmd, capture_output=True, text=True) - for line in result.stdout.splitlines(): - if line.strip(): - logger.info(line) - for line in result.stderr.splitlines(): - if line.strip(): - logger.warning(line) - if result.returncode != 0: - raise RuntimeError(f"Failed to install environment {env_id} (exit code {result.returncode})") - logger.info(f"Successfully installed environment {env_id}") - - -def get_env_ids_to_install(env_configs: list[EnvConfig] | list[EvalEnvConfig]) -> set[str]: - """Get the list of environment IDs to install.""" - env_ids_to_install = set() - for env_config in env_configs: - if "/" in env_config.id: - env_ids_to_install.add(env_config.id) - return env_ids_to_install diff --git a/tests/unit/orchestrator/test_advantage.py b/tests/unit/orchestrator/test_advantage.py index 0f5fd9bf92d..26d41f1e768 100644 --- a/tests/unit/orchestrator/test_advantage.py +++ b/tests/unit/orchestrator/test_advantage.py @@ -1,6 +1,5 @@ -import uuid - import pytest +import verifiers.v1 as vf from prime_rl.configs.orchestrator import ( CustomAdvantageConfig, @@ -16,17 +15,22 @@ setup_advantage_fn, ) from prime_rl.orchestrator.envs import Env, TrainEnv -from prime_rl.orchestrator.types import TrainRollout - - -def _make_rollout(reward: float, completion_len: int = 0) -> dict: - """Minimal single-turn rollout dict with ``completion_len`` completion tokens.""" - return { - "reward": reward, - "trajectory": [{"tokens": {"prompt_ids": [0], "completion_ids": list(range(completion_len))}}], - "env_name": "test", - "example_id": 0, - } +from prime_rl.orchestrator.types import Rollout + + +def _make_rollout(reward: float, completion_len: int = 0, env_name: str = "test") -> Rollout: + """Build a ``Rollout`` (message-graph trace): ``reward`` via the reward dict and + ``completion_len`` model-sampled completion tokens (a single assistant node).""" + node = vf.MessageNode( + message=vf.AssistantMessage(content="x"), + token_ids=list(range(completion_len)), + mask=[True] * completion_len, + logprobs=[0.0] * completion_len, + sampled=True, + ) + rollout = Rollout[vf.Task](task=vf.Task(idx=0, prompt=""), nodes=[node], rewards={"reward": reward}) + rollout.env_name = env_name + return rollout def _make_group(rewards, completion_lengths=None) -> AdvantageInputs: @@ -38,6 +42,11 @@ def _make_group(rewards, completion_lengths=None) -> AdvantageInputs: return AdvantageInputs(rollouts=rollouts) +def _train_rollouts(rewards: list[float]) -> list[Rollout]: + """One group's worth of Rollouts for ``assign_advantages`` (operates on a single group).""" + return [Rollout[vf.Task](task=vf.Task(idx=0, prompt=""), rewards={"reward": r}) for r in rewards] + + def test_default_advantage_fn_simple_mean(): inputs = _make_group(rewards=[1.0, 0.5, 0.8], completion_lengths=[10, 12, 8]) result = default_advantage_fn(inputs) @@ -105,14 +114,18 @@ def test_train_env_threads_max_seq_len_into_advantage_fn(monkeypatch): def fake_env_init(self, config): self.config = config - self._env = None - self._env_client = None - self._env_server_process = None monkeypatch.setattr(Env, "__init__", fake_env_init) config = OrchestratorConfig( seq_len=100, - advantage={"type": "default", "length_penalty": {"coef": 1.0}}, + train={ + "env": [ + { + "taskset": {"id": "reverse-text-v1"}, + "advantage": {"type": "default", "length_penalty": {"coef": 1.0}}, + } + ] + }, ) env = TrainEnv(config.train.env[0], max_seq_len=config.seq_len) @@ -128,12 +141,10 @@ def test_per_env_linear_advantage_uses_runtime_schema(): train={ "env": [ { - "advantage": { - "type": "default", - "length_penalty": {"coef": 1.0}, - }, + "taskset": {"id": "reverse-text-v1"}, + "advantage": {"type": "default", "length_penalty": {"coef": 1.0}}, } - ], + ] }, ) @@ -152,6 +163,7 @@ def test_per_env_custom_advantage_uses_runtime_schema(): train={ "env": [ { + "taskset": {"id": "reverse-text-v1"}, "advantage": { "type": "custom", "import_path": "tests.unit.orchestrator.test_advantage._dummy_custom_advantage", @@ -183,24 +195,6 @@ def test_length_weighted_baseline(): assert sum(length * adv for length, adv in zip((10, 30, 60), result.advantages)) == pytest.approx(0.0, abs=1e-5) -def _train_rollouts(rewards: list[float]) -> list[TrainRollout]: - """Wrap a list of rewards into ``TrainRollout``\\ s sharing a single - ``group_id`` — ``assign_advantages`` works on one group at a time - (the sink groups by ``group_id`` upstream).""" - gid = uuid.uuid4() - return [ - TrainRollout( - raw={"reward": r, "trajectory": []}, - env_name="test", - example_id=0, - group_id=gid, - policy_version=0, - off_policy_steps=0, - ) - for r in rewards - ] - - def test_assign_advantages_writes_field(): rollouts = _train_rollouts([1.0, 0.5, 0.8]) fn = setup_advantage_fn(DefaultAdvantageConfig()) @@ -240,4 +234,4 @@ def test_setup_advantage_fn_with_custom_config(): def _dummy_custom_advantage(inputs: AdvantageInputs, scale: float = 1.0) -> AdvantageOutputs: """A simple custom advantage for testing.""" - return AdvantageOutputs(advantages=[r["reward"] * scale for r in inputs.rollouts]) + return AdvantageOutputs(advantages=[r.reward * scale for r in inputs.rollouts]) diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index bbe0e617241..58bf154a79f 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -3,13 +3,9 @@ import numpy as np import pytest -from prime_rl.trainer.batch import ( - pad_micro_batch, - prepare_batch, - prepare_sample, -) +from prime_rl.trainer.batch import pad_micro_batch, prepare_batch, prepare_sample from prime_rl.trainer.utils import build_bin_cost -from prime_rl.transport.types import MicroBatch, RoutedExperts, TrainingSample +from prime_rl.transport.types import EncodedTensor, MicroBatch, RoutedExperts, TrainingSample def _routed_experts(data, dtype=np.uint8): @@ -29,12 +25,10 @@ def _make_training_example( env_name: str = "test-env", ) -> TrainingSample: return TrainingSample( - prompt_ids=[1, 2], - prompt_mask=[False, False], - completion_ids=[3, 4], - completion_mask=[True, True], - completion_logprobs=[-0.1, -0.2], - completion_temperatures=[temperature, temperature], # Per-token temperatures + token_ids=[1, 2, 3, 4], + mask=[False, False, True, True], + logprobs=[0.0, 0.0, -0.1, -0.2], + temperatures=[temperature, temperature, temperature, temperature], teacher_logprobs=[0.0, 0.0, 0.0, 0.0], advantage=1.0, env_name=env_name, @@ -48,12 +42,10 @@ def make_sized_training_example(length: int, env_name: str = "test-env") -> Trai assert length >= 1 prompt_len = length - 1 return TrainingSample( - prompt_ids=[1] * prompt_len, - prompt_mask=[False] * prompt_len, - completion_ids=[2], - completion_mask=[True], - completion_logprobs=[-0.1], - completion_temperatures=[1.0], + token_ids=[1] * prompt_len + [2], + mask=[False] * prompt_len + [True], + logprobs=[0.0] * prompt_len + [-0.1], + temperatures=[1.0] * length, advantage=1.0, env_name=env_name, ) @@ -120,35 +112,8 @@ def test_randomized_packing_invariants(): assert not any(batch.advantages) -def test_prepare_batch_packs_different_temperatures(make_training_example): - """With per-token temperatures, samples can be packed together regardless of their temperature values.""" - example1 = make_training_example(temperature=0.7, env_name="env-a") - example2 = make_training_example(temperature=1.1, env_name="env-b") - - batches_per_gpu = prepare_batch( - rollouts=[example1, example2], - seq_len=16, - num_train_workers=1, - idxs=[0, 0], - num_loras=1, - bin_cost=build_bin_cost(None), - ) - - flat_batches = _flatten_batches(batches_per_gpu) - # With per-token temperatures, samples can now be packed together - assert len(flat_batches) == 1 - # Each sample has 4 tokens (2 prompt + 2 completion), so 8 total tokens - assert len(flat_batches[0].temperatures) == 8 - # First sample (4 tokens): all get temp 0.7 - assert flat_batches[0].temperatures[:4] == [0.7, 0.7, 0.7, 0.7] - # Second sample (4 tokens): all get temp 1.1 - assert flat_batches[0].temperatures[4:8] == [1.1, 1.1, 1.1, 1.1] - assert flat_batches[0].env_names == ["env-a"] * 4 + ["env-b"] * 4 - assert flat_batches[0].sequence_lengths == [4, 4] - - -def test_pad_micro_batch_preserves_explicit_sequence_lengths(make_training_example): - micro_batch = prepare_sample(make_training_example(), seq_len=16) +def test_pad_micro_batch_preserves_explicit_sequence_lengths(): + micro_batch = prepare_sample(make_sized_training_example(4), seq_len=16) padded = pad_micro_batch(micro_batch, pad_to_multiple_of=6) @@ -158,25 +123,6 @@ def test_pad_micro_batch_preserves_explicit_sequence_lengths(make_training_examp assert padded.loss_mask[-2:] == [False, False] -def test_prepare_batch_does_not_pack_mixed_training_mode(make_training_example): - rl_example = make_training_example(training_mode="rl") - sft_example = make_training_example(training_mode="sft") - - batches_per_gpu = prepare_batch( - rollouts=[rl_example, sft_example], - seq_len=16, - num_train_workers=1, - idxs=[0, 0], - num_loras=1, - bin_cost=build_bin_cost(None), - ) - - flat_batches = _flatten_batches(batches_per_gpu) - assert len(flat_batches) == 2 - assert {batch.training_mode for batch in flat_batches} == {"rl", "sft"} - assert [batch.sequence_lengths for batch in flat_batches] == [[4], [4]] - - def test_split_to_align_avoids_dummy_micro_batches(): examples = [make_sized_training_example(length) for length in [6, 6, 5, 5, 4, 4]] @@ -244,18 +190,136 @@ def test_flop_aware_split_to_align_splits_heaviest_flop_bin(): assert sum(len(batch.sequence_lengths) > 1 for batch in real_batches) == 3 +def test_training_sample_requires_env_name(): + with pytest.raises(TypeError, match="env_name"): + TrainingSample( + token_ids=[1, 2, 3, 4], + mask=[False, False, True, True], + logprobs=[0.0, 0.0, -0.1, -0.2], + temperatures=[1.0, 1.0, 1.0, 1.0], + advantage=1.0, + ) + + +@pytest.mark.parametrize( + ("rollout_count", "num_train_workers", "expected_batches_per_worker"), [(4, 2, 2), (5, 2, 3), (7, 1, 7), (11, 4, 3)] +) +def test_prepare_batch_balances_micro_batches_across_workers( + make_training_example, rollout_count, num_train_workers, expected_batches_per_worker +): + examples = [make_training_example() for i in range(rollout_count)] + + batches_per_gpu = prepare_batch( + rollouts=examples, + seq_len=4, + num_train_workers=num_train_workers, + idxs=[0] * rollout_count, + num_loras=1, + bin_cost=build_bin_cost(None), + ) + + assert all(len(worker_batches) == expected_batches_per_worker for worker_batches in batches_per_gpu) + + flat_batches = [batch for worker_batches in batches_per_gpu for batch in worker_batches] + assert len(examples) <= len(flat_batches) < len(examples) + num_train_workers + print(flat_batches) + + # Real rollouts and dummy padding can be interleaved across workers by the + # FLOP-balanced partition, so classify by content rather than position. + real_batches = [b for b in flat_batches if any(advantage != 0.0 for advantage in b.advantages)] + assert len(real_batches) == len(examples) + for batch in real_batches: + assert sum(1 for advantage in batch.advantages if advantage != 0.0) == 4 + assert sum(1 for loss_mask in batch.loss_mask if loss_mask) == 2 + + # Dummy padding batches have zero advantages and loss mask. + for batch in flat_batches: + if all(advantage == 0.0 for advantage in batch.advantages): + assert sum(1 for loss_mask in batch.loss_mask if loss_mask) == 0 + + +def test_prepare_batch_packs_different_temperatures(make_training_example): + """With per-token temperatures, samples can be packed together regardless of their temperature values.""" + example1 = make_training_example(temperature=0.7, env_name="env-a") + example2 = make_training_example(temperature=1.1, env_name="env-b") + + batches_per_gpu = prepare_batch( + rollouts=[example1, example2], + seq_len=16, + num_train_workers=1, + idxs=[0, 0], + num_loras=1, + bin_cost=build_bin_cost(None), + ) + + flat_batches = [batch for worker_batches in batches_per_gpu for batch in worker_batches] + # With per-token temperatures, samples can now be packed together + assert len(flat_batches) == 1 + # Each sample has 4 tokens (2 prompt + 2 completion), so 8 total tokens + assert len(flat_batches[0].temperatures) == 8 + # First sample (4 tokens): all get temp 0.7 + assert flat_batches[0].temperatures[:4] == [0.7, 0.7, 0.7, 0.7] + # Second sample (4 tokens): all get temp 1.1 + assert flat_batches[0].temperatures[4:8] == [1.1, 1.1, 1.1, 1.1] + assert flat_batches[0].env_names == ["env-a"] * 4 + ["env-b"] * 4 + + +def test_prepare_sample_propagates_training_mode(make_training_example): + example = make_training_example(training_mode="sft") + + micro_batch = prepare_sample(example, seq_len=16) + + assert micro_batch.training_mode == "sft" + + +def test_prepare_batch_does_not_pack_mixed_training_mode(make_training_example): + rl_example = make_training_example(training_mode="rl") + sft_example = make_training_example(training_mode="sft") + + batches_per_gpu = prepare_batch( + rollouts=[rl_example, sft_example], + seq_len=16, + num_train_workers=1, + idxs=[0, 0], + num_loras=1, + bin_cost=build_bin_cost(None), + ) + + flat_batches = [batch for worker_batches in batches_per_gpu for batch in worker_batches] + assert len(flat_batches) == 2 + assert {batch.training_mode for batch in flat_batches} == {"rl", "sft"} + + +def test_prepare_sample_with_routed_experts(): + """Routed experts are passed through prepare_sample and match input_ids length.""" + # 2 prompt + 2 completion = 4 tokens, 2 layers, topk=2 + routed_experts = [[[0, 1], [2, 3]], [[4, 5], [6, 7]], [[0, 2], [1, 3]], [[1, 0], [3, 2]]] + routed_payload = _routed_experts(routed_experts) + sample = TrainingSample( + token_ids=[1, 2, 3, 4], + mask=[False, False, True, True], + logprobs=[0.0, 0.0, -0.1, -0.2], + temperatures=[1.0, 1.0, 1.0, 1.0], + advantage=1.0, + env_name="test-env", + routed_experts=routed_payload, + ) + + micro_batch = prepare_sample(sample, seq_len=8) + assert micro_batch.routed_experts is not None + assert micro_batch.routed_experts == routed_payload + + def test_prepare_sample_truncates_routed_experts(): """Routed experts are truncated to seq_len when input exceeds it.""" routed_experts = [[[0, 1]], [[2, 3]], [[4, 5]], [[6, 7]]] routed_payload = _routed_experts(routed_experts) expected_payload = _routed_experts(routed_experts[:3]) sample = TrainingSample( - prompt_ids=[1, 2], - prompt_mask=[False, False], - completion_ids=[3, 4], - completion_mask=[True, True], - completion_logprobs=[-0.1, -0.2], - completion_temperatures=[1.0, 1.0], + token_ids=[1, 2, 3, 4], + mask=[False, False, True, True], + logprobs=[0.0, 0.0, -0.1, -0.2], + temperatures=[1.0, 1.0, 1.0, 1.0], advantage=1.0, env_name="test-env", routed_experts=routed_payload, @@ -265,3 +329,55 @@ def test_prepare_sample_truncates_routed_experts(): assert micro_batch.routed_experts is not None assert micro_batch.routed_experts == expected_payload assert micro_batch.env_names == ["test-env"] * 3 + + +def _encoded(arr) -> EncodedTensor: + a = np.asarray(arr) + return EncodedTensor(data=a.tobytes(), shape=list(a.shape), dtype=str(a.dtype)) + + +def test_prepare_sample_truncates_mm_at_image_boundary(): + """Truncation never splits an image's placeholder block: it cuts to a whole-image boundary + and slices mm_kwargs to match, so image-token count stays == image-embedding count.""" + # Two 2-token images (patches-per-token = 1): image-pad at indices 1,2 (img0) and 4,5 (img1). + mm_token_type_ids = [0, 1, 1, 0, 1, 1, 0] + pixel_values = np.array([[1.0], [1.0], [2.0], [2.0]], dtype=np.float32) # img0=1.0, img1=2.0 + grid = np.array([[1, 2, 1], [1, 2, 1]], dtype=np.int64) + sample = TrainingSample( + token_ids=[10, 11, 12, 13, 14, 15, 16], + mask=[False, False, False, False, False, True, True], + logprobs=[0.0] * 7, + temperatures=[1.0] * 7, + advantage=1.0, + env_name="test-env", + mm_token_type_ids=mm_token_type_ids, + mm_kwargs={"pixel_values": _encoded(pixel_values), "image_grid_thw": _encoded(grid)}, + ) + + # seq_len=5 falls inside img1 (one of its two placeholders survives) -> drop img1 entirely. + mb = prepare_sample(sample, seq_len=5) + assert len(mb.input_ids) == 4 # cut back to img1's first placeholder (index 4) + assert len(mb.mm_token_type_ids) == len(mb.input_ids) + n_placeholders = sum(1 for t in mb.mm_token_type_ids if t) + assert n_placeholders == 2 # only img0's two placeholders remain + # No mismatch: placeholders == image embeddings, and only img0's pixels are kept. + assert mb.mm_kwargs["pixel_values"].shape == [2, 1] + assert mb.mm_kwargs["image_grid_thw"].shape == [1, 3] + kept = np.frombuffer(bytearray(mb.mm_kwargs["pixel_values"].data), dtype=np.float32) + assert kept.tolist() == [1.0, 1.0] + assert n_placeholders == mb.mm_kwargs["pixel_values"].shape[0] # ppt == 1 here + + +def test_prepare_sample_none_routed_experts(): + """When routed_experts is None, micro_batch.routed_experts is None.""" + sample = TrainingSample( + token_ids=[1, 2, 3, 4], + mask=[False, False, True, True], + logprobs=[0.0, 0.0, -0.1, -0.2], + temperatures=[1.0, 1.0, 1.0, 1.0], + advantage=1.0, + env_name="test-env", + ) + + micro_batch = prepare_sample(sample, seq_len=8) + assert micro_batch.routed_experts is None diff --git a/tests/unit/orchestrator/test_filters.py b/tests/unit/orchestrator/test_filters.py index 2643bf71bb5..1b195cb0d79 100644 --- a/tests/unit/orchestrator/test_filters.py +++ b/tests/unit/orchestrator/test_filters.py @@ -1,6 +1,8 @@ import math import uuid +import verifiers.v1 as vf + from prime_rl.configs.orchestrator import GibberishFilterConfig, RepetitionFilterConfig from prime_rl.orchestrator.filters import ( GibberishFilter, @@ -9,7 +11,18 @@ setup_filter, setup_filters, ) -from prime_rl.orchestrator.types import TrainRollout +from prime_rl.orchestrator.types import Rollout + + +def _assistant_node(token_ids: list[int], logprobs: list[float]) -> vf.MessageNode: + """An assistant node whose tokens are all model-sampled (the filters read each node's + masked-True tokens + logprobs).""" + return vf.MessageNode( + message=vf.AssistantMessage(content="x"), + token_ids=token_ids, + mask=[True] * len(token_ids), + logprobs=logprobs, + ) def _make_rollout( @@ -18,52 +31,21 @@ def _make_rollout( *, reward: float = 1.0, multi_step: bool = False, -) -> TrainRollout: - """Build a ``TrainRollout`` with a minimal ``vf.RolloutOutput``-shaped - raw payload — enough for the filters to inspect ``trajectory`` / - ``stop_condition`` / etc.""" +) -> Rollout: + """Build a ``Rollout`` (a message-graph trace) carrying the completion tokens — enough for + the filters to inspect each node's sampled tokens / logprobs.""" if multi_step: mid = len(completion_ids) // 2 - trajectory = [ - { - "tokens": { - "completion_ids": completion_ids[:mid], - "completion_logprobs": completion_logprobs[:mid], - "completion_mask": [1] * mid, - } - }, - { - "tokens": { - "completion_ids": completion_ids[mid:], - "completion_logprobs": completion_logprobs[mid:], - "completion_mask": [1] * (len(completion_ids) - mid), - } - }, + nodes = [ + _assistant_node(completion_ids[:mid], completion_logprobs[:mid]), + _assistant_node(completion_ids[mid:], completion_logprobs[mid:]), ] else: - trajectory = [ - { - "tokens": { - "completion_ids": completion_ids, - "completion_logprobs": completion_logprobs, - "completion_mask": [1] * len(completion_ids), - } - } - ] - raw = { - "trajectory": trajectory, - "reward": reward, - "stop_condition": None, - "metrics": {}, - } - return TrainRollout( - raw=raw, - env_name="test", - example_id=0, - group_id=uuid.uuid4(), - policy_version=0, - off_policy_steps=0, - ) + nodes = [_assistant_node(completion_ids, completion_logprobs)] + rollout = Rollout[vf.Task](task=vf.Task(idx=0, prompt=""), nodes=nodes, rewards={"reward": reward}) + rollout.env_name = "test" + rollout.group_id = uuid.uuid4() + return rollout def _make_gibberish_filter(vocab_size=128_000, token_id_threshold=100_000, logprob_offset=2.0, enforce=False): @@ -248,9 +230,9 @@ def test_apply_filters_enforced_flags_rollout(): apply_filters([gibberish_filter], [rollout]) assert rollout.reward == 1.0 - assert rollout.raw["trajectory"][0]["tokens"]["completion_ids"] == [120_000] - assert rollout.raw["trajectory"][0]["tokens"]["completion_mask"] == [1] - assert rollout.raw["stop_condition"] is None + assert rollout.nodes[0].token_ids == [120_000] + assert rollout.nodes[0].mask == [True] + assert rollout.stop_condition is None assert rollout.filter_results == {"gibberish": True} assert rollout.is_filtered is True @@ -267,9 +249,9 @@ def test_apply_filters_preserves_clean_rollouts(): apply_filters([gibberish_filter], [rollout]) assert rollout.reward == 1.0 - assert rollout.raw["trajectory"][0]["tokens"]["completion_ids"] == [50, 60, 70] - assert all(m == 1 for m in rollout.raw["trajectory"][0]["tokens"]["completion_mask"]) - assert rollout.raw["stop_condition"] is None + assert rollout.nodes[0].token_ids == [50, 60, 70] + assert all(rollout.nodes[0].mask) + assert rollout.stop_condition is None assert rollout.filter_results == {"gibberish": False} assert rollout.is_filtered is False @@ -286,7 +268,7 @@ def test_apply_filters_first_filter_wins(): apply_filters([gibberish_filter, repetition_filter], [rollout]) - assert rollout.raw["stop_condition"] is None + assert rollout.stop_condition is None assert rollout.filter_results == {"gibberish": True, "repetition": False} assert rollout.is_filtered is True @@ -329,13 +311,13 @@ def test_apply_filters_enforced_preserves_rollout_tokens(): apply_filters([gibberish_filter], [rollout]) - assert rollout.raw["trajectory"][0]["tokens"]["completion_ids"] == [10, 120_000, 30] - assert rollout.raw["trajectory"][0]["tokens"]["completion_logprobs"] == [ + assert rollout.nodes[0].token_ids == [10, 120_000, 30] + assert rollout.nodes[0].logprobs == [ -1.0, gibberish_filter.logprob_threshold - 1.0, -0.5, ] - assert rollout.raw["trajectory"][0]["tokens"]["completion_mask"] == [1, 1, 1] + assert rollout.nodes[0].mask == [True, True, True] assert rollout.is_filtered is True @@ -347,11 +329,11 @@ def test_apply_filters_preserves_existing_stop_condition(): completion_logprobs=[gibberish_filter.logprob_threshold - 1.0], reward=1.0, ) - rollout.raw["stop_condition"] = "generation_truncated" + rollout.stop_condition = "generation_truncated" apply_filters([gibberish_filter], [rollout]) - assert rollout.raw["stop_condition"] == "generation_truncated" + assert rollout.stop_condition == "generation_truncated" assert rollout.is_filtered is True @@ -370,8 +352,8 @@ def test_apply_filters_monitor_only_tracks_detection(): apply_filters([gibberish_filter], [rollout]) assert rollout.reward == 1.0 - assert all(m == 1 for m in rollout.raw["trajectory"][0]["tokens"]["completion_mask"]) - assert rollout.raw["stop_condition"] is None + assert all(rollout.nodes[0].mask) + assert rollout.stop_condition is None assert rollout.filter_results == {"gibberish": True} assert rollout.is_filtered is False diff --git a/tests/unit/orchestrator/test_orchestrator_setup.py b/tests/unit/orchestrator/test_orchestrator_setup.py index 2372b004fd7..07ac09500cb 100644 --- a/tests/unit/orchestrator/test_orchestrator_setup.py +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -48,44 +48,3 @@ async def run() -> None: ) asyncio.run(run()) - - -def test_setup_student_inference_pool_defaults_to_mito(): - """No renderer -> plain MITO chat completions.""" - - async def run() -> None: - tokenizer = object() - config = SimpleNamespace( - training_mode="rl", - renderer=None, - pool_size=None, - student=SimpleNamespace( - client=SimpleNamespace(base_url=["http://localhost:8000/v1"]), - model=SimpleNamespace(name="student-model"), - ), - ) - inference_pool = object() - - with ( - patch("renderers.base.create_renderer") as create_renderer_mock, - patch( - "prime_rl.orchestrator.utils.setup_inference_pool", - new=AsyncMock(return_value=inference_pool), - ) as setup_pool_mock, - ): - renderer, returned_pool = await setup_student_inference_pool( - config=config, - tokenizer=tokenizer, - ) - - assert renderer is None - assert returned_pool is inference_pool - create_renderer_mock.assert_not_called() - setup_pool_mock.assert_awaited_once_with( - config.student.client, - model_name="student-model", - train_client_type="openai_chat_completions", - eval_client_type="openai_chat_completions", - ) - - asyncio.run(run()) diff --git a/tests/unit/orchestrator/test_sft_trajectories.py b/tests/unit/orchestrator/test_sft_trajectories.py deleted file mode 100644 index 8252c6cff5c..00000000000 --- a/tests/unit/orchestrator/test_sft_trajectories.py +++ /dev/null @@ -1,141 +0,0 @@ -from unittest.mock import MagicMock - -import verifiers as vf - -from prime_rl.orchestrator.trajectories import backfill_rollout_tokens, interleave_rollout - - -class SimpleChatTokenizer: - def __init__(self): - self._tok2id: dict[str, int] = {} - self._next_id = 1 - - def _id(self, token: str) -> int: - if token not in self._tok2id: - self._tok2id[token] = self._next_id - self._next_id += 1 - return self._tok2id[token] - - def apply_chat_template(self, messages, add_generation_prompt=False, return_dict=False): - del return_dict - ids = [] - for message in messages: - role = message.get("role", "unknown") - ids.append(self._id(f"<|{role}|>")) - content = message.get("content", "") - if isinstance(content, str): - if content: - ids.append(self._id(content)) - else: - ids.append(self._id(str(content))) - if add_generation_prompt: - ids.append(self._id("<|assistant|>")) - return ids - - -def test_interleave_rollout_missing_tokens_returns_none(): - output = vf.RolloutOutput( - example_id=1, - trajectory=[ - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U1"}], - completion=[{"role": "assistant", "content": "A1"}], - response=MagicMock(), - tokens=None, - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ) - ], - sampling_args={"temperature": 1.0}, - error=None, - ) - - assert interleave_rollout(output) is None - - -def test_backfill_rollout_tokens_for_sft(): - tokenizer = SimpleChatTokenizer() - output = vf.RolloutOutput( - example_id=42, - env_name="test-env", - trajectory=[ - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U1"}], - completion=[{"role": "assistant", "content": "A1"}], - response=MagicMock(), - tokens=None, - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - vf.TrajectoryStep( - prompt=[ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "A1"}, - {"role": "user", "content": "U2"}, - ], - completion=[{"role": "assistant", "content": "A2"}], - response=MagicMock(), - tokens=None, - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - ], - sampling_args={"temperature": 1.0}, - error=None, - ) - - backfill_rollout_tokens(output, tokenizer) - - rollouts = interleave_rollout(output) - assert rollouts is not None - assert len(rollouts) == 1 - - rollout = rollouts[0] - step1_prompt_ids = tokenizer.apply_chat_template( - [{"role": "user", "content": "U1"}], - add_generation_prompt=True, - ) - step1_full_ids = tokenizer.apply_chat_template( - [{"role": "user", "content": "U1"}, {"role": "assistant", "content": "A1"}], - add_generation_prompt=False, - ) - step2_prompt_ids = tokenizer.apply_chat_template( - [ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "A1"}, - {"role": "user", "content": "U2"}, - ], - add_generation_prompt=True, - ) - step2_full_ids = tokenizer.apply_chat_template( - [ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "A1"}, - {"role": "user", "content": "U2"}, - {"role": "assistant", "content": "A2"}, - ], - add_generation_prompt=False, - ) - - prefix_len_1 = len(step1_prompt_ids) - prefix_len_2 = len(step2_prompt_ids) - step1_completion_ids = step1_full_ids[prefix_len_1:] - step2_completion_ids = step2_full_ids[prefix_len_2:] - step1_prefix = step1_prompt_ids + step1_completion_ids - step2_new_prompt_ids = step2_prompt_ids[len(step1_prefix) :] - - assert rollout.prompt_ids == step1_prompt_ids - assert rollout.completion_ids == step1_completion_ids + step2_new_prompt_ids + step2_completion_ids - assert rollout.completion_mask == ( - [True] * len(step1_completion_ids) + [False] * len(step2_new_prompt_ids) + [True] * len(step2_completion_ids) - ) - assert rollout.completion_logprobs == [0.0] * len(rollout.completion_ids) diff --git a/tests/unit/orchestrator/test_teacher_logprobs.py b/tests/unit/orchestrator/test_teacher_logprobs.py index d63fdce7924..5be1fe2a42d 100644 --- a/tests/unit/orchestrator/test_teacher_logprobs.py +++ b/tests/unit/orchestrator/test_teacher_logprobs.py @@ -2,7 +2,8 @@ import json import httpx -import verifiers as vf +import openai +from verifiers.v1.clients.config import EvalClientConfig from prime_rl.orchestrator import utils as orchestrator_utils from prime_rl.transport import TrainingSample @@ -41,20 +42,19 @@ async def _run(): "kv_transfer_params": None, } ) - monkeypatch.setattr(orchestrator_utils, "setup_openai_client", lambda _: fake_client) + # compute_teacher_logprobs constructs AsyncOpenAI directly; hand back the fake. + monkeypatch.setattr(openai, "AsyncOpenAI", lambda **kwargs: fake_client) sample = TrainingSample( - prompt_ids=[1], - prompt_mask=[True], - completion_ids=[2, 3], - completion_mask=[True, True], - completion_logprobs=[-0.1, -0.2], - completion_temperatures=[1.0, 1.0], + token_ids=[1, 2, 3], + mask=[False, True, True], + logprobs=[0.0, -0.1, -0.2], + temperatures=[1.0, 1.0, 1.0], env_name="test-env", ) result = await orchestrator_utils.compute_teacher_logprobs( - clients=[vf.ClientConfig()], + clients=[EvalClientConfig(base_url="http://fake-host:8000/v1")], model_name="teacher-model", samples=[sample], ) diff --git a/tests/unit/orchestrator/test_trajectories.py b/tests/unit/orchestrator/test_trajectories.py deleted file mode 100644 index bda129cd43e..00000000000 --- a/tests/unit/orchestrator/test_trajectories.py +++ /dev/null @@ -1,1380 +0,0 @@ -from unittest.mock import MagicMock - -import numpy as np -import pybase64 -import pytest -import verifiers as vf - -from prime_rl.orchestrator.trajectories import ( - _deserialize_tool_calls, - align_routed_experts, - interleave_rollout, -) - -_interleave_rollout = interleave_rollout - - -def interleave_rollout(output, *args, **kwargs): - kwargs.setdefault("env_name", output.get("env_name", "test-env")) - return _interleave_rollout(output, *args, **kwargs) - - -def _decode_mm_pixels(sample) -> list: - """Decode ``sample.mm_kwargs['pixel_values']`` to a nested list.""" - p = sample.mm_kwargs["pixel_values"] - return np.frombuffer(p.data, dtype=np.dtype(p.dtype)).reshape(p.shape).tolist() - - -def _decode_mm_thw(sample) -> list: - """Decode ``sample.mm_kwargs['image_grid_thw']`` to a nested list.""" - g = sample.mm_kwargs["image_grid_thw"] - return np.frombuffer(g.data, dtype=np.dtype(g.dtype)).reshape(g.shape).tolist() - - -def _routed_experts_payload(data, start: int = 0) -> dict: - arr = np.asarray(data, dtype=np.uint8) - return { - "data": pybase64.b64encode(memoryview(np.ascontiguousarray(arr))).decode("ascii"), - "shape": list(arr.shape), - "start": start, - } - - -def _sample_routed_experts(sample) -> np.ndarray: - assert sample.routed_experts is not None - return np.frombuffer(sample.routed_experts.data, dtype=np.dtype(sample.routed_experts.dtype)).reshape( - sample.routed_experts.shape - ) - - -def test_deserialize_tool_calls_does_not_inject_missing_key(): - messages = [{"role": "assistant", "content": "hello"}] - - deserialized = _deserialize_tool_calls(messages) - - assert "tool_calls" not in deserialized[0] - - -def test_deserialize_tool_calls_parses_arguments_when_present(): - messages = [ - { - "role": "assistant", - "tool_calls": [ - { - "id": "1", - "type": "function", - "function": {"name": "lookup", "arguments": '{"x": 1}'}, - } - ], - } - ] - - deserialized = _deserialize_tool_calls(messages) - - assert deserialized[0]["tool_calls"][0]["function"]["arguments"] == {"x": 1} - - -@pytest.fixture -def single_step_trajectory_output(): - output = vf.RolloutOutput( - example_id=0, - trajectory=[ - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U1"}], - completion=[{"role": "assistant", "content": "A1"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ) - ], - sampling_args={"temperature": 1.0}, - error=None, - ) - return output - - -@pytest.fixture -def multi_step_trajectory_output(): - output = vf.RolloutOutput( - example_id=0, - trajectory=[ - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U1"}], - completion=[{"role": "assistant", "content": "A1"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - vf.TrajectoryStep( - prompt=[ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "A1"}, - {"role": "user", "content": "U2"}, - ], - completion=[{"role": "assistant", "content": "A2"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2, 3, 4, 5, 6], - prompt_mask=[0, 0, 0, 0, 0, 0], - completion_ids=[7, 8], - completion_mask=[1, 1], - completion_logprobs=[-0.3, -0.4], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - ], - sampling_args={"temperature": 1.0}, - error=None, - ) - return output - - -@pytest.fixture -def multi_step_trajectory_with_tool_calls_output(): - output = vf.RolloutOutput( - example_id=0, - trajectory=[ - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U1"}], - completion=[{"role": "assistant", "content": "A1 + TC1"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - vf.TrajectoryStep( - prompt=[ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "A1 + TC1"}, - {"role": "tool", "tool_call_id": "TR1", "content": "TR1"}, - ], - completion=[{"role": "assistant", "content": "A2 + TC2"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2, 3, 4, 5, 6], - prompt_mask=[0, 0, 0, 0, 0, 0], - completion_ids=[7, 8], - completion_mask=[1, 1], - completion_logprobs=[-0.3, -0.4], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - ], - reward=1.0, - advantage=None, - stop_condition=None, - metrics={"has_error": 0.0, "tool_calls": 1.0}, - sampling_args={"temperature": 1.0}, - error=None, - ) - return output - - -@pytest.fixture -def multi_step_trajectory_extension_never_holds(): - """ - 2-step trajectory where extension NEVER holds (step 2 has completely different tokens). - This simulates e.g. a chat template that re-renders the entire conversation differently. - """ - output = vf.RolloutOutput( - example_id=0, - trajectory=[ - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U1"}], - completion=[{"role": "assistant", "content": "A1"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - vf.TrajectoryStep( - prompt=[ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "A1"}, - {"role": "user", "content": "U2"}, - ], - completion=[{"role": "assistant", "content": "A2"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - # Different tokens - extension breaks (e.g. thinking was stripped) - prompt_ids=[10, 20, 30, 40, 50, 60], - prompt_mask=[0, 0, 0, 0, 0, 0], - completion_ids=[7, 8], - completion_mask=[1, 1], - completion_logprobs=[-0.3, -0.4], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - ], - sampling_args={"temperature": 1.0}, - error=None, - ) - return output - - -@pytest.fixture -def multi_step_trajectory_with_tool_calls_extension_never_holds(): - """2-step trajectory with tool calls where extension NEVER holds.""" - output = vf.RolloutOutput( - example_id=0, - trajectory=[ - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U1"}], - completion=[{"role": "assistant", "content": "A1 + TC1"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - extras={}, - is_truncated=False, - trajectory_id="1", - ), - vf.TrajectoryStep( - prompt=[ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "A1 + TC1"}, - {"role": "tool", "tool_call_id": "TR1", "content": "TR1"}, - ], - completion=[{"role": "assistant", "content": "A2 + TC2"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - # Different tokens - extension breaks - prompt_ids=[10, 20, 30, 40, 50, 60], - prompt_mask=[0, 0, 0, 0, 0, 0], - completion_ids=[7, 8], - completion_mask=[1, 1], - completion_logprobs=[-0.3, -0.4], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - extras={}, - is_truncated=False, - trajectory_id="1", - ), - ], - reward=1.0, - advantage=None, - stop_condition=None, - sampling_args={"temperature": 1.0}, - metrics={"has_error": 0.0, "tool_calls": 1.0}, - error=None, - ) - return output - - -def test_branching_equivalent_multi_step_trajectory(multi_step_trajectory_extension_never_holds): - """When extension never holds, each step becomes its own sample (same as old branching).""" - rollouts = interleave_rollout(multi_step_trajectory_extension_never_holds) - assert rollouts is not None - assert len(rollouts) == 2 - - # first step - rollout = rollouts[0] - assert rollout.prompt_ids == [1, 2] - assert rollout.prompt_mask == [False, False] - assert rollout.completion_ids == [3, 4] - assert rollout.completion_mask == [True, True] - assert rollout.completion_logprobs == [-0.1, -0.2] - assert rollout.completion_temperatures == [] - - # second step - rollout = rollouts[1] - assert rollout.prompt_ids == [10, 20, 30, 40, 50, 60] - assert rollout.prompt_mask == [False, False, False, False, False, False] - assert rollout.completion_ids == [7, 8] - assert rollout.completion_mask == [True, True] - assert rollout.completion_logprobs == [-0.3, -0.4] - assert rollout.completion_temperatures == [] - - -def test_branching_equivalent_multi_step_trajectory_with_tool_calls( - multi_step_trajectory_with_tool_calls_extension_never_holds, -): - """When extension never holds (with tool calls), same as old branching.""" - rollouts = interleave_rollout(multi_step_trajectory_with_tool_calls_extension_never_holds) - assert rollouts is not None - assert len(rollouts) == 2 - - # first step - rollout = rollouts[0] - assert rollout.prompt_ids == [1, 2] - assert rollout.prompt_mask == [False, False] - assert rollout.completion_ids == [3, 4] - assert rollout.completion_mask == [True, True] - assert rollout.completion_logprobs == [-0.1, -0.2] - assert rollout.completion_temperatures == [] - - # second step - rollout = rollouts[1] - assert rollout.prompt_ids == [10, 20, 30, 40, 50, 60] - assert rollout.prompt_mask == [False, False, False, False, False, False] - assert rollout.completion_ids == [7, 8] - assert rollout.completion_mask == [True, True] - assert rollout.completion_logprobs == [-0.3, -0.4] - assert rollout.completion_temperatures == [] - - -def test_interleave_rollout_single_step_trajectory(single_step_trajectory_output): - single_step_trajectory_output["env_name"] = "test-env" - rollouts = interleave_rollout(single_step_trajectory_output) - assert rollouts is not None - assert len(rollouts) == 1 - rollout = rollouts[0] - - assert rollout.prompt_ids == [1, 2] - assert rollout.prompt_mask == [False, False] - assert rollout.completion_ids == [3, 4] - assert rollout.completion_mask == [True, True] - assert rollout.completion_logprobs == [-0.1, -0.2] - assert rollout.completion_temperatures == [] - assert rollout.env_name == "test-env" - - -def test_interleave_rollout_multi_step_trajectory(multi_step_trajectory_output): - rollouts = interleave_rollout(multi_step_trajectory_output) - assert rollouts is not None - assert len(rollouts) == 1 - rollout = rollouts[0] - - assert rollout.prompt_ids == [1, 2] - assert rollout.prompt_mask == [False, False] - assert rollout.completion_ids == [3, 4, 5, 6, 7, 8] - assert rollout.completion_mask == [True, True, False, False, True, True] - assert rollout.completion_logprobs == [-0.1, -0.2, 0, 0, -0.3, -0.4] - # ``completion_temperatures`` is filled by the orchestrator post-interleave; empty here. - assert rollout.completion_temperatures == [] - - -def test_interleave_rollout_multi_step_trajectory_with_tool_calls(multi_step_trajectory_with_tool_calls_output): - rollouts = interleave_rollout(multi_step_trajectory_with_tool_calls_output) - assert rollouts is not None - assert len(rollouts) == 1 - rollout = rollouts[0] - - assert rollout.prompt_ids == [1, 2] - assert rollout.prompt_mask == [False, False] - assert rollout.completion_ids == [3, 4, 5, 6, 7, 8] - assert rollout.completion_mask == [True, True, False, False, True, True] - assert rollout.completion_logprobs == [-0.1, -0.2, 0, 0, -0.3, -0.4] - # ``completion_temperatures`` is filled by the orchestrator post-interleave; empty here. - assert rollout.completion_temperatures == [] - - -@pytest.fixture -def five_step_trajectory_with_extension_break(): - """ - 5-step trajectory where extension property breaks at step 4. - - Steps 1-3: extension holds (tokens grow by appending) - Step 4: extension breaks (completely different prefix, e.g. context compaction) - Steps 4-5: extension holds again - - Expected: 2 samples (steps 1-3 merged, steps 4-5 merged) - """ - output = vf.RolloutOutput( - example_id=0, - trajectory=[ - # Step 1: initial prompt and completion - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U1"}], - completion=[{"role": "assistant", "content": "A1"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - extras={}, - is_truncated=False, - trajectory_id="1", - ), - # Step 2: extends step 1 (prefix [1,2,3,4] matches) - vf.TrajectoryStep( - prompt=[ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "A1"}, - {"role": "user", "content": "U2"}, - ], - completion=[{"role": "assistant", "content": "A2"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2, 3, 4, 5, 6], - prompt_mask=[0, 0, 0, 0, 0, 0], - completion_ids=[7, 8], - completion_mask=[1, 1], - completion_logprobs=[-0.3, -0.4], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - extras={}, - is_truncated=False, - trajectory_id="1", - ), - # Step 3: extends step 2 (prefix [1,2,3,4,5,6,7,8] matches) - vf.TrajectoryStep( - prompt=[ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "A1"}, - {"role": "user", "content": "U2"}, - {"role": "assistant", "content": "A2"}, - {"role": "user", "content": "U3"}, - ], - completion=[{"role": "assistant", "content": "A3"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - prompt_mask=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - completion_ids=[11, 12], - completion_mask=[1, 1], - completion_logprobs=[-0.5, -0.6], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - extras={}, - is_truncated=False, - trajectory_id="1", - ), - # Step 4: EXTENSION BREAKS - different prefix (e.g. thinking stripped, context compacted) - vf.TrajectoryStep( - prompt=[ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "A1"}, # thinking stripped - {"role": "user", "content": "U2"}, - {"role": "assistant", "content": "A2"}, - {"role": "user", "content": "U3"}, - {"role": "assistant", "content": "A3"}, - {"role": "user", "content": "U4"}, - ], - completion=[{"role": "assistant", "content": "A4"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[100, 101, 102, 103], # completely different tokens (re-rendered) - prompt_mask=[0, 0, 0, 0], - completion_ids=[104, 105], - completion_mask=[1, 1], - completion_logprobs=[-0.7, -0.8], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - extras={}, - is_truncated=False, - trajectory_id="1", - ), - # Step 5: extends step 4 (prefix [100,101,102,103,104,105] matches) - vf.TrajectoryStep( - prompt=[ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "A1"}, - {"role": "user", "content": "U2"}, - {"role": "assistant", "content": "A2"}, - {"role": "user", "content": "U3"}, - {"role": "assistant", "content": "A3"}, - {"role": "user", "content": "U4"}, - {"role": "assistant", "content": "A4"}, - {"role": "user", "content": "U5"}, - ], - completion=[{"role": "assistant", "content": "A5"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[100, 101, 102, 103, 104, 105, 106, 107], - prompt_mask=[0, 0, 0, 0, 0, 0, 0, 0], - completion_ids=[108, 109], - completion_mask=[1, 1], - completion_logprobs=[-0.9, -1.0], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - extras={}, - is_truncated=False, - trajectory_id="1", - ), - ], - sampling_args={"temperature": 1.0}, - error=None, - ) - return output - - -def test_interleave_rollout_extension_break_creates_multiple_samples(five_step_trajectory_with_extension_break): - """ - When extension property breaks mid-trajectory, interleave_rollout should: - - Merge steps 1-3 into first sample (extension held) - - Start new sample at step 4 (extension broke) - - Merge steps 4-5 into second sample (extension held again) - """ - rollouts = interleave_rollout(five_step_trajectory_with_extension_break) - - assert rollouts is not None - assert len(rollouts) == 2, "Should produce 2 samples when extension breaks at step 4" - - # First sample: steps 1-3 merged - sample1 = rollouts[0] - assert sample1.prompt_ids == [1, 2] - assert sample1.prompt_mask == [False, False] - # completion_ids: step1 completion [3,4] + step2 new prompt [5,6] + step2 completion [7,8] - # + step3 new prompt [9,10] + step3 completion [11,12] - assert sample1.completion_ids == [3, 4, 5, 6, 7, 8, 9, 10, 11, 12] - # completion_mask: step1 [T,T] + step2 prompt [F,F] + step2 completion [T,T] - # + step3 prompt [F,F] + step3 completion [T,T] - assert sample1.completion_mask == [True, True, False, False, True, True, False, False, True, True] - assert sample1.completion_logprobs == [-0.1, -0.2, 0, 0, -0.3, -0.4, 0, 0, -0.5, -0.6] - - # Second sample: steps 4-5 merged (fresh start after extension break) - sample2 = rollouts[1] - assert sample2.prompt_ids == [100, 101, 102, 103] - assert sample2.prompt_mask == [False, False, False, False] - # completion_ids: step4 completion [104,105] + step5 new prompt [106,107] + step5 completion [108,109] - assert sample2.completion_ids == [104, 105, 106, 107, 108, 109] - # completion_mask: step4 [T,T] + step5 prompt [F,F] + step5 completion [T,T] - assert sample2.completion_mask == [True, True, False, False, True, True] - assert sample2.completion_logprobs == [-0.7, -0.8, 0, 0, -0.9, -1.0] - - -@pytest.fixture -def interleaved_agents_trajectory(): - """ - Trajectory with interleaved agents: agent1 steps, then agent2 step, then agent1 continues. - This tests multi-prefix tracking where agent1-step3 should merge back with agent1 sample. - - agent1-step1: prompt=[1,2], completion=[3,4] - agent1-step2: prompt=[1,2,3,4,5,6], completion=[7,8] (extends agent1-step1) - agent2-step1: prompt=[100,101], completion=[102,103] (different prefix, new sample) - agent1-step3: prompt=[1,2,3,4,5,6,7,8,9,10], completion=[11,12] (extends agent1-step2!) - """ - output = vf.RolloutOutput( - example_id=1, - task="test", - trajectory=[ - # agent1-step1 - vf.TrajectoryStep( - prompt="agent1 turn 1", - completion="response 1", - response=None, - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="traj1", - extras={}, - ), - # agent1-step2 (extends agent1-step1) - vf.TrajectoryStep( - prompt="agent1 turn 2", - completion="response 2", - response=None, - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2, 3, 4, 5, 6], - prompt_mask=[0, 0, 0, 0, 0, 0], - completion_ids=[7, 8], - completion_mask=[1, 1], - completion_logprobs=[-0.3, -0.4], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="traj1", - extras={}, - ), - # agent2-step1 (different prefix, starts new sample) - vf.TrajectoryStep( - prompt="agent2 turn 1", - completion="agent2 response", - response=None, - tokens=vf.TrajectoryStepTokens( - prompt_ids=[100, 101], - prompt_mask=[0, 0], - completion_ids=[102, 103], - completion_mask=[1, 1], - completion_logprobs=[-0.5, -0.6], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="traj2", - extras={}, - ), - # agent1-step3 (extends agent1-step2, should merge back!) - vf.TrajectoryStep( - prompt="agent1 turn 3", - completion="response 3", - response=None, - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - prompt_mask=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - completion_ids=[11, 12], - completion_mask=[1, 1], - completion_logprobs=[-0.7, -0.8], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="traj1", - extras={}, - ), - ], - sampling_args={"temperature": 1.0}, - error=None, - ) - return output - - -def test_interleave_rollout_interleaved_agents(interleaved_agents_trajectory): - """ - When agents are interleaved (agent1, agent1, agent2, agent1), the multi-prefix - tracking should merge agent1-step3 back into the agent1 sample, not start a new one. - """ - rollouts = interleave_rollout(interleaved_agents_trajectory) - - assert rollouts is not None - assert len(rollouts) == 2, "Should produce 2 samples (agent1 merged, agent2 separate)" - - # First sample: agent1 steps 1, 2, 3 merged - agent1_sample = rollouts[0] - assert agent1_sample.prompt_ids == [1, 2] - assert agent1_sample.prompt_mask == [False, False] - # completion_ids: step1 [3,4] + step2 new prompt [5,6] + step2 completion [7,8] - # + step3 new prompt [9,10] + step3 completion [11,12] - assert agent1_sample.completion_ids == [3, 4, 5, 6, 7, 8, 9, 10, 11, 12] - assert agent1_sample.completion_mask == [True, True, False, False, True, True, False, False, True, True] - assert agent1_sample.completion_logprobs == [-0.1, -0.2, 0, 0, -0.3, -0.4, 0, 0, -0.7, -0.8] - - # Second sample: agent2 step 1 only - agent2_sample = rollouts[1] - assert agent2_sample.prompt_ids == [100, 101] - assert agent2_sample.prompt_mask == [False, False] - assert agent2_sample.completion_ids == [102, 103] - assert agent2_sample.completion_mask == [True, True] - assert agent2_sample.completion_logprobs == [-0.5, -0.6] - - -@pytest.fixture -def prefix_of_prefix_trajectory(): - """ - Trajectory where one active sample's prefix is a strict prefix of another's. - - Construction: - - step 0: prompt=[1,2], completion=[3,4] -> sample A, P_A=[1,2,3,4] - - step 1: extends A. prompt=[1,2,3,4,5], completion=[6] -> P_A=[1,2,3,4,5,6] - - step 2: rollback/regenerate. prompt=[1,2] (shorter than P_A so no match), - completion=[3,4,5,6,7] -> sample B, P_B=[1,2,3,4,5,6,7] - P_B starts with P_A. - - step 3: extends B. prompt=[1,2,3,4,5,6,7,8], completion=[9] - Both P_A and P_B are token-prefixes of the step's prompt. - - The correct match is the longer P_B. First-match-wins picks P_A and silently - folds B's generated tokens into A as user-input tokens (mask=False). - """ - output = vf.RolloutOutput( - example_id=2, - task="test", - trajectory=[ - vf.TrajectoryStep( - prompt="step 0", - completion="completion 0", - response=None, - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="traj_A", - extras={}, - ), - vf.TrajectoryStep( - prompt="step 1", - completion="completion 1", - response=None, - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2, 3, 4, 5], - prompt_mask=[0, 0, 0, 0, 0], - completion_ids=[6], - completion_mask=[1], - completion_logprobs=[-0.3], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="traj_A", - extras={}, - ), - vf.TrajectoryStep( - prompt="step 2 (rollback)", - completion="completion 2", - response=None, - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4, 5, 6, 7], - completion_mask=[1, 1, 1, 1, 1], - completion_logprobs=[-0.4, -0.5, -0.6, -0.7, -0.8], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="traj_B", - extras={}, - ), - vf.TrajectoryStep( - prompt="step 3 (extends B)", - completion="completion 3", - response=None, - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2, 3, 4, 5, 6, 7, 8], - prompt_mask=[0, 0, 0, 0, 0, 0, 0, 0], - completion_ids=[9], - completion_mask=[1], - completion_logprobs=[-0.9], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="traj_B", - extras={}, - ), - ], - sampling_args={"temperature": 1.0}, - error=None, - ) - return output - - -def test_interleave_rollout_picks_longest_matching_prefix(prefix_of_prefix_trajectory): - """ - When two active samples both match (one's prefix is a strict prefix of the - other's), the longer prefix is the correct extension. Previously the first- - match-wins loop folded the longer sample's generated tokens into the shorter - sample as user input (mask=False) and left the longer sample stale. - """ - rollouts = interleave_rollout(prefix_of_prefix_trajectory) - - assert rollouts is not None - assert len(rollouts) == 2 - - # Sample A: steps 0 and 1 only. Step 3 must NOT have been folded in here. - sample_a = rollouts[0] - assert sample_a.prompt_ids == [1, 2] - # step 0 completion [3,4] + step 1 new prompt [5] + step 1 completion [6] - assert sample_a.completion_ids == [3, 4, 5, 6] - assert sample_a.completion_mask == [True, True, False, True] - assert sample_a.completion_logprobs == [-0.1, -0.2, 0.0, -0.3] - - # Sample B: steps 2 and 3 merged. The token 7 (from step 2's completion) - # must remain masked as a generated token, not silently re-classified. - sample_b = rollouts[1] - assert sample_b.prompt_ids == [1, 2] - # step 2 completion [3,4,5,6,7] + step 3 new prompt [8] + step 3 completion [9] - assert sample_b.completion_ids == [3, 4, 5, 6, 7, 8, 9] - assert sample_b.completion_mask == [True, True, True, True, True, False, True] - assert sample_b.completion_logprobs == [-0.4, -0.5, -0.6, -0.7, -0.8, 0.0, -0.9] - - -def test_interleave_rollout_empty_trajectory(): - """Empty trajectory returns None.""" - output = vf.RolloutOutput( - example_id=1, - trajectory=[], - error=None, - ) - assert interleave_rollout(output) is None - - -def test_interleave_rollout_error_masks_all_false(): - """ - When rollout output has an error, all completion_mask values should be False - across both make_sample (step 0) and extend_sample (step 1). - """ - output = vf.RolloutOutput( - example_id=1, - trajectory=[ - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U1"}], - completion=[{"role": "assistant", "content": "A1"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U2"}], - completion=[{"role": "assistant", "content": "A2"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2, 3, 4, 5, 6], - prompt_mask=[0, 0, 0, 0, 0, 0], - completion_ids=[7, 8], - completion_mask=[1, 1], - completion_logprobs=[-0.3, -0.4], - overlong_prompt=False, - is_truncated=False, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - ], - error="timeout: environment exceeded time limit", - sampling_args={"temperature": 0.8}, - ) - - rollouts = interleave_rollout(output) - - assert rollouts is not None - assert len(rollouts) == 1 - rollout = rollouts[0] - # Extension holds so tokens merge, but ALL completion_mask should be False - assert rollout.completion_ids == [3, 4, 5, 6, 7, 8] - assert rollout.completion_mask == [False, False, False, False, False, False] - # Logprobs preserved; ``completion_temperatures`` is filled by the orchestrator post-interleave. - assert rollout.completion_logprobs == [-0.1, -0.2, 0.0, 0.0, -0.3, -0.4] - assert rollout.completion_temperatures == [] - - -def test_align_routed_experts_none(): - assert align_routed_experts(None, 10) is None - - -def test_align_routed_experts_empty(): - experts = np.empty((0, 2, 2), dtype=np.uint8) - result = align_routed_experts(experts, 10) - assert result is not None - assert result.shape == (10, 2, 2) - assert np.all(result == 0) - - -def test_align_routed_experts_no_deficit(): - # 3 tokens, 2 layers, topk=2 - experts = np.asarray([[[0, 1], [2, 3]], [[4, 5], [6, 7]], [[0, 2], [1, 3]]], dtype=np.uint8) - result = align_routed_experts(experts, expected_len=3) - np.testing.assert_array_equal(result, experts) - - -def test_align_routed_experts_with_deficit(): - # 2 tokens but expected 4 (deficit of 2) - experts = np.asarray([[[1, 2], [3, 4]], [[5, 6], [7, 0]]], dtype=np.uint8) - result = align_routed_experts(experts, expected_len=4) - assert result is not None - assert result.shape == (4, 2, 2) - np.testing.assert_array_equal(result[:2], experts) - # Padded entries should be zero-filled with same shape [layers=2, topk=2] - np.testing.assert_array_equal(result[2], [[0, 0], [0, 0]]) - np.testing.assert_array_equal(result[3], [[0, 0], [0, 0]]) - - -def test_align_routed_experts_excess_length(): - experts = np.asarray([[[1, 2]], [[3, 4]], [[5, 6]]], dtype=np.uint8) - result = align_routed_experts(experts, expected_len=2) - np.testing.assert_array_equal(result, experts[:2]) - - -def test_interleave_rollout_single_step_with_routed_experts(): - """Routed experts are aligned and passed through for a single-step trajectory.""" - # prompt_ids=[1,2], completion_ids=[3,4] -> total 4 tokens - # vLLM returns num_tokens-1 = 3 routed expert entries - routed_experts_from_vllm = np.asarray([[[0, 1]], [[2, 3]], [[4, 5]]], dtype=np.uint8) - output = vf.RolloutOutput( - example_id=0, - trajectory=[ - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U1"}], - completion=[{"role": "assistant", "content": "A1"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - routed_experts=_routed_experts_payload(routed_experts_from_vllm), - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ) - ], - sampling_args={"temperature": 1.0}, - error=None, - ) - - rollouts = interleave_rollout(output) - assert rollouts is not None - assert len(rollouts) == 1 - sample = rollouts[0] - - # Should be aligned to 4 tokens (2 prompt + 2 completion) - assert sample.routed_experts is not None - routed_experts = _sample_routed_experts(sample) - assert routed_experts.shape == (4, 1, 2) - # First 3 are original, last one is zero-padded - np.testing.assert_array_equal(routed_experts[:3], routed_experts_from_vllm) - np.testing.assert_array_equal(routed_experts[3], [[0, 0]]) - - -def test_interleave_rollout_multi_step_with_routed_experts(): - """Routed experts are extended and aligned across multi-step trajectories.""" - # Step 1: prompt=[1,2], completion=[3,4] -> 4 tokens, vLLM returns 3 - step1_experts = np.asarray([[[1, 2]], [[3, 4]], [[5, 6]]], dtype=np.uint8) - # Step 2: prompt=[1,2,3,4,5,6], completion=[7,8], bridged from prefix len 4. - # vLLM returns routed experts starting at row 3: boundary token 4, then 5, 6, 7. - step2_experts = np.asarray([[[40, 41]], [[50, 51]], [[60, 61]], [[70, 71]]], dtype=np.uint8) - - output = vf.RolloutOutput( - example_id=0, - trajectory=[ - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U1"}], - completion=[{"role": "assistant", "content": "A1"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - routed_experts=_routed_experts_payload(step1_experts), - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - vf.TrajectoryStep( - prompt=[ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "A1"}, - {"role": "user", "content": "U2"}, - ], - completion=[{"role": "assistant", "content": "A2"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2, 3, 4, 5, 6], - prompt_mask=[0, 0, 0, 0, 0, 0], - completion_ids=[7, 8], - completion_mask=[1, 1], - completion_logprobs=[-0.3, -0.4], - overlong_prompt=False, - is_truncated=False, - routed_experts=_routed_experts_payload(step2_experts, start=3), - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - ], - sampling_args={"temperature": 1.0}, - error=None, - ) - - rollouts = interleave_rollout(output) - assert rollouts is not None - assert len(rollouts) == 1 - sample = rollouts[0] - - # Merged sample: prompt=[1,2], completion=[3,4,5,6,7,8] -> 8 tokens total - assert len(sample.prompt_ids) + len(sample.completion_ids) == 8 - assert sample.routed_experts is not None - routed_experts = _sample_routed_experts(sample) - assert routed_experts.shape == (8, 1, 2) - np.testing.assert_array_equal( - routed_experts, - np.asarray( - [ - [[1, 2]], - [[3, 4]], - [[5, 6]], - [[40, 41]], - [[50, 51]], - [[60, 61]], - [[70, 71]], - [[0, 0]], - ], - dtype=np.uint8, - ), - ) - - -def test_interleave_rollout_branch_delta_uses_prior_routed_prefix(): - step1_experts = np.asarray([[[1, 2]], [[3, 4]], [[5, 6]]], dtype=np.uint8) - step2_experts = np.asarray([[[40, 41]], [[50, 51]], [[60, 61]], [[70, 71]]], dtype=np.uint8) - step3_experts = np.asarray([[[80, 81]], [[90, 91]], [[100, 101]], [[110, 111]]], dtype=np.uint8) - - output = vf.RolloutOutput( - example_id=0, - trajectory=[ - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U1"}], - completion=[{"role": "assistant", "content": "A1"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - routed_experts=_routed_experts_payload(step1_experts), - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - vf.TrajectoryStep( - prompt=[ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "A1"}, - {"role": "user", "content": "U2"}, - ], - completion=[{"role": "assistant", "content": "A2"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2, 3, 4, 5, 6], - prompt_mask=[0, 0, 0, 0, 0, 0], - completion_ids=[7, 8], - completion_mask=[1, 1], - completion_logprobs=[-0.3, -0.4], - overlong_prompt=False, - is_truncated=False, - routed_experts=_routed_experts_payload(step2_experts, start=3), - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - vf.TrajectoryStep( - prompt=[ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "A1"}, - {"role": "user", "content": "branch"}, - ], - completion=[{"role": "assistant", "content": "A3"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2, 3, 4, 9, 10], - prompt_mask=[0, 0, 0, 0, 0, 0], - completion_ids=[11, 12], - completion_mask=[1, 1], - completion_logprobs=[-0.5, -0.6], - overlong_prompt=False, - is_truncated=False, - routed_experts=_routed_experts_payload(step3_experts, start=3), - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - ], - sampling_args={"temperature": 1.0}, - error=None, - ) - - rollouts = interleave_rollout(output) - - assert rollouts is not None - assert len(rollouts) == 2 - branched = rollouts[1] - assert branched.prompt_ids == [1, 2, 3, 4, 9, 10] - assert branched.completion_ids == [11, 12] - routed_experts = _sample_routed_experts(branched) - assert routed_experts.shape == (8, 1, 2) - np.testing.assert_array_equal( - routed_experts, - np.asarray( - [ - [[1, 2]], - [[3, 4]], - [[5, 6]], - [[80, 81]], - [[90, 91]], - [[100, 101]], - [[110, 111]], - [[0, 0]], - ], - dtype=np.uint8, - ), - ) - - -def test_interleave_rollout_none_routed_experts_stays_none(): - """When routed_experts is None, sample.routed_experts remains None.""" - output = vf.RolloutOutput( - example_id=0, - trajectory=[ - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "U1"}], - completion=[{"role": "assistant", "content": "A1"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - routed_experts=None, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ) - ], - sampling_args={"temperature": 1.0}, - error=None, - ) - - rollouts = interleave_rollout(output) - assert rollouts is not None - assert rollouts[0].routed_experts is None - - -# ============================================================================= -# Renderer-emitted multimodal data -# ============================================================================= - - -def test_interleave_rollout_packs_pixels_from_renderer_mm_data(): - """``interleave_rollout`` packs renderer-emitted ``multi_modal_data`` - (pixel_values / image_grid_thw / mm_token_type_ids) onto the - TrainingSample. - - verifiers' ``_delta_intermediate_mm_data`` ships per-step *delta* - mm_data (each step contains only items not present in the prior - step's cumulative set). Prime-rl unions across the sample's step - range to recover the cumulative set in image-placeholder order. - """ - import torch as _torch - from renderers.base import MultiModalData, PlaceholderRange - - # Two synthetic single-image items — values are arbitrary, what - # matters is that the packer concatenates them correctly. - item1_pv = _torch.tensor([[1.0, 2.0]], dtype=_torch.float32) - item2_pv = _torch.tensor([[3.0, 4.0]], dtype=_torch.float32) - item1_thw = _torch.tensor([[1, 2, 3]], dtype=_torch.int64) - item2_thw = _torch.tensor([[1, 4, 4]], dtype=_torch.int64) - - # Step 0: image h1 (first time it's seen, included in delta). - mm_step_0 = MultiModalData( - mm_hashes={"image": ["h1"]}, - mm_placeholders={"image": [PlaceholderRange(offset=1, length=1)]}, - mm_items={"image": [{"pixel_values": item1_pv, "image_grid_thw": item1_thw}]}, - ) - # Step 1: post-delta — only h2 (h1 was dropped because it was in - # the prior step's cumulative set). Renderer's bridge would have - # produced cumulative [h1, h2] before verifiers' delta rewrite. - mm_step_1 = MultiModalData( - mm_hashes={"image": ["h2"]}, - mm_placeholders={"image": [PlaceholderRange(offset=4, length=1)]}, - mm_items={"image": [{"pixel_values": item2_pv, "image_grid_thw": item2_thw}]}, - ) - - output = vf.RolloutOutput( - example_id=1, - trajectory=[ - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "Turn 1"}], - completion=[{"role": "assistant", "content": "Response 1"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2], - prompt_mask=[0, 0], - completion_ids=[3, 4], - completion_mask=[1, 1], - completion_logprobs=[-0.1, -0.2], - overlong_prompt=False, - is_truncated=False, - multi_modal_data=mm_step_0, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - vf.TrajectoryStep( - prompt=[{"role": "user", "content": "Turn 2"}], - completion=[{"role": "assistant", "content": "Response 2"}], - response=MagicMock(), - tokens=vf.TrajectoryStepTokens( - prompt_ids=[1, 2, 3, 4, 5], - prompt_mask=[0, 0, 0, 0, 0], - completion_ids=[6, 7], - completion_mask=[1, 1], - completion_logprobs=[-0.3, -0.4], - overlong_prompt=False, - is_truncated=False, - multi_modal_data=mm_step_1, - ), - reward=None, - advantage=None, - is_truncated=False, - trajectory_id="1", - extras={}, - ), - ], - sampling_args={"temperature": 1.0}, - error=None, - ) - - # Token 2 is the image placeholder, token 5 is the video placeholder. - mm_mapping = {2: 1, 5: 2} - rollouts = interleave_rollout(output, mm_token_type_ids_mapping=mm_mapping) - - assert rollouts is not None and len(rollouts) == 1 - sample = rollouts[0] - # Extension holds; both steps merge into one sample. mm_data is - # the union of step 0's delta ([h1]) and step 1's delta ([h2]). - assert sample.prompt_ids == [1, 2] - assert sample.completion_ids == [3, 4, 5, 6, 7] - # Pixel values packed by concatenating step 0's item then step 1's. - assert _decode_mm_pixels(sample) == [ - [1.0, 2.0], - [3.0, 4.0], - ] - assert _decode_mm_thw(sample) == [[1, 2, 3], [1, 4, 4]] - # mm_token_type_ids: image at token 2, video at token 5, rest 0. - assert sample.mm_token_type_ids == [0, 1, 0, 0, 2, 0, 0] diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 0376fd84ae8..51c458fced9 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -215,7 +215,7 @@ def test_multi_node_auto_inference_client_dp_rank_count_uses_router_url(): def test_orchestrator_vlm_requires_renderer(): - with pytest.raises(ValidationError, match="orchestrator.renderer must be set when model.vlm is set"): + with pytest.raises(ValidationError, match="renderer"): OrchestratorConfig.model_validate( { "student": { @@ -259,7 +259,7 @@ def test_shared_model_name_propagates_to_subconfigs(): { "model": {"name": model_name}, "trainer": {}, - "orchestrator": {"renderer": None}, + "orchestrator": {"renderer": {"name": "default"}}, "inference": {}, } ) @@ -276,7 +276,7 @@ def test_shared_tokenizer_propagates_when_subconfigs_unset(): "model": {"name": "my-model"}, "tokenizer": {"name": "my-tokenizer"}, "trainer": {}, - "orchestrator": {"renderer": None}, + "orchestrator": {"renderer": {"name": "default"}}, } ) assert config.trainer.tokenizer.name == "my-tokenizer" @@ -293,7 +293,7 @@ def test_shared_and_sub_tokenizer_name_conflict_raises(): "model": {"name": "my-model"}, "tokenizer": {"name": "shared-tok"}, "trainer": {"tokenizer": {"name": "trainer-tok"}}, - "orchestrator": {"renderer": None}, + "orchestrator": {"renderer": {"name": "default"}}, } ) @@ -304,7 +304,7 @@ def test_tokenizer_name_falls_back_to_model_name_when_unset(): "model": {"name": "my-model"}, "tokenizer": {"trust_remote_code": True}, "trainer": {}, - "orchestrator": {"renderer": None}, + "orchestrator": {"renderer": {"name": "default"}}, } ) assert config.trainer.tokenizer.name == "my-model" @@ -332,7 +332,7 @@ def test_explicit_subconfig_tokenizer_name_survives_shared_model_propagation(): "model": {"name": "M"}, "trainer": {}, "orchestrator": { - "renderer": None, + "renderer": {"name": "default"}, "tokenizer": {"name": "explicit-orch-tok"}, }, } @@ -351,7 +351,7 @@ def test_tokenizer_chat_template_mismatch_raises(): RLConfig.model_validate( { "trainer": {"tokenizer": {"chat_template": "A"}}, - "orchestrator": {"renderer": None, "tokenizer": {"chat_template": "B"}}, + "orchestrator": {"renderer": {"name": "default"}, "tokenizer": {"chat_template": "B"}}, } ) @@ -361,7 +361,7 @@ def test_shared_seq_len_propagates_to_subconfigs(): { "seq_len": 4096, "trainer": {}, - "orchestrator": {"renderer": None}, + "orchestrator": {"renderer": {"name": "default"}}, } ) assert config.trainer.model.seq_len == 4096 @@ -377,7 +377,7 @@ def test_shared_and_sub_seq_len_conflict_raises(): { "seq_len": 4096, "trainer": {"model": {"seq_len": 8192}}, - "orchestrator": {"renderer": None}, + "orchestrator": {"renderer": {"name": "default"}}, } ) @@ -389,7 +389,7 @@ def test_shared_and_sub_model_name_conflict_raises(): { "model": {"name": "X"}, "trainer": {"model": {"name": "Y"}}, - "orchestrator": {"renderer": None}, + "orchestrator": {"renderer": {"name": "default"}}, } ) @@ -401,7 +401,7 @@ def test_shared_and_sub_max_steps_conflict_raises(): { "max_steps": 100, "trainer": {}, - "orchestrator": {"renderer": None, "max_steps": 200}, + "orchestrator": {"renderer": {"name": "default"}, "max_steps": 200}, } ) @@ -416,7 +416,7 @@ def test_trainer_chat_template_cascades_to_inference(): { "model": {"name": "Qwen/Qwen3-0.6B"}, "trainer": {"tokenizer": {"chat_template": "TPL"}}, - "orchestrator": {"renderer": None, "tokenizer": {"chat_template": "TPL"}}, + "orchestrator": {"renderer": {"name": "default"}, "tokenizer": {"chat_template": "TPL"}}, "inference": {}, } ) @@ -442,7 +442,7 @@ def test_shared_wandb_fields_propagate_to_subconfigs(): "offline": False, }, "trainer": {}, - "orchestrator": {"renderer": None}, + "orchestrator": {"renderer": {"name": "default"}}, } ) for component in (config.trainer.wandb, config.orchestrator.wandb): @@ -462,7 +462,7 @@ def test_empty_shared_ckpt_block_does_not_conflict_with_subconfig_ckpt(): { "ckpt": {}, # empty block, no field set "trainer": {"ckpt": {"interval": 50}}, - "orchestrator": {"renderer": None, "ckpt": {"interval": 50}}, + "orchestrator": {"renderer": {"name": "default"}, "ckpt": {"interval": 50}}, } ) assert config.trainer.ckpt is not None @@ -476,7 +476,7 @@ def test_shared_and_subconfig_disjoint_fields_coexist(): { "model": {"name": "Qwen/Qwen3-0.6B"}, "trainer": {"model": {"impl": "custom"}}, - "orchestrator": {"renderer": None}, + "orchestrator": {"renderer": {"name": "default"}}, } ) assert config.trainer.model.name == "Qwen/Qwen3-0.6B" @@ -528,15 +528,15 @@ def test_orchestrator_explicit_renderer_skips_unmapped_check(): assert config.renderer.name == "qwen3" -def test_orchestrator_renderer_none_skips_unmapped_check(): - """renderer=None (MITO mode) means the renderer client isn't used, so MODEL_RENDERER_MAP doesn't apply.""" - config = OrchestratorConfig.model_validate( - { - "model": {"name": "not-a-real-org/not-a-real-model"}, - "renderer": None, - } - ) - assert config.renderer is None +def test_orchestrator_renderer_none_rejected(): + """A renderer is required (training is renderer-only): the non-optional type rejects None.""" + with pytest.raises(ValidationError, match="renderer"): + OrchestratorConfig.model_validate( + { + "model": {"name": "not-a-real-org/not-a-real-model"}, + "renderer": None, + } + ) def test_orchestrator_explicit_default_renderer_with_unmapped_model(): @@ -561,7 +561,7 @@ def test_shared_model_name_resolves_inference_parsers(): { "model": {"name": "Qwen/Qwen3-Coder-30B-A3B-Instruct"}, "trainer": {}, - "orchestrator": {"renderer": None}, + "orchestrator": {"renderer": {"name": "default"}}, "inference": {}, } ) @@ -577,7 +577,7 @@ def test_explicit_inference_parser_wins_over_auto(): { "model": {"name": "Qwen/Qwen3-Coder-30B-A3B-Instruct"}, "trainer": {}, - "orchestrator": {"renderer": None}, + "orchestrator": {"renderer": {"name": "default"}}, "inference": {"model": {"tool_call_parser": "hermes"}}, } ) diff --git a/tests/unit/train/rl/test_packer.py b/tests/unit/train/rl/test_packer.py index 1cf6beac0d6..5ed2b45ce54 100644 --- a/tests/unit/train/rl/test_packer.py +++ b/tests/unit/train/rl/test_packer.py @@ -33,8 +33,8 @@ def create_run_with_config(output_dir: Path, run_name: str) -> Path: "group_size": 1, "env": [{"id": "test-env"}], "sampling": {"temperature": 1.0}, - # test-model isn't in MODEL_RENDERER_MAP; bypass the renderer-resolution validator. - "renderer": "None", + # test-model isn't in MODEL_RENDERER_MAP; use the explicit default renderer. + "renderer": {"name": "default"}, } with open(control_dir / "orch.toml", "wb") as f: tomli_w.dump(config, f) @@ -43,12 +43,10 @@ def create_run_with_config(output_dir: Path, run_name: str) -> Path: def make_training_sample() -> TrainingSample: return TrainingSample( - prompt_ids=[1], - prompt_mask=[False], - completion_ids=[2], - completion_mask=[True], - completion_logprobs=[-0.1], - completion_temperatures=[1.0], + token_ids=[1, 2], + mask=[False, True], + logprobs=[0.0, -0.1], + temperatures=[1.0, 1.0], env_name="test-env", ) diff --git a/tests/unit/train/test_runs.py b/tests/unit/train/test_runs.py index 6bb6623c06b..cb8d6e5f9e4 100644 --- a/tests/unit/train/test_runs.py +++ b/tests/unit/train/test_runs.py @@ -42,8 +42,8 @@ def create_run_with_config( "batch_size": 32, "group_size": 4, "env": [{"id": "test-env"}], - # test-model isn't in MODEL_RENDERER_MAP; bypass the renderer-resolution validator. - "renderer": "None", + # test-model isn't in MODEL_RENDERER_MAP; use the explicit default renderer. + "renderer": {"name": "default"}, } with open(config_dir / "orch.toml", "wb") as f: @@ -203,7 +203,7 @@ def test_config_loading(tmp_path: Path) -> None: "max_steps": 1000, "group_size": 4, "env": [{"id": "test-env"}], - "renderer": "None", + "renderer": {"name": "default"}, } create_run_with_config(tmp_path, "run_test123", config=test_config) @@ -248,7 +248,7 @@ def test_config_cleanup_on_deletion(tmp_path: Path) -> None: "batch_size": 16, "group_size": 4, "env": [{"id": "test-env"}], - "renderer": "None", + "renderer": {"name": "default"}, } run_dir = create_run_with_config(tmp_path, "run_delete_me", config=test_config) diff --git a/tests/unit/utils/test_client.py b/tests/unit/utils/test_client.py index 69325e6c4a3..40de4cfee66 100644 --- a/tests/unit/utils/test_client.py +++ b/tests/unit/utils/test_client.py @@ -3,7 +3,7 @@ from unittest.mock import AsyncMock, MagicMock import httpx -import verifiers as vf +from verifiers.v1.clients.config import EvalClientConfig from prime_rl.configs.shared import ClientConfig from prime_rl.utils.client import _is_retryable_lora_error, load_lora_adapter, setup_clients @@ -67,13 +67,12 @@ def test_setup_clients_assigns_renderer_and_dp_rank_headers(): renderer_config=renderer_settings, ) - assert [client.client_type for client in clients] == ["renderer", "renderer"] - assert [client.renderer_config for client in clients] == [renderer_settings, renderer_settings] + assert [client.type for client in clients] == ["train", "train"] + assert [client.renderer for client in clients] == [renderer_settings, renderer_settings] assert [client.renderer_model_name for client in clients] == [None, None] - assert [client.api_base_url for client in clients] == ["http://worker-a:8000/v1"] * 2 - assert [client.extra_headers["X-data-parallel-rank"] for client in clients] == ["0", "1"] - assert clients[0].extra_headers["X-Test"] == "test" - assert clients[0].extra_headers_from_state == {"X-Session-ID": "session_id"} + assert [client.base_url for client in clients] == ["http://worker-a:8000/v1"] * 2 + assert [client.headers["X-data-parallel-rank"] for client in clients] == ["0", "1"] + assert clients[0].headers["X-Test"] == "test" def test_setup_clients_assigns_renderer_model_name(): @@ -103,17 +102,9 @@ def test_setup_clients_preserves_chat_client_defaults(): clients = setup_clients(client_config) assert clients == [ - vf.ClientConfig( - client_idx=0, - client_type="openai_chat_completions", + EvalClientConfig( api_key_var="PRIME_API_KEY", - api_base_url="http://worker-a:8000/v1", - timeout=client_config.timeout, - connect_timeout=client_config.connect_timeout, - max_connections=8192, - max_keepalive_connections=8192, - max_retries=10, - extra_headers={}, - extra_headers_from_state={}, + base_url="http://worker-a:8000/v1", + headers={}, ) ] diff --git a/tests/unit/utils/test_elastic.py b/tests/unit/utils/test_elastic.py index 21490497e79..01e4afb91e3 100644 --- a/tests/unit/utils/test_elastic.py +++ b/tests/unit/utils/test_elastic.py @@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx -import verifiers as vf +from verifiers.v1.clients.config import TrainClientConfig from prime_rl.utils.elastic import ( AdapterState, @@ -411,8 +411,6 @@ def test_elastic_clients_preserve_renderer_model_name_when_model_name_updates(): client_config.elastic.port = 8000 client_config.elastic.sync_interval = 5.0 client_config.router_url = None - client_config.timeout = 1200 - client_config.connect_timeout = 30.0 client_config.api_key_var = "PRIME_API_KEY" client_config.headers = {} client_config.headers_from_env = {} @@ -435,20 +433,14 @@ def test_elastic_clients_preserve_renderer_model_name_when_model_name_updates(): pool.update_model_name("r8-smoke") clients = pool.train_clients + # `renderer_model_name` stays pinned to the original base model even after the + # served model name flips to the LoRA adapter (`r8-smoke`). assert clients == [ - vf.ClientConfig( - client_idx=0, - client_type="renderer", - renderer_config=renderer_settings, + TrainClientConfig( + renderer=renderer_settings, renderer_model_name="Qwen/Qwen3-VL-4B-Instruct", api_key_var="PRIME_API_KEY", - api_base_url="http://10.0.0.1:8000/v1", - timeout=1200, - connect_timeout=30.0, - max_connections=8192, - max_keepalive_connections=8192, - max_retries=10, - extra_headers={}, - extra_headers_from_state={}, + base_url="http://10.0.0.1:8000/v1", + headers={}, ) ] diff --git a/tests/unit/utils/test_prime_monitor.py b/tests/unit/utils/test_prime_monitor.py index f44065a7c68..aaccbff60d6 100644 --- a/tests/unit/utils/test_prime_monitor.py +++ b/tests/unit/utils/test_prime_monitor.py @@ -3,7 +3,9 @@ from unittest.mock import Mock import pyarrow.parquet as pq +import verifiers.v1 as vf +from prime_rl.orchestrator.types import Rollout from prime_rl.utils.monitor.prime import PrimeMonitor @@ -13,32 +15,34 @@ def _new_monitor() -> PrimeMonitor: return monitor -def _build_rollout(*, example_id: int, reward: float, task: str) -> dict: - return { - "example_id": example_id, - "prompt": [{"role": "user", "content": f"prompt-{example_id}"}], - "completion": [{"role": "assistant", "content": f"completion-{example_id}"}], - "trajectory": [ - { - "prompt": [{"role": "user", "content": f"prompt-{example_id}"}], - "completion": [{"role": "assistant", "content": f"completion-{example_id}"}], - "reward": reward, - "advantage": reward / 2, - "extras": {"source": "test"}, - "tokens": { - "prompt_ids": [1, 2, 3], - "completion_ids": [4, 5], - }, - } - ], - "answer": f"answer-{example_id}", - "task": task, - "info": {"difficulty": "easy"}, - "reward": reward, - "advantage": reward / 2, - "metrics": {"accuracy": reward}, - "timing": {"generation": {"start": 0.0, "end": 12.5, "duration": 12.5}}, - } +def _build_rollout(*, example_id: int, reward: float, task: str) -> Rollout: + """Build a v1 ``Rollout`` (message-graph trace). The user node carries the prompt and the + assistant node the completion; ``_rollouts_to_parquet_bytes`` reads the conversation off the + branches (its ``completion`` column is the last branch's messages, ``trajectory`` is one + message list per branch).""" + nodes = [ + vf.MessageNode( + message=vf.UserMessage(content=f"prompt-{example_id}"), + token_ids=[1, 2, 3], + mask=[False, False, False], + logprobs=[0.0, 0.0, 0.0], + ), + vf.MessageNode( + message=vf.AssistantMessage(content=f"completion-{example_id}"), + token_ids=[4, 5], + mask=[True, True], + logprobs=[-0.1, -0.2], + sampled=True, + ), + ] + rollout = Rollout[vf.Task]( + task=vf.Task(idx=example_id, prompt=f"prompt-{example_id}"), + nodes=nodes, + rewards={"reward": reward}, + ) + rollout.env_name = task + rollout.advantage = reward / 2 + return rollout def test_rollouts_to_parquet_bytes_preserves_all_rollouts_and_ids(): @@ -63,24 +67,22 @@ def test_rollouts_to_parquet_bytes_preserves_all_rollouts_and_ids(): assert [row["sample_id"] for row in rows] == [0, 1] assert all(row["run_id"] == "run-123" for row in rows) assert all(row["step"] == 7 for row in rows) - assert json.loads(rows[0]["prompt"])[0]["content"] == "prompt-101" + # `completion` is the last branch's messages; the prompt user message lives in `trajectory`. assert json.loads(rows[1]["completion"])[0]["content"] == "completion-202" + trajectory = json.loads(rows[0]["trajectory"]) + assert trajectory[0]["messages"][0]["content"] == "prompt-101" def test_rollouts_to_parquet_bytes_skips_rollouts_without_trajectory(): monitor = _new_monitor() monitor.run_id = "run-456" + rollout_with_branches = _build_rollout(example_id=1, reward=1.0, task="task-a") + rollout_without_branches = Rollout[vf.Task](task=vf.Task(idx=2, prompt="missing-trajectory")) + assert rollout_without_branches.branches == [] + parquet_bytes = monitor._rollouts_to_parquet_bytes( - [ - _build_rollout(example_id=1, reward=1.0, task="task-a"), - { - "example_id": 2, - "prompt": [{"role": "user", "content": "missing-trajectory"}], - "completion": [{"role": "assistant", "content": "ignored"}], - "trajectory": [], - }, - ], + [rollout_with_branches, rollout_without_branches], step=3, ) diff --git a/uv.lock b/uv.lock index f8be90899ec..db123c0f544 100644 --- a/uv.lock +++ b/uv.lock @@ -11,7 +11,7 @@ supported-markers = [ ] [options] -exclude-newer = "2026-06-10T23:37:49.383138959Z" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P7D" [options.exclude-newer-package] @@ -26,12 +26,10 @@ flash-attn-3 = false prime-tunnel = false deep-gemm = false prime-evals = false -tasksets = false tokenspeed-mla = false deep-ep = false prime-sandboxes = false renderers = false -harnesses = false prime = false [manifest] @@ -41,7 +39,7 @@ overrides = [ { name = "openenv-core" }, { name = "torch", specifier = ">=2.9.0", index = "https://download.pytorch.org/whl/cu128" }, { name = "transformers", specifier = "==5.6.2" }, - { name = "verifiers", extras = ["packages"], editable = "deps/verifiers" }, + { name = "verifiers", editable = "deps/verifiers" }, ] [[manifest.dependency-metadata]] @@ -73,7 +71,7 @@ version = "0.2.0" source = { editable = "deps/research-environments/environments/aime2024" } dependencies = [ { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -88,7 +86,7 @@ version = "0.2.0" source = { editable = "deps/research-environments/environments/aime2025" } dependencies = [ { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -98,15 +96,18 @@ requires-dist = [ ] [[package]] -name = "aiofile" -version = "3.11.1" -source = { registry = "https://pypi.org/simple" } +name = "aime24-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/aime24_v1" } dependencies = [ - { name = "caio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/48/41/2fea7e193e061ce54eacc3b7bc0e6a99e4fcff43c78cf0a76dd781ed8334/aiofile-3.11.1.tar.gz", hash = "sha256:1f91912c6643d2a4e49ca4ae3514f0bf3867ce948a36d99a6411b8f4755f4cf9", size = 19342, upload-time = "2026-05-16T08:18:33.538Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/67/cd/0d76dfc5de72bde52f55f53e925c7d152d9c7906634ec1e0cbc7e8d4ad93/aiofile-3.11.1-py3-none-any.whl", hash = "sha256:ce77d14ac07f77bc2b757834a5c129321f3f705c474593deed5ab209079a52c9", size = 20446, upload-time = "2026-05-16T08:18:32.051Z" }, + +[package.metadata] +requires-dist = [ + { name = "datasets" }, + { name = "verifiers" }, ] [[package]] @@ -175,12 +176,27 @@ name = "alphabet-sort" version = "0.1.12" source = { editable = "deps/verifiers/environments/alphabet_sort" } dependencies = [ - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] requires-dist = [{ name = "verifiers", specifier = ">=0.1.9" }] +[[package]] +name = "alphabet-sort-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/alphabet_sort_v1" } +dependencies = [ + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "datasets" }, + { name = "verifiers" }, +] + [[package]] name = "annotated-doc" version = "0.0.4" @@ -291,19 +307,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, ] -[[package]] -name = "authlib" -version = "1.7.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cryptography", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "joserfc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/36/98/7d93f30d029643c0275dbc0bd6d5a6f670661ee6c9a94d93af7ab4887600/authlib-1.7.2.tar.gz", hash = "sha256:2cea25fefcd4e7173bdf1372c0afc265c8034b23a8cd5dcb6a9164b826c64231", size = 176511, upload-time = "2026-05-06T08:10:23.116Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fb/95/adcb68e20c34162e9135f370d6e31737719c2b6f94bc953fe7ed1f10fe21/authlib-1.7.2-py2.py3-none-any.whl", hash = "sha256:3e1faedc9d87e7d56a164eca3ccb6ace0d61b94abe83e92242f8dc8bba9b4a9f", size = 259548, upload-time = "2026-05-06T08:10:21.436Z" }, -] - [[package]] name = "backoff" version = "2.2.1" @@ -384,18 +387,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0e/c6/d1fe8bdea4a6088bd54b5a58bc40aed89a4e784cd796af7722a06f74bae7/blake3-1.0.8-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:a25db3d36b55f5ed6a86470155cc749fc9c5b91c949b8d14f48658f9d960d9ec", size = 554211, upload-time = "2025-10-14T06:46:00.269Z" }, ] -[[package]] -name = "brotli" -version = "1.2.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f7/16/c92ca344d646e71a43b8bb353f0a6490d7f6e06210f8554c8f874e454285/brotli-1.2.0.tar.gz", hash = "sha256:e310f77e41941c13340a95976fe66a8a95b01e783d430eeaf7a2f87e0a57dd0a", size = 7388632, upload-time = "2025-11-05T18:39:42.86Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3d/6f/feba03130d5fceadfa3a1bb102cb14650798c848b1df2a808356f939bb16/brotli-1.2.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:acec55bb7c90f1dfc476126f9711a8e81c9af7fb617409a9ee2953115343f08d", size = 1528071, upload-time = "2025-11-05T18:38:26.081Z" }, - { url = "https://files.pythonhosted.org/packages/03/a7/03aa61fbc3c5cbf99b44d158665f9b0dd3d8059be16c460208d9e385c837/brotli-1.2.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:072e7624b1fc4d601036ab3f4f27942ef772887e876beff0301d261210bca97f", size = 1419762, upload-time = "2025-11-05T18:38:28.295Z" }, - { url = "https://files.pythonhosted.org/packages/21/1b/0374a89ee27d152a5069c356c96b93afd1b94eae83f1e004b57eb6ce2f10/brotli-1.2.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:adedc4a67e15327dfdd04884873c6d5a01d3e3b6f61406f99b1ed4865a2f6d28", size = 1484494, upload-time = "2025-11-05T18:38:29.29Z" }, - { url = "https://files.pythonhosted.org/packages/d5/3b/39e13ce78a8e9a621c5df3aeb5fd181fcc8caba8c48a194cd629771f6828/brotli-1.2.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:af43b8711a8264bb4e7d6d9a6d004c3a2019c04c01127a868709ec29962b6036", size = 1487913, upload-time = "2025-11-05T18:38:31.618Z" }, -] - [[package]] name = "build" version = "1.5.0" @@ -430,18 +421,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8c/7b/1fc1c09cc0756cf25861a3be10565915953876da48bb228fb9a672b20a42/cachetools-7.1.4-py3-none-any.whl", hash = "sha256:323dc4127934744db5b54eb4924482d7edafbf9554e820d1531c2e08c0e4ef54", size = 16761, upload-time = "2026-05-21T22:40:41.845Z" }, ] -[[package]] -name = "caio" -version = "0.9.25" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/92/88/b8527e1b00c1811db339a1df8bd1ae49d146fcea9d6a5c40e3a80aaeb38d/caio-0.9.25.tar.gz", hash = "sha256:16498e7f81d1d0f5a4c0ad3f2540e65fe25691376e0a5bd367f558067113ed10", size = 26781, upload-time = "2025-12-26T15:21:36.501Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a3/2b/21288691f16d479945968a0a4f2856818c1c5be56881d51d4dac9b255d26/caio-0.9.25-cp312-cp312-manylinux2010_x86_64.manylinux2014_x86_64.manylinux_2_12_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:97084e4e30dfa598449d874c4d8e0c8d5ea17d2f752ef5e48e150ff9d240cd64", size = 82012, upload-time = "2025-12-26T15:22:20.983Z" }, - { url = "https://files.pythonhosted.org/packages/03/c4/8a1b580875303500a9c12b9e0af58cb82e47f5bcf888c2457742a138273c/caio-0.9.25-cp312-cp312-manylinux_2_34_aarch64.whl", hash = "sha256:4fa69eba47e0f041b9d4f336e2ad40740681c43e686b18b191b6c5f4c5544bfb", size = 81502, upload-time = "2026-03-04T22:08:22.381Z" }, - { url = "https://files.pythonhosted.org/packages/d1/1c/0fe770b8ffc8362c48134d1592d653a81a3d8748d764bec33864db36319d/caio-0.9.25-cp312-cp312-manylinux_2_34_x86_64.whl", hash = "sha256:6bebf6f079f1341d19f7386db9b8b1f07e8cc15ae13bfdaff573371ba0575d69", size = 80200, upload-time = "2026-03-04T22:08:23.382Z" }, - { url = "https://files.pythonhosted.org/packages/86/93/1f76c8d1bafe3b0614e06b2195784a3765bbf7b0a067661af9e2dd47fc33/caio-0.9.25-py3-none-any.whl", hash = "sha256:06c0bb02d6b929119b1cfbe1ca403c768b2013a369e2db46bfa2a5761cf82e40", size = 19087, upload-time = "2025-12-26T15:22:00.221Z" }, -] - [[package]] name = "cbor2" version = "6.1.1" @@ -580,7 +559,7 @@ version = "0.3.2" source = { editable = "deps/research-environments/environments/code_env" } dependencies = [ { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -592,13 +571,24 @@ requires-dist = [ ] provides-extras = ["dev"] +[[package]] +name = "code-golf-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/code_golf_v1" } +dependencies = [ + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [{ name = "verifiers" }] + [[package]] name = "color-codeword" version = "0.1.0" source = { editable = "deps/research-environments/environments/color_codeword" } dependencies = [ { name = "pillow", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -607,6 +597,21 @@ requires-dist = [ { name = "verifiers", specifier = ">=0.1.15.dev17" }, ] +[[package]] +name = "color-codeword-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/color_codeword_v1" } +dependencies = [ + { name = "pillow", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "pillow" }, + { name = "verifiers" }, +] + [[package]] name = "comm" version = "0.2.3" @@ -781,21 +786,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, ] -[[package]] -name = "cyclopts" -version = "4.16.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "attrs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "docstring-parser", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "rich-rst", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/34/07/bf61d13de86d96a4c46aff00c9ca0eced44bcc8c3e16280605c1253e5720/cyclopts-4.16.1.tar.gz", hash = "sha256:8aa47bf92a5fb33abca5af05e576eecdb0d2f79893ad29238046df78370fc4a8", size = 181196, upload-time = "2026-05-25T15:29:08.518Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/00/8d/7f362c2fb8ef4decd2160bc24d4292c6ca658cc6d9a161b89ca5122bbdbf/cyclopts-4.16.1-py3-none-any.whl", hash = "sha256:617795392c4113a2c2cc7af716f20244900e87f23daa05442d1268d81472a592", size = 219020, upload-time = "2026-05-25T15:29:09.646Z" }, -] - [[package]] name = "dataclasses-json" version = "0.6.7" @@ -891,7 +881,7 @@ dependencies = [ { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "diskcache", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "pdfminer-six", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -903,6 +893,17 @@ requires-dist = [ { name = "verifiers", specifier = ">=0.1.15.dev17" }, ] +[[package]] +name = "deepwiki-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/deepwiki_v1" } +dependencies = [ + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [{ name = "verifiers" }] + [[package]] name = "depyf" version = "0.20.0" @@ -1033,18 +1034,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/de/15/545e2b6cf2e3be84bc1ed85613edd75b8aea69807a71c26f4ca6a9258e82/email_validator-2.3.0-py3-none-any.whl", hash = "sha256:80f13f623413e6b197ae73bb10bf4eb0908faf509ad8362c5edeb0be7fd450b4", size = 35604, upload-time = "2025-08-26T13:09:05.858Z" }, ] -[[package]] -name = "exceptiongroup" -version = "1.3.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/8a/0e/97c33bf5009bdbac74fd2beace167cab3f978feb69cc36f1ef79360d6c4e/exceptiongroup-1.3.1-py3-none-any.whl", hash = "sha256:a7a39a3bd276781e98394987d3a5701d0c4edffb633bb7a5144577f82c773598", size = 16740, upload-time = "2025-11-21T23:01:53.443Z" }, -] - [[package]] name = "executing" version = "2.2.1" @@ -1153,66 +1142,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/41/2c368f804bb9bd918da3b61324207fc4b410d0f32352c372c0680fc1f670/fastcore-1.13.2-py3-none-any.whl", hash = "sha256:2103c9e9e613311c0b36eab17299a221e778fd214ec526e8df1d32908928277c", size = 105060, upload-time = "2026-05-17T06:02:22.28Z" }, ] -[[package]] -name = "fastmcp" -version = "3.3.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "fastmcp-slim", extra = ["client", "server"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/3b/a9/5c5a01b6abd5346bf60b97cfd29e4a86661940c27dd562bfcda07fd03519/fastmcp-3.3.1.tar.gz", hash = "sha256:979362ea557de42a5f40342563c7e4b236bcc8e7cd192715f50030695d1a71cd", size = 28681699, upload-time = "2026-05-15T15:50:39.673Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9f/11/6b1bdada6ccfe647d615ae63f9106f8136aec17971e9361546af01c7d38e/fastmcp-3.3.1-py3-none-any.whl", hash = "sha256:862440c5c4d281363a5995eee59d77f0f7cac1f18869038729cecf03b02fc522", size = 7903, upload-time = "2026-05-15T15:50:36.424Z" }, -] - -[[package]] -name = "fastmcp-slim" -version = "3.3.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "platformdirs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pydantic", extra = ["email"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pydantic-settings", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "python-dotenv", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d1/a0/627103e517e1d0d6f1eec633d5662d13e776f01b45ad188e4f5f7478b438/fastmcp_slim-3.3.1.tar.gz", hash = "sha256:0957835fc59452e143ab2f4b7836d2d2df9b2d9958408edc79ba8b56232b2a88", size = 567007, upload-time = "2026-05-15T15:50:10.426Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7a/ee/97047f4cc2d7b1d46670d08d8ad01a96e7a748cc01c0b4b351ad8eddbc7a/fastmcp_slim-3.3.1-py3-none-any.whl", hash = "sha256:6cf1c2d77e3adb0d409d6825ed6b0b2a999062973e00b8eea03bd48bf9b4c043", size = 738644, upload-time = "2026-05-15T15:50:08.336Z" }, -] - -[package.optional-dependencies] -client = [ - { name = "authlib", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "exceptiongroup", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "mcp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "opentelemetry-api", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "py-key-value-aio", extra = ["filetree", "keyring", "memory"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -server = [ - { name = "authlib", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "cyclopts", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "exceptiongroup", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "griffelib", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "jsonref", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "jsonschema-path", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "mcp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "openapi-pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "opentelemetry-api", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "py-key-value-aio", extra = ["filetree", "keyring", "memory"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pyperclip", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "python-multipart", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "uncalled-for", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "uvicorn", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "watchfiles", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "websockets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] - [[package]] name = "fastokens" version = "0.2.0" @@ -1308,13 +1237,11 @@ dependencies = [ [[package]] name = "flash-linear-attention" -version = "0.5.1" -source = { git = "https://github.com/fla-org/flash-linear-attention#e8e791f314a2210e94dde06214ac8f1f0b35e7e3" } +version = "0.5.2" +source = { git = "https://github.com/fla-org/flash-linear-attention#bd67d11782696e4ff54ee6cb411483621ed3d93d" } dependencies = [ { name = "einops", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "torch", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "transformers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "triton", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [[package]] @@ -1419,7 +1346,7 @@ source = { editable = "deps/research-environments/environments/general_agent" } dependencies = [ { name = "mcp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "tyro", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -1434,6 +1361,23 @@ requires-dist = [ ] provides-extras = ["dev", "test"] +[[package]] +name = "general-agent-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/general_agent_v1" } +dependencies = [ + { name = "filelock", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "filelock" }, + { name = "pydantic" }, + { name = "verifiers" }, +] + [[package]] name = "gepa" version = "0.1.1" @@ -1501,43 +1445,15 @@ wheels = [ ] [[package]] -name = "google-auth" -version = "2.53.0" -source = { registry = "https://pypi.org/simple" } +name = "glossary-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/glossary_v1" } dependencies = [ - { name = "cryptography", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pyasn1-modules", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/c6/ad/ff781329bbbdc0974a098d996e89c9e1f7024262f9e3eec442fbb9ad1ac6/google_auth-2.53.0.tar.gz", hash = "sha256:e7e6aa16f6bee7b2b264830fd04f08087a1d5a836df516251a5d15327b246c9c", size = 335844, upload-time = "2026-05-15T20:53:07.928Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/4a/c9/db44165ba7c581268c6d46017ef63339110378305062830104fc7fa144cb/google_auth-2.53.0-py3-none-any.whl", hash = "sha256:6e7449917c599b35126a99ec268ec6880301f2fea41dce198fe8fd83ff642b68", size = 246071, upload-time = "2026-05-15T20:53:05.609Z" }, -] - -[package.optional-dependencies] -requests = [ - { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -[[package]] -name = "google-genai" -version = "2.7.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "anyio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "distro", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "google-auth", extra = ["requests"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "sniffio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "tenacity", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "websockets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/a7/7b/6eb3b3d545b6bb4c374acba1ccf91b0f33b605e551536a6243cfcef2f07f/google_genai-2.7.0.tar.gz", hash = "sha256:3c6f32f5ced9877ededd1b384b5e5b7f09c20046ec3390b662b16d8cd1882ac5", size = 555853, upload-time = "2026-05-28T15:39:24.58Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3c/dd/7a8be39e9d698e80e9db796514efbc6083dbd787bdb9a101e8ba47248e5e/google_genai-2.7.0-py3-none-any.whl", hash = "sha256:21cac381e09a869151706aba797b6a4f96cfe92c484e13204d092caee7ff11cb", size = 822545, upload-time = "2026-05-28T15:39:22.907Z" }, -] +[package.metadata] +requires-dist = [{ name = "verifiers" }] [[package]] name = "googleapis-common-protos" @@ -1557,7 +1473,7 @@ version = "0.1.5" source = { editable = "deps/research-environments/environments/gpqa" } dependencies = [ { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -1566,60 +1482,6 @@ requires-dist = [ { name = "verifiers", specifier = ">=0.1.15.dev17" }, ] -[[package]] -name = "gradio" -version = "6.15.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "anyio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "brotli", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "fastapi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "gradio-client", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "groovy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "hf-gradio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "huggingface-hub", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "jinja2", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "markupsafe", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "orjson", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pandas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pillow", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pydub", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "python-multipart", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pytz", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "safehttpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "semantic-version", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "starlette", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "tomlkit", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "typer", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "uvicorn", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/b9/7a/edc719d67beea0963721a7f552604faa0feaf218d4a51c6e0dacfb51ba6a/gradio-6.15.1.tar.gz", hash = "sha256:58be31be7b3aab53bbe61f21c20666b4a8a25a6737c399e02b7463d669625851", size = 36429761, upload-time = "2026-05-27T13:20:35.232Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/00/87/b6d08b4b0a5bf7bb219c4c44d1d267ce6dfbc1ff062ca05ce251a50d18a5/gradio-6.15.1-py3-none-any.whl", hash = "sha256:f4f50488f8da1137b8e0d65fe656348b11bac3cd2f5ccab60636eb45e5a6f39f", size = 20093799, upload-time = "2026-05-27T13:20:31.18Z" }, -] - -[[package]] -name = "gradio-client" -version = "2.5.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "fsspec", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "huggingface-hub", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "packaging", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e8/e6/6b6029f5fe2ad7f1211105d530e34d991014c2cae463f9223033031cfc4f/gradio_client-2.5.0.tar.gz", hash = "sha256:4cde99bad62149595c30c90876ca2e405e3a13687ecf895474f3412cb476673d", size = 59013, upload-time = "2026-04-20T23:16:21.518Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/78/81/0a861b8e1ff42960139c6cd4c7dd591292fa09ea1ae2d87677441cba4c00/gradio_client-2.5.0-py3-none-any.whl", hash = "sha256:d43e2179c29076292a76485ad7ed2e6eaa19d14ac58283bd7f5beabfe4ca958c", size = 59952, upload-time = "2026-04-20T23:16:20.186Z" }, -] - [[package]] name = "griffelib" version = "2.0.2" @@ -1629,15 +1491,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/11/8c/c9138d881c79aa0ea9ed83cbd58d5ca75624378b38cee225dcf5c42cc91f/griffelib-2.0.2-py3-none-any.whl", hash = "sha256:925c857658fb1ba40c0772c37acbc2ab650bd794d9c1b9726922e36ea4117ea1", size = 142357, upload-time = "2026-03-27T11:34:46.275Z" }, ] -[[package]] -name = "groovy" -version = "0.1.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/52/36/bbdede67400277bef33d3ec0e6a31750da972c469f75966b4930c753218f/groovy-0.1.2.tar.gz", hash = "sha256:25c1dc09b3f9d7e292458aa762c6beb96ea037071bf5e917fc81fb78d2231083", size = 17325, upload-time = "2025-02-28T20:24:56.068Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/28/27/3d6dcadc8a3214d8522c1e7f6a19554e33659be44546d44a2f7572ac7d2a/groovy-0.1.2-py3-none-any.whl", hash = "sha256:7f7975bab18c729a257a8b1ae9dcd70b7cafb1720481beae47719af57c35fa64", size = 14090, upload-time = "2025-02-28T20:24:55.152Z" }, -] - [[package]] name = "grpcio" version = "1.80.0" @@ -1666,6 +1519,17 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5c/90/b0cbbd9efcc82816c58f31a34963071aa19fb792a212a5d9caf8e0fc3097/grpclib-0.4.9-py3-none-any.whl", hash = "sha256:7762ec1c8ed94dfad597475152dd35cbd11aecaaca2f243e29702435ca24cf0e", size = 77063, upload-time = "2025-12-14T22:23:13.224Z" }, ] +[[package]] +name = "gsm8k-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/gsm8k_v1" } +dependencies = [ + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [{ name = "datasets" }] + [[package]] name = "gymnasium" version = "1.3.0" @@ -1703,34 +1567,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/69/b2/119f6e6dcbd96f9069ce9a2665e0146588dc9f88f29549711853645e736a/h2-4.3.0-py3-none-any.whl", hash = "sha256:c438f029a25f7945c69e0ccf0fb951dc3f73a5f6412981daee861431b70e2bdd", size = 61779, upload-time = "2025-08-23T18:12:17.779Z" }, ] -[[package]] -name = "harnesses" -source = { editable = "deps/verifiers/packages/harnesses" } -dependencies = [ - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] - -[package.metadata] -requires-dist = [ - { name = "aiohttp", marker = "extra == 'nemogym'", specifier = ">=3.9.0" }, - { name = "nemo-gym", marker = "python_full_version >= '3.12' and extra == 'nemogym'", specifier = ">=0.2.1,<0.3" }, - { name = "verifiers", specifier = ">=0.1.15.dev158" }, -] -provides-extras = ["nemogym"] - -[[package]] -name = "hf-gradio" -version = "0.4.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "gradio-client", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "typer", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/ce/86/c9694b7cfada5780e75769e60dc161a161f4dd7fc91b61db5e3a3338bef9/hf_gradio-0.4.1.tar.gz", hash = "sha256:a017d942618f0d495a58ee4563047fa04bef614c00e0cb789a9a6d0633cffa7b", size = 6560, upload-time = "2026-04-22T14:01:32.334Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/30/2d/afff2ee87e75d8eb85c92bb8cf0e15b05c23c2ebd8fd8dec781d8601ed7f/hf_gradio-0.4.1-py3-none-any.whl", hash = "sha256:76b8cb8be6abe62d74c1ad2d35b42f0629db89aa9e1a8d033cecfe7c856eeab3", size = 4482, upload-time = "2026-04-17T19:53:31.827Z" }, -] - [[package]] name = "hf-xet" version = "1.5.0" @@ -1749,7 +1585,7 @@ version = "0.2.1" source = { editable = "deps/research-environments/environments/hle" } dependencies = [ { name = "pillow", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -1900,7 +1736,7 @@ dependencies = [ { name = "immutabledict", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "langdetect", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nltk", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -2043,39 +1879,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/56/6d/0d9848617b9f753b87f214f1c682592f7ca42de085f564352f10f0843026/ipywidgets-8.1.8-py3-none-any.whl", hash = "sha256:ecaca67aed704a338f88f67b1181b58f821ab5dc89c1f0f5ef99db43c1c2921e", size = 139808, upload-time = "2025-11-01T21:18:10.956Z" }, ] -[[package]] -name = "jaraco-classes" -version = "3.4.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "more-itertools", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/06/c0/ed4a27bc5571b99e3cff68f8a9fa5b56ff7df1c2251cc715a652ddd26402/jaraco.classes-3.4.0.tar.gz", hash = "sha256:47a024b51d0239c0dd8c8540c6c7f484be3b8fcf0b2d85c13825780d3b3f3acd", size = 11780, upload-time = "2024-03-31T07:27:36.643Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7f/66/b15ce62552d84bbfcec9a4873ab79d993a1dd4edb922cbfccae192bd5b5f/jaraco.classes-3.4.0-py3-none-any.whl", hash = "sha256:f662826b6bed8cace05e7ff873ce0f9283b5c924470fe664fff1c2f00f581790", size = 6777, upload-time = "2024-03-31T07:27:34.792Z" }, -] - -[[package]] -name = "jaraco-context" -version = "6.1.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/af/50/4763cd07e722bb6285316d390a164bc7e479db9d90daa769f22578f698b4/jaraco_context-6.1.2.tar.gz", hash = "sha256:f1a6c9d391e661cc5b8d39861ff077a7dc24dc23833ccee564b234b81c82dfe3", size = 16801, upload-time = "2026-03-20T22:13:33.922Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f2/58/bc8954bda5fcda97bd7c19be11b85f91973d67a706ed4a3aec33e7de22db/jaraco_context-6.1.2-py3-none-any.whl", hash = "sha256:bf8150b79a2d5d91ae48629d8b427a8f7ba0e1097dd6202a9059f29a36379535", size = 7871, upload-time = "2026-03-20T22:13:32.808Z" }, -] - -[[package]] -name = "jaraco-functools" -version = "4.5.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "more-itertools", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/36/cf/ea4ef2920830dea3f5ab2ea4da6fb67724e6dca80ee2553788c3607243d0/jaraco_functools-4.5.0.tar.gz", hash = "sha256:3bb5665ea4a020cf78a7040e89154c77edadb3ca74f366479669c5999aa70b03", size = 20272, upload-time = "2026-05-15T21:34:10.025Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/96/9a/982e48afcffcd727a9144506720ffd4224b6b7e355c98641866f38b7c043/jaraco_functools-4.5.0-py3-none-any.whl", hash = "sha256:79ce39246eddbde4b3a03b77ea5f0f7878dc669b166a66cf3fa8e266aa3fa2f4", size = 10594, upload-time = "2026-05-15T21:34:08.595Z" }, -] - [[package]] name = "jaxtyping" version = "0.3.10" @@ -2100,15 +1903,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9a/93/242e2eab5fe682ffcb8b0084bde703a41d51e17ee0f3a31ff0d9d813620a/jedi-0.20.0-py2.py3-none-any.whl", hash = "sha256:7bdd9c2634f56713299976f4cbd59cb3fa92165cc5e05ea811fb253480728b67", size = 4884812, upload-time = "2026-05-01T23:38:43.919Z" }, ] -[[package]] -name = "jeepney" -version = "0.9.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/7b/6f/357efd7602486741aa73ffc0617fb310a29b588ed0fd69c2399acbb85b0c/jeepney-0.9.0.tar.gz", hash = "sha256:cf0e9e845622b81e4a28df94c40345400256ec608d0e55bb8a3feaa9163f5732", size = 106758, upload-time = "2025-02-27T18:51:01.684Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b2/a3/e137168c9c44d18eff0376253da9f1e9234d0239e0ee230d2fee6cea8e55/jeepney-0.9.0-py3-none-any.whl", hash = "sha256:97e5714520c16fc0a45695e5365a2e11b81ea79bba796e26f9f1d178cb182683", size = 49010, upload-time = "2025-02-27T18:51:00.104Z" }, -] - [[package]] name = "jinja2" version = "3.1.6" @@ -2153,27 +1947,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl", hash = "sha256:5fc3c5039fc5ca8c0276333a188bbd59d6b7ab37fe6632daa76bc7f9ec18e713", size = 309071, upload-time = "2025-12-15T08:41:44.973Z" }, ] -[[package]] -name = "joserfc" -version = "1.6.8" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cryptography", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/5d/ac/d4fd5b30f82900eac60d765f179f0ba005825ac462cc8ced6e13ec685ab3/joserfc-1.6.8.tar.gz", hash = "sha256:878620c553a6ebdd76ccdc356782fee3f735f21a356d079a546b42a4670ace5f", size = 232930, upload-time = "2026-05-27T03:22:37.819Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/98/8c/5cdce2cf3ce8155849baf9a5e2ce77e89dc87ec3bdb38259e5d85fbc45bd/joserfc-1.6.8-py3-none-any.whl", hash = "sha256:22fb31a69094a5e6f44632002a9df2c30c941fc6c8ce1b037e92c03de954cf9f", size = 70927, upload-time = "2026-05-27T03:22:35.796Z" }, -] - -[[package]] -name = "jsonref" -version = "1.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/aa/0d/c1f3277e90ccdb50d33ed5ba1ec5b3f0a242ed8c1b1a85d3afeb68464dca/jsonref-1.1.0.tar.gz", hash = "sha256:32fe8e1d85af0fdefbebce950af85590b22b60f9e95443176adbde4e1ecea552", size = 8814, upload-time = "2023-01-16T16:10:04.455Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/0c/ec/e1db9922bceb168197a558a2b8c03a7963f1afe93517ddd3cf99f202f996/jsonref-1.1.0-py3-none-any.whl", hash = "sha256:590dc7773df6c21cbf948b5dac07a72a251db28b0238ceecce0a2abfa8ec30a9", size = 9425, upload-time = "2023-01-16T16:10:02.255Z" }, -] - [[package]] name = "jsonschema" version = "4.26.0" @@ -2189,21 +1962,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/69/90/f63fb5873511e014207a475e2bb4e8b2e570d655b00ac19a9a0ca0a385ee/jsonschema-4.26.0-py3-none-any.whl", hash = "sha256:d489f15263b8d200f8387e64b4c3a75f06629559fb73deb8fdfb525f2dab50ce", size = 90630, upload-time = "2026-01-07T13:41:05.306Z" }, ] -[[package]] -name = "jsonschema-path" -version = "0.5.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "attrs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pathable", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "referencing", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/39/79/cd02a4df6d9270efdc7d3feefe6edd730b0820c39eeaa107a2faee8322d5/jsonschema_path-0.5.0.tar.gz", hash = "sha256:493b156ba895c97602655b620a8456caa2ce08c1aa389f5a7addec065e6e855c", size = 19597, upload-time = "2026-05-19T20:45:00.971Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/04/2c/9e69d73c4297508be9e3b64a970ea3971b3eb8db64ffc5802d40bd25981f/jsonschema_path-0.5.0-py3-none-any.whl", hash = "sha256:2790a070bc7abb08ea3dbe4d340ece4efadf639223001f020c7503229ba068e2", size = 24077, upload-time = "2026-05-19T20:44:59.225Z" }, -] - [[package]] name = "jsonschema-specifications" version = "2025.9.1" @@ -2282,22 +2040,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/87/ea/fe955e157094d93fe20919441dc709ba7afbcf933de5896b546d5c217938/kernels_data-0.14.1-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:ec694ae26ed8569f613e5cd4269222210cb57b1c16cf8b0c5acbf0082324804f", size = 1414055, upload-time = "2026-05-14T06:41:14.424Z" }, ] -[[package]] -name = "keyring" -version = "25.7.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "jaraco-classes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "jaraco-context", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "jaraco-functools", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "jeepney", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "secretstorage", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/43/4b/674af6ef2f97d56f0ab5153bf0bfa28ccb6c3ed4d1babf4305449668807b/keyring-25.7.0.tar.gz", hash = "sha256:fe01bd85eb3f8fb3dd0405defdeac9a5b4f6f0439edbb3149577f244a2e8245b", size = 63516, upload-time = "2025-11-16T16:26:09.482Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/81/db/e655086b7f3a705df045bf0933bdd9c2f79bb3c97bfef1384598bb79a217/keyring-25.7.0-py3-none-any.whl", hash = "sha256:be4a0b195f149690c166e850609a477c532ddbfbaed96a404d4e43f8d5e2689f", size = 39160, upload-time = "2025-11-16T16:26:08.402Z" }, -] - [[package]] name = "kiwisolver" version = "1.5.0" @@ -2313,7 +2055,7 @@ wheels = [ [[package]] name = "kubernetes" -version = "36.0.0" +version = "36.0.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohttp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -2327,9 +2069,9 @@ dependencies = [ { name = "urllib3", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "websocket-client", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/bf/59/dc635e4e9afb3884bc5c57f14fe23783e4c04601aa20b835ac75c41d1625/kubernetes-36.0.0.tar.gz", hash = "sha256:027b606bb8032e6c6464a53236bdd9bd9a94c237e1063bc45a303c25b304ced9", size = 2346728, upload-time = "2026-05-20T20:44:24.28Z" } +sdist = { url = "https://files.pythonhosted.org/packages/2f/57/8b538af5076bc3372949d76f70ba3449bdfe52f9e6488170fa5d4f7cbe70/kubernetes-36.0.2.tar.gz", hash = "sha256:03551fcb49cae1f708f63624041e37403545b7aaed10cbf54e2b01a37a5438e3", size = 2336738, upload-time = "2026-06-01T18:20:30.785Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/cd/d2/6f99ca9c7eb961dfdd45b9643101399a8ee20922c662c362c91e9cc7e832/kubernetes-36.0.0-py2.py3-none-any.whl", hash = "sha256:a766433357ec9f90db7565cccf52e28e7fca40b0ef366c80a6022adbc0ac0425", size = 4660469, upload-time = "2026-05-20T20:44:20.893Z" }, + { url = "https://files.pythonhosted.org/packages/46/2c/5c160dbdef7123f8cc97fd8ece7e0198627a426a2a49614845e9086feb8d/kubernetes-36.0.2-py2.py3-none-any.whl", hash = "sha256:faf9b5241b58de0c4a5069f2a0ffc8ac06fece7215156cd3d3ba081a78a858b6", size = 4617568, upload-time = "2026-06-01T18:20:28.737Z" }, ] [[package]] @@ -2343,7 +2085,7 @@ sdist = { url = "https://files.pythonhosted.org/packages/0e/72/a3add0e4eec4eb9e2 [[package]] name = "langfuse" -version = "4.6.1" +version = "4.7.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "backoff", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -2355,9 +2097,9 @@ dependencies = [ { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "wrapt", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a6/31/4b7157be23e7c8c3581ac5f6547c5c003e232e7044c92398c468ef78a809/langfuse-4.6.1.tar.gz", hash = "sha256:7f256c669e610909c2e93ca3e9e4168dbef344b753b6874f14b0edd673863f17", size = 281379, upload-time = "2026-05-08T14:08:15.909Z" } +sdist = { url = "https://files.pythonhosted.org/packages/2c/74/a6f1a99893ee6d1a69439ae7eb92f8fe8806103492dc26531d5942dbd3bf/langfuse-4.7.1.tar.gz", hash = "sha256:f9e262eceedb353b191c1da1f8452d1e8ebf52297ca20e160cda0206608e3a40", size = 320620, upload-time = "2026-05-29T18:06:22.435Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/4b/bf/3a6082f7809bdcc1269e9920c07d7c7f92a53cc265a4a879e59c92b23b36/langfuse-4.6.1-py3-none-any.whl", hash = "sha256:a696ac3089a0c8431bf7f1b47b7f6417da311f418dd04ce9ef62d63608fd8797", size = 481237, upload-time = "2026-05-08T14:08:17.141Z" }, + { url = "https://files.pythonhosted.org/packages/9f/9a/bd3368f46b6c72ee2068b80536826b02ae86df53eff1c79941344503098f/langfuse-4.7.1-py3-none-any.whl", hash = "sha256:a4e59c81ad5e5b16a65d3849f4923ebc3ad6e67ec803ada83d50c0cb66149490", size = 562571, upload-time = "2026-05-29T18:06:20.517Z" }, ] [[package]] @@ -2409,7 +2151,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.86.0" +version = "1.87.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohttp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -2425,9 +2167,9 @@ dependencies = [ { name = "tiktoken", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "tokenizers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/2f/a7/26b8b04e4fcff26b60200ffe7458a255552ae51014468188f5db45674eb2/litellm-1.86.0.tar.gz", hash = "sha256:eccab86e0820b60b3f9484b233fb8d818b97afb19d5b4fa08d0d045621350ba4", size = 15379195, upload-time = "2026-05-24T02:42:10.865Z" } +sdist = { url = "https://files.pythonhosted.org/packages/77/0d/ccdf682ccfd7f18bf0e179c39d85616b8f8ef05a798588285310412db13d/litellm-1.87.0.tar.gz", hash = "sha256:cafc1882cb0cbab8374c41180af86e4a067796e4524e15f59e99f6e689cd1bd8", size = 15453755, upload-time = "2026-06-02T03:53:29.076Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ff/0b/9a044c061a69e801de042e962c34f5bc2e094810e28b49ce0b3bedee9327/litellm-1.86.0-py3-none-any.whl", hash = "sha256:9d8171ca1a17705b7c7a6fdce8cfc07bbf641284b46c1b6047f83a779159990c", size = 17011225, upload-time = "2026-05-24T02:42:00.629Z" }, + { url = "https://files.pythonhosted.org/packages/98/20/88a372fa7e50fc2c33458c6eef94a79afcf7bdfa43610079531b82b484a3/litellm-1.87.0-py3-none-any.whl", hash = "sha256:fbbba7e47ae29b55f878fe1acc80effb92761bc168f6236bd81a0cb6e147d855", size = 17103948, upload-time = "2026-06-02T03:53:25.677Z" }, ] [[package]] @@ -2436,7 +2178,7 @@ version = "0.2.7" source = { editable = "deps/research-environments/environments/livecodebench" } dependencies = [ { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -2488,7 +2230,7 @@ dependencies = [ { name = "markdown", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "sympy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -2580,7 +2322,7 @@ source = { editable = "deps/research-environments/environments/math_env" } dependencies = [ { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -2590,6 +2332,17 @@ requires-dist = [ { name = "verifiers", specifier = ">=0.1.15.dev17" }, ] +[[package]] +name = "math-env-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/math_env_v1" } +dependencies = [ + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [{ name = "datasets" }] + [[package]] name = "math-python" version = "0.1.10" @@ -2597,7 +2350,7 @@ source = { editable = "deps/verifiers/environments/math_python" } dependencies = [ { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -2625,7 +2378,7 @@ version = "0.1.17" source = { editable = "deps/research-environments/environments/math500" } dependencies = [ { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -2721,7 +2474,7 @@ dependencies = [ { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "swebench", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "tenacity", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -2787,7 +2540,7 @@ version = "0.1.3" source = { editable = "deps/research-environments/environments/mmlu_pro" } dependencies = [ { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -2868,15 +2621,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3a/a4/a187adcd485ff27bdbdb5c2b4d9cf210427bc74bcaacfc8226409db17535/mooncake_transfer_engine-0.3.11.post1-cp312-cp312-manylinux_2_39_aarch64.whl", hash = "sha256:1ccad9f44cf1a67f4e0494bd02f505503139ab606ecbe76cd6050d7a069247d5", size = 18089789, upload-time = "2026-05-24T16:19:01.828Z" }, ] -[[package]] -name = "more-itertools" -version = "11.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/de/1d/f4da6f02cdffe04d6362210b807146a26044c88d839208aec273bb0d9184/more_itertools-11.1.0.tar.gz", hash = "sha256:48e8f4d9e7e5878571ecf6f2b4e57634f93cd474cc8cfbd2376f2d11b396e30d", size = 145772, upload-time = "2026-05-22T14:14:29.909Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e8/3d/1087453384dbde46a8c7f9356eead2c58be8a7bf156bca40243377c85715/more_itertools-11.1.0-py3-none-any.whl", hash = "sha256:4b65538ae22f6fed0ce4874efd317463a7489796a0939fa66824dd542125a192", size = 72226, upload-time = "2026-05-22T14:14:28.824Z" }, -] - [[package]] name = "mpmath" version = "1.3.0" @@ -2970,11 +2714,11 @@ wheels = [ [[package]] name = "narwhals" -version = "2.21.2" +version = "2.22.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/cf/a0/6198c56d42ef2f3c6ed0c42ba30dbcefdc86a91262d7d449010770ae085b/narwhals-2.21.2.tar.gz", hash = "sha256:5c5b2d0b47aef7c73ea412cfcbcd467f2f2d5be73e3c2ab19d78f4a97718790a", size = 632176, upload-time = "2026-05-16T08:49:08.314Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9c/1c/c80cb7719721a44846c6301ef118434bae30a423924bfad3a47f16bdc064/narwhals-2.22.0.tar.gz", hash = "sha256:6486282bb7e4b4ab55963efbd8be1451b764cc4874b74d1fd625eba9dc60b86f", size = 417565, upload-time = "2026-06-01T13:34:36.249Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1d/77/928ea2e70641ca177a11140062cc5840d421795f2e82749d408d0cce900a/narwhals-2.21.2-py3-none-any.whl", hash = "sha256:7bb57c3700486039215455b9bf2d64261915cc0fd845cc30272d631df696b251", size = 451201, upload-time = "2026-05-16T08:49:05.536Z" }, + { url = "https://files.pythonhosted.org/packages/1f/b6/e7cdde7b8e90d5dff25b622f95833ef26567ad184c977278b93a1cbd5717/narwhals-2.22.0-py3-none-any.whl", hash = "sha256:1421797ede01789cc1537619dbc3f36f840737240f748fdb24a60a0225fc80be", size = 453815, upload-time = "2026-06-01T13:34:34.127Z" }, ] [[package]] @@ -3394,20 +3138,8 @@ sdist = { url = "https://files.pythonhosted.org/packages/3e/92/2d038d096f29179c7 wheels = [ { url = "https://files.pythonhosted.org/packages/d3/d2/ce6953ca87db9cae3e775024184da7d1c5cb88cead19a2d75b42f00a959c/openai_harmony-0.0.8-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e4f709815924ec325b9a890e6ab2bbb0ceec8e319a4e257328eb752cf36b2efc", size = 2948463, upload-time = "2025-11-05T19:06:48.17Z" }, { url = "https://files.pythonhosted.org/packages/25/3f/1a192b93bb47c6b44cd98ba8cc1d3d2a9308f1bb700c3017e6352da11bda/openai_harmony-0.0.8-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c007d277218a50db8839e599ed78e0fffe5130f614c3f6d93ae257f282071a29", size = 2953260, upload-time = "2025-11-05T19:06:55.406Z" }, - { url = "https://files.pythonhosted.org/packages/5b/f8/93b582cad3531797c3db7c2db5400fd841538ccddfd9f5e3df61be99a630/openai_harmony-0.0.8-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:8565d4f5a0638da1bffde29832ed63c9e695c558611053add3b2dc0b56c92dbc", size = 3127044, upload-time = "2025-11-05T19:06:59.553Z" }, - { url = "https://files.pythonhosted.org/packages/60/c3/3d1e01e2dba517a91760e4a03e4f20ffc75039a6fe584d0e6f9b5c78fd15/openai_harmony-0.0.8-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:007b0476a1f331f8130783f901f1da6f5a7057af1a4891f1b6a31dec364189b5", size = 3205080, upload-time = "2025-11-05T19:07:05.078Z" }, -] - -[[package]] -name = "openapi-pydantic" -version = "0.5.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/02/2e/58d83848dd1a79cb92ed8e63f6ba901ca282c5f09d04af9423ec26c56fd7/openapi_pydantic-0.5.1.tar.gz", hash = "sha256:ff6835af6bde7a459fb93eb93bb92b8749b754fc6e51b2f1590a19dc3005ee0d", size = 60892, upload-time = "2025-01-08T19:29:27.083Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/12/cf/03675d8bd8ecbf4445504d8071adab19f5f993676795708e36402ab38263/openapi_pydantic-0.5.1-py3-none-any.whl", hash = "sha256:a3a09ef4586f5bd760a8df7f43028b60cafb6d9f61de2acba9574766255ab146", size = 96381, upload-time = "2025-01-08T19:29:25.275Z" }, + { url = "https://files.pythonhosted.org/packages/5b/f8/93b582cad3531797c3db7c2db5400fd841538ccddfd9f5e3df61be99a630/openai_harmony-0.0.8-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:8565d4f5a0638da1bffde29832ed63c9e695c558611053add3b2dc0b56c92dbc", size = 3127044, upload-time = "2025-11-05T19:06:59.553Z" }, + { url = "https://files.pythonhosted.org/packages/60/c3/3d1e01e2dba517a91760e4a03e4f20ffc75039a6fe584d0e6f9b5c78fd15/openai_harmony-0.0.8-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:007b0476a1f331f8130783f901f1da6f5a7057af1a4891f1b6a31dec364189b5", size = 3205080, upload-time = "2025-11-05T19:07:05.078Z" }, ] [[package]] @@ -3416,7 +3148,7 @@ version = "0.3.12" source = { editable = "deps/research-environments/environments/opencode_cp" } dependencies = [ { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -3432,7 +3164,7 @@ source = { editable = "deps/research-environments/environments/opencode_deepdive dependencies = [ { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -3448,7 +3180,7 @@ version = "0.4.13" source = { editable = "deps/research-environments/environments/opencode_math" } dependencies = [ { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -3463,7 +3195,7 @@ version = "0.3.13" source = { editable = "deps/research-environments/environments/opencode_science" } dependencies = [ { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -3479,7 +3211,7 @@ source = { editable = "deps/research-environments/environments/opencode_swe" } dependencies = [ { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "swebench", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -3503,56 +3235,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4b/33/b5db29a6c00eb8f50708110d8d453747ca125c8b805bc437b289dbdcc057/opencv_python_headless-4.13.0.92-cp37-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:0bd48544f77c68b2941392fcdf9bcd2b9cdf00e98cb8c29b2455d194763cf99e", size = 60391106, upload-time = "2026-02-05T10:30:14.236Z" }, ] -[[package]] -name = "openenv-core" -version = "0.3.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "fastapi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "fastmcp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "gradio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "huggingface-hub", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "openai", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "tomli", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "tomli-w", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "typer", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "uvicorn", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "websockets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/ce/d6/3bebe8afb55fcc3ea9251c4c2dfbab2879e31089bc91a8fe9696e5ce019b/openenv_core-0.3.0.tar.gz", hash = "sha256:c7fee2035badab5be497eb6f4afb2cb417de000f82cc19afd72fb5ec332c431d", size = 164720, upload-time = "2026-05-11T11:37:57.274Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f8/f5/aafa43138589bfd5d369a8d02ea365aae9d6fe55ac0b3894368d6d69bd03/openenv_core-0.3.0-py3-none-any.whl", hash = "sha256:859e875c9d5211b157c30fb9abc681606fcf0bf1b6ffcdf404678992823a1df0", size = 194313, upload-time = "2026-05-11T11:37:55.537Z" }, -] - -[[package]] -name = "openreward" -version = "0.1.125" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "aiohttp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "anthropic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "click", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "fastapi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "google-genai", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "openai", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pyyaml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "sse-starlette", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "structlog", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "tenacity", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "uvicorn", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/9d/dd/b69264c77fd9720f69443c5f6420a7ae9934b4e2799e276d8655f6dc0721/openreward-0.1.125.tar.gz", hash = "sha256:519687307f960ab3a395bf844d6c2fc018d8a0faad0fc367ad3b24331366d390", size = 138936, upload-time = "2026-05-21T10:24:10.075Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/72/11/51c5473d5b3d209ecaacee6ac5418eef15114d0817a98f9cee0a2d6364ee/openreward-0.1.125-py3-none-any.whl", hash = "sha256:784faeeef6aba2ce8f175bd9af5dff29b3bb0c07a4c3642eb56dc431c9af2924", size = 135521, upload-time = "2026-05-21T10:24:11.369Z" }, -] - [[package]] name = "opentelemetry-api" version = "1.42.1" @@ -3761,15 +3443,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/42/32/658973117bf0fd82a24abbfb94fe73a5e86216e49342985e10acce54775a/partial_json_parser-0.2.1.1.post7-py3-none-any.whl", hash = "sha256:145119e5eabcf80cbb13844a6b50a85c68bf99d376f8ed771e2a3c3b03e653ae", size = 10877, upload-time = "2025-11-17T07:27:40.457Z" }, ] -[[package]] -name = "pathable" -version = "0.6.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/66/f3/5a20387de9bcd0607871bfc2198ee0e15836da7baa4592ccd7f24c27c986/pathable-0.6.0.tar.gz", hash = "sha256:6404b8b82aef5ff0fd478934137128b99b12212ba35afdde5525ca4f8388ea58", size = 18970, upload-time = "2026-05-19T18:15:11.911Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a2/e8/6d75ffd9784bce2e93d1ae4415649427e39a53bb172d4672b2b59c6f0a7b/pathable-0.6.0-py3-none-any.whl", hash = "sha256:82c4ca6c98c502ad12e0d4e9779b6210afee93c38990988c8c5d1b49bdcdf566", size = 18983, upload-time = "2026-05-19T18:15:10.728Z" }, -] - [[package]] name = "pdfminer-six" version = "20260107" @@ -3876,7 +3549,7 @@ dependencies = [ { name = "toml", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "tomli", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "typer", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/99/62/3361ae244cfb65009a1a089c4101c02e953fe3b126b2337c44ac50aac51e/prime-0.6.10.tar.gz", hash = "sha256:eee8341905e5daed4f751a2bf0f03e5d52f420450e5f45c4d5aa1856b806faea", size = 629514, upload-time = "2026-05-27T20:53:45.841Z" } wheels = [ @@ -3942,10 +3615,12 @@ dependencies = [ { name = "liger-kernel", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "loguru", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "mooncake-transfer-engine", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "msgspec", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nvidia-ml-py", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "openai", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "orjson", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pandas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "prime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "prime-pydantic-config", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "prime-rl-configs", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -3965,7 +3640,7 @@ dependencies = [ { name = "torchvision", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "transformers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "uvloop", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "vllm", version = "0.22.0+cu129", source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0+cu129-cp38-abi3-manylinux_2_28_aarch64.whl" }, marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, { name = "vllm", version = "0.22.0+cu129", source = { url = "https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0+cu129-cp38-abi3-manylinux_2_28_x86_64.whl" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "wandb", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -3993,17 +3668,26 @@ disagg = [ envs = [ { name = "aime2024", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "aime2025", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "aime24-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "alphabet-sort", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "alphabet-sort-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "code-env", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "code-golf-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "color-codeword", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "color-codeword-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "deepdive", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "deepwiki-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "general-agent", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "general-agent-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "glossary-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "gpqa", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "gsm8k-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "hle", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "ifeval", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "livecodebench", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "logic-env", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "math-env", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "math-env-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "math-python", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "math500", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "mini-swe-agent-plus", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -4013,15 +3697,24 @@ envs = [ { name = "opencode-math", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "opencode-science", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "opencode-swe", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "r2e-gym-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "reverse-text", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "reverse-text-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "rlm-search", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "rlm-swe", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "rlm-uuid-ctf", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "scaleswe-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "science-env", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "simpleqa-verified", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "swebench-verified-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "swelego-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "tau2-bench", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "terminal-bench-2-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "wiki-search", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "wiki-search-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "wikispeedia-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "wordle", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "wordle-v1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] flash-attn = [ { name = "flash-attn", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, @@ -4059,22 +3752,30 @@ mamba-ssm = [ requires-dist = [ { name = "aime2024", marker = "extra == 'envs'", editable = "deps/research-environments/environments/aime2024" }, { name = "aime2025", marker = "extra == 'envs'", editable = "deps/research-environments/environments/aime2025" }, + { name = "aime24-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/aime24_v1" }, { name = "aiolimiter", specifier = ">=1.2.1" }, { name = "alphabet-sort", marker = "extra == 'envs'", editable = "deps/verifiers/environments/alphabet_sort" }, + { name = "alphabet-sort-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/alphabet_sort_v1" }, { name = "beartype", specifier = ">=0.21.0" }, { name = "code-env", marker = "extra == 'envs'", editable = "deps/research-environments/environments/code_env" }, + { name = "code-golf-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/code_golf_v1" }, { name = "color-codeword", marker = "extra == 'envs'", editable = "deps/research-environments/environments/color_codeword" }, + { name = "color-codeword-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/color_codeword_v1" }, { name = "datasets", specifier = ">=4.0.0" }, { name = "deep-ep", marker = "platform_machine == 'x86_64' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl" }, { name = "deep-gemm", marker = "platform_machine == 'x86_64' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_x86_64.whl" }, { name = "deepdive", marker = "extra == 'envs'", editable = "deps/research-environments/environments/deepdive" }, + { name = "deepwiki-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/deepwiki_v1" }, { name = "dion", git = "https://github.com/samsja/dion.git?rev=d891eeb" }, { name = "flash-attn", marker = "platform_machine == 'x86_64' and extra == 'flash-attn'", url = "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.4/flash_attn-2.8.3+cu128torch2.11-cp312-cp312-linux_x86_64.whl" }, { name = "flash-attn-3", marker = "extra == 'flash-attn-3'", index = "https://download.pytorch.org/whl/test/cu128" }, { name = "flash-attn-4", marker = "extra == 'flash-attn-cute'", git = "https://github.com/Dao-AILab/flash-attention.git?subdirectory=flash_attn%2Fcute&rev=96bd151" }, { name = "flash-linear-attention", git = "https://github.com/fla-org/flash-linear-attention" }, { name = "general-agent", marker = "extra == 'envs'", editable = "deps/research-environments/environments/general_agent" }, + { name = "general-agent-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/general_agent_v1" }, + { name = "glossary-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/glossary_v1" }, { name = "gpqa", marker = "extra == 'envs'", editable = "deps/research-environments/environments/gpqa" }, + { name = "gsm8k-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/gsm8k_v1" }, { name = "hle", marker = "extra == 'envs'", editable = "deps/research-environments/environments/hle" }, { name = "ifeval", marker = "extra == 'envs'", editable = "deps/research-environments/environments/ifeval" }, { name = "jaxtyping", specifier = ">=0.3.2" }, @@ -4084,12 +3785,14 @@ requires-dist = [ { name = "logic-env", marker = "extra == 'envs'", editable = "deps/research-environments/environments/logic_env" }, { name = "loguru", specifier = ">=0.7.3" }, { name = "math-env", marker = "extra == 'envs'", editable = "deps/research-environments/environments/math_env" }, + { name = "math-env-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/math_env_v1" }, { name = "math-python", marker = "extra == 'envs'", editable = "deps/verifiers/environments/math_python" }, { name = "math500", marker = "extra == 'envs'", editable = "deps/research-environments/environments/math500" }, { name = "mini-swe-agent-plus", marker = "extra == 'envs'", editable = "deps/research-environments/environments/mini_swe_agent_plus" }, { name = "mmlu-pro", marker = "extra == 'envs'", editable = "deps/research-environments/environments/mmlu_pro" }, { name = "modelexpress", marker = "extra == 'modelexpress'", specifier = "==0.3.0" }, { name = "mooncake-transfer-engine", specifier = ">=0.3.10.post2" }, + { name = "msgspec", specifier = ">=0.18" }, { name = "nixl", marker = "extra == 'disagg'" }, { name = "nixl-cu12", marker = "platform_machine == 'x86_64' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/nixl_cu12-0.10.1-cp312-cp312-linux_x86_64.whl" }, { name = "numpy", specifier = ">=2.2.6" }, @@ -4100,7 +3803,9 @@ requires-dist = [ { name = "opencode-math", marker = "extra == 'envs'", editable = "deps/research-environments/environments/opencode_math" }, { name = "opencode-science", marker = "extra == 'envs'", editable = "deps/research-environments/environments/opencode_science" }, { name = "opencode-swe", marker = "extra == 'envs'", editable = "deps/research-environments/environments/opencode_swe" }, + { name = "orjson", specifier = ">=3.10" }, { name = "orjson", specifier = ">=3.11.0" }, + { name = "pandas", specifier = ">=2.0" }, { name = "prime", specifier = ">=0.6.4" }, { name = "prime-pydantic-config", editable = "deps/pydantic-config" }, { name = "prime-rl", extras = ["disagg"], marker = "extra == 'all'" }, @@ -4113,18 +3818,24 @@ requires-dist = [ { name = "pybase64", specifier = ">=1.4.2" }, { name = "pyzmq", specifier = ">=27.1.0" }, { name = "quack-kernels", marker = "extra == 'quack'", specifier = ">=0.4.1" }, + { name = "r2e-gym-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/r2e_gym_v1" }, { name = "renderers", editable = "deps/renderers" }, { name = "reverse-text", marker = "extra == 'envs'", editable = "deps/verifiers/environments/reverse_text" }, + { name = "reverse-text-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/reverse_text_v1" }, { name = "rich", specifier = ">=14.0.0" }, { name = "ring-flash-attn", specifier = ">=0.1.8" }, { name = "rlm-search", marker = "extra == 'envs'", editable = "deps/research-environments/environments/rlm_search" }, { name = "rlm-swe", marker = "extra == 'envs'", editable = "deps/research-environments/environments/rlm_swe" }, { name = "rlm-uuid-ctf", marker = "extra == 'envs'", editable = "deps/research-environments/environments/rlm_uuid_ctf" }, + { name = "scaleswe-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/scaleswe_v1" }, { name = "science-env", marker = "extra == 'envs'", editable = "deps/research-environments/environments/science_env" }, { name = "setproctitle", specifier = ">=1.3.0" }, { name = "simpleqa-verified", marker = "extra == 'envs'", editable = "deps/research-environments/environments/simpleqa_verified" }, + { name = "swebench-verified-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/swebench_verified_v1" }, + { name = "swelego-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/swelego_v1" }, { name = "tau2-bench", marker = "extra == 'envs'", editable = "deps/research-environments/environments/tau2_bench" }, { name = "tenacity", specifier = ">=8.2.0" }, + { name = "terminal-bench-2-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/terminal_bench_2_v1" }, { name = "tilelang", specifier = ">=0.1.8" }, { name = "torch", specifier = ">=2.9.0", index = "https://download.pytorch.org/whl/cu128" }, { name = "torchaudio", index = "https://download.pytorch.org/whl/cu128" }, @@ -4140,7 +3851,10 @@ requires-dist = [ { name = "vllm-router", marker = "platform_machine == 'x86_64' and extra == 'disagg'", url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.1.26/vllm_router-0.1.26-cp38-abi3-manylinux_2_28_x86_64.whl" }, { name = "wandb", specifier = ">=0.26.1" }, { name = "wiki-search", marker = "extra == 'envs'", editable = "deps/verifiers/environments/wiki_search" }, + { name = "wiki-search-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/wiki_search_v1" }, + { name = "wikispeedia-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/wikispeedia_v1" }, { name = "wordle", marker = "extra == 'envs'", editable = "deps/verifiers/environments/wordle" }, + { name = "wordle-v1", marker = "extra == 'envs'", editable = "deps/verifiers/environments/wordle_v1" }, ] provides-extras = ["flash-attn", "flash-attn-3", "flash-attn-cute", "envs", "disagg", "gpt-oss", "modelexpress", "quack", "all"] @@ -4165,6 +3879,7 @@ dependencies = [ { name = "renderers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "tomli", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "tomli-w", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -4174,11 +3889,12 @@ requires-dist = [ { name = "renderers", specifier = ">=0.1.8.dev28" }, { name = "tomli", specifier = ">=2.2.1" }, { name = "tomli-w", specifier = ">=1.2.0" }, + { name = "verifiers", specifier = ">=0.1.15.dev357" }, ] [[package]] name = "prime-sandboxes" -version = "0.2.26" +version = "0.2.27" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiofiles", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -4188,23 +3904,23 @@ dependencies = [ { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "tenacity", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/9b/98/7ecfdd389664c5112eda8d83e287b266f627fdbcefe5510b1b7346effa4b/prime_sandboxes-0.2.26.tar.gz", hash = "sha256:87b0aa2d5ee44f201cb4d0fc1bdf52efa617a7f3cbf9e5a5ee7bc2de0f45ac98", size = 68496, upload-time = "2026-05-12T22:17:47.801Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9a/1f/b257b21f54e9b961bcac05335f6637c494ed0986fb4b5d2054f97f96faa0/prime_sandboxes-0.2.27.tar.gz", hash = "sha256:db4071387f4b2dc8bcd0c8916af4031f864dfd5dfd9818f56119347af20b1469", size = 68661, upload-time = "2026-06-05T21:55:35.175Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e8/51/b5c282dffd4eba73953cdf86c34df36dd31264464489fd14a1a565776e03/prime_sandboxes-0.2.26-py3-none-any.whl", hash = "sha256:03226e4fa81a8008bbdb3046f295a56a43e74bb89b8e894b001547643deca918", size = 34234, upload-time = "2026-05-12T22:17:46.494Z" }, + { url = "https://files.pythonhosted.org/packages/94/bc/a9142d0ef67d92672469cff15cadc04eb46ca9957b00156445d14c14b04f/prime_sandboxes-0.2.27-py3-none-any.whl", hash = "sha256:3fb227cc909c15475fb2874974d166d857021ee6ed9a6e0d4482b1b5ebfc50e2", size = 34402, upload-time = "2026-06-05T21:55:33.451Z" }, ] [[package]] name = "prime-tunnel" -version = "0.1.9" +version = "0.1.8" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "tenacity", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/36/de/0d17589b232708d1f5985372d1ca49e73493e0b08f860e6540d4885ce229/prime_tunnel-0.1.9.tar.gz", hash = "sha256:8d6c7fd7f5a98c47f50707edd606c4e7110e9dd1a08a1008b073be466408ac8d", size = 15162, upload-time = "2026-06-11T18:10:03.513Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9c/e7/557c7623bd8c9a7a9127d880dce2a3d8612d3a7df02bdde6f01c958a20b8/prime_tunnel-0.1.8.tar.gz", hash = "sha256:07803d5d5c6ec83c260bbef0ecd340f4f6acab0f6a254d5d34a4c6ddfc777c0b", size = 14576, upload-time = "2026-06-02T06:39:36.56Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/50/2a/4b370d5d7cf7cbd98ce2714b7d877d2a184cc1b00d599df498aca04bce2d/prime_tunnel-0.1.9-py3-none-any.whl", hash = "sha256:9ad888e26ba4d81e96c643d9d27687ed4737c873f3b80ed3da8f92d6d7a2a615", size = 16274, upload-time = "2026-06-11T18:10:01.463Z" }, + { url = "https://files.pythonhosted.org/packages/f2/5f/181fe0ca1ca8507c0a373e55363b411bdeab9ad904c68edec7d7b1bdea74/prime_tunnel-0.1.8-py3-none-any.whl", hash = "sha256:fab081f566e0887ce50c7b5872169e9f356fc3263d5f0493cfabb3150eb73b0a", size = 15843, upload-time = "2026-06-02T06:39:34.043Z" }, ] [[package]] @@ -4304,31 +4020,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e0/a9/023730ba63db1e494a271cb018dcd361bd2c917ba7004c3e49d5daf795a2/py_cpuinfo-9.0.0-py3-none-any.whl", hash = "sha256:859625bc251f64e21f077d099d4162689c762b5d6a4c3c97553d56241c9674d5", size = 22335, upload-time = "2022-10-25T20:38:27.636Z" }, ] -[[package]] -name = "py-key-value-aio" -version = "0.4.5" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "beartype", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/fb/e2/d689d922894a7ecde73b6daeaf9b13dab5aae06fe6aaaf7514722644d382/py_key_value_aio-0.4.5.tar.gz", hash = "sha256:c6563a2c6abe5da5e20f4f9e875c2a9b425a2244a54fadbf46cf140a9eea45d7", size = 107547, upload-time = "2026-05-27T16:37:08.107Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f6/95/b8ba862968712caa12a19666175334fa979e1f198b896a430adb3bacfe87/py_key_value_aio-0.4.5-py3-none-any.whl", hash = "sha256:ab862adbcb8c72547d1c57821f22cbbb71ab86509039c96f36e914e0336c8dd7", size = 170005, upload-time = "2026-05-27T16:37:06.629Z" }, -] - -[package.optional-dependencies] -filetree = [ - { name = "aiofile", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "anyio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -keyring = [ - { name = "keyring", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -memory = [ - { name = "cachetools", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] - [[package]] name = "pyarrow" version = "24.0.0" @@ -4341,27 +4032,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/36/7a/82c363caa145fff88fb475da50d3bf52bb024f61917be5424c3392eaf878/pyarrow-24.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:25ea65d868eb04015cd18e6df2fbe98f07e5bda2abefabcb88fce39a947716f6", size = 51929490, upload-time = "2026-04-21T10:47:55.981Z" }, ] -[[package]] -name = "pyasn1" -version = "0.6.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/5c/5f/6583902b6f79b399c9c40674ac384fd9cd77805f9e6205075f828ef11fb2/pyasn1-0.6.3.tar.gz", hash = "sha256:697a8ecd6d98891189184ca1fa05d1bb00e2f84b5977c481452050549c8a72cf", size = 148685, upload-time = "2026-03-17T01:06:53.382Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/5d/a0/7d793dce3fa811fe047d6ae2431c672364b462850c6235ae306c0efd025f/pyasn1-0.6.3-py3-none-any.whl", hash = "sha256:a80184d120f0864a52a073acc6fc642847d0be408e7c7252f31390c0f4eadcde", size = 83997, upload-time = "2026-03-17T01:06:52.036Z" }, -] - -[[package]] -name = "pyasn1-modules" -version = "0.4.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pyasn1", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e9/e6/78ebbb10a8c8e4b61a59249394a4a594c1a7af95593dc933a349c8d00964/pyasn1_modules-0.4.2.tar.gz", hash = "sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6", size = 307892, upload-time = "2025-03-28T02:41:22.17Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/47/8d/d529b5d697919ba8c11ad626e835d4039be708a35b0d22de83a269a6682c/pyasn1_modules-0.4.2-py3-none-any.whl", hash = "sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a", size = 181259, upload-time = "2025-03-28T02:41:19.028Z" }, -] - [[package]] name = "pybase64" version = "1.4.3" @@ -4475,15 +4145,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ae/8d/f1af3832f5e6eb13ba94ee809e72b8ecb5eef226d27ee0bef7d963d943c7/pydantic_settings-2.14.1-py3-none-any.whl", hash = "sha256:6e3c7edfd8277687cdc598f56e5cff0e9bfff0910a3749deaa8d4401c3a2b9de", size = 60964, upload-time = "2026-05-08T13:40:04.958Z" }, ] -[[package]] -name = "pydub" -version = "0.25.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/fe/9a/e6bca0eed82db26562c73b5076539a4a08d3cffd19c3cc5913a3e61145fd/pydub-0.25.1.tar.gz", hash = "sha256:980a33ce9949cab2a569606b65674d748ecbca4f0796887fd6f46173a7b0d30f", size = 38326, upload-time = "2021-03-10T02:09:54.659Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a6/53/d78dc063216e62fc55f6b2eebb447f6a4b0a59f55c8406376f76bf959b08/pydub-0.25.1-py2.py3-none-any.whl", hash = "sha256:65617e33033874b59d87db603aa1ed450633288aefead953b30bded59cb599a6", size = 32327, upload-time = "2021-03-10T02:09:53.503Z" }, -] - [[package]] name = "pyelftools" version = "0.32" @@ -4532,17 +4193,6 @@ crypto = [ { name = "cryptography", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -[[package]] -name = "pymupdf" -version = "1.27.2.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/22/32/708bedc9dde7b328d45abbc076091769d44f2f24ad151ad92d56a6ec142b/pymupdf-1.27.2.3.tar.gz", hash = "sha256:7a92faa25129e8bbec5e50eeb9214f187665428c31b05c4ef6e36c58c0b1c6d2", size = 85759618, upload-time = "2026-04-24T14:13:14.42Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c3/26/b7e5a70eb83bd189f8b5df87ec442746b992f2f632662839b288170d357d/pymupdf-1.27.2.3-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:1dd460a3ae4597a755f00a3bd9771f5ebf1531dc111f6a36bf05dd00a6b84425", size = 24333923, upload-time = "2026-04-24T14:09:47.341Z" }, - { url = "https://files.pythonhosted.org/packages/e4/a0/aa1ee2240f29481a04a827c313333b4ecd8a14d6ac3e15d3f41a30574781/pymupdf-1.27.2.3-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:857842b4888827bd6155a1131341b2822a7ebe9a8c15a975fd7d490d7a64a30c", size = 24963198, upload-time = "2026-04-24T14:10:07.408Z" }, - { url = "https://files.pythonhosted.org/packages/69/49/4f742451f980840829fc00ba158bebb25d389c846d8f4f8c65936ee55de8/pymupdf-1.27.2.3-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:580983849c64a08d08344ca3d1580e87c01f046a8392421797bc850efd72a5b6", size = 25184609, upload-time = "2026-04-24T14:10:22.911Z" }, -] - [[package]] name = "pynacl" version = "1.6.2" @@ -4571,15 +4221,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" }, ] -[[package]] -name = "pyperclip" -version = "1.11.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e8/52/d87eba7cb129b81563019d1679026e7a112ef76855d6159d24754dbd2a51/pyperclip-1.11.0.tar.gz", hash = "sha256:244035963e4428530d9e3a6101a1ef97209c6825edab1567beac148ccc1db1b6", size = 12185, upload-time = "2025-09-26T14:40:37.245Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/df/80/fc9d01d5ed37ba4c42ca2b55b4339ae6e200b456be3a1aaddf4a9fa99b8c/pyperclip-1.11.0-py3-none-any.whl", hash = "sha256:299403e9ff44581cb9ba2ffeed69c7aa96a008622ad0c46cb575ca75b5b84273", size = 11063, upload-time = "2025-09-26T14:40:36.069Z" }, -] - [[package]] name = "pypika" version = "0.51.1" @@ -4697,15 +4338,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8f/cb/769cfc37177252872a45a71f3fbdde9d51b471a3f3c14bfe95dde3407386/python_multipart-0.0.29-py3-none-any.whl", hash = "sha256:2ddcc971cef266225f54f552d8fa10bcfbb1f14446caec199060daac59ff2d69", size = 29640, upload-time = "2026-05-17T17:29:45.69Z" }, ] -[[package]] -name = "pytz" -version = "2026.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ff/46/dd499ec9038423421951e4fad73051febaa13d2df82b4064f87af8b8c0c3/pytz-2026.2.tar.gz", hash = "sha256:0e60b47b29f21574376f218fe21abc009894a2321ea16c6754f3cad6eb7cdd6a", size = 320861, upload-time = "2026-05-04T01:35:29.667Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ec/dd/96da98f892250475bdf2328112d7468abdd4acc7b902b6af23f4ed958ea0/pytz-2026.2-py2.py3-none-any.whl", hash = "sha256:04156e608bee23d3792fd45c94ae47fae1036688e75032eea2e3bf0323d1f126", size = 510141, upload-time = "2026-05-04T01:35:27.408Z" }, -] - [[package]] name = "pyyaml" version = "6.0.3" @@ -4749,13 +4381,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/38/e4/a6c3bbbe3d4242fa412454b8e8069a079e500be331aecf8f2aa666164e9c/quack_kernels-0.4.1-py3-none-any.whl", hash = "sha256:c1c8df2935bf5156ec47d2c5384ac08b411fd0ee702d80ae916dbf6d6f5ae813", size = 260827, upload-time = "2026-04-30T14:37:54.584Z" }, ] +[[package]] +name = "r2e-gym-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/r2e_gym_v1" } +dependencies = [ + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [{ name = "datasets" }] + [[package]] name = "redis" -version = "7.4.0" +version = "8.0.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/7b/7f/3759b1d0d72b7c92f0d70ffd9dc962b7b7b5ee74e135f9d7d8ab06b8a318/redis-7.4.0.tar.gz", hash = "sha256:64a6ea7bf567ad43c964d2c30d82853f8df927c5c9017766c55a1d1ed95d18ad", size = 4943913, upload-time = "2026-03-24T09:14:37.53Z" } +sdist = { url = "https://files.pythonhosted.org/packages/53/ae/ed461cca5780b5fc8b9fe8ca0ed98d89508645fb9d880c24cc42c087678f/redis-8.0.0.tar.gz", hash = "sha256:a00c5355432051ac14e593b8b197fc76c887ee12d55a0984f69328a1115fdc49", size = 5101591, upload-time = "2026-05-28T12:45:13.5Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/74/3a/95deec7db1eb53979973ebd156f3369a72732208d1391cd2e5d127062a32/redis-7.4.0-py3-none-any.whl", hash = "sha256:a9c74a5c893a5ef8455a5adb793a31bb70feb821c86eccb62eebef5a19c429ec", size = 409772, upload-time = "2026-03-24T09:14:35.968Z" }, + { url = "https://files.pythonhosted.org/packages/27/e3/b519734372d305bd547534a9f32e4ce9f98552af753dce72cf3483a0ff0b/redis-8.0.0-py3-none-any.whl", hash = "sha256:c938c18338585009f0bc310f4c7e4e4b4d37639356c4ac072cedf3af570c8dc7", size = 499870, upload-time = "2026-05-28T12:45:11.697Z" }, ] [[package]] @@ -4856,7 +4499,7 @@ version = "0.1.4" source = { editable = "deps/verifiers/environments/reverse_text" } dependencies = [ { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -4865,6 +4508,17 @@ requires-dist = [ { name = "verifiers", specifier = ">=0.1.5.post0" }, ] +[[package]] +name = "reverse-text-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/reverse_text_v1" } +dependencies = [ + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [{ name = "datasets" }] + [[package]] name = "rich" version = "15.0.0" @@ -4878,19 +4532,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" }, ] -[[package]] -name = "rich-rst" -version = "2.0.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pygments", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/57/56/3191bae66b08ccc637ea8120426068bcb361cc323c96404c310886937067/rich_rst-2.0.1.tar.gz", hash = "sha256:cbe236ed0901d1ec8427cc6a50bf0a34353ba28ad014dc24def68bfe7f3b9e68", size = 300570, upload-time = "2026-05-16T00:47:57.362Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a0/3d/55c17d3ebdf3cd81356002afe5bef9bb8af631db2819785b6eac845b925b/rich_rst-2.0.1-py3-none-any.whl", hash = "sha256:7ee15f345ce25fa02b582c272a6cdbaf0c21243e38061cea273cff659bf3ef61", size = 272922, upload-time = "2026-05-16T00:47:55.508Z" }, -] - [[package]] name = "rich-toolkit" version = "0.19.10" @@ -4935,7 +4576,7 @@ dependencies = [ { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "openai", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -4955,7 +4596,7 @@ dependencies = [ { name = "multi-swe-bench", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "swebench", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -4973,7 +4614,7 @@ source = { editable = "deps/research-environments/environments/rlm_uuid_ctf" } dependencies = [ { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -5007,18 +4648,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ed/c5/8864e4e7925b836ea354b31d57641ec03830564e281a8b6f061f8c3e0ec1/ruff-0.15.14-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:bb01d645694e3ec0102105d07ef2d53703970407d59c04e59d3ba0b7a1d53553", size = 11560214, upload-time = "2026-05-21T14:34:50.975Z" }, ] -[[package]] -name = "safehttpx" -version = "0.1.7" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/89/d1/4282284d9cf1ee873607a46442da977fc3c985059315ab23610be31d5885/safehttpx-0.1.7.tar.gz", hash = "sha256:db201c0978c41eddb8bb480f3eee59dd67304fdd91646035e9d9a720049a9d23", size = 10385, upload-time = "2025-10-24T18:30:09.783Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2e/a3/0f0b7d78e2f1eb9e8e1afbff1d2bff8d60144aee17aca51c065b516743dd/safehttpx-0.1.7-py3-none-any.whl", hash = "sha256:c4f4a162db6993464d7ca3d7cc4af0ffc6515a606dfd220b9f82c6945d869cde", size = 8959, upload-time = "2025-10-24T18:30:08.733Z" }, -] - [[package]] name = "safetensors" version = "0.7.0" @@ -5031,13 +4660,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4a/d8/0c8a7dc9b41dcac53c4cbf9df2b9c83e0e0097203de8b37a712b345c0be5/safetensors-0.7.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b0f6d66c1c538d5a94a73aa9ddca8ccc4227e6c9ff555322ea40bdd142391dd4", size = 677368, upload-time = "2025-11-19T15:18:41.627Z" }, ] +[[package]] +name = "scaleswe-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/scaleswe_v1" } +dependencies = [ + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [{ name = "datasets" }] + [[package]] name = "science-env" version = "0.1.4" source = { editable = "deps/research-environments/environments/science_env" } dependencies = [ { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -5048,18 +4688,19 @@ requires-dist = [ [[package]] name = "scikit-learn" -version = "1.8.0" +version = "1.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "joblib", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "narwhals", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "scipy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "threadpoolctl", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/0e/d4/40988bf3b8e34feec1d0e6a051446b1f66225f8529b9309becaeef62b6c4/scikit_learn-1.8.0.tar.gz", hash = "sha256:9bccbb3b40e3de10351f8f5068e105d0f4083b1a65fa07b6634fbc401a6287fd", size = 7335585, upload-time = "2025-12-10T07:08:53.618Z" } +sdist = { url = "https://files.pythonhosted.org/packages/fa/6f/37092bdb25f712817231799fc5674d8e704066a8a70c1d2d40517e18b4ab/scikit_learn-1.9.0.tar.gz", hash = "sha256:8833266989d3a5110178a9fae30783675460724d0e1efb13b14901d2c660c557", size = 7750767, upload-time = "2026-06-02T11:54:32.706Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/dd/47/f187b4636ff80cc63f21cd40b7b2d177134acaa10f6bb73746130ee8c2e5/scikit_learn-1.8.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4496bb2cf7a43ce1a2d7524a79e40bc5da45cf598dbf9545b7e8316ccba47bb4", size = 8660492, upload-time = "2025-12-10T07:07:55.574Z" }, - { url = "https://files.pythonhosted.org/packages/97/74/b7a304feb2b49df9fafa9382d4d09061a96ee9a9449a7cbea7988dda0828/scikit_learn-1.8.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a0bcfe4d0d14aec44921545fd2af2338c7471de9cb701f1da4c9d85906ab847a", size = 8931904, upload-time = "2025-12-10T07:07:57.666Z" }, + { url = "https://files.pythonhosted.org/packages/a0/ee/5adbc77656b71f9456a2f5a7a9fdb4bcf9207a6b962889f1c2f9323afa4e/scikit_learn-1.9.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5e50ed4da51974e86e940690e9a3d82e729b62b5a49f7c9bac534d515d39d86f", size = 8837603, upload-time = "2026-06-02T11:53:30.328Z" }, + { url = "https://files.pythonhosted.org/packages/6c/c2/63fdda36c56437eeb44aaf9493c8bcd62ce230ab1598924fc626ffbfa943/scikit_learn-1.9.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:056c92bb67ad4c28463c2f2653d9701449201e7e7a9e94e321be0f71c4fef2b8", size = 9132097, upload-time = "2026-06-02T11:53:33.456Z" }, ] [[package]] @@ -5091,28 +4732,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/83/11/00d3c3dfc25ad54e731d91449895a79e4bf2384dc3ac01809010ba88f6d5/seaborn-0.13.2-py3-none-any.whl", hash = "sha256:636f8336facf092165e27924f223d3c62ca560b1f2bb5dff7ab7fad265361987", size = 294914, upload-time = "2024-01-25T13:21:49.598Z" }, ] -[[package]] -name = "secretstorage" -version = "3.5.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cryptography", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "jeepney", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/1c/03/e834bcd866f2f8a49a85eaff47340affa3bfa391ee9912a952a1faa68c7b/secretstorage-3.5.0.tar.gz", hash = "sha256:f04b8e4689cbce351744d5537bf6b1329c6fc68f91fa666f60a380edddcd11be", size = 19884, upload-time = "2025-11-23T19:02:53.191Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/46/f5af3402b579fd5e11573ce652019a67074317e18c1935cc0b4ba9b35552/secretstorage-3.5.0-py3-none-any.whl", hash = "sha256:0ce65888c0725fcb2c5bc0fdb8e5438eece02c523557ea40ce0703c266248137", size = 15554, upload-time = "2025-11-23T19:02:51.545Z" }, -] - -[[package]] -name = "semantic-version" -version = "2.10.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/7d/31/f2289ce78b9b473d582568c234e104d2a342fd658cc288a7553d83bb8595/semantic_version-2.10.0.tar.gz", hash = "sha256:bdabb6d336998cbb378d4b9db3a4b56a1e3235701dc05ea2690d9a997ed5041c", size = 52289, upload-time = "2022-05-26T13:35:23.454Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/6a/23/8146aad7d88f4fcb3a6218f41a60f6c2d4e3a72de72da1825dc7c8f7877c/semantic_version-2.10.0-py2.py3-none-any.whl", hash = "sha256:de78a3b8e0feda74cabc54aab2da702113e33ac9d9eb9d2389bcf1f58b7d9177", size = 15552, upload-time = "2022-05-26T13:35:21.206Z" }, -] - [[package]] name = "sentencepiece" version = "0.2.1" @@ -5171,7 +4790,7 @@ name = "simpleqa-verified" version = "0.1.2" source = { editable = "deps/research-environments/environments/simpleqa_verified" } dependencies = [ - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -5215,14 +4834,15 @@ wheels = [ [[package]] name = "sse-starlette" -version = "2.3.6" +version = "3.4.4" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "starlette", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/8c/f4/989bc70cb8091eda43a9034ef969b25145291f3601703b82766e5172dfed/sse_starlette-2.3.6.tar.gz", hash = "sha256:0382336f7d4ec30160cf9ca0518962905e1b69b72d6c1c995131e0a703b436e3", size = 18284, upload-time = "2025-05-30T13:34:12.914Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f7/2b/58abc2d1fd397e7dde08e947e05c884d8ef2f78d5e2588c17a12d42d6994/sse_starlette-3.4.4.tar.gz", hash = "sha256:07e0fa0460138baf25cdd5fb28683472c3995dc1642225191b3832d62526bcb0", size = 31819, upload-time = "2026-05-12T17:37:17.019Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/81/05/78850ac6e79af5b9508f8841b0f26aa9fd329a1ba00bf65453c2d312bcc8/sse_starlette-2.3.6-py3-none-any.whl", hash = "sha256:d49a8285b182f6e2228e2609c350398b2ca2c36216c2675d875f81e93548f760", size = 10606, upload-time = "2025-05-30T13:34:11.703Z" }, + { url = "https://files.pythonhosted.org/packages/dc/67/805710444ea8cc75fbf70b920ed431a560c4bf9c57f7d5a3117213189399/sse_starlette-3.4.4-py3-none-any.whl", hash = "sha256:3f4dd50d8aed2771a091f3a83000323fc3844541c16b4fe585ae2420cc6df973", size = 16514, upload-time = "2026-05-12T17:37:15.601Z" }, ] [[package]] @@ -5252,15 +4872,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/0d/13d1d239a25cbfb19e740db83143e95c772a1fe10202dda4b76792b114dd/starlette-0.52.1-py3-none-any.whl", hash = "sha256:0029d43eb3d273bc4f83a08720b4912ea4b071087a3b48db01b7c839f7954d74", size = 74272, upload-time = "2026-01-18T13:34:09.188Z" }, ] -[[package]] -name = "structlog" -version = "25.5.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ef/52/9ba0f43b686e7f3ddfeaa78ac3af750292662284b3661e91ad5494f21dbc/structlog-25.5.0.tar.gz", hash = "sha256:098522a3bebed9153d4570c6d0288abf80a031dfdb2048d59a49e9dc2190fc98", size = 1460830, upload-time = "2025-10-27T08:28:23.028Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a8/45/a132b9074aa18e799b891b91ad72133c98d8042c70f6240e4c5f9dabee2f/structlog-25.5.0-py3-none-any.whl", hash = "sha256:a8453e9b9e636ec59bd9e79bbd4a72f025981b3ba0f5837aebf48f02f37a7f9f", size = 72510, upload-time = "2025-10-27T08:28:21.535Z" }, -] - [[package]] name = "supervisor" version = "4.3.0" @@ -5314,6 +4925,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/36/67/981d8b642ac3eac7c8a7b7832ff8b2fb74f96b28b5fcd9a8979879e5c46d/swebench-4.1.0-py3-none-any.whl", hash = "sha256:1243776f720047cc9e20a427f7a52b75c13a07abda6154fb60fe77f82ec8af57", size = 157231, upload-time = "2025-09-11T02:57:58.953Z" }, ] +[[package]] +name = "swebench-verified-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/swebench_verified_v1" } +dependencies = [ + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [{ name = "verifiers" }] + +[[package]] +name = "swelego-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/swelego_v1" } +dependencies = [ + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [{ name = "datasets" }] + [[package]] name = "sympy" version = "1.14.0" @@ -5328,14 +4961,14 @@ wheels = [ [[package]] name = "synchronicity" -version = "0.12.2" +version = "0.12.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f9/5e/50ea27817003665c7cc4f5bdad309f13d6329037f657848ee87fe06c3740/synchronicity-0.12.2.tar.gz", hash = "sha256:6fd605a5035d1ec74ce48fffaca80ea00345c84ca34223914e2436fb4f162ff9", size = 60018, upload-time = "2026-04-06T15:06:15.447Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ec/d5/e96e6082790c92480380f28aa53e111844cdac7b0f75846f4772cb535a43/synchronicity-0.12.3.tar.gz", hash = "sha256:0d4228b85eaf2805f23b4615b2039a9d24ea811646e2d9f8d0c033094eb85841", size = 60261, upload-time = "2026-05-28T12:33:50.206Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/41/44/4f6ba4e2c171847e6f9a460213b196bbf26edea43d0e66889c7ccc55d368/synchronicity-0.12.2-py3-none-any.whl", hash = "sha256:9dbaca81fb7f2b57c6dea326e514e1c80e9ccfd9c9618515e84fa6091026273b", size = 41312, upload-time = "2026-04-06T15:06:14.459Z" }, + { url = "https://files.pythonhosted.org/packages/57/ea/531a6ea751cbd989da386144810b1b8f529b0aae8c1a9beda8b40966c9c2/synchronicity-0.12.3-py3-none-any.whl", hash = "sha256:e476818cd14102136f41622c619de548f0000c024485fc18521c8fe908ea7574", size = 40982, upload-time = "2026-05-28T12:33:49.125Z" }, ] [[package]] @@ -5347,36 +4980,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/99/55/db07de81b5c630da5cbf5c7df646580ca26dfaefa593667fc6f2fe016d2e/tabulate-0.10.0-py3-none-any.whl", hash = "sha256:f0b0622e567335c8fabaaa659f1b33bcb6ddfe2e496071b743aa113f8774f2d3", size = 39814, upload-time = "2026-03-04T18:55:31.284Z" }, ] -[[package]] -name = "tasksets" -source = { editable = "deps/verifiers/packages/tasksets" } -dependencies = [ - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] - -[package.optional-dependencies] -openenv = [ - { name = "openenv-core", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -openreward = [ - { name = "openreward", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -ta = [ - { name = "nltk", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "textarena", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] - -[package.metadata] -requires-dist = [ - { name = "nemo-gym", marker = "python_full_version >= '3.12' and extra == 'nemogym'", specifier = ">=0.2.1,<0.3" }, - { name = "nltk", marker = "extra == 'ta'" }, - { name = "openenv-core", marker = "extra == 'openenv'", specifier = ">=0.3.0" }, - { name = "openreward", marker = "python_full_version >= '3.11' and extra == 'openreward'", specifier = ">=0.1.123" }, - { name = "textarena", marker = "extra == 'ta'" }, - { name = "verifiers", specifier = ">=0.1.15.dev11" }, -] -provides-extras = ["nemogym", "openenv", "openreward", "ta"] - [[package]] name = "tau2" version = "0.2.1.dev0" @@ -5416,7 +5019,7 @@ version = "0.2.3" source = { editable = "deps/research-environments/environments/tau2_bench" } dependencies = [ { name = "tau2", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -5463,6 +5066,17 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/73/c6/825dab04195756cf8ff2e12698f22513b3db2f64925bdd41671bfb33aaa5/tensorboard_data_server-0.7.2-py3-none-manylinux_2_31_x86_64.whl", hash = "sha256:ef687163c24185ae9754ed5650eb5bc4d84ff257aabdc33f0cc6f74d8ba54530", size = 6590363, upload-time = "2023-10-23T21:23:35.583Z" }, ] +[[package]] +name = "terminal-bench-2-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/terminal_bench_2_v1" } +dependencies = [ + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [{ name = "verifiers" }] + [[package]] name = "textarena" version = "0.7.4" @@ -5644,11 +5258,11 @@ wheels = [ [[package]] name = "tomlkit" -version = "0.14.0" +version = "0.15.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/c3/af/14b24e41977adb296d6bd1fb59402cf7d60ce364f90c890bd2ec65c43b5a/tomlkit-0.14.0.tar.gz", hash = "sha256:cf00efca415dbd57575befb1f6634c4f42d2d87dbba376128adb42c121b87064", size = 187167, upload-time = "2026-01-13T01:14:53.304Z" } +sdist = { url = "https://files.pythonhosted.org/packages/51/db/03eaf4331631ef6b27d6e3c9b68c54dc6f0d63d87201fed600cc409307fd/tomlkit-0.15.0.tar.gz", hash = "sha256:7d1a9ecba3086638211b13814ea79c90dd54dd11993564376f3aa92271f5c7a3", size = 161875, upload-time = "2026-05-10T07:38:22.245Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b5/11/87d6d29fb5d237229d67973a6c9e06e048f01cf4994dee194ab0ea841814/tomlkit-0.14.0-py3-none-any.whl", hash = "sha256:592064ed85b40fa213469f81ac584f67a4f2992509a7c3ea2d632208623a3680", size = 39310, upload-time = "2026-01-13T01:14:51.965Z" }, + { url = "https://files.pythonhosted.org/packages/6a/43/8bd850ee71a191bf072e31302c73a66be413fecdd98fdcd111ecbcce13ca/tomlkit-0.15.0-py3-none-any.whl", hash = "sha256:4dbc8f0fc024412b57ced8757ac7461305126a648ff8c2c807fcb8e133a78738", size = 41328, upload-time = "2026-05-10T07:38:23.517Z" }, ] [[package]] @@ -5912,15 +5526,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/61/73/d21edf5b204d1467e06500080a50f79d49ef2b997c79123a536d4a17d97c/uc_micro_py-2.0.0-py3-none-any.whl", hash = "sha256:3603a3859af53e5a39bc7677713c78ea6589ff188d70f4fee165db88e22b242c", size = 6383, upload-time = "2026-03-01T06:31:26.257Z" }, ] -[[package]] -name = "uncalled-for" -version = "0.3.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b5/82/345cc927f7fbdae6065e7768759932fcc827fc20b29b45dfbafa2f1f7da4/uncalled_for-0.3.2.tar.gz", hash = "sha256:89f5dbcd71e2b8f47c030b1fa302e6cce2ec795d1ac565eeb6525c5fe55cb8a2", size = 50032, upload-time = "2026-05-06T13:38:25.204Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3b/25/2c87754f3a9e692315f7b811244090e68f362979fc8886b3fbd2985a1d8c/uncalled_for-0.3.2-py3-none-any.whl", hash = "sha256:0ff60b142c7d1f8070bde9d42afaa70aedc77dcc10998c227687e9c15713418e", size = 11444, upload-time = "2026-05-06T13:38:24.025Z" }, -] - [[package]] name = "unidiff" version = "0.7.5" @@ -5981,40 +5586,30 @@ dependencies = [ { name = "aiohttp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "aiolimiter", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "anthropic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "certifi", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "gepa", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "httpx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "jinja2", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "loguru", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "math-verify", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "mcp", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "msgpack", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nest-asyncio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "openai", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "openai-agents", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pillow", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "prime-pydantic-config", extra = ["toml"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "prime-sandboxes", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "prime-tunnel", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "pydantic", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "pymupdf", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "pyzmq", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "regex", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "renderers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "requests", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "rich", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "setproctitle", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "tenacity", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "textual", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tomli-w", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "uvloop", marker = "(platform_machine == 'aarch64' and platform_python_implementation != 'PyPy' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'PyPy' and sys_platform == 'linux')" }, ] -[package.optional-dependencies] -packages = [ - { name = "harnesses", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "tasksets", extra = ["openenv", "openreward", "ta"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] - [package.metadata] requires-dist = [ { name = "accelerate", marker = "extra == 'rl'", specifier = ">=1.4.0" }, @@ -6022,51 +5617,42 @@ requires-dist = [ { name = "aiohttp", marker = "extra == 'browser'", specifier = ">=3.9.0" }, { name = "aiolimiter", specifier = ">=1.2.1" }, { name = "anthropic", specifier = ">=0.78.0" }, - { name = "certifi" }, + { name = "certifi", marker = "extra == 'quest'" }, { name = "datasets", specifier = ">=3.0.0,<4.7.0" }, { name = "deepspeed", marker = "extra == 'rl'", specifier = ">=0.17.6" }, { name = "flash-attn", marker = "extra == 'rl'", specifier = ">=2.8.3" }, { name = "gepa" }, - { name = "harnesses", marker = "extra == 'harnesses'", editable = "deps/verifiers/packages/harnesses" }, - { name = "harnesses", marker = "extra == 'packages'", editable = "deps/verifiers/packages/harnesses" }, - { name = "harnesses", extras = ["nemogym"], marker = "extra == 'nemogym'", editable = "deps/verifiers/packages/harnesses" }, { name = "httpx", specifier = ">=0.27.0" }, - { name = "jinja2", specifier = ">=3.1.6" }, { name = "liger-kernel", marker = "extra == 'rl'", specifier = ">=0.5.10" }, + { name = "loguru", specifier = ">=0.7.0" }, { name = "math-verify", specifier = ">=0.8.0" }, { name = "mcp", specifier = ">=1.14.1" }, + { name = "modal", marker = "extra == 'modal'", specifier = ">=1.4.0" }, { name = "msgpack", specifier = ">=1.1.2" }, - { name = "nest-asyncio", specifier = ">=1.6.0" }, - { name = "nltk", marker = "extra == 'ta'" }, + { name = "nest-asyncio", marker = "extra == 'notebook'", specifier = ">=1.6.0" }, + { name = "nltk", marker = "extra == 'ta'", specifier = ">=3.9.2" }, { name = "numpy" }, { name = "openai", specifier = ">=1.108.1" }, { name = "openai-agents", specifier = ">=0.0.7" }, { name = "peft", marker = "extra == 'rl'" }, - { name = "pillow" }, - { name = "prime-pydantic-config", extras = ["toml"] }, - { name = "prime-sandboxes", specifier = ">=0.2.25" }, + { name = "pillow", marker = "extra == 'quest'" }, + { name = "prime-pydantic-config", extras = ["toml"], specifier = ">=0.3.0.dev86" }, + { name = "prime-sandboxes", specifier = ">=0.2.27" }, { name = "prime-tunnel", specifier = ">=0.1.8" }, { name = "pydantic", specifier = ">=2.11.9" }, - { name = "pymupdf" }, + { name = "pymupdf", marker = "extra == 'quest'" }, { name = "python-dotenv", marker = "extra == 'browser'", specifier = ">=1.0.0" }, { name = "pyzmq", specifier = ">=27.1.0" }, { name = "reasoning-gym", marker = "extra == 'rg'" }, - { name = "regex", specifier = "<2026.4.4" }, - { name = "renderers", marker = "extra == 'renderers'", specifier = ">=0.1.8.dev28" }, + { name = "renderers", specifier = ">=0.1.8.dev40" }, { name = "requests" }, { name = "requests", marker = "extra == 'rl'" }, { name = "rich" }, { name = "setproctitle", specifier = ">=1.3.0" }, { name = "stagehand", marker = "extra == 'browser'", specifier = ">=3.0.0" }, - { name = "tasksets", extras = ["nemogym"], marker = "extra == 'nemogym'", editable = "deps/verifiers/packages/tasksets" }, - { name = "tasksets", extras = ["openenv"], marker = "extra == 'openenv'", editable = "deps/verifiers/packages/tasksets" }, - { name = "tasksets", extras = ["openenv", "openreward", "ta"], marker = "extra == 'packages'", editable = "deps/verifiers/packages/tasksets" }, - { name = "tasksets", extras = ["openenv", "openreward", "ta"], marker = "extra == 'tasksets'", editable = "deps/verifiers/packages/tasksets" }, - { name = "tasksets", extras = ["openreward"], marker = "extra == 'openreward'", editable = "deps/verifiers/packages/tasksets" }, { name = "tenacity", specifier = ">=8.5.0" }, - { name = "textarena", marker = "extra == 'ta'" }, - { name = "textual" }, - { name = "tomli", marker = "python_full_version < '3.11'" }, + { name = "textarena", marker = "extra == 'ta'", specifier = "==0.7.4" }, + { name = "tomli-w", specifier = ">=1.0.0" }, { name = "torch", marker = "extra == 'rl'", specifier = ">=2.8.0,<2.9.0" }, { name = "transformers", marker = "extra == 'rl'", specifier = ">=4.56.2" }, { name = "typing-extensions", marker = "python_full_version < '3.12'" }, @@ -6074,14 +5660,14 @@ requires-dist = [ { name = "vllm", marker = "extra == 'rl'", specifier = ">=0.10.0,<0.11.0" }, { name = "wandb", marker = "extra == 'rl'" }, ] -provides-extras = ["browser", "harnesses", "nemogym", "openenv", "openreward", "packages", "renderers", "rg", "rl", "ta", "tasksets"] +provides-extras = ["browser", "modal", "notebook", "quest", "rg", "rl", "ta"] [package.metadata.requires-dev] dev = [ { name = "aiohttp", specifier = ">=3.9.0" }, - { name = "harnesses", editable = "deps/verifiers/packages/harnesses" }, { name = "ipykernel" }, { name = "ipywidgets" }, + { name = "nltk", specifier = ">=3.9.2" }, { name = "pre-commit" }, { name = "pytest", specifier = ">=7.0.0" }, { name = "pytest-asyncio", specifier = ">=0.21.0" }, @@ -6089,13 +5675,34 @@ dev = [ { name = "pytest-xdist", specifier = ">=3.8.0" }, { name = "python-dotenv", specifier = ">=1.0.0" }, { name = "reasoning-gym" }, - { name = "renderers", specifier = ">=0.1.8.dev28" }, { name = "ruff" }, { name = "stagehand", specifier = ">=3.0.0" }, - { name = "tasksets", extras = ["openenv", "openreward", "ta"], editable = "deps/verifiers/packages/tasksets" }, + { name = "textarena", specifier = "==0.7.4" }, { name = "ty", specifier = ">=0.0.1a29,<0.0.22" }, ] -policy = [{ name = "semgrep", specifier = ">=1.150.0" }] +examples = [ + { name = "aime24-v1", editable = "deps/verifiers/environments/aime24_v1" }, + { name = "alphabet-sort-v1", editable = "deps/verifiers/environments/alphabet_sort_v1" }, + { name = "code-golf-v1", editable = "deps/verifiers/environments/code_golf_v1" }, + { name = "color-codeword-v1", editable = "deps/verifiers/environments/color_codeword_v1" }, + { name = "compact", editable = "deps/verifiers/environments/compact" }, + { name = "deepwiki-v1", editable = "deps/verifiers/environments/deepwiki_v1" }, + { name = "general-agent-v1", editable = "deps/verifiers/environments/general_agent_v1" }, + { name = "glossary-v1", editable = "deps/verifiers/environments/glossary_v1" }, + { name = "gsm8k-v1", editable = "deps/verifiers/environments/gsm8k_v1" }, + { name = "math-env-v1", editable = "deps/verifiers/environments/math_env_v1" }, + { name = "r2e-gym-v1", editable = "deps/verifiers/environments/r2e_gym_v1" }, + { name = "reverse-text-v1", editable = "deps/verifiers/environments/reverse_text_v1" }, + { name = "scaleswe-v1", editable = "deps/verifiers/environments/scaleswe_v1" }, + { name = "scratchpad-v1", editable = "deps/verifiers/environments/scratchpad_v1" }, + { name = "swebench-verified-v1", editable = "deps/verifiers/environments/swebench_verified_v1" }, + { name = "swelego-v1", editable = "deps/verifiers/environments/swelego_v1" }, + { name = "tau2-bench-v1", editable = "deps/verifiers/environments/tau2_bench_v1" }, + { name = "terminal-bench-2-v1", editable = "deps/verifiers/environments/terminal_bench_2_v1" }, + { name = "wiki-search-v1", editable = "deps/verifiers/environments/wiki_search_v1" }, + { name = "wikispeedia-v1", editable = "deps/verifiers/environments/wikispeedia_v1" }, + { name = "wordle-v1", editable = "deps/verifiers/environments/wordle_v1" }, +] [[package]] name = "virtualenv" @@ -6622,7 +6229,7 @@ dependencies = [ { name = "chromadb", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "openai", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -6633,6 +6240,34 @@ requires-dist = [ { name = "verifiers", specifier = ">=0.1.9" }, ] +[[package]] +name = "wiki-search-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/wiki_search_v1" } +dependencies = [ + { name = "chromadb", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "datasets", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "chromadb" }, + { name = "datasets" }, + { name = "verifiers" }, +] + +[[package]] +name = "wikispeedia-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/wikispeedia_v1" } +dependencies = [ + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [{ name = "verifiers" }] + [[package]] name = "wordle" version = "0.1.7" @@ -6640,7 +6275,7 @@ source = { editable = "deps/verifiers/environments/wordle" } dependencies = [ { name = "nltk", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "textarena", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "verifiers", extra = ["packages"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] [package.metadata] @@ -6650,6 +6285,17 @@ requires-dist = [ { name = "verifiers", specifier = ">=0.1.9.post3" }, ] +[[package]] +name = "wordle-v1" +version = "0.1.0" +source = { editable = "deps/verifiers/environments/wordle_v1" } +dependencies = [ + { name = "verifiers", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [{ name = "verifiers", extras = ["ta"] }] + [[package]] name = "wrapt" version = "1.17.3"