Repository navigation
refactor(python): async-native fleet, and handle release semantics - #88
Sunny (sunnylovestiramisu) wants to merge 1 commit into
Conversation
e8e8e65 to
34776c2
Compare
…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.
34776c2 to
8716da2
Compare
|
Addressed the nit pick in PR #89 (review) |
Tomer Glottmann (tomergee)
left a comment
There was a problem hiding this comment.
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._WarmPoolDictaccepts string keys. No code in the package uses a string key now. Use a plaindict[tuple[str, str | None], list[str]].backend/substrate.py:432. Withoutrootfs_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 thewaitargument. Remove the argument. Also removewait=Falsefrom__del__at line 60.async_fleet.py:34. The logger name issandbox_sdk.async_fleet. The other modules log underate_env.*. Useate_env.async_fleet.types.py:259-261. Three blank lines come beforeExecResult. 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_handleandtest_async_fleet_step_timeout_applies_to_handlecover the timeout pass-through.- Probe against the fake stack at
8716da2, with a digest-pinned image. Both fleets sendtemplate.name == "default-template"and the image asrootfs_image. A syncrelease()from a worker thread after an asyncacquire()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.
Summary
This PR introduces
AsyncSandboxFleet, anasyncio-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.
AsyncSandboxFleetAsyncSandboxFleetusingasynciotasks, semaphores, and thread pools to enable non-blocking batch acquisition and warm-pool sizing without blocking the caller's event loop.await fleet.setup(tasks=None)to run pre-flight, compute aFleetPlan, and warm actor pools concurrently across task images.warm_images()andunwarm_image()to support rolling or sliding-window allocation strategies for bounded footprint execution.AsyncSandboxFleetfromate_env.2. Handle Lifecycle & Release Semantics
SandboxHandle.exec()andexec_async()to returnExecResult.SandboxHandle.release_async()andrecycle_async(), explicitly documenting that recycling suspends the actor without resetting in-guest filesystem or memory state.SandboxHandlecontext manager (__exit__/__aexit__) to default torelease()(terminate) rather than destroying the entire fleet run.3. Task Image & Rootfs Resolution (Addressing Review Feedback)
_ensure_template_for_imageResolution: Returns(template, rootfs_image):config.image_templatesresolve to(mapped_template, None).(config.template or "default-template", image), passing the task image through asrootfs_imagewhile using the configured template as base.BackendDriverSignature Updates: Addedrootfs_image: Optional[str] = Noneparameter acrossacquire,acquire_async,warm_pool,warm_pool_async, andunwarm_poolinBackendDriver,SubstrateBackendDriver, andMockBackendDriver.Client.createPropagation: Passedrootfs_image=rootfs_imageat bothclient.createinvocations inSubstrateBackendDriver(on-demand creation in_acquire_internaland pre-warm inwarm_pool_async).(template, image): Keyed_warm_paused_poolby(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.rootfs_imageontoRawSandboxInstanceandPlanEntry, and correctly preserved the pool key uponrelease_async(..., recycle=True).Task.imagebehavior detailing lookup inimage_templatesversus propagation asrootfs_image.Review Feedback Addressed
FleetConfigwithTask(image="...")senttemplate.name == "default-template"and emptyrootfs_image, ignoring the task container image.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._warm_paused_poolby(template, image).Verification
Unit Tests:
test_substrate_rootfs_image_sent_to_serverintest_substrate_driver.pyverifying that:FleetConfigsendstemplate.name == "default-template"andlast_create_request.rootfs_image == image.AsyncSandboxFleetcorrectly propagatesrootfs_imageduring async acquisition.driver.warm_pool_async(..., rootfs_image=image)passesrootfs_imageand keys_warm_paused_pool[("custom-base", image)].image_templatessendrootfs_image == ""and use the mapped template.test_async_fleet.pycovering concurrency, timeouts, and batch acquisition.End to End