mrg.bench¶

class Sandbox:
    """A locked-down place to run a sandboxed agent, with a wired-up silicon path.

    Args:
        files: files copied into the container's ``/work`` (your ``design.py``,
            the agent entrypoint, any helpers).
        isolation: ``"locked"`` (default; untrusted, no net/key/ro-root),
            ``"dev"`` (trusted local poking), or a :class:`SandboxProfile`.
        silicon: ``"auto"`` (default: real cloud if a key is set, else a no-op),
            ``"cloud"``, ``"mock"``, or your own ``SiliconFn``.
        api_key: passed to the cloud silicon backend. Stays in this trusted
            process — never in the container.
        api_url: passed to the cloud silicon backend.
        sys_clk_freq: passed to the cloud silicon backend.
        guard: optional ``(design_bytes, report) -> reject-reason | None`` to vet
            promotes. Default ``None`` = every promote goes straight through.
        image: override the sandbox docker image.
    """

    def __init__(
        self,
        files: Sequence[str | Path] = (),
        *,
        isolation: str | SandboxProfile = "locked",
        silicon: str | SiliconFn = "auto",
        api_key: str | None = None,
        api_url: str | None = None,
        sys_clk_freq: int | None = None,
        guard: Guard | None = None,
        image: str | None = None,
        poll_interval: float = 0.2,
    ) -> None:
        self.files = [Path(f) for f in files]
        self.profile = _resolve_profile(isolation, image)
        self._broker = _PromoteBroker(
            _resolve_silicon(silicon, api_key, api_url, sys_clk_freq), guard
        )
        self.poll_interval = poll_interval

    def run(self, entrypoint: str | Path, *, timeout: int = 1800) -> SandboxResult:
        """Launch the agent in the container and broker its promotes to silicon.

        Creates a throwaway workspace, copies :attr:`files` in, runs
        ``python /work/<entrypoint>`` under the isolation profile while an
        internal loop answers promotes, then tears the workspace down.
        """
        workspace = Path(tempfile.mkdtemp(prefix="mrg_sandbox_"))
        try:
            for f in self.files:
                shutil.copy(f, workspace / f.name)

            entry = Path(entrypoint).name
            promotions: list[dict] = []
            stop = threading.Event()

            def poll_loop() -> None:
                while not stop.is_set():
                    promotions.extend(self._broker.poll_once(workspace))
                    time.sleep(self.poll_interval)

            poller = threading.Thread(target=poll_loop, daemon=True)
            poller.start()
            try:
                proc = run_sandbox(
                    ["python", f"/work/{entry}"],
                    workspace=workspace,
                    profile=self.profile,
                    timeout=timeout,
                )
            finally:
                stop.set()
                poller.join(timeout=2)
            # Drain any promote that landed in the final tick.
            promotions.extend(self._broker.poll_once(workspace))

            return SandboxResult(
                proc.returncode, proc.stdout, proc.stderr, promotions
            )
        finally:
            shutil.rmtree(workspace, ignore_errors=True)

@dataclass
class SandboxResult:
    """Outcome of a :meth:`Sandbox.run`."""

    returncode: int
    stdout: str
    stderr: str
    promotions: list[dict] = field(default_factory=list)

    @property
    def ok(self) -> bool:
        return self.returncode == 0

@dataclass
class SandboxProfile:
    """The docker run constraints applied to the container.

    Defaults are the locked-down profile: no egress, dropped capabilities,
    read-only root, bounded memory/cpu/pids. Only the workspace mount and tmpfs
    are writable. Use ``.locked()`` / ``.dev()`` for the named presets.
    """

    image: str = DEFAULT_IMAGE
    network: str = "none"  # no egress for untrusted builds
    memory: str = "8g"  # nextpnr is the memory hog
    cpus: str = "4"
    pids_limit: int = 512
    read_only_root: bool = True
    user: str | None = None  # e.g. "1000:1000"; None = image default
    # Writable scratch on the read-only rootfs. mrg_build builds in /tmp by
    # default, and HOME is pointed here so tool caches don't hit the ro root.
    tmpfs: tuple[str, ...] = ("/tmp:size=4g",)
    home: str = "/tmp"
    # Extra env vars passed into the container. EMPTY in the locked profile
    # (never inject a key into untrusted code); dev() forwards MRG_API_KEY here.
    env: dict[str, str] = field(default_factory=dict)

    def argv(
        self,
        *,
        command: list[str],
        workspace: Path | str | None = None,
        mounts: tuple[tuple[str, str, str], ...] = (),
    ) -> list[str]:
        """Build the full ``docker run`` argv for one invocation.

        ``workspace`` (if given) is mounted read-write at /work. ``mounts`` are
        extra (src, dst, mode) bind mounts.
        """
        argv = [
            "docker", "run", "--rm",
            f"--network={self.network}",
            f"--memory={self.memory}", f"--memory-swap={self.memory}",
            f"--cpus={self.cpus}", f"--pids-limit={self.pids_limit}",
            "--cap-drop=ALL",
            "--security-opt", "no-new-privileges",
            "-e", f"HOME={self.home}",
        ]
        for key, value in self.env.items():
            argv += ["-e", f"{key}={value}"]
        if self.read_only_root:
            argv.append("--read-only")
        for t in self.tmpfs:
            argv += ["--tmpfs", t]
        if self.user:
            argv += ["--user", self.user]
        if workspace is not None:
            argv += ["-v", f"{Path(workspace).resolve()}:/work"]
        for src, dst, mode in mounts:
            argv += ["-v", f"{Path(src).resolve()}:{dst}:{mode}"]
        argv += [self.image, *command]
        return argv

    @classmethod
    def locked(cls, **overrides) -> SandboxProfile:
        """The default profile — for UNTRUSTED agent code and benchmark eval.

        No network, no credentials, read-only root, dropped caps. Reproducible
        and benchmark-valid. ``overrides`` tweak individual fields if needed.
        """
        return cls(**overrides)

    @classmethod
    def dev(cls, *, forward_api_key: bool = True, **overrides) -> SandboxProfile:
        """TRUSTED experimentation only — relaxed, WITH internet.

        For a developer (or your own agent) poking around on your machine with
        your own key: network on, writable root, and your host ``MRG_API_KEY``
        forwarded in so ``mrg.cloud`` works directly.

        Do NOT run untrusted agent code in this profile, and do NOT treat its
        results as benchmark scores — internet access makes them non-reproducible
        by definition.
        """
        env = dict(overrides.pop("env", {}))
        if forward_api_key:
            key = os.environ.get("MRG_API_KEY")
            if key:
                env["MRG_API_KEY"] = key
        overrides.setdefault("network", "bridge")
        overrides.setdefault("read_only_root", False)
        return cls(env=env, **overrides)