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..27ede6f8 --- /dev/null +++ b/deploy/benchmark_docker/Dockerfile.candidate.dockerignore @@ -0,0 +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 +**/.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 +**/dist-static +**/.generated +**/.next +**/.vite +**/.wrangler +**/.cache +**/*.egg-info +**/build +**/logs +**/*.log +**/*.tmp +reme_studio/src/reme_studio/static +benchmark +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/README.md b/deploy/benchmark_docker/README.md new file mode 100644 index 00000000..438eb0e7 --- /dev/null +++ b/deploy/benchmark_docker/README.md @@ -0,0 +1,140 @@ +# 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 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`. `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 +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()`, +`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 +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/_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 new file mode 100644 index 00000000..e8d28909 --- /dev/null +++ b/deploy/benchmark_docker/candidate.py @@ -0,0 +1,186 @@ +"""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 +import os +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]] = [] + # 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: + 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..ea62749c --- /dev/null +++ b/deploy/benchmark_docker/direct_docker_workspace.py @@ -0,0 +1,118 @@ +"""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 + +import asyncio +import os +import sys +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.""" + + _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() + self._stopped = False + + 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.""" + 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 + 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 ( + 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], + timeout=10.0, + ) + except Exception: + pass + try: + # 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 asyncio.wait_for(self._client.close(), timeout=DOCKER_CLEANUP_TIMEOUT) + 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 new file mode 100644 index 00000000..ea360213 --- /dev/null +++ b/deploy/benchmark_docker/docker_sandbox.py @@ -0,0 +1,1142 @@ +"""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 ._cleanup import cleanup_on_error +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}$") +_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"}, +) + + +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 + + 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 + + +WorkspaceBuilder = Callable[[str, str, dict[str, str]], Workspace] + + +class SandboxCommandError(RuntimeError): + """A required command failed inside a case sandbox.""" + + +class SandboxCommandTimeoutError(SandboxCommandError): + """A command timed out; cleanup failures are attached as exception notes.""" + + +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 _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`, 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[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 + info = tarfile.TarInfo(path.as_posix()) + info.mode = item.stat().st_mode & 0o777 + if item.is_dir(): + info.type = tarfile.DIRTYPE + data = None + elif item.is_file(): + 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: dict[PurePosixPath, tarfile.TarInfo] = {} + for member in archive.getmembers(): + path = _workspace_member_path(member.name) + if path is None: + continue + 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 path == workspace_root: + if not member.isdir(): + raise ValueError("case export has a non-directory reme_workspace entry") + continue + path = path.relative_to(workspace_root) + info = tarfile.TarInfo(path.as_posix()) + info.mode = member.mode & 0o777 + if member.isdir(): + info.type = tarfile.DIRTYPE + 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.size = len(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 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 as exc: + await cleanup_on_error([(f"workspace {container_id!r}", workspace.close())], exc) + 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 | 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 as exc: + error = exc + for task in tasks: + if not task.done() and not task.cancelling(): + task.cancel() + results = await asyncio.gather(*tasks, return_exceptions=True) + created = [result for result in results if isinstance(result, DockerReMeSandbox)] + failures = [result for result in results if isinstance(result, BaseException)] + 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: + """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._reusable = True + self._operation_lock = asyncio.Lock() + self._build_log: bytes | None = None + + @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.""" + 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._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._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 _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 (Exception, asyncio.CancelledError): + # 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: + raise SandboxCommandError("sandbox container is no longer reusable") + started = _utc_now() + try: + async with asyncio.timeout(self.command_timeout): + result = await self.backend.exec_shell(command, cwd=cwd, timeout=self.command_timeout) + except TimeoutError as exc: + 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. + await cleanup_on_error([(f"container {self.container_id!r}", self._abort())], exc) + raise + 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(), + ), + ) + if result.exit_code == -1: + 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(): + 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.""" + 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.""" + 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._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, + [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: + 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 + + 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 = ( + { + "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") + 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: + 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.""" + async with self._operation_lock: + result = await self._run_job("agentic_answer", {"query": query, "query_time": query_time}) + 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 export( + self, + destination: str | Path, + *, + profile: Literal["analysis", "full"] = "analysis", + include_candidate: bool = False, + ) -> Path: + """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: + 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: + return await self._export_archive( + destination, + "export_workspace", + self.runtime_workspace, + ["."], + archive_path=WORKSPACE_EXPORT_ARCHIVE, + ) + + 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() + return self._save_host_file(destination, content) + + 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: + 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: + """Read a required sandbox artifact, preserving its export-specific error.""" + assert self.backend is not None, "sandbox is not initialized" + try: + return await self.backend.read_file(path) + except FileNotFoundError as exc: + raise SandboxCommandError(f"{description} is missing required artifact: {path}") from exc + + async def export_memory_construction(self, destination: str | Path) -> Path: + """Download the workspace and build log; reusable with upload_workspace.""" + async with self._operation_lock: + 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: + 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: + 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 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: + pass + else: + members.append("build_log") + 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, + 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 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. + The copy is made through a validated archive, so each sandbox receives + an independent filesystem tree and can run queries in parallel. + """ + 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: + 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._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 _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, + "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._write_file( + 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( + (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) + self._save_host_file(target, await self.backend.read_file(EXPORT_ARCHIVE)) + 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. + + 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: + 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 + + 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..77dea88a --- /dev/null +++ b/deploy/benchmark_docker/models.py @@ -0,0 +1,73 @@ +"""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) + answer_job: str = "agentic_answer" diff --git a/deploy/benchmark_docker/pool.py b/deploy/benchmark_docker/pool.py new file mode 100644 index 00000000..740cc058 --- /dev/null +++ b/deploy/benchmark_docker/pool.py @@ -0,0 +1,277 @@ +"""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 ._cleanup import finish_cleanup, format_error +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 and owned-worker cleanup succeeded.""" + 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) + 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 | None] = asyncio.Queue() + for worker in workers: + available.put_nowait(worker) + 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: + 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: + 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) + worker_results.setdefault(id(worker), []).append(result) + + def record_lifecycle_error(phase: str, exc: Exception) -> str: + detail = f"{phase} failed: {format_error(exc)}" + result.error = f"{result.error}; {detail}" if result.error else detail + return detail + + 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 = format_error(exc) + await self._retain_failed_build_log(worker, plan, result) + return result + if not result.build.get("success"): + result.error = "build failed" + 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) + # Capture before queries, since answering may change the + # workspace. Dependents receive construction state. + 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: + 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 = format_error(exc) + return result + finally: + if worker.reusable: + available.put_nowait(worker) + else: + # Confirm cleanup (or retry a failed removal) before + # allowing another case into this slot. + closed = False + try: + await worker.close() + except Exception as exc: + exhaustion_error = ( + f"no available workers: {record_lifecycle_error('worker close', exc)}" + ) + else: + closed = True + + replaced = False + if closed and unclaimed 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 unclaimed: + available.put_nowait(None) + + 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() 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 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: + 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..c77ebaa8 --- /dev/null +++ b/deploy/benchmark_docker/worker.py @@ -0,0 +1,545 @@ +"""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 +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +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.plugin import resolve_plugin_runtime + 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) + # 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``) + # 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. + result = { + "success": False, + "answer": "", + "metadata": {}, + "error": f"{type(exc).__name__}: {exc}", + "exception": _exception_details(exc), + "traceback": traceback.format_exc(), + } + 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: + """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) + 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"] + job_args = request.get("arguments") or {} + _, app_config = _prepare_runtime(request) + from reme.reme import ReMe + + app = ReMe(**app_config) + failed_results: list[dict[str, Any]] = [] + async with _application_scope(app, failed_results): + await app.start() + 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]: + """Execute all construction jobs in one Application and one 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() + + 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) + 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 + 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_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. + + ``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]] = [] + 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) + 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, + "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) + elif mode == "checkpoint": + operation = _run_checkpoint(request) + 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 97cc3588..7a0b8b79 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,6 +78,10 @@ dev = [ "pytest>=9.1.1", "pytest-asyncio>=0.23", ] +benchmark-docker = [ + "reme-ai[as]", + "aiodocker", +] full = [ "reme-ai[core]", "reme-ai[qwen-tokenizer]", @@ -95,12 +99,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..82a59fbd --- /dev/null +++ b/tests/integration/test_benchmark_docker.py @@ -0,0 +1,227 @@ +"""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_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.""" + 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..788a6711 --- /dev/null +++ b/tests/unit/test_benchmark_docker.py @@ -0,0 +1,1468 @@ +"""Pool ordering, reuse, and cleanup without Docker or model calls.""" + +# pylint: disable=missing-class-docstring,missing-function-docstring,protected-access,unused-argument + +from __future__ import annotations + +import asyncio +import io +import json +import os +from pathlib import Path +import subprocess +import tarfile + +import pytest + +from deploy.benchmark_docker import ( + BenchmarkDockerPool, + CasePlan, + DockerReMeSandboxFactory, + EvaluationQuery, + JobRequest, + SourceCandidate, + SourceSnapshot, +) +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 + + 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(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", + export_evaluation_to=tmp_path / "evaluation.tar.gz", + ), + 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 ("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)) + 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 +@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): + if raises: + raise TimeoutError("build exceeded command timeout") + 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 + 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_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_unblocks_waiting_and_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"),)), + 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 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 +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 "no available workers" in results[1].error + assert "close failed" in results[1].error + assert ("next", "build", 1) not in events + + +@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")) + + +@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") + 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) + + +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): + 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): + if failure_phase == "config": + + 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 +@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): + get_logger(log_to_console=False, log_to_file=False, force_init=True) + + async def start(self): + if failure_phase in {"start", "start_and_close"}: + raise startup_error + + async def close(self): + 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} + + 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 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] = [] + 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): + 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}} + 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"} + 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 +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 + + 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): + from types import SimpleNamespace + + self.commands.append(command) + output = b"" + exit_code = 0 + 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]]) + 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=exit_code, stdout=output, stderr=b"", ok=lambda: exit_code == 0) + + +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 + self.closed = False + self.stopped = False + + async def initialize(self) -> None: + pass + + def get_backend(self) -> FakeBackend: + return self.backend + + 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() + (tmp_path / "reme" / "__init__.py").write_text("ok", encoding="utf-8") + return DockerReMeSandboxFactory( + SourceCandidate(SourceSnapshot.from_directory(tmp_path), install_projects=install_projects), + plugins=("lme",), + 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): + 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"): + await case.reset_case("next") + finally: + await case.close() + 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 + 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) + 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 +@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() + + 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 +@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 operation == "write" and "/inbox/" in path: + raise RuntimeError("Docker stream failed") + await super().write_file(path, data) + + async def exec_shell(self, command, *, cwd=None, timeout=None): + 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(f"{operation}-error") + try: + 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"): + 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() + workspace = FakeWorkspace(backend) + factory = _fake_docker_factory( + tmp_path, + lambda _id, _image, _env: workspace, + install_projects=(".", "plugins/lme"), + ) + 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_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): + timed_out_workspace = FakeWorkspace(TimedOutBackend()) + replacement_workspace = FakeWorkspace(FakeBackend()) + workspaces = iter((timed_out_workspace, replacement_workspace)) + 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"),)), + ] + + 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 +@pytest.mark.parametrize( + ("method", "required", "description", "members"), + [ + ( + "export_memory_construction", + "build_log/build.log", + "memory construction export", + ["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"], + ), + ], +) +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) 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 ["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): + 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): + started = asyncio.Event() + + class BlockingWorkspace(FakeWorkspace): + async def initialize(self): + started.set() + await asyncio.Event().wait() + + workspace = BlockingWorkspace(FakeBackend()) + 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() + with pytest.raises(asyncio.CancelledError): + await task + assert workspace.closed + + +@pytest.mark.asyncio +async def test_cancelled_batch_creation_closes_completed_workspace_and_reports_cleanup_failure(tmp_path): + second_started = asyncio.Event() + + class BlockingWorkspace(FakeWorkspace): + async def initialize(self): + 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)) + 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) + task.cancel() + 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 +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") 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 + + pytest.importorskip("aiodocker") + from deploy.benchmark_docker.direct_docker_workspace import DirectDockerWorkspace + + workspaces = [] + container = SimpleNamespace( + kill=AsyncMock(), + wait=AsyncMock(), + delete=AsyncMock(side_effect=[RuntimeError("delete failed"), RuntimeError("delete failed"), None]), + ) + 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() + + 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 == 3 + assert workspaces[0]._container is workspaces[0]._client is None + client.close.assert_awaited_once() 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() diff --git a/tests/unit/test_benchmark_docker_workspace.py b/tests/unit/test_benchmark_docker_workspace.py new file mode 100644 index 00000000..09f8ee33 --- /dev/null +++ b/tests/unit/test_benchmark_docker_workspace.py @@ -0,0 +1,160 @@ +"""Direct workspace lifecycle without a Docker daemon.""" + +# pylint: disable=missing-function-docstring,protected-access,wrong-import-position + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +DockerError = pytest.importorskip("aiodocker.exceptions").DockerError +from deploy.benchmark_docker.direct_docker_workspace import DirectDockerWorkspace # noqa: E402 + + +def make_workspace(): + workspace = DirectDockerWorkspace(workspace_id="cleanup-test") + 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() + 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 diff --git a/tests/unit/test_package_versions.py b/tests/unit/test_package_versions.py index 07855130..2f7cd49f 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"]