feat: add GPU training automation for verl-agent E2E workflow

- Add GPU_VM_SIZE_FALLBACKS to azure_vm.py (NC48ads_A100_v4, NC24ads, NC12s_v3)
- Add GPU_INSTANCE_TYPE_FALLBACKS to aws_vm.py (p3.8xlarge, g5.12xlarge, p3.2xlarge)
- Update find_available_size_and_region(gpu=True) on both providers + protocol
- Add scripts/setup_gpu_training.sh: installs conda, vLLM, flash-attn, verl-agent
- Add scripts/train_verl_e2e.py: provisions GPU VM, uploads setup, launches training
- Add oa-vm gpu-setup and gpu-train CLI commands

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

Author: abrichr

openadapt_evals/benchmarks/vm_cli.py

@@ -7784,6 +7784,181 @@ def cmd_resources(args):
     return 0
 
 
+# =============================================================================
+# GPU Training Commands
+# =============================================================================
+
+
+def _get_gpu_vm_manager(cloud: str):
+    """Get VM manager for GPU training."""
+    if cloud == "azure":
+        from openadapt_evals.infrastructure.azure_vm import AzureVMManager
+        return AzureVMManager()
+    elif cloud == "aws":
+        from openadapt_evals.infrastructure.aws_vm import AWSVMManager
+        return AWSVMManager()
+    raise ValueError(f"Unknown cloud: {cloud}")
+
+
+def cmd_gpu_setup(args):
+    """Provision a GPU VM and install verl-agent for RL training."""
+    import time
+    from pathlib import Path
+
+    from openadapt_evals.infrastructure.azure_vm import ssh_run
+
+    cloud = getattr(args, "cloud", "azure")
+    vm = _get_gpu_vm_manager(cloud)
+    username = vm.ssh_username
+    gpu_vm_name = "verl-train-00"
+
+    if args.gpu_ip:
+        ip = args.gpu_ip
+        print(f"Using existing GPU VM: {ip}")
+    else:
+        print("Finding available GPU VM size...")
+        vm_size, region, cost = vm.find_available_size_and_region(gpu=True)
+        print(f"Selected: {vm_size} (${cost:.2f}/hr) in {region}")
+
+        if args.dry_run:
+            print(f"[DRY RUN] Would provision {vm_size} in {region}")
+            return 0
+
+        print(f"Creating GPU VM '{gpu_vm_name}'...")
+        info = vm.create_vm(name=gpu_vm_name, region=region, size=vm_size)
+        ip = info.get("publicIpAddress") or vm.get_vm_ip(gpu_vm_name)
+        vm.set_auto_shutdown(gpu_vm_name, hours=6)
+
+        # Wait for SSH
+        print("Waiting for SSH...")
+        for _ in range(30):
+            try:
+                result = ssh_run(ip, "echo ready", username=username, stream=False)
+                if result.returncode == 0:
+                    break
+            except Exception:
+                pass
+            time.sleep(10)
+        else:
+            print(f"ERROR: SSH not ready after 5 minutes: {ip}")
+            return 1
+
+    # Upload and run setup script
+    setup_script = Path(__file__).parent.parent.parent / "scripts" / "setup_gpu_training.sh"
+    if not setup_script.exists():
+        print(f"ERROR: Setup script not found: {setup_script}")
+        return 1
+
+    from openadapt_evals.infrastructure.azure_vm import SSH_OPTS
+    import subprocess
+
+    print("Uploading setup script...")
+    subprocess.run(
+        ["scp", *SSH_OPTS, str(setup_script), f"{username}@{ip}:/tmp/setup_gpu_training.sh"],
+        check=True,
+    )
+
+    print("Running setup (this may take 15-30 minutes)...")
+    result = ssh_run(ip, "bash /tmp/setup_gpu_training.sh", username=username, stream=True)
+    if result.returncode != 0:
+        print(f"ERROR: Setup failed with exit code {result.returncode}")
+        return 1
+
+    print(f"\nGPU VM ready at: {ip}")
+    print(f"SSH: ssh {username}@{ip}")
+    return 0
+
+
+def cmd_gpu_train(args):
+    """Launch verl-agent training on a GPU VM."""
+    import time
+    from pathlib import Path
+
+    from openadapt_evals.infrastructure.azure_vm import ssh_run
+
+    cloud = getattr(args, "cloud", "azure")
+    vm = _get_gpu_vm_manager(cloud)
+    username = vm.ssh_username
+    gpu_vm_name = "verl-train-00"
+
+    if args.gpu_ip:
+        ip = args.gpu_ip
+        print(f"Using existing GPU VM: {ip}")
+    else:
+        # Provision GPU VM
+        print("Finding available GPU VM size...")
+        vm_size, region, cost = vm.find_available_size_and_region(gpu=True)
+        print(f"Selected: {vm_size} (${cost:.2f}/hr) in {region}")
+        print(f"Creating GPU VM '{gpu_vm_name}'...")
+        info = vm.create_vm(name=gpu_vm_name, region=region, size=vm_size)
+        ip = info.get("publicIpAddress") or vm.get_vm_ip(gpu_vm_name)
+        vm.set_auto_shutdown(gpu_vm_name, hours=6)
+
+        # Wait for SSH
+        for _ in range(30):
+            try:
+                result = ssh_run(ip, "echo ready", username=username, stream=False)
+                if result.returncode == 0:
+                    break
+            except Exception:
+                pass
+            time.sleep(10)
+
+    # Setup if needed
+    if not args.skip_setup:
+        setup_script = Path(__file__).parent.parent.parent / "scripts" / "setup_gpu_training.sh"
+        if setup_script.exists():
+            from openadapt_evals.infrastructure.azure_vm import SSH_OPTS
+            import subprocess
+            subprocess.run(
+                ["scp", *SSH_OPTS, str(setup_script), f"{username}@{ip}:/tmp/setup_gpu_training.sh"],
+                check=True,
+            )
+            result = ssh_run(ip, "bash /tmp/setup_gpu_training.sh", username=username, stream=True)
+            if result.returncode != 0:
+                print(f"ERROR: Setup failed")
+                return 1
+
+    # Launch training
+    train_cmd = (
+        f"cd ~/verl-agent && "
+        f"conda activate verl-agent && "
+        f"python3 -m verl.trainer.main_ppo "
+        f"algorithm.adv_estimator={args.algorithm} "
+        f"actor_rollout_ref.model.path={args.model} "
+        f"actor_rollout_ref.rollout.name=vllm "
+        f"actor_rollout_ref.rollout.tensor_model_parallel_size={args.n_gpus} "
+        f"env.env_name=openadapt_evals.adapters.verl_env.WAADesktopEnv "
+        f"env.env_kwargs.server_url={args.waa_server} "
+        f"env.env_kwargs.task_id={args.task_id} "
+        f"env.env_kwargs.max_steps=15 "
+        f"env.max_steps=15 "
+        f"env.rollout.n=8 "
+        f"data.train_batch_size=8 "
+        f"data.max_prompt_length=2048 "
+        f"data.max_response_length=512 "
+        f"data.return_raw_chat=True "
+        f"trainer.n_gpus_per_node={args.n_gpus} "
+        f"trainer.nnodes=1 "
+        f"trainer.total_epochs={args.epochs} "
+        f"trainer.logger=['console','wandb'] "
+        f"trainer.project_name=openadapt-waa-rl"
+    )
+
+    print(f"Launching {args.algorithm} training on {args.n_gpus} GPU(s)...")
+    print(f"Model: {args.model}")
+    print(f"WAA server: {args.waa_server}")
+    print(f"Task: {args.task_id}")
+
+    try:
+        result = ssh_run(ip, train_cmd, username=username, stream=True)
+        return result.returncode
+    finally:
+        if args.cleanup and not args.gpu_ip:
+            print(f"Deallocating GPU VM '{gpu_vm_name}'...")
+            vm.deallocate_vm(gpu_vm_name)
+
+
 # =============================================================================
 # Main
 # =============================================================================
