diff --git a/deploy/benchmark_docker/docker_sandbox.py b/deploy/benchmark_docker/docker_sandbox.py index 994aa6f5..4c25c256 100644 --- a/deploy/benchmark_docker/docker_sandbox.py +++ b/deploy/benchmark_docker/docker_sandbox.py @@ -7,12 +7,13 @@ from dataclasses import asdict from datetime import datetime, timezone import gzip import io +import inspect import json from pathlib import Path from pathlib import PurePosixPath import re import tarfile -from typing import Any, BinaryIO, Callable, Literal, Protocol +from typing import Any, Awaitable, BinaryIO, Callable, Literal, Protocol import uuid from ._cleanup import cleanup_on_error @@ -272,11 +273,12 @@ class DockerReMeSandboxFactory: plugins: tuple[str, ...] | list[str] = (), command_timeout: float = 1800.0, workspace_builder: WorkspaceBuilder | None = None, - on_sandbox_started: Callable[[], None] | None = None, + on_sandbox_started: Callable[[], None | Awaitable[None]] | None = None, ) -> None: self.candidate = candidate self._on_sandbox_started = on_sandbox_started self._sandbox_started_reported = False + self._sandbox_started_lock = asyncio.Lock() self.env = dict(env or {}) self.config = config self.plugins = tuple(plugins) @@ -312,9 +314,13 @@ class DockerReMeSandboxFactory: ) try: await case.initialize() - if self._on_sandbox_started is not None and not self._sandbox_started_reported: - self._on_sandbox_started() - self._sandbox_started_reported = True + if self._on_sandbox_started is not None: + async with self._sandbox_started_lock: + if not self._sandbox_started_reported: + result = self._on_sandbox_started() + if inspect.isawaitable(result): + await result + self._sandbox_started_reported = True except BaseException as exc: await cleanup_on_error([(f"workspace {container_id!r}", workspace.close())], exc) raise diff --git a/plugins/home/src/reme_home/agent/candidate_gate.py b/plugins/home/src/reme_home/agent/candidate_gate.py index ecf706e5..5300cbae 100644 --- a/plugins/home/src/reme_home/agent/candidate_gate.py +++ b/plugins/home/src/reme_home/agent/candidate_gate.py @@ -7,10 +7,11 @@ from uuid import uuid4 from pydantic import Field +from reme.utils.async_utils import complete_in_thread as blocking + from ..code.repository import current_branch, validate_candidate, validate_repository from ..code.snapshot import export_blob, tree from ..evaluation.artifacts import contained, read_json -from ..evaluation.component import blocking from ..evaluation.repetition import previous_candidate_result from ..preparation.code import git from ..preparation.spec import SpecModel diff --git a/plugins/home/src/reme_home/agent/case_study.py b/plugins/home/src/reme_home/agent/case_study.py index 71d957d9..119b2652 100644 --- a/plugins/home/src/reme_home/agent/case_study.py +++ b/plugins/home/src/reme_home/agent/case_study.py @@ -11,13 +11,13 @@ import yaml from reme.schema.token_usage import TokenUsage from reme.steps.base_step import BaseStep, Ref +from reme.utils.async_utils import complete_in_thread as blocking from reme.utils.local_backend import ProcessTimeout from reme.utils.qwen_token_tools import DEFAULT_TOKENIZER, QwenTokenTools from ..auto_bash import AutoBash from ..code.repository import resolve_revision, validate_repository from ..evaluation.artifacts import contained -from ..evaluation.component import blocking from ..evaluation.evidence_step import _source_case from ..evaluation.inputs import entity_id from ..evaluation.memory_inspection_step import _repository diff --git a/plugins/home/src/reme_home/agent/compression.py b/plugins/home/src/reme_home/agent/compression.py index caaf3100..fdd67460 100644 --- a/plugins/home/src/reme_home/agent/compression.py +++ b/plugins/home/src/reme_home/agent/compression.py @@ -1,6 +1,5 @@ """Meta-ReMe compression recovery adapted to AgentScope's native context pipeline.""" -import asyncio import logging from pathlib import Path from uuid import uuid4 @@ -11,6 +10,8 @@ import jsonschema import openai from pydantic import ValidationError +from reme.utils.async_utils import finish_cleanup + from ..evaluation.artifacts import contained, write_json from .bounded import BoundedReActAgent @@ -127,9 +128,4 @@ class RecoverableCompressionAgent(BoundedReActAgent): self.state.context[-1].usage = usage logger.warning("Local compression offloaded %d message(s) to %s", len(old_messages), path) - task = asyncio.create_task(apply_change()) - try: - await asyncio.shield(task) - except asyncio.CancelledError: - await task - raise + await finish_cleanup(apply_change()) diff --git a/plugins/home/src/reme_home/agent/construction_equivalence.py b/plugins/home/src/reme_home/agent/construction_equivalence.py index 58e9a5f3..dc2cc0a6 100644 --- a/plugins/home/src/reme_home/agent/construction_equivalence.py +++ b/plugins/home/src/reme_home/agent/construction_equivalence.py @@ -5,8 +5,9 @@ from typing import Literal from pydantic import Field +from reme.utils.async_utils import complete_in_thread as blocking + from ..evaluation.artifacts import contained, file_sha256, inspect_archive, read_json, write_json -from ..evaluation.component import blocking from ..evaluation.inputs import reuse_inputs from ..preparation.code import tree_fingerprint from ..preparation.spec import SpecModel diff --git a/plugins/home/src/reme_home/agent/diagnostic.py b/plugins/home/src/reme_home/agent/diagnostic.py index 06d07174..7f457c3c 100644 --- a/plugins/home/src/reme_home/agent/diagnostic.py +++ b/plugins/home/src/reme_home/agent/diagnostic.py @@ -9,9 +9,9 @@ from agentscope.tool import FunctionTool from pydantic import Field, field_validator from reme.steps.base_step import BaseStep, Ref +from reme.utils.async_utils import complete_in_thread as blocking from ..evaluation.artifacts import contained, read_json, write_json -from ..evaluation.component import blocking from ..preparation.spec import SpecModel from ..workspace import HomeWorkspace from .draft import DraftStore, create_draft_tools diff --git a/plugins/home/src/reme_home/agent/draft.py b/plugins/home/src/reme_home/agent/draft.py index f87c15c3..3fe6966e 100644 --- a/plugins/home/src/reme_home/agent/draft.py +++ b/plugins/home/src/reme_home/agent/draft.py @@ -7,8 +7,7 @@ from threading import Lock from agentscope.tool import FunctionTool from reme.utils.qwen_token_tools import QwenTokenTools - -from ..evaluation.component import blocking +from reme.utils.async_utils import complete_in_thread as blocking MAX_DRAFTS = 6 RECOMMENDED_DRAFT_TOKENS = 200 diff --git a/plugins/home/src/reme_home/agent/optimizer.py b/plugins/home/src/reme_home/agent/optimizer.py index 5a29286a..513ff185 100644 --- a/plugins/home/src/reme_home/agent/optimizer.py +++ b/plugins/home/src/reme_home/agent/optimizer.py @@ -13,10 +13,10 @@ from agentscope.tool import Edit, FunctionTool from pydantic import Field from reme.steps.base_step import BaseStep, Ref +from reme.utils.async_utils import complete_in_thread as blocking from ..code.repository import validate_repository, require_clean, validate_tree, status from ..evaluation.artifacts import contained, read_json, write_json -from ..evaluation.component import blocking from ..evaluation.inputs import EvaluationSpec, load_inputs from ..evaluation.history import read_state from ..evaluation.budget import validate_max_evaluations, validation_budget_status diff --git a/plugins/home/src/reme_home/agent/workflow_common.py b/plugins/home/src/reme_home/agent/workflow_common.py index 6a55d4d1..f1edf289 100644 --- a/plugins/home/src/reme_home/agent/workflow_common.py +++ b/plugins/home/src/reme_home/agent/workflow_common.py @@ -8,10 +8,11 @@ from uuid import uuid4 from agentscope.tool import BackendBase, DirEntry, Glob, Grep, Read, Toolkit, Write import yaml +from reme.utils.async_utils import complete_in_thread as blocking + from ..auto_bash import AutoBash from ..code.repository import validate_repository from ..evaluation.artifacts import contained, write_json -from ..evaluation.component import blocking from .case_study import python_tool from .sandbox import create_case_study_sandbox diff --git a/plugins/home/src/reme_home/code/component.py b/plugins/home/src/reme_home/code/component.py index eaaecb03..f111259c 100644 --- a/plugins/home/src/reme_home/code/component.py +++ b/plugins/home/src/reme_home/code/component.py @@ -3,6 +3,7 @@ import asyncio from reme.components.base_component import BaseComponent +from reme.utils.async_utils import finish_cleanup from ..workspace import HomeWorkspace, _initialization_lock, _validate_directory from . import repository, snapshot @@ -45,22 +46,8 @@ class HomeCode(BaseComponent): arguments = ARGUMENTS[operation].model_validate(kwargs).model_dump() task = asyncio.create_task(asyncio.to_thread(self._run, operation, arguments)) self._pending.add(task) - cancelled = False try: - while True: - try: - result = await asyncio.shield(task) - if cancelled: - raise asyncio.CancelledError - return result - except asyncio.CancelledError: - if task.done(): - raise - cancelled = True - except Exception: - if cancelled: - raise asyncio.CancelledError from None - raise + return await finish_cleanup(task) finally: self._pending.discard(task) diff --git a/plugins/home/src/reme_home/evaluation/component.py b/plugins/home/src/reme_home/evaluation/component.py index dbfbd0ea..e09e35eb 100644 --- a/plugins/home/src/reme_home/evaluation/component.py +++ b/plugins/home/src/reme_home/evaluation/component.py @@ -9,6 +9,7 @@ from contextlib import nullcontext from deploy.benchmark_docker import BenchmarkDockerPool, CasePlan, DockerReMeSandboxFactory, SourceCandidate from reme.components.base_component import BaseComponent +from reme.utils.async_utils import complete_in_thread as blocking from reme.utils.async_utils import finish_cleanup from ..preparation.data.models import utc_now @@ -35,11 +36,6 @@ from .repetition import previous_candidate_result from .reuse import automatic_reuse, construction_fingerprints, reuse_context -async def blocking(function, *args, **kwargs): - """Await owned filesystem work even when the calling job is cancelled.""" - return await finish_cleanup(asyncio.to_thread(function, *args, **kwargs)) - - class HomeEvaluation(BaseComponent): """Own evaluation tasks; the application owns the host judge Job lifecycle.""" @@ -241,9 +237,6 @@ class HomeEvaluation(BaseComponent): # Preserve those fields when the caller later saves failure status. manifest.update(read_json(run / "manifest.json")) - def _finish_sampling(self, run): - self.history.record(run) - def _allocate(self, branch, commit, manifest): """Allocate ascending IDs under a collision-safe commit prefix.""" home = self.workspace.home_path @@ -351,7 +344,8 @@ class HomeEvaluation(BaseComponent): on_sandbox_started=( None if spec.debug - else lambda: self._sampling_started( + else lambda: blocking( + self._sampling_started, run, manifest, sampling_plan, @@ -409,11 +403,6 @@ class HomeEvaluation(BaseComponent): ) if not spec.debug and spec.record_history: await blocking(self.history.publish_completed, run, manifest, summary) - try: - await blocking(self._finish_sampling, run) - except Exception: - await blocking(self.history.invalidate) - self.logger.exception("Sampling registration failed; completed evaluation retained") else: write_json(run / "manifest.json", manifest) write_json(run / "summary.json", summary, compact_leaf_containers=True) diff --git a/plugins/home/src/reme_home/evaluation/history.py b/plugins/home/src/reme_home/evaluation/history.py index 9a52784b..03ce8706 100644 --- a/plugins/home/src/reme_home/evaluation/history.py +++ b/plugins/home/src/reme_home/evaluation/history.py @@ -306,12 +306,12 @@ class HomeEvaluationHistory(BaseComponent): def record_started(self, run: Path, policy: SamplingPolicy, plan: dict | None = None) -> dict: """Record a started round, independently of whether it later enters successful history.""" - with self._history_lock(): + with self._history_lock(wait=True): return start_validation(self.workspace.home_path, self._run_path(run), policy, plan) def record(self, run: Path) -> EvaluationRecord | None: """Register a successful result once; unsuccessful runs never train the sampler.""" - with self._history_lock(): + with self._history_lock(wait=True): home = self.workspace.home_path run = self._run_path(run) record = result_record(home, run) @@ -321,16 +321,25 @@ class HomeEvaluationHistory(BaseComponent): return record def publish_completed(self, run: Path, manifest: dict, summary: dict) -> None: - """Invalidate observations before publishing the durable completion boundary. + """Publish durable completion before best-effort sampling registration. - A process exit before registration leaves a missing cache, so the next - reader reconstructs all observations from source artifacts. + Invalidate observations before publishing so interruption before registration + leaves a missing cache that readers rebuild from source artifacts. Cache + failures after publication must not turn completed evaluations into failures. """ with self._history_lock(wait=True): run = self._run_path(run) contained(self.workspace.home_path, "sampling/state.json").unlink(missing_ok=True) write_json(run / "manifest.json", manifest) write_json(run / "summary.json", summary, compact_leaf_containers=True) + try: + self.record(run) + except Exception: + self.logger.exception("Sampling registration failed; completed evaluation retained") + try: + self.invalidate() + except Exception: + self.logger.exception("Sampling cache invalidation failed; completed evaluation retained") def list_results(self) -> list[EvaluationRecord]: """Return successful results in chronological order, including low-scoring candidates.""" @@ -346,12 +355,12 @@ class HomeEvaluationHistory(BaseComponent): def rebuild(self) -> list[EvaluationRecord]: """Rebuild only derived indexes and sampling observations; never rewrite source evidence.""" - with self._history_lock(): + with self._history_lock(wait=True): home = self.workspace.home_path rebuild_sampling_state(home) return self._records() def invalidate(self) -> None: """Discard only derived caches after an unsuccessful registration; retain all source evidence.""" - with self._history_lock(): + with self._history_lock(wait=True): contained(self.workspace.home_path, "sampling/state.json").unlink(missing_ok=True) diff --git a/plugins/home/src/reme_home/evaluation/inputs.py b/plugins/home/src/reme_home/evaluation/inputs.py index 8a79f545..276b283b 100644 --- a/plugins/home/src/reme_home/evaluation/inputs.py +++ b/plugins/home/src/reme_home/evaluation/inputs.py @@ -166,17 +166,22 @@ def snapshot_code(home: Path, record: dict): return snapshot, commit, branch +def _reuse_source(home: Path, value: str): + """Read current source artifacts within the managed evaluation tree.""" + evaluations = contained(home, "evaluations", directory=True) + source = contained(home, value, directory=True) + if not source.is_relative_to(evaluations) or source == evaluations: + raise ValueError("reuse_construction_from must be inside home/evaluations") + manifest = read_json(contained(source, "manifest.json")) + summary = read_json(contained(source, "summary.json")) + return source, manifest, summary + + def reuse_inputs(home: Path, spec: EvaluationSpec, cases: list, fingerprints: dict): """Preflight the entire requested selection; never silently skip or rebuild cases.""" if spec.reuse_construction_from is None: return None, {} - evaluations = contained(home, "evaluations", directory=True) - value = spec.reuse_construction_from - source = contained(home, value, directory=True) - if not source.is_relative_to(evaluations) or source == evaluations: - raise ValueError("reuse_construction_from must be inside home/evaluations") - manifest = read_json(source / "manifest.json") - summary = read_json(source / "summary.json") + source, manifest, summary = _reuse_source(home, spec.reuse_construction_from) if manifest.get("dataset_kind") != "training-set" or manifest.get("status", summary.get("status")) != "completed": raise ValueError("Reuse requires a completed training-set evaluation") for key in ("dataset", "construction_config"): @@ -205,12 +210,7 @@ def replay_sampling_eligibility(home: Path, spec: EvaluationSpec, cases: list[Ca """Port dev_meta's replay eligibility without restricting future full rebuilds.""" if spec.reuse_construction_from is None: return None - source = contained(home, spec.reuse_construction_from, directory=True) - evaluations = contained(home, "evaluations", directory=True) - if not source.is_relative_to(evaluations) or source == evaluations: - raise ValueError("reuse_construction_from must be inside home/evaluations") - manifest = read_json(contained(source, "manifest.json")) - summary = read_json(contained(source, "summary.json")) + source, manifest, summary = _reuse_source(home, spec.reuse_construction_from) if ( manifest.get("dataset_kind") != "training-set" or manifest.get("mode") == "debug" diff --git a/plugins/home/src/reme_home/evaluation/inspection_step.py b/plugins/home/src/reme_home/evaluation/inspection_step.py index c0ecd488..aa1059e5 100644 --- a/plugins/home/src/reme_home/evaluation/inspection_step.py +++ b/plugins/home/src/reme_home/evaluation/inspection_step.py @@ -6,6 +6,7 @@ import json from pydantic import Field, field_validator, model_validator from reme.steps.base_step import BaseStep, Ref +from reme.utils.async_utils import finish_cleanup from ..code.repository import current_branch, resolve_revision, validate_branch, validate_repository from ..preparation.code import git @@ -323,19 +324,8 @@ async def _execute(step, model, function): spec = model.model_validate({name: step.context.get(name) for name in model.model_fields if name in step.context}) if not step.workspace.is_started: raise RuntimeError("Start the Home workspace before inspecting evaluations") - task = asyncio.create_task(asyncio.to_thread(function, step.workspace.home_path, spec)) - cancelled = False - while True: - try: - result = await asyncio.shield(task) - if cancelled: - raise asyncio.CancelledError - step.context.response.answer = result - return step.context.response - except asyncio.CancelledError: - if task.done(): - raise - cancelled = True + step.context.response.answer = await finish_cleanup(asyncio.to_thread(function, step.workspace.home_path, spec)) + return step.context.response class HomeListValidationResultsStep(BaseStep): diff --git a/plugins/home/src/reme_home/evaluation/reuse.py b/plugins/home/src/reme_home/evaluation/reuse.py index ea3c2a91..9c4955d0 100644 --- a/plugins/home/src/reme_home/evaluation/reuse.py +++ b/plugins/home/src/reme_home/evaluation/reuse.py @@ -5,6 +5,8 @@ import random import subprocess import tarfile +from reme.utils.async_utils import complete_in_thread as blocking + from ..code.repository import current_branch, validate_candidate, validate_repository from ..code.snapshot import tree from ..preparation.code import git @@ -136,7 +138,6 @@ async def automatic_reuse(wrapper, home, prepared, *, memory_limit_mib=None): """Automatic failures rebuild; only a sandboxed unchanged review permits a probability draw.""" from ..agent.construction_equivalence import review_construction from ..agent.workflow_common import save_record - from .component import blocking if prepared is None: return None diff --git a/plugins/home/src/reme_home/preparation/component.py b/plugins/home/src/reme_home/preparation/component.py index 67a2d66a..48515601 100644 --- a/plugins/home/src/reme_home/preparation/component.py +++ b/plugins/home/src/reme_home/preparation/component.py @@ -12,6 +12,7 @@ import yaml from reme.components.base_component import BaseComponent from reme.steps.base_step import BaseStep, Ref +from reme.utils.async_utils import finish_cleanup from .code import prepare_code from .evaluation_source import EvaluationSource @@ -72,20 +73,21 @@ class HomePreparation(BaseComponent): asyncio.to_thread(self._prepare, spec=spec, spec_path=spec_path, stop=stop, code_source=code_source), ) self._pending.add(task) - cancelled = False try: - while True: - try: - return await asyncio.shield(task) - except asyncio.CancelledError: - if task.cancelled(): - raise - cancelled = True - stop.set() - except InterruptedError: - if cancelled: - raise asyncio.CancelledError from None + try: + return await asyncio.shield(task) + except asyncio.CancelledError as cancellation: + if task.cancelled(): raise + stop.set() + try: + await finish_cleanup(asyncio.gather(task, return_exceptions=True)) + except asyncio.CancelledError: + pass + if task.cancelled() or task.exception() is not None: + raise cancellation from None + # A worker past the durable commit boundary still returns success. + return task.result() finally: self._pending.discard(task) diff --git a/plugins/home/tests/test_code.py b/plugins/home/tests/test_code.py index 16eaf0a6..8ceb522d 100644 --- a/plugins/home/tests/test_code.py +++ b/plugins/home/tests/test_code.py @@ -239,7 +239,9 @@ async def test_lock_restrictions_and_unprepared(app): @pytest.mark.asyncio -async def test_cancel_waits_for_owned_git_thread(prepared, monkeypatch): +@pytest.mark.parametrize("cancelled", [False, True]) +@pytest.mark.parametrize("fail", [False, True]) +async def test_git_thread_drains_and_preserves_cancellation_or_error(prepared, monkeypatch, cancelled, fail): app, _, _ = prepared component = app.context.components["home.code"]["default"] entered, release, finished = Event(), Event(), Event() @@ -249,19 +251,35 @@ async def test_cancel_waits_for_owned_git_thread(prepared, monkeypatch): entered.set() assert release.wait(5) try: + if fail: + raise OSError("injected Git failure") return original(*args) finally: finished.set() monkeypatch.setattr(component, "_run", blocked) task = asyncio.create_task(component.run("status")) - assert await asyncio.to_thread(entered.wait, 5) - task.cancel() - await asyncio.sleep(0) - assert not task.done() - release.set() - with pytest.raises(asyncio.CancelledError): - await task + try: + assert await asyncio.to_thread(entered.wait, 5) + assert len(component._pending) == 1 + if cancelled: + task.cancel("first cancellation") + await asyncio.sleep(0) + task.cancel("second cancellation") + await asyncio.sleep(0) + assert not task.done() + release.set() + if cancelled: + with pytest.raises(asyncio.CancelledError, match="first cancellation"): + await task + elif fail: + with pytest.raises(OSError, match="injected Git failure"): + await task + else: + assert isinstance(await task, dict) + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) assert finished.is_set() and not component._pending diff --git a/plugins/home/tests/test_evaluation.py b/plugins/home/tests/test_evaluation.py index 98a5ff32..1c87688c 100644 --- a/plugins/home/tests/test_evaluation.py +++ b/plugins/home/tests/test_evaluation.py @@ -167,7 +167,7 @@ class FakeFactory: async def create_cases(self, ids, **kwargs): self.workers = [FakeWorker(self) for _ in ids] if self.workers and self.kwargs.get("on_sandbox_started"): - self.kwargs["on_sandbox_started"]() + await self.kwargs["on_sandbox_started"]() return self.workers diff --git a/plugins/home/tests/test_evaluation_history.py b/plugins/home/tests/test_evaluation_history.py index 6e3efddc..088039ef 100644 --- a/plugins/home/tests/test_evaluation_history.py +++ b/plugins/home/tests/test_evaluation_history.py @@ -98,6 +98,68 @@ async def test_completion_waits_for_concurrent_code_browsing(app, prepared, monk await asyncio.gather(evaluation, *([code_status] if code_status is not None else []), return_exceptions=True) +@pytest.mark.asyncio +@pytest.mark.parametrize("stage", ["record_started", "record", "invalidate"]) +async def test_history_writes_wait_for_code_read_without_losing_result(app, prepared, monkeypatch, stage): + from reme_home.code.component import OPERATIONS + + commit_candidate(prepared) + manager = history(app) + reached, browsing, acquiring, release = Event(), Event(), Event(), Event() + original, original_status = getattr(manager, stage), OPERATIONS["status"] + + def write(*args, **kwargs): + reached.set() + assert browsing.wait(5) + acquiring.set() + return original(*args, **kwargs) + + def status(*args, **kwargs): + browsing.set() # The real code Job already owns workspace.lock. + assert release.wait(5) + return original_status(*args, **kwargs) + + if stage == "invalidate": + + def fail_record(*args): + raise OSError("injected registration failure after completion publication") + + monkeypatch.setattr(manager, "record", fail_record) + monkeypatch.setattr(manager, stage, write) + monkeypatch.setitem(OPERATIONS, "status", status) + evaluation = asyncio.create_task(app.run_job("home_evaluate", max_evaluations=1)) + code_status = None + try: + assert await asyncio.to_thread(reached.wait, 5) + code_status = asyncio.create_task(app.run_job("home_code_status")) + assert await asyncio.to_thread(acquiring.wait, 5) + await asyncio.sleep(0.1) + assert not evaluation.done() + # This coroutine must remain runnable while the history worker waits. + release.set() + response = await evaluation + assert response.success, response.answer + assert (await code_status).success + payload = manifest(response) + assert payload["status"] == "completed" and payload["success"] is True + plan = json.loads((prepared / payload["sampling"]["plan_path"]).read_text()) + assert payload["sandbox_started_at"] == plan["sandbox_started_at"] + budget = validation_budget_status(prepared, 1) + assert budget["used_validations"] == 1 and budget["remaining_validations"] == 0 + assert response.answer["validation_budget"] == budget + state = read_state(prepared) + assert state["round"] == state["successful_round"] == 1 + assert state["started_validations"] == {plan["validation_key"]: 1} + assert len(manager.list_results()) == 1 + manager.rebuild() + assert read_state(prepared) == state + retry = await app.run_job("home_evaluate") + assert retry.answer["error_code"] == "UNCHANGED_CANDIDATE" + finally: + release.set() + await asyncio.gather(evaluation, *([code_status] if code_status is not None else []), return_exceptions=True) + + @pytest.mark.asyncio async def test_start_cache_failure_keeps_start_evidence_and_exhausts_budget(app, prepared, monkeypatch): commit_candidate(prepared) @@ -508,7 +570,14 @@ async def test_cancellation_after_durable_completion_keeps_history_consistent(ap @pytest.mark.asyncio @pytest.mark.parametrize("after_record", [False, True]) -async def test_registration_failure_invalidates_only_caches(app, prepared, monkeypatch, after_record): +@pytest.mark.parametrize("invalidation_fails", [False, True]) +async def test_registration_failure_invalidates_only_caches( + app, + prepared, + monkeypatch, + after_record, + invalidation_fails, +): manager = history(app) original = manager.record @@ -518,6 +587,12 @@ async def test_registration_failure_invalidates_only_caches(app, prepared, monke raise OSError("sampling observations could not be saved") monkeypatch.setattr(manager, "record", fail) + if invalidation_fails: + + def fail_invalidate(): + raise OSError("injected cache invalidation failure") + + monkeypatch.setattr(manager, "invalidate", fail_invalidate) response = await app.run_job("home_evaluate") assert response.success, response.answer runs = list((prepared / "evaluations").glob("*/*/*/manifest.json")) @@ -526,7 +601,7 @@ async def test_registration_failure_invalidates_only_caches(app, prepared, monke completed = json.loads(runs[0].read_text()) assert completed["status"] == "completed" and completed["success"] is True assert (run / "summary.json").is_file() and (run / "cases/1/queries/result.json").is_file() - assert not (prepared / "sampling/state.json").exists() + assert (prepared / "sampling/state.json").exists() is (invalidation_fails and after_record) assert read_state(prepared)["round"] == 1 and read_state(prepared)["successful_round"] == 1 monkeypatch.setattr(manager, "record", original) assert len(manager.list_results()) == 1 @@ -570,7 +645,7 @@ async def run(): await app.start() app.context.components["agent_wrapper"]["judge"].reply = AsyncMock( return_value={"result": '{"score": 0.5}'}) - app.context.components["home.evaluation"]["default"]._finish_sampling = lambda run: os._exit(23) + app.context.components["home.evaluation_history"]["default"].record = lambda run: os._exit(23) await app.run_job("home_evaluate") asyncio.run(run()) """ diff --git a/plugins/home/tests/test_evaluation_inspection.py b/plugins/home/tests/test_evaluation_inspection.py index 8fe604ed..0cb0c65d 100644 --- a/plugins/home/tests/test_evaluation_inspection.py +++ b/plugins/home/tests/test_evaluation_inspection.py @@ -291,22 +291,41 @@ async def test_hidden_case_scope_is_enforced_by_agent_job_injection(app, prepare @pytest.mark.asyncio -async def test_cancelled_job_drains_read_thread(app, prepared, monkeypatch): +@pytest.mark.parametrize("cancelled", [False, True]) +@pytest.mark.parametrize("fail", [False, True]) +async def test_read_thread_drains_and_preserves_cancellation_or_error(app, prepared, monkeypatch, cancelled, fail): entered, release, finished = Event(), Event(), Event() def read(*args): entered.set() assert release.wait(5) finished.set() + if fail: + raise OSError("injected inspection failure") return {} monkeypatch.setattr(inspection, "_list_results", read) task = asyncio.create_task(app.run_job("home_list_validation_results")) - assert await asyncio.to_thread(entered.wait, 5) - task.cancel() - await asyncio.sleep(0) - assert not task.done() - release.set() - with pytest.raises(asyncio.CancelledError): - await task + try: + assert await asyncio.to_thread(entered.wait, 5) + if cancelled: + task.cancel("first cancellation") + await asyncio.sleep(0) + task.cancel("second cancellation") + await asyncio.sleep(0) + assert not task.done() + release.set() + if cancelled: + with pytest.raises(asyncio.CancelledError, match="first cancellation"): + await task + else: + response = await task + assert response.success is (not fail) + if fail: + assert "injected inspection failure" in str(response.answer) + else: + assert response.answer == {} + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) assert finished.is_set() diff --git a/plugins/home/tests/test_optimizer_compression.py b/plugins/home/tests/test_optimizer_compression.py index d54b6bc1..86556e6a 100644 --- a/plugins/home/tests/test_optimizer_compression.py +++ b/plugins/home/tests/test_optimizer_compression.py @@ -182,7 +182,12 @@ async def test_offload_write_failure_keeps_context_and_does_not_count(tmp_path, @pytest.mark.asyncio -async def test_cancellation_during_apply_finishes_consistent_offload_then_propagates(tmp_path, monkeypatch): +@pytest.mark.parametrize("cancellations", [1, 3]) +async def test_cancellation_during_apply_finishes_consistent_offload_then_propagates( + tmp_path, + monkeypatch, + cancellations, +): agent = make_agent(tmp_path) started, release = asyncio.Event(), asyncio.Event() @@ -192,11 +197,18 @@ async def test_cancellation_during_apply_finishes_consistent_offload_then_propag monkeypatch.setattr(agent, "_clear_unreserved_read_cache", clear_cache) task = asyncio.create_task(agent.compress_context()) - await started.wait() - task.cancel() - release.set() - with pytest.raises(asyncio.CancelledError): - await task + try: + await asyncio.wait_for(started.wait(), 5) + for index in range(cancellations): + task.cancel(f"cancellation-{index}") + await asyncio.sleep(0) + assert not task.done() + release.set() + with pytest.raises(asyncio.CancelledError, match="cancellation-0"): + await task + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) assert len(list(tmp_path.rglob("context-*.json"))) == 1 assert "offloaded" in agent.state.summary assert len(agent.state.context) == 2 diff --git a/plugins/home/tests/test_preparation.py b/plugins/home/tests/test_preparation.py index be91daf3..467cb841 100644 --- a/plugins/home/tests/test_preparation.py +++ b/plugins/home/tests/test_preparation.py @@ -482,7 +482,9 @@ def test_bad_recovery_journal_does_not_remove_outside_paths(tmp_path): @pytest.mark.asyncio -async def test_cancellation_waits_for_rollback(app, inputs, monkeypatch): +@pytest.mark.parametrize("cancellations", [1, 3]) +@pytest.mark.parametrize("fail_stage", [False, True]) +async def test_cancellation_waits_for_rollback(app, inputs, monkeypatch, cancellations, fail_stage): """The event loop stays available, and cancellation waits for the worker to restore state.""" spec, _ = inputs root = Path(app.config.workspace_dir) @@ -496,26 +498,66 @@ async def test_cancellation_waits_for_rollback(app, inputs, monkeypatch): entered.set() if not proceed.wait(10): raise RuntimeError("test worker timed out") - return stage(*args) + result = stage(*args) + if fail_stage: + raise OSError("injected stage failure after cancellation") + return result monkeypatch.setattr(component, "_stage", blocked) task = asyncio.create_task(app.run_job("home_prepare", spec=spec)) try: assert await asyncio.to_thread(entered.wait, 10) - task.cancel() - await asyncio.sleep(0) + for index in range(cancellations): + task.cancel(f"cancellation-{index}") + await asyncio.sleep(0) assert not task.done() + assert component._pending # pylint: disable=protected-access proceed.set() - with pytest.raises(asyncio.CancelledError): + with pytest.raises(asyncio.CancelledError, match="cancellation-0"): await task finally: proceed.set() await asyncio.gather(task, return_exceptions=True) assert snapshot(root) == before + assert not component._pending # pylint: disable=protected-access monkeypatch.setattr(component, "_stage", stage) assert (await app.run_job("home_prepare", spec=spec)).success +@pytest.mark.asyncio +@pytest.mark.parametrize("cancellations", [1, 3]) +async def test_cancellation_after_preparation_commit_retains_success(app, inputs, monkeypatch, cancellations): + """A stop request after the last published target must not undo a completed transaction.""" + from reme_home.preparation import component as preparation + + entered, release = Event(), Event() + original = preparation.publish + + def committed(*args): + original(*args) + entered.set() + assert release.wait(5) + + monkeypatch.setattr(preparation, "publish", committed) + task = asyncio.create_task(app.run_job("home_prepare", spec=inputs[0])) + try: + assert await asyncio.to_thread(entered.wait, 5) + for index in range(cancellations): + task.cancel(f"cancellation-{index}") + await asyncio.sleep(0) + assert not task.done() + release.set() + response = await task + assert response.success, response.answer + home = Path(app.config.workspace_dir) / "home" + assert json.loads((home / "metadata/preparation.json").read_text())["preparation_id"] + assert not (home / "metadata" / TRANSACTION).exists() + assert not app.context.components["home.preparation"]["default"]._pending # pylint: disable=protected-access + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) + + @pytest.mark.asyncio @pytest.mark.parametrize("app", ["lme-judge"], indirect=True) async def test_duplicate_lme_ids_fail_without_mutation(app, inputs): diff --git a/plugins/home/tests/test_sampling.py b/plugins/home/tests/test_sampling.py index ef0c826d..8a5ba91f 100644 --- a/plugins/home/tests/test_sampling.py +++ b/plugins/home/tests/test_sampling.py @@ -147,7 +147,7 @@ async def test_baseline_then_adaptive_job_updates_only_selected_cases(app, prepa assert after["cases"][selected]["observed_count"] == 2 assert after["cases"][selected]["successful_observations_after_last_error"] == 1 assert after["cases"][unselected] == before["cases"][unselected] - evaluation._finish_sampling(run) + app.context.components["home.evaluation_history"]["default"].record(run) assert state(prepared) == after # A summary is consumed once, including after reopening. assert "home_validate" not in app.context.jobs assert app.context.jobs["home_evaluate"].parameters["properties"]["sampling"]["type"] == "boolean" diff --git a/reme/utils/async_utils.py b/reme/utils/async_utils.py index d2dd9a85..58c8697c 100644 --- a/reme/utils/async_utils.py +++ b/reme/utils/async_utils.py @@ -29,6 +29,6 @@ async def finish_cleanup(operation: Awaitable[T]) -> T: return task.result() -async def complete_in_thread(func: Callable[..., Any], /, *args) -> Any: +async def complete_in_thread(func: Callable[..., Any], /, *args, **kwargs) -> Any: """Finish a side-effecting thread call, including through repeated cancellation.""" - return await finish_cleanup(asyncio.to_thread(func, *args)) + return await finish_cleanup(asyncio.to_thread(func, *args, **kwargs)) diff --git a/tests/unit/test_async_utils.py b/tests/unit/test_async_utils.py index 70b90adc..95fcd0f8 100644 --- a/tests/unit/test_async_utils.py +++ b/tests/unit/test_async_utils.py @@ -17,17 +17,17 @@ async def test_complete_in_thread_finishes_before_propagating_result(cancellatio """Repeated cancellation waits for success or failure without detaching work.""" started, release, finished = threading.Event(), threading.Event(), threading.Event() - def operation(): + def operation(value, *, suffix): started.set() try: assert release.wait(5) if fail: raise RuntimeError("thread failed") - return "complete" + return value + suffix finally: finished.set() - task = asyncio.create_task(complete_in_thread(operation)) + task = asyncio.create_task(complete_in_thread(operation, "comp", suffix="lete")) try: assert await asyncio.to_thread(started.wait, 5) for _ in range(cancellations): diff --git a/tests/unit/test_benchmark_docker_start_callback.py b/tests/unit/test_benchmark_docker_start_callback.py index 2dfa5294..7a67261c 100644 --- a/tests/unit/test_benchmark_docker_start_callback.py +++ b/tests/unit/test_benchmark_docker_start_callback.py @@ -36,6 +36,37 @@ async def test_start_callback_only_once_after_initialization(monkeypatch): assert len(reported) == 1 +@pytest.mark.asyncio +async def test_async_start_callback_serializes_concurrent_initializations(monkeypatch): + """Other initialized workers await the same durable start before being returned.""" + monkeypatch.setattr(docker_sandbox.DockerReMeSandbox, "initialize", AsyncMock()) + entered, release = asyncio.Event(), asyncio.Event() + calls = [] + + async def started(): + calls.append(True) + entered.set() + await release.wait() + + factory = DockerReMeSandboxFactory( + ImageCandidate("fixture"), + workspace_builder=lambda *_: AsyncMock(), + on_sandbox_started=started, + ) + first = asyncio.create_task(factory.create_case("first")) + second = asyncio.create_task(factory.create_case("second")) + try: + await asyncio.wait_for(entered.wait(), 5) + await asyncio.sleep(0) + assert not first.done() and not second.done() + assert calls == [True] + finally: + release.set() + await asyncio.gather(first, second) + assert factory._sandbox_started_reported # pylint: disable=protected-access + assert calls == [True] + + @pytest.mark.asyncio async def test_start_callback_failure_closes_initialized_resource(monkeypatch): """Persisting a start marker must not orphan an already-created workspace on failure."""