From 4535861333fa1499c45e35297dc70c9c87ff85ac Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Thu, 12 Nov 2026 17:22:42 +0000 Subject: [PATCH 01/19] feat(aorlta): register a task's prompts, parse, checks and reward from a plugin A run config lists plugin modules under `plugins:`. Each is imported before the config is checked and registers its components with aorlta.plugins: a task (the prompt pool's records), a parse adapter, a feedback template, a compile check, a check profile, check adapters and a reward function. Every place that looked a name up in a fixed list now also looks in the registry, so a new task needs no edit to the package. The run plan records each plugin file's sha256. A plugin's data is handled as the stand-in source's, and its reward keeps the service's range and parts; both are stated in ADAPTING.md. --- aorlta/docs/ADAPTING.md | 75 +++- aorlta/src/aorlta/config.py | 62 ++- aorlta/src/aorlta/orchestrator/wiring.py | 90 +++- aorlta/src/aorlta/plugins.py | 210 ++++++++++ aorlta/src/aorlta/reward/factories.py | 14 +- aorlta/src/aorlta/reward/profiles.py | 11 +- aorlta/src/aorlta/reward/remote.py | 2 +- .../src/aorlta/reward/rewardfn/for_profile.py | 23 +- aorlta/src/aorlta/rollout/chain.py | 42 +- aorlta/src/aorlta/rollout/compile_check.py | 11 +- aorlta/src/aorlta/schemas.py | 5 +- .../tests/orchestrator/test_plugin_wired.py | 54 +++ aorlta/tests/test_config.py | 6 +- aorlta/tests/test_plugins.py | 396 ++++++++++++++++++ 14 files changed, 942 insertions(+), 59 deletions(-) create mode 100644 aorlta/src/aorlta/plugins.py create mode 100644 aorlta/tests/orchestrator/test_plugin_wired.py create mode 100644 aorlta/tests/test_plugins.py diff --git a/aorlta/docs/ADAPTING.md b/aorlta/docs/ADAPTING.md index 38c4dbb25..3789fac01 100644 --- a/aorlta/docs/ADAPTING.md +++ b/aorlta/docs/ADAPTING.md @@ -3,6 +3,52 @@ Code paths are relative to `src/aorlta/`; commands and the `configs/`, `examples/` and `tests/` paths are relative to `aorlta/`, and the commands assume the environment from the README. +## 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 +an edit to the package. Write a Python module that registers its components with `aorlta.plugins`, and list it in +the run config: + +```yaml +plugins: [examples/my-task/my_task.py] # a .py path (relative: from aorlta/), or an importable module name +data: + prompt_pools: + - {path: /data/my-task, source: standin, task: my-task} +rollout: + parse_adapter: my-task.parse # or the built-in standin.parse + feedback_template: my-task.feedback # or the built-in standin + compile_adapter: standin.hipcc # or one of yours +reward: + profile: my-task-profile + reward_fn: my-task.v1 +``` + +The module is imported before the rest of the config is checked, so every name above may be one it registered: + +| config key | register with | what it is | +|---|---|---| +| `data.prompt_pools[i].task` | `register_task(Task(name, records, probe_messages=None))` | `records(pool_dir, tokenizer, thinking)` returns the pool's `PromptRecord`s ("A new prompt source" below) | +| `rollout.parse_adapter` | `register_parse_adapter(name, factory)` | a `ParseAdapter` (`rollout/chain.py`) | +| `rollout.feedback_template` | `register_feedback_template(name, factory)` | a `FeedbackTemplate`, optionally with `build_apply` | +| `rollout.compile_adapter` | `register_compile_adapter(name, factory)` | `factory(workdir, evidence=..., structure=...)` returns a compiler | +| `reward.profile` | `register_profile(name, path)` | a profile YAML file ("Checks and profiles" below) | +| `reward.reward_fn` | `register_reward_fn(name, factory)` | `factory(profile)` returns a `RewardFn` ("Your own reward function") | +| a profile pool's `adapter` | `register_check_adapter(name, factory)` | `factory(ctx)` returns a `CheckAdapter` for one worker; `ctx` (`plugins.CheckContext`) holds its profile settings, worker id, clock, timeout and the run's pool paths | + +A built-in component may be mixed in by name (`standin.parse`, `standin.hipcc` and `standin` are the stand-in +parse, compile and feedback). A component's `name` must be the name it is registered under. The landing check and +the trainer's warm-up read `sync.probes` chats: the task's `probe_messages(pool_dir)`, or the first records' own +messages, which must then share one system message. The run plan records each plugin's file and its sha256. + +Limits of a plugin: its data is handled as the stand-in source's (public; text may appear in logs and tracebacks); +its reward stays within the stand-in range and parts (below); the episode is the compile-repair chain (a parse may +reject an answer with a message, and the feedback template's `build_apply` turns that into a repair turn); and +`reward.remote` and the run job's checker pre-flight serve only the stand-in sanitizers, so a plugin's GPU check +runs in the orchestrator's process, on GPUs its own settings name. A plugin file under `aorlta/` is in the run +job's snapshot of the committed tree; one elsewhere must be in a directory the job mounts (`AORLTA_MOUNTS`) and is +then not part of the snapshot. `tests/test_plugins.py` and `tests/orchestrator/test_plugin_wired.py` hold a small +plugin that registers one of each and trains on CPU through the real wiring. + ## Prompts in the stand-in format The stand-in source reads one dataset directory: @@ -40,27 +86,27 @@ The loop reads prompts through one method (`PromptPool` in `orchestrator/ports.p | `task` | a `TaskRef`: `contract_id`, `task_sha256` (a hash of what the checks read; part of the cache key), and `parent_id`, `parent_healthy` and `reference_kind`, which only `standin.v11`'s broke-healthy and frozen-reference terms read | | `source` | `"standin"`; it must equal `run.source` | -`standin_records` in `rollout/templates.py` builds these for the stand-in pool. To use a new source, pass it to the -wiring, `wiring.build(cfg, parts=wiring.Parts(pool=MyPool(...)))`, or build it inside `orchestrator/wiring.py:build`. -Three things there assume the stand-in pool: +`standin_records` in `rollout/templates.py` builds these for the stand-in pool. A new set of prompts is a task a +plugin registers ("Your own task, as a plugin" above): the pool's `task` names it, and its `records` builds them. +What the wiring still assumes: - `check_supported` accepts exactly one pool, with `source: standin`; -- `build` still reads the stand-in dataset at `data.prompt_pools[0].path` for the landing probes and the warm-up, - unless `Parts.probes` is given; +- a `standin` pool also gives the landing probes and the warm-up, from 40 of its kernels, so it needs at least 40 + kernels of 60 to 400 instructions; a plugin task gives its own probes; - the answer is parsed by `rollout.parse_adapter`. `standin.parse` takes the last fenced `cpp` block of the final turn, and so does the reward's parse gate (`reward/adapters/standin_parse.py`); the compile checks run `hipcc --offload-arch=gfx950 -c -O1` on it. Another answer format is a parse adapter and a feedback template (`ParseAdapter` and `FeedbackTemplate` in -`rollout/chain.py`, listed in `PARSE_ADAPTERS` and `FEEDBACK_TEMPLATES`), named by `rollout.parse_adapter` and -`rollout.feedback_template`. A parse returns `edits` (the code to compile), `rejected`, or `no_change`: a reply that +`rollout/chain.py`), registered by a plugin and named by `rollout.parse_adapter` and `rollout.feedback_template`. A parse returns `edits` (the code to compile), `rejected`, or `no_change`: a reply that leaves the kernel as shown ends the chain without a compile, end reason `no_change` (the stand-in parse never returns it). A feedback template builds the compile-feedback turn, and may give a rejected answer a repair turn of its own (`build_apply`). Repair turns are appended to the token ids (`rollout.repair_turn_ids: append`, the default), or rebuilt with `rerender`: the whole conversation rendered again with the chat template, the replies re-tokenised from their text (Qwen3.5's template also drops the earlier turns' think blocks and trims every turn). -A data source other than `standin` is a value of `Source` (`config.py`, `schemas.py`) and `SOURCES` +A data source other than `standin` (data whose text must stay out of logs) is still a change to the package: a +value of `Source` (`config.py`, `schemas.py`) and `SOURCES` (`episodes/schema.py`, `eval/hook.py`), with a component-name prefix that `component_source` (`config.py`) maps to it. `tests/orchestrator/test_wiring_e2e.py` builds the real loop on CPU with its own 7-kernel dataset, a tiny model and in-process engines; it is the template for testing a new source without a GPU. @@ -109,11 +155,9 @@ A clean `gate.compile` also records the launched kernels' LDS, read from their k "ConSan does not apply to a launch with no LDS" would read it in `not_applicable_legs`. The gate builds with the run image's compiler and ConSan with the checker image's, and that the two builds agree on LDS is not checked. -To add a reward function, write the class; add its name to `RewardSection.reward_fn` and to the profile-to-function -map `_PROFILE_REWARD_FN` in `config.py`, and set `reward_fn:` in the profile; build it in `standin_rewardfn` -(`reward/rewardfn/for_profile.py`), which the wiring and the offline evaluator call; and add its name to -`STANDIN_REWARD_FNS` in `orchestrator/wiring.py`, which `check_supported` reads. A new profile name goes in -`PROFILE_NAMES` (`reward/profiles.py`), `CheckProfile` (`schemas.py`) and `RewardSection.profile` (`config.py`). +To add a reward function, write the class, register it (`plugins.register_reward_fn(name, factory)`, where +`factory(profile)` builds it), and set `reward_fn:` to its name in your profile, which you register with +`plugins.register_profile`. The config then accepts both names and checks that they agree. ## Checks and profiles @@ -156,8 +200,9 @@ To add a check: a quick heartbeat. `make_result` and `unsettled` build results, `ProcessRunner` runs a subprocess the watchdog can kill, and the job's files are `job.package`, `job.inputs` and `job.output_dir`. The compile gate, `adapters/standin_hipcc.py`, is a short real one. -2. Register a factory in `standin_factories` (`reward/factories.py`): one adapter per worker, configured from the - profile's `adapters:` block. The checker servers build adapters through the same function. +2. Register a factory with `plugins.register_check_adapter(name, factory)`: `factory(ctx)` builds one adapter per + worker, from `ctx.settings` (the profile's `adapters:` entry for it). `standin_factories` (`reward/factories.py`) + builds the built-in ones and the registered ones the same way. 3. In the profile, add a pool that names the adapter, a stage for its leg, and its `requires`. 4. List the new leg in a reward function's `legs` and read it there. The service builds each candidate's final record from the reward function's legs, and for a clean reward checks every profile leg in that record, so diff --git a/aorlta/src/aorlta/config.py b/aorlta/src/aorlta/config.py index d407e3eaa..95ce241a8 100644 --- a/aorlta/src/aorlta/config.py +++ b/aorlta/src/aorlta/config.py @@ -339,13 +339,18 @@ def _absolute(name: str, value: str) -> None: def component_source(name: str) -> str | None: - """The data source a named component belongs to: standin* and fake* are stand-ins, and "none" (only allowed - for the leak check) belongs to none.""" + """The data source a named component belongs to: standin* and fake* are stand-ins, and so is a component a + plugin registered (`plugins.py`); "none" (only allowed for the leak check) belongs to none.""" if name == "none": return None + from . import plugins + + if plugins.registered(name): + return "standin" m = _COMPONENT_RE.fullmatch(name) if m is None: - raise ConfigError(f"component name {name!r} must be 'none' or start with 'standin' or 'fake'") + raise ConfigError(f"component name {name!r} must be 'none', start with 'standin' or 'fake', or be registered " + f"by a plugin") return "standin" @@ -384,9 +389,17 @@ def _validate(self) -> None: class PromptPool(_Section): path: str source: Source + # what reads the pool's directory: "standin" (a kernel dataset's listings, rollout/templates.py) or a task a + # plugin registers (plugins.register_task) + task: str = "standin" def _validate(self) -> None: _absolute("path", self.path) + if self.task != "standin": + from . import plugins + + _check(self.task in plugins.TASKS, f"task: {self.task!r} is not 'standin' or a task a plugin registered " + f"(registered: {', '.join(sorted(plugins.TASKS)) or 'none'})") @dataclass(frozen=True) @@ -499,6 +512,20 @@ def _validate(self) -> None: } +def profile_reward_fn(profile: str) -> str: + """The reward function a check profile names: the built-in map, or a registered profile's `reward_fn:`.""" + if profile in _PROFILE_REWARD_FN: + return _PROFILE_REWARD_FN[profile] + from . import plugins + + if profile not in plugins.PROFILES: + raise ConfigError(f"profile: {profile!r} is not one of {', '.join(_PROFILE_REWARD_FN)} or a profile a plugin " + f"registered (registered: {', '.join(sorted(plugins.PROFILES)) or 'none'})") + from .reward.profiles import load_profile + + return load_profile(profile).reward_fn + + @dataclass(frozen=True) class RewardRemote(_Section): """Adapters that run in their own checker process, reached over a Unix socket (`aorlta.reward.remote`). The @@ -536,8 +563,8 @@ def _validate(self) -> None: class RewardSection(_Section): """Section 6.2.""" - profile: Literal["standin-static", "standin-gpu", "fake"] - reward_fn: Literal["standin.v11"] + profile: str # standin-static, standin-gpu, fake, or a profile a plugin registered + reward_fn: str # standin.v11, or a reward function a plugin registered; must be the one the profile names cache: RewardCache = field(default_factory=RewardCache) sim_workers: int | None = None # None: sized by the reward service; "sized by measurement" speculative_parallel: bool = True @@ -554,7 +581,11 @@ class RewardSection(_Section): gpu_check_s: dict[str, float] = field(default_factory=dict) def _validate(self) -> None: - expected = _PROFILE_REWARD_FN[self.profile] + expected = profile_reward_fn(self.profile) + from . import plugins + + _check(self.reward_fn == "standin.v11" or self.reward_fn in plugins.REWARD_FNS, + f"reward_fn: {self.reward_fn!r} is not standin.v11 or a reward function a plugin registered") _check( self.reward_fn == expected, f"reward_fn: profile {self.profile!r} is scored by {expected!r}, got {self.reward_fn!r}", @@ -774,6 +805,9 @@ class RunConfig(_Section): selector: SelectorSection = field(default_factory=SelectorSection) trainer: TrainerSection = field(default_factory=TrainerSection) sync: SyncSection = field(default_factory=SyncSection) + # modules or .py files that register a task's components (plugins.py); config_from_dict imports them before it + # checks the sections + plugins: tuple[str, ...] = () def _validate(self) -> None: if self.trainer.loss_temperature is None: @@ -846,9 +880,20 @@ def source_mismatches(cfg: RunConfig) -> list[str]: return out +def _load_plugins(entries: Any, where: str = "plugins") -> None: + from . import plugins + + try: + plugins.load_plugins(entries) + except plugins.PluginError as exc: + raise ConfigError(f"{where}: {exc}") from None + + def config_from_dict(data: Any, source: str = "") -> RunConfig: if not isinstance(data, Mapping): raise ConfigError(f"{source}: the config must be a mapping, got {_describe(data)}") + if data.get("plugins"): # before the sections, whose names a plugin may register + _load_plugins(data["plugins"], f"{source}: plugins") return strict_decode(RunConfig, data, "") @@ -997,6 +1042,11 @@ def write_run_plan( }, "environment": {"python": sys.version.split()[0], "pyyaml": yaml.__version__}, } + if cfg.plugins: + from . import plugins + + loaded = plugins.identity() + plan["plugins"] = {name: loaded.get(name) for name in cfg.plugins} run_dir.mkdir(parents=True, exist_ok=True) try: atomic_write_json(plan_path, plan, replace=False) diff --git a/aorlta/src/aorlta/orchestrator/wiring.py b/aorlta/src/aorlta/orchestrator/wiring.py index 0d52802ec..ea0674fcb 100644 --- a/aorlta/src/aorlta/orchestrator/wiring.py +++ b/aorlta/src/aorlta/orchestrator/wiring.py @@ -11,7 +11,8 @@ What this wiring does not do, each refused by `check_supported` or stated where it bites: - streaming accumulation only, one step per iteration; - - stand-in prompts only, with the held-out kernels left out of training; + - one prompt pool: the stand-in prompts, with the held-out kernels left out of training, or a task a plugin + registered (plugins.py), whose prompts are its own; - no eval exporter (`orchestrator.eval_every` must be null); - the selector off or log-only with no planner endpoint; - the generation deadline stops chains from starting another turn, but does not abort a request already running @@ -81,8 +82,11 @@ def check_supported(cfg: RunConfig) -> None: pools = cfg.data.prompt_pools if len(pools) != 1 or pools[0].source != "standin": problems.append("data.prompt_pools must be exactly one stand-in pool") - if cfg.reward.reward_fn not in STANDIN_REWARD_FNS: - problems.append(f"reward.reward_fn is {cfg.reward.reward_fn!r}: only standin.v11 is wired") + from .. import plugins + + if cfg.reward.reward_fn not in STANDIN_REWARD_FNS and cfg.reward.reward_fn not in plugins.REWARD_FNS: + problems.append(f"reward.reward_fn is {cfg.reward.reward_fn!r}: only standin.v11 and reward functions a " + f"plugin registered are wired") if cfg.selector.endpoint is not None: problems.append("selector.endpoint is set: no planner is wired (log-only runs without one)") if cfg.rollout.P < SELF_TEST_PROMPTS: @@ -165,6 +169,59 @@ def standin_dataset_dir(path: str | Path) -> Path: raise WiringError(f"{path}: neither it nor its parent holds kernels.json and listings/") +class TaskPromptPool: + """The PromptPool port over a task a plugin registered (`plugins.Task`): its records from the pool's directory, + checked once: unique prompt ids, every record of the stand-in source.""" + + source = "standin" + + def __init__(self, task, tokenizer, thinking: str, pool_dir: str | Path): + self.task, self.tokenizer, self.thinking, self.pool_dir = task, tokenizer, thinking, Path(pool_dir) + self._records: list | None = None + + def records(self) -> list: + if self._records is None: + recs = list(self.task.records(self.pool_dir, self.tokenizer, self.thinking)) + if not recs: + raise WiringError(f"task {self.task.name!r} gave no prompts from {self.pool_dir}") + ids = [r.prompt_id for r in recs] + if len(set(ids)) != len(ids): + raise WiringError(f"task {self.task.name!r} gave a prompt id twice") + stray = sorted({r.source for r in recs} - {self.source}) + if stray: + raise WiringError(f"task {self.task.name!r} gave records of source {stray}, not {self.source!r}") + self._records = recs + return list(self._records) + + +def task_probes(cfg: RunConfig, tokenizer, task, pool_dir: str | Path, records: Sequence): + """The landing probes and warm-up ids of a task's run: `sync.probes` chats from the task's `probe_messages`, or + the first records' own messages. Each chat is a system and a user message (or a user message alone), and they + share one system message.""" + from ..sync.landing import build_probes + + chats = (list(task.probe_messages(Path(pool_dir))) if task.probe_messages is not None + else [r.messages for r in records if r.messages is not None]) + chats = chats[:cfg.sync.probes] + if len(chats) < cfg.sync.probes: + raise WiringError(f"task {task.name!r}: the landing check needs {cfg.sync.probes} probe chats " + f"(sync.probes), it has {len(chats)}") + systems, users = set(), [] + for chat in chats: + roles = [m["role"] for m in chat] + if roles == ["system", "user"]: + systems.add(chat[0]["content"]) + elif roles == ["user"]: + systems.add("") + else: + raise WiringError(f"task {task.name!r}: a probe chat is a system and a user message, got roles {roles}") + users.append(chat[-1]["content"]) + if len(systems) != 1: + raise WiringError(f"task {task.name!r}: the probe chats must share one system message") + return build_probes(getattr(tokenizer, "tok", tokenizer), users, system=systems.pop(), + tokens=cfg.sync.probe_tokens) + + def node_dir(cfg: RunConfig) -> Path: """This run's node-local directory, which also holds its weight slots: `/`.""" return Path(cfg.sync.dir) / cfg.run.run_id @@ -416,7 +473,8 @@ def __init__(self, cfg: RunConfig, *, clock, spans, factories: dict | None = Non from ..reward.factories import standin_factories factories = standin_factories(self.profile, clock, timeouts=dataclasses.asdict(cfg.reward.timeouts), - check_s=cfg.reward.gpu_check_s) + check_s=cfg.reward.gpu_check_s, + pools=tuple(p.path for p in cfg.data.prompt_pools)) factories.update(remote_factories(cfg, clock)) missing = sorted({p.adapter for p in self.profile.pools.values()} - set(factories)) if missing: @@ -1179,20 +1237,32 @@ def build(cfg: RunConfig, *, config_path=None, commit: str | None = None, parts: from ..spans import SpanRecorder from .loop import Orchestrator + from .. import plugins + parts = parts or Parts() check_supported(cfg) clock = Clock() spans = SpanRecorder(monotonic=clock.now, wall=clock.wall) dataset = None - if parts.pool is None or parts.probes is None: # the prompts and the landing probes come from the dataset - dataset = standin_dataset_dir(cfg.data.prompt_pools[0].path) + pool_cfg = cfg.data.prompt_pools[0] + task = None if pool_cfg.task == "standin" else plugins.task(pool_cfg.task) + if task is None and (parts.pool is None or parts.probes is None): # the prompts and probes come from the dataset + dataset = standin_dataset_dir(pool_cfg.path) tok = parts.tokenizer or HFTokenizer.load(cfg.model.base_dir) - pool = parts.pool if parts.pool is not None else StandinPromptPool(tok, cfg.rollout.thinking, dataset) + if parts.pool is not None: + pool = parts.pool + elif task is not None: + pool = TaskPromptPool(task, tok, cfg.rollout.thinking, pool_cfg.path) + else: + pool = StandinPromptPool(tok, cfg.rollout.thinking, dataset) try: records = pool.records() max_len = rollout_max_len(cfg, tok, records) except TemplateError as exc: - raise WiringError(f"stand-in prompts: {exc}") from None + raise WiringError(f"{pool_cfg.task} prompts: {exc}") from None + probes = parts.probes + if probes is None and task is not None: + probes = task_probes(cfg, tok, task, pool_cfg.path, records) compiler = parts.compiler if parts.compiler is not None else rollout_compiler(cfg) handles = parts.handles if parts.handles is not None else engine_handles(cfg, max_len) if cfg.orchestrator.mode == "score": # no trainer, no optimizer, no weight sync @@ -1200,7 +1270,7 @@ def build(cfg: RunConfig, *, config_path=None, commit: str | None = None, parts: rollout = RolloutAdapter(cfg, tok, handles, max_len=max_len, spans=spans, compiler=compiler, on_shutdown=[lambda: remove_node_dir(cfg)], clock=clock) - model = ServedModel(cfg, tok, rollout, spans=spans, dataset=dataset, probes=parts.probes) + model = ServedModel(cfg, tok, rollout, spans=spans, dataset=dataset, probes=probes) reward = RewardAdapter(cfg, clock=clock, spans=spans, factories=parts.reward_factories) return Scorer(cfg, pool=pool, rollout=rollout, reward=reward, model=model, clock=clock, spans=spans, commit=commit, config_path=config_path) @@ -1208,7 +1278,7 @@ def build(cfg: RunConfig, *, config_path=None, commit: str | None = None, parts: require_fast_path=parts.require_fast_path, amp=parts.amp) rollout = RolloutAdapter(cfg, tok, handles, max_len=max_len, spans=spans, compiler=compiler, on_shutdown=[trainer.shutdown, lambda: remove_node_dir(cfg)], clock=clock) - sync = SyncAdapter(cfg, tok, trainer, rollout, spans=spans, dataset=dataset, probes=parts.probes) + sync = SyncAdapter(cfg, tok, trainer, rollout, spans=spans, dataset=dataset, probes=probes) trainer.warmup_ids = lambda: sync.probes().ids # the landing probes: real stand-in prompt tokens reward = RewardAdapter(cfg, clock=clock, spans=spans, factories=parts.reward_factories) return Orchestrator(cfg, pool=pool, rollout=rollout, reward=reward, trainer=trainer, sync=sync, clock=clock, diff --git a/aorlta/src/aorlta/plugins.py b/aorlta/src/aorlta/plugins.py new file mode 100644 index 000000000..ce07d23fa --- /dev/null +++ b/aorlta/src/aorlta/plugins.py @@ -0,0 +1,210 @@ +"""Plugins: a task's own prompts, answer parsing, checks and reward function, by name, without editing this package. + +A run config lists its plugins under `plugins:`, each an importable module name or the path of a `.py` file (a +relative path is read from the package root, the directory that holds `src/aorlta`, so a plugin committed beside the +package runs from the run job's snapshot of the tree). Each is imported once, before the rest of the config is +checked, and registers its components with the `register_*` functions below. A config, a profile or the wiring +then finds a name among the built-in components first and the registered ones after: + +| config key | registered with | +|---|---| +| `data.prompt_pools[i].task` | `register_task`: the pool's prompt records, from the pool's directory | +| `rollout.parse_adapter` | `register_parse_adapter` (`rollout.chain.ParseAdapter`) | +| `rollout.feedback_template` | `register_feedback_template` (`rollout.chain.FeedbackTemplate`) | +| `rollout.compile_adapter` | `register_compile_adapter` (`rollout.compile_check.make_compiler`'s arguments) | +| `reward.profile` | `register_profile`: a profile YAML file (`reward/profiles.py`) | +| `reward.reward_fn` | `register_reward_fn` (`reward.rewardfn.base.RewardFn`) | +| a profile's pool `adapter` | `register_check_adapter` (`reward.adapters.base.CheckAdapter`), one per worker | + +A registered component belongs to the `standin` source: its data is public, and its text may appear in logs and +tracebacks as the stand-in data's does. Data that must not, and a reward outside the stand-in range and parts +(`reward/service.py`), still need a change to this package. A name is registered once; registering it again with +another object is refused, with the same object (a module imported twice) it is a no-op. +""" + +from __future__ import annotations + +import hashlib +import importlib +import importlib.util +import sys +import threading +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +PACKAGE_ROOT = Path(__file__).resolve().parents[2] # the directory that holds src/aorlta + + +class PluginError(ValueError): + """A plugin could not be loaded, or registered a name twice.""" + + +@dataclass(frozen=True) +class Task: + """A prompt source for `data.prompt_pools[i].task`. + + `records(pool_dir, tokenizer, thinking)` returns the pool's `PromptRecord`s (`schemas.py`), each with + `source="standin"`, a unique `prompt_id` and a `task.task_sha256` that hashes what the checks read for it. + `probe_messages(pool_dir)`, if given, returns the chats (a system and a user message) the landing check and the + trainer's warm-up use; without it they are the first records' own messages.""" + + name: str + records: Callable[[Path, Any, str], list] + probe_messages: Callable[[Path], list[list[dict]]] | None = None + + +@dataclass(frozen=True) +class CheckContext: + """What a registered check adapter's factory is given, once per worker: its `settings` (the profile's + `adapters:` entry under its name, with `${...}` tool roots expanded), the worker id, the service's wall clock, + `timeout_s` (a little under its pool's timeout, so the adapter reports before the watchdog), the run's prompt pool + directories, and `reward.gpu_check_s`.""" + + name: str + worker_id: str + settings: Mapping[str, Any] + clock: Callable[[], float] + timeout_s: float + pools: tuple[str, ...] = () + check_s: Mapping[str, float] = field(default_factory=dict) + + +TASKS: dict[str, Task] = {} +PARSE_ADAPTERS: dict[str, Callable[[], Any]] = {} +FEEDBACK_TEMPLATES: dict[str, Callable[[], Any]] = {} +COMPILE_ADAPTERS: dict[str, Callable[..., Any]] = {} +PROFILES: dict[str, Path] = {} +REWARD_FNS: dict[str, Callable[[Any], Any]] = {} +CHECK_ADAPTERS: dict[str, Callable[[CheckContext], Any]] = {} +LOADED: dict[str, dict] = {} # plugin as listed -> {"module", "file", "sha256"} +_LOCK = threading.RLock() + +_BUILT_IN = { + "task": ("standin",), + "parse adapter": ("standin.parse",), + "feedback template": ("standin",), + "compile adapter": ("standin.hipcc", "fake.compile"), + "profile": ("standin-static", "standin-gpu", "fake"), + "reward function": ("standin.v11",), + "check adapter": ("standin.parse", "standin.hipcc", "standin.rj_waitcheck", "standin.aorta_consan", + "standin.aorta_asan", "fake", "fake-asan"), +} + + +def _register(kind: str, table: dict, name: str, value: Any) -> None: + if not (isinstance(name, str) and name and name.strip() == name and name != "none"): + raise PluginError(f"a {kind} needs a non-empty name without surrounding spaces, other than 'none'; got {name!r}") + if name in _BUILT_IN[kind]: + raise PluginError(f"{kind} {name!r} is built in; register it under another name") + with _LOCK: + if name in table and table[name] is not value and table[name] != value: + raise PluginError(f"{kind} {name!r} is already registered") + table[name] = value + + +def register_task(task: Task) -> None: + _register("task", TASKS, task.name, task) + + +def register_parse_adapter(name: str, factory: Callable[[], Any]) -> None: + """`factory()` -> a ParseAdapter whose `name` is `name`.""" + _register("parse adapter", PARSE_ADAPTERS, name, factory) + + +def register_feedback_template(name: str, factory: Callable[[], Any]) -> None: + """`factory()` -> a FeedbackTemplate whose `name` is `name`.""" + _register("feedback template", FEEDBACK_TEMPLATES, name, factory) + + +def register_compile_adapter(name: str, factory: Callable[..., Any]) -> None: + """`factory(workdir, evidence=..., structure=...)` -> a compiler with `check(code, tag) -> CompileResult`.""" + _register("compile adapter", COMPILE_ADAPTERS, name, factory) + + +def register_profile(name: str, path: str | Path) -> None: + """A check profile YAML file whose `name:` is `name`.""" + p = Path(path).resolve() + if not p.is_file(): + raise PluginError(f"profile {name!r}: no file at {p}") + _register("profile", PROFILES, name, p) + + +def register_reward_fn(name: str, factory: Callable[[Any], Any]) -> None: + """`factory(profile)` -> a RewardFn whose `name` is `name`.""" + _register("reward function", REWARD_FNS, name, factory) + + +def register_check_adapter(name: str, factory: Callable[[CheckContext], Any]) -> None: + """`factory(ctx)` -> a CheckAdapter whose `name` is `name`, for one worker.""" + _register("check adapter", CHECK_ADAPTERS, name, factory) + + +def registered(name: str) -> bool: + """Whether a plugin registered `name` as any kind of component.""" + return any(name in t for t in (TASKS, PARSE_ADAPTERS, FEEDBACK_TEMPLATES, COMPILE_ADAPTERS, PROFILES, REWARD_FNS, + CHECK_ADAPTERS)) + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def load_plugins(plugins: Sequence[str]) -> None: + """Import every plugin listed, once each (see the module docstring for the two forms).""" + if isinstance(plugins, (str, bytes)) or not isinstance(plugins, (list, tuple)): + raise PluginError(f"plugins: expected a list of module names or .py paths, got {type(plugins).__name__}") + for entry in plugins: + if not (isinstance(entry, str) and entry): + raise PluginError(f"plugins: each entry is a module name or a .py path, got {entry!r}") + with _LOCK: + if entry in LOADED: + continue + LOADED[entry] = _load(entry) + + +def _load(entry: str) -> dict: + if entry.endswith(".py") or "/" in entry: + path = Path(entry) + path = (path if path.is_absolute() else PACKAGE_ROOT / path).resolve() + if not path.is_file(): + raise PluginError(f"plugins: no file at {path}") + name = "aorlta_plugin_" + hashlib.sha256(str(path).encode()).hexdigest()[:12] + if name in sys.modules: + module = sys.modules[name] + else: + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise PluginError(f"plugins: cannot import {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + try: + spec.loader.exec_module(module) + except Exception as exc: + del sys.modules[name] + raise PluginError(f"plugins: importing {path} failed: {type(exc).__name__}: {exc}") from exc + else: + try: + module = importlib.import_module(entry) + except Exception as exc: + raise PluginError(f"plugins: importing {entry} failed: {type(exc).__name__}: {exc}") from exc + file = getattr(module, "__file__", None) + return {"module": module.__name__, "file": file, "sha256": _sha256(Path(file)) if file else None} + + +def identity() -> dict: + """The loaded plugins as the run plan records them: each entry's module, file and the file's sha256.""" + with _LOCK: + return {k: dict(v) for k, v in sorted(LOADED.items())} + + +def profile_path(name: str) -> Path | None: + return PROFILES.get(name) + + +def task(name: str) -> Task: + try: + return TASKS[name] + except KeyError: + raise PluginError(f"no task named {name!r}; registered: {', '.join(sorted(TASKS)) or 'none'}") from None diff --git a/aorlta/src/aorlta/reward/factories.py b/aorlta/src/aorlta/reward/factories.py index 9dd3f1b1d..c50a611ce 100644 --- a/aorlta/src/aorlta/reward/factories.py +++ b/aorlta/src/aorlta/reward/factories.py @@ -17,10 +17,14 @@ def standin_factories(profile: Profile, clock: Clock, *, timeouts: dict[str, float] | None = None, - check_s: dict[str, float] | None = None) -> dict[str, Callable[[str], Any]]: - """Factories for every `standin.*` adapter the profile names. An adapter's own subprocess timeout - sits a little under the watchdog's, so the adapter reports first when it can. `check_s` is + check_s: dict[str, float] | None = None, + pools: tuple[str, ...] = ()) -> dict[str, Callable[[str], Any]]: + """Factories for every `standin.*` adapter the profile names, and every adapter it names that a plugin registered + (plugins.py; `pools` are the run's prompt pool paths, which such an adapter is given). An adapter's own subprocess + timeout sits a little under the watchdog's, so the adapter reports first when it can. `check_s` is `reward.gpu_check_s`: the GPU legs' per-check limit, by leg, inside that bound.""" + from .. import plugins + timeouts = timeouts or {} check_s = check_s or {} cfg = profile.adapters @@ -61,4 +65,8 @@ def own(name: str) -> float: aorta_src=c["aorta_src"], hipcc=c.get("hipcc", hipcc), target=c.get("target", "gfx950"), timeout_s=max(1.0, own("standin.aorta_asan") / 18), check_timeout_s=check_s.get("san.asan"), control_source=c.get("control_source"), clock=clock.wall), max_packages=16, inband_controls=True) + for name in sorted(names & set(plugins.CHECK_ADAPTERS)): + out[name] = lambda wid, name=name: plugins.CHECK_ADAPTERS[name](plugins.CheckContext( + name=name, worker_id=wid, settings=dict(cfg.get(name, {})), clock=clock.wall, timeout_s=own(name), + pools=tuple(pools), check_s=dict(check_s))) return out diff --git a/aorlta/src/aorlta/reward/profiles.py b/aorlta/src/aorlta/reward/profiles.py index b50412466..b0d0822a4 100644 --- a/aorlta/src/aorlta/reward/profiles.py +++ b/aorlta/src/aorlta/reward/profiles.py @@ -107,9 +107,11 @@ def stage_of(self, leg: str) -> StageSpec: def profile_problems(p: Profile) -> list[str]: + from .. import plugins + out = [] - if p.name not in PROFILE_NAMES: - out.append(f"name must be one of {PROFILE_NAMES}") + if p.name not in PROFILE_NAMES and p.name not in plugins.PROFILES: + out.append(f"name must be one of {PROFILE_NAMES} or a profile a plugin registered") if p.clean_repeats < 1: out.append("clean_repeats must be >= 1") if not p.stages: @@ -164,7 +166,10 @@ def profile_from_dict(data: Mapping[str, Any], source: str = "") -> Pro def load_profile(name_or_path: str | Path) -> Profile: + """A shipped profile by name, a profile a plugin registered by name, or a profile file by path.""" + from .. import plugins + path = Path(name_or_path) if path.suffix not in (".yaml", ".yml"): - path = PROFILE_DIR / f"{name_or_path}.yaml" + path = plugins.profile_path(str(name_or_path)) or PROFILE_DIR / f"{name_or_path}.yaml" return profile_from_dict(parse_yaml(path.read_bytes(), str(path)), str(path)) diff --git a/aorlta/src/aorlta/reward/remote.py b/aorlta/src/aorlta/reward/remote.py index 7c5c87d13..0afa580ad 100644 --- a/aorlta/src/aorlta/reward/remote.py +++ b/aorlta/src/aorlta/reward/remote.py @@ -415,7 +415,7 @@ def standin_factory(config_path: str, adapter: str) -> Callable[[str], Any]: cfg = load_config(config_path) profile = load_profile(cfg.reward.profile) factories = standin_factories(profile, Clock(), timeouts=dataclasses.asdict(cfg.reward.timeouts), - check_s=cfg.reward.gpu_check_s) + check_s=cfg.reward.gpu_check_s, pools=tuple(p.path for p in cfg.data.prompt_pools)) if adapter not in factories: raise SystemExit(f"profile {profile.name!r} has no stand-in adapter {adapter!r}") return factories[adapter] diff --git a/aorlta/src/aorlta/reward/rewardfn/for_profile.py b/aorlta/src/aorlta/reward/rewardfn/for_profile.py index ed3d306a0..589891b72 100644 --- a/aorlta/src/aorlta/reward/rewardfn/for_profile.py +++ b/aorlta/src/aorlta/reward/rewardfn/for_profile.py @@ -1,5 +1,6 @@ -"""The stand-in reward function a check profile names, built one way for a run (orchestrator/wiring.py) and for the -offline evaluators that re-score its records (selector/plan_eval.py).""" +"""The reward function a check profile names, built one way for a run (orchestrator/wiring.py) and for the +offline evaluators that re-score its records (selector/plan_eval.py): standin.v11, or one a plugin registered +(plugins.py).""" from __future__ import annotations @@ -7,13 +8,23 @@ class UnsupportedRewardFn(ValueError): - """The profile names a reward function that is not a stand-in one.""" + """The profile names a reward function that is neither standin.v11 nor registered by a plugin.""" def standin_rewardfn(profile: Profile): - """standin.v11, the only stand-in reward function.""" - from .standin_v11 import StandinV11 + """standin.v11, or the registered reward function the profile names, built for it.""" + from ... import plugins if profile.reward_fn == "standin.v11": + from .standin_v11 import StandinV11 + return StandinV11(clean_repeats=profile.clean_repeats) - raise UnsupportedRewardFn(f"profile {profile.name!r} is scored by {profile.reward_fn}, not a stand-in reward") + factory = plugins.REWARD_FNS.get(profile.reward_fn) + if factory is not None: + fn = factory(profile) + if getattr(fn, "name", None) != profile.reward_fn: + raise UnsupportedRewardFn(f"the reward function registered as {profile.reward_fn!r} is named " + f"{getattr(fn, 'name', None)!r}") + return fn + raise UnsupportedRewardFn(f"profile {profile.name!r} is scored by {profile.reward_fn}, which is neither standin.v11 " + f"nor registered by a plugin") diff --git a/aorlta/src/aorlta/rollout/chain.py b/aorlta/src/aorlta/rollout/chain.py index 51789c0fd..1d2341c15 100644 --- a/aorlta/src/aorlta/rollout/chain.py +++ b/aorlta/src/aorlta/rollout/chain.py @@ -132,22 +132,48 @@ def _check_built(kind: str, name: str, built: Sequence[str]) -> None: raise RolloutStop(f"{kind} {name!r} is not available; built: {', '.join(built)}") +def _parse_adapters() -> tuple[str, ...]: + from .. import plugins + + return PARSE_ADAPTERS + tuple(sorted(plugins.PARSE_ADAPTERS)) + + +def _feedback_templates() -> tuple[str, ...]: + from .. import plugins + + return FEEDBACK_TEMPLATES + tuple(sorted(plugins.FEEDBACK_TEMPLATES)) + + +def _named(kind: str, obj: Any, name: str) -> Any: + if getattr(obj, "name", None) != name: + raise RolloutStop(f"the {kind} registered as {name!r} is named {getattr(obj, 'name', None)!r}") + return obj + + def check_adapter_names(parse_adapter: str, feedback_template: str) -> None: - """Refuse a configured parse or feedback component that is not built, rather than silently use the - stand-in in its place.""" - _check_built("parse adapter", parse_adapter, PARSE_ADAPTERS) - _check_built("feedback template", feedback_template, FEEDBACK_TEMPLATES) + """Refuse a configured parse or feedback component that is not built or registered by a plugin, rather than + silently use the stand-in in its place.""" + _check_built("parse adapter", parse_adapter, _parse_adapters()) + _check_built("feedback template", feedback_template, _feedback_templates()) def make_parse_adapter(name: str) -> ParseAdapter: - """The configured parse.""" - _check_built("parse adapter", name, PARSE_ADAPTERS) + """The configured parse: the stand-in one, or one a plugin registered.""" + from .. import plugins + + _check_built("parse adapter", name, _parse_adapters()) + if name in plugins.PARSE_ADAPTERS: + return _named("parse adapter", plugins.PARSE_ADAPTERS[name](), name) return StandinParse() def make_feedback_template(name: str) -> FeedbackTemplate: - """The configured feedback.""" - _check_built("feedback template", name, FEEDBACK_TEMPLATES) + """The configured feedback: the stand-in one, or one a plugin registered.""" + from .. import plugins + + _check_built("feedback template", name, _feedback_templates()) + if name in plugins.FEEDBACK_TEMPLATES: + return _named("feedback template", plugins.FEEDBACK_TEMPLATES[name](), name) return StandinFeedback() diff --git a/aorlta/src/aorlta/rollout/compile_check.py b/aorlta/src/aorlta/rollout/compile_check.py index d9c79b951..d66fce6a2 100644 --- a/aorlta/src/aorlta/rollout/compile_check.py +++ b/aorlta/src/aorlta/rollout/compile_check.py @@ -220,12 +220,19 @@ def check(self, code: str, tag: str) -> CompileResult: def make_compiler(name: str, workdir: str | Path, evidence: EvidenceStore | None = None, structure: Structure = compile_errors) -> Compiler: - """The configured rollout compile check (`rollout.compile_adapter`).""" + """The configured rollout compile check (`rollout.compile_adapter`): a built-in one, or one a plugin registered + (`plugins.register_compile_adapter`), called with the same arguments.""" + from .. import plugins + if name == "standin.hipcc": return HipccCompiler(workdir, evidence=evidence, structure=structure) if name == "fake.compile": return FakeCompiler(evidence=evidence, structure=structure) - raise CompileUnavailable(f"compile adapter {name!r} is not available; built: standin.hipcc, fake.compile") + if name in plugins.COMPILE_ADAPTERS: + return plugins.COMPILE_ADAPTERS[name](workdir, evidence=evidence, structure=structure) + registered = "".join(f", {n}" for n in sorted(plugins.COMPILE_ADAPTERS)) + raise CompileUnavailable(f"compile adapter {name!r} is not available; built: standin.hipcc, fake.compile" + f"{registered}") class FakeCompiler: diff --git a/aorlta/src/aorlta/schemas.py b/aorlta/src/aorlta/schemas.py index 0831ba587..8edba69b2 100644 --- a/aorlta/src/aorlta/schemas.py +++ b/aorlta/src/aorlta/schemas.py @@ -18,8 +18,9 @@ Source = Literal["standin"] LegState = Literal["clean", "finding", "fail", "unsettled", "skipped_by_plan", "not_run"] -# "standin-gpu" adds aorta ConSan/ASan on a GPU to standin-static (C2.T11, and stand-in runs). -CheckProfile = Literal["standin-static", "standin-gpu", "fake"] +# "standin-static", "standin-gpu" (aorta ConSan/ASan on a GPU added, C2.T11), "fake", or a profile a plugin +# registered (plugins.py); the run config checks the name, and the service serves only its own profile. +CheckProfile = str EndReason = Literal[ "compiled", "no_change", # a parse that finds the kernel left as shown (an explicit no-change answer, or edits that change diff --git a/aorlta/tests/orchestrator/test_plugin_wired.py b/aorlta/tests/orchestrator/test_plugin_wired.py new file mode 100644 index 000000000..699e70b75 --- /dev/null +++ b/aorlta/tests/orchestrator/test_plugin_wired.py @@ -0,0 +1,54 @@ +"""A task that only a plugin knows, through the real wiring on CPU: its prompts, parse, feedback, compile check, +profile, check adapter and reward function are the plugin's (tests/test_plugins.py), the engines and the model are +test_wiring_e2e's. Nothing in the package names the task.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +import test_plugins as TP # noqa: E402 - the toy plugin +import test_wiring_e2e as E # noqa: E402 - the tiny model, its twin engines and the probes + +from aorlta.config import config_from_dict # noqa: E402 +from aorlta.orchestrator import wiring # noqa: E402 +from aorlta.rollout.engine_proc import InProcessEngine # noqa: E402 + +pytestmark = pytest.mark.heavy + +tiny_base = E.tiny_base +registry, plugin = TP.registry, TP.plugin + + +def build(cfg, base): + engines = [E.TwinEngine(base, E.FAKE, i) for i in range(len(cfg.rollout.engines))] + 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 wiring.build(cfg, parts=parts) + + +def test_a_plugin_task_trains_and_lands(tmp_path, tiny_base, plugin): + data = TP.toy_config(tmp_path, str(plugin)) + data["model"]["base_dir"] = str(tiny_base) + data["rollout"].update({"engines": [0, 1], "K": 4, "P": 4, "max_new_tokens": 512, "request_timeout_s": 300}) + data["trainer"] = {"ranks": 1, "lr": 1e-2, "micro_batch_tokens": 2048} + cfg = config_from_dict(data) + orch = build(cfg, tiny_base) + assert [r.prompt_id for r in orch.pool.records()] == [f"toy:p{i}" for i in range(6)] + 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.episodes(cfg, 1) + assert {e["prompt_id"] for e in eps} <= {f"toy:p{i}" for i in range(6)} + rewards = {e["reward"] for e in eps if e["status"] == "scored"} + assert rewards and rewards <= {1.0, -1.0} and len(rewards) == 2 # toy.v1's two values, from toy.check + assert any(len(e["turns"]) > 1 for e in eps) # a repair turn: toy.feedback's compile or apply turn + plan = json.loads((cfg.run_dir / "run_plan.json").read_text()) + assert plan["plugins"][str(plugin)]["sha256"] == TP.sha256_file(plugin) diff --git a/aorlta/tests/test_config.py b/aorlta/tests/test_config.py index d96d3cbb8..1ae490e54 100644 --- a/aorlta/tests/test_config.py +++ b/aorlta/tests/test_config.py @@ -209,7 +209,7 @@ def test_only_the_standin_source_is_known(): @pytest.mark.parametrize("dotted", ["rollout.compile_adapter", "rollout.parse_adapter", "rollout.feedback_template", "rollout.leak_check_adapter"]) def test_a_component_of_no_known_source_is_refused(dotted): - with pytest.raises(ConfigError, match="must be 'none' or start with 'standin' or 'fake'"): + with pytest.raises(ConfigError, match="must be 'none', start with 'standin' or 'fake', or be registered by a plugin"): config_from_dict(edit(standin(), dotted, "other.component")) @@ -276,7 +276,7 @@ def test_unusable_files_are_refused(bad): ("rollout.gpu_memory_utilization", 0, "rollout.gpu_memory_utilization"), ("rollout.compile_adapter", "hipcc", "rollout.compile_adapter: component name 'hipcc'"), ("rollout.parse_adapter", "none", "rollout.parse_adapter: cannot be 'none'"), - ("reward.reward_fn", "other.v1", "reward.reward_fn: expected one of 'standin.v11'"), + ("reward.reward_fn", "other.v1", "reward.reward_fn: 'other.v1' is not standin.v11 or a reward function a plugin"), ("reward.asan.max_packages", 17, "reward.asan.max_packages"), ("reward.timeouts.gate", 0, "reward.timeouts.gate: must be > 0"), ("selector.mode", "plan", "selector.mode: 'plan' is v2"), @@ -502,7 +502,7 @@ def test_a_standin_run_can_select_standin_gpu_with_its_checkers_remote(): ({"socket_dir": "relative/dir", "adapters": {"standin.aorta_asan": 5}}, "absolute path"), ({"socket_dir": "/dev/shm/c", "adapters": {}}, "name at least one adapter"), ({"socket_dir": "/dev/shm/c", "adapters": {"standin.aorta_asan": -1}}, "GPU index must be >= 0"), - ({"socket_dir": "/dev/shm/c", "adapters": {"nobody": 1}}, "must be 'none' or start with"), + ({"socket_dir": "/dev/shm/c", "adapters": {"nobody": 1}}, "must be 'none', start with 'standin' or 'fake'"), ({"socket_dir": "/dev/shm/c", "adapters": {"standin.aorta_asan": 5}, "connect_timeout_s": 0}, "must be > 0"), ({"socket_dir": "/dev/shm/c", "adapters": {"standin.aorta_asan": 5}, "extra": 1}, "unknown key"), ]) diff --git a/aorlta/tests/test_plugins.py b/aorlta/tests/test_plugins.py new file mode 100644 index 000000000..c76c16e2d --- /dev/null +++ b/aorlta/tests/test_plugins.py @@ -0,0 +1,396 @@ +"""Plugins (`aorlta.plugins`): a task's prompts, parse, feedback, compile check, profile, check adapter and reward +function, registered by a module the run config lists, and found by every place that looks a name up.""" + +from __future__ import annotations + +import json +import sys +import textwrap +from pathlib import Path + +import pytest + +from aorlta import plugins +from aorlta.config import ConfigError, config_from_dict, write_run_plan +from aorlta.ids import sha256_file + +ROOT = Path(__file__).resolve().parents[1] + +PLUGIN = ''' +from pathlib import Path + +from aorlta import plugins +from aorlta.ids import sha256_ids, sha256_text +from aorlta.reward.adapters.base import make_result +from aorlta.reward.rewardfn.base import Decided, NeedMore +from aorlta.rollout.compile_check import FakeCompiler, ParseOutcome +from aorlta.rollout.feedback import Feedback, build_feedback +from aorlta.rollout.templates import render_turn0 +from aorlta.schemas import PromptRecord, TaskRef + +HERE = Path(__file__).resolve().parent + + +def records(pool_dir, tokenizer, thinking): + out = [] + for line in (Path(pool_dir) / "prompts.jsonl").read_text().splitlines(): + row = __import__("json").loads(line) + messages = [{"role": "system", "content": "sys"}, {"role": "user", "content": row["ask"]}] + out.append(PromptRecord(prompt_id=f"toy:{row['id']}", turn0_ids=None, messages=messages, + prompt_ids_sha256=sha256_ids(render_turn0(tokenizer, messages, thinking)), + task=TaskRef(contract_id=f"toy:{row['id']}", parent_id=None, parent_healthy=None, + reference_kind=None, task_sha256=sha256_text(row["ask"])), + source="standin")) + return out + + +class Parse: + name = "toy.parse" + reply_stops = (248046,) + + def check_prompts(self, prompts, thinking, max_repair_turns=0): + return None + + def check(self, text, thinking, prompt): + return ParseOutcome("edits", code=text) if "```" in text else ParseOutcome("rejected", rejection=("no block", "x")) + + +class Fb: + name = "toy.feedback" + + def build(self, errors, max_chars): + return build_feedback(errors, max_chars) + + def build_apply(self, message, reason): + return Feedback(message, 0, 0) + + +class Adapter: + version = "1" + legs = ("ref.native",) + resource = "local_cpu" + + def __init__(self, ctx): + self.name, self.ctx = ctx.name, ctx + + def run(self, job): # half the sources pass, by hash: rewards differ within a group + t = self.ctx.clock() + if job.source_sha256[-1] in "02468ace": + return make_result(job, "clean", t_start=t, t_end=t) + return make_result(job, "fail", "wrong_output", t_start=t, t_end=t) + + def controls(self): + return [] + + def health(self): + return {"responsive": True} + + +class Reward: + name = "toy.v1" + version = "1" + code_sha256 = "0" * 64 + legs = ("ref.native",) + + def __init__(self, profile): + self.profile = profile + + def decide(self, partial, cand): + if "ref.native" not in partial: + return NeedMore(("ref.native",)) + return Decided(*self.compute(partial, cand)) + + def compute(self, legs, cand): + base = 1.0 if legs["ref.native"].state == "clean" else -1.0 + return base + 0.0 + 0.0, {"base": base, "broke_healthy": 0.0, "length": 0.0} + + def claims_clean(self, reward, parts): + return reward is not None and parts["base"] > 0 + + +plugins.register_task(plugins.Task("toy", records)) +plugins.register_parse_adapter("toy.parse", Parse) +plugins.register_feedback_template("toy.feedback", Fb) +plugins.register_compile_adapter("toy.compile", lambda workdir, evidence=None, structure=None: FakeCompiler( + evidence=evidence)) +plugins.register_profile("toy-profile", HERE / "toy-profile.yaml") +plugins.register_check_adapter("toy.check", Adapter) +plugins.register_reward_fn("toy.v1", Reward) +''' + +PROFILE = """ +name: toy-profile +reward_fn: toy.v1 +materialize: true +pools: + run: {resource: local_cpu, adapter: toy.check, timeout: compile, workers: 2} +stages: + - {name: S0, pool: run, legs: [ref.native]} +adapters: + toy.check: {tool: "${AORLTA_ROCM_ROOT}/bin/hipcc", gpus: [4, 5]} +""" + + +@pytest.fixture +def registry(): + """The registries as they were, restored after the test: a plugin's registrations are global to the process.""" + tables = [plugins.TASKS, plugins.PARSE_ADAPTERS, plugins.FEEDBACK_TEMPLATES, plugins.COMPILE_ADAPTERS, + plugins.PROFILES, plugins.REWARD_FNS, plugins.CHECK_ADAPTERS, plugins.LOADED] + saved = [dict(t) for t in tables] + yield + for entry in set(plugins.LOADED) - set(saved[-1]): # a plugin loaded here: imported afresh by the next test + sys.modules.pop(plugins.LOADED[entry]["module"], None) + for t, s in zip(tables, saved): + t.clear() + t.update(s) + + +@pytest.fixture +def plugin(tmp_path, registry) -> Path: + d = tmp_path / "toyplugin" + d.mkdir() + (d / "toy_plugin.py").write_text(PLUGIN) + (d / "toy-profile.yaml").write_text(PROFILE) + pool = tmp_path / "pool" + pool.mkdir() + (pool / "prompts.jsonl").write_text("".join(json.dumps({"id": f"p{i}", "ask": f"do {i}"}) + "\n" for i in range(6))) + return d / "toy_plugin.py" + + +def toy_config(tmp_path, plugin_entry: str, **over) -> dict: + cfg = { + "plugins": [plugin_entry], + "run": {"run_id": "toy-1", "source": "standin", "root": str(tmp_path / "runs")}, + "model": {"base_dir": str(tmp_path / "base"), "source": "standin"}, + "data": {"prompt_pools": [{"path": str(tmp_path / "pool"), "source": "standin", "task": "toy"}]}, + "rollout": {"K": 2, "P": 4, "compile_adapter": "toy.compile", "parse_adapter": "toy.parse", + "feedback_template": "toy.feedback"}, + "reward": {"profile": "toy-profile", "reward_fn": "toy.v1"}, + "orchestrator": {"iterations": 1, "run_seed": 1}, + "sync": {"dir": str(tmp_path / "shm")}, + } + for k, v in over.items(): + cfg[k] = v + return cfg + + +def test_a_config_names_what_its_plugin_registered(tmp_path, plugin): + cfg = config_from_dict(toy_config(tmp_path, str(plugin))) + assert cfg.plugins == (str(plugin),) and cfg.data.prompt_pools[0].task == "toy" + assert (cfg.reward.profile, cfg.reward.reward_fn) == ("toy-profile", "toy.v1") + assert (cfg.rollout.parse_adapter, cfg.rollout.feedback_template, cfg.rollout.compile_adapter) == ( + "toy.parse", "toy.feedback", "toy.compile") + + +def test_without_the_plugin_its_names_are_refused(tmp_path, registry): + data = toy_config(tmp_path, "unused") + del data["plugins"] + with pytest.raises(ConfigError, match="'toy' is not 'standin' or a task a plugin registered"): + config_from_dict(data) + data["data"]["prompt_pools"][0]["task"] = "standin" + with pytest.raises(ConfigError, match="'toy-profile' is not one of standin-static, standin-gpu, fake or a profile a " + "plugin registered"): + config_from_dict(data) + data["reward"] = {"profile": "fake", "reward_fn": "standin.v11"} + with pytest.raises(ConfigError, match="'toy.compile' must be 'none', start with 'standin' or 'fake', or be " + "registered by a plugin"): + config_from_dict(data) + + +@pytest.mark.parametrize("key, value, message", [ + ("task", "nobody", "task: 'nobody' is not 'standin' or a task a plugin registered"), +]) +def test_an_unregistered_task_is_refused(tmp_path, plugin, key, value, message): + data = toy_config(tmp_path, str(plugin)) + data["data"]["prompt_pools"][0][key] = value + with pytest.raises(ConfigError, match=message): + config_from_dict(data) + + +def test_an_unregistered_profile_and_reward_fn_are_refused(tmp_path, plugin): + data = toy_config(tmp_path, str(plugin)) + data["reward"]["profile"] = "nobody" + with pytest.raises(ConfigError, match="profile: 'nobody' is not one of"): + config_from_dict(data) + data = toy_config(tmp_path, str(plugin)) + data["reward"]["reward_fn"] = "standin.v11" + with pytest.raises(ConfigError, match="profile 'toy-profile' is scored by 'toy.v1'"): + config_from_dict(data) + + +def test_a_relative_plugin_path_is_read_from_the_package_root(tmp_path, registry, monkeypatch): + d = tmp_path / "pkg" + (d / "examples").mkdir(parents=True) + (d / "examples" / "p.py").write_text("from aorlta import plugins\nplugins.register_parse_adapter('rel.parse', " + "lambda: None)\n") + monkeypatch.setattr(plugins, "PACKAGE_ROOT", d) + plugins.load_plugins(["examples/p.py"]) + assert "rel.parse" in plugins.PARSE_ADAPTERS + assert plugins.identity()["examples/p.py"]["sha256"] == sha256_file(d / "examples" / "p.py") + + +def test_a_module_name_is_imported(tmp_path, registry, monkeypatch): + (tmp_path / "toy_mod_plugin.py").write_text("from aorlta import plugins\nplugins.register_feedback_template(" + "'mod.feedback', lambda: None)\n") + monkeypatch.syspath_prepend(str(tmp_path)) + plugins.load_plugins(["toy_mod_plugin"]) + assert "mod.feedback" in plugins.FEEDBACK_TEMPLATES and plugins.identity()["toy_mod_plugin"]["module"] == ( + "toy_mod_plugin") + + +@pytest.mark.parametrize("source, message", [ + ("raise RuntimeError('boom')", "importing .* failed: RuntimeError: boom"), + ("from aorlta import plugins\nplugins.register_parse_adapter('standin.parse', lambda: None)", "is built in"), + ("from aorlta import plugins\nplugins.register_profile('p', '/nonexistent/p.yaml')", "no file at"), +]) +def test_a_plugin_that_fails_stops_the_config(tmp_path, registry, source, message): + p = tmp_path / "bad.py" + p.write_text(source) + data = toy_config(tmp_path, str(p)) + with pytest.raises(ConfigError, match=message): + config_from_dict(data) + + +def test_a_missing_plugin_file_stops_the_config(tmp_path, registry): + with pytest.raises(ConfigError, match="no file at"): + config_from_dict(toy_config(tmp_path, str(tmp_path / "missing.py"))) + + +def test_a_name_is_registered_once(registry): + a, b = (lambda: None), (lambda: None) + plugins.register_parse_adapter("once.parse", a) + plugins.register_parse_adapter("once.parse", a) # the same object: a no-op + with pytest.raises(plugins.PluginError, match="already registered"): + plugins.register_parse_adapter("once.parse", b) + for bad in ("", " x", "none"): + with pytest.raises(plugins.PluginError, match="non-empty name"): + plugins.register_parse_adapter(bad, a) + + +def test_a_plugin_is_imported_once(tmp_path, registry): + p = tmp_path / "count.py" + p.write_text("import builtins\nbuiltins.__dict__.setdefault('_aorlta_count', []).append(1)\n") + import builtins + + builtins.__dict__.pop("_aorlta_count", None) + plugins.load_plugins([str(p)]) + plugins.load_plugins([str(p)]) + assert builtins.__dict__.pop("_aorlta_count") == [1] + + +def test_the_run_plan_records_each_plugin_file(tmp_path, plugin): + cfg = config_from_dict(toy_config(tmp_path, str(plugin))) + plan = json.loads(write_run_plan(cfg, commit="a" * 40).read_text()) + assert plan["plugins"][str(plugin)]["sha256"] == sha256_file(plugin) + assert plan["config"]["plugins"] == [str(plugin)] + + +def test_every_lookup_finds_the_registered_component(tmp_path, plugin, monkeypatch): + from aorlta.reward.factories import standin_factories + from aorlta.reward.profiles import load_profile + from aorlta.reward.rewardfn.for_profile import standin_rewardfn + from aorlta.rollout.chain import check_adapter_names, make_feedback_template, make_parse_adapter + from aorlta.rollout.compile_check import FakeCompiler, make_compiler + + monkeypatch.setenv("AORLTA_ROCM_ROOT", "/rocm-here") + config_from_dict(toy_config(tmp_path, str(plugin))) + check_adapter_names("toy.parse", "toy.feedback") + assert make_parse_adapter("toy.parse").name == "toy.parse" + assert make_feedback_template("toy.feedback").build_apply("m", "r").text == "m" + assert isinstance(make_compiler("toy.compile", tmp_path / "c"), FakeCompiler) + profile = load_profile("toy-profile") + assert profile.name == "toy-profile" and profile.adapters["toy.check"]["tool"] == "/rocm-here/bin/hipcc" + assert standin_rewardfn(profile).name == "toy.v1" + + class Clock: + @staticmethod + def wall(): + return 5.0 + + factories = standin_factories(profile, Clock(), timeouts={"compile": 100.0}, pools=("/pool/a",)) + adapter = factories["toy.check"]("w3") + ctx = adapter.ctx + assert (ctx.name, ctx.worker_id, ctx.pools, ctx.timeout_s) == ("toy.check", "w3", ("/pool/a",), 90.0) + assert ctx.settings["gpus"] == [4, 5] and ctx.clock() == 5.0 + + +def test_a_component_registered_under_another_name_is_refused(registry): + from aorlta.rollout.engine_pool import RolloutStop + from aorlta.rollout.chain import make_parse_adapter + + class Misnamed: + name = "other" + + plugins.register_parse_adapter("wrong.parse", Misnamed) + with pytest.raises(RolloutStop, match="is named 'other'"): + make_parse_adapter("wrong.parse") + + +def test_the_task_pool_and_its_probes(tmp_path, plugin): + from aorlta.orchestrator import wiring + + cfg = config_from_dict(toy_config(tmp_path, str(plugin))) + tok = CharChat() + pool = wiring.TaskPromptPool(plugins.task("toy"), tok, "off", tmp_path / "pool") + recs = pool.records() + assert [r.prompt_id for r in recs] == [f"toy:p{i}" for i in range(6)] + probes = wiring.task_probes(cfg, tok, plugins.task("toy"), tmp_path / "pool", recs) + assert len(probes.ids) == cfg.sync.probes == 4 and all(len(p) == cfg.sync.probe_tokens for p in probes.ids) + + +def test_the_task_pool_refuses_duplicate_ids_and_too_few_probes(tmp_path, plugin): + from aorlta.orchestrator import wiring + + cfg = config_from_dict(toy_config(tmp_path, str(plugin))) + (tmp_path / "pool" / "prompts.jsonl").write_text(json.dumps({"id": "a", "ask": "x"}) + "\n" + + json.dumps({"id": "a", "ask": "y"}) + "\n") + with pytest.raises(wiring.WiringError, match="gave a prompt id twice"): + wiring.TaskPromptPool(plugins.task("toy"), CharChat(), "off", tmp_path / "pool").records() + (tmp_path / "pool" / "prompts.jsonl").write_text(json.dumps({"id": "a", "ask": "x"}) + "\n") + recs = wiring.TaskPromptPool(plugins.task("toy"), CharChat(), "off", tmp_path / "pool").records() + with pytest.raises(wiring.WiringError, match="needs 4 probe chats"): + wiring.task_probes(cfg, CharChat(), plugins.task("toy"), tmp_path / "pool", recs) + + +def test_check_supported_takes_a_registered_reward_fn(tmp_path, plugin): + from aorlta.orchestrator import wiring + + cfg = config_from_dict(toy_config(tmp_path, str(plugin))) + wiring.check_supported(cfg) + + +class CharChat: + """The rollout's Tokenizer and the probes' Hugging Face calls, one id per character, with Qwen3.5's generation + prompt so `render_turn0` accepts it.""" + + SPECIAL = {"<|im_start|>": 248045, "<|im_end|>": 248046, "": 248068, "": 248069} + + def encode(self, text): + out, i = [], 0 + while i < len(text): + for s, t in self.SPECIAL.items(): + if text.startswith(s, i): + out.append(t) + i += len(s) + break + else: + out.append(ord(text[i]) + 10000) + i += 1 + return out + + def render(self, messages, thinking, add_generation_prompt=True): + return self.apply_chat_template(messages, enable_thinking=thinking == "on") + + def __call__(self, text, add_special_tokens=False): + return {"input_ids": self.encode(text)} + + def apply_chat_template(self, msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False): + body = "".join(f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n" for m in msgs) + return body + "<|im_start|>assistant\n" + ("\n" if enable_thinking else "\n\n\n\n") + + +def test_the_module_docstring_names_every_registry(): + doc = textwrap.dedent(plugins.__doc__) + for fn in ("register_task", "register_parse_adapter", "register_feedback_template", "register_compile_adapter", + "register_profile", "register_reward_fn", "register_check_adapter"): + assert fn in doc and callable(getattr(plugins, fn)) From 64f63c3dbcd396c2b91d6e9028c2ffe201231d31 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Thu, 12 Nov 2026 17:50:59 +0000 Subject: [PATCH 02/19] fix(aorlta): self-test a plugin task's engines on open-ended chats Each engine's start-up self-test samples 4 chats with two seeds and needs 3 of the 4 pairs to differ. On the first batch's own prompts that fails a sound sampler when the task is precise: on MI350X, the hip-kernels example's prompts (write a vector add, a SAXPY, ...) gave the same 128 tokens for both seeds on 2 of 4 prompts, on every engine. A plugin task now self-tests on fixed open-ended chats (wiring.SELF_TEST_CHATS), or on its own (Task.selftest_messages). The stand-in path is unchanged. --- aorlta/docs/ADAPTING.md | 5 +- aorlta/src/aorlta/orchestrator/wiring.py | 34 +++++++++++-- aorlta/src/aorlta/plugins.py | 7 ++- .../tests/orchestrator/test_plugin_wired.py | 51 +++++++++++++++++++ 4 files changed, 91 insertions(+), 6 deletions(-) diff --git a/aorlta/docs/ADAPTING.md b/aorlta/docs/ADAPTING.md index 3789fac01..a6f6c00e3 100644 --- a/aorlta/docs/ADAPTING.md +++ b/aorlta/docs/ADAPTING.md @@ -38,7 +38,10 @@ The module is imported before the rest of the config is checked, so every name a A built-in component may be mixed in by name (`standin.parse`, `standin.hipcc` and `standin` are the stand-in parse, compile and feedback). A component's `name` must be the name it is registered under. The landing check and the trainer's warm-up read `sync.probes` chats: the task's `probe_messages(pool_dir)`, or the first records' own -messages, which must then share one system message. The run plan records each plugin's file and its sha256. +messages, which must then share one system message. Each engine's start-up self-test samples 4 chats twice with +different seeds and needs at least 3 of the 4 pairs to differ, so for a plugin task it uses fixed open-ended chats +(`SELF_TEST_CHATS` in `orchestrator/wiring.py`), not the task's own prompts, whose answers may be near-unique; +`Task(selftest_messages=...)` gives others. The run plan records each plugin's file and its sha256. Limits of a plugin: its data is handled as the stand-in source's (public; text may appear in logs and tracebacks); its reward stays within the stand-in range and parts (below); the episode is the compile-repair chain (a parse may diff --git a/aorlta/src/aorlta/orchestrator/wiring.py b/aorlta/src/aorlta/orchestrator/wiring.py index ea0674fcb..512539093 100644 --- a/aorlta/src/aorlta/orchestrator/wiring.py +++ b/aorlta/src/aorlta/orchestrator/wiring.py @@ -48,6 +48,14 @@ PROBE_KERNELS = 40 # WP3's probes are the last `sync.probes` of these stand-in kernels WORKER_THREADS = 8 # CPU threads per spawned trainer rank: its CPU work is small, and the compile pool shares the node ENGINE_CALL_TIMEOUT_S = 600.0 # a probe, a greedy draw or a reload, as EngineProc's own reload default +# A plugin task's engine self-test chats (plugins.Task.selftest_messages): open-ended, so two seeds' first 128 tokens +# differ on a sound sampler whatever the task's own prompts are. +SELF_TEST_CHATS = tuple( + ({"role": "system", "content": "You are an expert AMD GPU kernel engineer."}, {"role": "user", "content": q}) + for q in ("Describe, in your own words, three different reasons a GPU kernel can run slower than expected.", + "Write a short story, of about a paragraph, about a GPU that is learning to schedule its own work.", + "List some ideas for small exercises that teach someone how GPU memory coalescing works.", + "Explain to a new engineer, with an example of your choosing, what a race condition in a kernel is.")) FAKE_FAIL_RATE = 0.05 # per leg, in the fake profile: enough failures that rewards differ within a group STANDIN_REWARD_FNS = ("standin.v11",) MERGE_TOOL = Path(__file__).resolve().parents[3] / "tools" / "merge_base_adapter.py" @@ -320,10 +328,11 @@ class RolloutAdapter: """C1 behind the RolloutPort.""" def __init__(self, cfg: RunConfig, tokenizer, handles: Sequence[Any], *, max_len: int, spans, compiler, - on_shutdown: Sequence[Callable[[], None]] = (), clock=None): + on_shutdown: Sequence[Callable[[], None]] = (), clock=None, selftest_messages: Sequence | None = None): from ..rollout.engine_pool import EnginePool self.cfg, self.tok, self.spans, self.max_len, self.compiler = cfg, tokenizer, spans, max_len, compiler + self.selftest_messages = None if selftest_messages is None else list(selftest_messages) # else the prompts self.clock = clock # the orchestrator's clock: its deadlines are read on it self.handles = [CountedHandle(h) for h in handles] self.pool = EnginePool(self.handles, run_seed=cfg.orchestrator.run_seed, @@ -336,9 +345,15 @@ def start(self, selftest_prompts: Sequence) -> dict: from ..rollout.engine_pool import RolloutStop from ..rollout.templates import turn0_ids + from ..rollout.templates import render_turn0 + r = self.cfg.rollout try: - self.pool.start([turn0_ids(p, self.tok, r.thinking) for p in selftest_prompts]) + if self.selftest_messages is not None: + ids = [render_turn0(self.tok, m, r.thinking) for m in self.selftest_messages] + else: + ids = [turn0_ids(p, self.tok, r.thinking) for p in selftest_prompts] + self.pool.start(ids) except RolloutStop as exc: raise RunStopped("engine_start_failed", {"error": describe_error(exc, self.cfg.run.source)}) from None settings = RolloutSettings.from_config(self.cfg, self.max_len) @@ -1263,13 +1278,23 @@ def build(cfg: RunConfig, *, config_path=None, commit: str | None = None, parts: probes = parts.probes if probes is None and task is not None: probes = task_probes(cfg, tok, task, pool_cfg.path, records) + selftest = None + if task is not None: + selftest = (list(task.selftest_messages(Path(pool_cfg.path))) if task.selftest_messages is not None + else list(SELF_TEST_CHATS)) + from ..rollout.engine_pool import SELF_TEST_PROMPTS + + if len(selftest) < SELF_TEST_PROMPTS: + raise WiringError(f"task {task.name!r}: the engine self-test needs {SELF_TEST_PROMPTS} chats, " + f"got {len(selftest)}") + selftest = selftest[:SELF_TEST_PROMPTS] compiler = parts.compiler if parts.compiler is not None else rollout_compiler(cfg) handles = parts.handles if parts.handles is not None else engine_handles(cfg, max_len) if cfg.orchestrator.mode == "score": # no trainer, no optimizer, no weight sync from .score import Scorer rollout = RolloutAdapter(cfg, tok, handles, max_len=max_len, spans=spans, compiler=compiler, - on_shutdown=[lambda: remove_node_dir(cfg)], clock=clock) + on_shutdown=[lambda: remove_node_dir(cfg)], clock=clock, selftest_messages=selftest) model = ServedModel(cfg, tok, rollout, spans=spans, dataset=dataset, probes=probes) reward = RewardAdapter(cfg, clock=clock, spans=spans, factories=parts.reward_factories) return Scorer(cfg, pool=pool, rollout=rollout, reward=reward, model=model, clock=clock, spans=spans, @@ -1277,7 +1302,8 @@ def build(cfg: RunConfig, *, config_path=None, commit: str | None = None, parts: trainer = TrainerAdapter(cfg, spans=spans, load_model=parts.load_model, device=parts.trainer_device, require_fast_path=parts.require_fast_path, amp=parts.amp) rollout = RolloutAdapter(cfg, tok, handles, max_len=max_len, spans=spans, compiler=compiler, - on_shutdown=[trainer.shutdown, lambda: remove_node_dir(cfg)], clock=clock) + on_shutdown=[trainer.shutdown, lambda: remove_node_dir(cfg)], clock=clock, + selftest_messages=selftest) sync = SyncAdapter(cfg, tok, trainer, rollout, spans=spans, dataset=dataset, probes=probes) trainer.warmup_ids = lambda: sync.probes().ids # the landing probes: real stand-in prompt tokens reward = RewardAdapter(cfg, clock=clock, spans=spans, factories=parts.reward_factories) diff --git a/aorlta/src/aorlta/plugins.py b/aorlta/src/aorlta/plugins.py index ce07d23fa..76689dc2e 100644 --- a/aorlta/src/aorlta/plugins.py +++ b/aorlta/src/aorlta/plugins.py @@ -48,11 +48,16 @@ class Task: `records(pool_dir, tokenizer, thinking)` returns the pool's `PromptRecord`s (`schemas.py`), each with `source="standin"`, a unique `prompt_id` and a `task.task_sha256` that hashes what the checks read for it. `probe_messages(pool_dir)`, if given, returns the chats (a system and a user message) the landing check and the - trainer's warm-up use; without it they are the first records' own messages.""" + trainer's warm-up use; without it they are the first records' own messages. `selftest_messages(pool_dir)`, if + given, returns the 4 chats each engine's start-up self-test samples twice with different seeds; at least 3 of the + 4 pairs must differ, so they must be open-ended requests. Without it the self-test uses + `orchestrator.wiring.SELF_TEST_CHATS`: a precise task (a short kernel with one obvious answer) can give the same + 128 tokens for both seeds and fail a sound sampler.""" name: str records: Callable[[Path, Any, str], list] probe_messages: Callable[[Path], list[list[dict]]] | None = None + selftest_messages: Callable[[Path], list[list[dict]]] | None = None @dataclass(frozen=True) diff --git a/aorlta/tests/orchestrator/test_plugin_wired.py b/aorlta/tests/orchestrator/test_plugin_wired.py index 699e70b75..cecc33c79 100644 --- a/aorlta/tests/orchestrator/test_plugin_wired.py +++ b/aorlta/tests/orchestrator/test_plugin_wired.py @@ -4,6 +4,7 @@ from __future__ import annotations +import dataclasses import json import sys from pathlib import Path @@ -15,6 +16,7 @@ import test_plugins as TP # noqa: E402 - the toy plugin import test_wiring_e2e as E # noqa: E402 - the tiny model, its twin engines and the probes +from aorlta import plugins # noqa: E402 from aorlta.config import config_from_dict # noqa: E402 from aorlta.orchestrator import wiring # noqa: E402 from aorlta.rollout.engine_proc import InProcessEngine # noqa: E402 @@ -52,3 +54,52 @@ def test_a_plugin_task_trains_and_lands(tmp_path, tiny_base, plugin): assert any(len(e["turns"]) > 1 for e in eps) # a repair turn: toy.feedback's compile or apply turn plan = json.loads((cfg.run_dir / "run_plan.json").read_text()) assert plan["plugins"][str(plugin)]["sha256"] == TP.sha256_file(plugin) + + +def started_on(monkeypatch) -> list: + """The turn-0 ids each engine pool's self-test was started on.""" + from aorlta.rollout import engine_pool + + seen, real = [], engine_pool.EnginePool.start + + def start(self, prompt_ids, *a, **kw): + seen.append([list(p) for p in prompt_ids]) + return real(self, prompt_ids, *a, **kw) + + monkeypatch.setattr(engine_pool.EnginePool, "start", start) + return seen + + +def score_config(tmp_path, base, plugin_path): + data = TP.toy_config(tmp_path, str(plugin_path)) + data["model"]["base_dir"] = str(base) + data["rollout"].update({"engines": [0, 1], "K": 4, "P": 4, "max_new_tokens": 512, "request_timeout_s": 300}) + data["orchestrator"] = {"mode": "score", "run_seed": 3, "score": {"prompts": 4, "passes": 1}} + return config_from_dict(data) + + +def test_a_plugin_task_self_tests_its_engines_on_open_ended_chats(tmp_path, tiny_base, plugin, monkeypatch): + """Not on the task's own prompts: a precise task can give two seeds the same first 128 tokens.""" + from aorlta.rollout.templates import render_turn0 + + seen = started_on(monkeypatch) + cfg = score_config(tmp_path, tiny_base, plugin) + result = build(cfg, tiny_base).run() + assert result["status"] == "completed", result + tok = E.ByteTokenizer() + assert seen == [[render_turn0(tok, list(m), "off") for m in wiring.SELF_TEST_CHATS]] + + +def test_a_task_may_give_its_own_self_test_chats(tmp_path, tiny_base, plugin, monkeypatch): + from aorlta.rollout.templates import render_turn0 + + chats = [[{"role": "user", "content": f"tell me anything at all, number {i}"}] for i in range(5)] + cfg = score_config(tmp_path, tiny_base, plugin) + toy = plugins.TASKS["toy"] + plugins.TASKS["toy"] = dataclasses.replace(toy, selftest_messages=lambda pool_dir: chats) + seen = started_on(monkeypatch) + assert build(cfg, tiny_base).run()["status"] == "completed" + assert seen == [[render_turn0(E.ByteTokenizer(), c, "off") for c in chats[:4]]] + plugins.TASKS["toy"] = dataclasses.replace(toy, selftest_messages=lambda pool_dir: chats[:3]) + with pytest.raises(wiring.WiringError, match="the engine self-test needs 4 chats, got 3"): + build(cfg, tiny_base) From ece5476d402769bee9601ebbd8ba6218ce4e06cc Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Thu, 12 Nov 2026 17:41:56 +0000 Subject: [PATCH 03/19] feat(aorlta): a worked example of a task as a plugin, and a tutorial examples/hip-kernels/ asks the model to write a HIP kernel from its description (vector add, SAXPY, ReLU, scale and shift, clamp, reverse, a per-block sum, a transpose). Its parse gives a repair turn when the reply has no code block or does not define the kernel asked for. Its check builds the answer with a generated host program, runs it on a GPU and compares every output element with a host reference; an element never written, and a write past an array's end, are caught. hipk.v1 scores the outcome. Each task comes with a correct reference kernel, which 'hip_kernels.py selfcheck' builds and runs. docs/TUTORIAL.md walks through it: what to write, where to register it, and how to check, test, score and train it. configs/hipk-score.yaml and hipk-train.yaml run it. It needs none of the stand-in sanitizer tools, only the run image's hipcc and GPUs. --- aorlta/README.md | 8 +- aorlta/configs/hipk-score.yaml | 36 ++ aorlta/configs/hipk-train.yaml | 38 ++ aorlta/docs/TUTORIAL.md | 150 ++++++ aorlta/examples/hip-kernels/README.md | 15 + aorlta/examples/hip-kernels/hip_kernels.py | 449 ++++++++++++++++++ aorlta/examples/hip-kernels/hipk-gpu.yaml | 26 + .../hip-kernels/reference/block_sum.hip | 16 + .../examples/hip-kernels/reference/clamp.hip | 6 + .../examples/hip-kernels/reference/relu.hip | 6 + .../hip-kernels/reference/reverse.hip | 6 + .../examples/hip-kernels/reference/saxpy.hip | 6 + .../hip-kernels/reference/scale_shift.hip | 6 + .../hip-kernels/reference/transpose.hip | 7 + .../hip-kernels/reference/vector_add.hip | 6 + aorlta/examples/hip-kernels/tasks.json | 99 ++++ aorlta/tests/examples/test_hip_kernels.py | 367 ++++++++++++++ 17 files changed, 1244 insertions(+), 3 deletions(-) create mode 100644 aorlta/configs/hipk-score.yaml create mode 100644 aorlta/configs/hipk-train.yaml create mode 100644 aorlta/docs/TUTORIAL.md create mode 100644 aorlta/examples/hip-kernels/README.md create mode 100644 aorlta/examples/hip-kernels/hip_kernels.py create mode 100644 aorlta/examples/hip-kernels/hipk-gpu.yaml create mode 100644 aorlta/examples/hip-kernels/reference/block_sum.hip create mode 100644 aorlta/examples/hip-kernels/reference/clamp.hip create mode 100644 aorlta/examples/hip-kernels/reference/relu.hip create mode 100644 aorlta/examples/hip-kernels/reference/reverse.hip create mode 100644 aorlta/examples/hip-kernels/reference/saxpy.hip create mode 100644 aorlta/examples/hip-kernels/reference/scale_shift.hip create mode 100644 aorlta/examples/hip-kernels/reference/transpose.hip create mode 100644 aorlta/examples/hip-kernels/reference/vector_add.hip create mode 100644 aorlta/examples/hip-kernels/tasks.json create mode 100644 aorlta/tests/examples/test_hip_kernels.py diff --git a/aorlta/README.md b/aorlta/README.md index e69ab3b81..af2a18d68 100644 --- a/aorlta/README.md +++ b/aorlta/README.md @@ -35,8 +35,9 @@ Not there yet: - No trained model worth sharing: only 3-iteration stand-in runs. The stand-in prompts show a kernel's original disassembly, with no bug planted, and ask for a HIP source, so they exercise the loop, not the task. - The stand-in checks are cheap, so the times above say little about a run with heavier checks. -- One prompt format, the stand-in one, and one node layout are wired; [docs/ADAPTING.md](docs/ADAPTING.md) says how - to add a prompt source, a reward function or a check. +- One node layout is wired. A task of your own (prompts, answer parsing, checks and reward) is a plugin, with no + edit to the package: [docs/ADAPTING.md](docs/ADAPTING.md) has the contract, and [docs/TUTORIAL.md](docs/TUTORIAL.md) + walks through `examples/hip-kernels/`, a kernel-writing task whose reward runs the kernel and checks its output. ## Layout @@ -58,7 +59,8 @@ and `configs/dev-standin.yaml` (one trainer rank and the CPU checks of `standin- compile and waitcheck failures, so a candidate that passes them stays unscored); `jobs/aorlta-run.sbatch` (a run), `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; `docs/ADAPTING.md`. Code comments cite sections and +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 labels (Section 6.x, C1-C9, D7, ...) of our internal design notes, which are not included. ## Requirements diff --git a/aorlta/configs/hipk-score.yaml b/aorlta/configs/hipk-score.yaml new file mode 100644 index 000000000..0db985b5b --- /dev/null +++ b/aorlta/configs/hipk-score.yaml @@ -0,0 +1,36 @@ +# Score the public Qwen3.5-4B on the example task of docs/TUTORIAL.md (examples/hip-kernels): 8 kernels to write, +# 4 samples each, 2 passes, up to 2 repair turns. 4 engines on GPUs 0-3; the kernels run on GPUs 4-7 (hipk-gpu.yaml). +# 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/hip-kernels. +plugins: [examples/hip-kernels/hip_kernels.py] +run: + run_id: hipk-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/hip-kernels + source: standin + task: hipk +rollout: + engines: [0, 1, 2, 3] + K: 4 + P: 8 + max_repair_turns: 2 + thinking: "off" + temperature: 0.7 + max_new_tokens: 2048 + compile_adapter: standin.hipcc + parse_adapter: hipk.parse + feedback_template: hipk.feedback +reward: + profile: hipk-gpu + reward_fn: hipk.v1 +orchestrator: + mode: score + run_seed: 20261112 + score: + passes: 2 diff --git a/aorlta/configs/hipk-train.yaml b/aorlta/configs/hipk-train.yaml new file mode 100644 index 000000000..51953aaf5 --- /dev/null +++ b/aorlta/configs/hipk-train.yaml @@ -0,0 +1,38 @@ +# Train on the example task of docs/TUTORIAL.md (examples/hip-kernels): 3 iterations of the 8 kernels x 8 samples, +# one LoRA rank beside engine 0, the kernels run on GPUs 4-7 (hipk-gpu.yaml). A starting point to copy, not a result: +# 8 prompts are far too few to learn from. +# 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/hip-kernels. +plugins: [examples/hip-kernels/hip_kernels.py] +run: + run_id: hipk-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/hip-kernels + source: standin + task: hipk +rollout: + engines: [0, 1, 2, 3] + K: 8 + P: 8 + max_repair_turns: 2 + thinking: "off" + temperature: 0.7 + max_new_tokens: 2048 + compile_adapter: standin.hipcc + parse_adapter: hipk.parse + feedback_template: hipk.feedback +reward: + profile: hipk-gpu + reward_fn: hipk.v1 +orchestrator: + iterations: 3 + run_seed: 20261112 +trainer: + ranks: 1 + micro_batch_tokens: 20480 diff --git a/aorlta/docs/TUTORIAL.md b/aorlta/docs/TUTORIAL.md new file mode 100644 index 000000000..6210144e1 --- /dev/null +++ b/aorlta/docs/TUTORIAL.md @@ -0,0 +1,150 @@ +# Tutorial: your own task, end to end + +This walks through `examples/hip-kernels/`, a small task built as a plugin: the model is asked to write a HIP kernel +from its description (vector add, SAXPY, ReLU, scale and shift, clamp, reverse, a per-block sum, a transpose), and +the reward builds its answer, runs it on a GPU and compares every output element with a host reference. Nothing in +`src/aorlta/` names this task. Copy the directory, change the four parts below, and you have your own. + +Paths are relative to `aorlta/`. The commands assume the environment from the README (a Python 3.12 venv with +`requirements-dev.txt` and `pip install -e . --no-deps`). + +## What you write + +| part | in the example | what it does | +|---|---|---| +| the prompts | `examples/hip-kernels/tasks.json`, and `records()` in `hip_kernels.py` | one prompt record per task: the chat, its prompt id, and a hash of what the checks read for it | +| how an answer is read | `HipkParse` and `HipkFeedback` | the last ```cpp block is the answer; no block, or no definition of the kernel asked for, gets a repair turn that says so; a compile error gets one with the compiler's errors | +| the checks | `HipkRun` and `examples/hip-kernels/hipk-gpu.yaml` | the stand-in parse and compile gates, then `ref.native`: build the answer with a generated host program, run it on a GPU, compare | +| the reward | `HipkV1` | -1.0 no parsable answer, -0.9 compile error, -0.8 wrong signature, -0.6 crash or hang, -0.5 wrong output, +1.0 correct, minus 0.05 x the share of the token budget used | + +The module registers each part under a name (`register()` at its end): + +```python +plugins.register_task(plugins.Task("hipk", records)) +plugins.register_parse_adapter("hipk.parse", HipkParse) +plugins.register_feedback_template("hipk.feedback", HipkFeedback) +plugins.register_profile("hipk-gpu", HERE / "hipk-gpu.yaml") +plugins.register_check_adapter("hipk.run", HipkRun) +plugins.register_reward_fn("hipk.v1", HipkV1) +``` + +and the run config lists the module and uses the names (`configs/hipk-score.yaml`): + +```yaml +plugins: [examples/hip-kernels/hip_kernels.py] +data: + prompt_pools: + - {path: /path/to/aorta/aorlta/examples/hip-kernels, source: standin, task: hipk} +rollout: + compile_adapter: standin.hipcc # a built-in part, mixed in by name + parse_adapter: hipk.parse + feedback_template: hipk.feedback +reward: + profile: hipk-gpu + reward_fn: hipk.v1 +``` + +[ADAPTING.md](ADAPTING.md) has the full contract of each part. The rules that matter most: + +- **Prompts.** `records(pool_dir, tokenizer, thinking)` returns `PromptRecord`s with `source="standin"` (a plugin's + data is handled as public stand-in data: its text may appear in logs). `task.task_sha256` must change whenever what + a check reads for that prompt changes: it is part of every cache key, and `hipk.run` finds the task by it. + Training takes `rollout.P` different prompts per iteration, so the pool needs at least `P` of them; the landing + check uses the first `sync.probes` (4) prompts' chats, which must share one system message. +- **A check never guesses.** Anything it cannot settle (no GPU, a tool missing, output it cannot read) is + `unsettled`, never `clean`; only the check may blame the candidate for a timeout. `verdict()` in the example is + the whole mapping, and `tests/examples/test_hip_kernels.py` tests every row of it on CPU. +- **The reward's parts.** `compute` returns `(reward, {"base", "broke_healthy", "length"})`, which add up to the + reward, within [-1.55, 1.0]; an unscored sample is `(None, {"reason": ...})` and is never counted clean. `decide` + answers early only when no pending leg could change the answer. The service recomputes and checks all of this, + and a reward that breaks it is unscored and raises an alarm. +- **The check's GPUs** are its own setting (`hipk.run.gpus` in the profile): keep them off `rollout.engines`. + +## 1. Check the tasks + +Each task in `tasks.json` comes with a correct reference kernel in `reference/.hip`. Build every task's harness +with its reference kernel, on CPU: + +```bash +python examples/hip-kernels/hip_kernels.py selfcheck --hipcc /opt/rocm/bin/hipcc --build-only +``` + +and run them on one GPU (one line per task, all `"pass": true`; on a cluster, inside a job): + +```bash +python examples/hip-kernels/hip_kernels.py selfcheck --hipcc /opt/rocm/bin/hipcc --gpu 0 +``` + +The harness is the candidate's source followed by a host `main` that fills the inputs from a fixed seed, runs the +kernel once with the launch the prompt states, and compares every output element with the task's `reference` loop +(host C++), within `atol + rtol * |want|`. Output elements start as a sentinel, so an element the kernel never +writes is a mismatch, and every array has a guard region past its end, so a write past the end is caught too. + +## 2. Run the tests + +```bash +OMP_NUM_THREADS=1 python -m pytest -q tests/examples/test_hip_kernels.py tests/test_plugins.py +``` + +They load the configs with the plugin, build the prompts and probes, run the parse and feedback, run `hipk.run` +against a stand-in `hipcc` (no GPU), and check the reward against the service's invariants for every combination of +leg outcomes. With hipcc installed, they also build every harness. + +## 3. The loop on CPU, with fake components + +`--components fake` loads the config and the plugin exactly as a real run does, then runs the loop with fake engines, +checks and trainer: it tells you the config is consistent and every name resolves, not that your task works (the +fakes use prompts and checks of their own). + +```bash +export SCRATCH=/path/to/scratch +cat > "$SCRATCH/site.yaml" <.hip` | a correct kernel for each task, used only by `selfcheck` | +| `hipk-gpu.yaml` | the check profile: parse gate, compile gate, `ref.native` on GPUs 4-7 | + +The run configs are `configs/hipk-score.yaml` and `configs/hipk-train.yaml`; the tests are +`tests/examples/test_hip_kernels.py`. diff --git a/aorlta/examples/hip-kernels/hip_kernels.py b/aorlta/examples/hip-kernels/hip_kernels.py new file mode 100644 index 000000000..3a9a1a2be --- /dev/null +++ b/aorlta/examples/hip-kernels/hip_kernels.py @@ -0,0 +1,449 @@ +"""A worked example of a task of your own, as a plugin (docs/TUTORIAL.md): write a HIP kernel from its description. + +Each task in `tasks.json` names a kernel, its exact signature, what it computes and how it is launched. The model +answers with a HIP source in one ```cpp block. The rollout compiles it with the stand-in compile check (hipcc -c) and +gives up to `rollout.max_repair_turns` repair turns with the compiler's errors; a reply that has no code block, or +whose code does not define the kernel, gets a repair turn that says so (this module's parse and feedback). The reward +service then runs three legs: the stand-in parse gate and compile gate, and `ref.native`, this module's check, which +builds the candidate with a generated host program, runs it on a GPU and compares every output element with a host +reference computed from the same inputs. The reward is `hipk.v1` (below). + +Registered names: task `hipk`; parse `hipk.parse`; feedback `hipk.feedback`; profile `hipk-gpu` (hipk-gpu.yaml); +check adapter `hipk.run`; reward function `hipk.v1`. + +Run as a script, it checks the tasks themselves with the reference kernels in reference/: + python examples/hip-kernels/hip_kernels.py selfcheck --hipcc /opt/rocm/bin/hipcc [--gpu 0] [--build-only] +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import shutil +import sys +import tempfile +import time +from pathlib import Path + +HERE = Path(__file__).resolve().parent +SYSTEM = "You are an expert AMD GPU kernel engineer." +USER = ( + "Write a HIP kernel with exactly this signature:\n\n```cpp\n{signature}\n```\n\n" + "It must compute {description}. It is launched {launch}.\n\n" + "Write a complete HIP C++ source file: include , define the kernel, and write no main " + "function. Reply with the source in a single ```cpp code block." +) +RESULT_TAG = "AORLTA_RESULT " +SENTINEL = "-7777.0f" # what an output element holds before the kernel runs: one it never writes is a mismatch +GUARD = 256 # elements past the end of every device array, which the kernel must leave alone +GUARD_VALUE = "31337.0f" + + +# --- the tasks ------------------------------------------------------------------------------------- + + +def load_tasks(pool_dir: str | Path) -> list[dict]: + """The pool's tasks, in file order: `/tasks.json`, a list of task specs.""" + tasks = json.loads((Path(pool_dir) / "tasks.json").read_text(encoding="utf-8")) + ids = [t["id"] for t in tasks] + if len(set(ids)) != len(ids): + raise ValueError(f"{pool_dir}/tasks.json lists a task id twice") + return tasks + + +def task_sha256(task: dict) -> str: + """What the check reads for a task, hashed: part of every check's cache key, and how the check finds the task.""" + return hashlib.sha256(json.dumps(task, sort_keys=True).encode("utf-8")).hexdigest() + + +def messages(task: dict) -> list[dict]: + user = USER.format(signature=task["signature"], description=task["description"], launch=task["launch"]) + return [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user}] + + +def records(pool_dir, tokenizer, thinking): + """The prompt records: one per task, prompt id `hipk:`.""" + from aorlta.ids import sha256_ids + from aorlta.rollout.templates import render_turn0 + from aorlta.schemas import PromptRecord, TaskRef + + out = [] + for task in load_tasks(pool_dir): + msgs = messages(task) + out.append(PromptRecord( + prompt_id=f"hipk:{task['id']}", turn0_ids=None, messages=msgs, + prompt_ids_sha256=sha256_ids(render_turn0(tokenizer, msgs, thinking)), + task=TaskRef(contract_id=f"hipk:{task['id']}", parent_id=None, parent_healthy=None, reference_kind=None, + task_sha256=task_sha256(task)), + source="standin")) + return out + + +# --- the rollout: parse and feedback ------------------------------------------------------------ + + +def defines_kernel(source: str, kernel: str) -> bool: + return re.search(r"__global__\s+(?:__launch_bounds__\s*\([^)]*\)\s*)?void\s+" + re.escape(kernel) + r"\s*\(", + source) is not None + + +class HipkParse: + """`hipk.parse`: the last ```cpp block (the stand-in parse's rule). No block, or a block that does not define the + task's kernel, is rejected with a message, which `hipk.feedback` turns into a repair turn.""" + + name = "hipk.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): + for p in prompts: + if not p.prompt_id.startswith("hipk:"): + raise ValueError(f"hipk.parse reads hipk prompts, got {p.prompt_id!r}") + + def check(self, text, thinking, prompt): + from aorlta.rollout.compile_check import ParseOutcome, parse_code + + code = parse_code(text, thinking) + if code is None: + return ParseOutcome("rejected", rejection=( + "Your reply has no complete ```cpp code block. Reply with the whole HIP source in one ```cpp block.", + "no_code_block")) + kernel = prompt.prompt_id.split(":", 1)[1] + if not defines_kernel(code, kernel): + return ParseOutcome("rejected", rejection=( + f"Your source does not define the kernel `{kernel}`. Define it as a __global__ function with exactly " + f"the signature asked for, and reply with the whole source in one ```cpp block.", "no_kernel")) + return ParseOutcome("edits", code=code) + + +class HipkFeedback: + """`hipk.feedback`: the stand-in compile feedback, and a rejected answer's message as its own repair turn.""" + + name = "hipk.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) + + +# --- the check: build, run, compare ------------------------------------------------------------ + + +def _f(x) -> str: + """A float as a C++ float literal.""" + return f"{float(x)!r}f" + + +def harness_source(task: dict, kernel_source: str) -> str: + """The candidate's source followed by a host program that fills the inputs, runs the kernel once, and prints one + `AORLTA_RESULT {json}` line: `ok`, `mismatch` (an output element differs from the host reference beyond the + tolerance), `out_of_bounds` (a write past the end of an array) or `hip_error`.""" + lines = [kernel_source.rstrip(), "", "// ---- AORLTA harness ----", "#include ", "#include ", + "#include ", "#include ", "", + "static void fill(std::vector& v, uint32_t seed, float lo, float hi) {", + " for (auto& x : v) { seed = seed * 1664525u + 1013904223u; " + "x = lo + (hi - lo) * float(seed >> 8) / 16777216.0f; }", + "}", "", + "#define CHECK(call) do { hipError_t e_ = (call); if (e_ != hipSuccess) { " + f"printf(\"{RESULT_TAG}{{\\\"status\\\": \\\"hip_error\\\", \\\"at\\\": \\\"%s\\\"}}\\n\", #call); " + "return 0; } } while (0)", "", + "int main() {"] + for ctype, name, value in task["scalars"]: + lines.append(f" const {ctype} {name} = {value};") + arrays = task["arrays"] + for i, a in enumerate(arrays): + n = a["size"] + lines.append(f" std::vector h_{a['name']}(({n}) + {GUARD}, {GUARD_VALUE});") + if a["role"] == "out": + lines.append(f" for (long i = 0; i < ({n}); ++i) h_{a['name']}[i] = {SENTINEL};") + else: + lines.append(f" {{ std::vector v({n}); fill(v, {1000 + 17 * i}u, {_f(a.get('lo', -1.0))}, " + f"{_f(a.get('hi', 1.0))}); for (long i = 0; i < ({n}); ++i) h_{a['name']}[i] = v[i]; }}") + lines.append(f" std::vector r_{a['name']} = h_{a['name']};") + lines.append(" {") + for a in arrays: + lines.append(f" float* {a['name']} = r_{a['name']}.data();") + lines.append(" " + task["reference"]) + lines.append(" }") + for a in arrays: + lines += [f" float* d_{a['name']} = nullptr;", + f" CHECK(hipMalloc(&d_{a['name']}, h_{a['name']}.size() * sizeof(float)));", + f" CHECK(hipMemcpy(d_{a['name']}, h_{a['name']}.data(), h_{a['name']}.size() * sizeof(float), " + "hipMemcpyHostToDevice));"] + names = {a["name"] for a in arrays} + args = ", ".join(f"d_{x}" if x in names else x for x in task["args"]) + lines += [f" {task['kernel']}<<<{task['grid']}, {task['block']}>>>({args});", + " CHECK(hipGetLastError());", " CHECK(hipDeviceSynchronize());", + " long checked = 0, mismatches = 0, first = -1, oob = 0;"] + for a in arrays: + n = a["size"] + lines += [f" CHECK(hipMemcpy(h_{a['name']}.data(), d_{a['name']}, h_{a['name']}.size() * sizeof(float), " + "hipMemcpyDeviceToHost));", + f" for (long i = ({n}); i < (long)h_{a['name']}.size(); ++i) if (h_{a['name']}[i] != {GUARD_VALUE}) " + "++oob;"] + if a["role"] != "in": + lines += [f" for (long i = 0; i < ({n}); ++i) {{", + f" float got = h_{a['name']}[i], want = r_{a['name']}[i];", + f" bool same = (std::isnan(got) && std::isnan(want)) || std::fabs(got - want) <= " + f"{_f(task['atol'])} + {_f(task['rtol'])} * std::fabs(want);", + " ++checked;", + " if (!same) { if (first < 0) first = checked - 1; ++mismatches; }", + " }"] + lines += [" const char* status = oob ? \"out_of_bounds\" : (mismatches ? \"mismatch\" : \"ok\");", + f" printf(\"{RESULT_TAG}{{\\\"status\\\": \\\"%s\\\", \\\"checked\\\": %ld, \\\"mismatches\\\": %ld, " + "\\\"first\\\": %ld, \\\"out_of_bounds\\\": %ld}\\n\", status, checked, mismatches, first, oob);", + " return 0;", "}", ""] + return "\n".join(lines) + + +def read_result(stdout: str) -> dict | None: + """The harness's result line, or None.""" + for line in reversed(stdout.splitlines()): + if line.startswith(RESULT_TAG): + try: + out = json.loads(line[len(RESULT_TAG):]) + except ValueError: + return None + return out if isinstance(out, dict) and isinstance(out.get("status"), str) else None + return None + + +def verdict(build_rc, build_stderr, run, task) -> tuple[str, str | None, list[dict]]: + """A leg's (state, reason, findings) from the harness's build and run. `run` is None when the build failed, else + a ProcResult. Only what is the candidate's is blamed on it: a build failure is the candidate's when the compiler + names its kernel or `main` (its compile alone passed in gate.compile, so the signature is what differs), a hang + past the limit and a crash are its, and anything else is unsettled, never clean.""" + if build_rc != 0: + if re.search(r"\b(" + re.escape(task["kernel"]) + r"|main)\b", build_stderr or ""): + return "fail", "bad_signature", [] + return "unsettled", "harness_build_failed", [] + if run.launch_error or run.killed: + return "unsettled", "run_not_started", [] + if run.timed_out: + return "fail", "timeout", [] + result = read_result(run.stdout) + if result is None: + return ("fail", "crash", []) if run.returncode not in (0, None) else ("unsettled", "no_result", []) + status = result["status"] + if status == "ok" and run.returncode == 0 and result.get("checked", 0) > 0: + return "clean", None, [] + if status in ("mismatch", "out_of_bounds"): + return "fail", "wrong_output", [] + if status == "hip_error": + return "fail", "crash", [] + return "unsettled", "no_result", [] + + +class HipkRun: + """`hipk.run`, leg `ref.native`: the candidate built with its task's harness (hipcc, settings `hipcc`, leg setting + `target`), run on this worker's GPU (settings `gpus`: worker n runs on `gpus[(n - 1) % len(gpus)]`, by + HIP_VISIBLE_DEVICES) for at most `run_timeout_s`, and its outputs compared with the host reference. Clean only + when every output element matches and nothing was written past an array's end.""" + + version = "1" + legs = ("ref.native",) + resource = "own_gpu" + + def __init__(self, ctx): + from aorlta.reward.adapters.base import ProcessRunner + + self.name, self.ctx = ctx.name, ctx + s = ctx.settings + self.hipcc = s.get("hipcc") or "/opt/rocm/bin/hipcc" + gpus = list(s.get("gpus") or [0]) + seq = int(ctx.worker_id.rsplit("/", 1)[-1]) if ctx.worker_id.rsplit("/", 1)[-1].isdigit() else 1 + self.gpu = int(gpus[(seq - 1) % len(gpus)]) + self.tasks = {} + for pool in ctx.pools: + if (Path(pool) / "tasks.json").is_file(): + for task in load_tasks(pool): + self.tasks[task_sha256(task)] = task + self.runner = ProcessRunner() + + def run(self, job): + from aorlta.reward.adapters.base import make_result, run_dir, unsettled + + t0 = self.ctx.clock() + try: + task = self.tasks.get(job.task_sha256) + source = job.inputs.get("source") + if job.leg != "ref.native" or task is None or not source or not job.output_dir: + return unsettled(job, "adapter_error", t_start=t0, t_end=self.ctx.clock()) + d = run_dir(job) + (d / "harness.hip").write_text(harness_source(task, Path(source).read_text(encoding="utf-8")), + encoding="utf-8") + target = job.leg_settings.get("target", "gfx950") + budget = max(1.0, self.ctx.timeout_s) + t_build = time.monotonic() + build = self.runner.run([self.hipcc, f"--offload-arch={target}", "-O2", "-o", str(d / "harness"), + str(d / "harness.hip")], timeout_s=budget, cwd=str(d)) + if build.timed_out or build.killed or build.launch_error: + return unsettled(job, "harness_build_failed", t_start=t0, t_end=self.ctx.clock()) + (d / "build.log").write_text(build.stderr[-20000:], encoding="utf-8") + run = None + if build.returncode == 0: + left = budget - (time.monotonic() - t_build) + limit = min(float(job.leg_settings.get("run_timeout_s", 20.0)), left) + if limit <= 0: + return unsettled(job, "no_time_left", t_start=t0, t_end=self.ctx.clock()) + env = dict(os.environ, HIP_VISIBLE_DEVICES=str(self.gpu)) + run = self.runner.run([str(d / "harness")], timeout_s=limit, env=env, cwd=str(d)) + (d / "run.log").write_text((run.stdout + "\n" + run.stderr)[-20000:], encoding="utf-8") + state, reason, findings = verdict(build.returncode, build.stderr, run, task) + return make_result(job, state, reason, findings=findings, 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, "gpu": self.gpu, "tasks": len(self.tasks)} + + def close(self): + self.runner.close() + + +# --- the reward -------------------------------------------------------------------------------- + + +class HipkV1: + """`hipk.v1`: -1.0 when the parse gate fails, -0.9 for a compile error, -0.8 when the kernel's signature is not + the one asked for, -0.6 for a crash or a hang, -0.5 for a wrong output, +1.0 when every output matches; then + -0.05 x the share of the token budget the answer used. A leg that did not settle leaves the sample unscored: + an unscored sample is never counted clean. The parts are the service's (`base`, `broke_healthy` (always 0), + `length`), so the reward stays within its range.""" + + name = "hipk.v1" + version = "1" + legs = ("gate.parse", "gate.compile", "ref.native") + RUN = {"wrong_output": -0.5, "crash": -0.6, "timeout": -0.6, "bad_signature": -0.8} + + 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) -> tuple[str, float | str]: + """("score", base), ("unscored", reason) or ("pending", leg), reading the legs in order.""" + for name in self.legs: + if name not in legs: + return "pending", name + leg = legs[name] + if leg.state in ("unsettled", "skipped_by_plan", "finding", "not_run"): + return "unscored", f"{name}:{leg.reason or leg.state}" + if leg.state == "fail": + if name == "gate.parse": + return "score", -1.0 + if name == "gate.compile": + return "score", -0.9 if leg.reason == "compile_error" else -1.0 + if leg.reason in self.RUN: + return "score", self.RUN[leg.reason] + return "unscored", f"{name}:{leg.reason}" + return "score", 1.0 + + @staticmethod + def _finish(base, cand): + length = -0.05 * (cand.response_tokens / cand.max_new_tokens) + parts = {"base": float(base), "broke_healthy": 0.0, "length": length} + return parts["base"] + parts["broke_healthy"] + parts["length"], parts + + def compute(self, legs, cand): + kind, value = self._base(legs) + if kind != "score": + return None, {"reason": str(value)} + return 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,)) + if kind == "unscored": + return Unscored(value) + return 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 + + +# --- registration ------------------------------------------------------------------------------ + + +def register() -> None: + from aorlta import plugins + + plugins.register_task(plugins.Task("hipk", records)) + plugins.register_parse_adapter("hipk.parse", HipkParse) + plugins.register_feedback_template("hipk.feedback", HipkFeedback) + plugins.register_profile("hipk-gpu", HERE / "hipk-gpu.yaml") + plugins.register_check_adapter("hipk.run", HipkRun) + plugins.register_reward_fn("hipk.v1", HipkV1) + + +if __name__ != "__main__": + register() + + +# --- the self-check ---------------------------------------------------------------------------- + + +def selfcheck(argv=None) -> int: + """Each task's harness with its reference kernel (reference/.hip): built, and unless --build-only run on the + GPU, where every task must report `ok`. Prints one JSON line per task; exit 0 when all pass.""" + ap = argparse.ArgumentParser(prog="hip_kernels.py selfcheck") + ap.add_argument("--hipcc", default="/opt/rocm/bin/hipcc") + ap.add_argument("--arch", default="gfx950") + ap.add_argument("--gpu", type=int, default=0) + ap.add_argument("--pool", default=str(HERE)) + ap.add_argument("--build-only", action="store_true", help="compile the harnesses, run nothing (no GPU)") + args = ap.parse_args(argv) + import subprocess + + failed = 0 + work = Path(tempfile.mkdtemp(prefix="hipk-selfcheck-")) + try: + for task in load_tasks(args.pool): + ref = (HERE / "reference" / f"{task['id']}.hip").read_text(encoding="utf-8") + src, exe = work / f"{task['id']}.hip", work / task["id"] + src.write_text(harness_source(task, ref), encoding="utf-8") + build = subprocess.run([args.hipcc, f"--offload-arch={args.arch}", "-O2", "-o", str(exe), str(src)], + capture_output=True, text=True) + row = {"task": task["id"], "build_rc": build.returncode} + if build.returncode == 0 and not args.build_only: + run = subprocess.run([str(exe)], capture_output=True, text=True, timeout=60, + env=dict(os.environ, HIP_VISIBLE_DEVICES=str(args.gpu))) + row["result"] = read_result(run.stdout) + ok = run.returncode == 0 and row["result"] is not None and row["result"]["status"] == "ok" + else: + ok = build.returncode == 0 + if not ok: + row["stderr"] = build.stderr[-2000:] + row["pass"] = ok + failed += not ok + print(json.dumps(row), flush=True) + finally: + shutil.rmtree(work, ignore_errors=True) + return 1 if failed else 0 + + +if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "selfcheck": + sys.exit(selfcheck(sys.argv[2:])) + print(__doc__) + sys.exit(2) diff --git a/aorlta/examples/hip-kernels/hipk-gpu.yaml b/aorlta/examples/hip-kernels/hipk-gpu.yaml new file mode 100644 index 000000000..1a7d33f5b --- /dev/null +++ b/aorlta/examples/hip-kernels/hipk-gpu.yaml @@ -0,0 +1,26 @@ +# `hipk-gpu`: the example task's checks (examples/hip-kernels/hip_kernels.py). The stand-in parse and compile gates, +# then `ref.native`: the candidate built with its task's harness and run on one of the GPUs in `hipk.run.gpus`, its +# outputs compared with a host reference. Scored by hipk.v1. +name: hipk-gpu +reward_fn: hipk.v1 +clean_repeats: 1 +materialize: true +pools: + gates: {resource: local_cpu, adapter: standin.parse, timeout: gate, workers: 4} + compile: {resource: local_cpu, adapter: standin.hipcc, timeout: compile, workers: 8} + run: {resource: own_gpu, adapter: hipk.run, timeout: sim_run, workers: 4} +stages: + - {name: S0, pool: gates, legs: [gate.parse]} + - {name: S1, pool: compile, legs: [gate.compile]} + - {name: S2, pool: run, legs: [ref.native]} +requires: + gate.compile: [gate.parse] + ref.native: [gate.compile] +leg_settings: + gate.compile: {offload_arch: gfx950, flags: ["-c", "-O1"]} + ref.native: {target: gfx950, run_timeout_s: 20} +# `${AORLTA_ROCM_ROOT}` is that environment variable, or /opt/rocm. Worker n of the run pool uses +# gpus[(n - 1) % len(gpus)]: keep these off the engines' GPUs (rollout.engines). +adapters: + standin.hipcc: {hipcc: "${AORLTA_ROCM_ROOT}/bin/hipcc", readelf: "${AORLTA_ROCM_ROOT}/llvm/bin/llvm-readelf"} + hipk.run: {hipcc: "${AORLTA_ROCM_ROOT}/bin/hipcc", gpus: [4, 5, 6, 7]} diff --git a/aorlta/examples/hip-kernels/reference/block_sum.hip b/aorlta/examples/hip-kernels/reference/block_sum.hip new file mode 100644 index 000000000..4607e3bf7 --- /dev/null +++ b/aorlta/examples/hip-kernels/reference/block_sum.hip @@ -0,0 +1,16 @@ +#include + +// One block of 256 threads per 256 elements: each thread loads one element (0 past the end), then a tree reduction +// in shared memory halves the active threads each step, with a barrier between steps. +__global__ void block_sum(const float* x, float* out, int n) { + __shared__ float s[256]; + int t = threadIdx.x; + int i = blockIdx.x * 256 + t; + s[t] = i < n ? x[i] : 0.0f; + __syncthreads(); + for (int stride = 128; stride > 0; stride >>= 1) { + if (t < stride) s[t] += s[t + stride]; + __syncthreads(); + } + if (t == 0) out[blockIdx.x] = s[0]; +} diff --git a/aorlta/examples/hip-kernels/reference/clamp.hip b/aorlta/examples/hip-kernels/reference/clamp.hip new file mode 100644 index 000000000..545e9f613 --- /dev/null +++ b/aorlta/examples/hip-kernels/reference/clamp.hip @@ -0,0 +1,6 @@ +#include + +__global__ void clamp(const float* x, float* y, float lo, float hi, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) y[i] = fminf(fmaxf(x[i], lo), hi); +} diff --git a/aorlta/examples/hip-kernels/reference/relu.hip b/aorlta/examples/hip-kernels/reference/relu.hip new file mode 100644 index 000000000..f1bec822b --- /dev/null +++ b/aorlta/examples/hip-kernels/reference/relu.hip @@ -0,0 +1,6 @@ +#include + +__global__ void relu(const float* x, float* y, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) y[i] = x[i] > 0.0f ? x[i] : 0.0f; +} diff --git a/aorlta/examples/hip-kernels/reference/reverse.hip b/aorlta/examples/hip-kernels/reference/reverse.hip new file mode 100644 index 000000000..aaa2fbf43 --- /dev/null +++ b/aorlta/examples/hip-kernels/reference/reverse.hip @@ -0,0 +1,6 @@ +#include + +__global__ void reverse(const float* x, float* y, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) y[i] = x[n - 1 - i]; +} diff --git a/aorlta/examples/hip-kernels/reference/saxpy.hip b/aorlta/examples/hip-kernels/reference/saxpy.hip new file mode 100644 index 000000000..ba56c006f --- /dev/null +++ b/aorlta/examples/hip-kernels/reference/saxpy.hip @@ -0,0 +1,6 @@ +#include + +__global__ void saxpy(float alpha, const float* x, float* y, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) y[i] = alpha * x[i] + y[i]; +} diff --git a/aorlta/examples/hip-kernels/reference/scale_shift.hip b/aorlta/examples/hip-kernels/reference/scale_shift.hip new file mode 100644 index 000000000..2e3e3a353 --- /dev/null +++ b/aorlta/examples/hip-kernels/reference/scale_shift.hip @@ -0,0 +1,6 @@ +#include + +__global__ void scale_shift(const float* x, float* y, float scale, float shift, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) y[i] = x[i] * scale + shift; +} diff --git a/aorlta/examples/hip-kernels/reference/transpose.hip b/aorlta/examples/hip-kernels/reference/transpose.hip new file mode 100644 index 000000000..83c7a4738 --- /dev/null +++ b/aorlta/examples/hip-kernels/reference/transpose.hip @@ -0,0 +1,7 @@ +#include + +__global__ void transpose(const float* in, float* out, int rows, int cols) { + int c = blockIdx.x * blockDim.x + threadIdx.x; + int r = blockIdx.y * blockDim.y + threadIdx.y; + if (r < rows && c < cols) out[c * rows + r] = in[r * cols + c]; +} diff --git a/aorlta/examples/hip-kernels/reference/vector_add.hip b/aorlta/examples/hip-kernels/reference/vector_add.hip new file mode 100644 index 000000000..967b8a668 --- /dev/null +++ b/aorlta/examples/hip-kernels/reference/vector_add.hip @@ -0,0 +1,6 @@ +#include + +__global__ void vector_add(const float* a, const float* b, float* c, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) c[i] = a[i] + b[i]; +} diff --git a/aorlta/examples/hip-kernels/tasks.json b/aorlta/examples/hip-kernels/tasks.json new file mode 100644 index 000000000..7174ea1da --- /dev/null +++ b/aorlta/examples/hip-kernels/tasks.json @@ -0,0 +1,99 @@ +[ + { + "id": "vector_add", + "kernel": "vector_add", + "signature": "__global__ void vector_add(const float* a, const float* b, float* c, int n)", + "description": "c[i] = a[i] + b[i] for every 0 <= i < n", + "launch": "with blockDim.x = 256 and gridDim.x = (n + 255) / 256, one thread per element", + "scalars": [["int", "n", "100003"]], + "arrays": [{"name": "a", "size": "n", "role": "in"}, {"name": "b", "size": "n", "role": "in"}, + {"name": "c", "size": "n", "role": "out"}], + "grid": "dim3((n + 255) / 256)", "block": "dim3(256)", "args": ["a", "b", "c", "n"], + "reference": "for (int i = 0; i < n; ++i) c[i] = a[i] + b[i];", + "rtol": 0.0, "atol": 0.0 + }, + { + "id": "saxpy", + "kernel": "saxpy", + "signature": "__global__ void saxpy(float alpha, const float* x, float* y, int n)", + "description": "y[i] = alpha * x[i] + y[i] for every 0 <= i < n, in place", + "launch": "with blockDim.x = 256 and gridDim.x = (n + 255) / 256, one thread per element", + "scalars": [["int", "n", "100003"], ["float", "alpha", "2.5f"]], + "arrays": [{"name": "x", "size": "n", "role": "in"}, {"name": "y", "size": "n", "role": "inout"}], + "grid": "dim3((n + 255) / 256)", "block": "dim3(256)", "args": ["alpha", "x", "y", "n"], + "reference": "for (int i = 0; i < n; ++i) y[i] = alpha * x[i] + y[i];", + "rtol": 1e-05, "atol": 1e-06 + }, + { + "id": "relu", + "kernel": "relu", + "signature": "__global__ void relu(const float* x, float* y, int n)", + "description": "y[i] = max(x[i], 0) for every 0 <= i < n", + "launch": "with blockDim.x = 256 and gridDim.x = (n + 255) / 256, one thread per element", + "scalars": [["int", "n", "100003"]], + "arrays": [{"name": "x", "size": "n", "role": "in"}, {"name": "y", "size": "n", "role": "out"}], + "grid": "dim3((n + 255) / 256)", "block": "dim3(256)", "args": ["x", "y", "n"], + "reference": "for (int i = 0; i < n; ++i) y[i] = x[i] > 0.0f ? x[i] : 0.0f;", + "rtol": 0.0, "atol": 0.0 + }, + { + "id": "scale_shift", + "kernel": "scale_shift", + "signature": "__global__ void scale_shift(const float* x, float* y, float scale, float shift, int n)", + "description": "y[i] = x[i] * scale + shift for every 0 <= i < n", + "launch": "with blockDim.x = 256 and gridDim.x = (n + 255) / 256, one thread per element", + "scalars": [["int", "n", "100003"], ["float", "scale", "-1.75f"], ["float", "shift", "0.25f"]], + "arrays": [{"name": "x", "size": "n", "role": "in"}, {"name": "y", "size": "n", "role": "out"}], + "grid": "dim3((n + 255) / 256)", "block": "dim3(256)", "args": ["x", "y", "scale", "shift", "n"], + "reference": "for (int i = 0; i < n; ++i) y[i] = x[i] * scale + shift;", + "rtol": 1e-05, "atol": 1e-06 + }, + { + "id": "clamp", + "kernel": "clamp", + "signature": "__global__ void clamp(const float* x, float* y, float lo, float hi, int n)", + "description": "y[i] = min(max(x[i], lo), hi) for every 0 <= i < n", + "launch": "with blockDim.x = 256 and gridDim.x = (n + 255) / 256, one thread per element", + "scalars": [["int", "n", "100003"], ["float", "lo", "-0.5f"], ["float", "hi", "0.75f"]], + "arrays": [{"name": "x", "size": "n", "role": "in"}, {"name": "y", "size": "n", "role": "out"}], + "grid": "dim3((n + 255) / 256)", "block": "dim3(256)", "args": ["x", "y", "lo", "hi", "n"], + "reference": "for (int i = 0; i < n; ++i) y[i] = x[i] < lo ? lo : (x[i] > hi ? hi : x[i]);", + "rtol": 0.0, "atol": 0.0 + }, + { + "id": "reverse", + "kernel": "reverse", + "signature": "__global__ void reverse(const float* x, float* y, int n)", + "description": "y[i] = x[n - 1 - i] for every 0 <= i < n", + "launch": "with blockDim.x = 256 and gridDim.x = (n + 255) / 256, one thread per element", + "scalars": [["int", "n", "100003"]], + "arrays": [{"name": "x", "size": "n", "role": "in"}, {"name": "y", "size": "n", "role": "out"}], + "grid": "dim3((n + 255) / 256)", "block": "dim3(256)", "args": ["x", "y", "n"], + "reference": "for (int i = 0; i < n; ++i) y[i] = x[n - 1 - i];", + "rtol": 0.0, "atol": 0.0 + }, + { + "id": "block_sum", + "kernel": "block_sum", + "signature": "__global__ void block_sum(const float* x, float* out, int n)", + "description": "out[b] = the sum of x[i] over b * 256 <= i < min((b + 1) * 256, n), for every block b", + "launch": "with blockDim.x = 256 and gridDim.x = (n + 255) / 256, one block per 256 elements", + "scalars": [["int", "n", "100003"], ["int", "blocks", "(n + 255) / 256"]], + "arrays": [{"name": "x", "size": "n", "role": "in"}, {"name": "out", "size": "blocks", "role": "out"}], + "grid": "dim3(blocks)", "block": "dim3(256)", "args": ["x", "out", "n"], + "reference": "for (int b = 0; b < blocks; ++b) { double s = 0.0; for (int i = b * 256; i < (b + 1) * 256 && i < n; ++i) s += x[i]; out[b] = (float)s; }", + "rtol": 0.0001, "atol": 0.002 + }, + { + "id": "transpose", + "kernel": "transpose", + "signature": "__global__ void transpose(const float* in, float* out, int rows, int cols)", + "description": "out[c * rows + r] = in[r * cols + c] for every 0 <= r < rows and 0 <= c < cols: `in` is a rows x cols row-major matrix and `out` its cols x rows transpose", + "launch": "with blockDim = (16, 16) and gridDim = ((cols + 15) / 16, (rows + 15) / 16), where threadIdx.x and blockIdx.x run along the columns of `in`", + "scalars": [["int", "rows", "300"], ["int", "cols", "517"]], + "arrays": [{"name": "in", "size": "rows * cols", "role": "in"}, {"name": "out", "size": "rows * cols", "role": "out"}], + "grid": "dim3((cols + 15) / 16, (rows + 15) / 16)", "block": "dim3(16, 16)", "args": ["in", "out", "rows", "cols"], + "reference": "for (int r = 0; r < rows; ++r) for (int c = 0; c < cols; ++c) out[c * rows + r] = in[r * cols + c];", + "rtol": 0.0, "atol": 0.0 + } +] diff --git a/aorlta/tests/examples/test_hip_kernels.py b/aorlta/tests/examples/test_hip_kernels.py new file mode 100644 index 000000000..4624f47f7 --- /dev/null +++ b/aorlta/tests/examples/test_hip_kernels.py @@ -0,0 +1,367 @@ +"""The worked example of docs/TUTORIAL.md (examples/hip-kernels), on CPU: its configs load with the plugin, its +prompts and probes build, its parse asks for a repair turn, its check maps each harness outcome to a leg state +without a GPU (a fake hipcc stands in for the compiler and the GPU run), its reward keeps the service's invariants, +and every task's harness compiles with its reference kernel where hipcc is installed.""" + +from __future__ import annotations + +import itertools +import json +import os +import stat +import subprocess +import sys +from pathlib import Path + +import pytest + +from aorlta import plugins +from aorlta.config import config_from_dict, load_config, parse_yaml + +ROOT = Path(__file__).resolve().parents[2] +EXAMPLE = ROOT / "examples" / "hip-kernels" +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, for the wired test +from test_plugins import CharChat, registry # noqa: E402,F401 - the probe tokenizer and the registry's restore + +tiny_base = E.tiny_base + +HIPCC = os.environ.get("HIPCC", "/opt/rocm/bin/hipcc") + + +@pytest.fixture +def hipk(registry): # noqa: F811 + """The example plugin, loaded as a run config loads it; its module.""" + plugins.load_plugins(["examples/hip-kernels/hip_kernels.py"]) + return sys.modules[plugins.identity()["examples/hip-kernels/hip_kernels.py"]["module"]] + + +def site_config(tmp_path, name: str): + site = tmp_path / "site.yaml" + site.write_text(json.dumps({"run": {"root": str(tmp_path / "runs")}, "model": {"base_dir": str(tmp_path / "base")}, + "data": {"prompt_pools": [str(EXAMPLE)]}})) + return load_config(ROOT / "configs" / f"{name}.yaml", site=site) + + +@pytest.mark.parametrize("name", ["hipk-score", "hipk-train"]) +def test_the_example_configs_load_with_the_plugin(tmp_path, registry, name): # noqa: F811 + cfg = site_config(tmp_path, name) + assert cfg.plugins == ("examples/hip-kernels/hip_kernels.py",) and cfg.data.prompt_pools[0].task == "hipk" + assert (cfg.reward.profile, cfg.reward.reward_fn, cfg.rollout.parse_adapter) == ("hipk-gpu", "hipk.v1", + "hipk.parse") + from aorlta.orchestrator import wiring + + wiring.check_supported(cfg) + + +def test_the_tasks_give_prompts_and_probes(tmp_path, hipk): + from aorlta.orchestrator import wiring + + cfg = site_config(tmp_path, "hipk-score") + tok = CharChat() + recs = wiring.TaskPromptPool(plugins.task("hipk"), tok, "off", EXAMPLE).records() + tasks = hipk.load_tasks(EXAMPLE) + assert [r.prompt_id for r in recs] == [f"hipk:{t['id']}" for t in tasks] and len(recs) == 8 + assert {r.task.task_sha256 for r in recs} == {hipk.task_sha256(t) for t in tasks} + for r, t in zip(recs, tasks): + assert t["signature"] in r.messages[1]["content"] + assert (EXAMPLE / "reference" / f"{t['id']}.hip").is_file() + probes = wiring.task_probes(cfg, tok, plugins.task("hipk"), EXAMPLE, recs) + assert len(probes.ids) == cfg.sync.probes + + +def prompt(hipk, kernel="saxpy"): + from aorlta.orchestrator import wiring + + return next(r for r in wiring.TaskPromptPool(plugins.task("hipk"), CharChat(), "off", EXAMPLE).records() + if r.prompt_id == f"hipk:{kernel}") + + +def test_the_parse_gives_a_repair_turn_for_no_block_or_no_kernel(hipk): + from aorlta.rollout.chain import make_feedback_template, make_parse_adapter + + parse, fb = make_parse_adapter("hipk.parse"), make_feedback_template("hipk.feedback") + p = prompt(hipk) + no_block = parse.check("I would write a loop.", "off", p) + assert no_block.kind == "rejected" and no_block.rejection[1] == "no_code_block" + other = parse.check("```cpp\n__global__ void axpy(float a, const float* x, float* y, int n) {}\n```", "off", p) + assert other.kind == "rejected" and other.rejection[1] == "no_kernel" and "`saxpy`" in other.rejection[0] + good = parse.check("x\n```cpp\n#include \n__global__ void saxpy(float a, const float* x, " + "float* y, int n) { }\n```\n", "off", p) + assert good.kind == "edits" and "__global__ void saxpy(" in good.code + turn = fb.build_apply(*other.rejection) + assert turn.text == other.rejection[0] and len(turn.text) <= 2000 + + +def proc(rc=0, stdout="", stderr="", timed_out=False, killed=False, launch_error=None): + from aorlta.reward.adapters.base import ProcResult + + return ProcResult(rc, stdout, stderr, 0.1, timed_out, killed, launch_error) + + +def result_line(status, **kw): + return "AORLTA_RESULT " + json.dumps({"status": status, "checked": 10, "mismatches": 0, **kw}) + + +@pytest.mark.parametrize("build, run, want", [ + ((0, ""), proc(stdout=result_line("ok")), ("clean", None)), + ((0, ""), proc(stdout=result_line("mismatch", mismatches=3)), ("fail", "wrong_output")), + ((0, ""), proc(stdout=result_line("out_of_bounds")), ("fail", "wrong_output")), + ((0, ""), proc(stdout=result_line("hip_error")), ("fail", "crash")), + ((0, ""), proc(rc=-11), ("fail", "crash")), + ((0, ""), proc(timed_out=True, rc=-9), ("fail", "timeout")), + ((0, ""), proc(stdout="nothing"), ("unsettled", "no_result")), + ((0, ""), proc(rc=1, stdout=result_line("ok")), ("unsettled", "no_result")), + ((0, ""), proc(stdout=result_line("ok", checked=0)), ("unsettled", "no_result")), + ((0, ""), proc(killed=True), ("unsettled", "run_not_started")), + ((0, ""), proc(launch_error="OSError"), ("unsettled", "run_not_started")), + ((1, "error: no matching function for call to 'saxpy'"), None, ("fail", "bad_signature")), + ((1, "error: redefinition of 'main'"), None, ("fail", "bad_signature")), + ((1, "ld.lld: error: unable to find library"), None, ("unsettled", "harness_build_failed")), +]) +def test_each_harness_outcome_has_one_leg_state(hipk, build, run, want): + task = next(t for t in hipk.load_tasks(EXAMPLE) if t["id"] == "saxpy") + state, reason, _ = hipk.verdict(build[0], build[1], run, task) + assert (state, reason) == want + + +FAKE_HIPCC = """#!/bin/bash +# a stand-in for hipcc: the 'compiled' program prints what FAKE_RUN says (ok, mismatch, crash, hang) +out=""; while [ $# -gt 0 ]; do [ "$1" = -o ] && { out=$2; shift; }; shift; done +[ "${FAKE_BUILD:-ok}" = sig ] && { echo "error: no matching function for call to 'saxpy'" >&2; exit 1; } +printf '#!/bin/bash\\necho "gpu $HIP_VISIBLE_DEVICES" >&2\\n' > "$out" +case "${FAKE_RUN:-ok}" in + ok) echo 'echo "AORLTA_RESULT {\\"status\\": \\"ok\\", \\"checked\\": 5}"' >> "$out" ;; + mismatch) echo 'echo "AORLTA_RESULT {\\"status\\": \\"mismatch\\", \\"checked\\": 5, \\"mismatches\\": 2}"' >> "$out" ;; + crash) echo 'kill -SEGV $$' >> "$out" ;; + hang) echo 'sleep 30' >> "$out" ;; +esac +chmod +x "$out" +""" + + +@pytest.fixture +def fake_hipcc(tmp_path) -> Path: + p = tmp_path / "hipcc" + p.write_text(FAKE_HIPCC) + p.chmod(p.stat().st_mode | stat.S_IXUSR) + return p + + +def run_job(hipk, tmp_path, fake_hipcc, *, worker="run/2", task_id="saxpy", timeout=5.0, run_timeout_s=2.0): + from aorlta.schemas import LegJob + + task = next(t for t in hipk.load_tasks(EXAMPLE) if t["id"] == task_id) + src = tmp_path / "source.hip" + src.write_text("__global__ void saxpy(float a, const float* x, float* y, int n) {}\n") + out = tmp_path / "out" + out.mkdir(exist_ok=True) + ctx = plugins.CheckContext(name="hipk.run", worker_id=worker, settings={"hipcc": str(fake_hipcc), "gpus": [4, 6]}, + clock=lambda: 0.0, timeout_s=timeout, pools=(str(EXAMPLE),)) + adapter = plugins.CHECK_ADAPTERS["hipk.run"](ctx) + job = LegJob(job_id="j", leg="ref.native", candidate_id="c", control_id=None, source_sha256="s" * 64, + task_sha256=hipk.task_sha256(task), leg_settings={"target": "gfx950", "run_timeout_s": run_timeout_s}, + inputs={"source": str(src)}, output_dir=str(out)) + return adapter, adapter.run(job) + + +@pytest.mark.parametrize("build, run, want", [ + ("ok", "ok", ("clean", None)), ("ok", "mismatch", ("fail", "wrong_output")), ("ok", "crash", ("fail", "crash")), + ("ok", "hang", ("fail", "timeout")), ("sig", "ok", ("fail", "bad_signature")), +]) +def test_the_check_runs_the_harness_on_its_workers_gpu(hipk, tmp_path, fake_hipcc, monkeypatch, build, run, want): + monkeypatch.setenv("FAKE_BUILD", build) + monkeypatch.setenv("FAKE_RUN", run) + adapter, res = run_job(hipk, tmp_path, fake_hipcc) + assert (res.state, res.reason) == want and res.leg == "ref.native" + assert adapter.gpu == 6 # worker 2 of gpus [4, 6] + if build == "ok" and run == "ok": + (log,) = tmp_path.glob("out/*/run.log") + assert "gpu 6" in log.read_text() + + +def test_a_task_the_check_does_not_know_is_unsettled(hipk, tmp_path, fake_hipcc): + from aorlta.schemas import LegJob + + ctx = plugins.CheckContext(name="hipk.run", worker_id="run/1", settings={"hipcc": str(fake_hipcc)}, + clock=lambda: 0.0, timeout_s=5.0, pools=(str(tmp_path),)) + adapter = plugins.CHECK_ADAPTERS["hipk.run"](ctx) + job = LegJob(job_id="j", leg="ref.native", candidate_id="c", control_id=None, source_sha256="s" * 64, + task_sha256="0" * 64, inputs={"source": str(tmp_path / "x.hip")}, output_dir=str(tmp_path)) + res = adapter.run(job) + assert (res.state, res.reason) == ("unsettled", "adapter_error") + + +def leg(name, state, reason=None): + from aorlta.schemas import LegResult + + return LegResult(leg=name, state=state, reason=reason, findings=[], repeats=1 if state != "unsettled" else 0, + attempts=1, worker_id=None, bracket=None, cache_hit=False, t_queued=0.0, t_start=0.0, t_end=0.0) + + +def cand(tokens=100): + from aorlta.schemas import Candidate, TaskRef + + return Candidate(candidate_id="c", run_id="r", iteration=1, prompt_id="hipk:saxpy", k=0, policy_version="v", + response_text="", response_tokens=tokens, max_new_tokens=2048, + task=TaskRef("hipk:saxpy", None, None, None, "t" * 64), check_profile="hipk-gpu", + rollout_compile=None) + + +OUTCOMES = { + "gate.parse": [("clean", None), ("fail", "parse"), ("unsettled", "timeout")], + "gate.compile": [("clean", None), ("fail", "compile_error"), ("unsettled", "adapter_error"), ("not_run", None)], + "ref.native": [("clean", None), ("fail", "wrong_output"), ("fail", "crash"), ("fail", "timeout"), + ("fail", "bad_signature"), ("fail", "other"), ("unsettled", "no_result"), ("not_run", None), + ("skipped_by_plan", None)], +} + + +def test_the_reward_ladder(hipk): + fn = plugins.REWARD_FNS["hipk.v1"](None) + c = cand(tokens=0) + + def score(**states): + legs = {n: leg(n, *states.get(n.replace(".", "_"), ("clean", None))) for n in fn.legs} + return fn.compute(legs, c)[0] + + assert score() == 1.0 + assert score(gate_parse=("fail", "parse")) == -1.0 + assert score(gate_compile=("fail", "compile_error")) == -0.9 + assert score(ref_native=("fail", "bad_signature")) == -0.8 + assert score(ref_native=("fail", "crash")) == score(ref_native=("fail", "timeout")) == -0.6 + assert score(ref_native=("fail", "wrong_output")) == -0.5 + assert score(ref_native=("unsettled", "no_result")) is None and score(ref_native=("fail", "other")) is None + + +def test_the_reward_keeps_the_services_invariants_and_decide_agrees_with_compute(hipk): + """Every combination of leg outcomes, and every prefix of it as a partial record: decide never contradicts + compute, and a score is a float in the service's range whose parts are base, broke_healthy and length and add + up to it (reward/service.py `_invariant_problems`).""" + from aorlta.reward.rewardfn.base import Decided, NeedMore, Unscored + from aorlta.reward.service import REWARD_MAX, REWARD_MIN + + fn = plugins.REWARD_FNS["hipk.v1"](None) + for tokens in (0, 1024, 2048): + c = cand(tokens) + for combo in itertools.product(*(OUTCOMES[n] for n in fn.legs)): + legs = {n: leg(n, *o) for n, o in zip(fn.legs, combo)} + reward, parts = fn.compute(legs, c) + if reward is not None: + assert type(reward) is float and REWARD_MIN <= reward <= REWARD_MAX + assert set(parts) == {"base", "broke_healthy", "length"} and all(type(v) is float for v in parts.values()) + assert parts["base"] + parts["broke_healthy"] + parts["length"] == reward + assert -0.05 <= parts["length"] <= 0.0 and parts["broke_healthy"] == 0.0 + assert fn.claims_clean(reward, parts) == all(legs[n].state == "clean" for n in fn.legs) + else: + assert isinstance(parts.get("reason"), str) and not fn.claims_clean(reward, parts) + for k in range(len(fn.legs) + 1): + d = fn.decide({n: legs[n] for n in fn.legs[:k]}, c) + if isinstance(d, Decided): + assert (d.reward, d.parts) == (reward, parts) + elif isinstance(d, Unscored): + assert reward is None + else: + assert isinstance(d, NeedMore) and d.legs == (fn.legs[k],) + + +@pytest.mark.skipif(not Path(HIPCC).exists(), reason=f"no hipcc at {HIPCC} (HIPCC)") +def test_every_harness_compiles_with_its_reference_kernel(tmp_path): + """The self-check's build step, CPU only: nothing is run.""" + proc_ = subprocess.run([sys.executable, str(EXAMPLE / "hip_kernels.py"), "selfcheck", "--hipcc", HIPCC, + "--build-only"], capture_output=True, text=True, env=dict(os.environ, TMPDIR=str(tmp_path))) + rows = [json.loads(x) for x in proc_.stdout.splitlines()] + assert proc_.returncode == 0 and len(rows) == 8 and all(r["pass"] for r in rows), proc_.stdout + proc_.stderr + + +def test_the_profile_names_only_what_the_plugin_and_the_package_provide(hipk): + from aorlta.reward.profiles import load_profile + + profile = load_profile("hipk-gpu") + assert profile.reward_fn == "hipk.v1" and profile.legs == ("gate.parse", "gate.compile", "ref.native") + assert {p.adapter for p in profile.pools.values()} == {"standin.parse", "standin.hipcc", "hipk.run"} + raw = parse_yaml((EXAMPLE / "hipk-gpu.yaml").read_bytes()) + assert raw["adapters"]["hipk.run"]["gpus"] == [4, 5, 6, 7] + + +def test_the_score_config_keeps_the_checks_off_the_engines(tmp_path, registry): # noqa: F811 + cfg = site_config(tmp_path, "hipk-score") + from aorlta.reward.profiles import load_profile + + gpus = set(load_profile(cfg.reward.profile).adapters["hipk.run"]["gpus"]) + assert not gpus & set(cfg.rollout.engines) + + +def test_the_tutorial_names_files_that_exist(): + text = (ROOT / "docs" / "TUTORIAL.md").read_text() + for rel in ("examples/hip-kernels/hip_kernels.py", "examples/hip-kernels/tasks.json", + "examples/hip-kernels/hipk-gpu.yaml", "configs/hipk-score.yaml", "configs/hipk-train.yaml"): + assert rel in text and (ROOT / rel).exists(), rel + + +def test_config_from_dict_also_takes_the_plugin(tmp_path, registry): # noqa: F811 + data = parse_yaml((ROOT / "configs" / "hipk-score.yaml").read_bytes()) + data["run"]["root"] = str(tmp_path) + data["data"]["prompt_pools"][0]["path"] = str(EXAMPLE) + assert config_from_dict(data).reward.profile == "hipk-gpu" + + +class _CpuLeg: + """A CPU stand-in for one of the example's checks that need hipcc or a GPU (the wired test below): the compile + gate is clean, and `ref.native` passes half the sources by hash, so the rewards in a group differ.""" + + version, resource = "1", "local_cpu" + + def __init__(self, name, leg): + self.name, self.legs = name, (leg,) + + def run(self, job): + from aorlta.reward.adapters.base import make_result + + if job.leg == "gate.compile" or job.source_sha256[-1] in "02468ace": + return make_result(job, "clean", t_start=0.0, t_end=0.0) + return make_result(job, "fail", "wrong_output", t_start=0.0, t_end=0.0) + + def controls(self): + return [] + + def health(self): + return {"responsive": True} + + +@pytest.mark.heavy +def test_the_example_trains_on_cpu_through_the_real_wiring(tmp_path, hipk, tiny_base): + """The example's prompts, parse, feedback, profile and reward in the real loop, with test_wiring_e2e's tiny model + and engines. The fake engine never writes the kernel asked for, so every chain gets hipk.parse's repair turn.""" + from aorlta.orchestrator import wiring + from aorlta.reward.adapters.standin_parse import StandinParse + from aorlta.rollout.engine_proc import InProcessEngine + + data = parse_yaml((ROOT / "configs" / "hipk-train.yaml").read_bytes()) + data["run"].update(root=str(tmp_path / "runs"), run_id="hipk-cpu") + data["model"]["base_dir"] = str(tiny_base) + data["data"]["prompt_pools"][0]["path"] = str(EXAMPLE) + data["rollout"].update(engines=[0, 1], K=4, P=4, max_new_tokens=512, request_timeout_s=300, + compile_adapter="fake.compile") + 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(tiny_base, E.FAKE, i) for i in range(2)] + factories = {"standin.parse": lambda w: StandinParse(), "standin.hipcc": lambda w: _CpuLeg("standin.hipcc", + "gate.compile"), + "hipk.run": lambda w: _CpuLeg("hipk.run", "ref.native")} + 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, reward_factories=factories) + orch = wiring.build(cfg, parts=parts) + 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.episodes(cfg, 1) + assert {e["prompt_id"].split(":")[0] for e in eps} == {"hipk"} + assert all(len(e["turns"]) == 3 and e["end_reason"] == "no_code_block" for e in eps if e["end_reason"] != "engine_error") + bases = {round(e["reward"] + 0.05, 1) for e in eps if e["status"] == "scored"} # base, give or take the length term + assert bases <= {1.0, -0.5, -1.0} and {1.0, -0.5} <= bases # correct, wrong output, and no code block at all From 198466f1927ff60175166108fcfbdf6a0a0f743b Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Thu, 12 Nov 2026 18:05:57 +0000 Subject: [PATCH 04/19] feat(aorlta): rollout.identical_groups_stop, the share of all-identical groups that stops a run The rollout stops an iteration when more than 5% of its groups have all K answers identical (sampler suspected). With few prompts and a precise task that is one group: the hip-kernels example's 3-iteration run on MI350X stopped in iteration 1 with 1 of 8 groups identical, on engines whose sampler had just passed the start-up self-test. The share is now a config key, default unchanged. --- aorlta/docs/ADAPTING.md | 4 +++- aorlta/src/aorlta/config.py | 6 ++++++ aorlta/src/aorlta/rollout/chain.py | 6 ++++-- aorlta/tests/rollout/test_chain.py | 22 ++++++++++++++++++++++ 4 files changed, 35 insertions(+), 3 deletions(-) diff --git a/aorlta/docs/ADAPTING.md b/aorlta/docs/ADAPTING.md index a6f6c00e3..9dad69b3b 100644 --- a/aorlta/docs/ADAPTING.md +++ b/aorlta/docs/ADAPTING.md @@ -41,7 +41,9 @@ the trainer's warm-up read `sync.probes` chats: the task's `probe_messages(pool_ messages, which must then share one system message. Each engine's start-up self-test samples 4 chats twice with different seeds and needs at least 3 of the 4 pairs to differ, so for a plugin task it uses fixed open-ended chats (`SELF_TEST_CHATS` in `orchestrator/wiring.py`), not the task's own prompts, whose answers may be near-unique; -`Task(selftest_messages=...)` gives others. The run plan records each plugin's file and its sha256. +`Task(selftest_messages=...)` gives others. For the same reason, a task with near-unique answers and few prompts may +need `rollout.identical_groups_stop` raised: by default an iteration stops when more than 5% of its groups have all +`K` answers identical, and with 8 prompts one such group is 12.5%. The run plan records each plugin's file and its sha256. Limits of a plugin: its data is handled as the stand-in source's (public; text may appear in logs and tracebacks); its reward stays within the stand-in range and parts (below); the episode is the compile-repair chain (a parse may diff --git a/aorlta/src/aorlta/config.py b/aorlta/src/aorlta/config.py index 95ce241a8..eca6b6a9a 100644 --- a/aorlta/src/aorlta/config.py +++ b/aorlta/src/aorlta/config.py @@ -435,6 +435,10 @@ class RolloutSection(_Section): # how a repair turn's ids are built: "append" (Appendix C) or "rerender" (the whole conversation through the chat # template again); null: "append" repair_turn_ids: Literal["append", "rerender"] | None = None + # stop when more than this share of an iteration's groups have all K answers identical (sampler suspected). A task + # with near-unique answers and few prompts may raise it: one such group in 8 is 12.5%. The engines' start-up + # self-test checks the sampler on its own. + identical_groups_stop: float = 0.05 def _validate(self) -> None: _check(len(self.engines) >= 1, "engines: at least one engine GPU is needed") @@ -449,6 +453,8 @@ def _validate(self) -> None: _check(self.max_new_tokens >= 1, f"max_new_tokens: must be >= 1, got {self.max_new_tokens}") _check(self.max_feedback_chars >= 1, f"max_feedback_chars: must be >= 1, got {self.max_feedback_chars}") _check(self.request_timeout_s > 0, f"request_timeout_s: must be > 0, got {self.request_timeout_s}") + _check(0 <= self.identical_groups_stop <= 1, + f"identical_groups_stop: must be in [0, 1], got {self.identical_groups_stop}") _check( 0 < self.gpu_memory_utilization <= 1, f"gpu_memory_utilization: must be in (0, 1], got {self.gpu_memory_utilization}", diff --git a/aorlta/src/aorlta/rollout/chain.py b/aorlta/src/aorlta/rollout/chain.py index 1d2341c15..4a3a28e83 100644 --- a/aorlta/src/aorlta/rollout/chain.py +++ b/aorlta/src/aorlta/rollout/chain.py @@ -65,7 +65,7 @@ SCORED_ENDS = frozenset({"compiled", "no_change", "no_code_block", "turns_exhausted", "context_limit", "logprob_error"}) TRAINED_ENDS = frozenset({"compiled", "no_change", "no_code_block", "turns_exhausted", "context_limit"}) -IDENTICAL_GROUPS_STOP = 0.05 # above this share of groups with all K answers identical, stop (sampler suspected) +IDENTICAL_GROUPS_STOP = 0.05 # the default of rollout.identical_groups_stop REPAIR_TURN_IDS = ("append", "rerender") @@ -203,6 +203,7 @@ class RolloutSettings: parse_adapter: str = "standin.parse" feedback_template: str = "standin" repair_turn_ids: str = "append" # REPAIR_TURN_IDS + identical_groups_stop: float = IDENTICAL_GROUPS_STOP # above this share of all-identical groups, stop @classmethod def from_config(cls, cfg: RunConfig, max_model_len: int) -> RolloutSettings: @@ -214,6 +215,7 @@ def from_config(cls, cfg: RunConfig, max_model_len: int) -> RolloutSettings: temperature=r.temperature, top_p=r.top_p, top_k=r.top_k, max_new_tokens=r.max_new_tokens, max_feedback_chars=r.max_feedback_chars, max_model_len=max_model_len, P=r.P, parse_adapter=r.parse_adapter, feedback_template=r.feedback_template, repair_turn_ids=r.repair_turn_ids, + identical_groups_stop=r.identical_groups_stop, ) def sampling_params(self, seed: int) -> dict: @@ -517,7 +519,7 @@ def chain_task(g: int, k: int) -> None: chains = [o.record for o in outcomes] metrics = self._metrics(iteration, prompts, outcomes, assignment, estimates, seeds, t_start, t_end) stop: list[str] = [] - if metrics["groups_all_identical_frac"] > IDENTICAL_GROUPS_STOP: + if metrics["groups_all_identical_frac"] > s.identical_groups_stop: stop.append(f"{metrics['groups_all_identical']} of {metrics['groups_compared']} groups have all " f"{s.K} answers identical at T {s.temperature} (sampler suspected)") if self.pool.stop_reason: diff --git a/aorlta/tests/rollout/test_chain.py b/aorlta/tests/rollout/test_chain.py index dd767b4da..d60c73683 100644 --- a/aorlta/tests/rollout/test_chain.py +++ b/aorlta/tests/rollout/test_chain.py @@ -318,6 +318,28 @@ def test_identical_groups_above_5_percent_stop_the_run(rt): assert healthy.metrics["distinct_turn0_mean"] == 4 and healthy.stop_reasons == [] +def test_rollout_identical_groups_stop_sets_the_share_that_stops_the_run(rt): + """`rollout.identical_groups_stop`: the same all-identical groups stop the run at the default and at 0.5, and do + not at 1.0; the config carries the value into the rollout's settings.""" + from aorlta.config import config_from_dict + from aorlta.rollout.chain import IDENTICAL_GROUPS_STOP, RolloutSettings + + for limit, stops in ((IDENTICAL_GROUPS_STOP, True), (0.5, True), (1.0, False)): + fake = rt.fake_engines(1, FakeEngineConfig(seed=1, **SHORT))[0] + pool = rt.make_pool([SeedBlind(0, fake)]) + result = rt.make_manager(pool, K=4, identical_groups_stop=limit).run( + rt.make_prompts(rt.SpecialByteTokenizer(), 3), 1, "v") + assert result.metrics["groups_all_identical"] == 3 + assert any("sampler suspected" in r for r in result.stop_reasons) == stops, limit + cfg = config_from_dict({"run": {"run_id": "r", "source": "standin", "root": "/x"}, + "model": {"base_dir": "/b", "source": "standin"}, + "data": {"prompt_pools": [{"path": "/p", "source": "standin"}]}, + "rollout": {"identical_groups_stop": 0.3}, + "reward": {"profile": "fake", "reward_fn": "standin.v11"}, + "orchestrator": {"iterations": 1, "run_seed": 1}}) + assert RolloutSettings.from_config(cfg, 4096).identical_groups_stop == 0.3 + + def test_prompts_go_to_engines_by_lpt_and_groups_stay_on_one_engine(rt): result, reqs, got, prompts, mgr = _run(rt, FakeEngineConfig(seed=1, compile_error_rate=0.5, **SHORT), P=6, K=3, engines=3) From af03c11be448bd5e236f762af38247bb4784acc7 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Thu, 12 Nov 2026 18:06:59 +0000 Subject: [PATCH 05/19] feat(aorlta): the example's configs allow identical groups, and the tutorial says why --- aorlta/configs/hipk-score.yaml | 3 +++ aorlta/configs/hipk-train.yaml | 3 +++ aorlta/docs/TUTORIAL.md | 4 ++++ 3 files changed, 10 insertions(+) diff --git a/aorlta/configs/hipk-score.yaml b/aorlta/configs/hipk-score.yaml index 0db985b5b..89c22f014 100644 --- a/aorlta/configs/hipk-score.yaml +++ b/aorlta/configs/hipk-score.yaml @@ -26,6 +26,9 @@ rollout: compile_adapter: standin.hipcc parse_adapter: hipk.parse feedback_template: hipk.feedback + # a short kernel can have one obvious answer: 8 identical answers to one prompt is not a broken sampler here + # (the engines' self-test checks that), and with 8 prompts one such group is 12.5% of them + identical_groups_stop: 0.5 reward: profile: hipk-gpu reward_fn: hipk.v1 diff --git a/aorlta/configs/hipk-train.yaml b/aorlta/configs/hipk-train.yaml index 51953aaf5..057b10fdb 100644 --- a/aorlta/configs/hipk-train.yaml +++ b/aorlta/configs/hipk-train.yaml @@ -27,6 +27,9 @@ rollout: compile_adapter: standin.hipcc parse_adapter: hipk.parse feedback_template: hipk.feedback + # a short kernel can have one obvious answer: 8 identical answers to one prompt is not a broken sampler here + # (the engines' self-test checks that), and with 8 prompts one such group is 12.5% of them + identical_groups_stop: 0.5 reward: profile: hipk-gpu reward_fn: hipk.v1 diff --git a/aorlta/docs/TUTORIAL.md b/aorlta/docs/TUTORIAL.md index 6210144e1..2bbf97bae 100644 --- a/aorlta/docs/TUTORIAL.md +++ b/aorlta/docs/TUTORIAL.md @@ -59,6 +59,10 @@ reward: answers early only when no pending leg could change the answer. The service recomputes and checks all of this, and a reward that breaks it is unscored and raises an alarm. - **The check's GPUs** are its own setting (`hipk.run.gpus` in the profile): keep them off `rollout.engines`. +- **Precise tasks.** An iteration stops when more than `rollout.identical_groups_stop` (5% by default) of its groups + have all `K` answers identical, which is how a broken sampler shows. A short kernel can have one obvious answer, + and with 8 prompts one such group is 12.5%, so the example's configs set it to 0.5; the engines' start-up + self-test still checks the sampler, on open-ended chats of its own. ## 1. Check the tasks From 2973a9626f2d9ebe900131c6ca72b19e126047a3 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Thu, 12 Nov 2026 23:33:44 +0000 Subject: [PATCH 06/19] feat(aorlta): tasks with no compile step, and plugins from an installed package's entry point rollout.compile_adapter: none is a built-in for a task whose answers are not compiled (a question and answer, a verdict): every parsed answer passes, and the reward's checks decide. A plugins: entry may name an entry point in the aorlta.plugins group, so a task can live in a private package of its own and be registered without copying it here; AORLTA_PYTHONPATH puts that package on the run job's import path, on the host and in its containers. --- aorlta/jobs/aorlta-run.sbatch | 18 +++++++++++------- aorlta/src/aorlta/config.py | 4 ++-- aorlta/src/aorlta/plugins.py | 22 ++++++++++++++++------ aorlta/src/aorlta/rollout/compile_check.py | 20 +++++++++++++++++++- 4 files changed, 48 insertions(+), 16 deletions(-) diff --git a/aorlta/jobs/aorlta-run.sbatch b/aorlta/jobs/aorlta-run.sbatch index 456e6c067..044e121e0 100644 --- a/aorlta/jobs/aorlta-run.sbatch +++ b/aorlta/jobs/aorlta-run.sbatch @@ -50,6 +50,9 @@ # which holds the registry credentials for pulling rocm/dev-ubuntu-24.04:7.0.2 # AORLTA_MOUNTS default none. More host directories, space-separated, mounted read-only into the # orchestrator's and the checkers' containers +# AORLTA_PYTHONPATH default none. More import paths, colon-separated, after the package's own src/: where a +# plugin package of your own lives (`plugins:` in the config), on the host and in every +# container. Mount its directories too (AORLTA_MOUNTS) # AORLTA_EXCLUDE_NODES default none. Nodes this job refuses to run on (exit 90), as a Slurm host list. #SBATCH # lines cannot read the environment, so give Slurm the same list when submitting: # sbatch --exclude="$AORLTA_EXCLUDE_NODES" jobs/aorlta-run.sbatch @@ -80,6 +83,7 @@ AORLTA_ASAN_LIB=${AORLTA_ASAN_LIB:-} AORLTA_HF_HOME=${AORLTA_HF_HOME:-${HF_HOME:-$HOME/.cache/huggingface}} AORLTA_TMPDIR=${AORLTA_TMPDIR:-${TMPDIR:-/tmp}} read -ra EXTRA_MOUNTS <<< "${AORLTA_MOUNTS:-}" +XPATH=${AORLTA_PYTHONPATH:+:$AORLTA_PYTHONPATH} export TMPDIR=$AORLTA_TMPDIR HF_HOME=$AORLTA_HF_HOME export DOCKER_CONFIG=${AORLTA_DOCKER_CONFIG:-${DOCKER_CONFIG:-$HOME/.docker}} P=$SLURM_SUBMIT_DIR @@ -153,7 +157,7 @@ if [ -n "${AORLTA_SITE:-}" ]; then fi /usr/bin/git -C "$P" archive HEAD | tar -x -C "$CODE" || { echo "[guard] cannot snapshot HEAD of $P" >&2; exit 93; } [ -f "$CODE/$CONFIG" ] || { echo "[guard] $CONFIG is not committed at $COMMIT" >&2; exit 94; } -read -r RUN_ID SYNC_DIR RUN_DIR RUN_ROOT < <(PYTHONPATH="$CODE/src" "$PY" -c " +read -r RUN_ID SYNC_DIR RUN_DIR RUN_ROOT < <(PYTHONPATH="$CODE/src$XPATH" "$PY" -c " from aorlta.config import load_config; c = load_config('$CODE/$CONFIG') print(c.run.run_id, c.sync.dir, c.run_dir, c.run.root)") || { echo "[guard] $CONFIG does not load" >&2; exit 94; } mkdir -p "$RUN_ROOT" || { echo "[guard] cannot create run.root ($RUN_ROOT)" >&2; exit 95; } @@ -161,7 +165,7 @@ mkdir -p "$RUN_ROOT" || { echo "[guard] cannot create run.root ($RUN_ROOT)" >&2; # (for a Hugging Face snapshot, the model's whole cache directory, so the snapshot's links into blobs/ resolve), then # each prompt pool's dataset directory (the pool's path, or its parent if that holds kernels.json and listings/, as the # wiring reads it) -mapfile -t DATA_DIRS < <(PYTHONPATH="$CODE/src" "$PY" -c " +mapfile -t DATA_DIRS < <(PYTHONPATH="$CODE/src$XPATH" "$PY" -c " from pathlib import Path from aorlta.config import load_config c = load_config('$CODE/$CONFIG') @@ -177,7 +181,7 @@ MODEL_DIR=${DATA_DIRS[0]:-} ADAPTER_DIR=${DATA_DIRS[1]:-} POOL_DIRS=("${DATA_DIRS[@]:2}") # reward.remote: the socket directory, then one 'adapter gpu socket' line per checker server -mapfile -t REMOTE < <(PYTHONPATH="$CODE/src" "$PY" -c " +mapfile -t REMOTE < <(PYTHONPATH="$CODE/src$XPATH" "$PY" -c " from aorlta.config import load_config; from aorlta.reward.remote import servers r = load_config('$CODE/$CONFIG').reward.remote if r is not None: @@ -261,7 +265,7 @@ if [ -n "$SOCK_DIR" ]; then -e SLURM_JOB_ID="$SLURM_JOB_ID" -e AORTA_ASAN_ROCM_LIB="$AORLTA_ASAN_LIB" -e TMPDIR="$SOCK_DIR/tmp" -e AORLTA_AORTA_ROOT="$AORLTA_AORTA_ROOT" -e AORLTA_ROCJITSU_ROOT="$AORLTA_ROCJITSU_ROOT" -e AORLTA_LIBSTDCXX_SHIM="$AORLTA_LIBSTDCXX_SHIM" -e AORLTA_SITE="${AORLTA_SITE:-}" - -e PYTHONPATH="$CODE/src" -e PYTHONDONTWRITEBYTECODE=1 -e PYTHONUNBUFFERED=1 -w "$CODE") + -e PYTHONPATH="$CODE/src$XPATH" -e PYTHONDONTWRITEBYTECODE=1 -e PYTHONUNBUFFERED=1 -w "$CODE") GPUS=$(printf '%s\n' "${REMOTE[@]:1}" | awk '{print $2}' | sort -un | paste -sd, -) stamp "CHECKERS: aorta's ASan node-health canary on GPUs $GPUS" docker run --rm --name "$CNAME-canary" "${CHK_DOCKER[@]}" -e ROCR_VISIBLE_DEVICES="$GPUS" --entrypoint bash "$CHK_IMG" \ @@ -341,7 +345,7 @@ if [ "$COMPONENTS" = real ]; then # every module the run imports, in its image, add_mounts IMPORT_MOUNTS rw "$P" add_mounts IMPORT_MOUNTS ro "$AORLTA_AORTA_ROOT" docker run --rm --name "$CNAME-imports" --network none --ulimit core=0 "${IMPORT_MOUNTS[@]}" \ - -e PYTHONPATH="$CODE/src" -e PYTHONDONTWRITEBYTECODE=1 -w "$CODE" --entrypoint python3 "$IMG" \ + -e PYTHONPATH="$CODE/src$XPATH" -e PYTHONDONTWRITEBYTECODE=1 -w "$CODE" --entrypoint python3 "$IMG" \ jobs/aorlta-imports-check.py "$AORLTA_AORTA_ROOT/src" > "$OUT/imports.log" 2>&1 & wait $! rc=$? @@ -377,7 +381,7 @@ orchestrator() { # extra arguments for aorlta.orchestrator.run; core=0, as ever --device=/dev/kfd --device=/dev/dri --group-add video --group-add render \ --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --ipc=host --ulimit memlock=-1 --ulimit core=0 \ --ulimit nofile=65535:65535 --network=host "${ORCH_MOUNTS[@]}" \ - -e SLURM_JOB_ID="$SLURM_JOB_ID" -e PYTHONUNBUFFERED=1 -e PYTHONDONTWRITEBYTECODE=1 -e PYTHONPATH="$CODE/src" \ + -e SLURM_JOB_ID="$SLURM_JOB_ID" -e PYTHONUNBUFFERED=1 -e PYTHONDONTWRITEBYTECODE=1 -e PYTHONPATH="$CODE/src$XPATH" \ -e HF_HOME="$AORLTA_HF_HOME" -e HF_HUB_OFFLINE=1 -e TMPDIR="$AORLTA_TMPDIR" -e AORLTA_SITE="${AORLTA_SITE:-}" \ -e AORLTA_AORTA_ROOT="$AORLTA_AORTA_ROOT" -e AORLTA_ROCJITSU_ROOT="$AORLTA_ROCJITSU_ROOT" \ -e AORLTA_LIBSTDCXX_SHIM="$AORLTA_LIBSTDCXX_SHIM" \ @@ -408,7 +412,7 @@ stamp "RUN rc=$rc (rc 0 is completed, rc 3 is stopped; see run_result.json)" tail -5 "$OUT/run.log" # Every committed iteration's summary names every span it must have; a score run prints its summary instead. -PYTHONPATH="$CODE/src" "$PY" - "$CODE/$CONFIG" <<'EOF' | tee "$OUT/spans-check.txt" +PYTHONPATH="$CODE/src$XPATH" "$PY" - "$CODE/$CONFIG" <<'EOF' | tee "$OUT/spans-check.txt" import sys from aorlta.config import load_config from aorlta.episodes.writer import complete_iterations, read_summary diff --git a/aorlta/src/aorlta/config.py b/aorlta/src/aorlta/config.py index eca6b6a9a..8c85980f4 100644 --- a/aorlta/src/aorlta/config.py +++ b/aorlta/src/aorlta/config.py @@ -428,7 +428,7 @@ class RolloutSection(_Section): max_feedback_chars: int = 2000 request_timeout_s: float = 1800.0 gpu_memory_utilization: float = 0.45 - compile_adapter: str = "standin.hipcc" + compile_adapter: str = "standin.hipcc" # "none": no compile step; a parsed answer ends its chain parse_adapter: str = "standin.parse" leak_check_adapter: str = "none" # the stand-in prompts need none (Section 6.1, Leak check) feedback_template: str = "standin" @@ -461,7 +461,7 @@ def _validate(self) -> None: ) for name in ("compile_adapter", "parse_adapter", "feedback_template", "leak_check_adapter"): value = getattr(self, name) - _check(value != "none" or name == "leak_check_adapter", f"{name}: cannot be 'none'") + _check(value != "none" or name in ("leak_check_adapter", "compile_adapter"), f"{name}: cannot be 'none'") try: component_source(value) except ConfigError as exc: diff --git a/aorlta/src/aorlta/plugins.py b/aorlta/src/aorlta/plugins.py index 76689dc2e..e2c97a8c6 100644 --- a/aorlta/src/aorlta/plugins.py +++ b/aorlta/src/aorlta/plugins.py @@ -1,8 +1,10 @@ """Plugins: a task's own prompts, answer parsing, checks and reward function, by name, without editing this package. -A run config lists its plugins under `plugins:`, each an importable module name or the path of a `.py` file (a -relative path is read from the package root, the directory that holds `src/aorlta`, so a plugin committed beside the -package runs from the run job's snapshot of the tree). Each is imported once, before the rest of the config is +A run config lists its plugins under `plugins:`. Each entry is, in this order: the path of a `.py` file (a relative +path is read from the package root, the directory that holds `src/aorlta`, so a plugin committed beside the package +runs from the run job's snapshot of the tree); the name of an entry point in the `aorlta.plugins` group, which an +installed package of your own declares (`[project.entry-points."aorlta.plugins"] my-task = "my_pkg.plugin"`), so a +private task never has to be copied into this repository; or an importable module name. Each is imported once, before the rest of the config is checked, and registers its components with the `register_*` functions below. A config, a profile or the wiring then finds a name among the built-in components first and the registered ones after: @@ -35,6 +37,7 @@ from typing import Any PACKAGE_ROOT = Path(__file__).resolve().parents[2] # the directory that holds src/aorlta +ENTRY_POINT_GROUP = "aorlta.plugins" class PluginError(ValueError): @@ -90,7 +93,7 @@ class CheckContext: "task": ("standin",), "parse adapter": ("standin.parse",), "feedback template": ("standin",), - "compile adapter": ("standin.hipcc", "fake.compile"), + "compile adapter": ("none", "standin.hipcc", "fake.compile"), "profile": ("standin-static", "standin-gpu", "fake"), "reward function": ("standin.v11",), "check adapter": ("standin.parse", "standin.hipcc", "standin.rj_waitcheck", "standin.aorta_consan", @@ -190,10 +193,17 @@ def _load(entry: str) -> dict: del sys.modules[name] raise PluginError(f"plugins: importing {path} failed: {type(exc).__name__}: {exc}") from exc else: + from importlib.metadata import entry_points + + named = [ep for ep in entry_points(group=ENTRY_POINT_GROUP) if ep.name == entry] + if len(named) > 1: + raise PluginError(f"plugins: {len(named)} installed packages declare the entry point {entry!r} in " + f"{ENTRY_POINT_GROUP}") + target = named[0].value.split(":", 1)[0].strip() if named else entry try: - module = importlib.import_module(entry) + module = importlib.import_module(target) except Exception as exc: - raise PluginError(f"plugins: importing {entry} failed: {type(exc).__name__}: {exc}") from exc + raise PluginError(f"plugins: importing {target} failed: {type(exc).__name__}: {exc}") from exc file = getattr(module, "__file__", None) return {"module": module.__name__, "file": file, "sha256": _sha256(Path(file)) if file else None} diff --git a/aorlta/src/aorlta/rollout/compile_check.py b/aorlta/src/aorlta/rollout/compile_check.py index d66fce6a2..ffb4694b9 100644 --- a/aorlta/src/aorlta/rollout/compile_check.py +++ b/aorlta/src/aorlta/rollout/compile_check.py @@ -224,6 +224,8 @@ def make_compiler(name: str, workdir: str | Path, evidence: EvidenceStore | None (`plugins.register_compile_adapter`), called with the same arguments.""" from .. import plugins + if name == "none": + return NoCompile(evidence=evidence) if name == "standin.hipcc": return HipccCompiler(workdir, evidence=evidence, structure=structure) if name == "fake.compile": @@ -231,10 +233,26 @@ def make_compiler(name: str, workdir: str | Path, evidence: EvidenceStore | None if name in plugins.COMPILE_ADAPTERS: return plugins.COMPILE_ADAPTERS[name](workdir, evidence=evidence, structure=structure) registered = "".join(f", {n}" for n in sorted(plugins.COMPILE_ADAPTERS)) - raise CompileUnavailable(f"compile adapter {name!r} is not available; built: standin.hipcc, fake.compile" + raise CompileUnavailable(f"compile adapter {name!r} is not available; built: none, standin.hipcc, fake.compile" f"{registered}") +class NoCompile: + """`rollout.compile_adapter: none`, for a task whose answers are not compiled (a question and answer, a verdict): + every parsed answer passes, so its chain ends after it and the reward's checks decide. Repair turns then come only + from a parse that rejects an answer with a message.""" + + name = "none" + + def __init__(self, evidence: EvidenceStore | None = None): + self.evidence = evidence + + def check(self, code: str, tag: str) -> CompileResult: + _check_tag(tag) + digest = self.evidence.put("") if self.evidence is not None else sha256_text("") + return CompileResult(ok=True, rc=0, seconds=0.0, errors=[], raw_stderr_sha256=digest) + + class FakeCompiler: """Fails exactly the code that holds `marker`, with one clang-format error per occurrence.""" From d525a806bc11cc77b65c4e4cc47853a9a63cf550 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Thu, 12 Nov 2026 23:53:48 +0000 Subject: [PATCH 07/19] test(aorlta): three task shapes from a package outside aorlta/, and a checklist for picking a vertical tests/test_plugins_out_of_tree.py installs a package of its own (an aorlta.plugins entry point) and registers from it: code rewarded by a compile gate and a reference check, a question with no compile step and no sanitizer, and a reward that is an external tool's verdict. Each loads as a run config, passes the wiring's checks and scores through its own check and reward. ADAPTING.md gains the checklist for picking a vertical (ending with the two-line summary), the three shapes, and the route for a task kept in a private package. --- aorlta/docs/ADAPTING.md | 52 +++- aorlta/tests/test_plugins_out_of_tree.py | 290 +++++++++++++++++++++++ 2 files changed, 341 insertions(+), 1 deletion(-) create mode 100644 aorlta/tests/test_plugins_out_of_tree.py diff --git a/aorlta/docs/ADAPTING.md b/aorlta/docs/ADAPTING.md index 9dad69b3b..1ab0df87c 100644 --- a/aorlta/docs/ADAPTING.md +++ b/aorlta/docs/ADAPTING.md @@ -3,6 +3,29 @@ Code paths are relative to `src/aorlta/`; commands and the `configs/`, `examples/` and `tests/` paths are relative to `aorlta/`, and the commands assume the environment from the README. +## Pick your vertical: a checklist + +1. **The task in one sentence.** What does the model get, and what must it produce? +2. **Is it automatically checkable?** A program must score every answer with no human in the loop: a compiler and + a reference run, a reference answer (exact or normalised match, or a judge command), or another tool's verdict. + If the check cannot always decide, it must say so: an answer it cannot settle is unscored, and **an unscored + sample is never counted clean**. Too many unscored samples stop a training iteration + (`orchestrator.min_scored_fraction`). +3. **Where the prompts come from,** and a held-out set kept out of training to evaluate on. Training takes + `rollout.P` prompts per iteration, so have many more than `P`. +4. **Single-turn or a repair loop?** The episode is the model's answer, then up to `rollout.max_repair_turns` turns + of feedback: compiler errors, or a message your parse gives a rejected answer. Anything else (tool calls, an + environment the model acts in) is not supported yet. +5. **The reward.** Give each failure you can name its own value, so the answers to one prompt get different rewards + (GRPO learns only from those differences); keep it in [-1.55, 1.0] with the parts `base`, `broke_healthy` and + `length`. If the base model already gets nearly everything right, or nothing, there is nothing to learn: score + it first. +6. **The compute budget.** One node; see "How many GPUs" in GETTING-STARTED.md. A short check (a string match) costs + nothing; a GPU run of every answer needs GPUs of its own and its own time limit. +7. **Score the base model** on the held-out set (`orchestrator.mode: score`), **train**, then **score the trained + 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. + ## 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 @@ -36,7 +59,34 @@ The module is imported before the rest of the config is checked, so every name a | a profile pool's `adapter` | `register_check_adapter(name, factory)` | `factory(ctx)` returns a `CheckAdapter` for one worker; `ctx` (`plugins.CheckContext`) holds its profile settings, worker id, clock, timeout and the run's pool paths | A built-in component may be mixed in by name (`standin.parse`, `standin.hipcc` and `standin` are the stand-in -parse, compile and feedback). A component's `name` must be the name it is registered under. The landing check and +parse, compile and feedback). Nothing in a profile is required: a task whose answers are not code sets +`rollout.compile_adapter: none` (every parsed answer ends its chain) and gives its profile only its own checks, with +no compile gate and no sanitizer. Three shapes, each tested from a package outside `aorlta/` +(`tests/test_plugins_out_of_tree.py`): + +| shape | rollout | the profile's legs | +|---|---|---| +| code, rewarded by compiling and matching a reference | a compile check (`standin.hipcc`) and its repair turns | the parse and compile gates, then your reference check (`examples/hip-kernels/`) | +| a question with a reference answer | `compile_adapter: none`; your parse may give a reply with no answer a repair turn | one leg comparing the answer, exactly or normalised, with an optional judge command (`examples/qa/`) | +| a verdict from an external tool or service | `compile_adapter: none` | one leg that calls the tool and reads its verdict; no verdict is unsettled | + +### Keeping your task in your own package + +The recommended route for a task that is private or not yet released: keep its module, prompts, profile and checks +in a package of your own, outside this repository, and declare an entry point: + +```toml +# your package's pyproject.toml +[project.entry-points."aorlta.plugins"] +my-task = "my_package.aorlta_plugin" +``` + +`plugins: [my-task]` then loads it wherever the package is installed (`pip install` into your venv, or your own +image). The run job's containers run the image's Python, which does not see your venv, so for the job name the +module instead, `plugins: [my_package.aorlta_plugin]`, and put the package on the containers' import path and mount +it: `AORLTA_PYTHONPATH=/path/to/my-package/src AORLTA_MOUNTS=/path/to/my-package`. Either way nothing is copied into +this repository, and the run plan records the plugin file's sha256. + A component's `name` must be the name it is registered under. The landing check and the trainer's warm-up read `sync.probes` chats: the task's `probe_messages(pool_dir)`, or the first records' own messages, which must then share one system message. Each engine's start-up self-test samples 4 chats twice with different seeds and needs at least 3 of the 4 pairs to differ, so for a plugin task it uses fixed open-ended chats diff --git a/aorlta/tests/test_plugins_out_of_tree.py b/aorlta/tests/test_plugins_out_of_tree.py new file mode 100644 index 000000000..38a5daf17 --- /dev/null +++ b/aorlta/tests/test_plugins_out_of_tree.py @@ -0,0 +1,290 @@ +"""Three task shapes registered from a package outside aorlta/, through an installed entry point: (a) code rewarded +by a compile gate and a reference check, (b) a question and answer with no compile step and no sanitizer, (c) a +reward that is an external tool's verdict. Each loads as a run config, passes the wiring's checks, and scores a +candidate through its own check adapter and reward function.""" + +from __future__ import annotations + +import json +import stat +import sys +from pathlib import Path + +import pytest + +from aorlta import plugins +from aorlta.config import config_from_dict + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from test_plugins import CharChat, registry # noqa: E402,F401 + +PLUGIN = ''' +import json +import subprocess +from pathlib import Path + +from aorlta import plugins +from aorlta.ids import sha256_ids, sha256_text +from aorlta.reward.adapters.base import make_result, unsettled +from aorlta.reward.rewardfn.base import Decided, NeedMore, Unscored +from aorlta.rollout.compile_check import ParseOutcome +from aorlta.rollout.templates import render_turn0 +from aorlta.schemas import PromptRecord, TaskRef + +HERE = Path(__file__).resolve().parent + + +def records(pool_dir, tokenizer, thinking): + out = [] + for i in range(6): + msgs = [{"role": "system", "content": "s"}, {"role": "user", "content": f"task {i}"}] + out.append(PromptRecord(f"team:{i}", None, msgs, sha256_ids(render_turn0(tokenizer, msgs, thinking)), + TaskRef(f"team:{i}", None, None, None, sha256_text(f"task {i}")), "standin")) + return out + + +class Parse: + name = "team.parse" + reply_stops = (248046,) + + def check_prompts(self, prompts, thinking, max_repair_turns=0): + return None + + def check(self, text, thinking, prompt): + return ParseOutcome("edits", code=text.strip()) + + +class Check: + version, resource = "1", "local_cpu" + + def __init__(self, ctx, leg): + self.name, self.ctx, self.legs = ctx.name, ctx, (leg,) + + def answer(self, job): + return Path(job.inputs["response"]).read_text().strip() + + def controls(self): + return [] + + def health(self): + return {"responsive": True} + + +class Reference(Check): # shape (a), after the stand-in compile gate: the source must hold the reference's marker + def __init__(self, ctx): + super().__init__(ctx, "ref.native") + + def run(self, job): + t = self.ctx.clock() + ok = "MARK" in Path(job.inputs["source"]).read_text() + return make_result(job, "clean" if ok else "fail", None if ok else "wrong_output", t_start=t, t_end=t) + + +class Exact(Check): # shape (b): the answer against a reference, no compile + def __init__(self, ctx): + super().__init__(ctx, "ref.answer") + + def run(self, job): + t = self.ctx.clock() + ok = self.answer(job) == "42" + return make_result(job, "clean" if ok else "fail", None if ok else "wrong_answer", t_start=t, t_end=t) + + +class Tool(Check): # shape (c): an external command's verdict; no verdict is unsettled + def __init__(self, ctx): + super().__init__(ctx, "ext.verdict") + + def run(self, job): + t = self.ctx.clock() + p = subprocess.run(list(self.ctx.settings["tool"]), input=self.answer(job), capture_output=True, text=True) + v = p.stdout.strip() + if v in ("PASS", "FAIL"): + return make_result(job, "clean" if v == "PASS" else "fail", None if v == "PASS" else "rejected", + t_start=t, t_end=t) + return unsettled(job, "no_verdict", t_start=t, t_end=t) + + +class Reward: + version, code_sha256 = "1", "0" * 64 + + def __init__(self, name, legs): + self.name, self.legs = name, legs + + def compute(self, legs, cand): + if any(legs[n].state in ("unsettled", "not_run", "skipped_by_plan", "finding") for n in self.legs): + return None, {"reason": "unsettled"} + base = 1.0 if all(legs[n].state == "clean" for n in self.legs) else -0.5 + return base + 0.0 + 0.0, {"base": base, "broke_healthy": 0.0, "length": 0.0} + + def decide(self, partial, cand): + pending = [n for n in self.legs if n not in partial] + if pending: + return NeedMore((pending[0],)) + r, p = self.compute(partial, cand) + return Unscored(p["reason"]) if r is None else Decided(r, p) + + def claims_clean(self, reward, parts): + return reward is not None and parts["base"] > 0 + + +plugins.register_task(plugins.Task("team", records)) +plugins.register_parse_adapter("team.parse", Parse) +for shape, adapter, leg in (("code", Reference, "ref.native"), ("qa", Exact, "ref.answer"), + ("tool", Tool, "ext.verdict")): + plugins.register_check_adapter(f"team.{shape}", adapter) + legs = ("gate.parse", "gate.compile", leg) if shape == "code" else (leg,) + plugins.register_reward_fn(f"team.{shape}.v1", lambda profile, legs=legs, shape=shape: Reward(f"team.{shape}.v1", legs)) + plugins.register_profile(f"team-{shape}", HERE / f"team-{shape}.yaml") +''' + +PROFILES = { + "code": """name: team-code +reward_fn: team.code.v1 +materialize: true +pools: + gates: {resource: local_cpu, adapter: standin.parse, timeout: gate} + compile: {resource: local_cpu, adapter: standin.hipcc, timeout: compile} + ref: {resource: local_cpu, adapter: team.code, timeout: sim_run} +stages: + - {name: S0, pool: gates, legs: [gate.parse]} + - {name: S1, pool: compile, legs: [gate.compile]} + - {name: S2, pool: ref, legs: [ref.native]} +requires: {gate.compile: [gate.parse], ref.native: [gate.compile]} +leg_settings: {gate.compile: {offload_arch: gfx950, flags: ["-c"]}} +""", + "qa": """name: team-qa +reward_fn: team.qa.v1 +materialize: true +pools: + ref: {resource: local_cpu, adapter: team.qa, timeout: gate} +stages: + - {name: S0, pool: ref, legs: [ref.answer]} +""", + "tool": """name: team-tool +reward_fn: team.tool.v1 +materialize: true +pools: + ext: {resource: local_cpu, adapter: team.tool, timeout: sim_run} +stages: + - {name: S0, pool: ext, legs: [ext.verdict]} +adapters: + team.tool: {tool: [TOOL]} +""", +} + + +@pytest.fixture +def installed(tmp_path, registry, monkeypatch): # noqa: F811 + """A package of a team's own, `team_tasks`, installed (a dist-info with an `aorlta.plugins` entry point) in a + directory outside aorlta/, and a stand-in for the team's verdict tool.""" + site = tmp_path / "site-packages" + pkg = site / "team_tasks" + pkg.mkdir(parents=True) + (pkg / "__init__.py").write_text("") + (pkg / "plugin.py").write_text(PLUGIN) + tool = tmp_path / "verdict-tool" + tool.write_text("#!/bin/bash\nread a\n[ \"$a\" = ok ] && echo PASS || { [ \"$a\" = bad ] && echo FAIL || echo ?; }\n") + tool.chmod(tool.stat().st_mode | stat.S_IXUSR) + for shape, text in PROFILES.items(): + (pkg / f"team-{shape}.yaml").write_text(text.replace("TOOL", str(tool))) + dist = site / "team_tasks-0.1.dist-info" + dist.mkdir() + (dist / "METADATA").write_text("Metadata-Version: 2.1\nName: team-tasks\nVersion: 0.1\n") + (dist / "entry_points.txt").write_text("[aorlta.plugins]\nteam = team_tasks.plugin\n") + monkeypatch.syspath_prepend(str(site)) + yield tmp_path + sys.modules.pop("team_tasks.plugin", None) + sys.modules.pop("team_tasks", None) + + +def config(tmp_path, shape): + compile_adapter = {"code": "fake.compile", "qa": "none", "tool": "none"}[shape] + return config_from_dict({ + "plugins": ["team"], + "run": {"run_id": f"team-{shape}", "source": "standin", "root": str(tmp_path / "runs")}, + "model": {"base_dir": str(tmp_path / "base"), "source": "standin"}, + "data": {"prompt_pools": [{"path": str(tmp_path / "pool"), "source": "standin", "task": "team"}]}, + "rollout": {"K": 2, "P": 4, "compile_adapter": compile_adapter, "parse_adapter": "team.parse"}, + "reward": {"profile": f"team-{shape}", "reward_fn": f"team.{shape}.v1"}, + "orchestrator": {"iterations": 1, "run_seed": 1}, + }) + + +def leg_job(tmp_path, leg, response, source=None): + from aorlta.schemas import LegJob + + (tmp_path / "response.txt").write_text(response) + inputs = {"response": str(tmp_path / "response.txt")} + if source is not None: + (tmp_path / "source.hip").write_text(source) + inputs["source"] = str(tmp_path / "source.hip") + return LegJob(job_id="j", leg=leg, candidate_id="c", control_id=None, source_sha256="s" * 64, task_sha256="t" * 64, + inputs=inputs, output_dir=str(tmp_path)) + + +@pytest.mark.parametrize("shape, leg, good, bad, extra", [ + ("code", "ref.native", ("x", "// MARK"), ("x", "// none"), ("gate.parse", "gate.compile")), + ("qa", "ref.answer", ("42", None), ("41", None), ()), + ("tool", "ext.verdict", ("ok", None), ("bad", None), ()), +]) +def test_a_shape_registered_from_its_own_package_scores_through_its_own_check(installed, shape, leg, good, bad, extra): + from aorlta.orchestrator import wiring + from aorlta.reward.factories import standin_factories + from aorlta.reward.profiles import load_profile + from aorlta.reward.rewardfn.for_profile import standin_rewardfn + from aorlta.schemas import Candidate, LegResult, TaskRef + + cfg = config(installed, shape) + assert plugins.identity()["team"]["module"] == "team_tasks.plugin" + wiring.check_supported(cfg) + profile = load_profile(cfg.reward.profile) + assert profile.legs == (*extra, leg) + if shape != "code": + assert not any(p.adapter.startswith("standin.") for p in profile.pools.values()) # no compile, no sanitizer + + class Clock: + wall = staticmethod(lambda: 0.0) + + adapter = standin_factories(profile, Clock(), pools=(str(installed / "pool"),))[f"team.{shape}"]("w/1") + fn = standin_rewardfn(profile) + cand = Candidate("c", "r", 1, "team:0", 0, "v", "", 1, 10, TaskRef("team:0", None, None, None, "t"), profile.name, + None) + clean = LegResult("x", "clean", None, [], 1, 1, None, None, False, 0.0, 0.0, 0.0) + for (answer, source), want in ((good, 1.0), (bad, -0.5)): + res = adapter.run(leg_job(installed, leg, answer, source)) + legs = {n: clean for n in extra} | {leg: res} + assert fn.compute(legs, cand)[0] == want, (shape, answer, res) + if shape == "tool": + assert adapter.run(leg_job(installed, leg, "huh")).state == "unsettled" # no verdict is never clean + + +def test_two_packages_declaring_one_entry_point_name_are_refused(installed, monkeypatch, tmp_path): + other = tmp_path / "other-site" + dist = other / "other-0.1.dist-info" + dist.mkdir(parents=True) + (dist / "METADATA").write_text("Metadata-Version: 2.1\nName: other\nVersion: 0.1\n") + (dist / "entry_points.txt").write_text("[aorlta.plugins]\nteam = other.plugin\n") + monkeypatch.syspath_prepend(str(other)) + with pytest.raises(plugins.PluginError, match="2 installed packages declare the entry point 'team'"): + plugins.load_plugins(["team"]) + + +def test_an_unknown_entry_point_name_falls_back_to_a_module_import(installed): + with pytest.raises(plugins.PluginError, match="importing no_such_team_module failed"): + plugins.load_plugins(["no_such_team_module"]) + + +def test_compile_adapter_none_passes_every_answer_and_only_it_may_be_none(tmp_path): + from aorlta.config import ConfigError + from aorlta.rollout.compile_check import make_compiler + + res = make_compiler("none", tmp_path).check("not code at all", "c0-0-t0") + assert res.ok and res.rc == 0 and res.errors == [] + base = {"run": {"run_id": "r", "source": "standin", "root": str(tmp_path)}, + "model": {"base_dir": str(tmp_path / "b"), "source": "standin"}, + "data": {"prompt_pools": [{"path": str(tmp_path / "p"), "source": "standin"}]}, + "reward": {"profile": "fake", "reward_fn": "standin.v11"}, + "orchestrator": {"iterations": 1, "run_seed": 1}} + assert config_from_dict({**base, "rollout": {"compile_adapter": "none"}}).rollout.compile_adapter == "none" + with pytest.raises(ConfigError, match="parse_adapter: cannot be 'none'"): + config_from_dict({**base, "rollout": {"parse_adapter": "none"}}) From 424ecdc342a5376213a6badc761b65772dd9f05c Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 07:51:25 +0000 Subject: [PATCH 08/19] docs(aorlta): a diagram of where each plugin component hooks in, and what it receives and returns Co-authored-by: Cursor --- aorlta/docs/ADAPTING.md | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/aorlta/docs/ADAPTING.md b/aorlta/docs/ADAPTING.md index 1ab0df87c..4afe41866 100644 --- a/aorlta/docs/ADAPTING.md +++ b/aorlta/docs/ADAPTING.md @@ -58,6 +58,47 @@ The module is imported before the rest of the config is checked, so every name a | `reward.reward_fn` | `register_reward_fn(name, factory)` | `factory(profile)` returns a `RewardFn` ("Your own reward function") | | a profile pool's `adapter` | `register_check_adapter(name, factory)` | `factory(ctx)` returns a `CheckAdapter` for one worker; `ctx` (`plugins.CheckContext`) holds its profile settings, worker id, clock, timeout and the run's pool paths | +Where each one hooks into a run, and what it receives and returns. Blue boxes are yours; grey boxes are the +framework, which calls them: + +```mermaid +flowchart TD + cfg["run config: plugins: [your module]
imported before the rest of the config is checked;
its register_* calls give each component its name"] + cfg --> task + task["task: register_task(Task(name, records, ...))
records(pool_dir, tokenizer, thinking) returns a list of PromptRecord:
prompt_id, messages or turn0_ids, prompt_ids_sha256,
task (a TaskRef whose task_sha256 hashes what the checks read), source standin"] + task -. "probe_messages(pool_dir): chats for the landing check and the warm-up;
selftest_messages(pool_dir): 4 open-ended chats (both optional)" .-> startup["start-up: engine self-test,
landing probes, trainer warm-up"] + task -- "P prompts per iteration" --> engines["TokenSpeed engines: turn 0, then repair turns
(up to rollout.max_repair_turns)"] + engines -- "the reply's text" --> parse + parse["parse adapter: register_parse_adapter(name, factory)
check(text, thinking, prompt) returns a ParseOutcome:
edits with code, rejected with (message, reason) or None, or no_change
also: reply_stops, and check_prompts(prompts, thinking, max_repair_turns) at start"] + parse -- "edits: the code" --> compile + compile["compile adapter: register_compile_adapter(name, factory)
factory(workdir, evidence=..., structure=...) returns a compiler;
check(code, tag) returns a CompileResult: ok, rc, seconds, errors"] + parse -- "rejected with a message, a turn left" --> feedback + compile -- "not ok, a turn left" --> feedback + feedback["feedback template: register_feedback_template(name, factory)
build(errors, max_chars) returns a Feedback for a failed compile;
build_apply(message, reason) for a rejected answer (optional)"] + feedback -- "Feedback.text, after the leak check" --> engines + compile -- "ok, or the last turn" --> cand + parse -- "no_change; or rejected on the last turn,
without a message, or with no build_apply" --> cand + cfg -. "rollout.compile_adapter: none: every answer ends its chain here" .-> cand + cand["Candidate: response_text, response_tokens, max_new_tokens, task, prompt and chain ids"] + cand --> service["reward service: the profile's stages in order,
settled legs cached by source, task, leg, settings and checker"] + profile["profile: register_profile(name, path to YAML)
pools (adapter, resource, timeout, workers), stages, requires,
leg_settings, adapters' settings, reward_fn, clean_repeats"] + profile -. "configures" .-> service + service -- "a LegJob per leg: leg, task_sha256, inputs (the response, parsed files),
package, output_dir, leg_settings" --> check + check["check adapter: register_check_adapter(name, factory)
factory(ctx: settings, worker_id, clock, timeout_s, pools) returns one per worker;
run(job) returns a LegResult and never raises: clean, finding or fail (settled),
or unsettled; with reason and findings. Also controls() and health()"] + check -- "LegResult" --> service + service -- "the legs finished so far, the candidate" --> rfn + rfn["reward function: register_reward_fn(name, factory); factory(profile)
decide(partial, cand) returns NeedMore(legs), Decided(reward, parts) or Unscored(reason);
compute(legs, cand) returns (reward, parts) or (None, reason); claims_clean(reward, parts)"] + rfn -- "NeedMore: run the next stage" --> service + rfn -- "Decided or Unscored" --> inv{"the service's checks: compute agrees with decide;
reward in -1.55..1.0 with parts base, broke_healthy, length;
a clean claim needs every profile leg clean"} + inv -- "pass" --> grpo["the group's GRPO advantages, then the LoRA trainer"] + inv -- "fail" --> unscored["unscored: reward_invariant_violated"] + + classDef hook fill:#e3eefc,stroke:#2b6cb0,color:#1a365d + classDef core fill:#f2f2f2,stroke:#888,color:#222 + class task,parse,compile,feedback,profile,check,rfn hook + class cfg,startup,engines,cand,service,inv,grpo,unscored core +``` + A built-in component may be mixed in by name (`standin.parse`, `standin.hipcc` and `standin` are the stand-in parse, compile and feedback). Nothing in a profile is required: a task whose answers are not code sets `rollout.compile_adapter: none` (every parsed answer ends its chain) and gives its profile only its own checks, with From 76ba70be6329a06866a3c6834b959c06a10b3cd1 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 07:53:00 +0000 Subject: [PATCH 09/19] docs(aorlta): wider boxes in the plugin diagram Co-authored-by: Cursor --- aorlta/docs/ADAPTING.md | 1 + 1 file changed, 1 insertion(+) diff --git a/aorlta/docs/ADAPTING.md b/aorlta/docs/ADAPTING.md index 4afe41866..61b2cb57b 100644 --- a/aorlta/docs/ADAPTING.md +++ b/aorlta/docs/ADAPTING.md @@ -62,6 +62,7 @@ Where each one hooks into a run, and what it receives and returns. Blue boxes ar framework, which calls them: ```mermaid +%%{init: {"flowchart": {"wrappingWidth": 520}}}%% flowchart TD cfg["run config: plugins: [your module]
imported before the rest of the config is checked;
its register_* calls give each component its name"] cfg --> task From e53bbc90ff6d6801338549fb78352885d1af2fd1 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 09:34:43 +0000 Subject: [PATCH 10/19] fix(aorlta): plugins load all or nothing; profile names, the engine context, the import pre-flight, offline re-scoring - A registered profile whose YAML name: differs from its registered name is refused at config load: a candidate carried one name and the reward service served the other, so every sample was profile_mismatch. - A plugin that raised halfway through its import left its registrations behind, so a retry in the same process was refused as already registered. What the failed import registered is now undone (file and module forms). - The engine context covered only the task's prompts and repairs. It now also fits the start-up requests: the self-test chats (or the pool's turn 0s) plus SELF_TEST_TOKENS, and the landing probes (sync.probe_tokens) plus the round trip's continuation (roundtrip.GREEDY_TOKENS). - The run job's import pre-flight loaded aorta's models and waitcheck for every real run, so a task with checks of its own could not start without an aorta checkout (exit 99). It now loads them only when the profile uses an aorta adapter (reward.profiles.aorta_preflight_modules), decided on the host from the committed config. - selector.plan_eval accepted only standin.v11 before asking the registry; it now builds a plugin's reward function, takes --plugins to load the plugin first, and refuses an unknown profile with a message instead of a traceback. --- aorlta/docs/ADAPTING.md | 6 +- aorlta/jobs/aorlta-imports-check.py | 9 +-- aorlta/jobs/aorlta-run.sbatch | 10 +++- aorlta/src/aorlta/orchestrator/wiring.py | 20 +++++++ aorlta/src/aorlta/plugins.py | 13 ++++- aorlta/src/aorlta/reward/profiles.py | 22 ++++++- aorlta/src/aorlta/selector/plan_eval.py | 26 +++++++-- aorlta/src/aorlta/sync/roundtrip.py | 7 ++- .../tests/orchestrator/test_plugin_wired.py | 33 +++++++++++ aorlta/tests/orchestrator/test_run_job.py | 55 ++++++++++++++++++ aorlta/tests/selector/test_plan_eval.py | 57 +++++++++++++++++++ aorlta/tests/test_plugins.py | 35 ++++++++++++ 12 files changed, 272 insertions(+), 21 deletions(-) diff --git a/aorlta/docs/ADAPTING.md b/aorlta/docs/ADAPTING.md index 61b2cb57b..15619630b 100644 --- a/aorlta/docs/ADAPTING.md +++ b/aorlta/docs/ADAPTING.md @@ -254,7 +254,11 @@ image's compiler and ConSan with the checker image's, and that the two builds ag To add a reward function, write the class, register it (`plugins.register_reward_fn(name, factory)`, where `factory(profile)` builds it), and set `reward_fn:` to its name in your profile, which you register with -`plugins.register_profile`. The config then accepts both names and checks that they agree. +`plugins.register_profile`. The config then accepts both names and checks that they agree, and that the profile's +`name:` is the name it is registered under (a candidate carries one and the reward service serves the other). +A profile that names none of the aorta-backed adapters (`standin.rj_waitcheck`, `standin.aorta_consan`, `standin.aorta_asan`) +needs no aorta checkout: the run job's import pre-flight loads aorta's backends only for one that does. To +re-score a run's plan log offline, `python -m aorlta.selector.plan_eval LOG --plugins `. ## Checks and profiles diff --git a/aorlta/jobs/aorlta-imports-check.py b/aorlta/jobs/aorlta-imports-check.py index 1184ec1ca..871b7cbad 100644 --- a/aorlta/jobs/aorlta-imports-check.py +++ b/aorlta/jobs/aorlta-imports-check.py @@ -21,7 +21,7 @@ DISTS = ("torch", "transformers", "peft", "safetensors", "xxhash", "fla-core", "causal-conv1d", "PyYAML", "click") -def main(aorta_src: str) -> int: +def main(aorta_src: str, aorta_modules: tuple[str, ...] = ("models", "waitcheck")) -> int: bad = {} for name in MODULES: try: @@ -30,7 +30,7 @@ def main(aorta_src: str) -> int: bad[name] = f"{type(exc).__name__}: {exc}"[:300] from aorlta.reward.adapters import aorta_import - for name in ("models", "waitcheck"): # the waitcheck leg runs in this image (the GPU legs do not) + for name in aorta_modules: # the waitcheck leg runs in this image (the GPU legs do not); none: no aorta adapter mod, why = aorta_import.load(name, aorta_src) if mod is None: bad[f"aorta.{name}"] = why @@ -40,9 +40,10 @@ def main(aorta_src: str) -> int: versions[d] = md.version(d) except md.PackageNotFoundError: versions[d] = None - print(json.dumps({"python": sys.version.split()[0], "versions": versions, "failed": bad}, sort_keys=True)) + print(json.dumps({"python": sys.version.split()[0], "versions": versions, "aorta": list(aorta_modules), + "failed": bad}, sort_keys=True)) return 1 if bad else 0 if __name__ == "__main__": - sys.exit(main(sys.argv[1])) + sys.exit(main(sys.argv[1], tuple(sys.argv[2:]))) diff --git a/aorlta/jobs/aorlta-run.sbatch b/aorlta/jobs/aorlta-run.sbatch index 044e121e0..8cd9c1659 100644 --- a/aorlta/jobs/aorlta-run.sbatch +++ b/aorlta/jobs/aorlta-run.sbatch @@ -33,7 +33,8 @@ # TS_SRC no default. TokenSpeed's source tree at commit 3b684074, needed only on a node without the # image aorlta-tokenspeed:3b684074, which is then built from its docker/Dockerfile.amd # AORLTA_AORTA_ROOT default /opt/aorta. aorta at commit d22f6116235c05395a0241a83d6f7e952a988454: its src/ for -# the checkers and the import check, its recipes/sanitizers/fixtures for the pre-flight +# the checkers and the import check, its recipes/sanitizers/fixtures for the pre-flight; needed +# only when the check profile uses an aorta adapter (reward.profiles.AORTA_ADAPTERS) # AORLTA_ROCJITSU_ROOT default /opt/rocjitsu. rocjitsu's bin/rj_waitcheck and lib/librocjitsu_dbi_hooks.so # AORLTA_LIBSTDCXX_SHIM default /opt/libstdcxx-shim. A libstdc++ with GLIBCXX_3.4.31, for rj_waitcheck # (these three are also where the check profiles look for the tools; the job passes them into @@ -341,12 +342,17 @@ if [ "$IMG" != "$BASE" ]; then fi echo "$IMG $(docker image inspect --format '{{.Id}}' "$IMG")" | tee -a "$OUT/image.txt" if [ "$COMPONENTS" = real ]; then # every module the run imports, in its image, before any engine starts + # aorta's backends only when the profile uses an aorta adapter: a task with checks of its own needs no checkout + AORTA_MODULES=$(PYTHONPATH="$CODE/src$XPATH" "$PY" -c " +from aorlta.config import load_config; from aorlta.reward.profiles import aorta_preflight_modules, load_profile +print(' '.join(aorta_preflight_modules(load_profile(load_config('$CODE/$CONFIG').reward.profile))))") || { + echo "[guard] cannot read the check profile of $CONFIG" >&2; exit 94; } IMPORT_MOUNTS=() add_mounts IMPORT_MOUNTS rw "$P" add_mounts IMPORT_MOUNTS ro "$AORLTA_AORTA_ROOT" docker run --rm --name "$CNAME-imports" --network none --ulimit core=0 "${IMPORT_MOUNTS[@]}" \ -e PYTHONPATH="$CODE/src$XPATH" -e PYTHONDONTWRITEBYTECODE=1 -w "$CODE" --entrypoint python3 "$IMG" \ - jobs/aorlta-imports-check.py "$AORLTA_AORTA_ROOT/src" > "$OUT/imports.log" 2>&1 & + jobs/aorlta-imports-check.py "$AORLTA_AORTA_ROOT/src" $AORTA_MODULES > "$OUT/imports.log" 2>&1 & wait $! rc=$? stamp "IMPORTS rc=$rc"; tail -3 "$OUT/imports.log" | cut -c1-2000 diff --git a/aorlta/src/aorlta/orchestrator/wiring.py b/aorlta/src/aorlta/orchestrator/wiring.py index 512539093..e22c46307 100644 --- a/aorlta/src/aorlta/orchestrator/wiring.py +++ b/aorlta/src/aorlta/orchestrator/wiring.py @@ -302,6 +302,22 @@ def rollout_max_len(cfg: RunConfig, tokenizer, records: Sequence) -> int: r.max_repair_turns) +def startup_max_len(cfg: RunConfig, tokenizer, records: Sequence, selftest_messages: Sequence | None) -> int: + """The context the start-up requests need, which the rollouts' budget does not cover when prompts are short and + replies small: the engine self-test (its chats, or the pool's turn 0s, plus `SELF_TEST_TOKENS`) and the landing + probes (`sync.probe_tokens`) plus the round trip's greedy continuation.""" + from ..rollout.engine_pool import SELF_TEST_TOKENS + from ..rollout.templates import render_turn0, turn0_ids + from ..sync.roundtrip import GREEDY_TOKENS + + thinking = cfg.rollout.thinking + if selftest_messages is not None: + selftest = max(len(render_turn0(tokenizer, m, thinking)) for m in selftest_messages) + else: + selftest = max(len(turn0_ids(p, tokenizer, thinking)) for p in records) + return max(selftest + SELF_TEST_TOKENS, cfg.sync.probe_tokens + GREEDY_TOKENS) + + def engine_handles(cfg: RunConfig, max_len: int) -> list: """One TokenSpeed engine process per configured GPU, as WP2's harness starts them, on `served_model_dir`.""" from ..rollout.engine_pool import engine_kwargs, engine_port @@ -1288,6 +1304,10 @@ def build(cfg: RunConfig, *, config_path=None, commit: str | None = None, parts: raise WiringError(f"task {task.name!r}: the engine self-test needs {SELF_TEST_PROMPTS} chats, " f"got {len(selftest)}") selftest = selftest[:SELF_TEST_PROMPTS] + try: + max_len = max(max_len, startup_max_len(cfg, tok, records, selftest)) + except TemplateError as exc: + raise WiringError(f"{pool_cfg.task} engine self-test chats: {exc}") from None compiler = parts.compiler if parts.compiler is not None else rollout_compiler(cfg) handles = parts.handles if parts.handles is not None else engine_handles(cfg, max_len) if cfg.orchestrator.mode == "score": # no trainer, no optimizer, no weight sync diff --git a/aorlta/src/aorlta/plugins.py b/aorlta/src/aorlta/plugins.py index e2c97a8c6..ff5cfd993 100644 --- a/aorlta/src/aorlta/plugins.py +++ b/aorlta/src/aorlta/plugins.py @@ -87,6 +87,7 @@ class CheckContext: REWARD_FNS: dict[str, Callable[[Any], Any]] = {} CHECK_ADAPTERS: dict[str, Callable[[CheckContext], Any]] = {} LOADED: dict[str, dict] = {} # plugin as listed -> {"module", "file", "sha256"} +_TABLES = (TASKS, PARSE_ADAPTERS, FEEDBACK_TEMPLATES, COMPILE_ADAPTERS, PROFILES, REWARD_FNS, CHECK_ADAPTERS) _LOCK = threading.RLock() _BUILT_IN = { @@ -151,8 +152,7 @@ def register_check_adapter(name: str, factory: Callable[[CheckContext], Any]) -> def registered(name: str) -> bool: """Whether a plugin registered `name` as any kind of component.""" - return any(name in t for t in (TASKS, PARSE_ADAPTERS, FEEDBACK_TEMPLATES, COMPILE_ADAPTERS, PROFILES, REWARD_FNS, - CHECK_ADAPTERS)) + return any(name in t for t in _TABLES) def _sha256(path: Path) -> str: @@ -169,7 +169,14 @@ def load_plugins(plugins: Sequence[str]) -> None: with _LOCK: if entry in LOADED: continue - LOADED[entry] = _load(entry) + before = {id(t): set(t) for t in _TABLES} + try: + LOADED[entry] = _load(entry) + except BaseException: + for t in _TABLES: # what this import registered before failing is undone, nothing else + for name in set(t) - before[id(t)]: + del t[name] + raise def _load(entry: str) -> dict: diff --git a/aorlta/src/aorlta/reward/profiles.py b/aorlta/src/aorlta/reward/profiles.py index b0d0822a4..9ab4f87f4 100644 --- a/aorlta/src/aorlta/reward/profiles.py +++ b/aorlta/src/aorlta/reward/profiles.py @@ -154,6 +154,17 @@ def profile_problems(p: Profile) -> list[str]: return out +# The adapters that import aorta's sanitizer backends from AORLTA_AORTA_ROOT; a profile naming none of them needs no +# aorta checkout. +AORTA_ADAPTERS = frozenset({"standin.rj_waitcheck", "standin.aorta_consan", "standin.aorta_asan"}) + + +def aorta_preflight_modules(profile: Profile) -> tuple[str, ...]: + """The aorta modules the run job's import pre-flight loads in the run image (`jobs/aorlta-imports-check.py`): + `models` and `waitcheck` when a pool of the profile uses an aorta adapter, none otherwise.""" + return ("models", "waitcheck") if any(p.adapter in AORTA_ADAPTERS for p in profile.pools.values()) else () + + def profile_from_dict(data: Mapping[str, Any], source: str = "") -> Profile: if not isinstance(data, Mapping): raise ConfigError(f"{source}: a profile is a mapping") @@ -170,6 +181,11 @@ def load_profile(name_or_path: str | Path) -> Profile: from .. import plugins path = Path(name_or_path) - if path.suffix not in (".yaml", ".yml"): - path = plugins.profile_path(str(name_or_path)) or PROFILE_DIR / f"{name_or_path}.yaml" - return profile_from_dict(parse_yaml(path.read_bytes(), str(path)), str(path)) + if path.suffix in (".yaml", ".yml"): + return profile_from_dict(parse_yaml(path.read_bytes(), str(path)), str(path)) + path = plugins.profile_path(str(name_or_path)) or PROFILE_DIR / f"{name_or_path}.yaml" + profile = profile_from_dict(parse_yaml(path.read_bytes(), str(path)), str(path)) + if profile.name != name_or_path: # candidates carry the name asked for; the service serves the file's + raise ConfigError(f"{path}: profile {str(name_or_path)!r} names itself {profile.name!r}; a candidate would " + f"carry one and the reward service serve the other (profile_mismatch on every sample)") + return profile diff --git a/aorlta/src/aorlta/selector/plan_eval.py b/aorlta/src/aorlta/selector/plan_eval.py index f5c41aa2e..2cd1acf4e 100644 --- a/aorlta/src/aorlta/selector/plan_eval.py +++ b/aorlta/src/aorlta/selector/plan_eval.py @@ -56,10 +56,12 @@ from typing import Any from ..atomic import atomic_write_bytes, atomic_write_json, read_text +from ..config import ConfigError from ..ids import canonical_json, candidate_id, policy_version from ..reward.profiles import Profile, load_profile from ..reward.rewardfn.base import Decided, NeedMore, Unscored -from ..reward.rewardfn.for_profile import standin_rewardfn +from .. import plugins +from ..reward.rewardfn.for_profile import UnsupportedRewardFn, standin_rewardfn from ..reward.rewardfn.standin_v11 import StandinV11 from ..schemas import SETTLED_STATES, Candidate, LegResult, RewardResult, SchemaError, TaskRef, from_dict from .base import Z95, SafetySwitch, audit_coin, audit_draw, is_defect, wilson_interval @@ -457,11 +459,18 @@ def evaluate(samples: Sequence[Sample], *, rewardfn: Any = None, cap: float = DE """The whole evaluation as JSON values (see the module docstring).""" names = {s.cand.check_profile for s in samples} _need(len(names) == 1, f"one check profile per log, got {sorted(names)}") - profile = load_profile(names.pop()) + name = names.pop() + try: + profile = load_profile(name) + except (OSError, ConfigError) as exc: + raise PlanEvalError(f"check profile {name!r}: {exc}; if a plugin registers it, load the plugin " + f"(--plugins)") from None if rewardfn is None: - _need(profile.reward_fn == "standin.v11", f"profile {profile.name!r} is scored by " - f"{profile.reward_fn!r}; pass its reward function") - rewardfn = standin_rewardfn(profile) + try: + rewardfn = standin_rewardfn(profile) + except UnsupportedRewardFn as exc: + raise PlanEvalError(f"{exc}; load the plugin that registers it (--plugins), or pass its reward " + f"function") from None identity = {"name": rewardfn.name, "version": rewardfn.version, "code_sha256": rewardfn.code_sha256} for s in samples: _need(s.result.reward_fn == identity, f"plan log line {s.line}: scored by {s.result.reward_fn}, not by the " @@ -715,6 +724,8 @@ def main(argv: list[str] | None = None) -> int: ap.add_argument("--max-miss-rate", type=float, default=0.02) ap.add_argument("--min-scored-per-group", type=int, default=2) ap.add_argument("--out", type=Path, help="write plan-eval.json and plan-eval.md here") + ap.add_argument("--plugins", default="", help="plugins to load first, as a run config's `plugins:` lists them " + "(module names or .py paths, comma separated): those of a log whose profile a plugin registered") args = ap.parse_args(argv) if (args.log is None) == (args.fake is None): ap.error("give a plan log, or --fake PATH") @@ -722,6 +733,11 @@ def main(argv: list[str] | None = None) -> int: if args.fake is not None: path = write_jsonl(args.fake, fake_log()) fractions = [float(x) for x in args.audit_fractions.split(",") if x.strip()] + try: + plugins.load_plugins([x.strip() for x in args.plugins.split(",") if x.strip()]) + except plugins.PluginError as exc: + print(f"plan_eval: {exc}", file=sys.stderr) + return 2 try: report = evaluate(read_plan_log(path), cap=args.cap, audit_fractions=fractions, max_miss_rate=args.max_miss_rate, min_scored_per_group=args.min_scored_per_group) diff --git a/aorlta/src/aorlta/sync/roundtrip.py b/aorlta/src/aorlta/sync/roundtrip.py index 27f56d94d..cc5406ba0 100644 --- a/aorlta/src/aorlta/sync/roundtrip.py +++ b/aorlta/src/aorlta/sync/roundtrip.py @@ -26,6 +26,7 @@ from aorlta.sync.reload import ReloadError, reload_all Draw = tuple[tuple[float, ...], tuple[int, ...]] +GREEDY_TOKENS = 16 # each draw's greedy continuation after its probe def draw_request(ids: Sequence[int], greedy_tokens: int) -> dict[str, Any]: @@ -33,7 +34,7 @@ def draw_request(ids: Sequence[int], greedy_tokens: int) -> dict[str, Any]: return_logprob=True, logprob_start_len=0) -def draw(engine, ids: Sequence[int], greedy_tokens: int = 16) -> Draw: +def draw(engine, ids: Sequence[int], greedy_tokens: int = GREEDY_TOKENS) -> Draw: out = engine.generate(**draw_request(ids, greedy_tokens)) lps = parse_input_logprobs(out, ids) if not all(math.isfinite(v) for v in lps): @@ -45,7 +46,7 @@ def draw(engine, ids: Sequence[int], greedy_tokens: int = 16) -> Draw: return tuple(lps), tuple(int(t) for t in gen) -def probe_draws(engines: Sequence[Any], probes: ProbeSet, reps: int, greedy_tokens: int = 16) -> list[list[Draw]]: +def probe_draws(engines: Sequence[Any], probes: ProbeSet, reps: int, greedy_tokens: int = GREEDY_TOKENS) -> list[list[Draw]]: """[rep][engine-major, probe-minor] draws, every probe scored alone and in sequence.""" return [[draw(e, ids, greedy_tokens) for e in engines for ids in probes.ids] for _ in range(reps)] @@ -88,7 +89,7 @@ def judge(base: list[list[Draw]], pert: list[list[Draw]], rest: list[list[Draw]] def round_trip(engines: Sequence[Any], probes: ProbeSet, *, perturbed_slot: str | Path, restore_slot: str | Path, - expected: int, in_flight: Callable[[], int], reps: int = 3, greedy_tokens: int = 16) -> RoundTripResult: + expected: int, in_flight: Callable[[], int], reps: int = 3, greedy_tokens: int = GREEDY_TOKENS) -> RoundTripResult: if not probes.ids: raise LandingError("no probes: the round trip would be judged on zero draws") t0 = time.monotonic() diff --git a/aorlta/tests/orchestrator/test_plugin_wired.py b/aorlta/tests/orchestrator/test_plugin_wired.py index cecc33c79..05ff5eb5a 100644 --- a/aorlta/tests/orchestrator/test_plugin_wired.py +++ b/aorlta/tests/orchestrator/test_plugin_wired.py @@ -103,3 +103,36 @@ def test_a_task_may_give_its_own_self_test_chats(tmp_path, tiny_base, plugin, mo plugins.TASKS["toy"] = dataclasses.replace(toy, selftest_messages=lambda pool_dir: chats[:3]) with pytest.raises(wiring.WiringError, match="the engine self-test needs 4 chats, got 3"): build(cfg, tiny_base) + + +def test_the_engine_context_fits_the_start_up_requests(tmp_path, plugin, monkeypatch): + """Short prompts, no repair turn and a small reply budget gave the engines a context shorter than the self-test's + requests (its chats plus SELF_TEST_TOKENS) and the round trip's (sync.probe_tokens plus its continuation), so + the run would have died at start.""" + from aorlta.rollout.engine_pool import SELF_TEST_TOKENS + from aorlta.rollout.templates import render_turn0 + from aorlta.sync.roundtrip import GREEDY_TOKENS + + tok = E.ByteTokenizer() + seen = [] + + class Sized(Exception): + pass + + def engine_handles(cfg, max_len): + seen.append(max_len) + raise Sized + + monkeypatch.setattr(wiring, "engine_handles", engine_handles) + selftest = max(len(render_turn0(tok, list(m), "off")) for m in wiring.SELF_TEST_CHATS) + SELF_TEST_TOKENS + for probe_tokens in (64, 4096): + data = TP.toy_config(tmp_path, str(plugin)) + data["rollout"].update({"max_new_tokens": 16, "max_repair_turns": 0}) + data["orchestrator"] = {"mode": "score", "run_seed": 3, "score": {"prompts": 4, "passes": 1}} + data["sync"]["probe_tokens"] = probe_tokens + cfg = config_from_dict(data) + with pytest.raises(Sized): + wiring.build(cfg, parts=wiring.Parts(tokenizer=tok, probes=E.PROBES)) + records = wiring.TaskPromptPool(plugins.task("toy"), tok, "off", str(tmp_path / "pool")).records() + assert wiring.rollout_max_len(cfg, tok, records) < min(selftest, probe_tokens + GREEDY_TOKENS) + assert seen == [selftest, 4096 + GREEDY_TOKENS] diff --git a/aorlta/tests/orchestrator/test_run_job.py b/aorlta/tests/orchestrator/test_run_job.py index c12102fc8..5fae50080 100644 --- a/aorlta/tests/orchestrator/test_run_job.py +++ b/aorlta/tests/orchestrator/test_run_job.py @@ -88,3 +88,58 @@ def test_an_adapter_inside_the_hugging_face_cache_is_mounted_read_only(tmp_path) out = subprocess.run(["bash", "-c", f"{funcs}{state}{mounting}printf '%s\\n' \"${{ORCH_MOUNTS[@]}}\""], capture_output=True, text=True, check=True, timeout=60).stdout.split() assert f"{hf}:{hf}" in out and f"{adapter}:{adapter}:ro" in out + + +def _aorta_modules_snippet() -> str: + """The Python the job runs on the host to pick the aorta modules of the import pre-flight.""" + block = JOB[JOB.index("AORTA_MODULES=$("):] + return block[block.index('-c "') + 4:block.index('") ||')] + + +def test_the_import_pre_flight_loads_aorta_only_for_a_profile_that_uses_it(tmp_path, monkeypatch): + """A task whose profile names no aorta adapter starts without an aorta checkout; a stand-in profile still has + its waitcheck backend imported in the run image, and a missing checkout still stops it there.""" + import json + import sys + + sys.path.insert(0, str(PKG / "tests")) + import test_plugins as TP + + from aorlta import plugins + from aorlta.reward.profiles import aorta_preflight_modules, load_profile + + assert aorta_preflight_modules(load_profile("standin-static")) == ("models", "waitcheck") + assert aorta_preflight_modules(load_profile("standin-gpu")) == ("models", "waitcheck") + assert aorta_preflight_modules(load_profile("fake")) == () + saved = {id(t): dict(t) for t in plugins._TABLES} + try: + for t in plugins._TABLES: + t.clear() + plugins.LOADED.clear() + d = tmp_path / "toyplugin" + d.mkdir() + (d / "toy_plugin.py").write_text(TP.PLUGIN) + (d / "toy-profile.yaml").write_text(TP.PROFILE) + (tmp_path / "pool").mkdir() + toy = TP.toy_config(tmp_path, str(d / "toy_plugin.py")) + standin = TP.toy_config(tmp_path, str(d / "toy_plugin.py")) + standin["reward"] = {"profile": "standin-static", "reward_fn": "standin.v11"} + picked = {} + for name, data in (("toy", toy), ("standin", standin)): + (tmp_path / f"{name}.yaml").write_text(json.dumps(data)) + code = _aorta_modules_snippet().replace("$CODE/$CONFIG", str(tmp_path / f"{name}.yaml")) + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, + env={"PYTHONPATH": str(PKG / "src"), "PATH": "/usr/bin:/bin"}) + assert out.returncode == 0, out.stderr + picked[name] = out.stdout.split() + assert picked == {"toy": [], "standin": ["models", "waitcheck"]} + finally: + for t in plugins._TABLES: + t.clear() + t.update(saved[id(t)]) + check = [sys.executable, str(PKG / "jobs" / "aorlta-imports-check.py"), str(tmp_path / "no-aorta" / "src")] + env = {"PYTHONPATH": str(PKG / "src"), "PATH": "/usr/bin:/bin"} + for extra, aorta in (([], []), (["models", "waitcheck"], ["aorta.models", "aorta.waitcheck"])): + report = json.loads(subprocess.run(check + extra, capture_output=True, text=True, env=env).stdout) + assert report["aorta"] == extra + assert sorted(k for k in report["failed"] if k.startswith("aorta.")) == aorta diff --git a/aorlta/tests/selector/test_plan_eval.py b/aorlta/tests/selector/test_plan_eval.py index b800ca71a..53c354b76 100644 --- a/aorlta/tests/selector/test_plan_eval.py +++ b/aorlta/tests/selector/test_plan_eval.py @@ -427,3 +427,60 @@ def test_the_cli_writes_the_report_and_refuses_a_bad_log(tmp_path, capsys): (tmp_path / "bad.jsonl").write_text('{"schema": "nope"}\n') assert main([str(tmp_path / "bad.jsonl")]) == 2 assert "plan_eval:" in capsys.readouterr().err + + +TOY_V11_PLUGIN = ''' +from pathlib import Path + +from aorlta import plugins +from aorlta.reward.profiles import PROFILE_DIR +from aorlta.reward.rewardfn.standin_v11 import StandinV11 + + +class ToyV11(StandinV11): + name = "toy.v11" + + +here = Path(__file__).parent +(here / "toy-fake.yaml").write_text((PROFILE_DIR / "fake.yaml").read_text().replace("name: fake", "name: toy-fake") + .replace("reward_fn: standin.v11", "reward_fn: toy.v11")) +plugins.register_profile("toy-fake", here / "toy-fake.yaml") +plugins.register_reward_fn("toy.v11", lambda profile: ToyV11(clean_repeats=profile.clean_repeats)) +''' + + +def test_a_log_whose_profile_and_reward_a_plugin_registered_is_evaluated_once_the_plugin_is_loaded(tmp_path, capsys): + """The evaluator refused any reward function but standin.v11 before asking the registry; now `--plugins` loads + the plugin and the log is re-scored by the reward it registered, with the same counts as the built-in one.""" + from aorlta import plugins + + saved = {id(t): dict(t) for t in plugins._TABLES} + loaded = dict(plugins.LOADED) + try: + (tmp_path / "toy_v11_plugin.py").write_text(TOY_V11_PLUGIN) + lines = fake_log(iterations=1, prompts=4, K=4) + ident = None + for line in lines: + line["candidate"]["check_profile"] = "toy-fake" + if ident is None: + from aorlta.reward.rewardfn.standin_v11 import StandinV11 + + ident = dict(line["result"]["reward_fn"], name="toy.v11") + assert line["result"]["reward_fn"]["name"] == StandinV11.name + line["result"]["reward_fn"] = ident + log = write_jsonl(tmp_path / "toy.jsonl", lines) + assert main([str(log), "--audit-fractions", ""]) == 2 + err = capsys.readouterr().err + assert "check profile 'toy-fake'" in err and "--plugins" in err + out = tmp_path / "out" + assert main([str(log), "--audit-fractions", "", "--plugins", str(tmp_path / "toy_v11_plugin.py"), + "--out", str(out)]) == 0 + report = json.loads((out / "plan-eval.json").read_text()) + assert report["log"]["reward_fn"]["name"] == "toy.v11" and report["log"]["profile"] == "toy-fake" + assert report["log"]["reproduced"] == len(lines) + finally: + for t in plugins._TABLES: + t.clear() + t.update(saved[id(t)]) + plugins.LOADED.clear() + plugins.LOADED.update(loaded) diff --git a/aorlta/tests/test_plugins.py b/aorlta/tests/test_plugins.py index c76c16e2d..85311fff4 100644 --- a/aorlta/tests/test_plugins.py +++ b/aorlta/tests/test_plugins.py @@ -278,6 +278,41 @@ def test_a_plugin_is_imported_once(tmp_path, registry): assert builtins.__dict__.pop("_aorlta_count") == [1] +@pytest.mark.parametrize("form", ["file", "module"]) +def test_a_plugin_that_raises_leaves_nothing_registered(tmp_path, registry, monkeypatch, form): + """What a plugin registered before raising is undone, so nothing sees its components and, fixed, it loads in + the same process (its factories are new objects, which a leftover registration would refuse).""" + body = ("from aorlta import plugins\nplugins.register_parse_adapter('half.parse', lambda: None)\n" + "plugins.register_reward_fn('half.v1', lambda profile: None)\n") + p = tmp_path / "half_plugin.py" + p.write_text(body + "raise RuntimeError('boom')\n") + monkeypatch.syspath_prepend(str(tmp_path)) + entry = str(p) if form == "file" else "half_plugin" + with pytest.raises(plugins.PluginError, match="RuntimeError: boom"): + plugins.load_plugins([entry]) + assert not plugins.registered("half.parse") and not plugins.registered("half.v1") and entry not in plugins.LOADED + p.write_text(body) + plugins.load_plugins([entry]) + assert plugins.registered("half.parse") and plugins.registered("half.v1") and entry in plugins.LOADED + + +def test_a_registered_profile_must_carry_its_registered_name(tmp_path, plugin): + """The candidates carry the registered name and the reward service serves the profile file's: a file naming + anything else (here a built-in profile's name) would have every candidate refused as profile_mismatch.""" + from aorlta.reward.profiles import load_profile + + other = tmp_path / "other.yaml" + other.write_text(PROFILE.replace("name: toy-profile", "name: fake")) + plugins.register_profile("toy-other", other) + with pytest.raises(ConfigError, match="'toy-other' names itself 'fake'"): + load_profile("toy-other") + data = toy_config(tmp_path, str(plugin)) + data["reward"]["profile"] = "toy-other" + with pytest.raises(ConfigError, match="'toy-other' names itself 'fake'"): + config_from_dict(data) + assert load_profile("toy-profile").name == "toy-profile" and load_profile(other).name == "fake" + + def test_the_run_plan_records_each_plugin_file(tmp_path, plugin): cfg = config_from_dict(toy_config(tmp_path, str(plugin))) plan = json.loads(write_run_plan(cfg, commit="a" * 40).read_text()) From cec0d31bd9975e153da2d47000e99dca73a3e8e2 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 09:35:35 +0000 Subject: [PATCH 11/19] fix(aorlta): the HIP example never runs the answer's host code; Q&A judge and eval_compare fixes HIP-kernel example (hip_kernels.py): - The candidate was compiled into the same program as the checker's main, so a static initializer could print the result line and exit before any comparison. The candidate is now compiled alone into a device code object (hipcc --genco): none of its host code is built or run. The harness is built from the task alone, loads that code object, launches the kernel by its exact mangled signature, and writes its verdict to a record file; stdout is never read. A code object without the kernel is bad_signature. - A hipMalloc, first copy or code-object load failure is the node's (unsettled setup_failed), not a candidate crash. - Workers lease GPUs from hipk.run.gpus and give them back on close(): a replacement worker's sequence number said nothing about which GPU its predecessor freed. - The checker version names the compiler hipcc runs and this file, so a cached ref.native leg is not reused across either. - hipk.gate: the reward's parse gate uses the rollout's parse rule, so a final reply whose last block was cut off is not scored on an earlier complete block. - The evidence archive keeps ref.* legs' logs, harness sources and records, not only the sanitizer legs'. Q&A example: the judge (command, timeout, script) is part of the checker version; it runs through ProcessRunner (which now takes stdin), so a timeout or close() kills its whole process group. tools/eval_compare.py refuses runs that differ in prompt order, batch options, reward, task hashes or per-prompt sample counts, and shows a sample with no reply instead of crashing. --- aorlta/docs/TUTORIAL.md | 23 +- aorlta/examples/hip-kernels/hip_kernels.py | 343 +++++++++++++++------ aorlta/examples/hip-kernels/hipk-gpu.yaml | 13 +- aorlta/examples/qa/qa_task.py | 27 +- aorlta/src/aorlta/reward/adapters/base.py | 8 +- aorlta/src/aorlta/reward/evidence.py | 12 +- aorlta/tests/examples/test_hip_kernels.py | 191 +++++++++--- aorlta/tests/examples/test_qa.py | 43 +++ aorlta/tests/reward/test_evidence.py | 4 + aorlta/tests/tools/test_eval_compare.py | 53 ++++ aorlta/tools/eval_compare.py | 26 +- 11 files changed, 563 insertions(+), 180 deletions(-) diff --git a/aorlta/docs/TUTORIAL.md b/aorlta/docs/TUTORIAL.md index 5e8805213..45c2da22a 100644 --- a/aorlta/docs/TUTORIAL.md +++ b/aorlta/docs/TUTORIAL.md @@ -14,7 +14,7 @@ Paths are relative to `aorlta/`. The commands assume the environment from the RE |---|---|---| | the prompts | `examples/hip-kernels/tasks.json`, and `records()` in `hip_kernels.py` | one prompt record per task: the chat, its prompt id, and a hash of what the checks read for it | | how an answer is read | `HipkParse` and `HipkFeedback` | the last ```cpp block is the answer; no block, or no definition of the kernel asked for, gets a repair turn that says so; a compile error gets one with the compiler's errors | -| the checks | `HipkRun` and `examples/hip-kernels/hipk-gpu.yaml` | the stand-in parse and compile gates, then `ref.native`: build the answer with a generated host program, run it on a GPU, compare | +| the checks | `HipkGate`, `HipkRun` and `examples/hip-kernels/hipk-gpu.yaml` | a parse gate with the rollout's own rule and the stand-in compile gate, then `ref.native`: compile the answer alone into a code object, run its kernel on a GPU from a harness built from the task, compare | | the reward | `HipkV1` | -1.0 no parsable answer, -0.9 compile error, -0.8 wrong signature, -0.6 crash or hang, -0.5 wrong output, +1.0 correct, minus 0.05 x the share of the token budget used | The module registers each part under a name (`register()` at its end): @@ -25,6 +25,7 @@ plugins.register_parse_adapter("hipk.parse", HipkParse) plugins.register_feedback_template("hipk.feedback", HipkFeedback) plugins.register_profile("hipk-gpu", HERE / "hipk-gpu.yaml") plugins.register_check_adapter("hipk.run", HipkRun) +plugins.register_check_adapter("hipk.gate", HipkGate) plugins.register_reward_fn("hipk.v1", HipkV1) ``` @@ -54,11 +55,15 @@ reward: - **A check never guesses.** Anything it cannot settle (no GPU, a tool missing, output it cannot read) is `unsettled`, never `clean`; only the check may blame the candidate for a timeout. `verdict()` in the example is the whole mapping, and `tests/examples/test_hip_kernels.py` tests every row of it on CPU. +- **The candidate is not trusted.** Its verdict must come from code it cannot change: the example never runs the + answer's host code and reads only the harness's own record file. A check's `version` must change whenever its + verdict could (the example's names its compiler and its own file), because cached legs are keyed by it. - **The reward's parts.** `compute` returns `(reward, {"base", "broke_healthy", "length"})`, which add up to the reward, within [-1.55, 1.0]; an unscored sample is `(None, {"reason": ...})` and is never counted clean. `decide` answers early only when no pending leg could change the answer. The service recomputes and checks all of this, and a reward that breaks it is unscored and raises an alarm. -- **The check's GPUs** are its own setting (`hipk.run.gpus` in the profile): keep them off `rollout.engines`. +- **The check's GPUs** are its own setting (`hipk.run.gpus` in the profile): keep them off `rollout.engines`. Each + live worker holds one of them, and a worker the service replaces takes the GPU its predecessor gave back. - **Precise tasks.** An iteration stops when more than `rollout.identical_groups_stop` (5% by default) of its groups have all `K` answers identical, which is how a broken sampler shows. A short kernel can have one obvious answer, and with 8 prompts one such group is 12.5%, so the example's configs set it to 0.5; the engines' start-up @@ -79,9 +84,12 @@ and run them on one GPU (one line per task, all `"pass": true`; on a cluster, in python examples/hip-kernels/hip_kernels.py selfcheck --hipcc /opt/rocm/bin/hipcc --gpu 0 ``` -The harness is the candidate's source followed by a host `main` that fills the inputs from a fixed seed, runs the -kernel once with the launch the prompt states, and compares every output element with the task's `reference` loop -(host C++), within `atol + rtol * |want|`. Output elements start as a sentinel, so an element the kernel never +The candidate is compiled on its own into a device code object (`hipcc --genco`), so none of its host code is built +or run. The harness is a host program built from the task alone: it loads that code object, fills the inputs from a +fixed seed, runs the task's kernel once with the launch the prompt states, and compares every output element with +the task's `reference` loop (host C++), within `atol + rtol * |want|`. It writes its verdict to a file the reward +reads; nothing the candidate prints is read, and a failure before the launch (an allocation, the first copy) leaves +the check unsettled rather than blaming the candidate. Output elements start as a sentinel, so an element the kernel never writes is a mismatch, and every array has a guard region past its end, so a write past the end is caught too. ## 2. Run the tests @@ -125,8 +133,9 @@ export AORLTA_SITE=/path/to/site.yaml AORLTA_VENV=/path/to/venv TS_SRC=/path/to/ mkdir -p runs && CONFIG=configs/hipk-score.yaml sbatch jobs/aorlta-run.sbatch ``` -This task needs none of the stand-in sanitizer tools (rocjitsu, aorta's sanitizer backends, the ASan runtime): only -the TokenSpeed image the job builds, whose hipcc builds and runs the kernels. The job prints `score/summary.txt` at +This task needs none of the stand-in sanitizer tools (rocjitsu, aorta's sanitizer backends, the ASan runtime) and no +aorta checkout: the job's import pre-flight loads aorta's backends only for a profile that uses one (#611). It needs +only the TokenSpeed image the job builds, whose hipcc builds and runs the kernels. The job prints `score/summary.txt` at the end; `score/report.jsonl` has one row per prompt. ## 5. Train diff --git a/aorlta/examples/hip-kernels/hip_kernels.py b/aorlta/examples/hip-kernels/hip_kernels.py index 3a9a1a2be..e7b1ad8d4 100644 --- a/aorlta/examples/hip-kernels/hip_kernels.py +++ b/aorlta/examples/hip-kernels/hip_kernels.py @@ -4,12 +4,13 @@ answers with a HIP source in one ```cpp block. The rollout compiles it with the stand-in compile check (hipcc -c) and gives up to `rollout.max_repair_turns` repair turns with the compiler's errors; a reply that has no code block, or whose code does not define the kernel, gets a repair turn that says so (this module's parse and feedback). The reward -service then runs three legs: the stand-in parse gate and compile gate, and `ref.native`, this module's check, which -builds the candidate with a generated host program, runs it on a GPU and compares every output element with a host -reference computed from the same inputs. The reward is `hipk.v1` (below). +service then runs three legs: this module's parse gate (the rollout's own parse rule) and the stand-in compile gate, +and `ref.native`, this module's check, which compiles the candidate alone into a device code object, runs its kernel +on a GPU from a host program built from the task alone, and compares every output element with a host reference +computed from the same inputs. The reward is `hipk.v1` (below). Registered names: task `hipk`; parse `hipk.parse`; feedback `hipk.feedback`; profile `hipk-gpu` (hipk-gpu.yaml); -check adapter `hipk.run`; reward function `hipk.v1`. +check adapters `hipk.gate` and `hipk.run`; reward function `hipk.v1`. Run as a script, it checks the tasks themselves with the reference kernels in reference/: python examples/hip-kernels/hip_kernels.py selfcheck --hipcc /opt/rocm/bin/hipcc [--gpu 0] [--build-only] @@ -25,6 +26,7 @@ import shutil import sys import tempfile +import threading import time from pathlib import Path @@ -36,7 +38,6 @@ "Write a complete HIP C++ source file: include , define the kernel, and write no main " "function. Reply with the source in a single ```cpp code block." ) -RESULT_TAG = "AORLTA_RESULT " SENTINEL = "-7777.0f" # what an output element holds before the kernel runs: one it never writes is a mismatch GUARD = 256 # elements past the end of every device array, which the kernel must leave alone GUARD_VALUE = "31337.0f" @@ -146,22 +147,41 @@ def _f(x) -> str: return f"{float(x)!r}f" -def harness_source(task: dict, kernel_source: str) -> str: - """The candidate's source followed by a host program that fills the inputs, runs the kernel once, and prints one - `AORLTA_RESULT {json}` line: `ok`, `mismatch` (an output element differs from the host reference beyond the - tolerance), `out_of_bounds` (a write past the end of an array) or `hip_error`.""" - lines = [kernel_source.rstrip(), "", "// ---- AORLTA harness ----", "#include ", "#include ", - "#include ", "#include ", "", +def harness_source(task: dict) -> str: + """The task's host program, which holds none of the candidate's code: it loads the candidate's device code object + (argv[1]), fills the inputs, runs the task's kernel from it once and compares every output element with a host + reference. It writes its record to the file argv[2], replacing it whole each time: `{"stage": "setup"}` at the + start, `{"stage": "launched"}` just before the kernel runs, then one with a `status`: `ok`, `mismatch` (an output + element differs from the host reference beyond the tolerance), `out_of_bounds` (a write past the end of an + array), `no_kernel` (the code object has no kernel of the task's exact signature) or `hip_error` (with the + `stage` it failed in).""" + lines = ["#include ", "#include ", "#include ", "#include ", + "#include ", "#include ", "", + "// The task's kernel, defined here only to name it (__builtin_get_device_side_mangled_name): the code", + "// that runs is the candidate's, from the code object in argv[1].", + f"{task['signature']} {{}}", "", + "static const char* g_path = nullptr;", + "static void put(const char* json) { // written aside, then renamed over the record", + " std::string tmp = std::string(g_path) + \".partial\";", + " FILE* f = std::fopen(tmp.c_str(), \"w\");", + " if (f == nullptr) return;", + " std::fputs(json, f);", + " if (std::fclose(f) == 0) std::rename(tmp.c_str(), g_path);", + "}", "", "static void fill(std::vector& v, uint32_t seed, float lo, float hi) {", " for (auto& x : v) { seed = seed * 1664525u + 1013904223u; " "x = lo + (hi - lo) * float(seed >> 8) / 16777216.0f; }", "}", "", - "#define CHECK(call) do { hipError_t e_ = (call); if (e_ != hipSuccess) { " - f"printf(\"{RESULT_TAG}{{\\\"status\\\": \\\"hip_error\\\", \\\"at\\\": \\\"%s\\\"}}\\n\", #call); " + "#define CHECK(stage, call) do { hipError_t e_ = (call); if (e_ != hipSuccess) { char b_[512]; " + "std::snprintf(b_, sizeof b_, \"{\\\"status\\\": \\\"hip_error\\\", \\\"stage\\\": \\\"%s\\\", " + "\\\"at\\\": \\\"%s\\\", \\\"error\\\": \\\"%s\\\"}\", stage, #call, hipGetErrorName(e_)); put(b_); " "return 0; } } while (0)", "", - "int main() {"] + "int main(int argc, char** argv) {", + " if (argc != 3) return 2;", + " g_path = argv[2];", + " put(\"{\\\"stage\\\": \\\"setup\\\"}\");"] for ctype, name, value in task["scalars"]: - lines.append(f" const {ctype} {name} = {value};") + lines.append(f" {ctype} {name} = {value};") arrays = task["arrays"] for i, a in enumerate(arrays): n = a["size"] @@ -179,18 +199,29 @@ def harness_source(task: dict, kernel_source: str) -> str: lines.append(" }") for a in arrays: lines += [f" float* d_{a['name']} = nullptr;", - f" CHECK(hipMalloc(&d_{a['name']}, h_{a['name']}.size() * sizeof(float)));", - f" CHECK(hipMemcpy(d_{a['name']}, h_{a['name']}.data(), h_{a['name']}.size() * sizeof(float), " - "hipMemcpyHostToDevice));"] + f" CHECK(\"setup\", hipMalloc(&d_{a['name']}, h_{a['name']}.size() * sizeof(float)));", + f" CHECK(\"setup\", hipMemcpy(d_{a['name']}, h_{a['name']}.data(), h_{a['name']}.size() * " + "sizeof(float), hipMemcpyHostToDevice));"] names = {a["name"] for a in arrays} - args = ", ".join(f"d_{x}" if x in names else x for x in task["args"]) - lines += [f" {task['kernel']}<<<{task['grid']}, {task['block']}>>>({args});", - " CHECK(hipGetLastError());", " CHECK(hipDeviceSynchronize());", + args = ", ".join(f"&d_{x}" if x in names else f"&{x}" for x in task["args"]) + lines += [" hipModule_t mod_;", " CHECK(\"setup\", hipModuleLoad(&mod_, argv[1]));", " hipFunction_t fn_;", + " {", + f" hipError_t e_ = hipModuleGetFunction(&fn_, mod_, __builtin_get_device_side_mangled_name(" + f"{task['kernel']}));", + " if (e_ == hipErrorNotFound || e_ == hipErrorInvalidDeviceFunction) {", + " put(\"{\\\"status\\\": \\\"no_kernel\\\", \\\"stage\\\": \\\"setup\\\"}\");", + " return 0;", " }", " CHECK(\"setup\", e_);", " }", + f" void* args_[] = {{{args}}};", + f" dim3 grid_ = {task['grid']};", f" dim3 block_ = {task['block']};", + " put(\"{\\\"stage\\\": \\\"launched\\\"}\");", + " CHECK(\"launch\", hipModuleLaunchKernel(fn_, grid_.x, grid_.y, grid_.z, block_.x, block_.y, block_.z, 0, " + "nullptr, args_, nullptr));", + " CHECK(\"run\", hipDeviceSynchronize());", " long checked = 0, mismatches = 0, first = -1, oob = 0;"] for a in arrays: n = a["size"] - lines += [f" CHECK(hipMemcpy(h_{a['name']}.data(), d_{a['name']}, h_{a['name']}.size() * sizeof(float), " - "hipMemcpyDeviceToHost));", + lines += [f" CHECK(\"run\", hipMemcpy(h_{a['name']}.data(), d_{a['name']}, h_{a['name']}.size() * " + "sizeof(float), hipMemcpyDeviceToHost));", f" for (long i = ({n}); i < (long)h_{a['name']}.size(); ++i) if (h_{a['name']}[i] != {GUARD_VALUE}) " "++oob;"] if a["role"] != "in": @@ -202,75 +233,131 @@ def harness_source(task: dict, kernel_source: str) -> str: " if (!same) { if (first < 0) first = checked - 1; ++mismatches; }", " }"] lines += [" const char* status = oob ? \"out_of_bounds\" : (mismatches ? \"mismatch\" : \"ok\");", - f" printf(\"{RESULT_TAG}{{\\\"status\\\": \\\"%s\\\", \\\"checked\\\": %ld, \\\"mismatches\\\": %ld, " - "\\\"first\\\": %ld, \\\"out_of_bounds\\\": %ld}\\n\", status, checked, mismatches, first, oob);", - " return 0;", "}", ""] + " char b_[512];", + " std::snprintf(b_, sizeof b_, \"{\\\"status\\\": \\\"%s\\\", \\\"checked\\\": %ld, " + "\\\"mismatches\\\": %ld, \\\"first\\\": %ld, \\\"out_of_bounds\\\": %ld}\", status, checked, " + "mismatches, first, oob);", + " put(b_);", " return 0;", "}", ""] return "\n".join(lines) -def read_result(stdout: str) -> dict | None: - """The harness's result line, or None.""" - for line in reversed(stdout.splitlines()): - if line.startswith(RESULT_TAG): - try: - out = json.loads(line[len(RESULT_TAG):]) - except ValueError: - return None - return out if isinstance(out, dict) and isinstance(out.get("status"), str) else None - return None - - -def verdict(build_rc, build_stderr, run, task) -> tuple[str, str | None, list[dict]]: - """A leg's (state, reason, findings) from the harness's build and run. `run` is None when the build failed, else - a ProcResult. Only what is the candidate's is blamed on it: a build failure is the candidate's when the compiler - names its kernel or `main` (its compile alone passed in gate.compile, so the signature is what differs), a hang - past the limit and a crash are its, and anything else is unsettled, never clean.""" - if build_rc != 0: - if re.search(r"\b(" + re.escape(task["kernel"]) + r"|main)\b", build_stderr or ""): - return "fail", "bad_signature", [] - return "unsettled", "harness_build_failed", [] +def read_record(path: str | Path) -> dict | None: + """The harness's record file, or None (absent, unreadable, or not a record).""" + try: + out = json.loads(Path(path).read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + return out if isinstance(out, dict) and isinstance(out.get("status", out.get("stage")), str) else None + + +def verdict(run, record) -> tuple[str, str | None, list[dict]]: + """A leg's (state, reason, findings) from the harness's run (a ProcResult) and its record (read_record). The + verdict comes only from the record the harness writes, never from the candidate's output, and only what is the + candidate's is blamed on it: a code object without the task's kernel (bad_signature), a wrong output, and a crash + or a hang once the kernel was launched. A failure before the launch (an allocation, the first copy, loading the + code object) is the node's, and leaves the leg unsettled, never clean.""" if run.launch_error or run.killed: return "unsettled", "run_not_started", [] - if run.timed_out: - return "fail", "timeout", [] - result = read_result(run.stdout) - if result is None: - return ("fail", "crash", []) if run.returncode not in (0, None) else ("unsettled", "no_result", []) - status = result["status"] - if status == "ok" and run.returncode == 0 and result.get("checked", 0) > 0: - return "clean", None, [] + record = record or {} + status, stage = record.get("status"), record.get("stage") + if status == "no_kernel": + return "fail", "bad_signature", [] + if status == "hip_error": + return ("unsettled", "setup_failed", []) if stage == "setup" else ("fail", "crash", []) if status in ("mismatch", "out_of_bounds"): return "fail", "wrong_output", [] - if status == "hip_error": - return "fail", "crash", [] + if status == "ok": + if run.returncode == 0 and not run.timed_out and isinstance(record.get("checked"), int) \ + and record["checked"] > 0: + return "clean", None, [] + return "unsettled", "no_result", [] + if stage == "launched": + if run.timed_out: + return "fail", "timeout", [] + if run.returncode not in (0, None): + return "fail", "crash", [] + if stage == "setup" or not record: + return "unsettled", ("setup_timeout" if run.timed_out else "setup_failed"), [] return "unsettled", "no_result", [] +_LEASE_LOCK = threading.Lock() +_LEASED: dict[int, int] = {} # GPU -> live hipk.run adapters on it, in this process + + +def lease_gpu(gpus: list[int], preferred: int) -> int: + """A GPU of `gpus` for a new worker: `preferred` if no live worker has it, else the first free one, else the + least shared. The service retires and replaces workers under new ids, so an id says nothing about which GPUs + the live workers hold.""" + with _LEASE_LOCK: + free = [g for g in gpus if not _LEASED.get(g)] + gpu = preferred if preferred in free else (free[0] if free else min(gpus, key=lambda g: _LEASED.get(g, 0))) + _LEASED[gpu] = _LEASED.get(gpu, 0) + 1 + return gpu + + +def release_gpu(gpu: int) -> None: + with _LEASE_LOCK: + n = _LEASED.get(gpu, 0) - 1 + if n > 0: + _LEASED[gpu] = n + else: + _LEASED.pop(gpu, None) + + class HipkRun: - """`hipk.run`, leg `ref.native`: the candidate built with its task's harness (hipcc, settings `hipcc`, leg setting - `target`), run on this worker's GPU (settings `gpus`: worker n runs on `gpus[(n - 1) % len(gpus)]`, by - HIP_VISIBLE_DEVICES) for at most `run_timeout_s`, and its outputs compared with the host reference. Clean only - when every output element matches and nothing was written past an array's end.""" + """`hipk.run`, leg `ref.native`: the candidate compiled alone into a device code object (hipcc --genco, settings + `hipcc`, leg setting `target`), so none of its host code is built or run; the task's harness, built from the task + alone, loads it and runs its kernel on this worker's GPU (settings `gpus`, leased to live workers: see + `lease_gpu`; by HIP_VISIBLE_DEVICES) for at most `run_timeout_s`, and compares its outputs with the host reference. + Clean only when every output element matches and nothing was written past an array's end. `version` names the + compiler hipcc runs and this file, so a leg cached under another compiler or harness is not reused.""" - version = "1" legs = ("ref.native",) resource = "own_gpu" def __init__(self, ctx): from aorlta.reward.adapters.base import ProcessRunner + from aorlta.reward.adapters.toolchain import toolchain_identity self.name, self.ctx = ctx.name, ctx s = ctx.settings self.hipcc = s.get("hipcc") or "/opt/rocm/bin/hipcc" - gpus = list(s.get("gpus") or [0]) - seq = int(ctx.worker_id.rsplit("/", 1)[-1]) if ctx.worker_id.rsplit("/", 1)[-1].isdigit() else 1 - self.gpu = int(gpus[(seq - 1) % len(gpus)]) + self.version = f"2+{toolchain_identity(self.hipcc)}+{hashlib.sha256(Path(__file__).read_bytes()).hexdigest()[:16]}" + gpus = [int(g) for g in (s.get("gpus") or [0])] + tail = ctx.worker_id.rsplit("/", 1)[-1] + seq = int(tail) if tail.isdigit() else 1 + self.gpu = lease_gpu(gpus, gpus[(seq - 1) % len(gpus)]) + self._leased = True self.tasks = {} for pool in ctx.pools: if (Path(pool) / "tasks.json").is_file(): for task in load_tasks(pool): self.tasks[task_sha256(task)] = task self.runner = ProcessRunner() + self._harness_dir: Path | None = None + self._harnesses: dict[tuple[str, str], Path] = {} + + def _harness(self, task: dict, target: str, timeout_s: float, evidence: Path) -> Path | None: + """The task's harness program for `target`, built once per worker (its source is also written to + `evidence`); None if it does not build.""" + src = harness_source(task) + (evidence / "harness.hip").write_text(src, encoding="utf-8") + key = (task_sha256(task), target) + if key in self._harnesses: + return self._harnesses[key] + if self._harness_dir is None: + self._harness_dir = Path(tempfile.mkdtemp(prefix="hipk-harness-")) + d = self._harness_dir / f"{key[0][:16]}-{target}" + d.mkdir(exist_ok=True) + (d / "harness.hip").write_text(src, encoding="utf-8") + build = self.runner.run([self.hipcc, f"--offload-arch={target}", "-O2", "-o", str(d / "harness"), + str(d / "harness.hip")], timeout_s=timeout_s, cwd=str(d)) + (evidence / "harness-build.log").write_text(build.stderr[-20000:], encoding="utf-8") + if build.returncode != 0 or build.timed_out or build.killed or build.launch_error: + return None + self._harnesses[key] = d / "harness" + return d / "harness" def run(self, job): from aorlta.reward.adapters.base import make_result, run_dir, unsettled @@ -282,26 +369,27 @@ def run(self, job): if job.leg != "ref.native" or task is None or not source or not job.output_dir: return unsettled(job, "adapter_error", t_start=t0, t_end=self.ctx.clock()) d = run_dir(job) - (d / "harness.hip").write_text(harness_source(task, Path(source).read_text(encoding="utf-8")), - encoding="utf-8") target = job.leg_settings.get("target", "gfx950") budget = max(1.0, self.ctx.timeout_s) t_build = time.monotonic() - build = self.runner.run([self.hipcc, f"--offload-arch={target}", "-O2", "-o", str(d / "harness"), - str(d / "harness.hip")], timeout_s=budget, cwd=str(d)) - if build.timed_out or build.killed or build.launch_error: + co = self.runner.run([self.hipcc, "--genco", f"--offload-arch={target}", "-O2", "-o", + str(d / "candidate.co"), str(source)], timeout_s=budget, cwd=str(d)) + (d / "build.log").write_text(co.stderr[-20000:], encoding="utf-8") + if co.returncode != 0 or co.timed_out or co.killed or co.launch_error: # it compiled in gate.compile + return unsettled(job, "candidate_build_failed", t_start=t0, t_end=self.ctx.clock()) + harness = self._harness(task, target, max(1.0, budget - (time.monotonic() - t_build)), d) + if harness is None: return unsettled(job, "harness_build_failed", t_start=t0, t_end=self.ctx.clock()) - (d / "build.log").write_text(build.stderr[-20000:], encoding="utf-8") - run = None - if build.returncode == 0: - left = budget - (time.monotonic() - t_build) - limit = min(float(job.leg_settings.get("run_timeout_s", 20.0)), left) - if limit <= 0: - return unsettled(job, "no_time_left", t_start=t0, t_end=self.ctx.clock()) - env = dict(os.environ, HIP_VISIBLE_DEVICES=str(self.gpu)) - run = self.runner.run([str(d / "harness")], timeout_s=limit, env=env, cwd=str(d)) - (d / "run.log").write_text((run.stdout + "\n" + run.stderr)[-20000:], encoding="utf-8") - state, reason, findings = verdict(build.returncode, build.stderr, run, task) + left = budget - (time.monotonic() - t_build) + limit = min(float(job.leg_settings.get("run_timeout_s", 20.0)), left) + if limit <= 0: + return unsettled(job, "no_time_left", t_start=t0, t_end=self.ctx.clock()) + env = dict(os.environ, HIP_VISIBLE_DEVICES=str(self.gpu)) + record = d / "record.json" + run = self.runner.run([str(harness), str(d / "candidate.co"), str(record)], timeout_s=limit, env=env, + cwd=str(d)) + (d / "run.log").write_text((run.stdout + "\n" + run.stderr)[-20000:], encoding="utf-8") + state, reason, findings = verdict(run, read_record(record)) return make_result(job, state, reason, findings=findings, 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()) @@ -314,6 +402,52 @@ def health(self): def close(self): self.runner.close() + if self._leased: + self._leased = False + release_gpu(self.gpu) + if self._harness_dir is not None: + shutil.rmtree(self._harness_dir, ignore_errors=True) + + +class HipkGate: + """`hipk.gate`, leg `gate.parse`: the reward's parse gate with the rollout's own rule (`hipk.parse`, which is + the stand-in `parse_code`): the last ```cpp block, and none when a later block never closes. The stand-in gate + takes the last complete block, so a final reply whose last block was cut off would be scored on an earlier one + the rollout had rejected. Writes `source.hip` as the stand-in gate does.""" + + name = "hipk.gate" + version = "1" + legs = ("gate.parse",) + resource = "local_cpu" + + def __init__(self, ctx=None): + self._clock = ctx.clock if ctx is not None else time.time + + def run(self, job): + from aorlta.atomic import atomic_write_bytes + from aorlta.reward.adapters.base import make_result, unsettled + from aorlta.rollout.compile_check import parse_code + + t0 = self._clock() + try: + response = job.inputs.get("response") or (str(Path(job.package) / "response.txt") if job.package else None) + if job.leg != "gate.parse" or not response or not job.output_dir: + return unsettled(job, "adapter_error", t_start=t0, t_end=self._clock()) + source = parse_code(Path(response).read_text(encoding="utf-8"), "off") + if source is None: + return make_result(job, "fail", "parse", t_start=t0, t_end=self._clock()) + out = Path(job.output_dir) + out.mkdir(parents=True, exist_ok=True) + atomic_write_bytes(out / "source.hip", source.encode("utf-8")) + return make_result(job, "clean", t_start=t0, t_end=self._clock()) + except Exception: # noqa: BLE001 - an adapter never raises + return unsettled(job, "adapter_error", t_start=t0, t_end=self._clock()) + + def controls(self): + return [] + + def health(self): + return {"name": self.name, "responsive": True} # --- the reward -------------------------------------------------------------------------------- @@ -393,6 +527,7 @@ def register() -> None: plugins.register_feedback_template("hipk.feedback", HipkFeedback) plugins.register_profile("hipk-gpu", HERE / "hipk-gpu.yaml") plugins.register_check_adapter("hipk.run", HipkRun) + plugins.register_check_adapter("hipk.gate", HipkGate) plugins.register_reward_fn("hipk.v1", HipkV1) @@ -404,36 +539,48 @@ def register() -> None: def selfcheck(argv=None) -> int: - """Each task's harness with its reference kernel (reference/.hip): built, and unless --build-only run on the - GPU, where every task must report `ok`. Prints one JSON line per task; exit 0 when all pass.""" + """Each task's reference kernel (reference/.hip), or with --sources DIR the kernel DIR/.hip where there + is one, through the check's own steps: compiled alone into a code object, the task's harness built, and unless + --build-only run on the GPU and judged by `verdict`. Prints one JSON line per task; exit 0 when all are clean + (built, with --build-only).""" + from aorlta.reward.adapters.base import ProcessRunner + ap = argparse.ArgumentParser(prog="hip_kernels.py selfcheck") ap.add_argument("--hipcc", default="/opt/rocm/bin/hipcc") ap.add_argument("--arch", default="gfx950") ap.add_argument("--gpu", type=int, default=0) ap.add_argument("--pool", default=str(HERE)) + ap.add_argument("--sources", default=None, help="a directory of .hip kernels to check instead") ap.add_argument("--build-only", action="store_true", help="compile the harnesses, run nothing (no GPU)") args = ap.parse_args(argv) - import subprocess - + runner = ProcessRunner() failed = 0 work = Path(tempfile.mkdtemp(prefix="hipk-selfcheck-")) try: for task in load_tasks(args.pool): - ref = (HERE / "reference" / f"{task['id']}.hip").read_text(encoding="utf-8") - src, exe = work / f"{task['id']}.hip", work / task["id"] - src.write_text(harness_source(task, ref), encoding="utf-8") - build = subprocess.run([args.hipcc, f"--offload-arch={args.arch}", "-O2", "-o", str(exe), str(src)], - capture_output=True, text=True) - row = {"task": task["id"], "build_rc": build.returncode} - if build.returncode == 0 and not args.build_only: - run = subprocess.run([str(exe)], capture_output=True, text=True, timeout=60, - env=dict(os.environ, HIP_VISIBLE_DEVICES=str(args.gpu))) - row["result"] = read_result(run.stdout) - ok = run.returncode == 0 and row["result"] is not None and row["result"]["status"] == "ok" + src = Path(args.sources or HERE / "reference") / f"{task['id']}.hip" + if not src.is_file(): + continue + d = work / task["id"] + d.mkdir() + (d / "harness.hip").write_text(harness_source(task), encoding="utf-8") + co = runner.run([args.hipcc, "--genco", f"--offload-arch={args.arch}", "-O2", "-o", str(d / "candidate.co"), + str(src)], timeout_s=300) + hb = runner.run([args.hipcc, f"--offload-arch={args.arch}", "-O2", "-o", str(d / "harness"), + str(d / "harness.hip")], timeout_s=300) + row = {"task": task["id"], "source": str(src), "build_rc": [co.returncode, hb.returncode]} + built = co.returncode == 0 and hb.returncode == 0 + if built and not args.build_only: + run = runner.run([str(d / "harness"), str(d / "candidate.co"), str(d / "record.json")], timeout_s=60, + env=dict(os.environ, HIP_VISIBLE_DEVICES=str(args.gpu))) + row["record"] = read_record(d / "record.json") + row["stdout"] = run.stdout[-500:] + row["verdict"] = verdict(run, row["record"])[:2] + ok = row["verdict"][0] == "clean" else: - ok = build.returncode == 0 + ok = built if not ok: - row["stderr"] = build.stderr[-2000:] + row["stderr"] = (co.stderr + hb.stderr)[-2000:] row["pass"] = ok failed += not ok print(json.dumps(row), flush=True) diff --git a/aorlta/examples/hip-kernels/hipk-gpu.yaml b/aorlta/examples/hip-kernels/hipk-gpu.yaml index 1a7d33f5b..5d49a121d 100644 --- a/aorlta/examples/hip-kernels/hipk-gpu.yaml +++ b/aorlta/examples/hip-kernels/hipk-gpu.yaml @@ -1,12 +1,13 @@ -# `hipk-gpu`: the example task's checks (examples/hip-kernels/hip_kernels.py). The stand-in parse and compile gates, -# then `ref.native`: the candidate built with its task's harness and run on one of the GPUs in `hipk.run.gpus`, its -# outputs compared with a host reference. Scored by hipk.v1. +# `hipk-gpu`: the example task's checks (examples/hip-kernels/hip_kernels.py). `hipk.gate` (the parse gate with the +# rollout's own rule) and the stand-in compile gate, then `ref.native`: the candidate compiled alone into a code object, +# run by its task's harness on one of the GPUs in `hipk.run.gpus`, its outputs compared with a host reference. Scored +# by hipk.v1. name: hipk-gpu reward_fn: hipk.v1 clean_repeats: 1 materialize: true pools: - gates: {resource: local_cpu, adapter: standin.parse, timeout: gate, workers: 4} + gates: {resource: local_cpu, adapter: hipk.gate, timeout: gate, workers: 4} compile: {resource: local_cpu, adapter: standin.hipcc, timeout: compile, workers: 8} run: {resource: own_gpu, adapter: hipk.run, timeout: sim_run, workers: 4} stages: @@ -19,8 +20,8 @@ requires: leg_settings: gate.compile: {offload_arch: gfx950, flags: ["-c", "-O1"]} ref.native: {target: gfx950, run_timeout_s: 20} -# `${AORLTA_ROCM_ROOT}` is that environment variable, or /opt/rocm. Worker n of the run pool uses -# gpus[(n - 1) % len(gpus)]: keep these off the engines' GPUs (rollout.engines). +# `${AORLTA_ROCM_ROOT}` is that environment variable, or /opt/rocm. Each live worker of the run pool holds one of +# `gpus` (a replaced worker's successor takes the one it gave back): keep these off the engines' GPUs (rollout.engines). adapters: standin.hipcc: {hipcc: "${AORLTA_ROCM_ROOT}/bin/hipcc", readelf: "${AORLTA_ROCM_ROOT}/llvm/bin/llvm-readelf"} hipk.run: {hipcc: "${AORLTA_ROCM_ROOT}/bin/hipcc", gpus: [4, 5, 6, 7]} diff --git a/aorlta/examples/qa/qa_task.py b/aorlta/examples/qa/qa_task.py index e7bc9f709..884f6f55e 100644 --- a/aorlta/examples/qa/qa_task.py +++ b/aorlta/examples/qa/qa_task.py @@ -22,7 +22,6 @@ import hashlib import json import re -import subprocess import unicodedata from pathlib import Path @@ -128,14 +127,23 @@ class QaMatch: `no_answer` when the reply has no `Answer:` line; unsettled when the task is unknown, the reply cannot be read, or the judge does not give a verdict.""" - version = "1" legs = ("ref.answer",) 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.judge = ctx.settings.get("judge") self.judge_timeout_s = float(ctx.settings.get("judge_timeout_s", 60.0)) + self.version = "1" # cached legs are keyed by it: the judge, which can turn a mismatch clean, is part of it + if self.judge: + files = {a: file_identity(a) for a in self.judge if isinstance(a, str) and Path(a).is_file()} + ident = json.dumps({"judge": list(self.judge), "timeout_s": self.judge_timeout_s, "files": files}, + sort_keys=True) + self.version = f"1+judge-{hashlib.sha256(ident.encode()).hexdigest()[:16]}" + self.runner = ProcessRunner() # kills the judge's whole process group on a timeout or at close() self.rows = {} for pool in ctx.pools: if (Path(pool) / "qa.jsonl").is_file(): @@ -169,12 +177,14 @@ def run(self, job): return unsettled(job, "adapter_error", t_start=t0, t_end=self.ctx.clock()) def _judge(self, row, answer) -> str | None: + proc = self.runner.run(list(self.judge), input=json.dumps({"question": row["question"], + "reference": row["answer"], "answer": answer}), + timeout_s=min(self.judge_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: - proc = subprocess.run(list(self.judge), input=json.dumps({"question": row["question"], - "reference": row["answer"], "answer": answer}), - capture_output=True, text=True, timeout=min(self.judge_timeout_s, self.ctx.timeout_s)) - out = json.loads(proc.stdout.strip().splitlines()[-1]) if proc.returncode == 0 and proc.stdout else {} - except (OSError, subprocess.TimeoutExpired, ValueError, IndexError): + 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 ("correct", "wrong") else None @@ -185,6 +195,9 @@ def controls(self): def health(self): return {"name": self.name, "responsive": True, "questions": len(self.rows)} + def close(self): + self.runner.close() + class QaV1: """`qa.v1`: +1.0 for a matching answer, -0.5 for a wrong one, -1.0 for no `Answer:` line; then -0.05 x the share diff --git a/aorlta/src/aorlta/reward/adapters/base.py b/aorlta/src/aorlta/reward/adapters/base.py index ea89c1d41..426980ec4 100644 --- a/aorlta/src/aorlta/reward/adapters/base.py +++ b/aorlta/src/aorlta/reward/adapters/base.py @@ -218,12 +218,14 @@ def __init__(self) -> None: self._closed = threading.Event() def run(self, argv: Sequence[str], *, timeout_s: float, env: Mapping[str, str] | None = None, - cwd: str | None = None) -> ProcResult: + cwd: str | None = None, input: str | None = None) -> ProcResult: + """`input`, if given, is written to the child's stdin (else stdin is not redirected).""" t0 = time.monotonic() if self._closed.is_set(): return ProcResult(None, "", "", 0.0, False, True, "runner closed") try: - proc = subprocess.Popen(list(argv), stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, + proc = subprocess.Popen(list(argv), stdin=subprocess.PIPE if input is not None else None, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, errors="replace", env=dict(env) if env is not None else None, cwd=cwd, start_new_session=True) except OSError as exc: @@ -234,7 +236,7 @@ def run(self, argv: Sequence[str], *, timeout_s: float, env: Mapping[str, str] | else: self._proc = proc try: - out, err = proc.communicate(timeout=timeout_s) + out, err = proc.communicate(input=input, timeout=timeout_s) timed_out = False except subprocess.TimeoutExpired: self._kill(proc) diff --git a/aorlta/src/aorlta/reward/evidence.py b/aorlta/src/aorlta/reward/evidence.py index 200a0319e..e97119aa4 100644 --- a/aorlta/src/aorlta/reward/evidence.py +++ b/aorlta/src/aorlta/reward/evidence.py @@ -1,7 +1,7 @@ """Keeping the GPU legs' evidence (run 2, step 4): the reward service's packages live in the run's node-local scratch, which the orchestrator removes at shutdown, so what explains an unsettled ConSan or ASan leg would go with it (run 1's -evidence archive came out empty). `archive_legs` copies the sanitizer legs' logs, generated harness sources and JSON, -and each candidate's parsed source, into one tarball in the run's sealed evidence store (Section 6.10 item 3) before +evidence archive came out empty). `archive_legs` copies the sanitizer and reference legs' logs, generated harness +sources and JSON, and each candidate's parsed source, into one tarball in the run's sealed evidence store (Section 6.10 item 3) before the scratch goes. Package layout: `pkg///-/...`. """ @@ -15,14 +15,14 @@ def kept(rel: Path) -> bool: - """Whether a file under the package root is evidence: a sanitizer leg's log, JSON or harness source, or the - source the parse gate wrote.""" + """Whether a file under the package root is evidence: a sanitizer leg's or a reference check's (`ref.*`, such as + a task plugin's own check) log, JSON or harness source, or the source the parse gate wrote.""" parts = rel.parts if len(parts) < 4: return False leg = parts[2] - return (leg.startswith("san.") and rel.suffix in SUFFIXES) or (leg.startswith("gate.parse-") - and rel.name == "source.hip") + return (leg.startswith(("san.", "ref.")) and rel.suffix in SUFFIXES) or (leg.startswith("gate.parse-") + and rel.name == "source.hip") def archive_legs(pkg_root: str | Path, dest: str | Path, *, max_file_bytes: int = MAX_FILE_BYTES) -> int: diff --git a/aorlta/tests/examples/test_hip_kernels.py b/aorlta/tests/examples/test_hip_kernels.py index 4624f47f7..49954e612 100644 --- a/aorlta/tests/examples/test_hip_kernels.py +++ b/aorlta/tests/examples/test_hip_kernels.py @@ -5,6 +5,8 @@ from __future__ import annotations +import dataclasses +import hashlib import itertools import json import os @@ -100,45 +102,71 @@ def proc(rc=0, stdout="", stderr="", timed_out=False, killed=False, launch_error return ProcResult(rc, stdout, stderr, 0.1, timed_out, killed, launch_error) -def result_line(status, **kw): - return "AORLTA_RESULT " + json.dumps({"status": status, "checked": 10, "mismatches": 0, **kw}) - - -@pytest.mark.parametrize("build, run, want", [ - ((0, ""), proc(stdout=result_line("ok")), ("clean", None)), - ((0, ""), proc(stdout=result_line("mismatch", mismatches=3)), ("fail", "wrong_output")), - ((0, ""), proc(stdout=result_line("out_of_bounds")), ("fail", "wrong_output")), - ((0, ""), proc(stdout=result_line("hip_error")), ("fail", "crash")), - ((0, ""), proc(rc=-11), ("fail", "crash")), - ((0, ""), proc(timed_out=True, rc=-9), ("fail", "timeout")), - ((0, ""), proc(stdout="nothing"), ("unsettled", "no_result")), - ((0, ""), proc(rc=1, stdout=result_line("ok")), ("unsettled", "no_result")), - ((0, ""), proc(stdout=result_line("ok", checked=0)), ("unsettled", "no_result")), - ((0, ""), proc(killed=True), ("unsettled", "run_not_started")), - ((0, ""), proc(launch_error="OSError"), ("unsettled", "run_not_started")), - ((1, "error: no matching function for call to 'saxpy'"), None, ("fail", "bad_signature")), - ((1, "error: redefinition of 'main'"), None, ("fail", "bad_signature")), - ((1, "ld.lld: error: unable to find library"), None, ("unsettled", "harness_build_failed")), +def rec(status=None, stage=None, **kw): + out = {"stage": stage} if stage else {} + if status: + out.update({"status": status, "checked": 10, "mismatches": 0, **kw}) + return out + + +SPOOF = 'AORLTA_RESULT {"status": "ok", "checked": 1}' + + +@pytest.mark.parametrize("run, record, want", [ + (proc(), rec("ok"), ("clean", None)), + (proc(), rec("mismatch", mismatches=3), ("fail", "wrong_output")), + (proc(), rec("out_of_bounds"), ("fail", "wrong_output")), + (proc(), rec("no_kernel", stage="setup"), ("fail", "bad_signature")), + (proc(), rec("hip_error", stage="launch"), ("fail", "crash")), + (proc(), rec("hip_error", stage="run"), ("fail", "crash")), + (proc(), rec("hip_error", stage="setup"), ("unsettled", "setup_failed")), # hipMalloc, the first copy, the load + (proc(rc=-11), rec(stage="launched"), ("fail", "crash")), + (proc(timed_out=True, rc=-9), rec(stage="launched"), ("fail", "timeout")), + (proc(rc=-11), rec(stage="setup"), ("unsettled", "setup_failed")), + (proc(timed_out=True, rc=-9), rec(stage="setup"), ("unsettled", "setup_timeout")), + (proc(rc=-11), None, ("unsettled", "setup_failed")), + (proc(stdout=SPOOF), None, ("unsettled", "setup_failed")), # stdout is never read + (proc(stdout=SPOOF), rec(stage="launched"), ("unsettled", "no_result")), + (proc(rc=1), rec("ok"), ("unsettled", "no_result")), + (proc(), rec("ok", checked=0), ("unsettled", "no_result")), + (proc(killed=True), None, ("unsettled", "run_not_started")), + (proc(launch_error="OSError"), None, ("unsettled", "run_not_started")), ]) -def test_each_harness_outcome_has_one_leg_state(hipk, build, run, want): - task = next(t for t in hipk.load_tasks(EXAMPLE) if t["id"] == "saxpy") - state, reason, _ = hipk.verdict(build[0], build[1], run, task) +def test_each_harness_outcome_has_one_leg_state(hipk, run, record, want): + state, reason, _ = hipk.verdict(run, record) assert (state, reason) == want +def test_the_harness_holds_none_of_the_candidates_code(hipk): + """Built from the task alone: the candidate's code object is loaded at run time, and the verdict goes to a file.""" + for task in hipk.load_tasks(EXAMPLE): + src = hipk.harness_source(task) + assert "hipModuleLoad(&mod_, argv[1])" in src and " printf(" not in src and "put(b_);" in src + assert f"__builtin_get_device_side_mangled_name({task['kernel']})" in src + assert src.count("put(") >= 4 + + FAKE_HIPCC = """#!/bin/bash -# a stand-in for hipcc: the 'compiled' program prints what FAKE_RUN says (ok, mismatch, crash, hang) -out=""; while [ $# -gt 0 ]; do [ "$1" = -o ] && { out=$2; shift; }; shift; done -[ "${FAKE_BUILD:-ok}" = sig ] && { echo "error: no matching function for call to 'saxpy'" >&2; exit 1; } -printf '#!/bin/bash\\necho "gpu $HIP_VISIBLE_DEVICES" >&2\\n' > "$out" +# a stand-in for hipcc: --genco writes a code object; otherwise the 'harness' writes the record FAKE_RUN says +out=""; genco=0; for a in "$@"; do [ "$a" = --genco ] && genco=1; done +while [ $# -gt 0 ]; do [ "$1" = -o ] && { out=$2; shift; }; shift; done +if [ "$genco" = 1 ]; then + [ "${FAKE_BUILD:-ok}" = co ] && { echo "error: cannot compile" >&2; exit 1; } + echo object > "$out"; exit 0 +fi +[ "${FAKE_BUILD:-ok}" = harness ] && { echo "error: harness" >&2; exit 1; } +printf '#!/bin/bash\necho "gpu $HIP_VISIBLE_DEVICES co $(cat $1)" >&2\necho %s\n' "'$SPOOF'" > "$out" case "${FAKE_RUN:-ok}" in - ok) echo 'echo "AORLTA_RESULT {\\"status\\": \\"ok\\", \\"checked\\": 5}"' >> "$out" ;; - mismatch) echo 'echo "AORLTA_RESULT {\\"status\\": \\"mismatch\\", \\"checked\\": 5, \\"mismatches\\": 2}"' >> "$out" ;; - crash) echo 'kill -SEGV $$' >> "$out" ;; - hang) echo 'sleep 30' >> "$out" ;; + ok) echo 'echo "{\\"status\\": \\"ok\\", \\"checked\\": 5}" > "$2"' >> "$out" ;; + mismatch) echo 'echo "{\\"status\\": \\"mismatch\\", \\"checked\\": 5, \\"mismatches\\": 2}" > "$2"' >> "$out" ;; + nokernel) echo 'echo "{\\"status\\": \\"no_kernel\\", \\"stage\\": \\"setup\\"}" > "$2"' >> "$out" ;; + setup) echo 'echo "{\\"status\\": \\"hip_error\\", \\"stage\\": \\"setup\\", \\"at\\": \\"hipMalloc\\"}" > "$2"' >> "$out" ;; + crash) echo 'echo "{\\"stage\\": \\"launched\\"}" > "$2"; kill -SEGV $$' >> "$out" ;; + hang) echo 'echo "{\\"stage\\": \\"launched\\"}" > "$2"; sleep 30' >> "$out" ;; + spoof) echo 'echo "{\\"stage\\": \\"launched\\"}" > "$2"' >> "$out" ;; esac chmod +x "$out" -""" +""".replace("$SPOOF", SPOOF.replace('"', '\\"')) @pytest.fixture @@ -149,6 +177,13 @@ def fake_hipcc(tmp_path) -> Path: return p +def make_adapter(hipk, fake_hipcc, worker="run/2", gpus=(4, 6), timeout=5.0): + ctx = plugins.CheckContext(name="hipk.run", worker_id=worker, settings={"hipcc": str(fake_hipcc), + "gpus": list(gpus)}, + clock=lambda: 0.0, timeout_s=timeout, pools=(str(EXAMPLE),)) + return plugins.CHECK_ADAPTERS["hipk.run"](ctx) + + def run_job(hipk, tmp_path, fake_hipcc, *, worker="run/2", task_id="saxpy", timeout=5.0, run_timeout_s=2.0): from aorlta.schemas import LegJob @@ -157,9 +192,7 @@ def run_job(hipk, tmp_path, fake_hipcc, *, worker="run/2", task_id="saxpy", time src.write_text("__global__ void saxpy(float a, const float* x, float* y, int n) {}\n") out = tmp_path / "out" out.mkdir(exist_ok=True) - ctx = plugins.CheckContext(name="hipk.run", worker_id=worker, settings={"hipcc": str(fake_hipcc), "gpus": [4, 6]}, - clock=lambda: 0.0, timeout_s=timeout, pools=(str(EXAMPLE),)) - adapter = plugins.CHECK_ADAPTERS["hipk.run"](ctx) + adapter = make_adapter(hipk, fake_hipcc, worker=worker, timeout=timeout) job = LegJob(job_id="j", leg="ref.native", candidate_id="c", control_id=None, source_sha256="s" * 64, task_sha256=hipk.task_sha256(task), leg_settings={"target": "gfx950", "run_timeout_s": run_timeout_s}, inputs={"source": str(src)}, output_dir=str(out)) @@ -168,17 +201,83 @@ def run_job(hipk, tmp_path, fake_hipcc, *, worker="run/2", task_id="saxpy", time @pytest.mark.parametrize("build, run, want", [ ("ok", "ok", ("clean", None)), ("ok", "mismatch", ("fail", "wrong_output")), ("ok", "crash", ("fail", "crash")), - ("ok", "hang", ("fail", "timeout")), ("sig", "ok", ("fail", "bad_signature")), + ("ok", "hang", ("fail", "timeout")), ("ok", "nokernel", ("fail", "bad_signature")), + ("ok", "setup", ("unsettled", "setup_failed")), ("ok", "spoof", ("unsettled", "no_result")), + ("co", "ok", ("unsettled", "candidate_build_failed")), ("harness", "ok", ("unsettled", "harness_build_failed")), ]) def test_the_check_runs_the_harness_on_its_workers_gpu(hipk, tmp_path, fake_hipcc, monkeypatch, build, run, want): monkeypatch.setenv("FAKE_BUILD", build) monkeypatch.setenv("FAKE_RUN", run) adapter, res = run_job(hipk, tmp_path, fake_hipcc) - assert (res.state, res.reason) == want and res.leg == "ref.native" - assert adapter.gpu == 6 # worker 2 of gpus [4, 6] - if build == "ok" and run == "ok": - (log,) = tmp_path.glob("out/*/run.log") - assert "gpu 6" in log.read_text() + try: + assert (res.state, res.reason) == want and res.leg == "ref.native" + assert adapter.gpu == 6 # worker 2 of gpus [4, 6], no other worker live + if build == "ok" and run == "ok": + (log,) = tmp_path.glob("out/*/run.log") + assert "gpu 6 co object" in log.read_text() + assert {p.name for p in log.parent.iterdir()} >= {"build.log", "harness.hip", "record.json", "run.log"} + finally: + adapter.close() + + +def test_a_replaced_worker_takes_the_gpu_its_predecessor_gave_back(hipk, fake_hipcc): + """The service names a replacement by the next sequence number (run/5 for run/2), so a GPU from the id alone + would land run/5 on run/1's GPU and leave run/2's idle.""" + probe = make_adapter(hipk, fake_hipcc, worker="run/probe", gpus=(4, 5, 6, 7)) + probe.close() + workers = {n: make_adapter(hipk, fake_hipcc, worker=f"run/{n}", gpus=(4, 5, 6, 7)) for n in (1, 2, 3, 4)} + assert [w.gpu for w in workers.values()] == [4, 5, 6, 7] + workers.pop(2).close() + workers[5] = make_adapter(hipk, fake_hipcc, worker="run/5", gpus=(4, 5, 6, 7)) + assert workers[5].gpu == 5 + workers[6] = make_adapter(hipk, fake_hipcc, worker="run/6", gpus=(4, 5, 6, 7)) # more workers than GPUs: shared + assert workers[6].gpu in (4, 5, 6, 7) + for w in workers.values(): + w.close() + workers[1].close() # a second close gives nothing back twice + assert hipk._LEASED == {} + + +def test_the_checker_version_names_the_compiler_and_this_file(hipk, fake_hipcc, tmp_path): + """Cached ref.native legs are keyed by the checker's version: another compiler, or an edited harness, is a new + version, so an old verdict is not reused.""" + a = make_adapter(hipk, fake_hipcc) + other = tmp_path / "other" / "hipcc" + other.parent.mkdir() + other.write_text(FAKE_HIPCC + "# another build\n") + other.chmod(0o755) + b = make_adapter(hipk, other) + again = make_adapter(hipk, fake_hipcc) + try: + assert a.version == again.version != b.version + assert a.version.endswith(hashlib.sha256(Path(hipk.__file__).read_bytes()).hexdigest()[:16]) + finally: + for x in (a, b, again): + x.close() + + +def test_the_reward_parse_gate_reads_the_answer_as_the_rollout_does(hipk, tmp_path): + """A final reply whose last ```cpp block never closes is no answer to the rollout; the stand-in gate would score + the earlier, complete block in its place.""" + from aorlta.reward.adapters.standin_parse import StandinParse + from aorlta.rollout.compile_check import ParseOutcome # noqa: F401 + from aorlta.schemas import LegJob + + cut = ("Here:\n```cpp\n__global__ void saxpy(float a, const float* x, float* y, int n) {}\n```\n" + "Better:\n```cpp\n__global__ void saxpy(float a, const") + whole = "```cpp\n__global__ void saxpy(float a, const float* x, float* y, int n) {}\n```\n" + gate = plugins.CHECK_ADAPTERS["hipk.gate"](plugins.CheckContext(name="hipk.gate", worker_id="gates/1", + settings={}, clock=lambda: 0.0, timeout_s=5.0)) + for i, (text, want) in enumerate([(cut, ("fail", "parse")), (whole, ("clean", None))]): + (tmp_path / f"r{i}.txt").write_text(text) + job = LegJob(job_id="j", leg="gate.parse", candidate_id="c", control_id=None, source_sha256="s" * 64, + task_sha256="t" * 64, inputs={"response": str(tmp_path / f"r{i}.txt")}, + output_dir=str(tmp_path / f"o{i}")) + res = gate.run(job) + assert (res.state, res.reason) == want + assert (tmp_path / f"o{i}" / "source.hip").exists() == (want[0] == "clean") + assert StandinParse().run(dataclasses.replace(job, output_dir=str(tmp_path / f"s{i}"))).state == "clean" + assert hipk.HipkParse().check(cut, "off", None).kind == "rejected" def test_a_task_the_check_does_not_know_is_unsettled(hipk, tmp_path, fake_hipcc): @@ -189,8 +288,11 @@ def test_a_task_the_check_does_not_know_is_unsettled(hipk, tmp_path, fake_hipcc) adapter = plugins.CHECK_ADAPTERS["hipk.run"](ctx) job = LegJob(job_id="j", leg="ref.native", candidate_id="c", control_id=None, source_sha256="s" * 64, task_sha256="0" * 64, inputs={"source": str(tmp_path / "x.hip")}, output_dir=str(tmp_path)) - res = adapter.run(job) - assert (res.state, res.reason) == ("unsettled", "adapter_error") + try: + res = adapter.run(job) + assert (res.state, res.reason) == ("unsettled", "adapter_error") + finally: + adapter.close() def leg(name, state, reason=None): @@ -280,7 +382,7 @@ def test_the_profile_names_only_what_the_plugin_and_the_package_provide(hipk): profile = load_profile("hipk-gpu") assert profile.reward_fn == "hipk.v1" and profile.legs == ("gate.parse", "gate.compile", "ref.native") - assert {p.adapter for p in profile.pools.values()} == {"standin.parse", "standin.hipcc", "hipk.run"} + assert {p.adapter for p in profile.pools.values()} == {"hipk.gate", "standin.hipcc", "hipk.run"} raw = parse_yaml((EXAMPLE / "hipk-gpu.yaml").read_bytes()) assert raw["adapters"]["hipk.run"]["gpus"] == [4, 5, 6, 7] @@ -335,7 +437,6 @@ def test_the_example_trains_on_cpu_through_the_real_wiring(tmp_path, hipk, tiny_ """The example's prompts, parse, feedback, profile and reward in the real loop, with test_wiring_e2e's tiny model and engines. The fake engine never writes the kernel asked for, so every chain gets hipk.parse's repair turn.""" from aorlta.orchestrator import wiring - from aorlta.reward.adapters.standin_parse import StandinParse from aorlta.rollout.engine_proc import InProcessEngine data = parse_yaml((ROOT / "configs" / "hipk-train.yaml").read_bytes()) @@ -349,7 +450,7 @@ def test_the_example_trains_on_cpu_through_the_real_wiring(tmp_path, hipk, tiny_ data["sync"] = {"dir": str(tmp_path / "shm")} cfg = config_from_dict(data) engines = [E.TwinEngine(tiny_base, E.FAKE, i) for i in range(2)] - factories = {"standin.parse": lambda w: StandinParse(), "standin.hipcc": lambda w: _CpuLeg("standin.hipcc", + factories = {"hipk.gate": lambda w: hipk.HipkGate(), "standin.hipcc": lambda w: _CpuLeg("standin.hipcc", "gate.compile"), "hipk.run": lambda w: _CpuLeg("hipk.run", "ref.native")} parts = wiring.Parts(tokenizer=E.ByteTokenizer(), handles=[InProcessEngine(i, e) for i, e in enumerate(engines)], diff --git a/aorlta/tests/examples/test_qa.py b/aorlta/tests/examples/test_qa.py index 8a15c6483..38eae9c9c 100644 --- a/aorlta/tests/examples/test_qa.py +++ b/aorlta/tests/examples/test_qa.py @@ -198,3 +198,46 @@ def test_the_qa_example_trains_on_cpu_through_the_real_wiring(tmp_path, qa, tiny assert all(len(e["turns"]) == 2 for e in eps) # the answer, and qa.parse's repair turn scored = [e for e in eps if e["status"] == "scored"] assert scored and all(e["legs"]["ref.answer"] == "fail" and -1.05 <= e["reward"] <= -1.0 for e in scored) + + +def test_the_judge_is_part_of_the_checkers_identity(qa, tmp_path): + """Cached ref.answer legs are keyed by the checker's version: turning the judge on, changing its command, its + timeout or its script is a new version, so a cached wrong_answer is not reused where the judge would decide.""" + judge = tmp_path / "judge" + judge.write_text("#!/bin/bash\necho '{\"verdict\": \"correct\"}'\n") + judge.chmod(0o755) + versions = [adapter(qa).version, adapter(qa, {"judge": [str(judge)]}).version, + adapter(qa, {"judge": [str(judge), "--strict"]}).version, + adapter(qa, {"judge": [str(judge)], "judge_timeout_s": 5}).version] + assert len(set(versions)) == 4 and versions[0] == "1" + assert adapter(qa, {"judge": [str(judge)]}).version == versions[1] + judge.write_text("#!/bin/bash\necho '{\"verdict\": \"wrong\"}'\n") + assert adapter(qa, {"judge": [str(judge)]}).version != versions[1] + + +def test_a_judge_that_times_out_or_is_retired_leaves_nothing_running(qa, tmp_path): + """The judge's whole process group is killed, on a timeout and when the worker is retired (close()), not only + the command itself: a wrapper's children went on running before.""" + import os + import threading + import time + + row = next(r for r in qa.load_rows(QA / "gpu-facts" / "train") if r["id"] == "wave-cdna") + for how in ("timeout", "close"): + pid = tmp_path / f"{how}.pid" + judge = tmp_path / f"judge-{how}" + judge.write_text(f"#!/bin/bash\nsleep 60 &\necho $! > {pid}\nwait\n") + judge.chmod(0o755) + a = adapter(qa, {"judge": [str(judge)], "judge_timeout_s": 1.0 if how == "timeout" else 60.0}) + if how == "close": + threading.Thread(target=lambda: (time.sleep(1.0), a.close()), daemon=True).start() + t0 = time.monotonic() + res = a.run(job_for(qa, tmp_path, row, "Answer: a wave of sixty-four")) + assert (res.state, res.reason) == ("unsettled", "judge_no_verdict") + assert time.monotonic() - t0 < 15, how # not the adapter's own 30 s, nor the child's 60 s + 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 + a.close() diff --git a/aorlta/tests/reward/test_evidence.py b/aorlta/tests/reward/test_evidence.py index 0ba995209..ef1a01071 100644 --- a/aorlta/tests/reward/test_evidence.py +++ b/aorlta/tests/reward/test_evidence.py @@ -23,6 +23,10 @@ "ab/abcd/san.asan-0123456789abcdef/asan/asan.log": True, "ab/abcd/san.waitcheck-0123456789abcdef/waitcheck-0.log": True, "ab/abcd/response.txt": False, + "ab/abcd/ref.native-0123456789abcdef/r1/run.log": True, # a task's own check, as the worked HIP example's + "ab/abcd/ref.native-0123456789abcdef/r1/harness.hip": True, + "ab/abcd/ref.native-0123456789abcdef/r1/record.json": True, + "ab/abcd/ref.native-0123456789abcdef/r1/candidate.co": False, } diff --git a/aorlta/tests/tools/test_eval_compare.py b/aorlta/tests/tools/test_eval_compare.py index 9436ed19b..8cf8aa0ef 100644 --- a/aorlta/tests/tools/test_eval_compare.py +++ b/aorlta/tests/tools/test_eval_compare.py @@ -64,3 +64,56 @@ def test_runs_that_sampled_differently_are_refused(tmp_path): def test_a_directory_that_is_not_a_score_run_is_refused(tmp_path): base = score_run(tmp_path, "base") assert X.main(["--base", str(base), "--trained", str(tmp_path), "--out", str(tmp_path / "out")]) == 2 + + +def test_runs_with_another_batch_layout_or_reward_are_refused(tmp_path, capsys): + """Seeds follow the iteration number, which a different P changes; a different reward scores another thing.""" + base = score_run(tmp_path, "base") + other = score_run(tmp_path, "p8") + _rewrite(other / "score" / "report.jsonl", + lambda r: dict(r, options=dict(r["options"], prompts_per_batch=8)) if r["kind"] == "run" else r) + assert X.main(["--base", str(base), "--trained", str(other), "--out", str(tmp_path / "out")]) == 2 + assert "differ in options" in capsys.readouterr().err + other = score_run(tmp_path, "rfn") + _rewrite(other / "score" / "report.jsonl", + lambda r: dict(r, reward=dict(r["reward"], reward_fn="other.v1")) if r["kind"] == "run" else r) + assert X.main(["--base", str(base), "--trained", str(other), "--out", str(tmp_path / "out")]) == 2 + assert "differ in reward" in capsys.readouterr().err + + +def _rewrite(report: Path, edit) -> None: + rows = [json.loads(x) for x in report.read_text().splitlines()] + report.write_text("".join(json.dumps(edit(r)) + "\n" for r in rows)) + + +def test_runs_whose_tasks_order_or_sample_counts_differ_are_refused(tmp_path, capsys): + base = score_run(tmp_path, "base") + first = next(json.loads(x)["prompt_id"] for x in (base / "score" / "report.jsonl").read_text().splitlines() + if json.loads(x)["kind"] == "prompt") + for name, edit, what in ( + ("task", lambda r: dict(r, task_sha256="0" * 64) if r.get("prompt_id") == first else r, "task_sha256"), + ("count", lambda r: dict(r, samples=r["samples"] - 1) if r.get("prompt_id") == first else r, "samples")): + other = score_run(tmp_path, name) + _rewrite(other / "score" / "report.jsonl", edit) + assert X.main(["--base", str(base), "--trained", str(other), "--out", str(tmp_path / "out")]) == 2 + assert what in capsys.readouterr().err + other = score_run(tmp_path, "order") + rep = other / "score" / "report.jsonl" + rows = [json.loads(x) for x in rep.read_text().splitlines()] + prompts = [r for r in rows if r["kind"] == "prompt"] + rows = [rows[0], *reversed(prompts), rows[-1]] + rep.write_text("".join(json.dumps(r) + "\n" for r in rows)) + assert X.main(["--base", str(base), "--trained", str(other), "--out", str(tmp_path / "out")]) == 2 + assert "same order" in capsys.readouterr().err + + +def test_a_sample_with_no_reply_is_shown_as_such(tmp_path): + """A chain that died before its first reply has no turns; the report shows it rather than crashing.""" + base, trained = score_run(tmp_path, "base"), score_run(tmp_path, "trained") + for f in trained.glob("iterations/iter-*/episodes.jsonl"): + eps = [json.loads(x) for x in f.read_text().splitlines()] + f.write_text("".join(json.dumps(dict(e, turns=[], final_turn=-1, end_reason="engine_error")) + "\n" + for e in eps)) + assert X.main(["--base", str(base), "--trained", str(trained), "--out", str(tmp_path / "out")]) == 0 + md = (tmp_path / "out" / "eval-compare.md").read_text() + assert "(no reply: the chain ended engine_error before its first turn)" in md diff --git a/aorlta/tools/eval_compare.py b/aorlta/tools/eval_compare.py index 00c0db1f3..1ad827706 100644 --- a/aorlta/tools/eval_compare.py +++ b/aorlta/tools/eval_compare.py @@ -9,8 +9,9 @@ It prints the pass rate (the share of samples whose reward claims clean: a correct answer), the mean reward and the scored fraction of each, and the per-prompt change, and writes `/eval-compare.json` and `/eval-compare.md`. The markdown holds, for the prompts that changed most, one sample of each model: its reward, its checks' states and, -with `--tokenizer`, its final reply decoded. It refuses two runs that did not score the same prompts with the same -sampling settings, since their numbers would not be comparable. An unscored sample is never counted as passing. +with `--tokenizer`, its final reply decoded. It refuses two runs that did not score the same prompts in the same +order, with the same sampling settings, batch layout, task hashes and reward, since their numbers would not be +comparable. An unscored sample is never counted as passing. """ from __future__ import annotations @@ -57,11 +58,18 @@ def headline(rep: dict) -> dict: def compare(base_dir: Path, trained_dir: Path) -> dict: base, trained = read_report(base_dir), read_report(trained_dir) - if set(base["prompts"]) != set(trained["prompts"]): - raise Refused("the two runs did not score the same prompts") - sb, st = base["run"].get("sampling"), trained["run"].get("sampling") - if sb != st: - raise Refused(f"the two runs sampled differently: {sb} against {st}") + if list(base["prompts"]) != list(trained["prompts"]): + raise Refused("the two runs did not score the same prompts in the same order") + for key in ("sampling", "options", "reward"): # options: P and the passes, which decide each prompt's seeds + sb, st = base["run"].get(key), trained["run"].get(key) + if sb != st: + raise Refused(f"the two runs differ in {key}: {sb} against {st}") + for pid, b in base["prompts"].items(): + t = trained["prompts"][pid] + for key in ("task_sha256", "samples"): + if b.get(key) != t.get(key): + raise Refused(f"{pid}: the two runs differ in {key}: {b.get(key)} against {t.get(key)}") + sb = base["run"].get("sampling") rows = [] for pid in sorted(base["prompts"]): b, t = base["prompts"][pid], trained["prompts"][pid] @@ -79,6 +87,8 @@ def compare(base_dir: Path, trained_dir: Path) -> dict: def sample_text(ep: dict, tok) -> str: + if not ep.get("turns"): # a chain that ended before its first reply (engine_error) is a legal episode + return f"(no reply: the chain ended {ep.get('end_reason')} before its first turn)" ids = ep["turns"][-1]["response_ids"] if tok is None: return f"({len(ids)} reply tokens; pass --tokenizer to decode)" @@ -105,7 +115,7 @@ def markdown(result: dict, base_dir: Path, trained_dir: Path, tok, samples: int) got = eps.get(r["prompt_id"]) or [] if not got: continue - e = got[0] + e = next((x for x in got if x.get("turns")), got[0]) # a sample with a reply, if there is one lines += ["", f"**{name}**: reward {fmt(e.get('reward'))}, {e['status']}, checks {e.get('legs')}", "", "```text", sample_text(e, tok)[:3000], "```"] lines.append("") From d673cbb550679165f3104fed7c308863b9016f92 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 09:53:13 +0000 Subject: [PATCH 12/19] fix(aorlta): the HIP harness's kernel lookup no longer reads its own uninitialised error CHECK declares its own e_, so CHECK("setup", e_) after the lookup expanded to hipError_t e_ = (e_). That is undefined, and at -O2 the compiler took the no_kernel branch for every lookup: on a GPU every reference kernel came back bad_signature. The lookup's result is now gf_, and the no_kernel record carries the HIP error name. selfcheck runs the check's own steps without importing aorlta, so it runs in a plain ROCm image. --- aorlta/examples/hip-kernels/hip_kernels.py | 42 ++++++++++++++-------- aorlta/tests/examples/test_hip_kernels.py | 2 ++ 2 files changed, 30 insertions(+), 14 deletions(-) diff --git a/aorlta/examples/hip-kernels/hip_kernels.py b/aorlta/examples/hip-kernels/hip_kernels.py index e7b1ad8d4..1f20871f7 100644 --- a/aorlta/examples/hip-kernels/hip_kernels.py +++ b/aorlta/examples/hip-kernels/hip_kernels.py @@ -206,11 +206,15 @@ def harness_source(task: dict) -> str: args = ", ".join(f"&d_{x}" if x in names else f"&{x}" for x in task["args"]) lines += [" hipModule_t mod_;", " CHECK(\"setup\", hipModuleLoad(&mod_, argv[1]));", " hipFunction_t fn_;", " {", - f" hipError_t e_ = hipModuleGetFunction(&fn_, mod_, __builtin_get_device_side_mangled_name(" + # gf_, not e_: CHECK declares its own e_, and `hipError_t e_ = (e_)` reads itself (undefined) + f" hipError_t gf_ = hipModuleGetFunction(&fn_, mod_, __builtin_get_device_side_mangled_name(" f"{task['kernel']}));", - " if (e_ == hipErrorNotFound || e_ == hipErrorInvalidDeviceFunction) {", - " put(\"{\\\"status\\\": \\\"no_kernel\\\", \\\"stage\\\": \\\"setup\\\"}\");", - " return 0;", " }", " CHECK(\"setup\", e_);", " }", + " if (gf_ == hipErrorNotFound || gf_ == hipErrorInvalidDeviceFunction) {", + " char b_[256];", + " std::snprintf(b_, sizeof b_, \"{\\\"status\\\": \\\"no_kernel\\\", \\\"stage\\\": \\\"setup\\\", " + "\\\"error\\\": \\\"%s\\\"}\", hipGetErrorName(gf_));", + " put(b_);", + " return 0;", " }", " CHECK(\"setup\", gf_);", " }", f" void* args_[] = {{{args}}};", f" dim3 grid_ = {task['grid']};", f" dim3 block_ = {task['block']};", " put(\"{\\\"stage\\\": \\\"launched\\\"}\");", @@ -543,7 +547,18 @@ def selfcheck(argv=None) -> int: is one, through the check's own steps: compiled alone into a code object, the task's harness built, and unless --build-only run on the GPU and judged by `verdict`. Prints one JSON line per task; exit 0 when all are clean (built, with --build-only).""" - from aorlta.reward.adapters.base import ProcessRunner + import subprocess + from types import SimpleNamespace + + def run(argv, timeout_s, env=None): # what `verdict` reads of a ProcResult; standalone, so no aorlta import + try: + p = subprocess.run(argv, capture_output=True, text=True, timeout=timeout_s, env=env) + return SimpleNamespace(returncode=p.returncode, stdout=p.stdout, stderr=p.stderr, timed_out=False, + killed=False, launch_error=None) + except subprocess.TimeoutExpired as exc: + out = exc.stdout.decode(errors="replace") if isinstance(exc.stdout, bytes) else (exc.stdout or "") + return SimpleNamespace(returncode=-9, stdout=out, stderr="", timed_out=True, killed=False, + launch_error=None) ap = argparse.ArgumentParser(prog="hip_kernels.py selfcheck") ap.add_argument("--hipcc", default="/opt/rocm/bin/hipcc") @@ -553,7 +568,6 @@ def selfcheck(argv=None) -> int: ap.add_argument("--sources", default=None, help="a directory of .hip kernels to check instead") ap.add_argument("--build-only", action="store_true", help="compile the harnesses, run nothing (no GPU)") args = ap.parse_args(argv) - runner = ProcessRunner() failed = 0 work = Path(tempfile.mkdtemp(prefix="hipk-selfcheck-")) try: @@ -564,18 +578,18 @@ def selfcheck(argv=None) -> int: d = work / task["id"] d.mkdir() (d / "harness.hip").write_text(harness_source(task), encoding="utf-8") - co = runner.run([args.hipcc, "--genco", f"--offload-arch={args.arch}", "-O2", "-o", str(d / "candidate.co"), - str(src)], timeout_s=300) - hb = runner.run([args.hipcc, f"--offload-arch={args.arch}", "-O2", "-o", str(d / "harness"), - str(d / "harness.hip")], timeout_s=300) + co = run([args.hipcc, "--genco", f"--offload-arch={args.arch}", "-O2", "-o", str(d / "candidate.co"), + str(src)], 300) + hb = run([args.hipcc, f"--offload-arch={args.arch}", "-O2", "-o", str(d / "harness"), + str(d / "harness.hip")], 300) row = {"task": task["id"], "source": str(src), "build_rc": [co.returncode, hb.returncode]} built = co.returncode == 0 and hb.returncode == 0 if built and not args.build_only: - run = runner.run([str(d / "harness"), str(d / "candidate.co"), str(d / "record.json")], timeout_s=60, - env=dict(os.environ, HIP_VISIBLE_DEVICES=str(args.gpu))) + res = run([str(d / "harness"), str(d / "candidate.co"), str(d / "record.json")], 60, + env=dict(os.environ, HIP_VISIBLE_DEVICES=str(args.gpu))) row["record"] = read_record(d / "record.json") - row["stdout"] = run.stdout[-500:] - row["verdict"] = verdict(run, row["record"])[:2] + row["stdout"] = res.stdout[-500:] + row["verdict"] = verdict(res, row["record"])[:2] ok = row["verdict"][0] == "clean" else: ok = built diff --git a/aorlta/tests/examples/test_hip_kernels.py b/aorlta/tests/examples/test_hip_kernels.py index 49954e612..57ccf64a8 100644 --- a/aorlta/tests/examples/test_hip_kernels.py +++ b/aorlta/tests/examples/test_hip_kernels.py @@ -144,6 +144,8 @@ def test_the_harness_holds_none_of_the_candidates_code(hipk): assert "hipModuleLoad(&mod_, argv[1])" in src and " printf(" not in src and "put(b_);" in src assert f"__builtin_get_device_side_mangled_name({task['kernel']})" in src assert src.count("put(") >= 4 + assert ", e_)" not in src # CHECK declares e_ itself: `hipError_t e_ = (e_)` is undefined, and at -O2 it + # made every lookup take the no_kernel branch (found on a GPU) FAKE_HIPCC = """#!/bin/bash From 6ba612a691c448bf7c55eadbb630f9c603db3191 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 09:58:45 +0000 Subject: [PATCH 13/19] fix(aorlta): a failed plugin's registering helper imports run again on a retry Undoing a failed plugin's registrations left a helper module it imported (one that registers on import) cached in sys.modules, so the fixed plugin's retry did not run it again and its components stayed missing. While a plugin loads, each registration records the module that made it; on failure, the modules first imported by that load that registered something are dropped from sys.modules. A module imported before the load is kept. --- aorlta/src/aorlta/plugins.py | 14 ++++++++++++++ aorlta/tests/test_plugins.py | 25 +++++++++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/aorlta/src/aorlta/plugins.py b/aorlta/src/aorlta/plugins.py index ff5cfd993..b218bdb7a 100644 --- a/aorlta/src/aorlta/plugins.py +++ b/aorlta/src/aorlta/plugins.py @@ -88,6 +88,7 @@ class CheckContext: CHECK_ADAPTERS: dict[str, Callable[[CheckContext], Any]] = {} LOADED: dict[str, dict] = {} # plugin as listed -> {"module", "file", "sha256"} _TABLES = (TASKS, PARSE_ADAPTERS, FEEDBACK_TEMPLATES, COMPILE_ADAPTERS, PROFILES, REWARD_FNS, CHECK_ADAPTERS) +_OWNERS: set[str] | None = None # while a plugin loads: the modules that registered something _LOCK = threading.RLock() _BUILT_IN = { @@ -111,6 +112,12 @@ def _register(kind: str, table: dict, name: str, value: Any) -> None: if name in table and table[name] is not value and table[name] != value: raise PluginError(f"{kind} {name!r} is already registered") table[name] = value + if _OWNERS is not None: # inside load_plugins: which module registered it, for a rollback + f = sys._getframe(1) + while f is not None and f.f_globals.get("__name__") == __name__: + f = f.f_back + if f is not None: + _OWNERS.add(f.f_globals.get("__name__")) def register_task(task: Task) -> None: @@ -169,14 +176,21 @@ def load_plugins(plugins: Sequence[str]) -> None: with _LOCK: if entry in LOADED: continue + global _OWNERS before = {id(t): set(t) for t in _TABLES} + modules = set(sys.modules) + _OWNERS = set() try: LOADED[entry] = _load(entry) except BaseException: for t in _TABLES: # what this import registered before failing is undone, nothing else for name in set(t) - before[id(t)]: del t[name] + for owner in _OWNERS - modules: # and a module it imported that registered runs again on a retry + sys.modules.pop(owner, None) raise + finally: + _OWNERS = None def _load(entry: str) -> dict: diff --git a/aorlta/tests/test_plugins.py b/aorlta/tests/test_plugins.py index 85311fff4..246729c3c 100644 --- a/aorlta/tests/test_plugins.py +++ b/aorlta/tests/test_plugins.py @@ -296,6 +296,31 @@ def test_a_plugin_that_raises_leaves_nothing_registered(tmp_path, registry, monk assert plugins.registered("half.parse") and plugins.registered("half.v1") and entry in plugins.LOADED +def test_a_helper_that_registered_runs_again_when_its_failed_plugin_is_retried(tmp_path, registry, monkeypatch): + """A plugin that imports a helper (which registers on import) and then raises: the helper's registrations are + undone, and so is its import, so the fixed plugin's retry imports it again and its components come back. A module + that was imported before the plugin is left alone.""" + import sys + + (tmp_path / "half_helper.py").write_text( + "from aorlta import plugins\nplugins.register_parse_adapter('helper.parse', lambda: None)\n") + (tmp_path / "pre_helper.py").write_text("from aorlta import plugins\n\n\ndef setup():\n" + " plugins.register_feedback_template('pre.feedback', lambda: None)\n") + (tmp_path / "half_plugin2.py").write_text("import half_helper # noqa: F401\nimport pre_helper\n" + "pre_helper.setup()\nraise RuntimeError('boom')\n") + monkeypatch.syspath_prepend(str(tmp_path)) + import pre_helper # noqa: F401 - imported before the plugin: it registers during the load, and stays imported + with pytest.raises(plugins.PluginError, match="RuntimeError: boom"): + plugins.load_plugins(["half_plugin2"]) + assert not plugins.registered("helper.parse") and not plugins.registered("pre.feedback") + assert "half_helper" not in sys.modules and "pre_helper" in sys.modules + (tmp_path / "half_plugin2.py").write_text("import half_helper # noqa: F401\n") + plugins.load_plugins(["half_plugin2"]) + assert plugins.registered("helper.parse") + for m in ("half_helper", "half_plugin2", "pre_helper"): + sys.modules.pop(m, None) + + def test_a_registered_profile_must_carry_its_registered_name(tmp_path, plugin): """The candidates carry the registered name and the reward service serves the profile file's: a file naming anything else (here a built-in profile's name) would have every candidate refused as profile_mismatch.""" From adde2b4025a4baa8ecba2a9071a0b96edb43ad33 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 10:16:37 +0000 Subject: [PATCH 14/19] fix(aorlta): eval_compare also refuses runs whose prompts, seeds, reward code or plugin files differ Names can agree while what they name changed: a Q&A template, a reward function's code, a plugin's check. The paired samples' first-prompt token hash and seed, the reward version and code hash of every scored sample, and the plugin files the run plans hashed must all agree, from what the runs already record. --- aorlta/tests/tools/test_eval_compare.py | 30 ++++++++++++++++++++++ aorlta/tools/eval_compare.py | 33 +++++++++++++++++++++++-- 2 files changed, 61 insertions(+), 2 deletions(-) diff --git a/aorlta/tests/tools/test_eval_compare.py b/aorlta/tests/tools/test_eval_compare.py index 8cf8aa0ef..8d88fd81b 100644 --- a/aorlta/tests/tools/test_eval_compare.py +++ b/aorlta/tests/tools/test_eval_compare.py @@ -117,3 +117,33 @@ def test_a_sample_with_no_reply_is_shown_as_such(tmp_path): assert X.main(["--base", str(base), "--trained", str(trained), "--out", str(tmp_path / "out")]) == 0 md = (tmp_path / "out" / "eval-compare.md").read_text() assert "(no reply: the chain ended engine_error before its first turn)" in md + + +def _edit_episodes(run_dir: Path, edit) -> None: + for f in run_dir.glob("iterations/iter-*/episodes.jsonl"): + eps = [json.loads(x) for x in f.read_text().splitlines()] + f.write_text("".join(json.dumps(edit(e)) + "\n" for e in eps)) + + +def test_runs_whose_prompts_seeds_reward_code_or_plugins_differ_are_refused(tmp_path, capsys): + """Names can agree while what they name changed: a Q&A template (the first prompt's token hash), the seeds, the + reward function's code, or a plugin's file.""" + base = score_run(tmp_path, "base") + + def first_turn(**kw): + return lambda e: dict(e, turns=[dict(e["turns"][0], **kw), *e["turns"][1:]]) if e["turns"] else e + + for name, edit, what in (("prompt", first_turn(prompt_ids_sha256="0" * 64), "first prompts or seeds"), + ("seed", first_turn(seed=7), "first prompts or seeds"), + ("code", lambda e: dict(e, reward_code_sha256="0" * 64) if e["status"] == "scored" + else e, "different reward code")): + other = score_run(tmp_path, name) + _edit_episodes(other, edit) + assert X.main(["--base", str(base), "--trained", str(other), "--out", str(tmp_path / "out")]) == 2, name + assert what in capsys.readouterr().err + other = score_run(tmp_path, "plugin") + plan = json.loads((other / "run_plan.json").read_text()) + plan["plugins"] = {"my_plugin.py": {"module": "m", "file": "f", "sha256": "1" * 64}} + (other / "run_plan.json").write_text(json.dumps(plan)) + assert X.main(["--base", str(base), "--trained", str(other), "--out", str(tmp_path / "out")]) == 2 + assert "different plugin files" in capsys.readouterr().err diff --git a/aorlta/tools/eval_compare.py b/aorlta/tools/eval_compare.py index 1ad827706..ea2349782 100644 --- a/aorlta/tools/eval_compare.py +++ b/aorlta/tools/eval_compare.py @@ -10,8 +10,8 @@ scored fraction of each, and the per-prompt change, and writes `/eval-compare.json` and `/eval-compare.md`. The markdown holds, for the prompts that changed most, one sample of each model: its reward, its checks' states and, with `--tokenizer`, its final reply decoded. It refuses two runs that did not score the same prompts in the same -order, with the same sampling settings, batch layout, task hashes and reward, since their numbers would not be -comparable. An unscored sample is never counted as passing. +order, with the same sampling settings, batch layout, task hashes, first prompts and seeds, reward code and plugin +files, since their numbers would not be comparable. An unscored sample is never counted as passing. """ from __future__ import annotations @@ -48,6 +48,34 @@ def episodes(run_dir: Path) -> dict[str, list[dict]]: return out +def fingerprints(base_dir: Path, trained_dir: Path) -> None: + """Refuse two runs whose names agree but whose inputs or scoring differ: the paired samples' first prompt (its + token hash, so a changed template shows) and seed, the reward function's version and code hash on every scored + sample, and the plugin files the run plans hashed (a task's own checks and reward live there).""" + eb, et = episodes(base_dir), episodes(trained_dir) + for pid in sorted(set(eb) & set(et)): + first = {(e["k"], e["iteration"]): (e["turns"][0]["prompt_ids_sha256"], e["turns"][0]["seed"]) + for e in eb[pid] if e.get("turns")} + for e in et[pid]: + got = first.get((e["k"], e["iteration"])) + if e.get("turns") and got is not None and got != (e["turns"][0]["prompt_ids_sha256"], + e["turns"][0]["seed"]): + raise Refused(f"{pid} sample {e['k']}: the two runs' first prompts or seeds differ") + rb, rt = ({(e.get("reward_version"), e.get("reward_code_sha256")) for eps in run.values() for e in eps + if e.get("status") == "scored"} for run in (eb, et)) + if rb and rt and rb != rt: + raise Refused(f"the two runs were scored by different reward code: {sorted(rb)} against {sorted(rt)}") + pb, pt = (_plugin_hashes(d) for d in (base_dir, trained_dir)) + if pb != pt: + raise Refused(f"the two runs loaded different plugin files: {pb} against {pt}") + + +def _plugin_hashes(run_dir: Path) -> dict: + path = Path(run_dir) / "run_plan.json" + plan = json.loads(path.read_text(encoding="utf-8")) if path.is_file() else {} + return {k: (v or {}).get("sha256") for k, v in sorted((plan.get("plugins") or {}).items())} + + def headline(rep: dict) -> dict: t = rep["total"] return {"samples": t["samples"], "scored": t["scored"], "scored_fraction": t["scored_fraction"], @@ -69,6 +97,7 @@ def compare(base_dir: Path, trained_dir: Path) -> dict: for key in ("task_sha256", "samples"): if b.get(key) != t.get(key): raise Refused(f"{pid}: the two runs differ in {key}: {b.get(key)} against {t.get(key)}") + fingerprints(base_dir, trained_dir) sb = base["run"].get("sampling") rows = [] for pid in sorted(base["prompts"]): From b17ccab25d17b554a44f1bf9f76b2557ac9c2a90 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 10:20:15 +0000 Subject: [PATCH 15/19] fix(aorlta): a Q&A judge named by a bare command is fingerprinted as PATH finds it The judge's identity hashed only its arguments that were files under the working directory, so a judge given as a bare name (resolved through PATH) could be replaced, or another installation selected, without a new checker version. The executable PATH selects is now resolved and its path and content are part of the version (raised by review on the vertical starters, whose ISA tool had the same code). --- aorlta/examples/qa/qa_task.py | 5 +++++ aorlta/tests/examples/test_qa.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/aorlta/examples/qa/qa_task.py b/aorlta/examples/qa/qa_task.py index 884f6f55e..f0a4e3e00 100644 --- a/aorlta/examples/qa/qa_task.py +++ b/aorlta/examples/qa/qa_task.py @@ -21,7 +21,9 @@ import hashlib import json +import os import re +import shutil import unicodedata from pathlib import Path @@ -140,6 +142,9 @@ def __init__(self, ctx): self.version = "1" # cached legs are keyed by it: the judge, which can turn a mismatch clean, is part of it if self.judge: files = {a: file_identity(a) for a in self.judge if isinstance(a, str) and Path(a).is_file()} + exe = shutil.which(self.judge[0]) if isinstance(self.judge[0], str) else None # as PATH finds it, if bare + if exe: + files[f"exe:{self.judge[0]}"] = f"{os.path.realpath(exe)}:{file_identity(exe)}" ident = json.dumps({"judge": list(self.judge), "timeout_s": self.judge_timeout_s, "files": files}, sort_keys=True) self.version = f"1+judge-{hashlib.sha256(ident.encode()).hexdigest()[:16]}" diff --git a/aorlta/tests/examples/test_qa.py b/aorlta/tests/examples/test_qa.py index 38eae9c9c..5ae920e0b 100644 --- a/aorlta/tests/examples/test_qa.py +++ b/aorlta/tests/examples/test_qa.py @@ -241,3 +241,20 @@ def test_a_judge_that_times_out_or_is_retired_leaves_nothing_running(qa, tmp_pat time.sleep(0.05) assert not os.path.exists(f"/proc/{child}"), how a.close() + + +def test_a_judge_found_through_path_is_part_of_the_checkers_identity(qa, 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.""" + bins = [tmp_path / "a", tmp_path / "b"] + for d in bins: + d.mkdir() + (d / "my-judge").write_text("#!/bin/bash\necho '{\"verdict\": \"correct\"}'\n") + (d / "my-judge").chmod(0o755) + monkeypatch.setenv("PATH", f"{bins[0]}:/usr/bin:/bin") + first = adapter(qa, {"judge": ["my-judge"]}).version + assert adapter(qa, {"judge": ["my-judge"]}).version == first + (bins[0] / "my-judge").write_text("#!/bin/bash\necho '{\"verdict\": \"wrong\"}'\n") + second = adapter(qa, {"judge": ["my-judge"]}).version + monkeypatch.setenv("PATH", f"{bins[1]}:/usr/bin:/bin") + assert len({first, second, adapter(qa, {"judge": ["my-judge"]}).version}) == 3 From 3748c97ba46868146cb747a3870d7ea121c8617e Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 14:34:04 +0000 Subject: [PATCH 16/19] fix(aorlta): eval_compare refuses a run with no run plan or episode records that do not cover its report A missing run_plan.json or missing iterations/ used to make the plugin and fingerprint checks compare empty collections and pass. Now both runs need a run plan and committed episode records (iter-/, never iter-.partial/) with exactly one record per sample their report counts, and the two runs' (iteration, k) pairs must match. A record with no turns stays legal. --- aorlta/tests/tools/test_eval_compare.py | 61 +++++++++++++++++++++++++ aorlta/tools/eval_compare.py | 57 +++++++++++++++++++---- 2 files changed, 110 insertions(+), 8 deletions(-) diff --git a/aorlta/tests/tools/test_eval_compare.py b/aorlta/tests/tools/test_eval_compare.py index 8d88fd81b..2a8f14653 100644 --- a/aorlta/tests/tools/test_eval_compare.py +++ b/aorlta/tests/tools/test_eval_compare.py @@ -119,6 +119,67 @@ def test_a_sample_with_no_reply_is_shown_as_such(tmp_path): assert "(no reply: the chain ended engine_error before its first turn)" in md +def _episode_file(run_dir: Path) -> Path: + return next(run_dir.glob("iterations/iter-*/episodes.jsonl")) + + +def _rewrite_episodes(run_dir: Path, edit) -> None: + f = _episode_file(run_dir) + eps = edit([json.loads(x) for x in f.read_text().splitlines()]) + f.write_text("".join(json.dumps(e) + "\n" for e in eps)) + + +def test_runs_without_the_evidence_the_checks_read_are_refused(tmp_path, capsys): + """A missing run plan, or episode records that do not cover the report one to one, would leave the fingerprint + checks nothing to compare; such a run is refused rather than passed.""" + base = score_run(tmp_path, "base") + + def drop_samples(eps): # three of the first prompt's four samples gone + pid = eps[0]["prompt_id"] + return [e for e in eps if e["prompt_id"] != pid or e["k"] == 0] + + def duplicate(eps): # one sample recorded twice in place of another, so the count still agrees + twin = next(e for e in eps[1:] if e["prompt_id"] == eps[0]["prompt_id"]) + return [e for e in eps if e is not twin] + [eps[0]] + + def renumber(eps): # one sample's k moved, so that sample has no partner in the other run + return [dict(eps[0], k=99), *eps[1:]] + + def unlisted(eps): + return [*eps, dict(eps[0], prompt_id="not-in-the-report")] + + for name, edit, what in (("dropped", drop_samples, "committed episode records; its report counts"), + ("duplicate", duplicate, "more than one episode record"), + ("unpaired", renumber, "do not pair"), + ("unlisted", unlisted, "does not list")): + other = score_run(tmp_path, name) + _rewrite_episodes(other, edit) + assert X.main(["--base", str(base), "--trained", str(other), "--out", str(tmp_path / "out")]) == 2, name + assert what in capsys.readouterr().err, name + assert X.main(["--base", str(other), "--trained", str(base), "--out", str(tmp_path / "out")]) == 2, name + assert what in capsys.readouterr().err, name + + other = score_run(tmp_path, "plan") + (other / "run_plan.json").unlink() + assert X.main(["--base", str(base), "--trained", str(other), "--out", str(tmp_path / "out")]) == 2 + assert "no run_plan.json" in capsys.readouterr().err + + other = score_run(tmp_path, "partial") # the batch never committed: only iter-.partial holds its records + d = _episode_file(other).parent + d.rename(d.with_name(d.name + ".partial")) + assert X.main(["--base", str(base), "--trained", str(other), "--out", str(tmp_path / "out")]) == 2 + assert "no committed episode records" in capsys.readouterr().err + assert not (tmp_path / "out" / "eval-compare.json").exists() + + other = score_run(tmp_path, "leftover") # a leftover .partial beside the committed batch is not read at all + f = _episode_file(other) + stray = f.parent.with_name(f.parent.name + ".partial") + stray.mkdir() + (stray / f.name).write_text("".join(json.dumps(dict(json.loads(x), k=99)) + "\n" + for x in f.read_text().splitlines())) + assert X.main(["--base", str(base), "--trained", str(other), "--out", str(tmp_path / "out")]) == 0 + + def _edit_episodes(run_dir: Path, edit) -> None: for f in run_dir.glob("iterations/iter-*/episodes.jsonl"): eps = [json.loads(x) for x in f.read_text().splitlines()] diff --git a/aorlta/tools/eval_compare.py b/aorlta/tools/eval_compare.py index ea2349782..2bcc18b1b 100644 --- a/aorlta/tools/eval_compare.py +++ b/aorlta/tools/eval_compare.py @@ -11,13 +11,16 @@ The markdown holds, for the prompts that changed most, one sample of each model: its reward, its checks' states and, with `--tokenizer`, its final reply decoded. It refuses two runs that did not score the same prompts in the same order, with the same sampling settings, batch layout, task hashes, first prompts and seeds, reward code and plugin -files, since their numbers would not be comparable. An unscored sample is never counted as passing. +files, since their numbers would not be comparable, and a run whose run plan is missing or whose committed episode +records do not cover every sample of its report, since those checks would have nothing to read. An unscored sample is +never counted as passing. """ from __future__ import annotations import argparse import json +import re import sys from pathlib import Path @@ -39,21 +42,58 @@ def read_report(run_dir: Path) -> dict: return {"run": run, "total": total, "prompts": prompts} +COMMITTED = re.compile(r"iter-\d+") # a batch is committed when iterations/iter-/ appears; iter-.partial is not + + def episodes(run_dir: Path) -> dict[str, list[dict]]: + """The committed episode records, by prompt.""" out: dict[str, list[dict]] = {} for f in sorted((Path(run_dir) / "iterations").glob("iter-*/episodes.jsonl")): + if not COMMITTED.fullmatch(f.parent.name): + continue for line in f.read_text(encoding="utf-8").splitlines(): e = json.loads(line) out.setdefault(e["prompt_id"], []).append(e) return out -def fingerprints(base_dir: Path, trained_dir: Path) -> None: +def paired_episodes(base_dir: Path, trained_dir: Path, base: dict, trained: dict) -> tuple[dict, dict]: + """Both runs' committed episodes, refused unless each report's every sample has exactly one record (its + iteration and k) and the two runs' records pair one to one. A record with no turns (a chain that ended before its + first reply) is legal.""" + runs = [] + for run_dir, rep in ((base_dir, base), (trained_dir, trained)): + eps = episodes(run_dir) + if not eps: + raise Refused(f"{run_dir}: no committed episode records (iterations/iter-/episodes.jsonl)") + if extra := sorted(set(eps) - set(rep["prompts"])): + raise Refused(f"{run_dir}: episode records for prompts its report does not list: {extra[:3]}") + keys = {} + for pid, row in rep["prompts"].items(): + got = [(e["iteration"], e["k"]) for e in eps.get(pid, [])] + if len(got) != len(set(got)): + raise Refused(f"{run_dir}: {pid} has more than one episode record for a sample") + if len(got) != row.get("samples"): + raise Refused(f"{run_dir}: {pid} has {len(got)} committed episode records; its report counts " + f"{row.get('samples')} samples") + keys[pid] = set(got) + runs.append((eps, keys)) + (eb, kb), (et, kt) = runs + for pid in base["prompts"]: + if kb[pid] != kt[pid]: + raise Refused(f"{pid}: the two runs' samples do not pair (iteration, k): " + f"{sorted(kb[pid] ^ kt[pid])[:3]} are in one run only") + return eb, et + + +def fingerprints(base_dir: Path, trained_dir: Path, base: dict, trained: dict) -> None: """Refuse two runs whose names agree but whose inputs or scoring differ: the paired samples' first prompt (its token hash, so a changed template shows) and seed, the reward function's version and code hash on every scored - sample, and the plugin files the run plans hashed (a task's own checks and reward live there).""" - eb, et = episodes(base_dir), episodes(trained_dir) - for pid in sorted(set(eb) & set(et)): + sample, and the plugin files the run plans hashed (a task's own checks and reward live there). A run with no + run plan, or whose episode records do not cover its report, is refused: there is nothing to check it against.""" + pb, pt = (_plugin_hashes(d) for d in (base_dir, trained_dir)) + eb, et = paired_episodes(base_dir, trained_dir, base, trained) + for pid in sorted(eb): first = {(e["k"], e["iteration"]): (e["turns"][0]["prompt_ids_sha256"], e["turns"][0]["seed"]) for e in eb[pid] if e.get("turns")} for e in et[pid]: @@ -65,14 +105,15 @@ def fingerprints(base_dir: Path, trained_dir: Path) -> None: if e.get("status") == "scored"} for run in (eb, et)) if rb and rt and rb != rt: raise Refused(f"the two runs were scored by different reward code: {sorted(rb)} against {sorted(rt)}") - pb, pt = (_plugin_hashes(d) for d in (base_dir, trained_dir)) if pb != pt: raise Refused(f"the two runs loaded different plugin files: {pb} against {pt}") def _plugin_hashes(run_dir: Path) -> dict: path = Path(run_dir) / "run_plan.json" - plan = json.loads(path.read_text(encoding="utf-8")) if path.is_file() else {} + if not path.is_file(): + raise Refused(f"{run_dir}: no run_plan.json, so its plugin files cannot be checked") + plan = json.loads(path.read_text(encoding="utf-8")) return {k: (v or {}).get("sha256") for k, v in sorted((plan.get("plugins") or {}).items())} @@ -97,7 +138,7 @@ def compare(base_dir: Path, trained_dir: Path) -> dict: for key in ("task_sha256", "samples"): if b.get(key) != t.get(key): raise Refused(f"{pid}: the two runs differ in {key}: {b.get(key)} against {t.get(key)}") - fingerprints(base_dir, trained_dir) + fingerprints(base_dir, trained_dir, base, trained) sb = base["run"].get("sampling") rows = [] for pid in sorted(base["prompts"]): From 6a1576d12ea4162eecce97fa2771123314bb42c7 Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 14:39:47 +0000 Subject: [PATCH 17/19] fix(aorlta): the Q&A example reads each Answer: line on its own The pattern's \s could cross a newline, so an empty 'Answer:' took the next line's text as its answer, and a final empty marker was passed over for an earlier one. Each line is now matched alone and the last marker decides. The checker's version moves to 2, since cached verdicts were read with the old rule. --- aorlta/examples/qa/qa_task.py | 13 +++++++------ aorlta/tests/examples/test_qa.py | 8 +++++++- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/aorlta/examples/qa/qa_task.py b/aorlta/examples/qa/qa_task.py index f0a4e3e00..29ef8ea81 100644 --- a/aorlta/examples/qa/qa_task.py +++ b/aorlta/examples/qa/qa_task.py @@ -30,7 +30,7 @@ HERE = Path(__file__).resolve().parent SYSTEM = "You are a helpful assistant." USER = "{question}\n\nEnd your reply with a line of the form 'Answer: '." -ANSWER_RE = re.compile(r"^\s*\**\s*answer\s*\**\s*[:：]\s*(.+?)\s*$", re.IGNORECASE | re.MULTILINE) +ANSWER_RE = re.compile(r"^[ \t]*\**[ \t]*answer[ \t]*\**[ \t]*[:：][ \t]*(.*?)[ \t]*$", re.IGNORECASE) # one line def load_rows(pool_dir) -> list[dict]: @@ -57,9 +57,10 @@ def normalise(text: str) -> str: def final_answer(reply: str) -> str | None: - """The last `Answer:` line's answer, or None.""" - found = ANSWER_RE.findall(reply) - return found[-1] if found else None + """The last `Answer:` line's answer, or None. Each line is read on its own, so an empty `Answer:` never takes the + next line as its answer, and the last marker decides, an empty one included.""" + found = [m.group(1) for line in reply.splitlines() if (m := ANSWER_RE.match(line))] + return (found[-1] or None) if found else None def messages(row: dict) -> list[dict]: @@ -139,7 +140,7 @@ def __init__(self, ctx): self.name, self.ctx = ctx.name, ctx self.judge = ctx.settings.get("judge") self.judge_timeout_s = float(ctx.settings.get("judge_timeout_s", 60.0)) - self.version = "1" # cached legs are keyed by it: the judge, which can turn a mismatch clean, is part of it + self.version = "2" # cached legs are keyed by it: the judge, which can turn a mismatch clean, is part of it if self.judge: files = {a: file_identity(a) for a in self.judge if isinstance(a, str) and Path(a).is_file()} exe = shutil.which(self.judge[0]) if isinstance(self.judge[0], str) else None # as PATH finds it, if bare @@ -147,7 +148,7 @@ def __init__(self, ctx): files[f"exe:{self.judge[0]}"] = f"{os.path.realpath(exe)}:{file_identity(exe)}" ident = json.dumps({"judge": list(self.judge), "timeout_s": self.judge_timeout_s, "files": files}, sort_keys=True) - self.version = f"1+judge-{hashlib.sha256(ident.encode()).hexdigest()[:16]}" + self.version = f"2+judge-{hashlib.sha256(ident.encode()).hexdigest()[:16]}" self.runner = ProcessRunner() # kills the judge's whole process group on a timeout or at close() self.rows = {} for pool in ctx.pools: diff --git a/aorlta/tests/examples/test_qa.py b/aorlta/tests/examples/test_qa.py index 5ae920e0b..b74c1e231 100644 --- a/aorlta/tests/examples/test_qa.py +++ b/aorlta/tests/examples/test_qa.py @@ -91,6 +91,12 @@ def test_no_answer_line_is_no_answer(qa): assert qa.final_answer("The answer is 42.") is None +@pytest.mark.parametrize("reply", ["Answer:\n42", "Answer: \nthe rest of my working", "Answer: 7\nmore working\nAnswer:"]) +def test_an_answer_line_is_read_on_its_own(qa, reply): + """An empty `Answer:` does not take the next line's text, and it is the last marker that decides.""" + assert qa.final_answer(reply) is None + + def test_the_parse_gives_a_repair_turn_to_a_reply_without_an_answer(qa): from aorlta.rollout.chain import make_feedback_template, make_parse_adapter from aorlta.rollout.compile_check import make_compiler @@ -209,7 +215,7 @@ def test_the_judge_is_part_of_the_checkers_identity(qa, tmp_path): versions = [adapter(qa).version, adapter(qa, {"judge": [str(judge)]}).version, adapter(qa, {"judge": [str(judge), "--strict"]}).version, adapter(qa, {"judge": [str(judge)], "judge_timeout_s": 5}).version] - assert len(set(versions)) == 4 and versions[0] == "1" + assert len(set(versions)) == 4 and versions[0] == "2" assert adapter(qa, {"judge": [str(judge)]}).version == versions[1] judge.write_text("#!/bin/bash\necho '{\"verdict\": \"wrong\"}'\n") assert adapter(qa, {"judge": [str(judge)]}).version != versions[1] From 9791e27085a6864496d516ef918396391e071ffb Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 14:39:47 +0000 Subject: [PATCH 18/19] docs(aorlta): VALIDATION.md section 6, learning on a real kernel task Qwen3.5-4B trained with GRPO on 54 HIP kernel-writing tasks: 47.4% to 76.6% on 39 held-out tasks (108 samples wrong to right against 17). The reward checks correctness only, so it learned correct kernels, not fast ones. The task set is described, not committed; Not yet shown is updated to match. --- aorlta/VALIDATION.md | 85 +++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 81 insertions(+), 4 deletions(-) diff --git a/aorlta/VALIDATION.md b/aorlta/VALIDATION.md index 49a6f4a78..d69dc856d 100644 --- a/aorlta/VALIDATION.md +++ b/aorlta/VALIDATION.md @@ -114,14 +114,91 @@ room to move (94.9% to 98.4% held out; 12 samples wrong to right, 3 right to wro What this shows and does not: the loop learns, the gain holds on prompts it never saw, and the checks that guard each step (landing, gap, scored fraction) passed throughout. It is one run, one seed, on a toy task the base model -was already fair at; it says nothing yet about a real vertical. +was already fair at; section 6 makes the same measurement on a real kernel task. + +## 6. Learning on a real kernel task + +The kernel example of section 4 has 8 tasks, far too few to learn from. This run used 93 HIP kernel-writing tasks in +eight families (elementwise, reduction, transpose, softmax, layernorm, scan, GEMM, gather/scatter), in float, double, +int and `__half`, with launches that include one thread per element, grid-stride loops with fewer threads than +elements, one block per row, 2-D tiles and a batch on `gridDim.z`. Each task gives a kernel's name, signature, launch +and what it computes, in the format of `examples/hip-kernels/tasks.json`. The model writes the kernel; the check +builds it, runs it on a GPU and compares every output element with a host reference loop, and also fails a write +past the end of an array or into an input. The reward is the `hipk.v1` table (`examples/hip-kernels/hip_kernels.py`): ++1.0 when the output matches, -0.5 wrong output, -0.6 crash or timeout, -0.8 bad signature, -0.9 compile error, -1.0 +no code block, minus a small length penalty. Each task has one correct reference kernel, used only to prove the check +(93 of 93 pass); no kernel carries a planted bug. + +The task set is not in this repository. It was generated by a script that writes each family's ops at several +shapes, with a reference kernel per task, and splits them by op and by shape: within each family the ops are divided +between 54 training and 39 held-out tasks, a held-out task uses no shape a training task of its family uses, and the +whole scan family is held out; the script refuses a split that breaks either rule. The check in +`examples/hip-kernels` handles float tasks only, so running this set needs that check extended to the other dtypes. +These runs used the check as it was before it compiled the candidate alone into a device code object, when the +candidate's host code was built into the harness. So every reply was searched afterwards: none of the trained model's +312 held-out replies prints the verdict line, defines `main` or runs host code of its own (`printf`, a constructor, +`exit`), and in the 5,120 training replies the only matches are in prose or a commented-out `main`. + +To reproduce, put the tasks in `/train/tasks.json` and `/heldout/tasks.json`, copy `configs/hipk-train.yaml` +and `configs/hipk-score.yaml` with `data.prompt_pools[0].path` set to those directories (and, for training, +`rollout.P: 16`, `rollout.temperature: 1.0`, `orchestrator.iterations: 40`, `trainer.lr: 2.0e-5`), then as in section 5: + +```bash +CONFIG= sbatch jobs/aorlta-run.sbatch +python tools/export_adapter.py --ckpt /ckpt/iter-000040 --out +CONFIG= AORLTA_SITE= sbatch jobs/aorlta-run.sbatch # run.root: +CONFIG= AORLTA_SITE= sbatch jobs/aorlta-run.sbatch # adds model.adapter_dir +python tools/eval_compare.py --base --trained --tokenizer \ + --out +``` + +Result (Qwen3.5-4B, 8 MI350X GPUs: 4 engines, one LoRA rank, kernels run on GPUs 4–7; a fresh LoRA, 40 iterations of +16 prompts × 8 samples at temperature 1.0, lr 2e-5, up to 2 compile-repair turns; 52 minutes). Every update LANDED +on all four engines, no sample came near the gap bound (largest 1.24 against 2.0, off-policy fraction 0), and the +scored fraction never fell below 0.984. On the training prompts the share of correct kernels rose from 59% in +iterations 1–5 to 88% in iterations 31–40, and the mean reward from 0.30 to 0.78. + +Held out (39 tasks never trained on, 4 samples × 2 passes each at temperature 0.7, up to 2048 new tokens), from +`tools/eval_compare.py`: + +| | base | trained | +|---|---|---| +| pass rate | 47.4% (148/312) | 76.6% (239/312) | +| mean reward | 0.130 | 0.633 | +| compiled | 84.0% | 97.1% | +| mean reply length | 966 tokens | 223 tokens | + +On the same task, pass and sample slot: 108 samples went from wrong to right and 17 from right to wrong (exact sign +test p = 2e-17). By task, 23 improved and 5 got worse (p = 0.0009). Part of the gain is shorter replies (the base hit +the 2048-token cap on 45 samples), but leaving out every slot where either model hit the cap, the pass rate still +rises from 56% to 77%, with 75 samples wrong to right against 17. The largest gains are on families the model saw +least or not at all: + +- scan, never trained on: 2% → 73%; +- transpose: 31% → 97%; +- softmax: 25% → 63%; +- layernorm: 54% → 92%. + +GEMM (43% → 43%) and gather/scatter (75% → 69%) did not improve. On a held-out running max, the base model spent its +budget in comments and then wrote a kernel that calls `hipMalloc` from device code; the trained model wrote a short +correct running-max loop. + +What this shows and does not: the loop teaches a real, checkable task, and the gain holds on ops, shapes and a +family it never trained on. **But the reward checks correctness only, not speed**, and the model learned exactly +that. On tasks launched with one block per row, none of the trained model's 101 correct kernels has the block's +threads share the row: one thread does the whole row. The base model already did this (45 of 48), and the reward +cannot tell the difference. It also learned one wrong habit: it drops the grid-stride loop where a launch has fewer +threads than elements (one task, 5/8 → 0/8). It is one run with one seed. Writing *fast* kernels needs a speed term +in the reward, which is not done yet. ## Not yet shown -- **No model has been shown to learn a real task.** The learning shown above is on generated arithmetic, a task - chosen because it is checkable and has room to improve; the stand-in kernel prompts exercise the loop, not a task. +- **Fast kernels.** Section 6 shows a model learning to write *correct* kernels on held-out tasks; its reward has no + speed term, the trained kernels leave most of a block idle, and GEMM and gather/scatter did not improve. Each + learning result here is one run with one seed. +- **The 93-task set of section 6 is not in the repository**, and the kernel example's check handles float tasks only. - **The examples' checks are cheap stand-ins** for the checks a real vertical would run (a full test suite, a - sanitizer on a real workload, a simulator). + sanitizer on a real workload, a simulator); section 6's check compares outputs and does not time them. - **Only Qwen3.5-4B, only gfx950, only one node.** Other Qwen3.5 sizes should work and have not been run; Qwen3 models (such as Qwen3-0.6B) need the changes listed under "Models" in GETTING-STARTED.md; fewer than 8 GPUs has run only in the 1-GPU arithmetic and Q&A layouts; MI300X and MI355X have not been run. From 6db2f3050f8ad72d40b2c03a8ceeaad60bbb233d Mon Sep 17 00:00:00 2001 From: Vivek Khandelwal Date: Fri, 13 Nov 2026 15:30:56 +0000 Subject: [PATCH 19/19] fix(aorlta): hipk reads the kernel from the signature; selfcheck fails a missing reference The parse took the kernel name from the prompt id, so a task whose id is not its kernel's name (a shape variant) had its correct answers rejected as no_kernel. It now reads the name from the signature in the prompt's user message, and refuses at start a prompt that has none. The self-check skipped a task with no reference kernel and could exit 0 having checked nothing; a missing reference now fails, and with --sources a run that checked no task fails. --- aorlta/examples/hip-kernels/hip_kernels.py | 30 ++++++++++++++++--- aorlta/tests/examples/test_hip_kernels.py | 34 ++++++++++++++++++++++ 2 files changed, 60 insertions(+), 4 deletions(-) diff --git a/aorlta/examples/hip-kernels/hip_kernels.py b/aorlta/examples/hip-kernels/hip_kernels.py index 1f20871f7..0aa6d1823 100644 --- a/aorlta/examples/hip-kernels/hip_kernels.py +++ b/aorlta/examples/hip-kernels/hip_kernels.py @@ -91,6 +91,16 @@ def defines_kernel(source: str, kernel: str) -> bool: source) is not None +def prompt_kernel(prompt) -> str | None: + """The kernel a prompt asks for, read from the signature in its user message: a task's id names the task and + need not be its kernel's name (shape variants of one kernel have their own ids).""" + for m in getattr(prompt, "messages", None) or (): + if m.get("role") == "user": + found = re.search(r"__global__\s+(?:__launch_bounds__\s*\([^)]*\)\s*)?void\s+(\w+)\s*\(", m["content"]) + return found.group(1) if found else None + return None + + class HipkParse: """`hipk.parse`: the last ```cpp block (the stand-in parse's rule). No block, or a block that does not define the task's kernel, is rejected with a message, which `hipk.feedback` turns into a repair turn.""" @@ -106,6 +116,8 @@ def check_prompts(self, prompts, thinking, max_repair_turns=0): for p in prompts: if not p.prompt_id.startswith("hipk:"): raise ValueError(f"hipk.parse reads hipk prompts, got {p.prompt_id!r}") + if prompt_kernel(p) is None: + raise ValueError(f"{p.prompt_id}: its user message has no __global__ signature to name the kernel") def check(self, text, thinking, prompt): from aorlta.rollout.compile_check import ParseOutcome, parse_code @@ -115,7 +127,9 @@ def check(self, text, thinking, prompt): return ParseOutcome("rejected", rejection=( "Your reply has no complete ```cpp code block. Reply with the whole HIP source in one ```cpp block.", "no_code_block")) - kernel = prompt.prompt_id.split(":", 1)[1] + kernel = prompt_kernel(prompt) + if kernel is None: + raise ValueError(f"{prompt.prompt_id}: its user message has no __global__ signature to name the kernel") if not defines_kernel(code, kernel): return ParseOutcome("rejected", rejection=( f"Your source does not define the kernel `{kernel}`. Define it as a __global__ function with exactly " @@ -546,7 +560,8 @@ def selfcheck(argv=None) -> int: """Each task's reference kernel (reference/.hip), or with --sources DIR the kernel DIR/.hip where there is one, through the check's own steps: compiled alone into a code object, the task's harness built, and unless --build-only run on the GPU and judged by `verdict`. Prints one JSON line per task; exit 0 when all are clean - (built, with --build-only).""" + (built, with --build-only). A task with no reference kernel fails; with --sources a task with no source is + skipped, and a run that checked no task at all fails.""" import subprocess from types import SimpleNamespace @@ -568,13 +583,18 @@ def run(argv, timeout_s, env=None): # what `verdict` reads of a ProcResult; sta ap.add_argument("--sources", default=None, help="a directory of .hip kernels to check instead") ap.add_argument("--build-only", action="store_true", help="compile the harnesses, run nothing (no GPU)") args = ap.parse_args(argv) - failed = 0 + failed = checked = 0 work = Path(tempfile.mkdtemp(prefix="hipk-selfcheck-")) try: for task in load_tasks(args.pool): src = Path(args.sources or HERE / "reference") / f"{task['id']}.hip" if not src.is_file(): + if args.sources is None: + failed += 1 + print(json.dumps({"task": task["id"], "source": str(src), "pass": False, + "error": "no reference kernel"}), flush=True) continue + checked += 1 d = work / task["id"] d.mkdir() (d / "harness.hip").write_text(harness_source(task), encoding="utf-8") @@ -600,7 +620,9 @@ def run(argv, timeout_s, env=None): # what `verdict` reads of a ProcResult; sta print(json.dumps(row), flush=True) finally: shutil.rmtree(work, ignore_errors=True) - return 1 if failed else 0 + if not checked: + print(f"selfcheck: no task of {args.pool} had a kernel to check", file=sys.stderr) + return 1 if failed or not checked else 0 if __name__ == "__main__": diff --git a/aorlta/tests/examples/test_hip_kernels.py b/aorlta/tests/examples/test_hip_kernels.py index 57ccf64a8..2e8760caf 100644 --- a/aorlta/tests/examples/test_hip_kernels.py +++ b/aorlta/tests/examples/test_hip_kernels.py @@ -96,6 +96,40 @@ def test_the_parse_gives_a_repair_turn_for_no_block_or_no_kernel(hipk): assert turn.text == other.rejection[0] and len(turn.text) <= 2000 +def test_the_parse_names_the_kernel_from_the_signature_not_the_task_id(hipk, tmp_path): + """Shape variants of one kernel have their own task ids; the kernel the reply must define is the signature's.""" + from aorlta.orchestrator import wiring + from aorlta.rollout.chain import make_parse_adapter + + task = dict(hipk.load_tasks(EXAMPLE)[1], id="saxpy_n1024") + (tmp_path / "tasks.json").write_text(json.dumps([task])) + (p,) = wiring.TaskPromptPool(plugins.task("hipk"), CharChat(), "off", tmp_path).records() + parse = make_parse_adapter("hipk.parse") + parse.check_prompts([p], "off") + ref = (EXAMPLE / "reference" / "saxpy.hip").read_text() + assert parse.check(f"```cpp\n{ref}\n```", "off", p).kind == "edits" + other = parse.check("```cpp\n__global__ void saxpy_n1024(float a, const float* x, float* y, int n) {}\n```", + "off", p) + assert other.kind == "rejected" and other.rejection[1] == "no_kernel" and "`saxpy`" in other.rejection[0] + unnamed = dataclasses.replace(p, messages=[p.messages[0], {"role": "user", "content": "Write a kernel."}]) + with pytest.raises(ValueError, match="no __global__ signature"): + parse.check_prompts([unnamed], "off") + + +def test_the_selfcheck_fails_a_task_with_no_reference_and_a_run_that_checked_nothing(hipk, tmp_path, capsys): + """By default every task needs its reference kernel; with --sources a task without one is skipped, but a run + that checked no task at all is not a pass.""" + (tmp_path / "pool").mkdir() + (tmp_path / "pool" / "tasks.json").write_text(json.dumps([dict(hipk.load_tasks(EXAMPLE)[0], id="no_reference")])) + assert hipk.selfcheck(["--pool", str(tmp_path / "pool"), "--build-only"]) == 1 + (row,) = [json.loads(x) for x in capsys.readouterr().out.splitlines()] + assert (row["task"], row["pass"], row["error"]) == ("no_reference", False, "no reference kernel") + (tmp_path / "sources").mkdir() + assert hipk.selfcheck(["--pool", str(EXAMPLE), "--sources", str(tmp_path / "sources"), "--build-only"]) == 1 + out = capsys.readouterr() + assert out.out == "" and "had a kernel to check" in out.err + + def proc(rc=0, stdout="", stderr="", timed_out=False, killed=False, launch_error=None): from aorlta.reward.adapters.base import ProcResult