diff --git a/AGENTS.md b/AGENTS.md index 265eada1..fa2ee8b7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -110,7 +110,7 @@ and deserialized through a registry. | `eval/` | `Eval` (`eval`), a versioned list of task references, and its store. | | `plugins/` | Entry-point plugins (see Extension points): the group names, `inventory()`, which `agent-env plugin list/show/check` report, `load_failures` and `settings`. The top level is public; the underscored modules are internal. | | `explorer/` | Optional local web UI: `agent-env up`, needs the `explorer` extra, binds loopback `:8234`. | -| `cli/` | Click CLI with the groups `a2a-agent`, `artifact`, `bundle`, `config`, `env`, `eval`, `plugin`, `run`, `task`, `up`. | +| `cli/` | Click CLI with the groups `a2a-agent`, `artifact`, `bundle`, `config`, `env`, `eval`, `plugin`, `run`, `task`, `up`. `task run` and `eval run` hand their runs to the configured `[runner]` and follow them. | | `examples/` | The bundles agent-env ships (`hello`, and `echo-mcp`, which needs Docker), registered under `agent_env.bundles` in `pyproject.toml`. Tests run them, but they are examples for users, not fixtures. | ## Extension points diff --git a/src/agent_env/cli/eval/run.py b/src/agent_env/cli/eval/run.py index 35380a7f..cbba1698 100644 --- a/src/agent_env/cli/eval/run.py +++ b/src/agent_env/cli/eval/run.py @@ -1,7 +1,5 @@ -"""Run all tasks in an eval concurrently, tearing down each run's sandboxes as it ends.""" +"""Run all tasks in an eval through the configured ``[runner]``, following each run to its end.""" -import asyncio -import contextlib import copy import json import os @@ -9,40 +7,14 @@ import click +from agent_env.cli import submitted_runs from agent_env.cli.banner import print_banner from agent_env.cli.identity import get_agent_env_client_id -from agent_env.cli.teardown_output import echo_teardown from agent_env.store.ids import fs_safe from agent_env.task.interrupts import Interrupts -from agent_env.task.teardown import teardown_run from agent_env.task_step.context import TaskStepContext -def _make_callbacks(tag: str): - """Return (on_step_start, on_step_complete) callbacks with a prefix tag.""" - last_score = None - - def on_step_start(i, total, step, context): - click.echo(click.style( - f"{tag} Running step [{i+1}/{total}]: {step.type} (id={step.id})...", - fg="yellow", - )) - - def on_step_complete(i, total, step, context, duration): - nonlocal last_score - click.echo(click.style( - f"{tag} Completed step [{i+1}/{total}]: {step.type} (id={step.id}) [{duration:.1f}s]", - fg="green", - )) - - score = context.metadata.get("score") - if score is not None and score != last_score: - click.echo(click.style(f"{tag} score: {score}", fg="green")) - last_score = score - - return on_step_start, on_step_complete - - def _write_context(context, task_id, output_dir, prefix=""): """Write context JSON to output_dir.""" os.makedirs(output_dir, exist_ok=True) @@ -53,26 +25,6 @@ def _write_context(context, task_id, output_dir, prefix=""): click.echo(click.style(f"{prefix}Output written to: {output_path}", fg="blue")) -async def _run_single(task, tag, output_dir, agent_model=None, agent_artifact_id=None, base_metadata=None): - """Execute a single task run with logging callbacks, then tear down what it deployed, even when it - raised or was cancelled.""" - on_start, on_complete = _make_callbacks(tag) - context = TaskStepContext(metadata=copy.deepcopy(base_metadata or {})) - try: - await task.run( - on_step_start=on_start, - on_step_complete=on_complete, - agent_model=agent_model, - agent_artifact_id=agent_artifact_id, - context=context, - ) - finally: - echo_teardown(tag, await teardown_run(context)) - if output_dir: - _write_context(context, task.id, output_dir, prefix=f"{tag} ") - return task.id, tag, context - - @click.command() @click.option("--id", "eval_id", required=True, help="Eval id") @click.option("--version", "eval_version", default=None, type=int, help="Eval version (defaults to latest)") @@ -84,7 +36,7 @@ async def _run_single(task, tag, output_dir, agent_model=None, agent_artifact_id @click.option("--judge-model", default=None, type=str, help="Override the model rubrics_verifier judges grade with") @click.option("--agent-artifact-id", default=None, type=str, help="Override the agent image artifact for deploy_agent steps") def run(eval_id: str, eval_version: int | None, output_dir: str | None, k: int, max_concurrency: int | None, agent_model: str | None, user_model: str | None, judge_model: str | None, agent_artifact_id: str | None): - """Run all tasks in an eval concurrently.""" + """Run all tasks in an eval through the configured [runner], concurrently.""" if k < 1: raise click.BadParameter("must be at least 1", param_hint="'--k'") if max_concurrency is not None and max_concurrency < 1: @@ -119,71 +71,43 @@ def run(eval_id: str, eval_version: int | None, output_dir: str | None, k: int, if overrides := {key: value for key, value in (("user_model", user_model), ("judge_model", judge_model)) if value}: base_metadata["user_overrides"] = overrides - async def _run_safe(task, tag, slot): - """Wrap _run_single so failures, and a cancel while it waits for its slot, carry task identity.""" - try: - async with slot: - return await _run_single( - task, - tag, - output_dir, - agent_model=agent_model, - agent_artifact_id=agent_artifact_id, - base_metadata=base_metadata, - ) - except BaseException as exc: - return (task.id, tag, exc) - - def _on_signal(count): - message = "Cancelling the runs and tearing them down (Ctrl-C again to stop now)" if count == 1 else ( - "Stopping the teardown now") - click.echo(click.style(message, fg="yellow")) - - async def _run_all(interrupts): - slot = asyncio.Semaphore(max_concurrency) if max_concurrency else contextlib.nullcontext() - coros = [] - for task in tasks: - for run_idx in range(1, k + 1): - if k > 1: - tag = f"[{task.id} run {run_idx}]" - else: - tag = f"[{task.id}]" - coros.append(_run_safe(task, tag, slot)) - return [run.result() for run in await interrupts.gather(coros, _on_signal)] - + submissions = [(f"[{task.id} run {i}]" if k > 1 else f"[{task.id}]", task, copy.deepcopy(base_metadata), None) + for task in tasks for i in range(1, k + 1)] + runs: list[submitted_runs.SubmittedRun] = [] with Interrupts() as interrupts: - results = interrupts.run(_run_all(interrupts)) + submitted_runs.follow(interrupts, submissions, runs, agent_model=agent_model, + agent_artifact_id=agent_artifact_id, limit=max_concurrency) click.echo() # Summarize results failures = [] successes = [] - for result in results: - task_id, tag, payload = result - if isinstance(payload, BaseException): - failures.append((task_id, tag, payload)) + for run in runs: + context = run.instance.context if run.instance else None + if output_dir and context is not None: + _write_context(TaskStepContext.from_dict(context), run.task.id, output_dir, prefix=f"{run.tag} ") + if run.completed: + successes.append((run.tag, (context or {}).get("metadata", {}).get("score"))) else: - score = payload.metadata.get("score") - successes.append((task_id, tag, score)) + failures.append((run.tag, run.outcome if run.done else "CANCELLED")) if successes: click.echo(click.style("Results:", fg="blue")) - for task_id, tag, score in successes: + for tag, score in successes: score_str = f" score={score}" if score is not None else "" click.echo(click.style(f" {tag} PASSED{score_str}", fg="green")) - if failures: - for task_id, tag, exc in failures: - outcome = "CANCELLED" if isinstance(exc, asyncio.CancelledError) else f"FAILED: {exc}" + if len(successes) < total_runs: + for tag, outcome in failures: click.echo(click.style(f" {tag} {outcome}", fg="red")) click.echo(click.style( - f"\n{len(successes)}/{total_runs} runs completed, {len(failures)}/{total_runs} failed.", + f"\n{len(successes)}/{total_runs} runs completed, {total_runs - len(successes)}/{total_runs} failed.", fg="red", )) else: click.echo(click.style(f"\nAll {total_runs} runs completed!", fg="blue")) if interrupts.count: raise SystemExit(128 + interrupts.signum) - if failures: + if len(successes) < total_runs: raise SystemExit(1) diff --git a/src/agent_env/cli/submitted_runs.py b/src/agent_env/cli/submitted_runs.py new file mode 100644 index 00000000..5ef2ef07 --- /dev/null +++ b/src/agent_env/cli/submitted_runs.py @@ -0,0 +1,148 @@ +"""Hand runs to the configured ``[runner]`` and follow them to the end. Progress comes from the run index and the +task instance, the records every runner keeps.""" + +from __future__ import annotations + +import asyncio +import functools +from dataclasses import dataclass, field + +import click + +from agent_env.config import get_runner +from agent_env.runner.runner import RunHandle, RunRecord, RunStatus +from agent_env.task import Task +from agent_env.task.interrupts import Interrupts, tearing_down +from agent_env.task.step_journal import replay_context +from agent_env.task.store import TaskInstance, TaskStepStatus, find_task_instance, get_task_instance_store +from agent_env.task.task import substitute_seed_in_steps +from agent_env.task_step.context import TaskStepContext + +POLL_SECONDS = 2.0 + + +@dataclass +class SubmittedRun: + tag: str + task: Task + handle: RunHandle + record: RunRecord | None = None + instance: TaskInstance | None = None + steps: set[tuple[str, str]] = field(default_factory=set) + definitions: list | None = None + + @property + def done(self) -> bool: + return self.record is not None and self.record.status.is_terminal + + @property + def completed(self) -> bool: + return self.record is not None and self.record.status is RunStatus.COMPLETED + + @property + def outcome(self) -> str: + return f"{self.record.status}" + (f": {self.record.error}" if self.record.error else "") + + +def follow(interrupts: Interrupts, submissions, runs: list[SubmittedRun], *, step_lines=None, **submit_kwargs) -> None: + """Submit each ``(tag, task, metadata, instance_id)`` into ``runs`` and follow the runs to the end, echoing each step + that ends (with ``step_lines(step, context)`` when given, else a verifier's score). A first Ctrl-C cancels the + runs; a second stops waiting for them. ``runs`` holds what was submitted even when following raises.""" + [outcome] = interrupts.run(interrupts.gather([_submit_and_follow(submissions, runs, step_lines, **submit_kwargs)], + _on_signal)) + if not outcome.cancelled() and outcome.exception() is not None: + raise outcome.exception() + + +async def _submit_and_follow(submissions, runs, step_lines, *, limit=None, **submit_kwargs) -> None: + """``stop()`` is always awaited, spared from a first Ctrl-C so a ``LocalRunner`` finishes tearing its runs down, + short of a second. Never ``start()``: that is the serving process's.""" + runner = get_runner() + submit = functools.partial(runner.submit, **submit_kwargs) + pending = list(submissions) + try: + try: + await _follow(runner, submit, pending, runs, limit, step_lines) + except asyncio.CancelledError: + pending.clear() + await asyncio.gather(*(runner.cancel(run.handle.run_id) for run in runs if not run.done)) + await _follow(runner, submit, pending, runs, limit, step_lines) + raise + finally: + if asyncio.current_task().cancelling() < 2: + tearing_down() + await runner.stop() + + +async def _follow(runner, submit, pending, runs: list[SubmittedRun], limit, step_lines) -> None: + while True: + await _poll(runner, runs, step_lines) + while pending and (limit is None or sum(not run.done for run in runs) < limit): + tag, task, metadata, instance_id = pending.pop(0) + handle = await submit(task.id, task.version, metadata=metadata, instance_id=instance_id) + _echo(tag, f"Submitted run {handle.run_id} (instance {handle.instance_id})", "blue") + runs.append(SubmittedRun(tag, task, handle)) + if not pending and all(run.done for run in runs): + return + await asyncio.sleep(POLL_SECONDS) + + +def _on_signal(count: int) -> None: + message = "Cancelling the runs (Ctrl-C again to stop waiting for them)" if count == 1 else ( + "Stopped waiting; what the runs deployed may still be up") + click.echo(click.style(message, fg="yellow")) + + +async def _poll(runner, runs: list[SubmittedRun], step_lines) -> None: + """Reads every run before raising the first read that failed, so a run that ended is seen as ended.""" + errors = [] + for run in runs: + if not run.done: + try: + await _poll_run(runner, run, step_lines) + except Exception as e: + errors.append(e) + if errors: + raise errors[0] + + +async def _poll_run(runner, run: SubmittedRun, step_lines) -> None: + record = await runner.status(run.handle.run_id) + run.instance = find_task_instance(run.handle.instance_id) or run.instance + for step in run.instance.completed_steps if run.instance else []: + if (step.step_id, step.status) not in run.steps: + run.steps.add((step.step_id, step.status)) + _echo_step(run, step.step_id, step.status == TaskStepStatus.SUCCESS, step_lines) + if record is not None and (run.record is None or record.status != run.record.status): + run.record = record + _echo(run.tag, f"Run {run.outcome}", + "green" if record.status is RunStatus.COMPLETED else "red" if record.status.is_terminal else "yellow") + + +def _echo_step(run: SubmittedRun, step_id: str, passed: bool, step_lines) -> None: + if run.definitions is None: + run.definitions = substitute_seed_in_steps( + run.task.steps, ((run.instance.context or {}).get("metadata") or {}).get("seed")) + i, step = next((i, step) for i, step in enumerate(run.definitions) if step.id == step_id) + _echo(run.tag, f"{'Completed' if passed else 'Failed'} step [{i + 1}/{len(run.definitions)}]: {step.type} " + f"(id={step.id})", "green" if passed else "red") + vid = getattr(step, "verifier_id", None) + if not passed or not (step_lines or vid): + return + context = _step_context(run.instance, step_id) + if step_lines: + for line in step_lines(step, context): + _echo(run.tag, line, "green") + elif (score := context.metadata.get("verifications", {}).get(vid, {}).get("score")) is not None: + _echo(run.tag, f"score (verifier_id={vid}): {score}", "green") + + +def _step_context(instance: TaskInstance, step_id: str) -> TaskStepContext: + """What ``step_id`` wrote, replayed from its journal entry; the instance's whole context when it has none.""" + entries = get_task_instance_store().journal_entries_sync(instance.instance_id) + entry = next((e for e in entries if e["step_id"] == step_id), None) + return TaskStepContext.from_dict(replay_context([entry], ()) if entry else instance.context or {}) + + +def _echo(tag: str, message: str, color: str) -> None: + click.echo(click.style(f"{tag} {message}" if tag else message, fg=color)) diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index c30406bf..c38a2cd1 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -12,13 +12,18 @@ import click from agent_env.cli.banner import print_banner +from agent_env.cli import submitted_runs from agent_env.cli.identity import get_agent_env_client_id from agent_env.cli.teardown_output import echo_teardown from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox +from agent_env.runner.runner import KEEP_KEY, RunStatus from agent_env.store.ids import derive_id, fs_safe, is_local_id, validate_local_id +from agent_env.store.routing import run_scope from agent_env.task.interrupts import Interrupts +from agent_env.task.store import get_task_instance_store, new_instance_id from agent_env.task.teardown import TeardownReport, kind, teardown_run -from agent_env.task_step.context import read_dual_keyed +from agent_env.task_step.context import TaskStepContext, read_dual_keyed +from agent_env.task_step.context_ops import build_context_update_ops from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep @@ -248,30 +253,6 @@ def _format_verifier_output(vid: str, vdata: dict) -> list[str]: return lines -# -- Single-run (k=1) callbacks: verbose output, no prefix -- - -async def _log_step_start(i, total, step, context): - async with _print_lock: - if i == 0 and context.instance_id: - click.echo(click.style(f"Task instance: {context.instance_id}", fg="blue")) - if i > 0: - click.echo() - click.echo(click.style( - f"Running step [{i+1}/{total}]: {step.type} (id={step.id})...", - fg="yellow", - )) - - -async def _log_step_complete(i, total, step, context, duration): - async with _print_lock: - click.echo(click.style( - f"Completed step [{i+1}/{total}]: {step.type} (id={step.id}) [{duration:.1f}s]", - fg="green", - )) - for line in _step_output_lines(step, context): - click.echo(click.style(line, fg="green")) - - # -- Parallel-run (k>1) callbacks: compact output with [run N] prefix -- def _make_parallel_callbacks(run_index): @@ -405,6 +386,24 @@ def _hold(tagged_contexts, interrupts: Interrupts) -> None: _tear_down_kept(kept) +def _ended(runs) -> list: + """Each run that ended on its own, with its stored context: what ``--keep`` leaves this command to tear down. The + runner tears a cancelled one down itself.""" + return [(run.tag, TaskStepContext.from_dict(run.instance.context)) for run in runs + if run.done and run.record.status is not RunStatus.CANCELED and run.instance and run.instance.context] + + +def _stored(task, context: TaskStepContext) -> str: + """A new instance holding ``context``, for the runner to continue a run from, recorded where the task's run records + go. A write that fails raises: the run would otherwise continue from an empty context.""" + instance_id = new_instance_id(task.id) + store = get_task_instance_store() + with run_scope(task.id): + store.upsert_instance(instance_id, task.id, task.version, len(task.steps), context.metadata["start_step"]) + store.seed_context(instance_id, build_context_update_ops(None, context)) + return instance_id + + @click.command() @click.option("--id", "task_id", required=True, help="Task id") @click.option("--version", "task_version", default=None, type=int, help="Task version (defaults to latest)") @@ -437,16 +436,15 @@ def _hold(tagged_contexts, interrupts: Interrupts) -> None: "another terminal can resume the run with --start-step and --context-json. Without it, each run " "is torn down as it ends, as Ctrl-C mid-run always does.") def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, agent_model: str | None, user_model: str | None, judge_model: str | None, agent_artifact_id: str | None, start_step: int, context_json: str | None, litellm_api_key: str | None, judge_litellm_api_key: str | None, apply_trajectory_filter: bool | None, a2a_agent_id: str | None, agent_sandbox: str | None, env_sandbox: str | None, gateway_env_id: str | None, service_db_env_id: str | None, env_state_type: str | None, env_state_instance_id: str | None, keep: bool): - """Run a task by executing its steps sequentially. + """Run a task through the configured [runner] and follow it. Each run's sandboxes are torn down as it ends, passed or failed; its instance and context JSON stay. - --keep holds them up until Ctrl-C instead. Ctrl-C or SIGTERM mid-run cancels the runs, tears them down - and exits 130 (143 for SIGTERM); a second one stops the teardown.""" + --keep holds them up until Ctrl-C instead. Ctrl-C or SIGTERM mid-run cancels the runs, which tears them down, + and exits 130 (143 for SIGTERM); a second one stops waiting.""" if k < 1: raise click.BadParameter("must be at least 1", param_hint="'--k'") from agent_env.task import Task - from agent_env.task_step.context import TaskStepContext if output_dir is None: output_dir = f"/tmp/taskid={fs_safe(task_id)}-run{uuid4().hex[:6]}" @@ -483,6 +481,8 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, initial_context.metadata.setdefault("user_overrides", {})["gateway_env_id"] = gateway_env_id if service_db_env_id: initial_context.metadata.setdefault("user_overrides", {})["service_db_env_id"] = service_db_env_id + initial_context.metadata["start_step"] = start_step + initial_context.metadata[KEEP_KEY] = keep initial_context.metadata["task_id"] = task_id click.echo(f"Fetching task: id={task_id} version={task_version or 'latest'}...") @@ -491,50 +491,43 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, print_banner() - run_kwargs = dict(agent_model=agent_model, agent_artifact_id=agent_artifact_id, start_step=start_step) - if k > 1: - # Parallel runs: compact output with [run N] prefixes - click.echo(click.style(f"Running {k} task runs in parallel...", fg="blue")) - click.echo() - tagged = [(f"[run {i}]", copy.deepcopy(initial_context)) for i in range(1, k + 1)] - callbacks = [_make_parallel_callbacks(i) for i in range(1, k + 1)] - else: - tagged = [("", initial_context)] - callbacks = [(_log_step_start, _log_step_complete)] - - async def _run_all(interrupts): - runs = [ - _run_and_settle(task, tag, context, keep=keep, output_dir=output_dir, output_name=task_id, - on_start=on_start, on_complete=on_complete, **run_kwargs) - for (tag, context), (on_start, on_complete) in zip(tagged, callbacks) - ] - return _outcomes(await interrupts.gather(runs, _on_signal)) - + tags = [f"[run {i}]" for i in range(1, k + 1)] if k > 1 else [""] + submissions = [(tag, task, copy.deepcopy(initial_context.metadata), + _stored(task, initial_context) if context_json else None) for tag in tags] + runs: list[submitted_runs.SubmittedRun] = [] with Interrupts() as interrupts: - results = interrupts.run(_run_all(interrupts)) + try: + submitted_runs.follow(interrupts, submissions, runs, agent_model=agent_model, + agent_artifact_id=agent_artifact_id, step_lines=_step_output_lines if k == 1 else None) + except Exception: + if keep and (kept := _kept(_ended(runs))): + _tear_down_kept(kept) + raise interrupted = interrupts.signum if interrupts.count else None click.echo() - failures = [(tag, r) for (tag, _), r in zip(tagged, results) if isinstance(r, BaseException)] - if k > 1: - for tag, exc in failures: - outcome = "CANCELLED" if isinstance(exc, asyncio.CancelledError) else f"FAILED: {exc}" - click.echo(click.style(f"{tag} {outcome}", fg="red")) - if failures: - click.echo(click.style( - f"{len(results) - len(failures)}/{k} runs completed, {len(failures)}/{k} failed.", fg="red")) - else: - click.echo(click.style(f"All {k} runs completed!", fg="blue")) - elif failures and interrupted is not None: - click.echo(click.style("Task cancelled.", fg="red")) - elif not failures: - click.echo(click.style("Task completed!", fg="blue")) + unexported = False + for run in runs: + if run.instance is None or run.instance.context is None: + continue + context = TaskStepContext.from_dict(run.instance.context) + prefix = f"{run.tag} " if run.tag else "" + try: + _write_context(context, task.id, output_dir, prefix=prefix) + except OSError as e: + unexported = True + click.echo(click.style(f"{prefix}Couldn't write the run's context: {e}", fg="red")) + completed = sum(run.completed for run in runs) + if interrupted is not None: + click.echo(click.style("Task cancelled." if k == 1 else f"{completed}/{k} runs completed.", fg="red")) + elif completed == k and not unexported: + click.echo(click.style("Task completed!" if k == 1 else f"All {k} runs completed!", fg="blue")) + elif k > 1: + click.echo(click.style(f"{completed}/{k} runs completed, {k - completed}/{k} failed.", fg="red")) if keep: - _settle_kept(tagged, interrupts, interrupted) + _settle_kept(_ended(runs), interrupts, interrupted) if interrupted is not None: raise SystemExit(128 + interrupted) - if failures and k == 1: - raise failures[0][1] # torn down already; the CLI reports it as it always has - if failures: + if completed < k or unexported: raise SystemExit(1) @@ -580,7 +573,6 @@ def run_batch(task_id: str, task_version: int | None, seeds: str, concurrency: i are rendered using the seed values before being sent to the agent. """ from agent_env.task import Task - from agent_env.task_step.context import TaskStepContext if output_dir is None: output_dir = f"/tmp/taskid={fs_safe(task_id)}-batch{uuid4().hex[:6]}" diff --git a/src/agent_env/runner/local_runner.py b/src/agent_env/runner/local_runner.py index 842f1012..81ced9ff 100644 --- a/src/agent_env/runner/local_runner.py +++ b/src/agent_env/runner/local_runner.py @@ -2,9 +2,9 @@ semaphore. ``submit()`` fires the run in the background and returns immediately; there is no persisted queue or lease, so the run's asyncio task lives only in this process. -Consequence: a run does not survive a hub restart. ``start()`` fails any run a previous -process left non-terminal (it can't still be executing), and ``stop()`` cancels the ones in -flight. A durable, lease-based runner that resurrects a run across a restart — and supports +Consequence: a run does not survive its process. Each run records the process that owns it, which marks it seen +while it runs; ``start()``, and a started runner from then on, fails a run left non-terminal by an owner that is +gone (it can't still be executing), and ``stop()`` cancels the ones in flight. A durable, lease-based runner that resurrects a run across a restart — and supports multiple worker processes — is a follow-up; configure an external durable ``[runner]`` where that matters. """ @@ -13,11 +13,15 @@ import asyncio import contextlib import logging +import os +import socket import uuid +from datetime import datetime, timedelta, timezone from typing import Optional from agent_env.runner import store as run_store from agent_env.runner.runner import KEEP_KEY, RunHandle, RunRecord, Runner, RunStatus +from agent_env.task.store import new_instance_id from agent_env.task.teardown import TERMINATE_TIMEOUT_SECONDS, teardown_run logger = logging.getLogger(__name__) @@ -34,6 +38,10 @@ class LocalRunner(Runner): STOP_WAIT_SECONDS = 60 # After those waits, how long stop() gives a run it cancels again before shutting down without it. ABANDON_WAIT_SECONDS = 1 + # How often a runner marks its runs seen, and a started one reconciles orphans. + OWNER_SEEN_SECONDS = 30 + # How long after it last marked a run seen an owner on another host still counts as alive. + OWNER_LEASE_SECONDS = 120 def __init__(self, workers: int = 2) -> None: if workers < 1: @@ -45,15 +53,41 @@ def __init__(self, workers: int = 2) -> None: self._contexts: dict = {} # each run's context, so stop() can still tear down a run it gives up on self._abandoned: set[str] = set() # runs stop() gave up on and tore down itself self._stopping = False + self._reconciles = False + self._watch: Optional[asyncio.Task] = None # --- lifecycle --------------------------------------------------------- async def start(self) -> None: - """Reconcile orphans from a previous process: a run left QUEUED/RUNNING has no live - task here, so fail it rather than leave it hanging. There is no queue to resume.""" + """Reconcile orphans now and every ``OWNER_SEEN_SECONDS`` after: a run left QUEUED/RUNNING by an owner that is + gone has no live task anywhere, so fail it rather than leave it hanging. One a live process owns, such as a + CLI following its runs, is left to it. There is no queue to resume.""" run_store.ensure_indexes() + self._reconciles = True + self._reconcile() + self._watching() + + def _watching(self) -> None: + if self._watch is None: + self._watch = asyncio.create_task(self._watch_runs(), name="agent-env-local-runner-watch") + + async def _watch_runs(self) -> None: + while True: + await asyncio.sleep(self.OWNER_SEEN_SECONDS) + try: + for run_id in list(self._inflight): + run_store.see(run_id) + if self._reconciles: + self._reconcile() + except Exception: + logger.warning("Couldn't mark runs seen or reconcile orphans; trying again in %ss", + self.OWNER_SEEN_SECONDS, exc_info=True) + + def _reconcile(self) -> None: orphaned = 0 for record in run_store.active_runs(self.type): + if _owner_alive(record.owner): + continue run_store.mark_terminal(record.run_id, RunStatus.FAILED, error="interrupted by a hub restart") orphaned += 1 if orphaned: @@ -64,6 +98,9 @@ async def stop(self) -> None: its own teardown by TEARDOWN_WAIT_SECONDS. A run still going after a run's whole allowance, a step that ignored its cancel or a teardown that never ended, is cancelled again and left behind.""" self._stopping = True + if self._watch is not None: + self._watch.cancel() + self._watch = None for run_id, task in list(self._inflight.items()): if run_id not in self._tearing_down: task.cancel() @@ -114,10 +151,11 @@ async def submit( instance_id: Optional[str] = None, ) -> RunHandle: run_id = f"local-{uuid.uuid4().hex}" - instance_id = instance_id or uuid.uuid4().hex # minted here so callers can link immediately + instance_id = instance_id or new_instance_id(task_id) # minted here so callers can link immediately record = RunRecord( run_id=run_id, runner=self.type, task_id=task_id, task_version=task_version, instance_id=instance_id, status=RunStatus.QUEUED, + owner={"host": socket.gethostname(), "pid": os.getpid(), "seen": datetime.now(timezone.utc).isoformat()}, overrides={ "agent_model": agent_model, "agent_artifact_id": agent_artifact_id, @@ -126,6 +164,7 @@ async def submit( ) run_store.insert_run(record) self._inflight[run_id] = asyncio.create_task(self._run(record), name=f"agent-env-run-{run_id}") + self._watching() logger.info("Submitted run %s for task %s v%s", run_id, task_id, task_version) return RunHandle(run_id=run_id, instance_id=instance_id) @@ -137,17 +176,15 @@ async def status(self, run_id: str) -> Optional[RunRecord]: return record async def cancel(self, run_id: str) -> bool: + """Cancel a run this runner is running; False for any other, whose record it leaves to the process running it.""" record = run_store.get_run(run_id) - if record is None or record.status.is_terminal: - return False task = self._inflight.get(run_id) - if task is not None: - # Signal only. `_run` is the sole writer of the terminal state, so a cancel that - # races the run's completion can't mislabel a finished run (see `_run`). This - # preempts Task.run() at its current await; it unwinds its in-flight step cleanly. - task.cancel() - else: - run_store.mark_terminal(run_id, RunStatus.CANCELED) # no live task here (defensive) + if record is None or record.status.is_terminal or task is None: + return False + # Signal only. `_run` is the sole writer of the terminal state, so a cancel that + # races the run's completion can't mislabel a finished run (see `_run`). This + # preempts Task.run() at its current await; it unwinds its in-flight step cleanly. + task.cancel() return True # --- execution --------------------------------------------------------- @@ -267,3 +304,21 @@ def _seed_context(self, record: RunRecord, stored): context.metadata["instance_id"] = record.instance_id context.metadata["runner"] = self.type return context + + +def _owner_alive(owner: Optional[dict]) -> bool: + """Whether the process that submitted a run may still be running it: on this host, whether its pid is alive; on + another, whether it marked the run seen within ``OWNER_LEASE_SECONDS``. A run naming no owner predates owners.""" + if not owner: + return False + if owner["host"] == socket.gethostname(): + try: + os.kill(owner["pid"], 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + seen = owner.get("seen") + lease = timedelta(seconds=LocalRunner.OWNER_LEASE_SECONDS) + return seen is not None and datetime.now(timezone.utc) - datetime.fromisoformat(seen) < lease diff --git a/src/agent_env/runner/runner.py b/src/agent_env/runner/runner.py index 761aea96..1fffd7d1 100644 --- a/src/agent_env/runner/runner.py +++ b/src/agent_env/runner/runner.py @@ -55,6 +55,7 @@ class RunRecord: finished_at_utc: Optional[str] = None error: Optional[str] = None overrides: dict[str, Any] = field(default_factory=dict) + owner: Optional[dict[str, Any]] = None def to_dict(self) -> dict: return { @@ -69,6 +70,7 @@ def to_dict(self) -> dict: "finished_at_utc": self.finished_at_utc, "error": self.error, "overrides": self.overrides, + "owner": self.owner, } @classmethod @@ -85,6 +87,7 @@ def from_dict(cls, data: dict) -> "RunRecord": finished_at_utc=data.get("finished_at_utc"), error=data.get("error"), overrides=data.get("overrides") or {}, + owner=data.get("owner"), ) diff --git a/src/agent_env/runner/store.py b/src/agent_env/runner/store.py index 5eaf2698..6443af86 100644 --- a/src/agent_env/runner/store.py +++ b/src/agent_env/runner/store.py @@ -26,15 +26,24 @@ def _now() -> str: return datetime.now(timezone.utc).isoformat() +_indexed = None + + def _docs(): - return get_config().get_document_store() + """The configured document store, with the indexes backing the two reads (status-by-run-id and the per-task + listing) made once per store: a caller that never ``start()``s a runner still gets them.""" + global _indexed + docs = get_config().get_document_store() + if _indexed is not docs: + docs.ensure_index(RUNS_COLLECTION, ["run_id"], unique=True) + docs.ensure_index(RUNS_COLLECTION, ["task_id", "created_at_utc"]) + _indexed = docs + return docs def ensure_indexes() -> None: - """Indexes backing the two reads: status-by-run-id and the per-task listing.""" - docs = _docs() - docs.ensure_index(RUNS_COLLECTION, ["run_id"], unique=True) - docs.ensure_index(RUNS_COLLECTION, ["task_id", "created_at_utc"]) + """Make the run index's indexes now rather than on first use.""" + _docs() def insert_run(record: RunRecord) -> RunRecord: @@ -86,6 +95,15 @@ def set_running(run_id: str) -> None: ) +def see(run_id: str) -> None: + """Mark a non-terminal run as still seen by its owner: what keeps an owner on another host counting as alive.""" + _docs().update( + RUNS_COLLECTION, + Filter(conditions={"run_id": [Eq(run_id)], "status": _NON_TERMINAL}), + UpdateSpec(set={"owner.seen": _now()}), + ) + + def mark_terminal(run_id: str, status: RunStatus, error: Optional[str] = None) -> None: """Move a run to a terminal status, but only from QUEUED/RUNNING — so a run that a concurrent cancel already terminated is left as-is (cancel wins, and it's idempotent).""" diff --git a/src/agent_env/task/store.py b/src/agent_env/task/store.py index 707b62c9..b944dbbb 100644 --- a/src/agent_env/task/store.py +++ b/src/agent_env/task/store.py @@ -299,8 +299,7 @@ def create_instance( completed_step_docs = [ dataclasses.asdict(e) for e in completed_steps or [] ] - suffix = "".join(secrets.choice(string.ascii_lowercase + string.digits) for _ in range(8)) - instance_id = f"{task_id}-{suffix}" + instance_id = new_instance_id(task_id) created_at_utc = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC") doc = _new_instance_doc( instance_id, task_id, task_version, total_steps, created_at_utc, @@ -804,6 +803,11 @@ def find_task_instance(instance_id: str) -> TaskInstance | None: return None +def new_instance_id(task_id: str) -> str: + """A fresh id for a run of ``task_id``, derived from it so the run keeps the task's namespace.""" + return f"{task_id}-{''.join(secrets.choice(string.ascii_lowercase + string.digits) for _ in range(8))}" + + def register_task_instance( task_id: str, task_version: int, diff --git a/tst/unit/cli/eval_run_teardown_test.py b/tst/unit/cli/eval_run_teardown_test.py index c9575317..e1025a91 100644 --- a/tst/unit/cli/eval_run_teardown_test.py +++ b/tst/unit/cli/eval_run_teardown_test.py @@ -1,8 +1,7 @@ -"""``agent-env eval run`` tears down each run's sandboxes, however the run ends, and a first Ctrl-C lets a -teardown already under way finish.""" +"""``agent-env eval run``'s runs, through the configured runner (here ``LocalRunner``): a first Ctrl-C cancels the +running ones, which are torn down, and lets a teardown already under way finish.""" import asyncio -import importlib import logging import signal import types @@ -10,43 +9,13 @@ import pytest from click.testing import CliRunner -from agent_env.cli import cli +from agent_env.cli import cli, submitted_runs from agent_env.eval import Eval from agent_env.task import Task, teardown -from agent_env.task.teardown import TeardownReport from agent_env.task_step.context import DeployedSandbox from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider -eval_run = importlib.import_module("agent_env.cli.eval.run") # the package's ``run`` is the click command -pytestmark = pytest.mark.usefixtures("sigint_handled") - - -class _Raising: - id = "t" - - def __init__(self, ending): - self.ending, self.context = ending, None - - async def run(self, context, **_): - self.context = context - raise self.ending - - -@pytest.mark.parametrize("ending", [RuntimeError("boom"), asyncio.CancelledError()], ids=["raised", "cancelled"]) -def test_a_run_that_raises_or_is_cancelled_is_still_torn_down(monkeypatch, ending): - torn_down = [] - - async def teardown_run(context): - torn_down.append(context) - return TeardownReport() - - monkeypatch.setattr(eval_run, "teardown_run", teardown_run) - task = _Raising(ending) - - with pytest.raises(type(ending)): - asyncio.run(eval_run._run_single(task, "[t]", None)) - - assert torn_down == [task.context] +pytestmark = pytest.mark.usefixtures("sigint_handled", "local_stores") class _Deploys: @@ -87,6 +56,7 @@ async def get_sandbox(self, sandbox_id): return _Sandbox(sandbox_id) monkeypatch.setattr(teardown, "build_sandbox_provider", lambda name: _Provider()) + monkeypatch.setattr(submitted_runs, "POLL_SECONDS", 0.02) logging.disable(logging.CRITICAL) try: result = CliRunner().invoke(cli, ["eval", "run", "--id", "e"]) @@ -95,6 +65,6 @@ async def get_sandbox(self, sandbox_id): assert result.exit_code == 130, result.output assert sorted(terminated) == ["sb-done", "sb-slow"] - assert "Cancelling the runs and tearing them down (Ctrl-C again to stop now)" in result.output + assert "Cancelling the runs (Ctrl-C again to stop waiting for them)" in result.output assert " [done] PASSED" in result.output - assert " [slow] CANCELLED" in result.output + assert " [slow] CANCELED" in result.output diff --git a/tst/unit/cli/run_role_models_test.py b/tst/unit/cli/run_role_models_test.py index 5e6d9fe2..d266d2fe 100644 --- a/tst/unit/cli/run_role_models_test.py +++ b/tst/unit/cli/run_role_models_test.py @@ -8,14 +8,14 @@ import pytest from click.testing import CliRunner -from agent_env.cli import cli +from agent_env.cli import cli, submitted_runs from agent_env.eval import Eval +from agent_env.runner import local_runner from agent_env.task import Task from agent_env.task.teardown import TeardownReport task_run = importlib.import_module("agent_env.cli.task.run") # the package's ``run`` is the click command -eval_run = importlib.import_module("agent_env.cli.eval.run") -pytestmark = pytest.mark.usefixtures("sigint_handled") +pytestmark = pytest.mark.usefixtures("sigint_handled", "local_stores") class _Recorded: @@ -37,7 +37,8 @@ async def teardown_run(context): return TeardownReport() monkeypatch.setattr(task_run, "teardown_run", teardown_run) - monkeypatch.setattr(eval_run, "teardown_run", teardown_run) + monkeypatch.setattr(local_runner, "teardown_run", teardown_run) + monkeypatch.setattr(submitted_runs, "POLL_SECONDS", 0.02) recorded = _Recorded() monkeypatch.setattr(Task, "get", classmethod(lambda cls, id, version=None: recorded)) monkeypatch.setattr(Eval, "get", classmethod(lambda cls, id, version=None: types.SimpleNamespace( diff --git a/tst/unit/cli/submit_runs_test.py b/tst/unit/cli/submit_runs_test.py new file mode 100644 index 00000000..115f3702 --- /dev/null +++ b/tst/unit/cli/submit_runs_test.py @@ -0,0 +1,187 @@ +"""``task run`` and ``eval run`` hand their runs to the configured ``[runner]`` and follow them from the run index: +here ``LocalRunner``, which runs them in this process and tears them down itself.""" + +import asyncio +import importlib +import json +import logging +import os +import signal +import threading +import types + +import pytest +from click.testing import CliRunner + +from agent_env.cli import cli, submitted_runs +from agent_env.eval import Eval +from agent_env.runner import local_runner +from agent_env.runner import store as run_store +from agent_env.runner.runner import RunHandle, Runner, RunRecord, RunStatus +from agent_env.task import Task +from agent_env.task.teardown import TeardownReport +from agent_env.task_step.context import TaskStepContext + +eval_run = importlib.import_module("agent_env.cli.eval.run") # the package's ``run`` is the click command + +pytestmark = pytest.mark.usefixtures("sigint_handled", "local_stores") + + +class _Task: + id, version, steps = "t", 1, [] + + def __init__(self, nap=0.0, failed_steps=None): + self.nap, self.failed_steps, self.contexts, self.calls = nap, failed_steps, [], [] + self.active = self.most_active = 0 + + async def run(self, context, **kwargs): + self.contexts.append(context) + self.calls.append(kwargs) + self.active += 1 + self.most_active = max(self.most_active, self.active) + try: + await asyncio.sleep(self.nap) + finally: + self.active -= 1 + if self.failed_steps: + context.metadata["failed_steps"] = self.failed_steps + return context + + +@pytest.fixture +def torn_down(monkeypatch): + torn = [] + + async def teardown_run(context): + torn.append(context) + return TeardownReport() + + monkeypatch.setattr(local_runner, "teardown_run", teardown_run) + monkeypatch.setattr(submitted_runs, "POLL_SECONDS", 0.02) + return torn + + +def _invoke(monkeypatch, task, *args, signal_after=None): + monkeypatch.setattr(Task, "get", classmethod(lambda cls, id, version=None: task)) + timer = threading.Timer(signal_after, lambda: os.kill(os.getpid(), signal.SIGINT)) if signal_after else None + if timer: + timer.start() + logging.disable(logging.CRITICAL) + try: + return CliRunner().invoke(cli, list(args)) + finally: + if timer: + timer.cancel() + logging.disable(logging.NOTSET) + + +def test_a_failed_submitted_run_fails_the_command(monkeypatch, tmp_path, torn_down): + task = _Task(failed_steps=[{"step_id": "s1"}]) + + result = _invoke(monkeypatch, task, "task", "run", "--id", "t", "--k", "2", + "--output-dir", str(tmp_path)) + + assert result.exit_code == 1, result.output + assert "0/2 runs completed, 2/2 failed." in result.output + assert {r.status for r in run_store.list_runs(task_id="t")} == {RunStatus.FAILED} + + +def test_a_context_json_is_stored_as_an_instance_the_runner_continues(monkeypatch, tmp_path, torn_down): + task = _Task() + saved = tmp_path / "context.json" + saved.write_text(json.dumps(TaskStepContext(metadata={"seeded": "yes"}).to_dict())) + + result = _invoke(monkeypatch, task, "task", "run", "--id", "t", "--context-json", str(saved), "--start-step", "1", + "--gateway-env-id", "gw-run", "--output-dir", str(tmp_path)) + + assert result.exit_code == 0, result.output + [call], [context] = task.calls, task.contexts + assert call["start_step"] == 1 and context.metadata["seeded"] == "yes" + assert context.metadata["user_overrides"]["gateway_env_id"] == "gw-run" + assert call["instance_id"] == run_store.list_runs(task_id="t")[0].instance_id + + +def test_an_eval_submits_every_task_k_times(monkeypatch, tmp_path, torn_down): + task = _Task() + monkeypatch.setattr(Eval, "get", classmethod(lambda cls, id, version=None: types.SimpleNamespace( + id="e", version=1, tasks=[types.SimpleNamespace(task_id="t", task_version=1)]))) + + monkeypatch.setattr(eval_run, "get_agent_env_client_id", lambda: "agent-env-cli/alex-dev") + + result = _invoke(monkeypatch, task, "eval", "run", "--id", "e", "--k", "2") + + assert result.exit_code == 0, result.output + assert result.output.count("PASSED") == 2 and "All 2 runs completed!" in result.output + assert len(run_store.list_runs(task_id="t")) == 2 + assert {c.metadata["agent_env_hub"]["caller"] for c in task.contexts} == {"agent-env-cli/alex-dev"} + + +def test_an_eval_keeps_at_most_max_concurrency_runs_in_flight(monkeypatch, torn_down): + task = _Task(nap=0.1) + monkeypatch.setattr(Eval, "get", classmethod(lambda cls, id, version=None: types.SimpleNamespace( + id="e", version=1, tasks=[types.SimpleNamespace(task_id="t", task_version=1)]))) + + result = _invoke(monkeypatch, task, "eval", "run", "--id", "e", "--k", "3", "--max-concurrency", "1") + + assert result.exit_code == 0, result.output + assert (len(task.contexts), task.most_active) == (3, 1) + + +class _Runner(Runner): + """Gives every run ``run_status`` at submit; ``stop()`` naps, then records that it finished.""" + + type = "fake" + + def __init__(self, status_error=None, cancel_error=None, stop_seconds=0.0, run_status=RunStatus.COMPLETED): + self.status_error, self.cancel_error, self.stop_seconds = status_error, cancel_error, stop_seconds + self.run_status, self.stopped, self.records = run_status, False, {} + + async def submit(self, task_id, task_version=None, **_): + record = RunRecord(run_id=f"r{len(self.records)}", runner=self.type, task_id=task_id, + task_version=task_version, instance_id=f"i{len(self.records)}", status=self.run_status) + self.records[record.run_id] = record + return RunHandle(record.run_id, record.instance_id) + + async def status(self, run_id): + if self.status_error: + raise self.status_error + return self.records[run_id] + + async def cancel(self, run_id): + if self.cancel_error: + raise self.cancel_error + return False + + async def stop(self): + await asyncio.sleep(self.stop_seconds) + self.stopped = True + + +def test_an_error_following_the_runs_still_stops_the_runner(monkeypatch, local_stores): + runner = _Runner(status_error=RuntimeError("run index unreachable")) + local_stores.set_runner(runner) + + result = _invoke(monkeypatch, _Task(), "task", "run", "--id", "t") + + assert isinstance(result.exception, RuntimeError) + assert runner.stopped + + +def test_a_first_ctrl_c_lets_the_runner_finish_stopping(monkeypatch, local_stores): + runner = _Runner(stop_seconds=2.0) + local_stores.set_runner(runner) + + result = _invoke(monkeypatch, _Task(), "task", "run", "--id", "t", signal_after=0.5) + + assert result.exit_code == 130, result.output + assert runner.stopped + + +def test_an_error_cancelling_after_ctrl_c_still_stops_the_runner(monkeypatch, local_stores): + runner = _Runner(cancel_error=RuntimeError("runner unreachable"), run_status=RunStatus.RUNNING) + local_stores.set_runner(runner) + + result = _invoke(monkeypatch, _Task(), "task", "run", "--id", "t", signal_after=0.5) + + assert isinstance(result.exception, RuntimeError) + assert runner.stopped diff --git a/tst/unit/cli/task_run_stored_test.py b/tst/unit/cli/task_run_stored_test.py new file mode 100644 index 00000000..cf2a5421 --- /dev/null +++ b/tst/unit/cli/task_run_stored_test.py @@ -0,0 +1,133 @@ +"""``task run`` on stored tasks, through ``LocalRunner`` and the real stores: an ``@local`` task's run stays in the +``@local`` namespace, seeded steps report under the ids they ran as, and a context that can't be exported fails +the command.""" + +import json +import logging + +import pytest +from click.testing import CliRunner + +from agent_env.cli import cli, submitted_runs +from agent_env.config import get_config +from agent_env.config.runtime import Config +from agent_env.runner import store as run_store +from agent_env.runner.store import RUNS_COLLECTION +from agent_env.store.document_store import Filter +from agent_env.task import Task +from agent_env.task.store import TASK_INSTANCES_COLLECTION, TASK_STEP_JOURNAL_COLLECTION, TaskInstanceStore +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_step import TaskStep + +pytestmark = pytest.mark.usefixtures("sigint_handled", "local_stores", "cli_routing") + +RUN_COLLECTIONS = (TASK_INSTANCES_COLLECTION, TASK_STEP_JOURNAL_COLLECTION, RUNS_COLLECTION) + + +class _Writes(TaskStep): + """Makes the check ``collect_artifacts`` makes on an id it derives from the run's instance, records a score under + its ``verifier_id``, and fails when told to.""" + + type = "writes_stored_run_test" + entity_refs = () + + def __init__(self, verifier_id=None, fail=False, **base): + super().__init__(**base) + self.verifier_id, self.fail = verifier_id, fail + + @classmethod + def from_dict(cls, data): + return cls(**cls._base_from_dict(data), verifier_id=data.get("verifier_id"), fail=data.get("fail", False)) + + def to_dict(self): + return {**super().to_dict(), "verifier_id": self.verifier_id, "fail": self.fail} + + async def execute(self, context): + get_config().check_local_run_write(context.instance_id) + if self.verifier_id: + context.metadata.setdefault("verifications", {})[self.verifier_id] = {"score": 1.0} + if self.fail: + raise RuntimeError("boom") + return context + + +@pytest.fixture(autouse=True) +def registries(monkeypatch): + steps = {**Config().task_step_registry(), _Writes.type: _Writes} + monkeypatch.setattr(Config, "task_step_registry", lambda self: steps) + monkeypatch.setattr(submitted_runs, "POLL_SECONDS", 0.02) + + +def _put(task_id, *steps): + return Task.put(id=task_id, steps=list(steps)) + + +def _invoke(*args): + logging.disable(logging.CRITICAL) + try: + return CliRunner().invoke(cli, ["task", "run", *args]) + finally: + logging.disable(logging.NOTSET) + + +def _context_file(tmp_path, **metadata): + path = tmp_path / "context.json" + path.write_text(json.dumps(TaskStepContext(metadata=metadata).to_dict())) + return str(path) + + +def _configured_run_records(): + configured = get_config().get_document_store().configured + return {collection: configured.query(collection, Filter()) for collection in RUN_COLLECTIONS if + configured.query(collection, Filter())} + + +@pytest.mark.parametrize("continued", [False, True], ids=["fresh", "context-json"]) +def test_an_at_local_tasks_run_stays_in_the_local_namespace(tmp_path, continued): + _put("@local/demo/t", _Writes(id="first", version=None), _Writes(id="second", version=None)) + resume = ["--context-json", _context_file(tmp_path, from_file="yes"), "--start-step", "1"] if continued else [] + + result = _invoke("--id", "@local/demo/t", "--output-dir", str(tmp_path / "out"), *resume) + + assert result.exit_code == 0, result.output + assert _configured_run_records() == {} + [instance] = get_config().local_namespace_document_store().query(TASK_INSTANCES_COLLECTION, Filter()) + assert instance["instance_id"].startswith("@local/demo/t-") and instance["status"] == "completed" + assert instance["context"]["metadata"].get("from_file") == ("yes" if continued else None) + + +@pytest.mark.parametrize("k", ["1", "2"]) +def test_a_seeded_step_reports_under_the_id_it_ran_as(tmp_path, k): + _put("seeded", _Writes(id="verify-", version=None, verifier_id="grade-")) + + result = _invoke("--id", "seeded", "--k", k, "--output-dir", str(tmp_path / "out"), + "--context-json", _context_file(tmp_path, seed={"week": "3"})) + + assert result.exit_code == 0, result.output + assert "Completed step [1/1]: writes_stored_run_test (id=verify-3)" in result.output + assert ("grade-3" if k == "1" else "score (verifier_id=grade-3): 1.0") in result.output + + +@pytest.mark.parametrize("fail", [False, True], ids=["passed", "failed"]) +def test_a_context_that_cant_be_exported_fails_the_command(tmp_path, fail): + _put("exported", _Writes(id="only", version=None, fail=fail)) + (tmp_path / "taken").write_text("") + + result = _invoke("--id", "exported", "--output-dir", str(tmp_path / "taken" / "out")) + + assert result.exit_code == 1, result.output + assert "Couldn't write the run's context" in result.output and "Task completed!" not in result.output + assert ("Run FAILED: RuntimeError: boom" in result.output) == fail + + +def test_a_context_json_that_cant_be_stored_fails_before_anything_is_submitted(tmp_path, monkeypatch): + def unwritable(*args, **kwargs): + raise ConnectionError("store unreachable") + + _put("resumed", _Writes(id="only", version=None)) + monkeypatch.setattr(TaskInstanceStore, "seed_context", unwritable) + + result = _invoke("--id", "resumed", "--context-json", _context_file(tmp_path, a=1), "--output-dir", str(tmp_path)) + + assert isinstance(result.exception, ConnectionError), result.output + assert run_store.list_runs(task_id="resumed") == [] diff --git a/tst/unit/cli/task_run_teardown_test.py b/tst/unit/cli/task_run_teardown_test.py index cde31098..24381292 100644 --- a/tst/unit/cli/task_run_teardown_test.py +++ b/tst/unit/cli/task_run_teardown_test.py @@ -1,5 +1,6 @@ """``agent-env task run`` and ``task run-batch`` tear down what each run deployed, however it ends; ``--keep`` holds -it up until Ctrl-C, and Ctrl-C mid-run tears it down and exits 130.""" +it up until Ctrl-C, and Ctrl-C mid-run tears it down and exits 130. ``task run`` goes through the configured runner, +here ``LocalRunner``: it tears a run down, and the CLI tears down what ``--keep`` held, from the stored instance.""" import asyncio import importlib @@ -7,17 +8,22 @@ import os import signal import threading +import types import pytest from click.testing import CliRunner -from agent_env.cli import cli +from agent_env.cli import cli, submitted_runs +from agent_env.runner import local_runner +from agent_env.runner import store as run_store +from agent_env.runner.runner import RunStatus from agent_env.task import Task +from agent_env.task.store import register_task_instance, seed_task_instance_context from agent_env.task.teardown import TORN_DOWN_KEY, RecordedSandbox, TeardownReport from agent_env.task_step.context import DeployedSandbox task_run = importlib.import_module("agent_env.cli.task.run") # the package's ``run`` is the click command -pytestmark = pytest.mark.usefixtures("sigint_handled") +pytestmark = pytest.mark.usefixtures("sigint_handled", "local_stores") class _Deploys: @@ -28,10 +34,13 @@ class _Deploys: def __init__(self, nap=0.0, ending=None): self.nap, self.ending, self.contexts = nap, ending, [] - async def run(self, context, **_): + async def run(self, context, instance_id=None, **_): self.contexts.append(context) context.deployed_sandboxes.append(DeployedSandbox( sandbox_name="box", sandbox_id=f"sb-{len(self.contexts)}", sandbox_mode="vm", sandbox_type="fake")) + if instance_id: + register_task_instance("t", 1, 1, instance_id=instance_id) + seed_task_instance_context(instance_id, context) await asyncio.sleep(self.nap) if self.ending: raise self.ending @@ -49,9 +58,15 @@ async def teardown_run(context): # records what it took down, as the real one d RecordedSandbox(s.sandbox_id, s.sandbox_type, "sandbox") for s in context.deployed_sandboxes)) monkeypatch.setattr(task_run, "teardown_run", teardown_run) + monkeypatch.setattr(local_runner, "teardown_run", teardown_run) + monkeypatch.setattr(submitted_runs, "POLL_SECONDS", 0.02) return torn +def _ids(contexts) -> list[str]: + return sorted(s.sandbox_id for context in contexts for s in context.deployed_sandboxes) + + def _invoke(monkeypatch, tmp_path, task, *args, signal_after=None): """Invoke the CLI, sending Ctrl-C ``signal_after`` seconds in; a Ctrl-C not yet sent when it returns never is.""" monkeypatch.setattr(Task, "get", classmethod(lambda cls, id, version=None: task)) @@ -73,18 +88,19 @@ def test_a_run_that_passes_is_torn_down(monkeypatch, tmp_path, torn_down): result = _invoke(monkeypatch, tmp_path, task, "run") assert result.exit_code == 0, result.output - assert torn_down == task.contexts - assert "Tore down 1 sandbox" in result.output + assert _ids(torn_down) == ["sb-1"] + [record] = run_store.list_runs(task_id="t") + assert record.status is RunStatus.COMPLETED and task.contexts[0].metadata["workflow_id"] == record.run_id -def test_a_run_that_fails_is_torn_down_and_reported_as_before(monkeypatch, tmp_path, torn_down): +def test_a_run_that_fails_is_torn_down_and_fails_the_command(monkeypatch, tmp_path, torn_down): task = _Deploys(ending=RuntimeError("boom")) result = _invoke(monkeypatch, tmp_path, task, "run") - assert result.exit_code != 0 - assert isinstance(result.exception, RuntimeError) - assert torn_down == task.contexts + assert result.exit_code == 1, result.output + assert "FAILED: RuntimeError: boom" in result.output + assert _ids(torn_down) == ["sb-1"] def test_each_parallel_run_is_torn_down(monkeypatch, tmp_path, torn_down): @@ -93,7 +109,7 @@ def test_each_parallel_run_is_torn_down(monkeypatch, tmp_path, torn_down): result = _invoke(monkeypatch, tmp_path, task, "run", "--k", "2") assert result.exit_code == 0, result.output - assert len(torn_down) == 2 and {id(c) for c in torn_down} == {id(c) for c in task.contexts} + assert _ids(torn_down) == ["sb-1", "sb-2"] def test_keep_holds_a_run_up_until_ctrl_c_then_tears_it_down(monkeypatch, tmp_path, torn_down): @@ -104,7 +120,7 @@ def test_keep_holds_a_run_up_until_ctrl_c_then_tears_it_down(monkeypatch, tmp_pa assert result.exit_code == 0, result.output assert "Kept up:" in result.output and "sb-1" in result.output assert "Holding 1 sandbox up; Ctrl-C tears it down." in result.output - assert torn_down == task.contexts + assert _ids(torn_down) == ["sb-1"] def test_keep_still_tears_down_a_run_ctrl_c_cancels(monkeypatch, tmp_path, torn_down): @@ -113,7 +129,7 @@ def test_keep_still_tears_down_a_run_ctrl_c_cancels(monkeypatch, tmp_path, torn_ result = _invoke(monkeypatch, tmp_path, task, "run", "--keep", signal_after=1.0) assert result.exit_code == 130, result.output - assert torn_down == task.contexts + assert _ids(torn_down) == ["sb-1"] assert "Holding" not in result.output @@ -123,9 +139,10 @@ def test_ctrl_c_mid_run_tears_down_every_run_and_exits_130(monkeypatch, tmp_path result = _invoke(monkeypatch, tmp_path, task, "run", "--k", "2", signal_after=1.0) assert result.exit_code == 130, result.output - assert len(torn_down) == 2 - assert "Cancelling and tearing down (Ctrl-C again to stop now)" in result.output - assert result.output.count("CANCELLED") == 2 + assert _ids(torn_down) == ["sb-1", "sb-2"] + assert "Cancelling the runs (Ctrl-C again to stop waiting for them)" in result.output + assert result.output.count("Run CANCELED") == 2 + assert {r.status for r in run_store.list_runs(task_id="t")} == {RunStatus.CANCELED} def test_each_seed_of_a_batch_is_torn_down(monkeypatch, tmp_path, torn_down): @@ -140,10 +157,14 @@ def test_each_seed_of_a_batch_is_torn_down(monkeypatch, tmp_path, torn_down): class _OneQuickOneSlow(_Deploys): - """Its first run ends at once; every later one naps.""" + """Its first run ends at once and every later one naps, or the other way round when ``quick_last``.""" + + def __init__(self, quick_last=False): + super().__init__() + self.quick_last = quick_last async def run(self, context, **kwargs): - self.nap = 0 if not self.contexts else 30 + self.nap = 0 if bool(self.contexts) == self.quick_last else 30 return await super().run(context, **kwargs) @@ -153,16 +174,70 @@ def test_ctrl_c_with_keep_also_tears_down_a_run_that_already_finished(monkeypatc result = _invoke(monkeypatch, tmp_path, task, "run", "--k", "2", "--keep", signal_after=1.0) assert result.exit_code == 130, result.output - assert {id(c) for c in torn_down} == {id(c) for c in task.contexts} + assert _ids(torn_down) == ["sb-1", "sb-2"] assert "Holding" not in result.output +class _CancelledElsewhere(_Deploys): + """Its first run is asked to cancel through another runner while it runs, then ends; the second ends only once + the first has.""" + + def __init__(self, other): + super().__init__() + self.other, self.first_ended, self.refused = other, asyncio.Event(), [] + + async def run(self, context, **kwargs): + first = not self.contexts + await super().run(context, **kwargs) + if first: + self.refused.append(await self.other.cancel(context.metadata["workflow_id"])) + self.first_ended.set() + else: + await self.first_ended.wait() + return context + + +def test_a_cancel_another_runner_cant_carry_out_leaves_a_kept_run_kept_and_torn_down(monkeypatch, tmp_path, torn_down): + def sleep(_): + signal.raise_signal(signal.SIGINT) + + monkeypatch.setattr(task_run, "time", types.SimpleNamespace(sleep=sleep)) + task = _CancelledElsewhere(local_runner.LocalRunner()) + + result = _invoke(monkeypatch, tmp_path, task, "run", "--k", "2", "--keep") + + assert result.exit_code == 0, result.output + assert task.refused == [False] + assert {r.status for r in run_store.list_runs(task_id="t")} == {RunStatus.COMPLETED} + assert "Holding 2 sandboxes up" in result.output and _ids(torn_down) == ["sb-1", "sb-2"] + + +@pytest.mark.parametrize("quick_last", [False, True], ids=["ended-first", "ended-last"]) +def test_an_error_following_kept_runs_still_tears_down_one_that_ended(monkeypatch, tmp_path, torn_down, quick_last): + real_status = local_runner.LocalRunner.status + + async def status(self, run_id): + record = await real_status(self, run_id) + if record.status is not RunStatus.COMPLETED and any( + run.status is RunStatus.COMPLETED for run in run_store.list_runs()): + raise RuntimeError("run index unreachable") + return record + + monkeypatch.setattr(local_runner.LocalRunner, "status", status) + task = _OneQuickOneSlow(quick_last) + + result = _invoke(monkeypatch, tmp_path, task, "run", "--k", "2", "--keep") + + assert isinstance(result.exception, RuntimeError), result.output + assert _ids(torn_down) == ["sb-1", "sb-2"] + + def test_keep_writes_the_context_of_a_run_that_failed(monkeypatch, tmp_path, torn_down): task = _Deploys(ending=RuntimeError("boom")) result = _invoke(monkeypatch, tmp_path, task, "run", "--keep", signal_after=1.0) - assert isinstance(result.exception, RuntimeError) + assert result.exit_code == 1, result.output assert list(tmp_path.glob("t_*.json")), "no context file to resume the kept run from" @@ -175,6 +250,6 @@ def disk_full(*_, **__): result = _invoke(monkeypatch, tmp_path, task, "run", "--keep", signal_after=1.0) - assert isinstance(result.exception, RuntimeError) and str(result.exception) == "boom" + assert result.exit_code == 1 and "FAILED: RuntimeError: boom" in result.output assert "Couldn't write the run's context: [Errno 28] No space left on device" in result.output - assert torn_down == task.contexts # the Ctrl-C that ended the hold + assert _ids(torn_down) == ["sb-1"] diff --git a/tst/unit/explorer/explorer_app_test.py b/tst/unit/explorer/explorer_app_test.py index b8d33e7f..86f55b9e 100644 --- a/tst/unit/explorer/explorer_app_test.py +++ b/tst/unit/explorer/explorer_app_test.py @@ -319,8 +319,9 @@ async def run(self, **kw): "total_steps": 1, "current_step": 1, "completed_steps": [{"step_id": "s", "status": "success"}], }) assert single["instance_id"] in {i["instance_id"] for i in client.get(f"{base}/instances").json()["items"]} - assert client.get(f"{base}/instances/{single['instance_id']}").json()["task_id"] == task_id - assert client.get(f"{base}/instances/{single['instance_id']}/progress").json()["instance_id"] == single["instance_id"] + instance = f"{base}/instances/{quote(single['instance_id'], safe='')}" + assert client.get(instance).json()["task_id"] == task_id + assert client.get(f"{instance}/progress").json()["instance_id"] == single["instance_id"] assert group["run_group_id"] in {g["run_group_id"] for g in client.get(f"{base}/run-groups").json()["items"]} assert client.get(f"{base}/run-groups/{group['run_group_id']}").json()["task_id"] == task_id assert "event: complete" in client.get(f"{base}/run-groups/{group['run_group_id']}/stream").text diff --git a/tst/unit/runner/local_runner_test.py b/tst/unit/runner/local_runner_test.py index 371cbd5a..59d8219a 100644 --- a/tst/unit/runner/local_runner_test.py +++ b/tst/unit/runner/local_runner_test.py @@ -5,6 +5,10 @@ """ import asyncio +import socket +import subprocess +import sys +from datetime import datetime, timedelta, timezone import pytest @@ -266,6 +270,120 @@ async def test_start_fails_a_run_left_running_by_a_previous_process(docs): assert "restart" in rec.error +@pytest.mark.asyncio +async def test_start_spares_the_runs_a_live_runner_owns_and_fails_a_dead_processes(docs, monkeypatch): + started = asyncio.Event() + + async def on_run(_): + started.set() + await asyncio.sleep(30) + + _install_task(monkeypatch, _FakeTask(on_run=on_run)) + gone = subprocess.Popen([sys.executable, "-c", "pass"]) + gone.wait() + orphan = RunRecord(run_id="orphan", runner="local", task_id="t", task_version=1, instance_id="i", + status=RunStatus.RUNNING) + docs.insert(run_store.RUNS_COLLECTION, {**orphan.to_dict(), "owner": {"host": socket.gethostname(), "pid": gone.pid}}) + owner = LocalRunner(workers=1) + handle = await owner.submit("t1", 1) + await asyncio.wait_for(started.wait(), 5) + explorer = LocalRunner(workers=1) + try: + await explorer.start() + assert (await explorer.status(handle.run_id)).status is RunStatus.RUNNING + assert run_store.get_run("orphan").status is RunStatus.FAILED + finally: + await explorer.stop() + await owner.stop() + + +@pytest.mark.asyncio +async def test_only_the_runner_running_a_run_can_cancel_it(docs, monkeypatch): + from agent_env.runner import local_runner + from agent_env.task.teardown import TeardownReport + + torn_down, started, release, cancelled = [], asyncio.Event(), asyncio.Event(), asyncio.Event() + + async def teardown_run(context): + torn_down.append(context.metadata["workflow_id"]) + return TeardownReport() + + async def on_run(kwargs): + started.set() + try: + await release.wait() + except asyncio.CancelledError: + cancelled.set() + raise + + monkeypatch.setattr(local_runner, "teardown_run", teardown_run) + _install_task(monkeypatch, _FakeTask(on_run=on_run)) + owner, other = LocalRunner(workers=2), LocalRunner(workers=1) + finishes, stopped = await owner.submit("t1", 1), await owner.submit("t1", 1) + try: + await asyncio.wait_for(started.wait(), 5) + assert await other.cancel(finishes.run_id) is False + assert run_store.get_run(finishes.run_id).status is RunStatus.RUNNING and not cancelled.is_set() + assert await owner.cancel(stopped.run_id) is True + assert (await _drain(owner, stopped.run_id)).status is RunStatus.CANCELED and cancelled.is_set() + release.set() + assert (await _drain(owner, finishes.run_id)).status is RunStatus.COMPLETED + await _settled(owner, finishes.run_id) + finally: + await other.stop() + await owner.stop() + + assert sorted(torn_down) == sorted([finishes.run_id, stopped.run_id]) + + +async def _until(condition, timeout=5.0): + deadline = asyncio.get_event_loop().time() + timeout + while not condition(): + if asyncio.get_event_loop().time() > deadline: + raise AssertionError("condition never held") + await asyncio.sleep(0.01) + + +@pytest.mark.asyncio +async def test_a_run_owned_on_another_host_is_failed_once_its_owner_stops_seeing_it(docs, monkeypatch): + monkeypatch.setattr(LocalRunner, "OWNER_SEEN_SECONDS", 0.01) + monkeypatch.setattr(LocalRunner, "OWNER_LEASE_SECONDS", 0.5) + now = datetime.now(timezone.utc) + for run_id, seen in (("gone", now - timedelta(seconds=10)), ("going", now)): + record = RunRecord(run_id=run_id, runner="local", task_id="t", task_version=1, instance_id=run_id, + status=RunStatus.RUNNING) + docs.insert(run_store.RUNS_COLLECTION, {**record.to_dict(), "owner": {"host": "elsewhere", "pid": 1, + "seen": seen.isoformat()}}) + runner = LocalRunner(workers=1) + try: + await runner.start() + assert (run_store.get_run("gone").status, run_store.get_run("going").status) == (RunStatus.FAILED, + RunStatus.RUNNING) + await _until(lambda: run_store.get_run("going").status is RunStatus.FAILED) + finally: + await runner.stop() + + +@pytest.mark.asyncio +async def test_a_runner_keeps_seeing_the_runs_it_is_running(docs, monkeypatch): + monkeypatch.setattr(LocalRunner, "OWNER_SEEN_SECONDS", 0.01) + started = asyncio.Event() + + async def on_run(_): + started.set() + await asyncio.sleep(30) + + _install_task(monkeypatch, _FakeTask(on_run=on_run)) + runner = LocalRunner(workers=1) + handle = await runner.submit("t1", 1) + try: + await asyncio.wait_for(started.wait(), 5) + first_seen = run_store.get_run(handle.run_id).owner["seen"] + await _until(lambda: run_store.get_run(handle.run_id).owner["seen"] > first_seen) + finally: + await runner.stop() + + @pytest.mark.asyncio async def test_start_reconciles_an_orphan_behind_many_terminal_runs(docs): # Reconcile is by status, not a window of the newest records: an orphan older than a diff --git a/tst/unit/test_client_identity.py b/tst/unit/test_client_identity.py index a339438f..e2c146a0 100644 --- a/tst/unit/test_client_identity.py +++ b/tst/unit/test_client_identity.py @@ -1,4 +1,3 @@ -import asyncio import importlib.util import subprocess import sys @@ -126,73 +125,3 @@ def test_task_cli_stamp_client_metadata_noops_without_client_id(monkeypatch): task_run_module._stamp_agent_env_client_metadata(metadata, None) assert metadata == {} - - -def test_eval_run_single_passes_client_metadata_to_task(monkeypatch): - task_step_package = types.ModuleType("agent_env.task_step") - task_step_package.__path__ = [] - context_module = types.ModuleType("agent_env.task_step.context") - - class FakeTaskStepContext: - def __init__(self, metadata=None): - self.metadata = metadata or {} - - context_module.TaskStepContext = FakeTaskStepContext - task_package = types.ModuleType("agent_env.task") - task_package.__path__ = [] - teardown_module = types.ModuleType("agent_env.task.teardown") - torn_down = [] - - async def teardown_run(context): - torn_down.append(context) - return teardown_module.TeardownReport() - - teardown_module.TeardownReport = lambda: types.SimpleNamespace(terminated=(), failed=(), left=()) - teardown_module.teardown_run = teardown_run - teardown_module.kind = lambda sandbox: sandbox.sandbox_type - interrupts_module = types.ModuleType("agent_env.task.interrupts") - interrupts_module.Interrupts = object - monkeypatch.setitem(sys.modules, "agent_env.task_step", task_step_package) - monkeypatch.setitem(sys.modules, "agent_env.task_step.context", context_module) - monkeypatch.setitem(sys.modules, "agent_env.task", task_package) - monkeypatch.setitem(sys.modules, "agent_env.task.teardown", teardown_module) - monkeypatch.setitem(sys.modules, "agent_env.task.interrupts", interrupts_module) - eval_run_module = _load_cli_run_module( - monkeypatch, - "src/agent_env/cli/eval/run.py", - "_agent_env_eval_run_under_test", - ) - - class FakeTask: - id = "fake-task" - - def __init__(self): - self.context = None - - async def run(self, **kwargs): - self.context = kwargs["context"] - return self.context - - task = FakeTask() - monkeypatch.setattr( - subprocess, - "run", - lambda args, **kwargs: _completed(args, stdout="Alex Dev\n"), - ) - - task_id, tag, context = asyncio.run( - eval_run_module._run_single( - task, - "[fake-task]", - None, - base_metadata={"agent_env_hub": {"caller": get_agent_env_client_id()}}, - ) - ) - - assert task_id == "fake-task" - assert tag == "[fake-task]" - assert context is task.context - assert context.metadata == { - "agent_env_hub": {"caller": "agent-env-cli/alex-dev"} - } - assert torn_down == [context]