feat: add multi-cloud VM support with AWS backend and VMProvider protocol (#66)

* feat: add multi-cloud VM support with AWS backend and VMProvider protocol

- Create VMProvider Protocol (typing.Protocol) for cloud-agnostic VM management
- Create AWSVMManager with boto3 for EC2 lifecycle (create, delete, start, stop)
- Add resource_scope/ssh_username properties to AzureVMManager
- Add list_pool_resources/cleanup_pool_resources to AzureVMManager
- Parameterize pool.py SSH calls and scripts with username/home_dir
- Add --cloud flag (azure|aws) to all pool CLI commands
- Add cloud_provider/aws_region to config.py settings
- Add boto3 optional dependency (openadapt-evals[aws])
- Update tests for WAA_START_SCRIPT_TEMPLATE rename

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: address review findings in AWS VM backend

- Fix DOCKER_SETUP_SCRIPT_WITH_ACR daemon.json double-brace corruption
  that produced invalid JSON ({{"data-root"...}}) breaking Docker start
- Use .metal instance types for AWS (KVM/nested virt required for QEMU)
- Fix region mismatch: update self.region and invalidate cached clients
  when create_vm uses a different region than the manager default
- Fix hardcoded "azureuser" in pool-wait diagnostic message
- Set AWSVMManager = None on ImportError so `import *` doesn't raise
- Only delete pool registry on successful cleanup (prevents orphaned
  cloud resources when deletion fails)
- Remove unused `time` import from aws_vm.py

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: address second review findings

- Fix pool-vnc/pool-logs/pool-exec hardcoded azureuser: read ssh_username
  from pool registry with backward-compatible default
- Store ssh_username in VMPool dataclass and persist to registry on create
- Move set_auto_shutdown after SSH is available (was racing with boot)
- Fix cleanup_pool_resources: handle raw instance IDs and allocation IDs
  for resources without Name tags (prevents orphaned resources)
- Narrow key pair exception handling: re-raise unless InvalidKeyPair.NotFound
- Add TODO for restricting SSH security group to user's IP

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: restore ssh_username on registry load, fix EIP disassociate API

- Add ssh_username to VMPoolRegistry.load() so it persists across
  process restarts (was silently reverting to "azureuser" default)
- Fix disassociate_address for raw allocation IDs: look up AssociationId
  via describe_addresses first (disassociate_address does not accept
  AllocationId parameter)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: abrichr

openadapt_evals/benchmarks/vm_cli.py

@@ -76,6 +76,44 @@ def _get_resource_group() -> str:
 
 
 RESOURCE_GROUP = _get_resource_group()
+
+
+def _get_default_cloud() -> str:
+    """Get default cloud provider from config."""
+    try:
+        from openadapt_evals.config import settings
+
+        return settings.cloud_provider
+    except Exception:
+        return "azure"
+
+
+def _get_pool_ssh_username(pool: dict) -> str:
+    """Get SSH username from pool registry, defaulting to azureuser for backward compat."""
+    return pool.get("ssh_username", "azureuser")
+
+
+def _create_vm_manager(cloud: str | None = None, resource_group: str | None = None):
+    """Factory to create the appropriate VM manager based on cloud provider.
+
+    Args:
+        cloud: Cloud provider ("azure" or "aws"). If None, uses config default.
+        resource_group: Azure resource group (ignored for AWS).
+
+    Returns:
+        VMProvider instance (AzureVMManager or AWSVMManager).
+    """
+    cloud = cloud or _get_default_cloud()
+    if cloud == "aws":
+        from openadapt_evals.infrastructure.aws_vm import AWSVMManager
+
+        return AWSVMManager()
+    else:
+        from openadapt_evals.infrastructure.azure_vm import AzureVMManager
+
+        return AzureVMManager(resource_group=resource_group or RESOURCE_GROUP)
+
+
 # Custom WAA image built from waa_deploy/Dockerfile
 # Uses dockurr/windows:latest as base (with proper ISO download) + WAA components
 DOCKER_IMAGE = "waa-auto:latest"
@@ -518,11 +556,10 @@ def cmd_pool_status(args):
     """Show status of all VMs in the current pool."""
     init_logging()
 
-    from openadapt_evals.infrastructure.azure_vm import AzureVMManager
     from openadapt_evals.infrastructure.pool import PoolManager
     from openadapt_evals.infrastructure.vm_monitor import VMMonitor, VMConfig
 
-    vm_manager = AzureVMManager(resource_group=RESOURCE_GROUP)
+    vm_manager = _create_vm_manager(getattr(args, "cloud", None))
     manager = PoolManager(vm_manager=vm_manager, log_fn=log)
     pool = manager.status()
 
@@ -592,10 +629,9 @@ def cmd_delete_pool(args):
     init_logging()
     from concurrent.futures import ThreadPoolExecutor, as_completed
 
-    from openadapt_evals.infrastructure.azure_vm import AzureVMManager
     from openadapt_evals.infrastructure.pool import PoolManager
 
-    vm_manager = AzureVMManager(resource_group=RESOURCE_GROUP)
+    vm_manager = _create_vm_manager(getattr(args, "cloud", None))
     manager = PoolManager(vm_manager=vm_manager, log_fn=log)
     pool = manager.status()
 
@@ -642,15 +678,14 @@ def cmd_pool_create(args):
     Uses ThreadPoolExecutor for concurrent VM creation.
     """
     init_logging()
-    from openadapt_evals.infrastructure.azure_vm import AzureVMManager
     from openadapt_evals.infrastructure.pool import PoolManager
 
     num_workers = getattr(args, "workers", 3)
     auto_shutdown_hours = getattr(args, "auto_shutdown_hours", 4)
     use_acr = getattr(args, "use_acr", False)
     image_id = getattr(args, "image", None)
 
-    vm_manager = AzureVMManager(resource_group=RESOURCE_GROUP)
+    vm_manager = _create_vm_manager(getattr(args, "cloud", None))
     manager = PoolManager(vm_manager=vm_manager, log_fn=log)
 
     try:
@@ -673,13 +708,12 @@ def cmd_pool_wait(args):
     and the WAA server to respond.
     """
     init_logging()
-    from openadapt_evals.infrastructure.azure_vm import AzureVMManager
     from openadapt_evals.infrastructure.pool import PoolManager
 
     timeout_minutes = getattr(args, "timeout", 30)
     no_start = getattr(args, "no_start", False)
 
-    vm_manager = AzureVMManager(resource_group=RESOURCE_GROUP)
+    vm_manager = _create_vm_manager(getattr(args, "cloud", None))
     manager = PoolManager(vm_manager=vm_manager, log_fn=log)
 
     try:
@@ -700,15 +734,14 @@ def cmd_pool_run(args):
     in parallel. Collects results from all workers.
     """
     init_logging()
-    from openadapt_evals.infrastructure.azure_vm import AzureVMManager
     from openadapt_evals.infrastructure.pool import PoolManager
 
     num_tasks = getattr(args, "tasks", 10)
     agent = getattr(args, "agent", "navi")
     model = getattr(args, "model", "gpt-4o-mini")
     api_key = getattr(args, "api_key", None)
 
-    vm_manager = AzureVMManager(resource_group=RESOURCE_GROUP)
+    vm_manager = _create_vm_manager(getattr(args, "cloud", None))
     manager = PoolManager(vm_manager=vm_manager, log_fn=log)
 
     try:
@@ -731,10 +764,9 @@ def cmd_pool_cleanup(args):
     weren't properly deleted.
     """
     init_logging()
-    from openadapt_evals.infrastructure.azure_vm import AzureVMManager
     from openadapt_evals.infrastructure.pool import PoolManager
 
-    vm_manager = AzureVMManager(resource_group=RESOURCE_GROUP)
+    vm_manager = _create_vm_manager(getattr(args, "cloud", None))
     manager = PoolManager(vm_manager=vm_manager, log_fn=log)
 
     confirm = not getattr(args, "yes", False)
@@ -752,7 +784,6 @@ def cmd_pool_auto(args):
     If a pool already exists, skips creation and resumes from wait → run.
     """
     init_logging()
-    from openadapt_evals.infrastructure.azure_vm import AzureVMManager
     from openadapt_evals.infrastructure.pool import PoolManager
 
     num_workers = getattr(args, "workers", 1)
@@ -763,7 +794,7 @@ def cmd_pool_auto(args):
     model = getattr(args, "model", "gpt-4o-mini")
     api_key = getattr(args, "api_key", None)
 
-    vm_manager = AzureVMManager(resource_group=RESOURCE_GROUP)
+    vm_manager = _create_vm_manager(getattr(args, "cloud", None))
     manager = PoolManager(vm_manager=vm_manager, log_fn=log)
 
     try:
@@ -820,10 +851,9 @@ def cmd_pool_pause(args):
     instead of recreating from scratch (~42 min). Idle cost ~$0.25/day.
     """
     init_logging()
-    from openadapt_evals.infrastructure.azure_vm import AzureVMManager
     from openadapt_evals.infrastructure.pool import PoolManager
 
-    vm_manager = AzureVMManager(resource_group=RESOURCE_GROUP)
+    vm_manager = _create_vm_manager(getattr(args, "cloud", None))
     manager = PoolManager(vm_manager=vm_manager, log_fn=log)
 
     try:
@@ -842,12 +872,11 @@ def cmd_pool_resume(args):
     (~5 min vs ~42 min).
     """
     init_logging()
-    from openadapt_evals.infrastructure.azure_vm import AzureVMManager
     from openadapt_evals.infrastructure.pool import PoolManager
 
     timeout_minutes = getattr(args, "timeout", 10)
 
-    vm_manager = AzureVMManager(resource_group=RESOURCE_GROUP)
+    vm_manager = _create_vm_manager(getattr(args, "cloud", None))
     manager = PoolManager(vm_manager=vm_manager, log_fn=log)
 
     try:
@@ -884,6 +913,8 @@ def cmd_pool_vnc(args):
     worker_name = getattr(args, "worker", None)
     all_workers = getattr(args, "all", False)
 
+    ssh_user = _get_pool_ssh_username(pool)
+
     if all_workers:
         # Set up tunnels for all workers
         log("POOL-VNC", f"Setting up VNC tunnels for {len(workers)} workers...")
@@ -912,7 +943,7 @@ def cmd_pool_vnc(args):
                     "-N",
                     "-L",
                     f"{local_port}:localhost:8006",
-                    f"azureuser@{ip}",
+                    f"{ssh_user}@{ip}",
                 ],
                 stdout=subprocess.DEVNULL,
                 stderr=subprocess.DEVNULL,
@@ -989,7 +1020,7 @@ def cmd_pool_vnc(args):
                 "-N",
                 "-L",
                 f"{local_port}:localhost:8006",
-                f"azureuser@{ip}",
+                f"{ssh_user}@{ip}",
             ],
             stdout=subprocess.DEVNULL,
             stderr=subprocess.DEVNULL,
@@ -1042,6 +1073,7 @@ def cmd_pool_logs(args):
         print("ERROR: Pool has no workers.")
         return 1
 
+    ssh_user = _get_pool_ssh_username(pool)
     pool_id = pool.get("pool_id", "unknown")
     print(f"[pool-logs] Streaming logs from {len(workers)} workers (pool: {pool_id})")
     print("[pool-logs] Press Ctrl+C to stop\n", flush=True)
@@ -1060,7 +1092,7 @@ def stream_worker_logs(worker_name: str, ip: str):
             "UserKnownHostsFile=/dev/null",
             "-o",
             "LogLevel=ERROR",
-            f"azureuser@{ip}",
+            f"{ssh_user}@{ip}",
             "docker logs -f winarena",
         ]
         try:
@@ -1132,6 +1164,7 @@ def cmd_pool_exec(args):
         log("POOL-EXEC", "ERROR: Pool has no workers.")
         return 1
 
+    ssh_user = _get_pool_ssh_username(pool)
     cmd = getattr(args, "cmd", None)
     docker = getattr(args, "docker", False)
     worker_filter = getattr(args, "worker", None)
@@ -1164,7 +1197,7 @@ def cmd_pool_exec(args):
 
         try:
             result = subprocess.run(
-                ["ssh", *SSH_OPTS, f"azureuser@{ip}", full_cmd],
+                ["ssh", *SSH_OPTS, f"{ssh_user}@{ip}", full_cmd],
                 capture_output=True,
                 text=True,
                 timeout=60,
@@ -7623,10 +7656,18 @@ def main():
     p_delete = subparsers.add_parser("delete", help="Delete VM and all resources")
     p_delete.set_defaults(func=cmd_delete)
 
+    # Shared --cloud argument for pool commands
+    _cloud_kwargs = {
+        "choices": ["azure", "aws"],
+        "default": None,
+        "help": "Cloud provider (default: from config, usually azure)",
+    }
+
     # pool-status
     p_pool_status = subparsers.add_parser(
         "pool-status", help="Show status of all VMs in the current pool"
     )
+    p_pool_status.add_argument("--cloud", **_cloud_kwargs)
     p_pool_status.add_argument(
         "--probe",
         action="store_true",
@@ -7636,13 +7677,15 @@ def main():
 
     # delete-pool
     p_delete_pool = subparsers.add_parser("delete-pool", help="Delete all VMs in the current pool")
+    p_delete_pool.add_argument("--cloud", **_cloud_kwargs)
     p_delete_pool.add_argument("-y", "--yes", action="store_true", help="Skip confirmation")
     p_delete_pool.set_defaults(func=cmd_delete_pool)
 
     # pool-create
     p_pool_create = subparsers.add_parser(
         "pool-create", help="Create a pool of VMs for parallel WAA evaluation"
     )
+    p_pool_create.add_argument("--cloud", **_cloud_kwargs)
     p_pool_create.add_argument(
         "--workers",
         "-n",
@@ -7671,6 +7714,7 @@ def main():
     p_pool_wait = subparsers.add_parser(
         "pool-wait", help="Wait for all pool workers to have WAA ready"
     )
+    p_pool_wait.add_argument("--cloud", **_cloud_kwargs)
     p_pool_wait.add_argument(
         "--timeout", "-t", type=int, default=30, help="Timeout in minutes (default: 30)"
     )
@@ -7685,6 +7729,7 @@ def main():
     p_pool_run = subparsers.add_parser(
         "pool-run", help="Run WAA benchmark tasks distributed across pool workers"
     )
+    p_pool_run.add_argument("--cloud", **_cloud_kwargs)
     p_pool_run.add_argument(
         "--tasks",
         "-n",
@@ -7703,6 +7748,7 @@ def main():
     p_pool_cleanup = subparsers.add_parser(
         "pool-cleanup", help="Clean up orphaned pool resources (VMs, NICs, IPs, disks)"
     )
+    p_pool_cleanup.add_argument("--cloud", **_cloud_kwargs)
     p_pool_cleanup.add_argument("-y", "--yes", action="store_true", help="Skip confirmation")
     p_pool_cleanup.set_defaults(func=cmd_pool_cleanup)
 
@@ -7711,6 +7757,7 @@ def main():
         "pool-auto",
         help="Fully automated: create VMs → wait for WAA → run benchmark",
     )
+    p_pool_auto.add_argument("--cloud", **_cloud_kwargs)
     p_pool_auto.add_argument(
         "--workers", "-w", type=int, default=1, help="Number of worker VMs (default: 1)"
     )
@@ -7738,13 +7785,15 @@ def main():
         "pool-pause",
         help="Deallocate pool VMs (stops compute billing, keeps disks ~$0.25/day)",
     )
+    p_pool_pause.add_argument("--cloud", **_cloud_kwargs)
     p_pool_pause.set_defaults(func=cmd_pool_pause)
 
     # pool-resume
     p_pool_resume = subparsers.add_parser(
         "pool-resume",
         help="Resume a paused pool (start VMs, wait for WAA ~5 min)",
     )
+    p_pool_resume.add_argument("--cloud", **_cloud_kwargs)
     p_pool_resume.add_argument(
         "--timeout",
         "-t",

openadapt_evals/config.py

@@ -45,6 +45,10 @@ class Settings(BaseSettings):
     azure_checkpoints_container: str = "checkpoints"
     azure_comparisons_container: str = "comparisons"
 
+    # Multi-cloud settings
+    cloud_provider: str = "azure"  # "azure" or "aws"
+    aws_region: str = "us-east-1"
+
     model_config = {
         "env_file": ".env",
         "env_file_encoding": "utf-8",

openadapt_evals/infrastructure/__init__.py

@@ -1,7 +1,9 @@
 """Infrastructure components for VM management and monitoring.
 
 This module provides:
+- VMProvider: Cloud-agnostic VM provider protocol
 - AzureVMManager: Azure VM lifecycle management (SDK + CLI fallback)
+- AWSVMManager: AWS EC2 lifecycle management (boto3)
 - PoolManager: Multi-VM pool orchestration
 - VMMonitor: Azure VM status monitoring
 - AzureOpsTracker: Azure operation logging
@@ -12,10 +14,15 @@
     ```python
     from openadapt_evals.infrastructure import AzureVMManager, PoolManager
 
-    # Manage VMs
+    # Manage VMs (Azure)
     vm = AzureVMManager()
     ip = vm.get_vm_ip("waa-eval-vm")
 
+    # Or use AWS
+    from openadapt_evals.infrastructure import AWSVMManager
+    vm = AWSVMManager(region="us-east-1")
+    pool = PoolManager(vm_manager=vm)
+
     # Create and manage pools
     pool = PoolManager()
     pool.create(workers=3)
@@ -42,15 +49,23 @@
 from openadapt_evals.infrastructure.ssh_tunnel import SSHTunnelManager, get_tunnel_manager
 from openadapt_evals.infrastructure.vm_ip import resolve_vm_ip
 from openadapt_evals.infrastructure.vm_monitor import VMMonitor, VMConfig
+from openadapt_evals.infrastructure.vm_provider import VMProvider
+
+try:
+    from openadapt_evals.infrastructure.aws_vm import AWSVMManager
+except ImportError:
+    AWSVMManager = None  # boto3 not installed; use `pip install openadapt-evals[aws]`
 
 __all__ = [
+    "AWSVMManager",
     "AzureOpsTracker",
     "AzureVMManager",
     "PoolManager",
     "PoolRunResult",
     "QEMUResetManager",
     "VMMonitor",
     "VMConfig",
+    "VMProvider",
     "SSHTunnelManager",
     "compare_screenshots",
     "get_tunnel_manager",