From 331f6394535f9ec9d9a919618b4670d8e5c2585f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Mon, 12 Oct 2026 16:12:52 +0800 Subject: [PATCH 01/14] Add Docker benchmark sandbox and pooled runner --- deploy/__init__.py | 1 + deploy/benchmark_docker/Dockerfile.base | 24 + deploy/benchmark_docker/Dockerfile.candidate | 12 + .../Dockerfile.candidate.dockerignore | 21 + deploy/benchmark_docker/README.md | 87 ++ deploy/benchmark_docker/__init__.py | 20 + deploy/benchmark_docker/candidate.py | 173 +++ .../direct_docker_workspace.py | 40 + deploy/benchmark_docker/docker_sandbox.py | 1105 +++++++++++++++++ deploy/benchmark_docker/models.py | 77 ++ deploy/benchmark_docker/pool.py | 197 +++ deploy/benchmark_docker/worker.py | 460 +++++++ pyproject.toml | 7 +- tests/integration/test_benchmark_docker.py | 188 +++ tests/unit/test_benchmark_docker.py | 383 ++++++ tests/unit/test_package_versions.py | 7 +- 16 files changed, 2800 insertions(+), 2 deletions(-) create mode 100644 deploy/__init__.py create mode 100644 deploy/benchmark_docker/Dockerfile.base create mode 100644 deploy/benchmark_docker/Dockerfile.candidate create mode 100644 deploy/benchmark_docker/Dockerfile.candidate.dockerignore create mode 100644 deploy/benchmark_docker/README.md create mode 100644 deploy/benchmark_docker/__init__.py create mode 100644 deploy/benchmark_docker/candidate.py create mode 100644 deploy/benchmark_docker/direct_docker_workspace.py create mode 100644 deploy/benchmark_docker/docker_sandbox.py create mode 100644 deploy/benchmark_docker/models.py create mode 100644 deploy/benchmark_docker/pool.py create mode 100644 deploy/benchmark_docker/worker.py create mode 100644 tests/integration/test_benchmark_docker.py create mode 100644 tests/unit/test_benchmark_docker.py diff --git a/deploy/__init__.py b/deploy/__init__.py new file mode 100644 index 00000000..b00027d6 --- /dev/null +++ b/deploy/__init__.py @@ -0,0 +1 @@ +"""Deployment and benchmark infrastructure.""" diff --git a/deploy/benchmark_docker/Dockerfile.base b/deploy/benchmark_docker/Dockerfile.base new file mode 100644 index 00000000..b7e3ac81 --- /dev/null +++ b/deploy/benchmark_docker/Dockerfile.base @@ -0,0 +1,24 @@ +# Build from repository root: +# docker build -f deploy/benchmark_docker/Dockerfile.base -t reme-benchmark-base:agentscope-2.0.9 . +FROM python:3.11-slim + +ENV PIP_DISABLE_PIP_VERSION_CHECK=1 \ + PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 + +RUN apt-get update \ + && apt-get install -y --no-install-recommends build-essential git ripgrep tar \ + && rm -rf /var/lib/apt/lists/* + +# Install dependencies with an expendable package stub. Each candidate is +# installed later with --no-deps so its code, not a fixed ReMe wheel, runs. +WORKDIR /tmp/reme-dependencies +COPY pyproject.toml README.md ./ +RUN mkdir reme \ + && printf '__version__ = "0.0.0"\n' > reme/__init__.py \ + && python -m pip install '.[core]' \ + && python -m pip uninstall -y reme-ai \ + && python -c 'import agentscope; assert agentscope.__version__ == "2.0.9"' \ + && rm -rf /tmp/reme-dependencies + +WORKDIR /workspace diff --git a/deploy/benchmark_docker/Dockerfile.candidate b/deploy/benchmark_docker/Dockerfile.candidate new file mode 100644 index 00000000..62c3b743 --- /dev/null +++ b/deploy/benchmark_docker/Dockerfile.candidate @@ -0,0 +1,12 @@ +# Build from repository root with this Dockerfile-specific ignore file. +ARG BASE_IMAGE=reme-benchmark-base:agentscope-2.0.9 +FROM ${BASE_IMAGE} + +COPY . /opt/reme +ARG REME_PLUGIN_PATH="" +RUN python -m pip install --no-deps --no-build-isolation /opt/reme \ + && if [ -n "${REME_PLUGIN_PATH}" ]; then \ + python -m pip install --no-deps --no-build-isolation "/opt/reme/${REME_PLUGIN_PATH}"; \ + fi + +WORKDIR /workspace diff --git a/deploy/benchmark_docker/Dockerfile.candidate.dockerignore b/deploy/benchmark_docker/Dockerfile.candidate.dockerignore new file mode 100644 index 00000000..c9628b66 --- /dev/null +++ b/deploy/benchmark_docker/Dockerfile.candidate.dockerignore @@ -0,0 +1,21 @@ +.git +.reme +.env +.env.* +.venv +**/__pycache__ +**/.pytest_cache +**/node_modules +**/dist +**/build +**/*.egg-info +benchmark +baselines +meta-reme +output +tests +docs +github-pages +integrations +skills +reme_studio diff --git a/deploy/benchmark_docker/README.md b/deploy/benchmark_docker/README.md new file mode 100644 index 00000000..d6f6259b --- /dev/null +++ b/deploy/benchmark_docker/README.md @@ -0,0 +1,87 @@ +# Benchmark Docker runner + +This package runs ReMe benchmark jobs directly inside isolated AgentScope +Docker workspaces. It does not start the ReMe HTTP service. Use AgentScope +**2.0.9** on the host and in the image. + +From the repository root: + +```bash +pip install -e '.[benchmark-docker]' +docker build -f deploy/benchmark_docker/Dockerfile.base \ + -t reme-benchmark-base:agentscope-2.0.9 . +``` + +`SourceCandidate` uploads an immutable source snapshot and installs the root +project and any explicitly selected plugin projects into a container-local +environment. The base image contains dependencies, but no ReMe candidate. +`ImageCandidate` instead uses a prebuilt image with ReMe and plugins installed. +The included `Dockerfile.candidate` is an example of that image mode. + +```python +import asyncio +from pathlib import Path + +from deploy.benchmark_docker import ( + BenchmarkDockerPool, CasePlan, DockerReMeSandboxFactory, + EvaluationQuery, JobRequest, SourceCandidate, SourceSnapshot, +) + + +async def main(): + snapshot = SourceSnapshot.from_directory(".") + factory = DockerReMeSandboxFactory( + SourceCandidate(snapshot, install_projects=(".", "plugins/lme")), + config="benchmark.yaml", + plugins=("lme",), + env={"LLM_API_KEY": "...", "EMBEDDING_API_KEY": "..."}, + ) + plans = [ + CasePlan( + case_id="case-1", + build_jobs=(JobRequest("auto_memory", {"session_id": "s1", "messages": []}),), + queries=(EvaluationQuery("q1", "Question?", "Private answer"),), + export_workspace_to=Path("artifacts/case-1-workspace.tar.gz"), + export_evaluation_to=Path("artifacts/case-1-evaluation.tar.gz"), + ), + ] + return await BenchmarkDockerPool(factory, concurrency=2).run(plans) + + +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. +`export_workspace_to` persists a portable snapshot beyond the pool run. +`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` before reusing +or closing the container. + +For direct control, call `create_case()` or `create_cases()`, then use +`run_job()`, `run_build()`, `run_queries()`, `export_workspace()`, +`upload_workspace()`, `reset_case()`, and `close()`. Always close a case in a +`finally` block or use its async context manager. Export before closing. + +The candidate receives only query ID, question, answer Job, and answer +arguments. Golden answers, rubrics, and scoring remain on the host. Workspace +archives reject links, special files, duplicate paths, and traversal paths. +Source snapshots exclude VCS state, credentials, benchmark data, and generated +artifacts; review custom source trees for other secrets before uploading. + +`tests/unit/test_benchmark_docker.py` runs without Docker. The credential-free +container smoke tests build the base image and temporary candidate images when +explicitly enabled: + +```bash +REME_DOCKER_INTEGRATION=1 pytest tests/integration/test_benchmark_docker.py -v +``` + +The host must have AgentScope 2.0.9 and a working Docker daemon. To use an +already built base image, set `REME_BENCHMARK_TEST_IMAGE` to its tag; the tests +do not remove a caller-supplied image. diff --git a/deploy/benchmark_docker/__init__.py b/deploy/benchmark_docker/__init__.py new file mode 100644 index 00000000..ba3e9bd3 --- /dev/null +++ b/deploy/benchmark_docker/__init__.py @@ -0,0 +1,20 @@ +"""Docker isolation for benchmark cases; independent of the ReMe service.""" + +from .candidate import ImageCandidate, SourceCandidate, SourceSnapshot +from .docker_sandbox import DockerReMeSandbox, DockerReMeSandboxFactory +from .models import EvaluationQuery, JobRequest, JobResult +from .pool import BenchmarkDockerPool, CasePlan, CaseResult + +__all__ = [ + "BenchmarkDockerPool", + "CasePlan", + "CaseResult", + "DockerReMeSandbox", + "DockerReMeSandboxFactory", + "EvaluationQuery", + "ImageCandidate", + "JobRequest", + "JobResult", + "SourceCandidate", + "SourceSnapshot", +] diff --git a/deploy/benchmark_docker/candidate.py b/deploy/benchmark_docker/candidate.py new file mode 100644 index 00000000..d7db4383 --- /dev/null +++ b/deploy/benchmark_docker/candidate.py @@ -0,0 +1,173 @@ +"""Candidate-code inputs shared by one or more isolated case sandboxes.""" + +from __future__ import annotations + +from dataclasses import dataclass +import fnmatch +import gzip +import hashlib +import io +from pathlib import Path +from pathlib import PurePosixPath +import tarfile + +DEFAULT_EXCLUDES = ( + ".git", + ".reme", + ".env", + ".env.*", + ".ssh", + ".codex", + ".claude", + ".venv", + "venv", + "__pycache__", + ".pytest_cache", + ".mypy_cache", + ".ruff_cache", + "datasets", + "benchmark", + "baselines", + "meta-reme", + "output", + "tests", + "cookbook", + "docs", + "github-pages", + "integrations", + "skills", + "reme_studio", + "node_modules", + "memory_workspaces", + "logs", + "artifacts", + "build", + "dist", + "*.egg-info", + "*.pyc", + "*.pyo", + "*.log", +) + + +def _is_excluded(relative_path: Path, patterns: tuple[str, ...]) -> bool: + """Return whether any path component or full POSIX path is excluded.""" + posix = relative_path.as_posix() + return any( + fnmatch.fnmatch(posix, pattern) or any(fnmatch.fnmatch(part, pattern) for part in relative_path.parts) + for pattern in patterns + ) + + +def _normalized_info(path: Path, arcname: str) -> tarfile.TarInfo: + """Create deterministic tar metadata for a regular source file.""" + stat = path.stat() + info = tarfile.TarInfo(arcname) + info.size = stat.st_size + info.mode = stat.st_mode & 0o777 + info.mtime = 0 + info.uid = 0 + info.gid = 0 + info.uname = "" + info.gname = "" + return info + + +@dataclass(frozen=True, slots=True) +class SourceSnapshot: + """A deterministic source archive reusable across many case sandboxes.""" + + archive: bytes + sha256: str + file_count: int + + @classmethod + def from_directory( + cls, + source_dir: str | Path, + *, + excludes: tuple[str, ...] = DEFAULT_EXCLUDES, + ) -> "SourceSnapshot": + """Snapshot a candidate directory without runtime or VCS state. + + Symbolic links are rejected instead of followed. This prevents a + candidate archive from silently capturing files outside its root. + """ + root = Path(source_dir).resolve(strict=True) + if not root.is_dir(): + 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)) + + raw = io.BytesIO() + with gzip.GzipFile(fileobj=raw, mode="wb", mtime=0) as compressed: + with tarfile.open(fileobj=compressed, mode="w") as archive: + for relative, path in sorted(files, key=lambda item: item[0].as_posix()): + info = _normalized_info(path, relative.as_posix()) + with path.open("rb") as source: + archive.addfile(info, source) + + payload = raw.getvalue() + return cls( + archive=payload, + sha256=hashlib.sha256(payload).hexdigest(), + file_count=len(files), + ) + + +@dataclass(frozen=True, slots=True) +class SourceCandidate: + """A changing source snapshot containing one or more Python projects. + + ``install_projects`` is a trusted, archive-relative installation plan. A + plugin directory merely being present in the snapshot is insufficient: + installing it is what publishes its ``reme.plugins`` entry point. + """ + + snapshot: SourceSnapshot + base_image: str = "reme-benchmark-base:agentscope-2.0.9" + install_projects: tuple[str, ...] = (".",) + + def __post_init__(self) -> None: + if not self.install_projects: + raise ValueError("source candidate requires at least one install project") + normalized: list[str] = [] + for value in self.install_projects: + if not isinstance(value, str) or not value: + raise TypeError("candidate install projects must be non-empty strings") + path = PurePosixPath(value) + if path.is_absolute() or ".." in path.parts or "\\" in value: + raise ValueError(f"candidate install project must stay inside the archive: {value!r}") + name = path.as_posix() + if name in normalized: + raise ValueError(f"candidate install project is duplicated: {value!r}") + normalized.append(name) + if normalized[0] != ".": + raise ValueError("the candidate root project must be installed before plugin projects") + object.__setattr__(self, "install_projects", tuple(normalized)) + + @property + def candidate_id(self) -> str: + """Stable identifier for the exact source snapshot.""" + return self.snapshot.sha256 + + +@dataclass(frozen=True, slots=True) +class ImageCandidate: + """A ReMe candidate already installed in an immutable Docker image.""" + + image: str + candidate_id: str | None = None + + @property + def resolved_candidate_id(self) -> str: + """Return the caller-supplied digest/ID or the image reference.""" + return self.candidate_id or self.image diff --git a/deploy/benchmark_docker/direct_docker_workspace.py b/deploy/benchmark_docker/direct_docker_workspace.py new file mode 100644 index 00000000..48782485 --- /dev/null +++ b/deploy/benchmark_docker/direct_docker_workspace.py @@ -0,0 +1,40 @@ +"""Minimal AgentScope 2.0.9 DockerWorkspace for direct ReMe jobs. + +The stock workspace builds and starts an MCP gateway. The benchmark harness +uses only DockerBackend's exec/read/write primitives, so building that gateway +would add an unrelated network-dependent image layer and process. +""" + +from __future__ import annotations + +from typing import Any + +from agentscope.workspace import DockerWorkspace + + +class DirectDockerWorkspace(DockerWorkspace): + """Use the supplied image directly and skip the unused MCP gateway.""" + + _image_tag: str + _backend: Any + is_alive: bool + + 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() + 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 diff --git a/deploy/benchmark_docker/docker_sandbox.py b/deploy/benchmark_docker/docker_sandbox.py new file mode 100644 index 00000000..04519558 --- /dev/null +++ b/deploy/benchmark_docker/docker_sandbox.py @@ -0,0 +1,1105 @@ +"""AgentScope DockerWorkspace orchestration for isolated ReMe cases.""" + +from __future__ import annotations + +import asyncio +from dataclasses import asdict +from datetime import datetime, timezone +import gzip +import io +import json +from pathlib import Path +from pathlib import PurePosixPath +import re +import tarfile +from typing import Any, Callable, Literal, Protocol +import uuid + +from .candidate import ImageCandidate, SourceCandidate +from .models import ActionRecord, EvaluationQuery, JobRequest, JobResult + +CASE_ROOT = "/workspace/case" +CANDIDATE_ROOT = "/workspace/candidate" +CANDIDATE_VENV = "/workspace/candidate-venv" +HARNESS_ROOT = "/workspace/harness" +SOURCE_ARCHIVE = "/workspace/candidate.tar.gz" +EXPORT_ARCHIVE = "/tmp/reme-case-export.tar.gz" +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}$") +_GIT_AUTHOR_NAME = "ReMe Sandbox" +_GIT_AUTHOR_EMAIL = "reme-sandbox@localhost" +_PRIVATE_EVALUATION_FIELDS = frozenset( + {"query", "question", "golden_answer", "question_type", "rubric", "judge_arguments"}, +) + + +class Backend(Protocol): + """Subset of AgentScope's backend API used by this module.""" + + async def exec_shell( + self, + command: list[str], + *, + cwd: str | None = None, + timeout: float | None = None, + ) -> Any: + """Execute an argv vector in the sandbox.""" + raise NotImplementedError + + async def read_file(self, path: str) -> bytes: + """Read a sandbox file.""" + raise NotImplementedError + + async def write_file(self, path: str, data: bytes) -> None: + """Write a sandbox file.""" + raise NotImplementedError + + +class Workspace(Protocol): + """Subset of AgentScope's DockerWorkspace used by this module.""" + + async def initialize(self) -> None: + """Start the workspace.""" + raise NotImplementedError + + async def close(self) -> None: + """Close the workspace.""" + raise NotImplementedError + + def get_backend(self) -> Backend: + """Return the execution backend.""" + raise NotImplementedError + + +WorkspaceBuilder = Callable[[str, str, dict[str, str]], Workspace] + + +class SandboxCommandError(RuntimeError): + """A required command failed inside a case sandbox.""" + + +def _validate_case_id(case_id: str) -> None: + """Reject identifiers that are unsafe for Docker names and manifests.""" + if not _CASE_ID_RE.fullmatch(case_id): + raise ValueError("case_id must contain only letters, numbers, '.', '_' or '-' and be <= 128 characters") + + +def _utc_now() -> str: + """Return an audit-friendly UTC timestamp.""" + return datetime.now(timezone.utc).isoformat() + + +def _decode(value: bytes | str | None) -> str: + """Decode backend process output without losing diagnostics.""" + if isinstance(value, bytes): + return value.decode("utf-8", errors="replace") + return value or "" + + +def _workspace_member_path(name: str) -> PurePosixPath | None: + """Validate and normalize one relative path in a workspace archive.""" + path = PurePosixPath(name) + if not name or path.is_absolute() or ".." in path.parts or "\\" in name: + raise ValueError(f"workspace archive contains unsafe path: {name!r}") + if path == PurePosixPath("."): + return None + return path + + +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. + """ + 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]] = [] + 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) + if item.is_symlink(): + raise ValueError(f"workspace source contains a symbolic link: {relative}") + path = _workspace_member_path(relative.as_posix()) + assert path is not None + 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)) + 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())) + else: + raise ValueError(f"workspace source contains unsupported file: {relative}") + 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: + path = _workspace_member_path(member.name) + if path is None: + continue + if case_export: + if not path.parts or path.parts[0] != "reme_workspace": + continue + if len(path.parts) == 1: + 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) + if member.isdir(): + info = tarfile.TarInfo(path.as_posix()) + info.type = tarfile.DIRTYPE + info.mode = member.mode & 0o777 + entries.append((path, info, 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}") + 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: + info.mtime = 0 + info.uid = info.gid = 0 + info.uname = info.gname = "" + archive.addfile(info, io.BytesIO(data) if data is not None else None) + return raw.getvalue() + + +def _default_workspace_builder(workspace_id: str, image: str, env: dict[str, str]) -> Workspace: + """Create the AgentScope 2.0.9 Docker backend lazily.""" + import agentscope + + from .direct_docker_workspace import DirectDockerWorkspace + + if agentscope.__version__ != "2.0.9": + raise RuntimeError(f"benchmark Docker host requires agentscope 2.0.9, got {agentscope.__version__}") + return DirectDockerWorkspace( + workspace_id=workspace_id, + base_image=image, + env=env, + ) + + +class DockerReMeSandboxFactory: + """Create independent case sandboxes that share one candidate definition. + + A :class:`SourceCandidate` stores one immutable source archive in host + memory. The same bytes are uploaded to every case, while each case gets a + new Docker container and a fresh ReMe runtime workspace. An + :class:`ImageCandidate` skips source upload and uses an image in which ReMe + is already installed. + """ + + def __init__( + self, + candidate: SourceCandidate | ImageCandidate, + *, + env: dict[str, str] | None = None, + config: str = "benchmark.yaml", + plugins: tuple[str, ...] | list[str] = (), + command_timeout: float = 1800.0, + workspace_builder: WorkspaceBuilder | None = None, + ) -> None: + self.candidate = candidate + self.env = dict(env or {}) + self.config = config + self.plugins = tuple(plugins) + if not all(isinstance(name, str) and name for name in self.plugins): + raise TypeError("sandbox plugins must be non-empty strings") + if len(self.plugins) != len(set(self.plugins)): + raise ValueError("sandbox plugins must be unique") + self.command_timeout = command_timeout + self._workspace_builder = workspace_builder or _default_workspace_builder + self._workspace_prefix = f"reme-benchmark-{uuid.uuid4().hex[:12]}" + self._workspace_sequence = 0 + + def _next_workspace_id(self) -> str: + """Allocate a factory-scoped container identity independent of case IDs.""" + self._workspace_sequence += 1 + return f"{self._workspace_prefix}-{self._workspace_sequence:04d}" + + async def create_case(self, case_id: str) -> "DockerReMeSandbox": + """Start and prepare a new container for one benchmark case.""" + _validate_case_id(case_id) + image = self.candidate.base_image if isinstance(self.candidate, SourceCandidate) else self.candidate.image + container_id = self._next_workspace_id() + workspace = self._workspace_builder(container_id, image, self.env) + case = DockerReMeSandbox( + case_id=case_id, + container_id=container_id, + candidate=self.candidate, + workspace=workspace, + config=self.config, + plugins=self.plugins, + command_timeout=self.command_timeout, + environment_names=sorted(self.env), + ) + try: + await case.initialize() + except BaseException: + await workspace.close() + raise + return case + + async def create_cases(self, case_ids: list[str], *, concurrency: int = 4) -> list["DockerReMeSandbox"]: + """Create several independent cases while bounding Docker pressure.""" + if concurrency < 1: + raise ValueError("concurrency must be at least 1") + if len(case_ids) != len(set(case_ids)): + 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) + + tasks = [asyncio.create_task(create(case_id)) for case_id in case_ids] + try: + results = await asyncio.gather(*tasks, return_exceptions=True) + except BaseException: + for task in tasks: + task.cancel() + results = await asyncio.gather(*tasks, return_exceptions=True) + await asyncio.gather( + *(result.close() for result in results if isinstance(result, DockerReMeSandbox)), + return_exceptions=True, + ) + raise + failures = [result for result in results if isinstance(result, BaseException)] + if failures: + await asyncio.gather( + *(result.close() for result in results if isinstance(result, DockerReMeSandbox)), + ) + raise failures[0] + return [result for result in results if isinstance(result, DockerReMeSandbox)] + + +class DockerReMeSandbox: + """One isolated Docker container and one fresh ReMe runtime workspace.""" + + def __init__( + self, + *, + case_id: str, + container_id: str, + candidate: SourceCandidate | ImageCandidate, + workspace: Workspace, + config: str, + plugins: tuple[str, ...], + command_timeout: float, + environment_names: list[str], + ) -> None: + self.case_id = case_id + self.container_id = container_id + self.candidate = candidate + self.workspace = workspace + self.config = config + self.plugins = plugins + self.command_timeout = command_timeout + self.environment_names = environment_names + self.backend: Backend | None = None + self.python = "python3" + self.actions: list[ActionRecord] = [] + self._closed = False + self._operation_lock = asyncio.Lock() + + @property + def runtime_workspace(self) -> str: + """ReMe-owned runtime workspace inside this case container.""" + return f"{CASE_ROOT}/reme_workspace" + + @property + def logs_dir(self) -> str: + """Directory containing ReMe and orchestration logs.""" + return f"{CASE_ROOT}/logs" + + @property + def results_dir(self) -> str: + """Directory containing job, answer and score artifacts.""" + return f"{CASE_ROOT}/results" + + async def initialize(self) -> None: + """Start Docker, upload the worker, and prepare the selected candidate.""" + await self.workspace.initialize() + 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) + + if isinstance(self.candidate, SourceCandidate): + await self._install_source_candidate() + else: + await self._validate_candidate(self.python) + + async def _create_case_layout(self) -> None: + """Create the disposable directories owned by the active case.""" + await self._exec_checked( + "create_case_layout", + [ + "mkdir", + "-p", + f"{CASE_ROOT}/inbox", + f"{CASE_ROOT}/tmp", + self.runtime_workspace, + self.logs_dir, + self.results_dir, + f"{CASE_ROOT}/build_log", + f"{CASE_ROOT}/queries", + HARNESS_ROOT, + ], + ) + await self._exec_checked( + "initialize_memory_history", + ["git", "init", "--quiet"], + cwd=self.runtime_workspace, + ) + + 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._exec_checked("create_candidate_dir", ["mkdir", "-p", CANDIDATE_ROOT]) + await self._exec_checked( + "extract_candidate", + ["tar", "-xzf", SOURCE_ARCHIVE, "-C", CANDIDATE_ROOT], + ) + await self._exec_checked( + "create_candidate_venv", + ["python3", "-m", "venv", "--system-site-packages", CANDIDATE_VENV], + ) + self.python = f"{CANDIDATE_VENV}/bin/python" + for index, project in enumerate(self.candidate.install_projects): + project_path = CANDIDATE_ROOT if project == "." else f"{CANDIDATE_ROOT}/{project}" + await self._exec_checked( + "install_candidate" if index == 0 else f"install_candidate:{project}", + [ + self.python, + "-m", + "pip", + "install", + "--no-deps", + "--no-build-isolation", + "--editable", + project_path, + ], + ) + await self._validate_candidate(self.python) + + async def _validate_candidate(self, python: str) -> None: + plugin_names = json.dumps(self.plugins) + script = ( + "import agentscope, importlib, json, reme; " + "from importlib.metadata import entry_points; " + "from reme.plugin import PluginManager; " + f"names={plugin_names}; " + "eps=entry_points().select(group='reme.plugins'); " + "selected={name: next((ep for ep in eps if ep.name == name), None) for name in names}; " + "PluginManager.discover(names); " + "paths={name: importlib.import_module(ep.value.split(':', 1)[0]).__file__ " + "for name, ep in selected.items() if ep is not None}; " + "print(json.dumps({'reme': reme.__file__, 'agentscope': agentscope.__version__, 'plugins': paths}))" + ) + result = await self._exec_checked("validate_candidate", [python, "-c", script]) + validation = json.loads(_decode(result.stdout)) + if validation.get("agentscope") != "2.0.9": + raise SandboxCommandError( + f"candidate environment requires agentscope 2.0.9, got {validation.get('agentscope')!r}", + ) + candidate_root = PurePosixPath(CANDIDATE_ROOT) + + def is_candidate_path(value: Any) -> bool: + return isinstance(value, str) and candidate_root in PurePosixPath(value).parents + + if isinstance(self.candidate, SourceCandidate) and not is_candidate_path(validation.get("reme")): + raise SandboxCommandError( + f"candidate import did not resolve under {CANDIDATE_ROOT}: {validation.get('reme')}", + ) + plugin_paths = validation.get("plugins") + if not isinstance(plugin_paths, dict) or set(plugin_paths) != set(self.plugins): + expected = list(self.plugins) + raise SandboxCommandError( + f"candidate did not expose the requested plugins: expected {expected!r}, got {plugin_paths!r}", + ) + if isinstance(self.candidate, SourceCandidate): + escaped = {name: path for name, path in plugin_paths.items() if not is_candidate_path(path)} + if escaped: + raise SandboxCommandError(f"candidate plugin imports did not resolve under {CANDIDATE_ROOT}: {escaped}") + + async def _exec(self, name: str, command: list[str], *, cwd: str | None = None) -> Any: + assert self.backend is not None, "sandbox is not initialized" + started = _utc_now() + result = await self.backend.exec_shell(command, cwd=cwd, timeout=self.command_timeout) + self.actions.append( + ActionRecord( + name=name, + exit_code=int(result.exit_code), + stdout=_decode(result.stdout), + stderr=_decode(result.stderr), + started_at=started, + finished_at=_utc_now(), + ), + ) + return result + + async def _exec_checked(self, name: str, command: list[str], *, cwd: str | None = None) -> Any: + result = await self._exec(name, command, cwd=cwd) + if not result.ok(): + raise SandboxCommandError( + f"sandbox action {name!r} failed with exit {result.exit_code}: {_decode(result.stderr)}", + ) + return result + + async def _run_job(self, job: str, arguments: dict[str, Any] | None = None) -> JobResult: + """Run one ReMe job while the caller holds the operation lock.""" + assert self.backend is not None, "sandbox is not initialized" + request_id = uuid.uuid4().hex + request_path = f"{CASE_ROOT}/inbox/{request_id}.json" + response_path = f"{self.results_dir}/{request_id}.json" + request = { + "job": job, + "arguments": arguments or {}, + "config": self.config, + "plugins": list(self.plugins), + "case_root": CASE_ROOT, + "workspace_dir": self.runtime_workspace, + } + await self.backend.write_file(request_path, json.dumps(request, ensure_ascii=False).encode("utf-8")) + await self._exec( + f"job:{job}", + [self.python, f"{HARNESS_ROOT}/worker.py", "--request", request_path, "--response", response_path], + cwd=CASE_ROOT, + ) + try: + payload = json.loads((await self.backend.read_file(response_path)).decode("utf-8")) + except FileNotFoundError as exc: + raise SandboxCommandError(f"job {job!r} did not produce {response_path}") from exc + return JobResult.from_dict(payload) + + async def _run_worker_request(self, mode: str, payload: dict[str, Any]) -> dict[str, Any]: + """Run one structured worker request and return its JSON response.""" + assert self.backend is not None, "sandbox is not initialized" + request_id = uuid.uuid4().hex + request_path = f"{CASE_ROOT}/inbox/{request_id}.json" + response_path = f"{self.results_dir}/{request_id}.json" + request = { + "mode": mode, + "config": self.config, + "plugins": list(self.plugins), + "case_id": self.case_id, + "case_root": CASE_ROOT, + "workspace_dir": self.runtime_workspace, + **payload, + } + await self.backend.write_file(request_path, json.dumps(request, ensure_ascii=False).encode("utf-8")) + await self._exec( + f"worker:{mode}", + [self.python, f"{HARNESS_ROOT}/worker.py", "--request", request_path, "--response", response_path], + cwd=CASE_ROOT, + ) + try: + result = json.loads((await self.backend.read_file(response_path)).decode("utf-8")) + except FileNotFoundError as exc: + raise SandboxCommandError(f"worker mode {mode!r} did not produce {response_path}") from exc + if not isinstance(result, dict): + raise SandboxCommandError(f"worker mode {mode!r} returned a non-object response") + return result + + async def run_job(self, job: str, arguments: dict[str, Any] | None = None) -> JobResult: + """Run one ReMe job without overlapping another operation.""" + async with self._operation_lock: + return await self._run_job(job, arguments) + + async def run_build(self, jobs: list[JobRequest | dict[str, Any]]) -> dict[str, Any]: + """Run all construction jobs in one Application and capture ``build.log``.""" + serialized: list[dict[str, Any]] = [] + for specification in jobs: + value = asdict(specification) if isinstance(specification, JobRequest) else dict(specification) + if not isinstance(value.get("job"), str) or not value["job"]: + raise ValueError("every build job requires a non-empty job name") + if not isinstance(value.get("arguments", {}), dict): + raise ValueError("build job arguments must be an object") + leaked = _PRIVATE_EVALUATION_FIELDS.intersection(value.get("arguments", {})) + if leaked: + names = ", ".join(sorted(leaked)) + raise ValueError(f"construction job arguments contain private query fields: {names}") + checkpoint = value.get("memory_checkpoint") + if checkpoint is not None and (not isinstance(checkpoint, str) or not checkpoint.strip()): + raise ValueError("build job memory_checkpoint must be a non-empty string or null") + serialized.append(value) + if not serialized: + raise ValueError("run_build requires at least one job") + async with self._operation_lock: + return await self._run_worker_request("build", {"jobs": serialized}) + + @staticmethod + def _validate_query_artifact_id(query_id: str) -> None: + """Keep verbatim query-directory names inside the artifact root.""" + if not query_id or len(query_id.encode("utf-8")) > 255: + raise ValueError("query_id must be non-empty and no longer than 255 UTF-8 bytes") + if query_id in {".", ".."} or "/" in query_id or "\\" in query_id or "\x00" in query_id: + raise ValueError(f"query_id is not a safe directory name: {query_id!r}") + if query_id == "summary.json": + raise ValueError("query_id conflicts with the queries summary: 'summary.json'") + + async def run_queries(self, queries: list[EvaluationQuery | dict[str, Any]]) -> dict[str, Any]: + """Answer multiple queries in one candidate Application with isolated logs.""" + serialized, query_ids = self._serialize_queries(queries) + if not serialized: + raise ValueError("run_queries requires at least one query") + if len(query_ids) != len(set(query_ids)): + raise ValueError("query IDs must be unique") + async with self._operation_lock: + return await self._run_worker_request("queries", {"queries": serialized}) + + async def run_query(self, query: EvaluationQuery | dict[str, Any]) -> dict[str, Any]: + """Answer one append-only query for a host-side lease.""" + + serialized, _ = self._serialize_queries([query]) + async with self._operation_lock: + response = await self._run_worker_request("query", {"queries": serialized}) + results = response.get("queries") + if not isinstance(results, list) or len(results) != 1 or not isinstance(results[0], dict): + raise SandboxCommandError("single-query worker returned an invalid result") + return results[0] + + def _serialize_queries( + self, + queries: list[EvaluationQuery | dict[str, Any]], + ) -> tuple[list[dict[str, Any]], list[str]]: + """Validate and serialize query payloads shared by batch and lease APIs.""" + + serialized: list[dict[str, Any]] = [] + query_ids: list[str] = [] + for query in queries: + value = asdict(query) if isinstance(query, EvaluationQuery) else dict(query) + query_id = value.get("query_id") + if not isinstance(query_id, str): + raise ValueError("every evaluation query requires a string query_id") + self._validate_query_artifact_id(query_id) + if not isinstance(value.get("question"), str): + raise ValueError(f"query {query_id!r} requires a string question") + answer_arguments = value.get("answer_arguments") or {} + if not isinstance(answer_arguments, dict): + raise ValueError(f"query {query_id!r} answer_arguments must be an object") + leaked = _PRIVATE_EVALUATION_FIELDS.intersection(answer_arguments) + if leaked: + names = ", ".join(sorted(leaked)) + raise ValueError(f"query {query_id!r} answer_arguments contain private judge fields: {names}") + query_ids.append(query_id) + # This allowlist is a security boundary. Golden answers, rubrics, + # question types and scoring configuration remain on the host and + # are not written to the candidate container's inbox. + serialized.append( + { + "query_id": query_id, + "question": value["question"], + "answer_arguments": dict(answer_arguments), + "answer_job": str(value.get("answer_job") or "agentic_answer"), + }, + ) + return serialized, query_ids + + async def ingest_session( + self, + *, + messages: list[dict[str, Any]], + session_id: str, + date: str = "", + memory_hint: str = "", + update_index: bool = True, + ) -> JobResult: + """Add one session, optionally making it searchable immediately.""" + async with self._operation_lock: + result = await self._run_job( + "auto_memory", + {"messages": messages, "session_id": session_id, "date": date, "memory_hint": memory_hint}, + ) + if result.success and update_index: + index_result = await self._run_job("index_update") + if not index_result.success: + return index_result + return result + + async def commit_memory_history(self, message: str) -> None: + """Create a host-requested checkpoint containing durable file-native memory.""" + if not isinstance(message, str) or not message.strip(): + raise ValueError("memory history commit message must be a non-empty string") + async with self._operation_lock: + await self._commit_memory_history(message) + + async def _daily_history_path(self) -> str: + """Return the validated workspace-relative daily directory from the runtime layout.""" + return (await self._memory_history_paths())[0] + + async def _memory_history_paths(self) -> list[str]: + """Return validated workspace-relative daily, digest, and session directories.""" + assert self.backend is not None + try: + payload = json.loads((await self.backend.read_file(RUNTIME_LAYOUT)).decode("utf-8")) + except (FileNotFoundError, UnicodeDecodeError, json.JSONDecodeError) as exc: + raise SandboxCommandError(f"cannot resolve daily history path from runtime layout: {exc}") from exc + + expected_root = PurePosixPath(self.runtime_workspace).relative_to(CASE_ROOT) + configured_paths = payload.get("configured_paths") if isinstance(payload, dict) else None + if not isinstance(configured_paths, dict): + raise SandboxCommandError("sandbox runtime layout has invalid configured_paths") + paths: list[str] = [] + for field in ("daily_dir", "digest_dir", "session_dir"): + case_relative = self._validated_workspace_path( + configured_paths.get(field), + expected_root, + allow_root=False, + ) + relative_path = PurePosixPath(case_relative).relative_to(expected_root).as_posix() + if relative_path not in paths: + paths.append(relative_path) + return paths + + async def _commit_memory_history(self, message: str) -> None: + """Commit one host-selected boundary across configured durable memory directories.""" + memory_paths = await self._memory_history_paths() + await self._exec_checked( + "stage_memory_history", + ["git", "add", "-A", "--", *memory_paths], + cwd=self.runtime_workspace, + ) + changed_paths: list[str] = [] + for path in memory_paths: + diff_result = await self._exec( + "check_memory_history", + ["git", "diff", "--cached", "--quiet", "--", path], + cwd=self.runtime_workspace, + ) + if diff_result.exit_code not in (0, 1): + raise SandboxCommandError( + f"cannot inspect staged memory changes: {_decode(diff_result.stderr)}", + ) + if diff_result.exit_code == 1: + changed_paths.append(path) + commit_command = [ + "git", + "-c", + f"user.name={_GIT_AUTHOR_NAME}", + "-c", + f"user.email={_GIT_AUTHOR_EMAIL}", + "commit", + "--quiet", + "--allow-empty", + "--only", + "-m", + message, + ] + if changed_paths: + commit_command.extend(["--", *changed_paths]) + await self._exec_checked( + "commit_memory_history", + commit_command, + cwd=self.runtime_workspace, + ) + + async def answer(self, *, query: str, query_time: str = "") -> JobResult: + """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) + 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( + f"{self.results_dir}/{name}", + json.dumps(asdict(result), ensure_ascii=False, default=str).encode("utf-8"), + ) + + async def export( + self, + destination: str | Path, + *, + profile: Literal["analysis", "full"] = "analysis", + include_candidate: bool = False, + ) -> Path: + """Download analysis artifacts, or the full case when requested.""" + if profile not in ("analysis", "full"): + raise ValueError("export profile must be 'analysis' or 'full'") + async with self._operation_lock: + return await self._export(destination, profile=profile, include_candidate=include_candidate) + + async def export_full(self, destination: str | Path, *, include_candidate: bool = False) -> Path: + """Download every file under the disposable case root.""" + return await self.export(destination, profile="full", include_candidate=include_candidate) + + async def export_workspace(self, destination: str | Path) -> Path: + """Download a portable snapshot of this runtime workspace to the host. + + The resulting gzip tar archive can be passed directly to + :meth:`upload_workspace` on any initialized case. It contains the + workspace contents (including its local Git history), not the enclosing + 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( + "export_workspace", + ["tar", "-czf", WORKSPACE_EXPORT_ARCHIVE, "-C", self.runtime_workspace, "."], + ) + 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 + + async def read_build_log(self) -> str: + """Return the build log, including when construction failed.""" + async with self._operation_lock: + return (await self._read_build_log()).decode("utf-8", errors="replace") + + async def _read_build_log(self) -> bytes: + assert self.backend is not None, "sandbox is not initialized" + source = f"{CASE_ROOT}/build_log/build.log" + try: + return await self.backend.read_file(source) + except FileNotFoundError as exc: + raise SandboxCommandError(f"build log is missing required artifact: {source}") from exc + + async def export_memory_construction(self, destination: str | Path) -> Path: + """Download the post-construction workspace and its build log.""" + async with self._operation_lock: + assert self.backend is not None, "sandbox is not initialized" + required = f"{CASE_ROOT}/build_log/build.log" + try: + await self.backend.read_file(required) + except FileNotFoundError as exc: + raise SandboxCommandError( + f"memory construction export is missing required artifact: {required}", + ) from exc + await self._exec_checked( + "export_memory_construction", + [ + "tar", + "-czf", + EXPORT_ARCHIVE, + "-C", + 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.""" + async with self._operation_lock: + assert self.backend is not None, "sandbox is not initialized" + required = f"{CASE_ROOT}/queries/summary.json" + try: + 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 + + 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: + assert self.backend is not None, "sandbox is not initialized" + query_root = f"{CASE_ROOT}/queries/{query_id}" + required = f"{query_root}/result.json" + try: + 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 + + async def export_evaluation(self, destination: str | Path) -> Path: + """Download workspace and query artifacts, plus the build log when present.""" + async with self._operation_lock: + assert self.backend is not None, "sandbox is not initialized" + summary = f"{CASE_ROOT}/queries/summary.json" + try: + await self.backend.read_file(summary) + except FileNotFoundError as exc: + raise SandboxCommandError(f"evaluation export is missing required artifact: {summary}") from exc + try: + await self.backend.read_file(f"{CASE_ROOT}/build_log/build.log") + except FileNotFoundError: + 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 + + async def upload_workspace(self, source: str | Path, *, clear: bool = True) -> None: + """Copy a host workspace directory or exported archive into this case. + + 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. + The copy is made through a validated archive, so each sandbox receives + an independent filesystem tree and can run queries in parallel. + """ + archive = _archive_workspace(source) + async with self._operation_lock: + assert self.backend is not None, "sandbox is not initialized" + 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._exec_checked( + "upload_workspace", + ["tar", "-xzf", WORKSPACE_ARCHIVE, "-C", self.runtime_workspace], + ) + # A directory copied from the host may not have Git metadata. Git + # init is idempotent and preserves history when the snapshot has it. + await self._exec_checked( + "initialize_uploaded_memory_history", + ["git", "init", "--quiet"], + cwd=self.runtime_workspace, + ) + await self._exec_checked("remove_workspace_upload", ["rm", "-f", WORKSPACE_ARCHIVE]) + + @staticmethod + def _validated_workspace_path(value: Any, workspace_root: PurePosixPath, *, allow_root: bool) -> str: + """Validate one case-relative path received from candidate-controlled state.""" + if not isinstance(value, str): + raise SandboxCommandError("sandbox runtime layout contains a non-string path") + path = PurePosixPath(value) + is_root = path == workspace_root + is_descendant = workspace_root in path.parents + is_unsafe = path.is_absolute() or ".." in path.parts + if is_unsafe or (is_root and not allow_root) or (not is_root and not is_descendant): + raise SandboxCommandError(f"sandbox runtime layout contains unsafe path: {value!r}") + return path.as_posix() + + async def _analysis_archive_spec(self) -> tuple[str, list[str]]: + """Read and validate the worker-resolved workspace export rules.""" + assert self.backend is not None + try: + payload = json.loads((await self.backend.read_file(RUNTIME_LAYOUT)).decode("utf-8")) + except FileNotFoundError as exc: + raise SandboxCommandError("analysis export requires at least one completed job") from exc + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise SandboxCommandError(f"invalid sandbox runtime layout: {exc}") from exc + + if not isinstance(payload, dict): + raise SandboxCommandError("sandbox runtime layout must be an object") + expected_root = PurePosixPath(self.runtime_workspace).relative_to(CASE_ROOT) + workspace_root_value = self._validated_workspace_path( + payload.get("workspace_root"), + expected_root, + allow_root=True, + ) + if PurePosixPath(workspace_root_value) != expected_root: + raise SandboxCommandError( + f"sandbox runtime layout changed workspace root: {workspace_root_value!r}", + ) + + values = payload.get("analysis_excludes") + if not isinstance(values, list): + raise SandboxCommandError("sandbox runtime layout has invalid analysis_excludes") + excludes: list[str] = [] + for value in values: + normalized = self._validated_workspace_path(value, expected_root, allow_root=False) + if normalized not in excludes: + 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" + if include_candidate and isinstance(self.candidate, ImageCandidate): + raise ValueError("an image candidate has no source directory to export") + candidate_id = ( + self.candidate.candidate_id + if isinstance(self.candidate, SourceCandidate) + else self.candidate.resolved_candidate_id + ) + manifest = { + "case_id": self.case_id, + "container_id": self.container_id, + "candidate_id": candidate_id, + "candidate_mode": "source" if isinstance(self.candidate, SourceCandidate) else "image", + "config": self.config, + "plugins": list(self.plugins), + "install_projects": ( + list(self.candidate.install_projects) if isinstance(self.candidate, SourceCandidate) else [] + ), + "environment_names": self.environment_names, + "export_profile": profile, + "exported_at": _utc_now(), + } + await self.backend.write_file( + f"{CASE_ROOT}/manifest.json", + json.dumps(manifest, ensure_ascii=False, indent=2).encode("utf-8"), + ) + await self.backend.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 + ), + ) + if profile == "analysis": + workspace_root, analysis_excludes = await self._analysis_archive_spec() + command = [ + "tar", + "-czf", + EXPORT_ARCHIVE, + "-C", + CASE_ROOT, + *(f"--exclude={path}" for path in analysis_excludes), + "manifest.json", + "runtime-layout.json", + "logs", + "results", + workspace_root, + ] + else: + command = ["tar", "-czf", EXPORT_ARCHIVE, "-C", CASE_ROOT, "."] + if include_candidate: + command.extend(["-C", "/workspace", "candidate"]) + await self._exec_checked("export_case", command) + payload = await self.backend.read_file(EXPORT_ARCHIVE) + target = Path(destination).resolve() + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(payload) + return target + + async def reset_case(self, case_id: str) -> None: + """Discard the active case and prepare this container for another. + + Candidate code, its virtual environment, and the uploaded worker stay + intact. The runtime workspace, requests, logs, results, manifest, and + previous export archive are removed. Operations are serialized so a + reset cannot overlap a job or export. + """ + _validate_case_id(case_id) + async with self._operation_lock: + assert self.backend is not None, "sandbox is not initialized" + await self._exec_checked( + "clear_case", + ["rm", "-rf", CASE_ROOT, EXPORT_ARCHIVE, WORKSPACE_ARCHIVE, WORKSPACE_EXPORT_ARCHIVE], + ) + await self._create_case_layout() + self.case_id = case_id + self.actions = [] + + async def close(self) -> None: + """Destroy the case container. Export before calling this method.""" + async with self._operation_lock: + if not self._closed: + self._closed = True + await self.workspace.close() + + async def __aenter__(self) -> "DockerReMeSandbox": + return self + + async def __aexit__(self, _exc_type: Any, _exc: Any, _tb: Any) -> None: + await self.close() diff --git a/deploy/benchmark_docker/models.py b/deploy/benchmark_docker/models.py new file mode 100644 index 00000000..53c40187 --- /dev/null +++ b/deploy/benchmark_docker/models.py @@ -0,0 +1,77 @@ +"""Small serializable contracts shared by the host and sandbox worker.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass(frozen=True, slots=True) +class JobResult: + """Structured result returned by a direct ReMe job invocation. + + ``token_usage`` contains the input/output/total token delta for each + configured agent wrapper during this invocation. Missing provider usage is + represented by ``None`` at the metric level. ``job_call_counts`` contains + non-zero ReMe job invocation deltas, including nested job tools. + """ + + success: bool + answer: Any = "" + metadata: dict[str, Any] = field(default_factory=dict) + error: str | None = None + token_usage: dict[str, dict[str, int | None]] = field(default_factory=dict) + job_call_counts: dict[str, int] = field(default_factory=dict) + + @classmethod + def from_dict(cls, value: dict[str, Any]) -> "JobResult": + """Build a result from the worker's JSON response.""" + return cls( + success=bool(value.get("success")), + answer=value.get("answer", ""), + metadata=value.get("metadata") if isinstance(value.get("metadata"), dict) else {}, + token_usage=value.get("token_usage") if isinstance(value.get("token_usage"), dict) else {}, + job_call_counts=value.get("job_call_counts") if isinstance(value.get("job_call_counts"), dict) else {}, + error=value.get("error") if isinstance(value.get("error"), str) else None, + ) + + +@dataclass(frozen=True, slots=True) +class ActionRecord: + """Audit record for one command executed inside a case sandbox.""" + + name: str + exit_code: int + stdout: str + stderr: str + started_at: str + finished_at: str + + +@dataclass(frozen=True, slots=True) +class JobRequest: + """One job executed as part of a shared build Application lifecycle.""" + + job: str + arguments: dict[str, Any] = field(default_factory=dict) + memory_checkpoint: str | None = None + + +@dataclass(frozen=True, slots=True) +class EvaluationQuery: + """One answer request plus host-only judging metadata. + + Only ``query_id``, ``question``, ``answer_arguments`` and ``answer_job`` + cross the candidate sandbox boundary. The remaining fields are retained + on the trusted host for post-answer evaluation. + """ + + query_id: str + 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 new file mode 100644 index 00000000..aea2161c --- /dev/null +++ b/deploy/benchmark_docker/pool.py @@ -0,0 +1,197 @@ +"""Bounded Docker workers for independent benchmark cases. + +Each case completes construction and its queries before its worker is assigned +another case. A dependent case may reuse the workspace captured immediately +after its parent's construction, without sharing mutable container state. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import Any + +from .docker_sandbox import DockerReMeSandbox, DockerReMeSandboxFactory +from .models import EvaluationQuery, JobRequest + + +@dataclass(frozen=True, slots=True) +class CasePlan: + """One case, optionally seeded by a host workspace or another case.""" + + case_id: str + build_jobs: tuple[JobRequest, ...] = () + queries: tuple[EvaluationQuery, ...] = () + workspace_source: str | Path | None = None + reuse_workspace_from: str | None = None + export_workspace_to: str | Path | None = None + export_build_log_to: str | Path | None = None + export_evaluation_to: str | Path | None = None + + +@dataclass(slots=True) +class CaseResult: + """Structured build and answer results; judging stays on the host.""" + + case_id: str + build: dict[str, Any] | None = None + queries: dict[str, Any] | None = None + workspace_snapshot: Path | None = None + build_log: str | None = None + error: str | None = None + + @property + def success(self) -> bool: + """Whether all requested phases completed successfully.""" + return ( + self.error is None + and (self.build is None or self.build.get("success") is True) + and (self.queries is None or self.queries.get("success") is True) + ) + + +class BenchmarkDockerPool: + """Reuse at most ``concurrency`` containers across case plans. + + ``run`` owns and closes its containers. It never stages all construction + phases ahead of queries: each available worker runs build then queries for + one case. Dependencies await a completed parent without holding a worker. + """ + + def __init__(self, factory: DockerReMeSandboxFactory, *, concurrency: int = 4) -> None: + if concurrency < 1: + raise ValueError("concurrency must be at least 1") + self.factory = factory + self.concurrency = concurrency + + async def run(self, plans: list[CasePlan]) -> list[CaseResult]: + """Return results in input order, closing every container on exit.""" + if not plans: + return [] + by_id = {plan.case_id: plan for plan in plans} + if len(by_id) != len(plans): + raise ValueError("case IDs must be unique") + for plan in plans: + if plan.workspace_source is not None and plan.reuse_workspace_from is not None: + raise ValueError(f"case {plan.case_id!r} has two workspace sources") + if not plan.build_jobs and plan.workspace_source is None and plan.reuse_workspace_from is None: + raise ValueError(f"case {plan.case_id!r} needs build jobs or a workspace source") + if plan.reuse_workspace_from is not None and plan.reuse_workspace_from not in by_id: + raise ValueError(f"unknown workspace source: {plan.reuse_workspace_from!r}") + if plan.export_build_log_to is not None and not plan.build_jobs: + raise ValueError("build log export requires build jobs") + if plan.export_evaluation_to is not None and not plan.queries: + raise ValueError("evaluation export requires queries") + self._check_acyclic(by_id) + + 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() + for worker in workers: + available.put_nowait(worker) + + try: + with TemporaryDirectory(prefix="reme-benchmark-pool-") as temporary: + snapshot_root = Path(temporary) + # Tasks are started together; awaiting a parent never consumes + # a container. The graph check above rules out deadlocks. + tasks: dict[str, asyncio.Task[CaseResult]] = {} + + async def execute(plan: CasePlan) -> CaseResult: + 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: + return CaseResult(plan.case_id, error=f"source case {parent.case_id!r} failed") + source = parent.workspace_snapshot + + worker = await available.get() + + 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: + await worker.upload_workspace(source) + if plan.build_jobs: + try: + 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() + return result + if not result.build.get("success"): + result.error = "build failed" + await retain_failed_build_log() + return result + if plan.export_build_log_to is not None: + 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 + ): + 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: + result.queries = await worker.run_queries(list(plan.queries)) + if not result.queries.get("success"): + result.error = "query failed" + if plan.export_evaluation_to is not None: + await worker.export_evaluation(plan.export_evaluation_to) + return result + except Exception as exc: # Keep other cases running and release this worker. + result.error = f"{type(exc).__name__}: {exc}" + return result + finally: + available.put_nowait(worker) + + for plan in plans: + tasks[plan.case_id] = asyncio.create_task(execute(plan)) + try: + results = await asyncio.gather(*(tasks[plan.case_id] for plan in plans)) + finally: + for task in tasks.values(): + if not task.done(): + 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 + finally: + await asyncio.gather(*(worker.close() for worker in workers), return_exceptions=True) + + @staticmethod + def _check_acyclic(plans: dict[str, CasePlan]) -> None: + visiting: set[str] = set() + visited: set[str] = set() + + def visit(case_id: str) -> None: + if case_id in visiting: + raise ValueError(f"workspace reuse cycle at {case_id!r}") + if case_id in visited: + return + visiting.add(case_id) + parent = plans[case_id].reuse_workspace_from + if parent is not None: + visit(parent) + visiting.remove(case_id) + visited.add(case_id) + + for case_id in plans: + visit(case_id) diff --git a/deploy/benchmark_docker/worker.py b/deploy/benchmark_docker/worker.py new file mode 100644 index 00000000..c4826881 --- /dev/null +++ b/deploy/benchmark_docker/worker.py @@ -0,0 +1,460 @@ +"""In-container worker that calls ReMe jobs directly, without HTTP. + +The host uploads this file into every case sandbox. It deliberately imports +ReMe only after process startup so source candidates can control exactly which +ReMe package is loaded. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import logging +import os +from pathlib import Path +import subprocess +import traceback +from typing import Any + +_WORKSPACE_PATH_FIELDS = ( + "metadata_dir", + "session_dir", + "mem_session_dir", + "resource_dir", + "daily_dir", + "digest_dir", + "dialog_dir", +) +_ANALYSIS_EXCLUDE_FIELDS = ("metadata_dir", "resource_dir") + + +def _exception_details(exc: BaseException) -> dict[str, Any]: + """Return small, structured exception metadata for host-side classification.""" + + details: dict[str, Any] = { + "module": type(exc).__module__, + "type": type(exc).__name__, + "message": str(exc), + } + response = getattr(exc, "response", None) + status_code = getattr(response, "status_code", None) + if isinstance(status_code, int): + details["status_code"] = status_code + error_number = getattr(exc, "errno", None) + if isinstance(error_number, int): + details["errno"] = error_number + return details + + +def _write_runtime_layout(case_root: Path, app_config: Any) -> None: + """Persist validated, non-secret paths needed by selective export.""" + resolved_case_root = case_root.resolve() + workspace_root = Path(app_config.workspace_dir).resolve() + try: + workspace_relative = workspace_root.relative_to(resolved_case_root) + except ValueError as exc: + raise ValueError(f"sandbox workspace must stay under {resolved_case_root}: {workspace_root}") from exc + + resolved_paths: dict[str, str | None] = {} + for field in _WORKSPACE_PATH_FIELDS: + configured = ( + str(Path(app_config.session_dir) / "dialog") + if field == "dialog_dir" + else str(getattr(app_config, field)).strip() + ) + if not configured: + resolved_paths[field] = None + continue + path = (workspace_root / configured).resolve() + try: + relative = path.relative_to(resolved_case_root) + except ValueError as exc: + raise ValueError(f"sandbox config path {field} must stay under {workspace_root}: {path}") from exc + if path != workspace_root and workspace_root not in path.parents: + raise ValueError(f"sandbox config path {field} must stay under {workspace_root}: {path}") + resolved_paths[field] = relative.as_posix() + + analysis_excludes = list( + dict.fromkeys(path for field in _ANALYSIS_EXCLUDE_FIELDS if (path := resolved_paths[field]) is not None), + ) + layout = { + "workspace_root": workspace_relative.as_posix(), + "configured_paths": resolved_paths, + "analysis_excludes": analysis_excludes, + } + (resolved_case_root / "runtime-layout.json").write_text( + json.dumps(layout, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + + +def _prepare_runtime(request: dict[str, Any]) -> tuple[Path, dict[str, Any]]: + """Resolve one request's config after confining process-local output.""" + case_root = Path(request["case_root"]) + case_tmp = case_root / "tmp" + case_tmp.mkdir(parents=True, exist_ok=True) + os.environ["TMPDIR"] = str(case_tmp) + os.chdir(case_root) + + # Import after setting TMPDIR so ReMe and its subprocesses cannot leave + # case-specific temporary files outside the disposable case root. + from reme.config import resolve_app_config + from reme.schema import ApplicationConfig + + plugins = request.get("plugins") or [] + if not isinstance(plugins, list) or not all(isinstance(name, str) and name for name in plugins): + raise TypeError("sandbox request plugins must be a list of non-empty strings") + app_config = resolve_app_config( + config=request.get("config") or "benchmark.yaml", + plugins=plugins, + workspace_dir=request["workspace_dir"], + log_to_console=False, + log_to_file=True, + ) + app_config["environment"] = dict(os.environ) + resolved_config = ApplicationConfig.model_validate(app_config) + _write_runtime_layout(case_root, resolved_config) + # The build checkpoint reads workspace paths from this mapping. Return the + # validated model's dump so optional config fields (including ``daily_dir``) + # have the same defaults as the Application instance. + return case_root, resolved_config.model_dump() + + +async def _run_job_on_app(app: Any, job: str, arguments: dict[str, Any]) -> dict[str, Any]: + """Run one job and return token and ReMe job invocation deltas.""" + from reme.enumeration import ComponentEnum + from reme.utils.evaluation_interface import track_agent_token_usage, track_job_counts + + agent_names = list(app.context.components.get(ComponentEnum.AGENT_WRAPPER, {})) + job_names = list(app.context.jobs) + token_usage: dict[str, dict[str, int | None]] = {} + job_call_counts: dict[str, int] = {} + try: + with ( + track_agent_token_usage(agent_names, app.context) as token_usage, + track_job_counts(job_names, app.context) as job_call_counts, + ): + response = await app.run_job(job, **arguments) + except Exception as exc: # Preserve usage accumulated before a job failure. + return { + "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), + } + + +def _atomic_write_json(path: Path, value: Any) -> None: + """Publish one JSON artifact without exposing a partially written file.""" + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2, default=str) + "\n", encoding="utf-8") + temporary.replace(path) + + +def _start_case_log(path: Path) -> Any: + """Attach an exact artifact path to ReMe's configured logger backend.""" + from reme.utils import get_logger + + path.parent.mkdir(parents=True, exist_ok=True) + logger = get_logger(log_to_console=False, log_to_file=False, force_init=True) + if hasattr(logger, "add"): # Loguru logger. + logger.add(str(path), level="INFO", encoding="utf-8") + else: # Standard library logger, selected by REME_DISABLE_LOGURU. + logger.addHandler(logging.FileHandler(path, encoding="utf-8")) + return logger + + +def _validate_query_id(query_id: Any) -> str: + """Validate an ID used verbatim as a directory name in the artifact.""" + if not isinstance(query_id, str) or not query_id or len(query_id.encode("utf-8")) > 255: + raise ValueError("query_id must be a non-empty string no longer than 255 UTF-8 bytes") + if query_id in {".", ".."} or "/" in query_id or "\\" in query_id or "\x00" in query_id: + raise ValueError(f"query_id is not a safe directory name: {query_id!r}") + if query_id == "summary.json": + raise ValueError("query_id conflicts with the queries summary: 'summary.json'") + return query_id + + +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) + + +async def _run(request: dict[str, Any]) -> dict[str, Any]: + """Construct an Application, run one job, and close it deterministically.""" + job = request["job"] + job_args = request.get("arguments") or {} + _, app_config = _prepare_runtime(request) + from reme.reme import ReMe + + app = ReMe(**app_config) + try: + await app.start() + return await _run_job_on_app(app, job, job_args) + finally: + await app.close() + + +async def _run_build(request: dict[str, Any]) -> dict[str, Any]: + """Execute all construction jobs in one Application and one build log.""" + case_root, app_config = _prepare_runtime(request) + from reme.reme import ReMe + from reme.utils import get_logger + + jobs = request.get("jobs") + if not isinstance(jobs, list) or not jobs: + raise ValueError("build request requires a non-empty jobs list") + build_log = case_root / "build_log" / "build.log" + if build_log.exists(): + raise FileExistsError(f"build artifact already exists: {build_log}") + + app = ReMe(**app_config) + 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, + } + + +def _commit_memory_checkpoint(case_root: Path, app_config: dict[str, Any], message: str) -> dict[str, Any]: + """Commit one host-declared boundary containing durable file-native memory.""" + + workspace_root = Path(app_config["workspace_dir"]).resolve() + try: + workspace_root.relative_to(case_root.resolve()) + except ValueError as exc: + raise ValueError("memory checkpoint path escapes the runtime workspace") from exc + relative_paths: list[str] = [] + for field in ("daily_dir", "digest_dir", "session_dir"): + path = (workspace_root / app_config[field]).resolve() + try: + relative_path = path.relative_to(workspace_root).as_posix() + except ValueError as exc: + raise ValueError(f"memory checkpoint {field} escapes the runtime workspace") from exc + if not relative_path or relative_path == ".": + raise ValueError(f"memory checkpoint {field} must not be the workspace root") + if relative_path not in relative_paths: + relative_paths.append(relative_path) + + def git(*arguments: str, allowed_exit_codes: tuple[int, ...] = (0,)) -> subprocess.CompletedProcess[str]: + result = subprocess.run( + ["git", *arguments], + cwd=workspace_root, + capture_output=True, + text=True, + check=False, + ) + if result.returncode not in allowed_exit_codes: + detail = (result.stderr or result.stdout).strip() + raise RuntimeError(f"git {' '.join(arguments)} failed: {detail}") + return result + + git("add", "-A", "--", *relative_paths) + changed_paths = [ + path + for path in relative_paths + if git("diff", "--cached", "--quiet", "--", path, allowed_exit_codes=(0, 1)).returncode == 1 + ] + commit = [ + "-c", + "user.name=ReMe Sandbox", + "-c", + "user.email=reme-sandbox@localhost", + "commit", + "--quiet", + "--allow-empty", + "--only", + "-m", + message, + ] + if changed_paths: + commit.extend(["--", *changed_paths]) + git(*commit) + return { + "commit_sha": git("rev-parse", "HEAD").stdout.strip(), + "message": message, + "paths": relative_paths, + } + + +async def _run_queries(request: dict[str, Any], *, write_summary: bool = True) -> dict[str, Any]: + """Execute answer requests in one candidate Application with isolated logs. + + ``write_summary=False`` is the append-only primitive used by the validation + scheduler. It runs exactly one independently leased query without sealing + the case-level artifact directory with ``summary.json``. + """ + case_root, app_config = _prepare_runtime(request) + from reme.reme import ReMe + from reme.utils import get_logger + + queries = request.get("queries") + if not isinstance(queries, list) or not queries: + raise ValueError("query request requires a non-empty queries list") + if not write_summary and len(queries) != 1: + raise ValueError("single-query request requires exactly one query") + query_ids = [_validate_query_id(query.get("query_id") if isinstance(query, dict) else None) for query in queries] + if len(query_ids) != len(set(query_ids)): + raise ValueError("query IDs must be unique") + queries_root = case_root / "queries" + if (queries_root / "summary.json").exists() or any((queries_root / query_id).exists() for query_id in query_ids): + raise FileExistsError("query artifacts already exist for this case") + queries_root.mkdir(parents=True, exist_ok=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) + + 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() + + summary = { + "schema_version": 1, + "case_id": request.get("case_id"), + "query_count": len(results), + "scored_count": 0, + "mean_score": None, + "queries": [ + { + "query_id": result["query_id"], + "score": result["score"], + **({"error": result["error"]} if result["error"] is not None else {}), + } + for result in results + ], + } + if write_summary: + _atomic_write_json(queries_root / "summary.json", summary) + return {"success": all(result["error"] is None for result in results), "summary": summary, "queries": results} + + +def main() -> int: + """Read one request file and always emit a structured response file.""" + parser = argparse.ArgumentParser() + parser.add_argument("--request", required=True) + parser.add_argument("--response", required=True) + args = parser.parse_args() + + request_path = Path(args.request) + response_path = Path(args.response) + response_path.parent.mkdir(parents=True, exist_ok=True) + try: + request = json.loads(request_path.read_text(encoding="utf-8")) + mode = request.get("mode", "job") + if mode == "job": + operation = _run(request) + elif mode == "build": + operation = _run_build(request) + elif mode == "queries": + operation = _run_queries(request) + elif mode == "query": + operation = _run_queries(request, write_summary=False) + else: + raise ValueError(f"unknown sandbox worker mode: {mode!r}") + result = asyncio.run(operation) + exit_code = 0 if result["success"] else 1 + except Exception as exc: # The host needs an artifact even for broken candidates. + result = { + "success": False, + "answer": "", + "metadata": {}, + "token_usage": {}, + "job_call_counts": {}, + "error": f"{type(exc).__name__}: {exc}", + "exception": _exception_details(exc), + "traceback": traceback.format_exc(), + } + exit_code = 1 + + response_path.write_text(json.dumps(result, ensure_ascii=False, default=str), encoding="utf-8") + return exit_code + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pyproject.toml b/pyproject.toml index 6b22bb6c..de87fd9c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,6 +72,10 @@ dev = [ "pytest>=9.1.1", "pytest-asyncio>=0.23", ] +benchmark-docker = [ + "reme-ai[as]", + "aiodocker", +] full = [ "reme-ai[core]", "reme-ai[dev]", @@ -87,12 +91,13 @@ Repository = "https://github.com/agentscope-ai/ReMe" reme = "reme.reme:main" [tool.setuptools] -packages = { find = { where = ["."], include = ["reme", "reme.*"], exclude = ["reme_studio*"] } } +packages = { find = { where = ["."], include = ["reme", "reme.*", "deploy", "deploy.*"], exclude = ["reme_studio*"] } } include-package-data = true [tool.setuptools.package-data] "*" = ["py.typed", "**/*.yaml", "**/*.json"] "reme.components.tokenizer" = ["stopwords"] +"deploy.benchmark_docker" = ["Dockerfile.*", "README.md"] [tool.setuptools.dynamic] version = { attr = "reme.__version__" } diff --git a/tests/integration/test_benchmark_docker.py b/tests/integration/test_benchmark_docker.py new file mode 100644 index 00000000..a22f7497 --- /dev/null +++ b/tests/integration/test_benchmark_docker.py @@ -0,0 +1,188 @@ +"""Opt-in, credential-free smoke test against real AgentScope Docker containers. + +Run with ``REME_DOCKER_INTEGRATION=1 pytest tests/integration/test_benchmark_docker.py -v``. +The base image is built from the repository Dockerfile unless +``REME_BENCHMARK_TEST_IMAGE`` names an existing one. Temporary image tags are +removed after the test; workspace files stay under pytest's temporary root. +""" + +from __future__ import annotations + +import asyncio +import os +from pathlib import Path +import shutil +import subprocess +import tarfile +import uuid + +import pytest + +from deploy.benchmark_docker import ( + BenchmarkDockerPool, + CasePlan, + DockerReMeSandboxFactory, + EvaluationQuery, + ImageCandidate, + JobRequest, + SourceCandidate, + SourceSnapshot, +) + +ROOT = Path(__file__).resolve().parents[2] +SMOKE_CONFIG = """\ +enable_logo: false +service: + backend: cli +jobs: + smoke_write: + backend: base + steps: + - backend: write_step + smoke_read: + backend: base + steps: + - backend: read_step +""" + + +def _docker(*arguments: str, cwd: Path | None = None) -> None: + subprocess.run(["docker", *arguments], cwd=cwd, check=True, timeout=1200) + + +@pytest.fixture(scope="module", name="base_image") +def _base_image(): + """Build the repository base image unless the caller supplies one.""" + if os.environ.get("REME_DOCKER_INTEGRATION") != "1": + pytest.skip("set REME_DOCKER_INTEGRATION=1 to run the Docker integration test") + supplied = os.environ.get("REME_BENCHMARK_TEST_IMAGE") + if supplied: + yield supplied + return + + tag = f"reme-benchmark-smoke:{uuid.uuid4().hex[:12]}" + _docker("build", "-f", "deploy/benchmark_docker/Dockerfile.base", "-t", tag, ".", cwd=ROOT) + try: + yield tag + finally: + _docker("image", "rm", tag) + + +def _candidate_source(tmp_path: Path) -> Path: + source = tmp_path / "candidate" + source.mkdir() + shutil.copytree( + ROOT / "reme", + source / "reme", + ignore=shutil.ignore_patterns("__pycache__", "*.pyc", "*.pyo"), + ) + shutil.copy2(ROOT / "pyproject.toml", source / "pyproject.toml") + shutil.copy2(ROOT / "README.md", source / "README.md") + (source / "smoke.yaml").write_text(SMOKE_CONFIG, encoding="utf-8") + return source + + +def _container_exists(name: str) -> bool: + result = subprocess.run( + ["docker", "container", "inspect", name], + capture_output=True, + check=False, + timeout=30, + ) + return result.returncode == 0 + + +def _query(query_id: str) -> EvaluationQuery: + return EvaluationQuery( + query_id=query_id, + question="Read the saved note", + golden_answer="docker smoke", + answer_job="smoke_read", + answer_arguments={"path": "note.md"}, + ) + + +@pytest.mark.asyncio +async def test_real_docker_source_pool_workspace_reuse(base_image, tmp_path): + """Run build, query, workspace reuse, and query-only export in a real container.""" + source = _candidate_source(tmp_path) + factory = DockerReMeSandboxFactory( + SourceCandidate(SourceSnapshot.from_directory(source), base_image=base_image), + config="/workspace/candidate/smoke.yaml", + ) + evaluation = tmp_path / "child-evaluation.tar.gz" + plans = [ + CasePlan( + "parent", + build_jobs=(JobRequest("smoke_write", {"path": "note.md", "content": "docker smoke"}),), + queries=(_query("parent-query"),), + ), + CasePlan( + "child", + reuse_workspace_from="parent", + queries=(_query("child-query"),), + export_evaluation_to=evaluation, + ), + ] + + results = await asyncio.wait_for(BenchmarkDockerPool(factory, concurrency=1).run(plans), timeout=360) + + assert all(result.success for result in results), results + assert "docker smoke" in str(results[1].queries["queries"][0]["answer"]) + with tarfile.open(evaluation, "r:gz") as archive: + names = archive.getnames() + assert any(name.endswith("reme_workspace/note.md") for name in names) + assert any(name.endswith("queries/child-query/answer.log") for name in names) + assert not any(name.endswith("build_log/build.log") for name in names) + + +@pytest.mark.asyncio +async def test_real_docker_source_plugin_install(base_image, tmp_path): + """Upload and install a plugin project in the candidate environment.""" + source = _candidate_source(tmp_path) + shutil.copytree( + ROOT / "plugins" / "lme", + source / "plugins" / "lme", + ignore=shutil.ignore_patterns("__pycache__", "*.pyc", "*.pyo"), + ) + factory = DockerReMeSandboxFactory( + SourceCandidate( + SourceSnapshot.from_directory(source), + base_image=base_image, + install_projects=(".", "plugins/lme"), + ), + plugins=("lme",), + config="/workspace/candidate/smoke.yaml", + ) + async with await asyncio.wait_for(factory.create_case("plugin-smoke"), timeout=180) as case: + container_name = f"as_ws_{case.container_id}" + assert _container_exists(container_name) + assert not _container_exists(container_name) + + +@pytest.mark.asyncio +async def test_real_docker_prebuilt_image(base_image, tmp_path): + """Install a fixed candidate into an image and execute its jobs.""" + source = _candidate_source(tmp_path) + dockerfile = source / "Dockerfile.smoke" + dockerfile.write_text( + f"FROM {base_image}\nCOPY . /opt/reme\nRUN python -m pip install --no-deps --no-build-isolation /opt/reme\n", + encoding="utf-8", + ) + image = f"reme-benchmark-candidate-smoke:{uuid.uuid4().hex[:12]}" + _docker("build", "-f", dockerfile.name, "-t", image, ".", cwd=source) + try: + factory = DockerReMeSandboxFactory(ImageCandidate(image), config="/opt/reme/smoke.yaml") + async with await factory.create_case("image-smoke") as case: + container_name = f"as_ws_{case.container_id}" + build = await asyncio.wait_for( + case.run_build([JobRequest("smoke_write", {"path": "note.md", "content": "image smoke"})]), + timeout=120, + ) + assert build["success"], build + query = await asyncio.wait_for(case.run_queries([_query("image-query")]), timeout=120) + assert query["success"], query + assert "image smoke" in str(query["queries"][0]["answer"]) + assert not _container_exists(container_name) + finally: + _docker("image", "rm", image) diff --git a/tests/unit/test_benchmark_docker.py b/tests/unit/test_benchmark_docker.py new file mode 100644 index 00000000..2a729974 --- /dev/null +++ b/tests/unit/test_benchmark_docker.py @@ -0,0 +1,383 @@ +"""Pool ordering, reuse, and cleanup without Docker or model calls.""" + +# pylint: disable=missing-class-docstring,missing-function-docstring,unused-argument + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest + +from deploy.benchmark_docker import ( + BenchmarkDockerPool, + CasePlan, + DockerReMeSandboxFactory, + EvaluationQuery, + JobRequest, + SourceCandidate, + SourceSnapshot, +) +from deploy.benchmark_docker.worker import _start_case_log + + +class FakeCase: + def __init__(self, name: str, events: list[tuple]) -> None: + self.name = name + self.events = events + self.case_id = name + + async def reset_case(self, case_id: str) -> None: + self.case_id = case_id + self.events.append((case_id, "reset")) + + async def upload_workspace(self, source: Path) -> None: + assert source.exists() + self.events.append((self.case_id, "upload", source.read_text(encoding="utf-8"))) + + async def run_build(self, jobs: list[JobRequest]) -> dict: + self.events.append((self.case_id, "build", len(jobs))) + await asyncio.sleep(0) + return {"success": True} + + async def export_workspace(self, destination: Path) -> Path: + destination = Path(destination) + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(self.case_id, encoding="utf-8") + self.events.append((self.case_id, "export_workspace")) + return destination + + async def run_queries(self, queries: list[EvaluationQuery]) -> dict: + self.events.append((self.case_id, "query", len(queries))) + return {"success": True} + + async def export_evaluation(self, destination: Path) -> Path: + self.events.append((self.case_id, "export_evaluation")) + return destination + + async def export_build_log(self, destination: Path) -> Path: + destination = Path(destination) + destination.write_text("build failed", encoding="utf-8") + self.events.append((self.case_id, "export_build_log")) + return destination + + async def read_build_log(self) -> str: + self.events.append((self.case_id, "read_build_log")) + return "build failed" + + async def close(self) -> None: + self.events.append((self.name, "close")) + + +class FakeFactory: + def __init__(self, events: list[tuple]) -> None: + self.events = events + + async def create_cases(self, case_ids: list[str], *, concurrency: int) -> list[FakeCase]: + assert len(case_ids) <= concurrency + return [FakeCase(name, self.events) for name in case_ids] + + +@pytest.mark.asyncio +async def test_pool_builds_and_queries_each_case_then_reuses_workspace(): + 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("parent", build_jobs=(JobRequest("build"),), queries=(EvaluationQuery("q1", "question", "secret"),)), + ] + results = await pool.run(plans) + + 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 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)) + assert events[-1] == ("pool-0", "close") + assert all(result.workspace_snapshot is None for result in results) + + +@pytest.mark.asyncio +async def test_pool_rejects_workspace_cycle_before_creating_containers(): + events: list[tuple] = [] + pool = BenchmarkDockerPool(FakeFactory(events), concurrency=2) + with pytest.raises(ValueError, match="cycle"): + await pool.run([CasePlan("a", reuse_workspace_from="b"), CasePlan("b", reuse_workspace_from="a")]) + 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 +async def test_pool_keeps_build_log_on_failure(tmp_path): + class FailedBuildCase(FakeCase): + async def run_build(self, jobs): + self.events.append((self.case_id, "build_failed")) + return {"success": False, "jobs": [{"error": "failed"}]} + + class FailedBuildFactory(FakeFactory): + async def create_cases(self, case_ids, *, concurrency): + return [FailedBuildCase(name, self.events) for name in case_ids] + + events: list[tuple] = [] + log_path = tmp_path / "build.log" + plan = CasePlan("failed", build_jobs=(JobRequest("build"),), export_build_log_to=log_path) + + result = (await BenchmarkDockerPool(FailedBuildFactory(events), concurrency=1).run([plan]))[0] + + assert not result.success + assert result.build == {"success": False, "jobs": [{"error": "failed"}]} + assert result.build_log == "build failed" + assert log_path.read_text(encoding="utf-8") == "build failed" + assert events.index(("failed", "read_build_log")) < events.index(("pool-0", "close")) + + +@pytest.mark.asyncio +async def test_pool_keeps_build_log_when_build_raises(tmp_path): + class TimedOutBuildCase(FakeCase): + async def run_build(self, jobs): + self.events.append((self.case_id, "build_timed_out")) + raise TimeoutError("build exceeded command timeout") + + class TimedOutBuildFactory(FakeFactory): + async def create_cases(self, case_ids, *, concurrency): + return [TimedOutBuildCase(name, self.events) for name in case_ids] + + events: list[tuple] = [] + log_path = tmp_path / "build.log" + plan = CasePlan("timed-out", build_jobs=(JobRequest("build"),), export_build_log_to=log_path) + + result = (await BenchmarkDockerPool(TimedOutBuildFactory(events), concurrency=1).run([plan]))[0] + + assert result.error == "TimeoutError: build exceeded command timeout" + assert result.build is None + assert result.build_log == "build failed" + assert log_path.read_text(encoding="utf-8") == "build failed" + assert events.index(("timed-out", "read_build_log")) < events.index(("pool-0", "close")) + + +@pytest.mark.asyncio +async def test_pool_cancels_case_before_closing_container(): + started = asyncio.Event() + events: list[tuple] = [] + + class BlockingCase(FakeCase): + async def run_queries(self, queries): + started.set() + try: + await asyncio.Event().wait() + finally: + self.events.append((self.case_id, "query_stopped")) + + class BlockingFactory(FakeFactory): + async def create_cases(self, case_ids, *, concurrency): + return [BlockingCase(name, self.events) for name in case_ids] + + plan = CasePlan("case", workspace_source=Path(__file__), queries=(EvaluationQuery("q", "question", "secret"),)) + task = asyncio.create_task(BenchmarkDockerPool(BlockingFactory(events), concurrency=1).run([plan])) + await asyncio.wait_for(started.wait(), timeout=2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert events.index(("case", "query_stopped")) < events.index(("pool-0", "close")) + + +def test_source_snapshot_excludes_unrelated_and_secret_files(tmp_path): + (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") + (tmp_path / ".env").write_text("secret", encoding="utf-8") + snapshot = SourceSnapshot.from_directory(tmp_path) + assert snapshot.file_count == 1 + + +def test_worker_writes_log_to_requested_artifact_path(tmp_path): + """The current ReMe logger accepts a directory, so bind the exact log file explicitly.""" + from reme.utils import get_logger + + path = tmp_path / "build_log" / "build.log" + logger = _start_case_log(path) + try: + logger.info("benchmark marker") + finally: + get_logger(log_to_console=False, log_to_file=False, force_init=True) + assert "benchmark marker" in path.read_text(encoding="utf-8") + + +class FakeBackend: + def __init__(self) -> None: + self.files: dict[str, bytes] = {} + self.commands: list[list[str]] = [] + + async def write_file(self, path: str, data: bytes) -> None: + self.files[path] = data + + async def read_file(self, path: str) -> bytes: + try: + return self.files[path] + except KeyError as exc: + raise FileNotFoundError(path) from exc + + async def exec_shell(self, command: list[str], *, cwd=None, timeout=None): + import json + from types import SimpleNamespace + + self.commands.append(command) + output = b"" + if len(command) > 2 and command[1] == "-c": + output = json.dumps( + { + "reme": "/workspace/candidate/reme/__init__.py", + "agentscope": "2.0.9", + "plugins": {"lme": "/workspace/candidate/plugins/lme/src/reme_lme/__init__.py"}, + }, + ).encode() + if "/workspace/harness/worker.py" in command: + request = json.loads(self.files[command[command.index("--request") + 1]]) + response = {"success": True, "jobs": []} if request.get("mode") == "build" else {"success": True} + self.files[command[command.index("--response") + 1]] = json.dumps(response).encode() + if command[:2] == ["tar", "-czf"]: + self.files[command[2]] = b"workspace archive" + return SimpleNamespace(exit_code=0, stdout=output, stderr=b"", ok=lambda: True) + + +class FakeWorkspace: + def __init__(self, backend: FakeBackend) -> None: + self.backend = backend + self.closed = False + + async def initialize(self) -> None: + pass + + def get_backend(self) -> FakeBackend: + return self.backend + + async def close(self) -> None: + self.closed = True + + +@pytest.mark.asyncio +async def test_source_candidate_installs_plugin_and_closes_container(tmp_path): + (tmp_path / "reme").mkdir() + (tmp_path / "reme" / "__init__.py").write_text("ok", encoding="utf-8") + backend = FakeBackend() + workspace = FakeWorkspace(backend) + factory = DockerReMeSandboxFactory( + SourceCandidate(SourceSnapshot.from_directory(tmp_path), install_projects=(".", "plugins/lme")), + plugins=("lme",), + workspace_builder=lambda _id, _image, _env: workspace, + ) + case = await factory.create_case("example") + try: + assert any(command[-1] == "/workspace/candidate/plugins/lme" for command in backend.commands) + assert "/workspace/candidate.tar.gz" in backend.files + assert "/workspace/harness/worker.py" in backend.files + assert (await case.run_build([JobRequest("auto_memory")]))["success"] + exported = await case.export_workspace(tmp_path / "workspace.tar.gz") + assert exported.read_bytes() == b"workspace archive" + with pytest.raises(ValueError, match="private"): + await case.run_queries( + [EvaluationQuery("q1", "question", "secret", answer_arguments={"golden_answer": "secret"})], + ) + finally: + await case.close() + assert workspace.closed + + +@pytest.mark.asyncio +async def test_query_only_evaluation_export_omits_missing_build_log(tmp_path): + (tmp_path / "reme").mkdir() + (tmp_path / "reme" / "__init__.py").write_text("ok", encoding="utf-8") + backend = FakeBackend() + workspace = FakeWorkspace(backend) + factory = DockerReMeSandboxFactory( + SourceCandidate(SourceSnapshot.from_directory(tmp_path)), + plugins=("lme",), + workspace_builder=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 +async def test_cancelled_case_initialization_closes_workspace(tmp_path): + (tmp_path / "reme").mkdir() + (tmp_path / "reme" / "__init__.py").write_text("ok", encoding="utf-8") + started = asyncio.Event() + + class BlockingWorkspace(FakeWorkspace): + async def initialize(self): + started.set() + await asyncio.Event().wait() + + workspace = BlockingWorkspace(FakeBackend()) + factory = DockerReMeSandboxFactory( + SourceCandidate(SourceSnapshot.from_directory(tmp_path)), + workspace_builder=lambda _id, _image, _env: workspace, + ) + task = asyncio.create_task(factory.create_case("cancelled")) + await asyncio.wait_for(started.wait(), timeout=2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert workspace.closed + + +@pytest.mark.asyncio +async def test_cancelled_batch_creation_closes_completed_workspace(tmp_path): + (tmp_path / "reme").mkdir() + (tmp_path / "reme" / "__init__.py").write_text("ok", encoding="utf-8") + second_started = asyncio.Event() + + class BlockingWorkspace(FakeWorkspace): + async def initialize(self): + second_started.set() + await asyncio.Event().wait() + + first = FakeWorkspace(FakeBackend()) + second = BlockingWorkspace(FakeBackend()) + workspaces = iter((first, second)) + factory = DockerReMeSandboxFactory( + SourceCandidate(SourceSnapshot.from_directory(tmp_path)), + plugins=("lme",), + workspace_builder=lambda _id, _image, _env: next(workspaces), + ) + task = asyncio.create_task(factory.create_cases(["first", "second"], concurrency=2)) + await asyncio.wait_for(second_started.wait(), timeout=2) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert first.closed and second.closed diff --git a/tests/unit/test_package_versions.py b/tests/unit/test_package_versions.py index 88f46847..698f9511 100644 --- a/tests/unit/test_package_versions.py +++ b/tests/unit/test_package_versions.py @@ -62,7 +62,12 @@ def test_studio_packages_have_independent_identity() -> None: assert auto_fin_config["project"]["version"] == "0.1.3" assert daily_paper_config["project"]["version"] == "0.1.3" assert dingtalk_config["project"]["version"] == "0.1.0" - assert main_config["tool"]["setuptools"]["packages"]["find"]["include"] == ["reme", "reme.*"] + assert main_config["tool"]["setuptools"]["packages"]["find"]["include"] == [ + "reme", + "reme.*", + "deploy", + "deploy.*", + ] assert "reme_studio*" in main_config["tool"]["setuptools"]["packages"]["find"]["exclude"] From 132a0ee7cc6037d939c9f2caf2a2de7e14e5983f 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 13:57:24 +0800 Subject: [PATCH 02/14] Fix benchmark worker reuse after Docker command timeout --- deploy/benchmark_docker/README.md | 5 +- deploy/benchmark_docker/docker_sandbox.py | 21 +++++++- deploy/benchmark_docker/pool.py | 15 +++++- tests/unit/test_benchmark_docker.py | 64 +++++++++++++++++++++++ 4 files changed, 101 insertions(+), 4 deletions(-) diff --git a/deploy/benchmark_docker/README.md b/deploy/benchmark_docker/README.md index d6f6259b..0ec0cb77 100644 --- a/deploy/benchmark_docker/README.md +++ b/deploy/benchmark_docker/README.md @@ -60,8 +60,9 @@ archive. A reused workspace may be queried without additional build jobs. `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` before reusing -or closing the container. +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. 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/docker_sandbox.py b/deploy/benchmark_docker/docker_sandbox.py index 04519558..b1c75e93 100644 --- a/deploy/benchmark_docker/docker_sandbox.py +++ b/deploy/benchmark_docker/docker_sandbox.py @@ -80,6 +80,10 @@ class SandboxCommandError(RuntimeError): """A required command failed inside a case sandbox.""" +class SandboxCommandTimeoutError(SandboxCommandError): + """A command timed out and may still be running inside the container.""" + + def _validate_case_id(case_id: str) -> None: """Reject identifiers that are unsafe for Docker names and manifests.""" if not _CASE_ID_RE.fullmatch(case_id): @@ -332,8 +336,14 @@ class DockerReMeSandbox: self.python = "python3" self.actions: list[ActionRecord] = [] self._closed = False + self._reusable = True self._operation_lock = asyncio.Lock() + @property + def reusable(self) -> bool: + """Whether another case can safely use this container.""" + return self._reusable and not self._closed + @property def runtime_workspace(self) -> str: """ReMe-owned runtime workspace inside this case container.""" @@ -457,8 +467,14 @@ class DockerReMeSandbox: 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: + raise SandboxCommandError("sandbox container is no longer reusable") started = _utc_now() - result = await self.backend.exec_shell(command, cwd=cwd, timeout=self.command_timeout) + try: + result = await self.backend.exec_shell(command, cwd=cwd, timeout=self.command_timeout) + except TimeoutError as exc: + self._reusable = False + raise SandboxCommandTimeoutError(f"sandbox action {name!r} timed out; container must be destroyed") from exc self.actions.append( ActionRecord( name=name, @@ -469,6 +485,9 @@ class DockerReMeSandbox: finished_at=_utc_now(), ), ) + if result.exit_code == -1: + self._reusable = False + raise SandboxCommandTimeoutError(f"sandbox action {name!r} timed out; container must be destroyed") return result async def _exec_checked(self, name: str, command: list[str], *, cwd: str | None = None) -> Any: diff --git a/deploy/benchmark_docker/pool.py b/deploy/benchmark_docker/pool.py index aea2161c..cdd41101 100644 --- a/deploy/benchmark_docker/pool.py +++ b/deploy/benchmark_docker/pool.py @@ -93,6 +93,7 @@ class BenchmarkDockerPool: available: asyncio.Queue[DockerReMeSandbox] = asyncio.Queue() for worker in workers: available.put_nowait(worker) + claimed = 0 try: with TemporaryDirectory(prefix="reme-benchmark-pool-") as temporary: @@ -102,6 +103,7 @@ class BenchmarkDockerPool: tasks: dict[str, asyncio.Task[CaseResult]] = {} async def execute(plan: CasePlan) -> CaseResult: + nonlocal claimed result = CaseResult(plan.case_id) source = plan.workspace_source if plan.reuse_workspace_from is not None: @@ -111,6 +113,7 @@ class BenchmarkDockerPool: source = parent.workspace_snapshot worker = await available.get() + claimed += 1 async def retain_failed_build_log() -> None: try: @@ -158,7 +161,17 @@ class BenchmarkDockerPool: result.error = f"{type(exc).__name__}: {exc}" return result finally: - available.put_nowait(worker) + 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) for plan in plans: tasks[plan.case_id] = asyncio.create_task(execute(plan)) diff --git a/tests/unit/test_benchmark_docker.py b/tests/unit/test_benchmark_docker.py index 2a729974..aefe3839 100644 --- a/tests/unit/test_benchmark_docker.py +++ b/tests/unit/test_benchmark_docker.py @@ -18,6 +18,7 @@ from deploy.benchmark_docker import ( SourceCandidate, SourceSnapshot, ) +from deploy.benchmark_docker.docker_sandbox import SandboxCommandError, SandboxCommandTimeoutError from deploy.benchmark_docker.worker import _start_case_log @@ -26,6 +27,7 @@ class FakeCase: self.name = name self.events = events self.case_id = name + self.reusable = True async def reset_case(self, case_id: str) -> None: self.case_id = case_id @@ -263,6 +265,17 @@ class FakeBackend: return SimpleNamespace(exit_code=0, stdout=output, stderr=b"", ok=lambda: True) +class TimedOutBackend(FakeBackend): + async def exec_shell(self, command: list[str], *, cwd=None, timeout=None): + if "/workspace/harness/worker.py" in command: + from types import SimpleNamespace + + self.commands.append(command) + self.files["/workspace/case/build_log/build.log"] = b"build still running" + return SimpleNamespace(exit_code=-1, stdout=b"", stderr=b"timed out", ok=lambda: False) + return await super().exec_shell(command, cwd=cwd, timeout=timeout) + + class FakeWorkspace: def __init__(self, backend: FakeBackend) -> None: self.backend = backend @@ -278,6 +291,30 @@ class FakeWorkspace: self.closed = True +@pytest.mark.asyncio +async def test_backend_timeout_prevents_direct_case_reset(tmp_path): + (tmp_path / "reme").mkdir() + (tmp_path / "reme" / "__init__.py").write_text("ok", encoding="utf-8") + workspace = FakeWorkspace(TimedOutBackend()) + factory = DockerReMeSandboxFactory( + SourceCandidate(SourceSnapshot.from_directory(tmp_path)), + plugins=("lme",), + workspace_builder=lambda _id, _image, _env: workspace, + ) + case = await factory.create_case("timed-out") + try: + with pytest.raises(SandboxCommandTimeoutError): + await case.run_build([JobRequest("build")]) + assert not case.reusable + assert case.actions[-1].exit_code == -1 + assert await case.read_build_log() == "build still running" + 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_source_candidate_installs_plugin_and_closes_container(tmp_path): (tmp_path / "reme").mkdir() @@ -306,6 +343,33 @@ async def test_source_candidate_installs_plugin_and_closes_container(tmp_path): assert workspace.closed +@pytest.mark.asyncio +async def test_pool_replaces_worker_after_backend_timeout(tmp_path): + (tmp_path / "reme").mkdir() + (tmp_path / "reme" / "__init__.py").write_text("ok", encoding="utf-8") + timed_out_workspace = FakeWorkspace(TimedOutBackend()) + replacement_workspace = FakeWorkspace(FakeBackend()) + workspaces = iter((timed_out_workspace, replacement_workspace)) + factory = DockerReMeSandboxFactory( + SourceCandidate(SourceSnapshot.from_directory(tmp_path)), + plugins=("lme",), + workspace_builder=lambda _id, _image, _env: next(workspaces), + ) + plans = [ + CasePlan("timed-out", build_jobs=(JobRequest("build"),), export_build_log_to=tmp_path / "build.log"), + CasePlan("next", build_jobs=(JobRequest("build"),)), + ] + + results = await BenchmarkDockerPool(factory, concurrency=1).run(plans) + + assert "SandboxCommandTimeoutError" in results[0].error + assert results[0].build_log == "build still running" + assert (tmp_path / "build.log").read_text(encoding="utf-8") == "build still running" + assert results[1].success + assert timed_out_workspace.closed and replacement_workspace.closed + 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): (tmp_path / "reme").mkdir() From 082b8f8ccbcac10ff5290591bd3a6bbada3c1d12 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 16:44:08 +0800 Subject: [PATCH 03/14] Simplify benchmark worker requests and preserve artifacts --- deploy/benchmark_docker/README.md | 6 + deploy/benchmark_docker/docker_sandbox.py | 158 +++++------- deploy/benchmark_docker/worker.py | 41 ++- tests/unit/test_benchmark_docker.py | 292 ++++++++++++++++------ 4 files changed, 307 insertions(+), 190 deletions(-) diff --git a/deploy/benchmark_docker/README.md b/deploy/benchmark_docker/README.md index 0ec0cb77..b1ec0ad4 100644 --- a/deploy/benchmark_docker/README.md +++ b/deploy/benchmark_docker/README.md @@ -63,11 +63,17 @@ and query-only cases; it includes `build_log` when that phase ran. Set 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. For direct control, call `create_case()` or `create_cases()`, then use `run_job()`, `run_build()`, `run_queries()`, `export_workspace()`, `upload_workspace()`, `reset_case()`, and `close()`. Always close a case in a `finally` block or use its async context manager. Export before closing. +The default `export("analysis.tar.gz")` also saves available construction and +query logs under `analysis.tar.gz.artifacts/build_log/` and +`analysis.tar.gz.artifacts/queries/`. These logs stay outside the archive; +choose a new destination if that artifacts directory already exists. The candidate receives only query ID, question, answer Job, and answer arguments. Golden answers, rubrics, and scoring remain on the host. Workspace diff --git a/deploy/benchmark_docker/docker_sandbox.py b/deploy/benchmark_docker/docker_sandbox.py index b1c75e93..a0928ee8 100644 --- a/deploy/benchmark_docker/docker_sandbox.py +++ b/deploy/benchmark_docker/docker_sandbox.py @@ -28,8 +28,6 @@ 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}$") -_GIT_AUTHOR_NAME = "ReMe Sandbox" -_GIT_AUTHOR_EMAIL = "reme-sandbox@localhost" _PRIVATE_EVALUATION_FIELDS = frozenset( {"query", "question", "golden_answer", "question_type", "rubric", "judge_arguments"}, ) @@ -500,29 +498,8 @@ class DockerReMeSandbox: async def _run_job(self, job: str, arguments: dict[str, Any] | None = None) -> JobResult: """Run one ReMe job while the caller holds the operation lock.""" - assert self.backend is not None, "sandbox is not initialized" - request_id = uuid.uuid4().hex - request_path = f"{CASE_ROOT}/inbox/{request_id}.json" - response_path = f"{self.results_dir}/{request_id}.json" - request = { - "job": job, - "arguments": arguments or {}, - "config": self.config, - "plugins": list(self.plugins), - "case_root": CASE_ROOT, - "workspace_dir": self.runtime_workspace, - } - await self.backend.write_file(request_path, json.dumps(request, ensure_ascii=False).encode("utf-8")) - await self._exec( - f"job:{job}", - [self.python, f"{HARNESS_ROOT}/worker.py", "--request", request_path, "--response", response_path], - cwd=CASE_ROOT, - ) - try: - payload = json.loads((await self.backend.read_file(response_path)).decode("utf-8")) - except FileNotFoundError as exc: - raise SandboxCommandError(f"job {job!r} did not produce {response_path}") from exc - return JobResult.from_dict(payload) + response = await self._run_worker_request("job", {"job": job, "arguments": arguments or {}}) + return JobResult.from_dict(response) async def _run_worker_request(self, mode: str, payload: dict[str, Any]) -> dict[str, Any]: """Run one structured worker request and return its JSON response.""" @@ -540,15 +517,17 @@ class DockerReMeSandbox: **payload, } await self.backend.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( - f"worker:{mode}", + action_name, [self.python, f"{HARNESS_ROOT}/worker.py", "--request", request_path, "--response", response_path], cwd=CASE_ROOT, ) try: result = json.loads((await self.backend.read_file(response_path)).decode("utf-8")) except FileNotFoundError as exc: - raise SandboxCommandError(f"worker mode {mode!r} did not produce {response_path}") from exc + description = f"job {payload['job']!r}" if mode == "job" else f"worker mode {mode!r}" + raise SandboxCommandError(f"{description} did not produce {response_path}") from exc if not isinstance(result, dict): raise SandboxCommandError(f"worker mode {mode!r} returned a non-object response") return result @@ -674,77 +653,9 @@ class DockerReMeSandbox: if not isinstance(message, str) or not message.strip(): raise ValueError("memory history commit message must be a non-empty string") async with self._operation_lock: - await self._commit_memory_history(message) - - async def _daily_history_path(self) -> str: - """Return the validated workspace-relative daily directory from the runtime layout.""" - return (await self._memory_history_paths())[0] - - async def _memory_history_paths(self) -> list[str]: - """Return validated workspace-relative daily, digest, and session directories.""" - assert self.backend is not None - try: - payload = json.loads((await self.backend.read_file(RUNTIME_LAYOUT)).decode("utf-8")) - except (FileNotFoundError, UnicodeDecodeError, json.JSONDecodeError) as exc: - raise SandboxCommandError(f"cannot resolve daily history path from runtime layout: {exc}") from exc - - expected_root = PurePosixPath(self.runtime_workspace).relative_to(CASE_ROOT) - configured_paths = payload.get("configured_paths") if isinstance(payload, dict) else None - if not isinstance(configured_paths, dict): - raise SandboxCommandError("sandbox runtime layout has invalid configured_paths") - paths: list[str] = [] - for field in ("daily_dir", "digest_dir", "session_dir"): - case_relative = self._validated_workspace_path( - configured_paths.get(field), - expected_root, - allow_root=False, - ) - relative_path = PurePosixPath(case_relative).relative_to(expected_root).as_posix() - if relative_path not in paths: - paths.append(relative_path) - return paths - - async def _commit_memory_history(self, message: str) -> None: - """Commit one host-selected boundary across configured durable memory directories.""" - memory_paths = await self._memory_history_paths() - await self._exec_checked( - "stage_memory_history", - ["git", "add", "-A", "--", *memory_paths], - cwd=self.runtime_workspace, - ) - changed_paths: list[str] = [] - for path in memory_paths: - diff_result = await self._exec( - "check_memory_history", - ["git", "diff", "--cached", "--quiet", "--", path], - cwd=self.runtime_workspace, - ) - if diff_result.exit_code not in (0, 1): - raise SandboxCommandError( - f"cannot inspect staged memory changes: {_decode(diff_result.stderr)}", - ) - if diff_result.exit_code == 1: - changed_paths.append(path) - commit_command = [ - "git", - "-c", - f"user.name={_GIT_AUTHOR_NAME}", - "-c", - f"user.email={_GIT_AUTHOR_EMAIL}", - "commit", - "--quiet", - "--allow-empty", - "--only", - "-m", - message, - ] - if changed_paths: - commit_command.extend(["--", *changed_paths]) - await self._exec_checked( - "commit_memory_history", - commit_command, - cwd=self.runtime_workspace, - ) + result = await self._run_worker_request("checkpoint", {"message": message}) + if result.get("success") is not True: + raise SandboxCommandError(result.get("error") or "memory history checkpoint failed") async def answer(self, *, query: str, query_time: str = "") -> JobResult: """Generate an answer with the candidate's agentic_answer job.""" @@ -799,7 +710,11 @@ class DockerReMeSandbox: profile: Literal["analysis", "full"] = "analysis", include_candidate: bool = False, ) -> Path: - """Download analysis artifacts, or the full case when requested.""" + """Download the case archive and, for analysis, separate build/query logs. + + Analysis logs are saved under ``.artifacts`` without + changing the archive's contents. + """ if profile not in ("analysis", "full"): raise ValueError("export profile must be 'analysis' or 'full'") async with self._operation_lock: @@ -1034,6 +949,10 @@ class DockerReMeSandbox: ) -> 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") candidate_id = ( @@ -1086,11 +1005,50 @@ class DockerReMeSandbox: command.extend(["-C", "/workspace", "candidate"]) await self._exec_checked("export_case", command) payload = await self.backend.read_file(EXPORT_ARCHIVE) - target = Path(destination).resolve() target.parent.mkdir(parents=True, exist_ok=True) target.write_bytes(payload) + if profile == "analysis": + await self._export_analysis_logs(artifacts) return target + async def _export_analysis_logs(self, destination: Path) -> None: + """Save available build and query artifacts alongside the analysis archive.""" + assert self.backend is not None + try: + build_log = await self.backend.read_file(f"{CASE_ROOT}/build_log/build.log") + except FileNotFoundError: + build_log = None + if build_log is not None: + log_path = destination / "build_log" / "build.log" + log_path.parent.mkdir(parents=True, exist_ok=True) + log_path.write_bytes(build_log) + + try: + await self.backend.read_file(f"{CASE_ROOT}/queries/summary.json") + except FileNotFoundError: + return + await self._exec_checked( + "export_analysis_queries", + ["tar", "-czf", EXPORT_ARCHIVE, "-C", CASE_ROOT, "queries"], + ) + archive = await self.backend.read_file(EXPORT_ARCHIVE) + with tarfile.open(fileobj=io.BytesIO(archive), mode="r:gz") as source: + for member in source: + path = _workspace_member_path(member.name) + if path is None or not path.parts or path.parts[0] != "queries": + raise SandboxCommandError(f"query export contains unsafe path: {member.name!r}") + output = destination.joinpath(*path.parts) + if member.isdir(): + output.mkdir(parents=True, exist_ok=True) + elif member.isfile(): + content = source.extractfile(member) + if content is None: + raise SandboxCommandError(f"query export cannot read: {member.name!r}") + output.parent.mkdir(parents=True, exist_ok=True) + output.write_bytes(content.read()) + else: + raise SandboxCommandError(f"query export contains unsupported member: {member.name!r}") + async def reset_case(self, case_id: str) -> None: """Discard the active case and prepare this container for another. diff --git a/deploy/benchmark_docker/worker.py b/deploy/benchmark_docker/worker.py index c4826881..c5f764ef 100644 --- a/deploy/benchmark_docker/worker.py +++ b/deploy/benchmark_docker/worker.py @@ -211,18 +211,30 @@ async def _run(request: dict[str, Any]) -> dict[str, Any]: async def _run_build(request: dict[str, Any]) -> dict[str, Any]: """Execute all construction jobs in one Application and one build log.""" - case_root, app_config = _prepare_runtime(request) - from reme.reme import ReMe - from reme.utils import get_logger - - jobs = request.get("jobs") - if not isinstance(jobs, list) or not jobs: - raise ValueError("build request requires a non-empty jobs list") - build_log = case_root / "build_log" / "build.log" + build_log = Path(request["case_root"]) / "build_log" / "build.log" if build_log.exists(): raise FileExistsError(f"build artifact already exists: {build_log}") + build_log.parent.mkdir(parents=True, exist_ok=True) + build_log.touch() - app = ReMe(**app_config) + try: + case_root, app_config = _prepare_runtime(request) + from reme.reme import ReMe + from reme.utils import get_logger + + jobs = request.get("jobs") + 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) + 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]] = [] @@ -313,6 +325,15 @@ def _commit_memory_checkpoint(case_root: Path, app_config: dict[str, Any], messa } +async def _run_checkpoint(request: dict[str, Any]) -> dict[str, Any]: + """Commit host-requested memory history using the build checkpoint rules.""" + message = request.get("message") + if not isinstance(message, str) or not message.strip(): + raise ValueError("memory history commit message must be a non-empty string") + case_root, app_config = _prepare_runtime(request) + return {"success": True, "memory_checkpoint": _commit_memory_checkpoint(case_root, app_config, message)} + + async def _run_queries(request: dict[str, Any], *, write_summary: bool = True) -> dict[str, Any]: """Execute answer requests in one candidate Application with isolated logs. @@ -435,6 +456,8 @@ def main() -> int: operation = _run_queries(request) elif mode == "query": operation = _run_queries(request, write_summary=False) + elif mode == "checkpoint": + operation = _run_checkpoint(request) else: raise ValueError(f"unknown sandbox worker mode: {mode!r}") result = asyncio.run(operation) diff --git a/tests/unit/test_benchmark_docker.py b/tests/unit/test_benchmark_docker.py index aefe3839..2c82ab48 100644 --- a/tests/unit/test_benchmark_docker.py +++ b/tests/unit/test_benchmark_docker.py @@ -1,11 +1,15 @@ """Pool ordering, reuse, and cleanup without Docker or model calls.""" -# pylint: disable=missing-class-docstring,missing-function-docstring,unused-argument +# pylint: disable=missing-class-docstring,missing-function-docstring,protected-access,unused-argument from __future__ import annotations import asyncio +import io +import json from pathlib import Path +import subprocess +import tarfile import pytest @@ -19,7 +23,7 @@ from deploy.benchmark_docker import ( SourceSnapshot, ) from deploy.benchmark_docker.docker_sandbox import SandboxCommandError, SandboxCommandTimeoutError -from deploy.benchmark_docker.worker import _start_case_log +from deploy.benchmark_docker import worker as sandbox_worker class FakeCase: @@ -131,10 +135,12 @@ async def test_pool_exports_query_only_case_from_reused_workspace(tmp_path): @pytest.mark.asyncio -async def test_pool_keeps_build_log_on_failure(tmp_path): +@pytest.mark.parametrize("raises", [False, True], ids=["failed-result", "exception"]) +async def test_pool_keeps_build_log_on_failure(tmp_path, raises): class FailedBuildCase(FakeCase): async def run_build(self, jobs): - self.events.append((self.case_id, "build_failed")) + if raises: + raise TimeoutError("build exceeded command timeout") return {"success": False, "jobs": [{"error": "failed"}]} class FailedBuildFactory(FakeFactory): @@ -148,36 +154,17 @@ async def test_pool_keeps_build_log_on_failure(tmp_path): result = (await BenchmarkDockerPool(FailedBuildFactory(events), concurrency=1).run([plan]))[0] assert not result.success - assert result.build == {"success": False, "jobs": [{"error": "failed"}]} + if raises: + assert result.error == "TimeoutError: build exceeded command timeout" + assert result.build is None + else: + assert result.error == "build failed" + assert result.build == {"success": False, "jobs": [{"error": "failed"}]} assert result.build_log == "build failed" assert log_path.read_text(encoding="utf-8") == "build failed" assert events.index(("failed", "read_build_log")) < events.index(("pool-0", "close")) -@pytest.mark.asyncio -async def test_pool_keeps_build_log_when_build_raises(tmp_path): - class TimedOutBuildCase(FakeCase): - async def run_build(self, jobs): - self.events.append((self.case_id, "build_timed_out")) - raise TimeoutError("build exceeded command timeout") - - class TimedOutBuildFactory(FakeFactory): - async def create_cases(self, case_ids, *, concurrency): - return [TimedOutBuildCase(name, self.events) for name in case_ids] - - events: list[tuple] = [] - log_path = tmp_path / "build.log" - plan = CasePlan("timed-out", build_jobs=(JobRequest("build"),), export_build_log_to=log_path) - - result = (await BenchmarkDockerPool(TimedOutBuildFactory(events), concurrency=1).run([plan]))[0] - - assert result.error == "TimeoutError: build exceeded command timeout" - assert result.build is None - assert result.build_log == "build failed" - assert log_path.read_text(encoding="utf-8") == "build failed" - assert events.index(("timed-out", "read_build_log")) < events.index(("pool-0", "close")) - - @pytest.mark.asyncio async def test_pool_cancels_case_before_closing_container(): started = asyncio.Event() @@ -215,23 +202,109 @@ def test_source_snapshot_excludes_unrelated_and_secret_files(tmp_path): assert snapshot.file_count == 1 -def test_worker_writes_log_to_requested_artifact_path(tmp_path): - """The current ReMe logger accepts a directory, so bind the exact log file explicitly.""" - from reme.utils import get_logger +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_phase", ["config", "application"]) +async def test_build_log_captures_initialization_failure(tmp_path, monkeypatch, failure_phase): + if failure_phase == "config": - path = tmp_path / "build_log" / "build.log" - logger = _start_case_log(path) - try: - logger.info("benchmark marker") - finally: - get_logger(log_to_console=False, log_to_file=False, force_init=True) - assert "benchmark marker" in path.read_text(encoding="utf-8") + def fail_config(_request): + raise ValueError("invalid benchmark config") + + monkeypatch.setattr(sandbox_worker, "_prepare_runtime", fail_config) + expected = "ValueError: invalid benchmark config" + else: + config = { + "workspace_dir": str(tmp_path / "workspace"), + "service": {"backend": "cli"}, + "jobs": {"broken": {"backend": "missing"}}, + "enable_logo": False, + "log_to_console": False, + "log_to_file": False, + } + monkeypatch.setattr(sandbox_worker, "_prepare_runtime", lambda _request: (tmp_path, config)) + expected = "ValueError: Unregistered backend 'missing' for Job 'broken'" + + with pytest.raises(ValueError, match="invalid benchmark config|Unregistered backend"): + await sandbox_worker._run_build({"case_root": str(tmp_path), "jobs": [{"job": "build"}]}) + + log_text = (tmp_path / "build_log" / "build.log").read_text(encoding="utf-8") + assert "Traceback (most recent call last):" in log_text + assert expected in log_text + + +@pytest.mark.asyncio +async def test_build_log_remains_bound_after_application_initializes_logger(tmp_path, monkeypatch): + import reme.reme as reme_module + + monkeypatch.setattr(sandbox_worker, "_prepare_runtime", lambda _request: (tmp_path, {})) + + 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 + + async def close(self): + pass + + async def run_job(_app, _job, _arguments): + return {"success": True} + + monkeypatch.setattr(reme_module, "ReMe", FakeReMe) + monkeypatch.setattr(sandbox_worker, "_run_job_on_app", run_job) + + result = await sandbox_worker._run_build({"case_root": str(tmp_path), "jobs": [{"job": "build"}]}) + + assert result["success"] + 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 + + +@pytest.mark.asyncio +async def test_worker_checkpoint_commits_only_memory_paths(tmp_path, monkeypatch): + workspace = tmp_path / "workspace" + workspace.mkdir() + subprocess.run(["git", "init", "--quiet", str(workspace)], check=True) + (workspace / "daily").mkdir() + (workspace / "daily" / "note.md").write_text("memory", encoding="utf-8") + (workspace / "digest").mkdir() + (workspace / "digest" / "note.md").write_text("digest", encoding="utf-8") + (workspace / "sessions").mkdir() + (workspace / "sessions" / "note.md").write_text("session", encoding="utf-8") + (workspace / "other.txt").write_text("unrelated", encoding="utf-8") + config = {"workspace_dir": str(workspace), "daily_dir": "daily", "digest_dir": "digest", "session_dir": "sessions"} + monkeypatch.setattr(sandbox_worker, "_prepare_runtime", lambda _request: (tmp_path, config)) + + result = await sandbox_worker._run_checkpoint({"message": "memory boundary"}) + + assert result["success"] is True + assert result["memory_checkpoint"]["message"] == "memory boundary" + committed = subprocess.run( + ["git", "show", "--pretty=format:", "--name-only", "HEAD"], + cwd=workspace, + capture_output=True, + text=True, + check=True, + ) + assert set(committed.stdout.splitlines()) == {"daily/note.md", "digest/note.md", "sessions/note.md"} + assert result["memory_checkpoint"]["paths"] == ["daily", "digest", "sessions"] + second = await sandbox_worker._run_checkpoint({"message": "empty boundary"}) + assert second["memory_checkpoint"]["commit_sha"] != result["memory_checkpoint"]["commit_sha"] + + config["daily_dir"] = "../outside" + with pytest.raises(ValueError, match="escapes the runtime workspace"): + await sandbox_worker._run_checkpoint({"message": "unsafe boundary"}) class FakeBackend: def __init__(self) -> None: self.files: dict[str, bytes] = {} self.commands: list[list[str]] = [] + self.worker_responses: dict[str, dict] = {} async def write_file(self, path: str, data: bytes) -> None: self.files[path] = data @@ -243,11 +316,11 @@ class FakeBackend: raise FileNotFoundError(path) from exc async def exec_shell(self, command: list[str], *, cwd=None, timeout=None): - import json from types import SimpleNamespace self.commands.append(command) output = b"" + exit_code = 0 if len(command) > 2 and command[1] == "-c": output = json.dumps( { @@ -258,11 +331,14 @@ class FakeBackend: ).encode() if "/workspace/harness/worker.py" in command: request = json.loads(self.files[command[command.index("--request") + 1]]) - response = {"success": True, "jobs": []} if request.get("mode") == "build" else {"success": True} + mode = request.get("mode", "job") + default = {"success": True, "jobs": []} if mode == "build" else {"success": True} + response = self.worker_responses.get(mode, default) self.files[command[command.index("--response") + 1]] = json.dumps(response).encode() + exit_code = 0 if response["success"] else 1 if command[:2] == ["tar", "-czf"]: self.files[command[2]] = b"workspace archive" - return SimpleNamespace(exit_code=0, stdout=output, stderr=b"", ok=lambda: True) + return SimpleNamespace(exit_code=exit_code, stdout=output, stderr=b"", ok=lambda: exit_code == 0) class TimedOutBackend(FakeBackend): @@ -291,16 +367,20 @@ class FakeWorkspace: self.closed = True -@pytest.mark.asyncio -async def test_backend_timeout_prevents_direct_case_reset(tmp_path): +def _fake_docker_factory(tmp_path, workspace_builder, *, install_projects=(".",)): (tmp_path / "reme").mkdir() (tmp_path / "reme" / "__init__.py").write_text("ok", encoding="utf-8") - workspace = FakeWorkspace(TimedOutBackend()) - factory = DockerReMeSandboxFactory( - SourceCandidate(SourceSnapshot.from_directory(tmp_path)), + return DockerReMeSandboxFactory( + SourceCandidate(SourceSnapshot.from_directory(tmp_path), install_projects=install_projects), plugins=("lme",), - workspace_builder=lambda _id, _image, _env: workspace, + workspace_builder=workspace_builder, ) + + +@pytest.mark.asyncio +async def test_backend_timeout_prevents_direct_case_reset(tmp_path): + workspace = FakeWorkspace(TimedOutBackend()) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: workspace) case = await factory.create_case("timed-out") try: with pytest.raises(SandboxCommandTimeoutError): @@ -317,14 +397,12 @@ async def test_backend_timeout_prevents_direct_case_reset(tmp_path): @pytest.mark.asyncio async def test_source_candidate_installs_plugin_and_closes_container(tmp_path): - (tmp_path / "reme").mkdir() - (tmp_path / "reme" / "__init__.py").write_text("ok", encoding="utf-8") backend = FakeBackend() workspace = FakeWorkspace(backend) - factory = DockerReMeSandboxFactory( - SourceCandidate(SourceSnapshot.from_directory(tmp_path), install_projects=(".", "plugins/lme")), - plugins=("lme",), - workspace_builder=lambda _id, _image, _env: workspace, + factory = _fake_docker_factory( + tmp_path, + lambda _id, _image, _env: workspace, + install_projects=(".", "plugins/lme"), ) case = await factory.create_case("example") try: @@ -343,18 +421,43 @@ async def test_source_candidate_installs_plugin_and_closes_container(tmp_path): assert workspace.closed +@pytest.mark.asyncio +async def test_job_and_checkpoint_share_worker_request_path(tmp_path): + backend = FakeBackend() + workspace = FakeWorkspace(backend) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: workspace) + case = await factory.create_case("example") + try: + result = await case.run_job("answer", {"query": "hello"}) + assert result.success + assert case.actions[-1].name == "job:answer" + job_request = json.loads(backend.files[next(path for path in backend.files if "/inbox/" in path)]) + assert job_request["mode"] == "job" + assert job_request["arguments"] == {"query": "hello"} + + await case.commit_memory_history("memory boundary") + assert case.actions[-1].name == "worker:checkpoint" + requests = [json.loads(content) for path, content in backend.files.items() if "/inbox/" in path] + assert requests[-1]["mode"] == "checkpoint" + assert requests[-1]["message"] == "memory boundary" + assert not any(command[0] == "git" and "commit" in command for command in backend.commands) + with pytest.raises(ValueError, match="non-empty"): + await case.commit_memory_history(" ") + + backend.worker_responses["checkpoint"] = {"success": False, "error": "git commit failed"} + with pytest.raises(SandboxCommandError, match="git commit failed"): + await case.commit_memory_history("next boundary") + assert case.actions[-1].exit_code == 1 + finally: + await case.close() + + @pytest.mark.asyncio async def test_pool_replaces_worker_after_backend_timeout(tmp_path): - (tmp_path / "reme").mkdir() - (tmp_path / "reme" / "__init__.py").write_text("ok", encoding="utf-8") timed_out_workspace = FakeWorkspace(TimedOutBackend()) replacement_workspace = FakeWorkspace(FakeBackend()) workspaces = iter((timed_out_workspace, replacement_workspace)) - factory = DockerReMeSandboxFactory( - SourceCandidate(SourceSnapshot.from_directory(tmp_path)), - plugins=("lme",), - workspace_builder=lambda _id, _image, _env: next(workspaces), - ) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: next(workspaces)) plans = [ CasePlan("timed-out", build_jobs=(JobRequest("build"),), export_build_log_to=tmp_path / "build.log"), CasePlan("next", build_jobs=(JobRequest("build"),)), @@ -372,15 +475,9 @@ async def test_pool_replaces_worker_after_backend_timeout(tmp_path): @pytest.mark.asyncio async def test_query_only_evaluation_export_omits_missing_build_log(tmp_path): - (tmp_path / "reme").mkdir() - (tmp_path / "reme" / "__init__.py").write_text("ok", encoding="utf-8") backend = FakeBackend() workspace = FakeWorkspace(backend) - factory = DockerReMeSandboxFactory( - SourceCandidate(SourceSnapshot.from_directory(tmp_path)), - plugins=("lme",), - workspace_builder=lambda _id, _image, _env: workspace, - ) + 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"{}" @@ -395,10 +492,52 @@ async def test_query_only_evaluation_export_omits_missing_build_log(tmp_path): await case.close() +@pytest.mark.asyncio +async def test_default_analysis_export_saves_logs_outside_archive(tmp_path): + class AnalysisBackend(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"] and command[-1] == "queries": + stream = io.BytesIO() + with tarfile.open(fileobj=stream, mode="w:gz") as archive: + for name, content in ( + ("queries/summary.json", b"{}"), + ("queries/q1/answer.log", b"answer log"), + ): + info = tarfile.TarInfo(name) + info.size = len(content) + archive.addfile(info, io.BytesIO(content)) + self.files[command[2]] = stream.getvalue() + return result + + backend = AnalysisBackend() + workspace = FakeWorkspace(backend) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: workspace) + case = await factory.create_case("example") + try: + backend.files["/workspace/case/runtime-layout.json"] = ( + b'{"workspace_root":"reme_workspace","analysis_excludes":[]}' + ) + backend.files["/workspace/case/build_log/build.log"] = b"build log" + backend.files["/workspace/case/queries/summary.json"] = b"{}" + archive = tmp_path / "analysis.tar.gz" + assert await case.export(archive) == archive + assert archive.read_bytes() == b"workspace archive" + artifacts = tmp_path / "analysis.tar.gz.artifacts" + assert (artifacts / "build_log" / "build.log").read_bytes() == b"build log" + assert (artifacts / "queries" / "summary.json").read_bytes() == b"{}" + assert (artifacts / "queries" / "q1" / "answer.log").read_bytes() == b"answer log" + archive_command = next( + command for command in backend.commands if command[:3] == ["tar", "-czf", "/tmp/reme-case-export.tar.gz"] + ) + assert "build_log" not in archive_command + assert "queries" not in archive_command + finally: + await case.close() + + @pytest.mark.asyncio async def test_cancelled_case_initialization_closes_workspace(tmp_path): - (tmp_path / "reme").mkdir() - (tmp_path / "reme" / "__init__.py").write_text("ok", encoding="utf-8") started = asyncio.Event() class BlockingWorkspace(FakeWorkspace): @@ -407,10 +546,7 @@ async def test_cancelled_case_initialization_closes_workspace(tmp_path): await asyncio.Event().wait() workspace = BlockingWorkspace(FakeBackend()) - factory = DockerReMeSandboxFactory( - SourceCandidate(SourceSnapshot.from_directory(tmp_path)), - workspace_builder=lambda _id, _image, _env: workspace, - ) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: workspace) task = asyncio.create_task(factory.create_case("cancelled")) await asyncio.wait_for(started.wait(), timeout=2) task.cancel() @@ -421,8 +557,6 @@ async def test_cancelled_case_initialization_closes_workspace(tmp_path): @pytest.mark.asyncio async def test_cancelled_batch_creation_closes_completed_workspace(tmp_path): - (tmp_path / "reme").mkdir() - (tmp_path / "reme" / "__init__.py").write_text("ok", encoding="utf-8") second_started = asyncio.Event() class BlockingWorkspace(FakeWorkspace): @@ -433,11 +567,7 @@ async def test_cancelled_batch_creation_closes_completed_workspace(tmp_path): first = FakeWorkspace(FakeBackend()) second = BlockingWorkspace(FakeBackend()) workspaces = iter((first, second)) - factory = DockerReMeSandboxFactory( - SourceCandidate(SourceSnapshot.from_directory(tmp_path)), - plugins=("lme",), - workspace_builder=lambda _id, _image, _env: next(workspaces), - ) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: next(workspaces)) task = asyncio.create_task(factory.create_cases(["first", "second"], concurrency=2)) await asyncio.wait_for(second_started.wait(), timeout=2) await asyncio.sleep(0) 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 04/14] 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) From 821163c542f8e93cc7415f2d36785a1da338706f 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:32:40 +0800 Subject: [PATCH 05/14] Deduplicate benchmark sandbox cleanup and artifact checks --- deploy/benchmark_docker/docker_sandbox.py | 54 ++++++++--------------- tests/unit/test_benchmark_docker.py | 19 +++++--- 2 files changed, 31 insertions(+), 42 deletions(-) diff --git a/deploy/benchmark_docker/docker_sandbox.py b/deploy/benchmark_docker/docker_sandbox.py index fe3c4d86..73f5b634 100644 --- a/deploy/benchmark_docker/docker_sandbox.py +++ b/deploy/benchmark_docker/docker_sandbox.py @@ -289,6 +289,12 @@ class DockerReMeSandboxFactory: 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, + ) + tasks = [asyncio.create_task(create(case_id)) for case_id in case_ids] try: results = await asyncio.gather(*tasks, return_exceptions=True) @@ -296,17 +302,11 @@ class DockerReMeSandboxFactory: for task in tasks: task.cancel() results = await asyncio.gather(*tasks, return_exceptions=True) - await asyncio.gather( - *(result.close() for result in results if isinstance(result, DockerReMeSandbox)), - return_exceptions=True, - ) + await close_created(results) raise failures = [result for result in results if isinstance(result, BaseException)] if failures: - await asyncio.gather( - *(result.close() for result in results if isinstance(result, DockerReMeSandbox)), - return_exceptions=True, - ) + await close_created(results) raise failures[0] return [result for result in results if isinstance(result, DockerReMeSandbox)] @@ -747,24 +747,20 @@ class DockerReMeSandbox: return (await self._read_build_log()).decode("utf-8", errors="replace") async def _read_build_log(self) -> bytes: + return await self._require_artifact(f"{CASE_ROOT}/build_log/build.log", "build log") + + async def _require_artifact(self, path: str, description: str) -> bytes: + """Read a required sandbox artifact, preserving its export-specific error.""" assert self.backend is not None, "sandbox is not initialized" - source = f"{CASE_ROOT}/build_log/build.log" try: - return await self.backend.read_file(source) + return await self.backend.read_file(path) except FileNotFoundError as exc: - raise SandboxCommandError(f"build log is missing required artifact: {source}") from exc + 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.""" async with self._operation_lock: - assert self.backend is not None, "sandbox is not initialized" - required = f"{CASE_ROOT}/build_log/build.log" - try: - await self.backend.read_file(required) - except FileNotFoundError as exc: - raise SandboxCommandError( - f"memory construction export is missing required artifact: {required}", - ) from exc + await self._require_artifact(f"{CASE_ROOT}/build_log/build.log", "memory construction export") return await self._export_archive( destination, "export_memory_construction", @@ -775,12 +771,7 @@ class DockerReMeSandbox: async def export_queries(self, destination: str | Path) -> Path: """Download per-query logs, results, and their aggregate summary.""" async with self._operation_lock: - assert self.backend is not None, "sandbox is not initialized" - required = f"{CASE_ROOT}/queries/summary.json" - try: - await self.backend.read_file(required) - except FileNotFoundError as exc: - raise SandboxCommandError(f"query export is missing required artifact: {required}") from exc + await self._require_artifact(f"{CASE_ROOT}/queries/summary.json", "query export") return await self._export_archive(destination, "export_queries", CASE_ROOT, ["queries"]) async def export_query(self, query_id: str, destination: str | Path) -> Path: @@ -788,24 +779,15 @@ class DockerReMeSandbox: self._validate_query_artifact_id(query_id) async with self._operation_lock: - assert self.backend is not None, "sandbox is not initialized" query_root = f"{CASE_ROOT}/queries/{query_id}" - required = f"{query_root}/result.json" - try: - await self.backend.read_file(required) - except FileNotFoundError as exc: - raise SandboxCommandError(f"query export is missing required artifact: {required}") from exc + 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}"]) async def export_evaluation(self, destination: str | Path) -> Path: """Download workspace and query artifacts, plus the build log when present.""" async with self._operation_lock: assert self.backend is not None, "sandbox is not initialized" - summary = f"{CASE_ROOT}/queries/summary.json" - try: - await self.backend.read_file(summary) - except FileNotFoundError as exc: - raise SandboxCommandError(f"evaluation export is missing required artifact: {summary}") from exc + await self._require_artifact(f"{CASE_ROOT}/queries/summary.json", "evaluation export") try: await self.backend.read_file(f"{CASE_ROOT}/build_log/build.log") except FileNotFoundError: diff --git a/tests/unit/test_benchmark_docker.py b/tests/unit/test_benchmark_docker.py index e4bc7e7d..321e1d31 100644 --- a/tests/unit/test_benchmark_docker.py +++ b/tests/unit/test_benchmark_docker.py @@ -786,22 +786,29 @@ async def test_query_only_evaluation_export_omits_missing_build_log(tmp_path): @pytest.mark.asyncio @pytest.mark.parametrize( - ("method", "required", "members"), + ("method", "required", "description", "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"]), + ( + "export_memory_construction", + "build_log/build.log", + "memory construction export", + ["reme_workspace", "build_log"], + ), + ("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, members): +async def test_archive_exports_check_required_artifacts(tmp_path, method, required, description, 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"): + with pytest.raises(SandboxCommandError) as failure: await export(*args) + 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" From 461e4c6fb7e3d894f8232203e2c2bd69e1a9f16c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Wed, 14 Oct 2026 21:16:45 +0800 Subject: [PATCH 06/14] Report benchmark cleanup failures and preserve diagnostic artifacts --- deploy/benchmark_docker/README.md | 46 +- deploy/benchmark_docker/_cleanup.py | 52 ++ deploy/benchmark_docker/candidate.py | 29 +- .../direct_docker_workspace.py | 76 ++- deploy/benchmark_docker/docker_sandbox.py | 254 +++++--- deploy/benchmark_docker/pool.py | 81 ++- deploy/benchmark_docker/worker.py | 255 +++++--- tests/unit/test_benchmark_docker.py | 615 ++++++++++++++---- tests/unit/test_benchmark_docker_workspace.py | 124 ++++ 9 files changed, 1176 insertions(+), 356 deletions(-) create mode 100644 deploy/benchmark_docker/_cleanup.py create mode 100644 tests/unit/test_benchmark_docker_workspace.py 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 From 1d210d0f4068c7d69b96c245a8bedfbb4d16a29b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Wed, 14 Oct 2026 21:23:19 +0800 Subject: [PATCH 07/14] Handle optional Docker dependencies in benchmark unit tests --- .github/workflows/ci-python-tests.yml | 2 +- tests/unit/test_benchmark_docker.py | 3 +++ tests/unit/test_benchmark_docker_workspace.py | 6 +++--- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci-python-tests.yml b/.github/workflows/ci-python-tests.yml index 5eeb5d29..10f5a6aa 100644 --- a/.github/workflows/ci-python-tests.yml +++ b/.github/workflows/ci-python-tests.yml @@ -38,7 +38,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip setuptools wheel - pip install -e reme_studio -e ".[dev,core]" + pip install -e reme_studio -e ".[dev,core,benchmark-docker]" pip install --no-deps -e plugins/auto-fin pip install -e plugins/daily_paper pip install -e plugins/lme -e plugins/beam diff --git a/tests/unit/test_benchmark_docker.py b/tests/unit/test_benchmark_docker.py index 07740e31..0cdb4792 100644 --- a/tests/unit/test_benchmark_docker.py +++ b/tests/unit/test_benchmark_docker.py @@ -1299,6 +1299,7 @@ async def test_pool_retries_strict_delete_without_replacing_uncertain_container( from types import SimpleNamespace from unittest.mock import AsyncMock + pytest.importorskip("aiodocker") from deploy.benchmark_docker.direct_docker_workspace import DirectDockerWorkspace workspaces = [] @@ -1306,6 +1307,8 @@ async def test_pool_retries_strict_delete_without_replacing_uncertain_container( client = SimpleNamespace(close=AsyncMock()) class StrictWorkspace(DirectDockerWorkspace): + # Provisioning deliberately populates the inherited Docker handles. + # pylint: disable=attribute-defined-outside-init async def _provision_backend(self): self._container, self._client, self._backend = container, client, TimedOutBackend() diff --git a/tests/unit/test_benchmark_docker_workspace.py b/tests/unit/test_benchmark_docker_workspace.py index 5817474c..a30521bc 100644 --- a/tests/unit/test_benchmark_docker_workspace.py +++ b/tests/unit/test_benchmark_docker_workspace.py @@ -1,15 +1,15 @@ """Direct workspace lifecycle without a Docker daemon.""" -# pylint: disable=missing-function-docstring,protected-access +# pylint: disable=missing-function-docstring,protected-access,wrong-import-position 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 +DockerError = pytest.importorskip("aiodocker.exceptions").DockerError +from deploy.benchmark_docker.direct_docker_workspace import DirectDockerWorkspace # noqa: E402 def make_workspace(): From 5e2de0025877ca0e9ef6fe29878a5a1a58e38f28 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Wed, 14 Oct 2026 21:28:10 +0800 Subject: [PATCH 08/14] Keep benchmark Docker dependencies optional in CI --- .github/workflows/ci-python-tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci-python-tests.yml b/.github/workflows/ci-python-tests.yml index 10f5a6aa..5eeb5d29 100644 --- a/.github/workflows/ci-python-tests.yml +++ b/.github/workflows/ci-python-tests.yml @@ -38,7 +38,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip setuptools wheel - pip install -e reme_studio -e ".[dev,core,benchmark-docker]" + pip install -e reme_studio -e ".[dev,core]" pip install --no-deps -e plugins/auto-fin pip install -e plugins/daily_paper pip install -e plugins/lme -e plugins/beam From 6a75c5db2514f4165c91dc724aa5a371a58e0645 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Thu, 15 Oct 2026 16:53:31 +0800 Subject: [PATCH 09/14] Fix cancellation of thread-pool background jobs --- reme/components/job/background_job.py | 51 +++++++++++++++---- tests/unit/test_job.py | 70 +++++++++++++++++++++++++++ 2 files changed, 112 insertions(+), 9 deletions(-) diff --git a/reme/components/job/background_job.py b/reme/components/job/background_job.py index ec8406e6..d4a99c99 100644 --- a/reme/components/job/background_job.py +++ b/reme/components/job/background_job.py @@ -5,6 +5,7 @@ import contextlib import random import threading import time +from concurrent.futures import Future from .base_job import BaseJob from ..component_registry import R @@ -23,7 +24,9 @@ class BackgroundJob(BaseJob): supervisor cannot trigger a restart. On close, the stop_event is set and the task is given up to - ``close_timeout`` seconds to exit gracefully; after that it is cancelled. + ``close_timeout`` seconds to exit gracefully; after that cancellation is + requested on its event loop and close waits for the actual worker to exit. + Job bodies must cooperate with cancellation, including in thread-pool mode. If a single run survives at least ``attempt_reset_after`` seconds before crashing, the backoff attempt counter resets — so a long-stable job that eventually crashes restarts quickly rather than at the capped delay. @@ -50,20 +53,49 @@ class BackgroundJob(BaseJob): self.attempt_reset_after: float = attempt_reset_after self.use_thread_pool: bool = use_thread_pool self._stop_event: asyncio.Event | threading.Event | None = None - self._task: asyncio.Task | None = None + self._task: asyncio.Future | None = None + self._executor_future: Future | None = None + self._worker_lock = threading.Lock() + self._worker_task: asyncio.Task | None = None async def _start(self) -> None: await super()._start() if self.use_thread_pool and self.app_context.thread_pool: self._stop_event = threading.Event() - self._task = asyncio.get_event_loop().run_in_executor( - self.app_context.thread_pool, - lambda: asyncio.run(self._run_with_supervisor()), - ) + self._executor_future = self.app_context.thread_pool.submit(self._run_in_thread) + self._task = asyncio.wrap_future(self._executor_future) else: self._stop_event = asyncio.Event() self._task = asyncio.create_task(self._run_with_supervisor()) + def _run_in_thread(self) -> None: + """Keep executor completion tied to actual event-loop shutdown.""" + with contextlib.suppress(asyncio.CancelledError): + asyncio.run(self._run_thread_task()) + + async def _run_thread_task(self) -> None: + with self._worker_lock: + self._worker_task = asyncio.current_task() + try: + # close() may have set stop before the worker published its task. + if not self._stop_event.is_set(): + await self._run_with_supervisor() + finally: + # Keep cancellation from targeting a loop that has already closed. + with self._worker_lock: + self._worker_task = None + + def _cancel_task(self) -> None: + if self._executor_future is None: + self._task.cancel() + return + # Cancel queued work without waiting for another long-running job to free + # a worker. If already running, cancel the coroutine on its own loop. + self._executor_future.cancel() + with self._worker_lock: + if self._worker_task is not None: + self._worker_task.get_loop().call_soon_threadsafe(self._worker_task.cancel) + async def _close(self) -> None: if self._stop_event is not None: self._stop_event.set() @@ -71,7 +103,7 @@ class BackgroundJob(BaseJob): await super()._close() async def _shutdown_task(self) -> None: - """Wait close_timeout for graceful exit, then force-cancel.""" + """Wait for graceful exit, then request cancellation and await completion.""" if self._task is None: return try: @@ -79,12 +111,13 @@ class BackgroundJob(BaseJob): # so a timeout here truly times out instead of cancelling silently. await asyncio.wait_for(asyncio.shield(self._task), timeout=self.close_timeout) except asyncio.TimeoutError: - self._task.cancel() + self._cancel_task() with contextlib.suppress(BaseException): - await self._task + await asyncio.shield(self._task) except Exception: self.logger.exception(f"Background task '{self.name}' raised during close") self._task = None + self._executor_future = None def _backoff_delay(self, attempt: int) -> float: """Exponential backoff with ±50% jitter, capped at backoff_cap.""" diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 191a27d9..1d68adbf 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -3,6 +3,8 @@ # pylint: disable=protected-access,missing-function-docstring,missing-class-docstring,no-self-argument,unused-argument import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -337,6 +339,74 @@ def test_shutdown_task_cancels_on_timeout(): asyncio.run(run()) +def test_thread_pool_shutdown_waits_for_worker_cleanup(): + async def run(): + started = threading.Event() + events = [] + + class WaitingJob(BackgroundJob): + async def __call__(self, **kwargs): + started.set() + try: + # Bounded fallback ensures a regression cannot hang pytest. + await asyncio.sleep(2) + except asyncio.CancelledError: + events.append("cancelled") + raise + finally: + await asyncio.sleep(0.02) + events.append("cleaned") + + with ThreadPoolExecutor(max_workers=1) as pool: + job = WaitingJob(app_context=SimpleNamespace(thread_pool=pool), use_thread_pool=True, close_timeout=0.01) + await job.start() + assert await asyncio.to_thread(started.wait, 1) + executor_future = job._executor_future + await asyncio.wait_for(job.close(), 1) + assert events == ["cancelled", "cleaned"] + assert executor_future.done() + assert job._task is None + + asyncio.run(run()) + + +@pytest.mark.parametrize("queued", [True, False], ids=["queued", "starting"]) +def test_thread_pool_shutdown_before_coroutine_starts(queued): + async def run(): + entered = threading.Event() + release = threading.Event() + + class DelayedJob(BackgroundJob): + def _run_in_thread(self): + entered.set() + release.wait(2) + super()._run_in_thread() + + def _cancel_task(self): + super()._cancel_task() + release.set() + + with ThreadPoolExecutor(max_workers=1) as pool: + if queued: + pool.submit(release.wait, 2) + job = DelayedJob(app_context=SimpleNamespace(thread_pool=pool), use_thread_pool=True, close_timeout=0.01) + body = AsyncMock() + job._run_with_supervisor = body + try: + await job.start() + if not queued: + assert await asyncio.to_thread(entered.wait, 1) + executor_future = job._executor_future + await asyncio.wait_for(job.close(), 1) + assert executor_future.done() + assert executor_future.cancelled() == queued + body.assert_not_awaited() + finally: + release.set() + + asyncio.run(run()) + + def test_cron_uses_configured_timezone(monkeypatch): from zoneinfo import ZoneInfo From 7241e1053e91ee71793e398eb1f2cbe5f2d80728 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Thu, 15 Oct 2026 17:12:36 +0800 Subject: [PATCH 10/14] Finish background worker cleanup before propagating cancellation --- reme/components/job/background_job.py | 43 ++++++++--- tests/unit/test_job.py | 103 ++++++++++++++++++++++++++ 2 files changed, 134 insertions(+), 12 deletions(-) diff --git a/reme/components/job/background_job.py b/reme/components/job/background_job.py index d4a99c99..5304c874 100644 --- a/reme/components/job/background_job.py +++ b/reme/components/job/background_job.py @@ -26,6 +26,7 @@ class BackgroundJob(BaseJob): On close, the stop_event is set and the task is given up to ``close_timeout`` seconds to exit gracefully; after that cancellation is requested on its event loop and close waits for the actual worker to exit. + Cancellation of close is propagated only after worker cleanup finishes. Job bodies must cooperate with cancellation, including in thread-pool mode. If a single run survives at least ``attempt_reset_after`` seconds before crashing, the backoff attempt counter resets — so a long-stable job that @@ -99,25 +100,43 @@ class BackgroundJob(BaseJob): async def _close(self) -> None: if self._stop_event is not None: self._stop_event.set() - await self._shutdown_task() - await super()._close() + try: + await self._shutdown_task() + finally: + await super()._close() async def _shutdown_task(self) -> None: """Wait for graceful exit, then request cancellation and await completion.""" if self._task is None: return - try: - # shield prevents wait_for's cancellation from propagating to the task itself, - # so a timeout here truly times out instead of cancelling silently. - await asyncio.wait_for(asyncio.shield(self._task), timeout=self.close_timeout) - except asyncio.TimeoutError: - self._cancel_task() - with contextlib.suppress(BaseException): - await asyncio.shield(self._task) - except Exception: - self.logger.exception(f"Background task '{self.name}' raised during close") + + async def wait_for_exit() -> None: + try: + try: + # A timeout must not cancel the executor's completion future. + await asyncio.wait_for(asyncio.shield(self._task), timeout=self.close_timeout) + except asyncio.TimeoutError: + self._cancel_task() + await self._task + except asyncio.CancelledError: + pass # The worker or queued executor future was cancelled. + except Exception: + self.logger.exception(f"Background task '{self.name}' raised during close") + + # Shield the whole shutdown so repeated caller cancellations cannot + # interrupt worker cleanup or cancel the executor's completion future. + shutdown = asyncio.create_task(wait_for_exit()) + cancelled = None + while not shutdown.done(): + try: + await asyncio.shield(shutdown) + except asyncio.CancelledError as exc: + cancelled = exc + shutdown.result() self._task = None self._executor_future = None + if cancelled is not None: + raise cancelled def _backoff_delay(self, attempt: int) -> float: """Exponential backoff with ±50% jitter, capped at backoff_cap.""" diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 1d68adbf..4fad942c 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -407,6 +407,109 @@ def test_thread_pool_shutdown_before_coroutine_starts(queued): asyncio.run(run()) +@pytest.mark.parametrize("use_thread_pool", [False, True], ids=["async", "thread"]) +@pytest.mark.parametrize("phase", ["graceful", "forced"]) +def test_close_finishes_worker_cleanup_before_propagating_repeated_cancellation(use_thread_pool, phase): + async def run(): + started, cleaning, release, finished = (threading.Event() for _ in range(4)) + + class WaitingJob(BackgroundJob): + async def __call__(self, **kwargs): + started.set() + try: + await asyncio.sleep(1) + finally: + cleaning.set() + deadline = asyncio.get_running_loop().time() + 1 + while not release.is_set() and asyncio.get_running_loop().time() < deadline: + await asyncio.sleep(0.001) + finished.set() + + with ThreadPoolExecutor(max_workers=1) as pool: + job = WaitingJob( + app_context=SimpleNamespace(thread_pool=pool), + use_thread_pool=use_thread_pool, + close_timeout=0.1 if phase == "graceful" else 0.01, + ) + await job.start() + worker = job._task + executor_future = job._executor_future + closing = None + try: + assert await asyncio.to_thread(started.wait, 1) + closing = asyncio.create_task(job.close()) + if phase == "forced": + assert await asyncio.to_thread(cleaning.wait, 1) + else: + await asyncio.sleep(0) + for _ in range(3): + closing.cancel() + await asyncio.sleep(0) + assert not closing.done() + assert not worker.done() + assert job.is_started + assert await asyncio.to_thread(cleaning.wait, 1) + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(closing, 1) + assert finished.is_set() + assert worker.done() + assert executor_future is None or executor_future.done() + assert job._task is None + assert job._executor_future is None + assert not job.is_started + finally: + release.set() + if closing is not None: + await asyncio.gather(closing, return_exceptions=True) + await asyncio.gather(worker, return_exceptions=True) + await job.close() + + asyncio.run(run()) + + +@pytest.mark.parametrize("use_thread_pool", [False, True], ids=["async", "thread"]) +def test_close_logs_worker_cleanup_failure_after_timeout(use_thread_pool): + async def run(): + started = threading.Event() + errors = [] + + class FailingCleanupJob(BackgroundJob): + async def __call__(self, **kwargs): + started.set() + try: + await asyncio.sleep(1) + finally: + raise RuntimeError("worker cleanup failed") + + def capture_error(_message): + import sys + + errors.append(sys.exception()) + + with ThreadPoolExecutor(max_workers=1) as pool: + job = FailingCleanupJob( + app_context=SimpleNamespace(thread_pool=pool), + use_thread_pool=use_thread_pool, + supervisor=False, + close_timeout=0.01, + ) + job.logger = MagicMock() + job.logger.exception.side_effect = capture_error + await job.start() + try: + assert await asyncio.to_thread(started.wait, 1) + finally: + await job.close() + assert len(errors) == 1 + assert isinstance(errors[0], RuntimeError) + assert str(errors[0]) == "worker cleanup failed" + assert job._task is None + assert job._executor_future is None + + asyncio.run(run()) + + def test_cron_uses_configured_timezone(monkeypatch): from zoneinfo import ZoneInfo From 4a2cf7d19c37ac7a06bc87bf51c87290048794d3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Thu, 15 Oct 2026 17:23:05 +0800 Subject: [PATCH 11/14] Revert background worker cancellation changes --- reme/components/job/background_job.py | 90 +++----------- tests/unit/test_job.py | 173 -------------------------- 2 files changed, 19 insertions(+), 244 deletions(-) diff --git a/reme/components/job/background_job.py b/reme/components/job/background_job.py index 5304c874..ec8406e6 100644 --- a/reme/components/job/background_job.py +++ b/reme/components/job/background_job.py @@ -5,7 +5,6 @@ import contextlib import random import threading import time -from concurrent.futures import Future from .base_job import BaseJob from ..component_registry import R @@ -24,10 +23,7 @@ class BackgroundJob(BaseJob): supervisor cannot trigger a restart. On close, the stop_event is set and the task is given up to - ``close_timeout`` seconds to exit gracefully; after that cancellation is - requested on its event loop and close waits for the actual worker to exit. - Cancellation of close is propagated only after worker cleanup finishes. - Job bodies must cooperate with cancellation, including in thread-pool mode. + ``close_timeout`` seconds to exit gracefully; after that it is cancelled. If a single run survives at least ``attempt_reset_after`` seconds before crashing, the backoff attempt counter resets — so a long-stable job that eventually crashes restarts quickly rather than at the capped delay. @@ -54,89 +50,41 @@ class BackgroundJob(BaseJob): self.attempt_reset_after: float = attempt_reset_after self.use_thread_pool: bool = use_thread_pool self._stop_event: asyncio.Event | threading.Event | None = None - self._task: asyncio.Future | None = None - self._executor_future: Future | None = None - self._worker_lock = threading.Lock() - self._worker_task: asyncio.Task | None = None + self._task: asyncio.Task | None = None async def _start(self) -> None: await super()._start() if self.use_thread_pool and self.app_context.thread_pool: self._stop_event = threading.Event() - self._executor_future = self.app_context.thread_pool.submit(self._run_in_thread) - self._task = asyncio.wrap_future(self._executor_future) + self._task = asyncio.get_event_loop().run_in_executor( + self.app_context.thread_pool, + lambda: asyncio.run(self._run_with_supervisor()), + ) else: self._stop_event = asyncio.Event() self._task = asyncio.create_task(self._run_with_supervisor()) - def _run_in_thread(self) -> None: - """Keep executor completion tied to actual event-loop shutdown.""" - with contextlib.suppress(asyncio.CancelledError): - asyncio.run(self._run_thread_task()) - - async def _run_thread_task(self) -> None: - with self._worker_lock: - self._worker_task = asyncio.current_task() - try: - # close() may have set stop before the worker published its task. - if not self._stop_event.is_set(): - await self._run_with_supervisor() - finally: - # Keep cancellation from targeting a loop that has already closed. - with self._worker_lock: - self._worker_task = None - - def _cancel_task(self) -> None: - if self._executor_future is None: - self._task.cancel() - return - # Cancel queued work without waiting for another long-running job to free - # a worker. If already running, cancel the coroutine on its own loop. - self._executor_future.cancel() - with self._worker_lock: - if self._worker_task is not None: - self._worker_task.get_loop().call_soon_threadsafe(self._worker_task.cancel) - async def _close(self) -> None: if self._stop_event is not None: self._stop_event.set() - try: - await self._shutdown_task() - finally: - await super()._close() + await self._shutdown_task() + await super()._close() async def _shutdown_task(self) -> None: - """Wait for graceful exit, then request cancellation and await completion.""" + """Wait close_timeout for graceful exit, then force-cancel.""" if self._task is None: return - - async def wait_for_exit() -> None: - try: - try: - # A timeout must not cancel the executor's completion future. - await asyncio.wait_for(asyncio.shield(self._task), timeout=self.close_timeout) - except asyncio.TimeoutError: - self._cancel_task() - await self._task - except asyncio.CancelledError: - pass # The worker or queued executor future was cancelled. - except Exception: - self.logger.exception(f"Background task '{self.name}' raised during close") - - # Shield the whole shutdown so repeated caller cancellations cannot - # interrupt worker cleanup or cancel the executor's completion future. - shutdown = asyncio.create_task(wait_for_exit()) - cancelled = None - while not shutdown.done(): - try: - await asyncio.shield(shutdown) - except asyncio.CancelledError as exc: - cancelled = exc - shutdown.result() + try: + # shield prevents wait_for's cancellation from propagating to the task itself, + # so a timeout here truly times out instead of cancelling silently. + await asyncio.wait_for(asyncio.shield(self._task), timeout=self.close_timeout) + except asyncio.TimeoutError: + self._task.cancel() + with contextlib.suppress(BaseException): + await self._task + except Exception: + self.logger.exception(f"Background task '{self.name}' raised during close") self._task = None - self._executor_future = None - if cancelled is not None: - raise cancelled def _backoff_delay(self, attempt: int) -> float: """Exponential backoff with ±50% jitter, capped at backoff_cap.""" diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 4fad942c..191a27d9 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -3,8 +3,6 @@ # pylint: disable=protected-access,missing-function-docstring,missing-class-docstring,no-self-argument,unused-argument import asyncio -import threading -from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -339,177 +337,6 @@ def test_shutdown_task_cancels_on_timeout(): asyncio.run(run()) -def test_thread_pool_shutdown_waits_for_worker_cleanup(): - async def run(): - started = threading.Event() - events = [] - - class WaitingJob(BackgroundJob): - async def __call__(self, **kwargs): - started.set() - try: - # Bounded fallback ensures a regression cannot hang pytest. - await asyncio.sleep(2) - except asyncio.CancelledError: - events.append("cancelled") - raise - finally: - await asyncio.sleep(0.02) - events.append("cleaned") - - with ThreadPoolExecutor(max_workers=1) as pool: - job = WaitingJob(app_context=SimpleNamespace(thread_pool=pool), use_thread_pool=True, close_timeout=0.01) - await job.start() - assert await asyncio.to_thread(started.wait, 1) - executor_future = job._executor_future - await asyncio.wait_for(job.close(), 1) - assert events == ["cancelled", "cleaned"] - assert executor_future.done() - assert job._task is None - - asyncio.run(run()) - - -@pytest.mark.parametrize("queued", [True, False], ids=["queued", "starting"]) -def test_thread_pool_shutdown_before_coroutine_starts(queued): - async def run(): - entered = threading.Event() - release = threading.Event() - - class DelayedJob(BackgroundJob): - def _run_in_thread(self): - entered.set() - release.wait(2) - super()._run_in_thread() - - def _cancel_task(self): - super()._cancel_task() - release.set() - - with ThreadPoolExecutor(max_workers=1) as pool: - if queued: - pool.submit(release.wait, 2) - job = DelayedJob(app_context=SimpleNamespace(thread_pool=pool), use_thread_pool=True, close_timeout=0.01) - body = AsyncMock() - job._run_with_supervisor = body - try: - await job.start() - if not queued: - assert await asyncio.to_thread(entered.wait, 1) - executor_future = job._executor_future - await asyncio.wait_for(job.close(), 1) - assert executor_future.done() - assert executor_future.cancelled() == queued - body.assert_not_awaited() - finally: - release.set() - - asyncio.run(run()) - - -@pytest.mark.parametrize("use_thread_pool", [False, True], ids=["async", "thread"]) -@pytest.mark.parametrize("phase", ["graceful", "forced"]) -def test_close_finishes_worker_cleanup_before_propagating_repeated_cancellation(use_thread_pool, phase): - async def run(): - started, cleaning, release, finished = (threading.Event() for _ in range(4)) - - class WaitingJob(BackgroundJob): - async def __call__(self, **kwargs): - started.set() - try: - await asyncio.sleep(1) - finally: - cleaning.set() - deadline = asyncio.get_running_loop().time() + 1 - while not release.is_set() and asyncio.get_running_loop().time() < deadline: - await asyncio.sleep(0.001) - finished.set() - - with ThreadPoolExecutor(max_workers=1) as pool: - job = WaitingJob( - app_context=SimpleNamespace(thread_pool=pool), - use_thread_pool=use_thread_pool, - close_timeout=0.1 if phase == "graceful" else 0.01, - ) - await job.start() - worker = job._task - executor_future = job._executor_future - closing = None - try: - assert await asyncio.to_thread(started.wait, 1) - closing = asyncio.create_task(job.close()) - if phase == "forced": - assert await asyncio.to_thread(cleaning.wait, 1) - else: - await asyncio.sleep(0) - for _ in range(3): - closing.cancel() - await asyncio.sleep(0) - assert not closing.done() - assert not worker.done() - assert job.is_started - assert await asyncio.to_thread(cleaning.wait, 1) - release.set() - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(closing, 1) - assert finished.is_set() - assert worker.done() - assert executor_future is None or executor_future.done() - assert job._task is None - assert job._executor_future is None - assert not job.is_started - finally: - release.set() - if closing is not None: - await asyncio.gather(closing, return_exceptions=True) - await asyncio.gather(worker, return_exceptions=True) - await job.close() - - asyncio.run(run()) - - -@pytest.mark.parametrize("use_thread_pool", [False, True], ids=["async", "thread"]) -def test_close_logs_worker_cleanup_failure_after_timeout(use_thread_pool): - async def run(): - started = threading.Event() - errors = [] - - class FailingCleanupJob(BackgroundJob): - async def __call__(self, **kwargs): - started.set() - try: - await asyncio.sleep(1) - finally: - raise RuntimeError("worker cleanup failed") - - def capture_error(_message): - import sys - - errors.append(sys.exception()) - - with ThreadPoolExecutor(max_workers=1) as pool: - job = FailingCleanupJob( - app_context=SimpleNamespace(thread_pool=pool), - use_thread_pool=use_thread_pool, - supervisor=False, - close_timeout=0.01, - ) - job.logger = MagicMock() - job.logger.exception.side_effect = capture_error - await job.start() - try: - assert await asyncio.to_thread(started.wait, 1) - finally: - await job.close() - assert len(errors) == 1 - assert isinstance(errors[0], RuntimeError) - assert str(errors[0]) == "worker cleanup failed" - assert job._task is None - assert job._executor_future is None - - asyncio.run(run()) - - def test_cron_uses_configured_timezone(monkeypatch): from zoneinfo import ZoneInfo From 37343bc580d200cdf732691c73cc8915fea239fc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Thu, 15 Oct 2026 18:56:15 +0800 Subject: [PATCH 12/14] Fix benchmark worker plugin defaults before config validation --- deploy/benchmark_docker/worker.py | 6 ++++- tests/unit/test_benchmark_docker.py | 37 +++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/deploy/benchmark_docker/worker.py b/deploy/benchmark_docker/worker.py index 3abeeea0..c77ebaa8 100644 --- a/deploy/benchmark_docker/worker.py +++ b/deploy/benchmark_docker/worker.py @@ -102,6 +102,7 @@ def _prepare_runtime(request: dict[str, Any]) -> tuple[Path, dict[str, Any]]: # Import after setting TMPDIR so ReMe and its subprocesses cannot leave # case-specific temporary files outside the disposable case root. from reme.config import resolve_app_config + from reme.plugin import resolve_plugin_runtime from reme.schema import ApplicationConfig plugins = request.get("plugins") or [] @@ -115,7 +116,10 @@ def _prepare_runtime(request: dict[str, Any]) -> tuple[Path, dict[str, Any]]: log_to_file=True, ) app_config["environment"] = dict(os.environ) - resolved_config = ApplicationConfig.model_validate(app_config) + # Apply plugin defaults before validation materializes schema defaults, + # otherwise those defaults would override plugin paths and partial jobs. + runtime = resolve_plugin_runtime(app_config) + resolved_config = ApplicationConfig.model_validate(runtime.config) _write_runtime_layout(case_root, resolved_config) # The build checkpoint reads workspace paths from this mapping. Return the # validated model's dump so optional config fields (including ``daily_dir``) diff --git a/tests/unit/test_benchmark_docker.py b/tests/unit/test_benchmark_docker.py index 0cdb4792..e997076c 100644 --- a/tests/unit/test_benchmark_docker.py +++ b/tests/unit/test_benchmark_docker.py @@ -421,6 +421,43 @@ def test_source_snapshot_preserves_custom_path_rules_and_symlink_checks(tmp_path SourceSnapshot.from_directory(tmp_path, excludes=patterns) +def test_worker_merges_plugin_defaults_before_validating_config(tmp_path, monkeypatch): + from reme.plugin import Plugin, PluginManager + + defaults = { + "daily_dir": "notes", + "digest_dir": "summaries", + "jobs": {"plugin_version": {"backend": "base", "steps": [{"backend": "version_step"}]}}, + } + manager = PluginManager([Plugin(name="example", config=defaults)]) + monkeypatch.setattr(PluginManager, "discover", classmethod(lambda cls, specs: manager)) + # _prepare_runtime changes cwd and TMPDIR in the disposable worker process. + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("TMPDIR", str(tmp_path)) + config = tmp_path / "benchmark.yaml" + config.write_text( + "digest_dir: custom_digest\njobs:\n plugin_version:\n enable_serve: false\n", + encoding="utf-8", + ) + case_root, prepared = sandbox_worker._prepare_runtime( + { + "case_root": str(tmp_path), + "workspace_dir": str(tmp_path / "reme_workspace"), + "config": str(config), + "plugins": ["example"], + }, + ) + + assert prepared["daily_dir"] == "notes" + assert prepared["digest_dir"] == "custom_digest" + job = prepared["jobs"]["plugin_version"] + assert job["backend"] == "base" + assert job["steps"] == [{"backend": "version_step"}] + assert job["enable_serve"] is False + layout = json.loads((case_root / "runtime-layout.json").read_text(encoding="utf-8")) + assert layout["configured_paths"]["daily_dir"] == "reme_workspace/notes" + + @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): From 83265fc3b819fcb61da193080b1474f27bf4bc94 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Thu, 15 Oct 2026 19:02:26 +0800 Subject: [PATCH 13/14] Stop interrupted benchmark containers before collecting diagnostics --- deploy/benchmark_docker/README.md | 14 ++- .../direct_docker_workspace.py | 36 +++++- deploy/benchmark_docker/docker_sandbox.py | 67 +++++++++-- deploy/benchmark_docker/pool.py | 8 +- tests/integration/test_benchmark_docker.py | 39 +++++++ tests/unit/test_benchmark_docker.py | 106 +++++++++++++++++- tests/unit/test_benchmark_docker_workspace.py | 38 ++++++- 7 files changed, 281 insertions(+), 27 deletions(-) diff --git a/deploy/benchmark_docker/README.md b/deploy/benchmark_docker/README.md index 8b74bdf1..438eb0e7 100644 --- a/deploy/benchmark_docker/README.md +++ b/deploy/benchmark_docker/README.md @@ -93,9 +93,17 @@ wrapper but no manifest are rejected; to reuse one, extract it and supply its `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. +build, the pool also returns its text in `CaseResult.build_log`. `command_timeout` +(default 1800 seconds) bounds the entire shell command, including worker startup, +jobs, and Application shutdown. On timeout, cancellation, or a shell backend +failure, the sandbox force-stops the container without candidate or plugin +cooperation, attempts to cache the build log for up to 5 seconds, and removes the +container before propagating the error. Docker stop and removal each have a +30-second timeout; cleanup failures include the container ID in exception notes. +Only the cached build log remains available after removal. The pool creates a +replacement when more cases need a worker. +Custom `workspace_builder` implementations must provide an async `stop()` that +force-stops execution while retaining files for diagnostics, in addition to `close()`. 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 diff --git a/deploy/benchmark_docker/direct_docker_workspace.py b/deploy/benchmark_docker/direct_docker_workspace.py index 5a6f8827..ea62749c 100644 --- a/deploy/benchmark_docker/direct_docker_workspace.py +++ b/deploy/benchmark_docker/direct_docker_workspace.py @@ -15,6 +15,8 @@ from typing import Any from aiodocker.exceptions import DockerError from agentscope.workspace import DockerWorkspace +DOCKER_CLEANUP_TIMEOUT = 30.0 + class DirectDockerWorkspace(DockerWorkspace): """Use the supplied image directly and skip the unused MCP gateway.""" @@ -28,6 +30,7 @@ class DirectDockerWorkspace(DockerWorkspace): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self._lifecycle_lock = asyncio.Lock() + self._stopped = False async def _build_or_reuse_image(self) -> None: """Select the caller-provided image instead of deriving a gateway image.""" @@ -56,13 +59,36 @@ class DirectDockerWorkspace(DockerWorkspace): exc.add_note(f"workspace close failed: {type(close_exc).__name__}: {close_exc}") raise self.is_alive = True + self._stopped = False + + async def stop(self) -> None: + """Kill all container processes without asking the candidate to shut down.""" + async with self._lifecycle_lock: + self.is_alive = False + if self._container is None or self._stopped: + return + async with asyncio.timeout(DOCKER_CLEANUP_TIMEOUT): + try: + await self._container.kill(signal="SIGKILL") + await self._container.wait(condition="not-running") + except DockerError as exc: + if exc.status != 404: + # Docker returns 409 for an already stopped container. + if exc.status != 409 or (await self._container.show())["State"]["Running"]: + raise + self._stopped = 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: + if ( + not self._stopped + and 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], @@ -71,16 +97,16 @@ class DirectDockerWorkspace(DockerWorkspace): 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) + # Also stop processes if an earlier explicit stop failed, and + # retain handles unless Docker confirms removal. + await asyncio.wait_for(self._container.delete(force=True), timeout=DOCKER_CLEANUP_TIMEOUT) except DockerError as exc: if exc.status != 404: raise self._container = None self._backend = None if self._client is not None: - await self._client.close() + await asyncio.wait_for(self._client.close(), timeout=DOCKER_CLEANUP_TIMEOUT) self._client = None self._backend = None diff --git a/deploy/benchmark_docker/docker_sandbox.py b/deploy/benchmark_docker/docker_sandbox.py index 92a4b632..7b3c48ce 100644 --- a/deploy/benchmark_docker/docker_sandbox.py +++ b/deploy/benchmark_docker/docker_sandbox.py @@ -30,6 +30,8 @@ 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" +_CLEANUP_TIMEOUT = 30.0 +_LOG_TIMEOUT = 5.0 _PRIVATE_EVALUATION_FIELDS = frozenset( {"query", "question", "golden_answer", "question_type", "rubric", "judge_arguments"}, ) @@ -68,6 +70,10 @@ class Workspace(Protocol): """Close the workspace.""" raise NotImplementedError + async def stop(self) -> None: + """Force-stop container processes while keeping files available.""" + raise NotImplementedError + def get_backend(self) -> Backend: """Return the execution backend.""" raise NotImplementedError @@ -81,7 +87,7 @@ class SandboxCommandError(RuntimeError): class SandboxCommandTimeoutError(SandboxCommandError): - """A command timed out and may still be running inside the container.""" + """A command timed out; cleanup failures are attached as exception notes.""" def _validate_case_id(case_id: str) -> None: @@ -379,6 +385,7 @@ class DockerReMeSandbox: self._closed = False self._reusable = True self._operation_lock = asyncio.Lock() + self._build_log: bytes | None = None @property def reusable(self) -> bool: @@ -524,13 +531,15 @@ class DockerReMeSandbox: raise SandboxCommandError("sandbox container is no longer reusable") started = _utc_now() try: - result = await self.backend.exec_shell(command, cwd=cwd, timeout=self.command_timeout) + async with asyncio.timeout(self.command_timeout): + result = await self.backend.exec_shell(command, cwd=cwd, timeout=self.command_timeout) except TimeoutError as exc: - self._reusable = False - raise SandboxCommandTimeoutError(f"sandbox action {name!r} timed out; container must be destroyed") from exc - except (Exception, asyncio.CancelledError): + error = SandboxCommandTimeoutError(f"sandbox action {name!r} timed out") + await cleanup_on_error([(f"container {self.container_id!r}", self._abort())], error) + raise error from exc + except (Exception, asyncio.CancelledError) as exc: # A backend failure may occur after Docker started the command. - self._reusable = False + await cleanup_on_error([(f"container {self.container_id!r}", self._abort())], exc) raise self.actions.append( ActionRecord( @@ -543,10 +552,36 @@ class DockerReMeSandbox: ), ) if result.exit_code == -1: - self._reusable = False - raise SandboxCommandTimeoutError(f"sandbox action {name!r} timed out; container must be destroyed") + error = SandboxCommandTimeoutError(f"sandbox action {name!r} timed out") + await cleanup_on_error([(f"container {self.container_id!r}", self._abort())], error) + raise error return result + async def _abort(self) -> None: + """Retire an uncertain command while its caller holds the operation lock.""" + self._reusable = False + error = None + try: + await asyncio.wait_for(self.workspace.stop(), timeout=_CLEANUP_TIMEOUT) + # Read files only after the candidate has stopped. Diagnostics must + # never delay termination or prevent container removal. + try: + assert self.backend is not None + self._build_log = await asyncio.wait_for( + self.backend.read_file(f"{CASE_ROOT}/build_log/build.log"), + timeout=_LOG_TIMEOUT, + ) + except Exception: + pass + except BaseException as exc: + error = exc + raise + finally: + if error is None: + await self._close_workspace() + else: + await cleanup_on_error([("container removal", self._close_workspace())], error) + async def _exec_checked(self, name: str, command: list[str], *, cwd: str | None = None) -> Any: result = await self._exec(name, command, cwd=cwd) if not result.ok(): @@ -776,6 +811,10 @@ class DockerReMeSandbox: return (await self._read_build_log()).decode("utf-8", errors="replace") async def _read_build_log(self) -> bytes: + if self._build_log is not None: + return self._build_log + if not self.reusable: + raise SandboxCommandError("build log unavailable after interrupted command") return await self._require_artifact(f"{CASE_ROOT}/build_log/build.log", "build log") async def _require_artifact(self, path: str, description: str) -> bytes: @@ -1078,10 +1117,14 @@ class DockerReMeSandbox: async def close(self) -> None: """Destroy the case container. Export before calling this method.""" async with self._operation_lock: - if not self._closed: - self._reusable = False - await self.workspace.close() - self._closed = True + await self._close_workspace() + + async def _close_workspace(self) -> None: + """Close without acquiring the operation lock a second time.""" + if not self._closed: + self._reusable = False + await asyncio.wait_for(self.workspace.close(), timeout=_CLEANUP_TIMEOUT) + self._closed = True async def __aenter__(self) -> "DockerReMeSandbox": return self diff --git a/deploy/benchmark_docker/pool.py b/deploy/benchmark_docker/pool.py index 92b9bca0..740cc058 100644 --- a/deploy/benchmark_docker/pool.py +++ b/deploy/benchmark_docker/pool.py @@ -140,7 +140,7 @@ class BenchmarkDockerPool: try: result.build = await worker.run_build(list(plan.build_jobs)) except Exception as exc: - result.error = f"{type(exc).__name__}: {exc}" + result.error = format_error(exc) await self._retain_failed_build_log(worker, plan, result) return result if not result.build.get("success"): @@ -162,14 +162,14 @@ class BenchmarkDockerPool: await worker.export_evaluation(plan.export_evaluation_to) return result except Exception as exc: # Keep other cases running and release this worker. - result.error = f"{type(exc).__name__}: {exc}" + result.error = format_error(exc) return result finally: if worker.reusable: available.put_nowait(worker) else: - # A timed-out command may still be running in the container. - # Retire it before allowing another case into this slot. + # Confirm cleanup (or retry a failed removal) before + # allowing another case into this slot. closed = False try: await worker.close() diff --git a/tests/integration/test_benchmark_docker.py b/tests/integration/test_benchmark_docker.py index a22f7497..82a59fbd 100644 --- a/tests/integration/test_benchmark_docker.py +++ b/tests/integration/test_benchmark_docker.py @@ -102,6 +102,45 @@ def _query(query_id: str) -> EvaluationQuery: ) +@pytest.mark.asyncio +async def test_real_docker_timeout_kills_blocked_thread_pool(base_image): + """A host deadline must terminate a worker blocked in executor shutdown.""" + from deploy.benchmark_docker.direct_docker_workspace import DirectDockerWorkspace + from deploy.benchmark_docker.docker_sandbox import DockerReMeSandbox, SandboxCommandTimeoutError + + container_id = f"timeout-{uuid.uuid4().hex[:12]}" + workspace = DirectDockerWorkspace(workspace_id=container_id, base_image=base_image) + case = DockerReMeSandbox( + case_id="timeout", + container_id=container_id, + candidate=ImageCandidate(base_image), + workspace=workspace, + config="unused", + plugins=(), + command_timeout=2, + environment_names=[], + ) + try: + await workspace.initialize() + case.backend = workspace.get_backend() + # Reproduce the shutdown condition without importing ReMe or a plugin. + script = ( + "import threading; from pathlib import Path; from concurrent.futures import ThreadPoolExecutor; " + "pool = ThreadPoolExecutor(1); pool.submit(threading.Event().wait); " + "log = Path('/workspace/case/build_log/build.log'); log.parent.mkdir(parents=True); " + "log.write_text('waiting for thread pool shutdown'); pool.shutdown(wait=True)" + ) + # pylint: disable=protected-access + async with case._operation_lock: + with pytest.raises(SandboxCommandTimeoutError): + await asyncio.wait_for(case._exec("blocked-close", ["python3", "-c", script]), timeout=75) + assert not _container_exists(f"as_ws_{container_id}") + assert not case.reusable + assert await case.read_build_log() == "waiting for thread pool shutdown" + finally: + await case.close() + + @pytest.mark.asyncio async def test_real_docker_source_pool_workspace_reuse(base_image, tmp_path): """Run build, query, workspace reuse, and query-only export in a real container.""" diff --git a/tests/unit/test_benchmark_docker.py b/tests/unit/test_benchmark_docker.py index e997076c..788a6711 100644 --- a/tests/unit/test_benchmark_docker.py +++ b/tests/unit/test_benchmark_docker.py @@ -803,6 +803,7 @@ class FakeWorkspace: def __init__(self, backend: FakeBackend) -> None: self.backend = backend self.closed = False + self.stopped = False async def initialize(self) -> None: pass @@ -813,6 +814,9 @@ class FakeWorkspace: async def close(self) -> None: self.closed = True + async def stop(self) -> None: + self.stopped = True + def _fake_docker_factory(tmp_path, workspace_builder, *, install_projects=(".",)): (tmp_path / "reme").mkdir() @@ -833,6 +837,7 @@ async def test_backend_timeout_prevents_direct_case_reset(tmp_path): with pytest.raises(SandboxCommandTimeoutError): await case.run_build([JobRequest("build")]) assert not case.reusable + assert workspace.stopped and workspace.closed assert case.actions[-1].exit_code == -1 assert await case.read_build_log() == "build still running" with pytest.raises(SandboxCommandError, match="no longer reusable"): @@ -863,6 +868,7 @@ async def test_external_timeout_prevents_case_reset(tmp_path): with pytest.raises(TimeoutError): await asyncio.wait_for(build, timeout=0.01) assert not case.reusable + assert workspace.stopped and workspace.closed 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) @@ -874,6 +880,98 @@ async def test_external_timeout_prevents_case_reset(tmp_path): assert workspace.closed +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["deadline", "backend-timeout", "cancel"]) +async def test_interrupted_worker_stops_before_reading_logs_and_removal(tmp_path, failure): + events = [] + started = asyncio.Event() + stopping = asyncio.Event() + release_stop = asyncio.Event() + + class BlockingBackend(FakeBackend): + async def exec_shell(self, command, *, cwd=None, timeout=None): + if "/workspace/harness/worker.py" in command: + self.files["/workspace/case/build_log/build.log"] = b"closing application" + started.set() + if failure == "backend-timeout": + raise TimeoutError("backend timed out") + await asyncio.Event().wait() # Includes a worker stuck in Application.close(). + return await super().exec_shell(command, cwd=cwd, timeout=timeout) + + async def read_file(self, path): + events.append("read") + assert workspace.stopped + return await super().read_file(path) + + class OrderedWorkspace(FakeWorkspace): + async def stop(self): + stopping.set() + await release_stop.wait() + events.append("stop") + await super().stop() + + async def close(self): + events.append("close") + await super().close() + + workspace = OrderedWorkspace(BlockingBackend()) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: workspace) + case = await factory.create_case("interrupted") + case.command_timeout = 0.01 if failure == "deadline" else 10 + task = asyncio.create_task(case.run_build([JobRequest("build")])) + await asyncio.wait_for(started.wait(), timeout=2) + if failure == "cancel": + task.cancel() + await asyncio.wait_for(stopping.wait(), timeout=2) + if failure == "cancel": + task.cancel() # Repeated cancellation must not interrupt container cleanup. + release_stop.set() + expected = asyncio.CancelledError if failure == "cancel" else SandboxCommandTimeoutError + with pytest.raises(expected): + await asyncio.wait_for(task, timeout=2) + assert events == ["stop", "read", "close"] + assert workspace.closed and not case.reusable + assert await case.read_build_log() == "closing application" + await case.close() + assert events == ["stop", "read", "close"] + + +@pytest.mark.asyncio +async def test_failed_stop_still_removes_container_and_preserves_timeout(tmp_path): + class BrokenStopWorkspace(FakeWorkspace): + async def stop(self): + raise RuntimeError("kill failed") + + workspace = BrokenStopWorkspace(TimedOutBackend()) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: workspace) + case = await factory.create_case("failed-stop") + with pytest.raises(SandboxCommandTimeoutError) as failure: + await case.run_build([JobRequest("build")]) + assert workspace.closed + assert any("kill failed" in note for note in failure.value.__notes__) + with pytest.raises(SandboxCommandError, match="log unavailable"): + await case.read_build_log() + + +@pytest.mark.asyncio +async def test_blocked_diagnostics_cannot_prevent_removal(tmp_path, monkeypatch): + from deploy.benchmark_docker import docker_sandbox + + monkeypatch.setattr(docker_sandbox, "_LOG_TIMEOUT", 0.01) + + class BlockedLogBackend(TimedOutBackend): + async def read_file(self, path): + assert workspace.stopped + await asyncio.Event().wait() + + workspace = FakeWorkspace(BlockedLogBackend()) + factory = _fake_docker_factory(tmp_path, lambda _id, _image, _env: workspace) + case = await factory.create_case("blocked-log") + with pytest.raises(SandboxCommandTimeoutError): + await asyncio.wait_for(case.run_build([JobRequest("build")]), timeout=2) + assert workspace.closed + + @pytest.mark.asyncio async def test_cancelled_request_upload_prevents_case_reset(tmp_path): started = asyncio.Event() @@ -1340,7 +1438,11 @@ async def test_pool_retries_strict_delete_without_replacing_uncertain_container( from deploy.benchmark_docker.direct_docker_workspace import DirectDockerWorkspace workspaces = [] - container = SimpleNamespace(delete=AsyncMock(side_effect=[RuntimeError("delete failed"), None])) + container = SimpleNamespace( + kill=AsyncMock(), + wait=AsyncMock(), + delete=AsyncMock(side_effect=[RuntimeError("delete failed"), RuntimeError("delete failed"), None]), + ) client = SimpleNamespace(close=AsyncMock()) class StrictWorkspace(DirectDockerWorkspace): @@ -1361,6 +1463,6 @@ async def test_pool_retries_strict_delete_without_replacing_uncertain_container( 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 container.delete.await_count == 3 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 index a30521bc..09f8ee33 100644 --- a/tests/unit/test_benchmark_docker_workspace.py +++ b/tests/unit/test_benchmark_docker_workspace.py @@ -14,13 +14,49 @@ from deploy.benchmark_docker.direct_docker_workspace import DirectDockerWorkspac def make_workspace(): workspace = DirectDockerWorkspace(workspace_id="cleanup-test") - workspace._container = SimpleNamespace(delete=AsyncMock()) + workspace._container = SimpleNamespace(kill=AsyncMock(), wait=AsyncMock(), 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_stop_kills_and_confirms_exit_before_removal(): + workspace = make_workspace() + container = workspace._container + await workspace.stop() + container.kill.assert_awaited_once_with(signal="SIGKILL") + container.wait.assert_awaited_once_with(condition="not-running") + container.delete.assert_not_awaited() + assert workspace._backend is not None + assert not workspace.is_alive + await workspace.stop() + assert container.kill.await_count == 1 + await workspace.close() + container.delete.assert_awaited_once_with(force=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["kill", "delete"]) +async def test_docker_cleanup_timeout_retains_handles(monkeypatch, operation): + from deploy.benchmark_docker import direct_docker_workspace + + monkeypatch.setattr(direct_docker_workspace, "DOCKER_CLEANUP_TIMEOUT", 0.01) + workspace = make_workspace() + container = workspace._container + + async def blocked(**_kwargs): + await asyncio.Event().wait() + + getattr(container, operation).side_effect = blocked + with pytest.raises(TimeoutError): + await (workspace.stop() if operation == "kill" else workspace.close()) + assert workspace._container is container + getattr(container, operation).side_effect = None + await workspace.close() + + @pytest.mark.asyncio async def test_delete_failure_retains_handles_and_can_be_retried(): workspace = make_workspace() From 28a384afb5d9c8a29468e235c9f910513150aaff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Fri, 16 Oct 2026 15:38:26 +0800 Subject: [PATCH 14/14] Fix candidate exclusions and benchmark upload serialization --- .../Dockerfile.candidate.dockerignore | 49 ++++++++-- deploy/benchmark_docker/docker_sandbox.py | 13 ++- tests/unit/test_benchmark_docker_queries.py | 65 +++++++++++++ tests/unit/test_benchmark_docker_upload.py | 95 +++++++++++++++++++ 4 files changed, 211 insertions(+), 11 deletions(-) create mode 100644 tests/unit/test_benchmark_docker_queries.py create mode 100644 tests/unit/test_benchmark_docker_upload.py diff --git a/deploy/benchmark_docker/Dockerfile.candidate.dockerignore b/deploy/benchmark_docker/Dockerfile.candidate.dockerignore index c9628b66..27ede6f8 100644 --- a/deploy/benchmark_docker/Dockerfile.candidate.dockerignore +++ b/deploy/benchmark_docker/Dockerfile.candidate.dockerignore @@ -1,21 +1,52 @@ +# Only candidate-side benchmark plugins are needed in this image. +# Add another explicit plugin directory here before selecting it with REME_PLUGIN_PATH. +plugins/* +!plugins/lme +!plugins/beam + +# Keep exclusions after plugin exceptions so nested secrets stay excluded. .git -.reme -.env -.env.* -.venv +**/.DS_Store +**/.ssh +**/.codex +**/.claude +**/private* +**/.env +**/.env.* +**/.reme +**/.venv +**/venv **/__pycache__ +**/*.py[cod] **/.pytest_cache +**/.mypy_cache +**/.ruff_cache +**/.coverage* +**/htmlcov **/node_modules **/dist -**/build +**/dist-static +**/.generated +**/.next +**/.vite +**/.wrangler +**/.cache **/*.egg-info +**/build +**/logs +**/*.log +**/*.tmp +reme_studio/src/reme_studio/static benchmark -baselines -meta-reme -output -tests +cookbook docs github-pages integrations skills +tests + +# Candidate build inputs do not need local benchmark results or Studio sources. +baselines +meta-reme +output reme_studio diff --git a/deploy/benchmark_docker/docker_sandbox.py b/deploy/benchmark_docker/docker_sandbox.py index 7b3c48ce..ea360213 100644 --- a/deploy/benchmark_docker/docker_sandbox.py +++ b/deploy/benchmark_docker/docker_sandbox.py @@ -693,7 +693,16 @@ class DockerReMeSandbox: serialized: list[dict[str, Any]] = [] query_ids: list[str] = [] for query in queries: - value = asdict(query) if isinstance(query, EvaluationQuery) else dict(query) + value = ( + { + "query_id": query.query_id, + "question": query.question, + "answer_arguments": query.answer_arguments, + "answer_job": query.answer_job, + } + if isinstance(query, EvaluationQuery) + else dict(query) + ) query_id = value.get("query_id") if not isinstance(query_id, str): raise ValueError("every evaluation query requires a string query_id") @@ -916,7 +925,7 @@ class DockerReMeSandbox: The copy is made through a validated archive, so each sandbox receives an independent filesystem tree and can run queries in parallel. """ - archive = _archive_workspace(source) + archive = await asyncio.to_thread(_archive_workspace, source) async with self._operation_lock: assert self.backend is not None, "sandbox is not initialized" if clear: diff --git a/tests/unit/test_benchmark_docker_queries.py b/tests/unit/test_benchmark_docker_queries.py new file mode 100644 index 00000000..0b87951a --- /dev/null +++ b/tests/unit/test_benchmark_docker_queries.py @@ -0,0 +1,65 @@ +"""Host-only query answers never participate in candidate serialization.""" + +# pylint: disable=missing-function-docstring,protected-access + +import json +import threading +from unittest.mock import AsyncMock + +import pytest + +from deploy.benchmark_docker import EvaluationQuery, ImageCandidate +from deploy.benchmark_docker.docker_sandbox import DockerReMeSandbox + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["run_query", "run_queries"]) +@pytest.mark.parametrize("as_mapping", [False, True]) +async def test_query_upload_ignores_uncopyable_golden_answer(monkeypatch, method, as_mapping): + case = DockerReMeSandbox( + case_id="query", + container_id="unused", + candidate=ImageCandidate("unused"), + workspace=AsyncMock(), + config="default", + plugins=(), + command_timeout=5, + environment_names=[], + ) + arguments = {"limit": 3} + golden_answer = threading.Lock() + query = EvaluationQuery("q1", "Question?", golden_answer, arguments, "custom_answer") + if as_mapping: + query = { + "query_id": "q1", + "question": "Question?", + "golden_answer": golden_answer, + "answer_arguments": arguments, + "answer_job": "custom_answer", + "rubric": golden_answer, + } + result = {"answer": "candidate answer"} + worker = AsyncMock(return_value={"queries": [result]}) + monkeypatch.setattr(case, "_run_worker_request", worker) + + if method == "run_query": + assert await case.run_query(query) == result + else: + assert await case.run_queries([query]) == {"queries": [result]} + + mode, payload = worker.await_args.args + assert mode == ("query" if method == "run_query" else "queries") + assert payload == { + "queries": [ + { + "query_id": "q1", + "question": "Question?", + "answer_arguments": {"limit": 3}, + "answer_job": "custom_answer", + }, + ], + } + assert payload["queries"][0]["answer_arguments"] is not arguments + # A host-only lock is neither deep-copied nor included in the JSON request. + json.dumps(payload) + assert not golden_answer.locked() diff --git a/tests/unit/test_benchmark_docker_upload.py b/tests/unit/test_benchmark_docker_upload.py new file mode 100644 index 00000000..92157dd5 --- /dev/null +++ b/tests/unit/test_benchmark_docker_upload.py @@ -0,0 +1,95 @@ +"""Candidate build exclusions and nonblocking workspace upload without Docker.""" + +# pylint: disable=missing-function-docstring,protected-access + +import asyncio +import io +from pathlib import Path +import tarfile +import threading +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from deploy.benchmark_docker import ImageCandidate +from deploy.benchmark_docker import docker_sandbox as sandbox_module +from deploy.benchmark_docker.docker_sandbox import _archive_workspace + + +def test_candidate_image_preserves_root_exclusions_after_plugin_exceptions(): + root = Path(__file__).resolve().parents[2] + root_rules = (root / ".dockerignore").read_text(encoding="utf-8").splitlines() + candidate_rules = ( + (root / "deploy/benchmark_docker/Dockerfile.candidate.dockerignore").read_text(encoding="utf-8").splitlines() + ) + exceptions = [rule for rule in candidate_rules if rule.startswith("!")] + assert set(exceptions) == {"!plugins/lme", "!plugins/beam"} + assert "plugins/*" in candidate_rules + last_exception = max(candidate_rules.index(rule) for rule in exceptions) + exclusions = candidate_rules[last_exception + 1 :] + # New root protections must also cover the Dockerfile-specific context, + # including secrets inside otherwise allowed plugin directories. + assert set(root_rules) - {"plugins", ""} <= set(exclusions) + assert {"baselines", "meta-reme", "output", "reme_studio"} <= set(exclusions) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("archive_error", [False, True]) +async def test_workspace_archiving_keeps_event_loop_responsive(tmp_path, monkeypatch, archive_error): + backend = AsyncMock() + backend.exec_shell.return_value = SimpleNamespace(exit_code=0, stdout=b"", stderr=b"", ok=lambda: True) + case = sandbox_module.DockerReMeSandbox( + case_id="archive-thread", + container_id="unused", + candidate=ImageCandidate("unused"), + workspace=AsyncMock(), + config="default", + plugins=(), + command_timeout=5, + environment_names=[], + ) + case.backend = backend + source = tmp_path / "memory" + source.mkdir() + (source / "note.md").write_text("saved memory", encoding="utf-8") + loop = asyncio.get_running_loop() + started = asyncio.Event() + release = threading.Event() + event_loop_thread = threading.get_ident() + + def blocking_archive(path): + assert threading.get_ident() != event_loop_thread + loop.call_soon_threadsafe(started.set) + if not release.wait(timeout=5): + raise TimeoutError("test did not release archive worker") + if archive_error: + raise ValueError("invalid workspace archive") + return _archive_workspace(path) + + monkeypatch.setattr(sandbox_module, "_archive_workspace", blocking_archive) + upload = asyncio.create_task(case.upload_workspace(source)) + try: + await asyncio.wait_for(started.wait(), timeout=2) + await asyncio.sleep(0) + assert not upload.done() + # No workspace clearing or upload may happen before validation finishes. + backend.exec_shell.assert_not_awaited() + backend.write_file.assert_not_awaited() + release.set() + if archive_error: + with pytest.raises(ValueError, match="invalid workspace archive"): + await upload + backend.exec_shell.assert_not_awaited() + backend.write_file.assert_not_awaited() + else: + await upload + backend.write_file.assert_awaited_once() + archive_path, payload = backend.write_file.await_args.args + assert archive_path == sandbox_module.WORKSPACE_ARCHIVE + with tarfile.open(fileobj=io.BytesIO(payload), mode="r:gz") as archive: + assert archive.extractfile("note.md").read() == b"saved memory" + finally: + release.set() + await asyncio.gather(upload, return_exceptions=True) + await case.close()