diff --git a/aorlta/README.md b/aorlta/README.md index 1f33dc7cb..327631e66 100644 --- a/aorlta/README.md +++ b/aorlta/README.md @@ -174,7 +174,8 @@ compile and waitcheck failures, so a candidate that passes them stays unscored); `jobs/aorlta-reward-standin.sbatch` (a test of the stand-in GPU checks) and two scripts the run job calls; `docker/` the trainer image; `tests/` the CPU suite; `tools/merge_base_adapter.py`, which merges a PEFT adapter into its base checkpoint; `examples/standin-pool/`, a synthetic prompt pool; `examples/hip-kernels/` with `configs/hipk-score.yaml` -and `configs/hipk-train.yaml`, the worked example of a task as a plugin; `docs/ADAPTING.md` and `docs/TUTORIAL.md`. Code comments cite sections and +and `configs/hipk-train.yaml`, the worked example of a task as a plugin; `examples/verticals/`, starter plugins to copy +for an ISA fuzzer, a Q&A expert and sanitizer triage; `docs/ADAPTING.md` and `docs/TUTORIAL.md`. Code comments cite sections and labels (Section 6.x, C1-C9, D7, ...) of our internal design notes, which are not included. ## Requirements diff --git a/aorlta/docs/ADAPTING.md b/aorlta/docs/ADAPTING.md index 15619630b..3c1adb614 100644 --- a/aorlta/docs/ADAPTING.md +++ b/aorlta/docs/ADAPTING.md @@ -26,6 +26,9 @@ Code paths are relative to `src/aorlta/`; commands and the `configs/`, `examples adapter** on the same set and compare. 8. **Write your two-line summary:** what the model is post-trained on, and how its answers are checked. +Starters to copy for three verticals (an ISA fuzzer, a Q&A expert, sanitizer triage), each a plugin package with its +configs and a CPU test, are in `examples/verticals/`. + ## Your own task, as a plugin A task of your own (its prompts, how an answer is read, its repair turns, its checks and its reward) does not need diff --git a/aorlta/examples/verticals/README.md b/aorlta/examples/verticals/README.md new file mode 100644 index 000000000..0e3950f63 --- /dev/null +++ b/aorlta/examples/verticals/README.md @@ -0,0 +1,17 @@ +# verticals: starter plugins, one per kind of task + +Three skeletons to copy and fill in, each a plugin package on its own ([docs/ADAPTING.md](../../docs/ADAPTING.md)): +its module, check profile, a toy prompt pool split into `train/` and `heldout/`, and a score and a train config. +Nothing in one imports another, so copy a directory, rename its prefix, and replace what is marked PLACEHOLDER and +TODO. The pools, formats and reports are toy placeholders written for these starters; none of them is a real task. + +| directory | the model gets | it produces | the placeholder check | +|---|---|---|---| +| `isa-fuzzer/` | a goal | a short test program in a declared toy format | the format's grammar, line count and required instructions; an optional external tool after that | +| `qa-expert/` | a question in your domain | an answer on an `Answer:` line | an exact match against the reference, after normalising | +| `sanitizer-triage/` | a tool's report | a JSON verdict: a category from a closed list and a rationale | the verdict's keys, category and rationale; not yet the reference label | + +Each runs on one GPU, with `rollout.compile_adapter: none` and one CPU leg, and its reward keeps the service's rules: +in [-1.55, 1.0], parts exactly `base`, `broke_healthy` and `length`, a value of its own for each failure, and an +unsettled check unscored, never clean. `tests/examples/test_verticals.py` loads each config, tests each placeholder, +and scores and trains each on CPU through the real wiring. diff --git a/aorlta/examples/verticals/isa-fuzzer/README.md b/aorlta/examples/verticals/isa-fuzzer/README.md new file mode 100644 index 000000000..bc3d9628c --- /dev/null +++ b/aorlta/examples/verticals/isa-fuzzer/README.md @@ -0,0 +1,24 @@ +# isa-fuzzer: a starter for an ISA-fuzzer-style task + +The model is given a goal and writes a short test program in a declared format; a check validates the program and, in +a real task, runs it through your own tool. Everything here is a PLACEHOLDER: `toy-asm` is a made-up format and the +goals are toy ones. It is a plugin ([docs/ADAPTING.md](../../../docs/ADAPTING.md)) and imports nothing from the other +examples, so the directory can be copied on its own. + +| file | what | what to replace | +|---|---|---| +| `isa_fuzzer.py` | the plugin: prompts, parse and feedback, the `ref.program` check, the `isa-fuzzer.v1` reward | `USER` and `LINE_RE` (your format), `program_problem` (your check), `IsaV1.BASE` (your reward: its TODO lists the rules) | +| `isa-fuzzer.yaml` | the check profile: one CPU leg | `isa-fuzzer.check: {tool: [...]}` sends the programs that pass to your tool | +| `pool/{train,heldout}/tasks.jsonl` | 4 + 4 toy goals | your goals: many more than `rollout.P` | +| `score.yaml`, `train.yaml` | score the held-out set, train on the rest; one GPU | the site paths | + +The tool reads `{"goal", "program"}` as JSON on stdin and prints `{"verdict": "pass"}` or `{"verdict": "fail"}`; +anything else leaves the sample unscored, never clean. + +From `aorlta/`, with the site paths set (`AORLTA_SITE`, docs/TUTORIAL.md section 4) and the config committed: + +```bash +OMP_NUM_THREADS=1 python -m pytest -q tests/examples/test_verticals.py -k isa # on CPU, a score and a training run +CONFIG=examples/verticals/isa-fuzzer/score.yaml sbatch jobs/aorlta-run.sbatch # score the base model +CONFIG=examples/verticals/isa-fuzzer/train.yaml sbatch jobs/aorlta-run.sbatch # train, then score the adapter +``` diff --git a/aorlta/examples/verticals/isa-fuzzer/isa-fuzzer.yaml b/aorlta/examples/verticals/isa-fuzzer/isa-fuzzer.yaml new file mode 100644 index 000000000..ac0939cad --- /dev/null +++ b/aorlta/examples/verticals/isa-fuzzer/isa-fuzzer.yaml @@ -0,0 +1,12 @@ +# `isa-fuzzer`: one leg, `ref.program`, on the CPU: the reply's last fenced block against the toy grammar +# (examples/verticals/isa-fuzzer/isa_fuzzer.py). No compile and no sanitizer leg. +name: isa-fuzzer +reward_fn: isa-fuzzer.v1 +clean_repeats: 1 +materialize: true +pools: + check: {resource: local_cpu, adapter: isa-fuzzer.check, timeout: gate, workers: 4} +stages: + - {name: S0, pool: check, legs: [ref.program]} +adapters: + isa-fuzzer.check: {} # {tool: [/path/to/your-tool, args...]} decides the programs that pass the toy grammar diff --git a/aorlta/examples/verticals/isa-fuzzer/isa_fuzzer.py b/aorlta/examples/verticals/isa-fuzzer/isa_fuzzer.py new file mode 100644 index 000000000..a183c6130 --- /dev/null +++ b/aorlta/examples/verticals/isa-fuzzer/isa_fuzzer.py @@ -0,0 +1,285 @@ +"""A starter for an ISA-fuzzer-style task (examples/verticals/README.md): the model writes a short test program for a +stated goal, and a check validates it. Every part marked PLACEHOLDER is to be replaced with your own. + +A pool is a directory holding `tasks.jsonl`, one task per line: `{"id", "goal", "min_lines", "max_lines", +"must_use"}`. The model is asked for a program in `toy-asm`, a made-up text format declared in `USER` (one +instruction per line: `NOP`, `MOV rD, rS|#imm`, or `ADD|SUB|MUL|AND|OR|XOR rD, rA, rB|#imm`, registers r0-r7), in one +fenced block. A reply with no fenced block gets one repair turn that says so (this module's parse and feedback); there +is no compile step (`rollout.compile_adapter: none`). The reward service runs one leg, `ref.program`, on the CPU: the +last fenced block against the toy grammar, the task's line range and its required instructions. With a command in the +profile's `isa-fuzzer.check.tool`, a program that passes those is then the tool's to decide: it reads `{"goal", +"program"}` as JSON on stdin and prints `{"verdict": "pass"}` or `{"verdict": "fail"}`; anything else leaves the +sample unscored. + +Registered names: task `isa-fuzzer`; parse `isa-fuzzer.parse`; feedback `isa-fuzzer.feedback`; profile `isa-fuzzer` +(isa-fuzzer.yaml); check adapter `isa-fuzzer.check`; reward function `isa-fuzzer.v1`. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import re +import shutil +from pathlib import Path + +HERE = Path(__file__).resolve().parent +# PLACEHOLDER: the format and the prompt. Declare your real program format here, or point the model at its manual. +SYSTEM = "You write short test programs in the toy-asm format." +USER = ("Goal: {goal}\n\n" + "Write a toy-asm program of {min_lines} to {max_lines} lines. toy-asm has one instruction per line: NOP; " + "MOV rD, rS or MOV rD, #imm; or OP rD, rA, rB or OP rD, rA, #imm with OP one of ADD, SUB, MUL, AND, OR, XOR. " + "Registers are r0 to r7; an immediate is # and a whole number of at most six digits, which may be negative. " + "Give the program in one fenced code block.") +BLOCK_RE = re.compile(r"```[^\n`]*\n(.*?)```", re.DOTALL) +REG, IMM = r"r[0-7]", r"#-?\d{1,6}" +LINE_RE = re.compile(rf"NOP|MOV\s+{REG}\s*,\s*(?:{REG}|{IMM})" + rf"|(?:ADD|SUB|MUL|AND|OR|XOR)\s+{REG}\s*,\s*{REG}\s*,\s*(?:{REG}|{IMM})") + + +def load_rows(pool_dir) -> list[dict]: + rows = [json.loads(line) for line in (Path(pool_dir) / "tasks.jsonl").read_text(encoding="utf-8").splitlines() + if line.strip()] + ids = [r["id"] for r in rows] + if len(set(ids)) != len(ids): + raise ValueError(f"{pool_dir}/tasks.jsonl lists an id twice") + return rows + + +def row_sha256(row: dict) -> str: + return hashlib.sha256(json.dumps(row, sort_keys=True).encode("utf-8")).hexdigest() + + +def final_program(reply: str) -> str | None: + """The last fenced block's text, or None.""" + found = BLOCK_RE.findall(reply) + return found[-1] if found else None + + +def program_problem(program: str, row: dict) -> str | None: + """PLACEHOLDER check: None for a program that keeps the toy grammar, the task's line range and its required + instructions; otherwise the failure's reason (`malformed`, `wrong_length`, `missing_op`). A real check runs the + program (an assembler, a simulator, your fuzzer's harness) instead, or after this, through the `tool` setting.""" + lines = [ln.strip() for ln in program.splitlines() if ln.strip()] + if not all(LINE_RE.fullmatch(ln) for ln in lines): + return "malformed" + if not row["min_lines"] <= len(lines) <= row["max_lines"]: + return "wrong_length" + if not set(row.get("must_use", ())) <= {ln.split()[0] for ln in lines}: + return "missing_op" + return None + + +def messages(row: dict) -> list[dict]: + return [{"role": "system", "content": SYSTEM}, + {"role": "user", "content": USER.format(goal=row["goal"], min_lines=row["min_lines"], + max_lines=row["max_lines"])}] + + +def records(pool_dir, tokenizer, thinking): + from aorlta.ids import sha256_ids + from aorlta.rollout.templates import render_turn0 + from aorlta.schemas import PromptRecord, TaskRef + + out = [] + for row in load_rows(pool_dir): + msgs = messages(row) + out.append(PromptRecord( + prompt_id=f"isa-fuzzer:{row['id']}", turn0_ids=None, messages=msgs, + prompt_ids_sha256=sha256_ids(render_turn0(tokenizer, msgs, thinking)), + task=TaskRef(contract_id=f"isa-fuzzer:{row['id']}", parent_id=None, parent_healthy=None, + reference_kind=None, task_sha256=row_sha256(row)), + source="standin")) + return out + + +class IsaParse: + """`isa-fuzzer.parse`: a reply with a fenced block is a program (nothing to compile); one without is rejected with + a message, which `isa-fuzzer.feedback` turns into a repair turn. Whether the program is any good is the check's to + say, not the parse's.""" + + name = "isa-fuzzer.parse" + + def __init__(self): + from aorlta.rollout.templates import IM_END_ID + + self.reply_stops = (IM_END_ID,) + + def check_prompts(self, prompts, thinking, max_repair_turns=0): + return None + + def check(self, text, thinking, prompt): + from aorlta.rollout.compile_check import ParseOutcome + + reply = text.rsplit("", 1)[-1] if thinking == "on" else text + program = final_program(reply) + if program is None: + return ParseOutcome("rejected", rejection=( + "Your reply has no fenced code block. Give the toy-asm program in one fenced code block.", + "no_program")) + return ParseOutcome("edits", code=program) + + +class IsaFeedback: + name = "isa-fuzzer.feedback" + + def build(self, errors, max_chars): + from aorlta.rollout.feedback import build_feedback + + return build_feedback(errors, max_chars) + + def build_apply(self, message, reason): + from aorlta.rollout.feedback import Feedback + + return Feedback(message, 0, 0) + + +class IsaCheck: + """`isa-fuzzer.check`, leg `ref.program`: `fail` `no_program` for a reply with no fenced block, `fail` with + `program_problem`'s reason for a program it rejects, then clean, or the tool's (settings `tool`, a command as a + list) `pass` clean and `fail` `tool_rejected`; unsettled when the task is unknown, the reply cannot be read, or the + tool gives no verdict.""" + + legs = ("ref.program",) + resource = "local_cpu" + + def __init__(self, ctx): + from aorlta.reward.adapters.base import ProcessRunner + from aorlta.reward.adapters.toolchain import file_identity + + self.name, self.ctx = ctx.name, ctx + self.tool = ctx.settings.get("tool") + self.tool_timeout_s = float(ctx.settings.get("tool_timeout_s", 60.0)) + self.version = "1" # cached legs are keyed by it: the tool, which decides what the format passes, is part of it + if self.tool: + files = {a: file_identity(a) for a in self.tool if isinstance(a, str) and Path(a).is_file()} + exe = shutil.which(self.tool[0]) if isinstance(self.tool[0], str) else None # as PATH finds it, if bare + if exe: + files[f"exe:{self.tool[0]}"] = f"{os.path.realpath(exe)}:{file_identity(exe)}" + ident = json.dumps({"tool": list(self.tool), "timeout_s": self.tool_timeout_s, "files": files}, + sort_keys=True) + self.version = f"1+tool-{hashlib.sha256(ident.encode()).hexdigest()[:16]}" + self.runner = ProcessRunner() # kills the tool's whole process group on a timeout or at close() + self.rows = {} + for pool in ctx.pools: + if (Path(pool) / "tasks.jsonl").is_file(): + for row in load_rows(pool): + self.rows[row_sha256(row)] = row + + def run(self, job): + from aorlta.reward.adapters.base import make_result, unsettled + + t0 = self.ctx.clock() + try: + row = self.rows.get(job.task_sha256) + response = job.inputs.get("response") + if job.leg != "ref.program" or row is None or not response: + return unsettled(job, "adapter_error", t_start=t0, t_end=self.ctx.clock()) + program = final_program(Path(response).read_text(encoding="utf-8").rsplit("", 1)[-1]) + if program is None: + return make_result(job, "fail", "no_program", t_start=t0, t_end=self.ctx.clock()) + problem = program_problem(program, row) + if problem is not None: + return make_result(job, "fail", problem, t_start=t0, t_end=self.ctx.clock()) + if not self.tool: + return make_result(job, "clean", t_start=t0, t_end=self.ctx.clock()) + verdict = self._tool(row, program) + if verdict == "pass": + return make_result(job, "clean", t_start=t0, t_end=self.ctx.clock()) + if verdict == "fail": + return make_result(job, "fail", "tool_rejected", t_start=t0, t_end=self.ctx.clock()) + return unsettled(job, "tool_no_verdict", t_start=t0, t_end=self.ctx.clock()) + except Exception: # noqa: BLE001 - an adapter never raises + return unsettled(job, "adapter_error", t_start=t0, t_end=self.ctx.clock()) + + def _tool(self, row, program) -> str | None: + proc = self.runner.run(list(self.tool), input=json.dumps({"goal": row["goal"], "program": program}), + timeout_s=min(self.tool_timeout_s, self.ctx.timeout_s)) + if proc.timed_out or proc.killed or proc.launch_error or proc.returncode != 0 or not proc.stdout.strip(): + return None + try: + out = json.loads(proc.stdout.strip().splitlines()[-1]) + except ValueError: + return None + v = out.get("verdict") if isinstance(out, dict) else None + return v if v in ("pass", "fail") else None + + def controls(self): + return [] + + def health(self): + return {"name": self.name, "responsive": True, "tasks": len(self.rows)} + + def close(self): + self.runner.close() + + +class IsaV1: + """`isa-fuzzer.v1`, a PLACEHOLDER reward: +1.0 for a program the check passes, a value of its own for each failure + it names (`BASE`), then -0.05 x the share of the token budget used. An unsettled leg, or a reason `BASE` does not + name, leaves the sample unscored.""" + + name = "isa-fuzzer.v1" + version = "1" + legs = ("ref.program",) + # TODO(your vertical): replace these values with your real reward, and the check's reasons with your tool's. Keep: + # - the total in [-1.55, 1.0], its parts exactly `base`, `broke_healthy` and `length`, adding up to it (the + # service recomputes it and leaves anything else unscored, with an alarm): so a base in [-1.5, 1.0]; + # - an unsettled or missing leg, or a reason you did not name, unscored: never clean, never a score; + # - each failure its own value, ordered by how close it came, so the answers to one prompt get different + # rewards (GRPO learns only from the differences); + # - `claims_clean` true only when the leg is clean. + BASE = {"tool_rejected": -0.3, "missing_op": -0.5, "wrong_length": -0.6, "malformed": -0.8, "no_program": -1.0} + + def __init__(self, profile=None): + from aorlta.reward.rewardfn.base import code_sha256 + + self.settings = {} + self.code_sha256 = code_sha256(__file__) + + def _base(self, legs): + leg = legs.get("ref.program") + if leg is None: + return "pending", "ref.program" + if leg.state == "clean": + return "score", 1.0 + if leg.state == "fail" and leg.reason in self.BASE: + return "score", self.BASE[leg.reason] + return "unscored", f"ref.program:{leg.reason or leg.state}" + + @staticmethod + def _finish(base, cand): + parts = {"base": float(base), "broke_healthy": 0.0, + "length": -0.05 * (cand.response_tokens / cand.max_new_tokens)} + return parts["base"] + parts["broke_healthy"] + parts["length"], parts + + def compute(self, legs, cand): + kind, value = self._base(legs) + return (None, {"reason": str(value)}) if kind != "score" else self._finish(value, cand) + + def decide(self, partial, cand): + from aorlta.reward.rewardfn.base import Decided, NeedMore, Unscored + + kind, value = self._base(partial) + if kind == "pending": + return NeedMore((value,)) + return Unscored(value) if kind == "unscored" else Decided(*self._finish(value, cand)) + + @staticmethod + def claims_clean(reward, parts): + return reward is not None and isinstance(parts, dict) and parts.get("base", 0.0) > 0.0 + + +def register() -> None: + from aorlta import plugins + + plugins.register_task(plugins.Task("isa-fuzzer", records)) + plugins.register_parse_adapter("isa-fuzzer.parse", IsaParse) + plugins.register_feedback_template("isa-fuzzer.feedback", IsaFeedback) + plugins.register_profile("isa-fuzzer", HERE / "isa-fuzzer.yaml") + plugins.register_check_adapter("isa-fuzzer.check", IsaCheck) + plugins.register_reward_fn("isa-fuzzer.v1", IsaV1) + + +register() diff --git a/aorlta/examples/verticals/isa-fuzzer/pool/heldout/tasks.jsonl b/aorlta/examples/verticals/isa-fuzzer/pool/heldout/tasks.jsonl new file mode 100644 index 000000000..406c1d3af --- /dev/null +++ b/aorlta/examples/verticals/isa-fuzzer/pool/heldout/tasks.jsonl @@ -0,0 +1,4 @@ +{"id": "placeholder-h1", "goal": "PLACEHOLDER held-out goal 1: any toy program that uses SUB.", "min_lines": 2, "max_lines": 4, "must_use": ["SUB"]} +{"id": "placeholder-h2", "goal": "PLACEHOLDER held-out goal 2: any toy program that uses AND and OR.", "min_lines": 3, "max_lines": 6, "must_use": ["AND", "OR"]} +{"id": "placeholder-h3", "goal": "PLACEHOLDER held-out goal 3: any toy program that uses MOV.", "min_lines": 1, "max_lines": 3, "must_use": ["MOV"]} +{"id": "placeholder-h4", "goal": "PLACEHOLDER held-out goal 4: any toy program that uses ADD and SUB.", "min_lines": 2, "max_lines": 5, "must_use": ["ADD", "SUB"]} diff --git a/aorlta/examples/verticals/isa-fuzzer/pool/train/tasks.jsonl b/aorlta/examples/verticals/isa-fuzzer/pool/train/tasks.jsonl new file mode 100644 index 000000000..c462713a2 --- /dev/null +++ b/aorlta/examples/verticals/isa-fuzzer/pool/train/tasks.jsonl @@ -0,0 +1,4 @@ +{"id": "placeholder-1", "goal": "PLACEHOLDER goal 1: any toy program that uses ADD.", "min_lines": 2, "max_lines": 4, "must_use": ["ADD"]} +{"id": "placeholder-2", "goal": "PLACEHOLDER goal 2: any toy program that uses XOR.", "min_lines": 2, "max_lines": 5, "must_use": ["XOR"]} +{"id": "placeholder-3", "goal": "PLACEHOLDER goal 3: any toy program that uses MOV and MUL.", "min_lines": 3, "max_lines": 6, "must_use": ["MOV", "MUL"]} +{"id": "placeholder-4", "goal": "PLACEHOLDER goal 4: any toy program that uses NOP.", "min_lines": 1, "max_lines": 3, "must_use": ["NOP"]} diff --git a/aorlta/examples/verticals/isa-fuzzer/score.yaml b/aorlta/examples/verticals/isa-fuzzer/score.yaml new file mode 100644 index 000000000..a621a11da --- /dev/null +++ b/aorlta/examples/verticals/isa-fuzzer/score.yaml @@ -0,0 +1,39 @@ +# Score a model on the isa-fuzzer starter's 4 held-out placeholder goals, 4 samples each +# (orchestrator.mode: score). +# One GPU: one TokenSpeed engine on GPU 0; the check runs on the CPU. +# Site-specific paths, to edit here or to set in a site file (AORLTA_SITE): run.root, model.base_dir (and +# model.adapter_dir to score a trained adapter) and data.prompt_pools[0].path, which is this checkout's +# aorlta/examples/verticals/isa-fuzzer/pool/heldout. +plugins: [examples/verticals/isa-fuzzer/isa_fuzzer.py] +run: + run_id: isa-fuzzer-score + source: standin + root: /scratch/aorlta/runs +model: + base_dir: /data/hf/hub/models--Qwen--Qwen3.5-4B/snapshots/851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a + source: standin +data: + prompt_pools: + - path: /opt/aorlta/examples/verticals/isa-fuzzer/pool/heldout + source: standin + task: isa-fuzzer +rollout: + engines: [0] + K: 4 + P: 4 + max_repair_turns: 1 + thinking: "off" + temperature: 0.7 + max_new_tokens: 512 + compile_adapter: none + parse_adapter: isa-fuzzer.parse + feedback_template: isa-fuzzer.feedback + # few prompts, each with a short answer, can give a group K identical replies with a sound sampler + # (the engines' start-up self-test checks that), so only stop when most groups do + identical_groups_stop: 0.75 +reward: + profile: isa-fuzzer + reward_fn: isa-fuzzer.v1 +orchestrator: + mode: score + run_seed: 20261113 diff --git a/aorlta/examples/verticals/isa-fuzzer/train.yaml b/aorlta/examples/verticals/isa-fuzzer/train.yaml new file mode 100644 index 000000000..788b1735e --- /dev/null +++ b/aorlta/examples/verticals/isa-fuzzer/train.yaml @@ -0,0 +1,42 @@ +# Train on the isa-fuzzer starter's 4 placeholder goals, 8 samples each, 3 iterations: a template to +# copy with your own pool, far too small to learn from. +# One GPU: one TokenSpeed engine on GPU 0 and one LoRA rank beside it; the check runs on the CPU. +# Site-specific paths, to edit here or to set in a site file (AORLTA_SITE): run.root, model.base_dir and +# data.prompt_pools[0].path, which is this checkout's aorlta/examples/verticals/isa-fuzzer/pool/train. +plugins: [examples/verticals/isa-fuzzer/isa_fuzzer.py] +run: + run_id: isa-fuzzer-train + source: standin + root: /scratch/aorlta/runs +model: + base_dir: /data/hf/hub/models--Qwen--Qwen3.5-4B/snapshots/851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a + source: standin +data: + prompt_pools: + - path: /opt/aorlta/examples/verticals/isa-fuzzer/pool/train + source: standin + task: isa-fuzzer +rollout: + engines: [0] + K: 8 + P: 4 + max_repair_turns: 1 + thinking: "off" + temperature: 1.0 + max_new_tokens: 512 + compile_adapter: none + parse_adapter: isa-fuzzer.parse + feedback_template: isa-fuzzer.feedback + # few prompts, each with a short answer, can give a group K identical replies with a sound sampler + # (the engines' start-up self-test checks that), so only stop when most groups do + identical_groups_stop: 0.75 +reward: + profile: isa-fuzzer + reward_fn: isa-fuzzer.v1 +orchestrator: + iterations: 3 + run_seed: 20261113 +trainer: + ranks: 1 + micro_batch_tokens: 20480 + lr: 2.0e-5 diff --git a/aorlta/examples/verticals/qa-expert/README.md b/aorlta/examples/verticals/qa-expert/README.md new file mode 100644 index 000000000..0733a5d6b --- /dev/null +++ b/aorlta/examples/verticals/qa-expert/README.md @@ -0,0 +1,23 @@ +# qa-expert: a starter for a Q&A expert model + +The model answers questions in your domain and is rewarded for matching a reference answer. Everything here is a +PLACEHOLDER: the questions are toy ones and the check is an exact match after normalising. It is a plugin +([docs/ADAPTING.md](../../../docs/ADAPTING.md)) and imports nothing from the other examples, so the directory can be +copied on its own; `examples/qa/` is a fuller example of the same shape, with a judge command. + +| file | what | what to replace | +|---|---|---| +| `qa_expert.py` | the plugin: prompts, parse and feedback, the `ref.answer` check, the `qa-expert.v1` reward | `SYSTEM` (your domain), `normalise` and `ExpertMatch.run` (your match or judge), `ExpertV1.BASE` (your reward: its TODO lists the rules) | +| `qa-expert.yaml` | the check profile: one CPU leg | your check's settings, if it has any | +| `pool/{train,heldout}/questions.jsonl` | 4 + 4 toy questions, `{"id", "question", "answer", "aliases"?}` | your questions: many more than `rollout.P` | +| `score.yaml`, `train.yaml` | score the held-out set, train on the rest; one GPU | the site paths | + +Score the base model first: if it already answers nearly everything, or nothing, there is nothing to learn. + +From `aorlta/`, with the site paths set (`AORLTA_SITE`, docs/TUTORIAL.md section 4) and the config committed: + +```bash +OMP_NUM_THREADS=1 python -m pytest -q tests/examples/test_verticals.py -k qa # on CPU, a score and a training run +CONFIG=examples/verticals/qa-expert/score.yaml sbatch jobs/aorlta-run.sbatch # score the base model +CONFIG=examples/verticals/qa-expert/train.yaml sbatch jobs/aorlta-run.sbatch # train, then score the adapter +``` diff --git a/aorlta/examples/verticals/qa-expert/pool/heldout/questions.jsonl b/aorlta/examples/verticals/qa-expert/pool/heldout/questions.jsonl new file mode 100644 index 000000000..b6ea3085f --- /dev/null +++ b/aorlta/examples/verticals/qa-expert/pool/heldout/questions.jsonl @@ -0,0 +1,4 @@ +{"id": "placeholder-h1", "question": "[placeholder held-out question] What is 4 + 4?", "answer": "8", "aliases": ["eight"]} +{"id": "placeholder-h2", "question": "[placeholder held-out question] How many hours are in a day?", "answer": "24", "aliases": ["twenty-four"]} +{"id": "placeholder-h3", "question": "[placeholder held-out question] What is the first letter of the English alphabet?", "answer": "a"} +{"id": "placeholder-h4", "question": "[placeholder held-out question] What is the opposite of up?", "answer": "down"} diff --git a/aorlta/examples/verticals/qa-expert/pool/train/questions.jsonl b/aorlta/examples/verticals/qa-expert/pool/train/questions.jsonl new file mode 100644 index 000000000..3a6b3ea9a --- /dev/null +++ b/aorlta/examples/verticals/qa-expert/pool/train/questions.jsonl @@ -0,0 +1,4 @@ +{"id": "placeholder-1", "question": "[placeholder question] What is 2 + 3?", "answer": "5", "aliases": ["five"]} +{"id": "placeholder-2", "question": "[placeholder question] What colour is a clear daytime sky?", "answer": "blue"} +{"id": "placeholder-3", "question": "[placeholder question] How many days are in a week?", "answer": "7", "aliases": ["seven"]} +{"id": "placeholder-4", "question": "[placeholder question] What is the opposite of hot?", "answer": "cold"} diff --git a/aorlta/examples/verticals/qa-expert/qa-expert.yaml b/aorlta/examples/verticals/qa-expert/qa-expert.yaml new file mode 100644 index 000000000..74dbca6ee --- /dev/null +++ b/aorlta/examples/verticals/qa-expert/qa-expert.yaml @@ -0,0 +1,12 @@ +# `qa-expert`: one leg, `ref.answer`, on the CPU: the reply's last `Answer:` line against the reference +# (examples/verticals/qa-expert/qa_expert.py). No compile and no sanitizer leg. +name: qa-expert +reward_fn: qa-expert.v1 +clean_repeats: 1 +materialize: true +pools: + match: {resource: local_cpu, adapter: qa-expert.match, timeout: gate, workers: 4} +stages: + - {name: S0, pool: match, legs: [ref.answer]} +adapters: + qa-expert.match: {} diff --git a/aorlta/examples/verticals/qa-expert/qa_expert.py b/aorlta/examples/verticals/qa-expert/qa_expert.py new file mode 100644 index 000000000..2b256155b --- /dev/null +++ b/aorlta/examples/verticals/qa-expert/qa_expert.py @@ -0,0 +1,230 @@ +"""A starter for a Q&A expert model (examples/verticals/README.md): questions in your domain with reference answers, +and a reward for matching the reference. Every part marked PLACEHOLDER is to be replaced with your own. + +A pool is a directory holding `questions.jsonl`, one question per line: `{"id", "question", "answer", "aliases"?}`. +The model is asked the question and told to end its reply with a line `Answer: `. A reply without that line +gets one repair turn that says so (this module's parse and feedback); there is no compile step +(`rollout.compile_adapter: none`). The reward service runs one leg, `ref.answer`, on the CPU: the last `Answer:` line, +normalised, against the reference and its aliases. + +Registered names: task `qa-expert`; parse `qa-expert.parse`; feedback `qa-expert.feedback`; profile `qa-expert` +(qa-expert.yaml); check adapter `qa-expert.match`; reward function `qa-expert.v1`. +""" + +from __future__ import annotations + +import hashlib +import json +import re +import unicodedata +from pathlib import Path + +HERE = Path(__file__).resolve().parent +# PLACEHOLDER: the prompt. A domain expert usually wants a system message that names the domain. +SYSTEM = "You are a helpful assistant." +USER = "{question}\n\nEnd your reply with a line of the form 'Answer: '." +ANSWER_RE = re.compile(r"^[ \t]*\**[ \t]*answer[ \t]*\**[ \t]*:[ \t]*(.*?)[ \t]*$", re.IGNORECASE) # one line + + +def load_rows(pool_dir) -> list[dict]: + rows = [json.loads(line) for line in (Path(pool_dir) / "questions.jsonl").read_text(encoding="utf-8").splitlines() + if line.strip()] + ids = [r["id"] for r in rows] + if len(set(ids)) != len(ids): + raise ValueError(f"{pool_dir}/questions.jsonl lists an id twice") + return rows + + +def row_sha256(row: dict) -> str: + return hashlib.sha256(json.dumps(row, sort_keys=True).encode("utf-8")).hexdigest() + + +def normalise(text: str) -> str: + """PLACEHOLDER normalisation: NFKC, lower case, surrounding quotes, backticks and asterisks, a trailing full stop + and runs of space removed. A real expert task may want more (units, numeric tolerance, synonyms) or a judge.""" + t = unicodedata.normalize("NFKC", text).strip().lower() + t = t.strip("*`'\" ").rstrip(".").strip() + return re.sub(r"\s+", " ", t) + + +def final_answer(reply: str) -> str | None: + """The last `Answer:` line's answer, or None.""" + found = [m.group(1) for line in reply.splitlines() if (m := ANSWER_RE.match(line))] + return (found[-1] or None) if found else None # the last marker decides, an empty one included + + +def messages(row: dict) -> list[dict]: + return [{"role": "system", "content": SYSTEM}, {"role": "user", "content": USER.format(question=row["question"])}] + + +def records(pool_dir, tokenizer, thinking): + from aorlta.ids import sha256_ids + from aorlta.rollout.templates import render_turn0 + from aorlta.schemas import PromptRecord, TaskRef + + out = [] + for row in load_rows(pool_dir): + msgs = messages(row) + out.append(PromptRecord( + prompt_id=f"qa-expert:{row['id']}", turn0_ids=None, messages=msgs, + prompt_ids_sha256=sha256_ids(render_turn0(tokenizer, msgs, thinking)), + task=TaskRef(contract_id=f"qa-expert:{row['id']}", parent_id=None, parent_healthy=None, + reference_kind=None, task_sha256=row_sha256(row)), + source="standin")) + return out + + +class ExpertParse: + """`qa-expert.parse`: a reply with an `Answer:` line is an answer (nothing to compile); one without is rejected + with a message, which `qa-expert.feedback` turns into a repair turn.""" + + name = "qa-expert.parse" + + def __init__(self): + from aorlta.rollout.templates import IM_END_ID + + self.reply_stops = (IM_END_ID,) + + def check_prompts(self, prompts, thinking, max_repair_turns=0): + return None + + def check(self, text, thinking, prompt): + from aorlta.rollout.compile_check import ParseOutcome + + reply = text.rsplit("", 1)[-1] if thinking == "on" else text + answer = final_answer(reply) + if answer is None: + return ParseOutcome("rejected", rejection=( + "Your reply has no final line of the form 'Answer: '. Give your answer on such a line.", + "no_answer")) + return ParseOutcome("edits", code=answer) + + +class ExpertFeedback: + name = "qa-expert.feedback" + + def build(self, errors, max_chars): + from aorlta.rollout.feedback import build_feedback + + return build_feedback(errors, max_chars) + + def build_apply(self, message, reason): + from aorlta.rollout.feedback import Feedback + + return Feedback(message, 0, 0) + + +class ExpertMatch: + """`qa-expert.match`, leg `ref.answer`: clean when the final answer matches the reference or an alias after + `normalise`; `fail` `wrong_answer`, or `no_answer` when the reply has no `Answer:` line; unsettled when the task is + unknown or the reply cannot be read.""" + + version = "1" + legs = ("ref.answer",) + resource = "local_cpu" + + def __init__(self, ctx): + self.name, self.ctx = ctx.name, ctx + self.rows = {} + for pool in ctx.pools: + if (Path(pool) / "questions.jsonl").is_file(): + for row in load_rows(pool): + self.rows[row_sha256(row)] = row + + def run(self, job): + from aorlta.reward.adapters.base import make_result, unsettled + + t0 = self.ctx.clock() + try: + row = self.rows.get(job.task_sha256) + response = job.inputs.get("response") + if job.leg != "ref.answer" or row is None or not response: + return unsettled(job, "adapter_error", t_start=t0, t_end=self.ctx.clock()) + answer = final_answer(Path(response).read_text(encoding="utf-8").rsplit("", 1)[-1]) + if answer is None: + return make_result(job, "fail", "no_answer", t_start=t0, t_end=self.ctx.clock()) + # PLACEHOLDER: an exact match after normalising. A judge (a command or a service that reads the + # question, the reference and the answer) goes here; a judge that gives no verdict is unsettled. + wanted = {normalise(a) for a in [row["answer"], *row.get("aliases", [])]} + if normalise(answer) in wanted: + return make_result(job, "clean", t_start=t0, t_end=self.ctx.clock()) + return make_result(job, "fail", "wrong_answer", t_start=t0, t_end=self.ctx.clock()) + except Exception: # noqa: BLE001 - an adapter never raises + return unsettled(job, "adapter_error", t_start=t0, t_end=self.ctx.clock()) + + def controls(self): + return [] + + def health(self): + return {"name": self.name, "responsive": True, "questions": len(self.rows)} + + +class ExpertV1: + """`qa-expert.v1`, a PLACEHOLDER reward: +1.0 for a matching answer, a value of its own for each failure the check + names (`BASE`), then -0.05 x the share of the token budget used. An unsettled leg, or a reason `BASE` does not + name, leaves the sample unscored.""" + + name = "qa-expert.v1" + version = "1" + legs = ("ref.answer",) + # TODO(your vertical): replace these values with your real reward, and the check's reasons with your own (a + # judge's "partly right", say). Keep: + # - the total in [-1.55, 1.0], its parts exactly `base`, `broke_healthy` and `length`, adding up to it (the + # service recomputes it and leaves anything else unscored, with an alarm): so a base in [-1.5, 1.0]; + # - an unsettled or missing leg, or a reason you did not name, unscored: never clean, never a score; + # - each failure its own value, so the answers to one prompt get different rewards (GRPO learns only from the + # differences); + # - `claims_clean` true only when the leg is clean. + BASE = {"wrong_answer": -0.5, "no_answer": -1.0} + + def __init__(self, profile=None): + from aorlta.reward.rewardfn.base import code_sha256 + + self.settings = {} + self.code_sha256 = code_sha256(__file__) + + def _base(self, legs): + leg = legs.get("ref.answer") + if leg is None: + return "pending", "ref.answer" + if leg.state == "clean": + return "score", 1.0 + if leg.state == "fail" and leg.reason in self.BASE: + return "score", self.BASE[leg.reason] + return "unscored", f"ref.answer:{leg.reason or leg.state}" + + @staticmethod + def _finish(base, cand): + parts = {"base": float(base), "broke_healthy": 0.0, + "length": -0.05 * (cand.response_tokens / cand.max_new_tokens)} + return parts["base"] + parts["broke_healthy"] + parts["length"], parts + + def compute(self, legs, cand): + kind, value = self._base(legs) + return (None, {"reason": str(value)}) if kind != "score" else self._finish(value, cand) + + def decide(self, partial, cand): + from aorlta.reward.rewardfn.base import Decided, NeedMore, Unscored + + kind, value = self._base(partial) + if kind == "pending": + return NeedMore((value,)) + return Unscored(value) if kind == "unscored" else Decided(*self._finish(value, cand)) + + @staticmethod + def claims_clean(reward, parts): + return reward is not None and isinstance(parts, dict) and parts.get("base", 0.0) > 0.0 + + +def register() -> None: + from aorlta import plugins + + plugins.register_task(plugins.Task("qa-expert", records)) + plugins.register_parse_adapter("qa-expert.parse", ExpertParse) + plugins.register_feedback_template("qa-expert.feedback", ExpertFeedback) + plugins.register_profile("qa-expert", HERE / "qa-expert.yaml") + plugins.register_check_adapter("qa-expert.match", ExpertMatch) + plugins.register_reward_fn("qa-expert.v1", ExpertV1) + + +register() diff --git a/aorlta/examples/verticals/qa-expert/score.yaml b/aorlta/examples/verticals/qa-expert/score.yaml new file mode 100644 index 000000000..c957f2094 --- /dev/null +++ b/aorlta/examples/verticals/qa-expert/score.yaml @@ -0,0 +1,39 @@ +# Score a model on the qa-expert starter's 4 held-out placeholder questions, 4 samples each +# (orchestrator.mode: score). +# One GPU: one TokenSpeed engine on GPU 0; the answer check runs on the CPU. +# Site-specific paths, to edit here or to set in a site file (AORLTA_SITE): run.root, model.base_dir (and +# model.adapter_dir to score a trained adapter) and data.prompt_pools[0].path, which is this checkout's +# aorlta/examples/verticals/qa-expert/pool/heldout. +plugins: [examples/verticals/qa-expert/qa_expert.py] +run: + run_id: qa-expert-score + source: standin + root: /scratch/aorlta/runs +model: + base_dir: /data/hf/hub/models--Qwen--Qwen3.5-4B/snapshots/851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a + source: standin +data: + prompt_pools: + - path: /opt/aorlta/examples/verticals/qa-expert/pool/heldout + source: standin + task: qa-expert +rollout: + engines: [0] + K: 4 + P: 4 + max_repair_turns: 1 + thinking: "off" + temperature: 0.7 + max_new_tokens: 512 + compile_adapter: none + parse_adapter: qa-expert.parse + feedback_template: qa-expert.feedback + # few prompts, each with a short answer, can give a group K identical replies with a sound sampler + # (the engines' start-up self-test checks that), so only stop when most groups do + identical_groups_stop: 0.75 +reward: + profile: qa-expert + reward_fn: qa-expert.v1 +orchestrator: + mode: score + run_seed: 20261113 diff --git a/aorlta/examples/verticals/qa-expert/train.yaml b/aorlta/examples/verticals/qa-expert/train.yaml new file mode 100644 index 000000000..59a8f4c33 --- /dev/null +++ b/aorlta/examples/verticals/qa-expert/train.yaml @@ -0,0 +1,42 @@ +# Train on the qa-expert starter's 4 placeholder questions, 8 samples each, 3 iterations: a template to +# copy with your own pool, far too small to learn from. +# One GPU: one TokenSpeed engine on GPU 0 and one LoRA rank beside it; the answer check runs on the CPU. +# Site-specific paths, to edit here or to set in a site file (AORLTA_SITE): run.root, model.base_dir and +# data.prompt_pools[0].path, which is this checkout's aorlta/examples/verticals/qa-expert/pool/train. +plugins: [examples/verticals/qa-expert/qa_expert.py] +run: + run_id: qa-expert-train + source: standin + root: /scratch/aorlta/runs +model: + base_dir: /data/hf/hub/models--Qwen--Qwen3.5-4B/snapshots/851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a + source: standin +data: + prompt_pools: + - path: /opt/aorlta/examples/verticals/qa-expert/pool/train + source: standin + task: qa-expert +rollout: + engines: [0] + K: 8 + P: 4 + max_repair_turns: 1 + thinking: "off" + temperature: 1.0 + max_new_tokens: 512 + compile_adapter: none + parse_adapter: qa-expert.parse + feedback_template: qa-expert.feedback + # few prompts, each with a short answer, can give a group K identical replies with a sound sampler + # (the engines' start-up self-test checks that), so only stop when most groups do + identical_groups_stop: 0.75 +reward: + profile: qa-expert + reward_fn: qa-expert.v1 +orchestrator: + iterations: 3 + run_seed: 20261113 +trainer: + ranks: 1 + micro_batch_tokens: 20480 + lr: 2.0e-5 diff --git a/aorlta/examples/verticals/sanitizer-triage/README.md b/aorlta/examples/verticals/sanitizer-triage/README.md new file mode 100644 index 000000000..2f1f36a4b --- /dev/null +++ b/aorlta/examples/verticals/sanitizer-triage/README.md @@ -0,0 +1,22 @@ +# sanitizer-triage: a starter for a sanitizer or NaN triage task + +The model reads a tool's report and answers with a structured verdict: a JSON object with a category from a closed +list and a short rationale. Everything here is a PLACEHOLDER: the reports are toy text written for this starter, not +any tool's output, and the check looks only at the verdict's form. It is a plugin +([docs/ADAPTING.md](../../../docs/ADAPTING.md)) and imports nothing from the other examples, so the directory can be +copied on its own. + +| file | what | what to replace | +|---|---|---| +| `sanitizer_triage.py` | the plugin: prompts, parse and feedback, the `ref.verdict` check, the `sanitizer-triage.v1` reward | `CATEGORIES` and `USER` (your verdicts), `TriageCheck.run` (compare with the report's `label`, or run the real tool: its TODO), `TriageV1.BASE` (your reward: its TODO lists the rules) | +| `sanitizer-triage.yaml` | the check profile: one CPU leg | your check's settings, if it has any | +| `pool/{train,heldout}/reports.jsonl` | 4 + 4 toy reports, `{"id", "report", "label"}` | your reports: many more than `rollout.P` | +| `score.yaml`, `train.yaml` | score the held-out set, train on the rest; one GPU | the site paths | + +From `aorlta/`, with the site paths set (`AORLTA_SITE`, docs/TUTORIAL.md section 4) and the config committed: + +```bash +OMP_NUM_THREADS=1 python -m pytest -q tests/examples/test_verticals.py -k triage # on CPU, a score and a training run +CONFIG=examples/verticals/sanitizer-triage/score.yaml sbatch jobs/aorlta-run.sbatch # score the base model +CONFIG=examples/verticals/sanitizer-triage/train.yaml sbatch jobs/aorlta-run.sbatch # train, then score the adapter +``` diff --git a/aorlta/examples/verticals/sanitizer-triage/pool/heldout/reports.jsonl b/aorlta/examples/verticals/sanitizer-triage/pool/heldout/reports.jsonl new file mode 100644 index 000000000..e755b3b4e --- /dev/null +++ b/aorlta/examples/verticals/sanitizer-triage/pool/heldout/reports.jsonl @@ -0,0 +1,4 @@ +{"id": "placeholder-h1", "report": "[placeholder report, synthetic] toy-checker 0.0: toy_case_h1 did not finish.", "label": "inconclusive"} +{"id": "placeholder-h2", "report": "[placeholder report, synthetic] toy-checker 0.0: toy_case_h2 finished, 0 findings.", "label": "no_issue"} +{"id": "placeholder-h3", "report": "[placeholder report, synthetic] toy-checker 0.0: toy_case_h3 finished, 1 finding of kind TOY-RACE.", "label": "race"} +{"id": "placeholder-h4", "report": "[placeholder report, synthetic] toy-checker 0.0: toy_case_h4 finished, 1 finding of kind TOY-NAN.", "label": "numeric"} diff --git a/aorlta/examples/verticals/sanitizer-triage/pool/train/reports.jsonl b/aorlta/examples/verticals/sanitizer-triage/pool/train/reports.jsonl new file mode 100644 index 000000000..219ccc726 --- /dev/null +++ b/aorlta/examples/verticals/sanitizer-triage/pool/train/reports.jsonl @@ -0,0 +1,4 @@ +{"id": "placeholder-1", "report": "[placeholder report, synthetic] toy-checker 0.0: toy_case_1 finished, 0 findings.", "label": "no_issue"} +{"id": "placeholder-2", "report": "[placeholder report, synthetic] toy-checker 0.0: toy_case_2 finished, 1 finding of kind TOY-MEMORY.", "label": "memory"} +{"id": "placeholder-3", "report": "[placeholder report, synthetic] toy-checker 0.0: toy_case_3 finished, 1 finding of kind TOY-RACE.", "label": "race"} +{"id": "placeholder-4", "report": "[placeholder report, synthetic] toy-checker 0.0: toy_case_4 finished, 1 finding of kind TOY-NAN.", "label": "numeric"} diff --git a/aorlta/examples/verticals/sanitizer-triage/sanitizer-triage.yaml b/aorlta/examples/verticals/sanitizer-triage/sanitizer-triage.yaml new file mode 100644 index 000000000..98b73702b --- /dev/null +++ b/aorlta/examples/verticals/sanitizer-triage/sanitizer-triage.yaml @@ -0,0 +1,12 @@ +# `sanitizer-triage`: one leg, `ref.verdict`, on the CPU: the reply's last JSON object against the verdict's form +# (examples/verticals/sanitizer-triage/sanitizer_triage.py). No compile and no sanitizer leg. +name: sanitizer-triage +reward_fn: sanitizer-triage.v1 +clean_repeats: 1 +materialize: true +pools: + check: {resource: local_cpu, adapter: sanitizer-triage.check, timeout: gate, workers: 4} +stages: + - {name: S0, pool: check, legs: [ref.verdict]} +adapters: + sanitizer-triage.check: {} diff --git a/aorlta/examples/verticals/sanitizer-triage/sanitizer_triage.py b/aorlta/examples/verticals/sanitizer-triage/sanitizer_triage.py new file mode 100644 index 000000000..3c8215648 --- /dev/null +++ b/aorlta/examples/verticals/sanitizer-triage/sanitizer_triage.py @@ -0,0 +1,252 @@ +"""A starter for a sanitizer- or NaN-triage task (examples/verticals/README.md): the model reads a tool's report and +answers with a structured verdict. Every part marked PLACEHOLDER is to be replaced with your own. + +A pool is a directory holding `reports.jsonl`, one report per line: `{"id", "report", "label"}`. The reports that come +with it are toy text written for this starter, not the output of any tool. The model is shown the report and asked +for a JSON object `{"category": ..., "rationale": ...}`, with the category one of `CATEGORIES`. A reply with no JSON +object gets one repair turn that says so (this module's parse and feedback); there is no compile step +(`rollout.compile_adapter: none`). The reward service runs one leg, `ref.verdict`, on the CPU; as a placeholder it +checks only that the verdict has exactly the two keys, a category from the list and a short rationale. It does not +yet compare the category with the report's `label`: that, or running the real tool, is the check to write. + +Registered names: task `sanitizer-triage`; parse `sanitizer-triage.parse`; feedback `sanitizer-triage.feedback`; +profile `sanitizer-triage` (sanitizer-triage.yaml); check adapter `sanitizer-triage.check`; reward function +`sanitizer-triage.v1`. +""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +HERE = Path(__file__).resolve().parent +# PLACEHOLDER: the closed list of verdicts, and the prompt. Use your own tool's categories. +CATEGORIES = ("no_issue", "memory", "race", "numeric", "inconclusive") +KEYS = ("category", "rationale") +MAX_RATIONALE_CHARS = 400 +SYSTEM = "You triage reports from a checking tool." +USER = ("Report:\n{report}\n\nReply with one JSON object with exactly the keys \"category\" and \"rationale\": the " + "category is one of " + ", ".join(CATEGORIES) + ", and the rationale is one or two sentences.") + + +def load_rows(pool_dir) -> list[dict]: + rows = [json.loads(line) for line in (Path(pool_dir) / "reports.jsonl").read_text(encoding="utf-8").splitlines() + if line.strip()] + ids = [r["id"] for r in rows] + if len(set(ids)) != len(ids): + raise ValueError(f"{pool_dir}/reports.jsonl lists an id twice") + return rows + + +def row_sha256(row: dict) -> str: + return hashlib.sha256(json.dumps(row, sort_keys=True).encode("utf-8")).hexdigest() + + +def final_verdict(reply: str) -> dict | None: + """The last JSON object in the reply that parses, or None. Objects are found by the JSON decoder, so a brace inside + a string (a rationale that quotes `}`) does not end one.""" + dec, last, i = json.JSONDecoder(), None, 0 + while (i := reply.find("{", i)) != -1: + try: + obj, end = dec.raw_decode(reply, i) + except ValueError: + i += 1 + continue + if isinstance(obj, dict): + last = obj + i = end + return last + + +def verdict_problem(verdict: dict) -> str | None: + """PLACEHOLDER check: None for a verdict with exactly `KEYS`, a category from `CATEGORIES` and a non-empty + rationale of at most `MAX_RATIONALE_CHARS`; otherwise the failure's reason (`bad_keys`, `bad_category`, + `bad_rationale`).""" + if set(verdict) != set(KEYS): + return "bad_keys" + if verdict["category"] not in CATEGORIES: + return "bad_category" + r = verdict["rationale"] + if not (isinstance(r, str) and r.strip() and len(r) <= MAX_RATIONALE_CHARS): + return "bad_rationale" + return None + + +def messages(row: dict) -> list[dict]: + return [{"role": "system", "content": SYSTEM}, {"role": "user", "content": USER.format(report=row["report"])}] + + +def records(pool_dir, tokenizer, thinking): + from aorlta.ids import sha256_ids + from aorlta.rollout.templates import render_turn0 + from aorlta.schemas import PromptRecord, TaskRef + + out = [] + for row in load_rows(pool_dir): + msgs = messages(row) + out.append(PromptRecord( + prompt_id=f"sanitizer-triage:{row['id']}", turn0_ids=None, messages=msgs, + prompt_ids_sha256=sha256_ids(render_turn0(tokenizer, msgs, thinking)), + task=TaskRef(contract_id=f"sanitizer-triage:{row['id']}", parent_id=None, parent_healthy=None, + reference_kind=None, task_sha256=row_sha256(row)), + source="standin")) + return out + + +class TriageParse: + """`sanitizer-triage.parse`: a reply with a JSON object is a verdict (nothing to compile); one without is rejected + with a message, which `sanitizer-triage.feedback` turns into a repair turn. Whether the verdict is well formed or + right is the check's to say, not the parse's.""" + + name = "sanitizer-triage.parse" + + def __init__(self): + from aorlta.rollout.templates import IM_END_ID + + self.reply_stops = (IM_END_ID,) + + def check_prompts(self, prompts, thinking, max_repair_turns=0): + return None + + def check(self, text, thinking, prompt): + from aorlta.rollout.compile_check import ParseOutcome + + reply = text.rsplit("", 1)[-1] if thinking == "on" else text + verdict = final_verdict(reply) + if verdict is None: + return ParseOutcome("rejected", rejection=( + "Your reply has no JSON object. Reply with one JSON object with the keys \"category\" and " + "\"rationale\".", "no_verdict")) + return ParseOutcome("edits", code=json.dumps(verdict, sort_keys=True)) + + +class TriageFeedback: + name = "sanitizer-triage.feedback" + + def build(self, errors, max_chars): + from aorlta.rollout.feedback import build_feedback + + return build_feedback(errors, max_chars) + + def build_apply(self, message, reason): + from aorlta.rollout.feedback import Feedback + + return Feedback(message, 0, 0) + + +class TriageCheck: + """`sanitizer-triage.check`, leg `ref.verdict`: `fail` `no_verdict` for a reply with no JSON object, `fail` with + `verdict_problem`'s reason for a verdict it rejects, clean otherwise; unsettled when the task is unknown or the + reply cannot be read.""" + + version = "1" + legs = ("ref.verdict",) + resource = "local_cpu" + + def __init__(self, ctx): + self.name, self.ctx = ctx.name, ctx + self.rows = {} + for pool in ctx.pools: + if (Path(pool) / "reports.jsonl").is_file(): + for row in load_rows(pool): + self.rows[row_sha256(row)] = row + + def run(self, job): + from aorlta.reward.adapters.base import make_result, unsettled + + t0 = self.ctx.clock() + try: + row = self.rows.get(job.task_sha256) + response = job.inputs.get("response") + if job.leg != "ref.verdict" or row is None or not response: + return unsettled(job, "adapter_error", t_start=t0, t_end=self.ctx.clock()) + verdict = final_verdict(Path(response).read_text(encoding="utf-8").rsplit("", 1)[-1]) + if verdict is None: + return make_result(job, "fail", "no_verdict", t_start=t0, t_end=self.ctx.clock()) + problem = verdict_problem(verdict) + if problem is not None: + return make_result(job, "fail", problem, t_start=t0, t_end=self.ctx.clock()) + # TODO(your vertical): the placeholder stops at the verdict's form. Compare `verdict["category"]` with + # `row["label"]` (a `fail` `wrong_category` of its own), or run the real tool on the case and read its + # result; a tool that gives no result is unsettled, never clean. + return make_result(job, "clean", t_start=t0, t_end=self.ctx.clock()) + except Exception: # noqa: BLE001 - an adapter never raises + return unsettled(job, "adapter_error", t_start=t0, t_end=self.ctx.clock()) + + def controls(self): + return [] + + def health(self): + return {"name": self.name, "responsive": True, "reports": len(self.rows)} + + +class TriageV1: + """`sanitizer-triage.v1`, a PLACEHOLDER reward: +1.0 for a verdict the check passes, a value of its own for each + failure it names (`BASE`), then -0.05 x the share of the token budget used. An unsettled leg, or a reason `BASE` + does not name, leaves the sample unscored.""" + + name = "sanitizer-triage.v1" + version = "1" + legs = ("ref.verdict",) + # TODO(your vertical): replace these values with your real reward (a wrong category from the check above, say, + # nearer +1.0 than a malformed verdict). Keep: + # - the total in [-1.55, 1.0], its parts exactly `base`, `broke_healthy` and `length`, adding up to it (the + # service recomputes it and leaves anything else unscored, with an alarm): so a base in [-1.5, 1.0]; + # - an unsettled or missing leg, or a reason you did not name, unscored: never clean, never a score; + # - each failure its own value, so the answers to one prompt get different rewards (GRPO learns only from the + # differences); + # - `claims_clean` true only when the leg is clean. + BASE = {"bad_rationale": -0.4, "bad_category": -0.6, "bad_keys": -0.8, "no_verdict": -1.0} + + def __init__(self, profile=None): + from aorlta.reward.rewardfn.base import code_sha256 + + self.settings = {} + self.code_sha256 = code_sha256(__file__) + + def _base(self, legs): + leg = legs.get("ref.verdict") + if leg is None: + return "pending", "ref.verdict" + if leg.state == "clean": + return "score", 1.0 + if leg.state == "fail" and leg.reason in self.BASE: + return "score", self.BASE[leg.reason] + return "unscored", f"ref.verdict:{leg.reason or leg.state}" + + @staticmethod + def _finish(base, cand): + parts = {"base": float(base), "broke_healthy": 0.0, + "length": -0.05 * (cand.response_tokens / cand.max_new_tokens)} + return parts["base"] + parts["broke_healthy"] + parts["length"], parts + + def compute(self, legs, cand): + kind, value = self._base(legs) + return (None, {"reason": str(value)}) if kind != "score" else self._finish(value, cand) + + def decide(self, partial, cand): + from aorlta.reward.rewardfn.base import Decided, NeedMore, Unscored + + kind, value = self._base(partial) + if kind == "pending": + return NeedMore((value,)) + return Unscored(value) if kind == "unscored" else Decided(*self._finish(value, cand)) + + @staticmethod + def claims_clean(reward, parts): + return reward is not None and isinstance(parts, dict) and parts.get("base", 0.0) > 0.0 + + +def register() -> None: + from aorlta import plugins + + plugins.register_task(plugins.Task("sanitizer-triage", records)) + plugins.register_parse_adapter("sanitizer-triage.parse", TriageParse) + plugins.register_feedback_template("sanitizer-triage.feedback", TriageFeedback) + plugins.register_profile("sanitizer-triage", HERE / "sanitizer-triage.yaml") + plugins.register_check_adapter("sanitizer-triage.check", TriageCheck) + plugins.register_reward_fn("sanitizer-triage.v1", TriageV1) + + +register() diff --git a/aorlta/examples/verticals/sanitizer-triage/score.yaml b/aorlta/examples/verticals/sanitizer-triage/score.yaml new file mode 100644 index 000000000..4a07589ca --- /dev/null +++ b/aorlta/examples/verticals/sanitizer-triage/score.yaml @@ -0,0 +1,39 @@ +# Score a model on the sanitizer-triage starter's 4 held-out placeholder reports, 4 samples each +# (orchestrator.mode: score). +# One GPU: one TokenSpeed engine on GPU 0; the verdict check runs on the CPU. +# Site-specific paths, to edit here or to set in a site file (AORLTA_SITE): run.root, model.base_dir (and +# model.adapter_dir to score a trained adapter) and data.prompt_pools[0].path, which is this checkout's +# aorlta/examples/verticals/sanitizer-triage/pool/heldout. +plugins: [examples/verticals/sanitizer-triage/sanitizer_triage.py] +run: + run_id: sanitizer-triage-score + source: standin + root: /scratch/aorlta/runs +model: + base_dir: /data/hf/hub/models--Qwen--Qwen3.5-4B/snapshots/851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a + source: standin +data: + prompt_pools: + - path: /opt/aorlta/examples/verticals/sanitizer-triage/pool/heldout + source: standin + task: sanitizer-triage +rollout: + engines: [0] + K: 4 + P: 4 + max_repair_turns: 1 + thinking: "off" + temperature: 0.7 + max_new_tokens: 512 + compile_adapter: none + parse_adapter: sanitizer-triage.parse + feedback_template: sanitizer-triage.feedback + # few prompts, each with a short answer, can give a group K identical replies with a sound sampler + # (the engines' start-up self-test checks that), so only stop when most groups do + identical_groups_stop: 0.75 +reward: + profile: sanitizer-triage + reward_fn: sanitizer-triage.v1 +orchestrator: + mode: score + run_seed: 20261113 diff --git a/aorlta/examples/verticals/sanitizer-triage/train.yaml b/aorlta/examples/verticals/sanitizer-triage/train.yaml new file mode 100644 index 000000000..873da6f30 --- /dev/null +++ b/aorlta/examples/verticals/sanitizer-triage/train.yaml @@ -0,0 +1,42 @@ +# Train on the sanitizer-triage starter's 4 placeholder reports, 8 samples each, 3 iterations: a template to +# copy with your own pool, far too small to learn from. +# One GPU: one TokenSpeed engine on GPU 0 and one LoRA rank beside it; the verdict check runs on the CPU. +# Site-specific paths, to edit here or to set in a site file (AORLTA_SITE): run.root, model.base_dir and +# data.prompt_pools[0].path, which is this checkout's aorlta/examples/verticals/sanitizer-triage/pool/train. +plugins: [examples/verticals/sanitizer-triage/sanitizer_triage.py] +run: + run_id: sanitizer-triage-train + source: standin + root: /scratch/aorlta/runs +model: + base_dir: /data/hf/hub/models--Qwen--Qwen3.5-4B/snapshots/851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a + source: standin +data: + prompt_pools: + - path: /opt/aorlta/examples/verticals/sanitizer-triage/pool/train + source: standin + task: sanitizer-triage +rollout: + engines: [0] + K: 8 + P: 4 + max_repair_turns: 1 + thinking: "off" + temperature: 1.0 + max_new_tokens: 512 + compile_adapter: none + parse_adapter: sanitizer-triage.parse + feedback_template: sanitizer-triage.feedback + # few prompts, each with a short answer, can give a group K identical replies with a sound sampler + # (the engines' start-up self-test checks that), so only stop when most groups do + identical_groups_stop: 0.75 +reward: + profile: sanitizer-triage + reward_fn: sanitizer-triage.v1 +orchestrator: + iterations: 3 + run_seed: 20261113 +trainer: + ranks: 1 + micro_batch_tokens: 20480 + lr: 2.0e-5 diff --git a/aorlta/tests/examples/test_verticals.py b/aorlta/tests/examples/test_verticals.py new file mode 100644 index 000000000..18d9785de --- /dev/null +++ b/aorlta/tests/examples/test_verticals.py @@ -0,0 +1,431 @@ +"""The vertical starters (examples/verticals): three plugin packages, each on its own. Their configs load with the +plugin, their pools keep the held-out prompts out of training, their parse gives a repair turn to a reply with nothing +to check, their placeholder check gives each failure its own reason and anything it cannot read is unsettled, their +placeholder reward keeps the service's invariants and leaves the unsettled unscored, and each scores and trains through +the real wiring on CPU with test_wiring_e2e's tiny model and engines, the engines' replies drawn from a few written for +the vertical.""" + +from __future__ import annotations + +import dataclasses +import json +import re +import stat +import sys +from dataclasses import dataclass +from pathlib import Path + +import pytest + +from aorlta import plugins +from aorlta.config import config_from_dict, load_config, parse_yaml +from aorlta.episodes.schema import LEG_RE + +ROOT = Path(__file__).resolve().parents[2] +VERTICALS = ROOT / "examples" / "verticals" +sys.path.insert(0, str(ROOT / "tests")) +sys.path.insert(0, str(ROOT / "tests" / "orchestrator")) +import test_wiring_e2e as E # noqa: E402 - the tiny model and its engines +from test_plugins import CharChat, registry # noqa: E402,F401 + +tiny_base = E.tiny_base + + +@dataclass(frozen=True) +class Vertical: + name: str # the directory, the task, the profile and every component's prefix + module: str # the plugin file, beside its profile and configs + pool_file: str + leg: str + check: str # the check adapter's name + good: str # a reply the placeholder check passes, for the first training prompt + bad: tuple[tuple[str, str], ...] # (reply, fail reason) for the first training prompt; each reason at least once + nothing: str # a reply with nothing to check: the parse's repair turn, and the check's first reason + + @property + def plugin(self) -> str: + return f"examples/verticals/{self.name}/{self.module}" + + @property + def reasons(self) -> set[str]: + return {r for _, r in self.bad} + + +ISA = Vertical( + "isa-fuzzer", "isa_fuzzer.py", "tasks.jsonl", "ref.program", "isa-fuzzer.check", + good="Here it is.\n```toy-asm\nMOV r1, #3\nADD r2, r1, r1\n```", + bad=(("No block here, sorry.", "no_program"), ("```\nMOV r1, #3\nADD r2, r1\n```", "malformed"), + ("```\nMOV r1, #3\nJMP r2\n```", "malformed"), ("```\nMOV r9, #3\nADD r2, r1, r1\n```", "malformed"), + ("```\nNOP r1\nADD r2, r1, r1\n```", "malformed"), + ("```\nADD r1, r1, #1\nADD r1, r1, #1\nADD r1, r1, #1\nADD r1, r1, #1\nADD r1, r1, #1\n```", "wrong_length"), + ("```\nMOV r1, #3\nSUB r2, r1, r1\n```", "missing_op")), + nothing="I would add two registers.") +QA = Vertical( + "qa-expert", "qa_expert.py", "questions.jsonl", "ref.answer", "qa-expert.match", + good="Two and three make five.\n**Answer:** `Five.`", + bad=(("It is five.", "no_answer"), ("Answer: 5\nOn second thought:\nAnswer: 6", "wrong_answer")), + nothing="It is five.") +TRIAGE = Vertical( + "sanitizer-triage", "sanitizer_triage.py", "reports.jsonl", "ref.verdict", "sanitizer-triage.check", + good='Looks clean.\n```json\n{"category": "no_issue", "rationale": "The report lists no findings."}\n```', + bad=(("No findings, so it is fine.", "no_verdict"), ('{"category": "no_issue"}', "bad_keys"), + ('{"category": "fine", "rationale": "No findings."}', "bad_category"), + ('{"category": "no_issue", "rationale": "Ok."}\nNo\n{"category": "fine", "rationale": "Ok."}', "bad_category"), + ('{"category": "no_issue", "rationale": "No findings.", "extra": 1}', "bad_keys"), + ('{"category": "no_issue", "rationale": " "}', "bad_rationale"), + ('{"category": "no_issue", "rationale": "%s"}' % ("x" * 401), "bad_rationale")), + nothing="No findings, so it is fine.") +ALL = (ISA, QA, TRIAGE) +BY_NAME = pytest.mark.parametrize("v", ALL, ids=lambda v: v.name) + + +@pytest.fixture +def load(registry): # noqa: F811 + """Load a vertical's plugin as a run config loads it; its module.""" + def _load(v): + plugins.load_plugins([v.plugin]) + return sys.modules[plugins.identity()[v.plugin]["module"]] + return _load + + +def site_config(tmp_path, v, kind): + site = tmp_path / "site.yaml" + split = "heldout" if kind == "score" else "train" + site.write_text(json.dumps({"run": {"root": str(tmp_path / "runs")}, "model": {"base_dir": str(tmp_path / "base")}, + "data": {"prompt_pools": [str(VERTICALS / v.name / "pool" / split)]}})) + return load_config(VERTICALS / v.name / f"{kind}.yaml", site=site) + + +@BY_NAME +@pytest.mark.parametrize("kind", ["score", "train"]) +def test_the_configs_load_with_the_plugin_and_run_on_one_gpu_with_no_compile(tmp_path, registry, v, kind): # noqa: F811 + from aorlta.orchestrator import wiring + + cfg = site_config(tmp_path, v, kind) + assert cfg.plugins == (v.plugin,) and cfg.data.prompt_pools[0].task == v.name + assert (cfg.reward.profile, cfg.reward.reward_fn) == (v.name, f"{v.name}.v1") + assert (cfg.rollout.parse_adapter, cfg.rollout.feedback_template) == (f"{v.name}.parse", f"{v.name}.feedback") + assert cfg.rollout.compile_adapter == "none" and cfg.rollout.engines == (0,) + assert cfg.orchestrator.mode == kind if kind == "score" else cfg.trainer.ranks == 1 + wiring.check_supported(cfg) + + +@BY_NAME +def test_the_profile_has_one_leg_and_no_compile_or_sanitizer(load, v): + from aorlta.reward.profiles import load_profile + + load(v) + profile = load_profile(v.name) + assert profile.legs == (v.leg,) and {p.adapter for p in profile.pools.values()} == {v.check} + assert LEG_RE.fullmatch(v.leg) # an episode record names its legs gate.*, ref.* or san.* + + +@BY_NAME +def test_the_pools_are_placeholders_and_the_held_out_prompts_are_not_trained_on(load, v): + mod = load(v) + train, held = (mod.load_rows(VERTICALS / v.name / "pool" / s) for s in ("train", "heldout")) + assert len(train) == len(held) == 4 # the self-test and a score batch need 4 prompts + assert not {r["id"] for r in train} & {r["id"] for r in held} + assert not {mod.row_sha256(r) for r in train} & {mod.row_sha256(r) for r in held} + assert all("placeholder" in json.dumps(r).lower() for r in train + held) + + +@BY_NAME +def test_the_records_name_the_task_and_hash_what_the_check_reads(load, v): + from aorlta.ids import sha256_ids + from aorlta.rollout.templates import render_turn0 + + mod = load(v) + pool = VERTICALS / v.name / "pool" / "train" + recs = plugins.task(v.name).records(pool, CharChat(), "off") + rows = mod.load_rows(pool) + assert [r.prompt_id for r in recs] == [f"{v.name}:{row['id']}" for row in rows] + assert [r.task.task_sha256 for r in recs] == [mod.row_sha256(row) for row in rows] + for key in rows[0]: # every field the check reads moves the hash, so a cached result is never another task's + assert mod.row_sha256({**rows[0], key: "changed"}) != mod.row_sha256(rows[0]), key + assert all(r.prompt_ids_sha256 == sha256_ids(render_turn0(CharChat(), list(r.messages), "off")) for r in recs) + assert len({r.messages[0]["content"] for r in recs}) == 1 # one system message, as the probes need + + +@BY_NAME +def test_the_parse_gives_a_repair_turn_to_a_reply_with_nothing_to_check(load, v): + from aorlta.rollout.chain import make_feedback_template, make_parse_adapter + + load(v) + parse, fb = make_parse_adapter(f"{v.name}.parse"), make_feedback_template(f"{v.name}.feedback") + out = parse.check(v.nothing, "off", None) + assert out.kind == "rejected" and out.rejection[1] == v.bad[0][1] + assert fb.build_apply(*out.rejection).text == out.rejection[0] + assert parse.check(v.good, "off", None).kind == "edits" + assert parse.check(f"{v.nothing}\n{v.good}", "on", None).kind == "edits" + assert parse.check(f"{v.good}\n{v.nothing}", "on", None).kind == "rejected" + for reply, _ in v.bad[1:]: # a reply that is wrong but present is the check's to score, not the parse's + assert parse.check(reply, "off", None).kind == "edits" + + +def job_for(mod, tmp_path, row, reply, leg): + from aorlta.schemas import LegJob + + resp = tmp_path / "response.txt" + resp.write_text(reply) + return LegJob(job_id="j", leg=leg, candidate_id="c", control_id=None, source_sha256="s" * 64, + task_sha256=mod.row_sha256(row), inputs={"response": str(resp)}, output_dir=str(tmp_path)) + + +def adapter(v, settings=None): + ctx = plugins.CheckContext(name=v.check, worker_id="w/1", settings=settings or {}, clock=lambda: 0.0, + timeout_s=30.0, pools=(str(VERTICALS / v.name / "pool" / "train"),)) + return plugins.CHECK_ADAPTERS[v.check](ctx) + + +@BY_NAME +def test_the_check_gives_each_failure_its_own_reason(load, tmp_path, v): + mod = load(v) + row = mod.load_rows(VERTICALS / v.name / "pool" / "train")[0] + check = adapter(v) + assert (check.run(job_for(mod, tmp_path, row, v.good, v.leg)).state, check.legs) == ("clean", (v.leg,)) + for reply, reason in v.bad: + res = check.run(job_for(mod, tmp_path, row, reply, v.leg)) + assert (res.state, res.reason) == ("fail", reason), reply + + +@BY_NAME +def test_what_the_check_cannot_read_is_unsettled_never_clean(load, tmp_path, v): + mod = load(v) + row = mod.load_rows(VERTICALS / v.name / "pool" / "train")[0] + check = adapter(v) + unknown = dict(row, id="not-in-any-pool") + held = mod.load_rows(VERTICALS / v.name / "pool" / "heldout")[0] # a pool the check was not given + jobs = [job_for(mod, tmp_path, unknown, v.good, v.leg), job_for(mod, tmp_path, held, v.good, v.leg), + job_for(mod, tmp_path, row, v.good, "some.other.leg")] + good = job_for(mod, tmp_path, row, v.good, v.leg) + jobs += [dataclasses.replace(good, inputs={}), + dataclasses.replace(good, inputs={"response": str(tmp_path / "missing.txt")})] + for job in jobs: + res = check.run(job) + assert (res.state, res.reason) == ("unsettled", "adapter_error"), job + + +def tool(tmp_path, out: str, rc: int = 0) -> list[str]: + t = tmp_path / "tool" + t.write_text(f"#!/bin/bash\ncat > {tmp_path}/tool-input.json\necho '{out}'\nexit {rc}\n") + t.chmod(t.stat().st_mode | stat.S_IXUSR) + return [str(t)] + + +@pytest.mark.parametrize("out, rc, want", [ + ('{"verdict": "pass"}', 0, ("clean", None)), ('{"verdict": "fail"}', 0, ("fail", "tool_rejected")), + ("not json", 0, ("unsettled", "tool_no_verdict")), ('{"verdict": "maybe"}', 0, ("unsettled", "tool_no_verdict")), + ('{"verdict": "pass"}', 3, ("unsettled", "tool_no_verdict")), +]) +def test_the_isa_tool_decides_a_well_formed_program_and_no_verdict_is_unsettled(load, tmp_path, out, rc, want): + mod = load(ISA) + row = mod.load_rows(VERTICALS / ISA.name / "pool" / "train")[0] + check = adapter(ISA, {"tool": tool(tmp_path, out, rc)}) + res = check.run(job_for(mod, tmp_path, row, ISA.good, ISA.leg)) + assert (res.state, res.reason) == want + assert json.loads((tmp_path / "tool-input.json").read_text()) == {"goal": row["goal"], + "program": "MOV r1, #3\nADD r2, r1, r1\n"} + (tmp_path / "tool-input.json").unlink() + malformed = check.run(job_for(mod, tmp_path, row, ISA.bad[1][0], ISA.leg)) + assert malformed.reason == "malformed" and not (tmp_path / "tool-input.json").exists() # the tool sees only these + + +def test_the_isa_tool_is_part_of_the_checkers_identity(load, tmp_path): + """Cached legs are keyed by the checker's version: the tool's command, timeout and script are part of it.""" + load(ISA) + t = tool(tmp_path, '{"verdict": "pass"}') + versions = [adapter(ISA).version, adapter(ISA, {"tool": t}).version, adapter(ISA, {"tool": t + ["-x"]}).version, + adapter(ISA, {"tool": t, "tool_timeout_s": 5}).version] + assert len(set(versions)) == 4 and versions[0] == "1" and adapter(ISA, {"tool": t}).version == versions[1] + Path(t[0]).write_text("#!/bin/bash\necho '{\"verdict\": \"fail\"}'\n") + assert adapter(ISA, {"tool": t}).version != versions[1] + + +def test_an_isa_tool_that_times_out_or_is_retired_leaves_nothing_running(load, tmp_path): + import os + import threading + import time + + mod = load(ISA) + row = mod.load_rows(VERTICALS / ISA.name / "pool" / "train")[0] + for how in ("timeout", "close"): + pid = tmp_path / f"{how}.pid" + t = tmp_path / f"tool-{how}" + t.write_text(f"#!/bin/bash\nsleep 60 &\necho $! > {pid}\nwait\n") + t.chmod(0o755) + check = adapter(ISA, {"tool": [str(t)], "tool_timeout_s": 1.0 if how == "timeout" else 60.0}) + if how == "close": + threading.Thread(target=lambda: (time.sleep(1.0), check.close()), daemon=True).start() + t0 = time.monotonic() + res = check.run(job_for(mod, tmp_path, row, ISA.good, ISA.leg)) + assert (res.state, res.reason) == ("unsettled", "tool_no_verdict") and time.monotonic() - t0 < 15, how + child = int(pid.read_text()) + deadline = time.monotonic() + 5 + while os.path.exists(f"/proc/{child}") and time.monotonic() < deadline: + time.sleep(0.05) + assert not os.path.exists(f"/proc/{child}"), how + check.close() + + +def test_a_triage_verdict_whose_rationale_quotes_a_brace_is_read(load): + mod = load(TRIAGE) + want = {"category": "no_issue", "rationale": "The report contains } but no finding."} + reply = 'An earlier draft: {"category": "race"}\nFinal:\n```json\n' + json.dumps(want) + "\n```" + assert mod.final_verdict(reply) == want and mod.verdict_problem(want) is None + assert mod.final_verdict('{"rationale": "a { and an escaped \\" quote", "category": "memory"}')["category"] == "memory" + assert mod.final_verdict("no object { here") is None + + +def test_the_qa_answer_is_the_last_answer_lines_own_value(load): + """Each line is matched alone: an empty marker does not take the next line's text, and an empty last marker is no + answer rather than an earlier one.""" + mod = load(QA) + assert mod.final_answer("Answer:\nAnswer: 5") == "5" + assert mod.final_answer("Answer: 5\nAnswer:") is None + assert mod.final_answer("Answer: 4\nthinking...\nAnswer: 5") == "5" and mod.final_answer("no marker") is None + + +def test_the_isa_prompt_states_the_immediates_the_check_accepts(load): + """The check's six-digit limit on an immediate is in the model's prompt, not only in the checker.""" + mod = load(ISA) + assert "at most six digits, which may be negative" in mod.USER + row = {"min_lines": 2, "max_lines": 4, "must_use": ["ADD"]} + for imm, ok in (("#999999", True), ("#-999999", True), ("#1000000", False)): + problem = mod.program_problem(f"MOV r1, {imm}\nADD r2, r1, r1\n", row) + assert (problem is None) == ok, (imm, problem) + + +def test_a_tool_found_through_path_is_part_of_the_checkers_identity(load, tmp_path, monkeypatch): + """A bare command name is resolved as PATH finds it: replacing that executable, or which one PATH selects, is a new + checker version.""" + load(ISA) + bins = [tmp_path / "a", tmp_path / "b"] + for d in bins: + d.mkdir() + (d / "my-checker").write_text("#!/bin/bash\necho '{\"verdict\": \"pass\"}'\n") + (d / "my-checker").chmod(0o755) + monkeypatch.setenv("PATH", f"{bins[0]}:/usr/bin:/bin") + first = adapter(ISA, {"tool": ["my-checker"]}).version + assert adapter(ISA, {"tool": ["my-checker"]}).version == first + (bins[0] / "my-checker").write_text("#!/bin/bash\necho '{\"verdict\": \"fail\"}'\n") + second = adapter(ISA, {"tool": ["my-checker"]}).version + monkeypatch.setenv("PATH", f"{bins[1]}:/usr/bin:/bin") + third = adapter(ISA, {"tool": ["my-checker"]}).version + assert len({first, second, third}) == 3 + + +@BY_NAME +def test_the_reward_keeps_the_services_invariants_and_leaves_the_unsettled_unscored(load, v): + from aorlta.reward.service import REWARD_MAX, REWARD_MIN + from aorlta.schemas import Candidate, LegResult, TaskRef + + load(v) + fn = plugins.REWARD_FNS[f"{v.name}.v1"](None) + assert fn.legs == (v.leg,) and fn.name == f"{v.name}.v1" and len(fn.code_sha256) == 64 + # every reason the check gives has its own value, below a clean answer's + named = v.reasons | ({"tool_rejected"} if v is ISA else set()) + assert set(fn.BASE) == named and len(set(fn.BASE.values())) == len(fn.BASE) + cases = [("clean", None, 1.0)] + [("fail", r, b) for r, b in fn.BASE.items()] + [ + ("fail", "other", None), ("fail", None, None), ("finding", None, None), ("unsettled", "x", None), + ("not_run", None, None), ("skipped_by_plan", None, None)] + for state, reason, base in cases: + for tokens in (0, 256, 512): + leg = LegResult(v.leg, state, reason, [], 1, 1, None, None, False, 0.0, 0.0, 0.0) + c = Candidate("c", "r", 1, f"{v.name}:x", 0, "v", "", tokens, 512, TaskRef("x", None, None, None, "t"), + v.name, None) + reward, parts = fn.compute({v.leg: leg}, c) + d = fn.decide({v.leg: leg}, c) + if base is None: + assert reward is None and type(d).__name__ == "Unscored" and not fn.claims_clean(reward, parts) + continue + assert set(parts) == {"base", "broke_healthy", "length"} and parts["base"] == base + assert REWARD_MIN <= reward <= REWARD_MAX and base <= 1.0 + assert parts["base"] + parts["broke_healthy"] + parts["length"] == reward + assert (d.reward, d.parts) == (reward, parts) + assert fn.claims_clean(reward, parts) == (state == "clean") + assert type(fn.decide({}, c)).__name__ == "NeedMore" and fn.compute({}, c)[0] is None + + +@BY_NAME +def test_each_package_stands_alone_and_its_readme_is_short(v): + d = VERTICALS / v.name + files = {p.name for p in d.iterdir()} - {"__pycache__"} + assert files == {v.module, f"{v.name}.yaml", "score.yaml", "train.yaml", "pool", "README.md"} + text = (d / v.module).read_text() + imports = re.findall(r"^\s*(?:from|import)\s+([\w.]+)", text, re.MULTILINE) + assert all(m.split(".")[0] in {"__future__", "hashlib", "json", "os", "re", "shutil", "subprocess", "unicodedata", + "pathlib", "aorlta"} for m in imports), imports + assert "TODO" in text and "PLACEHOLDER" in text + readme = (d / "README.md").read_text() + assert len(readme.splitlines()) < 30 + for name in (v.module, f"{v.name}.yaml", "score.yaml", "train.yaml", "tests/examples/test_verticals.py"): + assert name in readme + + +def replies_from(monkeypatch, replies): + """The fake engines answer with one of `replies`, drawn from the request's own random stream, after a line of its + own: two seeds then give different text, as the engines' start-up self-test needs.""" + from aorlta.rollout.fake_engine import FakeEngine + + monkeypatch.setattr(FakeEngine, "_answer_text", + lambda self, rng: f"({rng.hex('tag')})\n" + replies[int(rng.uniform("reply") * len(replies))]) + + +def wired(tmp_path, base, v, kind, monkeypatch): + from aorlta.orchestrator import wiring + from aorlta.rollout.engine_proc import InProcessEngine + + replies_from(monkeypatch, [v.good, *(r for r, _ in v.bad)]) + data = parse_yaml((VERTICALS / v.name / f"{kind}.yaml").read_bytes()) + data["run"].update(root=str(tmp_path / "runs"), run_id=f"{v.name}-cpu") + data["model"]["base_dir"] = str(base) + data["data"]["prompt_pools"][0]["path"] = str(VERTICALS / v.name / "pool" / ("heldout" if kind == "score" + else "train")) + data["rollout"].update(engines=[0, 1], K=4, P=4, request_timeout_s=300) + if kind == "train": + data["orchestrator"]["iterations"] = 1 + data["trainer"] = {"ranks": 1, "lr": 1e-2, "micro_batch_tokens": 2048} + data["sync"] = {"dir": str(tmp_path / "shm")} + cfg = config_from_dict(data) + engines = [E.TwinEngine(base, E.FAKE, i) for i in range(2)] + parts = wiring.Parts(tokenizer=E.ByteTokenizer(), handles=[InProcessEngine(i, e) for i, e in enumerate(engines)], + load_model=E.load_trainer_model, trainer_device="cpu", require_fast_path=False, + probes=E.PROBES) + return cfg, wiring.build(cfg, parts=parts) + + +def scored_bases(eps, v): + """Each scored episode's base reward (give or take the length term), and its leg state.""" + return {(round(e["reward"] + 0.05, 1), e["legs"][v.leg]) for e in eps if e["status"] == "scored"} + + +@pytest.mark.heavy +@BY_NAME +def test_a_vertical_trains_one_iteration_through_the_real_wiring(tmp_path, load, tiny_base, v, monkeypatch): + load(v) + cfg, orch = wired(tmp_path, tiny_base, v, "train", monkeypatch) + result = orch.run() + assert result["status"] == "completed", result + s = E.summary(cfg, 1) + assert s["trainer"]["step_taken"] and s["samples"]["trained"] > 0 and s["sync"]["landing"] == "LANDED" + eps = [e for e in E.episodes(cfg, 1) if e["end_reason"] != "engine_error"] + assert eps and {e["prompt_id"].partition(":")[0] for e in eps} == {v.name} + bases = scored_bases(eps, v) + fn = plugins.REWARD_FNS[f"{v.name}.v1"](None) + assert {b for b, _ in bases} <= {1.0, *(round(b, 1) for b in fn.BASE.values())} + assert len({b for b, _ in bases}) >= 2 # the answers get different rewards, so there is something to learn + + +@pytest.mark.heavy +@BY_NAME +def test_a_vertical_scores_its_held_out_prompts_through_the_real_wiring(tmp_path, load, tiny_base, v, monkeypatch): + load(v) + cfg, scorer = wired(tmp_path, tiny_base, v, "score", monkeypatch) + assert type(scorer).__name__ == "Scorer" + result = scorer.run() + assert result["status"] == "completed", result + rows = [json.loads(x) for x in (cfg.run_dir / "score" / "report.jsonl").read_text().splitlines()] + assert rows[-1]["samples"] == 16 + held = {f"{v.name}:{r['id']}" for r in load(v).load_rows(VERTICALS / v.name / "pool" / "heldout")} + assert {r["prompt_id"] for r in rows if "prompt_id" in r} == held