Skip to content

refactor(python): async-native fleet, and handle release semantics - #88

Open
Sunny (sunnylovestiramisu) wants to merge 1 commit into
agent-substrate:mainfrom
sunnylovestiramisu:pr2
Open

Sunny (sunnylovestiramisu) wants to merge 1 commit into
agent-substrate:mainfrom
sunnylovestiramisu:pr2

Conversation

@sunnylovestiramisu

@sunnylovestiramisu Sunny (sunnylovestiramisu) commented Oct 14, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

This PR introduces AsyncSandboxFleet, an asyncio-native orchestrator designed for high-concurrency workloads (such as reinforcement learning rollouts with VeRL, TorchRL, and Ray), standardizes sandbox handle lifecycle semantics, and integrates rootfs image resolution for unmodified container images via Substrate's derived template mechanism (#70).


Key Changes

1. AsyncSandboxFleet

  • Native Async Orchestrator: Introduces AsyncSandboxFleet using asyncio tasks, semaphores, and thread pools to enable non-blocking batch acquisition and warm-pool sizing without blocking the caller's event loop.
  • Async Setup & Sizing: Supports await fleet.setup(tasks=None) to run pre-flight, compute a FleetPlan, and warm actor pools concurrently across task images.
  • Batch & Windowed Warmup: Provides warm_images() and unwarm_image() to support rolling or sliding-window allocation strategies for bounded footprint execution.
  • Top-Level Export: Exported AsyncSandboxFleet from ate_env.

2. Handle Lifecycle & Release Semantics

  • Standardized Execution Results: Standardized SandboxHandle.exec() and exec_async() to return ExecResult.
  • Async Release & Recycle: Added SandboxHandle.release_async() and recycle_async(), explicitly documenting that recycling suspends the actor without resetting in-guest filesystem or memory state.
  • Safer Context Manager Defaults: Changed SandboxHandle context manager (__exit__ / __aexit__) to default to release() (terminate) rather than destroying the entire fleet run.

3. Task Image & Rootfs Resolution (Addressing Review Feedback)

  • _ensure_template_for_image Resolution: Returns (template, rootfs_image):
    • Mapped images in config.image_templates resolve to (mapped_template, None).
    • Unmapped task images resolve to (config.template or "default-template", image), passing the task image through as rootfs_image while using the configured template as base.
  • BackendDriver Signature Updates: Added rootfs_image: Optional[str] = None parameter across acquire, acquire_async, warm_pool, warm_pool_async, and unwarm_pool in BackendDriver, SubstrateBackendDriver, and MockBackendDriver.
  • Client.create Propagation: Passed rootfs_image=rootfs_image at both client.create invocations in SubstrateBackendDriver (on-demand creation in _acquire_internal and pre-warm in warm_pool_async).
  • Pool Keying by (template, image): Keyed _warm_paused_pool by (template, image) tuple to prevent warm pool collisions between environments sharing a base template with distinct rootfs images, while retaining transparent string key normalization for backwards compatibility.
  • Metadata Propagation: Propagated rootfs_image onto RawSandboxInstance and PlanEntry, and correctly preserved the pool key upon release_async(..., recycle=True).
  • Task Image Docstring: Documented Task.image behavior detailing lookup in image_templates versus propagation as rootfs_image.

Review Feedback Addressed

  • Image Propagation to Server: Fixed issue where FleetConfig with Task(image="...") sent template.name == "default-template" and empty rootfs_image, ignoring the task container image.
  • Per-Image Template Derivation: Leveraged Client.create(rootfs_image=...) introduced in Run unmodified rootfs images: guest injected as an OCI image volume, plus second-runtime layers #70 so the Substrate server derives actor templates per digest on the fly.
  • Pool Key Alignment: Keyed _warm_paused_pool by (template, image).

Verification

  • Unit Tests:

    • Added test_substrate_rootfs_image_sent_to_server in test_substrate_driver.py verifying that:
      1. Default FleetConfig sends template.name == "default-template" and last_create_request.rootfs_image == image.
      2. AsyncSandboxFleet correctly propagates rootfs_image during async acquisition.
      3. driver.warm_pool_async(..., rootfs_image=image) passes rootfs_image and keys _warm_paused_pool[("custom-base", image)].
      4. Mapped images in image_templates send rootfs_image == "" and use the mapped template.
    • Added test_async_fleet.py covering concurrency, timeouts, and batch acquisition.
    • Ran full test suite: 124 passed, 29 skipped.
  • End to End

    • Tested with creating 1000 sandboxes with only python-slim image

…R 2)

Introduce AsyncSandboxFleet as an asyncio-native pool manager and refine handle
lifecycle semantics:

- Introduce AsyncSandboxFleet using pure asyncio queues, locks, and tasks
  for non-blocking concurrent rollout acquisition.