@@ -8898,6 +9073,72 @@ def main():
     )
     p_view_pool.set_defaults(func=cmd_view_pool)
 
+    # --- GPU Training Commands ---
+
+    p_gpu_setup = subparsers.add_parser(
+        "gpu-setup",
+        help="Provision a GPU VM and install verl-agent for RL training",
+    )
+    p_gpu_setup.add_argument(
+        "--cloud", choices=["azure", "aws"], default="azure",
+        help="Cloud provider (default: azure)",
+    )
+    p_gpu_setup.add_argument(
+        "--gpu-ip", type=str, default=None,
+        help="Use an existing GPU VM (skip provisioning)",
+    )
+    p_gpu_setup.add_argument(
+        "--dry-run", action="store_true",
+        help="Show what would happen without doing it",
+    )
+    p_gpu_setup.set_defaults(func=cmd_gpu_setup)
+
+    p_gpu_train = subparsers.add_parser(
+        "gpu-train",
+        help="Launch verl-agent training on a GPU VM",
+    )
+    p_gpu_train.add_argument(
+        "--cloud", choices=["azure", "aws"], default="azure",
+        help="Cloud provider (default: azure)",
+    )
+    p_gpu_train.add_argument(
+        "--gpu-ip", type=str, default=None,
+        help="Use an existing GPU VM (skip provisioning)",
+    )
+    p_gpu_train.add_argument(
+        "--waa-server", type=str, default="http://localhost:5001",
+        help="WAA server URL accessible from GPU VM",
+    )
+    p_gpu_train.add_argument(
+        "--task-id", type=str, required=True,
+        help="WAA task UUID to train on",
+    )
+    p_gpu_train.add_argument(
+        "--algorithm", choices=["gigpo", "grpo", "ppo"], default="gigpo",
+        help="RL algorithm (default: gigpo)",
+    )
+    p_gpu_train.add_argument(
+        "--model", type=str, default="Qwen/Qwen2.5-VL-3B-Instruct",
+        help="Model to train",
+    )
+    p_gpu_train.add_argument(
+        "--n-gpus", type=int, default=2,
+        help="Number of GPUs (default: 2)",
+    )
+    p_gpu_train.add_argument(
+        "--epochs", type=int, default=100,
+        help="Training epochs (default: 100)",
+    )
+    p_gpu_train.add_argument(
+        "--skip-setup", action="store_true",
+        help="Skip setup (VM already configured)",
+    )
+    p_gpu_train.add_argument(
+        "--cleanup", action="store_true",
+        help="Deallocate GPU VM after training",
+    )
+    p_gpu_train.set_defaults(func=cmd_gpu_train)
+
     args = parser.parse_args()
 
     # Allow --resource-group to override the module-level constant

