EvalHarness/evalharness/sandbox/docker.py

"""Docker sandbox: one implementation serving both consumers.

exec():  untrusted model-generated code — hard isolation
         (--network none, cpu/mem/pids caps, read-only rootfs, tmpfs /tmp;
          writes escape only through explicit bind mounts)
serve(): trusted engine containers (vllm/sglang) — network ON (weights pull),
         GPU passthrough; consumed by the model Deployer via acquire()
"""

import os
import shlex
import subprocess
import tempfile
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

from .base import EnvHandle, ExecResult, Sandbox, acquire, register_sandbox


def _run(cmd: List[str], **kw) -> subprocess.CompletedProcess:
    return subprocess.run(cmd, capture_output=True, text=True, **kw)


def docker_available() -> bool:
    return _run(['docker', 'info']).returncode == 0


@register_sandbox('docker')
class DockerSandbox(Sandbox):
    name = 'docker'

    DEFAULT_EXEC_IMAGE = 'python:3.11-slim'

    def exec(
        self,
        files: Dict[str, str],
        entry: str = 'main.py',
        mounts: Optional[Dict[str, str]] = None,
        timeout_s: int = 60,
        image: str = '',
    ) -> ExecResult:
        img = image or self.DEFAULT_EXEC_IMAGE
        with tempfile.TemporaryDirectory(prefix='eh-sbx-') as host_dir:
            workdir = Path(host_dir) / 'work'
            workdir.mkdir()
            for fname, content in (files or {}).items():
                dest = workdir / fname
                dest.parent.mkdir(parents=True, exist_ok=True)
                dest.write_text(content, encoding='utf-8')
            cmd = [
                'docker', 'run', '--rm',
                '--network', 'none',              # untrusted code: no egress
                '--cpus', '2', '--memory', '2g', '--pids-limit', '256',
                '--read-only', '--tmpfs', '/tmp:rw,size=64m',
                # /work must be writable: BigCodeBench tasks write output
                # files (task_func_data/, matplotlib caches, etc.) to cwd;
                # the official Evaluate.Dockerfile runs with a writable fs
                '-v', f'{workdir}:/work:rw',
            ]
            out_host = None
            if mounts:
                for cpath, hpath in mounts.items():
                    out_host = Path(hpath).expanduser()
                    out_host.mkdir(parents=True, exist_ok=True)
                    cmd += ['-v', f'{out_host}:{cpath}:rw']
            runner = ['python', f'/work/{entry}'] if entry.endswith('.py') \
                else ['sh', f'/work/{entry}']
            # Named + retried runs. A timed-out/killed `docker run` only kills
            # the CLI client -- the container lives on (--rm fires on EXIT)
            # and leaked containers slowly drown the daemon; the name gives
            # us something to rm -f. Exit 125 is a DAEMON-side failure (shim
            # error etc.), not the model's code failing -- scoring it pass=0
            # would corrupt results, so retry it.
            import uuid

            t0 = time.time()
            proc = None
            for attempt in range(3):
                cname = f'eh-exec-{uuid.uuid4().hex[:10]}'
                full = cmd + ['--name', cname, img, *runner]
                try:
                    proc = _run(full, timeout=timeout_s + 30)
                except subprocess.TimeoutExpired:
                    _run(['docker', 'rm', '-f', cname])  # CLI died, container didn't
                    return ExecResult(exit_code=-1, timed_out=True, duration_s=timeout_s,
                                      error=f'sandbox timeout after {timeout_s}s')
                except BaseException:  # Ctrl+C / kill: reap, then propagate
                    _run(['docker', 'rm', '-f', cname])
                    raise
                if proc.returncode != 125 or attempt == 2:
                    break
                _run(['docker', 'rm', '-f', cname])  # clear any husk; fresh name next try
                time.sleep(2 * (attempt + 1))        # give the daemon a beat
            return ExecResult(
                exit_code=proc.returncode,
                stdout=proc.stdout,
                stderr=proc.stderr,
                duration_s=round(time.time() - t0, 2),
            )


# ---------------- serve-side (model deployment environments) ----------------


def docker_serve(
    engine: str,
    model: str,
    cfg: Dict[str, Any],
    default_image: str,
    engine_args: List[str],
) -> Dict[str, str]:
    """Start (or reuse) an OpenAI-protocol serving container.

    Used through sandbox.acquire() by the model Deployer — the deployment
    container is just another environment this layer provides. Bind-mounts
    the HF cache (weights stay on the host; containers come and go).
    """
    import socket

    image = cfg.get('image', default_image)
    port = int(cfg.get('port', 0))
    if not port:
        with socket.socket() as s:
            s.bind(('', 0))
            port = s.getsockname()[1]
    hf = cfg.get('hf_home', os.environ.get('HF_HOME', '~/.cache/huggingface'))
    gpus = str(cfg.get('gpus', 'all'))
    container = f'evalharness-{engine}-{model}'.replace('/', '-')[:120]

    _run(['docker', 'rm', '-f', container])  # stale instance from a crashed run
    cmd = [
        'docker', 'run', '-d',
        '--name', container,
        '--gpus', f'device={gpus}' if gpus.isdigit() else gpus,
        '--network', 'host',
        '-v', f'{Path(hf).expanduser()}:/root/.cache/huggingface',
        '-e', f'HF_ENDPOINT={os.environ.get("HF_ENDPOINT", "https://hf-mirror.com")}',
        image,
        '--model', cfg.get('model_id', cfg.get('model', model)),
        '--served-model-name', model,
        '--port', str(port),
        *engine_args,
        *shlex.split(cfg.get('extra_args', '')),
    ]
    for key, flag in (('gpu_mem_util', '--gpu-memory-utilization'),
                      ('max_model_len', '--max-model-len'),
                      ('tp_size', '--tensor-parallel-size'),
                      ('dtype', '--dtype')):
        if cfg.get(key):
            cmd += [flag, str(cfg[key])]
    r = _run(cmd)
    if r.returncode != 0:
        raise RuntimeError(f'docker run failed: {r.stderr[:500]}')
    api_base = f'http://localhost:{port}/v1'
    _wait_healthy(api_base, int(cfg.get('timeout_s', 1800)))
    return {'api_base': api_base, 'model': model, 'container': container, 'port': str(port)}


def _wait_healthy(api_base: str, timeout_s: int) -> None:
    import urllib.request

    deadline = time.time() + timeout_s
    while time.time() < deadline:
        try:
            with urllib.request.urlopen(f'{api_base}/models', timeout=5) as resp:
                if resp.status == 200:
                    return
        except Exception:
            time.sleep(5)
    raise TimeoutError(f'serving engine not healthy after {timeout_s}s at {api_base}')


def docker_stop(handle: EnvHandle) -> None:
    if handle.container:
        _run(['docker', 'rm', '-f', handle.container])


def serve_env(engine: str, model: str, cfg: Dict[str, Any],
              default_image: str, engine_args: Optional[List[str]] = None) -> EnvHandle:
    """acquire() wrapper: shared, refcounted serve environment."""
    return acquire(
        kind=engine,
        name=model,
        start_fn=lambda _m: docker_serve(engine, _m, cfg, default_image, engine_args or []),
        stop_fn=docker_stop,
    )