diff --git a/src/agent_env/bundle/__init__.py b/src/agent_env/bundle/__init__.py index 57290f61d..5c736a577 100644 --- a/src/agent_env/bundle/__init__.py +++ b/src/agent_env/bundle/__init__.py @@ -3,6 +3,7 @@ from .parse import Bundle, BundleEntry, BundleError, BundleKind, is_ignored, parse_bundle from .run import (BundleRun, DryRun, EvalRun, Outcome, PreparedBundle, RunInterrupted, TaskRun, dry_run_bundle, prepare_bundle, run_bundle) +from .variants import Variant __all__ = [ "Bundle", @@ -16,6 +17,7 @@ "PreparedBundle", "RunInterrupted", "TaskRun", + "Variant", "dry_run_bundle", "is_ignored", "parse_bundle", diff --git a/src/agent_env/bundle/plan.py b/src/agent_env/bundle/plan.py index 4dda85f85..9ec3aca68 100644 --- a/src/agent_env/bundle/plan.py +++ b/src/agent_env/bundle/plan.py @@ -66,6 +66,7 @@ from ._fs import fold, relative, with_article from .authoring import AuthoringContext, entry_files, staged_build_context from .parse import NAMED_BY, Bundle, BundleEntry, BundleError, BundleKind, parse_bundle +from .variants import Varied from .resolve import (BuiltImage, BundleImage, Reference, ResolvedBundle, ResolvedEntry, build_step, entity_class, resolve_bundle) @@ -194,6 +195,7 @@ class Plan: # Each alias the run names, by the kind it names and then by alias. alias_uses: Mapping[EntityKind, Mapping[str, AliasUse]] = field(default_factory=dict) agent: AgentOverride | None = None # the agent --agent names for the solver + variant: Varied | None = None # what --set changes, written under its variant's ids def bundle_of(self, entry: BundleEntry) -> Bundle: """The bundle ``entry`` comes from: the installed one for an entry pulled in from one, else the run's.""" diff --git a/src/agent_env/bundle/register.py b/src/agent_env/bundle/register.py index c65abfaa1..ca983e069 100644 --- a/src/agent_env/bundle/register.py +++ b/src/agent_env/bundle/register.py @@ -225,7 +225,7 @@ def _slug(name: str) -> str: # The fields of a Plan that hold no ``@local`` id, so a translated plan keeps them: the store's latest versions, by store # id, and the agent ``--agent`` names, which a registration plans without. -_KEPT = frozenset({"store_latest", "agent"}) +_KEPT = frozenset({"store_latest", "agent", "variant"}) def translated(plan: Plan, ids: Mapping[str, str]) -> Plan: diff --git a/src/agent_env/bundle/run.py b/src/agent_env/bundle/run.py index 3d5d94f25..0ad85b9f4 100644 --- a/src/agent_env/bundle/run.py +++ b/src/agent_env/bundle/run.py @@ -18,6 +18,7 @@ import shlex import time import uuid +from collections import Counter from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass, replace from enum import Enum @@ -45,6 +46,7 @@ from .requires import Declaring, declare, declaring_for from .resolve import (BuiltImage, BundleImage, Reference, ReferencedImage, ResolvedEntry, resolve_bundle, without_played_agents) +from .variants import Variant, with_variant logger = logging.getLogger(__name__) @@ -220,6 +222,7 @@ def run_bundle( keep: bool = False, emulate: bool = True, source: HubSource | None = None, + variant: Variant | None = None, ) -> BundleRun: """Write what the tasks and evals selected in ``root`` need, as ``plan_bundle`` selects them, then run each of their tasks once. ``model``, ``user_model`` and ``judge_model`` replace the models the solver, the @@ -231,7 +234,8 @@ def run_bundle( at each write and step. ``id_root`` roots the ids of a bundle an installed package provides (``InstalledBundle.id_root``) or one downloaded from the Hub (``HubBundle.id_root``); a folder's ids otherwise come from its path. ``source`` says where a Hub bundle came from (``HubBundle.source``), which each run's context - records under ``metadata["bundle_source"]``. + records under ``metadata["bundle_source"]``. ``variant`` sets env keys for this run (``with_variant``), writing + what it changes under ids of their own and leaving the bundle as written. Each run's sandboxes are torn down as it ends, so the contexts returned record sandboxes that are gone. ``keep`` leaves them up instead, for the caller to use and then remove with ``BundleRun.teardown()``. @@ -253,7 +257,8 @@ def run_bundle( say = _progress(on_progress) with namespace_routing(): prepared = _prepare(root, tasks, evals, sandbox, id_root, agent, False, say, - declaring_for(Path(root), installed=id_root is not None, dry_run=False), composing) + declaring_for(Path(root), installed=id_root is not None, dry_run=False), composing, + variant=variant) plan, entries = prepared.plan, [entry for entry, _ in prepared.tasks] with Interrupts() as interrupts: overrides = _overrides(sandbox, user_model, judge_model) @@ -339,6 +344,7 @@ def dry_run_bundle( on_progress: Callable[[str], None] | None = None, id_root: str | None = None, emulate: bool = True, + variant: Variant | None = None, ) -> DryRun: """Check what ``run_bundle`` with these arguments checks before its first task runs, and say what it would write and run, writing and running nothing: no entity, ledger row, lock or instance. Each write's version is @@ -352,7 +358,7 @@ def dry_run_bundle( composing = _composing(sandbox, emulate) say = _progress(on_progress) with namespace_routing(): - plan, preflight = _planned(root, tasks, evals, id_root, sandbox, agent) + plan, preflight = _planned(root, tasks, evals, id_root, sandbox, agent, variant=variant) _say_plan(plan, say, declaring_for(Path(root), installed=id_root is not None, dry_run=True)) materialization = materialize(plan, dry_run=True, on_write=lambda done: say(_written(plan, done)), composing=composing, sandbox=sandbox, @@ -385,9 +391,10 @@ def _bundle_run(plan: Plan, materialization: Materialization, runs: tuple[TaskRu def _prepare(root: Path | str, tasks: Sequence[str], evals: Sequence[str], sandbox: str | None, id_root: str | None, agent: str | None, played: bool, say: Callable[[str], None], declaring: Declaring, - composing: Composing = Composing(), *, solver_endpoint: bool = False) -> PreparedBundle: + composing: Composing = Composing(), *, solver_endpoint: bool = False, + variant: Variant | None = None) -> PreparedBundle: plan, preflight = _planned(root, tasks, evals, id_root, sandbox, agent, played=played, - solver_endpoint=solver_endpoint) + solver_endpoint=solver_endpoint, variant=variant) _say_plan(plan, say, declaring) materialization = _written_with_infra(plan, preflight, say, composing, sandbox) to_run = tuple((entry, Task.get(entry.id, materialization.version_of("task", entry.id))) for entry in _to_run(plan)) @@ -396,11 +403,19 @@ def _prepare(root: Path | str, tasks: Sequence[str], evals: Sequence[str], sandb def _planned(root: Path | str, tasks: Sequence[str], evals: Sequence[str], id_root: str | None, sandbox: str | None, agent: str | None = None, *, played: bool = False, - solver_endpoint: bool = False) -> tuple[Plan, Preflight]: - resolved = resolve_bundle(parse_bundle(Path(root), id_root=id_root)) + solver_endpoint: bool = False, variant: Variant | None = None) -> tuple[Plan, Preflight]: + bundle = parse_bundle(Path(root), id_root=id_root) + varied = with_variant(bundle, variant) if variant is not None else None + if varied is not None: + tasks, evals = varied.selected(tasks, BundleKind.TASK), varied.selected(evals, BundleKind.EVAL) + resolved = resolve_bundle(varied.bundle if varied is not None else bundle) plan = plan_bundle(without_played_agents(resolved) if played else resolved, tasks=tasks, evals=evals, installed=installed_bundles, agent=agent) refuse_store_tasks(plan) + if varied is not None: + plan = replace(plan, variant=varied) + if not _changed(plan): + raise BundleError([f"{variant.typed}: none of the tasks this run selects deploys an env it changes"]) to_run = set(_to_run(plan)) return plan, preflight_run(plan, [write.source for write in plan.writes if write.kind is BundleKind.TASK and write.source.entry in to_run], sandbox, @@ -452,8 +467,9 @@ def refuse_store_tasks(plan: Plan) -> None: def _say_plan(plan: Plan, say: Callable[[str], None], declaring: Declaring) -> None: - """Say what each agent, env and artifact alias names, what the solver deploys in place of the tasks' agents, and - what plugin packages the bundle uses that its bundle.toml doesn't require, as ``declaring`` says (``declare``).""" + """Say what each agent, env and artifact alias names, what the solver deploys in place of the tasks' agents, the + ids what --set changes is written under, and what plugin packages the bundle uses that its bundle.toml doesn't + require, as ``declaring`` says (``declare``).""" for use in (use for uses in plan.alias_uses.values() for use in uses.values()): say(_aliased(use)) if plan.agent is not None: @@ -461,10 +477,22 @@ def _say_plan(plan: Plan, say: Callable[[str], None], declaring: Declaring) -> N solved = sum(entry.id in plan.agent.steps for entry in runs) say(f"--agent {plan.agent.name}: the solver deploys {plan.agent.agent_id}, in {solved} of {len(runs)} " f"task{'s' if len(runs) != 1 else ''}") + if plan.variant is not None: + variant = plan.variant.variant + say(f"{variant.typed}: {_changed(plan)} written under ids ending in ({variant.label})") for line in declare(plan.bundle, declaring): say(line) +def _changed(plan: Plan) -> str: + """How many envs, tasks and evals that --set changes the run writes: ``1 env and 2 tasks``, or ``""``.""" + counts = Counter(write.kind for write in plan.writes + if isinstance(write.source, ResolvedEntry) and plan.variant.changes(write.kind, write.id)) + named = [f"{counts[kind]} {noun}{'s' if counts[kind] != 1 else ''}" for kind, noun in ( + (BundleKind.ENV, "env"), (BundleKind.TASK, "task"), (BundleKind.EVAL, "eval")) if counts[kind]] + return " and ".join([", ".join(named[:-1]), named[-1]] if len(named) > 1 else named) + + def _hint(id_root: str | None) -> Declaring: """How a service writing a bundle declares its plugins: nothing for an installed bundle, else a hint.""" return Declaring.QUIET if id_root is not None else Declaring.HINT diff --git a/src/agent_env/bundle/variants.py b/src/agent_env/bundle/variants.py new file mode 100644 index 000000000..ce13272d8 --- /dev/null +++ b/src/agent_env/bundle/variants.py @@ -0,0 +1,194 @@ +"""Run-only variants for ``agent-env run --set envs/NAME.KEY=VALUE``: the bundle as parsed, with env keys set for one +run. Each env changed, each env and task that names one, and each eval that names such a task is written under a +derived id, so a variant is written and composed beside the bundle's own. Entries keep their names, so every reference to +one by name reaches its variant; the bundle's files stay as written.""" +from __future__ import annotations + +import copy +import hashlib +import tomllib +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Any + +from agent_env.entity.refs import EntityKind, EntityRef, RefSite, ref_sites +from agent_env.task_step.registry import get_task_step_registry + +from .parse import Bundle, BundleEntry, BundleError, BundleKind, segment_problem +from .resolve import entity_class + +_TARGET = f"{BundleKind.ENV.value}/" +_FORM = "envs/NAME.KEY=VALUE, such as envs/desktop.driver=cua-driver" +#: The keys that say what an entry is rather than how it's set, which a run can't change. +_FIXED = ("type", "id") + + +@dataclass(frozen=True) +class Variant: + """The ``--set`` settings of a run, as typed: ``envs/NAME.KEY=VALUE``, KEY of the env in folder ``envs/NAME``, + dotted into a table, set to VALUE, TOML or a bare word.""" + + settings: tuple[str, ...] + + @property + def label(self) -> str: + """The settings as ids end in them: ``desktop.driver cua-driver, desktop.display 3f2a9c1e``, a value an id can't + hold hashed.""" + return ", ".join(f"{key.removeprefix(_TARGET)} {_shown(raw)}" + for key, raw in (setting.partition("=")[::2] for setting in self.settings)) + + @property + def typed(self) -> str: + """The settings as the command line gives them.""" + return " ".join(f"--set {setting}" for setting in self.settings) + + +@dataclass(frozen=True) +class Varied: + """A bundle with a variant's settings: ``bundle``, and ``renamed``, each entry it changes, by kind and id, to its + variant's id.""" + + variant: Variant + bundle: Bundle + renamed: Mapping[tuple[BundleKind, str], str] + + def changes(self, kind: BundleKind, entry_id: str) -> bool: + """Whether ``entry_id`` is the variant's id of an entry of ``kind`` it changes.""" + return entry_id in {new for (of, _), new in self.renamed.items() if of is kind} + + def selected(self, names: Sequence[str], kind: BundleKind) -> tuple[str, ...]: + """``names``, a run's selections of ``kind`` by name or id, each id of an entry it changes as its variant's.""" + return tuple(self.renamed.get((kind, name), name) for name in names) + + +@dataclass(frozen=True) +class _Setting: + env: str # the env it sets, by name + path: tuple[str, ...] + value: Any + + +def with_variant(bundle: Bundle, variant: Variant) -> Varied: + """``bundle`` with ``variant``'s settings in its envs' tomls, and each entry that changes under an id ending in + ``(