scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-11-14 09:01:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

fix(home): unify history writes and cancellation cleanup

Browse source
This commit is contained in:
酱牛肉 2026-11-04 17:30:33 +08:00
parent 0f50c086bc
commit 3441303d4d
26 changed files with 311 additions and 132 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
deploy/benchmark_docker/docker_sandbox.py
View file

@ -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

3
plugins/home/src/reme_home/agent/candidate_gate.py
View file

@ -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

2
plugins/home/src/reme_home/agent/case_study.py
View file

@ -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

10
plugins/home/src/reme_home/agent/compression.py
View file

@ -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())

3
plugins/home/src/reme_home/agent/construction_equivalence.py
View file

@ -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

2
plugins/home/src/reme_home/agent/diagnostic.py
View file

@ -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

3
plugins/home/src/reme_home/agent/draft.py
View file

@ -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

2
plugins/home/src/reme_home/agent/optimizer.py
View file

@ -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

3
plugins/home/src/reme_home/agent/workflow_common.py
View file

@ -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

17
plugins/home/src/reme_home/code/component.py
View file

@ -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)

17
plugins/home/src/reme_home/evaluation/component.py
View file

@ -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)

23
plugins/home/src/reme_home/evaluation/history.py
View file

@ -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)

26
plugins/home/src/reme_home/evaluation/inputs.py
View file

@ -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"

16
plugins/home/src/reme_home/evaluation/inspection_step.py
View file

@ -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):

3
plugins/home/src/reme_home/evaluation/reuse.py
View file

@ -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

26
plugins/home/src/reme_home/preparation/component.py
View file

@ -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)

34
plugins/home/tests/test_code.py
View file

@ -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

2
plugins/home/tests/test_evaluation.py
View file

@ -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

81
plugins/home/tests/test_evaluation_history.py
View file

@ -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())
"""

35
plugins/home/tests/test_evaluation_inspection.py
View file

@ -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()

24
plugins/home/tests/test_optimizer_compression.py
View file

@ -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

52
plugins/home/tests/test_preparation.py
View file

@ -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):

2
plugins/home/tests/test_sampling.py
View file

@ -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"

4
reme/utils/async_utils.py
View file

@ -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))

6
tests/unit/test_async_utils.py
View file

@ -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):

31
tests/unit/test_benchmark_docker_start_callback.py
View file

@ -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."""