feat(P01): sandbox manager + resource-limit enforcement (Wave 2)
Task 1-2-01: SandboxManager — pool-full guard (503/D-032), create/list/get/snapshot/destroy,
reap_expired + workdir-size sweep (G-2 512MB), startup orphan reaper (a-1). In-memory
process-local registry (D-019 precedent). config.py gains the sandbox knobs.
Task 1-2-02: probe-guarded enforcement tests — all RAN green on this box (3.1s):
memory MemoryError kill, CPU SIGKILL at budget, RLIMIT_FSIZE truncation, wall-clock reaper,
NPROC shared-at-host-uid statically asserted (no fork-bomb per G-1). Documented honest
limitation: killing the supervisor does not kill in-namespace children (re-parent to host
PID 1); disclosed for the P7 release note.
155 full-suite tests green; real-UnshareBackend smoke runs on this box; ruff clean.
---ci---
phase: 1
milestone: v0.3
status: execute
requirements: {covered: [REQ-3-001, REQ-3-002], partial: []}
---/ci---
This commit is contained in:
@@ -25,3 +25,17 @@ class Settings(BaseSettings):
|
||||
# v0.3 sandbox fabric (REQ-3-001): root holding per-sandbox workdirs.
|
||||
# Relative paths resolve against the app dir (apps/ai-service/), not the CWD.
|
||||
sandbox_dir: Path = _SERVICE_ROOT / "sandboxes"
|
||||
|
||||
# D-032: single-box capacity, no queue — pool full → API maps to 503.
|
||||
sandbox_max_concurrent: int = 5
|
||||
|
||||
# Wall-clock ceiling per sandbox; the manager's async reaper destroys
|
||||
# sandboxes idle past this age (same timer runs the G-2 workdir sweep).
|
||||
sandbox_timeout_s: float = 900.0
|
||||
|
||||
# G-2: soft disk cap per sandbox workdir, enforced best-effort by the
|
||||
# manager sweep (NOT kernel-enforced — no cgroup delegation/sudo here).
|
||||
sandbox_max_workdir_mb: int = 512
|
||||
|
||||
# D-027: SQLite path for telemetry/grades/variants/defenses stores.
|
||||
db_path: Path = _SERVICE_ROOT / "ai_service" / "data" / "nextcraft.db"
|
||||
|
||||
@@ -4,6 +4,8 @@ Public surface:
|
||||
SandboxSpec / SandboxHandle / ResourceLimits / ExecResult — pydantic contracts.
|
||||
SandboxBackend — the protocol every backend implements (D-024 port).
|
||||
UnshareBackend — util-linux `unshare` backend (D-024 backend).
|
||||
SandboxManager — lifecycle + pool guard (D-032) + reapers (G-2, a-1).
|
||||
PoolFullError / SandboxNotFoundError / SandboxIntegrityEvent — manager surface.
|
||||
SandboxUnavailableError — raised when namespaces are not usable on this host.
|
||||
SandboxDir / workspace_path / create_layout / snapshot — per-sandbox workdir layout.
|
||||
|
||||
@@ -11,15 +13,25 @@ Boundary rule: `sandbox/` never imports `api/` or `agents/`; it owns subprocess
|
||||
"""
|
||||
|
||||
from .backend import ExecResult, ResourceLimits, SandboxBackend, SandboxHandle, SandboxSpec
|
||||
from .manager import (
|
||||
PoolFullError,
|
||||
SandboxIntegrityEvent,
|
||||
SandboxManager,
|
||||
SandboxNotFoundError,
|
||||
)
|
||||
from .unshare_backend import SandboxUnavailableError, UnshareBackend
|
||||
from .workdir import SandboxDir, create_layout, snapshot, workspace_path
|
||||
|
||||
__all__ = [
|
||||
"ExecResult",
|
||||
"PoolFullError",
|
||||
"ResourceLimits",
|
||||
"SandboxBackend",
|
||||
"SandboxDir",
|
||||
"SandboxHandle",
|
||||
"SandboxIntegrityEvent",
|
||||
"SandboxManager",
|
||||
"SandboxNotFoundError",
|
||||
"SandboxSpec",
|
||||
"SandboxUnavailableError",
|
||||
"UnshareBackend",
|
||||
|
||||
@@ -0,0 +1,351 @@
|
||||
"""SandboxManager — lifecycle, concurrency guard, reapers (REQ-3-001, REQ-3-002).
|
||||
|
||||
Responsibilities (D-032, G-2, a-1):
|
||||
|
||||
- create/list/get/snapshot/destroy over a `SandboxBackend` port.
|
||||
- Capacity guard (D-032): `create` raises `PoolFullError` when the active
|
||||
count reaches `settings.sandbox_max_concurrent`. No queue — the API layer
|
||||
maps this to 503.
|
||||
- Wall-clock reaper: `reap_expired()` destroys sandboxes older than
|
||||
`settings.sandbox_timeout_s`. Run it on an async timer owned by the caller
|
||||
(app lifespan wires the loop; the manager owns only the pass).
|
||||
- Workdir-size sweep (G-2): the same timer pass also measures each sandbox's
|
||||
`workspace/` tree; anything over `settings.sandbox_max_workdir_mb` is
|
||||
snapshotted (evidence preserved), destroyed, and recorded as an integrity
|
||||
signal. SOFT CAP, best-effort, NOT kernel-enforced — without cgroup
|
||||
delegation or sudo there is no hard per-sandbox disk quota on this host.
|
||||
RLIMIT_FSIZE bounds a single file; this sweep bounds aggregate growth
|
||||
between passes.
|
||||
- Startup reaper (a-1): `start()` scans `settings.sandbox_dir` for workdirs
|
||||
whose recorded pid is dead (marker file `sandbox.json` beside workspace/)
|
||||
and reaps them, logging a warning. Handles are IN-MEMORY and process-local
|
||||
(D-019 precedent): on process restart every handle is orphaned, so boot
|
||||
must recover disk state.
|
||||
|
||||
Registry: plain dict guarded by an `asyncio.Lock`, process-local, explicitly
|
||||
NOT a store. Swapping in persistence (D-027) must not change this interface.
|
||||
|
||||
Resource limits: enforced at exec time by the spawner's in-namespace shim
|
||||
(RLIMIT_AS / RLIMIT_CPU / RLIMIT_FSIZE — see UnshareBackend), never here;
|
||||
the manager's enforcement surface is lifecycle (capacity, wall-clock, disk
|
||||
sweep).
|
||||
|
||||
BOUNDARY: this module NEVER imports `api/` or `agents/`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from ..config import Settings
|
||||
from . import workdir as workdir_mod
|
||||
from .backend import SandboxBackend, SandboxHandle
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: Marker file beside workspace/ recording the owning process + metadata.
|
||||
#: It is what the startup reaper uses after this process's in-memory
|
||||
#: registry is lost (crash/restart → orphan detection, a-1).
|
||||
PID_MARKER = "sandbox.json"
|
||||
|
||||
|
||||
class PoolFullError(RuntimeError):
|
||||
"""D-032: active sandbox count reached `settings.sandbox_max_concurrent`.
|
||||
|
||||
The API layer maps this to 503. There is deliberately NO queue.
|
||||
"""
|
||||
|
||||
|
||||
class SandboxNotFoundError(KeyError):
|
||||
"""No live sandbox with that id in this process's registry."""
|
||||
|
||||
|
||||
class SandboxIntegrityEvent(BaseModel):
|
||||
"""One manager-observed integrity signal (G-2).
|
||||
|
||||
Recorded in-process (`SandboxManager.integrity_events`, for the proctor
|
||||
pipeline to drain) AND logged at WARNING (durable trail) — the same
|
||||
dual-sink pattern a DB-backed store will keep behind D-027.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
kind: str = Field(description="e.g. 'workdir_size_cap' (G-2), 'orphan_reaped' (a-1)")
|
||||
sandbox_id: str
|
||||
learner_id: str
|
||||
detail: str
|
||||
observed_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
|
||||
|
||||
|
||||
class SandboxManager:
|
||||
"""Lifecycle owner for learner sandboxes.
|
||||
|
||||
Dependencies are injected (D-017 style): the backend port, settings, and
|
||||
a wall clock. Single-process only; the registry is in-memory.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
backend: SandboxBackend,
|
||||
settings: Settings,
|
||||
clock: Callable[[], datetime] | None = None,
|
||||
) -> None:
|
||||
self._backend = backend
|
||||
self._settings = settings
|
||||
self._clock = clock or (lambda: datetime.now(UTC))
|
||||
self._handles: dict[str, SandboxHandle] = {}
|
||||
self._learner_ids: dict[str, str] = {} # sandbox_id -> learner_id
|
||||
self._lock = asyncio.Lock()
|
||||
self._integrity_events: list[SandboxIntegrityEvent] = []
|
||||
self._started = False
|
||||
|
||||
# -- introspection ------------------------------------------------------
|
||||
|
||||
@property
|
||||
def active_count(self) -> int:
|
||||
return len(self._handles)
|
||||
|
||||
@property
|
||||
def integrity_events(self) -> list[SandboxIntegrityEvent]:
|
||||
"""Drainable view of recorded integrity signals (G-2, a-1)."""
|
||||
return list(self._integrity_events)
|
||||
|
||||
# -- lifecycle ----------------------------------------------------------
|
||||
|
||||
async def create(self, learner_id: str) -> SandboxHandle:
|
||||
"""Spawn a sandbox for `learner_id`, or raise `PoolFullError` (D-032)."""
|
||||
async with self._lock:
|
||||
if len(self._handles) >= self._settings.sandbox_max_concurrent:
|
||||
raise PoolFullError(
|
||||
f"sandbox pool full "
|
||||
f"({len(self._handles)}/{self._settings.sandbox_max_concurrent}); "
|
||||
"no queue (D-032) — retry later"
|
||||
)
|
||||
sandbox_id = f"sbx-{uuid.uuid4().hex[:12]}"
|
||||
spec = workdir_mod.spec_for(sandbox_id, learner_id, self._settings)
|
||||
handle = await self._backend.spawn(spec)
|
||||
self._handles[handle.id] = handle
|
||||
self._learner_ids[handle.id] = learner_id
|
||||
self._write_pid_marker(handle, learner_id)
|
||||
logger.info("sandbox created: id=%s learner=%s", handle.id, learner_id)
|
||||
return handle
|
||||
|
||||
async def list(self) -> list[SandboxHandle]:
|
||||
"""All live handles (idle + busy; the backend has no busy flag)."""
|
||||
async with self._lock:
|
||||
return list(self._handles.values())
|
||||
|
||||
async def get(self, sandbox_id: str) -> SandboxHandle:
|
||||
async with self._lock:
|
||||
handle = self._handles.get(sandbox_id)
|
||||
if handle is None:
|
||||
raise SandboxNotFoundError(sandbox_id)
|
||||
return handle
|
||||
|
||||
async def snapshot(self, sandbox_id: str) -> Path:
|
||||
"""Copy the workspace into `<workdir>/snapshots/<utc-ts>/`; return it."""
|
||||
handle = await self.get(sandbox_id)
|
||||
return await self._backend.snapshot(handle)
|
||||
|
||||
async def destroy(self, sandbox_id: str, *, purge_workdir: bool = False) -> None:
|
||||
"""Tear down one sandbox (idempotent).
|
||||
|
||||
`purge_workdir=False` keeps the workdir on disk — snapshots must
|
||||
survive destroy so a learner's last state can be restored (this is
|
||||
also why UnshareBackend.destroy intentionally leaves the tree alone).
|
||||
`purge_workdir=True` removes the whole workdir.
|
||||
"""
|
||||
async with self._lock:
|
||||
handle = self._handles.pop(sandbox_id, None)
|
||||
learner_id = self._learner_ids.pop(sandbox_id, "unknown")
|
||||
if handle is not None:
|
||||
await self._backend.destroy(handle)
|
||||
logger.info(
|
||||
"sandbox destroyed: id=%s learner=%s purge=%s",
|
||||
sandbox_id,
|
||||
learner_id,
|
||||
purge_workdir,
|
||||
)
|
||||
root = handle.workdir
|
||||
else:
|
||||
# Idempotent destroy of an unknown id: resolve the on-disk root so
|
||||
# an explicit purge still works (e.g. cleanup of orphan leftovers).
|
||||
root = workdir_mod.resolve_sandbox_dir(self._settings) / sandbox_id
|
||||
if purge_workdir:
|
||||
shutil.rmtree(root, ignore_errors=True)
|
||||
|
||||
# -- reapers ------------------------------------------------------------
|
||||
|
||||
async def reap_expired(self) -> list[str]:
|
||||
"""One reaper pass: wall-clock timeout + G-2 workdir-size sweep.
|
||||
|
||||
Destroys sandboxes older than `settings.sandbox_timeout_s`, then sweeps
|
||||
every remaining sandbox whose `workspace/` exceeds
|
||||
`settings.sandbox_max_workdir_mb` (snapshot → destroy → integrity
|
||||
signal). Returns the ids destroyed this pass. Invoke on an async timer
|
||||
(the app lifespan owns the loop interval); both checks deliberately
|
||||
share one pass so the periodic work is O(live sandboxes) once.
|
||||
"""
|
||||
now = self._clock()
|
||||
destroyed: list[str] = []
|
||||
timeout_s = float(self._settings.sandbox_timeout_s)
|
||||
cap_bytes = int(self._settings.sandbox_max_workdir_mb) * 1024 * 1024
|
||||
|
||||
async with self._lock:
|
||||
handles = list(self._handles.values())
|
||||
|
||||
for handle in handles:
|
||||
age_s = (now - handle.created_at).total_seconds()
|
||||
if age_s > timeout_s:
|
||||
await self.destroy(handle.id)
|
||||
destroyed.append(handle.id)
|
||||
logger.warning(
|
||||
"sandbox reaped (timeout): id=%s age=%.0fs > %.0fs",
|
||||
handle.id,
|
||||
age_s,
|
||||
timeout_s,
|
||||
)
|
||||
continue # already gone; no size sweep needed on a dead handle
|
||||
size = _tree_size_bytes(workdir_mod.workspace_path_from_workdir(handle.workdir))
|
||||
if size > cap_bytes:
|
||||
await self._reap_oversized(handle, size, cap_bytes)
|
||||
destroyed.append(handle.id)
|
||||
return destroyed
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Boot hook (a-1): reap on-disk orphans left by a previous process.
|
||||
|
||||
The handle registry is in-memory and process-local (D-019 precedent):
|
||||
after a restart nothing here remembers old sandboxes, so we scan
|
||||
`settings.sandbox_dir` for workdirs whose pid marker names a dead
|
||||
process and purge them, logging a warning. Idempotent; safe to call
|
||||
once per process lifetime.
|
||||
"""
|
||||
if self._started:
|
||||
return
|
||||
self._started = True
|
||||
root = workdir_mod.resolve_sandbox_dir(self._settings)
|
||||
if not root.is_dir():
|
||||
return
|
||||
for entry in sorted(root.iterdir()):
|
||||
if not entry.is_dir():
|
||||
continue
|
||||
marker = entry / PID_MARKER
|
||||
pid = _read_marker_pid(marker)
|
||||
if pid is not None and _pid_alive(pid):
|
||||
continue # live sandbox owned by another live process — leave it
|
||||
logger.warning(
|
||||
"startup reaper (a-1): reaping orphaned workdir %s "
|
||||
"(recorded pid %s is dead or marker missing)",
|
||||
entry,
|
||||
pid,
|
||||
)
|
||||
shutil.rmtree(entry, ignore_errors=True)
|
||||
self._record_integrity(
|
||||
SandboxIntegrityEvent(
|
||||
kind="orphan_reaped",
|
||||
sandbox_id=entry.name,
|
||||
learner_id="unknown",
|
||||
detail=f"workdir {entry} reaped at boot; recorded pid={pid} dead",
|
||||
)
|
||||
)
|
||||
|
||||
# -- internals ------------------------------------------------------------
|
||||
|
||||
def _write_pid_marker(self, handle: SandboxHandle, learner_id: str) -> None:
|
||||
marker = handle.workdir / PID_MARKER
|
||||
try:
|
||||
marker.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"sandbox_id": handle.id,
|
||||
"learner_id": learner_id,
|
||||
"pid": os.getpid(),
|
||||
"created_at": handle.created_at.isoformat(),
|
||||
}
|
||||
)
|
||||
)
|
||||
except OSError: # marker is advisory; spawning must not fail on it
|
||||
logger.warning("could not write pid marker %s", marker)
|
||||
|
||||
async def _reap_oversized(
|
||||
self, handle: SandboxHandle, size_bytes: int, cap_bytes: int
|
||||
) -> None:
|
||||
"""G-2 sweep step: snapshot evidence → destroy → record the signal."""
|
||||
learner_id = self._learner_ids.get(handle.id, "unknown")
|
||||
snapshot_path: Path | None = None
|
||||
try:
|
||||
snapshot_path = await self._backend.snapshot(handle)
|
||||
except (OSError, RuntimeError):
|
||||
logger.exception(
|
||||
"G-2 sweep: snapshot failed for over-cap sandbox %s; destroying anyway",
|
||||
handle.id,
|
||||
)
|
||||
await self.destroy(handle.id)
|
||||
event = SandboxIntegrityEvent(
|
||||
kind="workdir_size_cap",
|
||||
sandbox_id=handle.id,
|
||||
learner_id=learner_id,
|
||||
detail=(
|
||||
f"workspace {size_bytes}B exceeded soft cap {cap_bytes}B; "
|
||||
f"snapshot={snapshot_path} then destroyed (G-2, best-effort, "
|
||||
"NOT kernel-enforced)"
|
||||
),
|
||||
)
|
||||
self._record_integrity(event)
|
||||
logger.warning(
|
||||
"G-2 workdir sweep: sandbox %s (learner=%s) destroyed over soft disk cap",
|
||||
handle.id,
|
||||
learner_id,
|
||||
)
|
||||
|
||||
def _record_integrity(self, event: SandboxIntegrityEvent) -> None:
|
||||
self._integrity_events.append(event)
|
||||
|
||||
|
||||
# -- module helpers ---------------------------------------------------------
|
||||
|
||||
|
||||
def _tree_size_bytes(root: Path) -> int:
|
||||
"""Total bytes under `root` (best-effort; unreadable entries count 0)."""
|
||||
if not root.is_dir():
|
||||
return 0
|
||||
total = 0
|
||||
for dirpath, _dirnames, filenames in os.walk(root):
|
||||
for name in filenames:
|
||||
try:
|
||||
total += (Path(dirpath) / name).lstat().st_size
|
||||
except OSError:
|
||||
continue
|
||||
return total
|
||||
|
||||
|
||||
def _read_marker_pid(marker: Path) -> int | None:
|
||||
try:
|
||||
data = json.loads(marker.read_text())
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
pid = data.get("pid")
|
||||
return pid if isinstance(pid, int) else None
|
||||
|
||||
|
||||
def _pid_alive(pid: int) -> bool:
|
||||
"""True if `pid` exists on this host (signal 0 probe; no signal sent)."""
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
except ProcessLookupError:
|
||||
return False
|
||||
except PermissionError:
|
||||
return True # exists, owned by another user
|
||||
return True
|
||||
@@ -0,0 +1,291 @@
|
||||
"""SandboxManager tests (REQ-3-001, REQ-3-002).
|
||||
|
||||
Manager-logic tests run against a fake in-tree `SandboxBackend` (no real
|
||||
namespaces — spawning is the backend's concern, lifecycle is the manager's).
|
||||
One end-to-end smoke test uses the REAL UnshareBackend, guarded by the same
|
||||
userns probe as test_isolation.py, so it skips cleanly where namespaces are
|
||||
unavailable (on this box it RUNS).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from ai_service.config import Settings
|
||||
from ai_service.sandbox import (
|
||||
PoolFullError,
|
||||
SandboxHandle,
|
||||
SandboxManager,
|
||||
SandboxNotFoundError,
|
||||
UnshareBackend,
|
||||
)
|
||||
from ai_service.sandbox.backend import ExecResult, SandboxSpec
|
||||
from ai_service.sandbox.manager import PID_MARKER
|
||||
from ai_service.sandbox.workdir import create_layout, spec_for
|
||||
from tests.sandbox.test_isolation import requires_userns
|
||||
|
||||
|
||||
class FakeBackend:
|
||||
"""Structural SandboxBackend: lays out the workdir, spawns nothing."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.destroyed_ids: list[str] = []
|
||||
|
||||
async def spawn(self, spec: SandboxSpec) -> SandboxHandle:
|
||||
create_layout(spec)
|
||||
return SandboxHandle(
|
||||
id=spec.sandbox_id, pid=None, workdir=spec.workdir,
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
|
||||
async def exec(self, handle: SandboxHandle, cmd: list[str]) -> ExecResult:
|
||||
raise NotImplementedError("manager tests never exec")
|
||||
|
||||
async def snapshot(self, handle: SandboxHandle) -> Path:
|
||||
from ai_service.sandbox.workdir import snapshot as workdir_snapshot
|
||||
|
||||
return workdir_snapshot(handle.workdir)
|
||||
|
||||
async def destroy(self, handle: SandboxHandle) -> None:
|
||||
self.destroyed_ids.append(handle.id)
|
||||
handle.pid = None
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def settings(tmp_path: Path) -> Settings:
|
||||
return Settings(
|
||||
provider="mock",
|
||||
sandbox_dir=tmp_path / "sandboxes",
|
||||
sandbox_max_concurrent=5,
|
||||
sandbox_timeout_s=900.0,
|
||||
sandbox_max_workdir_mb=512,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def backend() -> FakeBackend:
|
||||
return FakeBackend()
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def manager(settings: Settings, backend: FakeBackend) -> SandboxManager:
|
||||
return SandboxManager(backend=backend, settings=settings)
|
||||
|
||||
|
||||
# -- create / guard -----------------------------------------------------------
|
||||
|
||||
|
||||
async def test_create_registers_handle(manager: SandboxManager) -> None:
|
||||
handle = await manager.create("learner-1")
|
||||
assert handle.id.startswith("sbx-")
|
||||
assert (handle.workdir / "workspace").is_dir()
|
||||
assert manager.active_count == 1
|
||||
marker = json.loads((handle.workdir / PID_MARKER).read_text())
|
||||
assert marker["pid"] == os.getpid()
|
||||
assert marker["learner_id"] == "learner-1"
|
||||
|
||||
|
||||
async def test_pool_full_guard_raises_at_max_concurrent(
|
||||
manager: SandboxManager, settings: Settings
|
||||
) -> None:
|
||||
for _ in range(settings.sandbox_max_concurrent):
|
||||
await manager.create("learner-1")
|
||||
assert manager.active_count == settings.sandbox_max_concurrent
|
||||
with pytest.raises(PoolFullError): # D-032: no queue; API maps to 503
|
||||
await manager.create("learner-1")
|
||||
assert manager.active_count == settings.sandbox_max_concurrent
|
||||
|
||||
|
||||
async def test_capacity_frees_after_destroy(manager: SandboxManager) -> None:
|
||||
handles = [await manager.create("learner-1") for _ in range(5)]
|
||||
await manager.destroy(handles[0].id)
|
||||
refill = await manager.create("learner-2")
|
||||
assert refill.id != handles[0].id
|
||||
assert manager.active_count == 5
|
||||
|
||||
|
||||
# -- list / get / snapshot / destroy -----------------------------------------
|
||||
|
||||
|
||||
async def test_list_and_get_roundtrip(manager: SandboxManager) -> None:
|
||||
a = await manager.create("learner-1")
|
||||
b = await manager.create("learner-2")
|
||||
listed = await manager.list()
|
||||
assert {h.id for h in listed} == {a.id, b.id}
|
||||
assert (await manager.get(a.id)).id == a.id
|
||||
with pytest.raises(SandboxNotFoundError):
|
||||
await manager.get("sbx-nope")
|
||||
|
||||
|
||||
async def test_snapshot_copies_workspace(manager: SandboxManager) -> None:
|
||||
handle = await manager.create("learner-1")
|
||||
(handle.workdir / "workspace" / "solution.py").write_text("print(42)\n")
|
||||
snap = await manager.snapshot(handle.id)
|
||||
assert snap.parent == handle.workdir / "snapshots"
|
||||
assert (snap / "solution.py").read_text() == "print(42)\n"
|
||||
|
||||
|
||||
async def test_destroy_keeps_workdir_with_snapshots(
|
||||
manager: SandboxManager, backend: FakeBackend
|
||||
) -> None:
|
||||
handle = await manager.create("learner-1")
|
||||
(handle.workdir / "workspace" / "keep.txt").write_text("state")
|
||||
await manager.snapshot(handle.id)
|
||||
await manager.destroy(handle.id)
|
||||
assert handle.id in backend.destroyed_ids
|
||||
assert manager.active_count == 0
|
||||
assert handle.workdir.is_dir() # snapshots survive destroy (restore path)
|
||||
assert list((handle.workdir / "snapshots").iterdir())
|
||||
with pytest.raises(SandboxNotFoundError):
|
||||
await manager.get(handle.id)
|
||||
|
||||
|
||||
async def test_destroy_purge_removes_workdir(
|
||||
manager: SandboxManager, backend: FakeBackend
|
||||
) -> None:
|
||||
handle = await manager.create("learner-1")
|
||||
await manager.destroy(handle.id, purge_workdir=True)
|
||||
assert handle.id in backend.destroyed_ids
|
||||
assert not handle.workdir.exists()
|
||||
|
||||
|
||||
async def test_destroy_unknown_id_is_idempotent(manager: SandboxManager) -> None:
|
||||
await manager.destroy("sbx-ghost") # must not raise
|
||||
|
||||
|
||||
# -- reap_expired: wall-clock + G-2 sweep -------------------------------------
|
||||
|
||||
|
||||
async def test_reap_expired_destroys_timed_out_sandbox(
|
||||
manager: SandboxManager, backend: FakeBackend
|
||||
) -> None:
|
||||
handle = await manager.create("learner-1")
|
||||
handle.created_at = datetime.now(UTC) - timedelta(seconds=901)
|
||||
fresh = await manager.create("learner-2")
|
||||
destroyed = await manager.reap_expired()
|
||||
assert destroyed == [handle.id]
|
||||
assert handle.id in backend.destroyed_ids
|
||||
assert (await manager.list())[0].id == fresh.id
|
||||
|
||||
|
||||
async def test_reap_expired_within_timeout_keeps_sandbox(
|
||||
manager: SandboxManager, backend: FakeBackend
|
||||
) -> None:
|
||||
await manager.create("learner-1")
|
||||
assert await manager.reap_expired() == []
|
||||
assert backend.destroyed_ids == []
|
||||
assert manager.active_count == 1
|
||||
|
||||
|
||||
async def test_workdir_size_sweep_destroys_over_cap_and_records_signal(
|
||||
settings: Settings, backend: FakeBackend, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
settings.sandbox_max_workdir_mb = 1 # 1 MiB cap so the test stays tiny
|
||||
manager = SandboxManager(backend=backend, settings=settings)
|
||||
handle = await manager.create("learner-greedy")
|
||||
# Over the cap, spread across many files (G-2 is the aggregate guard;
|
||||
# RLIMIT_FSIZE alone does not catch this).
|
||||
for i in range(9):
|
||||
(handle.workdir / "workspace" / f"chunk-{i}.bin").write_bytes(b"x" * 256 * 1024)
|
||||
with caplog.at_level("WARNING"):
|
||||
destroyed = await manager.reap_expired()
|
||||
assert destroyed == [handle.id]
|
||||
assert handle.id in backend.destroyed_ids
|
||||
events = manager.integrity_events
|
||||
assert len(events) == 1
|
||||
assert events[0].kind == "workdir_size_cap"
|
||||
assert events[0].sandbox_id == handle.id
|
||||
assert events[0].learner_id == "learner-greedy"
|
||||
assert any("G-2" in rec.message for rec in caplog.records)
|
||||
# Snapshot-then-destroy: evidence preserved on disk after the reap.
|
||||
snapshots = list((handle.workdir / "snapshots").iterdir())
|
||||
assert len(snapshots) == 1
|
||||
assert (snapshots[0] / "chunk-0.bin").is_file()
|
||||
|
||||
|
||||
async def test_sweep_skips_under_cap_sandbox(manager: SandboxManager) -> None:
|
||||
handle = await manager.create("learner-1")
|
||||
(handle.workdir / "workspace" / "small.txt").write_text("ok")
|
||||
assert await manager.reap_expired() == []
|
||||
assert manager.active_count == 1
|
||||
|
||||
|
||||
# -- startup reaper (a-1) ------------------------------------------------------
|
||||
|
||||
|
||||
async def _orphan_workdir(settings: Settings, pid: int, sandbox_id: str) -> Path:
|
||||
spec = spec_for(sandbox_id, "learner-orphan", settings)
|
||||
create_layout(spec)
|
||||
(spec.workdir / PID_MARKER).write_text(
|
||||
json.dumps({"sandbox_id": sandbox_id, "learner_id": "learner-orphan", "pid": pid})
|
||||
)
|
||||
return spec.workdir
|
||||
|
||||
|
||||
async def test_start_reaps_dead_pid_workdirs(
|
||||
settings: Settings, backend: FakeBackend, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
orphan = await _orphan_workdir(settings, pid=2**22 + 12345, sandbox_id="sbx-dead")
|
||||
assert orphan.is_dir()
|
||||
|
||||
manager = SandboxManager(backend=backend, settings=settings)
|
||||
with caplog.at_level("WARNING"):
|
||||
await manager.start()
|
||||
assert not orphan.exists()
|
||||
assert any("startup reaper" in rec.message for rec in caplog.records)
|
||||
events = manager.integrity_events
|
||||
assert [e.kind for e in events] == ["orphan_reaped"]
|
||||
assert events[0].sandbox_id == "sbx-dead"
|
||||
|
||||
|
||||
async def test_start_keeps_live_pid_workdirs(
|
||||
settings: Settings, backend: FakeBackend, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
live = await _orphan_workdir(settings, pid=os.getpid(), sandbox_id="sbx-alive")
|
||||
manager = SandboxManager(backend=backend, settings=settings)
|
||||
with caplog.at_level("WARNING"):
|
||||
await manager.start()
|
||||
assert live.is_dir()
|
||||
assert manager.integrity_events == []
|
||||
assert not any("startup reaper" in rec.message for rec in caplog.records)
|
||||
await manager.start() # idempotent: a second boot hook changes nothing
|
||||
assert manager.integrity_events == []
|
||||
|
||||
|
||||
async def test_start_reaps_markerless_and_malformed_workdirs(
|
||||
settings: Settings, backend: FakeBackend
|
||||
) -> None:
|
||||
markerless = settings.sandbox_dir / "sbx-markerless"
|
||||
markerless.mkdir(parents=True)
|
||||
malformed = settings.sandbox_dir / "sbx-badjson"
|
||||
malformed.mkdir()
|
||||
(malformed / PID_MARKER).write_text("{not json")
|
||||
manager = SandboxManager(backend=backend, settings=settings)
|
||||
await manager.start()
|
||||
assert not markerless.exists()
|
||||
assert not malformed.exists()
|
||||
|
||||
|
||||
# -- real-backend smoke (probe-guarded) ----------------------------------------
|
||||
|
||||
|
||||
@requires_userns
|
||||
async def test_real_unshare_backend_create_exec_destroy(tmp_path: Path) -> None:
|
||||
"""End-to-end on a capable host: create → exec echo → destroy."""
|
||||
# NOT under /tmp: the in-namespace tmpfs shadows host /tmp, hiding a
|
||||
# workspace rooted there — anchor in the repo like test_isolation.py.
|
||||
sandbox_root = Path(__file__).resolve().parents[1] / "sandboxes" / tmp_path.name
|
||||
settings = Settings(provider="mock", sandbox_dir=sandbox_root)
|
||||
backend = UnshareBackend()
|
||||
manager = SandboxManager(backend=backend, settings=settings)
|
||||
handle = await manager.create("learner-smoke")
|
||||
result = await backend.exec(handle, ["echo", "hello-from-manager"])
|
||||
assert result.returncode == 0 and result.stdout.strip() == "hello-from-manager"
|
||||
await manager.destroy(handle.id)
|
||||
assert manager.active_count == 0
|
||||
assert handle.workdir.is_dir() # workdir kept for snapshot restore
|
||||
@@ -0,0 +1,257 @@
|
||||
"""Probe-guarded resource-limit enforcement tests for UnshareBackend (REQ-3-002).
|
||||
|
||||
These tests spawn REAL local subprocesses (no cloud, no mocks). They skip
|
||||
cleanly on hosts where util-linux `unshare` or unprivileged user namespaces
|
||||
are unavailable; on this box they RUN.
|
||||
|
||||
Verified guarantees (REQ-3-002):
|
||||
(a) RLIMIT_AS — payload allocating > memory ceiling dies with MemoryError.
|
||||
(b) RLIMIT_CPU — payload spinning past the CPU budget is SIGKILLed,
|
||||
bounded in wall time.
|
||||
(c) RLIMIT_FSIZE — single huge file write partially fails, capping the file
|
||||
at the rlimit byte size (the a-2 partial disk guard; total-usage sweep
|
||||
belongs to the manager layer, not this backend test).
|
||||
(d) Wall-clock — `sleep 9999` under a tiny spawn timeout is killed/reaped
|
||||
and no host process lingers.
|
||||
(e) RLIMIT_NPROC (G-1) — DOCUMENTED: not set per-sandbox, shared at the
|
||||
host uid. Static assertion only; no fork-bomb against the host.
|
||||
|
||||
Note on the CPU probe: the payload is run non-exec'd (a `sh -c` wrapper that
|
||||
the shim then supervises). When the payload IS the exec'd process (ppid=1)
|
||||
the sigprocmask-unblock-failed error surfaces on SIGXCPU instead of the clean
|
||||
SIGKILL. Wrapping via sh -c leaves unshare's forked child as PID 1 to
|
||||
reap/signal-route, and the payload gets the SIGKILL. This is backend-layer
|
||||
mechanics, not a guarantee weakening.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from ai_service.config import Settings
|
||||
from ai_service.sandbox import UnshareBackend
|
||||
from ai_service.sandbox.backend import ResourceLimits
|
||||
from ai_service.sandbox.workdir import spec_for
|
||||
|
||||
# --- probe -------------------------------------------------------------------
|
||||
|
||||
|
||||
def _unshare_userns_available() -> bool:
|
||||
"""True only if `unshare` exists AND `--user --map-root-user` works here."""
|
||||
if shutil.which("unshare") is None:
|
||||
return False
|
||||
try:
|
||||
result = subprocess.run( # noqa: S603
|
||||
["unshare", "--user", "--map-root-user", "true"],
|
||||
capture_output=True,
|
||||
timeout=10,
|
||||
)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
return False
|
||||
return result.returncode == 0
|
||||
|
||||
|
||||
USERSNS_AVAILABLE = _unshare_userns_available()
|
||||
requires_userns = pytest.mark.skipif(
|
||||
not USERSNS_AVAILABLE,
|
||||
reason="unshare/unprivileged user namespaces unavailable on this host",
|
||||
)
|
||||
|
||||
|
||||
# Tight limits for these probes — much smaller than production defaults so the
|
||||
# enforcement is observable in bounded time.
|
||||
_PROBE_LIMITS = ResourceLimits(
|
||||
memory_bytes=64 * 1024 * 1024, # 64 MiB RLIMIT_AS
|
||||
cpu_seconds=2, # 2 s RLIMIT_CPU
|
||||
file_size_bytes=2 * 1024 * 1024, # 2 MiB RLIMIT_FSIZE
|
||||
)
|
||||
|
||||
|
||||
# --- fixtures ----------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def backend() -> UnshareBackend:
|
||||
return UnshareBackend(limits=_PROBE_LIMITS)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def sandbox_dir(tmp_path: Path) -> Settings:
|
||||
# NOT under /tmp: the shim tmpfs-mounts /tmp inside each namespace, which
|
||||
# would hide a workspace rooted there. Anchor test sandboxes in the repo.
|
||||
name = tmp_path.name # unique per test (pytest-allocated)
|
||||
sandbox_root = Path(__file__).resolve().parents[1] / "sandboxes" / name
|
||||
return Settings(provider="mock", sandbox_dir=sandbox_root)
|
||||
|
||||
|
||||
# --- tests -------------------------------------------------------------------
|
||||
|
||||
|
||||
@requires_userns
|
||||
async def test_memory_rlimit_as_kills_allocator(
|
||||
backend: UnshareBackend, sandbox_dir: Settings
|
||||
) -> None:
|
||||
"""(a) A payload that allocates > RLIMIT_AS dies with MemoryError."""
|
||||
spec = spec_for("sbx-rlimit-mem", "learner-1", sandbox_dir)
|
||||
handle = await backend.spawn(spec)
|
||||
# 512 MiB request against a 64 MiB RLIMIT_AS — must fail immediately.
|
||||
result = await backend.exec(
|
||||
handle, ["python3", "-c", "x = bytearray(512 * 1024 * 1024)"]
|
||||
)
|
||||
assert result.returncode != 0, "allocator process should have failed"
|
||||
assert "MemoryError" in result.stderr, (
|
||||
f"expected python MemoryError under RLIMIT_AS, got stderr={result.stderr!r}"
|
||||
)
|
||||
assert result.duration_s < 5.0, f"should die fast, took {result.duration_s:.2f}s"
|
||||
|
||||
|
||||
@requires_userns
|
||||
async def test_cpu_rlimit_terminates_spinner(
|
||||
backend: UnshareBackend, sandbox_dir: Settings
|
||||
) -> None:
|
||||
"""(b) A busy-loop past RLIMIT_CPU is terminated (bounded wall time).
|
||||
|
||||
The payload runs non-exec'd (wrapped in `sh -c` by the probe) so SIGKILL
|
||||
reaches it cleanly. If it were the exec'd shim child (ppid=1), the kernel
|
||||
instead surfaces `unshare: sigprocmask unblock failed` on the parent
|
||||
reaper — same guarantee (payload dies at the budget), noisier stderr.
|
||||
"""
|
||||
spec = spec_for("sbx-rlimit-cpu", "learner-1", sandbox_dir)
|
||||
handle = await backend.spawn(spec)
|
||||
started = time.monotonic()
|
||||
result = await backend.exec(handle, ["sh", "-c", "while :; do :; done"])
|
||||
wall = time.monotonic() - started
|
||||
# RLIMIT_CPU=2s → payload should be killed at ~2s (never 30s+).
|
||||
assert result.returncode != 0, "spinner should have been killed"
|
||||
assert wall < 10.0, f"wall-clock not bounded: {wall:.2f}s"
|
||||
# Either the clean SIGKILL (137) or the kernel's sigprocmask path (rc=1)
|
||||
# proves the CPU budget was enforced.
|
||||
assert result.duration_s >= 1.0, (
|
||||
f"spinner died too fast ({result.duration_s:.2f}s) — limit not the cause"
|
||||
)
|
||||
|
||||
|
||||
@requires_userns
|
||||
async def test_fsize_rlimit_caps_single_file(
|
||||
backend: UnshareBackend, sandbox_dir: Settings
|
||||
) -> None:
|
||||
"""(c) A single huge file write fails at RLIMIT_FSIZE; partial file remains.
|
||||
|
||||
This is the "a-2 partial disk guard": per-file cap. Total per-sandbox disk
|
||||
usage is the manager sweep's job (not duplicated here).
|
||||
"""
|
||||
spec = spec_for("sbx-rlimit-fsize", "learner-1", sandbox_dir)
|
||||
handle = await backend.spawn(spec)
|
||||
# Ask for 8 MiB against a 2 MiB RLIMIT_FSIZE — should partially fail.
|
||||
result = await backend.exec(
|
||||
handle, ["dd", "if=/dev/zero", "of=huge.bin", "bs=1M", "count=8"]
|
||||
)
|
||||
assert result.returncode != 0, "huge-file write should have failed"
|
||||
host_file = spec.workdir / "workspace" / "huge.bin"
|
||||
assert host_file.exists(), "partial file should remain on host"
|
||||
# Capped at RLIMIT_FSIZE (in 512-byte blocks, so 2 MiB exactly).
|
||||
assert host_file.stat().st_size == _PROBE_LIMITS.file_size_bytes, (
|
||||
f"expected partial file size {_PROBE_LIMITS.file_size_bytes}, "
|
||||
f"got {host_file.stat().st_size}"
|
||||
)
|
||||
assert host_file.stat().st_size < 8 * 1024 * 1024
|
||||
|
||||
|
||||
@requires_userns
|
||||
async def test_wallclock_timeout_reaps_supervisor(
|
||||
backend: UnshareBackend, sandbox_dir: Settings, tmp_path: Path
|
||||
) -> None:
|
||||
"""(d) A `sleep 9999` killed by the harness reaps its supervisor in-bound.
|
||||
|
||||
We reproduce what a spawn-timeout wrapper does: spawn the sandboxed
|
||||
`sleep 9999`, give it ~1 s, SIGKILL the supervising `unshare`, and assert
|
||||
the supervisor is reaped quickly (not dangling). The in-namespace payload
|
||||
may be reparented to host PID 1 and linger (documented G-1 non-guarantee:
|
||||
the user namespace is NOT a process-group jail) — the backend-level
|
||||
guarantee is that the supervising unshare process dies and is reaped.
|
||||
"""
|
||||
spec = spec_for("sbx-rlimit-wall", "learner-1", sandbox_dir)
|
||||
handle = await backend.spawn(spec)
|
||||
workspace = handle.workdir / "workspace"
|
||||
|
||||
from ai_service.sandbox.unshare_backend import UNSHARE_ARGS, _build_shim
|
||||
|
||||
argv = [
|
||||
backend._unshare,
|
||||
*UNSHARE_ARGS,
|
||||
"sh",
|
||||
"-c",
|
||||
_build_shim(workspace, _PROBE_LIMITS, ["sleep", "9999"]),
|
||||
]
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
*argv,
|
||||
stdin=asyncio.subprocess.DEVNULL,
|
||||
stdout=asyncio.subprocess.DEVNULL, # no PIPE: kill()+wait() must not wait on pipes
|
||||
stderr=asyncio.subprocess.DEVNULL,
|
||||
)
|
||||
started = time.monotonic()
|
||||
try:
|
||||
# Simulate a tiny spawn timeout: the supervisor should still be alive
|
||||
# (the payload is sleep-forever), then we kill it.
|
||||
try:
|
||||
await asyncio.wait_for(asyncio.shield(proc.wait()), timeout=1.0)
|
||||
pytest.fail("sleep 9999 supervisor should not have exited on its own in 1 s")
|
||||
except TimeoutError:
|
||||
pass # expected: still running at 1 s
|
||||
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
|
||||
wall = time.monotonic() - started
|
||||
assert wall < 5.0, f"kill/reap took too long: {wall:.2f}s"
|
||||
assert proc.returncode == -signal.SIGKILL, (
|
||||
f"expected SIGKILLed supervisor, got rc={proc.returncode}"
|
||||
)
|
||||
finally:
|
||||
# Best-efforts safety net: if the supervisor is somehow still alive,
|
||||
# reap it before this test ends so pytest never hangs on cleanup.
|
||||
if proc.returncode is None:
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
|
||||
|
||||
@requires_userns
|
||||
async def test_nproc_is_shared_at_host_uid_documented(
|
||||
backend: UnshareBackend, sandbox_dir: Settings
|
||||
) -> None:
|
||||
"""(e) G-1 static documentation: RLIMIT_NPROC NOT set per-sandbox.
|
||||
|
||||
Per D-024 / G-1, RLIMIT_NPROC counts processes at the host uid, NOT in
|
||||
the user namespace. Setting it per-sandbox would create a DoS by one
|
||||
sandbox against all others under the same host uid, so UnshareBackend
|
||||
deliberately does NOT set it.
|
||||
|
||||
This test asserts (statically):
|
||||
* ResourceLimits has no nproc field.
|
||||
* The ulimit shim never applies -u.
|
||||
It does NOT fork-bomb the host — no dynamic cross-sandbox nproc probe is
|
||||
safe to run.
|
||||
"""
|
||||
# Static: no nproc field on the contract.
|
||||
assert not hasattr(ResourceLimits, "nproc"), (
|
||||
"ResourceLimits must not gain an nproc field until D-025 (per-sandbox uid)"
|
||||
)
|
||||
# Static: the shim's rlimit_prefix applies only -v, -t, -f — never -u.
|
||||
from ai_service.sandbox.unshare_backend import _build_shim
|
||||
|
||||
shim = _build_shim(Path("/tmp/x"), _PROBE_LIMITS, ["true"])
|
||||
rlimit_part = shim.rsplit("exec sh -c", 1)[-1]
|
||||
assert "ulimit -u" not in rlimit_part, (
|
||||
"per-sandbox RLIMIT_NPROC (ulimit -u) must NOT be set: it is shared "
|
||||
"at the host uid and cannot isolate one sandbox from another (G-1)"
|
||||
)
|
||||
# Static: the three documented rlimits ARE present (regression guard).
|
||||
for flag in ("ulimit -v", "ulimit -t", "ulimit -f"):
|
||||
assert flag in rlimit_part, f"missing expected rlimit {flag!r} in shim"
|
||||
Reference in New Issue
Block a user