Make Ec2InstanceProvider an async context manager; default impl uses aiobotocore - #15
Draft
art-dsit wants to merge 2 commits into
Draft
Make Ec2InstanceProvider an async context manager; default impl uses aiobotocore#15art-dsit wants to merge 2 commits into
art-dsit wants to merge 2 commits into
Conversation
sample_cleanup(interrupted=True) now returns early — matching the docker / k8s / proxmox sandboxes — and task_cleanup sweeps any leftover tracked instances at the end of the task. This covers Ctrl-C, setup-script failures, and any other path that fires sample_cleanup with interrupted=True. - Process-global tracker (set of (instance_id, region)) registers in sample_init, deregisters in successful sample_cleanup, drains in task_cleanup. Keyed only by (id, region) since inspect_ai always calls task_cleanup with task_name="shutdown". - task_cleanup respects --no-sandbox-cleanup: clears the tracker but leaves instances running, and surfaces a hint pointing at `inspect sandbox cleanup ec2`. - task_cleanup tolerates per-item terminate failures (one bad call doesn't block the rest). - DefaultEc2InstanceProvider.create_instance terminates the launched instance if the post-launch wait (instance-running / SSM-ready) fails, so a cloud-init / SSM timeout doesn't leak. - Pytest gains a `req_aws` marker (with --strict-markers) so the AWS-touching modules can be deselected via -m "not req_aws". Existing AWS-touching modules marked; the redundant `has_aws_creds` helper is removed. - New `tests/ec2sandboxtest/integration/` package with four cleanup-behaviour scenarios (happy path, interrupted → task_cleanup sweep, --no-sandbox-cleanup, multiple samples with mixed cleanup paths). - New `test_cleanup_unit.py` with mock-based unit tests for the tracker + sample/task cleanup contracts. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…aiobotocore Ec2InstanceProvider is now an async context manager: ``__aenter__`` is the place to initialize any loop-bound resources (``aiobotocore`` / ``httpx`` clients bind to the event loop they were created in) and ``__aexit__`` releases them. The sandbox enters the registered provider in ``task_init`` and exits it in ``task_cleanup``, so loop-bound clients have a lifetime tied to inspect_ai's task lifecycle. DefaultEc2InstanceProvider converted to use ``aiobotocore`` for its EC2 / SSM control-plane calls so it no longer blocks the event loop during run_instances / wait-for-running / wait-for-SSM / terminate / describe-instances. Clients are created in ``__aenter__`` and closed in ``__aexit__``; methods raise a clear ``RuntimeError`` if called on an un-entered provider. Sandbox internals: - New ``_active_provider`` ClassVar plus a ``_entered`` async context manager that yields the active provider when present and enters an ad-hoc one when not — tests and CLI cleanup work without going through task_init/task_cleanup. - ``task_cleanup`` exits the active provider in a ``finally`` so loop-bound clients are released even if the cleanup body raises. Tests: - Unit tests for ``DefaultEc2InstanceProvider`` rewritten to mock the aiobotocore "async with create_client(...)" shape. - Fake provider in ``test_cleanup_unit.py`` gains async-context- manager methods. - New ``test_eval_cleanup_on_setup_failure`` integration test runs a real ``inspect_ai.eval()`` with a setup script that exits non-zero, then asserts no instances leaked. This was unstable before the lifecycle change because the previous provider client would bind to one event loop and break when reused on another. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
art-dsit
force-pushed
the
issue-12-cleanup
branch
from
July 8, 2026 11:33
4f31437 to
0b1442e
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
SLOP ALERT
Summary
Stacked on top of #14. Once #14 lands, the base will switch to
main.Reshapes the
Ec2InstanceProvidercontract to acknowledge that providers are inherently loop-bound — its methods are alreadyasync def, and any async provider that usesaiobotocore/httpxclients holds resources that bind to the running event loop.Ec2InstanceProvideris now an async context manager.__aenter__is where loop-bound resources get initialized;__aexit__releases them. The sandbox enters the registered provider intask_initand exits it intask_cleanup, so loop-bound clients have a well-defined lifetime tied to inspect_ai's task lifecycle. Sync providers may implement both as no-ops.DefaultEc2InstanceProvideris converted to useaiobotocorefor its EC2 / SSM control-plane calls so it no longer blocks the event loop. Clients are created in__aenter__and closed in__aexit__; methods raise a clearRuntimeErrorif called on an un-entered provider.Why
Before the lifecycle change, an
inspect_ai.eval()run that uses a custom async provider was fine on its own — but anything that touched the provider on a different event loop (e.g. a second eval in the same process, or pytest-asyncio's per-function loops) failed with "Event loop is closed". The async-context-manager shape makes the loop boundary explicit: every entry creates a fresh client, every exit releases it.The default provider's sync-boto3-in-async-def pattern was also blocking the event loop on every EC2 / SSM call. Converting to
aiobotocoreremoves that. The runtime sandbox methods (exec,read_file,write_file) still use sync boto3 against SSM/S3 — out of scope here.What changes
__aenter__/__aexit__toEc2InstanceProvider.DefaultEc2InstanceProvider: aiobotocore-based, lifecycle-managed clients, partial-failure cleanup preserved.Ec2SandboxEnvironment: new_active_providerClassVar;task_initenters the registered provider,task_cleanupexits it in afinally; a_enteredasync-context-manager helper yields the active provider when present and enters an ad-hoc one when not (so test / CLI code works without a task lifecycle).aiobotocoredependency, plus mypy override.DefaultEc2InstanceProviderrewritten to mock the aiobotocore "async with create_client(...)" shape.test_eval_cleanup_on_setup_failureintegration test runs a realinspect_ai.eval()with a setup script that exits non-zero, then asserts no instances leaked. This was unstable before the lifecycle change.Test plan
uv run ruff check,uv run ruff format --check,uv run mypy— all pass.uv run pytest -m \"not req_aws\"— 15 unit tests pass (8 cleanup + 6 provider + 1 schema).uv run pytest -m req_aws tests/ec2sandboxtest/integration/— 5 integration tests pass against real EC2 (~8 minutes), including the failing-setup eval.🤖 Generated with Claude Code