diff --git a/.github/workflows/harness.yml b/.github/workflows/harness.yml new file mode 100644 index 00000000..b4bc1000 --- /dev/null +++ b/.github/workflows/harness.yml @@ -0,0 +1,55 @@ +name: Harness contract + +# TRL's harness rollouts on a bundle's own agents (agent_env.training.harness, tst/unit/training/harness_test.py) +# through OpenEnv's capture proxy, on the uv.lock pins and on the latest OpenEnv. The harness reads the proxy's internals +# (its session registry, export and training trace), which OpenEnv can change in any minor release. It runs on pull +# requests that touch the harness or its tier, weekly, and on demand; it isn't a required check. + +on: + pull_request: + paths: + - "src/agent_env/training/**" + - "tst/unit/training/**" + - "tst/util/fake_vllm.py" + - "pyproject.toml" + - "uv.lock" + - ".github/workflows/harness.yml" + schedule: + - cron: "53 7 * * 1" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: harness-${{ github.event_name == 'pull_request' && github.ref || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +env: + PIP_DISABLE_PIP_VERSION_CHECK: "1" + +jobs: + capture: + name: capture (openenv ${{ matrix.openenv }}) + runs-on: ubuntu-latest + timeout-minutes: 15 + strategy: + fail-fast: false + matrix: + openenv: ["locked", "latest"] + steps: + - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + - name: Install from uv.lock, with the latest OpenEnv for latest + run: | + python -m pip install --quiet -r .github/ci-requirements.txt + uv sync --locked --extra dev --python "$(command -v python)" + if [ "${{ matrix.openenv }}" = latest ]; then uv pip install --python .venv/bin/python --upgrade openenv; fi + .venv/bin/python -c "import importlib.metadata as m; print('openenv', m.version('openenv'))" + - name: Harness tier + run: .venv/bin/python -m pytest tst/unit/training/harness_test.py -rA --timeout=600 --junitxml=junit-harness.xml + - name: No skips + if: always() + run: .venv/bin/python .github/scripts/check_skip_policy.py junit-harness.xml diff --git a/pyproject.toml b/pyproject.toml index 0720a1de..7f9ff35d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,12 @@ dependencies = [ openenv = [ "openenv>=0.9", ] +# agent_env.training.harness: TRL's harness rollouts on a bundle's own agents, through OpenEnv's capture proxy (0.10 +# and later). TRL and vLLM are the trainer's to pin. +harness = [ + "agentenv-framework[openenv]", + "openenv>=0.10", +] # agent_env.training.trl: TRL's trainers on a bundle served as an OpenEnv env. TRL itself is the trainer's to pin. trl = [ "agentenv-framework[openenv]", @@ -133,7 +139,7 @@ coreweave = [ "cwsandbox>=1.19.6,<2", # Reject exec streams that close without an exit status. ] dev = [ - "agentenv-framework[aws,boat,boxd,coreweave,daytona,gcp,openenv,sail,skyrl,tensorlake,tenki,trl,vercel,verl]", # exercise portable optional providers and stores in CI + "agentenv-framework[aws,boat,boxd,coreweave,daytona,gcp,harness,openenv,sail,skyrl,tensorlake,tenki,trl,vercel,verl]", # exercise portable optional providers and stores in CI "fastmcp>=3.4,<4", # the standalone-fastmcp mount path of the protocol; <4 keeps mcp 1.x "Pillow>=10.4", # the protocol's computer extra, whose tests the unit tier runs "av>=17", # the protocol's computer extra, whose tests the unit tier runs and which the recording tests decode with diff --git a/src/agent_env/bundle/preflight.py b/src/agent_env/bundle/preflight.py index beb026b4..b4c8e30d 100644 --- a/src/agent_env/bundle/preflight.py +++ b/src/agent_env/bundle/preflight.py @@ -80,7 +80,7 @@ from agent_env.utils.docker.build import DEFAULT_BUILD_PLATFORM, docker_plugin_missing, docker_unreachable from .parse import BundleError, BundleKind -from .plan import AgentOverride, Plan +from .plan import AgentOverride, Plan, solver_steps from .resolve import BuiltImage, Reference, ReferencedImage, ResolvedEntry, build_step, played_deploys _SHOWN = 3 @@ -96,13 +96,16 @@ class Preflight: built_here: frozenset[str] = frozenset() # the ids of the images the bundle builds that a deploy here runs -def preflight_run(plan: Plan, tasks: Iterable[ResolvedEntry], sandbox: str | None, *, played: bool = False) -> Preflight: +def preflight_run(plan: Plan, tasks: Iterable[ResolvedEntry], sandbox: str | None, *, played: bool = False, + solver_endpoint: bool = False) -> Preflight: """Check where each of ``tasks`` deploys, with ``sandbox`` overriding every deploy's provider, and name the infra - envs the run builds; ``played`` skips each task's played agent's deploys (``played_deploys``). Raises BundleError - listing every problem found. Reads the store but writes nothing.""" + envs the run builds; ``played`` skips each task's played agent's deploys (``played_deploys``), and with + ``solver_endpoint`` the caller gives the solver's deploys (``solver_steps``) their model endpoint. Raises + BundleError listing every problem found. Reads the store but writes nothing.""" walk = _Walk(plan, sandbox) for task in tasks: - walk.task(task, skip=played_deploys(task) if played else frozenset()) + walk.task(task, skip=played_deploys(task) if played else frozenset(), + given=frozenset(solver_steps(task.config)) if solver_endpoint else frozenset()) return walk.finish() @@ -148,7 +151,7 @@ def __init__(self, plan: Plan, sandbox: str | None): # Each deploy of an image the bundle builds: (where, what, its provider's links, the images it accepts on each) self.uses: dict[str, list[tuple[str, str, list[SandboxProvider], Callable[[SandboxProvider], Accepts]]]] = {} - def task(self, task: ResolvedEntry, skip: frozenset[str]) -> None: + def task(self, task: ResolvedEntry, skip: frozenset[str], given: frozenset[str] = frozenset()) -> None: steps = [build_step(config) for config in task.config] sandboxes = {step.sandbox_name: step for step in steps if isinstance(step, DeploySandboxTaskStep)} solver = self.plan.agent.steps.get(task.entry.id, ()) if self.plan.agent else () @@ -161,7 +164,7 @@ def task(self, task: ResolvedEntry, skip: frozenset[str]) -> None: if isinstance(step, DeployEnvTaskStep): self._env(where, step) elif isinstance(step, DeployAgentTaskStep): - self._agent(where, step, sandboxes, override) + self._agent(where, step, sandboxes, override, endpoint_given=step.id in given) elif isinstance(step, DeploySandboxTaskStep): self._sandbox(where, step) elif isinstance(step, RubricsVerifierTaskStep): @@ -289,7 +292,7 @@ def _planned_version(self, kind: EntityKind, entity_id: str, version: int | None return version if version is not None else self.plan.store_latest.get((kind, entity_id)) def _agent(self, where: str, step: DeployAgentTaskStep, sandboxes: dict[str, DeploySandboxTaskStep], - override: AgentOverride | None = None) -> None: + override: AgentOverride | None = None, endpoint_given: bool = False) -> None: """Check a deploy of the step's agent, or of ``override``, the agent ``--agent`` names for the solver.""" accepts: Callable[[SandboxProvider], Accepts] = A2AAgent.accepts if step.sandbox_name: @@ -303,7 +306,8 @@ def _agent(self, where: str, step: DeployAgentTaskStep, sandboxes: dict[str, Dep accepts = functools.partial(A2AAgent.accepts, sandbox_mode=linked.sandbox_mode) # the VM loads its image else: provider = _provider(self.sandbox or step.sandbox_type, get_agent_sandbox_provider) - given = frozenset({*step.env_vars, *([ENV_BASE_URL] if step.litellm_base_url else [])}) + given = frozenset({*step.env_vars, *([ENV_BASE_URL] if step.litellm_base_url else []), + *(_MODEL_VARS if endpoint_given else ())}) agent_id, version = ((override.agent_id, override.version) if override is not None else (step.a2a_agent_id, step.a2a_agent_version)) self._agent_deploy(where, provider, agent_id, version, "names no agent", accepts, given=given) diff --git a/src/agent_env/bundle/run.py b/src/agent_env/bundle/run.py index 07bea46e..9cb4207b 100644 --- a/src/agent_env/bundle/run.py +++ b/src/agent_env/bundle/run.py @@ -41,7 +41,8 @@ from .plan import AliasUse, Plan, Write, plan_bundle from .preflight import Preflight, preflight_run from .requires import Declaring, declare, declaring_for -from .resolve import BuiltImage, BundleImage, Reference, ReferencedImage, resolve_bundle, without_played_agents +from .resolve import (BuiltImage, BundleImage, Reference, ReferencedImage, ResolvedEntry, resolve_bundle, + without_played_agents) logger = logging.getLogger(__name__) @@ -251,7 +252,7 @@ def run_bundle( with Interrupts() as interrupts: overrides = _overrides(sandbox, user_model, judge_model) runs = interrupts.run(_run_all(plan, list(prepared.tasks), model, overrides, - _solver_params(plan, prepared.materialization), keep, say, interrupts)) + solver_params(plan, prepared.materialization), keep, say, interrupts)) _mark_cancelled(runs, interrupts.reason) result = _bundle_run(plan, prepared.materialization, runs, _skipped(plan, entries, every=not tasks and not evals)) @@ -274,17 +275,20 @@ def prepare_bundle( on_progress: Callable[[str], None] | None = None, id_root: str | None = None, played: bool = False, + agent: str | None = None, + solver_endpoint: bool = False, ) -> PreparedBundle: """Do what ``run_bundle`` with these arguments does before its first task runs, and stop there. ``played`` leaves out what each task's only prompted agent would deploy, for a caller that plays that agent itself - (``EpisodeSpec``), so the agent needn't exist. + (``EpisodeSpec``), so the agent needn't exist. With ``solver_endpoint`` the caller gives the solver's deploys their + model endpoint in each run's ``step_params``, so the config needn't set one. It doesn't turn namespace routing on: the caller keeps it on for as long as it runs the tasks (``agent_env.store.routing.namespace_routing``), which suits a service that runs nothing else. Raises what ``run_bundle`` raises before its first task runs, except for a running event loop, which it doesn't need.""" composing = _composing(sandbox) - return _prepare(root, tasks, evals, sandbox, id_root, None, played, _progress(on_progress), _hint(id_root), - composing) + return _prepare(root, tasks, evals, sandbox, id_root, agent, played, _progress(on_progress), _hint(id_root), + composing, solver_endpoint=solver_endpoint) def _write_bundle( @@ -373,8 +377,9 @@ 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()) -> PreparedBundle: - plan, preflight = _planned(root, tasks, evals, id_root, sandbox, agent, played=played) + composing: Composing = Composing(), *, solver_endpoint: bool = False) -> PreparedBundle: + plan, preflight = _planned(root, tasks, evals, id_root, sandbox, agent, played=played, + solver_endpoint=solver_endpoint) _say_plan(plan, say, declaring) materialization = _written_with_infra(plan, preflight, say, composing) to_run = tuple((entry, Task.get(entry.id, materialization.version_of("task", entry.id))) for entry in _to_run(plan)) @@ -382,7 +387,8 @@ 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) -> tuple[Plan, Preflight]: + 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)) plan = plan_bundle(without_played_agents(resolved) if played else resolved, tasks=tasks, evals=evals, installed=installed_bundles, agent=agent) @@ -390,7 +396,7 @@ def _planned(root: Path | str, tasks: Sequence[str], evals: Sequence[str], id_ro 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, - played=played) + played=played, solver_endpoint=solver_endpoint) def _composing(sandbox: str | None, emulate: bool = True) -> Composing: @@ -399,6 +405,16 @@ def _composing(sandbox: str | None, emulate: bool = True) -> Composing: return Composing(build_sandbox_provider(sandbox) if sandbox else None, emulate) +def planned_tasks(root: Path | str, *, tasks: Sequence[str] = (), evals: Sequence[str] = (), + agent: str | None = None) -> tuple[ResolvedEntry, ...]: + """The tasks ``run_bundle`` with these arguments would run, as resolved, without checking or writing anything. + Namespace routing must be on, as for ``plan_bundle``.""" + plan = plan_bundle(resolve_bundle(parse_bundle(Path(root))), tasks=tasks, evals=evals, installed=installed_bundles, + agent=agent) + selected = set(_to_run(plan)) + return tuple(write.source for write in plan.writes if write.kind is BundleKind.TASK and write.source.entry in selected) + + def _to_run(plan: Plan) -> tuple[BundleEntry, ...]: """The tasks selected on their own and every task a selected eval names, each once, in the plan's order.""" wanted = {entry.entry.id for entry in plan.tasks} @@ -445,7 +461,7 @@ def _hint(id_root: str | None) -> Declaring: return Declaring.QUIET if id_root is not None else Declaring.HINT -def _solver_params(plan: Plan, materialization: Materialization) -> dict[str, dict[str, dict]]: +def solver_params(plan: Plan, materialization: Materialization) -> dict[str, dict[str, dict]]: """Each task's ``step_params`` that deploy ``Plan.agent`` at its solver steps: at the version this run wrote, for an agent it writes, or the store's version the plan read.""" if plan.agent is None: diff --git a/src/agent_env/openenv/app.py b/src/agent_env/openenv/app.py index becd7fd5..11acaa64 100644 --- a/src/agent_env/openenv/app.py +++ b/src/agent_env/openenv/app.py @@ -23,6 +23,7 @@ from agent_env.task import EpisodeError, EpisodeSpec, Task from .environment import AgentEnvEnvironment, Served, ServedTask +from .reward import reward_of logger = logging.getLogger(__name__) @@ -142,11 +143,11 @@ def _served(splits: dict[str, list[tuple[str, Task]]], except EpisodeError as e: refused[task.id] = str(e) continue - reward_of, problem = _reward_of(task, reward) + rewarded, problem = reward_of(task, reward) if problem: unscorable[task.id] = problem continue - kept.append(ServedTask(task, name, spec, reward_of)) + kept.append(ServedTask(task, name, spec, rewarded)) served[split] = tuple(kept) for why in refused.values(): logger.warning("left out: %s", why) @@ -156,32 +157,3 @@ def _served(splits: dict[str, list[tuple[str, Task]]], raise BundleError([f"no task can be served as an OpenEnv env: {why}" for why in refused.values()] or ["no task can be served as an OpenEnv env: the bundle has none"]) return served - - -def _reward_of(task: Task, reward: str | Callable[[dict[str, float]], float] | None - ) -> tuple[Callable[[dict[str, float]], float | None] | None, str | None]: - """``task``'s reward as a function of its scores, or why ``reward`` can't make one.""" - if callable(reward): - return reward, None - scoring = [step.id for step in task.steps if getattr(step, "verifier_id", None)] - if reward is None: - if len(scoring) > 1: - return None, (f"task {task.id!r} records a score at each of {', '.join(scoring)}; name the one that is " - "its reward with reward=") - return (_score_at(scoring[0]) if scoring else _unscored), None - if reward not in scoring: - return None, f"task {task.id!r} records no score at step {reward!r}; it records them at {scoring}" - return _score_at(reward), None - - -def _score_at(step: str) -> Callable[[dict[str, float]], float]: - def score(scores: dict[str, float]) -> float: - if step not in scores: - raise LookupError(f"the run recorded no score at step {step!r}") - return scores[step] - - return score - - -def _unscored(scores: dict[str, float]) -> None: - return None diff --git a/src/agent_env/openenv/reward.py b/src/agent_env/openenv/reward.py new file mode 100644 index 00000000..08778e77 --- /dev/null +++ b/src/agent_env/openenv/reward.py @@ -0,0 +1,36 @@ +"""How a task's recorded scores make its reward: the rules ``build_app`` applies, which harness training shares.""" + +from __future__ import annotations + +from collections.abc import Callable + +from agent_env.task import Task + + +def reward_of(task: Task, reward: str | Callable[[dict[str, float]], float] | None + ) -> tuple[Callable[[dict[str, float]], float | None] | None, str | None]: + """``task``'s reward as a function of its scores, or why ``reward`` can't make one.""" + if callable(reward): + return reward, None + scoring = [step.id for step in task.steps if getattr(step, "verifier_id", None)] + if reward is None: + if len(scoring) > 1: + return None, (f"task {task.id!r} records a score at each of {', '.join(scoring)}; name the one that is " + "its reward with reward=") + return (_score_at(scoring[0]) if scoring else _unscored), None + if reward not in scoring: + return None, f"task {task.id!r} records no score at step {reward!r}; it records them at {scoring}" + return _score_at(reward), None + + +def _score_at(step: str) -> Callable[[dict[str, float]], float]: + def score(scores: dict[str, float]) -> float: + if step not in scores: + raise LookupError(f"the run recorded no score at step {step!r}") + return scores[step] + + return score + + +def _unscored(scores: dict[str, float]) -> None: + return None diff --git a/src/agent_env/task_step/task_steps/deploy_agent.py b/src/agent_env/task_step/task_steps/deploy_agent.py index 53baf931..8e763e70 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -365,13 +365,7 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas requested_disk = user_overrides.get("agent_disk_size_gb") or self.disk_size_gb min_disk = agent.metadata.get("min_disk_size_gb") agent_disk_size_gb = max(requested_disk or 0, min_disk or 0) or None - env_vars = dict(self.env_vars) if self.env_vars else {} - override_litellm_key = user_overrides.get("litellm_api_key") - if override_litellm_key: - env_vars["LITELLM_API_KEY"] = config.resolve_reference(override_litellm_key) - resolved_litellm_base_url = user_overrides.get("litellm_base_url") or self.litellm_base_url - if resolved_litellm_base_url: - env_vars["LITELLM_BASE_URL"] = resolved_litellm_base_url + env_vars = {**(self.env_vars or {}), **self._model_endpoint(context)} if self.role is not None: env_vars.setdefault(ROLE_ENV_VAR, self.role) # the generated MCP CLI sends its role from here @@ -509,6 +503,18 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas )) return context + def _model_endpoint(self, context: TaskStepContext) -> dict[str, str]: + """The agent's ``LITELLM_API_KEY`` and ``LITELLM_BASE_URL`` in this run: this step's ``step_params`` over the + run-wide overrides, so a trainer can point the solver alone at its own endpoint, and the base URL over the + step's.""" + overrides = {**context.metadata.get("user_overrides", {}), **self.step_param_overrides(context)} + endpoint = {} + if overrides.get("litellm_api_key"): + endpoint["LITELLM_API_KEY"] = get_config().resolve_reference(overrides["litellm_api_key"]) + if base_url := overrides.get("litellm_base_url") or self.litellm_base_url: + endpoint["LITELLM_BASE_URL"] = base_url + return endpoint + def _resolved_ttl_seconds(self, context: TaskStepContext) -> int: override = context.metadata.get("user_overrides", {}).get("ttl_seconds") if override is not None: diff --git a/src/agent_env/training/_openenv.py b/src/agent_env/training/_openenv.py index 1ceabbd2..8b96ae40 100644 --- a/src/agent_env/training/_openenv.py +++ b/src/agent_env/training/_openenv.py @@ -48,6 +48,8 @@ logger = logging.getLogger(__name__) _waited: set[str] = set() +#: The servers whose refused connections rollouts have been warned they wait out, once a process. +_restarting: set[str] = set() _loop: asyncio.AbstractEventLoop | None = None _loop_lock = threading.Lock() @@ -301,6 +303,10 @@ async def reopened(url: str, timeout: float, *, dropped: bool, patience: float, waited = max(patience, timeout) if restarting else patience if not (restarting or _transient(e)) or time.monotonic() + REOPEN_SECONDS > began + waited: raise Failed(f"{lead}{e}") from e + if restarting and url not in _restarting: + _restarting.add(url) + logger.warning("the server at %s refuses connections, so rollouts wait up to %gs for it to come back", + url, waited) await asyncio.sleep(REOPEN_SECONDS) diff --git a/src/agent_env/training/harness.py b/src/agent_env/training/harness.py new file mode 100644 index 00000000..807eddde --- /dev/null +++ b/src/agent_env/training/harness.py @@ -0,0 +1,240 @@ +"""Train a bundle's own agents with TRL's harness rollouts (``trl.experimental.async_grpo.HarnessRolloutWorker`` with +``harness_adapter=None``): each rollout runs a task as ``agent-env run`` does, its agents' model calls going through +OpenEnv's capture proxy to the trainer's vLLM, which records the tokens TRL trains on, and the task's verifiers score +it.""" + +from __future__ import annotations + +import asyncio +import logging +import re +import secrets +import socket +import threading +import uuid +from collections.abc import Callable, Sequence +from pathlib import Path +from typing import Any + +try: + from openenv.core.harness.capture import CaptureServer, require_llm, to_training_trace + from openenv.core.harness.capture.export import export_session + from openenv.core.harness.capture.sessions import Upstream + from openenv.core.harness.rollout import ResourceSession, ResourceSessionFactory, ToolResult, VerifyResult + from openenv.core.harness.training import TrainingTrace +except ModuleNotFoundError as exc: + if exc.name is None or exc.name.split(".")[0] != "openenv": + raise + raise ImportError("agent_env.training.harness requires pip install 'agentenv-framework[harness]'") from exc + +from agent_env.bundle.plan import solver_steps +from agent_env.bundle.resolve import build_step +from agent_env.bundle.run import planned_tasks, prepare_bundle, solver_params +from agent_env.store.routing import namespace_routing +from agent_env.openenv.reward import reward_of +from agent_env.task.task import Task, recorded_scores +from agent_env.task.teardown import teardown_run +from agent_env.task_step.context import TaskStepContext, sandbox_overrides +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.verifiers.rubrics_verifier import RubricsVerifierTaskStep +from agent_env.training._openenv import on_session_loop + +logger = logging.getLogger(__name__) +#: How a row's prompt names its task, since TRL hands the factory only the prompt. +_NAMED = re.compile(r"^AgentEnv task (\d+): (.+)$") +#: Where every deploy runs: local Docker alone reaches the capture proxy on this machine, a remote sandbox needing a +#: tunnel. +_SANDBOX = "local" + + +def prompt_rows(bundle: str | Path, *, tasks: Sequence[str] = (), evals: Sequence[str] = (), + agent: str | None = None) -> list[dict[str, Any]]: + """The trainer's dataset, a row per task ``AgentEnvSessionFactory`` plays with the same arguments: a system line + naming the task, which is how the factory finds it, then the task's first prompt, for TRL's logs. The model never + sees this prompt: the task's own agent builds its own.""" + return [{"prompt": [{"role": "system", "content": f"AgentEnv task {index}: {task.id}"}, + {"role": "user", "content": _first_prompt(task)}], "task_id": task.id} + for index, (task, _) in enumerate(_trained(bundle, tasks, evals, agent))] + + +class AgentEnvSessionFactory(ResourceSessionFactory): + """OpenEnv's session factory for TRL's ``HarnessRolloutWorker``, built in its rollout process. It prepares the + bundle to run on local Docker, checks that ``llm_url``'s engine returns token ids, and serves OpenEnv's capture + proxy in front of it, where the task's solver reaches it: the agent a step prompts, which it deploys, its other + agents and judges keeping the configured endpoint. TRL sets ``model`` to the policy it publishes; ``reward`` is as + ``build_app``'s, and ``agent`` the agent the solver deploys, as ``agent-env run --agent``'s.""" + + def __init__(self, bundle: str | Path, *, llm_url: str, model: str, sampling: dict[str, Any] | None = None, + tasks: Sequence[str] = (), evals: Sequence[str] = (), + reward: str | Callable[[dict[str, float]], float] | None = None, agent: str | None = None): + self.model = model + self._llm_url, self._sampling = llm_url, sampling + self._trained = _trained(bundle, tasks, evals, agent) + self._rewards = [_reward(task, reward) for task, _ in self._trained] + require_llm(llm_url, model, require_tokens=True) + self._proxy = CaptureServer(llm_url=llm_url, model=model, port=_free_port(), + admin_key=secrets.token_urlsafe(32)) + self._proxy.start() + + def create(self, task: Any, seed: int | None = None, episode_id: str | None = None) -> AgentEnvSession: + index = _index_named(task, [each.id for each, _ in self._trained]) + played, solver = self._trained[index] + capture = self._proxy.registry.create( + upstream=Upstream(llm_url=self._llm_url, model=self.model), capture_level="tokens", purpose="train", + sampling=self._sampling, task=played.id, episode=episode_id or uuid.uuid4().hex) + endpoint = {"litellm_base_url": f"http://127.0.0.1:{self._proxy.port}", "litellm_api_key": capture.session_id} + overrides = {**sandbox_overrides(_SANDBOX), + "step_params": {step: {**params, **endpoint} for step, params in solver.items()}} + return AgentEnvSession(task=played, overrides=overrides, reward=self._rewards[index], + registry=self._proxy.registry, capture=capture, model=self.model) + + +class AgentEnvSession(ResourceSession): + """One run of a task with its own agents, their model calls captured under this session's key.""" + + def __init__(self, *, task: Task, overrides: dict[str, Any], reward: Callable[[dict[str, float]], float | None], + registry: Any, capture: Any, model: str): + self._task, self._overrides, self._reward, self._model = task, overrides, reward, model + self._registry, self._capture = registry, capture + self._scores: dict[str, float] | None = None + self._lock = threading.Lock() + self._closed = False + self._task_run: asyncio.Task | None = None # Task.run, the one part of a run a stop cancels + self._done: threading.Event | None = None # set once the run is torn down + + def initial_messages(self) -> list[dict[str, Any]]: + return [{"role": "user", "content": _first_prompt(self._task)}] + + def list_tools(self) -> list[Any]: + return [] + + def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult: + return ToolResult(error="the task's own agent calls its tools in its sandbox; none are offered here") + + def wait_for_completion(self, timeout_s: float | None = None) -> int: + """Runs the task and tears it down: 0 when it ran to the end, 1 when it failed or ``close`` stopped it. Raises + TimeoutError, once the run is torn down, when the task runs past ``timeout_s``.""" + with self._lock: + if self._closed: + return 1 + self._done = threading.Event() + return on_session_loop(self._run(timeout_s), None) + + def fetch_training_trace(self) -> TrainingTrace: + document = export_session(self._capture, include_messages=True, capture_level="tokens") + return to_training_trace(self._capture.graph, document) + + def verify(self, transcript: list[dict[str, Any]], final_state: Any | None = None) -> VerifyResult: + """The task's reward from its recorded scores; ``None``, which TRL leaves out of the update, when the run + failed or recorded no score its reward reads.""" + try: + reward = None if self._scores is None else self._reward(self._scores) + except LookupError: + reward = None + return VerifyResult(env_reward=reward, done=True, metrics={f"score/{step}": score + for step, score in (self._scores or {}).items()}) + + def close(self) -> None: + """Stops a run still going, as TRL does to the rollouts in flight when it stops, and returns once the run is + torn down, so no sandbox outlives the trainer; then drops the session's captures.""" + with self._lock: + stopping, self._closed = not self._closed, True + task_run, done = self._task_run, self._done + if stopping and task_run is not None: + task_run.get_loop().call_soon_threadsafe(task_run.cancel) + if done is not None: + done.wait() + self._registry.delete(self._capture.session_id) + + async def _run(self, timeout: float | None) -> int: + """``Task.run`` in an asyncio task of its own, which a stop or the timeout cancels, then the teardown in this + coroutine, which nothing cancels, so a stop can't cut it short.""" + context = TaskStepContext(metadata={"task_id": self._task.id, "user_overrides": self._overrides}) + try: + with namespace_routing(): + with self._lock: + if self._closed: + return 1 + task_run = self._task_run = asyncio.ensure_future( + self._task.run(agent_model=self._model, context=context)) + try: + await asyncio.wait_for(task_run, timeout) + self._scores = recorded_scores(self._task.steps, context) + return 0 + except TimeoutError: + raise + except asyncio.CancelledError: + if not task_run.cancelled(): + raise + return 1 + except Exception as e: + logger.warning("task %s failed, so its rollout is left unscored: %s", self._task.id, e) + return 1 + finally: + await teardown_run(context) + finally: + self._done.set() + + +def _trained(bundle: str | Path, tasks: Sequence[str], evals: Sequence[str], + agent: str | None) -> list[tuple[Task, dict[str, dict[str, Any]]]]: + """The selected tasks that can train their solver, prepared as ``agent-env run`` prepares them, each with its + solver's deploys and their step params, ``agent``'s if it swaps the solver. The others are left out before they're + checked, so what they alone need can't stop the rest.""" + with namespace_routing(): + kept = [] + for entry in planned_tasks(bundle, tasks=tasks, evals=evals, agent=agent): + if why := _untrainable(entry.config): + logger.warning("%s is left out: %s", entry.entry.name, why) + else: + kept.append(entry.entry.name) + if not kept: + raise ValueError("no task of the bundle has a solver to train: an agent a step prompts, deployed by a " + "deploy_agent step, that isn't also its judge") + prepared = prepare_bundle(bundle, tasks=kept, sandbox=_SANDBOX, agent=agent, solver_endpoint=True) + swapped = solver_params(prepared.plan, prepared.materialization) + return [(task, {step: swapped.get(entry.id, {}).get(step, {}) + for step in solver_steps([step.to_dict() for step in task.steps])}) + for entry, task in prepared.tasks] + + +def _untrainable(config: list[dict[str, Any]]) -> str | None: + """Why a task's rollouts can't train its solver alone, or None.""" + solver = set(solver_steps(config)) + steps = [build_step(step) for step in config] + deploys = [step for step in steps if step.id in solver] + if not deploys or not all(isinstance(step, DeployAgentTaskStep) for step in deploys): + return "training needs the agent a step prompts deployed by a deploy_agent step" + agents = {step.agent_name for step in deploys} + if any(isinstance(step, RubricsVerifierTaskStep) and step.use_agent_judge and step.agent_name in agents + for step in steps): + return "its rubric judge is its solver, whose grading would be trained on" + return None + + +def _reward(task: Task, reward: str | Callable[[dict[str, float]], float] | None + ) -> Callable[[dict[str, float]], float | None]: + of, why = reward_of(task, reward) + if of is None: + raise ValueError(why) + return of + + +def _index_named(prompt: Any, ids: list[str]) -> int: + for message in prompt if isinstance(prompt, list) else []: + if isinstance(message, dict) and (named := _NAMED.match(str(message.get("content") or ""))): + index = int(named[1]) + if index < len(ids) and ids[index] == named[2]: + return index + raise KeyError("the row's prompt names no task of this factory's: build the dataset with prompt_rows, from the " + "same bundle, tasks and evals") + + +def _first_prompt(task: Task) -> str: + return next((step.prompt for step in task.steps if step.type == "prompt_agent"), "") + + +def _free_port() -> int: + with socket.socket() as probe: + probe.bind(("", 0)) + return probe.getsockname()[1] diff --git a/tst/integration/training/harness_test.py b/tst/integration/training/harness_test.py new file mode 100644 index 00000000..b22c36bc --- /dev/null +++ b/tst/integration/training/harness_test.py @@ -0,0 +1,65 @@ +"""``AgentEnvSessionFactory`` on the shipped agents' smoke tasks, as TRL's harness rollouts run them: the harness, in +local Docker, calls the model through OpenEnv's capture proxy, in front of a stand-in vLLM that has it run the task's +shell command and then answer, so the rollout is captured as trainable turns and scored by the task's verifier, with no +model endpoint configured. Needs a Docker daemon and network access for the agents' images.""" + +from pathlib import Path + +import pytest + +import agent_env +from agent_env.artifact.store import reset_artifact_store +from agent_env.config import configure, reset_config +from agent_env.training.harness import AgentEnvSessionFactory, prompt_rows +from tst.util.fake_vllm import fake_vllm, shell_then_says + +pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] + +MODEL = "Qwen/Qwen3-4B" + + +@pytest.fixture +def state(monkeypatch, tmp_path): + """Local stores under this test's folder and no model endpoint; no rollout may leave a sandbox work folder.""" + sandboxes = tmp_path / "sandboxes" + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "state")) + monkeypatch.setenv("AGENT_ENV_DOCUMENT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_OBJECT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(sandboxes)) + for name in ("LITELLM_API_KEY", "LITELLM_BASE_URL"): + monkeypatch.delenv(name, raising=False) + configure() + reset_artifact_store() + try: + yield tmp_path + assert not sandboxes.exists() or not any(sandboxes.iterdir()), "a rollout left a sandbox work folder" + finally: + reset_artifact_store() + reset_config() + + +@pytest.mark.parametrize("agent", ["opencode", "claude-code"]) +def test_a_harness_runs_its_smoke_task_through_the_capture_proxy_into_trainable_turns(state, agent): + shipped = Path(agent_env.__file__).parent / "agents" / "tasks" / f"{agent}-smoke.json" + (state / "smoke" / "tasks").mkdir(parents=True) + (state / "smoke" / "tasks" / shipped.name).write_text(shipped.read_text()) + + with fake_vllm(MODEL, shell_then_says("echo $((6 * 7))", "42")) as engine: + [row] = prompt_rows(state / "smoke") + factory = AgentEnvSessionFactory(state / "smoke", llm_url=engine.url, model=MODEL, + sampling={"temperature": 1.0}) + session = factory.create(row["prompt"]) + try: + assert session.wait_for_completion(900) == 0 + trace = session.fetch_training_trace() + verified = session.verify([]) + finally: + session.close() + + assert verified.env_reward == 1.0 + assert any(message.get("role") == "tool" for body in engine.requests for message in body["messages"]) + assert len(trace.turns) >= 2 + for turn in trace.turns: + completion = len(turn.completion_token_ids) + assert completion and len(turn.per_token_logps) == completion + assert turn.loss_mask == [0] * len(turn.prompt_token_ids) + [1] * completion diff --git a/tst/unit/bundle/preflight_test.py b/tst/unit/bundle/preflight_test.py index 4c92736b..d60fcda0 100644 --- a/tst/unit/bundle/preflight_test.py +++ b/tst/unit/bundle/preflight_test.py @@ -18,7 +18,7 @@ from agent_env.artifact.artifacts.docker_image import DockerImageArtifact from agent_env.artifact.store import get_artifact_store from agent_env.bundle import BundleError, BundleKind, dry_run_bundle, run_bundle -from agent_env.bundle.run import _write_bundle +from agent_env.bundle.run import _write_bundle, prepare_bundle from agent_env.bundle.resolve import build_step from agent_env.cli import cli from agent_env.config import get_config @@ -1176,12 +1176,18 @@ def test_an_agent_or_agent_judge_needs_a_model_endpoint_unless_its_step_gives_on "judge_a2a_agent_id": "solver", "criteria": [{"id": "c1", "description": "done"}]} direct = {**judge, "id": "direct", "use_agent_judge": False} given = {**AGENT, "env_vars": {"LITELLM_API_KEY": "sk-1"}, "litellm_base_url": "http://models:4000"} - _task(bundle_dir, [AGENT, judge], name="judged") + asks = {"id": "ask", "type": "prompt_agent", "prompt_id": "agent", "prompt": "go", "depends_on": ["agent"]} + _task(bundle_dir, [AGENT, asks, judge], name="judged") _task(bundle_dir, [given, direct], name="given") assert _problems(lambda: dry_run_bundle(bundle_dir, sandbox="local")) == [ "tasks/judged.json: step 'agent', tasks/judged.json: step 'judge' deploy an agent that calls the model endpoint, " "and no model endpoint is set: set LITELLM_BASE_URL, or [model] base_url in .agentenv/config.toml"] + with namespace_routing(): # a trainer gives the solver its endpoint, and the judge still needs the configured one + trained = lambda: prepare_bundle(bundle_dir, tasks=["judged"], sandbox="local", solver_endpoint=True) # noqa: E731 + assert _problems(trained) == [ + "tasks/judged.json: step 'judge' deploys an agent that calls the model endpoint, and no model endpoint is " + "set: set LITELLM_BASE_URL, or [model] base_url in .agentenv/config.toml"] assert [entry.name for entry in dry_run_bundle(bundle_dir, tasks=["given"], sandbox="local").runs] == ["given"] monkeypatch.setenv("LITELLM_BASE_URL", "http://models:4000") [problem] = _problems(lambda: dry_run_bundle(bundle_dir, sandbox="local")) diff --git a/tst/unit/bundle/run_test.py b/tst/unit/bundle/run_test.py index 31d3572d..c51efa6f 100644 --- a/tst/unit/bundle/run_test.py +++ b/tst/unit/bundle/run_test.py @@ -359,9 +359,8 @@ def broken(line): def test_a_run_says_before_it_builds_an_image_and_names_the_image_apart_from_its_agent(bundle_dir, docker, monkeypatch): """As it does for an image a deploy on this machine runs, which preflight names.""" - monkeypatch.setattr(run_module, "preflight_run", - lambda plan, tasks, sandbox, played: Preflight(frozenset(), (), - frozenset({f"{ROOT}/solver__agent_image"}))) + monkeypatch.setattr(run_module, "preflight_run", lambda plan, tasks, sandbox, played, solver_endpoint: Preflight( + frozenset(), (), frozenset({f"{ROOT}/solver__agent_image"}))) layout(bundle_dir, BUILT_AGENT) lines = [] diff --git a/tst/unit/openenv/environment_test.py b/tst/unit/openenv/environment_test.py index 046800be..7f9ac0f3 100644 --- a/tst/unit/openenv/environment_test.py +++ b/tst/unit/openenv/environment_test.py @@ -12,7 +12,7 @@ from openenv.core.env_server.mcp_types import CallToolAction, ListToolsAction, ToolErrorType import agent_env.task.episode.run as episode_run -from agent_env.openenv.app import _score_at +from agent_env.openenv.reward import _score_at from agent_env.openenv.environment import ENDED, NOT_RESET, TIMEOUT_METADATA, AgentEnvEnvironment, Served, ServedTask from agent_env.task import EpisodeError, EpisodeSpec, Task from tst.unit.task.episode.run_test import ALICE, ITEMS, REPLIES, TURN, _gate, _until diff --git a/tst/unit/task_step/test_step_overrides_agent_and_model.py b/tst/unit/task_step/test_step_overrides_agent_and_model.py index 7a7fd976..d9eb11de 100644 --- a/tst/unit/task_step/test_step_overrides_agent_and_model.py +++ b/tst/unit/task_step/test_step_overrides_agent_and_model.py @@ -71,6 +71,25 @@ def test_other_steps_overrides_are_ignored(self): assert _deploy_step()._resolve_agent_ref(ctx) == ("stored-agent", 7) +class TestDeployAgentModelEndpoint: + def test_the_steps_own_url_without_overrides(self): + assert _deploy_step(litellm_base_url="http://stored")._model_endpoint(_ctx()) == { + "LITELLM_BASE_URL": "http://stored"} + + def test_the_run_wide_endpoint_applies(self): + ctx = _ctx(run_overrides={"litellm_base_url": "http://run", "litellm_api_key": "sk-run"}) + assert _deploy_step(litellm_base_url="http://stored")._model_endpoint(ctx) == { + "LITELLM_API_KEY": "sk-run", "LITELLM_BASE_URL": "http://run"} + + def test_the_steps_override_wins_and_another_steps_is_ignored(self): + ctx = _ctx({"deploy-solver": {"litellm_base_url": "http://trainer", "litellm_api_key": "session"}, + "deploy-judge": {"litellm_base_url": "http://elsewhere"}}, + {"litellm_base_url": "http://run", "litellm_api_key": "sk-run"}) + assert _deploy_step()._model_endpoint(ctx) == {"LITELLM_API_KEY": "session", "LITELLM_BASE_URL": "http://trainer"} + assert _deploy_step(id="deploy-judge")._model_endpoint(ctx) == { + "LITELLM_API_KEY": "sk-run", "LITELLM_BASE_URL": "http://elsewhere"} + + class TestPromptModelAndEffort: def test_stored_values_without_overrides(self): assert _prompt_step()._resolve_model_and_effort(_ctx()) == ("stored-model", "stored-effort") diff --git a/tst/unit/test_optional_extras_stay_optional.py b/tst/unit/test_optional_extras_stay_optional.py index 79004823..1b441bef 100644 --- a/tst/unit/test_optional_extras_stay_optional.py +++ b/tst/unit/test_optional_extras_stay_optional.py @@ -32,6 +32,7 @@ def find_spec(self, fullname, path=None, target=None): _NOT_CORE = ( "agent_env.providers.sandbox_providers.sail_vm", "agent_env.providers.sandbox_providers.smol_vm", "agent_env.openenv", "agent_env.training.nemo_gym", "agent_env.training.verl", "agent_env.training.skyrl", + "agent_env.training.harness", "agent_env.env.gateway", "agent_env.task_step.task_steps.run_code_runner", ) @@ -173,7 +174,7 @@ def _run(script: str, blocked: list[str], *args: str) -> str: #: Extras every distribution of which another extra installs too, so blocking what only they install blocks nothing: #: verl's and SkyRL's pyarrow comes with trl's datasets. The test without any extra covers them. -_COVERED = {"verl": "trl", "skyrl": "trl"} +_COVERED = {"verl": "trl", "skyrl": "trl", "harness": "openenv"} @pytest.mark.parametrize("extra", sorted(set(_optional_requirements()) - set(_COVERED))) diff --git a/tst/unit/training/_support.py b/tst/unit/training/_support.py index 157db065..60237369 100644 --- a/tst/unit/training/_support.py +++ b/tst/unit/training/_support.py @@ -3,7 +3,8 @@ import pytest -from agent_env.openenv.app import _app, _score_at +from agent_env.openenv.app import _app +from agent_env.openenv.reward import _score_at from agent_env.openenv.environment import Served, ServedTask from agent_env.task import EpisodeSpec, Task from agent_env.training._openenv import Server diff --git a/tst/unit/training/harness_test.py b/tst/unit/training/harness_test.py new file mode 100644 index 00000000..defddc23 --- /dev/null +++ b/tst/unit/training/harness_test.py @@ -0,0 +1,248 @@ +"""``AgentEnvSessionFactory`` running a bundle's task with its own agents, as TRL's harness rollouts run it: the +solver's deploy alone is given OpenEnv's capture proxy as its model endpoint, in front of a stand-in vLLM that returns +token ids and processed logprobs, with no endpoint configured; its calls are captured and the task's verifier scores +the run.""" + +import asyncio +import json +import logging +import threading +from concurrent.futures import ThreadPoolExecutor + +import httpx +import pytest +from agent_env.a2a_agent import A2AAgent +from agent_env.artifact.artifacts.docker_image import DockerImageArtifact +from agent_env.artifact.store import get_artifact_store +from agent_env.bundle import preflight +from agent_env.config.runtime import Config +from agent_env.store.routing import namespace_routing +from agent_env.task_step.task_step import TaskStep +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep +from agent_env.training import harness +from agent_env.training.harness import AgentEnvSessionFactory, prompt_rows +from tst.unit.bundle._support import layout +from tst.unit.conftest import local_stores # noqa: F401 +from tst.util.fake_vllm import fake_vllm, says, token_ids + +MODEL = "Qwen/Qwen3-4B" +MODEL_ENDPOINT_MISSING = preflight.model_endpoint_missing # the probe, which the unit tier's conftest stubs +DEPLOY = {"id": "agent", "type": "deploy_agent", "env_ids": [], "a2a_agent_id": "theirs"} +ASK = {"id": "ask", "type": "prompt_agent", "prompt_id": "ask", "prompt": "What is 6 * 7?", "depends_on": ["agent"]} +CHECK = {"id": "check", "type": "harness_test_check", "verifier_id": "answer", "depends_on": ["ask"]} +JUDGE = {"id": "judge", "type": "rubrics_verifier", "verifier_id": "judge", "prompt_id": "ask", + "judge_a2a_agent_id": "theirs", "agent_name": DeployAgentTaskStep.DEFAULT_AGENT_NAME, + "criteria": [{"id": "c1", "description": "says 42"}], "depends_on": ["ask"]} +asked: list[dict] = [] +hanging = threading.Event() + + +class _Deploys(DeployAgentTaskStep): + """Deploys nothing: records the model endpoint the run gives the agent.""" + + async def execute(self, context): + if self.a2a_agent_id == "crashes": + raise RuntimeError("the agent crashed") + if self.a2a_agent_id == "hangs": + hanging.set() + await asyncio.sleep(3600) + context.metadata["endpoint"] = self._model_endpoint(context) + return context + + +class _Asks(PromptAgentTaskStep): + """The deployed agent's turn: it asks the model at its endpoint, with its key and the run's model.""" + + async def execute(self, context): + endpoint = context.metadata["endpoint"] + async with httpx.AsyncClient() as client: + answer = (await client.post( + f"{endpoint['LITELLM_BASE_URL']}/v1/chat/completions", + headers={"Authorization": f"Bearer {endpoint['LITELLM_API_KEY']}"}, + json={"model": context.agent_model, "messages": [{"role": "user", "content": self.prompt}]}, + )).json()["choices"][0]["message"]["content"] + asked.append({"model": context.agent_model, "key": endpoint["LITELLM_API_KEY"], "answer": answer, + "overrides": context.metadata["user_overrides"]}) + context.metadata["answer"] = answer + return context + + +class _Checks(TaskStep): + type = "harness_test_check" + entity_refs = () + + def __init__(self, id, version, verifier_id="answer", depends_on=None, fail_task_on_error=True): + super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) + self.verifier_id = verifier_id + + def to_dict(self): + return {**super().to_dict(), "verifier_id": self.verifier_id} + + @classmethod + def from_dict(cls, data): + return cls(**cls._base_from_dict(data), verifier_id=data["verifier_id"]) + + async def execute(self, context): + context.metadata.setdefault("verifications", {})[self.verifier_id] = { + "score": float("42" in context.metadata["answer"])} + return context + + +@pytest.fixture +def engine(): + with fake_vllm(MODEL, says("It is 42.")) as server: + yield server + + +@pytest.fixture(autouse=True) +def registries(monkeypatch): + steps = {**Config().task_step_registry(), "deploy_agent": _Deploys, "prompt_agent": _Asks, _Checks.type: _Checks} + monkeypatch.setattr(Config, "task_step_registry", lambda self: steps) + + +@pytest.fixture +def bundle(tmp_path, monkeypatch, local_stores): # noqa: F811 + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setattr(preflight, "model_endpoint_missing", MODEL_ENDPOINT_MISSING) + monkeypatch.delenv("LITELLM_BASE_URL", raising=False) + monkeypatch.delenv("LITELLM_API_KEY", raising=False) + asked.clear() + hanging.clear() + for agent_id in ("theirs", "crashes", "hangs", "mine"): + with namespace_routing(): + image = get_artifact_store().put_document(DockerImageArtifact( + id=f"{agent_id}-image", description="d", image_name="ghcr.io/acme/agent:1", + tar_gz_s3_url="s3://b/x.tar.gz")) + A2AAgent.put(id=agent_id, docker_image_artifact=image) + return layout(tmp_path / "quiz", {"tasks/sum.json": json.dumps([DEPLOY, ASK, CHECK]), + "tasks/broken.json": json.dumps([{**DEPLOY, "a2a_agent_id": "crashes"}, ASK]), + "tasks/hanging.json": json.dumps([{**DEPLOY, "a2a_agent_id": "hangs"}, ASK]), + "tasks/unprompted.json": json.dumps([DEPLOY]), + "tasks/self-judged.json": json.dumps([DEPLOY, ASK, JUDGE])}) + + +def test_the_solver_alone_calls_the_policy_through_the_proxy_and_the_run_is_captured_and_scored(engine, bundle, caplog): + with caplog.at_level(logging.WARNING, logger=harness.__name__): + rows = prompt_rows(bundle) + factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL, sampling={"temperature": 1.0}) + engine.served.append("trl-policy-v3") + factory.model = "trl-policy-v3" + session = factory.create(next(row["prompt"] for row in rows if row["task_id"].endswith("sum"))) + factory.model = "trl-policy-v4" + + assert session.wait_for_completion() == 0 + trace = session.fetch_training_trace() + verified = session.verify([]) + session.close() + + [call] = asked + [turn] = trace.turns + assert sorted(row["task_id"].rsplit("/", 1)[1] for row in rows) == ["broken", "hanging", "sum"] + assert "unprompted is left out: training needs the agent a step prompts deployed by a deploy_agent step" in ( + caplog.text) + assert "self-judged is left out: its rubric judge is its solver, whose grading would be trained on" in caplog.text + assert (call["model"], call["answer"]) == ("trl-policy-v3", "It is 42.") + assert call["overrides"]["step_params"] == {"agent": {"litellm_base_url": f"http://127.0.0.1:{factory._proxy.port}", + "litellm_api_key": call["key"]}} + assert not {"litellm_base_url", "litellm_api_key"} & call["overrides"].keys() + assert (verified.env_reward, verified.metrics) == (1.0, {"score/check": 1.0}) + assert turn.completion_token_ids == token_ids({"role": "assistant", "content": "It is 42."}) + assert turn.loss_mask == [0] * len(turn.prompt_token_ids) + [1] * len(turn.completion_token_ids) + assert engine.requests[-1]["model"] == "trl-policy-v3" + assert factory._proxy.registry.get(call["key"]) is None + + +def test_the_agent_named_is_the_one_the_solver_deploys(engine, bundle): + factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL, agent="mine") + session = factory.create(next(row["prompt"] for row in prompt_rows(bundle, agent="mine") + if row["task_id"].endswith("sum"))) + + assert session.wait_for_completion() == 0 + [call] = asked + assert call["overrides"]["step_params"]["agent"] == { + "a2a_agent_id": "mine", "a2a_agent_version": 1, "litellm_base_url": f"http://127.0.0.1:{factory._proxy.port}", + "litellm_api_key": call["key"]} + + +def test_a_run_that_fails_is_unscored_rather_than_zero(engine, bundle, caplog): + rows = prompt_rows(bundle) + factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL, sampling={"temperature": 1.0}) + session = factory.create(next(row["prompt"] for row in rows if row["task_id"].endswith("broken"))) + + with caplog.at_level(logging.WARNING, logger=harness.__name__): + assert session.wait_for_completion() == 1 + assert session.verify([]).env_reward is None + assert "failed, so its rollout is left unscored: " in caplog.text + + +@pytest.fixture +def torn(monkeypatch): + """The task ids of the runs torn down, each recorded once its teardown is done.""" + done, teardown = [], harness.teardown_run + + async def tears_down(context): + report = await teardown(context) + done.append(context.metadata["task_id"]) + return report + + monkeypatch.setattr(harness, "teardown_run", tears_down) + return done + + +def test_close_stops_a_rollout_in_flight_and_returns_once_it_is_torn_down(engine, bundle, torn): + factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL) + session = factory.create(next(row["prompt"] for row in prompt_rows(bundle) if row["task_id"].endswith("hanging"))) + + with ThreadPoolExecutor(1) as pool: + waiting = pool.submit(session.wait_for_completion) + assert hanging.wait(30) + session.close() + assert [task_id.rsplit("/", 1)[1] for task_id in torn] == ["hanging"] + assert waiting.result(30) == 1 + session.close() # as TRL closes it again once the rollout returns + assert len(torn) == 1 + assert session.verify([]).env_reward is None + + +def test_a_rollout_past_its_timeout_is_torn_down_before_the_timeout_is_raised(engine, bundle, torn): + factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL) + session = factory.create(next(row["prompt"] for row in prompt_rows(bundle) if row["task_id"].endswith("hanging"))) + + with pytest.raises(TimeoutError): + session.wait_for_completion(0.5) + assert [task_id.rsplit("/", 1)[1] for task_id in torn] == ["hanging"] + + +def test_a_session_closed_before_it_runs_never_runs(engine, bundle, torn): + factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL) + session = factory.create(next(row["prompt"] for row in prompt_rows(bundle) if row["task_id"].endswith("sum"))) + + session.close() + + assert session.wait_for_completion() == 1 + assert (asked, torn) == ([], []) + + +def test_a_prompt_naming_no_task_of_the_factorys_is_refused(engine, bundle): + factory = AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL) + + for prompt in ([{"role": "user", "content": "What is 6 * 7?"}], + [{"role": "system", "content": "AgentEnv task 0: another-bundle/sum"}]): + with pytest.raises(KeyError, match="build the dataset with prompt_rows"): + factory.create(prompt) + + +def test_a_bundle_with_no_solver_to_train_is_refused(engine, bundle): + with pytest.raises(ValueError, match="no task of the bundle has a solver to train"): + AgentEnvSessionFactory(bundle, llm_url=engine.url, model=MODEL, tasks=["unprompted"]) + + +def test_an_engine_without_token_ids_is_refused_before_any_rollout(bundle, monkeypatch): + def eval_only(llm_url, model, **_): + raise RuntimeError(f"{llm_url} is eval-only: no token ids") + + monkeypatch.setattr(harness, "require_llm", eval_only) + + with pytest.raises(RuntimeError, match="eval-only"): + AgentEnvSessionFactory(bundle, llm_url="http://127.0.0.1:9", model=MODEL) diff --git a/tst/unit/training/sessions_test.py b/tst/unit/training/sessions_test.py index f515e115..b74f27f9 100644 --- a/tst/unit/training/sessions_test.py +++ b/tst/unit/training/sessions_test.py @@ -11,16 +11,16 @@ from agent_env.training._openenv import Failed, reopened -@pytest.mark.parametrize(("cause", "waits_for_server", "patience", "opens", "attempts"), [ - (TimeoutError("timed out"), False, 5, True, 3), - (TimeoutError("timed out"), False, 0, False, 1), - (ConnectionRefusedError(61, "Connection refused"), False, 5, False, 1), - (ConnectionRefusedError(61, "Connection refused"), True, 5, True, 3), - (ConnectionRefusedError(61, "Connection refused"), True, 0, True, 3), - (ConnectionError("server rejected WebSocket connection: HTTP 403"), True, 5, False, 1), +@pytest.mark.parametrize(("cause", "waits_for_server", "patience", "opens", "attempts", "warned"), [ + (TimeoutError("timed out"), False, 5, True, 3, 0), + (TimeoutError("timed out"), False, 0, False, 1, 0), + (ConnectionRefusedError(61, "Connection refused"), False, 5, False, 1, 0), + (ConnectionRefusedError(61, "Connection refused"), True, 5, True, 3, 1), + (ConnectionRefusedError(61, "Connection refused"), True, 0, True, 3, 1), + (ConnectionError("server rejected WebSocket connection: HTTP 403"), True, 5, False, 1, 0), ], ids=["busy", "busy-past-patience", "refused", "restarting", "restarting-past-patience", "token-refused"]) def test_a_session_is_tried_again_while_its_server_is_busy_and_given_up_at_once_when_none_listens( - monkeypatch, cause, waits_for_server, patience, opens, attempts): + monkeypatch, caplog, cause, waits_for_server, patience, opens, attempts, warned): tried = [] async def flaky(url, timeout, *, token=None): @@ -33,6 +33,7 @@ async def flaky(url, timeout, *, token=None): monkeypatch.setattr(_openenv, "opened", flaky) monkeypatch.setattr(_openenv, "REOPEN_SECONDS", 0.01) + monkeypatch.setattr(_openenv, "_restarting", set()) if opens: assert asyncio.run(reopened("http://busy", 5, dropped=False, patience=patience, @@ -41,3 +42,4 @@ async def flaky(url, timeout, *, token=None): with pytest.raises(Failed, match="couldn't open a session on the server at http://busy"): asyncio.run(reopened("http://busy", 5, dropped=False, patience=patience, waits_for_server=waits_for_server)) assert len(tried) == attempts + assert caplog.text.count("refuses connections, so rollouts wait up to 5s for it to come back") == warned diff --git a/tst/util/fake_vllm.py b/tst/util/fake_vllm.py new file mode 100644 index 00000000..9a57e132 --- /dev/null +++ b/tst/util/fake_vllm.py @@ -0,0 +1,87 @@ +"""A stand-in for vLLM served with ``--return-tokens-as-token-ids`` and processed logprobs, as OpenEnv's capture proxy +needs its engine: it lists the models it serves and answers chat completions with token ids, a logprob per completion +token and the top two at each, scaled by the temperature as processed logprobs are. ``reply`` makes each answer's +message from the request body.""" + +import json +import uuid +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from typing import Any + +from starlette.applications import Starlette +from starlette.responses import JSONResponse +from starlette.routing import Route + +from agent_env.training._openenv import Server + +Reply = Callable[[dict[str, Any]], dict[str, Any]] + + +def says(content: str) -> Reply: + return lambda body: {"role": "assistant", "content": content} + + +def shell_then_says(command: str, content: str) -> Reply: + """Calls the request's shell tool (``bash`` in any case) with ``command`` until a tool result follows the last user + message, then says ``content``; a request offering no shell gets ``content`` at once, as a harness's side calls + do.""" + + def reply(body: dict[str, Any]) -> dict[str, Any]: + messages = body["messages"] + shell = next((tool["function"]["name"] for tool in body.get("tools") or [] + if tool["function"]["name"].lower() == "bash"), None) + last_user = max((i for i, message in enumerate(messages) if message.get("role") == "user"), default=-1) + if shell is None or any(message.get("role") == "tool" for message in messages[last_user + 1:]): + return {"role": "assistant", "content": content} + arguments = {"command": command, "description": "Run the command"} + return {"role": "assistant", "content": None, "tool_calls": [ + {"id": f"call_{uuid.uuid4().hex[:8]}", "type": "function", + "function": {"name": shell, "arguments": json.dumps(arguments)}}]} + + return reply + + +def token_ids(message: dict[str, Any]) -> list[int]: + """A message's tokens: one per character of its role, content and tool calls.""" + text = f"<{message.get('role')}>{json.dumps(message.get('content'))}{json.dumps(message.get('tool_calls') or [])}" + return [ord(char) % 50000 + 1 for char in text] + + +@contextmanager +def fake_vllm(model: str, reply: Reply) -> Iterator[Server]: + """The engine, served on loopback; ``.requests`` are the completion bodies it got, and ``.served`` the models it + lists, which a test appends to as vLLM lists a LoRA adapter once it's loaded.""" + requests: list[dict[str, Any]] = [] + served = [model] + + async def models(request): + return JSONResponse({"data": [{"id": name, "object": "model"} for name in served]}) + + async def completions(request): + body = await request.json() + requests.append(body) + temperature = body.get("temperature") or 1.0 + message = reply(body) + prompt = [token for each in body["messages"] for token in token_ids(each)] + completion = token_ids(message) + content = [{"token": f"token_id:{token}", "logprob": -0.5 / temperature, "bytes": None, + "top_logprobs": [{"token": f"token_id:{token}", "logprob": -0.5 / temperature}, + {"token": f"token_id:{token + 1}", "logprob": -3.0 / temperature}]} + for token in completion] + return JSONResponse({ + "id": f"chatcmpl-{uuid.uuid4().hex[:8]}", "object": "chat.completion", "model": body["model"], + "prompt_token_ids": prompt, + "choices": [{"index": 0, "message": message, "finish_reason": "tool_calls" if message.get("tool_calls") + else "stop", "token_ids": completion, "logprobs": {"content": content}}], + "usage": {"prompt_tokens": len(prompt), "completion_tokens": len(completion), + "total_tokens": len(prompt) + len(completion)}, + }) + + server = Server(Starlette(routes=[Route("/v1/models", models), + Route("/v1/chat/completions", completions, methods=["POST"])])) + server.requests, server.served = requests, served + try: + yield server + finally: + server.stop() diff --git a/uv.lock b/uv.lock index 17e19491..26252fae 100644 --- a/uv.lock +++ b/uv.lock @@ -137,6 +137,9 @@ gcp = [ { name = "google-cloud-storage" }, { name = "requests" }, ] +harness = [ + { name = "openenv" }, +] nemo-gym = [ { name = "nemo-gym", marker = "python_full_version >= '3.13.14'" }, { name = "omegaconf", marker = "python_full_version >= '3.13.14'" }, @@ -184,7 +187,8 @@ verl = [ [package.metadata] requires-dist = [ { name = "a2a-sdk", specifier = "==0.3.26" }, - { name = "agentenv-framework", extras = ["aws", "boat", "boxd", "coreweave", "daytona", "gcp", "openenv", "sail", "skyrl", "tensorlake", "tenki", "trl", "vercel", "verl"], marker = "extra == 'dev'" }, + { name = "agentenv-framework", extras = ["aws", "boat", "boxd", "coreweave", "daytona", "gcp", "harness", "openenv", "sail", "skyrl", "tensorlake", "tenki", "trl", "vercel", "verl"], marker = "extra == 'dev'" }, + { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'harness'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'nemo-gym'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'skyrl'" }, { name = "agentenv-framework", extras = ["openenv"], marker = "extra == 'trl'" }, @@ -219,6 +223,7 @@ requires-dist = [ { name = "moto", marker = "extra == 'dev'", specifier = ">=5.0.0" }, { name = "nemo-gym", marker = "python_full_version >= '3.13.14' and extra == 'nemo-gym'", specifier = ">=0.7.0rc0" }, { name = "omegaconf", marker = "python_full_version >= '3.13.14' and extra == 'nemo-gym'", specifier = ">=2.3" }, + { name = "openenv", marker = "extra == 'harness'", specifier = ">=0.10" }, { name = "openenv", marker = "extra == 'openenv'", specifier = ">=0.9" }, { name = "packaging", specifier = ">=22" }, { name = "pillow", marker = "extra == 'dev'", specifier = ">=10.4" }, @@ -250,7 +255,7 @@ requires-dist = [ { name = "uvicorn", specifier = ">=0.47" }, { name = "vercel-sandbox", marker = "extra == 'vercel'", specifier = ">=0.7.0,<0.8" }, ] -provides-extras = ["openenv", "trl", "nemo-gym", "verl", "skyrl", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] +provides-extras = ["openenv", "harness", "trl", "nemo-gym", "verl", "skyrl", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] [[package]] name = "agentenv-framework-protocol"