From 4a1707bfbbf61f2eb259b0c40300ef0bebf3afb5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Tue, 13 Oct 2026 19:22:50 +0800 Subject: [PATCH] Harden benchmark sandbox reuse and isolate query state --- deploy/benchmark_docker/docker_sandbox.py | 176 +++++------ deploy/benchmark_docker/models.py | 4 - deploy/benchmark_docker/pool.py | 58 +++- deploy/benchmark_docker/worker.py | 1 + tests/unit/test_benchmark_docker.py | 349 ++++++++++++++++++++++ 5 files changed, 468 insertions(+), 120 deletions(-) diff --git a/deploy/benchmark_docker/docker_sandbox.py b/deploy/benchmark_docker/docker_sandbox.py index a0928ee8..fe3c4d86 100644 --- a/deploy/benchmark_docker/docker_sandbox.py +++ b/deploy/benchmark_docker/docker_sandbox.py @@ -269,8 +269,11 @@ class DockerReMeSandboxFactory: ) try: await case.initialize() - except BaseException: - await workspace.close() + except BaseException as exc: + try: + await workspace.close() + except Exception as close_exc: + exc.add_note(f"workspace close failed: {type(close_exc).__name__}: {close_exc}") raise return case @@ -302,6 +305,7 @@ class DockerReMeSandboxFactory: if failures: await asyncio.gather( *(result.close() for result in results if isinstance(result, DockerReMeSandbox)), + return_exceptions=True, ) raise failures[0] return [result for result in results if isinstance(result, DockerReMeSandbox)] @@ -363,7 +367,7 @@ class DockerReMeSandbox: self.backend = self.workspace.get_backend() await self._create_case_layout() worker = Path(__file__).with_name("worker.py").read_bytes() - await self.backend.write_file(f"{HARNESS_ROOT}/worker.py", worker) + await self._write_file(f"{HARNESS_ROOT}/worker.py", worker) if isinstance(self.candidate, SourceCandidate): await self._install_source_candidate() @@ -395,7 +399,7 @@ class DockerReMeSandbox: async def _install_source_candidate(self) -> None: assert self.backend is not None - await self.backend.write_file(SOURCE_ARCHIVE, self.candidate.snapshot.archive) + await self._write_file(SOURCE_ARCHIVE, self.candidate.snapshot.archive) await self._exec_checked("create_candidate_dir", ["mkdir", "-p", CANDIDATE_ROOT]) await self._exec_checked( "extract_candidate", @@ -463,6 +467,21 @@ class DockerReMeSandbox: if escaped: raise SandboxCommandError(f"candidate plugin imports did not resolve under {CANDIDATE_ROOT}: {escaped}") + async def _write_file(self, path: str, data: bytes) -> None: + """Reject writes after an interrupted or uncertain Docker upload.""" + assert self.backend is not None, "sandbox is not initialized" + if not self.reusable: + raise SandboxCommandError("sandbox container is no longer reusable") + try: + await self.backend.write_file(path, data) + except asyncio.CancelledError: + self._reusable = False + raise + except Exception: + # Docker may have accepted the upload before the client failed. + self._reusable = False + raise + async def _exec(self, name: str, command: list[str], *, cwd: str | None = None) -> Any: assert self.backend is not None, "sandbox is not initialized" if not self.reusable: @@ -473,6 +492,13 @@ class DockerReMeSandbox: except TimeoutError as exc: self._reusable = False raise SandboxCommandTimeoutError(f"sandbox action {name!r} timed out; container must be destroyed") from exc + except asyncio.CancelledError: + self._reusable = False + raise + except Exception: + # A backend failure may occur after Docker started the command. + self._reusable = False + raise self.actions.append( ActionRecord( name=name, @@ -516,7 +542,7 @@ class DockerReMeSandbox: "workspace_dir": self.runtime_workspace, **payload, } - await self.backend.write_file(request_path, json.dumps(request, ensure_ascii=False).encode("utf-8")) + await self._write_file(request_path, json.dumps(request, ensure_ascii=False).encode("utf-8")) action_name = f"job:{payload['job']}" if mode == "job" else f"worker:{mode}" await self._exec( action_name, @@ -664,41 +690,9 @@ class DockerReMeSandbox: await self._write_named_result("answer.json", result) return result - async def judge( - self, - *, - query: str, - agent_answer: str, - golden_answer: str, - question_type: str = "", - ) -> JobResult: - """Run LLM-as-Judge and persist a normalized numeric score.""" - async with self._operation_lock: - result = await self._run_job( - "answer_judge", - { - "query": query, - "agent_answer": agent_answer, - "golden_answer": golden_answer, - "question_type": question_type, - }, - ) - verdict = str(result.answer).strip().lower() - score = 1 if verdict == "yes" else 0 if verdict == "no" else None - assert self.backend is not None - await self.backend.write_file( - f"{self.results_dir}/score.json", - json.dumps( - {"success": result.success, "score": score, "verdict": verdict, "raw": asdict(result)}, - ensure_ascii=False, - default=str, - ).encode("utf-8"), - ) - return result - async def _write_named_result(self, name: str, result: JobResult) -> None: assert self.backend is not None - await self.backend.write_file( + await self._write_file( f"{self.results_dir}/{name}", json.dumps(asdict(result), ensure_ascii=False, default=str).encode("utf-8"), ) @@ -733,24 +727,19 @@ class DockerReMeSandbox: case directory, logs, or candidate source. """ async with self._operation_lock: - assert self.backend is not None, "sandbox is not initialized" - await self._exec_checked( + return await self._export_archive( + destination, "export_workspace", - ["tar", "-czf", WORKSPACE_EXPORT_ARCHIVE, "-C", self.runtime_workspace, "."], + self.runtime_workspace, + ["."], + archive_path=WORKSPACE_EXPORT_ARCHIVE, ) - target = Path(destination).expanduser().resolve() - target.parent.mkdir(parents=True, exist_ok=True) - target.write_bytes(await self.backend.read_file(WORKSPACE_EXPORT_ARCHIVE)) - return target async def export_build_log(self, destination: str | Path) -> Path: """Download the memory-construction log without wrapping it in an archive.""" async with self._operation_lock: content = await self._read_build_log() - target = Path(destination).expanduser().resolve() - target.parent.mkdir(parents=True, exist_ok=True) - target.write_bytes(content) - return target + return self._save_host_file(destination, content) async def read_build_log(self) -> str: """Return the build log, including when construction failed.""" @@ -776,22 +765,12 @@ class DockerReMeSandbox: raise SandboxCommandError( f"memory construction export is missing required artifact: {required}", ) from exc - await self._exec_checked( + return await self._export_archive( + destination, "export_memory_construction", - [ - "tar", - "-czf", - EXPORT_ARCHIVE, - "-C", - CASE_ROOT, - "reme_workspace", - "build_log", - ], + CASE_ROOT, + ["reme_workspace", "build_log"], ) - target = Path(destination).expanduser().resolve() - target.parent.mkdir(parents=True, exist_ok=True) - target.write_bytes(await self.backend.read_file(EXPORT_ARCHIVE)) - return target async def export_queries(self, destination: str | Path) -> Path: """Download per-query logs, results, and their aggregate summary.""" @@ -802,14 +781,7 @@ class DockerReMeSandbox: await self.backend.read_file(required) except FileNotFoundError as exc: raise SandboxCommandError(f"query export is missing required artifact: {required}") from exc - await self._exec_checked( - "export_queries", - ["tar", "-czf", EXPORT_ARCHIVE, "-C", CASE_ROOT, "queries"], - ) - target = Path(destination).expanduser().resolve() - target.parent.mkdir(parents=True, exist_ok=True) - target.write_bytes(await self.backend.read_file(EXPORT_ARCHIVE)) - return target + return await self._export_archive(destination, "export_queries", CASE_ROOT, ["queries"]) async def export_query(self, query_id: str, destination: str | Path) -> Path: """Download one leased query's logs and result without a case summary.""" @@ -823,14 +795,7 @@ class DockerReMeSandbox: await self.backend.read_file(required) except FileNotFoundError as exc: raise SandboxCommandError(f"query export is missing required artifact: {required}") from exc - await self._exec_checked( - "export_query", - ["tar", "-czf", EXPORT_ARCHIVE, "-C", CASE_ROOT, f"queries/{query_id}"], - ) - target = Path(destination).expanduser().resolve() - target.parent.mkdir(parents=True, exist_ok=True) - target.write_bytes(await self.backend.read_file(EXPORT_ARCHIVE)) - return target + return await self._export_archive(destination, "export_query", CASE_ROOT, [f"queries/{query_id}"]) async def export_evaluation(self, destination: str | Path) -> Path: """Download workspace and query artifacts, plus the build log when present.""" @@ -847,23 +812,29 @@ class DockerReMeSandbox: has_build_log = False else: has_build_log = True - await self._exec_checked( - "export_evaluation", - [ - "tar", - "-czf", - EXPORT_ARCHIVE, - "-C", - CASE_ROOT, - "reme_workspace", - "queries", - *(["build_log"] if has_build_log else []), - ], - ) - target = Path(destination).expanduser().resolve() - target.parent.mkdir(parents=True, exist_ok=True) - target.write_bytes(await self.backend.read_file(EXPORT_ARCHIVE)) - return target + members = ["reme_workspace", "queries", *(["build_log"] if has_build_log else [])] + return await self._export_archive(destination, "export_evaluation", CASE_ROOT, members) + + @staticmethod + def _save_host_file(destination: str | Path, content: bytes) -> Path: + target = Path(destination).expanduser().resolve() + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(content) + return target + + async def _export_archive( + self, + destination: str | Path, + name: str, + root: str, + members: list[str], + *, + archive_path: str = EXPORT_ARCHIVE, + ) -> Path: + """Archive selected sandbox paths and download the archive while locked.""" + assert self.backend is not None, "sandbox is not initialized" + await self._exec_checked(name, ["tar", "-czf", archive_path, "-C", root, *members]) + return self._save_host_file(destination, await self.backend.read_file(archive_path)) async def upload_workspace(self, source: str | Path, *, clear: bool = True) -> None: """Copy a host workspace directory or exported archive into this case. @@ -880,7 +851,7 @@ class DockerReMeSandbox: if clear: await self._exec_checked("clear_runtime_workspace", ["rm", "-rf", self.runtime_workspace]) await self._exec_checked("create_runtime_workspace", ["mkdir", "-p", self.runtime_workspace]) - await self.backend.write_file(WORKSPACE_ARCHIVE, archive) + await self._write_file(WORKSPACE_ARCHIVE, archive) await self._exec_checked( "upload_workspace", ["tar", "-xzf", WORKSPACE_ARCHIVE, "-C", self.runtime_workspace], @@ -974,11 +945,11 @@ class DockerReMeSandbox: "export_profile": profile, "exported_at": _utc_now(), } - await self.backend.write_file( + await self._write_file( f"{CASE_ROOT}/manifest.json", json.dumps(manifest, ensure_ascii=False, indent=2).encode("utf-8"), ) - await self.backend.write_file( + await self._write_file( f"{self.logs_dir}/actions.jsonl", b"".join( (json.dumps(asdict(action), ensure_ascii=False) + "\n").encode("utf-8") for action in self.actions @@ -1004,9 +975,7 @@ class DockerReMeSandbox: if include_candidate: command.extend(["-C", "/workspace", "candidate"]) await self._exec_checked("export_case", command) - payload = await self.backend.read_file(EXPORT_ARCHIVE) - target.parent.mkdir(parents=True, exist_ok=True) - target.write_bytes(payload) + self._save_host_file(target, await self.backend.read_file(EXPORT_ARCHIVE)) if profile == "analysis": await self._export_analysis_logs(artifacts) return target @@ -1072,8 +1041,9 @@ class DockerReMeSandbox: """Destroy the case container. Export before calling this method.""" async with self._operation_lock: if not self._closed: - self._closed = True + self._reusable = False await self.workspace.close() + self._closed = True async def __aenter__(self) -> "DockerReMeSandbox": return self diff --git a/deploy/benchmark_docker/models.py b/deploy/benchmark_docker/models.py index 53c40187..77dea88a 100644 --- a/deploy/benchmark_docker/models.py +++ b/deploy/benchmark_docker/models.py @@ -70,8 +70,4 @@ class EvaluationQuery: question: str golden_answer: Any answer_arguments: dict[str, Any] = field(default_factory=dict) - judge_arguments: dict[str, Any] = field(default_factory=dict) answer_job: str = "agentic_answer" - judge_answer_argument: str = "agent_answer" - score_path: str = "answer" - score_mapping: dict[str, float] | None = field(default_factory=lambda: {"yes": 1.0, "no": 0.0}) diff --git a/deploy/benchmark_docker/pool.py b/deploy/benchmark_docker/pool.py index cdd41101..fd70e5ac 100644 --- a/deploy/benchmark_docker/pool.py +++ b/deploy/benchmark_docker/pool.py @@ -85,15 +85,19 @@ class BenchmarkDockerPool: if plan.export_evaluation_to is not None and not plan.queries: raise ValueError("evaluation export requires queries") self._check_acyclic(by_id) + reused_case_ids = {plan.reuse_workspace_from for plan in plans if plan.reuse_workspace_from is not None} workers = await self.factory.create_cases( [f"pool-{index}" for index in range(min(self.concurrency, len(plans)))], concurrency=self.concurrency, ) - available: asyncio.Queue[DockerReMeSandbox] = asyncio.Queue() + available: asyncio.Queue[DockerReMeSandbox | None] = asyncio.Queue() for worker in workers: available.put_nowait(worker) claimed = 0 + live_slots = len(workers) + next_worker_index = len(workers) + exhaustion_error = "no available workers" try: with TemporaryDirectory(prefix="reme-benchmark-pool-") as temporary: @@ -103,7 +107,7 @@ class BenchmarkDockerPool: tasks: dict[str, asyncio.Task[CaseResult]] = {} async def execute(plan: CasePlan) -> CaseResult: - nonlocal claimed + nonlocal claimed, live_slots, next_worker_index, exhaustion_error result = CaseResult(plan.case_id) source = plan.workspace_source if plan.reuse_workspace_from is not None: @@ -113,8 +117,15 @@ class BenchmarkDockerPool: source = parent.workspace_snapshot worker = await available.get() + if worker is None: + return CaseResult(plan.case_id, error=exhaustion_error) claimed += 1 + def record_lifecycle_error(phase: str, exc: Exception) -> str: + detail = f"{phase} failed: {type(exc).__name__}: {exc}" + result.error = f"{result.error}; {detail}" if result.error else detail + return detail + async def retain_failed_build_log() -> None: try: result.build_log = await worker.read_build_log() @@ -145,9 +156,7 @@ class BenchmarkDockerPool: await worker.export_build_log(plan.export_build_log_to) # Capture before queries, since answering may change the # workspace. Dependents receive construction state. - if plan.export_workspace_to is not None or any( - other.reuse_workspace_from == plan.case_id for other in plans - ): + if plan.export_workspace_to is not None or plan.case_id in reused_case_ids: destination = plan.export_workspace_to or snapshot_root / f"{plan.case_id}.tar.gz" result.workspace_snapshot = await worker.export_workspace(destination) if plan.queries: @@ -164,14 +173,37 @@ class BenchmarkDockerPool: if worker.reusable: available.put_nowait(worker) else: - # AgentScope reports a shell timeout while its process - # may still be running. Destroy the container before - # another case can reset and use this pool slot. - await worker.close() - if claimed < len(plans) and not asyncio.current_task().cancelling(): - replacement = await self.factory.create_case(f"pool-{len(workers)}") - workers.append(replacement) - available.put_nowait(replacement) + # A timed-out command may still be running in the container. + # Retire it before allowing another case into this slot. + closed = False + try: + await worker.close() + except Exception as exc: + record_lifecycle_error("worker close", exc) + else: + closed = True + + replaced = False + if closed and claimed < len(plans) and not asyncio.current_task().cancelling(): + worker_index = next_worker_index + next_worker_index += 1 + try: + replacement = await self.factory.create_case(f"pool-{worker_index}") + except Exception as exc: + exhaustion_error = ( + f"no available workers: {record_lifecycle_error('replacement', exc)}" + ) + else: + workers.append(replacement) + available.put_nowait(replacement) + replaced = True + if not replaced: + live_slots -= 1 + if live_slots == 0: + # Every future case, including one waiting on a + # parent, must be able to leave available.get(). + for _ in range(len(plans) - claimed): + available.put_nowait(None) for plan in plans: tasks[plan.case_id] = asyncio.create_task(execute(plan)) diff --git a/deploy/benchmark_docker/worker.py b/deploy/benchmark_docker/worker.py index c5f764ef..783879e9 100644 --- a/deploy/benchmark_docker/worker.py +++ b/deploy/benchmark_docker/worker.py @@ -192,6 +192,7 @@ def _validate_query_id(query_id: Any) -> str: def _clear_query_state(app: Any) -> None: """Drop request-scoped metadata that must not leak into the next query.""" app.context.metadata.pop("tool_contexts", None) + app.context.metadata.pop("__search_call_budgets", None) async def _run(request: dict[str, Any]) -> dict[str, Any]: diff --git a/tests/unit/test_benchmark_docker.py b/tests/unit/test_benchmark_docker.py index 2c82ab48..e4bc7e7d 100644 --- a/tests/unit/test_benchmark_docker.py +++ b/tests/unit/test_benchmark_docker.py @@ -165,6 +165,116 @@ async def test_pool_keeps_build_log_on_failure(tmp_path, raises): assert events.index(("failed", "read_build_log")) < events.index(("pool-0", "close")) +@pytest.mark.asyncio +async def test_pool_replacement_failure_preserves_healthy_cases(): + events: list[tuple] = [] + + class Worker(FakeCase): + async def run_build(self, jobs): + if self.case_id == "broken": + self.reusable = False + raise RuntimeError("command failed") + await asyncio.sleep(0.02) + self.events.append((self.case_id, "finished")) + return {"success": True} + + class Factory(FakeFactory): + async def create_cases(self, case_ids, *, concurrency): + return [Worker(name, self.events) for name in case_ids] + + async def create_case(self, case_id): + raise RuntimeError("replacement failed") + + plans = [CasePlan(case_id, build_jobs=(JobRequest("build"),)) for case_id in ("broken", "healthy", "waiting")] + results = await asyncio.wait_for(BenchmarkDockerPool(Factory(events), concurrency=2).run(plans), timeout=2) + + assert [result.case_id for result in results] == ["broken", "healthy", "waiting"] + assert "command failed" in results[0].error + assert "replacement failed" in results[0].error + assert results[1].success and results[2].success + assert ("healthy", "finished") in events and ("waiting", "finished") in events + + +@pytest.mark.asyncio +async def test_pool_exhaustion_fails_waiting_cases(): + events: list[tuple] = [] + + class BrokenCase(FakeCase): + async def run_build(self, jobs): + self.reusable = False + raise RuntimeError("command failed") + + class Factory(FakeFactory): + async def create_cases(self, case_ids, *, concurrency): + return [BrokenCase(name, self.events) for name in case_ids] + + async def create_case(self, case_id): + raise RuntimeError("replacement failed") + + plans = [CasePlan(case_id, build_jobs=(JobRequest("build"),)) for case_id in ("broken", "waiting-1", "waiting-2")] + results = await asyncio.wait_for(BenchmarkDockerPool(Factory(events), concurrency=1).run(plans), timeout=2) + + assert "replacement failed" in results[0].error + assert all("no available workers" in result.error for result in results[1:]) + assert not any(case_id.startswith("waiting") for case_id, action in events if action == "reset") + + +@pytest.mark.asyncio +async def test_pool_exhaustion_unblocks_dependent_cases(): + events: list[tuple] = [] + + class BrokenCase(FakeCase): + async def run_build(self, jobs): + self.reusable = False + raise RuntimeError("command failed") + + class Factory(FakeFactory): + async def create_cases(self, case_ids, *, concurrency): + return [BrokenCase(name, self.events) for name in case_ids] + + async def create_case(self, case_id): + raise RuntimeError("replacement failed") + + plans = [ + CasePlan("dependent", reuse_workspace_from="waiting"), + CasePlan("broken", build_jobs=(JobRequest("build"),)), + CasePlan("waiting", build_jobs=(JobRequest("build"),)), + ] + results = await asyncio.wait_for(BenchmarkDockerPool(Factory(events), concurrency=1).run(plans), timeout=2) + + assert results[0].error == "source case 'waiting' failed" + assert "replacement failed" in results[1].error + assert "no available workers" in results[2].error + + +@pytest.mark.asyncio +async def test_pool_close_failure_retires_slot_without_replacement(): + events: list[tuple] = [] + + class BrokenCase(FakeCase): + async def run_build(self, jobs): + self.reusable = False + raise RuntimeError("command failed") + + async def close(self): + self.events.append((self.name, "close failed")) + raise RuntimeError("close failed") + + class Factory(FakeFactory): + async def create_cases(self, case_ids, *, concurrency): + return [BrokenCase(name, self.events) for name in case_ids] + + async def create_case(self, case_id): + pytest.fail("replacement must not start while the old container may still be running") + + plans = [CasePlan(case_id, build_jobs=(JobRequest("build"),)) for case_id in ("broken", "next")] + results = await asyncio.wait_for(BenchmarkDockerPool(Factory(events), concurrency=1).run(plans), timeout=2) + + assert "close failed" in results[0].error + assert results[1].error == "no available workers" + assert ("next", "build", 1) not in events + + @pytest.mark.asyncio async def test_pool_cancels_case_before_closing_container(): started = asyncio.Event() @@ -264,6 +374,52 @@ async def test_build_log_remains_bound_after_application_initializes_logger(tmp_ assert "[sandbox] build phase finished" in log_text +@pytest.mark.asyncio +async def test_worker_clears_search_budget_between_queries(tmp_path, monkeypatch): + import reme.reme as reme_module + + monkeypatch.setattr(sandbox_worker, "_prepare_runtime", lambda _request: (tmp_path, {})) + seen_metadata: list[dict] = [] + apps = [] + + class FakeReMe: + def __init__(self, **_kwargs): + from types import SimpleNamespace + + self.context = SimpleNamespace(metadata={"shared": "retained"}) + apps.append(self) + + async def start(self): + pass + + async def close(self): + pass + + async def run_job(app, _job, _arguments): + seen_metadata.append(dict(app.context.metadata)) + app.context.metadata["tool_contexts"] = {"same-id": {"seen": True}} + app.context.metadata["__search_call_budgets"] = {"same-id": {"count": 1}} + return {"success": True, "answer": "ok", "metadata": {}, "error": None} + + monkeypatch.setattr(reme_module, "ReMe", FakeReMe) + monkeypatch.setattr(sandbox_worker, "_run_job_on_app", run_job) + + result = await sandbox_worker._run_queries( + { + "case_root": str(tmp_path), + "queries": [ + {"query_id": "q1", "question": "first"}, + {"query_id": "q2", "question": "second"}, + ], + }, + ) + + assert result["success"] + assert len(apps) == 1 + assert seen_metadata == [{"shared": "retained"}, {"shared": "retained"}] + assert apps[0].context.metadata == {"shared": "retained"} + + @pytest.mark.asyncio async def test_worker_checkpoint_commits_only_memory_paths(tmp_path, monkeypatch): workspace = tmp_path / "workspace" @@ -395,6 +551,142 @@ async def test_backend_timeout_prevents_direct_case_reset(tmp_path): assert workspace.closed +@pytest.mark.asyncio +async def test_external_timeout_prevents_case_reset(tmp_path): + started = asyncio.Event() + + class BlockingBackend(FakeBackend): + async def exec_shell(self, command, *, cwd=None, timeout=None): + if "/workspace/harness/worker.py" in command: + self.commands.append(command) + started.set() + await asyncio.Event().wait() + return await super().exec_shell(command, cwd=cwd, timeout=timeout) + + workspace = FakeWorkspace(BlockingBackend()) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: workspace) + case = await factory.create_case("timed-out") + try: + build = asyncio.create_task(case.run_build([JobRequest("build")])) + await asyncio.wait_for(started.wait(), timeout=2) + with pytest.raises(TimeoutError): + await asyncio.wait_for(build, timeout=0.01) + assert not case.reusable + with pytest.raises(SandboxCommandError, match="no longer reusable"): + await case.reset_case("next") + inbox_count = sum("/inbox/" in path for path in workspace.backend.files) + with pytest.raises(SandboxCommandError, match="no longer reusable"): + await case.run_build([JobRequest("build")]) + assert sum("/inbox/" in path for path in workspace.backend.files) == inbox_count + finally: + await case.close() + assert workspace.closed + + +@pytest.mark.asyncio +async def test_cancelled_request_upload_prevents_case_reset(tmp_path): + started = asyncio.Event() + + class BlockingBackend(FakeBackend): + async def write_file(self, path, data): + if "/inbox/" in path: + started.set() + await asyncio.Event().wait() + await super().write_file(path, data) + + workspace = FakeWorkspace(BlockingBackend()) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: workspace) + case = await factory.create_case("cancelled-upload") + try: + build = asyncio.create_task(case.run_build([JobRequest("build")])) + await asyncio.wait_for(started.wait(), timeout=2) + build.cancel() + with pytest.raises(asyncio.CancelledError): + await build + assert not case.reusable + with pytest.raises(SandboxCommandError, match="no longer reusable"): + await case.reset_case("next") + finally: + await case.close() + assert workspace.closed + + +@pytest.mark.asyncio +async def test_backend_write_error_prevents_case_reset(tmp_path): + class FailingBackend(FakeBackend): + async def write_file(self, path, data): + if "/inbox/" in path: + raise RuntimeError("Docker upload stream failed") + await super().write_file(path, data) + + workspace = FakeWorkspace(FailingBackend()) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: workspace) + case = await factory.create_case("write-error") + try: + with pytest.raises(RuntimeError, match="Docker upload stream failed"): + await case.run_build([JobRequest("build")]) + assert not case.reusable + with pytest.raises(SandboxCommandError, match="no longer reusable"): + await case.reset_case("next") + finally: + await case.close() + assert workspace.closed + + +@pytest.mark.asyncio +async def test_backend_exec_error_prevents_case_reset(tmp_path): + class FailingBackend(FakeBackend): + async def exec_shell(self, command, *, cwd=None, timeout=None): + if "/workspace/harness/worker.py" in command: + raise RuntimeError("Docker exec stream failed") + return await super().exec_shell(command, cwd=cwd, timeout=timeout) + + workspace = FakeWorkspace(FailingBackend()) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: workspace) + case = await factory.create_case("exec-error") + try: + with pytest.raises(RuntimeError, match="Docker exec stream failed"): + await case.run_build([JobRequest("build")]) + assert not case.reusable + with pytest.raises(SandboxCommandError, match="no longer reusable"): + await case.reset_case("next") + finally: + await case.close() + assert workspace.closed + + +@pytest.mark.asyncio +async def test_cancelled_close_can_be_retried(tmp_path): + started = asyncio.Event() + + class InterruptibleWorkspace(FakeWorkspace): + def __init__(self): + super().__init__(FakeBackend()) + self.close_calls = 0 + + async def close(self): + self.close_calls += 1 + if self.close_calls == 1: + started.set() + await asyncio.Event().wait() + self.closed = True + + workspace = InterruptibleWorkspace() + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: workspace) + case = await factory.create_case("close-retry") + closing = asyncio.create_task(case.close()) + await asyncio.wait_for(started.wait(), timeout=2) + closing.cancel() + with pytest.raises(asyncio.CancelledError): + await closing + assert not case.reusable + assert not workspace.closed + + await case.close() + assert workspace.closed + assert workspace.close_calls == 2 + + @pytest.mark.asyncio async def test_source_candidate_installs_plugin_and_closes_container(tmp_path): backend = FakeBackend() @@ -492,6 +784,32 @@ async def test_query_only_evaluation_export_omits_missing_build_log(tmp_path): await case.close() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("method", "required", "members"), + [ + ("export_memory_construction", "build_log/build.log", ["reme_workspace", "build_log"]), + ("export_queries", "queries/summary.json", ["queries"]), + ("export_query", "queries/q1/result.json", ["queries/q1"]), + ], +) +async def test_archive_exports_check_required_artifacts(tmp_path, method, required, members): + backend = FakeBackend() + case = await _fake_docker_factory(tmp_path, lambda _id, _image, _env: FakeWorkspace(backend)).create_case("example") + destination = tmp_path / "exports" / "result.tar.gz" + export = getattr(case, method) + args = ("q1", destination) if method == "export_query" else (destination,) + try: + with pytest.raises(SandboxCommandError, match="missing required artifact"): + await export(*args) + backend.files[f"/workspace/case/{required}"] = b"{}" + assert await export(*args) == destination + assert destination.read_bytes() == b"workspace archive" + assert ["tar", "-czf", "/tmp/reme-case-export.tar.gz", "-C", "/workspace/case", *members] in backend.commands + finally: + await case.close() + + @pytest.mark.asyncio async def test_default_analysis_export_saves_logs_outside_archive(tmp_path): class AnalysisBackend(FakeBackend): @@ -575,3 +893,34 @@ async def test_cancelled_batch_creation_closes_completed_workspace(tmp_path): with pytest.raises(asyncio.CancelledError): await task assert first.closed and second.closed + + +@pytest.mark.asyncio +async def test_failed_case_creation_preserves_initialization_error_when_close_fails(tmp_path): + class BrokenWorkspace(FakeWorkspace): + async def initialize(self): + raise RuntimeError("initialization failed") + + async def close(self): + raise RuntimeError("close failed") + + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: BrokenWorkspace(FakeBackend())) + with pytest.raises(RuntimeError, match="initialization failed") as failure: + await factory.create_case("broken") + assert any("close failed" in note for note in failure.value.__notes__) + + +@pytest.mark.asyncio +async def test_failed_batch_creation_preserves_creation_error_when_close_fails(tmp_path): + class CloseFailsWorkspace(FakeWorkspace): + async def close(self): + raise RuntimeError("close failed") + + class InitializeFailsWorkspace(FakeWorkspace): + async def initialize(self): + raise RuntimeError("initialization failed") + + workspaces = iter((CloseFailsWorkspace(FakeBackend()), InitializeFailsWorkspace(FakeBackend()))) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: next(workspaces)) + with pytest.raises(RuntimeError, match="initialization failed"): + await factory.create_cases(["first", "second"], concurrency=2)