From eaf262ffe88a310759af255444176c1ad2a0a164 Mon Sep 17 00:00:00 2001 From: huxy1225-cloud Date: Sat, 15 Aug 2026 21:12:05 +0800 Subject: [PATCH] feat(agentic): add GraphGPO ALFWorld recipe --- examples/graphgpo/README.md | 296 +++++++ examples/graphgpo/REPRODUCIBILITY.md | 60 ++ examples/graphgpo/__init__.py | 42 + examples/graphgpo/action_parser.py | 61 ++ examples/graphgpo/alfworld_env.py | 411 +++++++++ .../configs/alfworld_qwen2_5_1_5b.yaml | 30 + examples/graphgpo/custom_advantage.py | 394 +++++++++ examples/graphgpo/diagnostics.py | 353 ++++++++ examples/graphgpo/diagnostics_summary.py | 28 + examples/graphgpo/eval_logger.py | 133 +++ examples/graphgpo/graph_credit.py | 794 ++++++++++++++++++ examples/graphgpo/manifest.py | 244 ++++++ examples/graphgpo/preflight.py | 377 +++++++++ examples/graphgpo/prepare_alfworld.py | 401 +++++++++ examples/graphgpo/prompt.py | 139 +++ examples/graphgpo/reproducibility.py | 476 +++++++++++ examples/graphgpo/requirements-alfworld.txt | 4 + examples/graphgpo/reward.py | 37 + examples/graphgpo/rollout_agent.py | 605 +++++++++++++ examples/graphgpo/run_agent_app.sh | 15 + .../graphgpo/run_alfworld_qwen2_5_1_5b.sh | 358 ++++++++ examples/graphgpo/state.py | 236 ++++++ relax/agentic/pipeline/reward.py | 76 +- relax/agentic/pipeline/transfer.py | 654 ++++++++++++++- relax/backends/megatron/actor.py | 763 ++++++++++++++++- relax/backends/megatron/data.py | 385 ++++++++- relax/backends/megatron/loss.py | 10 +- relax/backends/megatron/model.py | 189 ++++- relax/core/controller.py | 68 +- ...test_agentic_variable_row_actor_helpers.py | 531 ++++++++++++ .../megatron/test_variable_row_loss_scale.py | 167 ++++ .../test_variable_row_metric_count.py | 44 + .../megatron/test_variable_row_plan.py | 260 ++++++ .../test_variable_row_scheduler_accounting.py | 191 +++++ tests/core/test_controller_variable_rows.py | 119 +++ tests/graphgpo/test_action_parser.py | 55 ++ tests/graphgpo/test_alfworld_env.py | 229 +++++ tests/graphgpo/test_custom_advantage.py | 407 +++++++++ tests/graphgpo/test_diagnostics.py | 175 ++++ .../graphgpo/test_dynamic_turn_accounting.py | 184 ++++ tests/graphgpo/test_eval_logger.py | 86 ++ tests/graphgpo/test_graph_credit.py | 569 +++++++++++++ tests/graphgpo/test_grpo_one_step_parity.py | 381 +++++++++ tests/graphgpo/test_manifest.py | 122 +++ tests/graphgpo/test_preflight.py | 173 ++++ tests/graphgpo/test_prepare_alfworld.py | 193 +++++ tests/graphgpo/test_prompt.py | 124 +++ tests/graphgpo/test_recipe_packaging.py | 133 +++ tests/graphgpo/test_reproducibility.py | 158 ++++ tests/graphgpo/test_reward.py | 38 + tests/graphgpo/test_rollout_agent.py | 570 +++++++++++++ tests/graphgpo/test_state.py | 169 ++++ tests/test_agentic_variable_row_transfer.py | 548 ++++++++++++ 53 files changed, 13184 insertions(+), 81 deletions(-) create mode 100644 examples/graphgpo/README.md create mode 100644 examples/graphgpo/REPRODUCIBILITY.md create mode 100644 examples/graphgpo/__init__.py create mode 100644 examples/graphgpo/action_parser.py create mode 100644 examples/graphgpo/alfworld_env.py create mode 100644 examples/graphgpo/configs/alfworld_qwen2_5_1_5b.yaml create mode 100644 examples/graphgpo/custom_advantage.py create mode 100644 examples/graphgpo/diagnostics.py create mode 100644 examples/graphgpo/diagnostics_summary.py create mode 100644 examples/graphgpo/eval_logger.py create mode 100644 examples/graphgpo/graph_credit.py create mode 100644 examples/graphgpo/manifest.py create mode 100644 examples/graphgpo/preflight.py create mode 100644 examples/graphgpo/prepare_alfworld.py create mode 100644 examples/graphgpo/prompt.py create mode 100644 examples/graphgpo/reproducibility.py create mode 100644 examples/graphgpo/requirements-alfworld.txt create mode 100644 examples/graphgpo/reward.py create mode 100644 examples/graphgpo/rollout_agent.py create mode 100644 examples/graphgpo/run_agent_app.sh create mode 100644 examples/graphgpo/run_alfworld_qwen2_5_1_5b.sh create mode 100644 examples/graphgpo/state.py create mode 100644 tests/backends/megatron/test_agentic_variable_row_actor_helpers.py create mode 100644 tests/backends/megatron/test_variable_row_loss_scale.py create mode 100644 tests/backends/megatron/test_variable_row_metric_count.py create mode 100644 tests/backends/megatron/test_variable_row_plan.py create mode 100644 tests/backends/megatron/test_variable_row_scheduler_accounting.py create mode 100644 tests/core/test_controller_variable_rows.py create mode 100644 tests/graphgpo/test_action_parser.py create mode 100644 tests/graphgpo/test_alfworld_env.py create mode 100644 tests/graphgpo/test_custom_advantage.py create mode 100644 tests/graphgpo/test_diagnostics.py create mode 100644 tests/graphgpo/test_dynamic_turn_accounting.py create mode 100644 tests/graphgpo/test_eval_logger.py create mode 100644 tests/graphgpo/test_graph_credit.py create mode 100644 tests/graphgpo/test_grpo_one_step_parity.py create mode 100644 tests/graphgpo/test_manifest.py create mode 100644 tests/graphgpo/test_preflight.py create mode 100644 tests/graphgpo/test_prepare_alfworld.py create mode 100644 tests/graphgpo/test_prompt.py create mode 100644 tests/graphgpo/test_recipe_packaging.py create mode 100644 tests/graphgpo/test_reproducibility.py create mode 100644 tests/graphgpo/test_reward.py create mode 100644 tests/graphgpo/test_rollout_agent.py create mode 100644 tests/graphgpo/test_state.py create mode 100644 tests/test_agentic_variable_row_transfer.py diff --git a/examples/graphgpo/README.md b/examples/graphgpo/README.md new file mode 100644 index 000000000..94dca37e5 --- /dev/null +++ b/examples/graphgpo/README.md @@ -0,0 +1,296 @@ +# GraphGPO on ALFWorld + +This recipe reproduces GraphGPO on ALFWorld text-only with +Qwen2.5-1.5B-Instruct. It keeps the Relax rollout and Megatron training path: +each environment turn is one named explicit export, and the group-level hook +returns one scalar advantage for each exported turn. + +The recipe supports three methods through the same data and launch path: + +- `METHOD=grpo`: episode-level GRPO advantage only. +- `METHOD=gigpo`: GiGPO same-state step advantage plus episode advantage. +- `METHOD=graphgpo`: graph edge advantage plus episode advantage. + +The formal reproduction uses the default +`EPISODE_WEIGHTING=trajectory_once`: each episode return participates in the +group statistics exactly once, then the normalized value is broadcast to that +trajectory's turns. This preserves native GRPO weighting when trajectories +have different lengths. `EPISODE_WEIGHTING=reference_cross_steps` is retained +only as an explicit historical-reference ablation; it must not be used for the +main GRPO, GiGPO, or GraphGPO comparison. + +## Frozen algorithm contract + +Each custom-advantage invocation contains one current rollout group: eight +trajectories from exactly one task. Episode return is +`10 * success - 0.1 * invalid_action_count`. The default episode advantage +standardizes those eight trajectory returns with sample standard deviation +(`ddof=1`) and `epsilon=1e-6`, once per trajectory, before broadcasting the +result to every real turn in that trajectory. + +The graph has one occurrence edge per real turn. Repeated and parallel edges +remain repeated when same-source graph returns are standardized, while the +shortest-path topology uses the minimum observed cost for each source/target +pair. The final real action of a successful trajectory points to the canonical +`__GRAPHGPO_SUCCESS__` sink. Self-loops and cycles are allowed. Nodes that +cannot reach success receive `max_finite_distance + 1`; an all-fail graph has +zero graph advantage, and singleton or zero-variance groups also return zero. + +The accepted Task 37/reference-commit reward is +`10 * omega ** d(next_state)` (unit costs, `omega=0.1`). Thus same-source edges +whose next-state distances are `0`, `1`, and `2` have the fixed oracle +`[10, 1, 0.1]`. This intentionally differs from paper Eq. 4, +`10 * omega ** (d(next_state) + cost)`, whose unit-cost values are smaller by a +factor of `omega`. Both conventions remain explicit in the graph kernel, but +training uses the Task/reference convention. + +An ALFWorld state key is SHA256 over `reference_anchor_v1`: the raw observation +plus the reference-visible tracker context and the complete sorted admissible +command list (including `help`). Tracker input is restricted to exactly +`location`, `holding`, `history_items`, and `item_location`. Top-level and +nested mappings are canonicalized with the fixed whitelist order and sorted +nested keys, finite JSON values, UTF-8 text, and no insignificant whitespace +before rendering and hashing. + +Graph diagnostics are disabled by default, including their timers. To record +them, set `GRAPHGPO_DIAGNOSTICS_JSONL` to an output path. Each task-local graph +then appends one JSONL record with `perf_counter_ns` durations for graph build, +reverse Dijkstra, and graph-advantage calculation; node and edge-occurrence +counts; exact `(source, action, target)` duplicate rate; cross-trajectory +shared-source rate; singleton-source rate; nonzero-advantage rate; all-fail +status; unreachable-node rate; and a node-distance histogram. The diagnostics +operate only on the pure-Python graph objects and do not synchronize GPU work. + +```bash +GRAPHGPO_DIAGNOSTICS_JSONL=/path/to/run/graph-diagnostics.jsonl \ + METHOD=graphgpo ... \ + bash examples/graphgpo/run_alfworld_qwen2_5_1_5b.sh +``` + +After the run, retain the raw JSONL and create a non-overwriting summary with +per-stage totals plus median and nearest-rank p95 values: + +```bash +python3 -m examples.graphgpo.diagnostics_summary \ + --input /path/to/run/graph-diagnostics.jsonl \ + --output /path/to/run/graph-diagnostics.summary.json +``` + +Every turn's metadata must carry unique `row_id` plus `rollout_group_id` and +`policy_version`. The custom adapter rejects mixed group/version inputs before +using the compact local `(trajectory_id, turn_index)` graph key. The latter +therefore remains compatible with existing graph math and output lookup while +the full identity tuple is preserved and checked at the transport boundary. + +The implementation follows the frozen Task 37 proposal. It does not claim the +paper's reported success rate until the registered multi-seed experiment has +finished. + +## Frozen inputs + +- Relax commit: `8a54679e971087566bfda939a5f75649e07fb861` +- GraphGPO reference commit: + `20bd331bdbc9026a5668e11362178e10ab7400c8` +- Model: + `Qwen/Qwen2.5-1.5B-Instruct@775b11afaf83e0dc75bd5abaf90133e47b3ec082` +- ALFWorld: `0.4.2`, text-only +- Expected base image for the external execution protocol: + `ghcr.io/redai-infra/relaxrl@sha256:3fa8ce578acda6c829b83016bde42c38fa892681e4f36ca330f545616fe578e2` + +Install ALFWorld in an isolated environment and download its text data: + +```bash +uv venv --python 3.10 /path/to/alfworld-0.4.2 +uv pip install --python /path/to/alfworld-0.4.2/bin/python \ + -r examples/graphgpo/requirements-alfworld.txt +export ALFWORLD_DATA=/path/to/alfworld-data +/path/to/alfworld-0.4.2/bin/alfworld-download \ + --data-dir "${ALFWORLD_DATA}" +uv pip freeze --python /path/to/alfworld-0.4.2/bin/python \ + > /path/to/locks/alfworld-0.4.2-python.freeze.txt +``` + +Do not start GPU training until importing ALFWorld and a real text-only +`reset`/`step` smoke both pass. + +If ALFWorld is installed in a separate interpreter, set +`ALFWORLD_PYTHON=/path/to/alfworld-0.4.2/bin/python`. That interpreter also +uses the pinned OpenAI Python SDK and `httpx` in +`requirements-alfworld.txt`, because the managed process talks to Relax's +local chat-completions endpoint. The training process itself still uses the +image's normal `python3`. + +## Prepare manifests and prompt rows + +The preparation command does not import ALFWorld. It filters the downloaded +games using ALFWorld 0.4.2's text-environment rules, hashes every selected game, +trajectory, PDDL, and grammar file, and writes one seed row per task. Each row +stores the complete `ALFWORLD_DATA`-relative `game.tw-pddl` path as `task_id`. +The environment adapter sorts all discovered game files, restricts the +single-task environment to that exact file, and verifies the ID returned by +`reset`. A missing, outside-root, or mismatched game fails the trajectory. + +```bash +export ALFWORLD_DATA=/path/to/alfworld-data +python3 -m examples.graphgpo.prepare_alfworld \ + --data-root "${ALFWORLD_DATA}" \ + --output-dir /path/to/task37-data-lock +``` + +The default outputs are: + +```text +train.manifest.json +train.prompts.jsonl +eval_in_distribution.manifest.json +eval_in_distribution.prompts.jsonl +prepare.lock.json +``` + +Every prompt row also freezes the request-level sampling contract consumed by +the managed agent: train uses `temperature=1.0`, validation uses +`temperature=0.4`, and both use `top_p=1.0`, `max_tokens=512`. These values +are recorded in `prepare.lock.json`; the OpenAI-compatible request does not +rely on an inference-server default. + +Training keeps all eligible train games. Validation takes the first 128 +eligible `valid_seen` games in normalized path order. Override either choice +explicitly: + +```bash +python3 -m examples.graphgpo.prepare_alfworld \ + --data-root "${ALFWORLD_DATA}" \ + --output-dir /path/to/task37-data-lock-custom \ + --limit train=512 \ + --limit eval_in_distribution=128 +``` + +Prepared files are content-addressed and idempotent. The command refuses to +replace an existing file with different bytes. + +Before a dry run or training launch, the recipe rechecks the lock, manifest, +and prompt hashes; split names; task counts and order; `MAX_STEPS`; and the +model revision. It also rehashes every ALFWorld game, trajectory, and shared +asset referenced by the selected manifests. A mismatch stops before Ray or a +GPU job is contacted. + +The model snapshot has an independent JSON lock with schema +`task37-huggingface-model-lock-v1`. It pins the repository revision and the +size and SHA256 of all nine snapshot files. Set `MODEL_LOCK` to that file; the +preflight rehashes the local checkpoint before launch. + +The environment config keeps `num_train_games` and `num_eval_games` at `-1`. +This is intentional: ALFWorld must not truncate its unsorted discovery list +before the adapter applies normalized sorting and the manifest task binding. + +## Dependency-free checks + +The graph kernel, action parser, state adapter, manifest builder, managed-agent +contract, and packaging tests run without ALFWorld or a GPU: + +```bash +python -m pytest tests/graphgpo -q +``` + +The ALFWorld import remains inside the real environment factory, so importing +`examples.graphgpo` and running fake-environment tests does not require the +optional dependency. + +Agentic group-RM evaluation still invokes a batched reward function. The +launcher therefore wires +`examples.graphgpo.reward.reward_func`, which validates and forwards the +finite environment reward already attached to every explicit turn row. It +does not rescore model text. + +## Inspect the launch without using a GPU + +Set paths to existing local fixtures and use `DRY_RUN=1`. The script validates +the frozen inputs and prints the fully expanded Relax command without +contacting Ray: + +```bash +DRY_RUN=1 \ +METHOD=graphgpo \ +ALFWORLD_DATA=/path/to/alfworld-data \ +HF_CHECKPOINT=/path/to/Qwen2.5-1.5B-Instruct-snapshot \ +SAVE_DIR=/path/to/task37-output \ +DATA_ARTIFACT_DIR=/path/to/task37-data-lock \ +DEPENDENCY_LOCK=/path/to/locks/alfworld-0.4.2-python.freeze.txt \ +MODEL_LOCK=/path/to/locks/qwen2.5-1.5b-775b11af-v1.lock.json \ +ALFWORLD_PYTHON=/path/to/alfworld-0.4.2/bin/python \ +bash examples/graphgpo/run_alfworld_qwen2_5_1_5b.sh +``` + +The default environment config is +`configs/alfworld_qwen2_5_1_5b.yaml`. It uses only the text environment and +expands all dataset paths from `ALFWORLD_DATA`. + +Evaluation is enabled by default with `ENABLE_EVAL=1`, preserving the formal +128-task evaluation path. For a rollout/training plumbing smoke only, set +`ENABLE_EVAL=0`; the launcher then omits the evaluation artifacts from +preflight and installs no evaluation arguments. The selected value is written +to the run lock. Formal runs must retain the default. + +## Train + +After the real environment smoke and the dry-run command pass: + +```bash +METHOD=graphgpo \ +SEED=0 \ +NUM_GPUS=2 \ +OUTER_EPOCHS=150 \ +TASK_GROUPS=16 \ +GROUP_SIZE=8 \ +GLOBAL_BATCH_SIZE=128 \ +ALFWORLD_DATA=/path/to/alfworld-data \ +HF_CHECKPOINT=/path/to/Qwen2.5-1.5B-Instruct-snapshot \ +SAVE_DIR=/path/to/task37-output \ +DATA_ARTIFACT_DIR=/path/to/task37-data-lock \ +DEPENDENCY_LOCK=/path/to/locks/alfworld-0.4.2-python.freeze.txt \ +MODEL_LOCK=/path/to/locks/qwen2.5-1.5b-775b11af-v1.lock.json \ +ALFWORLD_PYTHON=/path/to/alfworld-0.4.2/bin/python \ +bash examples/graphgpo/run_alfworld_qwen2_5_1_5b.sh +``` + +Change only `METHOD` and `SEED` for the registered GRPO/GiGPO/GraphGPO +comparison. The script fixes group size, sampling, batch, optimizer, KL, and +model settings across methods. Keep the default `trajectory_once` weighting +unchanged across the three methods. SGLang static memory fraction defaults to +the reference recipe's `0.50` for every method; any explicit +`SGLANG_MEM_FRACTION_STATIC` override is validated in `(0, 1]` and recorded in +the run lock. The script also records a local run lock and the +expanded command under `SAVE_DIR/runs/`. These raw files can contain local +paths and must be scrubbed before they are copied into public report material. +The run lock records the image digest expected by the outer execution +protocol, but the recipe script cannot prove which image launched it. The +outer container/Ray job runner must actually select that digest and retain its +own runtime inspection evidence. + +For the two registered ablations: + +```bash +# Episode-only parity. +METHOD=graphgpo BETA=0 BETA_EPISODE=1 ... + +# Graph-only. +METHOD=graphgpo BETA=1 BETA_EPISODE=0 ... +``` + +`RELAX_AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE` is set to `MAX_STEPS` by the +launcher. Variable turn counts therefore use the explicit variable-row path; +the final physical batch is padded only with zero-loss rows. + +## Output contract + +Each trajectory writes `turn_000`, `turn_001`, and so on as JSONL records. +Every record contains only the current user prompt and assistant response. +Metadata carries the task, trajectory, turn, state transition, validity, +terminal state, and episode return needed by +`examples.graphgpo.custom_advantage.compute_custom_advantage`. + +Reports must aggregate success and return after deduplicating by +`trajectory_id`; averaging per-turn rows would length-weight long episodes. +The launcher installs the custom eval hook that records episode count, +success rate, mean episode return, and truncation rate with exactly one vote +per trajectory. diff --git a/examples/graphgpo/REPRODUCIBILITY.md b/examples/graphgpo/REPRODUCIBILITY.md new file mode 100644 index 000000000..12fd65455 --- /dev/null +++ b/examples/graphgpo/REPRODUCIBILITY.md @@ -0,0 +1,60 @@ +# GraphGPO reproducibility evidence + +`examples.graphgpo.reproducibility` generates a content-addressed evidence +bundle without treating planned runs or operator-entered versions as verified +results. + +First retain the complete machine-readable test report and the complete +pre-commit output. For example: + +```bash +python -m pytest tests/graphgpo \ + --junitxml=/path/to/evidence/pytest-graphgpo.xml \ + 2>&1 | tee /path/to/evidence/pytest-graphgpo.log +pre-commit run --all-files \ + 2>&1 | tee /path/to/evidence/pre-commit.log +``` + +Then generate the bundle. Artifact and report labels are public identifiers; +local source paths are used only to read and hash files and are not serialized. + +```bash +python -m examples.graphgpo.reproducibility \ + --output-dir /path/to/evidence/bundle \ + --artifact train_manifest=/path/to/train.manifest.json \ + --artifact eval_manifest=/path/to/eval_in_distribution.manifest.json \ + --artifact expanded_command=/path/to/expanded_command.sh \ + --artifact run_lock=/path/to/run_lock.env \ + --junit pytest_graphgpo=/path/to/evidence/pytest-graphgpo.xml \ + --log pytest_graphgpo_full=/path/to/evidence/pytest-graphgpo.log \ + --log pre_commit_full=/path/to/evidence/pre-commit.log \ + --version relax_baseline_commit= \ + --version candidate_commit= \ + --version graphgpo_reference_commit= \ + --version container_image_digest=@sha256:<64-hex-digest> \ + --version model_revision= \ + --version alfworld_version= +``` + +Omit a `--version` argument when it is not known. In particular, do not use a +worktree label or the baseline commit as `candidate_commit`; leave it absent +until an immutable candidate commit exists. Missing values are emitted as +`null` with `verification: not_claimed`. Supplied values are marked +`operator_supplied` and `not_verified_by_generator` because a declaration by +itself is not execution evidence. + +The generated files are: + +- `seed_mapping.json`: the registered GRPO, GiGPO, GraphGPO, and graph-only + conditions for seeds 0, 1, and 2. Every entry remains `planned`; this file + never claims a run completed. +- `test_report.json`: test counts parsed from JUnit XML and SHA256 records for + full text logs. Outcomes are never inferred from unstructured text. +- `reproducibility.manifest.json`: SHA256 and byte size for supplied artifacts + and generated reports, declared revision fields, and versions observed in + the interpreter that generated the bundle. + +Runtime package observations describe the evidence-generator interpreter, +not a remote training image. Container-side package, CUDA, driver, hardware, +and image inspection records should therefore be supplied as hashed input +artifacts rather than inferred from this command. diff --git a/examples/graphgpo/__init__.py b/examples/graphgpo/__init__.py new file mode 100644 index 000000000..697d58fea --- /dev/null +++ b/examples/graphgpo/__init__.py @@ -0,0 +1,42 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""GraphGPO reproduction helpers.""" + +from .graph_credit import ( + SUCCESS, + DistanceResult, + EdgeOccurrence, + EpisodeWeighting, + OccurrenceGraph, + Turn, + build_occurrence_graph, + compute_method_advantages, + episode_advantages, + episode_return, + finalize_trajectory, + gigpo_discounted_returns, + graph_advantages, + graph_raw_returns, + reverse_shortest_distances, + standardize_by_key, +) + + +__all__ = [ + "SUCCESS", + "DistanceResult", + "EdgeOccurrence", + "EpisodeWeighting", + "OccurrenceGraph", + "Turn", + "build_occurrence_graph", + "compute_method_advantages", + "episode_advantages", + "episode_return", + "finalize_trajectory", + "gigpo_discounted_returns", + "graph_advantages", + "graph_raw_returns", + "reverse_shortest_distances", + "standardize_by_key", +] diff --git a/examples/graphgpo/action_parser.py b/examples/graphgpo/action_parser.py new file mode 100644 index 000000000..6baa67b63 --- /dev/null +++ b/examples/graphgpo/action_parser.py @@ -0,0 +1,61 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Reference-compatible action parsing helpers for the ALFWorld recipe.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Literal + + +_CHINESE_PATTERN = re.compile(r"[\u4e00-\u9fff]") +_ACTION_PATTERN = re.compile(r"(.*?)", flags=re.DOTALL) + + +@dataclass(frozen=True) +class ParsedAction: + """The environment action extracted from one model response.""" + + action: str + is_valid: bool + source: Literal["tag", "fallback"] + + +def _contains_chinese(text: str) -> bool: + return _CHINESE_PATTERN.search(text) is not None + + +def parse_action(response: str, *, fallback_chars: int = 30) -> ParsedAction: + """Extract the first ```` block after lowercasing the response. + + This intentionally mirrors the frozen GraphGPO reference parser instead of + adding stricter recipe-specific checks. A missing tag falls back to the + last ``fallback_chars`` characters for debugging, but that result is always + marked invalid. Duplicate tags and actions outside the admissible set are + not rejected. The reference validity flag additionally requires lowercase + ```` tags in the original response and rejects Chinese text anywhere + in that response. + """ + + if not isinstance(response, str): + raise TypeError("response must be a string") + if isinstance(fallback_chars, bool) or not isinstance(fallback_chars, int) or fallback_chars <= 0: + raise ValueError("fallback_chars must be a positive integer") + + lowered_response = response.lower() + match = _ACTION_PATTERN.search(lowered_response) + if match is None: + action = lowered_response[-fallback_chars:] + source: Literal["tag", "fallback"] = "fallback" + is_valid = False + else: + action = match.group(1).strip() + source = "tag" + is_valid = True + + if "" not in response or "" not in response: + is_valid = False + if _contains_chinese(response): + is_valid = False + return ParsedAction(action=action, is_valid=is_valid, source=source) diff --git a/examples/graphgpo/alfworld_env.py b/examples/graphgpo/alfworld_env.py new file mode 100644 index 000000000..80bfc2f69 --- /dev/null +++ b/examples/graphgpo/alfworld_env.py @@ -0,0 +1,411 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Single-environment ALFWorld text adapter for the GraphGPO recipe. + +The ALFWorld dependency is imported only when a real environment is created. +Tests and downstream recipes can inject an already constructed environment or +an environment factory without installing ALFWorld. +""" + +from __future__ import annotations + +import copy +import os +import re +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from examples.graphgpo.state import TrackerState + + +_ITEM_STATE_ACTION = re.compile(r"^(heat|cool|clean|slice)\s+([\w\s\d]+?)(?:\s+with\s+[\w\s\d]+)?$") +_GO_TO_ACTION = re.compile(r"^go to\s+(.+)$") +_TAKE_ACTION = re.compile(r"^take\s+(.+?)(?:\s+from\s+.+)?$") +_DROP_ACTION = re.compile(r"^drop\s+(.+)$") +_PUT_ACTION = re.compile(r"^(put|place)\s+(.+?)\s+(?:in|on)\s+.+$") +_MOVE_ACTION = re.compile(r"^move\s+(.+?)(?:\s+to\s+.+)?$") + + +@dataclass(frozen=True) +class AlfWorldSnapshot: + """The observable state after one ALFWorld reset or step.""" + + raw_observation: str + admissible_commands: tuple[str, ...] + won: bool + done: bool + gamefile: str | None + tracker: TrackerState + + +EnvironmentFactory = Callable[ + [str | Path, int, str, str | Path | None], + Any, +] + + +def _canonicalize_game_files(base_env: Any) -> None: + """Make ALFWorld's seed-to-task mapping independent of ``os.walk`` + order.""" + + game_files = getattr(base_env, "game_files", None) + if isinstance(game_files, (str, bytes)) or not isinstance(game_files, Sequence): + raise TypeError("ALFWorld base environment must expose game_files") + + sortable: list[tuple[str, Any]] = [] + for game_file in game_files: + try: + normalized = os.fspath(game_file).replace("\\", "/") + except TypeError as exc: + raise TypeError("every ALFWorld game file must be path-like") from exc + sortable.append((normalized, game_file)) + + normalized_paths = [normalized for normalized, _ in sortable] + if len(set(normalized_paths)) != len(normalized_paths): + raise ValueError("ALFWorld game_files contains duplicate normalized paths") + + base_env.game_files = [game_file for _, game_file in sorted(sortable, key=lambda item: item[0])] + if hasattr(base_env, "num_games"): + base_env.num_games = len(base_env.game_files) + + +def _select_game_file( + base_env: Any, + selected_game_file: str | Path, + *, + data_root: str | Path | None, +) -> None: + selected = os.fspath(selected_game_file).replace("\\", "/") + if not selected or ".." in Path(selected).parts: + raise ValueError("selected ALFWorld game file must be a safe path") + + root = None + if data_root is not None: + root = os.fspath(data_root).replace("\\", "/").rstrip("/") + if not root: + raise ValueError("ALFWORLD_DATA must not be empty") + + matches: list[Any] = [] + for game_file in base_env.game_files: + normalized = os.fspath(game_file).replace("\\", "/") + relative = None + if root is not None and normalized.startswith(f"{root}/"): + relative = normalized[len(root) + 1 :] + if normalized == selected or relative == selected: + matches.append(game_file) + if len(matches) != 1: + raise ValueError(f"selected ALFWorld game file must match exactly one discovered game; matched {len(matches)}") + base_env.game_files = matches + if hasattr(base_env, "num_games"): + base_env.num_games = 1 + + +def _default_environment_factory( + config_path: str | Path, + seed: int, + train_eval: str, + game_file: str | Path | None, +) -> Any: + try: + import yaml + from alfworld.agents.environment import get_environment + except ImportError as exc: + raise RuntimeError( + "ALFWorld is required for a real GraphGPO rollout; inject env= or env_factory= for dependency-free tests." + ) from exc + + path = Path(config_path) + if not path.is_file(): + raise FileNotFoundError(f"ALFWorld config does not exist: {path}") + with path.open(encoding="utf-8") as reader: + config = yaml.safe_load(reader) + if not isinstance(config, Mapping): + raise ValueError("ALFWorld config must contain a mapping") + try: + env_type = config["env"]["type"] + except (KeyError, TypeError) as exc: + raise ValueError("ALFWorld config must define env.type") from exc + + base_env = get_environment(env_type)(config, train_eval=train_eval) + _canonicalize_game_files(base_env) + if game_file is not None: + _select_game_file( + base_env, + game_file, + data_root=os.environ.get("ALFWORLD_DATA"), + ) + env = base_env.init_env(batch_size=1) + seed_method = getattr(env, "seed", None) + if not callable(seed_method): + raise TypeError("ALFWorld text environment must provide seed(seed)") + seed_method(seed) + return env + + +def _single_batch_item(value: Any, *, field: str) -> Any: + if isinstance(value, (str, bytes, bytearray, Mapping)): + raise TypeError(f"{field} must be a one-element batch") + if not isinstance(value, Sequence) and not (hasattr(value, "__len__") and hasattr(value, "__getitem__")): + raise TypeError(f"{field} must be a one-element batch") + if len(value) != 1: + raise ValueError(f"{field} must contain exactly one environment value") + return value[0] + + +def _normalize_info(infos: Any) -> dict[str, Any]: + if isinstance(infos, Mapping): + return {str(key): _single_batch_item(value, field=f"infos[{key!r}]") for key, value in infos.items()} + info = _single_batch_item(infos, field="infos") + if not isinstance(info, Mapping): + raise TypeError("the single environment info must be a mapping") + return dict(info) + + +def _normalize_commands(value: Any) -> tuple[str, ...]: + if isinstance(value, (str, bytes)) or not isinstance(value, Sequence): + raise TypeError("info['admissible_commands'] must be a sequence of strings") + commands: list[str] = [] + for command in value: + if not isinstance(command, str): + raise TypeError("every admissible command must be a string") + commands.append(command) + return tuple(commands) + + +class AlfWorldTextEnv: + """Adapt ALFWorld's list-shaped text API to one explicit environment.""" + + def __init__( + self, + *, + env: Any | None = None, + config_path: str | Path | None = None, + seed: int = 0, + train_eval: str = "train", + game_file: str | Path | None = None, + env_factory: EnvironmentFactory | None = None, + ) -> None: + if env is not None and (config_path is not None or game_file is not None or env_factory is not None): + raise ValueError("env cannot be combined with config_path, game_file, or env_factory") + if isinstance(seed, bool) or not isinstance(seed, int) or seed < 0: + raise ValueError("seed must be a non-negative integer") + if not isinstance(train_eval, str) or not train_eval: + raise ValueError("train_eval must be a non-empty string") + if env is None and config_path is None: + raise ValueError("config_path is required when env is not injected") + + self._env = env + self._config_path = config_path + self._seed = seed + self._train_eval = train_eval + self._game_file = game_file + self._env_factory = env_factory or _default_environment_factory + self._closed = False + self._done = False + self._location = "middle of a room" + self._holding = "nothing" + self._history_items: dict[str, dict[str, bool]] = {} + self._item_location: dict[str, dict[str, str]] = {} + self._history_list: list[tuple[str, ...]] = [] + + def _ensure_env(self) -> Any: + if self._closed: + raise RuntimeError("ALFWorld environment is already closed") + if self._env is None: + if self._config_path is None: + raise RuntimeError("ALFWorld config path is unavailable") + self._env = self._env_factory( + self._config_path, + self._seed, + self._train_eval, + self._game_file, + ) + return self._env + + def _tracker(self) -> TrackerState: + return TrackerState.from_mapping( + { + "location": self._location, + "holding": self._holding, + "history_items": copy.deepcopy(self._history_items), + "item_location": copy.deepcopy(self._item_location), + } + ) + + def _snapshot( + self, + *, + raw_observation: Any, + info: Mapping[str, Any], + done: bool, + ) -> AlfWorldSnapshot: + if not isinstance(raw_observation, str): + raise TypeError("ALFWorld observation must be a string") + if "admissible_commands" not in info: + raise KeyError("ALFWorld info is missing 'admissible_commands'") + gamefile = info.get("extra.gamefile") + if gamefile is not None and not isinstance(gamefile, str): + raise TypeError("info['extra.gamefile'] must be a string or null") + return AlfWorldSnapshot( + raw_observation=raw_observation, + admissible_commands=_normalize_commands(info["admissible_commands"]), + won=bool(info.get("won", False)), + done=done, + gamefile=gamefile, + tracker=self._tracker(), + ) + + def reset(self) -> AlfWorldSnapshot: + env = self._ensure_env() + result = env.reset() + if not isinstance(result, tuple) or len(result) != 2: + raise TypeError("ALFWorld reset() must return (observations, infos)") + observations, infos = result + raw_observation = _single_batch_item(observations, field="observations") + info = _normalize_info(infos) + + self._history_items = {} + self._location = "middle of a room" + self._holding = "nothing" + self._item_location = {} + self._history_list = [] + self._done = False + return self._snapshot( + raw_observation=raw_observation, + info=info, + done=False, + ) + + def _update_item_location(self, old_holding: str) -> None: + obj = old_holding if old_holding != "nothing" else self._holding + if obj == "nothing": + raise AssertionError("a holding transition must identify an object") + if obj not in self._item_location: + if self._holding != obj: + raise AssertionError("a new item location requires holding that item") + self._item_location[obj] = { + "old_location": self._location, + "new_location": self._location, + } + elif old_holding == obj: + self._item_location[obj]["new_location"] = self._location + elif self._holding == obj: + self._item_location[obj]["new_location"] = self._item_location[obj]["old_location"] + + def _update_item_state(self, action: str, observation: str) -> None: + match = _ITEM_STATE_ACTION.match(action.lower().strip()) + if match is None: + return + verb, obj = match.groups() + obj = obj.strip() + if verb not in observation: + return + if obj not in self._history_items: + # ``slice`` versus ``sliced`` below intentionally preserves the + # frozen AlfworldWorker tracker schema used by the reference run. + self._history_items[obj] = { + "heated": False, + "cooled": False, + "cleaned": False, + "slice": False, + } + if verb == "heat": + self._history_items[obj]["heated"] = True + elif verb == "cool": + self._history_items[obj]["cooled"] = True + elif verb == "clean": + self._history_items[obj]["cleaned"] = True + elif verb == "slice": + self._history_items[obj]["sliced"] = True + + def _update_position(self, action: str, observation: str) -> None: + match = _GO_TO_ACTION.match(action.lower().strip()) + if match is None: + return + target_location = match.group(1).strip() + if target_location not in observation: + raise AssertionError(f"target location {target_location!r} is absent from observation") + self._location = target_location + + def _update_held_items(self, action: str, observation: str) -> None: + action = action.lower().strip() + old_holding = self._holding + + match_take = _TAKE_ACTION.match(action) + if match_take is not None: + obj = match_take.group(1).strip() + if self._holding != "nothing": + raise AssertionError("take action requires empty hands") + if obj in observation: + self._holding = obj + + match_drop = _DROP_ACTION.match(action) + if match_drop is not None: + obj = match_drop.group(1).strip() + if obj in observation: + self._holding = "nothing" + + match_put = _PUT_ACTION.match(action) + if match_put is not None: + obj = match_put.group(2).strip() + if obj in observation: + self._holding = "nothing" + + match_move = _MOVE_ACTION.match(action) + if match_move is not None: + obj = match_move.group(1).strip() + if self._holding != obj: + raise AssertionError("move action requires holding the moved object") + if obj in observation: + self._holding = "nothing" + + if old_holding != self._holding: + self._update_item_location(old_holding) + + def step(self, action: str) -> AlfWorldSnapshot: + if not isinstance(action, str): + raise TypeError("action must be a string") + if self._done: + raise RuntimeError("cannot step a completed ALFWorld environment") + env = self._ensure_env() + result = env.step([action]) + if not isinstance(result, tuple) or len(result) != 4: + raise TypeError("ALFWorld step() must return (observations, scores, dones, infos)") + observations, scores, dones, infos = result + raw_observation = _single_batch_item(observations, field="observations") + _single_batch_item(scores, field="scores") + done = bool(_single_batch_item(dones, field="dones")) + info = _normalize_info(infos) + + self._history_list.append( + ( + f"action:{action}", + f"obs:{observations}", + f"self.location: {self._location}", + f"self.holding: {self._holding}", + f"self.history_items: {self._history_items}", + f"self.item_location: {self._item_location}", + ) + ) + if "Nothing happens" not in raw_observation and not self._done: + self._update_item_state(action, raw_observation) + self._update_position(action, raw_observation) + self._update_held_items(action, raw_observation) + self._done = done + return self._snapshot( + raw_observation=raw_observation, + info=info, + done=done, + ) + + def close(self) -> None: + if self._closed: + return + self._closed = True + env = self._env + self._env = None + close_method = getattr(env, "close", None) + if callable(close_method): + close_method() diff --git a/examples/graphgpo/configs/alfworld_qwen2_5_1_5b.yaml b/examples/graphgpo/configs/alfworld_qwen2_5_1_5b.yaml new file mode 100644 index 000000000..74aab08ae --- /dev/null +++ b/examples/graphgpo/configs/alfworld_qwen2_5_1_5b.yaml @@ -0,0 +1,30 @@ +# GraphGPO ALFWorld text-only environment configuration. +# Paths expand from ALFWORLD_DATA at runtime. + +dataset: + data_path: "$ALFWORLD_DATA/json_2.1.1/train" + eval_id_data_path: "$ALFWORLD_DATA/json_2.1.1/valid_seen" + eval_ood_data_path: "$ALFWORLD_DATA/json_2.1.1/valid_unseen" + num_train_games: -1 + num_eval_games: -1 + +logic: + domain: "$ALFWORLD_DATA/logic/alfred.pddl" + grammar: "$ALFWORLD_DATA/logic/alfred.twl2" + +env: + type: "AlfredTWEnv" + domain_randomization: false + task_types: [1, 2, 3, 4, 5, 6] + expert_timeout_steps: 150 + expert_type: "handcoded" + goal_desc_human_anns_prob: 0.0 + +general: + random_seed: 42 + use_cuda: false + training_method: "dqn" + +rl: + training: + max_nb_steps_per_episode: 50 diff --git a/examples/graphgpo/custom_advantage.py b/examples/graphgpo/custom_advantage.py new file mode 100644 index 000000000..ad43783c8 --- /dev/null +++ b/examples/graphgpo/custom_advantage.py @@ -0,0 +1,394 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Custom-advantage adapter for the GraphGPO reproduction recipe. + +Relax calls :func:`compute_custom_advantage` with the value returned by +``RewardWaitingGroup.metadata_by_slot()``. Each slot is one trajectory and +each named unit is one environment turn. This module validates that contract, +delegates the numerical work to :mod:`examples.graphgpo.graph_credit`, and +returns one scalar advantage for every input unit. +""" + +from __future__ import annotations + +import math +import os +import re +from collections.abc import Hashable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any, Literal + +from examples.graphgpo.diagnostics import diagnostics_callback_from_environment +from examples.graphgpo.graph_credit import ( + DEFAULT_EPISODE_WEIGHTING, + EPISODE_WEIGHTINGS, + EpisodeWeighting, + Turn, + compute_method_advantages, + episode_return, +) + + +Method = Literal["grpo", "gigpo", "graphgpo"] +_METHODS = frozenset(("grpo", "gigpo", "graphgpo")) +_TURN_NAME = re.compile(r"^turn_(\d{3,})$") +_REQUIRED_FIELDS = ( + "row_id", + "rollout_group_id", + "policy_version", + "task_id", + "trajectory_id", + "turn_index", + "state_key", + "action", + "next_state_key", + "is_action_valid", + "success", + "terminal", + "truncated", + "episode_return", +) + + +@dataclass(frozen=True) +class _UnitRef: + slot_index: int + unit_name: str + row_id: Hashable + rollout_group_id: Hashable + policy_version: Hashable + turn: Turn + declared_episode_return: float + + +def _require_positive_int(name: str, value: object) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + return value + + +def _require_non_negative_finite(name: str, value: object) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{name} must be a finite non-negative real number") + result = float(value) + if not math.isfinite(result) or result < 0.0: + raise ValueError(f"{name} must be a finite non-negative real number") + return result + + +def _require_finite(name: str, value: object) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{name} must be a finite real number") + result = float(value) + if not math.isfinite(result): + raise ValueError(f"{name} must be a finite real number") + return result + + +def _require_hashable(name: str, value: object) -> Hashable: + if isinstance(value, bool): + raise ValueError(f"{name} must be a hashable non-boolean value") + try: + hash(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be hashable") from exc + return value # type: ignore[return-value] + + +def _require_bool(name: str, value: object) -> bool: + if not isinstance(value, bool): + raise ValueError(f"{name} must be a boolean") + return value + + +def _required(metadata: Mapping[str, Any], field: str, context: str) -> Any: + if field not in metadata: + raise ValueError(f"{context} is missing required metadata field {field!r}") + return metadata[field] + + +def _parse_unit(slot_index: int, unit_name: object, metadata: object) -> _UnitRef: + if not isinstance(unit_name, str): + raise ValueError(f"slot {slot_index} contains a non-string unit name") + context = f"slot {slot_index} unit {unit_name!r}" + if not isinstance(metadata, Mapping): + raise ValueError(f"{context} metadata must be a mapping") + for field in _REQUIRED_FIELDS: + _required(metadata, field, context) + + turn_index = _required(metadata, "turn_index", context) + if isinstance(turn_index, bool) or not isinstance(turn_index, int) or turn_index < 0: + raise ValueError(f"{context} turn_index must be a non-negative integer") + match = _TURN_NAME.fullmatch(unit_name) + if match is None or int(match.group(1)) != turn_index: + raise ValueError(f"{context} must use the canonical name {f'turn_{turn_index:03d}'!r}") + + action = _required(metadata, "action", context) + if not isinstance(action, str): + raise ValueError(f"{context} action must be a string") + cost = _require_finite(f"{context} cost", metadata.get("cost", 1.0)) + if cost <= 0.0: + raise ValueError(f"{context} cost must be greater than zero") + + turn = Turn( + task_id=_require_hashable(f"{context} task_id", _required(metadata, "task_id", context)), + trajectory_id=_require_hashable( + f"{context} trajectory_id", + _required(metadata, "trajectory_id", context), + ), + turn_index=turn_index, + state_key=_require_hashable(f"{context} state_key", _required(metadata, "state_key", context)), + action=action, + next_state_key=_require_hashable( + f"{context} next_state_key", + _required(metadata, "next_state_key", context), + ), + success=_require_bool(f"{context} success", _required(metadata, "success", context)), + terminal=_require_bool(f"{context} terminal", _required(metadata, "terminal", context)), + truncated=_require_bool(f"{context} truncated", _required(metadata, "truncated", context)), + is_action_valid=_require_bool( + f"{context} is_action_valid", + _required(metadata, "is_action_valid", context), + ), + cost=cost, + ) + declared_episode_return = _require_finite( + f"{context} episode_return", + _required(metadata, "episode_return", context), + ) + return _UnitRef( + slot_index=slot_index, + unit_name=unit_name, + row_id=_require_hashable(f"{context} row_id", _required(metadata, "row_id", context)), + rollout_group_id=_require_hashable( + f"{context} rollout_group_id", + _required(metadata, "rollout_group_id", context), + ), + policy_version=_require_hashable( + f"{context} policy_version", + _required(metadata, "policy_version", context), + ), + turn=turn, + declared_episode_return=declared_episode_return, + ) + + +def _parse_group( + metadata_by_slot: object, + *, + expected_group_size: int, + success_reward: float, + invalid_penalty: float, +) -> tuple[tuple[_UnitRef, ...], ...]: + if isinstance(metadata_by_slot, (str, bytes)) or not isinstance(metadata_by_slot, Sequence): + raise ValueError("metadata_by_slot must be a sequence of slot mappings") + if len(metadata_by_slot) != expected_group_size: + raise ValueError(f"rollout group size mismatch: expected {expected_group_size}, got {len(metadata_by_slot)}") + + parsed_slots: list[tuple[_UnitRef, ...]] = [] + seen_trajectory_ids: set[Hashable] = set() + seen_row_ids: set[Hashable] = set() + unset_task_id = object() + task_id: Hashable | object = unset_task_id + unset_rollout_group_id = object() + rollout_group_id: Hashable | object = unset_rollout_group_id + unset_policy_version = object() + policy_version: Hashable | object = unset_policy_version + for slot_index, slot in enumerate(metadata_by_slot): + if not isinstance(slot, Mapping): + raise ValueError(f"slot {slot_index} must be a unit metadata mapping") + if not slot: + raise ValueError(f"slot {slot_index} must contain at least one turn unit") + + refs = tuple(_parse_unit(slot_index, unit_name, metadata) for unit_name, metadata in slot.items()) + for ref in refs: + if ref.row_id in seen_row_ids: + raise ValueError(f"duplicate row_id in rollout group: {ref.row_id!r}") + seen_row_ids.add(ref.row_id) + if rollout_group_id is unset_rollout_group_id: + rollout_group_id = ref.rollout_group_id + elif ref.rollout_group_id != rollout_group_id: + raise ValueError("one custom-advantage call cannot mix multiple rollout_group_id values") + if policy_version is unset_policy_version: + policy_version = ref.policy_version + elif ref.policy_version != policy_version: + raise ValueError("one custom-advantage call cannot mix multiple policy_version values") + + trajectory_ids = {ref.turn.trajectory_id for ref in refs} + if len(trajectory_ids) != 1: + raise ValueError(f"slot {slot_index} mixes multiple trajectory_id values") + trajectory_id = next(iter(trajectory_ids)) + if trajectory_id in seen_trajectory_ids: + raise ValueError(f"trajectory_id {trajectory_id!r} appears in more than one slot") + seen_trajectory_ids.add(trajectory_id) + + task_ids = {ref.turn.task_id for ref in refs} + if len(task_ids) != 1: + raise ValueError(f"slot {slot_index} mixes multiple task_id values") + slot_task_id = next(iter(task_ids)) + if task_id is unset_task_id: + task_id = slot_task_id + elif slot_task_id != task_id: + raise ValueError("one rollout group cannot mix multiple task_id values") + + refs_by_turn = tuple(sorted(refs, key=lambda ref: ref.turn.turn_index)) + actual_indices = [ref.turn.turn_index for ref in refs_by_turn] + expected_indices = list(range(len(refs_by_turn))) + if actual_indices != expected_indices: + raise ValueError(f"slot {slot_index} has non-contiguous turn indices: {actual_indices!r}") + expected_names = [f"turn_{turn_index:03d}" for turn_index in expected_indices] + actual_names = [ref.unit_name for ref in refs_by_turn] + if actual_names != expected_names: + raise ValueError(f"slot {slot_index} turn unit names do not match their indices") + + declared_returns = {ref.declared_episode_return for ref in refs_by_turn} + if len(declared_returns) != 1: + raise ValueError(f"slot {slot_index} has inconsistent episode_return values") + computed_return = episode_return( + [ref.turn for ref in refs_by_turn], + success_reward=success_reward, + invalid_penalty=invalid_penalty, + ) + declared_return = next(iter(declared_returns)) + if not math.isclose(declared_return, computed_return, rel_tol=1e-9, abs_tol=1e-9): + raise ValueError( + f"slot {slot_index} episode_return mismatch: declared {declared_return}, computed {computed_return}" + ) + parsed_slots.append(refs_by_turn) + + return tuple(parsed_slots) + + +def compute_group_advantages( + metadata_by_slot: Sequence[Mapping[str, Mapping[str, Any]]], + *, + method: Method | str = "graphgpo", + expected_group_size: int = 8, + omega: float = 0.1, + gamma: float = 0.95, + beta: float = 1.0, + beta_episode: float = 1.0, + success_reward: float = 10.0, + invalid_penalty: float = 0.1, + eps: float = 1e-6, + episode_weighting: EpisodeWeighting | str = DEFAULT_EPISODE_WEIGHTING, +) -> list[dict[str, float]]: + """Validate one complete rollout group and return per-unit advantages. + + ``beta`` weights the method-specific graph or step term. ``beta_episode`` + weights the episode-level GRPO term. For ``method="grpo"`` only + ``beta_episode`` is used. + + Invalid input raises a descriptive exception. This function never returns + ``None`` because the recipe does not silently drop malformed groups. + """ + + if not isinstance(method, str) or method not in _METHODS: + raise ValueError(f"method must be one of {sorted(_METHODS)!r}, got {method!r}") + if not isinstance(episode_weighting, str) or episode_weighting not in EPISODE_WEIGHTINGS: + raise ValueError(f"episode_weighting must be one of {sorted(EPISODE_WEIGHTINGS)!r}, got {episode_weighting!r}") + expected_group_size = _require_positive_int("expected_group_size", expected_group_size) + omega = _require_finite("omega", omega) + gamma = _require_finite("gamma", gamma) + beta = _require_finite("beta", beta) + beta_episode = _require_finite("beta_episode", beta_episode) + success_reward = _require_non_negative_finite("success_reward", success_reward) + invalid_penalty = _require_non_negative_finite("invalid_penalty", invalid_penalty) + eps = _require_non_negative_finite("eps", eps) + + parsed_slots = _parse_group( + metadata_by_slot, + expected_group_size=expected_group_size, + success_reward=success_reward, + invalid_penalty=invalid_penalty, + ) + turns = [ref.turn for refs in parsed_slots for ref in refs] + advantages = compute_method_advantages( + method, + turns, + expected_group_size=expected_group_size, + success_reward=success_reward, + invalid_penalty=invalid_penalty, + omega=omega, + gamma=gamma, + beta_episode=beta_episode, + beta_graph=beta, + beta_step=beta, + eps=eps, + episode_weighting=episode_weighting, + graph_diagnostics_callback=diagnostics_callback_from_environment(), + diagnostic_rollout_group_id=parsed_slots[0][0].rollout_group_id, + diagnostic_policy_version=parsed_slots[0][0].policy_version, + ) + + result: list[dict[str, float]] = [] + for refs in parsed_slots: + slot_result: dict[str, float] = {} + for ref in refs: + value = _require_finite( + f"advantage for slot {ref.slot_index} unit {ref.unit_name!r}", + advantages[ref.turn.key], + ) + slot_result[ref.unit_name] = value + result.append(slot_result) + + expected_shape = [set(slot) for slot in metadata_by_slot] + actual_shape = [set(slot) for slot in result] + if actual_shape != expected_shape: + raise ValueError("custom advantage output shape does not match its input") + return result + + +def _environment_value(primary: str, fallback: str | None, default: str) -> str: + if primary in os.environ: + return os.environ[primary] + if fallback is not None and fallback in os.environ: + return os.environ[fallback] + return default + + +def _environment_int(primary: str, fallback: str | None, default: int) -> int: + raw = _environment_value(primary, fallback, str(default)) + try: + value = int(raw) + except ValueError as exc: + raise ValueError(f"{primary} must be an integer, got {raw!r}") from exc + return _require_positive_int(primary, value) + + +def _environment_float(primary: str, fallback: str | None, default: float) -> float: + raw = _environment_value(primary, fallback, str(default)) + try: + value = float(raw) + except ValueError as exc: + raise ValueError(f"{primary} must be a real number, got {raw!r}") from exc + return _require_finite(primary, value) + + +def compute_custom_advantage( + metadata_by_slot: Sequence[Mapping[str, Mapping[str, Any]]], +) -> list[dict[str, float]]: + """Environment-configured hook loaded by ``--agentic-custom-advantage- + path``. + + Recipe-specific variables take precedence over the short names used by the + launch script. Defaults reproduce the proposal's GraphGPO group-of-eight + configuration, including one vote per trajectory for episode statistics. + """ + + method = _environment_value("GRAPHGPO_METHOD", "METHOD", "graphgpo") + episode_weighting = _environment_value( + "GRAPHGPO_EPISODE_WEIGHTING", + "EPISODE_WEIGHTING", + DEFAULT_EPISODE_WEIGHTING, + ) + return compute_group_advantages( + metadata_by_slot, + method=method, + episode_weighting=episode_weighting, + expected_group_size=_environment_int("GRAPHGPO_EXPECTED_GROUP_SIZE", "GROUP_SIZE", 8), + omega=_environment_float("GRAPHGPO_OMEGA", "OMEGA", 0.1), + gamma=_environment_float("GRAPHGPO_GAMMA", "GAMMA", 0.95), + beta=_environment_float("GRAPHGPO_BETA", "BETA", 1.0), + beta_episode=_environment_float("GRAPHGPO_BETA_EPISODE", "BETA_EPISODE", 1.0), + ) diff --git a/examples/graphgpo/diagnostics.py b/examples/graphgpo/diagnostics.py new file mode 100644 index 000000000..2f2c8fd48 --- /dev/null +++ b/examples/graphgpo/diagnostics.py @@ -0,0 +1,353 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Opt-in JSONL diagnostics for the pure-Python GraphGPO graph kernel.""" + +from __future__ import annotations + +import json +import math +import os +from collections import Counter, defaultdict +from dataclasses import asdict, dataclass +from math import ceil +from pathlib import Path +from statistics import median +from typing import Callable, Hashable, Mapping, Protocol, Sequence + + +DIAGNOSTICS_PATH_ENV = "GRAPHGPO_DIAGNOSTICS_JSONL" +DIAGNOSTICS_SCHEMA = "graphgpo-diagnostics-v1" +DIAGNOSTICS_SUMMARY_SCHEMA = "graphgpo-diagnostics-summary-v1" +TurnKey = tuple[Hashable, int] + + +class EdgeOccurrenceLike(Protocol): + """Structural subset used to keep this module independent of + graph_credit.""" + + turn_key: TurnKey + source: Hashable + action: str + target: Hashable + + +@dataclass(frozen=True) +class GraphDiagnostics: + """One task-local graph-credit diagnostic record.""" + + schema: str + task_id: str + rollout_group_id: str | None + policy_version: str | None + graph_build_ns: int + reverse_dijkstra_ns: int + graph_advantage_ns: int + graph_credit_total_ns: int + node_count: int + edge_occurrence_count: int + duplicate_edge_count: int + duplicate_edge_rate: float + shared_state_count: int + shared_state_rate: float + singleton_source_count: int + singleton_source_rate: float + nonzero_graph_advantage_count: int + nonzero_graph_advantage_rate: float + all_fail: bool + unreachable_node_count: int + unreachable_node_rate: float + max_finite_distance: float + distance_histogram: Mapping[str, int] + + def to_mapping(self) -> dict[str, object]: + return asdict(self) + + +DiagnosticsCallback = Callable[[GraphDiagnostics], None] + + +@dataclass(frozen=True) +class GraphDiagnosticsSummary: + """Aggregate median/p95 evidence derived from raw diagnostic records.""" + + schema: str + record_count: int + timing_ns: Mapping[str, Mapping[str, float | int]] + graph_metrics: Mapping[str, Mapping[str, float]] + all_fail_count: int + all_fail_rate: float + distance_histogram: Mapping[str, int] + + def to_mapping(self) -> dict[str, object]: + return asdict(self) + + +def _stable_identifier(value: object) -> str: + value_type = type(value) + return f"{value_type.__module__}.{value_type.__qualname__}:{value!r}" + + +def _rate(numerator: int, denominator: int) -> float: + if denominator == 0: + return 0.0 + return numerator / denominator + + +def _require_duration(name: str, value: int) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"{name} must be a non-negative integer") + return value + + +def _distance_label(value: float) -> str: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError("distances must be finite real numbers") + finite_value = float(value) + if not math.isfinite(finite_value) or finite_value < 0.0: + raise ValueError("distances must be finite non-negative real numbers") + return format(finite_value, ".17g") + + +def build_graph_diagnostics( + *, + task_id: object, + nodes: Sequence[Hashable] | frozenset[Hashable], + occurrences: Sequence[EdgeOccurrenceLike], + distances: Mapping[Hashable, float], + max_finite_distance: float, + has_success: bool, + advantages: Mapping[TurnKey, float], + graph_build_ns: int, + reverse_dijkstra_ns: int, + graph_advantage_ns: int, + rollout_group_id: object | None = None, + policy_version: object | None = None, +) -> GraphDiagnostics: + """Derive deterministic graph metrics from one completed task-local + graph.""" + + graph_build_ns = _require_duration("graph_build_ns", graph_build_ns) + reverse_dijkstra_ns = _require_duration("reverse_dijkstra_ns", reverse_dijkstra_ns) + graph_advantage_ns = _require_duration("graph_advantage_ns", graph_advantage_ns) + if not isinstance(has_success, bool): + raise ValueError("has_success must be a boolean") + + node_set = frozenset(nodes) + if not node_set or set(distances) != set(node_set): + raise ValueError("distances must contain exactly one value per graph node") + + occurrence_keys = {occurrence.turn_key for occurrence in occurrences} + if len(occurrence_keys) != len(occurrences): + raise ValueError("edge occurrence turn keys must be unique") + if set(advantages) != occurrence_keys: + raise ValueError("advantages must contain exactly one value per edge occurrence") + + edge_counts = Counter((occurrence.source, occurrence.action, occurrence.target) for occurrence in occurrences) + duplicate_edge_count = sum(count - 1 for count in edge_counts.values()) + + trajectories_by_source: dict[Hashable, set[Hashable]] = defaultdict(set) + occurrences_by_source: Counter[Hashable] = Counter() + for occurrence in occurrences: + trajectories_by_source[occurrence.source].add(occurrence.turn_key[0]) + occurrences_by_source[occurrence.source] += 1 + shared_state_count = sum(len(trajectory_ids) > 1 for trajectory_ids in trajectories_by_source.values()) + singleton_source_count = sum(count == 1 for count in occurrences_by_source.values()) + + finite_advantages: list[float] = [] + for value in advantages.values(): + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError("graph advantages must be finite real numbers") + finite_value = float(value) + if not math.isfinite(finite_value): + raise ValueError("graph advantages must be finite real numbers") + finite_advantages.append(finite_value) + nonzero_count = sum(value != 0.0 for value in finite_advantages) + + _distance_label(max_finite_distance) + distance_values = {node: float(value) for node, value in distances.items()} + if has_success: + unreachable_count = sum(value > float(max_finite_distance) for value in distance_values.values()) + else: + unreachable_count = len(node_set) + histogram = Counter(_distance_label(value) for value in distance_values.values()) + + edge_count = len(occurrences) + source_count = len(occurrences_by_source) + node_count = len(node_set) + return GraphDiagnostics( + schema=DIAGNOSTICS_SCHEMA, + task_id=_stable_identifier(task_id), + rollout_group_id=(None if rollout_group_id is None else _stable_identifier(rollout_group_id)), + policy_version=(None if policy_version is None else _stable_identifier(policy_version)), + graph_build_ns=graph_build_ns, + reverse_dijkstra_ns=reverse_dijkstra_ns, + graph_advantage_ns=graph_advantage_ns, + graph_credit_total_ns=(graph_build_ns + reverse_dijkstra_ns + graph_advantage_ns), + node_count=node_count, + edge_occurrence_count=edge_count, + duplicate_edge_count=duplicate_edge_count, + duplicate_edge_rate=_rate(duplicate_edge_count, edge_count), + shared_state_count=shared_state_count, + shared_state_rate=_rate(shared_state_count, source_count), + singleton_source_count=singleton_source_count, + singleton_source_rate=_rate(singleton_source_count, source_count), + nonzero_graph_advantage_count=nonzero_count, + nonzero_graph_advantage_rate=_rate(nonzero_count, edge_count), + all_fail=not has_success, + unreachable_node_count=unreachable_count, + unreachable_node_rate=_rate(unreachable_count, node_count), + max_finite_distance=float(max_finite_distance), + distance_histogram=dict(sorted(histogram.items())), + ) + + +class JsonlDiagnosticsWriter: + """Append one compact JSON object per graph when explicitly configured.""" + + def __init__(self, path: str | os.PathLike[str]) -> None: + self._path = Path(path) + + @property + def path(self) -> Path: + return self._path + + def __call__(self, record: GraphDiagnostics) -> None: + if not isinstance(record, GraphDiagnostics): + raise TypeError("record must be a GraphDiagnostics instance") + self._path.parent.mkdir(parents=True, exist_ok=True) + payload = json.dumps( + record.to_mapping(), + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + with self._path.open("a", encoding="utf-8", newline="\n") as stream: + stream.write(payload + "\n") + + +def diagnostics_callback_from_environment() -> DiagnosticsCallback | None: + """Return a writer only when ``GRAPHGPO_DIAGNOSTICS_JSONL`` is non- + empty.""" + + path = os.environ.get(DIAGNOSTICS_PATH_ENV, "").strip() + if not path: + return None + return JsonlDiagnosticsWriter(path) + + +def _nearest_rank_p95(values: Sequence[float]) -> float: + if not values: + raise ValueError("at least one value is required") + ordered = sorted(values) + return ordered[ceil(0.95 * len(ordered)) - 1] + + +def _distribution(values: Sequence[float]) -> dict[str, float]: + if not values: + raise ValueError("at least one value is required") + return { + "median": float(median(values)), + "p95": float(_nearest_rank_p95(values)), + } + + +def summarize_graph_diagnostics( + records: Sequence[GraphDiagnostics], +) -> GraphDiagnosticsSummary: + """Summarize raw task-local records with median and nearest-rank p95.""" + + if not records: + raise ValueError("at least one graph diagnostic record is required") + if any(not isinstance(record, GraphDiagnostics) for record in records): + raise TypeError("all records must be GraphDiagnostics instances") + + timing_fields = ( + "graph_build_ns", + "reverse_dijkstra_ns", + "graph_advantage_ns", + "graph_credit_total_ns", + ) + timing_ns: dict[str, dict[str, float | int]] = {} + for field in timing_fields: + values = [getattr(record, field) for record in records] + timing_ns[field] = { + **_distribution(values), + "total": sum(values), + } + + metric_fields = ( + "node_count", + "edge_occurrence_count", + "duplicate_edge_rate", + "shared_state_rate", + "singleton_source_rate", + "nonzero_graph_advantage_rate", + "unreachable_node_rate", + ) + graph_metrics = { + field: _distribution([float(getattr(record, field)) for record in records]) for field in metric_fields + } + + all_fail_count = sum(record.all_fail for record in records) + distance_histogram: Counter[str] = Counter() + for record in records: + distance_histogram.update(record.distance_histogram) + return GraphDiagnosticsSummary( + schema=DIAGNOSTICS_SUMMARY_SCHEMA, + record_count=len(records), + timing_ns=timing_ns, + graph_metrics=graph_metrics, + all_fail_count=all_fail_count, + all_fail_rate=_rate(all_fail_count, len(records)), + distance_histogram=dict(sorted(distance_histogram.items(), key=lambda item: float(item[0]))), + ) + + +def load_graph_diagnostics_jsonl( + path: str | os.PathLike[str], +) -> tuple[GraphDiagnostics, ...]: + """Load and schema-check raw diagnostic JSONL records.""" + + records: list[GraphDiagnostics] = [] + with Path(path).open("r", encoding="utf-8") as stream: + for line_number, line in enumerate(stream, start=1): + if not line.strip(): + continue + try: + payload = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError(f"invalid diagnostics JSON on line {line_number}") from exc + if not isinstance(payload, dict) or payload.get("schema") != DIAGNOSTICS_SCHEMA: + raise ValueError(f"unexpected diagnostics schema on line {line_number}") + try: + records.append(GraphDiagnostics(**payload)) + except TypeError as exc: + raise ValueError(f"invalid diagnostics record on line {line_number}") from exc + if not records: + raise ValueError("diagnostics JSONL contains no records") + return tuple(records) + + +def write_graph_diagnostics_summary( + input_path: str | os.PathLike[str], + output_path: str | os.PathLike[str], +) -> GraphDiagnosticsSummary: + """Write an immutable JSON summary next to retained raw JSONL evidence.""" + + source = Path(input_path) + destination = Path(output_path) + if source.resolve() == destination.resolve(): + raise ValueError("diagnostics summary output must differ from its input") + summary = summarize_graph_diagnostics(load_graph_diagnostics_jsonl(source)) + destination.parent.mkdir(parents=True, exist_ok=True) + payload = json.dumps( + summary.to_mapping(), + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + with destination.open("x", encoding="utf-8", newline="\n") as stream: + stream.write(payload + "\n") + return summary diff --git a/examples/graphgpo/diagnostics_summary.py b/examples/graphgpo/diagnostics_summary.py new file mode 100644 index 000000000..ac0f5211b --- /dev/null +++ b/examples/graphgpo/diagnostics_summary.py @@ -0,0 +1,28 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Command-line entry point for immutable GraphGPO diagnostic summaries.""" + +from __future__ import annotations + +import argparse +from pathlib import Path +from typing import Sequence + +from examples.graphgpo.diagnostics import write_graph_diagnostics_summary + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> int: + args = parse_args(argv) + write_graph_diagnostics_summary(args.input, args.output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/graphgpo/eval_logger.py b/examples/graphgpo/eval_logger.py new file mode 100644 index 000000000..b68f61b27 --- /dev/null +++ b/examples/graphgpo/eval_logger.py @@ -0,0 +1,133 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Episode-level evaluation metrics for variable-row agentic rollouts.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping, Sequence +from typing import Any, Hashable + + +def _sample_metadata(sample: object, *, context: str) -> Mapping[str, Any]: + metadata = getattr(sample, "metadata", None) + if not isinstance(metadata, Mapping): + raise ValueError(f"{context} metadata must be a mapping") + return metadata + + +def _required(metadata: Mapping[str, Any], key: str, *, context: str) -> Any: + if key not in metadata: + raise ValueError(f"{context} metadata is missing {key!r}") + return metadata[key] + + +def _trajectory_id(metadata: Mapping[str, Any], *, context: str) -> Hashable: + value = _required(metadata, "trajectory_id", context=context) + try: + hash(value) + except TypeError as exc: + raise ValueError(f"{context} trajectory_id must be hashable") from exc + return value + + +def _bool(metadata: Mapping[str, Any], key: str, *, context: str) -> bool: + value = _required(metadata, key, context=context) + if not isinstance(value, bool): + raise ValueError(f"{context} {key} must be a boolean") + return value + + +def _finite_return(metadata: Mapping[str, Any], *, context: str) -> float: + value = _required(metadata, "episode_return", context=context) + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{context} episode_return must be a finite number") + value = float(value) + if not math.isfinite(value): + raise ValueError(f"{context} episode_return must be a finite number") + return value + + +def episode_metrics(samples: Sequence[object]) -> dict[str, float | int]: + """Compute one vote per trajectory, regardless of exported turn count.""" + + if isinstance(samples, (str, bytes)) or not isinstance(samples, Sequence): + raise TypeError("samples must be a sequence") + if not samples: + raise ValueError("episode-level eval metrics require at least one sample") + + by_trajectory: dict[Hashable, list[Mapping[str, Any]]] = {} + for index, sample in enumerate(samples): + context = f"sample {index}" + metadata = _sample_metadata(sample, context=context) + by_trajectory.setdefault(_trajectory_id(metadata, context=context), []).append(metadata) + + successes: list[bool] = [] + episode_returns: list[float] = [] + truncated_episodes: list[bool] = [] + for trajectory_id, rows in by_trajectory.items(): + context = f"trajectory {trajectory_id!r}" + success_values = {_bool(row, "success", context=context) for row in rows} + return_values = {_finite_return(row, context=context) for row in rows} + if len(success_values) != 1: + raise ValueError(f"{context} has inconsistent success values") + if len(return_values) != 1: + raise ValueError(f"{context} has inconsistent episode_return values") + successes.append(next(iter(success_values))) + episode_returns.append(next(iter(return_values))) + truncated_episodes.append(any(_bool(row, "truncated", context=context) for row in rows)) + + episode_count = len(by_trajectory) + return { + "episode_count": episode_count, + "success_rate": sum(successes) / episode_count, + "episode_return_mean": sum(episode_returns) / episode_count, + "truncated_rate": sum(truncated_episodes) / episode_count, + } + + +def build_eval_metrics( + data: Mapping[str, Mapping[str, Any]], + extra_metrics: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Build namespaced episode metrics for every configured eval dataset.""" + + if not isinstance(data, Mapping) or not data: + raise ValueError("eval data must be a non-empty mapping") + metrics: dict[str, Any] = dict(extra_metrics or {}) + for dataset_name, dataset in data.items(): + if not isinstance(dataset_name, str) or not dataset_name: + raise ValueError("eval dataset names must be non-empty strings") + if not isinstance(dataset, Mapping): + raise ValueError(f"eval dataset {dataset_name!r} must be a mapping") + samples = dataset.get("samples") + if not isinstance(samples, Sequence) or isinstance(samples, (str, bytes)): + raise ValueError(f"eval dataset {dataset_name!r} must include a samples sequence") + for metric_name, value in episode_metrics(samples).items(): + metrics[f"eval/{dataset_name}/{metric_name}"] = value + return metrics + + +def _emit_metrics(rollout_id: int, args: Any, metrics: dict[str, Any]) -> None: + from relax.utils import tracking_utils + from relax.utils.logging_utils import get_logger + from relax.utils.metrics.metric_utils import compute_rollout_step + + step = compute_rollout_step(args, rollout_id) + metrics["eval/step"] = step + get_logger(__name__).info("eval %s: %s", rollout_id, metrics) + tracking_utils.log(args, metrics, step_key="eval/step") + tracking_utils.flush_metrics(args, step) + + +def log_eval_rollout_data( + rollout_id: int, + args: Any, + data: Mapping[str, Mapping[str, Any]], + extra_metrics: Mapping[str, Any] | None = None, +) -> bool: + """Relax custom eval hook that replaces row-weighted default logging.""" + + metrics = build_eval_metrics(data, extra_metrics) + _emit_metrics(rollout_id, args, metrics) + return True diff --git a/examples/graphgpo/graph_credit.py b/examples/graphgpo/graph_credit.py new file mode 100644 index 000000000..9d3f35210 --- /dev/null +++ b/examples/graphgpo/graph_credit.py @@ -0,0 +1,794 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Pure-Python credit assignment kernels for the GraphGPO example. + +The implementation intentionally has no dependency on Ray, PyTorch, or the +ALFWorld environment. It operates on finalized turn transitions so the same +kernel can be exercised by unit tests and by the custom-advantage adapter. +""" + +from __future__ import annotations + +import heapq +import math +from collections import defaultdict +from dataclasses import dataclass, replace +from itertools import count +from time import perf_counter_ns +from typing import Hashable, Literal, Mapping, Sequence, TypeVar + +from examples.graphgpo.diagnostics import ( + DiagnosticsCallback, + build_graph_diagnostics, +) + + +SUCCESS = "__GRAPHGPO_SUCCESS__" +# TurnKey stays compact because every public computation first validates one +# group-local identity namespace; rollout_group_id and policy_version are +# checked by the custom adapter before these keys are constructed. +TurnKey = tuple[Hashable, int] +TrajectoryId = Hashable +EpisodeWeighting = Literal["trajectory_once", "reference_cross_steps"] +DEFAULT_EPISODE_WEIGHTING: EpisodeWeighting = "trajectory_once" +EPISODE_WEIGHTINGS = frozenset(("trajectory_once", "reference_cross_steps")) +K = TypeVar("K", bound=Hashable) + + +@dataclass(frozen=True) +class Turn: + """One real environment action and its resulting state transition.""" + + task_id: Hashable + trajectory_id: TrajectoryId + turn_index: int + state_key: Hashable + action: str + next_state_key: Hashable + success: bool + terminal: bool + truncated: bool + is_action_valid: bool + cost: float = 1.0 + + @property + def key(self) -> TurnKey: + return (self.trajectory_id, self.turn_index) + + +@dataclass(frozen=True) +class EdgeOccurrence: + """An observed graph edge; repeated transitions remain repeated.""" + + turn_key: TurnKey + source: Hashable + action: str + target: Hashable + cost: float + + +@dataclass(frozen=True) +class OccurrenceGraph: + """One task-local graph with both topology and occurrence information.""" + + task_id: Hashable + nodes: frozenset[Hashable] + occurrences: tuple[EdgeOccurrence, ...] + reverse_edges: tuple[tuple[Hashable, Hashable, float], ...] + has_success: bool + + +@dataclass(frozen=True) +class DistanceResult: + """Finite distances used by the reference GraphGPO reward convention.""" + + distances: Mapping[Hashable, float] + max_finite_distance: float + has_success: bool + + +def _stable_key(value: object) -> str: + value_type = type(value) + return f"{value_type.__module__}.{value_type.__qualname__}:{value!r}" + + +def _turn_sort_key(turn: Turn) -> tuple[str, str, int]: + return (_stable_key(turn.task_id), _stable_key(turn.trajectory_id), turn.turn_index) + + +def _require_hashable(name: str, value: object) -> None: + try: + hash(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be hashable") from exc + + +def _require_finite(name: str, value: float) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{name} must be a finite real number") + result = float(value) + if not math.isfinite(result): + raise ValueError(f"{name} must be a finite real number") + return result + + +def _canonicalize_turns(turns: Sequence[Turn], *, require_finalized: bool) -> tuple[Turn, ...]: + if not turns: + raise ValueError("at least one turn is required") + + by_trajectory: dict[TrajectoryId, list[Turn]] = defaultdict(list) + task_by_trajectory: dict[TrajectoryId, Hashable] = {} + seen_keys: set[TurnKey] = set() + + for turn in turns: + if not isinstance(turn, Turn): + raise ValueError("all entries must be Turn instances") + for name, value in ( + ("task_id", turn.task_id), + ("trajectory_id", turn.trajectory_id), + ("state_key", turn.state_key), + ("next_state_key", turn.next_state_key), + ): + _require_hashable(name, value) + if turn.state_key == SUCCESS: + raise ValueError("SUCCESS is a sink and cannot be a source state") + if isinstance(turn.turn_index, bool) or not isinstance(turn.turn_index, int) or turn.turn_index < 0: + raise ValueError("turn_index must be a non-negative integer") + if not isinstance(turn.action, str): + raise ValueError("action must be a string") + if not all( + isinstance(flag, bool) + for flag in ( + turn.success, + turn.terminal, + turn.truncated, + turn.is_action_valid, + ) + ): + raise ValueError("success, terminal, truncated, and is_action_valid must be booleans") + if _require_finite("cost", turn.cost) <= 0.0: + raise ValueError("cost must be greater than zero") + if turn.key in seen_keys: + raise ValueError(f"duplicate turn key: {turn.key!r}") + seen_keys.add(turn.key) + + previous_task = task_by_trajectory.setdefault(turn.trajectory_id, turn.task_id) + if previous_task != turn.task_id: + raise ValueError("a trajectory_id cannot span multiple tasks") + by_trajectory[turn.trajectory_id].append(turn) + + for trajectory_id, trajectory in by_trajectory.items(): + trajectory.sort(key=lambda item: item.turn_index) + expected_indices = list(range(len(trajectory))) + actual_indices = [turn.turn_index for turn in trajectory] + if actual_indices != expected_indices: + raise ValueError(f"trajectory {trajectory_id!r} has non-contiguous turn indices") + + success_values = {turn.success for turn in trajectory} + if len(success_values) != 1: + raise ValueError(f"trajectory {trajectory_id!r} has inconsistent success flags") + + for current, following in zip(trajectory, trajectory[1:]): + if current.next_state_key != following.state_key: + raise ValueError(f"trajectory {trajectory_id!r} has a broken state transition chain") + if current.terminal or current.truncated: + raise ValueError(f"trajectory {trajectory_id!r} terminates before its final turn") + + final_turn = trajectory[-1] + if final_turn.terminal and final_turn.truncated: + raise ValueError(f"trajectory {trajectory_id!r} cannot be both terminal and truncated") + if not final_turn.terminal and not final_turn.truncated: + raise ValueError(f"trajectory {trajectory_id!r} must end as terminal or truncated") + + if require_finalized: + success = final_turn.success + success_targets = [turn.turn_index for turn in trajectory if turn.next_state_key == SUCCESS] + if success and success_targets != [final_turn.turn_index]: + raise ValueError("a successful trajectory must point its final real action to SUCCESS") + if not success and success_targets: + raise ValueError("an unsuccessful trajectory cannot point to SUCCESS") + + return tuple(sorted(turns, key=_turn_sort_key)) + + +def finalize_trajectory( + turns: Sequence[Turn], + *, + success: bool, + success_key: Hashable = SUCCESS, +) -> tuple[Turn, ...]: + """Finalize one trajectory without adding a synthetic transition. + + On success, only the last real action's ``next_state_key`` is replaced by + ``success_key``. The row count and all earlier transitions remain intact. + """ + + if not isinstance(success, bool): + raise ValueError("success must be a boolean") + _require_hashable("success_key", success_key) + if success_key != SUCCESS: + raise ValueError("the graph kernel requires the canonical SUCCESS key") + if not turns: + raise ValueError("at least one turn is required") + + trajectory_ids = {turn.trajectory_id for turn in turns if isinstance(turn, Turn)} + task_ids = {turn.task_id for turn in turns if isinstance(turn, Turn)} + if len(trajectory_ids) != 1 or len(task_ids) != 1 or len(trajectory_ids) != len(task_ids): + raise ValueError("finalize_trajectory accepts exactly one task-local trajectory") + + rewritten = [replace(turn, success=success) for turn in turns] + rewritten.sort(key=lambda item: item.turn_index) + last_index = len(rewritten) - 1 + for index, turn in enumerate(rewritten): + if turn.next_state_key == SUCCESS and index != last_index: + raise ValueError("SUCCESS may only appear on the final real action") + + if success: + rewritten[-1] = replace( + rewritten[-1], + next_state_key=SUCCESS, + success=True, + terminal=True, + truncated=False, + ) + elif rewritten[-1].next_state_key == SUCCESS: + raise ValueError("an unsuccessful trajectory cannot point to SUCCESS") + + return _canonicalize_turns(rewritten, require_finalized=True) + + +def episode_return( + turns: Sequence[Turn], + *, + success_reward: float = 10.0, + invalid_penalty: float = 0.1, +) -> float: + """Return ``success_reward * success - invalid_penalty * invalid_count``.""" + + canonical = _canonicalize_turns(turns, require_finalized=True) + if len({turn.trajectory_id for turn in canonical}) != 1: + raise ValueError("episode_return accepts exactly one trajectory") + success_reward = _require_finite("success_reward", success_reward) + invalid_penalty = _require_finite("invalid_penalty", invalid_penalty) + if success_reward < 0.0 or invalid_penalty < 0.0: + raise ValueError("reward and penalty magnitudes must be non-negative") + + success = canonical[0].success + invalid_count = sum(not turn.is_action_valid for turn in canonical) + result = success_reward * float(success) - invalid_penalty * invalid_count + return _require_finite("episode return", result) + + +def build_occurrence_graph(turns: Sequence[Turn]) -> OccurrenceGraph: + """Build one task-local graph while preserving every sampled edge.""" + + canonical = _canonicalize_turns(turns, require_finalized=True) + task_ids = {turn.task_id for turn in canonical} + if len(task_ids) != 1: + raise ValueError("build_occurrence_graph accepts turns from exactly one task") + task_id = next(iter(task_ids)) + + occurrences = tuple( + EdgeOccurrence( + turn_key=turn.key, + source=turn.state_key, + action=turn.action, + target=turn.next_state_key, + cost=float(turn.cost), + ) + for turn in canonical + ) + nodes = frozenset(node for occurrence in occurrences for node in (occurrence.source, occurrence.target)) + + minimum_costs: dict[tuple[Hashable, Hashable], float] = {} + for occurrence in occurrences: + topology_key = (occurrence.source, occurrence.target) + old_cost = minimum_costs.get(topology_key) + if old_cost is None or occurrence.cost < old_cost: + minimum_costs[topology_key] = occurrence.cost + + reverse_edges = tuple( + sorted( + ((target, source, cost) for (source, target), cost in minimum_costs.items()), + key=lambda item: (_stable_key(item[0]), _stable_key(item[1]), item[2]), + ) + ) + return OccurrenceGraph( + task_id=task_id, + nodes=nodes, + occurrences=occurrences, + reverse_edges=reverse_edges, + has_success=SUCCESS in nodes, + ) + + +def reverse_shortest_distances(graph: OccurrenceGraph) -> DistanceResult: + """Compute distances to SUCCESS with a finite unreachable replacement.""" + + if not isinstance(graph, OccurrenceGraph): + raise ValueError("graph must be an OccurrenceGraph") + if not graph.nodes: + raise ValueError("graph must contain at least one node") + + if not graph.has_success: + return DistanceResult( + distances={node: 1.0 for node in sorted(graph.nodes, key=_stable_key)}, + max_finite_distance=0.0, + has_success=False, + ) + + reverse_adjacency: dict[Hashable, list[tuple[Hashable, float]]] = defaultdict(list) + for target, source, cost in graph.reverse_edges: + reverse_adjacency[target].append((source, cost)) + + distances: dict[Hashable, float] = {SUCCESS: 0.0} + tie_breaker = count() + queue: list[tuple[float, int, Hashable]] = [(0.0, next(tie_breaker), SUCCESS)] + while queue: + distance, _, node = heapq.heappop(queue) + if distance != distances.get(node): + continue + for predecessor, edge_cost in reverse_adjacency.get(node, ()): + candidate = distance + edge_cost + if candidate < distances.get(predecessor, math.inf): + distances[predecessor] = candidate + heapq.heappush(queue, (candidate, next(tie_breaker), predecessor)) + + max_finite = max(distances.values()) + unreachable_distance = max_finite + 1.0 + effective_distances = { + node: distances.get(node, unreachable_distance) for node in sorted(graph.nodes, key=_stable_key) + } + return DistanceResult( + distances=effective_distances, + max_finite_distance=max_finite, + has_success=True, + ) + + +def _validate_distance_result(graph: OccurrenceGraph, distance_result: DistanceResult) -> None: + if distance_result.has_success != graph.has_success: + raise ValueError("distance result success flag does not match the graph") + if set(distance_result.distances) != set(graph.nodes): + raise ValueError("distance result keys do not match graph nodes") + max_finite = _require_finite("max_finite_distance", distance_result.max_finite_distance) + if max_finite < 0.0: + raise ValueError("max_finite_distance must be non-negative") + for node, value in distance_result.distances.items(): + if _require_finite(f"distance for {node!r}", value) < 0.0: + raise ValueError("node distances must be non-negative") + if graph.has_success and distance_result.distances[SUCCESS] != 0.0: + raise ValueError("SUCCESS must have distance zero") + + +def graph_raw_returns( + graph: OccurrenceGraph, + distance_result: DistanceResult, + *, + omega: float = 0.1, + success_reward: float = 10.0, + convention: Literal["reference", "paper"] = "reference", +) -> dict[TurnKey, float]: + """Compute raw graph returns for each observed transition occurrence. + + ``reference`` implements the frozen recipe used by the proposal: + ``success_reward * omega ** d(next)`` with unit edge costs. ``paper`` + implements GraphGPO Eq. 4: ``success_reward * omega ** (d(next) + cost)``. + """ + + if not isinstance(graph, OccurrenceGraph) or not isinstance(distance_result, DistanceResult): + raise ValueError("graph and distance_result have invalid types") + _validate_distance_result(graph, distance_result) + omega = _require_finite("omega", omega) + success_reward = _require_finite("success_reward", success_reward) + if not 0.0 < omega < 1.0: + raise ValueError("omega must be between zero and one") + if success_reward < 0.0: + raise ValueError("success_reward must be non-negative") + if convention not in {"reference", "paper"}: + raise ValueError(f"unknown graph reward convention: {convention!r}") + + result: dict[TurnKey, float] = {} + for occurrence in graph.occurrences: + if occurrence.target not in distance_result.distances: + raise ValueError(f"missing distance for node {occurrence.target!r}") + if convention == "reference": + if occurrence.cost != 1.0: + raise ValueError("the reference convention requires unit edge costs") + exponent = distance_result.distances[occurrence.target] + else: + exponent = distance_result.distances[occurrence.target] + occurrence.cost + value = success_reward * omega**exponent + result[occurrence.turn_key] = _require_finite("graph raw return", value) + + expected_keys = {occurrence.turn_key for occurrence in graph.occurrences} + if set(result) != expected_keys: + raise ValueError("graph raw return keys do not match edge occurrence keys") + return result + + +def standardize_by_key( + values: Mapping[K, float], + group_keys: Mapping[K, Hashable], + *, + eps: float = 1e-6, + ddof: int = 1, + singleton_zero: bool = True, +) -> dict[K, float]: + """Standardize values independently inside each group. + + The computation uses ``math.fsum`` and a deterministic key order. With + ``ddof=1`` it matches the frozen recipe's sample standard deviation. + """ + + if set(values) != set(group_keys): + raise ValueError("values and group_keys must contain exactly the same keys") + eps = _require_finite("eps", eps) + if eps < 0.0: + raise ValueError("eps must be non-negative") + if isinstance(ddof, bool) or not isinstance(ddof, int) or ddof < 0: + raise ValueError("ddof must be a non-negative integer") + if not isinstance(singleton_zero, bool): + raise ValueError("singleton_zero must be a boolean") + + grouped: dict[Hashable, list[K]] = defaultdict(list) + finite_values: dict[K, float] = {} + for key in sorted(values, key=_stable_key): + value = _require_finite(f"value for {key!r}", values[key]) + group_key = group_keys[key] + _require_hashable("group key", group_key) + grouped[group_key].append(key) + finite_values[key] = value + + result: dict[K, float] = {} + for group_key in sorted(grouped, key=_stable_key): + members = grouped[group_key] + if len(members) <= ddof: + if not singleton_zero: + raise ValueError(f"group {group_key!r} is too small for ddof={ddof}") + result.update((key, 0.0) for key in members) + continue + + mean = math.fsum(finite_values[key] for key in members) / len(members) + squared_deviations = math.fsum((finite_values[key] - mean) ** 2 for key in members) + variance = squared_deviations / (len(members) - ddof) + if variance <= 0.0: + result.update((key, 0.0) for key in members) + continue + + denominator = math.sqrt(variance) + eps + for key in members: + result[key] = _require_finite( + f"standardized value for {key!r}", + (finite_values[key] - mean) / denominator, + ) + + if set(result) != set(values): + raise ValueError("standardization did not return every input key") + return result + + +def graph_advantages( + graph: OccurrenceGraph, + distance_result: DistanceResult, + *, + omega: float = 0.1, + success_reward: float = 10.0, + eps: float = 1e-6, + convention: Literal["reference", "paper"] = "reference", +) -> dict[TurnKey, float]: + """Return occurrence-weighted, same-source graph advantages.""" + + omega = _require_finite("omega", omega) + success_reward = _require_finite("success_reward", success_reward) + eps = _require_finite("eps", eps) + if not 0.0 < omega < 1.0: + raise ValueError("omega must be between zero and one") + if success_reward < 0.0 or eps < 0.0: + raise ValueError("success_reward and eps must be non-negative") + if convention not in {"reference", "paper"}: + raise ValueError(f"unknown graph reward convention: {convention!r}") + + _validate_distance_result(graph, distance_result) + keys = {occurrence.turn_key for occurrence in graph.occurrences} + if not distance_result.has_success: + return {key: 0.0 for key in sorted(keys, key=_stable_key)} + + raw_returns = graph_raw_returns( + graph, + distance_result, + omega=omega, + success_reward=success_reward, + convention=convention, + ) + sources = {occurrence.turn_key: occurrence.source for occurrence in graph.occurrences} + return standardize_by_key(raw_returns, sources, eps=eps, ddof=1, singleton_zero=True) + + +def episode_advantages( + episode_returns: Mapping[TrajectoryId, float], + task_ids: Mapping[TrajectoryId, Hashable], + *, + expected_group_size: int | None = None, + eps: float = 1e-6, + weighting: EpisodeWeighting = DEFAULT_EPISODE_WEIGHTING, + trajectory_lengths: Mapping[TrajectoryId, int] | None = None, +) -> dict[TrajectoryId, float]: + """Normalize episode returns independently for every task. + + ``trajectory_once`` is the usual GRPO convention: every trajectory + contributes one value to the group statistics. ``reference_cross_steps`` + reproduces the frozen GraphGPO helper: a trajectory's episode return + contributes once for every real turn before the result is broadcast back to + all turns in that trajectory. + """ + + if not episode_returns: + raise ValueError("at least one episode return is required") + if set(episode_returns) != set(task_ids): + raise ValueError("episode_returns and task_ids must contain exactly the same trajectories") + if weighting not in EPISODE_WEIGHTINGS: + raise ValueError(f"unknown episode weighting: {weighting!r}") + if trajectory_lengths is not None and set(trajectory_lengths) != set(episode_returns): + raise ValueError("trajectory_lengths must contain exactly the same trajectories") + + finite_returns = { + trajectory_id: _require_finite(f"episode return for {trajectory_id!r}", episode_return) + for trajectory_id, episode_return in episode_returns.items() + } + finite_eps = _require_finite("eps", eps) + if finite_eps < 0.0: + raise ValueError("eps must be non-negative") + + normalized_lengths: dict[TrajectoryId, int] | None = None + if weighting == "reference_cross_steps": + if trajectory_lengths is None: + raise ValueError("trajectory_lengths are required for reference_cross_steps") + normalized_lengths = {} + for trajectory_id, length in trajectory_lengths.items(): + if isinstance(length, bool) or not isinstance(length, int) or length <= 0: + raise ValueError("trajectory lengths must be positive integers") + normalized_lengths[trajectory_id] = length + + if expected_group_size is not None: + if ( + isinstance(expected_group_size, bool) + or not isinstance(expected_group_size, int) + or expected_group_size <= 0 + ): + raise ValueError("expected_group_size must be a positive integer") + counts: dict[Hashable, int] = defaultdict(int) + for task_id in task_ids.values(): + _require_hashable("task_id", task_id) + counts[task_id] += 1 + wrong_sizes = {task_id: count for task_id, count in counts.items() if count != expected_group_size} + if wrong_sizes: + raise ValueError(f"trajectory group size mismatch: {wrong_sizes!r}") + + if weighting == "trajectory_once": + return standardize_by_key( + finite_returns, + task_ids, + eps=finite_eps, + ddof=1, + singleton_zero=True, + ) + + assert normalized_lengths is not None + trajectories_by_task: dict[Hashable, list[TrajectoryId]] = defaultdict(list) + for trajectory_id in sorted(finite_returns, key=_stable_key): + task_id = task_ids[trajectory_id] + _require_hashable("task_id", task_id) + trajectories_by_task[task_id].append(trajectory_id) + + result: dict[TrajectoryId, float] = {} + for task_id in sorted(trajectories_by_task, key=_stable_key): + trajectories = trajectories_by_task[task_id] + expanded_returns = [ + finite_returns[trajectory_id] + for trajectory_id in trajectories + for _ in range(normalized_lengths[trajectory_id]) + ] + if len(expanded_returns) == 1: + # Matches the frozen helper's explicit singleton branch. + mean = 0.0 + standard_deviation = 1.0 + else: + mean = math.fsum(expanded_returns) / len(expanded_returns) + squared_deviations = math.fsum((value - mean) ** 2 for value in expanded_returns) + variance = squared_deviations / (len(expanded_returns) - 1) + standard_deviation = math.sqrt(max(variance, 0.0)) + + denominator = standard_deviation + finite_eps + for trajectory_id in trajectories: + result[trajectory_id] = _require_finite( + f"episode advantage for {trajectory_id!r}", + (finite_returns[trajectory_id] - mean) / denominator, + ) + + if set(result) != set(finite_returns): + raise ValueError("episode normalization did not return every trajectory") + return result + + +def gigpo_discounted_returns( + turns: Sequence[Turn], + immediate_rewards: Mapping[TurnKey, float], + *, + gamma: float = 0.95, +) -> dict[TurnKey, float]: + """Compute GiGPO's discounted return from every real environment turn.""" + + canonical = _canonicalize_turns(turns, require_finalized=True) + expected_keys = {turn.key for turn in canonical} + if set(immediate_rewards) != expected_keys: + raise ValueError("immediate_rewards must contain exactly one value per turn") + gamma = _require_finite("gamma", gamma) + if not 0.0 <= gamma <= 1.0: + raise ValueError("gamma must be between zero and one") + finite_rewards = { + key: _require_finite(f"immediate reward for {key!r}", value) for key, value in immediate_rewards.items() + } + + by_trajectory: dict[TrajectoryId, list[Turn]] = defaultdict(list) + for turn in canonical: + by_trajectory[turn.trajectory_id].append(turn) + + result: dict[TurnKey, float] = {} + for trajectory_id in sorted(by_trajectory, key=_stable_key): + running_return = 0.0 + for turn in reversed(by_trajectory[trajectory_id]): + running_return = finite_rewards[turn.key] + gamma * running_return + result[turn.key] = _require_finite("discounted return", running_return) + + if set(result) != expected_keys: + raise ValueError("discounted returns do not match input turn keys") + return result + + +def _default_immediate_rewards( + turns: Sequence[Turn], + *, + success_reward: float, + invalid_penalty: float, +) -> dict[TurnKey, float]: + result: dict[TurnKey, float] = {} + for turn in turns: + reward = -invalid_penalty if not turn.is_action_valid else 0.0 + if turn.success and turn.next_state_key == SUCCESS: + reward += success_reward + result[turn.key] = reward + return result + + +def compute_method_advantages( + method: Literal["grpo", "gigpo", "graphgpo"] | str, + turns: Sequence[Turn], + *, + expected_group_size: int | None = None, + success_reward: float = 10.0, + invalid_penalty: float = 0.1, + omega: float = 0.1, + gamma: float = 0.95, + beta_episode: float = 1.0, + beta_graph: float = 1.0, + beta_step: float = 1.0, + eps: float = 1e-6, + immediate_rewards: Mapping[TurnKey, float] | None = None, + episode_weighting: EpisodeWeighting = DEFAULT_EPISODE_WEIGHTING, + graph_diagnostics_callback: DiagnosticsCallback | None = None, + diagnostic_rollout_group_id: object | None = None, + diagnostic_policy_version: object | None = None, +) -> dict[TurnKey, float]: + """Compute GRPO, GiGPO, or GraphGPO scalar advantages for every turn.""" + + if method not in {"grpo", "gigpo", "graphgpo"}: + raise ValueError(f"unknown method: {method!r}") + canonical = _canonicalize_turns(turns, require_finalized=True) + success_reward = _require_finite("success_reward", success_reward) + invalid_penalty = _require_finite("invalid_penalty", invalid_penalty) + beta_episode = _require_finite("beta_episode", beta_episode) + beta_graph = _require_finite("beta_graph", beta_graph) + beta_step = _require_finite("beta_step", beta_step) + if success_reward < 0.0 or invalid_penalty < 0.0: + raise ValueError("reward and penalty magnitudes must be non-negative") + + by_trajectory: dict[TrajectoryId, list[Turn]] = defaultdict(list) + task_by_trajectory: dict[TrajectoryId, Hashable] = {} + for turn in canonical: + by_trajectory[turn.trajectory_id].append(turn) + task_by_trajectory[turn.trajectory_id] = turn.task_id + + returns = { + trajectory_id: episode_return( + trajectory, + success_reward=success_reward, + invalid_penalty=invalid_penalty, + ) + for trajectory_id, trajectory in by_trajectory.items() + } + per_trajectory_episode = episode_advantages( + returns, + task_by_trajectory, + expected_group_size=expected_group_size, + eps=eps, + weighting=episode_weighting, + trajectory_lengths={trajectory_id: len(trajectory) for trajectory_id, trajectory in by_trajectory.items()}, + ) + episode_by_turn = {turn.key: per_trajectory_episode[turn.trajectory_id] for turn in canonical} + + if method == "grpo": + result = {key: beta_episode * value for key, value in episode_by_turn.items()} + elif method == "graphgpo": + graph_by_turn: dict[TurnKey, float] = {} + turns_by_task: dict[Hashable, list[Turn]] = defaultdict(list) + for turn in canonical: + turns_by_task[turn.task_id].append(turn) + for task_id in sorted(turns_by_task, key=_stable_key): + build_started_ns = perf_counter_ns() if graph_diagnostics_callback is not None else 0 + graph = build_occurrence_graph(turns_by_task[task_id]) + graph_build_ns = perf_counter_ns() - build_started_ns if graph_diagnostics_callback is not None else 0 + dijkstra_started_ns = perf_counter_ns() if graph_diagnostics_callback is not None else 0 + distance_result = reverse_shortest_distances(graph) + reverse_dijkstra_ns = ( + perf_counter_ns() - dijkstra_started_ns if graph_diagnostics_callback is not None else 0 + ) + advantage_started_ns = perf_counter_ns() if graph_diagnostics_callback is not None else 0 + task_graph_advantages = graph_advantages( + graph, + distance_result, + omega=omega, + success_reward=success_reward, + eps=eps, + convention="reference", + ) + graph_advantage_ns = ( + perf_counter_ns() - advantage_started_ns if graph_diagnostics_callback is not None else 0 + ) + graph_by_turn.update(task_graph_advantages) + if graph_diagnostics_callback is not None: + graph_diagnostics_callback( + build_graph_diagnostics( + task_id=task_id, + nodes=graph.nodes, + occurrences=graph.occurrences, + distances=distance_result.distances, + max_finite_distance=distance_result.max_finite_distance, + has_success=distance_result.has_success, + advantages=task_graph_advantages, + graph_build_ns=graph_build_ns, + reverse_dijkstra_ns=reverse_dijkstra_ns, + graph_advantage_ns=graph_advantage_ns, + rollout_group_id=diagnostic_rollout_group_id, + policy_version=diagnostic_policy_version, + ) + ) + result = { + turn.key: beta_episode * episode_by_turn[turn.key] + beta_graph * graph_by_turn[turn.key] + for turn in canonical + } + else: + if immediate_rewards is None: + immediate_rewards = _default_immediate_rewards( + canonical, + success_reward=success_reward, + invalid_penalty=invalid_penalty, + ) + discounted_returns = gigpo_discounted_returns(canonical, immediate_rewards, gamma=gamma) + state_groups = {turn.key: (turn.task_id, turn.state_key) for turn in canonical} + step_by_turn = standardize_by_key( + discounted_returns, + state_groups, + eps=eps, + ddof=1, + singleton_zero=True, + ) + result = { + turn.key: beta_episode * episode_by_turn[turn.key] + beta_step * step_by_turn[turn.key] + for turn in canonical + } + + expected_keys = {turn.key for turn in canonical} + if set(result) != expected_keys: + raise ValueError("computed advantages do not match input turn keys") + return {key: _require_finite(f"advantage for {key!r}", result[key]) for key in sorted(result, key=_stable_key)} diff --git a/examples/graphgpo/manifest.py b/examples/graphgpo/manifest.py new file mode 100644 index 000000000..a58f0c3f6 --- /dev/null +++ b/examples/graphgpo/manifest.py @@ -0,0 +1,244 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Deterministic, content-addressed ALFWorld task manifests.""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Sequence + + +MANIFEST_VERSION = "alfworld-task-manifest-v1" +TASK_FILENAMES = ("game.tw-pddl", "traj_data.json") +SHARED_FILENAMES = ("alfred.pddl", "alfred.twl2") +ALFWORLD_TASK_TYPES = ( + "look_at_obj_in_light", + "pick_and_place_simple", + "pick_clean_then_place_in_recep", + "pick_cool_then_place_in_recep", + "pick_heat_then_place_in_recep", + "pick_two_obj_and_place", +) + + +@dataclass(frozen=True) +class FileManifestEntry: + """One immutable file referenced by the task manifest.""" + + relative_path: str + sha256: str + size_bytes: int + + +@dataclass(frozen=True) +class TaskManifestEntry: + """One ALFWorld task and both files needed to reproduce it.""" + + task_id: str + split: str + task_type: str + game: FileManifestEntry + trajectory: FileManifestEntry + + +@dataclass(frozen=True) +class SharedAssetEntry: + """One shared ALFWorld grammar or PDDL asset.""" + + asset_name: str + file: FileManifestEntry + + +@dataclass(frozen=True) +class TaskManifest: + """A fixed-order collection of content-addressed ALFWorld tasks.""" + + schema_version: str + split: str + tasks: tuple[TaskManifestEntry, ...] + shared_assets: tuple[SharedAssetEntry, ...] + + +def sha256_file(path: Path, *, chunk_size: int = 1024 * 1024) -> str: + if isinstance(chunk_size, bool) or not isinstance(chunk_size, int) or chunk_size <= 0: + raise ValueError("chunk_size must be a positive integer") + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(chunk_size), b""): + digest.update(chunk) + return digest.hexdigest() + + +def infer_task_type(relative_path: str) -> str: + """Infer the official ALFWorld task family from a relative path.""" + + if not isinstance(relative_path, str) or not relative_path: + raise ValueError("relative_path must be a non-empty string") + for part in Path(relative_path).parts: + for task_type in ALFWORLD_TASK_TYPES: + if part == task_type or part.startswith(f"{task_type}-"): + return task_type + raise ValueError(f"cannot infer ALFWorld task type from {relative_path!r}") + + +def discover_task_files(root: Path, *, filename: str = "game.tw-pddl") -> list[Path]: + """Discover task files in stable relative-path order.""" + + root = root.resolve(strict=True) + if not root.is_dir(): + raise ValueError("root must be a directory") + if not isinstance(filename, str) or not filename: + raise ValueError("filename must be a non-empty string") + return sorted( + (path for path in root.rglob(filename) if path.is_file()), + key=lambda path: path.relative_to(root).as_posix(), + ) + + +def discover_shared_files(root: Path) -> list[Path]: + """Find exactly one copy of every required shared ALFWorld asset.""" + + root = root.resolve(strict=True) + shared_files: list[Path] = [] + for filename in SHARED_FILENAMES: + matches = sorted( + (path for path in root.rglob(filename) if path.is_file()), + key=lambda path: path.relative_to(root).as_posix(), + ) + if len(matches) != 1: + raise ValueError(f"expected exactly one {filename!r} below root, found {len(matches)}") + shared_files.append(matches[0]) + return shared_files + + +def _resolve_task_path(root: Path, path: Path) -> tuple[Path, str]: + candidate = path if path.is_absolute() else root / path + candidate = candidate.resolve(strict=True) + try: + relative_path = candidate.relative_to(root).as_posix() + except ValueError as exc: + raise ValueError(f"task path is outside root: {path}") from exc + if not candidate.is_file(): + raise ValueError(f"task path is not a regular file: {path}") + return candidate, relative_path + + +def _file_entry(path: Path, relative_path: str) -> FileManifestEntry: + return FileManifestEntry( + relative_path=relative_path, + sha256=sha256_file(path), + size_bytes=path.stat().st_size, + ) + + +def build_manifest( + root: Path, + *, + split: str, + task_files: Sequence[Path] | None = None, + shared_files: Sequence[Path] | None = None, +) -> TaskManifest: + """Hash every task and shared asset required for a reproducible run. + + ``task_files`` contains ``game.tw-pddl`` paths. The sibling + ``traj_data.json`` is mandatory and is frozen in the same task entry. + ``shared_files`` must contain one ``alfred.pddl`` and one ``alfred.twl2``; + when omitted, the builder discovers exactly one of each below ``root``. + """ + + root = root.resolve(strict=True) + if not root.is_dir(): + raise ValueError("root must be a directory") + if not isinstance(split, str) or not split.strip(): + raise ValueError("split must be a non-empty string") + split = split.strip() + + if task_files is None: + task_files = discover_task_files(root) + elif isinstance(task_files, (str, bytes)) or not isinstance(task_files, Sequence): + raise TypeError("task_files must be a sequence of paths") + + resolved: list[tuple[Path, str]] = [_resolve_task_path(root, Path(path)) for path in task_files] + resolved.sort(key=lambda item: item[1]) + relative_paths = [relative_path for _, relative_path in resolved] + if len(set(relative_paths)) != len(relative_paths): + raise ValueError("task_files contains duplicate paths") + + entries: list[TaskManifestEntry] = [] + for game_path, game_relative_path in resolved: + if game_path.name != TASK_FILENAMES[0]: + raise ValueError(f"task file must be named {TASK_FILENAMES[0]!r}: {game_path}") + trajectory_path = game_path.with_name(TASK_FILENAMES[1]) + trajectory_path, trajectory_relative_path = _resolve_task_path(root, trajectory_path) + task_parent = Path(game_relative_path).parent.as_posix() + entries.append( + TaskManifestEntry( + task_id=task_parent, + split=split, + task_type=infer_task_type(game_relative_path), + game=_file_entry(game_path, game_relative_path), + trajectory=_file_entry( + trajectory_path, + trajectory_relative_path, + ), + ) + ) + + if shared_files is None: + shared_files = discover_shared_files(root) + elif isinstance(shared_files, (str, bytes)) or not isinstance(shared_files, Sequence): + raise TypeError("shared_files must be a sequence of paths") + + shared_by_name: dict[str, SharedAssetEntry] = {} + for shared_file in shared_files: + path, relative_path = _resolve_task_path(root, Path(shared_file)) + if path.name not in SHARED_FILENAMES: + raise ValueError(f"unexpected shared asset: {path.name!r}") + if path.name in shared_by_name: + raise ValueError(f"duplicate shared asset: {path.name!r}") + shared_by_name[path.name] = SharedAssetEntry( + asset_name=path.name, + file=_file_entry(path, relative_path), + ) + missing_shared = set(SHARED_FILENAMES) - set(shared_by_name) + if missing_shared: + raise ValueError(f"missing shared assets: {sorted(missing_shared)!r}") + + return TaskManifest( + schema_version=MANIFEST_VERSION, + split=split, + tasks=tuple(entries), + shared_assets=tuple(shared_by_name[filename] for filename in SHARED_FILENAMES), + ) + + +def manifest_bytes(manifest: TaskManifest) -> bytes: + if not isinstance(manifest, TaskManifest): + raise TypeError("manifest must be a TaskManifest") + payload = { + "schema_version": manifest.schema_version, + "split": manifest.split, + "tasks": [asdict(task) for task in manifest.tasks], + "shared_assets": [asdict(asset) for asset in manifest.shared_assets], + } + return ( + json.dumps( + payload, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + + "\n" + ).encode("utf-8") + + +def manifest_sha256(manifest: TaskManifest) -> str: + return hashlib.sha256(manifest_bytes(manifest)).hexdigest() + + +def write_manifest(path: Path, manifest: TaskManifest) -> None: + path.write_bytes(manifest_bytes(manifest)) diff --git a/examples/graphgpo/preflight.py b/examples/graphgpo/preflight.py new file mode 100644 index 000000000..9c5abe0d0 --- /dev/null +++ b/examples/graphgpo/preflight.py @@ -0,0 +1,377 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Fail-closed consistency checks for prepared GraphGPO recipe inputs.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import re +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from examples.graphgpo.manifest import MANIFEST_VERSION +from examples.graphgpo.prepare_alfworld import PREPARE_SCHEMA_VERSION + + +MODEL_LOCK_SCHEMA = "task37-huggingface-model-lock-v1" +MODEL_REPO_ID = "Qwen/Qwen2.5-1.5B-Instruct" +MODEL_FILE_COUNT = 9 +_SHA256_HEX = re.compile(r"[0-9a-f]{64}") + + +def _read_json(path: Path, *, label: str) -> dict[str, Any]: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"cannot read {label} JSON from {path}") from exc + if not isinstance(payload, dict): + raise ValueError(f"{label} must contain a JSON object") + return payload + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _safe_locked_file(root: Path, relative_path: str, *, label: str) -> Path: + relative = Path(_non_empty_string(relative_path, label=label)) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError(f"{label} must be a safe relative path") + try: + candidate = (root / relative).resolve(strict=True) + candidate.relative_to(root) + except (OSError, ValueError) as exc: + raise ValueError(f"{label} is missing or outside its locked root") from exc + if not candidate.is_file(): + raise ValueError(f"{label} must reference a regular file") + return candidate + + +def _verify_locked_file( + *, + root: Path, + relative_path: str, + metadata: Mapping[str, Any], + size_key: str, + label: str, +) -> Path: + path = _safe_locked_file(root, relative_path, label=label) + expected_size = _positive_int(metadata.get(size_key), label=f"{label} {size_key}") + expected_sha = _non_empty_string(metadata.get("sha256"), label=f"{label} sha256") + if _SHA256_HEX.fullmatch(expected_sha) is None: + raise ValueError(f"{label} sha256 must be 64 lowercase hexadecimal characters") + if path.stat().st_size != expected_size: + raise ValueError(f"{label} size does not match its lock") + if _sha256(path) != expected_sha: + raise ValueError(f"{label} SHA256 does not match its lock") + return path + + +def verify_model_checkpoint( + *, + model_lock_path: Path, + checkpoint_path: Path, + model_revision: str, +) -> None: + """Hash every independently locked file in the local model snapshot.""" + + model_revision = _non_empty_string(model_revision, label="model_revision") + if re.fullmatch(r"[0-9a-f]{40}", model_revision) is None: + raise ValueError("model_revision must be a pinned 40-character commit SHA") + checkpoint_root = checkpoint_path.resolve(strict=True) + if not checkpoint_root.is_dir(): + raise ValueError("checkpoint_path must be a directory") + lock = _read_json(model_lock_path, label="model lock") + if lock.get("schema") != MODEL_LOCK_SCHEMA: + raise ValueError("model lock schema mismatch") + if lock.get("repo_id") != MODEL_REPO_ID: + raise ValueError("model lock repo_id mismatch") + if lock.get("revision") != model_revision: + raise ValueError("model revision does not match the model lock") + files = lock.get("files") + if not isinstance(files, Mapping) or len(files) != MODEL_FILE_COUNT: + raise ValueError(f"model lock must contain exactly {MODEL_FILE_COUNT} files") + for relative_path, metadata in files.items(): + if not isinstance(relative_path, str) or not isinstance(metadata, Mapping): + raise ValueError("model lock files must map relative paths to metadata") + _verify_locked_file( + root=checkpoint_root, + relative_path=relative_path, + metadata=metadata, + size_key="bytes", + label=f"model file {relative_path!r}", + ) + + +def _positive_int(value: object, *, label: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{label} must be a positive integer") + return value + + +def _non_empty_string(value: object, *, label: str) -> str: + if not isinstance(value, str) or not value: + raise ValueError(f"{label} must be a non-empty string") + return value + + +def validate_batch_arithmetic( + *, + task_groups: int, + group_size: int, + global_batch_size: int, +) -> None: + """Require every rollout's trajectory count to form complete train + batches.""" + + task_groups = _positive_int(task_groups, label="task_groups") + group_size = _positive_int(group_size, label="group_size") + global_batch_size = _positive_int(global_batch_size, label="global_batch_size") + rollout_sample_count = task_groups * group_size + if rollout_sample_count % global_batch_size != 0: + raise ValueError("task_groups * group_size must be divisible by global_batch_size") + + +def _prompt_rows(path: Path) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + try: + with path.open("r", encoding="utf-8") as stream: + for line_number, line in enumerate(stream, start=1): + if not line.strip(): + raise ValueError(f"prompt data has a blank row at line {line_number}") + row = json.loads(line) + if not isinstance(row, dict): + raise ValueError(f"prompt row {line_number} must be a JSON object") + rows.append(row) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"cannot read prompt JSONL from {path}") from exc + return rows + + +def _verify_manifest_assets( + manifest: Mapping[str, Any], + *, + data_root: Path, + split: str, +) -> None: + """Rehash every ALFWorld file referenced by one task manifest.""" + + data_root = data_root.resolve(strict=True) + if not data_root.is_dir(): + raise ValueError("ALFWorld data root must be a directory") + locked_entries: list[tuple[str, Mapping[str, Any], str]] = [] + tasks = manifest.get("tasks") + if not isinstance(tasks, list): + raise ValueError(f"{split} manifest tasks must be a list") + for index, task in enumerate(tasks): + if not isinstance(task, Mapping): + raise ValueError(f"{split} manifest task {index} must be an object") + for field in ("game", "trajectory"): + entry = task.get(field) + if not isinstance(entry, Mapping): + raise ValueError(f"{split} manifest task {index} is missing {field} metadata") + relative_path = _non_empty_string( + entry.get("relative_path"), + label=f"{split} manifest task {index} {field}.relative_path", + ) + locked_entries.append((relative_path, entry, f"{split} task {index} {field}")) + + shared_assets = manifest.get("shared_assets") + if not isinstance(shared_assets, list) or not shared_assets: + raise ValueError(f"{split} manifest shared_assets must be a non-empty list") + for index, asset in enumerate(shared_assets): + if not isinstance(asset, Mapping) or not isinstance(asset.get("file"), Mapping): + raise ValueError(f"{split} shared asset {index} is missing file metadata") + entry = asset["file"] + relative_path = _non_empty_string( + entry.get("relative_path"), + label=f"{split} shared asset {index} relative_path", + ) + locked_entries.append((relative_path, entry, f"{split} shared asset {index}")) + + relative_paths = [relative_path for relative_path, _, _ in locked_entries] + if len(relative_paths) != len(set(relative_paths)): + raise ValueError(f"{split} manifest references duplicate ALFWorld files") + for relative_path, metadata, label in locked_entries: + _verify_locked_file( + root=data_root, + relative_path=relative_path, + metadata=metadata, + size_key="size_bytes", + label=label, + ) + + +def verify_split( + *, + split: str, + prompt_path: Path, + manifest_path: Path, + lock_split: Mapping[str, Any], + max_steps: int, + data_root: Path, +) -> None: + expected_manifest_sha = _non_empty_string(lock_split.get("manifest_sha256"), label=f"{split} manifest_sha256") + expected_prompt_sha = _non_empty_string(lock_split.get("prompt_data_sha256"), label=f"{split} prompt_data_sha256") + if _sha256(manifest_path) != expected_manifest_sha: + raise ValueError(f"{split} manifest SHA256 does not match prepare.lock") + if _sha256(prompt_path) != expected_prompt_sha: + raise ValueError(f"{split} prompt SHA256 does not match prepare.lock") + + manifest = _read_json(manifest_path, label=f"{split} manifest") + if manifest.get("schema_version") != MANIFEST_VERSION: + raise ValueError(f"{split} manifest schema_version mismatch") + if manifest.get("split") != split: + raise ValueError(f"{split} manifest split mismatch") + _verify_manifest_assets(manifest, data_root=data_root, split=split) + tasks = manifest.get("tasks") + if not isinstance(tasks, list) or not tasks: + raise ValueError(f"{split} manifest tasks must be a non-empty list") + expected_count = _positive_int(lock_split.get("task_count"), label=f"{split} task_count") + if len(tasks) != expected_count: + raise ValueError(f"{split} manifest task count does not match prepare.lock") + + task_paths: list[str] = [] + for index, task in enumerate(tasks): + if not isinstance(task, Mapping): + raise ValueError(f"{split} manifest task {index} must be an object") + if task.get("split") != split: + raise ValueError(f"{split} manifest task {index} has the wrong split") + game = task.get("game") + if not isinstance(game, Mapping): + raise ValueError(f"{split} manifest task {index} is missing game metadata") + task_paths.append( + _non_empty_string( + game.get("relative_path"), + label=f"{split} manifest task {index} game.relative_path", + ) + ) + if task_paths != sorted(task_paths) or len(task_paths) != len(set(task_paths)): + raise ValueError(f"{split} manifest tasks are not in unique stable order") + + rows = _prompt_rows(prompt_path) + if len(rows) != expected_count: + raise ValueError(f"{split} prompt row count does not match prepare.lock") + prompt_task_paths: list[str] = [] + for index, row in enumerate(rows): + metadata = row.get("metadata") + if not isinstance(metadata, Mapping): + raise ValueError(f"{split} prompt row {index} is missing metadata") + if metadata.get("alfworld_train_eval") != split: + raise ValueError(f"{split} prompt row {index} has the wrong split") + if metadata.get("manifest_index") != index: + raise ValueError(f"{split} prompt row {index} has the wrong manifest_index") + if metadata.get("manifest_sha256") != expected_manifest_sha: + raise ValueError(f"{split} prompt row {index} has the wrong manifest SHA256") + if metadata.get("max_steps") != max_steps: + raise ValueError(f"{split} prompt row {index} has the wrong max_steps") + manifest_task_id = _non_empty_string( + metadata.get("manifest_task_id"), + label=f"{split} prompt row {index} manifest_task_id", + ) + if metadata.get("task_id") != manifest_task_id: + raise ValueError(f"{split} prompt row {index} task_id mismatch") + prompt_task_paths.append(manifest_task_id) + if prompt_task_paths != task_paths: + raise ValueError(f"{split} prompt task order does not match the manifest") + + +def verify_prepared_artifacts( + *, + prepare_lock_path: Path, + split_artifacts: Mapping[str, tuple[Path, Path]], + max_steps: int, + model_revision: str, + alfworld_data_root: Path, +) -> None: + """Verify the exact prepared artifacts selected by a launcher.""" + + max_steps = _positive_int(max_steps, label="max_steps") + model_revision = _non_empty_string(model_revision, label="model_revision") + if re.fullmatch(r"[0-9a-f]{40}", model_revision) is None: + raise ValueError("model_revision must be a pinned 40-character commit SHA") + if not split_artifacts: + raise ValueError("at least one split artifact binding is required") + + lock = _read_json(prepare_lock_path, label="prepare.lock") + if lock.get("schema_version") != PREPARE_SCHEMA_VERSION: + raise ValueError("prepare.lock schema_version mismatch") + if lock.get("max_steps") != max_steps: + raise ValueError("max_steps does not match prepare.lock") + if lock.get("model_revision") != model_revision: + raise ValueError("model_revision does not match prepare.lock") + splits = lock.get("splits") + if not isinstance(splits, Mapping): + raise ValueError("prepare.lock splits must be an object") + + for split, (prompt_path, manifest_path) in split_artifacts.items(): + lock_split = splits.get(split) + if not isinstance(lock_split, Mapping): + raise ValueError(f"prepare.lock does not contain split {split!r}") + verify_split( + split=split, + prompt_path=prompt_path, + manifest_path=manifest_path, + lock_split=lock_split, + max_steps=max_steps, + data_root=alfworld_data_root, + ) + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Verify prepared GraphGPO recipe inputs.") + parser.add_argument("--prepare-lock", type=Path, required=True) + parser.add_argument("--alfworld-data-root", type=Path, required=True) + parser.add_argument("--model-lock", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--max-steps", type=int, required=True) + parser.add_argument("--model-revision", required=True) + parser.add_argument("--task-groups", type=int, required=True) + parser.add_argument("--group-size", type=int, required=True) + parser.add_argument("--global-batch-size", type=int, required=True) + parser.add_argument( + "--split-artifact", + action="append", + nargs=3, + metavar=("SPLIT", "PROMPT_DATA", "MANIFEST"), + required=True, + ) + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> None: + args = parse_args(argv) + validate_batch_arithmetic( + task_groups=args.task_groups, + group_size=args.group_size, + global_batch_size=args.global_batch_size, + ) + verify_model_checkpoint( + model_lock_path=args.model_lock, + checkpoint_path=args.checkpoint, + model_revision=args.model_revision, + ) + split_artifacts: dict[str, tuple[Path, Path]] = {} + for split, prompt_data, manifest in args.split_artifact: + if split in split_artifacts: + raise ValueError(f"duplicate split artifact binding: {split!r}") + split_artifacts[split] = (Path(prompt_data), Path(manifest)) + verify_prepared_artifacts( + prepare_lock_path=args.prepare_lock, + split_artifacts=split_artifacts, + max_steps=args.max_steps, + model_revision=args.model_revision, + alfworld_data_root=args.alfworld_data_root, + ) + print("GraphGPO prepared-artifact preflight passed.") + + +if __name__ == "__main__": + main() diff --git a/examples/graphgpo/prepare_alfworld.py b/examples/graphgpo/prepare_alfworld.py new file mode 100644 index 000000000..3303b42fd --- /dev/null +++ b/examples/graphgpo/prepare_alfworld.py @@ -0,0 +1,401 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Prepare deterministic ALFWorld manifests and Relax prompt rows. + +This module intentionally does not import ALFWorld. It operates on an already- +downloaded text-only dataset and can therefore run in a lightweight CPU +environment before the real ``reset``/``step`` smoke test. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import re +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from examples.graphgpo.manifest import ( + ALFWORLD_TASK_TYPES, + TaskManifest, + build_manifest, + discover_task_files, + manifest_bytes, + manifest_sha256, +) + + +PREPARE_SCHEMA_VERSION = "graphgpo-alfworld-prepare-v1" +PINNED_ALFWORLD_VERSION = "0.4.2" +PINNED_MODEL_REVISION = "775b11afaf83e0dc75bd5abaf90133e47b3ec082" +DEFAULT_SPLIT_PATHS = { + "train": Path("json_2.1.1/train"), + "eval_in_distribution": Path("json_2.1.1/valid_seen"), +} +DEFAULT_LIMITS = { + "train": None, + "eval_in_distribution": 128, +} +TRAIN_TEMPERATURE = 1.0 +EVAL_TEMPERATURE = 0.4 +SAMPLING_TOP_P = 1.0 +SAMPLING_MAX_TOKENS = 512 +_SPLIT_NAME = re.compile(r"^[a-z][a-z0-9_]*$") +_PLACEHOLDER_MESSAGE = ( + "Start the ALFWorld episode. The managed agent builds every trainable turn prompt from the environment state." +) + + +@dataclass(frozen=True) +class PreparedSplit: + """Files and hashes produced for one ALFWorld split.""" + + split: str + manifest_name: str + prompt_name: str + task_count: int + manifest_sha256: str + prompt_sha256: str + temperature: float + top_p: float + max_tokens: int + + +def _read_json_object(path: Path, *, label: str) -> dict[str, Any]: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"cannot read {label} JSON from {path}") from exc + if not isinstance(payload, dict): + raise ValueError(f"{label} must contain a JSON object: {path}") + return payload + + +def is_supported_solvable_task(game_path: Path) -> bool: + """Match the filters used by ALFWorld 0.4.2's text environment.""" + + normalized_path = game_path.as_posix().lower() + if "movable" in normalized_path or "sliced" in normalized_path: + return False + trajectory_path = game_path.with_name("traj_data.json") + trajectory = _read_json_object(trajectory_path, label="trajectory") + if trajectory.get("task_type") not in ALFWORLD_TASK_TYPES: + return False + game = _read_json_object(game_path, label="game") + return game.get("solvable") is True + + +def eligible_task_files(split_root: Path) -> list[Path]: + """Return supported, solvable games in stable relative-path order.""" + + split_root = split_root.resolve(strict=True) + if not split_root.is_dir(): + raise ValueError(f"split root must be a directory: {split_root}") + return [path for path in discover_task_files(split_root) if is_supported_solvable_task(path)] + + +def select_task_files(task_files: Sequence[Path], *, limit: int | None) -> list[Path]: + if limit is not None: + if isinstance(limit, bool) or not isinstance(limit, int) or limit <= 0: + raise ValueError("split limit must be a positive integer or null") + if len(task_files) < limit: + raise ValueError(f"requested {limit} tasks but only {len(task_files)} eligible tasks are available") + return list(task_files[:limit]) + return list(task_files) + + +def prompt_rows_bytes( + manifest: TaskManifest, + *, + max_steps: int, + temperature: float | None = None, + top_p: float = SAMPLING_TOP_P, + max_tokens: int = SAMPLING_MAX_TOKENS, +) -> bytes: + """Serialize one deterministic Relax seed row per manifest task.""" + + if isinstance(max_steps, bool) or not isinstance(max_steps, int) or max_steps <= 0: + raise ValueError("max_steps must be a positive integer") + if temperature is None: + temperature = TRAIN_TEMPERATURE if manifest.split == "train" else EVAL_TEMPERATURE + if ( + isinstance(temperature, bool) + or not isinstance(temperature, (int, float)) + or not math.isfinite(float(temperature)) + or not 0.0 <= float(temperature) + ): + raise ValueError("temperature must be a non-negative real number") + if ( + isinstance(top_p, bool) + or not isinstance(top_p, (int, float)) + or not math.isfinite(float(top_p)) + or not 0.0 < float(top_p) <= 1.0 + ): + raise ValueError("top_p must be in (0, 1]") + if isinstance(max_tokens, bool) or not isinstance(max_tokens, int) or max_tokens <= 0: + raise ValueError("max_tokens must be a positive integer") + temperature = float(temperature) + top_p = float(top_p) + digest = manifest_sha256(manifest) + rows: list[str] = [] + for index, task in enumerate(manifest.tasks): + row = { + "messages": [{"role": "user", "content": _PLACEHOLDER_MESSAGE}], + "metadata": { + "alfworld_train_eval": manifest.split, + "manifest_index": index, + "manifest_sha256": digest, + "manifest_task_id": task.game.relative_path, + "max_tokens": max_tokens, + "max_steps": max_steps, + "task_id": task.game.relative_path, + "task_type": task.task_type, + "temperature": temperature, + "top_p": top_p, + }, + } + rows.append( + json.dumps( + row, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + ) + if not rows: + raise ValueError(f"split {manifest.split!r} contains no eligible tasks") + return ("\n".join(rows) + "\n").encode("utf-8") + + +def _write_idempotent(path: Path, content: bytes) -> None: + if path.exists(): + if not path.is_file(): + raise FileExistsError(f"output path exists and is not a file: {path}") + if path.read_bytes() == content: + return + raise FileExistsError(f"refusing to replace non-identical prepared artifact: {path}") + path.write_bytes(content) + + +def _sha256_bytes(content: bytes) -> str: + return hashlib.sha256(content).hexdigest() + + +def prepare_split( + *, + data_root: Path, + output_dir: Path, + split: str, + split_root: Path, + shared_files: Sequence[Path], + limit: int | None, + max_steps: int, +) -> PreparedSplit: + if not isinstance(split, str) or _SPLIT_NAME.fullmatch(split) is None: + raise ValueError(f"invalid split name: {split!r}") + selected = select_task_files( + eligible_task_files(split_root), + limit=limit, + ) + manifest = build_manifest( + data_root, + split=split, + task_files=selected, + shared_files=shared_files, + ) + manifest_content = manifest_bytes(manifest) + prompt_content = prompt_rows_bytes(manifest, max_steps=max_steps) + temperature = TRAIN_TEMPERATURE if manifest.split == "train" else EVAL_TEMPERATURE + manifest_name = f"{split}.manifest.json" + prompt_name = f"{split}.prompts.jsonl" + _write_idempotent(output_dir / manifest_name, manifest_content) + _write_idempotent(output_dir / prompt_name, prompt_content) + return PreparedSplit( + split=split, + manifest_name=manifest_name, + prompt_name=prompt_name, + task_count=len(manifest.tasks), + manifest_sha256=_sha256_bytes(manifest_content), + prompt_sha256=_sha256_bytes(prompt_content), + temperature=temperature, + top_p=SAMPLING_TOP_P, + max_tokens=SAMPLING_MAX_TOKENS, + ) + + +def prepare_artifacts( + *, + data_root: Path, + output_dir: Path, + split_roots: Mapping[str, Path], + shared_files: Sequence[Path] | None = None, + limits: Mapping[str, int | None] | None = None, + max_steps: int = 50, +) -> dict[str, Any]: + """Build all split artifacts and a path-free content lock.""" + + data_root = data_root.resolve(strict=True) + if not data_root.is_dir(): + raise ValueError(f"data_root must be a directory: {data_root}") + if not split_roots: + raise ValueError("at least one split root is required") + output_dir.mkdir(parents=True, exist_ok=True) + if shared_files is None: + shared_files = ( + data_root / "logic" / "alfred.pddl", + data_root / "logic" / "alfred.twl2", + ) + resolved_shared = tuple(Path(path).resolve(strict=True) for path in shared_files) + resolved_limits = dict(limits or {}) + + prepared: list[PreparedSplit] = [] + for split in sorted(split_roots): + prepared.append( + prepare_split( + data_root=data_root, + output_dir=output_dir, + split=split, + split_root=Path(split_roots[split]), + shared_files=resolved_shared, + limit=resolved_limits.get(split), + max_steps=max_steps, + ) + ) + + lock = { + "schema_version": PREPARE_SCHEMA_VERSION, + "alfworld_version": PINNED_ALFWORLD_VERSION, + "model_revision": PINNED_MODEL_REVISION, + "max_steps": max_steps, + "splits": { + item.split: { + "manifest": item.manifest_name, + "manifest_sha256": item.manifest_sha256, + "prompt_data": item.prompt_name, + "prompt_data_sha256": item.prompt_sha256, + "sampling": { + "max_tokens": item.max_tokens, + "temperature": item.temperature, + "top_p": item.top_p, + }, + "task_count": item.task_count, + } + for item in prepared + }, + } + lock_content = ( + json.dumps( + lock, + ensure_ascii=False, + sort_keys=True, + indent=2, + allow_nan=False, + ) + + "\n" + ).encode("utf-8") + _write_idempotent(output_dir / "prepare.lock.json", lock_content) + return lock + + +def _assignment(value: str, *, option: str) -> tuple[str, str]: + name, separator, raw_value = value.partition("=") + if not separator or not name or not raw_value: + raise argparse.ArgumentTypeError(f"{option} must use NAME=VALUE syntax") + if _SPLIT_NAME.fullmatch(name) is None: + raise argparse.ArgumentTypeError(f"{option} has an invalid split name: {name!r}") + return name, raw_value + + +def _split_assignments( + raw_values: Sequence[str] | None, + *, + data_root: Path, +) -> dict[str, Path]: + if not raw_values: + return {name: data_root / relative_path for name, relative_path in DEFAULT_SPLIT_PATHS.items()} + result: dict[str, Path] = {} + for raw_value in raw_values: + name, value = _assignment(raw_value, option="--split") + if name in result: + raise ValueError(f"duplicate split assignment: {name!r}") + result[name] = Path(value) + return result + + +def _limit_assignments(raw_values: Sequence[str] | None) -> dict[str, int | None]: + result: dict[str, int | None] = {} + for raw_value in raw_values or (): + name, value = _assignment(raw_value, option="--limit") + if name in result: + raise ValueError(f"duplicate limit assignment: {name!r}") + try: + parsed = int(value) + except ValueError as exc: + raise ValueError(f"limit for {name!r} must be an integer") from exc + if parsed < 0: + raise ValueError(f"limit for {name!r} must be non-negative") + result[name] = None if parsed == 0 else parsed + return result + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Prepare deterministic ALFWorld manifests and GraphGPO prompt rows.") + parser.add_argument("--data-root", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument( + "--split", + action="append", + help=( + "Split mapping in NAME=PATH form. May be repeated. Defaults to " + "train and eval_in_distribution below --data-root." + ), + ) + parser.add_argument( + "--limit", + action="append", + help="Per-split task limit in NAME=COUNT form; COUNT=0 means all.", + ) + parser.add_argument( + "--shared-file", + action="append", + type=Path, + help=( + "Explicit alfred.pddl/alfred.twl2 path. Repeat twice. Defaults " + "to --data-root/logic/{alfred.pddl,alfred.twl2}." + ), + ) + parser.add_argument("--max-steps", type=int, default=50) + args = parser.parse_args(argv) + if args.shared_file is not None and len(args.shared_file) != 2: + parser.error("--shared-file must be omitted or supplied exactly twice") + return args + + +def main(argv: Sequence[str] | None = None) -> None: + args = parse_args(argv) + data_root = args.data_root.resolve(strict=True) + split_roots = _split_assignments(args.split, data_root=data_root) + explicit_limits = _limit_assignments(args.limit) + unknown_limits = set(explicit_limits) - set(split_roots) + if unknown_limits: + raise ValueError(f"limits refer to unknown splits: {sorted(unknown_limits)!r}") + limits = {split: DEFAULT_LIMITS.get(split) for split in split_roots} + limits.update(explicit_limits) + lock = prepare_artifacts( + data_root=data_root, + output_dir=args.output_dir, + split_roots=split_roots, + shared_files=args.shared_file, + limits=limits, + max_steps=args.max_steps, + ) + print(json.dumps(lock, ensure_ascii=False, sort_keys=True, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/examples/graphgpo/prompt.py b/examples/graphgpo/prompt.py new file mode 100644 index 000000000..028541d9e --- /dev/null +++ b/examples/graphgpo/prompt.py @@ -0,0 +1,139 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Prompt construction matching the frozen ALFWorld history policy.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Sequence + +from examples.graphgpo.state import ( + TrackerState, +) +from examples.graphgpo.state import ( + format_reference_observation as _format_reference_observation, +) + + +HISTORY_LENGTH = 2 + + +@dataclass(frozen=True) +class HistoryTurn: + """One action/observation pair shown in a later prompt.""" + + action: str + observation: str + + +def format_reference_observation( + raw_observation: str, + tracker: TrackerState, +) -> str: + """Expose the shared frozen-reference display formatter from this + module.""" + + return _format_reference_observation(raw_observation, tracker) + + +def _visible_commands(admissible_commands: Sequence[str]) -> list[str]: + if isinstance(admissible_commands, (str, bytes)) or not isinstance(admissible_commands, Sequence): + raise TypeError("admissible_commands must be a sequence of strings") + + commands: list[str] = [] + for command in admissible_commands: + if not isinstance(command, str): + raise TypeError("every admissible command must be a string") + if command != "help": + commands.append(command) + return commands + + +def _history_context( + history: Sequence[HistoryTurn], + *, + step_index: int, +) -> tuple[str, int]: + recent_history = history[-HISTORY_LENGTH:] + if step_index < len(recent_history): + raise ValueError("step_index cannot be smaller than the visible history") + start_index = step_index - len(recent_history) + lines = [ + ( + f"[Observation {start_index + offset + 1}: '{turn.observation}', " + f"Action {start_index + offset + 1}: '{turn.action}']" + ) + for offset, turn in enumerate(recent_history) + ] + return "\n".join(lines), len(recent_history) + + +def build_prompt( + *, + task_description: str, + raw_observation: str, + admissible_commands: Sequence[str], + history: Sequence[HistoryTurn] = (), + step_index: int | None = None, +) -> str: + """Build a per-turn prompt with exactly two recent history entries. + + The initial ALFWorld observation already contains the task description, so + the first prompt does not duplicate it. Later prompts include the task, + current step, the last two action/observation pairs, the current + observation, and commands in environment order. ``help`` is hidden only + from the prompt; the state anchor keeps it. + """ + + if not isinstance(task_description, str): + raise TypeError("task_description must be a string") + if not isinstance(raw_observation, str): + raise TypeError("raw_observation must be a string") + if isinstance(history, (str, bytes)) or not isinstance(history, Sequence): + raise TypeError("history must be a sequence of HistoryTurn values") + if any(not isinstance(turn, HistoryTurn) for turn in history): + raise TypeError("every history entry must be a HistoryTurn") + + if step_index is None: + step_index = len(history) + if isinstance(step_index, bool) or not isinstance(step_index, int) or step_index < 0: + raise ValueError("step_index must be a non-negative integer") + + commands = "\n ".join(f"'{command}'" for command in _visible_commands(admissible_commands)) + if history: + action_history, valid_history_length = _history_context( + history, + step_index=step_index, + ) + return ( + "You are an expert agent operating in the ALFRED Embodied " + f"Environment. Your task is to: {task_description}\n\n" + f"Prior to this step, you have already taken {step_index} step(s). " + f"Below are the most recent {valid_history_length} observations and " + f"the corresponding actions you took: {action_history}\n" + f"You are now at step {step_index + 1} and your current observation " + f"is: {raw_observation}\n\n" + "Your admissible actions of the current situation are: " + f"[{commands}].\n\n" + "Now it's your turn to take an action.\n\n" + "You should first reason step-by-step about the current situation. " + "This reasoning process MUST be enclosed within " + "tags.\n" + "Once you've finished your reasoning, you should choose an " + "admissible action for current step and present it within " + " tags." + ) + + return ( + "You are an expert agent operating in the ALFRED Embodied " + "Environment.\n\n" + f"Your current observation is: {raw_observation}\n\n" + "Your admissible actions of the current situation are: " + f"[{commands}].\n\n" + "Now it's your turn to take an action.\n\n" + "You should first reason step-by-step about the current situation. " + "This reasoning process MUST be enclosed within " + "tags.\n" + "Once you've finished your reasoning, you should choose an admissible " + "action for current step and present it within tags." + ) diff --git a/examples/graphgpo/reproducibility.py b/examples/graphgpo/reproducibility.py new file mode 100644 index 000000000..b7cb7cff2 --- /dev/null +++ b/examples/graphgpo/reproducibility.py @@ -0,0 +1,476 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Build auditable Task 37 reproducibility evidence without inventing facts. + +The command writes three deterministic JSON files: + +``seed_mapping.json`` + The registered four-condition, three-seed plan. Entries are explicitly + labelled as planned runs and are not execution evidence. +``test_report.json`` + Counts parsed from JUnit XML plus content hashes for complete text logs. + A text log is never interpreted as a passing result. +``reproducibility.manifest.json`` + Content hashes for supplied artifacts, the two generated files, declared + source revisions, and versions observed from the current interpreter. + +Unknown revisions remain JSON ``null`` with ``verification="not_claimed"``. +Operator-supplied revisions are labelled as declarations rather than as facts +verified by this generator. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.metadata +import json +import platform +import re +import xml.etree.ElementTree as ET +from collections.abc import Mapping, Sequence +from pathlib import Path + + +REPRODUCIBILITY_SCHEMA_VERSION = "task37-reproducibility-manifest-v1" +SEED_MAPPING_SCHEMA_VERSION = "task37-seed-mapping-v1" +TEST_REPORT_SCHEMA_VERSION = "task37-test-report-v1" + +DECLARED_VERSION_FIELDS = ( + "relax_baseline_commit", + "candidate_commit", + "graphgpo_reference_commit", + "container_image_digest", + "model_revision", + "alfworld_version", +) + +RUNTIME_DISTRIBUTIONS = ( + ("relax", "relax"), + ("torch", "torch"), + ("megatron_core", "megatron-core"), + ("ray", "ray"), + ("sglang", "sglang"), + ("transformers", "transformers"), + ("alfworld", "alfworld"), +) + +_COMMIT_FIELDS = frozenset(("relax_baseline_commit", "candidate_commit", "graphgpo_reference_commit")) +_FULL_COMMIT = re.compile(r"^[0-9a-f]{40}$") +_IMAGE_DIGEST = re.compile(r"^\S+@sha256:[0-9a-f]{64}$") + +_REGISTERED_CONDITIONS = ( + ( + "grpo", + { + "METHOD": "grpo", + "BETA_EPISODE": "1", + }, + ), + ( + "gigpo", + { + "METHOD": "gigpo", + "BETA": "1", + "BETA_EPISODE": "1", + }, + ), + ( + "graphgpo", + { + "METHOD": "graphgpo", + "BETA": "1", + "BETA_EPISODE": "1", + }, + ), + ( + "graph_only", + { + "METHOD": "graphgpo", + "BETA": "1", + "BETA_EPISODE": "0", + }, + ), +) + + +def _json_bytes(payload: object) -> bytes: + return ( + json.dumps( + payload, + ensure_ascii=False, + sort_keys=True, + indent=2, + allow_nan=False, + ) + + "\n" + ).encode("utf-8") + + +def _sha256_file(path: Path, *, chunk_size: int = 1024 * 1024) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(chunk_size), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _artifact_record(label: str, path: Path) -> dict[str, object]: + if not isinstance(label, str) or not label.strip(): + raise ValueError("artifact label must be a non-empty string") + resolved = path.resolve(strict=True) + if not resolved.is_file(): + raise ValueError(f"artifact must be a regular file: {path}") + return { + "label": label.strip(), + "file_name": resolved.name, + "sha256": _sha256_file(resolved), + "size_bytes": resolved.stat().st_size, + } + + +def _validate_declared_version(field: str, value: object) -> str: + if field not in DECLARED_VERSION_FIELDS: + raise ValueError(f"unknown version field {field!r}; expected one of {list(DECLARED_VERSION_FIELDS)!r}") + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{field} must be a non-empty string") + normalized = value.strip() + if field in _COMMIT_FIELDS and _FULL_COMMIT.fullmatch(normalized) is None: + raise ValueError(f"{field} must be a full lowercase 40-character commit") + if field == "container_image_digest" and _IMAGE_DIGEST.fullmatch(normalized) is None: + raise ValueError("container_image_digest must be an image reference pinned by a sha256 digest") + return normalized + + +def build_version_inventory( + declared_versions: Mapping[str, str] | None = None, +) -> dict[str, object]: + """Record declarations separately from versions observed at runtime.""" + + supplied = dict(declared_versions or {}) + unexpected = sorted(set(supplied) - set(DECLARED_VERSION_FIELDS)) + if unexpected: + raise ValueError(f"unknown version fields: {unexpected!r}") + + declarations: dict[str, dict[str, str | None]] = {} + for field in DECLARED_VERSION_FIELDS: + if field not in supplied: + declarations[field] = { + "value": None, + "source": None, + "verification": "not_claimed", + } + continue + declarations[field] = { + "value": _validate_declared_version(field, supplied[field]), + "source": "operator_supplied", + "verification": "not_verified_by_generator", + } + + packages: dict[str, dict[str, str | None]] = {} + for label, distribution in RUNTIME_DISTRIBUTIONS: + try: + version = importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + packages[label] = { + "value": None, + "source": f"importlib.metadata:{distribution}", + "verification": "not_installed_in_generator_runtime", + } + else: + packages[label] = { + "value": version, + "source": f"importlib.metadata:{distribution}", + "verification": "observed_in_generator_runtime", + } + + return { + "declared": declarations, + "runtime_observed": { + "python": { + "value": platform.python_version(), + "source": "platform.python_version", + "verification": "observed_in_generator_runtime", + }, + "operating_system": { + "value": platform.platform(), + "source": "platform.platform", + "verification": "observed_in_generator_runtime", + }, + "packages": packages, + "scope_note": ( + "Runtime observations describe only the interpreter running this generator; " + "they are not evidence for a remote training container." + ), + }, + } + + +def build_seed_mapping( + seeds: Sequence[int] = (0, 1, 2), + *, + episode_weighting: str = "trajectory_once", +) -> dict[str, object]: + """Return the registered 4 x N run grid without claiming it was + executed.""" + + if isinstance(seeds, (str, bytes)) or not isinstance(seeds, Sequence): + raise TypeError("seeds must be a sequence of non-negative integers") + normalized_seeds: list[int] = [] + for seed in seeds: + if isinstance(seed, bool) or not isinstance(seed, int) or seed < 0: + raise ValueError("every seed must be a non-negative integer") + normalized_seeds.append(seed) + if not normalized_seeds: + raise ValueError("at least one seed is required") + if len(set(normalized_seeds)) != len(normalized_seeds): + raise ValueError("seeds must not contain duplicates") + if episode_weighting not in {"trajectory_once", "reference_cross_steps"}: + raise ValueError("unknown episode weighting") + + runs: list[dict[str, object]] = [] + for condition, condition_environment in _REGISTERED_CONDITIONS: + for seed in normalized_seeds: + environment = { + **condition_environment, + "EPISODE_WEIGHTING": episode_weighting, + "SEED": str(seed), + } + runs.append( + { + "run_id": f"{condition}-seed-{seed}", + "condition": condition, + "registered_seed": seed, + "launcher_environment": environment, + "execution_status": "planned", + "execution_evidence": None, + } + ) + + return { + "schema_version": SEED_MAPPING_SCHEMA_VERSION, + "mapping_scope": "registered_plan_not_execution_evidence", + "seed_semantics": { + "declared_entrypoint": "launcher environment variable SEED", + "declared_training_argument": "--seed", + "qualification": ( + "The mapping records the registered launcher input only; it does not claim that " + "every third-party component consumes that seed." + ), + }, + "runs": runs, + } + + +def _xml_local_name(tag: str) -> str: + return tag.rsplit("}", 1)[-1] + + +def parse_junit_report(path: Path) -> dict[str, object]: + """Parse outcomes from test cases instead of trusting suite totals.""" + + resolved = path.resolve(strict=True) + if not resolved.is_file(): + raise ValueError(f"JUnit report must be a regular file: {path}") + root = ET.parse(resolved).getroot() + cases = [element for element in root.iter() if _xml_local_name(element.tag) == "testcase"] + + failures = 0 + errors = 0 + skipped = 0 + duration_seconds = 0.0 + for case in cases: + raw_duration = case.attrib.get("time", "0") + try: + duration = float(raw_duration) + except ValueError as exc: + raise ValueError(f"invalid JUnit testcase duration {raw_duration!r}") from exc + if duration < 0: + raise ValueError("JUnit testcase duration must be non-negative") + duration_seconds += duration + + child_names = {_xml_local_name(child.tag) for child in case} + if "error" in child_names: + errors += 1 + elif "failure" in child_names: + failures += 1 + elif "skipped" in child_names: + skipped += 1 + + tests = len(cases) + passed = tests - failures - errors - skipped + if tests == 0: + outcome = "empty" + elif failures or errors: + outcome = "failed" + else: + outcome = "passed" + return { + "tests": tests, + "passed": passed, + "failures": failures, + "errors": errors, + "skipped": skipped, + "duration_seconds": duration_seconds, + "outcome": outcome, + } + + +def build_test_report( + junit_reports: Mapping[str, Path] | None = None, + text_logs: Mapping[str, Path] | None = None, +) -> dict[str, object]: + junit_entries: list[dict[str, object]] = [] + for label, path in sorted((junit_reports or {}).items()): + junit_entries.append( + { + "artifact": _artifact_record(label, path), + "summary": parse_junit_report(path), + } + ) + + log_entries: list[dict[str, object]] = [] + for label, path in sorted((text_logs or {}).items()): + log_entries.append( + { + "artifact": _artifact_record(label, path), + "outcome": "not_inferred_from_unstructured_log", + } + ) + + junit_outcomes = [entry["summary"]["outcome"] for entry in junit_entries] + if not junit_outcomes: + overall_outcome = "not_evaluated" + elif any(outcome == "failed" for outcome in junit_outcomes): + overall_outcome = "failed" + elif any(outcome == "empty" for outcome in junit_outcomes): + overall_outcome = "incomplete" + else: + overall_outcome = "passed" + return { + "schema_version": TEST_REPORT_SCHEMA_VERSION, + "overall_outcome_from_junit_only": overall_outcome, + "junit_reports": junit_entries, + "complete_text_logs": log_entries, + } + + +def _write_idempotent(path: Path, content: bytes) -> None: + if path.exists(): + if not path.is_file(): + raise ValueError(f"output path is not a regular file: {path}") + if path.read_bytes() != content: + raise FileExistsError(f"refusing to replace different evidence file: {path}") + return + path.write_bytes(content) + + +def write_reproducibility_bundle( + output_dir: Path, + *, + seeds: Sequence[int] = (0, 1, 2), + episode_weighting: str = "trajectory_once", + artifacts: Mapping[str, Path] | None = None, + junit_reports: Mapping[str, Path] | None = None, + text_logs: Mapping[str, Path] | None = None, + declared_versions: Mapping[str, str] | None = None, +) -> dict[str, Path]: + """Write a seed plan, parsed test report, and content-addressed + manifest.""" + + output_dir = output_dir.resolve() + output_dir.mkdir(parents=True, exist_ok=True) + if not output_dir.is_dir(): + raise ValueError("output_dir must be a directory") + + seed_path = output_dir / "seed_mapping.json" + test_path = output_dir / "test_report.json" + manifest_path = output_dir / "reproducibility.manifest.json" + + seed_payload = build_seed_mapping(seeds, episode_weighting=episode_weighting) + test_payload = build_test_report(junit_reports, text_logs) + _write_idempotent(seed_path, _json_bytes(seed_payload)) + _write_idempotent(test_path, _json_bytes(test_payload)) + + artifact_entries = [_artifact_record(label, path) for label, path in sorted((artifacts or {}).items())] + manifest_payload = { + "schema_version": REPRODUCIBILITY_SCHEMA_VERSION, + "claim_scope": { + "seed_mapping": "plan_only", + "test_results": "parsed_only_from_supplied_junit", + "text_logs": "content_addressed_without_outcome_inference", + "versions": "missing_values_are_not_claimed", + }, + "versions": build_version_inventory(declared_versions), + "input_artifacts": artifact_entries, + "generated_artifacts": [ + _artifact_record("seed_mapping", seed_path), + _artifact_record("test_report", test_path), + ], + } + _write_idempotent(manifest_path, _json_bytes(manifest_payload)) + return { + "seed_mapping": seed_path, + "test_report": test_path, + "manifest": manifest_path, + } + + +def _assignments(values: Sequence[str], *, kind: str) -> dict[str, str]: + result: dict[str, str] = {} + for value in values: + if "=" not in value: + raise ValueError(f"{kind} must use LABEL=VALUE syntax: {value!r}") + label, assigned_value = value.split("=", 1) + label = label.strip() + assigned_value = assigned_value.strip() + if not label or not assigned_value: + raise ValueError(f"{kind} must use non-empty LABEL=VALUE syntax") + if label in result: + raise ValueError(f"duplicate {kind} label: {label!r}") + result[label] = assigned_value + return result + + +def build_argument_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--seed", action="append", type=int, dest="seeds") + parser.add_argument( + "--episode-weighting", + choices=("trajectory_once", "reference_cross_steps"), + default="trajectory_once", + ) + parser.add_argument("--artifact", action="append", default=[], metavar="LABEL=PATH") + parser.add_argument("--junit", action="append", default=[], metavar="LABEL=PATH") + parser.add_argument("--log", action="append", default=[], metavar="LABEL=PATH") + parser.add_argument( + "--version", + action="append", + default=[], + metavar="FIELD=VALUE", + help=f"declared version; FIELD is one of {', '.join(DECLARED_VERSION_FIELDS)}", + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = build_argument_parser().parse_args(argv) + artifact_values = _assignments(args.artifact, kind="artifact") + junit_values = _assignments(args.junit, kind="junit") + log_values = _assignments(args.log, kind="log") + version_values = _assignments(args.version, kind="version") + write_reproducibility_bundle( + args.output_dir, + seeds=args.seeds or (0, 1, 2), + episode_weighting=args.episode_weighting, + artifacts={label: Path(path) for label, path in artifact_values.items()}, + junit_reports={label: Path(path) for label, path in junit_values.items()}, + text_logs={label: Path(path) for label, path in log_values.items()}, + declared_versions=version_values, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/graphgpo/requirements-alfworld.txt b/examples/graphgpo/requirements-alfworld.txt new file mode 100644 index 000000000..6ca94b200 --- /dev/null +++ b/examples/graphgpo/requirements-alfworld.txt @@ -0,0 +1,4 @@ +alfworld==0.4.2 +httpx==0.28.1 +openai==2.46.0 +PyYAML==6.0.3 diff --git a/examples/graphgpo/reward.py b/examples/graphgpo/reward.py new file mode 100644 index 000000000..c0e9611e6 --- /dev/null +++ b/examples/graphgpo/reward.py @@ -0,0 +1,37 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Reward passthrough for ALFWorld explicit turn rows.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from numbers import Real +from typing import Any + + +async def reward_func( + _args: Any, + samples: Sequence[Any], + **_kwargs: Any, +) -> list[float]: + """Return the environment reward already attached to every turn row. + + Agentic group-RM evaluation always invokes the configured batched reward + function, even when the managed agent exported a reward. GraphGPO's + environment is the reward source, so this adapter validates and forwards + those values without scoring the model response a second time. + """ + + if isinstance(samples, (str, bytes)) or not isinstance(samples, Sequence): + raise TypeError("samples must be a sequence") + if not samples: + raise ValueError("samples must not be empty") + + rewards: list[float] = [] + for index, sample in enumerate(samples): + reward = getattr(sample, "reward", None) + if isinstance(reward, bool) or not isinstance(reward, Real) or not math.isfinite(float(reward)): + raise ValueError(f"samples[{index}].reward must be a finite real number") + rewards.append(float(reward)) + return rewards diff --git a/examples/graphgpo/rollout_agent.py b/examples/graphgpo/rollout_agent.py new file mode 100644 index 000000000..cc6941c18 --- /dev/null +++ b/examples/graphgpo/rollout_agent.py @@ -0,0 +1,605 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Managed GraphGPO ALFWorld agent with explicit per-turn exports.""" + +from __future__ import annotations + +import argparse +import asyncio +import hashlib +import json +import math +import os +from collections.abc import Mapping, MutableMapping, Sequence +from pathlib import Path, PurePosixPath +from typing import Any, Protocol + +from examples.graphgpo.action_parser import parse_action +from examples.graphgpo.alfworld_env import AlfWorldSnapshot, AlfWorldTextEnv +from examples.graphgpo.graph_credit import SUCCESS +from examples.graphgpo.prompt import ( + HistoryTurn, + build_prompt, + format_reference_observation, +) +from examples.graphgpo.state import state_key_v1 + + +TASK_MARKER = "Your task is to: " +DEFAULT_MAX_STEPS = 50 +SUCCESS_REWARD = 10.0 +INVALID_PENALTY = 0.1 + + +class ChatCompletionClient(Protocol): + """Minimal injectable interface used by :func:`run_episode`.""" + + async def complete(self, *, messages: list[dict[str, Any]]) -> str: + """Return one assistant message for a fresh per-turn conversation.""" + + +class OpenAICompatibleChatClient: + """Lazy OpenAI SDK adapter for Relax's chat-completions service.""" + + def __init__( + self, + *, + base_url: str, + api_key: str, + model: str = "model", + timeout_s: float = 1200.0, + temperature: float = 1.0, + top_p: float = 1.0, + max_tokens: int = 512, + ) -> None: + if not isinstance(base_url, str) or not base_url: + raise ValueError("base_url must be a non-empty string") + if not isinstance(api_key, str) or not api_key: + raise ValueError("api_key must be a non-empty string") + if not isinstance(model, str) or not model: + raise ValueError("model must be a non-empty string") + if isinstance(timeout_s, bool) or not isinstance(timeout_s, (int, float)): + raise ValueError("timeout_s must be a positive number") + if timeout_s <= 0: + raise ValueError("timeout_s must be a positive number") + if ( + isinstance(temperature, bool) + or not isinstance(temperature, (int, float)) + or not math.isfinite(float(temperature)) + or temperature < 0 + ): + raise ValueError("temperature must be a finite non-negative number") + if ( + isinstance(top_p, bool) + or not isinstance(top_p, (int, float)) + or not math.isfinite(float(top_p)) + or not 0 < top_p <= 1 + ): + raise ValueError("top_p must be in (0, 1]") + if isinstance(max_tokens, bool) or not isinstance(max_tokens, int) or max_tokens <= 0: + raise ValueError("max_tokens must be a positive integer") + self._base_url = base_url.rstrip("/") + self._api_key = api_key + self._model = model + self._timeout_s = float(timeout_s) + self._temperature = float(temperature) + self._top_p = float(top_p) + self._max_tokens = max_tokens + self._client: Any | None = None + + def _ensure_client(self) -> Any: + if self._client is None: + try: + import httpx + from openai import AsyncOpenAI + except ImportError as exc: + raise RuntimeError("the OpenAI SDK is required by the managed GraphGPO agent") from exc + self._client = AsyncOpenAI( + api_key=self._api_key, + base_url=self._base_url, + timeout=httpx.Timeout(self._timeout_s, connect=30.0), + max_retries=0, + ) + return self._client + + async def complete(self, *, messages: list[dict[str, Any]]) -> str: + response = await self._ensure_client().chat.completions.create( + model=self._model, + messages=messages, + temperature=self._temperature, + top_p=self._top_p, + max_tokens=self._max_tokens, + ) + if not response.choices: + raise RuntimeError("chat completion returned no choices") + content = response.choices[0].message.content + if not isinstance(content, str): + raise RuntimeError("chat completion returned non-text content") + return content + + async def close(self) -> None: + if self._client is None: + return + client = self._client + self._client = None + close_method = getattr(client, "close", None) + if callable(close_method): + await close_method() + + +def stable_group_seed(group_id: str) -> int: + """Derive the same ALFWorld seed for every slot in one rollout group.""" + + if not isinstance(group_id, str) or not group_id: + raise ValueError("group_id must be a non-empty string") + digest = hashlib.sha256(group_id.encode("utf-8")).digest() + return int.from_bytes(digest[:8], byteorder="big") % (2**31) + + +def stable_row_id( + *, + task_id: object, + rollout_group_id: str, + trajectory_id: str, + turn_index: int, +) -> str: + """Return a deterministic identity for one explicit GraphGPO turn.""" + + if isinstance(task_id, bool) or not isinstance(task_id, (str, int)): + raise ValueError("task_id must be a JSON string or integer") + if isinstance(task_id, str) and not task_id: + raise ValueError("task_id must not be empty") + if not isinstance(rollout_group_id, str) or not rollout_group_id: + raise ValueError("rollout_group_id must be a non-empty string") + if not isinstance(trajectory_id, str) or not trajectory_id: + raise ValueError("trajectory_id must be a non-empty string") + if isinstance(turn_index, bool) or not isinstance(turn_index, int) or turn_index < 0: + raise ValueError("turn_index must be a non-negative integer") + encoded = json.dumps( + ["graphgpo-row-v1", task_id, rollout_group_id, trajectory_id, turn_index], + ensure_ascii=False, + allow_nan=False, + separators=(",", ":"), + ).encode("utf-8") + return f"graphgpo-row-v1:{hashlib.sha256(encoded).hexdigest()}" + + +def extract_task_description(initial_observation: str) -> str: + """Extract the frozen ALFWorld task suffix or fail explicitly.""" + + if not isinstance(initial_observation, str): + raise TypeError("initial_observation must be a string") + marker_index = initial_observation.find(TASK_MARKER) + if marker_index < 0: + raise ValueError("Task description not found in text observation.") + task_description = initial_observation[marker_index + len(TASK_MARKER) :].strip() + if not task_description: + raise ValueError("Task description is empty in text observation.") + return task_description + + +def _required_string( + name: str, + *, + explicit: str | None, + environment_key: str, + environ: Mapping[str, str] | None = None, +) -> str: + value = explicit + if value is None: + source = os.environ if environ is None else environ + value = source.get(environment_key) + if not isinstance(value, str) or not value: + raise ValueError(f"{name} is required; pass it explicitly or set {environment_key}") + return value + + +def _task_id_from_snapshot( + snapshot: AlfWorldSnapshot, + *, + fallback_task_id: object | None, + data_root: str | None, +) -> object: + if fallback_task_id is not None: + if isinstance(fallback_task_id, bool) or not isinstance(fallback_task_id, (str, int)): + raise ValueError("fallback task_id must be a JSON string or integer") + if isinstance(fallback_task_id, str) and not fallback_task_id: + raise ValueError("fallback task_id must not be empty") + + if snapshot.gamefile: + normalized_gamefile = snapshot.gamefile.replace("\\", "/") + gamefile_path = PurePosixPath(normalized_gamefile) + if ".." in gamefile_path.parts: + raise ValueError("ALFWorld gamefile must not contain parent traversal") + if gamefile_path.is_absolute(): + if not data_root: + raise ValueError("ALFWORLD_DATA is required to normalize an absolute gamefile") + normalized_root = data_root.replace("\\", "/") + data_root_path = PurePosixPath(normalized_root) + if not data_root_path.is_absolute(): + raise ValueError("ALFWORLD_DATA must be absolute when ALFWorld returns an absolute gamefile") + try: + gamefile_path = gamefile_path.relative_to(data_root_path) + except ValueError as exc: + raise ValueError("ALFWorld gamefile is outside ALFWORLD_DATA") from exc + resolved_task_id = gamefile_path.as_posix() + if fallback_task_id is not None: + declared_task_id = ( + fallback_task_id.replace("\\", "/") if isinstance(fallback_task_id, str) else fallback_task_id + ) + if declared_task_id != resolved_task_id: + raise ValueError("ALFWorld reset gamefile does not match the declared task_id") + return resolved_task_id + if fallback_task_id is None: + raise ValueError( + "task_id is unavailable: reset info has no extra.gamefile and session metadata has no explicit task_id" + ) + return fallback_task_id + + +def _require_max_steps(max_steps: int) -> int: + if isinstance(max_steps, bool) or not isinstance(max_steps, int) or max_steps <= 0: + raise ValueError("max_steps must be a positive integer") + return max_steps + + +def _sampling_float( + metadata: Mapping[str, Any], + source: Mapping[str, str], + *, + metadata_key: str, + environment_key: str, + default: float, + minimum_exclusive: float | None = None, + maximum: float | None = None, +) -> float: + raw: object = metadata.get( + metadata_key, + source.get(environment_key, default), + ) + if isinstance(raw, str): + try: + raw = float(raw) + except ValueError as exc: + raise ValueError(f"{environment_key} must be a real number") from exc + if isinstance(raw, bool) or not isinstance(raw, (int, float)) or not math.isfinite(float(raw)): + raise ValueError(f"{metadata_key} must be a finite real number") + value = float(raw) + if minimum_exclusive is not None and value <= minimum_exclusive: + raise ValueError(f"{metadata_key} must be greater than {minimum_exclusive}") + if minimum_exclusive is None and value < 0: + raise ValueError(f"{metadata_key} must be non-negative") + if maximum is not None and value > maximum: + raise ValueError(f"{metadata_key} must be at most {maximum}") + return value + + +def _sampling_max_tokens( + metadata: Mapping[str, Any], + source: Mapping[str, str], +) -> int: + raw: object = metadata.get( + "max_tokens", + source.get("GRAPHGPO_MAX_TOKENS", 512), + ) + if isinstance(raw, str): + try: + raw = int(raw) + except ValueError as exc: + raise ValueError("GRAPHGPO_MAX_TOKENS must be an integer") from exc + if isinstance(raw, bool) or not isinstance(raw, int) or raw <= 0: + raise ValueError("max_tokens must be a positive integer") + return raw + + +def _finalize_records( + records: list[dict[str, Any]], + *, + success: bool, + invalid_count: int, +) -> list[dict[str, Any]]: + if not records: + raise RuntimeError("an ALFWorld episode must export at least one turn") + episode_return = SUCCESS_REWARD * float(success) - INVALID_PENALTY * invalid_count + for record in records: + record["metadata"]["success"] = success + record["metadata"]["episode_return"] = episode_return + record["reward"] = episode_return + if success: + last_metadata = records[-1]["metadata"] + last_metadata["next_state_key"] = SUCCESS + last_metadata["terminal"] = True + last_metadata["truncated"] = False + return records + + +async def run_episode( + *, + chat_client: ChatCompletionClient, + env: AlfWorldTextEnv, + trajectory_id: str | None = None, + rollout_group_id: str | None = None, + fallback_task_id: object | None = None, + max_steps: int = DEFAULT_MAX_STEPS, + environ: Mapping[str, str] | None = None, +) -> list[dict[str, Any]]: + """Run one trajectory and return JSONL-ready explicit turn records.""" + + records: list[dict[str, Any]] = [] + history: list[HistoryTurn] = [] + invalid_count = 0 + final_success = False + source = os.environ if environ is None else environ + + try: + max_steps = _require_max_steps(max_steps) + resolved_trajectory_id = _required_string( + "trajectory_id", + explicit=trajectory_id, + environment_key="RELAX_SESSION_ID", + environ=environ, + ) + snapshot = env.reset() + task_description = extract_task_description(snapshot.raw_observation) + task_id = _task_id_from_snapshot( + snapshot, + fallback_task_id=fallback_task_id, + data_root=source.get("ALFWORLD_DATA"), + ) + if rollout_group_id is not None and (not isinstance(rollout_group_id, str) or not rollout_group_id): + raise ValueError("rollout_group_id must be a non-empty string") + + for turn_index in range(max_steps): + state_key = state_key_v1( + snapshot.raw_observation, + snapshot.tracker, + snapshot.admissible_commands, + ) + prompt = build_prompt( + task_description=task_description, + raw_observation=format_reference_observation( + snapshot.raw_observation, + snapshot.tracker, + ), + admissible_commands=snapshot.admissible_commands, + history=history, + step_index=turn_index, + ) + request_messages = [{"role": "user", "content": prompt}] + response_text = await chat_client.complete(messages=request_messages) + if not isinstance(response_text, str): + raise TypeError("chat_client.complete() must return a string") + parsed_action = parse_action(response_text) + if not parsed_action.is_valid: + invalid_count += 1 + + next_snapshot = env.step(parsed_action.action) + final_success = bool(next_snapshot.won) + reached_horizon = turn_index + 1 == max_steps + terminal = bool(final_success or (next_snapshot.done and not reached_horizon)) + truncated = bool(reached_horizon and not final_success) + next_state_key = ( + SUCCESS + if final_success + else state_key_v1( + next_snapshot.raw_observation, + next_snapshot.tracker, + next_snapshot.admissible_commands, + ) + ) + turn_id = f"turn_{turn_index:03d}" + identity_metadata: dict[str, Any] = {} + if rollout_group_id is not None: + identity_metadata = { + "rollout_group_id": rollout_group_id, + "turn_id": turn_id, + "row_id": stable_row_id( + task_id=task_id, + rollout_group_id=rollout_group_id, + trajectory_id=resolved_trajectory_id, + turn_index=turn_index, + ), + } + records.append( + { + "name": turn_id, + "messages": [ + dict(request_messages[0]), + {"role": "assistant", "content": response_text}, + ], + "metadata": { + "task_id": task_id, + "trajectory_id": resolved_trajectory_id, + "turn_index": turn_index, + "state_key": state_key, + "action": parsed_action.action, + "next_state_key": next_state_key, + "is_action_valid": parsed_action.is_valid, + "success": False, + "terminal": terminal, + "truncated": truncated, + "episode_return": 0.0, + **identity_metadata, + }, + "reward": 0.0, + } + ) + history.append( + HistoryTurn( + action=parsed_action.action, + observation=format_reference_observation( + next_snapshot.raw_observation, + next_snapshot.tracker, + ), + ) + ) + snapshot = next_snapshot + if terminal or truncated: + break + finally: + env.close() + + return _finalize_records( + records, + success=final_success, + invalid_count=invalid_count, + ) + + +def read_session_input(path: str | Path) -> dict[str, Any]: + payload = json.loads(Path(path).read_text(encoding="utf-8")) + if not isinstance(payload, dict): + raise TypeError("session input must be a JSON object") + metadata = payload.get("metadata", {}) + if not isinstance(metadata, dict): + raise TypeError("session input metadata must be a JSON object") + return payload + + +def write_session_output( + path: str | Path, + records: Sequence[Mapping[str, Any]], +) -> None: + """Write explicit records as JSONL, never as a JSON array.""" + + if isinstance(records, (str, bytes)) or not isinstance(records, Sequence): + raise TypeError("records must be a sequence of mappings") + lines: list[str] = [] + for record in records: + if not isinstance(record, Mapping): + raise TypeError("every output record must be a mapping") + lines.append( + json.dumps( + dict(record), + ensure_ascii=False, + allow_nan=False, + separators=(",", ":"), + ) + ) + if not lines: + raise ValueError("at least one explicit record is required") + try: + Path(path).write_text("\n".join(lines) + "\n", encoding="utf-8") + except FileNotFoundError: + # The managed runtime can remove its temporary directory after timeout. + pass + + +def _session_metadata(payload: Mapping[str, Any]) -> MutableMapping[str, Any]: + metadata = payload.get("metadata", {}) + if not isinstance(metadata, MutableMapping): + raise TypeError("session input metadata must be a mutable JSON object") + return metadata + + +async def run_managed_session( + *, + session_input: Mapping[str, Any], + chat_client: ChatCompletionClient | None = None, + env: AlfWorldTextEnv | None = None, + group_id: str | None = None, + trajectory_id: str | None = None, + task_id: object | None = None, + environ: Mapping[str, str] | None = None, +) -> list[dict[str, Any]]: + """Resolve managed-process inputs and run one ALFWorld episode.""" + + metadata = _session_metadata(session_input) + resolved_group_id = _required_string( + "group_id", + explicit=group_id, + environment_key="RELAX_GROUP_ID", + environ=environ, + ) + resolved_trajectory_id = _required_string( + "trajectory_id", + explicit=trajectory_id, + environment_key="RELAX_SESSION_ID", + environ=environ, + ) + source = os.environ if environ is None else environ + + owns_client = chat_client is None + if chat_client is None: + chat_client = OpenAICompatibleChatClient( + api_key=source.get("OPENAI_API_KEY") or source.get("RELAX_SESSION_ID", ""), + base_url=source.get("OPENAI_BASE_URL") or source.get("RELAX_BASE_URL", ""), + model=source.get("OPENAI_MODEL", "model"), + temperature=_sampling_float( + metadata, + source, + metadata_key="temperature", + environment_key="GRAPHGPO_TEMPERATURE", + default=1.0, + ), + top_p=_sampling_float( + metadata, + source, + metadata_key="top_p", + environment_key="GRAPHGPO_TOP_P", + default=1.0, + minimum_exclusive=0.0, + maximum=1.0, + ), + max_tokens=_sampling_max_tokens(metadata, source), + ) + + fallback_task_id = task_id if task_id is not None else metadata.get("task_id") + if env is None: + config_path = metadata.get("alfworld_config_path") or source.get("ALFWORLD_CONFIG_PATH") + if not isinstance(config_path, str) or not config_path: + raise ValueError("ALFWorld config is required in metadata.alfworld_config_path or ALFWORLD_CONFIG_PATH") + train_eval = metadata.get("alfworld_train_eval", "train") + if not isinstance(train_eval, str) or not train_eval: + raise ValueError("metadata.alfworld_train_eval must be a string") + env = AlfWorldTextEnv( + config_path=config_path, + seed=stable_group_seed(resolved_group_id), + train_eval=train_eval, + game_file=(fallback_task_id if isinstance(fallback_task_id, str) else None), + ) + + max_steps = metadata.get("max_steps", DEFAULT_MAX_STEPS) + try: + return await run_episode( + chat_client=chat_client, + env=env, + trajectory_id=resolved_trajectory_id, + rollout_group_id=resolved_group_id, + fallback_task_id=fallback_task_id, + max_steps=max_steps, + environ=environ, + ) + finally: + if owns_client and isinstance(chat_client, OpenAICompatibleChatClient): + await chat_client.close() + + +def parse_args( + argv: Sequence[str] | None = None, + *, + environ: Mapping[str, str] | None = None, +) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="GraphGPO ALFWorld managed agent session.") + parser.add_argument("--input-json") + parser.add_argument("--output-json") + args = parser.parse_args(argv) + source = os.environ if environ is None else environ + args.input_json = args.input_json or source.get("RELAX_INPUT_JSON") + args.output_json = args.output_json or source.get("RELAX_OUTPUT_JSON") + if not args.input_json: + parser.error("--input-json or the RELAX_INPUT_JSON environment variable is required") + if not args.output_json: + parser.error("--output-json or the RELAX_OUTPUT_JSON environment variable is required") + return args + + +def main() -> None: + args = parse_args() + session_input = read_session_input(args.input_json) + records = asyncio.run(run_managed_session(session_input=session_input)) + write_session_output(args.output_json, records) + + +if __name__ == "__main__": + main() diff --git a/examples/graphgpo/run_agent_app.sh b/examples/graphgpo/run_agent_app.sh new file mode 100644 index 000000000..5cdb37f43 --- /dev/null +++ b/examples/graphgpo/run_agent_app.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +PROJECT_ROOT="$(cd -- "${SCRIPT_DIR}/../.." >/dev/null 2>&1 && pwd)" + +export PYTHONPATH="${PROJECT_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" +export OPENAI_BASE_URL="${RELAX_BASE_URL:?RELAX_BASE_URL is required}" +export OPENAI_API_KEY="${RELAX_SESSION_ID:?RELAX_SESSION_ID is required}" + +exec "${ALFWORLD_PYTHON:-python3}" -m examples.graphgpo.rollout_agent \ + --input-json "${RELAX_INPUT_JSON:?RELAX_INPUT_JSON is required}" \ + --output-json "${RELAX_OUTPUT_JSON:?RELAX_OUTPUT_JSON is required}" diff --git a/examples/graphgpo/run_alfworld_qwen2_5_1_5b.sh b/examples/graphgpo/run_alfworld_qwen2_5_1_5b.sh new file mode 100644 index 000000000..8793a8122 --- /dev/null +++ b/examples/graphgpo/run_alfworld_qwen2_5_1_5b.sh @@ -0,0 +1,358 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +PROJECT_ROOT="$(cd -- "${SCRIPT_DIR}/../.." >/dev/null 2>&1 && pwd)" +[ -f "${SCRIPT_DIR}/env.sh" ] && source "${SCRIPT_DIR}/env.sh" + +METHOD="${METHOD:-graphgpo}" +case "${METHOD}" in + grpo|gigpo|graphgpo) ;; + *) + echo "ERROR: METHOD must be grpo, gigpo, or graphgpo (got ${METHOD})." >&2 + exit 2 + ;; +esac +ENABLE_EVAL="${ENABLE_EVAL:-1}" +case "${ENABLE_EVAL}" in + 0|1) ;; + *) + echo "ERROR: ENABLE_EVAL must be 0 or 1 (got ${ENABLE_EVAL})." >&2 + exit 2 + ;; +esac +EPISODE_WEIGHTING="${EPISODE_WEIGHTING:-trajectory_once}" +case "${EPISODE_WEIGHTING}" in + reference_cross_steps|trajectory_once) ;; + *) + echo "ERROR: EPISODE_WEIGHTING must be reference_cross_steps or trajectory_once (got ${EPISODE_WEIGHTING})." >&2 + exit 2 + ;; +esac + +SEED="${SEED:-0}" +NUM_GPUS="${NUM_GPUS:-2}" +OUTER_EPOCHS="${OUTER_EPOCHS:-150}" +TASK_GROUPS="${TASK_GROUPS:-16}" +GROUP_SIZE="${GROUP_SIZE:-8}" +GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-128}" +MAX_STEPS="${MAX_STEPS:-50}" +MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-32768}" +SGLANG_MEM_FRACTION_STATIC="${SGLANG_MEM_FRACTION_STATIC:-0.50}" +MODEL_REVISION="${MODEL_REVISION:-775b11afaf83e0dc75bd5abaf90133e47b3ec082}" +EXPECTED_EXTERNAL_IMAGE_DIGEST="${EXPECTED_EXTERNAL_IMAGE_DIGEST:-ghcr.io/redai-infra/relaxrl@sha256:3fa8ce578acda6c829b83016bde42c38fa892681e4f36ca330f545616fe578e2}" + +positive_integer_values=(TASK_GROUPS GROUP_SIZE GLOBAL_BATCH_SIZE MAX_STEPS) +for name in "${positive_integer_values[@]}"; do + if [[ ! "${!name}" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: ${name} must be a positive integer (got ${!name})." >&2 + exit 2 + fi +done +if [[ ! "${SGLANG_MEM_FRACTION_STATIC}" =~ ^(0\.[0-9]*[1-9][0-9]*|1(\.0+)?)$ ]]; then + echo "ERROR: SGLANG_MEM_FRACTION_STATIC must be greater than 0 and at most 1 (got ${SGLANG_MEM_FRACTION_STATIC})." >&2 + exit 2 +fi +ROLLOUT_SAMPLE_COUNT=$((TASK_GROUPS * GROUP_SIZE)) +if ((ROLLOUT_SAMPLE_COUNT % GLOBAL_BATCH_SIZE != 0)); then + echo "ERROR: TASK_GROUPS * GROUP_SIZE must be divisible by GLOBAL_BATCH_SIZE." >&2 + exit 2 +fi + +ALFWORLD_CONFIG_PATH="${ALFWORLD_CONFIG_PATH:-${SCRIPT_DIR}/configs/alfworld_qwen2_5_1_5b.yaml}" +DATA_ARTIFACT_DIR="${DATA_ARTIFACT_DIR:-}" +TRAIN_DATA="${TRAIN_DATA:-${DATA_ARTIFACT_DIR:+${DATA_ARTIFACT_DIR}/train.prompts.jsonl}}" +EVAL_DATA="${EVAL_DATA:-${DATA_ARTIFACT_DIR:+${DATA_ARTIFACT_DIR}/eval_in_distribution.prompts.jsonl}}" +TRAIN_MANIFEST="${TRAIN_MANIFEST:-${DATA_ARTIFACT_DIR:+${DATA_ARTIFACT_DIR}/train.manifest.json}}" +EVAL_MANIFEST="${EVAL_MANIFEST:-${DATA_ARTIFACT_DIR:+${DATA_ARTIFACT_DIR}/eval_in_distribution.manifest.json}}" +PREPARE_LOCK="${PREPARE_LOCK:-${DATA_ARTIFACT_DIR:+${DATA_ARTIFACT_DIR}/prepare.lock.json}}" +DEPENDENCY_LOCK="${DEPENDENCY_LOCK:-}" +MODEL_LOCK="${MODEL_LOCK:-}" +ALFWORLD_PYTHON="${ALFWORLD_PYTHON:-python3}" +HF_CHECKPOINT="${HF_CHECKPOINT:-}" +SAVE_DIR="${SAVE_DIR:-}" +LOAD_DIR="${LOAD_DIR:-}" + +if [ "${ENABLE_EVAL}" = "1" ] && { [ -z "${EVAL_DATA}" ] || [ -z "${EVAL_MANIFEST}" ]; }; then + echo "ERROR: EVAL_DATA and EVAL_MANIFEST must be set when ENABLE_EVAL=1." >&2 + exit 2 +fi + +required_values=( + ALFWORLD_DATA + HF_CHECKPOINT + SAVE_DIR + TRAIN_DATA + TRAIN_MANIFEST + PREPARE_LOCK + DEPENDENCY_LOCK + MODEL_LOCK +) +for name in "${required_values[@]}"; do + if [ -z "${!name:-}" ]; then + echo "ERROR: ${name} must be set." >&2 + exit 2 + fi +done + +required_files=( + "${ALFWORLD_CONFIG_PATH}" + "${TRAIN_DATA}" + "${TRAIN_MANIFEST}" + "${PREPARE_LOCK}" + "${DEPENDENCY_LOCK}" + "${MODEL_LOCK}" +) +if [ "${ENABLE_EVAL}" = "1" ]; then + required_files+=("${EVAL_DATA}" "${EVAL_MANIFEST}") +fi +for path in "${required_files[@]}"; do + if [ ! -f "${path}" ]; then + echo "ERROR: required file not found: ${path}" >&2 + exit 2 + fi +done +if [ ! -d "${HF_CHECKPOINT}" ]; then + echo "ERROR: HF_CHECKPOINT must be a local model snapshot directory: ${HF_CHECKPOINT}" >&2 + exit 2 +fi +if [ -n "${LOAD_DIR}" ] && [ ! -f "${LOAD_DIR}/latest_checkpointed_iteration.txt" ]; then + echo "ERROR: LOAD_DIR must be a native checkpoint root with latest_checkpointed_iteration.txt: ${LOAD_DIR}" >&2 + exit 2 +fi + +PREFLIGHT_COMMAND=( + python3 -m examples.graphgpo.preflight + --prepare-lock "${PREPARE_LOCK}" + --alfworld-data-root "${ALFWORLD_DATA}" + --model-lock "${MODEL_LOCK}" + --checkpoint "${HF_CHECKPOINT}" + --max-steps "${MAX_STEPS}" + --model-revision "${MODEL_REVISION}" + --task-groups "${TASK_GROUPS}" + --group-size "${GROUP_SIZE}" + --global-batch-size "${GLOBAL_BATCH_SIZE}" + --split-artifact train "${TRAIN_DATA}" "${TRAIN_MANIFEST}" +) +if [ "${ENABLE_EVAL}" = "1" ]; then + PREFLIGHT_COMMAND+=( + --split-artifact eval_in_distribution "${EVAL_DATA}" "${EVAL_MANIFEST}" + ) +fi +"${PREFLIGHT_COMMAND[@]}" + +export GRAPHGPO_METHOD="${METHOD}" +export GRAPHGPO_EXPECTED_GROUP_SIZE="${GROUP_SIZE}" +export GRAPHGPO_OMEGA="${OMEGA:-0.1}" +export GRAPHGPO_GAMMA="${GAMMA:-0.95}" +export GRAPHGPO_BETA="${BETA:-1.0}" +export GRAPHGPO_BETA_EPISODE="${BETA_EPISODE:-1.0}" +export GRAPHGPO_EPISODE_WEIGHTING="${EPISODE_WEIGHTING}" +if [ -n "${GRAPHGPO_DIAGNOSTICS_JSONL:-}" ]; then + export GRAPHGPO_DIAGNOSTICS_JSONL +fi +export RELAX_AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE="${MAX_STEPS}" + +TIMESTAMP="$(date -u '+%Y%m%dT%H%M%SZ')" +EXP_NAME="graphgpo-${METHOD}-qwen2.5-1.5b-s${SEED}-${TIMESTAMP}" +RUN_DIR="${SAVE_DIR}/runs/${EXP_NAME}" +CHECKPOINT_DIR="${SAVE_DIR}/checkpoints/${METHOD}/seed-${SEED}" + +MODEL_ARGS=( + --swiglu + --num-layers 28 + --hidden-size 1536 + --ffn-hidden-size 8960 + --num-attention-heads 12 + --group-query-attention + --num-query-groups 2 + --use-rotary-position-embeddings + --disable-bias-linear + --add-qkv-bias + --normalization RMSNorm + --norm-epsilon 1e-6 + --rotary-base 1000000 + --vocab-size 151936 + --kv-channels 128 +) + +RESOURCE_ARGS=( + --resource "{\"actor\":[1,${NUM_GPUS}],\"rollout\":[1,${NUM_GPUS}]}" + --max-staleness 0 + --num-data-storage-units 1 + --use-health-check + --colocate +) + +CHECKPOINT_ARGS=( + --hf-checkpoint "${HF_CHECKPOINT}" + --ref-load "${HF_CHECKPOINT}" + --save "${CHECKPOINT_DIR}" + --megatron-to-hf-mode bridge + --save-interval 10 + --max-actor-ckpt-to-keep 2 +) +if [ -n "${LOAD_DIR}" ]; then + CHECKPOINT_ARGS+=(--load "${LOAD_DIR}") +fi + +ROLLOUT_ARGS=( + --prompt-data "${TRAIN_DATA}" + --input-key messages + --metadata-key metadata + --use-agentic-rollout + --agent-command "bash ${SCRIPT_DIR}/run_agent_app.sh" + --agent-cwd "${PROJECT_ROOT}" + --agent-env + "ALFWORLD_CONFIG_PATH=${ALFWORLD_CONFIG_PATH}" + "ALFWORLD_DATA=${ALFWORLD_DATA}" + "ALFWORLD_PYTHON=${ALFWORLD_PYTHON}" + --agentic-custom-advantage-path examples.graphgpo.custom_advantage.compute_custom_advantage + --custom-rm-path examples.graphgpo.reward.reward_func + --group-rm + --num-rollout "${OUTER_EPOCHS}" + --rollout-batch-size "${TASK_GROUPS}" + --n-samples-per-prompt "${GROUP_SIZE}" + --rollout-max-prompt-len 2048 + --rollout-max-response-len 512 + --rollout-temperature 1.0 + --rollout-top-p 1.0 + --rollout-top-k -1 + --global-batch-size "${GLOBAL_BATCH_SIZE}" + --micro-batch-size 32 + --rollout-shuffle + --use-streaming-dataset + --agentic-prepare-pool-size 0 +) + +ALGORITHM_ARGS=( + --advantage-estimator grpo + --kl-coef 0 + --use-kl-loss + --kl-loss-coef 0.01 + --kl-loss-type low_var_kl + --eps-clip 0.2 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine "${NUM_GPUS}" + --sglang-mem-fraction-static "${SGLANG_MEM_FRACTION_STATIC}" +) + +MEGATRON_ARGS=( + --tensor-model-parallel-size "${NUM_GPUS}" + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + # TP ranks compile independently, so cold Dynamo/Triton progress can skew + # around collectives. Disable compilation before train actors import + # Megatron, and keep lazily imported JIT-fused kernels on the eager path. + --train-env-vars '{"TORCH_COMPILE_DISABLE":"1"}' + --disable-jit-fuser + --sequence-parallel + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --use-dynamic-batch-size + --max-tokens-per-gpu "${MAX_TOKENS_PER_GPU}" + --seed "${SEED}" +) + +EVAL_ARGS=() +if [ "${ENABLE_EVAL}" = "1" ]; then + EVAL_ARGS=( + --eval-interval 5 + --eval-prompt-data alfworld "${EVAL_DATA}" + --eval-input-key messages + --n-samples-per-eval-prompt 1 + --eval-max-response-len 512 + --eval-temperature 0.4 + --eval-top-p 1.0 + --eval-top-k -1 + --custom-eval-rollout-log-function-path examples.graphgpo.eval_logger.log_eval_rollout_data + ) +fi + +LOG_ARGS=( + --tb-project-name "${SAVE_DIR}/tensorboard" + --tb-experiment-name "${EXP_NAME}" +) + +TRAIN_COMMAND=( + python3 relax/entrypoints/train.py + "${RESOURCE_ARGS[@]}" + "${MODEL_ARGS[@]}" + "${CHECKPOINT_ARGS[@]}" + "${ROLLOUT_ARGS[@]}" + "${ALGORITHM_ARGS[@]}" + "${OPTIMIZER_ARGS[@]}" + "${SGLANG_ARGS[@]}" + "${MEGATRON_ARGS[@]}" + "${EVAL_ARGS[@]}" + "${LOG_ARGS[@]}" +) + +if [ "${DRY_RUN:-0}" = "1" ]; then + printf 'ray job submit --address=http://127.0.0.1:8265 -- ' + printf '%q ' "${TRAIN_COMMAND[@]}" + printf '\n' + exit 0 +fi + +if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then + set +u + source "${PROJECT_ROOT}/scripts/entrypoint/local.sh" + set -u +fi + +mkdir -p "${RUN_DIR}" "${CHECKPOINT_DIR}" +{ + echo "method=${METHOD}" + echo "enable_eval=${ENABLE_EVAL}" + echo "episode_weighting=${GRAPHGPO_EPISODE_WEIGHTING}" + echo "sglang_mem_fraction_static=${SGLANG_MEM_FRACTION_STATIC}" + if [ -n "${GRAPHGPO_DIAGNOSTICS_JSONL:-}" ]; then + echo "graph_diagnostics_jsonl=enabled" + else + echo "graph_diagnostics_jsonl=disabled" + fi + echo "seed=${SEED}" + echo "model_revision=${MODEL_REVISION}" + echo "expected_external_image_digest=${EXPECTED_EXTERNAL_IMAGE_DIGEST}" + echo "image_verification_scope=external_executor_required" + echo "alfworld_config_sha256=$(sha256sum "${ALFWORLD_CONFIG_PATH}" | cut -d' ' -f1)" + echo "train_manifest_sha256=$(sha256sum "${TRAIN_MANIFEST}" | cut -d' ' -f1)" + if [ "${ENABLE_EVAL}" = "1" ]; then + echo "eval_manifest_sha256=$(sha256sum "${EVAL_MANIFEST}" | cut -d' ' -f1)" + fi + echo "prepare_lock_sha256=$(sha256sum "${PREPARE_LOCK}" | cut -d' ' -f1)" + echo "dependency_lock_sha256=$(sha256sum "${DEPENDENCY_LOCK}" | cut -d' ' -f1)" + echo "model_lock_sha256=$(sha256sum "${MODEL_LOCK}" | cut -d' ' -f1)" +} >"${RUN_DIR}/run_lock.env" +{ + printf 'ray job submit --address=http://127.0.0.1:8265 -- ' + printf '%q ' "${TRAIN_COMMAND[@]}" + printf '\n' +} >"${RUN_DIR}/expanded_command.sh" + +ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://127.0.0.1:8265" \ + -- "${TRAIN_COMMAND[@]}" \ + 2>&1 | tee "${RUN_DIR}/train.log" diff --git a/examples/graphgpo/state.py b/examples/graphgpo/state.py new file mode 100644 index 000000000..6def6972f --- /dev/null +++ b/examples/graphgpo/state.py @@ -0,0 +1,236 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Deterministic ALFWorld state anchors without an environment dependency.""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass +from typing import Mapping, Sequence + + +STATE_ANCHOR_VERSION = "reference_anchor_v1" +ALFWORLD_TRACKER_FIELDS = ( + "location", + "holding", + "history_items", + "item_location", +) +_ALFWORLD_TRACKER_FIELD_SET = frozenset(ALFWORLD_TRACKER_FIELDS) + + +def _canonical_json(value: object) -> str: + try: + return json.dumps( + value, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + except (TypeError, ValueError) as exc: + raise ValueError("tracker values must be JSON-serializable and finite") from exc + + +@dataclass(frozen=True) +class TrackerState: + """An immutable snapshot of wrapper-maintained world-state fields.""" + + encoded_fields: tuple[tuple[str, str], ...] = () + + @classmethod + def from_mapping(cls, values: Mapping[str, object]) -> "TrackerState": + if not isinstance(values, Mapping): + raise TypeError("values must be a mapping") + + supplied_fields: set[str] = set() + for key in values: + if not isinstance(key, str) or not key: + raise ValueError("tracker field names must be non-empty strings") + supplied_fields.add(key) + + missing = _ALFWORLD_TRACKER_FIELD_SET - supplied_fields + unexpected = supplied_fields - _ALFWORLD_TRACKER_FIELD_SET + if missing or unexpected: + details = [] + if missing: + details.append(f"missing={sorted(missing)!r}") + if unexpected: + details.append(f"unexpected={sorted(unexpected)!r}") + raise ValueError( + f"tracker fields must be exactly {list(ALFWORLD_TRACKER_FIELDS)!r} ({', '.join(details)})" + ) + + encoded_fields = tuple((field, _canonical_json(values[field])) for field in ALFWORLD_TRACKER_FIELDS) + return cls(encoded_fields) + + def to_mapping(self) -> dict[str, object]: + return {key: json.loads(value) for key, value in self.encoded_fields} + + def render(self) -> str: + return "\n".join(f"{key}={value}" for key, value in self.encoded_fields) + + +def update_tracker( + tracker: TrackerState, + *, + raw_observation: str, + updates: Mapping[str, object], +) -> TrackerState: + """Apply structured wrapper updates unless the action changed nothing.""" + + if not isinstance(tracker, TrackerState): + raise TypeError("tracker must be a TrackerState") + if not isinstance(raw_observation, str): + raise TypeError("raw_observation must be a string") + if not isinstance(updates, Mapping): + raise TypeError("updates must be a mapping") + + update_fields: set[str] = set() + for key in updates: + if not isinstance(key, str) or not key: + raise ValueError("tracker update field names must be non-empty strings") + update_fields.add(key) + unexpected = update_fields - _ALFWORLD_TRACKER_FIELD_SET + if unexpected: + raise ValueError(f"tracker updates contain unexpected fields: {sorted(unexpected)!r}") + + if "nothing happens" in raw_observation.casefold(): + return tracker + + merged = tracker.to_mapping() + merged.update(updates) + return TrackerState.from_mapping(merged) + + +def _length_prefixed(value: str) -> str: + return f"{len(value)}:{value}" + + +def _active_item_states( + item: str, + history_items: Mapping[str, object], +) -> str: + state_value = history_items.get(item, {}) + if not isinstance(state_value, Mapping): + raise TypeError("every history_items value must be a mapping") + active_states: list[str] = [] + for state_name, active in state_value.items(): + if not isinstance(state_name, str) or not state_name: + raise TypeError("item state names must be non-empty strings") + if not isinstance(active, bool): + raise TypeError("item state values must be booleans") + if active: + active_states.append(state_name) + if active_states: + return f"({' '.join(active_states)})" + return "(unprocessed)" + + +def format_reference_observation( + raw_observation: str, + tracker: TrackerState, +) -> str: + """Append only tracker context displayed by the frozen ALFWorld worker.""" + + if not isinstance(raw_observation, str): + raise TypeError("raw_observation must be a string") + if not isinstance(tracker, TrackerState): + raise TypeError("tracker must be a TrackerState") + values = tracker.to_mapping() + + location = values.get("location") + holding = values.get("holding") + history_items = values.get("history_items") + item_location = values.get("item_location") + if not isinstance(location, str) or not isinstance(holding, str): + raise TypeError("tracker location and holding must be strings") + if not isinstance(history_items, Mapping): + raise TypeError("tracker history_items must be a mapping") + if not isinstance(item_location, Mapping): + raise TypeError("tracker item_location must be a mapping") + + holding_states: list[str] = [] + for item in history_items: + if not isinstance(item, str): + raise TypeError("history item names must be strings") + if item in holding: + active = _active_item_states(item, history_items) + if active != "(unprocessed)": + holding_states.append(active[1:-1]) + holding_suffix = "" + if holding_states: + holding_suffix = f"({' '.join(holding_states)})" + elif holding != "nothing": + holding_suffix = "(unprocessed)" + + movement_entries: list[str] = [] + for item, movement in item_location.items(): + if not isinstance(item, str): + raise TypeError("item_location names must be strings") + if not isinstance(movement, Mapping): + raise TypeError("every item_location value must be a mapping") + old_location = movement.get("old_location") + new_location = movement.get("new_location") + if not isinstance(old_location, str) or not isinstance(new_location, str): + raise TypeError("item locations must be strings") + item_state = _active_item_states(item, history_items) + if old_location != new_location or (item_state != "(unprocessed)" and holding != item): + movement_entries.append(f"{item}{item_state} from {old_location} move to {new_location};") + + movement_history = "" + if movement_entries: + movement_history = "History moving list by you: " + "".join(movement_entries) + location_holding = f"Location: {location}. Items in hand (status): {holding}{holding_suffix}." + return f"{raw_observation} {location_holding} {movement_history}" + + +def reference_anchor_v1( + raw_observation: str, + tracker: TrackerState, + admissible_commands: Sequence[str], +) -> str: + """Build the frozen reference-style, task-local state anchor. + + The anchor uses exactly the tracker-derived context displayed to the model, + rather than internal bookkeeping fields that the reference worker filters + out. It also sorts all admissible commands. In particular, ``help`` remains + part of the state even though it is hidden from the model prompt. Length + prefixes make the representation collision resistant without depending on + process-specific ``hash()``. + """ + + if not isinstance(raw_observation, str): + raise TypeError("raw_observation must be a string") + if not isinstance(tracker, TrackerState): + raise TypeError("tracker must be a TrackerState") + if isinstance(admissible_commands, (str, bytes)) or not isinstance(admissible_commands, Sequence): + raise TypeError("admissible_commands must be a sequence of strings") + + commands: list[str] = [] + for command in admissible_commands: + if not isinstance(command, str): + raise TypeError("every admissible command must be a string") + commands.append(command) + commands.sort() + + display_observation = format_reference_observation(raw_observation, tracker) + parts = [ + STATE_ANCHOR_VERSION, + f"observation={_length_prefixed(display_observation)}", + f"command_count={len(commands)}", + ] + parts.extend(f"command={_length_prefixed(command)}" for command in commands) + return "\n".join(parts) + + +def state_key_v1( + raw_observation: str, + tracker: TrackerState, + admissible_commands: Sequence[str], +) -> str: + """Return a stable SHA256 key for ``reference_anchor_v1``.""" + + anchor = reference_anchor_v1(raw_observation, tracker, admissible_commands) + return hashlib.sha256(anchor.encode("utf-8")).hexdigest() diff --git a/relax/agentic/pipeline/reward.py b/relax/agentic/pipeline/reward.py index c88a1edad..882fa51e2 100644 --- a/relax/agentic/pipeline/reward.py +++ b/relax/agentic/pipeline/reward.py @@ -63,13 +63,73 @@ def materialized_units(self) -> list[PendingExportUnit]: def materialized_group(self) -> list[Any]: return [unit.sample for unit in self.materialized_units()] - def metadata_by_slot(self) -> list[dict[Any, dict[str, Any]]]: + @staticmethod + def _resolved_policy_version(sample: Any) -> str: + metadata = sample.metadata + if not isinstance(metadata, dict): + raise RuntimeError("Agentic row metadata must be a dict before custom advantage evaluation.") + + weight_versions = getattr(sample, "weight_versions", None) + if weight_versions is None: + weight_versions = [] + if not isinstance(weight_versions, (list, tuple)): + raise RuntimeError("Sample.weight_versions must be a list or tuple.") + normalized_versions: list[str] = [] + for value in weight_versions: + if isinstance(value, bool) or not isinstance(value, (str, int)) or not str(value): + raise RuntimeError(f"Sample.weight_versions contains an invalid policy version: {value!r}.") + normalized_versions.append(str(value)) + unique_versions = set(normalized_versions) + if len(unique_versions) > 1: + raise RuntimeError( + f"One agentic row spans multiple policy versions: weight_versions={sorted(unique_versions)!r}." + ) + if unique_versions: + resolved = next(iter(unique_versions)) + else: + start_rollout_id = metadata.get("start_rollout_id") + if ( + isinstance(start_rollout_id, bool) + or not isinstance(start_rollout_id, (str, int)) + or not str(start_rollout_id) + ): + raise RuntimeError( + "Agentic row policy_version is unavailable: Sample.weight_versions is empty " + "and metadata.start_rollout_id is missing or invalid." + ) + resolved = str(start_rollout_id) + + declared = metadata.get("policy_version") + if declared is not None: + if isinstance(declared, bool) or not isinstance(declared, (str, int)) or not str(declared): + raise RuntimeError(f"metadata.policy_version is invalid: {declared!r}.") + if str(declared) != resolved: + raise RuntimeError( + "metadata.policy_version conflicts with the policy version captured by Relax: " + f"declared={declared!r}, captured={resolved!r}." + ) + return resolved + + def metadata_by_slot( + self, + *, + require_policy_version: bool = False, + ) -> list[dict[Any, dict[str, Any]]]: payload: list[dict[Any, dict[str, Any]]] = [] + missing_policy_versions: list[tuple[Any, str]] = [] for idx in range(self.expected_count): slot_payload: dict[Any, dict[str, Any]] = {} for unit in self.units_by_slot.get(idx, []): - slot_payload[unit.name] = copy.deepcopy(unit.sample.metadata) + metadata = copy.deepcopy(unit.sample.metadata) + if require_policy_version: + resolved = self._resolved_policy_version(unit.sample) + metadata["policy_version"] = resolved + if "policy_version" not in unit.sample.metadata: + missing_policy_versions.append((unit.sample, resolved)) + slot_payload[unit.name] = metadata payload.append(slot_payload) + for sample, resolved in missing_policy_versions: + sample.metadata["policy_version"] = resolved return payload def materialized_slot_count(self) -> int: @@ -93,6 +153,12 @@ def __init__( self.custom_advantage_func = ( load_function(custom_advantage_path) if custom_advantage_path is not None else None ) + if self.custom_advantage_func is None: + self._require_agentic_policy_version = False + else: + from relax.agentic.pipeline.transfer import use_agentic_variable_row_mode + + self._require_agentic_policy_version = use_agentic_variable_row_mode(args) self._waiting_groups: dict[GroupKey, RewardWaitingGroup] = {} self._inflight_sample_tasks: dict[SampleKey, asyncio.Task] = {} self._inflight_group_tasks: dict[GroupKey, asyncio.Task] = {} @@ -341,7 +407,11 @@ def _apply_custom_advantage(value: Any, unit: PendingExportUnit) -> None: raise TypeError("custom advantage output must be a number in this version.") async def _run_custom_advantage(self, waiting_group: RewardWaitingGroup) -> list[Any] | None: - result = self.custom_advantage_func(waiting_group.metadata_by_slot()) + result = self.custom_advantage_func( + waiting_group.metadata_by_slot( + require_policy_version=self._require_agentic_policy_version, + ) + ) if asyncio.iscoroutine(result): result = await result if result is None: diff --git a/relax/agentic/pipeline/transfer.py b/relax/agentic/pipeline/transfer.py index e10bfca4e..6ceb044d8 100644 --- a/relax/agentic/pipeline/transfer.py +++ b/relax/agentic/pipeline/transfer.py @@ -4,8 +4,12 @@ import asyncio import copy +import hashlib +import os import time from collections import deque +from collections.abc import Mapping +from typing import Any from relax.agentic.pipeline import GroupKey, sample_group_key from relax.agentic.profile import ( @@ -18,6 +22,316 @@ logger = get_logger(__name__) +AGENTIC_VARIABLE_ROW_PADDING_KEY = "agentic_variable_row_padding" +AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV = "RELAX_AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE" +AGENTIC_ROW_IDENTITY_KEY = "agentic_row_identity" +AGENTIC_ROW_IDENTITY_TAG_KEY = "agentic_row_identity_tag" +AGENTIC_ROW_IDENTITY_TAGS_FIELD = "agentic_row_identity_tags" +AGENTIC_ROW_IDENTITY_SCHEMA_VERSION = 1 +_AGENTIC_TRANSFER_IDENTITY_ATTR = "_agentic_transfer_identity" + + +def _agentic_max_exported_rows_per_sample() -> int | None: + raw_max_rows = os.environ.get(AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV) + if raw_max_rows is None or not raw_max_rows.strip(): + return None + normalized_max_rows = raw_max_rows.strip() + if not normalized_max_rows.isascii() or not normalized_max_rows.isdigit() or int(normalized_max_rows) <= 0: + raise ValueError( + f"{AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV} must be a positive integer, got {raw_max_rows!r}." + ) + return int(normalized_max_rows) + + +def use_agentic_variable_row_mode(args) -> bool: + """Return whether variable explicit-row transfer is enabled. + + The feature remains default-off and is limited to synchronous training. + Async and hybrid consumers have different lifecycle and accounting + contracts and must not silently enter this path. + """ + + base_enabled = bool( + getattr(args, "group_rm", False) + and getattr(args, "agentic_custom_advantage_path", None) + and getattr(args, "use_dynamic_batch_size", False) + ) + if not base_enabled: + return False + + max_rows_per_sample = _agentic_max_exported_rows_per_sample() + if max_rows_per_sample is None: + return False + + if getattr(args, "fully_async", False) or getattr(args, "hybrid", False): + raise ValueError("Agentic variable-row transfer supports synchronous training only.") + return True + + +def _flatten_transfer_samples(batch_samples: list) -> list: + flat_samples: list = [] + pending = deque(batch_samples) + while pending: + item = pending.popleft() + if isinstance(item, list): + pending.extendleft(reversed(item)) + else: + flat_samples.append(item) + return flat_samples + + +def _identity_scalar(name: str, value: Any) -> str: + if isinstance(value, bool) or not isinstance(value, (str, int)) or not str(value): + raise RuntimeError(f"Agentic row {name} must be a non-empty string or integer, got {value!r}.") + return str(value) + + +def _captured_policy_version(sample: Any, metadata: dict[str, Any]) -> str: + weight_versions = getattr(sample, "weight_versions", None) + if weight_versions is None: + weight_versions = [] + if not isinstance(weight_versions, (list, tuple)): + raise RuntimeError("Sample.weight_versions must be a list or tuple for agentic row transfer.") + normalized_versions = {_identity_scalar("weight_version", value) for value in weight_versions} + if len(normalized_versions) > 1: + raise RuntimeError( + "Agentic row spans multiple policy versions before transfer: " + f"weight_versions={sorted(normalized_versions)!r}." + ) + if normalized_versions: + captured = next(iter(normalized_versions)) + else: + captured = _identity_scalar("start_rollout_id", metadata.get("start_rollout_id")) + + declared = metadata.get("policy_version") + if declared is not None and _identity_scalar("policy_version", declared) != captured: + raise RuntimeError( + "Agentic row policy_version conflicts with Relax's captured version: " + f"declared={declared!r}, captured={captured!r}." + ) + if declared is None: + metadata["policy_version"] = captured + return captured + + +def _agentic_row_identity(sample: Any) -> dict[str, Any]: + metadata = sample.metadata + if not isinstance(metadata, dict): + raise RuntimeError("Agentic variable-row samples must carry dict metadata.") + row_id = metadata.get("row_id") + if not isinstance(row_id, str) or not row_id: + raise RuntimeError(f"Agentic row_id must be a non-empty string, got {row_id!r}.") + rollout_group_id = _identity_scalar("rollout_group_id", metadata.get("rollout_group_id")) + sample_group_id = _identity_scalar("Sample.group_index", getattr(sample, "group_index", None)) + if rollout_group_id != sample_group_id: + raise RuntimeError( + "Agentic row rollout_group_id does not match Sample.group_index: " + f"metadata={rollout_group_id!r}, sample={sample_group_id!r}." + ) + policy_version = _captured_policy_version(sample, metadata) + task_id = metadata.get("task_id") + if isinstance(task_id, bool) or not isinstance(task_id, (str, int)) or not str(task_id): + raise RuntimeError(f"Agentic row task_id must be a non-empty string or integer, got {task_id!r}.") + trajectory_id = metadata.get("trajectory_id") + if not isinstance(trajectory_id, str) or not trajectory_id: + raise RuntimeError(f"Agentic row trajectory_id must be a non-empty string, got {trajectory_id!r}.") + if getattr(sample, "session_id", None) != trajectory_id: + raise RuntimeError( + "Agentic row trajectory_id does not match Sample.session_id: " + f"metadata={trajectory_id!r}, sample={getattr(sample, 'session_id', None)!r}." + ) + turn_index = metadata.get("turn_index") + if isinstance(turn_index, bool) or not isinstance(turn_index, int) or turn_index < 0: + raise RuntimeError(f"Agentic row turn_index must be a non-negative integer, got {turn_index!r}.") + turn_id = metadata.get("turn_id") + expected_turn_id = f"turn_{turn_index:03d}" + if turn_id != expected_turn_id or getattr(sample, "_agentic_export_name", None) != expected_turn_id: + raise RuntimeError( + "Agentic row turn identity mismatch: " + f"turn_id={turn_id!r}, export_name={getattr(sample, '_agentic_export_name', None)!r}, " + f"expected={expected_turn_id!r}." + ) + sample_index = getattr(sample, "index", None) + if isinstance(sample_index, bool) or not isinstance(sample_index, int) or sample_index < 0: + raise RuntimeError(f"Agentic row Sample.index must be a non-negative integer, got {sample_index!r}.") + terminal = metadata.get("terminal") + truncated = metadata.get("truncated") + if type(terminal) is not bool or type(truncated) is not bool: + raise RuntimeError("Agentic row terminal and truncated markers must be bool values.") + + response_length = getattr(sample, "response_length", None) + if isinstance(response_length, bool) or not isinstance(response_length, int) or response_length <= 0: + raise RuntimeError(f"Agentic row response_length must be positive, got {response_length!r}.") + tokens = getattr(sample, "tokens", None) + if not isinstance(tokens, list) or len(tokens) < response_length: + raise RuntimeError("Agentic row tokens must contain the complete response token suffix.") + loss_mask = getattr(sample, "loss_mask", None) + if not isinstance(loss_mask, list) or len(loss_mask) != response_length: + raise RuntimeError( + "Agentic row loss_mask must contain one entry per response token: " + f"response_length={response_length}, mask_length={len(loss_mask) if isinstance(loss_mask, list) else None}." + ) + if any(type(value) is not int or value not in (0, 1) for value in loss_mask): + raise RuntimeError("Agentic row loss_mask must contain only integer 0/1 values.") + action_token_count = sum(loss_mask) + if action_token_count <= 0: + raise RuntimeError("Agentic row must contain at least one trainable assistant action token.") + rollout_log_probs = getattr(sample, "rollout_log_probs", None) + if not isinstance(rollout_log_probs, list) or len(rollout_log_probs) != response_length: + raise RuntimeError( + "Agentic row rollout_log_probs must contain one entry per response token: " + f"response_length={response_length}, " + f"log_prob_length={len(rollout_log_probs) if isinstance(rollout_log_probs, list) else None}." + ) + + return { + "schema_version": AGENTIC_ROW_IDENTITY_SCHEMA_VERSION, + "padding": False, + "row_id": row_id, + "rollout_group_id": rollout_group_id, + "policy_version": policy_version, + "task_id": task_id, + "trajectory_id": trajectory_id, + "turn_id": turn_id, + "turn_index": turn_index, + "sample_index": sample_index, + "terminal": terminal, + "truncated": truncated, + "total_length": len(tokens), + "response_length": response_length, + "action_token_count": action_token_count, + } + + +def _row_id_digest(row_ids: list[str]) -> str: + encoded = "\n".join(sorted(row_ids)).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _row_identity_tag(row_id: str) -> int: + digest = hashlib.sha256(row_id.encode("utf-8")).digest() + return int.from_bytes(digest[:8], byteorder="big") & ((1 << 63) - 1) + + +def _padding_row_identity( + *, + rollout_id: int, + row_ordinal: int, + policy_version: str, + total_length: int, + response_length: int, + expected_group_count: int, +) -> dict[str, Any]: + return { + "schema_version": AGENTIC_ROW_IDENTITY_SCHEMA_VERSION, + "padding": True, + "row_id": f"agentic-padding-v1:{rollout_id}:{row_ordinal}", + "rollout_group_id": None, + "policy_version": policy_version, + "task_id": None, + "trajectory_id": None, + "turn_id": None, + "turn_index": None, + "sample_index": None, + "terminal": False, + "truncated": False, + "total_length": total_length, + "response_length": response_length, + "action_token_count": 0, + "group_row_count": 0, + "group_trajectory_count": 0, + "group_row_ids_sha256": None, + "partition_expected_group_count": expected_group_count, + } + + +def _validate_agentic_materialized_group( + *, args: Any, group: list +) -> tuple[list[tuple[Any, dict[str, Any]]], dict[str, Any]]: + if not isinstance(group, list) or not group or any(isinstance(sample, list) for sample in group): + raise RuntimeError("Agentic variable-row transfer requires a non-empty flat materialized group.") + identities = [(sample, _agentic_row_identity(sample)) for sample in group] + row_ids = [identity["row_id"] for _, identity in identities] + if len(set(row_ids)) != len(row_ids): + raise RuntimeError("Agentic materialized group contains a duplicate row_id.") + + group_ids = {identity["rollout_group_id"] for _, identity in identities} + policy_versions = {identity["policy_version"] for _, identity in identities} + task_ids = {str(identity["task_id"]) for _, identity in identities} + if len(group_ids) != 1 or len(policy_versions) != 1 or len(task_ids) != 1: + raise RuntimeError( + "Agentic materialized group mixes rollout group, policy version, or task identity: " + f"groups={sorted(group_ids)!r}, policies={sorted(policy_versions)!r}, tasks={sorted(task_ids)!r}." + ) + + by_trajectory: dict[str, list[dict[str, Any]]] = {} + trajectory_by_sample_index: dict[int, str] = {} + for _, identity in identities: + trajectory_id = identity["trajectory_id"] + sample_index = identity["sample_index"] + previous = trajectory_by_sample_index.setdefault(sample_index, trajectory_id) + if previous != trajectory_id: + raise RuntimeError(f"Sample.index {sample_index} maps to multiple trajectory_id values.") + by_trajectory.setdefault(trajectory_id, []).append(identity) + if len(trajectory_by_sample_index) != len(by_trajectory): + raise RuntimeError("Multiple trajectory_id values share one Sample.index in an agentic group.") + expected_trajectories = int(args.n_samples_per_prompt) + if len(by_trajectory) != expected_trajectories: + raise RuntimeError( + "Agentic materialized group is incomplete: " + f"expected_trajectories={expected_trajectories}, got={len(by_trajectory)}." + ) + for trajectory_id, trajectory in by_trajectory.items(): + ordered = sorted(trajectory, key=lambda identity: identity["turn_index"]) + actual_turns = [identity["turn_index"] for identity in ordered] + if actual_turns != list(range(len(ordered))): + raise RuntimeError( + f"Agentic trajectory {trajectory_id!r} has missing or duplicate turn rows: {actual_turns!r}." + ) + for identity in ordered[:-1]: + if identity["terminal"] or identity["truncated"]: + raise RuntimeError(f"Agentic trajectory {trajectory_id!r} terminates before its final row.") + final = ordered[-1] + if final["terminal"] == final["truncated"]: + raise RuntimeError( + f"Agentic trajectory {trajectory_id!r} final row must be exactly one of terminal or truncated." + ) + + manifest = { + "group_row_count": len(identities), + "group_trajectory_count": len(by_trajectory), + "group_row_ids_sha256": _row_id_digest(row_ids), + "partition_expected_group_count": int(args.rollout_batch_size), + } + return identities, manifest + + +def _variable_row_padding_sample(*, args, template): + sample = copy.deepcopy(template) + sample.remove_sample = True + sample.loss_mask = None + sample.custom_advantage = 0.0 + reward_key = getattr(args, "reward_key", None) + sample.reward = {reward_key: 0.0} if reward_key else 0.0 + sample.metadata = copy.deepcopy(sample.metadata) if isinstance(sample.metadata, dict) else {} + sample.metadata.pop("raw_reward", None) + for key in ( + "row_id", + "rollout_group_id", + "policy_version", + "task_id", + "trajectory_id", + "turn_id", + "turn_index", + "terminal", + "truncated", + ): + sample.metadata.pop(key, None) + sample.metadata[AGENTIC_VARIABLE_ROW_PADDING_KEY] = True + if hasattr(sample, _AGENTIC_TRANSFER_IDENTITY_ATTR): + delattr(sample, _AGENTIC_TRANSFER_IDENTITY_ATTR) + return sample + async def _transfer_batch_to_data_system( *, @@ -27,30 +341,44 @@ async def _transfer_batch_to_data_system( rollout_id: int, data_system_client, is_last: bool = False, + timing_sink=None, ) -> list[str]: from relax.utils.utils import convert_samples_to_train_data - if batch_samples: - enqueue_at = time.time() - flat_samples = batch_samples - while flat_samples and isinstance(flat_samples[0], list): - flat_samples = sum(flat_samples, []) - for sample in flat_samples: - mark_sample_agentic_event(sample, "transfer_enqueue_at", enqueue_at) - else: + if not batch_samples: logger.warning( "transfer_batch_to_data_system called with empty batch_samples for rollout_id=%s, batch_count=%s", rollout_id, batch_count, ) return [] - batch_samples = sorted( + + variable_row_mode = use_agentic_variable_row_mode(args) + timing_record = { + "rollout_id": rollout_id, + "batch_count": batch_count, + "is_last": is_last, + "row_count": 0, + "reorder_ms": 0.0, + "identity_validation_ms": 0.0, + "serialization_ms": 0.0, + "queue_transfer_ms": 0.0, + "ok": False, + } + + reorder_started_at = time.perf_counter() + flat_samples = _flatten_transfer_samples(batch_samples) + enqueue_at = time.time() + for sample in flat_samples: + mark_sample_agentic_event(sample, "transfer_enqueue_at", enqueue_at) + ordered_groups = sorted( batch_samples, key=lambda group: group[0][0].index if isinstance(group[0], list) else group[0].index ) - while isinstance(batch_samples[0], list): - batch_samples = sum(batch_samples, []) + ordered_samples = _flatten_transfer_samples(ordered_groups) + timing_record["reorder_ms"] = (time.perf_counter() - reorder_started_at) * 1000.0 + transfer_samples = [] - for sample in batch_samples: + for sample in ordered_samples: if sample.reward is None: sample = copy.copy(sample) # Fall back to custom advantage when neither JSON reward nor custom RM supplies a reward. @@ -58,16 +386,112 @@ async def _transfer_batch_to_data_system( # still use the original sample and ignore this fallback. sample.reward = {args.reward_key: sample.custom_advantage} if args.reward_key else sample.custom_advantage transfer_samples.append(sample) + + identities: list[dict[str, Any]] = [] + if variable_row_mode: + identity_started_at = time.perf_counter() + actual_policy_versions: set[str] = set() + for sample in transfer_samples: + is_padding = bool( + sample.remove_sample + and isinstance(sample.metadata, dict) + and sample.metadata.get(AGENTIC_VARIABLE_ROW_PADDING_KEY) is True + ) + if is_padding: + identities.append({}) + continue + cached_identity = getattr(sample, _AGENTIC_TRANSFER_IDENTITY_ATTR, None) + if not isinstance(cached_identity, Mapping): + raise RuntimeError("Agentic row reached transfer without a validated group identity manifest.") + current_identity = _agentic_row_identity(sample) + for key, value in current_identity.items(): + if cached_identity.get(key) != value: + raise RuntimeError( + "Agentic row identity changed after group validation: " + f"row_id={current_identity['row_id']!r}, field={key!r}." + ) + identity = copy.deepcopy(dict(cached_identity)) + identities.append(identity) + actual_policy_versions.add(identity["policy_version"]) + if len(actual_policy_versions) != 1: + raise RuntimeError( + "One TransferQueue batch must contain exactly one policy version: " + f"versions={sorted(actual_policy_versions)!r}." + ) + policy_version = next(iter(actual_policy_versions)) + for row_ordinal, (sample, identity) in enumerate(zip(transfer_samples, identities, strict=True)): + if identity: + continue + identity.update( + _padding_row_identity( + rollout_id=rollout_id, + row_ordinal=row_ordinal, + policy_version=policy_version, + total_length=len(sample.tokens), + response_length=sample.response_length, + expected_group_count=int(args.rollout_batch_size), + ) + ) + sample.metadata.update( + { + "row_id": identity["row_id"], + "policy_version": policy_version, + } + ) + if len({identity["row_id"] for identity in identities}) != len(identities): + raise RuntimeError("TransferQueue batch contains duplicate physical row identities.") + timing_record["identity_validation_ms"] = (time.perf_counter() - identity_started_at) * 1000.0 + + serialization_started_at = time.perf_counter() rollout_batch = convert_samples_to_train_data(args, transfer_samples) + total_lengths = [int(length) for length in rollout_batch["total_lengths"]] + if len(total_lengths) != len(transfer_samples): + raise RuntimeError( + "Converted rollout batch lost row alignment: " + f"samples={len(transfer_samples)}, lengths={len(total_lengths)}." + ) + if variable_row_mode: + import torch + + for identity, total_length in zip(identities, total_lengths, strict=True): + if identity["total_length"] != total_length: + raise RuntimeError( + "Converted rollout row length does not match its identity: " + f"row_id={identity['row_id']!r}, identity={identity['total_length']}, batch={total_length}." + ) + rollout_batch[AGENTIC_ROW_IDENTITY_TAGS_FIELD] = torch.tensor( + [_row_identity_tag(identity["row_id"]) for identity in identities], + dtype=torch.int64, + ) + timing_record["serialization_ms"] = (time.perf_counter() - serialization_started_at) * 1000.0 + timing_record["row_count"] = len(transfer_samples) logger.info("Prepared rollout batch %s with %s samples for transfer", batch_count, rollout_batch.numel()) logger.info("Transferring batch rollout_batch: %s", rollout_batch) - custom_meta = [{"total_lengths": int(length)} for length in rollout_batch["total_lengths"]] - await data_system_client.async_put( - data=rollout_batch, - partition_id=f"train_{rollout_id}", - custom_meta=custom_meta, - is_last=is_last, - ) + if variable_row_mode: + custom_meta = [ + { + "total_lengths": total_length, + AGENTIC_VARIABLE_ROW_PADDING_KEY: identity["padding"], + AGENTIC_ROW_IDENTITY_KEY: copy.deepcopy(identity), + AGENTIC_ROW_IDENTITY_TAG_KEY: _row_identity_tag(identity["row_id"]), + } + for identity, total_length in zip(identities, total_lengths, strict=True) + ] + else: + custom_meta = [{"total_lengths": length} for length in total_lengths] + queue_started_at = time.perf_counter() + try: + await data_system_client.async_put( + data=rollout_batch, + partition_id=f"train_{rollout_id}", + custom_meta=custom_meta, + is_last=is_last, + ) + timing_record["ok"] = True + finally: + timing_record["queue_transfer_ms"] = (time.perf_counter() - queue_started_at) * 1000.0 + if variable_row_mode and timing_sink is not None: + timing_sink(copy.deepcopy(timing_record)) logger.info("Batch %s transferred successfully for rollout_id: %s", batch_count, rollout_id) return list(rollout_batch.keys()) @@ -93,6 +517,10 @@ def __init__( ) self._transfer_buffer: deque[list] = deque() self._transfer_tasks: list[asyncio.Task] = [] + self._variable_row_mode = use_agentic_variable_row_mode(args) + self._partition_actual_rows: dict[int, int] = {} + self._partition_identity_state: dict[int, dict[str, Any]] = {} + self._transfer_timing_records: deque[dict[str, Any]] = deque(maxlen=4096) self._reset_step_partition_state() def _reset_step_partition_state(self) -> None: @@ -117,6 +545,16 @@ def rebind_step( f"rollout_id={self.rollout_id}, next_rollout_id={rollout_id}, " f"buffer_groups={len(self._transfer_buffer)}, pending_tasks={len(self._transfer_tasks)}." ) + stale_identity_partitions = [ + partition_id + for partition_id in self._partition_identity_state + if partition_id < rollout_id - 1 or partition_id > rollout_id + ] + if stale_identity_partitions: + raise RuntimeError( + "TransferDomain cannot rebind with stale or future agentic identity state: " + f"next_rollout_id={rollout_id}, partitions={sorted(stale_identity_partitions)!r}." + ) self.rollout_id = rollout_id self._reset_step_partition_state() @@ -161,12 +599,141 @@ async def _dispatch_transfer_batch(self, *, groups, partition_rollout_id: int, i rollout_id=partition_rollout_id, data_system_client=self.data_system_client, is_last=is_last, + timing_sink=self._transfer_timing_records.append if self._variable_row_mode else None, ) + if self._variable_row_mode and is_last: + self._partition_actual_rows.pop(partition_rollout_id, None) + self._partition_identity_state.pop(partition_rollout_id, None) release_ended_at = time.time() for group in groups: for sample in group: mark_sample_agentic_event(sample, "transfer_release_end_at", release_ended_at) + async def _dispatch_transfer_batch_after( + self, + *, + predecessor: asyncio.Task | None, + groups, + partition_rollout_id: int, + is_last: bool, + ) -> None: + if predecessor is not None: + await predecessor + await self._dispatch_transfer_batch( + groups=groups, + partition_rollout_id=partition_rollout_id, + is_last=is_last, + ) + + def _prepare_variable_row_groups( + self, + *, + groups: list[list], + partition_rollout_id: int, + is_last: bool, + ) -> list[list]: + actual_samples = _flatten_transfer_samples(groups) + if not actual_samples: + raise RuntimeError( + "Agentic variable-row transfer cannot dispatch an empty materialized batch: " + f"partition_rollout_id={partition_rollout_id}." + ) + validated_groups = [_validate_agentic_materialized_group(args=self.args, group=group) for group in groups] + candidate_row_ids: set[str] = set() + candidate_group_manifests: dict[str, dict[str, Any]] = {} + candidate_policy_versions: set[str] = set() + for identities, manifest in validated_groups: + group_id = identities[0][1]["rollout_group_id"] + if group_id in candidate_group_manifests: + raise RuntimeError(f"Agentic transfer batch repeats rollout_group_id {group_id!r}.") + candidate_group_manifests[group_id] = manifest + for _, identity in identities: + row_id = identity["row_id"] + if row_id in candidate_row_ids: + raise RuntimeError(f"Agentic transfer batch repeats row_id {row_id!r}.") + candidate_row_ids.add(row_id) + candidate_policy_versions.add(identity["policy_version"]) + if len(candidate_policy_versions) != 1: + raise RuntimeError( + f"Agentic transfer batch mixes policy versions: versions={sorted(candidate_policy_versions)!r}." + ) + + previous_identity_state = self._partition_identity_state.get(partition_rollout_id) + existing_row_ids = set(previous_identity_state["row_ids"]) if previous_identity_state else set() + existing_group_manifests = ( + copy.deepcopy(previous_identity_state["group_manifests"]) if previous_identity_state else {} + ) + existing_policy_version = previous_identity_state["policy_version"] if previous_identity_state else None + duplicate_row_ids = existing_row_ids.intersection(candidate_row_ids) + if duplicate_row_ids: + raise RuntimeError( + "Agentic partition received duplicate row_id values across retries or batches: " + f"duplicates={sorted(duplicate_row_ids)!r}." + ) + duplicate_group_ids = set(existing_group_manifests).intersection(candidate_group_manifests) + if duplicate_group_ids: + raise RuntimeError( + f"Agentic partition received a rollout group more than once: groups={sorted(duplicate_group_ids)!r}." + ) + candidate_policy_version = next(iter(candidate_policy_versions)) + if existing_policy_version is not None and existing_policy_version != candidate_policy_version: + raise RuntimeError( + "Agentic partition mixes policy versions across transfer batches: " + f"existing={existing_policy_version!r}, candidate={candidate_policy_version!r}." + ) + next_group_manifests = {**existing_group_manifests, **candidate_group_manifests} + if len(next_group_manifests) > int(self.args.rollout_batch_size): + raise RuntimeError( + "Agentic partition contains more rollout groups than configured: " + f"groups={len(next_group_manifests)}, rollout_batch_size={self.args.rollout_batch_size}." + ) + if is_last and len(next_group_manifests) != int(self.args.rollout_batch_size): + raise RuntimeError( + "Agentic partition closed with a residual or missing rollout group: " + f"expected={self.args.rollout_batch_size}, got={len(next_group_manifests)}." + ) + + partition_actual_rows = self._partition_actual_rows.get(partition_rollout_id, 0) + len(actual_samples) + max_rows_per_sample = _agentic_max_exported_rows_per_sample() + if max_rows_per_sample is None: + raise RuntimeError( + f"{AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV} disappeared while " + "agentic variable-row transfer was active." + ) + max_partition_actual_rows = ( + int(self.args.rollout_batch_size) * int(self.args.n_samples_per_prompt) * max_rows_per_sample + ) + if partition_actual_rows > max_partition_actual_rows: + raise RuntimeError( + "Agentic variable-row export exceeded the configured partition bound: " + f"partition_rollout_id={partition_rollout_id}, " + f"actual_rows={partition_actual_rows}, " + f"max_rows={max_partition_actual_rows}." + ) + for identities, manifest in validated_groups: + for sample, identity in identities: + setattr( + sample, + _AGENTIC_TRANSFER_IDENTITY_ATTR, + {**copy.deepcopy(identity), **copy.deepcopy(manifest)}, + ) + self._partition_actual_rows[partition_rollout_id] = partition_actual_rows + self._partition_identity_state[partition_rollout_id] = { + "row_ids": existing_row_ids.union(candidate_row_ids), + "group_manifests": next_group_manifests, + "policy_version": existing_policy_version or candidate_policy_version, + } + if not is_last: + return groups + + padding_rows = (-partition_actual_rows) % self.args.global_batch_size + if padding_rows == 0: + return groups + padding_group = [ + _variable_row_padding_sample(args=self.args, template=actual_samples[-1]) for _ in range(padding_rows) + ] + return [*groups, padding_group] + def _spawn_transfer(self, *, force: bool = False) -> int: self._reap_completed_transfer_tasks() if not self._transfer_buffer: @@ -206,16 +773,39 @@ def _spawn_transfer(self, *, force: bool = False) -> int: groups = [] for _ in range(take_count): groups.append(self._transfer_buffer.popleft()) - if self.data_system_client is None: - return len(groups) - self._transfer_tasks.append( - asyncio.create_task( + transferred_group_count = len(groups) + if self._variable_row_mode: + groups = self._prepare_variable_row_groups( + groups=groups, + partition_rollout_id=partition_rollout_id, + is_last=is_last, + ) + if self.data_system_client is None: + if is_last: + self._partition_actual_rows.pop(partition_rollout_id, None) + self._partition_identity_state.pop(partition_rollout_id, None) + return transferred_group_count + predecessor = self._transfer_tasks[-1] if self._transfer_tasks else None + task = asyncio.create_task( + self._dispatch_transfer_batch_after( + predecessor=predecessor, + groups=groups, + partition_rollout_id=partition_rollout_id, + is_last=is_last, + ) + ) + else: + if self.data_system_client is None: + return transferred_group_count + task = asyncio.create_task( self._dispatch_transfer_batch( - groups=groups, partition_rollout_id=partition_rollout_id, is_last=is_last + groups=groups, + partition_rollout_id=partition_rollout_id, + is_last=is_last, ) ) - ) - return len(groups) + self._transfer_tasks.append(task) + return transferred_group_count def _spawn_ready_transfers(self) -> int: spawned_group_count = 0 @@ -272,6 +862,14 @@ def accounting_snapshot(self) -> dict[str, int]: "current_partition_quota": self._current_partition_quota, } + def transfer_timing_snapshot(self, *, clear: bool = False) -> list[dict[str, Any]]: + """Return raw per-batch transfer timings for median/p95 reporting.""" + + snapshot = [copy.deepcopy(record) for record in self._transfer_timing_records] + if clear: + self._transfer_timing_records.clear() + return snapshot + def resident_group_keys(self) -> set[GroupKey]: return {sample_group_key(group) for group in self.ready_group_buffer} @@ -304,6 +902,8 @@ async def discard_pending_transfers(self) -> tuple[int, int]: cancelled_tasks += 1 if tasks: await asyncio.gather(*tasks, return_exceptions=True) + self._partition_actual_rows.clear() + self._partition_identity_state.clear() return dropped_buffer_groups, cancelled_tasks async def drain_ready_group_payloads(self) -> tuple[list[list], int]: @@ -377,3 +977,5 @@ async def shutdown(self) -> None: self._committed_current_group_count = 0 self._dispatched_previous_group_count = 0 self._dispatched_current_group_count = 0 + self._partition_actual_rows.clear() + self._partition_identity_state.clear() diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 75143efa6..528a3b758 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -1,11 +1,14 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. +import hashlib import logging import os import random import socket import time from argparse import Namespace +from collections.abc import Mapping +from contextlib import nullcontext from functools import partial from typing import Any, List @@ -81,11 +84,14 @@ from .collective_utils import _agree_drained from .cp_utils import all_gather_with_cp, maybe_padded_total_lengths, slice_with_cp from .data import ( + ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY, ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY, + ROLLOUT_VALID_ROW_FLAGS_KEY, DataIterator, build_rollout_minibatch_plan, concat_rollout_batches, get_data_iterator, + get_rollout_valid_row_flags, log_perf_data, log_perf_data_fwd, log_rollout_data, @@ -105,6 +111,465 @@ ROLLOUT_MINI_BATCH_METAS_KEY = "rollout_mini_batch_metas" +AGENTIC_VARIABLE_ROW_PADDING_KEY = "agentic_variable_row_padding" +AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV = "RELAX_AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE" +AGENTIC_ROW_IDENTITY_KEY = "agentic_row_identity" +AGENTIC_ROW_IDENTITY_TAG_KEY = "agentic_row_identity_tag" +AGENTIC_ROW_IDENTITY_TAGS_FIELD = "agentic_row_identity_tags" +AGENTIC_ROW_IDENTITY_SCHEMA_VERSION = 1 + + +def _use_agentic_variable_row_mode(args: Namespace) -> bool: + """Whether this actor must consume explicit agentic turn rows. + + Keep this predicate aligned with the producer-side gate without importing + the agentic pipeline into the Megatron backend. The feature is + deliberately default-off and async/hybrid modes use different streaming + contracts. + """ + + base_enabled = bool( + getattr(args, "group_rm", False) + and getattr(args, "agentic_custom_advantage_path", None) + and getattr(args, "use_dynamic_batch_size", False) + ) + if not base_enabled: + return False + + raw_max_rows = os.environ.get(AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV) + if raw_max_rows is None or not raw_max_rows.strip(): + return False + normalized_max_rows = raw_max_rows.strip() + if not normalized_max_rows.isascii() or not normalized_max_rows.isdigit() or int(normalized_max_rows) <= 0: + raise ValueError( + f"{AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV} must be a positive integer, got {raw_max_rows!r}." + ) + + if getattr(args, "fully_async", False) or getattr(args, "hybrid", False): + raise ValueError("Agentic variable-row training supports synchronous mode only.") + return True + + +def _agentic_variable_row_max_padded_rows(args: Namespace) -> int: + """Return the maximum physical rows allowed in one padded partition.""" + + raw_value = os.environ.get(AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV) + if raw_value is None or not raw_value.strip(): + raise ValueError( + f"{AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV} must be set to a positive integer " + "for agentic variable-row training." + ) + try: + max_rows_per_sample = int(raw_value) + except ValueError as exc: + raise ValueError( + f"{AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV} must be a positive integer, got {raw_value!r}." + ) from exc + if max_rows_per_sample <= 0: + raise ValueError(f"{AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV} must be a positive integer, got {raw_value!r}.") + + global_batch_size = int(args.global_batch_size) + max_actual_rows = int(args.rollout_batch_size) * int(args.n_samples_per_prompt) * max_rows_per_sample + if global_batch_size <= 0 or max_actual_rows <= 0: + raise ValueError( + "Agentic variable-row limits require positive global_batch_size, rollout_batch_size, " + f"and n_samples_per_prompt, got global_batch_size={global_batch_size}, " + f"rollout_batch_size={args.rollout_batch_size}, n_samples_per_prompt={args.n_samples_per_prompt}." + ) + return ((max_actual_rows + global_batch_size - 1) // global_batch_size) * global_batch_size + + +def _extract_agentic_variable_row_padding_flags(batch_meta: Any, expected_samples: int) -> list[bool]: + """Strictly decode the producer's per-row padding marker.""" + + if isinstance(expected_samples, bool) or not isinstance(expected_samples, int) or expected_samples <= 0: + raise ValueError(f"expected_samples must be a positive integer, got {expected_samples!r}") + get_all_custom_meta = getattr(batch_meta, "get_all_custom_meta", None) + if not callable(get_all_custom_meta): + raise RuntimeError("Agentic variable-row BatchMeta must provide get_all_custom_meta().") + all_custom_meta = get_all_custom_meta() + if not isinstance(all_custom_meta, list): + raise RuntimeError("Agentic variable-row BatchMeta.get_all_custom_meta() must return a list.") + if len(all_custom_meta) != expected_samples: + raise RuntimeError( + f"Agentic variable-row metadata size mismatch: expected {expected_samples}, got {len(all_custom_meta)}." + ) + + flags: list[bool] = [] + for sample_idx, custom_meta in enumerate(all_custom_meta): + if not isinstance(custom_meta, Mapping): + raise RuntimeError( + "Agentic variable-row custom metadata must be a mapping: " + f"sample_idx={sample_idx}, got {type(custom_meta).__name__}." + ) + if AGENTIC_VARIABLE_ROW_PADDING_KEY not in custom_meta: + raise RuntimeError( + f"Agentic variable-row sample metadata is missing {AGENTIC_VARIABLE_ROW_PADDING_KEY!r}: " + f"sample_idx={sample_idx}." + ) + flag = custom_meta[AGENTIC_VARIABLE_ROW_PADDING_KEY] + if type(flag) is not bool: + raise RuntimeError( + f"Agentic variable-row padding marker must be bool: sample_idx={sample_idx}, got {flag!r}." + ) + flags.append(flag) + return flags + + +def _agentic_row_identity_tag(row_id: str) -> int: + digest = hashlib.sha256(row_id.encode("utf-8")).digest() + return int.from_bytes(digest[:8], byteorder="big") & ((1 << 63) - 1) + + +def _agentic_identity_tag_value(value: Any, *, context: str) -> int: + if isinstance(value, bool): + raise RuntimeError(f"{context} must be an int64 scalar, got {value!r}.") + item = getattr(value, "item", None) + if callable(item): + device = getattr(value, "device", None) + if device is not None and getattr(device, "type", "cpu") != "cpu": + raise RuntimeError(f"{context} must remain CPU-resident during identity validation.") + value = item() + if isinstance(value, bool) or not isinstance(value, int): + raise RuntimeError(f"{context} must be an int64 scalar, got {value!r}.") + return value + + +def _extract_agentic_variable_row_identities( + batch_meta: Any, + rollout_data: Mapping[str, Any], +) -> list[dict[str, Any]]: + """Decode and bind queue custom metadata to the fetched tensor rows.""" + + total_lengths = rollout_data.get("total_lengths") + response_lengths = rollout_data.get("response_lengths") + row_tags = rollout_data.get(AGENTIC_ROW_IDENTITY_TAGS_FIELD) + if not isinstance(total_lengths, list) or not isinstance(response_lengths, list): + raise RuntimeError("Agentic identity validation requires list total_lengths and response_lengths.") + expected_samples = len(total_lengths) + if expected_samples <= 0 or len(response_lengths) != expected_samples: + raise RuntimeError( + "Agentic identity data field size mismatch: " + f"total_lengths={expected_samples}, response_lengths={len(response_lengths)}." + ) + if not isinstance(row_tags, (list, tuple)) or len(row_tags) != expected_samples: + raise RuntimeError( + f"{AGENTIC_ROW_IDENTITY_TAGS_FIELD} must contain one CPU tag per fetched row: " + f"expected={expected_samples}, got={len(row_tags) if isinstance(row_tags, (list, tuple)) else None}." + ) + + get_all_custom_meta = getattr(batch_meta, "get_all_custom_meta", None) + if not callable(get_all_custom_meta): + raise RuntimeError("Agentic variable-row BatchMeta must provide get_all_custom_meta().") + all_custom_meta = get_all_custom_meta() + if not isinstance(all_custom_meta, list) or len(all_custom_meta) != expected_samples: + raise RuntimeError( + "Agentic identity custom metadata size mismatch: " + f"expected={expected_samples}, got={len(all_custom_meta) if isinstance(all_custom_meta, list) else None}." + ) + + identities: list[dict[str, Any]] = [] + for sample_idx, (custom_meta, data_tag) in enumerate(zip(all_custom_meta, row_tags, strict=True)): + if not isinstance(custom_meta, Mapping): + raise RuntimeError(f"Agentic identity custom metadata row {sample_idx} must be a mapping.") + identity = custom_meta.get(AGENTIC_ROW_IDENTITY_KEY) + if not isinstance(identity, Mapping): + raise RuntimeError(f"Agentic identity row {sample_idx} is missing {AGENTIC_ROW_IDENTITY_KEY!r}.") + identity = dict(identity) + if identity.get("schema_version") != AGENTIC_ROW_IDENTITY_SCHEMA_VERSION: + raise RuntimeError( + f"Agentic identity row {sample_idx} has unsupported schema_version {identity.get('schema_version')!r}." + ) + row_id = identity.get("row_id") + if not isinstance(row_id, str) or not row_id: + raise RuntimeError(f"Agentic identity row {sample_idx} has an invalid row_id {row_id!r}.") + padding = identity.get("padding") + if type(padding) is not bool or custom_meta.get(AGENTIC_VARIABLE_ROW_PADDING_KEY) is not padding: + raise RuntimeError(f"Agentic identity row {sample_idx} has inconsistent padding markers.") + expected_tag = _agentic_row_identity_tag(row_id) + meta_tag = _agentic_identity_tag_value( + custom_meta.get(AGENTIC_ROW_IDENTITY_TAG_KEY), + context=f"custom metadata row {sample_idx} identity tag", + ) + fetched_tag = _agentic_identity_tag_value( + data_tag, + context=f"fetched row {sample_idx} identity tag", + ) + if meta_tag != expected_tag or fetched_tag != expected_tag: + raise RuntimeError( + f"TransferQueue row/custom-metadata identity mismatch: sample_idx={sample_idx}, row_id={row_id!r}." + ) + total_length = int(total_lengths[sample_idx]) + response_length = int(response_lengths[sample_idx]) + if custom_meta.get("total_lengths") != total_length or identity.get("total_length") != total_length: + raise RuntimeError(f"Agentic identity row {sample_idx} has inconsistent total_length.") + if identity.get("response_length") != response_length: + raise RuntimeError(f"Agentic identity row {sample_idx} has inconsistent response_length.") + action_token_count = identity.get("action_token_count") + if isinstance(action_token_count, bool) or not isinstance(action_token_count, int) or action_token_count < 0: + raise RuntimeError(f"Agentic identity row {sample_idx} has invalid action_token_count.") + if padding: + if action_token_count != 0 or any( + identity.get(key) is not None + for key in ("rollout_group_id", "task_id", "trajectory_id", "turn_id", "turn_index") + ): + raise RuntimeError(f"Agentic padding row {sample_idx} leaks a real row identity.") + else: + required_keys = ( + "rollout_group_id", + "policy_version", + "task_id", + "trajectory_id", + "turn_id", + "turn_index", + "sample_index", + "terminal", + "truncated", + "group_row_count", + "group_trajectory_count", + "group_row_ids_sha256", + "partition_expected_group_count", + ) + missing = [key for key in required_keys if identity.get(key) is None] + if missing: + raise RuntimeError(f"Agentic identity row {sample_idx} is missing fields {missing!r}.") + if action_token_count <= 0: + raise RuntimeError(f"Agentic real row {sample_idx} contains no trainable action tokens.") + identities.append(identity) + return identities + + +def _validate_agentic_variable_row_partition_identities( + identities: list[Mapping[str, Any]], + *, + expected_group_count: int, + expected_trajectories_per_group: int, +) -> dict[str, Any]: + """Validate a globally gathered partition manifest independent of row + order.""" + + if not identities: + raise RuntimeError("Agentic identity partition is empty.") + row_ids = [identity.get("row_id") for identity in identities] + if any(not isinstance(row_id, str) or not row_id for row_id in row_ids): + raise RuntimeError("Agentic identity partition contains an invalid row_id.") + if len(set(row_ids)) != len(row_ids): + raise RuntimeError("Agentic identity partition contains duplicate row_id values.") + real_rows = [dict(identity) for identity in identities if identity.get("padding") is False] + padding_rows = [dict(identity) for identity in identities if identity.get("padding") is True] + if len(real_rows) + len(padding_rows) != len(identities) or not real_rows: + raise RuntimeError("Agentic identity partition contains invalid padding markers or no real rows.") + policy_versions = {identity.get("policy_version") for identity in identities} + if None in policy_versions or len(policy_versions) != 1: + raise RuntimeError(f"Agentic identity partition mixes policy versions: {policy_versions!r}.") + + groups: dict[str, list[dict[str, Any]]] = {} + for identity in real_rows: + group_id = identity.get("rollout_group_id") + if not isinstance(group_id, str) or not group_id: + raise RuntimeError(f"Agentic real row has invalid rollout_group_id {group_id!r}.") + groups.setdefault(group_id, []).append(identity) + if len(groups) != expected_group_count: + raise RuntimeError( + "Agentic identity partition has a residual or missing rollout group: " + f"expected={expected_group_count}, got={len(groups)}." + ) + + trajectory_to_group: dict[str, str] = {} + for group_id, rows in groups.items(): + tasks = {identity.get("task_id") for identity in rows} + policies = {identity.get("policy_version") for identity in rows} + expected_group_counts = {identity.get("partition_expected_group_count") for identity in rows} + row_counts = {identity.get("group_row_count") for identity in rows} + trajectory_counts = {identity.get("group_trajectory_count") for identity in rows} + digests = {identity.get("group_row_ids_sha256") for identity in rows} + if len(tasks) != 1 or len(policies) != 1: + raise RuntimeError(f"Agentic rollout group {group_id!r} mixes task or policy identity.") + if expected_group_counts != {expected_group_count}: + raise RuntimeError(f"Agentic rollout group {group_id!r} has a conflicting partition group count.") + if row_counts != {len(rows)}: + raise RuntimeError(f"Agentic rollout group {group_id!r} is missing or has duplicate physical rows.") + if trajectory_counts != {expected_trajectories_per_group}: + raise RuntimeError(f"Agentic rollout group {group_id!r} has a conflicting trajectory count.") + if digests != { + hashlib.sha256("\n".join(sorted(identity["row_id"] for identity in rows)).encode()).hexdigest() + }: + raise RuntimeError(f"Agentic rollout group {group_id!r} row manifest digest does not match.") + + by_trajectory: dict[str, list[dict[str, Any]]] = {} + sample_index_to_trajectory: dict[int, str] = {} + for identity in rows: + trajectory_id = identity.get("trajectory_id") + sample_index = identity.get("sample_index") + if not isinstance(trajectory_id, str) or not trajectory_id: + raise RuntimeError(f"Agentic rollout group {group_id!r} has an invalid trajectory_id.") + if isinstance(sample_index, bool) or not isinstance(sample_index, int) or sample_index < 0: + raise RuntimeError(f"Agentic rollout group {group_id!r} has an invalid sample_index.") + previous_group = trajectory_to_group.setdefault(trajectory_id, group_id) + if previous_group != group_id: + raise RuntimeError(f"trajectory_id {trajectory_id!r} appears in multiple rollout groups.") + previous_trajectory = sample_index_to_trajectory.setdefault(sample_index, trajectory_id) + if previous_trajectory != trajectory_id: + raise RuntimeError(f"Sample.index {sample_index} maps to multiple trajectories.") + by_trajectory.setdefault(trajectory_id, []).append(identity) + if len(by_trajectory) != expected_trajectories_per_group: + raise RuntimeError( + f"Agentic rollout group {group_id!r} expected {expected_trajectories_per_group} trajectories, " + f"got {len(by_trajectory)}." + ) + for trajectory_id, rows_for_trajectory in by_trajectory.items(): + ordered = sorted(rows_for_trajectory, key=lambda identity: identity.get("turn_index", -1)) + turns = [identity.get("turn_index") for identity in ordered] + if turns != list(range(len(ordered))): + raise RuntimeError(f"Agentic trajectory {trajectory_id!r} has a missing or duplicate turn row.") + if any(identity.get("turn_id") != f"turn_{identity['turn_index']:03d}" for identity in ordered): + raise RuntimeError(f"Agentic trajectory {trajectory_id!r} has a non-canonical turn_id.") + if any(identity.get("terminal") or identity.get("truncated") for identity in ordered[:-1]): + raise RuntimeError(f"Agentic trajectory {trajectory_id!r} terminates before its final row.") + final = ordered[-1] + if type(final.get("terminal")) is not bool or type(final.get("truncated")) is not bool: + raise RuntimeError(f"Agentic trajectory {trajectory_id!r} has invalid final status markers.") + if final["terminal"] == final["truncated"]: + raise RuntimeError( + f"Agentic trajectory {trajectory_id!r} final row must be exactly terminal or truncated." + ) + return { + "policy_version": next(iter(policy_versions)), + "group_count": len(groups), + "real_row_count": len(real_rows), + "padding_row_count": len(padding_rows), + } + + +def _validate_agentic_variable_row_window( + *, + actual_global_rows: int, + global_batch_size: int, + total_padded_rows_before: int, + max_padded_rows: int, +) -> int: + """Validate one consumed full window and return cumulative physical + rows.""" + + for name, value in ( + ("actual_global_rows", actual_global_rows), + ("global_batch_size", global_batch_size), + ("total_padded_rows_before", total_padded_rows_before), + ("max_padded_rows", max_padded_rows), + ): + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an int, got {type(value).__name__}") + if global_batch_size <= 0: + raise ValueError(f"global_batch_size must be positive, got {global_batch_size}") + if actual_global_rows <= 0 or actual_global_rows > global_batch_size: + raise RuntimeError( + "Agentic variable-row window must contain between 1 and global_batch_size real rows, " + f"got actual_global_rows={actual_global_rows}, global_batch_size={global_batch_size}." + ) + if total_padded_rows_before < 0: + raise ValueError(f"total_padded_rows_before must be non-negative, got {total_padded_rows_before}") + if max_padded_rows <= 0 or max_padded_rows % global_batch_size != 0: + raise ValueError( + "max_padded_rows must be a positive multiple of global_batch_size, " + f"got max_padded_rows={max_padded_rows}, global_batch_size={global_batch_size}." + ) + + total_padded_rows = total_padded_rows_before + global_batch_size + if total_padded_rows > max_padded_rows: + raise RuntimeError( + "Agentic variable-row consumer exceeded the configured partition bound: " + f"consumed_padded_rows={total_padded_rows}, max_padded_rows={max_padded_rows}." + ) + return total_padded_rows + + +def _validate_agentic_variable_row_partition( + *, + total_physical_rows: int, + total_actual_rows: int, + global_batch_size: int, +) -> int: + """Validate the producer's partition-level modulo-padding contract.""" + + for name, value in ( + ("total_physical_rows", total_physical_rows), + ("total_actual_rows", total_actual_rows), + ("global_batch_size", global_batch_size), + ): + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an int, got {type(value).__name__}") + if global_batch_size <= 0: + raise ValueError(f"global_batch_size must be positive, got {global_batch_size}") + if total_physical_rows <= 0 or total_physical_rows % global_batch_size != 0: + raise RuntimeError( + "Agentic variable-row partition must contain a positive multiple of global_batch_size " + f"physical rows, got total_physical_rows={total_physical_rows}, " + f"global_batch_size={global_batch_size}." + ) + if total_actual_rows <= 0 or total_actual_rows > total_physical_rows: + raise RuntimeError( + "Agentic variable-row partition has an invalid real-row total: " + f"total_actual_rows={total_actual_rows}, total_physical_rows={total_physical_rows}." + ) + total_padding_rows = total_physical_rows - total_actual_rows + if total_padding_rows >= global_batch_size: + raise RuntimeError( + "Agentic variable-row partition contains at least one global batch of padding: " + f"total_padding_rows={total_padding_rows}, global_batch_size={global_batch_size}." + ) + return total_padding_rows + + +def _agentic_variable_row_drain_action( + *, + ready_min: int, + ready_max: int, + stream_drained: bool, + accepted_windows: int, +) -> str: + """Pure state transition for the DP-aligned variable-row drain loop.""" + + if ready_min not in (0, 1) or ready_max not in (0, 1) or ready_min > ready_max: + raise ValueError(f"Invalid DP readiness range: min={ready_min}, max={ready_max}.") + if isinstance(stream_drained, bool) is False: + raise TypeError(f"stream_drained must be bool, got {type(stream_drained).__name__}") + if isinstance(accepted_windows, bool) or not isinstance(accepted_windows, int) or accepted_windows < 0: + raise ValueError(f"accepted_windows must be a non-negative integer, got {accepted_windows!r}.") + if ready_min == 1: + if stream_drained: + raise RuntimeError("A variable-row batch is pending after the partition reported drained.") + return "accept" + if ready_max == 1: + return "wait" + if stream_drained: + if accepted_windows == 0: + raise RuntimeError("Variable-row partition drained before yielding any training window.") + return "done" + return "retry" + + +def _agentic_variable_row_stream_drained_consensus( + *, + stream_drained: bool, + device: Any, + dp_group: Any, +) -> bool: + """Return true only after every data-parallel rank observes a drained + stream.""" + + if type(stream_drained) is not bool: + raise TypeError(f"stream_drained must be bool, got {type(stream_drained).__name__}") + drained_consensus = torch.tensor( + [int(stream_drained)], + dtype=torch.int64, + device=device, + ) + dist.all_reduce( + drained_consensus, + op=dist.ReduceOp.MIN, + group=dp_group, + ) + return bool(drained_consensus.item()) def _split_by_rollout_mini_counts(values: Any, counts: list[int]) -> list[Any]: @@ -609,7 +1074,217 @@ def _request_rollout_evaluation(self, rollout_id: int, *, end_update_weight: boo """Backward-compatible name kept for existing internal call sites.""" self._run_step_evaluation(rollout_id, end_update_weight=end_update_weight) + def _collect_agentic_variable_row_batches( + self, + *, + rollout_id: int, + task_name: str, + data_fields: list[str], + process_batch=None, + ) -> tuple[list[RolloutBatch], list, list[int], list[int], list[bool]]: + """Drain a padded variable-row partition into full local GBS + windows.""" + + if self.role != "actor": + raise NotImplementedError("Agentic variable-row training currently supports the GRPO actor only.") + if getattr(self.args, "debug_train_only", False): + raise NotImplementedError("debug_train_only does not support agentic variable-row training.") + if getattr(self.args, "advantage_estimator", None) != "grpo": + raise NotImplementedError( + "Agentic variable-row training currently supports advantage_estimator=grpo only." + ) + if getattr(self.args, "use_critic", False) or getattr(self.args, "use_opd", False): + raise NotImplementedError("Agentic variable-row training does not yet support critic or OPD paths.") + if getattr(self.args, "multimodal_keys", None) is not None: + raise NotImplementedError("Agentic variable-row training currently supports text-only samples.") + if mpu.get_context_parallel_world_size() != 1: + raise NotImplementedError("Agentic variable-row loss scaling currently requires context_parallel_size=1.") + + dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False) + global_batch_size = int(self.args.global_batch_size) + if global_batch_size <= 0 or global_batch_size % dp_size != 0: + raise ValueError( + "Agentic variable-row training requires global_batch_size divisible by data-parallel size, " + f"got global_batch_size={global_batch_size}, dp_size={dp_size}." + ) + local_batch_size = global_batch_size // dp_size + max_padded_rows = _agentic_variable_row_max_padded_rows(self.args) + partition_id = f"train_{rollout_id}" + + batches: list[RolloutBatch] = [] + batch_metas: list = [] + local_sample_counts: list[int] = [] + global_sample_counts: list[int] = [] + local_valid_row_flags: list[bool] = [] + local_partition_identities: list[dict[str, Any]] = [] + identity_data_fields = list(data_fields) + if AGENTIC_ROW_IDENTITY_TAGS_FIELD not in identity_data_fields: + identity_data_fields.append(AGENTIC_ROW_IDENTITY_TAGS_FIELD) + batch_index = 0 + total_padded_rows = 0 + empty_poll_sleep_s = float(os.environ.get("RELAX_EMPTY_POLL_SLEEP_MS", "50")) / 1000.0 + fetch_iter = 0 + timeout_s = max(float(getattr(self.args, "distributed_timeout_minutes", 30)) * 60.0, 1.0) + last_progress_at = time.monotonic() + pending_data = None + pending_meta = None + + while True: + if time.monotonic() - last_progress_at > timeout_s: + raise TimeoutError( + "Timed out draining agentic variable-row partition: " + f"partition_id={partition_id}, task_name={task_name}, batch_index={batch_index}." + ) + + if pending_data is None: + with timer("train_get_data"): + pending_data, pending_meta = self._get_data_from_transfer_queue( + task_name, + rollout_id, + identity_data_fields, + local_batch_size, + batch_index, + partition_id=partition_id, + ) + + has_pending = 1 if pending_data is not None else 0 + ready_min = torch.tensor( + [has_pending], + dtype=torch.int64, + device=device_utils.make_current_torch_device(), + ) + ready_max = ready_min.clone() + dp_group = mpu.get_data_parallel_group(with_context_parallel=False) + dist.all_reduce(ready_min, op=dist.ReduceOp.MIN, group=dp_group) + dist.all_reduce(ready_max, op=dist.ReduceOp.MAX, group=dp_group) + + ready_min_value = int(ready_min.item()) + ready_max_value = int(ready_max.item()) + stream_drained = False + if ready_max_value == 0: + stream_drained = bool( + self.all_consumed( + task_name, + rollout_id, + partition_id=partition_id, + streaming=True, + ) + ) + stream_drained = _agentic_variable_row_stream_drained_consensus( + stream_drained=stream_drained, + device=device_utils.make_current_torch_device(), + dp_group=dp_group, + ) + drain_action = _agentic_variable_row_drain_action( + ready_min=ready_min_value, + ready_max=ready_max_value, + stream_drained=stream_drained, + accepted_windows=len(batches), + ) + if drain_action == "done": + break + if drain_action != "accept": + if fetch_iter % 100 == 0: + logger.info( + "[agentic variable rows] rollout_id=%s batch_index=%s task=%s: " + "waiting until every DP rank has one full local batch.", + rollout_id, + batch_index, + task_name, + ) + fetch_iter += 1 + if empty_poll_sleep_s > 0: + time.sleep(empty_poll_sleep_s) + continue + + rollout_data, batch_meta = pending_data, pending_meta + pending_data = None + pending_meta = None + num_local_rows = len(rollout_data.get("total_lengths", [])) + if num_local_rows != local_batch_size: + raise RuntimeError( + "Agentic variable-row sampler returned a non-full local batch: " + f"rollout_id={rollout_id}, batch_index={batch_index}, " + f"expected={local_batch_size}, got={num_local_rows}." + ) + padding_flags = _extract_agentic_variable_row_padding_flags(batch_meta, num_local_rows) + row_identities = _extract_agentic_variable_row_identities(batch_meta, rollout_data) + rollout_data.pop(AGENTIC_ROW_IDENTITY_TAGS_FIELD, None) + if [identity["padding"] for identity in row_identities] != padding_flags: + raise RuntimeError("Agentic identity and padding metadata disagree after TransferQueue reorder.") + local_partition_identities.extend(row_identities) + local_actual_rows = num_local_rows - sum(padding_flags) + actual_rows_tensor = torch.tensor( + [local_actual_rows], + dtype=torch.int64, + device=device_utils.make_current_torch_device(), + ) + dist.all_reduce( + actual_rows_tensor, + op=dist.ReduceOp.SUM, + group=mpu.get_data_parallel_group(with_context_parallel=False), + ) + actual_global_rows = int(actual_rows_tensor.item()) + total_padded_rows = _validate_agentic_variable_row_window( + actual_global_rows=actual_global_rows, + global_batch_size=global_batch_size, + total_padded_rows_before=total_padded_rows, + max_padded_rows=max_padded_rows, + ) + + if process_batch is not None: + process_batch(rollout_data) + batches.append(rollout_data) + batch_metas.append(batch_meta) + local_sample_counts.append(num_local_rows) + global_sample_counts.append(actual_global_rows) + local_valid_row_flags.extend(not is_padding for is_padding in padding_flags) + # SeqlenBalancedSampler may reorder the producer's final padding rows + # across physical windows. Padding therefore does not identify a + # terminal window; the partition drain signal is the sole end marker. + batch_index += 1 + fetch_iter = 0 + last_progress_at = time.monotonic() + + if not batches: + raise RuntimeError( + "Agentic variable-row partition drained without yielding a training window: " + f"partition_id={partition_id}, task_name={task_name}." + ) + _validate_agentic_variable_row_partition( + total_physical_rows=total_padded_rows, + total_actual_rows=sum(global_sample_counts), + global_batch_size=global_batch_size, + ) + if dp_size == 1: + global_partition_identities = local_partition_identities + else: + gathered_identities: list[Any] = [None] * dp_size + dist.all_gather_object( + gathered_identities, + local_partition_identities, + group=mpu.get_data_parallel_group(with_context_parallel=False), + ) + global_partition_identities = [] + for dp_rank, rank_identities in enumerate(gathered_identities): + if not isinstance(rank_identities, list): + raise RuntimeError( + "Agentic identity all-gather returned a non-list payload: " + f"dp_rank={dp_rank}, type={type(rank_identities).__name__}." + ) + global_partition_identities.extend(rank_identities) + _validate_agentic_variable_row_partition_identities( + global_partition_identities, + expected_group_count=int(self.args.rollout_batch_size), + expected_trajectories_per_group=int(self.args.n_samples_per_prompt), + ) + return batches, batch_metas, local_sample_counts, global_sample_counts, local_valid_row_flags + def train(self, rollout_id: int) -> None: + variable_row_mode = _use_agentic_variable_row_mode(self.args) + if variable_row_mode and self.args.debug_train_only: + raise NotImplementedError("debug_train_only does not support agentic variable-row training.") + if self.args.offload_rollout and dist.get_rank() == 0: pre_train_offload_handles = [] if self.genrm_manager is not None: @@ -657,6 +1332,27 @@ def train(self, rollout_id: int) -> None: return self.train_actor(rollout_id, rollout_data) else: logger.info(f"start to get rollout_id: {rollout_id} data from transfer queue for train with mcore.") + if variable_row_mode: + task_name = sft_task_name(self.args, component="backend") + data_fields = build_data_fields(self.args, consumer="actor") + ( + rollout_mini_batches, + rollout_mini_batch_metas, + rollout_mini_local_sample_counts, + rollout_mini_global_sample_counts, + rollout_valid_row_flags, + ) = self._collect_agentic_variable_row_batches( + rollout_id=rollout_id, + task_name=task_name, + data_fields=data_fields, + ) + rollout_data = concat_rollout_batches(rollout_mini_batches) + rollout_data[ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY] = rollout_mini_local_sample_counts + rollout_data[ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY] = rollout_mini_global_sample_counts + rollout_data[ROLLOUT_MINI_BATCH_METAS_KEY] = rollout_mini_batch_metas + rollout_data[ROLLOUT_VALID_ROW_FLAGS_KEY] = rollout_valid_row_flags + return self.train_actor(rollout_id, rollout_data) + if is_sft_mode(self.args): batch_size = self.args.global_batch_size // mpu.get_data_parallel_world_size( with_context_parallel=False @@ -895,6 +1591,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: self.opt_param_scheduler, data_iterator, num_microbatches, + scheduler_increments=rollout_data.get(ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY), ) self.prof.step(rollout_id=rollout_id) @@ -920,7 +1617,20 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: logger.info(f"Updating ref model at rollout_id {rollout_id}") self.weights_backuper.backup("ref") - total_lengths = rollout_data["total_lengths"] + valid_row_flags = get_rollout_valid_row_flags(rollout_data) + if valid_row_flags is None: + metric_total_lengths = rollout_data["total_lengths"] + metric_loss_masks = rollout_data["loss_masks"] + else: + metric_total_lengths = [ + length + for length, is_valid in zip(rollout_data["total_lengths"], valid_row_flags, strict=True) + if is_valid + ] + metric_loss_masks = [ + mask for mask, is_valid in zip(rollout_data["loss_masks"], valid_row_flags, strict=True) if is_valid + ] + total_lengths = metric_total_lengths all_total_lengths = [None] * mpu.get_data_parallel_world_size(with_context_parallel=False) dist.all_gather_object( all_total_lengths, total_lengths, group=mpu.get_data_parallel_group(with_context_parallel=False) @@ -931,7 +1641,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: # assistant-only tokens; for RL it's the response-only mask sum. Avoids # the SFT `response_length == total_length` convention (see data.py:296). response_token_counts = [ - int(m.sum().item()) if isinstance(m, torch.Tensor) else int(sum(m)) for m in rollout_data["loss_masks"] + int(m.sum().item()) if isinstance(m, torch.Tensor) else int(sum(m)) for m in metric_loss_masks ] all_response_token_counts = [None] * mpu.get_data_parallel_world_size(with_context_parallel=False) dist.all_gather_object( @@ -1239,13 +1949,27 @@ def train_hybrid(self, rollout_id) -> None: global normalization across the full batch and DP group. """ logger.info(f"start to get rollout_id: {rollout_id} data from transfer queue for train_hybrid.") + variable_row_mode = _use_agentic_variable_row_mode(self.args) + if variable_row_mode and self.args.debug_train_only: + raise NotImplementedError("debug_train_only does not support agentic variable-row training.") + dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False) - plan = build_rollout_minibatch_plan(self.args, dp_size) - batch_size = plan.mini_local_sample_request + if variable_row_mode: + if self.args.global_batch_size % dp_size != 0: + raise ValueError( + "Agentic variable-row training requires global_batch_size divisible by data-parallel size, " + f"got global_batch_size={self.args.global_batch_size}, dp_size={dp_size}." + ) + plan = None + batch_size = self.args.global_batch_size // dp_size + else: + plan = build_rollout_minibatch_plan(self.args, dp_size) + batch_size = plan.mini_local_sample_request # ── Phase 1: Collect sub-batches and compute ref/actor forward in small chunks ── collected_batches: list[RolloutBatch] = [] rollout_mini_local_sample_counts: list[int] = [] + rollout_mini_global_sample_counts: list[int] = [] if self.args.debug_train_only: # Bypass the transfer queue and load the offline debug rollout dump # directly (mirrors `train`'s debug_train_only path). The dump holds @@ -1266,7 +1990,34 @@ def train_hybrid(self, rollout_id) -> None: collected_batches.append(sub_batch) rollout_mini_local_sample_counts.append(len(sub_batch["total_lengths"])) else: - batch_index = 0 + if variable_row_mode: + variable_data_fields = [ + "tokens", + "total_lengths", + "response_lengths", + "loss_masks", + "rollout_log_probs", + "rewards", + "raw_reward", + ] + variable_data_fields += ["rollout_routed_experts"] if self.args.use_rollout_routing_replay else [] + if self.args.multimodal_keys is not None: + variable_data_fields.append("multimodal_train_inputs") + if self.args.use_opd and self.args.opd_type == "sglang": + variable_data_fields.append("teacher_log_probs") + ( + collected_batches, + _batch_metas, + rollout_mini_local_sample_counts, + rollout_mini_global_sample_counts, + ) = self._collect_agentic_variable_row_batches( + rollout_id=rollout_id, + task_name="train", + data_fields=variable_data_fields, + process_batch=self._hybrid_forward_subbatch, + ) + plan = Namespace(num_rollout_minis=len(collected_batches)) + batch_index = plan.num_rollout_minis if variable_row_mode else 0 # Surface stuck-loop conditions: when the partition can never reach the # requested batch_size (e.g. rollout dropped samples without refilling), # `get_meta` keeps returning size=0 while `all_consumed` stays False, @@ -1344,6 +2095,8 @@ def train_hybrid(self, rollout_id) -> None: else: rollout_data[key].append(value) rollout_data[ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY] = rollout_mini_local_sample_counts + if variable_row_mode: + rollout_data[ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY] = rollout_mini_global_sample_counts with inverse_timer("train_wait"), timer("train"): if self.args.compute_advantages_and_returns: diff --git a/relax/backends/megatron/data.py b/relax/backends/megatron/data.py index 9aff3da6f..5232af845 100644 --- a/relax/backends/megatron/data.py +++ b/relax/backends/megatron/data.py @@ -41,6 +41,53 @@ ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY = "rollout_mini_local_sample_counts" ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY = "rollout_mini_global_sample_counts" ROLLOUT_MINI_PROMPT_GROUP_COUNTS_KEY = "rollout_mini_prompt_group_counts" +ROLLOUT_VALID_ROW_FLAGS_KEY = "agentic_variable_row_valid_rows" + + +def get_rollout_valid_row_flags(rollout_data: RolloutBatch) -> list[bool] | None: + """Return strict per-row validity flags for variable-row batches. + + Ordinary fixed-row batches do not carry the field and retain their legacy + accounting. Variable-row batches persist it after TransferQueue metadata + has been decoded so synthetic tail padding can never become a metric row. + """ + + flags = rollout_data.get(ROLLOUT_VALID_ROW_FLAGS_KEY) + if flags is None: + return None + if not isinstance(flags, list): + raise RuntimeError(f"{ROLLOUT_VALID_ROW_FLAGS_KEY} must be a list, got {type(flags).__name__}.") + expected_rows = len(rollout_data.get("total_lengths", [])) + if len(flags) != expected_rows: + raise RuntimeError( + f"{ROLLOUT_VALID_ROW_FLAGS_KEY} must contain one flag per row: expected={expected_rows}, got={len(flags)}." + ) + if any(type(flag) is not bool for flag in flags): + raise RuntimeError(f"{ROLLOUT_VALID_ROW_FLAGS_KEY} must contain only bool values.") + return flags + + +def _rollout_metrics_view(rollout_data: RolloutBatch) -> tuple[RolloutBatch, int | None]: + """Build a shallow metrics-only view without synthetic variable-row + pads.""" + + flags = get_rollout_valid_row_flags(rollout_data) + if flags is None: + return rollout_data, None + valid_indices = [idx for idx, is_valid in enumerate(flags) if is_valid] + row_count = len(flags) + metrics_view = {} + for key, value in rollout_data.items(): + if key == ROLLOUT_VALID_ROW_FLAGS_KEY: + continue + if isinstance(value, (list, tuple)) and len(value) == row_count: + metrics_view[key] = [value[idx] for idx in valid_indices] + elif isinstance(value, torch.Tensor) and value.ndim > 0 and value.size(0) == row_count: + index = torch.tensor(valid_indices, dtype=torch.long, device=value.device) + metrics_view[key] = value.index_select(0, index) + else: + metrics_view[key] = value + return metrics_view, len(valid_indices) @dataclass(frozen=True) @@ -52,6 +99,150 @@ class RolloutMiniBatchPlan: mini_local_sample_request: int | None +@dataclass(frozen=True) +class VariableRowMiniBatchWindow: + """One fixed-global-batch training window over variable rollout rows. + + ``actual_rows_per_dp`` describes the real rows assigned to each data + parallel rank. Callers must append ``padding_rows_per_dp`` rows whose loss + masks are all zero before training the window. + """ + + start_row: int + stop_row: int + actual_global_rows: int + padded_global_rows: int + actual_rows_per_dp: tuple[int, ...] + padded_rows_per_dp: tuple[int, ...] + padding_rows_per_dp: tuple[int, ...] + + +@dataclass(frozen=True) +class VariableRowMiniBatchPlan: + """Fixed-global-batch windows that consume every variable rollout row.""" + + actual_global_rows: int + global_batch_size: int + windows: tuple[VariableRowMiniBatchWindow, ...] + total_padded_global_rows: int + total_padding_rows: int + + +def _require_positive_int(name: str, value: int) -> None: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an int, got {type(value).__name__}") + if value <= 0: + raise ValueError(f"{name} must be positive, got {value}") + + +def build_variable_row_minibatch_plan( + actual_global_rows: int, + global_batch_size: int, + dp_size: int, + micro_batch_size: int, +) -> VariableRowMiniBatchPlan: + """Partition variable rollout rows into fixed-global-batch windows. + + The existing fixed-n rollout plan is intentionally not involved. Every + window keeps Megatron's configured global batch size; only the final window + may contain padding. Real rows are assigned to DP ranks as evenly as + possible, and the remaining per-rank rows must be materialized with a zero- + valued loss mask by the caller. + """ + + for name, value in ( + ("actual_global_rows", actual_global_rows), + ("global_batch_size", global_batch_size), + ("dp_size", dp_size), + ("micro_batch_size", micro_batch_size), + ): + _require_positive_int(name, value) + + padded_window_multiple = dp_size * micro_batch_size + if global_batch_size % padded_window_multiple != 0: + raise ValueError( + "global_batch_size must be divisible by dp_size * micro_batch_size, " + f"got global_batch_size={global_batch_size}, dp_size={dp_size}, " + f"micro_batch_size={micro_batch_size}" + ) + + padded_rows_per_rank = global_batch_size // dp_size + windows: list[VariableRowMiniBatchWindow] = [] + start_row = 0 + while start_row < actual_global_rows: + stop_row = min(start_row + global_batch_size, actual_global_rows) + actual_window_rows = stop_row - start_row + actual_rows_base, ranks_with_extra_row = divmod(actual_window_rows, dp_size) + actual_rows_per_dp = tuple(actual_rows_base + (rank < ranks_with_extra_row) for rank in range(dp_size)) + padded_rows_per_dp = (padded_rows_per_rank,) * dp_size + padding_rows_per_dp = tuple(padded_rows_per_rank - actual_rows for actual_rows in actual_rows_per_dp) + + windows.append( + VariableRowMiniBatchWindow( + start_row=start_row, + stop_row=stop_row, + actual_global_rows=actual_window_rows, + padded_global_rows=global_batch_size, + actual_rows_per_dp=actual_rows_per_dp, + padded_rows_per_dp=padded_rows_per_dp, + padding_rows_per_dp=padding_rows_per_dp, + ) + ) + start_row = stop_row + + total_padded_global_rows = len(windows) * global_batch_size + return VariableRowMiniBatchPlan( + actual_global_rows=actual_global_rows, + global_batch_size=global_batch_size, + windows=tuple(windows), + total_padded_global_rows=total_padded_global_rows, + total_padding_rows=total_padded_global_rows - actual_global_rows, + ) + + +def _build_rollout_mini_loss_scales( + *, + num_microbatches: Sequence[int], + actual_global_sample_counts: Sequence[int] | None, + global_batch_size: int, + dp_size: int, +) -> list[float] | None: + """Build one explicit loss scale per micro-batch for variable-row + windows.""" + + if actual_global_sample_counts is None: + return None + if len(actual_global_sample_counts) != len(num_microbatches): + raise ValueError( + f"{ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY} length must match num_microbatches windows, " + f"got counts={len(actual_global_sample_counts)}, windows={len(num_microbatches)}." + ) + if isinstance(global_batch_size, bool) or not isinstance(global_batch_size, int) or global_batch_size <= 0: + raise ValueError(f"global_batch_size must be a positive integer, got {global_batch_size!r}") + if isinstance(dp_size, bool) or not isinstance(dp_size, int) or dp_size <= 0: + raise ValueError(f"dp_size must be a positive integer, got {dp_size!r}") + + scales: list[float] = [] + for window_idx, (num_mbs, actual_rows) in enumerate( + zip(num_microbatches, actual_global_sample_counts, strict=True) + ): + if isinstance(num_mbs, bool) or not isinstance(num_mbs, int) or num_mbs <= 0: + raise ValueError(f"num_microbatches[{window_idx}] must be a positive integer, got {num_mbs!r}") + if ( + isinstance(actual_rows, bool) + or not isinstance(actual_rows, int) + or actual_rows <= 0 + or actual_rows > global_batch_size + ): + raise ValueError( + f"{ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY}[{window_idx}] must be in " + f"[1, {global_batch_size}], got {actual_rows!r}." + ) + scale = num_mbs * dp_size / actual_rows + scales.extend([scale] * num_mbs) + return scales + + def build_rollout_minibatch_plan(args: Namespace, dp_size: int) -> RolloutMiniBatchPlan: """Build a prompt-group based mini plan for one rollout partition.""" if dp_size <= 0: @@ -591,6 +782,7 @@ def gather_log_data( args: Namespace, rollout_id: int, log_dict: dict[str, float], + metric_row_count: int | None = None, ) -> dict[str, float] | None: """Gather per-rank metrics, reduce by mean on the DP source rank, and log. @@ -606,15 +798,27 @@ def gather_log_data( gathered_log_dict = [None] * dp_size # Not sure if this will be a performance bottleneck. dist.gather_object( - log_dict, + (log_dict, metric_row_count) if metric_row_count is not None else log_dict, gathered_log_dict, dst=mpu.get_data_parallel_src_rank(with_context_parallel=True), group=mpu.get_data_parallel_group_gloo(with_context_parallel=True), ) - reduced_log_dict = { - f"{metric_name}/{key}": sum([d[key] for d in gathered_log_dict]) / dp_size for key in log_dict - } + if metric_row_count is None: + reduced_log_dict = { + f"{metric_name}/{key}": sum([d[key] for d in gathered_log_dict]) / dp_size for key in log_dict + } + else: + if isinstance(metric_row_count, bool) or not isinstance(metric_row_count, int) or metric_row_count < 0: + raise ValueError(f"metric_row_count must be a non-negative integer, got {metric_row_count!r}") + total_rows = sum(row_count for _, row_count in gathered_log_dict) + if total_rows <= 0: + raise RuntimeError("Variable-row metrics contain no real rows across the data-parallel group.") + reduced_log_dict = { + f"{metric_name}/{key}": sum(metrics[key] * row_count for metrics, row_count in gathered_log_dict) + / total_rows + for key in log_dict + } logger.info(f"{metric_name} {rollout_id}: {reduced_log_dict}") # Calculate step once to avoid duplication @@ -625,7 +829,7 @@ def gather_log_data( return reduced_log_dict else: dist.gather_object( - log_dict, + (log_dict, metric_row_count) if metric_row_count is not None else log_dict, None, dst=mpu.get_data_parallel_src_rank(with_context_parallel=True), group=mpu.get_data_parallel_group_gloo(with_context_parallel=True), @@ -646,6 +850,7 @@ def __init__( micro_batch_size: int | None = None, micro_batch_indices: list[list[int]] | None = None, max_tokens_per_gpu: int | None = None, + micro_batch_loss_scales: list[float] | None = None, ) -> None: """Initialize an iterator over `rollout_data`. @@ -658,15 +863,38 @@ def __init__( reads it in get_batch to pick each mb's CP size consistently with how the micro-batches were packed (forward-only uses log_probs_max_tokens_per_gpu, training uses max_tokens_per_gpu). None falls back to args.max_tokens_per_gpu. + micro_batch_loss_scales: Optional explicit scale for every micro-batch. + Used only by agentic variable-row windows; ordinary paths leave it unset. """ self.rollout_data = rollout_data self.micro_batch_size = micro_batch_size self.micro_batch_indices = micro_batch_indices self.max_tokens_per_gpu = max_tokens_per_gpu + self.micro_batch_loss_scales = micro_batch_loss_scales assert micro_batch_size is None or micro_batch_indices is None self.offset = 0 + self.micro_batch_offset = 0 - def get_next(self, keys: Sequence[str]) -> dict[str, list[object] | None]: + if micro_batch_loss_scales is not None: + if micro_batch_indices is not None: + expected_microbatches = len(micro_batch_indices) + else: + if micro_batch_size is None or micro_batch_size <= 0: + raise ValueError("micro_batch_size must be positive when explicit loss scales are provided") + num_local_samples = len(rollout_data.get("total_lengths", [])) + if num_local_samples % micro_batch_size != 0: + raise ValueError( + "Local sample count must be divisible by micro_batch_size when explicit loss scales " + f"are provided, got samples={num_local_samples}, micro_batch_size={micro_batch_size}." + ) + expected_microbatches = num_local_samples // micro_batch_size + if len(micro_batch_loss_scales) != expected_microbatches: + raise ValueError( + "micro_batch_loss_scales length must equal the iterator micro-batch count, " + f"got scales={len(micro_batch_loss_scales)}, expected={expected_microbatches}." + ) + + def get_next(self, keys: Sequence[str]) -> dict[str, object]: """Return the next micro-batch for the requested keys. - If `micro_batch_indices` is provided, selects rows according to the current @@ -691,15 +919,25 @@ def get_next(self, keys: Sequence[str]) -> dict[str, list[object] | None]: ) batch[key] = vals[self.offset : self.offset + self.micro_batch_size] + if self.micro_batch_loss_scales is not None: + if self.micro_batch_offset >= len(self.micro_batch_loss_scales): + raise RuntimeError( + "DataIterator consumed more micro-batches than its explicit loss-scale schedule: " + f"offset={self.micro_batch_offset}, scales={len(self.micro_batch_loss_scales)}." + ) + batch["__loss_scale__"] = self.micro_batch_loss_scales[self.micro_batch_offset] + if self.micro_batch_indices is not None: self.offset += 1 else: self.offset += self.micro_batch_size + self.micro_batch_offset += 1 return batch def reset(self) -> "DataIterator": """Reset internal offset to the start and return self.""" self.offset = 0 + self.micro_batch_offset = 0 return self @@ -741,6 +979,7 @@ def get_data_iterator( global_batch_size = rollout_data.get("dynamic_global_batch_size", args.global_batch_size) num_local_gbs = global_batch_size // dp_size step_local_sample_counts = rollout_data.get(ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY) + step_global_sample_counts = rollout_data.get(ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY) if step_local_sample_counts is not None: if not isinstance(step_local_sample_counts, list) or not step_local_sample_counts: @@ -767,6 +1006,30 @@ def get_data_iterator( ) num_steps_per_rollout = num_local_samples // num_local_gbs + if step_global_sample_counts is not None: + if step_local_sample_counts is None: + raise ValueError( + f"{ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY} requires " + f"{ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY} to preserve rollout-window boundaries." + ) + if not isinstance(step_global_sample_counts, list) or not step_global_sample_counts: + raise ValueError(f"{ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY} must be a non-empty list") + if len(step_global_sample_counts) != num_steps_per_rollout: + raise ValueError( + f"{ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY} length must equal rollout windows, " + f"got counts={len(step_global_sample_counts)}, windows={num_steps_per_rollout}." + ) + invalid_global_counts = [ + count + for count in step_global_sample_counts + if isinstance(count, bool) or not isinstance(count, int) or count <= 0 or count > args.global_batch_size + ] + if invalid_global_counts: + raise ValueError( + f"{ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY} must contain integers in " + f"[1, {args.global_batch_size}], got {step_global_sample_counts}." + ) + # With balance_data, synchronise num_steps_per_rollout across DP ranks so # collectives stay aligned. Explicit rollout-mini boundaries or divisible # fixed-size local steps must already produce the same count on every rank; @@ -790,10 +1053,24 @@ def get_data_iterator( f"num_local_samples={num_local_samples}, num_steps_per_rollout={num_steps_per_rollout}" ) - def _generate_data_iterator(rollout_data, micro_batch_size, micro_batch_indices=None, max_tokens_per_gpu=None): + def _generate_data_iterator( + rollout_data, + micro_batch_size, + micro_batch_indices=None, + max_tokens_per_gpu=None, + micro_batch_loss_scales=None, + ): data_iterator = [] for _ in range(vpp_size): - data_iterator.append(DataIterator(rollout_data, micro_batch_size, micro_batch_indices, max_tokens_per_gpu)) + data_iterator.append( + DataIterator( + rollout_data, + micro_batch_size, + micro_batch_indices, + max_tokens_per_gpu, + micro_batch_loss_scales, + ) + ) return data_iterator if step_local_sample_counts is None: @@ -807,7 +1084,17 @@ def _generate_data_iterator(rollout_data, micro_batch_size, micro_batch_indices= f"got invalid_counts={invalid_counts}, micro_batch_size={args.micro_batch_size}" ) num_microbatches = [count // args.micro_batch_size for count in step_local_sample_counts] - data_iterator = _generate_data_iterator(rollout_data, args.micro_batch_size) + micro_batch_loss_scales = _build_rollout_mini_loss_scales( + num_microbatches=num_microbatches, + actual_global_sample_counts=step_global_sample_counts, + global_batch_size=int(args.global_batch_size), + dp_size=dp_size, + ) + data_iterator = _generate_data_iterator( + rollout_data, + args.micro_batch_size, + micro_batch_loss_scales=micro_batch_loss_scales, + ) else: _max_tokens = max_tokens_per_gpu if max_tokens_per_gpu is not None else args.max_tokens_per_gpu assert _max_tokens is not None @@ -862,7 +1149,19 @@ def _generate_data_iterator(rollout_data, micro_batch_size, micro_batch_indices= logger.info( f"After dynamic batching, num_microbatches: {num_microbatches}, micro_batch_indices: {micro_batch_indices}" ) - data_iterator = _generate_data_iterator(rollout_data, None, micro_batch_indices, _max_tokens) + micro_batch_loss_scales = _build_rollout_mini_loss_scales( + num_microbatches=num_microbatches, + actual_global_sample_counts=step_global_sample_counts, + global_batch_size=int(args.global_batch_size), + dp_size=dp_size, + ) + data_iterator = _generate_data_iterator( + rollout_data, + None, + micro_batch_indices, + _max_tokens, + micro_batch_loss_scales, + ) return ( data_iterator, @@ -884,6 +1183,8 @@ def log_rollout_data( - Non-tensor lists are averaged elementwise. - Scalars are converted to Python numbers. """ + source_rollout_data = rollout_data + rollout_data, metric_row_count = _rollout_metrics_view(source_rollout_data) if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage(): # Under dynamic CP, rollout_data was merged back to full-length responses # (dynamic_cp_merge_output), so metrics here run on full data — use cp=1 (no @@ -894,12 +1195,16 @@ def log_rollout_data( loss_masks = rollout_data["loss_masks"] total_lengths = rollout_data["total_lengths"] max_seq_lens = rollout_data.get("max_seq_lens", None) - padded_total_lengths = maybe_padded_total_lengths( - total_lengths, - args.qkv_format, - getattr(args, "is_vl_model", False) - or rollout_data.get("multimodal_train_inputs") is not None - or getattr(args, "uses_unsplit_forward", False), + padded_total_lengths = ( + maybe_padded_total_lengths( + total_lengths, + args.qkv_format, + getattr(args, "is_vl_model", False) + or rollout_data.get("multimodal_train_inputs") is not None + or getattr(args, "uses_unsplit_forward", False), + ) + if total_lengths + else None ) for key, val in rollout_data.items(): @@ -925,6 +1230,15 @@ def log_rollout_data( # There are the following assumptions: # - Each dp rank has the same number of samples if isinstance(val, (list, tuple)): + source_val = source_rollout_data.get(key) + if not val: + if ( + isinstance(source_val, (list, tuple)) + and source_val + and isinstance(source_val[0], (torch.Tensor, int, float)) + ): + log_dict[key] = 0.0 + continue if isinstance(val[0], torch.Tensor): # NOTE: Here we have to do the clone().detach(), otherwise the tensor will be # modified in place and will cause problem for the next rollout. @@ -964,23 +1278,25 @@ def log_rollout_data( continue val = sum(val) / len(val) elif isinstance(val, torch.Tensor): - val = val.float().mean() + val = 0.0 if metric_row_count == 0 and val.numel() == 0 else val.float().mean() else: continue log_dict[key] = val.item() if isinstance(val, torch.Tensor) else val - if total_lengths: + if total_lengths or metric_row_count is not None: dp_group = mpu.get_data_parallel_group(with_context_parallel=True) + local_max = max(total_lengths) if total_lengths else -1 + local_neg_min = -min(total_lengths) if total_lengths else -(1 << 62) stats = torch.tensor( - [max(total_lengths), -min(total_lengths)], + [local_max, local_neg_min], dtype=torch.int64, - device=loss_masks[0].device, + device=(loss_masks[0] if loss_masks else source_rollout_data["loss_masks"][0]).device, ) dist.all_reduce(stats, op=dist.ReduceOp.MAX, group=dp_group) log_dict["total_lengths/max"] = int(stats[0].item()) log_dict["total_lengths/min"] = -int(stats[1].item()) - reduced_log_dict = gather_log_data("rollout", args, rollout_id, log_dict) + reduced_log_dict = gather_log_data("rollout", args, rollout_id, log_dict, metric_row_count=metric_row_count) if args.ci_test and reduced_log_dict is not None: if ( rollout_id == 0 @@ -996,7 +1312,7 @@ def log_rollout_data( assert 0 < reduced_log_dict["rollout/entropy"] < 0.5 if args.log_multi_turn: - log_multi_turn_data(rollout_id, args, rollout_data) + log_multi_turn_data(rollout_id, args, source_rollout_data) if args.log_correct_samples: if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage(): @@ -1087,6 +1403,8 @@ def log_multi_turn_data(rollout_id: int, args: Namespace, rollout_data: RolloutB Operates only on PP last stage and TP rank 0. Uses GPU tensors when available to compute statistics without host transfers. """ + source_rollout_data = rollout_data + rollout_data, metric_row_count = _rollout_metrics_view(source_rollout_data) if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage(): log_dict = {} for key, val in rollout_data.items(): @@ -1110,13 +1428,26 @@ def log_multi_turn_data(rollout_id: int, args: Namespace, rollout_data: RolloutB log_dict["wo_obs_response_length/response_length_mean"] = wo_obs_response_lengths.mean().item() log_dict["wo_obs_response_length/response_length_max"] = wo_obs_response_lengths.max().item() log_dict["wo_obs_response_length/response_length_min"] = wo_obs_response_lengths.min().item() + elif source_rollout_data.get("loss_masks"): + log_dict["raw_response_length/response_length_mean"] = 0.0 + log_dict["raw_response_length/response_length_max"] = 0.0 + log_dict["raw_response_length/response_length_min"] = 0.0 + log_dict["raw_response_length/response_length_clip_ratio"] = 0.0 + log_dict["wo_obs_response_length/response_length_mean"] = 0.0 + log_dict["wo_obs_response_length/response_length_max"] = 0.0 + log_dict["wo_obs_response_length/response_length_min"] = 0.0 if key == "round_number": # Use numpy for vectorized round number statistics - round_number_array = np.array(val) - log_dict["multi_turn_metric/round_number_mean"] = np.mean(round_number_array) - log_dict["multi_turn_metric/round_number_max"] = np.max(round_number_array) - log_dict["multi_turn_metric/round_number_min"] = np.min(round_number_array) - gather_log_data("multi_turn", args, rollout_id, log_dict) + if val: + round_number_array = np.array(val) + log_dict["multi_turn_metric/round_number_mean"] = np.mean(round_number_array) + log_dict["multi_turn_metric/round_number_max"] = np.max(round_number_array) + log_dict["multi_turn_metric/round_number_min"] = np.min(round_number_array) + elif source_rollout_data.get("round_number"): + log_dict["multi_turn_metric/round_number_mean"] = 0.0 + log_dict["multi_turn_metric/round_number_max"] = 0.0 + log_dict["multi_turn_metric/round_number_min"] = 0.0 + gather_log_data("multi_turn", args, rollout_id, log_dict, metric_row_count=metric_row_count) def log_perf_data_fwd(args, rollout_id): diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py index 372a139cf..6b54ace37 100644 --- a/relax/backends/megatron/loss.py +++ b/relax/backends/megatron/loss.py @@ -42,6 +42,14 @@ maybe_padded_total_lengths, slice_log_prob_with_cp, ) +from .data import get_rollout_valid_row_flags + + +def _metric_num_samples(batch: RolloutBatch) -> int: + """Count only real rows when variable-row tail padding is present.""" + + valid_row_flags = get_rollout_valid_row_flags(batch) + return sum(valid_row_flags) if valid_row_flags is not None else len(batch["response_lengths"]) def get_responses( @@ -1357,7 +1365,7 @@ def loss_function( dynamic_cp_size=batch.get("dynamic_cp_size", None), dynamic_cp_rank=batch.get("dynamic_cp_rank", None), ) - num_samples = len(batch["response_lengths"]) + num_samples = _metric_num_samples(batch) sum_of_sample_mean = get_sum_of_sample_mean( batch["total_lengths"], diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 077e461d9..af1798d49 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -45,13 +45,15 @@ ) from .checkpoint import load_checkpoint, save_checkpoint -from .data import DataIterator, get_batch +from .data import ROLLOUT_VALID_ROW_FLAGS_KEY, DataIterator, get_batch from .loss import loss_function from .model_provider import get_model_provider_func, wrap_model_provider_with_freeze logger = get_logger(__name__) +AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV = "RELAX_AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE" + def _find_lm_output_layer(model: torch.nn.Module) -> torch.nn.Module | None: """Walk DDP / bridge-VL wrappers to the lm_head; None on non-last PP @@ -221,6 +223,42 @@ def _main_loss_has_tokens(batch: dict) -> bool: return bool(num_tokens.item() > 0) +def _optimizer_scheduler_training_plan(args: Namespace) -> tuple[int, int]: + """Return ``(train_iters, default_horizon_rows)`` for the scheduler. + + Fixed-row training preserves the legacy trajectory-count formula. The opt- + in synchronous variable-row path advances its scheduler by real turn rows, + so its default LR/WD horizon uses the configured maximum exported rows per + trajectory. This is a deterministic upper bound; early terminal + trajectories simply consume fewer scheduler rows. + """ + + global_batch_size = int(args.global_batch_size) + num_rollout = int(args.num_rollout) + trajectories_per_rollout = int(args.rollout_batch_size) * int(args.n_samples_per_prompt) + legacy_train_iters = num_rollout * trajectories_per_rollout // global_batch_size + + variable_row_enabled = bool( + getattr(args, "group_rm", False) + and getattr(args, "agentic_custom_advantage_path", None) + and getattr(args, "use_dynamic_batch_size", False) + ) + raw_max_rows = os.environ.get(AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV) + if not variable_row_enabled or raw_max_rows is None or not raw_max_rows.strip(): + return legacy_train_iters, legacy_train_iters * global_batch_size + if getattr(args, "fully_async", False) or getattr(args, "hybrid", False): + raise ValueError("Agentic variable-row scheduler supports synchronous training only.") + + normalized_max_rows = raw_max_rows.strip() + if not normalized_max_rows.isascii() or not normalized_max_rows.isdigit() or int(normalized_max_rows) <= 0: + raise ValueError( + f"{AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV} must be a positive integer, got {raw_max_rows!r}." + ) + max_rows_per_rollout = trajectories_per_rollout * int(normalized_max_rows) + max_windows_per_rollout = (max_rows_per_rollout + global_batch_size - 1) // global_batch_size + return num_rollout * max_windows_per_rollout, num_rollout * max_rows_per_rollout + + def get_optimizer_param_scheduler(args: Namespace, optimizer: MegatronOptimizer) -> OptimizerParamScheduler: """Create and configure the optimizer learning-rate/weight-decay scheduler. @@ -234,12 +272,15 @@ def get_optimizer_param_scheduler(args: Namespace, optimizer: MegatronOptimizer) Returns: OptimizerParamScheduler: Initialized scheduler bound to ``optimizer``. """ - # Iteration-based training. - args.train_iters = args.num_rollout * args.rollout_batch_size * args.n_samples_per_prompt // args.global_batch_size + # Iteration-based training. Variable-row mode uses real turn rows as the + # scheduler unit; fixed-row mode remains byte-for-byte equivalent. + args.train_iters, default_horizon_steps = _optimizer_scheduler_training_plan(args) if args.lr_decay_iters is None: args.lr_decay_iters = args.train_iters - lr_decay_steps = args.lr_decay_iters * args.global_batch_size - wd_incr_steps = args.train_iters * args.global_batch_size + lr_decay_steps = default_horizon_steps + else: + lr_decay_steps = args.lr_decay_iters * args.global_batch_size + wd_incr_steps = default_horizon_steps wsd_decay_steps = None if args.lr_wsd_decay_iters is not None: wsd_decay_steps = args.lr_wsd_decay_iters * args.global_batch_size @@ -952,6 +993,7 @@ def train_one_step( optimizer: MegatronOptimizer, opt_param_scheduler: OptimizerParamScheduler, num_microbatches: int, + scheduler_increment: int | None = None, ) -> tuple[dict[str, float], float]: """Execute a single pipeline-parallel training step. @@ -967,12 +1009,21 @@ def train_one_step( optimizer (MegatronOptimizer): Optimizer instance. opt_param_scheduler (OptimizerParamScheduler): LR/WD scheduler. num_microbatches (int): Number of microbatches to process. + scheduler_increment (int | None): Real global rows in this variable-row + window. ``None`` keeps the fixed-row global-batch increment. Returns: tuple[dict[str, float], float]: Reduced loss dictionary (last stage only) and gradient norm for logging. """ args = get_args() + resolved_scheduler_increment = args.global_batch_size if scheduler_increment is None else scheduler_increment + if ( + isinstance(resolved_scheduler_increment, bool) + or not isinstance(resolved_scheduler_increment, int) + or resolved_scheduler_increment <= 0 + ): + raise ValueError(f"scheduler increment must be a positive integer, got {resolved_scheduler_increment!r}") # Set grad to zero. for model_chunk in model: @@ -1029,6 +1080,7 @@ def forward_step( "returns", "rollout_log_probs", "max_seq_lens", + ROLLOUT_VALID_ROW_FLAGS_KEY, *_opd_keys, ], args.data_pad_size_multiplier, @@ -1214,7 +1266,7 @@ def forward_step( if valid_step: # Update learning rate. assert update_successful - opt_param_scheduler.step(increment=args.global_batch_size) + opt_param_scheduler.step(increment=resolved_scheduler_increment) else: grad_norm = float("nan") @@ -1265,6 +1317,95 @@ def should_disable_forward_pre_hook(args: Namespace) -> bool: return args.use_distributed_optimizer and args.overlap_param_gather +def _normalize_scheduler_increments( + scheduler_increments: Sequence[int] | None, + num_steps_per_rollout: int, +) -> tuple[int, ...] | None: + """Validate optional per-window real-row scheduler increments.""" + + if scheduler_increments is None: + return None + if len(scheduler_increments) != num_steps_per_rollout: + raise ValueError( + "scheduler_increments length must match the rollout training windows, " + f"got increments={len(scheduler_increments)}, windows={num_steps_per_rollout}." + ) + normalized = tuple(scheduler_increments) + if any( + isinstance(increment, bool) or not isinstance(increment, int) or increment <= 0 for increment in normalized + ): + raise ValueError(f"scheduler_increments must contain positive integers, got {scheduler_increments!r}") + return normalized + + +def _initial_variable_row_tracking_step(args: Namespace, opt_param_scheduler) -> int: + """Resume the monotonic row-based tracking cursor from actor or scheduler + state.""" + + tracking_step = getattr(args, "_agentic_variable_row_tracking_samples", None) + if tracking_step is None: + tracking_step = int(getattr(opt_param_scheduler, "num_steps", 0)) + if tracking_step < 0: + raise ValueError(f"_agentic_variable_row_tracking_samples must be non-negative, got {tracking_step}.") + return tracking_step + + +def _advance_variable_row_tracking_step(args: Namespace, tracking_step: int, increment: int) -> int: + """Consume one real-row window and persist the next collision-free + cursor.""" + + next_tracking_step = tracking_step + increment + setattr(args, "_agentic_variable_row_tracking_samples", next_tracking_step) + return next_tracking_step + + +def _restore_agentic_variable_row_resume_cursor( + args: Namespace, + opt_param_scheduler: OptimizerParamScheduler | None, + iteration: int, +) -> bool: + """Restore the real-row cursor from a loaded native scheduler checkpoint. + + Megatron persists ``OptimizerParamScheduler.num_steps`` in the native + checkpoint and restores it in ``load_state_dict``. That counter is the + exact real-row cursor used by agentic variable-row training. A resumed + variable-row run must reuse it and must not apply the legacy + ``iteration * global_batch_size`` catch-up step afterwards. + + Returns ``True`` only when a non-zero variable-row checkpoint was restored. + Missing, invalid, or conflicting state fails closed. + """ + + raw_max_rows = os.environ.get(AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV) + variable_row_enabled = bool( + getattr(args, "group_rm", False) + and getattr(args, "agentic_custom_advantage_path", None) + and getattr(args, "use_dynamic_batch_size", False) + and raw_max_rows is not None + and raw_max_rows.strip() + ) + if not variable_row_enabled or iteration <= 0: + return False + + if opt_param_scheduler is None: + raise RuntimeError("Agentic variable-row resume requires a restored optimizer scheduler.") + restored_step = getattr(opt_param_scheduler, "num_steps", None) + if isinstance(restored_step, bool) or not isinstance(restored_step, int) or restored_step <= 0: + raise RuntimeError( + "Agentic variable-row resume requires a positive checkpoint scheduler " + f"num_steps value, got {restored_step!r} at iteration {iteration}." + ) + + existing_step = getattr(args, "_agentic_variable_row_tracking_samples", None) + if existing_step is not None and existing_step != restored_step: + raise RuntimeError( + "Agentic variable-row resume cursor conflicts with the loaded scheduler: " + f"tracking={existing_step}, scheduler={restored_step}." + ) + setattr(args, "_agentic_variable_row_tracking_samples", restored_step) + return True + + def train( rollout_id: int, model: Sequence[DDP], @@ -1272,6 +1413,7 @@ def train( opt_param_scheduler: OptimizerParamScheduler, data_iterator: Sequence[DataIterator], num_microbatches: Sequence[int], + scheduler_increments: Sequence[int] | None = None, ) -> None: """Run training over a rollout consisting of multiple steps. @@ -1285,6 +1427,8 @@ def train( opt_param_scheduler (OptimizerParamScheduler): LR/WD scheduler. data_iterator (Sequence[DataIterator]): Iterable(s) yielding training batches. num_microbatches (Sequence[int]): Microbatches per step in the rollout. + scheduler_increments (Sequence[int] | None): Per-window real global row + counts for variable-row training. ``None`` preserves legacy behavior. """ args = get_args() is_data_iterator = isinstance(data_iterator[0], DataIterator) @@ -1371,6 +1515,9 @@ def train( ) num_steps_per_rollout = len(num_microbatches) + scheduler_increments = _normalize_scheduler_increments(scheduler_increments, num_steps_per_rollout) + if scheduler_increments is not None: + variable_row_tracking_step = _initial_variable_row_tracking_step(args, opt_param_scheduler) use_step_iterators = ( not is_data_iterator and len(data_iterator) > 1 and isinstance(data_iterator[0], StreamingTQIterator) ) @@ -1385,6 +1532,7 @@ def train( step_data_iterator = [data_iterator[step_id]] if use_step_iterators else data_iterator # Run training step. with timer(f"train_micro_batch_{step_id}", keep=False): + scheduler_increment = None if scheduler_increments is None else scheduler_increments[step_id] loss_dict, grad_norm = train_one_step( args, rollout_id, @@ -1394,6 +1542,14 @@ def train( optimizer, opt_param_scheduler, num_microbatches[step_id], + scheduler_increment=scheduler_increment, + ) + if scheduler_increments is None: + accumulated_step_id = rollout_id * num_steps_per_rollout + step_id + else: + accumulated_step_id = variable_row_tracking_step + variable_row_tracking_step = _advance_variable_row_tracking_step( + args, variable_row_tracking_step, scheduler_increment ) if keep_forward_pre_hook_disabled: force_param_sync(model) @@ -1434,7 +1590,6 @@ def train( and mpu.get_tensor_model_parallel_rank() == 0 and mpu.get_pipeline_model_parallel_rank() == mpu.get_pipeline_model_parallel_world_size() - 1 ): - accumulated_step_id = rollout_id * num_steps_per_rollout + step_id role = getattr(model[0], "role", "actor") role_tag = "" if role == "actor" else f"{role}-" log_dict = { @@ -1454,9 +1609,14 @@ def train( log_dict["train/step"] = accumulated_step_id num_per_epoch = getattr(args, "num_rollout_per_epoch", None) if num_per_epoch: - log_dict[f"train/{role_tag}cur_epoch"] = (accumulated_step_id + 1) / ( - num_per_epoch * num_steps_per_rollout - ) + if scheduler_increments is None: + log_dict[f"train/{role_tag}cur_epoch"] = (accumulated_step_id + 1) / ( + num_per_epoch * num_steps_per_rollout + ) + else: + log_dict[f"train/{role_tag}cur_epoch"] = ( + rollout_id + (step_id + 1) / num_steps_per_rollout + ) / num_per_epoch tracking_utils.log(args, log_dict, step_key="train/step") tracking_utils.flush_metrics(args, accumulated_step_id) @@ -1765,6 +1925,13 @@ def initialize_model_and_optimizer( checkpointing_context={}, skip_load_to_model_and_opt=False, ) + restored_variable_row_cursor = _restore_agentic_variable_row_resume_cursor(args, opt_param_scheduler, iteration) + if restored_variable_row_cursor: + logger.info( + "Restored agentic variable-row cursor %d from checkpoint iteration %d", + args._agentic_variable_row_tracking_samples, + iteration, + ) if role == "critic": release_critic_lm_heads(model) loaded_value_head_param_ids = validate_critic_value_head_registration(model, optimizer) @@ -1773,5 +1940,7 @@ def initialize_model_and_optimizer( ) install_critic_value_head_runtime_check(model) clear_memory() + if opt_param_scheduler is not None and not restored_variable_row_cursor: + opt_param_scheduler.step(increment=iteration * args.global_batch_size) return model, optimizer, opt_param_scheduler, iteration diff --git a/relax/core/controller.py b/relax/core/controller.py index e59f5cfd7..115719bbc 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -26,6 +26,7 @@ ) from e from relax.agentic.pipeline.runtime import clear_agentic_runtime_caches +from relax.agentic.pipeline.transfer import use_agentic_variable_row_mode from relax.agentic.session.service import ( deploy_agentic_chat_api_services, shutdown_agentic_chat_api_services, @@ -97,6 +98,39 @@ def _can_cleanup_s3_model_weights(config: Namespace, serve_dict: dict) -> bool: return getattr(config, "sglang_load_format", "auto") in _S3_MODEL_CLEANUP_SAFE_LOAD_FORMATS +AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV = "RELAX_AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE" + + +def _agentic_max_exported_rows_per_sample() -> int: + raw_value = os.environ.get(AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV) + if raw_value is None or not raw_value.strip(): + raise ValueError( + f"{AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV} must be set to a positive integer " + "when agentic variable-row transfer is enabled." + ) + normalized_value = raw_value.strip() + if not normalized_value.isascii() or not normalized_value.isdigit(): + raise ValueError(f"{AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV} must be a positive integer, got {raw_value!r}.") + try: + value = int(normalized_value) + except ValueError as exc: + raise ValueError( + f"{AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV} must be a positive integer, got {raw_value!r}." + ) from exc + if value <= 0: + raise ValueError(f"{AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV} must be a positive integer, got {raw_value!r}.") + return value + + +def _agentic_variable_row_storage_size(config: Namespace) -> int: + max_exported_rows = _agentic_max_exported_rows_per_sample() + max_actual_rows = config.rollout_batch_size * config.n_samples_per_prompt * max_exported_rows + padded_rows_per_partition = ( + (max_actual_rows + config.global_batch_size - 1) // config.global_batch_size + ) * config.global_batch_size + return padded_rows_per_partition * (config.max_staleness + 1) + + def _actor_rollout_pg_roles(config: Namespace) -> list[str]: """Roles that share the actor/rollout placement group. @@ -242,15 +276,31 @@ def _cleanup_s3_model_weights_after_init(self) -> None: def _initialize_data_system(self): algo_key = resolve_sft_algo_key(self.config) - batch_size_for_capacity = ( - self.config.over_sampling_batch_size - if self.config.partial_rollout and self.config.use_dynamic_global_batch_size - else self.config.rollout_batch_size - ) - total_storage_size = ( - batch_size_for_capacity * (self.config.max_staleness + 1) * self.config.n_samples_per_prompt - ) - if getattr(self.config, "fully_async", False) and getattr(self.config, "use_dynamic_batch_size", False): + variable_row_mode = use_agentic_variable_row_mode(self.config) + if variable_row_mode: + total_storage_size = _agentic_variable_row_storage_size(self.config) + else: + batch_size_for_capacity = ( + self.config.over_sampling_batch_size + if self.config.partial_rollout and self.config.use_dynamic_global_batch_size + else self.config.rollout_batch_size + ) + total_storage_size = ( + batch_size_for_capacity * (self.config.max_staleness + 1) * self.config.n_samples_per_prompt + ) + if variable_row_mode: + dp_size = compute_dp_size(self.config) + sampler = SeqlenBalancedSampler( + n_samples_per_prompt=1, + dp_size=dp_size, + ) + logger.info( + "Using SeqlenBalancedSampler for agentic variable-row transfer with " + "n_samples_per_prompt=1, dp_size=%s, total_storage_size=%s", + dp_size, + total_storage_size, + ) + elif getattr(self.config, "fully_async", False) and getattr(self.config, "use_dynamic_batch_size", False): # Fully-async + dynamic-batch path streams data per DP via token # budget; the controller-side sampler maintains per-DP buckets and # balances tokens at small-unit granularity. See diff --git a/tests/backends/megatron/test_agentic_variable_row_actor_helpers.py b/tests/backends/megatron/test_agentic_variable_row_actor_helpers.py new file mode 100644 index 000000000..569edf3c8 --- /dev/null +++ b/tests/backends/megatron/test_agentic_variable_row_actor_helpers.py @@ -0,0 +1,531 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import ast +import hashlib +import os +import time +from argparse import Namespace +from collections.abc import Mapping +from contextlib import nullcontext +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[3] +PADDING_KEY = "agentic_variable_row_padding" +IDENTITY_KEY = "agentic_row_identity" +IDENTITY_TAG_KEY = "agentic_row_identity_tag" +IDENTITY_TAGS_FIELD = "agentic_row_identity_tags" +MAX_ROWS_ENV = "RELAX_AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE" + + +def _load_helpers(): + path = REPO_ROOT / "relax" / "backends" / "megatron" / "actor.py" + names = { + "_use_agentic_variable_row_mode", + "_agentic_variable_row_max_padded_rows", + "_extract_agentic_variable_row_padding_flags", + "_agentic_row_identity_tag", + "_agentic_identity_tag_value", + "_extract_agentic_variable_row_identities", + "_validate_agentic_variable_row_partition_identities", + "_validate_agentic_variable_row_partition", + "_validate_agentic_variable_row_window", + "_agentic_variable_row_drain_action", + "_agentic_variable_row_stream_drained_consensus", + } + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + definitions = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in names] + assert {node.name for node in definitions} == names + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + *definitions, + ], + type_ignores=[], + ) + namespace = { + "AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV": MAX_ROWS_ENV, + "AGENTIC_VARIABLE_ROW_PADDING_KEY": PADDING_KEY, + "AGENTIC_ROW_IDENTITY_KEY": IDENTITY_KEY, + "AGENTIC_ROW_IDENTITY_TAG_KEY": IDENTITY_TAG_KEY, + "AGENTIC_ROW_IDENTITY_TAGS_FIELD": IDENTITY_TAGS_FIELD, + "AGENTIC_ROW_IDENTITY_SCHEMA_VERSION": 1, + "Any": Any, + "hashlib": hashlib, + "Mapping": Mapping, + "Namespace": Namespace, + "dist": None, + "os": os, + "torch": None, + } + exec(compile(ast.fix_missing_locations(module), str(path), "exec"), namespace) + return SimpleNamespace(**{name: namespace[name] for name in names}) + + +HELPERS = _load_helpers() + + +def _load_variable_row_collector(): + path = REPO_ROOT / "relax" / "backends" / "megatron" / "actor.py" + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + actor_class = next( + node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "MegatronTrainRayActor" + ) + collector = next( + node + for node in actor_class.body + if isinstance(node, ast.FunctionDef) and node.name == "_collect_agentic_variable_row_batches" + ) + host_class = ast.ClassDef( + name="_CollectorHost", + bases=[], + keywords=[], + body=[collector], + decorator_list=[], + ) + + class _Tensor: + def __init__(self, values): + self.value = values[0] + + def clone(self): + return _Tensor([self.value]) + + def item(self): + return self.value + + class _Torch: + int64 = object() + + @staticmethod + def tensor(values, *, dtype, device): + assert dtype is _Torch.int64 + assert device == "cuda:0" + return _Tensor(values) + + class _Dist: + class ReduceOp: + MIN = object() + MAX = object() + SUM = object() + + @staticmethod + def all_reduce(_tensor, *, op, group): + assert op in (_Dist.ReduceOp.MIN, _Dist.ReduceOp.MAX, _Dist.ReduceOp.SUM) + assert group == "dp-group" + + stream_drained_consensus = HELPERS._agentic_variable_row_stream_drained_consensus + stream_drained_consensus.__globals__["torch"] = _Torch + stream_drained_consensus.__globals__["dist"] = _Dist + + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + host_class, + ], + type_ignores=[], + ) + namespace = { + "_agentic_variable_row_drain_action": HELPERS._agentic_variable_row_drain_action, + "_agentic_variable_row_max_padded_rows": HELPERS._agentic_variable_row_max_padded_rows, + "_agentic_variable_row_stream_drained_consensus": stream_drained_consensus, + "_extract_agentic_variable_row_padding_flags": (HELPERS._extract_agentic_variable_row_padding_flags), + "_extract_agentic_variable_row_identities": (HELPERS._extract_agentic_variable_row_identities), + "_validate_agentic_variable_row_partition": HELPERS._validate_agentic_variable_row_partition, + "_validate_agentic_variable_row_partition_identities": ( + HELPERS._validate_agentic_variable_row_partition_identities + ), + "_validate_agentic_variable_row_window": HELPERS._validate_agentic_variable_row_window, + "AGENTIC_ROW_IDENTITY_TAGS_FIELD": IDENTITY_TAGS_FIELD, + "device_utils": SimpleNamespace(make_current_torch_device=lambda: "cuda:0"), + "dist": _Dist, + "logger": SimpleNamespace(info=lambda *args, **kwargs: None), + "mpu": SimpleNamespace( + get_context_parallel_world_size=lambda: 1, + get_data_parallel_group=lambda **_kwargs: "dp-group", + get_data_parallel_world_size=lambda **_kwargs: 1, + ), + "os": os, + "time": time, + "timer": lambda _name: nullcontext(), + "torch": _Torch, + } + exec(compile(ast.fix_missing_locations(module), str(path), "exec"), namespace) + return namespace["_CollectorHost"] + + +VARIABLE_ROW_COLLECTOR = _load_variable_row_collector() + + +def _mode_args(**overrides): + values = { + "group_rm": True, + "agentic_custom_advantage_path": "graphgpo.advantage", + "use_dynamic_batch_size": True, + "fully_async": False, + "hybrid": False, + } + values.update(overrides) + return Namespace(**values) + + +def _custom_meta_batch(custom_meta): + class _NonListSampleView: + pass + + class _BatchMeta: + samples = _NonListSampleView() + + def get_all_custom_meta(self): + return custom_meta + + return _BatchMeta() + + +def _batch_meta(flags): + return _custom_meta_batch([{PADDING_KEY: flag, "total_lengths": 7} for flag in flags]) + + +def _identity_tag(row_id: str) -> int: + digest = hashlib.sha256(row_id.encode("utf-8")).digest() + return int.from_bytes(digest[:8], byteorder="big") & ((1 << 63) - 1) + + +def _real_identity(*, group_id: str, turn_index: int, turn_count: int) -> dict[str, Any]: + row_ids = [f"row-{group_id}-{idx}" for idx in range(turn_count)] + return { + "schema_version": 1, + "padding": False, + "row_id": row_ids[turn_index], + "rollout_group_id": group_id, + "policy_version": "7", + "task_id": f"task-{group_id}", + "trajectory_id": f"trajectory-{group_id}", + "turn_id": f"turn_{turn_index:03d}", + "turn_index": turn_index, + "sample_index": int(group_id), + "terminal": turn_index == turn_count - 1, + "truncated": False, + "total_length": 7, + "response_length": 1, + "action_token_count": 1, + "group_row_count": turn_count, + "group_trajectory_count": 1, + "group_row_ids_sha256": hashlib.sha256("\n".join(sorted(row_ids)).encode()).hexdigest(), + "partition_expected_group_count": 2, + } + + +def _padding_identity(row_index: int) -> dict[str, Any]: + return { + "schema_version": 1, + "padding": True, + "row_id": f"padding-{row_index}", + "rollout_group_id": None, + "policy_version": "7", + "task_id": None, + "trajectory_id": None, + "turn_id": None, + "turn_index": None, + "sample_index": None, + "terminal": False, + "truncated": False, + "total_length": 7, + "response_length": 1, + "action_token_count": 0, + "group_row_count": 0, + "group_trajectory_count": 0, + "group_row_ids_sha256": None, + "partition_expected_group_count": 2, + } + + +def _identity_batch(identities: list[dict[str, Any]]): + custom_meta = [ + { + PADDING_KEY: identity["padding"], + "total_lengths": 7, + IDENTITY_KEY: identity, + IDENTITY_TAG_KEY: _identity_tag(identity["row_id"]), + } + for identity in identities + ] + rollout_data = { + "total_lengths": [7] * len(identities), + "response_lengths": [1] * len(identities), + IDENTITY_TAGS_FIELD: [_identity_tag(identity["row_id"]) for identity in identities], + } + return rollout_data, _custom_meta_batch(custom_meta) + + +def test_variable_row_mode_is_default_off_and_rejects_async_or_hybrid(monkeypatch): + monkeypatch.delenv(MAX_ROWS_ENV, raising=False) + assert HELPERS._use_agentic_variable_row_mode(_mode_args()) is False + + monkeypatch.setenv(MAX_ROWS_ENV, "50") + assert HELPERS._use_agentic_variable_row_mode(_mode_args()) is True + assert HELPERS._use_agentic_variable_row_mode(_mode_args(group_rm=False)) is False + assert HELPERS._use_agentic_variable_row_mode(_mode_args(agentic_custom_advantage_path=None)) is False + assert HELPERS._use_agentic_variable_row_mode(_mode_args(use_dynamic_batch_size=False)) is False + assert HELPERS._use_agentic_variable_row_mode(Namespace()) is False + + with pytest.raises(ValueError, match="synchronous mode only"): + HELPERS._use_agentic_variable_row_mode(_mode_args(fully_async=True)) + with pytest.raises(ValueError, match="synchronous mode only"): + HELPERS._use_agentic_variable_row_mode(_mode_args(fully_async=True, hybrid=True)) + with pytest.raises(ValueError, match="synchronous mode only"): + HELPERS._use_agentic_variable_row_mode(_mode_args(hybrid=True)) + + monkeypatch.setenv(MAX_ROWS_ENV, "not-a-number") + with pytest.raises(ValueError, match="positive integer"): + HELPERS._use_agentic_variable_row_mode(_mode_args()) + + +def test_max_padded_rows_uses_exported_row_bound(monkeypatch): + monkeypatch.setenv(MAX_ROWS_ENV, "3") + args = Namespace(global_batch_size=8, rollout_batch_size=2, n_samples_per_prompt=2) + assert HELPERS._agentic_variable_row_max_padded_rows(args) == 16 + + monkeypatch.delenv(MAX_ROWS_ENV) + with pytest.raises(ValueError, match=MAX_ROWS_ENV): + HELPERS._agentic_variable_row_max_padded_rows(args) + monkeypatch.setenv(MAX_ROWS_ENV, "0") + with pytest.raises(ValueError, match="positive integer"): + HELPERS._agentic_variable_row_max_padded_rows(args) + + +def test_padding_meta_parser_accepts_only_explicit_booleans(): + assert HELPERS._extract_agentic_variable_row_padding_flags(_batch_meta([False, True, False]), 3) == [ + False, + True, + False, + ] + + with pytest.raises(RuntimeError, match="metadata size mismatch"): + HELPERS._extract_agentic_variable_row_padding_flags(_batch_meta([False]), 2) + with pytest.raises(RuntimeError, match="missing"): + HELPERS._extract_agentic_variable_row_padding_flags(_custom_meta_batch([{}]), 1) + with pytest.raises(RuntimeError, match="must be bool"): + HELPERS._extract_agentic_variable_row_padding_flags(_batch_meta([1]), 1) + with pytest.raises(RuntimeError, match="must be a mapping"): + HELPERS._extract_agentic_variable_row_padding_flags(_custom_meta_batch([None]), 1) + with pytest.raises(RuntimeError, match="get_all_custom_meta"): + HELPERS._extract_agentic_variable_row_padding_flags(SimpleNamespace(samples=[]), 1) + + +def test_identity_parser_binds_custom_metadata_to_reordered_tensor_rows(): + identities = [ + _real_identity(group_id="0", turn_index=0, turn_count=1), + _padding_identity(0), + ] + rollout_data, batch_meta = _identity_batch(identities) + + assert HELPERS._extract_agentic_variable_row_identities(batch_meta, rollout_data) == identities + + rollout_data[IDENTITY_TAGS_FIELD] = list(reversed(rollout_data[IDENTITY_TAGS_FIELD])) + with pytest.raises(RuntimeError, match="identity mismatch"): + HELPERS._extract_agentic_variable_row_identities(batch_meta, rollout_data) + + +def test_partition_identity_validation_is_order_independent_and_rejects_restart_residue(): + identities = [ + *[_real_identity(group_id="0", turn_index=idx, turn_count=2) for idx in range(2)], + *[_real_identity(group_id="1", turn_index=idx, turn_count=3) for idx in range(3)], + _padding_identity(0), + ] + identities = [identities[idx] for idx in (3, 0, 5, 4, 2, 1)] + summary = HELPERS._validate_agentic_variable_row_partition_identities( + identities, + expected_group_count=2, + expected_trajectories_per_group=1, + ) + assert summary == { + "policy_version": "7", + "group_count": 2, + "real_row_count": 5, + "padding_row_count": 1, + } + + duplicate_after_restart = [*identities, identities[0]] + with pytest.raises(RuntimeError, match="duplicate row_id"): + HELPERS._validate_agentic_variable_row_partition_identities( + duplicate_after_restart, + expected_group_count=2, + expected_trajectories_per_group=1, + ) + + residual_group = [identity for identity in identities if identity.get("rollout_group_id") != "1"] + with pytest.raises(RuntimeError, match="residual or missing"): + HELPERS._validate_agentic_variable_row_partition_identities( + residual_group, + expected_group_count=2, + expected_trajectories_per_group=1, + ) + + incomplete_turn = [identity for identity in identities if identity.get("row_id") != "row-1-1"] + with pytest.raises(RuntimeError, match="missing or has duplicate physical rows"): + HELPERS._validate_agentic_variable_row_partition_identities( + incomplete_turn, + expected_group_count=2, + expected_trajectories_per_group=1, + ) + + +def test_window_progress_accepts_full_and_partial_then_rejects_overflow(): + consumed = HELPERS._validate_agentic_variable_row_window( + actual_global_rows=8, + global_batch_size=8, + total_padded_rows_before=0, + max_padded_rows=16, + ) + assert consumed == 8 + consumed = HELPERS._validate_agentic_variable_row_window( + actual_global_rows=5, + global_batch_size=8, + total_padded_rows_before=consumed, + max_padded_rows=16, + ) + assert consumed == 16 + + with pytest.raises(RuntimeError, match="between 1 and global_batch_size"): + HELPERS._validate_agentic_variable_row_window( + actual_global_rows=0, + global_batch_size=8, + total_padded_rows_before=0, + max_padded_rows=16, + ) + with pytest.raises(RuntimeError, match="exceeded"): + HELPERS._validate_agentic_variable_row_window( + actual_global_rows=1, + global_batch_size=8, + total_padded_rows_before=16, + max_padded_rows=16, + ) + + +def test_partition_padding_is_order_independent_but_still_bounded(): + assert ( + HELPERS._validate_agentic_variable_row_partition( + total_physical_rows=16, + total_actual_rows=13, + global_batch_size=8, + ) + == 3 + ) + + with pytest.raises(RuntimeError, match="at least one global batch of padding"): + HELPERS._validate_agentic_variable_row_partition( + total_physical_rows=16, + total_actual_rows=8, + global_batch_size=8, + ) + + +def test_collector_accepts_real_rows_after_a_window_containing_padding(monkeypatch): + monkeypatch.setenv(MAX_ROWS_ENV, "50") + batches = [ + _identity_batch( + [ + *[_real_identity(group_id="0", turn_index=idx, turn_count=5) for idx in range(5)], + *[_padding_identity(idx) for idx in range(3)], + ] + ), + _identity_batch([_real_identity(group_id="1", turn_index=idx, turn_count=8) for idx in range(8)]), + ] + + collector = VARIABLE_ROW_COLLECTOR() + collector.role = "actor" + collector.args = Namespace( + advantage_estimator="grpo", + distributed_timeout_minutes=1, + global_batch_size=8, + multimodal_keys=None, + n_samples_per_prompt=1, + rollout_batch_size=2, + use_critic=False, + use_opd=False, + ) + + def _get_data(*_args, **_kwargs): + if batches: + return batches.pop(0) + return None, None + + collector._get_data_from_transfer_queue = _get_data + collector.all_consumed = lambda *_args, **_kwargs: not batches + + result = collector._collect_agentic_variable_row_batches( + rollout_id=0, + task_name="actor_train", + data_fields=["total_lengths"], + ) + + rollout_batches, _metas, local_counts, global_counts, valid_flags = result + assert len(rollout_batches) == 2 + assert local_counts == [8, 8] + assert global_counts == [5, 8] + assert sum(valid_flags) == 13 + assert valid_flags == [True, True, True, True, True, False, False, False, *([True] * 8)] + + +def test_drain_state_machine_aligns_dp_readiness_and_requires_a_window(): + action = HELPERS._agentic_variable_row_drain_action + assert action(ready_min=1, ready_max=1, stream_drained=False, accepted_windows=0) == "accept" + assert action(ready_min=0, ready_max=1, stream_drained=False, accepted_windows=0) == "wait" + assert action(ready_min=0, ready_max=0, stream_drained=False, accepted_windows=0) == "retry" + assert action(ready_min=0, ready_max=0, stream_drained=True, accepted_windows=1) == "done" + with pytest.raises(RuntimeError, match="before yielding"): + action(ready_min=0, ready_max=0, stream_drained=True, accepted_windows=0) + with pytest.raises(RuntimeError, match="pending after"): + action(ready_min=1, ready_max=1, stream_drained=True, accepted_windows=1) + + +def test_stream_drained_requires_data_parallel_consensus(): + class _Tensor: + def __init__(self, value): + self.value = value + + def item(self): + return self.value + + class _Torch: + int64 = object() + + @staticmethod + def tensor(values, *, dtype, device): + assert dtype is _Torch.int64 + assert device == "cuda:0" + return _Tensor(values[0]) + + class _Dist: + class ReduceOp: + MIN = object() + + observed_group = None + + @classmethod + def all_reduce(cls, tensor, *, op, group): + assert op is cls.ReduceOp.MIN + cls.observed_group = group + # Simulate another DP rank that has not observed drained yet. + tensor.value = 0 + + consensus = HELPERS._agentic_variable_row_stream_drained_consensus + consensus.__globals__["torch"] = _Torch + consensus.__globals__["dist"] = _Dist + + assert ( + consensus( + stream_drained=True, + device="cuda:0", + dp_group="dp-group", + ) + is False + ) + assert _Dist.observed_group == "dp-group" diff --git a/tests/backends/megatron/test_variable_row_loss_scale.py b/tests/backends/megatron/test_variable_row_loss_scale.py new file mode 100644 index 000000000..9f9c81d70 --- /dev/null +++ b/tests/backends/megatron/test_variable_row_loss_scale.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import ast +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[3] +GLOBAL_COUNTS_KEY = "rollout_mini_global_sample_counts" +VALID_ROWS_KEY = "agentic_variable_row_valid_rows" + + +def _load_data_definitions(): + path = REPO_ROOT / "relax" / "backends" / "megatron" / "data.py" + names = { + "DataIterator", + "_build_rollout_mini_loss_scales", + "get_rollout_valid_row_flags", + "_rollout_metrics_view", + } + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + definitions = [ + node for node in tree.body if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name in names + ] + assert {node.name for node in definitions} == names + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + *definitions, + ], + type_ignores=[], + ) + + class _Torch: + class Tensor: + pass + + namespace = { + "ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY": GLOBAL_COUNTS_KEY, + "ROLLOUT_VALID_ROW_FLAGS_KEY": VALID_ROWS_KEY, + "RolloutBatch": dict, + "torch": _Torch, + } + exec(compile(ast.fix_missing_locations(module), str(path), "exec"), namespace) + return SimpleNamespace(**{name: namespace[name] for name in names}) + + +DATA = _load_data_definitions() + + +def test_variable_window_loss_scale_matches_full_gbs_and_actual_tail(): + scales = DATA._build_rollout_mini_loss_scales( + num_microbatches=[2, 2], + actual_global_sample_counts=[8, 5], + global_batch_size=8, + dp_size=2, + ) + assert scales == [0.5, 0.5, 0.8, 0.8] + assert ( + DATA._build_rollout_mini_loss_scales( + num_microbatches=[2], + actual_global_sample_counts=None, + global_batch_size=8, + dp_size=2, + ) + is None + ) + + +def test_data_iterator_injects_scale_per_microbatch_and_reset_replays_schedule(): + rollout_data = { + "total_lengths": [4] * 8, + "loss_masks": [[1], [1], [1], [1], [1], [1], [1], [0]], + } + iterator = DATA.DataIterator( + rollout_data, + micro_batch_size=2, + micro_batch_loss_scales=[0.5, 0.5, 0.8, 0.8], + ) + + batches = [iterator.get_next(["loss_masks"]) for _ in range(4)] + assert [batch["__loss_scale__"] for batch in batches] == [0.5, 0.5, 0.8, 0.8] + assert batches[-1]["loss_masks"] == [[1], [0]] + assert batches[-1]["loss_masks"][-1][0] * batches[-1]["__loss_scale__"] == 0 + + iterator.reset() + assert iterator.get_next(["loss_masks"])["__loss_scale__"] == 0.5 + + +def test_variable_row_validity_is_persisted_sliced_and_removed_from_metric_view(): + rollout_data = { + "total_lengths": [7, 8, 3, 3], + "response_lengths": [3, 4, 1, 1], + "raw_reward": [1.0, 0.0, 0.0, 0.0], + VALID_ROWS_KEY: [True, True, False, False], + } + assert DATA.get_rollout_valid_row_flags(rollout_data) == [True, True, False, False] + + iterator = DATA.DataIterator(rollout_data, micro_batch_size=2) + assert iterator.get_next([VALID_ROWS_KEY])[VALID_ROWS_KEY] == [True, True] + assert iterator.get_next([VALID_ROWS_KEY])[VALID_ROWS_KEY] == [False, False] + + metric_view, metric_rows = DATA._rollout_metrics_view(rollout_data) + assert metric_rows == 2 + assert metric_view["total_lengths"] == [7, 8] + assert metric_view["raw_reward"] == [1.0, 0.0] + assert VALID_ROWS_KEY not in metric_view + + +def test_variable_row_validity_is_strict_and_legacy_view_is_unchanged(): + legacy = {"total_lengths": [7], "response_lengths": [3]} + metric_view, metric_rows = DATA._rollout_metrics_view(legacy) + assert metric_view is legacy + assert metric_rows is None + + with pytest.raises(RuntimeError, match="one flag per row"): + DATA.get_rollout_valid_row_flags({"total_lengths": [7, 8], VALID_ROWS_KEY: [True]}) + with pytest.raises(RuntimeError, match="only bool"): + DATA.get_rollout_valid_row_flags({"total_lengths": [7], VALID_ROWS_KEY: [1]}) + + +def test_vpp_iterators_have_independent_scale_offsets(): + rollout_data = {"total_lengths": [4] * 4, "tokens": [[1]] * 4} + first = DATA.DataIterator( + rollout_data, + micro_batch_size=2, + micro_batch_loss_scales=[0.5, 0.8], + ) + second = DATA.DataIterator( + rollout_data, + micro_batch_size=2, + micro_batch_loss_scales=[0.5, 0.8], + ) + assert first.get_next(["tokens"])["__loss_scale__"] == 0.5 + assert first.get_next(["tokens"])["__loss_scale__"] == 0.8 + assert second.get_next(["tokens"])["__loss_scale__"] == 0.5 + + +def test_normal_iterator_does_not_inject_explicit_loss_scale(): + iterator = DATA.DataIterator( + {"total_lengths": [4, 4], "tokens": [[1], [2]]}, + micro_batch_size=1, + ) + assert "__loss_scale__" not in iterator.get_next(["tokens"]) + + +@pytest.mark.parametrize( + ("num_microbatches", "actual_counts"), + [ + ([2], [0]), + ([2], [9]), + ([2], [True]), + ([2, 2], [8]), + ], +) +def test_variable_window_loss_scale_rejects_invalid_counts(num_microbatches, actual_counts): + with pytest.raises(ValueError): + DATA._build_rollout_mini_loss_scales( + num_microbatches=num_microbatches, + actual_global_sample_counts=actual_counts, + global_batch_size=8, + dp_size=2, + ) diff --git a/tests/backends/megatron/test_variable_row_metric_count.py b/tests/backends/megatron/test_variable_row_metric_count.py new file mode 100644 index 000000000..fcaff422e --- /dev/null +++ b/tests/backends/megatron/test_variable_row_metric_count.py @@ -0,0 +1,44 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import ast +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[3] +VALID_ROWS_KEY = "agentic_variable_row_valid_rows" + + +def _load_metric_count_helper(): + path = REPO_ROOT / "relax" / "backends" / "megatron" / "loss.py" + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + definition = next( + node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "_metric_num_samples" + ) + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + definition, + ], + type_ignores=[], + ) + + def get_rollout_valid_row_flags(batch): + return batch.get(VALID_ROWS_KEY) + + namespace = {"RolloutBatch": dict, "get_rollout_valid_row_flags": get_rollout_valid_row_flags} + exec(compile(ast.fix_missing_locations(module), str(path), "exec"), namespace) + return namespace["_metric_num_samples"] + + +METRIC_NUM_SAMPLES = _load_metric_count_helper() + + +def test_train_metrics_exclude_synthetic_rows_and_legacy_count_is_unchanged(): + variable_rows = { + "response_lengths": [3, 4, 1, 1], + VALID_ROWS_KEY: [True, True, False, False], + } + assert METRIC_NUM_SAMPLES(variable_rows) == 2 + assert METRIC_NUM_SAMPLES({"response_lengths": [3, 4, 1, 1]}) == 4 diff --git a/tests/backends/megatron/test_variable_row_plan.py b/tests/backends/megatron/test_variable_row_plan.py new file mode 100644 index 000000000..5d1e7dd38 --- /dev/null +++ b/tests/backends/megatron/test_variable_row_plan.py @@ -0,0 +1,260 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import importlib +import sys +import types +from argparse import Namespace + +import pytest + + +def _load_data_module(monkeypatch): + megatron = types.ModuleType("megatron") + core = types.ModuleType("megatron.core") + mpu = types.ModuleType("megatron.core.mpu") + packed_seq_params = types.ModuleType("megatron.core.packed_seq_params") + training = types.ModuleType("megatron.training") + global_vars = types.ModuleType("megatron.training.global_vars") + tracking_utils = types.ModuleType("relax.utils.tracking_utils") + + class _PackedSeqParams: + pass + + core.mpu = mpu + packed_seq_params.PackedSeqParams = _PackedSeqParams + global_vars.get_args = lambda: None + + modules = { + "megatron": megatron, + "megatron.core": core, + "megatron.core.mpu": mpu, + "megatron.core.packed_seq_params": packed_seq_params, + "megatron.training": training, + "megatron.training.global_vars": global_vars, + "relax.utils.tracking_utils": tracking_utils, + } + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + + sys.modules.pop("relax.backends.megatron.data", None) + return importlib.import_module("relax.backends.megatron.data") + + +def _assert_window_invariants(window, dp_size, micro_batch_size): + assert window.stop_row - window.start_row == window.actual_global_rows + assert sum(window.actual_rows_per_dp) == window.actual_global_rows + assert sum(window.padded_rows_per_dp) == window.padded_global_rows + assert sum(window.padding_rows_per_dp) == window.padded_global_rows - window.actual_global_rows + assert len(window.actual_rows_per_dp) == dp_size + assert len(window.padded_rows_per_dp) == dp_size + assert len(window.padding_rows_per_dp) == dp_size + assert all(rows % micro_batch_size == 0 for rows in window.padded_rows_per_dp) + assert all( + actual + padding == padded + for actual, padding, padded in zip( + window.actual_rows_per_dp, + window.padding_rows_per_dp, + window.padded_rows_per_dp, + strict=True, + ) + ) + + +def test_variable_row_plan_exact_batches_have_no_padding(monkeypatch): + data_module = _load_data_module(monkeypatch) + + plan = data_module.build_variable_row_minibatch_plan( + actual_global_rows=512, + global_batch_size=256, + dp_size=8, + micro_batch_size=32, + ) + + assert [(window.start_row, window.stop_row) for window in plan.windows] == [(0, 256), (256, 512)] + assert [window.actual_global_rows for window in plan.windows] == [256, 256] + assert plan.total_padded_global_rows == 512 + assert plan.total_padding_rows == 0 + for window in plan.windows: + assert window.actual_rows_per_dp == (32,) * 8 + assert window.padding_rows_per_dp == (0,) * 8 + _assert_window_invariants(window, dp_size=8, micro_batch_size=32) + + +def test_variable_row_plan_pads_only_final_partial_batch(monkeypatch): + data_module = _load_data_module(monkeypatch) + + plan = data_module.build_variable_row_minibatch_plan( + actual_global_rows=513, + global_batch_size=256, + dp_size=8, + micro_batch_size=32, + ) + + assert [(window.start_row, window.stop_row) for window in plan.windows] == [ + (0, 256), + (256, 512), + (512, 513), + ] + assert [window.actual_global_rows for window in plan.windows] == [256, 256, 1] + assert [window.padded_global_rows for window in plan.windows] == [256, 256, 256] + assert plan.total_padded_global_rows == 768 + assert plan.total_padding_rows == 255 + + final = plan.windows[-1] + assert final.actual_rows_per_dp == (1, 0, 0, 0, 0, 0, 0, 0) + assert final.padded_rows_per_dp == (32,) * 8 + assert final.padding_rows_per_dp == (31, 32, 32, 32, 32, 32, 32, 32) + for window in plan.windows: + _assert_window_invariants(window, dp_size=8, micro_batch_size=32) + + +def test_variable_row_plan_balances_real_rows_across_dp_ranks(monkeypatch): + data_module = _load_data_module(monkeypatch) + + plan = data_module.build_variable_row_minibatch_plan( + actual_global_rows=5, + global_batch_size=8, + dp_size=2, + micro_batch_size=2, + ) + + window = plan.windows[0] + assert window.actual_rows_per_dp == (3, 2) + assert window.padded_rows_per_dp == (4, 4) + assert window.padding_rows_per_dp == (1, 2) + assert max(window.actual_rows_per_dp) - min(window.actual_rows_per_dp) == 1 + assert plan.total_padding_rows == 3 + _assert_window_invariants(window, dp_size=2, micro_batch_size=2) + + +def test_variable_row_plan_supports_fewer_real_rows_than_dp_ranks(monkeypatch): + data_module = _load_data_module(monkeypatch) + + plan = data_module.build_variable_row_minibatch_plan( + actual_global_rows=3, + global_batch_size=16, + dp_size=4, + micro_batch_size=2, + ) + + window = plan.windows[0] + assert window.actual_rows_per_dp == (1, 1, 1, 0) + assert window.padded_rows_per_dp == (4, 4, 4, 4) + assert window.padding_rows_per_dp == (3, 3, 3, 4) + assert plan.total_padding_rows == 13 + _assert_window_invariants(window, dp_size=4, micro_batch_size=2) + + +def test_variable_row_plan_grid_preserves_all_rows_and_contiguous_boundaries(monkeypatch): + data_module = _load_data_module(monkeypatch) + + for dp_size in (1, 2, 4): + for micro_batch_size in (1, 2, 4): + for local_microbatches in (1, 2, 3): + global_batch_size = dp_size * micro_batch_size * local_microbatches + for actual_global_rows in range(1, 2 * global_batch_size + 3): + plan = data_module.build_variable_row_minibatch_plan( + actual_global_rows=actual_global_rows, + global_batch_size=global_batch_size, + dp_size=dp_size, + micro_batch_size=micro_batch_size, + ) + + assert plan.windows[0].start_row == 0 + assert plan.windows[-1].stop_row == actual_global_rows + assert all( + left.stop_row == right.start_row + for left, right in zip(plan.windows, plan.windows[1:], strict=False) + ) + assert sum(window.actual_global_rows for window in plan.windows) == actual_global_rows + assert sum(window.padded_global_rows for window in plan.windows) == plan.total_padded_global_rows + assert plan.total_padding_rows == plan.total_padded_global_rows - actual_global_rows + assert all(window.actual_global_rows == global_batch_size for window in plan.windows[:-1]) + for window in plan.windows: + _assert_window_invariants(window, dp_size, micro_batch_size) + + +@pytest.mark.parametrize( + ("name", "kwargs"), + [ + ( + "actual_global_rows", + {"actual_global_rows": 0, "global_batch_size": 8, "dp_size": 2, "micro_batch_size": 2}, + ), + ( + "global_batch_size", + {"actual_global_rows": 1, "global_batch_size": 0, "dp_size": 2, "micro_batch_size": 2}, + ), + ( + "dp_size", + {"actual_global_rows": 1, "global_batch_size": 8, "dp_size": 0, "micro_batch_size": 2}, + ), + ( + "micro_batch_size", + {"actual_global_rows": 1, "global_batch_size": 8, "dp_size": 2, "micro_batch_size": -1}, + ), + ], +) +def test_variable_row_plan_rejects_non_positive_inputs(monkeypatch, name, kwargs): + data_module = _load_data_module(monkeypatch) + + with pytest.raises(ValueError, match=rf"{name} must be positive"): + data_module.build_variable_row_minibatch_plan(**kwargs) + + +@pytest.mark.parametrize( + ("name", "value"), + [ + ("actual_global_rows", True), + ("actual_global_rows", 1.0), + ("global_batch_size", "8"), + ("dp_size", False), + ("micro_batch_size", 2.5), + ], +) +def test_variable_row_plan_rejects_non_integer_inputs(monkeypatch, name, value): + data_module = _load_data_module(monkeypatch) + kwargs = { + "actual_global_rows": 1, + "global_batch_size": 8, + "dp_size": 2, + "micro_batch_size": 2, + } + kwargs[name] = value + + with pytest.raises(TypeError, match=rf"{name} must be an int"): + data_module.build_variable_row_minibatch_plan(**kwargs) + + +def test_variable_row_plan_rejects_non_divisible_fixed_global_batch(monkeypatch): + data_module = _load_data_module(monkeypatch) + + with pytest.raises(ValueError, match="global_batch_size must be divisible by dp_size \\* micro_batch_size"): + data_module.build_variable_row_minibatch_plan( + actual_global_rows=7, + global_batch_size=12, + dp_size=2, + micro_batch_size=4, + ) + + +def test_variable_row_helper_does_not_change_fixed_n_plan(monkeypatch): + data_module = _load_data_module(monkeypatch) + + fixed_plan = data_module.build_rollout_minibatch_plan( + Namespace( + rollout_batch_size=8, + n_samples_per_prompt=8, + global_batch_size=32, + num_steps_per_rollout=None, + ), + dp_size=2, + ) + + assert fixed_plan == data_module.RolloutMiniBatchPlan( + num_rollout_minis=2, + mini_rollout_batch_size=4, + fixed_n_samples_per_prompt=8, + mini_global_samples=32, + mini_local_sample_request=16, + ) diff --git a/tests/backends/megatron/test_variable_row_scheduler_accounting.py b/tests/backends/megatron/test_variable_row_scheduler_accounting.py new file mode 100644 index 000000000..d01a8db0e --- /dev/null +++ b/tests/backends/megatron/test_variable_row_scheduler_accounting.py @@ -0,0 +1,191 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import ast +import os +from argparse import Namespace +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[3] +MAX_ROWS_ENV = "RELAX_AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE" + + +def _load_model_helpers(): + path = REPO_ROOT / "relax" / "backends" / "megatron" / "model.py" + names = { + "_normalize_scheduler_increments", + "_initial_variable_row_tracking_step", + "_advance_variable_row_tracking_step", + "_restore_agentic_variable_row_resume_cursor", + "_optimizer_scheduler_training_plan", + "get_optimizer_param_scheduler", + } + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + definitions = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in names] + assert {node.name for node in definitions} == names + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + *definitions, + ], + type_ignores=[], + ) + + class _Scheduler: + def __init__(self, optimizer, **kwargs): + self.optimizer = optimizer + self.kwargs = kwargs + + namespace = { + "AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV": MAX_ROWS_ENV, + "Namespace": Namespace, + "OptimizerParamScheduler": _Scheduler, + "Sequence": list, + "os": os, + } + exec(compile(ast.fix_missing_locations(module), str(path), "exec"), namespace) + return SimpleNamespace(**{name: namespace[name] for name in names}) + + +HELPERS = _load_model_helpers() + + +def test_variable_row_scheduler_uses_each_window_actual_rows(): + assert HELPERS._normalize_scheduler_increments([8, 5, 3], 3) == (8, 5, 3) + assert HELPERS._normalize_scheduler_increments(None, 3) is None + + with pytest.raises(ValueError, match="length must match"): + HELPERS._normalize_scheduler_increments([8, 5], 3) + with pytest.raises(ValueError, match="positive integers"): + HELPERS._normalize_scheduler_increments([8, 0, 3], 3) + + +def test_variable_row_tracking_cursor_is_monotonic_across_different_window_counts(): + args = Namespace() + scheduler = SimpleNamespace(num_steps=100) + cursor = HELPERS._initial_variable_row_tracking_step(args, scheduler) + observed_steps = [] + + # Rollout A has three windows, rollout B only one. The old + # rollout_id * current_window_count formula would collide at step 1. + for increment in (8, 5, 3, 7): + observed_steps.append(cursor) + cursor = HELPERS._advance_variable_row_tracking_step(args, cursor, increment) + + assert observed_steps == [100, 108, 113, 116] + assert len(observed_steps) == len(set(observed_steps)) + assert observed_steps == sorted(observed_steps) + assert args._agentic_variable_row_tracking_samples == 123 + + # Re-entry uses the actor-persisted cursor instead of the scheduler fallback. + assert HELPERS._initial_variable_row_tracking_step(args, SimpleNamespace(num_steps=0)) == 123 + + +def test_variable_row_checkpoint_resume_restores_scheduler_cursor(monkeypatch): + monkeypatch.setenv(MAX_ROWS_ENV, "5") + args = _scheduler_args() + scheduler = SimpleNamespace(num_steps=321) + + assert not HELPERS._restore_agentic_variable_row_resume_cursor(args, scheduler, iteration=0) + assert HELPERS._restore_agentic_variable_row_resume_cursor(args, scheduler, iteration=3) + assert args._agentic_variable_row_tracking_samples == 321 + + +def test_variable_row_checkpoint_resume_fails_closed_on_missing_or_conflicting_cursor(monkeypatch): + monkeypatch.setenv(MAX_ROWS_ENV, "5") + + with pytest.raises(RuntimeError, match="positive checkpoint scheduler num_steps"): + HELPERS._restore_agentic_variable_row_resume_cursor( + _scheduler_args(), SimpleNamespace(num_steps=0), iteration=3 + ) + with pytest.raises(RuntimeError, match="requires a restored optimizer scheduler"): + HELPERS._restore_agentic_variable_row_resume_cursor(_scheduler_args(), None, iteration=3) + + args = _scheduler_args() + args._agentic_variable_row_tracking_samples = 320 + with pytest.raises(RuntimeError, match="cursor conflicts"): + HELPERS._restore_agentic_variable_row_resume_cursor(args, SimpleNamespace(num_steps=321), iteration=3) + + +def test_checkpoint_resume_preserves_legacy_paths(monkeypatch): + args = _scheduler_args() + scheduler = SimpleNamespace(num_steps=321) + + monkeypatch.delenv(MAX_ROWS_ENV, raising=False) + assert not HELPERS._restore_agentic_variable_row_resume_cursor(args, scheduler, iteration=3) + + monkeypatch.setenv(MAX_ROWS_ENV, "5") + assert not HELPERS._restore_agentic_variable_row_resume_cursor( + _scheduler_args(group_rm=False), scheduler, iteration=3 + ) + + +def _scheduler_args(**overrides): + values = { + "num_rollout": 3, + "rollout_batch_size": 2, + "n_samples_per_prompt": 4, + "global_batch_size": 8, + "group_rm": True, + "agentic_custom_advantage_path": "examples.graphgpo.custom_advantage.compute_custom_advantage", + "use_dynamic_batch_size": True, + "fully_async": False, + "hybrid": False, + "lr_decay_iters": None, + "lr_wsd_decay_iters": None, + "lr_warmup_fraction": 0.1, + "lr_warmup_iters": 0, + "lr_warmup_init": 0.0, + "lr": 1e-6, + "min_lr": 0.0, + "lr_decay_style": "constant", + "start_weight_decay": 0.1, + "end_weight_decay": 0.1, + "weight_decay_incr_style": "constant", + "use_checkpoint_opt_param_scheduler": False, + "override_opt_param_scheduler": False, + "lr_wsd_decay_style": "constant", + } + values.update(overrides) + return Namespace(**values) + + +def test_variable_row_scheduler_horizon_uses_max_real_turn_rows(monkeypatch): + monkeypatch.setenv(MAX_ROWS_ENV, "5") + args = _scheduler_args() + + # Per rollout: 2 * 4 trajectories * 5 rows = 40 rows = 5 windows. + assert HELPERS._optimizer_scheduler_training_plan(args) == (15, 120) + scheduler = HELPERS.get_optimizer_param_scheduler(args, optimizer=object()) + + assert args.train_iters == 15 + assert args.lr_decay_iters == 15 + assert scheduler.kwargs["lr_decay_steps"] == 120 + assert scheduler.kwargs["wd_incr_steps"] == 120 + assert scheduler.kwargs["lr_warmup_steps"] == 12 + + +def test_scheduler_plan_preserves_legacy_trajectory_units(monkeypatch): + monkeypatch.delenv(MAX_ROWS_ENV, raising=False) + args = _scheduler_args(group_rm=False) + + assert HELPERS._optimizer_scheduler_training_plan(args) == (3, 24) + scheduler = HELPERS.get_optimizer_param_scheduler(args, optimizer=object()) + assert args.train_iters == 3 + assert scheduler.kwargs["lr_decay_steps"] == 24 + assert scheduler.kwargs["wd_incr_steps"] == 24 + + +def test_scheduler_plan_rejects_invalid_row_bound_and_hybrid(monkeypatch): + monkeypatch.setenv(MAX_ROWS_ENV, "0") + with pytest.raises(ValueError, match="positive integer"): + HELPERS._optimizer_scheduler_training_plan(_scheduler_args()) + + monkeypatch.setenv(MAX_ROWS_ENV, "5") + with pytest.raises(ValueError, match="synchronous training only"): + HELPERS._optimizer_scheduler_training_plan(_scheduler_args(fully_async=True, hybrid=True)) diff --git a/tests/core/test_controller_variable_rows.py b/tests/core/test_controller_variable_rows.py new file mode 100644 index 000000000..34875720d --- /dev/null +++ b/tests/core/test_controller_variable_rows.py @@ -0,0 +1,119 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from argparse import Namespace + +import pytest + + +try: + import relax.core.controller as controller_module +except (ImportError, AssertionError) as _exc: + pytest.skip(f"relax.core.controller requires the Relax runtime image: {_exc}", allow_module_level=True) + + +class _Sampler: + def __init__(self, **kwargs): + self.kwargs = kwargs + + +def _config(**overrides): + values = { + "loss_type": "grpo", + "rollout_batch_size": 2, + "over_sampling_batch_size": 2, + "n_samples_per_prompt": 3, + "partial_rollout": False, + "use_dynamic_global_batch_size": False, + "max_staleness": 1, + "global_batch_size": 8, + "fully_async": False, + "hybrid": False, + "use_dynamic_batch_size": True, + "group_rm": True, + "agentic_custom_advantage_path": "examples.graphgpo.advantage.compute", + "balance_data": False, + "polling_mode": "default", + "num_data_storage_units": 1, + } + values.update(overrides) + return Namespace(**values) + + +def _initialize(monkeypatch, config): + captured = {} + + def _init(*, conf): + captured["conf"] = conf + return None + + monkeypatch.setattr(controller_module, "resolve_sft_algo_key", lambda _config: "grpo") + monkeypatch.setattr(controller_module, "compute_dp_size", lambda _config: 2) + monkeypatch.setattr(controller_module.tq, "init", _init) + monkeypatch.setattr(controller_module, "SeqlenBalancedSampler", _Sampler) + monkeypatch.setattr(controller_module, "GRPOGroupNSampler", _Sampler) + controller = object.__new__(controller_module.Controller) + controller.config = config + controller._initialize_data_system() + return captured["conf"] + + +def test_variable_row_controller_uses_row_sampler_and_padded_capacity(monkeypatch): + monkeypatch.setenv(controller_module.AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV, "5") + + conf = _initialize(monkeypatch, _config()) + + sampler = conf.controller.sampler + assert isinstance(sampler, _Sampler) + assert sampler.kwargs == {"n_samples_per_prompt": 1, "dp_size": 2} + # ceil((2 prompt groups * 3 trajectories * 5 rows) / GBS 8) * 8 * (1 + staleness 1) + assert conf.backend.SimpleStorage.total_storage_size == 64 + + +def test_variable_row_controller_rejects_hybrid_mode(monkeypatch): + monkeypatch.setenv(controller_module.AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV, "1") + + with pytest.raises(ValueError, match="synchronous training only"): + _initialize(monkeypatch, _config(fully_async=True, hybrid=True)) + + +def test_variable_row_controller_rejects_pure_fully_async_mode(monkeypatch): + monkeypatch.setenv(controller_module.AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV, "1") + + with pytest.raises(ValueError, match="synchronous training only"): + _initialize(monkeypatch, _config(fully_async=True, hybrid=False)) + + +@pytest.mark.parametrize("raw_value", ["0", "-1", "+1", "1_0", "1.5", "many"]) +def test_variable_row_controller_rejects_invalid_capacity_env(monkeypatch, raw_value): + monkeypatch.setenv(controller_module.AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV, raw_value) + + with pytest.raises(ValueError, match=controller_module.AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV): + _initialize(monkeypatch, _config()) + + +@pytest.mark.parametrize("raw_value", [None, ""]) +def test_variable_row_controller_is_explicitly_opt_in(monkeypatch, raw_value): + if raw_value is None: + monkeypatch.delenv(controller_module.AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV, raising=False) + else: + monkeypatch.setenv(controller_module.AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV, raw_value) + + conf = _initialize(monkeypatch, _config()) + + assert conf.controller.sampler.kwargs == {"n_samples_per_prompt": 3} + assert conf.backend.SimpleStorage.total_storage_size == 12 + + +def test_default_controller_path_preserves_grouped_sampler_and_capacity(monkeypatch): + monkeypatch.delenv(controller_module.AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE_ENV, raising=False) + + conf = _initialize( + monkeypatch, + _config( + group_rm=False, + agentic_custom_advantage_path=None, + ), + ) + + assert conf.controller.sampler.kwargs == {"n_samples_per_prompt": 3} + assert conf.backend.SimpleStorage.total_storage_size == 12 diff --git a/tests/graphgpo/test_action_parser.py b/tests/graphgpo/test_action_parser.py new file mode 100644 index 000000000..c3f99d158 --- /dev/null +++ b/tests/graphgpo/test_action_parser.py @@ -0,0 +1,55 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import unittest + +from examples.graphgpo.action_parser import parse_action + + +class ActionParserTest(unittest.TestCase): + def test_first_lowercase_action_tag_wins(self) -> None: + result = parse_action("choosego northopen door") + + self.assertEqual(result.action, "go north") + self.assertTrue(result.is_valid) + self.assertEqual(result.source, "tag") + + def test_response_is_lowercased_before_action_extraction(self) -> None: + response = "ChooseGo North" + result = parse_action(response) + + self.assertEqual(result.action, "go north") + self.assertTrue(result.is_valid) + self.assertEqual(result.source, "tag") + + def test_missing_tag_fallback_is_always_invalid(self) -> None: + result = parse_action("x" + "X" * 40) + + self.assertEqual(result.action, "x" * 30) + self.assertFalse(result.is_valid) + + def test_missing_think_or_chinese_response_is_invalid(self) -> None: + self.assertFalse(parse_action("go north").is_valid) + self.assertFalse(parse_action("我来选择go north").is_valid) + + def test_multiline_action_is_supported(self) -> None: + result = parse_action("choose\nput apple in fridge\n") + + self.assertEqual(result.action, "put apple in fridge") + self.assertTrue(result.is_valid) + + def test_closing_tag_before_opening_tag_is_not_a_match(self) -> None: + response = "chooseprefixgo north" + result = parse_action(response) + + self.assertEqual(result.action, "go north") + self.assertTrue(result.is_valid) + + def test_empty_action_matches_reference_validity(self) -> None: + result = parse_action("choose ") + + self.assertEqual(result.action, "") + self.assertTrue(result.is_valid) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/graphgpo/test_alfworld_env.py b/tests/graphgpo/test_alfworld_env.py new file mode 100644 index 000000000..30b8a983b --- /dev/null +++ b/tests/graphgpo/test_alfworld_env.py @@ -0,0 +1,229 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from examples.graphgpo.alfworld_env import ( + AlfWorldTextEnv, + _canonicalize_game_files, + _select_game_file, +) + + +def _batched_info( + commands: list[str], + *, + won: bool = False, + gamefile: str = "games/task-1/game.tw-pddl", +) -> dict[str, list[object]]: + return { + "admissible_commands": [commands], + "won": [won], + "extra.gamefile": [gamefile], + } + + +class FakeRawAlfWorldEnv: + def __init__(self, transitions: list[tuple[str, list[str], bool, bool]]) -> None: + self.transitions = list(transitions) + self.actions: list[list[str]] = [] + self.closed = False + self.seed_values: list[int] = [] + + def seed(self, seed: int) -> None: + self.seed_values.append(seed) + + def reset(self): + return ( + ["Room intro. Your task is to: put the apple in the fridge"], + _batched_info(["go to cabinet 1", "help"]), + ) + + def step(self, actions: list[str]): + self.actions.append(actions) + observation, commands, done, won = self.transitions.pop(0) + return ( + [observation], + [10.0 if won else 0.0], + [done], + _batched_info(commands, won=won), + ) + + def close(self) -> None: + self.closed = True + + +def test_alfworld_game_files_are_canonicalized_before_seeded_reset() -> None: + class FakeBaseEnv: + game_files = [ + r"root\z-task\game.tw-pddl", + "root/a-task/game.tw-pddl", + r"root\m-task\game.tw-pddl", + ] + num_games = 99 + + base_env = FakeBaseEnv() + _canonicalize_game_files(base_env) + + assert base_env.game_files == [ + "root/a-task/game.tw-pddl", + r"root\m-task\game.tw-pddl", + r"root\z-task\game.tw-pddl", + ] + assert base_env.num_games == 3 + + +def test_alfworld_selected_game_file_uses_manifest_relative_path() -> None: + class FakeBaseEnv: + game_files = [ + "/data/alfworld/json_2.1.1/train/task-b/game.tw-pddl", + "/data/alfworld/json_2.1.1/train/task-a/game.tw-pddl", + ] + num_games = 2 + + base_env = FakeBaseEnv() + _select_game_file( + base_env, + "json_2.1.1/train/task-a/game.tw-pddl", + data_root="/data/alfworld", + ) + + assert base_env.game_files == ["/data/alfworld/json_2.1.1/train/task-a/game.tw-pddl"] + assert base_env.num_games == 1 + + +def test_alfworld_env_tracks_reference_fields_and_nothing_happens() -> None: + raw_env = FakeRawAlfWorldEnv( + [ + ( + "You arrive at cabinet 1.", + ["take apple 1 from cabinet 1"], + False, + False, + ), + ( + "You take apple 1 from cabinet 1.", + ["heat apple 1 with microwave 1"], + False, + False, + ), + ( + "You heat apple 1 with microwave 1.", + ["go to fridge 1"], + False, + False, + ), + ( + "Nothing happens.", + ["look"], + True, + False, + ), + ] + ) + env = AlfWorldTextEnv(env=raw_env) + + initial = env.reset() + assert initial.raw_observation.startswith("Room intro.") + assert initial.admissible_commands == ("go to cabinet 1", "help") + assert initial.gamefile == "games/task-1/game.tw-pddl" + assert initial.tracker.to_mapping() == { + "location": "middle of a room", + "holding": "nothing", + "history_items": {}, + "item_location": {}, + } + + at_cabinet = env.step("go to cabinet 1") + assert at_cabinet.tracker.to_mapping()["location"] == "cabinet 1" + + holding = env.step("take apple 1 from cabinet 1") + holding_tracker = holding.tracker.to_mapping() + assert holding_tracker["holding"] == "apple 1" + assert holding_tracker["item_location"]["apple 1"] == { + "old_location": "cabinet 1", + "new_location": "cabinet 1", + } + + heated = env.step("heat apple 1 with microwave 1") + heated_tracker = heated.tracker.to_mapping() + assert heated_tracker["history_items"]["apple 1"] == { + "heated": True, + "cooled": False, + "cleaned": False, + "slice": False, + } + + unchanged = env.step("go to fridge 1") + assert unchanged.done is True + assert unchanged.tracker == heated.tracker + assert raw_env.actions == [ + ["go to cabinet 1"], + ["take apple 1 from cabinet 1"], + ["heat apple 1 with microwave 1"], + ["go to fridge 1"], + ] + + env.close() + env.close() + assert raw_env.closed is True + + +def test_alfworld_env_is_lazy_and_factory_receives_seed_and_split() -> None: + calls: list[tuple[str | Path, int, str, str | Path | None]] = [] + raw_env = FakeRawAlfWorldEnv([]) + + def factory( + config_path: str | Path, + seed: int, + split: str, + game_file: str | Path | None, + ): + calls.append((config_path, seed, split, game_file)) + return raw_env + + env = AlfWorldTextEnv( + config_path="config.yaml", + seed=123, + train_eval="eval_in_distribution", + game_file="json_2.1.1/valid_seen/task/game.tw-pddl", + env_factory=factory, + ) + assert calls == [] + env.reset() + assert calls == [ + ( + "config.yaml", + 123, + "eval_in_distribution", + "json_2.1.1/valid_seen/task/game.tw-pddl", + ) + ] + env.close() + + +def test_alfworld_env_rejects_step_after_done() -> None: + raw_env = FakeRawAlfWorldEnv([("Done.", ["look"], True, False)]) + env = AlfWorldTextEnv(env=raw_env) + env.reset() + env.step("look") + with pytest.raises(RuntimeError, match="completed"): + env.step("look") + env.close() + + +def test_alfworld_env_requires_one_element_batches() -> None: + class BadRawEnv: + def reset(self): + return ["one", "two"], _batched_info(["look"]) + + def close(self) -> None: + pass + + env = AlfWorldTextEnv(env=BadRawEnv()) + with pytest.raises(ValueError, match="exactly one"): + env.reset() + env.close() diff --git a/tests/graphgpo/test_custom_advantage.py b/tests/graphgpo/test_custom_advantage.py new file mode 100644 index 000000000..f6c93869e --- /dev/null +++ b/tests/graphgpo/test_custom_advantage.py @@ -0,0 +1,407 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import copy +import math + +import pytest + +from examples.graphgpo.custom_advantage import ( + compute_custom_advantage, + compute_group_advantages, +) +from examples.graphgpo.graph_credit import SUCCESS + + +def _slot( + trajectory_id: str, + transitions: list[tuple[str, str, str]], + *, + task_id: str = "task", + success: bool = False, + invalid_turns: set[int] | None = None, + rollout_group_id: str = "group-0", + policy_version: int = 0, +) -> dict[str, dict[str, object]]: + invalid_turns = invalid_turns or set() + declared_return = 10.0 * float(success) - 0.1 * len(invalid_turns) + result: dict[str, dict[str, object]] = {} + for turn_index, (state, action, next_state) in enumerate(transitions): + is_final = turn_index == len(transitions) - 1 + result[f"turn_{turn_index:03d}"] = { + "row_id": f"{rollout_group_id}:{trajectory_id}:{turn_index}", + "rollout_group_id": rollout_group_id, + "policy_version": policy_version, + "task_id": task_id, + "trajectory_id": trajectory_id, + "turn_index": turn_index, + "state_key": state, + "action": action, + "next_state_key": SUCCESS if success and is_final else next_state, + "is_action_valid": turn_index not in invalid_turns, + "success": success, + "terminal": success and is_final, + "truncated": not success and is_final, + "episode_return": declared_return, + } + return result + + +def _golden_group() -> list[dict[str, dict[str, object]]]: + return [ + _slot( + "tau0", + [("A", "to-b", "B"), ("B", "finish", "natural-terminal")], + success=True, + ), + _slot( + "tau1", + [("A", "to-c-1", "C"), ("C", "to-d", "D")], + ), + _slot( + "tau2", + [ + ("A", "to-c-2", "C"), + ("C", "to-b", "B"), + ("B", "back-to-c", "C"), + ], + ), + ] + + +def test_custom_advantage_three_method_golden_vectors(): + group = _golden_group() + + grpo = compute_group_advantages(group, method="grpo", expected_group_size=3) + gigpo = compute_group_advantages(group, method="gigpo", expected_group_size=3) + graphgpo = compute_group_advantages(group, method="graphgpo", expected_group_size=3) + + episode_high = 1.1547003383792862 + episode_low = -0.5773501691896431 + gigpo_a_high = 1.1547003278529744 + gigpo_a_low = -0.5773501639264869 + gigpo_b_high = 0.7071066811865616 + gigpo_b_low = -0.7071066811865616 + graph_a_high = 1.154698316161306 + graph_a_low = -0.577349158080653 + graph_b_high = 0.707106680176461 + graph_b_low = -0.707106680176461 + c_graph_high = 0.7071057710869802 + c_graph_low = -0.7071057710869802 + + assert grpo == [ + { + "turn_000": pytest.approx(episode_high), + "turn_001": pytest.approx(episode_high), + }, + { + "turn_000": pytest.approx(episode_low), + "turn_001": pytest.approx(episode_low), + }, + { + "turn_000": pytest.approx(episode_low), + "turn_001": pytest.approx(episode_low), + "turn_002": pytest.approx(episode_low), + }, + ] + assert gigpo == [ + { + "turn_000": pytest.approx(episode_high + gigpo_a_high), + "turn_001": pytest.approx(episode_high + gigpo_b_high), + }, + { + "turn_000": pytest.approx(episode_low + gigpo_a_low), + "turn_001": pytest.approx(episode_low), + }, + { + "turn_000": pytest.approx(episode_low + gigpo_a_low), + "turn_001": pytest.approx(episode_low), + "turn_002": pytest.approx(episode_low + gigpo_b_low), + }, + ] + assert graphgpo == [ + { + "turn_000": pytest.approx(episode_high + graph_a_high), + "turn_001": pytest.approx(episode_high + graph_b_high), + }, + { + "turn_000": pytest.approx(episode_low + graph_a_low), + "turn_001": pytest.approx(episode_low + c_graph_low), + }, + { + "turn_000": pytest.approx(episode_low + graph_a_low), + "turn_001": pytest.approx(episode_low + c_graph_high), + "turn_002": pytest.approx(episode_low + graph_b_low), + }, + ] + + +def test_custom_advantage_preserves_variable_length_slot_shape(): + group = [ + _slot("short", [("A", "wait", "A")]), + _slot( + "long", + [("A", "go", "B"), ("B", "go", "C"), ("C", "wait", "C")], + ), + ] + + result = compute_group_advantages(group, method="graphgpo", expected_group_size=2) + + assert [list(slot) for slot in result] == [ + ["turn_000"], + ["turn_000", "turn_001", "turn_002"], + ] + assert all(math.isfinite(value) for slot in result for value in slot.values()) + + +@pytest.mark.parametrize("method", ["grpo", "gigpo", "graphgpo"]) +def test_custom_advantage_all_fail_group_is_finite_zero(method): + group = [ + _slot("tau0", [("A", "wait-0", "A")]), + _slot("tau1", [("A", "wait-1", "A")]), + ] + + result = compute_group_advantages(group, method=method, expected_group_size=2) + + assert result == [{"turn_000": 0.0}, {"turn_000": 0.0}] + + +def test_custom_advantage_uses_slot_position_not_slot_metadata(): + group = _golden_group() + for slot_index, slot in enumerate(group): + for metadata in slot.values(): + metadata["slot_index"] = 99 - slot_index + metadata["group_generation"] = f"irrelevant-{slot_index}" + + result = compute_group_advantages(group, method="grpo", expected_group_size=3) + + assert len(result) == 3 + assert result[0]["turn_000"] > result[1]["turn_000"] + + +def test_custom_advantage_default_environment_entry_is_group_of_eight( + monkeypatch, +): + for name in ( + "GRAPHGPO_METHOD", + "METHOD", + "GRAPHGPO_EXPECTED_GROUP_SIZE", + "GROUP_SIZE", + "GRAPHGPO_OMEGA", + "OMEGA", + "GRAPHGPO_GAMMA", + "GAMMA", + "GRAPHGPO_BETA", + "BETA", + "GRAPHGPO_BETA_EPISODE", + "BETA_EPISODE", + "GRAPHGPO_EPISODE_WEIGHTING", + "EPISODE_WEIGHTING", + ): + monkeypatch.delenv(name, raising=False) + group = [_slot(f"tau{index}", [("A", f"wait-{index}", "A")]) for index in range(8)] + + result = compute_custom_advantage(group) + + assert result == [{"turn_000": 0.0} for _ in range(8)] + + +def test_custom_advantage_environment_entry_accepts_recipe_short_names( + monkeypatch, +): + monkeypatch.setenv("METHOD", "grpo") + monkeypatch.setenv("GROUP_SIZE", "3") + monkeypatch.setenv("BETA_EPISODE", "2") + monkeypatch.setenv("EPISODE_WEIGHTING", "trajectory_once") + + result = compute_custom_advantage(_golden_group()) + + assert result[0]["turn_000"] == pytest.approx(2.3080, rel=1e-3) + assert result[1]["turn_000"] == pytest.approx(-1.1547, rel=1e-3) + + +def test_custom_advantage_environment_defaults_to_trajectory_once( + monkeypatch, +): + monkeypatch.setenv("METHOD", "grpo") + monkeypatch.setenv("GROUP_SIZE", "2") + monkeypatch.delenv("GRAPHGPO_EPISODE_WEIGHTING", raising=False) + monkeypatch.delenv("EPISODE_WEIGHTING", raising=False) + group = [ + _slot( + "success", + [("A", "a", "B"), ("B", "b", "C"), ("C", "finish", "done")], + success=True, + ), + _slot("failure", [("A", "fail", "Z")]), + ] + + result = compute_custom_advantage(group) + + expected = 5.0 / (math.sqrt(50.0) + 1e-6) + assert result[0] == { + "turn_000": pytest.approx(expected), + "turn_001": pytest.approx(expected), + "turn_002": pytest.approx(expected), + } + assert result[1] == {"turn_000": pytest.approx(-expected)} + + +def test_custom_advantage_reference_cross_steps_is_explicit_opt_in(monkeypatch): + monkeypatch.setenv("METHOD", "grpo") + monkeypatch.setenv("GROUP_SIZE", "2") + monkeypatch.setenv("EPISODE_WEIGHTING", "reference_cross_steps") + group = [ + _slot( + "success", + [("A", "a", "B"), ("B", "b", "C"), ("C", "finish", "done")], + success=True, + ), + _slot("failure", [("A", "fail", "Z")]), + ] + + result = compute_custom_advantage(group) + + assert result[0] == { + "turn_000": pytest.approx(2.5 / (5.0 + 1e-6)), + "turn_001": pytest.approx(2.5 / (5.0 + 1e-6)), + "turn_002": pytest.approx(2.5 / (5.0 + 1e-6)), + } + assert result[1] == {"turn_000": pytest.approx(-7.5 / (5.0 + 1e-6))} + + +def test_custom_advantage_rejects_missing_metadata(): + group = _golden_group() + del group[0]["turn_000"]["state_key"] + + with pytest.raises(ValueError, match="missing required metadata field 'state_key'"): + compute_group_advantages(group, expected_group_size=3) + + +def test_custom_advantage_rejects_extra_slot(): + group = _golden_group() + group.append(_slot("tau3", [("A", "wait", "A")])) + + with pytest.raises(ValueError, match="rollout group size mismatch"): + compute_group_advantages(group, expected_group_size=3) + + +def test_custom_advantage_rejects_missing_slot(): + group = _golden_group()[:-1] + + with pytest.raises(ValueError, match="rollout group size mismatch"): + compute_group_advantages(group, expected_group_size=3) + + +def test_custom_advantage_rejects_duplicate_trajectory_id(): + group = _golden_group() + for metadata in group[1].values(): + metadata["trajectory_id"] = "tau0" + + with pytest.raises(ValueError, match="appears in more than one slot"): + compute_group_advantages(group, expected_group_size=3) + + +@pytest.mark.parametrize("non_finite", [math.nan, math.inf, -math.inf]) +def test_custom_advantage_rejects_non_finite_episode_return(non_finite): + group = _golden_group() + group[0]["turn_000"]["episode_return"] = non_finite + + with pytest.raises(ValueError, match="episode_return must be a finite"): + compute_group_advantages(group, expected_group_size=3) + + +def test_custom_advantage_rejects_non_finite_cost(): + group = _golden_group() + group[0]["turn_000"]["cost"] = math.nan + + with pytest.raises(ValueError, match="cost must be a finite"): + compute_group_advantages(group, expected_group_size=3) + + +def test_custom_advantage_rejects_turn_gap(): + group = _golden_group() + moved = group[2].pop("turn_001") + moved["turn_index"] = 3 + group[2]["turn_003"] = moved + + with pytest.raises(ValueError, match="non-contiguous turn indices"): + compute_group_advantages(group, expected_group_size=3) + + +def test_custom_advantage_rejects_noncanonical_turn_name(): + group = _golden_group() + group[0]["edge_000"] = group[0].pop("turn_000") + + with pytest.raises(ValueError, match="canonical name"): + compute_group_advantages(group, expected_group_size=3) + + +def test_custom_advantage_rejects_task_mix(): + group = _golden_group() + for metadata in group[2].values(): + metadata["task_id"] = "different-task" + + with pytest.raises(ValueError, match="cannot mix multiple task_id"): + compute_group_advantages(group, expected_group_size=3) + + +def test_custom_advantage_rejects_identity_mix_and_duplicate_row(): + mixed_group = _golden_group() + mixed_group[1]["turn_000"]["rollout_group_id"] = "other-group" + with pytest.raises(ValueError, match="multiple rollout_group_id"): + compute_group_advantages(mixed_group, expected_group_size=3) + + mixed_policy = _golden_group() + mixed_policy[1]["turn_000"]["policy_version"] = 1 + with pytest.raises(ValueError, match="multiple policy_version"): + compute_group_advantages(mixed_policy, expected_group_size=3) + + duplicate_row = _golden_group() + duplicate_row[1]["turn_000"]["row_id"] = duplicate_row[0]["turn_000"]["row_id"] + with pytest.raises(ValueError, match="duplicate row_id"): + compute_group_advantages(duplicate_row, expected_group_size=3) + + +def test_custom_advantage_rejects_inconsistent_episode_return(): + group = _golden_group() + group[0]["turn_001"]["episode_return"] = 9.9 + + with pytest.raises(ValueError, match="inconsistent episode_return"): + compute_group_advantages(group, expected_group_size=3) + + +def test_custom_advantage_rejects_episode_return_mismatch(): + group = _golden_group() + for metadata in group[0].values(): + metadata["episode_return"] = 9.9 + + with pytest.raises(ValueError, match="episode_return mismatch"): + compute_group_advantages(group, expected_group_size=3) + + +def test_custom_advantage_rejects_inconsistent_success_flags(): + group = _golden_group() + group[0]["turn_000"]["success"] = False + + with pytest.raises(ValueError, match="inconsistent success flags"): + compute_group_advantages(group, expected_group_size=3) + + +def test_custom_advantage_rejects_early_terminal(): + group = _golden_group() + group[0]["turn_000"]["terminal"] = True + + with pytest.raises(ValueError, match="terminates before its final turn"): + compute_group_advantages(group, expected_group_size=3) + + +def test_custom_advantage_rejects_unknown_method_without_mutating_input(): + group = _golden_group() + original = copy.deepcopy(group) + + with pytest.raises(ValueError, match="method must be one of"): + compute_group_advantages(group, method="unknown", expected_group_size=3) + + assert group == original diff --git a/tests/graphgpo/test_diagnostics.py b/tests/graphgpo/test_diagnostics.py new file mode 100644 index 000000000..f4d23a16b --- /dev/null +++ b/tests/graphgpo/test_diagnostics.py @@ -0,0 +1,175 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import json +from dataclasses import replace + +from examples.graphgpo.diagnostics import ( + DIAGNOSTICS_PATH_ENV, + DIAGNOSTICS_SCHEMA, + JsonlDiagnosticsWriter, + build_graph_diagnostics, + diagnostics_callback_from_environment, + summarize_graph_diagnostics, + write_graph_diagnostics_summary, +) +from examples.graphgpo.graph_credit import ( + SUCCESS, + DistanceResult, + EdgeOccurrence, + Turn, + compute_method_advantages, + finalize_trajectory, +) + + +def test_graph_diagnostics_fixed_metric_oracle(): + occurrences = ( + EdgeOccurrence(("tau0", 0), "S", "to-b", "B", 1.0), + EdgeOccurrence(("tau1", 0), "S", "to-b", "B", 1.0), + EdgeOccurrence(("tau2", 0), "S", "to-u", "U", 1.0), + EdgeOccurrence(("tau0", 1), "B", "finish", SUCCESS, 1.0), + ) + distances = DistanceResult( + distances={SUCCESS: 0.0, "B": 1.0, "S": 2.0, "U": 3.0}, + max_finite_distance=2.0, + has_success=True, + ) + + record = build_graph_diagnostics( + task_id="task", + nodes=frozenset((SUCCESS, "B", "S", "U")), + occurrences=occurrences, + distances=distances.distances, + max_finite_distance=distances.max_finite_distance, + has_success=distances.has_success, + advantages={ + ("tau0", 0): 1.0, + ("tau1", 0): -1.0, + ("tau2", 0): 0.0, + ("tau0", 1): 0.0, + }, + graph_build_ns=11, + reverse_dijkstra_ns=13, + graph_advantage_ns=17, + rollout_group_id="group-7", + policy_version=42, + ) + + assert record.schema == DIAGNOSTICS_SCHEMA + assert record.rollout_group_id == "builtins.str:'group-7'" + assert record.policy_version == "builtins.int:42" + assert record.graph_credit_total_ns == 41 + assert record.node_count == 4 + assert record.edge_occurrence_count == 4 + assert record.duplicate_edge_count == 1 + assert record.duplicate_edge_rate == 0.25 + assert record.shared_state_count == 1 + assert record.shared_state_rate == 0.5 + assert record.singleton_source_count == 1 + assert record.singleton_source_rate == 0.5 + assert record.nonzero_graph_advantage_count == 2 + assert record.nonzero_graph_advantage_rate == 0.5 + assert not record.all_fail + assert record.unreachable_node_count == 1 + assert record.unreachable_node_rate == 0.25 + assert record.distance_histogram == {"0": 1, "1": 1, "2": 1, "3": 1} + + +def test_diagnostics_are_default_off(monkeypatch): + monkeypatch.delenv(DIAGNOSTICS_PATH_ENV, raising=False) + + assert diagnostics_callback_from_environment() is None + + +def test_jsonl_writer_is_enabled_only_by_explicit_path(monkeypatch, tmp_path): + output_path = tmp_path / "nested" / "diagnostics.jsonl" + monkeypatch.setenv(DIAGNOSTICS_PATH_ENV, str(output_path)) + writer = diagnostics_callback_from_environment() + assert isinstance(writer, JsonlDiagnosticsWriter) + + turn = Turn( + "task", + "tau", + 0, + "S", + "finish", + "done", + True, + True, + False, + True, + ) + turns = finalize_trajectory((turn,), success=True) + compute_method_advantages( + "graphgpo", + turns, + expected_group_size=1, + graph_diagnostics_callback=writer, + ) + + payload = json.loads(output_path.read_text(encoding="utf-8")) + assert payload["schema"] == DIAGNOSTICS_SCHEMA + assert payload["node_count"] == 2 + assert payload["edge_occurrence_count"] == 1 + assert payload["all_fail"] is False + assert payload["graph_build_ns"] >= 0 + assert payload["reverse_dijkstra_ns"] >= 0 + assert payload["graph_advantage_ns"] >= 0 + + +def test_diagnostics_summary_reports_median_and_nearest_rank_p95(tmp_path): + occurrence = EdgeOccurrence(("tau", 0), "S", "finish", SUCCESS, 1.0) + base = build_graph_diagnostics( + task_id="task", + nodes=frozenset(("S", SUCCESS)), + occurrences=(occurrence,), + distances={"S": 1.0, SUCCESS: 0.0}, + max_finite_distance=1.0, + has_success=True, + advantages={("tau", 0): 0.0}, + graph_build_ns=10, + reverse_dijkstra_ns=20, + graph_advantage_ns=30, + ) + records = ( + base, + replace( + base, + graph_build_ns=20, + reverse_dijkstra_ns=30, + graph_advantage_ns=40, + graph_credit_total_ns=90, + ), + replace( + base, + graph_build_ns=100, + reverse_dijkstra_ns=200, + graph_advantage_ns=300, + graph_credit_total_ns=600, + all_fail=True, + ), + ) + + summary = summarize_graph_diagnostics(records) + + assert summary.record_count == 3 + assert summary.timing_ns["graph_build_ns"] == { + "median": 20.0, + "p95": 100.0, + "total": 130, + } + assert summary.all_fail_count == 1 + assert summary.all_fail_rate == 1 / 3 + assert summary.distance_histogram == {"0": 3, "1": 3} + + raw_path = tmp_path / "raw.jsonl" + writer = JsonlDiagnosticsWriter(raw_path) + for record in records: + writer(record) + output_path = tmp_path / "summary.json" + written = write_graph_diagnostics_summary(raw_path, output_path) + + assert written == summary + assert json.loads(output_path.read_text(encoding="utf-8"))["timing_ns"]["graph_build_ns"]["p95"] == 100.0 diff --git a/tests/graphgpo/test_dynamic_turn_accounting.py b/tests/graphgpo/test_dynamic_turn_accounting.py new file mode 100644 index 000000000..8aa9feac2 --- /dev/null +++ b/tests/graphgpo/test_dynamic_turn_accounting.py @@ -0,0 +1,184 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import ast +import copy +from argparse import Namespace +from collections import deque +from dataclasses import dataclass, field +from pathlib import Path +from types import SimpleNamespace +from typing import Any + + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _load_source_definitions(path: Path, names: set[str], namespace: dict[str, Any]) -> dict[str, Any]: + """Execute selected stdlib-only definitions without importing + Ray/Megatron.""" + source = path.read_text(encoding="utf-8") + tree = ast.parse(source, filename=str(path)) + definitions = [ + node for node in tree.body if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name in names + ] + found = {node.name for node in definitions} + if found != names: + raise AssertionError(f"Missing source definitions in {path}: {sorted(names - found)}") + + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + *definitions, + ], + type_ignores=[], + ) + exec(compile(ast.fix_missing_locations(module), str(path), "exec"), namespace) + return namespace + + +def _reward_waiting_group_type(): + namespace = _load_source_definitions( + REPO_ROOT / "relax" / "agentic" / "pipeline" / "reward.py", + {"RewardWaitingGroup"}, + { + "Any": Any, + "PendingExportUnit": object, + "copy": copy, + "dataclass": dataclass, + "field": field, + }, + ) + return namespace["RewardWaitingGroup"] + + +def _build_rollout_plan(args: Namespace): + namespace = _load_source_definitions( + REPO_ROOT / "relax" / "backends" / "megatron" / "data.py", + {"RolloutMiniBatchPlan", "build_rollout_minibatch_plan"}, + { + "Namespace": Namespace, + "dataclass": dataclass, + }, + ) + return namespace["build_rollout_minibatch_plan"](args, dp_size=1) + + +def _unit( + name: str, + *, + weight_versions: list[str] | None = None, + metadata: dict[str, Any] | None = None, +): + return SimpleNamespace( + name=name, + sample=SimpleNamespace( + metadata={"unit_name": name, **(metadata or {})}, + weight_versions=list(weight_versions or []), + ), + ) + + +def _variable_turn_group(): + waiting_group = _reward_waiting_group_type()(expected_count=2) + waiting_group.add_slot( + slot_idx=0, + units=[_unit("slot0_turn0"), _unit("slot0_turn1")], + ) + waiting_group.add_slot( + slot_idx=1, + units=[_unit("slot1_turn0"), _unit("slot1_turn1"), _unit("slot1_turn2")], + ) + assert waiting_group.is_complete() + return waiting_group + + +def _fixed_n_args(*, global_batch_size: int) -> Namespace: + return Namespace( + rollout_batch_size=1, + n_samples_per_prompt=2, + global_batch_size=global_batch_size, + num_steps_per_rollout=1, + ) + + +def test_fixed_n_windows_leave_variable_turn_rows_undrained(): + """Two trajectory slots with 2/3 turns produce five rows, but only two are + planned.""" + waiting_group = _variable_turn_group() + materialized_units = waiting_group.materialized_units() + assert len(materialized_units) == 5 + + plan = _build_rollout_plan(_fixed_n_args(global_batch_size=2)) + planned_capacity = plan.num_rollout_minis * plan.mini_global_samples + assert planned_capacity == 2 + + remaining = deque(materialized_units) + for _ in range(plan.num_rollout_minis): + for _ in range(min(plan.mini_global_samples, len(remaining))): + remaining.popleft() + + assert [unit.name for unit in remaining] == [ + "slot1_turn0", + "slot1_turn1", + "slot1_turn2", + ] + assert remaining # The fixed number of windows cannot report this partition drained. + + +def test_fixed_n_plan_rejects_using_actual_variable_row_count_as_global_batch(): + """The existing plan cannot be repaired by setting global_batch_size to the + five emitted rows.""" + try: + _build_rollout_plan(_fixed_n_args(global_batch_size=5)) + except ValueError as exc: + assert "fixed-n_samples rollout mini size" in str(exc) + else: + raise AssertionError("The fixed-n plan unexpectedly accepted five variable turn rows.") + + +def test_reward_metadata_injects_relax_captured_policy_version_without_overwriting_source(): + group = _reward_waiting_group_type()(expected_count=1) + unit = _unit("turn_000", weight_versions=["7", "7"]) + group.add_slot(slot_idx=0, units=[unit]) + + payload = group.metadata_by_slot(require_policy_version=True) + + assert payload[0]["turn_000"]["policy_version"] == "7" + assert unit.sample.metadata["policy_version"] == "7" + + +def test_reward_metadata_policy_version_fallback_and_conflicts_fail_closed(): + group_type = _reward_waiting_group_type() + + fallback = group_type(expected_count=1) + fallback_unit = _unit("turn_000", metadata={"start_rollout_id": 3}) + fallback.add_slot(slot_idx=0, units=[fallback_unit]) + assert fallback.metadata_by_slot(require_policy_version=True)[0]["turn_000"]["policy_version"] == "3" + + mixed = group_type(expected_count=1) + mixed_unit = _unit("turn_000", weight_versions=["7", "8"]) + mixed.add_slot(slot_idx=0, units=[mixed_unit]) + try: + mixed.metadata_by_slot(require_policy_version=True) + except RuntimeError as exc: + assert "multiple policy versions" in str(exc) + else: + raise AssertionError("mixed policy versions were accepted") + assert "policy_version" not in mixed_unit.sample.metadata + + conflicting = group_type(expected_count=1) + conflicting_unit = _unit( + "turn_000", + weight_versions=["7"], + metadata={"policy_version": "8"}, + ) + conflicting.add_slot(slot_idx=0, units=[conflicting_unit]) + try: + conflicting.metadata_by_slot(require_policy_version=True) + except RuntimeError as exc: + assert "conflicts" in str(exc) + else: + raise AssertionError("a user-declared conflicting policy version was overwritten") + assert conflicting_unit.sample.metadata["policy_version"] == "8" diff --git a/tests/graphgpo/test_eval_logger.py b/tests/graphgpo/test_eval_logger.py new file mode 100644 index 000000000..479f0b14d --- /dev/null +++ b/tests/graphgpo/test_eval_logger.py @@ -0,0 +1,86 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from examples.graphgpo.eval_logger import ( + build_eval_metrics, + episode_metrics, + log_eval_rollout_data, +) + + +def _sample( + trajectory_id: str, + *, + success: bool, + episode_return: float, + truncated: bool = False, +) -> SimpleNamespace: + return SimpleNamespace( + metadata={ + "trajectory_id": trajectory_id, + "success": success, + "episode_return": episode_return, + "truncated": truncated, + } + ) + + +class EvalLoggerTest(unittest.TestCase): + def test_episode_metrics_give_long_and_short_trajectories_one_vote_each(self) -> None: + samples = [ + _sample("long", success=True, episode_return=9.8), + _sample("long", success=True, episode_return=9.8), + _sample("long", success=True, episode_return=9.8), + _sample( + "short", + success=False, + episode_return=-0.1, + truncated=True, + ), + ] + + metrics = episode_metrics(samples) + + self.assertEqual(metrics["episode_count"], 2) + self.assertEqual(metrics["success_rate"], 0.5) + self.assertAlmostEqual(metrics["episode_return_mean"], 4.85) + self.assertEqual(metrics["truncated_rate"], 0.5) + + def test_episode_metrics_reject_inconsistent_turn_metadata(self) -> None: + samples = [ + _sample("trajectory", success=True, episode_return=10.0), + _sample("trajectory", success=False, episode_return=10.0), + ] + + with self.assertRaisesRegex(ValueError, "inconsistent success"): + episode_metrics(samples) + + def test_eval_hook_namespaces_metrics_and_replaces_default_row_logging(self) -> None: + data = { + "alfworld": { + "samples": [ + _sample("a", success=True, episode_return=10.0), + _sample("a", success=True, episode_return=10.0), + _sample("b", success=False, episode_return=0.0, truncated=True), + ] + } + } + expected = build_eval_metrics(data, {"eval/runtime_s": 1.25}) + + with patch("examples.graphgpo.eval_logger._emit_metrics") as emit: + handled = log_eval_rollout_data(7, object(), data, {"eval/runtime_s": 1.25}) + + self.assertTrue(handled) + emit.assert_called_once_with(7, unittest.mock.ANY, expected) + self.assertEqual(expected["eval/alfworld/episode_count"], 2) + self.assertEqual(expected["eval/alfworld/success_rate"], 0.5) + self.assertEqual(expected["eval/alfworld/truncated_rate"], 0.5) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/graphgpo/test_graph_credit.py b/tests/graphgpo/test_graph_credit.py new file mode 100644 index 000000000..ddcf5d6e8 --- /dev/null +++ b/tests/graphgpo/test_graph_credit.py @@ -0,0 +1,569 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import math +from dataclasses import replace + +import pytest + +from examples.graphgpo.graph_credit import ( + SUCCESS, + Turn, + build_occurrence_graph, + compute_method_advantages, + episode_advantages, + episode_return, + finalize_trajectory, + gigpo_discounted_returns, + graph_advantages, + graph_raw_returns, + reverse_shortest_distances, + standardize_by_key, +) + + +def _trajectory( + trajectory_id: str, + transitions: list[tuple[str, str, str]], + *, + success: bool, + task_id: str = "task", + invalid_turns: set[int] | None = None, +) -> tuple[Turn, ...]: + invalid_turns = invalid_turns or set() + turns = [] + for turn_index, (state, action, next_state) in enumerate(transitions): + is_final = turn_index == len(transitions) - 1 + turns.append( + Turn( + task_id=task_id, + trajectory_id=trajectory_id, + turn_index=turn_index, + state_key=state, + action=action, + next_state_key=next_state, + success=success, + terminal=success and is_final, + truncated=not success and is_final, + is_action_valid=turn_index not in invalid_turns, + ) + ) + return finalize_trajectory(turns, success=success) + + +def _golden_turns() -> tuple[Turn, ...]: + return ( + *_trajectory( + "tau0", + [("A", "to-b", "B"), ("B", "finish", "natural-terminal")], + success=True, + ), + *_trajectory( + "tau1", + [("A", "to-c-1", "C"), ("C", "to-d", "D")], + success=False, + ), + *_trajectory( + "tau2", + [("A", "to-c-2", "C"), ("C", "to-b", "B"), ("B", "back-to-c", "C")], + success=False, + ), + ) + + +def test_graph_credit_reference_golden_vector(): + graph = build_occurrence_graph(_golden_turns()) + distances = reverse_shortest_distances(graph) + + assert len(graph.occurrences) == 7 + assert sum(edge.source == "A" and edge.target == "C" for edge in graph.occurrences) == 2 + assert distances.has_success + assert distances.max_finite_distance == pytest.approx(2.0) + assert distances.distances == { + SUCCESS: 0.0, + "A": 2.0, + "B": 1.0, + "C": 2.0, + "D": 3.0, + } + + raw = graph_raw_returns(graph, distances) + assert raw == pytest.approx( + { + ("tau0", 0): 1.0, + ("tau0", 1): 10.0, + ("tau1", 0): 0.1, + ("tau1", 1): 0.01, + ("tau2", 0): 0.1, + ("tau2", 1): 1.0, + ("tau2", 2): 0.1, + } + ) + + advantages = graph_advantages(graph, distances) + assert advantages == pytest.approx( + { + ("tau0", 0): 1.15469831616131, + ("tau0", 1): 0.707106680176461, + ("tau1", 0): -0.577349158080653, + ("tau1", 1): -0.70710577108698, + ("tau2", 0): -0.577349158080653, + ("tau2", 1): 0.70710577108698, + ("tau2", 2): -0.707106680176461, + } + ) + + +def test_graph_credit_paper_convention_keeps_formula_difference_explicit(): + graph = build_occurrence_graph(_golden_turns()) + distances = reverse_shortest_distances(graph) + reference = graph_raw_returns(graph, distances, convention="reference") + paper = graph_raw_returns(graph, distances, convention="paper") + + assert paper == pytest.approx({key: value * 0.1 for key, value in reference.items()}) + + +def test_graph_credit_same_source_distance_oracle_is_10_1_point_1(): + turns = ( + *_trajectory("distance-0", [("S", "finish-now", "done")], success=True), + *_trajectory( + "distance-1", + [("S", "to-b", "B"), ("B", "finish", "done")], + success=True, + ), + *_trajectory( + "distance-2", + [ + ("S", "to-c", "C"), + ("C", "to-b", "B"), + ("B", "finish", "done"), + ], + success=True, + ), + ) + graph = build_occurrence_graph(turns) + distances = reverse_shortest_distances(graph) + raw = graph_raw_returns(graph, distances, convention="reference") + + assert [ + distances.distances[turns[0].next_state_key], + distances.distances[turns[1].next_state_key], + distances.distances[turns[3].next_state_key], + ] == pytest.approx([0.0, 1.0, 2.0]) + assert [ + raw[("distance-0", 0)], + raw[("distance-1", 0)], + raw[("distance-2", 0)], + ] == pytest.approx([10.0, 1.0, 0.1]) + + +def test_graph_credit_all_fail_group_has_zero_graph_advantage(): + turns = ( + *_trajectory("tau0", [("A", "left", "B")], success=False), + *_trajectory("tau1", [("A", "right", "C")], success=False), + ) + graph = build_occurrence_graph(turns) + distances = reverse_shortest_distances(graph) + + assert not distances.has_success + assert distances.max_finite_distance == 0.0 + assert set(distances.distances.values()) == {1.0} + assert graph_advantages(graph, distances) == {("tau0", 0): 0.0, ("tau1", 0): 0.0} + assert compute_method_advantages("graphgpo", turns, expected_group_size=2) == { + ("tau0", 0): 0.0, + ("tau1", 0): 0.0, + } + + +def test_graph_credit_all_success_zero_variance_is_finite_zero(): + turns = ( + *_trajectory("tau0", [("S", "finish-0", "done")], success=True), + *_trajectory("tau1", [("S", "finish-1", "done")], success=True), + ) + + result = compute_method_advantages( + "graphgpo", + turns, + expected_group_size=2, + ) + + assert result == {("tau0", 0): 0.0, ("tau1", 0): 0.0} + assert all(math.isfinite(value) for value in result.values()) + + +def test_graph_credit_cycle_self_loop_and_success_edge(): + turns = _trajectory( + "tau0", + [("S", "wait", "S"), ("S", "finish", "natural-terminal")], + success=True, + ) + graph = build_occurrence_graph(turns) + distances = reverse_shortest_distances(graph) + + assert distances.distances["S"] == pytest.approx(1.0) + assert graph_raw_returns(graph, distances) == pytest.approx({("tau0", 0): 1.0, ("tau0", 1): 10.0}) + assert graph_advantages(graph, distances) == pytest.approx( + {("tau0", 0): -0.707106680176461, ("tau0", 1): 0.707106680176461} + ) + + +def test_graph_credit_repeated_occurrences_change_group_statistics(): + graph = build_occurrence_graph(_golden_turns()) + distances = reverse_shortest_distances(graph) + advantages = graph_advantages(graph, distances) + + assert advantages[("tau0", 0)] == pytest.approx(1.15469831616131) + assert advantages[("tau1", 0)] == pytest.approx(-0.577349158080653) + assert advantages[("tau2", 0)] == pytest.approx(-0.577349158080653) + + +def test_finalize_trajectory_replaces_last_real_target_without_adding_a_row(): + unfinished = ( + Turn("task", "tau", 0, "A", "to-b", "B", False, False, False, True), + Turn( + "task", + "tau", + 1, + "B", + "finish", + "natural-terminal", + False, + False, + True, + True, + ), + ) + + finalized = finalize_trajectory(unfinished, success=True) + + assert len(finalized) == len(unfinished) + assert finalized[0].next_state_key == "B" + assert finalized[-1].next_state_key == SUCCESS + assert finalized[-1].terminal + assert not finalized[-1].truncated + assert all(turn.success for turn in finalized) + assert finalize_trajectory(finalized, success=True) == finalized + + +def test_episode_return_counts_invalid_actions_once_per_trajectory(): + clean = _trajectory( + "clean", + [("A", "a", "B"), ("B", "b", "C"), ("C", "c", "done")], + success=True, + ) + invalid = _trajectory( + "invalid", + [("A", "a", "B"), ("B", "b", "C"), ("C", "c", "done")], + success=True, + invalid_turns={0, 2}, + ) + + assert episode_return(clean) == pytest.approx(10.0) + assert episode_return(invalid) == pytest.approx(9.8) + normalized = episode_advantages( + {"clean": episode_return(clean), "invalid": episode_return(invalid)}, + {"clean": "task", "invalid": "task"}, + ) + expected = 0.1 / (math.sqrt(0.02) + 1e-6) + assert normalized == pytest.approx({"clean": expected, "invalid": -expected}) + + clean_raw = graph_raw_returns( + build_occurrence_graph(clean), + reverse_shortest_distances(build_occurrence_graph(clean)), + ) + invalid_graph = build_occurrence_graph(invalid) + invalid_raw = graph_raw_returns(invalid_graph, reverse_shortest_distances(invalid_graph)) + assert list(clean_raw.values()) == pytest.approx(list(invalid_raw.values())) + + +def test_grpo_uses_one_value_per_trajectory_and_broadcasts_to_turns(): + success = _trajectory( + "success", + [("A", "a", "B"), ("B", "b", "C"), ("C", "finish", "done")], + success=True, + ) + failure = _trajectory("failure", [("A", "fail", "Z")], success=False) + + result = compute_method_advantages( + "grpo", + (*success, *failure), + expected_group_size=2, + ) + expected = 5.0 / (math.sqrt(50.0) + 1e-6) + assert result == pytest.approx( + { + ("failure", 0): -expected, + ("success", 0): expected, + ("success", 1): expected, + ("success", 2): expected, + } + ) + + +def test_grpo_reference_cross_steps_matches_frozen_length_weighting(): + success = _trajectory( + "success", + [("A", "a", "B"), ("B", "b", "C"), ("C", "finish", "done")], + success=True, + ) + failure = _trajectory("failure", [("A", "fail", "Z")], success=False) + + result = compute_method_advantages( + "grpo", + (*success, *failure), + expected_group_size=2, + episode_weighting="reference_cross_steps", + ) + + assert result == pytest.approx( + { + ("failure", 0): -7.5 / (5.0 + 1e-6), + ("success", 0): 2.5 / (5.0 + 1e-6), + ("success", 1): 2.5 / (5.0 + 1e-6), + ("success", 2): 2.5 / (5.0 + 1e-6), + } + ) + + +def test_episode_reference_cross_steps_requires_valid_lengths(): + returns = {"success": 10.0, "failure": 0.0} + tasks = {"success": "task", "failure": "task"} + + with pytest.raises(ValueError, match="trajectory_lengths are required"): + episode_advantages( + returns, + tasks, + weighting="reference_cross_steps", + ) + with pytest.raises(ValueError, match="positive integers"): + episode_advantages( + returns, + tasks, + weighting="reference_cross_steps", + trajectory_lengths={"success": 0, "failure": 1}, + ) + with pytest.raises(ValueError, match="unknown episode weighting"): + episode_advantages(returns, tasks, weighting="unknown") + + +def test_grpo_eight_trajectory_golden_vector(): + returns = {f"tau{index}": 10.0 if index == 0 else 0.0 for index in range(8)} + tasks = {trajectory_id: "task" for trajectory_id in returns} + + result = episode_advantages(returns, tasks, expected_group_size=8) + + assert result["tau0"] == pytest.approx(2.47487303415311) + assert [result[f"tau{index}"] for index in range(1, 8)] == pytest.approx([-0.353553290593302] * 7) + + +def test_gigpo_discounted_return_and_same_state_advantage(): + success = _trajectory( + "success", + [("S", "to-x", "X"), ("X", "finish", "done")], + success=True, + ) + failure = _trajectory( + "failure", + [("S", "to-y", "Y"), ("Y", "fail", "Z")], + success=False, + ) + standalone = _trajectory( + "standalone", + [("A", "a", "B"), ("B", "invalid", "C"), ("C", "finish", "done")], + success=True, + invalid_turns={1}, + ) + immediate = { + ("standalone", 0): 0.0, + ("standalone", 1): -0.1, + ("standalone", 2): 10.0, + } + discounted = gigpo_discounted_returns(standalone, immediate, gamma=0.95) + assert discounted == pytest.approx( + { + ("standalone", 0): 8.93, + ("standalone", 1): 9.4, + ("standalone", 2): 10.0, + } + ) + + step_only = compute_method_advantages( + "gigpo", + (*success, *failure), + expected_group_size=2, + beta_episode=0.0, + ) + expected = 4.75 / (math.sqrt(45.125) + 1e-6) + assert step_only == pytest.approx( + { + ("failure", 0): -expected, + ("failure", 1): 0.0, + ("success", 0): expected, + ("success", 1): 0.0, + } + ) + + +def test_graphgpo_isolates_identical_state_keys_between_tasks(): + turns = ( + *_trajectory( + "task1-success", + [("A", "to-b", "B"), ("B", "finish", "done")], + success=True, + task_id="task1", + ), + *_trajectory( + "task1-fail", + [("A", "to-c", "C")], + success=False, + task_id="task1", + ), + *_trajectory( + "task2-success", + [("A", "to-d", "D"), ("D", "finish", "done")], + success=True, + task_id="task2", + ), + *_trajectory( + "task2-fail", + [("A", "to-e", "E")], + success=False, + task_id="task2", + ), + ) + + graph_only = compute_method_advantages( + "graphgpo", + turns, + expected_group_size=2, + beta_episode=0.0, + ) + + expected = 0.495 / (math.sqrt(0.49005) + 1e-6) + assert graph_only[("task1-success", 0)] == pytest.approx(expected) + assert graph_only[("task1-fail", 0)] == pytest.approx(-expected) + assert graph_only[("task2-success", 0)] == pytest.approx(expected) + assert graph_only[("task2-fail", 0)] == pytest.approx(-expected) + + +def test_compute_advantages_is_permutation_invariant_and_key_complete(): + turns = _golden_turns() + expected = compute_method_advantages( + "graphgpo", + turns, + expected_group_size=3, + ) + shuffled = compute_method_advantages( + "graphgpo", + tuple(reversed(turns)), + expected_group_size=3, + ) + + assert shuffled == expected + assert set(expected) == {turn.key for turn in turns} + + +def test_standardize_singleton_and_zero_variance_are_zero(): + assert standardize_by_key({"only": 3.0}, {"only": "group"}) == {"only": 0.0} + assert standardize_by_key({"a": 3.0, "b": 3.0}, {"a": "group", "b": "group"}) == { + "a": 0.0, + "b": 0.0, + } + + +@pytest.mark.parametrize( + "mutator,match", + [ + ( + lambda turns: (turns[0], replace(turns[1], turn_index=2)), + "non-contiguous", + ), + ( + lambda turns: (turns[0], replace(turns[1], task_id="other-task")), + "multiple tasks|span multiple tasks", + ), + ( + lambda turns: (replace(turns[0], next_state_key="wrong"), turns[1]), + "broken state transition", + ), + ( + lambda turns: (turns[0], replace(turns[1], next_state_key="not-success")), + "final real action", + ), + ( + lambda turns: (replace(turns[0], cost=math.inf), turns[1]), + "finite real", + ), + ( + lambda turns: (replace(turns[0], terminal=True), turns[1]), + "before its final turn", + ), + ( + lambda turns: ( + turns[0], + replace(turns[1], terminal=False, truncated=False), + ), + "must end", + ), + ], +) +def test_graph_credit_rejects_malformed_trajectories(mutator, match): + turns = _trajectory( + "tau", + [("A", "to-b", "B"), ("B", "finish", "done")], + success=True, + ) + malformed = mutator(turns) + + with pytest.raises(ValueError, match=match): + build_occurrence_graph(malformed) + + +def test_graph_credit_rejects_cross_task_graph_and_wrong_group_size(): + task1 = _trajectory("tau1", [("A", "a", "B")], success=False, task_id="task1") + task2 = _trajectory("tau2", [("A", "a", "B")], success=False, task_id="task2") + + with pytest.raises(ValueError, match="exactly one task"): + build_occurrence_graph((*task1, *task2)) + with pytest.raises(ValueError, match="group size mismatch"): + compute_method_advantages("grpo", (*task1, *task2), expected_group_size=2) + + +def test_graph_credit_rejects_unknown_method_bad_rewards_and_non_unit_reference_cost(): + turns = _trajectory("tau", [("A", "a", "B")], success=False) + + with pytest.raises(ValueError, match="unknown method"): + compute_method_advantages("unknown", turns) + with pytest.raises(ValueError, match="exactly one value per turn"): + gigpo_discounted_returns(turns, {}) + with pytest.raises(ValueError, match="finite real"): + gigpo_discounted_returns(turns, {("tau", 0): math.nan}) + + non_unit = (replace(turns[0], cost=2.0),) + graph = build_occurrence_graph(non_unit) + with pytest.raises(ValueError, match="unit edge costs"): + graph_raw_returns(graph, reverse_shortest_distances(graph), convention="reference") + + +def test_standardize_rejects_key_mismatch_and_non_finite_values(): + with pytest.raises(ValueError, match="same keys"): + standardize_by_key({"a": 1.0}, {"b": "group"}) + with pytest.raises(ValueError, match="finite real"): + standardize_by_key({"a": math.nan}, {"a": "group"}) + + +def test_graph_credit_rejects_tampered_distance_results(): + graph = build_occurrence_graph(_golden_turns()) + distances = reverse_shortest_distances(graph) + + with pytest.raises(ValueError, match="keys do not match"): + graph_raw_returns(graph, replace(distances, distances={SUCCESS: 0.0})) + with pytest.raises(ValueError, match="distance zero"): + graph_advantages( + graph, + replace( + distances, + distances={**distances.distances, SUCCESS: 1.0}, + ), + ) diff --git a/tests/graphgpo/test_grpo_one_step_parity.py b/tests/graphgpo/test_grpo_one_step_parity.py new file mode 100644 index 000000000..39b848e66 --- /dev/null +++ b/tests/graphgpo/test_grpo_one_step_parity.py @@ -0,0 +1,381 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""CPU-only one-step parity check for the GraphGPO GRPO adapter. + +The comparison deliberately has two independent sides: + +* the adapter side uses ``compute_group_advantages`` plus Relax's production + GRPO token broadcast and clipped policy-loss definitions; +* the oracle side uses Relax's ordinary trajectory-level reward + normalization, followed by a small, independently written PyTorch loss. + +This catches changes in row payloads as well as changes that only become +visible after backpropagation. Source definitions are loaded with ``ast`` so +the test does not import Ray or Megatron and remains runnable on a CPU host. +""" + +from __future__ import annotations + +import ast +import copy +from argparse import Namespace +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import pytest + +from examples.graphgpo.custom_advantage import compute_group_advantages +from examples.graphgpo.graph_credit import SUCCESS + + +REPO_ROOT = Path(__file__).resolve().parents[2] +RNG_SEED = 20260422 + + +class _Sample: + class Status: + TRUNCATED = "truncated" + + def __init__( + self, + *, + index: int, + reward: float, + group_index: int, + tokens: list[int] | None = None, + loss_mask: list[int] | None = None, + rollout_log_probs: list[float] | None = None, + custom_advantage: float | None = None, + ) -> None: + self.index = index + self.reward = reward + self.group_index = group_index + self.tokens = tokens or [0, 1] + self.response_length = len(loss_mask or [1]) + self.loss_mask = list(loss_mask or [1]) + self.rollout_log_probs = rollout_log_probs + self.custom_advantage = custom_advantage + self.status = "completed" + self.remove_sample = False + self.metadata = None + self.train_metadata = None + self.rollout_routed_experts = None + self.multimodal_train_inputs = None + + def get_reward_value(self, _args: Namespace) -> float: + return self.reward + + +@dataclass(frozen=True) +class _TurnRow: + trajectory_id: str + turn_index: int + episode_return: float + tokens: list[int] + loss_mask: list[int] + rollout_log_probs: list[float] + + +def _source_functions(torch: Any) -> Namespace: + """Load only the production functions needed by this dependency-light + test.""" + + utils_path = REPO_ROOT / "relax" / "utils" / "utils.py" + utils_tree = ast.parse(utils_path.read_text(encoding="utf-8"), filename=str(utils_path)) + utils_names = {"convert_samples_to_train_data", "post_process_rewards"} + utils_nodes = [node for node in utils_tree.body if isinstance(node, ast.FunctionDef) and node.name in utils_names] + assert {node.name for node in utils_nodes} == utils_names + utils_module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + *utils_nodes, + ], + type_ignores=[], + ) + utils_namespace = {"Any": Any, "Sample": _Sample, "torch": torch} + exec(compile(ast.fix_missing_locations(utils_module), str(utils_path), "exec"), utils_namespace) + + ppo_path = REPO_ROOT / "relax" / "utils" / "training" / "ppo_utils.py" + ppo_tree = ast.parse(ppo_path.read_text(encoding="utf-8"), filename=str(ppo_path)) + ppo_names = {"compute_policy_loss", "get_grpo_returns"} + ppo_nodes = [] + for node in ppo_tree.body: + if isinstance(node, ast.FunctionDef) and node.name in ppo_names: + copied_node = copy.deepcopy(node) + copied_node.decorator_list = [] + ppo_nodes.append(copied_node) + assert {node.name for node in ppo_nodes} == ppo_names + ppo_module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + *ppo_nodes, + ], + type_ignores=[], + ) + ppo_namespace = {"torch": torch} + exec(compile(ast.fix_missing_locations(ppo_module), str(ppo_path), "exec"), ppo_namespace) + + return Namespace( + convert_samples_to_train_data=utils_namespace["convert_samples_to_train_data"], + post_process_rewards=utils_namespace["post_process_rewards"], + compute_policy_loss=ppo_namespace["compute_policy_loss"], + get_grpo_returns=ppo_namespace["get_grpo_returns"], + ) + + +def _response_log_probs(torch: Any, model: Any, tokens: list[int], response_length: int): + token_tensor = torch.tensor(tokens, dtype=torch.long) + logits = model(token_tensor[:-1]) + target_tokens = token_tensor[1:] + all_log_probs = torch.log_softmax(logits, dim=-1).gather(1, target_tokens[:, None]).squeeze(1) + response_start = len(tokens) - response_length - 1 + return all_log_probs[response_start:] + + +def _build_fixture(torch: Any): + class TinyPolicy(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.embedding = torch.nn.Embedding(31, 7, dtype=torch.float64) + self.output = torch.nn.Linear(7, 31, bias=True, dtype=torch.float64) + + def forward(self, tokens): + return self.output(torch.tanh(self.embedding(tokens))) + + torch.manual_seed(RNG_SEED) + rollout_policy = TinyPolicy() + + trajectory_shapes = ( + ((2, [1, 1]), (3, [1, 0, 1])), + ((4, [1, 1, 1, 0]),), + ((1, [1]), (2, [0, 1]), (3, [1, 1, 1])), + ) + success_flags = (True, False, False) + invalid_turns = (frozenset(), frozenset(), frozenset({1})) + metadata_by_slot: list[dict[str, dict[str, object]]] = [] + rows: list[_TurnRow] = [] + + for trajectory_index, turns in enumerate(trajectory_shapes): + trajectory_id = f"trajectory-{trajectory_index}" + success = success_flags[trajectory_index] + episode_return = 10.0 * float(success) - 0.1 * len(invalid_turns[trajectory_index]) + slot: dict[str, dict[str, object]] = {} + for turn_index, (response_length, loss_mask) in enumerate(turns): + is_final = turn_index == len(turns) - 1 + unit_name = f"turn_{turn_index:03d}" + slot[unit_name] = { + "row_id": f"row-{trajectory_index}-{turn_index}", + "rollout_group_id": "rollout-group-0", + "policy_version": "policy-version-7", + "task_id": "shared-task", + "trajectory_id": trajectory_id, + "turn_index": turn_index, + "state_key": f"state-{trajectory_index}-{turn_index}", + "action": f"action-{trajectory_index}-{turn_index}", + "next_state_key": (SUCCESS if success and is_final else f"state-{trajectory_index}-{turn_index + 1}"), + "is_action_valid": turn_index not in invalid_turns[trajectory_index], + "success": success, + "terminal": success and is_final, + "truncated": not success and is_final, + "episode_return": episode_return, + } + + prompt_length = 2 + ((trajectory_index + turn_index) % 3) + tokens = torch.randint(0, 31, (prompt_length + response_length,), dtype=torch.long).tolist() + with torch.no_grad(): + rollout_log_probs = _response_log_probs( + torch, + rollout_policy, + tokens, + response_length, + ).tolist() + rows.append( + _TurnRow( + trajectory_id=trajectory_id, + turn_index=turn_index, + episode_return=episode_return, + tokens=tokens, + loss_mask=loss_mask, + rollout_log_probs=rollout_log_probs, + ) + ) + metadata_by_slot.append(slot) + + initial_policy = copy.deepcopy(rollout_policy) + torch.manual_seed(RNG_SEED + 1) + with torch.no_grad(): + for parameter in initial_policy.parameters(): + parameter.add_(0.015 * torch.randn_like(parameter)) + return metadata_by_slot, rows, initial_policy + + +def _args(*, custom_advantage: bool) -> Namespace: + return Namespace( + custom_reward_post_process_path=None, + agentic_custom_advantage_path="adapter" if custom_advantage else None, + advantage_estimator="grpo", + rewards_normalization=True, + grpo_std_normalization=True, + n_samples_per_prompt=3, + multimodal_keys=None, + use_opd=False, + debug_train_only=True, + ) + + +def _loss_grad_and_delta( + torch: Any, + model: Any, + rows: list[_TurnRow], + token_advantages: list[Any], + *, + production_loss: Any | None, +): + current_log_probs = torch.cat([_response_log_probs(torch, model, row.tokens, len(row.loss_mask)) for row in rows]) + rollout_log_probs = torch.cat([torch.tensor(row.rollout_log_probs, dtype=torch.float64) for row in rows]) + advantages = torch.cat(token_advantages).to(dtype=torch.float64) + loss_mask = torch.cat([torch.tensor(row.loss_mask, dtype=torch.float64) for row in rows]) + + if production_loss is not None: + element_loss, _ = production_loss( + rollout_log_probs - current_log_probs, + advantages, + 0.2, + 0.2, + ) + else: + ratio = torch.exp(current_log_probs - rollout_log_probs) + unclipped = -ratio * advantages + clipped = -torch.clamp(ratio, 0.8, 1.2) * advantages + element_loss = torch.maximum(unclipped, clipped) + loss = (element_loss * loss_mask).sum() / loss_mask.sum() + + before = {name: parameter.detach().clone() for name, parameter in model.named_parameters()} + loss.backward() + gradients = {name: parameter.grad.detach().clone() for name, parameter in model.named_parameters()} + with torch.no_grad(): + for parameter in model.parameters(): + parameter.add_(parameter.grad, alpha=-0.05) + deltas = {name: parameter.detach() - before[name] for name, parameter in model.named_parameters()} + return loss.detach(), gradients, deltas + + +def test_grpo_adapter_matches_native_grpo_through_one_cpu_optimizer_step(): + torch = pytest.importorskip("torch") + functions = _source_functions(torch) + metadata_by_slot, rows, initial_policy = _build_fixture(torch) + + adapter_by_slot = compute_group_advantages( + metadata_by_slot, + method="grpo", + expected_group_size=3, + episode_weighting="trajectory_once", + ) + adapter_by_turn = { + (f"trajectory-{slot_index}", int(unit_name.removeprefix("turn_"))): value + for slot_index, slot in enumerate(adapter_by_slot) + for unit_name, value in slot.items() + } + + trajectory_samples = [ + _Sample(index=index, reward=reward, group_index=0) for index, reward in enumerate((10.0, 0.0, -0.1)) + ] + _, native_trajectory_advantages = functions.post_process_rewards( + _args(custom_advantage=False), + trajectory_samples, + ) + native_by_trajectory = { + f"trajectory-{index}": advantage for index, advantage in enumerate(native_trajectory_advantages) + } + + adapter_samples = [ + _Sample( + index=index, + reward=row.episode_return, + group_index=0, + tokens=row.tokens, + loss_mask=row.loss_mask, + rollout_log_probs=row.rollout_log_probs, + custom_advantage=adapter_by_turn[(row.trajectory_id, row.turn_index)], + ) + for index, row in enumerate(rows) + ] + adapter_payload = functions.convert_samples_to_train_data( + _args(custom_advantage=True), + adapter_samples, + ) + oracle_payload = { + "tokens": [row.tokens for row in rows], + "loss_masks": [row.loss_mask for row in rows], + "rollout_log_probs": [row.rollout_log_probs for row in rows], + } + + assert adapter_payload["tokens"] == oracle_payload["tokens"] + assert adapter_payload["loss_masks"] == oracle_payload["loss_masks"] + assert any(mask_value == 0 for mask in adapter_payload["loss_masks"] for mask_value in mask) + for actual, expected in zip( + adapter_payload["rollout_log_probs"], + oracle_payload["rollout_log_probs"], + strict=True, + ): + torch.testing.assert_close( + torch.tensor(actual), + torch.tensor(expected), + rtol=0, + atol=0, + ) + + zero_kl = [torch.zeros(len(row.rollout_log_probs), dtype=torch.float64) for row in rows] + adapter_token_advantages = functions.get_grpo_returns( + torch.tensor(adapter_payload["rewards"], dtype=torch.float64), + zero_kl, + ) + oracle_token_advantages = [ + torch.full_like( + row_kl, + native_by_trajectory[row.trajectory_id], + ) + for row, row_kl in zip(rows, zero_kl, strict=True) + ] + for actual, expected in zip( + adapter_token_advantages, + oracle_token_advantages, + strict=True, + ): + torch.testing.assert_close(actual, expected, rtol=2e-6, atol=2e-6) + + adapter_model = copy.deepcopy(initial_policy) + oracle_model = copy.deepcopy(initial_policy) + adapter_loss, adapter_gradients, adapter_deltas = _loss_grad_and_delta( + torch, + adapter_model, + rows, + adapter_token_advantages, + production_loss=functions.compute_policy_loss, + ) + oracle_loss, oracle_gradients, oracle_deltas = _loss_grad_and_delta( + torch, + oracle_model, + rows, + oracle_token_advantages, + production_loss=None, + ) + + torch.testing.assert_close(adapter_loss, oracle_loss, rtol=2e-6, atol=2e-6) + assert set(adapter_gradients) == set(oracle_gradients) + assert set(adapter_deltas) == set(oracle_deltas) + for name in adapter_gradients: + torch.testing.assert_close( + adapter_gradients[name], + oracle_gradients[name], + rtol=2e-6, + atol=2e-6, + ) + torch.testing.assert_close( + adapter_deltas[name], + oracle_deltas[name], + rtol=2e-6, + atol=2e-6, + ) + assert any(torch.count_nonzero(delta).item() > 0 for delta in adapter_deltas.values()) diff --git a/tests/graphgpo/test_manifest.py b/tests/graphgpo/test_manifest.py new file mode 100644 index 000000000..450a05b89 --- /dev/null +++ b/tests/graphgpo/test_manifest.py @@ -0,0 +1,122 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import hashlib +import tempfile +import unittest +from pathlib import Path + +from examples.graphgpo.manifest import ( + MANIFEST_VERSION, + build_manifest, + infer_task_type, + manifest_bytes, + manifest_sha256, +) + + +class ManifestTest(unittest.TestCase): + def _make_task(self, root: Path, task_dir: str, content: bytes) -> Path: + path = root / task_dir / "trial_1" / "game.tw-pddl" + path.parent.mkdir(parents=True) + path.write_bytes(content) + path.with_name("traj_data.json").write_bytes(b'{"task":"metadata"}') + return path + + def _make_shared_assets(self, root: Path) -> None: + shared_root = root / "shared" + shared_root.mkdir() + (shared_root / "alfred.pddl").write_bytes(b"pddl") + (shared_root / "alfred.twl2").write_bytes(b"grammar") + + def test_manifest_has_fixed_path_order_and_content_hashes(self) -> None: + with tempfile.TemporaryDirectory() as temporary_directory: + root = Path(temporary_directory) + self._make_shared_assets(root) + later = self._make_task( + root, + "pick_heat_then_place_in_recep-Z", + b"later", + ) + earlier = self._make_task( + root, + "pick_and_place_simple-A", + b"earlier", + ) + + manifest = build_manifest( + root, + split="eval_in_distribution", + task_files=[later, earlier], + ) + + self.assertEqual(manifest.schema_version, MANIFEST_VERSION) + self.assertEqual( + [task.task_type for task in manifest.tasks], + ["pick_and_place_simple", "pick_heat_then_place_in_recep"], + ) + self.assertEqual( + manifest.tasks[0].game.sha256, + hashlib.sha256(b"earlier").hexdigest(), + ) + self.assertTrue(manifest.tasks[0].game.relative_path.endswith("game.tw-pddl")) + self.assertTrue(manifest.tasks[0].trajectory.relative_path.endswith("traj_data.json")) + self.assertEqual( + {task.split for task in manifest.tasks}, + {"eval_in_distribution"}, + ) + self.assertEqual( + [asset.asset_name for asset in manifest.shared_assets], + ["alfred.pddl", "alfred.twl2"], + ) + + def test_discovery_and_serialization_are_deterministic(self) -> None: + with tempfile.TemporaryDirectory() as temporary_directory: + root = Path(temporary_directory) + self._make_shared_assets(root) + self._make_task(root, "pick_cool_then_place_in_recep-B", b"b") + self._make_task(root, "look_at_obj_in_light-A", b"a") + + first = build_manifest(root, split="train") + second = build_manifest(root, split="train") + + self.assertEqual(first, second) + self.assertEqual(manifest_bytes(first), manifest_bytes(second)) + self.assertEqual(manifest_sha256(first), manifest_sha256(second)) + + def test_outside_root_is_rejected(self) -> None: + with tempfile.TemporaryDirectory() as root_directory: + with tempfile.TemporaryDirectory() as outside_directory: + root = Path(root_directory) + self._make_shared_assets(root) + outside = self._make_task( + Path(outside_directory), + "pick_and_place_simple-A", + b"outside", + ) + + with self.assertRaisesRegex(ValueError, "outside root"): + build_manifest(root, split="train", task_files=[outside]) + + def test_unknown_task_type_is_rejected(self) -> None: + with self.assertRaisesRegex(ValueError, "cannot infer"): + infer_task_type("unknown-task/trial/game.tw-pddl") + + def test_missing_trajectory_or_shared_asset_is_rejected(self) -> None: + with tempfile.TemporaryDirectory() as temporary_directory: + root = Path(temporary_directory) + game = root / "pick_and_place_simple-A" / "trial_1" / "game.tw-pddl" + game.parent.mkdir(parents=True) + game.write_bytes(b"game") + self._make_shared_assets(root) + + with self.assertRaises(FileNotFoundError): + build_manifest(root, split="train", task_files=[game]) + + game.with_name("traj_data.json").write_bytes(b"{}") + (root / "shared" / "alfred.twl2").unlink() + with self.assertRaisesRegex(ValueError, "expected exactly one"): + build_manifest(root, split="train", task_files=[game]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/graphgpo/test_preflight.py b/tests/graphgpo/test_preflight.py new file mode 100644 index 000000000..d29dbbce4 --- /dev/null +++ b/tests/graphgpo/test_preflight.py @@ -0,0 +1,173 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import hashlib +import json +import tempfile +import unittest +from pathlib import Path + +from examples.graphgpo.preflight import ( + MODEL_LOCK_SCHEMA, + MODEL_REPO_ID, + verify_model_checkpoint, + verify_prepared_artifacts, +) +from examples.graphgpo.prepare_alfworld import ( + PINNED_MODEL_REVISION, + prepare_artifacts, +) + + +class PreflightTest(unittest.TestCase): + def _write_task(self, root: Path, task_dir: str) -> None: + task_root = root / "json_2.1.1" / "train" / task_dir / "trial_1" + task_root.mkdir(parents=True) + (task_root / "game.tw-pddl").write_text(json.dumps({"solvable": True}), encoding="utf-8") + (task_root / "traj_data.json").write_text( + json.dumps({"task_type": "pick_and_place_simple"}), + encoding="utf-8", + ) + + def _prepare(self, root: Path) -> tuple[Path, Path, Path]: + logic = root / "logic" + logic.mkdir() + (logic / "alfred.pddl").write_text("domain", encoding="utf-8") + (logic / "alfred.twl2").write_text("grammar", encoding="utf-8") + self._write_task(root, "pick_and_place_simple-B") + self._write_task(root, "pick_and_place_simple-A") + output = root / "prepared" + prepare_artifacts( + data_root=root, + output_dir=output, + split_roots={"train": root / "json_2.1.1" / "train"}, + max_steps=50, + ) + return ( + output / "prepare.lock.json", + output / "train.prompts.jsonl", + output / "train.manifest.json", + ) + + def _verify(self, lock: Path, prompt: Path, manifest: Path) -> None: + verify_prepared_artifacts( + prepare_lock_path=lock, + split_artifacts={"train": (prompt, manifest)}, + max_steps=50, + model_revision=PINNED_MODEL_REVISION, + alfworld_data_root=lock.parent.parent, + ) + + def test_preflight_accepts_matching_lock_manifest_and_prompt(self) -> None: + with tempfile.TemporaryDirectory() as temporary_directory: + lock, prompt, manifest = self._prepare(Path(temporary_directory)) + self._verify(lock, prompt, manifest) + + def test_preflight_rejects_model_max_steps_and_prompt_sha_mismatches(self) -> None: + with tempfile.TemporaryDirectory() as temporary_directory: + lock, prompt, manifest = self._prepare(Path(temporary_directory)) + with self.assertRaisesRegex(ValueError, "model_revision"): + verify_prepared_artifacts( + prepare_lock_path=lock, + split_artifacts={"train": (prompt, manifest)}, + max_steps=50, + model_revision="0" * 40, + alfworld_data_root=lock.parent.parent, + ) + with self.assertRaisesRegex(ValueError, "max_steps"): + verify_prepared_artifacts( + prepare_lock_path=lock, + split_artifacts={"train": (prompt, manifest)}, + max_steps=49, + model_revision=PINNED_MODEL_REVISION, + alfworld_data_root=lock.parent.parent, + ) + prompt.write_text(prompt.read_text(encoding="utf-8") + "\n", encoding="utf-8") + with self.assertRaisesRegex(ValueError, "prompt SHA256"): + self._verify(lock, prompt, manifest) + + def test_preflight_rejects_split_count_and_task_order_mismatches(self) -> None: + with tempfile.TemporaryDirectory() as temporary_directory: + lock, prompt, manifest = self._prepare(Path(temporary_directory)) + lock_payload = json.loads(lock.read_text(encoding="utf-8")) + lock_payload["splits"]["train"]["task_count"] = 3 + lock.write_text(json.dumps(lock_payload), encoding="utf-8") + with self.assertRaisesRegex(ValueError, "task count"): + self._verify(lock, prompt, manifest) + + with tempfile.TemporaryDirectory() as temporary_directory: + lock, prompt, manifest = self._prepare(Path(temporary_directory)) + manifest_payload = json.loads(manifest.read_text(encoding="utf-8")) + manifest_payload["split"] = "eval_in_distribution" + manifest.write_text(json.dumps(manifest_payload), encoding="utf-8") + lock_payload = json.loads(lock.read_text(encoding="utf-8")) + lock_payload["splits"]["train"]["manifest_sha256"] = hashlib.sha256(manifest.read_bytes()).hexdigest() + lock.write_text(json.dumps(lock_payload), encoding="utf-8") + with self.assertRaisesRegex(ValueError, "manifest split"): + self._verify(lock, prompt, manifest) + + with tempfile.TemporaryDirectory() as temporary_directory: + lock, prompt, manifest = self._prepare(Path(temporary_directory)) + manifest_payload = json.loads(manifest.read_text(encoding="utf-8")) + manifest_payload["tasks"].reverse() + manifest.write_text(json.dumps(manifest_payload), encoding="utf-8") + lock_payload = json.loads(lock.read_text(encoding="utf-8")) + lock_payload["splits"]["train"]["manifest_sha256"] = hashlib.sha256(manifest.read_bytes()).hexdigest() + lock.write_text(json.dumps(lock_payload), encoding="utf-8") + with self.assertRaisesRegex(ValueError, "stable order"): + self._verify(lock, prompt, manifest) + + def test_preflight_rehashes_manifest_referenced_alfworld_files(self) -> None: + with tempfile.TemporaryDirectory() as temporary_directory: + root = Path(temporary_directory) + lock, prompt, manifest = self._prepare(root) + referenced_game = next((root / "json_2.1.1" / "train").rglob("game.tw-pddl")) + referenced_game.write_text("changed", encoding="utf-8") + + with self.assertRaisesRegex(ValueError, "size does not match"): + self._verify(lock, prompt, manifest) + + def test_model_lock_rehashes_all_nine_checkpoint_files(self) -> None: + with tempfile.TemporaryDirectory() as temporary_directory: + root = Path(temporary_directory) + checkpoint = root / "checkpoint" + checkpoint.mkdir() + files: dict[str, dict[str, object]] = {} + for index in range(9): + name = f"file-{index}.bin" + content = f"content-{index}".encode() + (checkpoint / name).write_bytes(content) + files[name] = { + "bytes": len(content), + "sha256": hashlib.sha256(content).hexdigest(), + } + model_lock = root / "model.lock.json" + model_lock.write_text( + json.dumps( + { + "schema": MODEL_LOCK_SCHEMA, + "repo_id": MODEL_REPO_ID, + "revision": PINNED_MODEL_REVISION, + "files": files, + } + ), + encoding="utf-8", + ) + + verify_model_checkpoint( + model_lock_path=model_lock, + checkpoint_path=checkpoint, + model_revision=PINNED_MODEL_REVISION, + ) + (checkpoint / "file-8.bin").write_bytes(b"changed-8") + with self.assertRaisesRegex(ValueError, "SHA256 does not match"): + verify_model_checkpoint( + model_lock_path=model_lock, + checkpoint_path=checkpoint, + model_revision=PINNED_MODEL_REVISION, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/graphgpo/test_prepare_alfworld.py b/tests/graphgpo/test_prepare_alfworld.py new file mode 100644 index 000000000..c152575b6 --- /dev/null +++ b/tests/graphgpo/test_prepare_alfworld.py @@ -0,0 +1,193 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import json +import tempfile +import unittest +from pathlib import Path + +from examples.graphgpo.prepare_alfworld import ( + PINNED_ALFWORLD_VERSION, + PREPARE_SCHEMA_VERSION, + eligible_task_files, + prepare_artifacts, + prompt_rows_bytes, +) + + +class PrepareAlfWorldTest(unittest.TestCase): + def _write_task( + self, + root: Path, + *, + split: str, + task_dir: str, + task_type: str = "pick_and_place_simple", + solvable: bool = True, + ) -> Path: + path = root / "json_2.1.1" / split / task_dir / "trial_1" / "game.tw-pddl" + path.parent.mkdir(parents=True) + path.write_text(json.dumps({"solvable": solvable}), encoding="utf-8") + path.with_name("traj_data.json").write_text( + json.dumps({"task_type": task_type}), + encoding="utf-8", + ) + return path + + def _write_shared(self, root: Path) -> None: + logic = root / "logic" + logic.mkdir() + (logic / "alfred.pddl").write_text("domain", encoding="utf-8") + (logic / "alfred.twl2").write_text("grammar", encoding="utf-8") + + def test_prepare_writes_stable_manifests_prompt_rows_and_lock(self) -> None: + with tempfile.TemporaryDirectory() as temporary_directory: + root = Path(temporary_directory) + output = root / "prepared" + self._write_shared(root) + self._write_task(root, split="train", task_dir="pick_and_place_simple-B") + self._write_task(root, split="train", task_dir="pick_and_place_simple-A") + self._write_task(root, split="valid_seen", task_dir="look_at_obj_in_light-A") + + first = prepare_artifacts( + data_root=root, + output_dir=output, + split_roots={ + "train": root / "json_2.1.1" / "train", + "eval_in_distribution": root / "json_2.1.1" / "valid_seen", + }, + limits={"eval_in_distribution": 1}, + max_steps=50, + ) + second = prepare_artifacts( + data_root=root, + output_dir=output, + split_roots={ + "train": root / "json_2.1.1" / "train", + "eval_in_distribution": root / "json_2.1.1" / "valid_seen", + }, + limits={"eval_in_distribution": 1}, + max_steps=50, + ) + + self.assertEqual(first, second) + self.assertEqual(first["schema_version"], PREPARE_SCHEMA_VERSION) + self.assertEqual(first["alfworld_version"], PINNED_ALFWORLD_VERSION) + self.assertEqual(first["splits"]["train"]["task_count"], 2) + lock = json.loads((output / "prepare.lock.json").read_text(encoding="utf-8")) + self.assertEqual(lock, first) + + train_rows = [ + json.loads(line) for line in (output / "train.prompts.jsonl").read_text(encoding="utf-8").splitlines() + ] + self.assertEqual(len(train_rows), 2) + self.assertEqual( + [row["metadata"]["manifest_index"] for row in train_rows], + [0, 1], + ) + self.assertTrue( + train_rows[0]["metadata"]["manifest_task_id"] < train_rows[1]["metadata"]["manifest_task_id"] + ) + self.assertEqual( + {row["metadata"]["alfworld_train_eval"] for row in train_rows}, + {"train"}, + ) + self.assertEqual( + { + ( + row["metadata"]["temperature"], + row["metadata"]["top_p"], + row["metadata"]["max_tokens"], + ) + for row in train_rows + }, + {(1.0, 1.0, 512)}, + ) + self.assertTrue(all(row["metadata"]["task_id"].endswith("/game.tw-pddl") for row in train_rows)) + self.assertEqual( + [row["metadata"]["task_id"] for row in train_rows], + [row["metadata"]["manifest_task_id"] for row in train_rows], + ) + self.assertEqual( + first["splits"]["eval_in_distribution"]["sampling"], + {"max_tokens": 512, "temperature": 0.4, "top_p": 1.0}, + ) + + def test_prepare_filters_tasks_the_text_environment_would_skip(self) -> None: + with tempfile.TemporaryDirectory() as temporary_directory: + root = Path(temporary_directory) + self._write_shared(root) + valid = self._write_task( + root, + split="train", + task_dir="pick_and_place_simple-A", + ) + self._write_task( + root, + split="train", + task_dir="pick_and_place_simple-movable", + ) + self._write_task( + root, + split="train", + task_dir="pick_and_place_simple-Sliced", + ) + self._write_task( + root, + split="train", + task_dir="pick_and_place_simple-unsolvable", + solvable=False, + ) + self._write_task( + root, + split="train", + task_dir="pick_and_place_simple-unknown", + task_type="not_supported", + ) + + self.assertEqual( + eligible_task_files(root / "json_2.1.1" / "train"), + [valid], + ) + + def test_prepare_refuses_to_replace_changed_output(self) -> None: + with tempfile.TemporaryDirectory() as temporary_directory: + root = Path(temporary_directory) + output = root / "prepared" + self._write_shared(root) + game = self._write_task( + root, + split="train", + task_dir="pick_and_place_simple-A", + ) + arguments = { + "data_root": root, + "output_dir": output, + "split_roots": {"train": root / "json_2.1.1" / "train"}, + } + prepare_artifacts(**arguments) + game.write_text(json.dumps({"solvable": True, "changed": True}), encoding="utf-8") + + with self.assertRaisesRegex(FileExistsError, "refusing to replace"): + prepare_artifacts(**arguments) + + def test_prompt_rows_reject_empty_manifest(self) -> None: + with tempfile.TemporaryDirectory() as temporary_directory: + root = Path(temporary_directory) + self._write_shared(root) + empty_split = root / "json_2.1.1" / "train" + empty_split.mkdir(parents=True) + with self.assertRaisesRegex(ValueError, "no eligible tasks"): + prepare_artifacts( + data_root=root, + output_dir=root / "prepared", + split_roots={"train": empty_split}, + ) + + with self.assertRaisesRegex(ValueError, "positive integer"): + prompt_rows_bytes(object(), max_steps=0) # type: ignore[arg-type] + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/graphgpo/test_prompt.py b/tests/graphgpo/test_prompt.py new file mode 100644 index 000000000..69cf3495b --- /dev/null +++ b/tests/graphgpo/test_prompt.py @@ -0,0 +1,124 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import unittest + +from examples.graphgpo.prompt import ( + HistoryTurn, + build_prompt, + format_reference_observation, +) +from examples.graphgpo.state import TrackerState + + +class PromptTest(unittest.TestCase): + def test_reference_observation_includes_tracker_context(self) -> None: + tracker = TrackerState.from_mapping( + { + "location": "fridge 1", + "holding": "apple 1", + "history_items": { + "apple 1": { + "heated": True, + "cooled": False, + }, + "plate 1": { + "cleaned": True, + }, + }, + "item_location": { + "apple 1": { + "old_location": "cabinet 1", + "new_location": "fridge 1", + }, + "plate 1": { + "old_location": "sink 1", + "new_location": "sink 1", + }, + }, + } + ) + + observation = format_reference_observation("You arrive.", tracker) + + self.assertIn("Location: fridge 1.", observation) + self.assertIn( + "Items in hand (status): apple 1(heated).", + observation, + ) + self.assertIn( + "apple 1(heated) from cabinet 1 move to fridge 1;", + observation, + ) + self.assertIn( + "plate 1(cleaned) from sink 1 move to sink 1;", + observation, + ) + + def test_first_prompt_uses_initial_observation_without_duplicate_task( + self, + ) -> None: + prompt = build_prompt( + task_description="Put the apple in the fridge.", + raw_observation="Your task is to put the apple in the fridge.", + admissible_commands=["help", "go north"], + ) + + self.assertNotIn("Task: Put the apple", prompt) + self.assertIn("Your task is to put the apple", prompt) + self.assertNotIn("'help'", prompt) + self.assertIn("'go north'", prompt) + + def test_later_prompt_shows_only_last_two_history_turns(self) -> None: + history = [ + HistoryTurn("old action", "old observation"), + HistoryTurn("take apple", "You take the apple."), + HistoryTurn("go north", "You enter the hall."), + ] + prompt = build_prompt( + task_description="Put the apple in the fridge.", + raw_observation="The fridge is here.", + admissible_commands=["open fridge", "help", "go south"], + history=history, + step_index=3, + ) + + self.assertNotIn("old action", prompt) + self.assertNotIn("old observation", prompt) + self.assertIn("take apple", prompt) + self.assertIn("go north", prompt) + self.assertIn( + "Your task is to: Put the apple in the fridge.", + prompt, + ) + self.assertIn("already taken 3 step(s)", prompt) + self.assertIn("now at step 4", prompt) + self.assertNotIn("'help'", prompt) + + def test_visible_commands_preserve_reference_order(self) -> None: + prompt = build_prompt( + task_description="task", + raw_observation="obs", + admissible_commands=["z command", "help", "a command"], + ) + + self.assertLess(prompt.index("'z command'"), prompt.index("'a command'")) + + def test_trimmed_history_uses_absolute_step_numbers(self) -> None: + prompt = build_prompt( + task_description="task", + raw_observation="current", + admissible_commands=["look"], + history=[ + HistoryTurn("action nine", "observation nine"), + HistoryTurn("action ten", "observation ten"), + ], + step_index=10, + ) + + self.assertIn("Observation 9", prompt) + self.assertIn("Action 10", prompt) + self.assertIn("now at step 11", prompt) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/graphgpo/test_recipe_packaging.py b/tests/graphgpo/test_recipe_packaging.py new file mode 100644 index 000000000..1b2641957 --- /dev/null +++ b/tests/graphgpo/test_recipe_packaging.py @@ -0,0 +1,133 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import unittest +from pathlib import Path + +import pytest + +from examples.graphgpo.preflight import validate_batch_arithmetic + + +REPOSITORY_ROOT = Path(__file__).resolve().parents[2] +RECIPE_ROOT = REPOSITORY_ROOT / "examples" / "graphgpo" + + +class RecipePackagingTest(unittest.TestCase): + def test_isolated_agent_dependencies_are_pinned(self) -> None: + content = (RECIPE_ROOT / "requirements-alfworld.txt").read_text(encoding="utf-8") + self.assertEqual( + content.splitlines(), + [ + "alfworld==0.4.2", + "httpx==0.28.1", + "openai==2.46.0", + "PyYAML==6.0.3", + ], + ) + + def test_alfworld_config_is_text_only_and_path_portable(self) -> None: + content = (RECIPE_ROOT / "configs" / "alfworld_qwen2_5_1_5b.yaml").read_text(encoding="utf-8") + self.assertIn('type: "AlfredTWEnv"', content) + self.assertIn("$ALFWORLD_DATA/json_2.1.1/train", content) + self.assertIn("$ALFWORLD_DATA/json_2.1.1/valid_seen", content) + self.assertIn("use_cuda: false", content) + self.assertIn('training_method: "dqn"', content) + self.assertIn("rl:", content) + self.assertNotIn('training_method: "dagger"', content) + self.assertIn("num_train_games: -1", content) + self.assertIn("num_eval_games: -1", content) + self.assertNotIn("/mnt/", content) + self.assertNotIn("C:\\", content) + + def test_launcher_keeps_method_switch_and_variable_row_contract(self) -> None: + content = (RECIPE_ROOT / "run_alfworld_qwen2_5_1_5b.sh").read_text(encoding="utf-8") + for method in ("grpo", "gigpo", "graphgpo"): + self.assertIn(method, content) + self.assertIn( + "--agentic-custom-advantage-path examples.graphgpo.custom_advantage.compute_custom_advantage", + content, + ) + self.assertIn( + "--custom-rm-path examples.graphgpo.reward.reward_func", + content, + ) + self.assertIn("--group-rm", content) + self.assertIn("--use-dynamic-batch-size", content) + self.assertIn('--train-env-vars \'{"TORCH_COMPILE_DISABLE":"1"}\'', content) + self.assertIn("--disable-jit-fuser", content) + # The frozen Qwen2.5-1.5B config has tie_word_embeddings=true. + self.assertNotIn("--untie-embeddings-and-output-weights", content) + self.assertIn("RELAX_AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE", content) + self.assertIn( + 'EPISODE_WEIGHTING="${EPISODE_WEIGHTING:-trajectory_once}"', + content, + ) + self.assertIn( + 'export GRAPHGPO_EPISODE_WEIGHTING="${EPISODE_WEIGHTING}"', + content, + ) + self.assertIn('ENABLE_EVAL="${ENABLE_EVAL:-1}"', content) + self.assertIn('if [ "${ENABLE_EVAL}" = "1" ]; then', content) + self.assertIn('echo "enable_eval=${ENABLE_EVAL}"', content) + self.assertIn( + 'SGLANG_MEM_FRACTION_STATIC="${SGLANG_MEM_FRACTION_STATIC:-0.50}"', + content, + ) + self.assertIn( + '--sglang-mem-fraction-static "${SGLANG_MEM_FRACTION_STATIC}"', + content, + ) + self.assertIn( + 'echo "sglang_mem_fraction_static=${SGLANG_MEM_FRACTION_STATIC}"', + content, + ) + self.assertIn("python3 -m examples.graphgpo.preflight", content) + self.assertIn('--model-lock "${MODEL_LOCK}"', content) + self.assertIn('--checkpoint "${HF_CHECKPOINT}"', content) + self.assertIn('--alfworld-data-root "${ALFWORLD_DATA}"', content) + self.assertIn( + "--custom-eval-rollout-log-function-path examples.graphgpo.eval_logger.log_eval_rollout_data", + content, + ) + self.assertIn('if [ "${DRY_RUN:-0}" = "1" ]', content) + self.assertIn("image_verification_scope=external_executor_required", content) + self.assertNotIn("image_digest=${IMAGE_DIGEST}", content) + + def test_launcher_default_batch_arithmetic_forms_one_complete_batch(self) -> None: + content = (RECIPE_ROOT / "run_alfworld_qwen2_5_1_5b.sh").read_text(encoding="utf-8") + self.assertIn('TASK_GROUPS="${TASK_GROUPS:-16}"', content) + self.assertIn('GROUP_SIZE="${GROUP_SIZE:-8}"', content) + self.assertIn('GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-128}"', content) + validate_batch_arithmetic( + task_groups=16, + group_size=8, + global_batch_size=128, + ) + with pytest.raises(ValueError, match="must be divisible"): + validate_batch_arithmetic( + task_groups=16, + group_size=8, + global_batch_size=256, + ) + + def test_launcher_temporarily_disables_nounset_while_sourcing_local_entrypoint(self) -> None: + content = (RECIPE_ROOT / "run_alfworld_qwen2_5_1_5b.sh").read_text(encoding="utf-8") + disable_index = content.index(" set +u\n") + source_index = content.index(' source "${PROJECT_ROOT}/scripts/entrypoint/local.sh"\n') + restore_index = content.index(" set -u\n", source_index) + self.assertLess(disable_index, source_index) + self.assertLess(source_index, restore_index) + + def test_managed_agent_entrypoint_uses_reserved_runtime_paths(self) -> None: + content = (RECIPE_ROOT / "run_agent_app.sh").read_text(encoding="utf-8") + self.assertIn("RELAX_INPUT_JSON", content) + self.assertIn("RELAX_OUTPUT_JSON", content) + self.assertIn("RELAX_BASE_URL", content) + self.assertIn("ALFWORLD_PYTHON", content) + self.assertIn("examples.graphgpo.rollout_agent", content) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/graphgpo/test_reproducibility.py b/tests/graphgpo/test_reproducibility.py new file mode 100644 index 000000000..f7bb795e5 --- /dev/null +++ b/tests/graphgpo/test_reproducibility.py @@ -0,0 +1,158 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +import pytest + +from examples.graphgpo.reproducibility import ( + DECLARED_VERSION_FIELDS, + build_seed_mapping, + build_test_report, + build_version_inventory, + main, + parse_junit_report, +) + + +BASELINE_COMMIT = "8a54679e971087566bfda939a5f75649e07fb861" + + +def _write_junit(path: Path, cases: str) -> None: + path.write_text( + f'\n{cases}\n', + encoding="utf-8", + ) + + +def test_registered_seed_mapping_is_a_plan_for_all_twelve_runs(): + mapping = build_seed_mapping() + + assert mapping["mapping_scope"] == "registered_plan_not_execution_evidence" + runs = mapping["runs"] + assert len(runs) == 12 + assert len({run["run_id"] for run in runs}) == 12 + assert {run["condition"] for run in runs} == { + "grpo", + "gigpo", + "graphgpo", + "graph_only", + } + for condition in {run["condition"] for run in runs}: + assert {run["registered_seed"] for run in runs if run["condition"] == condition} == {0, 1, 2} + assert all(run["execution_status"] == "planned" for run in runs) + assert all(run["execution_evidence"] is None for run in runs) + assert all(run["launcher_environment"]["EPISODE_WEIGHTING"] == "trajectory_once" for run in runs) + graph_only = next(run for run in runs if run["condition"] == "graph_only") + assert graph_only["launcher_environment"] == { + "METHOD": "graphgpo", + "BETA": "1", + "BETA_EPISODE": "0", + "EPISODE_WEIGHTING": "trajectory_once", + "SEED": "0", + } + + +def test_junit_counts_are_parsed_but_plain_log_outcome_is_not_inferred(tmp_path): + junit = tmp_path / "pytest.xml" + _write_junit( + junit, + """ + + + bad + boom + """, + ) + log = tmp_path / "pre-commit.log" + log.write_text("Passed\n", encoding="utf-8") + + summary = parse_junit_report(junit) + assert summary == { + "tests": 4, + "passed": 1, + "failures": 1, + "errors": 1, + "skipped": 1, + "duration_seconds": pytest.approx(1.0), + "outcome": "failed", + } + report = build_test_report({"pytest": junit}, {"pre_commit": log}) + assert report["overall_outcome_from_junit_only"] == "failed" + assert report["complete_text_logs"][0]["outcome"] == "not_inferred_from_unstructured_log" + assert report["complete_text_logs"][0]["artifact"]["sha256"] == hashlib.sha256(log.read_bytes()).hexdigest() + + +def test_cli_writes_content_addressed_bundle_without_fabricating_versions(tmp_path): + evidence = tmp_path / "expanded_command.sh" + evidence.write_text("python relax/entrypoints/train.py --seed 0\n", encoding="utf-8") + junit = tmp_path / "pytest.xml" + _write_junit( + junit, + '', + ) + log = tmp_path / "pre-commit.log" + log.write_text("ruff................................Passed\n", encoding="utf-8") + output_dir = tmp_path / "bundle" + + argv = [ + "--output-dir", + str(output_dir), + "--artifact", + f"expanded_command={evidence}", + "--junit", + f"pytest_graphgpo={junit}", + "--log", + f"pre_commit={log}", + "--version", + f"relax_baseline_commit={BASELINE_COMMIT}", + ] + assert main(argv) == 0 + assert main(argv) == 0 + + seed_path = output_dir / "seed_mapping.json" + test_path = output_dir / "test_report.json" + manifest_path = output_dir / "reproducibility.manifest.json" + assert seed_path.is_file() + assert test_path.is_file() + assert manifest_path.is_file() + + test_report = json.loads(test_path.read_text(encoding="utf-8")) + assert test_report["overall_outcome_from_junit_only"] == "passed" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + declarations = manifest["versions"]["declared"] + assert set(declarations) == set(DECLARED_VERSION_FIELDS) + assert declarations["relax_baseline_commit"] == { + "value": BASELINE_COMMIT, + "source": "operator_supplied", + "verification": "not_verified_by_generator", + } + assert declarations["candidate_commit"] == { + "value": None, + "source": None, + "verification": "not_claimed", + } + assert manifest["input_artifacts"] == [ + { + "label": "expanded_command", + "file_name": evidence.name, + "sha256": hashlib.sha256(evidence.read_bytes()).hexdigest(), + "size_bytes": evidence.stat().st_size, + } + ] + generated = {entry["label"]: entry for entry in manifest["generated_artifacts"]} + assert generated["seed_mapping"]["sha256"] == hashlib.sha256(seed_path.read_bytes()).hexdigest() + assert generated["test_report"]["sha256"] == hashlib.sha256(test_path.read_bytes()).hexdigest() + + serialized_bundle = "".join(path.read_text(encoding="utf-8") for path in (seed_path, test_path, manifest_path)) + assert str(tmp_path) not in serialized_bundle + + +def test_version_inventory_rejects_values_that_look_verified_but_are_ambiguous(): + with pytest.raises(ValueError, match="full lowercase 40-character commit"): + build_version_inventory({"candidate_commit": "deadbeef"}) + with pytest.raises(ValueError, match="pinned by a sha256 digest"): + build_version_inventory({"container_image_digest": "relax:latest"}) diff --git a/tests/graphgpo/test_reward.py b/tests/graphgpo/test_reward.py new file mode 100644 index 000000000..bcf587842 --- /dev/null +++ b/tests/graphgpo/test_reward.py @@ -0,0 +1,38 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import asyncio +import math +from types import SimpleNamespace + +import pytest + +from examples.graphgpo.reward import reward_func + + +def test_reward_func_forwards_environment_rewards_in_order() -> None: + samples = [ + SimpleNamespace(reward=10), + SimpleNamespace(reward=-0.2), + SimpleNamespace(reward=0.0), + ] + + assert asyncio.run(reward_func(None, samples)) == [10.0, -0.2, 0.0] + + +@pytest.mark.parametrize( + "samples", + [ + "not-a-group", + [], + [SimpleNamespace()], + [SimpleNamespace(reward=True)], + [SimpleNamespace(reward=math.nan)], + [SimpleNamespace(reward=math.inf)], + [SimpleNamespace(reward={"score": 1.0})], + ], +) +def test_reward_func_rejects_invalid_group_rewards(samples: object) -> None: + with pytest.raises((TypeError, ValueError)): + asyncio.run(reward_func(None, samples)) # type: ignore[arg-type] diff --git a/tests/graphgpo/test_rollout_agent.py b/tests/graphgpo/test_rollout_agent.py new file mode 100644 index 000000000..5fa31cfc4 --- /dev/null +++ b/tests/graphgpo/test_rollout_agent.py @@ -0,0 +1,570 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest + +import examples.graphgpo.rollout_agent as rollout_agent +from examples.graphgpo.alfworld_env import AlfWorldSnapshot +from examples.graphgpo.graph_credit import SUCCESS +from examples.graphgpo.rollout_agent import ( + extract_task_description, + parse_args, + run_episode, + run_managed_session, + stable_group_seed, + stable_row_id, + write_session_output, +) +from examples.graphgpo.state import TrackerState + + +def _snapshot( + observation: str, + commands: tuple[str, ...], + *, + done: bool = False, + won: bool = False, + gamefile: str | None = "games/task-alpha/game.tw-pddl", +) -> AlfWorldSnapshot: + return AlfWorldSnapshot( + raw_observation=observation, + admissible_commands=commands, + won=won, + done=done, + gamefile=gamefile, + tracker=TrackerState.from_mapping( + { + "location": "middle of a room", + "holding": "nothing", + "history_items": {}, + "item_location": {}, + } + ), + ) + + +class FakeEpisodeEnv: + def __init__( + self, + initial: AlfWorldSnapshot, + transitions: list[AlfWorldSnapshot], + ) -> None: + self.initial = initial + self.transitions = list(transitions) + self.actions: list[str] = [] + self.closed = False + + def reset(self) -> AlfWorldSnapshot: + return self.initial + + def step(self, action: str) -> AlfWorldSnapshot: + self.actions.append(action) + return self.transitions.pop(0) + + def close(self) -> None: + self.closed = True + + +class FakeChatClient: + def __init__(self, responses: list[str]) -> None: + self.responses = list(responses) + self.calls: list[list[dict[str, Any]]] = [] + + async def complete(self, *, messages: list[dict[str, Any]]) -> str: + self.calls.append(messages) + return self.responses.pop(0) + + +def _valid_response(action: str) -> str: + return f"reason{action}" + + +def test_stable_group_seed_is_group_shared_and_deterministic() -> None: + assert stable_group_seed("group-a") == stable_group_seed("group-a") + assert stable_group_seed("group-a") != stable_group_seed("group-b") + assert 0 <= stable_group_seed("group-a") < 2**31 + + +def test_stable_row_id_binds_task_group_trajectory_and_turn() -> None: + row_id = stable_row_id( + task_id="task-a", + rollout_group_id="group-a", + trajectory_id="trajectory-a", + turn_index=0, + ) + assert row_id == stable_row_id( + task_id="task-a", + rollout_group_id="group-a", + trajectory_id="trajectory-a", + turn_index=0, + ) + assert row_id != stable_row_id( + task_id="task-a", + rollout_group_id="group-a", + trajectory_id="trajectory-a", + turn_index=1, + ) + + +def test_extract_task_description_is_strict() -> None: + assert extract_task_description("Intro. Your task is to: cool the mug") == "cool the mug" + with pytest.raises(ValueError, match="not found"): + extract_task_description("No task marker") + with pytest.raises(ValueError, match="empty"): + extract_task_description("Your task is to: ") + + +def test_parse_args_uses_managed_runtime_environment_paths() -> None: + args = parse_args( + [], + environ={ + "RELAX_INPUT_JSON": "/tmp/session_input.json", + "RELAX_OUTPUT_JSON": "/tmp/session_output.json", + }, + ) + assert args.input_json == "/tmp/session_input.json" + assert args.output_json == "/tmp/session_output.json" + + explicit = parse_args( + [ + "--input-json", + "explicit-input.json", + "--output-json", + "explicit-output.json", + ], + environ={ + "RELAX_INPUT_JSON": "ignored-input.json", + "RELAX_OUTPUT_JSON": "ignored-output.json", + }, + ) + assert explicit.input_json == "explicit-input.json" + assert explicit.output_json == "explicit-output.json" + + +@pytest.mark.asyncio +async def test_run_episode_exports_fresh_per_turn_messages_and_last_two_history() -> None: + env = FakeEpisodeEnv( + _snapshot( + "Intro. Your task is to: inspect the room", + ("action zero", "help"), + ), + [ + _snapshot("observation one", ("action one",)), + _snapshot("observation two", ("action two",)), + _snapshot("observation three", ("action three",)), + _snapshot("observation four", ("look",), done=True), + ], + ) + client = FakeChatClient( + [ + _valid_response("action zero"), + _valid_response("action one"), + _valid_response("action two"), + _valid_response("action three"), + ] + ) + + records = await run_episode( + chat_client=client, + env=env, # type: ignore[arg-type] + trajectory_id="session-1", + ) + + assert len(records) == 4 + assert all(len(call) == 1 for call in client.calls) + assert all(call[0]["role"] == "user" for call in client.calls) + assert all(len(record["messages"]) == 2 for record in records) + assert all([message["role"] for message in record["messages"]] == ["user", "assistant"] for record in records) + fourth_prompt = client.calls[3][0]["content"] + assert "observation one" not in fourth_prompt + assert "observation two" in fourth_prompt + assert "observation three" in fourth_prompt + assert "'help'" not in client.calls[0][0]["content"] + assert env.closed is True + + +@pytest.mark.asyncio +async def test_parser_valid_action_outside_commands_stays_valid_and_is_executed() -> None: + env = FakeEpisodeEnv( + _snapshot( + "Intro. Your task is to: look around", + ("look",), + ), + [_snapshot("Nothing happens.", ("look",), done=True)], + ) + client = FakeChatClient([_valid_response("dance")]) + + records = await run_episode( + chat_client=client, + env=env, # type: ignore[arg-type] + trajectory_id="session-outside-commands", + ) + + assert env.actions == ["dance"] + assert records[0]["metadata"]["is_action_valid"] is True + assert records[0]["metadata"]["episode_return"] == 0.0 + + +@pytest.mark.asyncio +async def test_run_episode_success_backfills_return_and_success_sink() -> None: + env = FakeEpisodeEnv( + _snapshot( + "Intro. Your task is to: take the apple", + ("look",), + ), + [ + _snapshot("Nothing happens.", ("take apple",)), + _snapshot("You win.", ("look",), won=True, done=True), + ], + ) + client = FakeChatClient( + [ + "malformed response", + _valid_response("take apple"), + ] + ) + + records = await run_episode( + chat_client=client, + env=env, # type: ignore[arg-type] + trajectory_id="session-success", + ) + + assert [record["name"] for record in records] == ["turn_000", "turn_001"] + assert env.actions[0] == "malformed response"[-30:] + assert all(record["metadata"]["success"] is True for record in records) + assert all(record["metadata"]["episode_return"] == 9.9 for record in records) + assert all(record["reward"] == 9.9 for record in records) + assert records[0]["metadata"]["is_action_valid"] is False + assert records[0]["metadata"]["terminal"] is False + assert records[-1]["metadata"]["next_state_key"] == SUCCESS + assert records[-1]["metadata"]["terminal"] is True + assert records[-1]["metadata"]["truncated"] is False + assert records[-1]["metadata"]["task_id"] == ("games/task-alpha/game.tw-pddl") + + +@pytest.mark.asyncio +async def test_run_episode_marks_unsuccessful_done_and_max_steps() -> None: + done_env = FakeEpisodeEnv( + _snapshot("Your task is to: wait", ("look",)), + [_snapshot("Episode ended.", ("look",), done=True)], + ) + done_records = await run_episode( + chat_client=FakeChatClient([_valid_response("look")]), + env=done_env, # type: ignore[arg-type] + trajectory_id="done", + ) + assert done_records[-1]["metadata"]["terminal"] is True + assert done_records[-1]["metadata"]["truncated"] is False + assert done_records[-1]["metadata"]["success"] is False + + truncated_env = FakeEpisodeEnv( + _snapshot("Your task is to: wait", ("look",)), + [ + _snapshot("Still waiting one.", ("look",)), + _snapshot("Still waiting two.", ("look",)), + ], + ) + truncated_records = await run_episode( + chat_client=FakeChatClient([_valid_response("look"), _valid_response("look")]), + env=truncated_env, # type: ignore[arg-type] + trajectory_id="truncated", + max_steps=2, + ) + assert len(truncated_records) == 2 + assert truncated_records[-1]["metadata"]["terminal"] is False + assert truncated_records[-1]["metadata"]["truncated"] is True + + horizon_done_env = FakeEpisodeEnv( + _snapshot("Your task is to: wait", ("look",)), + [ + _snapshot("Still waiting one.", ("look",)), + _snapshot("Time limit reached.", ("look",), done=True), + ], + ) + horizon_done_records = await run_episode( + chat_client=FakeChatClient([_valid_response("look"), _valid_response("look")]), + env=horizon_done_env, # type: ignore[arg-type] + trajectory_id="horizon-done", + max_steps=2, + ) + assert horizon_done_records[-1]["metadata"]["terminal"] is False + assert horizon_done_records[-1]["metadata"]["truncated"] is True + + +@pytest.mark.asyncio +async def test_run_episode_uses_metadata_task_id_only_without_gamefile() -> None: + env = FakeEpisodeEnv( + _snapshot( + "Your task is to: wait", + ("look",), + gamefile=None, + ), + [_snapshot("Done.", ("look",), done=True, gamefile=None)], + ) + records = await run_episode( + chat_client=FakeChatClient([_valid_response("look")]), + env=env, # type: ignore[arg-type] + trajectory_id="session-fallback-task", + fallback_task_id="public-task-id", + ) + assert records[0]["metadata"]["task_id"] == "public-task-id" + + +@pytest.mark.asyncio +async def test_run_episode_uses_public_relative_task_id_for_absolute_gamefile() -> None: + gamefile = "/srv/task37/alfworld/json_2.1.1/train/pick_and_place_simple-Apple-None-Fridge-1/trial-1/game.tw-pddl" + env = FakeEpisodeEnv( + _snapshot("Your task is to: wait", ("look",), gamefile=gamefile), + [_snapshot("Done.", ("look",), done=True, gamefile=gamefile)], + ) + records = await run_episode( + chat_client=FakeChatClient([_valid_response("look")]), + env=env, # type: ignore[arg-type] + trajectory_id="session-relative-task", + environ={"ALFWORLD_DATA": "/srv/task37/alfworld"}, + ) + assert records[0]["metadata"]["task_id"] == ( + "json_2.1.1/train/pick_and_place_simple-Apple-None-Fridge-1/trial-1/game.tw-pddl" + ) + + +@pytest.mark.asyncio +async def test_run_episode_rejects_gamefile_outside_data_root() -> None: + env = FakeEpisodeEnv( + _snapshot( + "Your task is to: wait", + ("look",), + gamefile="/other/data/game.tw-pddl", + ), + [], + ) + with pytest.raises(ValueError, match="outside ALFWORLD_DATA"): + await run_episode( + chat_client=FakeChatClient([]), + env=env, # type: ignore[arg-type] + trajectory_id="session-outside-root", + environ={"ALFWORLD_DATA": "/srv/task37/alfworld"}, + ) + assert env.closed is True + + +@pytest.mark.asyncio +async def test_run_episode_rejects_declared_task_mismatch() -> None: + env = FakeEpisodeEnv( + _snapshot( + "Your task is to: wait", + ("look",), + gamefile="/srv/task37/alfworld/json_2.1.1/train/a/game.tw-pddl", + ), + [], + ) + with pytest.raises(ValueError, match="declared task_id"): + await run_episode( + chat_client=FakeChatClient([]), + env=env, # type: ignore[arg-type] + trajectory_id="session-task-mismatch", + fallback_task_id="json_2.1.1/train/b/game.tw-pddl", + environ={"ALFWORLD_DATA": "/srv/task37/alfworld"}, + ) + assert env.closed is True + + +@pytest.mark.asyncio +async def test_run_episode_closes_env_on_errors_and_missing_ids() -> None: + class RaisingClient: + async def complete(self, *, messages: list[dict[str, Any]]) -> str: + raise RuntimeError("chat failed") + + env = FakeEpisodeEnv( + _snapshot("Your task is to: wait", ("look",)), + [], + ) + with pytest.raises(RuntimeError, match="chat failed"): + await run_episode( + chat_client=RaisingClient(), + env=env, # type: ignore[arg-type] + trajectory_id="session-error", + ) + assert env.closed is True + + missing_id_env = FakeEpisodeEnv( + _snapshot("Your task is to: wait", ("look",)), + [], + ) + with pytest.raises(ValueError, match="RELAX_SESSION_ID"): + await run_episode( + chat_client=FakeChatClient([]), + env=missing_id_env, # type: ignore[arg-type] + environ={}, + ) + assert missing_id_env.closed is True + + +@pytest.mark.asyncio +async def test_run_managed_session_uses_relax_ids() -> None: + env = FakeEpisodeEnv( + _snapshot("Your task is to: wait", ("look",)), + [_snapshot("Done.", ("look",), done=True)], + ) + records = await run_managed_session( + session_input={"messages": [], "metadata": {}}, + chat_client=FakeChatClient([_valid_response("look")]), + env=env, # type: ignore[arg-type] + environ={ + "RELAX_GROUP_ID": "shared-group", + "RELAX_SESSION_ID": "slot-session", + }, + ) + metadata = records[0]["metadata"] + assert metadata["trajectory_id"] == "slot-session" + assert metadata["rollout_group_id"] == "shared-group" + assert metadata["turn_id"] == "turn_000" + assert metadata["row_id"] == stable_row_id( + task_id=metadata["task_id"], + rollout_group_id="shared-group", + trajectory_id="slot-session", + turn_index=0, + ) + + +@pytest.mark.asyncio +async def test_run_managed_session_uses_relax_chat_environment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + + class FakeDefaultClient(FakeChatClient): + def __init__(self, **kwargs: object) -> None: + super().__init__([_valid_response("look")]) + captured.update(kwargs) + self.was_closed = False + + async def close(self) -> None: + self.was_closed = True + + monkeypatch.setattr( + rollout_agent, + "OpenAICompatibleChatClient", + FakeDefaultClient, + ) + env = FakeEpisodeEnv( + _snapshot("Your task is to: wait", ("look",)), + [_snapshot("Done.", ("look",), done=True)], + ) + await run_managed_session( + session_input={ + "messages": [], + "metadata": { + "temperature": 0.8, + "top_p": 0.95, + "max_tokens": 512, + }, + }, + env=env, # type: ignore[arg-type] + environ={ + "RELAX_GROUP_ID": "shared-group", + "RELAX_SESSION_ID": "slot-session", + "RELAX_BASE_URL": "http://relax-chat/v1", + }, + ) + assert captured["base_url"] == "http://relax-chat/v1" + assert captured["api_key"] == "slot-session" + assert captured["temperature"] == 0.8 + assert captured["top_p"] == 0.95 + assert captured["max_tokens"] == 512 + + +@pytest.mark.asyncio +async def test_run_managed_session_allows_explicit_test_overrides() -> None: + env = FakeEpisodeEnv( + _snapshot( + "Your task is to: wait", + ("look",), + gamefile=None, + ), + [_snapshot("Done.", ("look",), done=True, gamefile=None)], + ) + records = await run_managed_session( + session_input={"messages": [], "metadata": {}}, + chat_client=FakeChatClient([_valid_response("look")]), + env=env, # type: ignore[arg-type] + group_id="test-group", + trajectory_id="test-trajectory", + task_id="test-task", + environ={}, + ) + assert records[0]["metadata"]["task_id"] == "test-task" + assert records[0]["metadata"]["trajectory_id"] == "test-trajectory" + assert records[0]["metadata"]["rollout_group_id"] == "test-group" + + +@pytest.mark.asyncio +async def test_same_group_seed_selects_same_fake_task() -> None: + task_by_seed: dict[int, str] = {} + + def make_env(seed: int) -> FakeEpisodeEnv: + task = task_by_seed.setdefault(seed, f"task-{seed}") + return FakeEpisodeEnv( + _snapshot( + "Your task is to: wait", + ("look",), + gamefile=f"games/{task}/game.tw-pddl", + ), + [ + _snapshot( + "Done.", + ("look",), + done=True, + gamefile=f"games/{task}/game.tw-pddl", + ) + ], + ) + + seed_a = stable_group_seed("one-group") + seed_b = stable_group_seed("one-group") + records_a = await run_episode( + chat_client=FakeChatClient([_valid_response("look")]), + env=make_env(seed_a), # type: ignore[arg-type] + trajectory_id="slot-a", + ) + records_b = await run_episode( + chat_client=FakeChatClient([_valid_response("look")]), + env=make_env(seed_b), # type: ignore[arg-type] + trajectory_id="slot-b", + ) + assert records_a[0]["metadata"]["task_id"] == records_b[0]["metadata"]["task_id"] + + +@pytest.mark.asyncio +async def test_jsonl_output_is_accepted_by_session_output(tmp_path: Path) -> None: + env = FakeEpisodeEnv( + _snapshot("Your task is to: wait", ("look",)), + [_snapshot("Done.", ("look",), done=True)], + ) + records = await run_episode( + chat_client=FakeChatClient([_valid_response("look")]), + env=env, # type: ignore[arg-type] + trajectory_id="session-jsonl", + ) + output_path = tmp_path / "session_output.json" + write_session_output(output_path, records) + + raw_text = output_path.read_text(encoding="utf-8") + assert raw_text.startswith("{") + assert not raw_text.lstrip().startswith("[") + parsed_records = [json.loads(line) for line in raw_text.splitlines() if line.strip()] + assert len(parsed_records) == 1 + + try: + from relax.agentic.pipeline.runtime import SessionOutput + except ImportError: + pytest.skip("Relax runtime dependencies are unavailable") + output = SessionOutput.from_records(parsed_records) + assert len(output.records) == 1 + assert output.records[0]["name"] == "turn_000" diff --git a/tests/graphgpo/test_state.py b/tests/graphgpo/test_state.py new file mode 100644 index 000000000..6a03b5d72 --- /dev/null +++ b/tests/graphgpo/test_state.py @@ -0,0 +1,169 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import unittest + +from examples.graphgpo.state import ( + ALFWORLD_TRACKER_FIELDS, + TrackerState, + reference_anchor_v1, + state_key_v1, + update_tracker, +) + + +class StateAnchorTest(unittest.TestCase): + @staticmethod + def _tracker(**updates: object) -> TrackerState: + values: dict[str, object] = { + "holding": "nothing", + "location": "kitchen", + "history_items": {}, + "item_location": {}, + } + values.update(updates) + return TrackerState.from_mapping(values) + + def test_anchor_contains_raw_observation_tracker_and_all_sorted_commands( + self, + ) -> None: + tracker = self._tracker(holding="apple") + anchor = reference_anchor_v1( + "You see a table.\nRaw line.", + tracker, + ["take apple", "help", "go north"], + ) + + self.assertIn("You see a table.\nRaw line.", anchor) + self.assertIn("Items in hand (status): apple(unprocessed)", anchor) + self.assertIn("Location: kitchen", anchor) + self.assertLess(anchor.index("go north"), anchor.index("help")) + self.assertLess(anchor.index("help"), anchor.index("take apple")) + + def test_nothing_happens_does_not_update_tracker(self) -> None: + tracker = self._tracker() + updated = update_tracker( + tracker, + raw_observation="Nothing happens.", + updates={"holding": "apple", "location": "hall"}, + ) + + self.assertIs(updated, tracker) + + def test_successful_observation_applies_structured_updates(self) -> None: + tracker = self._tracker() + updated = update_tracker( + tracker, + raw_observation="You pick up the apple.", + updates={"holding": "apple"}, + ) + + self.assertEqual( + updated.to_mapping(), + { + "location": "kitchen", + "holding": "apple", + "history_items": {}, + "item_location": {}, + }, + ) + + def test_state_key_is_stable_across_nested_mapping_and_command_order(self) -> None: + first = state_key_v1( + "obs", + self._tracker( + history_items={"mug": {"heated": True, "cooled": False}}, + item_location={ + "mug": { + "old_location": "kitchen", + "new_location": "hall", + } + }, + ), + ["help", "go north"], + ) + second = state_key_v1( + "obs", + TrackerState.from_mapping( + { + "item_location": { + "mug": { + "new_location": "hall", + "old_location": "kitchen", + } + }, + "history_items": {"mug": {"cooled": False, "heated": True}}, + "holding": "nothing", + "location": "kitchen", + } + ), + ["go north", "help"], + ) + + self.assertEqual(first, second) + self.assertEqual(len(first), 64) + + def test_tracker_fields_are_an_exact_whitelist(self) -> None: + self.assertEqual( + ALFWORLD_TRACKER_FIELDS, + ("location", "holding", "history_items", "item_location"), + ) + + with self.assertRaisesRegex(ValueError, "missing=.*item_location"): + TrackerState.from_mapping( + { + "location": "kitchen", + "holding": "nothing", + "history_items": {}, + } + ) + with self.assertRaisesRegex(ValueError, "unexpected=.*debug_history"): + TrackerState.from_mapping( + { + "location": "kitchen", + "holding": "nothing", + "history_items": {}, + "item_location": {}, + "debug_history": [], + } + ) + with self.assertRaisesRegex(ValueError, "unexpected fields"): + update_tracker( + self._tracker(), + raw_observation="You wait.", + updates={"debug_history": []}, + ) + + def test_raw_observation_changes_state_key(self) -> None: + tracker = self._tracker() + + self.assertNotEqual( + state_key_v1("first", tracker, ["help"]), + state_key_v1("second", tracker, ["help"]), + ) + + def test_hidden_internal_history_does_not_split_reference_state(self) -> None: + base = self._tracker() + with_hidden_history = self._tracker( + history_items={"mug": {"cleaned": False}}, + item_location={ + "mug": { + "old_location": "kitchen", + "new_location": "kitchen", + } + }, + ) + + self.assertEqual( + state_key_v1("obs", base, ["look"]), + state_key_v1("obs", with_hidden_history, ["look"]), + ) + + def test_displayed_tracker_difference_splits_reference_state(self) -> None: + self.assertNotEqual( + state_key_v1("obs", self._tracker(location="kitchen"), ["look"]), + state_key_v1("obs", self._tracker(location="hall"), ["look"]), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_agentic_variable_row_transfer.py b/tests/test_agentic_variable_row_transfer.py new file mode 100644 index 000000000..63e4da5c9 --- /dev/null +++ b/tests/test_agentic_variable_row_transfer.py @@ -0,0 +1,548 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from __future__ import annotations + +import asyncio +import copy +import random +from types import SimpleNamespace + +import pytest + +from relax.agentic.pipeline.transfer import ( + AGENTIC_ROW_IDENTITY_KEY, + AGENTIC_ROW_IDENTITY_TAG_KEY, + AGENTIC_ROW_IDENTITY_TAGS_FIELD, + AGENTIC_VARIABLE_ROW_PADDING_KEY, + TransferDomain, + _variable_row_padding_sample, +) +from relax.utils.types import Sample + + +MAX_ROWS_ENV = "RELAX_AGENTIC_MAX_EXPORTED_ROWS_PER_SAMPLE" + + +@pytest.fixture(autouse=True) +def _enable_variable_rows(monkeypatch): + monkeypatch.setenv(MAX_ROWS_ENV, "50") + + +class _FakeRolloutBatch(dict): + def numel(self): + return len(self["total_lengths"]) + + +class _RecordingDataClient: + def __init__(self, *, delay_first: bool = False): + self.calls: list[dict] = [] + self.events: list[str] = [] + self.delay_first = delay_first + + async def async_put(self, *, data, partition_id, custom_meta, is_last): + call_number = len(self.calls) + self.events.append(f"start-{call_number}") + if self.delay_first and call_number == 0: + await asyncio.sleep(0.01) + self.calls.append( + { + "data": data, + "partition_id": partition_id, + "custom_meta": custom_meta, + "is_last": is_last, + } + ) + self.events.append(f"end-{call_number}") + + +def _args(**overrides): + values = { + "rollout_batch_size": 2, + "over_sampling_batch_size": 2, + "n_samples_per_prompt": 1, + "colocate": False, + "global_batch_size": 4, + "num_iters_per_train_update": 4, + "group_rm": True, + "agentic_custom_advantage_path": "examples.graphgpo.advantage.compute", + "use_dynamic_batch_size": True, + "fully_async": False, + "hybrid": False, + "reward_key": None, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_variable_row_transfer_is_explicitly_opt_in(monkeypatch): + monkeypatch.delenv(MAX_ROWS_ENV) + + transfer = TransferDomain(args=_args(), data_system_client=None) + + assert transfer._variable_row_mode is False + + +def test_variable_row_transfer_rejects_export_above_partition_bound(monkeypatch): + monkeypatch.setenv(MAX_ROWS_ENV, "1") + transfer = TransferDomain(args=_args(), data_system_client=None) + + with pytest.raises(RuntimeError, match="exceeded the configured partition bound"): + transfer._prepare_variable_row_groups( + groups=[_group(0, 3)], + partition_rollout_id=0, + is_last=False, + ) + + assert transfer._partition_actual_rows == {} + + +def _group(group_index: int, rows: int) -> list[Sample]: + group: list[Sample] = [] + trajectory_id = f"trajectory-{group_index}" + for turn_index in range(rows): + sample = Sample( + group_index=group_index, + index=group_index, + session_id=trajectory_id, + tokens=[group_index + 1, turn_index + 11], + response_length=1, + reward=1.0, + custom_advantage=1.0, + loss_mask=[1], + rollout_log_probs=[-0.1], + weight_versions=["7"], + metadata={ + "row_id": f"row-{group_index}-{turn_index}", + "rollout_group_id": str(group_index), + "policy_version": "7", + "task_id": f"task-{group_index}", + "trajectory_id": trajectory_id, + "turn_id": f"turn_{turn_index:03d}", + "turn_index": turn_index, + "terminal": turn_index == rows - 1, + "truncated": False, + }, + ) + setattr(sample, "_agentic_export_name", f"turn_{turn_index:03d}") + group.append(sample) + return group + + +def _multi_trajectory_group( + *, + group_index: int, + task_id: str, + trajectory_lengths: list[int], + policy_version: str = "11", +) -> list[Sample]: + group: list[Sample] = [] + for slot_index, trajectory_length in enumerate(trajectory_lengths): + trajectory_id = f"trajectory-{group_index}-{slot_index}" + sample_index = group_index * 100 + slot_index + for turn_index in range(trajectory_length): + token_id = group_index * 10_000 + slot_index * 100 + turn_index + sample = Sample( + group_index=group_index, + index=sample_index, + session_id=trajectory_id, + tokens=[31, token_id], + response_length=1, + reward=1.0, + custom_advantage=1.0, + loss_mask=[1], + rollout_log_probs=[-0.25], + weight_versions=[policy_version], + metadata={ + "row_id": f"row-{group_index}-{slot_index}-{turn_index}", + "rollout_group_id": str(group_index), + "policy_version": policy_version, + "task_id": task_id, + "trajectory_id": trajectory_id, + "turn_id": f"turn_{turn_index:03d}", + "turn_index": turn_index, + "terminal": turn_index == trajectory_length - 1, + "truncated": False, + }, + ) + setattr(sample, "_agentic_export_name", f"turn_{turn_index:03d}") + group.append(sample) + random.Random(group_index).shuffle(group) + return group + + +def _install_fake_converter(monkeypatch, converted_samples: list[list[Sample]]) -> None: + def _convert(_args, samples): + converted_samples.append(list(samples)) + return _FakeRolloutBatch( + tokens=[list(sample.tokens) for sample in samples], + total_lengths=[len(sample.tokens) for sample in samples], + response_lengths=[sample.response_length for sample in samples], + loss_masks=[ + [0] * sample.response_length if sample.remove_sample else list(sample.loss_mask) for sample in samples + ], + rollout_log_probs=[list(sample.rollout_log_probs or []) for sample in samples], + ) + + monkeypatch.setattr("relax.utils.utils.convert_samples_to_train_data", _convert) + + +async def _run_partition( + transfer: TransferDomain, + *, + rollout_id: int, + previous_partition_quota: int, + current_partition_quota: int, + groups: list[list[Sample]], +) -> None: + transfer.rebind_step(rollout_id=rollout_id) + transfer.configure_transfer_quota( + previous_partition_quota=previous_partition_quota, + current_partition_quota=current_partition_quota, + ) + transfer.enqueue_ready_groups(groups) + await transfer.drain_ready_group_payloads() + await transfer.wait_for_pending_transfers() + + +async def _run_partition_in_batches( + transfer: TransferDomain, + *, + rollout_id: int, + current_partition_quota: int, + batches: list[list[list[Sample]]], +) -> None: + transfer.rebind_step(rollout_id=rollout_id) + transfer.configure_transfer_quota( + previous_partition_quota=0, + current_partition_quota=current_partition_quota, + ) + for groups in batches: + transfer.enqueue_ready_groups(groups) + await transfer.drain_ready_group_payloads() + await asyncio.sleep(0) + await transfer.wait_for_pending_transfers() + + +def test_variable_row_transfer_five_rows_adds_three_final_padding_rows_and_serializes_puts(monkeypatch): + converted_samples: list[list[Sample]] = [] + _install_fake_converter(monkeypatch, converted_samples) + client = _RecordingDataClient(delay_first=True) + transfer = TransferDomain(args=_args(), data_system_client=client) + + asyncio.run( + _run_partition_in_batches( + transfer, + rollout_id=7, + current_partition_quota=2, + batches=[[_group(0, 2)], [_group(1, 3)]], + ) + ) + + assert [len(samples) for samples in converted_samples] == [2, 6] + assert [call["partition_id"] for call in client.calls] == ["train_7", "train_7"] + assert [call["is_last"] for call in client.calls] == [False, True] + assert client.events == ["start-0", "end-0", "start-1", "end-1"] + assert [meta[AGENTIC_VARIABLE_ROW_PADDING_KEY] for call in client.calls for meta in call["custom_meta"]] == [ + False, + False, + False, + False, + False, + True, + True, + True, + ] + + final_samples = converted_samples[-1] + assert [sample.remove_sample for sample in final_samples] == [False, False, False, True, True, True] + for padding_sample in final_samples[-3:]: + assert padding_sample.custom_advantage == 0.0 + assert padding_sample.reward == 0.0 + assert padding_sample.metadata[AGENTIC_VARIABLE_ROW_PADDING_KEY] is True + assert padding_sample.metadata["row_id"].startswith("agentic-padding-v1:7:") + physical_row_ids = [ + meta[AGENTIC_ROW_IDENTITY_KEY]["row_id"] for call in client.calls for meta in call["custom_meta"] + ] + assert len(physical_row_ids) == len(set(physical_row_ids)) + assert all( + int(data_tag) == meta[AGENTIC_ROW_IDENTITY_TAG_KEY] + for call in client.calls + for data_tag, meta in zip( + call["data"][AGENTIC_ROW_IDENTITY_TAGS_FIELD], + call["custom_meta"], + strict=True, + ) + ) + timing = transfer.transfer_timing_snapshot() + assert len(timing) == 2 + assert all(record["ok"] is True for record in timing) + assert all( + record[key] >= 0 + for record in timing + for key in ("reorder_ms", "identity_validation_ms", "serialization_ms", "queue_transfer_ms") + ) + assert transfer._partition_actual_rows == {} + assert transfer._partition_identity_state == {} + + +def test_variable_row_transfer_exact_eight_rows_adds_no_padding(monkeypatch): + converted_samples: list[list[Sample]] = [] + _install_fake_converter(monkeypatch, converted_samples) + client = _RecordingDataClient() + transfer = TransferDomain(args=_args(), data_system_client=client) + + asyncio.run( + _run_partition_in_batches( + transfer, + rollout_id=3, + current_partition_quota=2, + batches=[[_group(0, 4)], [_group(1, 4)]], + ) + ) + + assert [len(samples) for samples in converted_samples] == [4, 4] + assert [call["is_last"] for call in client.calls] == [False, True] + assert all( + meta[AGENTIC_VARIABLE_ROW_PADDING_KEY] is False for call in client.calls for meta in call["custom_meta"] + ) + + +def test_variable_row_transfer_preserves_previous_partition_row_count_across_rebind(monkeypatch): + converted_samples: list[list[Sample]] = [] + _install_fake_converter(monkeypatch, converted_samples) + client = _RecordingDataClient() + transfer = TransferDomain(args=_args(), data_system_client=client) + + async def _scenario(): + await _run_partition( + transfer, + rollout_id=10, + previous_partition_quota=0, + current_partition_quota=2, + groups=[_group(0, 2)], + ) + assert transfer._partition_actual_rows == {10: 2} + await _run_partition( + transfer, + rollout_id=11, + previous_partition_quota=1, + current_partition_quota=0, + groups=[_group(1, 3)], + ) + + asyncio.run(_scenario()) + + assert [call["partition_id"] for call in client.calls] == ["train_10", "train_10"] + assert [call["is_last"] for call in client.calls] == [False, True] + assert [len(samples) for samples in converted_samples] == [2, 6] + assert sum(meta[AGENTIC_VARIABLE_ROW_PADDING_KEY] for call in client.calls for meta in call["custom_meta"]) == 3 + assert transfer._partition_actual_rows == {} + + +def test_two_tasks_eight_shuffled_unequal_trajectories_round_trip_identity_and_action_rows(monkeypatch): + converted_samples: list[list[Sample]] = [] + _install_fake_converter(monkeypatch, converted_samples) + client = _RecordingDataClient() + args = _args( + rollout_batch_size=2, + over_sampling_batch_size=2, + n_samples_per_prompt=8, + colocate=True, + global_batch_size=16, + num_iters_per_train_update=1, + ) + transfer = TransferDomain(args=args, data_system_client=client) + task_a = _multi_trajectory_group( + group_index=10, + task_id="task-a", + trajectory_lengths=[1, 3, 2, 4, 1, 2, 3, 5], + ) + task_b = _multi_trajectory_group( + group_index=20, + task_id="task-b", + trajectory_lengths=[2, 1, 4, 3, 2, 5, 1, 3], + ) + expected_tokens = {sample.metadata["row_id"]: list(sample.tokens) for sample in [*task_a, *task_b]} + + asyncio.run( + _run_partition( + transfer, + rollout_id=13, + previous_partition_quota=0, + current_partition_quota=2, + groups=[task_b, task_a], + ) + ) + + assert len(client.calls) == 1 + call = client.calls[0] + shuffled_round_trip = list( + zip( + call["custom_meta"], + call["data"]["tokens"], + call["data"]["loss_masks"], + call["data"][AGENTIC_ROW_IDENTITY_TAGS_FIELD], + strict=True, + ) + ) + random.Random(123).shuffle(shuffled_round_trip) + real_identities = [] + for custom_meta, tokens, loss_mask, row_tag in shuffled_round_trip: + identity = custom_meta[AGENTIC_ROW_IDENTITY_KEY] + assert int(row_tag) == custom_meta[AGENTIC_ROW_IDENTITY_TAG_KEY] + if identity["padding"]: + assert loss_mask == [0] + assert identity["rollout_group_id"] is None + continue + real_identities.append(identity) + assert tokens == expected_tokens[identity["row_id"]] + assert loss_mask == [1] + assert identity["action_token_count"] == 1 + assert {identity["task_id"] for identity in real_identities} == {"task-a", "task-b"} + assert {identity["rollout_group_id"] for identity in real_identities} == {"10", "20"} + assert {identity["policy_version"] for identity in real_identities} == {"11"} + assert len({identity["trajectory_id"] for identity in real_identities}) == 16 + assert len({identity["row_id"] for identity in real_identities}) == len(expected_tokens) + + +def test_variable_row_identity_rejects_duplicate_retry_incomplete_group_and_policy_mix(): + transfer = TransferDomain(args=_args(), data_system_client=None) + first_group = _group(0, 2) + transfer._prepare_variable_row_groups( + groups=[first_group], + partition_rollout_id=5, + is_last=False, + ) + with pytest.raises(RuntimeError, match="duplicate row_id|more than once"): + transfer._prepare_variable_row_groups( + groups=[copy.deepcopy(first_group)], + partition_rollout_id=5, + is_last=False, + ) + + different_policy = _group(1, 1) + for sample in different_policy: + sample.weight_versions = ["8"] + sample.metadata["policy_version"] = "8" + with pytest.raises(RuntimeError, match="mixes policy versions"): + transfer._prepare_variable_row_groups( + groups=[different_policy], + partition_rollout_id=5, + is_last=True, + ) + + incomplete_transfer = TransferDomain( + args=_args(n_samples_per_prompt=2), + data_system_client=None, + ) + with pytest.raises(RuntimeError, match="incomplete"): + incomplete_transfer._prepare_variable_row_groups( + groups=[_group(9, 2)], + partition_rollout_id=9, + is_last=False, + ) + + +def test_variable_row_identity_rejects_duplicate_row_inside_group_and_stale_restart_state(): + duplicate_group = _group(0, 2) + duplicate_group[1].metadata["row_id"] = duplicate_group[0].metadata["row_id"] + transfer = TransferDomain(args=_args(), data_system_client=None) + with pytest.raises(RuntimeError, match="duplicate row_id"): + transfer._prepare_variable_row_groups( + groups=[duplicate_group], + partition_rollout_id=0, + is_last=False, + ) + + transfer._partition_identity_state[1] = { + "row_ids": {"residual-row"}, + "group_manifests": {}, + "policy_version": "7", + } + with pytest.raises(RuntimeError, match="stale or future"): + transfer.rebind_step(rollout_id=3) + + +def test_default_transfer_mode_does_not_add_padding(monkeypatch): + converted_samples: list[list[Sample]] = [] + _install_fake_converter(monkeypatch, converted_samples) + client = _RecordingDataClient() + transfer = TransferDomain( + args=_args( + rollout_batch_size=1, + over_sampling_batch_size=1, + group_rm=False, + ), + data_system_client=client, + ) + + asyncio.run( + _run_partition( + transfer, + rollout_id=4, + previous_partition_quota=0, + current_partition_quota=1, + groups=[_group(0, 5)], + ) + ) + + assert [len(samples) for samples in converted_samples] == [5] + assert client.calls[0]["is_last"] is True + assert client.calls[0]["custom_meta"] == [{"total_lengths": 2}] * 5 + assert transfer._partition_actual_rows == {} + assert transfer.transfer_timing_snapshot() == [] + + +def test_variable_row_transfer_rejects_fully_async_mode(): + with pytest.raises(ValueError, match="synchronous training only"): + TransferDomain( + args=_args(fully_async=True, hybrid=False), + data_system_client=object(), + ) + + +def test_variable_row_transfer_rejects_hybrid_mode(): + with pytest.raises(ValueError, match="synchronous training only"): + TransferDomain( + args=_args(fully_async=True, hybrid=True), + data_system_client=None, + ) + + +def test_discard_pending_transfers_clears_variable_partition_accounting(): + transfer = TransferDomain(args=_args(), data_system_client=None) + transfer._partition_actual_rows[9] = 3 + transfer._partition_identity_state[9] = { + "row_ids": {"row"}, + "group_manifests": {}, + "policy_version": "7", + } + + asyncio.run(transfer.discard_pending_transfers()) + + assert transfer._partition_actual_rows == {} + assert transfer._partition_identity_state == {} + + +def test_padding_sample_converts_to_zero_loss_mask(): + from relax.utils.utils import convert_samples_to_train_data + + template = _group(0, 1)[0] + template.metadata["raw_reward"] = 9.0 + padding_sample = _variable_row_padding_sample(args=_args(), template=template) + convert_args = SimpleNamespace( + custom_reward_post_process_path=None, + agentic_custom_advantage_path="examples.graphgpo.advantage.compute", + reward_key=None, + multimodal_keys=None, + use_opd=False, + debug_train_only=True, + ) + + batch = convert_samples_to_train_data(convert_args, [padding_sample]) + + assert batch["loss_masks"] == [[0]] + assert batch["rewards"] == [0.0] + assert batch["raw_reward"] == [0.0] + assert template.remove_sample is False + assert template.metadata["raw_reward"] == 9.0