From 7dc622c627e199aa9653cf2c084ba5991a65829d Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Mon, 19 Oct 2026 16:20:07 +0800 Subject: [PATCH] Stream benchmark workspace archives and extract plan validation --- deploy/benchmark_docker/docker_sandbox.py | 120 ++++++++++----------- deploy/benchmark_docker/pool.py | 37 ++++--- tests/unit/test_benchmark_docker_upload.py | 51 +++++++++ 3 files changed, 132 insertions(+), 76 deletions(-) diff --git a/deploy/benchmark_docker/docker_sandbox.py b/deploy/benchmark_docker/docker_sandbox.py index 56cc17b5..c7452e00 100644 --- a/deploy/benchmark_docker/docker_sandbox.py +++ b/deploy/benchmark_docker/docker_sandbox.py @@ -12,7 +12,7 @@ from pathlib import Path from pathlib import PurePosixPath import re import tarfile -from typing import Any, Callable, Literal, Protocol +from typing import Any, BinaryIO, Callable, Literal, Protocol import uuid from ._cleanup import cleanup_on_error @@ -171,70 +171,70 @@ def _archive_workspace(source: str | Path) -> bytes: 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: + with tarfile.open(fileobj=compressed, mode="w") as output: + + def add(info: tarfile.TarInfo, content: BinaryIO | None = None) -> None: 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) + output.addfile(info, content) + + 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 + add(info) + elif item.is_file(): + with item.open("rb") as content: + info.size = item.stat().st_size + add(info, content) + else: + raise ValueError(f"workspace source contains unsupported file: {relative}") + elif source_path.is_file(): + with tarfile.open(source_path, mode="r:*") as archive: + members: 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 + add(info) + elif member.isfile(): + content = archive.extractfile(member) + if content is None: + raise ValueError(f"workspace archive cannot read file: {path}") + with content: + info.size = member.size + add(info, content) + else: + raise ValueError(f"workspace archive contains unsupported entry: {member.name!r}") + else: + raise ValueError(f"workspace source is neither a directory nor a file: {source_path}") return raw.getvalue() diff --git a/deploy/benchmark_docker/pool.py b/deploy/benchmark_docker/pool.py index 667bddd9..7e6f417a 100644 --- a/deploy/benchmark_docker/pool.py +++ b/deploy/benchmark_docker/pool.py @@ -83,21 +83,7 @@ class BenchmarkDockerPool: """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) + self._validate_plans(plans) 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( @@ -107,7 +93,7 @@ class BenchmarkDockerPool: available: asyncio.Queue[DockerReMeSandbox | None] = asyncio.Queue() for worker in workers: available.put_nowait(worker) - slots = _WorkerSlots(workers, available, set(by_id), len(workers), len(workers)) + slots = _WorkerSlots(workers, available, {plan.case_id for plan in plans}, len(workers), len(workers)) worker_results: dict[int, list[CaseResult]] = {} results: list[CaseResult] = [] run_error: BaseException | None = None @@ -259,6 +245,25 @@ class BenchmarkDockerPool: if cleanup_cancel is not None and run_error is None: raise cleanup_cancel + @staticmethod + def _validate_plans(plans: list[CasePlan]) -> None: + """Reject invalid plans and dependencies before creating containers.""" + 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") + BenchmarkDockerPool._check_acyclic(by_id) + @staticmethod def _check_acyclic(plans: dict[str, CasePlan]) -> None: visiting: set[str] = set() diff --git a/tests/unit/test_benchmark_docker_upload.py b/tests/unit/test_benchmark_docker_upload.py index 92157dd5..1f21d1a9 100644 --- a/tests/unit/test_benchmark_docker_upload.py +++ b/tests/unit/test_benchmark_docker_upload.py @@ -7,6 +7,7 @@ import io from pathlib import Path import tarfile import threading +import tracemalloc from types import SimpleNamespace from unittest.mock import AsyncMock @@ -17,6 +18,56 @@ from deploy.benchmark_docker import docker_sandbox as sandbox_module from deploy.benchmark_docker.docker_sandbox import _archive_workspace +@pytest.mark.parametrize("source_kind", ["directory", "archive"]) +def test_workspace_archive_streams_large_files(tmp_path, source_kind): + workspace = tmp_path / "workspace" + workspace.mkdir() + note = workspace / "large.md" + size = 4 * 1024 * 1024 + with note.open("wb") as content: + for _ in range(size // 65536): + content.write(b"x" * 65536) + note.chmod(0o640) + source = workspace + if source_kind == "archive": + source = tmp_path / "source.tar.gz" + with tarfile.open(source, "w:gz") as archive: + archive.add(note, arcname="large.md") + + # A compressible multi-megabyte file must not be materialized in host RAM. + # Leave ample room for gzip/tar buffers while catching whole-file reads. + tracemalloc.start() + try: + payload = _archive_workspace(source) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + assert peak < size // 2 + assert _archive_workspace(source) == payload + with tarfile.open(fileobj=io.BytesIO(payload), mode="r:gz") as archive: + member = archive.getmember("large.md") + assert (member.size, member.mode, member.mtime, member.uid, member.gid) == (size, 0o640, 0, 0, 0) + assert archive.extractfile(member).read() == b"x" * size + + +@pytest.mark.parametrize("kind", ["traversal", "symlink", "hardlink", "fifo", "duplicate"]) +def test_workspace_archive_rejects_unsafe_entries(tmp_path, kind): + source = tmp_path / "unsafe.tar.gz" + with tarfile.open(source, "w:gz") as archive: + info = tarfile.TarInfo("../escape" if kind == "traversal" else "note.md") + info.type = { + "symlink": tarfile.SYMTYPE, + "hardlink": tarfile.LNKTYPE, + "fifo": tarfile.FIFOTYPE, + }.get(kind, tarfile.REGTYPE) + info.linkname = "outside" + archive.addfile(info) + if kind == "duplicate": + archive.addfile(info) + with pytest.raises(ValueError, match="unsafe path|unsupported entry|duplicate path"): + _archive_workspace(source) + + 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()