Bundle ProvisionedInstance through the sandbox env

`region` was an explicit parameter on the env constructor, on the
sample-cleanup tracker (Set[Tuple[str, str]]), on the provider's
find_sandbox_instances(), and on the cli_cleanup fallback dance. All of
those were really one piece of information — instance metadata —
travelling alongside instance_id.

Pass ProvisionedInstance through whole:
- frozen so it can sit in a Set
- Ec2SandboxEnvironment.__init__ takes one (with flat self.X aliases
  for the rest of the file's reads)
- _tracked_instances is Set[ProvisionedInstance]
- find_sandbox_instances() drops its region param; the default impl
  derives it from its own config (with AWS_REGION fallback for the
  cli_cleanup path that builds an empty config)
- cli_cleanup drops the fallback_region env-var read

terminate_instance(instance_id, region) keeps its signature: region is
the instance's own region, which a control plane in a different region
(e.g. the Lambda-backed provider) needs to route the call.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: art-dsit

src/ec2sandbox/_ec2_sandbox_environment.py

@@ -9,7 +9,7 @@
 from importlib.metadata import entry_points
 from logging import getLogger
 from pathlib import Path
-from typing import Any, ClassVar, Dict, List, Set, Tuple, Union
+from typing import Any, ClassVar, Dict, List, Set, Union
 
 import boto3
 from botocore.exceptions import ClientError, WaiterError
@@ -32,6 +32,7 @@
     MARKER_TAG_KEY,
     DefaultEc2InstanceProvider,
     Ec2InstanceProvider,
+    ProvisionedInstance,
     get_ec2_instance_provider,
     get_provider_session,
 )
@@ -57,13 +58,13 @@ class Ec2SandboxEnvironment(SandboxEnvironment):
 
     _providers_loaded: ClassVar[bool] = False
 
-    # Process-global tracker of provisioned (instance_id, region) tuples.
+    # Process-global tracker of provisioned instances.
     # sample_init registers, successful sample_cleanup deregisters,
     # task_cleanup sweeps survivors (Ctrl-C, setup-script failures, etc.).
     # Matches the k8s / proxmox sandbox pattern of a shared tracker — note
     # that inspect_ai calls task_cleanup with task_name="shutdown" (a phase
     # marker, not the real task name) so we cannot key by task_name.
-    _tracked_instances: ClassVar[Set[Tuple[str, str]]] = set()
+    _tracked_instances: ClassVar[Set[ProvisionedInstance]] = set()
 
     @classmethod
     def set_session(cls, session: boto3.Session) -> None:
@@ -100,25 +101,22 @@ def _ensure_providers_loaded(cls) -> None:
                     "failed loading inspect_ai entry point %s", ep.value, exc_info=True
                 )
 
-    def __init__(
-        self,
-        instance_id: str,
-        region: str,
-        s3_bucket: str,
-        s3_key_prefix: str = "",
-    ):
-        self.instance_id = instance_id
-        self.region = region
-        self.s3_bucket = s3_bucket
-        self.s3_key_prefix = s3_key_prefix
+    def __init__(self, provisioned: ProvisionedInstance):
+        self.provisioned = provisioned
+        # Flat aliases so the rest of the class (and downstream readers)
+        # don't have to reach through self.provisioned for every call.
+        self.instance_id = provisioned.instance_id
+        self.region = provisioned.region
+        self.s3_bucket = provisioned.s3_bucket
+        self.s3_key_prefix = provisioned.s3_key_prefix
         session = self._get_session()
-        self.ssm_client = session.client("ssm", region_name=region)
+        self.ssm_client = session.client("ssm", region_name=self.region)
         self.s3_client = session.client(
             "s3",
-            region_name=region,
-            endpoint_url=f"https://s3.{region}.amazonaws.com",
+            region_name=self.region,
+            endpoint_url=f"https://s3.{self.region}.amazonaws.com",
         )
-        self.ec2_client = session.client("ec2", region_name=region)
+        self.ec2_client = session.client("ec2", region_name=self.region)
 
     @classmethod
     @override