openadapt_evals/infrastructure/aws_vm.py

@@ -45,6 +45,18 @@
     ("c5.metal", 4.080),
     ("m5a.xlarge", 0.172),  # Non-KVM fallback (won't run QEMU, for testing only)
 ]
+
+# GPU instance types for verl-agent RL training.
+# verl-agent requires 2+ GPUs for distributed VLM training.
+# p3.8xlarge: 4x V100 16GB NVLink — recommended for Qwen2.5-VL-3B training.
+# g5.12xlarge: 4x A10G 24GB — budget option (no NVLink).
+# p3.2xlarge: 1x V100 16GB — single-GPU baseline (tight for 3B).
+GPU_INSTANCE_TYPE_FALLBACKS = [
+    ("p3.8xlarge", 12.24),
+    ("g5.12xlarge", 7.48),
+    ("p3.2xlarge", 3.06),
+]
+
 # Regions to try in order of preference
 AWS_REGIONS = ["us-east-1", "us-west-2", "us-east-2", "eu-west-1"]
 
@@ -526,11 +538,17 @@ def set_auto_shutdown(self, name: str, hours: int = 4) -> bool:
             logger.warning(f"Failed to set auto-shutdown for {name}: {e}")
             return False
 
-    def find_available_size_and_region(self) -> tuple[str, str, float]:
+    def find_available_size_and_region(
+        self, gpu: bool = False,
+    ) -> tuple[str, str, float]:
         """Find a working EC2 instance type and region.
 
         Checks instance type availability in each region.
 
+        Args:
+            gpu: If True, try GPU instances (for verl-agent training).
+                Otherwise try CPU/metal instances (for WAA evaluation).
+
         Returns:
             Tuple of (instance_type, region, cost_per_hour).
 
@@ -539,7 +557,8 @@ def find_available_size_and_region(self) -> tuple[str, str, float]:
         """
         import boto3
 
