You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
feat(aorlta): add your own task as a plugin, without editing the package - #611
Adding a new task used to mean editing about ten places in the package: fixed name lists and Literals in the config, the schemas, the profiles, the wiring, the parse/feedback/compile lookups and the reward factories. This PR lets a task live in its own module instead.
A run config lists plugin modules under plugins:, either a .py path (relative to aorlta/) or a module name. Each plugin is imported before the config is checked, and registers its parts with aorlta.plugins:
a task: the prompt pool's records, and optionally the landing-check probes;
a parse adapter, a feedback template and a compile check;
a check profile, its check adapters and a reward function.
One behaviour change for plugin tasks only: each engine's start-up self-test now samples fixed open-ended chats, not the task's own first prompts. On MI350X the HIP-kernel example's prompts gave the same 128 tokens for both seeds on 2 of 4 prompts on every engine, which failed a sampler that was working fine. A task can give its own chats instead.
Every place that looked a name up in a fixed list now also checks the registry, so the built-in stand-in path doesn't change. The run plan records each plugin file's sha256. A plugin's data is handled as stand-in data, and its reward has to keep the service's range and parts. ADAPTING.md says both.
Tasks that aren't kernels are first-class too:
rollout.compile_adapter: none turns off the compile step for tasks like Q&A, a verdict or an ISA fuzzer.
A profile needs only the checks its task uses: no compile gate, no sanitizer.
rollout.identical_groups_stop makes the all-identical-groups stop configurable (default unchanged).
A team can keep its task in its own private package:
A plugins: entry can name an entry point in the aorlta.plugins group, or name a module.
AORLTA_PYTHONPATH puts that package on the run job's import path.
ADAPTING.md gains a "pick your vertical" checklist (it ends with the two-line summary), the three task shapes, and the private-package route.
Tests:
tests/test_plugins_out_of_tree.py: installs a package of its own through an entry point, and registers three shapes from it. They are code with a compile gate, Q&A with no compile, and an external tool's verdict. Each one scores through its own check and reward.
tests/test_plugins.py: registration rules, config and run-plan behaviour, every lookup, and the task pool and probes.
tests/orchestrator/test_plugin_wired.py (heavy): a plugin-only task through the real wiring on CPU. The tiny model and in-process engines train one iteration and land it.
All mutations of the new code are caught: 5 lookups, the self-test chats, the entry-point lookup and its duplicate check, and compile_adapter: none. Full CPU suite passes locally in all three parts.
The worked example that uses this (a HIP kernel-writing task) is stacked on top, in a separate PR.
…m 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.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Unresolved profile-name mismatches can cause every rollout candidate to be rejected, while context budgets that omit probe and continuation lengths can prevent startup.
This PR adds plugin-based task extension to aorlta, allowing custom prompts, adapters, profiles, and rewards without editing the package while preserving built-in stand-in behavior.
Changes:
Adds plugin registration and loading before config validation, with plugin file hashes recorded in run plans.
Connects registered components to rollout, reward, prompt-pool, and landing-probe wiring.
Documents the plugin contract and adds unit and CPU integration tests.
File
Description
aorlta/tests/test_plugins.py
Tests registration, loading, lookups, run plans, and task probes.
aorlta/tests/test_config.py
Updates expected validation messages.
aorlta/tests/orchestrator/test_plugin_wired.py
Tests a plugin task training and landing on CPU.
aorlta/src/aorlta/schemas.py
Allows custom check-profile names.
aorlta/src/aorlta/rollout/compile_check.py
Resolves registered compile adapters.
aorlta/src/aorlta/rollout/chain.py
Resolves registered parse and feedback components.
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.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Unresolved issues leave stale registrations after failed plugin imports, prevent offline evaluation of plugin rewards, and can cause startup failures from undersized engine contexts.
Check registered profiles before YAML suffix handling
aorlta/src/aorlta/reward/profiles.py:174
register_profile accepts names ending in .yaml or .yml, but this branch skips the registry for those names. For example, registering 'toy.yaml' against an existing profile file elsewhere succeeds, yet selecting reward.profile: toy.yaml tries to read ./toy.yaml instead and fails. Look up registered names before applying the file-suffix heuristic so every accepted profile name can be loaded.
…al 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.
…ed 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.
… 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.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
Unresolved profile-name mismatches can cause all rollout candidates to be rejected, while failed plugin imports leak registrations and undersized engine contexts can prevent valid configurations from starting.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
Unresolved defects can make valid plugin runs fail at startup, leave stale registrations after failed imports, and block offline evaluation of plugin reward functions.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
Engine context sizing omits startup self-tests and probes, causing otherwise valid runs to fail, while failed plugin imports leave stale registrations that prevent clean retries.
The reason will be displayed to describe this comment to others. Learn more.
Review summary
Verdict: fix-before-ship
The registry is the right cut: built-ins win, a plugin file's sha256 lands in the run plan, and compile_adapter: none is a real switch. Two startup failures will hit a task author on the first try, and a failed import poisons the process.
Reviewed in depth: plugins.py, profile loading, build(), the chain's rejection path, and the run job's import preflight. Skimmed the wired tests.
Findings: 2 High · 2 Medium · 0 Low
plan_eval.evaluate already accepts a plugin reward when you pass rewardfn; the error is only the default. I did not re-file that.
Still open from Copilot, confirmed
A registered profile whose YAML name: differs is accepted, then every candidate is profile_mismatch.
The engine context is sized from the task records alone, so the new self-test chats (and a longer probe) can exceed it.
The reason will be displayed to describe this comment to others. Learn more.
[High] A registered profile whose YAML name: differs rejects every candidate.
The candidate's check_profile is the config key. The reward service compares it with profile.name from the file (service.py, profile_mismatch). Nothing checks the two are the same, so the config loads and then no sample is scored.
Fix: when the name is a registry key, require the decoded profile's name to equal that key and fail config load otherwise. Copilot flagged this too.
The reason will be displayed to describe this comment to others. Learn more.
Fixed in e53bbc9: a profile loaded by name must carry that name, so a registered file naming itself anything else fails at config load. Test: test_a_registered_profile_must_carry_its_registered_name.
The reason will be displayed to describe this comment to others. Learn more.
[High] Every real run still fails closed without an aorta checkout.
The import preflight loads aorta.models and aorta.waitcheck even when the profile never names those adapters, and a missing AORLTA_AORTA_ROOT exits 99. A plugin task with its own checks cannot start the way the docs describe.
Fix: load those two modules only when the profile uses an aorta adapter.
The reason will be displayed to describe this comment to others. Learn more.
Fixed in e53bbc9: the job works out on the host, from the committed config, whether the profile uses an aorta adapter (standin.rj_waitcheck, standin.aorta_consan, standin.aorta_asan), and passes only those modules to the pre-flight, so a plugin profile loads none. The test runs the job's own snippet on a plugin config and a stand-in config, then runs the pre-flight with a missing checkout.
The reason will be displayed to describe this comment to others. Learn more.
[Medium] The context budget ignores the self-test chats and the probes.
rollout_max_len is the longest task prompt plus repairs. The self-test always draws 128 tokens from SELF_TEST_CHATS, which are longer than a short Q&A prompt. With max_repair_turns: 0 and max_new_tokens: 128 the engine's max_model_len can be shorter than that request, and the run dies at start. A custom probe chat has the same hole.
Fix: size max_len from the task prompts, the self-test chats (plus 128 tokens), and the probe chats. Copilot flagged this too.
The reason will be displayed to describe this comment to others. Learn more.
Fixed in e53bbc9: max_len is now the larger of the rollout budget, the self-test chats plus SELF_TEST_TOKENS, and sync.probe_tokens plus the round trip's 16-token continuation. The new test has a short-prompt, no-repair, 16-token config whose old budget was below both.
The reason will be displayed to describe this comment to others. Learn more.
[Medium] A plugin that raises halfway through import cannot be loaded again in this process.
register_* has already run. The except drops the module from sys.modules, not the registry. The retry builds new function objects and _register refuses them as already registered, so the original error is gone.
Fix: snapshot the registry, and on failure put back only the names this import added. Copilot flagged this too.
The reason will be displayed to describe this comment to others. Learn more.
Fixed in e53bbc9 as you suggested: the registry is snapshotted before each plugin's import and only the names that import added are removed on failure, for both file and module plugins. The test fails, then loads the fixed plugin in the same process.
…e or with main
- Every GPU run needs the aorta checkout (AORLTA_AORTA_ROOT): the import pre-flight loads two of its backends
whatever the profile, and exits 99 without it. Its own row, and out of the stand-in-only environment line.
- The score run and the training run are two submissions; give the training command.
- min_scored_fraction is over the candidates handed to the reward service, not every chain generated (also in the
README's correctness table).
- A scored_fraction_below_threshold stop writes no summary for the failing iteration: read stop.detail.policy and
stop.detail.unsettled in run_result.json.
- Name the pull requests that add the worked examples (#613), plugins (#611) and the adapter exporter (#612) where a
passage depends on them; drop hipk.run.gpus as if it were a config key; a checkpoint is not a PEFT adapter.
…e export and compare tools come with #612 and #613
- The README's iteration diagram drew every parse rejection with a turn left as a repair turn. A repair happens only
when the parser sets a rejection reason (the worked examples' parsers do); standin.parse sets none, so its reply
with no code block ends the chain (rollout/chain.py, no_code_block).
- The score-run paragraph names tools/export_adapter.py (#612) and tools/eval_compare.py (#613) as additions.
- The aorta checkout row says #611 makes the import pre-flight load aorta's backends only for profiles that use them.
…ontext, 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.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
A failed plugin load rolls back helper registrations but leaves the helper cached in sys.modules, so a successful retry can silently omit its components.
…n 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.
@amd-vivekag your four points are addressed in e53bbc9 (6ba612a covers a Copilot follow-up on the plugin rollback), with a reply in each thread. Could you take another look?
* docs(aorlta): a getting-started guide, and the README facts a new user needs first
docs/GETTING-STARTED.md puts the steps in order: what a run needs (and what is
untested: fewer GPUs, other GPUs, other models), the CPU tests as CI runs them,
the site file and the job's environment, a score run before a training run,
measured times, reading the outputs, the settings people change, and what each
of the job's exit codes and the common run stops mean.
The README now says the test suite must not run with -n auto on a many-core
machine, that only the Qwen3.5 family works as it stands, that the stand-in
dataset is not published and the example pool is too small to run, and that
the fake components do not read the configured pool or checks.
* docs(aorlta): point the guide's checkpoint note at the README's run outputs
* docs(aorlta): GPU counts and memory, models and checkpoints, getting help, and a pointer from the aorta README
GETTING-STARTED.md: how to place engines, trainer ranks and checks on 1, 2, 4 or
8 GPUs of one node, measured peak memory per GPU (labelled estimates beside
them), how to share a node, how to switch models and what limits that today,
and where checkpoints and adapters go. The README gains a Getting help section
(open an issue with the new AORLTA template, label aorlta; docs feedback as
docs: issues). The top-level README points to aorlta/ as early and evolving.
* docs(aorlta): diagrams of one iteration, of a score run and export, and of the GPU layout
* docs(aorlta): wider boxes in the README diagrams, so they read without scrolling
* docs(aorlta): getting started, where the guide disagreed with the code or with main
- Every GPU run needs the aorta checkout (AORLTA_AORTA_ROOT): the import pre-flight loads two of its backends
whatever the profile, and exits 99 without it. Its own row, and out of the stand-in-only environment line.
- The score run and the training run are two submissions; give the training command.
- min_scored_fraction is over the candidates handed to the reward service, not every chain generated (also in the
README's correctness table).
- A scored_fraction_below_threshold stop writes no summary for the failing iteration: read stop.detail.policy and
stop.detail.unsettled in run_result.json.
- Name the pull requests that add the worked examples (#613), plugins (#611) and the adapter exporter (#612) where a
passage depends on them; drop hipk.run.gpus as if it were a config key; a checkpoint is not a PEFT adapter.
* docs(aorlta): the stand-in parser ends a reply with no code block; the export and compare tools come with #612 and #613
- The README's iteration diagram drew every parse rejection with a turn left as a repair turn. A repair happens only
when the parser sets a rejection reason (the worked examples' parsers do); standin.parse sets none, so its reply
with no code block ends the chain (rollout/chain.py, no_code_block).
- The score-run paragraph names tools/export_adapter.py (#612) and tools/eval_compare.py (#613) as additions.
- The aorta checkout row says #611 makes the import pre-flight load aorta's backends only for profiles that use them.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Plugin rollback leaves evicted helpers attached to their parent package, causing retries to reuse cached helpers without restoring their registrations.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
amd-vivekag
left a comment
Adding a new task used to mean editing about ten places in the package: fixed name lists and Literals in the config, the schemas, the profiles, the wiring, the parse/feedback/compile lookups and the reward factories. This PR lets a task live in its own module instead.
A run config lists plugin modules under
plugins:, either a.pypath (relative toaorlta/) or a module name. Each plugin is imported before the config is checked, and registers its parts withaorlta.plugins:One behaviour change for plugin tasks only: each engine's start-up self-test now samples fixed open-ended chats, not the task's own first prompts. On MI350X the HIP-kernel example's prompts gave the same 128 tokens for both seeds on 2 of 4 prompts on every engine, which failed a sampler that was working fine. A task can give its own chats instead.
Every place that looked a name up in a fixed list now also checks the registry, so the built-in stand-in path doesn't change. The run plan records each plugin file's sha256. A plugin's data is handled as stand-in data, and its reward has to keep the service's range and parts. ADAPTING.md says both.
Tasks that aren't kernels are first-class too:
rollout.compile_adapter: noneturns off the compile step for tasks like Q&A, a verdict or an ISA fuzzer.rollout.identical_groups_stopmakes the all-identical-groups stop configurable (default unchanged).A team can keep its task in its own private package:
plugins:entry can name an entry point in theaorlta.pluginsgroup, or name a module.AORLTA_PYTHONPATHputs that package on the run job's import path.ADAPTING.md gains a "pick your vertical" checklist (it ends with the two-line summary), the three task shapes, and the private-package route.
Tests:
tests/test_plugins_out_of_tree.py: installs a package of its own through an entry point, and registers three shapes from it. They are code with a compile gate, Q&A with no compile, and an external tool's verdict. Each one scores through its own check and reward.tests/test_plugins.py: registration rules, config and run-plan behaviour, every lookup, and the task pool and probes.tests/orchestrator/test_plugin_wired.py(heavy): a plugin-only task through the real wiring on CPU. The tiny model and in-process engines train one iteration and land it.All mutations of the new code are caught: 5 lookups, the self-test chats, the entry-point lookup and its duplicate check, and
compile_adapter: none. Full CPU suite passes locally in all three parts.The worked example that uses this (a HIP kernel-writing task) is stacked on top, in a separate PR.