@@ -209,15 +207,8 @@ async def sample_init(
             result.region,
         )
 
-        cls._tracked_instances.add((result.instance_id, result.region))
-
-        environment = Ec2SandboxEnvironment(
-            instance_id=result.instance_id,
-            region=result.region,
-            s3_bucket=result.s3_bucket,
-            s3_key_prefix=result.s3_key_prefix,
-        )
-        return {"default": environment}
+        cls._tracked_instances.add(result)
+        return {"default": Ec2SandboxEnvironment(result)}
 
     @classmethod
     @override
@@ -241,7 +232,7 @@ async def sample_cleanup(
             if not isinstance(env, Ec2SandboxEnvironment):
                 continue
             await provider.terminate_instance(env.instance_id, env.region)
-            cls._tracked_instances.discard((env.instance_id, env.region))
+            cls._tracked_instances.discard(env.provisioned)
         return None
 
     @classmethod
@@ -270,15 +261,15 @@ async def task_cleanup(
             return None
 
         provider, _ = cls._resolve_provider(config)
-        for instance_id, region in tracked:
+        for inst in tracked:
             try:
-                await provider.terminate_instance(instance_id, region)
+                await provider.terminate_instance(inst.instance_id, inst.region)
             except Exception as e:
                 # Per-item: one bad termination must not block the others.
                 cls.logger.warning(
                     "task_cleanup: failed to terminate %s in %s: %s",
-                    instance_id,
-                    region,
+                    inst.instance_id,
+                    inst.region,
                     e,
                 )
         return None
@@ -347,8 +338,7 @@ async def cli_cleanup(cls, id: str | None) -> None:
 
         provider, _ = cls._resolve_provider(None)
 
-        fallback_region = os.getenv("AWS_REGION", os.getenv("AWS_DEFAULT_REGION", ""))
-        instances = await provider.find_sandbox_instances(fallback_region)
+        instances = await provider.find_sandbox_instances()
 
         if not instances:
             print("\nNo EC2 sandbox instances found to clean up.\n")
@@ -370,8 +360,7 @@ async def cli_cleanup(cls, id: str | None) -> None:
             return
 
         for inst in instances:
-            region = inst.region or fallback_region
-            await provider.terminate_instance(inst.instance_id, region)
+            await provider.terminate_instance(inst.instance_id, inst.region)
 
     @staticmethod
     def _confirm_cleanup() -> bool:

src/ec2sandbox/_instance_provider.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import logging
+import os
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, ClassVar, Protocol, runtime_checkable
 
@@ -35,12 +36,13 @@
 MARKER_TAG_KEY = "inspect_sandbox"
 
 
-@dataclass
+@dataclass(frozen=True)
 class ProvisionedInstance:
     """Result of provisioning an EC2 sandbox instance.
 
     Contains everything the :class:`Ec2SandboxEnvironment` needs for
     runtime operations (exec / read_file / write_file via SSM + S3).
+    Frozen so it can be used as a set element by the cleanup tracker.
     """
 
     instance_id: str
@@ -49,7 +51,7 @@ class ProvisionedInstance:
     s3_key_prefix: str = ""
 
 
-@dataclass
+@dataclass(frozen=True)
 class SandboxInstanceInfo:
     """Metadata for a discovered sandbox instance (used by cli_cleanup)."""
 
@@ -93,10 +95,15 @@ async def create_instance(
         ...
 
     async def terminate_instance(self, instance_id: str, region: str) -> None:
-        """Terminate an EC2 instance."""
+        """Terminate an EC2 instance.
+
+        ``region`` is the instance's own region (from ``create_instance``'s
+        :class:`ProvisionedInstance`); a control plane in a different
+        region (e.g. a Lambda) uses it to route the call.
+        """
         ...
 
-    async def find_sandbox_instances(self, region: str) -> list[SandboxInstanceInfo]:
+    async def find_sandbox_instances(self) -> list[SandboxInstanceInfo]:
         """Find running sandbox instances for interactive cleanup."""
         ...
 
@@ -203,7 +210,8 @@ class DefaultEc2InstanceProvider:
     etc.) from the supplied :class:`Ec2SandboxEnvironmentConfig` at the
     point they are needed — ``create_instance`` requires the full set,
     while ``terminate_instance`` and ``find_sandbox_instances`` only
-    need the supplied ``region``.
+    need a region (the instance's own region for terminate; the
+    configured region for find).
     """
 
     # Process-global Ubuntu 24.04 AMI cache keyed on region. Canonical's
@@ -315,7 +323,14 @@ async def terminate_instance(self, instance_id: str, region: str) -> None:
         ec2 = self._session.client("ec2", region_name=region or None)
         ec2.terminate_instances(InstanceIds=[instance_id])
 
-    async def find_sandbox_instances(self, region: str) -> list[SandboxInstanceInfo]:
+    async def find_sandbox_instances(self) -> list[SandboxInstanceInfo]:
+        # cli_cleanup builds this provider with an empty config (no region),
+        # so fall back to AWS_REGION / AWS_DEFAULT_REGION as the session does.
+        region = (
+            self._config.region
+            or os.getenv("AWS_REGION")
+            or os.getenv("AWS_DEFAULT_REGION")
+        )
         ec2 = self._session.client("ec2", region_name=region or None)
         response = ec2.describe_instances(
             Filters=[
@@ -338,7 +353,7 @@ async def find_sandbox_instances(self, region: str) -> list[SandboxInstanceInfo]
                     SandboxInstanceInfo(
                         instance_id=instance["InstanceId"],
                         name=name,
-                        region=region,
+                        region=region or "",
                     )
                 )
         return results

tests/ec2sandboxtest/integration/test_cleanup_behaviour.py

@@ -30,6 +30,10 @@
 REGION = "eu-west-2"
 
 
+def _tracked_ids() -> set[str]:
+    return {p.instance_id for p in Ec2SandboxEnvironment._tracked_instances}
+
+
 async def _provision(
     config: Ec2SandboxEnvironmentConfig,
     task_name: str,
@@ -47,7 +51,7 @@ async def test_happy_path_terminates_via_sample_cleanup(
     """sample_cleanup(interrupted=False) terminates the instance and clears the tracker."""  # noqa: E501
     envs, inst_id = await _provision(ec2_config, "test_happy")
     try:
-        assert (inst_id, REGION) in Ec2SandboxEnvironment._tracked_instances
+        assert inst_id in _tracked_ids()
 
         await Ec2SandboxEnvironment.sample_cleanup(
             task_name="test_happy",
@@ -56,7 +60,7 @@ async def test_happy_path_terminates_via_sample_cleanup(
             interrupted=False,
         )
 
-        assert (inst_id, REGION) not in Ec2SandboxEnvironment._tracked_instances
+        assert inst_id not in _tracked_ids()
         assert wait_until_terminated(inst_id) in ("terminated", "shutting-down", None)
     finally:
         # Defensive: even if assertions failed, don't leak the instance.
@@ -80,14 +84,14 @@ async def test_interrupted_sample_cleanup_skips_then_task_cleanup_sweeps(
 
         time.sleep(5)
         assert instance_state(inst_id) in ("pending", "running")
-        assert (inst_id, REGION) in Ec2SandboxEnvironment._tracked_instances
+        assert inst_id in _tracked_ids()
 
         # inspect_ai always calls task_cleanup with task_name="shutdown".
         await Ec2SandboxEnvironment.task_cleanup(
             task_name="shutdown", config=ec2_config, cleanup=True
         )
 
-        assert (inst_id, REGION) not in Ec2SandboxEnvironment._tracked_instances
+        assert inst_id not in _tracked_ids()
         assert wait_until_terminated(inst_id) in ("terminated", "shutting-down", None)
     finally:
         if instance_state(inst_id) in ("pending", "running"):
@@ -124,9 +128,7 @@ async def test_multiple_samples_one_interrupted(
     envs_a, inst_a = await _provision(ec2_config, "test_multi")
     envs_b, inst_b = await _provision(ec2_config, "test_multi")
     try:
-        assert {(inst_a, REGION), (inst_b, REGION)}.issubset(
-            Ec2SandboxEnvironment._tracked_instances
-        )
+        assert {inst_a, inst_b}.issubset(_tracked_ids())
 
         # Sample A succeeds.
         await Ec2SandboxEnvironment.sample_cleanup(
@@ -135,8 +137,8 @@ async def test_multiple_samples_one_interrupted(
             environments=envs_a,
             interrupted=False,
         )
-        assert (inst_a, REGION) not in Ec2SandboxEnvironment._tracked_instances
-        assert (inst_b, REGION) in Ec2SandboxEnvironment._tracked_instances
+        assert inst_a not in _tracked_ids()
+        assert inst_b in _tracked_ids()
 
         # Sample B interrupted.
         await Ec2SandboxEnvironment.sample_cleanup(
@@ -145,14 +147,14 @@ async def test_multiple_samples_one_interrupted(
             environments=envs_b,
             interrupted=True,
         )
-        assert (inst_b, REGION) in Ec2SandboxEnvironment._tracked_instances
+        assert inst_b in _tracked_ids()
 
         # task_cleanup sweeps the leftover.
         await Ec2SandboxEnvironment.task_cleanup(
             task_name="shutdown", config=ec2_config, cleanup=True
         )
 
-        assert (inst_b, REGION) not in Ec2SandboxEnvironment._tracked_instances
+        assert inst_b not in _tracked_ids()
         assert wait_until_terminated(inst_a) in ("terminated", "shutting-down", None)
         assert wait_until_terminated(inst_b) in ("terminated", "shutting-down", None)
     finally:

tests/ec2sandboxtest/test_cleanup_unit.py

@@ -31,7 +31,7 @@ async def terminate_instance(self, instance_id: str, region: str) -> None:
             raise RuntimeError(f"simulated terminate failure for {instance_id}")
         self.terminated.append((instance_id, region))
 
-    async def find_sandbox_instances(self, region):  # pragma: no cover
+    async def find_sandbox_instances(self):  # pragma: no cover
         return []
 
 
@@ -61,9 +61,15 @@ async def create_instance(**kwargs):
         )
 
 
+def _provisioned(instance_id: str, region: str = "eu-west-2") -> ProvisionedInstance:
+    return ProvisionedInstance(
+        instance_id=instance_id, region=region, s3_bucket="bucket-1"
+    )
+
+
 async def test_sample_init_registers_instance_in_tracker():
     await _seed_environment("task_a", instance_id="i-111")
-    assert Ec2SandboxEnvironment._tracked_instances == {("i-111", "eu-west-2")}
+    assert Ec2SandboxEnvironment._tracked_instances == {_provisioned("i-111")}
 
 
 async def test_sample_cleanup_interrupted_leaves_tracker_intact():
@@ -79,7 +85,7 @@ async def test_sample_cleanup_interrupted_leaves_tracker_intact():
         )
 
     assert terminator.terminated == []
-    assert Ec2SandboxEnvironment._tracked_instances == {("i-111", "eu-west-2")}
+    assert Ec2SandboxEnvironment._tracked_instances == {_provisioned("i-111")}
 
 
 async def test_sample_cleanup_success_terminates_and_deregisters():
@@ -107,8 +113,8 @@ async def test_task_cleanup_sweeps_remaining_tracked_instances():
     await _seed_environment("task_a", instance_id="i-111")
     await _seed_environment("task_b", instance_id="i-222")
     assert Ec2SandboxEnvironment._tracked_instances == {
-        ("i-111", "eu-west-2"),
-        ("i-222", "eu-west-2"),
+        _provisioned("i-111"),
+        _provisioned("i-222"),
     }
 
     terminator = _FakeProvider()
@@ -195,4 +201,4 @@ async def test_sample_cleanup_only_deregisters_its_own_environments():
         )
 
     assert terminator.terminated == [("i-aaa", "eu-west-2")]
-    assert Ec2SandboxEnvironment._tracked_instances == {("i-bbb", "eu-west-2")}
+    assert Ec2SandboxEnvironment._tracked_instances == {_provisioned("i-bbb")}