diff --git a/src/agent_env/runner/__init__.py b/src/agent_env/runner/__init__.py index cf24fddc..2356abdc 100644 --- a/src/agent_env/runner/__init__.py +++ b/src/agent_env/runner/__init__.py @@ -1,8 +1,8 @@ """Run dispatch: the seam between "run this task" and ``Task.run()``.""" -from agent_env.runner.runner import RunHandle, RunRecord, Runner, RunStatus +from agent_env.runner.runner import KEEP_KEY, RunHandle, RunRecord, Runner, RunStatus -__all__ = ["RunHandle", "RunRecord", "Runner", "RunStatus", "LocalRunner", "get_runner"] +__all__ = ["KEEP_KEY", "RunHandle", "RunRecord", "Runner", "RunStatus", "LocalRunner", "get_runner"] def __getattr__(name: str): diff --git a/src/agent_env/runner/local_runner.py b/src/agent_env/runner/local_runner.py index 434e1a70..842f1012 100644 --- a/src/agent_env/runner/local_runner.py +++ b/src/agent_env/runner/local_runner.py @@ -17,7 +17,7 @@ from typing import Optional from agent_env.runner import store as run_store -from agent_env.runner.runner import RunHandle, RunRecord, Runner, RunStatus +from agent_env.runner.runner import KEEP_KEY, RunHandle, RunRecord, Runner, RunStatus from agent_env.task.teardown import TERMINATE_TIMEOUT_SECONDS, teardown_run logger = logging.getLogger(__name__) @@ -111,9 +111,10 @@ async def submit( agent_model: Optional[str] = None, agent_artifact_id: Optional[str] = None, metadata: Optional[dict] = None, + instance_id: Optional[str] = None, ) -> RunHandle: run_id = f"local-{uuid.uuid4().hex}" - instance_id = uuid.uuid4().hex # minted here so callers can link immediately + instance_id = instance_id or uuid.uuid4().hex # 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, @@ -153,6 +154,7 @@ async def cancel(self, run_id: str) -> bool: async def _run(self, record: RunRecord) -> None: from agent_env.task import Task + from agent_env.task.store import find_task_instance run_task: Optional[asyncio.Task] = None context = None @@ -162,11 +164,16 @@ async def _run(self, record: RunRecord) -> None: return # canceled while it waited for a slot run_store.set_running(record.run_id) logger.info("Run %s starting (task=%s v%s)", record.run_id, record.task_id, record.task_version) - task = Task.get(record.task_id, record.task_version) + stored = find_task_instance(record.instance_id) + version = record.task_version if record.task_version is not None or stored is None else stored.task_version + if stored is not None and (stored.task_id, stored.task_version) != (record.task_id, version): + raise ValueError(f"Instance {record.instance_id} is a run of {stored.task_id} v{stored.task_version}, " + f"not of {record.task_id} v{version}") + task = Task.get(record.task_id, version) if task is None: - raise LookupError(f"Task {record.task_id} v{record.task_version} not found") + raise LookupError(f"Task {record.task_id} v{version} not found") - context = self._seed_context(record) + context = self._seed_context(record, stored) self._contexts[record.run_id] = context # start_step rides in metadata but Task.run takes it as a keyword; # without lifting it out here every resume re-ran from zero. @@ -202,7 +209,9 @@ async def _run(self, record: RunRecord) -> None: run_store.mark_terminal(record.run_id, RunStatus.FAILED, error=f"{type(e).__name__}: {e}") finally: try: - if context is not None and record.run_id not in self._abandoned: + kept = (context is not None and context.metadata.get(KEEP_KEY) and run_task is not None + and run_task.done() and not run_task.cancelled()) + if context is not None and record.run_id not in self._abandoned and not kept: self._tearing_down.add(record.run_id) try: async with asyncio.timeout(self.TEARDOWN_WAIT_SECONDS): @@ -243,15 +252,16 @@ def _finish(self, run_id: str, context, *, exc: Optional[BaseException] = None) run_store.mark_terminal(run_id, RunStatus.COMPLETED) logger.info("Run %s completed", run_id) - def _seed_context(self, record: RunRecord): - """Build the context ``Task.run()`` mutates, carrying the run identity in metadata. + def _seed_context(self, record: RunRecord, stored): + """Build the context ``Task.run()`` mutates, from the ``stored`` instance it continues if any, carrying the run + identity in metadata. ``workflow_id`` holds the run id — the hub, instance documents and UI all key off it — so nothing downstream needs a special case for the local runner. """ from agent_env.task_step import TaskStepContext - context = TaskStepContext() + context = TaskStepContext.from_dict(stored.context) if stored and stored.context else TaskStepContext() context.metadata.update(record.overrides.get("metadata") or {}) context.metadata["workflow_id"] = record.run_id context.metadata["instance_id"] = record.instance_id diff --git a/src/agent_env/runner/runner.py b/src/agent_env/runner/runner.py index 1771e68f..761aea96 100644 --- a/src/agent_env/runner/runner.py +++ b/src/agent_env/runner/runner.py @@ -12,6 +12,8 @@ from enum import StrEnum from typing import Any, Optional +KEEP_KEY = "keep" + class RunStatus(StrEnum): """Run lifecycle, engine-neutral: a backend maps its own states onto these. ``QUEUED`` @@ -109,8 +111,13 @@ async def submit( agent_model: Optional[str] = None, agent_artifact_id: Optional[str] = None, metadata: Optional[dict] = None, + instance_id: Optional[str] = None, ) -> RunHandle: - """Enqueue a run and return its handle immediately, without waiting for it.""" + """Enqueue a run and return its handle immediately, without waiting for it. + + With ``instance_id`` the run continues that instance from its stored context, at ``metadata["start_step"]``. + A run whose ``metadata[KEEP_KEY]`` is set and that ends on its own, passed or failed, is left up: whoever + submitted it tears it down. A cancelled run is torn down either way.""" @abstractmethod async def status(self, run_id: str) -> Optional[RunRecord]: diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index 8a73496d..7783a42d 100644 --- a/src/agent_env/task/task.py +++ b/src/agent_env/task/task.py @@ -471,6 +471,10 @@ async def _run( context.agent_model = agent_model if agent_artifact_id is not None: context.agent_artifact_id = agent_artifact_id + rerun = {step.id for step in steps[start_step:end_step]} + if context.metadata.get("failed_steps"): + context.metadata["failed_steps"] = [ + f for f in context.metadata["failed_steps"] if not (isinstance(f, dict) and f.get("step_id") in rerun)] # Pre-format a bracketed prefix for pass@k-correlated log lines. # agent-env's log pipeline doesn't index ad-hoc `key=value` fields, so diff --git a/tst/unit/runner/local_runner_test.py b/tst/unit/runner/local_runner_test.py index a25541e7..371cbd5a 100644 --- a/tst/unit/runner/local_runner_test.py +++ b/tst/unit/runner/local_runner_test.py @@ -11,7 +11,7 @@ from agent_env.config import configure, reset_config, set_document_store from agent_env.runner import store as run_store from agent_env.runner.local_runner import LocalRunner -from agent_env.runner.runner import RunRecord, RunStatus +from agent_env.runner.runner import KEEP_KEY, RunRecord, RunStatus from agent_env.store.document_store.sqlite_document_store import LocalSqliteDocumentStore @@ -328,6 +328,81 @@ async def on_run(_): assert torn_down == [handle.run_id] +@pytest.mark.asyncio +@pytest.mark.parametrize(("ending", "torn_down_count"), [("completes", 0), ("is cancelled", 1)]) +async def test_a_kept_run_is_left_up_unless_it_is_cancelled(docs, monkeypatch, ending, torn_down_count): + from agent_env.runner import local_runner + from agent_env.task.teardown import TeardownReport + + torn_down = [] + + async def teardown_run(context): + torn_down.append(context) + return TeardownReport() + + async def on_run(_): + started.set() + if ending == "is cancelled": + await asyncio.sleep(30) + + monkeypatch.setattr(local_runner, "teardown_run", teardown_run) + started = asyncio.Event() + _install_task(monkeypatch, _FakeTask(on_run=on_run)) + runner = LocalRunner(workers=1) + handle = await runner.submit("t1", 1, metadata={KEEP_KEY: True}) + try: + if ending == "is cancelled": + await asyncio.wait_for(started.wait(), 5) + await runner.cancel(handle.run_id) + await _drain(runner, handle.run_id) + await _settled(runner, handle.run_id) + finally: + await runner.stop() + + assert len(torn_down) == torn_down_count + + +@pytest.mark.asyncio +async def test_a_run_given_an_instance_continues_it_from_its_stored_context(docs, monkeypatch): + from agent_env.task.store import register_task_instance, seed_task_instance_context + from agent_env.task_step.context import TaskStepContext + + from agent_env.task import Task + + register_task_instance("t1", 3, 2, 1, instance_id="inst-1") + seed_task_instance_context("inst-1", TaskStepContext(metadata={"done": "s0"})) + task, versions = _FakeTask(), [] + monkeypatch.setattr(Task, "get", classmethod(lambda cls, tid, ver=None: versions.append(ver) or task)) + runner = LocalRunner(workers=1) + handle = await runner.submit("t1", metadata={"start_step": 1}, instance_id="inst-1") + try: + await _drain(runner, handle.run_id) + finally: + await runner.stop() + + [call] = task.calls + assert handle.instance_id == call["instance_id"] == "inst-1" + assert (call["start_step"], call["context"].metadata["done"], versions) == (1, "s0", [3]) + + +@pytest.mark.asyncio +async def test_a_run_given_an_instance_of_another_version_fails_without_running(docs, monkeypatch): + from agent_env.task.store import register_task_instance + + register_task_instance("t1", 3, 2, 1, instance_id="inst-1") + task = _FakeTask() + _install_task(monkeypatch, task) + runner = LocalRunner(workers=1) + handle = await runner.submit("t1", 4, instance_id="inst-1") + try: + record = await _drain(runner, handle.run_id) + finally: + await runner.stop() + + assert record.status is RunStatus.FAILED and "is a run of t1 v3, not of t1 v4" in record.error + assert task.calls == [] + + @pytest.mark.asyncio async def test_stop_lets_a_teardown_under_way_finish(docs, monkeypatch): from agent_env.runner import local_runner diff --git a/tst/unit/task_step/retry_test.py b/tst/unit/task_step/retry_test.py index 4d7341d2..b57fe692 100644 --- a/tst/unit/task_step/retry_test.py +++ b/tst/unit/task_step/retry_test.py @@ -744,6 +744,13 @@ async def test_recovered_failure_is_marked_retried(store): assert all(f.get("retried") for f in fs) +@pytest.mark.asyncio +async def test_a_run_drops_the_recorded_failures_of_the_steps_it_runs_again(store): + context = TaskStepContext(metadata={"failed_steps": [{"step_id": "s0", "error": "earlier"}, {"step_id": "x"}]}) + await Task(id="t", version=1, steps=[_FakeDeployEnv("s0")]).run(context=context) + assert context.metadata["failed_steps"] == [{"step_id": "x"}] + + def test_rolled_back_sandboxes_is_the_difference_between_stored_and_replayed(): before = { "deployed_envs": [