- Support async setup(tasks=None) to pre-warm the pool from a task list.
- Standardize SandboxHandle.exec() and exec_async() to return ExecResult.
- Add SandboxHandle.release_async() and recycle_async(), documenting that
  recycling suspends the actor without resetting in-guest filesystem or memory state.
- Change SandboxHandle context manager (__exit__ / __aexit__) to default
  to release() rather than teardown/destroy.
- Export AsyncSandboxFleet in ate_env top-level module.
- Add test_async_fleet.py covering concurrency, timeouts, and batch acquisition.
- Resolve task images to (template, rootfs_image) for Substrate backends and pass rootfs_image to Client.create.
@sunnylovestiramisu

Sunny (sunnylovestiramisu) commented Oct 14, 2026 •

Copy link
Copy Markdown
Collaborator Author

Addressed the nit pick in PR #89 (review)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review at 8716da2. Verdict: comment. No item blocks the merge. Five items are for before merge.

I read the diff and ran the suite: 126 tests pass, and ruff has no findings. CI passes. The branch merges without conflicts with main at 1c08a7a and with #91. The rootfs_image pass-through and the pool key (template, image) are correct. The default_timeout_s pass-through closes the #89 nit.

Blocking: none.

Before merge

M1. Digest requirement. types.py:249-251, README.md, tests/test_fleet_strategies.py:51.

This PR sends each unmapped Task.image to the server as rootfs_image. The server accepts only an image that is pinned by digest. Thus an image tag now fails at acquire. Probe against the fake stack:

Task(id="t-tag", image="docker.io/library/python:3.11")
SandboxStartError: Failed to create environment sb-…: image is not pinned by digest

Before this PR, the same task ran on default-template. This change is on the RL path. The VeRL example in #75 uses the tag sweb.eval.pytest-5221:latest, thus each acquire in that example fails.

Add "digest-pinned unless mapped in image_templates" to the Task.image docstring. Add a fleet section to the README. Show crane digest --platform linux/amd64 <tag> in that section. Use digest references in tests/test_fleet_strategies.py:51, :52 and :73, because users copy tests as examples. A later change can resolve tags to digests in the SDK.

M2. async_fleet.py:269 and :272. run() calls the synchronous handle.release() on the event-loop thread.

release() uses _sync, which blocks the calling thread until the delete RPC completes. Thus the event loop stops for one round trip for each task. An exception from release() in the except branch goes out of _worker and stops the full gather. Call await handle.release_async() in a finally block. Keep the exception of each task in results.

M3. async_fleet.py:212-218. acquire_batch uses a bare asyncio.gather.

When one acquire fails, the handles that succeeded stay allocated. The caller cannot get these handles. Call gather with return_exceptions=True. Release the handles that succeeded. Then raise the first exception.

The sync list form at fleet.py:196 has the same problem. Fix it in this PR or in a follow-up.

M4. async_fleet.py:93-106. setup() warms the plan entries one after the other.

The PR description says that setup() warms pools concurrently across task images. With rootfs images, the first use of each image creates a derived template. Thus N images give N sequential template creations. Run the entries with asyncio.gather, or correct the description. Each warm_pool_async call already starts its replicas in parallel.

M5. fleet.py:210 and async_fleet.py:161. Placeholder image for an unknown task id.

acquire("some-id") with an id that is not in the loaded tasks creates Task(id=task, image="default"). The SDK sends "default" to the server as rootfs_image, and the server rejects it. Probe: await fleet.acquire("no-such-task") gives the same image is not pinned by digest error. Before this PR, this call ran on default-template. The RL path loads Task objects with real images, thus this item does not affect it. Raise an error for an unknown id, or use an empty image that _ensure_template_for_image (fleet.py:291) resolves to (base, None).

Nits

  • backend/substrate.py:41. _WarmPoolDict accepts string keys. No code in the package uses a string key now. Use a plain dict[tuple[str, str | None], list[str]].
  • backend/substrate.py:432. Without rootfs_image, unwarm_pool_async(template_id) drains the pools of all images under that template. Write this in the docstring.
  • async_fleet.py:53. close() does not use the wait argument. Remove the argument. Also remove wait=False from __del__ at line 60.
  • async_fleet.py:34. The logger name is sandbox_sdk.async_fleet. The other modules log under ate_env.*. Use ate_env.async_fleet.
  • types.py:259-261. Three blank lines come before ExecResult. Use two.
  • PR description. It says "124 passed". The suite now gives 126 passed.

Verified

  • uv run --extra dev pytest: 126 passed, 29 skipped. The skipped tests need a live API. ruff check: no findings.
  • test_fleet_step_timeout_applies_to_handle and test_async_fleet_step_timeout_applies_to_handle cover the timeout pass-through.
  • Probe against the fake stack at 8716da2, with a digest-pinned image. Both fleets send template.name == "default-template" and the image as rootfs_image. A sync release() from a worker thread after an async acquire() works. run() over two tasks returns two results. close() returns.
  • The tag probe in M1 and the unknown-id probe in M5 ran at the same head.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants