diff --git a/tests/rl/test_sandbox_pool.py b/tests/rl/test_sandbox_pool.py new file mode 100644 index 0000000000..e4cf75b763 --- /dev/null +++ b/tests/rl/test_sandbox_pool.py @@ -0,0 +1,265 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Mapping +from typing import Any + +import pytest +from pydantic import ValidationError + +from xtuner.v1.rl.agent_loop.sandbox_agent_loop.sandbox import SandboxPool +from xtuner.v1.rl.agent_loop.sandbox_agent_loop.schemas import SandboxSpec, StageRecord + + +class FakeClient: + def __init__(self, name: str, events: list[str], *, health_gate: asyncio.Event | None = None): + self.name = name + self.url = f"http://{name}" + self.events = events + self.health_gate = health_gate + + async def health_check(self) -> dict[str, bool]: + self.events.append(f"health:{self.name}") + if self.health_gate is not None: + await self.health_gate.wait() + return {"ok": True} + + async def aclose(self) -> None: + self.events.append(f"close:{self.name}") + + +class FakeProvider: + def __init__( + self, + events: list[str], + *, + fail_image: str | None = None, + health_gate_image: str | None = None, + ) -> None: + self.events = events + self.fail_image = fail_image + self.health_gate_image = health_gate_image + self.health_gate = asyncio.Event() + self.created = 0 + + async def create(self, image_tag: str, ttl_seconds: int, **kwargs: Any) -> tuple[FakeClient, str]: + del ttl_seconds, kwargs + self.events.append(f"create:{image_tag}") + if image_tag == self.fail_image: + raise RuntimeError(f"create failed: {image_tag}") + self.created += 1 + env_id = f"env-{self.created}-{image_tag}" + gate = self.health_gate if image_tag == self.health_gate_image else None + return FakeClient(env_id, self.events, health_gate=gate), env_id + + async def delete(self, env_id: str) -> None: + self.events.append(f"delete:{env_id}") + + +class RecordingProvisioner: + def __init__(self, events: list[str], *, fail_once: bool = False) -> None: + self.events = events + self.fail_once = fail_once + self.calls = 0 + + async def __call__(self, primary: FakeClient, dependencies: Mapping[str, FakeClient]) -> None: + self.calls += 1 + self.events.append(f"provision:{primary.name}:{','.join(dependencies)}") + if self.fail_once and self.calls == 1: + raise RuntimeError("provision failed") + + +def spec(**overrides: Any) -> SandboxSpec: + return SandboxSpec(image="agent", **overrides) + + +def pool(provider: FakeProvider, sandbox_spec: SandboxSpec, **overrides: Any) -> SandboxPool: + return SandboxPool( + provider=provider, + specs={"main": sandbox_spec}, + creates_per_sec=None, + health_poll_interval_sec=0, + **overrides, + ) + + +def test_sandbox_spec_rejects_nested_dependencies_and_dependency_provisioner() -> None: + with pytest.raises(ValidationError, match="must not have dependencies"): + spec(dependencies={"target": spec(dependencies={"db": spec()})}) + + with pytest.raises(ValidationError, match="must not have a provisioner"): + spec(dependencies={"target": spec(provisioner=object())}) + + +@pytest.mark.asyncio +async def test_group_creation_provisioning_primary_api_and_release_order() -> None: + events: list[str] = [] + provider = FakeProvider(events) + provisioner_cfg = {"type": RecordingProvisioner, "events": events} + sandbox_pool = pool( + provider, + spec( + dependencies={ + "db": SandboxSpec(image="database"), + "target": SandboxSpec(image="target"), + }, + provisioner=provisioner_cfg, + ), + ) + + client = await sandbox_pool.get("main") + + assert [event for event in events if event.startswith("create:")] == [ + "create:database", + "create:target", + "create:agent", + ] + assert events.index("health:env-3-agent") < events.index("provision:env-3-agent:db,target") + assert client.name == "env-3-agent" + assert sandbox_pool.env_id("main") == "env-3-agent" + assert sandbox_pool.url("main") == "http://env-3-agent" + with pytest.raises(KeyError, match="unknown sandbox"): + sandbox_pool.validate_name("target") + + await sandbox_pool.release_all() + + assert [event for event in events if event.startswith("delete:")] == [ + "delete:env-3-agent", + "delete:env-2-target", + "delete:env-1-database", + ] + assert [event for event in events if event.startswith("close:")] == [ + "close:env-3-agent", + "close:env-2-target", + "close:env-1-database", + ] + + +@pytest.mark.asyncio +async def test_provision_failure_rolls_back_whole_attempt_before_retry(monkeypatch: pytest.MonkeyPatch) -> None: + events: list[str] = [] + provider = FakeProvider(events) + provisioner = RecordingProvisioner(events, fail_once=True) + + async def no_sleep(_: float) -> None: + return None + + monkeypatch.setattr(asyncio, "sleep", no_sleep) + sandbox_pool = pool( + provider, + spec( + dependencies={"target": SandboxSpec(image="target")}, + provisioner=provisioner, + ), + max_attempts=2, + ) + record = StageRecord() + + client = await sandbox_pool.get("main", record=record) + + assert client.name == "env-4-agent" + assert record.metadata["sandbox_create_attempts"] == 2 + assert [event for event in events if event.startswith("delete:")][:2] == [ + "delete:env-2-agent", + "delete:env-1-target", + ] + assert provisioner.calls == 2 + await sandbox_pool.release_all() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_image", ["target", "agent"]) +async def test_create_failure_rolls_back_only_returned_members(fail_image: str) -> None: + events: list[str] = [] + provider = FakeProvider(events, fail_image=fail_image) + sandbox_pool = pool( + provider, + spec(dependencies={"db": SandboxSpec(image="database"), "target": SandboxSpec(image="target")}), + max_attempts=1, + ) + + with pytest.raises(RuntimeError, match="could not acquire sandbox group"): + await sandbox_pool.get("main") + + created_ids = [event.removeprefix("health:") for event in events if event.startswith("health:")] + assert [event.removeprefix("delete:") for event in events if event.startswith("delete:")] == list( + reversed(created_ids) + ) + + +@pytest.mark.asyncio +async def test_unhealthy_member_rolls_back_the_group(monkeypatch: pytest.MonkeyPatch) -> None: + events: list[str] = [] + provider = FakeProvider(events) + sandbox_pool = pool( + provider, + spec(dependencies={"target": SandboxSpec(image="target")}), + max_attempts=1, + ) + + async def health(client: FakeClient) -> bool: + return not client.name.endswith("target") + + monkeypatch.setattr(sandbox_pool, "_wait_healthy", health) + with pytest.raises(RuntimeError, match="member 'target'.*unhealthy"): + await sandbox_pool.get("main") + + assert [event for event in events if event.startswith("delete:")] == ["delete:env-1-target"] + assert [event for event in events if event.startswith("close:")] == ["close:env-1-target"] + + +@pytest.mark.asyncio +async def test_cancellation_cleans_up_all_returned_members() -> None: + events: list[str] = [] + provider = FakeProvider(events, health_gate_image="agent") + sandbox_pool = pool( + provider, + spec(dependencies={"target": SandboxSpec(image="target")}), + max_attempts=1, + ) + + acquire = asyncio.create_task(sandbox_pool.get("main")) + while "health:env-2-agent" not in events: + await asyncio.sleep(0) + acquire.cancel() + with pytest.raises(asyncio.CancelledError): + await acquire + + assert [event for event in events if event.startswith("delete:")] == [ + "delete:env-2-agent", + "delete:env-1-target", + ] + assert [event for event in events if event.startswith("close:")] == [ + "close:env-2-agent", + "close:env-1-target", + ] + + +@pytest.mark.asyncio +async def test_rate_limiter_is_acquired_for_every_physical_create(monkeypatch: pytest.MonkeyPatch) -> None: + events: list[str] = [] + provider = FakeProvider(events) + + class Limiter: + calls = 0 + + async def acquire(self) -> None: + self.calls += 1 + + limiter = Limiter() + monkeypatch.setattr( + "xtuner.v1.rl.agent_loop.sandbox_agent_loop.sandbox.get_shared_async_token_bucket", + lambda *_args, **_kwargs: limiter, + ) + sandbox_pool = SandboxPool( + provider=provider, + specs={ + "main": spec(dependencies={"db": SandboxSpec(image="database"), "target": SandboxSpec(image="target")}) + }, + creates_per_sec=1.0, + ) + + await sandbox_pool.get("main") + + assert limiter.calls == 3 + await sandbox_pool.release_all() diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/sandbox.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/sandbox.py index d55399df2e..40b4029e01 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/sandbox.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/sandbox.py @@ -826,17 +826,29 @@ async def _sandbox_alive(client: Any, timeout_sec: float = 5.0) -> bool: # ───────────────────────────────────────────────────────────────── +@dataclass +class _SandboxLease: + name: str + client: Any + env_id: str + + class SandboxPool: - """Per-run sandbox client pool: lazily acquires + caches clients by name. + """Per-run sandbox pool with transactional dependency groups. One ``SandboxPool`` instance is created at the start of ``Runner.run`` and released in the ``finally`` block. The pool owns: - - ``provider.create`` / ``provider.delete`` lifecycle - - retry loop with health-check polling on acquire - - the ``client`` / ``env_id`` / ``url`` triplet per name + - every ``provider.create`` / ``provider.delete`` and client close + - rate limiting, health polling, group retry, and reverse-order rollback + - the primary ``client`` / ``env_id`` / ``url`` triplet per public name - sandbox-name validation against the configured spec map + Providers create and delete exactly one physical sandbox per call. They + must not manage dependency groups or close returned clients. An optional + spec provisioner runs only after every member is healthy and must restrict + itself to in-sandbox setup and application-level readiness. + ``get(name, record=...)`` writes the failure to ``record.error`` with ``category="acquire"`` so the caller does not need its own try/except. """ @@ -855,6 +867,10 @@ def __init__( ): self._provider = create_object(provider) self._specs: dict[str, SandboxSpec] = {name: create_object(spec) for name, spec in specs.items()} + self._provisioners = { + name: create_object(spec.provisioner) if spec.provisioner is not None else None + for name, spec in self._specs.items() + } self._max_attempts = max_attempts self._health_max_wait_sec = health_max_wait_sec self._health_poll_interval_sec = health_poll_interval_sec @@ -866,6 +882,7 @@ def __init__( self._clients: dict[str, Any] = {} self._env_ids: dict[str, str] = {} self._urls: dict[str, str | None] = {} + self._groups: dict[str, list[_SandboxLease]] = {} async def get(self, name: str, *, record: StageRecord | None = None) -> Any: if name in self._clients: @@ -873,7 +890,12 @@ async def get(self, name: str, *, record: StageRecord | None = None) -> Any: self.validate_name(name) spec = self._specs[name] try: - client, env_id = await self._acquire_ready(spec, record=record) + group = await self._acquire_ready( + name, + spec, + provisioner=self._provisioners[name], + record=record, + ) except Exception as exc: if record is not None: record.status = StageStatus.FAILED @@ -884,6 +906,10 @@ async def get(self, name: str, *, record: StageRecord | None = None) -> Any: message="".join(traceback.format_exception(type(exc), exc, exc.__traceback__)).rstrip(), ) raise + primary = group[-1] + client = primary.client + env_id = primary.env_id + self._groups[name] = group self._clients[name] = client self._env_ids[name] = env_id self._urls[name] = self._url_of(client) @@ -907,21 +933,11 @@ async def release_all(self) -> None: client = self._clients.pop(name) env_id = self._env_ids.pop(name, None) self._urls.pop(name, None) - if env_id is not None: - try: - await self._provider.delete(env_id) - except Exception as exc: - get_logger().warning( - f"gateway delete failed for sandbox {name} env_id={env_id}:\n" - f"{''.join(traceback.format_exception(type(exc), exc, exc.__traceback__)).rstrip()}" - ) - try: - await client.aclose() - except Exception as exc: - get_logger().warning( - f"client aclose failed for sandbox {name} env_id={env_id}:\n" - f"{''.join(traceback.format_exception(type(exc), exc, exc.__traceback__)).rstrip()}" - ) + group = self._groups.pop(name, None) + if group is None and env_id is not None: + group = [_SandboxLease(name="primary", client=client, env_id=env_id)] + if group: + await self._release_group(name, group) @staticmethod def _url_of(client: Any) -> str | None: @@ -931,69 +947,97 @@ def _url_of(client: Any) -> str | None: return str(val) return None - async def _acquire_ready(self, spec: SandboxSpec, *, record: StageRecord | None = None) -> tuple[Any, str]: + async def _acquire_ready( + self, + name: str, + spec: SandboxSpec, + *, + provisioner: Any | None, + record: StageRecord | None = None, + ) -> list[_SandboxLease]: last_err: Exception | None = None t_ready: float | None = None for attempt in range(1, self._max_attempts + 1): if record is not None: record.metadata["sandbox_create_attempts"] = attempt + group: list[_SandboxLease] = [] try: - create_kwargs: dict[str, Any] = {} - if spec.cluster_name: - create_kwargs["cluster_name"] = spec.cluster_name - if spec.key: - create_kwargs["key"] = spec.key - if spec.env_vars: - create_kwargs["env_vars"] = spec.env_vars - if spec.resources: - create_kwargs["resources"] = spec.resources - if self._create_limiter is not None: - t_limit = time.monotonic() - await self._create_limiter.acquire() - if record is not None: - record.metadata["sandbox_acquire_rate_limit_wait_s"] = ( - record.metadata.get("sandbox_acquire_rate_limit_wait_s", 0.0) + time.monotonic() - t_limit - ) if t_ready is None: t_ready = time.monotonic() - client, env_id = await self._provider.create( - image_tag=spec.image, - ttl_seconds=spec.ttl_seconds, - **create_kwargs, - ) - except Exception as exc: - last_err = exc - await asyncio.sleep(min(2**attempt, 8)) - continue + members = [*spec.dependencies.items(), ("primary", spec)] + for member_name, member_spec in members: + client, env_id = await self._create(member_spec, record=record) + lease = _SandboxLease(name=member_name, client=client, env_id=env_id) + group.append(lease) + if not await self._wait_healthy(client): + raise RuntimeError(f"sandbox group {name!r} member {member_name!r} ({env_id}) unhealthy") + + if provisioner is not None: + dependencies = {lease.name: lease.client for lease in group[:-1]} + await provisioner(group[-1].client, dependencies) - healthy = await self._wait_healthy(client) - if healthy: if record is not None and t_ready is not None: record.metadata["sandbox_create_to_ready_time_s"] = time.monotonic() - t_ready - return client, env_id + return group + except asyncio.CancelledError: + cleanup = asyncio.create_task(self._release_group(name, group)) + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError: + await cleanup + raise + except Exception as exc: + last_err = exc + await self._release_group(name, group) + if attempt < self._max_attempts: + await asyncio.sleep(min(2**attempt, 8)) + last_err_msg = ( + "".join(traceback.format_exception(type(last_err), last_err, last_err.__traceback__)).rstrip() + if last_err is not None + else "unknown" + ) + raise RuntimeError(f"could not acquire sandbox group after {self._max_attempts} attempts: {last_err_msg}") + + async def _create(self, spec: SandboxSpec, *, record: StageRecord | None = None) -> tuple[Any, str]: + create_kwargs: dict[str, Any] = {} + if spec.cluster_name: + create_kwargs["cluster_name"] = spec.cluster_name + if spec.key: + create_kwargs["key"] = spec.key + if spec.env_vars: + create_kwargs["env_vars"] = spec.env_vars + if spec.resources: + create_kwargs["resources"] = spec.resources + if self._create_limiter is not None: + t_limit = time.monotonic() + await self._create_limiter.acquire() + if record is not None: + record.metadata["sandbox_acquire_rate_limit_wait_s"] = ( + record.metadata.get("sandbox_acquire_rate_limit_wait_s", 0.0) + time.monotonic() - t_limit + ) + return await self._provider.create( + image_tag=spec.image, + ttl_seconds=spec.ttl_seconds, + **create_kwargs, + ) + + async def _release_group(self, name: str, group: list[_SandboxLease]) -> None: + for lease in reversed(group): try: - await self._provider.delete(env_id) + await self._provider.delete(lease.env_id) except Exception as exc: get_logger().warning( - f"delete of unhealthy sandbox env_id={env_id} failed:\n" + f"gateway delete failed for sandbox {name} member={lease.name} env_id={lease.env_id}:\n" f"{''.join(traceback.format_exception(type(exc), exc, exc.__traceback__)).rstrip()}" ) try: - await client.aclose() + await lease.client.aclose() except Exception as exc: get_logger().warning( - f"aclose of unhealthy sandbox env_id={env_id} failed:\n" + f"client aclose failed for sandbox {name} member={lease.name} env_id={lease.env_id}:\n" f"{''.join(traceback.format_exception(type(exc), exc, exc.__traceback__)).rstrip()}" ) - last_err = RuntimeError(f"sandbox {env_id} unhealthy") - - last_err_msg = ( - "".join(traceback.format_exception(type(last_err), last_err, last_err.__traceback__)).rstrip() - if last_err is not None - else "unknown" - ) - raise RuntimeError(f"could not acquire a healthy sandbox after {self._max_attempts} attempts: {last_err_msg}") async def _wait_healthy(self, client: Any) -> bool: deadline = time.monotonic() + self._health_max_wait_sec diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/schemas.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/schemas.py index 7065d290a3..069596d1b4 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/schemas.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/schemas.py @@ -18,7 +18,7 @@ from pathlib import Path from typing import Any -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, model_validator # ───────────────────────────────────────────────────────────────── @@ -219,7 +219,15 @@ class AgentRolloutItem(BaseModel): class SandboxSpec(BaseModel): - """Sandbox runtime config shared by infer and isolated-judger sandboxes.""" + """Sandbox runtime config shared by infer and isolated-judger sandboxes. + + A top-level spec describes the primary sandbox returned by + :class:`SandboxPool`. ``dependencies`` are private members of the same + lifecycle group: the pool creates them in declaration order, creates the + primary last, then calls the optional async ``provisioner`` with + ``(primary_client, dependency_clients)``. Dependency specs are deliberately + limited to one level. + """ model_config = ConfigDict(extra="forbid") @@ -230,6 +238,19 @@ class SandboxSpec(BaseModel): resources: dict[str, Any] = Field(default_factory=dict) key: str | None = None cluster_name: str | None = None + dependencies: dict[str, SandboxSpec] = Field(default_factory=dict) + provisioner: Any | None = None + + @model_validator(mode="after") + def validate_dependency_depth(self) -> SandboxSpec: + for name, dependency in self.dependencies.items(): + if not name: + raise ValueError("sandbox dependency names must not be empty") + if dependency.dependencies: + raise ValueError(f"sandbox dependency {name!r} must not have dependencies") + if dependency.provisioner is not None: + raise ValueError(f"sandbox dependency {name!r} must not have a provisioner") + return self class AgentSpec(BaseModel):