mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-27 08:28:00 +00:00
fix(sandbox): reconcile E2B sandboxes safely (#4443)
* fix(sandbox): reconcile E2B sandboxes safely * fix(sandbox): clear failed E2B adoption intent
This commit is contained in:
parent
1baa8ad696
commit
b22f85c686
@ -288,7 +288,7 @@ DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser
|
||||
>
|
||||
> With lease heartbeat enabled, a transient RunStore renewal error is retried only until the last confirmed lease expires; the stale worker then cancels local execution and suppresses checkpoint, completion-hook, delivery-receipt, and thread-status finalization. A remote tool side effect already in flight may still be outside local cancellation.
|
||||
>
|
||||
> Reconciliation uses an atomic takeover claim that re-checks the lease after candidate selection, so a successful owner renewal wins over orphan recovery and only one reconciler can report a run as recovered.
|
||||
> Reconciliation uses an atomic takeover claim that re-checks the lease after candidate selection, so a successful owner renewal wins over orphan recovery and only one reconciler can report a run as recovered. When multiple Gateway workers share the Docker/AIO or E2B sandbox backend, also configure `sandbox.ownership.type: redis`; E2B uses the leases during background startup and periodic reconciliation so duplicate/orphan cleanup cannot terminate a live peer's sandbox.
|
||||
|
||||
See [CONTRIBUTING.md](CONTRIBUTING.md) for detailed Docker development guide.
|
||||
|
||||
|
||||
@ -547,7 +547,15 @@ copying a raw checkpoint because delta state is not self-contained in one tuple.
|
||||
failed kill through a new client. An unconfirmed remote ID stays tracked.
|
||||
`reset()` uses full shutdown semantics. It destroys tracked active and warm
|
||||
VMs. It wakes capacity waiters. Callers cannot reuse the old provider
|
||||
instance.
|
||||
instance. A background startup pass and periodic reconciliation list
|
||||
provider-tagged remote sandboxes within page/item/time budgets, probe every
|
||||
candidate until a healthy canonical sandbox is found, adopt only through the
|
||||
shared ownership store, and reap duplicates/orphans only after their
|
||||
configured grace/TTL and an atomic `del:` claim. Lease renewal is independent
|
||||
of reconciliation. Failed canonical adoption clears both its capacity
|
||||
reservation and acquire-intent marker even if a peer takes ownership between
|
||||
the initial claim and bootstrap cleanup. Shutdown kills only IDs whose leases
|
||||
are owned by this provider instance; peer-owned clients are merely closed.
|
||||
- **Cross-instance ownership store** (`aio_sandbox/ownership/`, #4206): gateway instances sharing a container backend coordinate container ownership through a pluggable lease store, selected by `sandbox.ownership.type` (`memory` | `redis`) and resolved like `stream_bridge` (`factory.py`, lazy per-branch import, `redis` optional extra, `DEER_FLOW_SANDBOX_OWNERSHIP_REDIS_URL` env escape hatch; a set `DEER_FLOW_STREAM_BRIDGE_REDIS_URL` implies a multi-instance deployment and infers `redis`). `memory` is single-instance only and declares `supports_cross_process = False`.
|
||||
- **A lease answers "who reaps this container", not "who may use it".** That splits the interface in two: `take()` transfers ownership on the **acquire** path (a container is deterministic per user/thread, so consecutive turns legitimately land on different instances — a conditional claim there would strand the thread until the previous lease expired), while `claim()` succeeds only if the container is unowned or already ours and gates every **adopt/reap** path. `release()` never clears a peer's lease.
|
||||
- **A lease carries a state, and that is what makes the destroy window safe.** `own:` = responsible for this container; `del:` = tearing it down (`claim(..., for_destroy=True)`). `take()` is refused against a `del:` lease, so a container cannot be re-acquired between a destroy path's claim and its container stop. Without the two states an unconditional `take()` would silently overwrite the destroyer's claim and the peer's stop would land on a container the new owner had already handed to an agent — i.e. #4206 again. That pairing is what replaced the previous same-host `flock` guard, which is gone; Redis makes the scope genuinely multi-instance instead of same-host. A destroyer that dies mid-stop leaves a `del:` marker that lapses with the TTL. On the acquire path a refused take raises `SandboxBeingDestroyedError`: the reuse/reclaim paths drop the container and cold-start, and the discover path propagates (falling through to create would collide with the not-yet-removed container name).
|
||||
|
||||
@ -426,6 +426,15 @@ sandbox:
|
||||
home_dir: /home/user # /mnt/user-data is remapped under this directory
|
||||
idle_timeout: 600 # forwarded to e2b's server-side set_timeout()
|
||||
replicas: 3 # max concurrent sandboxes per gateway process
|
||||
ownership: # use Redis when more than one gateway shares E2B
|
||||
type: redis
|
||||
redis_url: $REDIS_URL
|
||||
reconciliation_interval_seconds: 60
|
||||
reconciliation_grace_seconds: 120
|
||||
reconciliation_orphan_ttl_seconds: 3600
|
||||
reconciliation_max_pages: 10
|
||||
reconciliation_max_items: 200
|
||||
reconciliation_max_seconds: 15
|
||||
mounts: # one-shot upload of host files at sandbox start
|
||||
- host_path: /path/on/host
|
||||
container_path: /home/user/shared
|
||||
@ -440,10 +449,19 @@ provider in `config.yaml`.
|
||||
|
||||
Notes specific to `E2BSandboxProvider`:
|
||||
|
||||
- Each DeerFlow thread is bound to its e2b sandbox via metadata
|
||||
(`deer_flow_user`, `deer_flow_thread`), so the same thread reuses the same
|
||||
sandbox across gateway restarts and across processes — no cross-process
|
||||
file lock is needed because the e2b control plane is the source of truth.
|
||||
- Each DeerFlow thread is bound to its E2B sandbox via metadata
|
||||
(`deer_flow_user`, `deer_flow_thread`). Startup and periodic reconciliation
|
||||
probe every bounded candidate, adopt one healthy canonical sandbox, and reap
|
||||
duplicates after a grace period. Provider-tagged entries without a complete
|
||||
user/thread identity are reaped only after the orphan TTL.
|
||||
- Ownership leases prevent one gateway from adopting or destroying a sandbox
|
||||
another live gateway is responsible for. The default in-memory store is safe
|
||||
only for one gateway process. Multi-worker/load-balanced deployments must use
|
||||
`sandbox.ownership.type: redis`; an existing Redis stream bridge configuration
|
||||
is inferred automatically.
|
||||
- Reconciliation is bounded by page, item, and wall-clock limits. Its summary log
|
||||
exposes discovered, adopted, duplicate, deferred, killed, dead, and budget-exhausted
|
||||
counts for operational monitoring.
|
||||
- Idle expiry is enforced server-side by e2b's `set_timeout()`. The provider
|
||||
refreshes the timeout on every release so warm sandboxes stay alive long
|
||||
enough for the next acquire.
|
||||
|
||||
@ -13,6 +13,15 @@ Configuration example (``config.yaml``)::
|
||||
domain: e2b.dev # optional e2b domain (e.g. self-hosted)
|
||||
idle_timeout: 600 # forwarded to e2b ``set_timeout`` (seconds)
|
||||
replicas: 3 # max concurrent sandboxes (LRU eviction beyond)
|
||||
ownership: # required for safe multi-worker reconciliation
|
||||
type: redis
|
||||
redis_url: $REDIS_URL
|
||||
reconciliation_interval_seconds: 60
|
||||
reconciliation_grace_seconds: 120
|
||||
reconciliation_orphan_ttl_seconds: 3600
|
||||
reconciliation_max_pages: 10
|
||||
reconciliation_max_items: 200
|
||||
reconciliation_max_seconds: 15
|
||||
mounts: # one-shot upload of host files into the sandbox
|
||||
- host_path: /path/on/host
|
||||
container_path: /path/in/sandbox
|
||||
|
||||
@ -38,6 +38,7 @@ import time
|
||||
import uuid
|
||||
from collections import OrderedDict
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from decimal import Decimal, InvalidOperation
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
@ -51,6 +52,14 @@ from deerflow.sandbox.exceptions import SandboxCapacityExceededError
|
||||
from deerflow.sandbox.sandbox import Sandbox
|
||||
from deerflow.sandbox.sandbox_provider import SandboxProvider
|
||||
|
||||
from ..aio_sandbox.ownership import (
|
||||
OwnershipBackendError,
|
||||
RenewOutcome,
|
||||
SandboxOwnershipStore,
|
||||
generate_owner_id,
|
||||
make_sandbox_ownership_store,
|
||||
resolve_ownership_config,
|
||||
)
|
||||
from .e2b_sandbox import DEFAULT_E2B_HOME_DIR, E2BSandbox, _is_sandbox_gone_error
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@ -62,6 +71,12 @@ DEFAULT_IDLE_TIMEOUT = 1800 # 30 minutes; passed to ``Sandbox.set_timeout``.
|
||||
DEFAULT_REPLICAS = 3
|
||||
DEFAULT_OVERFLOW_POLICY = "wait" # wait | reject | burst
|
||||
DEFAULT_ACQUIRE_TIMEOUT = 30 # seconds for wait policy
|
||||
DEFAULT_RECONCILIATION_INTERVAL_SECONDS = 60.0
|
||||
DEFAULT_RECONCILIATION_GRACE_SECONDS = 120.0
|
||||
DEFAULT_RECONCILIATION_ORPHAN_TTL_SECONDS = 3600.0
|
||||
DEFAULT_RECONCILIATION_MAX_PAGES = 10
|
||||
DEFAULT_RECONCILIATION_MAX_ITEMS = 200
|
||||
DEFAULT_RECONCILIATION_MAX_SECONDS = 15.0
|
||||
# Hard upper bound for ``set_timeout`` (e2b currently caps at 24h on the
|
||||
# free plan; passing an excessive value is rejected by the control-plane).
|
||||
MAX_E2B_TIMEOUT = 24 * 60 * 60
|
||||
@ -71,10 +86,25 @@ MAX_E2B_TIMEOUT = 24 * 60 * 60
|
||||
META_KEY_USER = "deer_flow_user"
|
||||
META_KEY_THREAD = "deer_flow_thread"
|
||||
META_KEY_PROVIDER = "deer_flow_provider"
|
||||
META_KEY_GATEWAY = "deer_flow_gateway"
|
||||
META_KEY_CREATED_AT = "deer_flow_created_at"
|
||||
META_VAL_PROVIDER = "e2b_sandbox_provider"
|
||||
E2B_EXTRA_CONFIG_KEYS = frozenset({"api_key", "domain", "home_dir", "template"})
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReconciliationStats:
|
||||
"""Bounded reconciliation outcome, also suitable for metrics/logging."""
|
||||
|
||||
discovered: int = 0
|
||||
adopted: int = 0
|
||||
duplicates: int = 0
|
||||
deferred: int = 0
|
||||
killed: int = 0
|
||||
dead: int = 0
|
||||
budget_exhausted: bool = False
|
||||
|
||||
|
||||
class E2BSandboxProvider(SandboxProvider):
|
||||
"""Sandbox provider backed by the e2b code-interpreter cloud SDK."""
|
||||
|
||||
@ -116,6 +146,13 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._transitioning_slots = 0
|
||||
self._capacity_cond = threading.Condition(self._lock)
|
||||
self._shutdown_called = False
|
||||
self._owned_sandbox_ids: set[str] = set()
|
||||
self._acquire_inflight: set[str] = set()
|
||||
self._orphan_first_seen: dict[str, float] = {}
|
||||
self._maintenance_stop = threading.Event()
|
||||
self._lease_thread: threading.Thread | None = None
|
||||
self._reconcile_thread: threading.Thread | None = None
|
||||
self._owner_id = generate_owner_id()
|
||||
|
||||
self._config = self._load_config()
|
||||
acquire_workers = max(4, min(32, self._capacity_limit() + 1))
|
||||
@ -123,9 +160,20 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
max_workers=acquire_workers,
|
||||
thread_name_prefix="e2b-sandbox-acquire",
|
||||
)
|
||||
self._ownership_config = resolve_ownership_config(
|
||||
self._config.get("ownership"),
|
||||
stream_bridge=self._config.get("stream_bridge"),
|
||||
)
|
||||
self._ownership: SandboxOwnershipStore = make_sandbox_ownership_store(
|
||||
self._ownership_config,
|
||||
owner_id=self._owner_id,
|
||||
)
|
||||
if not self._ownership.supports_cross_process:
|
||||
logger.warning("E2B sandbox ownership is process-local. Multi-worker gateways must configure sandbox.ownership.type: redis for safe reconciliation.")
|
||||
|
||||
atexit.register(self.shutdown)
|
||||
self._register_signal_handlers()
|
||||
self._start_maintenance_threads()
|
||||
|
||||
def _load_config(self) -> dict[str, Any]:
|
||||
"""Read e2b options off ``SandboxConfig`` (``extra="allow"``)."""
|
||||
@ -181,6 +229,32 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
"burst_limit": burst_limit,
|
||||
"mounts": _opt("mounts") or [],
|
||||
"environment": self._resolve_env_vars(_opt("environment") or {}),
|
||||
"ownership": _opt("ownership"),
|
||||
"stream_bridge": getattr(get_app_config(), "stream_bridge", None),
|
||||
"reconciliation_interval_seconds": max(
|
||||
1.0,
|
||||
float(_opt("reconciliation_interval_seconds", DEFAULT_RECONCILIATION_INTERVAL_SECONDS)),
|
||||
),
|
||||
"reconciliation_grace_seconds": max(
|
||||
0.0,
|
||||
float(_opt("reconciliation_grace_seconds", DEFAULT_RECONCILIATION_GRACE_SECONDS)),
|
||||
),
|
||||
"reconciliation_orphan_ttl_seconds": max(
|
||||
0.0,
|
||||
float(_opt("reconciliation_orphan_ttl_seconds", DEFAULT_RECONCILIATION_ORPHAN_TTL_SECONDS)),
|
||||
),
|
||||
"reconciliation_max_pages": max(
|
||||
1,
|
||||
int(_opt("reconciliation_max_pages", DEFAULT_RECONCILIATION_MAX_PAGES)),
|
||||
),
|
||||
"reconciliation_max_items": max(
|
||||
1,
|
||||
int(_opt("reconciliation_max_items", DEFAULT_RECONCILIATION_MAX_ITEMS)),
|
||||
),
|
||||
"reconciliation_max_seconds": max(
|
||||
0.1,
|
||||
float(_opt("reconciliation_max_seconds", DEFAULT_RECONCILIATION_MAX_SECONDS)),
|
||||
),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
@ -322,6 +396,9 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._refresh_remote_timeout(sandbox.client)
|
||||
except Exception as e: # pragma: no cover - defensive
|
||||
logger.debug("Failed to refresh timeout on reuse: %s", e)
|
||||
self._publish_ownership(sid)
|
||||
with self._lock:
|
||||
self._acquire_inflight.discard(sid)
|
||||
|
||||
logger.info(
|
||||
"Reusing in-process e2b sandbox %s for user/thread %s/%s",
|
||||
@ -370,13 +447,17 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._complete_transition_remote_op(target_id, remote_destroyed=True)
|
||||
return None
|
||||
|
||||
try:
|
||||
self._publish_ownership(target_id)
|
||||
except Exception:
|
||||
self._complete_transition_remote_op(target_id, remote_destroyed=False)
|
||||
self._safe_close_client(client)
|
||||
raise
|
||||
|
||||
self._refresh_remote_timeout(client)
|
||||
bootstrap_error, remote_destroyed = self._bootstrap_or_discard(client, target_id)
|
||||
if bootstrap_error is not None:
|
||||
if remote_destroyed:
|
||||
self._complete_transition_remote_op(target_id, remote_destroyed=True)
|
||||
else:
|
||||
self._complete_transition_remote_op(target_id, remote_destroyed=False)
|
||||
self._complete_transition_remote_op(target_id, remote_destroyed=remote_destroyed)
|
||||
return None
|
||||
|
||||
discard_after_shutdown = False
|
||||
@ -393,10 +474,11 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._end_transition_locked()
|
||||
|
||||
if discard_after_shutdown:
|
||||
self._kill_client(client)
|
||||
if self._claim_ownership(target_id, for_destroy=True):
|
||||
self._kill_client(client)
|
||||
self._release_ownership(target_id)
|
||||
self._safe_close_client(client)
|
||||
return None
|
||||
|
||||
logger.info(
|
||||
"Reclaimed warm-pool e2b sandbox %s for user/thread %s/%s",
|
||||
target_id,
|
||||
@ -414,81 +496,39 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
"""
|
||||
sandbox_cls = self._get_sandbox_cls()
|
||||
seed = self._stable_seed(thread_id, user_id)
|
||||
list_kwargs = self._common_kwargs()
|
||||
try:
|
||||
running = sandbox_cls.list( # type: ignore[attr-defined]
|
||||
query={
|
||||
"metadata": {
|
||||
META_KEY_PROVIDER: META_VAL_PROVIDER,
|
||||
META_KEY_USER: user_id,
|
||||
META_KEY_THREAD: thread_id,
|
||||
}
|
||||
},
|
||||
**list_kwargs,
|
||||
entries, _ = self._list_remote_entries(
|
||||
{
|
||||
META_KEY_PROVIDER: META_VAL_PROVIDER,
|
||||
META_KEY_USER: user_id,
|
||||
META_KEY_THREAD: thread_id,
|
||||
}
|
||||
)
|
||||
candidates = sorted(
|
||||
((sandbox_id, metadata) for entry in entries if (sandbox_id := self._entry_id(entry)) and (metadata := self._entry_metadata(entry)).get(META_KEY_USER) == user_id and metadata.get(META_KEY_THREAD) == thread_id),
|
||||
key=lambda item: (item[1].get(META_KEY_CREATED_AT, ""), item[0]),
|
||||
)
|
||||
for target_id, _metadata in candidates:
|
||||
adopted = self._adopt_remote_candidate(
|
||||
sandbox_cls,
|
||||
target_id,
|
||||
thread_id=thread_id,
|
||||
user_id=user_id,
|
||||
seed=seed,
|
||||
)
|
||||
except TypeError:
|
||||
try:
|
||||
running = sandbox_cls.list(
|
||||
metadata={
|
||||
META_KEY_PROVIDER: META_VAL_PROVIDER,
|
||||
META_KEY_USER: user_id,
|
||||
META_KEY_THREAD: thread_id,
|
||||
},
|
||||
**list_kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug("e2b Sandbox.list() unavailable, skipping discovery: %s", e)
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
"e2b Sandbox.list() raised while discovering thread %s: %s",
|
||||
thread_id,
|
||||
e,
|
||||
)
|
||||
return None
|
||||
if adopted is not None:
|
||||
return adopted
|
||||
return None
|
||||
|
||||
# Pick the first matching candidate; tolerate either ``SandboxInfo``
|
||||
# objects with ``sandbox_id`` or plain dicts.
|
||||
# Normalise the return value of ``Sandbox.list()``:
|
||||
# * Older SDKs (<= 1.x) returned a plain ``list[SandboxInfo]`` — directly iterable.
|
||||
# * e2b-code-interpreter >= 2.x returns a ``SandboxPaginator`` exposing
|
||||
# ``has_next: bool`` and ``next_items() -> list[SandboxInfo]`` instead
|
||||
# of being iterable. Walking pages keeps discovery correct when the
|
||||
# org has more sandboxes than fit in a single page.
|
||||
def _iter_running(obj):
|
||||
if obj is None:
|
||||
return
|
||||
if hasattr(obj, "next_items") and hasattr(obj, "has_next"):
|
||||
for _ in range(50):
|
||||
try:
|
||||
page = obj.next_items()
|
||||
except Exception as exc:
|
||||
logger.debug("SandboxPaginator.next_items() failed: %s", exc)
|
||||
return
|
||||
if not page:
|
||||
return
|
||||
yield from page
|
||||
if not getattr(obj, "has_next", False):
|
||||
return
|
||||
return
|
||||
try:
|
||||
yield from obj
|
||||
except TypeError:
|
||||
logger.debug("Sandbox.list() returned non-iterable %s; ignoring", type(obj).__name__)
|
||||
|
||||
target_id: str | None = None
|
||||
for entry in _iter_running(running):
|
||||
sid = getattr(entry, "sandbox_id", None) or (entry.get("sandbox_id") if isinstance(entry, dict) else None)
|
||||
metadata = getattr(entry, "metadata", None) or (entry.get("metadata") if isinstance(entry, dict) else {}) or {}
|
||||
if metadata.get(META_KEY_USER) != user_id:
|
||||
continue
|
||||
if metadata.get(META_KEY_THREAD) != thread_id:
|
||||
continue
|
||||
target_id = sid
|
||||
break
|
||||
|
||||
if not target_id:
|
||||
return None
|
||||
def _adopt_remote_candidate(
|
||||
self,
|
||||
sandbox_cls: type[E2BClientSandbox],
|
||||
target_id: str,
|
||||
*,
|
||||
thread_id: str,
|
||||
user_id: str,
|
||||
seed: str,
|
||||
) -> str | None:
|
||||
"""Try to adopt one discovered candidate without harming peer-owned VMs."""
|
||||
|
||||
try:
|
||||
client = self._reconnect_live_client(sandbox_cls, target_id)
|
||||
@ -528,6 +568,20 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._safe_close_client(client)
|
||||
raise
|
||||
|
||||
with self._lock:
|
||||
shutdown_before_ownership = self._shutdown_called
|
||||
if shutdown_before_ownership:
|
||||
self._complete_reserved_remote_op(target_id, remote_destroyed=False)
|
||||
self._safe_close_client(client)
|
||||
return None
|
||||
|
||||
try:
|
||||
self._publish_ownership(target_id)
|
||||
except Exception:
|
||||
self._complete_reserved_remote_op(target_id, remote_destroyed=False)
|
||||
self._safe_close_client(client)
|
||||
raise
|
||||
|
||||
self._refresh_remote_timeout(client)
|
||||
bootstrap_error, remote_destroyed = self._bootstrap_or_discard(client, target_id)
|
||||
if bootstrap_error is not None:
|
||||
@ -543,6 +597,9 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._register_connected_sandbox(target_id, client, thread_id=thread_id, user_id=user_id)
|
||||
self._commit_capacity()
|
||||
if discard_after_shutdown:
|
||||
if self._claim_ownership(target_id, for_destroy=True):
|
||||
self._kill_client(client)
|
||||
self._release_ownership(target_id)
|
||||
self._safe_close_client(client)
|
||||
return None
|
||||
logger.info(
|
||||
@ -787,6 +844,8 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
sandbox_cls = self._get_sandbox_cls()
|
||||
metadata: dict[str, str] = {
|
||||
META_KEY_PROVIDER: META_VAL_PROVIDER,
|
||||
META_KEY_GATEWAY: self._owner_id,
|
||||
META_KEY_CREATED_AT: str(time.time()),
|
||||
}
|
||||
if thread_id:
|
||||
metadata[META_KEY_USER] = user_id
|
||||
@ -849,6 +908,17 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
reason="shutdown",
|
||||
)
|
||||
|
||||
try:
|
||||
self._publish_ownership(sandbox_id)
|
||||
except Exception:
|
||||
remote_destroyed = False
|
||||
if self._claim_ownership(sandbox_id, for_destroy=True):
|
||||
remote_destroyed = self._kill_client(client) is None
|
||||
self._release_ownership(sandbox_id)
|
||||
self._safe_close_client(client)
|
||||
self._complete_reserved_remote_op(sandbox_id, remote_destroyed=remote_destroyed)
|
||||
raise
|
||||
|
||||
# Materialise DeerFlow's virtual path layout (/mnt/user-data/...) inside
|
||||
# the e2b VM. Without this step shell commands the agent emits — which
|
||||
# use the same /mnt/user-data prefix as LocalSandbox / AioSandbox — fail
|
||||
@ -883,7 +953,9 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._thread_sandboxes[self._thread_key(thread_id, user_id)] = sandbox_id
|
||||
|
||||
if should_kill:
|
||||
self._kill_client(client)
|
||||
if self._claim_ownership(sandbox_id, for_destroy=True):
|
||||
self._kill_client(client)
|
||||
self._release_ownership(sandbox_id)
|
||||
self._safe_close_client(client)
|
||||
raise SandboxCapacityExceededError(
|
||||
f"Sandbox provider shut down during sandbox creation; killed remote sandbox {sandbox_id}",
|
||||
@ -912,6 +984,339 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
kwargs["domain"] = self._config["domain"]
|
||||
return kwargs
|
||||
|
||||
@staticmethod
|
||||
def _entry_id(entry: Any) -> str | None:
|
||||
value = getattr(entry, "sandbox_id", None)
|
||||
if value is None and isinstance(entry, dict):
|
||||
value = entry.get("sandbox_id")
|
||||
return str(value) if value else None
|
||||
|
||||
@staticmethod
|
||||
def _entry_metadata(entry: Any) -> dict[str, Any]:
|
||||
value = getattr(entry, "metadata", None)
|
||||
if value is None and isinstance(entry, dict):
|
||||
value = entry.get("metadata")
|
||||
return value if isinstance(value, dict) else {}
|
||||
|
||||
def _list_remote_entries(self, metadata: dict[str, str]) -> tuple[list[Any], bool]:
|
||||
"""List matching E2B entries within configured page/item/time budgets."""
|
||||
sandbox_cls = self._get_sandbox_cls()
|
||||
try:
|
||||
result = sandbox_cls.list(query={"metadata": metadata}, **self._common_kwargs()) # type: ignore[attr-defined]
|
||||
except TypeError:
|
||||
try:
|
||||
result = sandbox_cls.list(metadata=metadata, **self._common_kwargs()) # type: ignore[attr-defined]
|
||||
except Exception as e:
|
||||
logger.warning("E2B reconciliation list failed: %s", e)
|
||||
return [], False
|
||||
except Exception as e:
|
||||
logger.warning("E2B reconciliation list failed: %s", e)
|
||||
return [], False
|
||||
|
||||
max_pages = int(self._config["reconciliation_max_pages"])
|
||||
max_items = int(self._config["reconciliation_max_items"])
|
||||
deadline = time.monotonic() + float(self._config["reconciliation_max_seconds"])
|
||||
entries: list[Any] = []
|
||||
exhausted = False
|
||||
|
||||
if hasattr(result, "next_items") and hasattr(result, "has_next"):
|
||||
for page_number in range(max_pages):
|
||||
if time.monotonic() >= deadline or len(entries) >= max_items:
|
||||
exhausted = True
|
||||
break
|
||||
try:
|
||||
page = result.next_items()
|
||||
except Exception as e:
|
||||
logger.warning("E2B reconciliation paginator failed: %s", e)
|
||||
break
|
||||
if not page:
|
||||
break
|
||||
room = max_items - len(entries)
|
||||
entries.extend(list(page)[:room])
|
||||
if len(page) > room:
|
||||
exhausted = True
|
||||
break
|
||||
if not getattr(result, "has_next", False):
|
||||
break
|
||||
if page_number + 1 >= max_pages:
|
||||
exhausted = True
|
||||
else:
|
||||
try:
|
||||
all_entries = list(result or [])
|
||||
except TypeError:
|
||||
logger.warning("E2B Sandbox.list returned non-iterable %s", type(result).__name__)
|
||||
return [], False
|
||||
entries = all_entries[:max_items]
|
||||
exhausted = len(all_entries) > max_items
|
||||
|
||||
if time.monotonic() >= deadline:
|
||||
exhausted = True
|
||||
return entries, exhausted
|
||||
|
||||
def _publish_ownership(self, sandbox_id: str) -> None:
|
||||
"""Publish acquire-side ownership before exposing a sandbox locally."""
|
||||
with self._lock:
|
||||
self._acquire_inflight.add(sandbox_id)
|
||||
try:
|
||||
if not self._ownership.take(sandbox_id):
|
||||
raise RuntimeError(f"E2B sandbox {sandbox_id} is being destroyed")
|
||||
except Exception:
|
||||
with self._lock:
|
||||
self._acquire_inflight.discard(sandbox_id)
|
||||
raise
|
||||
with self._lock:
|
||||
self._owned_sandbox_ids.add(sandbox_id)
|
||||
|
||||
def _claim_ownership(self, sandbox_id: str, *, for_destroy: bool = False) -> bool:
|
||||
"""Exclusively claim an unowned sandbox, failing closed on store errors."""
|
||||
try:
|
||||
claimed = self._ownership.claim(sandbox_id, for_destroy=for_destroy)
|
||||
except OwnershipBackendError as e:
|
||||
logger.warning("E2B ownership claim failed for %s: %s", sandbox_id, e)
|
||||
return False
|
||||
if claimed:
|
||||
with self._lock:
|
||||
self._owned_sandbox_ids.add(sandbox_id)
|
||||
return claimed
|
||||
|
||||
def _release_ownership(self, sandbox_id: str) -> None:
|
||||
try:
|
||||
self._ownership.release(sandbox_id)
|
||||
except OwnershipBackendError as e:
|
||||
logger.warning("Failed to release E2B ownership for %s: %s", sandbox_id, e)
|
||||
with self._lock:
|
||||
self._owned_sandbox_ids.discard(sandbox_id)
|
||||
self._acquire_inflight.discard(sandbox_id)
|
||||
|
||||
def _forget_local_sandbox(self, sandbox_id: str) -> None:
|
||||
"""Forget a lease taken by a peer without touching the remote VM."""
|
||||
with self._lock:
|
||||
if sandbox_id in self._acquire_inflight:
|
||||
return
|
||||
sandbox = self._sandboxes.pop(sandbox_id, None)
|
||||
self._warm_pool.pop(sandbox_id, None)
|
||||
self._owned_sandbox_ids.discard(sandbox_id)
|
||||
for key, sid in list(self._thread_sandboxes.items()):
|
||||
if sid == sandbox_id:
|
||||
self._thread_sandboxes.pop(key, None)
|
||||
if sandbox is not None:
|
||||
try:
|
||||
sandbox.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _refresh_owned_leases(self) -> None:
|
||||
with self._lock:
|
||||
sandbox_ids = list(self._owned_sandbox_ids)
|
||||
for sandbox_id in sandbox_ids:
|
||||
try:
|
||||
outcome = self._ownership.renew(sandbox_id)
|
||||
except OwnershipBackendError as e:
|
||||
logger.warning("Could not renew E2B ownership for %s; will retry: %s", sandbox_id, e)
|
||||
continue
|
||||
if outcome is RenewOutcome.RENEWED:
|
||||
continue
|
||||
if outcome is RenewOutcome.LAPSED:
|
||||
try:
|
||||
if self._ownership.claim(sandbox_id):
|
||||
continue
|
||||
except OwnershipBackendError as e:
|
||||
logger.warning("Could not re-establish E2B ownership for %s: %s", sandbox_id, e)
|
||||
continue
|
||||
logger.info("E2B sandbox %s ownership moved to a peer; forgetting local client", sandbox_id)
|
||||
self._forget_local_sandbox(sandbox_id)
|
||||
|
||||
def _start_maintenance_threads(self) -> None:
|
||||
def renew() -> None:
|
||||
interval = self._ownership_config.renewal_interval_seconds
|
||||
while not self._maintenance_stop.wait(interval):
|
||||
self._refresh_owned_leases()
|
||||
|
||||
def reconcile() -> None:
|
||||
interval = float(self._config["reconciliation_interval_seconds"])
|
||||
while not self._maintenance_stop.is_set():
|
||||
try:
|
||||
self._reconcile_remote_sandboxes()
|
||||
except Exception:
|
||||
logger.exception("Periodic E2B sandbox reconciliation failed")
|
||||
if self._maintenance_stop.wait(interval):
|
||||
break
|
||||
|
||||
self._lease_thread = threading.Thread(target=renew, name="e2b-lease-renewal", daemon=True)
|
||||
self._reconcile_thread = threading.Thread(target=reconcile, name="e2b-reconciliation", daemon=True)
|
||||
self._lease_thread.start()
|
||||
self._reconcile_thread.start()
|
||||
|
||||
def _reserve_reconciliation_capacity(self, sandbox_id: str) -> bool:
|
||||
"""Reserve one local slot for adoption without blocking maintenance."""
|
||||
with self._lock:
|
||||
if self._shutdown_called or self._total_capacity_used_locked() >= self._capacity_limit():
|
||||
return False
|
||||
self._reserved_slots += 1
|
||||
self._unowned_remote_ops_in_progress.add(sandbox_id)
|
||||
self._acquire_inflight.add(sandbox_id)
|
||||
return True
|
||||
|
||||
def _reconcile_remote_sandboxes(self, *, now: float | None = None) -> ReconciliationStats:
|
||||
"""Adopt canonical E2B sandboxes and safely reap duplicates/orphans."""
|
||||
observed_at = time.monotonic() if now is None else now
|
||||
deadline = time.monotonic() + float(self._config["reconciliation_max_seconds"])
|
||||
stats = ReconciliationStats()
|
||||
entries, stats.budget_exhausted = self._list_remote_entries({META_KEY_PROVIDER: META_VAL_PROVIDER})
|
||||
stats.discovered = len(entries)
|
||||
groups: dict[tuple[str, str], list[tuple[str, dict[str, Any]]]] = {}
|
||||
orphans: list[tuple[str, dict[str, Any]]] = []
|
||||
present_ids: set[str] = set()
|
||||
|
||||
for entry in entries:
|
||||
sandbox_id = self._entry_id(entry)
|
||||
if not sandbox_id:
|
||||
continue
|
||||
present_ids.add(sandbox_id)
|
||||
metadata = self._entry_metadata(entry)
|
||||
user_id = metadata.get(META_KEY_USER)
|
||||
thread_id = metadata.get(META_KEY_THREAD)
|
||||
if isinstance(user_id, str) and user_id and isinstance(thread_id, str) and thread_id:
|
||||
groups.setdefault((user_id, thread_id), []).append((sandbox_id, metadata))
|
||||
else:
|
||||
orphans.append((sandbox_id, metadata))
|
||||
|
||||
for (user_id, thread_id), candidates in groups.items():
|
||||
candidates.sort(key=lambda item: (item[1].get(META_KEY_CREATED_AT, ""), item[0]))
|
||||
with self._lock:
|
||||
local_id = self._thread_sandboxes.get((user_id, thread_id))
|
||||
if local_id:
|
||||
candidates.sort(key=lambda item: item[0] != local_id)
|
||||
|
||||
live: list[tuple[str, dict[str, Any], E2BClientSandbox]] = []
|
||||
for sandbox_id, metadata in candidates:
|
||||
if time.monotonic() >= deadline:
|
||||
stats.budget_exhausted = True
|
||||
break
|
||||
try:
|
||||
client = self._reconnect_live_client(self._get_sandbox_cls(), sandbox_id)
|
||||
except Exception as e:
|
||||
logger.debug("Could not probe E2B sandbox %s during reconciliation: %s", sandbox_id, e)
|
||||
continue
|
||||
if client is None:
|
||||
stats.dead += 1
|
||||
continue
|
||||
live.append((sandbox_id, metadata, client))
|
||||
|
||||
if not live:
|
||||
continue
|
||||
stats.duplicates += max(0, len(live) - 1)
|
||||
canonical_id, _metadata, canonical_client = live[0]
|
||||
with self._lock:
|
||||
already_local = canonical_id in self._sandboxes
|
||||
if already_local:
|
||||
self._safe_close_client(canonical_client)
|
||||
elif not self._reserve_reconciliation_capacity(canonical_id):
|
||||
self._safe_close_client(canonical_client)
|
||||
stats.deferred += 1
|
||||
elif not self._claim_ownership(canonical_id):
|
||||
self._complete_reserved_remote_op(canonical_id, remote_destroyed=False)
|
||||
with self._lock:
|
||||
self._acquire_inflight.discard(canonical_id)
|
||||
self._safe_close_client(canonical_client)
|
||||
stats.deferred += 1
|
||||
else:
|
||||
bootstrap_error, remote_destroyed = self._bootstrap_or_discard(canonical_client, canonical_id)
|
||||
if bootstrap_error is not None:
|
||||
self._complete_reserved_remote_op(canonical_id, remote_destroyed=remote_destroyed)
|
||||
with self._lock:
|
||||
self._acquire_inflight.discard(canonical_id)
|
||||
stats.deferred += 1
|
||||
else:
|
||||
discard_after_shutdown = False
|
||||
with self._lock:
|
||||
if self._shutdown_called:
|
||||
discard_after_shutdown = True
|
||||
else:
|
||||
self._owned_sandbox_ids.add(canonical_id)
|
||||
self._unowned_remote_ops_in_progress.discard(canonical_id)
|
||||
self._register_connected_sandbox(
|
||||
canonical_id,
|
||||
canonical_client,
|
||||
thread_id=thread_id,
|
||||
user_id=user_id,
|
||||
)
|
||||
self._commit_capacity()
|
||||
if discard_after_shutdown:
|
||||
if self._claim_ownership(canonical_id, for_destroy=True):
|
||||
self._kill_client(canonical_client)
|
||||
self._release_ownership(canonical_id)
|
||||
self._safe_close_client(canonical_client)
|
||||
else:
|
||||
stats.adopted += 1
|
||||
|
||||
for sandbox_id, _metadata, client in live[1:]:
|
||||
first_seen = self._orphan_first_seen.setdefault(sandbox_id, observed_at)
|
||||
if observed_at - first_seen < float(self._config["reconciliation_grace_seconds"]):
|
||||
self._safe_close_client(client)
|
||||
stats.deferred += 1
|
||||
continue
|
||||
if not self._claim_ownership(sandbox_id, for_destroy=True):
|
||||
self._safe_close_client(client)
|
||||
stats.deferred += 1
|
||||
continue
|
||||
if error := self._kill_client(client):
|
||||
logger.warning("Failed to kill duplicate E2B sandbox %s: %s", sandbox_id, error)
|
||||
self._release_ownership(sandbox_id)
|
||||
stats.deferred += 1
|
||||
else:
|
||||
stats.killed += 1
|
||||
self._orphan_first_seen.pop(sandbox_id, None)
|
||||
self._forget_local_sandbox(sandbox_id)
|
||||
self._release_ownership(sandbox_id)
|
||||
self._safe_close_client(client)
|
||||
|
||||
for sandbox_id, metadata in orphans:
|
||||
if time.monotonic() >= deadline:
|
||||
stats.budget_exhausted = True
|
||||
break
|
||||
first_seen = self._orphan_first_seen.setdefault(sandbox_id, observed_at)
|
||||
created_at = metadata.get(META_KEY_CREATED_AT)
|
||||
try:
|
||||
age = time.time() - float(created_at) if created_at is not None else observed_at - first_seen
|
||||
except (TypeError, ValueError):
|
||||
age = observed_at - first_seen
|
||||
if age < float(self._config["reconciliation_orphan_ttl_seconds"]):
|
||||
stats.deferred += 1
|
||||
continue
|
||||
if not self._claim_ownership(sandbox_id, for_destroy=True):
|
||||
stats.deferred += 1
|
||||
continue
|
||||
try:
|
||||
client = self._reconnect_client(self._get_sandbox_cls(), sandbox_id)
|
||||
except Exception:
|
||||
self._release_ownership(sandbox_id)
|
||||
continue
|
||||
if error := self._kill_client(client):
|
||||
logger.warning("Failed to kill orphan E2B sandbox %s: %s", sandbox_id, error)
|
||||
self._release_ownership(sandbox_id)
|
||||
stats.deferred += 1
|
||||
else:
|
||||
stats.killed += 1
|
||||
self._orphan_first_seen.pop(sandbox_id, None)
|
||||
self._forget_local_sandbox(sandbox_id)
|
||||
self._release_ownership(sandbox_id)
|
||||
self._safe_close_client(client)
|
||||
|
||||
for sandbox_id in set(self._orphan_first_seen) - present_ids:
|
||||
self._orphan_first_seen.pop(sandbox_id, None)
|
||||
logger.info(
|
||||
"E2B reconciliation: discovered=%d adopted=%d duplicates=%d deferred=%d killed=%d dead=%d budget_exhausted=%s",
|
||||
stats.discovered,
|
||||
stats.adopted,
|
||||
stats.duplicates,
|
||||
stats.deferred,
|
||||
stats.killed,
|
||||
stats.dead,
|
||||
stats.budget_exhausted,
|
||||
)
|
||||
return stats
|
||||
|
||||
def _reconnect_client(self, sandbox_cls: type[E2BClientSandbox], sandbox_id: str) -> E2BClientSandbox:
|
||||
"""Connect to an existing e2b sandbox by id, with consistent kwargs."""
|
||||
return sandbox_cls.connect(sandbox_id, **self._common_kwargs()) # type: ignore[attr-defined]
|
||||
@ -938,7 +1343,7 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
sandbox_id: str,
|
||||
client: E2BClientSandbox,
|
||||
*,
|
||||
thread_id: str,
|
||||
thread_id: str | None,
|
||||
user_id: str,
|
||||
) -> None:
|
||||
"""Track a live reconnected sandbox under its thread ownership.
|
||||
@ -947,7 +1352,10 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
"""
|
||||
sandbox = E2BSandbox(id=sandbox_id, client=client, home_dir=self._config["home_dir"])
|
||||
self._sandboxes[sandbox_id] = sandbox
|
||||
self._thread_sandboxes[self._thread_key(thread_id, user_id)] = sandbox_id
|
||||
self._warm_pool.pop(sandbox_id, None)
|
||||
if thread_id:
|
||||
self._thread_sandboxes[self._thread_key(thread_id, user_id)] = sandbox_id
|
||||
self._acquire_inflight.discard(sandbox_id)
|
||||
|
||||
def _refresh_remote_timeout(self, client: E2BClientSandbox) -> None:
|
||||
"""Push the configured idle timeout to the e2b control plane."""
|
||||
@ -1016,11 +1424,20 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._bootstrap_sandbox_paths(client)
|
||||
except Exception as e:
|
||||
logger.exception("Failed to bootstrap e2b sandbox %s. Discarding the unusable sandbox.", sandbox_id)
|
||||
kill_error = self._kill_client(client)
|
||||
if kill_error:
|
||||
logger.warning("Failed to kill e2b sandbox %s after bootstrap failure: %s", sandbox_id, kill_error)
|
||||
remote_destroyed = False
|
||||
if self._claim_ownership(sandbox_id, for_destroy=True):
|
||||
kill_error = self._kill_client(client)
|
||||
remote_destroyed = kill_error is None
|
||||
if kill_error:
|
||||
logger.warning("Failed to kill e2b sandbox %s after bootstrap failure: %s", sandbox_id, kill_error)
|
||||
self._release_ownership(sandbox_id)
|
||||
else:
|
||||
logger.info(
|
||||
"Not killing E2B sandbox %s after bootstrap failure because a peer owns it",
|
||||
sandbox_id,
|
||||
)
|
||||
self._safe_close_client(client)
|
||||
return e, kill_error is None
|
||||
return e, remote_destroyed
|
||||
return None, True
|
||||
|
||||
def _bootstrap_sandbox_paths(self, client: E2BClientSandbox) -> None:
|
||||
@ -1459,6 +1876,15 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
else:
|
||||
return None
|
||||
|
||||
if not self._claim_ownership(evict_id, for_destroy=True):
|
||||
logger.info("Deferring warm-pool eviction for %s because a peer owns it", evict_id)
|
||||
self._forget_local_sandbox(evict_id)
|
||||
with self._lock:
|
||||
if evict_id in self._evictions_in_progress:
|
||||
self._evictions_in_progress.discard(evict_id)
|
||||
self._eviction_tombstones.discard(evict_id)
|
||||
self._end_transition_locked()
|
||||
return evict_id
|
||||
try:
|
||||
client = self._reconnect_live_client(self._get_sandbox_cls(), evict_id)
|
||||
except Exception as e:
|
||||
@ -1472,6 +1898,7 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._evictions_in_progress.discard(evict_id)
|
||||
if not self._shutdown_called:
|
||||
self._eviction_tombstones.add(evict_id)
|
||||
self._release_ownership(evict_id)
|
||||
return None
|
||||
|
||||
if client is None:
|
||||
@ -1480,6 +1907,7 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._evictions_in_progress.discard(evict_id)
|
||||
self._eviction_tombstones.discard(evict_id)
|
||||
self._end_transition_locked()
|
||||
self._release_ownership(evict_id)
|
||||
logger.info("Evicted warm-pool e2b sandbox %s was already gone", evict_id)
|
||||
return evict_id
|
||||
|
||||
@ -1491,6 +1919,7 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._evictions_in_progress.discard(evict_id)
|
||||
if not self._shutdown_called:
|
||||
self._eviction_tombstones.add(evict_id)
|
||||
self._release_ownership(evict_id)
|
||||
return None
|
||||
|
||||
self._safe_close_client(client)
|
||||
@ -1499,6 +1928,7 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._evictions_in_progress.discard(evict_id)
|
||||
self._eviction_tombstones.discard(evict_id)
|
||||
self._end_transition_locked()
|
||||
self._release_ownership(evict_id)
|
||||
logger.info("Evicted warm-pool e2b sandbox %s", evict_id)
|
||||
return evict_id
|
||||
|
||||
@ -1608,8 +2038,10 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
"Provider shut down during release of sandbox %s; killing instead of parking in warm pool",
|
||||
sandbox_id,
|
||||
)
|
||||
if error := self._kill_client(client):
|
||||
logger.debug("Failed to kill e2b sandbox %s during release: %s", sandbox_id, error)
|
||||
if self._claim_ownership(sandbox_id, for_destroy=True):
|
||||
if error := self._kill_client(client):
|
||||
logger.debug("Failed to kill e2b sandbox %s during release: %s", sandbox_id, error)
|
||||
self._release_ownership(sandbox_id)
|
||||
self._safe_close_client(client)
|
||||
return
|
||||
|
||||
@ -1622,12 +2054,20 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._free_transitioning_slot()
|
||||
|
||||
def _kill_and_close(self, sandbox: E2BSandbox) -> None:
|
||||
if not self._claim_ownership(sandbox.id, for_destroy=True):
|
||||
logger.info("Not killing E2B sandbox %s because a peer owns it", sandbox.id)
|
||||
try:
|
||||
sandbox.close()
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
if error := self._kill_client(getattr(sandbox, "_client", None)):
|
||||
logger.debug(
|
||||
"kill() on e2b sandbox %s raised (probably already gone): %s",
|
||||
sandbox.id,
|
||||
error,
|
||||
)
|
||||
self._release_ownership(sandbox.id)
|
||||
try:
|
||||
sandbox.close()
|
||||
except Exception:
|
||||
@ -1658,8 +2098,14 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
if self._shutdown_called:
|
||||
return
|
||||
self._shutdown_called = True
|
||||
self._maintenance_stop.set()
|
||||
for thread in (self._lease_thread, self._reconcile_thread):
|
||||
if thread is not None and thread is not threading.current_thread():
|
||||
thread.join(timeout=max(5.0, float(self._config["reconciliation_max_seconds"]) + 1.0))
|
||||
with self._lock:
|
||||
active = list(self._sandboxes.items())
|
||||
warm_ids = list(self._warm_pool.keys() | self._eviction_tombstones | self._remote_ops_in_progress)
|
||||
owned_ids = set(self._owned_sandbox_ids)
|
||||
self._sandboxes.clear()
|
||||
self._warm_pool.clear()
|
||||
self._eviction_tombstones.clear()
|
||||
@ -1668,6 +2114,9 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._unowned_remote_ops_in_progress.clear()
|
||||
self._thread_sandboxes.clear()
|
||||
self._thread_locks.clear()
|
||||
self._owned_sandbox_ids.clear()
|
||||
self._acquire_inflight.clear()
|
||||
self._orphan_first_seen.clear()
|
||||
self._reserved_slots = 0
|
||||
self._transitioning_slots = 0
|
||||
self._capacity_cond.notify_all()
|
||||
@ -1683,12 +2132,14 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
)
|
||||
|
||||
for sandbox_id, sandbox in active:
|
||||
if error := self._kill_client(sandbox.client):
|
||||
logger.warning(
|
||||
"Failed to kill active e2b sandbox %s during shutdown: %s",
|
||||
sandbox_id,
|
||||
error,
|
||||
)
|
||||
if sandbox_id in owned_ids and self._claim_ownership(sandbox_id, for_destroy=True):
|
||||
if error := self._kill_client(sandbox.client):
|
||||
logger.warning(
|
||||
"Failed to kill active e2b sandbox %s during shutdown: %s",
|
||||
sandbox_id,
|
||||
error,
|
||||
)
|
||||
self._release_ownership(sandbox_id)
|
||||
try:
|
||||
sandbox.close()
|
||||
except Exception:
|
||||
@ -1696,6 +2147,10 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
|
||||
sandbox_cls = self._get_sandbox_cls()
|
||||
for sandbox_id in warm_ids:
|
||||
if sandbox_id not in owned_ids:
|
||||
continue
|
||||
if not self._claim_ownership(sandbox_id, for_destroy=True):
|
||||
continue
|
||||
try:
|
||||
client = self._reconnect_client(sandbox_cls, sandbox_id)
|
||||
except Exception as e:
|
||||
@ -1704,6 +2159,7 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
sandbox_id,
|
||||
e,
|
||||
)
|
||||
self._release_ownership(sandbox_id)
|
||||
continue
|
||||
if error := self._kill_client(client):
|
||||
logger.warning(
|
||||
@ -1711,9 +2167,14 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
sandbox_id,
|
||||
error,
|
||||
)
|
||||
self._release_ownership(sandbox_id)
|
||||
close = getattr(client, "close", None)
|
||||
if callable(close):
|
||||
try:
|
||||
close()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
self._ownership.close()
|
||||
except Exception as e:
|
||||
logger.warning("Failed to close E2B ownership store: %s", e)
|
||||
|
||||
@ -88,8 +88,10 @@ class SandboxConfig(BaseModel):
|
||||
port: Base port for sandbox containers (default: 8080)
|
||||
container_prefix: Prefix for container names (default: deer-flow-sandbox)
|
||||
mounts: List of volume mounts to share directories with the container
|
||||
ownership: Cross-instance container ownership store (memory | redis). Multi-instance
|
||||
deployments sharing a container backend need redis; see SandboxOwnershipConfig.
|
||||
|
||||
AioSandboxProvider and E2BSandboxProvider shared options:
|
||||
ownership: Cross-instance sandbox ownership store (memory | redis). Multi-instance
|
||||
deployments sharing a sandbox backend need redis; see SandboxOwnershipConfig.
|
||||
"""
|
||||
|
||||
use: str = Field(
|
||||
@ -143,8 +145,8 @@ class SandboxConfig(BaseModel):
|
||||
ownership: SandboxOwnershipConfig | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"AioSandboxProvider-only: where cross-instance container ownership is tracked (#4206). Omitted = memory (single-instance). "
|
||||
"Multi-worker / load-balanced gateways sharing one container backend must set type: redis, or peers will adopt and idle-destroy each other's live sandboxes."
|
||||
"AioSandboxProvider/E2BSandboxProvider: where cross-instance sandbox ownership is tracked (#4206, #4341). Omitted = memory (single-instance). "
|
||||
"Multi-worker / load-balanced gateways sharing one sandbox backend must set type: redis, or peers can adopt and destroy each other's live sandboxes."
|
||||
),
|
||||
)
|
||||
mounts: list[VolumeMountConfig] = Field(
|
||||
|
||||
@ -175,6 +175,65 @@ class FakeSandboxClass:
|
||||
return self.list_return
|
||||
|
||||
|
||||
class FakeOwnershipStore:
|
||||
"""Shared lease state with per-provider identities for reconciliation tests."""
|
||||
|
||||
supports_cross_process = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
leases: dict[str, tuple[str, str]],
|
||||
*,
|
||||
owner_id: str,
|
||||
lock: threading.Lock | None = None,
|
||||
) -> None:
|
||||
self._leases = leases
|
||||
self._lock = lock or threading.Lock()
|
||||
self.owner_id = owner_id
|
||||
|
||||
def take(self, sandbox_id: str) -> bool:
|
||||
with self._lock:
|
||||
current = self._leases.get(sandbox_id)
|
||||
if current is not None and current[1] == "del":
|
||||
return False
|
||||
self._leases[sandbox_id] = (self.owner_id, "own")
|
||||
return True
|
||||
|
||||
def claim(self, sandbox_id: str, *, for_destroy: bool = False) -> bool:
|
||||
with self._lock:
|
||||
current = self._leases.get(sandbox_id)
|
||||
if current is not None and current[0] != self.owner_id:
|
||||
return False
|
||||
if current is not None and current[1] == "del" and not for_destroy:
|
||||
return False
|
||||
self._leases[sandbox_id] = (self.owner_id, "del" if for_destroy else "own")
|
||||
return True
|
||||
|
||||
def renew(self, sandbox_id: str):
|
||||
from deerflow.community.aio_sandbox.ownership import RenewOutcome
|
||||
|
||||
with self._lock:
|
||||
current = self._leases.get(sandbox_id)
|
||||
if current is None:
|
||||
return RenewOutcome.LAPSED
|
||||
if current == (self.owner_id, "own"):
|
||||
return RenewOutcome.RENEWED
|
||||
return RenewOutcome.LOST
|
||||
|
||||
def release(self, sandbox_id: str) -> None:
|
||||
with self._lock:
|
||||
if self._leases.get(sandbox_id, (None,))[0] == self.owner_id:
|
||||
self._leases.pop(sandbox_id, None)
|
||||
|
||||
def owner(self, sandbox_id: str) -> str | None:
|
||||
with self._lock:
|
||||
current = self._leases.get(sandbox_id)
|
||||
return current[0] if current is not None else None
|
||||
|
||||
def close(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800, overflow_policy: str = "wait", acquire_timeout: int = 30, burst_limit: int = 0) -> Any:
|
||||
"""Build a ``E2BSandboxProvider`` instance bypassing ``__init__``."""
|
||||
mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox_provider")
|
||||
@ -192,6 +251,15 @@ def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800, overflow_poli
|
||||
provider._transitioning_slots = 0
|
||||
provider._capacity_cond = threading.Condition(provider._lock)
|
||||
provider._shutdown_called = False
|
||||
provider._owner_id = "owner-a"
|
||||
provider._ownership = FakeOwnershipStore({}, owner_id=provider._owner_id)
|
||||
provider._ownership_config = SimpleNamespace(renewal_interval_seconds=60.0)
|
||||
provider._owned_sandbox_ids = set()
|
||||
provider._acquire_inflight = set()
|
||||
provider._orphan_first_seen = {}
|
||||
provider._maintenance_stop = threading.Event()
|
||||
provider._lease_thread = None
|
||||
provider._reconcile_thread = None
|
||||
provider._config = {
|
||||
"api_key": "test-key",
|
||||
"template": "code-interpreter-v1",
|
||||
@ -204,6 +272,12 @@ def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800, overflow_poli
|
||||
"burst_limit": burst_limit,
|
||||
"mounts": [],
|
||||
"environment": {},
|
||||
"reconciliation_interval_seconds": 60.0,
|
||||
"reconciliation_grace_seconds": 30.0,
|
||||
"reconciliation_orphan_ttl_seconds": 3600.0,
|
||||
"reconciliation_max_pages": 10,
|
||||
"reconciliation_max_items": 100,
|
||||
"reconciliation_max_seconds": 10.0,
|
||||
}
|
||||
return provider
|
||||
|
||||
@ -412,6 +486,84 @@ def test_kill_and_close_swallows_kill_exceptions():
|
||||
|
||||
client.kill = explode
|
||||
p._kill_and_close(sb)
|
||||
assert "fake-sb-1" not in p._owned_sandbox_ids
|
||||
assert p._ownership.owner("fake-sb-1") is None
|
||||
|
||||
|
||||
def test_kill_and_close_skips_peer_owned_sandbox():
|
||||
p = _make_provider()
|
||||
client = FakeClient(sandbox_id="peer-owned")
|
||||
sandbox = _make_sandbox(client, sandbox_id="peer-owned")
|
||||
p._ownership = FakeOwnershipStore(
|
||||
{"peer-owned": ("owner-peer", "own")},
|
||||
owner_id=p._owner_id,
|
||||
)
|
||||
|
||||
p._kill_and_close(sandbox)
|
||||
|
||||
assert client.killed is False
|
||||
assert client.closed is True
|
||||
assert p._ownership.owner("peer-owned") == "owner-peer"
|
||||
|
||||
|
||||
def test_startup_reconciliation_runs_in_background_without_blocking_caller(monkeypatch):
|
||||
p = _make_provider()
|
||||
entered = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
def reconcile():
|
||||
entered.set()
|
||||
assert release.wait(timeout=2)
|
||||
|
||||
monkeypatch.setattr(p, "_reconcile_remote_sandboxes", reconcile)
|
||||
|
||||
started = time.monotonic()
|
||||
p._start_maintenance_threads()
|
||||
elapsed = time.monotonic() - started
|
||||
|
||||
assert entered.wait(timeout=1)
|
||||
assert elapsed < 0.5
|
||||
|
||||
release.set()
|
||||
p._maintenance_stop.set()
|
||||
for thread in (p._lease_thread, p._reconcile_thread):
|
||||
assert thread is not None
|
||||
thread.join(timeout=2)
|
||||
|
||||
|
||||
def test_refresh_owned_leases_reclaims_lapsed_lease():
|
||||
p = _make_provider()
|
||||
client = FakeClient(sandbox_id="sb-lapsed")
|
||||
sandbox = _make_sandbox(client, sandbox_id="sb-lapsed")
|
||||
p._sandboxes["sb-lapsed"] = sandbox
|
||||
p._thread_sandboxes[("u1", "t1")] = "sb-lapsed"
|
||||
p._owned_sandbox_ids.add("sb-lapsed")
|
||||
|
||||
p._refresh_owned_leases()
|
||||
|
||||
assert p._ownership.owner("sb-lapsed") == p._owner_id
|
||||
assert p.get("sb-lapsed") is sandbox
|
||||
assert client.closed is False
|
||||
|
||||
|
||||
def test_refresh_owned_leases_forgets_peer_owned_sandbox():
|
||||
p = _make_provider()
|
||||
client = FakeClient(sandbox_id="sb-lost")
|
||||
sandbox = _make_sandbox(client, sandbox_id="sb-lost")
|
||||
p._sandboxes["sb-lost"] = sandbox
|
||||
p._thread_sandboxes[("u1", "t1")] = "sb-lost"
|
||||
p._owned_sandbox_ids.add("sb-lost")
|
||||
p._ownership = FakeOwnershipStore(
|
||||
{"sb-lost": ("owner-peer", "own")},
|
||||
owner_id=p._owner_id,
|
||||
)
|
||||
|
||||
p._refresh_owned_leases()
|
||||
|
||||
assert p.get("sb-lost") is None
|
||||
assert ("u1", "t1") not in p._thread_sandboxes
|
||||
assert "sb-lost" not in p._owned_sandbox_ids
|
||||
assert client.closed is True
|
||||
|
||||
|
||||
def test_reuse_in_process_sandbox_returns_cached_id_on_healthy_reuse():
|
||||
@ -614,6 +766,239 @@ def test_discover_remote_sandbox_skips_dead_candidate(monkeypatch):
|
||||
assert client.closed is True
|
||||
|
||||
|
||||
def test_discover_remote_sandbox_tries_later_candidate_when_first_is_dead(monkeypatch):
|
||||
p = _make_provider()
|
||||
fake_cls = _install_fake_sdk(monkeypatch, p)
|
||||
fake_cls.list_return = [
|
||||
_info("sb-a-dead", "u1", "t1"),
|
||||
_info("sb-b-live", "u1", "t1"),
|
||||
]
|
||||
dead = FakeClient(
|
||||
sandbox_id="sb-a-dead",
|
||||
commands=FakeCommandsAPI([FakeCommandsAPI.GONE]),
|
||||
)
|
||||
live = FakeClient(sandbox_id="sb-b-live")
|
||||
fake_cls.connect_factory = lambda sid, **_kw: dead if sid == "sb-a-dead" else live
|
||||
|
||||
assert p._discover_remote_sandbox("t1", user_id="u1") == "sb-b-live"
|
||||
assert dead.closed is True
|
||||
assert p._thread_sandboxes[("u1", "t1")] == "sb-b-live"
|
||||
assert "sb-b-live" in p._owned_sandbox_ids
|
||||
|
||||
|
||||
def test_reconcile_defers_duplicate_with_live_peer_lease(monkeypatch):
|
||||
p = _make_provider()
|
||||
fake_cls = _install_fake_sdk(monkeypatch, p)
|
||||
fake_cls.list_return = [
|
||||
_info("sb-canonical", "u1", "t1"),
|
||||
_info("sb-duplicate", "u1", "t1"),
|
||||
]
|
||||
clients = {
|
||||
"sb-canonical": FakeClient(sandbox_id="sb-canonical"),
|
||||
"sb-duplicate": FakeClient(sandbox_id="sb-duplicate"),
|
||||
}
|
||||
fake_cls.connect_factory = lambda sid, **_kw: clients[sid]
|
||||
leases = {"sb-duplicate": ("owner-peer", "own")}
|
||||
p._ownership = FakeOwnershipStore(leases, owner_id=p._owner_id)
|
||||
p._config["reconciliation_grace_seconds"] = 0.0
|
||||
|
||||
stats = p._reconcile_remote_sandboxes(now=time.monotonic())
|
||||
|
||||
assert stats.adopted == 1
|
||||
assert stats.deferred == 1
|
||||
assert stats.killed == 0
|
||||
assert clients["sb-duplicate"].killed is False
|
||||
|
||||
|
||||
def test_reconcile_kills_unowned_duplicate_after_grace(monkeypatch):
|
||||
p = _make_provider()
|
||||
fake_cls = _install_fake_sdk(monkeypatch, p)
|
||||
fake_cls.list_return = [
|
||||
_info("sb-canonical", "u1", "t1"),
|
||||
_info("sb-duplicate", "u1", "t1"),
|
||||
]
|
||||
clients = {
|
||||
"sb-canonical": FakeClient(sandbox_id="sb-canonical"),
|
||||
"sb-duplicate": FakeClient(sandbox_id="sb-duplicate"),
|
||||
}
|
||||
fake_cls.connect_factory = lambda sid, **_kw: clients[sid]
|
||||
p._config["reconciliation_grace_seconds"] = 5.0
|
||||
|
||||
first = p._reconcile_remote_sandboxes(now=100.0)
|
||||
second = p._reconcile_remote_sandboxes(now=106.0)
|
||||
|
||||
assert first.deferred == 1
|
||||
assert first.killed == 0
|
||||
assert second.killed == 1
|
||||
assert clients["sb-duplicate"].killed is True
|
||||
|
||||
|
||||
def test_competing_reconcilers_only_one_kills_duplicate(monkeypatch):
|
||||
shared_leases: dict[str, tuple[str, str]] = {}
|
||||
shared_lock = threading.Lock()
|
||||
providers = [_make_provider(), _make_provider()]
|
||||
providers[0]._owner_id = "owner-a"
|
||||
providers[1]._owner_id = "owner-b"
|
||||
clients = {
|
||||
"sb-canonical": FakeClient(sandbox_id="sb-canonical"),
|
||||
"sb-duplicate": FakeClient(sandbox_id="sb-duplicate"),
|
||||
}
|
||||
kill_count = 0
|
||||
|
||||
def kill_duplicate() -> None:
|
||||
nonlocal kill_count
|
||||
kill_count += 1
|
||||
clients["sb-duplicate"].killed = True
|
||||
clients["sb-duplicate"].commands._responses.append(FakeCommandsAPI.GONE)
|
||||
|
||||
clients["sb-duplicate"].kill = kill_duplicate
|
||||
|
||||
for provider in providers:
|
||||
fake_cls = _install_fake_sdk(monkeypatch, provider)
|
||||
fake_cls.list_return = [
|
||||
_info("sb-canonical", "u1", "t1"),
|
||||
_info("sb-duplicate", "u1", "t1"),
|
||||
]
|
||||
fake_cls.connect_factory = lambda sid, **_kw: clients[sid]
|
||||
provider._ownership = FakeOwnershipStore(
|
||||
shared_leases,
|
||||
owner_id=provider._owner_id,
|
||||
lock=shared_lock,
|
||||
)
|
||||
provider._config["reconciliation_grace_seconds"] = 0.0
|
||||
|
||||
results = [provider._reconcile_remote_sandboxes(now=100.0) for provider in providers]
|
||||
|
||||
assert sum(result.killed for result in results) == 1
|
||||
assert kill_count == 1
|
||||
|
||||
|
||||
def test_reconcile_honors_item_budget(monkeypatch):
|
||||
p = _make_provider()
|
||||
fake_cls = _install_fake_sdk(monkeypatch, p)
|
||||
fake_cls.list_return = [_info(f"sb-{index}", "u1", f"t-{index}") for index in range(5)]
|
||||
p._config["reconciliation_max_items"] = 2
|
||||
|
||||
stats = p._reconcile_remote_sandboxes(now=100.0)
|
||||
|
||||
assert stats.discovered == 2
|
||||
assert len(fake_cls.connect_calls) == 2
|
||||
|
||||
|
||||
def test_reconcile_honors_page_budget(monkeypatch):
|
||||
p = _make_provider()
|
||||
fake_cls = _install_fake_sdk(monkeypatch, p)
|
||||
paginator = _FakePaginator(
|
||||
[
|
||||
[_info("sb-first", "u1", "t1")],
|
||||
[_info("sb-never-read", "u2", "t2")],
|
||||
]
|
||||
)
|
||||
fake_cls.list_return = paginator
|
||||
p._config["reconciliation_max_pages"] = 1
|
||||
|
||||
stats = p._reconcile_remote_sandboxes(now=100.0)
|
||||
|
||||
assert stats.discovered == 1
|
||||
assert stats.budget_exhausted is True
|
||||
assert paginator.calls == 1
|
||||
assert [call[0] for call in fake_cls.connect_calls] == ["sb-first"]
|
||||
|
||||
|
||||
def test_reconcile_honors_wall_clock_budget(monkeypatch):
|
||||
p = _make_provider()
|
||||
fake_cls = _install_fake_sdk(monkeypatch, p)
|
||||
fake_cls.list_return = [_info("sb-never-probed", "u1", "t1")]
|
||||
p._config["reconciliation_max_seconds"] = 0.5
|
||||
provider_mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox_provider")
|
||||
ticks = iter([0.0, 0.0, 1.0, 1.0])
|
||||
monkeypatch.setattr(provider_mod.time, "monotonic", lambda: next(ticks))
|
||||
|
||||
stats = p._reconcile_remote_sandboxes(now=100.0)
|
||||
|
||||
assert stats.budget_exhausted is True
|
||||
assert fake_cls.connect_calls == []
|
||||
|
||||
|
||||
def test_reconcile_adopts_canonical_after_restart_loses_local_state(monkeypatch):
|
||||
p = _make_provider()
|
||||
fake_cls = _install_fake_sdk(monkeypatch, p)
|
||||
fake_cls.list_return = [_info("sb-existing", "u1", "t1")]
|
||||
|
||||
stats = p._reconcile_remote_sandboxes(now=100.0)
|
||||
|
||||
assert stats.adopted == 1
|
||||
assert p._thread_sandboxes[("u1", "t1")] == "sb-existing"
|
||||
assert "sb-existing" in p._owned_sandbox_ids
|
||||
|
||||
|
||||
def test_reconcile_bootstrap_failure_clears_inflight_after_peer_take(monkeypatch):
|
||||
p = _make_provider()
|
||||
fake_cls = _install_fake_sdk(monkeypatch, p)
|
||||
sandbox_id = "sb-racy-bootstrap"
|
||||
fake_cls.list_return = [_info(sandbox_id, "u1", "t1")]
|
||||
shared_leases: dict[str, tuple[str, str]] = {}
|
||||
shared_lock = threading.Lock()
|
||||
p._ownership = FakeOwnershipStore(
|
||||
shared_leases,
|
||||
owner_id=p._owner_id,
|
||||
lock=shared_lock,
|
||||
)
|
||||
peer_ownership = FakeOwnershipStore(
|
||||
shared_leases,
|
||||
owner_id="owner-peer",
|
||||
lock=shared_lock,
|
||||
)
|
||||
|
||||
def peer_takes_before_bootstrap_fails(_command: str) -> SimpleNamespace:
|
||||
assert peer_ownership.take(sandbox_id) is True
|
||||
return SimpleNamespace(stdout="", stderr="permission denied", exit_code=1)
|
||||
|
||||
client = FakeClient(
|
||||
sandbox_id=sandbox_id,
|
||||
commands=FakeCommandsAPI(
|
||||
[
|
||||
SimpleNamespace(stdout="ok", stderr="", exit_code=0),
|
||||
peer_takes_before_bootstrap_fails,
|
||||
]
|
||||
),
|
||||
)
|
||||
fake_cls.connect_factory = lambda _sid, **_kw: client
|
||||
|
||||
stats = p._reconcile_remote_sandboxes(now=100.0)
|
||||
|
||||
assert stats.adopted == 0
|
||||
assert stats.deferred == 1
|
||||
assert sandbox_id not in p._acquire_inflight
|
||||
assert sandbox_id not in p._unowned_remote_ops_in_progress
|
||||
assert p._reserved_slots == 0
|
||||
assert p._ownership.owner(sandbox_id) == "owner-peer"
|
||||
assert client.killed is False
|
||||
assert client.closed is True
|
||||
|
||||
|
||||
def test_reconcile_kills_metadata_orphan_only_after_ttl(monkeypatch):
|
||||
p = _make_provider()
|
||||
fake_cls = _install_fake_sdk(monkeypatch, p)
|
||||
fake_cls.list_return = [
|
||||
SimpleNamespace(
|
||||
sandbox_id="sb-orphan",
|
||||
metadata={"deer_flow_provider": "e2b_sandbox_provider"},
|
||||
)
|
||||
]
|
||||
client = FakeClient(sandbox_id="sb-orphan")
|
||||
fake_cls.connect_factory = lambda _sid, **_kw: client
|
||||
p._config["reconciliation_orphan_ttl_seconds"] = 5.0
|
||||
|
||||
first = p._reconcile_remote_sandboxes(now=100.0)
|
||||
second = p._reconcile_remote_sandboxes(now=106.0)
|
||||
|
||||
assert first.deferred == 1
|
||||
assert first.killed == 0
|
||||
assert second.killed == 1
|
||||
assert client.killed is True
|
||||
|
||||
|
||||
def test_discover_remote_sandbox_discards_candidate_when_bootstrap_fails(monkeypatch):
|
||||
p = _make_provider()
|
||||
fake_cls = _install_fake_sdk(monkeypatch, p)
|
||||
@ -635,6 +1020,65 @@ def test_discover_remote_sandbox_discards_candidate_when_bootstrap_fails(monkeyp
|
||||
assert ("u1", "t1") not in p._thread_sandboxes
|
||||
|
||||
|
||||
def test_discovery_claims_ownership_before_bootstrap_cleanup(monkeypatch):
|
||||
events: list[str] = []
|
||||
|
||||
class RecordingOwnershipStore(FakeOwnershipStore):
|
||||
def take(self, sandbox_id: str) -> bool:
|
||||
events.append("take")
|
||||
return super().take(sandbox_id)
|
||||
|
||||
def claim(self, sandbox_id: str, *, for_destroy: bool = False) -> bool:
|
||||
events.append("claim-destroy" if for_destroy else "claim")
|
||||
return super().claim(sandbox_id, for_destroy=for_destroy)
|
||||
|
||||
def release(self, sandbox_id: str) -> None:
|
||||
events.append("release")
|
||||
super().release(sandbox_id)
|
||||
|
||||
p = _make_provider()
|
||||
p._ownership = RecordingOwnershipStore(
|
||||
{"sb-broken": ("owner-peer", "own")},
|
||||
owner_id=p._owner_id,
|
||||
)
|
||||
fake_cls = _install_fake_sdk(monkeypatch, p)
|
||||
fake_cls.list_return = [_info("sb-broken", "u1", "t1")]
|
||||
client = FakeClient(
|
||||
sandbox_id="sb-broken",
|
||||
commands=FakeCommandsAPI(
|
||||
[
|
||||
SimpleNamespace(stdout="ok", stderr="", exit_code=0),
|
||||
SimpleNamespace(stdout="", stderr="permission denied", exit_code=1),
|
||||
]
|
||||
),
|
||||
)
|
||||
client.kill = lambda: events.append("kill")
|
||||
fake_cls.connect_factory = lambda _sid, **_kw: client
|
||||
|
||||
assert p._discover_remote_sandbox("t1", user_id="u1") is None
|
||||
assert events == ["take", "claim-destroy", "kill", "release"]
|
||||
|
||||
|
||||
def test_bootstrap_failure_does_not_kill_without_destroy_lease(monkeypatch):
|
||||
p = _make_provider()
|
||||
client = FakeClient(
|
||||
sandbox_id="sb-peer",
|
||||
commands=FakeCommandsAPI([SimpleNamespace(stdout="", stderr="permission denied", exit_code=1)]),
|
||||
)
|
||||
p._ownership = FakeOwnershipStore(
|
||||
{"sb-peer": ("owner-peer", "own")},
|
||||
owner_id=p._owner_id,
|
||||
)
|
||||
|
||||
error, remote_destroyed = p._bootstrap_or_discard(client, "sb-peer")
|
||||
|
||||
assert error is not None
|
||||
assert remote_destroyed is False
|
||||
assert client.killed is False
|
||||
assert client.closed is True
|
||||
assert p._ownership.owner("sb-peer") == "owner-peer"
|
||||
|
||||
|
||||
def test_kill_client_returns_exception_without_raising():
|
||||
p = _make_provider()
|
||||
client = FakeClient()
|
||||
@ -728,6 +1172,48 @@ def test_evict_oldest_warm_keeps_slot_when_kill_lookup_raises(monkeypatch):
|
||||
assert client.closed is True
|
||||
assert p._eviction_tombstones == {"sb-warm"}
|
||||
assert p._transitioning_slots == 1
|
||||
assert "sb-warm" not in p._warm_pool
|
||||
assert "sb-warm" not in p._owned_sandbox_ids
|
||||
assert p._ownership.owner("sb-warm") is None
|
||||
|
||||
|
||||
def test_evict_oldest_warm_defers_peer_owned_sandbox(monkeypatch):
|
||||
p = _make_provider()
|
||||
fake_cls = _install_fake_sdk(monkeypatch, p)
|
||||
p._warm_pool["sb-peer"] = ("seed", 12345.0)
|
||||
p._ownership = FakeOwnershipStore(
|
||||
{"sb-peer": ("owner-peer", "own")},
|
||||
owner_id=p._owner_id,
|
||||
)
|
||||
|
||||
assert p._evict_oldest_warm() == "sb-peer"
|
||||
|
||||
assert fake_cls.connect_calls == []
|
||||
assert "sb-peer" not in p._warm_pool
|
||||
assert p._eviction_tombstones == set()
|
||||
assert p._evictions_in_progress == set()
|
||||
assert p._transitioning_slots == 0
|
||||
assert p._ownership.owner("sb-peer") == "owner-peer"
|
||||
|
||||
|
||||
def test_evict_oldest_warm_releases_destroy_lease_after_kill_failure(monkeypatch):
|
||||
p = _make_provider()
|
||||
fake_cls = _install_fake_sdk(monkeypatch, p)
|
||||
client = FakeClient(sandbox_id="sb-warm")
|
||||
client.kill = MagicMock(side_effect=RuntimeError("transient kill failure"))
|
||||
fake_cls.connect_factory = lambda _sid, **_kw: client
|
||||
p._warm_pool["sb-warm"] = ("seed", 12345.0)
|
||||
p._owned_sandbox_ids.add("sb-warm")
|
||||
p._ownership.take("sb-warm")
|
||||
|
||||
assert p._evict_oldest_warm() is None
|
||||
|
||||
assert client.closed is True
|
||||
assert "sb-warm" not in p._warm_pool
|
||||
assert p._eviction_tombstones == {"sb-warm"}
|
||||
assert p._transitioning_slots == 1
|
||||
assert "sb-warm" not in p._owned_sandbox_ids
|
||||
assert p._ownership.owner("sb-warm") is None
|
||||
|
||||
|
||||
def test_evict_oldest_warm_uses_kill_helper_and_closes_client(monkeypatch):
|
||||
@ -937,6 +1423,23 @@ def test_release_skips_warm_pool_when_sync_reveals_dead_vm(monkeypatch, tmp_path
|
||||
assert client.killed is True
|
||||
|
||||
|
||||
def test_shutdown_only_kills_sandboxes_owned_by_current_instance(monkeypatch):
|
||||
p = _make_provider()
|
||||
owned_client = FakeClient(sandbox_id="sb-owned")
|
||||
peer_client = FakeClient(sandbox_id="sb-peer")
|
||||
p._sandboxes = {
|
||||
"sb-owned": _make_sandbox(owned_client),
|
||||
"sb-peer": _make_sandbox(peer_client),
|
||||
}
|
||||
p._owned_sandbox_ids = {"sb-owned"}
|
||||
|
||||
p.shutdown()
|
||||
|
||||
assert owned_client.killed is True
|
||||
assert peer_client.killed is False
|
||||
assert peer_client.closed is True
|
||||
|
||||
|
||||
def _setup_paths(monkeypatch, tmp_path):
|
||||
paths_mod = importlib.import_module("deerflow.config.paths")
|
||||
monkeypatch.setattr(paths_mod, "get_paths", lambda: Paths(base_dir=tmp_path), raising=False)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user