diff --git a/deploy/benchmark_docker/README.md b/deploy/benchmark_docker/README.md index b1ec0ad4..8b74bdf1 100644 --- a/deploy/benchmark_docker/README.md +++ b/deploy/benchmark_docker/README.md @@ -54,17 +54,55 @@ asyncio.run(main()) Each worker completes one case's build and queries before taking the next case. A case can set `reuse_workspace_from="case-1"` to receive the source case's workspace **as it stood immediately after construction**, or provide -`workspace_source` pointing at a host directory or `export_workspace()` -archive. A reused workspace may be queried without additional build jobs. +`workspace_source` pointing at a host workspace directory or a saved archive. +A reused workspace may be queried without additional build jobs. `export_workspace_to` persists a portable snapshot beyond the pool run. +To build memory once and query it in later runs using the same factory: + +```python +memory = Path("artifacts/memory.tar.gz") +pool = BenchmarkDockerPool(factory, concurrency=2) +built = (await pool.run([ + CasePlan("build-memory", build_jobs=(JobRequest("auto_memory", { + "session_id": "s1", "messages": [], + }),), export_workspace_to=memory), +]))[0] +if not built.success: + raise RuntimeError(built.error or "memory build failed") + +# This may run later in another process with a compatible factory/config. +answers = await pool.run([ + CasePlan("query-memory", workspace_source=memory, + queries=(EvaluationQuery("q1", "Question?", "Private answer"),)), +]) +``` + +Use compatible candidate code, plugins, and workspace path configuration when +restoring memory. Uploading a snapshot does not change the receiving config. +`export_workspace()` archives contain workspace contents directly. Case archives +from `export_memory_construction()`, `export_evaluation()`, and `export()` carry +a `manifest.json` format marker and restore only `reme_workspace/`, including +its Git history; build logs and query artifacts are excluded from the upload. +Evaluation exports contain the workspace after queries; use a construction or +workspace export before queries to preserve the original built memory. +Query-only case exports cannot be uploaded as workspaces. Older full case +archives remain supported. Older partial archives with a `reme_workspace` +wrapper but no manifest are rejected; to reuse one, extract it and supply its +`reme_workspace` directory. + `export_evaluation_to` saves the workspace and query artifacts for both build and query-only cases; it includes `build_log` when that phase ran. Set `export_build_log_to` to retain the construction log separately. On a failed build, the pool also returns its text in `CaseResult.build_log`. A timed-out AgentScope shell command makes its container unusable: the pool closes it after retaining the log and creates a replacement when more cases need a worker. -The build log also records configuration and Application construction failures, -including their tracebacks. +Replacement requires successful cleanup of the old worker. Cleanup failures +are reported in `CaseResult.error` and make `success` false for cases that used +that worker; failures for an unused worker apply to the whole run. Cancellation +waits for owned-container cleanup before propagating. A failed `close()` retains +unfinished Docker handles so another call can retry removal. +The build log also records configuration, construction, build job, checkpoint, +and Application lifecycle failures, including available tracebacks. For direct control, call `create_case()` or `create_cases()`, then use `run_job()`, `run_build()`, `run_queries()`, `export_workspace()`, diff --git a/deploy/benchmark_docker/_cleanup.py b/deploy/benchmark_docker/_cleanup.py new file mode 100644 index 00000000..6de0cd40 --- /dev/null +++ b/deploy/benchmark_docker/_cleanup.py @@ -0,0 +1,52 @@ +"""Finish owned-resource cleanup before propagating caller cancellation.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Iterable + + +def format_error(error: BaseException) -> str: + """Include cleanup notes when an exception becomes a result string.""" + return "; ".join([f"{type(error).__name__}: {error}", *getattr(error, "__notes__", ())]) + + +async def finish_cleanup( + operations: Iterable[Awaitable[None]], +) -> tuple[list[None | BaseException], asyncio.CancelledError | None]: + """Collect all failures, shielding teardown from repeated cancellation. + + The caller must propagate the returned cancellation after reporting errors. + Individual cleanup operations may still fail or cancel themselves. + """ + + async def collect(operation: Awaitable[None]) -> None | BaseException: + try: + await operation + except BaseException as exc: + # gather(return_exceptions=True) replaces CancelledError and loses + # its notes. Keep the original exception for cleanup diagnostics. + return exc + return None + + cleanup = asyncio.gather(*(collect(operation) for operation in operations)) + cancelled = None + while not cleanup.done(): + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError as exc: + cancelled = exc + return cleanup.result(), cancelled + + +async def cleanup_on_error(operations: Iterable[tuple[str, Awaitable[None]]], error: BaseException) -> None: + """Attach teardown failures to the primary error before propagating cancellation.""" + operations = list(operations) + outcomes, cancelled = await finish_cleanup(operation for _, operation in operations) + for (name, _), outcome in zip(operations, outcomes): + if isinstance(outcome, BaseException): + error.add_note(f"{name} close failed: {format_error(outcome)}") + if cancelled is not None and not isinstance(error, asyncio.CancelledError): + for note in getattr(error, "__notes__", ()): + cancelled.add_note(note) + raise cancelled from error diff --git a/deploy/benchmark_docker/candidate.py b/deploy/benchmark_docker/candidate.py index d7db4383..e8d28909 100644 --- a/deploy/benchmark_docker/candidate.py +++ b/deploy/benchmark_docker/candidate.py @@ -7,6 +7,7 @@ import fnmatch import gzip import hashlib import io +import os from pathlib import Path from pathlib import PurePosixPath import tarfile @@ -98,14 +99,26 @@ class SourceSnapshot: raise ValueError(f"candidate source is not a directory: {root}") files: list[tuple[Path, Path]] = [] - for path in root.rglob("*"): - relative = path.relative_to(root) - if _is_excluded(relative, excludes): - continue - if path.is_symlink(): - raise ValueError(f"candidate source contains a symbolic link: {relative}") - if path.is_file(): - files.append((relative, path)) + # Component matches exclude every descendant. Full-path matches may + # exclude only the directory itself, so they must not prune traversal. + directory_patterns = tuple(pattern for pattern in excludes if "/" not in pattern) + for current, directories, filenames in os.walk(root, topdown=True, followlinks=False): + directory = Path(current) + for name in directories[:]: + path = directory / name + if any(fnmatch.fnmatch(name, pattern) for pattern in directory_patterns): + directories.remove(name) + elif path.is_symlink() and not _is_excluded(path.relative_to(root), excludes): + raise ValueError(f"candidate source contains a symbolic link: {path.relative_to(root)}") + for name in filenames: + path = directory / name + relative = path.relative_to(root) + if _is_excluded(relative, excludes): + continue + if path.is_symlink(): + raise ValueError(f"candidate source contains a symbolic link: {relative}") + if path.is_file(): + files.append((relative, path)) raw = io.BytesIO() with gzip.GzipFile(fileobj=raw, mode="wb", mtime=0) as compressed: diff --git a/deploy/benchmark_docker/direct_docker_workspace.py b/deploy/benchmark_docker/direct_docker_workspace.py index 48782485..5a6f8827 100644 --- a/deploy/benchmark_docker/direct_docker_workspace.py +++ b/deploy/benchmark_docker/direct_docker_workspace.py @@ -7,8 +7,12 @@ would add an unrelated network-dependent image layer and process. from __future__ import annotations +import asyncio +import os +import sys from typing import Any +from aiodocker.exceptions import DockerError from agentscope.workspace import DockerWorkspace @@ -17,24 +21,72 @@ class DirectDockerWorkspace(DockerWorkspace): _image_tag: str _backend: Any + _container: Any + _client: Any is_alive: bool + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._lifecycle_lock = asyncio.Lock() + async def _build_or_reuse_image(self) -> None: """Select the caller-provided image instead of deriving a gateway image.""" self._image_tag = self.base_image async def initialize(self) -> None: """Start one Docker container and expose its filesystem/exec backend.""" - if self.is_alive: - return - await self._provision_backend() - backend = self.get_backend() - result = await backend.exec_shell(["mkdir", "-p", self.workdir], cwd="/") - if not result.ok(): - await self._teardown_backend() + async with self._lifecycle_lock: + if self.is_alive: + return + if self._container is not None or self._client is not None: + raise RuntimeError("Docker workspace cleanup must complete before initialization") + try: + await self._provision_backend() + backend = self.get_backend() + result = await backend.exec_shell(["mkdir", "-p", self.workdir], cwd="/") + if not result.ok(): + raise RuntimeError( + f"failed to create Docker workspace {self.workdir!r}: " + f"{result.stderr.decode('utf-8', errors='replace')}", + ) + except BaseException as exc: + try: + await self._teardown_backend() + except Exception as close_exc: + exc.add_note(f"workspace close failed: {type(close_exc).__name__}: {close_exc}") + raise + self.is_alive = True + + async def _teardown_backend(self) -> None: + """Remove the container, preserving handles until cleanup is confirmed.""" + self.is_alive = False + if self._container is not None: + # Preserve AgentScope's bind-mount ownership repair on native Linux. + if self.host_workdir is not None and sys.platform == "linux" and self._backend is not None: + try: + await self._backend.exec_shell( + ["chown", "-R", f"{os.getuid()}:{os.getgid()}", self.workdir], + timeout=10.0, + ) + except Exception: + pass + try: + # Force deletion also stops a running container; a separate kill + # would add a failure point without proving that removal succeeded. + await self._container.delete(force=True) + except DockerError as exc: + if exc.status != 404: + raise + self._container = None self._backend = None - raise RuntimeError( - f"failed to create Docker workspace {self.workdir!r}: " - f"{result.stderr.decode('utf-8', errors='replace')}", - ) - self.is_alive = True + if self._client is not None: + await self._client.close() + self._client = None + self._backend = None + + async def close(self) -> None: + """Report cleanup errors and allow another call to retry unfinished work.""" + # The inherited close clears backend state in a finally block. This + # direct workspace has no gateway and must retain state on failure. + async with self._lifecycle_lock: + await self._teardown_backend() diff --git a/deploy/benchmark_docker/docker_sandbox.py b/deploy/benchmark_docker/docker_sandbox.py index 73f5b634..92a4b632 100644 --- a/deploy/benchmark_docker/docker_sandbox.py +++ b/deploy/benchmark_docker/docker_sandbox.py @@ -15,6 +15,7 @@ import tarfile from typing import Any, Callable, Literal, Protocol import uuid +from ._cleanup import cleanup_on_error from .candidate import ImageCandidate, SourceCandidate from .models import ActionRecord, EvaluationQuery, JobRequest, JobResult @@ -28,6 +29,7 @@ WORKSPACE_ARCHIVE = "/tmp/reme-workspace-upload.tar.gz" WORKSPACE_EXPORT_ARCHIVE = "/tmp/reme-workspace-export.tar.gz" RUNTIME_LAYOUT = f"{CASE_ROOT}/runtime-layout.json" _CASE_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$") +_CASE_ARCHIVE_FORMAT = "reme-case-v1" _PRIVATE_EVALUATION_FIELDS = frozenset( {"query", "question", "golden_answer", "question_type", "rubric", "judge_arguments"}, ) @@ -110,19 +112,60 @@ def _workspace_member_path(name: str) -> PurePosixPath | None: return path +def _case_archive_workspace_root( + archive: tarfile.TarFile, + members: dict[PurePosixPath, tarfile.TarInfo], +) -> PurePosixPath | None: + """Distinguish marked case artifacts from raw workspace contents.""" + workspace_root = PurePosixPath("reme_workspace") + has_workspace = any(path.parts[0] == workspace_root.name for path in members) + manifest = members.get(PurePosixPath("manifest.json")) + if manifest is None: + if has_workspace: + raise ValueError( + "unmarked case archive; upload its reme_workspace directory or recreate it with a current case export", + ) + return None + if not manifest.isfile(): + raise ValueError("workspace archive manifest.json must be a regular file") + content = archive.extractfile(manifest) + assert content is not None + try: + payload = json.loads(content.read()) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + if has_workspace: + raise ValueError("case archive contains an invalid manifest.json") from exc + return None # A raw workspace may contain an unrelated manifest file. + if not isinstance(payload, dict) or "archive_format" not in payload: + if has_workspace and not isinstance(payload, dict): + raise ValueError("case archive manifest.json must be an object") + # Older full case exports have a manifest without archive_format. + return workspace_root if has_workspace else None + if payload["archive_format"] != _CASE_ARCHIVE_FORMAT: + raise ValueError(f"unsupported case archive format: {payload['archive_format']!r}") + root = payload.get("workspace_root") + if root is None: + raise ValueError("case archive has no reusable workspace") + if root != workspace_root.as_posix(): + raise ValueError(f"case archive has an invalid workspace root: {root!r}") + if not has_workspace: + raise ValueError("case archive workspace is missing") + return workspace_root + + def _archive_workspace(source: str | Path) -> bytes: """Create a safe, deterministic gzip tar archive from a host workspace. ``source`` may be a workspace directory or an archive previously returned - by :meth:`DockerReMeSandbox.export_workspace`. Incoming archives are read - and rebuilt so links, special files, and traversal paths never reach the - sandbox extractor. + by :meth:`DockerReMeSandbox.export_workspace`, or a marked case archive + containing a workspace. Incoming archives are read and rebuilt so links, + special files, and traversal paths never reach the sandbox extractor. """ supplied_path = Path(source).expanduser() if supplied_path.is_symlink(): raise ValueError("workspace source must not be a symbolic link") source_path = supplied_path.resolve(strict=True) - entries: list[tuple[PurePosixPath, tarfile.TarInfo, bytes | None]] = [] + entries: list[tuple[tarfile.TarInfo, bytes | None]] = [] if source_path.is_dir(): for item in sorted(source_path.rglob("*"), key=lambda value: value.relative_to(source_path).as_posix()): relative = item.relative_to(source_path) @@ -130,67 +173,58 @@ def _archive_workspace(source: str | Path) -> bytes: raise ValueError(f"workspace source contains a symbolic link: {relative}") path = _workspace_member_path(relative.as_posix()) assert path is not None + info = tarfile.TarInfo(path.as_posix()) + info.mode = item.stat().st_mode & 0o777 if item.is_dir(): - info = tarfile.TarInfo(path.as_posix()) info.type = tarfile.DIRTYPE - info.mode = item.stat().st_mode & 0o777 - entries.append((path, info, None)) + data = None elif item.is_file(): - info = tarfile.TarInfo(path.as_posix()) - info.size = item.stat().st_size - info.mode = item.stat().st_mode & 0o777 - entries.append((path, info, item.read_bytes())) + data = item.read_bytes() + info.size = len(data) else: raise ValueError(f"workspace source contains unsupported file: {relative}") + entries.append((info, data)) elif source_path.is_file(): with tarfile.open(source_path, mode="r:*") as archive: - members = archive.getmembers() - normalized_paths = [path for member in members if (path := _workspace_member_path(member.name)) is not None] - # ``export()`` archives the whole case, while ``export_workspace`` - # archives only the workspace contents. Accept both forms so an - # existing case artifact can seed independent query sandboxes. - case_export = PurePosixPath("manifest.json") in normalized_paths and any( - path.parts and path.parts[0] == "reme_workspace" for path in normalized_paths - ) - seen: set[PurePosixPath] = set() - for member in members: + members: dict[PurePosixPath, tarfile.TarInfo] = {} + for member in archive.getmembers(): path = _workspace_member_path(member.name) if path is None: continue - if case_export: - if not path.parts or path.parts[0] != "reme_workspace": + if path in members: + raise ValueError(f"workspace archive contains duplicate path: {path}") + members[path] = member + workspace_root = _case_archive_workspace_root(archive, members) + for path, member in members.items(): + if workspace_root is not None: + if path != workspace_root and workspace_root not in path.parents: continue - if len(path.parts) == 1: + if path == workspace_root: if not member.isdir(): raise ValueError("case export has a non-directory reme_workspace entry") continue - path = PurePosixPath(*path.parts[1:]) - if path in seen: - raise ValueError(f"workspace archive contains duplicate path: {path}") - seen.add(path) + path = path.relative_to(workspace_root) + info = tarfile.TarInfo(path.as_posix()) + info.mode = member.mode & 0o777 if member.isdir(): - info = tarfile.TarInfo(path.as_posix()) info.type = tarfile.DIRTYPE - info.mode = member.mode & 0o777 - entries.append((path, info, None)) + data = None elif member.isfile(): content = archive.extractfile(member) if content is None: raise ValueError(f"workspace archive cannot read file: {path}") data = content.read() - info = tarfile.TarInfo(path.as_posix()) info.size = len(data) - info.mode = member.mode & 0o777 - entries.append((path, info, data)) else: raise ValueError(f"workspace archive contains unsupported entry: {member.name!r}") + entries.append((info, data)) else: raise ValueError(f"workspace source is neither a directory nor a file: {source_path}") raw = io.BytesIO() with gzip.GzipFile(fileobj=raw, mode="wb", mtime=0) as compressed: with tarfile.open(fileobj=compressed, mode="w") as archive: - for _path, info, data in entries: + for info, data in entries: info.mtime = 0 info.uid = info.gid = 0 info.uname = info.gname = "" @@ -270,10 +304,7 @@ class DockerReMeSandboxFactory: try: await case.initialize() 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}") + await cleanup_on_error([(f"workspace {container_id!r}", workspace.close())], exc) raise return case @@ -285,30 +316,38 @@ class DockerReMeSandboxFactory: raise ValueError("case_ids must be unique") semaphore = asyncio.Semaphore(concurrency) - async def create(case_id: str) -> DockerReMeSandbox: - async with semaphore: - return await self.create_case(case_id) - - async def close_created(results: list[DockerReMeSandbox | BaseException]) -> None: - await asyncio.gather( - *(result.close() for result in results if isinstance(result, DockerReMeSandbox)), - return_exceptions=True, - ) + async def create(case_id: str) -> DockerReMeSandbox | BaseException: + try: + async with semaphore: + return await self.create_case(case_id) + except BaseException as exc: + # Preserve cleanup notes on CancelledError instead of letting + # gather replace it with a fresh cancellation exception. + return exc tasks = [asyncio.create_task(create(case_id)) for case_id in case_ids] + error: BaseException | None = None try: results = await asyncio.gather(*tasks, return_exceptions=True) - except BaseException: + except BaseException as exc: + error = exc for task in tasks: - task.cancel() + if not task.done() and not task.cancelling(): + task.cancel() results = await asyncio.gather(*tasks, return_exceptions=True) - await close_created(results) - raise + created = [result for result in results if isinstance(result, DockerReMeSandbox)] failures = [result for result in results if isinstance(result, BaseException)] - if failures: - await close_created(results) - raise failures[0] - return [result for result in results if isinstance(result, DockerReMeSandbox)] + if error is None: + if not failures: + return created + error = failures.pop(0) + for failure in failures: + error.add_note(f"additional creation failure: {type(failure).__name__}: {failure}") + for failure in failures: + for note in getattr(failure, "__notes__", ()): + error.add_note(note) + await cleanup_on_error(((f"workspace {case.container_id!r}", case.close()) for case in created), error) + raise error class DockerReMeSandbox: @@ -474,10 +513,7 @@ class DockerReMeSandbox: 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: + except (Exception, asyncio.CancelledError): # Docker may have accepted the upload before the client failed. self._reusable = False raise @@ -492,10 +528,7 @@ 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: + except (Exception, asyncio.CancelledError): # A backend failure may occur after Docker started the command. self._reusable = False raise @@ -687,16 +720,12 @@ class DockerReMeSandbox: """Generate an answer with the candidate's agentic_answer job.""" async with self._operation_lock: result = await self._run_job("agentic_answer", {"query": query, "query_time": query_time}) - await self._write_named_result("answer.json", result) + await self._write_file( + f"{self.results_dir}/answer.json", + json.dumps(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._write_file( - f"{self.results_dir}/{name}", - json.dumps(asdict(result), ensure_ascii=False, default=str).encode("utf-8"), - ) - async def export( self, destination: str | Path, @@ -758,43 +787,52 @@ class DockerReMeSandbox: raise SandboxCommandError(f"{description} is missing required artifact: {path}") from exc async def export_memory_construction(self, destination: str | Path) -> Path: - """Download the post-construction workspace and its build log.""" + """Download the workspace and build log; reusable with upload_workspace.""" async with self._operation_lock: - await self._require_artifact(f"{CASE_ROOT}/build_log/build.log", "memory construction export") return await self._export_archive( destination, "export_memory_construction", CASE_ROOT, ["reme_workspace", "build_log"], + required=("build_log/build.log", "memory construction export"), ) async def export_queries(self, destination: str | Path) -> Path: """Download per-query logs, results, and their aggregate summary.""" async with self._operation_lock: - await self._require_artifact(f"{CASE_ROOT}/queries/summary.json", "query export") - return await self._export_archive(destination, "export_queries", CASE_ROOT, ["queries"]) + return await self._export_archive( + destination, + "export_queries", + CASE_ROOT, + ["queries"], + required=("queries/summary.json", "query export"), + ) async def export_query(self, query_id: str, destination: str | Path) -> Path: """Download one leased query's logs and result without a case summary.""" self._validate_query_artifact_id(query_id) async with self._operation_lock: - query_root = f"{CASE_ROOT}/queries/{query_id}" - await self._require_artifact(f"{query_root}/result.json", "query export") - return await self._export_archive(destination, "export_query", CASE_ROOT, [f"queries/{query_id}"]) + return await self._export_archive( + destination, + "export_query", + CASE_ROOT, + [f"queries/{query_id}"], + required=(f"queries/{query_id}/result.json", "query export"), + ) async def export_evaluation(self, destination: str | Path) -> Path: - """Download workspace and query artifacts, plus the build log when present.""" + """Download reusable workspace and query artifacts, plus any build log.""" async with self._operation_lock: assert self.backend is not None, "sandbox is not initialized" await self._require_artifact(f"{CASE_ROOT}/queries/summary.json", "evaluation export") + members = ["reme_workspace", "queries"] try: await self.backend.read_file(f"{CASE_ROOT}/build_log/build.log") except FileNotFoundError: - has_build_log = False + pass else: - has_build_log = True - members = ["reme_workspace", "queries", *(["build_log"] if has_build_log else [])] + members.append("build_log") return await self._export_archive(destination, "export_evaluation", CASE_ROOT, members) @staticmethod @@ -812,15 +850,27 @@ class DockerReMeSandbox: members: list[str], *, archive_path: str = EXPORT_ARCHIVE, + required: tuple[str, str] | None = None, ) -> Path: """Archive selected sandbox paths and download the archive while locked.""" assert self.backend is not None, "sandbox is not initialized" + if required is not None: + path, description = required + await self._require_artifact(f"{root}/{path}", description) + if root == CASE_ROOT: + await self._write_export_manifest( + name.removeprefix("export_"), + workspace_root="reme_workspace" if "reme_workspace" in members else None, + ) + members = ["manifest.json", *members] 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. + """Copy a host workspace directory or exported memory archive into this case. + Case archives restore only their workspace; logs and query artifacts + are excluded. Query-only archives cannot seed a runtime workspace. By default the existing runtime workspace is deleted first, ensuring a query case receives only the supplied memory state. Set ``clear=False`` only when intentionally merging files into the current workspace. @@ -893,27 +943,16 @@ class DockerReMeSandbox: excludes.append(normalized) return workspace_root_value, excludes - async def _export( - self, - destination: str | Path, - *, - profile: Literal["analysis", "full"], - include_candidate: bool, - ) -> Path: - """Export the active case while the caller holds the operation lock.""" - assert self.backend is not None, "sandbox is not initialized" - target = Path(destination).expanduser().resolve() - artifacts = Path(f"{target}.artifacts") - if profile == "analysis" and artifacts.exists(): - raise FileExistsError(f"analysis artifacts already exist: {artifacts}") - if include_candidate and isinstance(self.candidate, ImageCandidate): - raise ValueError("an image candidate has no source directory to export") + async def _write_export_manifest(self, profile: str, *, workspace_root: str | None) -> None: + """Identify case archive layout without adding artifacts to its workspace.""" candidate_id = ( self.candidate.candidate_id if isinstance(self.candidate, SourceCandidate) else self.candidate.resolved_candidate_id ) manifest = { + "archive_format": _CASE_ARCHIVE_FORMAT, + "workspace_root": workspace_root, "case_id": self.case_id, "container_id": self.container_id, "candidate_id": candidate_id, @@ -931,6 +970,23 @@ class DockerReMeSandbox: f"{CASE_ROOT}/manifest.json", json.dumps(manifest, ensure_ascii=False, indent=2).encode("utf-8"), ) + + async def _export( + self, + destination: str | Path, + *, + profile: Literal["analysis", "full"], + include_candidate: bool, + ) -> Path: + """Export the active case while the caller holds the operation lock.""" + assert self.backend is not None, "sandbox is not initialized" + target = Path(destination).expanduser().resolve() + artifacts = Path(f"{target}.artifacts") + if profile == "analysis" and artifacts.exists(): + raise FileExistsError(f"analysis artifacts already exist: {artifacts}") + if include_candidate and isinstance(self.candidate, ImageCandidate): + raise ValueError("an image candidate has no source directory to export") + await self._write_export_manifest(profile, workspace_root="reme_workspace") await self._write_file( f"{self.logs_dir}/actions.jsonl", b"".join( diff --git a/deploy/benchmark_docker/pool.py b/deploy/benchmark_docker/pool.py index fd70e5ac..92b9bca0 100644 --- a/deploy/benchmark_docker/pool.py +++ b/deploy/benchmark_docker/pool.py @@ -13,6 +13,7 @@ from pathlib import Path from tempfile import TemporaryDirectory from typing import Any +from ._cleanup import finish_cleanup, format_error from .docker_sandbox import DockerReMeSandbox, DockerReMeSandboxFactory from .models import EvaluationQuery, JobRequest @@ -44,7 +45,7 @@ class CaseResult: @property def success(self) -> bool: - """Whether all requested phases completed successfully.""" + """Whether all requested phases and owned-worker cleanup succeeded.""" return ( self.error is None and (self.build is None or self.build.get("success") is True) @@ -94,10 +95,13 @@ class BenchmarkDockerPool: available: asyncio.Queue[DockerReMeSandbox | None] = asyncio.Queue() for worker in workers: available.put_nowait(worker) - claimed = 0 + unclaimed = set(by_id) live_slots = len(workers) next_worker_index = len(workers) exhaustion_error = "no available workers" + worker_results: dict[int, list[CaseResult]] = {} + results: list[CaseResult] = [] + run_error: BaseException | None = None try: with TemporaryDirectory(prefix="reme-benchmark-pool-") as temporary: @@ -107,36 +111,27 @@ class BenchmarkDockerPool: tasks: dict[str, asyncio.Task[CaseResult]] = {} async def execute(plan: CasePlan) -> CaseResult: - nonlocal claimed, live_slots, next_worker_index, exhaustion_error + nonlocal live_slots, next_worker_index, exhaustion_error result = CaseResult(plan.case_id) source = plan.workspace_source if plan.reuse_workspace_from is not None: parent = await tasks[plan.reuse_workspace_from] if parent.workspace_snapshot is None: + unclaimed.discard(plan.case_id) return CaseResult(plan.case_id, error=f"source case {parent.case_id!r} failed") source = parent.workspace_snapshot worker = await available.get() + unclaimed.discard(plan.case_id) if worker is None: return CaseResult(plan.case_id, error=exhaustion_error) - claimed += 1 + worker_results.setdefault(id(worker), []).append(result) def record_lifecycle_error(phase: str, exc: Exception) -> str: - detail = f"{phase} failed: {type(exc).__name__}: {exc}" + detail = f"{phase} failed: {format_error(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() - except Exception as exc: - result.error += f"; build log unavailable: {type(exc).__name__}: {exc}" - if plan.export_build_log_to is not None: - try: - await worker.export_build_log(plan.export_build_log_to) - except Exception as exc: - result.error += f"; build log export failed: {type(exc).__name__}: {exc}" - try: await worker.reset_case(plan.case_id) if source is not None: @@ -146,11 +141,11 @@ class BenchmarkDockerPool: result.build = await worker.run_build(list(plan.build_jobs)) except Exception as exc: result.error = f"{type(exc).__name__}: {exc}" - await retain_failed_build_log() + await self._retain_failed_build_log(worker, plan, result) return result if not result.build.get("success"): result.error = "build failed" - await retain_failed_build_log() + await self._retain_failed_build_log(worker, plan, result) return result if plan.export_build_log_to is not None: await worker.export_build_log(plan.export_build_log_to) @@ -179,12 +174,14 @@ class BenchmarkDockerPool: try: await worker.close() except Exception as exc: - record_lifecycle_error("worker close", exc) + exhaustion_error = ( + f"no available workers: {record_lifecycle_error('worker close', exc)}" + ) else: closed = True replaced = False - if closed and claimed < len(plans) and not asyncio.current_task().cancelling(): + if closed and unclaimed and not asyncio.current_task().cancelling(): worker_index = next_worker_index next_worker_index += 1 try: @@ -202,7 +199,7 @@ class BenchmarkDockerPool: 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): + for _ in unclaimed: available.put_nowait(None) for plan in plans: @@ -211,15 +208,53 @@ class BenchmarkDockerPool: results = await asyncio.gather(*(tasks[plan.case_id] for plan in plans)) finally: for task in tasks.values(): - if not task.done(): + if not task.done() and not task.cancelling(): task.cancel() await asyncio.gather(*tasks.values(), return_exceptions=True) for plan, result in zip(plans, results): if plan.export_workspace_to is None: result.workspace_snapshot = None # The internal temporary archive is about to be deleted. return results + except BaseException as exc: + run_error = exc + raise finally: - await asyncio.gather(*(worker.close() for worker in workers), return_exceptions=True) + await self._close_workers(workers, worker_results, results, run_error) + + @staticmethod + async def _retain_failed_build_log(worker: DockerReMeSandbox, plan: CasePlan, result: CaseResult) -> None: + """Keep diagnostics without replacing the original build failure.""" + try: + result.build_log = await worker.read_build_log() + except Exception as exc: + result.error += f"; build log unavailable: {type(exc).__name__}: {exc}" + if plan.export_build_log_to is not None: + try: + await worker.export_build_log(plan.export_build_log_to) + except Exception as exc: + result.error += f"; build log export failed: {type(exc).__name__}: {exc}" + + @staticmethod + async def _close_workers( + workers: list[DockerReMeSandbox], + worker_results: dict[int, list[CaseResult]], + results: list[CaseResult], + run_error: BaseException | None, + ) -> None: + """Finish all teardown and attach failures to results or the run error.""" + outcomes, cleanup_cancel = await finish_cleanup(worker.close() for worker in workers) + for worker, outcome in zip(workers, outcomes): + if isinstance(outcome, BaseException): + detail = f"worker {worker.container_id!r} final close failed: {format_error(outcome)}" + error = run_error if run_error is not None else cleanup_cancel + if error is not None: + error.add_note(detail) + # An unused worker is still owned by this run; report its + # failure on all results rather than silently dropping it. + for result in worker_results.get(id(worker)) or results: + result.error = f"{result.error}; {detail}" if result.error else detail + if cleanup_cancel is not None and run_error is None: + raise cleanup_cancel @staticmethod def _check_acyclic(plans: dict[str, CasePlan]) -> None: diff --git a/deploy/benchmark_docker/worker.py b/deploy/benchmark_docker/worker.py index 783879e9..3abeeea0 100644 --- a/deploy/benchmark_docker/worker.py +++ b/deploy/benchmark_docker/worker.py @@ -9,6 +9,8 @@ from __future__ import annotations import argparse import asyncio +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager import json import logging import os @@ -137,24 +139,24 @@ async def _run_job_on_app(app: Any, job: str, arguments: dict[str, Any]) -> dict ): response = await app.run_job(job, **arguments) except Exception as exc: # Preserve usage accumulated before a job failure. - return { + result = { "success": False, "answer": "", "metadata": {}, - "token_usage": token_usage, - "job_call_counts": {name: count for name, count in job_call_counts.items() if count}, "error": f"{type(exc).__name__}: {exc}", "exception": _exception_details(exc), "traceback": traceback.format_exc(), } - return { - "success": bool(response.success), - "answer": response.answer, - "metadata": response.metadata, - "token_usage": token_usage, - "job_call_counts": {name: count for name, count in job_call_counts.items() if count}, - "error": None if response.success else str(response.answer), - } + else: + result = { + "success": bool(response.success), + "answer": response.answer, + "metadata": response.metadata, + "error": None if response.success else str(response.answer), + } + result["token_usage"] = token_usage + result["job_call_counts"] = {name: count for name, count in job_call_counts.items() if count} + return result def _atomic_write_json(path: Path, value: Any) -> None: @@ -195,6 +197,37 @@ def _clear_query_state(app: Any) -> None: app.context.metadata.pop("__search_call_budgets", None) +@asynccontextmanager +async def _application_scope(app: Any, failed_results: list[dict[str, Any]] | None = None) -> AsyncIterator[None]: + """Close the app, attaching cleanup diagnostics to already-failed results. + + Propagating exceptions retain cleanup notes; cleanup alone still raises. + """ + error: BaseException | None = None + try: + yield + except BaseException as exc: + error = exc + raise + finally: + try: + await app.close() + except BaseException as exc: + if error is None: + if isinstance(exc, Exception) and failed_results: + diagnostic = { + "error": f"{type(exc).__name__}: {exc}", + "exception": _exception_details(exc), + "traceback": traceback.format_exc(), + } + for result in failed_results: + result.setdefault("cleanup_errors", []).append(diagnostic) + else: + raise + else: + error.add_note("Application close failed:\n" + "".join(traceback.format_exception(exc, chain=False))) + + async def _run(request: dict[str, Any]) -> dict[str, Any]: """Construct an Application, run one job, and close it deterministically.""" job = request["job"] @@ -203,11 +236,19 @@ async def _run(request: dict[str, Any]) -> dict[str, Any]: from reme.reme import ReMe app = ReMe(**app_config) - try: + failed_results: list[dict[str, Any]] = [] + async with _application_scope(app, failed_results): await app.start() - return await _run_job_on_app(app, job, job_args) - finally: - await app.close() + result = await _run_job_on_app(app, job, job_args) + if not result["success"]: + failed_results.append(result) + return result + + +def _append_build_error(build_log: Path, detail: str) -> None: + """Keep build diagnostics even when the candidate reconfigures the logger.""" + with build_log.open("a", encoding="utf-8") as log_file: + log_file.write(detail.rstrip() + "\n") async def _run_build(request: dict[str, Any]) -> dict[str, Any]: @@ -227,38 +268,47 @@ async def _run_build(request: dict[str, Any]) -> dict[str, Any]: if not isinstance(jobs, list) or not jobs: raise ValueError("build request requires a non-empty jobs list") app = ReMe(**app_config) - except Exception: - # Config resolution and Application construction can fail before the - # ReMe logger is usable. Write the traceback to the durable artifact. - with build_log.open("a", encoding="utf-8") as log_file: - traceback.print_exc(file=log_file) + logger = None + results: list[dict[str, Any]] = [] + failed_results: list[dict[str, Any]] = [] + try: + async with _application_scope(app, failed_results): + # Application.__init__ resets the global logger. Attach afterwards. + logger = _start_case_log(build_log) + logger.info("[sandbox] build phase started") + await app.start() + for index, specification in enumerate(jobs): + if not isinstance(specification, dict) or not isinstance(specification.get("job"), str): + raise ValueError(f"invalid build job at index {index}") + arguments = specification.get("arguments") or {} + if not isinstance(arguments, dict): + raise ValueError(f"build job arguments must be an object at index {index}") + result = await _run_job_on_app(app, specification["job"], arguments) + if not result["success"]: + failed_results.append(result) + _append_build_error( + build_log, + f"[sandbox] build job {specification['job']!r} failed: " + f"{result.get('error') or 'job failed'}\n" + f"{result.get('traceback') or ''}", + ) + item = {"job": specification["job"], "result": result} + checkpoint = specification.get("memory_checkpoint") + if checkpoint is not None and (not isinstance(checkpoint, str) or not checkpoint.strip()): + raise ValueError(f"invalid memory_checkpoint at index {index}") + if result["success"] and checkpoint is not None: + item["memory_checkpoint"] = _commit_memory_checkpoint(case_root, app_config, checkpoint) + results.append(item) + for result in failed_results: + for diagnostic in result.get("cleanup_errors", []): + _append_build_error(build_log, diagnostic["traceback"]) + finally: + if logger is not None: + logger.info("[sandbox] build phase finished") + get_logger(log_to_console=False, log_to_file=False, force_init=True) + except BaseException: + _append_build_error(build_log, traceback.format_exc()) raise - - # Application.__init__ force-initializes the global logger, removing any - # sink bound before construction. Attach the build artifact afterwards. - logger = _start_case_log(build_log) - logger.info("[sandbox] build phase started") - results: list[dict[str, Any]] = [] - try: - await app.start() - for index, specification in enumerate(jobs): - if not isinstance(specification, dict) or not isinstance(specification.get("job"), str): - raise ValueError(f"invalid build job at index {index}") - arguments = specification.get("arguments") or {} - if not isinstance(arguments, dict): - raise ValueError(f"build job arguments must be an object at index {index}") - result = await _run_job_on_app(app, specification["job"], arguments) - item = {"job": specification["job"], "result": result} - checkpoint = specification.get("memory_checkpoint") - if checkpoint is not None and (not isinstance(checkpoint, str) or not checkpoint.strip()): - raise ValueError(f"invalid memory_checkpoint at index {index}") - if result["success"] and checkpoint is not None: - item["memory_checkpoint"] = _commit_memory_checkpoint(case_root, app_config, checkpoint) - results.append(item) - finally: - logger.info("[sandbox] build phase finished") - await app.close() - get_logger(log_to_console=False, log_to_file=False, force_init=True) return { "success": len(results) == len(jobs) and all(item["result"]["success"] for item in results), "jobs": results, @@ -361,60 +411,67 @@ async def _run_queries(request: dict[str, Any], *, write_summary: bool = True) - app = ReMe(**app_config) results: list[dict[str, Any]] = [] - try: - await app.start() - for query, query_id in zip(queries, query_ids): - query_dir = queries_root / query_id - query_dir.mkdir(parents=False, exist_ok=False) - log_path = query_dir / "answer.log" - logger = _start_case_log(log_path) - logger.info(f"[sandbox] query started query_id={query_id!r}") - answer_result: dict[str, Any] | None = None - error: str | None = None - exception: dict[str, Any] | None = None - try: - question = query.get("question") - if not isinstance(question, str): - raise ValueError(f"query {query_id!r} requires a string question") - answer_arguments = dict(query.get("answer_arguments") or {}) - answer_arguments["query"] = question - answer_result = await _run_job_on_app( - app, - str(query.get("answer_job") or "agentic_answer"), - answer_arguments, - ) - if not answer_result["success"]: - error = answer_result.get("error") or "answer job failed" - exception = answer_result.get("exception") - except Exception as exc: # Keep later queries runnable and preserve this query's log. - error = f"{type(exc).__name__}: {exc}" - exception = _exception_details(exc) - traceback_text = traceback.format_exc() - else: - traceback_text = None - finally: - logger.info(f"[sandbox] query finished query_id={query_id!r}") - _clear_query_state(app) - get_logger(log_to_console=False, log_to_file=False, force_init=True) + failed_results: list[dict[str, Any]] = [] + async with _application_scope(app, failed_results): + try: + await app.start() + for query, query_id in zip(queries, query_ids): + query_dir = queries_root / query_id + query_dir.mkdir(parents=False, exist_ok=False) + log_path = query_dir / "answer.log" + logger = _start_case_log(log_path) + logger.info(f"[sandbox] query started query_id={query_id!r}") + answer_result: dict[str, Any] | None = None + error: str | None = None + exception: dict[str, Any] | None = None + try: + question = query.get("question") + if not isinstance(question, str): + raise ValueError(f"query {query_id!r} requires a string question") + answer_arguments = dict(query.get("answer_arguments") or {}) + answer_arguments["query"] = question + answer_result = await _run_job_on_app( + app, + str(query.get("answer_job") or "agentic_answer"), + answer_arguments, + ) + if not answer_result["success"]: + error = answer_result.get("error") or "answer job failed" + exception = answer_result.get("exception") + except Exception as exc: # Keep later queries runnable and preserve this query's log. + error = f"{type(exc).__name__}: {exc}" + exception = _exception_details(exc) + traceback_text = traceback.format_exc() + else: + traceback_text = None + finally: + logger.info(f"[sandbox] query finished query_id={query_id!r}") + _clear_query_state(app) + get_logger(log_to_console=False, log_to_file=False, force_init=True) - result = { - "query_id": query_id, - "question": query.get("question"), - "answer": answer_result.get("answer") if answer_result is not None else None, - "score": None, - "answer_result": answer_result, - "judge_result": None, - "error": error, - } - if exception is not None: - result["exception"] = exception - if traceback_text is not None: - result["traceback"] = traceback_text - _atomic_write_json(query_dir / "result.json", result) - results.append(result) - finally: - get_logger(log_to_console=False, log_to_file=False, force_init=True) - await app.close() + result = { + "query_id": query_id, + "question": query.get("question"), + "answer": answer_result.get("answer") if answer_result is not None else None, + "score": None, + "answer_result": answer_result, + "judge_result": None, + "error": error, + } + if exception is not None: + result["exception"] = exception + if traceback_text is not None: + result["traceback"] = traceback_text + _atomic_write_json(query_dir / "result.json", result) + results.append(result) + if error is not None: + failed_results.append(result) + finally: + get_logger(log_to_console=False, log_to_file=False, force_init=True) + + for result in failed_results: + if result.get("cleanup_errors"): + _atomic_write_json(queries_root / result["query_id"] / "result.json", result) summary = { "schema_version": 1, diff --git a/tests/unit/test_benchmark_docker.py b/tests/unit/test_benchmark_docker.py index 321e1d31..07740e31 100644 --- a/tests/unit/test_benchmark_docker.py +++ b/tests/unit/test_benchmark_docker.py @@ -7,6 +7,7 @@ from __future__ import annotations import asyncio import io import json +import os from pathlib import Path import subprocess import tarfile @@ -22,13 +23,14 @@ from deploy.benchmark_docker import ( SourceCandidate, SourceSnapshot, ) -from deploy.benchmark_docker.docker_sandbox import SandboxCommandError, SandboxCommandTimeoutError +from deploy.benchmark_docker.docker_sandbox import SandboxCommandError, SandboxCommandTimeoutError, _archive_workspace from deploy.benchmark_docker import worker as sandbox_worker class FakeCase: def __init__(self, name: str, events: list[tuple]) -> None: self.name = name + self.container_id = name self.events = events self.case_id = name self.reusable = True @@ -85,12 +87,17 @@ class FakeFactory: @pytest.mark.asyncio -async def test_pool_builds_and_queries_each_case_then_reuses_workspace(): +async def test_pool_builds_and_queries_each_case_then_reuses_workspace(tmp_path): events: list[tuple] = [] pool = BenchmarkDockerPool(FakeFactory(events), concurrency=1) # Put the dependent first to exercise task scheduling independently of input order. plans = [ - CasePlan("child", queries=(EvaluationQuery("q2", "question", "secret"),), reuse_workspace_from="parent"), + CasePlan( + "child", + queries=(EvaluationQuery("q2", "question", "secret"),), + reuse_workspace_from="parent", + export_evaluation_to=tmp_path / "evaluation.tar.gz", + ), CasePlan("parent", build_jobs=(JobRequest("build"),), queries=(EvaluationQuery("q1", "question", "secret"),)), ] results = await pool.run(plans) @@ -98,6 +105,7 @@ async def test_pool_builds_and_queries_each_case_then_reuses_workspace(): assert [result.case_id for result in results] == ["child", "parent"] assert all(result.success for result in results) assert ("child", "upload", "parent") in events + assert ("child", "export_evaluation") in events assert events.index(("parent", "build", 1)) < events.index(("parent", "query", 1)) assert events.index(("parent", "export_workspace")) < events.index(("parent", "query", 1)) assert events.index(("parent", "query", 1)) < events.index(("child", "query", 1)) @@ -114,26 +122,6 @@ async def test_pool_rejects_workspace_cycle_before_creating_containers(): assert not events -@pytest.mark.asyncio -async def test_pool_exports_query_only_case_from_reused_workspace(tmp_path): - events: list[tuple] = [] - archive = tmp_path / "evaluation.tar.gz" - plans = [ - CasePlan("parent", build_jobs=(JobRequest("build"),)), - CasePlan( - "child", - queries=(EvaluationQuery("q1", "question", "secret"),), - reuse_workspace_from="parent", - export_evaluation_to=archive, - ), - ] - - results = await BenchmarkDockerPool(FakeFactory(events), concurrency=1).run(plans) - - assert all(result.success for result in results) - assert ("child", "export_evaluation") in events - - @pytest.mark.asyncio @pytest.mark.parametrize("raises", [False, True], ids=["failed-result", "exception"]) async def test_pool_keeps_build_log_on_failure(tmp_path, raises): @@ -196,31 +184,7 @@ async def test_pool_replacement_failure_preserves_healthy_cases(): @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(): +async def test_pool_exhaustion_unblocks_waiting_and_dependent_cases(): events: list[tuple] = [] class BrokenCase(FakeCase): @@ -239,12 +203,14 @@ async def test_pool_exhaustion_unblocks_dependent_cases(): CasePlan("dependent", reuse_workspace_from="waiting"), CasePlan("broken", build_jobs=(JobRequest("build"),)), CasePlan("waiting", build_jobs=(JobRequest("build"),)), + CasePlan("waiting-2", 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 + assert all("no available workers" in result.error for result in results[2:]) + assert not any(case_id.startswith("waiting") for case_id, action in events if action == "reset") @pytest.mark.asyncio @@ -271,7 +237,8 @@ async def test_pool_close_failure_retires_slot_without_replacement(): 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 "no available workers" in results[1].error + assert "close failed" in results[1].error assert ("next", "build", 1) not in events @@ -302,16 +269,259 @@ async def test_pool_cancels_case_before_closing_container(): assert events.index(("case", "query_stopped")) < events.index(("pool-0", "close")) -def test_source_snapshot_excludes_unrelated_and_secret_files(tmp_path): +@pytest.mark.asyncio +async def test_pool_final_close_failure_marks_only_cases_using_that_worker(): + events: list[tuple] = [] + + class Worker(FakeCase): + async def run_build(self, jobs): + if self.case_id == "failed-build": + return {"success": False} + return await super().run_build(jobs) + + async def close(self): + if self.name == "pool-0": + raise RuntimeError("delete failed") + await super().close() + + class Factory(FakeFactory): + async def create_cases(self, case_ids, *, concurrency): + return [Worker(name, self.events) for name in case_ids] + + plans = [ + CasePlan(case_id, build_jobs=(JobRequest("build"),)) for case_id in ("first", "healthy", "failed-build", "last") + ] + results = await BenchmarkDockerPool(Factory(events), concurrency=2).run(plans) + assert not results[0].success + assert "pool-0" in results[0].error and "delete failed" in results[0].error + assert results[1].success + assert results[2].error.startswith("build failed;") + assert "delete failed" in results[2].error + assert results[3].success + + +@pytest.mark.asyncio +async def test_pool_reports_unused_worker_cleanup_failure(): + events: list[tuple] = [] + + class Worker(FakeCase): + async def run_build(self, jobs): + return {"success": False} + + async def close(self): + if self.name == "pool-1": + raise RuntimeError("unused worker leaked") + + class Factory(FakeFactory): + async def create_cases(self, case_ids, *, concurrency): + return [Worker(name, self.events) for name in case_ids] + + plans = [CasePlan("parent", build_jobs=(JobRequest("build"),)), CasePlan("child", reuse_workspace_from="parent")] + results = await BenchmarkDockerPool(Factory(events), concurrency=2).run(plans) + assert results[0].error.startswith("build failed;") + assert results[1].error.startswith("source case 'parent' failed;") + assert all("unused worker leaked" in result.error for result in results) + + +@pytest.mark.asyncio +async def test_pool_finishes_final_cleanup_despite_repeated_cancellation(): + started, release = asyncio.Event(), asyncio.Event() + events: list[tuple] = [] + + class Worker(FakeCase): + async def close(self): + started.set() + await release.wait() + self.events.append((self.name, "closed")) + raise RuntimeError("delete failed") + + class Factory(FakeFactory): + async def create_cases(self, case_ids, *, concurrency): + return [Worker(name, self.events) for name in case_ids] + + task = asyncio.create_task( + BenchmarkDockerPool(Factory(events), concurrency=1).run([CasePlan("case", build_jobs=(JobRequest("build"),))]), + ) + await asyncio.wait_for(started.wait(), timeout=2) + for _ in range(2): + task.cancel() + await asyncio.sleep(0) + assert not task.done() + release.set() + with pytest.raises(asyncio.CancelledError) as failure: + await asyncio.wait_for(task, timeout=2) + assert events[-1] == ("pool-0", "closed") + assert any("delete failed" in note for note in failure.value.__notes__) + + +@pytest.mark.asyncio +async def test_pool_does_not_replace_worker_for_cases_already_skipped(): + events: list[tuple] = [] + + class Worker(FakeCase): + async def run_build(self, jobs): + if self.case_id == "parent": + return {"success": False} + # Let the dependent observe the failed parent before retirement. + await asyncio.sleep(0) + self.reusable = False + raise RuntimeError("command timed out") + + 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): + self.events.append((case_id, "replacement")) + return FakeCase(case_id, self.events) + + plans = [ + CasePlan("parent", build_jobs=(JobRequest("build"),)), + CasePlan("skipped", reuse_workspace_from="parent"), + CasePlan("last", build_jobs=(JobRequest("build"),)), + ] + results = await BenchmarkDockerPool(Factory(events), concurrency=1).run(plans) + assert results[1].error == "source case 'parent' failed" + assert "command timed out" in results[2].error + assert not any(action == "replacement" for _, action in events) + + +def test_source_snapshot_prunes_excluded_directories_and_secret_files(tmp_path, monkeypatch): (tmp_path / "reme").mkdir() (tmp_path / "reme" / "__init__.py").write_text("ok", encoding="utf-8") - (tmp_path / "baselines").mkdir() - (tmp_path / "baselines" / "large.json").write_text("private", encoding="utf-8") + excluded = [tmp_path / name for name in ("baselines", ".venv", "node_modules")] + for directory in excluded: + (directory / "nested").mkdir(parents=True) + (directory / "nested" / "large.json").write_text("private", encoding="utf-8") (tmp_path / ".env").write_text("secret", encoding="utf-8") + scandir = os.scandir + + def checked_scandir(path): + path = Path(path) + assert not any(path == directory or directory in path.parents for directory in excluded) + return scandir(path) + + monkeypatch.setattr(os, "scandir", checked_scandir) snapshot = SourceSnapshot.from_directory(tmp_path) assert snapshot.file_count == 1 +def test_source_snapshot_preserves_custom_path_rules_and_symlink_checks(tmp_path): + for name in ("pkg/cache/keep.py", "pkg/core.py", "pkg/cache/private.secret"): + path = tmp_path / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(name, encoding="utf-8") + (tmp_path / "pkg" / "ignored").symlink_to(tmp_path / "pkg" / "cache", target_is_directory=True) + patterns = ("pkg/cache", "pkg/ignored", "pkg/*.secret") + snapshot = SourceSnapshot.from_directory(tmp_path, excludes=patterns) + with tarfile.open(fileobj=io.BytesIO(snapshot.archive), mode="r:gz") as archive: + assert archive.getnames() == ["pkg/cache/keep.py", "pkg/core.py"] + (tmp_path / "link").symlink_to(tmp_path / "pkg" / "cache", target_is_directory=True) + with pytest.raises(ValueError, match="symbolic link: link"): + SourceSnapshot.from_directory(tmp_path, excludes=patterns) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["job", "queries"]) +async def test_worker_keeps_startup_error_or_cancellation_when_close_fails(tmp_path, monkeypatch, mode): + import reme.reme as reme_module + + primary = RuntimeError("startup exploded") if mode == "job" else asyncio.CancelledError("startup cancelled") + closed = False + + class FakeReMe: + def __init__(self, **_kwargs): + pass + + async def start(self): + raise primary + + async def close(self): + nonlocal closed + closed = True + raise RuntimeError("close exploded") + + monkeypatch.setattr(reme_module, "ReMe", FakeReMe) + monkeypatch.setattr(sandbox_worker, "_prepare_runtime", lambda _request: (tmp_path, {})) + runner = sandbox_worker._run if mode == "job" else sandbox_worker._run_queries + with pytest.raises(type(primary)) as failure: + await runner({"job": "build", "queries": [{"query_id": "q1", "question": "Question?"}]}) + assert failure.value is primary + assert closed + assert any("close exploded" in note for note in primary.__notes__) + + +@pytest.mark.parametrize("mode", ["job", "build", "queries"]) +@pytest.mark.parametrize("job_fails", [True, "result", False], ids=["exception", "failed-result", "close-only"]) +def test_worker_response_preserves_failed_job_when_close_fails(tmp_path, monkeypatch, mode, job_fails): + import reme.reme as reme_module + from types import SimpleNamespace + from reme.utils.counter import global_counter_inc + + class FakeReMe: + def __init__(self, **_kwargs): + self.context = SimpleNamespace(components={}, jobs={"build": object()}, metadata={}) + + async def start(self): + pass + + async def run_job(self, job, **_kwargs): + global_counter_inc(self.context.metadata, ["__job_counter", job]) + if job_fails == "result": + return SimpleNamespace(success=False, answer="job failed", metadata={"diagnostic": "original"}) + if job_fails: + raise RuntimeError("job exploded") + return SimpleNamespace(success=True, answer="done", metadata={}) + + async def close(self): + raise ValueError("close exploded") + + monkeypatch.setattr(reme_module, "ReMe", FakeReMe) + monkeypatch.setattr(sandbox_worker, "_prepare_runtime", lambda _request: (tmp_path, {})) + request = { + "mode": mode, + "case_root": str(tmp_path), + "job": "build", + "jobs": [{"job": "build"}], + "queries": [{"query_id": "q1", "question": "Question?", "answer_job": "build"}], + } + request_path = tmp_path / "request.json" + response_path = tmp_path / "response.json" + request_path.write_text(json.dumps(request), encoding="utf-8") + monkeypatch.setattr("sys.argv", ["worker", "--request", str(request_path), "--response", str(response_path)]) + assert sandbox_worker.main() == 1 + response = json.loads(response_path.read_text(encoding="utf-8")) + assert response["success"] is False + if not job_fails: + assert response["error"] == "ValueError: close exploded" + return + + result = response + if mode == "build": + result = response["jobs"][0]["result"] + log = (tmp_path / "build_log" / "build.log").read_text(encoding="utf-8") + assert ("job failed" if job_fails == "result" else "job exploded") in log + assert "close exploded" in log + elif mode == "queries": + query_result = response["queries"][0] + saved = json.loads((tmp_path / "queries" / "q1" / "result.json").read_text(encoding="utf-8")) + assert saved == query_result + result = query_result["answer_result"] + if job_fails == "result": + assert result["error"] == "job failed" + assert result["metadata"] == {"diagnostic": "original"} + else: + assert result["error"] == "RuntimeError: job exploded" + assert result["exception"]["type"] == "RuntimeError" + assert "job exploded" in result["traceback"] + assert result["job_call_counts"] == {"build": 1} + assert result["token_usage"] == {} + diagnostics = query_result["cleanup_errors"] if mode == "queries" else result["cleanup_errors"] + assert diagnostics[0]["error"] == "ValueError: close exploded" + assert diagnostics[0]["exception"]["type"] == "ValueError" + assert "close exploded" in diagnostics[0]["traceback"] + + @pytest.mark.asyncio @pytest.mark.parametrize("failure_phase", ["config", "application"]) async def test_build_log_captures_initialization_failure(tmp_path, monkeypatch, failure_phase): @@ -343,40 +553,79 @@ async def test_build_log_captures_initialization_failure(tmp_path, monkeypatch, @pytest.mark.asyncio -async def test_build_log_remains_bound_after_application_initializes_logger(tmp_path, monkeypatch): +@pytest.mark.parametrize("failure_phase", [None, "start", "checkpoint", "close", "start_and_close", "job"]) +async def test_build_log_captures_runtime_failures_and_preserves_primary_error(tmp_path, monkeypatch, failure_phase): import reme.reme as reme_module + from reme.utils import get_logger monkeypatch.setattr(sandbox_worker, "_prepare_runtime", lambda _request: (tmp_path, {})) + startup_error = RuntimeError("startup exploded") + checkpoint_error = RuntimeError("checkpoint exploded") + close_error = RuntimeError("close exploded") + closed = False class FakeReMe: def __init__(self, **_kwargs): - from reme.utils import get_logger - get_logger(log_to_console=False, log_to_file=False, force_init=True) async def start(self): - pass + if failure_phase in {"start", "start_and_close"}: + raise startup_error async def close(self): - pass + nonlocal closed + closed = True + if failure_phase in {"close", "start_and_close"}: + raise close_error async def run_job(_app, _job, _arguments): + if failure_phase == "job": + return {"success": False, "error": "RuntimeError: job exploded", "traceback": "job traceback"} return {"success": True} + def checkpoint(*_args): + raise checkpoint_error + monkeypatch.setattr(reme_module, "ReMe", FakeReMe) monkeypatch.setattr(sandbox_worker, "_run_job_on_app", run_job) + monkeypatch.setattr(sandbox_worker, "_commit_memory_checkpoint", checkpoint) + jobs = [{"job": "build"}] + if failure_phase == "checkpoint": + jobs[0]["memory_checkpoint"] = "boundary" + request = {"case_root": str(tmp_path), "jobs": jobs} - result = await sandbox_worker._run_build({"case_root": str(tmp_path), "jobs": [{"job": "build"}]}) + if failure_phase in {"start", "checkpoint", "close", "start_and_close"}: + expected_error = checkpoint_error if failure_phase == "checkpoint" else startup_error + if failure_phase == "close": + expected_error = close_error + with pytest.raises(RuntimeError) as failure: + await sandbox_worker._run_build(request) + assert failure.value is expected_error + else: + result = await sandbox_worker._run_build(request) + assert result["success"] is (failure_phase is None) - assert result["success"] + assert closed + get_logger().info("after build") log_text = (tmp_path / "build_log" / "build.log").read_text(encoding="utf-8") assert "[sandbox] build phase started" in log_text assert "[sandbox] build phase finished" in log_text + assert "after build" not in log_text + if failure_phase in {"start", "checkpoint", "close", "start_and_close"}: + assert "Traceback (most recent call last):" in log_text + assert str(expected_error) in log_text + if failure_phase == "start_and_close": + assert "close exploded" in log_text + assert any("close exploded" in note for note in failure.value.__notes__) + if failure_phase == "job": + assert "RuntimeError: job exploded" in log_text + assert "job traceback" in log_text @pytest.mark.asyncio async def test_worker_clears_search_budget_between_queries(tmp_path, monkeypatch): import reme.reme as reme_module + from reme.utils import get_logger monkeypatch.setattr(sandbox_worker, "_prepare_runtime", lambda _request: (tmp_path, {})) seen_metadata: list[dict] = [] @@ -396,6 +645,7 @@ async def test_worker_clears_search_budget_between_queries(tmp_path, monkeypatch pass async def run_job(app, _job, _arguments): + get_logger().info(f"answer body: {_arguments['query']}") 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}} @@ -418,6 +668,10 @@ async def test_worker_clears_search_budget_between_queries(tmp_path, monkeypatch assert len(apps) == 1 assert seen_metadata == [{"shared": "retained"}, {"shared": "retained"}] assert apps[0].context.metadata == {"shared": "retained"} + for query_id, own, other in (("q1", "first", "second"), ("q2", "second", "first")): + log = (tmp_path / "queries" / query_id / "answer.log").read_text(encoding="utf-8") + assert f"answer body: {own}" in log + assert f"answer body: {other}" not in log @pytest.mark.asyncio @@ -612,40 +866,24 @@ async def test_cancelled_request_upload_prevents_case_reset(tmp_path): @pytest.mark.asyncio -async def test_backend_write_error_prevents_case_reset(tmp_path): +@pytest.mark.parametrize("operation", ["write", "exec"]) +async def test_backend_error_prevents_case_reset(tmp_path, operation): class FailingBackend(FakeBackend): async def write_file(self, path, data): - if "/inbox/" in path: - raise RuntimeError("Docker upload stream failed") + if operation == "write" and "/inbox/" in path: + raise RuntimeError("Docker 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") + if operation == "exec" and "/workspace/harness/worker.py" in command: + raise RuntimeError("Docker 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") + case = await factory.create_case(f"{operation}-error") try: - with pytest.raises(RuntimeError, match="Docker exec stream failed"): + with pytest.raises(RuntimeError, match="Docker stream failed"): await case.run_build([JobRequest("build")]) assert not case.reusable with pytest.raises(SandboxCommandError, match="no longer reusable"): @@ -765,25 +1003,6 @@ async def test_pool_replaces_worker_after_backend_timeout(tmp_path): assert sum(command[:2] == ["rm", "-rf"] for command in timed_out_workspace.backend.commands) == 1 -@pytest.mark.asyncio -async def test_query_only_evaluation_export_omits_missing_build_log(tmp_path): - backend = FakeBackend() - workspace = FakeWorkspace(backend) - factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: workspace) - case = await factory.create_case("query-only") - try: - backend.files["/workspace/case/queries/summary.json"] = b"{}" - await case.export_evaluation(tmp_path / "evaluation.tar.gz") - export = next( - command for command in backend.commands if command[:3] == ["tar", "-czf", "/tmp/reme-case-export.tar.gz"] - ) - assert "reme_workspace" in export - assert "queries" in export - assert "build_log" not in export - finally: - await case.close() - - @pytest.mark.asyncio @pytest.mark.parametrize( ("method", "required", "description", "members"), @@ -792,11 +1011,16 @@ async def test_query_only_evaluation_export_omits_missing_build_log(tmp_path): "export_memory_construction", "build_log/build.log", "memory construction export", - ["reme_workspace", "build_log"], + ["manifest.json", "reme_workspace", "build_log"], + ), + ("export_queries", "queries/summary.json", "query export", ["manifest.json", "queries"]), + ("export_query", "queries/q1/result.json", "query export", ["manifest.json", "queries/q1"]), + ( + "export_evaluation", + "queries/summary.json", + "evaluation export", + ["manifest.json", "reme_workspace", "queries"], ), - ("export_queries", "queries/summary.json", "query export", ["queries"]), - ("export_query", "queries/q1/result.json", "query export", ["queries/q1"]), - ("export_evaluation", "queries/summary.json", "evaluation export", ["reme_workspace", "queries"]), ], ) async def test_archive_exports_check_required_artifacts(tmp_path, method, required, description, members): @@ -811,12 +1035,99 @@ async def test_archive_exports_check_required_artifacts(tmp_path, method, requir assert str(failure.value) == f"{description} is missing required artifact: /workspace/case/{required}" 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 +@pytest.mark.parametrize( + "method", + ["export_workspace", "export_memory_construction", "export_evaluation"], +) +async def test_exported_memory_can_seed_queries_without_build(tmp_path, method): + class ArchiveBackend(FakeBackend): + async def exec_shell(self, command, *, cwd=None, timeout=None): + result = await super().exec_shell(command, cwd=cwd, timeout=timeout) + if command[:2] == ["tar", "-czf"]: + root = command[command.index("-C") + 1] + "/" + members = command[command.index("-C") + 2 :] + stream = io.BytesIO() + with tarfile.open(fileobj=stream, mode="w:gz") as archive: + for path, content in self.files.items(): + if not path.startswith(root): + continue + name = path[len(root) :] + if not any( + member == "." or name == member or name.startswith(member + "/") for member in members + ): + continue + info = tarfile.TarInfo(name) + info.size = len(content) + archive.addfile(info, io.BytesIO(content)) + self.files[command[2]] = stream.getvalue() + return result + + source_backend, query_backend = ArchiveBackend(), FakeBackend() + workspaces = iter((FakeWorkspace(source_backend), FakeWorkspace(query_backend))) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: next(workspaces)) + source = await factory.create_case("build") + query_case = await factory.create_case("query") + try: + assert (await source.run_build([JobRequest("build")]))["success"] + source_backend.files.update( + { + "/workspace/case/reme_workspace/daily/note.md": b"saved memory", + "/workspace/case/reme_workspace/.git/HEAD": b"ref: refs/heads/main\n", + "/workspace/case/build_log/build.log": b"build log", + "/workspace/case/queries/summary.json": b"{}", + }, + ) + saved = await getattr(source, method)(tmp_path / "memory.tar.gz") + await query_case.upload_workspace(saved) + uploaded = query_backend.files["/tmp/reme-workspace-upload.tar.gz"] + with tarfile.open(fileobj=io.BytesIO(uploaded), mode="r:gz") as archive: + assert set(archive.getnames()) == {"daily/note.md", ".git/HEAD"} + assert archive.extractfile("daily/note.md").read() == b"saved memory" + assert (await query_case.run_queries([EvaluationQuery("q1", "Question?", "Answer")]))["success"] + requests = [json.loads(content) for path, content in query_backend.files.items() if "/inbox/" in path] + assert [request["mode"] for request in requests] == ["queries"] + finally: + await source.close() + await query_case.close() + + +@pytest.mark.parametrize( + ("manifest", "workspace_path", "error"), + [ + (None, "reme_workspace/daily/note.md", "unmarked case archive"), + ({"archive_format": "reme-case-v1", "workspace_root": None}, None, "no reusable workspace"), + ({"archive_format": "reme-case-v1", "workspace_root": "../escape"}, None, "workspace root"), + ({"archive_format": "reme-case-v2", "workspace_root": "reme_workspace"}, None, "unsupported.*format"), + ({"archive_format": "reme-case-v1", "workspace_root": "reme_workspace"}, None, "workspace is missing"), + ({}, "reme_workspace/daily/note.md", None), + ], +) +def test_workspace_import_checks_case_archive_format(tmp_path, manifest, workspace_path, error): + source = tmp_path / "case.tar.gz" + files = {"queries/q1/result.json": b"{}"} + if manifest is not None: + files["manifest.json"] = json.dumps(manifest).encode() + if workspace_path is not None: + files[workspace_path] = b"memory" + with tarfile.open(source, "w:gz") as archive: + for name, content in files.items(): + info = tarfile.TarInfo(name) + info.size = len(content) + archive.addfile(info, io.BytesIO(content)) + if error is not None: + with pytest.raises(ValueError, match=error): + _archive_workspace(source) + else: + with tarfile.open(fileobj=io.BytesIO(_archive_workspace(source)), mode="r:gz") as archive: + assert archive.getnames() == ["daily/note.md"] + + @pytest.mark.asyncio async def test_default_analysis_export_saves_logs_outside_archive(tmp_path): class AnalysisBackend(FakeBackend): @@ -881,7 +1192,7 @@ async def test_cancelled_case_initialization_closes_workspace(tmp_path): @pytest.mark.asyncio -async def test_cancelled_batch_creation_closes_completed_workspace(tmp_path): +async def test_cancelled_batch_creation_closes_completed_workspace_and_reports_cleanup_failure(tmp_path): second_started = asyncio.Event() class BlockingWorkspace(FakeWorkspace): @@ -889,6 +1200,10 @@ async def test_cancelled_batch_creation_closes_completed_workspace(tmp_path): second_started.set() await asyncio.Event().wait() + async def close(self): + await super().close() + raise RuntimeError("delete failed") + first = FakeWorkspace(FakeBackend()) second = BlockingWorkspace(FakeBackend()) workspaces = iter((first, second)) @@ -897,9 +1212,10 @@ async def test_cancelled_batch_creation_closes_completed_workspace(tmp_path): await asyncio.wait_for(second_started.wait(), timeout=2) await asyncio.sleep(0) task.cancel() - with pytest.raises(asyncio.CancelledError): + with pytest.raises(asyncio.CancelledError) as failure: await task assert first.closed and second.closed + assert any("delete failed" in note for note in failure.value.__notes__) @pytest.mark.asyncio @@ -929,5 +1245,82 @@ async def test_failed_batch_creation_preserves_creation_error_when_close_fails(t 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"): + with pytest.raises(RuntimeError, match="initialization failed") as failure: await factory.create_cases(["first", "second"], concurrency=2) + assert any("close failed" in note for note in failure.value.__notes__) + + +@pytest.mark.asyncio +async def test_factory_finishes_failed_creation_cleanup_before_cancellation(tmp_path): + started, release = asyncio.Event(), asyncio.Event() + + class BrokenWorkspace(FakeWorkspace): + async def initialize(self): + raise RuntimeError("initialization failed") + + async def close(self): + started.set() + await release.wait() + self.closed = True + + workspace = BrokenWorkspace(FakeBackend()) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: workspace) + task = asyncio.create_task(factory.create_case("failed")) + await asyncio.wait_for(started.wait(), timeout=2) + for _ in range(2): + task.cancel() + await asyncio.sleep(0) + assert not task.done() + release.set() + with pytest.raises(asyncio.CancelledError) as failure: + await asyncio.wait_for(task, timeout=2) + assert workspace.closed + assert "initialization failed" in str(failure.value.__cause__) + + +@pytest.mark.asyncio +async def test_batch_reports_all_initialization_and_cleanup_failures(tmp_path): + class BrokenWorkspace(FakeWorkspace): + async def initialize(self): + raise RuntimeError("initialization failed") + + async def close(self): + raise RuntimeError("delete 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_cases(["first", "second"], concurrency=2) + assert any("additional creation failure" in note for note in failure.value.__notes__) + assert sum("delete failed" in note for note in failure.value.__notes__) == 2 + + +@pytest.mark.asyncio +async def test_pool_retries_strict_delete_without_replacing_uncertain_container(tmp_path): + from types import SimpleNamespace + from unittest.mock import AsyncMock + + from deploy.benchmark_docker.direct_docker_workspace import DirectDockerWorkspace + + workspaces = [] + container = SimpleNamespace(delete=AsyncMock(side_effect=[RuntimeError("delete failed"), None])) + client = SimpleNamespace(close=AsyncMock()) + + class StrictWorkspace(DirectDockerWorkspace): + async def _provision_backend(self): + self._container, self._client, self._backend = container, client, TimedOutBackend() + + def build_workspace(workspace_id, image, env): + workspace = StrictWorkspace(workspace_id=workspace_id, base_image=image, env=env) + workspaces.append(workspace) + return workspace + + factory = _fake_docker_factory(tmp_path, build_workspace) + plans = [CasePlan(case_id, build_jobs=(JobRequest("build"),)) for case_id in ("timed-out", "waiting")] + results = await asyncio.wait_for(BenchmarkDockerPool(factory, concurrency=1).run(plans), timeout=2) + assert not any(result.success for result in results) + assert "worker close failed" in results[0].error + assert "delete failed" in results[1].error + assert len(workspaces) == 1 + assert container.delete.await_count == 2 + assert workspaces[0]._container is workspaces[0]._client is None + client.close.assert_awaited_once() diff --git a/tests/unit/test_benchmark_docker_workspace.py b/tests/unit/test_benchmark_docker_workspace.py new file mode 100644 index 00000000..5817474c --- /dev/null +++ b/tests/unit/test_benchmark_docker_workspace.py @@ -0,0 +1,124 @@ +"""Direct workspace lifecycle without a Docker daemon.""" + +# pylint: disable=missing-function-docstring,protected-access + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +from aiodocker.exceptions import DockerError +import pytest + +from deploy.benchmark_docker.direct_docker_workspace import DirectDockerWorkspace + + +def make_workspace(): + workspace = DirectDockerWorkspace(workspace_id="cleanup-test") + workspace._container = SimpleNamespace(delete=AsyncMock()) + workspace._client = SimpleNamespace(close=AsyncMock()) + workspace._backend = SimpleNamespace(exec_shell=AsyncMock()) + workspace.is_alive = True + return workspace + + +@pytest.mark.asyncio +async def test_delete_failure_retains_handles_and_can_be_retried(): + workspace = make_workspace() + container, client, backend = workspace._container, workspace._client, workspace._backend + container.delete.side_effect = [RuntimeError("delete failed"), None] + + with pytest.raises(RuntimeError, match="delete failed"): + await workspace.close() + assert workspace._container is container + assert workspace._client is client + assert workspace._backend is backend + assert not workspace.is_alive + client.close.assert_not_awaited() + with pytest.raises(RuntimeError, match="cleanup must complete"): + await workspace.initialize() + + await workspace.close() + await workspace.close() + assert container.delete.await_count == 2 + container.delete.assert_awaited_with(force=True) + client.close.assert_awaited_once() + assert workspace._container is workspace._client is workspace._backend is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [404, 500]) +async def test_only_missing_container_is_a_successful_delete(status): + workspace = make_workspace() + container, client = workspace._container, workspace._client + container.delete.side_effect = DockerError(status, {"message": "delete failed"}) + if status == 404: + await workspace.close() + assert workspace._container is None + client.close.assert_awaited_once() + else: + with pytest.raises(DockerError): + await workspace.close() + assert workspace._container is container + client.close.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_client_close_failure_retries_without_deleting_again(): + workspace = make_workspace() + container, client = workspace._container, workspace._client + client.close.side_effect = [RuntimeError("client close failed"), None] + with pytest.raises(RuntimeError, match="client close failed"): + await workspace.close() + assert workspace._container is workspace._backend is None + assert workspace._client is client + + await workspace.close() + container.delete.assert_awaited_once_with(force=True) + assert client.close.await_count == 2 + assert workspace._client is None + + +@pytest.mark.asyncio +async def test_cancelled_delete_retains_state_for_retry(): + workspace = make_workspace() + container, client = workspace._container, workspace._client + started = asyncio.Event() + + async def block_delete(**_kwargs): + started.set() + await asyncio.Event().wait() + + container.delete.side_effect = block_delete + task = asyncio.create_task(workspace.close()) + await asyncio.wait_for(started.wait(), timeout=2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert workspace._container is container + assert workspace._client is client + + container.delete.side_effect = None + await workspace.close() + assert workspace._container is workspace._client is None + + +@pytest.mark.asyncio +async def test_initialization_failure_keeps_primary_error_and_retries_cleanup(monkeypatch): + workspace = make_workspace() + container = workspace._container + container.delete.side_effect = [RuntimeError("delete failed"), None] + client, backend = workspace._client, workspace._backend + workspace._container = workspace._client = workspace._backend = None + workspace.is_alive = False + + async def provision(): + workspace._container, workspace._client, workspace._backend = container, client, backend + + monkeypatch.setattr(workspace, "_provision_backend", provision) + backend.exec_shell.return_value = SimpleNamespace(ok=lambda: False, stderr=b"mkdir failed") + with pytest.raises(RuntimeError, match="mkdir failed") as failure: + await workspace.initialize() + assert any("delete failed" in note for note in failure.value.__notes__) + assert workspace._container is container + await workspace.close() + assert workspace._container is workspace._client is workspace._backend is None