-        for instance_type, cost in INSTANCE_TYPE_FALLBACKS:
+        fallbacks = GPU_INSTANCE_TYPE_FALLBACKS if gpu else INSTANCE_TYPE_FALLBACKS
+        for instance_type, cost in fallbacks:
             for region in AWS_REGIONS:
                 try:
                     ec2 = boto3.client("ec2", region_name=region)

openadapt_evals/infrastructure/azure_vm.py

@@ -65,6 +65,18 @@
     ("Standard_D8ds_v4", 0.38),
     ("Standard_D8as_v5", 0.34),
 ]
+
+# GPU VM sizes for verl-agent RL training.
+# verl-agent requires 2+ GPUs for distributed VLM training.
+# NC48ads_A100_v4: 2x A100 80GB — recommended for Qwen2.5-VL-3B/7B.
+# NC24ads_A100_v4: 1x A100 80GB — single-GPU baseline.
+# NC12s_v3: 1x V100 16GB — budget option (tight for 3B).
+GPU_VM_SIZE_FALLBACKS = [
+    ("Standard_NC48ads_A100_v4", 11.04),
+    ("Standard_NC24ads_A100_v4", 5.52),
+    ("Standard_NC12s_v3", 1.50),
+]
+
 VM_REGIONS = ["centralus", "eastus", "westus2", "eastus2"]
 
 # Ubuntu 22.04 LTS image reference for Azure SDK
@@ -522,19 +534,25 @@ def _cli_set_auto_shutdown(self, name: str, hours: int) -> bool:
         )
         return result.returncode == 0
 
-    def find_available_size_and_region(self) -> tuple[str, str, float]:
+    def find_available_size_and_region(
+        self, gpu: bool = False,
+    ) -> tuple[str, str, float]:
         """Find a working VM size and region by creating a test VM.
 
         Tries size/region combinations until one succeeds, then cleans up
         the test VM.
 
+        Args:
+            gpu: If True, try GPU sizes (for verl-agent training).
+                Otherwise try CPU sizes (for WAA evaluation).
+
         Returns:
             Tuple of (vm_size, region, cost_per_hour).
 
         Raises:
             RuntimeError: If no available size/region found.
         """
-        sizes_to_try = VM_SIZE_FALLBACKS
+        sizes_to_try = GPU_VM_SIZE_FALLBACKS if gpu else VM_SIZE_FALLBACKS
 
         test_vm_to_cleanup = None
         try:

openadapt_evals/infrastructure/vm_provider.py

@@ -74,9 +74,14 @@ def set_auto_shutdown(self, name: str, hours: int = 4) -> bool:
         """Set auto-shutdown policy on a VM."""
         ...
 
-    def find_available_size_and_region(self) -> tuple[str, str, float]:
+    def find_available_size_and_region(
+        self, gpu: bool = False,
+    ) -> tuple[str, str, float]:
         """Find a working VM size and region.
 
+        Args:
+            gpu: If True, try GPU sizes for RL training.
+
         Returns:
             Tuple of (vm_size, region, cost_per_hour).
         """