mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-26 06:04:11 +00:00
Stream benchmark workspace archives and extract plan validation
This commit is contained in:
jinli.yl
parent
7af6f3c3ad
commit
7dc622c627
3 changed files with 132 additions and 76 deletions
|
|
@ -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()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue